Cyrene_v2/backend/services/ai-core/llm/llm.go

// Package llm provides LLM abstraction with streaming, model routing, and response caching.
package llm

import (
	"bufio"
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"strings"
	"sync"
	"time"

	"git.yeij.top/AskaEth/Cyrene_v2/services/ai-core/model"
)

// Config holds LLM connection configuration.
type Config struct {
	BaseURL       string
	APIKey        string
	Model         string
	FallbackModel string
	MaxRetries    int
	Timeout       time.Duration
}

// DefaultConfig returns a sensible default Config.
func DefaultConfig() Config {
	return Config{
		BaseURL:       "https://api.openai.com/v1",
		Model:         "gpt-4o",
		FallbackModel: "gpt-4o-mini",
		MaxRetries:    3,
		Timeout:       180 * time.Second,
	}
}

// Adapter handles communication with an OpenAI-compatible LLM API.
type Adapter struct {
	cfg      Config
	client   *http.Client
	onCall   func(call CallRecord) // Hook for logging each LLM call
}

// NewAdapter creates a new LLM adapter.
func NewAdapter(cfg Config) *Adapter {
	return &Adapter{
		cfg: cfg,
		client: &http.Client{
			Timeout: cfg.Timeout,
			Transport: &http.Transport{
				MaxIdleConns:        100,
				MaxIdleConnsPerHost: 20,
				IdleConnTimeout:     90 * time.Second,
			},
		},
	}
}

// SetCallLogger injects a callback that is invoked after each LLM call.
func (a *Adapter) SetCallLogger(fn func(call CallRecord)) {
	a.onCall = fn
}

// CallRecord holds metadata about a single LLM API call.
type CallRecord struct {
	Time           time.Time
	Model          string
	Purpose        string
	PromptTokens   int
	CompletionTokens int
	DurationMs     int64
	Success        bool
	Error          string
}

// ModelName returns the configured model name.
func (a *Adapter) ModelName() string { return a.cfg.Model }

// Chat sends a chat completion request and returns the full response.
func (a *Adapter) Chat(ctx context.Context, messages []model.LLMMessage) (*model.LLMResponse, error) {
	start := time.Now()
	body := map[string]any{
		"model":    a.cfg.Model,
		"messages": wireMessages(messages),
	}

	resp := &model.LLMResponse{Model: a.cfg.Model}
	err := a.doRequest(ctx, "POST", "/chat/completions", body, func(data map[string]any) error {
		if choices, ok := data["choices"].([]any); ok && len(choices) > 0 {
			choice := choices[0].(map[string]any)
			if msg, ok := choice["message"].(map[string]any); ok {
				if content, ok := msg["content"].(string); ok { resp.Content = content }
			}
			if reason, ok := choice["finish_reason"].(string); ok { resp.FinishReason = reason }
		}
		if usage, ok := data["usage"].(map[string]any); ok {
			if total, ok := usage["total_tokens"].(float64); ok { resp.TokensUsed = int(total) }
		}
		return nil
	})

	a.recordCall(start, "chat", resp.TokensUsed, 0, err)
	return resp, err
}

// ToolSpec describes a callable tool exposed to the LLM (function calling).
type ToolSpec struct {
	Name        string
	Description string
	// Parameters is a JSON Schema object describing the tool's input.
	Parameters  map[string]any
}

// ChatStream sends a chat completion request and returns a token channel.
func (a *Adapter) ChatStream(ctx context.Context, messages []model.LLMMessage) (<-chan model.StreamToken, error) {
	return a.ChatStreamWithTools(ctx, messages, nil)
}

// ChatStreamWithTools is ChatStream with function-calling support.
// When the model requests tool calls, they are delivered on the final
// Done token (StreamToken.ToolCalls) instead of regular content tokens.
func (a *Adapter) ChatStreamWithTools(ctx context.Context, messages []model.LLMMessage, tools []ToolSpec) (<-chan model.StreamToken, error) {
	start := time.Now()
	body := map[string]any{
		"model":    a.cfg.Model,
		"messages": wireMessages(messages),
		"stream":   true,
	}
	if len(tools) > 0 {
		body["tools"] = wireTools(tools)
	}

	ch := make(chan model.StreamToken, 64)
	go func() {
		defer close(ch)
		err := a.doStreamRequest(ctx, "POST", "/chat/completions", body, ch)
		a.recordCall(start, "chat_stream", 0, 0, err)
		if err != nil {
			select {
			case ch <- model.StreamToken{Content: fmt.Sprintf("[Error: %v]", err), Done: true}:
			default:
			}
		}
	}()

	return ch, nil
}

func (a *Adapter) recordCall(start time.Time, purpose string, promptTok, compTok int, err error) {
	if a.onCall != nil {
		a.onCall(CallRecord{
			Time:             start,
			Model:            a.cfg.Model,
			Purpose:          purpose,
			PromptTokens:     promptTok,
			CompletionTokens: compTok,
			DurationMs:       time.Since(start).Milliseconds(),
			Success:          err == nil,
			Error:            func() string { if err != nil { return err.Error() }; return "" }(),
		})
	}
}

