tracingapi

type ChatCompletionRecorder = SpanRecorder[openai.ChatCompletionRequest, openai.ChatCompletionResponse, openai.ChatCompletionResponseChunk]

type ChatCompletionSpan = Span[openai.ChatCompletionResponse, openai.ChatCompletionResponseChunk]

type ChatCompletionTracer = RequestTracer[openai.ChatCompletionRequest, openai.ChatCompletionResponse, openai.ChatCompletionResponseChunk]

type CompletionRecorder = SpanRecorder[openai.CompletionRequest, openai.CompletionResponse, openai.CompletionResponse]

type CompletionSpan = Span[openai.CompletionResponse, openai.CompletionResponse]

type CompletionTracer = RequestTracer[openai.CompletionRequest, openai.CompletionResponse, openai.CompletionResponse]

type EmbeddingsRecorder = SpanRecorder[openai.EmbeddingRequest, openai.EmbeddingResponse, struct{}]

type EmbeddingsSpan = Span[openai.EmbeddingResponse, struct{}]

type EmbeddingsTracer = RequestTracer[openai.EmbeddingRequest, openai.EmbeddingResponse, struct{}]

type ImageGenerationRecorder = SpanRecorder[openai.ImageGenerationRequest, openai.ImageGenerationResponse, struct{}]

type ImageGenerationSpan = Span[openai.ImageGenerationResponse, struct{}]

type ImageGenerationTracer = RequestTracer[openai.ImageGenerationRequest, openai.ImageGenerationResponse, struct{}]

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 MessageRecorder = SpanRecorder[anthropicschema.MessagesRequest, anthropicschema.MessagesResponse, anthropicschema.MessagesStreamChunk]

type MessageSpan = Span[anthropicschema.MessagesResponse, anthropicschema.MessagesStreamChunk]

type MessageTracer = RequestTracer[anthropicschema.MessagesRequest, anthropicschema.MessagesResponse, anthropicschema.MessagesStreamChunk]

type NoopChatCompletionTracer = NoopTracer[openai.ChatCompletionRequest, openai.ChatCompletionResponse, openai.ChatCompletionResponseChunk]

type NoopChunkRecorder[ChunkT any] struct{}

type NoopCompletionTracer = NoopTracer[openai.CompletionRequest, openai.CompletionResponse, openai.CompletionResponse]

type NoopEmbeddingsTracer = NoopTracer[openai.EmbeddingRequest, openai.EmbeddingResponse, struct{}]

type NoopImageGenerationTracer = NoopTracer[openai.ImageGenerationRequest, openai.ImageGenerationResponse, struct{}]

type NoopMCPTracer struct{}

type NoopMessageTracer = NoopTracer[anthropicschema.MessagesRequest, anthropicschema.MessagesResponse, anthropicschema.MessagesStreamChunk]

type NoopRerankTracer = NoopTracer[cohere.RerankV2Request, cohere.RerankV2Response, struct{}]

type NoopResponsesTracer = NoopTracer[openai.ResponseRequest, openai.Response, openai.ResponseStreamEventUnion]

type NoopTracer[ReqT any, RespT any, RespChunkT any] struct{}

type NoopTracing struct{}

type RequestTracer[ReqT any, RespT any, RespChunkT any] 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 span will have its context written to these headers unless NoopTracing is used.
	//   - req: The typed 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, carrier propagation.TextMapCarrier, req *ReqT, body []byte) Span[RespT, RespChunkT]
}

type RerankRecorder = SpanRecorder[cohere.RerankV2Request, cohere.RerankV2Response, struct{}]

type RerankSpan = Span[cohere.RerankV2Response, struct{}]

type RerankTracer = RequestTracer[cohere.RerankV2Request, cohere.RerankV2Response, struct{}]

type ResponsesRecorder = SpanRecorder[openai.ResponseRequest, openai.Response, openai.ResponseStreamEventUnion]

type ResponsesSpan = Span[openai.Response, openai.ResponseStreamEventUnion]

type ResponsesTracer = RequestTracer[openai.ResponseRequest, openai.Response, openai.ResponseStreamEventUnion]

type Span[RespT any, RespChunkT any] interface {
	// RecordResponseChunk records streaming response chunks. Implementations that do not support streaming should provide a no-op implementation.
	RecordResponseChunk(resp *RespChunkT)
	// RecordResponse records the response attributes to the span.
	RecordResponse(resp *RespT)
	// EndSpanOnError finalizes and ends the span with an error status.
	EndSpanOnError(statusCode int, body []byte)
	// EndSpan finalizes and ends the span.
	EndSpan()
}

type SpanRecorder[ReqT any, RespT any, RespChunkT any] interface {
	SpanResponseRecorder[RespT, RespChunkT]
	// StartParams returns the name and options to start the span with.
	//
	// Parameters:
	//   - req: contains the typed request.
	//   - body: contains the complete request body.
	//
	// Note: Avoid expensive data conversions since the span might not be sampled.
	StartParams(req *ReqT, body []byte) (spanName string, opts []trace.SpanStartOption)
	// RecordRequest records request attributes to the span.
	RecordRequest(span trace.Span, req *ReqT, body []byte)
}

type SpanResponseRecorder[RespT, RespChunkT any] interface {
	// RecordResponse records response attributes to the span.
	RecordResponse(span trace.Span, resp *RespT)
	// RecordResponseOnError ends recording the span with an error status.
	RecordResponseOnError(span trace.Span, statusCode int, body []byte)
	// RecordResponseChunks records response chunk attributes to the span for streaming response.
	RecordResponseChunks(span trace.Span, chunks []*RespChunkT)
}

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
	// ResponsesTracer creates spans for OpenAI responses requests on /v1/responses endpoint.
	ResponsesTracer() ResponsesTracer
	// RerankTracer creates spans for rerank requests.
	RerankTracer() RerankTracer
	// MessageTracer creates spans for Anthropic messages requests.
	MessageTracer() MessageTracer
	// MCPTracer creates spans for MCP requests.
	MCPTracer() MCPTracer
	// Shutdown shuts down the tracer, flushing any buffered spans.
	Shutdown(context.Context) error
}