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