github-featured/nevamind-ai--memu
2
0
Fork 0
Code Issues Pull Requests Deployments Models Agent Notes Packages Projects Releases Wiki Activity

feat: layered, track-aware workspace memorize & retrieve (#466)

Browse Source 
Co-authored-by: wu <evan299792458@outlook.com>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
nevamind-agent
2026-07-04 14:42:19 +09:00
committed by GitHub
parent 0afbcc550a
commit 094e22f0e1
44 changed files with 3057 additions and 588 deletions
Download Patch File Download Diff File Expand all files Collapse all files
scripts/db.py
+89
View File
@@ -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())
src/memu/app/crud.py
+2
View File
@@ -269,6 +269,8 @@ class CRUDMixin:
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
src/memu/app/memorize.py
+8 -294
View File
@@ -14,7 +14,6 @@ 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,11 +22,6 @@ from memu.prompts.category_summary import (
from memu.prompts.category_summary import (
PROMPT as CATEGORY_SUMMARY_PROMPT,
)
from memu.prompts.memory_fs import (
DESCRIPTIONS_PLACEHOLDER,
EXISTING_PLACEHOLDER,
SKILL_FILE_SYNTHESIS_PROMPT,
)
from memu.prompts.memory_type import (
CUSTOM_PROMPTS as MEMORY_TYPE_CUSTOM_PROMPTS,
)
@@ -107,6 +101,8 @@ class MemorizeMixin:
"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)
@@ -116,182 +112,6 @@ class MemorizeMixin:
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,
)
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,
) -> 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,
}
# 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
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.
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:
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(
@@ -346,6 +166,7 @@ class MemorizeMixin:
"modality",
"user",
"allow_new_categories",
"resource_track",
},
produces={"resources", "entries", "relations", "file_updates"},
capabilities={"db", "vector"},
@@ -383,32 +204,9 @@ class MemorizeMixin:
"category_ids",
"user",
"allow_new_categories",
"resource_track",
}
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.
"""
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
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})
@@ -481,6 +279,7 @@ class MemorizeMixin:
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 +290,7 @@ class MemorizeMixin:
store=store,
embed_client=embed_client,
user=user_scope,
track=track,
)
resources.append(res)
@@ -536,94 +336,6 @@ class MemorizeMixin:
)
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).
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.
"""
if not getattr(self.memory_files_config, "synthesize", False):
return state
content = self._format_skill_source_content(state.get("preprocessed_resources") or [])
if not content:
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)
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)
)
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
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)
@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 _parse_skill_files(self, raw: str) -> list[tuple[str, str, str]]:
"""Parse the skill-synthesis JSON array into ``(name, description, body)`` tuples."""
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 []
skills: list[tuple[str, str, str]] = []
for entry in parsed:
if not isinstance(entry, dict):
continue
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():
continue
description = entry.get("description")
description = description.strip() if isinstance(description, str) else ""
skills.append((name.strip(), description, body.strip()))
return skills
def _memorize_build_response(self, state: WorkflowState, step_context: Any) -> WorkflowState:
ctx = state["ctx"]
store = state["store"]
@@ -687,6 +399,7 @@ class MemorizeMixin:
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 +415,7 @@ class MemorizeMixin:
caption=caption_text,
embedding=caption_embedding,
user_data=dict(user or {}),
track=track,
)
# if caption:
# caption_text = caption.strip()
src/memu/app/memorize_workspace.py
+834
View File
@@ -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
src/memu/app/retrieve.py
+1 -169
View File
@@ -20,14 +20,13 @@ logger = logging.getLogger(__name__)
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,42 +85,6 @@ class RetrieveMixin:
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 across the file/entry/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``.
"""
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_entry": config.entry.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:
@@ -1462,134 +1425,3 @@ class RetrieveMixin:
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.
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.
"""
steps = [
WorkflowStep(
step_id="recall_files",
role="recall_files",
handler=self._ws_recall_files,
requires={"retrieve_file", "query", "store", "where"},
produces={"file_hits", "file_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"},
),
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={
"file_hits",
"file_pool",
"entry_hits",
"entry_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_entry", "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_files(self, state: WorkflowState, step_context: Any) -> WorkflowState:
if not state.get("retrieve_file"):
state["file_hits"] = []
state["file_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 {})
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)
qvec = await self._ws_query_vector(state, step_context)
state["file_hits"] = cosine_topk(
qvec,
[(fid, f.embedding) for fid, f in file_pool.items()],
k=self.retrieve_workspace_config.file.top_k,
)
state["file_pool"] = file_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
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
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"]
where_filters = state.get("where") or {}
resource_pool = store.resource_repo.list_resources(where_filters)
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
)
state["resource_pool"] = resource_pool
return state
def _ws_build_response(self, state: WorkflowState, _: Any) -> WorkflowState:
state["response"] = {
"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
src/memu/app/retrieve_workspace.py
+276
View File
@@ -0,0 +1,276 @@
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", "file_resource_urls"},
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",
"file_resource_urls",
"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
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"] = []
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:
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": files,
"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
src/memu/app/service.py
+4 -2
View File
@@ -10,8 +10,10 @@ from pydantic import BaseModel
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,
*,
@@ -343,7 +345,7 @@ class MemoryService(MemorizeMixin, RetrieveMixin, CRUDMixin):
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",
src/memu/app/settings.py
+1 -1
View File
@@ -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
src/memu/database/__init__.py
+14 -1
View File
@@ -6,9 +6,18 @@ from memu.database.interfaces import (
RecallEntryRecord,
RecallFileEntryRecord,
RecallFileRecord,
RecallFileResourceRecord,
RecallFileSegmentRecord,
ResourceRecord,
)
from memu.database.repositories import RecallEntryRepo, RecallFileEntryRepo, RecallFileRepo, ResourceRepo
from memu.database.repositories import (
RecallEntryRepo,
RecallFileEntryRepo,
RecallFileRepo,
RecallFileResourceRepo,
RecallFileSegmentRepo,
ResourceRepo,
)
__all__ = [
"Database",
@@ -18,6 +27,10 @@ __all__ = [
"RecallFileEntryRepo",
"RecallFileRecord",
"RecallFileRepo",
"RecallFileResourceRecord",
"RecallFileResourceRepo",
"RecallFileSegmentRecord",
"RecallFileSegmentRepo",
"ResourceRecord",
"ResourceRepo",
"build_database",
src/memu/database/inmemory/__init__.py
+10 -1
View File
@@ -12,13 +12,22 @@ 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,
recall_file_segment_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,
recall_file_segment_model=recall_file_segment_model,
)
src/memu/database/inmemory/models.py
+24 -1
View File
@@ -6,6 +6,8 @@ from memu.database.models import (
RecallEntry,
RecallFile,
RecallFileEntry,
RecallFileResource,
RecallFileSegment,
Resource,
merge_scope_model,
)
@@ -27,6 +29,14 @@ class InMemoryFileEntry(RecallFileEntry):
"""Concrete in-memory relation model."""
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[
@@ -34,6 +44,8 @@ def build_inmemory_models(
type[InMemoryRecallFile],
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.
@@ -42,11 +54,22 @@ 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")
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",
src/memu/database/inmemory/repo.py
+26 -1
View File
@@ -7,13 +7,22 @@ from pydantic import BaseModel
from memu.database.inmemory.models import build_inmemory_models
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, Resource
from memu.database.models import (
RecallEntry,
RecallFile,
RecallFileEntry,
RecallFileResource,
RecallFileSegment,
Resource,
)
from memu.database.repositories import RecallFileRepo, ResourceRepo
@@ -26,6 +35,8 @@ class InMemoryStore(Database):
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,
recall_file_segment_model: type[Any] | None = None,
state: InMemoryState | None = None,
) -> None:
self.scope_model = scope_model or BaseModel
@@ -34,6 +45,8 @@ class InMemoryStore(Database):
default_recall_file_model,
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()
@@ -41,11 +54,17 @@ class InMemoryStore(Database):
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
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
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
)
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(
@@ -55,6 +74,12 @@ class InMemoryStore(Database):
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
)
self.recall_file_segment_repo = InMemoryFileSegmentRepository(
state=self.state, recall_file_segment_model=recall_file_segment_model
)
def close(self) -> None:
return None
src/memu/database/inmemory/repositories/__init__.py
+12
View File
@@ -7,15 +7,27 @@ from memu.database.inmemory.repositories.recall_file_repo import (
InMemoryRecallFileRepository,
RecallFileRepo,
)
from memu.database.inmemory.repositories.recall_file_resource_repo import (
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",
"RecallEntryRepo",
"RecallFileEntryRepo",
"RecallFileRepo",
"RecallFileResourceRepo",
"RecallFileSegmentRepo",
"ResourceRepo",
]
src/memu/database/inmemory/repositories/recall_file_resource_repo.py
+67
View File
@@ -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.file_id == cat_id:
return rel
rel = self.recall_file_resource_model(
id=str(uuid.uuid4()), resource_id=resource_id, file_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.file_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"]
src/memu/database/inmemory/repositories/recall_file_segment_repo.py
+71
View File
@@ -0,0 +1,71 @@
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],
track: str = "memory",
) -> RecallFileSegment:
seg = self.recall_file_segment_model(
id=str(uuid.uuid4()),
recall_file_id=recall_file_id,
track=track,
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"]
src/memu/database/inmemory/repositories/resource_repo.py
+2
View File
@@ -44,6 +44,7 @@ class InMemoryResourceRepository(ResourceRepoProtocol):
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 @@ class InMemoryResourceRepository(ResourceRepoProtocol):
local_path=local_path,
caption=caption,
embedding=embedding,
track=track,
**user_data,
)
self.resources[rid] = res
src/memu/database/interfaces.py
+16 -1
View File
@@ -5,8 +5,17 @@ from typing import Protocol, runtime_checkable
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 RecallFileSegment as RecallFileSegmentRecord
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,
RecallFileSegmentRepo,
ResourceRepo,
)
@runtime_checkable
@@ -17,11 +26,15 @@ class Database(Protocol):
recall_file_repo: RecallFileRepo
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: ...
@@ -31,5 +44,7 @@ __all__ = [
"RecallEntryRecord",
"RecallFileEntryRecord",
"RecallFileRecord",
"RecallFileResourceRecord",
"RecallFileSegmentRecord",
"ResourceRecord",
]
src/memu/database/models.py
+48 -2
View File
@@ -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):
@@ -108,6 +111,31 @@ class RecallFileEntry(BaseRecord):
category_id: str
class RecallFileResource(BaseRecord):
resource_id: str
file_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.
``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
def merge_scope_model[TBaseRecord: BaseRecord](
user_model: type[BaseModel], core_model: type[TBaseRecord], *, name_suffix: str
) -> type[TBaseRecord]:
@@ -126,7 +154,14 @@ 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],
type[RecallFileSegment],
]:
"""
Build scoped interface models (Pydantic) that inherit from the base record models and user scope.
"""
@@ -134,7 +169,16 @@ 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")
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,
)
__all__ = [
@@ -143,6 +187,8 @@ __all__ = [
"RecallEntry",
"RecallFile",
"RecallFileEntry",
"RecallFileResource",
"RecallFileSegment",
"Resource",
"ToolCallResult",
"build_scoped_models",
src/memu/database/postgres/migration.py
+23 -39
View File
@@ -5,7 +5,7 @@ import re
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
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")
# Run any pending Alembic migrations
cfg = make_alembic_config(dsn=dsn, scope_model=scope_model)
command.upgrade(cfg, "head")
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"]
src/memu/database/postgres/migrations/env.py
+42 -1
View File
@@ -1,8 +1,11 @@
from __future__ import annotations
from logging.config import fileConfig
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
@@ -21,6 +24,36 @@ def get_target_metadata() -> MetaData | None:
target_metadata: MetaData | None = get_target_metadata()
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
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: AutogenContext) -> str | Literal[False]:
"""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 +62,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 +79,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()
src/memu/database/postgres/migrations/script.py.mako
+28
View File
@@ -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"}
src/memu/database/postgres/migrations/versions/20260703_0001_initial_schema.py
+161
View File
@@ -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 ###
src/memu/database/postgres/models.py
+47 -9
View File
@@ -17,7 +17,15 @@ from sqlalchemy import ForeignKey, MetaData, String, Text
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,
RecallFileSegment,
Resource,
)
class TZDateTime(DateTime):
@@ -49,6 +57,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):
@@ -69,12 +78,28 @@ 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("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("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))
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 [], {}
@@ -157,17 +182,28 @@ 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], 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(
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="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,
recall_entry_model,
recall_file_entry_model,
recall_file_resource_model,
recall_file_segment_model,
)
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
__all__ = [
@@ -175,6 +211,8 @@ __all__ = [
"RecallEntryModel",
"RecallFileEntryModel",
"RecallFileModel",
"RecallFileResourceModel",
"RecallFileSegmentModel",
"ResourceModel",
"build_scoped_models",
"build_table_model",
src/memu/database/postgres/postgres.py
+44 -2
View File
@@ -6,15 +6,31 @@ from typing import Any
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,
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
from memu.database.repositories import RecallEntryRepo, RecallFileEntryRepo, RecallFileRepo, ResourceRepo
from memu.database.repositories import (
RecallEntryRepo,
RecallFileEntryRepo,
RecallFileRepo,
RecallFileResourceRepo,
RecallFileSegmentRepo,
ResourceRepo,
)
from memu.database.state import DatabaseState
logger = logging.getLogger(__name__)
@@ -25,10 +41,14 @@ class PostgresStore(Database):
recall_file_repo: RecallFileRepo
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,
@@ -42,6 +62,8 @@ class PostgresStore(Database):
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,
recall_file_segment_model: type[Any] | None = None,
sqla_models: SQLAModels | None = None,
) -> None:
require_sqlalchemy()
@@ -60,6 +82,8 @@ class PostgresStore(Database):
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
recall_file_segment_model = recall_file_segment_model or self._sqla_models.RecallFileSegment
self.resource_repo = PostgresResourceRepo(
state=self._state,
@@ -90,11 +114,27 @@ class PostgresStore(Database):
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.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()
@@ -106,3 +146,5 @@ class PostgresStore(Database):
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()
self.recall_file_segment_repo.load_existing()
src/memu/database/postgres/repositories/__init__.py
+4
View File
@@ -1,11 +1,15 @@
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
__all__ = [
"PostgresRecallEntryRepo",
"PostgresRecallFileEntryRepo",
"PostgresRecallFileRepo",
"PostgresRecallFileResourceRepo",
"PostgresRecallFileSegmentRepo",
"PostgresResourceRepo",
]
src/memu/database/postgres/repositories/recall_file_resource_repo.py
+149
View File
@@ -0,0 +1,149 @@
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.file_id == cat_id:
return rel
now = self._now()
new_rel = self._recall_file_resource_model(
resource_id=resource_id,
file_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.file_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.file_id == cat_id,
)
)
session.commit()
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,
file_id=row.file_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"]
src/memu/database/postgres/repositories/recall_file_segment_repo.py
+133
View File
@@ -0,0 +1,133 @@
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,
track=row.track,
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],
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,
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"]
src/memu/database/postgres/repositories/resource_repo.py
+2
View File
@@ -81,6 +81,7 @@ class PostgresResourceRepo(PostgresRepoBase, ResourceRepo):
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 @@ class PostgresResourceRepo(PostgresRepoBase, ResourceRepo):
local_path=local_path,
caption=caption,
embedding=self._prepare_embedding(embedding),
track=track,
**user_data,
created_at=self._now(),
updated_at=self._now(),
src/memu/database/postgres/schema.py
+29 -4
View File
@@ -23,14 +23,23 @@ except ImportError as exc:
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,
RecallFileModel,
RecallFileResourceModel,
RecallFileSegmentModel,
ResourceModel,
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:
@@ -39,6 +48,8 @@ class SQLAModels:
RecallFile: type[Any]
RecallEntry: type[Any]
RecallFileEntry: type[Any]
RecallFileResource: type[Any]
RecallFileSegment: type[Any]
_MODEL_CACHE: dict[type[Any], SQLAModels] = {}
@@ -53,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:
@@ -70,19 +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="recall_file_resources",
metadata=metadata_obj,
)
recall_file_segment_model = build_table_model(
scope,
RecallFileSegmentModel,
tablename="recall_file_segments",
metadata=metadata_obj,
)
@@ -96,6 +119,8 @@ def get_sqlalchemy_models(*, scope_model: type[BaseModel] | None = None) -> SQLA
RecallFile=recall_file_model,
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
src/memu/database/repositories/__init__.py
+10 -1
View File
@@ -1,6 +1,15 @@
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.recall_file_segment import RecallFileSegmentRepo
from memu.database.repositories.resource import ResourceRepo
__all__ = ["RecallEntryRepo", "RecallFileEntryRepo", "RecallFileRepo", "ResourceRepo"]
__all__ = [
"RecallEntryRepo",
"RecallFileEntryRepo",
"RecallFileRepo",
"RecallFileResourceRepo",
"RecallFileSegmentRepo",
"ResourceRepo",
]
src/memu/database/repositories/recall_file_resource.py
+33
View File
@@ -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: ...
src/memu/database/repositories/recall_file_segment.py
+41
View File
@@ -0,0 +1,41 @@
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],
track: str = "memory",
) -> 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: ...
src/memu/database/repositories/resource.py
+1
View File
@@ -27,6 +27,7 @@ class ResourceRepo(Protocol):
caption: str | None,
embedding: list[float] | None,
user_data: dict[str, Any],
track: str | None = None,
) -> Resource: ...
def vector_search_resources(
src/memu/database/sqlite/models.py
+32 -1
View File
@@ -11,7 +11,15 @@ from pydantic import BaseModel
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,
RecallFileSegment,
Resource,
)
class TZDateTime(DateTime):
@@ -51,6 +59,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):
@@ -87,6 +96,26 @@ 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))
file_id: str = Field(sa_column=Column(String, nullable=False))
__table_args__ = (Index("idx_sqlite_recall_file_resources_unique", "resource_id", "file_id", unique=True),)
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).
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:
@@ -175,6 +204,8 @@ __all__ = [
"SQLiteRecallEntryModel",
"SQLiteRecallFileEntryModel",
"SQLiteRecallFileModel",
"SQLiteRecallFileResourceModel",
"SQLiteRecallFileSegmentModel",
"SQLiteResourceModel",
"build_sqlite_table_model",
]
src/memu/database/sqlite/repositories/__init__.py
+4
View File
@@ -4,12 +4,16 @@ from memu.database.sqlite.repositories.base import SQLiteRepoBase
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
__all__ = [
"SQLiteRecallEntryRepo",
"SQLiteRecallFileEntryRepo",
"SQLiteRecallFileRepo",
"SQLiteRecallFileResourceRepo",
"SQLiteRecallFileSegmentRepo",
"SQLiteRepoBase",
"SQLiteResourceRepo",
]
src/memu/database/sqlite/repositories/recall_file_resource_repo.py
+215
View File
@@ -0,0 +1,215 @@
"""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,
file_id=row.file_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, file_id: str, user_data: dict[str, Any]) -> RecallFileResource:
"""Create a link between a resource and a category.
Args:
resource_id: Resource ID.
file_id: File ID.
user_data: User scope data.
Returns:
Created RecallFileResource relation.
"""
# Check if relation already exists
where: dict[str, Any] = {
"resource_id": resource_id,
"file_id": file_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,
file_id=existing.file_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,
file_id=file_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,
file_id=row.file_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, file_id: str) -> None:
"""Remove a link between a resource and a category.
Args:
resource_id: Resource 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.file_id == file_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.file_id == file_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"]
src/memu/database/sqlite/repositories/recall_file_segment_repo.py
+142
View File
@@ -0,0 +1,142 @@
"""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,
track=row.track,
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],
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,
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"]
src/memu/database/sqlite/repositories/resource_repo.py
+6
View File
@@ -78,6 +78,7 @@ class SQLiteResourceRepo(SQLiteRepoBase, ResourceRepo):
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 @@ class SQLiteResourceRepo(SQLiteRepoBase, ResourceRepo):
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 @@ class SQLiteResourceRepo(SQLiteRepoBase, ResourceRepo):
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 @@ class SQLiteResourceRepo(SQLiteRepoBase, ResourceRepo):
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 @@ class SQLiteResourceRepo(SQLiteRepoBase, ResourceRepo):
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 @@ class SQLiteResourceRepo(SQLiteRepoBase, ResourceRepo):
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,
src/memu/database/sqlite/schema.py
+18
View File
@@ -13,6 +13,8 @@ from memu.database.sqlite.models import (
SQLiteRecallEntryModel,
SQLiteRecallFileEntryModel,
SQLiteRecallFileModel,
SQLiteRecallFileResourceModel,
SQLiteRecallFileSegmentModel,
SQLiteResourceModel,
build_sqlite_table_model,
)
@@ -27,6 +29,8 @@ class SQLiteSQLAModels:
RecallFile: type[Any]
RecallEntry: type[Any]
RecallFileEntry: type[Any]
RecallFileResource: type[Any]
RecallFileSegment: type[Any]
_MODEL_CACHE: dict[type[Any], SQLiteSQLAModels] = {}
@@ -75,6 +79,18 @@ 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,
)
recall_file_segment_model = build_sqlite_table_model(
scope,
SQLiteRecallFileSegmentModel,
tablename="memu_file_segments",
metadata=metadata_obj,
)
class SQLiteBase(SQLModel):
__abstract__ = True
@@ -86,6 +102,8 @@ def get_sqlite_sqlalchemy_models(*, scope_model: type[BaseModel] | None = None)
RecallFile=recall_file_model,
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
src/memu/database/sqlite/sqlite.py
+44 -2
View File
@@ -9,11 +9,27 @@ from pydantic import BaseModel
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,
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
@@ -44,10 +60,14 @@ class SQLiteStore(Database):
recall_file_repo: RecallFileRepo
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,
@@ -58,6 +78,8 @@ class SQLiteStore(Database):
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,
recall_file_segment_model: type[Any] | None = None,
sqla_models: SQLiteSQLAModels | None = None,
) -> None:
"""Initialize SQLite database store.
@@ -86,6 +108,8 @@ class SQLiteStore(Database):
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
recall_file_segment_model = recall_file_segment_model or self._sqla_models.RecallFileSegment
# Initialize repositories
self.resource_repo = SQLiteResourceRepo(
@@ -116,12 +140,28 @@ class SQLiteStore(Database):
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,
)
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
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
def _create_tables(self) -> None:
"""Create SQLite tables if they don't exist."""
@@ -140,6 +180,8 @@ class SQLiteStore(Database):
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()
self.recall_file_segment_repo.load_existing()
__all__ = ["SQLiteStore"]
src/memu/database/state.py
+10 -1
View File
@@ -2,7 +2,14 @@ from __future__ import annotations
from dataclasses import dataclass, field
from memu.database.models import RecallEntry, RecallFile, RecallFileEntry, Resource
from memu.database.models import (
RecallEntry,
RecallFile,
RecallFileEntry,
RecallFileResource,
RecallFileSegment,
Resource,
)
@dataclass
@@ -11,6 +18,8 @@ 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)
segments: list[RecallFileSegment] = field(default_factory=list)
__all__ = ["DatabaseState"]
src/memu/prompts/memory_fs/__init__.py
+133 -6
View File
@@ -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 @@ NEW source content:
__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": "<exact existing file name>"}
{"op": "create", "name": "<concise Title Case file 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": "<exact existing skill 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",
]
tests/test_folder_memorize.py
+7 -4
View File
@@ -159,7 +159,7 @@ async def test_memorize_workspace_sync_add_modify_delete(tmp_path: Path, monkeyp
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 test_memorize_workspace_sync_add_modify_delete(tmp_path: Path, monkeyp
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 test_memorize_workspace_exports_when_enabled(tmp_path: Path, monkeypat
caption="cap",
embedding=None,
user_data=dict(user_scope or {}),
track=track,
)
return {"resources": [res], "response": {"items": []}}
@@ -251,7 +253,7 @@ async def test_memorize_workspace_exports_when_enabled(tmp_path: Path, monkeypat
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()
@@ -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 test_memorize_workspace_export_failure_does_not_fail_sync(tmp_path: Pa
caption="cap",
embedding=None,
user_data=dict(user_scope or {}),
track=track,
)
return {"resources": [res], "response": {"items": []}}
tests/test_skill_track.py
+194 -45
View File
@@ -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 _FakeSkillClient:
"""Stand-in client exposing both chat (skill JSON) and embed (fixed vector)."""
class _FakeClient:
"""Fake LLM/embed client that answers the route step and the synthesis step.
def __init__(self, payload: str = _SKILLS_JSON) -> None:
self._payload = payload
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, 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,201 @@ 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].file_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"}
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)
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_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_skill_step(service, _FakeSkillClient(), state)
result = await _run_synthesize(service, _FakeClient(), track="workspace", 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"}) == {}
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_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]
# 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:
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."]
# 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:
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
user = {"user_id": "u1"}
result = await _run_synthesize(
service,
_FakeClient(route='[{"op": "update", "name": "does-not-exist"}]'),
track="skill",
user=user,
)
assert result["files"] == []
assert store.recall_file_repo.list_categories(where={"user_id": "u1", "track": "skill"}) == {}