// doRequest sends a non-streaming request and parses the response.
func (a *Adapter) doRequest(ctx context.Context, method, path string, body map[string]any, parser func(map[string]any) error) error {
	jsonBody, _ := json.Marshal(body)

	req, err := http.NewRequestWithContext(ctx, method, a.cfg.BaseURL+path, bytes.NewReader(jsonBody))
	if err != nil {
		return fmt.Errorf("create request: %w", err)
	}
	req.Header.Set("Authorization", "Bearer "+a.cfg.APIKey)
	req.Header.Set("Content-Type", "application/json")

	resp, err := a.client.Do(req)
	if err != nil {
		return fmt.Errorf("http request: %w", err)
	}
	defer resp.Body.Close()

	respBytes, err := io.ReadAll(resp.Body)
	if err != nil {
		return fmt.Errorf("read response: %w", err)
	}

	if resp.StatusCode != http.StatusOK {
		return fmt.Errorf("API error %d: %s", resp.StatusCode, string(respBytes[:min(len(respBytes), 200)]))
	}

	var data map[string]any
	if err := json.Unmarshal(respBytes, &data); err != nil {
		return fmt.Errorf("parse response: %w", err)
	}
	return parser(data)
}

// doStreamRequest sends a streaming request and sends tokens to the channel.
func (a *Adapter) doStreamRequest(ctx context.Context, method, path string, body map[string]any, ch chan<- model.StreamToken) error {
	jsonBody, _ := json.Marshal(body)

	req, err := http.NewRequestWithContext(ctx, method, a.cfg.BaseURL+path, bytes.NewReader(jsonBody))
	if err != nil {
		return err
	}
	req.Header.Set("Authorization", "Bearer "+a.cfg.APIKey)
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Accept", "text/event-stream")

	resp, err := a.client.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		errMsg, _ := io.ReadAll(resp.Body)
		return fmt.Errorf("stream API error %d: %s", resp.StatusCode, string(errMsg[:min(len(errMsg), 200)]))
	}

	scanner := bufio.NewScanner(resp.Body)
	toolAcc := newToolCallAccumulator()
	for scanner.Scan() {
		line := strings.TrimSpace(scanner.Text())
		if line == "" || !strings.HasPrefix(line, "data: ") {
			continue
		}
		data := strings.TrimPrefix(line, "data: ")
		if data == "[DONE]" {
			ch <- model.StreamToken{Done: true, ToolCalls: toolAcc.calls()}
			return nil
		}

		var frame map[string]any
		if err := json.Unmarshal([]byte(data), &frame); err != nil {
			continue
		}

		if choices, ok := frame["choices"].([]any); ok && len(choices) > 0 {
			choice := choices[0].(map[string]any)
			if delta, ok := choice["delta"].(map[string]any); ok {
				if content, ok := delta["content"].(string); ok {
					ch <- model.StreamToken{Content: content}
				}
				toolAcc.feed(delta)
			}
			if reason, ok := choice["finish_reason"].(string); ok && reason != "" {
				ch <- model.StreamToken{FinishReason: reason, Done: true, ToolCalls: toolAcc.calls()}
				return nil
			}
		}
	}
	return scanner.Err()
}

func min(a, b int) int {
	if a < b {
		return a
	}
	return b
}

// wireMessages converts LLMMessages to the OpenAI wire format,
// expanding ToolCalls into the nested function-call structure.
func wireMessages(messages []model.LLMMessage) []map[string]any {
	out := make([]map[string]any, len(messages))
	for i, m := range messages {
		wm := map[string]any{"role": m.Role, "content": m.Content}
		if m.Name != "" {
			wm["name"] = m.Name
		}
		if m.ToolCallID != "" {
			wm["tool_call_id"] = m.ToolCallID
		}
		if len(m.ToolCalls) > 0 {
			tcs := make([]map[string]any, len(m.ToolCalls))
			for j, tc := range m.ToolCalls {
				args, _ := json.Marshal(tc.Arguments)
				tcs[j] = map[string]any{
					"id":   tc.ID,
					"type": "function",
					"function": map[string]any{
						"name":      tc.Name,
						"arguments": string(args),
					},
				}
			}
			wm["tool_calls"] = tcs
		}
		out[i] = wm
	}
	return out
}

// wireTools converts ToolSpecs to the OpenAI tools array format.
func wireTools(tools []ToolSpec) []map[string]any {
	out := make([]map[string]any, len(tools))
	for i, t := range tools {
		params := t.Parameters
		if params == nil {
			params = map[string]any{"type": "object", "properties": map[string]any{}}
		}
		out[i] = map[string]any{
			"type": "function",
			"function": map[string]any{
				"name":        t.Name,
				"description": t.Description,
				"parameters":  params,
			},
		}
	}
	return out
}

