1package llm
  2
  3import (
  4	"context"
  5	"fmt"
  6	"os"
  7	"strings"
  8	"time"
  9
 10	"llm_aggregator/internal/progress"
 11
 12	"github.com/openai/openai-go/v3"
 13	"github.com/openai/openai-go/v3/option"
 14)
 15
 16// LLMClient is a client for interacting with LLM API.
 17type LLMClient struct {
 18	client      openai.Client
 19	model       string
 20	maxTokens   int
 21	temperature float64
 22	logger      *progress.Context
 23}
 24
 25// NewLLMClient creates a new LLM client.
 26// apiKey: LLM API key (or read from LLM_AGGREGATOR_API_KEY env var)
 27// baseURL: API base URL (defaults to "https://api.deepseek.com")
 28// model: Model to use (defaults to "deepseek-chat")
 29// maxTokens: Maximum tokens in response (defaults to 4000)
 30// temperature: Sampling temperature (0.0 to 1.0, defaults to 0.7)
 31func NewLLMClient(apiKey, baseURL, model string, maxTokens int, temperature float64) (*LLMClient, error) {
 32	// Get API key from parameter or environment variable
 33	if apiKey == "" {
 34		apiKey = os.Getenv("LLM_AGGREGATOR_API_KEY")
 35	}
 36	if apiKey == "" || strings.TrimSpace(apiKey) == "" {
 37		return nil, fmt.Errorf(
 38			"LLM API key is required. " +
 39				"Set LLM_AGGREGATOR_API_KEY environment variable or pass apiKey parameter",
 40		)
 41	}
 42
 43	// Set defaults
 44	if baseURL == "" {
 45		baseURL = "https://api.deepseek.com"
 46	}
 47	if model == "" {
 48		model = "deepseek-chat"
 49	}
 50	if maxTokens == 0 {
 51		maxTokens = 4000
 52	}
 53	if temperature == 0 {
 54		temperature = 0.7
 55	}
 56
 57	// Create OpenAI client configured for LLM
 58	client := openai.NewClient(
 59		option.WithAPIKey(apiKey),
 60		option.WithBaseURL(baseURL),
 61	)
 62
 63	return &LLMClient{
 64		client:      client,
 65		model:       model,
 66		maxTokens:   maxTokens,
 67		temperature: temperature,
 68	}, nil
 69}
 70
 71// SetLogger sets the logger for the LLM client
 72func (dc *LLMClient) SetLogger(logger *progress.Context) {
 73	dc.logger = logger
 74}
 75
 76// SummariseArticles summarises a list of articles based on user prompt.
 77// articles: List of article maps
 78// userPrompt: User's query/summarisation request
 79// systemPrompt: Optional system prompt (defaults to helpful assistant)
 80func (dc *LLMClient) SummariseArticles(
 81	articles []map[string]any,
 82	userPrompt string,
 83	systemPrompt string,
 84) (string, error) {
 85	if len(articles) == 0 {
 86		return "No articles to summarise.", nil
 87	}
 88
 89	// Prepare context from articles
 90	context := dc.prepareContext(articles)
 91
 92	// Create messages for chat completion API
 93	messages := dc.createMessages(context, userPrompt, systemPrompt)
 94
 95	// Call API with messages
 96	return dc.callAPIWithMessages(messages)
 97}
 98
 99func (dc *LLMClient) prepareContext(articles []map[string]any) string {
100	contextParts := []string{}
101
102	for i, article := range articles {
103		contextParts = append(contextParts, fmt.Sprintf("--- ARTICLE %d ---", i+1))
104		contextParts = append(contextParts, fmt.Sprintf("Title: %s", article["title"]))
105
106		if source, ok := article["source_feed"].(string); ok && source != "" {
107			contextParts = append(contextParts, fmt.Sprintf("Source: %s", source))
108		}
109
110		if published, ok := article["published"]; ok {
111			switch pub := published.(type) {
112			case time.Time:
113				if !pub.IsZero() {
114					contextParts = append(contextParts, fmt.Sprintf("Published: %s", pub.Format(time.RFC3339)))
115				}
116			case string:
117				contextParts = append(contextParts, fmt.Sprintf("Published: %s", pub))
118			default:
119				contextParts = append(contextParts, fmt.Sprintf("Published: %v", pub))
120			}
121		}
122
123		if author, ok := article["author"].(string); ok && author != "" {
124			contextParts = append(contextParts, fmt.Sprintf("Author: %s", author))
125		}
126
127		if link, ok := article["link"].(string); ok && link != "" {
128			contextParts = append(contextParts, fmt.Sprintf("Link: %s", link))
129		}
130
131		if content, ok := article["content"].(string); ok && content != "" {
132			// Truncate very long content
133			maxContentLen := 3000
134			if len(content) > maxContentLen {
135				content = content[:maxContentLen] + "... [truncated]"
136			}
137			contextParts = append(contextParts, fmt.Sprintf("Content: %s", content))
138		}
139
140		contextParts = append(contextParts, "") // Empty line between articles
141	}
142
143	return strings.Join(contextParts, "\n")
144}
145
146func (dc *LLMClient) createMessages(context, userPrompt, systemPrompt string) []openai.ChatCompletionMessageParamUnion {
147	if systemPrompt == "" {
148		systemPrompt = `You are an expert analyst and summariser.
149You analyse content from multiple sources and provide
150concise, insightful summaries based on user requests.
151Focus on key points, trends, and important information.`
152	}
153
154	// Combine context with user prompt
155	fullUserContent := fmt.Sprintf(`Here are articles from various RSS feeds:
156
157%s
158
159User request: %s
160
161Please provide a comprehensive summary/analysis addressing the user's request.
162Focus on key insights, trends, and important information from the articles.
163If relevant, note any patterns, contradictions, or notable developments.`,
164		context, userPrompt)
165
166	messages := []openai.ChatCompletionMessageParamUnion{
167		openai.SystemMessage(systemPrompt),
168		openai.UserMessage(fullUserContent),
169	}
170
171	return messages
172}
173
174func (dc *LLMClient) callAPIWithMessages(messages []openai.ChatCompletionMessageParamUnion) (string, error) {
175	ctx := context.Background()
176
177	if dc.logger != nil {
178		dc.logger.Logf("Calling LLM API with model: %s", dc.model)
179	}
180
181	response, err := dc.client.Chat.Completions.New(ctx, openai.ChatCompletionNewParams{
182		Model:       dc.model,
183		Messages:    messages,
184		MaxTokens:   openai.Int(int64(dc.maxTokens)),
185		Temperature: openai.Float(dc.temperature),
186	})
187
188	if err != nil {
189		// Provide more specific error messages
190		errStr := err.Error()
191		if strings.Contains(errStr, "401") {
192			return "", fmt.Errorf("invalid API key. Please check your LLM API key")
193		} else if strings.Contains(errStr, "429") {
194			return "", fmt.Errorf("rate limit exceeded. Please try again later")
195		} else if strings.Contains(errStr, "500") {
196			return "", fmt.Errorf("LLM API server error. Please try again later")
197		} else if strings.Contains(errStr, "404") {
198			return "", fmt.Errorf("API endpoint not found. Please check the base URL and endpoint. OpenAI API uses /chat/completions")
199		}
200		return "", fmt.Errorf("failed to connect to LLM API: %w", err)
201	}
202
203	// Extract text from response
204	if len(response.Choices) == 0 {
205		return "", fmt.Errorf("no response choices returned from API")
206	}
207
208	outputText := response.Choices[0].Message.Content
209
210	// Print token usage
211	if dc.logger != nil {
212		dc.logger.Logf(
213			"LLM API response: %d prompt tokens, %d completion tokens",
214			response.Usage.PromptTokens,
215			response.Usage.CompletionTokens,
216		)
217	}
218
219	return outputText, nil
220}
221
222// AnalyseWithCustomPrompt analyses articles with custom system and user prompts.
223func (dc *LLMClient) AnalyseWithCustomPrompt(
224	articles []map[string]any,
225	systemPrompt string,
226	userPrompt string,
227) (string, error) {
228	context := dc.prepareContext(articles)
229
230	// Create messages for chat completion API
231	messages := []openai.ChatCompletionMessageParamUnion{
232		openai.SystemMessage(systemPrompt),
233		openai.UserMessage(fmt.Sprintf(`Here are articles from various RSS feeds:
234
235%s
236
237%s`, context, userPrompt)),
238	}
239
240	return dc.callAPIWithMessages(messages)
241}