diff --git a/README.md b/README.md index e65fc45..e96f447 100644 --- a/README.md +++ b/README.md @@ -226,6 +226,7 @@ LocalRecall uses environment variables to configure its behavior. These variable | `VECTOR_ENGINE` | Vector database engine to use (`chromem` by default, `postgres` for PostgreSQL). | | `MAX_CHUNKING_SIZE` | Maximum size (in characters) for breaking down documents into chunks. Affects performance and accuracy. | | `CHUNK_OVERLAP` | Overlap in characters between consecutive chunks (word-aligned). Default: 0. Use to improve context across chunk boundaries. | +| `CHUNK_BY_PARAGRAPH` | When not `false` (default), files are chunked per blank-line-separated paragraph so chunk boundaries never cross a paragraph separator (e.g. one passage per line in mythological texts). Set to `false` for the legacy single-block chunking. | | `HYBRID_SEARCH_BM25_WEIGHT` | Weight for BM25 keyword search in hybrid search (default: 0.5, PostgreSQL only). | | `HYBRID_SEARCH_VECTOR_WEIGHT` | Weight for vector similarity search in hybrid search (default: 0.5, PostgreSQL only). | | `POSTGRES_LOCK_TIMEOUT` | Per-connection `lock_timeout` for the PostgreSQL engine (default: `30s`). Bounds how long a statement waits to acquire a lock so a single stuck operation cannot make every other statement on the table queue indefinitely. Set to `0`/`off` to disable. | @@ -255,13 +256,16 @@ curl -X POST $BASE_URL/collections \ -d '{"name":"myCollection"}' ``` -- **Upload File**: +- **Upload File** (with optional JSON metadata attached to every chunk): ```sh curl -X POST $BASE_URL/collections/myCollection/upload \ - -F "file=@/path/to/file.txt" + -F "file=@/path/to/file.txt" \ + -F 'metadata={"work":"Pyramid Texts","lang":"fr","form":"transliteration"}' ``` +The `metadata` form field is a JSON object of string values. Keys `source`, `type`, `file_name` and `memory_*` are reserved and rejected. + - **List Collections**: ```sh @@ -290,14 +294,16 @@ curl -X GET $BASE_URL/collections/myCollection/entries/file.pdf/raw Returns the original uploaded binary file with the appropriate Content-Type header. -- **Search Collection**: +- **Search Collection** (with optional metadata filters, AND semantics): ```sh curl -X POST $BASE_URL/collections/myCollection/search \ -H "Content-Type: application/json" \ - -d '{"query":"search term", "max_results":5}' + -d '{"query":"search term", "max_results":5, "filters":{"lang":"fr","form":"gardiner"}}' ``` +Filters are supported by the `chromem` and `postgres` engines; the `localai` engine returns an error when `filters` is non-empty. + - **Reset Collection**: ```sh diff --git a/pkg/chunk/chunk_test.go b/pkg/chunk/chunk_test.go index c25e8a5..aff5025 100644 --- a/pkg/chunk/chunk_test.go +++ b/pkg/chunk/chunk_test.go @@ -127,3 +127,19 @@ var _ = Describe("Chunk", func() { }) }) }) + +var _ = Describe("SplitIntoParagraphs", func() { + It("splits on blank lines and trims paragraphs", func() { + content := "first line\nsecond line\n\n\n third paragraph \n" + Expect(SplitIntoParagraphs(content)).To(Equal([]string{"first line\nsecond line", "third paragraph"})) + }) + + It("returns a single paragraph when there are no blank lines", func() { + Expect(SplitIntoParagraphs("one\ntwo\nthree")).To(Equal([]string{"one\ntwo\nthree"})) + }) + + It("returns no paragraphs for empty or whitespace-only content", func() { + Expect(SplitIntoParagraphs("")).To(BeEmpty()) + Expect(SplitIntoParagraphs(" \n \n\t\n")).To(BeEmpty()) + }) +}) diff --git a/pkg/chunk/chunking.go b/pkg/chunk/chunking.go index 31231a4..27ade1f 100644 --- a/pkg/chunk/chunking.go +++ b/pkg/chunk/chunking.go @@ -177,6 +177,31 @@ func SplitParagraphIntoChunksWithOptions(paragraph string, opts Options) []strin return chunks } +// SplitIntoParagraphs splits content into paragraphs separated by one or more +// blank lines. Each paragraph is trimmed; empty paragraphs are dropped. +// A content without blank lines yields a single paragraph. +func SplitIntoParagraphs(content string) []string { + raw := strings.Split(content, "\n") + var paragraphs []string + var current []string + flush := func() { + para := strings.TrimSpace(strings.Join(current, "\n")) + if para != "" { + paragraphs = append(paragraphs, para) + } + current = nil + } + for _, line := range raw { + if strings.TrimSpace(line) == "" { + flush() + continue + } + current = append(current, line) + } + flush() + return paragraphs +} + // SplitParagraphIntoChunks takes a paragraph and a maxChunkSize as input, // and returns a slice of strings where each string is a chunk of the paragraph // that is at most maxChunkSize long, ensuring that words are not split. diff --git a/rag/engine/chromem.go b/rag/engine/chromem.go index a5f0fe8..e9481b6 100644 --- a/rag/engine/chromem.go +++ b/rag/engine/chromem.go @@ -265,6 +265,10 @@ func (c *ChromemDB) GetBySource(source string) ([]types.Result, error) { } func (c *ChromemDB) Search(s string, similarEntries int) ([]types.Result, error) { + return c.SearchWithFilters(s, similarEntries, nil) +} + +func (c *ChromemDB) SearchWithFilters(s string, similarEntries int, where map[string]string) ([]types.Result, error) { // chromem rejects a query limit larger than the collection, including an // empty collection after complete memory deletion. if n := c.collection.Count(); similarEntries > n { @@ -273,7 +277,7 @@ func (c *ChromemDB) Search(s string, similarEntries int) ([]types.Result, error) } similarEntries = n } - res, err := c.query(context.Background(), s, similarEntries, nil) + res, err := c.query(context.Background(), s, similarEntries, where) if err != nil { return nil, err } diff --git a/rag/engine/localai.go b/rag/engine/localai.go index 4c81201..346a411 100644 --- a/rag/engine/localai.go +++ b/rag/engine/localai.go @@ -92,6 +92,12 @@ func (db *LocalAIRAGDB) GetBySource(source string) ([]types.Result, error) { return nil, fmt.Errorf("not implemented") } +// SearchWithFilters is unsupported: the LocalAI store find API has no metadata +// filter, and filtering client-side would silently change the top-k semantics. +func (db *LocalAIRAGDB) SearchWithFilters(s string, similarEntries int, where map[string]string) ([]types.Result, error) { + return nil, fmt.Errorf("metadata filters are not supported by the localai engine") +} + func (db *LocalAIRAGDB) Search(s string, similarEntries int) ([]types.Result, error) { resp, err := db.openaiClient.CreateEmbeddings(context.TODO(), openai.EmbeddingRequestStrings{ diff --git a/rag/engine/mock.go b/rag/engine/mock.go index 92bffe3..7b67e68 100644 --- a/rag/engine/mock.go +++ b/rag/engine/mock.go @@ -58,6 +58,43 @@ func (m *MockEngine) StoreDocuments(s []string, metadata map[string]string) ([]R return results, nil } +func (m *MockEngine) SearchWithFilters(s string, similarEntries int, where map[string]string) ([]types.Result, error) { + m.mu.Lock() + defer m.mu.Unlock() + + var results []types.Result + for _, doc := range m.docs { + match := true + for k, v := range where { + if doc.Metadata[k] != v { + match = false + break + } + } + if match && strings.Contains(strings.ToLower(doc.Content), strings.ToLower(s)) { + results = append(results, doc) + } + } + if len(results) == 0 { + for _, doc := range m.docs { + match := true + for k, v := range where { + if doc.Metadata[k] != v { + match = false + break + } + } + if match { + results = append(results, doc) + } + } + } + if len(results) > similarEntries { + results = results[:similarEntries] + } + return results, nil +} + func (m *MockEngine) Search(s string, similarEntries int) ([]types.Result, error) { m.mu.Lock() defer m.mu.Unlock() diff --git a/rag/engine/postgres.go b/rag/engine/postgres.go index f05f7a2..b207ed1 100644 --- a/rag/engine/postgres.go +++ b/rag/engine/postgres.go @@ -1055,7 +1055,59 @@ func buildScopedHybridSearchQuery(tableName, indexName string) string { `, tableName, candidatePool, rrfK, strings.ReplaceAll(indexName, "'", "''")) } +// buildScopedHybridSearchQueryFiltered is buildScopedHybridSearchQuery with an +// additional metadata containment filter ($6::jsonb) on the final join, so only +// documents whose metadata contains every filter entry are returned. +func buildScopedHybridSearchQueryFiltered(tableName, indexName string) string { + candidatePool := fmt.Sprintf("GREATEST($5 * %d, %d)", hybridCandidateMultiplier, hybridCandidateFloor) + return fmt.Sprintf(` + WITH bm25_results AS ( + SELECT id, ROW_NUMBER() OVER (ORDER BY full_text <@> to_bm25query($1, '%[4]s')) AS rank + FROM %[1]s + ORDER BY full_text <@> to_bm25query($1, '%[4]s') + LIMIT %[2]s + ), + vector_results AS ( + SELECT id, ROW_NUMBER() OVER (ORDER BY embedding <=> $3::vector) AS rank + FROM %[1]s + WHERE embedding IS NOT NULL + ORDER BY embedding <=> $3::vector + LIMIT %[2]s + ), + fused AS ( + SELECT + COALESCE(b.id, v.id) AS id, + COALESCE($2::float8 / (%[3]d + b.rank), 0) + COALESCE($4::float8 / (%[3]d + v.rank), 0) AS similarity + FROM bm25_results b + FULL OUTER JOIN vector_results v ON b.id = v.id + ) + SELECT + d.id::text, + COALESCE(d.title, '') as title, + d.content, + d.metadata, + f.similarity + FROM fused f + JOIN %[1]s d ON d.id = f.id AND d.metadata @> $6::jsonb + ORDER BY f.similarity DESC + LIMIT $5 + `, tableName, candidatePool, rrfK, strings.ReplaceAll(indexName, "'", "''")) +} + func (p *PostgresDB) Search(s string, similarEntries int) ([]types.Result, error) { + return p.search(s, similarEntries, nil) +} + +// SearchWithFilters restricts the hybrid search to rows whose JSONB metadata +// contains every entry of where (AND semantics, via the @> containment operator). +func (p *PostgresDB) SearchWithFilters(s string, similarEntries int, where map[string]string) ([]types.Result, error) { + if len(where) == 0 { + return p.Search(s, similarEntries) + } + return p.search(s, similarEntries, where) +} + +func (p *PostgresDB) search(s string, similarEntries int, where map[string]string) ([]types.Result, error) { ctx := context.Background() // Get query embedding @@ -1065,15 +1117,28 @@ func (p *PostgresDB) Search(s string, similarEntries int) ([]types.Result, error } queryEmbeddingStr := formatVector(queryEmbedding) + var filterJSON []byte + if len(where) > 0 { + filterJSON, err = json.Marshal(where) + if err != nil { + return nil, fmt.Errorf("failed to marshal metadata filters: %w", err) + } + } + // Build hybrid search query (BM25 + vector similarity) query := buildScopedHybridSearchQuery(p.tableName, p.scope.relation(p.indexBase("bm25"))) + args := []interface{}{s, p.bm25Weight, queryEmbeddingStr, p.vectorWeight, similarEntries} + if filterJSON != nil { + query = buildScopedHybridSearchQueryFiltered(p.tableName, p.scope.relation(p.indexBase("bm25"))) + args = append(args, string(filterJSON)) + } - rows, err := p.pool.Query(ctx, query, s, p.bm25Weight, queryEmbeddingStr, p.vectorWeight, similarEntries) + rows, err := p.pool.Query(ctx, query, args...) if err != nil { // If BM25 query fails, fallback to vector-only search xlog.Warn("BM25 search failed, falling back to vector search", "error", err) query = fmt.Sprintf(` - SELECT + SELECT id::text, COALESCE(title, '') as title, content, @@ -1084,7 +1149,23 @@ func (p *PostgresDB) Search(s string, similarEntries int) ([]types.Result, error ORDER BY embedding <=> $1::vector LIMIT $2 `, p.tableName) - rows, err = p.pool.Query(ctx, query, queryEmbeddingStr, similarEntries) + fallbackArgs := []interface{}{queryEmbeddingStr, similarEntries} + if filterJSON != nil { + query = fmt.Sprintf(` + SELECT + id::text, + COALESCE(title, '') as title, + content, + metadata, + (1 - (embedding <=> $1::vector)) as similarity + FROM %s + WHERE embedding IS NOT NULL AND metadata @> $3::jsonb + ORDER BY embedding <=> $1::vector + LIMIT $2 + `, p.tableName) + fallbackArgs = append(fallbackArgs, string(filterJSON)) + } + rows, err = p.pool.Query(ctx, query, fallbackArgs...) if err != nil { return nil, fmt.Errorf("failed to execute search: %w", err) } diff --git a/rag/engine/types.go b/rag/engine/types.go index d625c2e..e3f1267 100644 --- a/rag/engine/types.go +++ b/rag/engine/types.go @@ -1,5 +1,14 @@ package engine +import "github.com/mudler/localrecall/rag/types" + type Result struct { ID string `json:"id"` } + +// FilteredSearcher is an optional extension of the storage engines: when an +// engine implements it, callers can restrict search results to documents whose +// metadata matches all entries of the where map (AND semantics). +type FilteredSearcher interface { + SearchWithFilters(s string, similarEntries int, where map[string]string) ([]types.Result, error) +} diff --git a/rag/persistency.go b/rag/persistency.go index 235319a..c9b58b2 100644 --- a/rag/persistency.go +++ b/rag/persistency.go @@ -209,6 +209,23 @@ func (db *PersistentKB) Search(s string, similarEntries int) ([]types.Result, er return db.Engine.Search(s, similarEntries) } +// SearchWithFilters behaves like Search but restricts results to chunks whose +// metadata matches all entries of filters (AND semantics). Engines that do not +// implement engine.FilteredSearcher return an error when filters are non-empty. +func (db *PersistentKB) SearchWithFilters(s string, similarEntries int, filters map[string]string) ([]types.Result, error) { + db.Lock() + defer db.Unlock() + + if len(filters) == 0 { + return db.Engine.Search(s, similarEntries) + } + fs, ok := db.Engine.(engine.FilteredSearcher) + if !ok { + return nil, fmt.Errorf("metadata filters are not supported by engine %T", db.Engine) + } + return fs.SearchWithFilters(s, similarEntries, filters) +} + func (db *PersistentKB) Reset() error { db.Lock() defer db.Unlock() @@ -767,11 +784,27 @@ func chunkFile(fpath string, maxchunksize, chunkOverlap int) ([]string, error) { } opts := chunk.Options{MaxSize: maxchunksize, Overlap: chunkOverlap, SplitLongWords: true} - chunks := chunk.SplitParagraphIntoChunksWithOptions(content, opts) + var chunks []string + if paragraphChunkingEnabled() { + // Chunk each paragraph separately so chunk boundaries never cross a + // blank-line separator (e.g. one line per passage in mythological texts). + for _, para := range chunk.SplitIntoParagraphs(content) { + chunks = append(chunks, chunk.SplitParagraphIntoChunksWithOptions(para, opts)...) + } + } else { + chunks = chunk.SplitParagraphIntoChunksWithOptions(content, opts) + } xlog.Info("Chunked file", "file", fpath, "content_length", len(content), "max_chunk_size", maxchunksize, "chunk_overlap", chunkOverlap, "chunk_count", len(chunks)) return chunks, nil } +// paragraphChunkingEnabled reports whether files are chunked per paragraph +// (blank-line separated) instead of as a single block. Enabled by default; +// set CHUNK_BY_PARAGRAPH=false to restore the legacy single-block behavior. +func paragraphChunkingEnabled() bool { + return os.Getenv("CHUNK_BY_PARAGRAPH") != "false" +} + // migrateToUUIDLayout migrates flat files in assetDir (files not in UUID // subdirectories) to UUID subdirectory layout. This is a one-time migration. func (db *PersistentKB) migrateToUUIDLayout() error { diff --git a/routes.go b/routes.go index dc990c5..61a5bce 100644 --- a/routes.go +++ b/routes.go @@ -2,6 +2,7 @@ package main import ( "crypto/subtle" + "encoding/json" "fmt" "io" "net/http" @@ -281,8 +282,9 @@ func search(collections collectionList) func(c echo.Context) error { } type request struct { - Query string `json:"query"` - MaxResults int `json:"max_results"` + Query string `json:"query"` + MaxResults int `json:"max_results"` + Filters map[string]string `json:"filters"` } r := new(request) @@ -298,7 +300,7 @@ func search(collections collectionList) func(c echo.Context) error { } } - results, err := collection.Search(r.Query, r.MaxResults) + results, err := collection.SearchWithFilters(r.Query, r.MaxResults, r.Filters) if err != nil { return c.JSON(http.StatusInternalServerError, errorResponse(ErrCodeInternalError, "Failed to search collection", err.Error())) } @@ -448,8 +450,24 @@ func uploadFile(collections collectionList, fileAssets string) func(c echo.Conte now := time.Now().Format(time.RFC3339) + metadata := map[string]string{"created_at": now} + // Optional caller-supplied metadata (JSON object in the "metadata" + // form field), e.g. {"work": "Pyramid Texts", "lang": "fr", "form": "transliteration"}. + if raw := c.FormValue("metadata"); raw != "" { + extra := map[string]string{} + if err := json.Unmarshal([]byte(raw), &extra); err != nil { + return c.JSON(http.StatusBadRequest, errorResponse(ErrCodeInvalidRequest, "Invalid metadata", "metadata must be a JSON object of string values: "+err.Error())) + } + for k, v := range extra { + if k == "source" || k == "type" || k == "file_name" || strings.HasPrefix(k, "memory_") { + return c.JSON(http.StatusBadRequest, errorResponse(ErrCodeInvalidRequest, "Reserved metadata key", fmt.Sprintf("metadata key %q is reserved", k))) + } + metadata[k] = v + } + } + // Save the file to disk - key, err := collection.Store(uploadPath, map[string]string{"created_at": now}) + key, err := collection.Store(uploadPath, metadata) if err != nil { xlog.Error("Failed to store file", err) return c.JSON(http.StatusInternalServerError, errorResponse(ErrCodeInternalError, "Failed to store file", err.Error())) diff --git a/routes_test.go b/routes_test.go index 36edc50..0883d06 100644 --- a/routes_test.go +++ b/routes_test.go @@ -1,6 +1,9 @@ package main import ( + "bytes" + "encoding/json" + "mime/multipart" "net/http" "net/http/httptest" "strings" @@ -8,6 +11,7 @@ import ( "github.com/labstack/echo/v4" "github.com/mudler/localrecall/rag" + "github.com/mudler/localrecall/rag/engine" ) // Creating a collection that is already loaded must not build a new engine: @@ -38,3 +42,99 @@ func TestCreateCollectionReusesLoadedCollection(t *testing.T) { t.Fatal("the loaded collection was replaced") } } + +// The search endpoint must forward metadata filters to the engine and return +// only the matching chunks. +func TestSearchWithFilters(t *testing.T) { + mock := engine.NewMockEngine() + if _, err := mock.Store("Osiris texte en français", map[string]string{"lang": "fr", "form": "transliteration"}); err != nil { + t.Fatalf("store: %v", err) + } + if _, err := mock.Store("texte en codes Gardiner A1", map[string]string{"lang": "en", "form": "gardiner"}); err != nil { + t.Fatalf("store: %v", err) + } + + kb := &rag.PersistentKB{Engine: mock} + origLookup := lookupCollection + lookupCollection = func(name string) (*rag.PersistentKB, bool) { return kb, true } + defer func() { lookupCollection = origLookup }() + + e := echo.New() + body := `{"query":"texte","max_results":5,"filters":{"lang":"fr"}}` + req := httptest.NewRequest(http.MethodPost, "/api/collections/docs/search", strings.NewReader(body)) + req.Header.Set(echo.HeaderContentType, echo.MIMEApplicationJSON) + rec := httptest.NewRecorder() + + if err := search(collectionList{})(e.NewContext(req, rec)); err != nil { + t.Fatalf("handler error: %v", err) + } + if rec.Code != http.StatusOK { + t.Fatalf("status = %d, want %d; body: %s", rec.Code, http.StatusOK, rec.Body.String()) + } + + var resp struct { + Success bool `json:"success"` + Data struct { + Results []struct { + Content string `json:"Content"` + Metadata map[string]string `json:"Metadata"` + } `json:"results"` + } `json:"data"` + } + if err := json.Unmarshal(rec.Body.Bytes(), &resp); err != nil { + t.Fatalf("unmarshal response: %v; body: %s", err, rec.Body.String()) + } + if !resp.Success || len(resp.Data.Results) != 1 { + t.Fatalf("expected exactly 1 filtered result, got %+v; body: %s", resp.Data.Results, rec.Body.String()) + } + if resp.Data.Results[0].Metadata["lang"] != "fr" { + t.Fatalf("expected lang=fr result, got %+v", resp.Data.Results[0].Metadata) + } +} + +// Upload must reject malformed metadata JSON and reserved metadata keys before +// storing anything. +func TestUploadMetadataValidation(t *testing.T) { + kb := &rag.PersistentKB{} + origLookup := lookupCollection + lookupCollection = func(name string) (*rag.PersistentKB, bool) { return kb, true } + defer func() { lookupCollection = origLookup }() + + newUploadRequest := func(metadata string) (*httptest.ResponseRecorder, *http.Request) { + var buf bytes.Buffer + w := multipart.NewWriter(&buf) + fw, err := w.CreateFormFile("file", "spell.txt") + if err != nil { + t.Fatalf("create form file: %v", err) + } + fw.Write([]byte("jnṯ nṯjr ꜥnḫ ḏd w3s")) + if metadata != "" { + if err := w.WriteField("metadata", metadata); err != nil { + t.Fatalf("write metadata field: %v", err) + } + } + w.Close() + req := httptest.NewRequest(http.MethodPost, "/api/collections/docs/upload", &buf) + req.Header.Set(echo.HeaderContentType, w.FormDataContentType()) + return httptest.NewRecorder(), req + } + + e := echo.New() + handler := uploadFile(collectionList{}, t.TempDir()) + + rec, req := newUploadRequest(`{"lang":`) + if err := handler(e.NewContext(req, rec)); err != nil { + t.Fatalf("handler error: %v", err) + } + if rec.Code != http.StatusBadRequest { + t.Fatalf("malformed metadata: status = %d, want %d; body: %s", rec.Code, http.StatusBadRequest, rec.Body.String()) + } + + rec, req = newUploadRequest(`{"source":"evil"}`) + if err := handler(e.NewContext(req, rec)); err != nil { + t.Fatalf("handler error: %v", err) + } + if rec.Code != http.StatusBadRequest { + t.Fatalf("reserved key: status = %d, want %d; body: %s", rec.Code, http.StatusBadRequest, rec.Body.String()) + } +}