RomanNum3ral
/
honcho
mirror of https://github.com/plastic-labs/honcho.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

fix: LanceDB

This commit is contained in:
Rajat Ahuja 2025-12-04 15:55:22 -05:00
parent 2af9b1e1ca
commit f9a617a911
7 changed files with 137 additions and 67 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
.gitignore vendored
View File

@ -189,3 +189,4 @@ CRUSH.md
metrics.jsonl
AGENTS.md
lancedb_data/

54
migrations/versions/f1a2b3c4d5e6_add_chunk_index_to_message_embeddings.py
View File

@ -1,54 +0,0 @@
"""add chunk_index to message_embeddings
This migration adds the chunk_index column to message_embeddings table for tracking
chunked message embeddings in external vector stores (turbopuffer/lancedb).
Revision ID: f1a2b3c4d5e6
Revises: baa22cad81e2
Create Date: 2025-11-24 12:00:00.000000
"""
from collections.abc import Sequence
import sqlalchemy as sa
from alembic import op
from migrations.utils import column_exists, get_schema
# revision identifiers, used by Alembic.
revision: str = "f1a2b3c4d5e6"
down_revision: str | None = "baa22cad81e2"
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None
schema = get_schema()
def upgrade() -> None:
"""Add chunk_index column to message_embeddings for tracking chunked message embeddings."""
inspector = sa.inspect(op.get_bind())
# Add chunk_index column to message_embeddings if it doesn't exist
# This is needed to track which chunk of a message this embedding represents
# Vector ID format: {message_public_id}_{chunk_index}
if not column_exists("message_embeddings", "chunk_index", inspector):
op.add_column(
"message_embeddings",
sa.Column(
"chunk_index",
sa.Integer(),
nullable=False,
server_default="0",
),
schema=schema,
)
def downgrade() -> None:
"""Remove chunk_index column from message_embeddings."""
inspector = sa.inspect(op.get_bind())
# Remove chunk_index column if it exists
if column_exists("message_embeddings", "chunk_index", inspector):
op.drop_column("message_embeddings", "chunk_index", schema=schema)

97
migrations/versions/f1a2b3c4d5e6_make_embeddings_nullable.py Normal file
View File

@ -0,0 +1,97 @@
"""add chunk_index to message_embeddings and make embeddings nullable
This migration:
1. Adds the chunk_index column to message_embeddings table for tracking
chunked message embeddings in external vector stores (turbopuffer/lancedb).
2. Makes embedding columns nullable in both message_embeddings and documents tables
since embeddings are now stored in external vector stores instead of PostgreSQL.
Revision ID: f1a2b3c4d5e6
Revises: baa22cad81e2
Create Date: 2025-11-24 12:00:00.000000
"""
from collections.abc import Sequence
import sqlalchemy as sa
from alembic import op
from pgvector.sqlalchemy import Vector
from migrations.utils import column_exists, get_schema
# revision identifiers, used by Alembic.
revision: str = "f1a2b3c4d5e6"
down_revision: str | None = "baa22cad81e2"
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None
schema = get_schema()
def upgrade() -> None:
"""Add chunk_index column to message_embeddings and make embeddings nullable."""
inspector = sa.inspect(op.get_bind())
# Add chunk_index column to message_embeddings if it doesn't exist
# This is needed to track which chunk of a message this embedding represents
# Vector ID format: {message_public_id}_{chunk_index}
if not column_exists("message_embeddings", "chunk_index", inspector):
op.add_column(
"message_embeddings",
sa.Column(
"chunk_index",
sa.Integer(),
nullable=False,
server_default="0",
),
schema=schema,
)
# Make message_embeddings.embedding nullable since embeddings are now stored
# in external vector stores (turbopuffer/lancedb) instead of PostgreSQL
op.alter_column(
"message_embeddings",
"embedding",
existing_type=Vector(1536),
nullable=True,
schema=schema,
)
# Make documents.embedding nullable for the same reason
# (this should already be nullable, but ensure it for consistency)
op.alter_column(
"documents",
"embedding",
existing_type=Vector(1536),
nullable=True,
schema=schema,
)
def downgrade() -> None:
"""Remove chunk_index column and revert embedding columns to non-nullable."""
inspector = sa.inspect(op.get_bind())
# Revert documents.embedding back to nullable=True (it was originally nullable=True)
op.alter_column(
"documents",
"embedding",
existing_type=Vector(1536),
nullable=True, # Keep as nullable since it was nullable in the original schema
schema=schema,
)
# Revert message_embeddings.embedding back to nullable=False
# Note: This may fail if there are NULL values in the database
op.alter_column(
"message_embeddings",
"embedding",
existing_type=Vector(1536),
nullable=False,
schema=schema,
)
# Remove chunk_index column if it exists
if column_exists("message_embeddings", "chunk_index", inspector):
op.drop_column("message_embeddings", "chunk_index", schema=schema)

17
src/crud/document.py
View File

