From 33f8fd5af0f043dfecb3f199e9769a230ee87248 Mon Sep 17 00:00:00 2001 From: wu Date: Thu, 2 Jul 2026 10:55:19 +0900 Subject: [PATCH 1/9] feat: track resource provenance in workspace memorize MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add an optional `track` column to the Resource model that records which workspace subtree a file came from: files under chat/ → "chat", under agent/ → "skill", everything else → "workspace". Classification happens in memorize_workspace and threads through the workflow as the `resource_track` state key. Legacy single-file memorize leaves it None. The column is nullable across all three backends (in-memory, SQLite, Postgres), mirroring how RecallFile.track was added — created via metadata.create_all(), no migration. Co-Authored-By: Claude Opus 4.8 (1M context) --- src/memu/app/memorize.py | 27 +++++++++++++++++++ .../inmemory/repositories/resource_repo.py | 2 ++ src/memu/database/models.py | 3 +++ src/memu/database/postgres/models.py | 1 + .../postgres/repositories/resource_repo.py | 2 ++ src/memu/database/repositories/resource.py | 1 + src/memu/database/sqlite/models.py | 1 + .../sqlite/repositories/resource_repo.py | 6 +++++ tests/test_folder_memorize.py | 9 ++++--- 9 files changed, 49 insertions(+), 3 deletions(-) diff --git a/src/memu/app/memorize.py b/src/memu/app/memorize.py index 65978fa2..d5fcfc7a 100644 --- a/src/memu/app/memorize.py +++ b/src/memu/app/memorize.py @@ -107,6 +107,8 @@ async def memorize( "user": user_scope, # Legacy single-resource path: force only the provided categories. "allow_new_categories": False, + # Legacy path does not classify by workspace track. + "resource_track": None, } result = await self._run_workflow("memorize", state) @@ -161,6 +163,7 @@ async def memorize_workspace( user_scope=user_scope, ctx=ctx, store=store, + track=self._classify_track(scanned_file.rel_path), ) changed_resources.extend(cast("list[Resource]", result.get("resources") or [])) # The inner single-file ``memorize`` keeps its legacy response keys @@ -196,6 +199,7 @@ async def _memorize_one( user_scope: dict[str, Any] | None, ctx: Context, store: Database, + track: str | None = None, ) -> WorkflowState: """Run the memorize workflow for a single file (one file -> one Resource). @@ -215,6 +219,8 @@ async def _memorize_one( "user": user_scope, # Workspace sync path: let the extractor grow the taxonomy. "allow_new_categories": True, + # Which workspace track this file belongs to (chat/skill/workspace). + "resource_track": track, } # The workspace path runs its own workflow (memorize + per-file skill # generation); single-file ``memorize`` stays untouched (ADR 0006). @@ -224,6 +230,21 @@ async def _memorize_one( raise RuntimeError(msg) return result + @staticmethod + def _classify_track(rel_path: str) -> str: + """Classify a workspace file into a track by its top-level folder. + + Files under ``chat/`` are the ``"chat"`` track, files under ``agent/`` are + the ``"skill"`` track, and everything else is the ``"workspace"`` track. + ``rel_path`` is the posix path relative to the scanned folder root. + """ + top = rel_path.split("/", 1)[0] + if top == "chat": + return "chat" + if top == "agent": + return "skill" + return "workspace" + async def _cascade_delete_by_urls( self, urls: set[str], @@ -346,6 +367,7 @@ def _build_memorize_workflow(self) -> list[WorkflowStep]: "modality", "user", "allow_new_categories", + "resource_track", }, produces={"resources", "entries", "relations", "file_updates"}, capabilities={"db", "vector"}, @@ -383,6 +405,7 @@ def _list_memorize_initial_keys() -> set[str]: "category_ids", "user", "allow_new_categories", + "resource_track", } def _build_memorize_workspace_workflow(self) -> list[WorkflowStep]: @@ -481,6 +504,7 @@ async def _memorize_categorize_entries(self, state: WorkflowState, step_context: file_updates: dict[str, list[tuple[str, str]]] = {} user_scope = state.get("user", {}) allow_new_categories = state.get("allow_new_categories", False) + track = state.get("resource_track") for plan in state.get("resource_plans", []): res = await self._create_resource_with_caption( @@ -491,6 +515,7 @@ async def _memorize_categorize_entries(self, state: WorkflowState, step_context: store=store, embed_client=embed_client, user=user_scope, + track=track, ) resources.append(res) @@ -687,6 +712,7 @@ async def _create_resource_with_caption( store: Database, embed_client: Any | None = None, user: Mapping[str, Any] | None = None, + track: str | None = None, ) -> Resource: caption_text = caption.strip() if caption else None if caption_text: @@ -702,6 +728,7 @@ async def _create_resource_with_caption( caption=caption_text, embedding=caption_embedding, user_data=dict(user or {}), + track=track, ) # if caption: # caption_text = caption.strip() diff --git a/src/memu/database/inmemory/repositories/resource_repo.py b/src/memu/database/inmemory/repositories/resource_repo.py index 04c4e986..6116ec2f 100644 --- a/src/memu/database/inmemory/repositories/resource_repo.py +++ b/src/memu/database/inmemory/repositories/resource_repo.py @@ -44,6 +44,7 @@ def create_resource( caption: str | None, embedding: list[float] | None, user_data: dict[str, Any], + track: str | None = None, ) -> Resource: rid = str(uuid.uuid4()) res = self.resource_model( @@ -53,6 +54,7 @@ def create_resource( local_path=local_path, caption=caption, embedding=embedding, + track=track, **user_data, ) self.resources[rid] = res diff --git a/src/memu/database/models.py b/src/memu/database/models.py index dece135c..23361d5c 100644 --- a/src/memu/database/models.py +++ b/src/memu/database/models.py @@ -71,6 +71,9 @@ class Resource(BaseRecord): local_path: str caption: str | None = None embedding: list[float] | None = None + # Which workspace track this resource came from: "chat", "skill", or + # "workspace" (set by ``memorize_workspace``). None for legacy ``memorize``. + track: str | None = None class RecallEntry(BaseRecord): diff --git a/src/memu/database/postgres/models.py b/src/memu/database/postgres/models.py index e7d94faf..be8c11eb 100644 --- a/src/memu/database/postgres/models.py +++ b/src/memu/database/postgres/models.py @@ -49,6 +49,7 @@ class ResourceModel(BaseModelMixin, Resource): local_path: str = Field(sa_column=Column(String, nullable=False)) caption: str | None = Field(default=None, sa_column=Column(Text, nullable=True)) embedding: list[float] | None = Field(default=None, sa_column=Column(Vector(), nullable=True)) + track: str | None = Field(default=None, sa_column=Column(String, nullable=True)) class RecallEntryModel(BaseModelMixin, RecallEntry): diff --git a/src/memu/database/postgres/repositories/resource_repo.py b/src/memu/database/postgres/repositories/resource_repo.py index 2efcc848..e71e6da5 100644 --- a/src/memu/database/postgres/repositories/resource_repo.py +++ b/src/memu/database/postgres/repositories/resource_repo.py @@ -81,6 +81,7 @@ def create_resource( caption: str | None, embedding: list[float] | None, user_data: dict[str, Any], + track: str | None = None, ) -> Resource: res = self._resource_model( url=url, @@ -88,6 +89,7 @@ def create_resource( local_path=local_path, caption=caption, embedding=self._prepare_embedding(embedding), + track=track, **user_data, created_at=self._now(), updated_at=self._now(), diff --git a/src/memu/database/repositories/resource.py b/src/memu/database/repositories/resource.py index 34f4938e..a88ec01e 100644 --- a/src/memu/database/repositories/resource.py +++ b/src/memu/database/repositories/resource.py @@ -27,6 +27,7 @@ def create_resource( caption: str | None, embedding: list[float] | None, user_data: dict[str, Any], + track: str | None = None, ) -> Resource: ... def vector_search_resources( diff --git a/src/memu/database/sqlite/models.py b/src/memu/database/sqlite/models.py index 57ff6fe2..3fd00bb4 100644 --- a/src/memu/database/sqlite/models.py +++ b/src/memu/database/sqlite/models.py @@ -51,6 +51,7 @@ class SQLiteResourceModel(SQLiteBaseModelMixin, Resource): # Override inherited embedding field: SQLite has no native vector type, so store the # vector in a JSON column (a bare ``list`` annotation is not mappable by SQLModel). embedding: list[float] | None = Field(default=None, sa_column=Column(JSON, nullable=True)) + track: str | None = Field(default=None, sa_column=Column(String, nullable=True)) class SQLiteRecallEntryModel(SQLiteBaseModelMixin, RecallEntry): diff --git a/src/memu/database/sqlite/repositories/resource_repo.py b/src/memu/database/sqlite/repositories/resource_repo.py index 9d663c98..9dccb701 100644 --- a/src/memu/database/sqlite/repositories/resource_repo.py +++ b/src/memu/database/sqlite/repositories/resource_repo.py @@ -78,6 +78,7 @@ def list_resources(self, where: Mapping[str, Any] | None = None) -> dict[str, Re local_path=row.local_path, caption=row.caption, embedding=self._normalize_embedding(row.embedding), + track=row.track, created_at=row.created_at, updated_at=row.updated_at, **self._scope_kwargs_from(row), @@ -113,6 +114,7 @@ def clear_resources(self, where: Mapping[str, Any] | None = None) -> dict[str, R local_path=row.local_path, caption=row.caption, embedding=self._normalize_embedding(row.embedding), + track=row.track, created_at=row.created_at, updated_at=row.updated_at, **self._scope_kwargs_from(row), @@ -151,6 +153,7 @@ def create_resource( caption: str | None, embedding: list[float] | None, user_data: dict[str, Any], + track: str | None = None, ) -> Resource: """Create a new resource record. @@ -161,6 +164,7 @@ def create_resource( caption: Optional caption text. embedding: Optional embedding vector. user_data: User scope data. + track: Optional workspace track ("chat"/"skill"/"workspace"). Returns: Created Resource object. @@ -172,6 +176,7 @@ def create_resource( local_path=local_path, caption=caption, embedding=self._prepare_embedding(embedding), + track=track, created_at=now, updated_at=now, **user_data, @@ -188,6 +193,7 @@ def create_resource( local_path=row.local_path, caption=row.caption, embedding=embedding, + track=row.track, created_at=row.created_at, updated_at=row.updated_at, **user_data, diff --git a/tests/test_folder_memorize.py b/tests/test_folder_memorize.py index 2ed342c0..0bc8bc10 100644 --- a/tests/test_folder_memorize.py +++ b/tests/test_folder_memorize.py @@ -159,7 +159,7 @@ async def _noop_categories(*a, **k) -> None: async def _noop_patch(updates, *, ctx, store, llm_client=None) -> None: return None - async def _fake_memorize_one(*, resource_url, modality, user_scope, ctx, store) -> dict[str, Any]: + async def _fake_memorize_one(*, resource_url, modality, user_scope, ctx, store, track=None) -> dict[str, Any]: res = store.resource_repo.create_resource( url=resource_url, modality=modality, @@ -167,6 +167,7 @@ async def _fake_memorize_one(*, resource_url, modality, user_scope, ctx, store) caption="cap", embedding=None, user_data=dict(user_scope or {}), + track=track, ) store.recall_entry_repo.create_item( resource_id=res.id, @@ -225,7 +226,7 @@ async def test_memorize_workspace_exports_when_enabled(tmp_path: Path, monkeypat async def _noop_categories(*a, **k) -> None: return None - async def _fake_memorize_one(*, resource_url, modality, user_scope, ctx, store) -> dict[str, Any]: + async def _fake_memorize_one(*, resource_url, modality, user_scope, ctx, store, track=None) -> dict[str, Any]: res = store.resource_repo.create_resource( url=resource_url, modality=modality, @@ -233,6 +234,7 @@ async def _fake_memorize_one(*, resource_url, modality, user_scope, ctx, store) caption="cap", embedding=None, user_data=dict(user_scope or {}), + track=track, ) return {"resources": [res], "response": {"items": []}} @@ -267,7 +269,7 @@ async def test_memorize_workspace_export_failure_does_not_fail_sync(tmp_path: Pa async def _noop_categories(*a, **k) -> None: return None - async def _fake_memorize_one(*, resource_url, modality, user_scope, ctx, store) -> dict[str, Any]: + async def _fake_memorize_one(*, resource_url, modality, user_scope, ctx, store, track=None) -> dict[str, Any]: res = store.resource_repo.create_resource( url=resource_url, modality=modality, @@ -275,6 +277,7 @@ async def _fake_memorize_one(*, resource_url, modality, user_scope, ctx, store) caption="cap", embedding=None, user_data=dict(user_scope or {}), + track=track, ) return {"resources": [res], "response": {"items": []}} From f9afeedae2b0ff079f00f5c7557a377e2136d388 Mon Sep 17 00:00:00 2001 From: wu Date: Thu, 2 Jul 2026 12:09:40 +0900 Subject: [PATCH 2/9] =?UTF-8?q?feat:=20resource=20=E2=86=92=20file=20works?= =?UTF-8?q?pace=20memorize?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replace the per-file `resource → entry → file` workspace pipeline (plus the `generate_skills` bypass) with a direct `resource → file` synthesis for the chat and skill tracks: - Two-step extraction: (a) route the source to the set of files to update/create given existing files' names+descriptions, then (b) synthesize each target file's body in parallel. No RecallEntry is created on this path. - Track routing: chat → memory-track RecallFile, skill → skill-track RecallFile, workspace → resource-only (no file, no entry). - Add a RecallFileResource(resource_id, category_id) join table mirroring RecallFileEntry across all three backends (inmemory/sqlite/postgres), recording resource → file provenance; cascade-delete unlinks it. - Single-file `memorize` (resource → entry → file) is unchanged. Co-Authored-By: Claude Opus 4.8 (1M context) --- src/memu/app/memorize.py | 339 ++++++++++++++---- src/memu/database/__init__.py | 11 +- src/memu/database/inmemory/__init__.py | 9 +- src/memu/database/inmemory/models.py | 16 +- src/memu/database/inmemory/repo.py | 12 +- .../inmemory/repositories/__init__.py | 6 + .../repositories/recall_file_resource_repo.py | 67 ++++ src/memu/database/interfaces.py | 12 +- src/memu/database/models.py | 17 +- src/memu/database/postgres/models.py | 21 +- src/memu/database/postgres/postgres.py | 24 +- .../postgres/repositories/__init__.py | 2 + .../repositories/recall_file_resource_repo.py | 151 ++++++++ src/memu/database/postgres/schema.py | 9 + src/memu/database/repositories/__init__.py | 3 +- .../repositories/recall_file_resource.py | 33 ++ src/memu/database/sqlite/models.py | 12 +- .../database/sqlite/repositories/__init__.py | 2 + .../repositories/recall_file_resource_repo.py | 217 +++++++++++ src/memu/database/sqlite/schema.py | 9 + src/memu/database/sqlite/sqlite.py | 24 +- src/memu/database/state.py | 3 +- src/memu/prompts/memory_fs/__init__.py | 139 ++++++- tests/test_skill_track.py | 175 ++++++--- 24 files changed, 1170 insertions(+), 143 deletions(-) create mode 100644 src/memu/database/inmemory/repositories/recall_file_resource_repo.py create mode 100644 src/memu/database/postgres/repositories/recall_file_resource_repo.py create mode 100644 src/memu/database/repositories/recall_file_resource.py create mode 100644 src/memu/database/sqlite/repositories/recall_file_resource_repo.py diff --git a/src/memu/app/memorize.py b/src/memu/app/memorize.py index d5fcfc7a..9b0027d3 100644 --- a/src/memu/app/memorize.py +++ b/src/memu/app/memorize.py @@ -7,7 +7,7 @@ import pathlib import re from collections.abc import Awaitable, Callable, Mapping, Sequence -from typing import TYPE_CHECKING, Any, cast +from typing import TYPE_CHECKING, Any, ClassVar, cast from xml.etree.ElementTree import Element import defusedxml.ElementTree as ET @@ -24,9 +24,12 @@ PROMPT as CATEGORY_SUMMARY_PROMPT, ) from memu.prompts.memory_fs import ( - DESCRIPTIONS_PLACEHOLDER, + CONTENT_PLACEHOLDER, + DESCRIPTION_PLACEHOLDER, EXISTING_PLACEHOLDER, - SKILL_FILE_SYNTHESIS_PROMPT, + NAME_PLACEHOLDER, + ROUTE_PROMPTS, + SYNTHESIS_PROMPTS, ) from memu.prompts.memory_type import ( CUSTOM_PROMPTS as MEMORY_TYPE_CUSTOM_PROMPTS, @@ -266,7 +269,8 @@ async def _cascade_delete_by_urls( return [] target_ids = {res.id for res in targets} - # Discarded entry summaries per file, used to recompute summaries. + # Discarded entry summaries per file, used to recompute summaries. Only the + # legacy entry-plane path (single-file ``memorize``) populates these. file_discards: dict[str, list[str]] = {} for entry in store.recall_entry_repo.list_items(where=where).values(): if entry.resource_id not in target_ids: @@ -277,6 +281,11 @@ async def _cascade_delete_by_urls( store.recall_entry_repo.delete_item(entry.id) for res in targets: + # Drop the resource -> file provenance links for the new synthesis path. + # NOTE (ADR 0007 phase 1 open issue): we do not rebuild the affected files + # from their remaining linked resources, so their content may go stale after + # a source change/delete. Tolerated for now. + store.recall_file_resource_repo.unlink_resource(res.id) store.resource_repo.delete_resource(res.id) updates: dict[str, tuple[str | None, str | None]] = { @@ -409,28 +418,73 @@ def _list_memorize_initial_keys() -> set[str]: } def _build_memorize_workspace_workflow(self) -> list[WorkflowStep]: - """The workspace memorize pipeline: the memory steps plus skill generation. - - Identical to :meth:`_build_memorize_workflow` but inserts a per-file - ``generate_skills`` step (ADR 0006) before the response is emitted. It runs - on the ``memorize_workspace`` path only, so single-file ``memorize`` is - unchanged. The skill step has no data dependency on the memory persist - output — it consumes ``preprocessed_resources`` — so it slots in after - persist purely for sequencing. + """The workspace memorize pipeline: direct resource -> file synthesis (ADR 0007 phase 1). + + Unlike single-file :meth:`memorize` (``resource -> entry -> file``), the + workspace path synthesizes files straight from the preprocessed source and + creates no ``RecallEntry``. After ``preprocess`` it: + + - ``create_resource`` — one file maps to one :class:`Resource` (caption/embedding + for INDEX recall), for every track including ``workspace`` (resource-only). + - ``synthesize_files`` — for the ``chat`` and ``skill`` tracks only, route the + source to the files to update/create then synthesize each file's body, upserting + ``RecallFile`` and recording ``resource -> file`` provenance. ``workspace`` is a + no-op here. Retrieval over these files is deferred (ADR 0007 phase 2). """ - steps = self._build_memorize_workflow() - skill_step = WorkflowStep( - step_id="generate_skills", - role="generate_skills", - handler=self._memorize_generate_skills, - requires={"preprocessed_resources", "store", "user"}, - produces={"skills"}, - capabilities={"llm", "db", "vector"}, - config={"chat_llm_profile": getattr(self.memory_files_config, "synthesis_llm_profile", "default")}, - ) - # Insert just before the terminal build_response step. - steps.insert(-1, skill_step) - return steps + synthesis_profile = getattr(self.memory_files_config, "synthesis_llm_profile", "default") + return [ + WorkflowStep( + step_id="ingest_resource", + role="ingest", + handler=self._memorize_ingest_resource, + requires={"resource_url", "modality"}, + produces={"local_path", "raw_text"}, + capabilities={"io"}, + ), + WorkflowStep( + step_id="preprocess_multimodal", + role="preprocess", + handler=self._memorize_preprocess_multimodal, + requires={"local_path", "modality", "raw_text"}, + produces={"preprocessed_resources"}, + capabilities={"llm"}, + config={"chat_llm_profile": self.memorize_config.preprocess_llm_profile}, + ), + WorkflowStep( + step_id="create_resource", + role="persist", + handler=self._memorize_ws_create_resource, + requires={ + "preprocessed_resources", + "modality", + "local_path", + "resource_url", + "store", + "user", + "resource_track", + }, + produces={"resources"}, + capabilities={"db", "vector"}, + config={"embed_llm_profile": "embedding"}, + ), + WorkflowStep( + step_id="synthesize_files", + role="synthesize_files", + handler=self._memorize_ws_synthesize_files, + requires={"resources", "preprocessed_resources", "resource_track", "store", "user"}, + produces={"files"}, + capabilities={"llm", "db", "vector"}, + config={"chat_llm_profile": synthesis_profile, "embed_llm_profile": "embedding"}, + ), + WorkflowStep( + step_id="build_response", + role="emit", + handler=self._memorize_ws_build_response, + requires={"resources", "files"}, + produces={"response"}, + capabilities=set(), + ), + ] async def _memorize_ingest_resource(self, state: WorkflowState, step_context: Any) -> WorkflowState: local_path, raw_text = await self.fs.fetch(state["resource_url"], state["modality"]) @@ -561,69 +615,187 @@ async def _memorize_persist_and_index(self, state: WorkflowState, step_context: ) return state - async def _memorize_generate_skills(self, state: WorkflowState, step_context: Any) -> WorkflowState: - """Generate/patch skill-track ``RecallFile``s from this file's content (ADR 0006). + @staticmethod + def _format_skill_source_content(preprocessed_resources: list[dict[str, Any]]) -> str: + """Flatten a source's preprocessed segments into a single text block.""" + parts = [ + " ".join((prep.get("text") or "").split()) + for prep in preprocessed_resources + if (prep.get("text") or "").strip() + ] + return "\n\n".join(parts) + + # --- Workspace resource -> file path (ADR 0007 phase 1) ------------------- + + # Maps a workspace ``resource_track`` to the ``RecallFile.track`` it synthesizes + # into. ``workspace`` has no entry (resource-only), so it is absent. + _TRACK_TO_FILE_TRACK: ClassVar[dict[str, str]] = {"chat": "memory", "skill": "skill"} + + async def _memorize_ws_create_resource(self, state: WorkflowState, step_context: Any) -> WorkflowState: + """Create the single ``Resource`` for this file (one file -> one resource). + + Runs for every track; the ``workspace`` track stops here (resource-only). The + caption is the joined per-segment captions, embedded for INDEX/resource recall. + """ + embed_client = self._get_step_embedding_client(step_context) + store = state["store"] + preprocessed = state.get("preprocessed_resources") or [] + captions = [(prep.get("caption") or "").strip() for prep in preprocessed] + caption = "\n\n".join(c for c in captions if c) or None + res = await self._create_resource_with_caption( + resource_url=state["resource_url"], + modality=state["modality"], + local_path=state["local_path"], + caption=caption, + store=store, + embed_client=embed_client, + user=state.get("user", {}), + track=state.get("resource_track"), + ) + state["resources"] = [res] + return state + + async def _memorize_ws_synthesize_files(self, state: WorkflowState, step_context: Any) -> WorkflowState: + """Synthesize this source into ``RecallFile``s for the chat/skill tracks. - Gated behind ``memory_files_config.synthesize``. Reads the preprocessed - content of the current source plus the existing skill-track files (so file - *N* sees skills created by files *1..N-1*), asks the LLM for skills to add or - revise, and persists each directly as a ``RecallFile(track="skill")`` — - embedding ``name + description`` and storing the body as ``content``, bypassing - the ``RecallEntry`` plane. + Two steps: (a) route the source to the set of files to update/create given the + existing files' names+descriptions, and (b) synthesize each target file's body in + parallel. Persists each file and a ``resource -> file`` provenance link. The + ``workspace`` track (and any source with no content) is a no-op. """ - if not getattr(self.memory_files_config, "synthesize", False): - return state + track = state.get("resource_track") + file_track = self._TRACK_TO_FILE_TRACK.get(track or "") + resources = state.get("resources") or [] content = self._format_skill_source_content(state.get("preprocessed_resources") or []) - if not content: + if file_track is None or not resources or not content: + state["files"] = [] return state store = state["store"] user_scope = dict(state.get("user") or {}) llm_client = self._get_step_llm_client(step_context) embed_client = self._get_step_embedding_client(step_context) + resource = resources[0] - existing = store.recall_file_repo.list_categories(where={**user_scope, "track": "skill"}) - existing_text = self._format_existing_skills(existing) or "(none)" - prompt = SKILL_FILE_SYNTHESIS_PROMPT.replace(EXISTING_PLACEHOLDER, existing_text).replace( - DESCRIPTIONS_PLACEHOLDER, self._escape_prompt_value(content) + existing = store.recall_file_repo.list_categories(where={**user_scope, "track": file_track}) + ops = await self._route_source_to_files( + file_track=file_track, content=content, existing=existing, llm_client=llm_client ) - parsed = self._parse_skill_files(await llm_client.chat(prompt)) - - persisted: list[RecallFile] = [] - for name, description, body in parsed: - emb_text = f"{name}: {description}" if description else name - embedding = (await embed_client.embed([emb_text]))[0] - skill = store.recall_file_repo.get_or_create_category( - name=name, - description=description, - embedding=embedding, - user_data=user_scope, - track="skill", - ) - persisted.append(store.recall_file_repo.update_category(category_id=skill.id, content=body)) - state["skills"] = persisted + touched = await self._synthesize_file_ops( + ops=ops, + file_track=file_track, + content=content, + existing=existing, + resource=resource, + store=store, + user_scope=user_scope, + llm_client=llm_client, + embed_client=embed_client, + ) + state["files"] = touched return state - @staticmethod - def _format_skill_source_content(preprocessed_resources: list[dict[str, Any]]) -> str: - """Flatten a source's preprocessed segments into a single text block.""" - parts = [ - " ".join((prep.get("text") or "").split()) - for prep in preprocessed_resources - if (prep.get("text") or "").strip() + async def _route_source_to_files( + self, + *, + file_track: str, + content: str, + existing: Mapping[str, RecallFile], + llm_client: Any, + ) -> list[dict[str, str]]: + """Ask the model which existing files to update / what new files to create.""" + existing_text = self._format_existing_files(existing) or "(none)" + prompt = ( + ROUTE_PROMPTS[file_track] + .replace(EXISTING_PLACEHOLDER, existing_text) + .replace(CONTENT_PLACEHOLDER, self._escape_prompt_value(content)) + ) + return self._parse_file_ops(await llm_client.chat(prompt), existing) + + async def _synthesize_file_ops( + self, + *, + ops: list[dict[str, str]], + file_track: str, + content: str, + existing: Mapping[str, RecallFile], + resource: Resource, + store: Database, + user_scope: dict[str, Any], + llm_client: Any, + embed_client: Any, + ) -> list[RecallFile]: + """Synthesize each routed file's body (in parallel) and persist file + link.""" + existing_by_name = {f.name: f for f in existing.values()} + # Resolve ops to unique targets (dedup by name; last op's description wins). + targets: list[dict[str, Any]] = [] + by_name: dict[str, dict[str, Any]] = {} + for op in ops: + name = op["name"] + ex = existing_by_name.get(name) + description = (op.get("description") or (ex.description if ex else "") or "").strip() + target = by_name.get(name) + if target is None: + target = {"name": name, "description": description, "existing": ex} + by_name[name] = target + targets.append(target) + elif description: + target["description"] = description + if not targets: + return [] + + prompts = [ + SYNTHESIS_PROMPTS[file_track] + .replace(NAME_PLACEHOLDER, self._escape_prompt_value(t["name"])) + .replace(DESCRIPTION_PLACEHOLDER, self._escape_prompt_value(t["description"])) + .replace( + EXISTING_PLACEHOLDER, self._escape_prompt_value((t["existing"].content if t["existing"] else "") or "") + ) + .replace(CONTENT_PLACEHOLDER, self._escape_prompt_value(content)) + for t in targets ] - return "\n\n".join(parts) + bodies = await asyncio.gather(*[llm_client.chat(prompt) for prompt in prompts]) + + # Embed name+description for the files being created. + creates = [t for t in targets if t["existing"] is None] + create_vecs: dict[str, list[float]] = {} + if creates: + emb_texts = [f"{t['name']}: {t['description']}" if t["description"] else t["name"] for t in creates] + vecs = await embed_client.embed(emb_texts) + for t, vec in zip(creates, vecs, strict=True): + create_vecs[t["name"]] = vec + + touched: list[RecallFile] = [] + for target, body in zip(targets, bodies, strict=True): + cleaned = body.replace("```markdown", "").replace("```", "").strip() + file = target["existing"] + if file is None: + file = store.recall_file_repo.get_or_create_category( + name=target["name"], + description=target["description"], + embedding=create_vecs[target["name"]], + user_data=user_scope, + track=file_track, + ) + file = store.recall_file_repo.update_category(category_id=file.id, content=cleaned) + store.recall_file_resource_repo.link_resource_category(resource.id, file.id, user_data=dict(user_scope)) + touched.append(file) + return touched @staticmethod - def _format_existing_skills(existing: Mapping[str, RecallFile]) -> str: - """Render existing skill files as ``## name\\nbody`` blocks for the prompt.""" - return "\n\n".join( - f"## {skill.name}\n{(skill.content or '').strip()}".strip() - for skill in sorted(existing.values(), key=lambda s: s.name) + def _format_existing_files(existing: Mapping[str, RecallFile]) -> str: + """Render existing files as ``- name: description`` lines for the router prompt.""" + return "\n".join( + f"- {f.name}: {f.description}" if f.description else f"- {f.name}" + for f in sorted(existing.values(), key=lambda f: f.name) ) - def _parse_skill_files(self, raw: str) -> list[tuple[str, str, str]]: - """Parse the skill-synthesis JSON array into ``(name, description, body)`` tuples.""" + def _parse_file_ops(self, raw: str, existing: Mapping[str, RecallFile]) -> list[dict[str, str]]: + """Parse the router's JSON array into validated ``{op, name, description}`` dicts. + + ``update`` ops naming an unknown file are dropped (we never update a file that + does not exist); ``create``/``update`` are otherwise kept with a stripped name. + """ if not raw: return [] start = raw.find("[") @@ -636,18 +808,35 @@ def _parse_skill_files(self, raw: str) -> list[tuple[str, str, str]]: return [] if not isinstance(parsed, list): return [] - skills: list[tuple[str, str, str]] = [] + existing_names = {f.name for f in existing.values()} + ops: list[dict[str, str]] = [] for entry in parsed: if not isinstance(entry, dict): continue + op = entry.get("op") name = entry.get("name") - body = entry.get("body") - if not isinstance(name, str) or not name.strip() or not isinstance(body, str) or not body.strip(): + if op not in {"update", "create"} or not isinstance(name, str) or not name.strip(): + continue + name = name.strip() + if op == "update" and name not in existing_names: continue description = entry.get("description") description = description.strip() if isinstance(description, str) else "" - skills.append((name.strip(), description, body.strip())) - return skills + ops.append({"op": op, "name": name, "description": description}) + return ops + + def _memorize_ws_build_response(self, state: WorkflowState, step_context: Any) -> WorkflowState: + """Emit the workspace response (no entries; ``categories`` carries touched files).""" + resources = [self._model_dump_without_embeddings(r) for r in state.get("resources", [])] + files = [self._model_dump_without_embeddings(f) for f in state.get("files", [])] + # Keep the legacy response contract (``items``/``categories``); items is always + # empty on this path since the entry plane is gone. + base: dict[str, Any] = {"items": [], "categories": files, "relations": []} + if len(resources) == 1: + state["response"] = {"resource": resources[0], **base} + else: + state["response"] = {"resources": resources, **base} + return state def _memorize_build_response(self, state: WorkflowState, step_context: Any) -> WorkflowState: ctx = state["ctx"] diff --git a/src/memu/database/__init__.py b/src/memu/database/__init__.py index ac78d7f9..33b1c168 100644 --- a/src/memu/database/__init__.py +++ b/src/memu/database/__init__.py @@ -6,9 +6,16 @@ RecallEntryRecord, RecallFileEntryRecord, RecallFileRecord, + RecallFileResourceRecord, ResourceRecord, ) -from memu.database.repositories import RecallEntryRepo, RecallFileEntryRepo, RecallFileRepo, ResourceRepo +from memu.database.repositories import ( + RecallEntryRepo, + RecallFileEntryRepo, + RecallFileRepo, + RecallFileResourceRepo, + ResourceRepo, +) __all__ = [ "Database", @@ -18,6 +25,8 @@ "RecallFileEntryRepo", "RecallFileRecord", "RecallFileRepo", + "RecallFileResourceRecord", + "RecallFileResourceRepo", "ResourceRecord", "ResourceRepo", "build_database", diff --git a/src/memu/database/inmemory/__init__.py b/src/memu/database/inmemory/__init__.py index 083381da..2d7f267b 100644 --- a/src/memu/database/inmemory/__init__.py +++ b/src/memu/database/inmemory/__init__.py @@ -12,13 +12,20 @@ def build_inmemory_database( config: DatabaseConfig, user_model: type[BaseModel], ) -> InMemoryStore: - resource_model, recall_file_model, recall_entry_model, recall_file_entry_model = build_inmemory_models(user_model) + ( + resource_model, + recall_file_model, + recall_entry_model, + recall_file_entry_model, + recall_file_resource_model, + ) = build_inmemory_models(user_model) return InMemoryStore( scope_model=user_model, resource_model=resource_model, recall_entry_model=recall_entry_model, recall_file_model=recall_file_model, recall_file_entry_model=recall_file_entry_model, + recall_file_resource_model=recall_file_resource_model, ) diff --git a/src/memu/database/inmemory/models.py b/src/memu/database/inmemory/models.py index 648fe0fe..1a61fe01 100644 --- a/src/memu/database/inmemory/models.py +++ b/src/memu/database/inmemory/models.py @@ -6,6 +6,7 @@ RecallEntry, RecallFile, RecallFileEntry, + RecallFileResource, Resource, merge_scope_model, ) @@ -27,6 +28,10 @@ class InMemoryFileEntry(RecallFileEntry): """Concrete in-memory relation model.""" +class InMemoryFileResource(RecallFileResource): + """Concrete in-memory resource-category relation model.""" + + def build_inmemory_models( user_model: type[BaseModel], ) -> tuple[ @@ -34,6 +39,7 @@ def build_inmemory_models( type[InMemoryRecallFile], type[InMemoryRecallEntry], type[InMemoryFileEntry], + type[InMemoryFileResource], ]: """ Build scoped in-memory models that inherit from both the base interface and the user scope model. @@ -42,11 +48,19 @@ def build_inmemory_models( recall_file_model = merge_scope_model(user_model, InMemoryRecallFile, name_suffix="RecallFile") recall_entry_model = merge_scope_model(user_model, InMemoryRecallEntry, name_suffix="RecallEntry") recall_file_entry_model = merge_scope_model(user_model, InMemoryFileEntry, name_suffix="RecallFileEntry") - return resource_model, recall_file_model, recall_entry_model, recall_file_entry_model + recall_file_resource_model = merge_scope_model(user_model, InMemoryFileResource, name_suffix="RecallFileResource") + return ( + resource_model, + recall_file_model, + recall_entry_model, + recall_file_entry_model, + recall_file_resource_model, + ) __all__ = [ "InMemoryFileEntry", + "InMemoryFileResource", "InMemoryRecallEntry", "InMemoryRecallFile", "InMemoryResource", diff --git a/src/memu/database/inmemory/repo.py b/src/memu/database/inmemory/repo.py index 4cfefeda..44d1d155 100644 --- a/src/memu/database/inmemory/repo.py +++ b/src/memu/database/inmemory/repo.py @@ -7,13 +7,14 @@ from memu.database.inmemory.models import build_inmemory_models from memu.database.inmemory.repositories import ( InMemoryFileEntryRepository, + InMemoryFileResourceRepository, InMemoryRecallEntryRepository, InMemoryRecallFileRepository, InMemoryResourceRepository, ) from memu.database.inmemory.state import InMemoryState from memu.database.interfaces import Database -from memu.database.models import RecallEntry, RecallFile, RecallFileEntry, Resource +from memu.database.models import RecallEntry, RecallFile, RecallFileEntry, RecallFileResource, Resource from memu.database.repositories import RecallFileRepo, ResourceRepo @@ -26,6 +27,7 @@ def __init__( recall_entry_model: type[Any] | None = None, recall_file_model: type[Any] | None = None, recall_file_entry_model: type[Any] | None = None, + recall_file_resource_model: type[Any] | None = None, state: InMemoryState | None = None, ) -> None: self.scope_model = scope_model or BaseModel @@ -34,6 +36,7 @@ def __init__( default_recall_file_model, default_recall_entry_model, default_recall_file_entry_model, + default_recall_file_resource_model, ) = build_inmemory_models(self.scope_model) self.state = state or InMemoryState() @@ -41,11 +44,15 @@ def __init__( self.items: dict[str, RecallEntry] = self.state.items self.categories: dict[str, RecallFile] = self.state.categories self.relations: list[RecallFileEntry] = self.state.relations + self.resource_relations: list[RecallFileResource] = self.state.resource_relations resource_model = resource_model or default_resource_model or Resource recall_entry_model = recall_entry_model or default_recall_entry_model or RecallEntry recall_file_model = recall_file_model or default_recall_file_model or RecallFile recall_file_entry_model = recall_file_entry_model or default_recall_file_entry_model or RecallFileEntry + recall_file_resource_model = ( + recall_file_resource_model or default_recall_file_resource_model or RecallFileResource + ) self.resource_repo: ResourceRepo = InMemoryResourceRepository(state=self.state, resource_model=resource_model) self.recall_file_repo: RecallFileRepo = InMemoryRecallFileRepository( @@ -55,6 +62,9 @@ def __init__( self.recall_file_entry_repo = InMemoryFileEntryRepository( state=self.state, recall_file_entry_model=recall_file_entry_model ) + self.recall_file_resource_repo = InMemoryFileResourceRepository( + state=self.state, recall_file_resource_model=recall_file_resource_model + ) def close(self) -> None: return None diff --git a/src/memu/database/inmemory/repositories/__init__.py b/src/memu/database/inmemory/repositories/__init__.py index f835369c..ff79d762 100644 --- a/src/memu/database/inmemory/repositories/__init__.py +++ b/src/memu/database/inmemory/repositories/__init__.py @@ -7,15 +7,21 @@ InMemoryRecallFileRepository, RecallFileRepo, ) +from memu.database.inmemory.repositories.recall_file_resource_repo import ( + InMemoryFileResourceRepository, + RecallFileResourceRepo, +) from memu.database.inmemory.repositories.resource_repo import InMemoryResourceRepository, ResourceRepo __all__ = [ "InMemoryFileEntryRepository", + "InMemoryFileResourceRepository", "InMemoryRecallEntryRepository", "InMemoryRecallFileRepository", "InMemoryResourceRepository", "RecallEntryRepo", "RecallFileEntryRepo", "RecallFileRepo", + "RecallFileResourceRepo", "ResourceRepo", ] diff --git a/src/memu/database/inmemory/repositories/recall_file_resource_repo.py b/src/memu/database/inmemory/repositories/recall_file_resource_repo.py new file mode 100644 index 00000000..cccc251e --- /dev/null +++ b/src/memu/database/inmemory/repositories/recall_file_resource_repo.py @@ -0,0 +1,67 @@ +from __future__ import annotations + +import uuid +from collections.abc import Mapping +from typing import Any, override + +from memu.database.inmemory.repositories.filter import matches_where +from memu.database.inmemory.state import InMemoryState +from memu.database.models import RecallFileResource +from memu.database.repositories.recall_file_resource import RecallFileResourceRepo + + +class InMemoryFileResourceRepository(RecallFileResourceRepo): + def __init__(self, *, state: InMemoryState, recall_file_resource_model: type[RecallFileResource]) -> None: + self._state = state + self.recall_file_resource_model = recall_file_resource_model + self.relations: list[RecallFileResource] = self._state.resource_relations + + def list_relations(self, where: Mapping[str, Any] | None = None) -> list[RecallFileResource]: + if not where: + return list(self.relations) + return [rel for rel in self.relations if matches_where(rel, where)] + + def link_resource_category(self, resource_id: str, cat_id: str, user_data: dict[str, Any]) -> RecallFileResource: + _ = resource_id # enforced by caller via existing state + for rel in self.relations: + if rel.resource_id == resource_id and rel.category_id == cat_id: + return rel + rel = self.recall_file_resource_model( + id=str(uuid.uuid4()), resource_id=resource_id, category_id=cat_id, **user_data + ) + self.relations.append(rel) + return rel + + def load_existing(self) -> None: + return None + + @override + def get_resource_categories(self, resource_id: str) -> list[RecallFileResource]: + return [rel for rel in self.relations if rel.resource_id == resource_id] + + @override + def unlink_resource_category(self, resource_id: str, cat_id: str) -> None: + # Mutate the shared state list in place so the DatabaseState reference and + # this repo's view never diverge (rebinding self.relations would orphan the + # shared state.resource_relations list). + self.relations[:] = [ + rel for rel in self.relations if not (rel.resource_id == resource_id and rel.category_id == cat_id) + ] + + def unlink_resource(self, resource_id: str) -> list[RecallFileResource]: + removed = [rel for rel in self.relations if rel.resource_id == resource_id] + self.relations[:] = [rel for rel in self.relations if rel.resource_id != resource_id] + return removed + + def clear_relations(self, where: Mapping[str, Any] | None = None) -> list[RecallFileResource]: + if not where: + removed = list(self.relations) + self.relations.clear() + return removed + removed = [rel for rel in self.relations if matches_where(rel, where)] + removed_ids = {rel.id for rel in removed} + self.relations[:] = [rel for rel in self.relations if rel.id not in removed_ids] + return removed + + +__all__ = ["InMemoryFileResourceRepository"] diff --git a/src/memu/database/interfaces.py b/src/memu/database/interfaces.py index 52b9aa99..20b732e4 100644 --- a/src/memu/database/interfaces.py +++ b/src/memu/database/interfaces.py @@ -5,8 +5,15 @@ from memu.database.models import RecallEntry as RecallEntryRecord from memu.database.models import RecallFile as RecallFileRecord from memu.database.models import RecallFileEntry as RecallFileEntryRecord +from memu.database.models import RecallFileResource as RecallFileResourceRecord from memu.database.models import Resource as ResourceRecord -from memu.database.repositories import RecallEntryRepo, RecallFileEntryRepo, RecallFileRepo, ResourceRepo +from memu.database.repositories import ( + RecallEntryRepo, + RecallFileEntryRepo, + RecallFileRepo, + RecallFileResourceRepo, + ResourceRepo, +) @runtime_checkable @@ -17,11 +24,13 @@ class Database(Protocol): recall_file_repo: RecallFileRepo recall_entry_repo: RecallEntryRepo recall_file_entry_repo: RecallFileEntryRepo + recall_file_resource_repo: RecallFileResourceRepo resources: dict[str, ResourceRecord] items: dict[str, RecallEntryRecord] categories: dict[str, RecallFileRecord] relations: list[RecallFileEntryRecord] + resource_relations: list[RecallFileResourceRecord] def close(self) -> None: ... @@ -31,5 +40,6 @@ def close(self) -> None: ... "RecallEntryRecord", "RecallFileEntryRecord", "RecallFileRecord", + "RecallFileResourceRecord", "ResourceRecord", ] diff --git a/src/memu/database/models.py b/src/memu/database/models.py index 23361d5c..ed7d2596 100644 --- a/src/memu/database/models.py +++ b/src/memu/database/models.py @@ -111,6 +111,11 @@ class RecallFileEntry(BaseRecord): category_id: str +class RecallFileResource(BaseRecord): + resource_id: str + category_id: str + + def merge_scope_model[TBaseRecord: BaseRecord]( user_model: type[BaseModel], core_model: type[TBaseRecord], *, name_suffix: str ) -> type[TBaseRecord]: @@ -129,7 +134,7 @@ def merge_scope_model[TBaseRecord: BaseRecord]( def build_scoped_models( user_model: type[BaseModel], -) -> tuple[type[Resource], type[RecallFile], type[RecallEntry], type[RecallFileEntry]]: +) -> tuple[type[Resource], type[RecallFile], type[RecallEntry], type[RecallFileEntry], type[RecallFileResource]]: """ Build scoped interface models (Pydantic) that inherit from the base record models and user scope. """ @@ -137,7 +142,14 @@ def build_scoped_models( recall_file_model = merge_scope_model(user_model, RecallFile, name_suffix="RecallFile") recall_entry_model = merge_scope_model(user_model, RecallEntry, name_suffix="RecallEntry") recall_file_entry_model = merge_scope_model(user_model, RecallFileEntry, name_suffix="RecallFileEntry") - return resource_model, recall_file_model, recall_entry_model, recall_file_entry_model + recall_file_resource_model = merge_scope_model(user_model, RecallFileResource, name_suffix="RecallFileResource") + return ( + resource_model, + recall_file_model, + recall_entry_model, + recall_file_entry_model, + recall_file_resource_model, + ) __all__ = [ @@ -146,6 +158,7 @@ def build_scoped_models( "RecallEntry", "RecallFile", "RecallFileEntry", + "RecallFileResource", "Resource", "ToolCallResult", "build_scoped_models", diff --git a/src/memu/database/postgres/models.py b/src/memu/database/postgres/models.py index be8c11eb..e6628f34 100644 --- a/src/memu/database/postgres/models.py +++ b/src/memu/database/postgres/models.py @@ -17,7 +17,7 @@ from sqlalchemy.dialects.postgresql import JSONB from sqlmodel import Column, DateTime, Field, Index, SQLModel, func -from memu.database.models import EntryType, RecallEntry, RecallFile, RecallFileEntry, Resource +from memu.database.models import EntryType, RecallEntry, RecallFile, RecallFileEntry, RecallFileResource, Resource class TZDateTime(DateTime): @@ -76,6 +76,13 @@ class RecallFileEntryModel(BaseModelMixin, RecallFileEntry): __table_args__ = (Index("idx_recall_file_entries_unique", "item_id", "category_id", unique=True),) +class RecallFileResourceModel(BaseModelMixin, RecallFileResource): + resource_id: str = Field(sa_column=Column(ForeignKey("resources.id", ondelete="CASCADE"), nullable=False)) + category_id: str = Field(sa_column=Column(ForeignKey("memory_categories.id", ondelete="CASCADE"), nullable=False)) + + __table_args__ = (Index("idx_recall_file_resources_unique", "resource_id", "category_id", unique=True),) + + def _normalize_table_args(table_args: Any) -> tuple[list[Any], dict[str, Any]]: if table_args is None: return [], {} @@ -158,7 +165,7 @@ def build_table_model( def build_scoped_models( user_model: type[BaseModel], -) -> tuple[type[SQLModel], type[SQLModel], type[SQLModel], type[SQLModel]]: +) -> tuple[type[SQLModel], type[SQLModel], type[SQLModel], type[SQLModel], type[SQLModel]]: """ Build scoped SQLModel tables for each entity (resource, category, item, relation). """ @@ -168,7 +175,14 @@ def build_scoped_models( ) recall_entry_model = build_table_model(user_model, RecallEntryModel, tablename="memory_items") recall_file_entry_model = build_table_model(user_model, RecallFileEntryModel, tablename="category_items") - return resource_model, recall_file_model, recall_entry_model, recall_file_entry_model + recall_file_resource_model = build_table_model(user_model, RecallFileResourceModel, tablename="resource_categories") + return ( + resource_model, + recall_file_model, + recall_entry_model, + recall_file_entry_model, + recall_file_resource_model, + ) __all__ = [ @@ -176,6 +190,7 @@ def build_scoped_models( "RecallEntryModel", "RecallFileEntryModel", "RecallFileModel", + "RecallFileResourceModel", "ResourceModel", "build_scoped_models", "build_table_model", diff --git a/src/memu/database/postgres/postgres.py b/src/memu/database/postgres/postgres.py index cf168dbc..20b56853 100644 --- a/src/memu/database/postgres/postgres.py +++ b/src/memu/database/postgres/postgres.py @@ -6,15 +6,22 @@ from pydantic import BaseModel from memu.database.interfaces import Database -from memu.database.models import RecallEntry, RecallFile, RecallFileEntry, Resource +from memu.database.models import RecallEntry, RecallFile, RecallFileEntry, RecallFileResource, Resource from memu.database.postgres.migration import DDLMode, run_migrations from memu.database.postgres.repositories.recall_entry_repo import PostgresRecallEntryRepo from memu.database.postgres.repositories.recall_file_entry_repo import PostgresRecallFileEntryRepo from memu.database.postgres.repositories.recall_file_repo import PostgresRecallFileRepo +from memu.database.postgres.repositories.recall_file_resource_repo import PostgresRecallFileResourceRepo from memu.database.postgres.repositories.resource_repo import PostgresResourceRepo from memu.database.postgres.schema import SQLAModels, get_sqlalchemy_models, require_sqlalchemy from memu.database.postgres.session import SessionManager -from memu.database.repositories import RecallEntryRepo, RecallFileEntryRepo, RecallFileRepo, ResourceRepo +from memu.database.repositories import ( + RecallEntryRepo, + RecallFileEntryRepo, + RecallFileRepo, + RecallFileResourceRepo, + ResourceRepo, +) from memu.database.state import DatabaseState logger = logging.getLogger(__name__) @@ -25,10 +32,12 @@ class PostgresStore(Database): recall_file_repo: RecallFileRepo recall_entry_repo: RecallEntryRepo recall_file_entry_repo: RecallFileEntryRepo + recall_file_resource_repo: RecallFileResourceRepo resources: dict[str, Resource] items: dict[str, RecallEntry] categories: dict[str, RecallFile] relations: list[RecallFileEntry] + resource_relations: list[RecallFileResource] def __init__( self, @@ -42,6 +51,7 @@ def __init__( recall_file_model: type[Any] | None = None, recall_entry_model: type[Any] | None = None, recall_file_entry_model: type[Any] | None = None, + recall_file_resource_model: type[Any] | None = None, sqla_models: SQLAModels | None = None, ) -> None: require_sqlalchemy() @@ -60,6 +70,7 @@ def __init__( recall_file_model = recall_file_model or self._sqla_models.RecallFile recall_entry_model = recall_entry_model or self._sqla_models.RecallEntry recall_file_entry_model = recall_file_entry_model or self._sqla_models.RecallFileEntry + recall_file_resource_model = recall_file_resource_model or self._sqla_models.RecallFileResource self.resource_repo = PostgresResourceRepo( state=self._state, @@ -90,11 +101,19 @@ def __init__( sessions=self._sessions, scope_fields=self._scope_fields, ) + self.recall_file_resource_repo = PostgresRecallFileResourceRepo( + state=self._state, + recall_file_resource_model=recall_file_resource_model, + sqla_models=self._sqla_models, + sessions=self._sessions, + scope_fields=self._scope_fields, + ) self.resources = self._state.resources self.items = self._state.items self.categories = self._state.categories self.relations = self._state.relations + self.resource_relations = self._state.resource_relations # self._load_existing() @@ -106,3 +125,4 @@ def _load_existing(self) -> None: self.recall_file_repo.load_existing() self.recall_entry_repo.load_existing() self.recall_file_entry_repo.load_existing() + self.recall_file_resource_repo.load_existing() diff --git a/src/memu/database/postgres/repositories/__init__.py b/src/memu/database/postgres/repositories/__init__.py index 59bf0247..eea2e475 100644 --- a/src/memu/database/postgres/repositories/__init__.py +++ b/src/memu/database/postgres/repositories/__init__.py @@ -1,11 +1,13 @@ from memu.database.postgres.repositories.recall_entry_repo import PostgresRecallEntryRepo from memu.database.postgres.repositories.recall_file_entry_repo import PostgresRecallFileEntryRepo from memu.database.postgres.repositories.recall_file_repo import PostgresRecallFileRepo +from memu.database.postgres.repositories.recall_file_resource_repo import PostgresRecallFileResourceRepo from memu.database.postgres.repositories.resource_repo import PostgresResourceRepo __all__ = [ "PostgresRecallEntryRepo", "PostgresRecallFileEntryRepo", "PostgresRecallFileRepo", + "PostgresRecallFileResourceRepo", "PostgresResourceRepo", ] diff --git a/src/memu/database/postgres/repositories/recall_file_resource_repo.py b/src/memu/database/postgres/repositories/recall_file_resource_repo.py new file mode 100644 index 00000000..8045a5ca --- /dev/null +++ b/src/memu/database/postgres/repositories/recall_file_resource_repo.py @@ -0,0 +1,151 @@ +from __future__ import annotations + +from collections.abc import Mapping +from typing import Any + +from memu.database.models import RecallFileResource +from memu.database.postgres.repositories.base import PostgresRepoBase +from memu.database.postgres.session import SessionManager +from memu.database.repositories.recall_file_resource import RecallFileResourceRepo +from memu.database.state import DatabaseState + + +class PostgresRecallFileResourceRepo(PostgresRepoBase, RecallFileResourceRepo): + def __init__( + self, + *, + state: DatabaseState, + recall_file_resource_model: type[RecallFileResource], + sqla_models: Any, + sessions: SessionManager, + scope_fields: list[str], + ) -> None: + super().__init__(state=state, sqla_models=sqla_models, sessions=sessions, scope_fields=scope_fields) + self._recall_file_resource_model = recall_file_resource_model + self.relations: list[RecallFileResource] = self._state.resource_relations + + def list_relations(self, where: Mapping[str, Any] | None = None) -> list[RecallFileResource]: + from sqlmodel import select + + filters = self._build_filters(self._sqla_models.RecallFileResource, where) + with self._sessions.session() as session: + rows = session.scalars(select(self._sqla_models.RecallFileResource).where(*filters)).all() + return [self._cache_relation(row) for row in rows] + + def link_resource_category(self, resource_id: str, cat_id: str, user_data: dict[str, Any]) -> RecallFileResource: + from sqlmodel import select + + # Avoid duplicate inserts using local cache + for rel in self.relations: + if rel.resource_id == resource_id and rel.category_id == cat_id: + return rel + + now = self._now() + new_rel = self._recall_file_resource_model( + resource_id=resource_id, + category_id=cat_id, + **user_data, + created_at=now, + updated_at=now, + ) + + with self._sessions.session() as session: + existing = session.scalar( + select(self._sqla_models.RecallFileResource).where( + self._sqla_models.RecallFileResource.resource_id == resource_id, + self._sqla_models.RecallFileResource.category_id == cat_id, + ) + ) + if existing: + return self._cache_relation(existing) + + session.add(new_rel) + session.commit() + session.refresh(new_rel) + + return self._cache_relation(new_rel) + + def unlink_resource_category(self, resource_id: str, cat_id: str) -> None: + from sqlmodel import delete + + with self._sessions.session() as session: + session.exec( + delete(self._sqla_models.RecallFileResource).where( + self._sqla_models.RecallFileResource.resource_id == resource_id, + self._sqla_models.RecallFileResource.category_id == cat_id, + ) + ) + session.commit() + self.relations[:] = [ + r for r in self.relations if not (r.resource_id == resource_id and r.category_id == cat_id) + ] + + def _row_to_record(self, row: Any) -> RecallFileResource: + return RecallFileResource( + id=row.id, + resource_id=row.resource_id, + category_id=row.category_id, + created_at=row.created_at, + updated_at=row.updated_at, + **self._scope_kwargs_from(row), + ) + + def unlink_resource(self, resource_id: str) -> list[RecallFileResource]: + from sqlmodel import delete, select + + with self._sessions.session() as session: + rows = session.scalars( + select(self._sqla_models.RecallFileResource).where( + self._sqla_models.RecallFileResource.resource_id == resource_id + ) + ).all() + removed = [self._row_to_record(row) for row in rows] + if removed: + session.exec( + delete(self._sqla_models.RecallFileResource).where( + self._sqla_models.RecallFileResource.resource_id == resource_id + ) + ) + session.commit() + self.relations[:] = [r for r in self.relations if r.resource_id != resource_id] + return removed + + def clear_relations(self, where: Mapping[str, Any] | None = None) -> list[RecallFileResource]: + from sqlmodel import delete, select + + filters = self._build_filters(self._sqla_models.RecallFileResource, where) + with self._sessions.session() as session: + rows = session.scalars(select(self._sqla_models.RecallFileResource).where(*filters)).all() + removed = [self._row_to_record(row) for row in rows] + if removed: + session.exec(delete(self._sqla_models.RecallFileResource).where(*filters)) + session.commit() + removed_ids = {rel.id for rel in removed} + self.relations[:] = [r for r in self.relations if r.id not in removed_ids] + return removed + + def get_resource_categories(self, resource_id: str) -> list[RecallFileResource]: + from sqlmodel import select + + with self._sessions.session() as session: + rows = session.scalars( + select(self._sqla_models.RecallFileResource).where( + self._sqla_models.RecallFileResource.resource_id == resource_id + ) + ).all() + return [self._cache_relation(row) for row in rows] + + def load_existing(self) -> None: + from sqlmodel import select + + with self._sessions.session() as session: + rows = session.scalars(select(self._sqla_models.RecallFileResource)).all() + for row in rows: + self._cache_relation(row) + + def _cache_relation(self, rel: RecallFileResource) -> RecallFileResource: + self.relations.append(rel) + return rel + + +__all__ = ["PostgresRecallFileResourceRepo"] diff --git a/src/memu/database/postgres/schema.py b/src/memu/database/postgres/schema.py index fb26bdc2..40cd92d4 100644 --- a/src/memu/database/postgres/schema.py +++ b/src/memu/database/postgres/schema.py @@ -27,6 +27,7 @@ RecallEntryModel, RecallFileEntryModel, RecallFileModel, + RecallFileResourceModel, ResourceModel, build_table_model, ) @@ -39,6 +40,7 @@ class SQLAModels: RecallFile: type[Any] RecallEntry: type[Any] RecallFileEntry: type[Any] + RecallFileResource: type[Any] _MODEL_CACHE: dict[type[Any], SQLAModels] = {} @@ -85,6 +87,12 @@ def get_sqlalchemy_models(*, scope_model: type[BaseModel] | None = None) -> SQLA tablename="category_items", metadata=metadata_obj, ) + recall_file_resource_model = build_table_model( + scope, + RecallFileResourceModel, + tablename="resource_categories", + metadata=metadata_obj, + ) class Base(SQLModel): __abstract__ = True @@ -96,6 +104,7 @@ class Base(SQLModel): RecallFile=recall_file_model, RecallEntry=recall_entry_model, RecallFileEntry=recall_file_entry_model, + RecallFileResource=recall_file_resource_model, ) _MODEL_CACHE[cache_key] = models return models diff --git a/src/memu/database/repositories/__init__.py b/src/memu/database/repositories/__init__.py index ac91449d..95df1fd5 100644 --- a/src/memu/database/repositories/__init__.py +++ b/src/memu/database/repositories/__init__.py @@ -1,6 +1,7 @@ from memu.database.repositories.recall_entry import RecallEntryRepo from memu.database.repositories.recall_file import RecallFileRepo from memu.database.repositories.recall_file_entry import RecallFileEntryRepo +from memu.database.repositories.recall_file_resource import RecallFileResourceRepo from memu.database.repositories.resource import ResourceRepo -__all__ = ["RecallEntryRepo", "RecallFileEntryRepo", "RecallFileRepo", "ResourceRepo"] +__all__ = ["RecallEntryRepo", "RecallFileEntryRepo", "RecallFileRepo", "RecallFileResourceRepo", "ResourceRepo"] diff --git a/src/memu/database/repositories/recall_file_resource.py b/src/memu/database/repositories/recall_file_resource.py new file mode 100644 index 00000000..ae667e6d --- /dev/null +++ b/src/memu/database/repositories/recall_file_resource.py @@ -0,0 +1,33 @@ +from __future__ import annotations + +from collections.abc import Mapping +from typing import Any, Protocol, runtime_checkable + +from memu.database.models import RecallFileResource + + +@runtime_checkable +class RecallFileResourceRepo(Protocol): + """Repository contract for resource/category relations.""" + + relations: list[RecallFileResource] + + def list_relations(self, where: Mapping[str, Any] | None = None) -> list[RecallFileResource]: ... + + def link_resource_category( + self, resource_id: str, cat_id: str, user_data: dict[str, Any] + ) -> RecallFileResource: ... + + def unlink_resource_category(self, resource_id: str, cat_id: str) -> None: ... + + def unlink_resource(self, resource_id: str) -> list[RecallFileResource]: + """Remove all relations for a given resource. Returns the removed relations.""" + ... + + def clear_relations(self, where: Mapping[str, Any] | None = None) -> list[RecallFileResource]: + """Remove all relations matching the scope. Returns the removed relations.""" + ... + + def get_resource_categories(self, resource_id: str) -> list[RecallFileResource]: ... + + def load_existing(self) -> None: ... diff --git a/src/memu/database/sqlite/models.py b/src/memu/database/sqlite/models.py index 3fd00bb4..7546ac66 100644 --- a/src/memu/database/sqlite/models.py +++ b/src/memu/database/sqlite/models.py @@ -11,7 +11,7 @@ from sqlalchemy import JSON, MetaData, String, Text from sqlmodel import Column, DateTime, Field, Index, SQLModel, func -from memu.database.models import EntryType, RecallEntry, RecallFile, RecallFileEntry, Resource +from memu.database.models import EntryType, RecallEntry, RecallFile, RecallFileEntry, RecallFileResource, Resource class TZDateTime(DateTime): @@ -88,6 +88,15 @@ class SQLiteRecallFileEntryModel(SQLiteBaseModelMixin, RecallFileEntry): __table_args__ = (Index("idx_sqlite_recall_file_entries_unique", "item_id", "category_id", unique=True),) +class SQLiteRecallFileResourceModel(SQLiteBaseModelMixin, RecallFileResource): + """SQLite category-resource relation model.""" + + resource_id: str = Field(sa_column=Column(String, nullable=False)) + category_id: str = Field(sa_column=Column(String, nullable=False)) + + __table_args__ = (Index("idx_sqlite_recall_file_resources_unique", "resource_id", "category_id", unique=True),) + + def _normalize_table_args(table_args: Any) -> tuple[list[Any], dict[str, Any]]: """Normalize SQLAlchemy table args to a consistent format.""" if table_args is None: @@ -176,6 +185,7 @@ def build_sqlite_table_model( "SQLiteRecallEntryModel", "SQLiteRecallFileEntryModel", "SQLiteRecallFileModel", + "SQLiteRecallFileResourceModel", "SQLiteResourceModel", "build_sqlite_table_model", ] diff --git a/src/memu/database/sqlite/repositories/__init__.py b/src/memu/database/sqlite/repositories/__init__.py index 53b6efd7..49a3171a 100644 --- a/src/memu/database/sqlite/repositories/__init__.py +++ b/src/memu/database/sqlite/repositories/__init__.py @@ -4,12 +4,14 @@ from memu.database.sqlite.repositories.recall_entry_repo import SQLiteRecallEntryRepo from memu.database.sqlite.repositories.recall_file_entry_repo import SQLiteRecallFileEntryRepo from memu.database.sqlite.repositories.recall_file_repo import SQLiteRecallFileRepo +from memu.database.sqlite.repositories.recall_file_resource_repo import SQLiteRecallFileResourceRepo from memu.database.sqlite.repositories.resource_repo import SQLiteResourceRepo __all__ = [ "SQLiteRecallEntryRepo", "SQLiteRecallFileEntryRepo", "SQLiteRecallFileRepo", + "SQLiteRecallFileResourceRepo", "SQLiteRepoBase", "SQLiteResourceRepo", ] diff --git a/src/memu/database/sqlite/repositories/recall_file_resource_repo.py b/src/memu/database/sqlite/repositories/recall_file_resource_repo.py new file mode 100644 index 00000000..d894fc0d --- /dev/null +++ b/src/memu/database/sqlite/repositories/recall_file_resource_repo.py @@ -0,0 +1,217 @@ +"""SQLite category-resource relation repository implementation.""" + +from __future__ import annotations + +import logging +from collections.abc import Mapping +from typing import Any + +from sqlmodel import select + +from memu.database.models import RecallFileResource +from memu.database.repositories.recall_file_resource import RecallFileResourceRepo +from memu.database.sqlite.repositories.base import SQLiteRepoBase +from memu.database.sqlite.schema import SQLiteSQLAModels +from memu.database.sqlite.session import SQLiteSessionManager +from memu.database.state import DatabaseState + +logger = logging.getLogger(__name__) + + +class SQLiteRecallFileResourceRepo(SQLiteRepoBase, RecallFileResourceRepo): + """SQLite implementation of category-resource relation repository.""" + + def __init__( + self, + *, + state: DatabaseState, + recall_file_resource_model: type[Any], + sqla_models: SQLiteSQLAModels, + sessions: SQLiteSessionManager, + scope_fields: list[str], + ) -> None: + """Initialize category-resource repository. + + Args: + state: Shared database state for caching. + recall_file_resource_model: SQLModel class for category-resource relations. + sqla_models: SQLAlchemy model container. + sessions: Session manager for database connections. + scope_fields: List of user scope field names. + """ + super().__init__( + state=state, + sqla_models=sqla_models, + sessions=sessions, + scope_fields=scope_fields, + ) + self._recall_file_resource_model = recall_file_resource_model + self.relations = self._state.resource_relations + + def list_relations(self, where: Mapping[str, Any] | None = None) -> list[RecallFileResource]: + """List category-resource relations matching the where clause. + + Args: + where: Optional filter conditions. + + Returns: + List of RecallFileResource relations. + """ + with self._sessions.session() as session: + stmt = select(self._recall_file_resource_model) + filters = self._build_filters(self._recall_file_resource_model, where) + if filters: + stmt = stmt.where(*filters) + rows = session.exec(stmt).all() + + result: list[RecallFileResource] = [] + for row in rows: + rel = RecallFileResource( + id=row.id, + resource_id=row.resource_id, + category_id=row.category_id, + created_at=row.created_at, + updated_at=row.updated_at, + **self._scope_kwargs_from(row), + ) + result.append(rel) + # Update cache + if not any(r.id == rel.id for r in self.relations): + self.relations.append(rel) + + return result + + def link_resource_category( + self, resource_id: str, category_id: str, user_data: dict[str, Any] + ) -> RecallFileResource: + """Create a link between a resource and a category. + + Args: + resource_id: Resource ID. + category_id: Category ID. + user_data: User scope data. + + Returns: + Created RecallFileResource relation. + """ + # Check if relation already exists + where: dict[str, Any] = { + "resource_id": resource_id, + "category_id": category_id, + **user_data, + } + with self._sessions.session() as session: + stmt = select(self._recall_file_resource_model) + filters = self._build_filters(self._recall_file_resource_model, where) + if filters: + stmt = stmt.where(*filters) + existing = session.exec(stmt).first() + + if existing: + rel = RecallFileResource( + id=existing.id, + resource_id=existing.resource_id, + category_id=existing.category_id, + created_at=existing.created_at, + updated_at=existing.updated_at, + **self._scope_kwargs_from(existing), + ) + return rel + + # Create new relation + now = self._now() + row = self._recall_file_resource_model( + resource_id=resource_id, + category_id=category_id, + created_at=now, + updated_at=now, + **user_data, + ) + session.add(row) + session.commit() + session.refresh(row) + + rel = RecallFileResource( + id=row.id, + resource_id=row.resource_id, + category_id=row.category_id, + created_at=row.created_at, + updated_at=row.updated_at, + **user_data, + ) + self.relations.append(rel) + return rel + + def unlink_resource_category(self, resource_id: str, category_id: str) -> None: + """Remove a link between a resource and a category. + + Args: + resource_id: Resource ID. + category_id: Category ID. + """ + with self._sessions.session() as session: + stmt = select(self._recall_file_resource_model).where( + self._recall_file_resource_model.resource_id == resource_id, + self._recall_file_resource_model.category_id == category_id, + ) + row = session.exec(stmt).first() + if row: + session.delete(row) + session.commit() + # Remove from cache + self.relations[:] = [ + r for r in self.relations if not (r.resource_id == resource_id and r.category_id == category_id) + ] + + def unlink_resource(self, resource_id: str) -> list[RecallFileResource]: + """Remove all relations for a given resource (used on resource deletion).""" + from sqlmodel import delete + + removed = self.list_relations({"resource_id": resource_id}) + if not removed: + return [] + with self._sessions.session() as session: + session.exec( + delete(self._recall_file_resource_model).where( + self._recall_file_resource_model.resource_id == resource_id + ) + ) + session.commit() + self.relations[:] = [r for r in self.relations if r.resource_id != resource_id] + return removed + + def clear_relations(self, where: Mapping[str, Any] | None = None) -> list[RecallFileResource]: + """Remove all relations matching the scope (used on clear_memory).""" + from sqlmodel import delete + + removed = self.list_relations(where) + if not removed: + return [] + filters = self._build_filters(self._recall_file_resource_model, where) + with self._sessions.session() as session: + del_stmt = delete(self._recall_file_resource_model) + if filters: + del_stmt = del_stmt.where(*filters) + session.exec(del_stmt) + session.commit() + removed_ids = {rel.id for rel in removed} + self.relations[:] = [r for r in self.relations if r.id not in removed_ids] + return removed + + def get_resource_categories(self, resource_id: str) -> list[RecallFileResource]: + """Get all category relations for a given resource. + + Args: + resource_id: Resource ID. + + Returns: + List of RecallFileResource relations for the resource. + """ + return self.list_relations({"resource_id": resource_id}) + + def load_existing(self) -> None: + """Load all existing relations from database into cache.""" + self.list_relations() + + +__all__ = ["SQLiteRecallFileResourceRepo"] diff --git a/src/memu/database/sqlite/schema.py b/src/memu/database/sqlite/schema.py index 90396bec..e16242ba 100644 --- a/src/memu/database/sqlite/schema.py +++ b/src/memu/database/sqlite/schema.py @@ -13,6 +13,7 @@ SQLiteRecallEntryModel, SQLiteRecallFileEntryModel, SQLiteRecallFileModel, + SQLiteRecallFileResourceModel, SQLiteResourceModel, build_sqlite_table_model, ) @@ -27,6 +28,7 @@ class SQLiteSQLAModels: RecallFile: type[Any] RecallEntry: type[Any] RecallFileEntry: type[Any] + RecallFileResource: type[Any] _MODEL_CACHE: dict[type[Any], SQLiteSQLAModels] = {} @@ -75,6 +77,12 @@ def get_sqlite_sqlalchemy_models(*, scope_model: type[BaseModel] | None = None) tablename="memu_category_items", metadata=metadata_obj, ) + recall_file_resource_model = build_sqlite_table_model( + scope, + SQLiteRecallFileResourceModel, + tablename="memu_resource_categories", + metadata=metadata_obj, + ) class SQLiteBase(SQLModel): __abstract__ = True @@ -86,6 +94,7 @@ class SQLiteBase(SQLModel): RecallFile=recall_file_model, RecallEntry=recall_entry_model, RecallFileEntry=recall_file_entry_model, + RecallFileResource=recall_file_resource_model, ) _MODEL_CACHE[cache_key] = models return models diff --git a/src/memu/database/sqlite/sqlite.py b/src/memu/database/sqlite/sqlite.py index 3857e86e..13ef7009 100644 --- a/src/memu/database/sqlite/sqlite.py +++ b/src/memu/database/sqlite/sqlite.py @@ -9,11 +9,18 @@ from sqlmodel import SQLModel from memu.database.interfaces import Database -from memu.database.models import RecallEntry, RecallFile, RecallFileEntry, Resource -from memu.database.repositories import RecallEntryRepo, RecallFileEntryRepo, RecallFileRepo, ResourceRepo +from memu.database.models import RecallEntry, RecallFile, RecallFileEntry, RecallFileResource, Resource +from memu.database.repositories import ( + RecallEntryRepo, + RecallFileEntryRepo, + RecallFileRepo, + RecallFileResourceRepo, + ResourceRepo, +) from memu.database.sqlite.repositories.recall_entry_repo import SQLiteRecallEntryRepo from memu.database.sqlite.repositories.recall_file_entry_repo import SQLiteRecallFileEntryRepo from memu.database.sqlite.repositories.recall_file_repo import SQLiteRecallFileRepo +from memu.database.sqlite.repositories.recall_file_resource_repo import SQLiteRecallFileResourceRepo from memu.database.sqlite.repositories.resource_repo import SQLiteResourceRepo from memu.database.sqlite.schema import SQLiteSQLAModels, get_sqlite_sqlalchemy_models from memu.database.sqlite.session import SQLiteSessionManager @@ -44,10 +51,12 @@ class SQLiteStore(Database): recall_file_repo: RecallFileRepo recall_entry_repo: RecallEntryRepo recall_file_entry_repo: RecallFileEntryRepo + recall_file_resource_repo: RecallFileResourceRepo resources: dict[str, Resource] items: dict[str, RecallEntry] categories: dict[str, RecallFile] relations: list[RecallFileEntry] + resource_relations: list[RecallFileResource] def __init__( self, @@ -58,6 +67,7 @@ def __init__( recall_file_model: type[Any] | None = None, recall_entry_model: type[Any] | None = None, recall_file_entry_model: type[Any] | None = None, + recall_file_resource_model: type[Any] | None = None, sqla_models: SQLiteSQLAModels | None = None, ) -> None: """Initialize SQLite database store. @@ -86,6 +96,7 @@ def __init__( recall_file_model = recall_file_model or self._sqla_models.RecallFile recall_entry_model = recall_entry_model or self._sqla_models.RecallEntry recall_file_entry_model = recall_file_entry_model or self._sqla_models.RecallFileEntry + recall_file_resource_model = recall_file_resource_model or self._sqla_models.RecallFileResource # Initialize repositories self.resource_repo = SQLiteResourceRepo( @@ -116,12 +127,20 @@ def __init__( sessions=self._sessions, scope_fields=self._scope_fields, ) + self.recall_file_resource_repo = SQLiteRecallFileResourceRepo( + state=self._state, + recall_file_resource_model=recall_file_resource_model, + sqla_models=self._sqla_models, + sessions=self._sessions, + scope_fields=self._scope_fields, + ) # Set up cache references self.resources = self._state.resources self.items = self._state.items self.categories = self._state.categories self.relations = self._state.relations + self.resource_relations = self._state.resource_relations def _create_tables(self) -> None: """Create SQLite tables if they don't exist.""" @@ -140,6 +159,7 @@ def load_existing(self) -> None: self.recall_file_repo.load_existing() self.recall_entry_repo.load_existing() self.recall_file_entry_repo.load_existing() + self.recall_file_resource_repo.load_existing() __all__ = ["SQLiteStore"] diff --git a/src/memu/database/state.py b/src/memu/database/state.py index db8104cb..3d913189 100644 --- a/src/memu/database/state.py +++ b/src/memu/database/state.py @@ -2,7 +2,7 @@ from dataclasses import dataclass, field -from memu.database.models import RecallEntry, RecallFile, RecallFileEntry, Resource +from memu.database.models import RecallEntry, RecallFile, RecallFileEntry, RecallFileResource, Resource @dataclass @@ -11,6 +11,7 @@ class DatabaseState: items: dict[str, RecallEntry] = field(default_factory=dict) categories: dict[str, RecallFile] = field(default_factory=dict) relations: list[RecallFileEntry] = field(default_factory=list) + resource_relations: list[RecallFileResource] = field(default_factory=list) __all__ = ["DatabaseState"] diff --git a/src/memu/prompts/memory_fs/__init__.py b/src/memu/prompts/memory_fs/__init__.py index 3a65163d..2170a11b 100644 --- a/src/memu/prompts/memory_fs/__init__.py +++ b/src/memu/prompts/memory_fs/__init__.py @@ -1,16 +1,27 @@ -"""Prompts for the optional memory_fs synthesis bypass. +"""Prompts for the resource -> file memorize path (ADR 0007 phase 1) and the +legacy memory_fs synthesis bypass. -Both prompts consume the shared trunk — the per-source multimodal descriptions — -plus the current state of the artifact they maintain, and emit the updated artifact. -There is a single prompt per artifact: a from-scratch build is just the same prompt -with an empty ``__EXISTING__`` block. The literal tokens ``__DESCRIPTIONS__`` and -``__EXISTING__`` are replaced (not ``str.format``) so text containing braces is safe. +Two families live here: + +- The single-shot synthesis prompts (``*_SYNTHESIS_PROMPT``) — the legacy bypass + that consumes the shared trunk plus the current artifact and emits it wholesale. +- The two-step resource -> file prompts (``ROUTE_PROMPTS`` / ``SYNTHESIS_PROMPTS``), + keyed by track (``"memory"`` / ``"skill"``): step (a) routes a source to the set of + files to update/create; step (b) writes each target file's body. Both are track + parametric so the workspace workflow drives chat and skill through one code path. + +The literal placeholder tokens (``__DESCRIPTIONS__``, ``__EXISTING__``, ``__CONTENT__``, +``__NAME__``, ``__DESCRIPTION__``) are replaced (not ``str.format``) so source text +containing braces is safe. """ from __future__ import annotations DESCRIPTIONS_PLACEHOLDER = "__DESCRIPTIONS__" EXISTING_PLACEHOLDER = "__EXISTING__" +CONTENT_PLACEHOLDER = "__CONTENT__" +NAME_PLACEHOLDER = "__NAME__" +DESCRIPTION_PLACEHOLDER = "__DESCRIPTION__" MEMORY_SYNTHESIS_PROMPT = """You are maintaining an AI agent's long-term memory document about a user. @@ -77,10 +88,126 @@ __DESCRIPTIONS__ """ +# --- Two-step resource -> file prompts (ADR 0007 phase 1) --------------------- +# +# Step (a): route a single source to the set of files to update/create. The model +# sees the existing files (name + one-line description) and the source content, and +# returns a JSON plan. Step (b): given one target file (name + description + current +# body) and the source content, write the file's full body. + +_MEMORY_ROUTE_PROMPT = """You are maintaining an AI agent's long-term memory about a user, organized as a set +of memory files (each a themed document — e.g. Profile, Preferences, Goals, Work). + +Below are the EXISTING memory files (name + one-line description), followed by the +CONTENT of a single source the agent just processed. Decide which files this source +should update, and whether any new file should be created for facts that fit no +existing file. Capture durable facts, preferences, goals, and notable events; ignore +throwaway chatter. + +Return ONLY a JSON array of operations. Each element is an object: + {"op": "update", "name": ""} + {"op": "create", "name": "", "description": "one-line summary of the file"} +- Use "update" with a file's EXACT existing name to route the source there. +- Use "create" only when no existing file fits; give a reusable name and a description. +- Prefer updating an existing file over creating a near-duplicate. +- List a file at most once. If the source has nothing memory-worthy, return []. + +EXISTING memory files: +__EXISTING__ + +SOURCE content: +__CONTENT__ +""" + +_SKILL_ROUTE_PROMPT = """You are maintaining an AI agent's skill library — a set of skill files, each a +concrete, repeatable how-to (what worked, how to repeat it, what to avoid). + +Below are the EXISTING skills (name + one-line description), followed by the CONTENT of +a single source the agent just processed. Identify concrete, repeatable skills or tool +usage patterns in the content and decide which skills to update or create. Ignore +one-off facts, preferences, or trivia — those belong in memory, not here. + +Return ONLY a JSON array of operations. Each element is an object: + {"op": "update", "name": ""} + {"op": "create", "name": "kebab-case-skill-name", "description": "one-line summary of the skill"} +- Use "update" with a skill's EXACT existing name to revise it. +- Use "create" only when no existing skill fits; give a new kebab-case name and a description. +- Prefer updating an existing skill over creating a near-duplicate. +- List a skill at most once. If the source has no genuine skills, return []. + +EXISTING skills: +__EXISTING__ + +SOURCE content: +__CONTENT__ +""" + +_MEMORY_FILE_SYNTHESIS_PROMPT = """You are maintaining a single memory file about a user. + +FILE name: __NAME__ +FILE description: __DESCRIPTION__ + +Below is the CURRENT content of this file (empty if it is being created), followed by the +CONTENT of a new source. Produce the updated file. + +Requirements: +- Merge in facts from the source that belong in THIS file, revise statements the source + makes outdated, and keep existing content that is still valid. If the CURRENT content + is empty, synthesize a fresh document from the source alone. +- Only include material relevant to this file's topic; leave unrelated facts out. +- Output the FULL Markdown document only. Do not wrap it in code fences. +- Be concise and factual. Do not invent details not supported by the source. Write in the + same language as the source. + +CURRENT content: +__EXISTING__ + +NEW source content: +__CONTENT__ +""" + +_SKILL_FILE_SYNTHESIS_PROMPT = """You are maintaining a single skill file in an AI agent's skill library. + +SKILL name: __NAME__ +SKILL description: __DESCRIPTION__ + +Below is the CURRENT body of this skill (empty if it is being created), followed by the +CONTENT of a new source. Produce the updated skill body. + +Requirements: +- Capture the concrete, repeatable procedure this skill describes: what it accomplishes, + the steps to repeat it, and pitfalls to avoid. Merge in what the source adds and revise + anything it supersedes. If the CURRENT body is empty, write it fresh from the source. +- Output the FULL Markdown body only. Do not wrap it in code fences. +- Be concise and actionable. Do not invent steps not supported by the source. Write in the + same language as the source. + +CURRENT body: +__EXISTING__ + +NEW source content: +__CONTENT__ +""" + +# Track-keyed dispatch tables used by the workspace memorize workflow. +ROUTE_PROMPTS: dict[str, str] = { + "memory": _MEMORY_ROUTE_PROMPT, + "skill": _SKILL_ROUTE_PROMPT, +} +SYNTHESIS_PROMPTS: dict[str, str] = { + "memory": _MEMORY_FILE_SYNTHESIS_PROMPT, + "skill": _SKILL_FILE_SYNTHESIS_PROMPT, +} + __all__ = [ + "CONTENT_PLACEHOLDER", "DESCRIPTIONS_PLACEHOLDER", + "DESCRIPTION_PLACEHOLDER", "EXISTING_PLACEHOLDER", "MEMORY_SYNTHESIS_PROMPT", + "NAME_PLACEHOLDER", + "ROUTE_PROMPTS", "SKILL_FILE_SYNTHESIS_PROMPT", "SKILL_OVERVIEW_SYNTHESIS_PROMPT", + "SYNTHESIS_PROMPTS", ] diff --git a/tests/test_skill_track.py b/tests/test_skill_track.py index a1377170..241a1da9 100644 --- a/tests/test_skill_track.py +++ b/tests/test_skill_track.py @@ -1,23 +1,39 @@ +"""Tests for the resource -> file workspace memorize path (ADR 0007 phase 1). + +Exercises ``MemoryService._memorize_ws_synthesize_files`` — the two-step route + +per-file synthesis that replaces the entry plane for the chat/skill tracks — including +the ``RecallFile`` upsert and the ``resource -> file`` provenance link. +""" + from __future__ import annotations from pathlib import Path +from typing import Any from memu.app import MemoryService -# A skill-synthesis response in the per-file format: name + description + body. -_SKILLS_JSON = ( - '[{"name": "pour-over", "description": "Brew pour-over coffee", "body": "# Pour-over\\nUse a 1:16 ratio."}]' -) +# Router output (step a): which files to update/create for a source. +_SKILL_ROUTE = '[{"op": "create", "name": "pour-over", "description": "Brew pour-over coffee"}]' +_MEMORY_ROUTE = '[{"op": "create", "name": "Preferences", "description": "User preferences"}]' +# Synthesis output (step b): the file body. +_SKILL_BODY = "# Pour-over\nUse a 1:16 ratio." + +class _FakeClient: + """Fake LLM/embed client that answers the route step and the synthesis step. -class _FakeSkillClient: - """Stand-in client exposing both chat (skill JSON) and embed (fixed vector).""" + The two steps are distinguished by a marker only the route prompt contains, so a + single client can serve both ``chat`` calls in the workflow. + """ - def __init__(self, payload: str = _SKILLS_JSON) -> None: - self._payload = payload + def __init__(self, route: str = _SKILL_ROUTE, body: str = _SKILL_BODY) -> None: + self._route = route + self._body = body async def chat(self, prompt: str, system_prompt: str | None = None) -> str: - return self._payload + if "JSON array of operations" in prompt: + return self._route + return self._body async def embed(self, texts: list[str]) -> list[list[float]]: return [[0.1, 0.2, 0.3] for _ in texts] @@ -31,68 +47,137 @@ def _service(tmp_path: Path) -> MemoryService: ) -async def _run_skill_step(service: MemoryService, client: _FakeSkillClient, state: dict) -> dict: +def _seed_resource(service: MemoryService, *, track: str, user: dict[str, Any]) -> Any: + return service.database.resource_repo.create_resource( + url=f"/w/{track}/x.md", + modality="document", + local_path=f"/w/{track}/x.md", + caption=None, + embedding=None, + user_data=dict(user), + track=track, + ) + + +async def _run_synthesize( + service: MemoryService, + client: _FakeClient, + *, + track: str, + user: dict[str, Any], + text: str = "I brewed pour-over at a 1:16 ratio.", + resource: Any | None = None, +) -> dict: service._get_step_llm_client = lambda *a, **k: client # type: ignore[method-assign] service._get_step_embedding_client = lambda *a, **k: client # type: ignore[method-assign] - return await service._memorize_generate_skills(state, None) + res = resource if resource is not None else _seed_resource(service, track=track, user=user) + state = { + "resources": [res], + "preprocessed_resources": [{"text": text, "caption": None}], + "resource_track": track, + "store": service.database, + "user": user, + } + return await service._memorize_ws_synthesize_files(state, None) -async def test_skill_step_persists_skill_track_recall_file(tmp_path: Path) -> None: +async def test_skill_track_synthesizes_file_and_links_resource(tmp_path: Path) -> None: service = _service(tmp_path) store = service.database - state = { - "preprocessed_resources": [{"text": "I brewed pour-over at a 1:16 ratio.", "caption": None}], - "store": store, - "user": {"user_id": "u1"}, - } + user = {"user_id": "u1"} + res = _seed_resource(service, track="skill", user=user) - result = await _run_skill_step(service, _FakeSkillClient(), state) + result = await _run_synthesize(service, _FakeClient(), track="skill", user=user, resource=res) - skills = list(result["skills"]) - assert len(skills) == 1 - skill = skills[0] + files = list(result["files"]) + assert len(files) == 1 + skill = files[0] assert skill.name == "pour-over" assert skill.track == "skill" assert skill.description == "Brew pour-over coffee" - assert skill.content == "# Pour-over\nUse a 1:16 ratio." + assert skill.content == _SKILL_BODY # Persisted as a skill-track RecallFile, isolated from the memory track. skill_files = store.recall_file_repo.list_categories(where={"user_id": "u1", "track": "skill"}) assert [f.name for f in skill_files.values()] == ["pour-over"] - memory_files = store.recall_file_repo.list_categories(where={"user_id": "u1", "track": "memory"}) - assert memory_files == {} + assert store.recall_file_repo.list_categories(where={"user_id": "u1", "track": "memory"}) == {} + + # A resource -> file provenance link was recorded. + links = store.recall_file_resource_repo.list_relations(where=user) + assert len(links) == 1 + assert links[0].resource_id == res.id + assert links[0].category_id == skill.id -async def test_skill_step_revises_existing_skill_by_name(tmp_path: Path) -> None: +async def test_chat_track_routes_to_memory_track_file(tmp_path: Path) -> None: service = _service(tmp_path) store = service.database - state = { - "preprocessed_resources": [{"text": "pour-over notes", "caption": None}], - "store": store, - "user": {}, - } + user = {"user_id": "u1"} + + result = await _run_synthesize( + service, + _FakeClient(route=_MEMORY_ROUTE, body="## Preferences\nLikes strong coffee."), + track="chat", + user=user, + text="I really like strong coffee.", + ) + + files = list(result["files"]) + assert len(files) == 1 + assert files[0].name == "Preferences" + assert files[0].track == "memory" + assert store.recall_file_repo.list_categories(where={"user_id": "u1", "track": "skill"}) == {} + + +async def test_update_op_revises_existing_file_by_name(tmp_path: Path) -> None: + service = _service(tmp_path) + store = service.database + user: dict[str, Any] = {} - await _run_skill_step(service, _FakeSkillClient(), state) - revised = '[{"name": "pour-over", "description": "Brew pour-over coffee", "body": "# Pour-over\\nUpdated."}]' - await _run_skill_step(service, _FakeSkillClient(revised), state) + await _run_synthesize(service, _FakeClient(), track="skill", user=user) + # A second source updates the same skill by exact name. + revised = _FakeClient(route='[{"op": "update", "name": "pour-over"}]', body="# Pour-over\nUpdated.") + await _run_synthesize(service, revised, track="skill", user=user) skill_files = store.recall_file_repo.list_categories(where={"track": "skill"}) - # Same name -> revised in place, not duplicated. - assert len(skill_files) == 1 + assert len(skill_files) == 1 # revised in place, not duplicated assert next(iter(skill_files.values())).content == "# Pour-over\nUpdated." -async def test_skill_step_noop_when_synthesize_disabled(tmp_path: Path) -> None: +async def test_workspace_track_is_resource_only_noop(tmp_path: Path) -> None: service = _service(tmp_path) - service.memory_files_config.synthesize = False store = service.database - state = { - "preprocessed_resources": [{"text": "pour-over notes", "caption": None}], - "store": store, - "user": {}, - } + user = {"user_id": "u1"} + + result = await _run_synthesize(service, _FakeClient(), track="workspace", user=user) - result = await _run_skill_step(service, _FakeSkillClient(), state) + assert result["files"] == [] + assert store.recall_file_repo.list_categories(where={"user_id": "u1"}) == {} + assert store.recall_file_resource_repo.list_relations(where=user) == [] + + +async def test_empty_source_is_noop(tmp_path: Path) -> None: + service = _service(tmp_path) + store = service.database + user = {"user_id": "u1"} + + result = await _run_synthesize(service, _FakeClient(), track="skill", user=user, text=" ") + + assert result["files"] == [] + assert store.recall_file_repo.list_categories(where={"user_id": "u1", "track": "skill"}) == {} + + +async def test_update_op_for_unknown_file_is_dropped(tmp_path: Path) -> None: + service = _service(tmp_path) + store = service.database + user = {"user_id": "u1"} + + result = await _run_synthesize( + service, + _FakeClient(route='[{"op": "update", "name": "does-not-exist"}]'), + track="skill", + user=user, + ) - assert "skills" not in result - assert store.recall_file_repo.list_categories(where={"track": "skill"}) == {} + assert result["files"] == [] + assert store.recall_file_repo.list_categories(where={"user_id": "u1", "track": "skill"}) == {} From 79fce95d3b622d03c220fe89a5feea3f4858ef3c Mon Sep 17 00:00:00 2001 From: wu Date: Thu, 2 Jul 2026 13:15:50 +0900 Subject: [PATCH 3/9] feat: add RecallFileSegment layer for file retrieval MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Introduce RecallFileSegment (ADR 0007 L2 item): a searchable slice of a RecallFile carrying text + embedding + recall_file_id. Each file has 1..n segments, embedded as the retrieval unit that rolls up to its file. No ordinal — slicing is track-specific and not necessarily sequential. Threaded through the full stack like RecallFileResource: core model + inmemory/postgres/sqlite models and schemas (file_segments table), RecallFileSegmentRepo protocol and three backend repos, DatabaseState, interfaces, and store wiring. Segments are (re)built after workspace file synthesis: - skill track: one "name/description" segment per file - memory track: one segment per content line, skipping blanks and headings On update, diff existing vs new segment texts and drop-and-add only the difference so unchanged lines keep their embeddings. Cleared alongside categories in the crud clear path. --- src/memu/app/crud.py | 2 + src/memu/app/memorize.py | 66 +++++++++ src/memu/database/__init__.py | 4 + src/memu/database/inmemory/__init__.py | 2 + src/memu/database/inmemory/models.py | 9 ++ src/memu/database/inmemory/repo.py | 17 ++- .../inmemory/repositories/__init__.py | 6 + .../repositories/recall_file_segment_repo.py | 60 ++++++++ src/memu/database/interfaces.py | 5 + src/memu/database/models.py | 26 +++- src/memu/database/postgres/models.py | 25 +++- src/memu/database/postgres/postgres.py | 24 +++- .../postgres/repositories/__init__.py | 2 + .../repositories/recall_file_segment_repo.py | 125 ++++++++++++++++ src/memu/database/postgres/schema.py | 9 ++ src/memu/database/repositories/__init__.py | 10 +- .../repositories/recall_file_segment.py | 35 +++++ src/memu/database/sqlite/models.py | 21 ++- .../database/sqlite/repositories/__init__.py | 2 + .../repositories/recall_file_segment_repo.py | 134 ++++++++++++++++++ src/memu/database/sqlite/schema.py | 9 ++ src/memu/database/sqlite/sqlite.py | 24 +++- src/memu/database/state.py | 10 +- tests/test_skill_track.py | 60 ++++++++ 24 files changed, 677 insertions(+), 10 deletions(-) create mode 100644 src/memu/database/inmemory/repositories/recall_file_segment_repo.py create mode 100644 src/memu/database/postgres/repositories/recall_file_segment_repo.py create mode 100644 src/memu/database/repositories/recall_file_segment.py create mode 100644 src/memu/database/sqlite/repositories/recall_file_segment_repo.py diff --git a/src/memu/app/crud.py b/src/memu/app/crud.py index a83302a4..4ca35714 100644 --- a/src/memu/app/crud.py +++ b/src/memu/app/crud.py @@ -269,6 +269,8 @@ def _crud_clear_memory_relations(self, state: WorkflowState, step_context: Any) def _crud_clear_recall_files(self, state: WorkflowState, step_context: Any) -> WorkflowState: where_filters = state.get("where") or {} store = state["store"] + # Segments hang off files (ADR 0007 L2); clear them alongside their categories. + store.recall_file_segment_repo.clear_segments(where_filters) deleted = store.recall_file_repo.clear_categories(where_filters) state["deleted_categories"] = deleted return state diff --git a/src/memu/app/memorize.py b/src/memu/app/memorize.py index 9b0027d3..fcf3b95f 100644 --- a/src/memu/app/memorize.py +++ b/src/memu/app/memorize.py @@ -692,9 +692,75 @@ async def _memorize_ws_synthesize_files(self, state: WorkflowState, step_context llm_client=llm_client, embed_client=embed_client, ) + await self._sync_file_segments( + files=touched, + file_track=file_track, + store=store, + user_scope=user_scope, + embed_client=embed_client, + ) state["files"] = touched return state + @staticmethod + def _segment_texts_for_file(file: RecallFile, file_track: str) -> list[str]: + """Compute the searchable segment texts for a synthesized file (ADR 0007 L2 items). + + The slicing rule is track-specific: + + - ``skill``: a single ``name: ...\\ndescription: ...`` segment for the whole skill. + - ``memory``: one segment per content line, skipping blank lines and markdown + headings (lines starting with one or more ``#``). + + Texts are stripped and de-duplicated while preserving order so a repeated line is + embedded only once. + """ + if file_track == "skill": + return [f"name: {file.name}\ndescription: {file.description}"] + + texts: list[str] = [] + for line in (file.content or "").split("\n"): + stripped = line.strip() + if not stripped or stripped.startswith("#"): + continue + texts.append(stripped) + return list(dict.fromkeys(texts)) + + async def _sync_file_segments( + self, + *, + files: list[RecallFile], + file_track: str, + store: Database, + user_scope: dict[str, Any], + embed_client: Any, + ) -> None: + """Reconcile each file's stored segments with its freshly computed segment texts. + + Diffs the new segment texts against the existing ones and does a drop-and-add on the + difference only: segments whose text disappeared are deleted, and only genuinely new + texts are embedded and inserted. Unchanged lines keep their existing embedding, so an + edit that touches a few lines does not re-embed the whole file. + """ + for file in files: + new_texts = self._segment_texts_for_file(file, file_track) + existing = store.recall_file_segment_repo.list_segments_for_file(file.id) + existing_texts = {seg.text for seg in existing} + new_set = set(new_texts) + + for seg in existing: + if seg.text not in new_set: + store.recall_file_segment_repo.delete_segment(seg.id) + + to_add = [text for text in new_texts if text not in existing_texts] + if not to_add: + continue + vecs = await embed_client.embed(to_add) + for text, vec in zip(to_add, vecs, strict=True): + store.recall_file_segment_repo.create_segment( + recall_file_id=file.id, text=text, embedding=vec, user_data=dict(user_scope) + ) + async def _route_source_to_files( self, *, diff --git a/src/memu/database/__init__.py b/src/memu/database/__init__.py index 33b1c168..84960b0c 100644 --- a/src/memu/database/__init__.py +++ b/src/memu/database/__init__.py @@ -7,6 +7,7 @@ RecallFileEntryRecord, RecallFileRecord, RecallFileResourceRecord, + RecallFileSegmentRecord, ResourceRecord, ) from memu.database.repositories import ( @@ -14,6 +15,7 @@ RecallFileEntryRepo, RecallFileRepo, RecallFileResourceRepo, + RecallFileSegmentRepo, ResourceRepo, ) @@ -27,6 +29,8 @@ "RecallFileRepo", "RecallFileResourceRecord", "RecallFileResourceRepo", + "RecallFileSegmentRecord", + "RecallFileSegmentRepo", "ResourceRecord", "ResourceRepo", "build_database", diff --git a/src/memu/database/inmemory/__init__.py b/src/memu/database/inmemory/__init__.py index 2d7f267b..0de91a81 100644 --- a/src/memu/database/inmemory/__init__.py +++ b/src/memu/database/inmemory/__init__.py @@ -18,6 +18,7 @@ def build_inmemory_database( recall_entry_model, recall_file_entry_model, recall_file_resource_model, + recall_file_segment_model, ) = build_inmemory_models(user_model) return InMemoryStore( scope_model=user_model, @@ -26,6 +27,7 @@ def build_inmemory_database( recall_file_model=recall_file_model, recall_file_entry_model=recall_file_entry_model, recall_file_resource_model=recall_file_resource_model, + recall_file_segment_model=recall_file_segment_model, ) diff --git a/src/memu/database/inmemory/models.py b/src/memu/database/inmemory/models.py index 1a61fe01..0ff2412e 100644 --- a/src/memu/database/inmemory/models.py +++ b/src/memu/database/inmemory/models.py @@ -7,6 +7,7 @@ RecallFile, RecallFileEntry, RecallFileResource, + RecallFileSegment, Resource, merge_scope_model, ) @@ -32,6 +33,10 @@ class InMemoryFileResource(RecallFileResource): """Concrete in-memory resource-category relation model.""" +class InMemoryFileSegment(RecallFileSegment): + """Concrete in-memory file-segment model.""" + + def build_inmemory_models( user_model: type[BaseModel], ) -> tuple[ @@ -40,6 +45,7 @@ def build_inmemory_models( type[InMemoryRecallEntry], type[InMemoryFileEntry], type[InMemoryFileResource], + type[InMemoryFileSegment], ]: """ Build scoped in-memory models that inherit from both the base interface and the user scope model. @@ -49,18 +55,21 @@ def build_inmemory_models( recall_entry_model = merge_scope_model(user_model, InMemoryRecallEntry, name_suffix="RecallEntry") recall_file_entry_model = merge_scope_model(user_model, InMemoryFileEntry, name_suffix="RecallFileEntry") recall_file_resource_model = merge_scope_model(user_model, InMemoryFileResource, name_suffix="RecallFileResource") + recall_file_segment_model = merge_scope_model(user_model, InMemoryFileSegment, name_suffix="RecallFileSegment") return ( resource_model, recall_file_model, recall_entry_model, recall_file_entry_model, recall_file_resource_model, + recall_file_segment_model, ) __all__ = [ "InMemoryFileEntry", "InMemoryFileResource", + "InMemoryFileSegment", "InMemoryRecallEntry", "InMemoryRecallFile", "InMemoryResource", diff --git a/src/memu/database/inmemory/repo.py b/src/memu/database/inmemory/repo.py index 44d1d155..0da8228d 100644 --- a/src/memu/database/inmemory/repo.py +++ b/src/memu/database/inmemory/repo.py @@ -8,13 +8,21 @@ from memu.database.inmemory.repositories import ( InMemoryFileEntryRepository, InMemoryFileResourceRepository, + InMemoryFileSegmentRepository, InMemoryRecallEntryRepository, InMemoryRecallFileRepository, InMemoryResourceRepository, ) from memu.database.inmemory.state import InMemoryState from memu.database.interfaces import Database -from memu.database.models import RecallEntry, RecallFile, RecallFileEntry, RecallFileResource, Resource +from memu.database.models import ( + RecallEntry, + RecallFile, + RecallFileEntry, + RecallFileResource, + RecallFileSegment, + Resource, +) from memu.database.repositories import RecallFileRepo, ResourceRepo @@ -28,6 +36,7 @@ def __init__( recall_file_model: type[Any] | None = None, recall_file_entry_model: type[Any] | None = None, recall_file_resource_model: type[Any] | None = None, + recall_file_segment_model: type[Any] | None = None, state: InMemoryState | None = None, ) -> None: self.scope_model = scope_model or BaseModel @@ -37,6 +46,7 @@ def __init__( default_recall_entry_model, default_recall_file_entry_model, default_recall_file_resource_model, + default_recall_file_segment_model, ) = build_inmemory_models(self.scope_model) self.state = state or InMemoryState() @@ -45,6 +55,7 @@ def __init__( self.categories: dict[str, RecallFile] = self.state.categories self.relations: list[RecallFileEntry] = self.state.relations self.resource_relations: list[RecallFileResource] = self.state.resource_relations + self.segments: list[RecallFileSegment] = self.state.segments resource_model = resource_model or default_resource_model or Resource recall_entry_model = recall_entry_model or default_recall_entry_model or RecallEntry @@ -53,6 +64,7 @@ def __init__( recall_file_resource_model = ( recall_file_resource_model or default_recall_file_resource_model or RecallFileResource ) + recall_file_segment_model = recall_file_segment_model or default_recall_file_segment_model or RecallFileSegment self.resource_repo: ResourceRepo = InMemoryResourceRepository(state=self.state, resource_model=resource_model) self.recall_file_repo: RecallFileRepo = InMemoryRecallFileRepository( @@ -65,6 +77,9 @@ def __init__( self.recall_file_resource_repo = InMemoryFileResourceRepository( state=self.state, recall_file_resource_model=recall_file_resource_model ) + self.recall_file_segment_repo = InMemoryFileSegmentRepository( + state=self.state, recall_file_segment_model=recall_file_segment_model + ) def close(self) -> None: return None diff --git a/src/memu/database/inmemory/repositories/__init__.py b/src/memu/database/inmemory/repositories/__init__.py index ff79d762..fa6481bb 100644 --- a/src/memu/database/inmemory/repositories/__init__.py +++ b/src/memu/database/inmemory/repositories/__init__.py @@ -11,11 +11,16 @@ InMemoryFileResourceRepository, RecallFileResourceRepo, ) +from memu.database.inmemory.repositories.recall_file_segment_repo import ( + InMemoryFileSegmentRepository, + RecallFileSegmentRepo, +) from memu.database.inmemory.repositories.resource_repo import InMemoryResourceRepository, ResourceRepo __all__ = [ "InMemoryFileEntryRepository", "InMemoryFileResourceRepository", + "InMemoryFileSegmentRepository", "InMemoryRecallEntryRepository", "InMemoryRecallFileRepository", "InMemoryResourceRepository", @@ -23,5 +28,6 @@ "RecallFileEntryRepo", "RecallFileRepo", "RecallFileResourceRepo", + "RecallFileSegmentRepo", "ResourceRepo", ] diff --git a/src/memu/database/inmemory/repositories/recall_file_segment_repo.py b/src/memu/database/inmemory/repositories/recall_file_segment_repo.py new file mode 100644 index 00000000..4b475266 --- /dev/null +++ b/src/memu/database/inmemory/repositories/recall_file_segment_repo.py @@ -0,0 +1,60 @@ +from __future__ import annotations + +import uuid +from collections.abc import Mapping +from typing import Any + +from memu.database.inmemory.repositories.filter import matches_where +from memu.database.inmemory.state import InMemoryState +from memu.database.models import RecallFileSegment +from memu.database.repositories.recall_file_segment import RecallFileSegmentRepo + + +class InMemoryFileSegmentRepository(RecallFileSegmentRepo): + def __init__(self, *, state: InMemoryState, recall_file_segment_model: type[RecallFileSegment]) -> None: + self._state = state + self.recall_file_segment_model = recall_file_segment_model + self.segments: list[RecallFileSegment] = self._state.segments + + def list_segments(self, where: Mapping[str, Any] | None = None) -> list[RecallFileSegment]: + if not where: + return list(self.segments) + return [seg for seg in self.segments if matches_where(seg, where)] + + def list_segments_for_file(self, recall_file_id: str) -> list[RecallFileSegment]: + return [seg for seg in self.segments if seg.recall_file_id == recall_file_id] + + def create_segment( + self, *, recall_file_id: str, text: str, embedding: list[float] | None, user_data: dict[str, Any] + ) -> RecallFileSegment: + seg = self.recall_file_segment_model( + id=str(uuid.uuid4()), recall_file_id=recall_file_id, text=text, embedding=embedding, **user_data + ) + self.segments.append(seg) + return seg + + def delete_segment(self, segment_id: str) -> None: + # Mutate the shared state list in place so the DatabaseState reference and this + # repo's view never diverge. + self.segments[:] = [seg for seg in self.segments if seg.id != segment_id] + + def delete_segments_for_file(self, recall_file_id: str) -> list[RecallFileSegment]: + removed = [seg for seg in self.segments if seg.recall_file_id == recall_file_id] + self.segments[:] = [seg for seg in self.segments if seg.recall_file_id != recall_file_id] + return removed + + def clear_segments(self, where: Mapping[str, Any] | None = None) -> list[RecallFileSegment]: + if not where: + removed = list(self.segments) + self.segments.clear() + return removed + removed = [seg for seg in self.segments if matches_where(seg, where)] + removed_ids = {seg.id for seg in removed} + self.segments[:] = [seg for seg in self.segments if seg.id not in removed_ids] + return removed + + def load_existing(self) -> None: + return None + + +__all__ = ["InMemoryFileSegmentRepository"] diff --git a/src/memu/database/interfaces.py b/src/memu/database/interfaces.py index 20b732e4..1ebb5314 100644 --- a/src/memu/database/interfaces.py +++ b/src/memu/database/interfaces.py @@ -6,12 +6,14 @@ from memu.database.models import RecallFile as RecallFileRecord from memu.database.models import RecallFileEntry as RecallFileEntryRecord from memu.database.models import RecallFileResource as RecallFileResourceRecord +from memu.database.models import RecallFileSegment as RecallFileSegmentRecord from memu.database.models import Resource as ResourceRecord from memu.database.repositories import ( RecallEntryRepo, RecallFileEntryRepo, RecallFileRepo, RecallFileResourceRepo, + RecallFileSegmentRepo, ResourceRepo, ) @@ -25,12 +27,14 @@ class Database(Protocol): recall_entry_repo: RecallEntryRepo recall_file_entry_repo: RecallFileEntryRepo recall_file_resource_repo: RecallFileResourceRepo + recall_file_segment_repo: RecallFileSegmentRepo resources: dict[str, ResourceRecord] items: dict[str, RecallEntryRecord] categories: dict[str, RecallFileRecord] relations: list[RecallFileEntryRecord] resource_relations: list[RecallFileResourceRecord] + segments: list[RecallFileSegmentRecord] def close(self) -> None: ... @@ -41,5 +45,6 @@ def close(self) -> None: ... "RecallFileEntryRecord", "RecallFileRecord", "RecallFileResourceRecord", + "RecallFileSegmentRecord", "ResourceRecord", ] diff --git a/src/memu/database/models.py b/src/memu/database/models.py index ed7d2596..4df0c235 100644 --- a/src/memu/database/models.py +++ b/src/memu/database/models.py @@ -116,6 +116,20 @@ class RecallFileResource(BaseRecord): category_id: str +class RecallFileSegment(BaseRecord): + """A searchable slice (L2 item) of a ``RecallFile`` (ADR 0007). + + Each file has 1..n segments; ``text`` is the embed/search unit and ``embedding`` + its vector. Retrieval ranks segments and rolls the top hits up to their file via + ``recall_file_id``. Segments carry no ordinal: how a file is sliced is track-specific + and not necessarily sequential, so position would not be informative. + """ + + recall_file_id: str + text: str + embedding: list[float] | None = None + + def merge_scope_model[TBaseRecord: BaseRecord]( user_model: type[BaseModel], core_model: type[TBaseRecord], *, name_suffix: str ) -> type[TBaseRecord]: @@ -134,7 +148,14 @@ def merge_scope_model[TBaseRecord: BaseRecord]( def build_scoped_models( user_model: type[BaseModel], -) -> tuple[type[Resource], type[RecallFile], type[RecallEntry], type[RecallFileEntry], type[RecallFileResource]]: +) -> tuple[ + type[Resource], + type[RecallFile], + type[RecallEntry], + type[RecallFileEntry], + type[RecallFileResource], + type[RecallFileSegment], +]: """ Build scoped interface models (Pydantic) that inherit from the base record models and user scope. """ @@ -143,12 +164,14 @@ def build_scoped_models( recall_entry_model = merge_scope_model(user_model, RecallEntry, name_suffix="RecallEntry") recall_file_entry_model = merge_scope_model(user_model, RecallFileEntry, name_suffix="RecallFileEntry") recall_file_resource_model = merge_scope_model(user_model, RecallFileResource, name_suffix="RecallFileResource") + recall_file_segment_model = merge_scope_model(user_model, RecallFileSegment, name_suffix="RecallFileSegment") return ( resource_model, recall_file_model, recall_entry_model, recall_file_entry_model, recall_file_resource_model, + recall_file_segment_model, ) @@ -159,6 +182,7 @@ def build_scoped_models( "RecallFile", "RecallFileEntry", "RecallFileResource", + "RecallFileSegment", "Resource", "ToolCallResult", "build_scoped_models", diff --git a/src/memu/database/postgres/models.py b/src/memu/database/postgres/models.py index e6628f34..3441ab8b 100644 --- a/src/memu/database/postgres/models.py +++ b/src/memu/database/postgres/models.py @@ -17,7 +17,15 @@ from sqlalchemy.dialects.postgresql import JSONB from sqlmodel import Column, DateTime, Field, Index, SQLModel, func -from memu.database.models import EntryType, RecallEntry, RecallFile, RecallFileEntry, RecallFileResource, Resource +from memu.database.models import ( + EntryType, + RecallEntry, + RecallFile, + RecallFileEntry, + RecallFileResource, + RecallFileSegment, + Resource, +) class TZDateTime(DateTime): @@ -83,6 +91,14 @@ class RecallFileResourceModel(BaseModelMixin, RecallFileResource): __table_args__ = (Index("idx_recall_file_resources_unique", "resource_id", "category_id", unique=True),) +class RecallFileSegmentModel(BaseModelMixin, RecallFileSegment): + recall_file_id: str = Field( + sa_column=Column(ForeignKey("memory_categories.id", ondelete="CASCADE"), nullable=False, index=True) + ) + text: str = Field(sa_column=Column(Text, nullable=False)) + embedding: list[float] | None = Field(default=None, sa_column=Column(Vector(), nullable=True)) + + def _normalize_table_args(table_args: Any) -> tuple[list[Any], dict[str, Any]]: if table_args is None: return [], {} @@ -165,9 +181,9 @@ def build_table_model( def build_scoped_models( user_model: type[BaseModel], -) -> tuple[type[SQLModel], type[SQLModel], type[SQLModel], type[SQLModel], type[SQLModel]]: +) -> tuple[type[SQLModel], type[SQLModel], type[SQLModel], type[SQLModel], type[SQLModel], type[SQLModel]]: """ - Build scoped SQLModel tables for each entity (resource, category, item, relation). + Build scoped SQLModel tables for each entity (resource, category, item, relation, segment). """ resource_model = build_table_model(user_model, ResourceModel, tablename="resources") recall_file_model = build_table_model( @@ -176,12 +192,14 @@ def build_scoped_models( recall_entry_model = build_table_model(user_model, RecallEntryModel, tablename="memory_items") recall_file_entry_model = build_table_model(user_model, RecallFileEntryModel, tablename="category_items") recall_file_resource_model = build_table_model(user_model, RecallFileResourceModel, tablename="resource_categories") + recall_file_segment_model = build_table_model(user_model, RecallFileSegmentModel, tablename="file_segments") return ( resource_model, recall_file_model, recall_entry_model, recall_file_entry_model, recall_file_resource_model, + recall_file_segment_model, ) @@ -191,6 +209,7 @@ def build_scoped_models( "RecallFileEntryModel", "RecallFileModel", "RecallFileResourceModel", + "RecallFileSegmentModel", "ResourceModel", "build_scoped_models", "build_table_model", diff --git a/src/memu/database/postgres/postgres.py b/src/memu/database/postgres/postgres.py index 20b56853..3f785c0b 100644 --- a/src/memu/database/postgres/postgres.py +++ b/src/memu/database/postgres/postgres.py @@ -6,12 +6,20 @@ from pydantic import BaseModel from memu.database.interfaces import Database -from memu.database.models import RecallEntry, RecallFile, RecallFileEntry, RecallFileResource, Resource +from memu.database.models import ( + RecallEntry, + RecallFile, + RecallFileEntry, + RecallFileResource, + RecallFileSegment, + Resource, +) from memu.database.postgres.migration import DDLMode, run_migrations from memu.database.postgres.repositories.recall_entry_repo import PostgresRecallEntryRepo from memu.database.postgres.repositories.recall_file_entry_repo import PostgresRecallFileEntryRepo from memu.database.postgres.repositories.recall_file_repo import PostgresRecallFileRepo from memu.database.postgres.repositories.recall_file_resource_repo import PostgresRecallFileResourceRepo +from memu.database.postgres.repositories.recall_file_segment_repo import PostgresRecallFileSegmentRepo from memu.database.postgres.repositories.resource_repo import PostgresResourceRepo from memu.database.postgres.schema import SQLAModels, get_sqlalchemy_models, require_sqlalchemy from memu.database.postgres.session import SessionManager @@ -20,6 +28,7 @@ RecallFileEntryRepo, RecallFileRepo, RecallFileResourceRepo, + RecallFileSegmentRepo, ResourceRepo, ) from memu.database.state import DatabaseState @@ -33,11 +42,13 @@ class PostgresStore(Database): recall_entry_repo: RecallEntryRepo recall_file_entry_repo: RecallFileEntryRepo recall_file_resource_repo: RecallFileResourceRepo + recall_file_segment_repo: RecallFileSegmentRepo resources: dict[str, Resource] items: dict[str, RecallEntry] categories: dict[str, RecallFile] relations: list[RecallFileEntry] resource_relations: list[RecallFileResource] + segments: list[RecallFileSegment] def __init__( self, @@ -52,6 +63,7 @@ def __init__( recall_entry_model: type[Any] | None = None, recall_file_entry_model: type[Any] | None = None, recall_file_resource_model: type[Any] | None = None, + recall_file_segment_model: type[Any] | None = None, sqla_models: SQLAModels | None = None, ) -> None: require_sqlalchemy() @@ -71,6 +83,7 @@ def __init__( recall_entry_model = recall_entry_model or self._sqla_models.RecallEntry recall_file_entry_model = recall_file_entry_model or self._sqla_models.RecallFileEntry recall_file_resource_model = recall_file_resource_model or self._sqla_models.RecallFileResource + recall_file_segment_model = recall_file_segment_model or self._sqla_models.RecallFileSegment self.resource_repo = PostgresResourceRepo( state=self._state, @@ -108,12 +121,20 @@ def __init__( sessions=self._sessions, scope_fields=self._scope_fields, ) + self.recall_file_segment_repo = PostgresRecallFileSegmentRepo( + state=self._state, + recall_file_segment_model=recall_file_segment_model, + sqla_models=self._sqla_models, + sessions=self._sessions, + scope_fields=self._scope_fields, + ) self.resources = self._state.resources self.items = self._state.items self.categories = self._state.categories self.relations = self._state.relations self.resource_relations = self._state.resource_relations + self.segments = self._state.segments # self._load_existing() @@ -126,3 +147,4 @@ def _load_existing(self) -> None: self.recall_entry_repo.load_existing() self.recall_file_entry_repo.load_existing() self.recall_file_resource_repo.load_existing() + self.recall_file_segment_repo.load_existing() diff --git a/src/memu/database/postgres/repositories/__init__.py b/src/memu/database/postgres/repositories/__init__.py index eea2e475..1d2aeb55 100644 --- a/src/memu/database/postgres/repositories/__init__.py +++ b/src/memu/database/postgres/repositories/__init__.py @@ -2,6 +2,7 @@ from memu.database.postgres.repositories.recall_file_entry_repo import PostgresRecallFileEntryRepo from memu.database.postgres.repositories.recall_file_repo import PostgresRecallFileRepo from memu.database.postgres.repositories.recall_file_resource_repo import PostgresRecallFileResourceRepo +from memu.database.postgres.repositories.recall_file_segment_repo import PostgresRecallFileSegmentRepo from memu.database.postgres.repositories.resource_repo import PostgresResourceRepo __all__ = [ @@ -9,5 +10,6 @@ "PostgresRecallFileEntryRepo", "PostgresRecallFileRepo", "PostgresRecallFileResourceRepo", + "PostgresRecallFileSegmentRepo", "PostgresResourceRepo", ] diff --git a/src/memu/database/postgres/repositories/recall_file_segment_repo.py b/src/memu/database/postgres/repositories/recall_file_segment_repo.py new file mode 100644 index 00000000..e4faa7c5 --- /dev/null +++ b/src/memu/database/postgres/repositories/recall_file_segment_repo.py @@ -0,0 +1,125 @@ +from __future__ import annotations + +from collections.abc import Mapping +from typing import Any + +from memu.database.models import RecallFileSegment +from memu.database.postgres.repositories.base import PostgresRepoBase +from memu.database.postgres.session import SessionManager +from memu.database.repositories.recall_file_segment import RecallFileSegmentRepo +from memu.database.state import DatabaseState + + +class PostgresRecallFileSegmentRepo(PostgresRepoBase, RecallFileSegmentRepo): + def __init__( + self, + *, + state: DatabaseState, + recall_file_segment_model: type[RecallFileSegment], + sqla_models: Any, + sessions: SessionManager, + scope_fields: list[str], + ) -> None: + super().__init__(state=state, sqla_models=sqla_models, sessions=sessions, scope_fields=scope_fields) + self._recall_file_segment_model = recall_file_segment_model + self.segments: list[RecallFileSegment] = self._state.segments + + def _row_to_record(self, row: Any) -> RecallFileSegment: + return RecallFileSegment( + id=row.id, + recall_file_id=row.recall_file_id, + text=row.text, + embedding=self._normalize_embedding(row.embedding), + created_at=row.created_at, + updated_at=row.updated_at, + **self._scope_kwargs_from(row), + ) + + def _cache_segment(self, row: Any) -> RecallFileSegment: + seg = self._row_to_record(row) + self.segments.append(seg) + return seg + + def list_segments(self, where: Mapping[str, Any] | None = None) -> list[RecallFileSegment]: + from sqlmodel import select + + filters = self._build_filters(self._sqla_models.RecallFileSegment, where) + with self._sessions.session() as session: + rows = session.scalars(select(self._sqla_models.RecallFileSegment).where(*filters)).all() + return [self._cache_segment(row) for row in rows] + + def list_segments_for_file(self, recall_file_id: str) -> list[RecallFileSegment]: + return self.list_segments({"recall_file_id": recall_file_id}) + + def create_segment( + self, *, recall_file_id: str, text: str, embedding: list[float] | None, user_data: dict[str, Any] + ) -> RecallFileSegment: + now = self._now() + row = self._recall_file_segment_model( + recall_file_id=recall_file_id, + text=text, + embedding=self._prepare_embedding(embedding), + created_at=now, + updated_at=now, + **user_data, + ) + with self._sessions.session() as session: + session.add(row) + session.commit() + session.refresh(row) + return self._cache_segment(row) + + def delete_segment(self, segment_id: str) -> None: + from sqlmodel import delete + + with self._sessions.session() as session: + session.exec( + delete(self._sqla_models.RecallFileSegment).where(self._sqla_models.RecallFileSegment.id == segment_id) + ) + session.commit() + self.segments[:] = [seg for seg in self.segments if seg.id != segment_id] + + def delete_segments_for_file(self, recall_file_id: str) -> list[RecallFileSegment]: + from sqlmodel import delete, select + + with self._sessions.session() as session: + rows = session.scalars( + select(self._sqla_models.RecallFileSegment).where( + self._sqla_models.RecallFileSegment.recall_file_id == recall_file_id + ) + ).all() + removed = [self._row_to_record(row) for row in rows] + if removed: + session.exec( + delete(self._sqla_models.RecallFileSegment).where( + self._sqla_models.RecallFileSegment.recall_file_id == recall_file_id + ) + ) + session.commit() + self.segments[:] = [seg for seg in self.segments if seg.recall_file_id != recall_file_id] + return removed + + def clear_segments(self, where: Mapping[str, Any] | None = None) -> list[RecallFileSegment]: + from sqlmodel import delete, select + + filters = self._build_filters(self._sqla_models.RecallFileSegment, where) + with self._sessions.session() as session: + rows = session.scalars(select(self._sqla_models.RecallFileSegment).where(*filters)).all() + removed = [self._row_to_record(row) for row in rows] + if removed: + session.exec(delete(self._sqla_models.RecallFileSegment).where(*filters)) + session.commit() + removed_ids = {seg.id for seg in removed} + self.segments[:] = [seg for seg in self.segments if seg.id not in removed_ids] + return removed + + def load_existing(self) -> None: + from sqlmodel import select + + with self._sessions.session() as session: + rows = session.scalars(select(self._sqla_models.RecallFileSegment)).all() + for row in rows: + self._cache_segment(row) + + +__all__ = ["PostgresRecallFileSegmentRepo"] diff --git a/src/memu/database/postgres/schema.py b/src/memu/database/postgres/schema.py index 40cd92d4..d970baa9 100644 --- a/src/memu/database/postgres/schema.py +++ b/src/memu/database/postgres/schema.py @@ -28,6 +28,7 @@ RecallFileEntryModel, RecallFileModel, RecallFileResourceModel, + RecallFileSegmentModel, ResourceModel, build_table_model, ) @@ -41,6 +42,7 @@ class SQLAModels: RecallEntry: type[Any] RecallFileEntry: type[Any] RecallFileResource: type[Any] + RecallFileSegment: type[Any] _MODEL_CACHE: dict[type[Any], SQLAModels] = {} @@ -93,6 +95,12 @@ def get_sqlalchemy_models(*, scope_model: type[BaseModel] | None = None) -> SQLA tablename="resource_categories", metadata=metadata_obj, ) + recall_file_segment_model = build_table_model( + scope, + RecallFileSegmentModel, + tablename="file_segments", + metadata=metadata_obj, + ) class Base(SQLModel): __abstract__ = True @@ -105,6 +113,7 @@ class Base(SQLModel): RecallEntry=recall_entry_model, RecallFileEntry=recall_file_entry_model, RecallFileResource=recall_file_resource_model, + RecallFileSegment=recall_file_segment_model, ) _MODEL_CACHE[cache_key] = models return models diff --git a/src/memu/database/repositories/__init__.py b/src/memu/database/repositories/__init__.py index 95df1fd5..e6016d61 100644 --- a/src/memu/database/repositories/__init__.py +++ b/src/memu/database/repositories/__init__.py @@ -2,6 +2,14 @@ from memu.database.repositories.recall_file import RecallFileRepo from memu.database.repositories.recall_file_entry import RecallFileEntryRepo from memu.database.repositories.recall_file_resource import RecallFileResourceRepo +from memu.database.repositories.recall_file_segment import RecallFileSegmentRepo from memu.database.repositories.resource import ResourceRepo -__all__ = ["RecallEntryRepo", "RecallFileEntryRepo", "RecallFileRepo", "RecallFileResourceRepo", "ResourceRepo"] +__all__ = [ + "RecallEntryRepo", + "RecallFileEntryRepo", + "RecallFileRepo", + "RecallFileResourceRepo", + "RecallFileSegmentRepo", + "ResourceRepo", +] diff --git a/src/memu/database/repositories/recall_file_segment.py b/src/memu/database/repositories/recall_file_segment.py new file mode 100644 index 00000000..2b5a4a0c --- /dev/null +++ b/src/memu/database/repositories/recall_file_segment.py @@ -0,0 +1,35 @@ +from __future__ import annotations + +from collections.abc import Mapping +from typing import Any, Protocol, runtime_checkable + +from memu.database.models import RecallFileSegment + + +@runtime_checkable +class RecallFileSegmentRepo(Protocol): + """Repository contract for file segments (searchable L2 slices of a ``RecallFile``).""" + + segments: list[RecallFileSegment] + + def list_segments(self, where: Mapping[str, Any] | None = None) -> list[RecallFileSegment]: ... + + def list_segments_for_file(self, recall_file_id: str) -> list[RecallFileSegment]: + """Return all segments belonging to a given file.""" + ... + + def create_segment( + self, *, recall_file_id: str, text: str, embedding: list[float] | None, user_data: dict[str, Any] + ) -> RecallFileSegment: ... + + def delete_segment(self, segment_id: str) -> None: ... + + def delete_segments_for_file(self, recall_file_id: str) -> list[RecallFileSegment]: + """Remove all segments for a given file. Returns the removed segments.""" + ... + + def clear_segments(self, where: Mapping[str, Any] | None = None) -> list[RecallFileSegment]: + """Remove all segments matching the scope. Returns the removed segments.""" + ... + + def load_existing(self) -> None: ... diff --git a/src/memu/database/sqlite/models.py b/src/memu/database/sqlite/models.py index 7546ac66..acd1e2f6 100644 --- a/src/memu/database/sqlite/models.py +++ b/src/memu/database/sqlite/models.py @@ -11,7 +11,15 @@ from sqlalchemy import JSON, MetaData, String, Text from sqlmodel import Column, DateTime, Field, Index, SQLModel, func -from memu.database.models import EntryType, RecallEntry, RecallFile, RecallFileEntry, RecallFileResource, Resource +from memu.database.models import ( + EntryType, + RecallEntry, + RecallFile, + RecallFileEntry, + RecallFileResource, + RecallFileSegment, + Resource, +) class TZDateTime(DateTime): @@ -97,6 +105,16 @@ class SQLiteRecallFileResourceModel(SQLiteBaseModelMixin, RecallFileResource): __table_args__ = (Index("idx_sqlite_recall_file_resources_unique", "resource_id", "category_id", unique=True),) +class SQLiteRecallFileSegmentModel(SQLiteBaseModelMixin, RecallFileSegment): + """SQLite file-segment model.""" + + recall_file_id: str = Field(sa_column=Column(String, nullable=False, index=True)) + text: str = Field(sa_column=Column(Text, nullable=False)) + # Override inherited embedding field: SQLite has no native vector type, so store the + # vector in a JSON column (a bare ``list`` annotation is not mappable by SQLModel). + embedding: list[float] | None = Field(default=None, sa_column=Column(JSON, nullable=True)) + + def _normalize_table_args(table_args: Any) -> tuple[list[Any], dict[str, Any]]: """Normalize SQLAlchemy table args to a consistent format.""" if table_args is None: @@ -186,6 +204,7 @@ def build_sqlite_table_model( "SQLiteRecallFileEntryModel", "SQLiteRecallFileModel", "SQLiteRecallFileResourceModel", + "SQLiteRecallFileSegmentModel", "SQLiteResourceModel", "build_sqlite_table_model", ] diff --git a/src/memu/database/sqlite/repositories/__init__.py b/src/memu/database/sqlite/repositories/__init__.py index 49a3171a..31f980da 100644 --- a/src/memu/database/sqlite/repositories/__init__.py +++ b/src/memu/database/sqlite/repositories/__init__.py @@ -5,6 +5,7 @@ from memu.database.sqlite.repositories.recall_file_entry_repo import SQLiteRecallFileEntryRepo from memu.database.sqlite.repositories.recall_file_repo import SQLiteRecallFileRepo from memu.database.sqlite.repositories.recall_file_resource_repo import SQLiteRecallFileResourceRepo +from memu.database.sqlite.repositories.recall_file_segment_repo import SQLiteRecallFileSegmentRepo from memu.database.sqlite.repositories.resource_repo import SQLiteResourceRepo __all__ = [ @@ -12,6 +13,7 @@ "SQLiteRecallFileEntryRepo", "SQLiteRecallFileRepo", "SQLiteRecallFileResourceRepo", + "SQLiteRecallFileSegmentRepo", "SQLiteRepoBase", "SQLiteResourceRepo", ] diff --git a/src/memu/database/sqlite/repositories/recall_file_segment_repo.py b/src/memu/database/sqlite/repositories/recall_file_segment_repo.py new file mode 100644 index 00000000..bc9a2110 --- /dev/null +++ b/src/memu/database/sqlite/repositories/recall_file_segment_repo.py @@ -0,0 +1,134 @@ +"""SQLite file-segment repository implementation.""" + +from __future__ import annotations + +import logging +from collections.abc import Mapping +from typing import Any + +from sqlmodel import delete, select + +from memu.database.models import RecallFileSegment +from memu.database.repositories.recall_file_segment import RecallFileSegmentRepo +from memu.database.sqlite.repositories.base import SQLiteRepoBase +from memu.database.sqlite.schema import SQLiteSQLAModels +from memu.database.sqlite.session import SQLiteSessionManager +from memu.database.state import DatabaseState + +logger = logging.getLogger(__name__) + + +class SQLiteRecallFileSegmentRepo(SQLiteRepoBase, RecallFileSegmentRepo): + """SQLite implementation of the file-segment repository.""" + + def __init__( + self, + *, + state: DatabaseState, + recall_file_segment_model: type[Any], + sqla_models: SQLiteSQLAModels, + sessions: SQLiteSessionManager, + scope_fields: list[str], + ) -> None: + super().__init__( + state=state, + sqla_models=sqla_models, + sessions=sessions, + scope_fields=scope_fields, + ) + self._recall_file_segment_model = recall_file_segment_model + self.segments = self._state.segments + + def _row_to_record(self, row: Any) -> RecallFileSegment: + return RecallFileSegment( + id=row.id, + recall_file_id=row.recall_file_id, + text=row.text, + embedding=self._normalize_embedding(row.embedding), + created_at=row.created_at, + updated_at=row.updated_at, + **self._scope_kwargs_from(row), + ) + + def list_segments(self, where: Mapping[str, Any] | None = None) -> list[RecallFileSegment]: + with self._sessions.session() as session: + stmt = select(self._recall_file_segment_model) + filters = self._build_filters(self._recall_file_segment_model, where) + if filters: + stmt = stmt.where(*filters) + rows = session.exec(stmt).all() + + result: list[RecallFileSegment] = [] + for row in rows: + seg = self._row_to_record(row) + result.append(seg) + if not any(s.id == seg.id for s in self.segments): + self.segments.append(seg) + return result + + def list_segments_for_file(self, recall_file_id: str) -> list[RecallFileSegment]: + return self.list_segments({"recall_file_id": recall_file_id}) + + def create_segment( + self, *, recall_file_id: str, text: str, embedding: list[float] | None, user_data: dict[str, Any] + ) -> RecallFileSegment: + now = self._now() + with self._sessions.session() as session: + row = self._recall_file_segment_model( + recall_file_id=recall_file_id, + text=text, + embedding=self._prepare_embedding(embedding), + created_at=now, + updated_at=now, + **user_data, + ) + session.add(row) + session.commit() + session.refresh(row) + seg = self._row_to_record(row) + + self.segments.append(seg) + return seg + + def delete_segment(self, segment_id: str) -> None: + with self._sessions.session() as session: + session.exec( + delete(self._recall_file_segment_model).where(self._recall_file_segment_model.id == segment_id) + ) + session.commit() + self.segments[:] = [seg for seg in self.segments if seg.id != segment_id] + + def delete_segments_for_file(self, recall_file_id: str) -> list[RecallFileSegment]: + removed = self.list_segments_for_file(recall_file_id) + if not removed: + return [] + with self._sessions.session() as session: + session.exec( + delete(self._recall_file_segment_model).where( + self._recall_file_segment_model.recall_file_id == recall_file_id + ) + ) + session.commit() + self.segments[:] = [seg for seg in self.segments if seg.recall_file_id != recall_file_id] + return removed + + def clear_segments(self, where: Mapping[str, Any] | None = None) -> list[RecallFileSegment]: + removed = self.list_segments(where) + if not removed: + return [] + filters = self._build_filters(self._recall_file_segment_model, where) + with self._sessions.session() as session: + del_stmt = delete(self._recall_file_segment_model) + if filters: + del_stmt = del_stmt.where(*filters) + session.exec(del_stmt) + session.commit() + removed_ids = {seg.id for seg in removed} + self.segments[:] = [seg for seg in self.segments if seg.id not in removed_ids] + return removed + + def load_existing(self) -> None: + self.list_segments() + + +__all__ = ["SQLiteRecallFileSegmentRepo"] diff --git a/src/memu/database/sqlite/schema.py b/src/memu/database/sqlite/schema.py index e16242ba..fc3b67f7 100644 --- a/src/memu/database/sqlite/schema.py +++ b/src/memu/database/sqlite/schema.py @@ -14,6 +14,7 @@ SQLiteRecallFileEntryModel, SQLiteRecallFileModel, SQLiteRecallFileResourceModel, + SQLiteRecallFileSegmentModel, SQLiteResourceModel, build_sqlite_table_model, ) @@ -29,6 +30,7 @@ class SQLiteSQLAModels: RecallEntry: type[Any] RecallFileEntry: type[Any] RecallFileResource: type[Any] + RecallFileSegment: type[Any] _MODEL_CACHE: dict[type[Any], SQLiteSQLAModels] = {} @@ -83,6 +85,12 @@ def get_sqlite_sqlalchemy_models(*, scope_model: type[BaseModel] | None = None) tablename="memu_resource_categories", metadata=metadata_obj, ) + recall_file_segment_model = build_sqlite_table_model( + scope, + SQLiteRecallFileSegmentModel, + tablename="memu_file_segments", + metadata=metadata_obj, + ) class SQLiteBase(SQLModel): __abstract__ = True @@ -95,6 +103,7 @@ class SQLiteBase(SQLModel): RecallEntry=recall_entry_model, RecallFileEntry=recall_file_entry_model, RecallFileResource=recall_file_resource_model, + RecallFileSegment=recall_file_segment_model, ) _MODEL_CACHE[cache_key] = models return models diff --git a/src/memu/database/sqlite/sqlite.py b/src/memu/database/sqlite/sqlite.py index 13ef7009..2c3ebbf6 100644 --- a/src/memu/database/sqlite/sqlite.py +++ b/src/memu/database/sqlite/sqlite.py @@ -9,18 +9,27 @@ from sqlmodel import SQLModel from memu.database.interfaces import Database -from memu.database.models import RecallEntry, RecallFile, RecallFileEntry, RecallFileResource, Resource +from memu.database.models import ( + RecallEntry, + RecallFile, + RecallFileEntry, + RecallFileResource, + RecallFileSegment, + Resource, +) from memu.database.repositories import ( RecallEntryRepo, RecallFileEntryRepo, RecallFileRepo, RecallFileResourceRepo, + RecallFileSegmentRepo, ResourceRepo, ) from memu.database.sqlite.repositories.recall_entry_repo import SQLiteRecallEntryRepo from memu.database.sqlite.repositories.recall_file_entry_repo import SQLiteRecallFileEntryRepo from memu.database.sqlite.repositories.recall_file_repo import SQLiteRecallFileRepo from memu.database.sqlite.repositories.recall_file_resource_repo import SQLiteRecallFileResourceRepo +from memu.database.sqlite.repositories.recall_file_segment_repo import SQLiteRecallFileSegmentRepo from memu.database.sqlite.repositories.resource_repo import SQLiteResourceRepo from memu.database.sqlite.schema import SQLiteSQLAModels, get_sqlite_sqlalchemy_models from memu.database.sqlite.session import SQLiteSessionManager @@ -52,11 +61,13 @@ class SQLiteStore(Database): recall_entry_repo: RecallEntryRepo recall_file_entry_repo: RecallFileEntryRepo recall_file_resource_repo: RecallFileResourceRepo + recall_file_segment_repo: RecallFileSegmentRepo resources: dict[str, Resource] items: dict[str, RecallEntry] categories: dict[str, RecallFile] relations: list[RecallFileEntry] resource_relations: list[RecallFileResource] + segments: list[RecallFileSegment] def __init__( self, @@ -68,6 +79,7 @@ def __init__( recall_entry_model: type[Any] | None = None, recall_file_entry_model: type[Any] | None = None, recall_file_resource_model: type[Any] | None = None, + recall_file_segment_model: type[Any] | None = None, sqla_models: SQLiteSQLAModels | None = None, ) -> None: """Initialize SQLite database store. @@ -97,6 +109,7 @@ def __init__( recall_entry_model = recall_entry_model or self._sqla_models.RecallEntry recall_file_entry_model = recall_file_entry_model or self._sqla_models.RecallFileEntry recall_file_resource_model = recall_file_resource_model or self._sqla_models.RecallFileResource + recall_file_segment_model = recall_file_segment_model or self._sqla_models.RecallFileSegment # Initialize repositories self.resource_repo = SQLiteResourceRepo( @@ -134,6 +147,13 @@ def __init__( sessions=self._sessions, scope_fields=self._scope_fields, ) + self.recall_file_segment_repo = SQLiteRecallFileSegmentRepo( + state=self._state, + recall_file_segment_model=recall_file_segment_model, + sqla_models=self._sqla_models, + sessions=self._sessions, + scope_fields=self._scope_fields, + ) # Set up cache references self.resources = self._state.resources @@ -141,6 +161,7 @@ def __init__( self.categories = self._state.categories self.relations = self._state.relations self.resource_relations = self._state.resource_relations + self.segments = self._state.segments def _create_tables(self) -> None: """Create SQLite tables if they don't exist.""" @@ -160,6 +181,7 @@ def load_existing(self) -> None: self.recall_entry_repo.load_existing() self.recall_file_entry_repo.load_existing() self.recall_file_resource_repo.load_existing() + self.recall_file_segment_repo.load_existing() __all__ = ["SQLiteStore"] diff --git a/src/memu/database/state.py b/src/memu/database/state.py index 3d913189..a17297a9 100644 --- a/src/memu/database/state.py +++ b/src/memu/database/state.py @@ -2,7 +2,14 @@ from dataclasses import dataclass, field -from memu.database.models import RecallEntry, RecallFile, RecallFileEntry, RecallFileResource, Resource +from memu.database.models import ( + RecallEntry, + RecallFile, + RecallFileEntry, + RecallFileResource, + RecallFileSegment, + Resource, +) @dataclass @@ -12,6 +19,7 @@ class DatabaseState: categories: dict[str, RecallFile] = field(default_factory=dict) relations: list[RecallFileEntry] = field(default_factory=list) resource_relations: list[RecallFileResource] = field(default_factory=list) + segments: list[RecallFileSegment] = field(default_factory=list) __all__ = ["DatabaseState"] diff --git a/tests/test_skill_track.py b/tests/test_skill_track.py index 241a1da9..23d63bb0 100644 --- a/tests/test_skill_track.py +++ b/tests/test_skill_track.py @@ -167,6 +167,66 @@ async def test_empty_source_is_noop(tmp_path: Path) -> None: assert store.recall_file_repo.list_categories(where={"user_id": "u1", "track": "skill"}) == {} +async def test_skill_track_creates_single_name_description_segment(tmp_path: Path) -> None: + service = _service(tmp_path) + store = service.database + user = {"user_id": "u1"} + + result = await _run_synthesize(service, _FakeClient(), track="skill", user=user) + skill = next(iter(result["files"])) + + segments = store.recall_file_segment_repo.list_segments_for_file(skill.id) + assert len(segments) == 1 + assert segments[0].text == "name: pour-over\ndescription: Brew pour-over coffee" + assert segments[0].embedding == [0.1, 0.2, 0.3] + + +async def test_memory_track_segments_are_lines_skipping_headings(tmp_path: Path) -> None: + service = _service(tmp_path) + store = service.database + user = {"user_id": "u1"} + + result = await _run_synthesize( + service, + _FakeClient(route=_MEMORY_ROUTE, body="## Preferences\nLikes strong coffee.\n\nDrinks it black."), + track="chat", + user=user, + ) + file = next(iter(result["files"])) + + segments = store.recall_file_segment_repo.list_segments_for_file(file.id) + assert [s.text for s in segments] == ["Likes strong coffee.", "Drinks it black."] + + +async def test_memory_segments_drop_and_add_on_update(tmp_path: Path) -> None: + service = _service(tmp_path) + store = service.database + user = {"user_id": "u1"} + + first = await _run_synthesize( + service, + _FakeClient(route=_MEMORY_ROUTE, body="## P\nline a\nline b"), + track="chat", + user=user, + ) + file = next(iter(first["files"])) + before = {s.text: s.id for s in store.recall_file_segment_repo.list_segments_for_file(file.id)} + assert set(before) == {"line a", "line b"} + + # An update changes only one line: "line b" -> "line c". + await _run_synthesize( + service, + _FakeClient(route='[{"op": "update", "name": "Preferences"}]', body="## P\nline a\nline c"), + track="chat", + user=user, + ) + after = {s.text: s.id for s in store.recall_file_segment_repo.list_segments_for_file(file.id)} + assert set(after) == {"line a", "line c"} + # Unchanged line keeps its original segment (not re-embedded); changed line is fresh. + assert after["line a"] == before["line a"] + assert "line b" not in after + + async def test_update_op_for_unknown_file_is_dropped(tmp_path: Path) -> None: service = _service(tmp_path) store = service.database From ac0edefdab888f0d6abf55cb5d224e69cbdd4f6d Mon Sep 17 00:00:00 2001 From: wu Date: Thu, 2 Jul 2026 13:55:43 +0900 Subject: [PATCH 4/9] feat: workspace retrieve over segments; denormalize track onto segment Rework retrieve_workspace to return segments/files/resources: - segments: RecallFileSegment ranked by embedding, file.top_k - files: roll-up of the files those segments point to (max segment score), not a ranked search - resources: workspace-track resources ranked by embedding, resource.top_k The entry layer is off (config retained but ignored). Add a denormalized `track` to RecallFileSegment, mirrored from the owning file at creation, so segment retrieval filters by track with a plain column predicate instead of a join. Track is immutable (segments are drop-and-recreated on re-slice), so it never drifts from the file. Threaded through create_segment across the interface and all three backends; _ws_recall_segments applies file.tracks as a track__in filter. Co-Authored-By: Claude Opus 4.8 (1M context) --- src/memu/app/memorize.py | 2 +- src/memu/app/retrieve.py | 124 +++++++++++------- src/memu/app/service.py | 2 +- .../repositories/recall_file_segment_repo.py | 15 ++- src/memu/database/models.py | 6 + src/memu/database/postgres/models.py | 1 + .../repositories/recall_file_segment_repo.py | 10 +- .../repositories/recall_file_segment.py | 8 +- src/memu/database/sqlite/models.py | 1 + .../repositories/recall_file_segment_repo.py | 10 +- tests/test_skill_track.py | 4 + 11 files changed, 127 insertions(+), 56 deletions(-) diff --git a/src/memu/app/memorize.py b/src/memu/app/memorize.py index fcf3b95f..07c6b23e 100644 --- a/src/memu/app/memorize.py +++ b/src/memu/app/memorize.py @@ -758,7 +758,7 @@ async def _sync_file_segments( vecs = await embed_client.embed(to_add) for text, vec in zip(to_add, vecs, strict=True): store.recall_file_segment_repo.create_segment( - recall_file_id=file.id, text=text, embedding=vec, user_data=dict(user_scope) + recall_file_id=file.id, track=file_track, text=text, embedding=vec, user_data=dict(user_scope) ) async def _route_source_to_files( diff --git a/src/memu/app/retrieve.py b/src/memu/app/retrieve.py index e69858ba..feec70a0 100644 --- a/src/memu/app/retrieve.py +++ b/src/memu/app/retrieve.py @@ -91,14 +91,23 @@ async def retrieve_workspace( query: str, where: dict[str, Any] | None = None, ) -> dict[str, Any]: - """Single-shot, LLM-free retrieval across the file/entry/resource layers. + """Single-shot, LLM-free retrieval over the segment/file/resource layers. Mirrors the relation between :meth:`memorize` and ``memorize_workspace``: a simpler entry point built on the same store and workflow machinery. The - query is embedded once and each enabled layer is ranked by vector - similarity — no intention routing, sufficiency checks, or summarization. - When ``file.tracks`` is set, the file layer is filtered on the ``track`` - column. Returns ``files``, ``entries``, and ``resources``. + query is embedded once and used to rank two layers by vector similarity — + no intention routing, sufficiency checks, or summarization: + + * ``segments``: :class:`RecallFileSegment` slices ranked by embedding, + ``file.top_k`` of them. + * ``files``: the :class:`RecallFile`\\ s pointed to by those segments — not + a ranked search, just a roll-up. Each file's score is the max score of + the segments that point to it. + * ``resources``: workspace-track resources ranked by embedding, + ``resource.top_k`` of them. + + The entry layer is disabled here (its config is retained but ignored). + Returns ``segments``, ``files``, and ``resources``. """ if not query or not query.strip(): raise ValueError("empty_query") @@ -111,7 +120,6 @@ async def retrieve_workspace( "store": store, "where": where_filters, "retrieve_file": config.file.enabled, - "retrieve_entry": config.entry.enabled, "retrieve_resource": config.resource.enabled, } @@ -1466,28 +1474,30 @@ def _format_llm_resource_content(self, hits: list[dict[str, Any]]) -> str: def _build_retrieve_workspace_workflow(self) -> list[WorkflowStep]: """The simple embedding-only workspace retrieval pipeline. - Three recall steps (file/entry/resource) feeding a terminal response step, - with none of the routing/sufficiency machinery of ``retrieve_rag``. The - query vector is embedded by the first recall step and reused downstream. + A segment recall step ranks :class:`RecallFileSegment` slices by embedding; + a file roll-up step gathers the files those segments point to; a resource + recall step ranks workspace-track resources by embedding. A terminal step + assembles the response. None of the routing/sufficiency machinery of + ``retrieve_rag`` applies. The query vector is embedded by the first recall + step and reused downstream. """ steps = [ WorkflowStep( - step_id="recall_files", - role="recall_files", - handler=self._ws_recall_files, + step_id="recall_segments", + role="recall_segments", + handler=self._ws_recall_segments, requires={"retrieve_file", "query", "store", "where"}, - produces={"file_hits", "file_pool", "query_vector"}, + produces={"segment_hits", "segment_pool", "query_vector"}, capabilities={"vector"}, config={"embed_llm_profile": "embedding"}, ), WorkflowStep( - step_id="recall_entries", - role="recall_entries", - handler=self._ws_recall_entries, - requires={"retrieve_entry", "query", "store", "where", "query_vector"}, - produces={"entry_hits", "entry_pool", "query_vector"}, - capabilities={"vector"}, - config={"embed_llm_profile": "embedding"}, + step_id="collect_files", + role="collect_files", + handler=self._ws_collect_files, + requires={"retrieve_file", "segment_hits", "segment_pool", "store", "where"}, + produces={"file_hits", "file_pool"}, + capabilities=set(), ), WorkflowStep( step_id="recall_resources", @@ -1503,10 +1513,10 @@ def _build_retrieve_workspace_workflow(self) -> list[WorkflowStep]: role="build_context", handler=self._ws_build_response, requires={ + "segment_hits", + "segment_pool", "file_hits", "file_pool", - "entry_hits", - "entry_pool", "resource_hits", "resource_pool", }, @@ -1518,7 +1528,7 @@ def _build_retrieve_workspace_workflow(self) -> list[WorkflowStep]: @staticmethod def _list_retrieve_workspace_initial_keys() -> set[str]: - return {"query", "store", "where", "retrieve_file", "retrieve_entry", "retrieve_resource"} + return {"query", "store", "where", "retrieve_file", "retrieve_resource"} async def _ws_query_vector(self, state: WorkflowState, step_context: Any) -> list[float]: """Embed the query once and cache it on the state for reuse across steps.""" @@ -1530,44 +1540,58 @@ async def _ws_query_vector(self, state: WorkflowState, step_context: Any) -> lis state["query_vector"] = qvec return cast(list[float], qvec) - async def _ws_recall_files(self, state: WorkflowState, step_context: Any) -> WorkflowState: + async def _ws_recall_segments(self, state: WorkflowState, step_context: Any) -> WorkflowState: if not state.get("retrieve_file"): - state["file_hits"] = [] - state["file_pool"] = {} + state["segment_hits"] = [] + state["segment_pool"] = {} state.setdefault("query_vector", None) return state store = state["store"] - # The file repo has no vector search, so rank the stored file embeddings - # directly. Optionally scope to the requested tracks via the where filter. - file_where = dict(state.get("where") or {}) + # The segment repo has no vector search, so rank the stored segment + # embeddings directly, mirroring how files used to be ranked. Optionally + # scope to the requested tracks via the denormalized segment ``track``. + segment_where = dict(state.get("where") or {}) tracks = self.retrieve_workspace_config.file.tracks if tracks: - file_where["track__in"] = list(tracks) - file_pool = store.recall_file_repo.list_categories(file_where) + segment_where["track__in"] = list(tracks) + segment_pool = {seg.id: seg for seg in store.recall_file_segment_repo.list_segments(segment_where)} qvec = await self._ws_query_vector(state, step_context) - state["file_hits"] = cosine_topk( + state["segment_hits"] = cosine_topk( qvec, - [(fid, f.embedding) for fid, f in file_pool.items()], + [(sid, seg.embedding) for sid, seg in segment_pool.items()], k=self.retrieve_workspace_config.file.top_k, ) - state["file_pool"] = file_pool + state["segment_pool"] = segment_pool return state - async def _ws_recall_entries(self, state: WorkflowState, step_context: Any) -> WorkflowState: - if not state.get("retrieve_entry"): - state["entry_hits"] = [] - state["entry_pool"] = {} - return state + async def _ws_collect_files(self, state: WorkflowState, _: Any) -> WorkflowState: + """Roll the ranked segments up to their files (no ranked file search). + Every file pointed to by a top segment is returned; a file's score is the + max score across the segments that point to it. + """ + segment_hits = state.get("segment_hits") or [] + segment_pool = state.get("segment_pool") or {} store = state["store"] where_filters = state.get("where") or {} - entry_pool = store.recall_entry_repo.list_items(where_filters) - qvec = await self._ws_query_vector(state, step_context) - state["entry_hits"] = store.recall_entry_repo.vector_search_items( - qvec, self.retrieve_workspace_config.entry.top_k, where=where_filters - ) - state["entry_pool"] = entry_pool + file_pool = store.recall_file_repo.list_categories(where_filters) + + file_scores: dict[str, float] = {} + for seg_id, score in segment_hits: + seg = segment_pool.get(seg_id) + if seg is None: + continue + fid = seg.recall_file_id + if fid not in file_pool: + continue + score = float(score) + if fid not in file_scores or score > file_scores[fid]: + file_scores[fid] = score + + # Preserve descending-score order so the response reads best-first. + state["file_hits"] = sorted(file_scores.items(), key=lambda kv: kv[1], reverse=True) + state["file_pool"] = file_pool return state async def _ws_recall_resources(self, state: WorkflowState, step_context: Any) -> WorkflowState: @@ -1577,19 +1601,21 @@ async def _ws_recall_resources(self, state: WorkflowState, step_context: Any) -> return state store = state["store"] - where_filters = state.get("where") or {} - resource_pool = store.resource_repo.list_resources(where_filters) + # Workspace retrieval only surfaces resources ingested by + # ``memorize_workspace`` (track="workspace"); other tracks are excluded. + resource_where = {**(state.get("where") or {}), "track": "workspace"} + resource_pool = store.resource_repo.list_resources(resource_where) qvec = await self._ws_query_vector(state, step_context) state["resource_hits"] = store.resource_repo.vector_search_resources( - qvec, self.retrieve_workspace_config.resource.top_k, where=where_filters + qvec, self.retrieve_workspace_config.resource.top_k, where=resource_where ) state["resource_pool"] = resource_pool return state def _ws_build_response(self, state: WorkflowState, _: Any) -> WorkflowState: state["response"] = { + "segments": self._materialize_hits(state.get("segment_hits", []), state.get("segment_pool", {})), "files": self._materialize_hits(state.get("file_hits", []), state.get("file_pool", {})), - "entries": self._materialize_hits(state.get("entry_hits", []), state.get("entry_pool", {})), "resources": self._materialize_hits(state.get("resource_hits", []), state.get("resource_pool", {})), } return state diff --git a/src/memu/app/service.py b/src/memu/app/service.py index 05fb0a76..abb9865b 100644 --- a/src/memu/app/service.py +++ b/src/memu/app/service.py @@ -343,7 +343,7 @@ def _register_pipelines(self) -> None: self._pipelines.register("retrieve_rag", rag_workflow, initial_state_keys=retrieve_initial_keys) llm_workflow = self._build_llm_retrieve_workflow() self._pipelines.register("retrieve_llm", llm_workflow, initial_state_keys=retrieve_initial_keys) - # Simple embedding-only workspace retrieval: file/entry/resource recall + response. + # Simple embedding-only workspace retrieval: segment recall + file roll-up + resource recall. workspace_retrieve_workflow = self._build_retrieve_workspace_workflow() self._pipelines.register( "retrieve_workspace", diff --git a/src/memu/database/inmemory/repositories/recall_file_segment_repo.py b/src/memu/database/inmemory/repositories/recall_file_segment_repo.py index 4b475266..2eb238fd 100644 --- a/src/memu/database/inmemory/repositories/recall_file_segment_repo.py +++ b/src/memu/database/inmemory/repositories/recall_file_segment_repo.py @@ -25,10 +25,21 @@ def list_segments_for_file(self, recall_file_id: str) -> list[RecallFileSegment] return [seg for seg in self.segments if seg.recall_file_id == recall_file_id] def create_segment( - self, *, recall_file_id: str, text: str, embedding: list[float] | None, user_data: dict[str, Any] + self, + *, + recall_file_id: str, + text: str, + embedding: list[float] | None, + user_data: dict[str, Any], + track: str = "memory", ) -> RecallFileSegment: seg = self.recall_file_segment_model( - id=str(uuid.uuid4()), recall_file_id=recall_file_id, text=text, embedding=embedding, **user_data + id=str(uuid.uuid4()), + recall_file_id=recall_file_id, + track=track, + text=text, + embedding=embedding, + **user_data, ) self.segments.append(seg) return seg diff --git a/src/memu/database/models.py b/src/memu/database/models.py index 4df0c235..51171d47 100644 --- a/src/memu/database/models.py +++ b/src/memu/database/models.py @@ -123,9 +123,15 @@ class RecallFileSegment(BaseRecord): its vector. Retrieval ranks segments and rolls the top hits up to their file via ``recall_file_id``. Segments carry no ordinal: how a file is sliced is track-specific and not necessarily sequential, so position would not be informative. + + ``track`` mirrors the owning file's track ("memory"/"skill"), denormalized here so + retrieval can filter segments by track with a plain column predicate instead of a + join. It is immutable for a segment's lifetime (segments are drop-and-recreated when + a file is re-sliced), so it never drifts from the file. """ recall_file_id: str + track: str = "memory" text: str embedding: list[float] | None = None diff --git a/src/memu/database/postgres/models.py b/src/memu/database/postgres/models.py index 3441ab8b..9b7a89cc 100644 --- a/src/memu/database/postgres/models.py +++ b/src/memu/database/postgres/models.py @@ -95,6 +95,7 @@ class RecallFileSegmentModel(BaseModelMixin, RecallFileSegment): recall_file_id: str = Field( sa_column=Column(ForeignKey("memory_categories.id", ondelete="CASCADE"), nullable=False, index=True) ) + track: str = Field(default="memory", sa_column=Column(String, nullable=False, server_default="memory")) text: str = Field(sa_column=Column(Text, nullable=False)) embedding: list[float] | None = Field(default=None, sa_column=Column(Vector(), nullable=True)) diff --git a/src/memu/database/postgres/repositories/recall_file_segment_repo.py b/src/memu/database/postgres/repositories/recall_file_segment_repo.py index e4faa7c5..3ef234c2 100644 --- a/src/memu/database/postgres/repositories/recall_file_segment_repo.py +++ b/src/memu/database/postgres/repositories/recall_file_segment_repo.py @@ -28,6 +28,7 @@ def _row_to_record(self, row: Any) -> RecallFileSegment: return RecallFileSegment( id=row.id, recall_file_id=row.recall_file_id, + track=row.track, text=row.text, embedding=self._normalize_embedding(row.embedding), created_at=row.created_at, @@ -52,11 +53,18 @@ def list_segments_for_file(self, recall_file_id: str) -> list[RecallFileSegment] return self.list_segments({"recall_file_id": recall_file_id}) def create_segment( - self, *, recall_file_id: str, text: str, embedding: list[float] | None, user_data: dict[str, Any] + self, + *, + recall_file_id: str, + text: str, + embedding: list[float] | None, + user_data: dict[str, Any], + track: str = "memory", ) -> RecallFileSegment: now = self._now() row = self._recall_file_segment_model( recall_file_id=recall_file_id, + track=track, text=text, embedding=self._prepare_embedding(embedding), created_at=now, diff --git a/src/memu/database/repositories/recall_file_segment.py b/src/memu/database/repositories/recall_file_segment.py index 2b5a4a0c..b11a13c2 100644 --- a/src/memu/database/repositories/recall_file_segment.py +++ b/src/memu/database/repositories/recall_file_segment.py @@ -19,7 +19,13 @@ def list_segments_for_file(self, recall_file_id: str) -> list[RecallFileSegment] ... def create_segment( - self, *, recall_file_id: str, text: str, embedding: list[float] | None, user_data: dict[str, Any] + self, + *, + recall_file_id: str, + text: str, + embedding: list[float] | None, + user_data: dict[str, Any], + track: str = "memory", ) -> RecallFileSegment: ... def delete_segment(self, segment_id: str) -> None: ... diff --git a/src/memu/database/sqlite/models.py b/src/memu/database/sqlite/models.py index acd1e2f6..38b326f5 100644 --- a/src/memu/database/sqlite/models.py +++ b/src/memu/database/sqlite/models.py @@ -109,6 +109,7 @@ class SQLiteRecallFileSegmentModel(SQLiteBaseModelMixin, RecallFileSegment): """SQLite file-segment model.""" recall_file_id: str = Field(sa_column=Column(String, nullable=False, index=True)) + track: str = Field(default="memory", sa_column=Column(String, nullable=False, server_default="memory")) text: str = Field(sa_column=Column(Text, nullable=False)) # Override inherited embedding field: SQLite has no native vector type, so store the # vector in a JSON column (a bare ``list`` annotation is not mappable by SQLModel). diff --git a/src/memu/database/sqlite/repositories/recall_file_segment_repo.py b/src/memu/database/sqlite/repositories/recall_file_segment_repo.py index bc9a2110..3ba57144 100644 --- a/src/memu/database/sqlite/repositories/recall_file_segment_repo.py +++ b/src/memu/database/sqlite/repositories/recall_file_segment_repo.py @@ -43,6 +43,7 @@ def _row_to_record(self, row: Any) -> RecallFileSegment: return RecallFileSegment( id=row.id, recall_file_id=row.recall_file_id, + track=row.track, text=row.text, embedding=self._normalize_embedding(row.embedding), created_at=row.created_at, @@ -70,12 +71,19 @@ def list_segments_for_file(self, recall_file_id: str) -> list[RecallFileSegment] return self.list_segments({"recall_file_id": recall_file_id}) def create_segment( - self, *, recall_file_id: str, text: str, embedding: list[float] | None, user_data: dict[str, Any] + self, + *, + recall_file_id: str, + text: str, + embedding: list[float] | None, + user_data: dict[str, Any], + track: str = "memory", ) -> RecallFileSegment: now = self._now() with self._sessions.session() as session: row = self._recall_file_segment_model( recall_file_id=recall_file_id, + track=track, text=text, embedding=self._prepare_embedding(embedding), created_at=now, diff --git a/tests/test_skill_track.py b/tests/test_skill_track.py index 23d63bb0..daea3fbb 100644 --- a/tests/test_skill_track.py +++ b/tests/test_skill_track.py @@ -179,6 +179,8 @@ async def test_skill_track_creates_single_name_description_segment(tmp_path: Pat assert len(segments) == 1 assert segments[0].text == "name: pour-over\ndescription: Brew pour-over coffee" assert segments[0].embedding == [0.1, 0.2, 0.3] + # Segment track mirrors the owning file's track (denormalized for filtering). + assert segments[0].track == "skill" async def test_memory_track_segments_are_lines_skipping_headings(tmp_path: Path) -> None: @@ -196,6 +198,8 @@ async def test_memory_track_segments_are_lines_skipping_headings(tmp_path: Path) segments = store.recall_file_segment_repo.list_segments_for_file(file.id) assert [s.text for s in segments] == ["Likes strong coffee.", "Drinks it black."] + # Segment track mirrors the owning file's track (chat routes to the "memory" track). + assert all(s.track == "memory" for s in segments) async def test_memory_segments_drop_and_add_on_update(tmp_path: Path) -> None: From 1c690f1922d19f5621ef31662a76f975adb804ea Mon Sep 17 00:00:00 2001 From: wu Date: Fri, 3 Jul 2026 07:11:35 +0900 Subject: [PATCH 5/9] refactor: decouple legacy/workspace mixin --- src/memu/app/memorize.py | 570 +------------------- src/memu/app/memorize_workspace.py | 834 +++++++++++++++++++++++++++++ src/memu/app/retrieve.py | 196 +------ src/memu/app/retrieve_workspace.py | 249 +++++++++ src/memu/app/service.py | 4 +- 5 files changed, 1088 insertions(+), 765 deletions(-) create mode 100644 src/memu/app/memorize_workspace.py create mode 100644 src/memu/app/retrieve_workspace.py diff --git a/src/memu/app/memorize.py b/src/memu/app/memorize.py index 07c6b23e..2f6331f5 100644 --- a/src/memu/app/memorize.py +++ b/src/memu/app/memorize.py @@ -7,14 +7,13 @@ import pathlib import re from collections.abc import Awaitable, Callable, Mapping, Sequence -from typing import TYPE_CHECKING, Any, ClassVar, cast +from typing import TYPE_CHECKING, Any, cast from xml.etree.ElementTree import Element import defusedxml.ElementTree as ET from pydantic import BaseModel from memu.app.settings import CategoryConfig, CustomPrompt -from memu.blob.folder import diff_folder, load_manifest, manifest_from_scan, save_manifest, scan_folder from memu.database.models import EntryType, RecallEntry, RecallFile, RecallFileEntry, Resource from memu.preprocess import PreprocessContext, preprocess_resource from memu.prompts.category_summary import ( @@ -23,14 +22,6 @@ from memu.prompts.category_summary import ( PROMPT as CATEGORY_SUMMARY_PROMPT, ) -from memu.prompts.memory_fs import ( - CONTENT_PLACEHOLDER, - DESCRIPTION_PLACEHOLDER, - EXISTING_PLACEHOLDER, - NAME_PLACEHOLDER, - ROUTE_PROMPTS, - SYNTHESIS_PROMPTS, -) from memu.prompts.memory_type import ( CUSTOM_PROMPTS as MEMORY_TYPE_CUSTOM_PROMPTS, ) @@ -121,207 +112,6 @@ async def memorize( raise RuntimeError(msg) return response - async def memorize_workspace( - self, - *, - folder: str, - user: dict[str, Any] | None = None, - ) -> dict[str, Any]: - """Sync a folder of source files into memory by diffing an input manifest. - - Scans ``folder`` recursively, infers each file's modality by extension - (unsupported extensions are skipped), and diffs against the sidecar - ``.memu_manifest.json`` to find added/modified/deleted files. Modified and - deleted files have their previously extracted memory cascade-deleted (with - affected category summaries recomputed); added and modified files are - (re)memorized by submitting each one through the single-file - :meth:`memorize` workflow. The manifest is then rewritten. - - ``memorize`` itself is left untouched: this is purely an additive, - directory-oriented entry point built on top of it. - """ - ctx = self._get_context() - store = self._get_database() - user_scope = self.user_model(**user).model_dump() if user is not None else None - await self._ensure_categories_ready(ctx, store, user_scope) - - root = pathlib.Path(folder).resolve() - scanned = scan_folder(root) - manifest = load_manifest(root) - diff = diff_folder(scanned, manifest) - - # 1. Cascade-delete memory for files that were modified or removed. - stale_urls = {sf.abs_path for sf in diff.modified} - stale_urls.update(str(root / rel) for rel in diff.deleted) - removed_resources = await self._cascade_delete_by_urls(stale_urls, ctx=ctx, store=store, user_scope=user_scope) - - # 2. (Re)memorize added and modified files; each file maps to one Resource. - changed_resources: list[Resource] = [] - entries: list[dict[str, Any]] = [] - files: list[dict[str, Any]] = [] - for scanned_file in [*diff.added, *diff.modified]: - result = await self._memorize_one( - resource_url=scanned_file.abs_path, - modality=scanned_file.modality, - user_scope=user_scope, - ctx=ctx, - store=store, - track=self._classify_track(scanned_file.rel_path), - ) - changed_resources.extend(cast("list[Resource]", result.get("resources") or [])) - # The inner single-file ``memorize`` keeps its legacy response keys - # (``items``/``categories``); translate them to the new vocabulary here. - response = cast("dict[str, Any]", result.get("response") or {}) - entries.extend(response.get("items", [])) - # Files reflect the cumulative scoped state, so the latest wins. - if response.get("categories"): - files = response["categories"] - - # 3. Refresh the memory file tree (full rebuild when anything was removed). - await self._update_memory_files(changed_resources, user_scope, force_full=diff.has_removals) - - # 4. Persist the updated input manifest. - save_manifest(root, manifest_from_scan(scanned)) - - return { - "folder": str(root), - "added": [sf.rel_path for sf in diff.added], - "modified": [sf.rel_path for sf in diff.modified], - "deleted": list(diff.deleted), - "resources": [self._model_dump_without_embeddings(r) for r in changed_resources], - "removed_resources": [self._model_dump_without_embeddings(r) for r in removed_resources], - "entries": entries, - "files": files, - } - - async def _memorize_one( - self, - *, - resource_url: str, - modality: str, - user_scope: dict[str, Any] | None, - ctx: Context, - store: Database, - track: str | None = None, - ) -> WorkflowState: - """Run the memorize workflow for a single file (one file -> one Resource). - - This mirrors :meth:`memorize` but returns the full workflow state (so the - workspace sync can collect the created resources) and takes an already - resolved ``user_scope``/``ctx``/``store`` to avoid re-resolving them per file. - """ - memory_types = self._resolve_memory_types() - state: WorkflowState = { - "resource_url": resource_url, - "modality": modality, - "memory_types": memory_types, - "categories_prompt_str": self._category_prompt_str, - "ctx": ctx, - "store": store, - "category_ids": list(ctx.category_ids), - "user": user_scope, - # Workspace sync path: let the extractor grow the taxonomy. - "allow_new_categories": True, - # Which workspace track this file belongs to (chat/skill/workspace). - "resource_track": track, - } - # The workspace path runs its own workflow (memorize + per-file skill - # generation); single-file ``memorize`` stays untouched (ADR 0006). - result = await self._run_workflow("memorize_workspace", state) - if result.get("response") is None: - msg = "Memorize workflow failed to produce a response" - raise RuntimeError(msg) - return result - - @staticmethod - def _classify_track(rel_path: str) -> str: - """Classify a workspace file into a track by its top-level folder. - - Files under ``chat/`` are the ``"chat"`` track, files under ``agent/`` are - the ``"skill"`` track, and everything else is the ``"workspace"`` track. - ``rel_path`` is the posix path relative to the scanned folder root. - """ - top = rel_path.split("/", 1)[0] - if top == "chat": - return "chat" - if top == "agent": - return "skill" - return "workspace" - - async def _cascade_delete_by_urls( - self, - urls: set[str], - *, - ctx: Context, - store: Database, - user_scope: dict[str, Any] | None, - ) -> list[Resource]: - """Delete resources (and their items/relations) whose url is in ``urls``. - - Affected category summaries are recomputed so the structured memory stays - consistent after a source file is changed or removed. - """ - if not urls: - return [] - where = user_scope or None - targets = [res for res in store.resource_repo.list_resources(where=where).values() if res.url in urls] - if not targets: - return [] - target_ids = {res.id for res in targets} - - # Discarded entry summaries per file, used to recompute summaries. Only the - # legacy entry-plane path (single-file ``memorize``) populates these. - file_discards: dict[str, list[str]] = {} - for entry in store.recall_entry_repo.list_items(where=where).values(): - if entry.resource_id not in target_ids: - continue - for relation in store.recall_file_entry_repo.get_item_categories(entry.id): - store.recall_file_entry_repo.unlink_item_category(entry.id, relation.category_id) - file_discards.setdefault(relation.category_id, []).append(entry.summary) - store.recall_entry_repo.delete_item(entry.id) - - for res in targets: - # Drop the resource -> file provenance links for the new synthesis path. - # NOTE (ADR 0007 phase 1 open issue): we do not rebuild the affected files - # from their remaining linked resources, so their content may go stale after - # a source change/delete. Tolerated for now. - store.recall_file_resource_repo.unlink_resource(res.id) - store.resource_repo.delete_resource(res.id) - - updates: dict[str, tuple[str | None, str | None]] = { - cid: ("\n".join(s for s in summaries if s and s.strip()), None) - for cid, summaries in file_discards.items() - if any(s and s.strip() for s in summaries) - } - if updates: - await self._patch_category_summaries(updates, ctx=ctx, store=store, llm_client=self._get_llm_client()) - return targets - - async def _update_memory_files( - self, - changed_resources: list[Resource], - user_scope: dict[str, Any] | None, - *, - force_full: bool = False, - ) -> None: - """Refresh the memory file tree after a workspace sync (init or incremental). - - Gated behind ``memory_files_config.enabled`` so a sync without the export - feature configured is a no-op. When any file was modified or deleted - (``force_full``), the tree is rebuilt from the full scoped store so stale - skills/entries do not linger; otherwise an incremental update merges the - just-created resources. Best-effort: the structured memory is already - persisted, so an export error must not fail the sync. - """ - if not getattr(self.memory_files_config, "enabled", False): - return - if not changed_resources and not force_full: - return - try: - await self._build_memory_files(user_scope, changed=None if force_full else changed_resources) - except Exception: - logger.exception("Memory file export failed after workspace memorize") - def _build_memorize_workflow(self) -> list[WorkflowStep]: steps = [ WorkflowStep( @@ -417,75 +207,6 @@ def _list_memorize_initial_keys() -> set[str]: "resource_track", } - def _build_memorize_workspace_workflow(self) -> list[WorkflowStep]: - """The workspace memorize pipeline: direct resource -> file synthesis (ADR 0007 phase 1). - - Unlike single-file :meth:`memorize` (``resource -> entry -> file``), the - workspace path synthesizes files straight from the preprocessed source and - creates no ``RecallEntry``. After ``preprocess`` it: - - - ``create_resource`` — one file maps to one :class:`Resource` (caption/embedding - for INDEX recall), for every track including ``workspace`` (resource-only). - - ``synthesize_files`` — for the ``chat`` and ``skill`` tracks only, route the - source to the files to update/create then synthesize each file's body, upserting - ``RecallFile`` and recording ``resource -> file`` provenance. ``workspace`` is a - no-op here. Retrieval over these files is deferred (ADR 0007 phase 2). - """ - synthesis_profile = getattr(self.memory_files_config, "synthesis_llm_profile", "default") - return [ - WorkflowStep( - step_id="ingest_resource", - role="ingest", - handler=self._memorize_ingest_resource, - requires={"resource_url", "modality"}, - produces={"local_path", "raw_text"}, - capabilities={"io"}, - ), - WorkflowStep( - step_id="preprocess_multimodal", - role="preprocess", - handler=self._memorize_preprocess_multimodal, - requires={"local_path", "modality", "raw_text"}, - produces={"preprocessed_resources"}, - capabilities={"llm"}, - config={"chat_llm_profile": self.memorize_config.preprocess_llm_profile}, - ), - WorkflowStep( - step_id="create_resource", - role="persist", - handler=self._memorize_ws_create_resource, - requires={ - "preprocessed_resources", - "modality", - "local_path", - "resource_url", - "store", - "user", - "resource_track", - }, - produces={"resources"}, - capabilities={"db", "vector"}, - config={"embed_llm_profile": "embedding"}, - ), - WorkflowStep( - step_id="synthesize_files", - role="synthesize_files", - handler=self._memorize_ws_synthesize_files, - requires={"resources", "preprocessed_resources", "resource_track", "store", "user"}, - produces={"files"}, - capabilities={"llm", "db", "vector"}, - config={"chat_llm_profile": synthesis_profile, "embed_llm_profile": "embedding"}, - ), - WorkflowStep( - step_id="build_response", - role="emit", - handler=self._memorize_ws_build_response, - requires={"resources", "files"}, - produces={"response"}, - capabilities=set(), - ), - ] - async def _memorize_ingest_resource(self, state: WorkflowState, step_context: Any) -> WorkflowState: local_path, raw_text = await self.fs.fetch(state["resource_url"], state["modality"]) state.update({"local_path": local_path, "raw_text": raw_text}) @@ -615,295 +336,6 @@ async def _memorize_persist_and_index(self, state: WorkflowState, step_context: ) return state - @staticmethod - def _format_skill_source_content(preprocessed_resources: list[dict[str, Any]]) -> str: - """Flatten a source's preprocessed segments into a single text block.""" - parts = [ - " ".join((prep.get("text") or "").split()) - for prep in preprocessed_resources - if (prep.get("text") or "").strip() - ] - return "\n\n".join(parts) - - # --- Workspace resource -> file path (ADR 0007 phase 1) ------------------- - - # Maps a workspace ``resource_track`` to the ``RecallFile.track`` it synthesizes - # into. ``workspace`` has no entry (resource-only), so it is absent. - _TRACK_TO_FILE_TRACK: ClassVar[dict[str, str]] = {"chat": "memory", "skill": "skill"} - - async def _memorize_ws_create_resource(self, state: WorkflowState, step_context: Any) -> WorkflowState: - """Create the single ``Resource`` for this file (one file -> one resource). - - Runs for every track; the ``workspace`` track stops here (resource-only). The - caption is the joined per-segment captions, embedded for INDEX/resource recall. - """ - embed_client = self._get_step_embedding_client(step_context) - store = state["store"] - preprocessed = state.get("preprocessed_resources") or [] - captions = [(prep.get("caption") or "").strip() for prep in preprocessed] - caption = "\n\n".join(c for c in captions if c) or None - res = await self._create_resource_with_caption( - resource_url=state["resource_url"], - modality=state["modality"], - local_path=state["local_path"], - caption=caption, - store=store, - embed_client=embed_client, - user=state.get("user", {}), - track=state.get("resource_track"), - ) - state["resources"] = [res] - return state - - async def _memorize_ws_synthesize_files(self, state: WorkflowState, step_context: Any) -> WorkflowState: - """Synthesize this source into ``RecallFile``s for the chat/skill tracks. - - Two steps: (a) route the source to the set of files to update/create given the - existing files' names+descriptions, and (b) synthesize each target file's body in - parallel. Persists each file and a ``resource -> file`` provenance link. The - ``workspace`` track (and any source with no content) is a no-op. - """ - track = state.get("resource_track") - file_track = self._TRACK_TO_FILE_TRACK.get(track or "") - resources = state.get("resources") or [] - content = self._format_skill_source_content(state.get("preprocessed_resources") or []) - if file_track is None or not resources or not content: - state["files"] = [] - return state - - store = state["store"] - user_scope = dict(state.get("user") or {}) - llm_client = self._get_step_llm_client(step_context) - embed_client = self._get_step_embedding_client(step_context) - resource = resources[0] - - existing = store.recall_file_repo.list_categories(where={**user_scope, "track": file_track}) - ops = await self._route_source_to_files( - file_track=file_track, content=content, existing=existing, llm_client=llm_client - ) - touched = await self._synthesize_file_ops( - ops=ops, - file_track=file_track, - content=content, - existing=existing, - resource=resource, - store=store, - user_scope=user_scope, - llm_client=llm_client, - embed_client=embed_client, - ) - await self._sync_file_segments( - files=touched, - file_track=file_track, - store=store, - user_scope=user_scope, - embed_client=embed_client, - ) - state["files"] = touched - return state - - @staticmethod - def _segment_texts_for_file(file: RecallFile, file_track: str) -> list[str]: - """Compute the searchable segment texts for a synthesized file (ADR 0007 L2 items). - - The slicing rule is track-specific: - - - ``skill``: a single ``name: ...\\ndescription: ...`` segment for the whole skill. - - ``memory``: one segment per content line, skipping blank lines and markdown - headings (lines starting with one or more ``#``). - - Texts are stripped and de-duplicated while preserving order so a repeated line is - embedded only once. - """ - if file_track == "skill": - return [f"name: {file.name}\ndescription: {file.description}"] - - texts: list[str] = [] - for line in (file.content or "").split("\n"): - stripped = line.strip() - if not stripped or stripped.startswith("#"): - continue - texts.append(stripped) - return list(dict.fromkeys(texts)) - - async def _sync_file_segments( - self, - *, - files: list[RecallFile], - file_track: str, - store: Database, - user_scope: dict[str, Any], - embed_client: Any, - ) -> None: - """Reconcile each file's stored segments with its freshly computed segment texts. - - Diffs the new segment texts against the existing ones and does a drop-and-add on the - difference only: segments whose text disappeared are deleted, and only genuinely new - texts are embedded and inserted. Unchanged lines keep their existing embedding, so an - edit that touches a few lines does not re-embed the whole file. - """ - for file in files: - new_texts = self._segment_texts_for_file(file, file_track) - existing = store.recall_file_segment_repo.list_segments_for_file(file.id) - existing_texts = {seg.text for seg in existing} - new_set = set(new_texts) - - for seg in existing: - if seg.text not in new_set: - store.recall_file_segment_repo.delete_segment(seg.id) - - to_add = [text for text in new_texts if text not in existing_texts] - if not to_add: - continue - vecs = await embed_client.embed(to_add) - for text, vec in zip(to_add, vecs, strict=True): - store.recall_file_segment_repo.create_segment( - recall_file_id=file.id, track=file_track, text=text, embedding=vec, user_data=dict(user_scope) - ) - - async def _route_source_to_files( - self, - *, - file_track: str, - content: str, - existing: Mapping[str, RecallFile], - llm_client: Any, - ) -> list[dict[str, str]]: - """Ask the model which existing files to update / what new files to create.""" - existing_text = self._format_existing_files(existing) or "(none)" - prompt = ( - ROUTE_PROMPTS[file_track] - .replace(EXISTING_PLACEHOLDER, existing_text) - .replace(CONTENT_PLACEHOLDER, self._escape_prompt_value(content)) - ) - return self._parse_file_ops(await llm_client.chat(prompt), existing) - - async def _synthesize_file_ops( - self, - *, - ops: list[dict[str, str]], - file_track: str, - content: str, - existing: Mapping[str, RecallFile], - resource: Resource, - store: Database, - user_scope: dict[str, Any], - llm_client: Any, - embed_client: Any, - ) -> list[RecallFile]: - """Synthesize each routed file's body (in parallel) and persist file + link.""" - existing_by_name = {f.name: f for f in existing.values()} - # Resolve ops to unique targets (dedup by name; last op's description wins). - targets: list[dict[str, Any]] = [] - by_name: dict[str, dict[str, Any]] = {} - for op in ops: - name = op["name"] - ex = existing_by_name.get(name) - description = (op.get("description") or (ex.description if ex else "") or "").strip() - target = by_name.get(name) - if target is None: - target = {"name": name, "description": description, "existing": ex} - by_name[name] = target - targets.append(target) - elif description: - target["description"] = description - if not targets: - return [] - - prompts = [ - SYNTHESIS_PROMPTS[file_track] - .replace(NAME_PLACEHOLDER, self._escape_prompt_value(t["name"])) - .replace(DESCRIPTION_PLACEHOLDER, self._escape_prompt_value(t["description"])) - .replace( - EXISTING_PLACEHOLDER, self._escape_prompt_value((t["existing"].content if t["existing"] else "") or "") - ) - .replace(CONTENT_PLACEHOLDER, self._escape_prompt_value(content)) - for t in targets - ] - bodies = await asyncio.gather(*[llm_client.chat(prompt) for prompt in prompts]) - - # Embed name+description for the files being created. - creates = [t for t in targets if t["existing"] is None] - create_vecs: dict[str, list[float]] = {} - if creates: - emb_texts = [f"{t['name']}: {t['description']}" if t["description"] else t["name"] for t in creates] - vecs = await embed_client.embed(emb_texts) - for t, vec in zip(creates, vecs, strict=True): - create_vecs[t["name"]] = vec - - touched: list[RecallFile] = [] - for target, body in zip(targets, bodies, strict=True): - cleaned = body.replace("```markdown", "").replace("```", "").strip() - file = target["existing"] - if file is None: - file = store.recall_file_repo.get_or_create_category( - name=target["name"], - description=target["description"], - embedding=create_vecs[target["name"]], - user_data=user_scope, - track=file_track, - ) - file = store.recall_file_repo.update_category(category_id=file.id, content=cleaned) - store.recall_file_resource_repo.link_resource_category(resource.id, file.id, user_data=dict(user_scope)) - touched.append(file) - return touched - - @staticmethod - def _format_existing_files(existing: Mapping[str, RecallFile]) -> str: - """Render existing files as ``- name: description`` lines for the router prompt.""" - return "\n".join( - f"- {f.name}: {f.description}" if f.description else f"- {f.name}" - for f in sorted(existing.values(), key=lambda f: f.name) - ) - - def _parse_file_ops(self, raw: str, existing: Mapping[str, RecallFile]) -> list[dict[str, str]]: - """Parse the router's JSON array into validated ``{op, name, description}`` dicts. - - ``update`` ops naming an unknown file are dropped (we never update a file that - does not exist); ``create``/``update`` are otherwise kept with a stripped name. - """ - if not raw: - return [] - start = raw.find("[") - end = raw.rfind("]") - if start == -1 or end == -1 or end <= start: - return [] - try: - parsed = json.loads(raw[start : end + 1]) - except (json.JSONDecodeError, TypeError): - return [] - if not isinstance(parsed, list): - return [] - existing_names = {f.name for f in existing.values()} - ops: list[dict[str, str]] = [] - for entry in parsed: - if not isinstance(entry, dict): - continue - op = entry.get("op") - name = entry.get("name") - if op not in {"update", "create"} or not isinstance(name, str) or not name.strip(): - continue - name = name.strip() - if op == "update" and name not in existing_names: - continue - description = entry.get("description") - description = description.strip() if isinstance(description, str) else "" - ops.append({"op": op, "name": name, "description": description}) - return ops - - def _memorize_ws_build_response(self, state: WorkflowState, step_context: Any) -> WorkflowState: - """Emit the workspace response (no entries; ``categories`` carries touched files).""" - resources = [self._model_dump_without_embeddings(r) for r in state.get("resources", [])] - files = [self._model_dump_without_embeddings(f) for f in state.get("files", [])] - # Keep the legacy response contract (``items``/``categories``); items is always - # empty on this path since the entry plane is gone. - base: dict[str, Any] = {"items": [], "categories": files, "relations": []} - if len(resources) == 1: - state["response"] = {"resource": resources[0], **base} - else: - state["response"] = {"resources": resources, **base} - return state - def _memorize_build_response(self, state: WorkflowState, step_context: Any) -> WorkflowState: ctx = state["ctx"] store = state["store"] diff --git a/src/memu/app/memorize_workspace.py b/src/memu/app/memorize_workspace.py new file mode 100644 index 00000000..2df24c00 --- /dev/null +++ b/src/memu/app/memorize_workspace.py @@ -0,0 +1,834 @@ +from __future__ import annotations + +import asyncio +import contextlib +import json +import logging +import pathlib +from collections.abc import Awaitable, Callable, Mapping +from typing import TYPE_CHECKING, Any, ClassVar, cast + +from pydantic import BaseModel + +from memu.app.settings import CategoryConfig, CustomPrompt +from memu.blob.folder import diff_folder, load_manifest, manifest_from_scan, save_manifest, scan_folder +from memu.database.models import EntryType, RecallFile, Resource +from memu.preprocess import PreprocessContext, preprocess_resource +from memu.prompts.memory_fs import ( + CONTENT_PLACEHOLDER, + DESCRIPTION_PLACEHOLDER, + EXISTING_PLACEHOLDER, + NAME_PLACEHOLDER, + ROUTE_PROMPTS, + SYNTHESIS_PROMPTS, +) +from memu.prompts.memory_type import DEFAULT_MEMORY_TYPES +from memu.workflow.step import WorkflowState, WorkflowStep + +logger = logging.getLogger(__name__) + +if TYPE_CHECKING: + from memu.app.service import Context + from memu.app.settings import MemorizeConfig, MemoryFilesConfig + from memu.blob.local_fs import LocalFS + from memu.database.interfaces import Database + + +class MemorizeWorkspaceMixin: + if TYPE_CHECKING: + memorize_config: MemorizeConfig + category_configs: list[CategoryConfig] + _category_prompt_str: str + fs: LocalFS + _run_workflow: Callable[..., Awaitable[WorkflowState]] + _get_context: Callable[[], Context] + _get_database: Callable[[], Database] + _get_step_llm_client: Callable[[Mapping[str, Any] | None], Any] + _get_step_embedding_client: Callable[[Mapping[str, Any] | None], Any] + _get_embedding_client: Callable[..., Any] + _get_llm_client: Callable[..., Any] + _get_vlm_client: Callable[..., Any] + _model_dump_without_embeddings: Callable[[BaseModel], dict[str, Any]] + _extract_json_blob: Callable[[str], str] + _escape_prompt_value: Callable[[str], str] + user_model: type[BaseModel] + + # Memory file system export (provided by MemoryService). + memory_files_config: MemoryFilesConfig + _build_memory_files: Callable[..., Awaitable[dict[str, Any]]] + + # Provided by CRUDMixin (composed onto MemoryService). + async def _patch_category_summaries( + self, + updates: dict[str, tuple[str | None, str | None]], + ctx: Context, + store: Database, + llm_client: Any | None = None, + ) -> None: ... + + async def memorize_workspace( + self, + *, + folder: str, + user: dict[str, Any] | None = None, + ) -> dict[str, Any]: + """Sync a folder of source files into memory by diffing an input manifest. + + Scans ``folder`` recursively, infers each file's modality by extension + (unsupported extensions are skipped), and diffs against the sidecar + ``.memu_manifest.json`` to find added/modified/deleted files. Modified and + deleted files have their previously extracted memory cascade-deleted (with + affected category summaries recomputed); added and modified files are + (re)memorized by submitting each one through the single-file + :meth:`memorize` workflow. The manifest is then rewritten. + + ``memorize`` itself is left untouched: this is purely an additive, + directory-oriented entry point built on top of it. + """ + ctx = self._get_context() + store = self._get_database() + user_scope = self.user_model(**user).model_dump() if user is not None else None + await self._ensure_categories_ready(ctx, store, user_scope) + + root = pathlib.Path(folder).resolve() + scanned = scan_folder(root) + manifest = load_manifest(root) + diff = diff_folder(scanned, manifest) + + # 1. Cascade-delete memory for files that were modified or removed. + stale_urls = {sf.abs_path for sf in diff.modified} + stale_urls.update(str(root / rel) for rel in diff.deleted) + removed_resources = await self._cascade_delete_by_urls(stale_urls, ctx=ctx, store=store, user_scope=user_scope) + + # 2. (Re)memorize added and modified files; each file maps to one Resource. + changed_resources: list[Resource] = [] + entries: list[dict[str, Any]] = [] + files: list[dict[str, Any]] = [] + for scanned_file in [*diff.added, *diff.modified]: + result = await self._memorize_one( + resource_url=scanned_file.abs_path, + modality=scanned_file.modality, + user_scope=user_scope, + ctx=ctx, + store=store, + track=self._classify_track(scanned_file.rel_path), + ) + changed_resources.extend(cast("list[Resource]", result.get("resources") or [])) + # The inner single-file ``memorize`` keeps its legacy response keys + # (``items``/``categories``); translate them to the new vocabulary here. + response = cast("dict[str, Any]", result.get("response") or {}) + entries.extend(response.get("items", [])) + # Files reflect the cumulative scoped state, so the latest wins. + if response.get("categories"): + files = response["categories"] + + # 3. Refresh the memory file tree (full rebuild when anything was removed). + await self._update_memory_files(changed_resources, user_scope, force_full=diff.has_removals) + + # 4. Persist the updated input manifest. + save_manifest(root, manifest_from_scan(scanned)) + + return { + "folder": str(root), + "added": [sf.rel_path for sf in diff.added], + "modified": [sf.rel_path for sf in diff.modified], + "deleted": list(diff.deleted), + "resources": [self._model_dump_without_embeddings(r) for r in changed_resources], + "removed_resources": [self._model_dump_without_embeddings(r) for r in removed_resources], + "entries": entries, + "files": files, + } + + async def _memorize_one( + self, + *, + resource_url: str, + modality: str, + user_scope: dict[str, Any] | None, + ctx: Context, + store: Database, + track: str | None = None, + ) -> WorkflowState: + """Run the memorize workflow for a single file (one file -> one Resource). + + This mirrors :meth:`memorize` but returns the full workflow state (so the + workspace sync can collect the created resources) and takes an already + resolved ``user_scope``/``ctx``/``store`` to avoid re-resolving them per file. + """ + memory_types = self._resolve_memory_types() + state: WorkflowState = { + "resource_url": resource_url, + "modality": modality, + "memory_types": memory_types, + "categories_prompt_str": self._category_prompt_str, + "ctx": ctx, + "store": store, + "category_ids": list(ctx.category_ids), + "user": user_scope, + # Workspace sync path: let the extractor grow the taxonomy. + "allow_new_categories": True, + # Which workspace track this file belongs to (chat/skill/workspace). + "resource_track": track, + } + # The workspace path runs its own workflow (memorize + per-file skill + # generation); single-file ``memorize`` stays untouched (ADR 0006). + result = await self._run_workflow("memorize_workspace", state) + if result.get("response") is None: + msg = "Memorize workflow failed to produce a response" + raise RuntimeError(msg) + return result + + @staticmethod + def _classify_track(rel_path: str) -> str: + """Classify a workspace file into a track by its top-level folder. + + Files under ``chat/`` are the ``"chat"`` track, files under ``agent/`` are + the ``"skill"`` track, and everything else is the ``"workspace"`` track. + ``rel_path`` is the posix path relative to the scanned folder root. + """ + top = rel_path.split("/", 1)[0] + if top == "chat": + return "chat" + if top == "agent": + return "skill" + return "workspace" + + async def _cascade_delete_by_urls( + self, + urls: set[str], + *, + ctx: Context, + store: Database, + user_scope: dict[str, Any] | None, + ) -> list[Resource]: + """Delete resources (and their items/relations) whose url is in ``urls``. + + Affected category summaries are recomputed so the structured memory stays + consistent after a source file is changed or removed. + """ + if not urls: + return [] + where = user_scope or None + targets = [res for res in store.resource_repo.list_resources(where=where).values() if res.url in urls] + if not targets: + return [] + target_ids = {res.id for res in targets} + + # Discarded entry summaries per file, used to recompute summaries. Only the + # legacy entry-plane path (single-file ``memorize``) populates these. + file_discards: dict[str, list[str]] = {} + for entry in store.recall_entry_repo.list_items(where=where).values(): + if entry.resource_id not in target_ids: + continue + for relation in store.recall_file_entry_repo.get_item_categories(entry.id): + store.recall_file_entry_repo.unlink_item_category(entry.id, relation.category_id) + file_discards.setdefault(relation.category_id, []).append(entry.summary) + store.recall_entry_repo.delete_item(entry.id) + + for res in targets: + # Drop the resource -> file provenance links for the new synthesis path. + # NOTE (ADR 0007 phase 1 open issue): we do not rebuild the affected files + # from their remaining linked resources, so their content may go stale after + # a source change/delete. Tolerated for now. + store.recall_file_resource_repo.unlink_resource(res.id) + store.resource_repo.delete_resource(res.id) + + updates: dict[str, tuple[str | None, str | None]] = { + cid: ("\n".join(s for s in summaries if s and s.strip()), None) + for cid, summaries in file_discards.items() + if any(s and s.strip() for s in summaries) + } + if updates: + await self._patch_category_summaries(updates, ctx=ctx, store=store, llm_client=self._get_llm_client()) + return targets + + async def _update_memory_files( + self, + changed_resources: list[Resource], + user_scope: dict[str, Any] | None, + *, + force_full: bool = False, + ) -> None: + """Refresh the memory file tree after a workspace sync (init or incremental). + + Gated behind ``memory_files_config.enabled`` so a sync without the export + feature configured is a no-op. When any file was modified or deleted + (``force_full``), the tree is rebuilt from the full scoped store so stale + skills/entries do not linger; otherwise an incremental update merges the + just-created resources. Best-effort: the structured memory is already + persisted, so an export error must not fail the sync. + """ + if not getattr(self.memory_files_config, "enabled", False): + return + if not changed_resources and not force_full: + return + try: + await self._build_memory_files(user_scope, changed=None if force_full else changed_resources) + except Exception: + logger.exception("Memory file export failed after workspace memorize") + + def _build_memorize_workspace_workflow(self) -> list[WorkflowStep]: + """The workspace memorize pipeline: direct resource -> file synthesis (ADR 0007 phase 1). + + Unlike single-file :meth:`memorize` (``resource -> entry -> file``), the + workspace path synthesizes files straight from the preprocessed source and + creates no ``RecallEntry``. After ``preprocess`` it: + + - ``create_resource`` — one file maps to one :class:`Resource` (caption/embedding + for INDEX recall), for every track including ``workspace`` (resource-only). + - ``synthesize_files`` — for the ``chat`` and ``skill`` tracks only, route the + source to the files to update/create then synthesize each file's body, upserting + ``RecallFile`` and recording ``resource -> file`` provenance. ``workspace`` is a + no-op here. Retrieval over these files is deferred (ADR 0007 phase 2). + """ + synthesis_profile = getattr(self.memory_files_config, "synthesis_llm_profile", "default") + return [ + WorkflowStep( + step_id="ingest_resource", + role="ingest", + handler=self._memorize_ingest_resource, + requires={"resource_url", "modality"}, + produces={"local_path", "raw_text"}, + capabilities={"io"}, + ), + WorkflowStep( + step_id="preprocess_multimodal", + role="preprocess", + handler=self._memorize_preprocess_multimodal, + requires={"local_path", "modality", "raw_text"}, + produces={"preprocessed_resources"}, + capabilities={"llm"}, + config={"chat_llm_profile": self.memorize_config.preprocess_llm_profile}, + ), + WorkflowStep( + step_id="create_resource", + role="persist", + handler=self._memorize_ws_create_resource, + requires={ + "preprocessed_resources", + "modality", + "local_path", + "resource_url", + "store", + "user", + "resource_track", + }, + produces={"resources"}, + capabilities={"db", "vector"}, + config={"embed_llm_profile": "embedding"}, + ), + WorkflowStep( + step_id="synthesize_files", + role="synthesize_files", + handler=self._memorize_ws_synthesize_files, + requires={"resources", "preprocessed_resources", "resource_track", "store", "user"}, + produces={"files"}, + capabilities={"llm", "db", "vector"}, + config={"chat_llm_profile": synthesis_profile, "embed_llm_profile": "embedding"}, + ), + WorkflowStep( + step_id="build_response", + role="emit", + handler=self._memorize_ws_build_response, + requires={"resources", "files"}, + produces={"response"}, + capabilities=set(), + ), + ] + + @staticmethod + def _list_memorize_initial_keys() -> set[str]: + return { + "resource_url", + "modality", + "memory_types", + "categories_prompt_str", + "ctx", + "store", + "category_ids", + "user", + "allow_new_categories", + "resource_track", + } + + async def _memorize_ingest_resource(self, state: WorkflowState, step_context: Any) -> WorkflowState: + local_path, raw_text = await self.fs.fetch(state["resource_url"], state["modality"]) + state.update({"local_path": local_path, "raw_text": raw_text}) + return state + + # Modalities whose preprocessing analyzes media via the VLM (vision) client. + _VISION_MODALITIES = frozenset({"image", "video"}) + + async def _memorize_preprocess_multimodal(self, state: WorkflowState, step_context: Any) -> WorkflowState: + modality = state["modality"] + client = self._get_step_llm_client(step_context) + if modality in self._VISION_MODALITIES: + with contextlib.suppress(KeyError): + client = self._get_vlm_client(self.memorize_config.vlm_profile, step_context=step_context) + preprocessed = await self._preprocess_resource_url( + local_path=state["local_path"], + text=state.get("raw_text"), + modality=modality, + llm_client=client, + ) + if not preprocessed: + preprocessed = [{"text": state.get("raw_text"), "caption": None}] + state["preprocessed_resources"] = preprocessed + return state + + @staticmethod + def _format_skill_source_content(preprocessed_resources: list[dict[str, Any]]) -> str: + """Flatten a source's preprocessed segments into a single text block.""" + parts = [ + " ".join((prep.get("text") or "").split()) + for prep in preprocessed_resources + if (prep.get("text") or "").strip() + ] + return "\n\n".join(parts) + + # --- Workspace resource -> file path (ADR 0007 phase 1) ------------------- + + # Maps a workspace ``resource_track`` to the ``RecallFile.track`` it synthesizes + # into. ``workspace`` has no entry (resource-only), so it is absent. + _TRACK_TO_FILE_TRACK: ClassVar[dict[str, str]] = {"chat": "memory", "skill": "skill"} + + async def _memorize_ws_create_resource(self, state: WorkflowState, step_context: Any) -> WorkflowState: + """Create the single ``Resource`` for this file (one file -> one resource). + + Runs for every track; the ``workspace`` track stops here (resource-only). The + caption is the joined per-segment captions, embedded for INDEX/resource recall. + """ + embed_client = self._get_step_embedding_client(step_context) + store = state["store"] + preprocessed = state.get("preprocessed_resources") or [] + captions = [(prep.get("caption") or "").strip() for prep in preprocessed] + caption = "\n\n".join(c for c in captions if c) or None + res = await self._create_resource_with_caption( + resource_url=state["resource_url"], + modality=state["modality"], + local_path=state["local_path"], + caption=caption, + store=store, + embed_client=embed_client, + user=state.get("user", {}), + track=state.get("resource_track"), + ) + state["resources"] = [res] + return state + + async def _memorize_ws_synthesize_files(self, state: WorkflowState, step_context: Any) -> WorkflowState: + """Synthesize this source into ``RecallFile``s for the chat/skill tracks. + + Two steps: (a) route the source to the set of files to update/create given the + existing files' names+descriptions, and (b) synthesize each target file's body in + parallel. Persists each file and a ``resource -> file`` provenance link. The + ``workspace`` track (and any source with no content) is a no-op. + """ + track = state.get("resource_track") + file_track = self._TRACK_TO_FILE_TRACK.get(track or "") + resources = state.get("resources") or [] + content = self._format_skill_source_content(state.get("preprocessed_resources") or []) + if file_track is None or not resources or not content: + state["files"] = [] + return state + + store = state["store"] + user_scope = dict(state.get("user") or {}) + llm_client = self._get_step_llm_client(step_context) + embed_client = self._get_step_embedding_client(step_context) + resource = resources[0] + + existing = store.recall_file_repo.list_categories(where={**user_scope, "track": file_track}) + ops = await self._route_source_to_files( + file_track=file_track, content=content, existing=existing, llm_client=llm_client + ) + touched = await self._synthesize_file_ops( + ops=ops, + file_track=file_track, + content=content, + existing=existing, + resource=resource, + store=store, + user_scope=user_scope, + llm_client=llm_client, + embed_client=embed_client, + ) + await self._sync_file_segments( + files=touched, + file_track=file_track, + store=store, + user_scope=user_scope, + embed_client=embed_client, + ) + state["files"] = touched + return state + + @staticmethod + def _segment_texts_for_file(file: RecallFile, file_track: str) -> list[str]: + """Compute the searchable segment texts for a synthesized file (ADR 0007 L2 items). + + The slicing rule is track-specific: + + - ``skill``: a single ``name: ...\\ndescription: ...`` segment for the whole skill. + - ``memory``: one segment per content line, skipping blank lines and markdown + headings (lines starting with one or more ``#``). + + Texts are stripped and de-duplicated while preserving order so a repeated line is + embedded only once. + """ + if file_track == "skill": + return [f"name: {file.name}\ndescription: {file.description}"] + + texts: list[str] = [] + for line in (file.content or "").split("\n"): + stripped = line.strip() + if not stripped or stripped.startswith("#"): + continue + texts.append(stripped) + return list(dict.fromkeys(texts)) + + async def _sync_file_segments( + self, + *, + files: list[RecallFile], + file_track: str, + store: Database, + user_scope: dict[str, Any], + embed_client: Any, + ) -> None: + """Reconcile each file's stored segments with its freshly computed segment texts. + + Diffs the new segment texts against the existing ones and does a drop-and-add on the + difference only: segments whose text disappeared are deleted, and only genuinely new + texts are embedded and inserted. Unchanged lines keep their existing embedding, so an + edit that touches a few lines does not re-embed the whole file. + """ + for file in files: + new_texts = self._segment_texts_for_file(file, file_track) + existing = store.recall_file_segment_repo.list_segments_for_file(file.id) + existing_texts = {seg.text for seg in existing} + new_set = set(new_texts) + + for seg in existing: + if seg.text not in new_set: + store.recall_file_segment_repo.delete_segment(seg.id) + + to_add = [text for text in new_texts if text not in existing_texts] + if not to_add: + continue + vecs = await embed_client.embed(to_add) + for text, vec in zip(to_add, vecs, strict=True): + store.recall_file_segment_repo.create_segment( + recall_file_id=file.id, track=file_track, text=text, embedding=vec, user_data=dict(user_scope) + ) + + async def _route_source_to_files( + self, + *, + file_track: str, + content: str, + existing: Mapping[str, RecallFile], + llm_client: Any, + ) -> list[dict[str, str]]: + """Ask the model which existing files to update / what new files to create.""" + existing_text = self._format_existing_files(existing) or "(none)" + prompt = ( + ROUTE_PROMPTS[file_track] + .replace(EXISTING_PLACEHOLDER, existing_text) + .replace(CONTENT_PLACEHOLDER, self._escape_prompt_value(content)) + ) + return self._parse_file_ops(await llm_client.chat(prompt), existing) + + async def _synthesize_file_ops( + self, + *, + ops: list[dict[str, str]], + file_track: str, + content: str, + existing: Mapping[str, RecallFile], + resource: Resource, + store: Database, + user_scope: dict[str, Any], + llm_client: Any, + embed_client: Any, + ) -> list[RecallFile]: + """Synthesize each routed file's body (in parallel) and persist file + link.""" + existing_by_name = {f.name: f for f in existing.values()} + # Resolve ops to unique targets (dedup by name; last op's description wins). + targets: list[dict[str, Any]] = [] + by_name: dict[str, dict[str, Any]] = {} + for op in ops: + name = op["name"] + ex = existing_by_name.get(name) + description = (op.get("description") or (ex.description if ex else "") or "").strip() + target = by_name.get(name) + if target is None: + target = {"name": name, "description": description, "existing": ex} + by_name[name] = target + targets.append(target) + elif description: + target["description"] = description + if not targets: + return [] + + prompts = [ + SYNTHESIS_PROMPTS[file_track] + .replace(NAME_PLACEHOLDER, self._escape_prompt_value(t["name"])) + .replace(DESCRIPTION_PLACEHOLDER, self._escape_prompt_value(t["description"])) + .replace( + EXISTING_PLACEHOLDER, self._escape_prompt_value((t["existing"].content if t["existing"] else "") or "") + ) + .replace(CONTENT_PLACEHOLDER, self._escape_prompt_value(content)) + for t in targets + ] + bodies = await asyncio.gather(*[llm_client.chat(prompt) for prompt in prompts]) + + # Embed name+description for the files being created. + creates = [t for t in targets if t["existing"] is None] + create_vecs: dict[str, list[float]] = {} + if creates: + emb_texts = [f"{t['name']}: {t['description']}" if t["description"] else t["name"] for t in creates] + vecs = await embed_client.embed(emb_texts) + for t, vec in zip(creates, vecs, strict=True): + create_vecs[t["name"]] = vec + + touched: list[RecallFile] = [] + for target, body in zip(targets, bodies, strict=True): + cleaned = body.replace("```markdown", "").replace("```", "").strip() + file = target["existing"] + if file is None: + file = store.recall_file_repo.get_or_create_category( + name=target["name"], + description=target["description"], + embedding=create_vecs[target["name"]], + user_data=user_scope, + track=file_track, + ) + file = store.recall_file_repo.update_category(category_id=file.id, content=cleaned) + store.recall_file_resource_repo.link_resource_category(resource.id, file.id, user_data=dict(user_scope)) + touched.append(file) + return touched + + @staticmethod + def _format_existing_files(existing: Mapping[str, RecallFile]) -> str: + """Render existing files as ``- name: description`` lines for the router prompt.""" + return "\n".join( + f"- {f.name}: {f.description}" if f.description else f"- {f.name}" + for f in sorted(existing.values(), key=lambda f: f.name) + ) + + def _parse_file_ops(self, raw: str, existing: Mapping[str, RecallFile]) -> list[dict[str, str]]: + """Parse the router's JSON array into validated ``{op, name, description}`` dicts. + + ``update`` ops naming an unknown file are dropped (we never update a file that + does not exist); ``create``/``update`` are otherwise kept with a stripped name. + """ + if not raw: + return [] + start = raw.find("[") + end = raw.rfind("]") + if start == -1 or end == -1 or end <= start: + return [] + try: + parsed = json.loads(raw[start : end + 1]) + except (json.JSONDecodeError, TypeError): + return [] + if not isinstance(parsed, list): + return [] + existing_names = {f.name for f in existing.values()} + ops: list[dict[str, str]] = [] + for entry in parsed: + if not isinstance(entry, dict): + continue + op = entry.get("op") + name = entry.get("name") + if op not in {"update", "create"} or not isinstance(name, str) or not name.strip(): + continue + name = name.strip() + if op == "update" and name not in existing_names: + continue + description = entry.get("description") + description = description.strip() if isinstance(description, str) else "" + ops.append({"op": op, "name": name, "description": description}) + return ops + + def _memorize_ws_build_response(self, state: WorkflowState, step_context: Any) -> WorkflowState: + """Emit the workspace response (no entries; ``categories`` carries touched files).""" + resources = [self._model_dump_without_embeddings(r) for r in state.get("resources", [])] + files = [self._model_dump_without_embeddings(f) for f in state.get("files", [])] + # Keep the legacy response contract (``items``/``categories``); items is always + # empty on this path since the entry plane is gone. + base: dict[str, Any] = {"items": [], "categories": files, "relations": []} + if len(resources) == 1: + state["response"] = {"resource": resources[0], **base} + else: + state["response"] = {"resources": resources, **base} + return state + + async def _create_resource_with_caption( + self, + *, + resource_url: str, + modality: str, + local_path: str, + caption: str | None, + store: Database, + embed_client: Any | None = None, + user: Mapping[str, Any] | None = None, + track: str | None = None, + ) -> Resource: + caption_text = caption.strip() if caption else None + if caption_text: + client = embed_client or self._get_embedding_client() + caption_embedding = (await client.embed([caption_text]))[0] + else: + caption_embedding = None + + res = store.resource_repo.create_resource( + url=resource_url, + modality=modality, + local_path=local_path, + caption=caption_text, + embedding=caption_embedding, + user_data=dict(user or {}), + track=track, + ) + return res + + def _resolve_memory_types(self) -> list[EntryType]: + configured_types = self.memorize_config.memory_types or DEFAULT_MEMORY_TYPES + return [cast(EntryType, mtype) for mtype in configured_types] + + @staticmethod + def _resolve_custom_prompt(prompt: str | CustomPrompt, templates: Mapping[str, str]) -> str: + if isinstance(prompt, str): + return prompt + valid_blocks = [ + (block.ordinal, name, block.prompt or templates.get(name)) + for name, block in prompt.items() + if (block.ordinal >= 0 and (block.prompt or templates.get(name))) + ] + if not valid_blocks: + # raise ValueError(f"No valid blocks contained in custom prompt: {prompt}") + return "" + sorted_blocks = sorted(valid_blocks) + return "\n\n".join(block for (_, _, block) in sorted_blocks if block is not None) + + async def _preprocess_resource_url( + self, *, local_path: str, text: str | None, modality: str, llm_client: Any | None = None + ) -> list[dict[str, str | None]]: + """Preprocess a resource by delegating to the per-format ``preprocess`` package. + + Returns a list of preprocessed resources, each with 'text' and 'caption'. + """ + return await preprocess_resource( + modality=modality, + local_path=local_path, + text=text, + ctx=self._build_preprocess_context(), + llm_client=llm_client, + ) + + def _build_preprocess_context(self) -> PreprocessContext: + """Bundle the service dependencies the preprocessors need.""" + return PreprocessContext( + get_llm_client=self._get_llm_client, + get_vlm_client=lambda: self._get_vlm_client(self.memorize_config.vlm_profile), + escape_prompt_value=self._escape_prompt_value, + extract_json_blob=self._extract_json_blob, + resolve_custom_prompt=self._resolve_custom_prompt, + multimodal_preprocess_prompts=self.memorize_config.multimodal_preprocess_prompts, + ) + + async def _ensure_categories_ready( + self, ctx: Context, store: Database, user_scope: Mapping[str, Any] | None = None + ) -> None: + if ctx.categories_ready: + return + if ctx.category_init_task: + await ctx.category_init_task + ctx.category_init_task = None + return + await self._initialize_categories(ctx, store, user_scope) + + @staticmethod + def _classify_categories( + configs: list[CategoryConfig], + existing_by_name: dict[str, RecallFile], + ) -> tuple[ + list[tuple[int, CategoryConfig]], + list[tuple[int, CategoryConfig, RecallFile]], + dict[int, RecallFile], + ]: + to_create: list[tuple[int, CategoryConfig]] = [] + to_update: list[tuple[int, CategoryConfig, RecallFile]] = [] + ready: dict[int, RecallFile] = {} + for i, cfg in enumerate(configs): + name = cfg.name.strip() or "Untitled" + description = cfg.description.strip() + ex = existing_by_name.get(name) + if ex is None: + to_create.append((i, cfg)) + elif ex.embedding is None or (ex.description or "") != description: + to_update.append((i, cfg, ex)) + else: + ready[i] = ex + return to_create, to_update, ready + + async def _initialize_categories( + self, ctx: Context, store: Database, user: Mapping[str, Any] | None = None + ) -> None: + if ctx.categories_ready: + return + if not self.category_configs: + ctx.categories_ready = True + return + + user_data = dict(user or {}) + existing = store.recall_file_repo.list_categories(where={**user_data, "track": "memory"}) + existing_by_name: dict[str, RecallFile] = {c.name: c for c in existing.values()} + + to_create, to_update, ready = self._classify_categories(self.category_configs, existing_by_name) + + needs_embed: list[tuple[int, CategoryConfig]] = [] + needs_embed.extend(to_create) + needs_embed.extend((i, cfg) for i, cfg, _ in to_update) + + embed_map: dict[int, list[float]] = {} + if needs_embed: + texts = [self._category_embedding_text(cfg) for _, cfg in needs_embed] + vecs = await self._get_embedding_client("embedding").embed(texts) + for (i, _), vec in zip(needs_embed, vecs, strict=True): + embed_map[i] = vec + + cats: dict[int, RecallFile] = dict(ready) + + for i, cfg in to_create: + name = cfg.name.strip() or "Untitled" + description = cfg.description.strip() + cat = store.recall_file_repo.get_or_create_category( + name=name, description=description, embedding=embed_map[i], user_data=user_data + ) + cats[i] = cat + + for i, cfg, ex in to_update: + description = cfg.description.strip() + cat = store.recall_file_repo.update_category( + category_id=ex.id, description=description, embedding=embed_map[i] + ) + cats[i] = cat + + ctx.category_ids = [] + ctx.category_name_to_id = {} + for i in range(len(self.category_configs)): + cat = cats[i] + ctx.category_ids.append(cat.id) + name = self.category_configs[i].name.strip() or "Untitled" + ctx.category_name_to_id[name.lower()] = cat.id + ctx.categories_ready = True + + @staticmethod + def _category_embedding_text(cat: CategoryConfig) -> str: + name = cat.name.strip() or "Untitled" + desc = cat.description.strip() + return f"{name}: {desc}" if desc else name diff --git a/src/memu/app/retrieve.py b/src/memu/app/retrieve.py index feec70a0..1e17c824 100644 --- a/src/memu/app/retrieve.py +++ b/src/memu/app/retrieve.py @@ -20,14 +20,13 @@ if TYPE_CHECKING: from memu.app.service import Context - from memu.app.settings import RetrieveConfig, RetrieveWorkspaceConfig + from memu.app.settings import RetrieveConfig from memu.database.interfaces import Database class RetrieveMixin: if TYPE_CHECKING: retrieve_config: RetrieveConfig - retrieve_workspace_config: RetrieveWorkspaceConfig _run_workflow: Callable[..., Awaitable[WorkflowState]] _get_context: Callable[[], Context] _get_database: Callable[[], Database] @@ -86,50 +85,6 @@ async def retrieve( raise RuntimeError(msg) return response - async def retrieve_workspace( - self, - query: str, - where: dict[str, Any] | None = None, - ) -> dict[str, Any]: - """Single-shot, LLM-free retrieval over the segment/file/resource layers. - - Mirrors the relation between :meth:`memorize` and ``memorize_workspace``: - a simpler entry point built on the same store and workflow machinery. The - query is embedded once and used to rank two layers by vector similarity — - no intention routing, sufficiency checks, or summarization: - - * ``segments``: :class:`RecallFileSegment` slices ranked by embedding, - ``file.top_k`` of them. - * ``files``: the :class:`RecallFile`\\ s pointed to by those segments — not - a ranked search, just a roll-up. Each file's score is the max score of - the segments that point to it. - * ``resources``: workspace-track resources ranked by embedding, - ``resource.top_k`` of them. - - The entry layer is disabled here (its config is retained but ignored). - Returns ``segments``, ``files``, and ``resources``. - """ - if not query or not query.strip(): - raise ValueError("empty_query") - store = self._get_database() - where_filters = self._normalize_where(where) - config = self.retrieve_workspace_config - - state: WorkflowState = { - "query": query, - "store": store, - "where": where_filters, - "retrieve_file": config.file.enabled, - "retrieve_resource": config.resource.enabled, - } - - result = await self._run_workflow("retrieve_workspace", state) - response = cast(dict[str, Any] | None, result.get("response")) - if response is None: - msg = "Retrieve workspace workflow failed to produce a response" - raise RuntimeError(msg) - return response - def _normalize_where(self, where: Mapping[str, Any] | None) -> dict[str, Any]: """Validate and clean the `where` scope filters against the configured user model.""" if not where: @@ -1470,152 +1425,3 @@ def _format_llm_resource_content(self, hits: list[dict[str, Any]]) -> str: caption = res.get("caption", "") or f"Resource {res['url']}" lines.append(f"Resource: {caption}") return "\n\n".join(lines).strip() - - def _build_retrieve_workspace_workflow(self) -> list[WorkflowStep]: - """The simple embedding-only workspace retrieval pipeline. - - A segment recall step ranks :class:`RecallFileSegment` slices by embedding; - a file roll-up step gathers the files those segments point to; a resource - recall step ranks workspace-track resources by embedding. A terminal step - assembles the response. None of the routing/sufficiency machinery of - ``retrieve_rag`` applies. The query vector is embedded by the first recall - step and reused downstream. - """ - steps = [ - WorkflowStep( - step_id="recall_segments", - role="recall_segments", - handler=self._ws_recall_segments, - requires={"retrieve_file", "query", "store", "where"}, - produces={"segment_hits", "segment_pool", "query_vector"}, - capabilities={"vector"}, - config={"embed_llm_profile": "embedding"}, - ), - WorkflowStep( - step_id="collect_files", - role="collect_files", - handler=self._ws_collect_files, - requires={"retrieve_file", "segment_hits", "segment_pool", "store", "where"}, - produces={"file_hits", "file_pool"}, - capabilities=set(), - ), - WorkflowStep( - step_id="recall_resources", - role="recall_resources", - handler=self._ws_recall_resources, - requires={"retrieve_resource", "query", "store", "where", "query_vector"}, - produces={"resource_hits", "resource_pool", "query_vector"}, - capabilities={"vector"}, - config={"embed_llm_profile": "embedding"}, - ), - WorkflowStep( - step_id="build_response", - role="build_context", - handler=self._ws_build_response, - requires={ - "segment_hits", - "segment_pool", - "file_hits", - "file_pool", - "resource_hits", - "resource_pool", - }, - produces={"response"}, - capabilities=set(), - ), - ] - return steps - - @staticmethod - def _list_retrieve_workspace_initial_keys() -> set[str]: - return {"query", "store", "where", "retrieve_file", "retrieve_resource"} - - async def _ws_query_vector(self, state: WorkflowState, step_context: Any) -> list[float]: - """Embed the query once and cache it on the state for reuse across steps.""" - cached = state.get("query_vector") - if cached is not None: - return cast(list[float], cached) - embed_client = self._get_step_embedding_client(step_context) - qvec = (await embed_client.embed([state["query"]]))[0] - state["query_vector"] = qvec - return cast(list[float], qvec) - - async def _ws_recall_segments(self, state: WorkflowState, step_context: Any) -> WorkflowState: - if not state.get("retrieve_file"): - state["segment_hits"] = [] - state["segment_pool"] = {} - state.setdefault("query_vector", None) - return state - - store = state["store"] - # The segment repo has no vector search, so rank the stored segment - # embeddings directly, mirroring how files used to be ranked. Optionally - # scope to the requested tracks via the denormalized segment ``track``. - segment_where = dict(state.get("where") or {}) - tracks = self.retrieve_workspace_config.file.tracks - if tracks: - segment_where["track__in"] = list(tracks) - segment_pool = {seg.id: seg for seg in store.recall_file_segment_repo.list_segments(segment_where)} - qvec = await self._ws_query_vector(state, step_context) - state["segment_hits"] = cosine_topk( - qvec, - [(sid, seg.embedding) for sid, seg in segment_pool.items()], - k=self.retrieve_workspace_config.file.top_k, - ) - state["segment_pool"] = segment_pool - return state - - async def _ws_collect_files(self, state: WorkflowState, _: Any) -> WorkflowState: - """Roll the ranked segments up to their files (no ranked file search). - - Every file pointed to by a top segment is returned; a file's score is the - max score across the segments that point to it. - """ - segment_hits = state.get("segment_hits") or [] - segment_pool = state.get("segment_pool") or {} - store = state["store"] - where_filters = state.get("where") or {} - file_pool = store.recall_file_repo.list_categories(where_filters) - - file_scores: dict[str, float] = {} - for seg_id, score in segment_hits: - seg = segment_pool.get(seg_id) - if seg is None: - continue - fid = seg.recall_file_id - if fid not in file_pool: - continue - score = float(score) - if fid not in file_scores or score > file_scores[fid]: - file_scores[fid] = score - - # Preserve descending-score order so the response reads best-first. - state["file_hits"] = sorted(file_scores.items(), key=lambda kv: kv[1], reverse=True) - state["file_pool"] = file_pool - return state - - async def _ws_recall_resources(self, state: WorkflowState, step_context: Any) -> WorkflowState: - if not state.get("retrieve_resource"): - state["resource_hits"] = [] - state["resource_pool"] = {} - return state - - store = state["store"] - # Workspace retrieval only surfaces resources ingested by - # ``memorize_workspace`` (track="workspace"); other tracks are excluded. - resource_where = {**(state.get("where") or {}), "track": "workspace"} - resource_pool = store.resource_repo.list_resources(resource_where) - qvec = await self._ws_query_vector(state, step_context) - state["resource_hits"] = store.resource_repo.vector_search_resources( - qvec, self.retrieve_workspace_config.resource.top_k, where=resource_where - ) - state["resource_pool"] = resource_pool - return state - - def _ws_build_response(self, state: WorkflowState, _: Any) -> WorkflowState: - state["response"] = { - "segments": self._materialize_hits(state.get("segment_hits", []), state.get("segment_pool", {})), - "files": self._materialize_hits(state.get("file_hits", []), state.get("file_pool", {})), - "resources": self._materialize_hits(state.get("resource_hits", []), state.get("resource_pool", {})), - } - return state diff --git a/src/memu/app/retrieve_workspace.py b/src/memu/app/retrieve_workspace.py new file mode 100644 index 00000000..08406386 --- /dev/null +++ b/src/memu/app/retrieve_workspace.py @@ -0,0 +1,249 @@ +from __future__ import annotations + +import logging +from collections.abc import Awaitable, Callable, Mapping, Sequence +from typing import TYPE_CHECKING, Any, cast + +from pydantic import BaseModel + +from memu.vector import cosine_topk +from memu.workflow.step import WorkflowState, WorkflowStep + +logger = logging.getLogger(__name__) + +if TYPE_CHECKING: + from memu.app.settings import RetrieveWorkspaceConfig + from memu.database.interfaces import Database + + +class RetrieveWorkspaceMixin: + if TYPE_CHECKING: + retrieve_workspace_config: RetrieveWorkspaceConfig + _run_workflow: Callable[..., Awaitable[WorkflowState]] + _get_database: Callable[[], Database] + _get_step_embedding_client: Callable[[Mapping[str, Any] | None], Any] + _model_dump_without_embeddings: Callable[[BaseModel], dict[str, Any]] + user_model: type[BaseModel] + + async def retrieve_workspace( + self, + query: str, + where: dict[str, Any] | None = None, + ) -> dict[str, Any]: + """Single-shot, LLM-free retrieval over the segment/file/resource layers. + + Mirrors the relation between :meth:`memorize` and ``memorize_workspace``: + a simpler entry point built on the same store and workflow machinery. The + query is embedded once and used to rank two layers by vector similarity — + no intention routing, sufficiency checks, or summarization: + + * ``segments``: :class:`RecallFileSegment` slices ranked by embedding, + ``file.top_k`` of them. + * ``files``: the :class:`RecallFile`\\ s pointed to by those segments — not + a ranked search, just a roll-up. Each file's score is the max score of + the segments that point to it. + * ``resources``: workspace-track resources ranked by embedding, + ``resource.top_k`` of them. + + The entry layer is disabled here (its config is retained but ignored). + Returns ``segments``, ``files``, and ``resources``. + """ + if not query or not query.strip(): + raise ValueError("empty_query") + store = self._get_database() + where_filters = self._normalize_where(where) + config = self.retrieve_workspace_config + + state: WorkflowState = { + "query": query, + "store": store, + "where": where_filters, + "retrieve_file": config.file.enabled, + "retrieve_resource": config.resource.enabled, + } + + result = await self._run_workflow("retrieve_workspace", state) + response = cast(dict[str, Any] | None, result.get("response")) + if response is None: + msg = "Retrieve workspace workflow failed to produce a response" + raise RuntimeError(msg) + return response + + def _normalize_where(self, where: Mapping[str, Any] | None) -> dict[str, Any]: + """Validate and clean the `where` scope filters against the configured user model.""" + if not where: + return {} + + valid_fields = set(getattr(self.user_model, "model_fields", {}).keys()) + cleaned: dict[str, Any] = {} + + for raw_key, value in where.items(): + if value is None: + continue + field = raw_key.split("__", 1)[0] + if field not in valid_fields: + msg = f"Unknown filter field '{field}' for current user scope" + raise ValueError(msg) + cleaned[raw_key] = value + + return cleaned + + def _build_retrieve_workspace_workflow(self) -> list[WorkflowStep]: + """The simple embedding-only workspace retrieval pipeline. + + A segment recall step ranks :class:`RecallFileSegment` slices by embedding; + a file roll-up step gathers the files those segments point to; a resource + recall step ranks workspace-track resources by embedding. A terminal step + assembles the response. None of the routing/sufficiency machinery of + ``retrieve_rag`` applies. The query vector is embedded by the first recall + step and reused downstream. + """ + steps = [ + WorkflowStep( + step_id="recall_segments", + role="recall_segments", + handler=self._ws_recall_segments, + requires={"retrieve_file", "query", "store", "where"}, + produces={"segment_hits", "segment_pool", "query_vector"}, + capabilities={"vector"}, + config={"embed_llm_profile": "embedding"}, + ), + WorkflowStep( + step_id="collect_files", + role="collect_files", + handler=self._ws_collect_files, + requires={"retrieve_file", "segment_hits", "segment_pool", "store", "where"}, + produces={"file_hits", "file_pool"}, + capabilities=set(), + ), + WorkflowStep( + step_id="recall_resources", + role="recall_resources", + handler=self._ws_recall_resources, + requires={"retrieve_resource", "query", "store", "where", "query_vector"}, + produces={"resource_hits", "resource_pool", "query_vector"}, + capabilities={"vector"}, + config={"embed_llm_profile": "embedding"}, + ), + WorkflowStep( + step_id="build_response", + role="build_context", + handler=self._ws_build_response, + requires={ + "segment_hits", + "segment_pool", + "file_hits", + "file_pool", + "resource_hits", + "resource_pool", + }, + produces={"response"}, + capabilities=set(), + ), + ] + return steps + + @staticmethod + def _list_retrieve_workspace_initial_keys() -> set[str]: + return {"query", "store", "where", "retrieve_file", "retrieve_resource"} + + async def _ws_query_vector(self, state: WorkflowState, step_context: Any) -> list[float]: + """Embed the query once and cache it on the state for reuse across steps.""" + cached = state.get("query_vector") + if cached is not None: + return cast(list[float], cached) + embed_client = self._get_step_embedding_client(step_context) + qvec = (await embed_client.embed([state["query"]]))[0] + state["query_vector"] = qvec + return cast(list[float], qvec) + + async def _ws_recall_segments(self, state: WorkflowState, step_context: Any) -> WorkflowState: + if not state.get("retrieve_file"): + state["segment_hits"] = [] + state["segment_pool"] = {} + state.setdefault("query_vector", None) + return state + + store = state["store"] + # The segment repo has no vector search, so rank the stored segment + # embeddings directly, mirroring how files used to be ranked. Optionally + # scope to the requested tracks via the denormalized segment ``track``. + segment_where = dict(state.get("where") or {}) + tracks = self.retrieve_workspace_config.file.tracks + if tracks: + segment_where["track__in"] = list(tracks) + segment_pool = {seg.id: seg for seg in store.recall_file_segment_repo.list_segments(segment_where)} + qvec = await self._ws_query_vector(state, step_context) + state["segment_hits"] = cosine_topk( + qvec, + [(sid, seg.embedding) for sid, seg in segment_pool.items()], + k=self.retrieve_workspace_config.file.top_k, + ) + state["segment_pool"] = segment_pool + return state + + async def _ws_collect_files(self, state: WorkflowState, _: Any) -> WorkflowState: + """Roll the ranked segments up to their files (no ranked file search). + + Every file pointed to by a top segment is returned; a file's score is the + max score across the segments that point to it. + """ + segment_hits = state.get("segment_hits") or [] + segment_pool = state.get("segment_pool") or {} + store = state["store"] + where_filters = state.get("where") or {} + file_pool = store.recall_file_repo.list_categories(where_filters) + + file_scores: dict[str, float] = {} + for seg_id, score in segment_hits: + seg = segment_pool.get(seg_id) + if seg is None: + continue + fid = seg.recall_file_id + if fid not in file_pool: + continue + score = float(score) + if fid not in file_scores or score > file_scores[fid]: + file_scores[fid] = score + + # Preserve descending-score order so the response reads best-first. + state["file_hits"] = sorted(file_scores.items(), key=lambda kv: kv[1], reverse=True) + state["file_pool"] = file_pool + return state + + async def _ws_recall_resources(self, state: WorkflowState, step_context: Any) -> WorkflowState: + if not state.get("retrieve_resource"): + state["resource_hits"] = [] + state["resource_pool"] = {} + return state + + store = state["store"] + # Workspace retrieval only surfaces resources ingested by + # ``memorize_workspace`` (track="workspace"); other tracks are excluded. + resource_where = {**(state.get("where") or {}), "track": "workspace"} + resource_pool = store.resource_repo.list_resources(resource_where) + qvec = await self._ws_query_vector(state, step_context) + state["resource_hits"] = store.resource_repo.vector_search_resources( + qvec, self.retrieve_workspace_config.resource.top_k, where=resource_where + ) + state["resource_pool"] = resource_pool + return state + + def _ws_build_response(self, state: WorkflowState, _: Any) -> WorkflowState: + state["response"] = { + "segments": self._materialize_hits(state.get("segment_hits", []), state.get("segment_pool", {})), + "files": self._materialize_hits(state.get("file_hits", []), state.get("file_pool", {})), + "resources": self._materialize_hits(state.get("resource_hits", []), state.get("resource_pool", {})), + } + return state + + def _materialize_hits(self, hits: Sequence[tuple[str, float]], pool: dict[str, Any]) -> list[dict[str, Any]]: + out = [] + for _id, score in hits: + obj = pool.get(_id) + if not obj: + continue + data = self._model_dump_without_embeddings(obj) + data["score"] = float(score) + out.append(data) + return out diff --git a/src/memu/app/service.py b/src/memu/app/service.py index abb9865b..6e1a017a 100644 --- a/src/memu/app/service.py +++ b/src/memu/app/service.py @@ -10,8 +10,10 @@ from memu.app.client_pool import ClientPool from memu.app.crud import CRUDMixin from memu.app.memorize import MemorizeMixin +from memu.app.memorize_workspace import MemorizeWorkspaceMixin from memu.app.memory_files import MemoryFilesBuilder from memu.app.retrieve import RetrieveMixin +from memu.app.retrieve_workspace import RetrieveWorkspaceMixin from memu.app.settings import ( BlobConfig, CategoryConfig, @@ -57,7 +59,7 @@ class Context: category_init_task: asyncio.Task | None = None -class MemoryService(MemorizeMixin, RetrieveMixin, CRUDMixin): +class MemoryService(MemorizeMixin, MemorizeWorkspaceMixin, RetrieveMixin, RetrieveWorkspaceMixin, CRUDMixin): def __init__( self, *, From 816d910723eca91ff60e4eef5079a9157e85e5da Mon Sep 17 00:00:00 2001 From: wu Date: Fri, 3 Jul 2026 07:27:27 +0900 Subject: [PATCH 6/9] fix: add resource urls to file retrieve response --- src/memu/app/retrieve_workspace.py | 31 +++++++++++++++++-- .../repositories/recall_file_resource_repo.py | 6 ++-- src/memu/database/models.py | 2 +- src/memu/database/postgres/models.py | 4 +-- .../repositories/recall_file_resource_repo.py | 14 ++++----- src/memu/database/sqlite/models.py | 4 +-- .../repositories/recall_file_resource_repo.py | 24 +++++++------- tests/test_skill_track.py | 2 +- 8 files changed, 55 insertions(+), 32 deletions(-) diff --git a/src/memu/app/retrieve_workspace.py b/src/memu/app/retrieve_workspace.py index 08406386..b07be8f6 100644 --- a/src/memu/app/retrieve_workspace.py +++ b/src/memu/app/retrieve_workspace.py @@ -113,7 +113,7 @@ def _build_retrieve_workspace_workflow(self) -> list[WorkflowStep]: role="collect_files", handler=self._ws_collect_files, requires={"retrieve_file", "segment_hits", "segment_pool", "store", "where"}, - produces={"file_hits", "file_pool"}, + produces={"file_hits", "file_pool", "file_resource_urls"}, capabilities=set(), ), WorkflowStep( @@ -134,6 +134,7 @@ def _build_retrieve_workspace_workflow(self) -> list[WorkflowStep]: "segment_pool", "file_hits", "file_pool", + "file_resource_urls", "resource_hits", "resource_pool", }, @@ -209,8 +210,30 @@ async def _ws_collect_files(self, state: WorkflowState, _: Any) -> WorkflowState # Preserve descending-score order so the response reads best-first. state["file_hits"] = sorted(file_scores.items(), key=lambda kv: kv[1], reverse=True) state["file_pool"] = file_pool + state["file_resource_urls"] = self._ws_collect_file_resource_urls(store, where_filters, file_pool) return state + @staticmethod + def _ws_collect_file_resource_urls( + store: Database, where_filters: dict[str, Any], file_pool: dict[str, Any] + ) -> dict[str, list[str]]: + """Map each file id to the URLs of the resources linked to it. + + Resolves the ``RecallFileResource`` link table (file -> resource) and the + resource records (resource -> url) within the current scope, surfacing url + strings only — the raw resource/link ids are not exposed to callers. + """ + resources = store.resource_repo.list_resources(where_filters) + file_resource_urls: dict[str, list[str]] = {} + for rel in store.recall_file_resource_repo.list_relations(where_filters): + if rel.file_id not in file_pool: + continue + resource = resources.get(rel.resource_id) + if resource is None: + continue + file_resource_urls.setdefault(rel.file_id, []).append(resource.url) + return file_resource_urls + async def _ws_recall_resources(self, state: WorkflowState, step_context: Any) -> WorkflowState: if not state.get("retrieve_resource"): state["resource_hits"] = [] @@ -230,9 +253,13 @@ async def _ws_recall_resources(self, state: WorkflowState, step_context: Any) -> return state def _ws_build_response(self, state: WorkflowState, _: Any) -> WorkflowState: + files = self._materialize_hits(state.get("file_hits", []), state.get("file_pool", {})) + file_resource_urls = state.get("file_resource_urls", {}) + for file in files: + file["resource_urls"] = file_resource_urls.get(file["id"], []) state["response"] = { "segments": self._materialize_hits(state.get("segment_hits", []), state.get("segment_pool", {})), - "files": self._materialize_hits(state.get("file_hits", []), state.get("file_pool", {})), + "files": files, "resources": self._materialize_hits(state.get("resource_hits", []), state.get("resource_pool", {})), } return state diff --git a/src/memu/database/inmemory/repositories/recall_file_resource_repo.py b/src/memu/database/inmemory/repositories/recall_file_resource_repo.py index cccc251e..909c6855 100644 --- a/src/memu/database/inmemory/repositories/recall_file_resource_repo.py +++ b/src/memu/database/inmemory/repositories/recall_file_resource_repo.py @@ -24,10 +24,10 @@ def list_relations(self, where: Mapping[str, Any] | None = None) -> list[RecallF def link_resource_category(self, resource_id: str, cat_id: str, user_data: dict[str, Any]) -> RecallFileResource: _ = resource_id # enforced by caller via existing state for rel in self.relations: - if rel.resource_id == resource_id and rel.category_id == cat_id: + if rel.resource_id == resource_id and rel.file_id == cat_id: return rel rel = self.recall_file_resource_model( - id=str(uuid.uuid4()), resource_id=resource_id, category_id=cat_id, **user_data + id=str(uuid.uuid4()), resource_id=resource_id, file_id=cat_id, **user_data ) self.relations.append(rel) return rel @@ -45,7 +45,7 @@ def unlink_resource_category(self, resource_id: str, cat_id: str) -> None: # this repo's view never diverge (rebinding self.relations would orphan the # shared state.resource_relations list). self.relations[:] = [ - rel for rel in self.relations if not (rel.resource_id == resource_id and rel.category_id == cat_id) + rel for rel in self.relations if not (rel.resource_id == resource_id and rel.file_id == cat_id) ] def unlink_resource(self, resource_id: str) -> list[RecallFileResource]: diff --git a/src/memu/database/models.py b/src/memu/database/models.py index 51171d47..94f79707 100644 --- a/src/memu/database/models.py +++ b/src/memu/database/models.py @@ -113,7 +113,7 @@ class RecallFileEntry(BaseRecord): class RecallFileResource(BaseRecord): resource_id: str - category_id: str + file_id: str class RecallFileSegment(BaseRecord): diff --git a/src/memu/database/postgres/models.py b/src/memu/database/postgres/models.py index 9b7a89cc..cd95ed8f 100644 --- a/src/memu/database/postgres/models.py +++ b/src/memu/database/postgres/models.py @@ -86,9 +86,9 @@ class RecallFileEntryModel(BaseModelMixin, RecallFileEntry): class RecallFileResourceModel(BaseModelMixin, RecallFileResource): resource_id: str = Field(sa_column=Column(ForeignKey("resources.id", ondelete="CASCADE"), nullable=False)) - category_id: str = Field(sa_column=Column(ForeignKey("memory_categories.id", ondelete="CASCADE"), nullable=False)) + file_id: str = Field(sa_column=Column(ForeignKey("memory_categories.id", ondelete="CASCADE"), nullable=False)) - __table_args__ = (Index("idx_recall_file_resources_unique", "resource_id", "category_id", unique=True),) + __table_args__ = (Index("idx_recall_file_resources_unique", "resource_id", "file_id", unique=True),) class RecallFileSegmentModel(BaseModelMixin, RecallFileSegment): diff --git a/src/memu/database/postgres/repositories/recall_file_resource_repo.py b/src/memu/database/postgres/repositories/recall_file_resource_repo.py index 8045a5ca..5a4ba418 100644 --- a/src/memu/database/postgres/repositories/recall_file_resource_repo.py +++ b/src/memu/database/postgres/repositories/recall_file_resource_repo.py @@ -37,13 +37,13 @@ def link_resource_category(self, resource_id: str, cat_id: str, user_data: dict[ # Avoid duplicate inserts using local cache for rel in self.relations: - if rel.resource_id == resource_id and rel.category_id == cat_id: + if rel.resource_id == resource_id and rel.file_id == cat_id: return rel now = self._now() new_rel = self._recall_file_resource_model( resource_id=resource_id, - category_id=cat_id, + file_id=cat_id, **user_data, created_at=now, updated_at=now, @@ -53,7 +53,7 @@ def link_resource_category(self, resource_id: str, cat_id: str, user_data: dict[ existing = session.scalar( select(self._sqla_models.RecallFileResource).where( self._sqla_models.RecallFileResource.resource_id == resource_id, - self._sqla_models.RecallFileResource.category_id == cat_id, + self._sqla_models.RecallFileResource.file_id == cat_id, ) ) if existing: @@ -72,19 +72,17 @@ def unlink_resource_category(self, resource_id: str, cat_id: str) -> None: session.exec( delete(self._sqla_models.RecallFileResource).where( self._sqla_models.RecallFileResource.resource_id == resource_id, - self._sqla_models.RecallFileResource.category_id == cat_id, + self._sqla_models.RecallFileResource.file_id == cat_id, ) ) session.commit() - self.relations[:] = [ - r for r in self.relations if not (r.resource_id == resource_id and r.category_id == cat_id) - ] + self.relations[:] = [r for r in self.relations if not (r.resource_id == resource_id and r.file_id == cat_id)] def _row_to_record(self, row: Any) -> RecallFileResource: return RecallFileResource( id=row.id, resource_id=row.resource_id, - category_id=row.category_id, + file_id=row.file_id, created_at=row.created_at, updated_at=row.updated_at, **self._scope_kwargs_from(row), diff --git a/src/memu/database/sqlite/models.py b/src/memu/database/sqlite/models.py index 38b326f5..6832d61f 100644 --- a/src/memu/database/sqlite/models.py +++ b/src/memu/database/sqlite/models.py @@ -100,9 +100,9 @@ class SQLiteRecallFileResourceModel(SQLiteBaseModelMixin, RecallFileResource): """SQLite category-resource relation model.""" resource_id: str = Field(sa_column=Column(String, nullable=False)) - category_id: str = Field(sa_column=Column(String, nullable=False)) + file_id: str = Field(sa_column=Column(String, nullable=False)) - __table_args__ = (Index("idx_sqlite_recall_file_resources_unique", "resource_id", "category_id", unique=True),) + __table_args__ = (Index("idx_sqlite_recall_file_resources_unique", "resource_id", "file_id", unique=True),) class SQLiteRecallFileSegmentModel(SQLiteBaseModelMixin, RecallFileSegment): diff --git a/src/memu/database/sqlite/repositories/recall_file_resource_repo.py b/src/memu/database/sqlite/repositories/recall_file_resource_repo.py index d894fc0d..a82fd8ba 100644 --- a/src/memu/database/sqlite/repositories/recall_file_resource_repo.py +++ b/src/memu/database/sqlite/repositories/recall_file_resource_repo.py @@ -69,7 +69,7 @@ def list_relations(self, where: Mapping[str, Any] | None = None) -> list[RecallF rel = RecallFileResource( id=row.id, resource_id=row.resource_id, - category_id=row.category_id, + file_id=row.file_id, created_at=row.created_at, updated_at=row.updated_at, **self._scope_kwargs_from(row), @@ -81,14 +81,12 @@ def list_relations(self, where: Mapping[str, Any] | None = None) -> list[RecallF return result - def link_resource_category( - self, resource_id: str, category_id: str, user_data: dict[str, Any] - ) -> RecallFileResource: + def link_resource_category(self, resource_id: str, file_id: str, user_data: dict[str, Any]) -> RecallFileResource: """Create a link between a resource and a category. Args: resource_id: Resource ID. - category_id: Category ID. + file_id: File ID. user_data: User scope data. Returns: @@ -97,7 +95,7 @@ def link_resource_category( # Check if relation already exists where: dict[str, Any] = { "resource_id": resource_id, - "category_id": category_id, + "file_id": file_id, **user_data, } with self._sessions.session() as session: @@ -111,7 +109,7 @@ def link_resource_category( rel = RecallFileResource( id=existing.id, resource_id=existing.resource_id, - category_id=existing.category_id, + file_id=existing.file_id, created_at=existing.created_at, updated_at=existing.updated_at, **self._scope_kwargs_from(existing), @@ -122,7 +120,7 @@ def link_resource_category( now = self._now() row = self._recall_file_resource_model( resource_id=resource_id, - category_id=category_id, + file_id=file_id, created_at=now, updated_at=now, **user_data, @@ -134,7 +132,7 @@ def link_resource_category( rel = RecallFileResource( id=row.id, resource_id=row.resource_id, - category_id=row.category_id, + file_id=row.file_id, created_at=row.created_at, updated_at=row.updated_at, **user_data, @@ -142,17 +140,17 @@ def link_resource_category( self.relations.append(rel) return rel - def unlink_resource_category(self, resource_id: str, category_id: str) -> None: + def unlink_resource_category(self, resource_id: str, file_id: str) -> None: """Remove a link between a resource and a category. Args: resource_id: Resource ID. - category_id: Category ID. + file_id: File ID. """ with self._sessions.session() as session: stmt = select(self._recall_file_resource_model).where( self._recall_file_resource_model.resource_id == resource_id, - self._recall_file_resource_model.category_id == category_id, + self._recall_file_resource_model.file_id == file_id, ) row = session.exec(stmt).first() if row: @@ -160,7 +158,7 @@ def unlink_resource_category(self, resource_id: str, category_id: str) -> None: session.commit() # Remove from cache self.relations[:] = [ - r for r in self.relations if not (r.resource_id == resource_id and r.category_id == category_id) + r for r in self.relations if not (r.resource_id == resource_id and r.file_id == file_id) ] def unlink_resource(self, resource_id: str) -> list[RecallFileResource]: diff --git a/tests/test_skill_track.py b/tests/test_skill_track.py index daea3fbb..7e736301 100644 --- a/tests/test_skill_track.py +++ b/tests/test_skill_track.py @@ -106,7 +106,7 @@ async def test_skill_track_synthesizes_file_and_links_resource(tmp_path: Path) - links = store.recall_file_resource_repo.list_relations(where=user) assert len(links) == 1 assert links[0].resource_id == res.id - assert links[0].category_id == skill.id + assert links[0].file_id == skill.id async def test_chat_track_routes_to_memory_track_file(tmp_path: Path) -> None: From faeb041d5da3788bbad07a557497fe3fb0258162 Mon Sep 17 00:00:00 2001 From: wu Date: Fri, 3 Jul 2026 20:13:52 +0900 Subject: [PATCH 7/9] feat: add v2 tables & migration --- scripts/db.py | 89 ++++++++++ src/memu/app/settings.py | 2 +- src/memu/database/postgres/migration.py | 64 +++---- src/memu/database/postgres/migrations/env.py | 41 ++++- .../postgres/migrations/script.py.mako | 28 +++ .../versions/20260703_0001_initial_schema.py | 161 ++++++++++++++++++ src/memu/database/postgres/models.py | 20 ++- src/memu/database/postgres/schema.py | 19 ++- 8 files changed, 367 insertions(+), 57 deletions(-) create mode 100644 scripts/db.py create mode 100644 src/memu/database/postgres/migrations/script.py.mako create mode 100644 src/memu/database/postgres/migrations/versions/20260703_0001_initial_schema.py diff --git a/scripts/db.py b/scripts/db.py new file mode 100644 index 00000000..689e2866 --- /dev/null +++ b/scripts/db.py @@ -0,0 +1,89 @@ +#!/usr/bin/env python +"""Programmatic Alembic entry point for the Postgres backend. + +The migration environment is parameterized by a user *scope model* (see +``memu.database.postgres.schema.get_metadata``). The bare ``alembic`` CLI cannot +pass that in, so this wrapper builds the config via +``memu.database.postgres.migration.make_alembic_config`` and drives Alembic's +command API directly. + +The default scope model is ``None`` (the base schema with no scope columns), +which is the schema committed under ``migrations/versions``. + +Usage: + python scripts/db.py revision -m "add foo" # autogenerate a revision + python scripts/db.py upgrade [head] # apply migrations + python scripts/db.py downgrade -1 # revert one revision + python scripts/db.py current # show applied revision + python scripts/db.py history # show revision history + +DSN resolution: --dsn, else $MEMU_DB_DSN, else a localhost default. +""" + +from __future__ import annotations + +import argparse +import os +import sys + +from alembic import command + +from memu.database.postgres.migration import make_alembic_config + +DEFAULT_DSN = "postgresql+psycopg://postgres:postgres@localhost:5432/memu" + + +def _config(dsn: str): + # scope_model=None -> base schema (matches committed baseline revision). + return make_alembic_config(dsn=dsn, scope_model=None) + + +def main(argv: list[str] | None = None) -> int: + parser = argparse.ArgumentParser(description="MemU Postgres migrations") + parser.add_argument( + "--dsn", + default=os.environ.get("MEMU_DB_DSN", DEFAULT_DSN), + help="SQLAlchemy DSN (default: $MEMU_DB_DSN or localhost memu db)", + ) + sub = parser.add_subparsers(dest="cmd", required=True) + + p_rev = sub.add_parser("revision", help="create a new revision") + p_rev.add_argument("-m", "--message", required=True) + p_rev.add_argument( + "--no-autogenerate", + action="store_true", + help="create an empty revision instead of diffing against the DB", + ) + + p_up = sub.add_parser("upgrade", help="apply migrations") + p_up.add_argument("revision", nargs="?", default="head") + + p_down = sub.add_parser("downgrade", help="revert migrations") + p_down.add_argument("revision") + + sub.add_parser("current", help="show current revision") + + p_hist = sub.add_parser("history", help="show revision history") + p_hist.add_argument("-v", "--verbose", action="store_true") + + args = parser.parse_args(argv) + cfg = _config(args.dsn) + + if args.cmd == "revision": + command.revision(cfg, message=args.message, autogenerate=not args.no_autogenerate) + elif args.cmd == "upgrade": + command.upgrade(cfg, args.revision) + elif args.cmd == "downgrade": + command.downgrade(cfg, args.revision) + elif args.cmd == "current": + command.current(cfg, verbose=True) + elif args.cmd == "history": + command.history(cfg, verbose=args.verbose) + else: # pragma: no cover - argparse guards this + parser.error(f"unknown command: {args.cmd}") + + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/src/memu/app/settings.py b/src/memu/app/settings.py index 6c15e7eb..efd06ff1 100644 --- a/src/memu/app/settings.py +++ b/src/memu/app/settings.py @@ -485,7 +485,7 @@ class PatchConfig(BaseModel): class DefaultUserModel(BaseModel): user_id: str | None = None # Agent/session scoping for multi-agent and multi-session memory filtering - # agent_id: str | None = None + agent_id: str | None = None # session_id: str | None = None diff --git a/src/memu/database/postgres/migration.py b/src/memu/database/postgres/migration.py index 3d41bba6..e8a5f335 100644 --- a/src/memu/database/postgres/migration.py +++ b/src/memu/database/postgres/migration.py @@ -5,7 +5,7 @@ from pathlib import Path from typing import Any, Literal -from sqlalchemy import create_engine, inspect, text +from sqlalchemy import create_engine, inspect try: # Optional dependency for Postgres backend from alembic import command @@ -39,50 +39,34 @@ def run_migrations(*, dsn: str, scope_model: type[Any], ddl_mode: DDLMode = "cre Args: dsn: Database connection string scope_model: User scope model for scoped tables - ddl_mode: "create" to create missing tables, "validate" to only check schema + ddl_mode: "create" to apply migrations up to head, "validate" to only check schema + + Alembic is the source of truth for schema: "create" runs ``upgrade head`` + rather than ``metadata.create_all`` so that a fresh database is built from + the committed revisions. The pgvector extension is enabled by the initial + revision, so no separate bootstrap step is required here. """ from memu.database.postgres.schema import get_metadata - metadata = get_metadata(scope_model) - engine = create_engine(dsn) - if ddl_mode == "create": - # Enable pgvector extension if needed (requires superuser or extension already installed) - with engine.connect() as conn: - try: - conn.execute(text("CREATE EXTENSION IF NOT EXISTS vector")) - conn.commit() - logger.info("pgvector extension enabled") - except Exception as e: - # Check if extension already exists - result = conn.execute(text("SELECT 1 FROM pg_extension WHERE extname = 'vector'")).fetchone() - if result: - logger.info("pgvector extension already installed") - else: - msg = ( - "Failed to create pgvector extension. " - "Please run 'CREATE EXTENSION vector;' as a superuser first." - ) - raise RuntimeError(msg) from e - - # Create all tables that don't exist - metadata.create_all(engine) - logger.info("Database tables created/verified") + cfg = make_alembic_config(dsn=dsn, scope_model=scope_model) + command.upgrade(cfg, "head") + logger.info("Database migrated to head") elif ddl_mode == "validate": - # Validate that all expected tables exist - inspector = inspect(engine) - existing_tables = set(inspector.get_table_names()) - expected_tables = set(metadata.tables.keys()) - missing_tables = expected_tables - existing_tables - - if missing_tables: - msg = f"Database schema validation failed. Missing tables: {sorted(missing_tables)}" - raise RuntimeError(msg) - logger.info("Database schema validated successfully") - - # Run any pending Alembic migrations - cfg = make_alembic_config(dsn=dsn, scope_model=scope_model) - command.upgrade(cfg, "head") + metadata = get_metadata(scope_model) + engine = create_engine(dsn) + try: + inspector = inspect(engine) + existing_tables = set(inspector.get_table_names()) + expected_tables = set(metadata.tables.keys()) + missing_tables = expected_tables - existing_tables + + if missing_tables: + msg = f"Database schema validation failed. Missing tables: {sorted(missing_tables)}" + raise RuntimeError(msg) + logger.info("Database schema validated successfully") + finally: + engine.dispose() __all__ = ["DDLMode", "make_alembic_config", "run_migrations"] diff --git a/src/memu/database/postgres/migrations/env.py b/src/memu/database/postgres/migrations/env.py index 950eef1a..9fa75609 100644 --- a/src/memu/database/postgres/migrations/env.py +++ b/src/memu/database/postgres/migrations/env.py @@ -1,6 +1,7 @@ from __future__ import annotations from logging.config import fileConfig +from typing import Any from alembic import context from sqlalchemy import MetaData, engine_from_config, pool @@ -21,6 +22,36 @@ def get_target_metadata() -> MetaData | None: target_metadata: MetaData | None = get_target_metadata() +def include_name(name: str | None, type_: str, parent_names: dict[str, str | None]) -> bool: + """Only manage tables declared in our metadata. + + Keeps autogenerate from emitting drops for unrelated tables when the + reflection target happens to share a schema with other applications. + """ + if type_ == "table" and target_metadata is not None: + return name in target_metadata.tables + return True + + +def render_item(type_: str, obj: Any, autogen_context: Any) -> str | bool: + """Keep generated revisions self-contained (no app-model imports).""" + if type_ == "type": + module = obj.__class__.__module__ + if module.startswith("pgvector."): + autogen_context.imports.add("import pgvector.sqlalchemy") + return f"pgvector.sqlalchemy.{obj!r}" + # TZDateTime is just a timezone-aware DateTime; render it as such so + # the migration does not have to import memu app modules. + if obj.__class__.__name__ == "TZDateTime": + return "sa.DateTime(timezone=True)" + # SQLModel's AutoString (used for scope str columns) is a plain VARCHAR; + # render it as sa.String() for parity with the other string columns and + # to avoid an extra sqlmodel import in the migration. + if module.startswith("sqlmodel.") and obj.__class__.__name__ == "AutoString": + return "sa.String()" + return False + + def run_migrations_offline() -> None: url = config.get_main_option("sqlalchemy.url") context.configure( @@ -29,6 +60,8 @@ def run_migrations_offline() -> None: literal_binds=True, dialect_opts={"paramstyle": "named"}, compare_type=True, + include_name=include_name, + render_item=render_item, ) with context.begin_transaction(): @@ -44,7 +77,13 @@ def run_migrations_online() -> None: ) with connectable.connect() as connection: - context.configure(connection=connection, target_metadata=target_metadata, compare_type=True) + context.configure( + connection=connection, + target_metadata=target_metadata, + compare_type=True, + include_name=include_name, + render_item=render_item, + ) with context.begin_transaction(): context.run_migrations() diff --git a/src/memu/database/postgres/migrations/script.py.mako b/src/memu/database/postgres/migrations/script.py.mako new file mode 100644 index 00000000..8fb2b5ff --- /dev/null +++ b/src/memu/database/postgres/migrations/script.py.mako @@ -0,0 +1,28 @@ +"""${message} + +Revision ID: ${up_revision} +Revises: ${down_revision | comma,n} +Create Date: ${create_date} + +""" +from __future__ import annotations + +from collections.abc import Sequence + +import sqlalchemy as sa +from alembic import op +${imports if imports else ""} + +# revision identifiers, used by Alembic. +revision: str = ${repr(up_revision)} +down_revision: str | Sequence[str] | None = ${repr(down_revision)} +branch_labels: str | Sequence[str] | None = ${repr(branch_labels)} +depends_on: str | Sequence[str] | None = ${repr(depends_on)} + + +def upgrade() -> None: + ${upgrades if upgrades else "pass"} + + +def downgrade() -> None: + ${downgrades if downgrades else "pass"} diff --git a/src/memu/database/postgres/migrations/versions/20260703_0001_initial_schema.py b/src/memu/database/postgres/migrations/versions/20260703_0001_initial_schema.py new file mode 100644 index 00000000..74d968cb --- /dev/null +++ b/src/memu/database/postgres/migrations/versions/20260703_0001_initial_schema.py @@ -0,0 +1,161 @@ +"""initial schema + +Revision ID: 20260703_0001 +Revises: +Create Date: 2026-07-03 19:47:40.690785 + +""" + +from __future__ import annotations + +from collections.abc import Sequence + +import pgvector.sqlalchemy +import sqlalchemy as sa +from alembic import op +from sqlalchemy.dialects import postgresql + +# revision identifiers, used by Alembic. +revision: str = "20260703_0001" +down_revision: str | Sequence[str] | None = None +branch_labels: str | Sequence[str] | None = None +depends_on: str | Sequence[str] | None = None + + +def upgrade() -> None: + # pgvector must exist before any VECTOR column is created. + op.execute("CREATE EXTENSION IF NOT EXISTS vector") + # ### commands auto generated by Alembic - please adjust! ### + op.create_table( + "recall_files", + sa.Column("id", sa.String(), nullable=False), + sa.Column("created_at", sa.DateTime(timezone=True), server_default=sa.text("now()"), nullable=False), + sa.Column("updated_at", sa.DateTime(timezone=True), nullable=False), + sa.Column("name", sa.String(), nullable=False), + sa.Column("track", sa.String(), server_default="memory", nullable=False), + sa.Column("description", sa.Text(), nullable=False), + sa.Column("embedding", pgvector.sqlalchemy.VECTOR(), nullable=True), + sa.Column("content", sa.Text(), nullable=True), + sa.Column("user_id", sa.String(), nullable=True), + sa.Column("agent_id", sa.String(), nullable=True), + sa.PrimaryKeyConstraint("id"), + ) + op.create_index("ix_recall_files__scope", "recall_files", ["user_id", "agent_id"], unique=False) + op.create_index(op.f("ix_recall_files_id"), "recall_files", ["id"], unique=False) + op.create_index(op.f("ix_recall_files_name"), "recall_files", ["name"], unique=False) + op.create_table( + "resources", + sa.Column("id", sa.String(), nullable=False), + sa.Column("created_at", sa.DateTime(timezone=True), server_default=sa.text("now()"), nullable=False), + sa.Column("updated_at", sa.DateTime(timezone=True), nullable=False), + sa.Column("url", sa.String(), nullable=False), + sa.Column("modality", sa.String(), nullable=False), + sa.Column("local_path", sa.String(), nullable=False), + sa.Column("caption", sa.Text(), nullable=True), + sa.Column("embedding", pgvector.sqlalchemy.VECTOR(), nullable=True), + sa.Column("track", sa.String(), nullable=True), + sa.Column("user_id", sa.String(), nullable=True), + sa.Column("agent_id", sa.String(), nullable=True), + sa.PrimaryKeyConstraint("id"), + ) + op.create_index("ix_resources__scope", "resources", ["user_id", "agent_id"], unique=False) + op.create_index(op.f("ix_resources_id"), "resources", ["id"], unique=False) + op.create_table( + "recall_entries", + sa.Column("id", sa.String(), nullable=False), + sa.Column("created_at", sa.DateTime(timezone=True), server_default=sa.text("now()"), nullable=False), + sa.Column("updated_at", sa.DateTime(timezone=True), nullable=False), + sa.Column("resource_id", sa.String(), nullable=True), + sa.Column("memory_type", sa.String(), nullable=False), + sa.Column("summary", sa.Text(), nullable=False), + sa.Column("embedding", pgvector.sqlalchemy.VECTOR(), nullable=True), + sa.Column("happened_at", sa.DateTime(), nullable=True), + sa.Column("extra", postgresql.JSONB(astext_type=sa.Text()), nullable=True), + sa.Column("user_id", sa.String(), nullable=True), + sa.Column("agent_id", sa.String(), nullable=True), + sa.ForeignKeyConstraint(["resource_id"], ["resources.id"], ondelete="CASCADE"), + sa.PrimaryKeyConstraint("id"), + ) + op.create_index("ix_recall_entries__scope", "recall_entries", ["user_id", "agent_id"], unique=False) + op.create_index(op.f("ix_recall_entries_id"), "recall_entries", ["id"], unique=False) + op.create_table( + "recall_file_resources", + sa.Column("id", sa.String(), nullable=False), + sa.Column("created_at", sa.DateTime(timezone=True), server_default=sa.text("now()"), nullable=False), + sa.Column("updated_at", sa.DateTime(timezone=True), nullable=False), + sa.Column("resource_id", sa.String(), nullable=False), + sa.Column("file_id", sa.String(), nullable=False), + sa.Column("user_id", sa.String(), nullable=True), + sa.Column("agent_id", sa.String(), nullable=True), + sa.ForeignKeyConstraint(["file_id"], ["recall_files.id"], ondelete="CASCADE"), + sa.ForeignKeyConstraint(["resource_id"], ["resources.id"], ondelete="CASCADE"), + sa.PrimaryKeyConstraint("id"), + ) + op.create_index( + "idx_recall_file_resources_unique", "recall_file_resources", ["resource_id", "file_id"], unique=True + ) + op.create_index("ix_recall_file_resources__scope", "recall_file_resources", ["user_id", "agent_id"], unique=False) + op.create_index(op.f("ix_recall_file_resources_id"), "recall_file_resources", ["id"], unique=False) + op.create_table( + "recall_file_segments", + sa.Column("id", sa.String(), nullable=False), + sa.Column("created_at", sa.DateTime(timezone=True), server_default=sa.text("now()"), nullable=False), + sa.Column("updated_at", sa.DateTime(timezone=True), nullable=False), + sa.Column("recall_file_id", sa.String(), nullable=False), + sa.Column("track", sa.String(), server_default="memory", nullable=False), + sa.Column("text", sa.Text(), nullable=False), + sa.Column("embedding", pgvector.sqlalchemy.VECTOR(), nullable=True), + sa.Column("user_id", sa.String(), nullable=True), + sa.Column("agent_id", sa.String(), nullable=True), + sa.ForeignKeyConstraint(["recall_file_id"], ["recall_files.id"], ondelete="CASCADE"), + sa.PrimaryKeyConstraint("id"), + ) + op.create_index("ix_recall_file_segments__scope", "recall_file_segments", ["user_id", "agent_id"], unique=False) + op.create_index(op.f("ix_recall_file_segments_id"), "recall_file_segments", ["id"], unique=False) + op.create_index( + op.f("ix_recall_file_segments_recall_file_id"), "recall_file_segments", ["recall_file_id"], unique=False + ) + op.create_table( + "recall_file_entries", + sa.Column("id", sa.String(), nullable=False), + sa.Column("created_at", sa.DateTime(timezone=True), server_default=sa.text("now()"), nullable=False), + sa.Column("updated_at", sa.DateTime(timezone=True), nullable=False), + sa.Column("item_id", sa.String(), nullable=False), + sa.Column("category_id", sa.String(), nullable=False), + sa.Column("user_id", sa.String(), nullable=True), + sa.Column("agent_id", sa.String(), nullable=True), + sa.ForeignKeyConstraint(["category_id"], ["recall_files.id"], ondelete="CASCADE"), + sa.ForeignKeyConstraint(["item_id"], ["recall_entries.id"], ondelete="CASCADE"), + sa.PrimaryKeyConstraint("id"), + ) + op.create_index("idx_recall_file_entries_unique", "recall_file_entries", ["item_id", "category_id"], unique=True) + op.create_index("ix_recall_file_entries__scope", "recall_file_entries", ["user_id", "agent_id"], unique=False) + op.create_index(op.f("ix_recall_file_entries_id"), "recall_file_entries", ["id"], unique=False) + # ### end Alembic commands ### + + +def downgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + op.drop_index(op.f("ix_recall_file_entries_id"), table_name="recall_file_entries") + op.drop_index("ix_recall_file_entries__scope", table_name="recall_file_entries") + op.drop_index("idx_recall_file_entries_unique", table_name="recall_file_entries") + op.drop_table("recall_file_entries") + op.drop_index(op.f("ix_recall_file_segments_recall_file_id"), table_name="recall_file_segments") + op.drop_index(op.f("ix_recall_file_segments_id"), table_name="recall_file_segments") + op.drop_index("ix_recall_file_segments__scope", table_name="recall_file_segments") + op.drop_table("recall_file_segments") + op.drop_index(op.f("ix_recall_file_resources_id"), table_name="recall_file_resources") + op.drop_index("ix_recall_file_resources__scope", table_name="recall_file_resources") + op.drop_index("idx_recall_file_resources_unique", table_name="recall_file_resources") + op.drop_table("recall_file_resources") + op.drop_index(op.f("ix_recall_entries_id"), table_name="recall_entries") + op.drop_index("ix_recall_entries__scope", table_name="recall_entries") + op.drop_table("recall_entries") + op.drop_index(op.f("ix_resources_id"), table_name="resources") + op.drop_index("ix_resources__scope", table_name="resources") + op.drop_table("resources") + op.drop_index(op.f("ix_recall_files_name"), table_name="recall_files") + op.drop_index(op.f("ix_recall_files_id"), table_name="recall_files") + op.drop_index("ix_recall_files__scope", table_name="recall_files") + op.drop_table("recall_files") + # ### end Alembic commands ### diff --git a/src/memu/database/postgres/models.py b/src/memu/database/postgres/models.py index cd95ed8f..57f9b965 100644 --- a/src/memu/database/postgres/models.py +++ b/src/memu/database/postgres/models.py @@ -78,22 +78,22 @@ class RecallFileModel(BaseModelMixin, RecallFile): class RecallFileEntryModel(BaseModelMixin, RecallFileEntry): - item_id: str = Field(sa_column=Column(ForeignKey("memory_items.id", ondelete="CASCADE"), nullable=False)) - category_id: str = Field(sa_column=Column(ForeignKey("memory_categories.id", ondelete="CASCADE"), nullable=False)) + item_id: str = Field(sa_column=Column(ForeignKey("recall_entries.id", ondelete="CASCADE"), nullable=False)) + category_id: str = Field(sa_column=Column(ForeignKey("recall_files.id", ondelete="CASCADE"), nullable=False)) __table_args__ = (Index("idx_recall_file_entries_unique", "item_id", "category_id", unique=True),) class RecallFileResourceModel(BaseModelMixin, RecallFileResource): resource_id: str = Field(sa_column=Column(ForeignKey("resources.id", ondelete="CASCADE"), nullable=False)) - file_id: str = Field(sa_column=Column(ForeignKey("memory_categories.id", ondelete="CASCADE"), nullable=False)) + file_id: str = Field(sa_column=Column(ForeignKey("recall_files.id", ondelete="CASCADE"), nullable=False)) __table_args__ = (Index("idx_recall_file_resources_unique", "resource_id", "file_id", unique=True),) class RecallFileSegmentModel(BaseModelMixin, RecallFileSegment): recall_file_id: str = Field( - sa_column=Column(ForeignKey("memory_categories.id", ondelete="CASCADE"), nullable=False, index=True) + sa_column=Column(ForeignKey("recall_files.id", ondelete="CASCADE"), nullable=False, index=True) ) track: str = Field(default="memory", sa_column=Column(String, nullable=False, server_default="memory")) text: str = Field(sa_column=Column(Text, nullable=False)) @@ -188,12 +188,14 @@ def build_scoped_models( """ resource_model = build_table_model(user_model, ResourceModel, tablename="resources") recall_file_model = build_table_model( - user_model, RecallFileModel, tablename="memory_categories", unique_with_scope=["name", "track"] + user_model, RecallFileModel, tablename="recall_files", unique_with_scope=["name"] ) - recall_entry_model = build_table_model(user_model, RecallEntryModel, tablename="memory_items") - recall_file_entry_model = build_table_model(user_model, RecallFileEntryModel, tablename="category_items") - recall_file_resource_model = build_table_model(user_model, RecallFileResourceModel, tablename="resource_categories") - recall_file_segment_model = build_table_model(user_model, RecallFileSegmentModel, tablename="file_segments") + recall_entry_model = build_table_model(user_model, RecallEntryModel, tablename="recall_entries") + recall_file_entry_model = build_table_model(user_model, RecallFileEntryModel, tablename="recall_file_entries") + recall_file_resource_model = build_table_model( + user_model, RecallFileResourceModel, tablename="recall_file_resources" + ) + recall_file_segment_model = build_table_model(user_model, RecallFileSegmentModel, tablename="recall_file_segments") return ( resource_model, recall_file_model, diff --git a/src/memu/database/postgres/schema.py b/src/memu/database/postgres/schema.py index d970baa9..70d2cee1 100644 --- a/src/memu/database/postgres/schema.py +++ b/src/memu/database/postgres/schema.py @@ -23,6 +23,7 @@ msg = "pgvector is required for Postgres vector support" raise ImportError(msg) from exc +from memu.app.settings import DefaultUserModel from memu.database.postgres.models import ( RecallEntryModel, RecallFileEntryModel, @@ -33,6 +34,12 @@ build_table_model, ) +# Default user scope for the committed Alembic baseline. ``DefaultUserModel`` +# (in ``memu.app.settings``) is the single source of truth for the built-in +# scope columns (``user_id`` / ``agent_id``); this alias keeps the schema and +# the migration generated against it in sync with the app default. +DefaultScope = DefaultUserModel + @dataclass class SQLAModels: @@ -57,7 +64,7 @@ def get_sqlalchemy_models(*, scope_model: type[BaseModel] | None = None) -> SQLA Build (and cache) SQLModel ORM models for Postgres storage. """ require_sqlalchemy() - scope = scope_model or BaseModel + scope = scope_model or DefaultScope cache_key = scope cached = _MODEL_CACHE.get(cache_key) if cached: @@ -74,31 +81,31 @@ def get_sqlalchemy_models(*, scope_model: type[BaseModel] | None = None) -> SQLA recall_file_model = build_table_model( scope, RecallFileModel, - tablename="memory_categories", + tablename="recall_files", metadata=metadata_obj, ) recall_entry_model = build_table_model( scope, RecallEntryModel, - tablename="memory_items", + tablename="recall_entries", metadata=metadata_obj, ) recall_file_entry_model = build_table_model( scope, RecallFileEntryModel, - tablename="category_items", + tablename="recall_file_entries", metadata=metadata_obj, ) recall_file_resource_model = build_table_model( scope, RecallFileResourceModel, - tablename="resource_categories", + tablename="recall_file_resources", metadata=metadata_obj, ) recall_file_segment_model = build_table_model( scope, RecallFileSegmentModel, - tablename="file_segments", + tablename="recall_file_segments", metadata=metadata_obj, ) From 5a3d734d69c5f6739b2f7a12b589846362b5ac1e Mon Sep 17 00:00:00 2001 From: wu Date: Sat, 4 Jul 2026 14:36:28 +0900 Subject: [PATCH 8/9] fix: ruff type --- src/memu/database/postgres/migrations/env.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/src/memu/database/postgres/migrations/env.py b/src/memu/database/postgres/migrations/env.py index 9fa75609..621dcd6a 100644 --- a/src/memu/database/postgres/migrations/env.py +++ b/src/memu/database/postgres/migrations/env.py @@ -1,9 +1,11 @@ from __future__ import annotations from logging.config import fileConfig -from typing import Any +from typing import Any, Literal from alembic import context +from alembic.autogenerate.api import AutogenContext +from alembic.runtime.environment import NameFilterParentNames, NameFilterType from sqlalchemy import MetaData, engine_from_config, pool from memu.database.postgres.schema import get_metadata @@ -22,7 +24,7 @@ def get_target_metadata() -> MetaData | None: target_metadata: MetaData | None = get_target_metadata() -def include_name(name: str | None, type_: str, parent_names: dict[str, str | None]) -> bool: +def include_name(name: str | None, type_: NameFilterType, parent_names: NameFilterParentNames) -> bool: """Only manage tables declared in our metadata. Keeps autogenerate from emitting drops for unrelated tables when the @@ -33,7 +35,7 @@ def include_name(name: str | None, type_: str, parent_names: dict[str, str | Non return True -def render_item(type_: str, obj: Any, autogen_context: Any) -> str | bool: +def render_item(type_: str, obj: Any, autogen_context: AutogenContext) -> str | Literal[False]: """Keep generated revisions self-contained (no app-model imports).""" if type_ == "type": module = obj.__class__.__module__ From 015939303e76b1ab1729c73af0c533284f8d8245 Mon Sep 17 00:00:00 2001 From: wu Date: Sat, 4 Jul 2026 14:40:50 +0900 Subject: [PATCH 9/9] fix: test code --- tests/test_folder_memorize.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_folder_memorize.py b/tests/test_folder_memorize.py index 0bc8bc10..c30fdd05 100644 --- a/tests/test_folder_memorize.py +++ b/tests/test_folder_memorize.py @@ -253,7 +253,7 @@ def _spy_export(database, *, where=None, **kwargs): await service.memorize_workspace(folder=str(src_dir), user=user) # Export ran (scoped to the user) and produced the root index on disk. - assert exported == [user] + assert exported == [service.user_model(**user).model_dump()] assert (out_dir / "INDEX.md").exists()