all repos — llm_aggregator @ d3fad3374aee5ee549219c1387a0c5d87ce1a99b

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