Files
gochat/internal/service/rag_service.go
T
2026-06-04 15:44:48 +08:00

281 lines
9.4 KiB
Go

package service
import (
"context"
"fmt"
"strings"
"github.com/gochat/gochat/internal/llm"
"github.com/gochat/gochat/internal/model"
"github.com/pgvector/pgvector-go"
applogger "github.com/gochat/gochat/pkg/logger"
)
// RAGService implements Retrieval-Augmented Generation for Captain knowledge base Q&A.
// Reference: Chatwoot Captain::Llm::GraphService + Captain::Tools::CaptainKnowledgeBaseTool
//
// RAG flow:
// 1. Take user question → generate embedding via LLM provider
// 2. Search assistant response embeddings in pgvector for top-k relevant FAQs
// 3. Build prompt with retrieved context + question → call LLM for answer
// 4. Return answer with source document references
// RAGQueryRequest holds the input for a RAG Q&A query.
type RAGQueryRequest struct {
AssistantID uint `json:"assistant_id" validate:"required"`
Question string `json:"question" validate:"required,min=1"`
TopK int `json:"top_k,omitempty"` // number of FAQ responses to retrieve, default 5
}
// RAGQueryResult holds the answer from RAG Q&A.
type RAGQueryResult struct {
Answer string `json:"answer"`
Sources []RAGSource `json:"sources,omitempty"`
Confidence float64 `json:"confidence"`
AssistantID uint `json:"assistant_id"`
Conversation uint `json:"conversation_id,omitempty"`
}
// RAGSource references a source FAQ response used in the answer.
type RAGSource struct {
ResponseID uint `json:"response_id"`
Question string `json:"question"`
Answer string `json:"answer"`
DocumentID uint `json:"document_id,omitempty"`
Score float64 `json:"score"`
}
// AssistantRepoIface defines the repository interface RAGService needs from CaptainAssistantRepo.
type AssistantRepoIface interface {
GetByID(ctx context.Context, id uint) (*model.CaptainAssistant, error)
}
// ResponseRepoIface defines the repository interface RAGService needs from CaptainAssistantResponseRepo.
type ResponseRepoIface interface {
SearchByEmbedding(ctx context.Context, assistantID uint, embedding pgvector.Vector, limit int) ([]model.CaptainAssistantResponse, error)
GetByID(ctx context.Context, id uint) (*model.CaptainAssistantResponse, error)
Update(ctx context.Context, resp *model.CaptainAssistantResponse) error
}
// RAGService orchestrates embedding search + LLM generation for knowledge base Q&A.
type RAGService struct {
responseRepo ResponseRepoIface
assistantRepo AssistantRepoIface
llmProvider llm.Provider
promptBuilder *SystemPromptBuilder
}
// NewRAGService creates a new RAGService.
func NewRAGService(
responseRepo ResponseRepoIface,
assistantRepo AssistantRepoIface,
llmProvider llm.Provider,
) *RAGService {
return &RAGService{
responseRepo: responseRepo,
assistantRepo: assistantRepo,
llmProvider: llmProvider,
promptBuilder: NewSystemPromptBuilder(),
}
}
// Query performs a RAG Q&A: embed question → search FAQs → generate answer.
func (s *RAGService) Query(ctx context.Context, accountID uint, req *RAGQueryRequest) (*RAGQueryResult, error) {
if req.TopK <= 0 {
req.TopK = 5
}
// Step 1: Verify assistant exists and is active
assistant, err := s.assistantRepo.GetByID(ctx, req.AssistantID)
if err != nil {
return nil, fmt.Errorf("assistant not found: %w", err)
}
if assistant.Status != model.AssistantStatusActive {
return nil, fmt.Errorf("assistant is not active (status: %s)", assistant.Status)
}
// Get assistant config for prompt building and model selection
cfg, err := assistant.GetConfig()
if err != nil {
applogger.L().Warnf("RAG get assistant config: %v", err)
cfg = &model.AssistantConfig{}
}
// Step 2: Generate embedding for the question
embedResp, err := s.llmProvider.CreateEmbedding(ctx, llm.EmbeddingRequest{
Model: "text-embedding-3-small",
Input: []string{req.Question},
})
if err != nil {
applogger.L().Errorf("RAG embedding generation: %v", err)
return nil, fmt.Errorf("embedding generation failed: %w", err)
}
if len(embedResp.Data) == 0 {
return nil, fmt.Errorf("no embedding returned for question")
}
questionEmbedding := embedResp.Data[0].Embedding
// Convert float64 embedding to float32 for pgvector
float32Emb := make([]float32, len(questionEmbedding))
for i, v := range questionEmbedding {
float32Emb[i] = float32(v)
}
pgvectorEmb := pgvector.NewVector(float32Emb)
// Step 3: Search assistant responses by embedding similarity (pgvector cosine distance)
results, err := s.responseRepo.SearchByEmbedding(ctx, req.AssistantID, pgvectorEmb, req.TopK)
if err != nil {
applogger.L().Errorf("RAG response search: %v", err)
return nil, fmt.Errorf("FAQ search failed: %w", err)
}
if len(results) == 0 {
// No relevant FAQs found — fallback to LLM without context
return s.queryWithoutContext(ctx, assistant, cfg, req.Question)
}
// Step 4: Build RAG context from retrieved FAQ responses
var contextParts []string
var sources []RAGSource
for i, r := range results {
contextParts = append(contextParts, fmt.Sprintf("[FAQ %d]\nQ: %s\nA: %s", i+1, r.Question, r.Answer))
// Estimate score (pgvector cosine distance is already sorted by closest)
score := 1.0 - float64(i)*0.1 // rough heuristic since we don't have raw distance
if i == 0 {
score = 0.9 // top result is likely high-confidence
}
sources = append(sources, RAGSource{
ResponseID: r.ID,
Question: r.Question,
Answer: r.Answer,
Score: score,
})
if r.DocumentableID != nil {
sources[len(sources)-1].DocumentID = *r.DocumentableID
}
}
ragContext := strings.Join(contextParts, "\n\n")
// Step 5: Build prompt with RAG context + system prompt
systemPrompt := s.promptBuilder.BuildAssistantPrompt(assistant, cfg)
systemPrompt += fmt.Sprintf("\n\nKnowledge Base Context:\n%s", ragContext)
systemPrompt += "\n\nImportant: When answering, reference the FAQ entries by number. If the FAQ context doesn't contain enough information to answer, say so honestly."
messages := []llm.ChatMessage{
{Role: "system", Content: systemPrompt},
{Role: "user", Content: req.Question},
}
// Step 6: Call LLM for answer generation
modelName := cfg.Model
if modelName == "" {
modelName = "gpt-4"
}
temperature := cfg.Temperature
if temperature == 0 {
temperature = 0.3 // lower temp for factual answers
}
llmResp, err := s.llmProvider.ChatCompletion(ctx, llm.ChatRequest{
Model: modelName,
Messages: messages,
Temperature: temperature,
MaxTokens: 1024,
})
if err != nil {
applogger.L().Errorf("RAG LLM generation: %v", err)
return nil, fmt.Errorf("answer generation failed: %w", err)
}
answer := ""
confidence := 0.0
if len(llmResp.Choices) > 0 {
answer = llmResp.Choices[0].Message.Content
if len(sources) > 0 {
confidence = sources[0].Score
}
}
return &RAGQueryResult{
Answer: answer,
Sources: sources,
Confidence: confidence,
AssistantID: req.AssistantID,
}, nil
}
// queryWithoutContext handles the case when no relevant FAQ responses are found.
// Falls back to generating an answer purely from the assistant's configuration.
func (s *RAGService) queryWithoutContext(ctx context.Context, assistant *model.CaptainAssistant, cfg *model.AssistantConfig, question string) (*RAGQueryResult, error) {
systemPrompt := s.promptBuilder.BuildAssistantPrompt(assistant, cfg)
systemPrompt += "\n\nNote: No relevant FAQ entries were found for this question. Answer based on your general knowledge, but indicate that the answer may not be specific to the product."
modelName := cfg.Model
if modelName == "" {
modelName = "gpt-4"
}
temperature := cfg.Temperature
if temperature == 0 {
temperature = 0.5
}
llmResp, err := s.llmProvider.ChatCompletion(ctx, llm.ChatRequest{
Model: modelName,
Messages: []llm.ChatMessage{
{Role: "system", Content: systemPrompt},
{Role: "user", Content: question},
},
Temperature: temperature,
MaxTokens: 512,
})
if err != nil {
return nil, fmt.Errorf("fallback LLM generation failed: %w", err)
}
answer := ""
if len(llmResp.Choices) > 0 {
answer = llmResp.Choices[0].Message.Content
}
return &RAGQueryResult{
Answer: answer,
Confidence: 0.3, // low confidence when no context
AssistantID: assistant.ID,
}, nil
}
// IndexResponse generates embeddings for an assistant response and stores them for search.
// Called when a response is created or its question/answer is updated.
func (s *RAGService) IndexResponse(ctx context.Context, responseID uint) error {
resp, err := s.responseRepo.GetByID(ctx, responseID)
if err != nil {
return fmt.Errorf("response not found: %w", err)
}
// Generate embedding from question + answer for better semantic matching
inputText := fmt.Sprintf("Q: %s\nA: %s", resp.Question, resp.Answer)
embedResp, err := s.llmProvider.CreateEmbedding(ctx, llm.EmbeddingRequest{
Model: "text-embedding-3-small",
Input: []string{inputText},
})
if err != nil {
applogger.L().Errorf("RAG response embedding: %v", err)
return fmt.Errorf("response embedding failed: %w", err)
}
if len(embedResp.Data) > 0 {
float32Emb := make([]float32, len(embedResp.Data[0].Embedding))
for i, v := range embedResp.Data[0].Embedding {
float32Emb[i] = float32(v)
}
resp.Embedding = pgvector.NewVector(float32Emb)
if err := s.responseRepo.Update(ctx, resp); err != nil {
applogger.L().Errorf("RAG update response embedding: %v", err)
return fmt.Errorf("store embedding failed: %w", err)
}
}
applogger.L().Infof("RAG indexed response %d for assistant %d", resp.ID, resp.AssistantID)
return nil
}