rsspipe/internal/ai/client.go

package ai

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

	"git.pengzhan.dev/haopengzhan/rsspipe/internal/models"
)

type Config struct {
	Enabled     bool    `json:"enabled"`
	Provider    string  `json:"provider"`
	BaseURL     string  `json:"base_url"`
	APIKey      string  `json:"api_key"`
	Model       string  `json:"model"`
	Temperature float64 `json:"temperature"`
	MaxTokens   int     `json:"max_tokens"`
}

type PromptTemplate struct {
	ID          string `json:"id"`
	Name        string `json:"name"`
	Description string `json:"description,omitempty"`
	Prompt      string `json:"prompt"`
	IsDefault   bool   `json:"is_default,omitempty"`
}

type TestResult struct {
	Success   bool          `json:"success"`
	Message   string        `json:"message"`
	LatencyMs int64         `json:"latency_ms"`
	Model     string        `json:"model"`
	Duration  time.Duration `json:"duration"`
}

type AgentDecision struct {
	Reasoning string              `json:"reasoning"`
	Actions   []models.RuleAction `json:"actions"`
}

type Client struct {
	httpClient *http.Client
}

func NewClient(timeout time.Duration) *Client {
	if timeout <= 0 {
		timeout = 45 * time.Second
	}
	return &Client{
		httpClient: &http.Client{
			Timeout: timeout,
		},
	}
}

// DefaultPromptTemplates returns built-in prompt templates.
func DefaultPromptTemplates() []PromptTemplate {
	return []PromptTemplate{
		{
			ID:          "exec_summary",
			Name:        "3-Bullet Executive Brief",
			Description: "3 concise bullet points followed by a single key takeaway.",
			Prompt:      "You are an expert technical editor. Summarize the following article in exactly 3 concise bullet points, followed by a bold 'Key Takeaway:' sentence.\n\nTitle: {{title}}\nFeed: {{feed_title}}\nContent:\n{{content}}",
			IsDefault:   true,
		},
		{
			ID:          "tech_deep_dive",
			Name:        "Technical Deep Dive",
			Description: "Architecture, key technical concepts, and implementation highlights.",
			Prompt:      "Analyze the following technical post. Highlight:\n1. Core Architecture & Concepts\n2. Key Innovations / Techniques\n3. Practical Implications for Engineers\n\nTitle: {{title}}\nContent:\n{{content}}",
			IsDefault:   false,
		},
		{
			ID:          "news_general",
			Name:        "General News & Context",
			Description: "Accessible summary with necessary background context for general readers.",
			Prompt:      "Summarize this news article for a general audience. Explain what happened, why it matters, and provide any helpful context in 2-3 short paragraphs.\n\nTitle: {{title}}\nContent:\n{{content}}",
			IsDefault:   false,
		},
	}
}

type chatMessage struct {
	Role    string `json:"role"`
	Content string `json:"content"`
}

type chatCompletionRequest struct {
	Model       string        `json:"model"`
	Messages    []chatMessage `json:"messages"`
	Temperature float64       `json:"temperature,omitempty"`
	MaxTokens   int           `json:"max_tokens,omitempty"`
}

type chatCompletionResponse struct {
	Choices []struct {
		Message struct {
			Role    string `json:"role"`
			Content string `json:"content"`
		} `json:"message"`
		FinishReason string `json:"finish_reason"`
	} `json:"choices"`
	Error *struct {
		Message string `json:"message"`
		Type    string `json:"type"`
	} `json:"error,omitempty"`
}