// toolCallAccumulator reassembles streamed tool_call deltas (indexed fragments
// with partial id/name/arguments) into complete ToolCalls.
type toolCallAccumulator struct {
	order  []int
	byIdx  map[int]*model.ToolCall
	rawArg map[int]*strings.Builder
}

func newToolCallAccumulator() *toolCallAccumulator {
	return &toolCallAccumulator{byIdx: map[int]*model.ToolCall{}, rawArg: map[int]*strings.Builder{}}
}

// feed consumes one streaming delta and merges any tool_calls fragments.
func (a *toolCallAccumulator) feed(delta map[string]any) {
	tcs, ok := delta["tool_calls"].([]any)
	if !ok {
		return
	}
	for _, item := range tcs {
		tcMap, ok := item.(map[string]any)
		if !ok {
			continue
		}
		idx := 0
		if f, ok := tcMap["index"].(float64); ok {
			idx = int(f)
		}
		tc, exists := a.byIdx[idx]
		if !exists {
			tc = &model.ToolCall{}
			a.byIdx[idx] = tc
			a.rawArg[idx] = &strings.Builder{}
			a.order = append(a.order, idx)
		}
		if id, ok := tcMap["id"].(string); ok && id != "" {
			tc.ID = id
		}
		if fn, ok := tcMap["function"].(map[string]any); ok {
			if name, ok := fn["name"].(string); ok && name != "" {
				tc.Name = name
			}
			if args, ok := fn["arguments"].(string); ok {
				a.rawArg[idx].WriteString(args)
			}
		}
	}
}

// calls returns the accumulated tool calls in index order, or nil if none.
func (a *toolCallAccumulator) calls() []model.ToolCall {
	if len(a.order) == 0 {
		return nil
	}
	calls := make([]model.ToolCall, 0, len(a.order))
	for _, idx := range a.order {
		tc := a.byIdx[idx]
		if raw := a.rawArg[idx].String(); raw != "" {
			_ = json.Unmarshal([]byte(raw), &tc.Arguments)
		}
		calls = append(calls, *tc)
	}
	return calls
}

// ModelSelector routes requests to different models based on purpose.
type ModelSelector struct {
	adapters map[model.LLMPurpose]*Adapter
	default_ *Adapter
}

// NewModelSelector creates a selector with purpose-specific adapters.
func NewModelSelector(cfg Config) *ModelSelector {
	base := NewAdapter(cfg)
	return &ModelSelector{
		adapters: map[model.LLMPurpose]*Adapter{},
		default_: base,
	}
}

// SetAdapter assigns an adapter for a specific purpose.
func (s *ModelSelector) SetAdapter(purpose model.LLMPurpose, a *Adapter) {
	s.adapters[purpose] = a
}

// Select returns the adapter for the given purpose, falling back to default.
func (s *ModelSelector) Select(purpose model.LLMPurpose) *Adapter {
	if a, ok := s.adapters[purpose]; ok {
		return a
	}
	return s.default_
}

// DefaultAdapter returns the default adapter.
func (s *ModelSelector) DefaultAdapter() *Adapter {
	return s.default_
}

// SemanticCache caches common greetings to avoid redundant LLM calls.
type SemanticCache struct {
	entries map[string]CachedResponse
	mu      sync.RWMutex
}

// CachedResponse holds a pre-computed response for a specific input.
type CachedResponse struct {
	Content  string               `json:"content"`
	Segments []model.ReviewSegment `json:"segments"`
	ExpireAt time.Time            `json:"expire_at"`
}

// NewSemanticCache creates a cache preloaded with common greetings.
func NewSemanticCache() *SemanticCache {
	c := &SemanticCache{entries: make(map[string]CachedResponse)}
	ttl := 24 * time.Hour

	c.Set("你好", "你好呀~ 有什么可以帮你的吗？", nil, ttl)
	c.Set("早", "早上好♪ 今天有什么计划吗？", nil, ttl)
	c.Set("晚安", "晚安~ 做个好梦！", nil, ttl)
	c.Set("谢谢", "不客气~ 随时可以找我帮忙哦", nil, ttl)
	c.Set("hello", "Hello~ How can I help you today?", nil, ttl)

	return c
}

// Set adds a cached response.
func (c *SemanticCache) Set(input, content string, segments []model.ReviewSegment, ttl time.Duration) {
	c.mu.Lock()
	defer c.mu.Unlock()
	c.entries[strings.ToLower(strings.TrimSpace(input))] = CachedResponse{
		Content:  content,
		Segments: segments,
		ExpireAt: time.Now().Add(ttl),
	}
}

// Get retrieves a cached response, or nil if not found/expired.
func (c *SemanticCache) Get(input string) *CachedResponse {
	c.mu.RLock()
	defer c.mu.RUnlock()

	key := strings.ToLower(strings.TrimSpace(input))
	if entry, ok := c.entries[key]; ok && time.Now().Before(entry.ExpireAt) {
		return &entry
	}
	return nil
}