api

type ChatCompletionRecorder interface {
	// StartParams returns the name and options to start the span with.
	//
	// Parameters:
	//   - req: contains the completion request
	//   - body: contains the complete request body.
	//
	// Note: Do not do any expensive data conversions as the span might not be
	// sampled.
	StartParams(req *openai.ChatCompletionRequest, body []byte) (spanName string, opts []trace.SpanStartOption)

	// RecordRequest records request attributes to the span.
	//
	// Parameters:
	//   - req: contains the completion request
	//   - body: contains the complete request body.
	RecordRequest(span trace.Span, req *openai.ChatCompletionRequest, body []byte)

	// RecordResponseChunks records response chunk attributes to the span for streaming response.
	RecordResponseChunks(span trace.Span, chunks []*openai.ChatCompletionResponseChunk)

	// RecordResponse records response attributes to the span for non-streaming response.
	RecordResponse(span trace.Span, resp *openai.ChatCompletionResponse)

	// RecordResponseOnError ends recording the span with an error status.
	RecordResponseOnError(span trace.Span, statusCode int, body []byte)
}

type ChatCompletionSpan interface {
	// RecordResponseChunk records the response chunk attributes to the span for streaming response.
	RecordResponseChunk(resp *openai.ChatCompletionResponseChunk)

	// RecordResponse records the response attributes to the span for non-streaming response.
	RecordResponse(resp *openai.ChatCompletionResponse)

	// EndSpanOnError finalizes and ends the span with an error status.
	EndSpanOnError(statusCode int, body []byte)

	// EndSpan finalizes and ends the span.
	EndSpan()
}

type ChatCompletionTracer interface {
	// StartSpanAndInjectHeaders starts a span and injects trace context into
	// the header mutation.
	//
	// Parameters:
	//   - ctx: might include a parent span context.
	//   - headers: Incoming HTTP headers used to extract parent trace context.
	//   - headerMutation: The new LLM Span will have its context written to
	//     these headers unless NoopTracing is used.
	//   - req: The OpenAI chat completion request. Used to detect streaming
	//     and record request attributes.
	//
	// Returns nil unless the span is sampled.
	StartSpanAndInjectHeaders(ctx context.Context, headers map[string]string, headerMutation *extprocv3.HeaderMutation, req *openai.ChatCompletionRequest, body []byte) ChatCompletionSpan
}

type CompletionRecorder interface {
	// StartParams returns the name and options to start the span with.
	//
	// Parameters:
	//   - req: contains the completion request
	//   - body: contains the complete request body.
	//
	// Note: Do not do any expensive data conversions as the span might not be
	// sampled.
	StartParams(req *openai.CompletionRequest, body []byte) (spanName string, opts []trace.SpanStartOption)

	// RecordRequest records request attributes to the span.
	//
	// Parameters:
	//   - req: contains the completion request
	//   - body: contains the complete request body.
	RecordRequest(span trace.Span, req *openai.CompletionRequest, body []byte)

	// RecordResponseChunks records response chunk attributes to the span for streaming response.
	// Note: Completion chunks are full CompletionResponse objects, not deltas like chat.
	RecordResponseChunks(span trace.Span, chunks []*openai.CompletionResponse)

	// RecordResponse records response attributes to the span for non-streaming response.
	RecordResponse(span trace.Span, resp *openai.CompletionResponse)

	// RecordResponseOnError ends recording the span with an error status.
	RecordResponseOnError(span trace.Span, statusCode int, body []byte)
}

type CompletionSpan interface {
	// RecordResponseChunk records the response chunk attributes to the span for streaming response.
	// Note: Unlike chat completions, completion streaming uses full CompletionResponse objects, not deltas.
	RecordResponseChunk(resp *openai.CompletionResponse)

	// RecordResponse records the response attributes to the span for non-streaming response.
	RecordResponse(resp *openai.CompletionResponse)

	// EndSpanOnError finalizes and ends the span with an error status.
	EndSpanOnError(statusCode int, body []byte)

	// EndSpan finalizes and ends the span.
	EndSpan()
}

type CompletionTracer interface {
	// StartSpanAndInjectHeaders starts a span and injects trace context into
	// the header mutation.
	//
	// Parameters:
	//   - ctx: might include a parent span context.
	//   - headers: Incoming HTTP headers used to extract parent trace context.
	//   - headerMutation: The new LLM Span will have its context written to
	//     these headers unless NoopTracing is used.
	//   - req: The OpenAI completion request. Used to detect streaming
	//     and record request attributes.
	//   - body: contains the original raw request body as a byte slice.
	//
	// Returns nil unless the span is sampled.
	StartSpanAndInjectHeaders(ctx context.Context, headers map[string]string, headerMutation *extprocv3.HeaderMutation, req *openai.CompletionRequest, body []byte) CompletionSpan
}

type EmbeddingsRecorder interface {
	// StartParams returns the name and options to start the span with.
	//
	// Parameters:
	//   - req: contains the embeddings request
	//   - body: contains the complete request body.
	//
	// Note: Do not do any expensive data conversions as the span might not be
	// sampled.
	StartParams(req *openai.EmbeddingRequest, body []byte) (spanName string, opts []trace.SpanStartOption)

	// RecordRequest records request attributes to the span.
	//
	// Parameters:
	//   - req: contains the embeddings request
	//   - body: contains the complete request body.
	RecordRequest(span trace.Span, req *openai.EmbeddingRequest, body []byte)

	// RecordResponse records response attributes to the span.
	RecordResponse(span trace.Span, resp *openai.EmbeddingResponse)

	// RecordResponseOnError ends recording the span with an error status.
	RecordResponseOnError(span trace.Span, statusCode int, body []byte)
}

