From 9c8d5a9e8d11a521d93d370ee907e45d262586a9 Mon Sep 17 00:00:00 2001 From: Aska Date: Fri, 21 Aug 2026 18:40:01 +0800 Subject: [PATCH] feat(rag): pgvector knowledge base with API and LLM tools MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - services/rag: knowledge_docs store — vector search via the shared embedder (HNSW index, graceful ILIKE fallback), list/delete scoped to the user, search intentionally global for the single-user deployment - gateway: GET/POST /knowledge/docs, DELETE /knowledge/docs/:id, POST /knowledge/search - plugins/knowledge: knowledge_search and knowledge_ingest tools so the persona can save and recall information mid-conversation - main: migrations, store wiring, plugin registration --- backend/cmd/cyrene/main.go | 19 ++ backend/plugins/knowledge/plugin.go | 93 +++++++++ backend/plugins/knowledge/plugin_test.go | 68 +++++++ backend/services/gateway/api/handlers_more.go | 58 +++++- backend/services/gateway/api/router.go | 15 ++ backend/services/rag/rag.go | 181 ++++++++++++++++++ 6 files changed, 432 insertions(+), 2 deletions(-) create mode 100644 backend/plugins/knowledge/plugin.go create mode 100644 backend/plugins/knowledge/plugin_test.go create mode 100644 backend/services/rag/rag.go diff --git a/backend/cmd/cyrene/main.go b/backend/cmd/cyrene/main.go index 6004645..669a798 100644 --- a/backend/cmd/cyrene/main.go +++ b/backend/cmd/cyrene/main.go @@ -32,6 +32,7 @@ import ( "git.yeij.top/AskaEth/Cyrene_v2/plugins/text" webfetch "git.yeij.top/AskaEth/Cyrene_v2/plugins/web-fetch" websearch "git.yeij.top/AskaEth/Cyrene_v2/plugins/web-search" + "git.yeij.top/AskaEth/Cyrene_v2/plugins/knowledge" "git.yeij.top/AskaEth/Cyrene_v2/services/ai-core/background" "git.yeij.top/AskaEth/Cyrene_v2/services/ai-core/emotion" @@ -49,6 +50,7 @@ import ( "git.yeij.top/AskaEth/Cyrene_v2/services/gateway/store" "git.yeij.top/AskaEth/Cyrene_v2/services/gateway/ws" "git.yeij.top/AskaEth/Cyrene_v2/services/memory" + "git.yeij.top/AskaEth/Cyrene_v2/services/rag" bridge "git.yeij.top/AskaEth/Cyrene_v2/services/platform-bridge" bridgeqq "git.yeij.top/AskaEth/Cyrene_v2/services/platform-bridge/qq" bridgetg "git.yeij.top/AskaEth/Cyrene_v2/services/platform-bridge/telegram" @@ -146,6 +148,7 @@ func main() { // === 6. Memory Service === var memSvc *memory.Service + var ragStore *rag.Store if database != nil { memRepo := memory.NewRepository(database) memSvc = memory.NewService(memRepo) @@ -160,6 +163,19 @@ func main() { logger.Infof("[memory] hybrid search enabled (vector dims=%d, api=%v)", emb.Dims(), emb.IsAvailable()) } logger.Infof("[memory] connected") + + // === 6.5 Knowledge base (RAG) === + db.RunMigration(context.Background(), database, "knowledge", db.SplitSQL(rag.CreateTableSQL())) + ragStore = rag.NewStore(database, emb) + if err := db.RunMigration(context.Background(), database, "knowledge-vector", db.SplitSQL(rag.VectorMigrationSQL(emb.Dims()))); err != nil { + logger.Infof("[rag] pgvector unavailable, ILIKE-only search: %v", err) + } else { + ragStore.EnableVector() + } + if err := reg.Register(knowledge.New(ragStore), pluginCfg); err != nil { + logger.Infof("[plugin] knowledge: %v", err) + } + logger.Infof("[rag] knowledge base ready") } else { logger.Infof("[memory] disabled (no database)") } @@ -316,6 +332,9 @@ func main() { } gwDeps.VoiceSvc = voiceSvc gwDeps.Vision = visionAdapter{llmSelector.DefaultAdapter()} + if ragStore != nil { + gwDeps.Knowledge = ragStore + } wsHub := ws.NewHub(orch) gwDeps.WSHub = wsHub diff --git a/backend/plugins/knowledge/plugin.go b/backend/plugins/knowledge/plugin.go new file mode 100644 index 0000000..11f1907 --- /dev/null +++ b/backend/plugins/knowledge/plugin.go @@ -0,0 +1,93 @@ +// Package knowledge exposes the knowledge base as LLM tools. +package knowledge + +import ( + "context" + "fmt" + "strings" + + "git.yeij.top/AskaEth/Cyrene_v2/internal/plugin" + "git.yeij.top/AskaEth/Cyrene_v2/services/rag" +) + +// Store is the subset of rag.Store used by the tools. +type Store interface { + Ingest(ctx context.Context, userID, title, content string) (*rag.Document, error) + Search(ctx context.Context, userID, query string, topK int) ([]rag.Document, error) +} + +type Plugin struct{ store Store } + +func New(store Store) *Plugin { return &Plugin{store: store} } + +func (p *Plugin) Name() string { return "knowledge" } +func (p *Plugin) Version() string { return "1.0.0" } +func (p *Plugin) Description() string { return "知识库检索与写入" } +func (p *Plugin) Init(_ context.Context, _ plugin.PluginConfig) error { return nil } +func (p *Plugin) Close() error { return nil } +func (p *Plugin) Enabled() bool { return p.store != nil } + +func (p *Plugin) Tools() []plugin.ToolDefinition { + return []plugin.ToolDefinition{ + { + Name: "knowledge_search", + Description: "Search the knowledge base for relevant documents.", + Parameters: map[string]any{ + "type": "object", + "properties": map[string]any{ + "query": map[string]string{"type": "string", "description": "Search query"}, + }, + "required": []string{"query"}, + }, + }, + { + Name: "knowledge_ingest", + Description: "Save a piece of information into the knowledge base for later retrieval.", + Parameters: map[string]any{ + "type": "object", + "properties": map[string]any{ + "title": map[string]string{"type": "string", "description": "Document title"}, + "content": map[string]string{"type": "string", "description": "Document content"}, + }, + "required": []string{"title", "content"}, + }, + }, + } +} + +func (p *Plugin) Execute(ctx context.Context, toolName string, params map[string]any) (string, error) { + if p.store == nil { + return "", fmt.Errorf("knowledge store unavailable") + } + switch toolName { + case "knowledge_search": + query, _ := params["query"].(string) + if query == "" { + return "", fmt.Errorf("missing query") + } + docs, err := p.store.Search(ctx, "", query, 5) + if err != nil { + return "", err + } + if len(docs) == 0 { + return "知识库中没有相关内容", nil + } + var sb strings.Builder + for i, d := range docs { + fmt.Fprintf(&sb, "%d. 《%s》\n%s\n\n", i+1, d.Title, d.Content) + } + return sb.String(), nil + + case "knowledge_ingest": + title, _ := params["title"].(string) + content, _ := params["content"].(string) + if title == "" || content == "" { + return "", fmt.Errorf("missing title or content") + } + if _, err := p.store.Ingest(ctx, "", title, content); err != nil { + return "", err + } + return fmt.Sprintf("已收录《%s》", title), nil + } + return "", fmt.Errorf("unknown tool %q", toolName) +} diff --git a/backend/plugins/knowledge/plugin_test.go b/backend/plugins/knowledge/plugin_test.go new file mode 100644 index 0000000..e1ee6fb --- /dev/null +++ b/backend/plugins/knowledge/plugin_test.go @@ -0,0 +1,68 @@ +package knowledge + +import ( + "context" + "strings" + "testing" + + "git.yeij.top/AskaEth/Cyrene_v2/services/rag" +) + +type fakeStore struct { + docs []rag.Document +} + +func (f *fakeStore) Ingest(_ context.Context, userID, title, content string) (*rag.Document, error) { + d := &rag.Document{ID: int64(len(f.docs) + 1), UserID: userID, Title: title, Content: content} + f.docs = append(f.docs, *d) + return d, nil +} + +func (f *fakeStore) Search(_ context.Context, _, query string, topK int) ([]rag.Document, error) { + var out []rag.Document + for _, d := range f.docs { + if strings.Contains(d.Content, query) || strings.Contains(d.Title, query) { + out = append(out, d) + } + } + return out, nil +} + +func TestIngestAndSearch(t *testing.T) { + p := New(&fakeStore{}) + ctx := context.Background() + + out, err := p.Execute(ctx, "knowledge_ingest", map[string]any{"title": "火锅攻略", "content": "楼下那家火锅店周三半价"}) + if err != nil { + t.Fatalf("ingest: %v", err) + } + if !strings.Contains(out, "火锅攻略") { + t.Errorf("ingest output: %q", out) + } + + out, err = p.Execute(ctx, "knowledge_search", map[string]any{"query": "火锅"}) + if err != nil { + t.Fatalf("search: %v", err) + } + if !strings.Contains(out, "周三半价") { + t.Errorf("search output: %q", out) + } +} + +func TestSearchEmpty(t *testing.T) { + p := New(&fakeStore{}) + out, err := p.Execute(context.Background(), "knowledge_search", map[string]any{"query": "不存在"}) + if err != nil || !strings.Contains(out, "没有相关内容") { + t.Errorf("out=%q err=%v", out, err) + } +} + +func TestUnavailableStore(t *testing.T) { + p := New(nil) + if p.Enabled() { + t.Error("plugin should be disabled without store") + } + if _, err := p.Execute(context.Background(), "knowledge_search", map[string]any{"query": "x"}); err == nil { + t.Error("expected error with nil store") + } +} diff --git a/backend/services/gateway/api/handlers_more.go b/backend/services/gateway/api/handlers_more.go index 0ae4950..5d3daf7 100644 --- a/backend/services/gateway/api/handlers_more.go +++ b/backend/services/gateway/api/handlers_more.go @@ -258,8 +258,7 @@ func (r *Router) sessionExportHandler(w http.ResponseWriter, req *http.Request) } } -func (r *Router) imageAnalyzeHandler(w http.ResponseWriter, req *http.Request) { - if req.Method != "POST" { writeJSON(w, 405, M{"error": "POST required"}); return } +func (r *Router) imageAnalyzeHandler(w http.ResponseWriter, req *http.Request) { if req.Method != "POST" { writeJSON(w, 405, M{"error": "POST required"}); return } if r.deps.Vision == nil { writeJSON(w, 503, M{"error": "vision not available"}); return } var body struct { ImageURL string `json:"image_url"` @@ -305,3 +304,58 @@ func (r *Router) hostExecHandler(w http.ResponseWriter, req *http.Request) { "command": body.Command, }) } + +// ===== Knowledge Base ===== + +func (r *Router) knowledgeDocsHandler(w http.ResponseWriter, req *http.Request) { + if r.deps.Knowledge == nil { writeJSON(w, 503, M{"error": "knowledge base unavailable"}); return } + uid := r.userID(req) + switch req.Method { + case "GET": + docs, err := r.deps.Knowledge.List(req.Context(), uid, 50) + if err != nil { writeJSON(w, 500, M{"error": err.Error()}); return } + writeJSON(w, 200, M{"docs": docs}) + case "POST": + var body struct { + Title string `json:"title"` + Content string `json:"content"` + } + json.NewDecoder(req.Body).Decode(&body) + if body.Content == "" { writeJSON(w, 400, M{"error": "missing content"}); return } + doc, err := r.deps.Knowledge.Ingest(req.Context(), uid, body.Title, body.Content) + if err != nil { writeJSON(w, 500, M{"error": err.Error()}); return } + writeJSON(w, 201, M{"doc": doc}) + default: + writeJSON(w, 405, M{"error": "method not allowed"}) + } +} + +func (r *Router) knowledgeDocItemHandler(w http.ResponseWriter, req *http.Request) { + if r.deps.Knowledge == nil { writeJSON(w, 503, M{"error": "knowledge base unavailable"}); return } + if req.Method != "DELETE" { writeJSON(w, 405, M{"error": "method not allowed"}); return } + idStr := strings.TrimPrefix(req.URL.Path, "/api/v1/knowledge/docs/") + var id int64 + if _, err := fmt.Sscanf(idStr, "%d", &id); err != nil || id <= 0 { + writeJSON(w, 400, M{"error": "invalid id"}) + return + } + if err := r.deps.Knowledge.Delete(req.Context(), r.userID(req), id); err != nil { + writeJSON(w, 404, M{"error": err.Error()}) + return + } + writeJSON(w, 204, nil) +} + +func (r *Router) knowledgeSearchHandler(w http.ResponseWriter, req *http.Request) { + if r.deps.Knowledge == nil { writeJSON(w, 503, M{"error": "knowledge base unavailable"}); return } + if req.Method != "POST" { writeJSON(w, 405, M{"error": "POST required"}); return } + var body struct { + Query string `json:"query"` + TopK int `json:"top_k"` + } + json.NewDecoder(req.Body).Decode(&body) + if body.Query == "" { writeJSON(w, 400, M{"error": "missing query"}); return } + docs, err := r.deps.Knowledge.Search(req.Context(), r.userID(req), body.Query, body.TopK) + if err != nil { writeJSON(w, 500, M{"error": err.Error()}); return } + writeJSON(w, 200, M{"docs": docs, "total": len(docs)}) +} diff --git a/backend/services/gateway/api/router.go b/backend/services/gateway/api/router.go index 609a29e..d4227b4 100644 --- a/backend/services/gateway/api/router.go +++ b/backend/services/gateway/api/router.go @@ -13,6 +13,7 @@ import ( "git.yeij.top/AskaEth/Cyrene_v2/services/ai-core/model" "git.yeij.top/AskaEth/Cyrene_v2/services/gateway/store" + "git.yeij.top/AskaEth/Cyrene_v2/services/rag" ) // --- Interfaces for injected services --- @@ -96,6 +97,14 @@ type VisionAnalyzer interface { AnalyzeImage(ctx context.Context, prompt, imageRef string) (string, error) } +// KnowledgeStorer manages the knowledge base. +type KnowledgeStorer interface { + Ingest(ctx context.Context, userID, title, content string) (*rag.Document, error) + Search(ctx context.Context, userID, query string, topK int) ([]rag.Document, error) + List(ctx context.Context, userID string, limit int) ([]rag.Document, error) + Delete(ctx context.Context, userID string, id int64) error +} + // ChatStorer persists chat sessions and messages. type ChatStorer interface { CreateSession(ctx context.Context, id, userID, title string) (*store.Session, error) @@ -123,6 +132,7 @@ type Dependencies struct { AutomationEng AutomationEngine Reminders *ReminderStore Vision VisionAnalyzer + Knowledge KnowledgeStorer PluginList []map[string]string } @@ -184,6 +194,11 @@ func (r *Router) setupRoutes() { // Message search & session export r.mux.HandleFunc("/api/v1/messages/search", r.authMW(r.messageSearchHandler)) + // Knowledge base + r.mux.HandleFunc("/api/v1/knowledge/docs", r.authMW(r.knowledgeDocsHandler)) + r.mux.HandleFunc("/api/v1/knowledge/docs/", r.authMW(r.knowledgeDocItemHandler)) + r.mux.HandleFunc("/api/v1/knowledge/search", r.authMW(r.knowledgeSearchHandler)) + // Image analysis r.mux.HandleFunc("/api/v1/images/analyze", r.authMW(r.imageAnalyzeHandler)) diff --git a/backend/services/rag/rag.go b/backend/services/rag/rag.go new file mode 100644 index 0000000..c5f7b5d --- /dev/null +++ b/backend/services/rag/rag.go @@ -0,0 +1,181 @@ +// Package rag provides a minimal pgvector-backed knowledge base. +package rag + +import ( + "context" + "database/sql" + "fmt" + "strings" + "time" + + "git.yeij.top/AskaEth/Cyrene_v2/services/memory" +) + +// Document is a knowledge base entry. +type Document struct { + ID int64 `json:"id"` + UserID string `json:"user_id"` + Title string `json:"title"` + Content string `json:"content"` + CreatedAt time.Time `json:"created_at"` +} + +// Embedder is the subset of memory.Embedder used here. +type Embedder interface { + Embed(ctx context.Context, text string) ([]float64, error) + Dims() int +} + +// Store persists and searches documents. +type Store struct { + db *sql.DB + embedder Embedder + vectorEnabled bool +} + +// NewStore creates a knowledge store. Embedder may be nil (ILIKE-only search). +func NewStore(db *sql.DB, embedder Embedder) *Store { + return &Store{db: db, embedder: embedder} +} + +// EnableVector marks the embedding column as available (after migration). +func (s *Store) EnableVector() { s.vectorEnabled = true } + +// CreateTableSQL returns the base schema. +func CreateTableSQL() string { + return ` +CREATE TABLE IF NOT EXISTS knowledge_docs ( + id BIGSERIAL PRIMARY KEY, + user_id TEXT NOT NULL, + title TEXT NOT NULL DEFAULT '', + content TEXT NOT NULL, + created_at TIMESTAMPTZ NOT NULL DEFAULT now() +); +CREATE INDEX IF NOT EXISTS idx_knowledge_docs_user ON knowledge_docs(user_id, id DESC); +` +} + +// VectorMigrationSQL adds the embedding column (pgvector required). +func VectorMigrationSQL(dims int) string { + return fmt.Sprintf(` + ALTER TABLE knowledge_docs ADD COLUMN IF NOT EXISTS embedding vector(%d); + CREATE INDEX IF NOT EXISTS idx_knowledge_docs_embedding ON knowledge_docs USING hnsw (embedding vector_cosine_ops); + `, dims) +} + +// Ingest stores a document, embedding it when an embedder is available. +func (s *Store) Ingest(ctx context.Context, userID, title, content string) (*Document, error) { + doc := &Document{UserID: userID, Title: title, Content: content} + if s.embedder != nil && s.vectorEnabled { + if vec, err := s.embedder.Embed(ctx, title+"\n"+content); err == nil { + _, err := s.db.ExecContext(ctx, ` + INSERT INTO knowledge_docs (user_id, title, content, embedding) VALUES ($1, $2, $3, $4::vector) + `, userID, title, content, vectorString(vec)) + if err != nil { + return nil, fmt.Errorf("ingest: %w", err) + } + return doc, nil + } + } + err := s.db.QueryRowContext(ctx, ` + INSERT INTO knowledge_docs (user_id, title, content) VALUES ($1, $2, $3) RETURNING id, created_at + `, userID, title, content).Scan(&doc.ID, &doc.CreatedAt) + if err != nil { + return nil, fmt.Errorf("ingest: %w", err) + } + return doc, nil +} + +// Search finds documents by vector similarity when enabled, else ILIKE. +// The knowledge base is shared: userID is recorded on ingest but search +// is global (single-user self-hosted deployment). +func (s *Store) Search(ctx context.Context, userID, query string, topK int) ([]Document, error) { + if topK <= 0 || topK > 20 { + topK = 5 + } + if s.embedder != nil && s.vectorEnabled { + vec, err := s.embedder.Embed(ctx, query) + if err == nil { + return s.queryDocs(ctx, ` + SELECT id, user_id, title, content, created_at FROM knowledge_docs + WHERE embedding IS NOT NULL + ORDER BY embedding <=> $1::vector LIMIT $2 + `, vectorString(vec), topK) + } + } + return s.queryDocs(ctx, ` + SELECT id, user_id, title, content, created_at FROM knowledge_docs + WHERE title ILIKE '%' || $1 || '%' OR content ILIKE '%' || $1 || '%' + ORDER BY id DESC LIMIT $2 + `, query, topK) +} + +// List returns recent documents (without content payload). +func (s *Store) List(ctx context.Context, userID string, limit int) ([]Document, error) { + if limit <= 0 { + limit = 50 + } + rows, err := s.db.QueryContext(ctx, ` + SELECT id, user_id, title, created_at FROM knowledge_docs + WHERE user_id = $1 ORDER BY id DESC LIMIT $2 + `, userID, limit) + if err != nil { + return nil, err + } + defer rows.Close() + docs := []Document{} + for rows.Next() { + var d Document + if err := rows.Scan(&d.ID, &d.UserID, &d.Title, &d.CreatedAt); err != nil { + return nil, err + } + docs = append(docs, d) + } + return docs, rows.Err() +} + +// Delete removes a document owned by the user. +func (s *Store) Delete(ctx context.Context, userID string, id int64) error { + res, err := s.db.ExecContext(ctx, `DELETE FROM knowledge_docs WHERE id = $1 AND user_id = $2`, id, userID) + if err != nil { + return err + } + if n, _ := res.RowsAffected(); n == 0 { + return fmt.Errorf("document not found") + } + return nil +} + +func (s *Store) queryDocs(ctx context.Context, query string, args ...any) ([]Document, error) { + rows, err := s.db.QueryContext(ctx, query, args...) + if err != nil { + return nil, err + } + defer rows.Close() + docs := []Document{} + for rows.Next() { + var d Document + if err := rows.Scan(&d.ID, &d.UserID, &d.Title, &d.Content, &d.CreatedAt); err != nil { + return nil, err + } + docs = append(docs, d) + } + return docs, rows.Err() +} + +// vectorString formats an embedding for pgvector input. +func vectorString(vec []float64) string { + var sb strings.Builder + sb.WriteByte('[') + for i, v := range vec { + if i > 0 { + sb.WriteByte(',') + } + fmt.Fprintf(&sb, "%g", v) + } + sb.WriteByte(']') + return sb.String() +} + +// Ensure the memory embedder satisfies the interface at compile time. +var _ Embedder = (*memory.Embedder)(nil)