Compare commits

...

10 Commits

Author SHA1 Message Date
claude[bot] 339bf6cc31 fix(core): remove inline ALTER TABLE from postgres_search_repository
Runtime DDL causes deadlocks when two concurrent index jobs race to
ALTER TABLE search_vector_chunks simultaneously. The alembic migration
m6h7i8j9k0l1_add_vector_sync_fingerprints.py already handles these
schema changes — runtime code should never run ALTER TABLE.

Fixes #727

Co-authored-by: Paul Hernandez <phernandez@users.noreply.github.com>
Signed-off-by: claude[bot] <41898282+claude[bot]@users.noreply.github.com>
2026-04-08 04:44:39 +00:00
Paul Hernandez 8c81d3ce17 perf(core): reduce postgres vector sync work (#723)
Signed-off-by: phernandez <paul@basicmachines.co>
2026-04-07 19:00:19 -05:00
Drew Cain b35d594ef0 fix(core): preserve external_id during entity upsert on re-index (#724)
Signed-off-by: Drew Cain <groksrc@gmail.com>
Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
2026-04-07 10:26:10 -05:00
phernandez e982900084 fix(core): strip null bytes from markdown content before database insert
PostgreSQL rejects null bytes (0x00) in text columns, causing
CharacterNotInRepertoireError when syncing files like Claude agent
definitions that contain embedded nulls. SQLite silently accepts them,
so this only surfaces in cloud environments.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Signed-off-by: phernandez <paul@basicmachines.co>
2026-04-06 23:14:40 -05:00
Drew Cain b3403e96b3 fix: add workspace routing to cloud upload and API client (#704)
Signed-off-by: Drew Cain <groksrc@gmail.com>
Signed-off-by: phernandez <paul@basicmachines.co>
Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Co-authored-by: phernandez <paul@basicmachines.co>
2026-04-06 18:18:40 -05:00
Paul Hernandez fe04a0b2a2 fix(mcp): pass workspace parameter through client factory (#722)
Signed-off-by: phernandez <paul@basicmachines.co>
Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
2026-04-06 16:57:37 -05:00
jope-bm 86ad639890 fix(cli): show display_name instead of UUID for private projects in CLI (#718)
Signed-off-by: Joe P <joe@basicmemory.com>
Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
2026-04-06 10:14:46 -06:00
Paul Hernandez 88c8f18200 feat(core): add note_content tenant schema primitive (#719)
Signed-off-by: phernandez <paul@basicmachines.co>
2026-04-04 22:06:37 -05:00
Drew Cain 41a16b93cb fix: Increase brew outdated timeout from 15s to 60s (#695)
Signed-off-by: Drew Cain <groksrc@gmail.com>
Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
2026-04-04 20:38:31 -05:00
Drew Cain 367fcaac50 perf: eliminate redundant DB queries in upsert_entity_from_markdown (#714)
Signed-off-by: Drew Cain <groksrc@gmail.com>
Signed-off-by: phernandez <paul@basicmachines.co>
Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Co-authored-by: phernandez <paul@basicmachines.co>
2026-04-04 12:40:27 -05:00
43 changed files with 3440 additions and 250 deletions
+20 -1
View File
@@ -1,3 +1,22 @@
{
"enabledPlugins": {}
"$schema": "https://json.schemastore.org/claude-code-settings.json",
"env": {
"CLAUDE_BASH_MAINTAIN_PROJECT_WORKING_DIR": "1",
"CLAUDE_CODE_DISABLE_FEEDBACK_SURVEY": "1",
"DISABLE_TELEMETRY": "1",
"CLAUDE_CODE_NO_FLICKER": "1",
"CLAUDE_CODE_DISABLE_ADAPTIVE_THINKING": "1"
},
"permissions": {
"allow": [
"Bash(just fast-check)",
"Bash(just check)",
"Bash(just fix)",
"Bash(just typecheck)",
"Bash(just lint)",
"Bash(just test)"
],
"deny": []
},
"enableAllProjectMcpServers": true
}
+53 -25
View File
@@ -66,7 +66,7 @@ target_metadata = Base.metadata
# Add this function to tell Alembic what to include/exclude
def include_object(object, name, type_, reflected, compare_to):
def include_object(obj, name, type_, reflected, compare_to):
# Ignore SQLite FTS tables
if type_ == "table" and name.startswith("search_index"):
return False
@@ -118,6 +118,54 @@ async def run_async_migrations(connectable):
await connectable.dispose()
def _run_async_migrations_with_asyncio_run(connectable) -> None:
"""Run async migrations with asyncio.run while closing failed coroutines.
Trigger: asyncio.run() may reject execution when another event loop is already active.
Why: Python raises before awaiting the coroutine, which otherwise leaks a
RuntimeWarning about an un-awaited coroutine.
Outcome: close the pending coroutine before bubbling the RuntimeError to the
fallback path.
"""
migration_coro = run_async_migrations(connectable)
try:
asyncio.run(migration_coro)
except RuntimeError:
migration_coro.close()
raise
def _run_async_migrations_in_thread(connectable) -> None:
"""Run async migrations in a dedicated thread with its own event loop."""
import concurrent.futures
def run_in_thread():
"""Run async migrations in a new event loop in a separate thread."""
new_loop = asyncio.new_event_loop()
asyncio.set_event_loop(new_loop)
try:
new_loop.run_until_complete(run_async_migrations(connectable))
finally:
new_loop.close()
with concurrent.futures.ThreadPoolExecutor() as executor:
future = executor.submit(run_in_thread)
future.result() # Wait for completion and re-raise any exceptions
def _run_async_engine_migrations(connectable) -> None:
"""Run async-engine migrations with a running-loop fallback."""
try:
_run_async_migrations_with_asyncio_run(connectable)
except RuntimeError as e:
if "cannot be called from a running event loop" in str(e):
# We're in a running event loop (likely uvloop or Python 3.14+ tests).
# Switch to a dedicated thread so Alembic can finish without nesting loops.
_run_async_migrations_in_thread(connectable)
else:
raise
def run_migrations_online() -> None:
"""Run migrations in 'online' mode.
@@ -148,30 +196,10 @@ def run_migrations_online() -> None:
# Handle async engines (PostgreSQL with asyncpg)
if isinstance(connectable, AsyncEngine):
# Try to run async migrations
# nest_asyncio allows asyncio.run() from within event loops, but doesn't work with uvloop
try:
asyncio.run(run_async_migrations(connectable))
except RuntimeError as e:
if "cannot be called from a running event loop" in str(e):
# We're in a running event loop (likely uvloop) - need to use a different approach
# Create a new thread to run the async migrations
import concurrent.futures
def run_in_thread():
"""Run async migrations in a new event loop in a separate thread."""
new_loop = asyncio.new_event_loop()
asyncio.set_event_loop(new_loop)
try:
new_loop.run_until_complete(run_async_migrations(connectable))
finally:
new_loop.close()
with concurrent.futures.ThreadPoolExecutor() as executor:
future = executor.submit(run_in_thread)
future.result() # Wait for completion and re-raise any exceptions
else:
raise
# Trigger: async engines need Alembic work to cross the sync/async boundary.
# Why: most callers can use asyncio.run(), but running-loop contexts need a thread fallback.
# Outcome: migrations complete without leaking un-awaited coroutines.
_run_async_engine_migrations(connectable)
else:
# Handle sync engines (SQLite) or sync connections
if hasattr(connectable, "connect"):
@@ -0,0 +1,65 @@
"""Add note_content table
Revision ID: l5g6h7i8j9k0
Revises: k4e5f6g7h8i9
Create Date: 2026-04-04 12:00:00.000000
"""
from typing import Sequence, Union
import sqlalchemy as sa
from alembic import op
# revision identifiers, used by Alembic.
revision: str = "l5g6h7i8j9k0"
down_revision: Union[str, None] = "k4e5f6g7h8i9"
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None
def upgrade() -> None:
"""Create note_content for materialized note content and sync state."""
op.create_table(
"note_content",
sa.Column("entity_id", sa.Integer(), nullable=False),
sa.Column("project_id", sa.Integer(), nullable=False),
sa.Column("external_id", sa.String(), nullable=False),
sa.Column("file_path", sa.String(), nullable=False),
sa.Column("markdown_content", sa.Text(), nullable=False),
sa.Column("db_version", sa.BigInteger(), nullable=False),
sa.Column("db_checksum", sa.String(), nullable=False),
sa.Column("file_version", sa.BigInteger(), nullable=True),
sa.Column("file_checksum", sa.String(), nullable=True),
sa.Column("file_write_status", sa.String(), nullable=False),
sa.Column("last_source", sa.String(), nullable=True),
sa.Column("updated_at", sa.DateTime(timezone=True), nullable=False),
sa.Column("file_updated_at", sa.DateTime(timezone=True), nullable=True),
sa.Column("last_materialization_error", sa.Text(), nullable=True),
sa.Column("last_materialization_attempt_at", sa.DateTime(timezone=True), nullable=True),
sa.CheckConstraint(
"file_write_status IN ("
"'pending', "
"'writing', "
"'synced', "
"'failed', "
"'external_change_detected'"
")",
name="ck_note_content_file_write_status",
),
sa.ForeignKeyConstraint(["entity_id"], ["entity.id"], ondelete="CASCADE"),
sa.ForeignKeyConstraint(["project_id"], ["project.id"], ondelete="CASCADE"),
sa.PrimaryKeyConstraint("entity_id"),
)
op.create_index("ix_note_content_project_id", "note_content", ["project_id"], unique=False)
op.create_index("ix_note_content_file_path", "note_content", ["file_path"], unique=False)
op.create_index("ix_note_content_external_id", "note_content", ["external_id"], unique=True)
def downgrade() -> None:
"""Drop note_content and its supporting indexes."""
op.drop_index("ix_note_content_external_id", table_name="note_content")
op.drop_index("ix_note_content_file_path", table_name="note_content")
op.drop_index("ix_note_content_project_id", table_name="note_content")
op.drop_table("note_content")
@@ -0,0 +1,84 @@
"""Persist vector sync fingerprints on chunk metadata.
Revision ID: m6h7i8j9k0l1
Revises: l5g6h7i8j9k0
Create Date: 2026-04-07 00:00:00.000000
"""
from typing import Sequence, Union
from alembic import op
# revision identifiers, used by Alembic.
revision: str = "m6h7i8j9k0l1"
down_revision: Union[str, None] = "l5g6h7i8j9k0"
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None
def upgrade() -> None:
"""Add entity fingerprint + embedding model metadata to Postgres chunk rows.
Trigger: vector sync now fast-skips unchanged entities using persisted
semantic fingerprints.
Why: chunk rows already own the per-entity derived metadata we diff against,
so persisting the fingerprint on that table avoids a second sync-state table.
Outcome: existing rows get empty-string placeholders and will be refreshed on
the next vector sync before they become eligible for skip checks.
"""
connection = op.get_bind()
if connection.dialect.name != "postgresql":
return
op.execute(
"""
ALTER TABLE search_vector_chunks
ADD COLUMN IF NOT EXISTS entity_fingerprint TEXT
"""
)
op.execute(
"""
ALTER TABLE search_vector_chunks
ADD COLUMN IF NOT EXISTS embedding_model TEXT
"""
)
op.execute(
"""
UPDATE search_vector_chunks
SET entity_fingerprint = COALESCE(entity_fingerprint, ''),
embedding_model = COALESCE(embedding_model, '')
"""
)
op.execute(
"""
ALTER TABLE search_vector_chunks
ALTER COLUMN entity_fingerprint SET NOT NULL
"""
)
op.execute(
"""
ALTER TABLE search_vector_chunks
ALTER COLUMN embedding_model SET NOT NULL
"""
)
def downgrade() -> None:
"""Remove vector sync fingerprint columns from Postgres chunk rows."""
connection = op.get_bind()
if connection.dialect.name != "postgresql":
return
op.execute(
"""
ALTER TABLE search_vector_chunks
DROP COLUMN IF EXISTS embedding_model
"""
)
op.execute(
"""
ALTER TABLE search_vector_chunks
DROP COLUMN IF EXISTS entity_fingerprint
"""
)
+1 -1
View File
@@ -22,7 +22,7 @@ PACKAGE_NAME = "basic-memory"
PYPI_JSON_URL = "https://pypi.org/pypi/basic-memory/json"
PYPI_TIMEOUT_SECONDS = 5
BREW_OUTDATED_TIMEOUT_SECONDS = 15
BREW_OUTDATED_TIMEOUT_SECONDS = 60
UV_UPGRADE_TIMEOUT_SECONDS = 180
BREW_UPGRADE_TIMEOUT_SECONDS = 600
@@ -41,6 +41,10 @@ async def fetch_cloud_projects(
) -> CloudProjectList:
"""Fetch list of projects from cloud API.
Args:
project_name: Optional project name for workspace resolution
workspace: Cloud workspace tenant_id to list projects from
Returns:
CloudProjectList with projects from cloud
"""
@@ -112,12 +116,12 @@ async def sync_project(project_name: str, force_full: bool = False) -> None:
Args:
project_name: Name of project to sync
force_full: If True, force a full scan bypassing watermark optimization
force_full: ignored, kept for backwards compatibility
"""
try:
from basic_memory.cli.commands.command_utils import run_sync
await run_sync(project=project_name, force_full=force_full)
await run_sync(project=project_name)
except Exception as e:
raise CloudUtilsError(f"Failed to sync project '{project_name}': {e}") from e
@@ -106,6 +106,8 @@ def upload(
raise typer.Exit(1)
# Perform upload (or dry run)
if resolved_workspace:
console.print(f"[dim]Using workspace: {resolved_workspace}[/dim]")
if dry_run:
console.print(
f"[yellow]DRY RUN: Showing what would be uploaded to '{project}'[/yellow]"
@@ -140,7 +142,7 @@ def upload(
if sync and not dry_run:
console.print(f"[blue]Syncing project '{project}'...[/blue]")
try:
await sync_project(project, force_full=True)
await sync_project(project)
except Exception as e:
console.print(f"[yellow]Warning: Sync failed: {e}[/yellow]")
console.print("[dim]Files uploaded but may not be indexed yet[/dim]")
+1 -1
View File
@@ -101,7 +101,7 @@ async def run_doctor() -> None:
console.print("[green]OK[/green] Manual file written")
sync_data = await project_client.sync(
project_id, force_full=True, run_in_background=False
project_id, force_full=False, run_in_background=False
)
sync_report = SyncReportResponse.model_validate(sync_data)
if sync_report.total == 0:
+9 -1
View File
@@ -253,6 +253,12 @@ def list_projects(
if cloud_project is not None and cloud_ws_name:
ws_label = f"{cloud_ws_name} ({cloud_ws_type})" if cloud_ws_type else cloud_ws_name
# display_name is a human label for private UUID-named projects (e.g., "My Project").
# Keep "name" as the canonical identifier for scripting/JSON consumers;
# the Rich table uses display_name when available.
display_name = (
cloud_project.display_name if cloud_project and cloud_project.display_name else None
)
row_data = {
"name": project_name,
"permalink": permalink,
@@ -263,6 +269,8 @@ def list_projects(
"sync": has_sync,
"is_default": is_default,
}
if display_name:
row_data["display_name"] = display_name
if ws_label:
row_data["workspace"] = cloud_ws_name or ""
if cloud_ws_type:
@@ -278,7 +286,7 @@ def list_projects(
# --- Rich table output ---
for row_data in project_rows:
table.add_row(
row_data["name"],
row_data.get("display_name") or row_data["name"],
row_data["local_path"],
row_data["cloud_path"],
row_data.get("workspace", "")
+6
View File
@@ -198,6 +198,12 @@ class BasicMemoryConfig(BaseSettings):
description="Batch size for vector sync orchestration flushes.",
gt=0,
)
semantic_postgres_prepare_concurrency: int = Field(
default=4,
description="Number of Postgres entity prepare tasks to run concurrently during vector sync. Postgres only; keep this low to avoid overdriving the database connection pool.",
gt=0,
le=16,
)
semantic_embedding_cache_dir: str | None = Field(
default=None,
description="Optional cache directory for FastEmbed model artifacts.",
@@ -249,6 +249,10 @@ class EntityParser:
content = strip_bom(content)
# PostgreSQL rejects null bytes (0x00) in text columns.
# Some markdown files (e.g. Claude agent definitions) contain embedded nulls.
content = content.replace("\x00", "")
# Parse frontmatter with proper error handling for malformed YAML.
# We use frontmatter.parse() instead of frontmatter.loads() because
# loads() does Post(content, handler, **metadata), which crashes when
+7 -4
View File
@@ -128,11 +128,14 @@ async def get_cloud_control_plane_client(
yield client
# Optional factory override for dependency injection
_client_factory: Optional[Callable[[], AbstractAsyncContextManager[AsyncClient]]] = None
# Optional factory override for dependency injection.
# The factory accepts an optional workspace keyword argument so that MCP tools
# can route individual requests to a different workspace than the one set at
# connection time. See basic-memory-cloud main.py tenant_asgi_client_factory.
_client_factory: Optional[Callable[..., AbstractAsyncContextManager[AsyncClient]]] = None
def set_client_factory(factory: Callable[[], AbstractAsyncContextManager[AsyncClient]]) -> None:
def set_client_factory(factory: Callable[..., AbstractAsyncContextManager[AsyncClient]]) -> None:
"""Override the default client factory (for cloud app, testing, etc)."""
global _client_factory
_client_factory = factory
@@ -173,7 +176,7 @@ async def get_client(
4. Local ASGI transport by default.
"""
if _client_factory:
async with _client_factory() as client:
async with _client_factory(workspace=workspace) as client:
yield client
return
+5 -5
View File
@@ -622,10 +622,10 @@ async def get_project_client(
# Step 1b: Factory injection (in-process cloud server)
# Trigger: set_client_factory() was called (e.g., by cloud MCP server)
# Why: the transport layer already resolved workspace and tenant context;
# attempting cloud workspace resolution here would call the production
# control-plane API with no valid credentials and fail with 401
# Outcome: use the factory client directly, skip workspace resolution
# Why: the factory's transport layer handles auth and tenant resolution;
# we pass workspace through so the transport can route to the correct
# workspace when the tool specifies one different from the connection default
# Outcome: factory client with optional workspace override via inner request headers
if is_factory_mode():
route_mode = "factory"
with telemetry.scope(
@@ -635,7 +635,7 @@ async def get_project_client(
workspace_id=workspace,
):
logger.debug("Using injected client factory for project routing")
async with get_client() as client:
async with get_client(workspace=workspace) as client:
active_project = await get_active_project(client, resolved_project, context)
yield client, active_project
return
+2 -1
View File
@@ -2,12 +2,13 @@
import basic_memory
from basic_memory.models.base import Base
from basic_memory.models.knowledge import Entity, Observation, Relation
from basic_memory.models.knowledge import Entity, NoteContent, Observation, Relation
from basic_memory.models.project import Project
__all__ = [
"Base",
"Entity",
"NoteContent",
"Observation",
"Relation",
"Project",
+76
View File
@@ -6,6 +6,8 @@ from basic_memory.utils import ensure_timezone_aware
from typing import Optional
from sqlalchemy import (
BigInteger,
CheckConstraint,
Integer,
String,
Text,
@@ -116,6 +118,12 @@ class Entity(Base):
foreign_keys="[Relation.to_id]",
cascade="all, delete-orphan",
)
note_content = relationship(
"NoteContent",
back_populates="entity",
cascade="all, delete-orphan",
uselist=False,
)
@property
def relations(self):
@@ -141,6 +149,74 @@ class Entity(Base):
return f"Entity(id={self.id}, external_id='{self.external_id}', name='{self.title}', type='{self.note_type}', checksum='{self.checksum}')"
class NoteContent(Base):
"""Materialized markdown content and sync state for a note entity."""
__tablename__ = "note_content"
__table_args__ = (
CheckConstraint(
"file_write_status IN ("
"'pending', "
"'writing', "
"'synced', "
"'failed', "
"'external_change_detected'"
")",
name="ck_note_content_file_write_status",
),
Index("ix_note_content_project_id", "project_id"),
Index("ix_note_content_file_path", "file_path"),
Index("ix_note_content_external_id", "external_id", unique=True),
)
# Core identity mirrored from entity for hot note reads
entity_id: Mapped[int] = mapped_column(
Integer,
ForeignKey("entity.id", ondelete="CASCADE"),
primary_key=True,
)
project_id: Mapped[int] = mapped_column(
Integer,
ForeignKey("project.id", ondelete="CASCADE"),
nullable=False,
)
external_id: Mapped[str] = mapped_column(String, nullable=False)
file_path: Mapped[str] = mapped_column(String, nullable=False)
# Materialized content version tracked in the tenant database
markdown_content: Mapped[str] = mapped_column(Text, nullable=False)
db_version: Mapped[int] = mapped_column(BigInteger, nullable=False)
db_checksum: Mapped[str] = mapped_column(String, nullable=False)
# File materialization state tracked against the latest write attempts
file_version: Mapped[Optional[int]] = mapped_column(BigInteger, nullable=True)
file_checksum: Mapped[Optional[str]] = mapped_column(String, nullable=True)
file_write_status: Mapped[str] = mapped_column(String, nullable=False, default="pending")
last_source: Mapped[Optional[str]] = mapped_column(String, nullable=True)
updated_at: Mapped[datetime] = mapped_column(
DateTime(timezone=True),
default=lambda: datetime.now().astimezone(),
onupdate=lambda: datetime.now().astimezone(),
)
file_updated_at: Mapped[Optional[datetime]] = mapped_column(
DateTime(timezone=True),
nullable=True,
)
last_materialization_error: Mapped[Optional[str]] = mapped_column(Text, nullable=True)
last_materialization_attempt_at: Mapped[Optional[datetime]] = mapped_column(
DateTime(timezone=True),
nullable=True,
)
entity = relationship("Entity", back_populates="note_content")
def __repr__(self) -> str: # pragma: no cover
return (
f"NoteContent(entity_id={self.entity_id}, external_id='{self.external_id}', "
f"file_path='{self.file_path}', file_write_status='{self.file_write_status}')"
)
class Observation(Base):
"""An observation about an entity.
+4
View File
@@ -104,6 +104,8 @@ CREATE TABLE IF NOT EXISTS search_vector_chunks (
chunk_key TEXT NOT NULL,
chunk_text TEXT NOT NULL,
source_hash TEXT NOT NULL,
entity_fingerprint TEXT NOT NULL,
embedding_model TEXT NOT NULL,
updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
UNIQUE (project_id, entity_id, chunk_key)
)
@@ -124,6 +126,8 @@ CREATE TABLE IF NOT EXISTS search_vector_chunks (
chunk_key TEXT NOT NULL,
chunk_text TEXT NOT NULL,
source_hash TEXT NOT NULL,
entity_fingerprint TEXT NOT NULL,
embedding_model TEXT NOT NULL,
updated_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP
)
""")
+2
View File
@@ -1,10 +1,12 @@
from .entity_repository import EntityRepository
from .note_content_repository import NoteContentRepository
from .observation_repository import ObservationRepository
from .project_repository import ProjectRepository
from .relation_repository import RelationRepository
__all__ = [
"EntityRepository",
"NoteContentRepository",
"ObservationRepository",
"ProjectRepository",
"RelationRepository",
@@ -45,7 +45,17 @@ class EntityRepository(Repository[Entity]):
async with db.scoped_session(self.session_maker) as session:
return await self.select_by_id(session, entity_id)
async def get_by_external_id(self, external_id: str) -> Optional[Entity]:
async def _find_one_by_query(self, query, *, load_relations: bool) -> Optional[Entity]:
"""Return one entity row with optional eager loading."""
if load_relations:
return await self.find_one(query)
result = await self.execute_query(query, use_query_options=False)
return result.scalars().one_or_none()
async def get_by_external_id(
self, external_id: str, *, load_relations: bool = True
) -> Optional[Entity]:
"""Get entity by external UUID.
Args:
@@ -54,21 +64,21 @@ class EntityRepository(Repository[Entity]):
Returns:
Entity if found, None otherwise
"""
query = (
self.select().where(Entity.external_id == external_id).options(*self.get_load_options())
)
return await self.find_one(query)
query = self.select().where(Entity.external_id == external_id)
return await self._find_one_by_query(query, load_relations=load_relations)
async def get_by_permalink(self, permalink: str) -> Optional[Entity]:
async def get_by_permalink(
self, permalink: str, *, load_relations: bool = True
) -> Optional[Entity]:
"""Get entity by permalink.
Args:
permalink: Unique identifier for the entity
"""
query = self.select().where(Entity.permalink == permalink).options(*self.get_load_options())
return await self.find_one(query)
query = self.select().where(Entity.permalink == permalink)
return await self._find_one_by_query(query, load_relations=load_relations)
async def get_by_title(self, title: str) -> Sequence[Entity]:
async def get_by_title(self, title: str, *, load_relations: bool = True) -> Sequence[Entity]:
"""Get entities by title, ordered by shortest path first.
When multiple entities share the same title (in different folders),
@@ -82,23 +92,20 @@ class EntityRepository(Repository[Entity]):
self.select()
.where(Entity.title == title)
.order_by(func.length(Entity.file_path), Entity.file_path)
.options(*self.get_load_options())
)
result = await self.execute_query(query)
result = await self.execute_query(query, use_query_options=load_relations)
return list(result.scalars().all())
async def get_by_file_path(self, file_path: Union[Path, str]) -> Optional[Entity]:
async def get_by_file_path(
self, file_path: Union[Path, str], *, load_relations: bool = True
) -> Optional[Entity]:
"""Get entity by file_path.
Args:
file_path: Path to the entity file (will be converted to string internally)
"""
query = (
self.select()
.where(Entity.file_path == Path(file_path).as_posix())
.options(*self.get_load_options())
)
return await self.find_one(query)
query = self.select().where(Entity.file_path == Path(file_path).as_posix())
return await self._find_one_by_query(query, load_relations=load_relations)
# -------------------------------------------------------------------------
# Lightweight methods for permalink resolution (no eager loading)
@@ -381,6 +388,9 @@ class EntityRepository(Repository[Entity]):
# Use merge to avoid session state conflicts
# Set the ID to update existing entity
entity.id = existing_entity.id
# Preserve the stable external_id so that external references
# (e.g. public share links) survive re-indexing
entity.external_id = existing_entity.external_id
# Ensure observations reference the correct entity_id
for obs in entity.observations:
@@ -0,0 +1,191 @@
"""Repository for managing note materialization state."""
from pathlib import Path
from typing import Any, Mapping, Optional
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
from basic_memory import db
from basic_memory.models import Entity, NoteContent
from basic_memory.repository.repository import Repository
NOTE_CONTENT_MUTABLE_FIELDS = frozenset(
{
"markdown_content",
"db_version",
"db_checksum",
"file_version",
"file_checksum",
"file_write_status",
"last_source",
"updated_at",
"file_updated_at",
"last_materialization_error",
"last_materialization_attempt_at",
}
)
class NoteContentRepository(Repository[NoteContent]):
"""Repository for project-scoped note materialization state."""
def __init__(self, session_maker: async_sessionmaker[AsyncSession], project_id: int):
"""Initialize with session maker and project-scoped filtering."""
super().__init__(session_maker, NoteContent, project_id=project_id)
def _coerce_note_content(
self, data: Mapping[str, Any] | NoteContent
) -> tuple[NoteContent, set[str]]:
"""Convert input data to a NoteContent model and track explicit fields."""
if isinstance(data, NoteContent):
model_data = {
key: value for key, value in data.__dict__.items() if key in self.valid_columns
}
else:
model_data = {key: value for key, value in data.items() if key in self.valid_columns}
entity_id = model_data.get("entity_id")
if entity_id is None:
raise ValueError("entity_id is required for note_content writes")
return NoteContent(**model_data), set(model_data)
async def _load_entity_identity(self, session: AsyncSession, entity_id: int) -> Entity:
"""Load the owning entity so duplicated identity fields stay aligned."""
result = await session.execute(select(Entity).where(Entity.id == entity_id))
entity = result.scalar_one_or_none()
if entity is None:
raise ValueError(f"Entity {entity_id} does not exist")
if self.project_id is not None and entity.project_id != self.project_id:
raise ValueError(
f"Entity {entity_id} belongs to project {entity.project_id}, "
f"not repository project {self.project_id}"
)
return entity
async def _align_identity_fields(
self, session: AsyncSession, note_content: NoteContent
) -> None:
"""Mirror project identity from entity before persisting note content."""
entity = await self._load_entity_identity(session, note_content.entity_id)
note_content.project_id = entity.project_id
note_content.external_id = entity.external_id
note_content.file_path = Path(entity.file_path).as_posix()
async def get_by_entity_id(self, entity_id: int) -> Optional[NoteContent]:
"""Get note content by the owning entity identifier."""
return await self.find_by_id(entity_id)
async def get_by_external_id(self, external_id: str) -> Optional[NoteContent]:
"""Get note content by the mirrored entity external identifier."""
query = self.select().where(NoteContent.external_id == external_id)
return await self.find_one(query)
async def get_by_file_path(self, file_path: Path | str) -> Optional[NoteContent]:
"""Get note content by file path, preferring rows whose entity still owns that path."""
normalized_path = Path(file_path).as_posix()
# Trigger: note_content mirrors entity.file_path but does not enforce project-level uniqueness.
# Why: entity renames can leave stale mirrored paths behind until note_content realigns.
# Outcome: prefer the row whose current entity path still matches, then the newest mirror.
query = (
self.select()
.join(Entity, Entity.id == NoteContent.entity_id)
.where(NoteContent.file_path == normalized_path)
.order_by(
(Entity.file_path == normalized_path).desc(),
NoteContent.updated_at.desc(),
NoteContent.entity_id.desc(),
)
.limit(1)
.options(*self.get_load_options())
)
async with db.scoped_session(self.session_maker) as session:
result = await session.execute(query)
return result.scalars().first()
async def create(self, data: Mapping[str, Any] | NoteContent) -> NoteContent:
"""Create a note_content row aligned to its owning entity."""
note_content, _ = self._coerce_note_content(data)
async with db.scoped_session(self.session_maker) as session:
await self._align_identity_fields(session, note_content)
session.add(note_content)
await session.flush()
created = await self.select_by_id(session, note_content.entity_id)
if created is None: # pragma: no cover
raise ValueError(
f"Can't find NoteContent for entity {note_content.entity_id} after add"
)
return created
async def upsert(self, data: Mapping[str, Any] | NoteContent) -> NoteContent:
"""Insert or update note_content while keeping mirrored identity fields in sync."""
note_content, provided_fields = self._coerce_note_content(data)
async with db.scoped_session(self.session_maker) as session:
await self._align_identity_fields(session, note_content)
existing = await self.select_by_id(session, note_content.entity_id)
if existing is None:
session.add(note_content)
await session.flush()
created = await self.select_by_id(session, note_content.entity_id)
if created is None: # pragma: no cover
raise ValueError(
f"Can't find NoteContent for entity {note_content.entity_id} after upsert"
)
return created
fields_to_update = (provided_fields - {"entity_id"}) | {
"project_id",
"external_id",
"file_path",
}
for column_name in fields_to_update:
setattr(existing, column_name, getattr(note_content, column_name))
await session.flush()
updated = await self.select_by_id(session, existing.entity_id)
if updated is None: # pragma: no cover
raise ValueError(
f"Can't find NoteContent for entity {existing.entity_id} after upsert"
)
return updated
async def update_state_fields(self, entity_id: int, **updates: Any) -> Optional[NoteContent]:
"""Update sync fields and re-align project_id, external_id, and file_path from entity."""
invalid_fields = set(updates) - NOTE_CONTENT_MUTABLE_FIELDS
if invalid_fields:
invalid_list = ", ".join(sorted(invalid_fields))
raise ValueError(f"Unsupported note_content update fields: {invalid_list}")
async with db.scoped_session(self.session_maker) as session:
note_content = await self.select_by_id(session, entity_id)
if note_content is None:
return None
await self._align_identity_fields(session, note_content)
for field_name, value in updates.items():
setattr(note_content, field_name, value)
await session.flush()
updated = await self.select_by_id(session, entity_id)
if updated is None: # pragma: no cover
raise ValueError(f"Can't find NoteContent for entity {entity_id} after update")
return updated
async def delete_by_entity_id(self, entity_id: int) -> bool:
"""Delete note_content by entity identifier."""
async with db.scoped_session(self.session_maker) as session:
note_content = await self.select_by_id(session, entity_id)
if note_content is None:
return False
await session.delete(note_content)
return True
@@ -3,8 +3,9 @@
import asyncio
import json
import re
import time
from datetime import datetime
from typing import List, Optional
from typing import List, Optional, cast
from loguru import logger
from sqlalchemy import text
@@ -15,7 +16,10 @@ from basic_memory.config import BasicMemoryConfig, ConfigManager
from basic_memory.repository.embedding_provider import EmbeddingProvider
from basic_memory.repository.embedding_provider_factory import create_embedding_provider
from basic_memory.repository.search_index_row import SearchIndexRow
from basic_memory.repository.search_repository_base import SearchRepositoryBase
from basic_memory.repository.search_repository_base import (
SearchRepositoryBase,
_PreparedEntityVectorSync,
)
from basic_memory.repository.metadata_filters import parse_metadata_filters
from basic_memory.repository.semantic_errors import SemanticDependenciesMissingError
from basic_memory.schemas.search import SearchItemType, SearchRetrievalMode
@@ -61,6 +65,9 @@ class PostgresSearchRepository(SearchRepositoryBase):
self._semantic_embedding_sync_batch_size = (
self._app_config.semantic_embedding_sync_batch_size
)
self._semantic_postgres_prepare_concurrency = (
self._app_config.semantic_postgres_prepare_concurrency
)
self._embedding_provider = embedding_provider
self._vector_dimensions = 384
self._vector_tables_initialized = False
@@ -295,6 +302,8 @@ class PostgresSearchRepository(SearchRepositoryBase):
chunk_key TEXT NOT NULL,
chunk_text TEXT NOT NULL,
source_hash TEXT NOT NULL,
entity_fingerprint TEXT NOT NULL,
embedding_model TEXT NOT NULL,
updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
UNIQUE (project_id, entity_id, chunk_key)
)
@@ -441,35 +450,349 @@ class PostgresSearchRepository(SearchRepositoryBase):
)
return [dict(row) for row in vector_result.mappings().all()]
def _vector_prepare_window_size(self) -> int:
"""Use a bounded config-driven prepare window for Postgres vector sync."""
return self._semantic_postgres_prepare_concurrency
async def _prepare_entity_vector_jobs_window(
self, entity_ids: list[int]
) -> list[_PreparedEntityVectorSync | BaseException]:
"""Prepare one Postgres window concurrently to hide DB round-trip latency."""
prepared_window = await asyncio.gather(
*(self._prepare_entity_vector_jobs(entity_id) for entity_id in entity_ids),
return_exceptions=True,
)
return [
cast(_PreparedEntityVectorSync | BaseException, prepared)
for prepared in prepared_window
]
async def _prepare_entity_vector_jobs(self, entity_id: int) -> _PreparedEntityVectorSync:
"""Prepare chunk mutations with Postgres-specific bulk upserts."""
sync_start = time.perf_counter()
logger.info(
"Vector sync start: project_id={project_id} entity_id={entity_id}",
project_id=self.project_id,
entity_id=entity_id,
)
async with db.scoped_session(self.session_maker) as session:
await self._prepare_vector_session(session)
row_result = await session.execute(
text(
"SELECT id, type, title, permalink, content_stems, content_snippet, "
"category, relation_type "
"FROM search_index "
"WHERE entity_id = :entity_id AND project_id = :project_id "
"ORDER BY "
"CASE type "
"WHEN :entity_type THEN 0 "
"WHEN :observation_type THEN 1 "
"WHEN :relation_type_type THEN 2 "
"ELSE 3 END, id ASC"
),
{
"entity_id": entity_id,
"project_id": self.project_id,
"entity_type": SearchItemType.ENTITY.value,
"observation_type": SearchItemType.OBSERVATION.value,
"relation_type_type": SearchItemType.RELATION.value,
},
)
rows = row_result.fetchall()
source_rows_count = len(rows)
if not rows:
logger.info(
"Vector sync source prepared: project_id={project_id} entity_id={entity_id} "
"source_rows_count={source_rows_count} built_chunk_records_count=0",
project_id=self.project_id,
entity_id=entity_id,
source_rows_count=source_rows_count,
)
await self._delete_entity_chunks(session, entity_id)
await session.commit()
prepare_seconds = time.perf_counter() - sync_start
return _PreparedEntityVectorSync(
entity_id=entity_id,
sync_start=sync_start,
source_rows_count=source_rows_count,
embedding_jobs=[],
prepare_seconds=prepare_seconds,
)
chunk_records = self._build_chunk_records(rows)
built_chunk_records_count = len(chunk_records)
current_entity_fingerprint = self._build_entity_fingerprint(chunk_records)
current_embedding_model = self._embedding_model_key()
logger.info(
"Vector sync source prepared: project_id={project_id} entity_id={entity_id} "
"source_rows_count={source_rows_count} "
"built_chunk_records_count={built_chunk_records_count}",
project_id=self.project_id,
entity_id=entity_id,
source_rows_count=source_rows_count,
built_chunk_records_count=built_chunk_records_count,
)
if not chunk_records:
await self._delete_entity_chunks(session, entity_id)
await session.commit()
prepare_seconds = time.perf_counter() - sync_start
return _PreparedEntityVectorSync(
entity_id=entity_id,
sync_start=sync_start,
source_rows_count=source_rows_count,
embedding_jobs=[],
prepare_seconds=prepare_seconds,
)
existing_rows_result = await session.execute(
text(
"SELECT c.id, c.chunk_key, c.source_hash, c.entity_fingerprint, "
"c.embedding_model, "
"(e.chunk_id IS NOT NULL) AS has_embedding "
"FROM search_vector_chunks c "
"LEFT JOIN search_vector_embeddings e ON e.chunk_id = c.id "
"WHERE c.project_id = :project_id AND c.entity_id = :entity_id"
),
{"project_id": self.project_id, "entity_id": entity_id},
)
existing_rows = existing_rows_result.mappings().all()
existing_by_key = {str(row["chunk_key"]): row for row in existing_rows}
existing_chunks_count = len(existing_by_key)
incoming_chunk_keys = {record["chunk_key"] for record in chunk_records}
stale_ids = [
int(row["id"])
for chunk_key, row in existing_by_key.items()
if chunk_key not in incoming_chunk_keys
]
stale_chunks_count = len(stale_ids)
if stale_ids:
await self._delete_stale_chunks(session, stale_ids, entity_id)
orphan_ids = {int(row["id"]) for row in existing_rows if not bool(row["has_embedding"])}
orphan_chunks_count = len(orphan_ids)
skip_unchanged_entity = (
existing_chunks_count == built_chunk_records_count
and stale_chunks_count == 0
and orphan_chunks_count == 0
and existing_chunks_count > 0
and all(
row["entity_fingerprint"] == current_entity_fingerprint
and row["embedding_model"] == current_embedding_model
for row in existing_rows
)
)
if skip_unchanged_entity:
logger.info(
"Vector sync skipped unchanged entity: project_id={project_id} "
"entity_id={entity_id} chunks_skipped={chunks_skipped} "
"entity_fingerprint={entity_fingerprint} embedding_model={embedding_model}",
project_id=self.project_id,
entity_id=entity_id,
chunks_skipped=built_chunk_records_count,
entity_fingerprint=current_entity_fingerprint,
embedding_model=current_embedding_model,
)
prepare_seconds = time.perf_counter() - sync_start
return _PreparedEntityVectorSync(
entity_id=entity_id,
sync_start=sync_start,
source_rows_count=source_rows_count,
embedding_jobs=[],
chunks_total=built_chunk_records_count,
chunks_skipped=built_chunk_records_count,
entity_skipped=True,
prepare_seconds=prepare_seconds,
)
pending_records: list[dict[str, str]] = []
skipped_chunks_count = 0
for record in chunk_records:
current = existing_by_key.get(record["chunk_key"])
if current is None:
pending_records.append(record)
continue
row_id = int(current["id"])
is_orphan = row_id in orphan_ids
same_source_hash = current["source_hash"] == record["source_hash"]
same_entity_fingerprint = (
current["entity_fingerprint"] == current_entity_fingerprint
)
same_embedding_model = current["embedding_model"] == current_embedding_model
if same_source_hash and not is_orphan and same_embedding_model:
if not same_entity_fingerprint:
await session.execute(
text(
"UPDATE search_vector_chunks "
"SET entity_fingerprint = :entity_fingerprint, "
"embedding_model = :embedding_model, "
"updated_at = NOW() "
"WHERE id = :id"
),
{
"id": row_id,
"entity_fingerprint": current_entity_fingerprint,
"embedding_model": current_embedding_model,
},
)
skipped_chunks_count += 1
continue
pending_records.append(record)
shard_plan = self._plan_entity_vector_shard(pending_records)
self._log_vector_shard_plan(entity_id=entity_id, shard_plan=shard_plan)
scheduled_records = [
record
for record in sorted(pending_records, key=lambda record: record["chunk_key"])
if record["chunk_key"] in shard_plan.scheduled_chunk_keys
]
embedding_jobs: list[tuple[int, str]] = []
upsert_records = list(scheduled_records)
if upsert_records:
upsert_params: dict[str, object] = {
"project_id": self.project_id,
"entity_id": entity_id,
}
upsert_values: list[str] = []
# The SQL template is built from integer enumerate() indices only.
# No user-controlled text is interpolated into the statement.
for index, record in enumerate(upsert_records):
upsert_params[f"chunk_key_{index}"] = record["chunk_key"]
upsert_params[f"chunk_text_{index}"] = record["chunk_text"]
upsert_params[f"source_hash_{index}"] = record["source_hash"]
upsert_params[f"entity_fingerprint_{index}"] = current_entity_fingerprint
upsert_params[f"embedding_model_{index}"] = current_embedding_model
upsert_values.append(
"("
":entity_id, :project_id, "
f":chunk_key_{index}, :chunk_text_{index}, :source_hash_{index}, "
f":entity_fingerprint_{index}, :embedding_model_{index}, NOW()"
")"
)
upsert_result = await session.execute(
text(f"""
INSERT INTO search_vector_chunks (
entity_id,
project_id,
chunk_key,
chunk_text,
source_hash,
entity_fingerprint,
embedding_model,
updated_at
) VALUES {", ".join(upsert_values)}
ON CONFLICT (project_id, entity_id, chunk_key) DO UPDATE SET
chunk_text = EXCLUDED.chunk_text,
source_hash = EXCLUDED.source_hash,
entity_fingerprint = EXCLUDED.entity_fingerprint,
embedding_model = EXCLUDED.embedding_model,
updated_at = NOW()
RETURNING id, chunk_key
"""),
upsert_params,
)
upserted_ids_by_key = {
str(row["chunk_key"]): int(row["id"]) for row in upsert_result.mappings().all()
}
for record in upsert_records:
row_id = upserted_ids_by_key[record["chunk_key"]]
embedding_jobs.append((row_id, record["chunk_text"]))
logger.info(
"Vector sync diff complete: project_id={project_id} entity_id={entity_id} "
"existing_chunks_count={existing_chunks_count} "
"stale_chunks_count={stale_chunks_count} "
"orphan_chunks_count={orphan_chunks_count} "
"chunks_skipped={chunks_skipped} "
"embedding_jobs_count={embedding_jobs_count} "
"pending_jobs_total={pending_jobs_total} shard_index={shard_index} "
"shard_count={shard_count} remaining_jobs_after_shard={remaining_jobs_after_shard} "
"oversized_entity={oversized_entity} entity_complete={entity_complete}",
project_id=self.project_id,
entity_id=entity_id,
existing_chunks_count=existing_chunks_count,
stale_chunks_count=stale_chunks_count,
orphan_chunks_count=orphan_chunks_count,
chunks_skipped=skipped_chunks_count,
embedding_jobs_count=len(embedding_jobs),
pending_jobs_total=shard_plan.pending_jobs_total,
shard_index=shard_plan.shard_index,
shard_count=shard_plan.shard_count,
remaining_jobs_after_shard=shard_plan.remaining_jobs_after_shard,
oversized_entity=shard_plan.oversized_entity,
entity_complete=shard_plan.entity_complete,
)
prepare_seconds = time.perf_counter() - sync_start
return _PreparedEntityVectorSync(
entity_id=entity_id,
sync_start=sync_start,
source_rows_count=source_rows_count,
embedding_jobs=embedding_jobs,
chunks_total=built_chunk_records_count,
chunks_skipped=skipped_chunks_count,
entity_complete=shard_plan.entity_complete,
oversized_entity=shard_plan.oversized_entity,
pending_jobs_total=shard_plan.pending_jobs_total,
shard_index=shard_plan.shard_index,
shard_count=shard_plan.shard_count,
remaining_jobs_after_shard=shard_plan.remaining_jobs_after_shard,
prepare_seconds=prepare_seconds,
)
async def _write_embeddings(
self,
session: AsyncSession,
jobs: list[tuple[int, str]],
embeddings: list[list[float]],
) -> None:
for (row_id, _), vector in zip(jobs, embeddings, strict=True):
vector_literal = self._format_pgvector_literal(vector)
await session.execute(
text(
"INSERT INTO search_vector_embeddings ("
"chunk_id, project_id, embedding, embedding_dims, updated_at"
") VALUES ("
":chunk_id, :project_id, CAST(:embedding AS vector), :embedding_dims, NOW()"
") "
"ON CONFLICT (chunk_id) DO UPDATE SET "
"project_id = EXCLUDED.project_id, "
"embedding = EXCLUDED.embedding, "
"embedding_dims = EXCLUDED.embedding_dims, "
"updated_at = NOW()"
),
{
"chunk_id": row_id,
"project_id": self.project_id,
"embedding": vector_literal,
"embedding_dims": len(vector),
},
params: dict[str, object] = {"project_id": self.project_id}
value_rows: list[str] = []
# The SQL template is built from integer enumerate() indices only.
# No user-controlled text is interpolated into the statement.
for index, ((row_id, _), vector) in enumerate(zip(jobs, embeddings, strict=True)):
params[f"chunk_id_{index}"] = row_id
params[f"embedding_{index}"] = self._format_pgvector_literal(vector)
params[f"embedding_dims_{index}"] = len(vector)
value_rows.append(
"("
f":chunk_id_{index}, :project_id, CAST(:embedding_{index} AS vector), "
f":embedding_dims_{index}, NOW()"
")"
)
await session.execute(
text(f"""
INSERT INTO search_vector_embeddings (
chunk_id,
project_id,
embedding,
embedding_dims,
updated_at
) VALUES {", ".join(value_rows)}
ON CONFLICT (chunk_id) DO UPDATE SET
project_id = EXCLUDED.project_id,
embedding = EXCLUDED.embedding,
embedding_dims = EXCLUDED.embedding_dims,
updated_at = NOW()
"""),
params,
)
async def _delete_entity_chunks(
self,
session: AsyncSession,
@@ -506,9 +829,6 @@ class PostgresSearchRepository(SearchRepositoryBase):
stale_params,
)
async def _update_timestamp_sql(self) -> str:
return "NOW()" # pragma: no cover
def _distance_to_similarity(self, distance: float) -> float:
"""Convert pgvector cosine distance to cosine similarity.
@@ -2,6 +2,7 @@
import hashlib
import json
import math
import re
import time
from abc import ABC, abstractmethod
@@ -33,6 +34,7 @@ TOP_CHUNKS_PER_RESULT = 5
SMALL_NOTE_CONTENT_LIMIT = 2000
HEADER_LINE_PATTERN = re.compile(r"^\s*#{1,6}\s+")
BULLET_PATTERN = re.compile(r"^[\-\*]\s+")
OVERSIZED_ENTITY_VECTOR_SHARD_SIZE = 256
@dataclass
@@ -42,8 +44,14 @@ class VectorSyncBatchResult:
entities_total: int
entities_synced: int
entities_failed: int
entities_deferred: int = 0
entities_skipped: int = 0
failed_entity_ids: list[int] = field(default_factory=list)
chunks_total: int = 0
chunks_skipped: int = 0
embedding_jobs_total: int = 0
prepare_seconds_total: float = 0.0
queue_wait_seconds_total: float = 0.0
embed_seconds_total: float = 0.0
write_seconds_total: float = 0.0
@@ -56,6 +64,16 @@ class _PreparedEntityVectorSync:
sync_start: float
source_rows_count: int
embedding_jobs: list[tuple[int, str]]
chunks_total: int = 0
chunks_skipped: int = 0
entity_skipped: bool = False
entity_complete: bool = True
oversized_entity: bool = False
pending_jobs_total: int = 0
shard_index: int = 1
shard_count: int = 1
remaining_jobs_after_shard: int = 0
prepare_seconds: float = 0.0
@dataclass
@@ -75,10 +93,33 @@ class _EntitySyncRuntime:
source_rows_count: int
embedding_jobs_count: int
remaining_jobs: int
chunks_total: int = 0
chunks_skipped: int = 0
entity_skipped: bool = False
entity_complete: bool = True
oversized_entity: bool = False
pending_jobs_total: int = 0
shard_index: int = 1
shard_count: int = 1
remaining_jobs_after_shard: int = 0
prepare_seconds: float = 0.0
embed_seconds: float = 0.0
write_seconds: float = 0.0
@dataclass(frozen=True)
class _EntityVectorShardPlan:
"""Shard selection for one entity's pending embedding work."""
scheduled_chunk_keys: set[str]
pending_jobs_total: int
shard_index: int
shard_count: int
remaining_jobs_after_shard: int
oversized_entity: bool
entity_complete: bool
class SearchRepositoryBase(ABC):
"""Abstract base class for backend-specific search repository implementations.
@@ -247,11 +288,6 @@ class SearchRepositoryBase(ABC):
"""Delete stale chunk rows (and their embeddings) by ID."""
pass
@abstractmethod
async def _update_timestamp_sql(self) -> str:
"""Return the SQL expression for current timestamp in the backend."""
pass # pragma: no cover
@abstractmethod
def _distance_to_similarity(self, distance: float) -> float:
"""Convert a backend-specific vector distance to cosine similarity in [0, 1].
@@ -482,6 +518,108 @@ class SearchRepositoryBase(ABC):
return list(records_by_key.values())
def _build_entity_fingerprint(self, chunk_records: list[dict[str, str]]) -> str:
"""Hash the semantic chunk inputs for one entity.
Trigger: vector sync eligibility depends on the chunk records derived
from search_index rows, not raw file bytes.
Why: title/permalink/observation metadata can change vector inputs even
when unrelated file bytes do not, and vice versa.
Outcome: one deterministic fingerprint invalidates the entity-level skip
whenever the embeddable chunk set changes.
"""
canonical_records = [
{
"chunk_key": record["chunk_key"],
"source_hash": record["source_hash"],
}
for record in sorted(chunk_records, key=lambda record: record["chunk_key"])
]
payload = json.dumps(canonical_records, separators=(",", ":"), sort_keys=True)
return hashlib.sha256(payload.encode("utf-8")).hexdigest()
def _embedding_model_key(self) -> str:
"""Build a stable model identity for vector invalidation checks."""
assert self._embedding_provider is not None
return (
f"{type(self._embedding_provider).__name__}:"
f"{self._embedding_provider.model_name}:"
f"{self._embedding_provider.dimensions}"
)
def _plan_entity_vector_shard(
self,
pending_records: list[dict[str, str]],
) -> _EntityVectorShardPlan:
"""Select the bounded shard to process for one entity sync invocation."""
pending_jobs_total = len(pending_records)
if pending_jobs_total == 0:
return _EntityVectorShardPlan(
scheduled_chunk_keys=set(),
pending_jobs_total=0,
shard_index=1,
shard_count=1,
remaining_jobs_after_shard=0,
oversized_entity=False,
entity_complete=True,
)
ordered_pending_records = sorted(pending_records, key=lambda record: record["chunk_key"])
scheduled_records = ordered_pending_records[:OVERSIZED_ENTITY_VECTOR_SHARD_SIZE]
remaining_jobs_after_shard = pending_jobs_total - len(scheduled_records)
return _EntityVectorShardPlan(
scheduled_chunk_keys={record["chunk_key"] for record in scheduled_records},
pending_jobs_total=pending_jobs_total,
shard_index=1,
shard_count=max(
1,
math.ceil(pending_jobs_total / OVERSIZED_ENTITY_VECTOR_SHARD_SIZE),
),
remaining_jobs_after_shard=remaining_jobs_after_shard,
oversized_entity=pending_jobs_total > OVERSIZED_ENTITY_VECTOR_SHARD_SIZE,
entity_complete=remaining_jobs_after_shard == 0,
)
def _log_vector_shard_plan(
self,
*,
entity_id: int,
shard_plan: _EntityVectorShardPlan,
) -> None:
"""Emit shard planning logs once the pending work is known."""
if shard_plan.pending_jobs_total == 0:
return
scheduled_jobs_count = shard_plan.pending_jobs_total - shard_plan.remaining_jobs_after_shard
if shard_plan.oversized_entity:
logger.warning(
"Vector sync oversized entity detected: project_id={project_id} "
"entity_id={entity_id} pending_jobs_total={pending_jobs_total} "
"shard_size={shard_size} shard_count={shard_count}",
project_id=self.project_id,
entity_id=entity_id,
pending_jobs_total=shard_plan.pending_jobs_total,
shard_size=OVERSIZED_ENTITY_VECTOR_SHARD_SIZE,
shard_count=shard_plan.shard_count,
)
logger.info(
"Vector sync shard planned: project_id={project_id} entity_id={entity_id} "
"pending_jobs_total={pending_jobs_total} scheduled_jobs_count={scheduled_jobs_count} "
"shard_index={shard_index} shard_count={shard_count} "
"remaining_jobs_after_shard={remaining_jobs_after_shard} "
"oversized_entity={oversized_entity} entity_complete={entity_complete}",
project_id=self.project_id,
entity_id=entity_id,
pending_jobs_total=shard_plan.pending_jobs_total,
scheduled_jobs_count=scheduled_jobs_count,
shard_index=shard_plan.shard_index,
shard_count=shard_plan.shard_count,
remaining_jobs_after_shard=shard_plan.remaining_jobs_after_shard,
oversized_entity=shard_plan.oversized_entity,
entity_complete=shard_plan.entity_complete,
)
# --- Text splitting ---
def _split_text_into_chunks(self, text_value: str) -> list[str]:
@@ -664,91 +802,139 @@ class SearchRepositoryBase(ABC):
logger.info(
"Vector batch sync start: project_id={project_id} entities_total={entities_total} "
"sync_batch_size={sync_batch_size}",
"sync_batch_size={sync_batch_size} prepare_window_size={prepare_window_size}",
project_id=self.project_id,
entities_total=total_entities,
sync_batch_size=self._semantic_embedding_sync_batch_size,
prepare_window_size=self._vector_prepare_window_size(),
)
pending_jobs: list[_PendingEmbeddingJob] = []
entity_runtime: dict[int, _EntitySyncRuntime] = {}
failed_entity_ids: set[int] = set()
deferred_entity_ids: set[int] = set()
synced_entity_ids: set[int] = set()
for index, entity_id in enumerate(entity_ids):
prepare_window_size = self._vector_prepare_window_size()
for window_start in range(0, total_entities, prepare_window_size):
window_entity_ids = entity_ids[window_start : window_start + prepare_window_size]
if progress_callback is not None:
progress_callback(entity_id, index, total_entities)
# Trigger: Postgres prepares one bounded entity window concurrently.
# Why: callbacks still need per-entity progress positions before the gather starts.
# Outcome: progress advances in prepare_window_size bursts instead of strict one-by-one.
for offset, entity_id in enumerate(window_entity_ids, start=window_start):
progress_callback(entity_id, offset, total_entities)
try:
prepared = await self._prepare_entity_vector_jobs(entity_id)
except Exception as exc:
if not continue_on_error:
raise
failed_entity_ids.add(entity_id)
logger.warning(
"Vector batch sync entity prepare failed: project_id={project_id} "
"entity_id={entity_id} error={error}",
project_id=self.project_id,
entity_id=entity_id,
error=str(exc),
)
continue
prepared_window = await self._prepare_entity_vector_jobs_window(window_entity_ids)
embedding_jobs_count = len(prepared.embedding_jobs)
result.embedding_jobs_total += embedding_jobs_count
if embedding_jobs_count == 0:
synced_entity_ids.add(entity_id)
total_seconds = time.perf_counter() - prepared.sync_start
self._log_vector_sync_complete(
entity_id=entity_id,
total_seconds=total_seconds,
embed_seconds=0.0,
write_seconds=0.0,
source_rows_count=prepared.source_rows_count,
embedding_jobs_count=0,
)
continue
entity_runtime[entity_id] = _EntitySyncRuntime(
sync_start=prepared.sync_start,
source_rows_count=prepared.source_rows_count,
embedding_jobs_count=embedding_jobs_count,
remaining_jobs=embedding_jobs_count,
)
pending_jobs.extend(
_PendingEmbeddingJob(
entity_id=entity_id, chunk_row_id=row_id, chunk_text=chunk_text
)
for row_id, chunk_text in prepared.embedding_jobs
)
while len(pending_jobs) >= self._semantic_embedding_sync_batch_size:
flush_jobs = pending_jobs[: self._semantic_embedding_sync_batch_size]
pending_jobs = pending_jobs[self._semantic_embedding_sync_batch_size :]
try:
embed_seconds, write_seconds = await self._flush_embedding_jobs(
flush_jobs=flush_jobs,
entity_runtime=entity_runtime,
synced_entity_ids=synced_entity_ids,
)
result.embed_seconds_total += embed_seconds
result.write_seconds_total += write_seconds
except Exception as exc:
for entity_id, prepared in zip(window_entity_ids, prepared_window, strict=True):
if isinstance(prepared, BaseException):
if not continue_on_error:
raise
affected_entity_ids = sorted({job.entity_id for job in flush_jobs})
failed_entity_ids.update(affected_entity_ids)
for failed_entity_id in affected_entity_ids:
entity_runtime.pop(failed_entity_id, None)
raise prepared
failed_entity_ids.add(entity_id)
logger.warning(
"Vector batch sync flush failed: project_id={project_id} "
"affected_entities={affected_entities} chunk_count={chunk_count} error={error}",
"Vector batch sync entity prepare failed: project_id={project_id} "
"entity_id={entity_id} error={error}",
project_id=self.project_id,
affected_entities=affected_entity_ids,
chunk_count=len(flush_jobs),
error=str(exc),
entity_id=entity_id,
error=str(prepared),
)
continue
embedding_jobs_count = len(prepared.embedding_jobs)
result.chunks_total += prepared.chunks_total
result.chunks_skipped += prepared.chunks_skipped
if prepared.entity_skipped:
result.entities_skipped += 1
result.embedding_jobs_total += embedding_jobs_count
result.prepare_seconds_total += prepared.prepare_seconds
if embedding_jobs_count == 0:
if prepared.entity_complete:
synced_entity_ids.add(entity_id)
else:
deferred_entity_ids.add(entity_id)
total_seconds = time.perf_counter() - prepared.sync_start
queue_wait_seconds = max(0.0, total_seconds - prepared.prepare_seconds)
result.queue_wait_seconds_total += queue_wait_seconds
self._log_vector_sync_complete(
entity_id=entity_id,
total_seconds=total_seconds,
prepare_seconds=prepared.prepare_seconds,
queue_wait_seconds=queue_wait_seconds,
embed_seconds=0.0,
write_seconds=0.0,
source_rows_count=prepared.source_rows_count,
chunks_total=prepared.chunks_total,
chunks_skipped=prepared.chunks_skipped,
embedding_jobs_count=0,
entity_skipped=prepared.entity_skipped,
entity_complete=prepared.entity_complete,
oversized_entity=prepared.oversized_entity,
pending_jobs_total=prepared.pending_jobs_total,
shard_index=prepared.shard_index,
shard_count=prepared.shard_count,
remaining_jobs_after_shard=prepared.remaining_jobs_after_shard,
)
continue
entity_runtime[entity_id] = _EntitySyncRuntime(
sync_start=prepared.sync_start,
source_rows_count=prepared.source_rows_count,
embedding_jobs_count=embedding_jobs_count,
remaining_jobs=embedding_jobs_count,
chunks_total=prepared.chunks_total,
chunks_skipped=prepared.chunks_skipped,
entity_skipped=prepared.entity_skipped,
entity_complete=prepared.entity_complete,
oversized_entity=prepared.oversized_entity,
pending_jobs_total=prepared.pending_jobs_total,
shard_index=prepared.shard_index,
shard_count=prepared.shard_count,
remaining_jobs_after_shard=prepared.remaining_jobs_after_shard,
prepare_seconds=prepared.prepare_seconds,
)
pending_jobs.extend(
_PendingEmbeddingJob(
entity_id=entity_id, chunk_row_id=row_id, chunk_text=chunk_text
)
for row_id, chunk_text in prepared.embedding_jobs
)
while len(pending_jobs) >= self._semantic_embedding_sync_batch_size:
flush_jobs = pending_jobs[: self._semantic_embedding_sync_batch_size]
pending_jobs = pending_jobs[self._semantic_embedding_sync_batch_size :]
try:
embed_seconds, write_seconds = await self._flush_embedding_jobs(
flush_jobs=flush_jobs,
entity_runtime=entity_runtime,
synced_entity_ids=synced_entity_ids,
)
result.embed_seconds_total += embed_seconds
result.write_seconds_total += write_seconds
(result.queue_wait_seconds_total) += self._finalize_completed_entity_syncs(
entity_runtime=entity_runtime,
synced_entity_ids=synced_entity_ids,
deferred_entity_ids=deferred_entity_ids,
)
except Exception as exc:
if not continue_on_error:
raise
affected_entity_ids = sorted({job.entity_id for job in flush_jobs})
failed_entity_ids.update(affected_entity_ids)
synced_entity_ids.difference_update(affected_entity_ids)
deferred_entity_ids.difference_update(affected_entity_ids)
for failed_entity_id in affected_entity_ids:
entity_runtime.pop(failed_entity_id, None)
logger.warning(
"Vector batch sync flush failed: project_id={project_id} "
"affected_entities={affected_entities} chunk_count={chunk_count} error={error}",
project_id=self.project_id,
affected_entities=affected_entity_ids,
chunk_count=len(flush_jobs),
error=str(exc),
)
if pending_jobs:
flush_jobs = list(pending_jobs)
@@ -761,11 +947,18 @@ class SearchRepositoryBase(ABC):
)
result.embed_seconds_total += embed_seconds
result.write_seconds_total += write_seconds
(result.queue_wait_seconds_total) += self._finalize_completed_entity_syncs(
entity_runtime=entity_runtime,
synced_entity_ids=synced_entity_ids,
deferred_entity_ids=deferred_entity_ids,
)
except Exception as exc:
if not continue_on_error:
raise
affected_entity_ids = sorted({job.entity_id for job in flush_jobs})
failed_entity_ids.update(affected_entity_ids)
synced_entity_ids.difference_update(affected_entity_ids)
deferred_entity_ids.difference_update(affected_entity_ids)
for failed_entity_id in affected_entity_ids:
entity_runtime.pop(failed_entity_id, None)
logger.warning(
@@ -783,6 +976,8 @@ class SearchRepositoryBase(ABC):
if entity_runtime:
orphan_runtime_entities = sorted(entity_runtime.keys())
failed_entity_ids.update(orphan_runtime_entities)
synced_entity_ids.difference_update(orphan_runtime_entities)
deferred_entity_ids.difference_update(orphan_runtime_entities)
logger.warning(
"Vector batch sync left unfinished entities after flushes: "
"project_id={project_id} unfinished_entities={unfinished_entities}",
@@ -792,26 +987,59 @@ class SearchRepositoryBase(ABC):
# Keep result counters aligned with successful/failed terminal states.
synced_entity_ids.difference_update(failed_entity_ids)
deferred_entity_ids.difference_update(failed_entity_ids)
deferred_entity_ids.difference_update(synced_entity_ids)
result.failed_entity_ids = sorted(failed_entity_ids)
result.entities_failed = len(result.failed_entity_ids)
result.entities_deferred = len(deferred_entity_ids)
result.entities_synced = len(synced_entity_ids)
logger.info(
"Vector batch sync complete: project_id={project_id} entities_total={entities_total} "
"entities_synced={entities_synced} entities_failed={entities_failed} "
"embedding_jobs_total={embedding_jobs_total} embed_seconds_total={embed_seconds_total:.3f} "
"write_seconds_total={write_seconds_total:.3f}",
"entities_deferred={entities_deferred} "
"entities_skipped={entities_skipped} chunks_total={chunks_total} "
"chunks_skipped={chunks_skipped} embedding_jobs_total={embedding_jobs_total} "
"prepare_seconds_total={prepare_seconds_total:.3f} "
"queue_wait_seconds_total={queue_wait_seconds_total:.3f} "
"embed_seconds_total={embed_seconds_total:.3f} write_seconds_total={write_seconds_total:.3f}",
project_id=self.project_id,
entities_total=result.entities_total,
entities_synced=result.entities_synced,
entities_failed=result.entities_failed,
entities_deferred=result.entities_deferred,
entities_skipped=result.entities_skipped,
chunks_total=result.chunks_total,
chunks_skipped=result.chunks_skipped,
embedding_jobs_total=result.embedding_jobs_total,
prepare_seconds_total=result.prepare_seconds_total,
queue_wait_seconds_total=result.queue_wait_seconds_total,
embed_seconds_total=result.embed_seconds_total,
write_seconds_total=result.write_seconds_total,
)
return result
def _vector_prepare_window_size(self) -> int:
"""Return the number of entities to prepare in one orchestration window."""
return 1
async def _prepare_entity_vector_jobs_window(
self, entity_ids: list[int]
) -> list[_PreparedEntityVectorSync | BaseException]:
"""Prepare one window of entity vector jobs.
Default implementation is sequential to preserve backend behavior.
Postgres overrides this to use a bounded concurrent gather.
"""
prepared_window: list[_PreparedEntityVectorSync | BaseException] = []
for entity_id in entity_ids:
try:
prepared_window.append(await self._prepare_entity_vector_jobs(entity_id))
except Exception as exc:
prepared_window.append(exc)
return prepared_window
async def _prepare_entity_vector_jobs(self, entity_id: int) -> _PreparedEntityVectorSync:
"""Prepare chunk mutations and embedding jobs for one entity."""
sync_start = time.perf_counter()
@@ -863,15 +1091,19 @@ class SearchRepositoryBase(ABC):
)
await self._delete_entity_chunks(session, entity_id)
await session.commit()
prepare_seconds = time.perf_counter() - sync_start
return _PreparedEntityVectorSync(
entity_id=entity_id,
sync_start=sync_start,
source_rows_count=source_rows_count,
embedding_jobs=[],
prepare_seconds=prepare_seconds,
)
chunk_records = self._build_chunk_records(rows)
built_chunk_records_count = len(chunk_records)
current_entity_fingerprint = self._build_entity_fingerprint(chunk_records)
current_embedding_model = self._embedding_model_key()
logger.info(
"Vector sync source prepared: project_id={project_id} entity_id={entity_id} "
"source_rows_count={source_rows_count} "
@@ -884,17 +1116,19 @@ class SearchRepositoryBase(ABC):
if not chunk_records:
await self._delete_entity_chunks(session, entity_id)
await session.commit()
prepare_seconds = time.perf_counter() - sync_start
return _PreparedEntityVectorSync(
entity_id=entity_id,
sync_start=sync_start,
source_rows_count=source_rows_count,
embedding_jobs=[],
prepare_seconds=prepare_seconds,
)
# --- Diff existing chunks against incoming ---
existing_rows_result = await session.execute(
text(
"SELECT id, chunk_key, source_hash "
"SELECT id, chunk_key, source_hash, entity_fingerprint, embedding_model "
"FROM search_vector_chunks "
"WHERE project_id = :project_id AND entity_id = :entity_id"
),
@@ -927,9 +1161,48 @@ class SearchRepositoryBase(ABC):
orphan_ids = {int(row.id) for row in orphan_rows}
orphan_chunks_count = len(orphan_ids)
# --- Upsert changed / new chunks, collect embedding jobs ---
# Trigger: the persisted chunk metadata exactly matches the current
# semantic fingerprint/model and every chunk still has an embedding.
# Why: full reindex and embeddings-only runs should avoid reopening
# the expensive chunk diff + embed path for unchanged entities.
# Outcome: return immediately with skip counters and no writes.
skip_unchanged_entity = (
existing_chunks_count == built_chunk_records_count
and stale_chunks_count == 0
and orphan_chunks_count == 0
and existing_chunks_count > 0
and all(
row.entity_fingerprint == current_entity_fingerprint
and row.embedding_model == current_embedding_model
for row in existing_by_key.values()
)
)
if skip_unchanged_entity:
logger.info(
"Vector sync skipped unchanged entity: project_id={project_id} "
"entity_id={entity_id} chunks_skipped={chunks_skipped} "
"entity_fingerprint={entity_fingerprint} embedding_model={embedding_model}",
project_id=self.project_id,
entity_id=entity_id,
chunks_skipped=built_chunk_records_count,
entity_fingerprint=current_entity_fingerprint,
embedding_model=current_embedding_model,
)
prepare_seconds = time.perf_counter() - sync_start
return _PreparedEntityVectorSync(
entity_id=entity_id,
sync_start=sync_start,
source_rows_count=source_rows_count,
embedding_jobs=[],
chunks_total=built_chunk_records_count,
chunks_skipped=built_chunk_records_count,
entity_skipped=True,
prepare_seconds=prepare_seconds,
)
timestamp_expr = self._timestamp_now_expr()
embedding_jobs: list[tuple[int, str]] = []
pending_records: list[dict[str, str]] = []
skipped_chunks_count = 0
for record in chunk_records:
current = existing_by_key.get(record["chunk_key"])
@@ -937,16 +1210,64 @@ class SearchRepositoryBase(ABC):
# but chunk has no embedding (orphan from crash).
# Outcome: schedule re-embedding without touching chunk metadata.
is_orphan = current and int(current.id) in orphan_ids
if current and current.source_hash == record["source_hash"] and not is_orphan:
continue
if current:
row_id = int(current.id)
if current.source_hash != record["source_hash"]:
same_source_hash = current.source_hash == record["source_hash"]
same_entity_fingerprint = (
current.entity_fingerprint == current_entity_fingerprint
)
same_embedding_model = current.embedding_model == current_embedding_model
if same_source_hash and not is_orphan and same_embedding_model:
if not same_entity_fingerprint:
await session.execute(
text(
"UPDATE search_vector_chunks "
"SET entity_fingerprint = :entity_fingerprint, "
"embedding_model = :embedding_model, "
f"updated_at = {timestamp_expr} "
"WHERE id = :id"
),
{
"id": row_id,
"entity_fingerprint": current_entity_fingerprint,
"embedding_model": current_embedding_model,
},
)
skipped_chunks_count += 1
continue
pending_records.append(record)
shard_plan = self._plan_entity_vector_shard(pending_records)
self._log_vector_shard_plan(entity_id=entity_id, shard_plan=shard_plan)
# Trigger: oversized entities can accumulate thousands of pending chunks.
# Why: scheduling only one deterministic shard bounds memory and wall clock.
# Outcome: future runs resume from the remaining chunk rows without redoing completed work.
scheduled_records = [
record
for record in sorted(pending_records, key=lambda record: record["chunk_key"])
if record["chunk_key"] in shard_plan.scheduled_chunk_keys
]
# --- Upsert scheduled changed / new chunks, collect embedding jobs ---
embedding_jobs: list[tuple[int, str]] = []
for record in scheduled_records:
current = existing_by_key.get(record["chunk_key"])
if current:
row_id = int(current.id)
if (
current.source_hash != record["source_hash"]
or current.entity_fingerprint != current_entity_fingerprint
or current.embedding_model != current_embedding_model
):
await session.execute(
text(
"UPDATE search_vector_chunks "
"SET chunk_text = :chunk_text, source_hash = :source_hash, "
"entity_fingerprint = :entity_fingerprint, "
"embedding_model = :embedding_model, "
f"updated_at = {timestamp_expr} "
"WHERE id = :id"
),
@@ -954,6 +1275,8 @@ class SearchRepositoryBase(ABC):
"id": row_id,
"chunk_text": record["chunk_text"],
"source_hash": record["source_hash"],
"entity_fingerprint": current_entity_fingerprint,
"embedding_model": current_embedding_model,
},
)
embedding_jobs.append((row_id, record["chunk_text"]))
@@ -962,9 +1285,11 @@ class SearchRepositoryBase(ABC):
inserted = await session.execute(
text(
"INSERT INTO search_vector_chunks ("
"entity_id, project_id, chunk_key, chunk_text, source_hash, updated_at"
"entity_id, project_id, chunk_key, chunk_text, source_hash, "
"entity_fingerprint, embedding_model, updated_at"
") VALUES ("
f":entity_id, :project_id, :chunk_key, :chunk_text, :source_hash, "
":entity_fingerprint, :embedding_model, "
f"{timestamp_expr}"
") RETURNING id"
),
@@ -974,6 +1299,8 @@ class SearchRepositoryBase(ABC):
"chunk_key": record["chunk_key"],
"chunk_text": record["chunk_text"],
"source_hash": record["source_hash"],
"entity_fingerprint": current_entity_fingerprint,
"embedding_model": current_embedding_model,
},
)
row_id = int(inserted.scalar_one())
@@ -984,21 +1311,42 @@ class SearchRepositoryBase(ABC):
"existing_chunks_count={existing_chunks_count} "
"stale_chunks_count={stale_chunks_count} "
"orphan_chunks_count={orphan_chunks_count} "
"embedding_jobs_count={embedding_jobs_count}",
"chunks_skipped={chunks_skipped} "
"embedding_jobs_count={embedding_jobs_count} "
"pending_jobs_total={pending_jobs_total} shard_index={shard_index} "
"shard_count={shard_count} remaining_jobs_after_shard={remaining_jobs_after_shard} "
"oversized_entity={oversized_entity} entity_complete={entity_complete}",
project_id=self.project_id,
entity_id=entity_id,
existing_chunks_count=existing_chunks_count,
stale_chunks_count=stale_chunks_count,
orphan_chunks_count=orphan_chunks_count,
chunks_skipped=skipped_chunks_count,
embedding_jobs_count=len(embedding_jobs),
pending_jobs_total=shard_plan.pending_jobs_total,
shard_index=shard_plan.shard_index,
shard_count=shard_plan.shard_count,
remaining_jobs_after_shard=shard_plan.remaining_jobs_after_shard,
oversized_entity=shard_plan.oversized_entity,
entity_complete=shard_plan.entity_complete,
)
await session.commit()
prepare_seconds = time.perf_counter() - sync_start
return _PreparedEntityVectorSync(
entity_id=entity_id,
sync_start=sync_start,
source_rows_count=source_rows_count,
embedding_jobs=embedding_jobs,
chunks_total=built_chunk_records_count,
chunks_skipped=skipped_chunks_count,
entity_complete=shard_plan.entity_complete,
oversized_entity=shard_plan.oversized_entity,
pending_jobs_total=shard_plan.pending_jobs_total,
shard_index=shard_plan.shard_index,
shard_count=shard_plan.shard_count,
remaining_jobs_after_shard=shard_plan.remaining_jobs_after_shard,
prepare_seconds=prepare_seconds,
)
async def _flush_embedding_jobs(
@@ -1061,58 +1409,140 @@ class SearchRepositoryBase(ABC):
runtime.embed_seconds += embed_seconds * flush_share
runtime.write_seconds += write_seconds * flush_share
if runtime.remaining_jobs <= 0:
if runtime.remaining_jobs <= 0 and runtime.entity_complete:
synced_entity_ids.add(entity_id)
total_seconds = time.perf_counter() - runtime.sync_start
self._log_vector_sync_complete(
entity_id=entity_id,
total_seconds=total_seconds,
embed_seconds=runtime.embed_seconds,
write_seconds=runtime.write_seconds,
source_rows_count=runtime.source_rows_count,
embedding_jobs_count=runtime.embedding_jobs_count,
)
entity_runtime.pop(entity_id, None)
return embed_seconds, write_seconds
def _finalize_completed_entity_syncs(
self,
*,
entity_runtime: dict[int, _EntitySyncRuntime],
synced_entity_ids: set[int],
deferred_entity_ids: set[int],
) -> float:
"""Finalize completed entities and return cumulative queue wait seconds."""
queue_wait_seconds_total = 0.0
for entity_id, runtime in list(entity_runtime.items()):
if runtime.remaining_jobs > 0:
continue
if runtime.entity_complete:
synced_entity_ids.add(entity_id)
else:
deferred_entity_ids.add(entity_id)
total_seconds = time.perf_counter() - runtime.sync_start
queue_wait_seconds = max(
0.0,
total_seconds
- runtime.prepare_seconds
- runtime.embed_seconds
- runtime.write_seconds,
)
queue_wait_seconds_total += queue_wait_seconds
self._log_vector_sync_complete(
entity_id=entity_id,
total_seconds=total_seconds,
prepare_seconds=runtime.prepare_seconds,
queue_wait_seconds=queue_wait_seconds,
embed_seconds=runtime.embed_seconds,
write_seconds=runtime.write_seconds,
source_rows_count=runtime.source_rows_count,
chunks_total=runtime.chunks_total,
chunks_skipped=runtime.chunks_skipped,
embedding_jobs_count=runtime.embedding_jobs_count,
entity_skipped=runtime.entity_skipped,
entity_complete=runtime.entity_complete,
oversized_entity=runtime.oversized_entity,
pending_jobs_total=runtime.pending_jobs_total,
shard_index=runtime.shard_index,
shard_count=runtime.shard_count,
remaining_jobs_after_shard=runtime.remaining_jobs_after_shard,
)
entity_runtime.pop(entity_id, None)
return queue_wait_seconds_total
def _log_vector_sync_complete(
self,
*,
entity_id: int,
total_seconds: float,
prepare_seconds: float,
queue_wait_seconds: float,
embed_seconds: float,
write_seconds: float,
source_rows_count: int,
chunks_total: int,
chunks_skipped: int,
embedding_jobs_count: int,
entity_skipped: bool,
entity_complete: bool,
oversized_entity: bool,
pending_jobs_total: int,
shard_index: int,
shard_count: int,
remaining_jobs_after_shard: int,
) -> None:
"""Log completion and slow-entity warnings with a consistent format."""
logger.info(
"Vector sync complete: project_id={project_id} entity_id={entity_id} "
"total_seconds={total_seconds:.3f} embed_seconds={embed_seconds:.3f} "
"total_seconds={total_seconds:.3f} prepare_seconds={prepare_seconds:.3f} "
"queue_wait_seconds={queue_wait_seconds:.3f} embed_seconds={embed_seconds:.3f} "
"write_seconds={write_seconds:.3f} source_rows_count={source_rows_count} "
"embedding_jobs_count={embedding_jobs_count}",
"chunks_total={chunks_total} chunks_skipped={chunks_skipped} "
"embedding_jobs_count={embedding_jobs_count} entity_skipped={entity_skipped} "
"entity_complete={entity_complete} oversized_entity={oversized_entity} "
"pending_jobs_total={pending_jobs_total} shard_index={shard_index} "
"shard_count={shard_count} remaining_jobs_after_shard={remaining_jobs_after_shard}",
project_id=self.project_id,
entity_id=entity_id,
total_seconds=total_seconds,
prepare_seconds=prepare_seconds,
queue_wait_seconds=queue_wait_seconds,
embed_seconds=embed_seconds,
write_seconds=write_seconds,
source_rows_count=source_rows_count,
chunks_total=chunks_total,
chunks_skipped=chunks_skipped,
embedding_jobs_count=embedding_jobs_count,
entity_skipped=entity_skipped,
entity_complete=entity_complete,
oversized_entity=oversized_entity,
pending_jobs_total=pending_jobs_total,
shard_index=shard_index,
shard_count=shard_count,
remaining_jobs_after_shard=remaining_jobs_after_shard,
)
if total_seconds > 10:
logger.warning(
"Vector sync slow entity: project_id={project_id} entity_id={entity_id} "
"total_seconds={total_seconds:.3f} embed_seconds={embed_seconds:.3f} "
"total_seconds={total_seconds:.3f} prepare_seconds={prepare_seconds:.3f} "
"queue_wait_seconds={queue_wait_seconds:.3f} embed_seconds={embed_seconds:.3f} "
"write_seconds={write_seconds:.3f} source_rows_count={source_rows_count} "
"embedding_jobs_count={embedding_jobs_count}",
"chunks_total={chunks_total} chunks_skipped={chunks_skipped} "
"embedding_jobs_count={embedding_jobs_count} entity_skipped={entity_skipped} "
"entity_complete={entity_complete} oversized_entity={oversized_entity} "
"pending_jobs_total={pending_jobs_total} shard_index={shard_index} "
"shard_count={shard_count} remaining_jobs_after_shard={remaining_jobs_after_shard}",
project_id=self.project_id,
entity_id=entity_id,
total_seconds=total_seconds,
prepare_seconds=prepare_seconds,
queue_wait_seconds=queue_wait_seconds,
embed_seconds=embed_seconds,
write_seconds=write_seconds,
source_rows_count=source_rows_count,
chunks_total=chunks_total,
chunks_skipped=chunks_skipped,
embedding_jobs_count=embedding_jobs_count,
entity_skipped=entity_skipped,
entity_complete=entity_complete,
oversized_entity=oversized_entity,
pending_jobs_total=pending_jobs_total,
shard_index=shard_index,
shard_count=shard_count,
remaining_jobs_after_shard=remaining_jobs_after_shard,
)
async def _prepare_vector_session(self, session: AsyncSession) -> None:
@@ -398,10 +398,16 @@ class SQLiteSearchRepository(SearchRepositoryBase):
"chunk_key",
"chunk_text",
"source_hash",
"entity_fingerprint",
"embedding_model",
"updated_at",
}
schema_mismatch = bool(chunks_columns) and set(chunks_columns) != expected_columns
if schema_mismatch:
# Trigger: older SQLite installs are missing newly required chunk metadata columns.
# Why: vector tables store derived data only, so rebuilding them is safer than
# attempting piecemeal ALTER TABLE compatibility across sqlite-vec upgrades.
# Outcome: first startup after the schema change forces a clean re-embed.
logger.warning("search_vector_chunks schema mismatch, recreating vector tables")
await session.execute(text("DROP TABLE IF EXISTS search_vector_embeddings"))
await session.execute(text("DROP TABLE IF EXISTS search_vector_chunks"))
@@ -552,9 +558,6 @@ class SQLiteSearchRepository(SearchRepositoryBase):
stale_params,
)
async def _update_timestamp_sql(self) -> str:
return "CURRENT_TIMESTAMP" # pragma: no cover
def _distance_to_similarity(self, distance: float) -> float:
"""Convert L2 distance to cosine similarity for normalized embeddings.
+107 -34
View File
@@ -242,9 +242,17 @@ class EntityService(BaseService[EntityModel]):
# Try to find existing entity using strict resolution (no fuzzy search)
# This prevents incorrectly matching similar file paths like "Node A.md" and "Node C.md"
existing = await self.link_resolver.resolve_link(schema.file_path, strict=True)
existing = await self.link_resolver.resolve_link(
schema.file_path,
strict=True,
load_relations=False,
)
if not existing and schema.permalink:
existing = await self.link_resolver.resolve_link(schema.permalink, strict=True)
existing = await self.link_resolver.resolve_link(
schema.permalink,
strict=True,
load_relations=False,
)
if existing:
logger.debug(f"Found existing entity: {existing.file_path}")
@@ -863,10 +871,22 @@ class EntityService(BaseService[EntityModel]):
"""
logger.debug(f"Updating entity and observations: {file_path}")
db_entity = await self.repository.get_by_file_path(file_path.as_posix())
with telemetry.scope(
"upsert.update.fetch_entity",
domain="entity_service",
action="upsert",
phase="fetch_entity",
):
db_entity = await self.repository.get_by_file_path(file_path.as_posix())
# Clear observations for entity
await self.observation_repository.delete_by_fields(entity_id=db_entity.id)
with telemetry.scope(
"upsert.update.delete_observations",
domain="entity_service",
action="upsert",
phase="delete_observations",
):
await self.observation_repository.delete_by_fields(entity_id=db_entity.id)
# add new observations
observations = [
@@ -880,7 +900,14 @@ class EntityService(BaseService[EntityModel]):
)
for obs in markdown.observations
]
await self.observation_repository.add_all(observations)
with telemetry.scope(
"upsert.update.insert_observations",
domain="entity_service",
action="upsert",
phase="insert_observations",
count=len(observations),
):
await self.observation_repository.add_all(observations)
# update values from markdown
db_entity = entity_model_from_markdown(file_path, markdown, db_entity)
@@ -894,10 +921,16 @@ class EntityService(BaseService[EntityModel]):
db_entity.last_updated_by = user_id
# update entity
return await self.repository.update(
db_entity.id,
db_entity,
)
with telemetry.scope(
"upsert.update.save_entity",
domain="entity_service",
action="upsert",
phase="save_entity",
):
return await self.repository.update(
db_entity.id,
db_entity,
)
async def upsert_entity_from_markdown(
self,
@@ -911,20 +944,30 @@ class EntityService(BaseService[EntityModel]):
created = await self.create_entity_from_markdown(file_path, markdown)
else:
created = await self.update_entity_and_observations(file_path, markdown)
return await self.update_entity_relations(created.file_path, markdown)
# Pass entity directly — avoids redundant get_by_file_path inside update_entity_relations
return await self.update_entity_relations(created, markdown)
async def update_entity_relations(
self,
path: str,
entity: EntityModel,
markdown: EntityMarkdown,
) -> EntityModel:
"""Update relations for entity"""
logger.debug(f"Updating relations for entity: {path}")
"""Update relations for entity.
db_entity = await self.repository.get_by_file_path(path)
Accepts the entity object directly to avoid a redundant DB fetch.
Only entity.id and entity.permalink are used from the passed-in object.
"""
entity_id = entity.id
logger.debug(f"Updating relations for entity: {entity.file_path}")
# Clear existing relations first
await self.relation_repository.delete_outgoing_relations_from_entity(db_entity.id)
with telemetry.scope(
"upsert.relations.delete_existing",
domain="entity_service",
action="upsert",
phase="delete_relations",
):
await self.relation_repository.delete_outgoing_relations_from_entity(entity_id)
# Batch resolve all relation targets in parallel
if markdown.relations:
@@ -934,12 +977,23 @@ class EntityService(BaseService[EntityModel]):
# Use strict=True to disable fuzzy search - only exact matches should create resolved relations
# This ensures forward references (links to non-existent entities) remain unresolved (to_id=NULL)
lookup_tasks = [
self.link_resolver.resolve_link(rel.target, strict=True)
self.link_resolver.resolve_link(
rel.target,
strict=True,
load_relations=False,
)
for rel in markdown.relations
]
# Execute all lookups in parallel
resolved_entities = await asyncio.gather(*lookup_tasks, return_exceptions=True)
with telemetry.scope(
"upsert.relations.resolve_links",
domain="entity_service",
action="upsert",
phase="resolve_links",
count=len(lookup_tasks),
):
resolved_entities = await asyncio.gather(*lookup_tasks, return_exceptions=True)
# Process results and create relation records
relations_to_add = []
@@ -958,7 +1012,7 @@ class EntityService(BaseService[EntityModel]):
# Create the relation
relation = Relation(
project_id=self.relation_repository.project_id,
from_id=db_entity.id,
from_id=entity_id,
to_id=target_id,
to_name=target_name,
relation_type=rel.type,
@@ -968,22 +1022,37 @@ class EntityService(BaseService[EntityModel]):
# Batch insert all relations
if relations_to_add:
try:
await self.relation_repository.add_all(relations_to_add)
except IntegrityError:
# Some relations might be duplicates - fall back to individual inserts
logger.debug("Batch relation insert failed, trying individual inserts")
for relation in relations_to_add:
try:
await self.relation_repository.add(relation)
except IntegrityError:
# Unique constraint violation - relation already exists
logger.debug(
f"Skipping duplicate relation {relation.relation_type} from {db_entity.permalink}"
)
continue
with telemetry.scope(
"upsert.relations.insert_relations",
domain="entity_service",
action="upsert",
phase="insert_relations",
count=len(relations_to_add),
):
try:
await self.relation_repository.add_all(relations_to_add)
except IntegrityError:
# Some relations might be duplicates - fall back to individual inserts
logger.debug("Batch relation insert failed, trying individual inserts")
for relation in relations_to_add:
try:
await self.relation_repository.add(relation)
except IntegrityError:
# Unique constraint violation - relation already exists
logger.debug(
f"Skipping duplicate relation {relation.relation_type} from {entity.permalink}"
)
continue
return await self.repository.get_by_file_path(path)
# Reload entity with relations via PK lookup (faster than get_by_file_path string match)
with telemetry.scope(
"upsert.relations.reload_entity",
domain="entity_service",
action="upsert",
phase="reload_entity",
):
reloaded = await self.repository.find_by_ids([entity_id])
return reloaded[0]
async def edit_entity(
self,
@@ -1040,7 +1109,11 @@ class EntityService(BaseService[EntityModel]):
action="edit",
phase="resolve_entity",
):
entity = await self.link_resolver.resolve_link(identifier, strict=True)
entity = await self.link_resolver.resolve_link(
identifier,
strict=True,
load_relations=False,
)
if not entity:
raise EntityNotFoundError(f"Entity not found: {identifier}")
+46 -10
View File
@@ -47,6 +47,7 @@ class LinkResolver:
use_search: bool = True,
strict: bool = False,
source_path: Optional[str] = None,
load_relations: bool = True,
) -> Optional[Entity]:
"""Resolve a markdown link to a permalink.
@@ -56,6 +57,7 @@ class LinkResolver:
strict: If True, only exact matches are allowed (no fuzzy search fallback)
source_path: Optional path of the source file containing the link.
Used to prefer notes closer to the source (context-aware resolution).
load_relations: When False, skip eager loading and return a lightweight entity row.
"""
logger.trace(f"Resolving link: {link_text} (source: {source_path})")
@@ -70,7 +72,10 @@ class LinkResolver:
# UUIDs also match the stored external_id values.
try:
canonical_id = str(uuid_mod.UUID(clean_text))
entity = await self.entity_repository.get_by_external_id(canonical_id)
entity = await self.entity_repository.get_by_external_id(
canonical_id,
load_relations=load_relations,
)
if entity:
logger.debug(f"Found entity by external_id: {entity.permalink}")
return entity
@@ -98,6 +103,7 @@ class LinkResolver:
strict=strict,
source_path=None,
project_permalink=project.permalink,
load_relations=load_relations,
)
current_project_permalink = await self._get_current_project_permalink()
@@ -109,6 +115,7 @@ class LinkResolver:
strict=strict,
source_path=source_path,
project_permalink=current_project_permalink,
load_relations=load_relations,
)
if resolved:
return resolved
@@ -136,6 +143,7 @@ class LinkResolver:
strict=strict,
source_path=None,
project_permalink=project.permalink,
load_relations=load_relations,
)
def _normalize_link_text(self, link_text: str) -> Tuple[str, Optional[str]]:
@@ -176,6 +184,7 @@ class LinkResolver:
strict: bool,
source_path: Optional[str],
project_permalink: Optional[str],
load_relations: bool,
) -> Optional[Entity]:
"""Resolve a link within a specific project scope."""
clean_text = link_text
@@ -223,12 +232,18 @@ class LinkResolver:
# Try with .md extension
if not relative_path.endswith(".md"):
relative_path_md = f"{relative_path}.md"
entity = await entity_repository.get_by_file_path(relative_path_md)
entity = await entity_repository.get_by_file_path(
relative_path_md,
load_relations=load_relations,
)
if entity:
return entity
# Try as-is (already has extension or is a permalink)
entity = await entity_repository.get_by_file_path(relative_path)
entity = await entity_repository.get_by_file_path(
relative_path,
load_relations=load_relations,
)
if entity:
return entity
@@ -242,12 +257,18 @@ class LinkResolver:
# Check permalink match
for candidate_permalink in permalink_candidates:
permalink_entity = await entity_repository.get_by_permalink(candidate_permalink)
permalink_entity = await entity_repository.get_by_permalink(
candidate_permalink,
load_relations=load_relations,
)
if permalink_entity and permalink_entity.id not in [c.id for c in candidates]:
candidates.append(permalink_entity)
# Check title matches
title_entities = await entity_repository.get_by_title(clean_text)
title_entities = await entity_repository.get_by_title(
clean_text,
load_relations=load_relations,
)
for entity in title_entities:
# Avoid duplicates (permalink match might also be in title matches)
if entity.id not in [c.id for c in candidates]:
@@ -263,13 +284,19 @@ class LinkResolver:
# Standard resolution (no source context): permalink first, then title
# 1. Try exact permalink match first (most efficient)
for candidate_permalink in permalink_candidates:
entity = await entity_repository.get_by_permalink(candidate_permalink)
entity = await entity_repository.get_by_permalink(
candidate_permalink,
load_relations=load_relations,
)
if entity:
logger.debug(f"Found exact permalink match: {entity.permalink}")
return entity
# 2. Try exact title match
found = await entity_repository.get_by_title(clean_text)
found = await entity_repository.get_by_title(
clean_text,
load_relations=load_relations,
)
if found:
# Return first match (shortest path) if no source context
entity = found[0]
@@ -277,7 +304,10 @@ class LinkResolver:
return entity
# 3. Try file path
found_path = await entity_repository.get_by_file_path(clean_text)
found_path = await entity_repository.get_by_file_path(
clean_text,
load_relations=load_relations,
)
if found_path:
logger.debug(f"Found entity with path: {found_path.file_path}")
return found_path
@@ -285,7 +315,10 @@ class LinkResolver:
# 4. Try file path with .md extension if not already present
if not clean_text.endswith(".md") and "/" in clean_text:
file_path_with_md = f"{clean_text}.md"
found_path_md = await entity_repository.get_by_file_path(file_path_with_md)
found_path_md = await entity_repository.get_by_file_path(
file_path_with_md,
load_relations=load_relations,
)
if found_path_md:
logger.debug(f"Found entity with path (with .md): {found_path_md.file_path}")
return found_path_md
@@ -309,7 +342,10 @@ class LinkResolver:
f"Selected best match from {len(results)} results: {best_match.permalink}"
)
if best_match.permalink:
return await entity_repository.get_by_permalink(best_match.permalink)
return await entity_repository.get_by_permalink(
best_match.permalink,
load_relations=load_relations,
)
# if we couldn't find anything then return None
return None
+6
View File
@@ -65,6 +65,11 @@ ALL_COMBOS = [
SearchCombo("postgres-openai", DatabaseBackend.POSTGRES, "openai", 1536),
]
# Benchmark queries compare ranking quality across providers rather than enforcing
# the stricter production retrieval cutoff. OpenAI paraphrase matches cluster near
# ~0.37 in this corpus, so the default 0.55 filter hides otherwise-correct results.
BENCHMARK_MIN_SIMILARITY = 0.3
# --- Skip guards ---
@@ -229,6 +234,7 @@ async def create_search_service(
default_project="bench-project",
database_backend=combo.backend,
semantic_search_enabled=semantic_enabled,
semantic_min_similarity=BENCHMARK_MIN_SIMILARITY,
)
# Create search repository (backend-specific)
+1 -1
View File
@@ -51,7 +51,7 @@ RECALL_AT_5_THRESHOLDS: dict[tuple[str, str, str], float] = {
("sqlite-fastembed", "paraphrase", "hybrid"): 0.25,
("postgres-fastembed", "lexical", "hybrid"): 0.37,
("postgres-fastembed", "paraphrase", "hybrid"): 0.25,
# OpenAI hybrid should handle paraphrases better than FastEmbed
# OpenAI hybrid should handle paraphrases better than FastEmbed.
("postgres-openai", "lexical", "hybrid"): 0.37,
("postgres-openai", "paraphrase", "hybrid"): 0.25,
}
@@ -12,7 +12,7 @@ from basic_memory.config import ProjectMode
runner = CliRunner()
def test_cloud_upload_uses_control_plane_client(monkeypatch, tmp_path):
def test_cloud_upload_uses_control_plane_client(monkeypatch, tmp_path, config_manager):
"""Upload command should use control-plane cloud client for WebDAV PUT operations."""
import basic_memory.cli.commands.cloud.upload_command as upload_command
+4 -1
View File
@@ -28,7 +28,10 @@ def test_bm_help_exits_cleanly():
["uv", "run", "bm", "--help"],
capture_output=True,
text=True,
timeout=10,
# Help builds the full command tree, so use a looser timeout than the
# version fast path. This test is guarding against hangs, not enforcing
# a tight performance budget under full-suite load.
timeout=20,
cwd=Path(__file__).parent.parent.parent,
)
assert result.returncode == 0
+76
View File
@@ -139,6 +139,82 @@ def test_project_list_shows_local_cloud_presence_and_routes(
assert "/beta" in result.stdout
def test_project_list_shows_display_name_for_private_projects(
runner: CliRunner, write_config, mock_client, tmp_path, monkeypatch
):
"""Private projects should show display_name ('My Project') instead of raw UUID name."""
private_uuid = "f1df8f39-d5aa-4095-ae05-8c5a2883029a"
write_config(
{
"env": "dev",
"projects": {},
"default_project": "main",
"cloud_api_key": "bmc_test_key_123",
}
)
local_payload = {
"projects": [
{
"id": 1,
"external_id": "11111111-1111-1111-1111-111111111111",
"name": "main",
"path": "/main",
"is_default": True,
}
],
"default_project": "main",
}
cloud_payload = {
"projects": [
{
"id": 1,
"external_id": "11111111-1111-1111-1111-111111111111",
"name": "main",
"path": "/main",
"is_default": True,
},
{
"id": 2,
"external_id": "22222222-2222-2222-2222-222222222222",
"name": private_uuid,
"path": f"/{private_uuid}",
"is_default": False,
"display_name": "My Project",
"is_private": True,
},
],
"default_project": "main",
}
async def fake_list_projects(self):
if os.getenv("BASIC_MEMORY_FORCE_CLOUD", "").lower() in ("true", "1", "yes"):
return ProjectList.model_validate(cloud_payload)
return ProjectList.model_validate(local_payload)
monkeypatch.setattr(ProjectClient, "list_projects", fake_list_projects)
result = runner.invoke(app, ["project", "list"], env={"COLUMNS": "240"})
assert result.exit_code == 0, f"Exit code: {result.exit_code}, output: {result.stdout}"
# Rich table should show display_name in the Name column
assert "My Project" in result.stdout
lines = result.stdout.splitlines()
project_line = next(line for line in lines if "My Project" in line)
name_cell = project_line.split("")[1].strip()
assert name_cell == "My Project"
# JSON output should preserve canonical name for scripting, with display_name as separate field
json_result = runner.invoke(app, ["project", "list", "--json"], env={"COLUMNS": "240"})
assert json_result.exit_code == 0
data = json.loads(json_result.stdout)
private_project = next(p for p in data["projects"] if p.get("display_name") == "My Project")
assert private_project["name"] == private_uuid
assert private_project["display_name"] == "My Project"
def test_project_ls_local_mode_defaults_to_local_route(
runner: CliRunner, write_config, mock_client, tmp_path, monkeypatch
):
+21
View File
@@ -176,6 +176,27 @@ async def test_malformed_frontmatter(tmp_path):
assert entity.frontmatter.permalink is None
@pytest.mark.asyncio
async def test_null_bytes_stripped(tmp_path):
"""Test that null bytes are stripped from content before parsing.
PostgreSQL rejects null bytes (0x00) in text columns. Some files
(e.g. Claude agent definitions) can contain embedded nulls.
"""
content = "---\ntitle: Test\ntype: note\n---\n\nSome content\x00with nulls\x00inside\n"
parser = EntityParser(tmp_path)
entity = await parser.parse_markdown_content(
file_path=tmp_path / "nulls.md",
content=content,
)
assert "\x00" not in entity.content
assert "Some content" in entity.content
assert "with nulls" in entity.content
assert "inside" in entity.content
@pytest.mark.asyncio
async def test_file_not_found():
"""Test handling of non-existent files."""
+60 -2
View File
@@ -28,7 +28,7 @@ async def test_get_client_uses_injected_factory(monkeypatch):
seen = {"used": False}
@asynccontextmanager
async def factory():
async def factory(workspace=None):
seen["used"] = True
async with httpx.AsyncClient(base_url="https://example.test") as client:
yield client
@@ -185,7 +185,7 @@ async def test_get_client_factory_overrides_per_project_routing(config_manager):
config_manager.save_config(cfg)
@asynccontextmanager
async def factory():
async def factory(workspace=None):
async with httpx.AsyncClient(base_url="https://factory.test") as client:
yield client
@@ -316,6 +316,64 @@ async def test_get_cloud_control_plane_client_uses_oauth_token(config_manager):
assert client.headers.get("Authorization") == "Bearer oauth-control-123"
@pytest.mark.asyncio
async def test_get_cloud_control_plane_client_with_workspace(config_manager):
"""Control plane client passes X-Workspace-ID header when workspace is provided."""
cfg = config_manager.load_config()
cfg.cloud_host = "https://cloud.example.test"
cfg.cloud_api_key = "bmc_test_key_123"
config_manager.save_config(cfg)
async with get_cloud_control_plane_client(workspace="tenant-abc") as client:
assert client.headers.get("X-Workspace-ID") == "tenant-abc"
# Without workspace, header should not be present
async with get_cloud_control_plane_client() as client:
assert "X-Workspace-ID" not in client.headers
@pytest.mark.asyncio
async def test_get_client_auto_resolves_workspace_from_project_config(config_manager):
"""get_client resolves workspace from project entry when not explicitly passed."""
cfg = config_manager.load_config()
cfg.cloud_host = "https://cloud.example.test"
cfg.cloud_api_key = "bmc_test_key_123"
cfg.set_project_mode("research", ProjectMode.CLOUD)
cfg.projects["research"].workspace_id = "tenant-from-config"
config_manager.save_config(cfg)
async with get_client(project_name="research") as client:
assert client.headers.get("X-Workspace-ID") == "tenant-from-config"
@pytest.mark.asyncio
async def test_get_client_auto_resolves_workspace_from_default(config_manager):
"""get_client falls back to default_workspace when project has no workspace_id."""
cfg = config_manager.load_config()
cfg.cloud_host = "https://cloud.example.test"
cfg.cloud_api_key = "bmc_test_key_123"
cfg.set_project_mode("research", ProjectMode.CLOUD)
cfg.default_workspace = "default-tenant-456"
config_manager.save_config(cfg)
async with get_client(project_name="research") as client:
assert client.headers.get("X-Workspace-ID") == "default-tenant-456"
@pytest.mark.asyncio
async def test_get_client_explicit_workspace_overrides_config(config_manager):
"""Explicit workspace param takes priority over project config."""
cfg = config_manager.load_config()
cfg.cloud_host = "https://cloud.example.test"
cfg.cloud_api_key = "bmc_test_key_123"
cfg.set_project_mode("research", ProjectMode.CLOUD)
cfg.projects["research"].workspace_id = "tenant-from-config"
config_manager.save_config(cfg)
async with get_client(project_name="research", workspace="explicit-tenant") as client:
assert client.headers.get("X-Workspace-ID") == "explicit-tenant"
@pytest.mark.asyncio
async def test_get_cloud_control_plane_client_raises_without_credentials(config_manager):
cfg = config_manager.load_config()
+1 -1
View File
@@ -709,7 +709,7 @@ class TestGetProjectClientRoutingOrder:
# Set up a factory (simulates what cloud MCP server does)
@asynccontextmanager
async def fake_factory():
async def fake_factory(workspace=None):
from httpx import ASGITransport, AsyncClient
from basic_memory.api.app import app as fastapi_app
@@ -70,6 +70,50 @@ async def test_upsert_entity_same_file_update(entity_repository: EntityRepositor
assert result2.file_path == "test/test-entity.md"
@pytest.mark.asyncio
async def test_upsert_entity_preserves_external_id(entity_repository: EntityRepository):
"""Test that upserting an entity with the same file_path preserves the original external_id.
Trigger: force full re-index creates a new Entity model (with a fresh UUID)
for a file that already has a database record
Why: external_id is used by public share links if it changes, shares break
Outcome: the original external_id survives the upsert
"""
# Create initial entity
entity1 = Entity(
project_id=entity_repository.project_id,
title="Shared Note",
note_type="note",
permalink="test/shared-note",
file_path="test/shared-note.md",
content_type="text/markdown",
external_id="original-stable-id",
created_at=datetime.now(timezone.utc),
updated_at=datetime.now(timezone.utc),
)
result1 = await entity_repository.upsert_entity(entity1)
assert result1.external_id == "original-stable-id"
# Simulate re-index: new Entity model with a DIFFERENT external_id
entity2 = Entity(
project_id=entity_repository.project_id,
title="Shared Note (updated)",
note_type="note",
permalink="test/shared-note",
file_path="test/shared-note.md",
content_type="text/markdown",
external_id="newly-generated-uuid", # would break share links
created_at=datetime.now(timezone.utc),
updated_at=datetime.now(timezone.utc),
)
result2 = await entity_repository.upsert_entity(entity2)
# ID preserved, title updated, external_id stable
assert result2.id == result1.id
assert result2.title == "Shared Note (updated)"
assert result2.external_id == "original-stable-id"
@pytest.mark.asyncio
async def test_upsert_entity_permalink_conflict_different_file(entity_repository: EntityRepository):
"""Test upserting an entity with permalink conflict but different file_path."""
@@ -0,0 +1,425 @@
"""Tests for the NoteContentRepository."""
from datetime import datetime, timedelta, timezone
import pytest
from basic_memory import db
from basic_memory.models import NoteContent, Project
from basic_memory.repository.entity_repository import EntityRepository
from basic_memory.repository.note_content_repository import NoteContentRepository
from basic_memory.repository.project_repository import ProjectRepository
def build_note_content_payload(entity_id: int) -> dict:
"""Build a minimal payload for note_content writes."""
return {
"entity_id": entity_id,
"project_id": -1,
"external_id": "stale-external-id",
"file_path": "stale/path.md",
"markdown_content": "# Materialized content",
"db_version": 1,
"db_checksum": "db-checksum-1",
"file_version": None,
"file_checksum": None,
"file_write_status": "pending",
"last_source": "api",
"updated_at": datetime.now(timezone.utc),
"file_updated_at": None,
"last_materialization_error": None,
"last_materialization_attempt_at": None,
}
@pytest.mark.asyncio
async def test_create_and_lookup_note_content(
session_maker,
test_project: Project,
sample_entity,
):
"""Create note_content and read it back through each supported lookup."""
repository = NoteContentRepository(session_maker, project_id=test_project.id)
created = await repository.create(build_note_content_payload(sample_entity.id))
assert created.entity_id == sample_entity.id
assert created.project_id == sample_entity.project_id
assert created.external_id == sample_entity.external_id
assert created.file_path == sample_entity.file_path
by_entity = await repository.get_by_entity_id(sample_entity.id)
by_external = await repository.get_by_external_id(sample_entity.external_id)
by_path = await repository.get_by_file_path(sample_entity.file_path)
assert by_entity is not None
assert by_external is not None
assert by_path is not None
assert by_entity.entity_id == created.entity_id
assert by_external.entity_id == created.entity_id
assert by_path.entity_id == created.entity_id
@pytest.mark.asyncio
async def test_upsert_updates_existing_note_content(
session_maker,
test_project: Project,
sample_entity,
):
"""Upsert should update the existing row instead of inserting a duplicate."""
repository = NoteContentRepository(session_maker, project_id=test_project.id)
await repository.create(build_note_content_payload(sample_entity.id))
updated_at = datetime.now(timezone.utc)
updated = await repository.upsert(
NoteContent(
entity_id=sample_entity.id,
project_id=test_project.id,
external_id=sample_entity.external_id,
file_path=sample_entity.file_path,
markdown_content="# Updated materialized content",
db_version=2,
db_checksum="db-checksum-2",
file_version=7,
file_checksum="file-checksum-7",
file_write_status="synced",
last_source="reconciler",
updated_at=updated_at,
file_updated_at=updated_at,
last_materialization_error="transient failure",
last_materialization_attempt_at=updated_at,
)
)
assert updated.entity_id == sample_entity.id
assert updated.markdown_content == "# Updated materialized content"
assert updated.db_version == 2
assert updated.db_checksum == "db-checksum-2"
assert updated.file_version == 7
assert updated.file_checksum == "file-checksum-7"
assert updated.file_write_status == "synced"
assert updated.last_source == "reconciler"
assert updated.last_materialization_error == "transient failure"
@pytest.mark.asyncio
async def test_upsert_inserts_when_no_existing_row(
session_maker,
test_project: Project,
sample_entity,
):
"""Upsert should insert a new row when the entity has no note_content yet."""
repository = NoteContentRepository(session_maker, project_id=test_project.id)
created = await repository.upsert(build_note_content_payload(sample_entity.id))
assert created.entity_id == sample_entity.id
assert created.project_id == sample_entity.project_id
assert created.external_id == sample_entity.external_id
assert created.file_path == sample_entity.file_path
assert created.db_version == 1
@pytest.mark.asyncio
async def test_create_requires_entity_id(session_maker, test_project: Project):
"""Create should fail fast when note_content identity is missing."""
repository = NoteContentRepository(session_maker, project_id=test_project.id)
with pytest.raises(ValueError, match="entity_id is required"):
await repository.create({"markdown_content": "# Missing entity"})
@pytest.mark.asyncio
async def test_upsert_preserves_existing_fields_for_partial_payload(
session_maker,
test_project: Project,
sample_entity,
):
"""Partial upserts should only change explicit fields and preserve existing state."""
repository = NoteContentRepository(session_maker, project_id=test_project.id)
payload = build_note_content_payload(sample_entity.id)
payload["last_materialization_error"] = "stale failure"
created = await repository.create(payload)
updated_at = datetime.now(timezone.utc)
updated = await repository.upsert(
{
"entity_id": sample_entity.id,
"markdown_content": "# Partially updated content",
"db_version": 2,
"updated_at": updated_at,
"last_materialization_error": None,
}
)
assert updated.markdown_content == "# Partially updated content"
assert updated.db_version == 2
assert updated.db_checksum == created.db_checksum
assert updated.file_write_status == created.file_write_status
assert updated.last_source == created.last_source
assert updated.last_materialization_error is None
assert updated.file_path == sample_entity.file_path
@pytest.mark.asyncio
async def test_create_rejects_missing_entity(session_maker, test_project: Project):
"""Create should fail when the owning entity does not exist."""
repository = NoteContentRepository(session_maker, project_id=test_project.id)
with pytest.raises(ValueError, match="Entity 999999 does not exist"):
await repository.create(build_note_content_payload(999999))
@pytest.mark.asyncio
async def test_create_rejects_entity_from_another_project(session_maker, config_home):
"""Create should reject note_content writes across project boundaries."""
project_repository = ProjectRepository(session_maker)
project_one = await project_repository.create(
{
"name": "project-one-boundary",
"path": str(config_home / "project-one-boundary"),
"is_active": True,
}
)
project_two = await project_repository.create(
{
"name": "project-two-boundary",
"path": str(config_home / "project-two-boundary"),
"is_active": True,
}
)
entity_repository = EntityRepository(session_maker, project_id=project_two.id)
other_project_entity = await entity_repository.create(
{
"title": "Other Project Note",
"note_type": "test",
"permalink": "project-two/other-project-note",
"file_path": "notes/other-project-note.md",
"content_type": "text/markdown",
"created_at": datetime.now(timezone.utc),
"updated_at": datetime.now(timezone.utc),
}
)
repository = NoteContentRepository(session_maker, project_id=project_one.id)
with pytest.raises(
ValueError,
match=f"Entity {other_project_entity.id} belongs to project {project_two.id}",
):
await repository.create(build_note_content_payload(other_project_entity.id))
@pytest.mark.asyncio
async def test_update_state_fields_realigns_identity_with_entity(
session_maker,
test_project: Project,
sample_entity,
entity_repository: EntityRepository,
):
"""Sync-field updates should refresh mirrored identity from the owning entity."""
repository = NoteContentRepository(session_maker, project_id=test_project.id)
await repository.create(build_note_content_payload(sample_entity.id))
renamed_path = "renamed/test_entity.md"
await entity_repository.update(sample_entity.id, {"file_path": renamed_path})
updated = await repository.update_state_fields(
sample_entity.id,
file_write_status="failed",
file_version=3,
file_checksum="file-checksum-3",
last_materialization_error=None,
last_materialization_attempt_at=None,
)
assert updated is not None
assert updated.file_path == renamed_path
assert updated.external_id == sample_entity.external_id
assert updated.file_write_status == "failed"
assert updated.file_version == 3
assert updated.file_checksum == "file-checksum-3"
assert updated.last_materialization_error is None
assert updated.last_materialization_attempt_at is None
@pytest.mark.asyncio
async def test_update_state_fields_rejects_invalid_fields(
session_maker,
test_project: Project,
sample_entity,
):
"""Only the declared mutable sync fields should be accepted."""
repository = NoteContentRepository(session_maker, project_id=test_project.id)
await repository.create(build_note_content_payload(sample_entity.id))
with pytest.raises(ValueError, match="Unsupported note_content update fields: file_path"):
await repository.update_state_fields(sample_entity.id, file_path="renamed/note.md")
@pytest.mark.asyncio
async def test_update_state_fields_returns_none_for_missing_note_content(
session_maker,
test_project: Project,
):
"""Missing note_content rows should produce a clean None response."""
repository = NoteContentRepository(session_maker, project_id=test_project.id)
assert await repository.update_state_fields(999999, file_write_status="failed") is None
@pytest.mark.asyncio
async def test_delete_by_entity_id(session_maker, test_project: Project, sample_entity):
"""Delete note_content directly by entity identifier."""
repository = NoteContentRepository(session_maker, project_id=test_project.id)
await repository.create(build_note_content_payload(sample_entity.id))
deleted = await repository.delete_by_entity_id(sample_entity.id)
assert deleted is True
assert await repository.get_by_entity_id(sample_entity.id) is None
@pytest.mark.asyncio
async def test_delete_by_entity_id_returns_false_when_missing(
session_maker,
test_project: Project,
):
"""Delete should report False when the note_content row does not exist."""
repository = NoteContentRepository(session_maker, project_id=test_project.id)
assert await repository.delete_by_entity_id(999999) is False
@pytest.mark.asyncio
async def test_note_content_cascades_when_entity_is_deleted(
session_maker,
test_project: Project,
sample_entity,
entity_repository: EntityRepository,
):
"""Deleting the owning entity should cascade to note_content."""
repository = NoteContentRepository(session_maker, project_id=test_project.id)
await repository.create(build_note_content_payload(sample_entity.id))
deleted = await entity_repository.delete(sample_entity.id)
assert deleted is True
assert await repository.get_by_entity_id(sample_entity.id) is None
@pytest.mark.asyncio
async def test_note_content_file_path_lookup_is_project_scoped(session_maker, config_home):
"""Lookups by file_path should respect the repository project scope."""
project_repository = ProjectRepository(session_maker)
project_one = await project_repository.create(
{
"name": "project-one",
"path": str(config_home / "project-one"),
"is_active": True,
}
)
project_two = await project_repository.create(
{
"name": "project-two",
"path": str(config_home / "project-two"),
"is_active": True,
}
)
entity_one_repo = EntityRepository(session_maker, project_id=project_one.id)
entity_two_repo = EntityRepository(session_maker, project_id=project_two.id)
shared_file_path = "shared/note.md"
entity_one = await entity_one_repo.create(
{
"title": "Shared Note",
"note_type": "test",
"permalink": "project-one/shared-note",
"file_path": shared_file_path,
"content_type": "text/markdown",
"created_at": datetime.now(timezone.utc),
"updated_at": datetime.now(timezone.utc),
}
)
entity_two = await entity_two_repo.create(
{
"title": "Shared Note",
"note_type": "test",
"permalink": "project-two/shared-note",
"file_path": shared_file_path,
"content_type": "text/markdown",
"created_at": datetime.now(timezone.utc),
"updated_at": datetime.now(timezone.utc),
}
)
repository_one = NoteContentRepository(session_maker, project_id=project_one.id)
repository_two = NoteContentRepository(session_maker, project_id=project_two.id)
await repository_one.create(build_note_content_payload(entity_one.id))
await repository_two.create(build_note_content_payload(entity_two.id))
found_one = await repository_one.get_by_file_path(shared_file_path)
found_two = await repository_two.get_by_file_path(shared_file_path)
assert found_one is not None
assert found_two is not None
assert found_one.entity_id == entity_one.id
assert found_two.entity_id == entity_two.id
@pytest.mark.asyncio
async def test_note_content_file_path_lookup_prefers_entity_with_current_path(
session_maker,
config_home,
):
"""File-path lookup should prefer the entity whose current path still matches."""
project_repository = ProjectRepository(session_maker)
project = await project_repository.create(
{
"name": "project-path-drift",
"path": str(config_home / "project-path-drift"),
"is_active": True,
}
)
entity_repository = EntityRepository(session_maker, project_id=project.id)
stale_entity = await entity_repository.create(
{
"title": "Stale Note",
"note_type": "test",
"permalink": "project/stale-note",
"file_path": "archived/note.md",
"content_type": "text/markdown",
"created_at": datetime.now(timezone.utc),
"updated_at": datetime.now(timezone.utc),
}
)
current_entity = await entity_repository.create(
{
"title": "Current Note",
"note_type": "test",
"permalink": "project/current-note",
"file_path": "shared/note.md",
"content_type": "text/markdown",
"created_at": datetime.now(timezone.utc),
"updated_at": datetime.now(timezone.utc),
}
)
repository = NoteContentRepository(session_maker, project_id=project.id)
stale_payload = build_note_content_payload(stale_entity.id)
stale_payload["updated_at"] = datetime.now(timezone.utc) + timedelta(minutes=5)
await repository.create(stale_payload)
await repository.create(build_note_content_payload(current_entity.id))
async with db.scoped_session(session_maker) as session:
stale_note_content = await repository.select_by_id(session, stale_entity.id)
assert stale_note_content is not None
stale_note_content.file_path = "shared/note.md"
await session.flush()
found = await repository.get_by_file_path("shared/note.md")
assert found is not None
assert found.entity_id == current_entity.id
@@ -11,6 +11,7 @@ from sqlalchemy import text
from basic_memory import db
from basic_memory.config import BasicMemoryConfig, DatabaseBackend
import basic_memory.repository.search_repository_base as search_repository_base_module
from basic_memory.repository.postgres_search_repository import (
PostgresSearchRepository,
_strip_nul_from_row,
@@ -47,6 +48,19 @@ class StubEmbeddingProvider:
return [0.0, 0.0, 0.0, 1.0]
class StubEmbeddingProviderV2(StubEmbeddingProvider):
"""Same vectors, different model identity to force Postgres resync."""
model_name = "stub-v2"
def _oversized_entity_content(bullet_count: int) -> str:
"""Build deterministic content that produces many vector chunks."""
lines = ["# Oversized Entity"]
lines.extend(f"- embedding job {index}" for index in range(1, bullet_count + 1))
return "\n".join(lines)
async def _skip_if_pgvector_unavailable(session_maker) -> None:
"""Skip semantic pgvector tests when extension is not available in test Postgres image."""
async with db.scoped_session(session_maker) as session:
@@ -442,6 +456,203 @@ async def test_postgres_semantic_hybrid_search_combines_fts_and_vector(session_m
assert any(result.permalink == "specs/search-index" for result in results)
@pytest.mark.asyncio
async def test_postgres_vector_sync_skips_unchanged_and_reembeds_changed_content(
session_maker, test_project
):
"""Postgres vector sync tracks new, changed, unchanged, and model-changed entities."""
await _skip_if_pgvector_unavailable(session_maker)
app_config = BasicMemoryConfig(
env="test",
projects={"test-project": "/tmp/basic-memory-test"},
default_project="test-project",
database_backend=DatabaseBackend.POSTGRES,
semantic_search_enabled=True,
)
repo = PostgresSearchRepository(
session_maker,
project_id=test_project.id,
app_config=app_config,
embedding_provider=StubEmbeddingProvider(),
)
await repo.init_search_index()
now = datetime.now(timezone.utc)
await repo.index_item(
SearchIndexRow(
project_id=test_project.id,
id=421,
title="Auth and Schema Notes",
content_stems="# Overview\n- auth token rotation\n- schema migration planning",
content_snippet="# Overview\n- auth token rotation\n- schema migration planning",
permalink="specs/auth-and-schema",
file_path="specs/auth-and-schema.md",
type=SearchItemType.ENTITY.value,
entity_id=421,
metadata={"note_type": "spec"},
created_at=now,
updated_at=now,
)
)
new_result = await repo.sync_entity_vectors_batch([421])
assert new_result.entities_synced == 1
assert new_result.entities_skipped == 0
assert new_result.chunks_total >= 2
assert new_result.chunks_skipped == 0
assert new_result.embedding_jobs_total == new_result.chunks_total
async with db.scoped_session(session_maker) as session:
stored_rows = await session.execute(
text(
"SELECT entity_fingerprint, embedding_model "
"FROM search_vector_chunks "
"WHERE project_id = :project_id AND entity_id = :entity_id"
),
{"project_id": test_project.id, "entity_id": 421},
)
metadata_rows = stored_rows.fetchall()
assert metadata_rows
assert len({row.entity_fingerprint for row in metadata_rows}) == 1
assert len({row.embedding_model for row in metadata_rows}) == 1
assert metadata_rows[0].embedding_model == "StubEmbeddingProvider:stub:4"
unchanged_result = await repo.sync_entity_vectors_batch([421])
assert unchanged_result.entities_synced == 1
assert unchanged_result.entities_skipped == 1
assert unchanged_result.embedding_jobs_total == 0
assert unchanged_result.chunks_skipped == unchanged_result.chunks_total
await repo.index_item(
SearchIndexRow(
project_id=test_project.id,
id=421,
title="Auth and Schema Notes",
content_stems="# Overview\n- auth token rotation\n- database schema migration planning",
content_snippet="# Overview\n- auth token rotation\n- database schema migration planning",
permalink="specs/auth-and-schema",
file_path="specs/auth-and-schema.md",
type=SearchItemType.ENTITY.value,
entity_id=421,
metadata={"note_type": "spec"},
created_at=now,
updated_at=now,
)
)
changed_result = await repo.sync_entity_vectors_batch([421])
assert changed_result.entities_synced == 1
assert changed_result.entities_skipped == 0
assert changed_result.embedding_jobs_total >= 1
assert changed_result.chunks_skipped >= 1
assert changed_result.embedding_jobs_total < changed_result.chunks_total
repo_v2 = PostgresSearchRepository(
session_maker,
project_id=test_project.id,
app_config=app_config,
embedding_provider=StubEmbeddingProviderV2(),
)
await repo_v2.init_search_index()
model_changed_result = await repo_v2.sync_entity_vectors_batch([421])
assert model_changed_result.entities_synced == 1
assert model_changed_result.entities_skipped == 0
assert model_changed_result.chunks_skipped == 0
assert model_changed_result.embedding_jobs_total == model_changed_result.chunks_total
@pytest.mark.asyncio
async def test_postgres_vector_sync_shards_oversized_entity_and_resumes(
session_maker, test_project, monkeypatch
):
"""Oversized entities should sync one deterministic shard per run and resume cleanly."""
await _skip_if_pgvector_unavailable(session_maker)
monkeypatch.setattr(search_repository_base_module, "OVERSIZED_ENTITY_VECTOR_SHARD_SIZE", 2)
app_config = BasicMemoryConfig(
env="test",
projects={"test-project": "/tmp/basic-memory-test"},
default_project="test-project",
database_backend=DatabaseBackend.POSTGRES,
semantic_search_enabled=True,
)
repo = PostgresSearchRepository(
session_maker,
project_id=test_project.id,
app_config=app_config,
embedding_provider=StubEmbeddingProvider(),
)
await repo.init_search_index()
now = datetime.now(timezone.utc)
content = _oversized_entity_content(5)
await repo.index_item(
SearchIndexRow(
project_id=test_project.id,
id=430,
title="Oversized Vector Entity",
content_stems=content,
content_snippet=content,
permalink="specs/oversized-vector-entity",
file_path="specs/oversized-vector-entity.md",
type=SearchItemType.ENTITY.value,
entity_id=430,
metadata={"note_type": "spec"},
created_at=now,
updated_at=now,
)
)
first_result = await repo.sync_entity_vectors_batch([430])
assert first_result.entities_synced == 0
assert first_result.entities_deferred == 1
assert first_result.entities_failed == 0
assert first_result.embedding_jobs_total == 2
assert first_result.chunks_total == 6
assert first_result.chunks_skipped == 0
second_result = await repo.sync_entity_vectors_batch([430])
assert second_result.entities_synced == 0
assert second_result.entities_deferred == 1
assert second_result.entities_failed == 0
assert second_result.embedding_jobs_total == 2
assert second_result.chunks_total == 6
assert second_result.chunks_skipped == 2
third_result = await repo.sync_entity_vectors_batch([430])
assert third_result.entities_synced == 1
assert third_result.entities_deferred == 0
assert third_result.entities_failed == 0
assert third_result.embedding_jobs_total == 2
assert third_result.chunks_total == 6
assert third_result.chunks_skipped == 4
unchanged_result = await repo.sync_entity_vectors_batch([430])
assert unchanged_result.entities_synced == 1
assert unchanged_result.entities_deferred == 0
assert unchanged_result.entities_skipped == 1
assert unchanged_result.embedding_jobs_total == 0
assert unchanged_result.chunks_skipped == unchanged_result.chunks_total == 6
async with db.scoped_session(session_maker) as session:
chunk_count = await session.execute(
text(
"SELECT COUNT(*) FROM search_vector_chunks "
"WHERE project_id = :project_id AND entity_id = :entity_id"
),
{"project_id": test_project.id, "entity_id": 430},
)
embedding_count = await session.execute(
text(
"SELECT COUNT(*) FROM search_vector_embeddings e "
"JOIN search_vector_chunks c ON c.id = e.chunk_id "
"WHERE c.project_id = :project_id AND c.entity_id = :entity_id"
),
{"project_id": test_project.id, "entity_id": 430},
)
assert int(chunk_count.scalar_one()) == 6
assert int(embedding_count.scalar_one()) == 6
@pytest.mark.asyncio
async def test_postgres_vector_mode_rejects_non_text_query(session_maker, test_project):
"""Vector mode should fail fast for title-only queries."""
@@ -5,12 +5,15 @@ covering utility functions, formatting helpers, and constructor paths that
are difficult to reach in integration tests.
"""
import asyncio
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
import basic_memory.repository.search_repository_base as search_repository_base_module
from basic_memory.config import BasicMemoryConfig, DatabaseBackend
from basic_memory.repository.postgres_search_repository import PostgresSearchRepository
from basic_memory.repository.search_repository_base import _PreparedEntityVectorSync
from basic_memory.repository.semantic_errors import (
SemanticDependenciesMissingError,
SemanticSearchDisabledError,
@@ -37,6 +40,7 @@ def _make_repo(
*,
semantic_enabled: bool = False,
embedding_provider=None,
semantic_postgres_prepare_concurrency: int = 4,
) -> PostgresSearchRepository:
"""Build a PostgresSearchRepository with a no-op session maker."""
session_maker = MagicMock()
@@ -46,6 +50,7 @@ def _make_repo(
default_project="test-project",
database_backend=DatabaseBackend.POSTGRES,
semantic_search_enabled=semantic_enabled,
semantic_postgres_prepare_concurrency=semantic_postgres_prepare_concurrency,
)
return PostgresSearchRepository(
session_maker,
@@ -229,14 +234,196 @@ class TestWriteEmbeddings:
"""Cover _write_embeddings upsert logic."""
@pytest.mark.asyncio
async def test_write_embeddings_executes_per_job(self):
async def test_write_embeddings_executes_single_bulk_upsert(self):
repo = _make_repo()
session = AsyncMock()
jobs = [(100, "chunk text A"), (200, "chunk text B")]
embeddings = [[0.1, 0.2, 0.3, 0.4], [0.5, 0.6, 0.7, 0.8]]
await repo._write_embeddings(session, jobs, embeddings)
assert session.execute.call_count == 2
first_params = session.execute.call_args_list[0][0][1]
assert first_params["chunk_id"] == 100
assert first_params["project_id"] == repo.project_id
assert first_params["embedding_dims"] == 4
assert session.execute.call_count == 1
params = session.execute.call_args[0][1]
assert params["chunk_id_0"] == 100
assert params["chunk_id_1"] == 200
assert params["project_id"] == repo.project_id
assert params["embedding_dims_0"] == 4
assert params["embedding_dims_1"] == 4
class TestBatchPrepareConcurrency:
"""Cover the Postgres-specific concurrent prepare window."""
@pytest.mark.asyncio
async def test_sync_entity_vectors_batch_prepares_entities_concurrently(self, monkeypatch):
repo = _make_repo(
semantic_enabled=True,
embedding_provider=StubEmbeddingProvider(),
semantic_postgres_prepare_concurrency=2,
)
repo._semantic_embedding_sync_batch_size = 8
repo._vector_tables_initialized = True
active_prepares = 0
max_active_prepares = 0
async def _stub_prepare(entity_id: int) -> _PreparedEntityVectorSync:
nonlocal active_prepares, max_active_prepares
active_prepares += 1
max_active_prepares = max(max_active_prepares, active_prepares)
await asyncio.sleep(0)
active_prepares -= 1
return _PreparedEntityVectorSync(
entity_id=entity_id,
sync_start=float(entity_id),
source_rows_count=1,
embedding_jobs=[],
)
monkeypatch.setattr(repo, "_ensure_vector_tables", AsyncMock())
monkeypatch.setattr(repo, "_prepare_entity_vector_jobs", _stub_prepare)
result = await repo.sync_entity_vectors_batch([1, 2, 3, 4])
assert result.entities_total == 4
assert result.entities_synced == 4
assert result.entities_failed == 0
assert max_active_prepares == 2
@pytest.mark.asyncio
async def test_postgres_batch_sync_tracks_prepare_and_queue_wait(monkeypatch):
"""Postgres batch sync should separate queue wait from prepare/embed/write."""
repo = _make_repo(
semantic_enabled=True,
embedding_provider=StubEmbeddingProvider(),
)
repo._semantic_embedding_sync_batch_size = 2
repo._vector_tables_initialized = True
async def _stub_prepare(entity_id: int) -> _PreparedEntityVectorSync:
return _PreparedEntityVectorSync(
entity_id=entity_id,
sync_start=0.0,
source_rows_count=1,
embedding_jobs=[(200 + entity_id, f"chunk-{entity_id}")],
prepare_seconds=1.0,
)
async def _stub_flush(flush_jobs, entity_runtime, synced_entity_ids):
for job in flush_jobs:
runtime = entity_runtime[job.entity_id]
if job.entity_id == 1:
runtime.embed_seconds = 1.0
runtime.write_seconds = 0.5
else:
runtime.embed_seconds = 2.0
runtime.write_seconds = 0.5
runtime.remaining_jobs = 0
synced_entity_ids.add(job.entity_id)
return (3.0, 1.0)
completion_records: list[dict] = []
def _capture_log(**kwargs):
completion_records.append(kwargs)
perf_counter_values = iter([4.0, 5.0])
monkeypatch.setattr(repo, "_ensure_vector_tables", AsyncMock())
monkeypatch.setattr(repo, "_prepare_entity_vector_jobs", _stub_prepare)
monkeypatch.setattr(repo, "_flush_embedding_jobs", _stub_flush)
monkeypatch.setattr(repo, "_log_vector_sync_complete", _capture_log)
monkeypatch.setattr(
search_repository_base_module.time,
"perf_counter",
lambda: next(perf_counter_values),
)
result = await repo.sync_entity_vectors_batch([1, 2])
assert result.entities_total == 2
assert result.entities_synced == 2
assert result.entities_failed == 0
assert result.prepare_seconds_total == pytest.approx(2.0)
assert result.queue_wait_seconds_total == pytest.approx(3.0)
assert result.embed_seconds_total == pytest.approx(3.0)
assert result.write_seconds_total == pytest.approx(1.0)
assert len(completion_records) == 2
for record in completion_records:
assert record["prepare_seconds"] == pytest.approx(1.0)
assert record["queue_wait_seconds"] == pytest.approx(1.5)
@pytest.mark.asyncio
async def test_postgres_batch_sync_tracks_deferred_oversized_entities(monkeypatch):
"""Oversized shard runs should be deferred until the last shard completes."""
repo = _make_repo(
semantic_enabled=True,
embedding_provider=StubEmbeddingProvider(),
)
repo._semantic_embedding_sync_batch_size = 8
repo._vector_tables_initialized = True
async def _stub_prepare(entity_id: int) -> _PreparedEntityVectorSync:
if entity_id == 1:
return _PreparedEntityVectorSync(
entity_id=entity_id,
sync_start=0.0,
source_rows_count=1,
embedding_jobs=[(201, "chunk-1a"), (202, "chunk-1b")],
chunks_total=5,
pending_jobs_total=5,
entity_complete=False,
oversized_entity=True,
shard_index=1,
shard_count=3,
remaining_jobs_after_shard=3,
)
return _PreparedEntityVectorSync(
entity_id=entity_id,
sync_start=0.0,
source_rows_count=1,
embedding_jobs=[(301, "chunk-2a")],
chunks_total=1,
pending_jobs_total=1,
entity_complete=True,
shard_index=1,
shard_count=1,
remaining_jobs_after_shard=0,
)
async def _stub_flush(flush_jobs, entity_runtime, synced_entity_ids):
for job in flush_jobs:
runtime = entity_runtime[job.entity_id]
runtime.remaining_jobs -= 1
runtime.embed_seconds += 0.5
runtime.write_seconds += 0.25
return (1.5, 0.75)
completion_records: list[dict] = []
def _capture_log(**kwargs):
completion_records.append(kwargs)
monkeypatch.setattr(repo, "_ensure_vector_tables", AsyncMock())
monkeypatch.setattr(repo, "_prepare_entity_vector_jobs", _stub_prepare)
monkeypatch.setattr(repo, "_flush_embedding_jobs", _stub_flush)
monkeypatch.setattr(repo, "_log_vector_sync_complete", _capture_log)
result = await repo.sync_entity_vectors_batch([1, 2])
assert result.entities_total == 2
assert result.entities_synced == 1
assert result.entities_deferred == 1
assert result.entities_failed == 0
assert result.embedding_jobs_total == 3
deferred_record = next(record for record in completion_records if record["entity_id"] == 1)
assert deferred_record["entity_complete"] is False
assert deferred_record["oversized_entity"] is True
assert deferred_record["pending_jobs_total"] == 5
assert deferred_record["shard_count"] == 3
assert deferred_record["remaining_jobs_after_shard"] == 3
complete_record = next(record for record in completion_records if record["entity_id"] == 2)
assert complete_record["entity_complete"] is True
assert complete_record["oversized_entity"] is False
@@ -8,6 +8,7 @@ from types import SimpleNamespace
import pytest
import basic_memory.repository.search_repository_base as search_repository_base_module
from basic_memory.repository.search_repository_base import (
MAX_VECTOR_CHUNK_CHARS,
SearchRepositoryBase,
@@ -335,6 +336,8 @@ async def test_sync_entity_vectors_batch_flushes_at_configured_threshold(monkeyp
assert result.entities_failed == 0
assert result.failed_entity_ids == []
assert result.embedding_jobs_total == 3
assert result.prepare_seconds_total == pytest.approx(0.0)
assert result.queue_wait_seconds_total == pytest.approx(0.0)
assert result.embed_seconds_total == pytest.approx(0.2)
assert result.write_seconds_total == pytest.approx(0.4)
@@ -373,3 +376,65 @@ async def test_sync_entity_vectors_batch_continue_on_error(monkeypatch):
assert result.entities_synced == 1
assert result.entities_failed == 2
assert result.failed_entity_ids == [2, 3]
@pytest.mark.asyncio
async def test_sync_entity_vectors_batch_tracks_prepare_and_queue_wait_seconds(monkeypatch):
"""Queue wait should be reported separately from prepare/embed/write timings."""
repo = _ConcreteRepo()
repo._semantic_enabled = True
repo._embedding_provider = object()
repo._semantic_embedding_sync_batch_size = 2
async def _stub_prepare(entity_id: int) -> _PreparedEntityVectorSync:
return _PreparedEntityVectorSync(
entity_id=entity_id,
sync_start=0.0,
source_rows_count=1,
embedding_jobs=[(100 + entity_id, f"chunk-{entity_id}")],
prepare_seconds=1.0,
)
async def _stub_flush(flush_jobs, entity_runtime, synced_entity_ids):
assert len(flush_jobs) == 2
for job in flush_jobs:
runtime = entity_runtime[job.entity_id]
if job.entity_id == 1:
runtime.embed_seconds = 1.0
runtime.write_seconds = 0.5
else:
runtime.embed_seconds = 2.0
runtime.write_seconds = 0.5
runtime.remaining_jobs = 0
synced_entity_ids.add(job.entity_id)
return (3.0, 1.0)
logged_completion: list[dict] = []
def _capture_log(**kwargs):
logged_completion.append(kwargs)
perf_counter_values = iter([4.0, 5.0])
monkeypatch.setattr(repo, "_prepare_entity_vector_jobs", _stub_prepare)
monkeypatch.setattr(repo, "_flush_embedding_jobs", _stub_flush)
monkeypatch.setattr(repo, "_log_vector_sync_complete", _capture_log)
monkeypatch.setattr(
search_repository_base_module.time,
"perf_counter",
lambda: next(perf_counter_values),
)
result = await repo.sync_entity_vectors_batch([1, 2])
assert result.entities_total == 2
assert result.entities_synced == 2
assert result.entities_failed == 0
assert result.prepare_seconds_total == pytest.approx(2.0)
assert result.queue_wait_seconds_total == pytest.approx(3.0)
assert result.embed_seconds_total == pytest.approx(3.0)
assert result.write_seconds_total == pytest.approx(1.0)
assert len(logged_completion) == 2
for record in logged_completion:
assert record["prepare_seconds"] == pytest.approx(1.0)
assert record["queue_wait_seconds"] == pytest.approx(1.5)
@@ -36,6 +36,12 @@ class StubEmbeddingProvider:
return [0.0, 0.0, 0.0, 1.0]
class StubEmbeddingProviderV2(StubEmbeddingProvider):
"""Same vectors, different model identity to force resync."""
model_name = "stub-v2"
def _entity_row(
*,
project_id: int,
@@ -62,14 +68,17 @@ def _entity_row(
)
def _enable_semantic(search_repository: SQLiteSearchRepository) -> None:
def _enable_semantic(
search_repository: SQLiteSearchRepository,
embedding_provider: StubEmbeddingProvider | None = None,
) -> None:
try:
import sqlite_vec # noqa: F401
except ImportError:
pytest.skip("sqlite-vec dependency is required for sqlite vector repository tests.")
search_repository._semantic_enabled = True
search_repository._embedding_provider = StubEmbeddingProvider()
search_repository._embedding_provider = embedding_provider or StubEmbeddingProvider()
search_repository._vector_dimensions = search_repository._embedding_provider.dimensions
search_repository._vector_tables_initialized = False
@@ -102,6 +111,8 @@ async def test_sqlite_vec_tables_are_created_and_rebuilt(search_repository):
"chunk_key",
"chunk_text",
"source_hash",
"entity_fingerprint",
"embedding_model",
"updated_at",
}
@@ -182,6 +193,79 @@ async def test_sqlite_chunk_upsert_and_delete_lifecycle(search_repository):
assert int(embedding_count.scalar_one()) == 0
@pytest.mark.asyncio
async def test_sqlite_vector_sync_skips_unchanged_and_reembeds_changed_content(search_repository):
"""SQLite vector sync tracks new, changed, unchanged, and model-changed entities."""
if not isinstance(search_repository, SQLiteSearchRepository):
pytest.skip("sqlite-vec repository behavior is local SQLite-only.")
_enable_semantic(search_repository)
await search_repository.init_search_index()
await search_repository.index_item(
_entity_row(
project_id=search_repository.project_id,
row_id=111,
entity_id=111,
title="Auth and Schema Notes",
permalink="specs/auth-and-schema",
content_stems="# Overview\n- auth token rotation\n- schema migration planning",
)
)
new_result = await search_repository.sync_entity_vectors_batch([111])
assert new_result.entities_synced == 1
assert new_result.entities_skipped == 0
assert new_result.chunks_total >= 2
assert new_result.chunks_skipped == 0
assert new_result.embedding_jobs_total == new_result.chunks_total
async with db.scoped_session(search_repository.session_maker) as session:
stored_rows = await session.execute(
text(
"SELECT entity_fingerprint, embedding_model "
"FROM search_vector_chunks "
"WHERE project_id = :project_id AND entity_id = :entity_id"
),
{"project_id": search_repository.project_id, "entity_id": 111},
)
metadata_rows = stored_rows.fetchall()
assert metadata_rows
assert len({row.entity_fingerprint for row in metadata_rows}) == 1
assert len({row.embedding_model for row in metadata_rows}) == 1
assert metadata_rows[0].embedding_model == "StubEmbeddingProvider:stub:4"
unchanged_result = await search_repository.sync_entity_vectors_batch([111])
assert unchanged_result.entities_synced == 1
assert unchanged_result.entities_skipped == 1
assert unchanged_result.embedding_jobs_total == 0
assert unchanged_result.chunks_skipped == unchanged_result.chunks_total
await search_repository.index_item(
_entity_row(
project_id=search_repository.project_id,
row_id=111,
entity_id=111,
title="Auth and Schema Notes",
permalink="specs/auth-and-schema",
content_stems="# Overview\n- auth token rotation\n- database schema migration planning",
)
)
changed_result = await search_repository.sync_entity_vectors_batch([111])
assert changed_result.entities_synced == 1
assert changed_result.entities_skipped == 0
assert changed_result.embedding_jobs_total >= 1
assert changed_result.chunks_skipped >= 1
assert changed_result.embedding_jobs_total < changed_result.chunks_total
_enable_semantic(search_repository, StubEmbeddingProviderV2())
model_changed_result = await search_repository.sync_entity_vectors_batch([111])
assert model_changed_result.entities_synced == 1
assert model_changed_result.entities_skipped == 0
assert model_changed_result.chunks_skipped == 0
assert model_changed_result.embedding_jobs_total == model_changed_result.chunks_total
@pytest.mark.asyncio
async def test_sqlite_vector_search_returns_ranked_entities(search_repository):
"""Vector mode ranks entities using sqlite-vec nearest-neighbor search."""
@@ -105,8 +105,14 @@ async def test_embedding_status_orphaned_chunks(
await project_service.repository.execute_query(
text(
"INSERT INTO search_vector_chunks "
"(entity_id, project_id, chunk_key, chunk_text, source_hash) "
"VALUES (:entity_id, :project_id, 'chunk-1', 'test text', 'abc123')"
"("
"entity_id, project_id, chunk_key, chunk_text, source_hash, "
"entity_fingerprint, embedding_model"
") "
"VALUES ("
":entity_id, :project_id, 'chunk-1', 'test text', 'abc123', "
"'fp-abc123', 'bge-small-en-v1.5'"
")"
),
{"entity_id": entity_id, "project_id": test_project.id},
)
@@ -213,8 +219,14 @@ async def test_embedding_status_healthy(project_service: ProjectService, test_gr
await project_service.repository.execute_query(
text(
"INSERT INTO search_vector_chunks "
"(id, entity_id, project_id, chunk_key, chunk_text, source_hash) "
"VALUES (:id, :entity_id, :project_id, :key, 'text', 'hash')"
"("
"id, entity_id, project_id, chunk_key, chunk_text, source_hash, "
"entity_fingerprint, embedding_model"
") "
"VALUES ("
":id, :entity_id, :project_id, :key, 'text', 'hash', "
"'fp-hash', 'bge-small-en-v1.5'"
")"
),
{
"id": chunk_id,
@@ -0,0 +1,426 @@
"""Tests proving upsert_entity_from_markdown optimizations.
Verifies that:
1. Redundant get_by_file_path call is eliminated (entity passed directly)
2. Final reload uses find_by_ids (PK lookup) instead of get_by_file_path (string lookup)
3. Telemetry sub-spans are emitted for each DB phase
4. Correctness is preserved for create, update, and edit flows
"""
from __future__ import annotations
import importlib
from contextlib import contextmanager
from datetime import datetime, timezone
from pathlib import Path
from types import SimpleNamespace
from unittest.mock import AsyncMock
import pytest
from basic_memory.markdown.schemas import (
EntityFrontmatter,
EntityMarkdown,
Observation as MarkdownObservation,
Relation as MarkdownRelation,
)
from basic_memory.schemas import Entity as EntitySchema
from basic_memory.services.entity_service import EntityService
entity_service_module = importlib.import_module("basic_memory.services.entity_service")
# --- Helpers ---
def _make_markdown(
title: str = "Test Entity",
observations: list | None = None,
relations: list | None = None,
) -> EntityMarkdown:
frontmatter = EntityFrontmatter(metadata={"title": title, "type": "note"})
return EntityMarkdown(
frontmatter=frontmatter,
observations=observations or [],
relations=relations or [],
created=datetime.now(timezone.utc),
modified=datetime.now(timezone.utc),
)
def _capture_spans():
spans: list[tuple[str, dict]] = []
@contextmanager
def fake_span(name: str, **attrs):
spans.append((name, attrs))
yield
return spans, fake_span
# --- Optimization 1: No redundant get_by_file_path in update_entity_relations ---
@pytest.mark.asyncio
async def test_upsert_update_does_not_refetch_entity(entity_service: EntityService, monkeypatch):
"""update_entity_relations should NOT call get_by_file_path — entity is passed directly."""
# Create an entity first
entity = await entity_service.create_entity(
EntitySchema(
title="Refetch Test",
directory="notes",
note_type="note",
content="# Refetch Test\n\n## Observations\n- [fact] some fact",
)
)
# Spy on get_by_file_path calls
original_get_by_file_path = entity_service.repository.get_by_file_path
call_count = 0
async def spy_get_by_file_path(*args, **kwargs):
nonlocal call_count
call_count += 1
return await original_get_by_file_path(*args, **kwargs)
monkeypatch.setattr(entity_service.repository, "get_by_file_path", spy_get_by_file_path)
# Run upsert with is_new=False — this calls update_entity_and_observations + update_entity_relations
markdown = _make_markdown(
title="Refetch Test",
observations=[MarkdownObservation(content="updated fact", category="fact")],
)
await entity_service.upsert_entity_from_markdown(Path(entity.file_path), markdown, is_new=False)
# update_entity_and_observations calls get_by_file_path once (to load the entity)
# update_entity_relations should NOT call it at all (entity passed directly)
assert call_count == 1, (
f"Expected 1 get_by_file_path call (in update_entity_and_observations only), "
f"got {call_count}. update_entity_relations should not re-fetch."
)
# --- Optimization 2: Final reload uses find_by_ids (PK) not get_by_file_path ---
@pytest.mark.asyncio
async def test_update_entity_relations_uses_pk_reload(entity_service: EntityService, monkeypatch):
"""update_entity_relations should use find_by_ids for the final reload, not get_by_file_path."""
entity = await entity_service.create_entity(
EntitySchema(
title="PK Reload Test",
directory="notes",
note_type="note",
content="# PK Reload Test",
)
)
# Spy on find_by_ids calls
original_find_by_ids = entity_service.repository.find_by_ids
find_by_ids_calls = []
async def spy_find_by_ids(ids):
find_by_ids_calls.append(ids)
return await original_find_by_ids(ids)
monkeypatch.setattr(entity_service.repository, "find_by_ids", spy_find_by_ids)
markdown = _make_markdown(title="PK Reload Test")
await entity_service.upsert_entity_from_markdown(Path(entity.file_path), markdown, is_new=False)
# update_entity_relations should call find_by_ids once with the entity's PK
assert len(find_by_ids_calls) == 1
assert find_by_ids_calls[0] == [entity.id]
@pytest.mark.asyncio
async def test_create_or_update_entity_uses_lightweight_exact_resolution(
entity_service: EntityService, monkeypatch
):
"""create_or_update_entity should use strict lookups without eager relation loading."""
schema = EntitySchema(
title="Create Or Update",
directory="notes",
note_type="note",
content="# Create Or Update",
)
sentinel_entity = SimpleNamespace(file_path="notes/existing.md")
resolve_calls: list[tuple[str, dict]] = []
async def fake_resolve_link(link_text: str, **kwargs):
resolve_calls.append((link_text, kwargs))
if link_text == schema.file_path:
return None
return sentinel_entity
monkeypatch.setattr(entity_service.link_resolver, "resolve_link", fake_resolve_link)
monkeypatch.setattr(entity_service, "update_entity", AsyncMock(return_value=sentinel_entity))
entity, is_new = await entity_service.create_or_update_entity(schema)
assert entity is sentinel_entity
assert is_new is False
assert resolve_calls == [
(schema.file_path, {"strict": True, "load_relations": False}),
(schema.permalink, {"strict": True, "load_relations": False}),
]
# --- Telemetry sub-spans ---
@pytest.mark.asyncio
async def test_upsert_update_emits_sub_spans(entity_service: EntityService, monkeypatch):
"""upsert_entity_from_markdown (update path) should emit sub-spans for each DB phase."""
entity = await entity_service.create_entity(
EntitySchema(
title="Span Test",
directory="notes",
note_type="note",
content="# Span Test\n\n## Observations\n- [fact] original",
)
)
spans, fake_span = _capture_spans()
monkeypatch.setattr(entity_service_module.telemetry, "span", fake_span)
markdown = _make_markdown(
title="Span Test",
observations=[MarkdownObservation(content="updated", category="fact")],
)
await entity_service.upsert_entity_from_markdown(Path(entity.file_path), markdown, is_new=False)
span_names = [name for name, _ in spans]
# update_entity_and_observations sub-spans
assert "upsert.update.fetch_entity" in span_names
assert "upsert.update.delete_observations" in span_names
assert "upsert.update.insert_observations" in span_names
assert "upsert.update.save_entity" in span_names
# update_entity_relations sub-spans
assert "upsert.relations.delete_existing" in span_names
assert "upsert.relations.reload_entity" in span_names
@pytest.mark.asyncio
async def test_upsert_with_relations_emits_resolve_and_insert_spans(
entity_service: EntityService, monkeypatch
):
"""When relations exist, resolve_links and insert_relations spans should be emitted."""
# Create two entities so the relation can resolve
await entity_service.create_entity(
EntitySchema(
title="Target Entity",
directory="notes",
note_type="note",
content="# Target Entity",
)
)
source = await entity_service.create_entity(
EntitySchema(
title="Source Entity",
directory="notes",
note_type="note",
content="# Source Entity",
)
)
spans, fake_span = _capture_spans()
monkeypatch.setattr(entity_service_module.telemetry, "span", fake_span)
markdown = _make_markdown(
title="Source Entity",
relations=[MarkdownRelation(type="links_to", target="Target Entity")],
)
await entity_service.upsert_entity_from_markdown(Path(source.file_path), markdown, is_new=False)
span_names = [name for name, _ in spans]
assert "upsert.relations.resolve_links" in span_names
assert "upsert.relations.insert_relations" in span_names
@pytest.mark.asyncio
async def test_upsert_with_relations_uses_lightweight_exact_resolution(
entity_service: EntityService, monkeypatch
):
"""Relation target resolution should skip eager loading during upsert."""
target = await entity_service.create_entity(
EntitySchema(
title="Lightweight Target",
directory="notes",
note_type="note",
content="# Lightweight Target",
)
)
source = await entity_service.create_entity(
EntitySchema(
title="Lightweight Source",
directory="notes",
note_type="note",
content="# Lightweight Source",
)
)
resolve_calls: list[tuple[str, dict]] = []
async def fake_resolve_link(link_text: str, **kwargs):
resolve_calls.append((link_text, kwargs))
return target
monkeypatch.setattr(entity_service.link_resolver, "resolve_link", fake_resolve_link)
markdown = _make_markdown(
title="Lightweight Source",
relations=[MarkdownRelation(type="links_to", target="Lightweight Target")],
)
await entity_service.upsert_entity_from_markdown(Path(source.file_path), markdown, is_new=False)
assert resolve_calls == [
("Lightweight Target", {"strict": True, "load_relations": False}),
]
# --- Correctness: full round-trip ---
@pytest.mark.asyncio
async def test_upsert_update_preserves_observations(entity_service: EntityService):
"""After upsert (update path), observations should be correctly replaced."""
entity = await entity_service.create_entity(
EntitySchema(
title="Obs Test",
directory="notes",
note_type="note",
content="# Obs Test\n\n## Observations\n- [fact] original fact",
)
)
assert len(entity.observations) == 1
markdown = _make_markdown(
title="Obs Test",
observations=[
MarkdownObservation(content="new fact 1", category="fact"),
MarkdownObservation(content="new fact 2", category="idea"),
],
)
updated = await entity_service.upsert_entity_from_markdown(
Path(entity.file_path), markdown, is_new=False
)
assert updated.id == entity.id
assert len(updated.observations) == 2
obs_contents = {o.content for o in updated.observations}
assert obs_contents == {"new fact 1", "new fact 2"}
@pytest.mark.asyncio
async def test_upsert_update_preserves_relations(entity_service: EntityService):
"""After upsert (update path), relations should be correctly replaced."""
target = await entity_service.create_entity(
EntitySchema(
title="Relation Target",
directory="notes",
note_type="note",
content="# Relation Target",
)
)
source = await entity_service.create_entity(
EntitySchema(
title="Relation Source",
directory="notes",
note_type="note",
content="# Relation Source\n\n## Relations\n- links_to [[Relation Target]]",
)
)
assert len(source.relations) == 1
markdown = _make_markdown(
title="Relation Source",
relations=[MarkdownRelation(type="references", target="Relation Target")],
)
updated = await entity_service.upsert_entity_from_markdown(
Path(source.file_path), markdown, is_new=False
)
assert updated.id == source.id
# Old relation replaced with new one
outgoing = [r for r in updated.relations if r.from_id == source.id]
assert len(outgoing) == 1
assert outgoing[0].relation_type == "references"
assert outgoing[0].to_id == target.id
@pytest.mark.asyncio
async def test_upsert_create_path_works(entity_service: EntityService):
"""The is_new=True path should still work correctly."""
markdown = _make_markdown(
title="Create Path Test",
observations=[MarkdownObservation(content="a fact", category="fact")],
)
result = await entity_service.upsert_entity_from_markdown(
Path("notes/create-path-test.md"), markdown, is_new=True
)
assert result.title == "Create Path Test"
assert len(result.observations) == 1
assert result.observations[0].content == "a fact"
@pytest.mark.asyncio
async def test_edit_entity_end_to_end(entity_service: EntityService):
"""Full edit_entity flow uses optimized upsert and returns correct entity."""
entity = await entity_service.create_entity(
EntitySchema(
title="Edit E2E",
directory="notes",
note_type="note",
content="# Edit E2E\n\nOriginal content.",
)
)
updated = await entity_service.edit_entity(
entity.file_path,
operation="append",
content="\n\n## Observations\n- [fact] appended fact",
)
assert updated.id == entity.id
assert len(updated.observations) == 1
assert updated.observations[0].content == "appended fact"
# Checksum should be set (not None) after edit completes
assert updated.checksum is not None
@pytest.mark.asyncio
async def test_edit_entity_uses_lightweight_identifier_resolution(
entity_service: EntityService, monkeypatch
):
"""edit_entity should resolve the target note without eager relation loading."""
entity = await entity_service.create_entity(
EntitySchema(
title="Edit Lightweight",
directory="notes",
note_type="note",
content="# Edit Lightweight\n\nOriginal content.",
)
)
original_resolve_link = entity_service.link_resolver.resolve_link
resolve_calls: list[tuple[str, dict]] = []
async def spy_resolve_link(link_text: str, **kwargs):
resolve_calls.append((link_text, kwargs))
return await original_resolve_link(link_text, **kwargs)
monkeypatch.setattr(entity_service.link_resolver, "resolve_link", spy_resolve_link)
await entity_service.edit_entity(
entity.file_path,
operation="append",
content="\n\nNo relation changes here.",
)
assert resolve_calls[0] == (
entity.file_path,
{"strict": True, "load_relations": False},
)
+112
View File
@@ -0,0 +1,112 @@
"""Regression tests for Alembic env async migration helpers."""
import importlib.util
import uuid
from contextlib import nullcontext
from pathlib import Path
import pytest
class FakeAlembicConfig:
"""Minimal config object used while importing env.py under test."""
def __init__(self):
self.options = {"sqlalchemy.url": "sqlite:///:memory:"}
self.attributes = {}
self.config_file_name = None
self.config_ini_section = "alembic"
def get_main_option(self, name: str) -> str | None:
return self.options.get(name)
def set_main_option(self, name: str, value: str) -> None:
self.options[name] = value
def get_section(self, name: str, default=None):
return default or {}
class FakeCoroutine:
"""Track whether the migration coroutine gets closed on failure."""
def __init__(self):
self.closed = False
def close(self) -> None:
self.closed = True
def load_alembic_env_module(monkeypatch, tmp_path):
"""Import env.py with a fake Alembic context and isolated HOME."""
monkeypatch.setenv("HOME", str(tmp_path))
monkeypatch.setenv("BASIC_MEMORY_HOME", str(tmp_path / "basic-memory"))
from alembic import context as alembic_context
fake_config = FakeAlembicConfig()
monkeypatch.setattr(alembic_context, "config", fake_config, raising=False)
monkeypatch.setattr(alembic_context, "configure", lambda *args, **kwargs: None, raising=False)
monkeypatch.setattr(alembic_context, "begin_transaction", lambda: nullcontext(), raising=False)
monkeypatch.setattr(alembic_context, "run_migrations", lambda: None, raising=False)
monkeypatch.setattr(alembic_context, "is_offline_mode", lambda: True, raising=False)
env_path = Path(__file__).resolve().parents[1] / "src/basic_memory/alembic/env.py"
module_name = f"test_alembic_env_{uuid.uuid4().hex}"
spec = importlib.util.spec_from_file_location(module_name, env_path)
assert spec is not None
assert spec.loader is not None
module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(module)
return module
def test_asyncio_run_failure_closes_migration_coroutine(monkeypatch, tmp_path):
"""The running-loop fallback should not leak an un-awaited coroutine."""
env_module = load_alembic_env_module(monkeypatch, tmp_path)
fake_coro = FakeCoroutine()
monkeypatch.setattr(env_module, "run_async_migrations", lambda connectable: fake_coro)
def raising_asyncio_run(coro):
raise RuntimeError("asyncio.run() cannot be called from a running event loop")
monkeypatch.setattr(env_module.asyncio, "run", raising_asyncio_run)
with pytest.raises(RuntimeError, match="running event loop"):
env_module._run_async_migrations_with_asyncio_run(object())
assert fake_coro.closed is True
def test_running_loop_error_uses_thread_fallback(monkeypatch, tmp_path):
"""Async-engine helper should switch to the thread fallback for running-loop errors."""
env_module = load_alembic_env_module(monkeypatch, tmp_path)
connectable = object()
fallback_calls: list[object] = []
def raising_run(connectable):
raise RuntimeError("asyncio.run() cannot be called from a running event loop")
def record_fallback(target):
fallback_calls.append(target)
monkeypatch.setattr(env_module, "_run_async_migrations_with_asyncio_run", raising_run)
monkeypatch.setattr(env_module, "_run_async_migrations_in_thread", record_fallback)
env_module._run_async_engine_migrations(connectable)
assert fallback_calls == [connectable]
def test_non_loop_runtime_error_is_re_raised(monkeypatch, tmp_path):
"""Unexpected RuntimeError values should not be swallowed by the fallback path."""
env_module = load_alembic_env_module(monkeypatch, tmp_path)
def raising_run(connectable):
raise RuntimeError("different runtime failure")
monkeypatch.setattr(env_module, "_run_async_migrations_with_asyncio_run", raising_run)
with pytest.raises(RuntimeError, match="different runtime failure"):
env_module._run_async_engine_migrations(object())
+16
View File
@@ -882,6 +882,22 @@ class TestSemanticSearchConfig:
config = BasicMemoryConfig(semantic_embedding_dimensions=1536)
assert config.semantic_embedding_dimensions == 1536
def test_semantic_postgres_prepare_concurrency_defaults_to_4(self):
"""Postgres prepare concurrency should default to a conservative window of 4."""
config = BasicMemoryConfig()
assert config.semantic_postgres_prepare_concurrency == 4
def test_semantic_postgres_prepare_concurrency_validation(self):
"""Postgres prepare concurrency must stay within the bounded safe range."""
config = BasicMemoryConfig(semantic_postgres_prepare_concurrency=8)
assert config.semantic_postgres_prepare_concurrency == 8
with pytest.raises(Exception):
BasicMemoryConfig(semantic_postgres_prepare_concurrency=0)
with pytest.raises(Exception):
BasicMemoryConfig(semantic_postgres_prepare_concurrency=17)
def test_semantic_search_enabled_description_mentions_both_backends(self):
"""Description should not say 'SQLite only' anymore."""
field_info = BasicMemoryConfig.model_fields["semantic_search_enabled"]
+75
View File
@@ -0,0 +1,75 @@
"""Migration tests for note_content schema."""
import sqlite3
from pathlib import Path
from alembic import command
from alembic.config import Config
from basic_memory import db
def sqlite_alembic_config(database_path: Path) -> Config:
"""Build an Alembic config that upgrades a temporary SQLite database."""
alembic_dir = Path(db.__file__).parent / "alembic"
config = Config()
config.set_main_option("script_location", str(alembic_dir))
config.set_main_option(
"file_template",
"%%(year)d_%%(month).2d_%%(day).2d_%%(hour).2d%%(minute).2d-%%(rev)s_%%(slug)s",
)
config.set_main_option("timezone", "UTC")
config.set_main_option("revision_environment", "false")
config.set_main_option("sqlalchemy.url", f"sqlite:///{database_path}")
return config
def test_alembic_upgrade_creates_note_content_table(tmp_path, monkeypatch):
"""Running Alembic head should create note_content with its expected contract."""
monkeypatch.setenv("HOME", str(tmp_path))
monkeypatch.setenv("BASIC_MEMORY_HOME", str(tmp_path / "basic-memory"))
database_path = tmp_path / "note-content-migration.db"
command.upgrade(sqlite_alembic_config(database_path), "head")
connection = sqlite3.connect(database_path)
try:
columns = {
row[1] for row in connection.execute("PRAGMA table_info(note_content)").fetchall()
}
assert columns == {
"entity_id",
"project_id",
"external_id",
"file_path",
"markdown_content",
"db_version",
"db_checksum",
"file_version",
"file_checksum",
"file_write_status",
"last_source",
"updated_at",
"file_updated_at",
"last_materialization_error",
"last_materialization_attempt_at",
}
foreign_keys = connection.execute("PRAGMA foreign_key_list(note_content)").fetchall()
entity_fk = next(row for row in foreign_keys if row[3] == "entity_id")
project_fk = next(row for row in foreign_keys if row[3] == "project_id")
assert entity_fk[2] == "entity"
assert entity_fk[4] == "id"
assert entity_fk[6].upper() == "CASCADE"
assert project_fk[2] == "project"
assert project_fk[4] == "id"
assert project_fk[6].upper() == "CASCADE"
indexes = {
row[1] for row in connection.execute("PRAGMA index_list(note_content)").fetchall()
}
assert "ix_note_content_project_id" in indexes
assert "ix_note_content_file_path" in indexes
assert "ix_note_content_external_id" in indexes
finally:
connection.close()