all repos — llm_aggregator @ f8680f68d04b67101a791bdcffbbb96f0bc740a8

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