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 NoopTracer 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 NoopChatCompletionTracer struct{}

type NoopTracing struct{}

type Tracing interface {
	ChatCompletionTracer() ChatCompletionTracer
	Shutdown(context.Context) error
}

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