@ -153,7 +153,8 @@ async def create_documents(
Count of new documents
"""
honcho_documents: list[models.Document] = []
embeddings_to_store: list[tuple[str, list[float]]] = [] # [(doc_id, embedding)]
# Store (document_model, embedding) pairs - IDs aren't available until after commit
docs_with_embeddings: list[tuple[models.Document, list[float]]] = []
for doc in documents:
try:
@ -180,9 +181,9 @@ async def create_documents(
)
honcho_documents.append(new_doc)
# Track embedding for vector store (will use document's generated ID)
# Track embedding for vector store (ID will be available after commit)
if doc.embedding:
embeddings_to_store.append((new_doc.id, doc.embedding))
docs_with_embeddings.append((new_doc, doc.embedding))
except Exception as e:
logger.error(
@ -194,8 +195,8 @@ async def create_documents(
db.add_all(honcho_documents)
await db.commit()
# Store embeddings in vector store after documents are committed
if embeddings_to_store:
# Store embeddings in vector store after documents are committed (IDs now available)
if docs_with_embeddings:
vector_store = get_vector_store()
namespace = vector_store.get_document_namespace(
workspace_name, observer, observed
@ -203,12 +204,10 @@ async def create_documents(
# Build vector records with metadata for filtering
vector_records: list[VectorRecord] = []
doc_lookup = {doc.id: doc for doc in honcho_documents}
for doc_id, embedding in embeddings_to_store:
doc = doc_lookup[doc_id]
for doc, embedding in docs_with_embeddings:
vector_records.append(
VectorRecord(
id=doc_id,
id=doc.id,
embedding=embedding,
metadata={
"workspace_name": workspace_name,

31
src/vector_store/lancedb.py
View File

@ -29,6 +29,11 @@ class LanceDBVectorStore(VectorStore):
Uses LanceDB's embedded mode for local vector storage.
Each namespace corresponds to a LanceDB table.
Note: LanceDB table names can only contain alphanumeric characters,
underscores, hyphens, and periods. We use '.' as the namespace separator
instead of ':' (used by base class) since '.' is not allowed in
workspace/peer IDs.
"""
_db: lancedb.DBConnection
@ -38,6 +43,25 @@ class LanceDBVectorStore(VectorStore):
super().__init__()
self._db = lancedb.connect(settings.VECTOR_STORE.LANCEDB_PATH)
# === Namespace helpers (override to use LanceDB-compatible separator) ===
def get_document_namespace(
self, workspace_name: str, observer: str, observed: str
) -> str:
"""
Get the namespace for document embeddings (per collection).
Uses '.' as separator instead of ':' for LanceDB compatibility.
"""
return f"{self.namespace_prefix}.{workspace_name}.{observer}.{observed}"
def get_message_namespace(self, workspace_name: str) -> str:
"""
Get the namespace for message embeddings (per workspace).
Uses '.' as separator instead of ':' for LanceDB compatibility.
"""
return f"{self.namespace_prefix}.{workspace_name}.messages"
def _get_table(self, namespace: str) -> lancedb.table.Table | None:
"""Get a table if it exists, otherwise return None."""
if namespace in self._db.table_names():
@ -47,7 +71,7 @@ class LanceDBVectorStore(VectorStore):
def _get_or_create_table(
self, namespace: str, sample_data: list[dict[str, Any]] | None = None
) -> lancedb.table.Table:
"""
"""_get_or_create_table
Get existing table or create if not exists.
Args:
@ -127,10 +151,13 @@ class LanceDBVectorStore(VectorStore):
try:
rows = [self._row_to_dict(v) for v in vectors]
print(f"Rows: {rows}")
table = self._get_or_create_table(namespace, sample_data=rows)
# Use merge_insert for upsert behavior
table.merge_insert("id").when_matched_update_all().execute(rows)
table.merge_insert(
"id"
).when_matched_update_all().when_not_matched_insert_all().execute(rows)
logger.debug(f"Upserted {len(vectors)} vectors to namespace {namespace}")
except Exception:

4
tests/alembic/revisions/__init__.py
View File

@ -21,7 +21,7 @@ from . import (
test_d429de0e5338_adopt_peer_paradigm,
test_e9b705f9adf9_add_server_defaults_to_timestamp_,
test_ec8f94139b02_codify_workspace_name_and_message_id_in_,
test_f1a2b3c4d5e6_add_chunk_index_to_message_embeddings,
test_f1a2b3c4d5e6_make_embeddings_nullable,
)
__all__ = [
@ -45,5 +45,5 @@ __all__ = [
"test_d429de0e5338_adopt_peer_paradigm",
"test_e9b705f9adf9_add_server_defaults_to_timestamp_",
"test_ec8f94139b02_codify_workspace_name_and_message_id_in_",
"test_f1a2b3c4d5e6_add_chunk_index_to_message_embeddings",
"test_f1a2b3c4d5e6_make_embeddings_nullable",
]

0
tests/alembic/revisions/test_f1a2b3c4d5e6_add_chunk_index_to_message_embeddings.py → tests/alembic/revisions/test_f1a2b3c4d5e6_make_embeddings_nullable.py
View File