322 lines
9.5 KiB
Go
322 lines
9.5 KiB
Go
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,
|
||
},
|
||
}
|