Move ingestion to CLI and add server upsert API
This commit is contained in:
shenlan
parent
792eb04a71
commit
4f47c61837
110
cmd/cli/main.go
110
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)
|
||||
}
|
||||
}
|
||||
|
||||
@ -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) {
|
||||
|
||||
@ -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 {
|
||||
|
||||
@ -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
|
||||
|
||||
Loading…
Reference in New Issue
Block a user