graphsearch-rag 0.2.1__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -0,0 +1,21 @@
1
+ MIT License
2
+
3
+ Copyright (c) 2026 GraphSearch contributors
4
+
5
+ Permission is hereby granted, free of charge, to any person obtaining a copy
6
+ of this software and associated documentation files (the "Software"), to deal
7
+ in the Software without restriction, including without limitation the rights
8
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9
+ copies of the Software, and to permit persons to whom the Software is
10
+ furnished to do so, subject to the following conditions:
11
+
12
+ The above copyright notice and this permission notice shall be included in all
13
+ copies or substantial portions of the Software.
14
+
15
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21
+ SOFTWARE.
@@ -0,0 +1,227 @@
1
+ Metadata-Version: 2.4
2
+ Name: graphsearch-rag
3
+ Version: 0.2.1
4
+ Summary: GraphSearch: a GraphQL API server for Retrieval-Augmented Generation (RAG) over your documents
5
+ Author: GraphSearch contributors
6
+ License: MIT
7
+ Project-URL: Homepage, https://github.com/mohithgowdak/graphsearch
8
+ Project-URL: Issues, https://github.com/mohithgowdak/graphsearch/issues
9
+ Keywords: graphql,rag,semantic-search,llm,retrieval-augmented-generation
10
+ Classifier: Development Status :: 3 - Alpha
11
+ Classifier: Intended Audience :: Developers
12
+ Classifier: License :: OSI Approved :: MIT License
13
+ Classifier: Programming Language :: Python :: 3
14
+ Classifier: Topic :: Software Development :: Libraries :: Application Frameworks
15
+ Requires-Python: >=3.10
16
+ Description-Content-Type: text/markdown
17
+ License-File: LICENSE
18
+ Requires-Dist: fastapi>=0.110
19
+ Requires-Dist: strawberry-graphql[fastapi]>=0.220
20
+ Requires-Dist: uvicorn[standard]>=0.29
21
+ Requires-Dist: pydantic-settings>=2.2
22
+ Requires-Dist: numpy>=1.26
23
+ Provides-Extra: local
24
+ Requires-Dist: sentence-transformers>=2.7; extra == "local"
25
+ Provides-Extra: openai
26
+ Requires-Dist: openai>=1.30; extra == "openai"
27
+ Provides-Extra: anthropic
28
+ Requires-Dist: anthropic>=0.30; extra == "anthropic"
29
+ Provides-Extra: dev
30
+ Requires-Dist: pytest>=8.0; extra == "dev"
31
+ Requires-Dist: httpx>=0.27; extra == "dev"
32
+ Requires-Dist: ruff>=0.4; extra == "dev"
33
+ Dynamic: license-file
34
+
35
+ # GraphSearch
36
+
37
+ [![CI](https://github.com/mohithgowdak/graphsearch/actions/workflows/ci.yml/badge.svg)](https://github.com/mohithgowdak/graphsearch/actions/workflows/ci.yml)
38
+ [![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](LICENSE)
39
+ [![Python 3.10+](https://img.shields.io/badge/python-3.10%2B-blue.svg)](pyproject.toml)
40
+
41
+ **A GraphQL API server for Retrieval-Augmented Generation (RAG) over your documents.**
42
+
43
+ Upload documents through a GraphQL mutation, then ask questions through a GraphQL query.
44
+ GraphSearch chunks and embeds your documents, retrieves the most relevant passages with
45
+ vector similarity search, and feeds them to an LLM to generate a grounded answer — all
46
+ behind a single, typed, introspectable GraphQL endpoint.
47
+
48
+ ```graphql
49
+ query {
50
+ answer(question: "How do I reset my password?") {
51
+ text # cites sources as [1], [2], ...
52
+ sources { documentTitle text score }
53
+ }
54
+ }
55
+ ```
56
+
57
+ ![GraphSearch demo: asking "How do I get my money back?" in GraphiQL and getting the returns policy with ranked sources](docs/demo.gif)
58
+
59
+ *Semantic retrieval with the offline `local` backend: "How do I get my money back?" finds the returns policy — no shared keywords needed.*
60
+
61
+ ## Why GraphQL for RAG?
62
+
63
+ - **One endpoint, typed schema** — clients ask for exactly the fields they need
64
+ (answer text, source chunks, scores) instead of juggling REST routes.
65
+ - **Introspection & tooling for free** — GraphiQL, codegen'd TypeScript clients,
66
+ and schema docs come with the ecosystem.
67
+ - **Composable** — `answer`, `search`, and document management live in one schema
68
+ and can be combined in a single request.
69
+
70
+ ## Quickstart
71
+
72
+ ### Install
73
+
74
+ ```bash
75
+ pip install graphsearch-rag # from PyPI (imports as `graphsearch`)
76
+ # or run the prebuilt image:
77
+ docker run -p 8000:8000 ghcr.io/mohithgowdak/graphsearch:latest
78
+ # or from source:
79
+ git clone https://github.com/mohithgowdak/graphsearch && cd graphsearch
80
+ pip install -e ".[dev]"
81
+ ```
82
+
83
+ ### Run locally (no API keys needed)
84
+
85
+ The default configuration is fully offline: a hashing-trick embedder plus an
86
+ extractive answer mode. It exercises the entire pipeline and is perfect for
87
+ kicking the tires, tests, and CI.
88
+
89
+ ```bash
90
+ # Ingest some documents (bundled examples shown; any .md/.txt works)
91
+ graphsearch-ingest data/example_docs
92
+
93
+ # Start the server
94
+ graphsearch
95
+ ```
96
+
97
+ Open **http://localhost:8000/graphql** for the GraphiQL playground and try:
98
+
99
+ ```graphql
100
+ query {
101
+ answer(question: "What is the return policy?") {
102
+ text
103
+ sources { text score }
104
+ }
105
+ }
106
+ ```
107
+
108
+ ### Run with Docker
109
+
110
+ ```bash
111
+ docker compose up --build
112
+ ```
113
+
114
+ ### Semantic search without any API key
115
+
116
+ The `local` backend runs [sentence-transformers](https://www.sbert.net/)
117
+ on your CPU — real semantic retrieval ("How do I get my money back?" finds
118
+ the refund policy), still zero keys and zero external services:
119
+
120
+ ```bash
121
+ pip install -e ".[local]"
122
+ export GRAPHSEARCH_EMBEDDINGS=local # $env:GRAPHSEARCH_EMBEDDINGS='local' on Windows
123
+ graphsearch-ingest data/example_docs # re-ingest: embeddings are created at ingest time
124
+ graphsearch
125
+ ```
126
+
127
+ The default model (`all-MiniLM-L6-v2`, ~80 MB) downloads on first use.
128
+
129
+ ### Generated answers with citations
130
+
131
+ Point the LLM stage at OpenAI or Anthropic and answers become generated
132
+ text that cites its sources — `[1]`, `[2]`, … map 1:1 to the `sources`
133
+ list returned alongside the answer:
134
+
135
+ ```bash
136
+ export GRAPHSEARCH_LLM=anthropic # or openai
137
+ export ANTHROPIC_API_KEY=sk-ant-...
138
+ pip install -e ".[anthropic]"
139
+ graphsearch
140
+ ```
141
+
142
+ | Setting | Options | Default |
143
+ |---|---|---|
144
+ | `GRAPHSEARCH_EMBEDDINGS` | `hash` (offline), `local` (offline, semantic), `openai` | `hash` |
145
+ | `GRAPHSEARCH_LLM` | `extractive` (offline), `openai`, `anthropic` | `extractive` |
146
+
147
+ > **Note:** documents are embedded at ingest time, so re-ingest after switching
148
+ > embedding backends.
149
+
150
+ ## GraphQL API
151
+
152
+ ### Queries
153
+
154
+ ```graphql
155
+ answer(question: String!, topK: Int): Answer! # RAG: retrieve + generate
156
+ search(query: String!, topK: Int): [Chunk!]! # raw semantic search
157
+ documents(limit: Int = 20, offset: Int = 0): [Document!]!
158
+ document(id: ID!): Document
159
+ ```
160
+
161
+ ### Mutations
162
+
163
+ ```graphql
164
+ uploadDocument(content: String!, title: String, source: String): Document!
165
+ deleteDocument(id: ID!): Boolean!
166
+ ```
167
+
168
+ ### Example: ingest and ask in one session
169
+
170
+ ```graphql
171
+ mutation {
172
+ uploadDocument(
173
+ content: "Support hours are 9am-5pm PST, Monday through Friday."
174
+ title: "support-hours"
175
+ ) { id chunkCount }
176
+ }
177
+
178
+ query {
179
+ answer(question: "When is support available?") {
180
+ text
181
+ sources { documentId score }
182
+ }
183
+ }
184
+ ```
185
+
186
+ ## Architecture
187
+
188
+ ```
189
+ Client ── GraphQL (FastAPI + Strawberry) ── RagService
190
+ ├─ chunking (paragraph-aware splitter)
191
+ ├─ Embedder (hash | sentence-transformers | OpenAI)
192
+ ├─ VectorStore (SQLite-backed, in-memory cosine search)
193
+ ├─ Database (SQLite: documents, chunks, vectors)
194
+ └─ LLM (extractive | OpenAI | Anthropic)
195
+ ```
196
+
197
+ Every pipeline stage is an abstract interface (`Embedder`, `VectorStore`, `LLM`),
198
+ so new backends are drop-in additions — see [Contributing](#contributing).
199
+
200
+ ## Development
201
+
202
+ ```bash
203
+ pip install -e ".[dev]"
204
+ ruff check . # lint
205
+ pytest -v # tests run fully offline
206
+ ```
207
+
208
+ ## Roadmap / good first issues
209
+
210
+ - [ ] Additional vector store backends: Qdrant, Weaviate, Redis/Valkey, pgvector, FAISS
211
+ - [ ] Additional embedding backends: Cohere, Voyage
212
+ - [ ] Streaming answers via GraphQL subscriptions
213
+ - [ ] Metadata filters on `search` (tags, date ranges) and hybrid keyword+vector search
214
+ - [ ] Auth (API keys / JWT) and rate limiting
215
+ - [ ] Query/embedding caching
216
+ - [ ] Auto-generated TypeScript client (GraphQL Code Generator)
217
+ - [ ] Evaluation harness with known Q&A pairs; Prometheus metrics
218
+ - [ ] Advanced RAG: query rewriting, multi-hop retrieval, citation spans
219
+
220
+ ## Contributing
221
+
222
+ Contributions are welcome! See [CONTRIBUTING.md](CONTRIBUTING.md) for setup,
223
+ style, and PR guidelines, and the issue tracker for `good first issue` labels.
224
+
225
+ ## License
226
+
227
+ [MIT](LICENSE)
@@ -0,0 +1,193 @@
1
+ # GraphSearch
2
+
3
+ [![CI](https://github.com/mohithgowdak/graphsearch/actions/workflows/ci.yml/badge.svg)](https://github.com/mohithgowdak/graphsearch/actions/workflows/ci.yml)
4
+ [![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](LICENSE)
5
+ [![Python 3.10+](https://img.shields.io/badge/python-3.10%2B-blue.svg)](pyproject.toml)
6
+
7
+ **A GraphQL API server for Retrieval-Augmented Generation (RAG) over your documents.**
8
+
9
+ Upload documents through a GraphQL mutation, then ask questions through a GraphQL query.
10
+ GraphSearch chunks and embeds your documents, retrieves the most relevant passages with
11
+ vector similarity search, and feeds them to an LLM to generate a grounded answer — all
12
+ behind a single, typed, introspectable GraphQL endpoint.
13
+
14
+ ```graphql
15
+ query {
16
+ answer(question: "How do I reset my password?") {
17
+ text # cites sources as [1], [2], ...
18
+ sources { documentTitle text score }
19
+ }
20
+ }
21
+ ```
22
+
23
+ ![GraphSearch demo: asking "How do I get my money back?" in GraphiQL and getting the returns policy with ranked sources](docs/demo.gif)
24
+
25
+ *Semantic retrieval with the offline `local` backend: "How do I get my money back?" finds the returns policy — no shared keywords needed.*
26
+
27
+ ## Why GraphQL for RAG?
28
+
29
+ - **One endpoint, typed schema** — clients ask for exactly the fields they need
30
+ (answer text, source chunks, scores) instead of juggling REST routes.
31
+ - **Introspection & tooling for free** — GraphiQL, codegen'd TypeScript clients,
32
+ and schema docs come with the ecosystem.
33
+ - **Composable** — `answer`, `search`, and document management live in one schema
34
+ and can be combined in a single request.
35
+
36
+ ## Quickstart
37
+
38
+ ### Install
39
+
40
+ ```bash
41
+ pip install graphsearch-rag # from PyPI (imports as `graphsearch`)
42
+ # or run the prebuilt image:
43
+ docker run -p 8000:8000 ghcr.io/mohithgowdak/graphsearch:latest
44
+ # or from source:
45
+ git clone https://github.com/mohithgowdak/graphsearch && cd graphsearch
46
+ pip install -e ".[dev]"
47
+ ```
48
+
49
+ ### Run locally (no API keys needed)
50
+
51
+ The default configuration is fully offline: a hashing-trick embedder plus an
52
+ extractive answer mode. It exercises the entire pipeline and is perfect for
53
+ kicking the tires, tests, and CI.
54
+
55
+ ```bash
56
+ # Ingest some documents (bundled examples shown; any .md/.txt works)
57
+ graphsearch-ingest data/example_docs
58
+
59
+ # Start the server
60
+ graphsearch
61
+ ```
62
+
63
+ Open **http://localhost:8000/graphql** for the GraphiQL playground and try:
64
+
65
+ ```graphql
66
+ query {
67
+ answer(question: "What is the return policy?") {
68
+ text
69
+ sources { text score }
70
+ }
71
+ }
72
+ ```
73
+
74
+ ### Run with Docker
75
+
76
+ ```bash
77
+ docker compose up --build
78
+ ```
79
+
80
+ ### Semantic search without any API key
81
+
82
+ The `local` backend runs [sentence-transformers](https://www.sbert.net/)
83
+ on your CPU — real semantic retrieval ("How do I get my money back?" finds
84
+ the refund policy), still zero keys and zero external services:
85
+
86
+ ```bash
87
+ pip install -e ".[local]"
88
+ export GRAPHSEARCH_EMBEDDINGS=local # $env:GRAPHSEARCH_EMBEDDINGS='local' on Windows
89
+ graphsearch-ingest data/example_docs # re-ingest: embeddings are created at ingest time
90
+ graphsearch
91
+ ```
92
+
93
+ The default model (`all-MiniLM-L6-v2`, ~80 MB) downloads on first use.
94
+
95
+ ### Generated answers with citations
96
+
97
+ Point the LLM stage at OpenAI or Anthropic and answers become generated
98
+ text that cites its sources — `[1]`, `[2]`, … map 1:1 to the `sources`
99
+ list returned alongside the answer:
100
+
101
+ ```bash
102
+ export GRAPHSEARCH_LLM=anthropic # or openai
103
+ export ANTHROPIC_API_KEY=sk-ant-...
104
+ pip install -e ".[anthropic]"
105
+ graphsearch
106
+ ```
107
+
108
+ | Setting | Options | Default |
109
+ |---|---|---|
110
+ | `GRAPHSEARCH_EMBEDDINGS` | `hash` (offline), `local` (offline, semantic), `openai` | `hash` |
111
+ | `GRAPHSEARCH_LLM` | `extractive` (offline), `openai`, `anthropic` | `extractive` |
112
+
113
+ > **Note:** documents are embedded at ingest time, so re-ingest after switching
114
+ > embedding backends.
115
+
116
+ ## GraphQL API
117
+
118
+ ### Queries
119
+
120
+ ```graphql
121
+ answer(question: String!, topK: Int): Answer! # RAG: retrieve + generate
122
+ search(query: String!, topK: Int): [Chunk!]! # raw semantic search
123
+ documents(limit: Int = 20, offset: Int = 0): [Document!]!
124
+ document(id: ID!): Document
125
+ ```
126
+
127
+ ### Mutations
128
+
129
+ ```graphql
130
+ uploadDocument(content: String!, title: String, source: String): Document!
131
+ deleteDocument(id: ID!): Boolean!
132
+ ```
133
+
134
+ ### Example: ingest and ask in one session
135
+
136
+ ```graphql
137
+ mutation {
138
+ uploadDocument(
139
+ content: "Support hours are 9am-5pm PST, Monday through Friday."
140
+ title: "support-hours"
141
+ ) { id chunkCount }
142
+ }
143
+
144
+ query {
145
+ answer(question: "When is support available?") {
146
+ text
147
+ sources { documentId score }
148
+ }
149
+ }
150
+ ```
151
+
152
+ ## Architecture
153
+
154
+ ```
155
+ Client ── GraphQL (FastAPI + Strawberry) ── RagService
156
+ ├─ chunking (paragraph-aware splitter)
157
+ ├─ Embedder (hash | sentence-transformers | OpenAI)
158
+ ├─ VectorStore (SQLite-backed, in-memory cosine search)
159
+ ├─ Database (SQLite: documents, chunks, vectors)
160
+ └─ LLM (extractive | OpenAI | Anthropic)
161
+ ```
162
+
163
+ Every pipeline stage is an abstract interface (`Embedder`, `VectorStore`, `LLM`),
164
+ so new backends are drop-in additions — see [Contributing](#contributing).
165
+
166
+ ## Development
167
+
168
+ ```bash
169
+ pip install -e ".[dev]"
170
+ ruff check . # lint
171
+ pytest -v # tests run fully offline
172
+ ```
173
+
174
+ ## Roadmap / good first issues
175
+
176
+ - [ ] Additional vector store backends: Qdrant, Weaviate, Redis/Valkey, pgvector, FAISS
177
+ - [ ] Additional embedding backends: Cohere, Voyage
178
+ - [ ] Streaming answers via GraphQL subscriptions
179
+ - [ ] Metadata filters on `search` (tags, date ranges) and hybrid keyword+vector search
180
+ - [ ] Auth (API keys / JWT) and rate limiting
181
+ - [ ] Query/embedding caching
182
+ - [ ] Auto-generated TypeScript client (GraphQL Code Generator)
183
+ - [ ] Evaluation harness with known Q&A pairs; Prometheus metrics
184
+ - [ ] Advanced RAG: query rewriting, multi-hop retrieval, citation spans
185
+
186
+ ## Contributing
187
+
188
+ Contributions are welcome! See [CONTRIBUTING.md](CONTRIBUTING.md) for setup,
189
+ style, and PR guidelines, and the issue tracker for `good first issue` labels.
190
+
191
+ ## License
192
+
193
+ [MIT](LICENSE)
@@ -0,0 +1,3 @@
1
+ """GraphSearch: a GraphQL API for Retrieval-Augmented Generation over your documents."""
2
+
3
+ __version__ = "0.2.1"
@@ -0,0 +1,46 @@
1
+ """Split documents into overlapping chunks suitable for embedding."""
2
+
3
+ from __future__ import annotations
4
+
5
+
6
+ def chunk_text(text: str, chunk_size: int = 800, overlap: int = 100) -> list[str]:
7
+ """Split ``text`` into chunks of at most ``chunk_size`` characters.
8
+
9
+ Prefers paragraph boundaries, then falls back to a sliding window with
10
+ ``overlap`` characters of context carried between adjacent chunks.
11
+ """
12
+ if chunk_size <= 0:
13
+ raise ValueError("chunk_size must be positive")
14
+ if overlap >= chunk_size:
15
+ raise ValueError("overlap must be smaller than chunk_size")
16
+
17
+ text = text.strip()
18
+ if not text:
19
+ return []
20
+
21
+ paragraphs = [p.strip() for p in text.split("\n\n") if p.strip()]
22
+
23
+ chunks: list[str] = []
24
+ current = ""
25
+ for para in paragraphs:
26
+ candidate = f"{current}\n\n{para}" if current else para
27
+ if len(candidate) <= chunk_size:
28
+ current = candidate
29
+ continue
30
+ if current:
31
+ chunks.append(current)
32
+ if len(para) <= chunk_size:
33
+ current = para
34
+ else:
35
+ # Paragraph alone exceeds chunk_size: sliding window.
36
+ step = chunk_size - overlap
37
+ for start in range(0, len(para), step):
38
+ piece = para[start : start + chunk_size]
39
+ if len(piece) < overlap and chunks:
40
+ # Tiny tail already covered by the previous window's overlap.
41
+ break
42
+ chunks.append(piece)
43
+ current = ""
44
+ if current:
45
+ chunks.append(current)
46
+ return chunks
@@ -0,0 +1,54 @@
1
+ """Application configuration, loaded from environment variables / .env file.
2
+
3
+ All GraphSearch-specific settings use the ``GRAPHSEARCH_`` prefix
4
+ (e.g. ``GRAPHSEARCH_LLM=openai``). Provider API keys use their
5
+ conventional names: ``OPENAI_API_KEY`` and ``ANTHROPIC_API_KEY``.
6
+ """
7
+
8
+ from __future__ import annotations
9
+
10
+ from functools import lru_cache
11
+
12
+ from pydantic import Field
13
+ from pydantic_settings import BaseSettings, SettingsConfigDict
14
+
15
+
16
+ class Settings(BaseSettings):
17
+ model_config = SettingsConfigDict(
18
+ env_prefix="GRAPHSEARCH_", env_file=".env", extra="ignore"
19
+ )
20
+
21
+ # Storage
22
+ database_path: str = "graphsearch.db"
23
+
24
+ # Chunking
25
+ chunk_size: int = 800
26
+ chunk_overlap: int = 100
27
+
28
+ # Retrieval
29
+ default_top_k: int = 4
30
+
31
+ # Backends: which implementation to use for each pipeline stage.
32
+ # embeddings: "hash" (offline, no key needed), "local" (sentence-transformers), or "openai"
33
+ # llm: "extractive" (offline), "openai", or "anthropic"
34
+ embeddings: str = "hash"
35
+ llm: str = "extractive"
36
+
37
+ # Model names (only used when the matching backend is selected)
38
+ local_embedding_model: str = "sentence-transformers/all-MiniLM-L6-v2"
39
+ openai_embedding_model: str = "text-embedding-3-small"
40
+ openai_model: str = "gpt-4o-mini"
41
+ anthropic_model: str = "claude-sonnet-5"
42
+
43
+ # Provider keys (no prefix — conventional env var names)
44
+ openai_api_key: str = Field(default="", alias="OPENAI_API_KEY")
45
+ anthropic_api_key: str = Field(default="", alias="ANTHROPIC_API_KEY")
46
+
47
+ # Server
48
+ host: str = "0.0.0.0"
49
+ port: int = 8000
50
+
51
+
52
+ @lru_cache
53
+ def get_settings() -> Settings:
54
+ return Settings()
@@ -0,0 +1,139 @@
1
+ """SQLite persistence for documents, chunks, and their embeddings.
2
+
3
+ A single SQLite file stores everything, which keeps the default deployment
4
+ zero-dependency. Embeddings are stored as float32 blobs on the ``chunks``
5
+ table and searched in memory with numpy (see ``vectorstore.py``).
6
+ """
7
+
8
+ from __future__ import annotations
9
+
10
+ import sqlite3
11
+ import threading
12
+ import uuid
13
+ from dataclasses import dataclass
14
+ from datetime import datetime, timezone
15
+
16
+ import numpy as np
17
+
18
+ _SCHEMA = """
19
+ CREATE TABLE IF NOT EXISTS documents (
20
+ id TEXT PRIMARY KEY,
21
+ title TEXT,
22
+ source TEXT,
23
+ content TEXT NOT NULL,
24
+ created_at TEXT NOT NULL
25
+ );
26
+ CREATE TABLE IF NOT EXISTS chunks (
27
+ id TEXT PRIMARY KEY,
28
+ document_id TEXT NOT NULL REFERENCES documents(id) ON DELETE CASCADE,
29
+ idx INTEGER NOT NULL,
30
+ text TEXT NOT NULL,
31
+ embedding BLOB NOT NULL
32
+ );
33
+ CREATE INDEX IF NOT EXISTS idx_chunks_document ON chunks(document_id);
34
+ """
35
+
36
+
37
+ @dataclass
38
+ class DocumentRecord:
39
+ id: str
40
+ title: str | None
41
+ source: str | None
42
+ content: str
43
+ created_at: str
44
+ chunk_count: int
45
+
46
+
47
+ @dataclass
48
+ class ChunkRecord:
49
+ id: str
50
+ document_id: str
51
+ idx: int
52
+ text: str
53
+ document_title: str | None = None
54
+
55
+
56
+ class Database:
57
+ """Thread-safe wrapper around a SQLite connection."""
58
+
59
+ def __init__(self, path: str) -> None:
60
+ self._conn = sqlite3.connect(path, check_same_thread=False)
61
+ self._conn.execute("PRAGMA foreign_keys = ON")
62
+ self._conn.executescript(_SCHEMA)
63
+ self._lock = threading.Lock()
64
+
65
+ def close(self) -> None:
66
+ self._conn.close()
67
+
68
+ # -- documents ---------------------------------------------------------
69
+
70
+ def insert_document(
71
+ self,
72
+ content: str,
73
+ title: str | None,
74
+ source: str | None,
75
+ chunks: list[str],
76
+ embeddings: np.ndarray,
77
+ ) -> DocumentRecord:
78
+ if len(chunks) != len(embeddings):
79
+ raise ValueError("chunks and embeddings must have the same length")
80
+ doc_id = uuid.uuid4().hex
81
+ created_at = datetime.now(timezone.utc).isoformat()
82
+ with self._lock, self._conn:
83
+ self._conn.execute(
84
+ "INSERT INTO documents (id, title, source, content, created_at)"
85
+ " VALUES (?, ?, ?, ?, ?)",
86
+ (doc_id, title, source, content, created_at),
87
+ )
88
+ self._conn.executemany(
89
+ "INSERT INTO chunks (id, document_id, idx, text, embedding)"
90
+ " VALUES (?, ?, ?, ?, ?)",
91
+ [
92
+ (
93
+ uuid.uuid4().hex,
94
+ doc_id,
95
+ i,
96
+ chunk,
97
+ np.asarray(emb, dtype=np.float32).tobytes(),
98
+ )
99
+ for i, (chunk, emb) in enumerate(zip(chunks, embeddings, strict=True))
100
+ ],
101
+ )
102
+ return DocumentRecord(doc_id, title, source, content, created_at, len(chunks))
103
+
104
+ def get_document(self, doc_id: str) -> DocumentRecord | None:
105
+ row = self._conn.execute(
106
+ "SELECT d.id, d.title, d.source, d.content, d.created_at,"
107
+ " (SELECT COUNT(*) FROM chunks c WHERE c.document_id = d.id)"
108
+ " FROM documents d WHERE d.id = ?",
109
+ (doc_id,),
110
+ ).fetchone()
111
+ return DocumentRecord(*row) if row else None
112
+
113
+ def list_documents(self, limit: int = 20, offset: int = 0) -> list[DocumentRecord]:
114
+ rows = self._conn.execute(
115
+ "SELECT d.id, d.title, d.source, d.content, d.created_at,"
116
+ " (SELECT COUNT(*) FROM chunks c WHERE c.document_id = d.id)"
117
+ " FROM documents d ORDER BY d.created_at DESC LIMIT ? OFFSET ?",
118
+ (limit, offset),
119
+ ).fetchall()
120
+ return [DocumentRecord(*row) for row in rows]
121
+
122
+ def delete_document(self, doc_id: str) -> bool:
123
+ with self._lock, self._conn:
124
+ cursor = self._conn.execute("DELETE FROM documents WHERE id = ?", (doc_id,))
125
+ return cursor.rowcount > 0
126
+
127
+ # -- chunks / embeddings -------------------------------------------------
128
+
129
+ def all_embeddings(self) -> tuple[list[ChunkRecord], np.ndarray]:
130
+ """Return every chunk with its embedding matrix (rows align with chunks)."""
131
+ rows = self._conn.execute(
132
+ "SELECT c.id, c.document_id, c.idx, c.text, d.title, c.embedding"
133
+ " FROM chunks c JOIN documents d ON d.id = c.document_id"
134
+ ).fetchall()
135
+ if not rows:
136
+ return [], np.empty((0, 0), dtype=np.float32)
137
+ records = [ChunkRecord(r[0], r[1], r[2], r[3], r[4]) for r in rows]
138
+ matrix = np.vstack([np.frombuffer(r[5], dtype=np.float32) for r in rows])
139
+ return records, matrix