109 lines
3.4 KiB
Go
109 lines
3.4 KiB
Go
package main
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"encoding/json"
|
|
"net/http"
|
|
"strings"
|
|
|
|
"github.com/sirupsen/logrus"
|
|
)
|
|
|
|
// LLMClient abstracts LLM API calls
|
|
type LLMClient struct {
|
|
APIKey string
|
|
BaseURL string
|
|
}
|
|
|
|
// ExtractKeywords calls LLM to extract keywords from user message
|
|
func (llm *LLMClient) ExtractKeywords(ctx context.Context, message string) ([]string, error) {
|
|
prompt, err := renderPrompt(appConfig.LLM.ExtractKeywordsPrompt, map[string]string{"Message": message})
|
|
if err != nil {
|
|
logrus.WithError(err).Error("[CONFIG] Failed to render ExtractKeywords prompt")
|
|
return nil, err
|
|
}
|
|
logrus.WithField("prompt", prompt).Info("[LLM] ExtractKeywords prompt")
|
|
resp, err := llm.openAICompletion(ctx, prompt)
|
|
logrus.WithFields(logrus.Fields{"response": resp, "err": err}).Info("[LLM] ExtractKeywords response")
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
var keywords []string
|
|
if err := json.Unmarshal([]byte(resp), &keywords); err == nil {
|
|
return keywords, nil
|
|
}
|
|
// fallback: try splitting by comma
|
|
for _, k := range bytes.Split([]byte(resp), []byte{','}) {
|
|
kw := strings.TrimSpace(string(k))
|
|
if kw != "" {
|
|
keywords = append(keywords, kw)
|
|
}
|
|
}
|
|
return keywords, nil
|
|
}
|
|
|
|
// DisambiguateBestMatch calls LLM to pick best match from candidates
|
|
func (llm *LLMClient) DisambiguateBestMatch(ctx context.Context, message string, candidates []Reason) (string, error) {
|
|
entries, _ := json.Marshal(candidates)
|
|
prompt, err := renderPrompt(appConfig.LLM.DisambiguatePrompt, map[string]string{"Entries": string(entries), "Message": message})
|
|
if err != nil {
|
|
logrus.WithError(err).Error("[CONFIG] Failed to render Disambiguate prompt")
|
|
return "", err
|
|
}
|
|
logrus.WithField("prompt", prompt).Info("[LLM] DisambiguateBestMatch prompt")
|
|
resp, err := llm.openAICompletion(ctx, prompt)
|
|
logrus.WithFields(logrus.Fields{"response": resp, "err": err}).Info("[LLM] DisambiguateBestMatch response")
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
id := strings.TrimSpace(resp)
|
|
if id == "none" || id == "null" {
|
|
return "", nil
|
|
}
|
|
return id, nil
|
|
}
|
|
|
|
// openAICompletion is a minimal OpenAI API call (text-davinci-003 or gpt-3.5-turbo-instruct)
|
|
func (llm *LLMClient) openAICompletion(ctx context.Context, prompt string) (string, error) {
|
|
apiURL := llm.BaseURL
|
|
if apiURL == "" {
|
|
apiURL = "https://api.openai.com/v1/completions"
|
|
}
|
|
logrus.WithFields(logrus.Fields{"api_url": apiURL, "prompt": prompt}).Info("[LLM] openAICompletion POST")
|
|
body := map[string]interface{}{
|
|
"model": "text-davinci-003",
|
|
"prompt": prompt,
|
|
"max_tokens": 64,
|
|
"temperature": 0,
|
|
}
|
|
jsonBody, _ := json.Marshal(body)
|
|
req, _ := http.NewRequestWithContext(ctx, "POST", apiURL, bytes.NewBuffer(jsonBody))
|
|
if llm.APIKey != "" {
|
|
req.Header.Set("Authorization", "Bearer "+llm.APIKey)
|
|
}
|
|
req.Header.Set("Content-Type", "application/json")
|
|
client := &http.Client{}
|
|
resp, err := client.Do(req)
|
|
if err != nil {
|
|
logrus.WithError(err).Error("[LLM] openAICompletion error")
|
|
return "", err
|
|
}
|
|
defer resp.Body.Close()
|
|
var result struct {
|
|
Choices []struct {
|
|
Text string `json:"text"`
|
|
} `json:"choices"`
|
|
}
|
|
if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
|
|
logrus.WithError(err).Error("[LLM] openAICompletion decode error")
|
|
return "", err
|
|
}
|
|
if len(result.Choices) == 0 {
|
|
logrus.Warn("[LLM] openAICompletion: no choices returned")
|
|
return "", nil
|
|
}
|
|
logrus.WithField("text", result.Choices[0].Text).Info("[LLM] openAICompletion: got text")
|
|
return result.Choices[0].Text, nil
|
|
}
|