all repos — llm_aggregator @ 29a4710e981512a82abba77ec8d7b43e413453e7

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