Delete comments that restate the line below them
Sixty-three comments said what the next statement already said: # Connect to LanceDB above connect_lancedb, # Path object above isinstance(source, Path), # Get page numbers from provenance above the prov loop, # Clear and populate results above list_view.clear(). They cost a read and carry nothing. The line is whether a comment restates one statement or labels a phase. Phase labels stay: the migrations keep # Create staging table with new schema and # Copy from staging to final table in batches, each heading ten lines of a long procedure. So do comments carrying a fact the code cannot: the merge_insert update-only note on document_meta, why the poller builds sources eagerly, why create_document_from_source returns a list for directories, that indexes need training data, the field-group markers in the config models, and the file:// URL-encoding note in create_document_from_source. capabilities/ is untouched. Its docstrings sit next to prompt surface, and changing them needs an eval to back it. The cassette-recording docs were wrong three ways. They named tests/test_qa.py::test_qa_anthropic, which no longer exists; they targeted whole modules, so a rewrite would re-record cassettes for services the recorder is not running; and they used COHERE_API_KEY where the SDK reads CO_API_KEY. docs/development.md now names exact tests with -n0, and the keyed example is test_cohere_reranker, which owns the one cassette recording api.cohere.com.
This commit is contained in:
parent
476f8d07a0
commit
4967765878
22 changed files with 9 additions and 65 deletions
|
|
@ -109,8 +109,15 @@ uv run ty check
|
||||||
|
|
||||||
Tests automatically set mock API keys for providers that require them during client initialization. When running with VCR playback, these mock keys are sufficient since no real API calls are made.
|
Tests automatically set mock API keys for providers that require them during client initialization. When running with VCR playback, these mock keys are sufficient since no real API calls are made.
|
||||||
|
|
||||||
When recording new cassettes, set real API keys via environment variables:
|
Recording reaches the real service, so the recording command needs network
|
||||||
|
access and the keys that service reads. Name the exact test and pass `-n0`:
|
||||||
|
a module-wide `--record-mode=rewrite` re-records every cassette in it,
|
||||||
|
including ones whose service you do not have running.
|
||||||
|
|
||||||
```bash
|
```bash
|
||||||
ANTHROPIC_API_KEY=sk-ant-... uv run pytest tests/test_qa.py::test_qa_anthropic --record-mode=rewrite
|
# Ollama-backed cassettes need no key, only a running Ollama
|
||||||
|
uv run pytest tests/test_embedder.py::test_ollama_embedder -n0 --record-mode=rewrite
|
||||||
|
|
||||||
|
# A keyed provider reads its own variable. Cohere's SDK reads CO_API_KEY
|
||||||
|
CO_API_KEY=... uv run pytest tests/test_reranker.py::test_cohere_reranker -n0 --record-mode=rewrite
|
||||||
```
|
```
|
||||||
|
|
|
||||||
|
|
@ -55,7 +55,6 @@ class HaikuRAGApp:
|
||||||
)
|
)
|
||||||
return
|
return
|
||||||
|
|
||||||
# Create the database
|
|
||||||
async with HaikuRAG(db_path=self.db_path, config=self.config, create=True):
|
async with HaikuRAG(db_path=self.db_path, config=self.config, create=True):
|
||||||
pass
|
pass
|
||||||
self.console.print(
|
self.console.print(
|
||||||
|
|
@ -161,7 +160,6 @@ class HaikuRAGApp:
|
||||||
f"{tables['chunks'].num_versions}"
|
f"{tables['chunks'].num_versions}"
|
||||||
)
|
)
|
||||||
|
|
||||||
# Migration status
|
|
||||||
self.console.rule()
|
self.console.rule()
|
||||||
if info.pending_migrations:
|
if info.pending_migrations:
|
||||||
self.console.print(
|
self.console.print(
|
||||||
|
|
@ -747,7 +745,6 @@ class HaikuRAGApp:
|
||||||
)
|
)
|
||||||
return
|
return
|
||||||
|
|
||||||
# Check if index already exists
|
|
||||||
indices = await client.store.chunks_table.list_indices()
|
indices = await client.store.chunks_table.list_indices()
|
||||||
has_vector_index = any("vector" in str(idx).lower() for idx in indices)
|
has_vector_index = any("vector" in str(idx).lower() for idx in indices)
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -108,11 +108,9 @@ class DocumentFilterModal(ModalScreen):
|
||||||
"""Load all documents from the client."""
|
"""Load all documents from the client."""
|
||||||
docs = await self.client.list_documents()
|
docs = await self.client.list_documents()
|
||||||
|
|
||||||
# Remove loading indicator
|
|
||||||
loading = self.query_one("#loading-indicator", Static)
|
loading = self.query_one("#loading-indicator", Static)
|
||||||
loading.remove()
|
loading.remove()
|
||||||
|
|
||||||
# Add checkboxes for all documents
|
|
||||||
filter_list = self.query_one("#filter-list", VerticalScroll)
|
filter_list = self.query_one("#filter-list", VerticalScroll)
|
||||||
for doc in docs:
|
for doc in docs:
|
||||||
display_name = doc.title or doc.uri or str(doc.id)
|
display_name = doc.title or doc.uri or str(doc.id)
|
||||||
|
|
|
||||||
|
|
@ -127,13 +127,10 @@ class DoclingLocalChunker(DocumentChunker):
|
||||||
meta = cast("DocMeta | None", chunk.meta)
|
meta = cast("DocMeta | None", chunk.meta)
|
||||||
if meta and meta.doc_items:
|
if meta and meta.doc_items:
|
||||||
for doc_item in meta.doc_items:
|
for doc_item in meta.doc_items:
|
||||||
# Get JSON pointer reference
|
|
||||||
if doc_item.self_ref:
|
if doc_item.self_ref:
|
||||||
doc_item_refs.append(doc_item.self_ref)
|
doc_item_refs.append(doc_item.self_ref)
|
||||||
# Get label
|
|
||||||
if doc_item.label:
|
if doc_item.label:
|
||||||
labels.append(doc_item.label)
|
labels.append(doc_item.label)
|
||||||
# Get page numbers from provenance
|
|
||||||
if doc_item.prov:
|
if doc_item.prov:
|
||||||
for prov in doc_item.prov:
|
for prov in doc_item.prov:
|
||||||
if (
|
if (
|
||||||
|
|
@ -142,7 +139,6 @@ class DoclingLocalChunker(DocumentChunker):
|
||||||
):
|
):
|
||||||
page_numbers.append(prov.page_no)
|
page_numbers.append(prov.page_no)
|
||||||
|
|
||||||
# Get headings from chunk metadata
|
|
||||||
if meta and meta.headings:
|
if meta and meta.headings:
|
||||||
headings = list(meta.headings)
|
headings = list(meta.headings)
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -179,10 +179,8 @@ class DoclingServeChunker(DocumentChunker):
|
||||||
if "label" in item:
|
if "label" in item:
|
||||||
labels.append(item["label"])
|
labels.append(item["label"])
|
||||||
|
|
||||||
# Get headings directly from chunk
|
|
||||||
headings = chunk.get("headings")
|
headings = chunk.get("headings")
|
||||||
|
|
||||||
# Get page numbers directly from chunk
|
|
||||||
page_numbers = chunk.get("page_numbers", [])
|
page_numbers = chunk.get("page_numbers", [])
|
||||||
|
|
||||||
chunk_metadata = ChunkMetadata(
|
chunk_metadata = ChunkMetadata(
|
||||||
|
|
|
||||||
|
|
@ -113,7 +113,6 @@ def main(
|
||||||
loaded_config = AppConfig.model_validate(yaml_data)
|
loaded_config = AppConfig.model_validate(yaml_data)
|
||||||
set_config(loaded_config)
|
set_config(loaded_config)
|
||||||
|
|
||||||
# Configure logging for CLI context
|
|
||||||
configure_cli_logging()
|
configure_cli_logging()
|
||||||
|
|
||||||
from haiku.rag.telemetry import configure as configure_telemetry
|
from haiku.rag.telemetry import configure as configure_telemetry
|
||||||
|
|
|
||||||
|
|
@ -31,7 +31,6 @@ async def download_models(
|
||||||
- HuggingFace reranker models (cross-encoder, jina-local)
|
- HuggingFace reranker models (cross-encoder, jina-local)
|
||||||
- Ollama models
|
- Ollama models
|
||||||
"""
|
"""
|
||||||
# Docling models
|
|
||||||
try:
|
try:
|
||||||
from docling.utils.model_downloader import download_models
|
from docling.utils.model_downloader import download_models
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -115,7 +115,6 @@ async def convert(
|
||||||
)
|
)
|
||||||
return await converter.convert_file(file_path, source_uri=effective_uri)
|
return await converter.convert_file(file_path, source_uri=effective_uri)
|
||||||
|
|
||||||
# Path object - convert file directly
|
|
||||||
if isinstance(source, Path):
|
if isinstance(source, Path):
|
||||||
if not source.exists():
|
if not source.exists():
|
||||||
raise UnsupportedSourceError(f"File does not exist: {source}")
|
raise UnsupportedSourceError(f"File does not exist: {source}")
|
||||||
|
|
@ -126,7 +125,6 @@ async def convert(
|
||||||
_warn_if_descriptions_missing(config, doc, str(source))
|
_warn_if_descriptions_missing(config, doc, str(source))
|
||||||
return doc
|
return doc
|
||||||
|
|
||||||
# String - check if URL or text
|
|
||||||
parsed = urlparse(source)
|
parsed = urlparse(source)
|
||||||
|
|
||||||
if parsed.scheme in ("http", "https"):
|
if parsed.scheme in ("http", "https"):
|
||||||
|
|
@ -157,7 +155,6 @@ async def convert(
|
||||||
temp_path.unlink(missing_ok=True)
|
temp_path.unlink(missing_ok=True)
|
||||||
|
|
||||||
elif parsed.scheme == "file":
|
elif parsed.scheme == "file":
|
||||||
# file:// URI
|
|
||||||
file_path = Path(parsed.path)
|
file_path = Path(parsed.path)
|
||||||
if not file_path.exists():
|
if not file_path.exists():
|
||||||
raise UnsupportedSourceError(f"File does not exist: {file_path}")
|
raise UnsupportedSourceError(f"File does not exist: {file_path}")
|
||||||
|
|
|
||||||
|
|
@ -111,7 +111,6 @@ async def _rebuild_locked(
|
||||||
# destructive phase and are fine.
|
# destructive phase and are fine.
|
||||||
await client._await_vacuum_tasks()
|
await client._await_vacuum_tasks()
|
||||||
|
|
||||||
# Update settings to current config
|
|
||||||
settings_repo = SettingsRepository(client.store)
|
settings_repo = SettingsRepository(client.store)
|
||||||
await settings_repo.save_current_settings()
|
await settings_repo.save_current_settings()
|
||||||
|
|
||||||
|
|
@ -503,7 +502,6 @@ async def _rebuild_embed_only(
|
||||||
# rebuild discards harmlessly.
|
# rebuild discards harmlessly.
|
||||||
await _drop_staging_tables(client)
|
await _drop_staging_tables(client)
|
||||||
|
|
||||||
# Yield docs with no chunks
|
|
||||||
for doc in documents:
|
for doc in documents:
|
||||||
if doc.id and doc.id not in yielded_docs:
|
if doc.id and doc.id not in yielded_docs:
|
||||||
yield doc.id
|
yield doc.id
|
||||||
|
|
@ -622,7 +620,6 @@ async def _rebuild_rechunk(
|
||||||
)
|
)
|
||||||
embedded_chunks = await embed_chunks(chunks, embedder, client._config)
|
embedded_chunks = await embed_chunks(chunks, embedder, client._config)
|
||||||
|
|
||||||
# Prepare chunks with document_id and order
|
|
||||||
for order, chunk in enumerate(embedded_chunks):
|
for order, chunk in enumerate(embedded_chunks):
|
||||||
chunk.document_id = doc.id
|
chunk.document_id = doc.id
|
||||||
chunk.order = order
|
chunk.order = order
|
||||||
|
|
@ -636,13 +633,11 @@ async def _rebuild_rechunk(
|
||||||
# consistent with the rebuild already being non-atomic.
|
# consistent with the rebuild already being non-atomic.
|
||||||
yield doc.id
|
yield doc.id
|
||||||
|
|
||||||
# Flush batch when size reached
|
|
||||||
if len(pending_docs) >= _REBUILD_BATCH_SIZE:
|
if len(pending_docs) >= _REBUILD_BATCH_SIZE:
|
||||||
await _flush_rebuild_batch(client, pending_docs, pending_chunks)
|
await _flush_rebuild_batch(client, pending_docs, pending_chunks)
|
||||||
pending_chunks = []
|
pending_chunks = []
|
||||||
pending_docs = []
|
pending_docs = []
|
||||||
|
|
||||||
# Flush remaining
|
|
||||||
if pending_docs:
|
if pending_docs:
|
||||||
await _flush_rebuild_batch(client, pending_docs, pending_chunks)
|
await _flush_rebuild_batch(client, pending_docs, pending_chunks)
|
||||||
|
|
||||||
|
|
@ -858,7 +853,6 @@ async def _rebuild_full(
|
||||||
|
|
||||||
doc.set_docling(docling_document)
|
doc.set_docling(docling_document)
|
||||||
|
|
||||||
# Prepare chunks with document_id and order
|
|
||||||
for order, chunk in enumerate(embedded_chunks):
|
for order, chunk in enumerate(embedded_chunks):
|
||||||
chunk.document_id = doc.id
|
chunk.document_id = doc.id
|
||||||
chunk.order = order
|
chunk.order = order
|
||||||
|
|
@ -867,12 +861,10 @@ async def _rebuild_full(
|
||||||
pending_docs.append(doc)
|
pending_docs.append(doc)
|
||||||
yield doc.id
|
yield doc.id
|
||||||
|
|
||||||
# Flush batch when size reached
|
|
||||||
if len(pending_docs) >= _REBUILD_BATCH_SIZE:
|
if len(pending_docs) >= _REBUILD_BATCH_SIZE:
|
||||||
await _flush_rebuild_batch(client, pending_docs, pending_chunks)
|
await _flush_rebuild_batch(client, pending_docs, pending_chunks)
|
||||||
pending_chunks = []
|
pending_chunks = []
|
||||||
pending_docs = []
|
pending_docs = []
|
||||||
|
|
||||||
# Flush remaining
|
|
||||||
if pending_docs:
|
if pending_docs:
|
||||||
await _flush_rebuild_batch(client, pending_docs, pending_chunks)
|
await _flush_rebuild_batch(client, pending_docs, pending_chunks)
|
||||||
|
|
|
||||||
|
|
@ -379,7 +379,6 @@ async def visualize_chunk(
|
||||||
return []
|
return []
|
||||||
page_images = pages_doc.get_page_images(list(boxes_by_page.keys()))
|
page_images = pages_doc.get_page_images(list(boxes_by_page.keys()))
|
||||||
|
|
||||||
# Render each page with its bounding boxes
|
|
||||||
images = []
|
images = []
|
||||||
for page_no in sorted(boxes_by_page.keys()):
|
for page_no in sorted(boxes_by_page.keys()):
|
||||||
if page_no not in page_images:
|
if page_no not in page_images:
|
||||||
|
|
|
||||||
|
|
@ -466,7 +466,6 @@ def expand_with_items(
|
||||||
merged = _merge_ranges(ranges)
|
merged = _merge_ranges(ranges)
|
||||||
constituent_range = {id(result): (lo, hi) for lo, hi, result in ranges}
|
constituent_range = {id(result): (lo, hi) for lo, hi, result in ranges}
|
||||||
|
|
||||||
# Build results from the window items
|
|
||||||
pos_to_item = {item.position: item for item in window_items}
|
pos_to_item = {item.position: item for item in window_items}
|
||||||
final_results: list[SearchResult] = []
|
final_results: list[SearchResult] = []
|
||||||
for range_start, range_end, group in merged:
|
for range_start, range_end, group in merged:
|
||||||
|
|
|
||||||
|
|
@ -93,7 +93,6 @@ class InspectorApp(App):
|
||||||
await client.__aenter__()
|
await client.__aenter__()
|
||||||
self.client = client
|
self.client = client
|
||||||
|
|
||||||
# Load initial documents
|
|
||||||
doc_list = self.query_one(DocumentList)
|
doc_list = self.query_one(DocumentList)
|
||||||
await doc_list.load_documents(self.client)
|
await doc_list.load_documents(self.client)
|
||||||
|
|
||||||
|
|
@ -150,7 +149,6 @@ class InspectorApp(App):
|
||||||
doc_list = self.query_one(DocumentList)
|
doc_list = self.query_one(DocumentList)
|
||||||
chunk_list = self.query_one(ChunkList)
|
chunk_list = self.query_one(ChunkList)
|
||||||
|
|
||||||
# Find and select the document
|
|
||||||
for idx, d in enumerate(doc_list.documents):
|
for idx, d in enumerate(doc_list.documents):
|
||||||
if d.id == chunk.document_id:
|
if d.id == chunk.document_id:
|
||||||
doc_list.list_view.index = idx
|
doc_list.list_view.index = idx
|
||||||
|
|
@ -175,7 +173,6 @@ class InspectorApp(App):
|
||||||
if not self.client:
|
if not self.client:
|
||||||
return
|
return
|
||||||
|
|
||||||
# Show document details
|
|
||||||
detail_view = self.query_one(DetailView)
|
detail_view = self.query_one(DetailView)
|
||||||
await detail_view.show_document(message.document)
|
await detail_view.show_document(message.document)
|
||||||
|
|
||||||
|
|
@ -192,7 +189,6 @@ class InspectorApp(App):
|
||||||
Args:
|
Args:
|
||||||
message: Message containing selected chunk
|
message: Message containing selected chunk
|
||||||
"""
|
"""
|
||||||
# Show chunk details
|
|
||||||
detail_view = self.query_one(DetailView)
|
detail_view = self.query_one(DetailView)
|
||||||
await detail_view.show_chunk(message.chunk)
|
await detail_view.show_chunk(message.chunk)
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -69,7 +69,6 @@ class InfoModal(ModalScreen):
|
||||||
|
|
||||||
lines: list[str] = []
|
lines: list[str] = []
|
||||||
|
|
||||||
# Path
|
|
||||||
lines.append(f"[bold $accent]path[/bold $accent]: {self.db_path}")
|
lines.append(f"[bold $accent]path[/bold $accent]: {self.db_path}")
|
||||||
|
|
||||||
is_local = self.client.store._connection_mode == ConnectionMode.LOCAL
|
is_local = self.client.store._connection_mode == ConnectionMode.LOCAL
|
||||||
|
|
@ -88,10 +87,8 @@ class InfoModal(ModalScreen):
|
||||||
self._content_widget.update("\n".join(lines))
|
self._content_widget.update("\n".join(lines))
|
||||||
return
|
return
|
||||||
|
|
||||||
# Get versions
|
|
||||||
versions = get_package_versions()
|
versions = get_package_versions()
|
||||||
|
|
||||||
# Read settings
|
|
||||||
stored_version = "unknown"
|
stored_version = "unknown"
|
||||||
embed_provider: str | None = None
|
embed_provider: str | None = None
|
||||||
embed_model: str | None = None
|
embed_model: str | None = None
|
||||||
|
|
@ -130,7 +127,6 @@ class InfoModal(ModalScreen):
|
||||||
doc_versions = stats["documents"].get("num_versions", 0)
|
doc_versions = stats["documents"].get("num_versions", 0)
|
||||||
chunk_versions = stats["chunks"].get("num_versions", 0)
|
chunk_versions = stats["chunks"].get("num_versions", 0)
|
||||||
|
|
||||||
# Build output
|
|
||||||
lines.append(
|
lines.append(
|
||||||
f"[bold $accent]haiku.rag version (db)[/bold $accent]: {stored_version}"
|
f"[bold $accent]haiku.rag version (db)[/bold $accent]: {stored_version}"
|
||||||
)
|
)
|
||||||
|
|
@ -156,7 +152,6 @@ class InfoModal(ModalScreen):
|
||||||
f"[bold $accent]chunks[/bold $accent]: {num_chunks} ({format_bytes(chunk_bytes)})"
|
f"[bold $accent]chunks[/bold $accent]: {num_chunks} ({format_bytes(chunk_bytes)})"
|
||||||
)
|
)
|
||||||
|
|
||||||
# Vector index info
|
|
||||||
if has_vector_index:
|
if has_vector_index:
|
||||||
lines.append("[bold $accent]vector index[/bold $accent]: ✓ exists")
|
lines.append("[bold $accent]vector index[/bold $accent]: ✓ exists")
|
||||||
lines.append(
|
lines.append(
|
||||||
|
|
|
||||||
|
|
@ -101,10 +101,8 @@ class SearchModal(Screen):
|
||||||
status_label.update("Searching...")
|
status_label.update("Searching...")
|
||||||
|
|
||||||
try:
|
try:
|
||||||
# Perform search using client API
|
|
||||||
self.search_results = await self.client.search(query=query, limit=50)
|
self.search_results = await self.client.search(query=query, limit=50)
|
||||||
|
|
||||||
# Get chunks for the results
|
|
||||||
self.chunks = []
|
self.chunks = []
|
||||||
for result in self.search_results:
|
for result in self.search_results:
|
||||||
if result.chunk_id:
|
if result.chunk_id:
|
||||||
|
|
@ -112,7 +110,6 @@ class SearchModal(Screen):
|
||||||
if chunk:
|
if chunk:
|
||||||
self.chunks.append(chunk)
|
self.chunks.append(chunk)
|
||||||
|
|
||||||
# Clear and populate results
|
|
||||||
await list_view.clear()
|
await list_view.clear()
|
||||||
for result in self.search_results:
|
for result in self.search_results:
|
||||||
first_line = result.content.split("\n")[0][:60]
|
first_line = result.content.split("\n")[0][:60]
|
||||||
|
|
@ -125,7 +122,6 @@ class SearchModal(Screen):
|
||||||
item = ListItem(Static(f"[{score_str}]{page_info} {first_line}"))
|
item = ListItem(Static(f"[{score_str}]{page_info} {first_line}"))
|
||||||
await list_view.append(item)
|
await list_view.append(item)
|
||||||
|
|
||||||
# Update status
|
|
||||||
status_label.update(f"Found {len(self.chunks)} results")
|
status_label.update(f"Found {len(self.chunks)} results")
|
||||||
|
|
||||||
# Select first result, show in detail view, and focus list
|
# Select first result, show in detail view, and focus list
|
||||||
|
|
|
||||||
|
|
@ -47,7 +47,6 @@ def configure_cli_logging(level: int = logging.INFO) -> logging.Logger:
|
||||||
logging.getLogger(noisy).setLevel(logging.ERROR)
|
logging.getLogger(noisy).setLevel(logging.ERROR)
|
||||||
logging.getLogger(noisy).propagate = False
|
logging.getLogger(noisy).propagate = False
|
||||||
|
|
||||||
# Configure and return our app logger
|
|
||||||
logger = get_logger()
|
logger = get_logger()
|
||||||
logger.setLevel(level)
|
logger.setLevel(level)
|
||||||
logger.propagate = False
|
logger.propagate = False
|
||||||
|
|
|
||||||
|
|
@ -52,13 +52,11 @@ class VLLMReranker(RerankerBase):
|
||||||
|
|
||||||
result = response.json()
|
result = response.json()
|
||||||
|
|
||||||
# Extract scores and pair with chunks
|
|
||||||
scored_chunks = []
|
scored_chunks = []
|
||||||
for item in result.get("results", []):
|
for item in result.get("results", []):
|
||||||
index = item["index"]
|
index = item["index"]
|
||||||
score = item["relevance_score"]
|
score = item["relevance_score"]
|
||||||
scored_chunks.append((chunks[index], score))
|
scored_chunks.append((chunks[index], score))
|
||||||
|
|
||||||
# Sort by score (descending) and return top_n
|
|
||||||
scored_chunks.sort(key=lambda x: x[1], reverse=True)
|
scored_chunks.sort(key=lambda x: x[1], reverse=True)
|
||||||
return scored_chunks[:top_n]
|
return scored_chunks[:top_n]
|
||||||
|
|
|
||||||
|
|
@ -25,10 +25,8 @@ class ZeroEntropyReranker(RerankerBase): # pragma: no cover
|
||||||
async def _rerank(
|
async def _rerank(
|
||||||
self, query: str, chunks: list[Chunk], top_n: int = 10
|
self, query: str, chunks: list[Chunk], top_n: int = 10
|
||||||
) -> list[tuple[Chunk, float]]:
|
) -> list[tuple[Chunk, float]]:
|
||||||
# Prepare documents for Zero Entropy API
|
|
||||||
documents = [chunk.content for chunk in chunks]
|
documents = [chunk.content for chunk in chunks]
|
||||||
|
|
||||||
# Call Zero Entropy reranking API
|
|
||||||
model_name = self._model or "zerank-1"
|
model_name = self._model or "zerank-1"
|
||||||
response = await self._client.models.rerank(
|
response = await self._client.models.rerank(
|
||||||
model=model_name,
|
model=model_name,
|
||||||
|
|
|
||||||
|
|
@ -190,7 +190,6 @@ class Store:
|
||||||
self._rebuild_lock = asyncio.Lock()
|
self._rebuild_lock = asyncio.Lock()
|
||||||
self._is_new_db = False
|
self._is_new_db = False
|
||||||
|
|
||||||
# Check if database exists (for local filesystem only)
|
|
||||||
if self._connection_mode == ConnectionMode.LOCAL:
|
if self._connection_mode == ConnectionMode.LOCAL:
|
||||||
if not db_path.exists():
|
if not db_path.exists():
|
||||||
if not create:
|
if not create:
|
||||||
|
|
@ -208,7 +207,6 @@ class Store:
|
||||||
|
|
||||||
async def _initialize(self):
|
async def _initialize(self):
|
||||||
"""Perform async initialization: connect to LanceDB, init tables, validate."""
|
"""Perform async initialization: connect to LanceDB, init tables, validate."""
|
||||||
# Connect to LanceDB
|
|
||||||
self.db: lancedb.AsyncConnection = await connect_lancedb(
|
self.db: lancedb.AsyncConnection = await connect_lancedb(
|
||||||
self._config, self.db_path
|
self._config, self.db_path
|
||||||
)
|
)
|
||||||
|
|
@ -489,7 +487,6 @@ class Store:
|
||||||
self.settings_table = await self.db.create_table(
|
self.settings_table = await self.db.create_table(
|
||||||
"settings", schema=SettingsRecord
|
"settings", schema=SettingsRecord
|
||||||
)
|
)
|
||||||
# Save current settings to the new database
|
|
||||||
settings_data = self._config.model_dump(mode="json")
|
settings_data = self._config.model_dump(mode="json")
|
||||||
await self.settings_table.add(
|
await self.settings_table.add(
|
||||||
[SettingsRecord(id="settings", settings=json.dumps(settings_data))]
|
[SettingsRecord(id="settings", settings=json.dumps(settings_data))]
|
||||||
|
|
@ -582,7 +579,6 @@ class Store:
|
||||||
where="id = 'settings'",
|
where="id = 'settings'",
|
||||||
)
|
)
|
||||||
else:
|
else:
|
||||||
# Create new settings record
|
|
||||||
settings_data = self._config.model_dump(mode="json")
|
settings_data = self._config.model_dump(mode="json")
|
||||||
settings_data["version"] = version
|
settings_data["version"] = version
|
||||||
await self.settings_table.add(
|
await self.settings_table.add(
|
||||||
|
|
@ -603,7 +599,6 @@ class Store:
|
||||||
if "chunks" in (await self.db.list_tables()).tables:
|
if "chunks" in (await self.db.list_tables()).tables:
|
||||||
await self.db.drop_table("chunks")
|
await self.db.drop_table("chunks")
|
||||||
|
|
||||||
# Update the ChunkRecord model with new vector dimension
|
|
||||||
self.ChunkRecord = create_chunk_model(self.embedder._vector_dim)
|
self.ChunkRecord = create_chunk_model(self.embedder._vector_dim)
|
||||||
self.chunks_table = await self.db.create_table(
|
self.chunks_table = await self.db.create_table(
|
||||||
"chunks", schema=self.ChunkRecord
|
"chunks", schema=self.ChunkRecord
|
||||||
|
|
|
||||||
|
|
@ -65,7 +65,6 @@ class ChunkRepository:
|
||||||
chunks if needed.
|
chunks if needed.
|
||||||
"""
|
"""
|
||||||
self.store._assert_writable()
|
self.store._assert_writable()
|
||||||
# Handle single chunk
|
|
||||||
if isinstance(entity, Chunk):
|
if isinstance(entity, Chunk):
|
||||||
assert entity.document_id, "Chunk must have a document_id to be created"
|
assert entity.document_id, "Chunk must have a document_id to be created"
|
||||||
assert entity.embedding is not None, "Chunk must have an embedding"
|
assert entity.embedding is not None, "Chunk must have an embedding"
|
||||||
|
|
@ -78,7 +77,6 @@ class ChunkRepository:
|
||||||
entity.id = chunk_id
|
entity.id = chunk_id
|
||||||
return entity
|
return entity
|
||||||
|
|
||||||
# Handle batch of chunks
|
|
||||||
chunks = entity
|
chunks = entity
|
||||||
if not chunks:
|
if not chunks:
|
||||||
return []
|
return []
|
||||||
|
|
@ -88,7 +86,6 @@ class ChunkRepository:
|
||||||
assert chunk.document_id, "All chunks must have a document_id to be created"
|
assert chunk.document_id, "All chunks must have a document_id to be created"
|
||||||
assert chunk.embedding is not None, "All chunks must have embeddings"
|
assert chunk.embedding is not None, "All chunks must have embeddings"
|
||||||
|
|
||||||
# Prepare all chunk records
|
|
||||||
chunk_records = []
|
chunk_records = []
|
||||||
for chunk in chunks:
|
for chunk in chunks:
|
||||||
chunk_id = str(uuid4())
|
chunk_id = str(uuid4())
|
||||||
|
|
@ -97,7 +94,6 @@ class ChunkRepository:
|
||||||
chunk_records.append(chunk_record)
|
chunk_records.append(chunk_record)
|
||||||
chunk.id = chunk_id
|
chunk.id = chunk_id
|
||||||
|
|
||||||
# Single batch insert for all chunks
|
|
||||||
await self.store.chunks_table.add(chunk_records)
|
await self.store.chunks_table.add(chunk_records)
|
||||||
|
|
||||||
return chunks
|
return chunks
|
||||||
|
|
@ -305,7 +301,6 @@ class ChunkRepository:
|
||||||
|
|
||||||
results = await query_to_pydantic(query, self.store.ChunkRecord)
|
results = await query_to_pydantic(query, self.store.ChunkRecord)
|
||||||
|
|
||||||
# Get document info from the mutable attributes table
|
|
||||||
doc_rows = await (
|
doc_rows = await (
|
||||||
self.store.document_meta_table.query()
|
self.store.document_meta_table.query()
|
||||||
.select(["id", "uri", "title", "metadata"])
|
.select(["id", "uri", "title", "metadata"])
|
||||||
|
|
@ -401,10 +396,8 @@ class ChunkRepository:
|
||||||
|
|
||||||
df = await query_result.to_pandas()
|
df = await query_result.to_pandas()
|
||||||
|
|
||||||
# Extract scores
|
|
||||||
scores = extract_scores(df)
|
scores = extract_scores(df)
|
||||||
|
|
||||||
# Convert DataFrame rows to ChunkRecords
|
|
||||||
pydantic_results = [
|
pydantic_results = [
|
||||||
self.store.ChunkRecord(
|
self.store.ChunkRecord(
|
||||||
id=str(row["id"]),
|
id=str(row["id"]),
|
||||||
|
|
@ -433,7 +426,6 @@ class ChunkRepository:
|
||||||
)
|
)
|
||||||
documents_map = {str(row["id"]): row for row in doc_rows}
|
documents_map = {str(row["id"]): row for row in doc_rows}
|
||||||
|
|
||||||
# Build final results with document info
|
|
||||||
chunks_with_scores = []
|
chunks_with_scores = []
|
||||||
for i, chunk_record in enumerate(pydantic_results):
|
for i, chunk_record in enumerate(pydantic_results):
|
||||||
doc = documents_map.get(chunk_record.document_id)
|
doc = documents_map.get(chunk_record.document_id)
|
||||||
|
|
|
||||||
|
|
@ -307,7 +307,6 @@ class DocumentRepository:
|
||||||
"""Delete a document by its ID."""
|
"""Delete a document by its ID."""
|
||||||
self.store._assert_writable()
|
self.store._assert_writable()
|
||||||
|
|
||||||
# Check if document exists
|
|
||||||
doc = await self.get_by_id(entity_id)
|
doc = await self.get_by_id(entity_id)
|
||||||
if doc is None:
|
if doc is None:
|
||||||
return False
|
return False
|
||||||
|
|
@ -416,7 +415,6 @@ class DocumentRepository:
|
||||||
)
|
)
|
||||||
await ensure_indexes(self.store.document_items_table, "document_items")
|
await ensure_indexes(self.store.document_items_table, "document_items")
|
||||||
|
|
||||||
# Get count before deletion
|
|
||||||
count = len(
|
count = len(
|
||||||
await query_to_pydantic(
|
await query_to_pydantic(
|
||||||
self.store.documents_table.query().limit(1), DocumentRecord
|
self.store.documents_table.query().limit(1), DocumentRecord
|
||||||
|
|
|
||||||
|
|
@ -36,7 +36,6 @@ class SettingsRepository:
|
||||||
self.store._assert_writable()
|
self.store._assert_writable()
|
||||||
current_config = self.store._config.model_dump(mode="json")
|
current_config = self.store._config.model_dump(mode="json")
|
||||||
|
|
||||||
# Check if settings exist
|
|
||||||
existing = await query_to_pydantic(
|
existing = await query_to_pydantic(
|
||||||
self.store.settings_table.query().where("id = 'settings'").limit(1),
|
self.store.settings_table.query().where("id = 'settings'").limit(1),
|
||||||
SettingsRecord,
|
SettingsRecord,
|
||||||
|
|
@ -48,14 +47,12 @@ class SettingsRepository:
|
||||||
if "version" in existing_settings:
|
if "version" in existing_settings:
|
||||||
current_config["version"] = existing_settings["version"]
|
current_config["version"] = existing_settings["version"]
|
||||||
|
|
||||||
# Update existing settings
|
|
||||||
if existing_settings != current_config:
|
if existing_settings != current_config:
|
||||||
await self.store.settings_table.update(
|
await self.store.settings_table.update(
|
||||||
{"settings": json.dumps(current_config)},
|
{"settings": json.dumps(current_config)},
|
||||||
where="id = 'settings'",
|
where="id = 'settings'",
|
||||||
)
|
)
|
||||||
else:
|
else:
|
||||||
# Create new settings
|
|
||||||
settings_record = SettingsRecord(
|
settings_record = SettingsRecord(
|
||||||
id="settings", settings=json.dumps(current_config)
|
id="settings", settings=json.dumps(current_config)
|
||||||
)
|
)
|
||||||
|
|
|
||||||
|
|
@ -131,7 +131,6 @@ async def _apply_compress_docling_document(store: Store) -> None: # pragma: no
|
||||||
if records:
|
if records:
|
||||||
await store.documents_table.add(records)
|
await store.documents_table.add(records)
|
||||||
logger.info("Recovered batch %d/%d", batch_num, total_batches)
|
logger.info("Recovered batch %d/%d", batch_num, total_batches)
|
||||||
# Cleanup staging
|
|
||||||
await store.db.drop_table("documents_v4_staging")
|
await store.db.drop_table("documents_v4_staging")
|
||||||
logger.info("Recovery complete")
|
logger.info("Recovery complete")
|
||||||
return
|
return
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue