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": 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
}

// 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) {
	start := time.Now()
	body := map[string]any{
		"model":    a.cfg.Model,
		"messages": messages,
		"stream":   true,
	}

	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)
	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}
			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}
				}
			}
			if reason, ok := choice["finish_reason"].(string); ok && reason != "" {
				ch <- model.StreamToken{FinishReason: reason, Done: true}
				return nil
			}
		}
	}
	return scanner.Err()
}

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

// 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
}