diff --git a/cmd/cli/main.go b/cmd/cli/main.go index 279fe14..52819b0 100644 --- a/cmd/cli/main.go +++ b/cmd/cli/main.go @@ -1,49 +1,111 @@ package main import ( + "bytes" "context" + "encoding/json" "flag" - "io" "log" "net/http" "os" + "path/filepath" + "strings" "time" rconfig "xcontrol/server/rag/config" + "xcontrol/server/rag/embed" + "xcontrol/server/rag/ingest" + "xcontrol/server/rag/store" + rsync "xcontrol/server/rag/sync" ) -// main loads server RAG configuration and triggers a manual sync by -// calling the running API server's /api/rag/sync endpoint. +// main performs cloning, parsing and embedding before sending documents to the server. func main() { - configPath := flag.String("config", "", "Path to server RAG configuration file") + configPath := flag.String("config", "", "Path to RAG configuration file") flag.Parse() - if *configPath != "" { - if _, err := rconfig.Load(*configPath); err != nil { - log.Fatalf("load config: %v", err) - } + if *configPath == "" { + log.Fatalf("config path required") + } + cfg, err := rconfig.Load(*configPath) + if err != nil { + log.Fatalf("load config: %v", err) } baseURL := os.Getenv("SERVER_URL") if baseURL == "" { baseURL = "http://localhost:8080" } - ctx, cancel := context.WithTimeout(context.Background(), 10*time.Minute) + + ctx, cancel := context.WithTimeout(context.Background(), 30*time.Minute) defer cancel() - req, err := http.NewRequestWithContext(ctx, http.MethodPost, baseURL+"/api/rag/sync", nil) - if err != nil { - log.Fatalf("create request: %v", err) - } - resp, err := http.DefaultClient.Do(req) - if err != nil { - log.Fatalf("sync request: %v", err) - } - defer resp.Body.Close() - if resp.StatusCode != http.StatusOK { - b, _ := io.ReadAll(resp.Body) - log.Fatalf("sync failed: %s", string(b)) - } - if _, err := io.Copy(os.Stdout, resp.Body); err != nil { - log.Fatalf("read response: %v", err) + embCfg := cfg.ResolveEmbedding() + chunkCfg := cfg.ResolveChunking() + embedder := embed.NewOpenAI(embCfg.BaseURL, embCfg.APIKey, embCfg.Model, embCfg.Dimension) + + for _, ds := range cfg.Global.Datasources { + workdir := filepath.Join(os.TempDir(), "xcontrol", ds.Name) + if _, err := rsync.SyncRepo(ctx, ds.Repo, workdir); err != nil { + log.Fatalf("sync repo %s: %v", ds.Name, err) + } + root := filepath.Join(workdir, ds.Path) + files, err := ingest.ListMarkdown(root, chunkCfg.IncludeExts, chunkCfg.IgnoreDirs, 0) + if err != nil { + log.Fatalf("list markdown: %v", err) + } + var rows []store.DocRow + for _, f := range files { + secs, err := ingest.ParseMarkdown(f) + if err != nil { + log.Printf("parse %s: %v", f, err) + continue + } + chunks, err := ingest.BuildChunks(secs, chunkCfg) + if err != nil { + log.Printf("chunk %s: %v", f, err) + continue + } + texts := make([]string, len(chunks)) + rs := make([]store.DocRow, len(chunks)) + for i, ch := range chunks { + texts[i] = ch.Text + rs[i] = store.DocRow{ + Repo: ds.Repo, + Path: strings.TrimPrefix(f, workdir+"/"), + ChunkID: ch.ChunkID, + Content: ch.Text, + Metadata: ch.Meta, + ContentSHA: ch.SHA256, + } + } + vecs, _, err := embedder.Embed(ctx, texts) + if err != nil { + log.Printf("embed %s: %v", f, err) + continue + } + for i := range rs { + rs[i].Embedding = vecs[i] + } + rows = append(rows, rs...) + } + + payload := struct { + Docs []store.DocRow `json:"docs"` + }{Docs: rows} + b, _ := json.Marshal(payload) + req, err := http.NewRequestWithContext(ctx, http.MethodPost, baseURL+"/api/rag/upsert", bytes.NewReader(b)) + if err != nil { + log.Fatalf("create request: %v", err) + } + req.Header.Set("Content-Type", "application/json") + resp, err := http.DefaultClient.Do(req) + if err != nil { + log.Fatalf("upsert request: %v", err) + } + resp.Body.Close() + if resp.StatusCode != http.StatusOK { + log.Fatalf("upsert failed: %s", resp.Status) + } + log.Printf("ingested %d rows for %s", len(rows), ds.Name) } } diff --git a/server/api/rag.go b/server/api/rag.go index 1599e8a..8db465b 100644 --- a/server/api/rag.go +++ b/server/api/rag.go @@ -1,13 +1,13 @@ package api import ( - "context" "net/http" "github.com/gin-gonic/gin" "xcontrol/server/rag" rconfig "xcontrol/server/rag/config" + "xcontrol/server/rag/store" ) // ragSvc provides repository sync and retrieval operations. @@ -19,29 +19,29 @@ func initRAG() *rag.Service { if err != nil { return nil } - svc := rag.New(cfg.ToConfig()) - go svc.Sync(context.Background()) - go svc.Watch(context.Background()) - return svc + return rag.New(cfg.ToConfig()) } // registerRAGRoutes wires the /api/rag endpoints. func registerRAGRoutes(r *gin.RouterGroup) { - r.POST("/rag/sync", func(c *gin.Context) { + r.POST("/rag/upsert", func(c *gin.Context) { if ragSvc == nil { - c.String(http.StatusOK, "rag service not initialized\n") + c.JSON(http.StatusOK, gin.H{"rows": 0}) return } - c.Writer.Header().Set("Content-Type", "text/plain; charset=utf-8") - c.Status(http.StatusOK) - err := ragSvc.SyncWithProgress(c.Request.Context(), func(msg string) { - _, _ = c.Writer.Write([]byte(msg + "\n")) - c.Writer.Flush() - }) + var req struct { + Docs []store.DocRow `json:"docs"` + } + if err := c.BindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + n, err := ragSvc.Upsert(c.Request.Context(), req.Docs) if err != nil { - _, _ = c.Writer.Write([]byte("error: " + err.Error() + "\n")) + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) return } + c.JSON(http.StatusOK, gin.H{"rows": n}) }) r.POST("/rag/query", func(c *gin.Context) { diff --git a/server/rag/service.go b/server/rag/service.go index f18e6ef..fe40f11 100644 --- a/server/rag/service.go +++ b/server/rag/service.go @@ -3,14 +3,13 @@ package rag import ( "context" "encoding/json" - "time" "github.com/jackc/pgx/v5" pgvector "github.com/pgvector/pgvector-go" "xcontrol/server/rag/config" "xcontrol/server/rag/embed" - "xcontrol/server/rag/ingest" + "xcontrol/server/rag/store" ) type Service struct { @@ -21,49 +20,26 @@ func New(cfg *config.Config) *Service { return &Service{cfg: cfg} } -func (s *Service) Sync(ctx context.Context) error { - return s.SyncWithProgress(ctx, nil) -} +// Upsert stores pre-embedded documents into the vector database. +func (s *Service) Upsert(ctx context.Context, rows []store.DocRow) (int, error) { + if s == nil || s.cfg == nil || len(rows) == 0 { + return 0, nil + } + dsn := s.cfg.Global.VectorDB.DSN() + if dsn == "" { + return 0, nil + } + conn, err := pgx.Connect(ctx, dsn) + if err != nil { + return 0, err + } + defer conn.Close(ctx) -// SyncWithProgress performs a full sync while reporting progress via the provided callback. -// -// The progress callback may be nil. When non-nil it will receive human readable -// status messages as the sync operation progresses. -func (s *Service) SyncWithProgress(ctx context.Context, progress func(string)) error { - if s == nil || s.cfg == nil { - return nil - } - for _, ds := range s.cfg.Global.Datasources { - if progress != nil { - progress("syncing " + ds.Name) - } - if _, err := ingest.IngestRepo(ctx, s.cfg, ds, ingest.Options{}); err != nil { - if progress != nil { - progress("error syncing " + ds.Name + ": " + err.Error()) - } - return err - } - if progress != nil { - progress("completed " + ds.Name) - } - } - return nil -} - -func (s *Service) Watch(ctx context.Context) { - if s == nil || s.cfg == nil { - return - } - ticker := time.NewTicker(time.Minute) - defer ticker.Stop() - for { - select { - case <-ctx.Done(): - return - case <-ticker.C: - _ = s.Sync(ctx) - } + dim := len(rows[0].Embedding) + if err := store.EnsureSchema(ctx, conn, dim, false); err != nil { + return 0, err } + return store.UpsertDocuments(ctx, conn, rows) } type Document struct { diff --git a/server/rag/store/store.go b/server/rag/store/store.go index e8380a9..8e1401f 100644 --- a/server/rag/store/store.go +++ b/server/rag/store/store.go @@ -11,13 +11,13 @@ import ( // DocRow represents a row to be stored in the documents table. type DocRow struct { - Repo string - Path string - ChunkID int - Content string - Embedding []float32 - Metadata map[string]any - ContentSHA string + Repo string `json:"repo"` + Path string `json:"path"` + ChunkID int `json:"chunk_id"` + Content string `json:"content"` + Embedding []float32 `json:"embedding"` + Metadata map[string]any `json:"metadata"` + ContentSHA string `json:"content_sha"` } // EnsureSchema creates the documents table and indexes if they do not exist. It