func (c *Client) callChatCompletion(ctx context.Context, cfg Config, req chatCompletionRequest) (string, error) {
	baseURL := strings.TrimRight(cfg.BaseURL, "/")
	if baseURL == "" {
		baseURL = "https://api.openai.com/v1"
	}
	endpoint := baseURL + "/chat/completions"

	bodyBytes, err := json.Marshal(req)
	if err != nil {
		return "", err
	}

	httpReq, err := http.NewRequestWithContext(ctx, "POST", endpoint, bytes.NewReader(bodyBytes))
	if err != nil {
		return "", err
	}

	httpReq.Header.Set("Content-Type", "application/json")
	if cfg.APIKey != "" {
		httpReq.Header.Set("Authorization", "Bearer "+cfg.APIKey)
	}

	resp, err := c.httpClient.Do(httpReq)
	if err != nil {
		return "", fmt.Errorf("HTTP request failed: %w", err)
	}
	defer func() {
		_ = resp.Body.Close()
	}()

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

	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
		return "", fmt.Errorf("API error (status %d): %s", resp.StatusCode, string(respBody))
	}

	var chatResp chatCompletionResponse
	if err := json.Unmarshal(respBody, &chatResp); err != nil {
		return "", fmt.Errorf("failed to decode response JSON: %w", err)
	}

	if chatResp.Error != nil && chatResp.Error.Message != "" {
		return "", fmt.Errorf("LLM error: %s", chatResp.Error.Message)
	}

	if len(chatResp.Choices) == 0 {
		return "", fmt.Errorf("LLM returned no choices")
	}

	return chatResp.Choices[0].Message.Content, nil
}

// TestConnection pings the LLM with a strictly capped request (max_tokens: 15) to verify credentials.
func (c *Client) TestConnection(ctx context.Context, cfg Config) (*TestResult, error) {
	start := time.Now()
	req := chatCompletionRequest{
		Model: cfg.Model,
		Messages: []chatMessage{
			{Role: "user", Content: "Ping test. Reply with 'OK' in exactly two words."},
		},
		Temperature: 0.1,
		MaxTokens:   15, // Strictly capped to prevent token exhaustion
	}

	reply, err := c.callChatCompletion(ctx, cfg, req)
	duration := time.Since(start)

	if err != nil {
		return &TestResult{
			Success:   false,
			Message:   err.Error(),
			LatencyMs: duration.Milliseconds(),
			Model:     cfg.Model,
			Duration:  duration,
		}, err
	}

	return &TestResult{
		Success:   true,
		Message:   strings.TrimSpace(reply),
		LatencyMs: duration.Milliseconds(),
		Model:     cfg.Model,
		Duration:  duration,
	}, nil
}

// GenerateEnrichment interpolates the prompt template and queries the LLM to generate variant content.
func (c *Client) GenerateEnrichment(ctx context.Context, cfg Config, promptTemplate string, entry models.FeedEntry) (string, error) {
	if promptTemplate == "" {
		promptTemplate = DefaultPromptTemplates()[0].Prompt
	}

	content := entry.Content
	if entry.RawContent != "" && len(entry.RawContent) > len(content) {
		content = entry.RawContent
	}
	if entry.Variants != nil {
		if v, ok := entry.Variants["readability_full_text"]; ok && len(v.TextValue) > len(content) {
			content = v.TextValue
		}
	}

	// Truncate content to avoid exceeding token windows (e.g. max ~24,000 characters)
	if len(content) > 24000 {
		content = content[:24000] + "... [truncated]"
	}

	interpolated := promptTemplate
	interpolated = strings.ReplaceAll(interpolated, "{{title}}", entry.Title)
	interpolated = strings.ReplaceAll(interpolated, "{{content}}", content)
	interpolated = strings.ReplaceAll(interpolated, "{{feed_title}}", entry.FeedTitle)
	interpolated = strings.ReplaceAll(interpolated, "{{author}}", entry.Author)
	interpolated = strings.ReplaceAll(interpolated, "{{url}}", entry.URL)
	interpolated = strings.ReplaceAll(interpolated, "{{published_at}}", entry.PublishedAt.Format("2006-01-02 15:04:05"))

	maxTokens := cfg.MaxTokens
	if maxTokens <= 0 {
		maxTokens = 1000
	}
	temp := cfg.Temperature
	if temp <= 0 {
		temp = 0.3
	}

	req := chatCompletionRequest{
		Model: cfg.Model,
		Messages: []chatMessage{
			{Role: "user", Content: interpolated},
		},
		Temperature: temp,
		MaxTokens:   maxTokens,
	}

	return c.callChatCompletion(ctx, cfg, req)
}

