Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions src/basic_memory/deps/services.py
Original file line number Diff line number Diff line change
Expand Up @@ -41,6 +41,7 @@
LocalDirectoryDeleteRelationCleanupRefresher,
LocalDirectoryFileDeleteEnqueuer,
)
from basic_memory.repository import NoteContentRepository
from basic_memory.repository.accepted_note_repositories import AcceptedNoteRepositories
from basic_memory.repository.search_repository import create_search_repository
from basic_memory.index.local_project import (
Expand Down Expand Up @@ -429,10 +430,14 @@ async def get_relation_resolution_scheduler(
) -> RelationResolutionScheduler:
# Build the project-scoped resolution runtime. It owns its own sessions via
# session_maker, so it is safe to run from a detached background task.
project_id = entity_repository.project_id
if project_id is None:
raise RuntimeError("Relation resolution requires a project-scoped entity repository")
runtime = RepositoryRelationResolutionRuntime(
session_maker=session_maker,
relation_repository=relation_repository,
entity_repository=entity_repository,
note_content_repository=NoteContentRepository(project_id=project_id),
link_resolver=link_resolver,
entity_indexer=search_service,
)
Expand Down
2 changes: 2 additions & 0 deletions src/basic_memory/index/local_project.py
Original file line number Diff line number Diff line change
Expand Up @@ -80,6 +80,7 @@
resolve_project_index_completion_relations,
)
from basic_memory.models import Entity, Project
from basic_memory.repository import NoteContentRepository
from basic_memory.runtime.jobs import (
RuntimeIndexFileBatchJobRequest,
RuntimeJobId,
Expand Down Expand Up @@ -649,6 +650,7 @@ def runtime_from_dependencies(
session_maker=dependencies.session_maker,
relation_repository=dependencies.relation_repository,
entity_repository=dependencies.entity_repository,
note_content_repository=NoteContentRepository(project_id=dependencies.project_id),
link_resolver=dependencies.link_resolver,
entity_indexer=dependencies.search_service,
),
Expand Down
4 changes: 4 additions & 0 deletions src/basic_memory/index/local_runtime.py
Original file line number Diff line number Diff line change
Expand Up @@ -57,6 +57,7 @@
RepositoryExternalFileDeleteEntities,
)
from basic_memory.models import Entity, Project
from basic_memory.repository import NoteContentRepository
from basic_memory.runtime.projects import ProjectRuntimeReference
from basic_memory.runtime.storage import (
ProjectPath,
Expand Down Expand Up @@ -313,6 +314,9 @@ async def runtime_for_project(self, project: Project) -> StorageEventIndexRuntim
session_maker=dependencies.session_maker,
relation_repository=dependencies.relation_repository,
entity_repository=dependencies.entity_repository,
note_content_repository=NoteContentRepository(
project_id=dependencies.project_id
),
link_resolver=dependencies.link_resolver,
entity_indexer=dependencies.search_service,
),
Expand Down
57 changes: 53 additions & 4 deletions src/basic_memory/indexing/relation_resolution.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,13 +12,14 @@
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker

from basic_memory import db
from basic_memory.indexing.accepted_note_search import accepted_search_content_from_markdown
from basic_memory.indexing.models import IndexFileJobStatus
from basic_memory.services.exceptions import AmbiguousIdentifierError
from basic_memory.models import Entity
from basic_memory.repository.relation_repository import (
ResolvedRelationWrite,
ResolvedRelationWriteResult,
)
from basic_memory.services.exceptions import AmbiguousIdentifierError

