llmobs

type APMSpan interface {
	// Finish finishes the span with the given configuration.
	Finish(cfg FinishAPMSpanConfig)
	// AddLink adds a span link to this span.
	AddLink(link SpanLink)
	// SpanID returns the span ID.
	SpanID() string
	// TraceID returns the trace ID.
	TraceID() string
	// SetBaggageItem sets a baggage item on the span.
	SetBaggageItem(key string, value string)
	// BaggageItem returns the baggage item value for the given key.
	BaggageItem(key string) string
}

type EmbeddedDocument struct {
	// Text is the text content of the document.
	Text string `json:"text"`
	// Name is the name or title of the document.
	Name string `json:"name,omitempty"`
	// Score is the relevance score of the document (typically 0.0-1.0).
	Score float64 `json:"score,omitempty"`
	// ID is the unique identifier of the document.
	ID string `json:"id,omitempty"`
}

type EvaluationConfig struct {
	// Method 1: Direct span/trace ID join
	// SpanID is the span ID to evaluate.
	SpanID string
	// TraceID is the trace ID to evaluate.
	TraceID string

	// Method 2: Tag-based join
	// TagKey is the tag key to search for spans.
	TagKey string
	// TagValue is the tag value to match for spans.
	TagValue string

	// Required fields
	// Label is the name of the evaluation metric.
	Label string

	// Value fields (exactly one must be provided)
	// CategoricalValue is the categorical value of the evaluation metric.
	CategoricalValue *string
	// ScoreValue is the score value of the evaluation metric.
	ScoreValue *float64
	// BooleanValue is the boolean value of the evaluation metric.
	BooleanValue *bool

	// Optional fields
	// Tags are optional string key-value pairs to tag the evaluation metric.
	Tags []string
	// MLApp is the ML application name. If empty, uses the global config.
	MLApp string
	// TimestampMS is the timestamp in milliseconds. If zero, uses current time.
	TimestampMS int64
}

type ExperimentInfo struct {
	ID           string
	RunID        string
	RunIteration int
	ProjectID    string
}

type FinishAPMSpanConfig struct {
	// FinishTime is the finish time for the span.
	FinishTime time.Time
	// Error is an error to set on the span when finishing.
	Error error
}

type FinishSpanConfig struct {
	// FinishTime sets a custom finish time for the span. If zero, uses current time.
	FinishTime time.Time
	// Error sets an error on the span when finishing.
	Error error
}

type LLMMessage struct {
	// Role is the role of the message sender (e.g., "user", "assistant", "system").
	Role string `json:"role"`
	// Content is the text content of the message.
	Content string `json:"content"`
	// ToolCalls are the tool calls made in this message.
	ToolCalls []ToolCall `json:"tool_calls,omitempty"`
	// ToolResults are the results of tool calls in this message.
	ToolResults []ToolResult `json:"tool_results,omitempty"`
}

type LLMObs struct {
	// Config contains the LLMObs configuration.
	Config *config.Config
	// Transport handles sending data to the Datadog backend.
	Transport *transport.Transport
	// Tracer is the underlying APM tracer.
	Tracer Tracer
	// contains filtered or unexported fields
}

type Prompt struct {
	// ID is the unique identifier for the prompt within the ML app.
	ID string `json:"id,omitempty"`
	// Version is the version of the prompt.
	Version string `json:"version,omitempty"`
	// Label is the deployment label (e.g., "production", "staging").
	Label string `json:"label,omitempty"`
	// Template is the prompt template string.
	// Mutually exclusive with ChatTemplate; if both are set, Template is dropped and ChatTemplate is used.
	Template string `json:"template,omitempty"`
	// ChatTemplate is a list of messages forming the prompt.
	// Mutually exclusive with Template; if both are set, Template is dropped and ChatTemplate is used.
	ChatTemplate []LLMMessage `json:"chat_template,omitempty"`
	// Variables contains the variables used in the prompt template.
	Variables map[string]string `json:"variables,omitempty"`
	// Tags contains custom tags for the prompt.
	Tags map[string]string `json:"tags,omitempty"`
	// RAGContextVariables specifies which variables contain RAG context.
	RAGContextVariables []string `json:"_dd_context_variable_keys,omitempty"`
	// RAGQueryVariables specifies which variables contain RAG queries.
	RAGQueryVariables []string `json:"_dd_query_variable_keys,omitempty"`
}

