all repos — llm_aggregator @ 26eb79c063daa25ab326497558d2b3bf17a675bd

A CLI tool to aggregate RSS feeds and summarise them with LLMs

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		dc.logger.Logf("Max tokens: %d, Temperature: %.2f", dc.maxTokens, dc.temperature)
183		dc.logger.Logf("Messages count: %d", len(messages))
184	}
185
186	response, err := dc.client.Chat.Completions.New(ctx, openai.ChatCompletionNewParams{
187		Model:       dc.model,
188		Messages:    messages,
189		MaxTokens:   openai.Int(int64(dc.maxTokens)),
190		Temperature: openai.Float(dc.temperature),
191	})
192
193	if err != nil {
194		// Provide more specific error messages based on error type
195		errStr := err.Error()
196
197		// Log the full error for debugging
198		if dc.logger != nil {
199			dc.logger.Logf("API error: %s", errStr)
200		}
201
202		if strings.Contains(errStr, "401") {
203			return "", fmt.Errorf("invalid API key. Please check your LLM API key")
204		} else if strings.Contains(errStr, "429") {
205			return "", fmt.Errorf("rate limit exceeded. Please try again later")
206		} else if strings.Contains(errStr, "500") {
207			return "", fmt.Errorf("LLM API server error. Please try again later")
208		} else if strings.Contains(errStr, "404") {
209			return "", fmt.Errorf("API endpoint not found. Please check the base URL and endpoint. OpenAI API uses /chat/completions")
210		}
211		return "", fmt.Errorf("failed to connect to LLM API: %w", err)
212	}
213
214	// Extract text from response
215	if len(response.Choices) == 0 {
216		return "", fmt.Errorf("no response choices returned from API")
217	}
218
219	outputText := response.Choices[0].Message.Content
220
221	// Print token usage
222	if dc.logger != nil {
223		dc.logger.Logf(
224			"LLM API response: %d prompt tokens, %d completion tokens",
225			response.Usage.PromptTokens,
226			response.Usage.CompletionTokens,
227		)
228	}
229
230	return outputText, nil
231}
232
233// AnalyseWithCustomPrompt analyses articles with custom system and user prompts.
234func (dc *LLMClient) AnalyseWithCustomPrompt(
235	articles []map[string]any,
236	systemPrompt string,
237	userPrompt string,
238) (string, error) {
239	context := dc.prepareContext(articles)
240
241	// Create messages for chat completion API
242	messages := []openai.ChatCompletionMessageParamUnion{
243		openai.SystemMessage(systemPrompt),
244		openai.UserMessage(fmt.Sprintf(`Here are articles from various RSS feeds:
245
246%s
247
248%s`, context, userPrompt)),
249	}
250
251	return dc.callAPIWithMessages(messages)
252}