type EntityId = int
type AffectedEntityIds = set[EntityId]
Expand Down Expand Up @@ -93,6 +94,25 @@ async def find_by_ids(
"""Return source entities by database id."""


class RelationResolutionNoteContent(Protocol):
"""Accepted note-content fields needed for a state-aware search refresh."""

entity_id: EntityId
markdown_content: str
file_write_status: str


class BatchRelationResolutionNoteContentRepository(Protocol):
"""Repository capability for loading accepted content for affected sources."""

async def find_by_ids(
self,
session: AsyncSession,
ids: list[EntityId],
) -> Sequence[RelationResolutionNoteContent]:
"""Return accepted note content by source entity id."""


class RelationResolutionLinkResolver(Protocol):
"""Capability for resolving a relation target by link text."""

Expand All @@ -116,17 +136,34 @@ async def index_entity(self, entity: Entity) -> None:
class BatchRelationResolutionEntityIndexer(Protocol):
"""Capability for refreshing a batch of derived entity search rows."""

async def index_entities(self, entities: Sequence[Entity]) -> None:
async def index_entities(
self,
entities: Sequence[Entity],
*,
content_by_entity_id: Mapping[EntityId, str],
) -> None:
"""Refresh derived index rows for a group of entities."""


def accepted_search_content_for_pending_notes(
note_contents: Sequence[RelationResolutionNoteContent],
) -> dict[EntityId, str]:
"""Return accepted search content while Markdown projections are not synchronized."""
return {
note_content.entity_id: accepted_search_content_from_markdown(note_content.markdown_content)
for note_content in note_contents
if note_content.file_write_status != "synced"
}


@dataclass(frozen=True, slots=True)
class RepositoryRelationResolutionRuntime:
"""Resolve forward references with project-scoped repositories and services."""

session_maker: async_sessionmaker[AsyncSession]
relation_repository: RelationResolutionRelationRepository
entity_repository: BatchRelationResolutionEntityRepository
note_content_repository: BatchRelationResolutionNoteContentRepository
link_resolver: RelationResolutionLinkResolver
entity_indexer: BatchRelationResolutionEntityIndexer

Expand Down Expand Up @@ -217,13 +254,25 @@ async def resolve_relations(
)

if affected_entity_ids:
sorted_affected_entity_ids = sorted(affected_entity_ids)
async with db.scoped_session(self.session_maker) as session:
source_entities = await self.entity_repository.find_by_ids(
session,
sorted(affected_entity_ids),
sorted_affected_entity_ids,
)
note_contents = await self.note_content_repository.find_by_ids(
session,
sorted_affected_entity_ids,
)

# Trigger: an accepted note has not reached a synchronized Markdown projection.
# Why: relation repair must refresh its search row without racing the async file
# materializer, while synchronized notes still need missing files to surface.
# Outcome: pending sources use accepted DB content; all others retain disk fallback.
content_by_entity_id = accepted_search_content_for_pending_notes(note_contents)
await self.entity_indexer.index_entities(
sorted(source_entities, key=lambda entity: entity.id)
sorted(source_entities, key=lambda entity: entity.id),
content_by_entity_id=content_by_entity_id,
)

return affected_entity_ids
Expand Down
2 changes: 2 additions & 0 deletions src/basic_memory/services/composition.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@
from basic_memory.markdown import EntityParser
from basic_memory.repository import (
EntityRepository,
NoteContentRepository,
ObservationRepository,
ProjectRepository,
RelationRepository,
Expand Down Expand Up @@ -164,6 +165,7 @@ def build_default_project_runtime_bundle(
session_maker=session_maker,
relation_repository=relation_resolution_repository,
entity_repository=entity_repository,
note_content_repository=NoteContentRepository(project_id=project_id),
link_resolver=link_resolver,
entity_indexer=search_service,
)
Expand Down
16 changes: 13 additions & 3 deletions src/basic_memory/services/search_service.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@
import asyncio
import ast
import re
from collections.abc import Sequence
from collections.abc import Mapping, Sequence
from dataclasses import dataclass
from datetime import datetime
from typing import Any, List, Optional, Set, Dict
Expand Down Expand Up @@ -398,15 +398,25 @@ async def index_entity(
else:
await self.index_entity_data(entity, content)

async def index_entities(self, entities: Sequence[Entity]) -> None:
async def index_entities(
self,
entities: Sequence[Entity],
*,
content_by_entity_id: Mapping[int, str],
) -> None:
"""Refresh a group of entity search rows through one batch entry point.

Index writes stay sequential because local SQLite connections cannot
safely run these mutations concurrently. Callers still avoid reopening
repository sessions and dispatching one indexing API call per entity.
Accepted content bypasses the disk read while its Markdown projection
remains pending.
"""
for entity in entities:
await self.index_entity_data(entity)
await self.index_entity_data(
entity,
content=content_by_entity_id.get(entity.id),
)

async def index_entity_data(
self,
Expand Down
116 changes: 114 additions & 2 deletions tests/index/test_local_project_index.py
Original file line number Diff line number Diff line change
Expand Up @@ -60,7 +60,10 @@
)
from basic_memory.models import Entity, Project, Relation
from basic_memory.repository import EntityRepository
from basic_memory.repository.note_content_repository import NoteContentRepository
from basic_memory.repository.note_content_repository import (
AcceptedNoteContentWrite,
NoteContentRepository,
)
from basic_memory.repository.relation_repository import (
ResolvedRelationWrite,
ResolvedRelationWriteResult,
Expand Down Expand Up @@ -1712,6 +1715,110 @@ async def test_local_project_index_resolves_order_dependent_relations_after_batc
assert {row.to_id for row in relation_search_rows} == {entity_a.id, entity_b.id, entity_c.id}


async def test_local_relation_resolution_refreshes_pending_source_without_markdown_file(
test_project: Project,
project_config,
entity_repository,
session_maker: async_sessionmaker[AsyncSession],
search_service,
config_manager,
monkeypatch,
) -> None:
"""Accepted content keeps relation repair independent of async file materialization."""
del config_manager

source_path = project_config.home / "pending-source.md"
source_path.write_text(
"# Pending Source\n\n- relates_to [[Pending Target]]\n",
encoding="utf-8",
)
initial = await run_local_project_index_for_project(
test_project,
runtime_factory=LocalProjectIndexRuntimeFactory(batch_size=10),
force_full=True,
)
assert initial.enqueued_files == 1

accepted_markdown = (
source_path.read_text(encoding="utf-8")
+ "\nAccepted search content before materialization.\n"
)
accepted_checksum = sha256(accepted_markdown.encode("utf-8")).hexdigest()
accepted_at = datetime.now(timezone.utc)
note_content_repository = NoteContentRepository(test_project.id)

async with db.scoped_session(session_maker) as session:
source = await entity_repository.get_by_file_path(session, "pending-source.md")
assert source is not None
current_note_content = await note_content_repository.get_by_entity_id(
session,
source.id,
)
assert current_note_content is not None

target = await entity_repository.add(
session,
Entity(
permalink=f"{test_project.permalink}/pending-target",
title="Pending Target",
note_type="note",
file_path="pending-target.md",
checksum="target-checksum",
content_type="text/markdown",
created_at=accepted_at,
updated_at=accepted_at,
),
)
await entity_repository.update(
session,
source.id,
{
"checksum": accepted_checksum,
"updated_at": accepted_at,
},
)
await note_content_repository.accept_write(
session,
AcceptedNoteContentWrite(
entity_id=source.id,
markdown_content=accepted_markdown,
db_version=current_note_content.db_version + 1,
db_checksum=accepted_checksum,
last_source="test",
updated_at=accepted_at,
),
)
source_id = source.id
target_id = target.id

# The accepted database state is durable, but its Markdown projection is
# intentionally absent when relation resolution refreshes the source.
source_path.unlink()
runtime = await LocalProjectIndexRuntimeFactory().runtime_for_project(test_project)
assert isinstance(runtime.completion_relation_runtime, RepositoryRelationResolutionRuntime)

affected = await runtime.completion_relation_runtime.resolve_relations()

assert affected == {source_id}
search_results = await search_service.search(
SearchQuery(text="Accepted search content before materialization")
)
assert [result.entity_id for result in search_results] == [source_id]

async with db.scoped_session(session_maker) as session:
resolved_source = await entity_repository.find_by_id(session, source_id)
relation_search_rows = await search_service.repository.search(
search_item_types=[SearchItemType.RELATION],
session=session,
)

assert resolved_source is not None
assert [relation.to_id for relation in resolved_source.outgoing_relations] == [target_id]
source_relation_rows = [row for row in relation_search_rows if row.entity_id == source_id]
assert len(source_relation_rows) == 1
assert source_relation_rows[0].to_id == target_id


async def test_local_project_index_deduplicates_relations_by_type(
test_project: Project,
project_config,
Expand Down Expand Up @@ -2449,7 +2556,12 @@ async def handle_delete(self, entity: Entity) -> None:
async def index_entity(self, entity: Entity) -> None:
return None

async def index_entities(self, entities: Sequence[Entity]) -> None:
async def index_entities(
self,
entities: Sequence[Entity],
*,
content_by_entity_id: Mapping[int, str],
) -> None:
return None

async def sync_entity_vectors_batch(
Expand Down
Loading
Loading