init: 数字员工平台初始代码
包含前端（Vue3 + VueFlow 画布）、后端（Go）、文档体系。 - 工作台画布：节点拖放、连线模式、右键菜单、AI 助手 - 后端：连接器 API、专员种子数据 - 导航：左侧导航、工坊、市场、控制台
This commit is contained in:
eaiadmin
@@ -0,0 +1,82 @@
|
||||
package ai
|
||||
|
||||
import (
|
||||
"time"
|
||||
|
||||
"gorm.io/gorm"
|
||||
|
||||
"eaisalestrain/backend/internal/model"
|
||||
"eaisalestrain/backend/internal/store"
|
||||
)
|
||||
|
||||
// ──────────────────────────────────────────────
|
||||
// 按用户算力点计费（对齐 pj034 router.py 的 compute_credits / log_ai_call）
|
||||
// ──────────────────────────────────────────────
|
||||
|
||||
// AI 能力常量（pj034 AiCapability 的精简子集）
|
||||
const (
|
||||
CapabilityAIChat = "ai_chat" // PathCoach 对话（每轮扣 1 点）
|
||||
CapabilityTextGen = "text_gen" // 快捷动作（情景演练/查佣金/产品对比，扣 1 点）
|
||||
CapabilityEmbed = "embed" // 知识检索内部 embedding（不扣点，仅审计）
|
||||
CapabilityEssayGrade = "essay_grade" // 简答题 LLM 评分（系统自动，不扣点，仅审计）
|
||||
)
|
||||
|
||||
// CapabilityCredits 各能力扣点成本
|
||||
var CapabilityCredits = map[string]int{
|
||||
CapabilityAIChat: 1,
|
||||
CapabilityTextGen: 1,
|
||||
CapabilityEmbed: 0,
|
||||
CapabilityEssayGrade: 0,
|
||||
}
|
||||
|
||||
// ComputeCredits 扣点决策：仅「成功」才扣点；失败不扣。
|
||||
func ComputeCredits(capability string, success bool) int {
|
||||
if !success {
|
||||
return 0
|
||||
}
|
||||
return CapabilityCredits[capability]
|
||||
}
|
||||
|
||||
// LogEntry 一次 AI 调用的审计入参
|
||||
type LogEntry struct {
|
||||
UserID uint
|
||||
Capability string
|
||||
Provider string
|
||||
RouteID string
|
||||
Model string
|
||||
TokensInput int
|
||||
TokensOutput int
|
||||
Success bool
|
||||
ErrorMessage string
|
||||
LatencyMs int
|
||||
}
|
||||
|
||||
// LogCall 写 ai_call_log；成功且需扣点时从用户余额扣点。审计写入失败不阻断主流程。
|
||||
func LogCall(e LogEntry) {
|
||||
credits := ComputeCredits(e.Capability, e.Success)
|
||||
status := "success"
|
||||
if !e.Success {
|
||||
status = "failed"
|
||||
}
|
||||
rec := model.AiCallLog{
|
||||
UserID: e.UserID,
|
||||
Capability: e.Capability,
|
||||
Provider: e.Provider,
|
||||
RouteID: e.RouteID,
|
||||
Model: e.Model,
|
||||
TokensInput: e.TokensInput,
|
||||
TokensOutput: e.TokensOutput,
|
||||
CreditsCharged: credits,
|
||||
Status: status,
|
||||
ErrorMessage: e.ErrorMessage,
|
||||
LatencyMs: e.LatencyMs,
|
||||
CreatedAt: time.Now(),
|
||||
}
|
||||
if err := store.DB.Create(&rec).Error; err != nil {
|
||||
return
|
||||
}
|
||||
if credits > 0 {
|
||||
store.DB.Model(&model.User{}).Where("id = ?", e.UserID).
|
||||
UpdateColumn("ai_points", gorm.Expr("ai_points - ?", credits))
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,397 @@
|
||||
package ai
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"eaisalestrain/backend/internal/config"
|
||||
"eaisalestrain/backend/internal/model"
|
||||
"eaisalestrain/backend/internal/store"
|
||||
)
|
||||
|
||||
// ──────────────────────────────────────────────
|
||||
// 基础类型
|
||||
// ──────────────────────────────────────────────
|
||||
|
||||
// Message OpenAI 兼容消息
|
||||
type Message struct {
|
||||
Role string `json:"role"`
|
||||
Content string `json:"content"`
|
||||
}
|
||||
|
||||
// ChatResult 非流式调用结果
|
||||
type ChatResult struct {
|
||||
Content string
|
||||
Model string
|
||||
Provider string
|
||||
FinishReason string
|
||||
Usage struct {
|
||||
PromptTokens int `json:"prompt_tokens"`
|
||||
CompletionTokens int `json:"completion_tokens"`
|
||||
TotalTokens int `json:"total_tokens"`
|
||||
}
|
||||
}
|
||||
|
||||
// ──────────────────────────────────────────────
|
||||
// OpenAI 兼容 HTTP 客户端
|
||||
// ──────────────────────────────────────────────
|
||||
|
||||
// Client 低层 HTTP 客户端（基于 RouteConfig）
|
||||
type Client struct {
|
||||
baseURL string
|
||||
apiKey string
|
||||
model string
|
||||
maxTokens int
|
||||
temperature float64
|
||||
hc *http.Client
|
||||
}
|
||||
|
||||
// NewClient 从 RouteConfig 创建客户端
|
||||
func NewClient(route *config.RouteConfig) *Client {
|
||||
return &Client{
|
||||
baseURL: strings.TrimRight(route.BaseURL, "/"),
|
||||
apiKey: route.APIKey,
|
||||
model: route.Model,
|
||||
maxTokens: route.MaxTokens,
|
||||
temperature: route.Temperature,
|
||||
hc: &http.Client{Timeout: 120 * time.Second},
|
||||
}
|
||||
}
|
||||
|
||||
// NewClientLegacy 兼容旧接口（从 LLMConfig 创建）
|
||||
func NewClientLegacy(cfg LLMConfig) *Client {
|
||||
return &Client{
|
||||
baseURL: strings.TrimRight(cfg.BaseURL, "/"),
|
||||
apiKey: cfg.APIKey,
|
||||
model: cfg.Model,
|
||||
maxTokens: cfg.MaxTokens,
|
||||
temperature: cfg.Temperature,
|
||||
hc: &http.Client{Timeout: 120 * time.Second},
|
||||
}
|
||||
}
|
||||
|
||||
func (c *Client) url(path string) string {
|
||||
return c.baseURL + path
|
||||
}
|
||||
|
||||
func (c *Client) headers() map[string]string {
|
||||
h := map[string]string{"Content-Type": "application/json"}
|
||||
if c.apiKey != "" {
|
||||
h["Authorization"] = "Bearer " + c.apiKey
|
||||
}
|
||||
return h
|
||||
}
|
||||
|
||||
func (c *Client) post(path string, body any) (*http.Response, error) {
|
||||
j, err := json.Marshal(body)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
req, err := http.NewRequest(http.MethodPost, c.url(path), bytes.NewReader(j))
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
for k, v := range c.headers() {
|
||||
req.Header.Set(k, v)
|
||||
}
|
||||
return c.hc.Do(req)
|
||||
}
|
||||
|
||||
// ──────────────────────────────────────────────
|
||||
// 非流式调用
|
||||
// ──────────────────────────────────────────────
|
||||
|
||||
// Generate 非流式对话，返回完整正文
|
||||
func (c *Client) Generate(messages []Message) (string, error) {
|
||||
resp, err := c.post("/chat/completions", map[string]any{
|
||||
"model": c.model,
|
||||
"messages": messages,
|
||||
"stream": false,
|
||||
"temperature": c.temperature,
|
||||
"max_tokens": c.maxTokens,
|
||||
})
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("LLM 服务不可达: %w", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
data, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("读取 LLM 响应失败: %w", err)
|
||||
}
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
return "", fmt.Errorf("LLM 返回 %d: %s", resp.StatusCode, truncate(string(data), 200))
|
||||
}
|
||||
|
||||
var out struct {
|
||||
Choices []struct {
|
||||
Message struct {
|
||||
Content string `json:"content"`
|
||||
} `json:"message"`
|
||||
} `json:"choices"`
|
||||
}
|
||||
if err := json.Unmarshal(data, &out); err != nil {
|
||||
return "", fmt.Errorf("LLM 响应解析失败: %w", err)
|
||||
}
|
||||
if len(out.Choices) == 0 || strings.TrimSpace(out.Choices[0].Message.Content) == "" {
|
||||
return "", fmt.Errorf("LLM 返回空正文")
|
||||
}
|
||||
return out.Choices[0].Message.Content, nil
|
||||
}
|
||||
|
||||
// GenerateFull 非流式调用，返回完整 ChatResult（含 usage）
|
||||
func (c *Client) GenerateFull(messages []Message) (*ChatResult, error) {
|
||||
resp, err := c.post("/chat/completions", map[string]any{
|
||||
"model": c.model,
|
||||
"messages": messages,
|
||||
"stream": false,
|
||||
"temperature": c.temperature,
|
||||
"max_tokens": c.maxTokens,
|
||||
})
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("LLM 服务不可达: %w", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
data, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("读取 LLM 响应失败: %w", err)
|
||||
}
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
return nil, fmt.Errorf("LLM 返回 %d: %s", resp.StatusCode, truncate(string(data), 200))
|
||||
}
|
||||
|
||||
var raw struct {
|
||||
Model string `json:"model"`
|
||||
Choices []struct {
|
||||
Message struct {
|
||||
Content string `json:"content"`
|
||||
} `json:"message"`
|
||||
FinishReason string `json:"finish_reason"`
|
||||
} `json:"choices"`
|
||||
Usage *struct {
|
||||
PromptTokens int `json:"prompt_tokens"`
|
||||
CompletionTokens int `json:"completion_tokens"`
|
||||
TotalTokens int `json:"total_tokens"`
|
||||
} `json:"usage"`
|
||||
}
|
||||
if err := json.Unmarshal(data, &raw); err != nil {
|
||||
return nil, fmt.Errorf("LLM 响应解析失败: %w", err)
|
||||
}
|
||||
if len(raw.Choices) == 0 || strings.TrimSpace(raw.Choices[0].Message.Content) == "" {
|
||||
return nil, fmt.Errorf("LLM 返回空正文")
|
||||
}
|
||||
|
||||
result := &ChatResult{
|
||||
Content: raw.Choices[0].Message.Content,
|
||||
Model: raw.Model,
|
||||
FinishReason: raw.Choices[0].FinishReason,
|
||||
}
|
||||
if raw.Usage != nil {
|
||||
result.Usage.PromptTokens = raw.Usage.PromptTokens
|
||||
result.Usage.CompletionTokens = raw.Usage.CompletionTokens
|
||||
result.Usage.TotalTokens = raw.Usage.TotalTokens
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
|
||||
// ──────────────────────────────────────────────
|
||||
// 流式调用
|
||||
// ──────────────────────────────────────────────
|
||||
|
||||
// GenerateStream SSE 流式调用，逐 chunk 回调
|
||||
func (c *Client) GenerateStream(messages []Message, onChunk func(string)) error {
|
||||
resp, err := c.post("/chat/completions", map[string]any{
|
||||
"model": c.model,
|
||||
"messages": messages,
|
||||
"stream": true,
|
||||
"temperature": c.temperature,
|
||||
"max_tokens": c.maxTokens,
|
||||
"stream_options": map[string]any{"include_usage": true},
|
||||
})
|
||||
if err != nil {
|
||||
return fmt.Errorf("LLM 服务不可达: %w", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
data, _ := io.ReadAll(resp.Body)
|
||||
return fmt.Errorf("LLM 返回 %d: %s", resp.StatusCode, truncate(string(data), 200))
|
||||
}
|
||||
|
||||
scanner := bufio.NewScanner(resp.Body)
|
||||
scanner.Buffer(make([]byte, 64*1024), 1024*1024)
|
||||
for scanner.Scan() {
|
||||
line := scanner.Text()
|
||||
if !strings.HasPrefix(line, "data:") {
|
||||
continue
|
||||
}
|
||||
payload := strings.TrimSpace(strings.TrimPrefix(line, "data:"))
|
||||
if payload == "[DONE]" {
|
||||
break
|
||||
}
|
||||
var chunk struct {
|
||||
Choices []struct {
|
||||
Delta struct {
|
||||
Content string `json:"content"`
|
||||
} `json:"delta"`
|
||||
} `json:"choices"`
|
||||
}
|
||||
if err := json.Unmarshal([]byte(payload), &chunk); err != nil {
|
||||
continue
|
||||
}
|
||||
if len(chunk.Choices) > 0 && chunk.Choices[0].Delta.Content != "" {
|
||||
onChunk(chunk.Choices[0].Delta.Content)
|
||||
}
|
||||
}
|
||||
return scanner.Err()
|
||||
}
|
||||
|
||||
// ──────────────────────────────────────────────
|
||||
// Embedding
|
||||
// ──────────────────────────────────────────────
|
||||
|
||||
// Embed 批量向量化
|
||||
func (c *Client) Embed(inputs []string) ([][]float64, error) {
|
||||
resp, err := c.post("/embeddings", map[string]any{
|
||||
"model": c.model, "input": inputs,
|
||||
})
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("embedding 服务不可达: %w", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
data, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("读取 embedding 响应失败: %w", err)
|
||||
}
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
return nil, fmt.Errorf("embedding 返回 %d: %s", resp.StatusCode, truncate(string(data), 200))
|
||||
}
|
||||
|
||||
var out struct {
|
||||
Data []struct {
|
||||
Embedding []float64 `json:"embedding"`
|
||||
} `json:"data"`
|
||||
}
|
||||
if err := json.Unmarshal(data, &out); err != nil {
|
||||
return nil, fmt.Errorf("embedding 响应解析失败: %w", err)
|
||||
}
|
||||
res := make([][]float64, len(out.Data))
|
||||
for i, d := range out.Data {
|
||||
res[i] = d.Embedding
|
||||
}
|
||||
return res, nil
|
||||
}
|
||||
|
||||
// ──────────────────────────────────────────────
|
||||
// Fallback 链调用（参考 pj034 chat_completion_with_fallback）
|
||||
// ──────────────────────────────────────────────
|
||||
|
||||
// GenerateWithFallback 依次尝试主路由 + 回退链，首个成功返回
|
||||
func GenerateWithFallback(primary *config.RouteConfig, messages []Message) (string, error) {
|
||||
// 构建路由链
|
||||
chain := []*config.RouteConfig{primary}
|
||||
fallbacks, err := config.GetFallbackRoutes(primary.RouteID)
|
||||
if err == nil && len(fallbacks) > 0 {
|
||||
chain = append(chain, fallbacks...)
|
||||
}
|
||||
|
||||
var lastErr error
|
||||
for _, route := range chain {
|
||||
client := NewClient(route)
|
||||
content, err := client.Generate(messages)
|
||||
if err == nil {
|
||||
return content, nil
|
||||
}
|
||||
lastErr = fmt.Errorf("[%s] %w", route.RouteID, err)
|
||||
}
|
||||
return "", fmt.Errorf("所有路由均失败: %w", lastErr)
|
||||
}
|
||||
|
||||
// GenerateFullWithFallback 非流式 + 回退链：返回结果与「实际命中」的路由（用于审计）
|
||||
func GenerateFullWithFallback(primary *config.RouteConfig, messages []Message) (*ChatResult, *config.RouteConfig, error) {
|
||||
chain := []*config.RouteConfig{primary}
|
||||
if fallbacks, err := config.GetFallbackRoutes(primary.RouteID); err == nil && len(fallbacks) > 0 {
|
||||
chain = append(chain, fallbacks...)
|
||||
}
|
||||
|
||||
var lastErr error
|
||||
for _, route := range chain {
|
||||
client := NewClient(route)
|
||||
res, err := client.GenerateFull(messages)
|
||||
if err == nil {
|
||||
res.Provider = route.Provider
|
||||
return res, route, nil
|
||||
}
|
||||
lastErr = fmt.Errorf("[%s] %w", route.RouteID, err)
|
||||
}
|
||||
return nil, nil, fmt.Errorf("所有路由均失败: %w", lastErr)
|
||||
}
|
||||
|
||||
// ──────────────────────────────────────────────
|
||||
// 配置解析（兼容旧接口）
|
||||
// ──────────────────────────────────────────────
|
||||
|
||||
// LLMConfig 简单的 LLM 连接配置（旧版，逐步淘汰）
|
||||
type LLMConfig struct {
|
||||
BaseURL string
|
||||
APIKey string
|
||||
Model string
|
||||
EmbedModel string
|
||||
MaxTokens int
|
||||
Temperature float64
|
||||
}
|
||||
|
||||
// ResolveLLM 旧版：从 DB/环境 解析 LLM 配置
|
||||
func ResolveLLM(cfg *config.Config) (LLMConfig, bool) {
|
||||
get := func(key, def string) string {
|
||||
var sc model.SystemConfig
|
||||
if err := store.DB.Where("config_key = ?", key).First(&sc).Error; err == nil && strings.TrimSpace(sc.ConfigValue) != "" {
|
||||
return sc.ConfigValue
|
||||
}
|
||||
return def
|
||||
}
|
||||
|
||||
baseURL := get("llm_base_url", cfg.LLMBaseURL)
|
||||
apiKey := get("llm_api_key", cfg.LLMAPIKey)
|
||||
modelName := get("llm_model", cfg.LLMModel)
|
||||
embedModel := get("embed_model", cfg.EmbedModel)
|
||||
|
||||
// DB 空时降级到 JSON secrets（flat 格式，pj034 兼容）
|
||||
if baseURL == "" || apiKey == "" {
|
||||
if baseURL == "" {
|
||||
baseURL = config.GetProviderBaseURL("ollama")
|
||||
}
|
||||
if apiKey == "" {
|
||||
apiKey = config.GetProviderAPIKey("ollama")
|
||||
}
|
||||
}
|
||||
|
||||
c := LLMConfig{
|
||||
BaseURL: baseURL,
|
||||
APIKey: apiKey,
|
||||
Model: modelName,
|
||||
EmbedModel: embedModel,
|
||||
MaxTokens: 2048,
|
||||
Temperature: 0.7,
|
||||
}
|
||||
if c.BaseURL == "" || c.Model == "" {
|
||||
return c, false
|
||||
}
|
||||
return c, true
|
||||
}
|
||||
|
||||
func truncate(s string, n int) string {
|
||||
r := []rune(s)
|
||||
if len(r) <= n {
|
||||
return s
|
||||
}
|
||||
return string(r[:n]) + "..."
|
||||
}
|
||||
@@ -0,0 +1,163 @@
|
||||
package ai
|
||||
|
||||
import (
|
||||
"math"
|
||||
"sort"
|
||||
"strings"
|
||||
|
||||
"eaisalestrain/backend/internal/config"
|
||||
"eaisalestrain/backend/internal/model"
|
||||
"eaisalestrain/backend/internal/store"
|
||||
)
|
||||
|
||||
// Retrieve 混合检索：向量 brute-force 余弦 + 关键词兜底，去重合并取 topK
|
||||
func Retrieve(cfg *config.Config, query string, topK int) []string {
|
||||
if topK <= 0 {
|
||||
topK = 5
|
||||
}
|
||||
var vectorResults []string
|
||||
if route, err := config.GetRoute("embed_gen"); err == nil {
|
||||
if v, err := vectorRetrieve(route, query, topK); err == nil {
|
||||
vectorResults = v
|
||||
}
|
||||
}
|
||||
keywordResults := keywordRetrieve(query, topK*2)
|
||||
|
||||
seen := map[string]bool{}
|
||||
out := make([]string, 0, topK)
|
||||
for _, s := range append(vectorResults, keywordResults...) {
|
||||
s = strings.TrimSpace(s)
|
||||
if s == "" || seen[s] {
|
||||
continue
|
||||
}
|
||||
seen[s] = true
|
||||
out = append(out, s)
|
||||
if len(out) >= topK {
|
||||
break
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// vectorRetrieve 向量检索：query + 所有 chunk 一次批量 embedding，brute-force 余弦 topK
|
||||
func vectorRetrieve(route *config.RouteConfig, query string, topK int) ([]string, error) {
|
||||
var chunks []model.KnowledgeChunk
|
||||
store.DB.Order("id ASC").Find(&chunks)
|
||||
if len(chunks) == 0 {
|
||||
return nil, nil
|
||||
}
|
||||
client := NewClient(route)
|
||||
inputs := make([]string, 0, len(chunks)+1)
|
||||
inputs = append(inputs, query)
|
||||
for _, ch := range chunks {
|
||||
inputs = append(inputs, ch.Content)
|
||||
}
|
||||
vecs, err := client.Embed(inputs)
|
||||
if err != nil || len(vecs) != len(inputs) {
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return nil, nil
|
||||
}
|
||||
qv := vecs[0]
|
||||
type scored struct {
|
||||
idx int
|
||||
sim float64
|
||||
}
|
||||
ss := make([]scored, 0, len(chunks))
|
||||
for i := 1; i < len(vecs); i++ {
|
||||
ss = append(ss, scored{i - 1, cosine(qv, vecs[i])})
|
||||
}
|
||||
sort.Slice(ss, func(a, b int) bool { return ss[a].sim > ss[b].sim })
|
||||
|
||||
seen := map[string]bool{}
|
||||
out := make([]string, 0, topK)
|
||||
for _, s := range ss {
|
||||
content := chunks[s.idx].Content
|
||||
if seen[content] {
|
||||
continue
|
||||
}
|
||||
seen[content] = true
|
||||
out = append(out, content)
|
||||
if len(out) >= topK {
|
||||
break
|
||||
}
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// keywordRetrieve 关键词兜底：term 命中数排序
|
||||
func keywordRetrieve(query string, topK int) []string {
|
||||
terms := splitTerms(query)
|
||||
var chunks []model.KnowledgeChunk
|
||||
store.DB.Order("id ASC").Find(&chunks)
|
||||
|
||||
type scored struct {
|
||||
content string
|
||||
score int
|
||||
}
|
||||
var ss []scored
|
||||
for _, ch := range chunks {
|
||||
s := 0
|
||||
for _, t := range terms {
|
||||
if strings.Contains(ch.Content, t) {
|
||||
s++
|
||||
}
|
||||
}
|
||||
if s > 0 {
|
||||
ss = append(ss, scored{ch.Content, s})
|
||||
}
|
||||
}
|
||||
sort.Slice(ss, func(a, b int) bool { return ss[a].score > ss[b].score })
|
||||
|
||||
seen := map[string]bool{}
|
||||
out := make([]string, 0, topK)
|
||||
for _, s := range ss {
|
||||
if seen[s.content] {
|
||||
continue
|
||||
}
|
||||
seen[s.content] = true
|
||||
out = append(out, s.content)
|
||||
if len(out) >= topK {
|
||||
break
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func splitTerms(q string) []string {
|
||||
f := func(r rune) bool {
|
||||
switch r {
|
||||
case ' ', '，', '。', '？', '！', '、', ',', '.', '?', '!', ':', '：', ';', '；':
|
||||
return true
|
||||
}
|
||||
return false
|
||||
}
|
||||
terms := strings.FieldsFunc(q, f)
|
||||
var out []string
|
||||
for _, t := range terms {
|
||||
if len([]rune(t)) >= 2 {
|
||||
out = append(out, t)
|
||||
}
|
||||
}
|
||||
if len(out) == 0 {
|
||||
out = []string{q}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func cosine(a, b []float64) float64 {
|
||||
if len(a) == 0 || len(a) != len(b) {
|
||||
return 0
|
||||
}
|
||||
var dot, na, nb float64
|
||||
for i := range a {
|
||||
dot += a[i] * b[i]
|
||||
na += a[i] * a[i]
|
||||
nb += b[i] * b[i]
|
||||
}
|
||||
if na == 0 || nb == 0 {
|
||||
return 0
|
||||
}
|
||||
return dot / (math.Sqrt(na) * math.Sqrt(nb))
|
||||
}
|
||||
@@ -0,0 +1,42 @@
|
||||
package ai
|
||||
|
||||
import (
|
||||
"reflect"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestSplitTerms(t *testing.T) {
|
||||
got := splitTerms("这是 测试 内容")
|
||||
want := []string{"这是", "测试", "内容"}
|
||||
if !reflect.DeepEqual(got, want) {
|
||||
t.Errorf("splitTerms = %v want %v", got, want)
|
||||
}
|
||||
// 单字词被过滤
|
||||
got = splitTerms("我 喜欢 它")
|
||||
if !reflect.DeepEqual(got, []string{"喜欢"}) {
|
||||
t.Errorf("splitTerms single-char = %v", got)
|
||||
}
|
||||
// 全部被过滤时回退为整句
|
||||
got = splitTerms("我 你 他")
|
||||
if !reflect.DeepEqual(got, []string{"我 你 他"}) {
|
||||
t.Errorf("splitTerms fallback = %v", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCosine(t *testing.T) {
|
||||
if got := cosine([]float64{1, 0}, []float64{1, 0}); got != 1 {
|
||||
t.Errorf("cosine identical = %v", got)
|
||||
}
|
||||
if got := cosine([]float64{1, 0}, []float64{0, 1}); got != 0 {
|
||||
t.Errorf("cosine orthogonal = %v", got)
|
||||
}
|
||||
if got := cosine([]float64{1}, []float64{1, 0}); got != 0 {
|
||||
t.Errorf("cosine length-mismatch = %v", got)
|
||||
}
|
||||
if got := cosine([]float64{}, []float64{}); got != 0 {
|
||||
t.Errorf("cosine empty = %v", got)
|
||||
}
|
||||
if got := cosine([]float64{0, 0}, []float64{1, 1}); got != 0 {
|
||||
t.Errorf("cosine zero-norm = %v", got)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user