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