vetrag/chat_service.go

123 lines
3.1 KiB
Go

package main
import (
"context"
"net/http"
"github.com/gin-gonic/gin"
"github.com/sirupsen/logrus"
)
type ChatService struct {
LLM *LLMClient
}
func NewChatService(llm *LLMClient) *ChatService {
return &ChatService{LLM: llm}
}
func (cs *ChatService) HandleChat(c *gin.Context) {
ctx := context.Background()
req, err := cs.parseRequest(c)
if err != nil {
return
}
keywords, err := cs.extractKeywords(ctx, req.Message)
if err != nil {
cs.logChat(req, keywords, nil, "", err)
c.JSON(http.StatusOK, ChatResponse{Match: nil})
return
}
kwArr := cs.keywordsToStrings(keywords["keyword"])
best, bestID, candidates, err := cs.findBestCandidate(ctx, req, kwArr)
cs.logChat(req, keywords, candidates, bestID, err)
resp := cs.buildResponse(best)
c.JSON(http.StatusOK, resp)
}
func (cs *ChatService) parseRequest(c *gin.Context) (ChatRequest, error) {
var req ChatRequest
if err := c.ShouldBindJSON(&req); err != nil {
logrus.WithError(err).Error("Invalid request")
c.JSON(http.StatusBadRequest, gin.H{"error": "Invalid request"})
return req, err
}
return req, nil
}
func (cs *ChatService) extractKeywords(ctx context.Context, message string) (map[string]interface{}, error) {
keywords, err := cs.LLM.ExtractKeywords(ctx, message)
return keywords, err
}
func (cs *ChatService) keywordsToStrings(kwIface interface{}) []string {
var kwArr []string
switch v := kwIface.(type) {
case []interface{}:
for _, item := range v {
if s, ok := item.(string); ok {
kwArr = append(kwArr, s)
}
}
case []string:
kwArr = v
}
return kwArr
}
func (cs *ChatService) findBestCandidate(ctx context.Context, req ChatRequest, kwArr []string) (*Reason, string, []Reason, error) {
candidates := findCandidates(kwArr)
bestID := ""
var err error
if len(candidates) > 0 {
bestID, err = cs.LLM.DisambiguateBestMatch(ctx, req.Message, candidates)
}
var best *Reason
for i := range candidates {
if candidates[i].ID == bestID {
best = &candidates[i]
break
}
}
if err != nil || len(kwArr) == 0 || len(candidates) == 0 || bestID == "" || best == nil {
return nil, bestID, candidates, err
}
return best, bestID, candidates, nil
}
func (cs *ChatService) buildResponse(best *Reason) ChatResponse {
if best == nil {
return ChatResponse{Match: nil}
}
totalPrice, totalDuration := sumProcedures(best.Procedures)
return ChatResponse{
Match: &best.ID,
Procedures: best.Procedures,
TotalPrice: totalPrice,
TotalDuration: totalDuration,
Notes: best.Notes,
}
}
func (cs *ChatService) logChat(req ChatRequest, keywords map[string]interface{}, candidates []Reason, bestID string, err error) {
logRequest(req, keywords, candidates, bestID, err)
if candidates != nil && bestID != "" {
var best *Reason
for i := range candidates {
if candidates[i].ID == bestID {
best = &candidates[i]
break
}
}
if best != nil {
totalPrice, totalDuration := sumProcedures(best.Procedures)
logrus.WithFields(logrus.Fields{
"match": best.ID,
"total_price": totalPrice,
"total_duration": totalDuration,
"notes": best.Notes,
}).Info("Responding with match")
}
}
}