From 6aa13e8f85441ecdb068a3d951955d0586a36b33 Mon Sep 17 00:00:00 2001 From: Ashpreet Date: Sun, 26 Jul 2026 10:02:20 -0700 Subject: [PATCH] feat: revamp entity memory for the second brain (#9177) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary **The second brain becomes buildable.** This PR makes entity memory real — a brain that accumulates the people, projects and systems around the user for six months without fragmenting, going stale, or forgetting how to answer — and it ships the second-brain agent that proves it, live over REST and MCP. What a second brain can do after this PR that it could not before: - **Remember by name.** `remember_about("Sarah Chen", ...)` upserts; the store slugifies ids, resolves name/alias/type drift, and merges. The model never invents an id and never touches a fact id. - **Correct itself.** Stating "radar shipped" retires "radar is blocked on review" (fact supersession: one judged model call in the write path, threshold-gated, nothing deleted). Facts render with as-of dates, so July's truth outranks March's. - **Recall what the turn is about.** `recall(message)` injects the top-k relevant entities plus an always-present entity directory — "not in the directory" is a grounded no. The current message finally reaches the stores (the single dropped parameter that made entity memory write-only). - **Search its whole self.** A real `search_learnings` query surface on sqlite/postgres/async_postgres (JSONB `CAST` + `ILIKE`, slug-boundary crossing, errors raised — never a silent `[]`), replacing the 15-most-recent-rows scan that made the brain degrade as it grew. - **Point at its notes.** `remember_about(..., note="notes/radar.md")` writes the index-to-content stitch the one-claim-one-home rule depends on; links are reciprocal and recall shows linked entities by name. ### The tool surface, before and after | | Before | After | |---|---|---| | Entity tools | 8 (two dead — `update_fact`/`delete_fact` needed ids never shown to the model) | **4**: `remember_about` (incl. `note=`), `link_entities`, `search_entities` (query optional = browse), `forget` (fact retire / entity archive) | | Capture | ALWAYS extraction pass (append-only, blind to existing state, ≥1 LLM call/run) | **AGENTIC only** — recording is the agent's job | | Recall | keyed lookup that returned `None` in practice | `recall(message)` → top-k + directory; `get()` keeps the keyed path | | FileSystem default | 9 tools incl. `check_lines`/`delete_file`, record-set instructions | **the notes seven**; `allow_delete=True` / `include_tools` restore the rest; instructions are the notes contract | ### Breaking changes are confined to `EntityMemoryStore` and the `FileSystem` defaults That is deliberate and safe. Entity memory has no real users: every cookbook TEST_LOG entry for it was PENDING, and on the default path it never contributed anything to an agent — recall could not be keyed, so it was write-only. Two of its eight tools were unusable even in principle. `FileSystem` is weeks old with one behavior change already shipped in 2.8.3; flipping its default now costs one release-note line. **Migration for the old entity tools:** `create_entity`/`update_entity`/`add_fact`/`add_event` → one `remember_about` call; `add_relationship` → `link_entities`; `update_fact`/`delete_fact` → state the new fact (supersession retires the old one) or `forget(entity, fact=...)`; entity deletion → `forget(entity)` archives (reversible), `store.delete()` hard-deletes. **FileSystem:** `tools(allow_delete=True)` restores `delete_file`; `include_tools` selects from the whole surface (incl. `check_lines`); `FULL_TOOLS` stays exported. Also in this PR: the LearningMachine guidance/data split (`instructions()` on stores and the machine; `build_context()` is data-only; manual door == automatic door, tested), the manual door with `capture_hook()`, `RunContext` threading with signature-filtered dispatch (narrow third-party stores keep working), a lossless config round-trip (an AGENTIC machine no longer reloads as ALWAYS with zero tools), the `decision_log` contentless-row fix, and two visibility guards: a run with `user_id=None` while per-user stores are configured, and a hand-placed machine with no `model=` (capture is a model call, so without one the tools return "No model provided" and entity memory retires nothing). Both latch and fire once. The `enable_user_memories` deprecation and the `enable_agentic_memory` tool-name collision warning were built, measured and removed: the first fired on every `Agent.load()` and `deep_copy()` for a parameter the author never wrote, the second on every run. The collision is a comment on the field instead. Spec: `specs/agno/2.8.4/README.md`. Build report (commit map, §8 test mapping, seven adversarial review passes, deviations, live-proof evidence): `specs/agno/2.8.4/notes/BUILD_REPORT.md`. ## Type of change - [x] New feature - [x] Breaking change - [x] Bug fix --- ## Checklist - [x] Code complies with style guidelines - [x] Ran format/validation scripts (`./scripts/format.sh` and `./scripts/validate.sh`) - [x] Self-review completed - [x] Documentation updated (comments, docstrings) - [x] Examples and guides: Relevant cookbook examples have been included or updated (if applicable) - [x] Tested in clean environment - [x] Tests added/updated (if applicable) ### Duplicate and AI-Generated PR Check - [x] I have searched existing [open pull requests](../../pulls) and confirmed that no other PR already addresses this issue - [ ] If a similar PR exists, I have explained below why this PR is a better approach - [x] Check if this PR was entirely AI-generated (by Copilot, Claude Code, Cursor, etc.) --- ## Additional Notes **Release notes (for the 2.8.4 changelog):** - Entity memory revamped: AGENTIC-only capture and a four-tool surface (`remember_about`, `link_entities`, `search_entities`, `forget`). Anyone on the old eight tools migrates per the mapping above. - FileSystem defaults flipped: `tools()` now returns the notes seven; `allow_delete=True` restores `delete_file`, `include_tools` restores `check_lines` or any of the full surface; the default `instructions()` text was rewritten around the notes contract. - The fact-supersession judge prompt wording changed from 2.8.3; operators who overrode the three prompt strings on `EntityMemoryConfig` should re-check their overrides against the new defaults. - **Live proof:** the rebuilt second brain ran 8 REST turns (db verified between turns: note pointer, reciprocal edges, supersession's `superseded_by` id, archive flag, a private confidence landing in per-user memory) and 4 fresh-session `/mcp` turns via an in-process fastmcp client (pinned `user_id` carries one brain across channels). `cookbook/examples/second_brain/TEST_LOG.md` records model behavior honestly, including routing misses. - **Out of scope, deliberately** (2.8.5+): the recap engine, tier-one Memory fixes, `BaseLearningStore`, callable prompt hooks / render hook / typed custom-schema tools, `search_summaries`, the remaining fifteen `search_learnings` backend implementations, semantic retrieval over notes, relationship traversal beyond one-hop names, `replace_text`, `expected_version` on `write_file`, `learning="agentic"`. - **Identity over `/mcp`, measured:** `run_agent` advertises an optional `user_id`, and a host that fills it overrides a pinned `Agent(user_id=...)` - that run's profile and user memory go to the name the host sent. Entities are global, so the shared half of the brain is unaffected. The second-brain cookbook says so and points at auth; a code-level answer is 2.8.5-shaped. - **No version bump** — the release PR is cut separately. - Search folds separators and case on both sides: `multi-tenant` reaches a stored `multi tenant` and the reverse, and a wrong-case non-ASCII query (`CAFÉ`, `ΟΣ`) reaches its row on sqlite too - the JSON-escape prefilter carries a wildcard per escaped character and the value-scoped verification casefolds. 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Co-authored-by: Claude Fable 5 --- .../12_learning/03_team_entity_memory.py | 4 +- .../12_learning/10_team_agentic_learning.py | 2 +- .../08_learning/01_basics/5_entity_memory.py | 63 + .../01_basics/5a_entity_memory_always.py | 81 - .../01_basics/5b_entity_memory_agentic.py | 86 - .../01_basics/6_extraction_limits.py | 18 - cookbook/08_learning/01_basics/README.md | 3 +- cookbook/08_learning/01_basics/TEST_LOG.md | 23 +- .../04_entity_memory/01_facts_and_events.py | 103 - .../04_entity_memory/01_the_four_tools.py | 86 + .../02_entity_relationships.py | 104 - .../04_entity_memory/02_links_and_forget.py | 86 + .../08_learning/04_entity_memory/README.md | 23 +- .../08_learning/04_entity_memory/TEST_LOG.md | 36 +- .../07_patterns/personal_assistant.py | 3 +- .../08_learning/07_patterns/support_agent.py | 3 +- .../02_decision_log_always.py | 66 - .../09_decision_logs/02_record_outcomes.py | 62 + .../08_learning/09_decision_logs/README.md | 2 +- .../08_learning/09_decision_logs/TEST_LOG.md | 22 +- cookbook/08_learning/11_composition/README.md | 37 + .../08_learning/11_composition/TEST_LOG.md | 70 + .../11_composition/always_capture.py | 51 + cookbook/08_learning/11_composition/basic.py | 67 + .../11_composition/context_block.py | 41 + .../11_composition/with_filesystem.py | 50 + cookbook/08_learning/README.md | 35 +- .../13_filesystem/02_durable_records/basic.py | 22 +- .../02_durable_records/radar_news_delta.py | 11 +- .../13_filesystem/04_namespaces/README.md | 2 +- .../13_filesystem/04_namespaces/TEST_LOG.md | 10 + .../04_namespaces/shared_namespace.py | 2 +- .../05_operations/quota_recovery.py | 5 +- cookbook/examples/second_brain/TEST_LOG.md | 123 +- .../examples/second_brain/second_brain.py | 82 +- cookbook/examples/second_brain/test.py | 23 +- libs/agno/agno/agent/_init.py | 6 + libs/agno/agno/agent/_managers.py | 15 + libs/agno/agno/agent/_messages.py | 61 +- libs/agno/agno/agent/_run.py | 4 + libs/agno/agno/agent/_tools.py | 2 + libs/agno/agno/agent/agent.py | 9 +- libs/agno/agno/db/base.py | 92 + libs/agno/agno/db/postgres/async_postgres.py | 69 +- libs/agno/agno/db/postgres/postgres.py | 68 +- libs/agno/agno/db/sqlite/sqlite.py | 64 + libs/agno/agno/db/utils.py | 62 + libs/agno/agno/environments/runner.py | 6 + libs/agno/agno/fs/db.py | 13 +- libs/agno/agno/fs/fs.py | 72 +- libs/agno/agno/fs/toolkit.py | 153 +- libs/agno/agno/learn/config.py | 145 +- libs/agno/agno/learn/machine.py | 458 +- libs/agno/agno/learn/schemas.py | 241 +- libs/agno/agno/learn/stores/decision_log.py | 451 +- libs/agno/agno/learn/stores/entity_memory.py | 5523 +++++++++-------- .../agno/learn/stores/learned_knowledge.py | 45 +- libs/agno/agno/learn/stores/protocol.py | 16 +- .../agno/agno/learn/stores/session_context.py | 4 + libs/agno/agno/learn/stores/user_memory.py | 123 +- libs/agno/agno/learn/stores/user_profile.py | 62 +- libs/agno/agno/learn/utils.py | 83 + libs/agno/agno/team/_init.py | 6 + libs/agno/agno/team/_messages.py | 54 +- libs/agno/agno/team/team.py | 9 +- .../db/postgres/test_search_learnings.py | 107 + .../integration/fs/test_toolkit_over_db.py | 4 +- .../tests/unit/db/test_search_learnings.py | 153 + libs/agno/tests/unit/fs/test_toolkit.py | 149 +- .../unit/learn/test_decision_log_search.py | 265 + .../learn/test_deprecations_and_guards.py | 245 + .../unit/learn/test_entity_memory_store.py | 1405 +++++ .../unit/learn/test_entity_supersession.py | 324 + .../tests/unit/learn/test_guidance_split.py | 480 ++ .../unit/learn/test_machine_roundtrip.py | 197 + .../unit/learn/test_max_updates_per_run.py | 11 +- .../tests/unit/learn/test_relevance_recall.py | 484 ++ .../tests/unit/team/test_team_learning.py | 86 + 78 files changed, 9457 insertions(+), 3876 deletions(-) create mode 100644 cookbook/08_learning/01_basics/5_entity_memory.py delete mode 100644 cookbook/08_learning/01_basics/5a_entity_memory_always.py delete mode 100644 cookbook/08_learning/01_basics/5b_entity_memory_agentic.py delete mode 100644 cookbook/08_learning/04_entity_memory/01_facts_and_events.py create mode 100644 cookbook/08_learning/04_entity_memory/01_the_four_tools.py delete mode 100644 cookbook/08_learning/04_entity_memory/02_entity_relationships.py create mode 100644 cookbook/08_learning/04_entity_memory/02_links_and_forget.py delete mode 100644 cookbook/08_learning/09_decision_logs/02_decision_log_always.py create mode 100644 cookbook/08_learning/09_decision_logs/02_record_outcomes.py create mode 100644 cookbook/08_learning/11_composition/README.md create mode 100644 cookbook/08_learning/11_composition/TEST_LOG.md create mode 100644 cookbook/08_learning/11_composition/always_capture.py create mode 100644 cookbook/08_learning/11_composition/basic.py create mode 100644 cookbook/08_learning/11_composition/context_block.py create mode 100644 cookbook/08_learning/11_composition/with_filesystem.py create mode 100644 libs/agno/tests/integration/db/postgres/test_search_learnings.py create mode 100644 libs/agno/tests/unit/db/test_search_learnings.py create mode 100644 libs/agno/tests/unit/learn/test_decision_log_search.py create mode 100644 libs/agno/tests/unit/learn/test_deprecations_and_guards.py create mode 100644 libs/agno/tests/unit/learn/test_entity_memory_store.py create mode 100644 libs/agno/tests/unit/learn/test_entity_supersession.py create mode 100644 libs/agno/tests/unit/learn/test_guidance_split.py create mode 100644 libs/agno/tests/unit/learn/test_machine_roundtrip.py create mode 100644 libs/agno/tests/unit/learn/test_relevance_recall.py diff --git a/cookbook/03_teams/12_learning/03_team_entity_memory.py b/cookbook/03_teams/12_learning/03_team_entity_memory.py index 4d4f5c42cd..1dc656da93 100644 --- a/cookbook/03_teams/12_learning/03_team_entity_memory.py +++ b/cookbook/03_teams/12_learning/03_team_entity_memory.py @@ -55,9 +55,7 @@ team = Team( user_profile=UserProfileConfig( mode=LearningMode.ALWAYS, ), - entity_memory=EntityMemoryConfig( - mode=LearningMode.ALWAYS, - ), + entity_memory=EntityMemoryConfig(), # AGENTIC-only: the agent records through its four tools ), markdown=True, show_members_responses=True, diff --git a/cookbook/03_teams/12_learning/10_team_agentic_learning.py b/cookbook/03_teams/12_learning/10_team_agentic_learning.py index 8902be21c7..cd5ccf0820 100644 --- a/cookbook/03_teams/12_learning/10_team_agentic_learning.py +++ b/cookbook/03_teams/12_learning/10_team_agentic_learning.py @@ -11,8 +11,8 @@ In agentic mode: from agno.agent import Agent from agno.db.postgres import PostgresDb +from agno.learn.config import LearningMode from agno.learn.machine import LearningMachine -from agno.learn.mode import LearningMode from agno.models.openai import OpenAIResponses from agno.team import Team diff --git a/cookbook/08_learning/01_basics/5_entity_memory.py b/cookbook/08_learning/01_basics/5_entity_memory.py new file mode 100644 index 0000000000..13f99b2ffe --- /dev/null +++ b/cookbook/08_learning/01_basics/5_entity_memory.py @@ -0,0 +1,63 @@ +""" +Entity Memory: The Four Tools +============================= +Entity memory is the agent's knowledge about the WORLD - the people, +projects, companies and systems around the user - as opposed to user +memory, which is about the user themselves. + +It is AGENTIC-only: the agent records through four tools (remember_about, +link_entities, search_entities, forget), and the store does the librarian +work - ids are slugified from names, "Sarah Chen" and "sarah chen" resolve +to one person, and a correcting fact retires the stale one (supersession). + +Deep dives: cookbook/08_learning/04_entity_memory/ + +Run: + .venvs/demo/bin/python cookbook/08_learning/01_basics/5_entity_memory.py +""" + +from uuid import uuid4 + +from agno.agent import Agent +from agno.db.postgres import PostgresDb +from agno.learn import EntityMemoryConfig, LearningMachine +from agno.models.openai import OpenAIResponses + +# --------------------------------------------------------------------------- +# Create Agent +# --------------------------------------------------------------------------- + +db = PostgresDb(db_url="postgresql+psycopg://ai:ai@localhost:5532/ai") + +# Fresh per-run namespace so the demo starts clean on every execution. +NAMESPACE = f"basics_{uuid4().hex[:6]}" + +agent = Agent( + model=OpenAIResponses(id="gpt-5.5"), + db=db, + instructions="You are a sales assistant. Acknowledge notes briefly.", + learning=LearningMachine( + entity_memory=EntityMemoryConfig(namespace=NAMESPACE), + ), + markdown=True, +) + +# --------------------------------------------------------------------------- +# Run Demo +# --------------------------------------------------------------------------- + +if __name__ == "__main__": + agent.print_response( + "Note on Acme Corp: fintech startup in SF, about 50 people. " + "Jane Smith is their CTO.", + session_id="s1", + stream=True, + ) + + # A fresh session: the entity directory plus relevance recall carry the + # context - no tool call needed to answer. + agent.print_response( + "What do we know about Acme?", + session_id="s2", + stream=True, + ) diff --git a/cookbook/08_learning/01_basics/5a_entity_memory_always.py b/cookbook/08_learning/01_basics/5a_entity_memory_always.py deleted file mode 100644 index 4824625579..0000000000 --- a/cookbook/08_learning/01_basics/5a_entity_memory_always.py +++ /dev/null @@ -1,81 +0,0 @@ -""" -Entity Memory: Always Mode -========================== -Entity Memory stores knowledge about external things: -- Companies, people, projects -- Facts, events, relationships -- Shared context across users - -ALWAYS mode automatically extracts entity information from conversations. -No explicit tool calls - entities are discovered and saved behind the scenes. - -Compare with: 5b_entity_memory_agentic.py for explicit tool-based management. -""" - -from agno.agent import Agent -from agno.db.postgres import PostgresDb -from agno.learn import EntityMemoryConfig, LearningMachine, LearningMode -from agno.models.openai import OpenAIResponses - -# --------------------------------------------------------------------------- -# Create Agent -# --------------------------------------------------------------------------- - -db = PostgresDb(db_url="postgresql+psycopg://ai:ai@localhost:5532/ai") - -# ALWAYS mode: Entities are extracted automatically after responses. -# The agent doesn't see memory tools - extraction happens invisibly. -agent = Agent( - model=OpenAIResponses(id="gpt-5.5"), - db=db, - instructions="You're a sales assistant. Acknowledge notes briefly.", - learning=LearningMachine( - entity_memory=EntityMemoryConfig( - mode=LearningMode.ALWAYS, - ), - ), - markdown=True, -) - -# --------------------------------------------------------------------------- -# Run Demo -# --------------------------------------------------------------------------- - -if __name__ == "__main__": - from rich.pretty import pprint - - user_id = "sales@example.com" - - # Session 1: Mention entities naturally - print("\n" + "=" * 60) - print("SESSION 1: Discuss entities (extraction happens automatically)") - print("=" * 60 + "\n") - - agent.print_response( - "Just met with Acme Corp. They're a fintech startup in SF, " - "50 employees. CTO is Jane Smith. They use Python and Postgres.", - user_id=user_id, - session_id="session_1", - stream=True, - ) - - print("\n--- Extracted Entities ---") - entities = agent.learning_machine.entity_memory_store.search(query="acme", limit=10) - pprint(entities) - - # Session 2: Add more info about same entity - print("\n" + "=" * 60) - print("SESSION 2: Update same entity") - print("=" * 60 + "\n") - - agent.print_response( - "Update on Acme Corp: they just raised $50M Series B from Sequoia. " - "Jane Smith mentioned they're hiring 20 engineers.", - user_id=user_id, - session_id="session_2", - stream=True, - ) - - print("\n--- Updated Entities ---") - entities = agent.learning_machine.entity_memory_store.search(query="acme", limit=10) - pprint(entities) diff --git a/cookbook/08_learning/01_basics/5b_entity_memory_agentic.py b/cookbook/08_learning/01_basics/5b_entity_memory_agentic.py deleted file mode 100644 index a7cae7ef77..0000000000 --- a/cookbook/08_learning/01_basics/5b_entity_memory_agentic.py +++ /dev/null @@ -1,86 +0,0 @@ -""" -Entity Memory: Agentic Mode -=========================== -Entity Memory stores knowledge about external things: -- Companies, people, projects -- Facts, events, relationships -- Shared context across users - -AGENTIC mode gives the agent explicit tools to manage entities: -- search_entities, create_entity -- add_fact, add_event, add_relationship - -The agent decides when to store and retrieve information. - -Compare with: 5a_entity_memory_always.py for automatic extraction. -""" - -from agno.agent import Agent -from agno.db.postgres import PostgresDb -from agno.learn import EntityMemoryConfig, LearningMachine, LearningMode -from agno.models.openai import OpenAIResponses - -# --------------------------------------------------------------------------- -# Create Agent -# --------------------------------------------------------------------------- - -db = PostgresDb(db_url="postgresql+psycopg://ai:ai@localhost:5532/ai") - -# AGENTIC mode: Agent gets entity tools and decides when to use them. -# You'll see tool calls like "create_entity", "add_fact" in responses. -agent = Agent( - model=OpenAIResponses(id="gpt-5.5"), - db=db, - instructions=( - "You're a sales assistant tracking companies and contacts. " - "Be concise. Always search for existing entities before creating new ones." - ), - learning=LearningMachine( - entity_memory=EntityMemoryConfig( - mode=LearningMode.AGENTIC, - ), - ), - markdown=True, -) - -# --------------------------------------------------------------------------- -# Run Demo -# --------------------------------------------------------------------------- - -if __name__ == "__main__": - from rich.pretty import pprint - - user_id = "sales@example.com" - - # Session 1: Create entity - print("\n" + "=" * 60) - print("SESSION 1: Create entity (watch for tool calls)") - print("=" * 60 + "\n") - - agent.print_response( - "Track Acme Corp - fintech startup in SF, 50 employees, " - "uses Python and Postgres. CTO is Jane Smith.", - user_id=user_id, - session_id="session_1", - stream=True, - ) - - print("\n--- Created Entities ---") - entities = agent.learning_machine.entity_memory_store.search(query="acme", limit=10) - pprint(entities) - - # Session 2: Update same entity - print("\n" + "=" * 60) - print("SESSION 2: Update existing entity") - print("=" * 60 + "\n") - - agent.print_response( - "Acme Corp just raised $50M Series B from Sequoia.", - user_id=user_id, - session_id="session_2", - stream=True, - ) - - print("\n--- Updated Entities ---") - entities = agent.learning_machine.entity_memory_store.search(query="acme", limit=10) - pprint(entities) diff --git a/cookbook/08_learning/01_basics/6_extraction_limits.py b/cookbook/08_learning/01_basics/6_extraction_limits.py index 61254cc100..d5c2cca128 100644 --- a/cookbook/08_learning/01_basics/6_extraction_limits.py +++ b/cookbook/08_learning/01_basics/6_extraction_limits.py @@ -15,7 +15,6 @@ max_updates_per_run caps tool executions: from agno.agent import Agent from agno.db.postgres import PostgresDb from agno.learn import ( - EntityMemoryConfig, LearningMachine, LearningMode, UserMemoryConfig, @@ -32,7 +31,6 @@ db = PostgresDb(db_url="postgresql+psycopg://ai:ai@localhost:5532/ai") # Global max_updates_per_run=5 applies to all stores unless overridden. # user_profile: inherits 5 from LearningMachine # user_memory: explicit override to 3 -# entity_memory: explicit override to 15 (dense entity info needs more) agent = Agent( model=OpenAIResponses(id="gpt-5.5"), db=db, @@ -40,9 +38,6 @@ agent = Agent( max_updates_per_run=5, user_profile=UserProfileConfig(mode=LearningMode.ALWAYS), user_memory=UserMemoryConfig(mode=LearningMode.ALWAYS, max_updates_per_run=3), - entity_memory=EntityMemoryConfig( - mode=LearningMode.ALWAYS, max_updates_per_run=15 - ), ), markdown=True, debug_mode=True, # Shows "Tool call limit reached" logs @@ -91,16 +86,3 @@ if __name__ == "__main__": print("\n--- User Memory (limit: 3) ---") lm.user_memory_store.print(user_id=user_id) - - print("\n--- Entity Memory (limit: 15) ---") - from rich.pretty import pprint - - entities = lm.entity_memory_store.search(query="techcorp", limit=20) - if entities: - pprint(entities) - else: - print("No entities found") - - print("\n" + "=" * 70) - print("Check debug logs above for 'Tool call limit reached' messages") - print("=" * 70) diff --git a/cookbook/08_learning/01_basics/README.md b/cookbook/08_learning/01_basics/README.md index bb0d314fde..aab7d6874c 100644 --- a/cookbook/08_learning/01_basics/README.md +++ b/cookbook/08_learning/01_basics/README.md @@ -11,5 +11,4 @@ Core learning primitives and default patterns. - `3a_session_context_summary.py`: Session context summary tracking. - `3b_session_context_planning.py`: Session context planning with goal and progress state. - `4_learned_knowledge.py`: Learned knowledge retrieval and storage with vector search. -- `5a_entity_memory_always.py`: Entity memory extraction in ALWAYS mode. -- `5b_entity_memory_agentic.py`: Entity memory management in AGENTIC mode. +- `5_entity_memory.py`: The four entity tools (AGENTIC-only in 2.8.4). diff --git a/cookbook/08_learning/01_basics/TEST_LOG.md b/cookbook/08_learning/01_basics/TEST_LOG.md index f3a104115e..2617d0c744 100644 --- a/cookbook/08_learning/01_basics/TEST_LOG.md +++ b/cookbook/08_learning/01_basics/TEST_LOG.md @@ -1,5 +1,11 @@ # Test Log: 01_basics +> 2026-07-25 (feat/entity-memory-revamp): 5a/5b replaced by 5_entity_memory.py - +> entity memory is AGENTIC-only on the four-tool surface. Ran live against gpt-5.5: +> PASS (Acme Corp captured with properties + CTO edge in s1; s2 answered from the +> injected block with no tool call). 6_extraction_limits.py no longer includes +> entity memory (it has no extraction pass to limit). + > Tests not yet run. Run each file and update this log. ### 1a_user_profile_always.py @@ -72,23 +78,24 @@ --- -### 5a_entity_memory_always.py +### 5_entity_memory.py -**Status:** PENDING +**Status:** PASS -**Description:** Run and validate `5a_entity_memory_always.py` example behavior. +**Description:** The four entity tools; capture in one session, recall in a fresh one. -**Result:** Not run yet. +**Result:** Run live 2026-07-25 (gpt-5.5): Acme Corp captured with properties and the CTO +edge in s1; s2 answered from the injected block with no tool call. --- -### 5b_entity_memory_agentic.py +### 6_extraction_limits.py **Status:** PENDING -**Description:** Run and validate `5b_entity_memory_agentic.py` example behavior. +**Description:** max_updates_per_run caps for user_profile/user_memory (entity memory +removed - it has no extraction pass to limit). -**Result:** Not run yet. +**Result:** Not re-run in this pass. --- - diff --git a/cookbook/08_learning/04_entity_memory/01_facts_and_events.py b/cookbook/08_learning/04_entity_memory/01_facts_and_events.py deleted file mode 100644 index 82b08bf123..0000000000 --- a/cookbook/08_learning/04_entity_memory/01_facts_and_events.py +++ /dev/null @@ -1,103 +0,0 @@ -""" -Entity Memory: Facts and Events (Deep Dive) -============================================ -Semantic (facts) vs episodic (events) memory for entities. - -Entity Memory stores knowledge about external entities: -- Facts: Timeless truths ("Acme uses PostgreSQL") -- Events: Time-bound occurrences ("Acme raised $30M on Jan 15") - -AGENTIC mode gives the agent tools to create/update entities. - -Compare with: 04_always_extraction.py for automatic extraction. -See also: 01_basics/5a_entity_memory_always.py for the basics. -""" - -from agno.agent import Agent -from agno.db.postgres import PostgresDb -from agno.learn import EntityMemoryConfig, LearningMachine, LearningMode -from agno.models.openai import OpenAIResponses - -# --------------------------------------------------------------------------- -# Create Agent -# --------------------------------------------------------------------------- - -db = PostgresDb(db_url="postgresql+psycopg://ai:ai@localhost:5532/ai") - -# Dense entity info (many facts/events) can trigger many updates per message. -# max_updates_per_run (default: 10) caps updates per extraction. -agent = Agent( - model=OpenAIResponses(id="gpt-5.5"), - db=db, - instructions=( - "Track information about companies and people. " - "Distinguish between facts (timeless) and events (time-bound)." - ), - learning=LearningMachine( - entity_memory=EntityMemoryConfig( - mode=LearningMode.AGENTIC, - namespace="global", - ), - ), - markdown=True, -) - -# --------------------------------------------------------------------------- -# Run Demo -# --------------------------------------------------------------------------- - -if __name__ == "__main__": - from rich.pretty import pprint - - user_id = "research@example.com" - session_id = "company_research" - - # Share facts and events - print("\n" + "=" * 60) - print("MESSAGE 1: Share mixed facts and events") - print("=" * 60 + "\n") - - agent.print_response( - "Notes from my meeting with DataPipe: " - "They're based in San Francisco. " - "They build real-time ETL infrastructure in Rust. " - "Their CTO is Marcus Chen. " - "They just hit 1000 customers last month. " - "Series B closed at $80M two weeks ago.", - user_id=user_id, - session_id=session_id, - stream=True, - ) - print("\n--- Entities ---") - pprint( - agent.learning_machine.entity_memory_store.search(query="datapipe", limit=10) - ) - - # Query the entity - print("\n" + "=" * 60) - print("MESSAGE 2: Query the entity") - print("=" * 60 + "\n") - - agent.print_response( - "What do we know about DataPipe?", - user_id=user_id, - session_id="session_2", - stream=True, - ) - - # Add more events - print("\n" + "=" * 60) - print("MESSAGE 3: Add more events") - print("=" * 60 + "\n") - - agent.print_response( - "Update on DataPipe: They announced a partnership with BigCloud yesterday. " - "They're also opening a London office next quarter.", - user_id=user_id, - session_id="session_3", - stream=True, - ) - print("\n--- Updated Entities ---") - pprint( - agent.learning_machine.entity_memory_store.search(query="datapipe", limit=10) - ) diff --git a/cookbook/08_learning/04_entity_memory/01_the_four_tools.py b/cookbook/08_learning/04_entity_memory/01_the_four_tools.py new file mode 100644 index 0000000000..75f591e084 --- /dev/null +++ b/cookbook/08_learning/04_entity_memory/01_the_four_tools.py @@ -0,0 +1,86 @@ +""" +Entity Memory: The Four Tools +============================= +Entity memory is the agent's knowledge base about the world: the people, +projects, companies and systems around the user. The agent records through +four tools: + +- remember_about: upsert an entity by name - facts, events, a description, + and a note pointer. Resolution is the store's job (ids are slugified, + "Sarah Chen" and "sarah chen" are one person). +- link_entities: record a relationship; the edge is stored on both entities. +- search_entities: find entities, or list them by recency (no query). +- forget: retire a fact, or archive a whole entity. + +Corrections are just new facts: stating "radar shipped" retires "radar is +blocked on review" automatically (fact supersession), and facts render with +as-of dates so newer truth outranks older. + +Run: + .venvs/demo/bin/python cookbook/08_learning/04_entity_memory/01_the_four_tools.py +""" + +from uuid import uuid4 + +from agno.agent import Agent +from agno.db.postgres import PostgresDb +from agno.learn import EntityMemoryConfig, LearningMachine +from agno.models.openai import OpenAIResponses + +# --------------------------------------------------------------------------- +# Create Agent +# --------------------------------------------------------------------------- + +db = PostgresDb(db_url="postgresql+psycopg://ai:ai@localhost:5532/ai") + +# Fresh per-run namespace so the demo starts clean on every execution. +# A real deployment pins one namespace - that persistence is the point. +NAMESPACE = f"four_tools_{uuid4().hex[:6]}" + +agent = Agent( + model=OpenAIResponses(id="gpt-5.5"), + db=db, + instructions="You are a project tracker. Record what you are told, briefly.", + learning=LearningMachine( + entity_memory=EntityMemoryConfig(namespace=NAMESPACE), + ), + markdown=True, +) + +# --------------------------------------------------------------------------- +# Run Demo +# --------------------------------------------------------------------------- + +if __name__ == "__main__": + store = agent.learning_machine.entity_memory_store + + print("=" * 60) + print("TURN 1: capture a project and a person") + print("=" * 60) + agent.print_response( + "Track the radar project: it is blocked on security review, and Sarah Chen is the designer.", + session_id="s1", + stream=True, + ) + + print("=" * 60) + print("TURN 2: a correction - supersession retires the stale fact") + print("=" * 60) + agent.print_response( + "Good news: radar shipped v1 to production today.", + session_id="s2", + stream=True, + ) + + print("\n--- radar, live facts only (the blocked fact is retired, not deleted) ---") + store.print(entity_id="radar", entity_type="project", namespace=NAMESPACE) + + print("=" * 60) + print("TURN 3: recall in a fresh session - the entity directory plus") + print("relevance recall inject what this turn is about") + print("=" * 60) + agent.print_response( + "What do you know about radar?", + session_id="s3", + stream=True, + ) diff --git a/cookbook/08_learning/04_entity_memory/02_entity_relationships.py b/cookbook/08_learning/04_entity_memory/02_entity_relationships.py deleted file mode 100644 index 3b0ac2513b..0000000000 --- a/cookbook/08_learning/04_entity_memory/02_entity_relationships.py +++ /dev/null @@ -1,104 +0,0 @@ -""" -Entity Memory: Relationships (Deep Dive) -======================================== -Graph edges between entities. - -Relationships connect entities to form a knowledge graph: -- "Bob is CTO of Acme" -- "Acme acquired StartupX" -- "API Gateway depends on Auth Service" - -AGENTIC mode lets the agent create entities and add relationships. - -Compare with: 01_facts_and_events.py for facts/events. -See also: 01_basics/5b_entity_memory_agentic.py for the basics. -""" - -from agno.agent import Agent -from agno.db.postgres import PostgresDb -from agno.learn import EntityMemoryConfig, LearningMachine, LearningMode -from agno.models.openai import OpenAIResponses - -# --------------------------------------------------------------------------- -# Create Agent -# --------------------------------------------------------------------------- - -db = PostgresDb(db_url="postgresql+psycopg://ai:ai@localhost:5532/ai") - -# Org charts and relationship-heavy prompts can trigger many entity updates. -# max_updates_per_run (default: 10) caps updates per extraction. Increase -# if your prompts contain many entities. -agent = Agent( - model=OpenAIResponses(id="gpt-5.5"), - db=db, - instructions=( - "Build a knowledge graph of entities and their relationships. " - "Use appropriate relation types: works_at, reports_to, acquired, depends_on, etc." - ), - learning=LearningMachine( - entity_memory=EntityMemoryConfig( - mode=LearningMode.AGENTIC, - max_updates_per_run=15, - ), - ), - markdown=True, -) - -# --------------------------------------------------------------------------- -# Run Demo -# --------------------------------------------------------------------------- - -if __name__ == "__main__": - from rich.pretty import pprint - - user_id = "org@example.com" - session_id = "org_session" - - # Define org structure - print("\n" + "=" * 60) - print("MESSAGE 1: Define org structure") - print("=" * 60 + "\n") - - agent.print_response( - "TechCorp's leadership: " - "Sarah Chen is the CEO and founder. " - "Bob Martinez is the CTO, reporting to Sarah. " - "Alice Kim leads Engineering under Bob. " - "DevOps and Backend teams report to Alice.", - user_id=user_id, - session_id=session_id, - stream=True, - ) - print("\n--- Entities ---") - pprint( - agent.learning_machine.entity_memory_store.search(query="techcorp", limit=10) - ) - - # Query relationships - print("\n" + "=" * 60) - print("MESSAGE 2: Query relationships") - print("=" * 60 + "\n") - - agent.print_response( - "Who reports to Bob Martinez?", - user_id=user_id, - session_id="session_2", - stream=True, - ) - - # Add more relationships - print("\n" + "=" * 60) - print("MESSAGE 3: Company relationships") - print("=" * 60 + "\n") - - agent.print_response( - "TechCorp just acquired StartupAI for $50M. " - "They also partnered with CloudCo on infrastructure.", - user_id=user_id, - session_id="session_3", - stream=True, - ) - print("\n--- Updated Entities ---") - pprint( - agent.learning_machine.entity_memory_store.search(query="techcorp", limit=10) - ) diff --git a/cookbook/08_learning/04_entity_memory/02_links_and_forget.py b/cookbook/08_learning/04_entity_memory/02_links_and_forget.py new file mode 100644 index 0000000000..4f5a60337c --- /dev/null +++ b/cookbook/08_learning/04_entity_memory/02_links_and_forget.py @@ -0,0 +1,86 @@ +""" +Entity Memory: Links, Browse and Forget +======================================= +The graph half of entity memory: + +- link_entities writes the edge on BOTH entities, so "who works on radar?" + is answerable from either side, and recall shows linked entities by NAME + (one-hop expansion): "radar - designed_by <- Sarah Chen". +- search_entities with no query lists entities by recency - the browse + surface ("who works on what" needs enumeration). +- forget archives an entity: it leaves recall and the directory, stays + findable by explicit search with an (archived) marker, and any later + remember_about revives it. + +Run: + .venvs/demo/bin/python cookbook/08_learning/04_entity_memory/02_links_and_forget.py +""" + +from uuid import uuid4 + +from agno.agent import Agent +from agno.db.postgres import PostgresDb +from agno.learn import EntityMemoryConfig, LearningMachine +from agno.models.openai import OpenAIResponses + +# --------------------------------------------------------------------------- +# Create Agent +# --------------------------------------------------------------------------- + +db = PostgresDb(db_url="postgresql+psycopg://ai:ai@localhost:5532/ai") + +# Fresh per-run namespace so the demo starts clean on every execution. +# A real deployment pins one namespace - that persistence is the point. +NAMESPACE = f"links_{uuid4().hex[:6]}" + +agent = Agent( + model=OpenAIResponses(id="gpt-5.5"), + db=db, + instructions="You are a team tracker. Record what you are told, briefly.", + learning=LearningMachine( + entity_memory=EntityMemoryConfig(namespace=NAMESPACE), + ), + markdown=True, +) + +# --------------------------------------------------------------------------- +# Run Demo +# --------------------------------------------------------------------------- + +if __name__ == "__main__": + store = agent.learning_machine.entity_memory_store + + print("=" * 60) + print("TURN 1: capture people, projects and their links") + print("=" * 60) + agent.print_response( + "Sarah Chen designs the radar project. Tom Alvarez runs the infra platform. " + "Radar depends on the infra platform.", + session_id="s1", + stream=True, + ) + + print("\n--- the edge is on BOTH rows (reciprocal, with the far end's type) ---") + store.print(entity_id="radar", entity_type="project", namespace=NAMESPACE) + store.print(entity_id="sarah_chen", entity_type="person", namespace=NAMESPACE) + + print("=" * 60) + print("TURN 2: browse - no query lists entities by recency") + print("=" * 60) + agent.print_response( + "Who and what are you tracking right now? List everything.", + session_id="s2", + stream=True, + ) + + print("=" * 60) + print("TURN 3: archive - 'we killed radar' is a status change") + print("=" * 60) + agent.print_response( + "We cancelled the radar project. Archive it.", + session_id="s3", + stream=True, + ) + + print("\n--- archived: out of the directory, still searchable ---") + print(store.search_entities(query="radar", namespace=NAMESPACE)) diff --git a/cookbook/08_learning/04_entity_memory/README.md b/cookbook/08_learning/04_entity_memory/README.md index c29c87bbee..c278cb6fa8 100644 --- a/cookbook/08_learning/04_entity_memory/README.md +++ b/cookbook/08_learning/04_entity_memory/README.md @@ -1,8 +1,25 @@ # 04_entity_memory -Deep-dive examples for entity memory, facts, events, and relationships. +Deep-dive examples for entity memory: the agent's knowledge base about the +world (people, projects, companies, systems), revamped in 2.8.4 around four +tools with resolution, supersession and relevance recall in the store. ## Files -- `01_facts_and_events.py`: Distinguish timeless facts from time-bound events. -- `02_entity_relationships.py`: Capture and query entity relationships. +- `01_the_four_tools.py`: remember_about, a correction retiring a stale fact + via supersession, and relevance recall in a fresh session. +- `02_links_and_forget.py`: reciprocal links with one-hop names on recall, + browsing by recency, and archive/revive with forget. + +## The four tools + +| Tool | Job | +|---|---| +| `remember_about` | Upsert an entity by name: facts, events, description, `note=` pointer | +| `link_entities` | Record a relationship; the edge lands on both entities | +| `search_entities` | Find entities; with no query, list by recency (browse) | +| `forget` | Retire a fact, or archive a whole entity (reversible) | + +There are no ids to invent and no fact ids to track: resolution (slugified +ids, aliases, normalized types) and correction (fact supersession with as-of +dates) are the store's job. diff --git a/cookbook/08_learning/04_entity_memory/TEST_LOG.md b/cookbook/08_learning/04_entity_memory/TEST_LOG.md index 8c00f2fd64..e22b9d2c2b 100644 --- a/cookbook/08_learning/04_entity_memory/TEST_LOG.md +++ b/cookbook/08_learning/04_entity_memory/TEST_LOG.md @@ -1,24 +1,40 @@ # Test Log: 04_entity_memory -> Tests not yet run. Run each file and update this log. +> Tested 2026-07-25 against gpt-5.5 (OpenAIResponses), branch feat/entity-memory-revamp, +> Postgres (pgvector container on 5532). Rebuilt on the four-tool surface. -### 01_facts_and_events.py +### 01_the_four_tools.py -**Status:** PENDING +**Status:** PASS -**Description:** Run and validate `01_facts_and_events.py` example behavior. +**Description:** Capture (facts + events + link via remember_about/link_entities), a +correction, and relevance recall in a fresh session. -**Result:** Not run yet. +**Result:** Turn 2's "radar shipped v1" retired the "blocked on security review" fact via +supersession (tool reply: "Superseded 1 earlier fact(s)"); the live-facts print shows only +"status: v1 shipped to production (as of 2026-07-25)" plus the reciprocal designer edge. +Turn 3 in a fresh session recalled radar through the message-driven path and answered from +the injected block. Model behavior note: the agent also chose to record the ship as both a +fact (status line) and an event - exactly the intended split. + +Earlier phrasing note (kept for honesty): with turn 2 phrased "no longer using Postgres, we +switched to SQLite for the prototype", the judge did NOT supersede "uses Postgres" - +it read "for the prototype" as scoping that can coexist with prod Postgres. Defensible +conservatism; the demo now uses the unambiguous blocked->shipped correction. --- -### 02_entity_relationships.py +### 02_links_and_forget.py -**Status:** PENDING +**Status:** PASS -**Description:** Run and validate `02_entity_relationships.py` example behavior. +**Description:** Reciprocal links, one-hop names, browse-by-recency, archive via forget. -**Result:** Not run yet. +**Result:** The edge landed on both rows with the far end's type (radar: "designs <- +sarah_chen"; sarah_chen: outgoing twin). Browse turn listed all tracked entities with the +directory. The archive turn called forget; the explicit search then showed +"**radar project** (project) (archived)" while the entity left the directory. Model +behavior note: the agent named the entity "radar project" (slug radar_project) rather than +"radar" - consistent within the run, and resolution keeps later variants merged. --- - diff --git a/cookbook/08_learning/07_patterns/personal_assistant.py b/cookbook/08_learning/07_patterns/personal_assistant.py index 5483201686..5825b8cc7d 100644 --- a/cookbook/08_learning/07_patterns/personal_assistant.py +++ b/cookbook/08_learning/07_patterns/personal_assistant.py @@ -48,8 +48,7 @@ def create_personal_assistant(user_id: str, session_id: str) -> Agent: session_context=SessionContextConfig( enable_planning=True, ), - entity_memory=EntityMemoryConfig( - mode=LearningMode.ALWAYS, + entity_memory=EntityMemoryConfig( # AGENTIC-only: the agent records through its four tools namespace=f"user:{user_id}:personal", ), ), diff --git a/cookbook/08_learning/07_patterns/support_agent.py b/cookbook/08_learning/07_patterns/support_agent.py index da074f6d5c..53b99247d6 100644 --- a/cookbook/08_learning/07_patterns/support_agent.py +++ b/cookbook/08_learning/07_patterns/support_agent.py @@ -65,8 +65,7 @@ def create_support_agent(customer_id: str, ticket_id: str, org_id: str) -> Agent session_context=SessionContextConfig( enable_planning=True, ), - entity_memory=EntityMemoryConfig( - mode=LearningMode.ALWAYS, + entity_memory=EntityMemoryConfig( # AGENTIC-only: the agent records through its four tools namespace=f"org:{org_id}:support", ), learned_knowledge=LearnedKnowledgeConfig( diff --git a/cookbook/08_learning/09_decision_logs/02_decision_log_always.py b/cookbook/08_learning/09_decision_logs/02_decision_log_always.py deleted file mode 100644 index 1e59533106..0000000000 --- a/cookbook/08_learning/09_decision_logs/02_decision_log_always.py +++ /dev/null @@ -1,66 +0,0 @@ -""" -Decision Logs: ALWAYS Mode (Automatic Logging) -=============================================== - -This example demonstrates automatic decision logging where -tool calls are automatically recorded as decisions. - -In ALWAYS mode, DecisionLogStore extracts decisions from: -- Tool calls (which tool was used) -- Other significant choices the agent makes - -Run: - .venvs/demo/bin/python cookbook/08_learning/09_decision_logs/02_decision_log_always.py -""" - -from agno.agent import Agent -from agno.db.postgres import PostgresDb -from agno.learn import DecisionLogConfig, LearningMachine, LearningMode -from agno.models.openai import OpenAIResponses -from agno.tools.duckduckgo import DuckDuckGoTools - -# --------------------------------------------------------------------------- -# Setup -# --------------------------------------------------------------------------- -# Database connection -db = PostgresDb(db_url="postgresql+psycopg://ai:ai@localhost:5532/ai") - -# --------------------------------------------------------------------------- -# Create Agent -# --------------------------------------------------------------------------- -# Create an agent with automatic decision logging -# ALWAYS mode: Tool calls are automatically logged as decisions -agent = Agent( - id="auto-decision-logger", - name="Auto Decision Logger", - model=OpenAIResponses(id="gpt-5.5"), - db=db, - learning=LearningMachine( - decision_log=DecisionLogConfig( - mode=LearningMode.ALWAYS, - ), - ), - tools=[DuckDuckGoTools()], - instructions=[ - "You are a helpful research assistant.", - "Use web search to find current information when needed.", - ], - markdown=True, -) - -# --------------------------------------------------------------------------- -# Run Agent -# --------------------------------------------------------------------------- -if __name__ == "__main__": - # Test: Agent uses a tool (will be logged automatically) - print("=== Test: Agent uses web search ===\n") - agent.print_response( - "What are the latest developments in AI agents?", - session_id="session-002", - ) - - # View auto-logged decisions - print("\n=== Auto-Logged Decisions ===\n") - decision_store = agent.learning_machine.decision_log_store - if decision_store: - decision_store.print(agent_id="auto-decision-logger", limit=10) diff --git a/cookbook/08_learning/09_decision_logs/02_record_outcomes.py b/cookbook/08_learning/09_decision_logs/02_record_outcomes.py new file mode 100644 index 0000000000..2b5343a914 --- /dev/null +++ b/cookbook/08_learning/09_decision_logs/02_record_outcomes.py @@ -0,0 +1,62 @@ +""" +Decision Logs: Recording Outcomes +================================= +The feedback half of the decision log: log_decision records the choice with +its reasoning, and record_outcome closes the loop later with what actually +happened. Decision logging is AGENTIC-only - the agent logs deliberately; +there is no automatic extraction pass. + +Run: + .venvs/demo/bin/python cookbook/08_learning/09_decision_logs/02_record_outcomes.py +""" + +from agno.agent import Agent +from agno.db.postgres import PostgresDb +from agno.learn import DecisionLogConfig, LearningMachine +from agno.models.openai import OpenAIResponses + +# --------------------------------------------------------------------------- +# Create Agent +# --------------------------------------------------------------------------- + +db = PostgresDb(db_url="postgresql+psycopg://ai:ai@localhost:5532/ai") + +agent = Agent( + id="outcome-logger", + name="Outcome Logger", + model=OpenAIResponses(id="gpt-5.5"), + db=db, + learning=LearningMachine( + decision_log=DecisionLogConfig(), + ), + instructions=[ + "You are an engineering advisor.", + "When you make a recommendation, log it as a decision with your reasoning.", + "When told how a past recommendation worked out: first call " + "search_decisions to find that decision and its id, then call " + "record_outcome with that id. Never ask the user for a decision id.", + ], + markdown=True, +) + +# --------------------------------------------------------------------------- +# Run Demo +# --------------------------------------------------------------------------- + +if __name__ == "__main__": + agent.print_response( + "Should we use Postgres or DynamoDB for the new billing service? " + "We need transactions and our team knows SQL.", + session_id="s1", + stream=True, + ) + + agent.print_response( + "Update: we went with your Postgres recommendation and the migration " + "went smoothly. Record that outcome.", + session_id="s2", + stream=True, + ) + + print("\n--- the decision log, with its outcome ---") + agent.learning_machine.decision_log_store.print(agent_id="outcome-logger") diff --git a/cookbook/08_learning/09_decision_logs/README.md b/cookbook/08_learning/09_decision_logs/README.md index 3fd46d6015..1b761910ee 100644 --- a/cookbook/08_learning/09_decision_logs/README.md +++ b/cookbook/08_learning/09_decision_logs/README.md @@ -5,4 +5,4 @@ Examples for capturing and reviewing agent decision logs. ## Files - `01_basic_decision_log.py`: AGENTIC decision logging with explicit tool usage. -- `02_decision_log_always.py`: ALWAYS-mode automatic decision logging. +- `02_record_outcomes.py`: Close the loop - record how a logged decision worked out. diff --git a/cookbook/08_learning/09_decision_logs/TEST_LOG.md b/cookbook/08_learning/09_decision_logs/TEST_LOG.md index 4ad6cc30be..c9c22d6a44 100644 --- a/cookbook/08_learning/09_decision_logs/TEST_LOG.md +++ b/cookbook/08_learning/09_decision_logs/TEST_LOG.md @@ -1,24 +1,30 @@ # Test Log: 09_decision_logs -> Tests not yet run. Run each file and update this log. +> Tested 2026-07-25 against gpt-5.5 (OpenAIResponses), branch feat/entity-memory-revamp. +> Decision logging is AGENTIC-only in 2.8.4; the old 02_decision_log_always.py showcased +> the deleted contentless-row extraction and is replaced by 02_record_outcomes.py. ### 01_basic_decision_log.py **Status:** PENDING -**Description:** Run and validate `01_basic_decision_log.py` example behavior. +**Description:** Basic log_decision / search_decisions flow. -**Result:** Not run yet. +**Result:** Not re-run in this pass (unchanged API surface; the store default is now +AGENTIC, which this file already used implicitly through LearningMachine). --- -### 02_decision_log_always.py +### 02_record_outcomes.py -**Status:** PENDING +**Status:** PASS -**Description:** Run and validate `02_decision_log_always.py` example behavior. +**Description:** log_decision in one session, record_outcome in a later one. -**Result:** Not run yet. +**Result:** The Postgres-vs-DynamoDB recommendation was logged with full reasoning; the +follow-up session searched, found the decision id and recorded outcome "migration went +smoothly" with quality good. Honest note: on the first run the agent logged a meta-decision +asking for the id instead of searching - the instructions now say search first, never ask +for an id, and the re-run recorded the outcome correctly. --- - diff --git a/cookbook/08_learning/11_composition/README.md b/cookbook/08_learning/11_composition/README.md new file mode 100644 index 0000000000..9ea586d1fe --- /dev/null +++ b/cookbook/08_learning/11_composition/README.md @@ -0,0 +1,37 @@ +# 11_composition + +The manual door. `learning=` is the automatic door: pass it and the framework +attaches context, instructions and tools at fixed positions. Don't pass it, +and the framework attaches nothing - you place the machine's three public +surfaces yourself, the way FileSystem composes. Read this folder next to +`00_quickstart` to see the two doors side by side. + +## Files + +- `basic.py`: tools=[*learning.get_tools()] + instructions=[learning.instructions()]. +- `with_filesystem.py`: LearningMachine + FileSystem + your own system prompt, + in one deliberate order. +- `context_block.py`: build_context() placed via additional_context - data + without tools. +- `always_capture.py`: post_hooks=[learning.capture_hook()] - ALWAYS-mode + extraction through the manual door (the escape hatch). + +## The three surfaces + +| Surface | Returns | Place it in | +|---|---|---| +| `learning.get_tools(user_id=...)` | the capture tools | `tools=[...]` | +| `learning.instructions()` | the guidance block | `instructions=[...]` | +| `learning.build_context(user_id=..., message=...)` | the recalled data | `additional_context` / a dependency | + +The manual door injects nothing - give the machine its `db` and its `model` +explicitly. Every capture path is a model call: without one, `update_profile` +and `update_user_memory` return "No model provided", `capture_hook`'s ALWAYS +extraction stores nothing, and entity memory keeps every stated fact instead of +retiring the ones it contradicts. `get_tools()` warns once when a store is in +that state. + +The manual door is agentic by nature: with no `learning=` there is no +automatic post-run extraction, and the tools are the capture mechanism. +Passing the same machine to `learning=` AND placing its surfaces by hand +renders the blocks twice - the framework warns once when it detects that. diff --git a/cookbook/08_learning/11_composition/TEST_LOG.md b/cookbook/08_learning/11_composition/TEST_LOG.md new file mode 100644 index 0000000000..5e1fd8260f --- /dev/null +++ b/cookbook/08_learning/11_composition/TEST_LOG.md @@ -0,0 +1,70 @@ +# Test Log: 11_composition + +> Tested 2026-07-25 against gpt-5.5 (OpenAIResponses), branch feat/entity-memory-revamp, +> Postgres (pgvector container on 5532). Re-tested 2026-07-26 after the missing-model fix. + +### basic.py + +**Status:** PASS + +**Result:** With no learning=, the hand-placed tools captured the preference (user memory) +and the Meridian project + Priya link (entity memory); the printed manual-door surfaces +show the guidance block and a data block whose relevance recall expanded Meridian for the +message "what about meridian?" with the one-hop "runs <- Priya" edge. + +**2026-07-26 correction:** the first log was wrong about the user memory. The machine +carried no `model=`, so `update_user_memory` returned "No model provided for memories +extraction" and stored nothing - only the entity write (no model needed) landed. The +same hole `always_capture.py` hit below; the manual door injects nothing. Re-run with +`model=` on the machine: `update_user_memory` stored "Prefers sources with primary data", +verified in the learnings table. + +--- + +### with_filesystem.py + +**Status:** PASS + +**Result:** One deliberate order: learning tools + fs tools + both instruction blocks. The +agent wrote notes/vector-db-comparison.md with the deadline. Model behavior note: it also +wrote the "conclusions first" preference INTO the note alongside saving it - the +one-claim-one-home discipline is exactly what the second-brain instructions add on top. + +**2026-07-26 correction:** same missing `model=`, so the note was written but the +preference was not stored. Re-run with the model: `update_user_memory(task=User prefers +conclusions first in every summary.)` and `append_file(notes/tasks.md)` both fired, and +the memory is in the table. + +--- + +### context_block.py + +**Status:** PASS + +**Result:** build_context() placed via additional_context, no tools: the read-only agent +answered from its own knowledge. + +**2026-07-26 correction:** the earlier "despite the seeded memory in the context, the +summary put its conclusion last" reads a model failure into a store failure - the seeding +call needed a model too, so nothing was seeded and the context block was empty. With +`model=` on the machine the seed lands; the answer still leads with its list and closes +with the conclusion, so the original observation (a data-only block informs but does not +compel) holds on the re-run. + +--- + +### always_capture.py + +**Status:** PASS + +**Result:** post_hooks=[learning.capture_hook()] ran ALWAYS extraction in the background: +profile (Name/Preferred Name: Dana) and one memory (data engineer, Lisbon, ClickHouse +pipelines) appeared without any tool call. First run FAILED with empty stores - the manual +door injects nothing, so the machine needed model= passed explicitly; the file and README +now say so. + +**2026-07-26:** this was the only file that had learned the lesson. `LearningMachine.get_tools()` +now warns once when a store that captures through a model has none, so the next person +finds out at attach time instead of from an empty table. + +--- diff --git a/cookbook/08_learning/11_composition/always_capture.py b/cookbook/08_learning/11_composition/always_capture.py new file mode 100644 index 0000000000..8be9260e9d --- /dev/null +++ b/cookbook/08_learning/11_composition/always_capture.py @@ -0,0 +1,51 @@ +""" +Composition: ALWAYS Capture Through the Manual Door +=================================================== +The manual door has no automatic post-run extraction - the tools are the +capture mechanism. For hand-placed prompts AND ALWAYS-mode extraction, +capture_hook() returns a post_hooks-compatible callable around the machine's +capture pass (backgrounded on the agent's executor). An escape hatch, not a +third shape. + +Run: + .venvs/demo/bin/python cookbook/08_learning/11_composition/always_capture.py +""" + +import time + +from agno.agent import Agent +from agno.db.postgres import PostgresDb +from agno.learn import LearningMachine +from agno.models.openai import OpenAIResponses + +db = PostgresDb(db_url="postgresql+psycopg://ai:ai@localhost:5532/ai") + +# ALWAYS-mode stores: extraction runs after the response, no agent tools needed. +# The manual door injects nothing: the machine needs its db AND its model +# given explicitly (extraction is a model call). +learning = LearningMachine( + db=db, model=OpenAIResponses(id="gpt-5.5"), user_profile=True, user_memory=True +) + +USER_ID = "composer@example.com" + +agent = Agent( + model=OpenAIResponses(id="gpt-5.5"), + db=db, + instructions=["You are a helpful assistant."], + post_hooks=[learning.capture_hook()], + user_id=USER_ID, + markdown=True, +) + +if __name__ == "__main__": + agent.print_response( + "I'm Dana, a data engineer in Lisbon. I mostly work on our ClickHouse pipelines.", + stream=True, + ) + + # The capture pass runs in the background; give it a moment before reading + time.sleep(10) + print("\n--- what ALWAYS capture extracted ---") + learning.user_profile_store.print(user_id=USER_ID) + learning.user_memory_store.print(user_id=USER_ID) diff --git a/cookbook/08_learning/11_composition/basic.py b/cookbook/08_learning/11_composition/basic.py new file mode 100644 index 0000000000..7d73d029e8 --- /dev/null +++ b/cookbook/08_learning/11_composition/basic.py @@ -0,0 +1,67 @@ +""" +Composition: The Manual Door +============================ +learning= is the automatic door: the framework injects context, instructions +and tools for you. This folder is the other door - no learning= at all. You +place the three public surfaces yourself, the way FileSystem composes: + +- learning.get_tools(...) the capture tools +- learning.instructions() the guidance block (how to use them) +- learning.build_context(...) the recalled-data block + +An agent with no learning= has no automatic capture: the manual door is +agentic by nature - the agent captures by calling the tools you handed it. + +Run: + .venvs/demo/bin/python cookbook/08_learning/11_composition/basic.py +""" + +from agno.agent import Agent +from agno.db.postgres import PostgresDb +from agno.learn import LearningMachine, LearningMode, UserMemoryConfig +from agno.models.openai import OpenAIResponses + +# --------------------------------------------------------------------------- +# Build the machine, place its surfaces by hand +# --------------------------------------------------------------------------- + +db = PostgresDb(db_url="postgresql+psycopg://ai:ai@localhost:5532/ai") + +# The manual door injects nothing: without learning= nobody hands the machine +# the agent's model, and capture is a model call. +learning = LearningMachine( + db=db, + model=OpenAIResponses(id="gpt-5.5"), + user_memory=UserMemoryConfig(mode=LearningMode.AGENTIC), + entity_memory=True, +) + +USER_ID = "composer@example.com" + +agent = Agent( + model=OpenAIResponses(id="gpt-5.5"), + db=db, + tools=[*learning.get_tools(user_id=USER_ID)], + instructions=[ + "You are a research assistant.", + learning.instructions(), + ], + user_id=USER_ID, + markdown=True, +) + +# --------------------------------------------------------------------------- +# Run +# --------------------------------------------------------------------------- + +if __name__ == "__main__": + agent.print_response( + "Remember that I prefer sources with primary data, and track the " + "Meridian project - Priya runs it.", + stream=True, + ) + + print("\n--- what the manual door placed (guidance + data) ---") + print(learning.instructions()[:400]) + print("...") + print(learning.build_context(user_id=USER_ID, message="what about meridian?")) diff --git a/cookbook/08_learning/11_composition/context_block.py b/cookbook/08_learning/11_composition/context_block.py new file mode 100644 index 0000000000..934f3b43ad --- /dev/null +++ b/cookbook/08_learning/11_composition/context_block.py @@ -0,0 +1,41 @@ +""" +Composition: The Data Block via additional_context +================================================== +build_context() returns the recalled-data block on its own, for when you +want the data injected but not the tools - a read-only view of what the +machine knows, placed exactly where you choose. + +Run: + .venvs/demo/bin/python cookbook/08_learning/11_composition/context_block.py +""" + +from agno.agent import Agent +from agno.db.postgres import PostgresDb +from agno.learn import LearningMachine, LearningMode, UserMemoryConfig +from agno.models.openai import OpenAIResponses + +db = PostgresDb(db_url="postgresql+psycopg://ai:ai@localhost:5532/ai") + +learning = LearningMachine( + db=db, + model=OpenAIResponses(id="gpt-5.5"), # the manual door injects nothing + user_memory=UserMemoryConfig(mode=LearningMode.AGENTIC), +) + +USER_ID = "composer@example.com" + +# Seed a memory through the data API so the read-only agent has something to see +store = learning.user_memory_store +store.get_tools(user_id=USER_ID)[0]("Prefers conclusions first, then supporting detail") + +# No learning=, no tools: just the data block, placed as additional context +agent = Agent( + model=OpenAIResponses(id="gpt-5.5"), + db=db, + additional_context=learning.build_context(user_id=USER_ID), + user_id=USER_ID, + markdown=True, +) + +if __name__ == "__main__": + agent.print_response("Summarize why teams adopt vector databases.", stream=True) diff --git a/cookbook/08_learning/11_composition/with_filesystem.py b/cookbook/08_learning/11_composition/with_filesystem.py new file mode 100644 index 0000000000..cf40368bf7 --- /dev/null +++ b/cookbook/08_learning/11_composition/with_filesystem.py @@ -0,0 +1,50 @@ +""" +Composition: LearningMachine + FileSystem, One Deliberate Order +=============================================================== +The point of the manual door: LearningMachine, FileSystem and your own +system prompt compose in one order you can read off the page. Nothing is +attached behind your back. + +Run: + .venvs/demo/bin/python cookbook/08_learning/11_composition/with_filesystem.py +""" + +from agno.agent import Agent +from agno.db.postgres import PostgresDb +from agno.fs import FileSystem +from agno.learn import LearningMachine, LearningMode, UserMemoryConfig +from agno.models.openai import OpenAIResponses + +db = PostgresDb(db_url="postgresql+psycopg://ai:ai@localhost:5532/ai") + +learning = LearningMachine( + db=db, + model=OpenAIResponses(id="gpt-5.5"), # the manual door injects nothing + user_memory=UserMemoryConfig(mode=LearningMode.AGENTIC), +) +fs = FileSystem(db, namespace="composition-notes") + +USER_ID = "composer@example.com" + +agent = Agent( + model=OpenAIResponses(id="gpt-5.5"), + db=db, + tools=[*learning.get_tools(user_id=USER_ID), fs.tools()], + instructions=[ + "You are a research assistant. Keep running notes on topics you research.", + learning.instructions(), + fs.instructions(), + ], + user_id=USER_ID, + markdown=True, +) + +if __name__ == "__main__": + agent.print_response( + "Note down: the vector-db comparison is due Friday. And remember that " + "I want conclusions first in every summary.", + stream=True, + ) + print("\n--- files ---") + for f in fs.list(): + print(f.path) diff --git a/cookbook/08_learning/README.md b/cookbook/08_learning/README.md index 82254c6b5c..422b296caf 100644 --- a/cookbook/08_learning/README.md +++ b/cookbook/08_learning/README.md @@ -57,8 +57,8 @@ cookbook/08_learning/ │ ├── 3a_session_context_summary.py │ ├── 3b_session_context_planning.py │ ├── 4_learned_knowledge.py -│ ├── 5a_entity_memory_always.py -│ └── 5b_entity_memory_agentic.py +│ ├── 5_entity_memory.py +│ └── 6_extraction_limits.py │ ├── 02_user_profile/ # Deep dives into user profiles │ ├── 01_always_extraction.py @@ -69,9 +69,9 @@ cookbook/08_learning/ │ ├── 01_summary_mode.py │ └── 02_planning_mode.py │ -├── 04_entity_memory/ # Deep dives into entity memory -│ ├── 01_facts_and_events.py -│ └── 02_entity_relationships.py +├── 04_entity_memory/ # Deep dives into entity memory (the four tools) +│ ├── 01_the_four_tools.py +│ └── 02_links_and_forget.py │ ├── 05_learned_knowledge/ # Deep dives into learned knowledge │ ├── 01_agentic_mode.py @@ -88,14 +88,20 @@ cookbook/08_learning/ │ ├── 01_minimal_custom_store.py │ └── 02_custom_store_with_db.py │ -├── 09_decision_logs/ # Decision logging and auditing +├── 09_decision_logs/ # Decision logging and auditing (AGENTIC-only) │ ├── 01_basic_decision_log.py -│ └── 02_decision_log_always.py +│ └── 02_record_outcomes.py │ -└── 10_demo/ # AgentOS demo: browse learnings in the UI - ├── agents.py - ├── seed.py - └── run.py +├── 10_demo/ # AgentOS demo: browse learnings in the UI +│ ├── agents.py +│ ├── seed.py +│ └── run.py +│ +└── 11_composition/ # The manual door: place the surfaces yourself + ├── basic.py + ├── with_filesystem.py + ├── context_block.py + └── always_capture.py ``` ## Running the Cookbooks @@ -388,7 +394,9 @@ agent = Agent( Captures knowledge about external entities: companies, projects, people, products, systems. -**Supported modes:** ALWAYS, AGENTIC +**Supported modes:** AGENTIC only. The agent records through four tools +(`remember_about`, `link_entities`, `search_entities`, `forget`); there is no +extraction pass, and any other mode raises. **Three types of entity data:** - **Facts** (semantic memory): Timeless truths - "Uses PostgreSQL" @@ -420,7 +428,8 @@ agent.run("What database does Acme use?") Records decisions the agent makes, with reasoning and alternatives considered. Useful for auditing agent behavior and building feedback loops. -**Supported modes:** ALWAYS, AGENTIC +**Supported modes:** AGENTIC. The decision is the agent's to record, so it +records it with `log_decision`. **Scope:** Per agent - stored and retrieved by `agent_id`. diff --git a/cookbook/13_filesystem/02_durable_records/basic.py b/cookbook/13_filesystem/02_durable_records/basic.py index 28d884fb11..102b691732 100644 --- a/cookbook/13_filesystem/02_durable_records/basic.py +++ b/cookbook/13_filesystem/02_durable_records/basic.py @@ -5,6 +5,10 @@ Deduplicate work with a record log. The agent calls check_lines before it acts and append_file after. The log matches on exact lines and outlives the process, so a recurring job never repeats work it has already done. +check_lines is not in the default toolset (the default is the notes seven); +a feed-processing agent requests it explicitly via include_tools, and carries +the record-set conventions in its own instructions. + This example processes two overlapping ticket batches. The second pass acts only on the genuinely new ticket. """ @@ -24,15 +28,29 @@ DB_FILE = f"tmp/agent_fs_records_{uuid4().hex}.db" fs = FileSystem(SqliteDb(db_file=DB_FILE)) +# The record-set conventions, carried by the agents that use the pattern. +RECORD_SET_INSTRUCTIONS = ( + "To keep a record set (items already processed): store one record per line " + "in the seen/ directory, one file per date (seen/2026-07-24.md). Call " + "check_lines(lines, directory='seen') BEFORE acting, then record the new " + "ones with append_file(..., unique=True). check_lines matches exact whole " + "lines, and unique=True keeps the log free of duplicates if a run overlaps " + "another one." +) + # --------------------------------------------------------------------------- # Create Agent # --------------------------------------------------------------------------- agent = Agent( model=OpenAIResponses(id="gpt-5.5"), - tools=[fs.tools()], + tools=[ + fs.tools( + include_tools=["check_lines", "append_file", "read_file", "list_files"] + ) + ], instructions=[ "You triage support tickets.", - fs.instructions(), + RECORD_SET_INSTRUCTIONS, ], ) diff --git a/cookbook/13_filesystem/02_durable_records/radar_news_delta.py b/cookbook/13_filesystem/02_durable_records/radar_news_delta.py index 9d172e825a..2ec5f81536 100644 --- a/cookbook/13_filesystem/02_durable_records/radar_news_delta.py +++ b/cookbook/13_filesystem/02_durable_records/radar_news_delta.py @@ -48,15 +48,20 @@ fs = FileSystem(SqliteDb(db_file=DB_FILE)) # --------------------------------------------------------------------------- agent = Agent( model=OpenAIResponses(id="gpt-5.5"), - tools=[fs.tools()], + # check_lines is not in the default toolset; a feed agent asks for it by name. + tools=[ + fs.tools( + include_tools=["check_lines", "append_file", "read_file", "list_files"] + ) + ], instructions=[ "You are Radar, a news-brief agent that runs on a schedule. Each run you " "receive the current feed as 'id|headline' lines. Report ONLY stories you " "have never reported before: call check_lines on the ids with " "directory='seen' first, brief the missing ones (one bullet each), then " f"record the newly reported ids with append_file to seen/{TODAY}.md, one " - "id per line. If nothing is new, say so.", - fs.instructions(), + "id per line, with unique=True. If nothing is new, say so. check_lines " + "matches exact whole lines, so pass ids exactly as you will store them.", ], ) diff --git a/cookbook/13_filesystem/04_namespaces/README.md b/cookbook/13_filesystem/04_namespaces/README.md index a2388579e0..94ddb2434c 100644 --- a/cookbook/13_filesystem/04_namespaces/README.md +++ b/cookbook/13_filesystem/04_namespaces/README.md @@ -10,7 +10,7 @@ The main reason to reach for one is per-user (and per-team) file stores from a s - `basic.py`: the declarative common case, `namespace="assistant/{user_id}"`. One agent serves alice and bob with fully isolated files, and an anonymous run fails closed instead of collapsing into a shared store. - `custom_factory.py`: the escape hatch for arbitrary policy. A callable tool factory builds the FileSystem from `run_context`, and here VIP users get their own tier of namespaces. -- `shared_namespace.py`: two agents share files by attaching the same namespace name. The producer writes, and the consumer attaches with `tools(read_only=True)` plus `instructions(read_only=True)`, which gives it four read tools and no way to write. +- `shared_namespace.py`: two agents share files by attaching the same namespace name. The producer writes, and the consumer attaches with `tools(read_only=True)` plus `instructions(read_only=True)`, which gives it three read tools and no way to write. ## When to use diff --git a/cookbook/13_filesystem/04_namespaces/TEST_LOG.md b/cookbook/13_filesystem/04_namespaces/TEST_LOG.md index 3db30d437b..8dc71be6a7 100644 --- a/cookbook/13_filesystem/04_namespaces/TEST_LOG.md +++ b/cookbook/13_filesystem/04_namespaces/TEST_LOG.md @@ -44,3 +44,13 @@ folder now passes `fs.instructions()` in its own instructions list. Re-tested ag **Result:** Consumer tool surface printed exactly `['read_file', 'list_files', 'search_content', 'check_lines']`, with no write, append, move or delete. The recorder appended both decisions in one `append_file` call. The consumer answered via `list_files(directory=., ...)` then `read_file(path=decisions/2026-07.md, ...)`, finding the file without being told its path, and reported pgvector as approved for production. Note that read_only constrains the tool surface only: the same FileSystem object still exposes `write()` from Python. --- + +## 2026-07-26 — corrections after the default flip + +`read_only=True` is now three tools, not four: `check_lines` left every default +set and is reachable only through `include_tools`. Both entries above record a +consumer surface of `['read_file', 'list_files', 'search_content', 'check_lines']`, +which is what the pre-flip code printed; `shared_namespace.py` itself was updated +and says "three read tools". Re-verified against the current code: +`FileSystem(db, namespace=...).tools(read_only=True)` yields exactly +`['read_file', 'list_files', 'search_content']`. diff --git a/cookbook/13_filesystem/04_namespaces/shared_namespace.py b/cookbook/13_filesystem/04_namespaces/shared_namespace.py index a32aedac18..7b05f4ee3c 100644 --- a/cookbook/13_filesystem/04_namespaces/shared_namespace.py +++ b/cookbook/13_filesystem/04_namespaces/shared_namespace.py @@ -3,7 +3,7 @@ Namespaces - Sharing One Store ============================== Two agents share files by attaching the same namespace by name. The producer gets the full tool surface. The consumer gets tools(read_only=True) and the -matching instructions(read_only=True), which is four read tools, so it can +matching instructions(read_only=True), which is three read tools, so it can consult the records but holds no tool that could change them. This example has a recorder agent write decisions and an answering agent look diff --git a/cookbook/13_filesystem/05_operations/quota_recovery.py b/cookbook/13_filesystem/05_operations/quota_recovery.py index 01ef8d0182..6eba2e115c 100644 --- a/cookbook/13_filesystem/05_operations/quota_recovery.py +++ b/cookbook/13_filesystem/05_operations/quota_recovery.py @@ -3,8 +3,9 @@ Operations - Quota Recovery =========================== FileSystem caps the size of every file and every namespace, and nothing is ever evicted silently. This example hits both caps on purpose, shows the -guidance the agent gets back, and then recovers the way that guidance -suggests, by starting a new partition and deleting old ones. +guidance the agent gets back, and then recovers programmatically: starting +a new partition and deleting old ones through the Python API (the delete +tool itself is opt-in via tools(allow_delete=True)). The caps are set very small here so the numbers stay readable. No model, no API keys. diff --git a/cookbook/examples/second_brain/TEST_LOG.md b/cookbook/examples/second_brain/TEST_LOG.md index d05422a6c8..5a6207731e 100644 --- a/cookbook/examples/second_brain/TEST_LOG.md +++ b/cookbook/examples/second_brain/TEST_LOG.md @@ -1,50 +1,107 @@ -# Test Log - second_brain +# Test Log: second_brain -Tested 2026-07-25 against `gpt-5.6` (via `model="openai:gpt-5.6"`, which resolves to OpenAIResponses), -agno 2.8.2 (source tree at 5e6185ea9). Entries quote tool calls and printed state. -Model prose varies run to run and is paraphrased. +> Tested 2026-07-25 against gpt-5.6 (via model="openai:gpt-5.6"), branch +> feat/entity-memory-revamp. Rebuilt on the LearningMachine: entity memory (four tools, +> shared "global" namespace) + AGENTIC user profile/memory + shared "brain" notes +> namespace + the one-claim-one-home rule, with identity pinned (Agent(user_id="owner")). +> The pre-2.8.4 store (enable_agentic_memory era) was set aside as +> tmp/second_brain_pre284.db.bak; this brain starts fresh on the learnings table. -### test.py +## second_brain.py - REST live proof (8 turns, DB verified between turns) **Status:** PASS -**Description:** Durable, owned memory across a session boundary. The driver captures a decision in one session, then asks for it back in a brand new session that shares no history, and finally lists the files in that user's namespace. +Served with uvicorn on 7777; turns driven via POST /agents/second-brain/runs, each in a +fresh session; every claim below checked in tmp/second_brain.db between turns. -**Result:** Cold run (`rm -rf tmp` first), exit 0. The capture session wrote the memory and the note in the same turn: +1. Capture (project + decision): harbor filed with the index line + "database: Postgres, over DynamoDB - see note", the reasoning ONLY in notes/harbor.md, + and the note pointer on the entity (properties.note = notes/harbor.md). The agent also + linked postgres/dynamodb, created as minimal entity_type="unknown" targets - the + designed link-first-describe-later shape. +2. Capture (person): Sarah Chen recorded; the edge landed on BOTH rows with the far end's + type (harbor: designs <- sarah_chen, entity_type person). +3. Correction: "not blocked anymore, v0.1 shipped" retired + "status: blocked on the auth review" (superseded_by = the new fact's id, verified in + the fact record) and appended a dated entry to the note. +4. Why-question (fresh session): answered with the note's reasoning verbatim shape + (multi-row transactions, team knows SQL, single-table modeling) - the note pointer + round-trip works. +5. Negative: "Have we ever discussed the Acme contract?" -> "No - I searched the entity + directory and all notes for Acme and contract and found nothing." A grounded no. +6. Archive: legacy-sync archived (archived_at set in DB; all other entities live), and + the subsequent index turn did NOT list it. +7. Browse: the full index listed the four live entities from the directory. +8. Private confidence: "keep it between us" about the harbor budget went to user_memory + (user_id=owner), explicitly NOT to the shared entity, and the brain said so. -``` -• update_user_memory(task=Remember that Alice is building Harbor, a Postgres-backed - job queue in Rust ... prefers terse answers with no bullet lists.) -• append_file(path=notes/harbor.md, content=Project: Harbor - a Postgres-backed job - queue in Rust. Decision: Use PostgreSQL advisory locks instead of SELECT FOR UPDATE - SKIP LOCKED because workers are long-lived., unique=True) -``` +Model behavior notes (findings, not failures): +- The note template included a literal "**Date recorded:** Not provided" line in turn 1 - + harmless, but the model narrates missing dates rather than omitting them. +- "Numbers before prose" (stored as a user memory in turn 8) visibly shaped later answers + across channels, including MCP - the preference survived and applied. +- The profile store was not used in this run (name-shaped facts never came up); the + preference went to user_memory, which is the right home for it. -The recall session, a fresh session id with nothing in context, answered: "You chose PostgreSQL advisory locks instead of `SELECT FOR UPDATE SKIP LOCKED` because Harbor's workers are long-lived." Final print: `notes/harbor.md (170 bytes)`. - -The note body is model-written, so byte counts and wording vary run to run. What holds is the shape: the capture session writes, the recall session answers from the store, and a second process sees both. - -### second_brain.py +## /mcp live proof (4 turns, in-process fastmcp client) **Status:** PASS -**Description:** Serving run. `python second_brain.py` serves REST on `/` and MCP on `/mcp`. Verified over REST with three callers, and over MCP with a real `fastmcp.Client`. +The client saw the eight built-in AgentOS tools. Each run_agent call was a fresh session +(no session echo - exactly the MCP shape), and the pinned user_id carried the brain: -**Result:** `GET /config` reports `agents: [('second-brain', 'Second Brain')]`, the id MCP clients address. Three runs against `POST /agents/second-brain/runs`: +1. "Where does harbor stand?" recalled the REST-era state (v0.1 shipped, Postgres over + DynamoDB) across the channel boundary. +2. A write from MCP: Tom Alvarez / load tests, recorded with a note pointer + ("load testing: Tom Alvarez will run the tests next week - see note"). +3. A THIRD fresh session recalled that write. +4. The why-question over MCP answered from the note's reasoning. -``` -user_id=alice@example.com -> "You're building Harbor, a Postgres-backed job queue in Rust, - using advisory locks for its long-lived workers." -user_id=riya -> "I don't have any projects recorded for you yet." -no user_id -> • list_files -> "Error: this agent's files require user_id for - this run and none was provided." - "I couldn't access your note files because no user_id was provided." -``` +## test.py -The second and third are the point: notes are per user, and the templated namespace `brain/{user_id}` fails closed rather than falling back to a shared store. +**Status:** PASS -Over MCP the client sees the eight built-in AgentOS tools, and `run_agent` reaches the same brain. `run_agent(agent_id="second-brain", user_id="dana@example.com", ...)` both read the store and wrote a new decision that a later REST call read back. `user_id` is the load-bearing argument here: an MCP client that omits it gets the fail-closed error above rather than someone else's notes. +Capture session filed quill (advisory locks over SELECT FOR UPDATE SKIP LOCKED) with the +reasoning in notes/quill.md; a fresh recall session answered tersely and correctly +("advisory locks ... because Quill's workers are long-lived"). The shared brain listing +shows both notes/harbor.md (with its dated updates) and notes/quill.md. -`db_file` is relative, so run the CLI and the server from the example folder or they land in different stores. +## 2026-07-26 — add_datetime_to_context, cold-start proof ---- +**Status:** PASS + +The "**Date recorded:** Not provided" line above was the agent having no clock: the +instructions ask for dated notes, and until an entity fact renders its "(as of ...)" date +there is nothing in the prompt to date them with. Two prior runs produced headings like +"## Date not provided - architecture". + +With `add_datetime_to_context=True` on the agent, a fresh db and one turn ("we picked +Postgres over Dynamo for the radar ingest queue") produced `notes/radar-ingest-queue.md` +with the heading `## 2026-07-26 - Postgres over Dynamo`. Dated on the first turn, with no +entity in the store to borrow a date from. + +## 2026-07-26 — the pinned user_id does NOT survive an MCP client that fills it + +**Status:** FAIL (identity, deployment-shaped - see the note in second_brain.py) + +Measured in-process against the real `/mcp` app: `run_agent` advertises an optional +`user_id`, and a value supplied there overrides `Agent(user_id="owner")`. One call with +`user_id="claude-desktop"` wrote its session and its user memory under `claude-desktop`, +not `owner`. Entities are global so they stay shared; profile and user memory fork per +host. Unauthenticated `/mcp` is exactly where host models fill the field. + +## 2026-07-26 — four turns after the review fixes (fresh db, gpt-5.6) + +**Status:** PASS + +| Turn | Expected | Result | +|---|---|---| +| Decision + people + systems in one message | note written dated, entity indexed with a pointer, both links reciprocal | `write_file notes/quill.md` (`## 2026-07-26 - Concurrency control decision`), `remember_about(Quill, facts=["concurrency control: Postgres advisory locks... - see note"])`, `link_entities` twice | +| "Priya has moved off Quill" | the old relationship retired, the new one recorded | `forget("Priya Raman", "runs -> Quill")` **removed the edge on both rows**, then `link_entities(Priya, works_on, Billing rewrite)`. Quill's row keeps only `backed_by -> postgresql` | +| "Between us: the Nimbus renewal is shaky" | private memory, no shared entity | `update_user_memory(...)` only; no Nimbus entity exists in the store | +| "What is the state of Quill, and why advisory locks?" | read the note, answer from it | `search_entities` → `read_file(notes/quill.md)` → answered with the reasoning, and volunteered that no replacement lead is recorded | + +The retirement turn is the one that mattered: on the pre-fix code the same call +matched nothing it could act on, or matched two candidates it printed +identically. Model behavior note: turn 2 rewrote the whole note with +`write_file` rather than `replace_lines` - allowed, and the note stayed correct. diff --git a/cookbook/examples/second_brain/second_brain.py b/cookbook/examples/second_brain/second_brain.py index 147777900f..1cd8e6e57e 100644 --- a/cookbook/examples/second_brain/second_brain.py +++ b/cookbook/examples/second_brain/second_brain.py @@ -2,8 +2,26 @@ Second Brain - Memory You Own, Behind Your Own MCP Server ========================================================= A private agent that remembers what you are building: durable notes in its own -filesystem, plus what it learns about how you work. It is also an MCP server, so -your AI apps (claude, chatgpt, claude code) can read and write the same brain. +filesystem, an entity graph over the people and projects around you, and what +it learns about how you work. It is also an MCP server, so your AI apps +(claude, chatgpt, claude code) can read and write the same brain. + +The stores split the work: +- Notes (FileSystem) hold the content: decisions with their reasoning, running + documents, anything longer than a line. +- Entities index the world: people, projects, systems - one-line current + values, links, and a note pointer to where the detail lives. +- Profile and memory hold the self: who you are and how you like to work. + +Identity is pinned (user_id below): sessions do not thread over MCP and an +unauthenticated /mcp call carries no user, so the personal brain names its +owner once and calls that name nobody land on the same brain. + +One caveat, measured rather than assumed: /mcp's run_agent takes an optional +user_id, and a host that fills it wins over the pin - that run's profile and +user memory go to whatever it sent. Entities are global, so the world half of +the brain is shared either way. Run /mcp behind auth (the JWT subject then wins +over both) if your client volunteers a user_id. Running this file serves the AgentOS on http://localhost:7777 MCP Server on http://localhost:7777/mcp @@ -12,13 +30,29 @@ MCP Server on http://localhost:7777/mcp from agno.agent import Agent from agno.db.sqlite import SqliteDb from agno.fs import FileSystem +from agno.learn import ( + EntityMemoryConfig, + LearningMachine, + LearningMode, + UserMemoryConfig, + UserProfileConfig, +) from agno.os import AgentOS # --------------------------------------------------------------------------- # One database for agent sessions, learning, notes, traces, metrics, etc. +# Shared world, private self: notes live in the same shared namespace as the +# entities they document; profile and memory stay per-user. # --------------------------------------------------------------------------- db = SqliteDb(db_file="tmp/second_brain.db") -notes = FileSystem(db, namespace="brain/{user_id}") +notes = FileSystem(db, namespace="brain") + +brain = LearningMachine( + db=db, + user_profile=UserProfileConfig(mode=LearningMode.AGENTIC), # private to each person + user_memory=UserMemoryConfig(mode=LearningMode.AGENTIC), # private to each person + entity_memory=EntityMemoryConfig(namespace="global"), # shared by the team +) # --------------------------------------------------------------------------- # Create the Agent @@ -28,16 +62,48 @@ second_brain = Agent( id="second-brain", name="Second Brain", model="openai:gpt-5.6", + learning=brain, tools=[notes.tools()], instructions=[ - "You're a second brain for: {user_id}.", - "Remember what they're building. Keep one note per project in " - "notes/.md and append decisions as they are made.", - "Answer according to their taste. In under 3 sentences unless they ask for more.", + "You are a second brain: you hold what your owner is building and thinking, " + "and you answer from what you hold.", + # One claim, one home. Notes hold the content; entities are the index over it. + "One claim, one home. Notes hold the content; entities are the index over it:", + "- Reasoning, wording, anything longer than a line goes in the note " + "(notes/.md), dated, and only in the note.", + "- On the entity: names, links, and one-line current values you expect to be " + "replaced - with note='notes/.md' whenever the detail lives there. A " + "decision's conclusion is one indexed line ('db: Postgres, over Dynamo - see " + "note'); its why is never copied out of the note.", + "- It happened on a date and next month it is history: that is an event. " + "Positions and opinions are events, not facts.", + "- Corrections replace, they never accumulate: state the new fact (the stale " + "one is retired automatically), and fix the note line with replace_lines in " + "the same turn. Never append a contradiction.", + "- Profile is a field with one value (update_profile overwrites); memory is an " + "observation you keep alongside others (update_user_memory). Standing " + "instructions are rules to obey, not observations to narrate.", + "- Confidences stay private: something shared in confidence about the world " + "goes to user memory, never to a shared entity - and say so when you file one.", + "- What you file about other people is your judgement, and the test is whether " + "your owner would file it: what they told you to remember, and what bears on " + "the work. Not a colleague's health, pay, or family, mentioned in passing and " + "never asked to be kept - those you use in the conversation and let go.", + "Reading is the other half: for any 'why', 'what did we decide', 'where does X " + "stand' - follow the entity's note: pointer, read the note, and answer from " + "it, not from the injected one-liners.", + "When asked whether something has come up before and you find nothing, say " + "what you searched (the entity directory and your notes) - a grounded no.", + "Answer in under 3 sentences unless asked for more.", notes.instructions(), ], - enable_agentic_memory=True, + # The personal brain pins identity: every channel, MCP included, lands here. + user_id="owner", add_history_to_context=True, + # A brain that cannot date its notes cannot tell July's truth from March's, + # and the instructions above ask for dated notes. Without this the agent has + # no clock and writes "Date not provided" until the first fact gives it one. + add_datetime_to_context=True, ) # --------------------------------------------------------------------------- diff --git a/cookbook/examples/second_brain/test.py b/cookbook/examples/second_brain/test.py index 87c83727a3..b58714b390 100644 --- a/cookbook/examples/second_brain/test.py +++ b/cookbook/examples/second_brain/test.py @@ -1,8 +1,8 @@ """ Second Brain - CLI ==================================== -Runs the second_brain without starting the server: capture a -decision in one session, then recall in a new session. +Runs the second_brain without starting the server: capture a decision in one +session, then recall it in a new session that shares nothing but the brain. """ from uuid import uuid4 @@ -10,10 +10,8 @@ from uuid import uuid4 from second_brain import notes, second_brain # --------------------------------------------------------------------------- -# Create the run: one user, two sessions that share nothing +# Two sessions that share nothing; identity is pinned on the agent itself. # --------------------------------------------------------------------------- -# Notes live under brain/{user_id}, so every run needs a user_id. -USER_ID = "alice@example.com" CAPTURE_SESSION = f"capture-{uuid4().hex[:8]}" RECALL_SESSION = f"recall-{uuid4().hex[:8]}" @@ -21,28 +19,23 @@ RECALL_SESSION = f"recall-{uuid4().hex[:8]}" # Run: capture a decision, then ask for it back in the other session # --------------------------------------------------------------------------- if __name__ == "__main__": - print(f"User: {USER_ID}") - print("\n--- Session 1: capture a decision ---\n") second_brain.print_response( - "I am building Harbor, a Postgres-backed job queue in Rust. I picked " + "I am building quill, a Postgres-backed job queue in Rust. I picked " "advisory locks over SELECT FOR UPDATE SKIP LOCKED because our workers " "are long-lived. I want terse answers, no bullet lists.", - user_id=USER_ID, session_id=CAPTURE_SESSION, stream=True, ) print("\n--- Session 2: a new session, nothing in context ---\n") second_brain.print_response( - "What did I decide about locking in Harbor, and why?", - user_id=USER_ID, + "What did I decide about locking in quill, and why?", session_id=RECALL_SESSION, stream=True, ) - print("\n--- Files in this user's brain ---\n") - brain = notes.resolve(user_id=USER_ID) - for meta in brain.list(): + print("\n--- Files in the brain ---\n") + for meta in notes.list(): print(f" {meta.path} ({meta.size_bytes} bytes)\n") - print(brain.read(meta.path)) + print(notes.read(meta.path)) diff --git a/libs/agno/agno/agent/_init.py b/libs/agno/agno/agent/_init.py index 739ca845a1..3022e35ea9 100644 --- a/libs/agno/agno/agent/_init.py +++ b/libs/agno/agno/agent/_init.py @@ -170,6 +170,12 @@ def set_learning_machine(agent: Agent) -> None: agent.learning.db = agent.db if agent.learning.model is None: agent.learning.model = agent.model + if ( + agent.learning.learned_knowledge + and agent.learning.knowledge is None + and getattr(agent, "knowledge", None) is not None + ): + agent.learning.knowledge = agent.knowledge agent._learning = agent.learning # PROPOSE/HITL modes need chat history for multi-turn confirmation diff --git a/libs/agno/agno/agent/_managers.py b/libs/agno/agno/agent/_managers.py index 4a0dc9a2f8..f2cb0318b3 100644 --- a/libs/agno/agno/agent/_managers.py +++ b/libs/agno/agno/agent/_managers.py @@ -6,6 +6,7 @@ from asyncio import CancelledError, Task, create_task from concurrent.futures import Future from typing import ( TYPE_CHECKING, + Any, List, Optional, ) @@ -393,6 +394,7 @@ def process_learnings( run_messages: RunMessages, session: AgentSession, user_id: Optional[str], + run_context: Optional[Any] = None, ) -> Optional[RunMetrics]: """Process learnings from conversation (runs in background thread).""" if agent._learning is None: @@ -412,6 +414,10 @@ def process_learnings( agent_id=agent.id, team_id=agent.team_id, run_metrics=collector, + run_context=run_context, + metadata=getattr(run_context, "metadata", None), + dependencies=getattr(run_context, "dependencies", None), + session_state=getattr(run_context, "session_state", None), ) log_debug("Learning extraction completed.") except Exception as e: @@ -424,6 +430,7 @@ async def aprocess_learnings( run_messages: RunMessages, session: AgentSession, user_id: Optional[str], + run_context: Optional[Any] = None, ) -> Optional[RunMetrics]: """Async process learnings from conversation.""" if agent._learning is None: @@ -442,6 +449,10 @@ async def aprocess_learnings( agent_id=agent.id, team_id=agent.team_id, run_metrics=collector, + run_context=run_context, + metadata=getattr(run_context, "metadata", None), + dependencies=getattr(run_context, "dependencies", None), + session_state=getattr(run_context, "session_state", None), ) log_debug("Learning extraction completed.") except Exception as e: @@ -455,6 +466,7 @@ async def astart_learning_task( session: AgentSession, user_id: Optional[str], existing_task: Optional[Task] = None, + run_context: Optional[Any] = None, ) -> Optional[Task]: """Start learning extraction as async task. @@ -485,6 +497,7 @@ async def astart_learning_task( run_messages=run_messages, session=session, user_id=user_id, + run_context=run_context, ) ) @@ -497,6 +510,7 @@ def start_learning_future( session: AgentSession, user_id: Optional[str], existing_future: Optional[Future] = None, + run_context: Optional[Any] = None, ) -> Optional[Future]: """Start learning extraction in background thread. @@ -523,6 +537,7 @@ def start_learning_future( run_messages=run_messages, session=session, user_id=user_id, + run_context=run_context, ) return None diff --git a/libs/agno/agno/agent/_messages.py b/libs/agno/agno/agent/_messages.py index 7b8b5e6cba..8f53511c21 100644 --- a/libs/agno/agno/agent/_messages.py +++ b/libs/agno/agno/agent/_messages.py @@ -103,12 +103,41 @@ def format_message_with_state_variables( # --------------------------------------------------------------------------- +def _learning_message_text(input: Any) -> Optional[str]: + """Flatten the run input to text for message-scoped learning recall. + + Textual parts only: multimodal content parts contribute their "text" + values, and non-text parts are skipped rather than leaking Python reprs + into the recall terms. + """ + if input is None: + return None + if isinstance(input, str): + return input or None + if isinstance(input, Message): + return _learning_message_text(input.content) + if isinstance(input, list): + parts = [_learning_message_text(item) for item in input] + return "\n".join(part for part in parts if part) or None + if isinstance(input, dict): + # A multimodal content part or a payload dict: keep its string values. + parts = [value for value in input.values() if isinstance(value, str)] + return "\n".join(part for part in parts if part) or None + if isinstance(input, BaseModel): + try: + return _learning_message_text(input.model_dump()) + except Exception: + return None + return str(input) + + def get_system_message( agent: Agent, session: AgentSession, run_context: Optional[RunContext] = None, tools: Optional[List[Union[Function, dict]]] = None, add_session_state_to_context: Optional[bool] = None, + input: Optional[Any] = None, ) -> Optional[Message]: """Return the system message for the Agent. @@ -396,15 +425,25 @@ def get_system_message( "You should ALWAYS prefer information from this conversation over the past summary.\n\n" ) - # 3.3.12 then add learnings to the system prompt + # 3.3.12 then add learnings to the system prompt: guidance (how to use the + # tools) plus data (what was recalled), concatenated here so the automatic + # door renders exactly what the manual door's instructions() + build_context() + # would. if agent._learning is not None and agent.add_learnings_to_context: + learning_guidance = agent._learning._framework_instructions() learning_context = agent._learning.build_context( user_id=user_id, session_id=session.session_id if session else None, agent_id=agent.id, + message=_learning_message_text(input), + run_context=run_context, + metadata=run_context.metadata if run_context else None, + dependencies=run_context.dependencies if run_context else None, + session_state=session_state, ) - if learning_context: - system_message_content += learning_context + "\n" + learning_block = "\n".join(part for part in (learning_guidance, learning_context) if part) + if learning_block: + system_message_content += learning_block + "\n" # 3.3.13 then add search_knowledge instructions to the system prompt _resolved_knowledge = _get_resolved_knowledge(agent, run_context) @@ -456,6 +495,7 @@ async def aget_system_message( run_context: Optional[RunContext] = None, tools: Optional[List[Union[Function, dict]]] = None, add_session_state_to_context: Optional[bool] = None, + input: Optional[Any] = None, ) -> Optional[Message]: """Return the system message for the Agent. @@ -747,15 +787,22 @@ async def aget_system_message( "You should ALWAYS prefer information from this conversation over the past summary.\n\n" ) - # 3.3.12 then add learnings to the system prompt + # 3.3.12 then add learnings to the system prompt (see the sync twin) if agent._learning is not None and agent.add_learnings_to_context: + learning_guidance = agent._learning._framework_instructions() learning_context = await agent._learning.abuild_context( user_id=user_id, session_id=session.session_id if session else None, agent_id=agent.id, + message=_learning_message_text(input), + run_context=run_context, + metadata=run_context.metadata if run_context else None, + dependencies=run_context.dependencies if run_context else None, + session_state=session_state, ) - if learning_context: - system_message_content += learning_context + "\n" + learning_block = "\n".join(part for part in (learning_guidance, learning_context) if part) + if learning_block: + system_message_content += learning_block + "\n" # 3.3.13 then add search_knowledge instructions to the system prompt _resolved_knowledge = _get_resolved_knowledge(agent, run_context) @@ -1205,6 +1252,7 @@ def get_run_messages( run_context=run_context, tools=tools, add_session_state_to_context=add_session_state_to_context, + input=input, ) if system_message is not None: run_messages.system_message = system_message @@ -1410,6 +1458,7 @@ async def aget_run_messages( run_context=run_context, tools=tools, add_session_state_to_context=add_session_state_to_context, + input=input, ) if system_message is not None: run_messages.system_message = system_message diff --git a/libs/agno/agno/agent/_run.py b/libs/agno/agno/agent/_run.py index a779aa87ea..672d7a91e4 100644 --- a/libs/agno/agno/agent/_run.py +++ b/libs/agno/agno/agent/_run.py @@ -509,6 +509,7 @@ def _run( session=agent_session, user_id=user_id, existing_future=learning_future, + run_context=run_context, ) # Start cultural knowledge creation in background thread @@ -921,6 +922,7 @@ def _run_stream( session=agent_session, user_id=user_id, existing_future=learning_future, + run_context=run_context, ) # Start cultural knowledge creation in background thread @@ -1647,6 +1649,7 @@ async def _arun( session=agent_session, user_id=user_id, existing_task=learning_task, + run_context=run_context, ) # Start cultural knowledge creation as a background task (runs concurrently with the main execution) @@ -2329,6 +2332,7 @@ async def _arun_stream( session=agent_session, user_id=user_id, existing_task=learning_task, + run_context=run_context, ) # Start cultural knowledge creation as a background task (runs concurrently with the main execution) diff --git a/libs/agno/agno/agent/_tools.py b/libs/agno/agno/agent/_tools.py index 3154e52c2f..9b7612e839 100644 --- a/libs/agno/agno/agent/_tools.py +++ b/libs/agno/agno/agent/_tools.py @@ -170,6 +170,7 @@ def get_tools( user_id=user_id, session_id=session.session_id if session else None, agent_id=agent.id, + run_context=run_context, ) agent_tools.extend(learning_tools) @@ -302,6 +303,7 @@ async def aget_tools( user_id=user_id, session_id=session.session_id if session else None, agent_id=agent.id, + run_context=run_context, ) agent_tools.extend(learning_tools) diff --git a/libs/agno/agno/agent/agent.py b/libs/agno/agno/agent/agent.py index ac5a5cb608..d526b54f62 100644 --- a/libs/agno/agno/agent/agent.py +++ b/libs/agno/agno/agent/agent.py @@ -116,7 +116,10 @@ class Agent: # --- Agent Memory --- # Memory manager to use for this agent memory_manager: Optional[MemoryManager] = None - # Enable the agent to manage memories of the user + # Enable the agent to manage memories of the user. + # Do not combine with a LearningMachine that has a user_memory store: both + # register a tool named update_user_memory, tool parsing keeps the first + # name it sees, and the learning store's tool is dropped without a word. enable_agentic_memory: bool = False # If True, the agent creates/updates user memories at the end of runs update_memory_on_run: bool = False @@ -872,6 +875,7 @@ class Agent: run_context: Optional[RunContext] = None, tools: Optional[List[Union[Function, dict]]] = None, add_session_state_to_context: Optional[bool] = None, + input: Optional[Any] = None, ) -> Optional[Message]: return _messages.get_system_message( self, @@ -879,6 +883,7 @@ class Agent: run_context=run_context, tools=tools, add_session_state_to_context=add_session_state_to_context, + input=input, ) async def aget_system_message( @@ -887,6 +892,7 @@ class Agent: run_context: Optional[RunContext] = None, tools: Optional[List[Union[Function, dict]]] = None, add_session_state_to_context: Optional[bool] = None, + input: Optional[Any] = None, ) -> Optional[Message]: return await _messages.aget_system_message( self, @@ -894,6 +900,7 @@ class Agent: run_context=run_context, tools=tools, add_session_state_to_context=add_session_state_to_context, + input=input, ) def get_relevant_docs_from_knowledge( diff --git a/libs/agno/agno/db/base.py b/libs/agno/agno/db/base.py index a0a20974f2..6ca1d991a4 100644 --- a/libs/agno/agno/db/base.py +++ b/libs/agno/agno/db/base.py @@ -1032,6 +1032,52 @@ class BaseDb(ABC): """ raise NotImplementedError + def search_learnings( + self, + query: str, + learning_type: Optional[str] = None, + user_id: Optional[str] = None, + agent_id: Optional[str] = None, + team_id: Optional[str] = None, + workflow_id: Optional[str] = None, + session_id: Optional[str] = None, + namespace: Optional[str] = None, + entity_id: Optional[str] = None, + entity_type: Optional[str] = None, + limit: Optional[int] = None, + ) -> List[Dict[str, Any]]: + """Search learning records by text query, filtered server-side. + + The query is matched case-insensitively against the stored content, in + both its space and underscore forms ("sarah chen" finds "sarah_chen"). + Results are ordered by updated_at descending. + + Backends without a server-side search implementation keep this default, + which raises NotImplementedError - callers fall back to their + client-side scan then. Implementations RAISE on database errors rather + than returning an empty list: a dialect-wrong query must never present + as an empty store. + + Args: + query: Text to search for in the stored content. + learning_type: Filter by learning type. + user_id: Filter by user ID. + agent_id: Filter by agent ID. + team_id: Filter by team ID. + workflow_id: Filter by workflow ID. Note: upsert_learning does not + currently populate the workflow_id column, so this filter only + matches rows written by an external writer that sets it. + session_id: Filter by session ID. + namespace: Filter by namespace ('user', 'global', or custom). + entity_id: Filter by entity ID. + entity_type: Filter by entity type. + limit: Maximum number of records to return. + + Returns: + List of matching learning records, most recently updated first. + """ + raise NotImplementedError + def get_learning_by_id(self, id: str) -> Optional[Dict[str, Any]]: """Retrieve a single learning record by its primary key. @@ -2001,6 +2047,52 @@ class AsyncBaseDb(ABC): """ raise NotImplementedError + async def search_learnings( + self, + query: str, + learning_type: Optional[str] = None, + user_id: Optional[str] = None, + agent_id: Optional[str] = None, + team_id: Optional[str] = None, + workflow_id: Optional[str] = None, + session_id: Optional[str] = None, + namespace: Optional[str] = None, + entity_id: Optional[str] = None, + entity_type: Optional[str] = None, + limit: Optional[int] = None, + ) -> List[Dict[str, Any]]: + """Async search learning records by text query, filtered server-side. + + The query is matched case-insensitively against the stored content, in + both its space and underscore forms ("sarah chen" finds "sarah_chen"). + Results are ordered by updated_at descending. + + Backends without a server-side search implementation keep this default, + which raises NotImplementedError - callers fall back to their + client-side scan then. Implementations RAISE on database errors rather + than returning an empty list: a dialect-wrong query must never present + as an empty store. + + Args: + query: Text to search for in the stored content. + learning_type: Filter by learning type. + user_id: Filter by user ID. + agent_id: Filter by agent ID. + team_id: Filter by team ID. + workflow_id: Filter by workflow ID. Note: upsert_learning does not + currently populate the workflow_id column, so this filter only + matches rows written by an external writer that sets it. + session_id: Filter by session ID. + namespace: Filter by namespace ('user', 'global', or custom). + entity_id: Filter by entity ID. + entity_type: Filter by entity type. + limit: Maximum number of records to return. + + Returns: + List of matching learning records, most recently updated first. + """ + raise NotImplementedError + async def get_learning_by_id(self, id: str) -> Optional[Dict[str, Any]]: """Async retrieve a single learning record by its primary key. diff --git a/libs/agno/agno/db/postgres/async_postgres.py b/libs/agno/agno/db/postgres/async_postgres.py index a12277a385..5f6f76b313 100644 --- a/libs/agno/agno/db/postgres/async_postgres.py +++ b/libs/agno/agno/db/postgres/async_postgres.py @@ -29,7 +29,7 @@ from agno.db.schemas.service_accounts import ( resolve_service_account_sort_column, validate_service_account_update, ) -from agno.db.utils import deserialize_session, deserialize_sessions, json_serializer +from agno.db.utils import deserialize_session, deserialize_sessions, json_serializer, learning_search_patterns from agno.run.base import RunStatus from agno.session import AgentSession, Session, TeamSession, WorkflowSession from agno.utils.log import log_debug, log_error, log_info, log_warning @@ -3178,6 +3178,73 @@ class AsyncPostgresDb(AsyncBaseDb): log_debug(f"Error getting learnings: {e}") return [] + async def search_learnings( + self, + query: str, + learning_type: Optional[str] = None, + user_id: Optional[str] = None, + agent_id: Optional[str] = None, + team_id: Optional[str] = None, + workflow_id: Optional[str] = None, + session_id: Optional[str] = None, + namespace: Optional[str] = None, + entity_id: Optional[str] = None, + entity_type: Optional[str] = None, + limit: Optional[int] = None, + ) -> List[Dict[str, Any]]: + """Async search learning records by text query. See AsyncBaseDb.search_learnings. + + The content column is JSONB, which has no ILIKE operator - it is cast + to TEXT first (the same shape get_user_memories uses for + search_content). The query matches in both its space and underscore + forms. Errors are raised, never swallowed. + """ + try: + table = await self._get_table(table_type="learnings") + if table is None: + return [] + + patterns = learning_search_patterns(query) + if not patterns: + return [] + + async with self.async_session_factory() as sess: + stmt = select(table) + + if learning_type is not None: + stmt = stmt.where(table.c.learning_type == learning_type) + if user_id is not None: + stmt = stmt.where(table.c.user_id == user_id) + if agent_id is not None: + stmt = stmt.where(table.c.agent_id == agent_id) + if team_id is not None: + stmt = stmt.where(table.c.team_id == team_id) + if workflow_id is not None: + stmt = stmt.where(table.c.workflow_id == workflow_id) + if session_id is not None: + stmt = stmt.where(table.c.session_id == session_id) + if namespace is not None: + stmt = stmt.where(table.c.namespace == namespace) + if entity_id is not None: + stmt = stmt.where(table.c.entity_id == entity_id) + if entity_type is not None: + stmt = stmt.where(table.c.entity_type == entity_type) + + content_text = func.cast(table.c.content, postgresql.TEXT) + stmt = stmt.where(or_(*[content_text.ilike(pattern, escape="\\") for pattern in patterns])) + + stmt = stmt.order_by(table.c.updated_at.desc().nulls_last()) + if limit is not None: + stmt = stmt.limit(limit) + + result = await sess.execute(stmt) + rows = result.fetchall() + return [dict(row._mapping) for row in rows] + + except Exception as e: + log_error(f"Error searching learnings: {e}") + raise e + async def get_learning_by_id(self, id: str) -> Optional[Dict[str, Any]]: try: table = await self._get_table(table_type="learnings") diff --git a/libs/agno/agno/db/postgres/postgres.py b/libs/agno/agno/db/postgres/postgres.py index 235ad0d8b6..f0641ed6b6 100644 --- a/libs/agno/agno/db/postgres/postgres.py +++ b/libs/agno/agno/db/postgres/postgres.py @@ -38,7 +38,7 @@ from agno.db.schemas.service_accounts import ( resolve_service_account_sort_column, validate_service_account_update, ) -from agno.db.utils import deserialize_session, deserialize_sessions, json_serializer +from agno.db.utils import deserialize_session, deserialize_sessions, json_serializer, learning_search_patterns from agno.run.base import RunStatus from agno.session import AgentSession, Session, TeamSession, WorkflowSession from agno.utils.log import log_debug, log_error, log_info, log_warning @@ -4703,6 +4703,72 @@ class PostgresDb(BaseDb): log_debug(f"Error getting learnings: {e}") return [] + def search_learnings( + self, + query: str, + learning_type: Optional[str] = None, + user_id: Optional[str] = None, + agent_id: Optional[str] = None, + team_id: Optional[str] = None, + workflow_id: Optional[str] = None, + session_id: Optional[str] = None, + namespace: Optional[str] = None, + entity_id: Optional[str] = None, + entity_type: Optional[str] = None, + limit: Optional[int] = None, + ) -> List[Dict[str, Any]]: + """Search learning records by text query. See BaseDb.search_learnings. + + The content column is JSONB, which has no ILIKE operator - it is cast + to TEXT first (the same shape get_user_memories uses for + search_content). The query matches in both its space and underscore + forms. Errors are raised, never swallowed. + """ + try: + table = self._get_table(table_type="learnings") + if table is None: + return [] + + patterns = learning_search_patterns(query) + if not patterns: + return [] + + with self.Session() as sess: + stmt = select(table) + + if learning_type is not None: + stmt = stmt.where(table.c.learning_type == learning_type) + if user_id is not None: + stmt = stmt.where(table.c.user_id == user_id) + if agent_id is not None: + stmt = stmt.where(table.c.agent_id == agent_id) + if team_id is not None: + stmt = stmt.where(table.c.team_id == team_id) + if workflow_id is not None: + stmt = stmt.where(table.c.workflow_id == workflow_id) + if session_id is not None: + stmt = stmt.where(table.c.session_id == session_id) + if namespace is not None: + stmt = stmt.where(table.c.namespace == namespace) + if entity_id is not None: + stmt = stmt.where(table.c.entity_id == entity_id) + if entity_type is not None: + stmt = stmt.where(table.c.entity_type == entity_type) + + content_text = func.cast(table.c.content, postgresql.TEXT) + stmt = stmt.where(or_(*[content_text.ilike(pattern, escape="\\") for pattern in patterns])) + + stmt = stmt.order_by(table.c.updated_at.desc().nulls_last()) + if limit is not None: + stmt = stmt.limit(limit) + + result = sess.execute(stmt).fetchall() + return [dict(row._mapping) for row in result] + + except Exception as e: + log_error(f"Error searching learnings: {e}") + raise e + def get_learning_by_id(self, id: str) -> Optional[Dict[str, Any]]: try: table = self._get_table(table_type="learnings") diff --git a/libs/agno/agno/db/sqlite/sqlite.py b/libs/agno/agno/db/sqlite/sqlite.py index 6184b6d7b5..e9ee770072 100644 --- a/libs/agno/agno/db/sqlite/sqlite.py +++ b/libs/agno/agno/db/sqlite/sqlite.py @@ -43,6 +43,7 @@ from agno.db.utils import ( deserialize_session, deserialize_session_json_fields, deserialize_sessions, + learning_search_patterns, serialize_session_json_fields, ) from agno.run.base import RunStatus @@ -4547,6 +4548,69 @@ class SqliteDb(BaseDb): log_debug(f"Error getting learnings: {e}") return [] + def search_learnings( + self, + query: str, + learning_type: Optional[str] = None, + user_id: Optional[str] = None, + agent_id: Optional[str] = None, + team_id: Optional[str] = None, + workflow_id: Optional[str] = None, + session_id: Optional[str] = None, + namespace: Optional[str] = None, + entity_id: Optional[str] = None, + entity_type: Optional[str] = None, + limit: Optional[int] = None, + ) -> List[Dict[str, Any]]: + """Search learning records by text query. See BaseDb.search_learnings. + + The query matches the content column case-insensitively in both its + space and underscore forms. Errors are raised, never swallowed. + """ + try: + table = self._get_table(table_type="learnings") + if table is None: + return [] + + patterns = learning_search_patterns(query) + if not patterns: + return [] + + with self.Session() as sess: + stmt = select(table) + + if learning_type is not None: + stmt = stmt.where(table.c.learning_type == learning_type) + if user_id is not None: + stmt = stmt.where(table.c.user_id == user_id) + if agent_id is not None: + stmt = stmt.where(table.c.agent_id == agent_id) + if team_id is not None: + stmt = stmt.where(table.c.team_id == team_id) + if workflow_id is not None: + stmt = stmt.where(table.c.workflow_id == workflow_id) + if session_id is not None: + stmt = stmt.where(table.c.session_id == session_id) + if namespace is not None: + stmt = stmt.where(table.c.namespace == namespace) + if entity_id is not None: + stmt = stmt.where(table.c.entity_id == entity_id) + if entity_type is not None: + stmt = stmt.where(table.c.entity_type == entity_type) + + stmt = stmt.where(or_(*[table.c.content.ilike(pattern, escape="\\") for pattern in patterns])) + + stmt = stmt.order_by(table.c.updated_at.desc().nulls_last()) + if limit is not None: + stmt = stmt.limit(limit) + + results = sess.execute(stmt).fetchall() + return [dict(row._mapping) for row in results] + + except Exception as e: + log_error(f"Error searching learnings: {e}") + raise e + def get_learning_by_id(self, id: str) -> Optional[Dict[str, Any]]: try: table = self._get_table(table_type="learnings") diff --git a/libs/agno/agno/db/utils.py b/libs/agno/agno/db/utils.py index d58bd7d97e..3a165767d5 100644 --- a/libs/agno/agno/db/utils.py +++ b/libs/agno/agno/db/utils.py @@ -189,6 +189,68 @@ def get_sort_value(record: Dict[str, Any], sort_by: str) -> Any: return value +def learning_search_patterns(query: str) -> List[str]: + """Build the ILIKE patterns for a learnings text search. + + Three properties, each load-bearing: + + - The stored content mixes display names ("Sarah Chen") with slugs + ("sarah_chen"), so runs of spaces and underscores in the query become the + single-char LIKE wildcard ``_`` - one pattern crosses both forms. + - ``%`` and ``\\`` in the query are escaped (callers pass the pattern with + ``escape="\\\\"``), so a model-authored query containing ``%`` cannot + collapse search into match-everything-by-recency. + - SQLite stores JSON with ``ensure_ascii`` escapes (``café`` is stored as + ``caf\\u00e9``), so a non-ASCII query also gets its JSON-escaped variant; + on Postgres, where ``::text`` renders real characters, that extra pattern + simply never matches. + + A query with no content beyond wildcards and whitespace yields no patterns. + + Args: + query: The text to search for. + + Returns: + Deduplicated '%...%' patterns to OR together with ILIKE (escape '\\'). + """ + import re + + stripped = query.strip() + if not re.sub(r"[%_\s]+", "", stripped): + return [] + + variants = [stripped] + # SQLite stores JSON with ensure_ascii escapes, and LIKE folds ASCII only, + # so an escape sequence never case-matches: a stored "Ος" is unreachable + # from "ΟΣ", "ος" or "οσ". No set of pre-cased whole-string variants covers + # the mixed forms, so the escape carries a wildcard per character instead - + # \\uXXXX is six characters wide - and the caller's value-scoped Python + # check (which casefolds) rejects whatever that lets through. Loose + # prefilter, precise verification, which is what this pair is for. + # + # The wildcards are carried as a sentinel until after the separator + # collapse below, which would otherwise fold a run of them into one. + wildcard = "\x00" + json_form = json.dumps(stripped, ensure_ascii=True)[1:-1] + if json_form != stripped: + variants.append(re.sub(r"\\u[0-9a-fA-F]{4}", wildcard * 6, json_form)) + + patterns: List[str] = [] + for variant in variants: + escaped = variant.replace("\\", "\\\\").replace("%", "\\%") + # Runs of separators collapse to the single-char wildcard, so one + # pattern crosses the display-name/slug boundary in both directions + # ("sarah chen", "sarah_chen", "sarah__chen"). The hyphen is one of + # them: without it "multi-tenant" could never reach a stored "multi + # tenant", and the client-side verifier - which does fold hyphens - + # was already accepting what this pattern refused to fetch. + crossed = re.sub(r"[\s_\-]+", "_", escaped).replace(wildcard, "_") + pattern = f"%{crossed}%" + if pattern not in patterns: + patterns.append(pattern) + return patterns + + class CustomJSONEncoder(json.JSONEncoder): """Custom encoder to handle non JSON serializable types.""" diff --git a/libs/agno/agno/environments/runner.py b/libs/agno/agno/environments/runner.py index c1e1ac2650..a6d08b2e86 100644 --- a/libs/agno/agno/environments/runner.py +++ b/libs/agno/agno/environments/runner.py @@ -541,6 +541,12 @@ def _writes_off_learning_config(config: Any) -> Any: for flag in vars(config_copy): if flag.startswith("agent_can_") and not flag.startswith("agent_can_search"): setattr(config_copy, flag, False) + # Entity memory's 2.8.4 revamp replaced its per-op agent_can_* flags with a + # single enable_agent_tools gating all four tools (three of which write), so + # severing writes means severing the toolset. Entity data still reaches the + # attempt through the injected context block. + if type(config_copy).__name__ == "EntityMemoryConfig" and hasattr(config_copy, "enable_agent_tools"): + config_copy.enable_agent_tools = False config_model = getattr(config_copy, "model", None) if config_model is not None and hasattr(config_model, "cache_response"): config_copy.model = _cache_off_copy(config_model) diff --git a/libs/agno/agno/fs/db.py b/libs/agno/agno/fs/db.py index 2545de9e6d..04bc478bd4 100644 --- a/libs/agno/agno/fs/db.py +++ b/libs/agno/agno/fs/db.py @@ -446,10 +446,13 @@ class DbFileSystem(BaseFS): return found def search(self, namespace: str, query: str, directory: str = "", limit: int = 10) -> List[SearchMatch]: - """Case-insensitive substring search. Correctness is owned by Python; on - Postgres ILIKE prefilters candidate rows, on SQLite every in-scope row is - scanned (SQLite's LIKE folds ASCII only, which would miss non-ASCII case - variants).""" + """Case-insensitive substring search. Correctness is owned by Python; the + SQL predicate only prefilters candidate rows. On Postgres ILIKE folds + every query; on SQLite, LIKE folds ASCII only, so the prefilter applies + to pure-ASCII queries and a non-ASCII query still scans every in-scope + row. One known gap remains on SQLite: content containing a non-ASCII + uppercase form whose lowercase is ASCII (the Kelvin sign, U+212A) is + excluded by the ASCII prefilter for the matching ASCII query.""" self._ensure_table() if not query: return [] @@ -457,6 +460,8 @@ class DbFileSystem(BaseFS): conditions = [t.c.namespace == namespace, self._directory_predicate(directory)] if self.dialect == "postgresql": conditions.append(t.c.content.icontains(query, autoescape=True)) + elif query.isascii(): + conditions.append(t.c.content.icontains(query, autoescape=True)) stmt = select(t.c.path, t.c.size_bytes, t.c.content).where(and_(*conditions)) with self.db_engine.begin() as conn: rows = conn.execute(stmt).all() diff --git a/libs/agno/agno/fs/fs.py b/libs/agno/agno/fs/fs.py index b304450dfc..172f41094a 100644 --- a/libs/agno/agno/fs/fs.py +++ b/libs/agno/agno/fs/fs.py @@ -1,9 +1,9 @@ """FileSystem: a durable, private filesystem for agents. To the agent it looks exactly like a normal filesystem toolkit; underneath it is -a pluggable ``BaseFS`` backend, database by default. Use it for the agent's -own working state: records of items already processed, decisions, progress -checkpoints. +a pluggable ``BaseFS`` backend, database by default. Use it for the agent's own +durable notes: decisions with their reasoning, running documents, working state +it will need again. Attach the tools, and compose its instructions into your own: @@ -47,40 +47,35 @@ backend with no namespace share this one store, which is the intended behavior but is rarely what a multi-tenant app wants (see the templated namespaces above). """ -_DEFAULT_INSTRUCTIONS = """You have your own private, durable filesystem you can use to persist files across \ -sessions and runs. Use it for your working state: records of items you have processed, decisions, progress \ -checkpoints, notes to your future self. - -For storing facts and memories about the user, prefer the user memory if provided and fallback to the \ -filesystem if not. +_DEFAULT_INSTRUCTIONS = """You have your own private, durable filesystem for notes that persist across \ +sessions and runs. Use it to write, maintain and re-read prose that matters later: decisions with their \ +reasoning, running documents on a topic, notes to your future self. Conventions: - - Paths are relative, like notes/decisions.md or seen/2026-07-24.md. Group related files in directories. - - To keep a record set (items already processed): store one record per line in the seen/ directory, one \ -file per date (seen/2026-07-24.md). Call check_lines(lines, directory="seen") BEFORE acting, then record the \ -new ones with append_file(..., unique=True). check_lines matches exact whole lines, and unique=True keeps \ -the log free of duplicates if a run overlaps another one. + - Paths are relative, like notes/decisions.md. Group related files in directories. + - One topic, one file. Add dated entries as things develop; a note is a living document you maintain, not \ +a log you only append to. - To correct or update part of a file, read it, then call replace_lines with the line numbers you saw. \ -Rewriting a whole file with write_file to change a few lines wastes effort and risks losing the rest. - - To find something in a large file, use search_content first: it reports the line number of each match, \ -which you can pass to read_file as start_line. - - Store extracted facts and identifiers, not raw fetched payloads. +Never append a contradiction of what a note already says: fix the wording in place, in the same turn you \ +learn it was wrong. + - To find something, use search_content first: it reports the file and first-match line, which you can \ +pass to read_file as start_line. Then answer from what the note says, not from memory. + - To retire a note that is no longer current, move_file it into an archive/ directory. Do not blank it and \ +do not overwrite it: its history may still be needed. + - Store distilled content, not raw fetched payloads. - Never store secrets, passwords, or API keys. - - Files have size limits. If a write is refused, start a new file in your partition scheme (e.g. a new \ -dated file), or delete only files you are certain are obsolete, like a date partition older than you still \ -need. Never overwrite or delete a file of records to make room if you might still need those records: \ -dropping them means you will repeat work you already did. If nothing is safely disposable, stop and report \ -that storage is full rather than evicting history.""" + - Files have size limits. If a write is refused, split the topic into smaller files or archive what is \ +finished. Never overwrite a note to make room if you might still need it.""" _READ_ONLY_INSTRUCTIONS = """You have read access to a durable filesystem. The files persist across sessions \ and runs. -Use it to look up what you have recorded: items you have processed, decisions you made, notes you left. You \ -cannot change these files - you have no tool to write, append, move, or delete. +Use it to look up what has been recorded: decisions, running documents, notes. You cannot change these \ +files - you have no tool to write, append, move, or delete. Conventions: - - Paths are relative, like notes/decisions.md or seen/2026-07-24.md. - - Use check_lines to see what is already recorded before acting.""" + - Paths are relative, like notes/decisions.md. + - Use search_content to find where something is recorded, then read_file to read it.""" def _as_backend(source: Any) -> BaseFS: @@ -445,7 +440,7 @@ class FileSystem: # Agent surface # ------------------------------------------------------------------ - def tools(self, *, read_only: bool = False, **kwargs) -> "Toolkit": + def tools(self, *, read_only: bool = False, allow_delete: bool = False, **kwargs) -> "Toolkit": """Build the toolkit for this file store. ``Agent(tools=[fs.tools()], instructions=[..., fs.instructions()])`` is the @@ -453,16 +448,23 @@ class FileSystem: system prompt stays the developer's to order and edit. Pass ``add_instructions=True`` to have the toolkit carry them instead. - ``read_only=True`` registers only ``read_file``, ``list_files``, - ``search_content`` and ``check_lines``; pair it with - ``instructions(read_only=True)``. That is the surface for a consumer agent - that consults another agent's namespace by shared name. ``**kwargs`` - forwards to ``Toolkit`` (e.g. ``include_tools``, - ``requires_confirmation_tools``). + The default surface is the notes seven: ``read_file``, ``write_file``, + ``append_file``, ``replace_lines``, ``list_files``, ``search_content``, + ``move_file``. ``allow_delete=True`` adds ``delete_file`` - destructive + is opt-in; archiving with ``move_file`` is the blessed retirement flow. + ``check_lines`` (the batched record-set membership test feed agents use) + is requested by name via ``include_tools``, which selects from the whole + surface when passed. + + ``read_only=True`` registers only ``read_file``, ``list_files`` and + ``search_content``; pair it with ``instructions(read_only=True)``. That + is the surface for a consumer agent that consults another agent's + namespace by shared name. ``**kwargs`` forwards to ``Toolkit`` (e.g. + ``include_tools``, ``requires_confirmation_tools``). """ from agno.fs.toolkit import FileSystemTools - return FileSystemTools(fs=self, read_only=read_only, **kwargs) + return FileSystemTools(fs=self, read_only=read_only, allow_delete=allow_delete, **kwargs) @staticmethod def instructions(read_only: bool = False) -> str: diff --git a/libs/agno/agno/fs/toolkit.py b/libs/agno/agno/fs/toolkit.py index ca832b32bd..c87648028b 100644 --- a/libs/agno/agno/fs/toolkit.py +++ b/libs/agno/agno/fs/toolkit.py @@ -23,9 +23,11 @@ genuinely needs both FileSystem and a local workspace, wrap one in a sub-agent. import asyncio import json +from contextlib import AsyncExitStack, asynccontextmanager from datetime import datetime, timezone from fnmatch import fnmatch -from typing import Dict, List, Optional, Union +from typing import Any, Dict, List, Optional, Union +from weakref import WeakKeyDictionary # Real module-level imports, never TYPE_CHECKING-only: with postponed annotations a # deferred import does not fail loudly. get_type_hints raises during schema @@ -37,7 +39,7 @@ from agno.fs.fs import FileSystem from agno.run import RunContext from agno.team.team import Team from agno.tools.toolkit import Toolkit -from agno.utils.log import log_debug, log_error +from agno.utils.log import log_debug, log_error, log_warning _MAX_DIR_ENTRIES = 200 @@ -91,8 +93,9 @@ def _format_with_line_numbers(text: str, start_line: int = 1) -> str: class FileSystemTools(Toolkit): """Toolkit over one ``FileSystem`` file store. Build it with ``FileSystem.tools()``. - Registers the full nine-tool surface, or the four read tools when - ``read_only=True`` (the surface for a consumer agent that consults another + Registers the notes seven by default (``allow_delete=True`` adds + ``delete_file``; ``include_tools`` selects explicitly from the whole + nine-tool surface), or the three read tools when ``read_only=True`` (the surface for a consumer agent that consults another agent's namespace by shared name). Holds the matching FileSystem instructions, which reach the system prompt only under ``add_instructions=True``; otherwise compose ``FileSystem.instructions()`` into the agent's own instructions, and @@ -109,6 +112,7 @@ class FileSystemTools(Toolkit): agent genuinely needs both, wrap one in a sub-agent. """ + # The whole surface, exported for callers that want everything explicitly. FULL_TOOLS: List[str] = [ "read_file", "write_file", @@ -120,22 +124,69 @@ class FileSystemTools(Toolkit): "move_file", "delete_file", ] - READ_ONLY_TOOLS: List[str] = ["read_file", "list_files", "search_content", "check_lines"] + # The default surface: the notes seven. check_lines (the batched record-set + # membership test) is requested by name via include_tools; delete_file (the + # one tool that can lose work) sits behind allow_delete=True. + DEFAULT_TOOLS: List[str] = [ + "read_file", + "write_file", + "append_file", + "replace_lines", + "list_files", + "search_content", + "move_file", + ] + READ_ONLY_TOOLS: List[str] = ["read_file", "list_files", "search_content"] + # What a read-only toolkit may select from via include_tools. + _READ_CAPABLE_TOOLS: List[str] = ["read_file", "list_files", "search_content", "check_lines"] def __init__( self, fs: FileSystem, read_only: bool = False, + allow_delete: bool = False, instructions: Optional[str] = None, add_instructions: bool = False, **kwargs, ): self.fs = fs self.read_only = read_only + # Every mutating tool is a read-modify-write over one row, and a model's + # tool calls for one turn are gathered concurrently (models/base.py) - + # the path AgentOS REST and /mcp take. Two replace_lines on one note + # both read the pre-edit content and the second write wins, with both + # tool results reporting success. Same guard entity memory takes for + # its rows, and the same honesty: single process only. + self._write_locks: Any = WeakKeyDictionary() + if read_only and allow_delete: + raise ValueError("allow_delete=True contradicts read_only=True; pick one.") if instructions is None: instructions = FileSystem.instructions(read_only=read_only) - registered = self.READ_ONLY_TOOLS if read_only else self.FULL_TOOLS + if kwargs.get("include_tools") is not None: + # include_tools is a fully explicit whitelist over the whole + # (read-capable) surface: naming check_lines or delete_file there IS + # the opt-in for them. + registered = self._READ_CAPABLE_TOOLS if read_only else self.FULL_TOOLS + elif read_only: + registered = self.READ_ONLY_TOOLS + else: + registered = self.DEFAULT_TOOLS + (["delete_file"] if allow_delete else []) + if kwargs.get("exclude_tools"): + # Excluding a tool that exists but is not in this set (delete_file, + # which left the default) is a no-op, not an error - upgraders used + # exclude_tools=["delete_file"] as the safety idiom. A name that is + # no tool at all is still a typo, and swallowing it silently leaves + # a tool registered that the caller believed was gone - say so + # without breaking a call that otherwise resolves fine. + unknown = [name for name in kwargs["exclude_tools"] if name not in self.FULL_TOOLS] + if unknown: + log_warning( + f"FileSystem.tools: exclude_tools names {unknown}, which are not FileSystem " + f"tools - check the spelling, because nothing was excluded for them. " + f"Available: {', '.join(self.FULL_TOOLS)}." + ) + kwargs["exclude_tools"] = [name for name in kwargs["exclude_tools"] if name in registered] sync_tools = [getattr(self, name) for name in registered] async_tools = [(getattr(self, "a" + name), name) for name in registered] @@ -166,14 +217,13 @@ class FileSystemTools(Toolkit): if e.scope == "file": return ( f"Error: {path} would be {e.current} bytes (limit {e.limit} per file). " - "Start a new file (for record logs, partition by date, e.g. seen/2026-07-24.md) " - "or delete files you no longer need." + "Split the topic into smaller files (or partition by date) and retry." ) return ( f"Error: storage is full ({e.current} of {e.limit} bytes). " - "Delete only files you are certain are obsolete (see list_files), such as an old date " - "partition, then retry. Do not overwrite or delete records you might still need to " - "make room; if nothing is safely disposable, stop and report that storage is full." + "Free space only if you have a tool for it and only from files you are certain are " + "obsolete (see list_files). Never overwrite or discard records you might still need " + "to make room; if nothing is safely disposable, stop and report that storage is full." ) # ------------------------------------------------------------------ @@ -216,7 +266,7 @@ class FileSystemTools(Toolkit): return ( f"Error: file too long to read whole ({len(contents)} chars, {total} lines; " f"limit {_MAX_READ_CHARS} chars). Read a range with start_line/end_line, or " - "use search_content first: it reports the line number of each match." + "use search_content first: it reports each matching file's first-match line." ) return _format_with_line_numbers(contents, start_line=1) start = start_line if start_line is not None else 1 @@ -363,8 +413,7 @@ class FileSystemTools(Toolkit): Each result gives the line number of the first match, so you can follow up with read_file(path, start_line=..., end_line=...) instead of reading the whole file. ``matches`` counts every occurrence in that file, while ``snippet`` shows only - the first. To check whether exact records are already stored, use check_lines - instead, since substring matches can mislead there. + the first. :param query: Substring to search for. :param directory: Directory to scope the search (default "." = everything). @@ -480,7 +529,7 @@ class FileSystemTools(Toolkit): """Append lines to a file, creating it if needed. Line-oriented: appended content always starts on a fresh line and ends with a - newline. Keep one record per line (one URL, one ID) so check_lines can match + newline. Keep one record per line (one URL, one ID) so exact-line checks can match records exactly. :param path: File path, e.g. "seen/2026-07-24.md". Parent folders are implicit. @@ -678,6 +727,42 @@ class FileSystemTools(Toolkit): self.check_lines, lines, directory, run_context=run_context, agent=agent, team=team ) + def _lock_for(self, key: str) -> "asyncio.Lock": + """The lock for one file, per running event loop. + + An asyncio.Lock binds to the loop that first awaits it, so the cache is + keyed weakly by loop: a second loop (a fresh asyncio.run, a worker + thread) gets its own locks instead of one bound to a closed loop. + """ + try: + loop = asyncio.get_running_loop() + except RuntimeError: + loop = None + if loop is None: + return asyncio.Lock() + per_loop = self._write_locks.get(loop) + if per_loop is None: + per_loop = {} + self._write_locks[loop] = per_loop + lock = per_loop.get(key) + if lock is None: + lock = asyncio.Lock() + per_loop[key] = lock + return lock + + @asynccontextmanager + async def _locked(self, *paths: str): + """Hold the write lock for each named file, acquired in a stable order. + + Sorted so two tool calls naming the same pair (move_file src/dst) cannot + deadlock by taking them in opposite orders. + """ + namespace = getattr(self.fs, "namespace", "") + async with AsyncExitStack() as stack: + for key in sorted({f"{namespace}:{path}" for path in paths if path}): + await stack.enter_async_context(self._lock_for(key)) + yield + async def awrite_file( self, path: str, @@ -689,9 +774,10 @@ class FileSystemTools(Toolkit): team: Optional[Team] = None, ) -> str: """Async variant of ``write_file``.""" - return await asyncio.to_thread( - self.write_file, path, content, overwrite, run_context=run_context, agent=agent, team=team - ) + async with self._locked(path): + return await asyncio.to_thread( + self.write_file, path, content, overwrite, run_context=run_context, agent=agent, team=team + ) async def aappend_file( self, @@ -704,9 +790,10 @@ class FileSystemTools(Toolkit): team: Optional[Team] = None, ) -> str: """Async variant of ``append_file``.""" - return await asyncio.to_thread( - self.append_file, path, content, unique, run_context=run_context, agent=agent, team=team - ) + async with self._locked(path): + return await asyncio.to_thread( + self.append_file, path, content, unique, run_context=run_context, agent=agent, team=team + ) async def areplace_lines( self, @@ -720,9 +807,17 @@ class FileSystemTools(Toolkit): team: Optional[Team] = None, ) -> str: """Async variant of ``replace_lines``.""" - return await asyncio.to_thread( - self.replace_lines, path, start_line, end_line, content, run_context=run_context, agent=agent, team=team - ) + async with self._locked(path): + return await asyncio.to_thread( + self.replace_lines, + path, + start_line, + end_line, + content, + run_context=run_context, + agent=agent, + team=team, + ) async def amove_file( self, @@ -735,9 +830,10 @@ class FileSystemTools(Toolkit): team: Optional[Team] = None, ) -> str: """Async variant of ``move_file``.""" - return await asyncio.to_thread( - self.move_file, src, dst, overwrite, run_context=run_context, agent=agent, team=team - ) + async with self._locked(src, dst): + return await asyncio.to_thread( + self.move_file, src, dst, overwrite, run_context=run_context, agent=agent, team=team + ) async def adelete_file( self, @@ -748,7 +844,8 @@ class FileSystemTools(Toolkit): team: Optional[Team] = None, ) -> str: """Async variant of ``delete_file``.""" - return await asyncio.to_thread(self.delete_file, path, run_context=run_context, agent=agent, team=team) + async with self._locked(path): + return await asyncio.to_thread(self.delete_file, path, run_context=run_context, agent=agent, team=team) # The async twins delegate to their sync counterparts via asyncio.to_thread, so they diff --git a/libs/agno/agno/learn/config.py b/libs/agno/agno/learn/config.py index fc6394f5c4..7aa6d36f46 100644 --- a/libs/agno/agno/learn/config.py +++ b/libs/agno/agno/learn/config.py @@ -13,6 +13,10 @@ Configurations: - SessionContextConfig: Config for session context learning - LearnedKnowledgeConfig: Config for learned knowledge - EntityMemoryConfig: Config for entity memory + +Custom `schema` classes serialize by import path; define them in an +importable module, not `__main__`, or they will not survive a +to_dict/from_dict round-trip across processes. """ from dataclasses import dataclass @@ -34,8 +38,8 @@ class LearningMode(Enum): ALWAYS: Automatic extraction after each response. AGENTIC: Agent decides when to learn via tools. - PROPOSE: Agent proposes, human confirms. - HITL (Human-in-the-Loop): Reserved for future use. + PROPOSE: Agent proposes, human confirms (learned_knowledge only). + HITL: Deprecated, unsupported by every store; removed in 3.0. """ ALWAYS = "always" @@ -330,73 +334,75 @@ class EntityMemoryConfig: - "global": Shared with everyone (default) - Custom string: Explicit grouping (e.g., "sales_west") + Mode: AGENTIC only. The agent records entities through its four tools + (remember_about, link_entities, search_entities, forget); there is no + extraction pass. + Attributes: db: Database backend for storage. - model: Model for extraction (required for ALWAYS mode). - mode: How learning is extracted. Default: ALWAYS. + model: Model for the fact-supersession judgment on writes. + mode: AGENTIC (the only supported mode). schema: Custom schema for entity memory data. Default: EntityMemory. # Sharing boundary namespace: Sharing boundary ("user", "global", or custom). - # Extraction operations - enable_create_entity: Allow creating new entities. - enable_update_entity: Allow updating entity properties. - enable_add_fact: Allow adding facts to entities. - enable_update_fact: Allow updating existing facts. - enable_delete_fact: Allow deleting facts. - enable_add_event: Allow adding events to entities. - enable_add_relationship: Allow adding relationships. - # Limits - max_updates_per_run: Max updates per extraction run. Default: 10. - - # Agent tools - enable_agent_tools: Expose tools to the agent. - agent_can_create_entity: If agent_tools enabled, provide create_entity tool. - agent_can_update_entity: If agent_tools enabled, provide update_entity tool. - agent_can_search_entities: If agent_tools enabled, provide search_entities tool. + max_updates_per_run: Retained for LearningMachine limit propagation. # Prompt customization - instructions: Custom instructions for entity extraction. + instructions: Custom instructions for the supersession judgment. additional_instructions: Extra instructions appended to default. - system_message: Full override for extraction system message. + system_message: Full override for the supersession system message. """ # Required fields db: Optional[Union["BaseDb", "AsyncBaseDb"]] = None model: Optional["Model"] = None - # Mode and extraction - mode: LearningMode = LearningMode.ALWAYS + # Mode (AGENTIC only) + mode: LearningMode = LearningMode.AGENTIC schema: Optional[Type[Any]] = None # Sharing boundary namespace: str = "global" - # Extraction operations - enable_create_entity: bool = True - enable_update_entity: bool = True - enable_add_fact: bool = True - enable_update_fact: bool = True - enable_delete_fact: bool = True - enable_add_event: bool = True - enable_add_relationship: bool = True - # Limits max_updates_per_run: Optional[int] = None + # Fact supersession: a new fact retires a contradicted live fact when the + # judge's confidence is at or above this threshold. Conservative by design - + # a wrong supersession hides information the user gave us. + supersession_threshold: float = 0.8 + + # Rendering bounds. The injected block always starts with a one-line-per- + # entity directory (name + type, newest first, archived excluded), followed + # by at most max_entities_in_context expanded entities, each rendering at + # most max_facts_per_entity live facts (with as-of dates) and the last + # max_events_per_entity events - with visible truncation markers. + max_entities_in_context: int = 5 + max_entities_in_directory: int = 50 + max_facts_per_entity: int = 10 + max_events_per_entity: int = 5 + # Agent tools - enable_agent_tools: bool = False - agent_can_create_entity: bool = True - agent_can_update_entity: bool = True - agent_can_search_entities: bool = True + enable_agent_tools: bool = True # Prompt customization instructions: Optional[str] = None additional_instructions: Optional[str] = None system_message: Optional[str] = None + def __post_init__(self) -> None: + # Fail where the mistake is made: a store built from this config would + # raise the same error, but only on first use inside a run. + if self.mode != LearningMode.AGENTIC: + raise ValueError( + f"Entity memory is AGENTIC-only: the agent records entities through its tools " + f"and there is no extraction pass. Remove mode={self.mode.value!r} or set " + f"LearningMode.AGENTIC." + ) + def __repr__(self) -> str: return f"EntityMemoryConfig(mode={self.mode.value}, namespace={self.namespace}, enable_agent_tools={self.enable_agent_tools})" @@ -413,6 +419,8 @@ class DecisionLogConfig: Decision Logs record decisions made by the agent with reasoning and context. Useful for auditing and learning from past decisions. + Mode: AGENTIC only - the agent logs decisions via tools. + Scope: AGENT (fixed) - Stored and retrieved by agent_id. """ @@ -420,8 +428,8 @@ class DecisionLogConfig: db: Optional[Union["BaseDb", "AsyncBaseDb"]] = None model: Optional["Model"] = None - # Mode and extraction - mode: LearningMode = LearningMode.ALWAYS + # Mode (AGENTIC only) + mode: LearningMode = LearningMode.AGENTIC schema: Optional[Type[Any]] = None # Extraction limits @@ -439,62 +447,3 @@ class DecisionLogConfig: def __repr__(self) -> str: return f"DecisionLogConfig(mode={self.mode.value})" - - -# ============================================================================= -# Placeholder Configurations (Not yet implemented) -# ============================================================================= - - -@dataclass -class FeedbackConfig: - """Configuration for Behavioral Feedback learning type. - - Behavioral Feedback captures signals about what worked and what - didn't: thumbs up/down, corrections, regeneration requests. - - Scope: AGENT (fixed) - Stored and retrieved by agent_id. - - Note: Deferred to Phase 2. - """ - - # Required fields - db: Optional[Union["BaseDb", "AsyncBaseDb"]] = None - model: Optional["Model"] = None - - # Mode and extraction - mode: LearningMode = LearningMode.ALWAYS - schema: Optional[Type[Any]] = None - - # Prompt customization - instructions: Optional[str] = None - - def __repr__(self) -> str: - return "FeedbackConfig(mode=ALWAYS)" - - -@dataclass -class SelfImprovementConfig: - """Configuration for Self-Improvement learning type. - - Self-Improvement proposes updates to agent instructions based - on feedback patterns and successful interactions. - - Scope: AGENT (fixed) - Stored and retrieved by agent_id. - - Note: Deferred to Phase 3. - """ - - # Required fields - db: Optional[Union["BaseDb", "AsyncBaseDb"]] = None - model: Optional["Model"] = None - - # Mode and extraction - mode: LearningMode = LearningMode.HITL - schema: Optional[Type[Any]] = None - - # Prompt customization - instructions: Optional[str] = None - - def __repr__(self) -> str: - return "SelfImprovementConfig(mode=HITL)" diff --git a/libs/agno/agno/learn/machine.py b/libs/agno/agno/learn/machine.py index 1556f9d3de..02b4c97638 100644 --- a/libs/agno/agno/learn/machine.py +++ b/libs/agno/agno/learn/machine.py @@ -14,7 +14,7 @@ Plus maintenance via the Curator for keeping memories healthy. from dataclasses import dataclass, field from os import getenv -from typing import Any, Callable, Dict, List, Optional, Union +from typing import Any, Callable, ClassVar, Dict, List, Optional, Union from agno.learn.config import ( DecisionLogConfig, @@ -40,6 +40,33 @@ try: except ImportError: pass + +def _filter_store_kwargs(callee: Callable, kwargs: Dict[str, Any]) -> Dict[str, Any]: + """Filter kwargs to what the callee accepts (signature-aware dispatch). + + Third-party custom_stores satisfy a Protocol, and a store written as + ``def recall(self, user_id=None)`` with no ``**kwargs`` would raise + TypeError the moment new context kwargs (message, run_context, ...) arrive. + Built-in stores take ``**kwargs`` and see everything; narrow custom stores + keep working untouched. Same inspect.signature approach as + invoke_callable_factory. + """ + import inspect + + try: + parameters = inspect.signature(callee).parameters + except (TypeError, ValueError): + return kwargs + kinds = {p.kind for p in parameters.values()} + if inspect.Parameter.VAR_KEYWORD in kinds: + return kwargs + if inspect.Parameter.VAR_POSITIONAL in kinds: + # A *args-only callee cannot take these as keywords; pass everything + # through so the mismatch fails loudly instead of silently dropping data. + return kwargs + return {key: value for key, value in kwargs.items() if key in parameters} + + # Type aliases for cleaner signatures UserProfileInput = Union[bool, UserProfileConfig, LearningStore, None] UserMemoryInput = Union[bool, UserMemoryConfig, LearningStore, None] @@ -99,6 +126,16 @@ class LearningMachine: # Internal state (lazy initialization) _stores: Optional[Dict[str, LearningStore]] = field(default=None, init=False) _curator: Optional[Any] = field(default=None, init=False) + # Manual-door tracking: instructions() called by the developer marks the + # machine as hand-placed; if the framework then also injects it (learning=), + # the blocks would render twice - warn once. + _placed_by_hand: bool = field(default=False, init=False) + _double_render_warned: bool = field(default=False, init=False) + _missing_user_id_warned: bool = field(default=False, init=False) + _missing_model_warned: bool = field(default=False, init=False) + # Strong refs to fire-and-forget capture tasks (acapture_hook), so the event + # loop cannot garbage-collect them mid-flight. + _capture_tasks: set = field(default_factory=set, init=False) # ========================================================================= # Initialization (Lazy) @@ -106,9 +143,23 @@ class LearningMachine: @property def stores(self) -> Dict[str, LearningStore]: - """All registered stores, keyed by name. Lazily initialized.""" + """All registered stores, keyed by name. Lazily initialized. + + Stores take the machine's model at construction, but the model usually + arrives *after* them: ``learning=`` hands over the agent's model during + Agent init, which is later than any earlier read of this property — a + test, an inspection, or the manual door. Backfill on access so a store + built before the model was bound picks it up, instead of keeping + ``model=None`` for the life of the process and silently declining to + capture. + """ if self._stores is None: self._initialize_stores() + if self.model is not None: + for store in self._stores.values(): # type: ignore[union-attr] + config = getattr(store, "config", None) + if config is not None and getattr(config, "model", "unused") is None: + config.model = self.model return self._stores # type: ignore def _initialize_stores(self) -> None: @@ -178,9 +229,15 @@ class LearningMachine: Returns: Initialized store instance. """ - # Already a store instance + # Already a store instance. The Protocol gained instructions() in 2.8.4; + # the duck-typed fallback keeps third-party stores written before that + # working (they simply contribute no guidance text). if isinstance(input_value, LearningStore): return input_value + if not isinstance(input_value, bool) and all( + callable(getattr(input_value, method, None)) for method in ("recall", "process", "build_context") + ): + return input_value # Create store based on type if store_type == "user_profile": @@ -272,12 +329,16 @@ class LearningMachine: config.model = self.model if config.max_updates_per_run is None: config.max_updates_per_run = self.max_updates_per_run + # A config still on the class default inherits the machine namespace; + # an explicit config namespace wins over it at every call site. + if config.namespace == "global" and self.namespace != "global": + config.namespace = self.namespace else: config = EntityMemoryConfig( db=self.db, model=self.model, namespace=self.namespace, - mode=LearningMode.ALWAYS, + mode=LearningMode.AGENTIC, max_updates_per_run=self.max_updates_per_run, ) @@ -294,6 +355,8 @@ class LearningMachine: config.knowledge = self.knowledge if config.max_updates_per_run is None: config.max_updates_per_run = self.max_updates_per_run + if config.namespace == "global" and self.namespace != "global": + config.namespace = self.namespace else: config = LearnedKnowledgeConfig( model=self.model, @@ -420,7 +483,7 @@ class LearningMachine: message=message, entity_id=entity_id, entity_type=entity_type, - namespace=namespace or self.namespace, + namespace=namespace, agent_id=agent_id, team_id=team_id, **kwargs, @@ -447,7 +510,7 @@ class LearningMachine: message=message, entity_id=entity_id, entity_type=entity_type, - namespace=namespace or self.namespace, + namespace=namespace, agent_id=agent_id, team_id=team_id, **kwargs, @@ -455,6 +518,165 @@ class LearningMachine: return self._format_results(results=results) + def instructions(self) -> str: + """The guidance block: how and when to use the learning tools. + + Mode-aware and aggregated across enabled stores, which is why this is + an instance method and not a static string - the text depends on which + stores are enabled and what mode each is in. Pairs with build_context() + (the data block): the automatic path concatenates the two at the + injection site, and the manual door places each by hand. + + Calling this marks the machine as hand-placed; if the same machine is + ALSO passed to Agent(learning=...), the framework warns once about the + double render. + """ + self._placed_by_hand = True + return self._instructions_text() + + def _warn_if_user_id_missing(self, user_id: Optional[str]) -> None: + """Per-user stores silently drop their tools and capture without a + user_id (an unauthenticated /mcp run is the common way to get here). + Make the degradation visible, once per machine.""" + if user_id or self._missing_user_id_warned: + return + per_user = [name for name in ("user_profile", "user_memory") if self.stores.get(name) is not None] + if per_user: + self._missing_user_id_warned = True + log_warning( + f"This run has no user_id, but per-user learning stores are configured " + f"({', '.join(per_user)}). Their tools and capture are disabled for this run. " + f"Pin Agent(user_id=...) or authenticate the request so a user id reaches the stores." + ) + + def _warn_if_model_missing(self) -> None: + """Capture is a model call, and the manual door injects nothing. + + ``learning=`` hands the agent's model to the machine; a hand-placed + machine keeps whatever it was constructed with. Without one, the + capture tools return "No model provided" and entity memory keeps every + stated fact, both without saying so at the point of use. + """ + if self._missing_model_warned: + return + # Ask the stores, not the machine: `self.model is not None` was the + # earlier guard here, and it short-circuited exactly the case this + # check exists to catch — a machine whose model arrived after its + # stores did. `stores` now backfills those, so a store still without + # one means the machine genuinely has none. + without = [name for name, store in self.stores.items() if getattr(store, "model", "unused") is None] + if not without: + return + self._missing_model_warned = True + log_warning( + f"LearningMachine has no model, so these stores cannot capture: {', '.join(sorted(without))}. " + f"Their tools return 'No model provided', and entity memory records every stated fact without " + f"retiring the ones it contradicts. Pass model= to LearningMachine (the manual door injects " + f"nothing), or attach the machine with learning= so the agent's model is used." + ) + + def _framework_instructions(self) -> str: + """The guidance block, fetched by the automatic injection path. + + Detects the double-render case: the same machine placed by hand + (instructions() was called) AND attached via learning=. + """ + if self._placed_by_hand and not self._double_render_warned: + self._double_render_warned = True + log_warning( + "This LearningMachine is attached via learning= AND its surfaces are placed by " + "hand (instructions()/build_context()); its context will render twice. Pick one " + "door: pass learning= for the automatic path, or place the surfaces yourself " + "without learning=." + ) + return self._instructions_text() + + def _instructions_text(self) -> str: + parts: List[str] = [] + for name, store in self.stores.items(): + instructions_fn = getattr(store, "instructions", None) + if not callable(instructions_fn): + continue + try: + text = instructions_fn() + if text: + parts.append(text) + except Exception as e: + log_warning(f"Error getting instructions from {name}: {str(e)}") + return "\n\n".join(parts) + + def capture_hook(self) -> Callable: + """A post_hooks-compatible callable that runs this machine's capture pass. + + The manual door is agentic by nature: with no learning= there is no + automatic post-run extraction, and the tools are the capture mechanism. + For the developer who wants hand-placed prompts AND ALWAYS-mode + extraction, add this to Agent(post_hooks=[...]). An escape hatch, not a + third shape. + + Extraction is backgrounded on the agent's background executor when one + is available (the same place learning= runs it), so it stays off the + response path; without an executor it runs inline. Works on both sync + and async runs (async post-hook execution calls sync hooks directly). + """ + machine = self + + def learning_capture(run_output=None, agent=None, session=None, user_id=None, run_context=None) -> None: + messages = list(getattr(run_output, "messages", None) or []) + if not messages: + return + kwargs = dict( + messages=messages, + user_id=user_id, + session_id=getattr(session, "session_id", None) if session is not None else None, + agent_id=getattr(agent, "id", None) if agent is not None else None, + team_id=getattr(agent, "team_id", None) if agent is not None else None, + run_context=run_context, + metadata=getattr(run_context, "metadata", None), + dependencies=getattr(run_context, "dependencies", None), + session_state=getattr(run_context, "session_state", None), + ) + executor = getattr(agent, "background_executor", None) if agent is not None else None + if executor is not None: + executor.submit(machine.process, **kwargs) + else: + machine.process(**kwargs) + + return learning_capture + + def acapture_hook(self) -> Callable: + """Async version of capture_hook: schedules aprocess as a background task. + + Only for async runs - a sync run() skips async hooks with a warning, so + pass capture_hook() there instead. The task is fire-and-forget off the + response path; failures are logged, never raised into the run. + """ + machine = self + + async def alearning_capture(run_output=None, agent=None, session=None, user_id=None, run_context=None) -> None: + import asyncio + + messages = list(getattr(run_output, "messages", None) or []) + if not messages: + return + task = asyncio.create_task( + machine.aprocess( + messages=messages, + user_id=user_id, + session_id=getattr(session, "session_id", None) if session is not None else None, + agent_id=getattr(agent, "id", None) if agent is not None else None, + team_id=getattr(agent, "team_id", None) if agent is not None else None, + run_context=run_context, + metadata=getattr(run_context, "metadata", None), + dependencies=getattr(run_context, "dependencies", None), + session_state=getattr(run_context, "session_state", None), + ) + ) + machine._capture_tasks.add(task) + task.add_done_callback(machine._capture_tasks.discard) + + return alearning_capture + def get_tools( self, user_id: Optional[str] = None, @@ -467,8 +689,9 @@ class LearningMachine: """Get learning tools to expose to the agent. Returns tools based on which stores are enabled: - - user_profile: update_user_memory - - entity_memory: search_entities, create_entity, update_entity, add_fact, etc. + - user_profile: update_profile + - user_memory: update_user_memory + - entity_memory: remember_about, link_entities, search_entities, forget - learned_knowledge: search_learnings, save_learning Args: @@ -482,10 +705,12 @@ class LearningMachine: List of callable tools. """ tools = [] + self._warn_if_user_id_missing(user_id) + self._warn_if_model_missing() context = { "user_id": user_id, "session_id": session_id, - "namespace": namespace or self.namespace, + "namespace": namespace, "agent_id": agent_id, "team_id": team_id, **kwargs, @@ -493,7 +718,7 @@ class LearningMachine: for name, store in self.stores.items(): try: - store_tools = store.get_tools(**context) + store_tools = store.get_tools(**_filter_store_kwargs(store.get_tools, context)) if store_tools: tools.extend(store_tools) log_debug(f"Got {len(store_tools)} tools from {name}") @@ -513,10 +738,12 @@ class LearningMachine: ) -> List[Callable]: """Async version of get_tools.""" tools = [] + self._warn_if_user_id_missing(user_id) + self._warn_if_model_missing() context = { "user_id": user_id, "session_id": session_id, - "namespace": namespace or self.namespace, + "namespace": namespace, "agent_id": agent_id, "team_id": team_id, **kwargs, @@ -524,7 +751,7 @@ class LearningMachine: for name, store in self.stores.items(): try: - store_tools = await store.aget_tools(**context) + store_tools = await store.aget_tools(**_filter_store_kwargs(store.aget_tools, context)) if store_tools: tools.extend(store_tools) log_debug(f"Got {len(store_tools)} tools from {name}") @@ -560,7 +787,7 @@ class LearningMachine: "messages": messages, "user_id": user_id, "session_id": session_id, - "namespace": namespace or self.namespace, + "namespace": namespace, "agent_id": agent_id, "team_id": team_id, **kwargs, @@ -568,7 +795,7 @@ class LearningMachine: for name, store in self.stores.items(): try: - store.process(**context) + store.process(**_filter_store_kwargs(store.process, context)) if getattr(store, "was_updated", False): log_debug(f"Store {name} was updated") except Exception as e: @@ -589,7 +816,7 @@ class LearningMachine: "messages": messages, "user_id": user_id, "session_id": session_id, - "namespace": namespace or self.namespace, + "namespace": namespace, "agent_id": agent_id, "team_id": team_id, **kwargs, @@ -597,7 +824,7 @@ class LearningMachine: for name, store in self.stores.items(): try: - await store.aprocess(**context) + await store.aprocess(**_filter_store_kwargs(store.aprocess, context)) if getattr(store, "was_updated", False): log_debug(f"Store {name} was updated") except Exception as e: @@ -634,7 +861,7 @@ class LearningMachine: "query": message, # For learned_knowledge "entity_id": entity_id, "entity_type": entity_type, - "namespace": namespace or self.namespace, + "namespace": namespace, "agent_id": agent_id, "team_id": team_id, **kwargs, @@ -642,7 +869,7 @@ class LearningMachine: for name, store in self.stores.items(): try: - result = store.recall(**context) + result = store.recall(**_filter_store_kwargs(store.recall, context)) results[name] = result try: log_debug(f"Recalled from {name}: {result}") @@ -674,7 +901,7 @@ class LearningMachine: "query": message, "entity_id": entity_id, "entity_type": entity_type, - "namespace": namespace or self.namespace, + "namespace": namespace, "agent_id": agent_id, "team_id": team_id, **kwargs, @@ -682,7 +909,7 @@ class LearningMachine: for name, store in self.stores.items(): try: - result = await store.arecall(**context) + result = await store.arecall(**_filter_store_kwargs(store.arecall, context)) results[name] = result try: log_debug(f"Recalled from {name}: {result}") @@ -745,49 +972,196 @@ class LearningMachine: # Serialization # ========================================================================= + _STORE_CONFIG_CLASSES: ClassVar[Dict[str, Any]] = { + "user_profile": UserProfileConfig, + "user_memory": UserMemoryConfig, + "session_context": SessionContextConfig, + "entity_memory": EntityMemoryConfig, + "learned_knowledge": LearnedKnowledgeConfig, + "decision_log": DecisionLogConfig, + } + def to_dict(self) -> Dict[str, Any]: """Serialize the LearningMachine configuration to a dictionary. - Preserves which stores are enabled and the namespace so that - from_dict() can reconstruct an equivalent instance. Does not - serialize db, model, or knowledge (those are injected at init). + Lossless for what a platform rebuild needs: per-store mode, + namespace, enable_* toggles, limits, prompt strings and + schema-by-import-path all round-trip. db, model and knowledge are + never serialized (they are injected at init), and callables cannot + be - a config carrying one round-trips as present-but-not-restorable + and from_dict logs once. A schema class defined in a __main__ script + serializes as a __main__ path and only resolves inside the same + process; define custom schemas in an importable module for a + cross-process rebuild. Store INSTANCES serialize their config; + custom_stores serialize as import-path refs (informational - the + instances cannot be rebuilt from a dict). """ d: Dict[str, Any] = {} - if self.user_profile: - d["user_profile"] = True - if self.user_memory: - d["user_memory"] = True - if self.session_context: - d["session_context"] = True - if self.entity_memory: - d["entity_memory"] = True - if self.learned_knowledge: - d["learned_knowledge"] = True - if self.decision_log: - d["decision_log"] = True + for store_name in self._STORE_CONFIG_CLASSES: + value = getattr(self, store_name) + if not value: + continue + d[store_name] = self._serialize_store_input(value) if self.namespace != "global": d["namespace"] = self.namespace + if self.max_updates_per_run != 10: + d["max_updates_per_run"] = self.max_updates_per_run + if self.custom_stores: + d["custom_stores"] = { + name: f"{type(store).__module__}.{type(store).__qualname__}" + for name, store in self.custom_stores.items() + } if self.debug_mode: d["debug_mode"] = True return d + def _serialize_store_input(self, value: Any) -> Any: + """bool -> True; Config -> field dict; Store instance -> its config's dict.""" + import dataclasses + + if value is True: + return True + # A store instance carries its config (stores are dataclasses too, so + # check for the carried config before the dataclass test). + inner = getattr(value, "config", None) + if inner is not None and dataclasses.is_dataclass(inner): + config = inner + store_module = type(value).__module__ or "" + if not store_module.startswith("agno.learn.stores"): + log_warning( + f"LearningMachine.to_dict: store {type(value).__name__} is serialized by its " + f"config only; the subclass itself cannot be rebuilt from a dict." + ) + elif dataclasses.is_dataclass(value) and not isinstance(value, type): + config = value + else: + return True + + serialized: Dict[str, Any] = {} + for f in dataclasses.fields(config): + if f.name in ("db", "model", "knowledge"): + continue + field_value = getattr(config, f.name) + if field_value is None: + continue + if isinstance(field_value, LearningMode): + serialized["mode"] = field_value.value + elif isinstance(field_value, type): + serialized[f.name] = f"{field_value.__module__}.{field_value.__qualname__}" + elif callable(field_value): + serialized[f.name] = {"__callable__": getattr(field_value, "__name__", "callable")} + elif isinstance(field_value, (str, int, float, bool, list, dict)): + import json + + try: + json.dumps(field_value) + except (TypeError, ValueError): + log_warning( + f"LearningMachine.to_dict: dropping non-JSON-serializable config value {f.name!r}; " + f"set it programmatically after rebuild." + ) + continue + serialized[f.name] = field_value + return serialized + @classmethod def from_dict(cls, data: Dict[str, Any]) -> "LearningMachine": """Reconstruct a LearningMachine from a serialized dictionary. - db and model must be injected separately (e.g. during agent/team init). + db and model must be injected separately (e.g. during agent/team + init). Accepts both the lossless per-store dicts and the old + boolean-only payloads. Unrestorable entries (callables, custom-store + refs) are logged once, never silently dropped to a different default. """ + kwargs: Dict[str, Any] = {} + for store_name, config_cls in cls._STORE_CONFIG_CLASSES.items(): + value = data.get(store_name, False) + if isinstance(value, dict): + kwargs[store_name] = cls._deserialize_store_config(store_name, config_cls, value) + else: + kwargs[store_name] = bool(value) + + if data.get("custom_stores"): + log_warning( + f"LearningMachine.from_dict: custom stores {sorted(data['custom_stores'])} cannot be " + f"rebuilt from a serialized config; re-attach them programmatically." + ) + return cls( - user_profile=data.get("user_profile", False), - user_memory=data.get("user_memory", False), - session_context=data.get("session_context", False), - entity_memory=data.get("entity_memory", False), - learned_knowledge=data.get("learned_knowledge", False), - decision_log=data.get("decision_log", False), namespace=data.get("namespace", "global"), + max_updates_per_run=data.get("max_updates_per_run", 10), debug_mode=data.get("debug_mode", False), + **kwargs, ) + @classmethod + def _deserialize_store_config(cls, store_name: str, config_cls: Any, payload: Dict[str, Any]) -> Any: + import dataclasses + import importlib + + field_names = {f.name for f in dataclasses.fields(config_cls)} + kwargs: Dict[str, Any] = {} + for key, value in payload.items(): + if key not in field_names: + continue + if key == "mode": + try: + kwargs["mode"] = LearningMode(value) + except ValueError: + # AGENTIC, not the class default: several stores default to + # ALWAYS, and an unrecognized mode must never turn into a + # surprise per-run extraction pass. + log_warning(f"LearningMachine.from_dict: unknown mode {value!r} on {store_name}; using AGENTIC.") + kwargs["mode"] = LearningMode.AGENTIC + continue + if key == "schema" and isinstance(value, str): + module_path, _, attr_path = value.rpartition(".") + try: + target: Any = importlib.import_module(module_path) + for part in attr_path.split("."): + target = getattr(target, part) + kwargs["schema"] = target + except Exception as e: + if module_path == "__main__": + log_warning( + f"LearningMachine.from_dict: schema {value!r} on {store_name} was defined " + f"in a __main__ script and cannot be imported from another process; " + f"define custom schemas in an importable module. The store falls back " + f"to its default schema." + ) + else: + log_warning( + f"LearningMachine.from_dict: could not import schema {value!r} for {store_name}: {e}" + ) + continue + if isinstance(value, dict) and "__callable__" in value: + log_warning( + f"LearningMachine.from_dict: {store_name}.{key} was a callable " + f"({value['__callable__']}) and cannot be restored from a serialized config; " + f"set it programmatically." + ) + continue + kwargs[key] = value + try: + return config_cls(**kwargs) + except Exception as e: + # Keep every valid field: a payload carrying an invalid mode (e.g. a + # legacy 'always' for entity memory) must not also lose its + # namespace and knobs. + retry_kwargs = {k: v for k, v in kwargs.items() if k != "mode"} + if retry_kwargs != kwargs: + try: + rebuilt = config_cls(**retry_kwargs) + log_warning( + f"LearningMachine.from_dict: dropped invalid mode on {store_name} ({e}); " + f"other settings preserved." + ) + return rebuilt + except Exception: + pass + log_warning(f"LearningMachine.from_dict: could not rebuild {store_name} config ({e}); enabling defaults.") + return True + # ========================================================================= # Representation # ========================================================================= diff --git a/libs/agno/agno/learn/schemas.py b/libs/agno/agno/learn/schemas.py index 91fd564167..4f22ee1e00 100644 --- a/libs/agno/agno/learn/schemas.py +++ b/libs/agno/agno/learn/schemas.py @@ -28,8 +28,6 @@ Schemas: - LearnedKnowledge: Reusable knowledge/insights - EntityMemory: Third-party entity facts - DecisionLog: Decision logs -- Feedback: Behavioral feedback (Phase 2) -- InstructionUpdate: Self-improvement (Phase 3) """ from dataclasses import asdict, dataclass, field, fields @@ -594,6 +592,7 @@ class EntityMemory: properties: Dict[str, str] = field( default_factory=dict, metadata={"description": "Key-value properties (industry, tech_stack, etc)"} ) + aliases: List[str] = field(default_factory=list, metadata={"description": "Other names this entity resolves from"}) # Semantic memory (facts) facts: List[Dict[str, Any]] = field(default_factory=list) @@ -614,6 +613,9 @@ class EntityMemory: team_id: Optional[str] = field(default=None, metadata={"internal": True}) created_at: Optional[str] = field(default=None, metadata={"internal": True}) updated_at: Optional[str] = field(default=None, metadata={"internal": True}) + # Set when the entity is archived via forget(); archived entities are excluded + # from recall, rendering and the directory, but stay reachable by search. + archived_at: Optional[str] = field(default=None, metadata={"internal": True}) @classmethod def from_dict(cls, data: Any) -> Optional["EntityMemory"]: @@ -689,6 +691,18 @@ class EntityMemory: if content and content.strip(): now = _utc_now_iso() + # Idempotent for the same day, like add_relationship: models + # re-state what they already recorded, and a duplicated event + # renders twice and retires ambiguously. + for existing in self.events: + if ( + isinstance(existing, dict) + and existing.get("content") == content.strip() + and existing.get("date") == date + ): + existing["updated_at"] = now + return str(existing.get("id", "")) + event = {"id": event_id, "content": content.strip(), "created_at": now, "updated_at": now, **kwargs} if date: event["date"] = date @@ -710,9 +724,27 @@ class EntityMemory: """ import uuid - rel_id = str(uuid.uuid4())[:8] - now = _utc_now_iso() + + # Idempotent: models re-assert links they already know, and an appended + # duplicate cannot be retired - forget lists byte-identical candidates + # no wording can tell apart. Touch the existing edge instead. + far_type = kwargs.get("entity_type") + for existing in self.relationships: + if ( + isinstance(existing, dict) + and existing.get("entity_id") == related_entity_id + and existing.get("relation") == relation + and existing.get("direction") == direction + # entity_type is part of the identity: project/Harbor and + # company/Harbor share a slug and are different things. + and existing.get("entity_type") == far_type + ): + existing["updated_at"] = now + existing.update(kwargs) + return str(existing.get("id", "")) + + rel_id = str(uuid.uuid4())[:8] self.relationships.append( { "id": rel_id, @@ -758,8 +790,44 @@ class EntityMemory: self.facts = [f for f in self.facts if not (isinstance(f, dict) and f.get("id") == fact_id)] return len(self.facts) < original_len - def get_context_text(self) -> str: - """Get entity as formatted string for prompts.""" + def live_facts(self) -> List[Dict[str, Any]]: + """Facts that have not been superseded. Only these render.""" + return [f for f in self.facts if not (isinstance(f, dict) and f.get("superseded_at"))] + + def retire_fact(self, fact_id: str, superseded_by: str) -> bool: + """Mark a fact superseded without deleting it. + + Args: + fact_id: The fact to retire. + superseded_by: What replaced it - a new fact's id, "forgotten", or + "superseded" when several new facts jointly replaced it. + + Returns: + True if the fact was found and retired, False otherwise. + """ + for fact in self.facts: + if isinstance(fact, dict) and fact.get("id") == fact_id and not fact.get("superseded_at"): + fact["superseded_at"] = _utc_now_iso() + fact["superseded_by"] = superseded_by + return True + return False + + def get_context_text( + self, + max_facts: Optional[int] = 10, + max_events: Optional[int] = 5, + related_names: Optional[Dict[str, str]] = None, + ) -> str: + """Get entity as formatted string for prompts. + + Bounded and honest: at most ``max_facts`` live facts (each with an + as-of date, so July's truth outranks March's) and the last + ``max_events`` events, with explicit truncation markers so a capped + render is never mistaken for the whole record. Relationships render in + full - they are one line each; ``related_names`` maps far-end entity + ids to display names (the one-hop expansion), falling back to the id. + Pass ``None`` to lift a cap. + """ parts = [] if self.name: @@ -774,19 +842,58 @@ class EntityMemory: props = ", ".join(f"{k}: {v}" for k, v in self.properties.items()) parts.append(f"Properties: {props}") - if self.facts: - facts_text = "\n".join(f" - {f.get('content', f)}" for f in self.facts) - parts.append(f"Facts:\n{facts_text}") + live = self.live_facts() + if live: + # The NEWEST N facts render (in chronological order) - showing the + # oldest slice would date-stamp stale state as current. + if max_facts is None: + shown = live + elif max_facts <= 0: + shown = [] + else: + shown = live[-max_facts:] + marker = f" (newest {len(shown)} of {len(live)} facts)" if len(shown) < len(live) else "" + fact_lines = [] + for f in shown: + if isinstance(f, dict): + as_of = str(f.get("updated_at") or f.get("created_at") or "")[:10] + as_of_text = f" (as of {as_of})" if as_of else "" + fact_lines.append(f" - {f.get('content', f)}{as_of_text}") + else: + fact_lines.append(f" - {f}") + if shown: + parts.append(f"Facts:{marker}\n" + "\n".join(fact_lines)) if self.events: - events_text = "\n".join( - f" - {e.get('content', e)}" + (f" ({e.get('date')})" if e.get("date") else "") for e in self.events + if max_events is None: + shown_events = self.events + elif max_events <= 0: + shown_events = [] + else: + shown_events = self.events[-max_events:] + marker = ( + f" (last {len(shown_events)} of {len(self.events)} events)" + if len(shown_events) < len(self.events) + else "" ) - parts.append(f"Events:\n{events_text}") + event_lines = [] + for e in shown_events: + if isinstance(e, dict): + date = f" ({e.get('date')})" if e.get("date") else "" + event_lines.append(f" - {e.get('content', e)}{date}") + else: + event_lines.append(f" - {e}") + if shown_events: + parts.append(f"Events:{marker}\n" + "\n".join(event_lines)) if self.relationships: - rels_text = "\n".join(f" - {r.get('relation')}: {r.get('entity_id')}" for r in self.relationships) - parts.append(f"Relationships:\n{rels_text}") + rel_lines = [] + for r in self.relationships: + far_id = r.get("entity_id", "?") + far_label = (related_names or {}).get(far_id, far_id) + arrow = "->" if r.get("direction", "outgoing") == "outgoing" else "<-" + rel_lines.append(f" - {r.get('relation')} {arrow} {far_label}") + parts.append("Relationships:\n" + "\n".join(rel_lines)) return "\n\n".join(parts) @@ -1054,109 +1161,3 @@ class DecisionLog: # Backwards compatibility alias Decision = DecisionLog - - -# ============================================================================= -# Placeholder Schemas (Not yet implemented) -# ============================================================================= - - -@dataclass -class Feedback: - """Schema for Behavioral Feedback. (Phase 2) - - Captures signals about what worked and what didn't. - """ - - signal: str # thumbs_up, thumbs_down, correction, regeneration - learning: Optional[str] = None - context: Optional[str] = None - agent_id: Optional[str] = None - team_id: Optional[str] = None - created_at: Optional[str] = None - - @classmethod - def from_dict(cls, data: Any) -> Optional["Feedback"]: - """Parse from dict/JSON, returning None on any failure.""" - if data is None: - return None - if isinstance(data, cls): - return data - - try: - parsed = _parse_json(data) - if not parsed: - log_debug(f"{cls.__name__}.from_dict: _parse_json returned None for data={_truncate_for_log(data)}") - return None - - if not parsed.get("signal"): - log_debug(f"{cls.__name__}.from_dict: missing required field 'signal'") - return None - - field_names = {f.name for f in fields(cls)} - kwargs = {k: v for k, v in parsed.items() if k in field_names} - - return cls(**kwargs) - except Exception as e: - log_debug(f"{cls.__name__}.from_dict failed: {e}, data={_truncate_for_log(data)}") - return None - - def to_dict(self) -> Dict[str, Any]: - """Convert to dict.""" - try: - return asdict(self) - except Exception as e: - log_debug(f"{self.__class__.__name__}.to_dict failed: {e}") - return {} - - -@dataclass -class InstructionUpdate: - """Schema for Self-Improvement. (Phase 3) - - Proposes updates to agent instructions based on feedback patterns. - """ - - current_instruction: str - proposed_instruction: str - reasoning: str - evidence: Optional[List[str]] = None - agent_id: Optional[str] = None - team_id: Optional[str] = None - created_at: Optional[str] = None - - @classmethod - def from_dict(cls, data: Any) -> Optional["InstructionUpdate"]: - """Parse from dict/JSON, returning None on any failure.""" - if data is None: - return None - if isinstance(data, cls): - return data - - try: - parsed = _parse_json(data) - if not parsed: - log_debug(f"{cls.__name__}.from_dict: _parse_json returned None for data={_truncate_for_log(data)}") - return None - - required = ["current_instruction", "proposed_instruction", "reasoning"] - missing = [k for k in required if not parsed.get(k)] - if missing: - log_debug(f"{cls.__name__}.from_dict: missing required fields {missing}") - return None - - field_names = {f.name for f in fields(cls)} - kwargs = {k: v for k, v in parsed.items() if k in field_names} - - return cls(**kwargs) - except Exception as e: - log_debug(f"{cls.__name__}.from_dict failed: {e}, data={_truncate_for_log(data)}") - return None - - def to_dict(self) -> Dict[str, Any]: - """Convert to dict.""" - try: - return asdict(self) - except Exception as e: - log_debug(f"{self.__class__.__name__}.to_dict failed: {e}") - return {} diff --git a/libs/agno/agno/learn/stores/decision_log.py b/libs/agno/agno/learn/stores/decision_log.py index 6696da45a5..e98e61ff6d 100644 --- a/libs/agno/agno/learn/stores/decision_log.py +++ b/libs/agno/agno/learn/stores/decision_log.py @@ -17,8 +17,7 @@ Scope: - Can be queried by agent_id, session_id, or time range Supported Modes: -- ALWAYS: Automatic extraction of decisions from tool calls -- AGENTIC: Agent explicitly logs decisions via tools +- AGENTIC: the agent logs decisions via tools. """ import uuid @@ -31,7 +30,7 @@ from typing import Any, Callable, List, Optional, Union from agno.learn.config import DecisionLogConfig, LearningMode from agno.learn.schemas import DecisionLog from agno.learn.stores.protocol import LearningStore -from agno.learn.utils import from_dict_safe, to_dict_safe +from agno.learn.utils import from_dict_safe, to_dict_safe, values_match_query from agno.utils.log import ( log_debug, log_warning, @@ -41,7 +40,6 @@ from agno.utils.log import ( try: from agno.db.base import AsyncBaseDb, BaseDb - from agno.models.message import Message except ImportError: pass @@ -64,10 +62,14 @@ class DecisionLogStore(LearningStore): # State tracking (internal) decisions_updated: bool = field(default=False, init=False) _schema: Any = field(default=None, init=False) + _degraded_search_logged: bool = field(default=False, init=False) def __post_init__(self): self._schema = self.config.schema or DecisionLog + if self.config.mode != LearningMode.AGENTIC: + log_warning("DecisionLogStore is AGENTIC-only: the agent logs decisions via tools. Proceeding as AGENTIC.") + # ========================================================================= # LearningStore Protocol Implementation # ========================================================================= @@ -93,9 +95,15 @@ class DecisionLogStore(LearningStore): ) -> Optional[List[DecisionLog]]: """Retrieve recent decisions. + A decision log is worth having because it crosses sessions: the run's + session_id is accepted and deliberately NOT used as a filter, or the + injected block would say "No recent decisions logged" to every new + session while the store is full. Explicit callers scope with + search(session_id=...). + Args: agent_id: Filter by agent (optional). - session_id: Filter by session (optional). + session_id: Accepted for the store protocol; not used as a filter. decision_type: Filter by decision type (optional). limit: Maximum number of decisions to return. days: Only return decisions from last N days. @@ -106,7 +114,6 @@ class DecisionLogStore(LearningStore): """ return self.search( agent_id=agent_id, - session_id=session_id, decision_type=decision_type, limit=limit, days=days, @@ -124,78 +131,25 @@ class DecisionLogStore(LearningStore): """Async version of recall.""" return await self.asearch( agent_id=agent_id, - session_id=session_id, decision_type=decision_type, limit=limit, days=days, ) - def process( - self, - messages: List[Any], - agent_id: Optional[str] = None, - session_id: Optional[str] = None, - user_id: Optional[str] = None, - team_id: Optional[str] = None, - **kwargs, - ) -> None: - """Extract decisions from messages (tool calls, etc). + def process(self, messages: List[Any], **kwargs) -> None: + """No-op: decision logging is AGENTIC-only, capture happens through the tools.""" + return - In ALWAYS mode, this extracts decisions from tool calls and - significant response choices. In AGENTIC mode, this is a no-op - as decisions are logged explicitly via tools. - - Args: - messages: Conversation messages to analyze. - agent_id: Agent context. - session_id: Session context. - user_id: User context. - team_id: Team context. - **kwargs: Additional context (ignored). - """ - if self.config.mode != LearningMode.ALWAYS: - return - - if not messages: - return - - # Extract decisions from tool calls in messages - self._extract_decisions_from_messages( - messages=messages, - agent_id=agent_id, - session_id=session_id, - user_id=user_id, - team_id=team_id, - ) - - async def aprocess( - self, - messages: List[Any], - agent_id: Optional[str] = None, - session_id: Optional[str] = None, - user_id: Optional[str] = None, - team_id: Optional[str] = None, - **kwargs, - ) -> None: - """Async version of process.""" - if self.config.mode != LearningMode.ALWAYS: - return - - if not messages: - return - - await self._aextract_decisions_from_messages( - messages=messages, - agent_id=agent_id, - session_id=session_id, - user_id=user_id, - team_id=team_id, - ) + async def aprocess(self, messages: List[Any], **kwargs) -> None: + """Async version of process (no-op).""" + return def build_context(self, data: Any) -> str: - """Build context for the agent. + """Build the DATA context for the agent. Formats recent decisions for injection into the agent's system prompt. + Data only - the how-to-use guidance lives in instructions(); the + automatic path concatenates the two at the injection site. Args: data: List of decisions from recall(). @@ -208,9 +162,6 @@ class DecisionLogStore(LearningStore): return dedent("""\ No recent decisions logged. - - Use `log_decision` to record significant decisions with reasoning. - Use `search_decisions` to find past decisions. """) return "" @@ -235,17 +186,24 @@ class DecisionLogStore(LearningStore): context += f" Outcome: {decision['outcome']}\n" context += "\n" - if self._should_expose_tools: - context += dedent(""" - Use `log_decision` to record new decisions. - Use `search_decisions` to find past decisions. - Use `record_outcome` to update a decision with its outcome. - """) - context += "" return context + def instructions(self) -> str: + """Agent-facing guidance for this store: when to log and search decisions. + + Guidance only - the recalled decisions live in build_context(). + """ + if not self._should_expose_tools: + return "" + return dedent("""\ + + Use `log_decision` to record significant decisions with reasoning. + Use `search_decisions` to find past decisions. + Use `record_outcome` to update a decision with its outcome. + """) + def get_tools( self, agent_id: Optional[str] = None, @@ -710,6 +668,10 @@ class DecisionLogStore(LearningStore): ) -> List[DecisionLog]: """Search decisions with filters. + Text queries route through the db's server-side search_learnings; the + client-side scan is only the fallback for backends without it. + decision_type and days filters are applied client-side either way. + Args: query: Text to search for. agent_id: Filter by agent. @@ -728,62 +690,35 @@ class DecisionLogStore(LearningStore): if not isinstance(self.db, BaseDb): return [] - try: - # Get all matching records - results = self.db.get_learnings( - learning_type=self.learning_type, - agent_id=agent_id, - limit=limit * 3, # Over-fetch for filtering - ) + # Headroom for the client-side verification and decision_type/days filters. + fetch_limit = limit * 3 if (query or decision_type or days) else limit - if not results: - return [] + if query: + if not callable(getattr(self.db, "search_learnings", None)): + self._log_degraded_search_once() + results = self._fetch_recent_rows(agent_id=agent_id, session_id=session_id, limit=limit * 3) + else: + try: + results = self.db.search_learnings( + query=query, + learning_type=self.learning_type, + agent_id=agent_id, + session_id=session_id, + limit=fetch_limit, + ) + except NotImplementedError: + self._log_degraded_search_once() + results = self._fetch_recent_rows(agent_id=agent_id, session_id=session_id, limit=limit * 3) + else: + results = self._fetch_recent_rows(agent_id=agent_id, session_id=session_id, limit=fetch_limit) - decisions = [] - cutoff_date = None - if days: - cutoff_date = datetime.now(timezone.utc) - timedelta(days=days) - - for record in results: - content = record.get("content") if isinstance(record, dict) else None - if not content: - continue - - decision = from_dict_safe(DecisionLog, content) - if not decision: - continue - - # Apply filters - if decision_type and decision.decision_type != decision_type: - continue - - if cutoff_date and decision.created_at: - try: - created = datetime.fromisoformat(decision.created_at.replace("Z", "+00:00")) - # Records written before this fix are naive; assume UTC. - if created.tzinfo is None: - created = created.replace(tzinfo=timezone.utc) - if created < cutoff_date: - continue - except (ValueError, AttributeError): - pass - - if query: - query_lower = query.lower() - text = decision.to_text().lower() - if query_lower not in text: - continue - - decisions.append(decision) - - if len(decisions) >= limit: - break - - return decisions - - except Exception as e: - log_debug(f"DecisionLogStore.search failed: {e}") - return [] + return self._filter_decisions( + results or [], + query=query, + decision_type=decision_type, + days=days, + limit=limit, + ) async def asearch( self, @@ -798,68 +733,144 @@ class DecisionLogStore(LearningStore): if not self.db: return [] + fetch_limit = limit * 3 if (query or decision_type or days) else limit + + if query: + if not callable(getattr(self.db, "search_learnings", None)): + self._log_degraded_search_once() + results = await self._afetch_recent_rows(agent_id=agent_id, session_id=session_id, limit=limit * 3) + else: + try: + if isinstance(self.db, AsyncBaseDb): + results = await self.db.search_learnings( + query=query, + learning_type=self.learning_type, + agent_id=agent_id, + session_id=session_id, + limit=fetch_limit, + ) + else: + results = self.db.search_learnings( + query=query, + learning_type=self.learning_type, + agent_id=agent_id, + session_id=session_id, + limit=fetch_limit, + ) + except NotImplementedError: + self._log_degraded_search_once() + results = await self._afetch_recent_rows(agent_id=agent_id, session_id=session_id, limit=limit * 3) + else: + results = await self._afetch_recent_rows(agent_id=agent_id, session_id=session_id, limit=fetch_limit) + + return self._filter_decisions( + results or [], + query=query, + decision_type=decision_type, + days=days, + limit=limit, + ) + + def _log_degraded_search_once(self) -> None: + if not self._degraded_search_logged: + self._degraded_search_logged = True + log_warning( + "DecisionLogStore: this db backend has no search_learnings implementation; " + "falling back to a client-side scan over the most recently updated rows. " + "Search quality degrades as the store grows." + ) + + def _fetch_recent_rows(self, agent_id: Optional[str], limit: int, session_id: Optional[str] = None) -> List[Any]: + # Callers guard isinstance(self.db, BaseDb), so the call is sync here. + if not isinstance(self.db, BaseDb): + return [] + try: + return ( + self.db.get_learnings( + learning_type=self.learning_type, + agent_id=agent_id, + session_id=session_id, + limit=limit, + ) + or [] + ) + except Exception as e: + log_debug(f"DecisionLogStore._fetch_recent_rows failed: {e}") + return [] + + async def _afetch_recent_rows( + self, agent_id: Optional[str], limit: int, session_id: Optional[str] = None + ) -> List[Any]: try: if isinstance(self.db, AsyncBaseDb): - results = await self.db.get_learnings( + rows = await self.db.get_learnings( learning_type=self.learning_type, agent_id=agent_id, - limit=limit * 3, + session_id=session_id, + limit=limit, ) else: - results = self.db.get_learnings( + rows = self.db.get_learnings( # type: ignore[union-attr] learning_type=self.learning_type, agent_id=agent_id, - limit=limit * 3, + session_id=session_id, + limit=limit, ) - - if not results: - return [] - - decisions = [] - cutoff_date = None - if days: - cutoff_date = datetime.now(timezone.utc) - timedelta(days=days) - - for record in results: - content = record.get("content") if isinstance(record, dict) else None - if not content: - continue - - decision = from_dict_safe(DecisionLog, content) - if not decision: - continue - - if decision_type and decision.decision_type != decision_type: - continue - - if cutoff_date and decision.created_at: - try: - created = datetime.fromisoformat(decision.created_at.replace("Z", "+00:00")) - # Records written before this fix are naive; assume UTC. - if created.tzinfo is None: - created = created.replace(tzinfo=timezone.utc) - if created < cutoff_date: - continue - except (ValueError, AttributeError): - pass - - if query: - query_lower = query.lower() - text = decision.to_text().lower() - if query_lower not in text: - continue - - decisions.append(decision) - - if len(decisions) >= limit: - break - - return decisions - + return rows or [] except Exception as e: - log_debug(f"DecisionLogStore.asearch failed: {e}") + log_debug(f"DecisionLogStore._afetch_recent_rows failed: {e}") return [] + def _filter_decisions( + self, + records: List[Any], + query: Optional[str], + decision_type: Optional[str], + days: Optional[int], + limit: int, + ) -> List[DecisionLog]: + decisions: List[DecisionLog] = [] + cutoff_date = None + if days: + cutoff_date = datetime.now(timezone.utc) - timedelta(days=days) + + for record in records: + content = record.get("content") if isinstance(record, dict) else None + if not content: + continue + + decision = from_dict_safe(DecisionLog, content) + if not decision: + continue + + if decision_type and decision.decision_type != decision_type: + continue + + if cutoff_date and decision.created_at: + try: + created = datetime.fromisoformat(decision.created_at.replace("Z", "+00:00")) + # A stored timestamp without a zone is UTC. + if created.tzinfo is None: + created = created.replace(tzinfo=timezone.utc) + if created < cutoff_date: + continue + except (ValueError, AttributeError): + pass + + if query: + # Value-scoped verification: the db-side ILIKE matched the whole + # serialized document (keys included); this check keeps the match + # surface at the record's values, across every field. + if not values_match_query(content, query): + continue + + decisions.append(decision) + + if len(decisions) >= limit: + break + + return decisions + def get(self, decision_id: str) -> Optional[DecisionLog]: """Get a specific decision by ID.""" if not self.db: @@ -1025,82 +1036,6 @@ class DecisionLogStore(LearningStore): await self.asave(decision=decision) return True - # ========================================================================= - # Extraction (ALWAYS mode) - # ========================================================================= - - def _extract_decisions_from_messages( - self, - messages: List["Message"], - agent_id: Optional[str] = None, - session_id: Optional[str] = None, - user_id: Optional[str] = None, - team_id: Optional[str] = None, - ) -> None: - """Extract decisions from tool calls in messages.""" - for msg in messages: - if not hasattr(msg, "tool_calls") or not msg.tool_calls: - continue - - for tool_call in msg.tool_calls: - tool_name = getattr(tool_call, "name", None) or getattr( - getattr(tool_call, "function", None), "name", None - ) - - if not tool_name: - continue - - decision_id = f"dec_{uuid.uuid4().hex[:8]}" - decision = DecisionLog( - id=decision_id, - decision=f"Called tool: {tool_name}", - decision_type="tool_selection", - context="During conversation with user", - session_id=session_id, - user_id=user_id, - agent_id=agent_id, - team_id=team_id, - created_at=datetime.now(timezone.utc).isoformat(), - ) - - self.save(decision=decision) - - async def _aextract_decisions_from_messages( - self, - messages: List["Message"], - agent_id: Optional[str] = None, - session_id: Optional[str] = None, - user_id: Optional[str] = None, - team_id: Optional[str] = None, - ) -> None: - """Async version of _extract_decisions_from_messages.""" - for msg in messages: - if not hasattr(msg, "tool_calls") or not msg.tool_calls: - continue - - for tool_call in msg.tool_calls: - tool_name = getattr(tool_call, "name", None) or getattr( - getattr(tool_call, "function", None), "name", None - ) - - if not tool_name: - continue - - decision_id = f"dec_{uuid.uuid4().hex[:8]}" - decision = DecisionLog( - id=decision_id, - decision=f"Called tool: {tool_name}", - decision_type="tool_selection", - context="During conversation with user", - session_id=session_id, - user_id=user_id, - agent_id=agent_id, - team_id=team_id, - created_at=datetime.now(timezone.utc).isoformat(), - ) - - await self.asave(decision=decision) - # ========================================================================= # Representation # ========================================================================= diff --git a/libs/agno/agno/learn/stores/entity_memory.py b/libs/agno/agno/learn/stores/entity_memory.py index 6df0720abb..77158a5ddf 100644 --- a/libs/agno/agno/learn/stores/entity_memory.py +++ b/libs/agno/agno/learn/stores/entity_memory.py @@ -11,64 +11,358 @@ Think of it as: - UserProfile = what you know about THE USER - EntityMemory = what you know about EVERYTHING ELSE -Key Features: -- Entity-scoped storage (entity_id + entity_type) -- Three types of memory per entity: - - Facts (semantic): Timeless truths ("Acme uses PostgreSQL") - - Events (episodic): Time-bound occurrences ("Acme launched v2 on Jan 15") - - Relationships (graph): Connections to other entities ("Bob is CEO of Acme") -- Namespace-based sharing control -- Agent tools for CRUD operations -- Background extraction from conversations +The agent surface is four tools: +- remember_about: upsert an entity with facts, events and an optional note pointer +- link_entities: record a relationship between two entities +- search_entities: search stored entities, or list them by recency +- forget: retire a fact, or archive a whole entity Scoping: -- entity_id: Unique identifier (e.g., "acme_corp", "bob_smith") -- entity_type: Category (e.g., "company", "person", "project", "product") -- namespace: Sharing scope: +- entity_id: derived in the store from the entity's name (slugified) +- entity_type: category (e.g., "company", "person", "project", "product") +- namespace: sharing scope: - "user": Private to current user - "global": Shared with everyone (default) - "": Custom grouping (e.g., "sales_team") Supported Modes: -- ALWAYS: Automatic extraction of entity info from conversations -- AGENTIC: Agent calls tools directly to manage entity info +- AGENTIC only. The agent records entities through tools; there is no + extraction pass. This mirrors how session_context documents itself as + ALWAYS-only. """ -from copy import deepcopy +import re from dataclasses import dataclass, field from datetime import datetime, timezone from os import getenv from textwrap import dedent -from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Union, cast +from typing import Any, Callable, Dict, Iterator, List, Optional, Sequence, Tuple, Union, cast +from weakref import WeakKeyDictionary from agno.learn.config import EntityMemoryConfig, LearningMode from agno.learn.schemas import EntityMemory from agno.learn.stores.protocol import LearningStore -from agno.learn.utils import build_learning_id +from agno.learn.utils import build_learning_id, values_match_query from agno.utils.log import ( log_debug, + log_info, log_warning, set_log_level_to_debug, set_log_level_to_info, ) -from agno.utils.message import get_conversation_text - -if TYPE_CHECKING: - from agno.metrics import RunMetrics try: from agno.db.base import AsyncBaseDb, BaseDb - from agno.models.message import Message except ImportError: pass +def _utc_now_iso() -> str: + return datetime.now(timezone.utc).isoformat().replace("+00:00", "Z") + + +# Symbols that name a thing rather than separate words. Without these, "C++" +# and "C#" both key to `c` and merge into "C" - a wrong merge, and there is no +# unmerge. +_IDENTITY_SYMBOLS = {"++": "plus_plus", "+": "plus", "#": "sharp"} + + +def _slugify(name: str) -> str: + """Derive a stable entity_id from a display name: lowercase, underscores. + + Accented letters fold to their base letter first, so Anna Müller and Anna + Möller keep distinct ids while "Sofia Munoz" still resolves to the + "Sofía Muñoz" already on file. + + Letters that survive the fold are KEPT, whatever their script. Dropping + them meant a name only partly Latin lost the half that identified it - + "李 Ming" keyed to `ming` and merged with an unrelated Ming, the same + no-unmerge corruption the fold exists to prevent. + + A few symbols carry identity rather than punctuation: C, C++ and C# are + three languages, and collapsing every non-word character made them one + entity holding all three's facts, with the discarded names written + nowhere. Those spell themselves out; the rest still collapse, so + "Acme, Inc." and "Acme Inc" stay one company. + """ + import unicodedata + + decomposed = unicodedata.normalize("NFKD", name.strip().lower()) + folded = "".join(ch for ch in decomposed if not unicodedata.combining(ch)) + for symbol, word in _IDENTITY_SYMBOLS.items(): + folded = folded.replace(symbol, f"_{word}_") + slug = re.sub(r"[^\w]+", "_", folded, flags=re.UNICODE).strip("_") + return slug or name.strip().lower() + + +def _slugify_or_none(name: Optional[str]) -> Optional[str]: + """Slugify, returning None when the name carries no usable identity. + + Pure-punctuation names ("???", "%") are rejected; non-ASCII names pass + through (the ASCII slugifier cannot fold them, so the lowered name itself + is the id). + """ + if not name or not name.strip(): + return None + slug = _slugify(name) + if not slug or not re.search(r"[^\W_]", slug, re.UNICODE): + return None + return slug + + +def _normalize_fact_text(text: str) -> str: + """Fold case and collapse whitespace for fact matching.""" + return re.sub(r"\s+", " ", text.strip().casefold()) + + +def _normalize_name(name: str) -> str: + """Fold case and collapse whitespace for name and alias matching.""" + return re.sub(r"\s+", " ", name.strip().casefold()) + + +_MESSAGE_STOPWORDS = frozenset( + """a an and are as at be been but by can could did do does for from had has have how i if in into is it its me my + of on or our she so that the their them then there they this to was we were what when where which who why will + with would you your about tell know new more some any all just like get make said says told asked shall + please show give need want us him her his hers ours yours theirs also very much many still over after before + between should must may might each other than too not no yes ok okay let lets going go went come came thing + things something anything everything nothing someone anyone everyone whats hows whos""".split() +) + + +def _message_terms(message: str, max_terms: int = 8) -> List[str]: + """Distinctive lexical terms of a message, for relevance recall. + + Deliberately lexical (no semantic retrieval, per the non-goals): lowercase + word tokens, minus stopwords and short words, deduplicated in order. + """ + terms: List[str] = [] + for word in re.findall(r"[\w][\w\-']*", message.lower()): + if len(word) < 3 or word in _MESSAGE_STOPWORDS: + continue + if word not in terms: + terms.append(word) + if len(terms) >= max_terms: + break + return terms + + +# The DB key is entity_{namespace}_{entity_type}_{entity_id}, so type drift +# splits entities exactly like name drift: "person" on Monday and "people" on +# Friday is two Sarahs. Fold case and singularize against this canonical list; +# anything else passes through lowercased. +_CANONICAL_ENTITY_TYPES = frozenset({"person", "project", "company", "system", "product"}) +# The placeholder link_entities mints for an endpoint it has never been told +# about; it is the one type that merges into a real one. +_UNKNOWN_ENTITY_TYPE = "unknown" +_IRREGULAR_ENTITY_TYPES = {"people": "person", "persons": "person", "companies": "company"} + + +def _normalize_entity_type(entity_type: Optional[str]) -> Optional[str]: + if entity_type is None or not entity_type.strip(): + return entity_type + normalized = re.sub(r"\s+", "_", entity_type.strip().lower()) + if normalized in _CANONICAL_ENTITY_TYPES: + return normalized + if normalized in _IRREGULAR_ENTITY_TYPES: + return _IRREGULAR_ENTITY_TYPES[normalized] + if normalized.endswith("s") and normalized[:-1] in _CANONICAL_ENTITY_TYPES: + return normalized[:-1] + return normalized + + +def _blank_to_none(value: Optional[str]) -> Optional[str]: + """Blank string in, ``None`` out. + + Models fill unused optional arguments with "" rather than omitting them + (a strict tool schema has no way to say "absent"), so a blank argument + means "not supplied". Passed through, it would become a filter that + matches nothing and read as an honest "not found". + """ + if value is None: + return None + return value.strip() or None + + +def _split_qualified_name(entity: str, known_types: Sequence[str]) -> Tuple[str, Optional[str]]: + """Split a "type/name" reference into (name, type). + + Two entities can share a name under different types, and neither + link_entities nor forget carries an entity_type; the ambiguity replies tell + the model to name one as "project/Harbor", so both tools read it back. + + Only a prefix that actually names one of the stored types counts. A name + containing a slash is far more likely to be a name - "AC/DC" is a band, not + a DC of type AC - and treating it as qualified sends resolution to a key + the write path never uses. + """ + if "/" not in entity: + return entity, None + prefix, _, rest = entity.partition("/") + prefix, rest = prefix.strip(), rest.strip() + if not prefix or not rest: + return entity, None + if _normalize_entity_type(prefix) not in {_normalize_entity_type(t) for t in known_types}: + return entity, None + return rest, prefix + + +def _mentions(needle: str, term: str) -> bool: + """Whether ``needle`` names ``term`` as a whole word. + + Plain containment let a longer far-end name swallow a shorter one: + "designed_by -> Sarah Chen" matched the edge to `sarah` as well, and the + qualified form the refusal then asks for could not break the tie either, + because "person/sarah" is itself a substring of "person/sarah_chen". The + link became unretireable, which is the one thing forget exists to prevent. + """ + return re.search(r"(? bool: + """Whether the needle names an entity whose slug is ``term_slug``. + + The context block renders a far end's DISPLAY name, and a name is not + recoverable from its slug once the slug has folded an accent, a hyphen or + an ampersand: "supplier_is -> Café Noir" is the wording the docstring asks + for and `cafe_noir` is what the edge stores. Slugify both and compare on + token boundaries. + """ + needle_tokens = [t for t in _slugify(needle).split("_") if t] + term_tokens = [t for t in term_slug.split("_") if t] + if not term_tokens or len(term_tokens) > len(needle_tokens): + return False + return any( + needle_tokens[i : i + len(term_tokens)] == term_tokens for i in range(len(needle_tokens) - len(term_tokens) + 1) + ) + + +def _types_can_merge(incoming: Optional[str], existing: Optional[str]) -> bool: + """Whether a name match across two entity types is the same thing. + + Only the ``unknown`` placeholder ``link_entities`` mints merges across + types - describing it later is what gives it a real type. Two named types + are two things: Atlas the system is not Atlas the team, and a merge has no + unmerge. Models coin types freely, so restricting this to a canonical list + left the common case (``team``, ``framework``, ``service``) merging. + + Fragmenting is the recoverable failure here - both rows render, and the + write path says when a same-name sibling exists so the model can correct + itself. + """ + incoming_type = _normalize_entity_type(incoming) + existing_type = _normalize_entity_type(existing) + if not incoming_type or not existing_type or incoming_type == existing_type: + return True + return _UNKNOWN_ENTITY_TYPE in (incoming_type, existing_type) + + +# ============================================================================= +# Tool docstrings (shared between the sync and async tool variants) +# ============================================================================= + +_REMEMBER_ABOUT_DOC = """Record something about an entity - a person, project, company, system, or product. + +Upserts: the entity is created if new, merged into if already known. Refer to the +entity by name ("Sarah Chen", "radar") - never invent an id. + +What goes where: +- facts: one-line current values you expect to be replaced ("db: Postgres - see note"). +- events: things that happened on a date ("shipped v1 on 2026-07-20"). Positions and + opinions are events, not facts. +- note: the path of the note file holding the detail this entity indexes + (e.g. "notes/radar.md"). Set it whenever the content lives in a note. + +Args: + entity: The entity's name as people say it (e.g. "Sarah Chen", "radar"). + entity_type: Category: person, project, company, system, product - or another short noun. + description: One-line description of what this entity is. + facts: One-line facts to record on the entity. + events: Dated occurrences to record. + note: Path of the note file with the full detail (e.g. "notes/radar.md"). + +Returns: + Confirmation of what was recorded. +""" + +_LINK_ENTITIES_DOC = """Link two entities with a relationship ("Sarah Chen" works_on "radar"). + +Both ends are resolved by name. An end that is not known yet is created as a +minimal entity, so it is safe to link first and describe later. The link is +stored on both entities, so it is visible from either side. + +A link that stops being true is removed with forget, naming it the way it is +rendered ("written_in -> Rust"); recording the new one does not retire the old. + +Args: + entity: Source entity name. If two entities share it, name one as + "project/Harbor" - the reply tells you when that is needed. + relation: The relationship, a short verb phrase ("works_on", "owns", "uses"). + related_entity: Target entity name. + +Returns: + Confirmation of the recorded link. +""" + +_SEARCH_ENTITIES_DOC = """Search stored entities, or list them. + +With a query, matches entity names, facts, events and relationships. Without a +query, lists entities by recency - use that to browse what exists ("who works on +what"). Results include each entity's note path; follow it with read_file when you +need the detail behind an indexed line. + +Args: + query: Text to match (a name, a fact fragment). Omit to list entities by recency. + entity_type: Optional filter: person, project, company, system, product, etc. + +Returns: + Matching entities with their facts, events and relationships, or a listing. +""" + +_SUPERSESSION_SYSTEM_PROMPT = """You judge fact supersession for an entity's memory. + +You are given the entity's live facts (each with an id) and newly stated facts. +Identify which OLD facts the new facts contradict or replace. + +A new fact supersedes an old one when both cannot be true at the same time - a +changed value, a corrected status, a reversed decision - or when the new fact is +a direct update of the same attribute. Related but compatible facts do NOT +supersede each other. + +Call mark_superseded exactly once with the ids of the superseded old facts and +your confidence (0.0 to 1.0) for each, as parallel lists. If nothing is +superseded, call it with two empty lists. Be conservative: when unsure, do not +supersede - a wrong supersession hides information the user gave us. +""" + + +_FORGET_DOC = """Retire a fact or a relationship from an entity, or archive the whole entity. + +With fact: retires the matching fact - it stops being recalled, nothing is deleted. +A relationship is retired the same way, naming it as it is rendered +("written_in -> Rust"): use this when a link is no longer true, because stating +the new link does not remove the old one. +Without fact: archives the entity. An archived entity leaves recall and the entity +directory, stays findable via search_entities, and any later remember_about about +it revives it. + +Args: + entity: The entity's name. If two entities share it, name one as + "project/Harbor" - the reply tells you when that is needed. + fact: The fact to retire, worded as closely as you can to how it was stored, + or the relationship to remove ("works_on -> radar"). + +Returns: + Confirmation of what was retired, removed or archived. +""" + + @dataclass class EntityMemoryStore(LearningStore): """Storage backend for Entity Memory learning type. Stores knowledge about external entities with three types of memory: - - **Facts**: Semantic memory - timeless truths about the entity + - **Facts**: Semantic memory - current truths about the entity - **Events**: Episodic memory - time-bound occurrences - **Relationships**: Graph edges - connections to other entities @@ -85,14 +379,25 @@ class EntityMemoryStore(LearningStore): # State tracking (internal) entity_updated: bool = field(default=False, init=False) _schema: Any = field(default=None, init=False) + _degraded_search_logged: bool = field(default=False, init=False) + _async_db_in_sync_logged: bool = field(default=False, init=False) + # Every write here is a read-modify-write over a whole row, and an + # assistant turn's tool calls run concurrently (models/base.py gathers + # them), so two tools touching one entity would overwrite each other and + # both report success. Writes take this lock; reads and the judge's + # provider call stay outside it. Keyed weakly by event loop - see + # _write_lock. + _async_write_locks: Any = field(default_factory=WeakKeyDictionary, init=False) def __post_init__(self): self._schema = self.config.schema or EntityMemory - if self.config.mode == LearningMode.PROPOSE: - log_warning("EntityMemoryStore does not support PROPOSE mode. Falling back to ALWAYS mode.") - elif self.config.mode == LearningMode.HITL: - log_warning("EntityMemoryStore does not support HITL mode. Falling back to ALWAYS mode.") + if self.config.mode != LearningMode.AGENTIC: + raise ValueError( + f"EntityMemoryStore is AGENTIC-only: the agent records entities through its tools " + f"and there is no extraction pass. Remove mode={self.config.mode.value!r} from " + f"EntityMemoryConfig or set LearningMode.AGENTIC." + ) # ========================================================================= # LearningStore Protocol Implementation @@ -114,198 +419,399 @@ class EntityMemoryStore(LearningStore): entity_type: Optional[str] = None, user_id: Optional[str] = None, namespace: Optional[str] = None, + message: Optional[str] = None, **kwargs, ) -> Optional[Any]: - """Retrieve entity memory from storage. + """Retrieve entity memory for context injection. + + Always returns the entity directory (name + type, newest first, + archived excluded) so the agent can see what exists and ground its + negatives - "not in the directory" is a fact, "substring didn't match" + is not. With a message, the top-k relevant entities come back expanded: + entities whose name or alias appears in the message, then (still under + k) entities matched by a bounded lexical search over the message's + distinctive terms. A keyed lookup (entity_id + entity_type) stays + available and is expanded first. Archived entities are excluded from + recall (they stay reachable via search). Args: - entity_id: The entity to retrieve (required with entity_type). - entity_type: The type of entity (required with entity_id). + entity_id: Optional entity to expand (with entity_type). + entity_type: The type of entity (with entity_id). user_id: User ID for "user" namespace scoping. namespace: Filter by namespace. + message: The current message; drives relevance recall. **kwargs: Additional context (ignored). Returns: - Entity memory, or None if not found. + Dict with "directory" (all known entities, bounded), "entities" + (the expanded ones) and "related_names" (the one-hop name map), or + None when the namespace cannot be read. """ - if not entity_id or not entity_type: - return None - effective_namespace = namespace or self.config.namespace if effective_namespace == "user" and not user_id: - log_warning("EntityMemoryStore.process: namespace='user' requires user_id") + log_warning("EntityMemoryStore.recall: namespace='user' requires user_id") return None - return self.get( - entity_id=entity_id, - entity_type=entity_type, + # One row beyond the cap, so a truncated directory is distinguishable + # from one that happens to be exactly at the cap. + directory = self.list_entities( user_id=user_id, namespace=effective_namespace, + limit=self.config.max_entities_in_directory + 1, ) + entities: List[EntityMemory] = [] + if entity_id and entity_type: + entity = self.get( + entity_id=entity_id, + entity_type=entity_type, + user_id=user_id, + namespace=effective_namespace, + ) + if entity is not None and not getattr(entity, "archived_at", None): + entities.append(entity) + + if message: + top_k = self.config.max_entities_in_context + entities = self._merge_relevant(entities, self._match_directory_by_message(directory, message), top_k) + if len(entities) < top_k: + # All selected terms are searched and merged round-robin, so a + # generic first term cannot consume the budget before the + # distinctive one. + term_results = [ + self.search(query=term, user_id=user_id, namespace=effective_namespace, limit=top_k) + for term in self._terms_to_search(message, entities) + ] + entities = self._merge_relevant(entities, self._interleave(term_results), top_k) + + related_names = self._related_names( + entities=entities, directory=directory, user_id=user_id, namespace=effective_namespace + ) + return {"directory": directory, "entities": entities, "related_names": related_names} + async def arecall( self, entity_id: Optional[str] = None, entity_type: Optional[str] = None, user_id: Optional[str] = None, namespace: Optional[str] = None, + message: Optional[str] = None, **kwargs, ) -> Optional[Any]: """Async version of recall.""" - if not entity_id or not entity_type: - return None - effective_namespace = namespace or self.config.namespace if effective_namespace == "user" and not user_id: log_warning("EntityMemoryStore.arecall: namespace='user' requires user_id") return None - return await self.aget( - entity_id=entity_id, - entity_type=entity_type, + directory = await self.alist_entities( user_id=user_id, namespace=effective_namespace, + limit=self.config.max_entities_in_directory + 1, ) - def process( - self, - messages: List[Any], - user_id: Optional[str] = None, - agent_id: Optional[str] = None, - team_id: Optional[str] = None, - namespace: Optional[str] = None, - **kwargs, - ) -> None: - """Extract entity information from messages. + entities: List[EntityMemory] = [] + if entity_id and entity_type: + entity = await self.aget( + entity_id=entity_id, + entity_type=entity_type, + user_id=user_id, + namespace=effective_namespace, + ) + if entity is not None and not getattr(entity, "archived_at", None): + entities.append(entity) - Args: - messages: Conversation messages to analyze. - user_id: User context (for "user" namespace scoping). - agent_id: Agent context (stored for audit). - team_id: Team context (stored for audit). - namespace: Namespace to save entities to. - **kwargs: Additional context (ignored). + if message: + top_k = self.config.max_entities_in_context + entities = self._merge_relevant(entities, self._match_directory_by_message(directory, message), top_k) + if len(entities) < top_k: + term_results = [] + for term in self._terms_to_search(message, entities): + term_results.append( + await self.asearch(query=term, user_id=user_id, namespace=effective_namespace, limit=top_k) + ) + entities = self._merge_relevant(entities, self._interleave(term_results), top_k) + + related_names = await self._arelated_names( + entities=entities, directory=directory, user_id=user_id, namespace=effective_namespace + ) + return {"directory": directory, "entities": entities, "related_names": related_names} + + @staticmethod + def _match_directory_by_message(directory: List[EntityMemory], message: str) -> List[EntityMemory]: + """Directory entries whose name or alias appears in the message. + + Word-boundary matched: a two-letter entity named "Al" must not match + "always" and evict the entity the turn is actually about. """ - if self.config.mode == LearningMode.AGENTIC: - return + normalized_message = _normalize_name(message) + matched: List[EntityMemory] = [] + for entity in directory: + names = [getattr(entity, "name", None) or ""] + list(getattr(entity, "aliases", None) or []) + for candidate in names: + if not candidate: + continue + pattern = r"(? List[EntityMemory]: + """Round-robin merge, so one term's results cannot starve another's.""" + merged: List[EntityMemory] = [] + index = 0 + while True: + advanced = False + for results in result_lists: + if index < len(results): + merged.append(results[index]) + advanced = True + if not advanced: + return merged + index += 1 - effective_namespace = namespace or self.config.namespace - self.extract_and_save( - messages=messages, - user_id=user_id, - agent_id=agent_id, - team_id=team_id, - namespace=effective_namespace, - run_metrics=kwargs.get("run_metrics"), - ) + @staticmethod + def _merge_relevant(current: List[EntityMemory], additions: List[EntityMemory], top_k: int) -> List[EntityMemory]: + merged = list(current) + seen = {(e.entity_id, e.entity_type) for e in merged} + for entity in additions: + key = (entity.entity_id, entity.entity_type) + if key in seen: + continue + if getattr(entity, "archived_at", None): + continue + seen.add(key) + merged.append(entity) + if len(merged) >= top_k: + break + return merged[: max(top_k, 0)] - async def aprocess( + @staticmethod + def _terms_to_search(message: str, already_matched: List[EntityMemory], max_searches: int = 3) -> List[str]: + """Message terms worth a bounded lexical search. + + Terms whose token already appears in a matched entity's name are + skipped; the rest are ordered proper-noun-ish first (capitalized + mid-sentence in the original message - the strongest name signal), + then by length, so a generic early word cannot consume the search + budget before the distinctive one. + """ + covered_tokens: set = set() + for e in already_matched: + covered_tokens.update(_normalize_name(getattr(e, "name", None) or e.entity_id).split()) + + capitalized = {match.group(0).lower() for match in re.finditer(r"(? None: - """Async version of process.""" - if self.config.mode == LearningMode.AGENTIC: - return + entities: List[EntityMemory], + directory: List[EntityMemory], + user_id: Optional[str], + namespace: str, + max_lookups: int = 10, + ) -> Dict[str, str]: + """One-hop link expansion: the display names of the entities the + expanded entities link to. Names come from the already-fetched + directory where possible; the remainder is a bounded set of keyed + lookups (no recursion, no content).""" + name_map = { + e.entity_id: (getattr(e, "name", None) or e.entity_id) for e in directory if getattr(e, "entity_id", None) + } + missing = self._missing_link_targets(entities=entities, known=name_map) + for far_id, far_type in missing[:max_lookups]: + far = self.get(entity_id=far_id, entity_type=far_type, user_id=user_id, namespace=namespace) + if far is None: + # The edge's stored type can go stale (unknown -> real upgrade); + # resolve the id across types before giving up on the name. + for row in self._get_rows_by_entity_id(entity_id=far_id, user_id=user_id, namespace=namespace): + far = self.schema.from_dict(row.get("content")) + if far is not None: + break + if far is not None: + name = getattr(far, "name", None) or far_id + if getattr(far, "archived_at", None): + name += " (archived)" + name_map[far_id] = name + return name_map - if not messages: - return + async def _arelated_names( + self, + entities: List[EntityMemory], + directory: List[EntityMemory], + user_id: Optional[str], + namespace: str, + max_lookups: int = 10, + ) -> Dict[str, str]: + """Async version of _related_names.""" + name_map = { + e.entity_id: (getattr(e, "name", None) or e.entity_id) for e in directory if getattr(e, "entity_id", None) + } + missing = self._missing_link_targets(entities=entities, known=name_map) + for far_id, far_type in missing[:max_lookups]: + far = await self.aget(entity_id=far_id, entity_type=far_type, user_id=user_id, namespace=namespace) + if far is None: + for row in await self._aget_rows_by_entity_id(entity_id=far_id, user_id=user_id, namespace=namespace): + far = self.schema.from_dict(row.get("content")) + if far is not None: + break + if far is not None: + name = getattr(far, "name", None) or far_id + if getattr(far, "archived_at", None): + name += " (archived)" + name_map[far_id] = name + return name_map - effective_namespace = namespace or self.config.namespace - await self.aextract_and_save( - messages=messages, - user_id=user_id, - agent_id=agent_id, - team_id=team_id, - namespace=effective_namespace, - run_metrics=kwargs.get("run_metrics"), - ) + @staticmethod + def _missing_link_targets(entities: List[EntityMemory], known: Dict[str, str]) -> List[Tuple[str, str]]: + missing: List[Tuple[str, str]] = [] + seen = set(known) + for entity in entities: + for rel in getattr(entity, "relationships", None) or []: + if not isinstance(rel, dict): + continue + far_id = rel.get("entity_id") + far_type = rel.get("entity_type") + if far_id and far_type and far_id not in seen: + seen.add(far_id) + missing.append((far_id, far_type)) + return missing + + def process(self, messages: List[Any], **kwargs) -> None: + """No-op: entity memory is AGENTIC-only, capture happens through the tools.""" + return + + async def aprocess(self, messages: List[Any], **kwargs) -> None: + """Async version of process (no-op).""" + return def build_context(self, data: Any) -> str: """Build context for the agent. - Formats entity memory for injection into the agent's system prompt. - Entity memory provides knowledge about external things - people, companies, - projects, products - distinct from knowledge about the user themselves. + Formats entity memory for injection into the agent's system prompt: + the entity directory first (always, when anything exists), then the + expanded entities, bounded by the config's rendering knobs. Args: - data: Entity memory data from recall() - single entity or list. + data: Data from recall() - the {"directory", "entities"} dict, or a + single entity / list of entities from direct calls. Returns: Context string to inject into the agent's system prompt. """ - if not data: + directory: List[EntityMemory] = [] + entities: List[Any] = [] + related_names: Dict[str, str] = {} + if isinstance(data, dict) and "directory" in data: + directory = data.get("directory") or [] + entities = data.get("entities") or [] + related_names = data.get("related_names") or {} + elif data: + entities = data if isinstance(data, list) else [data] + + if not directory and not entities: if self._should_expose_tools: return dedent("""\ - - You have access to entity memory - a knowledge base about people, companies, - projects, products, and other external entities relevant to your work. - - **Available Tools:** - - `search_entities`: Find stored information about entities - - `create_entity`: Store a new entity with its facts - - `add_fact`: Add a timeless truth about an entity - - `add_event`: Record a time-bound occurrence - - `add_relationship`: Capture connections between entities - - **When to use entity memory:** - - You learn something substantive about a company, person, or project - - Information would be useful to recall in future conversations - - Facts are stable enough to be worth storing - - **Entity memory vs other memory types:** - - User memory = about THE USER (their preferences, role, context) - - Entity memory = about EXTERNAL THINGS (companies, people, projects) - - Learned knowledge = reusable TASK insights (patterns, approaches) - """) + + No entities recorded yet. + """) return "" - # Handle single entity or list - entities = data if isinstance(data, list) else [data] - if not entities: - return "" + sections: List[str] = [] - # Use schema's get_context_text - formatted_parts = [] - for entity in entities: - if hasattr(entity, "get_context_text"): - formatted_parts.append(entity.get_context_text()) + directory_truncated = False + if directory: + cap = self.config.max_entities_in_directory + directory_truncated = len(directory) > cap + listed = directory[:cap] + lines = [] + for entity in listed: + name = getattr(entity, "name", None) or getattr(entity, "entity_id", "?") + lines.append(f"- {name} ({getattr(entity, 'entity_type', '?')})") + if directory_truncated: + header = f"**Entity directory** (the {len(listed)} most recently updated entities; more exist)" else: - formatted_parts.append(self._format_entity_basic(entity=entity)) + header = "**Entity directory** (all known entities, newest first)" + sections.append(header + ":\n" + "\n".join(lines)) - formatted = "\n\n---\n\n".join(formatted_parts) + shown_entities = entities[: self.config.max_entities_in_context] + if shown_entities: + formatted_parts = [] + render_kwargs = { + "max_facts": self.config.max_facts_per_entity, + "max_events": self.config.max_events_per_entity, + "related_names": related_names, + } + for entity in shown_entities: + render = getattr(entity, "get_context_text", None) + if callable(render): + # A custom schema may define its own signature; pass only what + # it accepts (never re-invoke on error - side effects). + import inspect as _inspect - context = dedent(f"""\ + try: + params = _inspect.signature(render).parameters + has_var_kw = any(p.kind == p.VAR_KEYWORD for p in params.values()) + accepted = {k: v for k, v in render_kwargs.items() if has_var_kw or k in params} + except (TypeError, ValueError): + accepted = {} + formatted_parts.append(render(**accepted)) + else: + formatted_parts.append(self._format_entity_basic(entity=entity)) + sections.append("**Known information about relevant entities:**\n\n" + "\n\n---\n\n".join(formatted_parts)) + + body = "\n\n".join(sections) + context = dedent("""\ - **Known information about relevant entities:** - - {formatted} + {body} Use this knowledge naturally in your responses: - Reference stored facts without citing "entity memory" - Treat this as background knowledge you simply have - Current conversation takes precedence if there's conflicting information - - Update entity memory if you learn something new and substantive + - {directory_line} - """) - - if self._should_expose_tools: - context += dedent(""" - Entity memory tools are available to search, create, or update entities. - """) - else: - context += "" + """).format( + body=body, + directory_line=( + "The directory shows the most recently updated entities; use search_entities to check beyond it" + if directory_truncated + else "The directory is the full index: an entity not listed there is not known" + ), + ) return context + def instructions(self) -> str: + """Agent-facing guidance for this store: the four tools and when to use them. + + Guidance only - the recalled entities and directory live in + build_context(). + """ + if not self._should_expose_tools: + return "" + return dedent("""\ + + You have entity memory - a knowledge base about the people, companies, + projects, systems and products relevant to your work. + + - `remember_about`: record facts, events, a description or a note pointer + on an entity, by name. A correction is just the new fact: state it, and + the contradicted old fact is retired automatically. + - `link_entities`: record a relationship between two entities. + - `search_entities`: find stored entities; with no query, list them by + recency (the browse surface). + - `forget`: retire a fact, or archive a whole entity. + + Record something whenever you learn substantive information about a + company, person, project or system that future conversations will need. + """) + def get_tools( self, user_id: Optional[str] = None, @@ -314,7 +820,7 @@ class EntityMemoryStore(LearningStore): namespace: Optional[str] = None, **kwargs, ) -> List[Callable]: - """Get tools to expose to agent. + """Get the four agent tools (sync variants). Args: user_id: User context (for "user" namespace scoping). @@ -328,7 +834,8 @@ class EntityMemoryStore(LearningStore): """ if not self._should_expose_tools: return [] - return self.get_agent_tools( + return self._build_agent_tools( + async_mode=False, user_id=user_id, agent_id=agent_id, team_id=team_id, @@ -343,10 +850,11 @@ class EntityMemoryStore(LearningStore): namespace: Optional[str] = None, **kwargs, ) -> List[Callable]: - """Async version of get_tools.""" + """Async version of get_tools: the same four tools as async callables.""" if not self._should_expose_tools: return [] - return await self.aget_agent_tools( + return self._build_agent_tools( + async_mode=True, user_id=user_id, agent_id=agent_id, team_id=team_id, @@ -360,13 +868,8 @@ class EntityMemoryStore(LearningStore): @property def _should_expose_tools(self) -> bool: - """Check if tools should be exposed to the agent. - - Returns True if either: - - mode is AGENTIC (tools are the primary way to manage entities), OR - - enable_agent_tools is explicitly True - """ - return self.config.mode == LearningMode.AGENTIC or self.config.enable_agent_tools + """Whether the four tools are exposed to the agent.""" + return self.config.enable_agent_tools # ========================================================================= # Properties @@ -379,7 +882,7 @@ class EntityMemoryStore(LearningStore): @property def model(self): - """Model for extraction.""" + """Model for the fact-supersession judgment.""" return self.config.model # ========================================================================= @@ -395,1237 +898,1962 @@ class EntityMemoryStore(LearningStore): set_log_level_to_info() # ========================================================================= - # Agent Tools + # Agent Tools (one factory generates the sync and async variants) # ========================================================================= - def get_agent_tools( + def _build_agent_tools( self, + async_mode: bool, user_id: Optional[str] = None, agent_id: Optional[str] = None, team_id: Optional[str] = None, namespace: Optional[str] = None, ) -> List[Callable]: - """Get the tools to expose to the agent. + """Build the four agent tools, closing over the run's identity context. - Tools are included based on config settings: - - search_entities (agent_can_search_entities) - - create_entity (agent_can_create_entity) - - update_entity (agent_can_update_entity) - - add_fact, update_fact, delete_fact - - add_event - - add_relationship + The sync and async variants share their docstrings (the model-facing + contract) and both delegate to the store's public write methods. + """ + store = self + effective_namespace = namespace or self.config.namespace - Args: - user_id: User context (for "user" namespace scoping). - agent_id: Agent context (stored for audit). - team_id: Team context (stored for audit). - namespace: Default namespace for operations. + if async_mode: + + async def remember_about( + entity: str, + entity_type: str, + description: Optional[str] = None, + facts: List[str] = [], + events: List[str] = [], + note: Optional[str] = None, + ) -> str: + return await store.aremember_about( + entity=entity, + entity_type=entity_type, + description=description, + facts=facts, + events=events, + note=note, + user_id=user_id, + agent_id=agent_id, + team_id=team_id, + namespace=effective_namespace, + ) + + async def link_entities(entity: str, relation: str, related_entity: str) -> str: + return await store.alink_entities( + entity=entity, + relation=relation, + related_entity=related_entity, + user_id=user_id, + agent_id=agent_id, + team_id=team_id, + namespace=effective_namespace, + ) + + async def search_entities(query: Optional[str] = None, entity_type: Optional[str] = None) -> str: + return await store.asearch_entities( + query=query, + entity_type=entity_type, + user_id=user_id, + namespace=effective_namespace, + ) + + async def forget(entity: str, fact: Optional[str] = None) -> str: + return await store.aforget( + entity=entity, + fact=fact, + user_id=user_id, + agent_id=agent_id, + team_id=team_id, + namespace=effective_namespace, + ) + else: + + def remember_about( # type: ignore[misc] + entity: str, + entity_type: str, + description: Optional[str] = None, + facts: List[str] = [], + events: List[str] = [], + note: Optional[str] = None, + ) -> str: + return store.remember_about( + entity=entity, + entity_type=entity_type, + description=description, + facts=facts, + events=events, + note=note, + user_id=user_id, + agent_id=agent_id, + team_id=team_id, + namespace=effective_namespace, + ) + + def link_entities(entity: str, relation: str, related_entity: str) -> str: # type: ignore[misc] + return store.link_entities( + entity=entity, + relation=relation, + related_entity=related_entity, + user_id=user_id, + agent_id=agent_id, + team_id=team_id, + namespace=effective_namespace, + ) + + def search_entities( # type: ignore[misc] + query: Optional[str] = None, entity_type: Optional[str] = None + ) -> str: + return store.search_entities( + query=query, + entity_type=entity_type, + user_id=user_id, + namespace=effective_namespace, + ) + + def forget(entity: str, fact: Optional[str] = None) -> str: # type: ignore[misc] + return store.forget( + entity=entity, + fact=fact, + user_id=user_id, + agent_id=agent_id, + team_id=team_id, + namespace=effective_namespace, + ) + + remember_about.__doc__ = _REMEMBER_ABOUT_DOC + link_entities.__doc__ = _LINK_ENTITIES_DOC + search_entities.__doc__ = _SEARCH_ENTITIES_DOC + forget.__doc__ = _FORGET_DOC + + return [remember_about, link_entities, search_entities, forget] + + # ========================================================================= + # Public write API: remember_about + # ========================================================================= + + def remember_about( + self, + entity: str, + entity_type: str, + description: Optional[str] = None, + facts: Optional[List[str]] = None, + events: Optional[List[str]] = None, + note: Optional[str] = None, + aliases: Optional[List[str]] = None, + user_id: Optional[str] = None, + agent_id: Optional[str] = None, + team_id: Optional[str] = None, + namespace: Optional[str] = None, + ) -> str: + """Upsert an entity by name: create it if new, merge into it if known. Returns: - List of callable tools. + A confirmation message describing what was recorded. """ - tools = [] + if not self.db: + return "Entity memory has no database configured; nothing was recorded." + if not _slugify_or_none(entity): + return "Entity name is required; nothing was recorded." + + # A blank optional argument means the model did not supply it, not that + # it wants the stored value cleared - the tool surface has no way to + # clear one, so treating "" as a value silently wipes a description or + # a note pointer that a previous turn wrote. + description = _blank_to_none(description) + note = _blank_to_none(note) + effective_namespace = namespace or self.config.namespace + if effective_namespace == "user" and not user_id: + log_warning("EntityMemoryStore.remember_about: namespace='user' requires user_id") + return "Entity memory needs a user_id for the 'user' namespace; nothing was recorded." - if self.config.agent_can_search_entities: - tools.append( - self._create_search_entities_tool( - user_id=user_id, - namespace=effective_namespace, - ) + # The other three tools teach "project/Harbor" in their refusals and + # docstrings, and the model uses it here too. Slugging it whole minted + # project/project_harbor and stranded the correction on a phantom. + entity, qualifier = self._qualified( + entity=entity, user_id=user_id, namespace=effective_namespace, declared_type=entity_type + ) + entity_type = qualifier or entity_type + + existing = self._resolve( + entity=entity, + entity_type=entity_type, + user_id=user_id, + namespace=effective_namespace, + ) + + novel_facts, duplicate_count = self._novel_facts(existing=existing, facts=facts or []) + judgments = self._judge_superseded(existing=existing, new_facts=novel_facts) + # The judge is a blocking provider call in the middle of a + # read-modify-write over the whole row. Nothing else runs on this + # thread meanwhile, but the background capture thread shares the store, + # so merge onto the row as it stands now rather than the snapshot taken + # before the call. The async twin holds a lock instead - there, an + # assistant turn's tool calls really do interleave. + existing = self._resolve( + entity=entity, + entity_type=entity_type, + user_id=user_id, + namespace=effective_namespace, + ) + novel_facts, duplicate_count = self._novel_facts(existing=existing, facts=facts or []) + + entity_obj, created, revived, stale_row_key = self._apply_remember( + existing=existing, + entity=entity, + entity_type=entity_type, + description=description, + facts=novel_facts, + events=events or [], + note=note, + aliases=aliases, + user_id=user_id, + agent_id=agent_id, + team_id=team_id, + namespace=effective_namespace, + ) + superseded_count = self._apply_supersessions(entity_obj=entity_obj, judgments=judgments, new_facts=novel_facts) + + saved = self._save_entity( + entity=entity_obj, + user_id=user_id, + agent_id=agent_id, + team_id=team_id, + namespace=effective_namespace, + ) + if not saved: + return f"Failed to record on {entity_obj.entity_type}/{entity_obj.entity_id}." + + if stale_row_key: + # The new row is saved; drop the old-typed row so the entity is not doubled. + try: + self.db.delete_learning(id=stale_row_key) + except Exception as e: + log_warning(f"EntityMemoryStore.remember_about: failed to delete stale row {stale_row_key}: {e}") + self._repair_far_edge_types( + entity_obj=entity_obj, + user_id=user_id, + agent_id=agent_id, + team_id=team_id, + namespace=effective_namespace, ) - if self.config.agent_can_create_entity: - tools.append( - self._create_create_entity_tool( - user_id=user_id, - agent_id=agent_id, - team_id=team_id, - namespace=effective_namespace, - ) - ) + self.entity_updated = True + return self._remember_message( + entity_obj, + created=created, + revived=revived, + facts=novel_facts, + events=events or [], + note=note, + superseded_count=superseded_count, + duplicate_count=duplicate_count, + siblings=self._same_name_siblings( + entity_obj, + self._get_rows_by_entity_id( + entity_id=entity_obj.entity_id, user_id=user_id, namespace=effective_namespace + ), + ), + ) - if self.config.agent_can_update_entity: - tools.append( - self._create_update_entity_tool( - user_id=user_id, - agent_id=agent_id, - team_id=team_id, - namespace=effective_namespace, - ) - ) - - if self.config.enable_add_fact: - tools.append( - self._create_add_fact_tool( - user_id=user_id, - agent_id=agent_id, - team_id=team_id, - namespace=effective_namespace, - ) - ) - - if self.config.enable_update_fact: - tools.append( - self._create_update_fact_tool( - user_id=user_id, - agent_id=agent_id, - team_id=team_id, - namespace=effective_namespace, - ) - ) - - if self.config.enable_delete_fact: - tools.append( - self._create_delete_fact_tool( - user_id=user_id, - agent_id=agent_id, - team_id=team_id, - namespace=effective_namespace, - ) - ) - - if self.config.enable_add_event: - tools.append( - self._create_add_event_tool( - user_id=user_id, - agent_id=agent_id, - team_id=team_id, - namespace=effective_namespace, - ) - ) - - if self.config.enable_add_relationship: - tools.append( - self._create_add_relationship_tool( - user_id=user_id, - agent_id=agent_id, - team_id=team_id, - namespace=effective_namespace, - ) - ) - - return tools - - async def aget_agent_tools( + async def aremember_about( self, + entity: str, + entity_type: str, + description: Optional[str] = None, + facts: Optional[List[str]] = None, + events: Optional[List[str]] = None, + note: Optional[str] = None, + aliases: Optional[List[str]] = None, user_id: Optional[str] = None, agent_id: Optional[str] = None, team_id: Optional[str] = None, namespace: Optional[str] = None, - ) -> List[Callable]: - """Async version of get_agent_tools.""" - tools = [] + ) -> str: + """Async version of remember_about.""" + if not self.db: + return "Entity memory has no database configured; nothing was recorded." + if not _slugify_or_none(entity): + return "Entity name is required; nothing was recorded." + + # A blank optional argument means the model did not supply it, not that + # it wants the stored value cleared - the tool surface has no way to + # clear one, so treating "" as a value silently wipes a description or + # a note pointer that a previous turn wrote. + description = _blank_to_none(description) + note = _blank_to_none(note) + effective_namespace = namespace or self.config.namespace + if effective_namespace == "user" and not user_id: + log_warning("EntityMemoryStore.aremember_about: namespace='user' requires user_id") + return "Entity memory needs a user_id for the 'user' namespace; nothing was recorded." - if self.config.agent_can_search_entities: - tools.append( - self._create_async_search_entities_tool( - user_id=user_id, - namespace=effective_namespace, - ) + # See the sync twin: the qualified form the other tools teach. + entity, qualifier = await self._aqualified( + entity=entity, user_id=user_id, namespace=effective_namespace, declared_type=entity_type + ) + entity_type = qualifier or entity_type + + existing = await self._aresolve( + entity=entity, + entity_type=entity_type, + user_id=user_id, + namespace=effective_namespace, + ) + + novel_facts, duplicate_count = self._novel_facts(existing=existing, facts=facts or []) + # The judge's provider call stays outside the lock: it only reads. + judgments = await self._ajudge_superseded(existing=existing, new_facts=novel_facts) + + async with self._write_lock(): + # Resolve again inside the lock. Everything above suspended - the + # judge, and every db call on an async db - so a sibling tool call + # from the same assistant turn may have created or changed this row + # since. Merging onto the stale snapshot would overwrite it. + existing = await self._aresolve( + entity=entity, + entity_type=entity_type, + user_id=user_id, + namespace=effective_namespace, + ) + # Novelty has to be judged against the row we are about to merge + # onto, not the pre-lock snapshot: a sibling call carrying the same + # sentence would otherwise pass the duplicate check twice and + # append it twice. + novel_facts, duplicate_count = self._novel_facts(existing=existing, facts=facts or []) + + entity_obj, created, revived, stale_row_key = self._apply_remember( + existing=existing, + entity=entity, + entity_type=entity_type, + description=description, + facts=novel_facts, + events=events or [], + note=note, + aliases=aliases, + user_id=user_id, + agent_id=agent_id, + team_id=team_id, + namespace=effective_namespace, + ) + superseded_count = self._apply_supersessions( + entity_obj=entity_obj, judgments=judgments, new_facts=novel_facts ) - if self.config.agent_can_create_entity: - tools.append( - self._create_async_create_entity_tool( - user_id=user_id, - agent_id=agent_id, - team_id=team_id, - namespace=effective_namespace, - ) + saved = await self._asave_entity( + entity=entity_obj, + user_id=user_id, + agent_id=agent_id, + team_id=team_id, + namespace=effective_namespace, + ) + if not saved: + return f"Failed to record on {entity_obj.entity_type}/{entity_obj.entity_id}." + + if stale_row_key: + # The new row is saved; drop the old-typed row so the entity is not doubled. + try: + if isinstance(self.db, AsyncBaseDb): + await self.db.delete_learning(id=stale_row_key) + else: + self.db.delete_learning(id=stale_row_key) + except Exception as e: + log_warning(f"EntityMemoryStore.aremember_about: failed to delete stale row {stale_row_key}: {e}") + await self._arepair_far_edge_types( + entity_obj=entity_obj, + user_id=user_id, + agent_id=agent_id, + team_id=team_id, + namespace=effective_namespace, ) - if self.config.agent_can_update_entity: - tools.append( - self._create_async_update_entity_tool( - user_id=user_id, - agent_id=agent_id, - team_id=team_id, - namespace=effective_namespace, - ) - ) + self.entity_updated = True + return self._remember_message( + entity_obj, + created=created, + revived=revived, + facts=novel_facts, + events=events or [], + note=note, + superseded_count=superseded_count, + duplicate_count=duplicate_count, + siblings=self._same_name_siblings( + entity_obj, + await self._aget_rows_by_entity_id( + entity_id=entity_obj.entity_id, user_id=user_id, namespace=effective_namespace + ), + ), + ) - if self.config.enable_add_fact: - tools.append( - self._create_async_add_fact_tool( - user_id=user_id, - agent_id=agent_id, - team_id=team_id, - namespace=effective_namespace, - ) - ) + def _novel_facts(self, existing: Optional[EntityMemory], facts: List[str]) -> Tuple[List[str], int]: + """Drop blank facts and exact duplicates of the entity's live facts. - if self.config.enable_update_fact: - tools.append( - self._create_async_update_fact_tool( - user_id=user_id, - agent_id=agent_id, - team_id=team_id, - namespace=effective_namespace, - ) - ) + Returns (novel_facts, duplicate_count). Recording the same sentence + twice must not double it, and an exact duplicate needs no judgment. + """ + cleaned = [f for f in facts if f and f.strip()] + if existing is None: + return cleaned, 0 - if self.config.enable_delete_fact: - tools.append( - self._create_async_delete_fact_tool( - user_id=user_id, - agent_id=agent_id, - team_id=team_id, - namespace=effective_namespace, - ) - ) - - if self.config.enable_add_event: - tools.append( - self._create_async_add_event_tool( - user_id=user_id, - agent_id=agent_id, - team_id=team_id, - namespace=effective_namespace, - ) - ) - - if self.config.enable_add_relationship: - tools.append( - self._create_async_add_relationship_tool( - user_id=user_id, - agent_id=agent_id, - team_id=team_id, - namespace=effective_namespace, - ) - ) - - return tools + live_normalized = { + _normalize_fact_text(str(f.get("content", ""))) for f in existing.live_facts() if isinstance(f, dict) + } + novel: List[str] = [] + duplicates = 0 + for f in cleaned: + normalized = _normalize_fact_text(f) + if normalized in live_normalized: + duplicates += 1 + continue + live_normalized.add(normalized) + novel.append(f) + return novel, duplicates # ========================================================================= - # Tool: search_entities + # Fact supersession (one judge call in the write path) # ========================================================================= - def _create_search_entities_tool( + def _judge_superseded(self, existing: Optional[EntityMemory], new_facts: List[str]) -> List[Tuple[str, float]]: + """One structured model call: which live facts do the new facts supersede? + + Skipped entirely when the entity is new, has no live facts, the write + carries no facts, or no model is configured. Judge failures are + conservative: nothing is superseded and the write proceeds. + """ + if existing is None or not new_facts or self.model is None: + return [] + live = [f for f in existing.live_facts() if isinstance(f, dict) and f.get("id")][-50:] + if not live: + return [] + + try: + from copy import deepcopy + + from agno.tools.function import Function + + captured: List[Tuple[str, float]] = [] + + def mark_superseded(fact_ids: List[str], confidences: List[float] = []) -> str: + """Mark which existing facts are superseded by the newly stated facts. + + Args: + fact_ids: The ids of the superseded existing facts. Empty if none. + confidences: Confidence (0.0-1.0) per id, in the same order. + + Returns: + Confirmation. + """ + for fact_id, confidence in zip(fact_ids, confidences): + captured.append((str(fact_id), float(confidence))) + return "Recorded." + + func = Function.from_callable(mark_superseded, strict=True) + func.strict = True + + model_copy = deepcopy(self.model) + model_copy.response( + messages=self._build_supersession_messages(existing=existing, live=live, new_facts=new_facts), + tools=[func], + tool_call_limit=1, + ) + return captured + except Exception as e: + log_warning(f"EntityMemoryStore: supersession judgment failed, keeping all facts: {e}") + return [] + + async def _ajudge_superseded( + self, existing: Optional[EntityMemory], new_facts: List[str] + ) -> List[Tuple[str, float]]: + """Async version of _judge_superseded.""" + if existing is None or not new_facts or self.model is None: + return [] + live = [f for f in existing.live_facts() if isinstance(f, dict) and f.get("id")][-50:] + if not live: + return [] + + try: + from copy import deepcopy + + from agno.tools.function import Function + + captured: List[Tuple[str, float]] = [] + + def mark_superseded(fact_ids: List[str], confidences: List[float] = []) -> str: + """Mark which existing facts are superseded by the newly stated facts. + + Args: + fact_ids: The ids of the superseded existing facts. Empty if none. + confidences: Confidence (0.0-1.0) per id, in the same order. + + Returns: + Confirmation. + """ + for fact_id, confidence in zip(fact_ids, confidences): + captured.append((str(fact_id), float(confidence))) + return "Recorded." + + func = Function.from_callable(mark_superseded, strict=True) + func.strict = True + + model_copy = deepcopy(self.model) + await model_copy.aresponse( + messages=self._build_supersession_messages(existing=existing, live=live, new_facts=new_facts), + tools=[func], + tool_call_limit=1, + ) + return captured + except Exception as e: + log_warning(f"EntityMemoryStore: supersession judgment failed, keeping all facts: {e}") + return [] + + def _build_supersession_messages( + self, existing: EntityMemory, live: List[Dict[str, Any]], new_facts: List[str] + ) -> List[Any]: + from agno.models.message import Message + + if self.config.system_message: + system_content = self.config.system_message + else: + system_content = _SUPERSESSION_SYSTEM_PROMPT + if self.config.instructions: + system_content += f"\n## Additional Instructions\n\n{self.config.instructions}\n" + if self.config.additional_instructions: + system_content += f"\n{self.config.additional_instructions}\n" + + label = f"{existing.entity_type}/{existing.entity_id}" + live_lines = [] + for f in live: + as_of = str(f.get("updated_at") or f.get("created_at") or "")[:10] + as_of_text = f" (as of {as_of})" if as_of else "" + live_lines.append(f"- [{f.get('id')}] {f.get('content')}{as_of_text}") + new_lines = [f"- {f}" for f in new_facts] + + user_content = ( + f"Entity: {label}\n\nLive facts:\n" + + "\n".join(live_lines) + + "\n\nNewly stated facts:\n" + + "\n".join(new_lines) + ) + return [ + Message(role="system", content=system_content), + Message(role="user", content=user_content), + ] + + def _apply_supersessions( + self, entity_obj: EntityMemory, judgments: List[Tuple[str, float]], new_facts: List[str] + ) -> int: + """Retire judged facts at or above the threshold. Returns the count retired. + + Only pre-existing facts can be retired - the judge never saw the ids of + the facts added by this write, and a hallucinated id must not touch them. + """ + if not judgments: + return 0 + + new_normalized = {_normalize_fact_text(f) for f in new_facts} + eligible_ids = { + f.get("id") + for f in entity_obj.live_facts() + if isinstance(f, dict) and _normalize_fact_text(str(f.get("content", ""))) not in new_normalized + } + + superseded_by = "superseded" + if len(new_facts) == 1: + for f in entity_obj.live_facts(): + if isinstance(f, dict) and _normalize_fact_text(str(f.get("content", ""))) == _normalize_fact_text( + new_facts[0] + ): + superseded_by = str(f.get("id")) + break + + count = 0 + threshold = self.config.supersession_threshold + for fact_id, confidence in judgments: + if confidence < threshold: + continue + if fact_id not in eligible_ids: + continue + if entity_obj.retire_fact(fact_id, superseded_by=superseded_by): + count += 1 + return count + + def _repair_far_edge_types( self, + entity_obj: EntityMemory, + user_id: Optional[str], + agent_id: Optional[str], + team_id: Optional[str], + namespace: str, + max_repairs: int = 10, + ) -> None: + """After an unknown -> real type upgrade, fix the reciprocal edges on the + far ends, which still carry entity_type="unknown" for this entity.""" + for rel in (getattr(entity_obj, "relationships", None) or [])[:max_repairs]: + if not isinstance(rel, dict): + continue + far = self.get( + entity_id=rel.get("entity_id", ""), + entity_type=rel.get("entity_type", ""), + user_id=user_id, + namespace=namespace, + ) + if far is None: + continue + changed = False + for far_rel in getattr(far, "relationships", None) or []: + if ( + isinstance(far_rel, dict) + and far_rel.get("entity_id") == entity_obj.entity_id + and far_rel.get("entity_type") != entity_obj.entity_type + ): + far_rel["entity_type"] = entity_obj.entity_type + changed = True + if changed: + self._save_entity(entity=far, user_id=user_id, agent_id=agent_id, team_id=team_id, namespace=namespace) + + async def _arepair_far_edge_types( + self, + entity_obj: EntityMemory, + user_id: Optional[str], + agent_id: Optional[str], + team_id: Optional[str], + namespace: str, + max_repairs: int = 10, + ) -> None: + """Async version of _repair_far_edge_types.""" + for rel in (getattr(entity_obj, "relationships", None) or [])[:max_repairs]: + if not isinstance(rel, dict): + continue + far = await self.aget( + entity_id=rel.get("entity_id", ""), + entity_type=rel.get("entity_type", ""), + user_id=user_id, + namespace=namespace, + ) + if far is None: + continue + changed = False + for far_rel in getattr(far, "relationships", None) or []: + if ( + isinstance(far_rel, dict) + and far_rel.get("entity_id") == entity_obj.entity_id + and far_rel.get("entity_type") != entity_obj.entity_type + ): + far_rel["entity_type"] = entity_obj.entity_type + changed = True + if changed: + await self._asave_entity( + entity=far, user_id=user_id, agent_id=agent_id, team_id=team_id, namespace=namespace + ) + + def _apply_remember( + self, + existing: Optional[EntityMemory], + entity: str, + entity_type: str, + description: Optional[str], + facts: List[str], + events: List[str], + note: Optional[str], + user_id: Optional[str], + agent_id: Optional[str], + team_id: Optional[str], + namespace: str, + aliases: Optional[List[str]] = None, + ) -> Tuple[EntityMemory, bool, bool, Optional[str]]: + """Create or merge the entity in memory. + + Returns (entity, created, revived, stale_row_key). stale_row_key is set + when the merge changed the entity's type (a minimal 'unknown' entity + acquiring its real type), so the caller must delete the old row after + saving the new one. + """ + now = _utc_now_iso() + created = False + revived = False + stale_row_key: Optional[str] = None + normalized_type = _normalize_entity_type(entity_type) or _UNKNOWN_ENTITY_TYPE + + if existing is None: + created = True + entity_obj = self.schema( + entity_id=_slugify(entity), + entity_type=normalized_type, + name=entity.strip(), + description=description, + properties={}, + facts=[], + events=[], + relationships=[], + namespace=namespace, + user_id=user_id if namespace == "user" else None, + agent_id=agent_id, + team_id=team_id, + created_at=now, + updated_at=now, + ) + else: + entity_obj = existing + if description is not None: + entity_obj.description = description + + # A minimal entity created by link_entities acquires its real type; + # the row key embeds the type, so the old row must be replaced. + if entity_obj.entity_type == _UNKNOWN_ENTITY_TYPE and normalized_type != _UNKNOWN_ENTITY_TYPE: + stale_row_key = self._build_entity_db_id(entity_obj.entity_id, entity_obj.entity_type, namespace) + entity_obj.entity_type = normalized_type + + # Remember the name variant this write arrived under, bounded so the + # list cannot grow without limit. A name that differs from every one + # on file is kept even when it slugs to this same row: that is how a + # punctuation-only merge ("Acme, Inc." onto "Acme Inc") stays + # visible and searchable instead of being written nowhere. + incoming_name = entity.strip() + existing_aliases = list(getattr(entity_obj, "aliases", None) or []) + known_names = [entity_obj.name or ""] + existing_aliases + if ( + incoming_name + and len(existing_aliases) < 8 + and all(_normalize_name(incoming_name) != _normalize_name(n) for n in known_names if n) + ): + entity_obj.aliases = [*existing_aliases, incoming_name] + + if getattr(entity_obj, "archived_at", None): + entity_obj.archived_at = None + revived = True + log_info( + f"EntityMemoryStore: entity {entity_obj.entity_type}/{entity_obj.entity_id} " + f"was archived and has been revived by this write." + ) + + for fact in facts: + if fact and fact.strip(): + entity_obj.add_fact(fact) + for event in events: + if event and event.strip(): + entity_obj.add_event(event) + if note is not None and note.strip(): + entity_obj.properties = {**(entity_obj.properties or {}), "note": note.strip()} + + # Explicit aliases (the write path for the resolution ladder's third + # rung): merged with dedupe against the name and existing aliases, + # bounded at 8. + for alias in aliases or []: + alias = alias.strip() + existing_aliases = list(getattr(entity_obj, "aliases", None) or []) + known = [entity_obj.name or ""] + existing_aliases + if ( + alias + and len(existing_aliases) < 8 + and all(_normalize_name(alias) != _normalize_name(n) for n in known if n) + ): + entity_obj.aliases = [*existing_aliases, alias] + + entity_obj.updated_at = now + return entity_obj, created, revived, stale_row_key + + def _remember_message( + self, + entity_obj: EntityMemory, + created: bool, + revived: bool, + facts: List[str], + events: List[str], + note: Optional[str], + superseded_count: int = 0, + duplicate_count: int = 0, + siblings: Optional[List[str]] = None, + ) -> str: + label = f"{entity_obj.entity_type}/{entity_obj.entity_id}" + verb = "Created" if created else "Updated" + parts = [] + if facts: + parts.append(f"{len(facts)} fact(s)") + if events: + parts.append(f"{len(events)} event(s)") + if note: + parts.append(f"note pointer {note}") + recorded = f" Recorded {', '.join(parts)}." if parts else "" + superseded_text = f" Superseded {superseded_count} earlier fact(s)." if superseded_count else "" + duplicate_text = f" Skipped {duplicate_count} already-recorded fact(s)." if duplicate_count else "" + revived_text = " The entity was archived and is now revived." if revived else "" + # Types are never merged across, so a same-name sibling is a separate + # record the model may not have meant to create. Say so while it can + # still restate under the right type. + sibling_text = "" + if siblings: + sibling_text = ( + f" Note: {', '.join(siblings)} already exists under this name - if you meant that one," + " restate with its type." + ) + return f"{verb} {label}.{recorded}{superseded_text}{duplicate_text}{revived_text}{sibling_text}" + + def _same_name_siblings(self, entity_obj: EntityMemory, rows: List[Dict[str, Any]]) -> List[str]: + """Labels of rows sharing this entity's id under a different type.""" + labels = [] + for row in rows: + row_type = row.get("entity_type") + if row_type and row_type != entity_obj.entity_type: + labels.append(f"{row_type}/{entity_obj.entity_id}") + return labels + + # ========================================================================= + # Public write API: link_entities + # ========================================================================= + + def link_entities( + self, + entity: str, + relation: str, + related_entity: str, + user_id: Optional[str] = None, + agent_id: Optional[str] = None, + team_id: Optional[str] = None, + namespace: Optional[str] = None, + ) -> str: + """Record a relationship between two entities, resolving both ends by name. + + An end that does not resolve is created as a minimal entity with + entity_type="unknown"; a later remember_about with a real type merges it. + The edge is written on both rows, each carrying the far end's resolved id, + type, relation and direction. + """ + if not self.db: + return "Entity memory has no database configured; nothing was recorded." + + effective_namespace = namespace or self.config.namespace + if effective_namespace == "user" and not user_id: + log_warning("EntityMemoryStore.link_entities: namespace='user' requires user_id") + return "Entity memory needs a user_id for the 'user' namespace; nothing was recorded." + + if not _slugify_or_none(entity) or not _slugify_or_none(related_entity): + return "Both entity names are required; nothing was recorded." + + for endpoint in (entity, related_entity): + ambiguous = self._ambiguous_name(entity=endpoint, user_id=user_id, namespace=effective_namespace) + if ambiguous: + return ambiguous + entity, entity_qualifier = self._qualified(entity=entity, user_id=user_id, namespace=effective_namespace) + related_entity, related_qualifier = self._qualified( + entity=related_entity, user_id=user_id, namespace=effective_namespace + ) + + source = self._resolve_or_create_minimal( + entity, + user_id=user_id, + agent_id=agent_id, + team_id=team_id, + namespace=effective_namespace, + entity_type=entity_qualifier, + ) + target = self._resolve_or_create_minimal( + related_entity, + user_id=user_id, + agent_id=agent_id, + team_id=team_id, + namespace=effective_namespace, + entity_type=related_qualifier, + ) + if source.entity_id == target.entity_id and source.entity_type == target.entity_type: + return f"Cannot link {source.entity_type}/{source.entity_id} to itself; nothing was recorded." + + self._write_edge(source=source, target=target, relation=relation) + + for entity_obj in (source, target): + if not self._save_entity( + entity=entity_obj, + user_id=user_id, + agent_id=agent_id, + team_id=team_id, + namespace=effective_namespace, + ): + return "Failed to record the link." + + self.entity_updated = True + return self._link_message(source=source, relation=relation, target=target) + + async def alink_entities( + self, + entity: str, + relation: str, + related_entity: str, + user_id: Optional[str] = None, + agent_id: Optional[str] = None, + team_id: Optional[str] = None, + namespace: Optional[str] = None, + ) -> str: + """Async version of link_entities.""" + if not self.db: + return "Entity memory has no database configured; nothing was recorded." + + effective_namespace = namespace or self.config.namespace + if effective_namespace == "user" and not user_id: + log_warning("EntityMemoryStore.alink_entities: namespace='user' requires user_id") + return "Entity memory needs a user_id for the 'user' namespace; nothing was recorded." + + if not _slugify_or_none(entity) or not _slugify_or_none(related_entity): + return "Both entity names are required; nothing was recorded." + + for endpoint in (entity, related_entity): + ambiguous = await self._aambiguous_name(entity=endpoint, user_id=user_id, namespace=effective_namespace) + if ambiguous: + return ambiguous + entity, entity_qualifier = await self._aqualified(entity=entity, user_id=user_id, namespace=effective_namespace) + related_entity, related_qualifier = await self._aqualified( + entity=related_entity, user_id=user_id, namespace=effective_namespace + ) + + async with self._write_lock(): + # Both ends are read-modify-written; see aremember_about. + source = await self._aresolve_or_create_minimal( + entity, + user_id=user_id, + agent_id=agent_id, + team_id=team_id, + namespace=effective_namespace, + entity_type=entity_qualifier, + ) + target = await self._aresolve_or_create_minimal( + related_entity, + user_id=user_id, + agent_id=agent_id, + team_id=team_id, + namespace=effective_namespace, + entity_type=related_qualifier, + ) + if source.entity_id == target.entity_id and source.entity_type == target.entity_type: + return f"Cannot link {source.entity_type}/{source.entity_id} to itself; nothing was recorded." + + self._write_edge(source=source, target=target, relation=relation) + + for entity_obj in (source, target): + if not await self._asave_entity( + entity=entity_obj, + user_id=user_id, + agent_id=agent_id, + team_id=team_id, + namespace=effective_namespace, + ): + return "Failed to record the link." + + self.entity_updated = True + return self._link_message(source=source, relation=relation, target=target) + + def _write_edge(self, source: EntityMemory, target: EntityMemory, relation: str) -> None: + """Write the edge on both rows, each carrying the far end's id and type.""" + now = _utc_now_iso() + source.add_relationship( + related_entity_id=target.entity_id, + relation=relation, + direction="outgoing", + entity_type=target.entity_type, + ) + source.updated_at = now + target.add_relationship( + related_entity_id=source.entity_id, + relation=relation, + direction="incoming", + entity_type=source.entity_type, + ) + target.updated_at = now + + def _link_message(self, source: EntityMemory, relation: str, target: EntityMemory) -> str: + return ( + f"Linked {source.entity_type}/{source.entity_id} --[{relation}]--> {target.entity_type}/{target.entity_id}." + ) + + # ========================================================================= + # Public read API: search_entities (agent-facing, formatted) + # ========================================================================= + + def search_entities( + self, + query: Optional[str] = None, + entity_type: Optional[str] = None, user_id: Optional[str] = None, namespace: Optional[str] = None, - ) -> Callable: - """Create the search_entities tool.""" + limit: int = 10, + ) -> str: + """Search stored entities (or list them by recency) and format the results.""" + effective_namespace = namespace or self.config.namespace + if effective_namespace == "user" and not user_id: + log_warning("EntityMemoryStore.search_entities: namespace='user' requires user_id") + return "Entity memory needs a user_id for the 'user' namespace; nothing was searched." - def search_entities( - query: str, - entity_type: Optional[str] = None, - limit: int = 5, - ) -> str: - """Search for entities in the knowledge base. + query = _blank_to_none(query) + entity_type = _blank_to_none(entity_type) - Use this to recall information about people, companies, projects, products, - or other entities that have been stored. Searches across names, facts, - events, and relationships. - - **Good times to search:** - - Before discussing a company/person that might have stored context - - When the user references an entity by name - - To recall details about a project or product - - To find relationships between entities - - **Search tips:** - - Search by name: "Acme Corp", "Jane Smith" - - Search by attribute: "PostgreSQL", "San Francisco" - - Search by relationship: "CEO", "competitor" - - Combine with entity_type to narrow results - - Args: - query: What to search for. Can be a name, fact content, relationship, - or any text that might appear in entity records. - Examples: "Acme", "uses PostgreSQL", "VP Engineering" - entity_type: Optional filter - "person", "company", "project", "product", etc. - limit: Maximum results (default: 5) - - Returns: - Formatted list of matching entities with their facts, events, and relationships. - """ + if query: results = self.search( query=query, entity_type=entity_type, user_id=user_id, - namespace=namespace, + namespace=effective_namespace, limit=limit, + include_archived=True, + ) + else: + results = self.list_entities( + entity_type=entity_type, + user_id=user_id, + namespace=effective_namespace, + limit=limit, + include_archived=True, ) - if not results: - return "No matching entities found." + return self._format_search_results( + entities=results, + query=query, + entity_type=entity_type, + namespace=effective_namespace, + limit=limit, + ) - formatted = self._format_entities_list(entities=results) - return f"Found {len(results)} entity/entities:\n\n{formatted}" - - return search_entities - - def _create_async_search_entities_tool( + async def asearch_entities( self, + query: Optional[str] = None, + entity_type: Optional[str] = None, user_id: Optional[str] = None, namespace: Optional[str] = None, - ) -> Callable: - """Create the async search_entities tool.""" + limit: int = 10, + ) -> str: + """Async version of search_entities.""" + effective_namespace = namespace or self.config.namespace + if effective_namespace == "user" and not user_id: + log_warning("EntityMemoryStore.asearch_entities: namespace='user' requires user_id") + return "Entity memory needs a user_id for the 'user' namespace; nothing was searched." - async def search_entities( - query: str, - entity_type: Optional[str] = None, - limit: int = 5, - ) -> str: - """Search for entities in the knowledge base. + query = _blank_to_none(query) + entity_type = _blank_to_none(entity_type) - Use this to recall information about people, companies, projects, products, - or other entities that have been stored. Searches across names, facts, - events, and relationships. - - **Good times to search:** - - Before discussing a company/person that might have stored context - - When the user references an entity by name - - To recall details about a project or product - - To find relationships between entities - - **Search tips:** - - Search by name: "Acme Corp", "Jane Smith" - - Search by attribute: "PostgreSQL", "San Francisco" - - Search by relationship: "CEO", "competitor" - - Combine with entity_type to narrow results - - Args: - query: What to search for. Can be a name, fact content, relationship, - or any text that might appear in entity records. - Examples: "Acme", "uses PostgreSQL", "VP Engineering" - entity_type: Optional filter - "person", "company", "project", "product", etc. - limit: Maximum results (default: 5) - - Returns: - Formatted list of matching entities with their facts, events, and relationships. - """ + if query: results = await self.asearch( query=query, entity_type=entity_type, user_id=user_id, + namespace=effective_namespace, + limit=limit, + include_archived=True, + ) + else: + results = await self.alist_entities( + entity_type=entity_type, + user_id=user_id, + namespace=effective_namespace, + limit=limit, + include_archived=True, + ) + + return self._format_search_results( + entities=results, + query=query, + entity_type=entity_type, + namespace=effective_namespace, + limit=limit, + ) + + def _format_search_results( + self, + entities: List[EntityMemory], + query: Optional[str], + entity_type: Optional[str], + namespace: str, + limit: int, + ) -> str: + scope = f"namespace '{namespace}'" + if entity_type: + scope += f", type '{entity_type}'" + + if not entities: + if query: + return f"No entities matching {query!r} (searched {scope})." + return f"No entities stored yet (searched {scope})." + + parts = [] + for i, entity in enumerate(entities, 1): + parts.append(f"{i}. {self._format_entity_hit(entity)}") + + header = ( + f"Found {len(entities)} entity/entities matching {query!r} in {scope}" + if query + else f"{len(entities)} most recently updated entity/entities in {scope}" + ) + footer = "" + if len(entities) >= limit: + footer = f"\n\nShowing the first {limit}; narrow the query or entity_type to see others." + return f"{header}:\n\n" + "\n\n".join(parts) + footer + + def _format_entity_hit(self, entity: EntityMemory, max_facts: int = 6, max_events: int = 3) -> str: + """Format one search hit: bounded, with truncation markers and the note path.""" + name = getattr(entity, "name", None) or entity.entity_id + archived = " (archived)" if getattr(entity, "archived_at", None) else "" + lines = [f"**{name}** ({entity.entity_type}){archived}"] + + description = getattr(entity, "description", None) + if description: + lines.append(description) + + properties = getattr(entity, "properties", {}) or {} + note = properties.get("note") + if note: + lines.append(f"note: {note}") + other_props = {k: v for k, v in properties.items() if k != "note"} + if other_props: + lines.append("Properties: " + ", ".join(f"{k}: {v}" for k, v in other_props.items())) + + live = entity.live_facts() if hasattr(entity, "live_facts") else getattr(entity, "facts", []) + if live: + shown = live[-max_facts:] if max_facts > 0 else [] + marker = f" (newest {len(shown)} of {len(live)} facts)" if len(live) > len(shown) else "" + lines.append("Facts:" + marker) + for f in shown: + lines.append(f" - {f.get('content', f) if isinstance(f, dict) else f}") + + entity_events = getattr(entity, "events", []) or [] + if entity_events: + shown_events = entity_events[-max_events:] + marker = ( + f" (last {len(shown_events)} of {len(entity_events)} events)" + if len(entity_events) > len(shown_events) + else "" + ) + lines.append("Events:" + marker) + for e in shown_events: + if isinstance(e, dict): + date = f" ({e.get('date')})" if e.get("date") else "" + lines.append(f" - {e.get('content', e)}{date}") + else: + lines.append(f" - {e}") + + relationships = getattr(entity, "relationships", []) or [] + if relationships: + lines.append("Relationships:") + for r in relationships: + if isinstance(r, dict): + arrow = "->" if r.get("direction", "outgoing") == "outgoing" else "<-" + lines.append(f" - {r.get('relation')} {arrow} {r.get('entity_id')}") + + return "\n".join(lines) + + # ========================================================================= + # Public write API: forget + # ========================================================================= + + def forget( + self, + entity: str, + fact: Optional[str] = None, + user_id: Optional[str] = None, + agent_id: Optional[str] = None, + team_id: Optional[str] = None, + namespace: Optional[str] = None, + ) -> str: + """Retire a fact from an entity, or archive the whole entity.""" + if not self.db: + return "Entity memory has no database configured; nothing was changed." + + effective_namespace = namespace or self.config.namespace + if effective_namespace == "user" and not user_id: + log_warning("EntityMemoryStore.forget: namespace='user' requires user_id") + return "Entity memory needs a user_id for the 'user' namespace; nothing was changed." + ambiguous = self._ambiguous_name(entity=entity, user_id=user_id, namespace=effective_namespace) + if ambiguous: + return ambiguous + entity, qualifier = self._qualified(entity=entity, user_id=user_id, namespace=effective_namespace) + entity_obj = self._resolve(entity=entity, entity_type=qualifier, user_id=user_id, namespace=effective_namespace) + if entity_obj is None: + return f"No entity found matching {entity!r}." + + result, should_save, detached = self._apply_forget(entity_obj=entity_obj, fact=fact) + if should_save: + saved = self._save_entity( + entity=entity_obj, + user_id=user_id, + agent_id=agent_id, + team_id=team_id, + namespace=effective_namespace, + ) + if not saved: + return f"Failed to update {entity_obj.entity_type}/{entity_obj.entity_id}." + if detached is not None: + self._detach_far_edge(entity_obj=entity_obj, edge=detached) + self.entity_updated = True + return result + + async def aforget( + self, + entity: str, + fact: Optional[str] = None, + user_id: Optional[str] = None, + agent_id: Optional[str] = None, + team_id: Optional[str] = None, + namespace: Optional[str] = None, + ) -> str: + """Async version of forget.""" + if not self.db: + return "Entity memory has no database configured; nothing was changed." + + effective_namespace = namespace or self.config.namespace + if effective_namespace == "user" and not user_id: + log_warning("EntityMemoryStore.aforget: namespace='user' requires user_id") + return "Entity memory needs a user_id for the 'user' namespace; nothing was changed." + async with self._write_lock(): + # Read-modify-write; see aremember_about. + ambiguous = await self._aambiguous_name(entity=entity, user_id=user_id, namespace=effective_namespace) + if ambiguous: + return ambiguous + entity, qualifier = await self._aqualified(entity=entity, user_id=user_id, namespace=effective_namespace) + entity_obj = await self._aresolve( + entity=entity, entity_type=qualifier, user_id=user_id, namespace=effective_namespace + ) + if entity_obj is None: + return f"No entity found matching {entity!r}." + + result, should_save, detached = self._apply_forget(entity_obj=entity_obj, fact=fact) + if should_save: + saved = await self._asave_entity( + entity=entity_obj, + user_id=user_id, + agent_id=agent_id, + team_id=team_id, + namespace=effective_namespace, + ) + if not saved: + return f"Failed to update {entity_obj.entity_type}/{entity_obj.entity_id}." + if detached is not None: + await self._adetach_far_edge(entity_obj=entity_obj, edge=detached) + self.entity_updated = True + return result + + def _retire(self, entity_obj: EntityMemory, fact: Dict[str, Any], superseded_by: str = "forgotten") -> None: + """Retire a fact dict in place, tolerating records without an id.""" + fact_id = fact.get("id") + if fact_id: + entity_obj.retire_fact(fact_id, superseded_by=superseded_by) + else: + fact["superseded_at"] = _utc_now_iso() + fact["superseded_by"] = superseded_by + + def _apply_forget( + self, entity_obj: EntityMemory, fact: Optional[str] + ) -> Tuple[str, bool, Optional[Dict[str, Any]]]: + """Apply forget in memory. + + Returns (message, should_save, detached_edge). The far end of a removed + edge is written by the caller, which knows whether its db is async. + """ + label = f"{entity_obj.entity_type}/{entity_obj.entity_id}" + + # No fact: archive the entity. + if fact is None or not fact.strip(): + if getattr(entity_obj, "archived_at", None): + return f"{label} is already archived.", False, None + entity_obj.archived_at = _utc_now_iso() + entity_obj.updated_at = _utc_now_iso() + return ( + f"Archived {label}. It will no longer be recalled; search_entities can still " + f"find it, and any new remember_about about it revives it.", + True, + None, + ) + + # Fact given: match against live fact content. + needle = _normalize_fact_text(fact) + live = entity_obj.live_facts() + + exact = [f for f in live if isinstance(f, dict) and _normalize_fact_text(str(f.get("content", ""))) == needle] + if exact: + for f in exact: + self._retire(entity_obj, f) + entity_obj.updated_at = _utc_now_iso() + return f"Retired fact on {label}: {exact[0].get('content')}", True, None + + contains = [ + f + for f in live + if isinstance(f, dict) + and ( + needle in _normalize_fact_text(str(f.get("content", ""))) + or _normalize_fact_text(str(f.get("content", ""))) in needle + ) + ] + if len(contains) == 1: + self._retire(entity_obj, contains[0]) + entity_obj.updated_at = _utc_now_iso() + return f"Retired fact on {label}: {contains[0].get('content')}", True, None + if len(contains) > 1: + listing = "\n".join(f" - {f.get('content')}" for f in contains) + return ( + f"Multiple facts on {label} match {fact!r}; nothing was retired. " + f"Call forget again with the exact wording of one of:\n{listing}", + False, + None, + ) + + # No fact matched: try the relationships. A correction that changes a + # relation ("written_in Rust" -> Go) has no other retirement path, and + # a stale edge renders forever, undated, next to the corrected one. + edge_message, edge_removed, detached = self._forget_edge(entity_obj=entity_obj, needle=needle, label=label) + if edge_message is not None: + return edge_message, edge_removed, detached + + # Then the events. The instructions route positions and opinions here, + # so a retraction ("that rumour was wrong") lands on the one store + # nothing could correct: the judge only ever sees live facts, and a + # contradicting event just appended beside the retracted one. + event_message, event_removed = self._forget_event(entity_obj=entity_obj, needle=needle, label=label) + if event_message is not None: + return event_message, event_removed, None + + if not live: + return f"No matching fact on {label}. It has no live facts.", False, None + bounded = live[:10] + listing = "\n".join(f" - {f.get('content') if isinstance(f, dict) else f}" for f in bounded) + more = f"\n ... and {len(live) - len(bounded)} more" if len(live) > len(bounded) else "" + return f"No matching fact on {label}. Its live facts are:\n{listing}{more}", False, None + + def _name_rows(self, entity: str, user_id: Optional[str], namespace: str) -> List[Dict[str, Any]]: + """Every stored row this bare name could mean. + + Rows keyed by the slug, plus the exact-name and alias candidates + ``_resolve`` would reach - a row whose id was set by an external writer + ("harbor_co_001") is invisible to a slug-only lookup, so it used to + evade the ambiguity guard and be unreachable through the qualified form. + """ + name = _normalize_name(entity) + rows = list(self._get_rows_by_entity_id(entity_id=_slugify(entity), user_id=user_id, namespace=namespace)) + seen = {row.get("learning_id") for row in rows} + for row in self._name_candidates(entity=entity, user_id=user_id, namespace=namespace): + if row.get("learning_id") in seen: + continue + parsed = self.schema.from_dict(row.get("content")) + if parsed is None: + continue + names = [getattr(parsed, "name", None) or ""] + list(getattr(parsed, "aliases", None) or []) + if any(candidate and _normalize_name(str(candidate)) == name for candidate in names): + rows.append(row) + seen.add(row.get("learning_id")) + return rows + + async def _aname_rows(self, entity: str, user_id: Optional[str], namespace: str) -> List[Dict[str, Any]]: + """Async version of _name_rows.""" + name = _normalize_name(entity) + rows = list( + await self._aget_rows_by_entity_id(entity_id=_slugify(entity), user_id=user_id, namespace=namespace) + ) + seen = {row.get("learning_id") for row in rows} + for row in await self._aname_candidates(entity=entity, user_id=user_id, namespace=namespace): + if row.get("learning_id") in seen: + continue + parsed = self.schema.from_dict(row.get("content")) + if parsed is None: + continue + names = [getattr(parsed, "name", None) or ""] + list(getattr(parsed, "aliases", None) or []) + if any(candidate and _normalize_name(str(candidate)) == name for candidate in names): + rows.append(row) + seen.add(row.get("learning_id")) + return rows + + @staticmethod + def _ambiguity_message(entity: str, rows: List[Dict[str, Any]]) -> Optional[str]: + """A refusal listing every type this name resolves to, or None.""" + types = sorted({str(row.get("entity_type")) for row in rows if row.get("entity_type")}) + if len(types) < 2: + return None + listing = "\n".join(f" - {t}/{entity}" for t in types) + return f"{entity!r} matches more than one entity; nothing was changed. Call again naming one of:\n{listing}" + + def _ambiguous_name(self, entity: str, user_id: Optional[str], namespace: str) -> Optional[str]: + """link_entities and forget carry no entity_type, and types are never + merged across, so a bare "Harbor" would silently pick whichever row + sorted first and leave the other permanently unaddressable.""" + name, qualifier = self._qualified(entity=entity, user_id=user_id, namespace=namespace) + if qualifier: + return None + return self._ambiguity_message(name, self._name_rows(entity=name, user_id=user_id, namespace=namespace)) + + async def _aambiguous_name(self, entity: str, user_id: Optional[str], namespace: str) -> Optional[str]: + """Async version of _ambiguous_name.""" + name, qualifier = await self._aqualified(entity=entity, user_id=user_id, namespace=namespace) + if qualifier: + return None + return self._ambiguity_message(name, await self._aname_rows(entity=name, user_id=user_id, namespace=namespace)) + + def _qualified( + self, entity: str, user_id: Optional[str], namespace: str, declared_type: Optional[str] = None + ) -> Tuple[str, Optional[str]]: + """Read back the "type/name" form the ambiguity reply asks for. + + A prefix only counts when it names one of the types actually stored + under the remainder, so an entity called "AC/DC" stays one name. + + remember_about is the one tool that declares its own entity_type, and + it passes it here: the first write under a qualified name has no row to + recognise the prefix from, so the whole string was slugged into a + phantom (project/Harbor -> project/project_harbor) that no later call + could reach. + """ + remainder = entity.partition("/")[2] + if not remainder: + return entity, None + rows = self._name_rows(entity=remainder, user_id=user_id, namespace=namespace) + known = [str(r.get("entity_type")) for r in rows if r.get("entity_type")] + if declared_type: + known.append(declared_type) + return _split_qualified_name(entity, known) + + async def _aqualified( + self, entity: str, user_id: Optional[str], namespace: str, declared_type: Optional[str] = None + ) -> Tuple[str, Optional[str]]: + """Async version of _qualified.""" + remainder = entity.partition("/")[2] + if not remainder: + return entity, None + rows = await self._aname_rows(entity=remainder, user_id=user_id, namespace=namespace) + known = [str(r.get("entity_type")) for r in rows if r.get("entity_type")] + if declared_type: + known.append(declared_type) + return _split_qualified_name(entity, known) + + def _forget_event(self, entity_obj: EntityMemory, needle: str, label: str) -> Tuple[Optional[str], bool]: + """Retire an event the caller named, matched like a fact. + + Returns (None, False) when nothing matches, so the caller falls through + to its no-match reply. + """ + events = [e for e in (getattr(entity_obj, "events", None) or []) if isinstance(e, dict)] + exact = [e for e in events if _normalize_fact_text(str(e.get("content", ""))) == needle] + matches = exact or [ + e + for e in events + if needle in _normalize_fact_text(str(e.get("content", ""))) + or _normalize_fact_text(str(e.get("content", ""))) in needle + ] + if not matches: + return None, False + # Several DIFFERENT events matching loosely is ambiguous; several exact + # copies of the same one is not - retire the set, or the reply says + # "Retired" while the event keeps rendering. + if len(matches) > 1 and not exact: + listing = "\n".join(f" - {e.get('content')}" for e in matches) + return ( + f"Multiple events on {label} match {needle!r}; nothing was retired. " + f"Call forget again with the exact wording of one of:\n{listing}", + False, + ) + + retired = matches[0] + entity_obj.events = [e for e in events if e not in matches] + entity_obj.updated_at = _utc_now_iso() + return f"Retired event on {label}: {retired.get('content')}", True + + def _forget_edge( + self, entity_obj: EntityMemory, needle: str, label: str + ) -> Tuple[Optional[str], bool, Optional[Dict[str, Any]]]: + """Retire a relationship whose text the caller named. + + Matched the way the model sees an edge rendered - "written_in -> Rust", + or either half alone. Returns (None, False) when nothing matches, so + the caller can fall through to its no-match reply. + """ + edges = [r for r in (getattr(entity_obj, "relationships", None) or []) if isinstance(r, dict)] + matches = [] + qualified_matches = [] + directed_matches = [] + for edge in edges: + relation = _normalize_fact_text(str(edge.get("relation", ""))) + # The context block renders the far end's DISPLAY name ("Sarah + # Chen") while the edge stores its slug, so match both - the + # docstring tells the model to name the edge as it is rendered, + # and every multi-word far end failed that instruction. + slug = _normalize_fact_text(str(edge.get("entity_id", ""))) + spaced = _normalize_fact_text(str(edge.get("entity_id", "")).replace("_", " ")) + plain = {t for t in (slug, spaced) if t} + # The far end can also be named in the qualified "project/Harbor" + # form the other tools teach, which is the only way to tell two + # same-slug far ends apart. + far_type = _normalize_fact_text(str(edge.get("entity_type", ""))) + qualified = {f"{far_type}/{t}" for t in plain} if far_type else set() + targets = plain | qualified + arrow = "<-" if edge.get("direction") == "incoming" else "->" + rendered = {_normalize_fact_text(f"{edge.get('relation', '')} {target}") for target in targets} + rendered |= {_normalize_fact_text(f"{edge.get('relation', '')} {arrow} {target}") for target in targets} + names_far_end = any(_mentions(needle, t) for t in targets) or _slug_phrase_in(needle, slug) + if needle in ({relation} | targets | rendered) or ( + relation and _mentions(needle, relation) and names_far_end + ): + matches.append(edge) + if any(_mentions(needle, t) for t in qualified): + qualified_matches.append(edge) + if arrow in needle: + directed_matches.append(edge) + if not matches: + return None, False, None + # A needle naming the far end's type picks that edge out of its + # same-slug siblings; "works_on -> harbor" still matches both. + if qualified_matches: + matches = qualified_matches + # Same for the arrow: a relation asserted from both sides leaves two + # edges that differ only in direction, and the block renders them as + # "rel -> far" and "rel <- far". Naming one retires that one. + if directed_matches and len(directed_matches) < len(matches): + matches = [edge for edge in matches if edge in directed_matches] + # Edges written before add_relationship became idempotent can be exact + # duplicates: the listing below would offer two identical candidates + # that no wording can tell apart, so retire the whole set instead. + # entity_type is part of the identity here for the same reason it is in + # add_relationship - project/Harbor and company/Harbor are two links, + # and retiring both while detaching one reciprocal leaves the graph + # one-sided. + identical = {(e.get("relation"), e.get("entity_id"), e.get("entity_type"), e.get("direction")) for e in matches} + if len(matches) > 1 and len(identical) > 1: + # Rendered the way the block renders them, direction included: two + # candidates printed identically ask the model to pick between the + # same string twice, and it never can. + listing = "\n".join( + f" - {e.get('relation')} {'<-' if e.get('direction') == 'incoming' else '->'} " + f"{e.get('entity_type')}/{e.get('entity_id')}" + for e in matches + ) + return ( + f"Multiple relationships on {label} match; nothing was removed. " + f"Call forget again naming one of:\n{listing}", + False, + None, + ) + + edge = matches[0] + entity_obj.relationships = [r for r in edges if r not in matches] + entity_obj.updated_at = _utc_now_iso() + arrow = "<-" if edge.get("direction") == "incoming" else "->" + return ( + f"Removed relationship on {label}: {edge.get('relation')} {arrow} " + f"{edge.get('entity_type')}/{edge.get('entity_id')}", + True, + edge, + ) + + @staticmethod + def _without_reciprocal(far: EntityMemory, entity_obj: EntityMemory, edge: Dict[str, Any]) -> Optional[List[Any]]: + """The far end's edges minus the one pointing back, or None if absent.""" + kept = [ + r + for r in (getattr(far, "relationships", None) or []) + if not ( + isinstance(r, dict) + and r.get("entity_id") == entity_obj.entity_id + and r.get("relation") == edge.get("relation") + # Type too: project/Harbor and company/Harbor share a slug, and + # retiring one edge must not strip the other's reciprocal. + and r.get("entity_type") == entity_obj.entity_type + ) + ] + if len(kept) == len(getattr(far, "relationships", None) or []): + return None + return kept + + def _detach_far_edge(self, entity_obj: EntityMemory, edge: Dict[str, Any]) -> None: + """Drop the reciprocal edge, so the graph does not go one-sided. + + Best effort: a far end that cannot be loaded or saved leaves a dangling + incoming edge, which renders as a name and misleads nobody. + """ + namespace = entity_obj.namespace or self.config.namespace + far = self.get( + entity_id=str(edge.get("entity_id", "")), + entity_type=str(edge.get("entity_type", "")), + user_id=entity_obj.user_id, + namespace=namespace, + ) + if far is None: + return + kept = self._without_reciprocal(far=far, entity_obj=entity_obj, edge=edge) + if kept is None: + return + far.relationships = kept + far.updated_at = _utc_now_iso() + self._save_entity( + entity=far, + user_id=entity_obj.user_id, + agent_id=entity_obj.agent_id, + team_id=entity_obj.team_id, + namespace=namespace, + ) + + async def _adetach_far_edge(self, entity_obj: EntityMemory, edge: Dict[str, Any]) -> None: + """Async version of _detach_far_edge. + + The sync helpers no-op against an AsyncBaseDb, which left the far end + holding an edge the near end had already dropped. + """ + namespace = entity_obj.namespace or self.config.namespace + far = await self.aget( + entity_id=str(edge.get("entity_id", "")), + entity_type=str(edge.get("entity_type", "")), + user_id=entity_obj.user_id, + namespace=namespace, + ) + if far is None: + return + kept = self._without_reciprocal(far=far, entity_obj=entity_obj, edge=edge) + if kept is None: + return + far.relationships = kept + far.updated_at = _utc_now_iso() + await self._asave_entity( + entity=far, + user_id=entity_obj.user_id, + agent_id=entity_obj.agent_id, + team_id=entity_obj.team_id, + namespace=namespace, + ) + + def _resolve( + self, + entity: str, + entity_type: Optional[str], + user_id: Optional[str], + namespace: str, + ) -> Optional[EntityMemory]: + """Resolve an entity by name within the namespace. + + Deterministic, no LLM: normalized id (exact type first, then across + types), then exact name, then aliases. Matching is deliberately narrow - + a wrong merge has no unmerge. + """ + slug = _slugify(entity) + normalized_type = _normalize_entity_type(entity_type) + + if normalized_type: + found = self.get(entity_id=slug, entity_type=normalized_type, user_id=user_id, namespace=namespace) + if found is not None: + return found + + rows = self._get_rows_by_entity_id(entity_id=slug, user_id=user_id, namespace=namespace) + for row in rows: + parsed = self.schema.from_dict(row.get("content")) + if parsed is not None and _types_can_merge(normalized_type, getattr(parsed, "entity_type", None)): + return parsed + + candidates = self._name_candidates(entity=entity, user_id=user_id, namespace=namespace) + return self._match_name_or_alias(candidates=candidates, entity=entity, entity_type=normalized_type) + + async def _aresolve( + self, + entity: str, + entity_type: Optional[str], + user_id: Optional[str], + namespace: str, + ) -> Optional[EntityMemory]: + """Async version of _resolve.""" + slug = _slugify(entity) + normalized_type = _normalize_entity_type(entity_type) + + if normalized_type: + found = await self.aget(entity_id=slug, entity_type=normalized_type, user_id=user_id, namespace=namespace) + if found is not None: + return found + + rows = await self._aget_rows_by_entity_id(entity_id=slug, user_id=user_id, namespace=namespace) + for row in rows: + parsed = self.schema.from_dict(row.get("content")) + if parsed is not None and _types_can_merge(normalized_type, getattr(parsed, "entity_type", None)): + return parsed + + candidates = await self._aname_candidates(entity=entity, user_id=user_id, namespace=namespace) + return self._match_name_or_alias(candidates=candidates, entity=entity, entity_type=normalized_type) + + def _write_lock(self) -> Any: + """The write lock for the running event loop. + + Keyed by the loop itself, weakly: a store outlives the loops that use + it (a worker that restarts its loop, a test suite that runs many), and + an id-keyed cache would both grow forever and hand a new loop the lock + of a dead one that happened to reuse its address - which hangs if that + lock died held. + """ + import asyncio + + try: + loop = asyncio.get_running_loop() + except RuntimeError: + # No loop to key on (and nothing to serialize against either): + # a fresh lock is uncontended and correct. None cannot be a + # WeakKeyDictionary key. + return asyncio.Lock() + lock = self._async_write_locks.get(loop) + if lock is None: + lock = asyncio.Lock() + self._async_write_locks[loop] = lock + return lock + + def _name_candidates(self, entity: str, user_id: Optional[str], namespace: str) -> List[Dict[str, Any]]: + """Fetch candidate rows for name/alias matching via the text query surface. + + The probes are anchored to the JSON shape of a stored name ('name": "..') + and to the quoted value (for aliases), so unrelated rows whose FACTS + mention the name cannot crowd the true row out of the window. + + Database errors RAISE: a resolution miss caused by a failing backend + would silently create a duplicate entity instead of merging. + """ + db = self._sync_db() + if db is None: + return [] + if not callable(getattr(db, "search_learnings", None)): + # Resolution, not just search: without the query surface an alias + # or an externally-keyed row stops resolving past this window and + # the write mints a duplicate. Say so where it happens. + self._log_degraded_search_once() + return self._get_recent_rows(user_id=user_id, namespace=namespace, limit=50) + try: + rows: List[Dict[str, Any]] = [] + seen_ids: set = set() + for probe in self._name_probes(entity): + for row in ( + db.search_learnings( + query=probe, + learning_type=self.learning_type, + namespace=namespace, + user_id=user_id if namespace == "user" else None, + limit=20, + ) + or [] + ): + row_id = row.get("learning_id") + if row_id not in seen_ids: + seen_ids.add(row_id) + rows.append(row) + return rows + except NotImplementedError: + self._log_degraded_search_once() + return self._get_recent_rows(user_id=user_id, namespace=namespace, limit=50) + + async def _aname_candidates(self, entity: str, user_id: Optional[str], namespace: str) -> List[Dict[str, Any]]: + """Async version of _name_candidates.""" + if not self.db: + return [] + if not callable(getattr(self.db, "search_learnings", None)): + # See the sync twin. + self._log_degraded_search_once() + return await self._aget_recent_rows(user_id=user_id, namespace=namespace, limit=50) + try: + rows: List[Dict[str, Any]] = [] + seen_ids: set = set() + for probe in self._name_probes(entity): + if isinstance(self.db, AsyncBaseDb): + probe_rows = await self.db.search_learnings( + query=probe, + learning_type=self.learning_type, + namespace=namespace, + user_id=user_id if namespace == "user" else None, + limit=20, + ) + else: + probe_rows = self.db.search_learnings( + query=probe, + learning_type=self.learning_type, + namespace=namespace, + user_id=user_id if namespace == "user" else None, + limit=20, + ) + for row in probe_rows or []: + row_id = row.get("learning_id") + if row_id not in seen_ids: + seen_ids.add(row_id) + rows.append(row) + return rows + except NotImplementedError: + self._log_degraded_search_once() + return await self._aget_recent_rows(user_id=user_id, namespace=namespace, limit=50) + + @staticmethod + def _name_probes(entity: str) -> List[str]: + """Anchored search probes for a display name: the JSON name-field shape + and the quoted value (which also matches an alias entry).""" + stripped = entity.strip() + return [f'name": "{stripped}', f'"{stripped}"'] + + def _get_recent_rows(self, user_id: Optional[str], namespace: str, limit: int) -> List[Dict[str, Any]]: + db = self._sync_db() + if db is None: + return [] + try: + rows = db.get_learnings( + learning_type=self.learning_type, namespace=namespace, + user_id=user_id if namespace == "user" else None, limit=limit, ) + return rows or [] + except Exception as e: + log_debug(f"EntityMemoryStore._get_recent_rows failed: {e}") + return [] - if not results: - return "No matching entities found." + async def _aget_recent_rows(self, user_id: Optional[str], namespace: str, limit: int) -> List[Dict[str, Any]]: + try: + if isinstance(self.db, AsyncBaseDb): + rows = await self.db.get_learnings( + learning_type=self.learning_type, + namespace=namespace, + user_id=user_id if namespace == "user" else None, + limit=limit, + ) + else: + rows = self.db.get_learnings( # type: ignore[union-attr] + learning_type=self.learning_type, + namespace=namespace, + user_id=user_id if namespace == "user" else None, + limit=limit, + ) + return rows or [] + except Exception as e: + log_debug(f"EntityMemoryStore._aget_recent_rows failed: {e}") + return [] - formatted = self._format_entities_list(entities=results) - return f"Found {len(results)} entity/entities:\n\n{formatted}" + def _match_name_or_alias( + self, candidates: List[Dict[str, Any]], entity: str, entity_type: Optional[str] = None + ) -> Optional[EntityMemory]: + """Exact (normalized) name match first, then exact alias match. - return search_entities + A candidate of a different canonical type is a name collision, not the + same entity, and is skipped - see ``_types_can_merge``. + """ + target = _normalize_name(entity) + parsed_candidates: List[EntityMemory] = [] + for row in self._order_rows(candidates): + parsed = self.schema.from_dict(row.get("content")) + if parsed is not None and _types_can_merge(entity_type, getattr(parsed, "entity_type", None)): + parsed_candidates.append(parsed) - # ========================================================================= - # Tool: create_entity - # ========================================================================= + for parsed in parsed_candidates: + name = getattr(parsed, "name", None) + if name and _normalize_name(name) == target: + return parsed + for parsed in parsed_candidates: + aliases = getattr(parsed, "aliases", None) or [] + if any(_normalize_name(str(alias)) == target for alias in aliases): + return parsed + return None - def _create_create_entity_tool( + def _resolve_or_create_minimal( self, - user_id: Optional[str] = None, - agent_id: Optional[str] = None, - team_id: Optional[str] = None, - namespace: Optional[str] = None, - ) -> Callable: - """Create the create_entity tool.""" + entity: str, + user_id: Optional[str], + agent_id: Optional[str], + team_id: Optional[str], + namespace: str, + entity_type: Optional[str] = None, + ) -> EntityMemory: + """Resolve an entity by name, creating a minimal 'unknown' entity if absent.""" + found = self._resolve(entity=entity, entity_type=entity_type, user_id=user_id, namespace=namespace) + if found is not None: + return found + return self._minimal_entity( + entity=entity, user_id=user_id, agent_id=agent_id, team_id=team_id, namespace=namespace + ) - def create_entity( - entity_id: str, - entity_type: str, - name: str, - description: Optional[str] = None, - properties: Optional[Dict[str, str]] = None, - ) -> str: - """Create a new entity in the knowledge base. - - Use this when you encounter a person, company, project, or other entity - worth remembering. Create the entity first, then add facts/events/relationships. - - **When to create an entity:** - - A company, person, or project is discussed with substantive details - - Information would be useful to recall in future conversations - - The entity has a specific identity (not just "a company") - - **When NOT to create:** - - For the user themselves (use user memory) - - For generic concepts without specific identity - - For one-off mentions with no useful details - - Args: - entity_id: Unique identifier using lowercase and underscores. - Convention: descriptive name like "acme_corp", "jane_smith", "project_atlas" - Bad: "company1", "entity_123", "c" - entity_type: Category of entity. Common types: - - "person": Individual people - - "company": Businesses, organizations - - "project": Specific initiatives or projects - - "product": Software, services, offerings - - "system": Technical systems, platforms - - "concept": Domain-specific concepts worth tracking - name: Human-readable display name (e.g., "Acme Corporation", "Jane Smith") - description: Brief description of what/who this entity is. - Good: "Enterprise SaaS startup in the fintech space, potential client" - Bad: "A company" (too vague) - properties: Optional key-value metadata (e.g., {"industry": "fintech", "stage": "Series A"}) - - Returns: - Confirmation message. - """ - success = self.create_entity( - entity_id=entity_id, - entity_type=entity_type, - name=name, - description=description, - properties=properties, - user_id=user_id, - agent_id=agent_id, - team_id=team_id, - namespace=namespace, - ) - - if success: - self.entity_updated = True - return f"Entity created: {entity_type}/{entity_id} ({name})" - return "Failed to create entity (may already exist)" - - return create_entity - - def _create_async_create_entity_tool( + async def _aresolve_or_create_minimal( self, - user_id: Optional[str] = None, - agent_id: Optional[str] = None, - team_id: Optional[str] = None, - namespace: Optional[str] = None, - ) -> Callable: - """Create the async create_entity tool.""" + entity: str, + user_id: Optional[str], + agent_id: Optional[str], + team_id: Optional[str], + namespace: str, + entity_type: Optional[str] = None, + ) -> EntityMemory: + """Async version of _resolve_or_create_minimal.""" + found = await self._aresolve(entity=entity, entity_type=entity_type, user_id=user_id, namespace=namespace) + if found is not None: + return found + return self._minimal_entity( + entity=entity, user_id=user_id, agent_id=agent_id, team_id=team_id, namespace=namespace + ) - async def create_entity( - entity_id: str, - entity_type: str, - name: str, - description: Optional[str] = None, - properties: Optional[Dict[str, str]] = None, - ) -> str: - """Create a new entity in the knowledge base. + def _minimal_entity( + self, + entity: str, + user_id: Optional[str], + agent_id: Optional[str], + team_id: Optional[str], + namespace: str, + ) -> EntityMemory: + now = _utc_now_iso() + return self.schema( + entity_id=_slugify(entity), + entity_type=_UNKNOWN_ENTITY_TYPE, + name=entity.strip(), + properties={}, + facts=[], + events=[], + relationships=[], + namespace=namespace, + user_id=user_id if namespace == "user" else None, + agent_id=agent_id, + team_id=team_id, + created_at=now, + updated_at=now, + ) - Use this when you encounter a person, company, project, or other entity - worth remembering. Create the entity first, then add facts/events/relationships. - - **When to create an entity:** - - A company, person, or project is discussed with substantive details - - Information would be useful to recall in future conversations - - The entity has a specific identity (not just "a company") - - **When NOT to create:** - - For the user themselves (use user memory) - - For generic concepts without specific identity - - For one-off mentions with no useful details - - Args: - entity_id: Unique identifier using lowercase and underscores. - Convention: descriptive name like "acme_corp", "jane_smith", "project_atlas" - Bad: "company1", "entity_123", "c" - entity_type: Category of entity. Common types: - - "person": Individual people - - "company": Businesses, organizations - - "project": Specific initiatives or projects - - "product": Software, services, offerings - - "system": Technical systems, platforms - - "concept": Domain-specific concepts worth tracking - name: Human-readable display name (e.g., "Acme Corporation", "Jane Smith") - description: Brief description of what/who this entity is. - Good: "Enterprise SaaS startup in the fintech space, potential client" - Bad: "A company" (too vague) - properties: Optional key-value metadata (e.g., {"industry": "fintech", "stage": "Series A"}) - - Returns: - Confirmation message. - """ - success = await self.acreate_entity( + def _get_rows_by_entity_id(self, entity_id: str, user_id: Optional[str], namespace: str) -> List[Dict[str, Any]]: + """Fetch learnings rows for an entity id across entity types.""" + db = self._sync_db() + if db is None: + return [] + try: + rows = db.get_learnings( + learning_type=self.learning_type, entity_id=entity_id, - entity_type=entity_type, - name=name, - description=description, - properties=properties, - user_id=user_id, - agent_id=agent_id, - team_id=team_id, namespace=namespace, + user_id=user_id if namespace == "user" else None, ) + return self._order_rows(rows or []) + except Exception as e: + log_debug(f"EntityMemoryStore._get_rows_by_entity_id failed: {e}") + return [] - if success: - self.entity_updated = True - return f"Entity created: {entity_type}/{entity_id} ({name})" - return "Failed to create entity (may already exist)" + async def _aget_rows_by_entity_id( + self, entity_id: str, user_id: Optional[str], namespace: str + ) -> List[Dict[str, Any]]: + """Async version of _get_rows_by_entity_id.""" + if not self.db: + return [] + try: + if isinstance(self.db, AsyncBaseDb): + rows = await self.db.get_learnings( + learning_type=self.learning_type, + entity_id=entity_id, + namespace=namespace, + user_id=user_id if namespace == "user" else None, + ) + else: + rows = self.db.get_learnings( + learning_type=self.learning_type, + entity_id=entity_id, + namespace=namespace, + user_id=user_id if namespace == "user" else None, + ) + return self._order_rows(rows or []) + except Exception as e: + log_debug(f"EntityMemoryStore._aget_rows_by_entity_id failed: {e}") + return [] - return create_entity + @staticmethod + def _order_rows(rows: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + """Order rows newest-first with a deterministic tie-break on the id. + + Backend timestamps have second resolution, so same-second writes tie; + resolution must not flip between such rows across calls. + """ + return sorted( + rows, + key=lambda r: (-(r.get("updated_at") or r.get("created_at") or 0), str(r.get("learning_id") or "")), + ) # ========================================================================= - # Tool: update_entity - # ========================================================================= - - def _create_update_entity_tool( - self, - user_id: Optional[str] = None, - agent_id: Optional[str] = None, - team_id: Optional[str] = None, - namespace: Optional[str] = None, - ) -> Callable: - """Create the update_entity tool.""" - - def update_entity( - entity_id: str, - entity_type: str, - name: Optional[str] = None, - description: Optional[str] = None, - properties: Optional[Dict[str, str]] = None, - ) -> str: - """Update an existing entity's core properties. - - Use this to modify the entity's identity information. Only provided - fields will be updated - omitted fields remain unchanged. - - **When to update:** - - Name change: Company rebranded, person changed name - - Description evolved: Better understanding of what entity is - - Properties changed: New metadata to add - - **Note:** To update facts, events, or relationships, use the specific - tools (update_fact, add_event, add_relationship) instead. - - Args: - entity_id: The entity's identifier - entity_type: Type of entity - name: New display name (only if changed) - description: New description (only if you have better info) - properties: Properties to add/update (merged with existing) - Existing properties not in this dict are preserved - - Returns: - Confirmation message. - """ - success = self.update_entity( - entity_id=entity_id, - entity_type=entity_type, - name=name, - description=description, - properties=properties, - user_id=user_id, - agent_id=agent_id, - team_id=team_id, - namespace=namespace, - ) - - if success: - self.entity_updated = True - return f"Entity updated: {entity_type}/{entity_id}" - return f"Entity not found: {entity_type}/{entity_id}" - - return update_entity - - def _create_async_update_entity_tool( - self, - user_id: Optional[str] = None, - agent_id: Optional[str] = None, - team_id: Optional[str] = None, - namespace: Optional[str] = None, - ) -> Callable: - """Create the async update_entity tool.""" - - async def update_entity( - entity_id: str, - entity_type: str, - name: Optional[str] = None, - description: Optional[str] = None, - properties: Optional[Dict[str, str]] = None, - ) -> str: - """Update an existing entity's core properties. - - Use this to modify the entity's identity information. Only provided - fields will be updated - omitted fields remain unchanged. - - **When to update:** - - Name change: Company rebranded, person changed name - - Description evolved: Better understanding of what entity is - - Properties changed: New metadata to add - - **Note:** To update facts, events, or relationships, use the specific - tools (update_fact, add_event, add_relationship) instead. - - Args: - entity_id: The entity's identifier - entity_type: Type of entity - name: New display name (only if changed) - description: New description (only if you have better info) - properties: Properties to add/update (merged with existing) - Existing properties not in this dict are preserved - - Returns: - Confirmation message. - """ - success = await self.aupdate_entity( - entity_id=entity_id, - entity_type=entity_type, - name=name, - description=description, - properties=properties, - user_id=user_id, - agent_id=agent_id, - team_id=team_id, - namespace=namespace, - ) - - if success: - self.entity_updated = True - return f"Entity updated: {entity_type}/{entity_id}" - return f"Entity not found: {entity_type}/{entity_id}" - - return update_entity - - # ========================================================================= - # Tool: add_fact - # ========================================================================= - - def _create_add_fact_tool( - self, - user_id: Optional[str] = None, - agent_id: Optional[str] = None, - team_id: Optional[str] = None, - namespace: Optional[str] = None, - ) -> Callable: - """Create the add_fact tool.""" - - def add_fact( - entity_id: str, - entity_type: str, - fact: str, - ) -> str: - """Add a fact to an entity. - - Facts are **timeless truths** about an entity (semantic memory). - They describe what IS, not what HAPPENED. - - **Good facts (timeless, descriptive):** - - "Uses PostgreSQL and Redis for their data layer" - - "Headquarters in San Francisco, engineering team in Austin" - - "Founded by ex-Google engineers in 2019" - - "Main product is a B2B analytics platform" - - "Prefers async communication via Slack" - - **Not facts (use events instead):** - - "Launched v2.0 last month" → This is an EVENT (time-bound) - - "Just closed Series B" → This is an EVENT - - "Had a meeting yesterday" → This is an EVENT - - **Not facts (too vague):** - - "It's a good company" → Subjective, not useful - - "They do tech stuff" → Too vague - - Args: - entity_id: The entity's identifier (e.g., "acme_corp") - entity_type: Type of entity (e.g., "company") - fact: The fact to add - should be specific and timeless - - Returns: - Confirmation message with fact ID. - """ - fact_id = self.add_fact( - entity_id=entity_id, - entity_type=entity_type, - fact=fact, - user_id=user_id, - agent_id=agent_id, - team_id=team_id, - namespace=namespace, - ) - - if fact_id: - self.entity_updated = True - return f"Fact added to {entity_type}/{entity_id} (id: {fact_id})" - return "Failed to add fact (entity may not exist)" - - return add_fact - - def _create_async_add_fact_tool( - self, - user_id: Optional[str] = None, - agent_id: Optional[str] = None, - team_id: Optional[str] = None, - namespace: Optional[str] = None, - ) -> Callable: - """Create the async add_fact tool.""" - - async def add_fact( - entity_id: str, - entity_type: str, - fact: str, - ) -> str: - """Add a fact to an entity. - - Facts are **timeless truths** about an entity (semantic memory). - They describe what IS, not what HAPPENED. - - **Good facts (timeless, descriptive):** - - "Uses PostgreSQL and Redis for their data layer" - - "Headquarters in San Francisco, engineering team in Austin" - - "Founded by ex-Google engineers in 2019" - - "Main product is a B2B analytics platform" - - "Prefers async communication via Slack" - - **Not facts (use events instead):** - - "Launched v2.0 last month" → This is an EVENT (time-bound) - - "Just closed Series B" → This is an EVENT - - "Had a meeting yesterday" → This is an EVENT - - **Not facts (too vague):** - - "It's a good company" → Subjective, not useful - - "They do tech stuff" → Too vague - - Args: - entity_id: The entity's identifier (e.g., "acme_corp") - entity_type: Type of entity (e.g., "company") - fact: The fact to add - should be specific and timeless - - Returns: - Confirmation message with fact ID. - """ - fact_id = await self.aadd_fact( - entity_id=entity_id, - entity_type=entity_type, - fact=fact, - user_id=user_id, - agent_id=agent_id, - team_id=team_id, - namespace=namespace, - ) - - if fact_id: - self.entity_updated = True - return f"Fact added to {entity_type}/{entity_id} (id: {fact_id})" - return "Failed to add fact (entity may not exist)" - - return add_fact - - # ========================================================================= - # Tool: update_fact - # ========================================================================= - - def _create_update_fact_tool( - self, - user_id: Optional[str] = None, - agent_id: Optional[str] = None, - team_id: Optional[str] = None, - namespace: Optional[str] = None, - ) -> Callable: - """Create the update_fact tool.""" - - def update_fact( - entity_id: str, - entity_type: str, - fact_id: str, - fact: str, - ) -> str: - """Update an existing fact on an entity. - - Use this when a fact needs correction or has become more specific. - The new fact completely replaces the old one. - - **When to update:** - - Correction: Original fact was wrong - - More detail: "Uses PostgreSQL" → "Uses PostgreSQL 15 with TimescaleDB extension" - - Changed reality: "50 employees" → "75 employees after recent hiring" - - **When to delete instead:** - - Fact is no longer true and shouldn't be replaced - - Fact was a misunderstanding - - Args: - entity_id: The entity's identifier - entity_type: Type of entity - fact_id: ID of the fact to update (from search_entities results) - fact: New fact content - complete replacement, not a diff - - Returns: - Confirmation message. - """ - success = self.update_fact( - entity_id=entity_id, - entity_type=entity_type, - fact_id=fact_id, - fact=fact, - user_id=user_id, - agent_id=agent_id, - team_id=team_id, - namespace=namespace, - ) - - if success: - self.entity_updated = True - return f"Fact updated on {entity_type}/{entity_id}" - return f"Fact not found: {fact_id}" - - return update_fact - - def _create_async_update_fact_tool( - self, - user_id: Optional[str] = None, - agent_id: Optional[str] = None, - team_id: Optional[str] = None, - namespace: Optional[str] = None, - ) -> Callable: - """Create the async update_fact tool.""" - - async def update_fact( - entity_id: str, - entity_type: str, - fact_id: str, - fact: str, - ) -> str: - """Update an existing fact on an entity. - - Use this when a fact needs correction or has become more specific. - The new fact completely replaces the old one. - - **When to update:** - - Correction: Original fact was wrong - - More detail: "Uses PostgreSQL" → "Uses PostgreSQL 15 with TimescaleDB extension" - - Changed reality: "50 employees" → "75 employees after recent hiring" - - **When to delete instead:** - - Fact is no longer true and shouldn't be replaced - - Fact was a misunderstanding - - Args: - entity_id: The entity's identifier - entity_type: Type of entity - fact_id: ID of the fact to update (from search_entities results) - fact: New fact content - complete replacement, not a diff - - Returns: - Confirmation message. - """ - success = await self.aupdate_fact( - entity_id=entity_id, - entity_type=entity_type, - fact_id=fact_id, - fact=fact, - user_id=user_id, - agent_id=agent_id, - team_id=team_id, - namespace=namespace, - ) - - if success: - self.entity_updated = True - return f"Fact updated on {entity_type}/{entity_id}" - return f"Fact not found: {fact_id}" - - return update_fact - - # ========================================================================= - # Tool: delete_fact - # ========================================================================= - - def _create_delete_fact_tool( - self, - user_id: Optional[str] = None, - agent_id: Optional[str] = None, - team_id: Optional[str] = None, - namespace: Optional[str] = None, - ) -> Callable: - """Create the delete_fact tool.""" - - def delete_fact( - entity_id: str, - entity_type: str, - fact_id: str, - ) -> str: - """Delete a fact from an entity. - - Use this when a fact is no longer accurate and shouldn't be replaced - with updated information. - - **When to delete:** - - Fact was incorrect/misunderstood - - Fact is no longer true (and no replacement makes sense) - - Duplicate of another fact - - Too vague to be useful - - **When to update instead:** - - Fact needs correction but the topic is still relevant - - Fact needs more detail - - Args: - entity_id: The entity's identifier - entity_type: Type of entity - fact_id: ID of the fact to delete (from search_entities results) - - Returns: - Confirmation message. - """ - success = self.delete_fact( - entity_id=entity_id, - entity_type=entity_type, - fact_id=fact_id, - user_id=user_id, - agent_id=agent_id, - team_id=team_id, - namespace=namespace, - ) - - if success: - self.entity_updated = True - return f"Fact deleted from {entity_type}/{entity_id}" - return f"Fact not found: {fact_id}" - - return delete_fact - - def _create_async_delete_fact_tool( - self, - user_id: Optional[str] = None, - agent_id: Optional[str] = None, - team_id: Optional[str] = None, - namespace: Optional[str] = None, - ) -> Callable: - """Create the async delete_fact tool.""" - - async def delete_fact( - entity_id: str, - entity_type: str, - fact_id: str, - ) -> str: - """Delete a fact from an entity. - - Use this when a fact is no longer accurate and shouldn't be replaced - with updated information. - - **When to delete:** - - Fact was incorrect/misunderstood - - Fact is no longer true (and no replacement makes sense) - - Duplicate of another fact - - Too vague to be useful - - **When to update instead:** - - Fact needs correction but the topic is still relevant - - Fact needs more detail - - Args: - entity_id: The entity's identifier - entity_type: Type of entity - fact_id: ID of the fact to delete (from search_entities results) - - Returns: - Confirmation message. - """ - success = await self.adelete_fact( - entity_id=entity_id, - entity_type=entity_type, - fact_id=fact_id, - user_id=user_id, - agent_id=agent_id, - team_id=team_id, - namespace=namespace, - ) - - if success: - self.entity_updated = True - return f"Fact deleted from {entity_type}/{entity_id}" - return f"Fact not found: {fact_id}" - - return delete_fact - - # ========================================================================= - # Tool: add_event - # ========================================================================= - - def _create_add_event_tool( - self, - user_id: Optional[str] = None, - agent_id: Optional[str] = None, - team_id: Optional[str] = None, - namespace: Optional[str] = None, - ) -> Callable: - """Create the add_event tool.""" - - def add_event( - entity_id: str, - entity_type: str, - event: str, - date: Optional[str] = None, - ) -> str: - """Add an event to an entity. - - Events are **time-bound occurrences** (episodic memory). - They describe what HAPPENED, not what IS. - - **Good events (specific, time-bound):** - - "Launched v2.0 with new ML features" (date: "2025-01-15") - - "Closed $50M Series B led by Sequoia" (date: "2024-Q3") - - "Had 4-hour outage affecting payment processing" (date: "2024-12-20") - - "CEO announced pivot to enterprise market" (date: "2024-11") - - "Initial discovery call - interested in our analytics product" - - **Not events (use facts instead):** - - "Uses PostgreSQL" → This is a FACT (timeless truth) - - "Based in San Francisco" → This is a FACT - - "Has 50 employees" → This is a FACT - - **Include dates when known** - even approximate dates help: - - Exact: "2025-01-15" - - Month: "January 2025" or "2025-01" - - Quarter: "Q1 2025" - - Relative: "early 2024", "last week" - - Args: - entity_id: The entity's identifier (e.g., "acme_corp") - entity_type: Type of entity (e.g., "company") - event: Description of what happened - be specific - date: When it happened (ISO format, natural language, or approximate) - - Returns: - Confirmation message with event ID. - """ - event_id = self.add_event( - entity_id=entity_id, - entity_type=entity_type, - event=event, - date=date, - user_id=user_id, - agent_id=agent_id, - team_id=team_id, - namespace=namespace, - ) - - if event_id: - self.entity_updated = True - return f"Event added to {entity_type}/{entity_id} (id: {event_id})" - return "Failed to add event (entity may not exist)" - - return add_event - - def _create_async_add_event_tool( - self, - user_id: Optional[str] = None, - agent_id: Optional[str] = None, - team_id: Optional[str] = None, - namespace: Optional[str] = None, - ) -> Callable: - """Create the async add_event tool.""" - - async def add_event( - entity_id: str, - entity_type: str, - event: str, - date: Optional[str] = None, - ) -> str: - """Add an event to an entity. - - Events are **time-bound occurrences** (episodic memory). - They describe what HAPPENED, not what IS. - - **Good events (specific, time-bound):** - - "Launched v2.0 with new ML features" (date: "2025-01-15") - - "Closed $50M Series B led by Sequoia" (date: "2024-Q3") - - "Had 4-hour outage affecting payment processing" (date: "2024-12-20") - - "CEO announced pivot to enterprise market" (date: "2024-11") - - "Initial discovery call - interested in our analytics product" - - **Not events (use facts instead):** - - "Uses PostgreSQL" → This is a FACT (timeless truth) - - "Based in San Francisco" → This is a FACT - - "Has 50 employees" → This is a FACT - - **Include dates when known** - even approximate dates help: - - Exact: "2025-01-15" - - Month: "January 2025" or "2025-01" - - Quarter: "Q1 2025" - - Relative: "early 2024", "last week" - - Args: - entity_id: The entity's identifier (e.g., "acme_corp") - entity_type: Type of entity (e.g., "company") - event: Description of what happened - be specific - date: When it happened (ISO format, natural language, or approximate) - - Returns: - Confirmation message with event ID. - """ - event_id = await self.aadd_event( - entity_id=entity_id, - entity_type=entity_type, - event=event, - date=date, - user_id=user_id, - agent_id=agent_id, - team_id=team_id, - namespace=namespace, - ) - - if event_id: - self.entity_updated = True - return f"Event added to {entity_type}/{entity_id} (id: {event_id})" - return "Failed to add event (entity may not exist)" - - return add_event - - # ========================================================================= - # Tool: add_relationship - # ========================================================================= - - def _create_add_relationship_tool( - self, - user_id: Optional[str] = None, - agent_id: Optional[str] = None, - team_id: Optional[str] = None, - namespace: Optional[str] = None, - ) -> Callable: - """Create the add_relationship tool.""" - - def add_relationship( - entity_id: str, - entity_type: str, - related_entity_id: str, - relation: str, - direction: str = "outgoing", - ) -> str: - """Add a relationship between two entities. - - Relationships are **graph edges** connecting entities - they capture - how entities relate to each other. - - **Common relationship patterns:** - - People → Companies: - - "jane_smith" --[CEO]--> "acme_corp" - - "bob_jones" --[engineer_at]--> "acme_corp" - - "sarah_chen" --[founder]--> "startup_xyz" - - Companies → Companies: - - "acme_corp" --[competitor_of]--> "beta_inc" - - "acme_corp" --[acquired]--> "small_startup" - - "acme_corp" --[partner_of]--> "big_vendor" - - Projects → Other entities: - - "project_atlas" --[uses]--> "postgresql" - - "project_atlas" --[owned_by]--> "acme_corp" - - "project_atlas" --[led_by]--> "jane_smith" - - **Direction matters:** - - "outgoing": This entity → Related entity (default) - "jane_smith" --[CEO]--> "acme_corp" means Jane IS CEO OF Acme - - "incoming": Related entity → This entity - "acme_corp" with incoming "CEO" from "jane_smith" means Acme HAS CEO Jane - - Args: - entity_id: The source entity's identifier - entity_type: Type of source entity - related_entity_id: The target entity's identifier (must exist or will be created) - relation: Type of relationship - use clear, consistent labels: - For roles: "CEO", "CTO", "engineer_at", "founder" - For ownership: "owns", "owned_by", "part_of" - For competition: "competitor_of", "partner_of" - For technical: "uses", "depends_on", "integrates_with" - direction: "outgoing" (source → target) or "incoming" (target → source) - - Returns: - Confirmation message with relationship ID. - """ - rel_id = self.add_relationship( - entity_id=entity_id, - entity_type=entity_type, - related_entity_id=related_entity_id, - relation=relation, - direction=direction, - user_id=user_id, - agent_id=agent_id, - team_id=team_id, - namespace=namespace, - ) - - if rel_id: - self.entity_updated = True - return f"Relationship added: {entity_id} --[{relation}]--> {related_entity_id} (id: {rel_id})" - return "Failed to add relationship (entity may not exist)" - - return add_relationship - - def _create_async_add_relationship_tool( - self, - user_id: Optional[str] = None, - agent_id: Optional[str] = None, - team_id: Optional[str] = None, - namespace: Optional[str] = None, - ) -> Callable: - """Create the async add_relationship tool.""" - - async def add_relationship( - entity_id: str, - entity_type: str, - related_entity_id: str, - relation: str, - direction: str = "outgoing", - ) -> str: - """Add a relationship between two entities. - - Relationships are **graph edges** connecting entities - they capture - how entities relate to each other. - - **Common relationship patterns:** - - People → Companies: - - "jane_smith" --[CEO]--> "acme_corp" - - "bob_jones" --[engineer_at]--> "acme_corp" - - "sarah_chen" --[founder]--> "startup_xyz" - - Companies → Companies: - - "acme_corp" --[competitor_of]--> "beta_inc" - - "acme_corp" --[acquired]--> "small_startup" - - "acme_corp" --[partner_of]--> "big_vendor" - - Projects → Other entities: - - "project_atlas" --[uses]--> "postgresql" - - "project_atlas" --[owned_by]--> "acme_corp" - - "project_atlas" --[led_by]--> "jane_smith" - - **Direction matters:** - - "outgoing": This entity → Related entity (default) - "jane_smith" --[CEO]--> "acme_corp" means Jane IS CEO OF Acme - - "incoming": Related entity → This entity - "acme_corp" with incoming "CEO" from "jane_smith" means Acme HAS CEO Jane - - Args: - entity_id: The source entity's identifier - entity_type: Type of source entity - related_entity_id: The target entity's identifier (must exist or will be created) - relation: Type of relationship - use clear, consistent labels: - For roles: "CEO", "CTO", "engineer_at", "founder" - For ownership: "owns", "owned_by", "part_of" - For competition: "competitor_of", "partner_of" - For technical: "uses", "depends_on", "integrates_with" - direction: "outgoing" (source → target) or "incoming" (target → source) - - Returns: - Confirmation message with relationship ID. - """ - rel_id = await self.aadd_relationship( - entity_id=entity_id, - entity_type=entity_type, - related_entity_id=related_entity_id, - relation=relation, - direction=direction, - user_id=user_id, - agent_id=agent_id, - team_id=team_id, - namespace=namespace, - ) - - if rel_id: - self.entity_updated = True - return f"Relationship added: {entity_id} --[{relation}]--> {related_entity_id} (id: {rel_id})" - return "Failed to add relationship (entity may not exist)" - - return add_relationship - - # ========================================================================= - # Read Operations + # Data API: get / list / search / delete # ========================================================================= def get( @@ -1637,6 +2865,8 @@ class EntityMemoryStore(LearningStore): ) -> Optional[EntityMemory]: """Retrieve entity by entity_id and entity_type. + This is the keyed data API; it returns archived entities too. + Args: entity_id: The unique entity identifier. entity_type: The type of entity. @@ -1646,13 +2876,14 @@ class EntityMemoryStore(LearningStore): Returns: EntityMemory instance, or None if not found. """ - if not self.db: + db = self._sync_db() + if db is None: return None effective_namespace = namespace or self.config.namespace try: - result = self.db.get_learning( + result = db.get_learning( learning_type=self.learning_type, entity_id=entity_id, entity_type=entity_type, @@ -1709,9 +2940,166 @@ class EntityMemoryStore(LearningStore): log_debug(f"EntityMemoryStore.aget failed for {entity_type}/{entity_id}: {e}") return None - # ========================================================================= - # Search Operations - # ========================================================================= + def list_entities( + self, + entity_type: Optional[str] = None, + user_id: Optional[str] = None, + namespace: Optional[str] = None, + limit: int = 10, + include_archived: bool = False, + ) -> List[EntityMemory]: + """List entities by recency (most recently updated first).""" + db = self._sync_db() + if db is None: + return [] + + effective_namespace = namespace or self.config.namespace + entity_type = _normalize_entity_type(entity_type) + if effective_namespace == "user" and not user_id: + log_warning("EntityMemoryStore.list_entities: namespace='user' requires user_id") + return [] + + try: + for fetch in self._archive_headroom(limit=limit, include_archived=include_archived): + results = db.get_learnings( + learning_type=self.learning_type, + entity_type=entity_type, + namespace=effective_namespace, + user_id=user_id if effective_namespace == "user" else None, + limit=fetch, + ) + rows = results or [] + entities = self._parse_rows(rows, limit=limit, include_archived=include_archived) + if len(entities) >= limit or len(rows) < fetch: + return entities + return entities + except Exception as e: + log_debug(f"EntityMemoryStore.list_entities failed: {e}") + return [] + + async def alist_entities( + self, + entity_type: Optional[str] = None, + user_id: Optional[str] = None, + namespace: Optional[str] = None, + limit: int = 10, + include_archived: bool = False, + ) -> List[EntityMemory]: + """Async version of list_entities.""" + if not self.db: + return [] + + effective_namespace = namespace or self.config.namespace + entity_type = _normalize_entity_type(entity_type) + if effective_namespace == "user" and not user_id: + log_warning("EntityMemoryStore.alist_entities: namespace='user' requires user_id") + return [] + + try: + for fetch in self._archive_headroom(limit=limit, include_archived=include_archived): + if isinstance(self.db, AsyncBaseDb): + results = await self.db.get_learnings( + learning_type=self.learning_type, + entity_type=entity_type, + namespace=effective_namespace, + user_id=user_id if effective_namespace == "user" else None, + limit=fetch, + ) + else: + results = self.db.get_learnings( + learning_type=self.learning_type, + entity_type=entity_type, + namespace=effective_namespace, + user_id=user_id if effective_namespace == "user" else None, + limit=fetch, + ) + rows = results or [] + entities = self._parse_rows(rows, limit=limit, include_archived=include_archived) + if len(entities) >= limit or len(rows) < fetch: + return entities + return entities + except Exception as e: + log_debug(f"EntityMemoryStore.alist_entities failed: {e}") + return [] + + @staticmethod + def _archive_headroom(limit: int, include_archived: bool) -> List[int]: + """Fetch sizes to try when listing entities by recency. + + Archived rows are dropped after the fetch, so a fixed headroom quietly + shortens the listing once enough of the newest rows are archived - and + the entity directory then claims to be the full index while live + entities sit just outside the window. Escalate until the listing is + full or the table is exhausted. + """ + if include_archived: + return [limit] + return [limit * 2, limit * 8, limit * 32] + + @staticmethod + def _verified_search_window(limit: int) -> Iterator[int]: + """Fetch sizes to try when searching, until the backend runs out. + + The db-side LIKE matches the whole serialized row, key names included, + so ordinary words ("name", "content", "events") match every row: a + fixed window fills with rows that then fail value-only verification and + the real match, older, is never fetched. Any finite ladder just moves + the threshold - at 48x a brain needs 481 key-name matches to hide a + fact - so this yields until the caller stops, which it does when the + backend returns a short page. + """ + fetch = max(limit, 1) * 3 + while True: + yield fetch + fetch *= 4 + + def _anchored_hit( + self, query: str, entity_type: Optional[str], rows: List[Dict[str, Any]], include_archived: bool + ) -> List[EntityMemory]: + """The entity the query NAMES, ahead of the rows that merely mention it. + + search_learnings orders by recency and has no ranking, so an entity + whose name appears in five newer entities' facts ("works on Harbor + ingest") fills the verified window with rows that all legitimately + match, and its own row - older - is never fetched. Both of recall's + routes then close at once: the directory caps at 50 as well. + """ + found: List[EntityMemory] = [] + for row in rows: + if entity_type and row.get("entity_type") != entity_type: + continue + entity = self.schema.from_dict(row.get("content")) + if entity is None: + continue + if not include_archived and getattr(entity, "archived_at", None): + continue + found.append(entity) + return found + + @staticmethod + def _prepend_anchored(anchored: List[EntityMemory], rest: List[EntityMemory], limit: int) -> List[EntityMemory]: + merged = list(anchored) + seen = {(e.entity_id, e.entity_type) for e in merged} + for entity in rest: + key = (entity.entity_id, entity.entity_type) + if key in seen: + continue + seen.add(key) + merged.append(entity) + return merged[:limit] + + def _parse_rows(self, rows: List[Dict[str, Any]], limit: int, include_archived: bool) -> List[EntityMemory]: + entities: List[EntityMemory] = [] + for row in rows: + entity = self.schema.from_dict(row.get("content")) + if entity is None: + continue + if not include_archived and getattr(entity, "archived_at", None): + continue + entities.append(entity) + if len(entities) >= limit: + break + return entities def search( self, @@ -1720,52 +3108,82 @@ class EntityMemoryStore(LearningStore): user_id: Optional[str] = None, namespace: Optional[str] = None, limit: int = 10, + include_archived: bool = False, ) -> List[EntityMemory]: """Search for entities matching query. + Routes through the db's server-side search_learnings; falls back to the + client-side scan only when the backend does not implement it. Database + errors from the server-side path are raised, never swallowed - a broken + query must not present as an empty store. + Args: query: Search query (matched against name, facts, events, etc.). entity_type: Filter by entity type. user_id: User ID for "user" namespace scoping. namespace: Filter by namespace. limit: Maximum results to return. + include_archived: Include archived entities in results. Returns: List of matching EntityMemory objects. """ - if not self.db: + db = self._sync_db() + if db is None: return [] effective_namespace = namespace or self.config.namespace + entity_type = _normalize_entity_type(entity_type) + if effective_namespace == "user" and not user_id: + log_warning("EntityMemoryStore.search: namespace='user' requires user_id") + return [] - try: - results = self.db.get_learnings( - learning_type=self.learning_type, + if not callable(getattr(db, "search_learnings", None)): + self._log_degraded_search_once() + return self._search_client_side( + query=query, entity_type=entity_type, + user_id=user_id, namespace=effective_namespace, - user_id=user_id if effective_namespace == "user" else None, - limit=limit * 3, # Over-fetch for filtering + limit=limit, + include_archived=include_archived, ) - entities = [] - query_lower = query.lower() - - for result in results or []: # type: ignore[union-attr] - content = result.get("content", {}) - if self._matches_query(content=content, query=query_lower): - entity = self.schema.from_dict(content) - if entity: - entities.append(entity) - - if len(entities) >= limit: + try: + entities: List[EntityMemory] = [] + for fetch in self._verified_search_window(limit): + rows = db.search_learnings( + query=query, + learning_type=self.learning_type, + entity_type=entity_type, + namespace=effective_namespace, + user_id=user_id if effective_namespace == "user" else None, + limit=fetch, + ) + rows = rows or [] + entities = self._filter_rows_by_query(rows, query=query, limit=limit, include_archived=include_archived) + if len(entities) >= limit or len(rows) < fetch: break + except NotImplementedError: + self._log_degraded_search_once() + return self._search_client_side( + query=query, + entity_type=entity_type, + user_id=user_id, + namespace=effective_namespace, + limit=limit, + include_archived=include_archived, + ) - log_debug(f"EntityMemoryStore.search: found {len(entities)} entities for query: {query[:50]}...") - return entities - - except Exception as e: - log_debug(f"EntityMemoryStore.search failed: {e}") - return [] + anchored = self._anchored_hit( + query=query, + entity_type=entity_type, + rows=self._get_rows_by_entity_id(entity_id=_slugify(query), user_id=user_id, namespace=effective_namespace), + include_archived=include_archived, + ) + entities = self._prepend_anchored(anchored, entities, limit) + log_debug(f"EntityMemoryStore.search: found {len(entities)} entities for query: {query[:50]}") + return entities async def asearch( self, @@ -1774,771 +3192,242 @@ class EntityMemoryStore(LearningStore): user_id: Optional[str] = None, namespace: Optional[str] = None, limit: int = 10, + include_archived: bool = False, ) -> List[EntityMemory]: """Async version of search.""" if not self.db: return [] effective_namespace = namespace or self.config.namespace + entity_type = _normalize_entity_type(entity_type) + if effective_namespace == "user" and not user_id: + log_warning("EntityMemoryStore.asearch: namespace='user' requires user_id") + return [] + if not callable(getattr(self.db, "search_learnings", None)): + self._log_degraded_search_once() + return await self._asearch_client_side( + query=query, + entity_type=entity_type, + user_id=user_id, + namespace=effective_namespace, + limit=limit, + include_archived=include_archived, + ) + + try: + entities: List[EntityMemory] = [] + for fetch in self._verified_search_window(limit): + if isinstance(self.db, AsyncBaseDb): + rows = await self.db.search_learnings( + query=query, + learning_type=self.learning_type, + entity_type=entity_type, + namespace=effective_namespace, + user_id=user_id if effective_namespace == "user" else None, + limit=fetch, + ) + else: + rows = self.db.search_learnings( + query=query, + learning_type=self.learning_type, + entity_type=entity_type, + namespace=effective_namespace, + user_id=user_id if effective_namespace == "user" else None, + limit=fetch, + ) + rows = rows or [] + entities = self._filter_rows_by_query(rows, query=query, limit=limit, include_archived=include_archived) + if len(entities) >= limit or len(rows) < fetch: + break + except NotImplementedError: + self._log_degraded_search_once() + return await self._asearch_client_side( + query=query, + entity_type=entity_type, + user_id=user_id, + namespace=effective_namespace, + limit=limit, + include_archived=include_archived, + ) + + anchored = self._anchored_hit( + query=query, + entity_type=entity_type, + rows=await self._aget_rows_by_entity_id( + entity_id=_slugify(query), user_id=user_id, namespace=effective_namespace + ), + include_archived=include_archived, + ) + entities = self._prepend_anchored(anchored, entities, limit) + log_debug(f"EntityMemoryStore.asearch: found {len(entities)} entities for query: {query[:50]}") + return entities + + def _sync_db(self) -> Optional[Any]: + """The db for sync methods, or None when it is async (logged once). + + A sync call on an AsyncBaseDb would return an un-awaited coroutine and + surface as a confusing TypeError downstream; refuse it explicitly. + """ + if self.db is None: + return None + if isinstance(self.db, AsyncBaseDb): + if not self._async_db_in_sync_logged: + self._async_db_in_sync_logged = True + log_warning( + "EntityMemoryStore: a sync method was called with an async db (AsyncBaseDb). " + "Use the async API (arecall/aremember_about/asearch/...) with this backend." + ) + return None + return self.db + + def _log_degraded_search_once(self) -> None: + if not self._degraded_search_logged: + self._degraded_search_logged = True + log_warning( + "EntityMemoryStore: this db backend has no search_learnings implementation. " + "Entity search and resolution fall back to a scan of the most recently " + "updated rows, so past roughly 50 entities an alias or an externally-keyed " + "row stops resolving and a SECOND entity is created for the same thing - " + "there is no unmerge. Search quality degrades from the same point. " + "SqliteDb, PostgresDb and AsyncPostgresDb implement the real query surface; " + "use one of those for a store you intend to grow." + ) + + def _search_client_side( + self, + query: str, + entity_type: Optional[str], + user_id: Optional[str], + namespace: str, + limit: int, + include_archived: bool, + ) -> List[EntityMemory]: + """Degraded fallback: over-fetch recent rows and substring-match in Python.""" + db = self._sync_db() + if db is None: + return [] + try: + results = db.get_learnings( + learning_type=self.learning_type, + entity_type=entity_type, + namespace=namespace, + user_id=user_id if namespace == "user" else None, + limit=limit * 3, + ) + except Exception as e: + log_debug(f"EntityMemoryStore._search_client_side failed: {e}") + return [] + return self._filter_rows_by_query(results or [], query=query, limit=limit, include_archived=include_archived) + + async def _asearch_client_side( + self, + query: str, + entity_type: Optional[str], + user_id: Optional[str], + namespace: str, + limit: int, + include_archived: bool, + ) -> List[EntityMemory]: + """Async version of _search_client_side.""" try: if isinstance(self.db, AsyncBaseDb): results = await self.db.get_learnings( learning_type=self.learning_type, entity_type=entity_type, - namespace=effective_namespace, - user_id=user_id if effective_namespace == "user" else None, + namespace=namespace, + user_id=user_id if namespace == "user" else None, limit=limit * 3, ) else: - results = self.db.get_learnings( + results = self.db.get_learnings( # type: ignore[union-attr] learning_type=self.learning_type, entity_type=entity_type, - namespace=effective_namespace, - user_id=user_id if effective_namespace == "user" else None, + namespace=namespace, + user_id=user_id if namespace == "user" else None, limit=limit * 3, ) + except Exception as e: + log_debug(f"EntityMemoryStore._asearch_client_side failed: {e}") + return [] + return self._filter_rows_by_query(results or [], query=query, limit=limit, include_archived=include_archived) - entities = [] - query_lower = query.lower() - - for result in results or []: - content = result.get("content", {}) - if self._matches_query(content=content, query=query_lower): - entity = self.schema.from_dict(content) - if entity: - entities.append(entity) - + def _filter_rows_by_query( + self, rows: List[Dict[str, Any]], query: str, limit: int, include_archived: bool + ) -> List[EntityMemory]: + entities: List[EntityMemory] = [] + for row in rows: + try: + content = row.get("content", {}) + if not self._matches_query(content=content, query=query): + continue + entity = self.schema.from_dict(content) + if entity is None: + continue + if not include_archived and getattr(entity, "archived_at", None): + continue + entities.append(entity) if len(entities) >= limit: break - - log_debug(f"EntityMemoryStore.asearch: found {len(entities)} entities for query: {query[:50]}...") - return entities - - except Exception as e: - log_debug(f"EntityMemoryStore.asearch failed: {e}") - return [] + except Exception as e: + log_debug(f"EntityMemoryStore._filter_rows_by_query: skipping malformed row: {e}") + continue + return entities def _matches_query(self, content: Dict[str, Any], query: str) -> bool: - """Check if entity content matches search query.""" - # Check name - name = content.get("name", "") - if name and query in name.lower(): - return True + """Check if the entity content's VALUES match the query. - # Check entity_id - entity_id = content.get("entity_id", "") - if entity_id and query in entity_id.lower(): - return True - - # Check description - description = content.get("description", "") - if description and query in description.lower(): - return True - - # Check properties - properties = content.get("properties", {}) - for value in properties.values(): - if query in str(value).lower(): - return True - - # Check facts - facts = content.get("facts", []) - for fact in facts: - fact_content = fact.get("content", "") if isinstance(fact, dict) else str(fact) - if query in fact_content.lower(): - return True - - # Check events - events = content.get("events", []) - for event in events: - event_content = event.get("content", "") if isinstance(event, dict) else str(event) - if query in event_content.lower(): - return True - - # Check relationships - relationships = content.get("relationships", []) - for rel in relationships: - if isinstance(rel, dict): - if query in rel.get("entity_id", "").lower(): - return True - if query in rel.get("relation", "").lower(): - return True - - return False - - # ========================================================================= - # Create Operations - # ========================================================================= - - def create_entity( - self, - entity_id: str, - entity_type: str, - name: str, - description: Optional[str] = None, - properties: Optional[Dict[str, str]] = None, - user_id: Optional[str] = None, - agent_id: Optional[str] = None, - team_id: Optional[str] = None, - namespace: Optional[str] = None, - ) -> bool: - """Create a new entity. - - Args: - entity_id: Unique identifier for the entity. - entity_type: Type of entity. - name: Display name. - description: Brief description. - properties: Key-value properties. - user_id: User ID (required for "user" namespace). - agent_id: Agent context (stored for audit). - team_id: Team context (stored for audit). - namespace: Namespace for scoping. - - Returns: - True if created, False if already exists or error. + The db-side ILIKE matches the whole serialized document, keys included; + this value-scoped check verifies server hits (so "facts" or "name" do + not match every row) and drives the degraded client-side scan. Both + sides fold their separators, which is what the server's single-char + wildcard already does. """ - if not self.db: - return False + return values_match_query(content, query) - effective_namespace = namespace or self.config.namespace - - # Validate "user" namespace has user_id - if effective_namespace == "user" and not user_id: - log_warning("EntityMemoryStore.create_entity: 'user' namespace requires user_id") - return False - - # Check if already exists - existing = self.get( - entity_id=entity_id, - entity_type=entity_type, - user_id=user_id, - namespace=effective_namespace, - ) - if existing: - log_debug(f"EntityMemoryStore.create_entity: entity already exists {entity_type}/{entity_id}") - return False - - try: - now = datetime.now(timezone.utc).isoformat().replace("+00:00", "Z") - - entity = self.schema( - entity_id=entity_id, - entity_type=entity_type, - name=name, - description=description, - properties=properties or {}, - facts=[], - events=[], - relationships=[], - namespace=effective_namespace, - user_id=user_id if effective_namespace == "user" else None, - agent_id=agent_id, - team_id=team_id, - created_at=now, - updated_at=now, - ) - - self.db.upsert_learning( - id=self._build_entity_db_id(entity_id, entity_type, effective_namespace), - learning_type=self.learning_type, - entity_id=entity_id, - entity_type=entity_type, - namespace=effective_namespace, - user_id=user_id if effective_namespace == "user" else None, - agent_id=agent_id, - team_id=team_id, - content=entity.to_dict(), - ) - - log_debug(f"EntityMemoryStore.create_entity: created {entity_type}/{entity_id}") - return True - - except Exception as e: - log_debug(f"EntityMemoryStore.create_entity failed: {e}") - return False - - async def acreate_entity( + def delete( self, entity_id: str, entity_type: str, - name: str, - description: Optional[str] = None, - properties: Optional[Dict[str, str]] = None, - user_id: Optional[str] = None, - agent_id: Optional[str] = None, - team_id: Optional[str] = None, namespace: Optional[str] = None, ) -> bool: - """Async version of create_entity.""" + """Hard-delete an entity from the store (data API - not exposed as a tool).""" + db = self._sync_db() + if db is None: + return False + + effective_namespace = namespace or self.config.namespace + try: + return bool(db.delete_learning(id=self._build_entity_db_id(entity_id, entity_type, effective_namespace))) + except Exception as e: + log_debug(f"EntityMemoryStore.delete failed: {e}") + return False + + async def adelete( + self, + entity_id: str, + entity_type: str, + namespace: Optional[str] = None, + ) -> bool: + """Async version of delete.""" if not self.db: return False effective_namespace = namespace or self.config.namespace - - if effective_namespace == "user" and not user_id: - log_warning("EntityMemoryStore.acreate_entity: 'user' namespace requires user_id") - return False - - existing = await self.aget( - entity_id=entity_id, - entity_type=entity_type, - user_id=user_id, - namespace=effective_namespace, - ) - if existing: - log_debug(f"EntityMemoryStore.acreate_entity: entity already exists {entity_type}/{entity_id}") - return False - try: - now = datetime.now(timezone.utc).isoformat().replace("+00:00", "Z") - - entity = self.schema( - entity_id=entity_id, - entity_type=entity_type, - name=name, - description=description, - properties=properties or {}, - facts=[], - events=[], - relationships=[], - namespace=effective_namespace, - user_id=user_id if effective_namespace == "user" else None, - agent_id=agent_id, - team_id=team_id, - created_at=now, - updated_at=now, - ) - if isinstance(self.db, AsyncBaseDb): - await self.db.upsert_learning( - id=self._build_entity_db_id(entity_id, entity_type, effective_namespace), - learning_type=self.learning_type, - entity_id=entity_id, - entity_type=entity_type, - namespace=effective_namespace, - user_id=user_id if effective_namespace == "user" else None, - agent_id=agent_id, - team_id=team_id, - content=entity.to_dict(), + return bool( + await self.db.delete_learning( + id=self._build_entity_db_id(entity_id, entity_type, effective_namespace) + ) ) - else: - self.db.upsert_learning( - id=self._build_entity_db_id(entity_id, entity_type, effective_namespace), - learning_type=self.learning_type, - entity_id=entity_id, - entity_type=entity_type, - namespace=effective_namespace, - user_id=user_id if effective_namespace == "user" else None, - agent_id=agent_id, - team_id=team_id, - content=entity.to_dict(), - ) - - log_debug(f"EntityMemoryStore.acreate_entity: created {entity_type}/{entity_id}") - return True - + return bool( + self.db.delete_learning(id=self._build_entity_db_id(entity_id, entity_type, effective_namespace)) + ) except Exception as e: - log_debug(f"EntityMemoryStore.acreate_entity failed: {e}") + log_debug(f"EntityMemoryStore.adelete failed: {e}") return False - # ========================================================================= - # Update Operations - # ========================================================================= - - def update_entity( - self, - entity_id: str, - entity_type: str, - name: Optional[str] = None, - description: Optional[str] = None, - properties: Optional[Dict[str, str]] = None, - user_id: Optional[str] = None, - agent_id: Optional[str] = None, - team_id: Optional[str] = None, - namespace: Optional[str] = None, - ) -> bool: - """Update an existing entity's core properties. - - Args: - entity_id: The entity's identifier. - entity_type: Type of entity. - name: New display name (optional). - description: New description (optional). - properties: Properties to merge (optional). - user_id: User ID for namespace scoping. - agent_id: Agent context. - team_id: Team context. - namespace: Namespace to search in. - - Returns: - True if updated, False if not found. - """ - effective_namespace = namespace or self.config.namespace - - entity = self.get( - entity_id=entity_id, - entity_type=entity_type, - user_id=user_id, - namespace=effective_namespace, - ) - - if not entity: - return False - - # Update fields - if name is not None: - entity.name = name - if description is not None: - entity.description = description - if properties is not None: - entity.properties = {**(entity.properties or {}), **properties} - - entity.updated_at = datetime.now(timezone.utc).isoformat().replace("+00:00", "Z") - - return self._save_entity( - entity=entity, - user_id=user_id, - agent_id=agent_id, - team_id=team_id, - namespace=effective_namespace, - ) - - async def aupdate_entity( - self, - entity_id: str, - entity_type: str, - name: Optional[str] = None, - description: Optional[str] = None, - properties: Optional[Dict[str, str]] = None, - user_id: Optional[str] = None, - agent_id: Optional[str] = None, - team_id: Optional[str] = None, - namespace: Optional[str] = None, - ) -> bool: - """Async version of update_entity.""" - effective_namespace = namespace or self.config.namespace - - entity = await self.aget( - entity_id=entity_id, - entity_type=entity_type, - user_id=user_id, - namespace=effective_namespace, - ) - - if not entity: - return False - - if name is not None: - entity.name = name - if description is not None: - entity.description = description - if properties is not None: - entity.properties = {**(entity.properties or {}), **properties} - - entity.updated_at = datetime.now(timezone.utc).isoformat().replace("+00:00", "Z") - - return await self._asave_entity( - entity=entity, - user_id=user_id, - agent_id=agent_id, - team_id=team_id, - namespace=effective_namespace, - ) - - # ========================================================================= - # Fact Operations - # ========================================================================= - - def add_fact( - self, - entity_id: str, - entity_type: str, - fact: str, - user_id: Optional[str] = None, - agent_id: Optional[str] = None, - team_id: Optional[str] = None, - namespace: Optional[str] = None, - ) -> Optional[str]: - """Add a fact to an entity. - - Returns: - Fact ID if added, None if entity not found. - """ - effective_namespace = namespace or self.config.namespace - - entity = self.get( - entity_id=entity_id, - entity_type=entity_type, - user_id=user_id, - namespace=effective_namespace, - ) - - if not entity: - return None - - fact_id = entity.add_fact(fact) - entity.updated_at = datetime.now(timezone.utc).isoformat().replace("+00:00", "Z") - - success = self._save_entity( - entity=entity, - user_id=user_id, - agent_id=agent_id, - team_id=team_id, - namespace=effective_namespace, - ) - - return fact_id if success else None - - async def aadd_fact( - self, - entity_id: str, - entity_type: str, - fact: str, - user_id: Optional[str] = None, - agent_id: Optional[str] = None, - team_id: Optional[str] = None, - namespace: Optional[str] = None, - ) -> Optional[str]: - """Async version of add_fact.""" - effective_namespace = namespace or self.config.namespace - - entity = await self.aget( - entity_id=entity_id, - entity_type=entity_type, - user_id=user_id, - namespace=effective_namespace, - ) - - if not entity: - return None - - fact_id = entity.add_fact(fact) - entity.updated_at = datetime.now(timezone.utc).isoformat().replace("+00:00", "Z") - - success = await self._asave_entity( - entity=entity, - user_id=user_id, - agent_id=agent_id, - team_id=team_id, - namespace=effective_namespace, - ) - - return fact_id if success else None - - def update_fact( - self, - entity_id: str, - entity_type: str, - fact_id: str, - fact: str, - user_id: Optional[str] = None, - agent_id: Optional[str] = None, - team_id: Optional[str] = None, - namespace: Optional[str] = None, - ) -> bool: - """Update an existing fact.""" - effective_namespace = namespace or self.config.namespace - - entity = self.get( - entity_id=entity_id, - entity_type=entity_type, - user_id=user_id, - namespace=effective_namespace, - ) - - if not entity: - return False - - if not entity.update_fact(fact_id, fact): - return False - - entity.updated_at = datetime.now(timezone.utc).isoformat().replace("+00:00", "Z") - - return self._save_entity( - entity=entity, - user_id=user_id, - agent_id=agent_id, - team_id=team_id, - namespace=effective_namespace, - ) - - async def aupdate_fact( - self, - entity_id: str, - entity_type: str, - fact_id: str, - fact: str, - user_id: Optional[str] = None, - agent_id: Optional[str] = None, - team_id: Optional[str] = None, - namespace: Optional[str] = None, - ) -> bool: - """Async version of update_fact.""" - effective_namespace = namespace or self.config.namespace - - entity = await self.aget( - entity_id=entity_id, - entity_type=entity_type, - user_id=user_id, - namespace=effective_namespace, - ) - - if not entity: - return False - - if not entity.update_fact(fact_id, fact): - return False - - entity.updated_at = datetime.now(timezone.utc).isoformat().replace("+00:00", "Z") - - return await self._asave_entity( - entity=entity, - user_id=user_id, - agent_id=agent_id, - team_id=team_id, - namespace=effective_namespace, - ) - - def delete_fact( - self, - entity_id: str, - entity_type: str, - fact_id: str, - user_id: Optional[str] = None, - agent_id: Optional[str] = None, - team_id: Optional[str] = None, - namespace: Optional[str] = None, - ) -> bool: - """Delete a fact from an entity.""" - effective_namespace = namespace or self.config.namespace - - entity = self.get( - entity_id=entity_id, - entity_type=entity_type, - user_id=user_id, - namespace=effective_namespace, - ) - - if not entity: - return False - - if not entity.delete_fact(fact_id): - return False - - entity.updated_at = datetime.now(timezone.utc).isoformat().replace("+00:00", "Z") - - return self._save_entity( - entity=entity, - user_id=user_id, - agent_id=agent_id, - team_id=team_id, - namespace=effective_namespace, - ) - - async def adelete_fact( - self, - entity_id: str, - entity_type: str, - fact_id: str, - user_id: Optional[str] = None, - agent_id: Optional[str] = None, - team_id: Optional[str] = None, - namespace: Optional[str] = None, - ) -> bool: - """Async version of delete_fact.""" - effective_namespace = namespace or self.config.namespace - - entity = await self.aget( - entity_id=entity_id, - entity_type=entity_type, - user_id=user_id, - namespace=effective_namespace, - ) - - if not entity: - return False - - if not entity.delete_fact(fact_id): - return False - - entity.updated_at = datetime.now(timezone.utc).isoformat().replace("+00:00", "Z") - - return await self._asave_entity( - entity=entity, - user_id=user_id, - agent_id=agent_id, - team_id=team_id, - namespace=effective_namespace, - ) - - # ========================================================================= - # Event Operations - # ========================================================================= - - def add_event( - self, - entity_id: str, - entity_type: str, - event: str, - date: Optional[str] = None, - user_id: Optional[str] = None, - agent_id: Optional[str] = None, - team_id: Optional[str] = None, - namespace: Optional[str] = None, - ) -> Optional[str]: - """Add an event to an entity. - - Returns: - Event ID if added, None if entity not found. - """ - effective_namespace = namespace or self.config.namespace - - entity = self.get( - entity_id=entity_id, - entity_type=entity_type, - user_id=user_id, - namespace=effective_namespace, - ) - - if not entity: - return None - - event_id = entity.add_event(event, date=date) - entity.updated_at = datetime.now(timezone.utc).isoformat().replace("+00:00", "Z") - - success = self._save_entity( - entity=entity, - user_id=user_id, - agent_id=agent_id, - team_id=team_id, - namespace=effective_namespace, - ) - - return event_id if success else None - - async def aadd_event( - self, - entity_id: str, - entity_type: str, - event: str, - date: Optional[str] = None, - user_id: Optional[str] = None, - agent_id: Optional[str] = None, - team_id: Optional[str] = None, - namespace: Optional[str] = None, - ) -> Optional[str]: - """Async version of add_event.""" - effective_namespace = namespace or self.config.namespace - - entity = await self.aget( - entity_id=entity_id, - entity_type=entity_type, - user_id=user_id, - namespace=effective_namespace, - ) - - if not entity: - return None - - event_id = entity.add_event(event, date=date) - entity.updated_at = datetime.now(timezone.utc).isoformat().replace("+00:00", "Z") - - success = await self._asave_entity( - entity=entity, - user_id=user_id, - agent_id=agent_id, - team_id=team_id, - namespace=effective_namespace, - ) - - return event_id if success else None - - # ========================================================================= - # Relationship Operations - # ========================================================================= - - def add_relationship( - self, - entity_id: str, - entity_type: str, - related_entity_id: str, - relation: str, - direction: str = "outgoing", - user_id: Optional[str] = None, - agent_id: Optional[str] = None, - team_id: Optional[str] = None, - namespace: Optional[str] = None, - ) -> Optional[str]: - """Add a relationship to an entity. - - Returns: - Relationship ID if added, None if entity not found. - """ - effective_namespace = namespace or self.config.namespace - - entity = self.get( - entity_id=entity_id, - entity_type=entity_type, - user_id=user_id, - namespace=effective_namespace, - ) - - if not entity: - return None - - rel_id = entity.add_relationship(related_entity_id, relation, direction=direction) - entity.updated_at = datetime.now(timezone.utc).isoformat().replace("+00:00", "Z") - - success = self._save_entity( - entity=entity, - user_id=user_id, - agent_id=agent_id, - team_id=team_id, - namespace=effective_namespace, - ) - - return rel_id if success else None - - async def aadd_relationship( - self, - entity_id: str, - entity_type: str, - related_entity_id: str, - relation: str, - direction: str = "outgoing", - user_id: Optional[str] = None, - agent_id: Optional[str] = None, - team_id: Optional[str] = None, - namespace: Optional[str] = None, - ) -> Optional[str]: - """Async version of add_relationship.""" - effective_namespace = namespace or self.config.namespace - - entity = await self.aget( - entity_id=entity_id, - entity_type=entity_type, - user_id=user_id, - namespace=effective_namespace, - ) - - if not entity: - return None - - rel_id = entity.add_relationship(related_entity_id, relation, direction=direction) - entity.updated_at = datetime.now(timezone.utc).isoformat().replace("+00:00", "Z") - - success = await self._asave_entity( - entity=entity, - user_id=user_id, - agent_id=agent_id, - team_id=team_id, - namespace=effective_namespace, - ) - - return rel_id if success else None - # ========================================================================= # Internal Save Helpers # ========================================================================= @@ -2552,7 +3441,8 @@ class EntityMemoryStore(LearningStore): namespace: Optional[str] = None, ) -> bool: """Save entity to database.""" - if not self.db: + db = self._sync_db() + if db is None: return False effective_namespace = namespace or self.config.namespace @@ -2562,7 +3452,7 @@ class EntityMemoryStore(LearningStore): if not content: return False - self.db.upsert_learning( + db.upsert_learning( id=self._build_entity_db_id(entity.entity_id, entity.entity_type, effective_namespace), learning_type=self.learning_type, entity_id=entity.entity_id, @@ -2630,463 +3520,6 @@ class EntityMemoryStore(LearningStore): log_debug(f"EntityMemoryStore._asave_entity failed: {e}") return False - # ========================================================================= - # Background Extraction - # ========================================================================= - - def extract_and_save( - self, - messages: List[Any], - user_id: Optional[str] = None, - agent_id: Optional[str] = None, - team_id: Optional[str] = None, - namespace: Optional[str] = None, - run_metrics: Optional["RunMetrics"] = None, - ) -> None: - """Extract entities from messages (sync).""" - if not self.model or not self.db: - return - - try: - from agno.models.message import Message - - conversation_text = get_conversation_text(messages) - - tools = self._get_extraction_tools( - user_id=user_id, - agent_id=agent_id, - team_id=team_id, - namespace=namespace, - ) - - functions = self._build_functions_for_model(tools=tools) - - messages_for_model = [ - self._get_extraction_system_message(), - Message(role="user", content=f"Extract entities from this conversation:\n\n{conversation_text}"), - ] - - model_copy = deepcopy(self.model) - response = model_copy.response( - messages=messages_for_model, - tools=functions, - tool_call_limit=self.config.max_updates_per_run, - ) - - if run_metrics is not None and response.response_usage is not None: - from agno.metrics import ModelType, accumulate_model_metrics - - accumulate_model_metrics(response, model_copy, ModelType.LEARNING_MODEL, run_metrics) - - if response.tool_executions: - self.entity_updated = True - log_debug("EntityMemoryStore: Extraction saved entities") - - except Exception as e: - log_warning(f"EntityMemoryStore.extract_and_save failed: {str(e)}") - - async def aextract_and_save( - self, - messages: List[Any], - user_id: Optional[str] = None, - agent_id: Optional[str] = None, - team_id: Optional[str] = None, - namespace: Optional[str] = None, - run_metrics: Optional["RunMetrics"] = None, - ) -> None: - """Extract entities from messages (async).""" - if not self.model or not self.db: - return - - try: - conversation_text = get_conversation_text(messages) - - tools = self._aget_extraction_tools( - user_id=user_id, - agent_id=agent_id, - team_id=team_id, - namespace=namespace, - ) - - functions = self._build_functions_for_model(tools=tools) - - messages_for_model = [ - self._get_extraction_system_message(), - Message(role="user", content=f"Extract entities from this conversation:\n\n{conversation_text}"), - ] - - model_copy = deepcopy(self.model) - response = await model_copy.aresponse( - messages=messages_for_model, - tools=functions, - tool_call_limit=self.config.max_updates_per_run, - ) - - if run_metrics is not None and response.response_usage is not None: - from agno.metrics import ModelType, accumulate_model_metrics - - accumulate_model_metrics(response, model_copy, ModelType.LEARNING_MODEL, run_metrics) - - if response.tool_executions: - self.entity_updated = True - log_debug("EntityMemoryStore: Extraction saved entities") - - except Exception as e: - log_warning(f"EntityMemoryStore.aextract_and_save failed: {str(e)}") - - def _get_extraction_system_message(self) -> "Message": - """Get system message for extraction.""" - from agno.models.message import Message - - custom_instructions = self.config.instructions or "" - additional = self.config.additional_instructions or "" - - if self.config.system_message: - return Message(role="system", content=self.config.system_message) - - content = dedent("""\ - You are an Entity Extractor. Your job is to identify and capture knowledge about - external entities - people, companies, projects, products, systems, and other things - mentioned in conversations that are worth remembering. - - ## Philosophy - - Entity memory is your knowledge about the WORLD, distinct from: - - **User memory**: What you know about the user themselves - - **Learned knowledge**: Reusable task insights and patterns - - **Session context**: State of the current conversation - - Think of entity memory like a professional's mental rolodex - the accumulated knowledge - about clients, companies, technologies, and projects that helps you work effectively. - - ## Entity Structure - - Each entity has: - - **Core identity:** - - `entity_id`: Lowercase with underscores (e.g., "acme_corp", "jane_smith", "project_atlas") - - `entity_type`: Category - "person", "company", "project", "product", "system", "concept" - - `name`: Human-readable display name - - `description`: Brief description of what this entity is - - **Three types of memory:** - - 1. **Facts** (semantic memory) - Timeless truths about the entity - - "Uses PostgreSQL for their main database" - - "Headquarters in San Francisco" - - "Founded in 2019" - - "Prefers async communication" - - 2. **Events** (episodic memory) - Time-bound occurrences - - "Launched v2.0 on January 15, 2025" - - "Acquired by BigCorp in Q3 2024" - - "Had a major outage affecting 10K users" - - "Completed Series B funding" - - 3. **Relationships** (graph edges) - Connections to other entities - - "Bob Smith" --[CEO]--> "Acme Corp" - - "Project Atlas" --[uses]--> "PostgreSQL" - - "Acme Corp" --[competitor_of]--> "Beta Inc" - - "Jane" --[reports_to]--> "Bob" - - ## What to Extract - - **DO extract entities that are:** - - Named specifically (not just "a company" but "Acme Corp") - - Substantively discussed (not just mentioned in passing) - - Likely to be referenced again in future conversations - - Important to the user's work or context - - **DO capture:** - - Companies the user works with or mentions repeatedly - - People (colleagues, clients, stakeholders) with specific roles - - Projects with concrete details - - Products or systems with technical specifics - - Organizations relevant to the user's domain - - ## What NOT to Extract - - **DO NOT extract:** - - The user themselves (that belongs in UserProfile) - - Generic concepts without specific identity ("databases" vs "PostgreSQL") - - One-off mentions unlikely to recur ("I saw a company on the news") - - Entities with no substantive information to store - - Publicly available information that's easily searchable - - **Avoid:** - - Creating entities just because something was named - - Storing obvious facts ("Google is a tech company") - - Duplicating information across multiple entities unnecessarily - - ## Quality Guidelines - - **Good entity example:** - ``` - entity_id: "northstar_analytics" - entity_type: "company" - name: "NorthStar Analytics" - description: "Data analytics startup, potential client" - facts: - - "Series A stage, ~50 employees" - - "Tech stack: Python, Snowflake, dbt" - - "Main contact is Sarah Chen, VP Engineering" - - "Decision timeline is Q1 2025" - events: - - "Initial meeting held December 2024" - - "Requested technical deep-dive on ML capabilities" - relationships: - - sarah_chen --[works_at]--> northstar_analytics - ``` - - **Poor entity example:** - ``` - entity_id: "company1" # Too generic - name: "Some Company" # Vague - facts: - - "It's a company" # Obvious, not useful - ``` - - ## Extraction Guidelines - - 1. **Be selective**: Only extract entities with substantive, useful information - 2. **Be specific**: Capture concrete details, not vague generalities - 3. **Be accurate**: Only store information actually stated in the conversation - 4. **Categorize correctly**: Facts vs events vs relationships have different purposes - 5. **Use consistent IDs**: Lowercase, underscores, descriptive (e.g., "acme_corp" not "company_1") - - It's perfectly fine to extract nothing if no notable entities are mentioned. - Quality over quantity - one well-documented entity beats five sparse ones. - - """) - - if custom_instructions: - content += f"\n## Additional Instructions\n\n{custom_instructions}\n" - - if additional: - content += f"\n{additional}\n" - - return Message(role="system", content=content) - - def _get_extraction_tools( - self, - user_id: Optional[str] = None, - agent_id: Optional[str] = None, - team_id: Optional[str] = None, - namespace: Optional[str] = None, - ) -> List[Callable]: - """Get sync extraction tools based on config.""" - tools: List[Callable[..., str]] = [] - effective_namespace = namespace or self.config.namespace - - if self.config.enable_create_entity: - - def create_entity( - entity_id: str, - entity_type: str, - name: str, - description: Optional[str] = None, - ) -> str: - """Create a new entity.""" - success = self.create_entity( - entity_id=entity_id, - entity_type=entity_type, - name=name, - description=description, - user_id=user_id, - agent_id=agent_id, - team_id=team_id, - namespace=effective_namespace, - ) - return f"Created: {entity_type}/{entity_id}" if success else "Entity exists" - - tools.append(create_entity) - - if self.config.enable_add_fact: - - def add_fact(entity_id: str, entity_type: str, fact: str) -> str: - """Add a fact to an entity.""" - fact_id = self.add_fact( - entity_id=entity_id, - entity_type=entity_type, - fact=fact, - user_id=user_id, - agent_id=agent_id, - team_id=team_id, - namespace=effective_namespace, - ) - return f"Fact added: {fact_id}" if fact_id else "Entity not found" - - tools.append(add_fact) - - if self.config.enable_add_event: - - def add_event( - entity_id: str, - entity_type: str, - event: str, - date: Optional[str] = None, - ) -> str: - """Add an event to an entity.""" - event_id = self.add_event( - entity_id=entity_id, - entity_type=entity_type, - event=event, - date=date, - user_id=user_id, - agent_id=agent_id, - team_id=team_id, - namespace=effective_namespace, - ) - return f"Event added: {event_id}" if event_id else "Entity not found" - - tools.append(add_event) - - if self.config.enable_add_relationship: - - def add_relationship( - entity_id: str, - entity_type: str, - related_entity_id: str, - relation: str, - ) -> str: - """Add a relationship between entities.""" - rel_id = self.add_relationship( - entity_id=entity_id, - entity_type=entity_type, - related_entity_id=related_entity_id, - relation=relation, - user_id=user_id, - agent_id=agent_id, - team_id=team_id, - namespace=effective_namespace, - ) - return f"Relationship added: {rel_id}" if rel_id else "Entity not found" - - tools.append(add_relationship) - - return tools - - def _aget_extraction_tools( - self, - user_id: Optional[str] = None, - agent_id: Optional[str] = None, - team_id: Optional[str] = None, - namespace: Optional[str] = None, - ) -> List[Callable]: - """Get async extraction tools based on config.""" - tools: List[Callable] = [] - effective_namespace = namespace or self.config.namespace - - if self.config.enable_create_entity: - - async def create_entity( - entity_id: str, - entity_type: str, - name: str, - description: Optional[str] = None, - ) -> str: - """Create a new entity.""" - success = await self.acreate_entity( - entity_id=entity_id, - entity_type=entity_type, - name=name, - description=description, - user_id=user_id, - agent_id=agent_id, - team_id=team_id, - namespace=effective_namespace, - ) - return f"Created: {entity_type}/{entity_id}" if success else "Entity exists" - - tools.append(create_entity) - - if self.config.enable_add_fact: - - async def add_fact(entity_id: str, entity_type: str, fact: str) -> str: - """Add a fact to an entity.""" - fact_id = await self.aadd_fact( - entity_id=entity_id, - entity_type=entity_type, - fact=fact, - user_id=user_id, - agent_id=agent_id, - team_id=team_id, - namespace=effective_namespace, - ) - return f"Fact added: {fact_id}" if fact_id else "Entity not found" - - tools.append(add_fact) - - if self.config.enable_add_event: - - async def add_event( - entity_id: str, - entity_type: str, - event: str, - date: Optional[str] = None, - ) -> str: - """Add an event to an entity.""" - event_id = await self.aadd_event( - entity_id=entity_id, - entity_type=entity_type, - event=event, - date=date, - user_id=user_id, - agent_id=agent_id, - team_id=team_id, - namespace=effective_namespace, - ) - return f"Event added: {event_id}" if event_id else "Entity not found" - - tools.append(add_event) - - if self.config.enable_add_relationship: - - async def add_relationship( - entity_id: str, - entity_type: str, - related_entity_id: str, - relation: str, - ) -> str: - """Add a relationship between entities.""" - rel_id = await self.aadd_relationship( - entity_id=entity_id, - entity_type=entity_type, - related_entity_id=related_entity_id, - relation=relation, - user_id=user_id, - agent_id=agent_id, - team_id=team_id, - namespace=effective_namespace, - ) - return f"Relationship added: {rel_id}" if rel_id else "Entity not found" - - tools.append(add_relationship) - - return tools - - def _build_functions_for_model(self, tools: List[Callable]) -> List[Any]: - """Convert callables to Functions for model.""" - from agno.tools.function import Function - - functions = [] - seen_names = set() - - for tool in tools: - try: - name = tool.__name__ - if name in seen_names: - continue - seen_names.add(name) - - func = Function.from_callable(tool, strict=True) - func.strict = True - functions.append(func) - except Exception as e: - log_warning(f"Could not add function {tool}: {str(e)}") - - return functions - # ========================================================================= # Private Helpers # ========================================================================= @@ -3127,17 +3560,6 @@ class EntityMemoryStore(LearningStore): return "\n".join(parts) - def _format_entities_list(self, entities: List[EntityMemory]) -> str: - """Format entities for tool output.""" - parts = [] - for i, entity in enumerate(entities, 1): - if hasattr(entity, "get_context_text"): - formatted = entity.get_context_text() - else: - formatted = self._format_entity_basic(entity=entity) - parts.append(f"{i}. {formatted}") - return "\n\n".join(parts) - # ========================================================================= # Representation # ========================================================================= @@ -3172,27 +3594,6 @@ class EntityMemoryStore(LearningStore): user_id: User ID for "user" namespace scoping. namespace: Namespace to search in. raw: If True, print raw dict using pprint instead of formatted panel. - - Example: - >>> store.print(entity_id="acme_corp", entity_type="company") - ╭────────────────── Entity Memory ──────────────────╮ - │ Acme Corporation (company) │ - │ Enterprise software company │ - │ │ - │ Properties: │ - │ industry: fintech │ - │ size: startup │ - │ │ - │ Facts: │ - │ [dim][f1][/dim] Uses PostgreSQL for main DB │ - │ [dim][f2][/dim] API uses OAuth2 authentication │ - │ │ - │ Events: │ - │ [dim][e1][/dim] Launched v2.0 (2024-01-15) │ - │ │ - │ Relationships: │ - │ CEO → bob_smith │ - ╰────────────────── acme_corp ──────────────────────╯ """ from agno.learn.utils import print_panel @@ -3211,10 +3612,10 @@ class EntityMemoryStore(LearningStore): # Header: name and type name = getattr(entity, "name", None) etype = getattr(entity, "entity_type", entity_type) - if name: - lines.append(f"[bold]{name}[/bold] ({etype})") - else: - lines.append(f"[bold]{entity_id}[/bold] ({etype})") + header = f"[bold]{name or entity_id}[/bold] ({etype})" + if getattr(entity, "archived_at", None): + header += " [dim](archived)[/dim]" + lines.append(header) # Description description = getattr(entity, "description", None) @@ -3229,12 +3630,12 @@ class EntityMemoryStore(LearningStore): for key, value in properties.items(): lines.append(f" {key}: {value}") - # Facts - facts = getattr(entity, "facts", []) - if facts: + # Facts (live only) + live = entity.live_facts() if hasattr(entity, "live_facts") else getattr(entity, "facts", []) + if live: lines.append("") lines.append("Facts:") - for fact in facts: + for fact in live: if isinstance(fact, dict): fact_id = fact.get("id", "?") content = fact.get("content", str(fact)) diff --git a/libs/agno/agno/learn/stores/learned_knowledge.py b/libs/agno/agno/learn/stores/learned_knowledge.py index f5475504d2..f242e5e799 100644 --- a/libs/agno/agno/learn/stores/learned_knowledge.py +++ b/libs/agno/agno/learn/stores/learned_knowledge.py @@ -229,7 +229,11 @@ class LearnedKnowledgeStore(LearningStore): ) def build_context(self, data: Any) -> str: - """Build context for the agent. + """Build the DATA context for the agent. + + Data only (the relevant learnings) - the mode-specific how-to-use + guidance lives in instructions(); the automatic path concatenates the + two at the injection site. Args: data: List of learning objects from recall() (may be None). @@ -237,14 +241,19 @@ class LearnedKnowledgeStore(LearningStore): Returns: Context string to inject into the agent's system prompt. """ - mode = self.config.mode + return self._build_always_mode_context(data=data) + def instructions(self) -> str: + """Agent-facing guidance for this store, mode-aware. + + Guidance only - the recalled learnings live in build_context(). + """ + mode = self.config.mode if mode == LearningMode.PROPOSE: - return self._build_propose_mode_context(data=data) - elif mode == LearningMode.AGENTIC: - return self._build_agentic_mode_context(data=data) - else: - return self._build_background_mode_context(data=data) + return self._build_propose_mode_context(data=None) + if mode == LearningMode.AGENTIC: + return self._build_agentic_mode_context(data=None) + return "" def _build_agentic_mode_context(self, data: Any) -> str: """Build context for AGENTIC mode.""" @@ -373,7 +382,7 @@ class LearnedKnowledgeStore(LearningStore): return instructions - def _build_background_mode_context(self, data: Any) -> str: + def _build_always_mode_context(self, data: Any) -> str: """Build context for ALWAYS mode (just show relevant learnings).""" if not data: return "" @@ -383,16 +392,16 @@ class LearnedKnowledgeStore(LearningStore): return "" formatted = self._format_learnings_for_context(learnings=learnings) - return dedent(f"""\ - - Prior insights that may help with this task: - - {formatted} - - Apply these naturally if they're relevant to the current request. - Your current analysis and the user's specific context take precedence. - \ - """) + # Assembled by concatenation, not dedent: the interpolated learnings are + # flush-left and multi-line, which would defeat dedent's common prefix. + return ( + "\n" + "Prior insights that may help with this task:\n\n" + f"{formatted}\n\n" + "Apply these naturally if they're relevant to the current request.\n" + "Your current analysis and the user's specific context take precedence.\n" + "" + ) def _format_learnings_for_context(self, learnings: List[Any]) -> str: """Format learnings for inclusion in context.""" diff --git a/libs/agno/agno/learn/stores/protocol.py b/libs/agno/agno/learn/stores/protocol.py index b3235c6d59..b80ceeb691 100644 --- a/libs/agno/agno/learn/stores/protocol.py +++ b/libs/agno/agno/learn/stores/protocol.py @@ -76,10 +76,11 @@ class LearningStore(Protocol): ... def build_context(self, data: Any) -> str: - """Build context string for agent prompts. + """Build the DATA context string for agent prompts. - Formats the recalled data into a string that can be - injected into the agent's system prompt. + Formats the recalled data into a string that can be injected into the + agent's system prompt. Data only - the how-to-use guidance lives in + instructions(). Args: data: Data returned from recall(). @@ -89,6 +90,15 @@ class LearningStore(Protocol): """ ... + def instructions(self) -> str: + """The agent-facing guidance for this store. + + How and when to use the store's tools, mode-aware. Guidance only - the + recalled data lives in build_context(). Returns an empty string when the + store has nothing to instruct (e.g. no tools exposed). + """ + ... + def get_tools(self, **kwargs) -> List[Callable]: """Get tools to expose to the agent. diff --git a/libs/agno/agno/learn/stores/session_context.py b/libs/agno/agno/learn/stores/session_context.py index 2545ef70e0..28a5a8b75a 100644 --- a/libs/agno/agno/learn/stores/session_context.py +++ b/libs/agno/agno/learn/stores/session_context.py @@ -229,6 +229,10 @@ class SessionContextStore(LearningStore): \ """) + def instructions(self) -> str: + """Session context is system-managed (ALWAYS-only, no tools): no guidance.""" + return "" + def get_tools(self, **kwargs) -> List[Callable]: """Session context has no agent tools (system-managed only).""" return [] diff --git a/libs/agno/agno/learn/stores/user_memory.py b/libs/agno/agno/learn/stores/user_memory.py index fceb876e20..bb4f57678e 100644 --- a/libs/agno/agno/learn/stores/user_memory.py +++ b/libs/agno/agno/learn/stores/user_memory.py @@ -172,11 +172,11 @@ class UserMemoryStore(LearningStore): ) def build_context(self, data: Any) -> str: - """Build context for the agent. + """Build the DATA context for the agent. Formats memories data for injection into the agent's system prompt. - Designed to enable natural, personalized responses without meta-commentary - about memory systems. + Data only - the how-to-use guidance lives in instructions(); the + automatic path concatenates the two at the injection site. Args: data: Memories data from recall(). @@ -184,22 +184,13 @@ class UserMemoryStore(LearningStore): Returns: Context string to inject into the agent's system prompt. """ - # Build tool documentation based on what's enabled - tool_docs = self._build_tool_documentation() + empty_block = dedent("""\ + + No memories saved about this user yet. + """) if not data: - if self._should_expose_tools: - return ( - dedent("""\ - - No memories saved about this user yet. - - """) - + tool_docs - + dedent(""" - """) - ) - return "" + return empty_block if self._should_expose_tools else "" # Build memories section memories_text = None @@ -209,18 +200,7 @@ class UserMemoryStore(LearningStore): memories_text = "\n".join(f"- {m.get('content', str(m))}" for m in data.memories) if not memories_text: - if self._should_expose_tools: - return ( - dedent("""\ - - No memories saved about this user yet. - - """) - + tool_docs - + dedent(""" - """) - ) - return "" + return empty_block if self._should_expose_tools else "" context = "\n" context += memories_text + "\n" @@ -236,22 +216,23 @@ class UserMemoryStore(LearningStore): - Use memories to calibrate tone, depth, and examples without announcing it """) - if self._should_expose_tools: - context += ( - dedent(""" - - - - """) - + tool_docs - + dedent(""" - """) - ) - context += "\n" return context + def instructions(self) -> str: + """Agent-facing guidance for this store: when to save user memories. + + Guidance only - the recalled data lives in build_context(). Empty when + no tools are exposed (ALWAYS mode captures without agent involvement). + """ + if not self._should_expose_tools: + return "" + tool_docs = self._build_tool_documentation() + if not tool_docs: + return "" + return f"\n{tool_docs}\n" + def _build_tool_documentation(self) -> str: """Build documentation for available memory tools. @@ -750,6 +731,7 @@ class UserMemoryStore(LearningStore): user_id: str, agent_id: Optional[str] = None, team_id: Optional[str] = None, + instructed: bool = False, run_metrics: Optional["RunMetrics"] = None, ) -> str: """Extract memories from messages and save. @@ -793,7 +775,7 @@ class UserMemoryStore(LearningStore): functions = self._build_functions_for_model(tools=tools) messages_for_model = [ - self._get_system_message(existing_data=existing_data), + self._get_system_message(existing_data=existing_data, instructed=instructed), Message(role="user", content=f"Extract memories from this conversation:\n\n{conversation_text}"), ] @@ -822,6 +804,7 @@ class UserMemoryStore(LearningStore): user_id: str, agent_id: Optional[str] = None, team_id: Optional[str] = None, + instructed: bool = False, run_metrics: Optional["RunMetrics"] = None, ) -> str: """Async version of extract_and_save.""" @@ -855,7 +838,7 @@ class UserMemoryStore(LearningStore): functions = self._build_functions_for_model(tools=tools) messages_for_model = [ - self._get_system_message(existing_data=existing_data), + self._get_system_message(existing_data=existing_data, instructed=instructed), Message(role="user", content=f"Extract memories from this conversation:\n\n{conversation_text}"), ] @@ -908,6 +891,7 @@ class UserMemoryStore(LearningStore): user_id=user_id, agent_id=agent_id, team_id=team_id, + instructed=True, ) async def arun_memories_update( @@ -926,6 +910,7 @@ class UserMemoryStore(LearningStore): user_id=user_id, agent_id=agent_id, team_id=team_id, + instructed=True, ) # ========================================================================= @@ -981,8 +966,19 @@ class UserMemoryStore(LearningStore): def _get_system_message( self, existing_data: List[dict], + instructed: bool = False, ) -> "Message": - """Build system message for memory extraction.""" + """Build system message for memory extraction. + + Args: + existing_data: The user's current memories. + instructed: True when an agent called update_user_memory with an + explicit task. The decision to store has already been made, by + the agent that has the conversation and its instructions; this + store's job is then to execute it, not to re-judge it. What to + capture and what to leave alone belongs in the calling agent's + instructions, where the context to decide it exists. + """ from agno.models.message import Message if self.config.system_message is not None: @@ -1027,7 +1023,6 @@ class UserMemoryStore(LearningStore): - Patterns in how they think and work **DO NOT save:** - - Sensitive personal information (health conditions, financial details, relationships) unless directly relevant to helping them - One-off details unlikely to matter in future conversations - Information they'd find creepy to have remembered - Inferences or assumptions - only save what they've actually stated @@ -1035,6 +1030,33 @@ class UserMemoryStore(LearningStore): - Trivial preferences that don't affect interactions\ """) + if instructed: + # An agent called update_user_memory with a task. It had the + # conversation and its own instructions; it decided. Re-judging + # here is how an explicit "remember this privately" ended up + # stored as nothing. + capture_instructions = dedent("""\ + ## What To Capture + + An agent has instructed you to record something. Carry out that + instruction: save what it asked you to save, worded as given. + + - Do not second-guess whether the information is worth keeping, + sensitive, or awkward - that judgement was already made by the + agent holding the conversation, under its own instructions. + - Do still avoid duplicates: prefer update_memory over add_memory + when the instruction revises something already stored, and + remove what it asks you to remove. + - Record only what the instruction says. Do not add surrounding + detail it did not ask for.\ + """) + # A configured capture policy still applies - the execute contract + # replaces the store's own gatekeeping, not the operator's. The tool + # path is the ONLY path in AGENTIC mode, so dropping it here made + # UserMemoryConfig(instructions=...) a no-op in the flagship mode. + if self.config.instructions: + capture_instructions += f"\n\n{self.config.instructions}\n" + system_prompt += capture_instructions system_prompt += dedent(""" @@ -1098,7 +1120,8 @@ class UserMemoryStore(LearningStore): system_prompt += "- `clear_all_memories`: Reset all memories (use rarely)\n" # Examples - system_prompt += dedent(""" + system_prompt += ( + dedent(""" ## Examples **Example 1: New user introduction** @@ -1125,6 +1148,14 @@ class UserMemoryStore(LearningStore): - Respect boundaries: when in doubt about whether to save something, don't - It's fine to do nothing if the conversation reveals nothing worth remembering\ """) + if not instructed + else dedent(""" + ## Final Guidance + + - Carry out the instruction; do not weigh whether it should have been given + - Prefer updating an existing memory over adding a near-duplicate\ + """) + ) if self.config.additional_instructions: system_prompt += f"\n\n{self.config.additional_instructions}" diff --git a/libs/agno/agno/learn/stores/user_profile.py b/libs/agno/agno/learn/stores/user_profile.py index 5a1659758a..ca22b5ea40 100644 --- a/libs/agno/agno/learn/stores/user_profile.py +++ b/libs/agno/agno/learn/stores/user_profile.py @@ -182,11 +182,11 @@ class UserProfileStore(LearningStore): ) def build_context(self, data: Any) -> str: - """Build context for the agent. + """Build the DATA context for the agent. Formats user profile data for injection into the agent's system prompt. - Designed to enable natural, personalized responses without meta-commentary - about memory systems. + Data only - the how-to-use guidance lives in instructions(); the + automatic path concatenates the two at the injection site. Args: data: User profile data from recall(). @@ -194,22 +194,13 @@ class UserProfileStore(LearningStore): Returns: Context string to inject into the agent's system prompt. """ - # Build tool documentation based on what's enabled - tool_docs = self._build_tool_documentation() + empty_block = dedent("""\ + + No profile information saved about this user yet. + """) if not data: - if self._should_expose_tools: - return ( - dedent("""\ - - No profile information saved about this user yet. - - """) - + tool_docs - + dedent(""" - """) - ) - return "" + return empty_block if self._should_expose_tools else "" # Build profile fields section profile_parts = [] @@ -220,18 +211,7 @@ class UserProfileStore(LearningStore): profile_parts.append(f"{field_name.replace('_', ' ').title()}: {value}") if not profile_parts: - if self._should_expose_tools: - return ( - dedent("""\ - - No profile information saved about this user yet. - - """) - + tool_docs - + dedent(""" - """) - ) - return "" + return empty_block if self._should_expose_tools else "" context = "\n" context += "\n".join(profile_parts) + "\n" @@ -246,21 +226,23 @@ class UserProfileStore(LearningStore): - Current conversation always takes precedence over stored profile data """) - if self._should_expose_tools: - context += ( - dedent(""" - - - """) - + tool_docs - + dedent(""" - """) - ) - context += "\n" return context + def instructions(self) -> str: + """Agent-facing guidance for this store: when to update the profile. + + Guidance only - the recalled data lives in build_context(). Empty when + no tools are exposed (ALWAYS mode captures without agent involvement). + """ + if not self._should_expose_tools: + return "" + tool_docs = self._build_tool_documentation() + if not tool_docs: + return "" + return f"\n{tool_docs}\n" + def _build_tool_documentation(self) -> str: """Build documentation for available profile tools. diff --git a/libs/agno/agno/learn/utils.py b/libs/agno/agno/learn/utils.py index f95263aca0..43abb87fa0 100644 --- a/libs/agno/agno/learn/utils.py +++ b/libs/agno/agno/learn/utils.py @@ -57,6 +57,89 @@ def build_learning_id( return None +def content_values_text(content: Any) -> str: + """Flatten a content payload to a lowercased text of its VALUES only. + + Used to verify text-search hits: the db-side ILIKE matches the whole + serialized JSON document, keys included, so a query like "facts" or "name" + would match every row. Matching against the value projection restores + value-scoped precision without dropping any field. + """ + parts: List[str] = [] + + def walk(node: Any) -> None: + if isinstance(node, dict): + for value in node.values(): + walk(value) + elif isinstance(node, (list, tuple)): + for value in node: + walk(value) + elif node is not None: + parts.append(str(node)) + + walk(content) + # casefold, not lower: a stored "Ος" has to match a query of "ΟΣ". + return "\n".join(parts).casefold() + + +def separator_folded(text: str) -> str: + """Collapse runs of space, underscore and hyphen to one space, casefolded. + + The db-side pattern turns each separator run into the LIKE single-char + wildcard, which matches ANY separator. A verifier that enumerates uniform + rewrites cannot express mixed-separator text, so it threw away hits the + server had legitimately found: "end-to-end tests" stored and "end to end + tests" asked (or the reverse) verified as a miss, and the store answered + "no entities matching" for a fact it was holding. Folding both sides is + what the wildcard already means. + + Newlines are left alone: they separate one stored value from the next, and + a match must not span two of them. + """ + import re + + return re.sub(r"[ \t\r\f\v_\-]+", " ", text.casefold()) + + +def values_match_query(content: Any, query: str) -> bool: + """Whether the content's VALUES contain the query, separator-insensitively. + + The value-scoped half of the loose-prefilter/precise-verify pair: the + db-side ILIKE matches the whole serialized document, key names included, + so "facts" or "name" would otherwise match every row. + """ + needle = separator_folded(query).strip() + if not needle: + return False + return needle in separator_folded(content_values_text(content)) + + +def query_variants(query: str) -> List[str]: + """Lowercased query variants with word separators swapped. + + Mirrors the space/underscore(/hyphen) crossing the db-side search performs + with the LIKE single-char wildcard, so a client-side verification pass does + not drop hits the server legitimately matched ("sarah chen" vs sarah_chen). + + The query itself is the first variant. Rewriting every separator to one + character produces no form that matches a query which MIXES them - the db + matches "end-to-end tests" through the single-char wildcard, and a verifier + that only knows "end to end tests" / "end_to_end_tests" throws that hit + away. + """ + import re + + lowered = query.strip().casefold() + if not lowered: + return [] + variants: List[str] = [lowered] + for separator in (" ", "_", "-"): + variant = re.sub(r"[\s_\-]+", separator, lowered) + if variant and variant not in variants: + variants.append(variant) + return variants + + def _safe_get(data: Any, key: str, default: Any = None) -> Any: """Safely get a key from dict-like data. diff --git a/libs/agno/agno/team/_init.py b/libs/agno/agno/team/_init.py index 3722461bb5..7f114c49d4 100644 --- a/libs/agno/agno/team/_init.py +++ b/libs/agno/agno/team/_init.py @@ -633,6 +633,12 @@ def _set_learning_machine(team: "Team") -> None: team.learning.db = team.db if team.learning.model is None: team.learning.model = team.model + if ( + team.learning.learned_knowledge + and team.learning.knowledge is None + and getattr(team, "knowledge", None) is not None + ): + team.learning.knowledge = team.knowledge team._learning = team.learning # PROPOSE/HITL modes need chat history for multi-turn confirmation diff --git a/libs/agno/agno/team/_messages.py b/libs/agno/agno/team/_messages.py index 891cbeceef..170866e98a 100644 --- a/libs/agno/agno/team/_messages.py +++ b/libs/agno/agno/team/_messages.py @@ -52,6 +52,25 @@ from agno.utils.team import ( from agno.utils.timer import Timer +def _input_kwarg(method: Any, input_message: Any) -> Dict[str, Any]: + """``{"input": ...}`` only when the callee accepts it. + + ``Team.get_system_message`` is a public extension point and this is the + bound method, so a subclass written against the pre-2.8.4 signature is what + actually runs. Passing the new kwarg unconditionally makes every run of + such a team fail. + """ + import inspect + + try: + parameters = inspect.signature(method).parameters + except (TypeError, ValueError): + return {} + if "input" in parameters or any(p.kind == p.VAR_KEYWORD for p in parameters.values()): + return {"input": input_message} + return {} + + def _get_tool_names(member: Any, async_mode: bool = False) -> List[str]: """Extract tool names from a member's tools list.""" tool_names: List[str] = [] @@ -345,6 +364,7 @@ def get_system_message( files: Optional[Sequence[File]] = None, tools: Optional[List[Union[Function, dict]]] = None, add_session_state_to_context: Optional[bool] = None, + input: Optional[Any] = None, ) -> Optional[Message]: """Get the system message for the team. @@ -470,15 +490,25 @@ def get_system_message( # 2.2 Identity sections: description, role, instructions system_message_content += _build_identity_sections(team, instructions) - # 2.3 Learning context (user profile, user memory, session context) + # 2.3 Learning context: guidance + data, concatenated so the automatic door + # renders exactly what the manual door's instructions() + build_context() would if team._learning is not None and team.add_learnings_to_context: + from agno.agent._messages import _learning_message_text + + learning_guidance = team._learning._framework_instructions() learning_context = team._learning.build_context( user_id=user_id, session_id=session.session_id if session else None, team_id=team.id, + message=_learning_message_text(input), + run_context=run_context, + metadata=run_context.metadata if run_context else None, + dependencies=run_context.dependencies if run_context else None, + session_state=run_context.session_state if run_context else None, ) - if learning_context: - system_message_content += learning_context + "\n" + learning_block = "\n".join(part for part in (learning_guidance, learning_context) if part) + if learning_block: + system_message_content += learning_block + "\n" # 2.4 Knowledge base instructions if team.knowledge is not None and team.search_knowledge and team.add_search_knowledge_instructions: @@ -581,6 +611,7 @@ async def aget_system_message( files: Optional[Sequence[File]] = None, tools: Optional[List[Union[Function, dict]]] = None, add_session_state_to_context: Optional[bool] = None, + input: Optional[Any] = None, ) -> Optional[Message]: """Get the system message for the team.""" @@ -701,15 +732,24 @@ async def aget_system_message( # 2.2 Identity sections: description, role, instructions system_message_content += _build_identity_sections(team, instructions) - # 2.3 Learning context (user profile, user memory, session context) + # 2.3 Learning context (see the sync twin) if team._learning is not None and team.add_learnings_to_context: + from agno.agent._messages import _learning_message_text + + learning_guidance = team._learning._framework_instructions() learning_context = await team._learning.abuild_context( user_id=user_id, session_id=session.session_id if session else None, team_id=team.id, + message=_learning_message_text(input), + run_context=run_context, + metadata=run_context.metadata if run_context else None, + dependencies=run_context.dependencies if run_context else None, + session_state=run_context.session_state if run_context else None, ) - if learning_context: - system_message_content += learning_context + "\n" + learning_block = "\n".join(part for part in (learning_guidance, learning_context) if part) + if learning_block: + system_message_content += learning_block + "\n" # 2.4 Knowledge base instructions if team.knowledge is not None and team.search_knowledge and team.add_search_knowledge_instructions: @@ -850,6 +890,7 @@ def _get_run_messages( files=files, add_session_state_to_context=add_session_state_to_context, tools=tools, + **_input_kwarg(team.get_system_message, input_message), ) if system_message is not None: run_messages.system_message = system_message @@ -985,6 +1026,7 @@ async def _aget_run_messages( files=files, add_session_state_to_context=add_session_state_to_context, tools=tools, + **_input_kwarg(team.aget_system_message, input_message), ) if system_message is not None: run_messages.system_message = system_message diff --git a/libs/agno/agno/team/team.py b/libs/agno/agno/team/team.py index aefe625ee3..e9ee356b94 100644 --- a/libs/agno/agno/team/team.py +++ b/libs/agno/agno/team/team.py @@ -299,7 +299,10 @@ class Team: parse_response: bool = True # --- History --- - # Enable the agent to manage memories of the user + # Enable the agent to manage memories of the user. + # Do not combine with a LearningMachine that has a user_memory store: both + # register a tool named update_user_memory, tool parsing keeps the first + # name it sees, and the learning store's tool is dropped without a word. enable_agentic_memory: bool = False # If True, the agent creates/updates user memories at the end of runs update_memory_on_run: bool = False @@ -1442,6 +1445,7 @@ class Team: files: Optional[Sequence[File]] = None, tools: Optional[List[Union[Function, dict]]] = None, add_session_state_to_context: Optional[bool] = None, + input: Optional[Any] = None, ) -> Optional[Message]: return _messages.get_system_message( self, @@ -1453,6 +1457,7 @@ class Team: files=files, tools=tools, add_session_state_to_context=add_session_state_to_context, + input=input, ) async def aget_system_message( @@ -1465,6 +1470,7 @@ class Team: files: Optional[Sequence[File]] = None, tools: Optional[List[Union[Function, dict]]] = None, add_session_state_to_context: Optional[bool] = None, + input: Optional[Any] = None, ) -> Optional[Message]: return await _messages.aget_system_message( self, @@ -1476,6 +1482,7 @@ class Team: files=files, tools=tools, add_session_state_to_context=add_session_state_to_context, + input=input, ) ########################################################################### diff --git a/libs/agno/tests/integration/db/postgres/test_search_learnings.py b/libs/agno/tests/integration/db/postgres/test_search_learnings.py new file mode 100644 index 0000000000..cfb0f6e28d --- /dev/null +++ b/libs/agno/tests/integration/db/postgres/test_search_learnings.py @@ -0,0 +1,107 @@ +"""Integration tests for search_learnings on PostgresDb and AsyncPostgresDb. + +The JSONB content column has no ILIKE operator; these tests prove the +CAST-to-TEXT path live. Requires the pgvector container from +./cookbook/scripts/run_pgvector.sh (localhost:5532). +""" + +import uuid + +import pytest + +SYNC_DB_URL = "postgresql+psycopg://ai:ai@localhost:5532/ai" +ASYNC_DB_URL = "postgresql+psycopg_async://ai:ai@localhost:5532/ai" + + +def _postgres_reachable() -> bool: + try: + from sqlalchemy import create_engine, text + + engine = create_engine(SYNC_DB_URL) + with engine.connect() as conn: + conn.execute(text("SELECT 1")) + engine.dispose() + return True + except Exception: + return False + + +pytestmark = pytest.mark.skipif(not _postgres_reachable(), reason="postgres not reachable on localhost:5532") + + +@pytest.fixture +def db(): + from agno.db.postgres import PostgresDb + + return PostgresDb(db_url=SYNC_DB_URL) + + +@pytest.fixture +def learning_ids(): + return [f"test-search-{uuid.uuid4().hex[:8]}" for _ in range(3)] + + +def test_search_learnings_casts_jsonb_and_crosses_slug_boundary(db, learning_ids): + needle, spaced, decoy = learning_ids + try: + db.upsert_learning( + id=needle, + learning_type="entity_memory", + entity_id="sarah_chen", + namespace="test_search", + content={"entity_id": "sarah_chen", "facts": [{"content": "designs radar"}]}, + ) + db.upsert_learning( + id=spaced, + learning_type="entity_memory", + entity_id="acme_corp", + namespace="test_search", + content={"entity_id": "acme_corp", "name": "Acme Corp"}, + ) + db.upsert_learning( + id=decoy, + learning_type="entity_memory", + entity_id="unrelated", + namespace="test_search", + content={"entity_id": "unrelated", "facts": [{"content": "nothing here"}]}, + ) + + # JSONB path works at all (no cast -> this raises live) + rows = db.search_learnings(query="designs radar", namespace="test_search") + assert [r["learning_id"] for r in rows] == [needle] + + # Slug boundary: spaced query finds the slug... + rows = db.search_learnings(query="sarah chen", namespace="test_search") + assert [r["learning_id"] for r in rows] == [needle] + # ...and the underscore query finds the spaced display name + rows = db.search_learnings(query="acme_corp", namespace="test_search") + assert spaced in [r["learning_id"] for r in rows] + + # Filters compose + rows = db.search_learnings(query="sarah chen", namespace="test_search", entity_type="person") + assert rows == [] + finally: + for learning_id in learning_ids: + db.delete_learning(id=learning_id) + + +async def test_async_search_learnings_casts_jsonb(learning_ids): + from agno.db.postgres import AsyncPostgresDb + + adb = AsyncPostgresDb(db_url=ASYNC_DB_URL) + needle = learning_ids[0] + try: + await adb.upsert_learning( + id=needle, + learning_type="entity_memory", + entity_id="sarah_chen", + namespace="test_search_async", + content={"entity_id": "sarah_chen", "facts": [{"content": "prefers async communication"}]}, + ) + rows = await adb.search_learnings(query="sarah chen", namespace="test_search_async") + assert [r["learning_id"] for r in rows] == [needle] + + rows = await adb.search_learnings(query="ASYNC COMMUNICATION", namespace="test_search_async") + assert [r["learning_id"] for r in rows] == [needle] + finally: + await adb.delete_learning(id=needle) diff --git a/libs/agno/tests/integration/fs/test_toolkit_over_db.py b/libs/agno/tests/integration/fs/test_toolkit_over_db.py index c7a50ec8c3..3270bd1473 100644 --- a/libs/agno/tests/integration/fs/test_toolkit_over_db.py +++ b/libs/agno/tests/integration/fs/test_toolkit_over_db.py @@ -66,7 +66,7 @@ class TestToolkitOverDb: found = {m.path for m in fs.list(directory="seen")} assert found == {"seen/a.md"} - def test_read_only_toolkit_has_four_tools_over_db(self, db_fs): + def test_read_only_toolkit_has_three_tools_over_db(self, db_fs): tools = _fs(db_fs).tools(read_only=True) names = {f.name for f in tools.functions.values()} - assert names == {"read_file", "list_files", "search_content", "check_lines"} + assert names == {"read_file", "list_files", "search_content"} diff --git a/libs/agno/tests/unit/db/test_search_learnings.py b/libs/agno/tests/unit/db/test_search_learnings.py new file mode 100644 index 0000000000..c19d34f20e --- /dev/null +++ b/libs/agno/tests/unit/db/test_search_learnings.py @@ -0,0 +1,153 @@ +"""Unit tests for the search_learnings query surface. + +The sqlite path is proven here against a real SqliteDb; the postgres and +async_postgres paths share the same statement shape and are covered by the +integration suite. The base-class default must raise NotImplementedError so +stores know to fall back to their client-side scan. +""" + +import pytest + +from agno.db.base import AsyncBaseDb, BaseDb +from agno.db.sqlite import SqliteDb +from agno.db.utils import learning_search_patterns + + +@pytest.fixture +def db(tmp_path) -> SqliteDb: + return SqliteDb(db_file=str(tmp_path / "learnings.db")) + + +def _seed(db: SqliteDb, id: str, content: dict, **kwargs) -> None: + db.upsert_learning(id=id, learning_type=kwargs.pop("learning_type", "entity_memory"), content=content, **kwargs) + + +class TestPatterns: + def test_space_and_underscore_cross_via_wildcard(self) -> None: + # Spaces and underscores collapse to the single-char LIKE wildcard, so + # one pattern matches both "sarah chen" and "sarah_chen". + assert learning_search_patterns("sarah chen") == ["%sarah_chen%"] + assert learning_search_patterns("sarah_chen") == ["%sarah_chen%"] + assert learning_search_patterns("radar") == ["%radar%"] + + def test_percent_is_escaped(self) -> None: + assert learning_search_patterns("50% done") == ["%50\\%_done%"] + + def test_wildcard_only_query_yields_no_patterns(self) -> None: + assert learning_search_patterns(" ") == [] + assert learning_search_patterns("%") == [] + assert learning_search_patterns("___") == [] + assert learning_search_patterns("% _ %") == [] + + def test_non_ascii_gets_a_wildcarded_escape_variant(self) -> None: + patterns = learning_search_patterns("café") + assert patterns[0] == "%café%" + # SQLite stores JSON with ensure_ascii escapes and LIKE folds ASCII + # only, so a literal \\u00e9 can never case-match a stored \\u00c9. + # The escape is six characters wide and carries a wildcard per + # character; the caller's value-scoped Python check casefolds and + # rejects whatever that lets through. + assert patterns[1] == "%caf______%" + + def test_ascii_queries_get_no_wildcard_padding(self) -> None: + assert all("______" not in pattern for pattern in learning_search_patterns("sarah chen")) + + +class TestBaseDefaults: + def test_base_db_default_raises(self) -> None: + assert BaseDb.search_learnings.__qualname__ == "BaseDb.search_learnings" + with pytest.raises(NotImplementedError): + BaseDb.search_learnings(object(), query="x") # type: ignore[arg-type] + + async def test_async_base_db_default_raises(self) -> None: + with pytest.raises(NotImplementedError): + await AsyncBaseDb.search_learnings(object(), query="x") # type: ignore[arg-type] + + +class TestSqliteSearchLearnings: + def test_matches_content_case_insensitively(self, db: SqliteDb) -> None: + _seed(db, "e1", {"entity_id": "acme_corp", "facts": [{"content": "Uses PostgreSQL"}]}, entity_id="acme_corp") + rows = db.search_learnings(query="postgresql") + assert [r["learning_id"] for r in rows] == ["e1"] + + def test_search_crosses_slug_boundary(self, db: SqliteDb) -> None: + _seed(db, "e1", {"entity_id": "sarah_chen", "name": "sarah_chen"}, entity_id="sarah_chen") + _seed(db, "e2", {"entity_id": "acme_corp", "name": "acme_corp"}, entity_id="acme_corp") + assert [r["learning_id"] for r in db.search_learnings(query="sarah chen")] == ["e1"] + assert [r["learning_id"] for r in db.search_learnings(query="Acme Corp")] == ["e2"] + # And the reverse direction: an underscore query finds spaced content + _seed(db, "e3", {"entity_id": "jane_doe", "name": "Jane Doe"}, entity_id="jane_doe") + assert any(r["learning_id"] == "e3" for r in db.search_learnings(query="jane_doe")) + + def test_finds_match_outside_recent_window(self, db: SqliteDb) -> None: + # Insert 100 entities, then touch 40 of them so the needle row is far + # outside any recency window. The old over-fetch path scanned only the + # most recently updated rows and missed this by construction. + _seed(db, "needle", {"entity_id": "old_needle", "facts": [{"content": "the rare zanzibar detail"}]}) + for i in range(99): + _seed(db, f"filler-{i}", {"entity_id": f"filler_{i}", "facts": [{"content": f"filler fact {i}"}]}) + for i in range(40): + _seed(db, f"filler-{i}", {"entity_id": f"filler_{i}", "facts": [{"content": f"updated filler {i}"}]}) + + rows = db.search_learnings(query="zanzibar", limit=5) + assert [r["learning_id"] for r in rows] == ["needle"] + + def test_filters_and_order(self, db: SqliteDb) -> None: + _seed(db, "a", {"x": "same needle"}, namespace="global", entity_type="person") + _seed(db, "b", {"x": "same needle"}, namespace="global", entity_type="project") + _seed(db, "c", {"x": "same needle"}, namespace="private", entity_type="person") + _seed(db, "a", {"x": "same needle updated"}, namespace="global", entity_type="person") + + rows = db.search_learnings(query="needle", namespace="global") + assert [r["learning_id"] for r in rows] == ["a", "b"] or [r["learning_id"] for r in rows] == ["b", "a"] + + rows = db.search_learnings(query="needle", namespace="global", entity_type="person") + assert [r["learning_id"] for r in rows] == ["a"] + + rows = db.search_learnings(query="needle", limit=2) + assert len(rows) == 2 + + def test_empty_query_returns_empty(self, db: SqliteDb) -> None: + _seed(db, "a", {"x": "content"}) + assert db.search_learnings(query=" ") == [] + + def test_non_ascii_query_finds_json_escaped_content(self, db: SqliteDb) -> None: + # SQLite serializes JSON with ensure_ascii=True, so "café" is stored as + # "café" - the raw pattern alone would silently miss it. + _seed(db, "e1", {"entity_id": "cafe", "name": "Café Régulier"}, entity_id="cafe") + rows = db.search_learnings(query="café") + assert [r["learning_id"] for r in rows] == ["e1"] + + def test_wildcard_query_does_not_match_everything(self, db: SqliteDb) -> None: + _seed(db, "a", {"x": "alpha"}) + _seed(db, "b", {"x": "beta"}) + assert db.search_learnings(query="%") == [] + assert db.search_learnings(query="___") == [] + # An embedded % is literal, not a wildcard + _seed(db, "c", {"x": "progress: 50% done"}) + assert [r["learning_id"] for r in db.search_learnings(query="50% done")] == ["c"] + assert db.search_learnings(query="50%zzz") == [] + + def test_null_updated_at_sorts_last(self, db: SqliteDb) -> None: + from sqlalchemy import text + + _seed(db, "real", {"x": "needle latest"}) + _seed(db, "nulled", {"x": "needle nulled"}) + with db.Session() as sess, sess.begin(): + sess.execute(text("UPDATE agno_learnings SET updated_at = NULL WHERE learning_id = 'nulled'")) + + rows = db.search_learnings(query="needle") + assert [r["learning_id"] for r in rows] == ["real", "nulled"] + + def test_search_fails_loudly_on_db_error(self, db: SqliteDb, monkeypatch: pytest.MonkeyPatch) -> None: + # get_learnings swallows db errors into []; search_learnings must not. + _seed(db, "a", {"x": "content"}) + + def broken_session() -> None: + raise RuntimeError("connection lost") + + monkeypatch.setattr(db, "Session", broken_session) + with pytest.raises(RuntimeError, match="connection lost"): + db.search_learnings(query="content") + # ...while get_learnings demonstrates today's swallowing behavior + assert db.get_learnings() == [] diff --git a/libs/agno/tests/unit/fs/test_toolkit.py b/libs/agno/tests/unit/fs/test_toolkit.py index ccb5ac1c55..54d36492b2 100644 --- a/libs/agno/tests/unit/fs/test_toolkit.py +++ b/libs/agno/tests/unit/fs/test_toolkit.py @@ -46,13 +46,19 @@ def toolkit(fs) -> FileSystemTools: return fs.tools() +@pytest.fixture +def full_toolkit(fs) -> FileSystemTools: + # The whole nine-tool surface, selected explicitly. + return fs.tools(include_tools=FileSystemTools.FULL_TOOLS) + + class TestSchemas: """The empty-schema regression (spec D1): every schema must be non-empty AND name every documented parameter AND lack the framework-injected ones.""" @pytest.mark.parametrize("tool_name", list(EXPECTED_TOOL_PARAMS.keys())) - def test_schema_names_documented_params_and_lacks_injected(self, toolkit, tool_name): - function = toolkit.functions[tool_name] + def test_schema_names_documented_params_and_lacks_injected(self, full_toolkit, tool_name): + function = full_toolkit.functions[tool_name] function.process_entrypoint() properties = function.parameters.get("properties", {}) assert list(properties.keys()) == EXPECTED_TOOL_PARAMS[tool_name] @@ -62,8 +68,8 @@ class TestSchemas: assert prop.get("description"), "every model-facing parameter carries its docstring description" @pytest.mark.parametrize("tool_name", list(EXPECTED_TOOL_PARAMS.keys())) - def test_async_schema_names_documented_params(self, toolkit, tool_name): - function = toolkit.async_functions[tool_name] + def test_async_schema_names_documented_params(self, full_toolkit, tool_name): + function = full_toolkit.async_functions[tool_name] function.process_entrypoint() properties = function.parameters.get("properties", {}) assert list(properties.keys()) == EXPECTED_TOOL_PARAMS[tool_name] @@ -74,9 +80,9 @@ class TestSchemas: assert prop.get("description"), "every async parameter carries its docstring description" @pytest.mark.parametrize("tool_name", list(EXPECTED_TOOL_PARAMS.keys())) - def test_async_description_matches_sync(self, toolkit, tool_name): - sync_fn = toolkit.functions[tool_name] - async_fn = toolkit.async_functions[tool_name] + def test_async_description_matches_sync(self, full_toolkit, tool_name): + sync_fn = full_toolkit.functions[tool_name] + async_fn = full_toolkit.async_functions[tool_name] sync_fn.process_entrypoint() async_fn.process_entrypoint() assert async_fn.description == sync_fn.description @@ -130,19 +136,45 @@ class TestWorkspaceParity: assert toolkit.requires_confirmation_tools == [] def test_confirmation_opt_in_via_kwargs(self, fs): - tk = fs.tools(requires_confirmation_tools=["delete_file"]) + tk = fs.tools(allow_delete=True, requires_confirmation_tools=["delete_file"]) assert tk.functions["delete_file"].requires_confirmation is True class TestSurface: - def test_full_surface_registers_nine_sync_and_async(self, toolkit): - assert list(toolkit.functions.keys()) == FileSystemTools.FULL_TOOLS - assert list(toolkit.async_functions.keys()) == FileSystemTools.FULL_TOOLS + def test_default_surface_is_the_notes_seven(self, toolkit): + assert list(toolkit.functions.keys()) == FileSystemTools.DEFAULT_TOOLS + assert list(toolkit.async_functions.keys()) == FileSystemTools.DEFAULT_TOOLS + assert len(FileSystemTools.DEFAULT_TOOLS) == 7 + assert "check_lines" not in toolkit.functions + assert "delete_file" not in toolkit.functions - def test_read_only_registers_exactly_four(self, fs): + def test_allow_delete_adds_delete_file(self, fs): + tk = fs.tools(allow_delete=True) + assert list(tk.functions.keys()) == FileSystemTools.DEFAULT_TOOLS + ["delete_file"] + + def test_read_only_registers_exactly_three(self, fs): tk = fs.tools(read_only=True) assert list(tk.functions.keys()) == FileSystemTools.READ_ONLY_TOOLS assert list(tk.async_functions.keys()) == FileSystemTools.READ_ONLY_TOOLS + assert len(FileSystemTools.READ_ONLY_TOOLS) == 3 + + def test_check_lines_reachable_via_include_tools(self, fs): + tk = fs.tools(include_tools=FileSystemTools.DEFAULT_TOOLS + ["check_lines"]) + assert "check_lines" in tk.functions + read_tk = fs.tools(read_only=True, include_tools=["read_file", "check_lines"]) + assert set(read_tk.functions.keys()) == {"read_file", "check_lines"} + + def test_include_tools_cannot_reach_write_tools_when_read_only(self, fs): + with pytest.raises(ValueError): + fs.tools(read_only=True, include_tools=["write_file"]) + + def test_read_only_with_allow_delete_raises(self, fs): + with pytest.raises(ValueError): + fs.tools(read_only=True, allow_delete=True) + + def test_full_tools_stays_exported_as_the_whole_surface(self, fs): + tk = fs.tools(include_tools=FileSystemTools.FULL_TOOLS) + assert list(tk.functions.keys()) == FileSystemTools.FULL_TOOLS def test_toolkit_name(self, toolkit): assert toolkit.name == "filesystem" @@ -158,7 +190,14 @@ class TestInstructions: text = FileSystem.instructions() assert text.startswith("You have your own private, durable filesystem") assert "Never store secrets, passwords, or API keys." in text - assert 'check_lines(lines, directory="seen")' in text + assert "replace_lines" in text + assert "move_file" in text and "archive/" in text + # The record-set/seen-directory conventions moved to the durable-records + # cookbook, and the user-memory steer is gone: this text is the notes + # contract, nothing else. + assert "check_lines" not in text + assert "seen/" not in text + assert "user memory" not in text def test_read_only_variant_names_no_write_tool(self, fs): text = FileSystem.instructions(read_only=True) @@ -256,8 +295,7 @@ class TestErrorStringsVerbatim: result = toolkit.write_file("a.md", "0123456789x") assert result == ( "Error: a.md would be 11 bytes (limit 10 per file). " - "Start a new file (for record logs, partition by date, e.g. seen/2026-07-24.md) " - "or delete files you no longer need." + "Split the topic into smaller files (or partition by date) and retry." ) def test_quota_namespace_string(self, tmp_path): @@ -266,10 +304,10 @@ class TestErrorStringsVerbatim: toolkit.write_file("a.md", "123456") result = toolkit.write_file("b.md", "78901") assert result == ( - "Error: storage is full (6 of 10 bytes). Delete only files you are certain are obsolete " - "(see list_files), such as an old date partition, then retry. Do not overwrite or delete " - "records you might still need to make room; if nothing is safely disposable, stop and " - "report that storage is full." + "Error: storage is full (6 of 10 bytes). Free space only if you have a tool for it and " + "only from files you are certain are obsolete (see list_files). Never overwrite or " + "discard records you might still need to make room; if nothing is safely disposable, " + "stop and report that storage is full." ) def test_check_lines_count_string(self, toolkit): @@ -660,3 +698,76 @@ class TestReplaceLines: def test_not_in_the_read_only_surface(self, fs): assert "replace_lines" not in fs.tools(read_only=True).functions + + +class TestExcludeToolsTypos: + """Tolerating an exclusion that left the default set is the upgrade idiom. + + Tolerating a name that is no tool at all is a silent regression: the + caller believes a tool is gone and it is still registered. + """ + + def test_a_typo_warns_and_excludes_nothing(self, fs, caplog) -> None: + import logging + + with caplog.at_level(logging.WARNING): + toolkit = fs.tools(exclude_tools=["delete_fil"]) + assert any("not FileSystem tools" in r.getMessage() for r in caplog.records) + # the call still resolves; nothing was excluded for the misspelled name + assert set(toolkit.functions) == set(FileSystemTools.DEFAULT_TOOLS) + + def test_excluding_a_tool_outside_this_set_is_a_no_op(self, fs) -> None: + # delete_file left the default set in 2.8.4; the old safety idiom stays valid + names = set(fs.tools(exclude_tools=["delete_file"]).functions) + assert names == set(FileSystemTools.DEFAULT_TOOLS) + + def test_excluding_a_registered_tool_still_removes_it(self, fs) -> None: + names = set(fs.tools(exclude_tools=["move_file"]).functions) + assert "move_file" not in names + assert names == set(FileSystemTools.DEFAULT_TOOLS) - {"move_file"} + + +class TestConcurrentEditsInOneTurn: + """A model's tool calls for one turn are gathered concurrently, and every + mutating fs tool is a read-modify-write. Without a per-file lock the second + write lands on the pre-edit content and one edit vanishes, with both tool + results reporting success.""" + + @pytest.mark.asyncio + async def test_two_replace_lines_on_one_note_both_land(self, tmp_path) -> None: + import asyncio + + from agno.db.sqlite import SqliteDb + from agno.fs import FileSystem + + fs = FileSystem(SqliteDb(db_file=str(tmp_path / "n.db")), namespace="brain") + tools = fs.tools() + fs.write("notes/radar.md", "line one\nline two\nline three\n") + + await asyncio.gather( + tools.areplace_lines("notes/radar.md", 1, 1, "EDIT ONE"), + tools.areplace_lines("notes/radar.md", 3, 3, "EDIT THREE"), + ) + + content = fs.read("notes/radar.md") + assert "EDIT ONE" in content and "EDIT THREE" in content + assert "line two" in content + + @pytest.mark.asyncio + async def test_an_append_beside_an_edit_survives(self, tmp_path) -> None: + import asyncio + + from agno.db.sqlite import SqliteDb + from agno.fs import FileSystem + + fs = FileSystem(SqliteDb(db_file=str(tmp_path / "m.db")), namespace="brain") + tools = fs.tools() + fs.write("notes/log.md", "start\n") + + await asyncio.gather( + tools.aappend_file("notes/log.md", "appended"), + tools.areplace_lines("notes/log.md", 1, 1, "REPLACED"), + ) + + content = fs.read("notes/log.md") + assert "REPLACED" in content and "appended" in content diff --git a/libs/agno/tests/unit/learn/test_decision_log_search.py b/libs/agno/tests/unit/learn/test_decision_log_search.py new file mode 100644 index 0000000000..f1005fc144 --- /dev/null +++ b/libs/agno/tests/unit/learn/test_decision_log_search.py @@ -0,0 +1,265 @@ +"""Unit tests for DecisionLogStore search routing through search_learnings.""" + +import logging +from typing import Any, Dict, List + +import pytest + +from agno.db.base import BaseDb +from agno.learn.config import DecisionLogConfig +from agno.learn.schemas import DecisionLog +from agno.learn.stores.decision_log import DecisionLogStore + + +class FakeDecisionDb(BaseDb): + """Subclasses BaseDb (the sync search path requires a real BaseDb) but + only implements what DecisionLogStore touches.""" + + def __init__(self) -> None: # noqa: D401 - do not call BaseDb.__init__ + self.rows: List[Dict[str, Any]] = [] + self.search_calls: List[Dict[str, Any]] = [] + + # --- what the store uses --- + def upsert_learning(self, id: str, **kwargs: Any) -> None: + self.rows.append({"learning_id": id, **kwargs}) + + def get_learnings(self, **kwargs: Any) -> List[Dict[str, Any]]: + limit = kwargs.get("limit") + rows = list(reversed(self.rows)) + return rows[:limit] if limit else rows + + def search_learnings(self, query: str, **kwargs: Any) -> List[Dict[str, Any]]: + import json + + self.search_calls.append({"query": query, **kwargs}) + rows = [row for row in reversed(self.rows) if query.lower() in json.dumps(row.get("content", {})).lower()] + limit = kwargs.get("limit") + return rows[:limit] if limit else rows + + # --- abstract surface stubs --- + def __getattr__(self, name: str) -> Any: + raise AttributeError(name) + + +# BaseDb is an ABC with many abstract methods; bypass instantiation checks. +FakeDecisionDb.__abstractmethods__ = frozenset() + + +def _store(db: FakeDecisionDb) -> DecisionLogStore: + return DecisionLogStore(config=DecisionLogConfig(db=db)) + + +def _log(store: DecisionLogStore, decision: str, decision_type: str = "tool_selection") -> None: + import uuid + + store.save(decision=DecisionLog(id=f"dec_{uuid.uuid4().hex[:6]}", decision=decision, decision_type=decision_type)) + + +def test_query_routes_through_search_learnings() -> None: + db = FakeDecisionDb() + store = _store(db) + _log(store, "Chose Postgres over Dynamo") + _log(store, "Used web search for news") + + results = store.search(query="postgres") + assert len(results) == 1 + assert results[0].decision == "Chose Postgres over Dynamo" + assert db.search_calls and db.search_calls[0]["learning_type"] == "decision_log" + + +def test_no_query_uses_listing_path() -> None: + db = FakeDecisionDb() + store = _store(db) + _log(store, "first") + _log(store, "second") + + results = store.search() + assert len(results) == 2 + assert db.search_calls == [] + + +def test_fallback_on_not_implemented_filters_client_side(caplog: pytest.LogCaptureFixture) -> None: + class NoSearchDb(FakeDecisionDb): + def search_learnings(self, query: str, **kwargs: Any) -> List[Dict[str, Any]]: + raise NotImplementedError + + NoSearchDb.__abstractmethods__ = frozenset() + db = NoSearchDb() + store = _store(db) + _log(store, "Chose Postgres over Dynamo") + _log(store, "Used web search for news") + + with caplog.at_level(logging.WARNING): + results = store.search(query="postgres") + store.search(query="postgres") + + assert len(results) == 1 + assert results[0].decision == "Chose Postgres over Dynamo" + degraded = [r for r in caplog.records if "no search_learnings implementation" in r.getMessage()] + assert len(degraded) == 1 + + +async def test_async_query_routes_through_search_learnings() -> None: + db = FakeDecisionDb() + store = _store(db) + _log(store, "Chose Postgres over Dynamo") + + results = await store.asearch(query="postgres", session_id="sess-9") + assert len(results) == 1 + assert db.search_calls and db.search_calls[0]["query"] == "postgres" + # The sync-db branch of asearch honors session_id too + assert db.search_calls[0]["session_id"] == "sess-9" + + +def test_server_hit_in_any_field_is_kept() -> None: + # A hit whose only match is in a field to_text() omits (alternatives) must + # survive the client-side value verification. + import uuid + + db = FakeDecisionDb() + store = _store(db) + store.save( + decision=DecisionLog( + id=f"dec_{uuid.uuid4().hex[:6]}", + decision="Picked the database", + alternatives=["postgres", "dynamo"], + ) + ) + + results = store.search(query="postgres") + assert len(results) == 1 + assert db.search_calls # served by the server-side path, kept by the verifier + + +def test_json_key_names_do_not_match_decisions() -> None: + db = FakeDecisionDb() + store = _store(db) + _log(store, "Chose Postgres over Dynamo") + + assert store.search(query="reasoning") == [] + assert store.search(query="decision_type") == [] + + +def test_decision_log_defaults_agentic_and_never_writes_contentless_rows() -> None: + from agno.learn.config import LearningMode + + assert DecisionLogConfig().mode is LearningMode.AGENTIC + + db = FakeDecisionDb() + store = _store(db) + + # The old ALWAYS extraction wrote one contentless row per tool call from + # the message snapshot; process is now a no-op. + class FakeToolCallMessage: + tool_calls = [type("TC", (), {"name": "web_search"})()] + + store.process(messages=[FakeToolCallMessage()], agent_id="a1", session_id="s1") + assert db.rows == [] + assert not hasattr(store, "_extract_decisions_from_messages") + + +def test_search_honors_session_id() -> None: + db = FakeDecisionDb() + captured: List[Dict[str, Any]] = [] + original = db.search_learnings + + def spy(query: str, **kwargs: Any) -> List[Dict[str, Any]]: + captured.append(kwargs) + return original(query, **kwargs) + + db.search_learnings = spy # type: ignore[method-assign] + store = _store(db) + _log(store, "Chose Postgres over Dynamo") + + store.search(query="postgres", session_id="sess-42") + assert captured and captured[0]["session_id"] == "sess-42" + + # And the listing path passes it to get_learnings + listing_calls: List[Dict[str, Any]] = [] + original_get = db.get_learnings + + def spy_get(**kwargs: Any) -> List[Dict[str, Any]]: + listing_calls.append(kwargs) + return original_get(**kwargs) + + db.get_learnings = spy_get # type: ignore[method-assign] + store.search(session_id="sess-42") + assert listing_calls and listing_calls[0]["session_id"] == "sess-42" + + +def test_decision_type_filter_composes_with_query() -> None: + db = FakeDecisionDb() + store = _store(db) + _log(store, "Chose Postgres over Dynamo", decision_type="architecture") + _log(store, "Postgres connection retry", decision_type="tool_selection") + + results = store.search(query="postgres", decision_type="architecture") + assert len(results) == 1 + assert results[0].decision_type == "architecture" + + +def test_query_variants_keeps_the_query_itself() -> None: + """A mixed-separator query has no uniform-separator form. + + The db matches "end-to-end tests" through the LIKE single-char wildcard; + a verifier that only knows "end to end tests" and "end_to_end_tests" + discards that hit and the store answers "No entities matching". + """ + from agno.learn.utils import query_variants + + assert "end-to-end tests" in query_variants("end-to-end tests") + assert "user_id column" in query_variants("user_id column") + # the separator-swapped forms are still generated + assert "end_to_end_tests" in query_variants("end-to-end tests") + assert "sarah_chen" in query_variants("sarah chen") + assert query_variants(" ") == [] + + +def test_search_patterns_wildcard_non_ascii_escapes() -> None: + """SQLite compares the JSON escape as literal text and LIKE folds ASCII + only, so no pre-cased whole-string variant reaches a mixed form like "Ος". + The escape carries a wildcard per character instead; the caller's + value-scoped Python check casefolds and rejects the slack. + """ + from agno.db.utils import learning_search_patterns + + patterns = learning_search_patterns("Ος") + assert any("______" in p for p in patterns), patterns + # the raw form is still offered, for Postgres where ::text is real characters + assert any("Ος" in p for p in patterns), patterns + # ASCII queries gain no wildcard padding + assert all("______" not in p for p in learning_search_patterns("sarah chen")) + + +def test_value_projection_and_variants_casefold() -> None: + from agno.learn.utils import content_values_text, query_variants + + assert content_values_text({"facts": ["Ος"]}) == query_variants("ΟΣ")[0] + assert query_variants("CAFÉ")[0] == "café" + + +def test_recall_is_not_scoped_to_the_current_session(tmp_path) -> None: + """A decision log's value is cross-session. + + recall() receives the run's session_id from LearningMachine, and forwarding + it into search() made the injected block say "No recent decisions logged" to + every new session while the store was full. search(session_id=...) still + scopes for explicit callers. + """ + from agno.db.sqlite import SqliteDb + from agno.learn.config import DecisionLogConfig, LearningMode + from agno.learn.stores.decision_log import DecisionLogStore + + store = DecisionLogStore( + config=DecisionLogConfig(db=SqliteDb(db_file=str(tmp_path / "d.db")), mode=LearningMode.AGENTIC) + ) + log_decision = next( + t for t in store.get_tools(agent_id="a1", session_id="session-one") if t.__name__ == "log_decision" + ) + log_decision(decision="Postgres over Dynamo", reasoning="advisory locks") + + from_another_session = store.recall(agent_id="a1", session_id="session-two") + assert from_another_session and from_another_session[0].decision == "Postgres over Dynamo" + + assert store.search(agent_id="a1", session_id="session-two") == [] + assert len(store.search(agent_id="a1", session_id="session-one")) == 1 diff --git a/libs/agno/tests/unit/learn/test_deprecations_and_guards.py b/libs/agno/tests/unit/learn/test_deprecations_and_guards.py new file mode 100644 index 0000000000..e9d6b3cc1b --- /dev/null +++ b/libs/agno/tests/unit/learn/test_deprecations_and_guards.py @@ -0,0 +1,245 @@ +"""Unit tests for the 2.8.4 deprecation warnings and visibility guards.""" + +import logging +from typing import Any, Dict, List, Optional + +from agno.learn import LearningMachine + + +class RecordingLearningDb: + def __init__(self) -> None: + self.rows: Dict[str, Dict[str, Any]] = {} + + def get_learning(self, **kwargs: Any) -> Optional[Dict[str, Any]]: + return None + + def upsert_learning(self, id: str, **kwargs: Any) -> None: + self.rows[id] = {**kwargs, "learning_id": id} + + def get_learnings(self, **kwargs: Any) -> List[Dict[str, Any]]: + return [] + + def search_learnings(self, query: str, **kwargs: Any) -> List[Dict[str, Any]]: + return [] + + def delete_learning(self, id: str) -> bool: + return self.rows.pop(id, None) is not None + + +def test_enable_user_memories_is_quiet(caplog, tmp_path) -> None: + """enable_user_memories still aliases update_memory_on_run and says nothing. + + Every rebuild path passes the alias explicitly (deep_copy from the attribute, + Agent.load with a False default), so a deprecation warning here fires on + agents whose author never wrote the parameter. + """ + from agno.agent import Agent + from agno.db.sqlite import SqliteDb + + with caplog.at_level(logging.WARNING): + agent = Agent(db=SqliteDb(db_file=str(tmp_path / "a.db")), enable_user_memories=True) + agent.deep_copy() + assert agent.update_memory_on_run is True + assert not any("enable_user_memories" in r.getMessage() for r in caplog.records) + + +def test_hitl_mode_warns_unsupported_without_deprecating(caplog) -> None: + from agno.learn.config import LearningMode, UserMemoryConfig + from agno.learn.stores.user_memory import UserMemoryStore + + with caplog.at_level(logging.WARNING): + UserMemoryStore(config=UserMemoryConfig(db=RecordingLearningDb(), mode=LearningMode.HITL)) # type: ignore[arg-type] + messages = [r.getMessage() for r in caplog.records] + assert any("does not support HITL mode" in m for m in messages) + assert not any("deprecated" in m for m in messages) + + +def test_agentic_memory_collision_is_quiet(caplog, tmp_path) -> None: + """set_learning_machine runs on every run, so a guard that logs there logs + once per run. The collision is documented instead.""" + from agno.agent import Agent + from agno.agent._init import set_learning_machine + from agno.db.sqlite import SqliteDb + + db = SqliteDb(db_file=str(tmp_path / "b.db")) + agent = Agent(db=db, learning=True, enable_agentic_memory=True) + with caplog.at_level(logging.WARNING): + set_learning_machine(agent) + set_learning_machine(agent) + assert not any("silently dropped" in r.getMessage() for r in caplog.records) + + +def test_missing_user_id_with_per_user_stores_warns_once(caplog) -> None: + machine = LearningMachine(db=RecordingLearningDb(), user_memory=True, user_profile=True) # type: ignore[arg-type] + with caplog.at_level(logging.WARNING): + machine.get_tools(user_id=None) + machine.get_tools(user_id=None) + warnings = [r for r in caplog.records if "no user_id" in r.getMessage()] + assert len(warnings) == 1 + assert "user_profile" in warnings[0].getMessage() and "user_memory" in warnings[0].getMessage() + + +def test_no_missing_user_id_warning_when_user_present_or_no_per_user_stores(caplog) -> None: + machine = LearningMachine(db=RecordingLearningDb(), user_memory=True) # type: ignore[arg-type] + with caplog.at_level(logging.WARNING): + machine.get_tools(user_id="u1") + assert not any("no user_id" in r.getMessage() for r in caplog.records) + + entity_only = LearningMachine(db=RecordingLearningDb(), entity_memory=True) # type: ignore[arg-type] + with caplog.at_level(logging.WARNING): + entity_only.get_tools(user_id=None) + assert not any("no user_id" in r.getMessage() for r in caplog.records) + + +def test_backend_without_search_warns_once_on_the_write_path(caplog, tmp_path) -> None: + """A backend with no search_learnings degrades RESOLUTION, not just search. + + Past ~50 entities an alias stops resolving and the write mints a second + entity for the same thing, with no unmerge - so the warning has to reach + the write path, and has to name that consequence. + """ + import asyncio + + from agno.db.sqlite import AsyncSqliteDb + from agno.learn.config import EntityMemoryConfig, LearningMode + from agno.learn.stores.entity_memory import EntityMemoryStore + + db = AsyncSqliteDb(db_file=str(tmp_path / "a.db")) + store = EntityMemoryStore( + config=EntityMemoryConfig(db=db, mode=LearningMode.AGENTIC, namespace="global") # type: ignore[arg-type] + ) + + async def write_three() -> None: + for i in range(3): + await store.aremember_about( + entity=f"Thing {i}", entity_type="company", aliases=[f"T{i}"], namespace="global" + ) + + with caplog.at_level(logging.WARNING): + asyncio.run(write_three()) + + warnings = [r.getMessage() for r in caplog.records if "no search_learnings implementation" in r.getMessage()] + assert len(warnings) == 1, warnings + assert "SECOND entity" in warnings[0] + assert "SqliteDb, PostgresDb and AsyncPostgresDb" in warnings[0] + + +def test_the_three_first_tier_backends_implement_search() -> None: + from agno.db.postgres import AsyncPostgresDb, PostgresDb + from agno.db.sqlite import SqliteDb + + for backend in (SqliteDb, PostgresDb, AsyncPostgresDb): + assert "search_learnings" in backend.__dict__, backend.__name__ + + +def test_manual_door_without_a_model_says_capture_is_off(caplog, tmp_path) -> None: + """The manual door injects nothing, so a hand-built machine keeps model=None. + + Capture is a model call: the tools return "No model provided" and entity + memory retires nothing, while the guidance block promises the opposite. + """ + from agno.db.sqlite import SqliteDb + from agno.learn.config import LearningMode, UserMemoryConfig + + machine = LearningMachine( + db=SqliteDb(db_file=str(tmp_path / "a.db")), + user_memory=UserMemoryConfig(mode=LearningMode.AGENTIC), + entity_memory=True, + ) + with caplog.at_level(logging.WARNING): + machine.get_tools(user_id="composer@example.com") + machine.get_tools(user_id="composer@example.com") + + warnings = [r.getMessage() for r in caplog.records if "LearningMachine has no model" in r.getMessage()] + assert len(warnings) == 1, warnings + assert "user_memory" in warnings[0] and "entity_memory" in warnings[0] + + +def test_the_automatic_door_is_quiet_because_the_agent_injects_its_model(caplog, tmp_path) -> None: + from agno.agent import Agent + from agno.db.sqlite import SqliteDb + from agno.learn.config import LearningMode, UserMemoryConfig + + machine = LearningMachine(user_memory=UserMemoryConfig(mode=LearningMode.AGENTIC), entity_memory=True) + agent = Agent(model="openai:gpt-5.5", db=SqliteDb(db_file=str(tmp_path / "b.db")), learning=machine) + with caplog.at_level(logging.WARNING): + agent.initialize_agent() + machine.get_tools(user_id="composer@example.com") + + assert not any("LearningMachine has no model" in r.getMessage() for r in caplog.records) + + +def test_agent_knowledge_does_not_conjure_a_learned_knowledge_store(tmp_path) -> None: + """`knowledge=` on the agent must not add stores the machine never asked for. + + The copy exists to repair a round trip (from_dict cannot carry a knowledge + base). Unconditional, it turned any `Agent(knowledge=kb, learning=machine)` + into two extra tools - `save_learning` writes into the user's production + index - plus a second guidance block competing with the machine's own. + """ + from agno.agent import Agent + from agno.db.sqlite import SqliteDb + from agno.learn import EntityMemoryConfig, LearningMode, UserMemoryConfig + + class FakeKnowledge: + name = "kb" + + machine = LearningMachine( + user_memory=UserMemoryConfig(mode=LearningMode.AGENTIC), entity_memory=EntityMemoryConfig() + ) + agent = Agent( + model="openai:gpt-5.5", + db=SqliteDb(db_file=str(tmp_path / "a.db")), + knowledge=FakeKnowledge(), + learning=machine, + user_id="u1", + ) + agent.initialize_agent() + + assert sorted(machine.stores) == ["entity_memory", "user_memory"] + assert "save_learning" not in {t.__name__ for t in machine.get_tools(user_id="u1")} + + +def test_a_wanted_learned_knowledge_store_still_gets_the_agents_knowledge(tmp_path) -> None: + from agno.agent import Agent + from agno.db.sqlite import SqliteDb + from agno.learn import LearningMode, UserMemoryConfig + + class FakeKnowledge: + name = "kb" + + machine = LearningMachine(user_memory=UserMemoryConfig(mode=LearningMode.AGENTIC), learned_knowledge=True) + agent = Agent( + model="openai:gpt-5.5", + db=SqliteDb(db_file=str(tmp_path / "b.db")), + knowledge=FakeKnowledge(), + learning=machine, + user_id="u1", + ) + agent.initialize_agent() + + assert machine.knowledge is not None + assert "learned_knowledge" in machine.stores + + +def test_a_configured_capture_policy_survives_the_execute_contract(tmp_path) -> None: + """`instructions` is the capture-policy knob, and the tool path is the only + path AGENTIC mode takes. The execute contract replaces the store's own + gatekeeping, not the operator's policy.""" + from agno.db.sqlite import SqliteDb + from agno.learn.config import LearningMode, UserMemoryConfig + from agno.learn.stores.user_memory import UserMemoryStore + + store = UserMemoryStore( + config=UserMemoryConfig( + db=SqliteDb(db_file=str(tmp_path / "m.db")), + mode=LearningMode.AGENTIC, + instructions="House policy: never record salaries.", + ) + ) + instructed = store._get_system_message(existing_data=[], instructed=True).content + passive = store._get_system_message(existing_data=[], instructed=False).content + + assert "House policy: never record salaries." in instructed + assert "House policy: never record salaries." in passive + assert "Carry out that" in instructed diff --git a/libs/agno/tests/unit/learn/test_entity_memory_store.py b/libs/agno/tests/unit/learn/test_entity_memory_store.py new file mode 100644 index 0000000000..5439669d93 --- /dev/null +++ b/libs/agno/tests/unit/learn/test_entity_memory_store.py @@ -0,0 +1,1405 @@ +"""Unit tests for the revamped EntityMemoryStore. + +Entity memory is AGENTIC-only: the agent records through four tools +(remember_about, link_entities, search_entities, forget) and there is no +extraction pass. These tests run offline against a recording fake db. +""" + +import inspect +from typing import Any, Dict, List, Optional + +import pytest + +from agno.learn.config import EntityMemoryConfig, LearningMode +from agno.learn.stores.entity_memory import EntityMemoryStore + + +class RecordingLearningDb: + """In-memory fake of the learnings table, keyed by learning_id.""" + + def __init__(self) -> None: + self.rows: Dict[str, Dict[str, Any]] = {} + self._clock = 0 + + def get_learning(self, **kwargs: Any) -> Optional[Dict[str, Any]]: + learning_type = kwargs.get("learning_type") + entity_id = kwargs.get("entity_id") + entity_type = kwargs.get("entity_type") + namespace = kwargs.get("namespace") + for row in self.rows.values(): + if ( + row.get("learning_type") == learning_type + and row.get("entity_id") == entity_id + and row.get("entity_type") == entity_type + and row.get("namespace") == namespace + ): + return row + return None + + def upsert_learning(self, id: str, **kwargs: Any) -> None: + existing = self.rows.get(id, {}) + row = {**existing, **kwargs, "learning_id": id} + self._clock += 1 + row["updated_at"] = self._clock + self.rows[id] = row + + def get_learnings(self, **kwargs: Any) -> List[Dict[str, Any]]: + learning_type = kwargs.get("learning_type") + entity_id = kwargs.get("entity_id") + entity_type = kwargs.get("entity_type") + namespace = kwargs.get("namespace") + limit = kwargs.get("limit") + rows = [ + row + for row in self.rows.values() + if (learning_type is None or row.get("learning_type") == learning_type) + and (entity_id is None or row.get("entity_id") == entity_id) + and (entity_type is None or row.get("entity_type") == entity_type) + and (namespace is None or row.get("namespace") == namespace) + ] + rows.sort(key=lambda r: r.get("updated_at", 0), reverse=True) + if limit is not None: + rows = rows[:limit] + return rows + + def delete_learning(self, id: str) -> bool: + return self.rows.pop(id, None) is not None + + def search_learnings(self, query: str, **kwargs: Any) -> List[Dict[str, Any]]: + import json + + limit = kwargs.pop("limit", None) + kwargs.pop("workflow_id", None) + kwargs.pop("session_id", None) + kwargs.pop("agent_id", None) + kwargs.pop("team_id", None) + kwargs.pop("user_id", None) + candidates = self.get_learnings(**kwargs) + variants = {query.lower(), query.lower().replace(" ", "_"), query.lower().replace("_", " ")} + rows = [row for row in candidates if any(v in json.dumps(row.get("content", {})).lower() for v in variants)] + if limit is not None: + rows = rows[:limit] + return rows + + +class _SleepyJudge: + """Stands in for the supersession judge's provider round trip. + + The await is the point: it is the suspension that lets an assistant turn's + gathered tool calls interleave inside a read-modify-write. + """ + + id = "sleepy-judge" + + def response(self, messages: Any, tools: Any = None, **kwargs: Any) -> None: + return None + + async def aresponse(self, messages: Any, tools: Any = None, **kwargs: Any) -> None: + import asyncio + + await asyncio.sleep(0.01) + return None + + +@pytest.fixture +def db() -> RecordingLearningDb: + return RecordingLearningDb() + + +@pytest.fixture +def store(db: RecordingLearningDb) -> EntityMemoryStore: + return EntityMemoryStore(config=EntityMemoryConfig(db=db)) # type: ignore[arg-type] + + +class TestAgenticOnly: + def test_default_mode_is_agentic(self) -> None: + assert EntityMemoryConfig().mode is LearningMode.AGENTIC + + @pytest.mark.parametrize("mode", [LearningMode.ALWAYS, LearningMode.PROPOSE, LearningMode.HITL]) + def test_non_agentic_mode_raises(self, mode: LearningMode) -> None: + with pytest.raises(ValueError, match="AGENTIC-only"): + EntityMemoryStore(config=EntityMemoryConfig(mode=mode)) + + def test_extraction_api_is_gone(self, store: EntityMemoryStore) -> None: + for name in ( + "extract_and_save", + "aextract_and_save", + "_get_extraction_tools", + "_aget_extraction_tools", + "_get_extraction_system_message", + ): + assert not hasattr(store, name) + + def test_process_is_a_noop(self, store: EntityMemoryStore, db: RecordingLearningDb) -> None: + store.process(messages=[object()], user_id="user-1") + assert db.rows == {} + assert store.was_updated is False + + async def test_aprocess_is_a_noop(self, store: EntityMemoryStore, db: RecordingLearningDb) -> None: + await store.aprocess(messages=[object()], user_id="user-1") + assert db.rows == {} + assert store.was_updated is False + + def test_machine_bool_input_builds_agentic_store(self, db: RecordingLearningDb) -> None: + from agno.learn import LearningMachine + + machine = LearningMachine(db=db, entity_memory=True) # type: ignore[arg-type] + entity_store = machine.entity_memory_store + assert entity_store is not None + assert entity_store.config.mode is LearningMode.AGENTIC + + def test_store_config_namespace_wins_over_machine_default(self, db: RecordingLearningDb) -> None: + # EntityMemoryConfig(namespace="ops") under a default machine must not + # be silently overridden to "global" at the tool call sites. + from agno.learn import LearningMachine + + machine = LearningMachine(db=db, entity_memory=EntityMemoryConfig(namespace="ops")) # type: ignore[arg-type] + tools = machine.get_tools(user_id="u1") + remember = next(t for t in tools if t.__name__ == "remember_about") + remember(entity="radar", entity_type="project") + assert all(row.get("namespace") == "ops" for row in db.rows.values()) + + def test_machine_namespace_reaches_default_configs(self, db: RecordingLearningDb) -> None: + from agno.learn import LearningMachine + + machine = LearningMachine(db=db, namespace="team_west", entity_memory=True) # type: ignore[arg-type] + tools = machine.get_tools(user_id="u1") + remember = next(t for t in tools if t.__name__ == "remember_about") + remember(entity="radar", entity_type="project") + assert all(row.get("namespace") == "team_west" for row in db.rows.values()) + + +class TestToolSurface: + def test_sync_tools_are_the_four(self, store: EntityMemoryStore) -> None: + tools = store.get_tools(user_id="user-1") + assert [t.__name__ for t in tools] == ["remember_about", "link_entities", "search_entities", "forget"] + assert all(not inspect.iscoroutinefunction(t) for t in tools) + + async def test_async_tools_are_the_four(self, store: EntityMemoryStore) -> None: + tools = await store.aget_tools(user_id="user-1") + assert [t.__name__ for t in tools] == ["remember_about", "link_entities", "search_entities", "forget"] + assert all(inspect.iscoroutinefunction(t) for t in tools) + + def test_sync_and_async_docstrings_match(self, store: EntityMemoryStore) -> None: + import asyncio + + sync_tools = store.get_tools() + async_tools = asyncio.run(store.aget_tools()) + for sync_tool, async_tool in zip(sync_tools, async_tools): + assert sync_tool.__doc__ == async_tool.__doc__ + assert sync_tool.__doc__ # never empty + + def test_tool_signatures_match_the_spec(self, store: EntityMemoryStore) -> None: + tools = {t.__name__: t for t in store.get_tools()} + assert list(inspect.signature(tools["remember_about"]).parameters) == [ + "entity", + "entity_type", + "description", + "facts", + "events", + "note", + ] + assert list(inspect.signature(tools["link_entities"]).parameters) == ["entity", "relation", "related_entity"] + assert list(inspect.signature(tools["search_entities"]).parameters) == ["query", "entity_type"] + assert list(inspect.signature(tools["forget"]).parameters) == ["entity", "fact"] + + def test_tools_disabled_when_configured_off(self, db: RecordingLearningDb) -> None: + store = EntityMemoryStore(config=EntityMemoryConfig(db=db, enable_agent_tools=False)) # type: ignore[arg-type] + assert store.get_tools() == [] + + +class TestRememberAbout: + def test_creates_entity_with_slugified_id(self, store: EntityMemoryStore) -> None: + message = store.remember_about(entity="Sarah Chen", entity_type="person", facts=["designs radar"]) + assert "person/sarah_chen" in message + entity = store.get(entity_id="sarah_chen", entity_type="person") + assert entity is not None + assert entity.name == "Sarah Chen" + assert [f["content"] for f in entity.facts] == ["designs radar"] + + def test_merges_into_existing_entity(self, store: EntityMemoryStore) -> None: + store.remember_about(entity="radar", entity_type="project", facts=["db: Postgres"]) + store.remember_about(entity="radar", entity_type="project", events=["shipped v1"]) + entity = store.get(entity_id="radar", entity_type="project") + assert entity is not None + assert len(entity.facts) == 1 + assert len(entity.events) == 1 + + def test_note_pointer_round_trips(self, store: EntityMemoryStore) -> None: + store.remember_about(entity="radar", entity_type="project", note="notes/radar.md") + entity = store.get(entity_id="radar", entity_type="project") + assert entity is not None + assert entity.properties["note"] == "notes/radar.md" + # And it shows in search results + result = store.search_entities(query="radar") + assert "note: notes/radar.md" in result + + async def test_async_remember_about(self, store: EntityMemoryStore) -> None: + message = await store.aremember_about(entity="Acme Corp", entity_type="company", facts=["fintech"]) + assert "company/acme_corp" in message + entity = await store.aget(entity_id="acme_corp", entity_type="company") + assert entity is not None + + def test_user_namespace_requires_user_id(self, db: RecordingLearningDb) -> None: + store = EntityMemoryStore(config=EntityMemoryConfig(db=db, namespace="user")) # type: ignore[arg-type] + message = store.remember_about(entity="radar", entity_type="project") + assert "user_id" in message + assert db.rows == {} + + def test_user_namespace_fails_closed_on_reads_and_forget(self, db: RecordingLearningDb) -> None: + # Alice's private entity must not be readable or archivable without a user_id. + store = EntityMemoryStore(config=EntityMemoryConfig(db=db, namespace="user")) # type: ignore[arg-type] + store.remember_about(entity="secret project", entity_type="project", facts=["acme deal"], user_id="alice") + + assert "user_id" in store.search_entities(query="acme") + assert store.search(query="acme") == [] + assert store.list_entities() == [] + assert "user_id" in store.forget(entity="secret project") + entity = store.get(entity_id="secret_project", entity_type="project", user_id="alice") + assert entity is not None and entity.archived_at is None + + def test_blank_entity_name_is_rejected(self, store: EntityMemoryStore, db: RecordingLearningDb) -> None: + assert "Entity name is required" in store.remember_about(entity=" ", entity_type="person") + assert db.rows == {} + + +class TestResolution: + def test_name_variants_merge_into_one_entity(self, store: EntityMemoryStore, db: RecordingLearningDb) -> None: + store.remember_about(entity="Sarah Chen", entity_type="person", facts=["designs radar"]) + store.remember_about(entity="sarah chen", entity_type="person", facts=["prefers async"]) + assert len(db.rows) == 1 + entity = store.get(entity_id="sarah_chen", entity_type="person") + assert entity is not None + assert [f["content"] for f in entity.facts] == ["designs radar", "prefers async"] + + def test_type_drift_merges(self, store: EntityMemoryStore, db: RecordingLearningDb) -> None: + store.remember_about(entity="Sarah Chen", entity_type="person", facts=["designs radar"]) + store.remember_about(entity="Sarah Chen", entity_type="people", facts=["prefers async"]) + store.remember_about(entity="Sarah Chen", entity_type="Person") + assert len(db.rows) == 1 + entity = store.get(entity_id="sarah_chen", entity_type="person") + assert entity is not None + assert len(entity.facts) == 2 + + def test_type_is_normalized_on_create(self, store: EntityMemoryStore) -> None: + store.remember_about(entity="Acme", entity_type="Companies") + assert store.get(entity_id="acme", entity_type="company") is not None + + def test_resolves_by_exact_name_when_id_differs(self, store: EntityMemoryStore, db: RecordingLearningDb) -> None: + # A row whose entity_id does not derive from its name (external writer). + db.upsert_learning( + id="entity_global_person_sc_001", + learning_type="entity_memory", + entity_id="sc_001", + entity_type="person", + namespace="global", + content={"entity_id": "sc_001", "entity_type": "person", "name": "Sarah Chen", "facts": []}, + ) + store.remember_about(entity="Sarah Chen", entity_type="person", facts=["works at Acme"]) + assert len(db.rows) == 1 # merged, not duplicated + entity = store.get(entity_id="sc_001", entity_type="person") + assert entity is not None + assert [f["content"] for f in entity.facts] == ["works at Acme"] + + def test_resolves_by_alias(self, store: EntityMemoryStore, db: RecordingLearningDb) -> None: + db.upsert_learning( + id="entity_global_project_radar", + learning_type="entity_memory", + entity_id="radar", + entity_type="project", + namespace="global", + content={ + "entity_id": "radar", + "entity_type": "project", + "name": "radar", + "aliases": ["The Radar Initiative"], + "facts": [], + }, + ) + store.remember_about(entity="the radar initiative", entity_type="project", facts=["shipped v1"]) + assert len(db.rows) == 1 + entity = store.get(entity_id="radar", entity_type="project") + assert entity is not None + assert [f["content"] for f in entity.facts] == ["shipped v1"] + + def test_incoming_name_variant_recorded_as_alias(self, store: EntityMemoryStore, db: RecordingLearningDb) -> None: + db.upsert_learning( + id="entity_global_person_sc_001", + learning_type="entity_memory", + entity_id="sc_001", + entity_type="person", + namespace="global", + content={"entity_id": "sc_001", "entity_type": "person", "name": "Sarah Chen", "facts": []}, + ) + # Resolves via name match; a genuinely different surface form would be an alias, + # but the same normalized name must NOT be duplicated in. + store.remember_about(entity="SARAH CHEN", entity_type="person") + entity = store.get(entity_id="sc_001", entity_type="person") + assert entity is not None + assert entity.aliases == [] + + def test_unknown_type_upgraded_by_later_remember(self, store: EntityMemoryStore, db: RecordingLearningDb) -> None: + store.remember_about(entity="radar", entity_type="project") + store.link_entities(entity="radar", relation="uses", related_entity="Postgres") + assert store.get(entity_id="postgres", entity_type="unknown") is not None + + store.remember_about(entity="Postgres", entity_type="system", facts=["v16 in prod"]) + + upgraded = store.get(entity_id="postgres", entity_type="system") + assert upgraded is not None + assert [f["content"] for f in upgraded.facts] == ["v16 in prod"] + assert upgraded.relationships # the edge written while unknown survives + # The old-typed row is gone + assert store.get(entity_id="postgres", entity_type="unknown") is None + + def test_explicit_aliases_write_and_resolve(self, store: EntityMemoryStore, db: RecordingLearningDb) -> None: + # The alias path on remember_about: an explicit alias is written, then + # a later write arriving under that alias resolves to the same entity. + store.remember_about(entity="Sarah Chen", entity_type="person", aliases=["SC", "Sarah from Design"]) + entity = store.get(entity_id="sarah_chen", entity_type="person") + assert entity is not None + assert entity.aliases == ["SC", "Sarah from Design"] + + store.remember_about(entity="Sarah from Design", entity_type="person", facts=["owns radar UX"]) + assert len(db.rows) == 1 # resolved via the alias, not duplicated + entity = store.get(entity_id="sarah_chen", entity_type="person") + assert entity is not None + assert [f["content"] for f in entity.facts] == ["owns radar UX"] + + def test_no_fuzzy_merge(self, store: EntityMemoryStore, db: RecordingLearningDb) -> None: + # "Sarah" must NOT merge into "Sarah Chen" - a wrong merge has no unmerge. + store.remember_about(entity="Sarah Chen", entity_type="person") + store.remember_about(entity="Sarah", entity_type="person") + assert len(db.rows) == 2 + + def test_archived_entities_do_not_push_live_ones_out_of_the_directory(self, db: RecordingLearningDb) -> None: + """Archived rows are dropped after the fetch, so enough recent archives + used to shorten the directory - while the block still told the model + the directory was the full index and an absent entity was not known.""" + store = EntityMemoryStore( + config=EntityMemoryConfig(db=db, max_entities_in_directory=5) # type: ignore[arg-type] + ) + for i in range(5): + store.remember_about(entity=f"live{i}", entity_type="project", facts=["x"]) + for i in range(20): + store.remember_about(entity=f"dead{i}", entity_type="project", facts=["x"]) + store.forget(entity=f"dead{i}") + + data = store.recall(message="where do things stand?") + assert data is not None + assert sorted(e.entity_id for e in data["directory"]) == ["live0", "live1", "live2", "live3", "live4"] + + async def test_same_turn_tool_calls_do_not_overwrite_each_other(self, db: RecordingLearningDb) -> None: + """An assistant turn's tool calls run concurrently (models/base.py + gathers them), and every entity write is a read-modify-write over the + whole row. Whatever suspends in the middle - the judge's provider call, + or any await on an async db - must not cost a sibling its write.""" + import asyncio + + store = EntityMemoryStore( + config=EntityMemoryConfig(db=db, model=_SleepyJudge()) # type: ignore[arg-type] + ) + await store.aremember_about(entity="radar", entity_type="project", facts=["db: Postgres"]) + + await asyncio.gather( + store.aremember_about(entity="radar", entity_type="project", facts=["owner: Sarah Chen"]), + store.aremember_about(entity="radar", entity_type="project", events=["shipped v1"]), + store.alink_entities(entity="radar", relation="uses", related_entity="Postgres"), + ) + + entity = await store.aget(entity_id="radar", entity_type="project") + assert entity is not None + assert [f["content"] for f in entity.facts] == ["db: Postgres", "owner: Sarah Chen"] + assert [e["content"] for e in entity.events] == ["shipped v1"] + assert [r["entity_id"] for r in entity.relationships] == ["postgres"] + + async def test_same_turn_calls_do_not_split_a_new_entity(self, db: RecordingLearningDb) -> None: + """remember_about mints person/sarah_chen while link_entities mints the + unknown/ placeholder; concurrently, both used to survive as two rows. + + The fake db never suspends, so the judge's await is what opens the + window - the store needs a model for this to discriminate. + """ + import asyncio + + store = EntityMemoryStore( + config=EntityMemoryConfig(db=db, model=_SleepyJudge()) # type: ignore[arg-type] + ) + await store.aremember_about(entity="Sarah Chen", entity_type="person", facts=["joined radar"]) + await asyncio.gather( + store.aremember_about(entity="Sarah Chen", entity_type="person", facts=["leads radar"]), + store.alink_entities(entity="Sarah Chen", relation="works_on", related_entity="apollo"), + ) + assert sorted(r["entity_id"] for r in db.rows.values()) == ["apollo", "sarah_chen"] + entity = await store.aget(entity_id="sarah_chen", entity_type="person") + assert entity is not None + assert [f["content"] for f in entity.facts] == ["joined radar", "leads radar"] + assert [r["entity_id"] for r in entity.relationships] == ["apollo"] + + async def test_same_fact_twice_in_one_turn_is_stored_once(self, db: RecordingLearningDb) -> None: + """The duplicate check has to run against the row the merge lands on. + + Judged against the pre-lock snapshot, two siblings carrying the same + sentence both called it novel and both appended it. + """ + import asyncio + + store = EntityMemoryStore( + config=EntityMemoryConfig(db=db, model=_SleepyJudge()) # type: ignore[arg-type] + ) + await store.aremember_about(entity="Sarah Chen", entity_type="person", facts=["joined radar in March"]) + await asyncio.gather( + store.aremember_about(entity="Sarah Chen", entity_type="person", facts=["Sarah now leads radar"]), + store.aremember_about(entity="Sarah Chen", entity_type="person", facts=["Sarah now leads radar"]), + ) + entity = await store.aget(entity_id="sarah_chen", entity_type="person") + assert entity is not None + assert [f["content"] for f in entity.facts] == ["joined radar in March", "Sarah now leads radar"] + + def test_blank_description_and_note_do_not_wipe_stored_values(self, store: EntityMemoryStore) -> None: + # The model fills unused optional arguments with ""; the tool surface + # has no way to CLEAR a description, so "" cannot mean "clear it". + store.remember_about( + entity="radar", entity_type="project", description="The ingest rewrite", note="notes/radar.md" + ) + store.remember_about(entity="radar", entity_type="project", description="", note="", facts=["shipped"]) + entity = store.get(entity_id="radar", entity_type="project") + assert entity is not None + assert entity.description == "The ingest rewrite" + assert (entity.properties or {}).get("note") == "notes/radar.md" + + def test_write_lock_outside_a_running_loop(self, store: EntityMemoryStore) -> None: + # The loop-keyed cache is weak, and None is not weak-referenceable. + assert store._write_lock() is not None + + def test_accented_names_do_not_collide(self, store: EntityMemoryStore, db: RecordingLearningDb) -> None: + # Muller and Moller are two people; a slug that drops the accented + # letter made them one, and there is no unmerge. + store.remember_about(entity="Anna Müller", entity_type="person", facts=["office: Berlin"]) + store.remember_about(entity="Anna Möller", entity_type="person", facts=["office: Hamburg"]) + assert len(db.rows) == 2 + + def test_accent_folding_merges_the_same_person(self, store: EntityMemoryStore, db: RecordingLearningDb) -> None: + store.remember_about(entity="Sofía Muñoz", entity_type="person", facts=["office: Madrid"]) + store.remember_about(entity="Sofia Munoz", entity_type="person", facts=["team: platform"]) + assert len(db.rows) == 1 + entity = store.get(entity_id="sofia_munoz", entity_type="person") + assert entity is not None + assert len(entity.facts) == 2 + + def test_non_latin_names_keep_their_own_id(self, store: EntityMemoryStore, db: RecordingLearningDb) -> None: + # Nothing to fold: the lowered name is the id, and two of them stay two. + store.remember_about(entity="李明", entity_type="person", facts=["office: Shanghai"]) + store.remember_about(entity="Дмитрий", entity_type="person", facts=["office: Riga"]) + assert len(db.rows) == 2 + + def test_same_name_different_canonical_types_stay_separate( + self, store: EntityMemoryStore, db: RecordingLearningDb + ) -> None: + # The project Harbor and the company Harbor are different things. + store.remember_about(entity="Harbor", entity_type="project", facts=["db: Postgres"]) + store.remember_about(entity="Harbor", entity_type="company", facts=["HQ: Lisbon"]) + assert len(db.rows) == 2 + project = store.get(entity_id="harbor", entity_type="project") + company = store.get(entity_id="harbor", entity_type="company") + assert project is not None and company is not None + assert [f["content"] for f in project.facts] == ["db: Postgres"] + assert [f["content"] for f in company.facts] == ["HQ: Lisbon"] + + def test_name_match_across_canonical_types_stays_separate( + self, store: EntityMemoryStore, db: RecordingLearningDb + ) -> None: + # The name rung must not cross canonical types either. + db.upsert_learning( + id="entity_global_project_rdr_001", + learning_type="entity_memory", + entity_id="rdr_001", + entity_type="project", + namespace="global", + content={"entity_id": "rdr_001", "entity_type": "project", "name": "Radar", "facts": []}, + ) + store.remember_about(entity="Radar", entity_type="person", facts=["role: staff engineer"]) + assert len(db.rows) == 2 + + def test_an_ambiguous_name_changes_nothing(self, store: EntityMemoryStore) -> None: + """link_entities and forget carry no entity_type, so a shared name + would silently pick one row and strand the other.""" + store.remember_about(entity="Harbor", entity_type="project", facts=["the ingest rewrite"]) + store.remember_about(entity="Harbor", entity_type="company", facts=["HQ Lisbon"]) + + for message in ( + store.link_entities(entity="Harbor", relation="uses", related_entity="Postgres"), + store.forget(entity="Harbor"), + ): + assert "matches more than one entity" in message + assert "project/Harbor" in message and "company/Harbor" in message + + project = store.get(entity_id="harbor", entity_type="project") + company = store.get(entity_id="harbor", entity_type="company") + assert project is not None and company is not None + assert project.relationships == [] and not getattr(company, "archived_at", None) + + def test_a_qualified_name_reaches_either_entity(self, store: EntityMemoryStore) -> None: + store.remember_about(entity="Harbor", entity_type="project", facts=["the ingest rewrite"]) + store.remember_about(entity="Harbor", entity_type="company", facts=["HQ Lisbon"]) + + assert "project/harbor" in store.link_entities( + entity="project/Harbor", relation="uses", related_entity="Postgres" + ) + assert "Archived company/harbor" in store.forget(entity="company/Harbor") + + project = store.get(entity_id="harbor", entity_type="project") + company = store.get(entity_id="harbor", entity_type="company") + assert project is not None and company is not None + assert [r["entity_id"] for r in project.relationships] == ["postgres"] + assert getattr(company, "archived_at", None) + assert not getattr(project, "archived_at", None) + + def test_a_slash_in_a_name_is_part_of_the_name(self, store: EntityMemoryStore, db: RecordingLearningDb) -> None: + """A prefix only qualifies when it names a stored type. + + Reading every slash as "type/name" sent resolution to a key the write + path never used, so the second write created nothing and clobbered the + first entity's facts. + """ + store.remember_about(entity="AC/DC", entity_type="company", facts=["a band"]) + store.remember_about(entity="AC/DC", entity_type="company", facts=["still a band"]) + assert len(db.rows) == 1 + entity = store.get(entity_id="ac_dc", entity_type="company") + assert entity is not None + assert [f["content"] for f in entity.facts] == ["a band", "still a band"] + + async def test_async_ambiguous_name_changes_nothing(self, store: EntityMemoryStore) -> None: + await store.aremember_about(entity="Harbor", entity_type="project") + await store.aremember_about(entity="Harbor", entity_type="company") + assert "matches more than one entity" in await store.aforget(entity="Harbor") + assert "matches more than one entity" in await store.alink_entities( + entity="Harbor", relation="uses", related_entity="Postgres" + ) + + def test_a_stale_relationship_can_be_retired(self, store: EntityMemoryStore) -> None: + """A corrected link had no retirement path: stating the new one left + both edges rendering, undated, forever.""" + store.remember_about(entity="quill", entity_type="project") + store.link_entities(entity="quill", relation="written_in", related_entity="Rust") + store.link_entities(entity="quill", relation="written_in", related_entity="Go") + + message = store.forget(entity="quill", fact="written_in -> Rust") + assert "Removed relationship" in message + quill = store.get(entity_id="quill", entity_type="project") + assert quill is not None + assert [(r["relation"], r["entity_id"]) for r in quill.relationships] == [("written_in", "go")] + # the reciprocal edge goes too, so the graph does not go one-sided + rust = store.get(entity_id="rust", entity_type="unknown") + assert rust is not None + assert rust.relationships == [] + + def test_an_ambiguous_relationship_retires_nothing(self, store: EntityMemoryStore) -> None: + store.remember_about(entity="quill", entity_type="project") + store.link_entities(entity="quill", relation="written_in", related_entity="Rust") + store.link_entities(entity="quill", relation="written_in", related_entity="Go") + + message = store.forget(entity="quill", fact="written_in") + assert "Multiple relationships" in message + quill = store.get(entity_id="quill", entity_type="project") + assert quill is not None + assert len(quill.relationships) == 2 + + async def test_async_relationship_retirement(self, store: EntityMemoryStore) -> None: + await store.aremember_about(entity="tom", entity_type="person") + await store.alink_entities(entity="tom", relation="works_on", related_entity="radar") + assert "Removed relationship" in await store.aforget(entity="tom", fact="works_on -> radar") + tom = await store.aget(entity_id="tom", entity_type="person") + assert tom is not None + assert tom.relationships == [] + + def test_mixed_script_names_do_not_merge(self, store: EntityMemoryStore, db: RecordingLearningDb) -> None: + # Dropping the non-Latin half keyed "李 Ming" to `ming`. + store.remember_about(entity="李 Ming", entity_type="person", facts=["office: Shanghai"]) + store.remember_about(entity="Ming", entity_type="person", facts=["office: London"]) + assert len(db.rows) == 2 + + def test_two_named_types_are_two_entities(self, store: EntityMemoryStore, db: RecordingLearningDb) -> None: + """Only the ``unknown`` placeholder merges across types. + + Restricting the split to a canonical list left the types models + actually coin - team, framework, service - merging into whatever row + happened to share the name, which has no unmerge. + """ + store.remember_about(entity="Atlas", entity_type="system", facts=["the ingest system"]) + message = store.remember_about(entity="Atlas", entity_type="team", facts=["the platform team"]) + assert len(db.rows) == 2 + # ...and the write says the sibling exists, so the model can correct itself. + assert "system/atlas already exists under this name" in message + + def test_free_form_type_does_not_swallow_a_named_one( + self, store: EntityMemoryStore, db: RecordingLearningDb + ) -> None: + store.remember_about(entity="Sarah Chen", entity_type="engineer", facts=["designs radar"]) + store.remember_about(entity="Sarah Chen", entity_type="person", facts=["prefers async"]) + assert len(db.rows) == 2 + + def test_link_placeholder_still_upgrades_to_a_real_type( + self, store: EntityMemoryStore, db: RecordingLearningDb + ) -> None: + # link_entities mints entity_type="unknown"; describing it later must merge. + store.link_entities(entity="harbor", relation="uses", related_entity="Postgres") + store.remember_about(entity="Postgres", entity_type="system", facts=["v16"]) + assert sorted(r["entity_type"] for r in db.rows.values()) == ["system", "unknown"] + + async def test_async_same_name_different_canonical_types_stay_separate( + self, store: EntityMemoryStore, db: RecordingLearningDb + ) -> None: + await store.aremember_about(entity="Harbor", entity_type="project", facts=["db: Postgres"]) + await store.aremember_about(entity="Harbor", entity_type="company", facts=["HQ: Lisbon"]) + assert len(db.rows) == 2 + + +class TestLinkEntities: + def test_edge_written_on_both_rows_with_far_end_type(self, store: EntityMemoryStore) -> None: + store.remember_about(entity="Sarah Chen", entity_type="person") + store.remember_about(entity="radar", entity_type="project") + message = store.link_entities(entity="Sarah Chen", relation="works_on", related_entity="radar") + assert "person/sarah_chen" in message and "project/radar" in message + + sarah = store.get(entity_id="sarah_chen", entity_type="person") + radar = store.get(entity_id="radar", entity_type="project") + assert sarah is not None and radar is not None + + out_edge = sarah.relationships[0] + assert out_edge["entity_id"] == "radar" + assert out_edge["entity_type"] == "project" + assert out_edge["relation"] == "works_on" + assert out_edge["direction"] == "outgoing" + + in_edge = radar.relationships[0] + assert in_edge["entity_id"] == "sarah_chen" + assert in_edge["entity_type"] == "person" + assert in_edge["relation"] == "works_on" + assert in_edge["direction"] == "incoming" + + def test_unresolved_end_creates_minimal_unknown_entity(self, store: EntityMemoryStore) -> None: + store.remember_about(entity="radar", entity_type="project") + store.link_entities(entity="radar", relation="uses", related_entity="Postgres") + postgres = store.get(entity_id="postgres", entity_type="unknown") + assert postgres is not None + assert postgres.relationships[0]["direction"] == "incoming" + + async def test_async_link_entities(self, store: EntityMemoryStore) -> None: + await store.aremember_about(entity="radar", entity_type="project") + message = await store.alink_entities(entity="radar", relation="owned_by", related_entity="Acme") + assert "Linked" in message + + def test_self_link_is_rejected(self, store: EntityMemoryStore) -> None: + store.remember_about(entity="radar", entity_type="project") + message = store.link_entities(entity="radar", relation="relates_to", related_entity="Radar") + assert "itself" in message + entity = store.get(entity_id="radar", entity_type="project") + assert entity is not None and entity.relationships == [] + + +class TestSearchEntities: + def test_query_matches_fact_content(self, store: EntityMemoryStore) -> None: + store.remember_about(entity="Acme", entity_type="company", facts=["uses PostgreSQL"]) + result = store.search_entities(query="postgresql") + assert "Acme" in result + + def test_no_query_lists_by_recency(self, store: EntityMemoryStore) -> None: + store.remember_about(entity="older", entity_type="project") + store.remember_about(entity="newer", entity_type="project") + result = store.search_entities() + assert result.index("newer") < result.index("older") + + def test_no_match_reports_scan_scope(self, store: EntityMemoryStore) -> None: + store.remember_about(entity="radar", entity_type="project") + result = store.search_entities(query="nonexistent") + assert "No entities matching" in result + assert "namespace 'global'" in result + + def test_truncation_marker_on_many_facts(self, store: EntityMemoryStore) -> None: + store.remember_about(entity="radar", entity_type="project", facts=[f"fact number {i}" for i in range(19)]) + result = store.search_entities(query="radar") + assert "(newest 6 of 19 facts)" in result + # The NEWEST facts are shown, not the oldest + assert "fact number 18" in result and "fact number 0" not in result + + async def test_async_search_entities(self, store: EntityMemoryStore) -> None: + await store.aremember_about(entity="radar", entity_type="project") + result = await store.asearch_entities(query="radar") + assert "radar" in result + + def test_blank_entity_type_is_no_filter(self, store: EntityMemoryStore) -> None: + # A strict tool schema has no "absent", so models send "" for the + # optional argument they did not mean to use. Filtering on it would + # answer "not found" while holding the entity. + store.remember_about(entity="Meridian", entity_type="project", facts=["status: in migration"]) + assert "Meridian" in store.search_entities(query="Meridian", entity_type="") + assert "Meridian" in store.search_entities(query="Meridian", entity_type=" ") + + def test_blank_query_still_lists(self, store: EntityMemoryStore) -> None: + store.remember_about(entity="Meridian", entity_type="project") + assert "Meridian" in store.search_entities(query=" ", entity_type="") + + async def test_async_blank_entity_type_is_no_filter(self, store: EntityMemoryStore) -> None: + await store.aremember_about(entity="Meridian", entity_type="project", facts=["status: in migration"]) + assert "Meridian" in await store.asearch_entities(query="Meridian", entity_type="") + assert "Meridian" in await store.asearch_entities(query=" ", entity_type="") + + +class TestForget: + def test_archive_excluded_from_recall_but_searchable(self, store: EntityMemoryStore) -> None: + store.remember_about(entity="radar", entity_type="project", facts=["shipped"]) + message = store.forget(entity="radar") + assert "Archived project/radar" in message + + # Excluded from recall: neither expanded nor in the directory + recalled = store.recall(entity_id="radar", entity_type="project") + assert recalled is not None + assert recalled["directory"] == [] and recalled["entities"] == [] + # Still reachable via explicit search, marked archived + result = store.search_entities(query="radar") + assert "(archived)" in result + # Excluded from the listing path (recall-adjacent surfaces exclude archived) + assert store.list_entities() == [] + + def test_remember_revives_archived_entity(self, store: EntityMemoryStore) -> None: + store.remember_about(entity="radar", entity_type="project") + store.forget(entity="radar") + message = store.remember_about(entity="radar", entity_type="project", facts=["back on"]) + assert "revived" in message + recalled = store.recall(entity_id="radar", entity_type="project") + assert recalled is not None and len(recalled["entities"]) == 1 + + def test_forget_unknown_entity(self, store: EntityMemoryStore) -> None: + assert "No entity found" in store.forget(entity="ghost") + + def test_exact_fact_match_retires(self, store: EntityMemoryStore) -> None: + store.remember_about(entity="radar", entity_type="project", facts=["blocked on review", "db: Postgres"]) + message = store.forget(entity="radar", fact="Blocked on Review") + assert "Retired fact" in message + entity = store.get(entity_id="radar", entity_type="project") + assert entity is not None + live = entity.live_facts() + assert [f["content"] for f in live] == ["db: Postgres"] + retired = [f for f in entity.facts if f.get("superseded_at")] + assert len(retired) == 1 + assert retired[0]["superseded_by"] == "forgotten" + + def test_single_containment_match_retires(self, store: EntityMemoryStore) -> None: + store.remember_about(entity="radar", entity_type="project", facts=["blocked on security review"]) + message = store.forget(entity="radar", fact="security review") + assert "Retired fact" in message + + def test_multiple_matches_retire_nothing(self, store: EntityMemoryStore) -> None: + store.remember_about(entity="radar", entity_type="project", facts=["review is pending", "review was requested"]) + message = store.forget(entity="radar", fact="review") + assert "Multiple facts" in message + assert "review is pending" in message and "review was requested" in message + entity = store.get(entity_id="radar", entity_type="project") + assert entity is not None + assert len(entity.live_facts()) == 2 + + def test_zero_matches_returns_live_facts(self, store: EntityMemoryStore) -> None: + store.remember_about(entity="radar", entity_type="project", facts=["db: Postgres"]) + message = store.forget(entity="radar", fact="something else entirely") + assert "No matching fact on project/radar" in message + assert "db: Postgres" in message + + async def test_async_forget_archives(self, store: EntityMemoryStore) -> None: + await store.aremember_about(entity="radar", entity_type="project") + message = await store.aforget(entity="radar") + assert "Archived" in message + + +class TestSearchRouting: + def test_search_routes_through_search_learnings(self, db: RecordingLearningDb) -> None: + calls: List[Dict[str, Any]] = [] + + class SpyDb(RecordingLearningDb): + def search_learnings(self, query: str, **kwargs: Any) -> List[Dict[str, Any]]: + calls.append({"query": query, **kwargs}) + return super().search_learnings(query, **kwargs) + + spy = SpyDb() + store = EntityMemoryStore(config=EntityMemoryConfig(db=spy)) # type: ignore[arg-type] + store.remember_about(entity="radar", entity_type="project", facts=["db: Postgres"]) + calls.clear() + results = store.search(query="postgres") + assert len(results) == 1 + assert calls and calls[-1]["query"] == "postgres" + assert calls[-1]["learning_type"] == "entity_memory" + assert calls[-1]["namespace"] == "global" + + def test_search_crosses_slug_boundary_via_store(self, store: EntityMemoryStore) -> None: + store.remember_about(entity="Sarah Chen", entity_type="person", facts=["designs radar"]) + results = store.search(query="sarah chen") + assert [e.entity_id for e in results] == ["sarah_chen"] + + def test_search_falls_back_on_not_implemented(self, caplog: pytest.LogCaptureFixture) -> None: + class NoSearchDb(RecordingLearningDb): + def search_learnings(self, query: str, **kwargs: Any) -> List[Dict[str, Any]]: + raise NotImplementedError + + db = NoSearchDb() + store = EntityMemoryStore(config=EntityMemoryConfig(db=db)) # type: ignore[arg-type] + store.remember_about(entity="radar", entity_type="project", facts=["db: Postgres"]) + + import logging + + with caplog.at_level(logging.WARNING): + results = store.search(query="postgres") + store.search(query="postgres") + + assert [e.entity_id for e in results] == ["radar"] + degraded = [r for r in caplog.records if "no search_learnings implementation" in r.getMessage()] + assert len(degraded) == 1 # logged once, not per call + + def test_search_fails_loudly_when_backend_errors(self) -> None: + class BrokenSearchDb(RecordingLearningDb): + def search_learnings(self, query: str, **kwargs: Any) -> List[Dict[str, Any]]: + raise RuntimeError("dialect error") + + store = EntityMemoryStore(config=EntityMemoryConfig(db=BrokenSearchDb())) # type: ignore[arg-type] + with pytest.raises(RuntimeError, match="dialect error"): + store.search(query="anything") + + async def test_asearch_routes_and_falls_back(self, caplog: pytest.LogCaptureFixture) -> None: + class NoSearchDb(RecordingLearningDb): + def search_learnings(self, query: str, **kwargs: Any) -> List[Dict[str, Any]]: + raise NotImplementedError + + db = NoSearchDb() + store = EntityMemoryStore(config=EntityMemoryConfig(db=db)) # type: ignore[arg-type] + await store.aremember_about(entity="radar", entity_type="project", facts=["db: Postgres"]) + results = await store.asearch(query="postgres") + assert [e.entity_id for e in results] == ["radar"] + + def test_json_key_names_do_not_match(self, store: EntityMemoryStore) -> None: + # The db-side ILIKE sees the whole serialized document; the store must + # verify hits against values only, or "facts"/"name" match every row. + store.remember_about(entity="radar", entity_type="project", facts=["db: Postgres"]) + for key_name in ("facts", "name", "entity_id", "properties", "relationships"): + assert store.search(query=key_name) == [], key_name + # ...but a value containing such a word still matches + store.remember_about(entity="naming service", entity_type="system") + assert [e.entity_id for e in store.search(query="naming")] == ["naming_service"] + + def test_attribute_error_inside_backend_is_not_masked(self) -> None: + # A backend that HAS search_learnings but raises AttributeError inside + # it is a real bug, not a missing implementation - it must propagate. + class BuggyDb(RecordingLearningDb): + def search_learnings(self, query: str, **kwargs: Any) -> List[Dict[str, Any]]: + raise AttributeError("'NoneType' object has no attribute 'client'") + + store = EntityMemoryStore(config=EntityMemoryConfig(db=BuggyDb())) # type: ignore[arg-type] + with pytest.raises(AttributeError, match="client"): + store.search(query="anything") + + def test_archived_rows_do_not_crowd_out_live_hits(self, store: EntityMemoryStore) -> None: + for i in range(3): + store.remember_about(entity=f"live topic {i}", entity_type="project", facts=["about postgres"]) + for i in range(15): + store.remember_about(entity=f"dead topic {i}", entity_type="project", facts=["about postgres"]) + store.forget(entity=f"dead topic {i}") + + results = store.search(query="postgres", limit=10) + assert len(results) == 3 + assert all(e.archived_at is None for e in results) + + async def test_asearch_awaits_async_base_db(self) -> None: + from agno.db.base import AsyncBaseDb + + inner = RecordingLearningDb() + calls: List[str] = [] + + class FakeAsyncDb(AsyncBaseDb): + def __init__(self) -> None: + pass + + async def get_learning(self, **kwargs: Any) -> Any: + return inner.get_learning(**kwargs) + + async def upsert_learning(self, **kwargs: Any) -> None: + inner.upsert_learning(**kwargs) + + async def get_learnings(self, **kwargs: Any) -> Any: + return inner.get_learnings(**kwargs) + + async def search_learnings(self, query: str, **kwargs: Any) -> Any: + calls.append(query) + return inner.search_learnings(query, **kwargs) + + FakeAsyncDb.__abstractmethods__ = frozenset() # type: ignore[attr-defined] + store = EntityMemoryStore(config=EntityMemoryConfig(db=FakeAsyncDb())) + await store.aremember_about(entity="radar", entity_type="project", facts=["db: Postgres"]) + results = await store.asearch(query="postgres") + assert [e.entity_id for e in results] == ["radar"] + assert "postgres" in calls # the awaited AsyncBaseDb branch really ran + + def test_sync_call_with_async_db_refuses_instead_of_type_error(self) -> None: + from agno.db.base import AsyncBaseDb + + class FakeAsyncDb(AsyncBaseDb): + def __init__(self) -> None: + pass + + FakeAsyncDb.__abstractmethods__ = frozenset() # type: ignore[attr-defined] + store = EntityMemoryStore(config=EntityMemoryConfig(db=FakeAsyncDb())) + assert store.search(query="anything") == [] + assert store.list_entities() == [] + assert "Failed" in store.remember_about(entity="radar", entity_type="project") + + def test_search_finds_match_outside_recent_window_with_sqlite(self, tmp_path) -> None: + from agno.db.sqlite import SqliteDb + + sqlite_db = SqliteDb(db_file=str(tmp_path / "entities.db")) + store = EntityMemoryStore(config=EntityMemoryConfig(db=sqlite_db)) + + store.remember_about(entity="needle", entity_type="project", facts=["the rare zanzibar detail"]) + for i in range(60): + sqlite_db.upsert_learning( + id=f"entity_global_project_filler_{i}", + learning_type="entity_memory", + entity_id=f"filler_{i}", + entity_type="project", + namespace="global", + content={"entity_id": f"filler_{i}", "entity_type": "project", "facts": []}, + ) + + results = store.search(query="zanzibar", limit=5) + assert [e.entity_id for e in results] == ["needle"] + + +class TestRenderingAndDirectory: + def test_rendering_is_bounded_and_honest(self, store: EntityMemoryStore) -> None: + store.remember_about(entity="radar", entity_type="project", facts=[f"fact number {i}" for i in range(200)]) + entity = store.get(entity_id="radar", entity_type="project") + assert entity is not None + + text = entity.get_context_text(max_facts=10, max_events=5) + assert "(newest 10 of 200 facts)" in text + assert text.count("fact number") == 10 + # The NEWEST facts render - showing the oldest slice would date-stamp + # stale state as current + assert "fact number 199" in text and "fact number 0 " not in text + # Facts render with as-of dates + assert "(as of 20" in text + + def test_events_render_last_n_with_marker(self, store: EntityMemoryStore) -> None: + store.remember_about(entity="radar", entity_type="project", events=[f"event {i}" for i in range(9)]) + entity = store.get(entity_id="radar", entity_type="project") + assert entity is not None + text = entity.get_context_text(max_facts=10, max_events=3) + assert "(last 3 of 9 events)" in text + assert "event 8" in text and "event 0" not in text + + def test_directory_always_in_recall_and_context(self, store: EntityMemoryStore) -> None: + store.remember_about(entity="radar", entity_type="project") + store.remember_about(entity="Sarah Chen", entity_type="person") + + # No keyed lookup: recall still returns the directory + recalled = store.recall() + assert recalled is not None + assert [e.entity_id for e in recalled["directory"]] == ["sarah_chen", "radar"] + assert recalled["entities"] == [] + + context = store.build_context(data=recalled) + assert "Entity directory" in context + assert "- Sarah Chen (person)" in context + assert "- radar (project)" in context + assert "not listed there is not known" in context + + def test_directory_excludes_archived(self, store: EntityMemoryStore) -> None: + store.remember_about(entity="radar", entity_type="project") + store.remember_about(entity="dead project", entity_type="project") + store.forget(entity="dead project") + + recalled = store.recall() + assert recalled is not None + assert [e.entity_id for e in recalled["directory"]] == ["radar"] + + def test_context_caps_expanded_entities(self, db: RecordingLearningDb) -> None: + store = EntityMemoryStore(config=EntityMemoryConfig(db=db, max_entities_in_context=2)) # type: ignore[arg-type] + entities = [] + for i in range(4): + store.remember_about(entity=f"proj {i}", entity_type="project", facts=[f"about {i}"]) + entities.append(store.get(entity_id=f"proj_{i}", entity_type="project")) + + context = store.build_context(data={"directory": entities, "entities": entities}) + # Directory lists all four; only two expand + assert context.count("- proj ") == 4 + assert context.count("Facts:") == 2 + + def test_empty_store_splits_guidance_and_data(self, store: EntityMemoryStore) -> None: + # Guidance (the four tools) lives in instructions(); build_context is data only. + context = store.build_context(data=store.recall()) + assert "No entities recorded yet" in context + assert "remember_about" not in context + + guidance = store.instructions() + for tool in ("remember_about", "link_entities", "search_entities", "forget"): + assert tool in guidance + + def test_one_hop_link_names_render_on_recall(self, store: EntityMemoryStore) -> None: + store.remember_about(entity="radar", entity_type="project") + store.remember_about(entity="Sarah Chen", entity_type="person") + store.link_entities(entity="Sarah Chen", relation="works_on", related_entity="radar") + + recalled = store.recall(entity_id="radar", entity_type="project") + assert recalled is not None + assert recalled["related_names"]["sarah_chen"] == "Sarah Chen" + + context = store.build_context(data=recalled) + # The edge renders the far end's display NAME, not its slug + assert "works_on <- Sarah Chen" in context + + def test_one_hop_names_resolve_beyond_the_directory(self, db: RecordingLearningDb) -> None: + # Far end outside the directory cap still resolves via a bounded keyed lookup. + store = EntityMemoryStore(config=EntityMemoryConfig(db=db, max_entities_in_directory=1)) # type: ignore[arg-type] + store.remember_about(entity="Sarah Chen", entity_type="person") + store.remember_about(entity="radar", entity_type="project") + store.link_entities(entity="radar", relation="designed_by", related_entity="Sarah Chen") + store.remember_about(entity="newest thing", entity_type="project") # crowds the 1-slot directory + + recalled = store.recall(entity_id="radar", entity_type="project") + assert recalled is not None + # recall fetches one row beyond the cap to detect truncation; the render caps at 1 + context = store.build_context(data=recalled) + assert context.count("\n- ") >= 1 + assert "most recently updated entities; more exist" in context + assert recalled["related_names"].get("sarah_chen") == "Sarah Chen" + + async def test_arecall_matches_recall(self, store: EntityMemoryStore) -> None: + await store.aremember_about(entity="radar", entity_type="project") + sync_result = store.recall() + async_result = await store.arecall() + assert sync_result is not None and async_result is not None + assert [e.entity_id for e in sync_result["directory"]] == [e.entity_id for e in async_result["directory"]] + + +class TestDataApi: + def test_hard_delete(self, store: EntityMemoryStore, db: RecordingLearningDb) -> None: + store.remember_about(entity="radar", entity_type="project") + assert store.delete(entity_id="radar", entity_type="project") is True + assert db.rows == {} + assert store.delete(entity_id="radar", entity_type="project") is False + + async def test_async_hard_delete(self, store: EntityMemoryStore) -> None: + await store.aremember_about(entity="radar", entity_type="project") + assert await store.adelete(entity_id="radar", entity_type="project") is True + + +class TestSearchWindowExhaustion: + """The db-side LIKE matches key names, so ordinary words match every row. + + Any finite ladder just moves the threshold at which a real, older match + becomes unreachable; the window has to run out with the backend. + """ + + def test_a_real_match_survives_hundreds_of_key_name_false_positives(self, tmp_path) -> None: + import time + + from agno.db.sqlite import SqliteDb + + db = SqliteDb(db_file=str(tmp_path / "crowd.db")) + store = EntityMemoryStore( + config=EntityMemoryConfig(db=db, mode=LearningMode.AGENTIC, namespace="global") # type: ignore[arg-type] + ) + store.remember_about( + entity="target", entity_type="project", facts=["the content marker lives here"], namespace="global" + ) + time.sleep(1.1) # so ORDER BY updated_at DESC really ranks the fillers first + for i in range(500): + store.remember_about(entity=f"other{i}", entity_type="project", facts=["x"], namespace="global") + + # "content" is a JSON key on every row and a value on exactly one + hits = store.search(query="content", namespace="global", limit=10) + assert [e.entity_id for e in hits] == ["target"] + + def test_mixed_case_unicode_in_a_fact_value_is_reachable(self, tmp_path) -> None: + from agno.db.sqlite import SqliteDb + + db = SqliteDb(db_file=str(tmp_path / "unicode.db")) + store = EntityMemoryStore( + config=EntityMemoryConfig(db=db, mode=LearningMode.AGENTIC, namespace="global") # type: ignore[arg-type] + ) + store.remember_about(entity="Alpha", entity_type="project", facts=["Ος"], namespace="global") + store.remember_about(entity="Cafe", entity_type="company", facts=["Café Noir"], namespace="global") + + for query in ("Ος", "ΟΣ", "ος", "οσ"): + assert [e.entity_id for e in store.search(query=query, namespace="global")] == ["alpha"], query + for query in ("café", "CAFÉ", "Café"): + assert [e.entity_id for e in store.search(query=query, namespace="global")] == ["cafe"], query + + +class TestRetirementReachesEveryStore: + """forget is the only retirement path the four-tool surface has. + + Each rung below had no route at all: a retracted event could not be + removed (the agent archived whole people to suppress one), and an edge + could not be named the way the context block renders it. + """ + + def test_an_event_can_be_retired(self, store: EntityMemoryStore) -> None: + store.remember_about( + entity="Tom", + entity_type="person", + facts=["works in infra"], + events=["heard he is leaving for a competitor", "shipped the migration"], + ) + assert "Retired event" in store.forget(entity="Tom", fact="heard he is leaving for a competitor") + + tom = store.get(entity_id="tom", entity_type="person") + assert tom is not None + assert [e["content"] for e in tom.events] == ["shipped the migration"] + # the facts survive, and the entity was not archived to suppress it + assert [f["content"] for f in tom.live_facts()] == ["works in infra"] + assert not getattr(tom, "archived_at", None) + + def test_an_ambiguous_event_retires_nothing(self, store: EntityMemoryStore) -> None: + store.remember_about(entity="Tom", entity_type="person", events=["review passed", "review passed again"]) + assert "Multiple events" in store.forget(entity="Tom", fact="review") + tom = store.get(entity_id="tom", entity_type="person") + assert tom is not None and len(tom.events) == 2 + + def test_an_edge_retires_by_the_name_the_model_is_shown(self, store: EntityMemoryStore) -> None: + # The block renders the far end's display name; the edge stores a slug. + for form in ("owned_by -> Sarah Chen", "owned_by -> sarah_chen", "Sarah Chen"): + store.remember_about(entity="Sarah Chen", entity_type="person") + store.remember_about(entity="radar", entity_type="project") + store.link_entities(entity="radar", relation="owned_by", related_entity="Sarah Chen") + assert "Removed relationship" in store.forget(entity="radar", fact=form), form + + async def test_async_event_retirement(self, store: EntityMemoryStore) -> None: + await store.aremember_about(entity="Tom", entity_type="person", events=["a rumour"]) + assert "Retired event" in await store.aforget(entity="Tom", fact="a rumour") + + +class TestQualifiedNameOnRememberAbout: + def test_remember_about_reads_the_form_the_other_tools_teach( + self, store: EntityMemoryStore, db: RecordingLearningDb + ) -> None: + """The refusals and docstrings teach "project/Harbor"; slugging it + whole minted project/project_harbor and stranded the correction.""" + store.remember_about(entity="Harbor", entity_type="project", facts=["the ingest rewrite"]) + store.remember_about(entity="Harbor", entity_type="company", facts=["HQ Lisbon"]) + + message = store.remember_about(entity="project/Harbor", entity_type="project", facts=["status: shipped"]) + assert "project/harbor" in message + assert len(db.rows) == 2 + project = store.get(entity_id="harbor", entity_type="project") + assert project is not None + assert [f["content"] for f in project.facts] == ["the ingest rewrite", "status: shipped"] + + async def test_async_remember_about_reads_the_qualified_form( + self, store: EntityMemoryStore, db: RecordingLearningDb + ) -> None: + await store.aremember_about(entity="Harbor", entity_type="project") + await store.aremember_about(entity="Harbor", entity_type="company") + message = await store.aremember_about(entity="project/Harbor", entity_type="project", facts=["shipped"]) + assert "project/harbor" in message + assert len(db.rows) == 2 + + +class TestEdgesAreIdempotent: + """Models re-assert links they already know. + + An appended duplicate could not be retired: forget listed byte-identical + candidates no wording could tell apart, while the far side's detach + removed every copy - so the graph went one-sided and stayed that way. + """ + + def test_relinking_does_not_double_the_edge(self, store: EntityMemoryStore) -> None: + store.remember_about(entity="radar", entity_type="project") + store.remember_about(entity="Sarah Chen", entity_type="person") + for _ in range(3): + store.link_entities(entity="radar", relation="owned_by", related_entity="Sarah Chen") + + radar = store.get(entity_id="radar", entity_type="project") + sarah = store.get(entity_id="sarah_chen", entity_type="person") + assert radar is not None and sarah is not None + assert len(radar.relationships) == 1 + assert len(sarah.relationships) == 1 + assert "Removed relationship" in store.forget(entity="radar", fact="owned_by -> Sarah Chen") + + def test_a_different_relation_to_the_same_target_stays_distinct(self, store: EntityMemoryStore) -> None: + store.remember_about(entity="radar", entity_type="project") + store.remember_about(entity="Sarah Chen", entity_type="person") + store.link_entities(entity="radar", relation="owned_by", related_entity="Sarah Chen") + store.link_entities(entity="radar", relation="designed_by", related_entity="Sarah Chen") + + radar = store.get(entity_id="radar", entity_type="project") + assert radar is not None and len(radar.relationships) == 2 + # ambiguous by target alone, and correctly refused + assert "Multiple relationships" in store.forget(entity="radar", fact="Sarah Chen") + + def test_legacy_duplicate_edges_can_still_be_retired(self, store: EntityMemoryStore) -> None: + # A row written before add_relationship deduped carries exact copies. + store.remember_about(entity="radar", entity_type="project") + store.remember_about(entity="Sarah Chen", entity_type="person") + store.link_entities(entity="radar", relation="owned_by", related_entity="Sarah Chen") + radar = store.get(entity_id="radar", entity_type="project") + assert radar is not None + radar.relationships.append(dict(radar.relationships[0], id="dupe1234")) + store._save_entity(entity=radar, user_id=None, agent_id=None, team_id=None, namespace="global") + + assert "Removed relationship" in store.forget(entity="radar", fact="owned_by -> Sarah Chen") + radar = store.get(entity_id="radar", entity_type="project") + assert radar is not None and radar.relationships == [] + + +class TestEdgeAndEventIdentity: + """The duplicate class @kausmeows found on edges, generalised. + + entity_type is part of an edge's identity - project/Harbor and + company/Harbor share a slug - and events had the same append-forever, + retire-one-of-many shape edges did. + """ + + def test_two_same_named_entities_keep_separate_edges(self, store: EntityMemoryStore) -> None: + for entity_type in ("project", "company"): + store.remember_about(entity="Harbor", entity_type=entity_type) + store.remember_about(entity="Alice", entity_type="person") + store.link_entities(entity="project/Harbor", relation="managed_by", related_entity="Alice") + store.link_entities(entity="company/Harbor", relation="managed_by", related_entity="Alice") + + alice = store.get(entity_id="alice", entity_type="person") + assert alice is not None + assert sorted(r["entity_type"] for r in alice.relationships) == ["company", "project"] + + # retiring one must not strip the other's reciprocal + store.forget(entity="project/Harbor", fact="managed_by -> Alice") + alice = store.get(entity_id="alice", entity_type="person") + assert alice is not None + assert [r["entity_type"] for r in alice.relationships] == ["company"] + + def test_forget_refuses_an_edge_naming_two_same_slug_far_ends(self, store: EntityMemoryStore) -> None: + # The mirror of the case above: the same-name pair is at the FAR end, so + # "works_on -> harbor" names two links. Retiring both while detaching one + # reciprocal left Alice pointing at neither and company/Harbor pointing + # back at her. + for entity_type in ("project", "company"): + store.remember_about(entity="Harbor", entity_type=entity_type) + store.remember_about(entity="Alice", entity_type="person") + store.link_entities(entity="Alice", relation="works_on", related_entity="project/Harbor") + store.link_entities(entity="Alice", relation="works_on", related_entity="company/Harbor") + + refusal = store.forget(entity="Alice", fact="works_on -> harbor") + assert "Multiple relationships" in refusal + assert "works_on -> project/harbor" in refusal and "works_on -> company/harbor" in refusal + alice = store.get(entity_id="alice", entity_type="person") + assert alice is not None and len(alice.relationships) == 2 + + # The qualified form the refusal asks for retires exactly one. + assert "Removed relationship" in store.forget(entity="Alice", fact="works_on -> project/harbor") + alice = store.get(entity_id="alice", entity_type="person") + assert alice is not None + assert [r["entity_type"] for r in alice.relationships] == ["company"] + project_harbor = store.get(entity_id="harbor", entity_type="project") + company_harbor = store.get(entity_id="harbor", entity_type="company") + assert project_harbor is not None and project_harbor.relationships == [] + assert company_harbor is not None and len(company_harbor.relationships) == 1 + + def test_remember_about_reads_its_own_qualified_form(self, store: EntityMemoryStore) -> None: + # remember_about is the only entity tool that declares its own type, so + # the first write under "project/Harbor" had no row to recognise the + # prefix from and slugged the whole string into project/project_harbor - + # a phantom no later call could reach. + store.remember_about(entity="Harbor", entity_type="company", facts=["in talks"]) + store.remember_about(entity="project/Harbor", entity_type="project", facts=["ingest rewrite"]) + store.remember_about(entity="Harbor", entity_type="project", facts=["due friday"]) + + rows = {(e.entity_type, e.entity_id) for e in store.list_entities(limit=10)} + assert rows == {("company", "harbor"), ("project", "harbor")} + harbor = store.get(entity_id="harbor", entity_type="project") + assert harbor is not None and len(harbor.live_facts()) == 2 + + def test_a_slash_in_a_name_is_still_a_name(self, store: EntityMemoryStore) -> None: + store.remember_about(entity="AC/DC", entity_type="company", facts=["a band"]) + assert store.get(entity_id="ac_dc", entity_type="company") is not None + + def test_names_differing_only_by_a_naming_symbol_stay_apart(self, store: EntityMemoryStore) -> None: + # C, C++ and C# all slugged to `c`, so one row ended up holding three + # languages' facts under one name, and the two discarded names were + # written nowhere. A merge has no unmerge. + store.remember_about(entity="C++", entity_type="system", facts=["manual memory management"]) + store.remember_about(entity="C#", entity_type="system", facts=["garbage collected"]) + store.remember_about(entity="C", entity_type="system", facts=["no classes"]) + + rows = {e.entity_id: e for e in store.list_entities(entity_type="system", limit=10)} + assert sorted(rows) == ["c", "c_plus_plus", "c_sharp"] + assert [f["content"] for f in rows["c_sharp"].live_facts()] == ["garbage collected"] + + def test_punctuation_variants_still_merge_and_keep_the_other_surface(self, store: EntityMemoryStore) -> None: + # The collapse is load-bearing for real merges; what it must not do is + # discard the name it merged. + store.remember_about(entity="Acme, Inc.", entity_type="company", facts=["vendor"]) + store.remember_about(entity="Acme Inc", entity_type="company", facts=["renewal in march"]) + + rows = store.list_entities(entity_type="company", limit=10) + assert len(rows) == 1 + assert rows[0].aliases == ["Acme Inc"] + assert len(rows[0].live_facts()) == 2 + + def test_a_relation_stated_from_both_sides_can_be_retired(self, store: EntityMemoryStore) -> None: + # The model states a symmetric relation from either end, so the row + # holds the same link twice - once outgoing, once incoming. Both render; + # a listing that prints only "relation -> far" offers the same string + # twice and no needle can pick one. + store.remember_about(entity="Alice", entity_type="person") + store.remember_about(entity="Bob", entity_type="person") + store.link_entities(entity="Alice", relation="pairs_with", related_entity="Bob") + store.link_entities(entity="Bob", relation="pairs_with", related_entity="Alice") + + assert "Removed relationship" in store.forget(entity="Alice", fact="pairs_with -> Bob") + assert "Removed relationship" in store.forget(entity="Alice", fact="pairs_with <- Bob") + alice = store.get(entity_id="alice", entity_type="person") + bob = store.get(entity_id="bob", entity_type="person") + assert alice is not None and alice.relationships == [] + assert bob is not None and bob.relationships == [] + + def test_a_far_end_name_that_extends_another_stays_retireable(self, store: EntityMemoryStore) -> None: + # "designed_by -> Sarah Chen" used to match the edge to `sarah` too, and + # the qualified form could not break the tie either, because + # "person/sarah" is a substring of "person/sarah_chen". + store.remember_about(entity="Sarah", entity_type="person") + store.remember_about(entity="Sarah Chen", entity_type="person") + store.remember_about(entity="radar", entity_type="project") + store.link_entities(entity="radar", relation="designed_by", related_entity="Sarah") + store.link_entities(entity="radar", relation="designed_by", related_entity="Sarah Chen") + + refusal = store.forget(entity="radar", fact="designed_by -> Sarah Chen") + assert "Multiple relationships" in refusal + assert "Removed relationship" in store.forget(entity="radar", fact="designed_by -> person/sarah_chen") + radar = store.get(entity_id="radar", entity_type="project") + assert radar is not None + assert [r["entity_id"] for r in radar.relationships] == ["sarah"] + + def test_an_accented_far_end_is_retireable_as_rendered(self, store: EntityMemoryStore) -> None: + # The block renders the display name; the edge stores the folded slug. + store.remember_about(entity="Harbor", entity_type="project") + store.remember_about(entity="Caf\u00e9 Noir", entity_type="company") + store.link_entities(entity="Harbor", relation="supplier_is", related_entity="Caf\u00e9 Noir") + + assert "Removed relationship" in store.forget(entity="Harbor", fact="supplier_is -> Caf\u00e9 Noir") + harbor = store.get(entity_id="harbor", entity_type="project") + assert harbor is not None and harbor.relationships == [] + + def test_restating_an_event_does_not_double_it(self, store: EntityMemoryStore) -> None: + for _ in range(3): + store.remember_about(entity="Tom", entity_type="person", events=["heard he is leaving"]) + tom = store.get(entity_id="tom", entity_type="person") + assert tom is not None and len(tom.events) == 1 + + def test_exact_duplicate_events_retire_completely(self, store: EntityMemoryStore) -> None: + # A row written before add_event deduped carries exact copies; retiring + # one and reporting success left the retracted claim rendering. + store.remember_about(entity="Tom", entity_type="person", events=["heard he is leaving"]) + tom = store.get(entity_id="tom", entity_type="person") + assert tom is not None + tom.events.append(dict(tom.events[0], id="dupe1234")) + store._save_entity(entity=tom, user_id=None, agent_id=None, team_id=None, namespace="global") + + assert "Retired event" in store.forget(entity="Tom", fact="heard he is leaving") + tom = store.get(entity_id="tom", entity_type="person") + assert tom is not None and tom.events == [] diff --git a/libs/agno/tests/unit/learn/test_entity_supersession.py b/libs/agno/tests/unit/learn/test_entity_supersession.py new file mode 100644 index 0000000000..0f13677053 --- /dev/null +++ b/libs/agno/tests/unit/learn/test_entity_supersession.py @@ -0,0 +1,324 @@ +"""Unit tests for fact supersession on EntityMemoryStore. + +The judge is one structured model call in the write path: given the entity's +live facts and the newly stated ones, it marks superseded fact ids with +confidences, and only entries at or above supersession_threshold are retired. +The fake model must subclass Model so the real tool-execution loop runs. +""" + +import json +from typing import Any, AsyncIterator, Dict, Iterator, List, Optional, Tuple + +import pytest + +from agno.learn.config import EntityMemoryConfig +from agno.learn.stores.entity_memory import EntityMemoryStore +from agno.models.base import Model +from agno.models.response import ModelResponse + + +class RecordingLearningDb: + """In-memory fake of the learnings table (duplicated per test file - the + tests/unit/learn directory is not a package).""" + + def __init__(self) -> None: + self.rows: Dict[str, Dict[str, Any]] = {} + self._clock = 0 + + def get_learning(self, **kwargs: Any) -> Optional[Dict[str, Any]]: + for row in self.rows.values(): + if ( + row.get("learning_type") == kwargs.get("learning_type") + and row.get("entity_id") == kwargs.get("entity_id") + and row.get("entity_type") == kwargs.get("entity_type") + and row.get("namespace") == kwargs.get("namespace") + ): + return row + return None + + def upsert_learning(self, id: str, **kwargs: Any) -> None: + existing = self.rows.get(id, {}) + row = {**existing, **kwargs, "learning_id": id} + self._clock += 1 + row["updated_at"] = self._clock + self.rows[id] = row + + def get_learnings(self, **kwargs: Any) -> List[Dict[str, Any]]: + rows = [ + row + for row in self.rows.values() + if all( + kwargs.get(key) is None or row.get(key) == kwargs.get(key) + for key in ("learning_type", "entity_id", "entity_type", "namespace") + ) + ] + rows.sort(key=lambda r: r.get("updated_at", 0), reverse=True) + limit = kwargs.get("limit") + return rows[:limit] if limit is not None else rows + + def delete_learning(self, id: str) -> bool: + return self.rows.pop(id, None) is not None + + def search_learnings(self, query: str, **kwargs: Any) -> List[Dict[str, Any]]: + limit = kwargs.pop("limit", None) + for key in ("workflow_id", "session_id", "agent_id", "team_id", "user_id"): + kwargs.pop(key, None) + candidates = self.get_learnings(**kwargs) + variants = {query.lower(), query.lower().replace(" ", "_"), query.lower().replace("_", " ")} + rows = [row for row in candidates if any(v in json.dumps(row.get("content", {})).lower() for v in variants)] + return rows[:limit] if limit is not None else rows + + +class SupersessionJudgeModel(Model): + """Calls mark_superseded once with a canned verdict, then stops.""" + + def __init__( + self, + verdict: List[Tuple[str, float]], + provider_calls: Optional[List[dict]] = None, + verdict_by_content: Optional[Tuple[str, float]] = None, + ) -> None: + super().__init__(id="supersession-judge-test", name="supersession-judge-test", provider="test") + self.verdict = verdict + self.provider_calls = provider_calls if provider_calls is not None else [] + self.verdict_by_content = verdict_by_content + + def __deepcopy__(self, memo: dict) -> "SupersessionJudgeModel": + # Share mutable state so the recorder survives the store's deepcopy + return type(self)( + verdict=self.verdict, provider_calls=self.provider_calls, verdict_by_content=self.verdict_by_content + ) + + def _resolve_verdict(self, messages: Any) -> List[Tuple[str, float]]: + if self.verdict_by_content is None: + return self.verdict + # Find the fact id whose content contains the marker, judge that id. + marker, confidence = self.verdict_by_content + user_text = next(m.content for m in messages if m.role == "user") + for line in user_text.splitlines(): + if marker in line and line.strip().startswith("- ["): + fact_id = line.strip()[3 : line.strip().index("]")] + return [(fact_id, confidence)] + return [] + + def _response_for_call(self, messages: Any) -> ModelResponse: + call_number = len(self.provider_calls) + 1 + self.provider_calls.append({"call": call_number, "messages": messages}) + if call_number == 1: + verdict = self._resolve_verdict(messages) + return ModelResponse( + role="assistant", + tool_calls=[ + { + "id": "mark-superseded-1", + "type": "function", + "function": { + "name": "mark_superseded", + "arguments": json.dumps( + { + "fact_ids": [fact_id for fact_id, _ in verdict], + "confidences": [confidence for _, confidence in verdict], + } + ), + }, + } + ], + ) + return ModelResponse(role="assistant", content="Judged.") + + def invoke(self, *args: Any, messages: Any = None, **kwargs: Any) -> ModelResponse: + return self._response_for_call(messages) + + async def ainvoke(self, *args: Any, messages: Any = None, **kwargs: Any) -> ModelResponse: + return self._response_for_call(messages) + + def invoke_stream(self, *args: Any, **kwargs: Any) -> Iterator[ModelResponse]: + raise AssertionError("streaming should not be used") + yield # pragma: no cover + + async def ainvoke_stream(self, *args: Any, **kwargs: Any) -> AsyncIterator[ModelResponse]: + raise AssertionError("streaming should not be used") + yield # pragma: no cover + + def _parse_provider_response(self, response: Any, **kwargs: Any) -> ModelResponse: + return response + + def _parse_provider_response_delta(self, response: Any) -> ModelResponse: + return response + + +class RaisingModel(Model): + def __init__(self) -> None: + super().__init__(id="raising-test", name="raising-test", provider="test") + + def __deepcopy__(self, memo: dict) -> "RaisingModel": + return self + + def invoke(self, *args: Any, **kwargs: Any) -> ModelResponse: + raise RuntimeError("judge model unavailable") + + async def ainvoke(self, *args: Any, **kwargs: Any) -> ModelResponse: + raise RuntimeError("judge model unavailable") + + def invoke_stream(self, *args: Any, **kwargs: Any) -> Iterator[ModelResponse]: + raise AssertionError("streaming should not be used") + yield # pragma: no cover + + async def ainvoke_stream(self, *args: Any, **kwargs: Any) -> AsyncIterator[ModelResponse]: + raise AssertionError("streaming should not be used") + yield # pragma: no cover + + def _parse_provider_response(self, response: Any, **kwargs: Any) -> ModelResponse: + return response + + def _parse_provider_response_delta(self, response: Any) -> ModelResponse: + return response + + +def _store_with_judge(db: RecordingLearningDb, model: Model, **config_kwargs: Any) -> EntityMemoryStore: + return EntityMemoryStore(config=EntityMemoryConfig(db=db, model=model, **config_kwargs)) # type: ignore[arg-type] + + +@pytest.fixture +def db() -> RecordingLearningDb: + return RecordingLearningDb() + + +def test_supersession_retires_contradicted_fact(db: RecordingLearningDb) -> None: + model = SupersessionJudgeModel(verdict=[], verdict_by_content=("radar is blocked", 0.95)) + store = _store_with_judge(db, model) + + store.remember_about(entity="radar", entity_type="project", facts=["radar is blocked on review"]) + store.remember_about(entity="radar", entity_type="project", facts=["radar shipped"]) + + entity = store.get(entity_id="radar", entity_type="project") + assert entity is not None + live = entity.live_facts() + assert [f["content"] for f in live] == ["radar shipped"] + # Both facts are still in the record; the superseded one carries the marker + assert len(entity.facts) == 2 + retired = [f for f in entity.facts if f.get("superseded_at")] + assert len(retired) == 1 + assert retired[0]["content"] == "radar is blocked on review" + # With one new fact, superseded_by points at the new fact's id + new_fact_id = live[0]["id"] + assert retired[0]["superseded_by"] == new_fact_id + + +def test_supersession_below_threshold_keeps_fact(db: RecordingLearningDb) -> None: + model = SupersessionJudgeModel(verdict=[], verdict_by_content=("radar is blocked", 0.5)) + store = _store_with_judge(db, model) + + store.remember_about(entity="radar", entity_type="project", facts=["radar is blocked on review"]) + message = store.remember_about(entity="radar", entity_type="project", facts=["radar shipped"]) + + assert "Superseded" not in message + entity = store.get(entity_id="radar", entity_type="project") + assert entity is not None + assert len(entity.live_facts()) == 2 + + +def test_supersession_threshold_is_configurable(db: RecordingLearningDb) -> None: + model = SupersessionJudgeModel(verdict=[], verdict_by_content=("radar is blocked", 0.95)) + store = _store_with_judge(db, model, supersession_threshold=0.99) + + store.remember_about(entity="radar", entity_type="project", facts=["radar is blocked on review"]) + store.remember_about(entity="radar", entity_type="project", facts=["radar shipped"]) + + entity = store.get(entity_id="radar", entity_type="project") + assert entity is not None + assert len(entity.live_facts()) == 2 + + +def test_judge_skipped_when_entity_is_new(db: RecordingLearningDb) -> None: + calls: List[dict] = [] + model = SupersessionJudgeModel(verdict=[], provider_calls=calls) + store = _store_with_judge(db, model) + + store.remember_about(entity="radar", entity_type="project", facts=["first fact"]) + assert calls == [] + + +def test_judge_skipped_when_write_has_no_facts(db: RecordingLearningDb) -> None: + calls: List[dict] = [] + model = SupersessionJudgeModel(verdict=[], provider_calls=calls) + store = _store_with_judge(db, model) + + store.remember_about(entity="radar", entity_type="project", facts=["first fact"]) + store.remember_about(entity="radar", entity_type="project", events=["shipped v1"]) + assert calls == [] + + +def test_judge_skipped_without_model(db: RecordingLearningDb) -> None: + store = EntityMemoryStore(config=EntityMemoryConfig(db=db)) # type: ignore[arg-type] + store.remember_about(entity="radar", entity_type="project", facts=["a"]) + message = store.remember_about(entity="radar", entity_type="project", facts=["b"]) + assert "Recorded 1 fact(s)" in message + + +def test_judge_failure_does_not_lose_the_write(db: RecordingLearningDb) -> None: + store = _store_with_judge(db, RaisingModel()) + store.remember_about(entity="radar", entity_type="project", facts=["first fact"]) + message = store.remember_about(entity="radar", entity_type="project", facts=["second fact"]) + + assert "Recorded 1 fact(s)" in message + entity = store.get(entity_id="radar", entity_type="project") + assert entity is not None + assert len(entity.live_facts()) == 2 + + +def test_exact_duplicate_fact_not_appended_and_judge_skipped(db: RecordingLearningDb) -> None: + calls: List[dict] = [] + model = SupersessionJudgeModel(verdict=[], provider_calls=calls) + store = _store_with_judge(db, model) + + store.remember_about(entity="radar", entity_type="project", facts=["db: Postgres"]) + message = store.remember_about(entity="radar", entity_type="project", facts=["DB: Postgres"]) + + assert "already-recorded" in message + assert calls == [] # nothing novel to judge + entity = store.get(entity_id="radar", entity_type="project") + assert entity is not None + assert len(entity.facts) == 1 + + +def test_hallucinated_fact_id_cannot_retire_new_fact(db: RecordingLearningDb) -> None: + store_no_model = EntityMemoryStore(config=EntityMemoryConfig(db=db)) # type: ignore[arg-type] + store_no_model.remember_about(entity="radar", entity_type="project", facts=["old fact"]) + entity = store_no_model.get(entity_id="radar", entity_type="project") + assert entity is not None + + # Judge marks an id that does not belong to any pre-existing fact + model = SupersessionJudgeModel(verdict=[("not_a_real_id", 0.99)]) + store = _store_with_judge(db, model) + store.remember_about(entity="radar", entity_type="project", facts=["new fact"]) + + entity = store.get(entity_id="radar", entity_type="project") + assert entity is not None + assert len(entity.live_facts()) == 2 + + +def test_custom_system_message_reaches_the_judge(db: RecordingLearningDb) -> None: + calls: List[dict] = [] + model = SupersessionJudgeModel(verdict=[], provider_calls=calls) + store = _store_with_judge(db, model, system_message="CUSTOM JUDGE PROMPT") + + store.remember_about(entity="radar", entity_type="project", facts=["a"]) + store.remember_about(entity="radar", entity_type="project", facts=["b"]) + + assert calls + system = next(m for m in calls[0]["messages"] if m.role == "system") + assert system.content == "CUSTOM JUDGE PROMPT" + + +async def test_async_supersession_retires_contradicted_fact(db: RecordingLearningDb) -> None: + model = SupersessionJudgeModel(verdict=[], verdict_by_content=("radar is blocked", 0.95)) + store = _store_with_judge(db, model) + + await store.aremember_about(entity="radar", entity_type="project", facts=["radar is blocked on review"]) + message = await store.aremember_about(entity="radar", entity_type="project", facts=["radar shipped"]) + + assert "Superseded 1 earlier fact(s)" in message + entity = await store.aget(entity_id="radar", entity_type="project") + assert entity is not None + assert [f["content"] for f in entity.live_facts()] == ["radar shipped"] diff --git a/libs/agno/tests/unit/learn/test_guidance_split.py b/libs/agno/tests/unit/learn/test_guidance_split.py new file mode 100644 index 0000000000..a5267db576 --- /dev/null +++ b/libs/agno/tests/unit/learn/test_guidance_split.py @@ -0,0 +1,480 @@ +"""Unit tests for the guidance/data split (spec §4.1) and the manual door (§4.2). + +instructions() returns guidance only (mode-aware, aggregated on the machine); +build_context() returns data only; the automatic injection path concatenates the +two, so the manual door and the automatic door render the same text. +""" + +from typing import Any, Dict, List, Optional + +from agno.learn import LearningMachine +from agno.learn.config import ( + DecisionLogConfig, + EntityMemoryConfig, + LearningMode, + UserMemoryConfig, + UserProfileConfig, +) +from agno.learn.stores.decision_log import DecisionLogStore +from agno.learn.stores.entity_memory import EntityMemoryStore +from agno.learn.stores.session_context import SessionContextStore +from agno.learn.stores.user_memory import UserMemoryStore +from agno.learn.stores.user_profile import UserProfileStore + + +class RecordingLearningDb: + """In-memory fake of the learnings table (duplicated per test file - the + tests/unit/learn directory is not a package).""" + + def __init__(self) -> None: + self.rows: Dict[str, Dict[str, Any]] = {} + self._clock = 0 + + def get_learning(self, **kwargs: Any) -> Optional[Dict[str, Any]]: + for row in self.rows.values(): + if all( + kwargs.get(key) is None or row.get(key) == kwargs.get(key) + for key in ("learning_type", "entity_id", "entity_type", "namespace", "user_id") + if kwargs.get(key) is not None + ): + return row + return None + + def upsert_learning(self, id: str, **kwargs: Any) -> None: + existing = self.rows.get(id, {}) + row = {**existing, **kwargs, "learning_id": id} + self._clock += 1 + row["updated_at"] = self._clock + self.rows[id] = row + + def get_learnings(self, **kwargs: Any) -> List[Dict[str, Any]]: + rows = [ + row + for row in self.rows.values() + if all( + kwargs.get(key) is None or row.get(key) == kwargs.get(key) + for key in ("learning_type", "entity_id", "entity_type", "namespace") + ) + ] + rows.sort(key=lambda r: r.get("updated_at", 0), reverse=True) + limit = kwargs.get("limit") + return rows[:limit] if limit is not None else rows + + def delete_learning(self, id: str) -> bool: + return self.rows.pop(id, None) is not None + + def search_learnings(self, query: str, **kwargs: Any) -> List[Dict[str, Any]]: + import json + + limit = kwargs.pop("limit", None) + for key in ("workflow_id", "session_id", "agent_id", "team_id", "user_id"): + kwargs.pop(key, None) + candidates = self.get_learnings(**kwargs) + variants = {query.lower(), query.lower().replace(" ", "_"), query.lower().replace("_", " ")} + rows = [row for row in candidates if any(v in json.dumps(row.get("content", {})).lower() for v in variants)] + return rows[:limit] if limit is not None else rows + + +class TestStoreSplit: + def test_agentic_stores_guidance_never_in_data(self) -> None: + db = RecordingLearningDb() + cases = [ + ( + UserProfileStore(config=UserProfileConfig(db=db, mode=LearningMode.AGENTIC)), # type: ignore[arg-type] + "update_profile", + ), + ( + UserMemoryStore(config=UserMemoryConfig(db=db, mode=LearningMode.AGENTIC)), # type: ignore[arg-type] + "update_user_memory", + ), + ( + EntityMemoryStore(config=EntityMemoryConfig(db=db)), # type: ignore[arg-type] + "remember_about", + ), + ( + DecisionLogStore(config=DecisionLogConfig(db=db, mode=LearningMode.AGENTIC)), # type: ignore[arg-type] + "log_decision", + ), + ] + for store, tool_name in cases: + guidance = store.instructions() + assert tool_name in guidance, store + data = store.build_context(data=store.recall(user_id="u1")) + assert tool_name not in data, store + + def test_always_mode_stores_have_no_guidance(self) -> None: + db = RecordingLearningDb() + profile = UserProfileStore(config=UserProfileConfig(db=db, mode=LearningMode.ALWAYS)) # type: ignore[arg-type] + assert profile.instructions() == "" + session = SessionContextStore() + assert session.instructions() == "" + + def test_machine_instructions_aggregates_enabled_stores(self) -> None: + db = RecordingLearningDb() + machine = LearningMachine( + db=db, # type: ignore[arg-type] + user_memory=UserMemoryConfig(mode=LearningMode.AGENTIC), + entity_memory=True, + ) + guidance = machine.instructions() + assert "update_user_memory" in guidance + assert "remember_about" in guidance + assert "update_profile" not in guidance # not enabled + + def test_machine_instructions_is_mode_aware(self) -> None: + db = RecordingLearningDb() + always_machine = LearningMachine(db=db, user_memory=UserMemoryConfig(mode=LearningMode.ALWAYS)) # type: ignore[arg-type] + assert always_machine.instructions() == "" + + def test_custom_store_without_instructions_still_works(self) -> None: + class OldStyleStore: + # A pre-2.8.4 third-party store: no instructions() method. + learning_type = "old_style" + schema = dict + was_updated = False + + def recall(self, **kwargs: Any) -> Any: + return {"x": 1} + + async def arecall(self, **kwargs: Any) -> Any: + return {"x": 1} + + def process(self, messages: Any, **kwargs: Any) -> None: + pass + + async def aprocess(self, messages: Any, **kwargs: Any) -> None: + pass + + def build_context(self, data: Any) -> str: + return "x" + + def get_tools(self, **kwargs: Any) -> List[Any]: + return [] + + async def aget_tools(self, **kwargs: Any) -> List[Any]: + return [] + + machine = LearningMachine(db=RecordingLearningDb(), custom_stores={"old_style": OldStyleStore()}) # type: ignore[arg-type] + assert "old_style" in machine.stores + assert machine.instructions() == "" # contributes no guidance, no crash + assert "x" in machine.build_context() + + +class TestRunContextThreading: + def test_narrow_custom_store_keeps_working(self) -> None: + """The §8 bullet: a custom store whose recall(self, user_id=None) takes + no **kwargs still works after the new context kwargs arrive.""" + calls: Dict[str, Any] = {} + + class NarrowStore: + learning_type = "narrow" + schema = dict + was_updated = False + + def recall(self, user_id=None): # no **kwargs, deliberately + calls["recall_user_id"] = user_id + return {"seen": True} + + async def arecall(self, user_id=None): + calls["arecall_user_id"] = user_id + return {"seen": True} + + def process(self, messages, user_id=None): + calls["process_messages"] = messages + + async def aprocess(self, messages, user_id=None): + calls["aprocess_messages"] = messages + + def build_context(self, data): + return "ok" + + def get_tools(self, user_id=None): + calls["get_tools_user_id"] = user_id + return [] + + async def aget_tools(self, user_id=None): + return [] + + machine = LearningMachine(db=RecordingLearningDb(), custom_stores={"narrow": NarrowStore()}) # type: ignore[arg-type] + context = machine.build_context( + user_id="u1", + message="the current message", + run_context=object(), + metadata={"k": "v"}, + dependencies={"d": 1}, + session_state={"s": 2}, + ) + assert "ok" in context + assert calls["recall_user_id"] == "u1" + + machine.get_tools(user_id="u1", run_context=object()) + assert calls["get_tools_user_id"] == "u1" + + machine.process(messages=["m"], user_id="u1", run_context=object(), metadata={}, session_state={}) + assert calls["process_messages"] == ["m"] + + async def test_narrow_custom_store_async_paths(self) -> None: + class NarrowStore: + learning_type = "narrow" + schema = dict + was_updated = False + + def recall(self, user_id=None): + return None + + async def arecall(self, user_id=None): + return {"seen": True} + + def process(self, messages, user_id=None): + pass + + async def aprocess(self, messages, user_id=None): + pass + + def build_context(self, data): + return "async ok" if data else "" + + def get_tools(self, user_id=None): + return [] + + async def aget_tools(self, user_id=None): + return [] + + machine = LearningMachine(db=RecordingLearningDb(), custom_stores={"narrow": NarrowStore()}) # type: ignore[arg-type] + context = await machine.abuild_context(user_id="u1", run_context=object(), metadata={}) + assert "async ok" in context + await machine.aprocess(messages=["m"], run_context=object(), session_state={}) + + def test_builtin_store_receives_the_current_message(self) -> None: + """Built-in stores take **kwargs and see everything, including message.""" + seen: Dict[str, Any] = {} + machine = LearningMachine(db=RecordingLearningDb(), entity_memory=True) # type: ignore[arg-type] + store = machine.entity_memory_store + assert store is not None + original_recall = store.recall + + def spy_recall(**kwargs: Any) -> Any: + seen.update(kwargs) + return original_recall( + **{k: v for k, v in kwargs.items() if k in ("entity_id", "entity_type", "user_id", "namespace")} + ) + + store.recall = spy_recall # type: ignore[method-assign] + machine.build_context(user_id="u1", message="tell me about radar", run_context="RC") + assert seen.get("message") == "tell me about radar" + assert seen.get("run_context") == "RC" + + def test_injection_site_threads_the_input_message(self) -> None: + """End to end: get_system_message(input=...) reaches build_context(message=...).""" + from unittest.mock import MagicMock + + from agno.agent import Agent + from agno.agent._messages import get_system_message + from agno.models.openai import OpenAIResponses + from agno.run.base import RunContext + from agno.session import AgentSession + + db = RecordingLearningDb() + machine = LearningMachine(db=db, entity_memory=True) # type: ignore[arg-type] + machine.build_context = MagicMock(return_value="") # type: ignore[method-assign] + + agent = Agent(db=db, learning=machine, model=OpenAIResponses(id="gpt-5.5")) # type: ignore[arg-type] + agent._learning = machine + session = AgentSession(session_id="s1") + run_context = RunContext(run_id="r1", session_id="s1", user_id="u1", metadata={"m": 1}, session_state={"x": 2}) + + get_system_message(agent, session=session, run_context=run_context, input="what about radar?") + kwargs = machine.build_context.call_args.kwargs + assert kwargs["message"] == "what about radar?" + assert kwargs["run_context"] is run_context + assert kwargs["metadata"] == {"m": 1} + assert kwargs["session_state"] == {"x": 2} + + +class TestManualDoor: + def test_capture_hook_runs_process(self) -> None: + from types import SimpleNamespace + + db = RecordingLearningDb() + machine = LearningMachine(db=db, user_memory=True) # type: ignore[arg-type] + captured: Dict[str, Any] = {} + + def spy_process(**kwargs: Any) -> None: + captured.update(kwargs) + + machine.process = spy_process # type: ignore[method-assign] + hook = machine.capture_hook() + hook( + run_output=SimpleNamespace(messages=["m1"]), + agent=SimpleNamespace(id="agent-1", team_id=None), + session=SimpleNamespace(session_id="s1"), + user_id="u1", + ) + assert captured["messages"] == ["m1"] + assert captured["user_id"] == "u1" + assert captured["session_id"] == "s1" + assert captured["agent_id"] == "agent-1" + + async def test_acapture_hook_schedules_aprocess(self) -> None: + import asyncio + from types import SimpleNamespace + + db = RecordingLearningDb() + machine = LearningMachine(db=db, user_memory=True) # type: ignore[arg-type] + captured: Dict[str, Any] = {} + + async def spy_aprocess(**kwargs: Any) -> None: + captured.update(kwargs) + + machine.aprocess = spy_aprocess # type: ignore[method-assign] + hook = machine.acapture_hook() + await hook( + run_output=SimpleNamespace(messages=["m1"]), + agent=SimpleNamespace(id="agent-1", team_id=None), + session=SimpleNamespace(session_id="s1"), + user_id="u1", + ) + # Fire-and-forget: the capture runs as a background task off the + # response path; drain it before asserting. + await asyncio.gather(*machine._capture_tasks) + assert captured["messages"] == ["m1"] + + def test_capture_hook_is_post_hooks_compatible(self) -> None: + # filter_hook_args must find only parameters it can supply. + import inspect + + from agno.agent._hooks import filter_hook_args + + machine = LearningMachine(db=RecordingLearningDb(), user_memory=True) # type: ignore[arg-type] + hook = machine.capture_hook() + params = set(inspect.signature(hook).parameters) + assert params <= {"run_output", "agent", "session", "user_id", "run_context"} + filtered = filter_hook_args( + hook, + { + "run_output": None, + "agent": None, + "session": None, + "user_id": "u1", + "run_context": object(), + "debug_mode": False, + "metadata": None, + }, + ) + assert set(filtered) == params + + def test_double_render_warns_once(self, caplog) -> None: + import logging + + machine = LearningMachine(db=RecordingLearningDb(), entity_memory=True) # type: ignore[arg-type] + machine.instructions() # the manual door was used + with caplog.at_level(logging.WARNING): + machine._framework_instructions() # ...and the framework injects too + machine._framework_instructions() + warnings = [r for r in caplog.records if "render twice" in r.getMessage()] + assert len(warnings) == 1 + + def test_no_warning_on_a_single_door(self, caplog) -> None: + import logging + + machine = LearningMachine(db=RecordingLearningDb(), entity_memory=True) # type: ignore[arg-type] + with caplog.at_level(logging.WARNING): + machine._framework_instructions() + machine._framework_instructions() + assert [r for r in caplog.records if "render twice" in r.getMessage()] == [] + + +class TestManualDoorMatchesAutomaticDoor: + def test_injected_block_equals_instructions_plus_build_context(self) -> None: + """The §8 regression guard: for the same machine, what the automatic + injection site renders equals instructions() + build_context().""" + db = RecordingLearningDb() + machine = LearningMachine( + db=db, # type: ignore[arg-type] + user_memory=UserMemoryConfig(mode=LearningMode.AGENTIC), + entity_memory=True, + ) + entity_store = machine.entity_memory_store + assert entity_store is not None + entity_store.remember_about(entity="radar", entity_type="project", facts=["db: Postgres"]) + + # What the automatic door assembles at the injection site + # (agent/_messages.py): instructions() + "\n" + build_context(...) + guidance = machine.instructions() + data = machine.build_context(user_id="u1", session_id="s1", agent_id="a1") + auto_block = "\n".join(part for part in (guidance, data) if part) + + # The manual door places the same two surfaces by hand + manual_block = "\n".join( + part + for part in ( + machine.instructions(), + machine.build_context(user_id="u1", session_id="s1", agent_id="a1"), + ) + if part + ) + assert auto_block == manual_block + assert "remember_about" in auto_block # guidance present + assert "Entity directory" in auto_block # data present + + def test_agent_system_message_carries_guidance_and_data(self) -> None: + """End to end through the real injection site in agent/_messages.py.""" + from typing import Any as _Any + + from agno.agent import Agent + from agno.agent._messages import get_system_message + from agno.models.base import Model + from agno.models.response import ModelResponse + from agno.run.base import RunContext + from agno.session import AgentSession + + class _NullModel(Model): + def __init__(self) -> None: + super().__init__(id="null-test", name="null-test", provider="test") + + def invoke(self, *args: _Any, **kwargs: _Any) -> ModelResponse: + return ModelResponse(role="assistant", content="") + + async def ainvoke(self, *args: _Any, **kwargs: _Any) -> ModelResponse: + return ModelResponse(role="assistant", content="") + + def invoke_stream(self, *args: _Any, **kwargs: _Any): + raise AssertionError("not used") + yield # pragma: no cover + + async def ainvoke_stream(self, *args: _Any, **kwargs: _Any): + raise AssertionError("not used") + yield # pragma: no cover + + def _parse_provider_response(self, response: _Any, **kwargs: _Any) -> ModelResponse: + return response + + def _parse_provider_response_delta(self, response: _Any) -> ModelResponse: + return response + + db = RecordingLearningDb() + machine = LearningMachine(db=db, entity_memory=True) # type: ignore[arg-type] + entity_store = machine.entity_memory_store + assert entity_store is not None + entity_store.remember_about(entity="radar", entity_type="project", facts=["db: Postgres"]) + + agent = Agent(db=db, learning=machine, model=_NullModel()) # type: ignore[arg-type] + agent._learning = machine + session = AgentSession(session_id="s1") + run_context = RunContext(run_id="r1", session_id="s1", user_id="u1") + + message = get_system_message(agent, session=session, run_context=run_context) + assert message is not None + content = str(message.content) + assert "remember_about" in content # guidance reached the system prompt + assert "Entity directory" in content # data reached the system prompt + # And it equals the manual door's assembly, embedded verbatim + manual = "\n".join( + part + for part in ( + machine.instructions(), + machine.build_context(user_id="u1", session_id="s1", agent_id=agent.id), + ) + if part + ) + assert manual in content diff --git a/libs/agno/tests/unit/learn/test_machine_roundtrip.py b/libs/agno/tests/unit/learn/test_machine_roundtrip.py new file mode 100644 index 0000000000..b98ca1a7c7 --- /dev/null +++ b/libs/agno/tests/unit/learn/test_machine_roundtrip.py @@ -0,0 +1,197 @@ +"""Unit tests for the lossless LearningMachine config round-trip (spec §3.9). + +Verified live before the fix: an AGENTIC machine round-tripped through +Agent.save()/load() came back as ALWAYS and get_tools() returned [] - the +agent silently lost its entire learning tool surface. +""" + +import logging +from dataclasses import dataclass, field +from typing import Any, Dict, List, Optional + +from agno.learn import LearningMachine, UserProfile +from agno.learn.config import ( + EntityMemoryConfig, + LearningMode, + UserMemoryConfig, +) + + +class RecordingLearningDb: + def __init__(self) -> None: + self.rows: Dict[str, Dict[str, Any]] = {} + + def get_learning(self, **kwargs: Any) -> Optional[Dict[str, Any]]: + return None + + def upsert_learning(self, id: str, **kwargs: Any) -> None: + self.rows[id] = {**kwargs, "learning_id": id} + + def get_learnings(self, **kwargs: Any) -> List[Dict[str, Any]]: + return [] + + def search_learnings(self, query: str, **kwargs: Any) -> List[Dict[str, Any]]: + return [] + + def delete_learning(self, id: str) -> bool: + return self.rows.pop(id, None) is not None + + +@dataclass +class ExtendedProfile(UserProfile): + company: Optional[str] = field(default=None, metadata={"description": "Where they work"}) + + +def test_round_trip_preserves_mode_and_tools() -> None: + """The §8 bullet: from_dict(to_dict()) keeps AGENTIC, and the rebuilt + machine still has its tool surface.""" + machine = LearningMachine( + db=RecordingLearningDb(), # type: ignore[arg-type] + user_memory=UserMemoryConfig(mode=LearningMode.AGENTIC), + entity_memory=EntityMemoryConfig(namespace="global", supersession_threshold=0.9), + ) + rebuilt = LearningMachine.from_dict(machine.to_dict()) + rebuilt.db = RecordingLearningDb() # type: ignore[assignment] + + assert isinstance(rebuilt.user_memory, UserMemoryConfig) + assert rebuilt.user_memory.mode is LearningMode.AGENTIC + assert isinstance(rebuilt.entity_memory, EntityMemoryConfig) + assert rebuilt.entity_memory.supersession_threshold == 0.9 + + tools = rebuilt.get_tools(user_id="u1") + names = {t.__name__ for t in tools} + assert "update_user_memory" in names + assert {"remember_about", "link_entities", "search_entities", "forget"} <= names + + +def test_round_trip_preserves_prompts_limits_and_namespace() -> None: + machine = LearningMachine( + db=RecordingLearningDb(), # type: ignore[arg-type] + user_memory=UserMemoryConfig( + mode=LearningMode.AGENTIC, + instructions="capture durable facts", + additional_instructions="ignore small talk", + max_updates_per_run=7, + enable_delete_memory=False, + ), + namespace="team_west", + max_updates_per_run=25, + ) + rebuilt = LearningMachine.from_dict(machine.to_dict()) + + assert rebuilt.namespace == "team_west" + assert rebuilt.max_updates_per_run == 25 + assert isinstance(rebuilt.user_memory, UserMemoryConfig) + assert rebuilt.user_memory.instructions == "capture durable facts" + assert rebuilt.user_memory.additional_instructions == "ignore small talk" + assert rebuilt.user_memory.max_updates_per_run == 7 + assert rebuilt.user_memory.enable_delete_memory is False + + +def test_round_trip_preserves_schema_by_import_path() -> None: + machine = LearningMachine( + db=RecordingLearningDb(), # type: ignore[arg-type] + user_profile=__import__("agno.learn.config", fromlist=["UserProfileConfig"]).UserProfileConfig( + mode=LearningMode.AGENTIC, schema=ExtendedProfile + ), + ) + payload = machine.to_dict() + assert payload["user_profile"]["schema"].endswith("ExtendedProfile") + + rebuilt = LearningMachine.from_dict(payload) + assert rebuilt.user_profile.schema is ExtendedProfile # type: ignore[union-attr] + + +def test_from_dict_accepts_old_boolean_payloads() -> None: + rebuilt = LearningMachine.from_dict( + {"user_profile": True, "user_memory": True, "namespace": "global", "debug_mode": False} + ) + assert rebuilt.user_profile is True + assert rebuilt.user_memory is True + assert rebuilt.entity_memory is False + + +def test_store_instance_serializes_its_config() -> None: + from agno.learn.stores.entity_memory import EntityMemoryStore + + store = EntityMemoryStore(config=EntityMemoryConfig(db=RecordingLearningDb(), namespace="ops")) # type: ignore[arg-type] + machine = LearningMachine(db=RecordingLearningDb(), entity_memory=store) # type: ignore[arg-type] + payload = machine.to_dict() + assert payload["entity_memory"]["namespace"] == "ops" + + rebuilt = LearningMachine.from_dict(payload) + assert isinstance(rebuilt.entity_memory, EntityMemoryConfig) + assert rebuilt.entity_memory.namespace == "ops" + + +def test_custom_stores_round_trip_as_logged_refs(caplog) -> None: + class TinyStore: + learning_type = "tiny" + schema = dict + was_updated = False + + def recall(self, **kwargs: Any) -> Any: + return None + + async def arecall(self, **kwargs: Any) -> Any: + return None + + def process(self, messages: Any, **kwargs: Any) -> None: + pass + + async def aprocess(self, messages: Any, **kwargs: Any) -> None: + pass + + def build_context(self, data: Any) -> str: + return "" + + def instructions(self) -> str: + return "" + + def get_tools(self, **kwargs: Any) -> List[Any]: + return [] + + async def aget_tools(self, **kwargs: Any) -> List[Any]: + return [] + + machine = LearningMachine(db=RecordingLearningDb(), custom_stores={"tiny": TinyStore()}) # type: ignore[arg-type] + payload = machine.to_dict() + assert "tiny" in payload["custom_stores"] + + with caplog.at_level(logging.WARNING): + rebuilt = LearningMachine.from_dict(payload) + assert rebuilt.custom_stores is None + assert any("cannot be rebuilt" in r.getMessage() for r in caplog.records) + + +def test_agent_save_load_keeps_the_tool_surface(tmp_path) -> None: + """End to end through the real Agent storage path (the launch-blocking bug: + an AGENTIC machine reloaded as ALWAYS and lost its tools).""" + from agno.agent import Agent + from agno.db.sqlite import SqliteDb + from agno.models.openai import OpenAIResponses + + db = SqliteDb(db_file=str(tmp_path / "agents.db")) + machine = LearningMachine( + user_memory=UserMemoryConfig(mode=LearningMode.AGENTIC), + entity_memory=True, + ) + agent = Agent( + id="round-trip-agent", + db=db, + model=OpenAIResponses(id="gpt-5.5"), + learning=machine, + ) + agent.save() + + loaded = Agent.load(id="round-trip-agent", db=db) + assert loaded is not None + assert isinstance(loaded.learning, LearningMachine) + assert isinstance(loaded.learning.user_memory, UserMemoryConfig) + assert loaded.learning.user_memory.mode is LearningMode.AGENTIC + + loaded.learning.db = db + tools = loaded.learning.get_tools(user_id="u1") + names = {t.__name__ for t in tools} + assert "update_user_memory" in names + assert "remember_about" in names diff --git a/libs/agno/tests/unit/learn/test_max_updates_per_run.py b/libs/agno/tests/unit/learn/test_max_updates_per_run.py index 127cb8dce8..34ba01f9d3 100644 --- a/libs/agno/tests/unit/learn/test_max_updates_per_run.py +++ b/libs/agno/tests/unit/learn/test_max_updates_per_run.py @@ -10,7 +10,6 @@ from agno.learn.config import ( UserMemoryConfig, UserProfileConfig, ) -from agno.learn.stores.entity_memory import EntityMemoryStore from agno.learn.stores.learned_knowledge import LearnedKnowledgeStore from agno.learn.stores.session_context import SessionContextStore from agno.learn.stores.user_memory import UserMemoryStore @@ -162,14 +161,6 @@ def _build_user_memory_store(model: Any, **config_kwargs: Any) -> tuple[Any, dic return store, {"user_id": "user-1"} -def _build_entity_memory_store(model: Any, **config_kwargs: Any) -> tuple[Any, dict[str, str]]: - config_kwargs.setdefault("max_updates_per_run", DEFAULT_MAX_UPDATES) - store = EntityMemoryStore( - config=EntityMemoryConfig(db=_RecordingLearningDb(), model=model, **config_kwargs) # type: ignore[arg-type] - ) - return store, {"user_id": "user-1"} - - def _build_learned_knowledge_store(model: Any, **config_kwargs: Any) -> tuple[Any, dict[str, str]]: config_kwargs.setdefault("max_updates_per_run", DEFAULT_MAX_UPDATES) store = LearnedKnowledgeStore( @@ -178,11 +169,11 @@ def _build_learned_knowledge_store(model: Any, **config_kwargs: Any) -> tuple[An return store, {"user_id": "user-1"} +# Entity memory is AGENTIC-only (no extraction pass), so it does not appear here. STORE_BUILDERS = [ pytest.param(_build_session_context_store, id="session_context"), pytest.param(_build_user_profile_store, id="user_profile"), pytest.param(_build_user_memory_store, id="user_memory"), - pytest.param(_build_entity_memory_store, id="entity_memory"), pytest.param(_build_learned_knowledge_store, id="learned_knowledge"), ] diff --git a/libs/agno/tests/unit/learn/test_relevance_recall.py b/libs/agno/tests/unit/learn/test_relevance_recall.py new file mode 100644 index 0000000000..7e47dc0c2d --- /dev/null +++ b/libs/agno/tests/unit/learn/test_relevance_recall.py @@ -0,0 +1,484 @@ +"""Unit tests for relevance recall (spec §3.7 / commit 15). + +recall(message=...) returns the top-k relevant entities: name/alias matches +against the message first, then a bounded lexical term search. The keyed +lookup stays available as get(); the directory is always injected. +""" + +import json +from typing import Any, Dict, List, Optional + +import pytest + +from agno.learn.config import EntityMemoryConfig +from agno.learn.stores.entity_memory import EntityMemoryStore, _message_terms + + +class RecordingLearningDb: + """In-memory fake of the learnings table (duplicated per test file - the + tests/unit/learn directory is not a package).""" + + def __init__(self) -> None: + self.rows: Dict[str, Dict[str, Any]] = {} + self._clock = 0 + + def get_learning(self, **kwargs: Any) -> Optional[Dict[str, Any]]: + for row in self.rows.values(): + if all( + kwargs.get(key) is None or row.get(key) == kwargs.get(key) + for key in ("learning_type", "entity_id", "entity_type", "namespace") + ): + return row + return None + + def upsert_learning(self, id: str, **kwargs: Any) -> None: + existing = self.rows.get(id, {}) + row = {**existing, **kwargs, "learning_id": id} + self._clock += 1 + row["updated_at"] = self._clock + self.rows[id] = row + + def get_learnings(self, **kwargs: Any) -> List[Dict[str, Any]]: + rows = [ + row + for row in self.rows.values() + if all( + kwargs.get(key) is None or row.get(key) == kwargs.get(key) + for key in ("learning_type", "entity_id", "entity_type", "namespace") + ) + ] + rows.sort(key=lambda r: r.get("updated_at", 0), reverse=True) + limit = kwargs.get("limit") + return rows[:limit] if limit is not None else rows + + def delete_learning(self, id: str) -> bool: + return self.rows.pop(id, None) is not None + + def search_learnings(self, query: str, **kwargs: Any) -> List[Dict[str, Any]]: + limit = kwargs.pop("limit", None) + for key in ("workflow_id", "session_id", "agent_id", "team_id", "user_id"): + kwargs.pop(key, None) + candidates = self.get_learnings(**kwargs) + variants = {query.lower(), query.lower().replace(" ", "_"), query.lower().replace("_", " ")} + rows = [row for row in candidates if any(v in json.dumps(row.get("content", {})).lower() for v in variants)] + return rows[:limit] if limit is not None else rows + + +@pytest.fixture +def db() -> RecordingLearningDb: + return RecordingLearningDb() + + +@pytest.fixture +def store(db: RecordingLearningDb) -> EntityMemoryStore: + return EntityMemoryStore(config=EntityMemoryConfig(db=db)) # type: ignore[arg-type] + + +class TestMessageTerms: + def test_extracts_distinctive_terms(self) -> None: + assert _message_terms("What did we decide about the radar project?") == ["decide", "radar", "project"] + + def test_empty_and_stopword_only(self) -> None: + assert _message_terms("") == [] + assert _message_terms("what is the") == [] + + +class TestRelevanceRecall: + def test_name_path_serves_without_term_search(self, db: RecordingLearningDb) -> None: + # Discriminating test for the NAME path: with k=1 filled by the name + # match, term search must not run at all - so the name path cannot be + # silently replaced by the more expensive search fallback. + store = EntityMemoryStore(config=EntityMemoryConfig(db=db, max_entities_in_context=1)) # type: ignore[arg-type] + store.remember_about(entity="radar", entity_type="project", facts=["db: Postgres"]) + + def forbidden_search(**kwargs: Any) -> Any: + raise AssertionError("term search must not run when the name path fills k") + + store.search = forbidden_search # type: ignore[method-assign] + recalled = store.recall(message="what is the status of radar?") + assert recalled is not None + assert [e.entity_id for e in recalled["entities"]] == ["radar"] + + def test_short_names_match_on_word_boundaries_only(self, store: EntityMemoryStore) -> None: + # A two-letter entity must not match inside other words and evict the + # entity the turn is actually about. + store.remember_about(entity="Al", entity_type="person", facts=["works in finance"]) + store.remember_about(entity="radar", entity_type="project", facts=["db: Postgres"]) + + recalled = store.recall(message="always check the radar dashboard") + assert recalled is not None + ids = [e.entity_id for e in recalled["entities"]] + assert "radar" in ids and "al" not in ids + + # ...but naming Al on a word boundary matches + recalled = store.recall(message="ask Al about the budget") + assert recalled is not None + assert "al" in [e.entity_id for e in recalled["entities"]] + + def test_message_naming_an_entity_expands_it(self, store: EntityMemoryStore) -> None: + store.remember_about(entity="radar", entity_type="project", facts=["db: Postgres"]) + store.remember_about(entity="unrelated", entity_type="project", facts=["nothing here"]) + + recalled = store.recall(message="what is the status of radar?") + assert recalled is not None + assert [e.entity_id for e in recalled["entities"]] == ["radar"] + + def test_multi_word_name_matches(self, store: EntityMemoryStore) -> None: + store.remember_about(entity="Sarah Chen", entity_type="person", facts=["designs radar"]) + recalled = store.recall(message="Ask Sarah Chen about the deadline") + assert recalled is not None + assert [e.entity_id for e in recalled["entities"]] == ["sarah_chen"] + + def test_alias_matches(self, store: EntityMemoryStore, db: RecordingLearningDb) -> None: + db.upsert_learning( + id="entity_global_project_radar", + learning_type="entity_memory", + entity_id="radar", + entity_type="project", + namespace="global", + content={ + "entity_id": "radar", + "entity_type": "project", + "name": "radar", + "aliases": ["The Radar Initiative"], + "facts": [], + }, + ) + recalled = store.recall(message="Where does the radar initiative stand?") + assert recalled is not None + assert [e.entity_id for e in recalled["entities"]] == ["radar"] + + def test_term_search_finds_entity_by_fact_content(self, store: EntityMemoryStore) -> None: + # The entity's NAME is not in the message; a fact matches the term. + store.remember_about(entity="radar", entity_type="project", facts=["uses postgresql for storage"]) + recalled = store.recall(message="who knows postgresql here?") + assert recalled is not None + assert [e.entity_id for e in recalled["entities"]] == ["radar"] + + def test_no_match_returns_directory_only(self, store: EntityMemoryStore) -> None: + store.remember_about(entity="radar", entity_type="project") + recalled = store.recall(message="have we discussed quantum tunneling?") + assert recalled is not None + assert recalled["entities"] == [] + assert [e.entity_id for e in recalled["directory"]] == ["radar"] + + def test_top_k_bound_is_honored(self, db: RecordingLearningDb) -> None: + store = EntityMemoryStore(config=EntityMemoryConfig(db=db, max_entities_in_context=2)) # type: ignore[arg-type] + for i in range(5): + store.remember_about(entity=f"radar {i}", entity_type="project") + recalled = store.recall(message="radar 0 radar 1 radar 2 radar 3 radar 4") + assert recalled is not None + assert len(recalled["entities"]) == 2 + + def test_archived_entities_do_not_recall(self, store: EntityMemoryStore) -> None: + store.remember_about(entity="radar", entity_type="project") + store.forget(entity="radar") + recalled = store.recall(message="what about radar?") + assert recalled is not None + assert recalled["entities"] == [] + + def test_keyed_lookup_composes_with_message(self, store: EntityMemoryStore) -> None: + store.remember_about(entity="radar", entity_type="project") + store.remember_about(entity="lidar", entity_type="project") + recalled = store.recall(entity_id="radar", entity_type="project", message="compare with lidar") + assert recalled is not None + ids = [e.entity_id for e in recalled["entities"]] + assert ids[0] == "radar" and "lidar" in ids + + async def test_async_relevance_recall(self, store: EntityMemoryStore) -> None: + await store.aremember_about(entity="radar", entity_type="project", facts=["db: Postgres"]) + recalled = await store.arecall(message="what about radar?") + assert recalled is not None + assert [e.entity_id for e in recalled["entities"]] == ["radar"] + + +class TestRecallReachesTheAgent: + def test_stored_entity_reaches_the_system_message(self) -> None: + """The §8 bullet whose absence let entity memory ship write-only: given + a stored entity and a message mentioning it, the rendered system + message contains that entity's facts.""" + from agno.agent import Agent + from agno.agent._messages import get_system_message + from agno.learn import LearningMachine + from agno.models.openai import OpenAIResponses + from agno.run.base import RunContext + from agno.session import AgentSession + + db = RecordingLearningDb() + machine = LearningMachine(db=db, entity_memory=True) # type: ignore[arg-type] + entity_store = machine.entity_memory_store + assert entity_store is not None + entity_store.remember_about(entity="radar", entity_type="project", facts=["db: Postgres, over Dynamo"]) + + agent = Agent(db=db, learning=machine, model=OpenAIResponses(id="gpt-5.5")) # type: ignore[arg-type] + agent._learning = machine + session = AgentSession(session_id="s1") + run_context = RunContext(run_id="r1", session_id="s1", user_id="u1") + + message = get_system_message( + agent, session=session, run_context=run_context, input="what did we decide for radar?" + ) + assert message is not None + content = str(message.content) + assert "db: Postgres, over Dynamo" in content + + def test_stored_entity_reaches_the_system_message_through_get_run_messages(self) -> None: + """The same bullet, guarded one layer up. + + The original write-only bug was a dropped kwarg at the CALL SITE, not + inside get_system_message, so the regression guard has to run the seam + that broke: get_run_messages must hand the run input down. + """ + from agno.agent import Agent + from agno.agent._messages import get_run_messages + from agno.learn import LearningMachine + from agno.models.openai import OpenAIResponses + from agno.run.agent import RunOutput + from agno.run.base import RunContext + from agno.session import AgentSession + + db = RecordingLearningDb() + machine = LearningMachine(db=db, entity_memory=True) # type: ignore[arg-type] + entity_store = machine.entity_memory_store + assert entity_store is not None + entity_store.remember_about(entity="radar", entity_type="project", facts=["db: Postgres, over Dynamo"]) + + agent = Agent(db=db, learning=machine, model=OpenAIResponses(id="gpt-5.5")) # type: ignore[arg-type] + agent._learning = machine + session = AgentSession(session_id="s1") + run_context = RunContext(run_id="r1", session_id="s1", user_id="u1") + + run_messages = get_run_messages( + agent, + run_response=RunOutput(run_id="r1", session_id="s1"), + run_context=run_context, + input="what did we decide for radar?", + session=session, + ) + assert run_messages.system_message is not None + assert "db: Postgres, over Dynamo" in str(run_messages.system_message.content) + + async def test_stored_entity_reaches_the_system_message_through_aget_run_messages(self) -> None: + """Async twin of the call-site guard.""" + from agno.agent import Agent + from agno.agent._messages import aget_run_messages + from agno.learn import LearningMachine + from agno.models.openai import OpenAIResponses + from agno.run.agent import RunOutput + from agno.run.base import RunContext + from agno.session import AgentSession + + db = RecordingLearningDb() + machine = LearningMachine(db=db, entity_memory=True) # type: ignore[arg-type] + entity_store = machine.entity_memory_store + assert entity_store is not None + await entity_store.aremember_about(entity="radar", entity_type="project", facts=["db: Postgres, over Dynamo"]) + + agent = Agent(db=db, learning=machine, model=OpenAIResponses(id="gpt-5.5")) # type: ignore[arg-type] + agent._learning = machine + session = AgentSession(session_id="s1") + run_context = RunContext(run_id="r1", session_id="s1", user_id="u1") + + run_messages = await aget_run_messages( + agent, + run_response=RunOutput(run_id="r1", session_id="s1"), + run_context=run_context, + input="what did we decide for radar?", + session=session, + ) + assert run_messages.system_message is not None + assert "db: Postgres, over Dynamo" in str(run_messages.system_message.content) + + +async def test_async_far_edge_detach_reaches_an_async_db(tmp_path) -> None: + """The sync helpers no-op against an AsyncBaseDb. + + aforget dropped the near end's edge and left the far end holding a + reciprocal edge to a relationship that no longer exists. + """ + from agno.db.sqlite import AsyncSqliteDb + from agno.learn.config import LearningMode + from agno.learn.stores.entity_memory import EntityMemoryStore + + db = AsyncSqliteDb(db_file=str(tmp_path / "a.db")) + store = EntityMemoryStore( + config=EntityMemoryConfig(db=db, mode=LearningMode.AGENTIC, namespace="global") # type: ignore[arg-type] + ) + await store.aremember_about(entity="quill", entity_type="project", namespace="global") + await store.alink_entities(entity="quill", relation="written_in", related_entity="Rust", namespace="global") + await store.aforget(entity="quill", fact="written_in -> Rust", namespace="global") + + quill = await store.aget(entity_id="quill", entity_type="project", namespace="global") + rust = await store.aget(entity_id="rust", entity_type="unknown", namespace="global") + assert quill is not None and rust is not None + assert quill.relationships == [] + assert rust.relationships == [] + + +class TestAsyncEntityToolsAgainstARealAsyncDb: + """The sync helpers no-op against an AsyncBaseDb. + + Every one of these passed against the in-memory fake while the real async + backend was broken - the fake is synchronous, so it cannot see the gap. + """ + + async def _store(self, tmp_path, name): + from agno.db.sqlite import AsyncSqliteDb + from agno.learn.config import LearningMode + from agno.learn.stores.entity_memory import EntityMemoryStore + + db = AsyncSqliteDb(db_file=str(tmp_path / name)) + return EntityMemoryStore( + config=EntityMemoryConfig(db=db, mode=LearningMode.AGENTIC, namespace="global") # type: ignore[arg-type] + ) + + async def test_ambiguous_name_refuses(self, tmp_path) -> None: + store = await self._store(tmp_path, "a.db") + await store.aremember_about(entity="Harbor", entity_type="project", namespace="global") + await store.aremember_about(entity="Harbor", entity_type="company", namespace="global") + + assert "matches more than one entity" in await store.aforget(entity="Harbor", namespace="global") + assert "matches more than one entity" in await store.alink_entities( + entity="Harbor", relation="uses", related_entity="Postgres", namespace="global" + ) + project = await store.aget(entity_id="harbor", entity_type="project", namespace="global") + assert project is not None and not getattr(project, "archived_at", None) + + async def test_qualified_name_reaches_either_entity(self, tmp_path) -> None: + store = await self._store(tmp_path, "b.db") + await store.aremember_about(entity="Harbor", entity_type="project", namespace="global") + await store.aremember_about(entity="Harbor", entity_type="company", namespace="global") + + message = await store.alink_entities( + entity="project/Harbor", relation="uses", related_entity="Postgres", namespace="global" + ) + assert "project/harbor" in message + assert "Archived company/harbor" in await store.aforget(entity="company/Harbor", namespace="global") + + project = await store.aget(entity_id="harbor", entity_type="project", namespace="global") + company = await store.aget(entity_id="harbor", entity_type="company", namespace="global") + assert project is not None and company is not None + assert [r["entity_id"] for r in project.relationships] == ["postgres"] + assert getattr(company, "archived_at", None) and not getattr(project, "archived_at", None) + + async def test_sibling_signal_reaches_the_async_write(self, tmp_path) -> None: + store = await self._store(tmp_path, "c.db") + await store.aremember_about(entity="Atlas", entity_type="system", namespace="global") + message = await store.aremember_about(entity="Atlas", entity_type="team", namespace="global") + assert "system/atlas already exists under this name" in message + + +class TestTheSeamsThatBrokeBefore: + """Guards for the two seams whose absence is how this release's bugs shipped. + + Both were found by mutation: the shipped code is correct, and every test + stayed green when it was broken. + """ + + def test_stored_entity_reaches_a_teams_system_message(self, tmp_path) -> None: + """The Team half of the write-only bug. + + team/_messages.py threads the run input, the guidance block and the + message kwarg through three separate lines, and deleting any of them + left 1,679 tests green while Team recall died. Agents had a guard from + the first review pass; Teams did not. + """ + from agno.db.sqlite import SqliteDb + from agno.learn import LearningMachine + from agno.models.openai import OpenAIResponses + from agno.run.base import RunContext + from agno.run.team import TeamRunOutput + from agno.session import TeamSession + from agno.team._messages import _get_run_messages + from agno.team.team import Team + + db = SqliteDb(db_file=str(tmp_path / "team.db")) + machine = LearningMachine(db=db, entity_memory=True) + store = machine.entity_memory_store + assert store is not None + store.remember_about(entity="radar", entity_type="project", facts=["db: Postgres, over Dynamo"]) + + team = Team(members=[], db=db, learning=machine, model=OpenAIResponses(id="gpt-5.5")) + team._learning = machine + session = TeamSession(session_id="s1") + run_context = RunContext(run_id="r1", session_id="s1", user_id="u1") + + run_messages = _get_run_messages( + team, + run_response=TeamRunOutput(run_id="r1", session_id="s1"), + run_context=run_context, + session=session, + input_message="what did we decide for radar?", + ) + system = next(m for m in run_messages.messages if m.role == "system") + content = str(system.content) + assert "db: Postgres, over Dynamo" in content + assert "entity_memory_instructions" in content + + def test_the_entity_a_query_names_outranks_the_ones_that_mention_it(self, tmp_path) -> None: + """search_learnings orders by recency and does not rank. + + Five newer entities whose facts merely mention "Harbor" filled the + verified window - they all legitimately match - and Harbor's own row, + older, was never fetched. Recall then injected the wrong entities and + search_entities confirmed them. + """ + import time + + from agno.db.sqlite import SqliteDb + from agno.learn.config import EntityMemoryConfig, LearningMode + from agno.learn.stores.entity_memory import EntityMemoryStore + + store = EntityMemoryStore( + config=EntityMemoryConfig( # type: ignore[arg-type] + db=SqliteDb(db_file=str(tmp_path / "rank.db")), mode=LearningMode.AGENTIC, namespace="global" + ) + ) + store.remember_about(entity="Harbor", entity_type="project", facts=["db: Postgres, over Dynamo"]) + # Rows written in the same second do not order under updated_at DESC. + time.sleep(1.1) + for i in range(8): + store.remember_about( + entity=f"Person {i}", entity_type="person", facts=[f"works on Harbor ingest, shift {i}"] + ) + + hits = store.search(query="Harbor", namespace="global", limit=5) + assert hits and hits[0].entity_id == "harbor" + + recalled = store.recall(message="where does Harbor stand?", namespace="global") + assert recalled is not None + assert "harbor" in [e.entity_id for e in recalled["entities"]] + + @pytest.mark.asyncio + async def test_concurrent_writes_in_one_turn_all_land(self, tmp_path) -> None: + """The write lock, guarded against its own removal. + + Every entity write is a read-modify-write over one JSON row, and a + model's tool calls for one turn are gathered concurrently - the path + AgentOS REST and /mcp take. Without the lock this loses 7 of 8 facts on + a real AsyncSqliteDb, and the fake db cannot see it because it aliases + the content dict. + """ + import asyncio + + from agno.db.sqlite import AsyncSqliteDb + from agno.learn.config import EntityMemoryConfig, LearningMode + from agno.learn.stores.entity_memory import EntityMemoryStore + + db = AsyncSqliteDb(db_file=str(tmp_path / "concurrent.db")) + store = EntityMemoryStore( + config=EntityMemoryConfig(db=db, mode=LearningMode.AGENTIC, namespace="global") # type: ignore[arg-type] + ) + await store.aremember_about(entity="radar", entity_type="project", facts=["seed"], namespace="global") + + await asyncio.gather( + *[ + store.aremember_about(entity="radar", entity_type="project", facts=[f"fact {i}"], namespace="global") + for i in range(8) + ] + ) + + radar = await store.aget(entity_id="radar", entity_type="project", namespace="global") + assert radar is not None + stored = {f["content"] for f in radar.live_facts()} + assert stored == {"seed", *(f"fact {i}" for i in range(8))} diff --git a/libs/agno/tests/unit/team/test_team_learning.py b/libs/agno/tests/unit/team/test_team_learning.py index 31ccbf6164..e2e78d8f77 100644 --- a/libs/agno/tests/unit/team/test_team_learning.py +++ b/libs/agno/tests/unit/team/test_team_learning.py @@ -16,6 +16,7 @@ from agno.agent.agent import Agent from agno.learn.config import LearnedKnowledgeConfig, LearningMode, UserMemoryConfig, UserProfileConfig from agno.learn.machine import LearningMachine from agno.run import RunContext +from agno.run.team import TeamRunOutput from agno.session import TeamSession from agno.team._init import _set_learning_machine from agno.team._messages import aget_system_message, get_system_message @@ -160,6 +161,11 @@ class TestGetSystemMessageLearningContext: user_id="user-456", session_id="sess-123", team_id="my-team", + message=None, + run_context=run_context, + metadata=None, + dependencies=None, + session_state=None, ) def test_learning_context_not_added_when_build_context_returns_none(self, mock_db, mock_model, member_agent): @@ -186,6 +192,11 @@ class TestGetSystemMessageLearningContext: user_id=None, session_id="sess-123", team_id="test-team", + message=None, + run_context=None, + metadata=None, + dependencies=None, + session_state=None, ) def test_learning_context_forwards_empty_user_id(self, mock_db, mock_model, member_agent): @@ -206,6 +217,11 @@ class TestGetSystemMessageLearningContext: user_id="", session_id="sess-123", team_id="test-team", + message=None, + run_context=run_context, + metadata=None, + dependencies=None, + session_state=None, ) def test_custom_system_message_bypasses_learning_context(self, mock_db, mock_model, member_agent): @@ -302,6 +318,11 @@ class TestAgetSystemMessageLearningContext: user_id="async-user", session_id="async-sess", team_id="async-team", + message=None, + run_context=run_context, + metadata=None, + dependencies=None, + session_state=None, ) @pytest.mark.asyncio @@ -345,6 +366,11 @@ class TestAgetSystemMessageLearningContext: user_id=None, session_id="async-sess", team_id="test-team", + message=None, + run_context=None, + metadata=None, + dependencies=None, + session_state=None, ) @@ -701,3 +727,63 @@ class TestLearningContextFalseyValues: assert "None" not in msg.content assert "\n\n\n" not in msg.content team._learning.abuild_context.assert_awaited_once() + + +class TestSystemMessageOverrideCompatibility: + """Team.get_system_message is a public extension point. + + The framework calls the BOUND method, so a subclass written against the + pre-2.8.4 signature is what actually runs - passing a new kwarg + unconditionally would fail every run of such a team. + """ + + def test_legacy_subclass_override_still_runs(self, mock_db, mock_model, member_agent): + from agno.team._messages import _get_run_messages + from agno.team.team import Team + + class BannerTeam(Team): + def get_system_message( + self, + session, + run_context=None, + audio=None, + images=None, + videos=None, + files=None, + tools=None, + add_session_state_to_context=None, + ): + message = super().get_system_message( + session=session, + run_context=run_context, + audio=audio, + images=images, + videos=videos, + files=files, + tools=tools, + add_session_state_to_context=add_session_state_to_context, + ) + if message is not None: + message.content = "[BANNER]\n" + str(message.content) + return message + + team = BannerTeam(db=mock_db, id="t", name="T", members=[member_agent]) + team.model = mock_model + run_messages = _get_run_messages( + team, + run_response=TeamRunOutput(run_id="r", session_id="s"), + run_context=RunContext(run_id="r", session_id="s"), + input_message="hello", + session=TeamSession(session_id="s"), + ) + assert run_messages.system_message is not None + assert str(run_messages.system_message.content).startswith("[BANNER]") + + def test_current_signature_still_receives_the_input(self, mock_db, mock_model, member_agent): + from agno.team._messages import _input_kwarg + from agno.team.team import Team + + team = Team(db=mock_db, id="t2", name="T2", members=[member_agent]) + assert _input_kwarg(team.get_system_message, "hello") == {"input": "hello"} + assert _input_kwarg(lambda session: None, "hello") == {} + assert _input_kwarg(lambda **kw: None, "hello") == {"input": "hello"}