type EmbeddingsSpan interface {
	// RecordResponse records the response attributes to the span.
	RecordResponse(resp *openai.EmbeddingResponse)

	// EndSpanOnError finalizes and ends the span with an error status.
	EndSpanOnError(statusCode int, body []byte)

	// EndSpan finalizes and ends the span.
	EndSpan()
}

type EmbeddingsTracer interface {
	// StartSpanAndInjectHeaders starts a span and injects trace context into
	// the header mutation.
	//
	// Parameters:
	//   - ctx: might include a parent span context.
	//   - headers: Incoming HTTP headers used to extract parent trace context.
	//   - headerMutation: The new Embeddings Span will have its context
	//     written to these headers unless NoopTracing is used.
	//   - req: The OpenAI embeddings request. Used to record request attributes.
	//   - body: contains the original raw request body as a byte slice.
	//
	// Returns nil unless the span is sampled.
	StartSpanAndInjectHeaders(ctx context.Context, headers map[string]string, headerMutation *extprocv3.HeaderMutation, req *openai.EmbeddingRequest, body []byte) EmbeddingsSpan
}

type ImageGenerationRecorder interface {
	// StartParams returns the name and options to start the span with.
	//
	// Parameters:
	//   - req: contains the image generation request
	//   - body: contains the complete request body.
	//
	// Note: Do not do any expensive data conversions as the span might not be
	// sampled.
	StartParams(req *openaisdk.ImageGenerateParams, body []byte) (spanName string, opts []trace.SpanStartOption)

	// RecordRequest records request attributes to the span.
	//
	// Parameters:
	//   - req: contains the image generation request
	//   - body: contains the complete request body.
	RecordRequest(span trace.Span, req *openaisdk.ImageGenerateParams, body []byte)

	// RecordResponse records response attributes to the span.
	RecordResponse(span trace.Span, resp *openaisdk.ImagesResponse)

	// RecordResponseOnError ends recording the span with an error status.
	RecordResponseOnError(span trace.Span, statusCode int, body []byte)
}

type ImageGenerationSpan interface {
	// RecordResponse records the response attributes to the span.
	RecordResponse(resp *openaisdk.ImagesResponse)

	// EndSpanOnError finalizes and ends the span with an error status.
	EndSpanOnError(statusCode int, body []byte)

	// EndSpan finalizes and ends the span.
	EndSpan()
}

type ImageGenerationTracer interface {
	// StartSpanAndInjectHeaders starts a span and injects trace context into
	// the header mutation.
	//
	// Parameters:
	//   - ctx: might include a parent span context.
	//   - headers: Incoming HTTP headers used to extract parent trace context.
	//   - headerMutation: The new LLM Span will have its context written to
	//     these headers unless NoopTracer is used.
	//   - req: The OpenAI image generation request. Used to record request attributes.
	//
	// Returns nil unless the span is sampled.
	StartSpanAndInjectHeaders(ctx context.Context, headers map[string]string, headerMutation *extprocv3.HeaderMutation, req *openaisdk.ImageGenerateParams, body []byte) ImageGenerationSpan
}

type MCPSpan interface {
	// RecordRouteToBackend records the backend that was routed to.
	RecordRouteToBackend(backend string, session string, isNew bool)
	// EndSpan finalizes and ends the span.
	EndSpan()
	// EndSpanOnError finalizes and ends the span with an error status.
	EndSpanOnError(errType string, err error)
}

type MCPTracer interface {
	// StartSpanAndInjectMeta starts a span and injects trace context into
	// the _meta mutation.
	//
	// Parameters:
	//   - ctx: might include a parent span context.
	//   - req: Incoming MCP request message.
	//   - param: Incoming MCP parameter used to extract parent trace context.
	//   - headers: Request HTTP request headers.
	//
	// Returns nil unless the span is sampled.
	StartSpanAndInjectMeta(ctx context.Context, req *jsonrpc.Request, param mcp.Params, headers http.Header) MCPSpan
}

type NoopChatCompletionTracer struct{}

type NoopCompletionTracer struct{}

type NoopEmbeddingsTracer struct{}

type NoopImageGenerationTracer struct{}

type NoopMCPTracer struct{}

type NoopTracing struct{}

type Tracing interface {
	// ChatCompletionTracer creates spans for OpenAI chat completion requests on /chat/completions endpoint.
	ChatCompletionTracer() ChatCompletionTracer
	// ImageGenerationTracer creates spans for OpenAI image generation requests.
	ImageGenerationTracer() ImageGenerationTracer
	// CompletionTracer creates spans for OpenAI completion requests on /completions endpoint.
	CompletionTracer() CompletionTracer
	// EmbeddingsTracer creates spans for OpenAI embeddings requests on /embeddings endpoint.
	EmbeddingsTracer() EmbeddingsTracer
	// MCPTracer creates spans for MCP requests.
	MCPTracer() MCPTracer
	// Shutdown shuts down the tracer, flushing any buffered spans.
	Shutdown(context.Context) error
}

type TracingConfig struct {
	Tracer                  trace.Tracer
	Propagator              propagation.TextMapPropagator
	ChatCompletionRecorder  ChatCompletionRecorder
	CompletionRecorder      CompletionRecorder
	ImageGenerationRecorder ImageGenerationRecorder
	EmbeddingsRecorder      EmbeddingsRecorder
}