Files
ai-agent/internal/ai/embedding.go
T
mlogclub 5d7c10aeab refactor: rename agent widget references to AI agent for consistency
- Updated runtime configuration to use __CS_AI_AGENT_WIDGET_CONFIG__ instead of __CS_AGENT_WIDGET_CONFIG__.
- Changed message types in support host bridge from "cs-agent" to "cs-ai-agent".
- Minified SDK script updated to reflect new AI agent naming conventions.
- Adjusted scrollbar styles in main.scss to use .cs-ai-agent-scrollbar instead of .cs-agent-scrollbar.
2026-05-30 21:19:36 +08:00

102 lines
2.5 KiB
Go

package ai
import (
"context"
"fmt"
openai "github.com/openai/openai-go/v3"
"cs-ai-agent/internal/models"
"cs-ai-agent/internal/pkg/enums"
"cs-ai-agent/internal/pkg/errorsx"
)
type EmbeddingResult struct {
Vector []float32
TokensUsed int
ModelName string
Dimension int
}
type embedding struct{}
var Embedding = &embedding{}
func (s *embedding) GetModel(ctx context.Context) (*models.AIConfig, error) {
config, err := GetEnabledAIConfig(enums.AIModelTypeEmbedding)
if err != nil {
return nil, errorsx.BusinessError(2001, "未配置可用的 Embedding 模型")
}
return config, nil
}
func (s *embedding) GenerateEmbedding(ctx context.Context, text string) (*EmbeddingResult, error) {
if text == "" {
return nil, errorsx.InvalidParam("文本内容不能为空")
}
result, err := s.callEmbeddingAPI(ctx, text)
if err != nil {
return nil, err
}
return result, nil
}
func (s *embedding) GenerateBatchEmbeddings(ctx context.Context, texts []string) ([]EmbeddingResult, error) {
if len(texts) == 0 {
return nil, errorsx.InvalidParam("文本列表不能为空")
}
results := make([]EmbeddingResult, 0, len(texts))
for _, text := range texts {
result, err := s.callEmbeddingAPI(ctx, text)
if err != nil {
return nil, fmt.Errorf("failed to generate embedding for text: %w", err)
}
results = append(results, *result)
}
return results, nil
}
func (s *embedding) callEmbeddingAPI(ctx context.Context, text string) (*EmbeddingResult, error) {
config, err := GetEnabledAIConfig(enums.AIModelTypeEmbedding)
if err != nil {
return nil, err
}
client := newOpenAIClient(*config)
embeddingResp, err := client.Embeddings.New(ctx, openai.EmbeddingNewParams{
Input: openai.EmbeddingNewParamsInputUnion{
OfString: openai.String(text),
},
Model: openai.EmbeddingModel(config.ModelName),
})
if err != nil {
return nil, fmt.Errorf("failed to call embedding api: %w", err)
}
if len(embeddingResp.Data) == 0 {
return nil, fmt.Errorf("no embedding data in response")
}
vector := make([]float32, 0, len(embeddingResp.Data[0].Embedding))
for _, item := range embeddingResp.Data[0].Embedding {
vector = append(vector, float32(item))
}
return &EmbeddingResult{
Vector: vector,
TokensUsed: int(embeddingResp.Usage.TotalTokens),
ModelName: embeddingResp.Model,
Dimension: len(vector),
}, nil
}
func (s *embedding) GetDimension(ctx context.Context) (int, error) {
model, err := s.GetModel(ctx)
if err != nil {
return 0, err
}
return model.Dimension, nil
}