update
This commit is contained in:
lychang
321
internal/api/llm.go
Normal file
321
internal/api/llm.go
Normal file
@@ -0,0 +1,321 @@
|
||||
package api
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// LLMClient 大语言模型客户端
|
||||
type LLMClient struct {
|
||||
httpClient *HTTPClient
|
||||
config LLMConfig
|
||||
}
|
||||
|
||||
// LLMConfig LLM配置
|
||||
type LLMConfig struct {
|
||||
APIKey string `json:"api_key"` // API密钥
|
||||
BaseURL string `json:"base_url"` // API基础URL
|
||||
Model string `json:"model"` // 模型名称
|
||||
Provider string `json:"provider"` // 提供商: "openai", "ollama", "custom"
|
||||
Temperature float64 `json:"temperature"` // 温度参数
|
||||
MaxTokens int `json:"max_tokens"` // 最大token数
|
||||
}
|
||||
|
||||
// ChatMessage 聊天消息
|
||||
type ChatMessage struct {
|
||||
Role string `json:"role"` // 角色: "system", "user", "assistant"
|
||||
Content string `json:"content"` // 消息内容
|
||||
}
|
||||
|
||||
// ChatRequest 聊天请求 - OpenAI格式
|
||||
type ChatRequest struct {
|
||||
Model string `json:"model"` // 模型名称
|
||||
Messages []ChatMessage `json:"messages"` // 消息列表
|
||||
Temperature float64 `json:"temperature,omitempty"` // 温度参数
|
||||
MaxTokens int `json:"max_tokens,omitempty"` // 最大token数
|
||||
TopP float64 `json:"top_p,omitempty"` // 采样参数
|
||||
FrequencyPenalty float64 `json:"frequency_penalty,omitempty"` // 频率惩罚
|
||||
PresencePenalty float64 `json:"presence_penalty,omitempty"` // 存在惩罚
|
||||
Stream bool `json:"stream,omitempty"` // 是否流式输出
|
||||
}
|
||||
|
||||
// OllamaMessage Ollama消息格式
|
||||
type OllamaMessage struct {
|
||||
Role string `json:"role"` // 角色: "system", "user", "assistant"
|
||||
Content string `json:"content"` // 消息内容
|
||||
}
|
||||
|
||||
// OllamaRequest Ollama请求格式
|
||||
type OllamaRequest struct {
|
||||
Model string `json:"model"` // 模型名称
|
||||
Messages []OllamaMessage `json:"messages"` // 消息列表
|
||||
Stream bool `json:"stream"` // 是否流式输出
|
||||
Options OllamaOptions `json:"options,omitempty"` // 选项
|
||||
KeepAlive string `json:"keep_alive,omitempty"` // 保持连接时间
|
||||
Think bool `json:"think"` // 是否思考模式
|
||||
}
|
||||
|
||||
// OllamaOptions Ollama选项
|
||||
type OllamaOptions struct {
|
||||
Temperature float64 `json:"temperature,omitempty"` // 温度参数
|
||||
MaxTokens int `json:"num_predict,omitempty"` // 最大预测token数
|
||||
TopP float64 `json:"top_p,omitempty"` // 采样参数
|
||||
}
|
||||
|
||||
// ChatResponse 聊天响应 - OpenAI格式
|
||||
type ChatResponse struct {
|
||||
ID string `json:"id"` // 响应ID
|
||||
Object string `json:"object"` // 对象类型
|
||||
Created int64 `json:"created"` // 创建时间
|
||||
Model string `json:"model"` // 模型名称
|
||||
Choices []Choice `json:"choices"` // 选择列表
|
||||
Usage Usage `json:"usage"` // 使用情况
|
||||
}
|
||||
|
||||
// OllamaResponse Ollama响应格式
|
||||
type OllamaResponse struct {
|
||||
Model string `json:"model"` // 模型名称
|
||||
CreatedAt string `json:"created_at"` // 创建时间
|
||||
Message OllamaMessage `json:"message"` // 消息
|
||||
Done bool `json:"done"` // 是否完成
|
||||
TotalDuration int64 `json:"total_duration,omitempty"` // 总耗时
|
||||
LoadDuration int64 `json:"load_duration,omitempty"` // 加载耗时
|
||||
PromptEvalCount int `json:"prompt_eval_count,omitempty"` // 提示评估计数
|
||||
PromptEvalDuration int64 `json:"prompt_eval_duration,omitempty"` // 提示评估耗时
|
||||
EvalCount int `json:"eval_count,omitempty"` // 评估计数
|
||||
EvalDuration int64 `json:"eval_duration,omitempty"` // 评估耗时
|
||||
Response string `json:"response,omitempty"` // 响应内容（流式响应时）
|
||||
}
|
||||
|
||||
// Choice 选择 - OpenAI格式
|
||||
type Choice struct {
|
||||
Index int `json:"index"` // 索引
|
||||
Message ChatMessage `json:"message"` // 消息
|
||||
FinishReason string `json:"finish_reason"` // 结束原因
|
||||
}
|
||||
|
||||
// Usage 使用情况
|
||||
type Usage struct {
|
||||
PromptTokens int `json:"prompt_tokens"` // 提示token数
|
||||
CompletionTokens int `json:"completion_tokens"` // 完成token数
|
||||
TotalTokens int `json:"total_tokens"` // 总token数
|
||||
}
|
||||
|
||||
// NewLLMClient 创建LLM客户端
|
||||
func NewLLMClient(config LLMConfig) *LLMClient {
|
||||
httpClient := NewHTTPClient()
|
||||
|
||||
return &LLMClient{
|
||||
httpClient: httpClient,
|
||||
config: config,
|
||||
}
|
||||
}
|
||||
|
||||
// NewOpenAIClient 创建OpenAI客户端
|
||||
func NewOpenAIClient(apiKey, model string) *LLMClient {
|
||||
config := LLMConfig{
|
||||
APIKey: apiKey,
|
||||
BaseURL: "https://api.openai.com/v1",
|
||||
Model: model,
|
||||
Provider: "openai",
|
||||
Temperature: 0.1,
|
||||
MaxTokens: 16000,
|
||||
}
|
||||
return NewLLMClient(config)
|
||||
}
|
||||
|
||||
// NewOllamaClient 创建Ollama客户端
|
||||
func NewOllamaClient(baseURL, model string) *LLMClient {
|
||||
config := LLMConfig{
|
||||
BaseURL: baseURL,
|
||||
Model: model,
|
||||
Provider: "ollama",
|
||||
Temperature: 0.1,
|
||||
MaxTokens: 16000,
|
||||
}
|
||||
return NewLLMClient(config)
|
||||
}
|
||||
|
||||
// Chat 发送聊天请求
|
||||
func (c *LLMClient) Chat(messages []ChatMessage) (*ChatResponse, error) {
|
||||
if c.config.Provider == "ollama" {
|
||||
// 对于Ollama，我们需要转换格式
|
||||
return c.chatWithOllama(messages)
|
||||
} else {
|
||||
// 对于OpenAI和其他提供商，使用标准格式
|
||||
request := ChatRequest{
|
||||
Model: c.config.Model,
|
||||
Messages: messages,
|
||||
Temperature: c.config.Temperature,
|
||||
MaxTokens: c.config.MaxTokens,
|
||||
}
|
||||
|
||||
return c.sendChatRequest(request)
|
||||
}
|
||||
}
|
||||
|
||||
// chatWithOllama 使用Ollama格式发送请求
|
||||
func (c *LLMClient) chatWithOllama(messages []ChatMessage) (*ChatResponse, error) {
|
||||
// 转换消息格式
|
||||
ollamaMessages := make([]OllamaMessage, len(messages))
|
||||
for i, msg := range messages {
|
||||
ollamaMessages[i] = OllamaMessage{
|
||||
Role: msg.Role,
|
||||
Content: msg.Content,
|
||||
}
|
||||
}
|
||||
|
||||
// 创建Ollama请求
|
||||
request := OllamaRequest{
|
||||
Model: c.config.Model,
|
||||
Messages: ollamaMessages,
|
||||
Stream: false,
|
||||
Think: false,
|
||||
Options: OllamaOptions{
|
||||
Temperature: c.config.Temperature,
|
||||
MaxTokens: c.config.MaxTokens,
|
||||
},
|
||||
}
|
||||
|
||||
// 发送Ollama请求
|
||||
apiURL := c.buildAPIURL("/api/chat")
|
||||
headers := c.buildHeaders()
|
||||
|
||||
response, err := c.httpClient.PostJSON(apiURL, request, headers)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("发送Ollama请求失败: %v", err)
|
||||
}
|
||||
|
||||
// 解析Ollama响应
|
||||
var ollamaResp OllamaResponse
|
||||
err = json.Unmarshal(response, &ollamaResp)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("解析Ollama响应失败: %v", err)
|
||||
}
|
||||
|
||||
// 转换Ollama响应为统一的ChatResponse格式
|
||||
chatResponse := &ChatResponse{
|
||||
ID: "ollama-" + c.config.Model,
|
||||
Object: "chat.completion",
|
||||
Model: c.config.Model,
|
||||
Choices: []Choice{
|
||||
{
|
||||
Index: 0,
|
||||
Message: ChatMessage{
|
||||
Role: ollamaResp.Message.Role,
|
||||
Content: ollamaResp.Message.Content,
|
||||
},
|
||||
FinishReason: "stop",
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
return chatResponse, nil
|
||||
}
|
||||
|
||||
// ChatWithSystem 发送带系统消息的聊天请求
|
||||
func (c *LLMClient) ChatWithSystem(systemPrompt string, userMessage string) (*ChatResponse, error) {
|
||||
messages := []ChatMessage{
|
||||
{Role: "system", Content: systemPrompt},
|
||||
{Role: "user", Content: userMessage},
|
||||
}
|
||||
return c.Chat(messages)
|
||||
}
|
||||
|
||||
// ChatConversation 发送多轮对话
|
||||
func (c *LLMClient) ChatConversation(conversation []ChatMessage) (*ChatResponse, error) {
|
||||
return c.Chat(conversation)
|
||||
}
|
||||
|
||||
// sendChatRequest 发送标准聊天请求 (OpenAI格式)
|
||||
func (c *LLMClient) sendChatRequest(request ChatRequest) (*ChatResponse, error) {
|
||||
apiURL := c.buildAPIURL("/chat/completions")
|
||||
|
||||
headers := c.buildHeaders()
|
||||
|
||||
response, err := c.httpClient.PostJSON(apiURL, request, headers)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("发送请求失败: %v", err)
|
||||
}
|
||||
|
||||
var result ChatResponse
|
||||
err = json.Unmarshal(response, &result)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("解析响应失败: %v", err)
|
||||
}
|
||||
|
||||
return &result, nil
|
||||
}
|
||||
|
||||
// buildAPIURL 构建API URL
|
||||
func (c *LLMClient) buildAPIURL(endpoint string) string {
|
||||
baseURL := strings.TrimSuffix(c.config.BaseURL, "/")
|
||||
return baseURL + endpoint
|
||||
}
|
||||
|
||||
// buildHeaders 构建请求头
|
||||
func (c *LLMClient) buildHeaders() map[string]string {
|
||||
headers := map[string]string{
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
|
||||
if c.config.APIKey != "" {
|
||||
headers["Authorization"] = "Bearer " + c.config.APIKey
|
||||
}
|
||||
|
||||
return headers
|
||||
}
|
||||
|
||||
// SetTemperature 设置温度参数
|
||||
func (c *LLMClient) SetTemperature(temperature float64) {
|
||||
c.config.Temperature = temperature
|
||||
}
|
||||
|
||||
// SetMaxTokens 设置最大token数
|
||||
func (c *LLMClient) SetMaxTokens(maxTokens int) {
|
||||
c.config.MaxTokens = maxTokens
|
||||
}
|
||||
|
||||
// SetModel 设置模型
|
||||
func (c *LLMClient) SetModel(model string) {
|
||||
c.config.Model = model
|
||||
}
|
||||
|
||||
// GetConfig 获取配置
|
||||
func (c *LLMClient) GetConfig() LLMConfig {
|
||||
return c.config
|
||||
}
|
||||
|
||||
// CommonConfigs 预定义的常用配置
|
||||
var CommonConfigs = struct {
|
||||
OpenAI LLMConfig
|
||||
Ollama LLMConfig
|
||||
Groq LLMConfig
|
||||
DeepSeek LLMConfig
|
||||
}{
|
||||
OpenAI: LLMConfig{
|
||||
BaseURL: "https://api.openai.com/v1",
|
||||
Provider: "openai",
|
||||
Temperature: 0.1,
|
||||
MaxTokens: 16000,
|
||||
},
|
||||
Ollama: LLMConfig{
|
||||
BaseURL: "http://localhost:11434",
|
||||
Provider: "ollama",
|
||||
Temperature: 0.1,
|
||||
MaxTokens: 16000,
|
||||
},
|
||||
Groq: LLMConfig{
|
||||
BaseURL: "https://api.groq.com/openai/v1",
|
||||
Provider: "groq",
|
||||
Temperature: 0.1,
|
||||
MaxTokens: 16000,
|
||||
},
|
||||
DeepSeek: LLMConfig{
|
||||
BaseURL: "https://api.deepseek.com/v1",
|
||||
Provider: "deepseek",
|
||||
Temperature: 0.1,
|
||||
MaxTokens: 16000,
|
||||
},
|
||||
}
|
||||
Reference in New Issue
Block a user