all repos — llm_aggregator @ e3e7be8461bb50a7fe674b4055286d34dddb0c5c

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	"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}