70 lines
2.1 KiB
Go
70 lines
2.1 KiB
Go
package llm
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"knowledge-graph-backend/internal/config"
|
|
"log"
|
|
|
|
"github.com/tmc/langchaingo/llms"
|
|
"github.com/tmc/langchaingo/llms/openai"
|
|
)
|
|
|
|
type Client struct {
|
|
llm llms.Model
|
|
config config.LLMConfig
|
|
}
|
|
|
|
// 创建 LLM 实例
|
|
func NewClient(cfg config.LLMConfig) (*Client, error) {
|
|
model, err := openai.New(
|
|
openai.WithToken(cfg.APIKey),
|
|
openai.WithBaseURL(cfg.BaseURL),
|
|
openai.WithModel(cfg.Model),
|
|
)
|
|
|
|
if err != nil {
|
|
return nil, fmt.Errorf("failed to create llm client: %w", err)
|
|
}
|
|
|
|
return &Client{config: cfg, llm: model}, nil
|
|
}
|
|
|
|
// Client简单对话调用
|
|
func (c *Client) Chat(ctx context.Context, messages []llms.MessageContent) (string, error) {
|
|
resp, err := c.llm.GenerateContent(ctx, messages)
|
|
|
|
if err != nil {
|
|
log.Printf("[ERROR] GenerateContent failed: %v", err)
|
|
return "", fmt.Errorf("generate content error: %w", err)
|
|
}
|
|
if resp == nil || len(resp.Choices) == 0 {
|
|
log.Printf("[WARN] Response is empty or has no choices")
|
|
return "", fmt.Errorf("no content choices returned")
|
|
}
|
|
log.Printf("[INFO] Chat completed successfully, choice length: %d", len(resp.Choices[0].Content))
|
|
|
|
return resp.Choices[0].Content, nil
|
|
}
|
|
|
|
// Client包含系统提示词的调用
|
|
func (c *Client) ChatWithSystemPrompt(ctx context.Context, systemPrompt string, messages []llms.MessageContent) (string, error) {
|
|
// 构造系统提示词消息
|
|
systemMessage := llms.TextParts(llms.ChatMessageTypeSystem, systemPrompt)
|
|
|
|
// 将系统提示词插入到消息列表头部
|
|
allMessages := append([]llms.MessageContent{systemMessage}, messages...)
|
|
|
|
resp, err := c.llm.GenerateContent(ctx, allMessages)
|
|
if err != nil {
|
|
log.Printf("[ERROR] GenerateContent with system prompt failed: %v", err)
|
|
return "", fmt.Errorf("generate content error: %w", err)
|
|
}
|
|
if resp == nil || len(resp.Choices) == 0 {
|
|
log.Printf("[WARN] Response with system prompt is empty or has no choices")
|
|
return "", fmt.Errorf("no content choices returned")
|
|
}
|
|
log.Printf("[INFO] Chat with system prompt completed successfully, choice length: %d", len(resp.Choices[0].Content))
|
|
|
|
return resp.Choices[0].Content, nil
|
|
} |