lian/mchat
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

service 修改 Redis 存储 KV

Browse Source
This commit is contained in:
1iaan
2026-04-10 11:12:10 +00:00
parent c888ca8844
commit bc82e3e708
25 changed files with 322 additions and 3666 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
.gitignore vendored
View File

@@ -2,3 +2,5 @@ docs
.workspace.codex
.env
.vscode
faiss

20
ai-chat-service/chat-server/chat-context/chat_context.go
View File

@@ -1,20 +0,0 @@
package chat_context
import "github.com/sashabaranov/go-openai"
type ChatMessage struct {
//当前记录ID
ID string `json:"id,omitempty"`
//上一条记录ID
PID string `json:"pid,omitempty"`
//消息内容
Message openai.ChatCompletionMessage `json:"message"`
//该消息tokens数
Tokens int `json:"tokens,omitempty"`
}
type ContextCache interface {
Get(key string) (*ChatMessage, error)
Set(key string, value *ChatMessage, ttl int) error
Close()
}

51
ai-chat-service/chat-server/chat-context/redis.go
View File

@@ -1,51 +0,0 @@
package chat_context
import (
predis "ai-chat-service/pkg/db/redis"
"context"
"encoding/json"
"time"
"github.com/redis/go-redis/v9"
)
type redisCache struct {
redisClient *redis.Client
}
func NewRedisCache() ContextCache {
pool := predis.GetPool()
return &redisCache{
redisClient: pool.Get(),
}
}
func getRedisKey(key string) string {
return predis.GetKey(key)
}
func (c *redisCache) Get(key string) (*ChatMessage, error) {
key = getRedisKey(key)
str, err := c.redisClient.Get(context.Background(), key).Result()
if err == redis.Nil {
return nil, nil
}
if err != nil {
return nil, err
}
value := &ChatMessage{}
err = json.Unmarshal([]byte(str), value)
return value, err
}
func (c *redisCache) Set(key string, value *ChatMessage, ttl int) error {
key = getRedisKey(key)
bytes, err := json.Marshal(value)
if err != nil {
return err
}
str := string(bytes)
return c.redisClient.SetEx(context.Background(), key, str, time.Duration(ttl)*time.Second).Err()
}
func (c *redisCache) Close() {
pool := predis.GetPool()
pool.Put(c.redisClient)
}

69
ai-chat-service/chat-server/data/chat_records.go
View File