// GenerateAgentDecision asks the LLM to evaluate the article and propose actions from allowedActions only.
func (c *Client) GenerateAgentDecision(ctx context.Context, cfg Config, policyPrompt string, allowedActions []string, entry models.FeedEntry) (*AgentDecision, error) {
	allowedListStr := strings.Join(allowedActions, ", ")

	content := entry.Content
	if entry.RawContent != "" && len(entry.RawContent) > len(content) {
		content = entry.RawContent
	}
	if entry.Variants != nil {
		if v, ok := entry.Variants["readability_full_text"]; ok && len(v.TextValue) > len(content) {
			content = v.TextValue
		}
	}
	if len(content) > 16000 {
		content = content[:16000] + "... [truncated]"
	}

	systemInstruction := fmt.Sprintf(`You are an autonomous RSS triage agent.
Your goal is to evaluate the provided article and propose actions strictly from the Allowed Action Pool.

Allowed Action Pool:
%s

Rules:
1. You may ONLY propose actions that exist in the Allowed Action Pool.
2. Output ONLY a valid JSON object matching this schema:
{
  "reasoning": "brief explanation of your decision",
  "actions": [
    {"type": "action_type", "params": {"key": "value"}}
  ]
}
If no action is warranted, return {"reasoning": "...", "actions": []}.
Do not wrap in markdown backticks or other text. Return raw JSON only.`, allowedListStr)

	userContent := fmt.Sprintf("Evaluation Policy:\n%s\n\nArticle Title: %s\nFeed: %s\nContent:\n%s", policyPrompt, entry.Title, entry.FeedTitle, content)

	req := chatCompletionRequest{
		Model: cfg.Model,
		Messages: []chatMessage{
			{Role: "system", Content: systemInstruction},
			{Role: "user", Content: userContent},
		},
		Temperature: 0.1,
		MaxTokens:   800,
	}

	rawJSON, err := c.callChatCompletion(ctx, cfg, req)
	if err != nil {
		return nil, err
	}

	// Clean code fence formatting if returned
	rawJSON = strings.TrimSpace(rawJSON)
	rawJSON = strings.TrimPrefix(rawJSON, "```json")
	rawJSON = strings.TrimPrefix(rawJSON, "```")
	rawJSON = strings.TrimSuffix(rawJSON, "```")
	rawJSON = strings.TrimSpace(rawJSON)

	var parsed struct {
		Reasoning string `json:"reasoning"`
		Actions   []struct {
			Type   string                 `json:"type"`
			Params map[string]interface{} `json:"params,omitempty"`
		} `json:"actions"`
	}

	if err := json.Unmarshal([]byte(rawJSON), &parsed); err != nil {
		return nil, fmt.Errorf("failed to parse AI decision JSON: %w (raw: %s)", err, rawJSON)
	}

	// Strict Guardrail: Validate all actions against allowedActions
	allowedSet := make(map[string]bool)
	for _, a := range allowedActions {
		allowedSet[strings.ToLower(strings.TrimSpace(a))] = true
	}

	var validatedActions []models.RuleAction
	for _, act := range parsed.Actions {
		actTypeLower := strings.ToLower(strings.TrimSpace(act.Type))

		// Check basic action type or composite key (e.g. set_variant:tags:xxx)
		isAllowed := allowedSet[actTypeLower]
		if !isAllowed && actTypeLower == "set_variant" && act.Params != nil {
			vKey := fmt.Sprintf("%v", act.Params["variant_key"])
			vVal := fmt.Sprintf("%v", act.Params["value"])
			compositeKey := strings.ToLower(fmt.Sprintf("set_variant:%s:%s", vKey, vVal))
			isAllowed = allowedSet[compositeKey] || allowedSet["set_variant"]
		}

		if isAllowed {
			validatedActions = append(validatedActions, models.RuleAction{
				Type:   act.Type,
				Params: act.Params,
			})
		}
	}

	return &AgentDecision{
		Reasoning: parsed.Reasoning,
		Actions:   validatedActions,
	}, nil
}