362 lines
9.1 KiB
Go
362 lines
9.1 KiB
Go
package knowledge
|
||
|
||
import (
|
||
"context"
|
||
"encoding/json"
|
||
"fmt"
|
||
|
||
"github.com/google/uuid"
|
||
"resume-platform/internal/model"
|
||
"resume-platform/internal/repository"
|
||
"resume-platform/pkg/logger"
|
||
"strings"
|
||
)
|
||
|
||
type KnowledgeBaseService struct {
|
||
embeddingService *EmbeddingService
|
||
vectorService *VectorService
|
||
repo repository.ResumeRepository
|
||
}
|
||
|
||
func NewKnowledgeBaseService(embeddingService *EmbeddingService, vectorService *VectorService, repo repository.ResumeRepository) *KnowledgeBaseService {
|
||
return &KnowledgeBaseService{
|
||
embeddingService: embeddingService,
|
||
vectorService: vectorService,
|
||
repo: repo,
|
||
}
|
||
}
|
||
|
||
func (s *KnowledgeBaseService) AddDocument(ctx context.Context, documentID, content, userID string) error {
|
||
if content == "" {
|
||
return fmt.Errorf("empty content")
|
||
}
|
||
|
||
chunks := s.chunkContent(content, 500, 50)
|
||
|
||
doc, _ := s.repo.GetDocumentByID(documentID)
|
||
sourceName := documentID[:8]
|
||
if doc != nil && doc.Name != "" {
|
||
sourceName = doc.Name
|
||
}
|
||
|
||
for i, chunk := range chunks {
|
||
documentChunk := &model.DocumentChunk{
|
||
ID: uuid.New().String(),
|
||
DocumentID: documentID,
|
||
ChunkIndex: i,
|
||
Content: chunk,
|
||
Metadata: "",
|
||
}
|
||
|
||
err := s.vectorService.StoreChunk(ctx, documentChunk, userID, sourceName)
|
||
if err != nil {
|
||
return err
|
||
}
|
||
}
|
||
|
||
logger.Infof("Added %d chunks to document %s", len(chunks), documentID)
|
||
return nil
|
||
}
|
||
|
||
func (s *KnowledgeBaseService) AddResume(ctx context.Context, resumeID string, resume *model.Resume) error {
|
||
if resume == nil {
|
||
return fmt.Errorf("nil resume")
|
||
}
|
||
|
||
resumeText := s.serializeResume(resume)
|
||
if resumeText == "" {
|
||
return fmt.Errorf("empty resume content")
|
||
}
|
||
|
||
chunks := s.chunkContent(resumeText, 500, 50)
|
||
sourceName := resume.BasicInfo.Name
|
||
if sourceName == "" {
|
||
sourceName = resumeID[:8]
|
||
}
|
||
|
||
for i, chunk := range chunks {
|
||
documentChunk := &model.DocumentChunk{
|
||
ID: uuid.New().String(),
|
||
DocumentID: resumeID,
|
||
ChunkIndex: i,
|
||
Content: chunk,
|
||
Metadata: "",
|
||
}
|
||
|
||
err := s.vectorService.StoreChunk(ctx, documentChunk, resume.UserID, sourceName)
|
||
if err != nil {
|
||
return err
|
||
}
|
||
}
|
||
|
||
err := s.vectorService.Store(ctx, resumeText, "resume", resumeID, sourceName, resume.UserID)
|
||
if err != nil {
|
||
return err
|
||
}
|
||
|
||
logger.Infof("Added resume %s to knowledge base with %d chunks", resumeID, len(chunks))
|
||
return nil
|
||
}
|
||
|
||
func (s *KnowledgeBaseService) AddText(ctx context.Context, text, sourceType, sourceID, sourceName, userID string) error {
|
||
if text == "" {
|
||
return fmt.Errorf("empty text")
|
||
}
|
||
|
||
return s.vectorService.Store(ctx, text, sourceType, sourceID, sourceName, userID)
|
||
}
|
||
|
||
func (s *KnowledgeBaseService) Retrieve(ctx context.Context, query string, topK int, userID string) ([]model.ChunkWithScore, error) {
|
||
return s.vectorService.SearchWithChunks(ctx, query, topK, userID)
|
||
}
|
||
|
||
func (s *KnowledgeBaseService) RetrieveAll(ctx context.Context, query string, topK int) ([]model.ChunkWithScore, error) {
|
||
return s.vectorService.Search(ctx, query, topK, "")
|
||
}
|
||
|
||
func (s *KnowledgeBaseService) DeleteDocument(ctx context.Context, documentID string) error {
|
||
err := s.repo.DeleteEmbeddingsBySource("document", documentID)
|
||
if err != nil {
|
||
return err
|
||
}
|
||
return s.vectorService.DeleteDocumentChunks(ctx, documentID)
|
||
}
|
||
|
||
func (s *KnowledgeBaseService) DeleteResume(ctx context.Context, resumeID string) error {
|
||
return s.repo.DeleteEmbeddingsBySource("resume", resumeID)
|
||
}
|
||
|
||
func (s *KnowledgeBaseService) GetStats(ctx context.Context, userID string) (map[string]int, error) {
|
||
var stats = map[string]int{}
|
||
|
||
if userID != "" {
|
||
embeddings, err := s.repo.GetEmbeddingsByUserID(userID)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
stats["embeddings"] = len(embeddings)
|
||
|
||
chunks, err := s.repo.GetDocumentChunksByUserID(userID)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
stats["chunks"] = len(chunks)
|
||
} else {
|
||
embeddings, err := s.repo.GetAllEmbeddings()
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
stats["embeddings"] = len(embeddings)
|
||
|
||
chunks, err := s.repo.GetAllDocumentChunks()
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
stats["chunks"] = len(chunks)
|
||
}
|
||
|
||
return stats, nil
|
||
}
|
||
|
||
func (s *KnowledgeBaseService) chunkContent(content string, chunkSize, overlap int) []string {
|
||
if content == "" {
|
||
return nil
|
||
}
|
||
|
||
content = strings.ReplaceAll(content, "\r\n", "\n")
|
||
content = strings.ReplaceAll(content, "\r", "\n")
|
||
|
||
var chunks []string
|
||
start := 0
|
||
contentLen := len(content)
|
||
|
||
for start < contentLen {
|
||
end := start + chunkSize
|
||
if end > contentLen {
|
||
end = contentLen
|
||
}
|
||
|
||
if end < contentLen {
|
||
for i := end; i > start && i > start+chunkSize-overlap; i-- {
|
||
c := rune(content[i])
|
||
if c == '\n' || c == ';' || c == '\r' {
|
||
end = i + 1
|
||
break
|
||
}
|
||
}
|
||
}
|
||
|
||
chunk := strings.TrimSpace(content[start:end])
|
||
if chunk != "" {
|
||
chunks = append(chunks, chunk)
|
||
}
|
||
|
||
start = end - overlap
|
||
if start < 0 {
|
||
start = 0
|
||
}
|
||
if start >= contentLen {
|
||
break
|
||
}
|
||
}
|
||
|
||
return chunks
|
||
}
|
||
|
||
func (s *KnowledgeBaseService) serializeResume(resume *model.Resume) string {
|
||
var builder strings.Builder
|
||
|
||
if resume.BasicInfo.Name != "" {
|
||
builder.WriteString("姓名：")
|
||
builder.WriteString(resume.BasicInfo.Name)
|
||
builder.WriteString("\n")
|
||
}
|
||
if resume.BasicInfo.Title != "" {
|
||
builder.WriteString("职位：")
|
||
builder.WriteString(resume.BasicInfo.Title)
|
||
builder.WriteString("\n")
|
||
}
|
||
if resume.BasicInfo.Email != "" {
|
||
builder.WriteString("邮箱：")
|
||
builder.WriteString(resume.BasicInfo.Email)
|
||
builder.WriteString("\n")
|
||
}
|
||
if resume.BasicInfo.Phone != "" {
|
||
builder.WriteString("电话：")
|
||
builder.WriteString(resume.BasicInfo.Phone)
|
||
builder.WriteString("\n")
|
||
}
|
||
if resume.BasicInfo.Location != "" {
|
||
builder.WriteString("所在地：")
|
||
builder.WriteString(resume.BasicInfo.Location)
|
||
builder.WriteString("\n")
|
||
}
|
||
if resume.BasicInfo.Summary != "" {
|
||
builder.WriteString("个人简介：")
|
||
builder.WriteString(resume.BasicInfo.Summary)
|
||
builder.WriteString("\n")
|
||
}
|
||
if resume.BasicInfo.JobTarget != "" {
|
||
builder.WriteString("求职目标：")
|
||
builder.WriteString(resume.BasicInfo.JobTarget)
|
||
builder.WriteString("\n")
|
||
}
|
||
|
||
if len(resume.Experience) > 0 {
|
||
builder.WriteString("\n【工作经历】\n")
|
||
for _, exp := range resume.Experience {
|
||
builder.WriteString("公司：")
|
||
builder.WriteString(exp.Company)
|
||
builder.WriteString("\n职位：")
|
||
builder.WriteString(exp.Position)
|
||
builder.WriteString("\n时间：")
|
||
builder.WriteString(exp.StartDate)
|
||
if exp.EndDate != "" {
|
||
builder.WriteString(" - ")
|
||
builder.WriteString(exp.EndDate)
|
||
}
|
||
builder.WriteString("\n职责：")
|
||
builder.WriteString(exp.Description)
|
||
builder.WriteString("\n")
|
||
if len(exp.Highlights) > 0 {
|
||
builder.WriteString("亮点：")
|
||
builder.WriteString(strings.Join(exp.Highlights, "；"))
|
||
builder.WriteString("\n")
|
||
}
|
||
builder.WriteString("\n")
|
||
}
|
||
}
|
||
|
||
if len(resume.Education) > 0 {
|
||
builder.WriteString("\n【教育背景】\n")
|
||
for _, edu := range resume.Education {
|
||
builder.WriteString("学校：")
|
||
builder.WriteString(edu.School)
|
||
builder.WriteString("\n学位：")
|
||
builder.WriteString(edu.Degree)
|
||
builder.WriteString("\n专业：")
|
||
builder.WriteString(edu.Major)
|
||
builder.WriteString("\n时间：")
|
||
builder.WriteString(edu.StartDate)
|
||
if edu.EndDate != "" {
|
||
builder.WriteString(" - ")
|
||
builder.WriteString(edu.EndDate)
|
||
}
|
||
builder.WriteString("\n")
|
||
}
|
||
}
|
||
|
||
if len(resume.Skills) > 0 {
|
||
builder.WriteString("\n【专业技能】\n")
|
||
for _, skill := range resume.Skills {
|
||
builder.WriteString(skill.Name)
|
||
if skill.Level != "" {
|
||
builder.WriteString("（")
|
||
builder.WriteString(skill.Level)
|
||
builder.WriteString("）")
|
||
}
|
||
if skill.Category != "" {
|
||
builder.WriteString(" - ")
|
||
builder.WriteString(skill.Category)
|
||
}
|
||
builder.WriteString("\n")
|
||
}
|
||
}
|
||
|
||
if len(resume.Projects) > 0 {
|
||
builder.WriteString("\n【项目经验】\n")
|
||
for _, proj := range resume.Projects {
|
||
builder.WriteString("项目名称：")
|
||
builder.WriteString(proj.Name)
|
||
builder.WriteString("\n描述：")
|
||
builder.WriteString(proj.Description)
|
||
builder.WriteString("\n")
|
||
if len(proj.TechStack) > 0 {
|
||
builder.WriteString("技术栈：")
|
||
builder.WriteString(strings.Join(proj.TechStack, "、"))
|
||
builder.WriteString("\n")
|
||
}
|
||
if len(proj.Highlights) > 0 {
|
||
builder.WriteString("亮点：")
|
||
builder.WriteString(strings.Join(proj.Highlights, "；"))
|
||
builder.WriteString("\n")
|
||
}
|
||
if len(proj.Achievements) > 0 {
|
||
builder.WriteString("成果：")
|
||
builder.WriteString(strings.Join(proj.Achievements, "；"))
|
||
builder.WriteString("\n")
|
||
}
|
||
builder.WriteString("\n")
|
||
}
|
||
}
|
||
|
||
result := builder.String()
|
||
if len(result) > 15000 {
|
||
result = result[:15000]
|
||
}
|
||
return result
|
||
}
|
||
|
||
func (s *KnowledgeBaseService) BuildIndex(ctx context.Context) error {
|
||
documents, err := s.repo.GetAllDocumentChunks()
|
||
if err != nil {
|
||
return err
|
||
}
|
||
|
||
for _, chunk := range documents {
|
||
if chunk.Embedding == "" && chunk.Content != "" {
|
||
embedding, err := s.embeddingService.Generate(ctx, chunk.Content)
|
||
if err != nil {
|
||
return err
|
||
}
|
||
embeddingJSON, _ := json.Marshal(embedding)
|
||
chunk.Embedding = string(embeddingJSON)
|
||
err = s.repo.CreateDocumentChunk(chunk)
|
||
if err != nil {
|
||
return err
|
||
}
|
||
}
|
||
}
|
||
|
||
return nil
|
||
} |