type PropagatedLLMSpan struct {
	// MLApp is the ML application name.
	MLApp string
	// TraceID is the LLMObs trace ID.
	TraceID string
	// SpanID is the span ID.
	SpanID string
	// SessionID is the session ID.
	SessionID string
}

type RetrievedDocument struct {
	// Text is the text content of the document.
	Text string `json:"text"`
	// Name is the name or title of the document.
	Name string `json:"name,omitempty"`
	// Score is the relevance score of the document (typically 0.0-1.0).
	Score float64 `json:"score,omitempty"`
	// ID is the unique identifier of the document.
	ID string `json:"id,omitempty"`
}

type Span struct {
	// contains filtered or unexported fields
}

type SpanAnnotations struct {
	// InputText is the text input for the span.
	InputText string
	// InputMessages are the input messages for LLM spans.
	InputMessages []LLMMessage
	// InputEmbeddedDocs are the input documents for embedding spans.
	InputEmbeddedDocs []EmbeddedDocument

	// OutputText is the text output for the span.
	OutputText string
	// OutputMessages are the output messages for LLM spans.
	OutputMessages []LLMMessage
	// OutputRetrievedDocs are the output documents for retrieval spans.
	OutputRetrievedDocs []RetrievedDocument

	// ExperimentInput is the input data for experiment spans.
	ExperimentInput any
	// ExperimentOutput is the output data for experiment spans.
	ExperimentOutput any
	// ExperimentExpectedOutput is the expected output for experiment spans.
	ExperimentExpectedOutput any

	// Prompt is the prompt information for LLM spans.
	Prompt *Prompt
	// ToolDefinitions are the tool definitions for LLM spans.
	ToolDefinitions []ToolDefinition

	// Intent is a description of a reason for calling an MCP tool on tool spans
	Intent string

	// AgentManifest is the agent manifest for agent spans.
	AgentManifest string

	// Metadata contains arbitrary metadata key-value pairs.
	Metadata map[string]any
	// Metrics contains numeric metrics key-value pairs.
	Metrics map[string]float64
	// Tags contains string tags key-value pairs.
	Tags map[string]string
	// CostTags contains tag keys to propagate to LLMObs cost and token metrics.
	// Each key must reference a tag already present on the span.
	CostTags []string
}

type SpanKind string

type SpanLink = transport.SpanLink

type StartAPMSpanConfig struct {
	// SpanType is the type of the APM span.
	SpanType string
	// StartTime is the start time for the span.
	StartTime time.Time
}

type StartSpanConfig struct {
	// SessionID sets the session ID for the span.
	SessionID string
	// ModelName sets the model name for LLM and embedding spans.
	ModelName string
	// ModelProvider sets the model provider for LLM and embedding spans.
	ModelProvider string
	// MLApp sets the ML application name for the span.
	MLApp string
	// StartTime sets a custom start time for the span. If zero, uses current time.
	StartTime time.Time
	// Name of the tracing integration.
	Integration string
}

type ToolCall struct {
	// Name is the name of the tool being called.
	Name string `json:"name"`
	// Arguments are the JSON-encoded arguments passed to the tool.
	Arguments json.RawMessage `json:"arguments"`
	// ToolID is the unique identifier for this tool call.
	ToolID string `json:"tool_id,omitempty"`
	// Type is the type of the tool call.
	Type string `json:"type,omitempty"`
}

type ToolDefinition struct {
	// Name is the name of the tool.
	Name string `json:"name"`
	// Description is the description of what the tool does.
	Description string `json:"description,omitempty"`
	// ToolVersion is the version of the tool.
	ToolVersion string `json:"version,omitempty"`
	// Schema is the JSON schema defining the tool's parameters.
	Schema json.RawMessage `json:"schema,omitempty"`
}

type ToolResult struct {
	// Result is the result returned by the tool.
	Result any `json:"result"`
	// Name is the name of the tool that was called.
	Name string `json:"name,omitempty"`
	// ToolID is the unique identifier for the tool call this result corresponds to.
	ToolID string `json:"tool_id,omitempty"`
	// Type is the type of the tool result.
	Type string `json:"type,omitempty"`
}

type Tracer interface {
	// StartSpan starts a new APM span with the given name and configuration.
	StartSpan(ctx context.Context, name string, cfg StartAPMSpanConfig) (APMSpan, context.Context)
}