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