llm

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

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

type CompletionTokensDetails struct {
	AudioTokens              int
	ReasoningTokens          int
	AcceptedPredictionTokens int
	RejectedPredictionTokens int
}

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

type FewShotExample struct {
	User      string
	Assistant string
}

type LLMDetails struct {
	RequestID               string
	Provider                string
	BaseURL                 string
	Model                   string
	Input                   string
	SystemPrompt            string
	N                       int
	MaxTokens               int
	EstimatedMaxOutputChars int
	FilterEmoji             bool
	RequestedOutputFormat   string
	AppliedResponseFormat   string
	ResponseFormatApplied   bool
	ResponseID              string
	Object                  string
	Created                 int64
	SystemFingerprint       string
	PromptFilterResultsJSON string
	ServiceTierJSON         string
	ChoicesCount            int
	Choices                 []QueryChoice
	Usage                   *TokenUsage
	UsageRawJSON            string
	ChoicesRawJSON          string
	RawResponseJSON         string
}

type LLMHandler interface {
	Query(text, model string) (string, error)

	QueryWithOptions(text string, options *QueryOptions) (*QueryResponse, error)

	QueryStream(text string, options *QueryOptions, callback func(segment string, isComplete bool) error) (*QueryResponse, error)

	Provider() string

	Interrupt()

	ResetMemory()

	SummarizeMemory(model string) (string, error)

	SetMaxMemoryMessages(n int)

	GetMaxMemoryMessages() int
}

type LLMOptions struct {
	Provider        string
	ApiKey          string
	BaseURL         string
	SystemPrompt    string
	FewShotExamples []FewShotExample
	Logger          *zap.Logger // Logger is optional; used for async memory summarization warnings and diagnostics.
}

type LLMProvider string

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

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

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

type PromptTokensDetails struct {
	AudioTokens  int
	CachedTokens int
}

type QueryChoice struct {
	Index        int
	Content      string
	FinishReason string
}

type QueryExpansion struct {
	Original string
	Expanded string
	Terms    []string
	Debug    map[string]any
}

type QueryOptions struct {
	Model                string
	N                    int
	MaxTokens            int
	Temperature          float32
	TopP                 float32
	LogitBias            map[string]int
	FilterEmoji          bool
	EnableJSONOutput     bool
	OutputFormat         string
	EmotionalTone        bool   // EmotionalTone, when true, appends a short instruction so replies read warmer and more human (still factual).
	EnableQueryExpansion bool   // EnableQueryExpansion enables automatic query expansion using LLM
	ExpansionMaxTerms    int    // ExpansionMaxTerms maximum number of expansion terms
	ExpansionSeparator   string // ExpansionSeparator separator for expanded terms

	// EnableQueryRewrite rewrites the user message with a stateless LLM call before expansion/main query.
	EnableQueryRewrite bool
	// QueryRewriteModel overrides the model for the rewrite call only (empty = use Model, then handler default).
	QueryRewriteModel string
	// QueryRewriteInstruction is appended to the rewrite prompt as extra constraints.
	QueryRewriteInstruction string
	// EnableSelfQueryJSONOutput requests strict JSON object replies (response_format json_object on OpenaiHandler).
	// SelfQueryExtractor sets this by default; other handlers may ignore it and still return parseable text.
	EnableSelfQueryJSONOutput bool
}

type QueryResponse struct {
	Provider string
	Model    string
	Choices  []QueryChoice
	Usage    *TokenUsage
	// Expansion contains query expansion results if enabled
	Expansion *QueryExpansion
	// Rewrite contains query rewrite results if enabled
	Rewrite *QueryRewrite
}

type QueryRewrite struct {
	Original  string
	Rewritten string
}

type SelfQueryExtractor struct {
	LLM LLMHandler

	AllowedFields []string
}

type SelfQueryFilterSpec struct {
	Namespace string   `json:"namespace,omitempty"`
	Source    string   `json:"source,omitempty"`
	DocType   string   `json:"doc_type,omitempty"`
	Location  string   `json:"location,omitempty"`
	Years     []string `json:"years,omitempty"`
	Dates     []string `json:"dates,omitempty"`
	TagsAny   []string `json:"tags_any,omitempty"`
}

type SelfQueryOptions struct {
	Model string

	// AllowedFields restricts which filter keys the model may use (prompt hint).
	AllowedFields []string

	// MaxJSONChars caps input length before JSON extraction (0 = default 16000).
	MaxJSONChars int

	// UsePlainQuery uses Query instead of QueryWithOptions with JSON object mode.
	// Set true for backends that handle JSON poorly in structured-output mode.
	UsePlainQuery bool
}

type SelfQueryResult struct {
	Query   string
	Filters map[string]any

	Spec SelfQueryFilterSpec
	Raw  string
}

type TokenUsage struct {
	PromptTokens            int
	CompletionTokens        int
	TotalTokens             int
	PromptTokensDetails     *PromptTokensDetails
	CompletionTokensDetails *CompletionTokensDetails
}