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