Files
PR-Helper/services/llm.go
T

153 lines
3.7 KiB
Go

package services
import (
"context"
"database/sql"
"fmt"
"io"
"strings"
goopenai "github.com/sashabaranov/go-openai"
)
// LLMConfig holds LLM API configuration.
type LLMConfig struct {
Endpoint string
APIKey string
Model string
}
// GetLLMConfig reads LLM settings from the user_settings table for the given user.
func GetLLMConfig(db *sql.DB, userID int64) (LLMConfig, error) {
config := LLMConfig{}
rows, err := db.Query("SELECT `key`, value FROM user_settings WHERE user_id = ? AND `key` IN ('llm.endpoint', 'llm.api_key', 'llm.model')", userID)
if err != nil {
return config, fmt.Errorf("read settings: %w", err)
}
defer rows.Close()
for rows.Next() {
var key, value string
if rows.Scan(&key, &value) == nil {
switch key {
case "llm.endpoint":
config.Endpoint = value
case "llm.api_key":
config.APIKey = value
case "llm.model":
config.Model = value
}
}
}
if config.APIKey == "" {
return config, fmt.Errorf("LLM API key not configured — please set it in the settings page")
}
if config.Model == "" {
config.Model = "deepseek-v4-pro"
}
return config, nil
}
// StreamCallback is called for each SSE event during LLM streaming.
type StreamCallback func(event string, data interface{})
// ChatStream sends a streaming chat completion request to an OpenAI-compatible API.
// It calls callback with "content" events for each chunk received.
// Returns the full concatenated response text.
func ChatStream(config LLMConfig, messages []goopenai.ChatCompletionMessage, callback StreamCallback) (string, error) {
clientConfig := goopenai.DefaultConfig(config.APIKey)
if config.Endpoint != "" {
clientConfig.BaseURL = config.Endpoint
}
client := goopenai.NewClientWithConfig(clientConfig)
ctx := context.Background()
stream, err := client.CreateChatCompletionStream(ctx, goopenai.ChatCompletionRequest{
Model: config.Model,
Messages: messages,
Stream: true,
})
if err != nil {
return "", fmt.Errorf("create stream: %w", err)
}
defer stream.Close()
var fullResponse strings.Builder
for {
response, err := stream.Recv()
if err == io.EOF {
break
}
if err != nil {
return fullResponse.String(), fmt.Errorf("stream recv: %w", err)
}
if len(response.Choices) > 0 {
content := response.Choices[0].Delta.Content
if content != "" {
fullResponse.WriteString(content)
if callback != nil {
callback("content", map[string]interface{}{"content": content})
}
}
}
}
return fullResponse.String(), nil
}
// extractJSON attempts to extract a JSON object or array from a string
// that may contain markdown code blocks or extra text.
func extractJSON(s string) string {
// Try to find JSON in markdown code block
if idx := strings.Index(s, "```json"); idx >= 0 {
start := idx + 7
if end := strings.Index(s[start:], "```"); end >= 0 {
return strings.TrimSpace(s[start : start+end])
}
}
if idx := strings.Index(s, "```"); idx >= 0 {
start := idx + 3
if nl := strings.Index(s[start:], "\n"); nl >= 0 {
start += nl + 1
}
if end := strings.Index(s[start:], "```"); end >= 0 {
return strings.TrimSpace(s[start : start+end])
}
}
// Find first { or [
startObj := strings.Index(s, "{")
startArr := strings.Index(s, "[")
var start int
var endChar byte
if startObj >= 0 && (startArr < 0 || startObj < startArr) {
start = startObj
endChar = '}'
} else if startArr >= 0 {
start = startArr
endChar = ']'
} else {
return s
}
// Find matching closing bracket
depth := 0
for i := start; i < len(s); i++ {
if s[i] == '{' || s[i] == '[' {
depth++
} else if s[i] == '}' || s[i] == ']' {
depth--
if depth == 0 && s[i] == endChar {
return s[start : i+1]
}
}
}
return s[start:]
}