@@ -1,54 +1,59 @@
package data
import (
"database/sql"
"strings"
predis "ai-chat-service/pkg/db/redis"
"context"
"encoding/json"
redis "github.com/redis/go-redis/v9"
)
type IChatRecordsData interface {
Add(record *ChatRecord) error
GetById(id int64) (record *ChatRecord, err error)
GetById(id string) (record *ChatRecord, err error)
}
type ChatRecord struct {
ID int64 `json:"id"`
UserMsg string `json:"user_msg"`
UserMsgTokens int `json:"user_msg_tokens"`
UserMsgKeywords []string `json:"user_msg_keywords"`
AIMsg string `json:"ai_msg"`
AIMsgTokens int `json:"ai_msg_tokens"`
ReqTokens int `json:"req_tokens"`
CreateAt int64 `json:"create_at"`
ID string `json:"-"`
Question string `json:"q"`
Answer string `json:"a"`
}
type chatRecordsData struct {
db *sql.DB
type chatRecordsData struct{}
func NewChatRecordsData() IChatRecordsData {
return &chatRecordsData{}
}
func NewChatRecordsData(db *sql.DB) IChatRecordsData {
return &chatRecordsData{
db: db,
}
}
func (data *chatRecordsData) Add(record *ChatRecord) error {
client := predis.GetPool().Get()
defer predis.GetPool().Put(client)
func (data *chatRecordsData) Add(cr *ChatRecord) (err error) {
sqlStr := "insert into chat_records(user_msg,user_msg_tokens,user_msg_keywords,ai_msg,ai_msg_tokens,req_tokens,create_at)values(?,?,?,?,?,?,?)"
res, err := data.db.Exec(sqlStr, cr.UserMsg, cr.UserMsgTokens, strings.Join(cr.UserMsgKeywords, ","), cr.AIMsg, cr.AIMsgTokens, cr.ReqTokens, cr.CreateAt)
payload, err := json.Marshal(&ChatRecord{
Question: record.Question,
Answer: record.Answer,
})
if err != nil {
return
return err
}
cr.ID, _ = res.LastInsertId()
return
return client.Set(context.Background(), predis.GetKey("qa", record.ID), string(payload), 0).Err()
}
func (data *chatRecordsData) GetById(id int64) (cr *ChatRecord, err error) {
sqlStr := "select id,user_msg,user_msg_tokens,user_msg_keywords,ai_msg,ai_msg_tokens,req_tokens,create_at from chat_records where id = ?"
row := data.db.QueryRow(sqlStr, id)
cr = &ChatRecord{}
var keywords string
err = row.Scan(&cr.ID, &cr.UserMsg, &cr.UserMsgTokens, &keywords, &cr.AIMsg, &cr.AIMsgTokens, &cr.ReqTokens, &cr.CreateAt)
func (data *chatRecordsData) GetById(id string) (*ChatRecord, error) {
client := predis.GetPool().Get()
defer predis.GetPool().Put(client)
value, err := client.Get(context.Background(), predis.GetKey("qa", id)).Result()
if err == redis.Nil {
return nil, nil
}
if err != nil {
return nil, err
}
cr.UserMsgKeywords = strings.Split(keywords, ",")
return cr, err
record := &ChatRecord{ID: id}
if err = json.Unmarshal([]byte(value), record); err != nil {
return nil, err
}
return record, nil
}

13
ai-chat-service/chat-server/main.go
View File

@@ -5,13 +5,13 @@ import (
metrics_app "ai-chat-service/chat-server/metrics-app"
metrics_bus "ai-chat-service/chat-server/metrics-bus"
"ai-chat-service/chat-server/server"
vector_data "ai-chat-service/chat-server/vector-data"
"ai-chat-service/interceptor"
"ai-chat-service/pkg/config"
"ai-chat-service/pkg/db/mysql"
"ai-chat-service/pkg/db/redis"
"ai-chat-service/pkg/log"
"ai-chat-service/proto"
"ai-chat-service/services/embedding"
"ai-chat-service/services/faiss"
"flag"
"fmt"
"net/http"
@@ -52,23 +52,22 @@ func main() {
logger.SetOutput(log.GetRotateWriter(cnf.Log.LogPath))
logger.SetPrintCaller(true)
// 初始化Mysql
mysql.InitMysql(cnf)
// 初始化redis
redis.InitRedisPool(cnf)
recordsData := data.NewChatRecordsData(mysql.GetDB())
vectorRecordsData, err := vector_data.NewChatRecordsData(cnf)
recordsData := data.NewChatRecordsData()
embedder, err := embedding.NewEmbedder(cnf)
if err != nil {
log.Fatal(err)
}
faissClient := faiss.NewClient(cnf)
lis, err := net.Listen("tcp", fmt.Sprintf("%s:%d", cnf.Server.IP, cnf.Server.Port))
if err != nil {
log.Fatal(err)
}
s := grpc.NewServer(grpc.UnaryInterceptor(interceptor.UnaryAuthInterceptor), grpc.StreamInterceptor(metrics_app.NewStreamMiddleware(registry).WrapHandler()))
service := server.NewChatService(recordsData, vectorRecordsData, cnf, logger, busMetrics)
service := server.NewChatService(recordsData, embedder, faissClient, cnf, logger, busMetrics)
proto.RegisterChatServer(s, service)
healthCheckSrv := health.NewServer()

138
ai-chat-service/chat-server/server/app.go
View File

@@ -1,7 +1,6 @@
package server
import (
chat_context "ai-chat-service/chat-server/chat-context"
"ai-chat-service/pkg/config"
"ai-chat-service/pkg/log"
"ai-chat-service/pkg/zerror"
@@ -29,18 +28,15 @@ type openaiConf struct {
PresencePenalty float32
FrequencyPenalty float32
BotDesc string
ContextTTL int
ContextLen int
MinResponseTokens int
}
type app struct {
openaiConf *openaiConf
log log.ILogger
// TODO 内容上下文对象
contextCache chat_context.ContextCache
}
func (s *chatService) newApp(in *proto.ChatCompletionRequest, contextCache chat_context.ContextCache) *app {
func (s *chatService) newApp(in *proto.ChatCompletionRequest) *app {
conf := &openaiConf{
ApiKey: s.config.Chat.ApiKey,
BaseUrl: s.config.Chat.BaseUrl,
@@ -51,8 +47,6 @@ func (s *chatService) newApp(in *proto.ChatCompletionRequest, contextCache chat_
PresencePenalty: s.config.Chat.PresencePenalty,
FrequencyPenalty: s.config.Chat.FrequencyPenalty,
BotDesc: s.config.Chat.BotDesc,
ContextTTL: s.config.Chat.ContextTTL,
ContextLen: s.config.Chat.ContextLen,
MinResponseTokens: s.config.Chat.MinResponseTokens,
}
if in.ChatParam != nil {
@@ -69,40 +63,29 @@ func (s *chatService) newApp(in *proto.ChatCompletionRequest, contextCache chat_
if in.ChatParam.MaxTokens != 0 {
conf.MaxTokens = int(in.ChatParam.MaxTokens)
}
if in.ChatParam.ContextTTL != 0 {
conf.ContextTTL = int(in.ChatParam.ContextTTL)
}
if in.ChatParam.ContextLen != 0 {
conf.ContextLen = int(in.ChatParam.ContextLen)
}
if in.ChatParam.MinResponseTokens != 0 {
conf.MinResponseTokens = int(in.ChatParam.MinResponseTokens)
}
}
return &app{
openaiConf: conf,
log: s.log,
contextCache: contextCache,
openaiConf: conf,
log: s.log,
}
}
func (a *app) getOpenaiClient() *openai.Client {
accessToken := a.openaiConf.ApiKey
config := openai.DefaultConfig(accessToken)
config.BaseURL = a.openaiConf.BaseUrl
client := openai.NewClientWithConfig(config)
return client
conf := openai.DefaultConfig(a.openaiConf.ApiKey)
conf.BaseURL = a.openaiConf.BaseUrl
return openai.NewClientWithConfig(conf)
}
func (a *app) buildChatCompletionRequest(in *proto.ChatCompletionRequest, stream bool) (req openai.ChatCompletionRequest, tokens, currTokens int, currMessage openai.ChatCompletionMessage, err error) {
//当前消息
currMessage = openai.ChatCompletionMessage{
Role: openai.ChatMessageRoleUser,
Content: in.Message,
}
req = openai.ChatCompletionRequest{
Model: a.openaiConf.Model,
Messages: []openai.ChatCompletionMessage{
currMessage,
},
Model: a.openaiConf.Model,
MaxTokens: a.openaiConf.MinResponseTokens,
Temperature: a.openaiConf.Temperature,
TopP: a.openaiConf.TopP,
@@ -110,13 +93,7 @@ func (a *app) buildChatCompletionRequest(in *proto.ChatCompletionRequest, stream
FrequencyPenalty: a.openaiConf.FrequencyPenalty,
Stream: stream,
}
contextList := make([]*chat_context.ChatMessage, 0)
if in.EnableContext {
//从缓存中获取上下文信息
contextList = a.getContext(in.Pid)
}
//重构req.Messages
tokens, currTokens, req.Messages, err = a.rebuildMessages(contextList, currMessage)
tokens, currTokens, req.Messages, err = a.rebuildMessages(currMessage)
if err != nil {
a.log.Error(err)
return
@@ -124,51 +101,37 @@ func (a *app) buildChatCompletionRequest(in *proto.ChatCompletionRequest, stream
req.MaxTokens = a.openaiConf.MaxTokens - tokens
return
}
func (a *app) rebuildMessages(contextList []*chat_context.ChatMessage, currMessage openai.ChatCompletionMessage) (tokens, currTokens int, messages []openai.ChatCompletionMessage, err error) {
var sysMessage openai.ChatCompletionMessage
func (a *app) rebuildMessages(currMessage openai.ChatCompletionMessage) (tokens, currTokens int, messages []openai.ChatCompletionMessage, err error) {
messages = make([]openai.ChatCompletionMessage, 0, 2)
botTokens := 0
if a.openaiConf.BotDesc != "" {
sysMessage = openai.ChatCompletionMessage{
sysMessage := openai.ChatCompletionMessage{
Role: openai.ChatMessageRoleSystem,
Content: a.openaiConf.BotDesc,
}
botTokens, err = tokenizer.GetTokens(&sysMessage, a.openaiConf.Model)
if err != nil {
a.log.Error(err)
return
}
messages = append(messages, sysMessage)
}
messages = []openai.ChatCompletionMessage{currMessage}
currTokens, err = tokenizer.GetTokens(&currMessage, a.openaiConf.Model)
if err != nil {
a.log.Error(err)
return
}
if currTokens > a.openaiConf.MaxTokens-a.openaiConf.MinResponseTokens-botTokens-ChatPrimedTokens {
err = zerror.NewByMsg("请求消息超限")
a.log.Error(err)
return
return 0, 0, nil, zerror.NewByMsg("请求消息超限")
}
tokens = currTokens + botTokens + ChatPrimedTokens
if contextList != nil {
for _, item := range contextList {
if tokens+item.Tokens+ChatPrimedTokens > a.openaiConf.MaxTokens-a.openaiConf.MinResponseTokens {
break
}
messages = append(messages, item.Message)
tokens += item.Tokens + ChatPrimedTokens
}
}
for i, j := 0, len(messages)-1; i < j; i, j = i+1, j-1 {
messages[i], messages[j] = messages[j], messages[i]
}
if botTokens > 0 {
messages = append([]openai.ChatCompletionMessage{sysMessage}, messages...)
}
messages = append(messages, currMessage)
return
}
func (a *app) buildChatCompletionResponse(msg string) *proto.ChatCompletionResponse {
res := &proto.ChatCompletionResponse{
return &proto.ChatCompletionResponse{
Id: uuid.New().String(),
Object: "chat.completion",
Created: time.Now().Unix(),
@@ -182,17 +145,12 @@ func (a *app) buildChatCompletionResponse(msg string) *proto.ChatCompletionRespo
FinishReason: "stop",
},
},
Usage: &proto.Usage{
PromptTokens: 0,
CompletionTokens: 0,
TotalTokens: 0,
},
Usage: &proto.Usage{},
}
return res
}
func (a *app) buildChatCompletionStreamResponse(id, delta, finishReason string) *proto.ChatCompletionStreamResponse {
res := &proto.ChatCompletionStreamResponse{
return &proto.ChatCompletionStreamResponse{
Id: id,
Object: "chat.completion.chunk",
Created: time.Now().Unix(),
@@ -208,79 +166,49 @@ func (a *app) buildChatCompletionStreamResponse(id, delta, finishReason string)
},
},
}
return res
}
func (a *app) buildChatCompletionStreamResponseList(id, msg string) []*proto.ChatCompletionStreamResponse {
list := make([]*proto.ChatCompletionStreamResponse, 0)
list := make([]*proto.ChatCompletionStreamResponse, 0, len(msg))
for _, delta := range msg {
list = append(list, a.buildChatCompletionStreamResponse(id, string(delta), ""))
}
return list
}
func (a *app) getContext(id string) []*chat_context.ChatMessage {
maxLen := a.openaiConf.ContextLen
list := make([]*chat_context.ChatMessage, 0, maxLen)
key := id
for i := 0; i < maxLen; i++ {
value, err := a.contextCache.Get(key)
if err != nil {
a.log.Error(err)
return nil
}
if value == nil {
break
}
list = append(list, value)
key = value.PID
}
return list
}
func (a *app) saveContext(value *chat_context.ChatMessage) error {
err := a.contextCache.Set(value.ID, value, a.openaiConf.ContextTTL)
if err != nil {
a.log.Error(err)
return err
}
return nil
}
func (a *app) keywords(in *proto.ChatCompletionRequest) []string {
pool := keywords_filter.GetKeywordsClientPool()
conn := pool.Get()
defer pool.Put(conn)
accessToken := config.GetConfig().DependOn.Keywords.AccessToken
client := keywords_proto.NewFilterClient(conn)
ctx := services.AppendBearerTokenToContext(context.Background(), accessToken)
req := &keywords_proto.FilterReq{
Text: in.Message,
}
req := &keywords_proto.FilterReq{Text: in.Message}
res, err := client.FindAll(ctx, req)
if err != nil {
a.log.Error(err)
return []string{}
}
return res.Keywords
}
func (a *app) sensitive(in *proto.ChatCompletionRequest) (ok bool, msg string, err error) {
pool := keywords_filter.GetSensitiveClientPool()
conn := pool.Get()
defer pool.Put(conn)
accessToken := config.GetConfig().DependOn.Sensitive.AccessToken
client := keywords_proto.NewFilterClient(conn)
ctx := services.AppendBearerTokenToContext(context.Background(), accessToken)
req := &keywords_proto.FilterReq{
Text: in.Message,
}
req := &keywords_proto.FilterReq{Text: in.Message}
res, err := client.Validate(ctx, req)
if err != nil {
a.log.Error(err)
return false, "", err
}
ok = res.Ok
if !ok {
msg = "触发到了知识盲区，请换个问题再问"
if !res.Ok {
return false, "触发到了知识盲区，请换个问题再问", nil
}
return
return true, "", nil
}

412
ai-chat-service/chat-server/server/server.go
View File

@@ -1,20 +1,18 @@
package server
import (
chat_context "ai-chat-service/chat-server/chat-context"
"ai-chat-service/chat-server/data"
metrics_bus "ai-chat-service/chat-server/metrics-bus"
vector_data "ai-chat-service/chat-server/vector-data"
"ai-chat-service/pkg/config"
"ai-chat-service/pkg/log"
"ai-chat-service/proto"
"ai-chat-service/services/embedding"
"ai-chat-service/services/faiss"
"ai-chat-service/services/tokenizer"
"context"
"encoding/json"
"io"
"strconv"
"strings"
"time"
"github.com/golang/protobuf/jsonpb"
"github.com/google/uuid"
@@ -26,218 +24,127 @@ type chatService struct {
config *config.Config
log log.ILogger
data data.IChatRecordsData
vectorData vector_data.IChatRecordsData
embedder embedding.Embedder
faiss faiss.Client
busMetrics *metrics_bus.BusMetrics
}
func NewChatService(data data.IChatRecordsData, vectorData vector_data.IChatRecordsData, config *config.Config, log log.ILogger, busMetrics *metrics_bus.BusMetrics) proto.ChatServer {
func NewChatService(data data.IChatRecordsData, embedder embedding.Embedder, faissClient faiss.Client, config *config.Config, log log.ILogger, busMetrics *metrics_bus.BusMetrics) proto.ChatServer {
return &chatService{
config: config,
log: log,
data: data,
vectorData: vectorData,
embedder: embedder,
faiss: faissClient,
busMetrics: busMetrics,
}
}
func (s *chatService) ChatCompletion(ctx context.Context, in *proto.ChatCompletionRequest) (*proto.ChatCompletionResponse, error) {
redisContextCache := chat_context.NewRedisCache()
defer redisContextCache.Close()
app := s.newApp(in)
app := s.newApp(in, redisContextCache)
//敏感词过滤
ok, msg, err := app.sensitive(in)
if err != nil {
s.log.Error(err)
return nil, err
}
if !ok {
res := app.buildChatCompletionResponse(msg)
return res, nil
}
//关键词提取
keywords := app.keywords(in)
if len(keywords) > 0 {
idStr, score, err := s.vectorData.QueryData(context.Background(), map[string][]string{"keywords": {strings.Join(keywords, ",")}})
if err != nil {
s.log.Error(err)
} else if score > s.config.Vector.Threshold {
id, err := strconv.ParseInt(idStr, 10, 64)
if err != nil {
s.log.Error(err)
} else {
record, err := s.data.GetById(id)
if err != nil {
s.log.Error(err)
} else {
res := app.buildChatCompletionResponse(record.AIMsg)
return res, nil
}
}
}
}
client := app.getOpenaiClient()
req, tokens, currTokens, currMessage, err := app.buildChatCompletionRequest(in, false)
resp, err := client.CreateChatCompletion(ctx, req)
if err != nil {
s.log.Error(err)
return nil, err
}
res := &proto.ChatCompletionResponse{}
bytes, err := json.Marshal(resp)
if err != nil {
s.log.Error(err)
return nil, err
}
err = jsonpb.UnmarshalString(string(bytes), res)
if err != nil {
s.log.Error(err)
return nil, err
}
go func() {
reqContext := &chat_context.ChatMessage{
ID: in.Id,
PID: in.Pid,
Message: currMessage,
Tokens: currTokens,
}
err := app.saveContext(reqContext)
if err != nil {
s.log.Error(err)
return
}
resContext := &chat_context.ChatMessage{
ID: resp.ID,
PID: reqContext.ID,
Message: resp.Choices[0].Message,
Tokens: resp.Usage.CompletionTokens,
}
err = app.saveContext(resContext)
if err != nil {
s.log.Error(err)
return
}
}()
go func() {
records := &data.ChatRecord{
UserMsg: in.Message,
UserMsgTokens: currTokens,
UserMsgKeywords: keywords,
AIMsg: resp.Choices[0].Message.Content,
AIMsgTokens: resp.Usage.CompletionTokens,
ReqTokens: tokens,
CreateAt: time.Now().Unix(),
}
err := s.data.Add(records)
if err != nil {
s.log.Error(err)
return
}
//保存到向量数据库
if len(keywords) > 0 {
list := []*vector_data.ChatRecord{
{
ID: strconv.FormatInt(records.ID, 10),
KVs: map[string]string{
"keywords": strings.Join(keywords, ","),
},
},
}
err = s.vectorData.UpsertData(context.Background(), list)
if err != nil {
s.log.Error(err)
return
}
}
}()
return res, err
}
func (s *chatService) ChatCompletionStream(in *proto.ChatCompletionRequest, stream proto.Chat_ChatCompletionStreamServer) error {
redisContextCache := chat_context.NewRedisCache()
defer redisContextCache.Close()
app := s.newApp(in, redisContextCache)
//敏感词过滤
ok, msg, err := app.sensitive(in)
if err != nil {
s.busMetrics.ErrQuestionsTotalCounter.Inc()
return nil, err
}
if !ok {
s.busMetrics.SensitiveQuestionsTotalCounter.Inc()
return app.buildChatCompletionResponse(msg), nil
}
keywords := app.keywords(in)
if len(keywords) > 0 {
s.busMetrics.KeywordsQuestionsTotalCounter.Inc()
}
req, _, _, _, err := app.buildChatCompletionRequest(in, false)
if err != nil {
s.busMetrics.ErrQuestionsTotalCounter.Inc()
return nil, err
}
questionEmbedding, cachedRecord := s.searchCachedAnswer(ctx, in.Message)
if cachedRecord != nil {
return app.buildChatCompletionResponse(cachedRecord.Answer), nil
}
client := app.getOpenaiClient()
resp, err := client.CreateChatCompletion(ctx, req)
if err != nil {
s.busMetrics.ErrQuestionsTotalCounter.Inc()
s.log.Error(err)
return nil, err
}
res := &proto.ChatCompletionResponse{}
bytes, err := json.Marshal(resp)
if err != nil {
s.busMetrics.ErrQuestionsTotalCounter.Inc()
return nil, err
}
if err = jsonpb.UnmarshalString(string(bytes), res); err != nil {
s.busMetrics.ErrQuestionsTotalCounter.Inc()
return nil, err
}
if len(resp.Choices) > 0 {
if err = s.persistQA(ctx, questionEmbedding, in.Message, resp.Choices[0].Message.Content); err != nil {
s.log.Error(err)
} else {
s.busMetrics.QuestionsTotalCounter.Inc()
}
}
return res, nil
}
func (s *chatService) ChatCompletionStream(in *proto.ChatCompletionRequest, stream proto.Chat_ChatCompletionStreamServer) error {
app := s.newApp(in)
ok, msg, err := app.sensitive(in)
if err != nil {
s.busMetrics.ErrQuestionsTotalCounter.Inc()
return err
}
if !ok {
s.busMetrics.SensitiveQuestionsTotalCounter.Inc()
resId := uuid.New().String()
startRes := app.buildChatCompletionStreamResponse(resId, "", "")
endRes := app.buildChatCompletionStreamResponse(resId, "", "stop")
err = stream.Send(startRes)
if err != nil {
s.log.Error(err)
resID := uuid.New().String()
if err = stream.Send(app.buildChatCompletionStreamResponse(resID, "", "")); err != nil {
return err
}
resList := app.buildChatCompletionStreamResponseList(resId, msg)
for _, res := range resList {
err = stream.Send(res)
if err != nil {
s.log.Error(err)
for _, res := range app.buildChatCompletionStreamResponseList(resID, msg) {
if err = stream.Send(res); err != nil {
return err
}
}
err = stream.Send(endRes)
if err != nil {
s.log.Error(err)
return err
}
return nil
return stream.Send(app.buildChatCompletionStreamResponse(resID, "", "stop"))
}
//关键词提取
keywords := app.keywords(in)
if len(keywords) > 0 {
s.busMetrics.KeywordsQuestionsTotalCounter.Inc()
idStr, score, err := s.vectorData.QueryData(context.Background(), map[string][]string{"keywords": {strings.Join(keywords, ",")}})
if err != nil {
s.log.Error(err)
} else if score > s.config.Vector.Threshold {
id, err := strconv.ParseInt(idStr, 10, 64)
if err != nil {
s.log.Error(err)
} else {
record, err := s.data.GetById(id)
if err != nil {
s.log.Error(err)
} else {
resId := uuid.New().String()
startRes := app.buildChatCompletionStreamResponse(resId, "", "")
endRes := app.buildChatCompletionStreamResponse(resId, "", "stop")
err = stream.Send(startRes)
if err != nil {
s.log.Error(err)
return err
}
resList := app.buildChatCompletionStreamResponseList(resId, record.AIMsg)
for _, res := range resList {
err = stream.Send(res)
if err != nil {
s.log.Error(err)
return err
}
}
err = stream.Send(endRes)
if err != nil {
s.log.Error(err)
return err
}
return nil
}
}
req, _, _, _, err := app.buildChatCompletionRequest(in, true)
if err != nil {
s.busMetrics.ErrQuestionsTotalCounter.Inc()
return err
}
questionEmbedding, cachedRecord := s.searchCachedAnswer(stream.Context(), in.Message)
if cachedRecord != nil {
if err = stream.Send(app.buildChatCompletionStreamResponse(cachedRecord.ID, "", "")); err != nil {
return err
}
for _, res := range app.buildChatCompletionStreamResponseList(cachedRecord.ID, cachedRecord.Answer) {
if err = stream.Send(res); err != nil {
return err
}
}
return stream.Send(app.buildChatCompletionStreamResponse(cachedRecord.ID, "", "stop"))
}
client := app.getOpenaiClient()
req, tokens, currTokens, currMessage, err := app.buildChatCompletionRequest(in, false)
chatStream, err := client.CreateChatCompletionStream(stream.Context(), req)
if err != nil {
s.busMetrics.ErrQuestionsTotalCounter.Inc()
@@ -245,109 +152,106 @@ func (s *chatService) ChatCompletionStream(in *proto.ChatCompletionRequest, stre
return err
}
defer chatStream.Close()
completionContent := ""
resultID := ""
for {
resp, err := chatStream.Recv()
if err != nil && err != io.EOF {
s.busMetrics.ErrQuestionsTotalCounter.Inc()
s.log.Error(err)
return err
}
if err == io.EOF {
break
}
if resultID == "" {
resultID = resp.ID
}
completionContent += resp.Choices[0].Delta.Content
res := &proto.ChatCompletionStreamResponse{}
bytes, err := json.Marshal(resp)
if err != nil {
s.log.Error(err)
s.busMetrics.ErrQuestionsTotalCounter.Inc()
return err
}
err = jsonpb.UnmarshalString(string(bytes), res)
if err != nil {
s.log.Error(err)
if err = jsonpb.UnmarshalString(string(bytes), res); err != nil {
s.busMetrics.ErrQuestionsTotalCounter.Inc()
return err
}
err = stream.Send(res)
if err != nil {
s.log.Error(err)
if err = stream.Send(res); err != nil {
return err
}
}
model := s.config.Chat.Model
if in.ChatParam != nil && in.ChatParam.Model != "" {
model = in.ChatParam.Model
}
resultMessage := openai.ChatCompletionMessage{
Role: openai.ChatMessageRoleAssistant,
Content: completionContent,
}
model := s.config.Chat.Model
if in.ChatParam != nil && in.ChatParam.Model != "" {
model = in.ChatParam.Model
}
resultTokens, err := tokenizer.GetTokens(&resultMessage, model)
if err != nil {
if _, err = tokenizer.GetTokens(&resultMessage, model); err != nil {
s.busMetrics.ErrQuestionsTotalCounter.Inc()
s.log.Error(err)
return err
}
go func() {
reqContext := &chat_context.ChatMessage{
ID: in.Id,
PID: in.Pid,
Message: currMessage,
Tokens: currTokens,
}
err := app.saveContext(reqContext)
if err != nil {
s.log.Error(err)
return
}
resContext := &chat_context.ChatMessage{
ID: resultID,
PID: reqContext.ID,
Message: resultMessage,
Tokens: resultTokens,
}
err = app.saveContext(resContext)
if err != nil {
s.log.Error(err)
return
}
}()
go func() {
if err = s.persistQA(stream.Context(), questionEmbedding, in.Message, completionContent); err != nil {
s.log.Error(err)
} else {
s.busMetrics.QuestionsTotalCounter.Inc()
records := &data.ChatRecord{
UserMsg: in.Message,
UserMsgTokens: currTokens,
UserMsgKeywords: keywords,
AIMsg: completionContent,
AIMsgTokens: resultTokens,
ReqTokens: tokens,
CreateAt: time.Now().Unix(),
}
err := s.data.Add(records)
if err != nil {
s.log.Error(err)
return
}
//保存到向量数据库
if len(keywords) > 0 {
list := []*vector_data.ChatRecord{
{
ID: strconv.FormatInt(records.ID, 10),
KVs: map[string]string{
"keywords": strings.Join(keywords, ","),
},
},
}
err = s.vectorData.UpsertData(context.Background(), list)
if err != nil {
s.log.Error(err)
return
}
}
}()
}
return nil
}
func (s *chatService) searchCachedAnswer(ctx context.Context, question string) ([]float32, *data.ChatRecord) {
embeddingVector, err := s.embedder.Embed(ctx, question)
if err != nil {
s.log.Error(err)
return nil, nil
}
searchRes, err := s.faiss.Search(ctx, embeddingVector, s.config.Faiss.SearchK)
if err != nil {
s.log.Error(err)
return embeddingVector, nil
}
if searchRes == nil || len(searchRes.IDs) == 0 || len(searchRes.SimilarityScores) == 0 {
return embeddingVector, nil
}
limit := len(searchRes.IDs)
if len(searchRes.SimilarityScores) < limit {
limit = len(searchRes.SimilarityScores)
}
for i := 0; i < limit; i++ {
if searchRes.IDs[i] < 0 || searchRes.SimilarityScores[i] < s.config.Faiss.SimilarityThreshold {
continue
}
record, err := s.data.GetById(strconv.FormatInt(searchRes.IDs[i], 10))
if err != nil {
s.log.Error(err)
return embeddingVector, nil
}
if record != nil {
return embeddingVector, record
}
}
return embeddingVector, nil
}
func (s *chatService) persistQA(ctx context.Context, questionEmbedding []float32, question, answer string) error {
if len(questionEmbedding) == 0 {
vector, err := s.embedder.Embed(ctx, question)
if err != nil {
return err
}
questionEmbedding = vector
}
id, err := s.faiss.Insert(ctx, questionEmbedding)
if err != nil {
return err
}
return s.data.Add(&data.ChatRecord{
ID: id,
Question: question,
Answer: answer,
})
}

29
ai-chat-service/chat-server/vector-data/chat_records.go
View File

@@ -1,29 +0,0 @@
package vector_data
import (
"ai-chat-service/pkg/config"
"context"
"fmt"
)
const CHAT_RECORDS = "chat_records"
type ChatRecord struct {
ID string
KVs map[string]string
}
type IChatRecordsData interface {
UpsertData(ctx context.Context, list []*ChatRecord) error
QueryData(ctx context.Context, text map[string][]string) (id string, score float32, err error)
}
func NewChatRecordsData(config *config.Config) (IChatRecordsData, error) {
switch config.Vector.Provider {
case "tencent", "":
return newTencentChatRecordsData(config)
case "pgvector":
return newPgvectorChatRecordsData(config)
default:
return nil, fmt.Errorf("unsupported vector provider: %s", config.Vector.Provider)
}
}

121
ai-chat-service/chat-server/vector-data/pgvector.go
View File

@@ -1,121 +0,0 @@
package vector_data
import (
"ai-chat-service/pkg/config"
"ai-chat-service/services/embedding"
"context"
"fmt"
"strconv"
"strings"
"time"
"github.com/jackc/pgx"
)
type pgvectorChatRecordsData struct {
config *config.Config
pool *pgx.ConnPool
embedder embedding.Embedder
}
func newPgvectorChatRecordsData(config *config.Config) (IChatRecordsData, error) {
connConfig, err := pgx.ParseConnectionString(config.Vector.Pgvector.DSN)
if err != nil {
return nil, err
}
pool, err := pgx.NewConnPool(pgx.ConnPoolConfig{
ConnConfig: connConfig,
MaxConnections: config.Vector.Pgvector.MaxOpenConn,
})
if err != nil {
return nil, err
}
embedder, err := embedding.NewEmbedder(config)
if err != nil {
pool.Close()
return nil, err
}
return &pgvectorChatRecordsData{
config: config,
pool: pool,
embedder: embedder,
}, nil
}
func (data *pgvectorChatRecordsData) UpsertData(ctx context.Context, list []*ChatRecord) error {
table := data.config.Vector.Pgvector.Table
if table == "" {
table = "chat_record_vectors"
}
for _, item := range list {
recordID, err := strconv.ParseInt(item.ID, 10, 64)
if err != nil {
return err
}
keywordsText := embedding.BuildText(item.KVs["keywords"])
if keywordsText == "" {
continue
}
vector, err := data.embedder.Embed(ctx, keywordsText)
if err != nil {
return err
}
commandTag, err := data.pool.Exec(
fmt.Sprintf(
"INSERT INTO %s (record_id, keywords_text, embedding, created_at) VALUES ($1, $2, $3::vector, $4) ON CONFLICT (record_id) DO UPDATE SET keywords_text = EXCLUDED.keywords_text, embedding = EXCLUDED.embedding, created_at = EXCLUDED.created_at",
table,
),
recordID,
keywordsText,
vectorLiteral(vector),
time.Now().Unix(),
)
if err != nil {
return err
}
if commandTag.RowsAffected() == 0 {
return fmt.Errorf("pgvector upsert affected 0 rows for record_id=%d", recordID)
}
}
return nil
}
func (data *pgvectorChatRecordsData) QueryData(ctx context.Context, text map[string][]string) (id string, score float32, err error) {
keywordsText := embedding.BuildText(text["keywords"]...)
if keywordsText == "" {
return "", 0, nil
}
vector, err := data.embedder.Embed(ctx, keywordsText)
if err != nil {
return "", 0, err
}
table := data.config.Vector.Pgvector.Table
if table == "" {
table = "chat_record_vectors"
}
var recordID int64
err = data.pool.QueryRowEx(
ctx,
fmt.Sprintf(
"SELECT record_id, CAST(1 - (embedding <=> $1::vector) AS real) AS score FROM %s ORDER BY embedding <=> $1::vector LIMIT 1",
table,
),
nil,
vectorLiteral(vector),
).Scan(&recordID, &score)
if err != nil {
if err == pgx.ErrNoRows {
return "", 0, nil
}
return "", 0, err
}
return strconv.FormatInt(recordID, 10), score, nil
}
func vectorLiteral(values []float32) string {
parts := make([]string, 0, len(values))
for _, value := range values {
parts = append(parts, strconv.FormatFloat(float64(value), 'f', -1, 32))
}
return "[" + strings.Join(parts, ",") + "]"
}

66
ai-chat-service/chat-server/vector-data/tencent.go
View File

@@ -1,66 +0,0 @@
package vector_data
import (
"ai-chat-service/pkg/config"
"context"
"time"
"github.com/tencent/vectordatabase-sdk-go/tcvectordb"
)
type tencentChatRecordsData struct {
config *config.Config
vectorDB *tcvectordb.Client
}
func newTencentChatRecordsData(config *config.Config) (IChatRecordsData, error) {
option := &tcvectordb.ClientOption{
Timeout: time.Second * time.Duration(config.Vector.Tencent.Timeout),
MaxIdldConnPerHost: config.Vector.Tencent.MaxIdleConnPerHost,
IdleConnTimeout: time.Second * time.Duration(config.Vector.Tencent.IdleConnTimeout),
ReadConsistency: tcvectordb.ReadConsistency(config.Vector.Tencent.ReadConsistency),
}
client, err := tcvectordb.NewClient(config.Vector.Tencent.Url, config.Vector.Tencent.Username, config.Vector.Tencent.Pwd, option)
if err != nil {
return nil, err
}
return &tencentChatRecordsData{
config: config,
vectorDB: client,
}, nil
}
func (data *tencentChatRecordsData) UpsertData(ctx context.Context, list []*ChatRecord) error {
database := data.config.Vector.Tencent.Database
collection := CHAT_RECORDS
coll := data.vectorDB.Database(database).Collection(collection)
documentList := make([]tcvectordb.Document, 0, len(list))
for _, l := range list {
doc := tcvectordb.Document{Id: l.ID}
doc.Fields = make(map[string]tcvectordb.Field, len(l.KVs))
for k, v := range l.KVs {
doc.Fields[k] = tcvectordb.Field{Val: v}
}
documentList = append(documentList, doc)
}
_, err := coll.Upsert(ctx, documentList)
return err
}
func (data *tencentChatRecordsData) QueryData(ctx context.Context, text map[string][]string) (id string, score float32, err error) {
database := data.config.Vector.Tencent.Database
collection := CHAT_RECORDS
coll := data.vectorDB.Database(database).Collection(collection)
result, err := coll.SearchByText(ctx, text, &tcvectordb.SearchDocumentParams{
Params: &tcvectordb.SearchDocParams{Ef: 100},
Limit: 1,
})
if err != nil {
return "", 0, err
}
if len(result.Documents) > 0 && len(result.Documents[0]) > 0 {
doc := result.Documents[0][0]
return doc.Id, doc.Score, nil
}
return "", 0, nil
}

52
ai-chat-service/dev.config.yaml
View File

@@ -32,19 +32,10 @@ chat:
bot_desc: "你是一个AI助手，我需要你模拟一名资深的软件工程师来回答我的问题"
# 单次请求，保留的响应tokens数量
min_response_tokens: 2048
# 上下文缓存时长，单位s
context_ttl: 1800
# 上下文消息条数
context_len: 4
redis:
host: "127.0.0.1"
port: 8888
pwd: "123456"
mysql:
dsn: "root:root@tcp(127.0.0.1:3306)/ai_chat?collation=utf8mb4_unicode_ci&charset=utf8mb4"
maxLifeTime: 3600
maxOpenConn: 10
maxIdleConn: 10
dependOn:
sensitive:
address: "localhost:50053"
@@ -54,49 +45,16 @@ dependOn:
accessToken: "ang1chubdev1ozhome256487d22sapguuv1ozhom"
tokenizer:
address: "http://127.0.0.1:3002"
vector:
# 向量后端：tencent / pgvector
provider: "pgvector"
# 历史问答命中阈值
threshold: 0.99
tencent:
url: "http://lb-4u4r1fk4-1ys6gv3rpmdan420.clb.ap-guangzhou.tencentclb.com:60000"
username: "root"
pwd: "YaUfVueWZJ20e4ghyLlBT8Dou5OapwpFTUq50oft"
database: "ai-chat"
timeout: 5
maxIdleConnPerHost: 2
readConsistency: "eventualConsistency"
idleConnTimeout: 60
pgvector:
dsn: "postgres://postgres:postgres@127.0.0.1:15432/ai_chat?sslmode=disable"
table: "chat_record_vectors"
dimensions: 1024
maxLifeTime: 3600
maxOpenConn: 10
maxIdleConn: 10
embedding:
provider: "openai-compatible"
# 智谱 OpenAI 兼容网关；可被项目根目录 .env 覆盖
base_url: "https://open.bigmodel.cn/api/paas/v4"
# 默认故意设成错误值，真实 key 请放到项目根目录 .env
api_key: "__INVALID_SET_AI_CHAT_EMBEDDING_API_KEY__"
# embedding-2 固定 1024 维，和当前 pgvector 表结构一致
model: "embedding-2"
timeout: 10
vectorDB:
# 访问地址
url: "http://lb-4u4r1fk4-1ys6gv3rpmdan420.clb.ap-guangzhou.tencentclb.com:60000"
# 用户名
username: "root"
# 密码
pwd: "YaUfVueWZJ20e4ghyLlBT8Dou5OapwpFTUq50oft"
database: "ai-chat"
# 请求超时时长s
timeout: 5
# 最大空闲连接数
maxIdleConnPerHost: 2
# 读一致性： strongConsistency(强一致性)，eventualConsistency(最终一致性)
readConsistency: "eventualConsistency"
# 空闲连接超时时长s
idleConnTimeout: 60
faiss:
base_url: "http://127.0.0.1:8451"
search_k: 1
similarity_threshold: 0.9
timeout: 10

22
ai-chat-service/docker.config.yaml
View File

@@ -16,17 +16,10 @@ chat:
frequency_penalty: 0
bot_desc: "你是一个AI助手，我需要你模拟一名资深的软件工程师来回答我的问题"
min_response_tokens: 600
context_ttl: 1800
context_len: 4
redis:
host: "host.docker.internal"
port: 8888
pwd: "123456"
mysql:
dsn: "root:root@tcp(mysql:3306)/ai_chat?collation=utf8mb4_unicode_ci&charset=utf8mb4"
maxLifeTime: 3600
maxOpenConn: 10
maxIdleConn: 10
dependOn:
sensitive:
address: "sensitive-filter:50053"
@@ -36,19 +29,14 @@ dependOn:
accessToken: "ang1chubdev1ozhome256487d22sapguuv1ozhom"
tokenizer:
address: "http://tokenizer:3002"
vector:
provider: "pgvector"
threshold: 0.99
pgvector:
dsn: "postgres://postgres:postgres@pgvector:5432/ai_chat?sslmode=disable"
table: "chat_record_vectors"
dimensions: 1024
maxLifeTime: 3600
maxOpenConn: 10
maxIdleConn: 10
embedding:
provider: "openai-compatible"
base_url: "https://open.bigmodel.cn/api/paas/v4"
api_key: "d51b903546814cc9981d3649a4a899a3.NQOtz3ocRtQwimh9"
model: "embedding-2"
timeout: 10
faiss:
base_url: "http://host.docker.internal:8451"
search_k: 1
similarity_threshold: 0.9
timeout: 10

17
ai-chat-service/go.mod
View File

@@ -3,58 +3,43 @@ module ai-chat-service
go 1.25.0
require (
github.com/go-sql-driver/mysql v1.8.1
github.com/golang/protobuf v1.5.4
github.com/google/uuid v1.6.0
github.com/jackc/pgx v3.6.2+incompatible
github.com/prometheus/client_golang v1.20.4
github.com/redis/go-redis/v9 v9.6.1
github.com/sashabaranov/go-openai v1.9.4
github.com/sirupsen/logrus v1.9.3
github.com/spf13/viper v1.19.0
github.com/tencent/vectordatabase-sdk-go v1.3.5
google.golang.org/grpc v1.65.0
google.golang.org/protobuf v1.34.2
gopkg.in/natefinch/lumberjack.v2 v2.2.1
)
require (
filippo.io/edwards25519 v1.1.0 // indirect
github.com/beorn7/perks v1.0.1 // indirect
github.com/cespare/xxhash/v2 v2.3.0 // indirect
github.com/clbanning/mxj v1.8.4 // indirect
github.com/cockroachdb/apd v1.1.0 // indirect
github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f // indirect
github.com/fsnotify/fsnotify v1.7.0 // indirect
github.com/gofrs/uuid v4.4.0+incompatible // indirect
github.com/google/go-querystring v1.0.0 // indirect
github.com/hashicorp/hcl v1.0.0 // indirect
github.com/jackc/fake v0.0.0-20150926172116-812a484cc733 // indirect
github.com/klauspost/compress v1.17.9 // indirect
github.com/lib/pq v1.12.3 // indirect
github.com/magiconair/properties v1.8.7 // indirect
github.com/mitchellh/mapstructure v1.5.0 // indirect
github.com/mozillazg/go-httpheader v0.2.1 // indirect
github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 // indirect
github.com/pelletier/go-toml/v2 v2.2.2 // indirect
github.com/pkg/errors v0.9.1 // indirect
github.com/prometheus/client_model v0.6.1 // indirect
github.com/prometheus/common v0.55.0 // indirect
github.com/prometheus/procfs v0.15.1 // indirect
github.com/sagikazarmark/locafero v0.4.0 // indirect
github.com/sagikazarmark/slog-shim v0.1.0 // indirect
github.com/shopspring/decimal v1.4.0 // indirect
github.com/sourcegraph/conc v0.3.0 // indirect
github.com/spf13/afero v1.11.0 // indirect
github.com/spf13/cast v1.6.0 // indirect
github.com/spf13/pflag v1.0.5 // indirect
github.com/stretchr/testify v1.11.1 // indirect
github.com/subosito/gotenv v1.6.0 // indirect
github.com/tencentyun/cos-go-sdk-v5 v0.7.54 // indirect
go.uber.org/atomic v1.9.0 // indirect
go.uber.org/multierr v1.9.0 // indirect
golang.org/x/crypto v0.24.0 // indirect
golang.org/x/exp v0.0.0-20231006140011-7918f672742d // indirect
golang.org/x/exp v0.0.0-20230905200255-921286631fa9 // indirect
golang.org/x/net v0.26.0 // indirect
golang.org/x/sys v0.22.0 // indirect
golang.org/x/text v0.29.0 // indirect

2648
ai-chat-service/go.sum
View File

File diff suppressed because it is too large Load Diff

92
ai-chat-service/pkg/config/config.go
View File

@@ -28,14 +28,6 @@ type Config struct {
FrequencyPenalty float32 `mapstructure:"frequency_penalty"`
BotDesc string `mapstructure:"bot_desc"`
MinResponseTokens int `mapstructure:"min_response_tokens"`
ContextTTL int `mapstructure:"context_ttl"`
ContextLen int `mapstructure:"context_len"`
}
Mysql struct {
DSN string
MaxLifeTime int
MaxOpenConn int
MaxIdleConn int
}
Redis struct {
Host string
@@ -55,28 +47,6 @@ type Config struct {
Address string
}
}
Vector struct {
Provider string
Threshold float32
Tencent struct {
Url string
Username string
Pwd string
Database string
Timeout int
MaxIdleConnPerHost int
ReadConsistency string
IdleConnTimeout int
}
Pgvector struct {
DSN string `mapstructure:"dsn"`
Table string `mapstructure:"table"`
Dimensions int `mapstructure:"dimensions"`
MaxLifeTime int `mapstructure:"maxLifeTime"`
MaxOpenConn int `mapstructure:"maxOpenConn"`
MaxIdleConn int `mapstructure:"maxIdleConn"`
}
}
Embedding struct {
Provider string
BaseUrl string `mapstructure:"base_url"`
@@ -84,15 +54,11 @@ type Config struct {
Model string `mapstructure:"model"`
Timeout int
}
VectorDB struct {
Url string
Username string
Pwd string
Database string
Timeout int
MaxIdleConnPerHost int
ReadConsistency string
IdleConnTimeout int
Faiss struct {
BaseUrl string `mapstructure:"base_url"`
SearchK int `mapstructure:"search_k"`
SimilarityThreshold float32 `mapstructure:"similarity_threshold"`
Timeout int
}
}
@@ -123,39 +89,6 @@ func GetConfig() *Config {
}
func normalizeConfig(conf *Config) {
if conf.Vector.Provider == "" {
conf.Vector.Provider = "tencent"
}
if conf.Vector.Threshold == 0 {
conf.Vector.Threshold = 0.99
}
// Backfill the new vector.tencent block from the legacy vectorDB config.
if conf.Vector.Tencent.Url == "" {
conf.Vector.Tencent.Url = conf.VectorDB.Url
}
if conf.Vector.Tencent.Username == "" {
conf.Vector.Tencent.Username = conf.VectorDB.Username
}
if conf.Vector.Tencent.Pwd == "" {
conf.Vector.Tencent.Pwd = conf.VectorDB.Pwd
}
if conf.Vector.Tencent.Database == "" {
conf.Vector.Tencent.Database = conf.VectorDB.Database
}
if conf.Vector.Tencent.Timeout == 0 {
conf.Vector.Tencent.Timeout = conf.VectorDB.Timeout
}
if conf.Vector.Tencent.MaxIdleConnPerHost == 0 {
conf.Vector.Tencent.MaxIdleConnPerHost = conf.VectorDB.MaxIdleConnPerHost
}
if conf.Vector.Tencent.ReadConsistency == "" {
conf.Vector.Tencent.ReadConsistency = conf.VectorDB.ReadConsistency
}
if conf.Vector.Tencent.IdleConnTimeout == 0 {
conf.Vector.Tencent.IdleConnTimeout = conf.VectorDB.IdleConnTimeout
}
if conf.Embedding.Provider == "" {
conf.Embedding.Provider = "openai-compatible"
}
@@ -168,6 +101,18 @@ func normalizeConfig(conf *Config) {
if conf.Embedding.Timeout == 0 {
conf.Embedding.Timeout = 10
}
if conf.Faiss.BaseUrl == "" {
conf.Faiss.BaseUrl = "http://127.0.0.1:8451"
}
if conf.Faiss.SearchK == 0 {
conf.Faiss.SearchK = 1
}
if conf.Faiss.SimilarityThreshold == 0 {
conf.Faiss.SimilarityThreshold = 0.9
}
if conf.Faiss.Timeout == 0 {
conf.Faiss.Timeout = 10
}
}
func applySecretEnvOverrides(conf *Config) {
@@ -177,6 +122,9 @@ func applySecretEnvOverrides(conf *Config) {
if v := os.Getenv("AI_CHAT_EMBEDDING_API_KEY"); v != "" {
conf.Embedding.ApiKey = v
}
if v := os.Getenv("AI_CHAT_FAISS_BASE_URL"); v != "" {
conf.Faiss.BaseUrl = v
}
if v := os.Getenv("REDIS_PASSWORD"); v != "" {
conf.Redis.Pwd = v
}

28
ai-chat-service/pkg/db/mysql/mysql.go
View File

@@ -1,28 +0,0 @@
package mysql
import (
"ai-chat-service/pkg/config"
"database/sql"
_ "github.com/go-sql-driver/mysql"
"time"
)
var db *sql.DB
func InitMysql(cnf *config.Config) {
var err error
if cnf.Mysql.DSN == "" {
panic("数据库连接字符串不能为空")
}
db, err = sql.Open("mysql", cnf.Mysql.DSN)
if err != nil {
panic(err)
}
db.SetMaxOpenConns(cnf.Mysql.MaxOpenConn)
db.SetMaxIdleConns(cnf.Mysql.MaxIdleConn)
db.SetConnMaxLifetime(time.Second * time.Duration(cnf.Mysql.MaxLifeTime))
}
func GetDB() *sql.DB {
return db
}

29
ai-chat-service/pkg/db/vector/vector.go
View File

@@ -1,29 +0,0 @@
package vector
import (
"ai-chat-service/pkg/config"
"ai-chat-service/pkg/log"
"github.com/tencent/vectordatabase-sdk-go/tcvectordb"
"time"
)
var vdb *tcvectordb.Client
func InitDB(config *config.Config) {
var defaultOption = &tcvectordb.ClientOption{
Timeout: time.Second * time.Duration(config.VectorDB.Timeout),
MaxIdldConnPerHost: config.VectorDB.MaxIdleConnPerHost,
IdleConnTimeout: time.Second * time.Duration(config.VectorDB.IdleConnTimeout),
ReadConsistency: tcvectordb.ReadConsistency(config.VectorDB.ReadConsistency),
}
var err error
vdb, err = tcvectordb.NewClient(config.VectorDB.Url, config.VectorDB.Username, config.VectorDB.Pwd, defaultOption)
if err != nil {
log.Error(err)
return
}
}
func GetVdb() *tcvectordb.Client {
return vdb
}

5
ai-chat-service/services/embedding/embedding.go
View File

@@ -19,7 +19,6 @@ type openAICompatibleEmbedder struct {
baseURL string
apiKey string
model string
dimensions int
httpClient *http.Client
}
@@ -41,7 +40,6 @@ func NewEmbedder(cnf *config.Config) (Embedder, error) {
baseURL: strings.TrimRight(cnf.Embedding.BaseUrl, "/"),
apiKey: cnf.Embedding.ApiKey,
model: cnf.Embedding.Model,
dimensions: cnf.Vector.Pgvector.Dimensions,
httpClient: &http.Client{Timeout: time.Duration(cnf.Embedding.Timeout) * time.Second},
}, nil
default:
@@ -108,8 +106,5 @@ func (e *openAICompatibleEmbedder) Embed(ctx context.Context, text string) ([]fl
if len(result.Data) == 0 || len(result.Data[0].Embedding) == 0 {
return nil, fmt.Errorf("embedding response is empty")
}
if e.dimensions > 0 && len(result.Data[0].Embedding) != e.dimensions {
return nil, fmt.Errorf("embedding dimension mismatch: got=%d want=%d", len(result.Data[0].Embedding), e.dimensions)
}
return result.Data[0].Embedding, nil
}

33
ai-chat-stack/compose.yaml
View File

@@ -1,35 +1,4 @@
services:
mysql:
image: mysql:8.0
container_name: ai-chat-mysql
restart: unless-stopped
environment:
MYSQL_ROOT_PASSWORD: root
command:
- --default-authentication-plugin=mysql_native_password
volumes:
- /data/mysql:/var/lib/mysql
- /home/lian/share/aichat/init/create_db.sql:/docker-entrypoint-initdb.d/create_db.sql:ro
healthcheck:
test: ["CMD", "mysqladmin", "ping", "-h", "127.0.0.1", "-proot"]
interval: 15s
timeout: 5s
retries: 10
pgvector:
image: pgvector/pgvector:pg16
container_name: ai-chat-pgvector
restart: unless-stopped
environment:
POSTGRES_DB: ai_chat
POSTGRES_USER: postgres
POSTGRES_PASSWORD: postgres
ports:
- "15432:5432"
volumes:
- /data/pgvector:/var/lib/postgresql/data
- /home/lian/share/aichat/init/pgvector-init.sql:/docker-entrypoint-initdb.d/pgvector-init.sql:ro
tokenizer:
build:
context: ../tokenizer
@@ -83,11 +52,9 @@ services:
ports:
- "50055:50055"
depends_on:
- mysql
- tokenizer
- sensitive-filter
- keywords-filter
- pgvector
healthcheck:
test: ["CMD", "grpc_health_probe", "-addr=:50055"]
interval: 15s

22
ai-chat-stack/configs/ai-chat-service.yaml
View File

@@ -16,17 +16,10 @@ chat:
frequency_penalty: 0
bot_desc: "你是一个AI助手，我需要你模拟一名资深的软件工程师来回答我的问题"
min_response_tokens: 600
context_ttl: 1800
context_len: 4
redis:
host: "host.docker.internal"
port: 8888
pwd: "123456"
mysql:
dsn: "root:root@tcp(mysql:3306)/ai_chat?collation=utf8mb4_unicode_ci&charset=utf8mb4"
maxLifeTime: 3600
maxOpenConn: 10
maxIdleConn: 10
dependOn:
sensitive:
address: "sensitive-filter:50053"
@@ -36,19 +29,14 @@ dependOn:
accessToken: "ang1chubdev1ozhome256487d22sapguuv1ozhom"
tokenizer:
address: "http://tokenizer:3002"
vector:
provider: "pgvector"
threshold: 0.99
pgvector:
dsn: "postgres://postgres:postgres@pgvector:5432/ai_chat?sslmode=disable"
table: "chat_record_vectors"
dimensions: 1024
maxLifeTime: 3600
maxOpenConn: 10
maxIdleConn: 10
embedding:
provider: "openai-compatible"
base_url: "https://open.bigmodel.cn/api/paas/v4"
api_key: "__SET_FROM_ENV__"
model: "embedding-2"
timeout: 10
faiss:
base_url: "http://host.docker.internal:8451"
search_k: 1
similarity_threshold: 0.9
timeout: 10

4
ai-chat-web/src/locales/en-US.ts
View File

@@ -50,8 +50,8 @@ export default {
showRawText: 'Show as raw text',
sourceSemantic: 'Semantic Match',
sourceLlm: 'LLM Output',
promptTokens: 'Prompt {count} tokens',
completionTokens: 'Completion {count} tokens',
inputTokens: 'Input {count}',
outputTokens: 'Output {count}',
sessionTokens: 'Session {count} tokens',
},
setting: {

4
ai-chat-web/src/locales/zh-CN.ts
View File

@@ -50,8 +50,8 @@ export default {
showRawText: '显示原文',
sourceSemantic: '语义匹配',
sourceLlm: '大模型输出',
promptTokens: '问题 {count} tokens',
completionTokens: '回答 {count} tokens',
inputTokens: 'Input {count}',
outputTokens: 'Output {count}',
sessionTokens: '本轮消耗 {count} tokens',
},
setting: {

4
ai-chat-web/src/locales/zh-TW.ts
View File

@@ -50,8 +50,8 @@ export default {
showRawText: '顯示原文',
sourceSemantic: '語義匹配',
sourceLlm: '大模型輸出',
promptTokens: '問題 {count} tokens',
completionTokens: '回答 {count} tokens',
inputTokens: 'Input {count}',
outputTokens: 'Output {count}',
sessionTokens: '本輪消耗 {count} tokens',
},
setting: {

57
ai-chat-web/src/views/chat/components/Message/index.vue
View File

@@ -62,15 +62,20 @@ const sourceClass = computed(() => {
}
})
const usageLabel = computed(() => {
const inputUsageLabel = computed(() => {
const usage = props.messageMeta?.usage
if (!usage)
if (!usage || props.inversion)
return ''
if (props.inversion)
return t('chat.promptTokens', { count: usage.prompt_tokens })
return t('chat.inputTokens', { count: usage.prompt_tokens })
})
return t('chat.completionTokens', { count: usage.completion_tokens })
const outputUsageLabel = computed(() => {
const usage = props.messageMeta?.usage
if (!usage || props.inversion)
return ''
return t('chat.outputTokens', { count: usage.completion_tokens })
})
const options = computed(() => {
@@ -130,21 +135,35 @@ function handleRegenerate() {
<AvatarComponent :image="inversion" />
</div>
<div class="overflow-hidden text-sm " :class="[inversion ? 'items-end' : 'items-start']">
<div class="flex flex-wrap items-center gap-2 text-xs" :class="[inversion ? 'justify-end' : 'justify-start']">
<span class="font-medium text-[#9aa4af] dark:text-neutral-500">{{ dateTime }}</span>
<span
v-if="sourceLabel"
class="rounded-full border px-2.5 py-1 text-[11px] font-medium leading-none"
:class="sourceClass"
<div class="flex flex-col gap-2" :class="[inversion ? 'items-end' : 'items-start']">
<span class="text-xs font-medium text-[#9aa4af] dark:text-neutral-500">{{ dateTime }}</span>
<div
v-if="sourceLabel || inputUsageLabel || outputUsageLabel"
class="flex flex-wrap items-center gap-2 rounded-2xl border border-[#e6edf3] bg-white/85 px-2.5 py-2 shadow-[0_10px_25px_rgba(15,23,42,0.06)] backdrop-blur-sm dark:border-neutral-800 dark:bg-[#14161a]/90"
>
{{ sourceLabel }}
</span>
<span
v-if="usageLabel"
class="rounded-full border border-violet-200 bg-violet-50 px-2.5 py-1 text-[11px] font-medium leading-none text-violet-700 dark:border-violet-900/60 dark:bg-violet-950/40 dark:text-violet-300"
>
{{ usageLabel }}
</span>
<div
v-if="sourceLabel"
class="inline-flex items-center gap-1.5 rounded-full border px-2.5 py-1 text-[11px] font-semibold leading-none"
:class="sourceClass"
>
<SvgIcon icon="ri:radar-line" />
<span>{{ sourceLabel }}</span>
</div>
<div
v-if="inputUsageLabel"
class="inline-flex items-center gap-1.5 rounded-full border border-violet-200 bg-violet-50 px-2.5 py-1 text-[11px] font-semibold leading-none text-violet-700 dark:border-violet-900/60 dark:bg-violet-950/40 dark:text-violet-300"
>
<SvgIcon icon="ri:login-circle-line" />
<span>{{ inputUsageLabel }}</span>
</div>
<div
v-if="outputUsageLabel"
class="inline-flex items-center gap-1.5 rounded-full border border-fuchsia-200 bg-fuchsia-50 px-2.5 py-1 text-[11px] font-semibold leading-none text-fuchsia-700 dark:border-fuchsia-900/60 dark:bg-fuchsia-950/40 dark:text-fuchsia-300"
>
<SvgIcon icon="ri:logout-circle-r-line" />
<span>{{ outputUsageLabel }}</span>
</div>
</div>
</div>
<div
class="flex items-end gap-1 mt-2"

50
ai-chat-web/src/views/chat/index.vue
View File

@@ -95,45 +95,7 @@ function normalizeResponseMeta(data: Chat.ConversationResponse): {
}
}
function buildPromptUsage(usage: Chat.TokenUsage): Chat.TokenUsage {
return {
prompt_tokens: usage.prompt_tokens,
completion_tokens: 0,
total_tokens: usage.prompt_tokens,
}
}
function buildCompletionUsage(usage: Chat.TokenUsage): Chat.TokenUsage {
return {
prompt_tokens: 0,
completion_tokens: usage.completion_tokens,
total_tokens: usage.completion_tokens,
}
}
function findQuestionIndex(answerIndex: number) {
for (let current = answerIndex - 1; current >= 0; current -= 1) {
if (dataSources.value[current]?.inversion)
return current
}
return -1
}
function updateMessageMeta(questionIndex: number, answerIndex: number, usage: Chat.TokenUsage, source: Chat.ReplySource | null, tokenUsed?: boolean) {
if (questionIndex >= 0) {
updateChatSome(
+uuid,
questionIndex,
{
messageMeta: {
tokenUsed,
usage: buildPromptUsage(usage),
},
},
)
}
function updateMessageMeta(answerIndex: number, usage: Chat.TokenUsage, source: Chat.ReplySource | null, tokenUsed?: boolean) {
updateChatSome(
+uuid,
answerIndex,
@@ -141,7 +103,7 @@ function updateMessageMeta(questionIndex: number, answerIndex: number, usage: Ch
messageMeta: {
source,
tokenUsed,
usage: buildCompletionUsage(usage),
usage,
},
},
)
@@ -182,8 +144,6 @@ async function onConversation() {
)
scrollToBottom()
const questionIndex = dataSources.value.length - 1
loading.value = true
prompt.value = ''
@@ -250,7 +210,7 @@ async function onConversation() {
)
if (responseMeta.source || responseMeta.usage || responseMeta.tokenUsed !== undefined)
updateMessageMeta(questionIndex, answerIndex, nextUsage, responseMeta.source, responseMeta.tokenUsed)
updateMessageMeta(answerIndex, nextUsage, responseMeta.source, responseMeta.tokenUsed)
if (responseMeta.usage && !usageApplied) {
accumulatedUsage = nextUsage
@@ -343,9 +303,7 @@ async function onRegenerate(index: number) {
options = { ...requestOptions.options }
loading.value = true
let accumulatedUsage = createEmptyUsage()
const questionIndex = findQuestionIndex(index)
updateChat(
+uuid,
@@ -401,7 +359,7 @@ async function onRegenerate(index: number) {
)
if (responseMeta.source || responseMeta.usage || responseMeta.tokenUsed !== undefined)
updateMessageMeta(questionIndex, index, nextUsage, responseMeta.source, responseMeta.tokenUsed)
updateMessageMeta(index, nextUsage, responseMeta.source, responseMeta.tokenUsed)
if (responseMeta.usage && !usageApplied) {
accumulatedUsage = nextUsage