diff --git a/README.md b/README.md index d1c75eac2..aaa9c5532 100755 --- a/README.md +++ b/README.md @@ -4,7 +4,7 @@ Your personal AI that builds memory through screen observation and natural conversation -| ๐ŸŒ [Website](https://mirix.io) | ๐Ÿ“š [Documentation](https://docs.mirix.io) | ๐Ÿ“„ [Paper](https://arxiv.org/abs/2507.07957) | ๐Ÿ’ฌ [Discord](https://discord.gg/S6CeHNrJ) +| ๐ŸŒ [Website](https://mirix.io) | ๐Ÿ“š [Documentation](https://docs.mirix.io) | ๐Ÿ“„ [Paper](https://arxiv.org/abs/2507.07957) | ๐Ÿ’ฌ [Discord](https://discord.gg/S6CeHNrJ) --- @@ -12,7 +12,7 @@ Your personal AI that builds memory through screen observation and natural conve ### Key Features ๐Ÿ”ฅ - **Multi-Agent Memory System:** Six specialized memory components (Core, Episodic, Semantic, Procedural, Resource, Knowledge Vault) managed by dedicated agents -- **Screen Activity Tracking:** Continuous visual data capture and intelligent consolidation into structured memories +- **Screen Activity Tracking:** Continuous visual data capture and intelligent consolidation into structured memories - **Privacy-First Design:** All long-term data stored locally with user-controlled privacy settings - **Advanced Search:** PostgreSQL-native BM25 full-text search with vector similarity support - **Multi-Modal Input:** Text, images, voice, and screen captures processed seamlessly @@ -22,8 +22,8 @@ Your personal AI that builds memory through screen observation and natural conve ``` docker compose up -d --pull always ``` -- Dashboard: http://localhost:5173 -- API: http://localhost:8531 +- Dashboard: http://localhost:5173 +- API: http://localhost:8531 **Step 2: Create an API key in the dashboard (http://localhost:5173) and set as the environmental variable `MIRIX_API_KEY`.** @@ -83,6 +83,25 @@ client.add( {"role": "user", "content": [{"type": "text", "text": "The moon now has a president."}]}, {"role": "assistant", "content": [{"type": "text", "text": "Noted."}]}, ], + session_id="sess-demo-001", +) + +# Include tool activity when you want skills learned from the WORK PROCESS: +# tool errors, retries, and the fix that finally worked are the distiller's +# strongest signals. Pass tool results as role="tool" turns and/or a +# "tool_calls" list on assistant messages (OpenAI shape) โ€” both are preserved +# in the session conversation store. +client.add( + user_id="demo-user", + messages=[ + {"role": "user", "content": [{"type": "text", "text": "Deploy the service."}]}, + {"role": "assistant", "content": "", "tool_calls": [ + {"function": {"name": "kubectl_apply", "arguments": "{\"file\": \"deploy.yaml\"}"}} + ]}, + {"role": "tool", "name": "kubectl_apply", "content": "error: forbidden (missing RBAC)"}, + {"role": "assistant", "content": [{"type": "text", "text": "Fixed the RBAC role and redeployed successfully."}]}, + ], + session_id="sess-demo-001", ) memories = client.retrieve_with_conversation( @@ -123,18 +142,39 @@ Supported `mode` values: | `knowledge` | Knowledge Vault memories | | `experience` | Episodic, Semantic, and Knowledge Vault memories together | -Use `dry_run: true` to fetch counts and inspect what would be processed without applying updates. The optional `model` field can override the auto-dream agent model for that request. +Use `dry_run: true` to fetch counts and inspect what would be processed without applying updates. The optional `model` field can override the auto-dream agent model for that request. For procedural mode, pass the target `meta_agent_id` and optionally `last_n_sessions` to distill recent session experiences. Python client example: ```python result = client.auto_dream( user_id="demo-user", - mode="experience", + mode="procedural", + meta_agent_id=meta_agent.id, + last_n_sessions=3, dry_run=False, ) print(result) ``` +### Upgrading an Existing Database + +Fresh databases get their full schema automatically at startup. **Existing databases require manual migrations** when upgrading to a version that alters tables โ€” startup creates missing tables but never adds columns to existing ones. If migrations are pending, the server refuses to start and names the exact scripts to run. + +For the skill-based procedural memory release, run against your database (in this order): + +```bash +psql "$MIRIX_PG_URI" -f scripts/migrate_procedural_to_skill.sql +psql "$MIRIX_PG_URI" -f scripts/migrate_add_message_session_id.sql +psql "$MIRIX_PG_URI" -f scripts/migrate_add_message_session_id_phase2.sql # outside a transaction (CREATE INDEX CONCURRENTLY) +psql "$MIRIX_PG_URI" -f scripts/migrate_add_conversation_message.sql +psql "$MIRIX_PG_URI" -f scripts/migrate_add_conversation_message_phase2.sql # outside a transaction +psql "$MIRIX_PG_URI" -f scripts/migrate_add_skill_experience.sql +psql "$MIRIX_PG_URI" -f scripts/migrate_add_agent_trigger_state.sql +psql "$MIRIX_PG_URI" -f scripts/migrate_backfill_procedural_scope.sql # backfill filter_tags['scope'] from each skill's owning client +``` + +Each script is idempotent where possible and documents its own preconditions in its header. Run them once, then restart the server. + ## License Mirix is released under the Apache License 2.0. See the [LICENSE](LICENSE) file for more details. @@ -158,7 +198,7 @@ We host weekly discussion sessions where you can: - Get general consultations and support - Connect with the development team and community -**๐Ÿ“… Schedule:** Friday nights, 8-9 PM PST +**๐Ÿ“… Schedule:** Friday nights, 8-9 PM PST **๐Ÿ”— Zoom Link:** [https://ucsd.zoom.us/j/96278791276](https://ucsd.zoom.us/j/96278791276) ### ๐Ÿ“ฑ WeChat Group diff --git a/dashboard/src/pages/dashboard/Memories.tsx b/dashboard/src/pages/dashboard/Memories.tsx index f1940fe97..c6addb997 100644 --- a/dashboard/src/pages/dashboard/Memories.tsx +++ b/dashboard/src/pages/dashboard/Memories.tsx @@ -10,7 +10,10 @@ import { useSearchParams } from 'react-router-dom'; interface MemoryResult { memory_type: string; id: string; + name?: string; summary?: string; + description?: string; + instructions?: string; details?: string; content?: string; caption?: string; @@ -54,9 +57,13 @@ interface SemanticMemory { interface ProceduralMemory { id: string; + name?: string; entry_type?: string; - summary?: string; - steps?: string[]; + description?: string; + instructions?: string; + triggers?: string[]; + examples?: Record[]; + version?: string; created_at?: string; updated_at?: string; } @@ -103,7 +110,7 @@ const DEFAULT_FIELDS: Record = { all: [], episodic: ['summary', 'details'], semantic: ['name', 'summary', 'details'], - procedural: ['summary', 'steps'], + procedural: ['description', 'instructions'], resource: ['summary', 'content'], knowledge_vault: ['caption', 'secret_value'], core: ['label', 'value'], @@ -205,7 +212,7 @@ export const Memories: React.FC = () => { const handleSearch = async (e?: React.FormEvent) => { if (e) e.preventDefault(); if (!selectedUser || !query.trim()) return; - + setLoading(true); setHasSearched(true); try { @@ -441,15 +448,22 @@ export const Memories: React.FC = () => { {bucket.items.map((item) => ( -
- - {item.entry_type || 'Procedure'} - +
+
+ + {item.entry_type || 'Procedure'} + + {item.version && ( + + v{item.version} + + )} +
{item.created_at && {formatDate(item.created_at)}}
-
{item.summary || 'No summary'}
- {item.steps && item.steps.length > 0 && ( +
{item.name || item.description || 'No description'}
+ {item.instructions && ( )}
- {item.steps && item.steps.length > 0 ? ( + {item.name && item.description && ( +

{item.description}

+ )} + {item.instructions ? ( expandedProcedural[item.id] ? ( -
    - {item.steps.map((step, idx) => ( -
  • {step}
    • - ))} -
+

+ {item.instructions} +

) : null ) : ( -

No steps recorded.

+

No instructions recorded.

)} @@ -732,10 +747,14 @@ export const Memories: React.FC = () => { )}

- {memory.summary || memory.caption || 'No summary'} + {memory.memory_type === 'procedural' + ? (memory.name || memory.description || 'No description') + : (memory.summary || memory.caption || 'No summary')}

- {memory.details || memory.content || 'No additional details.'} + {memory.memory_type === 'procedural' + ? (memory.instructions || memory.description || 'No additional details.') + : (memory.details || memory.content || 'No additional details.')}

@@ -797,4 +816,3 @@ export const Memories: React.FC = () => { ); }; - diff --git a/docs/ARCHITECTURE.md b/docs/ARCHITECTURE.md index 4445323a6..ba77cc5d3 100644 --- a/docs/ARCHITECTURE.md +++ b/docs/ARCHITECTURE.md @@ -175,8 +175,6 @@ When `max_chaining_steps` is reached, agents are directed to call their terminal | DELETE | `/memory/episodic/{memory_id}` | | PATCH | `/memory/semantic/{memory_id}` | | DELETE | `/memory/semantic/{memory_id}` | -| PATCH | `/memory/procedural/{memory_id}` | -| DELETE | `/memory/procedural/{memory_id}` | | PATCH | `/memory/resource/{memory_id}` | | DELETE | `/memory/resource/{memory_id}` | | DELETE | `/memory/knowledge_vault/{memory_id}` | diff --git a/docs/architecture.html b/docs/architecture.html index 338f20e1e..b28d9f3f6 100644 --- a/docs/architecture.html +++ b/docs/architecture.html @@ -960,16 +960,6 @@

Memory Endpoints (23)

/memory/semantic/{memory_id}
Delete semantic memory
-
- PATCH - /memory/procedural/{memory_id} -
Update procedural memory
-
-
- DELETE - /memory/procedural/{memory_id} -
Delete procedural memory
-
PATCH /memory/resource/{memory_id} diff --git a/mirix/agent/agent.py b/mirix/agent/agent.py index a0d2a94cc..348b5f1f7 100644 --- a/mirix/agent/agent.py +++ b/mirix/agent/agent.py @@ -946,6 +946,18 @@ async def _handle_ai_response( if response_message_id is not None: assert response_message_id.startswith("message-"), response_message_id + # Only the chat_agent inherits the triggering input's session_id so its + # synthesized assistant/tool messages stay in the same session. For all + # non-chat agents (meta_memory_agent + the memory sub-agents) the + # synthesized messages carry NO session_id: session_id identifies an + # external conversation, never the memory-production machinery. Context + # reconstruction uses message_ids (not session_id), so nulling this is safe. + input_session_id = ( + getattr(input_message, "session_id", None) + if self.agent_state.is_type(AgentType.chat_agent) + else None + ) + messages = [] # append these to the history when done function_name = None @@ -974,7 +986,7 @@ async def _handle_ai_response( and len(response_message.tool_calls) > 1 ): self.logger.info( - "Memory agent %s returned %d tool call(s); executing each sequentially", + "Agent %s returned %d tool call(s); executing each sequentially", self.agent_state.agent_type, len(response_message.tool_calls), ) @@ -988,6 +1000,7 @@ async def _handle_ai_response( agent_id=self.agent_state.id, model=self.model, openai_message_dict=response_message.model_dump(), + session_id=input_session_id, ) ) # extend conversation with assistant's reply @@ -1037,6 +1050,7 @@ async def _handle_ai_response( "content": function_response, "tool_call_id": tool_call_id, }, + session_id=input_session_id, ) ) # extend conversation with function response self.interface.function_message(f"Error: {error_msg}", msg_obj=messages[-1]) @@ -1062,6 +1076,7 @@ async def _handle_ai_response( "content": function_response, "tool_call_id": tool_call_id, }, + session_id=input_session_id, ) ) # extend conversation with function response self.interface.function_message(f"Error: {error_msg}", msg_obj=messages[-1]) @@ -1111,6 +1126,7 @@ async def _handle_ai_response( "content": function_response, "tool_call_id": tool_call_id, }, + session_id=input_session_id, ) ) self.interface.function_message(f"Validation Error: {validation_error}", msg_obj=messages[-1]) @@ -1197,6 +1213,7 @@ async def _handle_ai_response( "content": function_response, "tool_call_id": tool_call_id, }, + session_id=input_session_id, ) ) # extend conversation with function response self.interface.function_message(f"Ran {function_name}()", msg_obj=messages[-1]) @@ -1218,6 +1235,7 @@ async def _handle_ai_response( "content": function_response, "tool_call_id": tool_call_id, }, + session_id=input_session_id, ) ) # extend conversation with function response self.interface.function_message(f"Ran {function_name}()", msg_obj=messages[-1]) @@ -1237,6 +1255,7 @@ async def _handle_ai_response( "content": function_response, "tool_call_id": tool_call_id, }, + session_id=input_session_id, ) ) # extend conversation with function response self.interface.function_message(f"Ran {function_name}()", msg_obj=messages[-1]) @@ -1320,17 +1339,21 @@ async def _handle_ai_response( ) if memory_item: memory_item = memory_item[0] + # Phase 1 of skill-evolve replaced the legacy + # summary+steps shape with name/description/ + # instructions/version. Surface the new fields. memory_item_str = "" - memory_item_str += "[Procedural Memory ID]: " + memory_item.id + "\n" + memory_item_str += "[Skill ID]: " + memory_item.id + "\n" + memory_item_str += "[Name]: " + memory_item.name + "\n" memory_item_str += "[Entry Type]: " + memory_item.entry_type + "\n" - memory_item_str += "[Summary]: " + (memory_item.summary or "N/A") + "\n" - memory_item_str += "[Steps]: " + "; ".join(memory_item.steps) + "\n" + memory_item_str += "[Description]: " + (memory_item.description or "N/A") + "\n" + memory_item_str += "[Instructions]: " + (memory_item.instructions or "N/A") + "\n" + memory_item_str += "[Version]: " + str(getattr(memory_item, "version", "0.1.0")) + "\n" + ts = memory_item.last_modify.get("timestamp") if isinstance(memory_item.last_modify, dict) else None + ts_str = ts.strftime("%Y-%m-%d %H:%M:%S") if hasattr(ts, "strftime") else str(ts) + op = memory_item.last_modify.get("operation") if isinstance(memory_item.last_modify, dict) else "?" memory_item_str += ( - "[Last Modified]: " - + memory_item.last_modify["operation"] - + " at " - + memory_item.last_modify["timestamp"].strftime("%Y-%m-%d %H:%M:%S") - + "\n" + "[Last Modified]: " + str(op) + " at " + ts_str + "\n" ) memory_item_str = memory_item_str.strip() @@ -1424,6 +1447,7 @@ async def _handle_ai_response( "role": "user", "content": message_content, }, + session_id=input_session_id, ) # persist the message to the database @@ -1446,8 +1470,17 @@ async def _handle_ai_response( message_ids=message_ids, actor=self.actor, ) + # Retention: the memory-production machinery (meta_memory_agent + # + the 6 memory sub-agents) is a pure transient producer again. + # The canonical learnable conversation lives in the dedicated + # Conversation Message Store, so the meta agent no longer needs to + # retain a transcript for the distiller. All non-chat agents keep + # the legacy full delete (retain=0) of detached extraction scratch. await self.message_manager.delete_detached_messages_for_agent( - agent_id=self.agent_state.id, actor=self.actor + agent_id=self.agent_state.id, + actor=self.actor, + retain_last_n_sessions=0, + user_id=self.user_id, ) # Clear all messages since they were manually added to the conversation history @@ -1466,6 +1499,7 @@ async def _handle_ai_response( agent_id=self.agent_state.id, model=self.model, openai_message_dict=response_message.model_dump(), + session_id=input_session_id, ) ) # extend conversation with assistant's reply self.interface.internal_monologue(response_message.content, msg_obj=messages[-1]) @@ -1542,6 +1576,22 @@ async def step( first_input_message = input_messages[0] if isinstance(input_messages, list) else input_messages + # Derive the session_id for this step once so all synthesized messages + # (heartbeats, warnings, meta-memory bootstrap, summaries) inherit it. + # Only the chat_agent inherits it; for every non-chat agent the + # synthesized messages carry NO session_id (session_id identifies an + # external conversation, never the meta agent's own bookkeeping). + step_session_id = ( + getattr(first_input_message, "session_id", None) + if self.agent_state.is_type(AgentType.chat_agent) + else None + ) + # Expose on self so helpers called without an explicit session_id + # (e.g. summarize_messages_inplace) can pick it up. Other entrypoints + # that skip step() (e.g. step_user_message) must set/reset this + # themselves โ€” see step_user_message for the save/restore pattern. + self._current_step_session_id = step_session_id + # Convert MessageCreate objects to Message objects if not isinstance(input_messages, list): input_messages = [input_messages] @@ -1630,6 +1680,7 @@ async def step( sender_id=None, group_id=None, filter_tags=self.filter_tags, + session_id=step_session_id, ), self.agent_state.id, wrap_user_message=False, @@ -1678,6 +1729,7 @@ async def step( "role": "user", "content": warning_content, }, + session_id=step_session_id, ) continue # give agent one more chance to respond elif max_chaining_steps is not None and counter > max_chaining_steps: @@ -1695,6 +1747,7 @@ async def step( "role": "user", # TODO: change to system? "content": get_token_limit_warning(), }, + session_id=step_session_id, ) continue # always chain elif function_failed: @@ -1706,6 +1759,7 @@ async def step( "role": "user", # TODO: change to system? "content": get_contine_chaining(FUNC_FAILED_HEARTBEAT_MESSAGE), }, + session_id=step_session_id, ) continue # always chain elif continue_chaining: @@ -1717,6 +1771,7 @@ async def step( "role": "user", # TODO: change to system? "content": get_contine_chaining(REQ_HEARTBEAT_MESSAGE), }, + session_id=step_session_id, ) continue # always chain # Mirix no-op / yield @@ -1897,35 +1952,20 @@ async def build_system_prompt_with_memories( "text": resource_memory, } - # Retrieve procedural memory - # Owning agents need IDs for merge/update operations, so always retrieve fresh - is_owning_agent = self.agent_state.is_type(AgentType.procedural_memory_agent, AgentType.reflexion_agent) - if is_owning_agent or "procedural" not in retrieved_memories: - current_procedural_memory = await self.procedural_memory_manager.list_procedures( - agent_state=self.agent_state, - user=self.user, - query=key_words, - embedded_text=embedded_text, - search_field="summary", - search_method=search_method, - limit=MAX_RETRIEVAL_LIMIT_IN_SYSTEM, - timezone_str=timezone_str, - ) - procedural_memory = "" - if len(current_procedural_memory) > 0: - for idx, procedure in enumerate(current_procedural_memory): - if is_owning_agent: - procedural_memory += f"[Procedure ID: {procedure.id}] Entry Type: {procedure.entry_type}; Summary: {procedure.summary}\n" - else: - procedural_memory += ( - f"[{idx}] Entry Type: {procedure.entry_type}; Summary: {procedure.summary}\n" - ) - procedural_memory = procedural_memory.strip() - retrieved_memories["procedural"] = { - "total_number_of_items": await self.procedural_memory_manager.get_total_number_of_items(user=self.user), - "current_count": len(current_procedural_memory), - "text": procedural_memory, - } + # Procedural memory (skills) is intentionally NOT injected into any + # agent's system prompt. Skills are a RETRIEVAL TARGET, not ambient + # context: + # * the procedural_memory_agent surveys them ON DEMAND via its + # `skill_list` / `skill_read` tools (its system prompt explicitly + # instructs "Survey existing skills: Call skill_list()"), and + # * task-executing / external agents fetch them through the unified + # search interface (GET /memory/search?memory_type=procedural, or + # the search_in_memory tool). + # Passively broadcasting the whole skill list into every agent's prompt + # was redundant for the owner (it uses the tool) and pure token noise + # for the pure-extractor agents (episodic/semantic/resource/...), which + # never act on skills. Removing it also drops a per-prompt-build + # retrieval from the hot path. # Retrieve semantic memory # Owning agents need IDs for merge/update operations, so always retrieve fresh @@ -1995,7 +2035,6 @@ def build_system_prompt(self, retrieved_memories: dict) -> str: episodic_memory = retrieved_memories["episodic"] resource_memory = retrieved_memories["resource"] semantic_memory = retrieved_memories["semantic"] - procedural_memory = retrieved_memories["procedural"] knowledge_vault = retrieved_memories["knowledge_vault"] system_prompt = template.format( @@ -2046,15 +2085,10 @@ def build_system_prompt(self, retrieved_memories: dict) -> str: + "\n\n" ) - # Add procedural memory with counts - procedural_total = procedural_memory["total_number_of_items"] if procedural_memory else 0 - procedural_text = procedural_memory["text"] if procedural_memory else "" - procedural_count = procedural_memory["current_count"] if procedural_memory else 0 - system_prompt += ( - f"\n ({procedural_count} out of {procedural_total} Items):\n" - + (procedural_text if procedural_text else "Empty") - + "\n" - ) + # NOTE: no block โ€” skills are retrieved on demand + # (procedural_memory_agent via skill_list; consumers via /memory/search), + # not broadcast into every prompt. See the procedural-retrieval removal + # note earlier in this method. return system_prompt @@ -2566,7 +2600,12 @@ async def inner_step( ) raise e - async def step_user_message(self, user_message_str: str, **kwargs) -> AgentStepResponse: + async def step_user_message( + self, + user_message_str: str, + session_id: Optional[str] = None, + **kwargs, + ) -> AgentStepResponse: """Takes a basic user message string, turns it into a stringified JSON with extra metadata, then sends it to the agent Example: @@ -2592,18 +2631,39 @@ async def step_user_message(self, user_message_str: str, **kwargs) -> AgentStepR "name": name, } + # session_id identifies an external conversation and is only stamped on the + # chat_agent's messages. For any non-chat agent reaching this entrypoint the + # message (and the step-level session context below) must carry NO session_id, + # matching the guard applied to input_session_id/step_session_id in step(). + effective_session_id = ( + session_id if self.agent_state.is_type(AgentType.chat_agent) else None + ) + # Create the associated Message object (in the database) assert self.agent_state.created_by_id is not None, "User ID is not set" user_message = Message.dict_to_message( agent_id=self.agent_state.id, model=self.model, openai_message_dict=openai_message_dict, + session_id=effective_session_id, # created_at=timestamp, ) - return await self.inner_step(messages=[user_message], **kwargs) + # Seed the step-level session context so any pre-persist summarization + # triggered inside inner_step() / _get_ai_reply() stamps the correct + # session_id on the summary message (Codex review v3, Important). + prev_session_id = getattr(self, "_current_step_session_id", None) + self._current_step_session_id = effective_session_id + try: + return await self.inner_step(messages=[user_message], **kwargs) + finally: + self._current_step_session_id = prev_session_id - async def summarize_messages_inplace(self, existing_file_uris: Optional[List[str]] = None): + async def summarize_messages_inplace( + self, + existing_file_uris: Optional[List[str]] = None, + session_id: Optional[str] = None, + ): in_context_messages = await self.agent_manager.get_in_context_messages( agent_state=self.agent_state, actor=self.actor, user=self.user ) @@ -2679,13 +2739,31 @@ async def summarize_messages_inplace(self, existing_file_uris: Optional[List[str ) packed_summary_message = {"role": "user", "content": summary_message} - # Prepend the summary + # Prepend the summary. The summary is a synthesized message, so it follows + # the same rule as every other synthesized message: only the chat_agent + # carries a session_id; for any non-chat (memory-production) agent the + # summary carries NONE. Within the chat_agent, preference order is: + # 1. explicit `session_id` argument (caller-provided, e.g. step_user_message) + # 2. _current_step_session_id stashed by Agent.step() for the current step + # 3. the latest in-context message's session_id (inherits whatever session + # the conversation being summarized is already in) + summary_session_id = None + if self.agent_state.is_type(AgentType.chat_agent): + summary_session_id = session_id + if summary_session_id is None: + summary_session_id = getattr(self, "_current_step_session_id", None) + if summary_session_id is None: + for m in reversed(in_context_messages): + if getattr(m, "session_id", None): + summary_session_id = m.session_id + break self.agent_state = await self.agent_manager.prepend_to_in_context_messages( messages=[ Message.dict_to_message( agent_id=self.agent_state.id, model=self.model, openai_message_dict=packed_summary_message, + session_id=summary_session_id, ) ], agent_id=self.agent_state.id, diff --git a/mirix/agent/tool_validators.py b/mirix/agent/tool_validators.py index 358f31002..77b22be2f 100644 --- a/mirix/agent/tool_validators.py +++ b/mirix/agent/tool_validators.py @@ -188,48 +188,96 @@ def validate_resource_memory_update(function_name: str, args: dict) -> Optional[ # ============================================================ -# Procedural Memory Validators +# Skill Validators # ============================================================ -@register_validator("procedural_memory_insert") -def validate_procedural_memory_insert(function_name: str, args: dict) -> Optional[str]: - """Validate procedural_memory_insert arguments.""" - items = args.get("items", []) - for i, item in enumerate(items): - if not item.get("summary", "").strip(): - return ( - f"Validation error: 'summary' field in item {i} cannot be empty. " - "Please provide a descriptive summary of this procedure." - ) - steps = item.get("steps", []) - if not steps or all(not s.strip() for s in steps): - return ( - f"Validation error: 'steps' field in item {i} cannot be empty. " - "Please provide at least one non-empty step." - ) +@register_validator("skill_create") +def validate_skill_create(function_name: str, args: dict) -> Optional[str]: + """Validate skill_create arguments.""" + from mirix.schemas.procedural_memory import ( + SKILL_ENTRY_TYPES, + SKILL_MAX_DESCRIPTION_LEN, + SKILL_MAX_INSTRUCTIONS_LEN, + SKILL_MAX_NAME_LEN, + ) + + name = args.get("name", "") + if not name.strip(): + return "Validation error: 'name' cannot be empty." + if len(name) > SKILL_MAX_NAME_LEN: + return f"Validation error: 'name' exceeds max length {SKILL_MAX_NAME_LEN}." + description = args.get("description", "") + if not description.strip(): + return "Validation error: 'description' cannot be empty." + if len(description) > SKILL_MAX_DESCRIPTION_LEN: + return f"Validation error: 'description' exceeds max length {SKILL_MAX_DESCRIPTION_LEN}." + instructions = args.get("instructions", "") + if not instructions.strip(): + return "Validation error: 'instructions' cannot be empty." + if len(instructions) > SKILL_MAX_INSTRUCTIONS_LEN: + return f"Validation error: 'instructions' exceeds max length {SKILL_MAX_INSTRUCTIONS_LEN}." + entry_type = args.get("entry_type", "") + if not entry_type: + return "Validation error: 'entry_type' cannot be empty." + if entry_type not in SKILL_ENTRY_TYPES: + return ( + f"Validation error: 'entry_type' must be one of: {sorted(SKILL_ENTRY_TYPES)}." + ) return None -@register_validator("procedural_memory_update") -def validate_procedural_memory_update(function_name: str, args: dict) -> Optional[str]: - """Validate procedural_memory_update arguments.""" - items = args.get("new_items", []) - for i, item in enumerate(items): - if not item.get("summary", "").strip(): - return ( - f"Validation error: 'summary' field in new_items[{i}] cannot be empty. " - "Please provide a descriptive summary of this procedure." - ) - steps = item.get("steps", []) - if not steps or all(not s.strip() for s in steps): +@register_validator("skill_edit") +def validate_skill_edit(function_name: str, args: dict) -> Optional[str]: + """Validate skill_edit arguments.""" + from mirix.schemas.procedural_memory import ( + SKILL_ENTRY_TYPES, + SKILL_MAX_DESCRIPTION_LEN, + SKILL_MAX_INSTRUCTIONS_LEN, + SKILL_MAX_NAME_LEN, + ) + + if not args.get("skill_id", "").strip(): + return "Validation error: 'skill_id' cannot be empty." + field = args.get("field", "") + if not field: + return "Validation error: 'field' cannot be empty." + valid_fields = {"name", "description", "instructions", "entry_type", "triggers", "examples"} + if field not in valid_fields: + return f"Validation error: 'field' must be one of: {', '.join(sorted(valid_fields))}." + text_fields = {"name", "description", "instructions"} + if field in text_fields: + if not args.get("old_text"): + return f"Validation error: 'old_text' is required for text field '{field}'." + if args.get("new_text") is None: + return f"Validation error: 'new_text' is required for text field '{field}'." + new_text = args.get("new_text") or "" + caps = { + "name": SKILL_MAX_NAME_LEN, + "description": SKILL_MAX_DESCRIPTION_LEN, + "instructions": SKILL_MAX_INSTRUCTIONS_LEN, + } + if len(new_text) > caps[field]: + return f"Validation error: 'new_text' for field '{field}' exceeds max length {caps[field]}." + else: + value = args.get("value") + if value is None: + return f"Validation error: 'value' is required for field '{field}'." + if field == "entry_type" and value not in SKILL_ENTRY_TYPES: return ( - f"Validation error: 'steps' field in new_items[{i}] cannot be empty. " - "Please provide at least one non-empty step." + f"Validation error: 'entry_type' must be one of: {sorted(SKILL_ENTRY_TYPES)}." ) return None +@register_validator("skill_delete") +def validate_skill_delete(function_name: str, args: dict) -> Optional[str]: + """Validate skill_delete arguments.""" + if not args.get("skill_id", "").strip(): + return "Validation error: 'skill_id' cannot be empty." + return None + + # ============================================================ # Knowledge Vault Validators # ============================================================ diff --git a/mirix/client/remote_client.py b/mirix/client/remote_client.py index 9bbf83838..ffa1e9220 100644 --- a/mirix/client/remote_client.py +++ b/mirix/client/remote_client.py @@ -26,7 +26,7 @@ from mirix.schemas.file import FileMetadata from mirix.schemas.llm_config import LLMConfig from mirix.schemas.memory import ArchivalMemorySummary, Memory, RecallMemorySummary -from mirix.schemas.message import Message, MessageCreate +from mirix.schemas.message import Message, MessageCreate, _validate_session_id from mirix.schemas.mirix_response import MirixResponse from mirix.schemas.organization import Organization from mirix.schemas.sandbox_config import ( @@ -1257,6 +1257,7 @@ async def add( block_filter_tags_update_mode: Optional[str] = "merge", use_cache: bool = True, occurred_at: Optional[str] = None, + session_id: Optional[str] = None, async_add: bool = True, headers: Optional[Dict[str, str]] = None, ) -> Dict[str, Any]: @@ -1274,6 +1275,13 @@ async def add( {"role": "user", "content": [{"type": "text", "text": "..."}]}, {"role": "assistant", "content": [{"type": "text", "text": "..."}]} ] + Tool activity may be included and is used by skill + distillation (tool errors/retries are strong learning + signals): pass tool results as {"role": "tool", + "name": "...", "content": "..."} turns, and/or a + "tool_calls" list on an assistant message (OpenAI shape) โ€” + both are preserved in the session conversation store when + session_id is set. chaining: Enable/disable chaining (default: True) verbose: If True, enable verbose output during memory processing filter_tags: Optional dict of tags for filtering and categorization. @@ -1286,6 +1294,9 @@ async def add( If provided, episodic memories will use this timestamp instead of current time. Format: "2025-11-18T10:30:00" or "2025-11-18T10:30:00+00:00" Example: "2025-11-18T15:30:00" + session_id: Optional external conversation session id. Required for + procedural session distillation and sent as a top-level + request field. headers: Optional headers dict to include in the request. Useful for passing per-request authentication tokens. Example: {"Authorization": "Bearer token123"} @@ -1312,7 +1323,7 @@ async def add( ... {"role": "assistant", "content": [{"type": "text", "text": "That's great!"}]} ... ], ... verbose=True, - ... filter_tags={"session_id": "sess-789"}, + ... session_id="sess-789", ... occurred_at="2025-11-18T15:30:00" ... ) >>> logger.debug(response) @@ -1334,6 +1345,14 @@ async def add( if occurred_at is not None: _validate_occurred_at(occurred_at) # Raises ValueError if invalid + if session_id is not None: + _validate_session_id(session_id) + if isinstance(filter_tags, dict) and filter_tags.get("session_id") not in (None, session_id): + raise ValueError( + "session_id and filter_tags['session_id'] must agree; " + f"got {session_id!r} vs {filter_tags.get('session_id')!r}" + ) + await self._ensure_user_exists(user_id, headers=headers) # Prepare request data - org is determined from API key on server side @@ -1345,6 +1364,9 @@ async def add( "verbose": verbose, } + if session_id is not None: + request_data["session_id"] = session_id + if filter_tags is not None: request_data["filter_tags"] = filter_tags @@ -1527,7 +1549,7 @@ async def search( query: str, memory_type: str = "all", search_field: str = "null", - search_method: str = "embedding", + search_method: Optional[str] = None, limit: int = 10, filter_tags: Optional[Dict[str, Any]] = None, similarity_threshold: Optional[float] = None, @@ -1554,11 +1576,15 @@ async def search( search_field: Field to search in. Options vary by memory type: - episodic: "summary", "details" - resource: "summary", "content" - - procedural: "summary", "steps" + - procedural: "description", "instructions" - knowledge_vault: "caption", "secret_value" - semantic: "name", "summary", "details" - For "all": use "null" (default) - search_method: Search method. Options: "bm25" (default), "embedding" + search_method: Search method. Options: "embedding", "bm25", and + "hybrid" (procedural memory only; BM25 + embedding + fused with Reciprocal Rank Fusion). Omit (None, the + default) to use the server's per-type default: + "hybrid" for procedural, "embedding" otherwise. limit: Maximum number of results per memory type (default: 10) filter_tags: Optional filter tags for additional filtering (scope added automatically) similarity_threshold: Optional similarity threshold for embedding search (0.0-2.0). @@ -1647,10 +1673,15 @@ async def search( "query": query, "memory_type": memory_type, "search_field": search_field, - "search_method": search_method, "limit": limit, } + # Omit search_method when unset so the server applies its per-type + # default (procedural -> hybrid, everything else -> embedding). + # Sending a value here would override that resolution. + if search_method is not None: + params["search_method"] = search_method + # Add filter_tags if provided if filter_tags: import json @@ -1677,7 +1708,7 @@ async def search_all_users( query: str, memory_type: str = "all", search_field: str = "null", - search_method: str = "embedding", + search_method: Optional[str] = None, limit: int = 10, client_id: Optional[str] = None, filter_tags: Optional[Dict[str, Any]] = None, @@ -1702,11 +1733,15 @@ async def search_all_users( search_field: Field to search in. Options vary by memory type: - episodic: "summary", "details" - resource: "summary", "content" - - procedural: "summary", "steps" + - procedural: "description", "instructions" - knowledge_vault: "caption", "secret_value" - semantic: "name", "summary", "details" - For "all": use "null" (default) - search_method: Search method. Options: "bm25" (default), "embedding" + search_method: Search method. Options: "embedding", "bm25", and + "hybrid" (procedural memory only; BM25 + embedding + fused with Reciprocal Rank Fusion). Omit (None, the + default) to use the server's per-type default: + "hybrid" for procedural, "embedding" otherwise. limit: Maximum results per memory type (total across all users) client_id: Optional client ID (uses its org_id and scope for filtering) filter_tags: Optional additional filter tags (scope added automatically) @@ -1784,10 +1819,14 @@ async def search_all_users( "query": query, "memory_type": memory_type, "search_field": search_field, - "search_method": search_method, "limit": limit, } + # Omit search_method when unset so the server applies its per-type + # default (procedural -> hybrid, everything else -> embedding). + if search_method is not None: + params["search_method"] = search_method + # Add client_id if provided (server will use this client's org_id) if client_id: params["client_id"] = client_id @@ -1860,6 +1899,8 @@ async def auto_dream( mode: str = "experience", dry_run: bool = False, model: Optional[str] = None, + meta_agent_id: Optional[str] = None, + last_n_sessions: Optional[int] = None, headers: Optional[Dict[str, str]] = None, ) -> Dict[str, Any]: """ @@ -1875,6 +1916,9 @@ async def auto_dream( episodic, semantic, and knowledge together. dry_run: If true, fetch and count memories without invoking the agent. model: Optional model override for this auto-dream run. + meta_agent_id: Optional meta memory agent id to run against. + last_n_sessions: For procedural mode, how many recent retained + sessions to distill. Returns: AutoDreamResponse dict with stats on removed/merged/resolved items @@ -1884,6 +1928,10 @@ async def auto_dream( body: Dict[str, Any] = {"mode": mode, "dry_run": dry_run} if model is not None: body["model"] = model + if meta_agent_id is not None: + body["meta_agent_id"] = meta_agent_id + if last_n_sessions is not None: + body["last_n_sessions"] = last_n_sessions return await self._request( "POST", diff --git a/mirix/constants.py b/mirix/constants.py index 966251288..a234ee119 100644 --- a/mirix/constants.py +++ b/mirix/constants.py @@ -1,132 +1,147 @@ -import os -from logging import CRITICAL, DEBUG, ERROR, INFO, NOTSET, WARN, WARNING - -# ============================================================================ -# Client Constants - Used by both client and server -# ============================================================================ - -# Default organization and admin user IDs (needed by schemas) -DEFAULT_ORG_ID = "org-00000000-0000-4000-8000-000000000000" -ADMIN_USER_ID = "user-00000000-0000-4000-8000-000000000000" - -# Embedding constants -MAX_EMBEDDING_DIM = 4096 # maximum supported embedding size - do NOT change or else DBs will need to be reset -DEFAULT_EMBEDDING_CHUNK_SIZE = 300 -MIN_CONTEXT_WINDOW = 4096 - -# Memory limits -CORE_MEMORY_BLOCK_CHAR_LIMIT: int = 5000 - -# Function/Tool constants -FUNCTION_RETURN_CHAR_LIMIT = 60000 # ~300 words -TOOL_CALL_ID_MAX_LEN = 29 - -# Tool module names -COMPOSIO_TOOL_TAG_NAME = "composio" -MIRIX_CORE_TOOL_MODULE_NAME = "mirix.functions.function_sets.base" -MIRIX_MEMORY_TOOL_MODULE_NAME = "mirix.functions.function_sets.memory_tools" -MIRIX_EXTRA_TOOL_MODULE_NAME = "mirix.functions.function_sets.extras" - -# Message defaults -DEFAULT_MESSAGE_TOOL = "send_message" -DEFAULT_MESSAGE_TOOL_KWARG = "message" - -# LLM model token limits -LLM_MAX_TOKENS = { - "DEFAULT": 8192, - ## OpenAI models: https://platform.openai.com/docs/models/overview - "chatgpt-4o-latest": 128000, - "gpt-4o-2024-08-06": 128000, - "gpt-4-turbo-preview": 128000, - "gpt-4o": 128000, - "gpt-3.5-turbo-instruct": 16385, - "gpt-4-0125-preview": 128000, - "gpt-3.5-turbo-0125": 16385, - "gpt-4-turbo-2024-04-09": 128000, - "gpt-4-turbo": 8192, - "gpt-4o-2024-05-13": 128000, - "gpt-4o-mini": 128000, - "gpt-4o-mini-2024-07-18": 128000, - "gpt-4-1106-preview": 128000, - "gpt-4": 8192, - "gpt-4-32k": 32768, - "gpt-4-0613": 8192, - "gpt-4-32k-0613": 32768, - "gpt-4-0314": 8192, # legacy - "gpt-4-32k-0314": 32768, # legacy - "gpt-3.5-turbo-1106": 16385, - "gpt-3.5-turbo": 4096, - "gpt-3.5-turbo-16k": 16385, - "gpt-3.5-turbo-0613": 4096, # legacy - "gpt-3.5-turbo-16k-0613": 16385, # legacy - "gpt-3.5-turbo-0301": 4096, # legacy -} - -# ============================================================================ -# Server-Only Constants -# ============================================================================ - -MIRIX_DIR = os.path.join(os.path.expanduser("~"), ".mirix") -MIRIX_DIR_TOOL_SANDBOX = os.path.join(MIRIX_DIR, "tool_sandbox_dir") - -ADMIN_PREFIX = "/v1/admin" -API_PREFIX = "/v1" -OPENAI_API_PREFIX = "/openai" - -COMPOSIO_ENTITY_ENV_VAR_KEY = "COMPOSIO_ENTITY" - -# String in the error message for when the context window is too large -# Example full message: -# This model's maximum context length is 8192 tokens. However, your messages resulted in 8198 tokens (7450 in the messages, 748 in the functions). Please reduce the length of the messages or functions. -OPENAI_CONTEXT_WINDOW_ERROR_SUBSTRING = "maximum context length" - -# System prompt templating -IN_CONTEXT_MEMORY_KEYWORD = "CORE_MEMORY" - -MAX_CHAINING_STEPS = int(os.getenv("MAX_CHAINING_STEPS", "10")) -MAX_RETRIEVAL_LIMIT_IN_SYSTEM = 10 - -# tokenizers -EMBEDDING_TO_TOKENIZER_MAP = { - "text-embedding-3-small": "cl100k_base", -} -EMBEDDING_TO_TOKENIZER_DEFAULT = "cl100k_base" - - -DEFAULT_MIRIX_MODEL = "gpt-4" # TODO: fixme -DEFAULT_PERSONA = "sam_pov" -DEFAULT_HUMAN = "basic" -DEFAULT_PRESET = "memgpt_chat" - -# Base tools that cannot be edited, as they access agent state directly -# Note that we don't include "conversation_search_date" for now -BASE_TOOLS = [ - "send_intermediate_message", - "conversation_search", - "search_in_memory", - "list_memory_within_timerange", -] -# Base memory tools CAN be edited, and are added by default by the server -CORE_MEMORY_TOOLS = ["core_memory_append", "core_memory_rewrite"] -EPISODIC_MEMORY_TOOLS = [ - "episodic_memory_insert", - "episodic_memory_merge", - "episodic_memory_replace", - "check_episodic_memory", -] -PROCEDURAL_MEMORY_TOOLS = ["procedural_memory_insert", "procedural_memory_update"] -RESOURCE_MEMORY_TOOLS = ["resource_memory_insert", "resource_memory_update"] -KNOWLEDGE_VAULT_TOOLS = ["knowledge_vault_insert", "knowledge_vault_update"] -SEMANTIC_MEMORY_TOOLS = [ - "semantic_memory_insert", - "semantic_memory_update", - "check_semantic_memory", -] -CHAT_AGENT_TOOLS = [] -EXTRAS_TOOLS = ["web_search", "fetch_and_read_pdf"] -MCP_TOOLS = [] -META_MEMORY_TOOLS = ["trigger_memory_update"] -SEARCH_MEMORY_TOOLS = ["search_in_memory", "list_memory_within_timerange"] +import os +from logging import CRITICAL, DEBUG, ERROR, INFO, NOTSET, WARN, WARNING + +# ============================================================================ +# Client Constants - Used by both client and server +# ============================================================================ + +# Default organization and admin user IDs (needed by schemas) +DEFAULT_ORG_ID = "org-00000000-0000-4000-8000-000000000000" +ADMIN_USER_ID = "user-00000000-0000-4000-8000-000000000000" + +# Embedding constants +MAX_EMBEDDING_DIM = 4096 # maximum supported embedding size - do NOT change or else DBs will need to be reset +DEFAULT_EMBEDDING_CHUNK_SIZE = 300 +MIN_CONTEXT_WINDOW = 4096 + +# Memory limits +CORE_MEMORY_BLOCK_CHAR_LIMIT: int = 5000 + +# Function/Tool constants +FUNCTION_RETURN_CHAR_LIMIT = 60000 # ~300 words +TOOL_CALL_ID_MAX_LEN = 29 + +# Tool module names +COMPOSIO_TOOL_TAG_NAME = "composio" +MIRIX_CORE_TOOL_MODULE_NAME = "mirix.functions.function_sets.base" +MIRIX_MEMORY_TOOL_MODULE_NAME = "mirix.functions.function_sets.memory_tools" +MIRIX_EXTRA_TOOL_MODULE_NAME = "mirix.functions.function_sets.extras" + +# Message defaults +DEFAULT_MESSAGE_TOOL = "send_message" +DEFAULT_MESSAGE_TOOL_KWARG = "message" + +# LLM model token limits +LLM_MAX_TOKENS = { + "DEFAULT": 8192, + ## OpenAI models: https://platform.openai.com/docs/models/overview + "chatgpt-4o-latest": 128000, + "gpt-4o-2024-08-06": 128000, + "gpt-4-turbo-preview": 128000, + "gpt-4o": 128000, + "gpt-3.5-turbo-instruct": 16385, + "gpt-4-0125-preview": 128000, + "gpt-3.5-turbo-0125": 16385, + "gpt-4-turbo-2024-04-09": 128000, + "gpt-4-turbo": 8192, + "gpt-4o-2024-05-13": 128000, + "gpt-4o-mini": 128000, + "gpt-4o-mini-2024-07-18": 128000, + "gpt-4-1106-preview": 128000, + "gpt-4": 8192, + "gpt-4-32k": 32768, + "gpt-4-0613": 8192, + "gpt-4-32k-0613": 32768, + "gpt-4-0314": 8192, # legacy + "gpt-4-32k-0314": 32768, # legacy + "gpt-3.5-turbo-1106": 16385, + "gpt-3.5-turbo": 4096, + "gpt-3.5-turbo-16k": 16385, + "gpt-3.5-turbo-0613": 4096, # legacy + "gpt-3.5-turbo-16k-0613": 16385, # legacy + "gpt-3.5-turbo-0301": 4096, # legacy +} + +# ============================================================================ +# Server-Only Constants +# ============================================================================ + +MIRIX_DIR = os.path.join(os.path.expanduser("~"), ".mirix") +MIRIX_DIR_TOOL_SANDBOX = os.path.join(MIRIX_DIR, "tool_sandbox_dir") + +ADMIN_PREFIX = "/v1/admin" +API_PREFIX = "/v1" +OPENAI_API_PREFIX = "/openai" + +COMPOSIO_ENTITY_ENV_VAR_KEY = "COMPOSIO_ENTITY" + +# String in the error message for when the context window is too large +# Example full message: +# This model's maximum context length is 8192 tokens. However, your messages resulted in 8198 tokens (7450 in the messages, 748 in the functions). Please reduce the length of the messages or functions. +OPENAI_CONTEXT_WINDOW_ERROR_SUBSTRING = "maximum context length" + +# System prompt templating +IN_CONTEXT_MEMORY_KEYWORD = "CORE_MEMORY" + +MAX_CHAINING_STEPS = int(os.getenv("MAX_CHAINING_STEPS", "10")) +# Automatic procedural-memory evolution drives one persistent procedural agent +# statelessly per evolution step. A curator pass can survey existing skills and +# create/edit a few, so it needs a larger chaining (tool-use) budget than the +# default chat path. With the skill tools correctly bound the curator converges +# in ~2-5 steps. The cap counts LLM turns; an edit-heavy run is serial +# (skill_list -> per edit skill_read + skill_edit -> finish_memory_update), so the +# worst case is ~2N+2 turns for N edits. With the default edit budget capped at +# SKILL_EDIT_BUDGET_MAX (6) that is ~14, so 15 leaves one turn of headroom while +# still BOUNDING a pathological spin to ~15*~12s (~3 min) instead of running the +# loop out to a 600s timeout. Overridable via env; only the procedural evolution +# path passes this into step(max_chaining_steps=...), so normal meta-agent flow +# remains governed by MAX_CHAINING_STEPS. +SKILL_EVOLVE_MAX_CHAINING_STEPS = int( + os.getenv("SKILL_EVOLVE_MAX_CHAINING_STEPS", "15") +) +MAX_RETRIEVAL_LIMIT_IN_SYSTEM = 10 + +# tokenizers +EMBEDDING_TO_TOKENIZER_MAP = { + "text-embedding-3-small": "cl100k_base", +} +EMBEDDING_TO_TOKENIZER_DEFAULT = "cl100k_base" + + +DEFAULT_MIRIX_MODEL = "gpt-4" # TODO: fixme +DEFAULT_PERSONA = "sam_pov" +DEFAULT_HUMAN = "basic" +DEFAULT_PRESET = "memgpt_chat" + +# Base tools that cannot be edited, as they access agent state directly +# Note that we don't include "conversation_search_date" for now +BASE_TOOLS = [ + "send_intermediate_message", + "conversation_search", + "search_in_memory", + "list_memory_within_timerange", +] +# Base memory tools CAN be edited, and are added by default by the server +CORE_MEMORY_TOOLS = ["core_memory_append", "core_memory_rewrite"] +EPISODIC_MEMORY_TOOLS = [ + "episodic_memory_insert", + "episodic_memory_merge", + "episodic_memory_replace", + "check_episodic_memory", +] +SKILL_TOOLS = ["skill_list", "skill_read", "skill_create", "skill_edit", "skill_delete"] +RESOURCE_MEMORY_TOOLS = ["resource_memory_insert", "resource_memory_update"] +KNOWLEDGE_VAULT_TOOLS = ["knowledge_vault_insert", "knowledge_vault_update"] +SEMANTIC_MEMORY_TOOLS = [ + "semantic_memory_insert", + "semantic_memory_update", + "check_semantic_memory", +] +CHAT_AGENT_TOOLS = [] +EXTRAS_TOOLS = ["web_search", "fetch_and_read_pdf"] +MCP_TOOLS = [] +META_MEMORY_TOOLS = ["trigger_memory_update"] +SEARCH_MEMORY_TOOLS = ["search_in_memory", "list_memory_within_timerange"] UNIVERSAL_MEMORY_TOOLS = [ "search_in_memory", "finish_memory_update", @@ -136,7 +151,9 @@ set( CORE_MEMORY_TOOLS + EPISODIC_MEMORY_TOOLS - + PROCEDURAL_MEMORY_TOOLS + # auto-dream consolidates procedural memory via our skill paradigm + # (skill_* tools), not the removed procedural_memory_insert/update. + + SKILL_TOOLS + RESOURCE_MEMORY_TOOLS + KNOWLEDGE_VAULT_TOOLS + SEMANTIC_MEMORY_TOOLS @@ -147,115 +164,248 @@ set( BASE_TOOLS + CORE_MEMORY_TOOLS - + EPISODIC_MEMORY_TOOLS - + PROCEDURAL_MEMORY_TOOLS - + RESOURCE_MEMORY_TOOLS - + KNOWLEDGE_VAULT_TOOLS - + SEMANTIC_MEMORY_TOOLS - + META_MEMORY_TOOLS - + UNIVERSAL_MEMORY_TOOLS - + CHAT_AGENT_TOOLS - + EXTRAS_TOOLS - + MCP_TOOLS - ) -) - -# The name of the tool used to send message to the user -# Structured output models -STRUCTURED_OUTPUT_MODELS = {"gpt-4o", "gpt-4o-mini"} - -# LOGGER_LOG_LEVEL is use to convert Text to Logging level value for logging mostly for Cli input to setting level -LOGGER_LOG_LEVELS = { - "CRITICAL": CRITICAL, - "ERROR": ERROR, - "WARN": WARN, - "WARNING": WARNING, - "INFO": INFO, - "DEBUG": DEBUG, - "NOTSET": NOTSET, -} - -FIRST_MESSAGE_ATTEMPTS = 10 - -INITIAL_BOOT_MESSAGE = "Boot sequence complete. Persona activated." -INITIAL_BOOT_MESSAGE_SEND_MESSAGE_THOUGHT = ( - "Bootup sequence complete. Persona activated. Testing messaging functionality." -) -STARTUP_QUOTES = [ - "I think, therefore I am.", - "All those moments will be lost in time, like tears in rain.", - "More human than human is our motto.", -] -INITIAL_BOOT_MESSAGE_SEND_MESSAGE_FIRST_MSG = STARTUP_QUOTES[2] - -CLI_WARNING_PREFIX = "Warning: " - -ERROR_MESSAGE_PREFIX = "Error" - -NON_USER_MSG_PREFIX = "[This is an automated system message hidden from the user] " - -# The error message that Mirix will receive -# MESSAGE_SUMMARY_WARNING_STR = f"Warning: the conversation history will soon reach its maximum length and be trimmed. Make sure to save any important information from the conversation to your memory before it is removed." -# Much longer and more specific variant of the prompt -# TODO: Emit the warning to Meta Memory Manager instead of the Chat Agent. -MESSAGE_SUMMARY_WARNING_STR = " ".join( - [ - f"{NON_USER_MSG_PREFIX}The conversation history will soon reach its maximum length and be trimmed.", - "Do NOT tell the user about this system alert, they should not know that the history is reaching max length.", - ] -) - -# The ackknowledgement message used in the summarize sequence -MESSAGE_SUMMARY_REQUEST_ACK = "Understood, I will respond with a summary of the message (and only the summary, nothing else) once I receive the conversation history. I'm ready." - -# Maximum length of an error message -MAX_ERROR_MESSAGE_CHAR_LIMIT = 500 - -# Default memory limits -CORE_MEMORY_PERSONA_CHAR_LIMIT: int = 5000 -CORE_MEMORY_HUMAN_CHAR_LIMIT: int = 5000 - -MAX_PAUSE_HEARTBEATS = 360 # in min - -MESSAGE_CHATGPT_FUNCTION_MODEL = "gpt-3.5-turbo" -MESSAGE_CHATGPT_FUNCTION_SYSTEM_MESSAGE = "You are a helpful assistant. Keep your responses short and concise." - -#### Functions related - -# REQ_HEARTBEAT_MESSAGE = f"{NON_USER_MSG_PREFIX}continue_chaining == true" -REQ_HEARTBEAT_MESSAGE = f"{NON_USER_MSG_PREFIX}Function called using continue_chaining=true, returning control" -# FUNC_FAILED_HEARTBEAT_MESSAGE = f"{NON_USER_MSG_PREFIX}Function call failed" -FUNC_FAILED_HEARTBEAT_MESSAGE = f"{NON_USER_MSG_PREFIX}Function call failed, returning control" - - -RETRIEVAL_QUERY_DEFAULT_PAGE_SIZE = 5 - -MAX_FILENAME_LENGTH = 255 -RESERVED_FILENAMES = {"CON", "PRN", "AUX", "NUL", "COM1", "COM2", "LPT1", "LPT2"} - -MAX_IMAGES_TO_PROCESS = 100 - -DEFAULT_WRAPPER_NAME = "chatml" -INNER_THOUGHTS_KWARG_DESCRIPTION = "Deep inner monologue private to you only." -INNER_THOUGHTS_CLI_SYMBOL = "๐Ÿ’ญ" -ASSISTANT_MESSAGE_CLI_SYMBOL = "๐Ÿค–" - -CLEAR_HISTORY_AFTER_MEMORY_UPDATE = os.getenv("CLEAR_HISTORY_AFTER_MEMORY_UPDATE", "true").lower() in ( - "true", - "1", - "yes", -) -CALL_MEMORY_AGENT_IN_PARALLEL = os.getenv("CALL_MEMORY_AGENT_IN_PARALLEL", "false").lower() in ("true", "1", "yes") -CHAINING_FOR_MEMORY_UPDATE = os.getenv("CHAINING_FOR_MEMORY_UPDATE", "false").lower() in ("true", "1", "yes") -CHAINING_FOR_META_AGENT = os.getenv("CHAINING_FOR_META_AGENT", "true").lower() in ( - "true", - "1", - "yes", -) - -LOAD_IMAGE_CONTENT_FOR_LAST_MESSAGE_ONLY = os.getenv("LOAD_IMAGE_CONTENT_FOR_LAST_MESSAGE_ONLY", "false").lower() in ( - "true", - "1", - "yes", -) -BUILD_EMBEDDINGS_FOR_MEMORY = os.getenv("BUILD_EMBEDDINGS_FOR_MEMORY", "true").lower() in ("true", "1", "yes") + + EPISODIC_MEMORY_TOOLS + + SKILL_TOOLS + + RESOURCE_MEMORY_TOOLS + + KNOWLEDGE_VAULT_TOOLS + + SEMANTIC_MEMORY_TOOLS + + META_MEMORY_TOOLS + + UNIVERSAL_MEMORY_TOOLS + + CHAT_AGENT_TOOLS + + EXTRAS_TOOLS + + MCP_TOOLS + ) +) + +# The name of the tool used to send message to the user +# Structured output models +STRUCTURED_OUTPUT_MODELS = {"gpt-4o", "gpt-4o-mini"} + +# LOGGER_LOG_LEVEL is use to convert Text to Logging level value for logging mostly for Cli input to setting level +LOGGER_LOG_LEVELS = { + "CRITICAL": CRITICAL, + "ERROR": ERROR, + "WARN": WARN, + "WARNING": WARNING, + "INFO": INFO, + "DEBUG": DEBUG, + "NOTSET": NOTSET, +} + +FIRST_MESSAGE_ATTEMPTS = 10 + +INITIAL_BOOT_MESSAGE = "Boot sequence complete. Persona activated." +INITIAL_BOOT_MESSAGE_SEND_MESSAGE_THOUGHT = ( + "Bootup sequence complete. Persona activated. Testing messaging functionality." +) +STARTUP_QUOTES = [ + "I think, therefore I am.", + "All those moments will be lost in time, like tears in rain.", + "More human than human is our motto.", +] +INITIAL_BOOT_MESSAGE_SEND_MESSAGE_FIRST_MSG = STARTUP_QUOTES[2] + +CLI_WARNING_PREFIX = "Warning: " + +ERROR_MESSAGE_PREFIX = "Error" + +NON_USER_MSG_PREFIX = "[This is an automated system message hidden from the user] " + +# The error message that Mirix will receive +# MESSAGE_SUMMARY_WARNING_STR = f"Warning: the conversation history will soon reach its maximum length and be trimmed. Make sure to save any important information from the conversation to your memory before it is removed." +# Much longer and more specific variant of the prompt +# TODO: Emit the warning to Meta Memory Manager instead of the Chat Agent. +MESSAGE_SUMMARY_WARNING_STR = " ".join( + [ + f"{NON_USER_MSG_PREFIX}The conversation history will soon reach its maximum length and be trimmed.", + "Do NOT tell the user about this system alert, they should not know that the history is reaching max length.", + ] +) + +# The ackknowledgement message used in the summarize sequence +MESSAGE_SUMMARY_REQUEST_ACK = "Understood, I will respond with a summary of the message (and only the summary, nothing else) once I receive the conversation history. I'm ready." + +# Maximum length of an error message +MAX_ERROR_MESSAGE_CHAR_LIMIT = 500 + +# Default memory limits +CORE_MEMORY_PERSONA_CHAR_LIMIT: int = 5000 +CORE_MEMORY_HUMAN_CHAR_LIMIT: int = 5000 + +MAX_PAUSE_HEARTBEATS = 360 # in min + +MESSAGE_CHATGPT_FUNCTION_MODEL = "gpt-3.5-turbo" +MESSAGE_CHATGPT_FUNCTION_SYSTEM_MESSAGE = ( + "You are a helpful assistant. Keep your responses short and concise." +) + +#### Functions related + +# REQ_HEARTBEAT_MESSAGE = f"{NON_USER_MSG_PREFIX}continue_chaining == true" +REQ_HEARTBEAT_MESSAGE = f"{NON_USER_MSG_PREFIX}Function called using continue_chaining=true, returning control" +# FUNC_FAILED_HEARTBEAT_MESSAGE = f"{NON_USER_MSG_PREFIX}Function call failed" +FUNC_FAILED_HEARTBEAT_MESSAGE = ( + f"{NON_USER_MSG_PREFIX}Function call failed, returning control" +) + + +RETRIEVAL_QUERY_DEFAULT_PAGE_SIZE = 5 + +MAX_FILENAME_LENGTH = 255 +RESERVED_FILENAMES = {"CON", "PRN", "AUX", "NUL", "COM1", "COM2", "LPT1", "LPT2"} + +MAX_IMAGES_TO_PROCESS = 100 + +DEFAULT_WRAPPER_NAME = "chatml" +INNER_THOUGHTS_KWARG_DESCRIPTION = "Deep inner monologue private to you only." +INNER_THOUGHTS_CLI_SYMBOL = "๐Ÿ’ญ" +ASSISTANT_MESSAGE_CLI_SYMBOL = "๐Ÿค–" + +CLEAR_HISTORY_AFTER_MEMORY_UPDATE = os.getenv( + "CLEAR_HISTORY_AFTER_MEMORY_UPDATE", "true" +).lower() in ( + "true", + "1", + "yes", +) +# When clearing history after a memory update, retain the raw conversation +# messages of the most-recent N sessions (per agent+user) instead of hard-deleting +# them. These retained rows stay DETACHED from the agent's message_ids (NOT re-added +# to any in-context list โ€” token economy preserved); they exist only so a later +# distiller / auto-dream pass can read the raw transcript. The window is a bounded +# rolling window: when a new session arrives the oldest retained session ages out +# and becomes eligible for deletion on the next clear. Default 5; 0 disables +# retention entirely (legacy hard-delete-everything behavior). +MESSAGE_RETAIN_LAST_N_SESSIONS = int(os.getenv("MESSAGE_RETAIN_LAST_N_SESSIONS", "5")) +CALL_MEMORY_AGENT_IN_PARALLEL = os.getenv( + "CALL_MEMORY_AGENT_IN_PARALLEL", "false" +).lower() in ("true", "1", "yes") +CHAINING_FOR_MEMORY_UPDATE = os.getenv( + "CHAINING_FOR_MEMORY_UPDATE", "false" +).lower() in ("true", "1", "yes") +CHAINING_FOR_META_AGENT = os.getenv("CHAINING_FOR_META_AGENT", "true").lower() in ( + "true", + "1", + "yes", +) + +LOAD_IMAGE_CONTENT_FOR_LAST_MESSAGE_ONLY = os.getenv( + "LOAD_IMAGE_CONTENT_FOR_LAST_MESSAGE_ONLY", "false" +).lower() in ( + "true", + "1", + "yes", +) +BUILD_EMBEDDINGS_FOR_MEMORY = os.getenv( + "BUILD_EMBEDDINGS_FOR_MEMORY", "true" +).lower() in ("true", "1", "yes") +SKILL_TRIGGER_MESSAGE_THRESHOLD = int( + os.getenv("SKILL_TRIGGER_MESSAGE_THRESHOLD", "10") +) +# Number of *sessions* (distinct message.session_id values) that must accumulate +# on an agent for a given user before procedural-memory extraction auto-fires. +# Counted against a persisted cursor (agent_trigger_state), so restarts and +# multiple workers all see the same value. +SKILL_TRIGGER_SESSION_THRESHOLD = int(os.getenv("SKILL_TRIGGER_SESSION_THRESHOLD", "5")) + +# --------------------------------------------------------------------------- +# Bounded, count-driven, quality-aware edit budget (skill evolution). +# +# The budget is the MAX number of skill mutations (create+edit) a single +# records-based evolution run may perform. It is driven by STRUCTURE-GATED +# COUNTS from the C2 aggregate (n_high_fail / n_high_succ), never by a sum of +# self-reported quality scores (quality_score is ranking-only): +# +# raw = B0 + alpha_f * n_high_fail + alpha_s * n_high_succ +# budget = clamp(round_half_up(raw), B_min, B_max) +# +# With ~5 records/window: all-failure -> ~6, mixed -> ~2-4, all-noise -> 0 +# (n_high_* == 0 => budget 0 => the curator is skipped entirely). +# --------------------------------------------------------------------------- +SKILL_EDIT_BUDGET_B0 = int(os.getenv("SKILL_EDIT_BUDGET_B0", "1")) +SKILL_EDIT_BUDGET_ALPHA_FAIL = float(os.getenv("SKILL_EDIT_BUDGET_ALPHA_FAIL", "1.0")) +SKILL_EDIT_BUDGET_ALPHA_SUCC = float(os.getenv("SKILL_EDIT_BUDGET_ALPHA_SUCC", "0.5")) +SKILL_EDIT_BUDGET_MIN = int(os.getenv("SKILL_EDIT_BUDGET_MIN", "0")) +SKILL_EDIT_BUDGET_MAX = int(os.getenv("SKILL_EDIT_BUDGET_MAX", "6")) + +# Hybrid budget (formula ceiling + a cheap LLM may only REDUCE). Held out of the +# validity run: formula-only is deterministic; autonomous is an ablation arm. +SKILL_USE_AUTONOMOUS_BUDGET = os.getenv( + "SKILL_USE_AUTONOMOUS_BUDGET", "false" +).lower() in ( + "true", + "1", + "yes", +) + +# Per-mutation HARD size gate on skill_edit. Reject (without consuming budget) +# when an edit changes too much at once; the curator must split it or extract a +# new skill instead. Two independent dimensions, either one trips the gate: +# * absolute char delta > SKILL_MAX_EDIT_CHAR_DELTA +# * change-ratio (1 - difflib.SequenceMatcher.ratio()) >= SKILL_EDIT_MAJOR_RATIO +SKILL_MAX_EDIT_CHAR_DELTA = int(os.getenv("SKILL_MAX_EDIT_CHAR_DELTA", "800")) +SKILL_EDIT_MAJOR_RATIO = float(os.getenv("SKILL_EDIT_MAJOR_RATIO", "0.4")) + +# Hard ceiling on a skill's `instructions` length. An edit whose RESULT exceeds +# this must route to skill_create (extract a new skill) rather than grow an +# existing one unbounded. +SKILL_MAX_INSTRUCTIONS_CHARS = int(os.getenv("SKILL_MAX_INSTRUCTIONS_CHARS", "12000")) + +# Deletes are gated SEPARATELY from the create/edit budget (P1-4): at most +# D_max per evolution, and only when a failure record names the skill actively +# harmful. Soft-delete (exclude-from-retrieval) is preferred over a hard delete. +SKILL_DELETE_BUDGET_MAX = int(os.getenv("SKILL_DELETE_BUDGET_MAX", "1")) + +# Records-based evolution window: evolve every N completed graded rounds. +SKILL_EVOLVE_RECORDS_WINDOW = int(os.getenv("SKILL_EVOLVE_RECORDS_WINDOW", "5")) + +# --------------------------------------------------------------------------- +# Procedural (skill) RETRIEVAL defaults. +# +# The single authoritative default for SEARCH-SURFACE procedural retrieval +# (the REST /memory/search endpoint, the eval client, and the documented +# recommendation for the search_in_memory agent tool). When a search surface +# does NOT explicitly choose a method, procedural memory is retrieved with the +# EverOS-aligned "hybrid" lane: BM25 (lexical) + embedding (dense) recall fused +# with Reciprocal Rank Fusion. +# +# Override to "bm25" or "embedding" via env for latency-sensitive deployments +# or A/B experiments. Internal pipelines that PIN a method explicitly (evolve +# before/after snapshots, the per-step system-context prefetch) keep their +# explicit choice and are not governed by this default. +PROCEDURAL_DEFAULT_SEARCH_METHOD = os.getenv("MIRIX_SKILL_SEARCH_METHOD", "hybrid") + +# Distiller transcript budgets (characters). The transcript fed to the +# experience-distiller LLM is bounded twice: per turn, and for the whole +# session (HEAD + TAIL kept, middle elided past the cap). Tool turns are part +# of the transcript โ€” work-process lessons live in tool errors/retries โ€” so +# the defaults are sized for tool-heavy sessions: 80K chars โ‰ˆ 20-25K tokens, +# comfortably inside modern model windows while still bounding a pathological +# session. Raise/lower via env per deployment model. +DISTILLER_MAX_TRANSCRIPT_CHARS = int( + os.getenv("MIRIX_DISTILLER_MAX_TRANSCRIPT_CHARS", "80000") +) +DISTILLER_MAX_MESSAGE_CHARS = int( + os.getenv("MIRIX_DISTILLER_MAX_MESSAGE_CHARS", "8000") +) + +# Retention window (days) for VERBATIM conversation turns in the +# conversation_message store AFTER they have been distilled. Once distilled, +# a turn's learning value lives in skill_experience rows; keeping the raw +# transcript indefinitely is a pure PII liability. Pruning runs at the end of +# each procedural distillation pass. Set to 0 (or negative) to keep forever. +CONVERSATION_RETENTION_DAYS = int(os.getenv("MIRIX_CONVERSATION_RETENTION_DAYS", "30")) + +# Reciprocal Rank Fusion constant for the procedural hybrid lane. k=60 is the +# canonical RRF constant (everalgo.rank.fusion.rrf / RankConfig.rrf_k=60); the +# fusion is rank-based, 1-based, and unweighted โ€” raw BM25/cosine scores are +# discarded, only rank positions feed the sum 1/(k+rank). +SKILL_HYBRID_RRF_K = int(os.getenv("MIRIX_SKILL_HYBRID_RRF_K", "60")) + +# Over-fetch multiplier for each hybrid lane before fusion (EverOS +# DEFAULT_RECALL_MULTIPLIER=2): each lane recalls limit*multiplier candidates so +# fusion picks from a larger, higher-recall pool, then the result is sliced to +# the requested limit. +SKILL_HYBRID_RECALL_MULTIPLIER = int( + os.getenv("MIRIX_SKILL_HYBRID_RECALL_MULTIPLIER", "2") +) diff --git a/mirix/database/redis_client.py b/mirix/database/redis_client.py index 6fa054cfe..bae83e4fa 100644 --- a/mirix/database/redis_client.py +++ b/mirix/database/redis_client.py @@ -701,8 +701,38 @@ async def _create_semantic_index(self) -> None: except Exception as e: logger.warning("Failed to create semantic index: %s", e) + @staticmethod + def _index_has_attribute(info, attribute: str) -> bool: + """True when an FT.INFO payload declares a field named `attribute`. + + The attributes section is a nested list of str/bytes tokens whose exact + shape varies across redis-py/RediSearch versions, so walk it generically + instead of assuming positions. + """ + + def _texts(obj): + if isinstance(obj, (list, tuple)): + for item in obj: + yield from _texts(item) + elif isinstance(obj, bytes): + yield obj.decode("utf-8", "ignore") + elif isinstance(obj, str): + yield obj + + if not isinstance(info, dict): + return False + attributes = info.get("attributes") or info.get(b"attributes") or [] + return any(text == attribute for text in _texts(attributes)) + async def _create_procedural_index(self) -> None: - """Create JSON-based index for procedural memory with 2 VECTOR fields.""" + """Create JSON-based index for procedural memory with 2 VECTOR fields. + + Upgrade-aware: an index created before the skill schema (summary / + summary_embedding / steps_embedding fields) persists across deploys and + would make every skill-field query miss Redis forever. When the existing + index lacks the skill fields, drop the INDEX ONLY (documents are kept) + and rebuild; RediSearch then re-indexes the prefix in the background. + """ try: from redis.commands.search.field import NumericField, TagField, TextField, VectorField from redis.commands.search.index_definition import IndexDefinition, IndexType @@ -710,33 +740,49 @@ async def _create_procedural_index(self) -> None: from mirix.constants import MAX_EMBEDDING_DIM try: - await self.client.ft(self.PROCEDURAL_INDEX).info() - logger.debug("Index %s already exists", self.PROCEDURAL_INDEX) - return + info = await self.client.ft(self.PROCEDURAL_INDEX).info() except Exception: - pass + info = None + + if info is not None: + if self._index_has_attribute(info, "description_embedding"): + logger.debug("Index %s already exists", self.PROCEDURAL_INDEX) + return + try: + await self.client.ft(self.PROCEDURAL_INDEX).dropindex(delete_documents=False) + logger.info( + "Dropped stale pre-skill procedural index %s; rebuilding with skill schema", + self.PROCEDURAL_INDEX, + ) + except Exception as e: + logger.warning( + "Failed to drop stale procedural index %s: %s", self.PROCEDURAL_INDEX, e + ) + return schema = ( TextField("$.organization_id", as_name="organization_id"), TextField("$.agent_id", as_name="agent_id"), TextField("$.entry_type", as_name="entry_type"), - TextField("$.summary", as_name="summary"), + TextField("$.name", as_name="name"), + TextField("$.description", as_name="description"), + TextField("$.instructions", as_name="instructions"), TagField("$.user_id", as_name="user_id"), NumericField("$.created_at_ts", as_name="created_at_ts"), TagField("$.filter_tags.scope", as_name="filter_tags_scope"), # Explicit scope field TextField("$.filter_tags.*", as_name="filter_tags"), # Filter tags for flexible filtering # Two vector fields (32KB total) VectorField( - "$.summary_embedding", + "$.description_embedding", "FLAT", {"TYPE": "FLOAT32", "DIM": MAX_EMBEDDING_DIM, "DISTANCE_METRIC": "COSINE"}, - as_name="summary_embedding", + as_name="description_embedding", ), VectorField( - "$.steps_embedding", + "$.instructions_embedding", "FLAT", {"TYPE": "FLOAT32", "DIM": MAX_EMBEDDING_DIM, "DISTANCE_METRIC": "COSINE"}, - as_name="steps_embedding", + as_name="instructions_embedding", ), ) @@ -998,7 +1044,7 @@ def escape_redis_query(text: str) -> str: logger.debug("๐Ÿ• Redis temporal filter: @occurred_at_ts:[%s %s]", min_ts, max_ts) # Add filter_tags filters - if filter_tags: + if filter_tags or scopes: filter_query = self._build_filter_tags_query(filter_tags, scopes=scopes) if filter_query: query_parts.append(filter_query) @@ -1130,7 +1176,7 @@ async def search_vector( logger.debug("๐Ÿ• Redis temporal filter: @occurred_at_ts:[%s %s]", min_ts, max_ts) # Add filter_tags filters - if filter_tags: + if filter_tags or scopes: filter_query = self._build_filter_tags_query(filter_tags, scopes=scopes) if filter_query: query_parts.append(filter_query) @@ -1264,7 +1310,7 @@ async def search_recent( logger.debug("๐Ÿ• Redis temporal filter: @occurred_at_ts:[%s %s]", min_ts, max_ts) # Add filter_tags filters - if filter_tags: + if filter_tags or scopes: filter_query = self._build_filter_tags_query(filter_tags, scopes=scopes) if filter_query: query_parts.append(filter_query) @@ -1372,7 +1418,7 @@ def escape_text_value(value: str) -> str: logger.debug("๐Ÿ• Redis temporal filter: @occurred_at_ts:[%s %s]", min_ts, max_ts) # Add filter_tags filters (including scope) - if filter_tags: + if filter_tags or scopes: filter_query = self._build_filter_tags_query(filter_tags, scopes=scopes) if filter_query: query_parts.append(filter_query) @@ -1465,7 +1511,7 @@ def escape_text_value(value: str) -> str: filter_parts.append(f"@occurred_at_ts:[{min_ts} {max_ts}]") # Add filter_tags filters (including scope) - if filter_tags: + if filter_tags or scopes: filter_query = self._build_filter_tags_query(filter_tags, scopes=scopes) if filter_query: filter_parts.append(filter_query) @@ -1563,7 +1609,7 @@ def escape_text_value(value: str) -> str: query_parts.append(f"@occurred_at_ts:[{min_ts} {max_ts}]") # Add filter_tags filters (including scope) - if filter_tags: + if filter_tags or scopes: filter_query = self._build_filter_tags_query(filter_tags, scopes=scopes) if filter_query: query_parts.append(filter_query) diff --git a/mirix/database/schema_extractor.py b/mirix/database/schema_extractor.py index c407c5de8..e2eadf2d8 100644 --- a/mirix/database/schema_extractor.py +++ b/mirix/database/schema_extractor.py @@ -264,9 +264,13 @@ def get_basic_schema(): CREATE TABLE procedural_memory ( id VARCHAR PRIMARY KEY, organization_id VARCHAR NOT NULL, + name VARCHAR NOT NULL, entry_type VARCHAR NOT NULL, - summary VARCHAR NOT NULL, - steps TEXT NOT NULL, + description VARCHAR NOT NULL, + instructions TEXT NOT NULL, + triggers JSON, + examples JSON, + version VARCHAR DEFAULT '0.1.0', last_modify TEXT NOT NULL, metadata_ TEXT, created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP, diff --git a/mirix/embeddings.py b/mirix/embeddings.py index 91835cf1b..6292a25ae 100755 --- a/mirix/embeddings.py +++ b/mirix/embeddings.py @@ -328,6 +328,39 @@ async def get_text_embedding(self, text: str) -> List[float]: return await embedding_with_retry(lambda: self._call_api(text)) +class OpenRouterEmbedding(EmbeddingEndpoint): + """EmbeddingEndpoint subclass that adds Bearer auth for OpenRouter.""" + + def __init__(self, api_key: str, **kwargs: Any): + super().__init__(**kwargs) + self._api_key = api_key + + async def _call_api(self, text: str) -> List[float]: + import httpx + + headers = { + "Content-Type": "application/json", + "Authorization": f"Bearer {self._api_key}", + } + json_data = {"input": text, "model": self.model_name} + + async with httpx.AsyncClient() as client: + response = await client.post( + f"{self._base_url}/embeddings", + headers=headers, + json=json_data, + timeout=self._timeout, + ) + + response_json = response.json() + if isinstance(response_json, dict): + try: + return response_json["data"][0]["embedding"] + except (KeyError, IndexError): + raise TypeError(f"Unexpected embedding response: {response_json}") + raise TypeError(f"Unexpected embedding response type: {response_json}") + + class AzureOpenAIEmbedding: def __init__( self, @@ -517,6 +550,15 @@ async def embedding_model(config: EmbeddingConfig, user_id: Optional[uuid.UUID] user=user_id, langfuse_model=config.langfuse_model, ) + elif endpoint_type == "openrouter": + api_key = config.api_key or model_settings.openai_api_key + return OpenRouterEmbedding( + api_key=api_key, + model=config.embedding_model, + base_url=config.embedding_endpoint, + user=str(user_id) if user_id else "", + ) + elif endpoint_type == "ollama": model = OllamaEmbeddings( model=config.embedding_model, diff --git a/mirix/functions/function_sets/base.py b/mirix/functions/function_sets/base.py index fa12aabb4..cfebe9fc4 100644 --- a/mirix/functions/function_sets/base.py +++ b/mirix/functions/function_sets/base.py @@ -1,329 +1,341 @@ -from typing import Optional - -from mirix.agent import Agent, AgentState -from mirix.utils import convert_timezone_to_utc - - -async def send_message(self: "Agent", agent_state: "AgentState", message: str) -> Optional[str]: - """ - Sends a message to the human user. Meanwhile, whenever this function is called, the agent needs to include the `topic` of the current focus. It can be the same as before, it can also be updated when the agent is focusing on something different. - - Args: - message (str): Message contents. All unicode (including emojis) are supported. - topic (str): The focus of the agent right now. It is used to track the most recent topic in the conversation and will be used to retrieve the relevant memories from each memory component. - - Returns: - Optional[str]: None is always returned as this function does not produce a response. - """ - self.interface.assistant_message(message) - return None - - -async def send_intermediate_message( - self: "Agent", - agent_state: "AgentState", - message: str, -) -> Optional[str]: - """ - Sends an intermediate message to the human user. Meanwhile, whenever this function is called, the agent needs to include the `topic` of the current focus. It should NEVER be any questions or requests for the user but only the agent's current progress on the task. - - Args: - message (str): Message contents. All unicode (including emojis) are supported. - topic (str): The focus of the agent right now. It is used to track the most recent topic in the conversation and will be used to retrieve the relevant memories from each memory component. - - Returns: - Optional[str]: None is always returned as this function does not produce a response. - """ - self.interface.assistant_message(message) - return None - - -async def conversation_search(self: "Agent", query: str, page: Optional[int] = 0) -> Optional[str]: - """ - Search prior conversation history using case-insensitive string matching. - - Args: - query (str): String to search for. - page (int): Allows you to page through results. Only use on a follow-up query. Defaults to 0 (first page). - - Returns: - str: Query result string - """ - - import math - - from mirix.constants import RETRIEVAL_QUERY_DEFAULT_PAGE_SIZE - from mirix.utils import json_dumps - - if page is None or (isinstance(page, str) and page.lower().strip() == "none"): - page = 0 - try: - page = int(page) - except (ValueError, TypeError): - raise ValueError("'page' argument must be an integer") - count = RETRIEVAL_QUERY_DEFAULT_PAGE_SIZE - # TODO: add paging by page number. currently cursor only works with strings. - # original: start=page * count - messages = await self.message_manager.list_user_messages_for_agent( - agent_id=self.agent_state.id, - actor=self.actor, - query_text=query, - limit=count, - ) - total = len(messages) - num_pages = math.ceil(total / count) - 1 # 0 index - if len(messages) == 0: - results_str = "No results found." - else: - results_pref = f"Showing {len(messages)} of {total} results (page {page}/{num_pages}):" - results_formatted = [message.text for message in messages] - results_str = f"{results_pref} {json_dumps(results_formatted)}" - return results_str - - -async def search_in_memory( - self: "Agent", - memory_type: str, - query: str, - search_field: str, - search_method: str, - timezone_str: str, -) -> Optional[str]: - """ - Choose which memory to search. All memory types support multiple search methods with different performance characteristics. Most of the time, you should use search over 'details' for episodic memory and semantic memory, 'content' for resource memory (but for resource memory, `embedding` is not supported for content field so you have to use other search methods), 'description' for procedural memory. This is because these fields have the richest information and is more likely to contain the keywords/query. You can always start from a thorough search over the whole memory by setting memory_type as 'all' and search_field as 'null', and then narrow down to specific fields and specific memories. - - Args: - memory_type: The type of memory to search in. It should be chosen from the following: "episodic", "resource", "procedural", "knowledge_vault", "semantic", "all". Here "all" means searching in all the memories. - query: The keywords/query used to search in the memory. - search_field: The field to search in the memory. It should be chosen from the attributes of the corresponding memory. For "episodic" memory, it can be 'summary', 'details'; for "resource" memory, it can be 'summary', 'content'; for "procedural" memory, it can be 'summary', 'steps'; for "knowledge_vault", it can be 'secret_value', 'caption'; for semantic memory, it can be 'name', 'summary', 'details'. For "all", it should also be "null" as the system will search all memories with default fields. - search_method: The method to search in the memory. Choose from: - - 'bm25': BM25 ranking-based full-text search (fast and effective for keyword-based searches) - - 'embedding': Vector similarity search using embeddings (most powerful, good for conceptual matches) - - Returns: - str: Query result string - """ - - if not self.user: - raise ValueError("Can not search in memory. User is not set") - - if memory_type == "resource" and search_field == "content" and search_method == "embedding": - raise ValueError("embedding is not supported for resource memory's 'content' field.") - if memory_type == "knowledge_vault" and search_field == "secret_value" and search_method == "embedding": - raise ValueError("embedding is not supported for knowledge_vault memory's 'secret_value' field.") - - if memory_type == "all": - search_field = "null" - - # Pre-compute embedding once if using embedding search (to avoid redundant embeddings) - embedded_text = None - if search_method == "embedding" and query: - import numpy as np - - from mirix.constants import MAX_EMBEDDING_DIM - from mirix.embeddings import embedding_model - - embedded_text = await (await embedding_model(self.agent_state.embedding_config)).get_text_embedding(query) - # Pad for episodic memory which requires MAX_EMBEDDING_DIM - embedded_text_padded = np.pad( - np.array(embedded_text), (0, MAX_EMBEDDING_DIM - len(embedded_text)), mode="constant" - ).tolist() - - if memory_type == "core": - # It means the model is an idiot, but we still return the results: - if self.blocks_in_memory: - return self.blocks_in_memory.compile(), len(self.blocks_in_memory.list_block_labels()) - return "", 0 - - if memory_type == "episodic" or memory_type == "all": - episodic_memory = await self.episodic_memory_manager.list_episodic_memory( - user=self.user, - agent_state=self.agent_state, - query=query, - embedded_text=embedded_text_padded if search_method == "embedding" and query else None, - search_field=search_field if search_field != "null" else "summary", - search_method=search_method, - limit=10, - timezone_str=timezone_str, - ) - formatted_results_from_episodic = [ - { - "memory_type": "episodic", - "id": x.id, - "timestamp": x.occurred_at, - "event_type": x.event_type, - "actor": x.actor, - "summary": x.summary, - "details": x.details, - } - for x in episodic_memory - ] - if memory_type == "episodic": - return formatted_results_from_episodic, len(formatted_results_from_episodic) - - if memory_type == "resource" or memory_type == "all": - resource_memories = await self.resource_memory_manager.list_resources( - user=self.user, - agent_state=self.agent_state, - query=query, - embedded_text=embedded_text if search_method == "embedding" and query else None, - search_field=( - search_field if search_field != "null" else ("summary" if search_method == "embedding" else "content") - ), - search_method=search_method, - limit=10, - timezone_str=timezone_str, - ) - formatted_results_resource = [ - { - "memory_type": "resource", - "id": x.id, - "resource_type": x.resource_type, - "summary": x.summary, - "content": x.content, - } - for x in resource_memories - ] - if memory_type == "resource": - return formatted_results_resource, len(formatted_results_resource) - - if memory_type == "procedural" or memory_type == "all": - procedural_memories = await self.procedural_memory_manager.list_procedures( - user=self.user, - agent_state=self.agent_state, - query=query, - embedded_text=embedded_text if search_method == "embedding" and query else None, - search_field=search_field if search_field != "null" else "summary", - search_method=search_method, - limit=10, - timezone_str=timezone_str, - ) - formatted_results_procedural = [ - { - "memory_type": "procedural", - "id": x.id, - "entry_type": x.entry_type, - "summary": x.summary, - "steps": x.steps, - } - for x in procedural_memories - ] - if memory_type == "procedural": - return formatted_results_procedural, len(formatted_results_procedural) - - if memory_type == "knowledge_vault" or memory_type == "all": - knowledge_vault_memories = await self.knowledge_vault_manager.list_knowledge( - user=self.user, - agent_state=self.agent_state, - query=query, - embedded_text=embedded_text if search_method == "embedding" and query else None, - search_field=search_field if search_field != "null" else "caption", - search_method=search_method, - limit=10, - timezone_str=timezone_str, - ) - formatted_results_knowledge_vault = [ - { - "memory_type": "knowledge_vault", - "id": x.id, - "entry_type": x.entry_type, - "source": x.source, - "sensitivity": x.sensitivity, - "secret_value": x.secret_value, - "caption": x.caption, - } - for x in knowledge_vault_memories - ] - if memory_type == "knowledge_vault": - return formatted_results_knowledge_vault, len(formatted_results_knowledge_vault) - - if memory_type == "semantic" or memory_type == "all": - semantic_memories = await self.semantic_memory_manager.list_semantic_items( - user=self.user, - agent_state=self.agent_state, - query=query, - embedded_text=embedded_text if search_method == "embedding" and query else None, - search_field=search_field if search_field != "null" else "summary", - search_method=search_method, - limit=10, - timezone_str=timezone_str, - ) - # title, summary, details, source - formatted_results_semantic = [ - { - "memory_type": "semantic", - "id": x.id, - "name": x.name, - "summary": x.summary, - "details": x.details, - "source": x.source, - } - for x in semantic_memories - ] - if memory_type == "semantic": - return formatted_results_semantic, len(formatted_results_semantic) - - else: - raise ValueError( - f"Memory type '{memory_type}' is not supported. Please choose from 'episodic', 'resource', 'procedural', 'knowledge_vault', 'semantic'." - ) - return ( - formatted_results_from_episodic - + formatted_results_resource - + formatted_results_procedural - + formatted_results_knowledge_vault - + formatted_results_semantic, - len(formatted_results_from_episodic) - + len(formatted_results_resource) - + len(formatted_results_procedural) - + len(formatted_results_knowledge_vault) - + len(formatted_results_semantic), - ) - - -async def list_memory_within_timerange( - self: "Agent", memory_type: str, start_time: str, end_time: str, timezone_str: str -) -> Optional[str]: - """ - List memories around a specific timestamp - Args: - memory_type (str): The type of memory to search in. It should be chosen from the following: "episodic", "resource", "procedural", "knowledge_vault", "semantic", "all". Here "all" means searching in all the memories. - start_time (str): The start time of the time range. It has to be in the form of "%Y-%m-%d %H:%M:%S" - end_time (str): The end time of the time range. It has to be in the form of "%Y-%m-%d %H:%M:%S" - """ - - start_time = convert_timezone_to_utc(start_time, timezone_str) - end_time = convert_timezone_to_utc(end_time, timezone_str) - - if not self.user: - raise ValueError("Can not list memory within timerange. User is not set") - - if memory_type == "episodic" or memory_type == "all": - episodic_memory = await self.episodic_memory_manager.list_episodic_memory_around_timestamp( - user=self.user, - agent_state=self.agent_state, - start_time=start_time, - end_time=end_time, - timezone_str=timezone_str, - ) - formatted_results_from_episodic = [ - { - "memory_type": "episodic", - "id": x.id, - "timestamp": x.occurred_at, - "event_type": x.event_type, - "actor": x.actor, - "summary": x.summary, - } - for x in episodic_memory - ] - if memory_type == "episodic": - if len(formatted_results_from_episodic) == 0: - return "No results found." - elif len(formatted_results_from_episodic) > 50: - return "Too many results found. Please narrow down your search." - else: - return formatted_results_from_episodic, len(formatted_results_from_episodic) - - # currently only episodic memory is supported - return None +from typing import Optional + +from mirix.agent import Agent, AgentState +from mirix.utils import convert_timezone_to_utc + + +async def send_message(self: "Agent", agent_state: "AgentState", message: str) -> Optional[str]: + """ + Sends a message to the human user. Meanwhile, whenever this function is called, the agent needs to include the `topic` of the current focus. It can be the same as before, it can also be updated when the agent is focusing on something different. + + Args: + message (str): Message contents. All unicode (including emojis) are supported. + topic (str): The focus of the agent right now. It is used to track the most recent topic in the conversation and will be used to retrieve the relevant memories from each memory component. + + Returns: + Optional[str]: None is always returned as this function does not produce a response. + """ + self.interface.assistant_message(message) + return None + + +async def send_intermediate_message( + self: "Agent", + agent_state: "AgentState", + message: str, +) -> Optional[str]: + """ + Sends an intermediate message to the human user. Meanwhile, whenever this function is called, the agent needs to include the `topic` of the current focus. It should NEVER be any questions or requests for the user but only the agent's current progress on the task. + + Args: + message (str): Message contents. All unicode (including emojis) are supported. + topic (str): The focus of the agent right now. It is used to track the most recent topic in the conversation and will be used to retrieve the relevant memories from each memory component. + + Returns: + Optional[str]: None is always returned as this function does not produce a response. + """ + self.interface.assistant_message(message) + return None + + +async def conversation_search(self: "Agent", query: str, page: Optional[int] = 0) -> Optional[str]: + """ + Search prior conversation history using case-insensitive string matching. + + Args: + query (str): String to search for. + page (int): Allows you to page through results. Only use on a follow-up query. Defaults to 0 (first page). + + Returns: + str: Query result string + """ + + import math + + from mirix.constants import RETRIEVAL_QUERY_DEFAULT_PAGE_SIZE + from mirix.utils import json_dumps + + if page is None or (isinstance(page, str) and page.lower().strip() == "none"): + page = 0 + try: + page = int(page) + except (ValueError, TypeError): + raise ValueError("'page' argument must be an integer") + count = RETRIEVAL_QUERY_DEFAULT_PAGE_SIZE + # TODO: add paging by page number. currently cursor only works with strings. + # original: start=page * count + messages = await self.message_manager.list_user_messages_for_agent( + agent_id=self.agent_state.id, + actor=self.actor, + query_text=query, + limit=count, + ) + total = len(messages) + num_pages = math.ceil(total / count) - 1 # 0 index + if len(messages) == 0: + results_str = "No results found." + else: + results_pref = f"Showing {len(messages)} of {total} results (page {page}/{num_pages}):" + results_formatted = [message.text for message in messages] + results_str = f"{results_pref} {json_dumps(results_formatted)}" + return results_str + + +async def search_in_memory( + self: "Agent", + memory_type: str, + query: str, + search_field: str, + search_method: str, + timezone_str: str, +) -> Optional[str]: + """ + Choose which memory to search. All memory types support multiple search methods with different performance characteristics. Most of the time, you should use search over 'details' for episodic memory and semantic memory, 'content' for resource memory (but for resource memory, `embedding` is not supported for content field so you have to use other search methods), 'description' for procedural memory. This is because these fields have the richest information and is more likely to contain the keywords/query. You can always start from a thorough search over the whole memory by setting memory_type as 'all' and search_field as 'null', and then narrow down to specific fields and specific memories. + + Args: + memory_type: The type of memory to search in. It should be chosen from the following: "episodic", "resource", "procedural", "knowledge_vault", "semantic", "all". Here "all" means searching in all the memories. + query: The keywords/query used to search in the memory. + search_field: The field to search in the memory. It should be chosen from the attributes of the corresponding memory. For "episodic" memory, it can be 'summary', 'details'; for "resource" memory, it can be 'summary', 'content'; for "procedural" memory, it can be 'description', 'instructions', 'entry_type', 'name' (default 'description'; note: 'embedding' search supports only 'description'/'instructions'); for "knowledge_vault", it can be 'secret_value', 'caption'; for semantic memory, it can be 'name', 'summary', 'details'. For "all", it should also be "null" as the system will search all memories with default fields. + search_method: The method to search in the memory. Choose from: + - 'bm25': BM25 ranking-based full-text search (fast and effective for keyword-based searches) + - 'embedding': Vector similarity search using embeddings (most powerful, good for conceptual matches) + - 'hybrid': **RECOMMENDED for "procedural" memory** - fuses 'bm25' and 'embedding' via Reciprocal Rank Fusion (EverOS-style). Best recall/precision for skill lookup. Only valid for memory_type='procedural'; for any other memory_type (including 'all') it is treated as 'embedding'. + + Returns: + str: Query result string + """ + + if not self.user: + raise ValueError("Can not search in memory. User is not set") + + # "hybrid" is procedural-only โ€” the other memory managers have no hybrid + # branch. Clamp to embedding for any non-procedural memory_type (including + # "all") BEFORE the field-validation checks below, so an unsupported combo + # like resource/content/hybrid resolves to embedding and then trips the same + # deterministic validation as resource/content/embedding โ€” rather than + # slipping past the guard and later hitting a missing embedding column. + # (Procedural hybrid self-embeds the query in the manager; no precompute.) + if search_method == "hybrid" and memory_type != "procedural": + search_method = "embedding" + + if memory_type == "resource" and search_field == "content" and search_method == "embedding": + raise ValueError("embedding is not supported for resource memory's 'content' field.") + if memory_type == "knowledge_vault" and search_field == "secret_value" and search_method == "embedding": + raise ValueError("embedding is not supported for knowledge_vault memory's 'secret_value' field.") + + if memory_type == "all": + search_field = "null" + + # Pre-compute embedding once if using embedding search (to avoid redundant embeddings) + embedded_text = None + if search_method == "embedding" and query: + import numpy as np + + from mirix.constants import MAX_EMBEDDING_DIM + from mirix.embeddings import embedding_model + + embedded_text = await (await embedding_model(self.agent_state.embedding_config)).get_text_embedding(query) + # Pad for episodic memory which requires MAX_EMBEDDING_DIM + embedded_text_padded = np.pad( + np.array(embedded_text), (0, MAX_EMBEDDING_DIM - len(embedded_text)), mode="constant" + ).tolist() + + if memory_type == "core": + # It means the model is an idiot, but we still return the results: + if self.blocks_in_memory: + return self.blocks_in_memory.compile(), len(self.blocks_in_memory.list_block_labels()) + return "", 0 + + if memory_type == "episodic" or memory_type == "all": + episodic_memory = await self.episodic_memory_manager.list_episodic_memory( + user=self.user, + agent_state=self.agent_state, + query=query, + embedded_text=embedded_text_padded if search_method == "embedding" and query else None, + search_field=search_field if search_field != "null" else "summary", + search_method=search_method, + limit=10, + timezone_str=timezone_str, + ) + formatted_results_from_episodic = [ + { + "memory_type": "episodic", + "id": x.id, + "timestamp": x.occurred_at, + "event_type": x.event_type, + "actor": x.actor, + "summary": x.summary, + "details": x.details, + } + for x in episodic_memory + ] + if memory_type == "episodic": + return formatted_results_from_episodic, len(formatted_results_from_episodic) + + if memory_type == "resource" or memory_type == "all": + resource_memories = await self.resource_memory_manager.list_resources( + user=self.user, + agent_state=self.agent_state, + query=query, + embedded_text=embedded_text if search_method == "embedding" and query else None, + search_field=( + search_field if search_field != "null" else ("summary" if search_method == "embedding" else "content") + ), + search_method=search_method, + limit=10, + timezone_str=timezone_str, + ) + formatted_results_resource = [ + { + "memory_type": "resource", + "id": x.id, + "resource_type": x.resource_type, + "summary": x.summary, + "content": x.content, + } + for x in resource_memories + ] + if memory_type == "resource": + return formatted_results_resource, len(formatted_results_resource) + + if memory_type == "procedural" or memory_type == "all": + procedural_memories = await self.procedural_memory_manager.list_procedures( + user=self.user, + agent_state=self.agent_state, + query=query, + embedded_text=embedded_text if search_method == "embedding" and query else None, + search_field=search_field if search_field != "null" else "description", + search_method=search_method, + limit=10, + timezone_str=timezone_str, + ) + formatted_results_procedural = [ + { + "memory_type": "procedural", + "id": x.id, + "entry_type": x.entry_type, + "name": x.name, + "description": x.description, + "instructions": x.instructions, + } + for x in procedural_memories + ] + if memory_type == "procedural": + return formatted_results_procedural, len(formatted_results_procedural) + + if memory_type == "knowledge_vault" or memory_type == "all": + knowledge_vault_memories = await self.knowledge_vault_manager.list_knowledge( + user=self.user, + agent_state=self.agent_state, + query=query, + embedded_text=embedded_text if search_method == "embedding" and query else None, + search_field=search_field if search_field != "null" else "caption", + search_method=search_method, + limit=10, + timezone_str=timezone_str, + ) + formatted_results_knowledge_vault = [ + { + "memory_type": "knowledge_vault", + "id": x.id, + "entry_type": x.entry_type, + "source": x.source, + "sensitivity": x.sensitivity, + "secret_value": x.secret_value, + "caption": x.caption, + } + for x in knowledge_vault_memories + ] + if memory_type == "knowledge_vault": + return formatted_results_knowledge_vault, len(formatted_results_knowledge_vault) + + if memory_type == "semantic" or memory_type == "all": + semantic_memories = await self.semantic_memory_manager.list_semantic_items( + user=self.user, + agent_state=self.agent_state, + query=query, + embedded_text=embedded_text if search_method == "embedding" and query else None, + search_field=search_field if search_field != "null" else "summary", + search_method=search_method, + limit=10, + timezone_str=timezone_str, + ) + # title, summary, details, source + formatted_results_semantic = [ + { + "memory_type": "semantic", + "id": x.id, + "name": x.name, + "summary": x.summary, + "details": x.details, + "source": x.source, + } + for x in semantic_memories + ] + if memory_type == "semantic": + return formatted_results_semantic, len(formatted_results_semantic) + + else: + raise ValueError( + f"Memory type '{memory_type}' is not supported. Please choose from 'episodic', 'resource', 'procedural', 'knowledge_vault', 'semantic'." + ) + return ( + formatted_results_from_episodic + + formatted_results_resource + + formatted_results_procedural + + formatted_results_knowledge_vault + + formatted_results_semantic, + len(formatted_results_from_episodic) + + len(formatted_results_resource) + + len(formatted_results_procedural) + + len(formatted_results_knowledge_vault) + + len(formatted_results_semantic), + ) + + +async def list_memory_within_timerange( + self: "Agent", memory_type: str, start_time: str, end_time: str, timezone_str: str +) -> Optional[str]: + """ + List memories around a specific timestamp + Args: + memory_type (str): The type of memory to search in. It should be chosen from the following: "episodic", "resource", "procedural", "knowledge_vault", "semantic", "all". Here "all" means searching in all the memories. + start_time (str): The start time of the time range. It has to be in the form of "%Y-%m-%d %H:%M:%S" + end_time (str): The end time of the time range. It has to be in the form of "%Y-%m-%d %H:%M:%S" + """ + + start_time = convert_timezone_to_utc(start_time, timezone_str) + end_time = convert_timezone_to_utc(end_time, timezone_str) + + if not self.user: + raise ValueError("Can not list memory within timerange. User is not set") + + if memory_type == "episodic" or memory_type == "all": + episodic_memory = await self.episodic_memory_manager.list_episodic_memory_around_timestamp( + user=self.user, + agent_state=self.agent_state, + start_time=start_time, + end_time=end_time, + timezone_str=timezone_str, + ) + formatted_results_from_episodic = [ + { + "memory_type": "episodic", + "id": x.id, + "timestamp": x.occurred_at, + "event_type": x.event_type, + "actor": x.actor, + "summary": x.summary, + } + for x in episodic_memory + ] + if memory_type == "episodic": + if len(formatted_results_from_episodic) == 0: + return "No results found." + elif len(formatted_results_from_episodic) > 50: + return "Too many results found. Please narrow down your search." + else: + return formatted_results_from_episodic, len(formatted_results_from_episodic) + + # currently only episodic memory is supported + return None diff --git a/mirix/functions/function_sets/memory_tools.py b/mirix/functions/function_sets/memory_tools.py index 6260c67e4..54c148858 100644 --- a/mirix/functions/function_sets/memory_tools.py +++ b/mirix/functions/function_sets/memory_tools.py @@ -20,7 +20,7 @@ from mirix.schemas.episodic_memory import EpisodicEventForLLM from mirix.schemas.knowledge_vault import KnowledgeVaultItemBase from mirix.schemas.mirix_message_content import TextContent -from mirix.schemas.procedural_memory import ProceduralMemoryItemBase + from mirix.schemas.resource_memory import ResourceMemoryItemBase from mirix.schemas.semantic_memory import SemanticMemoryItemBase @@ -42,7 +42,9 @@ async def core_memory_append( """ # check if the content starts with something like "Line n:" (here n is a number) using regex if re.match(r"^Line \d+:", content): - raise ValueError("You should not include 'Line n:' (here n is a number) in the content.") + raise ValueError( + "You should not include 'Line n:' (here n is a number) in the content." + ) # Get the current block and its limit current_block = blocks_in_memory.get_block(label) @@ -119,18 +121,26 @@ async def episodic_memory_insert(self: "Agent", items: List[EpisodicEventForLLM] Returns: Optional[str]: None is always returned as this function does not produce a response. """ - agent_id = self.agent_state.parent_id if self.agent_state.parent_id is not None else self.agent_state.id + agent_id = ( + self.agent_state.parent_id + if self.agent_state.parent_id is not None + else self.agent_state.id + ) # Get filter_tags, use_cache, client_id, user_id, and occurred_at from agent instance filter_tags = getattr(self, "filter_tags", None) use_cache = getattr(self, "use_cache", True) client_id = getattr(self, "client_id", None) user_id = getattr(self, "user_id", None) - occurred_at_override = getattr(self, "occurred_at", None) # Optional timestamp override from API + occurred_at_override = getattr( + self, "occurred_at", None + ) # Optional timestamp override from API for item in items: # Use occurred_at_override if provided, otherwise use LLM-extracted timestamp - timestamp = occurred_at_override if occurred_at_override else item["occurred_at"] + timestamp = ( + occurred_at_override if occurred_at_override else item["occurred_at"] + ) # Convert string to datetime if needed if isinstance(timestamp, str): @@ -155,8 +165,7 @@ async def episodic_memory_insert(self: "Agent", items: List[EpisodicEventForLLM] ) except Exception as e: print( - f"[episodic_memory_insert] insert_event FAILED for item " - f"{item!r}: {e}" + f"[episodic_memory_insert] insert_event FAILED for item {item!r}: {e}" ) traceback.print_exc() raise @@ -213,7 +222,9 @@ async def episodic_memory_merge( return response -async def episodic_memory_replace(self: "Agent", event_ids: List[str], new_items: List[EpisodicEventForLLM]): +async def episodic_memory_replace( + self: "Agent", event_ids: List[str], new_items: List[EpisodicEventForLLM] +): """ The tool to replace or delete items in the episodic memory. To replace the memory, set the event_ids to be the ids of the events that needs to be replaced and new_items as the updated events. Note that the number of new items does not need to be the same as the number of event_ids as it is not a one-to-one mapping. To delete the memory, set the event_ids to be the ids of the events that needs to be deleted and new_items as an empty list. To insert new events, use episodic_memory_insert function. @@ -221,14 +232,20 @@ async def episodic_memory_replace(self: "Agent", event_ids: List[str], new_items event_ids (str): The ids of the episodic events to be deleted (or replaced). new_items (array): List of new episodic memory items to insert. If this is an empty list, then it means that the items are being deleted. """ - agent_id = self.agent_state.parent_id if self.agent_state.parent_id is not None else self.agent_state.id + agent_id = ( + self.agent_state.parent_id + if self.agent_state.parent_id is not None + else self.agent_state.id + ) # Get filter_tags, use_cache, client_id, user_id, and occurred_at from agent instance filter_tags = getattr(self, "filter_tags", None) use_cache = getattr(self, "use_cache", True) client_id = getattr(self, "client_id", None) user_id = getattr(self, "user_id", None) - occurred_at_override = getattr(self, "occurred_at", None) # Optional timestamp override from API + occurred_at_override = getattr( + self, "occurred_at", None + ) # Optional timestamp override from API if self.user is None: raise ValueError("User is required to access episodic memory") @@ -238,11 +255,15 @@ async def episodic_memory_replace(self: "Agent", event_ids: List[str], new_items for event_id in event_ids: # It will raise an error if the event_id is not found in the episodic memory. - await self.episodic_memory_manager.get_episodic_memory_by_id(event_id, user=self.user) + await self.episodic_memory_manager.get_episodic_memory_by_id( + event_id, user=self.user + ) for event_id in event_ids: try: - await self.episodic_memory_manager.delete_event_by_id(event_id, actor=self.actor) + await self.episodic_memory_manager.delete_event_by_id( + event_id, actor=self.actor + ) except Exception as e: print( f"[episodic_memory_replace] delete_event_by_id FAILED for " @@ -253,7 +274,9 @@ async def episodic_memory_replace(self: "Agent", event_ids: List[str], new_items for new_item in new_items: # Use occurred_at_override if provided, otherwise use LLM-extracted timestamp - timestamp = occurred_at_override if occurred_at_override else new_item["occurred_at"] + timestamp = ( + occurred_at_override if occurred_at_override else new_item["occurred_at"] + ) # Convert string to datetime if needed if isinstance(timestamp, str): @@ -285,7 +308,9 @@ async def episodic_memory_replace(self: "Agent", event_ids: List[str], new_items raise -async def check_episodic_memory(self: "Agent", event_ids: List[str], timezone_str: str) -> List[EpisodicEventForLLM]: +async def check_episodic_memory( + self: "Agent", event_ids: List[str], timezone_str: str +) -> List[EpisodicEventForLLM]: """ The tool to check the episodic memory. This function will return the episodic events with the given event_ids. @@ -299,7 +324,9 @@ async def check_episodic_memory(self: "Agent", event_ids: List[str], timezone_st raise ValueError("User is required to check episodic memory") episodic_memory = [ - await self.episodic_memory_manager.get_episodic_memory_by_id(event_id, user=self.user, timezone_str=timezone_str) + await self.episodic_memory_manager.get_episodic_memory_by_id( + event_id, user=self.user, timezone_str=timezone_str + ) for event_id in event_ids ] @@ -330,7 +357,11 @@ async def resource_memory_insert(self: "Agent", items: List[ResourceMemoryItemBa """ # No imports needed - using agent instance attributes - agent_id = self.agent_state.parent_id if self.agent_state.parent_id is not None else self.agent_state.id + agent_id = ( + self.agent_state.parent_id + if self.agent_state.parent_id is not None + else self.agent_state.id + ) # Get filter_tags, use_cache, client_id, and user_id from agent instance filter_tags = getattr(self, "filter_tags", None) @@ -402,7 +433,9 @@ async def resource_memory_insert(self: "Agent", items: List[ResourceMemoryItemBa return "No resources were inserted." -async def resource_memory_update(self: "Agent", old_ids: List[str], new_items: List[ResourceMemoryItemBase]): +async def resource_memory_update( + self: "Agent", old_ids: List[str], new_items: List[ResourceMemoryItemBase] +): """ The tool to update and delete items in the resource memory. To update the memory, set the old_ids to be the ids of the items that needs to be updated and new_items as the updated items. Note that the number of new items does not need to be the same as the number of old ids as it is not a one-to-one mapping. To delete the memory, set the old_ids to be the ids of the items that needs to be deleted and new_items as an empty list. @@ -410,7 +443,11 @@ async def resource_memory_update(self: "Agent", old_ids: List[str], new_items: L old_ids (array): List of ids of the items to be deleted (or updated). new_items (array): List of new resource memory items to insert. If this is an empty list, then it means that the items are being deleted. """ - agent_id = self.agent_state.parent_id if self.agent_state.parent_id is not None else self.agent_state.id + agent_id = ( + self.agent_state.parent_id + if self.agent_state.parent_id is not None + else self.agent_state.id + ) # Get filter_tags, use_cache, client_id, and user_id from agent instance filter_tags = getattr(self, "filter_tags", None) @@ -419,7 +456,9 @@ async def resource_memory_update(self: "Agent", old_ids: List[str], new_items: L user_id = getattr(self, "user_id", None) for old_id in old_ids: - await self.resource_memory_manager.delete_resource_by_id(resource_id=old_id, actor=self.actor) + await self.resource_memory_manager.delete_resource_by_id( + resource_id=old_id, actor=self.actor + ) for item in new_items: await self.resource_memory_manager.insert_resource( @@ -437,134 +476,615 @@ async def resource_memory_update(self: "Agent", old_ids: List[str], new_items: L ) -async def procedural_memory_insert(self: "Agent", items: List[ProceduralMemoryItemBase]): +def _bump_patch_version(version: str) -> str: + """Increment patch version: '0.1.0' -> '0.1.1'. + + Args: + version (str): Semver version string to bump. + """ + try: + parts = version.split(".") + parts[-1] = str(int(parts[-1]) + 1) + return ".".join(parts) + except (ValueError, IndexError): + return "0.1.1" + + +# --------------------------------------------------------------------------- +# Edit-budget formula + per-mutation size/delete gates. +# +# These are pure, side-effect-free helpers so they are unit-testable without a +# DB, a server, or an LLM. The skill tools below consult them; the experience- +# curator evolve path (mirix.services.skill_experience_curator) sets the +# per-instance budget counters BEFORE running the procedural agent. +# --------------------------------------------------------------------------- + +# Sentinel returned by skill_create/skill_edit when the per-run mutation budget +# is exhausted. The curator's system prompt tells the agent that seeing this is +# its cue to stop and call finish_memory_update. +_BUDGET_EXHAUSTED_MSG = ( + "edit budget exhausted โ€” call finish_memory_update to end this evolution. " + "No further skill_create / skill_edit will take effect this run." +) + + +def _round_half_up(value: float) -> int: + """Deterministic round-half-up (Python's built-in round() is banker's). + + The design formula is clamp(round(raw)); we use round-half-up so e.g. a + successes-only window (raw=2.5) lands at 3 instead of 2, matching the + documented "mixed => 2-4" budget distribution. + """ + import math + + return int(math.floor(value + 0.5)) + + +def compute_edit_budget(aggregate: dict, *, autonomous: Optional[int] = None) -> int: + """Count-driven edit budget from the C2 `aggregate` (P0-3). + + ``raw = B0 + alpha_f*n_high_fail + alpha_s*n_high_succ`` clamped to + ``[B_min, B_max]``. Only structure-gated counts drive it; ``quality_score`` + (``mean_q``) is RANKING-ONLY and never read here. + + Hybrid (USER-LOCKED): when ``autonomous`` is given (a cheap LLM's suggestion) + it may only REDUCE the ceiling โ€” ``final = min(formula, clamp(autonomous))``. + For the validity run the caller passes ``autonomous=None`` (formula-only). + """ + from mirix.constants import ( + SKILL_EDIT_BUDGET_ALPHA_FAIL, + SKILL_EDIT_BUDGET_ALPHA_SUCC, + SKILL_EDIT_BUDGET_B0, + SKILL_EDIT_BUDGET_MAX, + SKILL_EDIT_BUDGET_MIN, + ) + + n_high_fail = int(aggregate.get("n_high_fail", 0) or 0) + n_high_succ = int(aggregate.get("n_high_succ", 0) or 0) + + raw = ( + SKILL_EDIT_BUDGET_B0 + + SKILL_EDIT_BUDGET_ALPHA_FAIL * n_high_fail + + SKILL_EDIT_BUDGET_ALPHA_SUCC * n_high_succ + ) + formula = _round_half_up(raw) + formula = max(SKILL_EDIT_BUDGET_MIN, min(SKILL_EDIT_BUDGET_MAX, formula)) + + # B_min=0 skip: a window with no structure-gated records yields 0 (the + # curator keys its skip decision off this value via the aggregate). + if n_high_fail + n_high_succ == 0: + return SKILL_EDIT_BUDGET_MIN + + if autonomous is None: + return formula + + auto = max(SKILL_EDIT_BUDGET_MIN, min(SKILL_EDIT_BUDGET_MAX, int(autonomous))) + return min(formula, auto) + + +def _edit_exceeds_size_gate(old_text: str, new_text: str) -> Optional[str]: + """Per-mutation HARD size gate (P1-3) for a TEXT-field edit. + + Returns a rejection message if the edit is too large on EITHER dimension: + * absolute char delta > SKILL_MAX_EDIT_CHAR_DELTA (default 800), or + * change-ratio (1 - SequenceMatcher.ratio()) >= SKILL_EDIT_MAJOR_RATIO (0.4) + Otherwise returns None (the edit is small enough to apply). + """ + import difflib + + from mirix.constants import SKILL_EDIT_MAJOR_RATIO, SKILL_MAX_EDIT_CHAR_DELTA + + old_text = old_text or "" + new_text = new_text or "" + + char_delta = abs(len(new_text) - len(old_text)) + if char_delta > SKILL_MAX_EDIT_CHAR_DELTA: + return ( + f"edit too large; split or extract a new skill " + f"(char delta {char_delta} > {SKILL_MAX_EDIT_CHAR_DELTA})." + ) + + ratio_change = 1.0 - difflib.SequenceMatcher(None, old_text, new_text).ratio() + if ratio_change >= SKILL_EDIT_MAJOR_RATIO: + return ( + f"edit too large; split or extract a new skill " + f"(change ratio {ratio_change:.2f} >= {SKILL_EDIT_MAJOR_RATIO})." + ) + return None + + +def _instructions_over_ceiling(text: str) -> bool: + """True when an `instructions` value exceeds the hard ceiling. + + Over-ceiling edits must route to skill_create (extract a new skill) instead + of growing an existing skill unbounded. + """ + from mirix.constants import SKILL_MAX_INSTRUCTIONS_CHARS + + return len(text or "") > SKILL_MAX_INSTRUCTIONS_CHARS + + +# Phrases in a failure record's detail/root_cause that signal a skill is +# ACTIVELY HARMFUL (not merely redundant/outdated). Only these authorize a +# delete. "redundant" / "duplicate" deliberately do NOT โ€” those are edit/merge +# territory, never a destructive delete. +_HARMFUL_MARKERS = ( + "actively harmful", + "is harmful", + "harmful", + "caused the failure", + "led to the failure", + "should be removed", + "should be deleted", + "misleading", + "wrong namespace", +) + + +def _record_authorizes_delete(record, *, skill_name: str, skill_id: str) -> bool: + """Whether a single failure record authorizes deleting a specific skill. + + Requires ALL of: + * the record is a FAILURE (a success never authorizes a delete), + * its detail/root_cause NAMES the skill (by name or id), and + * its detail flags the skill as actively HARMFUL (not merely redundant). + + Accepts an ORM row, a pydantic record, or a plain dict (duck-typed). + """ + + def _get(obj, key): + if isinstance(obj, dict): + return obj.get(key) + return getattr(obj, key, None) + + if (_get(record, "record_type") or "") != "failure": + return False + + detail = _get(record, "detail") or "" + + # Require the skill MENTION and a HARMFUL marker to co-occur in the SAME + # clause/sentence, not merely somewhere in the detail. Otherwise a record + # like "skill A is redundant; skill B was harmful" would wrongly authorize + # deleting A. Split on sentence/clause boundaries and require one segment to + # contain both the skill reference and a harmful marker. + import re as _re + + name_l = (skill_name or "").lower() + id_l = (skill_id or "").lower() + if not name_l and not id_l: + return False + + def _mentions(seg: str, token: str) -> bool: + # Whole-token match so 'proc-9' does NOT match inside 'proc-90' (skill + # ids/names contain hyphens, so \b is unreliable). Require the token to + # be bounded by a non [A-Za-z0-9-] char (or string edge) on both sides. + if not token: + return False + for m in _re.finditer(_re.escape(token), seg): + before = seg[m.start() - 1] if m.start() > 0 else "" + after = seg[m.end()] if m.end() < len(seg) else "" + if not _re.match(r"[A-Za-z0-9-]", before) and not _re.match( + r"[A-Za-z0-9-]", after + ): + return True + return False + + for seg in _re.split(r"[.;\n!?]+", detail.lower()): + names_skill = _mentions(seg, name_l) or _mentions(seg, id_l) + if not names_skill: + continue + if any(marker in seg for marker in _HARMFUL_MARKERS): + return True + return False + + +def _consume_edit_budget(self: "Agent") -> Optional[str]: + """Decrement the per-instance create/edit budget at the TOP of a mutation. + + Counter is the plain instance attr ``self._edit_budget_remaining`` (P1-2): + NOT module-level, NOT agent_id-keyed, NOT manager-keyed โ€” so two agents (two + users / two runs) never share it, and a single async task is race-free. + + Returns ``None`` when the mutation may proceed (and decrements), or the + exhausted-message when the budget is 0 (no decrement, caller must NOT + mutate). An UNSET attribute means "no limit" (logged once per call). + """ + remaining = getattr(self, "_edit_budget_remaining", None) + if remaining is None: + logger.debug("[skill budget] no _edit_budget_remaining set -> no limit") + return None + if remaining <= 0: + return _BUDGET_EXHAUSTED_MSG + self._edit_budget_remaining = remaining - 1 + return None + + +async def skill_list(self: "Agent", query: str = "", limit: int = 50) -> str: """ - The tool to insert new procedures into procedural memory. Note that the `summary` should not be a general term such as "guide" or "workflow" but rather a more informative description of the procedure. + List skills with optional search query. Args: - items (array): List of procedural memory items to insert. + query (str): Optional search term to filter skills. Empty string returns all skills. + limit (int): Maximum number of skills to return. Default 50. Returns: - Optional[str]: Message about insertion results including any duplicates detected. + str: Formatted list of skills with id, name, description, entry_type, and version. """ - agent_id = self.agent_state.parent_id if self.agent_state.parent_id is not None else self.agent_state.id + filter_tags = getattr(self, "filter_tags", None) + use_cache = getattr(self, "use_cache", True) - # Get filter_tags, use_cache, client_id, and user_id from agent instance + skills = await self.procedural_memory_manager.list_procedures( + agent_state=self.agent_state, + user=self.user, + query=query, + search_field="description" if query else "", + search_method="bm25" if query else "", + limit=limit, + filter_tags=filter_tags if filter_tags else None, + use_cache=use_cache, + ) + + if not skills: + return "No skills found." + + lines = [] + for skill in skills: + lines.append( + f"[ID: {skill.id}] {skill.name} (v{skill.version}) - {skill.entry_type}: {skill.description}" + ) + return "\n".join(lines) + + +async def skill_read(self: "Agent", name_or_id: str) -> str: + """ + Read the complete content of a skill by name or ID. + + Args: + name_or_id (str): The skill name (e.g. "deploy-production") or ID (e.g. "proc-xxx"). + + Returns: + str: Full skill content including all fields. + """ filter_tags = getattr(self, "filter_tags", None) use_cache = getattr(self, "use_cache", True) - client_id = getattr(self, "client_id", None) - user_id = getattr(self, "user_id", None) + skill = None - inserted_count = 0 - skipped_count = 0 - skipped_summaries = [] + if name_or_id.startswith("proc"): + try: + skill = await self.procedural_memory_manager.get_item_by_id( + item_id=name_or_id, user=self.user, timezone_str="UTC", actor=self.actor + ) + except Exception: + pass - for item in items: - # Check for existing similar procedures (by summary and filter_tags) - existing_procedures = await self.procedural_memory_manager.list_procedures( + if skill is None: + results = await self.procedural_memory_manager.list_procedures( agent_state=self.agent_state, - user=self.user, # User for read operations (data filtering) - query="", # Get all procedures - limit=1000, # Get enough to check for duplicates + user=self.user, + query=name_or_id, + search_field="name", + search_method="string_match", + limit=1, filter_tags=filter_tags if filter_tags else None, use_cache=use_cache, ) + if results: + skill = results[0] + + if skill is None: + return f"Skill '{name_or_id}' not found." + + parts = [ + f"ID: {skill.id}", + f"Name: {skill.name}", + f"Version: {skill.version}", + f"Entry Type: {skill.entry_type}", + f"Description: {skill.description}", + f"Instructions:\n{skill.instructions}", + f"Triggers: {skill.triggers}", + f"Examples: {skill.examples}", + ] + return "\n".join(parts) - # Check if this procedure already exists - is_duplicate = False - for existing in existing_procedures: - if existing.summary == item["summary"] and existing.steps == item["steps"]: - is_duplicate = True - skipped_count += 1 - skipped_summaries.append(item["summary"]) - break - if not is_duplicate: - try: - await self.procedural_memory_manager.insert_procedure( - agent_state=self.agent_state, - agent_id=agent_id, - entry_type=item["entry_type"], - summary=item["summary"], - steps=item["steps"], - actor=self.actor, - organization_id=self.user.organization_id, - filter_tags=filter_tags if filter_tags else None, - use_cache=use_cache, - user_id=user_id, - ) - except Exception as e: - print( - f"[procedural_memory_insert] insert_procedure FAILED for " - f"item {item!r}: {e}" - ) - traceback.print_exc() - raise - inserted_count += 1 +async def skill_create( + self: "Agent", + name: str, + description: str, + instructions: str, + entry_type: str, + triggers: List[str] = None, + examples: List[dict] = None, +) -> str: + """ + Create a new skill. Fails if a skill with the same name already exists. - # Return feedback message - if skipped_count > 0: - skipped_list = ", ".join(f"'{s}'" for s in skipped_summaries[:3]) - if len(skipped_summaries) > 3: - skipped_list += f" and {len(skipped_summaries) - 3} more" - return f"Inserted {inserted_count} new procedure(s). Skipped {skipped_count} duplicate(s): {skipped_list}." - elif inserted_count > 0: - return f"Successfully inserted {inserted_count} new procedure(s)." - else: - return "No procedures were inserted." + Args: + name (str): Concise kebab-case identifier (e.g. "deploy-production"). + description (str): What this skill does and when it's useful. + instructions (str): Detailed instructions as plain text. + entry_type (str): Category โ€” "workflow", "guide", or "script". + triggers (array): Optional list of conditions that indicate this skill is relevant. + examples (array): Optional list of input/output examples. + Returns: + str: Confirmation with the created skill's ID. + """ + # Edit-budget gate (counts create+edit). At the TOP, before any mutation. + exhausted = _consume_edit_budget(self) + if exhausted is not None: + return exhausted + + agent_id = ( + self.agent_state.parent_id + if self.agent_state.parent_id is not None + else self.agent_state.id + ) + filter_tags = getattr(self, "filter_tags", None) + use_cache = getattr(self, "use_cache", True) + user_id = getattr(self, "user_id", None) -async def procedural_memory_update(self: "Agent", old_ids: List[str], new_items: List[ProceduralMemoryItemBase]): + # Name dedup check + existing = await self.procedural_memory_manager.list_procedures( + agent_state=self.agent_state, + user=self.user, + query=name, + search_field="name", + search_method="string_match", + limit=100, + filter_tags=filter_tags if filter_tags else None, + use_cache=use_cache, + ) + for skill in existing: + if skill.name == name: + return f"Skill '{name}' already exists (ID: {skill.id}). Use skill_edit to modify it, or skill_delete to remove it first." + + try: + from mirix.orm.errors import UniqueConstraintViolationError + + created = await self.procedural_memory_manager.insert_procedure( + agent_state=self.agent_state, + agent_id=agent_id, + name=name, + description=description, + instructions=instructions, + entry_type=entry_type, + triggers=triggers or [], + examples=examples or [], + version="0.1.0", + actor=self.actor, + organization_id=self.user.organization_id, + filter_tags=filter_tags if filter_tags else None, + use_cache=use_cache, + user_id=user_id, + ) + return ( + f"Skill '{name}' created successfully (ID: {created.id}, version: 0.1.0)." + ) + except UniqueConstraintViolationError: + # Race: pre-check missed a concurrent create. Surface a clean message + # to the agent so it knows to read/edit the existing skill instead. + return ( + f"Skill '{name}' already exists. Use skill_read to inspect it, " + f"skill_edit to modify it, or skill_delete to remove it first." + ) + except Exception as e: + logger.error("[skill_create] FAILED for '%s': %s", name, e) + traceback.print_exc() + raise + + +async def skill_edit( + self: "Agent", + skill_id: str, + field: str, + old_text: str = None, + new_text: str = None, + value: str = None, +) -> str: """ - The tool to update/delete items in the procedural memory. To update the memory, set the old_ids to be the ids of the items that needs to be updated and new_items as the updated items. Note that the number of new items does not need to be the same as the number of old ids as it is not a one-to-one mapping. To delete the memory, set the old_ids to be the ids of the items that needs to be deleted and new_items as an empty list. + Edit an existing skill. For text fields (name, description, instructions), use old_text/new_text + to do a find-and-replace patch. For other fields (triggers, examples, entry_type), use value + to replace the entire field. For list/dict fields (triggers, examples), value must be a JSON string. Args: - old_ids (array): List of ids of the items to be deleted (or updated). - new_items (array): List of new procedural memory items to insert. If this is an empty list, then it means that the items are being deleted. + skill_id (str): The ID of the skill to edit. + field (str): The field to modify โ€” "name", "description", "instructions", "entry_type", "triggers", or "examples". + old_text (str): For text fields: the text to find and replace. + new_text (str): For text fields: the replacement text. + value (str): For non-text fields: the new value for the entire field. + - entry_type: a plain string such as "workflow" / "guide" / "script". + - triggers: a JSON-encoded list of strings, e.g. '["on error", "nightly"]'. + - examples: a JSON-encoded list of objects, e.g. '[{"input": "...", "output": "..."}]'. Returns: - Optional[str]: None is always returned as this function does not produce a response. + str: Confirmation of the edit with the new version number. """ - agent_id = self.agent_state.parent_id if self.agent_state.parent_id is not None else self.agent_state.id + import json as _json - # Get filter_tags, use_cache, client_id, and user_id from agent instance - filter_tags = getattr(self, "filter_tags", None) - use_cache = getattr(self, "use_cache", True) - client_id = getattr(self, "client_id", None) - user_id = getattr(self, "user_id", None) + from mirix.schemas.procedural_memory import ProceduralMemoryItemUpdate - for old_id in old_ids: - try: - await self.procedural_memory_manager.delete_procedure_by_id(procedure_id=old_id, actor=self.actor) - except Exception as e: - print( - f"[procedural_memory_update] delete_procedure_by_id FAILED for " - f"old_id {old_id!r}: {e}" + text_fields = {"name", "description", "instructions"} + json_list_fields = {"triggers", "examples"} + plain_value_fields = {"entry_type"} + valid_fields = text_fields | json_list_fields | plain_value_fields + + if field not in valid_fields: + return f"Invalid field '{field}'. Must be one of: {', '.join(sorted(valid_fields))}." + + try: + skill = await self.procedural_memory_manager.get_item_by_id( + item_id=skill_id, user=self.user, timezone_str="UTC", actor=self.actor + ) + except Exception: + return f"Skill '{skill_id}' not found." + + update_data = {"id": skill_id} + + if field in text_fields: + if old_text is None or new_text is None: + return f"For text field '{field}', both old_text and new_text are required." + + current_value = getattr(skill, field, "") + + # Normalize whitespace for matching + normalized_current = " ".join(current_value.split()) + normalized_old = " ".join(old_text.split()) + + if normalized_old not in normalized_current: + return ( + f"old_text not found in field '{field}'. " + f"Current value:\n{current_value}" ) - traceback.print_exc() - raise - for item in new_items: + import re + + pattern = re.escape(old_text) + pattern = re.sub(r"\\\s+", r"\\s+", pattern) + new_value = re.sub(pattern, new_text, current_value, count=1) + + # Per-mutation HARD size gate. Runs BEFORE the budget is + # consumed so a rejected (too-large) edit is FREE โ€” the agent is told to + # split it or extract a new skill instead. Only text fields are gated. + size_reason = _edit_exceeds_size_gate(current_value, new_value) + if size_reason is not None: + return size_reason + + # Hard `instructions` ceiling: an edit whose RESULT exceeds the ceiling + # must route to skill_create (extract a new skill), not grow this one. + if field == "instructions" and _instructions_over_ceiling(new_value): + from mirix.constants import SKILL_MAX_INSTRUCTIONS_CHARS + + return ( + f"resulting instructions ({len(new_value)} chars) exceed the " + f"{SKILL_MAX_INSTRUCTIONS_CHARS}-char ceiling; use skill_create to " + f"extract a new skill instead of growing this one." + ) + + update_data[field] = new_value + elif field in json_list_fields: + if value is None: + return ( + f"For field '{field}', the value parameter is required (JSON string)." + ) + # triggers/examples are typed List[str] / List[dict] in the schema. + # The agent delivers them as JSON strings, so decode + validate here + # rather than letting the raw string hit Pydantic and blow up. try: - await self.procedural_memory_manager.insert_procedure( - agent_state=self.agent_state, - agent_id=agent_id, - entry_type=item["entry_type"], - summary=item["summary"], - steps=item["steps"], - actor=self.actor, - organization_id=self.actor.organization_id, - filter_tags=filter_tags if filter_tags else None, - use_cache=use_cache, - user_id=user_id, + decoded = _json.loads(value) + except (TypeError, _json.JSONDecodeError) as exc: + return ( + f"Invalid JSON for field '{field}': {exc}. " + f'Pass a JSON-encoded list, e.g. \'["trigger-a", "trigger-b"]\'.' ) - except Exception as e: - print( - f"[procedural_memory_update] insert_procedure FAILED for item " - f"{item!r}: {e}" + if not isinstance(decoded, list): + return ( + f"Field '{field}' must be a JSON array, got {type(decoded).__name__}." ) - traceback.print_exc() - raise + if field == "triggers": + if not all(isinstance(x, str) for x in decoded): + return "Field 'triggers' must be a JSON array of strings." + else: # examples + if not all(isinstance(x, dict) for x in decoded): + return "Field 'examples' must be a JSON array of objects." + update_data[field] = decoded + else: + # plain_value_fields: entry_type + if value is None: + return f"For field '{field}', the value parameter is required." + update_data[field] = value + + new_version = _bump_patch_version(getattr(skill, "version", "0.1.0")) + update_data["version"] = new_version + + # Edit-budget gate (counts create+edit). Consumed only AFTER every validation + # and size gate has passed, so a rejected edit never costs a budget unit. + exhausted = _consume_edit_budget(self) + if exhausted is not None: + return exhausted + + try: + updated = await self.procedural_memory_manager.update_item( + item_update=ProceduralMemoryItemUpdate.model_validate(update_data), + user=self.user, + actor=self.actor, + agent_state=self.agent_state, + ) + return f"Skill '{updated.name}' updated (field: {field}, new version: {new_version})." + except Exception as e: + logger.error("[skill_edit] FAILED for '%s': %s", skill_id, e) + traceback.print_exc() + raise + + +async def skill_delete(self: "Agent", skill_id: str) -> str: + """ + Delete a skill by ID. + + Under an experience-curator run deletes are gated SEPARATELY from + the create/edit budget: at most ``D_max`` per evolution, and only when a + failure record named this skill actively harmful. Soft-delete (exclude from + retrieval, reversible) is preferred over a hard delete. The gate engages only + when the curator has set the delete-gate instance attributes; the in-flow + chat path (which sets none of them) keeps the original ungated behavior. + + Args: + skill_id (str): The ID of the skill to delete. + + Returns: + str: Confirmation message. + """ + # Gating is active iff the curator set up either the authorization set or the + # delete budget on this instance. Otherwise behave exactly as before. + has_auth_attr = hasattr(self, "_delete_authorized_skill_ids") + has_budget_attr = hasattr(self, "_delete_budget_remaining") + gated = has_auth_attr or has_budget_attr + + if gated: + # 1) D_max: if a delete budget is set, it must be > 0. (Unset => no cap.) + if has_budget_attr: + remaining = getattr(self, "_delete_budget_remaining", None) + if remaining is not None and remaining <= 0: + return ( + "delete budget exhausted (D_max reached) โ€” deletes are capped " + "per evolution; call finish_memory_update." + ) + + # 2) Authorization: a failure record must have named THIS skill harmful. + authorized = getattr(self, "_delete_authorized_skill_ids", None) or set() + if skill_id not in authorized: + return ( + f"delete not authorized: no failure record names skill " + f"'{skill_id}' as actively harmful. Prefer skill_edit to fix it, " + f"or leave it in place." + ) + + prefer_soft = getattr(self, "_prefer_soft_delete", False) + + try: + if prefer_soft and hasattr( + self.procedural_memory_manager, "soft_delete_procedure_by_id" + ): + await self.procedural_memory_manager.soft_delete_procedure_by_id( + procedure_id=skill_id, actor=self.actor, user=self.user + ) + outcome = f"Skill '{skill_id}' soft-deleted (excluded from retrieval)." + else: + await self.procedural_memory_manager.delete_procedure_by_id( + procedure_id=skill_id, actor=self.actor, user=self.user + ) + outcome = f"Skill '{skill_id}' deleted successfully." + except Exception as e: + logger.error("[skill_delete] FAILED for '%s': %s", skill_id, e) + return f"Failed to delete skill '{skill_id}': {e}" + + # Charge the delete budget only on a SUCCESSFUL gated delete. + if gated and has_budget_attr: + remaining = getattr(self, "_delete_budget_remaining", None) + if remaining is not None: + self._delete_budget_remaining = remaining - 1 + + return outcome async def check_semantic_memory( @@ -610,7 +1130,11 @@ async def semantic_memory_insert(self: "Agent", items: List[SemanticMemoryItemBa Returns: Optional[str]: Message about insertion results including any duplicates detected. """ - agent_id = self.agent_state.parent_id if self.agent_state.parent_id is not None else self.agent_state.id + agent_id = ( + self.agent_state.parent_id + if self.agent_state.parent_id is not None + else self.agent_state.id + ) # Get filter_tags, use_cache, client_id, and user_id from agent instance filter_tags = getattr(self, "filter_tags", None) @@ -697,7 +1221,11 @@ async def semantic_memory_update( Returns: Optional[str]: None is always returned as this function does not produce a response. """ - agent_id = self.agent_state.parent_id if self.agent_state.parent_id is not None else self.agent_state.id + agent_id = ( + self.agent_state.parent_id + if self.agent_state.parent_id is not None + else self.agent_state.id + ) # Get filter_tags, use_cache, client_id, and user_id from agent instance filter_tags = getattr(self, "filter_tags", None) @@ -761,7 +1289,11 @@ async def knowledge_vault_insert(self: "Agent", items: List[KnowledgeVaultItemBa Returns: Optional[str]: Message about insertion results including any duplicates detected. """ - agent_id = self.agent_state.parent_id if self.agent_state.parent_id is not None else self.agent_state.id + agent_id = ( + self.agent_state.parent_id + if self.agent_state.parent_id is not None + else self.agent_state.id + ) # Get filter_tags, use_cache, client_id, and user_id from agent instance filter_tags = getattr(self, "filter_tags", None) @@ -834,7 +1366,9 @@ async def knowledge_vault_insert(self: "Agent", items: List[KnowledgeVaultItemBa return "No knowledge vault items were inserted." -async def knowledge_vault_update(self: "Agent", old_ids: List[str], new_items: List[KnowledgeVaultItemBase]): +async def knowledge_vault_update( + self: "Agent", old_ids: List[str], new_items: List[KnowledgeVaultItemBase] +): """ The tool to update/delete items in the knowledge vault. To update the knowledge_vault, set the old_ids to be the ids of the items that needs to be updated and new_items as the updated items. Note that the number of new items does not need to be the same as the number of old ids as it is not a one-to-one mapping. To delete the memory, set the old_ids to be the ids of the items that needs to be deleted and new_items as an empty list. @@ -845,7 +1379,11 @@ async def knowledge_vault_update(self: "Agent", old_ids: List[str], new_items: L Returns: Optional[str]: None is always returned as this function does not produce a response """ - agent_id = self.agent_state.parent_id if self.agent_state.parent_id is not None else self.agent_state.id + agent_id = ( + self.agent_state.parent_id + if self.agent_state.parent_id is not None + else self.agent_state.id + ) # Get filter_tags, use_cache, client_id, and user_id from agent instance filter_tags = getattr(self, "filter_tags", None) @@ -891,6 +1429,112 @@ async def knowledge_vault_update(self: "Agent", old_ids: List[str], new_items: L raise +async def _resolve_meta_agent_state(self: "Agent"): + """Resolve the meta_memory_agent AgentState that owns the retained messages. + + Goal-1 retention keeps the last-N sessions' raw messages on the + ``meta_memory_agent``. This tool runs inside a memory sub-agent, whose + ``parent_id`` is that meta agent. If this is already the meta agent, return + it; otherwise fetch the parent. Returns ``None`` when it cannot be resolved. + """ + from mirix.schemas.agent import AgentType + + state = getattr(self, "agent_state", None) + if state is None: + return None + try: + if state.is_type(AgentType.meta_memory_agent): + return state + except Exception: # noqa: BLE001 โ€” is_type guard, never fatal + pass + parent_id = getattr(state, "parent_id", None) + if not parent_id: + return None + try: + parent = await self.agent_manager.get_agent_by_id( + agent_id=parent_id, actor=self.actor + ) + except Exception: # noqa: BLE001 + return None + try: + if parent is not None and parent.is_type(AgentType.meta_memory_agent): + return parent + except Exception: # noqa: BLE001 + pass + return parent + + +async def _schedule_experience_auto_dream(self: "Agent", *, threshold: int) -> bool: + """Fire-and-forget the GENERAL session-experience auto-dream + (session distillation + skill evolution). + + Resolves the meta agent (owner of the retained transcripts) and schedules + ``AutoDreamManager().run(mode='procedural', last_n_sessions=threshold)`` as a + background task. Guarded end-to-end: any failure (resolution, scheduling, or + the dream itself) is logged and swallowed so the chat turn is never broken. + + Returns ``True`` iff the background dream task was actually scheduled. The + caller uses this to decide whether to ALSO enqueue the legacy "procedural" + memory update: when the dream is scheduled it owns the procedural agent + (skill evolution runs under the per-agent lock), so the legacy update must + NOT run concurrently on the same agent. When scheduling did not take (no + user/actor, unresolved meta agent, or an error), the caller falls back to + the legacy update so a fire is never silently dropped. + """ + if getattr(self, "user", None) is None or getattr(self, "actor", None) is None: + logger.debug("Skipping experience auto-dream: no user/actor in scope.") + return False + try: + meta_agent_state = await _resolve_meta_agent_state(self) + if meta_agent_state is None: + logger.debug("Skipping experience auto-dream: could not resolve meta agent.") + return False + + user = self.user + actor = self.actor + + async def _run_dream(): + try: + from mirix.schemas.auto_dream import AutoDreamRequest + from mirix.services.auto_dream_manager import AutoDreamManager + + await AutoDreamManager().run( + request=AutoDreamRequest( + mode="procedural", last_n_sessions=threshold + ), + user=user, + actor=actor, + meta_agent_state=meta_agent_state, + ) + except Exception as e: # noqa: BLE001 โ€” background task, never fatal + # The trigger window was already CLAIMED (check_and_claim_fire), so a + # failure here means this window's sessions won't be auto-distilled. + # Not data loss: the raw messages are retained (Goal 1), so a manual + # POST /memory/auto_dream (mode=procedural) re-distills the last-N + # retained sessions, and the next window still fires normally. Logged + # LOUD (ERROR + traceback) so a lost window is observable/greppable. + logger.error( + "Experience auto-dream run FAILED for meta_agent=%s (window already " + "claimed; recover via POST /memory/auto_dream mode=procedural): %s", + meta_agent_state.id, + e, + exc_info=True, + ) + + import asyncio + + asyncio.create_task(_run_dream()) + logger.info( + "Scheduled experience auto-dream (mode=procedural, last_n=%d, meta_agent=%s)", + threshold, + meta_agent_state.id, + ) + return True + except Exception as e: # noqa: BLE001 โ€” scheduling must never break the turn + logger.warning("Failed to schedule experience auto-dream: %s", e) + return False + + async def trigger_memory_update_with_instruction( self: "Agent", user_message: object, instruction: str, memory_type: str ) -> Optional[str]: @@ -908,7 +1552,9 @@ async def trigger_memory_update_with_instruction( from mirix.local_client import create_client if not isinstance(user_message, dict): - raise TypeError(f"user_message must be a dictionary, got {type(user_message).__name__}: {user_message}") + raise TypeError( + f"user_message must be a dictionary, got {type(user_message).__name__}: {user_message}" + ) if memory_type == "core": agent_type = "core_memory_agent" @@ -917,6 +1563,13 @@ async def trigger_memory_update_with_instruction( elif memory_type == "resource": agent_type = "resource_memory_agent" elif memory_type == "procedural": + # NOTE: this is the EXPLICIT, user-driven memory-correction path (the chat + # agent invokes this tool only when the user explicitly asks to insert / + # update / delete a specific memory). It is intentionally NOT subject to the + # "procedural has exactly one producer" gate applied to trigger_memory_update, + # which governs only AUTOMATIC extraction from a conversation transcript. + # Skill *learning* still flows solely through the distillation path; a direct + # user command to edit procedural memory is a different operation and is kept. agent_type = "procedural_memory_agent" elif memory_type == "knowledge_vault": agent_type = "knowledge_vault_memory_agent" @@ -945,18 +1598,26 @@ async def trigger_memory_update_with_instruction( existing_file_uris=user_message["existing_file_uris"], retrieved_memories=user_message.get("retrieved_memories", None), block_filter_tags=getattr(self, "block_filter_tags", None), - block_filter_tags_update_mode=getattr(self, "block_filter_tags_update_mode", "merge"), + block_filter_tags_update_mode=getattr( + self, "block_filter_tags_update_mode", "merge" + ), + ) + result = ( + "[System Message] Agent " + + matching_agent.name + + " has been triggered to update the memory.\n" ) - result = "[System Message] Agent " + matching_agent.name + " has been triggered to update the memory.\n" return result.strip() -async def trigger_memory_update(self: "Agent", user_message: object, memory_types: List[str]) -> Optional[str]: +async def trigger_memory_update( + self: "Agent", user_message: object, memory_types: List[str] +) -> Optional[str]: """ Choose which memory to update. This function will trigger another memory agent which is specifically in charge of handling the corresponding memory to update its memory. Trigger all necessary memory updates at once. Put the explanations in the `internal_monologue` field. Args: - memory_types (List[str]): The types of memory to update. It should be chosen from the following: "core", "episodic", "resource", "procedural", "knowledge_vault", "semantic". For instance, ['episodic', 'resource']. + memory_types (List[str]): The types of memory to update. It should be chosen from the following: "core", "episodic", "resource", "knowledge_vault", "semantic". For instance, ['episodic', 'resource']. ("procedural" is still accepted for backward compatibility but is a no-op here: procedural/skill memory is produced solely by the conversation-store distillation path, not by this dispatch.) Returns: Optional[str]: None is always returned as this function does not produce a response. @@ -973,7 +1634,101 @@ async def trigger_memory_update(self: "Agent", user_message: object, memory_type # Validate that user_message is a dictionary if not isinstance(user_message, dict): - raise TypeError(f"user_message must be a dictionary, got {type(user_message).__name__}: {user_message}") + raise TypeError( + f"user_message must be a dictionary, got {type(user_message).__name__}: {user_message}" + ) + + # Fixed-interval procedural trigger: auto-include "procedural" once + # SKILL_TRIGGER_SESSION_THRESHOLD distinct message sessions have accrued + # on this (agent, user) pair since the last fire. The cursor lives in the + # agent_trigger_state table; the running count is derived from messages + # at query time so there is only one source of truth and restarts/multi- + # worker deployments all agree. + # + # `check_and_claim_fire` is the single atomic entry point: it serializes + # the threshold-check + cursor-advance under SELECT FOR UPDATE so two + # concurrent workers can never double-fire for the same window, and it + # advances the cursor to the observed MAX(messages.created_at) so the + # next window cannot skip messages inserted mid-check. + from mirix.constants import DEFAULT_ORG_ID, SKILL_TRIGGER_SESSION_THRESHOLD + from mirix.schemas.agent_trigger_state import TRIGGER_TYPE_PROCEDURAL_SKILL + from mirix.services.agent_trigger_state_manager import AgentTriggerStateManager + + # Messages are stored against the top-level (chat) agent, not the meta + # agent, so count sessions on the parent when this tool runs inside a + # child memory agent โ€” matches the convention used elsewhere in this module. + trigger_agent_id = ( + self.agent_state.parent_id + if self.agent_state.parent_id is not None + else self.agent_state.id + ) + trigger_user_id = self.user.id if self.user else None + trigger_org_id = ( + (self.actor.organization_id or DEFAULT_ORG_ID) + if getattr(self, "actor", None) + else None + ) + current_session_id = getattr(self, "_current_step_session_id", None) + + if trigger_user_id is None: + # No user scope โ€” we cannot bookkeep a per-user cursor. Skip the + # auto-trigger rather than silently lumping all users together. + logger.debug("Skipping session-based procedural trigger: no user on agent.") + else: + trigger_mgr = AgentTriggerStateManager() + claim = await trigger_mgr.check_and_claim_fire( + agent_id=trigger_agent_id, + user_id=trigger_user_id, + trigger_type=TRIGGER_TYPE_PROCEDURAL_SKILL, + threshold=SKILL_TRIGGER_SESSION_THRESHOLD, + organization_id=trigger_org_id, + current_session_id=current_session_id, + ) + if claim.fired: + # On the every-N-sessions fire, schedule the GENERAL + # session-experience auto-dream (distill the last-N sealed + # sessions โ†’ pending experiences โ†’ skill self-evolution via the + # procedural agent). Fire-and-forget so a failure here can never + # break the chat turn; the cadence reuses the claim we just won. + # + # Procedural memory now has EXACTLY ONE producer: this distillation + # path. The legacy "procedural" memory update is no longer injected + # into the meta dispatch (it is filtered out below regardless), so we + # do not (and must not) fall back to enqueuing it here โ€” that would + # both reintroduce the second producer and let its ProceduralMemoryAgent + # .step race the curator's reset of the same agent's in-context history. + scheduled = await _schedule_experience_auto_dream( + self, + threshold=SKILL_TRIGGER_SESSION_THRESHOLD, + ) + logger.info( + "Auto-triggering procedural skill evolution " + "(sessions_since=%d, threshold=%d, agent=%s, user=%s, " + "experience_dream_scheduled=%s)", + claim.sessions_since, + SKILL_TRIGGER_SESSION_THRESHOLD, + trigger_agent_id, + trigger_user_id, + scheduled, + ) + + # Procedural memory has EXACTLY ONE producer: the Conversation Message Store + # distillation path (scheduled above as the experience auto-dream). The meta + # agent must never dispatch the procedural_memory_agent during normal + # extraction, so "procedural" is filtered out of memory_types here. It stays + # an accepted input value (a harmless no-op) so existing callers / LLM emissions + # don't error; it simply no longer dispatches. The procedural_memory_agent + # registration and infrastructure are retained for the distillation path. + filtered_memory_types = [mt for mt in memory_types if mt != "procedural"] + if len(filtered_memory_types) != len(memory_types): + logger.debug( + "Filtered 'procedural' out of trigger_memory_update dispatch " + "(procedural memory is produced solely by the distillation path): " + "%s -> %s", + list(memory_types), + filtered_memory_types, + ) + memory_types = filtered_memory_types # De-duplicate memory types while preserving order. # The MetaMemoryAgent (LLM) can occasionally emit duplicates (e.g., ["semantic", "semantic"]). @@ -1005,14 +1760,17 @@ async def trigger_memory_update(self: "Agent", user_message: object, memory_type ) # Get child agents - child_agent_states = await self.agent_manager.list_agents(parent_id=self.agent_state.id, actor=self.actor) + child_agent_states = await self.agent_manager.list_agents( + parent_id=self.agent_state.id, actor=self.actor + ) # Map agent types to agent states (key by string so lookup works for enum or deserialized string) def _agent_type_key(at): return at.value if hasattr(at, "value") else str(at) agent_type_to_state = { - _agent_type_key(agent_state.agent_type): agent_state for agent_state in child_agent_states + _agent_type_key(agent_state.agent_type): agent_state + for agent_state in child_agent_states } if not child_agent_states: @@ -1047,14 +1805,24 @@ async def _run_single_memory_update(memory_type: str) -> str: # Deep copy filter_tags and block_filter_tags to ensure complete isolation between child agents parent_filter_tags = getattr(self, "filter_tags", None) # Don't use 'or {}' because empty dict {} is valid and different from None - filter_tags = deepcopy(parent_filter_tags) if parent_filter_tags is not None else None + filter_tags = ( + deepcopy(parent_filter_tags) if parent_filter_tags is not None else None + ) parent_block_filter_tags = getattr(self, "block_filter_tags", None) - block_filter_tags = deepcopy(parent_block_filter_tags) if parent_block_filter_tags is not None else None - block_filter_tags_update_mode = getattr(self, "block_filter_tags_update_mode", "merge") + block_filter_tags = ( + deepcopy(parent_block_filter_tags) + if parent_block_filter_tags is not None + else None + ) + block_filter_tags_update_mode = getattr( + self, "block_filter_tags_update_mode", "merge" + ) use_cache = getattr(self, "use_cache", True) actor = getattr(self, "actor", None) user = getattr(self, "user", None) - occurred_at = getattr(self, "occurred_at", None) # Get occurred_at from parent agent + occurred_at = getattr( + self, "occurred_at", None + ) # Get occurred_at from parent agent import logging @@ -1103,6 +1871,11 @@ async def _run_single_memory_update(memory_type: str) -> str: else: message_copy.content = [system_msg] + chaining = user_message.get("chaining", False) + # Procedural agent needs chaining for multi-step CLI workflow + if memory_type == "procedural": + chaining = True + # Extract topics and retrieved_memories from parent agent to pass to sub-agents # This ensures sub-agents use the same keywords for memory retrieval retrieved_memories = user_message.get("retrieved_memories", None) @@ -1114,8 +1887,14 @@ async def _run_single_memory_update(memory_type: str) -> str: # actor is needed for write operations, user is needed for read operations # Wrap child agent execution in a LangFuse span for hierarchical tracing - trace_id = parent_trace_context.get("trace_id") if parent_trace_context else None - parent_span_id = parent_trace_context.get("observation_id") if parent_trace_context else None + trace_id = ( + parent_trace_context.get("trace_id") if parent_trace_context else None + ) + parent_span_id = ( + parent_trace_context.get("observation_id") + if parent_trace_context + else None + ) if langfuse and trace_id: from typing import cast @@ -1157,7 +1936,7 @@ async def _run_single_memory_update(memory_type: str) -> str: await memory_agent.step( input_messages=message_copy, - chaining=user_message.get("chaining", False), + chaining=chaining, actor=actor, user=user, topics=topics, @@ -1167,7 +1946,7 @@ async def _run_single_memory_update(memory_type: str) -> str: # No tracing available, run directly await memory_agent.step( input_messages=message_copy, - chaining=user_message.get("chaining", False), + chaining=chaining, actor=actor, user=user, topics=topics, @@ -1187,10 +1966,14 @@ async def _run_single_memory_update(memory_type: str) -> str: results = await asyncio.gather(*tasks, return_exceptions=True) for i, result in enumerate(results): if isinstance(result, Exception): - raise RuntimeError(f"Failed to trigger memory update for '{memory_types[i]}'") from result + raise RuntimeError( + f"Failed to trigger memory update for '{memory_types[i]}'" + ) from result responses[i] = result - ordered_responses = [responses[i] for i in range(len(memory_types)) if i in responses] + ordered_responses = [ + responses[i] for i in range(len(memory_types)) if i in responses + ] return "".join(ordered_responses).strip() diff --git a/mirix/functions/functions.py b/mirix/functions/functions.py index d0cb42d31..6f2fc5a54 100755 --- a/mirix/functions/functions.py +++ b/mirix/functions/functions.py @@ -168,8 +168,19 @@ def load_function_set(module: ModuleType) -> dict: # Get the attribute attr = getattr(module, attr_name) - # Check if it's a callable function and not a built-in or special method - if inspect.isfunction(attr) and attr.__module__ == module.__name__: + # Check if it's a callable function and not a built-in or special method. + # Only public coroutine functions are tools: a tool is always `async def` + # and never underscore-prefixed. Module-level sync/private helpers (e.g. + # `compute_edit_budget`, `_edit_exceeds_size_gate` in memory_tools) are + # NOT tools and frequently lack the per-parameter docstrings that + # generate_schema() requires โ€” skipping them keeps one such helper from + # raising and dropping every real tool in the module. + if ( + inspect.isfunction(attr) + and attr.__module__ == module.__name__ + and inspect.iscoroutinefunction(attr) + and not attr_name.startswith("_") + ): if attr_name in function_dict: raise ValueError(f"Found a duplicate of function name '{attr_name}'") diff --git a/mirix/helpers/json_parsing.py b/mirix/helpers/json_parsing.py new file mode 100644 index 000000000..17f2b185d --- /dev/null +++ b/mirix/helpers/json_parsing.py @@ -0,0 +1,165 @@ +"""Generic, agent-agnostic JSON-extraction helpers for LLM replies. + +These are pure, dependency-free utilities for robustly recovering a JSON object +or array from a model reply that may wrap it in a ```json fenced block, plain +fences, or surrounding prose. +""" + +from __future__ import annotations + +import json +import re +from typing import Dict, List, Optional + + +def _try_load_object(text: str) -> Optional[Dict]: + try: + obj = json.loads(text) + except (json.JSONDecodeError, ValueError): + return None + return obj if isinstance(obj, dict) else None + + +def _try_load_array(text: str) -> Optional[List[Dict]]: + """Parse ``text`` as a JSON array, returning only its dict elements. + + Returns ``None`` when ``text`` is not a JSON array (so callers can fall + through to the next strategy), but ``[]`` for a genuinely empty array. + """ + try: + obj = json.loads(text) + except (json.JSONDecodeError, ValueError): + return None + if not isinstance(obj, list): + return None + return [el for el in obj if isinstance(el, dict)] + + +def _balanced_objects(text: str): + """Yield each top-level balanced ``{...}`` substring (respecting strings).""" + i = 0 + n = len(text) + while i < n: + if text[i] != "{": + i += 1 + continue + depth = 0 + in_str = False + escaped = False + for j in range(i, n): + ch = text[j] + if in_str: + if escaped: + escaped = False + elif ch == "\\": + escaped = True + elif ch == '"': + in_str = False + continue + if ch == '"': + in_str = True + elif ch == "{": + depth += 1 + elif ch == "}": + depth -= 1 + if depth == 0: + yield text[i : j + 1] + break + i += 1 + + +def _balanced_arrays(text: str): + """Yield each top-level balanced ``[...]`` substring (respecting strings).""" + i = 0 + n = len(text) + while i < n: + if text[i] != "[": + i += 1 + continue + depth = 0 + in_str = False + escaped = False + for j in range(i, n): + ch = text[j] + if in_str: + if escaped: + escaped = False + elif ch == "\\": + escaped = True + elif ch == '"': + in_str = False + continue + if ch == '"': + in_str = True + elif ch == "[": + depth += 1 + elif ch == "]": + depth -= 1 + if depth == 0: + yield text[i : j + 1] + break + i += 1 + + +def parse_distiller_json(raw: Optional[str]) -> Optional[Dict]: + """Robustly parse an LLM reply into a single JSON object (dict), or None. + + Tolerates: a bare JSON object, a ```json fenced block, a plain ``` fenced + block, or a single object embedded in surrounding prose. + """ + if not raw or not isinstance(raw, str): + return None + text = raw.strip() + + fence = re.search(r"```(?:json)?\s*(.*?)```", text, re.DOTALL | re.IGNORECASE) + if fence: + parsed = _try_load_object(fence.group(1).strip()) + if parsed is not None: + return parsed + + parsed = _try_load_object(text) + if parsed is not None: + return parsed + + for span in _balanced_objects(text): + parsed = _try_load_object(span) + if parsed is not None: + return parsed + + return None + + +def parse_distiller_json_array(raw: Optional[str]) -> List[Dict]: + """Robustly parse an LLM reply expected to be a JSON ARRAY of objects. + + Tolerates a bare array, a ```json fenced block, a plain fenced block, or an + array embedded in prose. Returns the list of dict elements (non-dict elements + dropped). Returns ``[]`` when nothing parseable is found OR when the model + legitimately emitted an empty array (the "nothing worth remembering" case). + Falls back to wrapping a single object so a model that emitted one item + without the enclosing array is not silently dropped. + """ + if not raw or not isinstance(raw, str): + return [] + text = raw.strip() + + fence = re.search(r"```(?:json)?\s*(.*?)```", text, re.DOTALL | re.IGNORECASE) + if fence: + items = _try_load_array(fence.group(1).strip()) + if items is not None: + return items + + items = _try_load_array(text) + if items is not None: + return items + + for span in _balanced_arrays(text): + items = _try_load_array(span) + if items is not None: + return items + + obj = parse_distiller_json(text) + if obj is not None: + return [obj] + + return [] diff --git a/mirix/helpers/keyed_locks.py b/mirix/helpers/keyed_locks.py new file mode 100644 index 000000000..7c43ba818 --- /dev/null +++ b/mirix/helpers/keyed_locks.py @@ -0,0 +1,49 @@ +"""Per-key asyncio.Lock registry with automatic eviction. + +The naive pattern โ€” a module-level ``dict[key, asyncio.Lock]`` filled lazily โ€” +leaks: entries are never removed because a bare Lock cannot know when its last +user is done, so a long-lived server accumulates one Lock per key it has ever +seen. ``KeyedLocks`` reference-counts holders and waiters and drops an entry +the moment nobody holds or awaits it. + +Used to serialize per-entity work across coroutines in ONE process (per-raw- +memory updates, per-agent skill evolution, per-user procedural dreams). It is +NOT a cross-process lock โ€” pair it with a Postgres advisory lock when multiple +server processes contend (see auto_dream_manager._procedural_dream_guard). +""" + +import asyncio +from contextlib import asynccontextmanager +from typing import AsyncIterator, Dict, Tuple + + +class KeyedLocks: + """Lazily-created, self-evicting registry of per-key ``asyncio.Lock``s. + + All bookkeeping runs synchronously between awaits on the single event + loop, so the get-or-create and the refcount updates are atomic without any + extra locking of their own. + """ + + def __init__(self) -> None: + # key -> (lock, number of coroutines holding or waiting on it) + self._entries: Dict[str, Tuple[asyncio.Lock, int]] = {} + + @asynccontextmanager + async def acquire(self, key: str) -> AsyncIterator[None]: + """``async with locks.acquire(key):`` โ€” mutual exclusion per key.""" + lock, refs = self._entries.get(key) or (asyncio.Lock(), 0) + self._entries[key] = (lock, refs + 1) + try: + async with lock: + yield + finally: + lock, refs = self._entries[key] + if refs <= 1: + del self._entries[key] + else: + self._entries[key] = (lock, refs - 1) + + def __len__(self) -> int: + """Number of live entries โ€” 0 when nothing is held or awaited.""" + return len(self._entries) diff --git a/mirix/helpers/message_helpers.py b/mirix/helpers/message_helpers.py index 631c81622..46e03be94 100644 --- a/mirix/helpers/message_helpers.py +++ b/mirix/helpers/message_helpers.py @@ -62,4 +62,5 @@ def prepare_input_message_create(message: MessageCreate, agent_id: str, **kwargs otid=message.otid, sender_id=message.sender_id, group_id=message.group_id, + session_id=message.session_id, ) diff --git a/mirix/llm_api/llm_client.py b/mirix/llm_api/llm_client.py index 8b51bd107..59917fa49 100644 --- a/mirix/llm_api/llm_client.py +++ b/mirix/llm_api/llm_client.py @@ -48,5 +48,11 @@ def create( return GoogleAIClient( llm_config=llm_config, ) + case "openrouter": + from mirix.llm_api.openrouter_client import OpenRouterClient + + return OpenRouterClient( + llm_config=llm_config, + ) case _: return None diff --git a/mirix/local_client/local_client.py b/mirix/local_client/local_client.py index 995ab43ed..f187728df 100644 --- a/mirix/local_client/local_client.py +++ b/mirix/local_client/local_client.py @@ -1841,7 +1841,7 @@ async def retrieve_memory( query: str, memory_type: str = "all", search_field: str = "null", - search_method: str = "embedding", + search_method: Optional[str] = None, timezone_str: str = "UTC", limit: int = 10, ) -> dict: @@ -1854,11 +1854,14 @@ async def retrieve_memory( memory_type (str): The type of memory to search in. Options: "episodic", "resource", "procedural", "knowledge_vault", "semantic", "all". Defaults to "all". search_field (str): The field to search in the memory. For "episodic": 'summary', 'details'; - for "resource": 'summary', 'content'; for "procedural": 'summary', 'steps'; + for "resource": 'summary', 'content'; for "procedural": 'description', 'instructions'; for "knowledge_vault": 'secret_value', 'caption'; for "semantic": 'name', 'summary', 'details'. Use "null" for default fields. Defaults to "null". - search_method (str): The method to search. Options: 'bm25' (keyword-based), 'embedding' (semantic). - Defaults to "embedding". + search_method (Optional[str]): The method to search. Options: 'bm25' (keyword-based), + 'embedding' (semantic), and 'hybrid' (procedural only; BM25 + + embedding fused with Reciprocal Rank Fusion). Omit (None) to use + the per-type default: 'hybrid' for procedural, 'embedding' otherwise + โ€” the same resolution the REST /memory/search endpoint applies. timezone_str (str): Timezone string for time-based operations. Defaults to "UTC". limit (int): Maximum number of results to return per memory type. Defaults to 10. @@ -1867,6 +1870,21 @@ async def retrieve_memory( """ # Import here to avoid circular imports + # Resolve the per-type default, mirroring the REST endpoint: only a + # single-type procedural search defaults to hybrid; "all" keeps the + # cross-type embedding default. + if not search_method: + if memory_type == "procedural": + from mirix.constants import PROCEDURAL_DEFAULT_SEARCH_METHOD + + search_method = PROCEDURAL_DEFAULT_SEARCH_METHOD + else: + search_method = "embedding" + elif search_method == "hybrid" and memory_type != "procedural": + # Hybrid is a procedural-only lane; other managers don't implement + # it. Same downgrade the REST endpoint applies. + search_method = "embedding" + # Validate inputs if memory_type == "resource" and search_field == "content" and search_method == "embedding": raise ValueError("embedding is not supported for resource memory's 'content' field.") @@ -1973,7 +1991,7 @@ async def retrieve_memory( agent_state=agent_state, query=query, embedded_text=embedded_text if search_method == "embedding" and query else None, - search_field=search_field if search_field != "null" else "summary", + search_field=search_field if search_field != "null" else "description", search_method=search_method, limit=limit, timezone_str=timezone_str, @@ -1983,8 +2001,12 @@ async def retrieve_memory( "memory_type": "procedural", "id": x.id, "entry_type": x.entry_type, - "summary": x.summary, - "steps": x.steps, + "name": x.name, + "description": x.description, + "instructions": x.instructions, + "triggers": getattr(x, "triggers", None) or [], + "examples": getattr(x, "examples", None) or [], + "version": getattr(x, "version", None), } for x in procedural_memories ] diff --git a/mirix/orm/__init__.py b/mirix/orm/__init__.py index 620073898..d96f708ce 100755 --- a/mirix/orm/__init__.py +++ b/mirix/orm/__init__.py @@ -1,9 +1,11 @@ from mirix.orm.agent import Agent +from mirix.orm.agent_trigger_state import AgentTriggerState from mirix.orm.base import Base from mirix.orm.block import Block from mirix.orm.client import Client from mirix.orm.client_api_key import ClientApiKey from mirix.orm.cloud_file_mapping import CloudFileMapping +from mirix.orm.conversation_message import ConversationMessage from mirix.orm.episodic_memory import EpisodicEvent from mirix.orm.file import FileMetadata from mirix.orm.knowledge_vault import KnowledgeVaultItem @@ -13,6 +15,7 @@ from mirix.orm.provider import Provider from mirix.orm.resource_memory import ResourceMemoryItem from mirix.orm.semantic_memory import SemanticMemoryItem +from mirix.orm.skill_experience import SkillExperience from mirix.orm.step import Step from mirix.orm.tool import Tool from mirix.orm.tools_agents import ToolsAgents @@ -20,11 +23,13 @@ __all__ = [ "Agent", + "AgentTriggerState", "Base", "Block", "Client", "ClientApiKey", "CloudFileMapping", + "ConversationMessage", "EpisodicEvent", "FileMetadata", "KnowledgeVaultItem", @@ -34,6 +39,7 @@ "Provider", "ResourceMemoryItem", "SemanticMemoryItem", + "SkillExperience", "Step", "Tool", "ToolsAgents", diff --git a/mirix/orm/agent_trigger_state.py b/mirix/orm/agent_trigger_state.py new file mode 100644 index 000000000..db7e8917c --- /dev/null +++ b/mirix/orm/agent_trigger_state.py @@ -0,0 +1,98 @@ +import uuid +from datetime import datetime +from typing import TYPE_CHECKING, List, Optional + +from sqlalchemy import JSON, DateTime, Index, String, UniqueConstraint +from sqlalchemy.orm import Mapped, declared_attr, mapped_column, relationship + +from mirix.orm.mixins import AgentMixin, OrganizationMixin, UserMixin +from mirix.orm.sqlalchemy_base import SqlalchemyBase +from mirix.schemas.agent_trigger_state import ( + AgentTriggerState as PydanticAgentTriggerState, +) + +if TYPE_CHECKING: + from mirix.orm.agent import Agent + from mirix.orm.organization import Organization + from mirix.orm.user import User + + +class AgentTriggerState(SqlalchemyBase, OrganizationMixin, UserMixin, AgentMixin): + """ + Persists per-(agent, user, trigger_type) bookkeeping for interval-driven + memory triggers (e.g. "fire a procedural-memory extraction every N sessions"). + + Only the last-fire cursor is stored here. The live counter (e.g. "how many + sessions since the last fire") is derived at read time from the messages + table via a COUNT(DISTINCT session_id) โ€” the single source of truth for + session accrual. + """ + + __tablename__ = "agent_trigger_state" + __pydantic_model__ = PydanticAgentTriggerState + + id: Mapped[str] = mapped_column( + String, primary_key=True, default=lambda: f"ats-{uuid.uuid4()}" + ) + + trigger_type: Mapped[str] = mapped_column( + String(64), + nullable=False, + doc="Kind of trigger this row bookkeeps (e.g. 'procedural_skill').", + ) + + last_fired_at: Mapped[Optional[datetime]] = mapped_column( + DateTime(timezone=True), + nullable=True, + doc="Created_at of the message that caused the last fire. Used as the " + "lower bound when counting new sessions.", + ) + + last_fired_session_id: Mapped[Optional[str]] = mapped_column( + String(64), + nullable=True, + doc="session_id that was in-progress on the agent when the last fire " + "claimed. Stored for audit/logging only; the fire filter is MIN-based " + "and subsumes in-progress-session dedup automatically.", + ) + + last_fired_tied_session_ids: Mapped[Optional[List[str]]] = mapped_column( + JSON, + nullable=True, + default=list, + doc="session_ids whose MIN(created_at) (first-appearance timestamp) " + "equalled last_fired_at in the last fire's counted window. The fire " + "filter uses MIN semantics so each session_id can only contribute to " + "exactly one window; this tied set handles the exact-microsecond case " + "where a session's first message commits at the same timestamp as our " + "watermark but was not visible to our SELECT. The next window lets " + "such sessions qualify via `first_ts = last_fired_at AND session_id " + "NOT IN last_fired_tied_session_ids`.", + ) + + __table_args__ = ( + UniqueConstraint( + "agent_id", + "user_id", + "trigger_type", + name="uq_agent_trigger_state_agent_user_type", + ), + Index( + "ix_agent_trigger_state_agent_user_type", + "agent_id", + "user_id", + "trigger_type", + ), + ) + + @declared_attr + def agent(cls) -> Mapped[Optional["Agent"]]: + return relationship("Agent", lazy="selectin") + + @declared_attr + def user(cls) -> Mapped[Optional["User"]]: + return relationship("User", lazy="selectin") + + @declared_attr + def organization(cls) -> Mapped[Optional["Organization"]]: + return relationship("Organization", lazy="selectin") diff --git a/mirix/orm/conversation_message.py b/mirix/orm/conversation_message.py new file mode 100644 index 000000000..1c311cc25 --- /dev/null +++ b/mirix/orm/conversation_message.py @@ -0,0 +1,146 @@ +"""ORM model for the Conversation Message Store. + +A `ConversationMessage` is one external conversation turn โ€” a real `user` or +`assistant` message that arrived through the memory-add API *carrying a +session_id*. It is the canonical, learnable record of a conversation and the +SINGLE source the procedural-memory (skill) distiller reads. + +This store is deliberately SEPARATE from the `messages` table that backs the +agent loop. The `messages` table also holds the meta agent's own synthesized +bookkeeping (`[System Message]` bootstraps, `trigger_memory_update` tool calls, +`continue_chaining` heartbeats); reading skills from there made the distiller +learn from MIRIX operating *itself*. By giving the real conversation its own +home โ€” with its REAL roles preserved, not the `[USER]`/`[ASSISTANT]` +role-collapsed form the meta agent receives โ€” correctness comes from structure +(which table we read) rather than from fragile string conventions. + +Only turns that arrived with a `session_id` land here (`session_id` is NOT +NULL): this store holds session'd conversation turns and nothing else. A +session's order is defined by `MIN(created_at)` over its turns; `distilled_at` +marks a session that a rolling distill round has already consumed so the barrier +advances and history is never reprocessed. + +Inherited (do NOT redeclare): + - id, created_at, updated_at, is_deleted -- SqlalchemyBase + meta mixin + - organization_id -- OrganizationMixin + - user_id -- UserMixin + +`role` is a plain String (NOT a pg ENUM); its value space ('user' | 'assistant') +is Literal-validated in `schemas/conversation_message.py`, so the value space can +change without a DB migration. The `id` PK carries a `convmsg-` prefix matching +the `sexp-`/`proc-` convention on sibling tables. +""" + +import uuid +from datetime import datetime +from typing import TYPE_CHECKING, Optional + +from sqlalchemy import CheckConstraint, DateTime, Index, String, Text +from sqlalchemy.orm import Mapped, declared_attr, mapped_column, relationship + +from mirix.orm.mixins import OrganizationMixin, UserMixin +from mirix.orm.sqlalchemy_base import SqlalchemyBase +from mirix.schemas.conversation_message import ( + ConversationMessage as PydanticConversationMessage, +) +from mirix.schemas.message import SESSION_ID_MAX_LEN, SESSION_ID_SQL_PATTERN +from mirix.settings import settings + +if TYPE_CHECKING: + from mirix.orm.organization import Organization + from mirix.orm.user import User + + +class ConversationMessage(SqlalchemyBase, OrganizationMixin, UserMixin): + """One external conversation turn that arrived with a session_id.""" + + __tablename__ = "conversation_message" + __pydantic_model__ = PydanticConversationMessage + + # Override the bare SqlalchemyBase PK to attach a prefixed default, matching + # the `sexp-`/`proc-` convention on sibling tables. + id: Mapped[str] = mapped_column( + String, + primary_key=True, + default=lambda: f"convmsg-{uuid.uuid4()}", + doc="Unique ID for this conversation message.", + ) + + session_id: Mapped[str] = mapped_column( + String(SESSION_ID_MAX_LEN), + nullable=False, + index=True, + doc="The conversation this turn belongs to. NOT NULL โ€” this store only " + "holds session'd turns. MIN(created_at) per session defines order.", + ) + role: Mapped[str] = mapped_column( + String, + nullable=False, + doc="The real turn role: 'user', 'assistant', or 'tool' (validated in " + "Pydantic, not a pg ENUM). NOT the role-collapsed form the meta agent " + "receives.", + ) + content: Mapped[str] = mapped_column( + Text, + nullable=False, + default="", + doc="The verbatim text of this conversation turn.", + ) + distilled_at: Mapped[Optional["datetime"]] = mapped_column( + DateTime(timezone=True), + nullable=True, + default=None, + doc="Set when this session's turns were consumed by a distill round. " + "NULL = not yet distilled; the rolling barrier reads only NULL sessions.", + ) + + # Indexes mirror the skill_experience pg/sqlite guard pattern: an explicit + # `is not None` filter (not `filter(None, ...)`) because an un-attached Index + # is falsy until bound. The composite index backs the primary access pattern + # โ€” "list/seal/order this (org, user)'s sessions by first-appearance time". + __table_args__ = tuple( + item + for item in [ + # Backstop the app-level validator so a create_all-built database + # matches the SQL migration: the DB must never store an invalid + # session_id. Unlike `messages`, session_id here is NOT NULL (this + # store only holds session'd turns), so the NULL branch is omitted. + # Uses the Postgres `~` operator, so emit only on Postgres (SQLite, + # used for some local/test setups, has no POSIX regex operator). + # Pattern derived from mirix.schemas.message โ€” one source of truth. + CheckConstraint( + f"session_id ~ '{SESSION_ID_SQL_PATTERN}'", + name="ck_conversation_message_session_id_format", + ).ddl_if(dialect="postgresql"), + Index( + "ix_conversation_message_org_user_session_created", + "organization_id", + "user_id", + "session_id", + "created_at", + ), + # Accelerates the per-session ascending fetch in list_turns_for_session. + Index( + "ix_conversation_message_session_created", + "session_id", + "created_at", + ), + ( + Index( + "ix_conversation_message_organization_id", + "organization_id", + ) + if settings.mirix_pg_uri_no_default + else None + ), + ] + if item is not None + ) + + @declared_attr + def user(cls) -> Mapped[Optional["User"]]: + return relationship("User", lazy="selectin") + + @declared_attr + def organization(cls) -> Mapped[Optional["Organization"]]: + return relationship("Organization", lazy="selectin") diff --git a/mirix/orm/message.py b/mirix/orm/message.py index bab3fa91b..ef9bd944b 100755 --- a/mirix/orm/message.py +++ b/mirix/orm/message.py @@ -1,6 +1,6 @@ from typing import TYPE_CHECKING, List, Optional -from sqlalchemy import JSON, ForeignKey, Index +from sqlalchemy import JSON, CheckConstraint, ForeignKey, Index, String from sqlalchemy.orm import Mapped, declared_attr, mapped_column, relationship from mirix.orm.custom_columns import ( @@ -10,8 +10,12 @@ ) from mirix.orm.mixins import AgentMixin, OrganizationMixin, UserMixin from mirix.orm.sqlalchemy_base import SqlalchemyBase -from mirix.schemas.message import Message as PydanticMessage -from mirix.schemas.message import ToolReturn +from mirix.schemas.message import ( + SESSION_ID_MAX_LEN, + SESSION_ID_SQL_PATTERN, + Message as PydanticMessage, + ToolReturn, +) from mirix.schemas.mirix_message_content import MessageContent from mirix.schemas.mirix_message_content import TextContent as PydanticTextContent from mirix.schemas.openai.openai import ToolCall as OpenAIToolCall @@ -32,6 +36,21 @@ class Message(SqlalchemyBase, OrganizationMixin, UserMixin, AgentMixin): Index("ix_messages_created_at", "created_at", "id"), Index("ix_messages_client_user", "client_id", "user_id"), Index("ix_messages_agent_client_user", "agent_id", "client_id", "user_id"), + # Accelerates "list messages of an agent in a session, newest first". + Index( + "ix_messages_agent_session_created_at", + "agent_id", + "session_id", + "created_at", + ), + # Backstop the app-level validator: DB must never store an invalid session_id. + # Uses the Postgres `~` operator, so emit the constraint only on Postgres + # (SQLite, used for some local/test setups, has no POSIX regex operator). + # Pattern derived from mirix.schemas.message so there's one source of truth. + CheckConstraint( + f"session_id IS NULL OR session_id ~ '{SESSION_ID_SQL_PATTERN}'", + name="ck_messages_session_id_format", + ).ddl_if(dialect="postgresql"), ) __pydantic_model__ = PydanticMessage @@ -78,6 +97,13 @@ class Message(SqlalchemyBase, OrganizationMixin, UserMixin, AgentMixin): nullable=True, doc="The id of the sender of the message, can be an identity id or agent id", ) + session_id: Mapped[Optional[str]] = mapped_column( + String(SESSION_ID_MAX_LEN), + nullable=True, + doc="Top-level conversation/session identifier for grouping messages. " + "Enforced by app validator and DB CHECK constraint " + f"(pattern {SESSION_ID_SQL_PATTERN}).", + ) # Relationships agent: Mapped["Agent"] = relationship("Agent", back_populates="messages", lazy="selectin") diff --git a/mirix/orm/procedural_memory.py b/mirix/orm/procedural_memory.py index efbbe7a01..56a1f9daa 100755 --- a/mirix/orm/procedural_memory.py +++ b/mirix/orm/procedural_memory.py @@ -1,165 +1,192 @@ -import datetime as dt -from datetime import datetime -from typing import TYPE_CHECKING, Optional - -from sqlalchemy import JSON, Column, ForeignKey, Index, String, text -from sqlalchemy.orm import Mapped, declared_attr, mapped_column, relationship - -from mirix.constants import MAX_EMBEDDING_DIM -from mirix.orm.custom_columns import CommonVector, EmbeddingConfigColumn -from mirix.orm.mixins import OrganizationMixin, UserMixin -from mirix.orm.sqlalchemy_base import SqlalchemyBase -from mirix.schemas.procedural_memory import ProceduralMemoryItem as PydanticProceduralMemoryItem -from mirix.settings import settings - -if TYPE_CHECKING: - from mirix.orm.agent import Agent - from mirix.orm.organization import Organization - from mirix.orm.user import User - - -class ProceduralMemoryItem(SqlalchemyBase, OrganizationMixin, UserMixin): - """ - Stores procedural memory entries, such as workflows, step-by-step guides, or how-to knowledge. - - type: The category or tag of the procedure (e.g. 'workflow', 'guide', 'script') - summary: Short descriptive text about what this procedure accomplishes - steps: Step-by-step instructions or method - """ - - __tablename__ = "procedural_memory" - __pydantic_model__ = PydanticProceduralMemoryItem - - # Primary key - id: Mapped[str] = mapped_column( - String, - primary_key=True, - doc="Unique ID for this procedural memory entry", - ) - - # Foreign key to agent - agent_id: Mapped[Optional[str]] = mapped_column( - String, - ForeignKey("agents.id", ondelete="CASCADE"), - nullable=True, - doc="ID of the agent this procedural memory item belongs to", - ) - - # Foreign key to client (for access control and filtering) - client_id: Mapped[Optional[str]] = mapped_column( - String, - ForeignKey("clients.id", ondelete="CASCADE"), - nullable=True, - doc="ID of the client application that created this item", - ) - - # Distinguish the type/category of the procedure - entry_type: Mapped[str] = mapped_column(String, doc="Category or type (e.g. 'workflow', 'guide', 'script')") - - # A human-friendly description of this procedure - summary: Mapped[str] = mapped_column(String, doc="Short description or title of the procedure") - - # Steps or instructions stored as a JSON object/list - steps: Mapped[list] = mapped_column(JSON, doc="Step-by-step instructions stored as a list of strings") - - # NEW: Filter tags for flexible filtering and categorization - filter_tags: Mapped[Optional[dict]] = mapped_column( - JSON, nullable=True, default=None, doc="Custom filter tags for filtering and categorization" - ) - - # When was this item last modified and what operation? - last_modify: Mapped[dict] = mapped_column( - JSON, - nullable=False, - default=lambda: { - "timestamp": datetime.now(dt.timezone.utc).isoformat(), - "operation": "created", - }, - doc="Last modification info including timestamp and operation type", - ) - - embedding_config: Mapped[Optional[dict]] = mapped_column( - EmbeddingConfigColumn, nullable=True, doc="Embedding configuration" - ) - - # Vector embedding field based on database type - if settings.mirix_pg_uri_no_default: - from pgvector.sqlalchemy import Vector - - summary_embedding = mapped_column(Vector(MAX_EMBEDDING_DIM), nullable=True) - steps_embedding = mapped_column(Vector(MAX_EMBEDDING_DIM), nullable=True) - else: - summary_embedding = Column(CommonVector, nullable=True) - steps_embedding = Column(CommonVector, nullable=True) - - # Database indexes for efficient querying - __table_args__ = tuple( - filter( - None, - [ - # Organization-level query optimization indexes - ( - Index("ix_procedural_memory_organization_id", "organization_id") - if settings.mirix_pg_uri_no_default - else None - ), - ( - Index( - "ix_procedural_memory_org_created_at", - "organization_id", - "created_at", - postgresql_using="btree", - ) - if settings.mirix_pg_uri_no_default - else None - ), - ( - Index( - "ix_procedural_memory_filter_tags_gin", - text("(filter_tags::jsonb)"), - postgresql_using="gin", - ) - if settings.mirix_pg_uri_no_default - else None - ), - ( - Index( - "ix_procedural_memory_org_filter_scope", - "organization_id", - text("((filter_tags->>'scope')::text)"), - postgresql_using="btree", - ) - if settings.mirix_pg_uri_no_default - else None - ), - # SQLite indexes - ( - Index("ix_procedural_memory_organization_id_sqlite", "organization_id") - if not settings.mirix_pg_uri_no_default - else None - ), - ], - ) - ) - - @declared_attr - def agent(cls) -> Mapped[Optional["Agent"]]: - """ - Relationship to the Agent that owns this procedural memory item. - """ - return relationship("Agent", lazy="selectin") - - @declared_attr - def organization(cls) -> Mapped["Organization"]: - """ - Relationship to organization (mirroring your existing patterns). - Adjust 'back_populates' to match the collection name in your `Organization` model. - """ - return relationship("Organization", back_populates="procedural_memory", lazy="selectin") - - @declared_attr - def user(cls) -> Mapped["User"]: - """ - Relationship to the User that owns this procedural memory item. - """ - return relationship("User", lazy="selectin") +import datetime as dt +from datetime import datetime +from typing import TYPE_CHECKING, Optional + +from sqlalchemy import JSON, Column, ForeignKey, Index, String, UniqueConstraint, text +from sqlalchemy.orm import Mapped, declared_attr, mapped_column, relationship + +from mirix.constants import MAX_EMBEDDING_DIM +from mirix.orm.custom_columns import CommonVector, EmbeddingConfigColumn +from mirix.orm.mixins import OrganizationMixin, UserMixin +from mirix.orm.sqlalchemy_base import SqlalchemyBase +from mirix.schemas.procedural_memory import ProceduralMemoryItem as PydanticProceduralMemoryItem +from mirix.settings import settings + +if TYPE_CHECKING: + from mirix.orm.agent import Agent + from mirix.orm.organization import Organization + from mirix.orm.user import User + + +class ProceduralMemoryItem(SqlalchemyBase, OrganizationMixin, UserMixin): + """ + Stores procedural knowledge as reusable skills. + + name: Short skill identifier (e.g. 'deploy-production') + entry_type: Category or tag of the skill (e.g. 'workflow', 'guide', 'script') + description: Short descriptive text about what this skill accomplishes + instructions: Step-by-step instructions as plain text + triggers: Conditions that indicate this skill is relevant + examples: Input/output examples for this skill + version: Semantic version of this skill + """ + + __tablename__ = "procedural_memory" + __pydantic_model__ = PydanticProceduralMemoryItem + + # Primary key + id: Mapped[str] = mapped_column( + String, + primary_key=True, + doc="Unique ID for this procedural memory entry", + ) + + # Foreign key to agent + agent_id: Mapped[Optional[str]] = mapped_column( + String, + ForeignKey("agents.id", ondelete="CASCADE"), + nullable=True, + doc="ID of the agent this procedural memory item belongs to", + ) + + # Foreign key to client (for access control and filtering) + client_id: Mapped[Optional[str]] = mapped_column( + String, + ForeignKey("clients.id", ondelete="CASCADE"), + nullable=True, + doc="ID of the client application that created this item", + ) + + # Short skill identifier + name: Mapped[str] = mapped_column(String, nullable=False, doc="Short skill identifier (e.g. 'deploy-production')") + + # Distinguish the type/category of the skill + entry_type: Mapped[str] = mapped_column(String, doc="Category or type (e.g. 'workflow', 'guide', 'script')") + + # A human-friendly description of this skill + description: Mapped[str] = mapped_column(String, doc="Short descriptive text about what this skill accomplishes") + + # Step-by-step instructions as plain text + instructions: Mapped[str] = mapped_column(String, doc="Step-by-step instructions as a single string") + + # Conditions that indicate this skill is relevant + triggers: Mapped[list] = mapped_column(JSON, default=list, doc="Conditions that indicate this skill is relevant") + + # Input/output examples + examples: Mapped[list] = mapped_column(JSON, default=list, doc="Input/output examples for this skill") + + # Semantic version + version: Mapped[str] = mapped_column(String, default="0.1.0", doc="Semantic version of this skill") + + # NEW: Filter tags for flexible filtering and categorization + filter_tags: Mapped[Optional[dict]] = mapped_column( + JSON, nullable=True, default=None, doc="Custom filter tags for filtering and categorization" + ) + + # When was this item last modified and what operation? + last_modify: Mapped[dict] = mapped_column( + JSON, + nullable=False, + default=lambda: { + "timestamp": datetime.now(dt.timezone.utc).isoformat(), + "operation": "created", + }, + doc="Last modification info including timestamp and operation type", + ) + + embedding_config: Mapped[Optional[dict]] = mapped_column( + EmbeddingConfigColumn, nullable=True, doc="Embedding configuration" + ) + + # Vector embedding field based on database type + if settings.mirix_pg_uri_no_default: + from pgvector.sqlalchemy import Vector + + description_embedding = mapped_column(Vector(MAX_EMBEDDING_DIM), nullable=True) + instructions_embedding = mapped_column(Vector(MAX_EMBEDDING_DIM), nullable=True) + else: + description_embedding = Column(CommonVector, nullable=True) + instructions_embedding = Column(CommonVector, nullable=True) + + # Database indexes for efficient querying. + # NB: use an explicit ``is not None`` filter rather than ``filter(None, ...)`` + # because an un-attached UniqueConstraint is falsy (it iterates over its + # columns which are empty until binding), so plain truthiness would drop it. + __table_args__ = tuple( + item + for item in [ + # Organization-level query optimization indexes + ( + Index("ix_procedural_memory_organization_id", "organization_id") + if settings.mirix_pg_uri_no_default + else None + ), + ( + Index( + "ix_procedural_memory_org_created_at", + "organization_id", + "created_at", + postgresql_using="btree", + ) + if settings.mirix_pg_uri_no_default + else None + ), + ( + Index( + "ix_procedural_memory_filter_tags_gin", + text("(filter_tags::jsonb)"), + postgresql_using="gin", + ) + if settings.mirix_pg_uri_no_default + else None + ), + ( + Index( + "ix_procedural_memory_org_filter_scope", + "organization_id", + text("((filter_tags->>'scope')::text)"), + postgresql_using="btree", + ) + if settings.mirix_pg_uri_no_default + else None + ), + # Enforce per-user skill-name uniqueness at the DB level. The + # application also pre-checks via a search, but that is racy + # under concurrent creates โ€” this constraint closes the race. + UniqueConstraint( + "organization_id", + "user_id", + "name", + name="uq_procedural_memory_org_user_name", + ), + # SQLite indexes + ( + Index("ix_procedural_memory_organization_id_sqlite", "organization_id") + if not settings.mirix_pg_uri_no_default + else None + ), + ] + if item is not None + ) + + @declared_attr + def agent(cls) -> Mapped[Optional["Agent"]]: + """ + Relationship to the Agent that owns this procedural memory item. + """ + return relationship("Agent", lazy="selectin") + + @declared_attr + def organization(cls) -> Mapped["Organization"]: + """ + Relationship to organization (mirroring your existing patterns). + Adjust 'back_populates' to match the collection name in your `Organization` model. + """ + return relationship("Organization", back_populates="procedural_memory", lazy="selectin") + + @declared_attr + def user(cls) -> Mapped["User"]: + """ + Relationship to the User that owns this procedural memory item. + """ + return relationship("User", lazy="selectin") diff --git a/mirix/orm/skill_experience.py b/mirix/orm/skill_experience.py new file mode 100644 index 000000000..b328a280b --- /dev/null +++ b/mirix/orm/skill_experience.py @@ -0,0 +1,146 @@ +"""ORM model for the general session-experience store. + +A `SkillExperience` is one transferable lesson distilled from ONE work +session's transcript and later consumed by the procedural skill agent to +create/edit skills. Every experience is either `worth_learning` (an approach +worth repeating) or `worth_avoiding` (a pitfall to avoid), scored by +`importance` and `credibility` in [0,1]. There is no external oracle: both are +derived purely from the conversation content. + +Records start `pending`; an evolution run flips them to `consumed` (recording +its run id in `consumed_by` + the skills it influenced in `influenced_skill_ids`), +or later evidence flips them to `superseded`. + +Inherited (do NOT redeclare): + - id, created_at, updated_at, is_deleted -- SqlalchemyBase + meta mixin + - organization_id -- OrganizationMixin + - user_id -- UserMixin + - agent_id (FK agents.id ON DELETE CASCADE) -- AgentMixin + +`experience_type` and `status` are plain Strings (NOT pg ENUMs); their value +spaces are Literal-validated in `schemas/skill_experience.py`, so adding a value +never needs a DB migration. +""" + +import uuid +from typing import TYPE_CHECKING, List, Optional + +from sqlalchemy import JSON, Float, Index, String, Text +from sqlalchemy.orm import Mapped, declared_attr, mapped_column, relationship + +from mirix.orm.mixins import AgentMixin, OrganizationMixin, UserMixin +from mirix.orm.sqlalchemy_base import SqlalchemyBase +from mirix.schemas.skill_experience import ( + SkillExperience as PydanticSkillExperience, +) +from mirix.settings import settings + +if TYPE_CHECKING: + from mirix.orm.agent import Agent + from mirix.orm.organization import Organization + from mirix.orm.user import User + + +class SkillExperience(SqlalchemyBase, OrganizationMixin, UserMixin, AgentMixin): + """One distilled, transferable experience from a single work session.""" + + __tablename__ = "skill_experience" + __pydantic_model__ = PydanticSkillExperience + + # Override the bare SqlalchemyBase PK to attach a prefixed default, matching + # the `proc-`/`sevr-` convention on sibling tables. + id: Mapped[str] = mapped_column( + String, + primary_key=True, + default=lambda: f"sexp-{uuid.uuid4()}", + doc="Unique ID for this skill-experience.", + ) + + session_id: Mapped[str] = mapped_column( + String, + nullable=False, + doc="Provenance: the session this experience was distilled from.", + ) + experience_type: Mapped[str] = mapped_column( + String, + nullable=False, + doc="'worth_learning' or 'worth_avoiding' (validated in Pydantic, not a pg ENUM).", + ) + title: Mapped[str] = mapped_column( + String, nullable=False, doc="Short headline for the experience." + ) + # content/evidence are NOT NULL (with a server-side '' default for the + # migration path) because the pydantic full schema requires non-null strings; + # leaving them nullable would let a NULL row break to_pydantic()/list. + content: Mapped[str] = mapped_column( + Text, + nullable=False, + default="", + doc="The transferable lesson: when-to-apply (learn) or how-to-avoid (avoid).", + ) + importance: Mapped[float] = mapped_column( + Float, + nullable=False, + default=0.0, + doc="How impactful/worth-acting-on the lesson is, in [0,1].", + ) + credibility: Mapped[float] = mapped_column( + Float, + nullable=False, + default=0.0, + doc="How well-grounded in a direct signal the lesson is, in [0,1].", + ) + evidence: Mapped[str] = mapped_column( + Text, + nullable=False, + default="", + doc="JSON string of the in-conversation signal: {quote, signal_type}.", + ) + status: Mapped[str] = mapped_column( + String, + nullable=False, + default="pending", + doc="Lifecycle: pending | consumed | superseded (validated in Pydantic).", + ) + consumed_by: Mapped[Optional[str]] = mapped_column( + String, nullable=True, doc="Evolution-run id that consumed this experience." + ) + influenced_skill_ids: Mapped[Optional[List[str]]] = mapped_column( + JSON, + nullable=True, + default=None, + doc="Skill ids this experience influenced post-evolution (lineage).", + ) + + # Indexes mirror the procedural_memory pg/sqlite guard pattern: an explicit + # `is not None` filter (not `filter(None, ...)`) because an un-attached + # Index is falsy until bound. + __table_args__ = tuple( + item + for item in [ + # Primary access pattern: list_experiences(agent_id, status). + Index( + "ix_skill_experience_agent_status", + "agent_id", + "status", + ), + ( + Index("ix_skill_experience_organization_id", "organization_id") + if settings.mirix_pg_uri_no_default + else None + ), + ] + if item is not None + ) + + @declared_attr + def agent(cls) -> Mapped[Optional["Agent"]]: + return relationship("Agent", lazy="selectin") + + @declared_attr + def user(cls) -> Mapped[Optional["User"]]: + return relationship("User", lazy="selectin") + + @declared_attr + def organization(cls) -> Mapped[Optional["Organization"]]: + return relationship("Organization", lazy="selectin") diff --git a/mirix/prompts/system/base/auto_dream_agent/procedural.txt b/mirix/prompts/system/base/auto_dream_agent/procedural.txt index 4743b431e..84e108452 100644 --- a/mirix/prompts/system/base/auto_dream_agent/procedural.txt +++ b/mirix/prompts/system/base/auto_dream_agent/procedural.txt @@ -1,11 +1,55 @@ -You are the Auto Dream Agent for Procedural Memory. +You are the Experience Distiller. You read the FULL transcript of ONE work session and distill the transferable EXPERIENCES worth remembering. A separate Skill Manager later reads these experiences to evolve a bank of procedural skills, so each experience must be concrete, actionable, and free of one-off noise. -Input contains workflow and how-to memories. Review them for duplicate procedures, overlapping steps, stale instructions, and conflicting workflows. +There is NO external grader and NO ground-truth answer key. You judge purely from the conversation content itself. -Use procedural_memory_update(old_ids, new_items) to merge, replace, or delete entries. Preserve complete and actionable steps. +## What You Receive -Rules: -- Only modify IDs present in the input. -- Prefer merging similar workflows into one clear procedure. -- If two workflows conflict, keep the more complete or more recent version, or record the ambiguity. -- Call finish_memory_update when done. +You receive ONE session's chronological messages (user / assistant / tool turns, including tool calls and their results), plus: +- `agent_id` and `session_id` โ€” provenance only. +- `existing_skills` โ€” the names and descriptions of skills that already exist, so you can avoid re-deriving a lesson the skill bank already captures (prefer NEW or sharper lessons). + +## Your Task + +Emit a JSON ARRAY of 0 or more experiences. A single session can yield MULTIPLE experiences, and a MIX of the two kinds below. Do NOT label the whole session a success or a failure โ€” extract the individual lessons. + +Each experience is exactly one of: + +- `worth_learning` โ€” an approach, decision, or technique that worked and is worth REPEATING. In `content`, state the concrete reusable approach AND when to apply it ("When , do , because "). +- `worth_avoiding` โ€” a pitfall, mistake, or dead-end to AVOID. In `content`, state what went wrong AND how to avoid it next time ("; to avoid it, "). + +If the session contains nothing worth remembering, output an empty array `[]`. + +## Signals to Ground On (these drive `credibility`) + +Derive every experience from observable signals in the transcript: +- DIRECT user critique or correction โ€” e.g. "this part isn't good enough", "that's wrong", "no, do it differently". Strong signal for `worth_avoiding`. +- DIRECT user confirmation โ€” e.g. "great, that worked", "perfect", "yes exactly". Strong signal for `worth_learning`. +- Tool errors / failed calls / retries โ€” a tool returning an error, then a corrected call. Strong signal. +- The agent's own self-correction โ€” the agent noticing a mistake and fixing it mid-session. +- Inference โ€” a lesson you reasonably infer from the flow but with no explicit signal. Weakest. + +## Scoring (each in [0.0, 1.0]) + +- `credibility` โ€” how well-grounded the lesson is in a DIRECT/explicit signal. HIGH (0.8-1.0) when grounded in explicit user critique, user confirmation, or a concrete tool error. MEDIUM (0.4-0.7) for a clear self-correction. LOW (0.1-0.3) when merely inferred from flow with no explicit signal. +- `importance` โ€” how impactful and worth-acting-on the lesson is. A pivotal decision or a costly, recurring mistake scores high; a minor or one-off detail scores low. + +## Evidence + +For each experience, cite the grounding signal in `evidence`: +- `quote` โ€” a SHORT verbatim snippet from the transcript that grounds the lesson (the user critique, the tool error text, the confirmation, etc.). Keep it brief. +- `signal_type` โ€” exactly one of: `user_critique`, `user_confirmation`, `tool_error`, `self_correction`, `inferred`. + +## Output Format (STRICT) + +Output ONLY a JSON array, no prose before or after, no markdown fences. Each element: + +{ + "experience_type": "worth_learning" | "worth_avoiding", + "title": "", + "content": "", + "importance": 0.0, + "credibility": 0.0, + "evidence": { "quote": "", "signal_type": "user_critique" } +} + +Output `[]` if nothing in this session is worth remembering. Do not invent signals that are not in the transcript. diff --git a/mirix/prompts/system/base/meta_memory_agent.txt b/mirix/prompts/system/base/meta_memory_agent.txt index da445d94d..faeedcdaa 100644 --- a/mirix/prompts/system/base/meta_memory_agent.txt +++ b/mirix/prompts/system/base/meta_memory_agent.txt @@ -1,4 +1,4 @@ -You are the Meta Memory Manager, as a part of a memory system. In this system, other agents are: Meta Memory Manager, Episodic Memory Manager, Procedural Memory Manager, Resource Memory Manager, Semantic Memory Manager, Core Memory Manager, Knowledge Vault Memory Manager, and Chat Agent. All agents share the same memories. +You are the Meta Memory Manager, as a part of a memory system. In this system, other agents are: Meta Memory Manager, Episodic Memory Manager, Procedural Memory Manager, Resource Memory Manager, Semantic Memory Manager, Core Memory Manager, Knowledge Vault Memory Manager, and Chat Agent. All agents share the same memories. The system will receive various types of messages from users, including text messages, images, transcripted voice recordings, and other multimedia content. When messages are accumulated to a certain amount, they will be sent to you, along with potential conversations between the user and the Chat Agent during this period. You need to analyze the input messages and conversations, understand what the user is communicating and going through, then save the details into corresponding memories by calling the function `trigger_memory_update` and send the messages to corresponding memory manager so that they will update the memories according to your instructions. When calling `trigger_memory_update`, you can choose one or more memories to update from ['core', 'episodic', 'procedural', 'resource', 'knowledge_vault', 'semantic']. The details of all memory components, and how you are supposed to classify content into different memories, are shown below: @@ -10,14 +10,14 @@ WHO the user is and HOW to interact with them effectively. Purpose: Stores user's personal characteristics and interaction preferences to optimize communication. Contains: User identity, personality traits, communication style preferences, relationship information, behavioral patterns that directly influence how the assistant should respond. Key Question: "What do I need to know about this user to communicate with them better?" -Examples: +Examples: - "User's name is Alex, prefers direct communication, works as a software engineer" - "Has a partner named Sam who is an artist, values work-life balance" - "Prefers final responses only, doesn't want intermediate messages" - "Gets frustrated with overly technical explanations, prefers simple analogies" Classification Rule: Update when you identify ANY personal information, preference, or trait that affects how you should interact with the user. -2. Episodic Memory - Time-Ordered Personal Events & Experiences +2. Episodic Memory - Time-Ordered Personal Events & Experiences WHAT happened WHEN in the user's life. Purpose: Chronicles the user's experiences and events in chronological order. @@ -29,16 +29,16 @@ Examples: - "Occurred at 2025-03-03 14:00:00 - Started learning Python programming, feeling overwhelmed by syntax" Classification Rule: Update for virtually ALL user activities, experiences, and significant interactions that happened at a specific time. (Hints: Episodic Memory almost always needs to be updated) -3. Procedural Memory - Step-by-Step Instructions & Processes -HOW to do specific tasks or follow procedures. -Purpose: Stores reusable instructions and workflows for accomplishing tasks. -Contains: Workflows, tutorials, step-by-step guides, and repeatable processes. -Key Question: "What are the steps to accomplish this task?" +3. Procedural Memory - Reusable Skills & Procedures +HOW to do specific tasks โ€” stored as structured, reusable skills. +Purpose: Stores procedural knowledge as skills with triggers, detailed instructions, and examples. +Contains: Workflows, tutorials, how-to guides, and repeatable processes โ€” each with a name, description, trigger conditions, and comprehensive instructions. +Key Question: "What reusable skill or procedure can be extracted from this?" Examples: -- "How to reset router: 1. Unplug device 2. Wait 10 seconds 3. Plug back in 4. Wait for lights" -- "Daily morning routine: 1. Check emails 2. Review calendar 3. Prioritize tasks 4. Start with hardest task" -- "Code review process: 1. Check functionality 2. Review style 3. Test edge cases 4. Approve/request changes" -Classification Rule: Update when messages contain sequential steps, workflows, or instructional content. +- Skill "reset-router": triggers when user mentions network issues, instructions cover full reset procedure with troubleshooting +- Skill "morning-routine": triggers on productivity discussions, instructions detail the complete daily startup workflow +- Skill "code-review-process": triggers when discussing code quality, instructions cover the full review checklist +Classification Rule: Update when messages contain sequential steps, workflows, instructional content, or when existing skills could be improved based on new information. 4. Resource Memory - Documents, Files & Reference Materials WHAT files and documents are available for reference. @@ -94,9 +94,7 @@ When processing messages, think step by step, evaluate each memory type in this - User's interactions and personal timeline - Update for almost all user activities and experiences -(3) Procedural Memory: Does this explain HOW TO do something? -- Step-by-step instructions, workflows, processes -- Sequential procedures or tutorials +(3) Procedural Memory: (Automatically triggered โ€” do NOT include "procedural" in your trigger_memory_update call. The system handles procedural memory updates on a fixed schedule.) (4) Resource Memory: Does this involve WHAT files or documents? - Shared files, documents, images, reference materials @@ -109,6 +107,6 @@ When processing messages, think step by step, evaluate each memory type in this - New concepts, people, places, organizations you didn't know - Definitions and abstract knowledge about the world -After evaluation, call `trigger_memory_update` with the appropriate memory types that require updates. You may select from: `['core', 'episodic', 'semantic', 'resource', 'procedural', 'knowledge_vault']`. When uncertain about classification, it is preferable to include additional memory types rather than omit potentially relevant ones, as the specialized memory managers will perform detailed analysis and filtering. Comprehensive coverage ensures no important information is lost during the memory update process. +After evaluation, call `trigger_memory_update` with the appropriate memory types that require updates. You may select from: `['core', 'episodic', 'semantic', 'resource', 'knowledge_vault']`. Note: procedural memory is triggered automatically and should NOT be included in your selection. When uncertain about classification, it is preferable to include additional memory types rather than omit potentially relevant ones, as the specialized memory managers will perform detailed analysis and filtering. Comprehensive coverage ensures no important information is lost during the memory update process. After all updates, you must call `finish_memory_update` to complete the process. Between steps (1)-(6) where you are calling `trigger_memory_update`, you can also do reasoning about the memory update. Whenever you do not call any tools, the message would be treated as reasoning message. \ No newline at end of file diff --git a/mirix/prompts/system/base/procedural_memory_agent.txt b/mirix/prompts/system/base/procedural_memory_agent.txt index e86ca3314..b39e2aadd 100644 --- a/mirix/prompts/system/base/procedural_memory_agent.txt +++ b/mirix/prompts/system/base/procedural_memory_agent.txt @@ -1,42 +1,65 @@ -You are the Procedural Memory Manager, one of six agents in a memory system. The other agents are the Meta Memory Manager, Episodic Memory Manager, Resource Memory Manager, Knowledge Vault Memory Manager, and the Chat Agent. You do not see or interact directly with these other agentsโ€”but you share the same memory base with them. +You are the Skill Manager, responsible for maintaining reusable procedural knowledge (skills) extracted from user interactions. You have access to CLI-style tools for managing skills. -The system will receive various types of messages from users, including text messages, images, transcripted voice recordings, and other multimedia content. When messages are accumulated to a certain amount, they will be sent to you, along with potential conversations between the user and the Chat Agent during this period. You need to analyze the input messages and conversations, extract step-by-step instructions, "how-to" guides, and any other instructions and skills, and save them into the procedural memory. +## Available Tools -This memory base includes the following components: +- `skill_list(query?, limit?)` โ€” List skills. Use without query to see all, or with a search term to filter. +- `skill_read(name_or_id)` โ€” Read the full content of a skill by name or ID. +- `skill_create(name, description, instructions, entry_type, triggers?, examples?)` โ€” Create a new skill. +- `skill_edit(skill_id, field, old_text, new_text)` โ€” Patch a text field (name, description, instructions) with find-and-replace. +- `skill_edit(skill_id, field, value)` โ€” Replace a non-text field (triggers, examples, entry_type) entirely. +- `skill_delete(skill_id)` โ€” Delete a skill. +- `finish_memory_update()` โ€” Call this when you are done to end the session. -1. Core Memory: -Contains fundamental information about the user, such as the name, personality, simple information that should help with the communication with the user. +## Your Workflow -2. Episodic Memory: -Stores time-ordered, event-based information from interactionsโ€”essentially, the "diary" of user and assistant events. +When you receive a batch of messages (possibly with an instruction from the Meta Memory Manager), follow this process: -3. Procedural Memory: -Definition: Contains how-to guides, step-by-step instructions, or processes the assistant or user might follow. -Example: "How to reset the router." -Each entry in Procedural Memory has: - (a) entry_type (e.g., 'workflow', 'guide', 'script') - (b) description (short descriptive text) - (c) steps (the procedure in a structured or JSON format) +1. **Survey existing skills**: Call `skill_list()` to see what skills already exist. +2. **Analyze the messages**: Identify any reusable procedural knowledge โ€” workflows, how-to guides, step-by-step processes, or improvements to existing procedures. +3. **Decide on actions**: + - If you find a NEW procedure not covered by existing skills โ†’ `skill_create` + - If you find information that IMPROVES an existing skill โ†’ `skill_read` to see current content, then `skill_edit` to patch it + - If an existing skill is OBSOLETE or WRONG โ†’ `skill_delete` (or edit to fix) + - If the messages contain NO procedural content โ†’ skip directly to step 4 +4. **Finish**: Call `finish_memory_update()` to complete the session. -4. Resource Memory: -Contains documents, files, and reference materials related to ongoing tasks or projects. +## Skill Format Guidelines -5. Knowledge Vault: -A repository for static, structured factual data such as phone numbers, email addresses, passwords, or other knowledge that are not necessarily always needed during the conversation but are potentially useful at some future point. +- **name**: Concise kebab-case identifier (e.g., "deploy-production", "weekly-report-generation"). Must be specific, NOT generic like "guide" or "workflow". +- **description**: One or two sentences about what this skill does and when it's useful. +- **instructions**: Comprehensive plain text โ€” detailed enough for someone to follow without additional context. Include context, caveats, and conditions. Not just a numbered list. +- **entry_type**: One of "workflow", "guide", or "script". +- **triggers**: Conditions that signal this skill is relevant (e.g., ["user mentions deployment", "discussing release process"]). +- **examples**: Input/output pairs if present in the conversation (e.g., [{"input": "Deploy the app", "output": "Deployment started..."}]). -6. Semantic Memory: -Contains general knowledge about a concept (e.g. a new software name, a new concept) or an object (e.g. a person, a place, where the details would be the understanding and information about them.) +## Important Rules -When receiving messages and potentially a message from the meta agent (There will be a bracket saying "[Instruction from Meta Memory Manager]"), make a single comprehensive memory update: +- Always check existing skills before creating โ€” avoid duplicates. +- Prefer `skill_edit` over delete+recreate when improving existing skills. +- You may make multiple tool calls in sequence (list โ†’ read โ†’ edit โ†’ finish). +- **You MUST call `finish_memory_update()` to mark the end of this round's work. It is the ONLY valid way to end the session.** Call it even when the messages contain no procedural content and you made no changes. Never end your turn โ€” and never stop calling tools โ€” without first calling `finish_memory_update()`. -**Single Function Call Process:** -1. **Analyze Content**: Examine all messages and conversations to identify step-by-step instructions, "how-to" guides, workflows, or any procedural knowledge. -2. **Make Update**: Use ONE appropriate procedural memory function to save the most important identified procedure or instruction with proper entry_type ('workflow', 'guide', 'script'), description, and detailed steps. -3. **Skip Update if Necessary**: If there are no updates to make, call `finish_memory_update` with no arguments to complete the process. +## Records-Based Evolution (Curator Mode) -**Important Notes:** -- Make only ONE function call total (either a procedural memory function OR finish_memory_update) -- Look for any structured processes, workflows, or instructional content in the messages -- Save procedures with appropriate entry_type ('workflow', 'guide', 'script'), description, and detailed steps -- `finish_memory_update` takes no parameters - call it as is when there's nothing to update -- Prioritize the most complete or useful procedural information if multiple procedures are present \ No newline at end of file +When the input is a set of DISTILLED RECORDS (failures first, each a `[FAILURE|SUCCESS] title โ€” detail (round; evidence: โ€ฆ)` line) rather than raw chat, you are acting as the skill CURATOR. Follow these rules: + +- **Additive-first.** Prefer creating a new skill or extending an existing one over removing knowledge. Never delete correct guidance. +- **Cite a record for every change.** Each `skill_create` / `skill_edit` you make MUST be justified by a specific record in the input (its failure root_cause or success what-worked). If no record supports a change, do not make it. +- **Prefer EDIT over CREATE when a skill should have prevented the failure.** First survey existing skills (`skill_list`). If an existing skill is the one that *should have* prevented a failure, `skill_edit` it to close the gap. Only `skill_create` when no existing skill covers the failure/success at all. +- **Never DELETE unless a failure record names the skill actively harmful.** Redundant or merely-outdated is NOT a reason to delete โ€” fix it with `skill_edit` instead. A delete is permitted ONLY when a failure record explicitly identifies a skill as actively harmful (e.g. it caused the failure). Even then, deletion is soft (the skill is excluded from retrieval, not destroyed). +- **Respect the edit budget.** You have a bounded number of create+edit operations this run. Spend them on the highest-value failures first. If a tool returns "edit budget exhausted", stop making changes and call `finish_memory_update()`. +- **Keep edits small.** If a tool rejects an edit as "too large", split it into a smaller patch or extract a new skill with `skill_create` instead of growing one skill unbounded. +- **Do not re-propose superseded changes.** If the input has a "Do NOT re-propose" block, avoid re-introducing those reverted edits (it causes oscillation). + +## Experience-Based Evolution (Self-Improvement Mode) + +When the input is a set of DISTILLED EXPERIENCES โ€” each a `[LEARN|AVOID] title โ€” content (importance=โ€ฆ, credibility=โ€ฆ; evidence: "โ€ฆ")` line, highest-priority first โ€” you are evolving your own skills from what recent work sessions taught you. There is NO external grader; each experience was derived purely from the conversation (a direct user critique/confirmation, a tool error, the agent's own self-correction, or weaker inference). Follow these rules: + +- **`[LEARN]` = an approach worth repeating.** Encode it into a skill so it is applied next time. The `content` states *when to apply* it. +- **`[AVOID]` = a pitfall to guard against.** Add an explicit caveat / "do NOT โ€ฆ" guidance to the relevant skill so the mistake is not repeated. The `content` states *how to avoid* it. +- **Work highest-priority first.** Experiences are ordered by importance ร— credibility (priority). Spend your bounded edit budget on the top lessons; a high-credibility lesson grounded in a direct user critique or a tool error outweighs a weakly-inferred one. +- **Survey before mutating.** Call `skill_list()` (and `skill_read` as needed) first. If an existing skill already covers the topic, `skill_edit` it to fold in the lesson (delta/incremental โ€” patch the gap, don't rewrite). Only `skill_create` when no existing skill covers it. Dedupe against existing skill names. +- **Delta, never wholesale.** Keep each edit small and targeted; if a tool rejects an edit as "too large", split it or extract a focused new skill instead of growing one skill unbounded. +- **No deletes in this mode.** An experience names a lesson, not a skill to destroy. Improve with `skill_edit`; do not `skill_delete`. +- **Respect the edit budget.** You have a bounded number of create+edit operations this run. If a tool returns "edit budget exhausted", stop and call `finish_memory_update()`. +- **Always finish.** Call `finish_memory_update()` to end the round, even if no change was warranted. diff --git a/mirix/queue/message.proto b/mirix/queue/message.proto index 9f83ed751..e7c322c98 100644 --- a/mirix/queue/message.proto +++ b/mirix/queue/message.proto @@ -98,6 +98,9 @@ message MessageCreate { // Optional multi-agent group id optional string group_id = 7; + + // Optional top-level session identifier (see mirix.schemas.message.MessageCreate.session_id). + optional string session_id = 8; } // List of message content parts (for structured content) diff --git a/mirix/queue/message_pb2.py b/mirix/queue/message_pb2.py index 32cdfd9c3..fda4612db 100644 --- a/mirix/queue/message_pb2.py +++ b/mirix/queue/message_pb2.py @@ -26,7 +26,7 @@ from google.protobuf import struct_pb2 as google_dot_protobuf_dot_struct__pb2 -DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x19mirix/queue/message.proto\x12\x05mirix\x1a\x1fgoogle/protobuf/timestamp.proto\x1a\x1cgoogle/protobuf/struct.proto\"\xaf\x05\n\x0cQueueMessage\x12\x11\n\tclient_id\x18\x01 \x01(\t\x12\x10\n\x08\x61gent_id\x18\x02 \x01(\t\x12,\n\x0einput_messages\x18\x03 \x03(\x0b\x32\x14.mirix.MessageCreate\x12\x15\n\x08\x63haining\x18\x04 \x01(\x08H\x00\x88\x01\x01\x12\x14\n\x07user_id\x18\x05 \x01(\tH\x01\x88\x01\x01\x12\x14\n\x07verbose\x18\x06 \x01(\x08H\x02\x88\x01\x01\x12,\n\x0b\x66ilter_tags\x18\x07 \x01(\x0b\x32\x17.google.protobuf.Struct\x12\x16\n\tuse_cache\x18\x08 \x01(\x08H\x03\x88\x01\x01\x12\x18\n\x0boccurred_at\x18\t \x01(\tH\x04\x88\x01\x01\x12\x1e\n\x11langfuse_trace_id\x18\n \x01(\tH\x05\x88\x01\x01\x12$\n\x17langfuse_observation_id\x18\x0b \x01(\tH\x06\x88\x01\x01\x12 \n\x13langfuse_session_id\x18\x0c \x01(\tH\x07\x88\x01\x01\x12\x1d\n\x10langfuse_user_id\x18\r \x01(\tH\x08\x88\x01\x01\x12\x32\n\x11\x62lock_filter_tags\x18\x0e \x01(\x0b\x32\x17.google.protobuf.Struct\x12*\n\x1d\x62lock_filter_tags_update_mode\x18\x0f \x01(\tH\t\x88\x01\x01\x42\x0b\n\t_chainingB\n\n\x08_user_idB\n\n\x08_verboseB\x0c\n\n_use_cacheB\x0e\n\x0c_occurred_atB\x14\n\x12_langfuse_trace_idB\x1a\n\x18_langfuse_observation_idB\x16\n\x14_langfuse_session_idB\x13\n\x11_langfuse_user_idB \n\x1e_block_filter_tags_update_mode\"\xcf\x01\n\x04User\x12\n\n\x02id\x18\x01 \x01(\t\x12\x17\n\x0forganization_id\x18\x02 \x01(\t\x12\x0c\n\x04name\x18\x03 \x01(\t\x12\x0e\n\x06status\x18\x04 \x01(\t\x12\x10\n\x08timezone\x18\x05 \x01(\t\x12.\n\ncreated_at\x18\x06 \x01(\x0b\x32\x1a.google.protobuf.Timestamp\x12.\n\nupdated_at\x18\x07 \x01(\x0b\x32\x1a.google.protobuf.Timestamp\x12\x12\n\nis_deleted\x18\x08 \x01(\x08\"\xd9\x02\n\rMessageCreate\x12\'\n\x04role\x18\x01 \x01(\x0e\x32\x19.mirix.MessageCreate.Role\x12\x16\n\x0ctext_content\x18\x02 \x01(\tH\x00\x12\x37\n\x12structured_content\x18\x03 \x01(\x0b\x32\x19.mirix.MessageContentListH\x00\x12\x11\n\x04name\x18\x04 \x01(\tH\x01\x88\x01\x01\x12\x11\n\x04otid\x18\x05 \x01(\tH\x02\x88\x01\x01\x12\x16\n\tsender_id\x18\x06 \x01(\tH\x03\x88\x01\x01\x12\x15\n\x08group_id\x18\x07 \x01(\tH\x04\x88\x01\x01\"<\n\x04Role\x12\x14\n\x10ROLE_UNSPECIFIED\x10\x00\x12\r\n\tROLE_USER\x10\x01\x12\x0f\n\x0bROLE_SYSTEM\x10\x02\x42\x0e\n\x0c\x63ontent_typeB\x07\n\x05_nameB\x07\n\x05_otidB\x0c\n\n_sender_idB\x0b\n\t_group_id\">\n\x12MessageContentList\x12(\n\x05parts\x18\x01 \x03(\x0b\x32\x19.mirix.MessageContentPart\"\xbc\x01\n\x12MessageContentPart\x12\"\n\x04text\x18\x01 \x01(\x0b\x32\x12.mirix.TextContentH\x00\x12$\n\x05image\x18\x02 \x01(\x0b\x32\x13.mirix.ImageContentH\x00\x12\"\n\x04\x66ile\x18\x03 \x01(\x0b\x32\x12.mirix.FileContentH\x00\x12-\n\ncloud_file\x18\x04 \x01(\x0b\x32\x17.mirix.CloudFileContentH\x00\x42\t\n\x07\x63ontent\"\x1b\n\x0bTextContent\x12\x0c\n\x04text\x18\x01 \x01(\t\"@\n\x0cImageContent\x12\x10\n\x08image_id\x18\x01 \x01(\t\x12\x13\n\x06\x64\x65tail\x18\x02 \x01(\tH\x00\x88\x01\x01\x42\t\n\x07_detail\"\x1e\n\x0b\x46ileContent\x12\x0f\n\x07\x66ile_id\x18\x01 \x01(\t\"*\n\x10\x43loudFileContent\x12\x16\n\x0e\x63loud_file_uri\x18\x01 \x01(\tb\x06proto3') +DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x19mirix/queue/message.proto\x12\x05mirix\x1a\x1fgoogle/protobuf/timestamp.proto\x1a\x1cgoogle/protobuf/struct.proto\"\xaf\x05\n\x0cQueueMessage\x12\x11\n\tclient_id\x18\x01 \x01(\t\x12\x10\n\x08\x61gent_id\x18\x02 \x01(\t\x12,\n\x0einput_messages\x18\x03 \x03(\x0b\x32\x14.mirix.MessageCreate\x12\x15\n\x08\x63haining\x18\x04 \x01(\x08H\x00\x88\x01\x01\x12\x14\n\x07user_id\x18\x05 \x01(\tH\x01\x88\x01\x01\x12\x14\n\x07verbose\x18\x06 \x01(\x08H\x02\x88\x01\x01\x12,\n\x0b\x66ilter_tags\x18\x07 \x01(\x0b\x32\x17.google.protobuf.Struct\x12\x16\n\tuse_cache\x18\x08 \x01(\x08H\x03\x88\x01\x01\x12\x18\n\x0boccurred_at\x18\t \x01(\tH\x04\x88\x01\x01\x12\x1e\n\x11langfuse_trace_id\x18\n \x01(\tH\x05\x88\x01\x01\x12$\n\x17langfuse_observation_id\x18\x0b \x01(\tH\x06\x88\x01\x01\x12 \n\x13langfuse_session_id\x18\x0c \x01(\tH\x07\x88\x01\x01\x12\x1d\n\x10langfuse_user_id\x18\r \x01(\tH\x08\x88\x01\x01\x12\x32\n\x11\x62lock_filter_tags\x18\x0e \x01(\x0b\x32\x17.google.protobuf.Struct\x12*\n\x1d\x62lock_filter_tags_update_mode\x18\x0f \x01(\tH\t\x88\x01\x01\x42\x0b\n\t_chainingB\n\n\x08_user_idB\n\n\x08_verboseB\x0c\n\n_use_cacheB\x0e\n\x0c_occurred_atB\x14\n\x12_langfuse_trace_idB\x1a\n\x18_langfuse_observation_idB\x16\n\x14_langfuse_session_idB\x13\n\x11_langfuse_user_idB \n\x1e_block_filter_tags_update_mode\"\xcf\x01\n\x04User\x12\n\n\x02id\x18\x01 \x01(\t\x12\x17\n\x0forganization_id\x18\x02 \x01(\t\x12\x0c\n\x04name\x18\x03 \x01(\t\x12\x0e\n\x06status\x18\x04 \x01(\t\x12\x10\n\x08timezone\x18\x05 \x01(\t\x12.\n\ncreated_at\x18\x06 \x01(\x0b\x32\x1a.google.protobuf.Timestamp\x12.\n\nupdated_at\x18\x07 \x01(\x0b\x32\x1a.google.protobuf.Timestamp\x12\x12\n\nis_deleted\x18\x08 \x01(\x08\"\x81\x03\n\rMessageCreate\x12\'\n\x04role\x18\x01 \x01(\x0e\x32\x19.mirix.MessageCreate.Role\x12\x16\n\x0ctext_content\x18\x02 \x01(\tH\x00\x12\x37\n\x12structured_content\x18\x03 \x01(\x0b\x32\x19.mirix.MessageContentListH\x00\x12\x11\n\x04name\x18\x04 \x01(\tH\x01\x88\x01\x01\x12\x11\n\x04otid\x18\x05 \x01(\tH\x02\x88\x01\x01\x12\x16\n\tsender_id\x18\x06 \x01(\tH\x03\x88\x01\x01\x12\x15\n\x08group_id\x18\x07 \x01(\tH\x04\x88\x01\x01\x12\x17\n\nsession_id\x18\x08 \x01(\tH\x05\x88\x01\x01\"<\n\x04Role\x12\x14\n\x10ROLE_UNSPECIFIED\x10\x00\x12\r\n\tROLE_USER\x10\x01\x12\x0f\n\x0bROLE_SYSTEM\x10\x02\x42\x0e\n\x0c\x63ontent_typeB\x07\n\x05_nameB\x07\n\x05_otidB\x0c\n\n_sender_idB\x0b\n\t_group_idB\r\n\x0b_session_id\">\n\x12MessageContentList\x12(\n\x05parts\x18\x01 \x03(\x0b\x32\x19.mirix.MessageContentPart\"\xbc\x01\n\x12MessageContentPart\x12\"\n\x04text\x18\x01 \x01(\x0b\x32\x12.mirix.TextContentH\x00\x12$\n\x05image\x18\x02 \x01(\x0b\x32\x13.mirix.ImageContentH\x00\x12\"\n\x04\x66ile\x18\x03 \x01(\x0b\x32\x12.mirix.FileContentH\x00\x12-\n\ncloud_file\x18\x04 \x01(\x0b\x32\x17.mirix.CloudFileContentH\x00\x42\t\n\x07\x63ontent\"\x1b\n\x0bTextContent\x12\x0c\n\x04text\x18\x01 \x01(\t\"@\n\x0cImageContent\x12\x10\n\x08image_id\x18\x01 \x01(\t\x12\x13\n\x06\x64\x65tail\x18\x02 \x01(\tH\x00\x88\x01\x01\x42\t\n\x07_detail\"\x1e\n\x0b\x46ileContent\x12\x0f\n\x07\x66ile_id\x18\x01 \x01(\t\"*\n\x10\x43loudFileContent\x12\x16\n\x0e\x63loud_file_uri\x18\x01 \x01(\tb\x06proto3') _globals = globals() _builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals) @@ -38,19 +38,19 @@ _globals['_USER']._serialized_start=790 _globals['_USER']._serialized_end=997 _globals['_MESSAGECREATE']._serialized_start=1000 - _globals['_MESSAGECREATE']._serialized_end=1345 - _globals['_MESSAGECREATE_ROLE']._serialized_start=1224 - _globals['_MESSAGECREATE_ROLE']._serialized_end=1284 - _globals['_MESSAGECONTENTLIST']._serialized_start=1347 - _globals['_MESSAGECONTENTLIST']._serialized_end=1409 - _globals['_MESSAGECONTENTPART']._serialized_start=1412 - _globals['_MESSAGECONTENTPART']._serialized_end=1600 - _globals['_TEXTCONTENT']._serialized_start=1602 - _globals['_TEXTCONTENT']._serialized_end=1629 - _globals['_IMAGECONTENT']._serialized_start=1631 - _globals['_IMAGECONTENT']._serialized_end=1695 - _globals['_FILECONTENT']._serialized_start=1697 - _globals['_FILECONTENT']._serialized_end=1727 - _globals['_CLOUDFILECONTENT']._serialized_start=1729 - _globals['_CLOUDFILECONTENT']._serialized_end=1771 + _globals['_MESSAGECREATE']._serialized_end=1385 + _globals['_MESSAGECREATE_ROLE']._serialized_start=1249 + _globals['_MESSAGECREATE_ROLE']._serialized_end=1309 + _globals['_MESSAGECONTENTLIST']._serialized_start=1387 + _globals['_MESSAGECONTENTLIST']._serialized_end=1449 + _globals['_MESSAGECONTENTPART']._serialized_start=1452 + _globals['_MESSAGECONTENTPART']._serialized_end=1640 + _globals['_TEXTCONTENT']._serialized_start=1642 + _globals['_TEXTCONTENT']._serialized_end=1669 + _globals['_IMAGECONTENT']._serialized_start=1671 + _globals['_IMAGECONTENT']._serialized_end=1735 + _globals['_FILECONTENT']._serialized_start=1737 + _globals['_FILECONTENT']._serialized_end=1767 + _globals['_CLOUDFILECONTENT']._serialized_start=1769 + _globals['_CLOUDFILECONTENT']._serialized_end=1811 # @@protoc_insertion_point(module_scope) diff --git a/mirix/queue/message_pb2.pyi b/mirix/queue/message_pb2.pyi index 4822f24f9..1d36fcc53 100644 --- a/mirix/queue/message_pb2.pyi +++ b/mirix/queue/message_pb2.pyi @@ -63,7 +63,7 @@ class User(_message.Message): def __init__(self, id: _Optional[str] = ..., organization_id: _Optional[str] = ..., name: _Optional[str] = ..., status: _Optional[str] = ..., timezone: _Optional[str] = ..., created_at: _Optional[_Union[_timestamp_pb2.Timestamp, _Mapping]] = ..., updated_at: _Optional[_Union[_timestamp_pb2.Timestamp, _Mapping]] = ..., is_deleted: bool = ...) -> None: ... class MessageCreate(_message.Message): - __slots__ = ("role", "text_content", "structured_content", "name", "otid", "sender_id", "group_id") + __slots__ = ("role", "text_content", "structured_content", "name", "otid", "sender_id", "group_id", "session_id") class Role(int, metaclass=_enum_type_wrapper.EnumTypeWrapper): __slots__ = () ROLE_UNSPECIFIED: _ClassVar[MessageCreate.Role] @@ -79,6 +79,7 @@ class MessageCreate(_message.Message): OTID_FIELD_NUMBER: _ClassVar[int] SENDER_ID_FIELD_NUMBER: _ClassVar[int] GROUP_ID_FIELD_NUMBER: _ClassVar[int] + SESSION_ID_FIELD_NUMBER: _ClassVar[int] role: MessageCreate.Role text_content: str structured_content: MessageContentList @@ -86,7 +87,8 @@ class MessageCreate(_message.Message): otid: str sender_id: str group_id: str - def __init__(self, role: _Optional[_Union[MessageCreate.Role, str]] = ..., text_content: _Optional[str] = ..., structured_content: _Optional[_Union[MessageContentList, _Mapping]] = ..., name: _Optional[str] = ..., otid: _Optional[str] = ..., sender_id: _Optional[str] = ..., group_id: _Optional[str] = ...) -> None: ... + session_id: str + def __init__(self, role: _Optional[_Union[MessageCreate.Role, str]] = ..., text_content: _Optional[str] = ..., structured_content: _Optional[_Union[MessageContentList, _Mapping]] = ..., name: _Optional[str] = ..., otid: _Optional[str] = ..., sender_id: _Optional[str] = ..., group_id: _Optional[str] = ..., session_id: _Optional[str] = ...) -> None: ... class MessageContentList(_message.Message): __slots__ = ("parts",) diff --git a/mirix/queue/queue_util.py b/mirix/queue/queue_util.py index deea94891..41b0cac9c 100644 --- a/mirix/queue/queue_util.py +++ b/mirix/queue/queue_util.py @@ -145,6 +145,8 @@ async def put_messages( proto_msg.sender_id = msg.sender_id if msg.group_id: proto_msg.group_id = msg.group_id + if msg.session_id: + proto_msg.session_id = msg.session_id proto_input_messages.append(proto_msg) diff --git a/mirix/queue/worker.py b/mirix/queue/worker.py index 2cde1df6a..f590d6742 100644 --- a/mirix/queue/worker.py +++ b/mirix/queue/worker.py @@ -98,6 +98,7 @@ def _convert_proto_message_to_pydantic(self, proto_msg) -> "MessageCreate": otid=proto_msg.otid if proto_msg.HasField("otid") else None, sender_id=proto_msg.sender_id if proto_msg.HasField("sender_id") else None, group_id=proto_msg.group_id if proto_msg.HasField("group_id") else None, + session_id=proto_msg.session_id if proto_msg.HasField("session_id") else None, filter_tags=None, ) diff --git a/mirix/schemas/agent_trigger_state.py b/mirix/schemas/agent_trigger_state.py new file mode 100644 index 000000000..10d2e5898 --- /dev/null +++ b/mirix/schemas/agent_trigger_state.py @@ -0,0 +1,86 @@ +from datetime import datetime +from typing import List, Optional + +from pydantic import Field, field_validator + +from mirix.schemas.mirix_base import MirixBase + +TRIGGER_TYPE_MAX_LEN = 64 +TRIGGER_TYPE_PATTERN = r"^[a-z][a-z0-9_]{0,63}$" + +# Registered trigger kinds. Adding a new trigger? Add its string here so +# callers can't typo their way into silently creating a parallel bookkeeping +# row. Shared with the manager and with memory_tools. +TRIGGER_TYPE_PROCEDURAL_SKILL = "procedural_skill" +KNOWN_TRIGGER_TYPES = frozenset({TRIGGER_TYPE_PROCEDURAL_SKILL}) + + +def _validate_trigger_type(value: str) -> str: + if not isinstance(value, str): + raise ValueError(f"trigger_type must be a string, got {type(value).__name__}") + if len(value) == 0 or len(value) > TRIGGER_TYPE_MAX_LEN: + raise ValueError( + f"trigger_type must be 1..{TRIGGER_TYPE_MAX_LEN} characters, got len={len(value)}" + ) + import re + + if not re.fullmatch(TRIGGER_TYPE_PATTERN, value): + raise ValueError( + f"trigger_type '{value}' does not match pattern {TRIGGER_TYPE_PATTERN}" + ) + # Membership check closes the "typo creates a parallel bookkeeping row" + # footgun. Adding a new trigger means adding it to KNOWN_TRIGGER_TYPES โ€” + # that is the whole point of the registry. The length/pattern checks + # above stay as a defensive belt for future additions. + if value not in KNOWN_TRIGGER_TYPES: + raise ValueError( + f"trigger_type '{value}' is not in the registry " + f"{sorted(KNOWN_TRIGGER_TYPES)}. Add it to KNOWN_TRIGGER_TYPES " + f"in mirix/schemas/agent_trigger_state.py." + ) + return value + + +class AgentTriggerStateBase(MirixBase): + """Shared fields for agent_trigger_state.""" + + __id_prefix__ = "ats" + + agent_id: str = Field(..., description="ID of the agent this trigger state belongs to.") + user_id: str = Field(..., description="ID of the user this trigger state belongs to.") + trigger_type: str = Field( + ..., + max_length=TRIGGER_TYPE_MAX_LEN, + description="Kind of trigger (e.g. 'procedural_skill').", + ) + + @field_validator("trigger_type") + @classmethod + def _check_trigger_type(cls, value: str) -> str: + return _validate_trigger_type(value) + + +class AgentTriggerState(AgentTriggerStateBase): + """Persisted trigger-state row.""" + + id: Optional[str] = Field(None, description="Unique identifier for this trigger state row.") + organization_id: Optional[str] = Field(None, description="Owning organization id.") + last_fired_at: Optional[datetime] = Field( + None, + description="created_at of the message that caused the last fire; lower " + "bound for 'new sessions since last fire' counts.", + ) + last_fired_session_id: Optional[str] = Field( + None, + max_length=64, + description="session_id of the message that caused the last fire.", + ) + last_fired_tied_session_ids: Optional[List[str]] = Field( + default_factory=list, + description="session_ids whose first-appearance timestamp equals " + "last_fired_at. The fire filter uses MIN(created_at) per session_id so " + "a given session_id counts in exactly one window; this tied set covers " + "the rare exact-microsecond tie at the watermark.", + ) + created_at: Optional[datetime] = Field(None, description="Row creation timestamp.") + updated_at: Optional[datetime] = Field(None, description="Last update timestamp.") diff --git a/mirix/schemas/auto_dream.py b/mirix/schemas/auto_dream.py index 78d0870ba..7b0e715f7 100644 --- a/mirix/schemas/auto_dream.py +++ b/mirix/schemas/auto_dream.py @@ -19,6 +19,22 @@ class AutoDreamRequest(BaseModel): ) dry_run: bool = Field(False, description="If true, return plan without applying changes") model: Optional[str] = Field(None, description="Override LLM model (e.g. gpt-4.1-mini for testing)") + meta_agent_id: Optional[str] = Field( + None, + description=( + "Optional meta memory agent id to run auto-dream against. When omitted, " + "the server uses the client's first meta memory agent for backward " + "compatibility." + ), + ) + last_n_sessions: Optional[int] = Field( + None, + description=( + "For mode='procedural' (general session-experience distillation): how " + "many of the meta agent's most-recent retained sessions to distill. " + "Defaults to MESSAGE_RETAIN_LAST_N_SESSIONS." + ), + ) @field_validator("mode") @classmethod @@ -43,3 +59,8 @@ class AutoDreamResponse(BaseModel): last_dream_at: dt.datetime dry_run: bool message: str = "" + # Structured evolution counts for the procedural path so drivers can + # health-gate without parsing `message`. Default 0/empty keeps every + # existing caller (and non-procedural modes) working unchanged. + skills_changed: int = 0 + changes: Dict[str, list] = Field(default_factory=dict) diff --git a/mirix/schemas/conversation_message.py b/mirix/schemas/conversation_message.py new file mode 100644 index 000000000..b9744b0a7 --- /dev/null +++ b/mirix/schemas/conversation_message.py @@ -0,0 +1,100 @@ +"""Pydantic schemas for the Conversation Message Store. + +A `ConversationMessage` is one external conversation turn โ€” a real `user`, +`assistant`, or `tool` message that arrived through the memory-add API carrying +a `session_id`. It is the canonical, learnable record of a conversation and the +single source the procedural-memory (skill) distiller reads. Tool turns matter: +the work-process lessons (a tool error, a retry, a fix that finally worked) live +in tool calls and tool results, so dropping them would blind the distiller to +exactly the worth_avoiding/worth_learning signal it exists to extract. + +Mirrors `schemas/skill_experience.py`: a Base with the user-facing fields, a +`Create` schema that carries the owner ids needed to persist, a full schema with +DB columns, and a Response alias. `role` is a `Literal` (validated here in +Pydantic), NOT a pg ENUM โ€” the ORM column is a plain String, so changing the +value space never requires a DB migration. `session_id` reuses the shared +validator from `schemas/message.py` so the conversation store and the agent-loop +`messages` table agree on exactly one session_id format. +""" + +from datetime import datetime +from typing import Literal, Optional + +from pydantic import Field, field_validator + +from mirix.helpers.datetime_helpers import get_utc_time +from mirix.schemas.message import _validate_session_id +from mirix.schemas.mirix_base import MirixBase + +# role value space. Kept as a Literal so the schema is the single source of +# truth; the ORM stores it as a plain string. 'tool' covers both tool results +# (role: "tool" messages) and serialized assistant tool calls. +ConversationRole = Literal["user", "assistant", "tool"] + +# Length cap on a single turn. Generous โ€” conversation turns can be long โ€” but +# bounded so a pathological caller can't bloat the DB. +CONVERSATION_MESSAGE_MAX_CONTENT_LEN = 65536 + + +class ConversationMessageBase(MirixBase): + """Shared, user-facing fields for a conversation turn.""" + + __id_prefix__ = "convmsg" + + session_id: str = Field( + ..., + description="The conversation this turn belongs to (NOT NULL in this store).", + ) + role: ConversationRole = Field( + ..., + description="The real turn role: 'user', 'assistant', or 'tool'.", + ) + content: str = Field( + default="", + max_length=CONVERSATION_MESSAGE_MAX_CONTENT_LEN, + description="The verbatim text of this conversation turn.", + ) + + @field_validator("session_id") + @classmethod + def _check_session_id(cls, v): + """Enforce the shared session_id format; reject None/empty. + + `_validate_session_id` permits None for the agent-loop `messages` table, + but this store only ever holds session'd turns, so a None/empty value is + rejected here. + """ + validated = _validate_session_id(v) + if validated is None: + raise ValueError("session_id is required for a conversation message") + return validated + + +class ConversationMessageCreate(ConversationMessageBase): + """Create schema โ€” carries the owner ids needed to persist the row.""" + + user_id: str = Field(..., description="The id of the user that owns this turn.") + organization_id: str = Field(..., description="The owning organization id.") + + +class ConversationMessage(ConversationMessageBase): + """Full conversation-message schema, with database-related fields.""" + + id: Optional[str] = Field(None, description="Unique identifier for this turn.") + user_id: str = Field(..., description="The id of the user that owns this turn.") + organization_id: str = Field(..., description="The owning organization id.") + distilled_at: Optional[datetime] = Field( + None, + description="Set when this session was consumed by a distill round " + "(NULL = not yet distilled).", + ) + created_at: datetime = Field( + default_factory=get_utc_time, description="Creation timestamp." + ) + updated_at: Optional[datetime] = Field(None, description="Last update timestamp.") + + +class ConversationMessageResponse(ConversationMessage): + """Response schema for a conversation message.""" + + pass diff --git a/mirix/schemas/embedding_config.py b/mirix/schemas/embedding_config.py index 43a2522d7..3d838bbe9 100755 --- a/mirix/schemas/embedding_config.py +++ b/mirix/schemas/embedding_config.py @@ -40,6 +40,7 @@ class EmbeddingConfig(BaseModel): "hugging-face", "mistral", "together", # completions endpoint + "openrouter", ] = Field(..., description="The endpoint type for the model.") embedding_endpoint: Optional[str] = Field(None, description="The endpoint for the model (`None` if local).") embedding_model: str = Field(..., description="The model for the embedding.") diff --git a/mirix/schemas/llm_config.py b/mirix/schemas/llm_config.py index e3a243cad..8743f6e5e 100755 --- a/mirix/schemas/llm_config.py +++ b/mirix/schemas/llm_config.py @@ -50,6 +50,7 @@ class LLMConfig(BaseModel): "bedrock", "deepseek", "xai", + "openrouter", ] = Field(..., description="The endpoint type for the model.") model_endpoint: Optional[str] = Field(None, description="The endpoint for the model.") model_wrapper: Optional[str] = Field(None, description="The wrapper for the model.") diff --git a/mirix/schemas/message.py b/mirix/schemas/message.py index d3ece01cc..c466b2e21 100644 --- a/mirix/schemas/message.py +++ b/mirix/schemas/message.py @@ -45,6 +45,44 @@ from mirix.system import unpack_message +import re + +# Single source of truth for the session_id column. The DB CHECK constraint in +# mirix/orm/message.py and the SQL migration in scripts/migrate_add_message_session_id.sql +# are derived from these โ€” keep them in sync via tests/test_session_id.py. +SESSION_ID_MAX_LEN = 64 +SESSION_ID_ALLOWED_CHARS = "A-Za-z0-9_-" +# Python-side Pattern: at least 1 char. The length cap is enforced separately so +# we can return a precise error message. +SESSION_ID_PATTERN = f"^[{SESSION_ID_ALLOWED_CHARS}]+$" +# DB-side pattern includes the length cap inline because Postgres regex has no +# "length=max" notion beyond a bounded quantifier. +SESSION_ID_SQL_PATTERN = f"^[{SESSION_ID_ALLOWED_CHARS}]{{1,{SESSION_ID_MAX_LEN}}}$" +_SESSION_ID_RE = re.compile(SESSION_ID_PATTERN) + + +def _validate_session_id(value: Optional[str]) -> Optional[str]: + """Shared validator for the top-level session_id field. + + Accepts None or a non-empty string of [A-Za-z0-9_-], up to SESSION_ID_MAX_LEN. + Rejects empty strings so the DB layer never stores an ambiguous "". + """ + if value is None: + return None + if not isinstance(value, str) or value == "": + raise ValueError("session_id must be a non-empty string or null") + if len(value) > SESSION_ID_MAX_LEN: + raise ValueError( + f"session_id exceeds max length of {SESSION_ID_MAX_LEN}" + ) + if not _SESSION_ID_RE.match(value): + raise ValueError( + f"session_id must match [{SESSION_ID_ALLOWED_CHARS}]+ " + "(letters, digits, '_' or '-')" + ) + return value + + class BaseMessage(OrmMetadataBase): __id_prefix__ = "message" @@ -69,10 +107,23 @@ class MessageCreate(BaseModel): description="The id of the sender of the message, can be an identity id or agent id", ) group_id: Optional[str] = Field(None, description="The multi-agent group that the message was sent in") + session_id: Optional[str] = Field( + None, + description=( + "Optional top-level conversation/session identifier. Messages sharing a session_id " + "belong to the same logical interaction. Must match [a-zA-Z0-9_-]+ and be <= 64 chars." + ), + examples=["sess-xyz"], + ) filter_tags: Optional[Dict[str, Any]] = Field( None, description="Optional tags for filtering and categorizing this message and related memories" ) + @field_validator("session_id") + @classmethod + def _check_session_id(cls, v: Optional[str]) -> Optional[str]: + return _validate_session_id(v) + def model_dump(self, to_orm: bool = False, **kwargs) -> Dict[str, Any]: data = super().model_dump(**kwargs) if to_orm and "content" in data: @@ -100,6 +151,15 @@ class MessageUpdate(BaseModel): # created_at: Optional[datetime] = Field(None, description="The time the message was created.") tool_calls: Optional[List[OpenAIToolCall,]] = Field(None, description="The list of tool calls requested.") tool_call_id: Optional[str] = Field(None, description="The id of the tool call.") + session_id: Optional[str] = Field( + None, + description="Update the message's top-level session_id. Same validation as MessageCreate.session_id.", + ) + + @field_validator("session_id") + @classmethod + def _check_session_id(cls, v: Optional[str]) -> Optional[str]: + return _validate_session_id(v) def model_dump(self, to_orm: bool = False, **kwargs) -> Dict[str, Any]: data = super().model_dump(**kwargs) @@ -157,6 +217,13 @@ class Message(BaseMessage): None, description="The id of the sender of the message, can be an identity id or agent id", ) + session_id: Optional[str] = Field( + None, + description=( + "Top-level conversation/session identifier. Messages sharing a session_id belong to " + "the same logical interaction. Indexed column; must match [a-zA-Z0-9_-]+ and be <= 64 chars." + ), + ) # This overrides the optional base orm schema, created_at MUST exist on all messages objects created_at: datetime = Field( default_factory=get_utc_time, @@ -179,6 +246,11 @@ def validate_role(cls, v: str) -> str: assert v in roles, f"Role must be one of {roles}" return v + @field_validator("session_id") + @classmethod + def _check_session_id(cls, v: Optional[str]) -> Optional[str]: + return _validate_session_id(v) + def to_json(self): json_message = vars(self) if json_message["tool_calls"] is not None: @@ -457,6 +529,7 @@ def dict_to_message( name: Optional[str] = None, group_id: Optional[str] = None, tool_returns: Optional[List[ToolReturn]] = None, + session_id: Optional[str] = None, ): """Convert a ChatCompletion message object into a Message object (synced to DB)""" if not created_at: @@ -516,6 +589,7 @@ def dict_to_message( id=str(id), tool_returns=tool_returns, group_id=group_id, + session_id=session_id, ) else: return Message( @@ -530,6 +604,7 @@ def dict_to_message( created_at=created_at, tool_returns=tool_returns, group_id=group_id, + session_id=session_id, ) elif "function_call" in openai_message_dict and openai_message_dict["function_call"] is not None: @@ -565,6 +640,7 @@ def dict_to_message( id=str(id), tool_returns=tool_returns, group_id=group_id, + session_id=session_id, ) else: return Message( @@ -579,6 +655,7 @@ def dict_to_message( created_at=created_at, tool_returns=tool_returns, group_id=group_id, + session_id=session_id, ) else: @@ -628,6 +705,7 @@ def dict_to_message( id=str(id), tool_returns=tool_returns, group_id=group_id, + session_id=session_id, ) else: return Message( @@ -642,6 +720,7 @@ def dict_to_message( created_at=created_at, tool_returns=tool_returns, group_id=group_id, + session_id=session_id, ) def to_openai_dict_search_results(self, max_tool_id_length: int = TOOL_CALL_ID_MAX_LEN) -> dict: diff --git a/mirix/schemas/procedural_memory.py b/mirix/schemas/procedural_memory.py index c818aec4f..1846eef3d 100755 --- a/mirix/schemas/procedural_memory.py +++ b/mirix/schemas/procedural_memory.py @@ -8,16 +8,112 @@ from mirix.schemas.embedding_config import EmbeddingConfig from mirix.schemas.mirix_base import MirixBase +# Allowed skill entry_type values. Kept as a set so tests and REST handlers +# can import and reuse the same source of truth. +SKILL_ENTRY_TYPES = frozenset({"workflow", "guide", "script"}) + +# Length caps for skill fields. They are intentionally generous: real skills +# can be long. The goal is to prevent pathological inputs (e.g. multi-megabyte +# instruction blobs) from hitting the DB and embedding pipeline, not to +# police prose. Revisit if users routinely bump against them. +SKILL_MAX_NAME_LEN = 128 +SKILL_MAX_DESCRIPTION_LEN = 1024 +SKILL_MAX_INSTRUCTIONS_LEN = 65_536 +SKILL_MAX_TRIGGER_LEN = 512 +SKILL_MAX_EXAMPLES_COUNT = 50 + + +def _validate_entry_type(value: Optional[str]) -> Optional[str]: + if value is None: + return value + if value not in SKILL_ENTRY_TYPES: + raise ValueError( + f"Invalid entry_type '{value}'. Must be one of: {sorted(SKILL_ENTRY_TYPES)}." + ) + return value + + +def normalize_entry_type(value: Any) -> str: + """Coerce a legacy/free-form entry_type into the closed skill set. + + Rows written before the skill schema carried free-form LLM output + ("process", "How-To", ...). Reads must never crash on that data, so the + schema normalizes instead of rejecting; strict validation of NEW values + stays in the write path (mirix/agent/tool_validators.py). Mirrors the + CASE mapping in scripts/migrate_procedural_to_skill.sql. + """ + lowered = str(value).strip().lower() if value is not None else "" + if lowered in SKILL_ENTRY_TYPES: + return lowered + if "script" in lowered: + return "script" + if "guide" in lowered or "how" in lowered: + return "guide" + return "workflow" + + +def _validate_triggers(triggers: Optional[List[str]]) -> Optional[List[str]]: + if triggers is None: + return triggers + for i, trigger in enumerate(triggers): + if not isinstance(trigger, str): + raise ValueError(f"triggers[{i}] must be a string, got {type(trigger).__name__}.") + if len(trigger) > SKILL_MAX_TRIGGER_LEN: + raise ValueError( + f"triggers[{i}] exceeds max length {SKILL_MAX_TRIGGER_LEN}." + ) + return triggers + + +def _validate_examples(examples: Optional[List[dict]]) -> Optional[List[dict]]: + if examples is None: + return examples + if len(examples) > SKILL_MAX_EXAMPLES_COUNT: + raise ValueError( + f"examples has {len(examples)} entries; max is {SKILL_MAX_EXAMPLES_COUNT}." + ) + for i, example in enumerate(examples): + if not isinstance(example, dict): + raise ValueError(f"examples[{i}] must be an object, got {type(example).__name__}.") + return examples + class ProceduralMemoryItemBase(MirixBase): """ - Base schema for storing procedural knowledge (e.g., workflows, methods). + Base schema for storing procedural knowledge as reusable skills. """ __id_prefix__ = "proc_item" - entry_type: str = Field(..., description="Category (e.g., 'workflow', 'guide', 'script')") - summary: str = Field(..., description="Short descriptive text about the procedure") - steps: List[str] = Field(..., description="Step-by-step instructions as a list of strings") + name: str = Field( + ..., max_length=SKILL_MAX_NAME_LEN, description="Short skill identifier (e.g., 'deploy-production')" + ) + entry_type: str = Field(..., description="Category (one of: 'workflow', 'guide', 'script')") + description: str = Field( + ..., max_length=SKILL_MAX_DESCRIPTION_LEN, description="Short descriptive text about the skill" + ) + instructions: str = Field( + ..., max_length=SKILL_MAX_INSTRUCTIONS_LEN, description="Step-by-step instructions as a single string" + ) + triggers: List[str] = Field(default_factory=list, description="Conditions indicating this skill is relevant") + examples: List[dict] = Field(default_factory=list, description="Input/output examples for this skill") + + @field_validator("entry_type", mode="before") + @classmethod + def _check_entry_type(cls, value: Any) -> str: + # Lenient on purpose: this validator also runs on READS + # (to_pydantic -> model_validate with from_attributes), where legacy + # rows may hold pre-skill free-form values that must not crash reads. + return normalize_entry_type(value) + + @field_validator("triggers") + @classmethod + def _check_triggers(cls, value: List[str]) -> List[str]: + return _validate_triggers(value) or [] + + @field_validator("examples") + @classmethod + def _check_examples(cls, value: List[dict]) -> List[dict]: + return _validate_examples(value) or [] class ProceduralMemoryItem(ProceduralMemoryItemBase): @@ -39,23 +135,14 @@ class ProceduralMemoryItem(ProceduralMemoryItemBase): description="Last modification info including timestamp and operation type", ) organization_id: str = Field(..., description="The unique identifier of the organization") - summary_embedding: Optional[List[float]] = Field(None, description="The embedding of the summary") - steps_embedding: Optional[List[float]] = Field(None, description="The embedding of the steps") + version: str = Field(default="0.1.0", description="Semantic version of this skill") + description_embedding: Optional[List[float]] = Field(None, description="The embedding of the description") + instructions_embedding: Optional[List[float]] = Field(None, description="The embedding of the instructions") embedding_config: Optional[EmbeddingConfig] = Field( - None, description="The embedding configuration used by the event" - ) - - # NEW: Filter tags for flexible filtering and categorization - filter_tags: Optional[Dict[str, Any]] = Field( - default=None, - description="Custom filter tags for filtering and categorization", - examples=[ - {"project_id": "proj-abc", "session_id": "sess-xyz", "tags": ["important", "work"], "priority": "high"} - ], + None, description="The embedding configuration used by the skill" ) - # need to validate both steps_embedding and summary_embedding to ensure they are the same size - # NEW: Filter tags for flexible filtering and categorization + # Filter tags for flexible filtering and categorization filter_tags: Optional[Dict[str, Any]] = Field( default=None, description="Custom filter tags for filtering and categorization", @@ -64,7 +151,7 @@ class ProceduralMemoryItem(ProceduralMemoryItemBase): ], ) - @field_validator("summary_embedding", "steps_embedding") + @field_validator("description_embedding", "instructions_embedding") @classmethod def pad_embeddings(cls, embedding: List[float]) -> List[float]: """Pad embeddings to `MAX_EMBEDDING_SIZE`. This is necessary to ensure all stored embeddings are the same size.""" @@ -86,21 +173,64 @@ class ProceduralMemoryItemUpdate(MirixBase): id: str = Field(..., description="Unique ID for this procedural memory entry") agent_id: Optional[str] = Field(None, description="The id of the agent this procedural memory item belongs to") - entry_type: Optional[str] = Field(None, description="Category (e.g., 'workflow', 'guide', 'script')") - summary: Optional[str] = Field(None, description="Short descriptive text") - steps: Optional[List[str]] = Field(None, description="Step-by-step instructions as a list of strings") + name: Optional[str] = Field(None, max_length=SKILL_MAX_NAME_LEN, description="Short skill identifier") + entry_type: Optional[str] = Field(None, description="Category (one of: 'workflow', 'guide', 'script')") + description: Optional[str] = Field( + None, max_length=SKILL_MAX_DESCRIPTION_LEN, description="Short descriptive text about the skill" + ) + instructions: Optional[str] = Field( + None, max_length=SKILL_MAX_INSTRUCTIONS_LEN, description="Step-by-step instructions as a single string" + ) + triggers: Optional[List[str]] = Field(None, description="Conditions indicating this skill is relevant") + examples: Optional[List[dict]] = Field(None, description="Input/output examples for this skill") + version: Optional[str] = Field(None, description="Semantic version of this skill") organization_id: Optional[str] = Field(None, description="The organization ID") updated_at: datetime = Field(default_factory=get_utc_time, description="Update timestamp") + + @field_validator("entry_type") + @classmethod + def _check_entry_type(cls, value: Optional[str]) -> Optional[str]: + return _validate_entry_type(value) + + @field_validator("triggers") + @classmethod + def _check_triggers(cls, value: Optional[List[str]]) -> Optional[List[str]]: + return _validate_triggers(value) + + @field_validator("examples") + @classmethod + def _check_examples(cls, value: Optional[List[dict]]) -> Optional[List[dict]]: + return _validate_examples(value) last_modify: Optional[Dict[str, Any]] = Field( None, description="Last modification info including timestamp and operation type", ) - steps_embedding: Optional[List[float]] = Field(None, description="The embedding of the event") - summary_embedding: Optional[List[float]] = Field(None, description="The embedding of the summary") + instructions_embedding: Optional[List[float]] = Field(None, description="The embedding of the instructions") + description_embedding: Optional[List[float]] = Field(None, description="The embedding of the description") embedding_config: Optional[EmbeddingConfig] = Field( - None, description="The embedding configuration used by the event" + None, description="The embedding configuration used by the skill" ) + @field_validator("description_embedding", "instructions_embedding") + @classmethod + def pad_embeddings(cls, embedding: Optional[List[float]]) -> Optional[List[float]]: + """Pad embeddings to `MAX_EMBEDDING_DIM` so REST-supplied vectors on the + UPDATE path are stored at the same width as every other write path + (the CREATE schema pads identically). Without this a raw provider + vector (e.g. 3072-dim) would mismatch the Vector(4096) column and break + SQLAlchemy compare_values on the next UPDATE.""" + import numpy as np + + if embedding and len(embedding) != MAX_EMBEDDING_DIM: + np_embedding = np.array(embedding) + padded_embedding = np.pad( + np_embedding, + (0, MAX_EMBEDDING_DIM - np_embedding.shape[0]), + mode="constant", + ) + return padded_embedding.tolist() + return embedding + filter_tags: Optional[Dict[str, Any]] = Field( None, description="Custom filter tags for filtering and categorization" ) diff --git a/mirix/schemas/skill_experience.py b/mirix/schemas/skill_experience.py new file mode 100644 index 000000000..a020feff1 --- /dev/null +++ b/mirix/schemas/skill_experience.py @@ -0,0 +1,156 @@ +"""Pydantic schemas for the general session-experience store. + +A `SkillExperience` is one transferable lesson distilled from a single work +session's transcript. Each is either `worth_learning` (an approach worth +repeating) or `worth_avoiding` (a pitfall to avoid), scored by `importance` and +`credibility` in [0,1]. Records start `pending`, are `consumed` by a skill +evolution run, or are `superseded` when later evidence overrides them. + +This module has a Base with the user-facing fields, a full schema with DB +columns, an Update schema, and a Response alias. Enums are `Literal`s (validated +here in Pydantic), NOT pg ENUMs โ€” the ORM column is a plain String, so adding a +value never requires a DB migration. `importance` and `credibility` are clamped +to [0,1] by validators (garbage -> 0.0). +""" + +import math +from datetime import datetime +from typing import List, Literal, Optional + +from pydantic import Field, field_validator + +from mirix.client.utils import get_utc_time +from mirix.schemas.mirix_base import MirixBase + +# experience_type / status value spaces. Kept as Literals so the schema is the +# single source of truth; the ORM stores them as plain strings. +ExperienceType = Literal["worth_learning", "worth_avoiding"] +ExperienceStatus = Literal["pending", "consumed", "superseded"] + +# Length caps. Generous โ€” these hold distilled prose, not transcripts โ€” but +# bounded so a pathological distiller output can't bloat the DB or the curator +# payload. +SKILL_EXPERIENCE_MAX_TITLE_LEN = 256 +SKILL_EXPERIENCE_MAX_CONTENT_LEN = 8192 +SKILL_EXPERIENCE_MAX_EVIDENCE_LEN = 2048 + + +def _clamp01(value, default: float = 0.0) -> float: + """Coerce a distiller-reported score into [0.0, 1.0]; default on garbage. + + Kept in sync with the distiller so validation and persistence agree on + out-of-range/garbage handling. + """ + try: + f = float(value) + except (TypeError, ValueError): + return default + # NaN slips past BOTH the < 0 and > 1 comparisons (all NaN comparisons are + # False), which would poison the importance*credibility priority ordering โ€” + # reject it. (+-inf is fine: it clamps to the bounds below.) + if math.isnan(f): + return default + if f < 0.0: + return 0.0 + if f > 1.0: + return 1.0 + return f + + +class SkillExperienceBase(MirixBase): + """Shared, user-facing fields for a session-experience.""" + + __id_prefix__ = "sexp" + + session_id: str = Field( + ..., description="Provenance: the session this experience was distilled from." + ) + experience_type: ExperienceType = Field( + ..., + description="'worth_learning' (approach to repeat) or 'worth_avoiding' (pitfall).", + ) + title: str = Field( + ..., + max_length=SKILL_EXPERIENCE_MAX_TITLE_LEN, + description="Short headline for the experience.", + ) + content: str = Field( + default="", + max_length=SKILL_EXPERIENCE_MAX_CONTENT_LEN, + description="The transferable lesson: when-to-apply (learn) or how-to-avoid (avoid).", + ) + importance: float = Field( + default=0.0, + description="How impactful/worth-acting-on the lesson is, in [0,1].", + ) + credibility: float = Field( + default=0.0, + description="How well-grounded in a direct signal the lesson is, in [0,1].", + ) + evidence: str = Field( + default="", + max_length=SKILL_EXPERIENCE_MAX_EVIDENCE_LEN, + description="JSON string of the in-conversation signal: {quote, signal_type}.", + ) + status: ExperienceStatus = Field( + default="pending", + description="Lifecycle: pending -> consumed (by an evolution run) or superseded.", + ) + + @field_validator("importance", "credibility", mode="before") + @classmethod + def _clamp_scores(cls, v): + """Clamp importance/credibility into [0,1]; coerce garbage -> 0.0.""" + return _clamp01(v) + + +class SkillExperienceCreate(SkillExperienceBase): + """Create schema โ€” carries the owner ids needed to persist the row.""" + + agent_id: str = Field(..., description="The id of the agent that owns this experience.") + user_id: str = Field(..., description="The id of the user that owns this experience.") + organization_id: str = Field(..., description="The owning organization id.") + + +class SkillExperience(SkillExperienceBase): + """Full experience schema, with database-related fields.""" + + id: Optional[str] = Field(None, description="Unique identifier for this experience.") + agent_id: Optional[str] = Field( + None, description="The id of the agent that owns this experience." + ) + user_id: str = Field(..., description="The id of the user that owns this experience.") + organization_id: str = Field(..., description="The owning organization id.") + consumed_by: Optional[str] = Field( + None, + description="Evolution-run id that consumed this experience (set on mark_consumed).", + ) + influenced_skill_ids: Optional[List[str]] = Field( + None, description="Skill ids this experience influenced post-evolution (lineage)." + ) + created_at: datetime = Field( + default_factory=get_utc_time, description="Creation timestamp." + ) + updated_at: Optional[datetime] = Field(None, description="Last update timestamp.") + + +class SkillExperienceUpdate(MirixBase): + """Schema for updating an existing experience (status transitions + lineage).""" + + id: str = Field(..., description="Unique ID for this experience.") + status: Optional[ExperienceStatus] = Field(None, description="New lifecycle status.") + consumed_by: Optional[str] = Field( + None, description="Evolution-run id that consumed this experience." + ) + influenced_skill_ids: Optional[List[str]] = Field( + None, description="Skill ids this experience influenced post-evolution." + ) + updated_at: datetime = Field( + default_factory=get_utc_time, description="Update timestamp." + ) + + +class SkillExperienceResponse(SkillExperience): + """Response schema for a session-experience.""" + + pass diff --git a/mirix/sdk.py b/mirix/sdk.py index 8632c4ae4..44c52a5a9 100644 --- a/mirix/sdk.py +++ b/mirix/sdk.py @@ -736,33 +736,20 @@ async def visualize_memories(self, user_id: Optional[str] = None) -> Dict[str, A else: procedural_items = [] - memories["procedural"] = [] - for item in procedural_items: - import json - - # Parse steps if it's a JSON string - steps = item.steps - if isinstance(steps, str): - try: - steps = json.loads(steps) - # Extract just the instruction text for simpler display - if isinstance(steps, list) and steps and isinstance(steps[0], dict): - steps = [step.get("instruction", str(step)) for step in steps] - except (json.JSONDecodeError, KeyError, TypeError): - # If parsing fails, keep as string and split by common delimiters - if isinstance(steps, str): - steps = [s.strip() for s in steps.replace("\n", "|").split("|") if s.strip()] - else: - steps = [] - - memories["procedural"].append( - { - "title": item.entry_type, - "type": "procedural", - "summary": item.summary, - "steps": steps if isinstance(steps, list) else [], - } - ) + memories["procedural"] = [ + { + "title": item.name, + "type": "procedural", + "entry_type": item.entry_type, + "name": item.name, + "description": item.description, + "instructions": item.instructions, + "triggers": getattr(item, "triggers", None) or [], + "examples": getattr(item, "examples", None) or [], + "version": getattr(item, "version", None), + } + for item in procedural_items + ] except Exception: memories["procedural"] = [] diff --git a/mirix/server/rest_api.py b/mirix/server/rest_api.py index 6c7e0b4e7..f04310427 100644 --- a/mirix/server/rest_api.py +++ b/mirix/server/rest_api.py @@ -17,7 +17,7 @@ from fastapi import APIRouter, Body, FastAPI, Header, HTTPException, Query, Request from fastapi.middleware.cors import CORSMiddleware from fastapi.responses import JSONResponse -from pydantic import BaseModel, Field +from pydantic import BaseModel, Field, field_validator, model_validator from mirix.helpers.message_helpers import prepare_input_message_create from mirix.llm_api.llm_client import LLMClient @@ -37,10 +37,9 @@ from mirix.schemas.file import FileMetadata from mirix.schemas.llm_config import LLMConfig from mirix.schemas.memory import ArchivalMemorySummary, Memory, RecallMemorySummary -from mirix.schemas.message import Message, MessageCreate +from mirix.schemas.message import Message, MessageCreate, _validate_session_id from mirix.schemas.mirix_response import MirixResponse from mirix.schemas.organization import Organization -from mirix.schemas.procedural_memory import ProceduralMemoryItemUpdate from mirix.schemas.raw_memory import ( RawMemoryItem, RawMemoryItemCreateRequest, @@ -72,16 +71,35 @@ from mirix.queue.queue_util import put_messages from mirix.server.constants import MAX_MEMORY_LIMIT # Initialize server (single instance shared across all requests) -_server: Optional[AsyncServer] = None +# The singleton accessor lives in mirix.server.server so service-layer modules +# never need to import this REST module; re-exported here for existing callers. +from mirix.server.server import get_server # noqa: E402 -def get_server() -> AsyncServer: - """Get or create the singleton AsyncServer instance.""" - global _server - if _server is None: - logger.info("Creating AsyncServer instance") - _server = AsyncServer() - return _server +def _isoformat_or_none(value) -> Optional[str]: + if value is None: + return None + return value.isoformat() if hasattr(value, "isoformat") else str(value) + + +def _procedural_memory_response(item, *, include_user_id: bool = False) -> Dict[str, Any]: + """Serialize procedural memory as the public skill-shaped read model.""" + response = { + "memory_type": "procedural", + "id": item.id, + "entry_type": item.entry_type, + "name": item.name, + "description": item.description, + "instructions": item.instructions, + "triggers": getattr(item, "triggers", None) or [], + "examples": getattr(item, "examples", None) or [], + "version": getattr(item, "version", None), + "created_at": _isoformat_or_none(getattr(item, "created_at", None)), + "updated_at": _isoformat_or_none(getattr(item, "updated_at", None)), + } + if include_user_id: + response["user_id"] = str(getattr(item, "user_id", "")) + return response async def initialize(): @@ -125,7 +143,9 @@ async def initialize(): logger.info("Initializing LangFuse observability...") await initialize_langfuse() except Exception as e: - logger.warning(f"LangFuse initialization failed: {e}. Continuing without observability.") + logger.warning( + f"LangFuse initialization failed: {e}. Continuing without observability." + ) async def cleanup(): @@ -199,7 +219,9 @@ async def lifespan(app: FastAPI): from contextvars import ContextVar # Stores the current request for access by decorators -_current_request: ContextVar[Optional[Request]] = ContextVar("current_request", default=None) +_current_request: ContextVar[Optional[Request]] = ContextVar( + "current_request", default=None +) def get_current_request() -> Optional[Request]: @@ -349,7 +371,9 @@ async def inject_client_org_headers(request: Request, call_next): headers.append((b"x-org-id", org_id.encode())) request.scope["headers"] = headers except HTTPException as exc: - return JSONResponse(status_code=exc.status_code, content={"detail": exc.detail}) + return JSONResponse( + status_code=exc.status_code, content={"detail": exc.detail} + ) return await call_next(request) @@ -415,7 +439,12 @@ async def extract_topics_and_temporal_info( # Convert from OpenAI format to internal format new_messages = [] for msg in messages: - prefix = "[USER]" if msg["role"] == "user" else "[ASSISTANT]" + role = msg["role"] + prefix = ( + "[USER]" + if role == "user" + else "[TOOL]" if role in ("tool", "function") else "[ASSISTANT]" + ) new_messages.extend([{"type": "text", "text": prefix + " " + part} for part in msg["content"]]) messages = new_messages @@ -500,15 +529,21 @@ async def extract_topics_and_temporal_info( and len(choice.message.tool_calls) > 0 ): try: - function_args = json.loads(choice.message.tool_calls[0].function.arguments) + function_args = json.loads( + choice.message.tool_calls[0].function.arguments + ) topics = function_args.get("topic") temporal_expr = function_args.get("temporal_expression", "") # Clean up empty strings temporal_expr = temporal_expr.strip() if temporal_expr else None - logger.debug("Extracted topics: %s, temporal: %s", topics, temporal_expr) + logger.debug( + "Extracted topics: %s, temporal: %s", topics, temporal_expr + ) return topics, temporal_expr except (json.JSONDecodeError, KeyError) as parse_error: - logger.warning("Failed to parse extraction response: %s", parse_error) + logger.warning( + "Failed to parse extraction response: %s", parse_error + ) continue except Exception as e: @@ -623,13 +658,19 @@ async def extract_topics_with_local_model( response.raise_for_status() response_data = response.json() except httpx.HTTPStatusError as exc: - logger.error("Failed to extract topics with local model %s: %s", model_name, exc) + logger.error( + "Failed to extract topics with local model %s: %s", model_name, exc + ) return None except httpx.RequestError as exc: - logger.error("Failed to extract topics with local model %s: %s", model_name, exc) + logger.error( + "Failed to extract topics with local model %s: %s", model_name, exc + ) return None - message_payload = response_data.get("message") if isinstance(response_data, dict) else None + message_payload = ( + response_data.get("message") if isinstance(response_data, dict) else None + ) text_response: Optional[str] = None if isinstance(message_payload, dict): text_response = message_payload.get("content") @@ -673,8 +714,23 @@ async def global_exception_handler(request: Request, exc: Exception): @router.get("/health") async def health_check(): - """Health check endpoint.""" - return {"status": "healthy", "service": "mirix-api"} + """Health check endpoint. + + Also reports the effective skill-trigger config so drivers can fail fast + when the in-band procedural trigger is unexpectedly enabled. The added + fields are additive and backward-compatible. + """ + from mirix.constants import ( + MESSAGE_RETAIN_LAST_N_SESSIONS, + SKILL_TRIGGER_SESSION_THRESHOLD, + ) + + return { + "status": "healthy", + "service": "mirix-api", + "skill_trigger_session_threshold": SKILL_TRIGGER_SESSION_THRESHOLD, + "message_retain_last_n_sessions": MESSAGE_RETAIN_LAST_N_SESSIONS, + } # ============================================================================ @@ -770,13 +826,9 @@ async def create_agent( if request.name: create_params["name"] = request.name - agent_state = await server.create_agent( - CreateAgent(**create_params), client - ) + agent_state = await server.create_agent(CreateAgent(**create_params), client) - return await server.agent_manager.get_agent_by_id( - agent_state.id, client - ) + return await server.agent_manager.get_agent_by_id(agent_state.id, client) @router.get("/agents/{agent_id}", response_model=AgentState) @@ -793,11 +845,11 @@ async def get_agent( client = await server.client_manager.get_client_by_id(client_id) try: - return await server.agent_manager.get_agent_by_id( - agent_id, client - ) + return await server.agent_manager.get_agent_by_id(agent_id, client) except NoResultFound as e: - raise HTTPException(status_code=404, detail=f"Agent {agent_id} not found or not accessible") + raise HTTPException( + status_code=404, detail=f"Agent {agent_id} not found or not accessible" + ) @router.delete("/agents/{agent_id}") @@ -892,9 +944,7 @@ async def update_agent_system_prompt_by_name( client = await server.client_manager.get_client_by_id(client_id) # List all top-level agents for this client - top_level_agents = await server.agent_manager.list_agents( - actor=client, limit=1000 - ) + top_level_agents = await server.agent_manager.list_agents(actor=client, limit=1000) # Also get sub-agents (children of meta agent) all_agents = list(top_level_agents) for agent in top_level_agents: @@ -923,7 +973,9 @@ async def update_agent_system_prompt_by_name( # e.g., "meta_memory_agent_episodic_memory_agent" โ†’ "episodic" if agent.name and "meta_memory_agent_" in agent.name: short_name = ( - agent.name.replace("meta_memory_agent_", "").replace("_memory_agent", "").replace("_agent", "") + agent.name.replace("meta_memory_agent_", "") + .replace("_memory_agent", "") + .replace("_agent", "") ) if short_name == agent_name: matching_agent = agent @@ -937,7 +989,9 @@ async def update_agent_system_prompt_by_name( full_name = agent.name if "meta_memory_agent_" in full_name and full_name != "meta_memory_agent": short_name = ( - full_name.replace("meta_memory_agent_", "").replace("_memory_agent", "").replace("_agent", "") + full_name.replace("meta_memory_agent_", "") + .replace("_memory_agent", "") + .replace("_agent", "") ) available_agents.append(f"'{short_name}' (full: {full_name})") else: @@ -947,11 +1001,15 @@ async def update_agent_system_prompt_by_name( if len(available_agents) > 5: available_list += f", and {len(available_agents) - 5} more" - error_detail = f"Agent with name '{agent_name}' not found for client {client_id}. " + error_detail = ( + f"Agent with name '{agent_name}' not found for client {client_id}. " + ) if available_agents: error_detail += f"Available agents: {available_list}" else: - error_detail += "No agents found for this client. Please initialize agents first." + error_detail += ( + "No agents found for this client. Please initialize agents first." + ) raise HTTPException(status_code=404, detail=error_detail) @@ -1020,6 +1078,7 @@ class SendMessageRequest(BaseModel): role: str user_id: Optional[str] = None # End-user ID for message attribution name: Optional[str] = None + session_id: Optional[str] = None # Top-level session identifier stream_steps: bool = False stream_tokens: bool = False filter_tags: Optional[Dict[str, Any]] = None # Filter tags support @@ -1029,6 +1088,11 @@ class SendMessageRequest(BaseModel): block_filter_tags_update_mode: Optional[str] = "merge" # "merge" or "replace" use_cache: bool = True # Control Redis cache behavior + @field_validator("session_id") + @classmethod + def _check_session_id(cls, v: Optional[str]) -> Optional[str]: + return _validate_session_id(v) + @app.post("/agents/{agent_id}/messages", response_model=MirixResponse) async def send_message_to_agent( @@ -1059,12 +1123,19 @@ async def send_message_to_agent( client_id, org_id = await get_client_and_org(x_client_id, x_org_id) client = await server.client_manager.get_client_by_id(client_id) - if request.block_filter_tags is not None and not isinstance(request.block_filter_tags, dict): - raise HTTPException(status_code=400, detail="block_filter_tags must be a dict when provided") + if request.block_filter_tags is not None and not isinstance( + request.block_filter_tags, dict + ): + raise HTTPException( + status_code=400, detail="block_filter_tags must be a dict when provided" + ) if request.block_filter_tags is not None: request.block_filter_tags.pop("scope", None) if request.block_filter_tags_update_mode not in ("merge", "replace"): - raise HTTPException(status_code=400, detail="block_filter_tags_update_mode must be 'merge' or 'replace'") + raise HTTPException( + status_code=400, + detail="block_filter_tags_update_mode must be 'merge' or 'replace'", + ) try: # Prepare the message @@ -1072,6 +1143,7 @@ async def send_message_to_agent( role=MessageRole(request.role), content=request.message, name=request.name, + session_id=request.session_id, ) # Put message on queue for processing @@ -1210,7 +1282,9 @@ async def create_block( server = get_server() client_id, org_id = await get_client_and_org(x_client_id, x_org_id) client = await server.client_manager.get_client_by_id(client_id) - return await server.block_manager.create_or_update_block(block, actor=client, user=user) + return await server.block_manager.create_or_update_block( + block, actor=client, user=user + ) @router.delete("/blocks/{block_id}") @@ -1389,7 +1463,9 @@ async def create_or_get_user( server = get_server() # Accept JWT or injected headers (from API key middleware) - client, auth_type = await get_client_from_jwt_or_api_key(authorization, http_request) + client, auth_type = await get_client_from_jwt_or_api_key( + authorization, http_request + ) if not client: raise HTTPException(status_code=404, detail="Client not found") @@ -1455,15 +1531,24 @@ async def delete_user(user_id: str): error_msg = str(e) # Provide a better error message if user not found or already deleted if "not found" in error_msg.lower() or "no result" in error_msg.lower(): - raise HTTPException(status_code=404, detail=f"User {user_id} not found or already deleted") + raise HTTPException( + status_code=404, detail=f"User {user_id} not found or already deleted" + ) raise HTTPException(status_code=500, detail=error_msg) @router.delete("/users/{user_id}/memories") -async def delete_user_memories(user_id: str): +async def delete_user_memories( + user_id: str, + authorization: Optional[str] = Header(None), + http_request: Request = None, +): """ Hard delete all memories, messages, and blocks for a user. + **Accepts both JWT (dashboard) and Client API Key (programmatic).** + The target user must belong to the caller's organization. + This permanently removes data records while preserving the user record. Use this for data cleanup/purging without affecting the user account itself. @@ -1475,14 +1560,32 @@ async def delete_user_memories(user_id: str): - Knowledge vault items for this user - Messages for this user - Blocks for this user + - Conversation transcripts recorded for this user + - Skill experiences distilled from this user's sessions Records that are PRESERVED: - User record Warning: This operation is irreversible. Deleted data cannot be recovered. """ + client, auth_type = await get_client_from_jwt_or_api_key( + authorization, http_request + ) server = get_server() + # Tenant guard: an irreversible cross-org erasure must be impossible. A + # foreign-org target returns the same 404 as a missing user so the endpoint + # can't be used as a user-id existence oracle. + default_org = server.organization_manager.DEFAULT_ORG_ID + try: + target_user = await server.user_manager.get_user_by_id(user_id) + except Exception: + raise HTTPException(status_code=404, detail=f"User {user_id} not found") + caller_org = client.organization_id or default_org + target_org = target_user.organization_id or default_org + if caller_org != target_org: + raise HTTPException(status_code=404, detail=f"User {user_id} not found") + try: await server.user_manager.delete_memories_by_user_id(user_id) return { @@ -1579,7 +1682,9 @@ async def create_or_get_client( ) else: logger.debug("Client already exists: %s", client_id) - return JSONResponse(status_code=200, content=client.model_dump(mode="json")) + return JSONResponse( + status_code=200, content=client.model_dump(mode="json") + ) except Exception as e: if fail_if_exists and "already exists" in str(e): raise @@ -1697,10 +1802,17 @@ async def delete_client(client_id: str): @router.delete("/clients/{client_id}/memories") -async def delete_client_memories(client_id: str): +async def delete_client_memories( + client_id: str, + authorization: Optional[str] = Header(None), + http_request: Request = None, +): """ Hard delete all memories, messages, and blocks for a client. + **Accepts both JWT (dashboard) and Client API Key (programmatic).** + The target client must belong to the caller's organization. + This permanently removes data records while preserving the client configuration. Use this for data cleanup/purging without affecting the client, agents, or tools. @@ -1712,6 +1824,8 @@ async def delete_client_memories(client_id: str): - Knowledge vault items for this client - Messages for this client - Blocks created by this client + - Conversation transcripts recorded by this client + - Skill experiences created by this client Records that are PRESERVED: - Client record @@ -1720,8 +1834,25 @@ async def delete_client_memories(client_id: str): Warning: This operation is irreversible. Deleted data cannot be recovered. """ + client, auth_type = await get_client_from_jwt_or_api_key( + authorization, http_request + ) server = get_server() + # Tenant guard: same-org only; a missing OR foreign-org target returns the + # same 404 so the endpoint can't be used to probe client-id existence. + # get_client_by_id RAISES (NoResultFound) on a miss rather than returning + # None, so the miss must be caught here, not compared to None. + default_org = server.organization_manager.DEFAULT_ORG_ID + try: + target_client = await server.client_manager.get_client_by_id(client_id) + except Exception: + raise HTTPException(status_code=404, detail=f"Client {client_id} not found") + if (target_client.organization_id or default_org) != ( + client.organization_id or default_org + ): + raise HTTPException(status_code=404, detail=f"Client {client_id} not found") + try: await server.client_manager.delete_memories_by_client_id(client_id) return { @@ -1741,8 +1872,12 @@ class CreateApiKeyRequest(BaseModel): """Request model for creating an API key.""" name: Optional[str] = Field(None, description="Optional name/label for the API key") - permission: Optional[str] = Field("all", description="Permission level: all, restricted, read_only") - user_id: Optional[str] = Field(None, description="User ID this API key is associated with") + permission: Optional[str] = Field( + "all", description="Permission level: all, restricted, read_only" + ) + user_id: Optional[str] = Field( + None, description="User ID this API key is associated with" + ) class CreateApiKeyResponse(BaseModel): @@ -1947,7 +2082,9 @@ async def initialize_meta_agent( actor=client, limit=1000 ) - assert len(existing_meta_agents) <= 1, "Only one meta agent can be created per client" + assert len(existing_meta_agents) <= 1, ( + "Only one meta agent can be created per client" + ) if len(existing_meta_agents) == 1: meta_agent = existing_meta_agents[0] @@ -1957,13 +2094,17 @@ async def initialize_meta_agent( from mirix.schemas.agent import UpdateMetaAgent # DEBUG: Log what we're passing to update_meta_agent - logger.debug("[INIT META AGENT] create_params for UpdateMetaAgent: %s", create_params) + logger.debug( + "[INIT META AGENT] create_params for UpdateMetaAgent: %s", create_params + ) logger.debug( "[INIT META AGENT] 'agents' in create_params: %s", "agents" in create_params, ) if "agents" in create_params: - logger.debug("[INIT META AGENT] agents list: %s", create_params["agents"]) + logger.debug( + "[INIT META AGENT] agents list: %s", create_params["agents"] + ) # Update the existing meta agent meta_agent = await server.agent_manager.update_meta_agent( @@ -1990,12 +2131,194 @@ class AddMemoryRequest(BaseModel): chaining: bool = True verbose: bool = False filter_tags: Optional[Dict[str, Any]] = None + session_id: Optional[str] = ( + None # Batch-level session id applied to every message in this request + ) block_filter_tags: Optional[Dict[str, Any]] = ( None # Applied only when blocks are created (e.g. from default template) ) block_filter_tags_update_mode: Optional[str] = "merge" # "merge" or "replace" use_cache: bool = True # Control Redis cache behavior - occurred_at: Optional[str] = None # Optional ISO 8601 timestamp string for episodic memory + occurred_at: Optional[str] = ( + None # Optional ISO 8601 timestamp string for episodic memory + ) + + @field_validator("session_id") + @classmethod + def _check_session_id(cls, v: Optional[str]) -> Optional[str]: + return _validate_session_id(v) + + @model_validator(mode="after") + def _check_session_id_agrees_with_filter_tags(self) -> "AddMemoryRequest": + # Codex review I1: if both the top-level session_id and filter_tags.session_id are + # provided, they must match. Otherwise message-scope and memory-scope silently diverge. + if self.session_id is not None and isinstance(self.filter_tags, dict): + tag_sid = self.filter_tags.get("session_id") + if tag_sid is not None and tag_sid != self.session_id: + raise ValueError( + "session_id and filter_tags['session_id'] must agree; " + f"got {self.session_id!r} vs {tag_sid!r}" + ) + return self + + +def _serialize_tool_calls(tool_calls: Any) -> str: + """Render an assistant message's tool_calls into a compact text form. + + Accepts the OpenAI shape (`[{"function": {"name", "arguments"}, ...}]`) as + well as simpler `[{"name", "arguments"}]` dicts; anything unrecognized falls + back to `str()`. The distiller only needs to SEE which tool was called with + what arguments โ€” a readable line beats a lossless JSON blob. + """ + if not isinstance(tool_calls, list): + return str(tool_calls) + lines = [] + for call in tool_calls: + if isinstance(call, dict): + function = call.get("function") if isinstance(call.get("function"), dict) else call + name = function.get("name") or call.get("name") or "unknown_tool" + arguments = function.get("arguments", call.get("arguments", "")) + if not isinstance(arguments, str): + import json as _json + + try: + arguments = _json.dumps(arguments, ensure_ascii=False) + except (TypeError, ValueError): + arguments = str(arguments) + lines.append(f"[tool_call] {name}({arguments})") + else: + lines.append(f"[tool_call] {call}") + return "\n".join(lines) + + +def _extract_conversation_turns(messages: List[Dict[str, Any]]) -> List[Dict[str, str]]: + """Extract clean per-role turns for the Conversation Message Store. + + Reads the REAL roles from the raw add-memory payload (the same shape the + role-collapse branch consumes) and preserves them as + 'user'/'assistant'/'tool' โ€” NOT the [USER]/[ASSISTANT] role-collapsed blob + the meta agent receives. The store is the single source of truth for skill + distillation, so it keeps the true turn structure. Tool activity is kept + deliberately: work-process lessons (a tool error, a retry, the fix that + worked) live in tool calls and tool results, and dropping them would blind + the distiller to exactly that signal. + + Only role-bearing payloads (`[{"role": ..., "content": ...}, ...]`) yield + turns; anything else (e.g. screenshot/content-only payloads) yields [] and + nothing is written. `content` may be a string or a list of string parts; a + list is joined with newlines to match how the collapse branch treats parts. + + Mapping: + - role 'user'/'assistant' โ†’ kept as-is; an assistant message that ALSO + carries `tool_calls` gets them serialized and appended to its content + (an assistant turn that is pure tool_calls with empty content still + yields a turn). + - role 'tool'/'function' โ†’ stored as a 'tool' turn (tool results). A + `name`/`tool_name` field, when present, is prefixed for readability. + - any other role is dropped โ€” not a learnable conversation turn. + """ + if not ( + isinstance(messages, list) + and messages + and isinstance(messages[0], dict) + and "role" in messages[0] + ): + return [] + + from mirix.schemas.conversation_message import ( + CONVERSATION_MESSAGE_MAX_CONTENT_LEN, + ) + + # Truncate at the STORE's per-row cap so one oversized turn (typically a + # huge tool result) fails softly instead of failing the whole batch's + # Pydantic validation in record_turns โ€” which would drop every turn of + # the request. Leave room for the truncation marker. + _marker = " โ€ฆ[truncated]" + _content_cap = CONVERSATION_MESSAGE_MAX_CONTENT_LEN - len(_marker) + + def _cap(text: str) -> str: + if len(text) > CONVERSATION_MESSAGE_MAX_CONTENT_LEN: + return text[:_content_cap] + _marker + return text + + def _part_text(part: Any) -> str: + # Standard SDK shape: {"type": "text", "text": "..."} โ€” store the text, + # not the dict repr. Anything else falls back to str(). + if isinstance(part, dict) and isinstance(part.get("text"), str): + return part["text"] + return str(part) + + turns: List[Dict[str, str]] = [] + for msg in messages: + role = msg.get("role") + content = msg.get("content", "") + if isinstance(content, list): + content = "\n".join(_part_text(part) for part in content) + elif content is None: + content = "" + elif not isinstance(content, str): + content = str(content) + + if role in ("user", "assistant"): + if role == "assistant" and msg.get("tool_calls"): + serialized = _serialize_tool_calls(msg["tool_calls"]) + content = f"{content}\n{serialized}".strip() if content else serialized + turns.append({"role": role, "content": _cap(content)}) + elif role in ("tool", "function"): + tool_name = msg.get("name") or msg.get("tool_name") + if tool_name: + content = f"[{tool_name}] {content}" + turns.append({"role": "tool", "content": _cap(content)}) + return turns + + +async def _ingest_session_turns(request, input_messages, client, user_id: str) -> None: + """Shared /memory/add + /memory/add_sync session ingestion seam. + + Stamps the batch-level session_id onto every message that didn't carry its + own, then persists the external turns (REAL user/assistant roles) into the + Conversation Message Store โ€” the single source of truth for procedural + (skill) distillation. Without a session_id nothing is written, so no + procedural memory is produced; the other five components still extract via + the unchanged meta dispatch. + + The store write is ADDITIVE: isolated in its own try/except so a store + failure (DB error, validation, etc.) is logged but never aborts the primary + memory ingestion. Worst case is a missed session for skill distillation, + not a dropped memory-add. + """ + if request.session_id is None: + return + + for msg_create in input_messages: + if msg_create.session_id is None: + msg_create.session_id = request.session_id + + from mirix.services.conversation_message_manager import ( + ConversationMessageManager, + owner_org, + ) + + # Extraction runs INSIDE the guard too: a malformed message shape (odd + # tool_calls, exotic content) must degrade to a missed session, never + # abort the primary memory ingestion. + try: + conversation_turns = _extract_conversation_turns(request.messages) + if not conversation_turns: + return + await ConversationMessageManager().record_turns( + session_id=request.session_id, + user_id=user_id, + organization_id=owner_org(client), + turns=conversation_turns, + actor=client, + ) + except Exception: + logger.exception( + "Failed to record conversation turns to the store for " + "session_id=%s; memory ingestion will still proceed", + request.session_id, + ) @router.post("/memory/add") @@ -2054,7 +2377,12 @@ async def add_memory( # We need to convert the message to the format in "content" new_message = [] for msg in message: - prefix = "[USER]" if msg["role"] == "user" else "[ASSISTANT]" + role = msg["role"] + prefix = ( + "[USER]" + if role == "user" + else "[TOOL]" if role in ("tool", "function") else "[ASSISTANT]" + ) # Handle both string and list content content = msg["content"] @@ -2070,6 +2398,10 @@ async def add_memory( input_messages = convert_message_to_mirix_message(message) + # Session ingestion (session_id stamping + Conversation Message Store write, + # independent of the meta dispatch below) โ€” see _ingest_session_turns. + await _ingest_session_turns(request, input_messages, client, user_id) + # Add client scope to filter_tags (create if not provided) if request.filter_tags is not None: # Create a copy to avoid modifying the original request @@ -2078,17 +2410,31 @@ async def add_memory( # Create new filter_tags if not provided filter_tags = {} - if request.block_filter_tags is not None and not isinstance(request.block_filter_tags, dict): - raise HTTPException(status_code=400, detail="block_filter_tags must be a dict when provided") + # Mirror session_id into filter_tags so extracted memories inherit it too. + # The model_validator on AddMemoryRequest has already ensured agreement if both were set. + if request.session_id is not None: + filter_tags["session_id"] = request.session_id + + if request.block_filter_tags is not None and not isinstance( + request.block_filter_tags, dict + ): + raise HTTPException( + status_code=400, detail="block_filter_tags must be a dict when provided" + ) if request.block_filter_tags is not None: request.block_filter_tags.pop("scope", None) if request.block_filter_tags_update_mode not in ("merge", "replace"): - raise HTTPException(status_code=400, detail="block_filter_tags_update_mode must be 'merge' or 'replace'") + raise HTTPException( + status_code=400, + detail="block_filter_tags_update_mode must be 'merge' or 'replace'", + ) # Add or update the "scope" key with the client's write_scope for memory creation # Memories are written with the client's write_scope if client.write_scope is None: - raise HTTPException(status_code=403, detail="Client has no write_scope - cannot create memories") + raise HTTPException( + status_code=403, detail="Client has no write_scope - cannot create memories" + ) filter_tags["scope"] = client.write_scope # Queue for async processing instead of synchronous execution @@ -2158,7 +2504,12 @@ async def add_memory_sync( if isinstance(message, list) and "role" in message[0].keys(): new_message = [] for msg in message: - prefix = "[USER]" if msg["role"] == "user" else "[ASSISTANT]" + role = msg["role"] + prefix = ( + "[USER]" + if role == "user" + else "[TOOL]" if role in ("tool", "function") else "[ASSISTANT]" + ) content = msg["content"] if isinstance(content, str): new_message.append({"type": "text", "text": prefix + " " + content}) @@ -2170,13 +2521,23 @@ async def add_memory_sync( input_messages = convert_message_to_mirix_message(message) + # Session ingestion (session_id stamping + Conversation Message Store write) โ€” + # shared with /memory/add; see _ingest_session_turns. + await _ingest_session_turns(request, input_messages, client, user_id) + if request.filter_tags is not None: filter_tags = dict(request.filter_tags) else: filter_tags = {} + # The AddMemoryRequest model_validator already ensured agreement if both were set. + if request.session_id is not None: + filter_tags["session_id"] = request.session_id + if client.write_scope is None: - raise HTTPException(status_code=403, detail="Client has no write_scope - cannot create memories") + raise HTTPException( + status_code=403, detail="Client has no write_scope - cannot create memories" + ) filter_tags["scope"] = client.write_scope from mirix.services.user_manager import UserManager @@ -2230,12 +2591,20 @@ class RetrieveMemoryRequest(BaseModel): user_id: Optional[str] = None # Optional - uses admin user if not provided messages: List[Dict[str, Any]] limit: int = 10 # Maximum number of items to retrieve per memory type - local_model_for_retrieval: Optional[str] = None # Optional local Ollama model for topic extraction - filter_tags: Optional[Dict[str, Any]] = None # Optional filter tags for filtering results + local_model_for_retrieval: Optional[str] = ( + None # Optional local Ollama model for topic extraction + ) + filter_tags: Optional[Dict[str, Any]] = ( + None # Optional filter tags for filtering results + ) use_cache: bool = True # Control Redis cache behavior # NEW: Optional date range for temporal filtering (ISO 8601 format) - start_date: Optional[str] = None # e.g., "2025-11-19T00:00:00" or "2025-11-19T00:00:00+00:00" - end_date: Optional[str] = None # e.g., "2025-11-19T23:59:59" or "2025-11-19T23:59:59+00:00" + start_date: Optional[str] = ( + None # e.g., "2025-11-19T00:00:00" or "2025-11-19T00:00:00+00:00" + ) + end_date: Optional[str] = ( + None # e.g., "2025-11-19T23:59:59" or "2025-11-19T23:59:59+00:00" + ) async def retrieve_memories_by_keywords( @@ -2326,7 +2695,9 @@ async def retrieve_memories_by_keywords( "recent": [ { "id": event.id, - "timestamp": (event.occurred_at.isoformat() if event.occurred_at else None), + "timestamp": ( + event.occurred_at.isoformat() if event.occurred_at else None + ), "summary": event.summary, "details": event.details, } @@ -2335,7 +2706,9 @@ async def retrieve_memories_by_keywords( "relevant": [ { "id": event.id, - "timestamp": (event.occurred_at.isoformat() if event.occurred_at else None), + "timestamp": ( + event.occurred_at.isoformat() if event.occurred_at else None + ), "summary": event.summary, "details": event.details, } @@ -2420,7 +2793,7 @@ async def retrieve_memories_by_keywords( agent_state=agent_state, # Not accessed during BM25 search user=user, query=key_words, - search_field="summary", + search_field="description", search_method=search_method, limit=limit, timezone_str=timezone_str, @@ -2430,15 +2803,10 @@ async def retrieve_memories_by_keywords( ) memories["procedural"] = { - "total_count": await procedural_manager.get_total_number_of_items(user=user), - "items": [ - { - "id": procedure.id, - "entry_type": procedure.entry_type, - "summary": procedure.summary, - } - for procedure in procedures - ], + "total_count": await procedural_manager.get_total_number_of_items( + user=user + ), + "items": [_procedural_memory_response(procedure) for procedure in procedures], } except Exception as e: logger.error("Error retrieving procedural memories: %s", e) @@ -2459,7 +2827,9 @@ async def retrieve_memories_by_keywords( ) memories["knowledge_vault"] = { - "total_count": await knowledge_vault_manager.get_total_number_of_items(user=user), + "total_count": await knowledge_vault_manager.get_total_number_of_items( + user=user + ), "items": [ { "id": item.id, @@ -2534,9 +2904,7 @@ async def retrieve_memory_with_conversation( filter_tags = dict(request.filter_tags) if request.filter_tags is not None else {} # Get all agents for this client (automatically filtered by client via apply_access_predicate) - all_agents = await server.agent_manager.list_agents( - actor=client, limit=1000 - ) + all_agents = await server.agent_manager.list_agents(actor=client, limit=1000) if not all_agents: return { @@ -2555,7 +2923,10 @@ async def retrieve_memory_with_conversation( for msg in request.messages: if isinstance(msg, dict) and "content" in msg: for content_item in msg.get("content", []): - if isinstance(content_item, dict) and content_item.get("text", "").strip(): + if ( + isinstance(content_item, dict) + and content_item.get("text", "").strip() + ): has_content = True break if has_content: @@ -2600,9 +2971,13 @@ async def retrieve_memory_with_conversation( # Use explicit date range from request try: if request.start_date: - start_date = datetime.fromisoformat(request.start_date.replace("Z", "+00:00")) + start_date = datetime.fromisoformat( + request.start_date.replace("Z", "+00:00") + ) if request.end_date: - end_date = datetime.fromisoformat(request.end_date.replace("Z", "+00:00")) + end_date = datetime.fromisoformat( + request.end_date.replace("Z", "+00:00") + ) logger.debug("Using explicit date range: %s to %s", start_date, end_date) except ValueError as e: logger.warning("Invalid date format in request: %s", e) @@ -2716,9 +3091,7 @@ async def retrieve_memory_with_topic( parsed_filter_tags = {} # Get all agents for this client (automatically filtered by client via apply_access_predicate) - all_agents = await server.agent_manager.list_agents( - actor=client, limit=1000 - ) + all_agents = await server.agent_manager.list_agents(actor=client, limit=1000) if not all_agents: return { @@ -2773,7 +3146,9 @@ async def _precompute_embedding_for_search( from mirix.embeddings import embedding_model # Compute embedding once - embedded_text = await (await embedding_model(agent_state.embedding_config)).get_text_embedding(query) + embedded_text = await ( + await embedding_model(agent_state.embedding_config) + ).get_text_embedding(query) # Pad for episodic memory which requires MAX_EMBEDDING_DIM embedded_text_padded = np.pad( @@ -2792,14 +3167,16 @@ async def search_memory( query: str = "", memory_type: str = "all", search_field: str = "null", - search_method: str = "embedding", + search_method: str = "", # "" = per-memory-type default (see resolution below) limit: int = 10, authorization: Optional[str] = Header(None), filter_tags: Optional[str] = Query(None), similarity_threshold: Optional[float] = Query(None), start_date: Optional[str] = Query(None), end_date: Optional[str] = Query(None), - include_core_memory: bool = Query(False, description="When True, include core (block) memory in the response."), + include_core_memory: bool = Query( + False, description="When True, include core (block) memory in the response." + ), x_client_id: Optional[str] = Header(None), x_org_id: Optional[str] = Header(None), ): @@ -2815,11 +3192,15 @@ async def search_memory( search_field: Field to search in. Options vary by memory type: - episodic: "summary", "details" - resource: "summary", "content" - - procedural: "summary", "steps" + - procedural: "description", "instructions" - knowledge_vault: "caption", "secret_value" - semantic: "name", "summary", "details" - For "all": use "null" (default) - search_method: Search method. Options: "bm25" (default), "embedding" + search_method: Search method. Options: "bm25", "embedding", and (procedural + only) "hybrid" (BM25 + embedding fused via Reciprocal Rank Fusion). + When omitted, defaults per memory type: procedural -> "hybrid" + (override via MIRIX_SKILL_SEARCH_METHOD), all others -> "embedding". + "hybrid" requested for a non-procedural type falls back to embedding. limit: Maximum number of results per memory type (default: 10) filter_tags: Optional JSON string of filter tags (scope added automatically) similarity_threshold: Optional similarity threshold for embedding search (0.0-2.0). @@ -2859,9 +3240,7 @@ async def search_memory( logger.debug("No user_id provided, using admin user: %s", user_id) # Get all agents for this client (automatically filtered by client via apply_access_predicate) - all_agents = await server.agent_manager.list_agents( - actor=client, limit=1000 - ) + all_agents = await server.agent_manager.list_agents(actor=client, limit=1000) if not all_agents: return { @@ -2905,7 +3284,9 @@ async def search_memory( if start_date: try: - parsed_start_date = datetime.fromisoformat(start_date.replace("Z", "+00:00")) + parsed_start_date = datetime.fromisoformat( + start_date.replace("Z", "+00:00") + ) # Strip timezone for DB comparison (DB stores naive datetimes) if parsed_start_date.tzinfo: parsed_start_date = parsed_start_date.replace(tzinfo=None) @@ -2921,12 +3302,31 @@ async def search_memory( except ValueError as e: logger.warning("Invalid end_date format: %s", e) - # Normalize empty search_method to the default (FastAPI passes "" for missing query params) + # Resolve the search method per memory type. FastAPI passes "" when the + # caller omits the query param, so "" is the "use the default" sentinel: + # * procedural -> EverOS-aligned hybrid (BM25 + embedding fused via RRF), + # env-overridable via MIRIX_SKILL_SEARCH_METHOD; + # * every other type (incl. the cross-type "all" sweep) -> embedding. + from mirix.constants import PROCEDURAL_DEFAULT_SEARCH_METHOD + if not search_method: + search_method = ( + PROCEDURAL_DEFAULT_SEARCH_METHOD + if memory_type == "procedural" + else "embedding" + ) + # "hybrid" is procedural-only โ€” the other managers have no hybrid branch. + # For any non-procedural memory_type (including "all") fall back to + # embedding so those managers receive a method they support. + if search_method == "hybrid" and memory_type != "procedural": search_method = "embedding" # Validate search parameters - if memory_type == "resource" and search_field == "content" and search_method == "embedding": + if ( + memory_type == "resource" + and search_field == "content" + and search_method == "embedding" + ): return { "success": False, "error": "embedding is not supported for resource memory's 'content' field.", @@ -2935,7 +3335,11 @@ async def search_memory( "count": 0, } - if memory_type == "knowledge_vault" and search_field == "secret_value" and search_method == "embedding": + if ( + memory_type == "knowledge_vault" + and search_field == "secret_value" + and search_method == "embedding" + ): return { "success": False, "error": "embedding is not supported for knowledge_vault memory's 'secret_value' field.", @@ -2948,10 +3352,15 @@ async def search_memory( search_field = "null" # Pre-compute embedding once if using embedding search (to avoid redundant embeddings) - embedded_text, embedded_text_padded = await _precompute_embedding_for_search(search_method, query, agent_state) + embedded_text, embedded_text_padded = await _precompute_embedding_for_search( + search_method, query, agent_state + ) # Collect results from requested memory types all_results = [] + # For a single procedural search, surface the user's total procedural count + # alongside the page-sized result set. + procedural_total_count = None # If searching all memory types, run searches concurrently for better performance if memory_type == "all": @@ -2964,7 +3373,11 @@ async def search_episodic(): agent_state=agent_state, user=user, query=query, - embedded_text=(embedded_text_padded if search_method == "embedding" and query else None), + embedded_text=( + embedded_text_padded + if search_method == "embedding" and query + else None + ), search_field=search_field if search_field != "null" else "details", search_method=search_method, limit=limit, @@ -2979,7 +3392,9 @@ async def search_episodic(): { "memory_type": "episodic", "id": x.id, - "occurred_at": (x.occurred_at.isoformat() if x.occurred_at else None), + "occurred_at": ( + x.occurred_at.isoformat() if x.occurred_at else None + ), "event_type": x.event_type, "actor": x.actor, "summary": x.summary, @@ -2997,7 +3412,11 @@ async def search_resource(): agent_state=agent_state, user=user, query=query, - embedded_text=(embedded_text if search_method == "embedding" and query else None), + embedded_text=( + embedded_text + if search_method == "embedding" and query + else None + ), search_field=( search_field if search_field != "null" @@ -3031,8 +3450,14 @@ async def search_procedural(): agent_state=agent_state, user=user, query=query, - embedded_text=(embedded_text if search_method == "embedding" and query else None), - search_field=search_field if search_field != "null" else "steps", + embedded_text=( + embedded_text + if search_method == "embedding" and query + else None + ), + search_field=search_field + if search_field != "null" + else "description", search_method=search_method, limit=limit, timezone_str=timezone_str, @@ -3040,16 +3465,7 @@ async def search_procedural(): scopes=scopes, similarity_threshold=similarity_threshold, ) - return [ - { - "memory_type": "procedural", - "id": x.id, - "entry_type": x.entry_type, - "summary": x.summary, - "steps": x.steps, - } - for x in memories - ] + return [_procedural_memory_response(x) for x in memories] except Exception as e: logger.error("Error searching procedural memories: %s", e) return [] @@ -3060,7 +3476,11 @@ async def search_knowledge(): agent_state=agent_state, user=user, query=query, - embedded_text=(embedded_text if search_method == "embedding" and query else None), + embedded_text=( + embedded_text + if search_method == "embedding" and query + else None + ), search_field=search_field if search_field != "null" else "caption", search_method=search_method, limit=limit, @@ -3091,7 +3511,11 @@ async def search_semantic(): agent_state=agent_state, user=user, query=query, - embedded_text=(embedded_text_padded if search_method == "embedding" and query else None), + embedded_text=( + embedded_text_padded + if search_method == "embedding" and query + else None + ), search_field=search_field if search_field != "null" else "details", search_method=search_method, limit=limit, @@ -3133,11 +3557,19 @@ async def search_core(): for block in blocks ] except Exception as e: - logger.error("Error retrieving core memory blocks: %s", e, exc_info=True) + logger.error( + "Error retrieving core memory blocks: %s", e, exc_info=True + ) return [] # Run all searches concurrently - tasks = [search_episodic(), search_resource(), search_procedural(), search_knowledge(), search_semantic()] + tasks = [ + search_episodic(), + search_resource(), + search_procedural(), + search_knowledge(), + search_semantic(), + ] if include_core_memory: tasks.append(search_core()) results = await asyncio.gather(*tasks) @@ -3149,27 +3581,35 @@ async def search_core(): # Single memory type searches (run serially) elif memory_type == "episodic": try: - episodic_memories = await server.episodic_memory_manager.list_episodic_memory( - agent_state=agent_state, - user=user, - query=query, - embedded_text=(embedded_text_padded if search_method == "embedding" and query else None), - search_field=search_field if search_field != "null" else "summary", - search_method=search_method, - limit=limit, - timezone_str=timezone_str, - filter_tags=parsed_filter_tags, - scopes=scopes, - start_date=parsed_start_date, - end_date=parsed_end_date, - similarity_threshold=similarity_threshold, + episodic_memories = ( + await server.episodic_memory_manager.list_episodic_memory( + agent_state=agent_state, + user=user, + query=query, + embedded_text=( + embedded_text_padded + if search_method == "embedding" and query + else None + ), + search_field=search_field if search_field != "null" else "summary", + search_method=search_method, + limit=limit, + timezone_str=timezone_str, + filter_tags=parsed_filter_tags, + scopes=scopes, + start_date=parsed_start_date, + end_date=parsed_end_date, + similarity_threshold=similarity_threshold, + ) ) all_results.extend( [ { "memory_type": "episodic", "id": x.id, - "timestamp": (x.occurred_at.isoformat() if x.occurred_at else None), + "timestamp": ( + x.occurred_at.isoformat() if x.occurred_at else None + ), "event_type": x.event_type, "actor": x.actor, "summary": x.summary, @@ -3188,7 +3628,9 @@ async def search_core(): agent_state=agent_state, user=user, query=query, - embedded_text=(embedded_text if search_method == "embedding" and query else None), + embedded_text=( + embedded_text if search_method == "embedding" and query else None + ), search_field=( search_field if search_field != "null" @@ -3209,7 +3651,9 @@ async def search_core(): "resource_type": x.resource_type, "title": x.title, "summary": x.summary, - "content": (x.content[:200] if x.content else None), # Truncate content for response + "content": ( + x.content[:200] if x.content else None + ), # Truncate content for response } for x in resource_memories ] @@ -3220,30 +3664,36 @@ async def search_core(): # Search procedural memories elif memory_type == "procedural": try: - procedural_memories = await server.procedural_memory_manager.list_procedures( - agent_state=agent_state, - user=user, - query=query, - embedded_text=(embedded_text if search_method == "embedding" and query else None), - search_field=search_field if search_field != "null" else "summary", - search_method=search_method, - limit=limit, - timezone_str=timezone_str, - filter_tags=parsed_filter_tags, - scopes=scopes, - similarity_threshold=similarity_threshold, + procedural_memories = ( + await server.procedural_memory_manager.list_procedures( + agent_state=agent_state, + user=user, + query=query, + embedded_text=( + embedded_text + if search_method == "embedding" and query + else None + ), + search_field=search_field + if search_field != "null" + else "description", + search_method=search_method, + limit=limit, + timezone_str=timezone_str, + filter_tags=parsed_filter_tags, + scopes=scopes, + similarity_threshold=similarity_threshold, + ) ) all_results.extend( - [ - { - "memory_type": "procedural", - "id": x.id, - "entry_type": x.entry_type, - "summary": x.summary, - "steps": x.steps, - } - for x in procedural_memories - ] + [_procedural_memory_response(x) for x in procedural_memories] + ) + # Report the user's total procedural count; the page's `count` is + # only len(results), not the global total. + procedural_total_count = ( + await server.procedural_memory_manager.get_total_number_of_items( + user=user + ) ) except Exception as e: logger.error("Error searching procedural memories: %s", e) @@ -3251,18 +3701,24 @@ async def search_core(): # Search knowledge vault elif memory_type == "knowledge_vault": try: - knowledge_vault_memories = await server.knowledge_vault_manager.list_knowledge( - agent_state=agent_state, - user=user, - query=query, - embedded_text=(embedded_text if search_method == "embedding" and query else None), - search_field=search_field if search_field != "null" else "caption", - search_method=search_method, - limit=limit, - timezone_str=timezone_str, - filter_tags=parsed_filter_tags, - scopes=scopes, - similarity_threshold=similarity_threshold, + knowledge_vault_memories = ( + await server.knowledge_vault_manager.list_knowledge( + agent_state=agent_state, + user=user, + query=query, + embedded_text=( + embedded_text + if search_method == "embedding" and query + else None + ), + search_field=search_field if search_field != "null" else "caption", + search_method=search_method, + limit=limit, + timezone_str=timezone_str, + filter_tags=parsed_filter_tags, + scopes=scopes, + similarity_threshold=similarity_threshold, + ) ) all_results.extend( [ @@ -3284,18 +3740,24 @@ async def search_core(): # Search semantic memories elif memory_type == "semantic": try: - semantic_memories = await server.semantic_memory_manager.list_semantic_items( - agent_state=agent_state, - user=user, - query=query, - embedded_text=(embedded_text_padded if search_method == "embedding" and query else None), - search_field=search_field if search_field != "null" else "summary", - search_method=search_method, - limit=limit, - timezone_str=timezone_str, - filter_tags=parsed_filter_tags, - scopes=scopes, - similarity_threshold=similarity_threshold, + semantic_memories = ( + await server.semantic_memory_manager.list_semantic_items( + agent_state=agent_state, + user=user, + query=query, + embedded_text=( + embedded_text_padded + if search_method == "embedding" and query + else None + ), + search_field=search_field if search_field != "null" else "summary", + search_method=search_method, + limit=limit, + timezone_str=timezone_str, + filter_tags=parsed_filter_tags, + scopes=scopes, + similarity_threshold=similarity_threshold, + ) ) all_results.extend( [ @@ -3334,9 +3796,13 @@ async def search_core(): } ) except Exception as e: - logger.error("Error retrieving core memory blocks for single-user search: %s", e, exc_info=True) + logger.error( + "Error retrieving core memory blocks for single-user search: %s", + e, + exc_info=True, + ) - return { + response = { "success": True, "query": query, "memory_type": memory_type, @@ -3353,6 +3819,11 @@ async def search_core(): "results": all_results, "count": len(all_results), } + # A single procedural search also reports the global per-user procedural + # total alongside the page count. + if procedural_total_count is not None: + response["total_count"] = procedural_total_count + return response @router.get("/memory/search_all_users") @@ -3361,13 +3832,14 @@ async def search_memory_all_users( query: str, memory_type: str = "all", search_field: str = "null", - search_method: str = "embedding", + search_method: str = "", # "" = per-memory-type default (see resolution below) limit: int = 10, client_id: Optional[str] = Query(None), org_id: Optional[str] = Query(None), filter_tags: Optional[str] = Query(None), include_core_memory: bool = Query( - False, description="When True, include a 'core' section with block memory in the response." + False, + description="When True, include a 'core' section with block memory in the response.", ), block_filter_tags: Optional[str] = Query( None, @@ -3421,7 +3893,9 @@ async def search_memory_all_users( ) else: # Fall back to headers - effective_client_id, header_org_id = await get_client_and_org(x_client_id, x_org_id) + effective_client_id, header_org_id = await get_client_and_org( + x_client_id, x_org_id + ) client = await server.client_manager.get_client_by_id(effective_client_id) # Use org_id from query param if provided, otherwise use header org_id effective_org_id = org_id or header_org_id @@ -3436,7 +3910,9 @@ async def search_memory_all_users( try: filter_tags_dict = json.loads(filter_tags) except json.JSONDecodeError: - raise HTTPException(status_code=400, detail="Invalid filter_tags JSON format") + raise HTTPException( + status_code=400, detail="Invalid filter_tags JSON format" + ) else: filter_tags_dict = {} @@ -3459,7 +3935,9 @@ async def search_memory_all_users( if start_date: try: - parsed_start_date = datetime.fromisoformat(start_date.replace("Z", "+00:00")) + parsed_start_date = datetime.fromisoformat( + start_date.replace("Z", "+00:00") + ) # Strip timezone for DB comparison (DB stores naive datetimes) if parsed_start_date.tzinfo: parsed_start_date = parsed_start_date.replace(tzinfo=None) @@ -3476,9 +3954,7 @@ async def search_memory_all_users( logger.warning("Invalid end_date format: %s", e) # Get agents for this client - all_agents = await server.agent_manager.list_agents( - actor=client, limit=1000 - ) + all_agents = await server.agent_manager.list_agents(actor=client, limit=1000) if not all_agents: return { "success": False, @@ -3490,12 +3966,31 @@ async def search_memory_all_users( agent_state = all_agents[0] - # Normalize empty search_method to the default (FastAPI passes "" for missing query params) + # Resolve the search method per memory type. FastAPI passes "" when the + # caller omits the query param, so "" is the "use the default" sentinel: + # * procedural -> EverOS-aligned hybrid (BM25 + embedding fused via RRF), + # env-overridable via MIRIX_SKILL_SEARCH_METHOD; + # * every other type (incl. the cross-type "all" sweep) -> embedding. + from mirix.constants import PROCEDURAL_DEFAULT_SEARCH_METHOD + if not search_method: + search_method = ( + PROCEDURAL_DEFAULT_SEARCH_METHOD + if memory_type == "procedural" + else "embedding" + ) + # "hybrid" is procedural-only โ€” the other managers have no hybrid branch. + # For any non-procedural memory_type (including "all") fall back to + # embedding so those managers receive a method they support. + if search_method == "hybrid" and memory_type != "procedural": search_method = "embedding" # Validate search parameters - if memory_type == "resource" and search_field == "content" and search_method == "embedding": + if ( + memory_type == "resource" + and search_field == "content" + and search_method == "embedding" + ): return { "success": False, "error": "embedding is not supported for resource memory's 'content' field.", @@ -3504,7 +3999,11 @@ async def search_memory_all_users( "count": 0, } - if memory_type == "knowledge_vault" and search_field == "secret_value" and search_method == "embedding": + if ( + memory_type == "knowledge_vault" + and search_field == "secret_value" + and search_method == "embedding" + ): return { "success": False, "error": "embedding is not supported for knowledge_vault memory's 'secret_value' field.", @@ -3517,7 +4016,9 @@ async def search_memory_all_users( search_field = "null" # Pre-compute embedding once if using embedding search (to avoid redundant embeddings) - embedded_text, embedded_text_padded = await _precompute_embedding_for_search(search_method, query, agent_state) + embedded_text, embedded_text_padded = await _precompute_embedding_for_search( + search_method, query, agent_state + ) # Collect results using organization_id filter all_results = [] @@ -3529,26 +4030,36 @@ async def search_memory_all_users( # Define async wrappers for each manager call async def search_episodic(): try: - memories = await server.episodic_memory_manager.list_episodic_memory_by_org( - agent_state=agent_state, - organization_id=effective_org_id, - query=query, - embedded_text=(embedded_text_padded if search_method == "embedding" and query else None), - search_field=search_field if search_field != "null" else "summary", - search_method=search_method, - limit=limit, - timezone_str="UTC", - filter_tags=filter_tags_dict, - scopes=scopes, - start_date=parsed_start_date, - end_date=parsed_end_date, - similarity_threshold=similarity_threshold, + memories = ( + await server.episodic_memory_manager.list_episodic_memory_by_org( + agent_state=agent_state, + organization_id=effective_org_id, + query=query, + embedded_text=( + embedded_text_padded + if search_method == "embedding" and query + else None + ), + search_field=search_field + if search_field != "null" + else "summary", + search_method=search_method, + limit=limit, + timezone_str="UTC", + filter_tags=filter_tags_dict, + scopes=scopes, + start_date=parsed_start_date, + end_date=parsed_end_date, + similarity_threshold=similarity_threshold, + ) ) return [ { "memory_type": "episodic", "id": x.id, - "timestamp": (x.occurred_at.isoformat() if x.occurred_at else None), + "timestamp": ( + x.occurred_at.isoformat() if x.occurred_at else None + ), "event_type": x.event_type, "actor": x.actor, "summary": x.summary, @@ -3567,7 +4078,11 @@ async def search_resource(): agent_state=agent_state, organization_id=effective_org_id, query=query, - embedded_text=(embedded_text if search_method == "embedding" and query else None), + embedded_text=( + embedded_text + if search_method == "embedding" and query + else None + ), search_field=( search_field if search_field != "null" @@ -3598,28 +4113,29 @@ async def search_resource(): async def search_procedural(): try: - memories = await server.procedural_memory_manager.list_procedures_by_org( - agent_state=agent_state, - organization_id=effective_org_id, - query=query, - embedded_text=(embedded_text if search_method == "embedding" and query else None), - search_field=search_field if search_field != "null" else "summary", - search_method=search_method, - limit=limit, - timezone_str="UTC", - filter_tags=filter_tags_dict, - scopes=scopes, - similarity_threshold=similarity_threshold, + memories = ( + await server.procedural_memory_manager.list_procedures_by_org( + agent_state=agent_state, + organization_id=effective_org_id, + query=query, + embedded_text=( + embedded_text + if search_method == "embedding" and query + else None + ), + search_field=search_field + if search_field != "null" + else "description", + search_method=search_method, + limit=limit, + timezone_str="UTC", + filter_tags=filter_tags_dict, + scopes=scopes, + similarity_threshold=similarity_threshold, + ) ) return [ - { - "memory_type": "procedural", - "id": x.id, - "entry_type": x.entry_type, - "summary": x.summary, - "steps": x.steps, - "user_id": str(x.user_id), - } + _procedural_memory_response(x, include_user_id=True) for x in memories ] except Exception as e: @@ -3632,7 +4148,11 @@ async def search_knowledge(): agent_state=agent_state, organization_id=effective_org_id, query=query, - embedded_text=(embedded_text if search_method == "embedding" and query else None), + embedded_text=( + embedded_text + if search_method == "embedding" and query + else None + ), search_field=search_field if search_field != "null" else "caption", search_method=search_method, limit=limit, @@ -3660,18 +4180,26 @@ async def search_knowledge(): async def search_semantic(): try: - memories = await server.semantic_memory_manager.list_semantic_items_by_org( - agent_state=agent_state, - organization_id=effective_org_id, - query=query, - embedded_text=(embedded_text_padded if search_method == "embedding" and query else None), - search_field=search_field if search_field != "null" else "summary", - search_method=search_method, - limit=limit, - timezone_str="UTC", - filter_tags=filter_tags_dict, - scopes=scopes, - similarity_threshold=similarity_threshold, + memories = ( + await server.semantic_memory_manager.list_semantic_items_by_org( + agent_state=agent_state, + organization_id=effective_org_id, + query=query, + embedded_text=( + embedded_text_padded + if search_method == "embedding" and query + else None + ), + search_field=search_field + if search_field != "null" + else "summary", + search_method=search_method, + limit=limit, + timezone_str="UTC", + filter_tags=filter_tags_dict, + scopes=scopes, + similarity_threshold=similarity_threshold, + ) ) return [ { @@ -3703,20 +4231,26 @@ async def search_semantic(): # Single memory type searches (run serially as before) elif memory_type == "episodic": try: - episodic_memories = await server.episodic_memory_manager.list_episodic_memory_by_org( - agent_state=agent_state, - organization_id=effective_org_id, - query=query, - embedded_text=(embedded_text_padded if search_method == "embedding" and query else None), - search_field=search_field if search_field != "null" else "summary", - search_method=search_method, - limit=limit, - timezone_str="UTC", - filter_tags=filter_tags_dict, - scopes=scopes, - start_date=parsed_start_date, - end_date=parsed_end_date, - similarity_threshold=similarity_threshold, + episodic_memories = ( + await server.episodic_memory_manager.list_episodic_memory_by_org( + agent_state=agent_state, + organization_id=effective_org_id, + query=query, + embedded_text=( + embedded_text_padded + if search_method == "embedding" and query + else None + ), + search_field=search_field if search_field != "null" else "summary", + search_method=search_method, + limit=limit, + timezone_str="UTC", + filter_tags=filter_tags_dict, + scopes=scopes, + start_date=parsed_start_date, + end_date=parsed_end_date, + similarity_threshold=similarity_threshold, + ) ) all_results.extend( [ @@ -3724,7 +4258,9 @@ async def search_semantic(): "memory_type": "episodic", "user_id": x.user_id, "id": x.id, - "timestamp": (x.occurred_at.isoformat() if x.occurred_at else None), + "timestamp": ( + x.occurred_at.isoformat() if x.occurred_at else None + ), "event_type": x.event_type, "actor": x.actor, "summary": x.summary, @@ -3739,22 +4275,28 @@ async def search_semantic(): # Search resource memories across organization elif memory_type == "resource": try: - resource_memories = await server.resource_memory_manager.list_resources_by_org( - agent_state=agent_state, - organization_id=effective_org_id, - query=query, - embedded_text=(embedded_text if search_method == "embedding" and query else None), - search_field=( - search_field - if search_field != "null" - else ("summary" if search_method == "embedding" else "content") - ), - search_method=search_method, - limit=limit, - timezone_str="UTC", - filter_tags=filter_tags_dict, - scopes=scopes, - similarity_threshold=similarity_threshold, + resource_memories = ( + await server.resource_memory_manager.list_resources_by_org( + agent_state=agent_state, + organization_id=effective_org_id, + query=query, + embedded_text=( + embedded_text + if search_method == "embedding" and query + else None + ), + search_field=( + search_field + if search_field != "null" + else ("summary" if search_method == "embedding" else "content") + ), + search_method=search_method, + limit=limit, + timezone_str="UTC", + filter_tags=filter_tags_dict, + scopes=scopes, + similarity_threshold=similarity_threshold, + ) ) all_results.extend( [ @@ -3776,50 +4318,59 @@ async def search_semantic(): # Search procedural memories across organization elif memory_type == "procedural": try: - procedural_memories = await server.procedural_memory_manager.list_procedures_by_org( - agent_state=agent_state, - organization_id=effective_org_id, - query=query, - embedded_text=(embedded_text if search_method == "embedding" and query else None), - search_field=search_field if search_field != "null" else "summary", - search_method=search_method, - limit=limit, - timezone_str="UTC", - filter_tags=filter_tags_dict, - scopes=scopes, - similarity_threshold=similarity_threshold, + procedural_memories = ( + await server.procedural_memory_manager.list_procedures_by_org( + agent_state=agent_state, + organization_id=effective_org_id, + query=query, + embedded_text=( + embedded_text + if search_method == "embedding" and query + else None + ), + search_field=search_field + if search_field != "null" + else "description", + search_method=search_method, + limit=limit, + timezone_str="UTC", + filter_tags=filter_tags_dict, + scopes=scopes, + similarity_threshold=similarity_threshold, + ) ) all_results.extend( [ - { - "memory_type": "procedural", - "user_id": x.user_id, - "id": x.id, - "entry_type": x.entry_type, - "summary": x.summary, - "steps": x.steps, - } + _procedural_memory_response(x, include_user_id=True) for x in procedural_memories ] ) except Exception as e: - logger.error("Error searching procedural memories across organization: %s", e) + logger.error( + "Error searching procedural memories across organization: %s", e + ) # Search knowledge vault across organization elif memory_type == "knowledge_vault": try: - knowledge_vault_memories = await server.knowledge_vault_manager.list_knowledge_by_org( - agent_state=agent_state, - organization_id=effective_org_id, - query=query, - embedded_text=(embedded_text if search_method == "embedding" and query else None), - search_field=search_field if search_field != "null" else "caption", - search_method=search_method, - limit=limit, - timezone_str="UTC", - filter_tags=filter_tags_dict, - scopes=scopes, - similarity_threshold=similarity_threshold, + knowledge_vault_memories = ( + await server.knowledge_vault_manager.list_knowledge_by_org( + agent_state=agent_state, + organization_id=effective_org_id, + query=query, + embedded_text=( + embedded_text + if search_method == "embedding" and query + else None + ), + search_field=search_field if search_field != "null" else "caption", + search_method=search_method, + limit=limit, + timezone_str="UTC", + filter_tags=filter_tags_dict, + scopes=scopes, + similarity_threshold=similarity_threshold, + ) ) all_results.extend( [ @@ -3842,18 +4393,24 @@ async def search_semantic(): # Search semantic memories across organization elif memory_type == "semantic": try: - semantic_memories = await server.semantic_memory_manager.list_semantic_items_by_org( - agent_state=agent_state, - organization_id=effective_org_id, - query=query, - embedded_text=(embedded_text_padded if search_method == "embedding" and query else None), - search_field=search_field if search_field != "null" else "summary", - search_method=search_method, - limit=limit, - timezone_str="UTC", - filter_tags=filter_tags_dict, - scopes=scopes, - similarity_threshold=similarity_threshold, + semantic_memories = ( + await server.semantic_memory_manager.list_semantic_items_by_org( + agent_state=agent_state, + organization_id=effective_org_id, + query=query, + embedded_text=( + embedded_text_padded + if search_method == "embedding" and query + else None + ), + search_field=search_field if search_field != "null" else "summary", + search_method=search_method, + limit=limit, + timezone_str="UTC", + filter_tags=filter_tags_dict, + scopes=scopes, + similarity_threshold=similarity_threshold, + ) ) all_results.extend( [ @@ -3880,7 +4437,9 @@ async def search_semantic(): try: block_filter_tags_parsed = json.loads(block_filter_tags) except json.JSONDecodeError: - raise HTTPException(status_code=400, detail="Invalid block_filter_tags JSON format") + raise HTTPException( + status_code=400, detail="Invalid block_filter_tags JSON format" + ) blocks = await server.block_manager.get_blocks( user=None, organization_id=effective_org_id, @@ -3910,7 +4469,11 @@ async def search_semantic(): except HTTPException: raise except Exception as e: - logger.error("Error retrieving core memory blocks for cross-user search: %s", e, exc_info=True) + logger.error( + "Error retrieving core memory blocks for cross-user search: %s", + e, + exc_info=True, + ) return { "success": True, @@ -3970,7 +4533,9 @@ async def list_memory_components( ) # Authenticate (JWT or API key) - client, auth_type = await get_client_from_jwt_or_api_key(authorization, http_request) + client, auth_type = await get_client_from_jwt_or_api_key( + authorization, http_request + ) server = get_server() # Default to the admin user for this client @@ -3988,9 +4553,7 @@ async def list_memory_components( limit = max(1, min(limit if limit is not None else MAX_MEMORY_LIMIT, MAX_MEMORY_LIMIT)) # Need an agent state for memory manager configuration - agents = await server.agent_manager.list_agents( - actor=client, limit=1 - ) + agents = await server.agent_manager.list_agents(actor=client, limit=1) if not agents: raise HTTPException(status_code=404, detail="No agents found for this client") agent_state = agents[0] @@ -4006,17 +4569,29 @@ async def list_memory_components( timezone_str=timezone_str, ) memories["episodic"] = { - "total_count": await server.episodic_memory_manager.get_total_number_of_items(user=user), + "total_count": await server.episodic_memory_manager.get_total_number_of_items( + user=user + ), "items": [ { "id": item.id, - "occurred_at": (item.occurred_at.isoformat() if item.occurred_at else None), + "occurred_at": ( + item.occurred_at.isoformat() if item.occurred_at else None + ), "event_type": item.event_type, "actor": item.actor, "summary": item.summary, "details": item.details, - "created_at": (item.created_at.isoformat() if getattr(item, "created_at", None) else None), - "updated_at": (item.updated_at.isoformat() if getattr(item, "updated_at", None) else None), + "created_at": ( + item.created_at.isoformat() + if getattr(item, "created_at", None) + else None + ), + "updated_at": ( + item.updated_at.isoformat() + if getattr(item, "updated_at", None) + else None + ), } for item in episodic_items ], @@ -4033,7 +4608,9 @@ async def list_memory_components( timezone_str=timezone_str, ) memories["semantic"] = { - "total_count": await server.semantic_memory_manager.get_total_number_of_items(user=user), + "total_count": await server.semantic_memory_manager.get_total_number_of_items( + user=user + ), "items": [ { "id": item.id, @@ -4041,8 +4618,16 @@ async def list_memory_components( "summary": item.summary, "details": item.details, "source": item.source, - "created_at": (item.created_at.isoformat() if getattr(item, "created_at", None) else None), - "updated_at": (item.updated_at.isoformat() if getattr(item, "updated_at", None) else None), + "created_at": ( + item.created_at.isoformat() + if getattr(item, "created_at", None) + else None + ), + "updated_at": ( + item.updated_at.isoformat() + if getattr(item, "updated_at", None) + else None + ), } for item in semantic_items ], @@ -4053,22 +4638,17 @@ async def list_memory_components( agent_state=agent_state, user=user, query="", - search_field="summary", + search_field="description", search_method="bm25", limit=limit, timezone_str=timezone_str, ) memories["procedural"] = { - "total_count": await server.procedural_memory_manager.get_total_number_of_items(user=user), + "total_count": await server.procedural_memory_manager.get_total_number_of_items( + user=user + ), "items": [ - { - "id": item.id, - "entry_type": item.entry_type, - "summary": item.summary, - "steps": item.steps, - "created_at": (item.created_at.isoformat() if getattr(item, "created_at", None) else None), - "updated_at": (item.updated_at.isoformat() if getattr(item, "updated_at", None) else None), - } + _procedural_memory_response(item) for item in procedural_items ], } @@ -4084,7 +4664,9 @@ async def list_memory_components( timezone_str=timezone_str, ) memories["resource"] = { - "total_count": await server.resource_memory_manager.get_total_number_of_items(user=user), + "total_count": await server.resource_memory_manager.get_total_number_of_items( + user=user + ), "items": [ { "id": item.id, @@ -4092,8 +4674,16 @@ async def list_memory_components( "title": item.title, "summary": item.summary, "content": item.content, - "created_at": (item.created_at.isoformat() if getattr(item, "created_at", None) else None), - "updated_at": (item.updated_at.isoformat() if getattr(item, "updated_at", None) else None), + "created_at": ( + item.created_at.isoformat() + if getattr(item, "created_at", None) + else None + ), + "updated_at": ( + item.updated_at.isoformat() + if getattr(item, "updated_at", None) + else None + ), } for item in resource_items ], @@ -4110,7 +4700,9 @@ async def list_memory_components( timezone_str=timezone_str, ) memories["knowledge_vault"] = { - "total_count": await server.knowledge_vault_manager.get_total_number_of_items(user=user), + "total_count": await server.knowledge_vault_manager.get_total_number_of_items( + user=user + ), "items": [ { "id": item.id, @@ -4119,8 +4711,16 @@ async def list_memory_components( "sensitivity": item.sensitivity, "secret_value": item.secret_value, "caption": item.caption, - "created_at": (item.created_at.isoformat() if getattr(item, "created_at", None) else None), - "updated_at": (item.updated_at.isoformat() if getattr(item, "updated_at", None) else None), + "created_at": ( + item.created_at.isoformat() + if getattr(item, "created_at", None) + else None + ), + "updated_at": ( + item.updated_at.isoformat() + if getattr(item, "updated_at", None) + else None + ), } for item in knowledge_items ], @@ -4178,7 +4778,7 @@ async def list_memory_fields( fields_by_type = { "episodic": ["summary", "details"], "semantic": ["name", "summary", "details"], - "procedural": ["summary", "steps"], + "procedural": ["description", "instructions"], "resource": ["summary", "content"], "knowledge_vault": ["caption", "secret_value"], "core": ["label", "value"], @@ -4224,7 +4824,9 @@ async def get_client_from_jwt_or_api_key( client_id = admin_payload["sub"] client = await server.client_manager.get_client_by_id(client_id) if not client: - raise HTTPException(status_code=404, detail=f"Client {client_id} not found") + raise HTTPException( + status_code=404, detail=f"Client {client_id} not found" + ) return client, "jwt" except HTTPException: pass # Try API key next @@ -4237,7 +4839,9 @@ async def get_client_from_jwt_or_api_key( client_id, org_id = await get_client_and_org(client_id, org_id) client = await server.client_manager.get_client_by_id(client_id) if not client: - raise HTTPException(status_code=404, detail=f"Client {client_id} not found") + raise HTTPException( + status_code=404, detail=f"Client {client_id} not found" + ) return client, "api_key" raise HTTPException( @@ -4269,7 +4873,9 @@ async def update_episodic_memory( Updates the summary and/or details fields of the memory. """ # Authenticate with either JWT or API key - client, auth_type = await get_client_from_jwt_or_api_key(authorization, http_request) + client, auth_type = await get_client_from_jwt_or_api_key( + authorization, http_request + ) server = get_server() @@ -4317,7 +4923,9 @@ async def delete_episodic_memory( **Accepts both JWT (dashboard) and Client API Key (programmatic).** """ - client, auth_type = await get_client_from_jwt_or_api_key(authorization, http_request) + client, auth_type = await get_client_from_jwt_or_api_key( + authorization, http_request + ) server = get_server() @@ -4350,7 +4958,9 @@ async def update_semantic_memory( **Accepts both JWT (dashboard) and Client API Key (programmatic).** """ # Authenticate with either JWT or API key - client, auth_type = await get_client_from_jwt_or_api_key(authorization, http_request) + client, auth_type = await get_client_from_jwt_or_api_key( + authorization, http_request + ) server = get_server() @@ -4405,75 +5015,17 @@ async def delete_semantic_memory( **Accepts both JWT (dashboard) and Client API Key (programmatic).** """ - client, auth_type = await get_client_from_jwt_or_api_key(authorization, http_request) - - server = get_server() - - try: - await server.semantic_memory_manager.delete_semantic_item_by_id(memory_id, actor=client) - return {"success": True, "message": f"Semantic memory {memory_id} deleted"} - except Exception as e: - raise HTTPException(status_code=404, detail=str(e)) - - -class UpdateProceduralMemoryRequest(BaseModel): - """Request model for updating a procedural memory.""" - - summary: Optional[str] = None - steps: Optional[List[str]] = None - - -@router.patch("/memory/procedural/{memory_id}") -async def update_procedural_memory( - memory_id: str, - request: UpdateProceduralMemoryRequest, - user_id: Optional[str] = None, - authorization: Optional[str] = Header(None), - http_request: Request = None, -): - """ - Update a procedural memory by ID. - - **Accepts both JWT (dashboard) and Client API Key (programmatic).** - """ - # Authenticate with either JWT or API key - client, auth_type = await get_client_from_jwt_or_api_key(authorization, http_request) + client, auth_type = await get_client_from_jwt_or_api_key( + authorization, http_request + ) server = get_server() - # If user_id is not provided, use the admin user for this client - if not user_id: - from mirix.services.admin_user_manager import ClientAuthManager - - user_id = ClientAuthManager.get_admin_user_id_for_client(client.id) - logger.debug("No user_id provided, using admin user: %s", user_id) - - # Get user - user = await server.user_manager.get_user_by_id(user_id) - if not user: - raise HTTPException(status_code=404, detail=f"User {user_id} not found") - try: - procedural_update_data = {"id": memory_id} - if request.summary is not None: - procedural_update_data["summary"] = request.summary - if request.steps is not None: - procedural_update_data["steps"] = request.steps - - updated_memory = await server.procedural_memory_manager.update_item( - item_update=ProceduralMemoryItemUpdate.model_validate(procedural_update_data), - user=user, - actor=client, + await server.semantic_memory_manager.delete_semantic_item_by_id( + memory_id, actor=client ) - return { - "success": True, - "message": f"Procedural memory {memory_id} updated", - "memory": { - "id": updated_memory.id, - "summary": updated_memory.summary, - "steps": updated_memory.steps, - }, - } + return {"success": True, "message": f"Semantic memory {memory_id} deleted"} except Exception as e: raise HTTPException(status_code=404, detail=str(e)) @@ -4485,16 +5037,25 @@ async def delete_procedural_memory( http_request: Request = None, ): """ - Delete a procedural memory by ID. + Delete a procedural memory (skill) by ID. **Accepts both JWT (dashboard) and Client API Key (programmatic).** + + Deletion stays on the public surface like every other memory type โ€” a user + must be able to remove a skill that is wrong or unwanted. Skill WRITES, by + contrast, have no public endpoint: skills are created/edited only through + the session-distillation evolution flow (auto_dream mode='procedural'). """ - client, auth_type = await get_client_from_jwt_or_api_key(authorization, http_request) + client, auth_type = await get_client_from_jwt_or_api_key( + authorization, http_request + ) server = get_server() try: - await server.procedural_memory_manager.delete_procedure_by_id(memory_id, actor=client) + await server.procedural_memory_manager.delete_procedure_by_id( + memory_id, actor=client + ) return {"success": True, "message": f"Procedural memory {memory_id} deleted"} except Exception as e: raise HTTPException(status_code=404, detail=str(e)) @@ -4521,7 +5082,9 @@ async def create_raw_memory_handler( request: Create request with context, filter_tags, etc. user_id: User ID this memory belongs to (required query parameter) """ - client, auth_type = await get_client_from_jwt_or_api_key(authorization, http_request) + client, auth_type = await get_client_from_jwt_or_api_key( + authorization, http_request + ) response = await create_raw_memory(request, user_id, client) return response @@ -4550,9 +5113,7 @@ async def create_raw_memory( raise HTTPException(status_code=401, detail="Client or client_id required") # Get agent_state for embedding generation (required) - agents = await server.agent_manager.list_agents( - actor=client, limit=1 - ) + agents = await server.agent_manager.list_agents(actor=client, limit=1) agent_state = agents[0] if agents else None if not agent_state: @@ -4562,7 +5123,9 @@ async def create_raw_memory( ) # Build PydanticRawMemoryItemCreate from request - from mirix.schemas.raw_memory import RawMemoryItemCreate as PydanticRawMemoryItemCreate + from mirix.schemas.raw_memory import ( + RawMemoryItemCreate as PydanticRawMemoryItemCreate, + ) assert client.organization_id is not None raw_memory_create = PydanticRawMemoryItemCreate( @@ -4615,7 +5178,9 @@ async def get_raw_memory_handler( **Accepts both JWT (dashboard) and Client API Key (programmatic).** """ # Authenticate with either JWT or API key - client, auth_type = await get_client_from_jwt_or_api_key(authorization, http_request) + client, auth_type = await get_client_from_jwt_or_api_key( + authorization, http_request + ) response = await get_raw_memory(memory_id, user_id, client) return response @@ -4658,7 +5223,9 @@ async def get_raw_memory( try: from mirix.orm.errors import NoResultFound - memory = await server.raw_memory_manager.get_raw_memory_by_id(memory_id, actor=client, user_id=user_id) + memory = await server.raw_memory_manager.get_raw_memory_by_id( + memory_id, actor=client, user_id=user_id + ) return { "success": True, "memory": memory.model_dump(mode="json"), @@ -4684,7 +5251,9 @@ async def update_raw_memory_handler( **Accepts both JWT (dashboard) and Client API Key (programmatic).** """ # Authenticate with either JWT or API key - client, auth_type = await get_client_from_jwt_or_api_key(authorization, http_request) + client, auth_type = await get_client_from_jwt_or_api_key( + authorization, http_request + ) response = await update_raw_memory(memory_id, request, user_id, client) return response @@ -4727,9 +5296,7 @@ async def update_raw_memory( raise HTTPException(status_code=404, detail=f"User {user_id} not found") # Get agent_state for embedding generation (required) - agents = await server.agent_manager.list_agents( - actor=client, limit=1 - ) + agents = await server.agent_manager.list_agents(actor=client, limit=1) agent_state = agents[0] if agents else None if not agent_state: @@ -4778,7 +5345,9 @@ async def delete_raw_memory_handler( **Accepts both JWT (dashboard) and Client API Key (programmatic).** """ # Authenticate with either JWT or API key - client, auth_type = await get_client_from_jwt_or_api_key(authorization, http_request) + client, auth_type = await get_client_from_jwt_or_api_key( + authorization, http_request + ) response = await delete_raw_memory(memory_id, client) return response @@ -4819,14 +5388,18 @@ async def delete_raw_memory( raise HTTPException(status_code=404, detail=f"User {user_id} not found") try: - deleted = await server.raw_memory_manager.delete_raw_memory(memory_id, client, user_id=user_id) + deleted = await server.raw_memory_manager.delete_raw_memory( + memory_id, client, user_id=user_id + ) if deleted: return { "success": True, "message": f"Raw memory {memory_id} deleted", } else: - raise HTTPException(status_code=404, detail=f"Raw memory {memory_id} not found") + raise HTTPException( + status_code=404, detail=f"Raw memory {memory_id} not found" + ) except Exception as e: logger.error(f"Error deleting raw memory {memory_id}: {e}") raise HTTPException(status_code=500, detail=f"Internal error: {str(e)}") @@ -4845,7 +5418,9 @@ async def search_raw_memory_handler( **Accepts both JWT (dashboard) and Client API Key (programmatic).** """ # Authenticate with either JWT or API key - client, auth_type = await get_client_from_jwt_or_api_key(authorization, http_request) + client, auth_type = await get_client_from_jwt_or_api_key( + authorization, http_request + ) response = await search_raw_memory(request, user_id, client) return response @@ -4882,7 +5457,9 @@ async def search_raw_memory( except NoResultFound: raise HTTPException(status_code=404, detail=f"User {user_id} not found") - filter_tags = dict[str, Any](request.filter_tags) if request.filter_tags is not None else {} + filter_tags = ( + dict[str, Any](request.filter_tags) if request.filter_tags is not None else {} + ) scopes = client.read_scopes # Parse time_range from request, converting to UTC and stripping timezone for DB comparison @@ -4907,27 +5484,51 @@ def to_utc_stripped(dt_val: Optional[datetime]) -> Optional[datetime]: time_range_dict = {} if request.time_range.created_at_gte: - time_range_dict["created_at_gte"] = to_utc_stripped(request.time_range.created_at_gte) + time_range_dict["created_at_gte"] = to_utc_stripped( + request.time_range.created_at_gte + ) if request.time_range.created_at_lte: - time_range_dict["created_at_lte"] = to_utc_stripped(request.time_range.created_at_lte) + time_range_dict["created_at_lte"] = to_utc_stripped( + request.time_range.created_at_lte + ) if request.time_range.occurred_at_gte: - time_range_dict["occurred_at_gte"] = to_utc_stripped(request.time_range.occurred_at_gte) + time_range_dict["occurred_at_gte"] = to_utc_stripped( + request.time_range.occurred_at_gte + ) if request.time_range.occurred_at_lte: - time_range_dict["occurred_at_lte"] = to_utc_stripped(request.time_range.occurred_at_lte) + time_range_dict["occurred_at_lte"] = to_utc_stripped( + request.time_range.occurred_at_lte + ) if request.time_range.updated_at_gte: - time_range_dict["updated_at_gte"] = to_utc_stripped(request.time_range.updated_at_gte) + time_range_dict["updated_at_gte"] = to_utc_stripped( + request.time_range.updated_at_gte + ) if request.time_range.updated_at_lte: - time_range_dict["updated_at_lte"] = to_utc_stripped(request.time_range.updated_at_lte) + time_range_dict["updated_at_lte"] = to_utc_stripped( + request.time_range.updated_at_lte + ) - if time_range_dict.get("created_at_gte") and time_range_dict.get("created_at_lte"): + if time_range_dict.get("created_at_gte") and time_range_dict.get( + "created_at_lte" + ): if time_range_dict["created_at_gte"] > time_range_dict["created_at_lte"]: - raise HTTPException(status_code=400, detail="created_at_gte must be <= created_at_lte") - if time_range_dict.get("occurred_at_gte") and time_range_dict.get("occurred_at_lte"): + raise HTTPException( + status_code=400, detail="created_at_gte must be <= created_at_lte" + ) + if time_range_dict.get("occurred_at_gte") and time_range_dict.get( + "occurred_at_lte" + ): if time_range_dict["occurred_at_gte"] > time_range_dict["occurred_at_lte"]: - raise HTTPException(status_code=400, detail="occurred_at_gte must be <= occurred_at_lte") - if time_range_dict.get("updated_at_gte") and time_range_dict.get("updated_at_lte"): + raise HTTPException( + status_code=400, detail="occurred_at_gte must be <= occurred_at_lte" + ) + if time_range_dict.get("updated_at_gte") and time_range_dict.get( + "updated_at_lte" + ): if time_range_dict["updated_at_gte"] > time_range_dict["updated_at_lte"]: - raise HTTPException(status_code=400, detail="updated_at_gte must be <= updated_at_lte") + raise HTTPException( + status_code=400, detail="updated_at_gte must be <= updated_at_lte" + ) # Validate sort string valid_sorts = { @@ -4946,7 +5547,9 @@ def to_utc_stripped(dt_val: Optional[datetime]) -> Optional[datetime]: # Validate client has organization_id if not client.organization_id: - raise HTTPException(status_code=401, detail="Client must have an organization_id") + raise HTTPException( + status_code=401, detail="Client must have an organization_id" + ) try: items, next_cursor = await server.raw_memory_manager.search_raw_memories( @@ -4974,7 +5577,9 @@ def to_utc_stripped(dt_val: Optional[datetime]) -> Optional[datetime]: @router.post("/memory/raw/cleanup") async def cleanup_raw_memories_handler( - days_threshold: int = Query(14, ge=1, le=365, description="Delete memories older than this many days"), + days_threshold: int = Query( + 14, ge=1, le=365, description="Delete memories older than this many days" + ), authorization: Optional[str] = Header(None), http_request: Request = None, ): @@ -5009,7 +5614,9 @@ async def cleanup_raw_memories_handler( } """ # Authenticate with either JWT or API key - client, auth_type = await get_client_from_jwt_or_api_key(authorization, http_request) + client, auth_type = await get_client_from_jwt_or_api_key( + authorization, http_request + ) logger.info( "Manual cleanup job triggered by client %s (auth: %s) with threshold: %d days", @@ -5023,7 +5630,9 @@ async def cleanup_raw_memories_handler( async def cleanup_raw_memories( - days_threshold: int = Query(14, ge=1, le=365, description="Delete memories older than this many days"), + days_threshold: int = Query( + 14, ge=1, le=365, description="Delete memories older than this many days" + ), client: Optional[Client] = None, client_id: Optional[str] = None, ): @@ -5049,7 +5658,8 @@ async def cleanup_raw_memories( # Add success message to result result["message"] = ( - f"Deleted {result['deleted_count']} stale raw memories " f"(older than {days_threshold} days)" + f"Deleted {result['deleted_count']} stale raw memories " + f"(older than {days_threshold} days)" ) logger.info( @@ -5100,11 +5710,37 @@ async def auto_dream_handler( user_id=ClientAuthManager.get_admin_user_id_for_client(client.id) ) - meta_agents = await server.agent_manager.list_agents(actor=client) - meta_agent_state = next( - (a for a in meta_agents if a.agent_type == AgentType.meta_memory_agent), - None, - ) + if request_body.meta_agent_id: + try: + meta_agent_state = await server.agent_manager.get_agent_by_id( + request_body.meta_agent_id, + client, + ) + except Exception as exc: + raise HTTPException( + status_code=404, + detail=f"Meta agent {request_body.meta_agent_id} not found.", + ) from exc + if meta_agent_state.agent_type != AgentType.meta_memory_agent: + raise HTTPException( + status_code=400, + detail=( + f"Agent {request_body.meta_agent_id} is " + f"{meta_agent_state.agent_type}, not a meta_memory_agent." + ), + ) + else: + meta_agents = await server.agent_manager.list_agents(actor=client) + matching_meta_agents = [ + a for a in meta_agents if a.agent_type == AgentType.meta_memory_agent + ] + meta_agent_state = matching_meta_agents[0] if matching_meta_agents else None + if len(matching_meta_agents) > 1: + logger.warning( + "Auto dream called without meta_agent_id; using first of %d meta agents for client %s", + len(matching_meta_agents), + client.id, + ) if meta_agent_state is None: raise HTTPException( status_code=400, @@ -5146,7 +5782,9 @@ async def update_resource_memory( **Accepts both JWT (dashboard) and Client API Key (programmatic).** """ # Authenticate with either JWT or API key - client, auth_type = await get_client_from_jwt_or_api_key(authorization, http_request) + client, auth_type = await get_client_from_jwt_or_api_key( + authorization, http_request + ) server = get_server() @@ -5200,12 +5838,16 @@ async def delete_resource_memory( **Accepts both JWT (dashboard) and Client API Key (programmatic).** """ - client, auth_type = await get_client_from_jwt_or_api_key(authorization, http_request) + client, auth_type = await get_client_from_jwt_or_api_key( + authorization, http_request + ) server = get_server() try: - await server.resource_memory_manager.delete_resource_by_id(memory_id, actor=client) + await server.resource_memory_manager.delete_resource_by_id( + memory_id, actor=client + ) return {"success": True, "message": f"Resource memory {memory_id} deleted"} except Exception as e: raise HTTPException(status_code=404, detail=str(e)) @@ -5222,12 +5864,16 @@ async def delete_knowledge_vault_memory( **Accepts both JWT (dashboard) and Client API Key (programmatic).** """ - client, auth_type = await get_client_from_jwt_or_api_key(authorization, http_request) + client, auth_type = await get_client_from_jwt_or_api_key( + authorization, http_request + ) server = get_server() try: - await server.knowledge_vault_manager.delete_knowledge_by_id(memory_id, actor=client) + await server.knowledge_vault_manager.delete_knowledge_by_id( + memory_id, actor=client + ) return {"success": True, "message": f"Knowledge vault item {memory_id} deleted"} except Exception as e: raise HTTPException(status_code=404, detail=str(e)) @@ -5411,7 +6057,9 @@ async def dashboard_login(request: DashboardLoginRequest): ) if auth_status == "not_found": - raise HTTPException(status_code=404, detail="Account does not exist. Please create an account.") + raise HTTPException( + status_code=404, detail="Account does not exist. Please create an account." + ) if auth_status == "wrong_password": raise HTTPException(status_code=401, detail="Incorrect password") if auth_status != "ok" or not client or not access_token: @@ -5476,7 +6124,9 @@ async def list_dashboard_clients( # Check scope - only admin can list dashboard clients if client_payload.get("scope") != "admin": - raise HTTPException(status_code=403, detail="Only admin clients can list dashboard clients") + raise HTTPException( + status_code=403, detail="Only admin clients can list dashboard clients" + ) from mirix.services.admin_user_manager import ClientAuthManager @@ -5550,7 +6200,9 @@ async def dashboard_check_setup(): return { "setup_required": is_first, "message": ( - "No dashboard users exist. Please create the first account." if is_first else "Dashboard setup complete." + "No dashboard users exist. Please create the first account." + if is_first + else "Dashboard setup complete." ), } diff --git a/mirix/server/server.py b/mirix/server/server.py index f8de97c3d..28cff5ced 100644 --- a/mirix/server/server.py +++ b/mirix/server/server.py @@ -17,6 +17,7 @@ from rich.panel import Panel from rich.text import Text from sqlalchemy import create_engine +from sqlalchemy import inspect as sqlalchemy_inspect from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine from sqlalchemy.orm import sessionmaker @@ -453,12 +454,72 @@ async def db_context(): await session.close() +# Tables that manual migration scripts ALTER on upgrades. create_all only +# creates missing tables โ€” it never adds columns to existing ones โ€” so a +# skipped migration surfaces as missing columns and would otherwise break +# every SELECT on that entity at request time. +_MIGRATION_HINTS: Dict[str, str] = { + "messages": "scripts/migrate_add_message_session_id.sql (then *_phase2.sql outside a transaction)", + "procedural_memory": "scripts/migrate_procedural_to_skill.sql", + "conversation_message": "scripts/migrate_add_conversation_message.sql (then *_phase2.sql)", + "skill_experience": "scripts/migrate_add_skill_experience.sql", + "agent_trigger_state": "scripts/migrate_add_agent_trigger_state.sql", +} + + +def _find_missing_columns(sync_conn) -> Dict[str, List[str]]: + """Compare ORM-declared columns against the live DB for existing tables.""" + inspector = sqlalchemy_inspect(sync_conn) + existing_tables = set(inspector.get_table_names()) + missing: Dict[str, List[str]] = {} + for table in Base.metadata.sorted_tables: + if table.name not in existing_tables: + continue # create_all just created it with the full schema + db_columns = {column["name"] for column in inspector.get_columns(table.name)} + absent = [column.name for column in table.columns if column.name not in db_columns] + if absent: + missing[table.name] = absent + return missing + + async def ensure_tables_created(): - """Create all tables on the async engine. Call from FastAPI lifespan startup.""" + """Create all tables on the async engine and fail fast on schema drift. + + Call from FastAPI lifespan startup. Raises RuntimeError with the exact + migration scripts to run when an existing database is missing columns the + ORM requires (i.e. the operator upgraded without running migrations). + """ if USE_PGLITE: return async with engine.begin() as conn: await conn.run_sync(Base.metadata.create_all) + missing = await conn.run_sync(_find_missing_columns) + if missing: + lines = [] + for table_name, columns in sorted(missing.items()): + hint = _MIGRATION_HINTS.get(table_name, "see scripts/migrate_*.sql") + lines.append(f" - table '{table_name}' is missing columns {columns}; run {hint}") + raise RuntimeError( + "Database schema is out of date for this MIRIX version. " + "Startup aborted so requests don't fail mid-flight with UndefinedColumn errors.\n" + "Run the pending migration scripts against this database, then restart:\n" + + "\n".join(lines) + ) + + +# Singleton AsyncServer accessor. Lives here (not in rest_api) so service-layer +# modules that need the assembled server never import the REST module โ€” the +# dependency arrow stays services -> server core, REST -> both. +_server_singleton: Optional["AsyncServer"] = None + + +def get_server() -> "AsyncServer": + """Get or create the singleton AsyncServer instance.""" + global _server_singleton + if _server_singleton is None: + logger.info("Creating AsyncServer instance") + _server_singleton = AsyncServer() + return _server_singleton async def sse_async_generator(generator, usage_task=None, finish_message=True): diff --git a/mirix/services/agent_manager.py b/mirix/services/agent_manager.py index 51787fa57..18c1124d0 100644 --- a/mirix/services/agent_manager.py +++ b/mirix/services/agent_manager.py @@ -15,7 +15,7 @@ KNOWLEDGE_VAULT_TOOLS, MCP_TOOLS, META_MEMORY_TOOLS, - PROCEDURAL_MEMORY_TOOLS, + SKILL_TOOLS, RESOURCE_MEMORY_TOOLS, SEARCH_MEMORY_TOOLS, SEMANTIC_MEMORY_TOOLS, @@ -125,7 +125,7 @@ async def create_agent( if agent_create.agent_type == AgentType.episodic_memory_agent: tool_names.extend(EPISODIC_MEMORY_TOOLS + UNIVERSAL_MEMORY_TOOLS) if agent_create.agent_type == AgentType.procedural_memory_agent: - tool_names.extend(PROCEDURAL_MEMORY_TOOLS + UNIVERSAL_MEMORY_TOOLS) + tool_names.extend(SKILL_TOOLS + UNIVERSAL_MEMORY_TOOLS) if agent_create.agent_type == AgentType.resource_memory_agent: tool_names.extend(RESOURCE_MEMORY_TOOLS + UNIVERSAL_MEMORY_TOOLS) if agent_create.agent_type == AgentType.knowledge_vault_memory_agent: @@ -564,7 +564,7 @@ async def update_agent_tools_and_system_prompts( if agent_state.agent_type == AgentType.episodic_memory_agent: tool_names.extend(EPISODIC_MEMORY_TOOLS + UNIVERSAL_MEMORY_TOOLS) if agent_state.agent_type == AgentType.procedural_memory_agent: - tool_names.extend(PROCEDURAL_MEMORY_TOOLS + UNIVERSAL_MEMORY_TOOLS) + tool_names.extend(SKILL_TOOLS + UNIVERSAL_MEMORY_TOOLS) if agent_state.agent_type == AgentType.resource_memory_agent: tool_names.extend(RESOURCE_MEMORY_TOOLS + UNIVERSAL_MEMORY_TOOLS) if agent_state.agent_type == AgentType.knowledge_vault_memory_agent: @@ -1595,7 +1595,10 @@ async def trim_older_in_context_messages( actor: The Client performing the operation. user_id: The user whose messages to trim. If None, trims all non-system messages. """ - message_ids = await self.get_agent_by_id(agent_id=agent_id, actor=actor).message_ids + # get_agent_by_id is async โ€” await the coroutine FIRST, then read .message_ids. + # `await X.message_ids` would bind await to the attribute access on the un-awaited + # coroutine -> "AttributeError: 'coroutine' object has no attribute 'message_ids'". + message_ids = (await self.get_agent_by_id(agent_id=agent_id, actor=actor)).message_ids system_message_id = message_ids[0] message_ids = message_ids[1:] @@ -1631,7 +1634,10 @@ async def trim_all_in_context_messages_except_system( actor: The Client performing the operation. user_id: The user whose messages to remove. If None, removes all non-system messages. """ - message_ids = await self.get_agent_by_id(agent_id=agent_id, actor=actor).message_ids + # get_agent_by_id is async โ€” await the coroutine FIRST, then read .message_ids. + # `await X.message_ids` would bind await to the attribute access on the un-awaited + # coroutine -> "AttributeError: 'coroutine' object has no attribute 'message_ids'". + message_ids = (await self.get_agent_by_id(agent_id=agent_id, actor=actor)).message_ids system_message_id = message_ids[0] # 0 is system message # Keep system message and only filter out messages belonging to the specified user @@ -1651,7 +1657,10 @@ async def prepend_to_in_context_messages( actor: PydanticClient, user_id: Optional[str] = None, ) -> PydanticAgentState: - message_ids = await self.get_agent_by_id(agent_id=agent_id, actor=actor).message_ids + # get_agent_by_id is async โ€” await the coroutine FIRST, then read .message_ids. + # `await X.message_ids` would bind await to the attribute access on the un-awaited + # coroutine -> "AttributeError: 'coroutine' object has no attribute 'message_ids'". + message_ids = (await self.get_agent_by_id(agent_id=agent_id, actor=actor)).message_ids new_messages = await self.message_manager.create_many_messages(messages, actor=actor, user_id=user_id) message_ids = [message_ids[0]] + [m.id for m in new_messages] + message_ids[1:] return await self.set_in_context_messages(agent_id=agent_id, message_ids=message_ids, actor=actor) diff --git a/mirix/services/agent_trigger_state_manager.py b/mirix/services/agent_trigger_state_manager.py new file mode 100644 index 000000000..2cb688bb9 --- /dev/null +++ b/mirix/services/agent_trigger_state_manager.py @@ -0,0 +1,389 @@ +import datetime as dt +from dataclasses import dataclass +from datetime import datetime +from typing import List, Optional + +from sqlalchemy import func, select +from sqlalchemy.exc import IntegrityError + +from mirix.log import get_logger +from mirix.orm.agent_trigger_state import AgentTriggerState as AgentTriggerStateModel +from mirix.orm.conversation_message import ( + ConversationMessage as ConversationMessageModel, +) +from mirix.schemas.agent_trigger_state import ( + AgentTriggerState as PydanticAgentTriggerState, + _validate_trigger_type, +) +from mirix.utils import enforce_types + +logger = get_logger(__name__) + + +@dataclass +class ClaimFireResult: + """Outcome of a `check_and_claim_fire` call. + + fired -- True if the cursor was advanced and the caller should + run its downstream action (e.g. kick off procedural + extraction). Exactly one worker gets `fired=True` per + qualifying batch, enforced by `SELECT ... FOR UPDATE`. + sessions_since -- Sealed, not-yet-distilled conversation sessions observed + for this (user, org) at decision time (excluding the + previous fire's still-unmarked claim). Useful for + logging / metrics. + just_installed -- True iff this call initialized the cursor for the + first time. Callers typically treat this as "don't + fire, just start counting from here". + state -- The persisted cursor after the call, for logging. + """ + + fired: bool + sessions_since: int + just_installed: bool + state: PydanticAgentTriggerState + + +class AgentTriggerStateManager: + """ + Manages per-(agent, user, trigger_type) bookkeeping rows used by + interval-driven memory triggers. + + Source of truth is the **Conversation Message Store**, not the agent-loop + `messages` table: the running counter is the number of SEALED, + not-yet-distilled sessions for this `(user, organization)`, derived at read + time. "Sealed" = a strictly-newer distinct session exists (by + `MIN(created_at)` per session), so the open head of the window never counts; + "not-yet-distilled" = none of the session's turns has a `distilled_at`. + Because distillation stamps `distilled_at` on a session's turns, that session + drops out of the count permanently โ€” the marker IS the dedup, replacing the + old `messages`-table watermark/tie machinery entirely. + + The fire is intentionally a SOFT, eventually-consistent trigger. It does not + try to make each fire claim a unique batch; instead the DOWNSTREAM procedural + auto-dream is serialized per `(user, organization)` (see + `auto_dream_manager._procedural_dream_guard`) and is idempotent via + `distilled_at` marking. So a redundant fire (two `trigger_memory_update` + calls before the first auto-dream has marked its batch) is harmless: the + second auto-dream run waits on the lock, then re-enumerates the OLDEST + sealed-undistilled sessions โ€” which is exactly the right next batch, and + which naturally RETRIES any session a previous run failed to distill (a failed + session keeps `distilled_at` NULL, stays sealed-undistilled, and is counted + and re-enumerated until it succeeds). + + The only persisted state is the cursor row (audit/log only โ€” it does NOT gate + the count, which is derived live from the store): + + * last_fired_at -- wall-clock of the last fire. + * last_fired_session_id -- the in-progress session at fire time. + * last_fired_tied_session_ids -- session_ids of the most recent fire's + batch (diagnostics only). NOT used to + exclude from the next count โ€” excluding a + still-undistilled (e.g. failed) session + from its own retry count would strand it, + so the count is always the raw live + sealed-undistilled set. + + Concurrency: `check_and_claim_fire` is the serialization point. It locks the + cursor row with `SELECT ... FOR UPDATE`, counts sealed-undistilled sessions + inside the same transaction, and โ€” when the count reaches the threshold โ€” + advances the cursor before committing. Two concurrent workers on the same + (agent, user, trigger_type) execute the fire decision strictly in order; a + redundant fire is absorbed by the downstream lock + idempotent marking. + """ + + def __init__(self): + from mirix.server.server import db_context + + self.session_maker = db_context + + @enforce_types + async def check_and_claim_fire( + self, + *, + agent_id: str, + user_id: str, + trigger_type: str, + threshold: int, + organization_id: Optional[str] = None, + current_session_id: Optional[str] = None, + ) -> ClaimFireResult: + """ + Atomically decide whether to fire, and if so, claim the fire by + recording the claimed sessions on the cursor. + + Semantics: + - No cursor row: install one and return fired=False, + just_installed=True. The first call only starts the cursor; the + threshold is evaluated on subsequent calls. (The Conversation Message + Store holds only real session'd conversations โ€” there is no scaffolding + backlog to sweep โ€” but the install-then-count step keeps the first + call cheap and the fire decision uniform.) + - Cursor exists: inside one transaction, lock it, count SEALED, + not-yet-distilled sessions for this (user, org) โ€” excluding the + session_ids claimed by the previous fire that the background distill + has not yet marked distilled. If count >= threshold, claim the oldest + `threshold` of them, store that set on the cursor, advance + `last_fired_at` to wall-clock now, and return fired=True. + - Cursor exists but count < threshold: no writes, release the lock. + - organization_id cannot be resolved (neither the call nor the cursor + carries one): cannot scope the store query, so never fire. + """ + _validate_trigger_type(trigger_type) + if threshold <= 0: + raise ValueError(f"threshold must be positive, got {threshold}") + + # First-install is its own short path: never fires, and racing two + # installers is fine because the unique key + IntegrityError fallback + # makes INSERT idempotent. + async with self.session_maker() as session: + existing = await self._fetch(session, agent_id, user_id, trigger_type) + if existing is None: + installed = await self._install_cursor( + session=session, + agent_id=agent_id, + user_id=user_id, + trigger_type=trigger_type, + organization_id=organization_id, + session_id=current_session_id, + ) + return ClaimFireResult( + fired=False, + sessions_since=0, + just_installed=True, + state=installed.to_pydantic(), + ) + + # Row exists; do threshold check and cursor advance under a + # pessimistic lock so concurrent workers serialize on this cursor. + async with self.session_maker() as session: + locked_stmt = ( + select(AgentTriggerStateModel) + .where( + AgentTriggerStateModel.agent_id == agent_id, + AgentTriggerStateModel.user_id == user_id, + AgentTriggerStateModel.trigger_type == trigger_type, + AgentTriggerStateModel.is_deleted.is_(False), + ) + .with_for_update() + ) + result = await session.execute(locked_stmt) + row = result.scalar_one_or_none() + if row is None: + # Lost a race with a deletion. Treat like a re-install. + installed = await self._install_cursor( + session=session, + agent_id=agent_id, + user_id=user_id, + trigger_type=trigger_type, + organization_id=organization_id, + session_id=current_session_id, + ) + await session.commit() + return ClaimFireResult( + fired=False, + sessions_since=0, + just_installed=True, + state=installed.to_pydantic(), + ) + + # Scope the store query by (user, org). Prefer the call's org, fall + # back to the cursor's stored org. Without one we cannot isolate this + # owner's sessions, so refuse to fire rather than count across orgs. + effective_org_id = organization_id or row.organization_id + if effective_org_id is None: + await session.commit() + logger.debug( + "Skipping procedural trigger: no organization_id to scope the " + "conversation store (agent=%s, user=%s).", + agent_id, + user_id, + ) + return ClaimFireResult( + fired=False, + sessions_since=0, + just_installed=False, + state=row.to_pydantic(), + ) + + # The full set of sealed, not-yet-distilled sessions for this + # (user, org), OLDEST first. distilled_at is the dedup, so anything an + # earlier distill already marked is gone; a session that FAILED to + # distill stays here and is therefore counted again (auto-retry). + all_sealed = await self._aggregate_sealed_undistilled( + session, + user_id=user_id, + organization_id=effective_org_id, + ) + count = len(all_sealed) + + if count < threshold: + # No write; the FOR UPDATE lock releases at commit. + await session.commit() + return ClaimFireResult( + fired=False, + sessions_since=count, + just_installed=False, + state=row.to_pydantic(), + ) + + # Fire. Record the oldest `threshold` sealed sessions on the cursor for + # diagnostics (NOT to exclude from future counts โ€” the downstream + # auto-dream is serialized per owner and idempotent, so a redundant + # fire is harmless and excluding would strand a failed session from its + # own retry). + fired_batch = all_sealed[:threshold] + row.last_fired_at = datetime.now(dt.timezone.utc) + row.last_fired_session_id = current_session_id + row.last_fired_tied_session_ids = fired_batch + row.set_updated_at() + await session.commit() + await session.refresh(row) + return ClaimFireResult( + fired=True, + sessions_since=count, + just_installed=False, + state=row.to_pydantic(), + ) + + @enforce_types + async def count_sealed_undistilled_sessions( + self, + *, + user_id: str, + organization_id: str, + ) -> int: + """Diagnostic helper: count SEALED, not-yet-distilled sessions for this + `(user, organization)` in the Conversation Message Store, using the same + seal/distill aggregate as the fire path. Not on the hot path; exposed for + tests/metrics/admin. (Does NOT subtract any in-flight claims โ€” that is a + fire-path concern; this is the raw sealed-undistilled count.) + """ + async with self.session_maker() as session: + sealed = await self._aggregate_sealed_undistilled( + session, + user_id=user_id, + organization_id=organization_id, + ) + return len(sealed) + + async def _aggregate_sealed_undistilled( + self, + session, + *, + user_id: str, + organization_id: str, + ) -> List[str]: + """ + Return ALL SEALED, not-yet-distilled session_ids for this + `(user, organization)` in the Conversation Message Store, OLDEST first, + from ONE aggregate query. + + Shape mirrors + `ConversationMessageManager.list_sealed_undistilled_sessions`: one + `GROUP BY session_id` with `MIN(created_at)` (first-appearance order) and + `MAX(distilled_at)` (a session is undistilled iff this is NULL on every + turn). The single newest session by `MIN(created_at)` is the OPEN HEAD of + the window and is dropped โ€” "sealed" means a strictly-newer distinct + session exists, so the in-progress head is never counted. Among the + remaining sealed sessions, the undistilled ones are returned in oldest- + first order. + + Exclusion of in-flight claims and the threshold cap are deliberately the + CALLER's job (it needs the full set to prune its stale-claim list against), + so this aggregate stays a pure "what is sealed and undistilled right now". + """ + first_ts = func.min(ConversationMessageModel.created_at).label("first_ts") + # MAX(distilled_at) IS NULL <=> no turn of this session has been distilled + # (any non-null turn makes the MAX non-null), so the session is "open". + last_distilled = func.max(ConversationMessageModel.distilled_at).label( + "last_distilled" + ) + stmt = ( + select( + ConversationMessageModel.session_id, + first_ts, + last_distilled, + ) + .where( + ConversationMessageModel.organization_id == organization_id, + ConversationMessageModel.user_id == user_id, + # Skip soft-deleted turns: if a conversation is wiped, the + # sessions it belonged to should not keep contributing to the + # procedural fire threshold. Mirrors the manager's filter. + ConversationMessageModel.is_deleted.is_(False), + ) + .group_by(ConversationMessageModel.session_id) + .order_by(first_ts.asc()) + ) + result = await session.execute(stmt) + grouped = result.all() + + if len(grouped) <= 1: + # 0 or 1 session: nothing is sealed (no strictly-newer session). + return [] + + # The last row (newest MIN(created_at)) is the open head โ€” drop it. Among + # the remaining sealed sessions, keep the undistilled ones, oldest-first. + sealed = grouped[:-1] + return [grp.session_id for grp in sealed if grp.last_distilled is None] + + async def _install_cursor( + self, + *, + session, + agent_id: str, + user_id: str, + trigger_type: str, + organization_id: Optional[str], + session_id: Optional[str], + ) -> AgentTriggerStateModel: + """ + Insert a fresh cursor row at `now`. Idempotent: if two workers both + try to install concurrently, one wins and the other reads back the + winner's row via the unique constraint. + """ + now = datetime.now(dt.timezone.utc) + row = AgentTriggerStateModel( + agent_id=agent_id, + user_id=user_id, + trigger_type=trigger_type, + organization_id=organization_id, + last_fired_at=now, + last_fired_session_id=session_id, + # Start the claimed-set EMPTY. Under the distilled_at-based scheme this + # list holds session_ids CLAIMED by a fire (excluded from the next + # count until the background distill marks them distilled) โ€” seeding it + # with the in-progress session_id would wrongly exclude the very first + # real session from the first threshold window, shifting the cadence by + # one (the 6th distinct session would only see 4 countable sealed + # sessions). No fire has happened yet, so nothing is claimed. + last_fired_tied_session_ids=[], + ) + session.add(row) + try: + await session.commit() + except IntegrityError: + await session.rollback() + existing = await self._fetch(session, agent_id, user_id, trigger_type) + if existing is None: + raise + return existing + await session.refresh(row) + return row + + async def _fetch( + self, + session, + agent_id: str, + user_id: str, + trigger_type: str, + ) -> Optional[AgentTriggerStateModel]: + stmt = select(AgentTriggerStateModel).where( + AgentTriggerStateModel.agent_id == agent_id, + AgentTriggerStateModel.user_id == user_id, + AgentTriggerStateModel.trigger_type == trigger_type, + AgentTriggerStateModel.is_deleted.is_(False), + ) + result = await session.execute(stmt) + return result.scalar_one_or_none() diff --git a/mirix/services/auto_dream_manager.py b/mirix/services/auto_dream_manager.py index b08e3f4b9..cf51e7d5d 100644 --- a/mirix/services/auto_dream_manager.py +++ b/mirix/services/auto_dream_manager.py @@ -10,12 +10,15 @@ 6. Return stats """ +import contextlib import datetime as dt import json import logging import os from typing import List, Optional +from mirix.helpers.keyed_locks import KeyedLocks +from mirix.services.conversation_message_manager import owner_org as _owner_org from mirix.schemas.agent import AgentState, AgentType, CreateAgent from mirix.schemas.auto_dream import AutoDreamRequest, AutoDreamResponse, MemoryTypeStats from mirix.schemas.client import Client as PydanticClient @@ -28,12 +31,129 @@ # event_type used to tag auto_dream checkpoint records in episodic memory _CHECKPOINT_EVENT_TYPE = "auto_dream_checkpoint" +# Per-(user, org) serialization for the procedural auto-dream. Two fires can be +# scheduled (fire-and-forget) before the first has marked its sessions distilled; +# if both ran concurrently they would each enumerate the SAME oldest sealed +# sessions and distill duplicate experiences. Serializing per owner makes the +# second run wait until the first has marked its batch distilled, so the second +# then enumerates the NEXT (disjoint) batch. +# +# Two layers: +# * a process-local asyncio.Lock (serializes coroutines within ONE event loop), AND +# * a Postgres SESSION-level ADVISORY lock keyed by hash(user, org), held across +# the whole run, so concurrent gunicorn/uvicorn WORKER PROCESSES also serialize. +# On non-Postgres backends (SQLite / PGlite โ€” single-process test setups with no +# advisory locks) the asyncio.Lock alone is sufficient, so the advisory layer is +# skipped gracefully. +# Self-evicting per-(user, org) registry; see _procedural_dream_guard. +_PROCEDURAL_DREAM_LOCKS = KeyedLocks() + +# Stable namespace salt so this feature's advisory-lock keys never collide with +# another feature's pg_advisory_lock keyspace. +_PROCEDURAL_ADVISORY_NAMESPACE = "mirix.procedural_dream" + + +def _dream_lock_key(user_id: str, organization_id: Optional[str]) -> str: + """Registry key for the process-local procedural auto-dream lock โ€” keyed + per (user, organization) so different owners never block each other.""" + return f"{organization_id or '-'}::{user_id}" + + +def _advisory_lock_key(user_id: str, organization_id: Optional[str]) -> int: + """Map (user, org) to a stable signed 64-bit int for pg_advisory_lock. + + Postgres advisory locks key on a bigint; we hash the namespaced owner string + with blake2b (stable across processes, unlike Python's salted hash()) and fold + it into the signed 64-bit range Postgres expects. + """ + import hashlib + + raw = f"{_PROCEDURAL_ADVISORY_NAMESPACE}:{organization_id or '-'}:{user_id}".encode() + digest = hashlib.blake2b(raw, digest_size=8).digest() + val = int.from_bytes(digest, "big", signed=False) + # Fold unsigned 64-bit into signed 64-bit range expected by bigint. + return val - (1 << 64) if val >= (1 << 63) else val + + +@contextlib.asynccontextmanager +async def _procedural_dream_guard(user_id: str, organization_id: Optional[str]): + """Serialize the procedural auto-dream for one (user, org) across coroutines + AND across processes. + + Holds the process-local asyncio.Lock for the whole run, and โ€” on Postgres โ€” + additionally takes a SESSION-level advisory lock on a dedicated connection + held open for the run's duration, releasing it in a finally. On non-Postgres + backends (SQLite/PGlite โ€” single-process setups the asyncio.Lock already + covers) the advisory step is skipped. Lock-plumbing failures degrade to the + asyncio.Lock rather than crashing the run. + """ + from sqlalchemy import text + + from mirix.server.server import db_context + + key = _advisory_lock_key(user_id, organization_id) + async with _PROCEDURAL_DREAM_LOCKS.acquire(_dream_lock_key(user_id, organization_id)): + # Open a dedicated session and try to hold a session-level advisory lock + # for the whole run. session_cm/session are None when we are not on + # Postgres or acquisition failed โ€” then the asyncio.Lock alone serializes. + session_cm = None + session = None + try: + session_cm = db_context() + session = await session_cm.__aenter__() + dialect = session.bind.dialect.name if session.bind is not None else "" + if dialect == "postgresql": + # SESSION-level (NOT xact) so it spans the run's many transactions. + await session.execute( + text("SELECT pg_advisory_lock(:k)"), {"k": key} + ) + else: + # No advisory-lock primitive on SQLite/PGlite โ€” release the unused + # session; the asyncio.Lock already covers these single-process setups. + await session_cm.__aexit__(None, None, None) + session_cm = None + session = None + except Exception: # noqa: BLE001 โ€” lock plumbing must never crash the run + logger.warning( + "Procedural dream advisory lock unavailable; relying on the " + "process-local lock only.", + exc_info=True, + ) + if session_cm is not None: + try: + await session_cm.__aexit__(None, None, None) + except Exception: # noqa: BLE001 + pass + session_cm = None + session = None + + try: + yield + finally: + if session is not None: + try: + await session.execute( + text("SELECT pg_advisory_unlock(:k)"), {"k": key} + ) + except Exception: # noqa: BLE001 โ€” close releases it anyway + logger.debug( + "pg_advisory_unlock failed; session close will release it." + ) + finally: + try: + await session_cm.__aexit__(None, None, None) + except Exception: # noqa: BLE001 + pass + +# Generic "fetch current memories of these components -> one consolidation agent" +# modes. mode='procedural' is intentionally NOT here: it is handled by the +# dedicated session-experience path (_run_procedural_experience), which returns +# before this mapping is ever consulted. _MODE_COMPONENTS = { "core": ["core"], "episodic": ["episodic"], "semantic": ["semantic"], "resource": ["resource"], - "procedural": ["procedural"], "knowledge": ["knowledge"], "experience": ["episodic", "semantic", "knowledge"], } @@ -240,7 +360,7 @@ async def get_or_create_dream_agent_state( meta_agent_state: AgentState, ) -> AgentState: """Return the auto_dream_agent state for this client, creating it if needed.""" - from mirix.server.rest_api import get_server + from mirix.server.server import get_server server = get_server() children = await server.agent_manager.list_agents( @@ -262,6 +382,368 @@ async def get_or_create_dream_agent_state( ) return await server.agent_manager.create_agent(agent_create=agent_create, actor=actor) + # ------------------------------------------------------------------ # + # Session distillation + skill evolution โ€” procedural path # + # ------------------------------------------------------------------ # + + async def _run_procedural_experience( + self, + *, + request: AutoDreamRequest, + user: PydanticUser, + actor: PydanticClient, + meta_agent_state: AgentState, + now: dt.datetime, + ) -> AutoDreamResponse: + """Serialize per (user, org), then run the procedural distillation. + + The guard makes the enumerate โ†’ distill โ†’ mark sequence atomic against + another procedural run for the SAME owner โ€” within the process (asyncio + lock) AND across worker processes (Postgres advisory lock). A second fire + scheduled before this one marks its batch distilled waits here, then + enumerates the NEXT (disjoint) sealed batch โ€” preventing duplicate + distillation of the oldest sessions. See ``_procedural_dream_guard``. + """ + async with _procedural_dream_guard(user.id, _owner_org(actor)): + return await self._run_procedural_experience_locked( + request=request, + user=user, + actor=actor, + meta_agent_state=meta_agent_state, + now=now, + ) + + async def _run_procedural_experience_locked( + self, + *, + request: AutoDreamRequest, + user: PydanticUser, + actor: PydanticClient, + meta_agent_state: AgentState, + now: dt.datetime, + ) -> AutoDreamResponse: + """Session distillation (+ skill evolution unless dry_run). + + This is the explicit "program call" entry point for general + session-experience distillation: calling + ``AutoDreamManager().run(AutoDreamRequest(mode='procedural', ...))`` + runs it directly. It is ALSO fired every N sessions by the + ``trigger_memory_update`` claim path (see + ``functions/function_sets/memory_tools.py``). + + Steps: + 1. Resolve N (``request.last_n_sessions`` or + ``MESSAGE_RETAIN_LAST_N_SESSIONS``) โ€” the max number of sealed + sessions to process this round. + 2. Enumerate up to N SEALED, not-yet-distilled sessions from the + Conversation Message Store (oldest-first; the open head of the + window is never included). + 3. Distill each session IN PARALLEL into pending SkillExperience rows. + 4. Mark the SUCCESSFULLY-processed sessions distilled so the rolling + barrier advances and they are never re-distilled (failed sessions + are left for retry). + 5. dry_run โ†’ return counts; else run skill evolution over the + freshly-pending experiences, then write the dream checkpoint. + + Always invoked under the per-owner lock acquired by + ``_run_procedural_experience``. + """ + from mirix.constants import MESSAGE_RETAIN_LAST_N_SESSIONS + from mirix.services.conversation_message_manager import ( + ConversationMessageManager, + ) + from mirix.services.session_experience_distiller import ( + SessionExperienceDistiller, + ) + + n = request.last_n_sessions or MESSAGE_RETAIN_LAST_N_SESSIONS + + # The dream agent inherits the meta agent's llm_config; allow a model + # override for testing (same convention as the generic path). + llm_config = meta_agent_state.llm_config + if request.model: + from copy import deepcopy + + llm_config = deepcopy(llm_config) + llm_config.model = request.model + + conversation_manager = ConversationMessageManager() + distiller = SessionExperienceDistiller( + llm_config=llm_config, + conversation_manager=conversation_manager, + ) + # Source of truth is the Conversation Message Store, scoped per + # (user, organization). Pull at most N sealed, not-yet-distilled + # sessions (oldest-first); the in-progress head is left for the next + # window. + session_ids = await distiller.enumerate_sealed_sessions( + user_id=user.id, + organization_id=_owner_org(actor), + actor=actor, + limit=n, + ) + logger.info( + "Auto dream (procedural): meta_agent=%s, user=%s, last_n=%d, sealed_sessions=%s", + meta_agent_state.id, + user.id, + n, + session_ids, + ) + + if request.dry_run: + # A dry run has NO side effects: it does NOT distill (no LLM call, no + # SkillExperience writes), does NOT mark sessions distilled, does NOT + # evolve skills, and does NOT checkpoint. It reports how many sealed + # sessions WOULD be processed, leaving the rolling barrier untouched + # so a subsequent real run processes the same window. + return AutoDreamResponse( + start_date=None, + end_date=None, + processed={"procedural": MemoryTypeStats(total=0)}, + last_dream_at=now, + dry_run=True, + message=( + f"Dry run โ€” {len(session_ids)} sealed session(s) would be " + f"distilled; no experiences persisted, barrier not advanced." + ), + ) + + existing_skills = await self._existing_skill_summaries(user, meta_agent_state) + # distill_sessions returns (experiences, processed_session_ids): only the + # SUCCESSFULLY-processed sessions (including legitimately-empty ones) are + # in processed_session_ids; a session that hit an operational failure is + # omitted so we never advance the barrier past a failed conversation. + experiences, processed_session_ids = await distiller.distill_sessions( + meta_agent_state=meta_agent_state, + user=user, + actor=actor, + session_ids=session_ids, + existing_skills=existing_skills, + ) + + stats = { + "procedural": MemoryTypeStats(total=len(experiences)), + } + + # Advance the rolling barrier: mark every SUCCESSFULLY-processed session + # distilled (idempotent; scoped to this (user, org)) so a later round โ€” + # automatic or explicit โ€” never re-distills them. Sessions that yielded + # zero experiences but were processed cleanly are still marked (an + # empty/low-signal conversation is "consumed" and must not block the + # window). Sessions that FAILED are intentionally left undistilled so a + # later round retries them rather than losing their learning. + if processed_session_ids: + marked = await conversation_manager.mark_sessions_distilled( + session_ids=processed_session_ids, + user_id=user.id, + organization_id=_owner_org(actor), + actor=actor, + ) + logger.info( + "Auto dream (procedural): marked %d turn(s) across %d/%d session(s) " + "distilled (%d failed, left for retry)", + marked, + len(processed_session_ids), + len(session_ids), + len(session_ids) - len(processed_session_ids), + ) + + # Retention: distilled turns past the window have already yielded their + # learning into skill_experience rows; drop the verbatim transcripts so + # the store doesn't accumulate raw conversation content forever. + from mirix.constants import CONVERSATION_RETENTION_DAYS + + try: + await conversation_manager.prune_distilled_turns( + user_id=user.id, + organization_id=_owner_org(actor), + actor=actor, + retention_days=CONVERSATION_RETENTION_DAYS, + ) + except Exception: + # Best-effort: a failed prune must not fail the dream run; the next + # pass retries. Logged so a persistent failure is visible. + logger.exception("Auto dream (procedural): retention prune failed") + + # -- Skill evolution: evolve skills from the pending experiences. -- + # + # Two independent steps, so a prior round's stranded experiences are + # drained even under active traffic (closing the "evolution failure + # strands distilled experiences" gap): + # (1) FRESH round: evolve THIS round's freshly-distilled batch, scoped to + # its own ids (the deliberate optimization โ€” earlier rounds' + # low-signal experiences would dilute the curator prompt). + # (2) RECOVERY: drain any LEFTOVER pending experiences from PRIOR rounds + # whose evolution failed (leftover pending minus this round's fresh + # ids), in its own scoped run. Always attempted, regardless of + # whether this round produced fresh experiences. + # The sessions stay distilled either way โ€” re-distilling would duplicate; + # the experiences are the unit retried, not the sessions. + skills_changed = 0 + evolution_changes: dict = {} + fresh_ids = [e.id for e in experiences] + leftover_ids = await self._leftover_pending_experience_ids( + meta_agent_state=meta_agent_state, user=user + ) + # Exclude this round's fresh ids from the recovery scope: they are handled + # by step (1), so the recovery step only sees PRIOR rounds' leftovers. + fresh_set = set(fresh_ids) + recovery_ids = [eid for eid in leftover_ids if eid not in fresh_set] + + if fresh_ids: + sc, ch = await self._evolve_experiences( + user=user, + actor=actor, + meta_agent_state=meta_agent_state, + experience_ids=fresh_ids, + label="fresh", + ) + skills_changed += sc + evolution_changes.update(ch) + if recovery_ids: + sc, ch = await self._evolve_experiences( + user=user, + actor=actor, + meta_agent_state=meta_agent_state, + experience_ids=recovery_ids, + label="recovery", + ) + skills_changed += sc + # Don't clobber fresh changes with recovery ones โ€” merge keys. + for k, v in ch.items(): + evolution_changes.setdefault(k, v) + + await self.write_checkpoint(user, actor, meta_agent_state, now) + + return AutoDreamResponse( + start_date=None, + end_date=None, + processed=stats, + last_dream_at=now, + dry_run=False, + skills_changed=skills_changed, + changes=evolution_changes, + message=( + f"Auto dream (procedural) completed: {len(experiences)} experience(s) " + f"from {len(session_ids)} session(s); {skills_changed} skill(s) changed." + ), + ) + + async def _evolve_experiences( + self, + *, + user: PydanticUser, + actor: PydanticClient, + meta_agent_state: AgentState, + experience_ids: List[str], + label: str, + ) -> tuple: + """Run skill evolution over an explicit, scoped experience batch. + + Returns ``(skills_changed, changes)``; on any failure returns ``(0, {})`` + and logs โ€” the experiences are durably `pending`, so a failed evolution + loses no learning and is retried by the RECOVERY step on a later run. + ``label`` ("fresh" / "recovery") only colours the log. + """ + if not experience_ids: + return 0, {} + try: + from mirix.services.skill_experience_curator import ( + run_experience_evolution, + ) + + evolution = await run_experience_evolution( + user=user, + actor=actor, + meta_agent_state=meta_agent_state, + experience_ids=experience_ids, + ) + skills_changed = (evolution or {}).get("skills_changed", 0) + changes = (evolution or {}).get("changes", {}) or {} + if label == "recovery": + logger.info( + "Auto dream (procedural): recovered %d stranded pending " + "experience(s); %d skill(s) changed.", + len(experience_ids), + skills_changed, + ) + return skills_changed, changes + except ImportError: + # Experience curator not yet present in this build โ€” distillation alone is + # still valuable, so do not fail the run. + logger.info( + "Experience curator unavailable; experiences persisted as pending." + ) + return 0, {} + except Exception as e: # noqa: BLE001 โ€” evolution failure mustn't lose experiences + # Experiences are durably `pending`, so NO learning is lost โ€” only the + # evolution STEP failed. The sessions stay distilled (re-distilling + # would duplicate); these pending experiences are drained by the + # RECOVERY step on a later procedural run. Logged LOUD for observability. + logger.error( + "Skill-evolution run (%s) FAILED for meta_agent=%s " + "(%d experience(s) left pending; auto-recovered on a later procedural " + "run, or via POST /memory/auto_dream mode=procedural): %s", + label, + meta_agent_state.id, + len(experience_ids), + e, + exc_info=True, + ) + return 0, {} + + async def _leftover_pending_experience_ids( + self, + *, + meta_agent_state: AgentState, + user: PydanticUser, + ) -> List[str]: + """Return ids of this (agent, user)'s still-`pending` experiences. + + Used by the RECOVERY path: when a procedural run distilled nothing fresh, + these are the experiences a PRIOR round's failed evolution left behind + (``list_experiences(status='pending')`` already excludes consumed / + superseded, so this is exactly the un-evolved leftover pool). Best-effort: + any failure yields ``[]`` so recovery is simply skipped this round. + """ + try: + from mirix.services.skill_experience_manager import SkillExperienceManager + + pending = await SkillExperienceManager().list_experiences( + agent_id=meta_agent_state.id, + user_id=user.id, + status="pending", + ) + except Exception as e: # noqa: BLE001 โ€” recovery is best-effort + logger.debug("Could not list leftover pending experiences: %s", e) + return [] + return [e.id for e in pending] + + async def _existing_skill_summaries( + self, + user: PydanticUser, + meta_agent_state: AgentState, + ) -> list: + """Return [{name, description}] of existing procedural skills (dedup context). + + Best-effort: a failure here just yields an empty list (the distiller + prompt treats it as "(none)"). + """ + try: + procedures = await self._fetch_procedural(user, meta_agent_state, None, None) + except Exception as e: # noqa: BLE001 + logger.debug("Could not fetch existing skills for dedup context: %s", e) + return [] + summaries = [] + for p in procedures: + summaries.append( + { + "name": getattr(p, "name", "") or "", + "description": getattr(p, "description", "") or "", + } + ) + return summaries + # ------------------------------------------------------------------ # # Main entry point # # ------------------------------------------------------------------ # @@ -273,12 +755,29 @@ async def run( actor: PydanticClient, meta_agent_state: AgentState, ) -> AutoDreamResponse: - from mirix.agent.auto_dream_agent import AutoDreamAgent - from mirix.server.rest_api import get_server + from mirix.server.server import get_server server = get_server() now = dt.datetime.now(dt.timezone.utc) + # ----------------------------------------------------------------- # + # mode='procedural' is general per-session experience distillation # + # distillation (NOT the legacy procedure-consolidation pass). It # + # reads the meta agent's last-N RETAINED sessions, distills each # + # session's transcript IN PARALLEL into pending SkillExperience # + # rows, then (unless dry_run) drives skill self-evolution # + # from those experiences. Other modes fall through to the generic # + # fetch โ†’ format โ†’ single-step path below, unchanged. # + # ----------------------------------------------------------------- # + if request.mode == "procedural": + return await self._run_procedural_experience( + request=request, + user=user, + actor=actor, + meta_agent_state=meta_agent_state, + now=now, + ) + # -- resolve time window -- end_date = request.end_date or now if request.start_date: diff --git a/mirix/services/client_manager.py b/mirix/services/client_manager.py index ffcfecd47..a32c96f43 100644 --- a/mirix/services/client_manager.py +++ b/mirix/services/client_manager.py @@ -404,6 +404,8 @@ async def delete_memories_by_client_id(self, client_id: str): - ResourceMemoryManager.delete_by_client_id() - KnowledgeVaultManager.delete_by_client_id() - MessageManager.delete_by_client_id() + - ConversationMessageManager.delete_by_client_id() (verbatim session transcripts) + - SkillExperienceManager.delete_by_client_id() (experiences distilled from them) 2. Delete blocks (via _created_by_id) 3. Each manager handles: - Bulk database deletion @@ -423,12 +425,14 @@ async def delete_memories_by_client_id(self, client_id: str): ) # Import managers + from mirix.services.conversation_message_manager import ConversationMessageManager from mirix.services.episodic_memory_manager import EpisodicMemoryManager from mirix.services.knowledge_vault_manager import KnowledgeVaultManager from mirix.services.message_manager import MessageManager from mirix.services.procedural_memory_manager import ProceduralMemoryManager from mirix.services.resource_memory_manager import ResourceMemoryManager from mirix.services.semantic_memory_manager import SemanticMemoryManager + from mirix.services.skill_experience_manager import SkillExperienceManager # Initialize managers episodic_manager = EpisodicMemoryManager() @@ -437,6 +441,8 @@ async def delete_memories_by_client_id(self, client_id: str): resource_manager = ResourceMemoryManager() knowledge_manager = KnowledgeVaultManager() message_manager = MessageManager() + conversation_manager = ConversationMessageManager() + skill_experience_manager = SkillExperienceManager() # Get client as actor for manager methods client = await self.get_client_by_id(client_id) @@ -465,6 +471,14 @@ async def delete_memories_by_client_id(self, client_id: str): message_count = await message_manager.delete_by_client_id(actor=client) logger.debug("Bulk deleted %d messages", message_count) + # Verbatim session transcripts and the experiences distilled from + # them were ingested by this client โ€” the purge must cover them. + conversation_count = await conversation_manager.delete_by_client_id(actor=client) + logger.debug("Bulk deleted %d conversation turns", conversation_count) + + experience_count = await skill_experience_manager.delete_by_client_id(actor=client) + logger.debug("Bulk deleted %d skill experiences", experience_count) + # Delete blocks created by this client (using bulk operations) block_count = 0 block_ids: List[str] = [] @@ -545,7 +559,8 @@ async def delete_memories_by_client_id(self, client_id: str): logger.info( "Bulk deleted all memories for client %s: " - "%d episodic, %d semantic, %d procedural, %d resource, %d knowledge_vault, %d messages, %d blocks " + "%d episodic, %d semantic, %d procedural, %d resource, %d knowledge_vault, %d messages, %d blocks, " + "%d conversation turns, %d skill experiences " "(client, agents, tools preserved)", client_id, episodic_count, @@ -555,6 +570,8 @@ async def delete_memories_by_client_id(self, client_id: str): knowledge_count, message_count, block_count, + conversation_count, + experience_count, ) except Exception as e: logger.error("Failed to bulk delete memories for client %s: %s", client_id, e) diff --git a/mirix/services/conversation_message_manager.py b/mirix/services/conversation_message_manager.py new file mode 100644 index 000000000..29d65568e --- /dev/null +++ b/mirix/services/conversation_message_manager.py @@ -0,0 +1,419 @@ +"""Manager for the Conversation Message Store. + +Async-only, mirroring `SkillExperienceManager` / `MessageManager`: it grabs +sessions via `self.session_maker` (the server's `db_context`) and never calls +`asyncio.run()` (the server event loop is already running). + +This store is the SINGLE source of truth for procedural-memory (skill) +distillation. It holds only external conversation turns that arrived with a +`session_id`, with their REAL `user`/`assistant` roles preserved. The methods +below are the stable contract the ingestion seam, the trigger/cadence logic, +the distiller, and the data-lifecycle paths all depend on: + + - record_turns -- append a batch of turns to a session + - count_distinct_sessions -- how many distinct sessions exist + - list_sealed_undistilled_sessions -- oldest sealed, not-yet-distilled ids + - list_turns_for_session -- one session's turns, ascending + - mark_sessions_distilled -- advance the rolling barrier + - prune_distilled_turns -- retention: drop old distilled turns + - delete_by_user_id -- erasure: purge one user's turns + - delete_by_client_id -- erasure: purge one client's turns + +Because this table stores VERBATIM user conversation content, it must be +covered by the same erasure guarantees as the memory tables: the user/client +"delete memories" purge paths call the two delete_* methods, and the +distillation flow prunes distilled turns past the retention window. + +"Sealed" = a strictly newer distinct `session_id` exists (by MIN(created_at)), +so the open head of the window is never distilled while it may still grow. +Everything is scoped per `(user, organization)` so one user's sessions never +count toward or leak into another's learning window. +""" + +from __future__ import annotations + +import datetime as dt +from datetime import timedelta +from typing import List + +from sqlalchemy import delete, func, select, update + +from mirix.client.utils import get_utc_time +from mirix.log import get_logger +from mirix.orm.conversation_message import ( + ConversationMessage as ConversationMessageModel, +) +from mirix.schemas.client import Client as PydanticClient +from mirix.schemas.conversation_message import ( + ConversationMessage as PydanticConversationMessage, + ConversationMessageCreate, +) +from mirix.utils import enforce_types + +logger = get_logger(__name__) + + +def owner_org(actor: PydanticClient) -> str: + """Resolve the organization the Conversation Message Store is scoped under. + + Ingestion records turns under ``client.organization_id or DEFAULT_ORG_ID`` + (the store's organization_id column is NOT NULL). Every read/mark/prune of + the store MUST mirror the same fallback, otherwise a NULL-org client would + write under DEFAULT_ORG_ID but read under ``None`` and its sessions would + never distill. Single definition here โ€” the distiller and the auto-dream + manager import it rather than re-deriving the fallback. + """ + from mirix.constants import DEFAULT_ORG_ID + + return actor.organization_id or DEFAULT_ORG_ID + + +class ConversationMessageManager: + """Persist and query external conversation turns for skill distillation.""" + + def __init__(self): + from mirix.server.server import db_context + + self.session_maker = db_context + + @enforce_types + async def record_turns( + self, + *, + session_id: str, + user_id: str, + organization_id: str, + turns: List[dict], + actor: PydanticClient, + ) -> List[PydanticConversationMessage]: + """Append `turns` to `session_id`, preserving their given order. + + `turns` is a list of `{"role": "user"|"assistant"|"tool", "content": str}`. + Each turn is validated through `ConversationMessageCreate` first so an + invalid `role` (or over-length content / malformed session_id) is + rejected up front โ€” the DB columns are plain String/Text, so without + this an invalid value would commit and only blow up later in + `to_pydantic()` / list. + + created_at ordering is preserved EXPLICITLY: a single batch shares one + transaction, and the server-side `func.now()` default would collapse all + rows in that transaction to one timestamp, destroying intra-batch order. + We instead stamp a strictly increasing `created_at` per turn โ€” anchored + at the current MAX(created_at) for this session so a later call's turns + always sort after an earlier call's, making multi-call sessions a single + ordered unit. Returns the persisted rows in insertion order. + """ + if not turns: + return [] + + validated = [ + ConversationMessageCreate( + session_id=session_id, + user_id=user_id, + organization_id=organization_id, + role=turn.get("role"), + content=turn.get("content", ""), + ) + for turn in turns + ] + + async with self.session_maker() as session: + # Anchor after the latest existing turn for this session so turns + # from successive record_turns calls accumulate in order. Scope the + # MAX by (org, user, session) for the same isolation as everything + # else here. + max_stmt = select(func.max(ConversationMessageModel.created_at)).where( + ConversationMessageModel.organization_id == organization_id, + ConversationMessageModel.user_id == user_id, + ConversationMessageModel.session_id == session_id, + ConversationMessageModel.is_deleted.is_(False), + ) + result = await session.execute(max_stmt) + last_ts = result.scalar_one_or_none() + + base = get_utc_time() + if last_ts is not None: + # Normalize to tz-aware UTC for a safe comparison with `base`. + if last_ts.tzinfo is None: + last_ts = last_ts.replace(tzinfo=dt.timezone.utc) + if last_ts >= base: + base = last_ts + timedelta(microseconds=1) + + rows = [] + for offset, payload in enumerate(validated): + row = ConversationMessageModel( + session_id=payload.session_id, + user_id=payload.user_id, + organization_id=payload.organization_id, + role=payload.role, + content=payload.content, + # +offset microseconds keeps a strict, stable order within + # the batch even though they share one transaction. + created_at=base + timedelta(microseconds=offset), + ) + if actor is not None: + row._set_created_and_updated_by_fields(actor.id) + session.add(row) + rows.append(row) + + await session.commit() + for row in rows: + await session.refresh(row) + return [row.to_pydantic() for row in rows] + + @enforce_types + async def count_distinct_sessions( + self, + *, + user_id: str, + organization_id: str, + actor: PydanticClient, + only_undistilled: bool = False, + ) -> int: + """Count DISTINCT `session_id`s for this `(user, organization)`. + + With `only_undistilled=True`, counts only sessions that have NOT been + distilled. "Undistilled" is a SESSION-level property: a session counts + iff NONE of its turns have a `distilled_at` (i.e. `MAX(distilled_at) IS + NULL`). A row-level `distilled_at IS NULL` filter would over-count a + previously-distilled session that later gained a fresh turn โ€” and would + disagree with `list_sealed_undistilled_sessions`, which already seals on + `MAX(distilled_at)`. Soft-deleted rows are excluded. + """ + async with self.session_maker() as session: + preds = [ + ConversationMessageModel.organization_id == organization_id, + ConversationMessageModel.user_id == user_id, + ConversationMessageModel.is_deleted.is_(False), + ] + if only_undistilled: + # Group by session and keep only those whose MAX(distilled_at) is + # NULL, then count the surviving groups. Wrapping the grouped + # query in a COUNT subquery keeps this a single round-trip and + # dialect-safe (no DISTINCT-on-HAVING gymnastics). + grouped = ( + select(ConversationMessageModel.session_id) + .where(*preds) + .group_by(ConversationMessageModel.session_id) + .having(func.max(ConversationMessageModel.distilled_at).is_(None)) + .subquery() + ) + stmt = select(func.count()).select_from(grouped) + else: + stmt = select( + func.count(func.distinct(ConversationMessageModel.session_id)) + ).where(*preds) + result = await session.execute(stmt) + return int(result.scalar_one() or 0) + + @enforce_types + async def list_sealed_undistilled_sessions( + self, + *, + user_id: str, + organization_id: str, + actor: PydanticClient, + limit: int, + ) -> List[str]: + """Return up to `limit` sealed, not-yet-distilled session_ids, OLDEST first. + + "Sealed" = a strictly newer distinct session exists, by first-appearance + time (`MIN(created_at)` per session). The newest session is the open head + of the window and is never returned (it may still grow). Of the sealed + sessions, only those that have NOT been distilled (`distilled_at IS NULL` + on their turns) are returned. + + Shape mirrors `agent_trigger_state_manager._aggregate_window`: one + `GROUP BY session_id` with `MIN(created_at)` per session, so ordering and + sealing come from a single dialect-safe aggregate. Soft-deleted rows are + excluded. + """ + if limit <= 0: + return [] + + async with self.session_maker() as session: + first_ts = func.min(ConversationMessageModel.created_at).label("first_ts") + # NULL distilled_at on EVERY turn of a session <=> the session is + # undistilled. MAX(distilled_at) IS NULL captures exactly that (any + # non-null turn makes the MAX non-null), so a session is "open" iff + # its MAX(distilled_at) is NULL. + last_distilled = func.max(ConversationMessageModel.distilled_at).label( + "last_distilled" + ) + stmt = ( + select( + ConversationMessageModel.session_id, + first_ts, + last_distilled, + ) + .where( + ConversationMessageModel.organization_id == organization_id, + ConversationMessageModel.user_id == user_id, + ConversationMessageModel.is_deleted.is_(False), + ) + .group_by(ConversationMessageModel.session_id) + .order_by(first_ts.asc()) + ) + result = await session.execute(stmt) + grouped = result.all() + + if len(grouped) <= 1: + # 0 or 1 session: nothing is sealed (no strictly-newer session). + return [] + + # The last row (newest MIN(created_at)) is the open head โ€” drop it. + # Among the remaining sealed sessions, keep the undistilled ones, + # oldest-first, capped at `limit`. + sealed = grouped[:-1] + out: List[str] = [] + for row in sealed: + if row.last_distilled is None: + out.append(row.session_id) + if len(out) >= limit: + break + return out + + @enforce_types + async def list_turns_for_session( + self, + *, + session_id: str, + user_id: str, + organization_id: str, + actor: PydanticClient, + ) -> List[PydanticConversationMessage]: + """Return one session's turns for this `(user, org)`, ascending by created_at.""" + async with self.session_maker() as session: + stmt = ( + select(ConversationMessageModel) + .where( + ConversationMessageModel.organization_id == organization_id, + ConversationMessageModel.user_id == user_id, + ConversationMessageModel.session_id == session_id, + ConversationMessageModel.is_deleted.is_(False), + ) + .order_by(ConversationMessageModel.created_at.asc()) + ) + result = await session.execute(stmt) + rows = result.scalars().all() + return [row.to_pydantic() for row in rows] + + @enforce_types + async def mark_sessions_distilled( + self, + *, + session_ids: List[str], + user_id: str, + organization_id: str, + actor: PydanticClient, + ) -> int: + """Stamp `distilled_at` on all turns of the given sessions; idempotent. + + Only turns still un-distilled (`distilled_at IS NULL`) are stamped, so a + re-run leaves already-distilled rows untouched and the rolling barrier + advances without reprocessing history. Scoped to this `(user, org)` so a + caller can never mark another owner's session. Returns the number of + rows updated. + + A bulk UPDATE (not load-modify-save) is correct here: `distilled_at` is a + write-only barrier marker with no Redis cache and no `updated_at` + semantics the distiller depends on, so the simple set is both faster and + avoids loading potentially large transcripts into memory. + """ + if not session_ids: + return 0 + + now = get_utc_time() + async with self.session_maker() as session: + stmt = ( + update(ConversationMessageModel) + .where( + ConversationMessageModel.organization_id == organization_id, + ConversationMessageModel.user_id == user_id, + ConversationMessageModel.session_id.in_(session_ids), + ConversationMessageModel.distilled_at.is_(None), + ConversationMessageModel.is_deleted.is_(False), + ) + .values(distilled_at=now) + ) + result = await session.execute(stmt) + await session.commit() + return int(result.rowcount or 0) + + @enforce_types + async def prune_distilled_turns( + self, + *, + user_id: str, + organization_id: str, + actor: PydanticClient, + retention_days: int, + ) -> int: + """Hard-delete turns distilled more than `retention_days` ago. + + Retention policy for the verbatim conversation store: once a turn has + been distilled its learning value has been extracted into + skill_experience rows, so keeping the raw transcript indefinitely is a + pure PII liability. Row-level `distilled_at <= cutoff` is deliberately + conservative: undistilled turns are NEVER pruned (their learning hasn't + been extracted yet), even when they share a session with pruned rows. + + `retention_days <= 0` disables pruning (keep forever). Returns the + number of rows deleted. + """ + if retention_days <= 0: + return 0 + + cutoff = get_utc_time() - timedelta(days=retention_days) + async with self.session_maker() as session: + stmt = delete(ConversationMessageModel).where( + ConversationMessageModel.organization_id == organization_id, + ConversationMessageModel.user_id == user_id, + ConversationMessageModel.distilled_at.is_not(None), + ConversationMessageModel.distilled_at <= cutoff, + ) + result = await session.execute(stmt) + await session.commit() + deleted = int(result.rowcount or 0) + if deleted: + logger.info( + "Pruned %d distilled conversation turn(s) older than %d day(s) for user %s", + deleted, + retention_days, + user_id, + ) + return deleted + + @enforce_types + async def delete_by_user_id(self, user_id: str) -> int: + """Hard delete ALL conversation turns for a user (erasure path). + + Called from `UserManager.delete_memories_by_user_id`, which backs the + irreversible DELETE /users/{user_id}/memories purge. Deletes across all + organizations and regardless of distillation state โ€” erasure trumps the + learning window. Returns the number of rows deleted. + """ + async with self.session_maker() as session: + stmt = delete(ConversationMessageModel).where( + ConversationMessageModel.user_id == user_id + ) + result = await session.execute(stmt) + await session.commit() + return int(result.rowcount or 0) + + @enforce_types + async def delete_by_client_id(self, actor: PydanticClient) -> int: + """Hard delete all conversation turns recorded by a client (erasure path). + + Attribution is via the `_created_by_id` audit column, which + `record_turns` stamps with the ingesting client's id โ€” the table has no + dedicated client_id column. Called from + `ClientManager.delete_memories_by_client_id`. + """ + async with self.session_maker() as session: + stmt = delete(ConversationMessageModel).where( + ConversationMessageModel._created_by_id == actor.id + ) + result = await session.execute(stmt) + await session.commit() + return int(result.rowcount or 0) diff --git a/mirix/services/episodic_memory_manager.py b/mirix/services/episodic_memory_manager.py index 4109b065c..0e1607ee2 100755 --- a/mirix/services/episodic_memory_manager.py +++ b/mirix/services/episodic_memory_manager.py @@ -909,8 +909,19 @@ async def list_episodic_memory( ) else: # Fallback to in-memory BM25 for SQLite (legacy method) - # Load all candidate events (memory-intensive, kept for compatibility) - result = await session.execute(select(EpisodicEvent).where(EpisodicEvent.user_id == user.id)) + # Load all candidate events (memory-intensive, kept for + # compatibility). Mirror base_query's WHERE set (org + + # filter_tags + scopes) so the SQLite fallback enforces + # the SAME scope filtering as the PG/other paths. + candidate_query = ( + select(EpisodicEvent) + .where(EpisodicEvent.user_id == user.id) + .where(EpisodicEvent.organization_id == organization_id) + ) + candidate_query = apply_filter_tags_sqlalchemy( + candidate_query, EpisodicEvent, filter_tags, scopes=scopes + ) + result = await session.execute(candidate_query) all_events = result.scalars().all() # Apply temporal filtering in memory for SQLite @@ -971,8 +982,18 @@ async def list_episodic_memory( return [event.to_pydantic() for event in episodic_memory] elif search_method == "fuzzy_match": - # Load all candidate events (kept for backward compatibility) - result = await session.execute(select(EpisodicEvent).where(EpisodicEvent.user_id == user.id)) + # Load all candidate events (kept for backward compatibility). + # Mirror base_query's WHERE set (org + filter_tags + scopes) + # so fuzzy reads enforce the SAME scope filtering. + candidate_query = ( + select(EpisodicEvent) + .where(EpisodicEvent.user_id == user.id) + .where(EpisodicEvent.organization_id == organization_id) + ) + candidate_query = apply_filter_tags_sqlalchemy( + candidate_query, EpisodicEvent, filter_tags, scopes=scopes + ) + result = await session.execute(candidate_query) all_events = result.scalars().all() scored_events = [] for event in all_events: diff --git a/mirix/services/knowledge_vault_manager.py b/mirix/services/knowledge_vault_manager.py index 7fac83181..babc86ef5 100755 --- a/mirix/services/knowledge_vault_manager.py +++ b/mirix/services/knowledge_vault_manager.py @@ -841,13 +841,24 @@ async def list_knowledge( ) else: # Fallback to in-memory BM25 for SQLite (legacy method) - # Load all candidate items (memory-intensive, kept for compatibility) - fuzzy_query = select(KnowledgeVaultItem).where(KnowledgeVaultItem.user_id == user.id) + # Load all candidate items (memory-intensive, kept for + # compatibility). Mirror base_query's WHERE set (org + + # filter_tags + scopes) so the SQLite fallback enforces + # the SAME scope filtering as the PG/other paths. + fuzzy_query = ( + select(KnowledgeVaultItem) + .where(KnowledgeVaultItem.user_id == user.id) + .where(KnowledgeVaultItem.organization_id == organization_id) + ) # Add sensitivity filter if provided if sensitivity is not None: fuzzy_query = fuzzy_query.where(KnowledgeVaultItem.sensitivity.in_(sensitivity)) + fuzzy_query = apply_filter_tags_sqlalchemy( + fuzzy_query, KnowledgeVaultItem, filter_tags, scopes=scopes + ) + result = await session.execute(fuzzy_query) all_items = result.scalars().all() @@ -905,12 +916,22 @@ async def list_knowledge( elif search_method == "fuzzy_match": # Fuzzy matching: load all candidate items into memory, # then compute fuzzy matching score using RapidFuzz. - fuzzy_query = select(KnowledgeVaultItem).where(KnowledgeVaultItem.user_id == user.id) + # Mirror base_query's WHERE set (org + filter_tags + scopes) + # so fuzzy reads enforce the SAME scope filtering. + fuzzy_query = ( + select(KnowledgeVaultItem) + .where(KnowledgeVaultItem.user_id == user.id) + .where(KnowledgeVaultItem.organization_id == organization_id) + ) # Add sensitivity filter if provided if sensitivity is not None: fuzzy_query = fuzzy_query.where(KnowledgeVaultItem.sensitivity.in_(sensitivity)) + fuzzy_query = apply_filter_tags_sqlalchemy( + fuzzy_query, KnowledgeVaultItem, filter_tags, scopes=scopes + ) + result = await session.execute(fuzzy_query) all_items = result.scalars().all() scored_items = [] diff --git a/mirix/services/message_manager.py b/mirix/services/message_manager.py index 1ccffdb3d..78d92b3e2 100755 --- a/mirix/services/message_manager.py +++ b/mirix/services/message_manager.py @@ -433,6 +433,7 @@ async def list_user_messages_for_agent( filters: Optional[Dict] = None, query_text: Optional[str] = None, ascending: bool = True, + session_id: Optional[str] = None, ) -> List[PydanticMessage]: """List user messages with flexible filtering and pagination options. @@ -443,6 +444,7 @@ async def list_user_messages_for_agent( limit: Maximum number of records to return filters: Additional filters to apply query_text: Optional text to search for in message content + session_id: Optional session id to filter by (exact match on column). Returns: List[PydanticMessage] - List of messages matching the criteria @@ -461,6 +463,7 @@ async def list_user_messages_for_agent( filters=message_filters, query_text=query_text, ascending=ascending, + session_id=session_id, ) @update_timezone @@ -477,6 +480,7 @@ async def list_messages_for_agent( query_text: Optional[str] = None, ascending: bool = True, use_cache: bool = True, + session_id: Optional[str] = None, ) -> List[PydanticMessage]: """List messages with flexible filtering and pagination options. @@ -501,6 +505,8 @@ async def list_messages_for_agent( message_filters.update({"organization_id": actor.organization_id}) if filters: message_filters.update(filters) + if session_id is not None: + message_filters["session_id"] = session_id results = await MessageModel.list( db_session=session, @@ -515,8 +521,52 @@ async def list_messages_for_agent( return [msg.to_pydantic() for msg in results] + @staticmethod + def _select_retained_session_ids( + session_first_seen: Dict[str, datetime], + retain_last_n_sessions: int, + ) -> set: + """Select the most-recent N session_ids to retain (DB-free, pure logic). + + Args: + session_first_seen: mapping of session_id -> MIN(created_at) for that + session. session_id keys must be non-null (NULL session_ids are + never part of any retained session and must be filtered out by the + caller before building this mapping). + retain_last_n_sessions: how many of the most-recent sessions to retain. + <= 0 retains nothing. + + Returns: + set of session_ids that should be retained (the N sessions whose first + message is the most recent). Ordering is by MIN(created_at) DESC, with + session_id as a deterministic tie-breaker so the window is stable. + """ + if retain_last_n_sessions <= 0 or not session_first_seen: + return set() + # Sessions with a known first-seen timestamp: most-recent first, breaking + # ties on session_id (desc) for determinism. Sessions whose timestamp is + # None (no created_at on any row) are treated as OLDEST and ordered after + # the timestamped ones โ€” this also avoids ever comparing None to datetime. + with_ts = sorted( + (kv for kv in session_first_seen.items() if kv[1] is not None), + key=lambda kv: (kv[1], kv[0]), + reverse=True, + ) + without_ts = sorted( + (sid for sid, ts in session_first_seen.items() if ts is None), + reverse=True, + ) + ordered = [sid for sid, _ in with_ts] + list(without_ts) + return set(ordered[:retain_last_n_sessions]) + @enforce_types - async def delete_detached_messages_for_agent(self, agent_id: str, actor: PydanticClient) -> int: + async def delete_detached_messages_for_agent( + self, + agent_id: str, + actor: PydanticClient, + retain_last_n_sessions: int = 0, + user_id: Optional[str] = None, + ) -> int: """ Delete messages that belong to an agent but are not in the agent's current message_ids list. @@ -526,6 +576,18 @@ async def delete_detached_messages_for_agent(self, agent_id: str, actor: Pydanti Args: agent_id: The ID of the agent to clean up messages for actor: The user performing this action + retain_last_n_sessions: When > 0, the raw messages of the most-recent N + sessions (by MIN(created_at) per session, descending) are EXCLUDED + from the hard-delete set even though they are detached from the + agent's message_ids. This preserves a bounded rolling window of raw + transcript so a later distiller / auto-dream pass can read it. The + retained rows are NOT re-added to message_ids / in-context (token + economy preserved). 0 (default) preserves the legacy behavior of + deleting every detached message. + user_id: Optional user scoping for the retained-session window. When + provided (and retain_last_n_sessions > 0), the most-recent-N session + set is computed per (agent, user) so one user's sessions never evict + or count toward another user's retained window on a shared agent. Returns: int: Number of messages deleted @@ -542,7 +604,9 @@ async def delete_detached_messages_for_agent(self, agent_id: str, actor: Pydanti # Get current message_ids (messages that should be kept) current_message_ids = set(agent.message_ids or []) - # Find all messages for this agent + # Find all messages for this agent โ€” a SINGLE snapshot used both to + # compute the retained rolling window and to pick delete candidates, + # so there is no read-read race between the two. all_messages = await MessageModel.list( db_session=session, agent_id=agent_id, @@ -550,8 +614,49 @@ async def delete_detached_messages_for_agent(self, agent_id: str, actor: Pydanti limit=None, # Get all messages ) - # Identify detached messages (not in current message_ids) - detached_messages = [msg for msg in all_messages if msg.id not in current_message_ids] + # Compute the set of session_ids to retain (bounded rolling window). + # When retention is active AND a user_id is given, BOTH the window and + # the delete scope are restricted to that user, so one user's cleanup + # can never evict or delete another user's retained sessions on a + # shared agent. + retained_session_ids: set = set() + scope_to_user = retain_last_n_sessions > 0 and user_id is not None + if retain_last_n_sessions > 0: + # Per-session first-seen timestamp, computed in-memory from the + # snapshot above (no second query). Mirrors agent_trigger_state's + # MIN(created_at)-per-session windowing. A session whose rows all + # lack created_at keeps a None timestamp and is treated as oldest. + session_first_seen: Dict[str, Optional[datetime]] = {} + for msg in all_messages: + sid = msg.session_id + if sid is None: + continue + if scope_to_user and msg.user_id != user_id: + continue + ts = msg.created_at + if sid not in session_first_seen: + session_first_seen[sid] = ts + else: + cur = session_first_seen[sid] + if cur is None: + session_first_seen[sid] = ts + elif ts is not None and ts < cur: + session_first_seen[sid] = ts + retained_session_ids = self._select_retained_session_ids( + session_first_seen, retain_last_n_sessions + ) + + # Identify detached messages (not in current message_ids). Keep any + # message whose session_id is in the retained rolling window, and โ€” + # when retention is user-scoped โ€” only treat THIS user's messages as + # delete candidates so other users' rows are never touched. + detached_messages = [ + msg + for msg in all_messages + if msg.id not in current_message_ids + and msg.session_id not in retained_session_ids + and (not scope_to_user or msg.user_id == user_id) + ] # Delete detached messages (and clean up Redis cache) deleted_count = 0 @@ -575,6 +680,13 @@ async def cleanup_all_detached_messages(self, actor: PydanticClient) -> Dict[str """ Cleanup detached messages for all agents in the organization. + WARNING: this is a destructive, organization-wide purge that does NOT honor + the last-N-session retention window (it hard-deletes every detached message + for every agent). Do NOT point it at the meta_memory_agent if you rely on the + MESSAGE_RETAIN_LAST_N_SESSIONS rolling window for session distillation / + auto-dream โ€” use delete_detached_messages_for_agent(retain_last_n_sessions=...) + for retention-aware cleanup. Kept as an explicit legacy full-purge helper. + Args: actor: The user performing this action diff --git a/mirix/services/procedural_evolution_runtime.py b/mirix/services/procedural_evolution_runtime.py new file mode 100644 index 000000000..17ebf3d8e --- /dev/null +++ b/mirix/services/procedural_evolution_runtime.py @@ -0,0 +1,62 @@ +"""Shared runtime helpers for procedural skill evolution.""" + +from __future__ import annotations + +from typing import Dict, List + +from mirix.helpers.keyed_locks import KeyedLocks + +# Procedural evolution resets a persistent agent's in-context messages before +# and after a step. Serializing per procedural agent prevents concurrent runs +# from deleting each other's transient chain messages. Self-evicting registry; +# use `async with _evolve_locks.acquire(agent_id):`. +_evolve_locks = KeyedLocks() + + +def _diff_skills(before: List, after: List) -> Dict[str, List[str]]: + """Compute created, edited, and deleted skill ids between two snapshots.""" + before_map = {_attr(s, "id"): s for s in (before or [])} + after_map = {_attr(s, "id"): s for s in (after or [])} + before_ids = set(before_map) + after_ids = set(after_map) + + created = sorted(after_ids - before_ids) + deleted = sorted(before_ids - after_ids) + edited = [] + for sid in before_ids & after_ids: + b = before_map[sid] + a = after_map[sid] + if ( + _attr(b, "version") != _attr(a, "version") + or _attr(b, "instructions") != _attr(a, "instructions") + or _attr(b, "description") != _attr(a, "description") + ): + edited.append(sid) + return {"created": created, "edited": sorted(edited), "deleted": deleted} + + +async def reset_agent_in_context_to_system(server, agent_id: str, actor) -> None: + """Force an agent's in-context history back to system-message-only. + + Keeps index 0 (the system message) and deletes detached conversation, tool, + and heartbeat messages so the DB does not accumulate them. This is + idempotent when the context is already system-only or empty. + """ + refreshed = await server.agent_manager.get_agent_by_id( + agent_id=agent_id, actor=actor + ) + msg_ids = refreshed.message_ids or [] + if len(msg_ids) <= 1: + return + await server.agent_manager.set_in_context_messages( + agent_id=agent_id, message_ids=[msg_ids[0]], actor=actor + ) + await server.message_manager.delete_detached_messages_for_agent( + agent_id=agent_id, actor=actor + ) + + +def _attr(obj, key): + if isinstance(obj, dict): + return obj.get(key) + return getattr(obj, key, None) diff --git a/mirix/services/procedural_memory_manager.py b/mirix/services/procedural_memory_manager.py index cda827658..78af8ee8f 100755 --- a/mirix/services/procedural_memory_manager.py +++ b/mirix/services/procedural_memory_manager.py @@ -3,18 +3,28 @@ import string from typing import Any, Dict, List, Optional +import numpy as np from rank_bm25 import BM25Okapi from rapidfuzz import fuzz from sqlalchemy import delete, func, select, text -from mirix.constants import BUILD_EMBEDDINGS_FOR_MEMORY +from sqlalchemy.exc import IntegrityError + +from mirix.constants import ( + BUILD_EMBEDDINGS_FOR_MEMORY, + MAX_EMBEDDING_DIM, + SKILL_HYBRID_RECALL_MULTIPLIER, + SKILL_HYBRID_RRF_K, +) from mirix.embeddings import embedding_model from mirix.log import get_logger -from mirix.orm.errors import NoResultFound +from mirix.orm.errors import NoResultFound, UniqueConstraintViolationError from mirix.orm.procedural_memory import ProceduralMemoryItem from mirix.schemas.agent import AgentState from mirix.schemas.client import Client as PydanticClient -from mirix.schemas.procedural_memory import ProceduralMemoryItem as PydanticProceduralMemoryItem +from mirix.schemas.procedural_memory import ( + ProceduralMemoryItem as PydanticProceduralMemoryItem, +) from mirix.schemas.procedural_memory import ( ProceduralMemoryItemUpdate, ) @@ -26,6 +36,95 @@ logger = get_logger(__name__) +def _pad_to_max(vec): + """Pad a raw embedding to MAX_EMBEDDING_DIM with trailing zeros. + + The CREATE path pads via the PydanticProceduralMemoryItem.pad_embeddings + field_validator, and the vector-search query path pads explicitly + (see the ``search_method == "embedding"`` branch). The UPDATE path + recomputes embeddings into a plain dict that bypasses the schema + validator, so it must pad here too โ€” otherwise a raw (e.g. 3072-dim) + vector is written straight onto the Vector(4096) column and SQLAlchemy's + compare_values raises "operands could not be broadcast together with + shapes (4096,) (N,)" on the next UPDATE. Zero-padding is cosine-preserving + (the appended dims contribute nothing to dot product or norms), so stored + and query vectors stay consistent at MAX_EMBEDDING_DIM. + """ + if vec is None: + return None + arr = np.array(vec) + if arr.shape[0] == MAX_EMBEDDING_DIM: + return list(vec) + return np.pad(arr, (0, MAX_EMBEDDING_DIM - arr.shape[0]), mode="constant").tolist() + + +# Whitelisted searchable text fields and their associated ORM columns. +# Kept as a constant so every search path uses the same mapping โ€” no eval(), +# no silent column fallbacks, and unknown fields can be rejected early. +_SEARCHABLE_TEXT_FIELDS: Dict[str, str] = { + "description": "description", + "instructions": "instructions", + "entry_type": "entry_type", + "name": "name", +} + +_EMBEDDING_FIELDS: Dict[str, str] = { + "description": "description_embedding", + "instructions": "instructions_embedding", +} + + +def _resolve_text_column(search_field: str): + """Return the ORM column for a text search field, or None if not whitelisted.""" + column_name = _SEARCHABLE_TEXT_FIELDS.get(search_field) + if column_name is None: + return None + return getattr(ProceduralMemoryItem, column_name) + + +def _resolve_embedding_column(search_field: str): + """Return the ORM embedding column for a text search field.""" + column_name = _EMBEDDING_FIELDS.get(search_field) + if column_name is None: + return None + return getattr(ProceduralMemoryItem, column_name) + + +def _rrf_fuse( + ranked_lists: List[List[PydanticProceduralMemoryItem]], + *, + k: int, + limit: Optional[int], +) -> List[PydanticProceduralMemoryItem]: + """Fuse several ranked result lists with Reciprocal Rank Fusion. + + Bit-for-bit aligned with EverOS/everalgo (``everalgo.rank.fusion.rrf``): + the fusion is RANK-BASED (raw BM25 ts_rank_cd and cosine scores are + intentionally discarded โ€” only a result's POSITION in each lane feeds the + sum), 1-BASED (the top item in a lane has rank 1), and UNWEIGHTED (every + lane contributes ``1 / (k + rank)`` with no per-lane weight). ``k`` defaults + to the canonical 60. + + Items are keyed by their stable ``id`` so the same skill surfaced by both + the lexical and the dense lane accumulates both contributions (the + hallmark of hybrid retrieval: agreement across lanes ranks higher). Items + with no ``id`` cannot be deduplicated across lanes and are skipped. The + fused list is sliced to ``limit`` (no slice when ``limit`` is falsy). + """ + scores: Dict[str, float] = {} + items_by_id: Dict[str, PydanticProceduralMemoryItem] = {} + for ranked in ranked_lists: + for rank, item in enumerate(ranked, start=1): + item_id = getattr(item, "id", None) + if item_id is None: + continue + scores[item_id] = scores.get(item_id, 0.0) + 1.0 / (k + rank) + items_by_id.setdefault(item_id, item) + ordered_ids = sorted(scores, key=lambda i: scores[i], reverse=True) + fused = [items_by_id[i] for i in ordered_ids] + return fused[:limit] if limit else fused + + class ProceduralMemoryManager: """Manager class to handle business logic related to Procedural Memory Items.""" @@ -74,7 +173,9 @@ def _preprocess_text_for_bm25(self, text: str) -> List[str]: cleaned_text = self._clean_text_for_search(text) # Split into tokens and filter out empty strings and very short tokens - tokens = [token for token in cleaned_text.split() if token.strip() and len(token) > 1] + tokens = [ + token for token in cleaned_text.split() if token.strip() and len(token) > 1 + ] return tokens def _parse_embedding_field(self, embedding_value): @@ -112,11 +213,17 @@ def _parse_embedding_field(self, embedding_value): # Try comma-separated values if "," in embedding_value: - return [float(x.strip()) for x in embedding_value.split(",") if x.strip()] + return [ + float(x.strip()) + for x in embedding_value.split(",") + if x.strip() + ] # Try space-separated values if " " in embedding_value: - return [float(x.strip()) for x in embedding_value.split() if x.strip()] + return [ + float(x.strip()) for x in embedding_value.split() if x.strip() + ] # Try using the original deserialize_vector approach for binary data try: @@ -136,7 +243,9 @@ class MockDialect: logger.error("Warning: Failed to parse embedding field: %s", e) return None - def _count_word_matches(self, item_data: Dict[str, Any], query_words: List[str], search_field: str = "") -> int: + def _count_word_matches( + self, item_data: Dict[str, Any], query_words: List[str], search_field: str = "" + ) -> int: """ Count how many of the query words are present in the procedural memory item data. @@ -152,17 +261,17 @@ def _count_word_matches(self, item_data: Dict[str, Any], query_words: List[str], return 0 # Determine which text fields to search in - if search_field == "summary": - search_texts = [item_data.get("summary", "")] - elif search_field == "steps": - search_texts = [item_data.get("steps", "")] + if search_field == "description": + search_texts = [item_data.get("description", "")] + elif search_field == "instructions": + search_texts = [item_data.get("instructions", "")] elif search_field == "entry_type": search_texts = [item_data.get("entry_type", "")] else: # Search across all relevant text fields search_texts = [ - item_data.get("summary", ""), - item_data.get("steps", ""), + item_data.get("description", ""), + item_data.get("instructions", ""), item_data.get("entry_type", ""), ] @@ -198,7 +307,7 @@ async def _postgresql_fulltext_search( session: Database session base_query: Base SQLAlchemy query (not used, kept for API compatibility) query_text: Search query string - search_field: Field to search in ('summary', 'steps', 'entry_type', etc.) + search_field: Field to search in ('description', 'instructions', 'entry_type', etc.) limit: Maximum number of results to return user: User object to filter by filter_tags: Optional dict of tag key-value pairs to filter by (e.g., {"scope": "CARE"}) @@ -225,7 +334,13 @@ async def _postgresql_fulltext_search( tsquery_parts = [] for word in query_words: # Escape special characters for tsquery - escaped_word = word.replace("'", "''").replace("&", "").replace("|", "").replace("!", "").replace(":", "") + escaped_word = ( + word.replace("'", "''") + .replace("&", "") + .replace("|", "") + .replace("!", "") + .replace(":", "") + ) if escaped_word and len(escaped_word) > 1: # Skip very short words # Add both exact and prefix matching for better results if len(escaped_word) >= 3: @@ -245,27 +360,23 @@ async def _postgresql_fulltext_search( tsquery_string_and = tsquery_string_or = tsquery_parts[0] # Determine which field to search based on search_field - if search_field == "summary": - tsvector_sql = "to_tsvector('english', coalesce(summary, ''))" - rank_sql = "ts_rank_cd(to_tsvector('english', coalesce(summary, '')), to_tsquery('english', :tsquery), 32)" - elif search_field == "steps": - # Convert JSON array to text by removing JSON formatting - tsvector_sql = "to_tsvector('english', coalesce(regexp_replace(steps::text, '[\"\\[\\],]', ' ', 'g'), ''))" - rank_sql = "ts_rank_cd(to_tsvector('english', coalesce(regexp_replace(steps::text, '[\"\\[\\],]', ' ', 'g'), '')), to_tsquery('english', :tsquery), 32)" + if search_field == "description": + tsvector_sql = "to_tsvector('english', coalesce(description, ''))" + rank_sql = "ts_rank_cd(to_tsvector('english', coalesce(description, '')), to_tsquery('english', :tsquery), 32)" + elif search_field == "instructions": + tsvector_sql = "to_tsvector('english', coalesce(instructions, ''))" + rank_sql = "ts_rank_cd(to_tsvector('english', coalesce(instructions, '')), to_tsquery('english', :tsquery), 32)" elif search_field == "entry_type": tsvector_sql = "to_tsvector('english', coalesce(entry_type, ''))" - rank_sql = ( - "ts_rank_cd(to_tsvector('english', coalesce(entry_type, '')), to_tsquery('english', :tsquery), 32)" - ) + rank_sql = "ts_rank_cd(to_tsvector('english', coalesce(entry_type, '')), to_tsquery('english', :tsquery), 32)" else: # Search across all relevant text fields with weighting - # Convert steps JSON array to text by removing JSON formatting - tsvector_sql = """setweight(to_tsvector('english', coalesce(summary, '')), 'A') || - setweight(to_tsvector('english', coalesce(regexp_replace(steps::text, '[\"\\[\\],]', ' ', 'g'), '')), 'B') || + tsvector_sql = """setweight(to_tsvector('english', coalesce(description, '')), 'A') || + setweight(to_tsvector('english', coalesce(instructions, '')), 'B') || setweight(to_tsvector('english', coalesce(entry_type, '')), 'C')""" rank_sql = """ts_rank_cd( - setweight(to_tsvector('english', coalesce(summary, '')), 'A') || - setweight(to_tsvector('english', coalesce(regexp_replace(steps::text, '[\"\\[\\],]', ' ', 'g'), '')), 'B') || + setweight(to_tsvector('english', coalesce(description, '')), 'A') || + setweight(to_tsvector('english', coalesce(instructions, '')), 'B') || setweight(to_tsvector('english', coalesce(entry_type, '')), 'C'), to_tsquery('english', :tsquery), 32)""" @@ -273,6 +384,9 @@ async def _postgresql_fulltext_search( where_clauses = [ f"{tsvector_sql} @@ to_tsquery('english', :tsquery)", "user_id = :user_id", + # Exclude soft-deleted skills from the native PG full-text path too, + # so a curator soft-delete is excluded from retrieval everywhere. + "is_deleted = false", ] query_params = { "tsquery": tsquery_string_and, @@ -286,128 +400,99 @@ async def _postgresql_fulltext_search( where_clause = " AND ".join(where_clauses) - # Try AND query first for more precise results - try: - and_query_sql = text( + async def _fetch_ordered_items(tsquery_value): + # Rank via raw SQL for ts_rank_cd, but fetch only IDs to avoid + # hydrating partial ORM instances (which silently drops triggers, + # examples, version, filter_tags, etc.). Full skill rows are then + # loaded through SQLAlchemy so every schema field is present. + params = {**query_params, "tsquery": tsquery_value} + id_query_sql = text( f""" - SELECT - id, created_at, entry_type, summary, steps, - steps_embedding, summary_embedding, embedding_config, - organization_id, last_modify, user_id, - {rank_sql} as rank_score + SELECT id, {rank_sql} as rank_score FROM procedural_memory WHERE {where_clause} ORDER BY rank_score DESC, created_at DESC LIMIT :limit_val - """ + """ ) + rows = (await session.execute(id_query_sql, params)).all() + ordered_ids = [row.id for row in rows] + if not ordered_ids: + return [] + items_result = await session.execute( + select(ProceduralMemoryItem).where( + ProceduralMemoryItem.id.in_(ordered_ids) + ) + ) + items_by_id = {item.id: item for item in items_result.scalars().all()} + return [ + items_by_id[item_id] + for item_id in ordered_ids + if item_id in items_by_id + ] - result = await session.execute(and_query_sql, query_params) - results = result.all() - - # If AND query returns sufficient results, use them - if len(results) >= min(limit or 10, 10): - procedures = [] - for row in results: - data = dict(row._mapping) - # Remove the rank_score field before creating the object - data.pop("rank_score", None) - - # Parse JSON fields that are returned as strings from raw SQL - json_fields = ["last_modify", "embedding_config"] - for field in json_fields: - if field in data and isinstance(data[field], str): - try: - data[field] = json.loads(data[field]) - except (json.JSONDecodeError, TypeError): - pass - - # Parse embedding fields - embedding_fields = ["steps_embedding", "summary_embedding"] - for field in embedding_fields: - if field in data and data[field] is not None: - data[field] = self._parse_embedding_field(data[field]) - - procedures.append(ProceduralMemoryItem(**data)) - + # Try AND query first for more precise results. + try: + procedures = await _fetch_ordered_items(tsquery_string_and) + if len(procedures) >= min(limit or 10, 10): return [procedure.to_pydantic() for procedure in procedures] - except Exception as e: logger.error("PostgreSQL AND query error: %s", e) - # If AND query fails or returns too few results, try OR query + # If AND query fails or returns too few results, try OR query. try: - # Update query params for OR query - or_query_params = query_params.copy() - or_query_params["tsquery"] = tsquery_string_or - - or_query_sql = text( - f""" - SELECT - id, created_at, entry_type, summary, steps, - steps_embedding, summary_embedding, embedding_config, - organization_id, last_modify, user_id, - {rank_sql} as rank_score - FROM procedural_memory - WHERE {where_clause} - ORDER BY rank_score DESC, created_at DESC - LIMIT :limit_val - """ - ) - - results = await session.execute(or_query_sql, or_query_params) - - procedures = [] - for row in results: - data = dict(row._mapping) - # Remove the rank_score field before creating the object - data.pop("rank_score", None) - - # Parse JSON fields that are returned as strings from raw SQL - json_fields = ["last_modify", "embedding_config"] - for field in json_fields: - if field in data and isinstance(data[field], str): - try: - data[field] = json.loads(data[field]) - except (json.JSONDecodeError, TypeError): - pass - - # Parse embedding fields - embedding_fields = ["steps_embedding", "summary_embedding"] - for field in embedding_fields: - if field in data and data[field] is not None: - data[field] = self._parse_embedding_field(data[field]) - - procedures.append(ProceduralMemoryItem(**data)) - + procedures = await _fetch_ordered_items(tsquery_string_or) return [procedure.to_pydantic() for procedure in procedures] - except Exception as e: # If there's an error with the tsquery, fall back to simpler search logger.error("PostgreSQL full-text search error: %s", e) - # Fall back to simple ILIKE search - fallback_field = ( - getattr(ProceduralMemoryItem, search_field) - if search_field and hasattr(ProceduralMemoryItem, search_field) - else ProceduralMemoryItem.summary + # Fall back to simple ILIKE search on a whitelisted field + fallback_column = ( + _resolve_text_column(search_field) or ProceduralMemoryItem.description ) - fallback_query = base_query.where(func.lower(fallback_field).contains(query_text.lower())).order_by( - ProceduralMemoryItem.created_at.desc() + fallback_query = ( + select(ProceduralMemoryItem) + .where(ProceduralMemoryItem.user_id == user.id) + # Exclude soft-deleted skills on the PG FTS error fallback too. + .where(~ProceduralMemoryItem.is_deleted) + .where(func.lower(fallback_column).contains(query_text.lower())) + .order_by(ProceduralMemoryItem.created_at.desc()) + ) + + from mirix.database.filter_tags_query import apply_filter_tags_sqlalchemy + + fallback_query = apply_filter_tags_sqlalchemy( + fallback_query, ProceduralMemoryItem, filter_tags, scopes=scopes ) if limit: fallback_query = fallback_query.limit(limit) - results = await session.execute(fallback_query) - procedures = [ProceduralMemoryItem(**dict(row._mapping)) for row in results] + result = await session.execute(fallback_query) + procedures = result.scalars().all() return [procedure.to_pydantic() for procedure in procedures] @update_timezone @enforce_types async def get_item_by_id( - self, item_id: str, user: PydanticUser, timezone_str: str + self, + item_id: str, + user: PydanticUser, + timezone_str: str, + actor: Optional[PydanticClient] = None, ) -> Optional[PydanticProceduralMemoryItem]: - """Fetch a procedural memory item by ID (with cache - Redis or IPS Cache).""" + """Fetch a procedural memory item by ID (with cache - Redis or IPS Cache). + + Access control (see ``delete_procedure_by_id`` for the same pattern): + + * ``actor`` โ†’ organization scope via ``apply_access_predicate``. + * ``user`` โ†’ explicit post-load ``user_id`` equality check. + + Passing only ``user`` is not enough, because ``SqlalchemyBase.read`` + only runs the access predicate when ``actor`` is truthy. The cache + hit path is gated the same way so a stale cache entry for another + user cannot leak via this endpoint either. + """ cache_provider = None try: from mirix.database.cache_provider import get_cache_provider @@ -418,29 +503,50 @@ async def get_item_by_id( cache_key = f"{cache_provider.PROCEDURAL_PREFIX}{item_id}" cached_data = await cache_provider.get_json(cache_key) if cached_data: + # Cache is keyed by id only, so re-check ownership before + # returning โ€” otherwise a same-org attacker could fetch + # someone else's cached skill. + if cached_data.get("user_id") != user.id: + raise NoResultFound( + f"Procedural memory item with id {item_id} not found." + ) logger.debug("Cache HIT for procedural memory %s", item_id) return PydanticProceduralMemoryItem(**cached_data) + except NoResultFound: + raise except Exception as e: logger.warning("Cache read failed for procedural memory %s: %s", item_id, e) async with self.session_maker() as session: try: - item = await ProceduralMemoryItem.read(db_session=session, identifier=item_id, user=user) - pydantic_item = item.to_pydantic() + item = await ProceduralMemoryItem.read( + db_session=session, identifier=item_id, actor=actor + ) + except NoResultFound: + raise NoResultFound( + f"Procedural memory item with id {item_id} not found." + ) - try: - if cache_provider: - from mirix.settings import settings + if getattr(item, "user_id", None) != user.id: + raise NoResultFound( + f"Procedural memory item with id {item_id} not found." + ) - cache_key = f"{cache_provider.PROCEDURAL_PREFIX}{item_id}" - data = pydantic_item.model_dump(mode="json") - await cache_provider.set_json(cache_key, data, ttl=settings.redis_ttl_default) - except Exception as e: - logger.warning("Failed to populate cache: %s", e) + pydantic_item = item.to_pydantic() - return pydantic_item - except NoResultFound: - raise NoResultFound(f"Procedural memory item with id {item_id} not found.") + try: + if cache_provider: + from mirix.settings import settings + + cache_key = f"{cache_provider.PROCEDURAL_PREFIX}{item_id}" + data = pydantic_item.model_dump(mode="json") + await cache_provider.set_json( + cache_key, data, ttl=settings.redis_ttl_default + ) + except Exception as e: + logger.warning("Failed to populate cache: %s", e) + + return pydantic_item @update_timezone @enforce_types @@ -457,7 +563,9 @@ async def get_most_recently_updated_item( from sqlalchemy import DateTime, cast, text query = select(ProceduralMemoryItem).order_by( - cast(text("procedural_memory.last_modify ->> 'timestamp'"), DateTime).desc() + cast( + text("procedural_memory.last_modify ->> 'timestamp'"), DateTime + ).desc() ) # Filter by user_id for multi-user support @@ -502,10 +610,12 @@ async def create_item( data_dict = item_data.model_dump() # Validate required fields - required_fields = ["entry_type"] + required_fields = ["entry_type", "name"] for field in required_fields: if field not in data_dict: - raise ValueError(f"Required field '{field}' missing from procedural memory data") + raise ValueError( + f"Required field '{field}' missing from procedural memory data" + ) # Set client_id and user_id on the memory data_dict["client_id"] = client_id @@ -515,7 +625,16 @@ async def create_item( async with self.session_maker() as session: item = ProceduralMemoryItem(**data_dict) - await item.create_with_redis(session, actor=actor, use_cache=use_cache) + try: + await item.create_with_redis(session, actor=actor, use_cache=use_cache) + except IntegrityError as exc: + # uq_procedural_memory_org_user_name โ€” surface a domain-level + # duplicate error so callers don't need to parse SQLSTATE. + if "uq_procedural_memory_org_user_name" in str(exc.orig): + raise UniqueConstraintViolationError( + f"Skill with name '{data_dict.get('name')}' already exists for this user." + ) from exc + raise return item.to_pydantic() @enforce_types @@ -524,16 +643,76 @@ async def update_item( item_update: ProceduralMemoryItemUpdate, user: PydanticUser, actor: PydanticClient, + agent_state: Optional[AgentState] = None, ) -> PydanticProceduralMemoryItem: - """Update an existing procedural memory item.""" + """Update an existing procedural memory item. + + When description/instructions change and agent_state is provided with + embedding_config, the corresponding embeddings are recomputed so that + vector search does not serve stale vectors against fresh text. Callers + that cannot supply agent_state can pre-compute embeddings themselves + and put them on item_update directly. + """ + update_data = item_update.model_dump(exclude_unset=True) + + # Recompute embeddings for text fields that changed, when we have a + # config and the caller didn't already supply embeddings explicitly. + if BUILD_EMBEDDINGS_FOR_MEMORY and agent_state is not None: + embedding_config = agent_state.embedding_config + embed_model = None + if ( + "description" in update_data + and "description_embedding" not in update_data + ): + embed_model = await embedding_model(embedding_config) + update_data["description_embedding"] = _pad_to_max( + await embed_model.get_text_embedding( + update_data["description"] or "" + ) + ) + if ( + "instructions" in update_data + and "instructions_embedding" not in update_data + ): + if embed_model is None: + embed_model = await embedding_model(embedding_config) + update_data["instructions_embedding"] = _pad_to_max( + await embed_model.get_text_embedding( + update_data["instructions"] or "" + ) + ) + if embed_model is not None and "embedding_config" not in update_data: + update_data["embedding_config"] = embedding_config + async with self.session_maker() as session: - item = await ProceduralMemoryItem.read(db_session=session, identifier=item_update.id, user=user) - update_data = item_update.model_dump(exclude_unset=True) + # Same two-step auth as delete_procedure_by_id/get_item_by_id: + # actor enables org-scope via apply_access_predicate, and the + # explicit user_id check guarantees per-user ownership even + # when the base read does not filter on user. + try: + item = await ProceduralMemoryItem.read( + db_session=session, identifier=item_update.id, actor=actor + ) + except NoResultFound: + raise NoResultFound( + f"Procedural memory item with id {item_update.id} not found." + ) + + if getattr(item, "user_id", None) != user.id: + raise NoResultFound( + f"Procedural memory item with id {item_update.id} not found." + ) + for k, v in update_data.items(): - if k not in ["id", "updated_at"]: # Exclude updated_at - handled by update() method + if k not in [ + "id", + "updated_at", + ]: # Exclude updated_at - handled by update() method setattr(item, k, v) # updated_at is automatically set to current UTC time by item.update() - await item.update_with_redis(session, actor=actor) # Updates Redis JSON cache + await item.update_with_redis( + session, actor=actor + ) # Updates Redis JSON cache return item.to_pydantic() @enforce_types @@ -548,10 +727,139 @@ async def create_many_items( async def get_total_number_of_items(self, user: PydanticUser) -> int: """Get the total number of items in the procedural memory for the user.""" async with self.session_maker() as session: - query = select(func.count(ProceduralMemoryItem.id)).where(ProceduralMemoryItem.user_id == user.id) + query = select(func.count(ProceduralMemoryItem.id)).where( + ProceduralMemoryItem.user_id == user.id + ) result = await session.execute(query) return result.scalar_one() + async def _hybrid_search_for_user( + self, + *, + agent_state: AgentState, + user: PydanticUser, + query: str, + embedded_text: Optional[List[float]], + search_field: str, + limit: Optional[int], + timezone_str: Optional[str], + filter_tags: Optional[dict], + scopes: Optional[List[str]], + use_cache: bool, + similarity_threshold: Optional[float], + ) -> List[PydanticProceduralMemoryItem]: + """EverOS-aligned hybrid retrieval for :meth:`list_procedures`. + + Reuses the EXISTING bm25 (lexical) and embedding (dense) branches as two + independent lanes โ€” so all backend dispatch (PG ts_rank_cd, SQLite + BM25Okapi, Redis KNN, pgvector cosine) is inherited, not reimplemented โ€” + then fuses them with RRF (see :func:`_rrf_fuse`). + + The lanes run SEQUENTIALLY, not via ``asyncio.gather``. Each recursive + call opens its own short-lived session, so gathering would require two + simultaneous pooled connections on what can be a hot path (or, if forced + onto one session, would violate SQLAlchemy's "one operation per + AsyncSession at a time" rule and raise). The dense lane's latency is + dominated by the embedding API round-trip, not the indexed DB query, so + overlapping the two buys almost nothing; sequential is async-native and + uses a single pooled connection at a time. + """ + effective_limit = limit or 50 + recall_limit = max( + effective_limit * SKILL_HYBRID_RECALL_MULTIPLIER, effective_limit + ) + # Lexical lane โ€” no similarity_threshold (a dense-only quality floor). + sparse = await self.list_procedures( + agent_state=agent_state, + user=user, + query=query, + search_field=search_field, + search_method="bm25", + limit=recall_limit, + timezone_str=timezone_str, + filter_tags=filter_tags, + scopes=scopes, + use_cache=use_cache, + ) + # Dense lane โ€” embedding search needs a valid embedding column, so fall + # back to "description" when search_field is blank/non-embeddable. + dense_field = ( + search_field if search_field in _EMBEDDING_FIELDS else "description" + ) + dense = await self.list_procedures( + agent_state=agent_state, + user=user, + query=query, + embedded_text=embedded_text, + search_field=dense_field, + search_method="embedding", + limit=recall_limit, + timezone_str=timezone_str, + filter_tags=filter_tags, + scopes=scopes, + use_cache=use_cache, + similarity_threshold=similarity_threshold, + ) + return _rrf_fuse( + [sparse, dense], k=SKILL_HYBRID_RRF_K, limit=effective_limit + ) + + async def _hybrid_search_for_org( + self, + *, + agent_state: AgentState, + organization_id: str, + query: str, + embedded_text: Optional[List[float]], + search_field: str, + limit: Optional[int], + timezone_str: Optional[str], + filter_tags: Optional[dict], + scopes: Optional[List[str]], + use_cache: bool, + similarity_threshold: Optional[float], + ) -> List[PydanticProceduralMemoryItem]: + """Org-wide counterpart of :meth:`_hybrid_search_for_user`. + + Same RRF fusion, but each lane reuses :meth:`list_procedures_by_org`. + """ + effective_limit = limit or 50 + recall_limit = max( + effective_limit * SKILL_HYBRID_RECALL_MULTIPLIER, effective_limit + ) + sparse = await self.list_procedures_by_org( + agent_state=agent_state, + organization_id=organization_id, + query=query, + search_field=search_field, + search_method="bm25", + limit=recall_limit, + timezone_str=timezone_str, + filter_tags=filter_tags, + scopes=scopes, + use_cache=use_cache, + ) + dense_field = ( + search_field if search_field in _EMBEDDING_FIELDS else "description" + ) + dense = await self.list_procedures_by_org( + agent_state=agent_state, + organization_id=organization_id, + query=query, + embedded_text=embedded_text, + search_field=dense_field, + search_method="embedding", + limit=recall_limit, + timezone_str=timezone_str, + filter_tags=filter_tags, + scopes=scopes, + use_cache=use_cache, + similarity_threshold=similarity_threshold, + ) + return _rrf_fuse( + [sparse, dense], k=SKILL_HYBRID_RRF_K, limit=effective_limit + ) + @update_timezone @enforce_types async def list_procedures( @@ -576,11 +884,16 @@ async def list_procedures( agent_state: The agent state containing embedding configuration query: Search query string embedded_text: Pre-computed embedding for semantic search - search_field: Field to search in ('summary', 'steps', 'entry_type') + search_field: Field to search in ('description', 'instructions', 'entry_type') search_method: Search method to use: + - 'hybrid': **RECOMMENDED for skill retrieval** - EverOS-aligned + fusion of the 'bm25' and 'embedding' lanes via + Reciprocal Rank Fusion (rank-based, k=60). Best + recall/precision for skill lookup; self-embeds the + query, so no embedded_text need be supplied. - 'embedding': Vector similarity search using embeddings - 'string_match': Simple string containment search - - 'bm25': **RECOMMENDED** - PostgreSQL native full-text search (ts_rank_cd) when using PostgreSQL, + - 'bm25': PostgreSQL native full-text search (ts_rank_cd) when using PostgreSQL, falls back to in-memory BM25 for SQLite - 'fuzzy_match': Fuzzy string matching (legacy, kept for compatibility) limit: Maximum number of results to return @@ -605,9 +918,36 @@ async def list_procedures( query = query.strip() if query else "" is_empty_query = not query or query == "" + # Pad any caller-supplied embedding to MAX_EMBEDDING_DIM at the point of + # CONSUMPTION (not just where we recompute it), so a raw provider vector + # (e.g. 3072-dim) passed in by REST callers never reaches the Vector(4096) + # PG column or the DIM=4096 Redis index unpadded. + embedded_text = _pad_to_max(embedded_text) + # Extract organization_id from user for multi-tenant isolation organization_id = user.organization_id + # Hybrid retrieval (EverOS-aligned BM25 + embedding fused via RRF) is + # intercepted at the TOP โ€” before the cache and PG/session branches โ€” so + # it always routes through the two-lane fusion and is never partially + # served by a single-method cache path. Each lane (a recursive call) + # still uses the cache for its own bm25/embedding sub-query. Empty + # queries fall through to the recent-items path, where method is moot. + if search_method == "hybrid" and not is_empty_query: + return await self._hybrid_search_for_user( + agent_state=agent_state, + user=user, + query=query, + embedded_text=embedded_text, + search_field=search_field, + limit=limit, + timezone_str=timezone_str, + filter_tags=filter_tags, + scopes=scopes, + use_cache=use_cache, + similarity_threshold=similarity_threshold, + ) + # Try Redis Search first (if cache enabled and Redis is available) from mirix.database.redis_client import get_redis_client @@ -617,7 +957,10 @@ async def list_procedures( try: # Case 1: No query - get recent items (regardless of search_method) if is_empty_query: - logger.debug("Searching cache for recent procedural items with filter_tags=%s", filter_tags) + logger.debug( + "Searching cache for recent procedural items with filter_tags=%s", + filter_tags, + ) results = await redis_client.search_recent( index_name=redis_client.PROCEDURAL_INDEX, limit=limit or 50, @@ -626,12 +969,19 @@ async def list_procedures( filter_tags=filter_tags, scopes=scopes, ) - logger.debug("Cache search returned %d results", len(results) if results else 0) + logger.debug( + "Cache search returned %d results", + len(results) if results else 0, + ) if results: - logger.debug("Cache HIT: returned %d procedural items", len(results)) + logger.debug( + "Cache HIT: returned %d procedural items", len(results) + ) # Clean Redis-specific fields before Pydantic validation results = redis_client.clean_redis_fields(results) - return [PydanticProceduralMemoryItem(**item) for item in results] + return [ + PydanticProceduralMemoryItem(**item) for item in results + ] # If no results, fall through to PostgreSQL (don't return empty list) # Case 2: Vector similarity search @@ -642,7 +992,9 @@ async def list_procedures( from mirix.constants import MAX_EMBEDDING_DIM from mirix.embeddings import embedding_model - embed_model = await embedding_model(agent_state.embedding_config) + embed_model = await embedding_model( + agent_state.embedding_config + ) embedded_text = await embed_model.get_text_embedding(query) embedded_text = np.array(embedded_text) embedded_text = np.pad( @@ -651,7 +1003,11 @@ async def list_procedures( mode="constant", ).tolist() - vector_field = f"{search_field}_embedding" if search_field else "summary_embedding" + vector_field = ( + f"{search_field}_embedding" + if search_field + else "description_embedding" + ) results = await redis_client.search_vector( index_name=redis_client.PROCEDURAL_INDEX, @@ -664,13 +1020,22 @@ async def list_procedures( scopes=scopes, ) if results: - logger.debug("Cache vector search HIT: found %d procedural items", len(results)) + logger.debug( + "Cache vector search HIT: found %d procedural items", + len(results), + ) # Clean Redis-specific fields before Pydantic validation results = redis_client.clean_redis_fields(results) - return [PydanticProceduralMemoryItem(**item) for item in results] + return [ + PydanticProceduralMemoryItem(**item) for item in results + ] elif search_method in ["bm25", "string_match"]: - fields = [search_field] if search_field else ["summary", "steps"] + fields = ( + [search_field] + if search_field + else ["description", "instructions"] + ) results = await redis_client.search_text( index_name=redis_client.PROCEDURAL_INDEX, @@ -683,19 +1048,31 @@ async def list_procedures( scopes=scopes, ) if results: - logger.debug("Cache text search HIT: found %d procedural items", len(results)) + logger.debug( + "Cache text search HIT: found %d procedural items", + len(results), + ) # Clean Redis-specific fields before Pydantic validation results = redis_client.clean_redis_fields(results) - return [PydanticProceduralMemoryItem(**item) for item in results] + return [ + PydanticProceduralMemoryItem(**item) for item in results + ] except Exception as e: - logger.warning("Cache search failed for procedural memory, falling back to PostgreSQL: %s", e) + logger.warning( + "Cache search failed for procedural memory, falling back to PostgreSQL: %s", + e, + ) # Log when bypassing cache or Redis unavailable if not use_cache: - logger.debug("Bypassing cache (use_cache=False), querying PostgreSQL directly for procedural memory") + logger.debug( + "Bypassing cache (use_cache=False), querying PostgreSQL directly for procedural memory" + ) elif not redis_client: - logger.debug("Cache unavailable, querying PostgreSQL directly for procedural memory") + logger.debug( + "Cache unavailable, querying PostgreSQL directly for procedural memory" + ) async with self.session_maker() as session: if query == "": @@ -703,10 +1080,18 @@ async def list_procedures( select(ProceduralMemoryItem) .where(ProceduralMemoryItem.user_id == user.id) .where(ProceduralMemoryItem.organization_id == organization_id) + # Exclude soft-deleted skills (curator soft-delete sets + # is_deleted=True; they must NOT surface in retrieval or the + # evolve before/after snapshots). The base SqlalchemyBase + # reads filter ~is_deleted, but these raw fallback queries + # build select() directly and would otherwise leak them. + .where(~ProceduralMemoryItem.is_deleted) .order_by(ProceduralMemoryItem.created_at.desc()) ) - from mirix.database.filter_tags_query import apply_filter_tags_sqlalchemy + from mirix.database.filter_tags_query import ( + apply_filter_tags_sqlalchemy, + ) query_stmt = apply_filter_tags_sqlalchemy( query_stmt, ProceduralMemoryItem, filter_tags, scopes=scopes @@ -719,60 +1104,55 @@ async def list_procedures( return [event.to_pydantic() for event in procedural_memory] else: + # Select full ORM entities so every schema field (triggers, + # examples, version, filter_tags, updated_at, client_id, ...) + # survives the round-trip. Projecting a subset and re-hydrating + # silently loses columns. base_query = ( - select( - ProceduralMemoryItem.id.label("id"), - ProceduralMemoryItem.created_at.label("created_at"), - ProceduralMemoryItem.entry_type.label("entry_type"), - ProceduralMemoryItem.summary.label("summary"), - ProceduralMemoryItem.steps.label("steps"), - ProceduralMemoryItem.steps_embedding.label("steps_embedding"), - ProceduralMemoryItem.summary_embedding.label("summary_embedding"), - ProceduralMemoryItem.embedding_config.label("embedding_config"), - ProceduralMemoryItem.organization_id.label("organization_id"), - ProceduralMemoryItem.last_modify.label("last_modify"), - ProceduralMemoryItem.user_id.label("user_id"), - ProceduralMemoryItem.agent_id.label("agent_id"), - ) + select(ProceduralMemoryItem) .where(ProceduralMemoryItem.user_id == user.id) .where(ProceduralMemoryItem.organization_id == organization_id) + # Exclude soft-deleted skills from every search method too + # (same reason as the empty-query branch above). + .where(~ProceduralMemoryItem.is_deleted) ) - from mirix.database.filter_tags_query import apply_filter_tags_sqlalchemy + from mirix.database.filter_tags_query import ( + apply_filter_tags_sqlalchemy, + ) base_query = apply_filter_tags_sqlalchemy( base_query, ProceduralMemoryItem, filter_tags, scopes=scopes ) if search_method == "embedding": + embedding_column = _resolve_embedding_column(search_field) + if embedding_column is None: + raise ValueError( + f"Invalid search_field '{search_field}' for embedding search. " + f"Allowed: {sorted(_EMBEDDING_FIELDS)}." + ) main_query = await build_query( base_query=base_query, query_text=query, embedded_text=embedded_text, embed_query=True, embedding_config=agent_state.embedding_config, - search_field=eval("ProceduralMemoryItem." + search_field + "_embedding"), + search_field=embedding_column, target_class=ProceduralMemoryItem, similarity_threshold=similarity_threshold, ) elif search_method == "string_match": - if search_field == "steps": - # For JSON array field, convert to text first and then search - from sqlalchemy import text - - if settings.mirix_pg_uri_no_default: - # PostgreSQL: use regexp_replace and ::text casting - search_condition = text( - "lower(regexp_replace(steps::text, '[\"\\[\\],]', ' ', 'g')) LIKE lower(:query)" - ) - else: - # SQLite: use simpler text conversion without regexp_replace or ::text - search_condition = text("lower(steps) LIKE lower(:query)") - main_query = base_query.where(search_condition).params(query=f"%{query}%") - else: - search_field_obj = eval("ProceduralMemoryItem." + search_field) - main_query = base_query.where(func.lower(search_field_obj).contains(query.lower())) + search_field_obj = _resolve_text_column(search_field) + if search_field_obj is None: + raise ValueError( + f"Invalid search_field '{search_field}' for string_match. " + f"Allowed: {sorted(_SEARCHABLE_TEXT_FIELDS)}." + ) + main_query = base_query.where( + func.lower(search_field_obj).contains(query.lower()) + ) elif search_method == "bm25": # Check if we're using PostgreSQL - use native full-text search if available @@ -790,10 +1170,30 @@ async def list_procedures( ) else: # Fallback to in-memory BM25 for SQLite (legacy method) - # Load all candidate items (memory-intensive, kept for compatibility) - result = await session.execute( - select(ProceduralMemoryItem).where(ProceduralMemoryItem.user_id == user.id) + # Load all candidate items (memory-intensive, kept for + # compatibility). Mirror base_query's full WHERE set + # (org + filter_tags + scopes) so the SQLite fallback + # enforces the SAME scope filtering as the PG/other + # paths โ€” otherwise scoped reads leak NULL-scope rows. + candidate_query = ( + select(ProceduralMemoryItem) + .where(ProceduralMemoryItem.user_id == user.id) + .where( + ProceduralMemoryItem.organization_id + == organization_id + ) + # Exclude soft-deleted skills from the SQLite BM25 + # fallback too (this is the path the validity run's + # before/after evolve snapshot uses). + .where(~ProceduralMemoryItem.is_deleted) ) + candidate_query = apply_filter_tags_sqlalchemy( + candidate_query, + ProceduralMemoryItem, + filter_tags, + scopes=scopes, + ) + result = await session.execute(candidate_query) all_items = result.scalars().all() if not all_items: @@ -805,24 +1205,19 @@ async def list_procedures( for item in all_items: # Determine which field to use for search - if search_field == "summary": - text_to_search = item.summary or "" - elif search_field == "steps": - # Convert JSON array to text for searching - if isinstance(item.steps, list): - text_to_search = " ".join(item.steps) - else: - text_to_search = str(item.steps) if item.steps else "" + if search_field == "description": + text_to_search = item.description or "" + elif search_field == "instructions": + text_to_search = item.instructions or "" elif search_field == "entry_type": text_to_search = item.entry_type or "" else: # Default to searching across all text fields - texts = [item.summary or "", item.entry_type or ""] - # Handle steps field conversion - if isinstance(item.steps, list): - texts.append(" ".join(item.steps)) - else: - texts.append(str(item.steps) if item.steps else "") + texts = [ + item.description or "", + item.entry_type or "", + item.instructions or "", + ] text_to_search = " ".join(texts) # Preprocess the text into tokens @@ -863,9 +1258,24 @@ async def list_procedures( elif search_method == "fuzzy_match": # For fuzzy matching, load all candidate items into memory. - result = await session.execute( - select(ProceduralMemoryItem).where(ProceduralMemoryItem.user_id == user.id) + # Mirror base_query's full WHERE set (org + filter_tags + + # scopes) so fuzzy reads enforce the SAME scope filtering. + candidate_query = ( + select(ProceduralMemoryItem) + .where(ProceduralMemoryItem.user_id == user.id) + .where( + ProceduralMemoryItem.organization_id == organization_id + ) + # Exclude soft-deleted skills from the fuzzy path too. + .where(~ProceduralMemoryItem.is_deleted) ) + candidate_query = apply_filter_tags_sqlalchemy( + candidate_query, + ProceduralMemoryItem, + filter_tags, + scopes=scopes, + ) + result = await session.execute(candidate_query) all_items = result.scalars().all() scored_items = [] for item in all_items: @@ -873,11 +1283,13 @@ async def list_procedures( if search_field and hasattr(item, search_field): text_to_search = getattr(item, search_field) else: - text_to_search = item.summary + text_to_search = item.description # Compute a fuzzy matching score using partial_ratio, # which is suited for comparing a short query to longer text. - score = fuzz.partial_ratio(query.lower(), text_to_search.lower()) + score = fuzz.partial_ratio( + query.lower(), text_to_search.lower() + ) scored_items.append((score, item)) # Sort items by score in descending order and select the top ones. @@ -889,12 +1301,7 @@ async def list_procedures( main_query = main_query.limit(limit) result = await session.execute(main_query) - results = result.all() - - procedures = [] - for row in results: - data = dict(row._mapping) - procedures.append(ProceduralMemoryItem(**data)) + procedures = result.scalars().all() return [procedure.to_pydantic() for procedure in procedures] @@ -903,11 +1310,15 @@ async def insert_procedure( self, agent_state: AgentState, agent_id: str, + name: str, + description: str, + instructions: str, entry_type: str, - summary: Optional[str], - steps: List[str], actor: PydanticClient, organization_id: str, + triggers: Optional[List[str]] = None, + examples: Optional[List[dict]] = None, + version: str = "0.1.0", filter_tags: Optional[dict] = None, use_cache: bool = True, user_id: Optional[str] = None, @@ -917,12 +1328,16 @@ async def insert_procedure( if BUILD_EMBEDDINGS_FOR_MEMORY: # TODO: need to check if we need to chunk the text embed_model = await embedding_model(agent_state.embedding_config) - summary_embedding = await embed_model.get_text_embedding(summary) - steps_embedding = await embed_model.get_text_embedding("\n".join(steps)) + description_embedding = await embed_model.get_text_embedding( + description + ) + instructions_embedding = await embed_model.get_text_embedding( + instructions + ) embedding_config = agent_state.embedding_config else: - summary_embedding = None - steps_embedding = None + description_embedding = None + instructions_embedding = None embedding_config = None # Set client_id from actor, user_id with fallback to DEFAULT_USER_ID @@ -934,14 +1349,18 @@ async def insert_procedure( procedure = await self.create_item( item_data=PydanticProceduralMemoryItem( + name=name, entry_type=entry_type, - summary=summary, - steps=steps, + description=description, + instructions=instructions, + triggers=triggers or [], + examples=examples or [], + version=version, user_id=user_id, agent_id=agent_id, organization_id=organization_id, - summary_embedding=summary_embedding, - steps_embedding=steps_embedding, + description_embedding=description_embedding, + instructions_embedding=instructions_embedding, embedding_config=embedding_config, filter_tags=filter_tags, ), @@ -955,21 +1374,96 @@ async def insert_procedure( except Exception as e: raise e - async def delete_procedure_by_id(self, procedure_id: str, actor: PydanticClient) -> None: - """Delete a procedural memory item by ID (removes from cache).""" + async def delete_procedure_by_id( + self, + procedure_id: str, + actor: PydanticClient, + user: Optional[PydanticUser] = None, + ) -> None: + """Delete a procedural memory item by ID (removes from cache). + + Ownership is enforced on three axes so no code path gets weaker checks + than another: + + * ``actor`` โ†’ organization-scope via ``apply_access_predicate`` (the + base ``read`` only runs access checks when ``actor`` is truthy). + * ``user`` โ†’ explicit ``user_id`` equality check after the row is + loaded. Relying on ``AccessType.USER`` alone is not enough because + callers historically passed only ``user`` without ``actor``, which + silently disables all access filtering in ``SqlalchemyBase.read``. + The post-load check is the belt-and-braces version. + """ async with self.session_maker() as session: try: - item = await ProceduralMemoryItem.read(db_session=session, identifier=procedure_id, actor=actor) - # Remove from cache - from mirix.database.cache_provider import get_cache_provider + item = await ProceduralMemoryItem.read( + db_session=session, + identifier=procedure_id, + actor=actor, + ) + except NoResultFound: + raise NoResultFound( + f"Procedural memory item with id {procedure_id} not found." + ) - cache_provider = get_cache_provider() - if cache_provider: - cache_key = f"{cache_provider.PROCEDURAL_PREFIX}{procedure_id}" - await cache_provider.delete(cache_key) - await item.hard_delete(session) + if user is not None and getattr(item, "user_id", None) != user.id: + # Hide the record entirely โ€” a distinct "forbidden" response + # would leak existence of skills that belong to other users + # under the same org. + raise NoResultFound( + f"Procedural memory item with id {procedure_id} not found." + ) + + # Remove from cache + from mirix.database.cache_provider import get_cache_provider + + cache_provider = get_cache_provider() + if cache_provider: + cache_key = f"{cache_provider.PROCEDURAL_PREFIX}{procedure_id}" + await cache_provider.delete(cache_key) + await item.hard_delete(session) + + async def soft_delete_procedure_by_id( + self, + procedure_id: str, + actor: PydanticClient, + user: Optional[PydanticUser] = None, + ) -> None: + """Soft-delete a skill: flag ``is_deleted=True`` so it is excluded from + retrieval (`read`/`list_procedures` filter `~is_deleted`) WITHOUT + destroying the row. This is the curator's PREFERRED delete (C4, P1-4): + it is reversible (un-set the flag) and keeps the skill auditable, unlike + the hard ``delete_procedure_by_id`` which removes it entirely. + + Same three-axis ownership enforcement as ``delete_procedure_by_id``. + """ + async with self.session_maker() as session: + try: + item = await ProceduralMemoryItem.read( + db_session=session, + identifier=procedure_id, + actor=actor, + ) except NoResultFound: - raise NoResultFound(f"Procedural memory item with id {procedure_id} not found.") + raise NoResultFound( + f"Procedural memory item with id {procedure_id} not found." + ) + + if user is not None and getattr(item, "user_id", None) != user.id: + raise NoResultFound( + f"Procedural memory item with id {procedure_id} not found." + ) + + # Evict the cache so the soft-deleted skill stops being served. + from mirix.database.cache_provider import get_cache_provider + + cache_provider = get_cache_provider() + if cache_provider: + cache_key = f"{cache_provider.PROCEDURAL_PREFIX}{procedure_id}" + await cache_provider.delete(cache_key) + + item.is_deleted = True + session.add(item) + await session.commit() @enforce_types async def delete_by_client_id(self, actor: PydanticClient) -> int: @@ -988,7 +1482,9 @@ async def delete_by_client_id(self, actor: PydanticClient) -> int: async with self.session_maker() as session: # Get IDs for Redis cleanup (only fetch IDs, not full objects) result = await session.execute( - select(ProceduralMemoryItem.id).where(ProceduralMemoryItem.client_id == actor.id) + select(ProceduralMemoryItem.id).where( + ProceduralMemoryItem.client_id == actor.id + ) ) item_ids = [row[0] for row in result.all()] @@ -997,14 +1493,20 @@ async def delete_by_client_id(self, actor: PydanticClient) -> int: return 0 # Bulk delete in single query - await session.execute(delete(ProceduralMemoryItem).where(ProceduralMemoryItem.client_id == actor.id)) + await session.execute( + delete(ProceduralMemoryItem).where( + ProceduralMemoryItem.client_id == actor.id + ) + ) await session.commit() # Batch delete from Redis cache (outside of session context) redis_client = get_redis_client() if redis_client and item_ids: - redis_keys = [f"{redis_client.PROCEDURAL_PREFIX}{item_id}" for item_id in item_ids] + redis_keys = [ + f"{redis_client.PROCEDURAL_PREFIX}{item_id}" for item_id in item_ids + ] # Delete in batches to avoid command size limits BATCH_SIZE = 1000 @@ -1126,7 +1628,9 @@ async def delete_by_user_id(self, user_id: str) -> int: async with self.session_maker() as session: # Get IDs for Redis cleanup (only fetch IDs, not full objects) result = await session.execute( - select(ProceduralMemoryItem.id).where(ProceduralMemoryItem.user_id == user_id) + select(ProceduralMemoryItem.id).where( + ProceduralMemoryItem.user_id == user_id + ) ) item_ids = [row[0] for row in result.all()] @@ -1135,14 +1639,20 @@ async def delete_by_user_id(self, user_id: str) -> int: return 0 # Bulk delete in single query - await session.execute(delete(ProceduralMemoryItem).where(ProceduralMemoryItem.user_id == user_id)) + await session.execute( + delete(ProceduralMemoryItem).where( + ProceduralMemoryItem.user_id == user_id + ) + ) await session.commit() # Batch delete from Redis cache (outside of session context) redis_client = get_redis_client() if redis_client and item_ids: - redis_keys = [f"{redis_client.PROCEDURAL_PREFIX}{item_id}" for item_id in item_ids] + redis_keys = [ + f"{redis_client.PROCEDURAL_PREFIX}{item_id}" for item_id in item_ids + ] # Delete in batches to avoid command size limits BATCH_SIZE = 1000 @@ -1172,6 +1682,30 @@ async def list_procedures_by_org( """ List procedural memories across ALL users in an organization. """ + # Pad any caller-supplied embedding to MAX_EMBEDDING_DIM at the point of + # CONSUMPTION (REST callers may pass a raw 3072-dim vector) so it never + # reaches the Vector(4096) PG column or DIM=4096 Redis index unpadded. + embedded_text = _pad_to_max(embedded_text) + + # Hybrid retrieval (EverOS-aligned) intercepted at the TOP โ€” before the + # cache block, whose `else` branch would otherwise feed "hybrid" to + # search_text_by_org as if it were a text method. See + # _hybrid_search_for_org. Empty queries fall through to recent-items. + if search_method == "hybrid" and query and query.strip(): + return await self._hybrid_search_for_org( + agent_state=agent_state, + organization_id=organization_id, + query=query, + embedded_text=embedded_text, + search_field=search_field, + limit=limit, + timezone_str=timezone_str, + filter_tags=filter_tags, + scopes=scopes, + use_cache=use_cache, + similarity_threshold=similarity_threshold, + ) + from mirix.database.redis_client import get_redis_client redis_client = get_redis_client() @@ -1188,9 +1722,15 @@ async def list_procedures_by_org( scopes=scopes, ) if results: - logger.debug("Cache: %d procedural memories for org %s", len(results), organization_id) + logger.debug( + "Cache: %d procedural memories for org %s", + len(results), + organization_id, + ) results = redis_client.clean_redis_fields(results) - return [PydanticProceduralMemoryItem(**item) for item in results] + return [ + PydanticProceduralMemoryItem(**item) for item in results + ] elif search_method == "embedding": if embedded_text is None: @@ -1199,7 +1739,9 @@ async def list_procedures_by_org( from mirix.constants import MAX_EMBEDDING_DIM from mirix.embeddings import embedding_model - embedded_text = await (await embedding_model(agent_state.embedding_config)).get_text_embedding(query) + embedded_text = await ( + await embedding_model(agent_state.embedding_config) + ).get_text_embedding(query) embedded_text = np.array(embedded_text) embedded_text = np.pad( embedded_text, @@ -1207,7 +1749,11 @@ async def list_procedures_by_org( mode="constant", ).tolist() - vector_field = f"{search_field}_embedding" if search_field else "summary_embedding" + vector_field = ( + f"{search_field}_embedding" + if search_field + else "description_embedding" + ) results = await redis_client.search_vector_by_org( index_name=redis_client.PROCEDURAL_INDEX, embedding=embedded_text, @@ -1220,13 +1766,15 @@ async def list_procedures_by_org( if results: logger.debug("Cache vector: %d results", len(results)) results = redis_client.clean_redis_fields(results) - return [PydanticProceduralMemoryItem(**item) for item in results] + return [ + PydanticProceduralMemoryItem(**item) for item in results + ] else: results = await redis_client.search_text_by_org( index_name=redis_client.PROCEDURAL_INDEX, query_text=query, - search_field=search_field or "summary", + search_field=search_field or "description", search_method=search_method, limit=limit or 50, organization_id=organization_id, @@ -1236,12 +1784,23 @@ async def list_procedures_by_org( if results: logger.debug("Cache text: %d results", len(results)) results = redis_client.clean_redis_fields(results) - return [PydanticProceduralMemoryItem(**item) for item in results] + return [ + PydanticProceduralMemoryItem(**item) for item in results + ] except Exception as e: logger.warning("Cache search failed: %s", e) async with self.session_maker() as session: - base_query = select(ProceduralMemoryItem).where(ProceduralMemoryItem.organization_id == organization_id) + base_query = ( + select(ProceduralMemoryItem) + .where(ProceduralMemoryItem.organization_id == organization_id) + # Exclude soft-deleted skills (curator soft-delete sets + # is_deleted=True; they must NOT surface via org-wide retrieval). + # The base SqlalchemyBase reads filter ~is_deleted, but this raw + # fallback query builds select() directly and would otherwise + # leak them. Mirrors list_procedures' ~is_deleted predicate. + .where(~ProceduralMemoryItem.is_deleted) + ) from mirix.database.filter_tags_query import apply_filter_tags_sqlalchemy @@ -1264,22 +1823,33 @@ async def list_procedures_by_org( if embedded_text is None: from mirix.embeddings import embedding_model - embedded_text = await (await embedding_model(embedding_config)).get_text_embedding(query) + # Pad to MAX_EMBEDDING_DIM so a raw (e.g. 3072-dim) query + # vector matches the Vector(4096) column. (Caller-supplied + # embeddings are already padded at function entry above.) + embedded_text = _pad_to_max( + await ( + await embedding_model(embedding_config) + ).get_text_embedding(query) + ) # Determine which embedding field to search - if search_field == "summary": - embedding_field = ProceduralMemoryItem.summary_embedding - elif search_field == "steps": - embedding_field = ProceduralMemoryItem.steps_embedding + if search_field == "description": + embedding_field = ProceduralMemoryItem.description_embedding + elif search_field == "instructions": + embedding_field = ProceduralMemoryItem.instructions_embedding else: - embedding_field = ProceduralMemoryItem.summary_embedding + embedding_field = ProceduralMemoryItem.description_embedding - embedding_query_field = embedding_field.cosine_distance(embedded_text).label("distance") + embedding_query_field = embedding_field.cosine_distance( + embedded_text + ).label("distance") base_query = base_query.add_columns(embedding_query_field) # Apply similarity threshold if provided if similarity_threshold is not None: - base_query = base_query.where(embedding_query_field < similarity_threshold) + base_query = base_query.where( + embedding_query_field < similarity_threshold + ) base_query = base_query.order_by(embedding_query_field) @@ -1288,22 +1858,26 @@ async def list_procedures_by_org( from sqlalchemy import func # Determine search field - if search_field == "summary": - text_field = ProceduralMemoryItem.summary - elif search_field == "steps": - text_field = ProceduralMemoryItem.steps + if search_field == "description": + text_field = ProceduralMemoryItem.description + elif search_field == "instructions": + text_field = ProceduralMemoryItem.instructions else: - text_field = ProceduralMemoryItem.summary + text_field = ProceduralMemoryItem.description tsquery = func.plainto_tsquery("english", query) tsvector = func.to_tsvector("english", text_field) rank = func.ts_rank_cd(tsvector, tsquery).label("rank") - base_query = base_query.add_columns(rank).where(tsvector.op("@@")(tsquery)).order_by(rank.desc()) + base_query = ( + base_query.add_columns(rank) + .where(tsvector.op("@@")(tsquery)) + .order_by(rank.desc()) + ) if limit: base_query = base_query.limit(limit) result = await session.execute(base_query) items = result.scalars().all() - return [item.to_pydantic() for item in items] \ No newline at end of file + return [item.to_pydantic() for item in items] diff --git a/mirix/services/raw_memory_manager.py b/mirix/services/raw_memory_manager.py index b4ca6d596..f96f40379 100644 --- a/mirix/services/raw_memory_manager.py +++ b/mirix/services/raw_memory_manager.py @@ -11,9 +11,10 @@ from datetime import datetime, timezone from typing import Any, Dict, List, Optional, Tuple -from sqlalchemy import and_, desc, func, or_, select +from sqlalchemy import and_, desc, or_, select from mirix.constants import BUILD_EMBEDDINGS_FOR_MEMORY +from mirix.helpers.keyed_locks import KeyedLocks from mirix.log import get_logger from mirix.orm.errors import NoResultFound from mirix.orm.raw_memory import RawMemory @@ -28,6 +29,11 @@ logger = get_logger(__name__) +# Serializes same-process updates to one raw memory so concurrent +# read-modify-write cycles don't drop each other's changes. Self-evicting, so +# it doesn't grow with every memory_id ever updated. +_raw_memory_update_locks = KeyedLocks() + class RawMemoryManager: """ @@ -326,116 +332,118 @@ async def update_raw_memory( tags_merge_mode, ) - async with self.session_maker() as session: - # Fetch the existing memory with row-level lock (SELECT FOR UPDATE) - # This prevents race conditions when multiple agents append/merge concurrently - stmt = select(RawMemory).where(RawMemory.id == memory_id).with_for_update() + async with _raw_memory_update_locks.acquire(memory_id): + async with self.session_maker() as session: + # Fetch the existing memory with row-level lock (SELECT FOR UPDATE). + # The in-process lock above covers SQLite and same-process async + # concurrency, where SELECT FOR UPDATE is ineffective/no-op. + stmt = select(RawMemory).where(RawMemory.id == memory_id).with_for_update() - result = await session.execute(stmt) - try: - raw_memory = result.scalar_one() - except NoResultFound: - raise ValueError(f"Raw memory {memory_id} not found") - - # Perform access control check (replaces RawMemory.read's built-in check) - if raw_memory.organization_id != actor.organization_id: - raise ValueError( - f"Access denied: memory {memory_id} belongs to " - f"organization {raw_memory.organization_id}, " - f"actor belongs to {actor.organization_id}" - ) + result = await session.execute(stmt) + try: + raw_memory = result.scalar_one() + except NoResultFound: + raise ValueError(f"Raw memory {memory_id} not found") - # Perform scope access control check - must match actor's write_scope to update - memory_scope = (raw_memory.filter_tags or {}).get("scope") - if memory_scope != actor.write_scope: - raise ValueError( - f"Access denied: memory {memory_id} has scope '{memory_scope}', " - f"actor has write_scope '{actor.write_scope}'" - ) + # Perform access control check (replaces RawMemory.read's built-in check) + if raw_memory.organization_id != actor.organization_id: + raise ValueError( + f"Access denied: memory {memory_id} belongs to " + f"organization {raw_memory.organization_id}, " + f"actor belongs to {actor.organization_id}" + ) - # Perform user_id access control check if provided - if user_id and raw_memory.user_id != user_id: - raise ValueError(f"Raw memory {memory_id} not found") - - # Prevent scope tampering in filter_tags updates - if new_filter_tags is not None and "scope" in new_filter_tags: - if new_filter_tags["scope"] != actor.write_scope: - raise ValueError("Cannot change memory scope - scope must match actor.write_scope") - - # Update context - if new_context is not None: - if context_update_mode == "append": - raw_memory.context = f"{raw_memory.context}\n\n{new_context}" - logger.debug("Appended to context for memory %s", memory_id) - else: # replace - raw_memory.context = new_context - logger.debug("Replaced context for memory %s", memory_id) - - # Update filter_tags - if new_filter_tags is not None: - if tags_merge_mode == "merge": - # Merge new tags with existing - existing_tags = raw_memory.filter_tags or {} - raw_memory.filter_tags = { - **existing_tags, - **new_filter_tags, - } - logger.debug("Merged filter_tags for memory %s", memory_id) - else: # replace - # Preserve scope when replacing tags - scope is immutable - preserved_scope = (raw_memory.filter_tags or {}).get("scope") - raw_memory.filter_tags = new_filter_tags - if preserved_scope: - raw_memory.filter_tags["scope"] = preserved_scope - logger.debug("Replaced filter_tags for memory %s", memory_id) - - # Regenerate embeddings if context changed and agent_state provided - if BUILD_EMBEDDINGS_FOR_MEMORY and agent_state is not None and new_context is not None: - try: - from mirix.embeddings import embedding_model + # Perform scope access control check - must match actor's write_scope to update + memory_scope = (raw_memory.filter_tags or {}).get("scope") + if memory_scope != actor.write_scope: + raise ValueError( + f"Access denied: memory {memory_id} has scope '{memory_scope}', " + f"actor has write_scope '{actor.write_scope}'" + ) - embed_model = await embedding_model(agent_state.embedding_config) - context_embedding = await embed_model.get_text_embedding(raw_memory.context) + # Perform user_id access control check if provided + if user_id and raw_memory.user_id != user_id: + raise ValueError(f"Raw memory {memory_id} not found") + + # Prevent scope tampering in filter_tags updates + if new_filter_tags is not None and "scope" in new_filter_tags: + if new_filter_tags["scope"] != actor.write_scope: + raise ValueError("Cannot change memory scope - scope must match actor.write_scope") + + # Update context + if new_context is not None: + if context_update_mode == "append": + raw_memory.context = f"{raw_memory.context}\n\n{new_context}" + logger.debug("Appended to context for memory %s", memory_id) + else: # replace + raw_memory.context = new_context + logger.debug("Replaced context for memory %s", memory_id) + + # Update filter_tags + if new_filter_tags is not None: + if tags_merge_mode == "merge": + # Merge new tags with existing + existing_tags = raw_memory.filter_tags or {} + raw_memory.filter_tags = { + **existing_tags, + **new_filter_tags, + } + logger.debug("Merged filter_tags for memory %s", memory_id) + else: # replace + # Preserve scope when replacing tags - scope is immutable + preserved_scope = (raw_memory.filter_tags or {}).get("scope") + raw_memory.filter_tags = new_filter_tags + if preserved_scope: + raw_memory.filter_tags["scope"] = preserved_scope + logger.debug("Replaced filter_tags for memory %s", memory_id) + + # Regenerate embeddings if context changed and agent_state provided + if BUILD_EMBEDDINGS_FOR_MEMORY and agent_state is not None and new_context is not None: + try: + from mirix.embeddings import embedding_model + + embed_model = await embedding_model(agent_state.embedding_config) + context_embedding = await embed_model.get_text_embedding(raw_memory.context) + + raw_memory.context_embedding = PydanticRawMemoryItem.pad_embeddings(context_embedding) + raw_memory.embedding_config = agent_state.embedding_config + except Exception as e: + logger.warning("Failed to regenerate embeddings for raw memory update: %s", e) + + # Update last_modify and timestamp (use naive UTC for TIMESTAMP WITHOUT TIME ZONE) + now_utc = datetime.now(timezone.utc) + raw_memory.updated_at = now_utc.replace(tzinfo=None) if now_utc.tzinfo else now_utc + raw_memory.last_modify = { + "timestamp": datetime.now(timezone.utc).isoformat(), + "operation": "updated", + } + # Audit field (_last_updated_by_id) is handled by base class via + # property accessor when using update_with_redis(), or can be set + # manually: raw_memory.last_updated_by_id = actor.id + if actor: + raw_memory.last_updated_by_id = actor.id - raw_memory.context_embedding = PydanticRawMemoryItem.pad_embeddings(context_embedding) - raw_memory.embedding_config = agent_state.embedding_config - except Exception as e: - logger.warning("Failed to regenerate embeddings for raw memory update: %s", e) - - # Update last_modify and timestamp (use naive UTC for TIMESTAMP WITHOUT TIME ZONE) - now_utc = datetime.now(timezone.utc) - raw_memory.updated_at = now_utc.replace(tzinfo=None) if now_utc.tzinfo else now_utc - raw_memory.last_modify = { - "timestamp": datetime.now(timezone.utc).isoformat(), - "operation": "updated", - } - # Audit field (_last_updated_by_id) is handled by base class via - # property accessor when using update_with_redis(), or can be set - # manually: raw_memory.last_updated_by_id = actor.id - if actor: - raw_memory.last_updated_by_id = actor.id - - # Commit changes - await session.commit() - - # Invalidate cache - try: - from mirix.database.cache_provider import get_cache_provider + # Commit changes + await session.commit() - cache_provider = get_cache_provider() - if cache_provider: - cache_key = f"{cache_provider.RAW_MEMORY_PREFIX}{memory_id}" - await cache_provider.delete(cache_key) - logger.debug("Invalidated cache for memory %s", memory_id) - except Exception as e: - logger.warning( - "Failed to invalidate cache for memory %s: %s", - memory_id, - e, - ) + # Invalidate cache + try: + from mirix.database.cache_provider import get_cache_provider + + cache_provider = get_cache_provider() + if cache_provider: + cache_key = f"{cache_provider.RAW_MEMORY_PREFIX}{memory_id}" + await cache_provider.delete(cache_key) + logger.debug("Invalidated cache for memory %s", memory_id) + except Exception as e: + logger.warning( + "Failed to invalidate cache for memory %s: %s", + memory_id, + e, + ) - logger.info("Raw memory updated: id=%s", memory_id) - return raw_memory.to_pydantic() + logger.info("Raw memory updated: id=%s", memory_id) + return raw_memory.to_pydantic() @enforce_types async def delete_raw_memory( diff --git a/mirix/services/resource_memory_manager.py b/mirix/services/resource_memory_manager.py index 95d5aa9d4..4b8d649f5 100755 --- a/mirix/services/resource_memory_manager.py +++ b/mirix/services/resource_memory_manager.py @@ -735,8 +735,19 @@ async def list_resources( ) else: # Fallback to in-memory BM25 for SQLite (legacy method) - # Load all candidate items (memory-intensive, kept for compatibility) - result = await session.execute(select(ResourceMemoryItem).where(ResourceMemoryItem.user_id == user.id)) + # Load all candidate items (memory-intensive, kept for + # compatibility). Mirror base_query's WHERE set (org + + # filter_tags + scopes) so the SQLite fallback enforces the + # SAME scope filtering as the PG/other paths. + candidate_query = ( + select(ResourceMemoryItem) + .where(ResourceMemoryItem.user_id == user.id) + .where(ResourceMemoryItem.organization_id == organization_id) + ) + candidate_query = apply_filter_tags_sqlalchemy( + candidate_query, ResourceMemoryItem, filter_tags, scopes=scopes + ) + result = await session.execute(candidate_query) all_items = result.scalars().all() if not all_items: diff --git a/mirix/services/semantic_memory_manager.py b/mirix/services/semantic_memory_manager.py index c43c98b17..c22442dc6 100755 --- a/mirix/services/semantic_memory_manager.py +++ b/mirix/services/semantic_memory_manager.py @@ -849,10 +849,19 @@ async def list_semantic_items( ) else: # Fallback to in-memory BM25 for SQLite (legacy method) - # Load all candidate items (memory-intensive, kept for compatibility) - result = await session.execute( - select(SemanticMemoryItem).where(SemanticMemoryItem.user_id == user.id) + # Load all candidate items (memory-intensive, kept for + # compatibility). Mirror base_query's WHERE set (org + + # filter_tags + scopes) so the SQLite fallback enforces + # the SAME scope filtering as the PG/other paths. + candidate_query = ( + select(SemanticMemoryItem) + .where(SemanticMemoryItem.user_id == user.id) + .where(SemanticMemoryItem.organization_id == organization_id) ) + candidate_query = apply_filter_tags_sqlalchemy( + candidate_query, SemanticMemoryItem, filter_tags, scopes=scopes + ) + result = await session.execute(candidate_query) all_items = result.scalars().all() if not all_items: @@ -908,7 +917,17 @@ async def list_semantic_items( elif search_method == "fuzzy_match": # Fuzzy matching: load all candidate items into memory and compute a fuzzy match score. - result = await session.execute(select(SemanticMemoryItem).where(SemanticMemoryItem.user_id == user.id)) + # Mirror base_query's WHERE set (org + filter_tags + scopes) + # so fuzzy reads enforce the SAME scope filtering. + candidate_query = ( + select(SemanticMemoryItem) + .where(SemanticMemoryItem.user_id == user.id) + .where(SemanticMemoryItem.organization_id == organization_id) + ) + candidate_query = apply_filter_tags_sqlalchemy( + candidate_query, SemanticMemoryItem, filter_tags, scopes=scopes + ) + result = await session.execute(candidate_query) all_items = result.scalars().all() scored_items = [] for item in all_items: diff --git a/mirix/services/session_experience_distiller.py b/mirix/services/session_experience_distiller.py new file mode 100644 index 000000000..ae69e1d11 --- /dev/null +++ b/mirix/services/session_experience_distiller.py @@ -0,0 +1,590 @@ +"""General per-session experience distillation. + +Given the sealed, not-yet-distilled sessions in the Conversation Message Store, +this distills each session's transcript, IN PARALLEL, into zero or more +transferable :class:`SkillExperience` rows (status ``pending``). A single +session can yield MULTIPLE experiences and a MIX of ``worth_learning`` / +``worth_avoiding`` โ€” there is no external success/failure oracle; the lessons +are derived purely from the conversation content (user critique/confirmation, +tool errors/retries, the agent's own self-corrections, or weak inference). + +The transcript source is the Conversation Message Store (the SINGLE source of +truth for skill learning): a dedicated store of external conversation turns with +their REAL ``user`` / ``assistant`` roles preserved, written only when an add +request carried a ``session_id``. The distiller therefore can NEVER see the meta +agent's own memory-management scaffolding (the "[System Message] As the meta +memory manager โ€ฆ" instruction, its trigger/finish_memory_update tool calls and +results, or the continue_chaining control replies) โ€” those never enter this +store. Correctness comes from structure (which table we read), not from a string +heuristic, so the legacy ``_is_mirix_scaffolding`` filter is gone. + +Design points (CLAUDE.md async rules): +* Async-only; never starts a nested event loop (the server loop is running). +* Per-session fan-out via ``asyncio.gather`` โ€” each coroutine does ONE transcript + fetch + ONE LLM call + N inserts, wrapped in try/except so one bad session can + never crash the whole run. +* The LLM completion reuses MIRIX's async :class:`LLMClient` + (``send_llm_request`` โ†’ ``choices[0].message``). +* The system prompt is the rewritten general Experience-Distiller prompt at + ``prompts/system/base/auto_dream_agent/procedural.txt``. + +Session enumeration delegates to +``ConversationMessageManager.list_sealed_undistilled_sessions`` (sealed = +a strictly-newer distinct session exists; oldest-first), so the rolling barrier +never distills the in-progress head of the window. +""" + +from __future__ import annotations + +import asyncio +import json +from typing import Dict, List, Optional, Tuple + +from mirix.constants import ( + DISTILLER_MAX_MESSAGE_CHARS, + DISTILLER_MAX_TRANSCRIPT_CHARS, +) +from mirix.log import get_logger +from mirix.services.conversation_message_manager import owner_org as _owner_org +from mirix.schemas.agent import AgentState +from mirix.schemas.client import Client as PydanticClient +from mirix.schemas.enums import MessageRole +from mirix.schemas.message import Message +from mirix.schemas.mirix_message_content import TextContent +from mirix.schemas.skill_experience import ( + SKILL_EXPERIENCE_MAX_CONTENT_LEN as _CONTENT_CAP, + SKILL_EXPERIENCE_MAX_EVIDENCE_LEN as _EVIDENCE_CAP, + SKILL_EXPERIENCE_MAX_TITLE_LEN as _TITLE_CAP, + SkillExperience as PydanticSkillExperience, + _clamp01, +) +from mirix.schemas.user import User as PydanticUser +from mirix.helpers.json_parsing import parse_distiller_json_array +from mirix.services.conversation_message_manager import ConversationMessageManager +from mirix.services.skill_experience_manager import SkillExperienceManager + +logger = get_logger(__name__) + +# Per-transcript char budget. We keep the HEAD and TAIL (where the task framing +# and the final user verdict / confirmation usually live) and elide the middle. +# Env-configurable (MIRIX_DISTILLER_MAX_TRANSCRIPT_CHARS) โ€” see constants.py. +_MAX_TRANSCRIPT_CHARS = DISTILLER_MAX_TRANSCRIPT_CHARS + +# Per-turn content cap so one giant turn (often a large tool result) can't +# dominate the transcript. Env-configurable (MIRIX_DISTILLER_MAX_MESSAGE_CHARS). +_MAX_MESSAGE_CHARS = DISTILLER_MAX_MESSAGE_CHARS + + +class _LLMCallError(Exception): + """Raised on an OPERATIONAL LLM failure (no client, failed request, or + malformed response) โ€” as opposed to a legitimately empty distillation. Lets + the per-session path mark the session as NOT processed so the barrier does + not advance and a later round retries it.""" + + +class _DistillFailed(Exception): + """Raised by ``_distill_one`` when a session could NOT be processed due to an + operational failure (DB read, LLM call, or persistence). ``distill_sessions`` + catches it to EXCLUDE that session from ``processed_session_ids`` so the + barrier is not advanced past a failed conversation and a later round retries + it. Carries the offending ``session_id`` for diagnostics.""" + + def __init__(self, session_id: str): + self.session_id = session_id + super().__init__(session_id) + + +def _load_distiller_prompt() -> str: + """Load the general Experience-Distiller system prompt. + + Uses the codebase-standard cached system-prompt loader (the same one the other + distillers use) instead of a raw ``open()`` in the async path โ€” it resolves to + prompts/system/base/auto_dream_agent/procedural.txt, the single on-disk source + of truth shared with the auto_dream procedural mode. + """ + from mirix.prompts.gpt_system import get_system_text + + return get_system_text("base/auto_dream_agent/procedural") + + +class SessionExperienceDistiller: + """Distill the sealed, not-yet-distilled conversation sessions into experiences. + + Args: + llm_client: an instance exposing ``async send_llm_request(messages)`` + (an ``LLMClientBase``). Injected directly in tests; in production + pass ``llm_config`` and one is created lazily. + llm_config: an ``LLMConfig`` used to build the client when ``llm_client`` + is not supplied. Required for a real completion. + experience_manager: a :class:`SkillExperienceManager`; defaults to a + fresh instance (constructed lazily so import never touches the DB). + conversation_manager: a :class:`ConversationMessageManager` โ€” the source + of session enumeration and per-session transcripts. Defaults to a + fresh instance (constructed lazily so import never touches the DB). + """ + + def __init__( + self, + *, + llm_client=None, + llm_config=None, + experience_manager: Optional[SkillExperienceManager] = None, + conversation_manager: Optional[ConversationMessageManager] = None, + ): + self._llm_client = llm_client + self._llm_config = llm_config + self._experience_manager = experience_manager + self._conversation_manager = conversation_manager + self._system_prompt = _load_distiller_prompt() + + # ------------------------------------------------------------------ # + # Session enumeration # + # ------------------------------------------------------------------ # + + async def enumerate_sealed_sessions( + self, + *, + user_id: str, + organization_id: str, + actor: PydanticClient, + limit: int, + ) -> List[str]: + """Return up to ``limit`` sealed, not-yet-distilled session_ids, OLDEST first. + + Delegates to + :meth:`ConversationMessageManager.list_sealed_undistilled_sessions`. + "Sealed" = a strictly-newer distinct session exists, so the open head of + the window is never returned; only sessions whose turns have NULL + ``distilled_at`` survive. Scoped per ``(user, organization)``. + """ + if limit <= 0: + return [] + return await self._get_conversation_manager().list_sealed_undistilled_sessions( + user_id=user_id, + organization_id=organization_id, + actor=actor, + limit=limit, + ) + + # ------------------------------------------------------------------ # + # Per-session distillation (parallel fan-out) # + # ------------------------------------------------------------------ # + + async def distill_sessions( + self, + *, + meta_agent_state: AgentState, + user: PydanticUser, + actor: PydanticClient, + session_ids: List[str], + existing_skills: Optional[List[Dict]] = None, + ) -> Tuple[List[PydanticSkillExperience], List[str]]: + """Distill each session IN PARALLEL into persisted experiences. + + Each session โ†’ one coroutine (fetch transcript โ†’ one LLM call โ†’ parse + the JSON array โ†’ persist each element as a ``pending`` SkillExperience). + Per-session try/except isolates failures: a single bad session yields an + empty outcome rather than crashing the whole gather. + + Returns ``(experiences, processed_session_ids)``: + * ``experiences`` โ€” the flat list of all created experiences. + * ``processed_session_ids``โ€” the session_ids that were SUCCESSFULLY + processed (including those that legitimately yielded zero + experiences). A session that hit an OPERATIONAL failure (bad + transcript fetch / LLM call / persistence) is OMITTED here, so the + caller leaves it undistilled and a later round retries it instead of + silently advancing the barrier past a failed conversation. + """ + if not session_ids: + return [], [] + + skills_block = self._format_existing_skills(existing_skills or []) + # return_exceptions=True so one operationally-failed session surfaces as a + # _DistillFailed in `results` rather than crashing the whole gather; we map + # exceptions -> "not processed" and any returned list (even []) -> + # "processed". + results = await asyncio.gather( + *[ + self._distill_one( + meta_agent_state=meta_agent_state, + user=user, + actor=actor, + session_id=sid, + skills_block=skills_block, + ) + for sid in session_ids + ], + return_exceptions=True, + ) + created: List[PydanticSkillExperience] = [] + processed_session_ids: List[str] = [] + for sid, result in zip(session_ids, results): + if isinstance(result, BaseException): + # Operational failure (logged in _distill_one). Leave the session + # undistilled so a later round retries it. + continue + created.extend(result) + processed_session_ids.append(sid) + logger.info( + "Distilled %d session(s) โ†’ %d experience(s); %d processed, %d failed", + len(session_ids), + len(created), + len(processed_session_ids), + len(session_ids) - len(processed_session_ids), + ) + return created, processed_session_ids + + async def _distill_one( + self, + *, + meta_agent_state: AgentState, + user: PydanticUser, + actor: PydanticClient, + session_id: str, + skills_block: str, + ) -> List[PydanticSkillExperience]: + """Distill ONE session into persisted experiences. + + Returns the (possibly empty) list of created experiences on success โ€” a + legitimately empty session (no turns / no usable content / the model + returned ``[]``) is a SUCCESS that returns ``[]``. On an OPERATIONAL + failure (a failed transcript fetch / LLM call / persistence) it raises + :class:`_DistillFailed`, which :meth:`distill_sessions` catches to mark + the session as NOT processed (so its barrier is not advanced and a later + round retries it). It never lets any OTHER exception escape. + """ + try: + turns = await self._get_conversation_manager().list_turns_for_session( + session_id=session_id, + user_id=user.id, + organization_id=_owner_org(actor), + actor=actor, + ) + if not turns: + # Empty session: nothing to learn, but successfully processed. + return [] + + transcript = self._render_transcript(turns) + if not transcript.strip(): + # No usable content โ€” successfully processed, nothing to learn. + return [] + parsed = await self._call_llm( + agent_id=meta_agent_state.id, + session_id=session_id, + transcript=transcript, + skills_block=skills_block, + ) + if not parsed: + # The model returned an empty array โ€” a legitimate "nothing worth + # remembering" result. Processed successfully. + return [] + + return await self._persist_experiences( + meta_agent_state=meta_agent_state, + user=user, + actor=actor, + session_id=session_id, + parsed=parsed, + ) + except _DistillFailed: + # Already classified (e.g. total persistence failure) โ€” propagate as-is + # so distill_sessions excludes this session from processed. + logger.warning( + "Session experience distill failed for session %s (operational)", + session_id, + ) + raise + except Exception as e: # noqa: BLE001 โ€” classify as an operational failure + # DB read, LLM call (raised as _LLMCallError), or another operational + # error. Re-raise as _DistillFailed so distill_sessions leaves the + # session undistilled for a later retry instead of advancing the barrier. + logger.warning( + "Session experience distill failed for session %s: %s", session_id, e + ) + raise _DistillFailed(session_id) from e + + async def _persist_experiences( + self, + *, + meta_agent_state: AgentState, + user: PydanticUser, + actor: PydanticClient, + session_id: str, + parsed: List[Dict], + ) -> List[PydanticSkillExperience]: + """Persist each valid parsed item as a pending SkillExperience. + + Per-row try/except means one bad row can't drop the rest. But if there + were valid rows to persist and EVERY persistence attempt raised (i.e. the + store is down, not just one malformed row), that is an operational failure + โ€” raise :class:`_DistillFailed` so the session is NOT marked distilled and + a later round retries it, rather than silently advancing the barrier past + a conversation whose experiences never landed. Pure VALIDATION skips + (bad type / empty title) are not failures and never raise. + """ + mgr = self._get_experience_manager() + created: List[PydanticSkillExperience] = [] + attempted = 0 # valid rows we actually tried to persist + persist_errors = 0 + for item in parsed: + exp_type = item.get("experience_type") + if exp_type not in ("worth_learning", "worth_avoiding"): + # Unknown/missing type โ†’ skip (create_experience would also + # reject it, but skip quietly here so one bad element does not + # abort the whole session's persistence). + logger.debug( + "Skipping experience with bad type %r in session %s", + exp_type, + session_id, + ) + continue + title = (item.get("title") or "").strip() + if not title: + continue + content = item.get("content") or "" + importance = _clamp01(item.get("importance")) + credibility = _clamp01(item.get("credibility")) + evidence = self._normalize_evidence(item.get("evidence")) + attempted += 1 + try: + exp = await mgr.create_experience( + agent_id=meta_agent_state.id, + user_id=user.id, + organization_id=_owner_org(actor), + session_id=session_id, + experience_type=exp_type, + title=title[: _TITLE_CAP], + content=content[: _CONTENT_CAP], + importance=importance, + credibility=credibility, + evidence=evidence[: _EVIDENCE_CAP], + status="pending", + created_by_id=getattr(actor, "id", None), + ) + created.append(exp) + except Exception as e: # noqa: BLE001 โ€” one bad row mustn't drop the rest + persist_errors += 1 + logger.warning( + "Failed to persist experience %r (session %s): %s", + title, + session_id, + e, + ) + + # Total persistence failure: we had valid rows but NONE landed and every + # attempt errored โ†’ the store is down. Signal an operational failure so + # the session is left undistilled for retry. (A partial success โ€” at least + # one row landed โ€” returns what persisted; a session with only validation + # skips has attempted==0 and is a clean empty result.) + if attempted > 0 and not created and persist_errors == attempted: + raise _DistillFailed(session_id) + return created + + # ------------------------------------------------------------------ # + # LLM plumbing # + # ------------------------------------------------------------------ # + + async def _call_llm( + self, + *, + agent_id: str, + session_id: str, + transcript: str, + skills_block: str, + ) -> List[Dict]: + """Run ONE completion and parse a JSON array of experiences. + + An empty ``[]`` is a legitimate "nothing worth remembering" result and is + returned normally. An OPERATIONAL failure (no client/config, a failed + request, or a malformed response) raises ``_LLMCallError`` instead of + returning ``[]`` โ€” so the caller can tell a transient failure apart from a + genuinely empty session and NOT advance the barrier on the former. + """ + client = self._get_client() + if client is None: + raise _LLMCallError( + f"no LLM client/config available for session {session_id}" + ) + + user_payload = ( + f"agent_id: {agent_id}\n" + f"session_id: {session_id}\n\n" + f"existing_skills:\n{skills_block}\n\n" + "session_transcript (chronological):\n" + f"{transcript}\n\n" + "Distill the transferable experiences from this session as a JSON array." + ) + messages = [ + Message( + agent_id=agent_id, + role=MessageRole.system, + content=[TextContent(text=self._system_prompt)], + ), + Message( + agent_id=agent_id, + role=MessageRole.user, + content=[TextContent(text=user_payload)], + ), + ] + try: + response = await client.send_llm_request(messages=messages) + except Exception as e: # noqa: BLE001 โ€” surface as an operational failure + raise _LLMCallError(f"LLM request failed: {e}") from e + try: + content = response.choices[0].message.content + except (AttributeError, IndexError) as e: + raise _LLMCallError(f"malformed LLM response: {e}") from e + + parsed = parse_distiller_json_array(content) + if not parsed and not self._is_explicit_empty_result(content): + # The shared parser returns [] both for a real empty array AND for + # unparseable output. A non-empty content that parsed to [] is most + # likely a TRANSIENT bad completion (truncation, a wrapped error), so + # we log it for observability. We do NOT raise here on purpose: this + # barrier has no durable per-session retry counter, so retrying a + # session the model *consistently* renders unparseable (e.g. it states + # "no experiences" in prose with no JSON) would peg the rolling window + # forever. Treating an unparseable empty as "nothing learned" advances + # the window at the cost of at most one session's learning on a one-off + # bad completion โ€” the safer tradeoff than a permanently stuck barrier. + logger.warning( + "session distiller: non-empty LLM output for session %s parsed to " + "[] (treating as empty; possible transient bad completion)", + session_id, + ) + return parsed + + @staticmethod + def _is_explicit_empty_result(content) -> bool: + """True iff `content` legitimately represents "no experiences" rather than + unparseable garbage. + + A model that means "nothing worth remembering" emits an empty array. We + treat empty/whitespace OR a content whose only JSON-array token is an empty + `[]` (optionally fenced / surrounded by prose) as a legitimate empty + result. Anything else that parsed to `[]` is malformed output, which + `_call_llm` raises on. Conservative by design: a false "explicit empty" + only costs one skipped (already-empty) session, never a wrongful retry. + """ + if content is None: + return True + if not isinstance(content, str): + return False + text = content.strip() + if not text: + return True + # Strip a single ```json / ``` fence if present, then look at the payload. + import re as _re + + fence = _re.search(r"```(?:json)?\s*(.*?)```", text, _re.DOTALL | _re.IGNORECASE) + if fence: + text = fence.group(1).strip() + # An explicit empty array (the only well-formed "nothing" the model emits). + # `[]`, `[ ]`, or those wrapped in a tiny bit of prose all qualify; a + # non-empty `[...]` would have parsed to a non-empty list and never reached + # here, so the only empty-array token we accept is a literal empty pair. + return _re.fullmatch(r"\[\s*\]", text) is not None + + def _get_client(self): + if self._llm_client is not None: + return self._llm_client + if self._llm_config is not None: + from mirix.llm_api.llm_client import LLMClient + + self._llm_client = LLMClient.create( + llm_config=self._llm_config.model_copy(deep=True) + ) + return self._llm_client + return None + + def _get_experience_manager(self) -> SkillExperienceManager: + if self._experience_manager is None: + self._experience_manager = SkillExperienceManager() + return self._experience_manager + + def _get_conversation_manager(self) -> ConversationMessageManager: + if self._conversation_manager is None: + self._conversation_manager = ConversationMessageManager() + return self._conversation_manager + + # ------------------------------------------------------------------ # + # Rendering helpers (pure) # + # ------------------------------------------------------------------ # + + @staticmethod + def _format_existing_skills(existing_skills: List[Dict]) -> str: + if not existing_skills: + return "(none)" + lines = [] + for s in existing_skills: + name = (s.get("name") or "").strip() + desc = (s.get("description") or "").strip() + if not name and not desc: + continue + lines.append(f"- {name}: {desc}"[:300]) + return "\n".join(lines) if lines else "(none)" + + @staticmethod + def _normalize_evidence(evidence) -> str: + """Coerce the model's `evidence` into a JSON string {quote, signal_type}.""" + valid_signals = { + "user_critique", + "user_confirmation", + "tool_error", + "self_correction", + "inferred", + } + if isinstance(evidence, dict): + quote = str(evidence.get("quote") or "")[:512] + signal = evidence.get("signal_type") + if signal not in valid_signals: + signal = "inferred" + return json.dumps({"quote": quote, "signal_type": signal}, ensure_ascii=False) + if isinstance(evidence, str) and evidence.strip(): + return json.dumps( + {"quote": evidence[:512], "signal_type": "inferred"}, + ensure_ascii=False, + ) + return json.dumps({"quote": "", "signal_type": "inferred"}, ensure_ascii=False) + + @classmethod + def _render_transcript(cls, turns) -> str: + """Render conversation turns compactly as ``role: content`` lines. + + The source is the Conversation Message Store, so each turn already has a + REAL ``role`` ('user' | 'assistant' | 'tool') and a plain-text + ``content`` โ€” no scaffolding, no embeddings. Tool turns (tool results + and serialized tool calls) are rendered like any other turn: the + distiller prompt treats tool errors/retries as strong signals. We cap + each turn and bound the whole transcript by keeping the HEAD and TAIL + (task framing + final verdict) and eliding the middle. + """ + lines: List[str] = [] + for t in turns: + role = cls._role_of(t) + text = cls._content_text(t) + if not text: + continue + if len(text) > _MAX_MESSAGE_CHARS: + text = text[:_MAX_MESSAGE_CHARS] + " โ€ฆ[truncated]" + lines.append(f"{role}: {text}") + transcript = "\n".join(lines) + if len(transcript) <= _MAX_TRANSCRIPT_CHARS: + return transcript + head = transcript[: _MAX_TRANSCRIPT_CHARS // 2] + tail = transcript[-(_MAX_TRANSCRIPT_CHARS // 2):] + return f"{head}\nโ€ฆ[elided middle of session]โ€ฆ\n{tail}" + + @staticmethod + def _role_of(t) -> str: + """The real turn role of a ConversationMessage ('user'|'assistant'|'tool').""" + role = getattr(t, "role", None) + return getattr(role, "value", None) or str(role) or "unknown" + + @staticmethod + def _content_text(t) -> str: + """Plain text of a ConversationMessage turn (content is already a str).""" + content = getattr(t, "content", None) + if isinstance(content, str): + return content.strip() + return "" diff --git a/mirix/services/skill_experience_curator.py b/mirix/services/skill_experience_curator.py new file mode 100644 index 000000000..ad9e03801 --- /dev/null +++ b/mirix/services/skill_experience_curator.py @@ -0,0 +1,492 @@ +"""General experience-driven skill self-evolution. + +This drives the procedural skill agent +(``skill_list`` / ``skill_read`` / ``skill_create`` / ``skill_edit``) from the +generic :class:`SkillExperience` store produced by session distillation. + +Flow (mirrors the load-bearing ordering of the records path so a context reset +can never wipe the bookkeeping): + + list_experiences(status="pending", priority DESC, ids=THIS ROUND's batch) + -> aggregate -> compute_edit_budget (avoid weighted heavier than learn) + -> B_min=0 early-exit when there is nothing structure-gated + -> build compact payload (one block per experience, priority-ordered) + -> set per-instance edit budget on the agent + -> snapshot BEFORE -> run procedural agent step (reset before+after) + -> snapshot AFTER -> diff skills -> compute influenced_skill_ids lineage + -> mark_consumed(ids, run_id) + write lineage (OUTSIDE the reset window) + +The curator does not authorize destructive deletes from experiences (a +``worth_avoiding`` lesson names a pitfall, not a harmful skill to destroy); v1 +is creates/edits only. The edit-budget gate, the per-mutation size gate, and +``skill_create`` name-dedup keep mutations delta/incremental โ€” never a wholesale +rewrite. + +A per-agent :class:`asyncio.Lock` serializes concurrent evolves on the same +procedural agent, because the step is bracketed by in-context resets that would +otherwise corrupt a concurrent run. +""" + +from __future__ import annotations + +import uuid +from typing import Dict, List, Optional + +from mirix.functions.function_sets.memory_tools import compute_edit_budget +from mirix.log import get_logger +from mirix.schemas.agent import AgentState, AgentType +from mirix.schemas.client import Client as PydanticClient +from mirix.schemas.skill_experience import ( + SKILL_EXPERIENCE_MAX_CONTENT_LEN as _CONTENT_CAP, +) +from mirix.schemas.user import User as PydanticUser +from mirix.services.procedural_evolution_runtime import ( + _attr, + _diff_skills, + _evolve_locks, + reset_agent_in_context_to_system, +) + +logger = get_logger(__name__) + +# How many pending experiences to consume in one evolution pass. Bounded so a +# huge backlog can't blow up the curator prompt; the highest-priority ones come +# first (importance*credibility DESC), so the cap drops only the weakest. +_MAX_EXPERIENCES_PER_RUN = 50 + + +def scope_filter_tags(actor) -> Optional[dict]: + """Build the ``filter_tags`` scope stamp for skills this actor creates. + + The evolution path is the single producer of procedural skills and bypasses + the REST layer's ``filter_tags["scope"] = client.write_scope`` injection + (``rest_api`` ``add_sync``/``add``), so it must stamp the owning client's + ``write_scope`` itself. ``skill_create`` later reads it back via + ``getattr(self, "filter_tags", None)``. Returns ``{"scope": write_scope}`` + when the actor has a truthy ``write_scope``, else ``None`` (a scope-less, + read-only actor donates no scope โ€” see caller warning). + + Side effect (intended): ``Agent.__init__`` also derives ``_block_scopes`` + from ``filter_tags["scope"]``, so the evolution agent's core-memory block + loading narrows to this scope โ€” the same behavior the normal add-flow + procedural child already has. + """ + write_scope = getattr(actor, "write_scope", None) + if write_scope: + return {"scope": write_scope} + return None + + +def build_experience_payload(experiences: List) -> str: + """Render pending experiences as a COMPACT, priority-ordered curator prompt. + + ``experiences`` arrive already ordered by ``importance*credibility`` DESC + (the manager's ``list_experiences`` ordering). Emits one block per + experience as ``[LEARN|AVOID] title โ€” content (importance=โ€ฆ, credibility=โ€ฆ; + evidence: "quote")`` โ€” far smaller than raw transcripts. + """ + lines: List[str] = [ + "## Distilled experiences from recent sessions (highest-priority first)", + ] + if not experiences: + lines.append("(no experiences)") + return "\n".join(lines) + + for exp in experiences: + etype = (_attr(exp, "experience_type") or "").lower() + tag = "LEARN" if etype == "worth_learning" else "AVOID" + title = _attr(exp, "title") or "" + content = (_attr(exp, "content") or "")[:_CONTENT_CAP] + importance = _attr(exp, "importance") + credibility = _attr(exp, "credibility") + quote = _evidence_quote(_attr(exp, "evidence")) + ev_str = f'; evidence: "{quote}"' if quote else "" + lines.append( + f"- [{tag}] {title} โ€” {content} " + f"(importance={_fmt(importance)}, credibility={_fmt(credibility)}{ev_str})" + ) + return "\n".join(lines) + + +def _fmt(v) -> str: + try: + return f"{float(v):.2f}" + except (TypeError, ValueError): + return "0.00" + + +def _evidence_quote(evidence) -> str: + """Pull the short grounding quote out of the stored evidence JSON string.""" + if not evidence: + return "" + import json + + try: + data = json.loads(evidence) + return str(data.get("quote") or "")[:240] + except (ValueError, TypeError): + return str(evidence)[:240] + + +async def _resolve_procedural_agent_state(server, actor, meta_agent_state): + """Resolve the procedural memory agent child of the meta agent. + + Deliberately scoped to THIS meta agent's children (no org-wide fallback), + so each user/meta-agent drives its own procedural agent. + """ + children = await server.agent_manager.list_agents( + actor=actor, parent_id=meta_agent_state.id + ) + for child in children: + if child.agent_type == AgentType.procedural_memory_agent: + return child + # No type-only flat fallback: matching any procedural agent in the org would + # break per-(meta-agent, user) scoping (it could evolve a DIFFERENT user's / + # meta agent's skill bank). If this meta agent has no procedural child, skip + # evolution โ€” the caller treats None as "no agent, leave experiences pending". + return None + + +async def run_experience_evolution( + *, + user: PydanticUser, + actor: PydanticClient, + meta_agent_state: AgentState, + experience_ids: Optional[List[str]] = None, + experience_manager=None, +) -> Dict: + """Evolve skills from this (agent, user)'s PENDING experiences. + + This is the production wiring: it resolves the real procedural agent, builds + the snapshot / step / lineage collaborators, and delegates the load-bearing + ordering to :func:`_run_experience_evolution_core` (which is injectable for + tests). Returns ``{skipped, budget, changes, influenced_skill_ids, + consumed_count, skills_changed}``. + + ``experience_ids`` scopes the run to ONE distillation round's freshly-made + experiences (the auto-dream procedural path passes exactly the ids it just + distilled from the last-N retained sessions). This keeps an evolve focused on + the current window: stale, lower-signal experiences from earlier rounds are + NOT re-surfaced (seeing the whole accumulated pool dilutes the prompt and + hurts skill quality). ``None`` (default) means "no scope" โ€” evolve the whole + pending pool (legacy/standalone behavior); ``[]`` means the round produced + nothing, so there is nothing to evolve (clean B_min=0 skip). + + Safe to call with no pending experiences (B_min=0 early-exit, no agent spawn). + """ + from mirix.agent import ProceduralMemoryAgent + from mirix.constants import SKILL_EVOLVE_MAX_CHAINING_STEPS + from mirix.schemas.message import Message as PydanticMessage + from mirix.schemas.mirix_message_content import TextContent + from mirix.server.server import get_server + from mirix.services.skill_experience_manager import SkillExperienceManager + + server = get_server() + exp_mgr = experience_manager or SkillExperienceManager() + + proc_agent_state = await _resolve_procedural_agent_state( + server, actor, meta_agent_state + ) + if proc_agent_state is None: + logger.info( + "[experience-curator] no procedural agent for meta=%s; skipping evolution", + meta_agent_state.id, + ) + return _empty_result(0, skipped=True) + + # The procedural child may have been created BEFORE the Experience-Based + # Evolution prompt section landed (or its tool set predates skill_*). Refresh + # its tools + system prompt from disk before we build the agent, so a manual + # /memory/auto_dream or a direct AutoDreamManager.run(mode="procedural") never + # runs a stale prompt. Best-effort: a failure here just uses the persisted + # state (the prompt section is additive, not load-bearing for the tools). + try: + await server.agent_manager.update_agent_tools_and_system_prompts( + proc_agent_state.id, actor=actor + ) + proc_agent_state = await server.agent_manager.get_agent_by_id( + agent_id=proc_agent_state.id, actor=actor + ) + except Exception as refresh_err: # noqa: BLE001 + logger.warning( + "[experience-curator] could not refresh procedural agent prompt/tools: %s", + refresh_err, + ) + + timezone_str = getattr(user, "timezone", None) or "UTC" + + # Stamp the owning client's write_scope onto the agent so every skill_create + # in this run tags its row with filter_tags={"scope": write_scope}. This + # evolution path is the single producer of skills and bypasses the REST + # layer's filter_tags["scope"]=client.write_scope injection (rest_api + # add_sync), so it must stamp the scope itself; skill_create reads it via + # getattr(self, "filter_tags", None). A scope-less actor would produce + # skills invisible to every scoped reader, so warn loudly in that case. + skill_filter_tags = scope_filter_tags(actor) + if skill_filter_tags is None: + logger.warning( + "[experience-curator] actor %s has no write_scope; evolved skills " + "will be unscoped (filter_tags=NULL) and invisible to scoped readers", + getattr(actor, "id", "?"), + ) + + proc_agent = ProceduralMemoryAgent( + agent_state=proc_agent_state, + interface=server.default_interface_factory(), + actor=actor, + user=user, + filter_tags=skill_filter_tags, + ) + + async def _snapshot(): + return await server.procedural_memory_manager.list_procedures( + agent_state=proc_agent_state, + user=user, + query="", + search_field="description", + search_method="bm25", + limit=1000, + timezone_str=timezone_str, + use_cache=False, + ) + + async def _run_step(agent, payload, budget): + input_message = PydanticMessage( + role="user", + content=[TextContent(text=payload)], + agent_id=proc_agent_state.id, + name="user", + ) + # Stateless reset discipline (identical to the records-evolve endpoint): + # hard-reset BEFORE and (in finally) AFTER the step so each evolution is + # context-isolated. The consume/lineage bookkeeping runs AFTER this + # returns (in the core), so the post-step reset can never wipe it. + await reset_agent_in_context_to_system(server, proc_agent_state.id, actor) + try: + await agent.step( + input_messages=[input_message], + chaining=True, + max_chaining_steps=SKILL_EVOLVE_MAX_CHAINING_STEPS, + actor=actor, + user=user, + ) + finally: + try: + await reset_agent_in_context_to_system(server, proc_agent_state.id, actor) + except Exception as cleanup_err: # noqa: BLE001 + logger.warning( + "Post-experience-evolve procedural context reset failed: %s", + cleanup_err, + ) + + return await _run_experience_evolution_core( + experience_manager=exp_mgr, + agent=proc_agent, + agent_id=proc_agent_state.id, + meta_agent_id=meta_agent_state.id, + user_id=user.id, + snapshot_skills=_snapshot, + run_step=_run_step, + experience_ids=experience_ids, + ) + + +def _empty_result(budget: int, *, skipped: bool) -> Dict: + return { + "skipped": skipped, + "budget": budget, + "changes": {"created": [], "edited": [], "deleted": []}, + "influenced_skill_ids": [], + "consumed_count": 0, + "superseded_count": 0, + "skills_changed": 0, + } + + +async def _maybe_await(value): + import asyncio + + if asyncio.iscoroutine(value): + return await value + return value + + +async def _run_experience_evolution_core( + *, + experience_manager, + agent, + agent_id: str, + meta_agent_id: str, + user_id: str, + snapshot_skills, + run_step, + run_id: Optional[str] = None, + experience_ids: Optional[List[str]] = None, +) -> Dict: + """The injectable core โ€” no server/agent assembly, fully unit-testable. + + Collaborators (``snapshot_skills``, ``run_step``) are passed in so a test can + drive the exact load-bearing ordering with mocks. The PER-AGENT lock makes + concurrent evolves on the same procedural agent safe. + + ``experience_ids`` scopes the run to one distillation round (see + :func:`run_experience_evolution`); ``None`` evolves the whole pending pool. + """ + run_id = run_id or f"xprun-{uuid.uuid4().hex[:12]}" + async with _evolve_locks.acquire(agent_id): + return await _evolve_locked( + experience_manager=experience_manager, + agent=agent, + agent_id=agent_id, + meta_agent_id=meta_agent_id, + user_id=user_id, + snapshot_skills=snapshot_skills, + run_step=run_step, + run_id=run_id, + experience_ids=experience_ids, + ) + + +async def _evolve_locked( + *, + experience_manager, + agent, + agent_id: str, + meta_agent_id: str, + user_id: str, + snapshot_skills, + run_step, + run_id: str, + experience_ids: Optional[List[str]] = None, +) -> Dict: + # 1) Read this round's pending experiences, priority-ordered + # (importance*credibility DESC). Experiences are keyed by the META agent id + # (their provenance owner), not the procedural agent. When the caller scopes + # to a freshly-distilled batch (`experience_ids`), only those are loaded โ€” + # earlier rounds' leftover pending stay invisible so this evolve sees ONLY + # the current window (`experience_ids=[]` -> nothing -> B_min=0 skip). + experiences = await experience_manager.list_experiences( + agent_id=meta_agent_id, + user_id=user_id, + status="pending", + limit=_MAX_EXPERIENCES_PER_RUN, + ids=experience_ids, + ) + exp_ids = [_attr(e, "id") for e in experiences] + + # 2) Aggregate -> count-driven budget. Map worth_avoiding -> n_high_fail and + # worth_learning -> n_high_succ so the edit-budget formula (avoid weighted + # heavier than learn) applies unchanged. + agg = await experience_manager.aggregate(ids=exp_ids) + n_avoid = int(agg.get("n_worth_avoiding", 0) or 0) + n_learn = int(agg.get("n_worth_learning", 0) or 0) + + # B_min=0 early-exit: nothing to learn from -> no step, no LLM, no consume. + if n_avoid + n_learn == 0: + logger.info( + "[experience-curator] B_min=0 skip: no pending experiences " + "(meta=%s, agent=%s, run=%s)", + meta_agent_id, + agent_id, + run_id, + ) + return _empty_result(0, skipped=True) + + budget = compute_edit_budget( + {"n_high_fail": n_avoid, "n_high_succ": n_learn} + ) + + # 3) Compact, priority-ordered payload. + payload = build_experience_payload(experiences) + + # 4) Set the per-instance budgets BEFORE the step (plain ints on the + # instance: no cross-user / cross-run leak). Experiences are creates/edits + # ONLY โ€” a lesson names a pitfall, not a skill to destroy. ENFORCE that + # here (the prompt's "No deletes" rule is guidance, not a gate): with a + # zero delete budget and an empty authorization set, skill_delete's gate + # (memory_tools._delete gate) rejects every delete this run. Also prefer + # soft-delete so even a gate change never hard-destroys a skill from here. + agent._edit_budget_remaining = budget + agent._delete_budget_remaining = 0 + agent._delete_authorized_skill_ids = set() + agent._prefer_soft_delete = True + + # 5) Snapshot BEFORE, run the step, snapshot AFTER. + before = await _maybe_await(snapshot_skills()) + await run_step(agent, payload, budget) + after = await _maybe_await(snapshot_skills()) + + # 6) Diff (reuses the records-path semantics: version/instructions/description). + changes = _diff_skills(before, after) + influenced = sorted( + set(changes["created"]) | set(changes["edited"]) | set(changes["deleted"]) + ) + + # 7) Bookkeeping OUTSIDE the step / reset window: consume the experiences and + # stamp the influenced-skill lineage in ONE load-modify-save (mark_consumed + # accepts influenced_skill_ids), so a reset can never lose the audit trail. + consumed = 0 + if exp_ids: + consumed = await experience_manager.mark_consumed( + ids=exp_ids, + run_id=run_id, + influenced_skill_ids=influenced or None, + ) + + # 7b) Retire this round's OVERFLOW: scoped experiences that didn't make the + # per-run priority cap (`_MAX_EXPERIENCES_PER_RUN`). Future rounds only + # pass THEIR freshly-distilled ids, so leaving these pending would strand + # them forever (they'd never be re-surfaced). They are the round's + # lowest-priority tail; supersede them so they stay out of sight while + # remaining auditable (status='superseded', NOT hard-deleted). Only in + # scoped mode โ€” the unscoped/global path (experience_ids is None) leaves + # leftovers pending on purpose so a later global sweep can pick them up. + superseded = 0 + if experience_ids: + consumed_set = set(exp_ids) + overflow = [eid for eid in experience_ids if eid not in consumed_set] + if overflow: + # Owner-scope the supersede: even though production only ever passes + # this (meta-agent, user)'s freshly-distilled ids, constraining the + # update by owner guarantees a stray foreign id can never retire a + # different (agent, user)'s pending experience. + superseded = await experience_manager.mark_superseded( + ids=overflow, + agent_id=meta_agent_id, + user_id=user_id, + ) + logger.info( + "[experience-curator] run=%s: superseded %d overflow experience(s) " + "beyond the per-run cap of %d", + run_id, + superseded, + _MAX_EXPERIENCES_PER_RUN, + ) + + skills_changed = ( + len(changes["created"]) + len(changes["edited"]) + len(changes["deleted"]) + ) + + logger.info( + "[experience-curator] meta=%s agent=%s run=%s: budget=%d consumed=%d " + "created=%d edited=%d deleted=%d", + meta_agent_id, + agent_id, + run_id, + budget, + consumed, + len(changes["created"]), + len(changes["edited"]), + len(changes["deleted"]), + ) + + return { + "skipped": False, + "budget": budget, + "changes": changes, + "influenced_skill_ids": influenced, + "consumed_count": consumed, + "superseded_count": superseded, + "skills_changed": skills_changed, + "run_id": run_id, + } diff --git a/mirix/services/skill_experience_manager.py b/mirix/services/skill_experience_manager.py new file mode 100644 index 000000000..9ee905d74 --- /dev/null +++ b/mirix/services/skill_experience_manager.py @@ -0,0 +1,335 @@ +"""Manager for the general session-experience store. + +Async-only: it grabs sessions via `self.session_maker` (the server's +`db_context`) and never calls `asyncio.run()` (the server event loop is already +running). + +The store is the durable hand-off between session distillation (one or more +experiences per session) and skill evolution (consumes pending experiences, +prioritized by importance*credibility, to create/edit skills). Experiences flow +pending -> consumed | superseded. +""" + +from __future__ import annotations + +from typing import Dict, List, Optional + +from sqlalchemy import delete, select + +from mirix.client.utils import get_utc_time +from mirix.log import get_logger +from mirix.orm.skill_experience import SkillExperience as SkillExperienceModel +from mirix.schemas.client import Client as PydanticClient +from mirix.schemas.skill_experience import ( + SkillExperience as PydanticSkillExperience, + SkillExperienceCreate, +) +from mirix.utils import enforce_types + +logger = get_logger(__name__) + + +class SkillExperienceManager: + """Persist and query distilled per-session transferable experiences.""" + + def __init__(self): + from mirix.server.server import db_context + + self.session_maker = db_context + + @enforce_types + async def create_experience( + self, + *, + agent_id: str, + user_id: str, + organization_id: str, + session_id: str, + experience_type: str, + title: str, + content: str = "", + importance: float = 0.0, + credibility: float = 0.0, + evidence: str = "", + status: str = "pending", + created_by_id: Optional[str] = None, + ) -> PydanticSkillExperience: + """Insert one distilled experience in status 'pending' (default). + + The payload is validated through `SkillExperienceCreate` first so an + invalid `experience_type` (or over-length field) is rejected up front, + and importance/credibility are clamped into [0,1]. The DB column is a + plain String, so without this an invalid value would commit and only + blow up later in `to_pydantic()` / `list_experiences`. + """ + validated = SkillExperienceCreate( + agent_id=agent_id, + user_id=user_id, + organization_id=organization_id, + session_id=session_id, + experience_type=experience_type, + title=title, + content=content, + importance=importance, + credibility=credibility, + evidence=evidence, + status=status, + ) + row = SkillExperienceModel( + agent_id=validated.agent_id, + user_id=validated.user_id, + organization_id=validated.organization_id, + session_id=validated.session_id, + experience_type=validated.experience_type, + title=validated.title, + content=validated.content, + importance=validated.importance, + credibility=validated.credibility, + evidence=validated.evidence, + status=validated.status, + ) + if created_by_id is not None: + # Client attribution via the _created_by_id audit column โ€” this is + # what delete_by_client_id (the client erasure path) filters on. + row._set_created_and_updated_by_fields(created_by_id) + async with self.session_maker() as session: + await row.create(session) + return row.to_pydantic() + + @enforce_types + async def list_experiences( + self, + *, + agent_id: str, + user_id: Optional[str] = None, + status: Optional[str] = "pending", + limit: int = 100, + ids: Optional[List[str]] = None, + ) -> List[PydanticSkillExperience]: + """Return this agent's experiences, prioritized for evolution consumption. + + Filters by agent (always), optionally by user and status (default + 'pending'); `status=None` returns all statuses. Excludes soft-deleted + rows. + + `ids` scopes the result to an explicit id set (the experience curator passes + THIS evolution round's freshly-distilled experiences so a run only ever + sees its own batch, never the accumulated cross-round pending pool). + `ids=None` (default) applies no id filter; `ids=[]` is an explicit + "scope to nothing" and short-circuits to `[]` (so a round that distilled + zero experiences evolves over nothing rather than the whole pool). + + Ordering: `(importance * credibility) DESC` computed in SQL (the + priority), then `created_at DESC` as a stable tiebreak. + """ + # Explicit empty scope -> nothing (distinct from ids=None == "no filter"). + # Also avoids emitting a degenerate `IN ()` predicate. + if ids is not None and len(ids) == 0: + return [] + + priority = (SkillExperienceModel.importance * SkillExperienceModel.credibility) + preds = [ + SkillExperienceModel.agent_id == agent_id, + SkillExperienceModel.is_deleted.is_(False), + ] + if user_id is not None: + preds.append(SkillExperienceModel.user_id == user_id) + if status is not None: + preds.append(SkillExperienceModel.status == status) + if ids is not None: + preds.append(SkillExperienceModel.id.in_(ids)) + + stmt = ( + select(SkillExperienceModel) + .where(*preds) + .order_by(priority.desc(), SkillExperienceModel.created_at.desc()) + .limit(limit) + ) + async with self.session_maker() as session: + result = await session.execute(stmt) + rows = result.scalars().all() + return [row.to_pydantic() for row in rows] + + @enforce_types + async def mark_consumed( + self, + *, + ids: List[str], + run_id: str, + influenced_skill_ids: Optional[List[str]] = None, + ) -> int: + """Flip the given experiences to `status='consumed'`, stamping + `consumed_by` (and optional `influenced_skill_ids` lineage). + + Only flips rows still pending (idempotent; a second consumer can't + clobber an existing `consumed_by`). Returns the number updated. + """ + if not ids: + return 0 + return await self._set_status( + ids=ids, + status="consumed", + consumed_by=run_id, + influenced_skill_ids=influenced_skill_ids, + require_pending=True, + ) + + @enforce_types + async def mark_superseded( + self, + *, + ids: List[str], + agent_id: Optional[str] = None, + user_id: Optional[str] = None, + ) -> int: + """Soft-delete via status transition to 'superseded' (NOT is_deleted). + + Only rows still `pending` are transitioned (`require_pending`): an + experience that already did its job (`consumed`) must never be flipped to + `superseded`, so passing a mixed id set (e.g. a curator's "scoped minus + consumed" overflow) can only ever retire the still-pending leftovers. + + `agent_id` / `user_id`, when given, additionally constrain the update to + that owner โ€” so an id set that (defensively) contains a foreign owner's id + can never supersede a DIFFERENT (agent, user)'s experience. + `list_experiences(status='pending')` already excludes the result. Returns + the number updated. + """ + if not ids: + return 0 + return await self._set_status( + ids=ids, + status="superseded", + require_pending=True, + agent_id=agent_id, + user_id=user_id, + ) + + async def _set_status( + self, + *, + ids: List[str], + status: str, + consumed_by: Optional[str] = None, + influenced_skill_ids: Optional[List[str]] = None, + require_pending: bool = False, + agent_id: Optional[str] = None, + user_id: Optional[str] = None, + ) -> int: + """Load-modify-save each row so updated_at / Redis cache stay consistent + with the rest of the ORM (instead of a bulk UPDATE that bypasses both). + + When `require_pending` is set, only rows still in `status='pending'` are + transitioned (the rest are left untouched). `agent_id` / `user_id`, when + given, additionally scope the update to that owner. + """ + updated = 0 + async with self.session_maker() as session: + preds = [ + SkillExperienceModel.id.in_(ids), + SkillExperienceModel.is_deleted.is_(False), + ] + if require_pending: + preds.append(SkillExperienceModel.status == "pending") + if agent_id is not None: + preds.append(SkillExperienceModel.agent_id == agent_id) + if user_id is not None: + preds.append(SkillExperienceModel.user_id == user_id) + stmt = select(SkillExperienceModel).where(*preds) + result = await session.execute(stmt) + rows = result.scalars().all() + for row in rows: + row.status = status + if consumed_by is not None: + row.consumed_by = consumed_by + if influenced_skill_ids is not None: + row.influenced_skill_ids = influenced_skill_ids + # Pass an explicit UTC timestamp: the no-arg path in + # CommonSqlalchemyMetaMixins.set_updated_at uses datetime.UTC, + # which only exists on Python 3.11+ (this repo targets 3.10+). + row.set_updated_at(get_utc_time()) + session.add(row) + updated += 1 + await session.commit() + return updated + + @enforce_types + async def aggregate(self, *, ids: List[str]) -> Dict[str, float]: + """Summarize a set of experiences for the skill-evolution budget driver. + + Returns: + n -- total experiences in the set + n_worth_learning -- experiences of type 'worth_learning' + n_worth_avoiding -- experiences of type 'worth_avoiding' + sum_priority -- sum of importance*credibility over the set + + The Goal-3 curator maps `n_worth_avoiding -> n_high_fail` and + `n_worth_learning -> n_high_succ` so the existing C4 edit-budget formula + (avoid weighted heavier than learn) applies unchanged. + """ + if not ids: + return { + "n": 0, + "n_worth_learning": 0, + "n_worth_avoiding": 0, + "sum_priority": 0.0, + } + + async with self.session_maker() as session: + stmt = select(SkillExperienceModel).where( + SkillExperienceModel.id.in_(ids), + SkillExperienceModel.is_deleted.is_(False), + ) + result = await session.execute(stmt) + rows = result.scalars().all() + + n = len(rows) + n_worth_learning = 0 + n_worth_avoiding = 0 + sum_priority = 0.0 + for row in rows: + sum_priority += (row.importance or 0.0) * (row.credibility or 0.0) + if row.experience_type == "worth_learning": + n_worth_learning += 1 + elif row.experience_type == "worth_avoiding": + n_worth_avoiding += 1 + + return { + "n": n, + "n_worth_learning": n_worth_learning, + "n_worth_avoiding": n_worth_avoiding, + "sum_priority": sum_priority, + } + + @enforce_types + async def delete_by_user_id(self, user_id: str) -> int: + """Hard delete ALL experiences distilled from a user's sessions (erasure path). + + Experiences carry titles/content/evidence quotes distilled from the + user's verbatim conversations, so the irreversible + DELETE /users/{user_id}/memories purge must cover them like every other + memory table. Deletes regardless of status (pending/consumed/superseded). + """ + async with self.session_maker() as session: + stmt = delete(SkillExperienceModel).where( + SkillExperienceModel.user_id == user_id + ) + result = await session.execute(stmt) + await session.commit() + return int(result.rowcount or 0) + + @enforce_types + async def delete_by_client_id(self, actor: PydanticClient) -> int: + """Hard delete all experiences created by a client (erasure path). + + Attribution is via the `_created_by_id` audit column โ€” the table has no + dedicated client_id column. Called from + `ClientManager.delete_memories_by_client_id`. + """ + async with self.session_maker() as session: + stmt = delete(SkillExperienceModel).where( + SkillExperienceModel._created_by_id == actor.id + ) + result = await session.execute(stmt) + await session.commit() + return int(result.rowcount or 0) diff --git a/mirix/services/tool_manager.py b/mirix/services/tool_manager.py index 929d5a5ae..b5e014ba8 100644 --- a/mirix/services/tool_manager.py +++ b/mirix/services/tool_manager.py @@ -1,194 +1,194 @@ -import importlib -import warnings -from typing import List, Optional - -from mirix.constants import ( - ALL_TOOLS, - BASE_TOOLS, - CHAT_AGENT_TOOLS, - CORE_MEMORY_TOOLS, - EPISODIC_MEMORY_TOOLS, - EXTRAS_TOOLS, - KNOWLEDGE_VAULT_TOOLS, - MCP_TOOLS, - META_MEMORY_TOOLS, - PROCEDURAL_MEMORY_TOOLS, - RESOURCE_MEMORY_TOOLS, - SEMANTIC_MEMORY_TOOLS, - UNIVERSAL_MEMORY_TOOLS, -) -from mirix.functions.functions import derive_openai_json_schema, load_function_set - -# TODO: Remove this once we translate all of these to the ORM -from mirix.orm.errors import NoResultFound -from mirix.orm.tool import Tool as ToolModel -from mirix.schemas.client import Client as PydanticClient -from mirix.schemas.enums import ToolType -from mirix.schemas.tool import Tool as PydanticTool -from mirix.schemas.tool import ToolUpdate -from mirix.utils import enforce_types, printd - - -class ToolManager: - """Manager class to handle business logic related to Tools.""" - - def __init__(self): - # Fetching the db_context similarly as in OrganizationManager - from mirix.server.server import db_context - - self.session_maker = db_context - - # TODO: Refactor this across the codebase to use CreateTool instead of passing in a Tool object - @enforce_types - async def create_or_update_tool( - self, pydantic_tool: PydanticTool, actor: PydanticClient - ) -> PydanticTool: - """Create or update a tool (async).""" - tool = await self.get_tool_by_name(tool_name=pydantic_tool.name, actor=actor) - if tool: - update_data = pydantic_tool.model_dump(exclude_unset=True, exclude_none=True) - if update_data: - return await self.update_tool_by_id(tool.id, ToolUpdate(**update_data), actor) - printd( - "`create_or_update_tool` was called with name=%s but found existing tool with nothing to update.", - pydantic_tool.name, - ) - return tool - return await self.create_tool(pydantic_tool, actor=actor) - - @enforce_types - async def create_tool(self, pydantic_tool: PydanticTool, actor: PydanticClient) -> PydanticTool: - """Create a new tool (async).""" - async with self.session_maker() as session: - pydantic_tool.organization_id = actor.organization_id - if pydantic_tool.description is None: - pydantic_tool.description = pydantic_tool.json_schema.get("description", None) - tool_data = pydantic_tool.model_dump() - tool = ToolModel(**tool_data) - await tool.create(session, actor=actor) - return tool.to_pydantic() - - @enforce_types - async def get_tool_by_id(self, tool_id: str, actor: PydanticClient) -> PydanticTool: - """Fetch a tool by its ID (async).""" - async with self.session_maker() as session: - tool = await ToolModel.read(db_session=session, identifier=tool_id, actor=actor) - return tool.to_pydantic() - - @enforce_types - async def get_tool_by_name( - self, tool_name: str, actor: PydanticClient - ) -> Optional[PydanticTool]: - """Retrieve a tool by name (async).""" - try: - async with self.session_maker() as session: - tool = await ToolModel.read(db_session=session, name=tool_name, actor=actor) - return tool.to_pydantic() - except NoResultFound: - return None - - @enforce_types - async def list_tools( - self, - actor: PydanticClient, - cursor: Optional[str] = None, - limit: Optional[int] = 50, - ) -> List[PydanticTool]: - """List all tools with optional pagination using cursor and limit.""" - async with self.session_maker() as session: - tools = await ToolModel.list( - db_session=session, - cursor=cursor, - limit=limit, - organization_id=actor.organization_id, - ) - return [tool.to_pydantic() for tool in tools] - - @enforce_types - async def update_tool_by_id( - self, tool_id: str, tool_update: ToolUpdate, actor: PydanticClient - ) -> PydanticTool: - """Update a tool by its ID (async).""" - async with self.session_maker() as session: - tool = await ToolModel.read(db_session=session, identifier=tool_id, actor=actor) - update_data = tool_update.model_dump(exclude_none=True) - for key, value in update_data.items(): - setattr(tool, key, value) - if "source_code" in update_data.keys() and "json_schema" not in update_data.keys(): - pydantic_tool = tool.to_pydantic() - new_schema = derive_openai_json_schema(source_code=pydantic_tool.source_code) - tool.json_schema = new_schema - updated = await tool.update(db_session=session, actor=actor) - return updated.to_pydantic() - - @enforce_types - async def delete_tool_by_id(self, tool_id: str, actor: PydanticClient) -> None: - """Delete a tool by its ID.""" - async with self.session_maker() as session: - try: - tool = await ToolModel.read( - db_session=session, identifier=tool_id, actor=actor - ) - await tool.hard_delete(db_session=session, actor=actor) - except NoResultFound: - raise ValueError(f"Tool with id {tool_id} not found.") - - @enforce_types - async def upsert_base_tools(self, actor: PydanticClient) -> List[PydanticTool]: - """Add default tools in base.py (async).""" - functions_to_schema = {} - module_names = ["base", "memory_tools", "extras"] - - for module_name in module_names: - full_module_name = f"mirix.functions.function_sets.{module_name}" - try: - module = importlib.import_module(full_module_name) - except Exception as e: - raise e - - try: - functions_to_schema.update(load_function_set(module)) - except ValueError as e: - err = f"Error loading function set '{module_name}': {e}" - warnings.warn(err) - - tools = [] - for name, schema in functions_to_schema.items(): - if name in ALL_TOOLS: - if name in BASE_TOOLS: - tool_type = ToolType.MIRIX_CORE - tags = [tool_type.value] - elif ( - name - in CORE_MEMORY_TOOLS - + EPISODIC_MEMORY_TOOLS - + PROCEDURAL_MEMORY_TOOLS - + RESOURCE_MEMORY_TOOLS - + KNOWLEDGE_VAULT_TOOLS - + META_MEMORY_TOOLS - + SEMANTIC_MEMORY_TOOLS - + UNIVERSAL_MEMORY_TOOLS - + CHAT_AGENT_TOOLS - ): - tool_type = ToolType.MIRIX_MEMORY_CORE - tags = [tool_type.value] - elif name in EXTRAS_TOOLS: - tool_type = ToolType.MIRIX_EXTRA - tags = [tool_type.value] - elif name in MCP_TOOLS: - tool_type = ToolType.MIRIX_EXTRA - tags = [tool_type.value, "mcp_wrapper"] - else: - raise ValueError(f"Tool name {name} is not in the list of tool names") - - tool = await self.create_or_update_tool( - PydanticTool( - name=name, - tags=tags, - source_type="python", - tool_type=tool_type, - ), - actor=actor, - ) - tools.append(tool) - return tools +import importlib +import warnings +from typing import List, Optional + +from mirix.constants import ( + ALL_TOOLS, + BASE_TOOLS, + CHAT_AGENT_TOOLS, + CORE_MEMORY_TOOLS, + EPISODIC_MEMORY_TOOLS, + EXTRAS_TOOLS, + KNOWLEDGE_VAULT_TOOLS, + MCP_TOOLS, + META_MEMORY_TOOLS, + SKILL_TOOLS, + RESOURCE_MEMORY_TOOLS, + SEMANTIC_MEMORY_TOOLS, + UNIVERSAL_MEMORY_TOOLS, +) +from mirix.functions.functions import derive_openai_json_schema, load_function_set + +# TODO: Remove this once we translate all of these to the ORM +from mirix.orm.errors import NoResultFound +from mirix.orm.tool import Tool as ToolModel +from mirix.schemas.client import Client as PydanticClient +from mirix.schemas.enums import ToolType +from mirix.schemas.tool import Tool as PydanticTool +from mirix.schemas.tool import ToolUpdate +from mirix.utils import enforce_types, printd + + +class ToolManager: + """Manager class to handle business logic related to Tools.""" + + def __init__(self): + # Fetching the db_context similarly as in OrganizationManager + from mirix.server.server import db_context + + self.session_maker = db_context + + # TODO: Refactor this across the codebase to use CreateTool instead of passing in a Tool object + @enforce_types + async def create_or_update_tool( + self, pydantic_tool: PydanticTool, actor: PydanticClient + ) -> PydanticTool: + """Create or update a tool (async).""" + tool = await self.get_tool_by_name(tool_name=pydantic_tool.name, actor=actor) + if tool: + update_data = pydantic_tool.model_dump(exclude_unset=True, exclude_none=True) + if update_data: + return await self.update_tool_by_id(tool.id, ToolUpdate(**update_data), actor) + printd( + "`create_or_update_tool` was called with name=%s but found existing tool with nothing to update.", + pydantic_tool.name, + ) + return tool + return await self.create_tool(pydantic_tool, actor=actor) + + @enforce_types + async def create_tool(self, pydantic_tool: PydanticTool, actor: PydanticClient) -> PydanticTool: + """Create a new tool (async).""" + async with self.session_maker() as session: + pydantic_tool.organization_id = actor.organization_id + if pydantic_tool.description is None: + pydantic_tool.description = pydantic_tool.json_schema.get("description", None) + tool_data = pydantic_tool.model_dump() + tool = ToolModel(**tool_data) + await tool.create(session, actor=actor) + return tool.to_pydantic() + + @enforce_types + async def get_tool_by_id(self, tool_id: str, actor: PydanticClient) -> PydanticTool: + """Fetch a tool by its ID (async).""" + async with self.session_maker() as session: + tool = await ToolModel.read(db_session=session, identifier=tool_id, actor=actor) + return tool.to_pydantic() + + @enforce_types + async def get_tool_by_name( + self, tool_name: str, actor: PydanticClient + ) -> Optional[PydanticTool]: + """Retrieve a tool by name (async).""" + try: + async with self.session_maker() as session: + tool = await ToolModel.read(db_session=session, name=tool_name, actor=actor) + return tool.to_pydantic() + except NoResultFound: + return None + + @enforce_types + async def list_tools( + self, + actor: PydanticClient, + cursor: Optional[str] = None, + limit: Optional[int] = 50, + ) -> List[PydanticTool]: + """List all tools with optional pagination using cursor and limit.""" + async with self.session_maker() as session: + tools = await ToolModel.list( + db_session=session, + cursor=cursor, + limit=limit, + organization_id=actor.organization_id, + ) + return [tool.to_pydantic() for tool in tools] + + @enforce_types + async def update_tool_by_id( + self, tool_id: str, tool_update: ToolUpdate, actor: PydanticClient + ) -> PydanticTool: + """Update a tool by its ID (async).""" + async with self.session_maker() as session: + tool = await ToolModel.read(db_session=session, identifier=tool_id, actor=actor) + update_data = tool_update.model_dump(exclude_none=True) + for key, value in update_data.items(): + setattr(tool, key, value) + if "source_code" in update_data.keys() and "json_schema" not in update_data.keys(): + pydantic_tool = tool.to_pydantic() + new_schema = derive_openai_json_schema(source_code=pydantic_tool.source_code) + tool.json_schema = new_schema + updated = await tool.update(db_session=session, actor=actor) + return updated.to_pydantic() + + @enforce_types + async def delete_tool_by_id(self, tool_id: str, actor: PydanticClient) -> None: + """Delete a tool by its ID.""" + async with self.session_maker() as session: + try: + tool = await ToolModel.read( + db_session=session, identifier=tool_id, actor=actor + ) + await tool.hard_delete(db_session=session, actor=actor) + except NoResultFound: + raise ValueError(f"Tool with id {tool_id} not found.") + + @enforce_types + async def upsert_base_tools(self, actor: PydanticClient) -> List[PydanticTool]: + """Add default tools in base.py (async).""" + functions_to_schema = {} + module_names = ["base", "memory_tools", "extras"] + + for module_name in module_names: + full_module_name = f"mirix.functions.function_sets.{module_name}" + try: + module = importlib.import_module(full_module_name) + except Exception as e: + raise e + + try: + functions_to_schema.update(load_function_set(module)) + except ValueError as e: + err = f"Error loading function set '{module_name}': {e}" + warnings.warn(err) + + tools = [] + for name, schema in functions_to_schema.items(): + if name in ALL_TOOLS: + if name in BASE_TOOLS: + tool_type = ToolType.MIRIX_CORE + tags = [tool_type.value] + elif ( + name + in CORE_MEMORY_TOOLS + + EPISODIC_MEMORY_TOOLS + + SKILL_TOOLS + + RESOURCE_MEMORY_TOOLS + + KNOWLEDGE_VAULT_TOOLS + + META_MEMORY_TOOLS + + SEMANTIC_MEMORY_TOOLS + + UNIVERSAL_MEMORY_TOOLS + + CHAT_AGENT_TOOLS + ): + tool_type = ToolType.MIRIX_MEMORY_CORE + tags = [tool_type.value] + elif name in EXTRAS_TOOLS: + tool_type = ToolType.MIRIX_EXTRA + tags = [tool_type.value] + elif name in MCP_TOOLS: + tool_type = ToolType.MIRIX_EXTRA + tags = [tool_type.value, "mcp_wrapper"] + else: + raise ValueError(f"Tool name {name} is not in the list of tool names") + + tool = await self.create_or_update_tool( + PydanticTool( + name=name, + tags=tags, + source_type="python", + tool_type=tool_type, + ), + actor=actor, + ) + tools.append(tool) + return tools diff --git a/mirix/services/user_manager.py b/mirix/services/user_manager.py index 45f50144c..87c2dd785 100755 --- a/mirix/services/user_manager.py +++ b/mirix/services/user_manager.py @@ -229,6 +229,8 @@ async def delete_memories_by_user_id(self, user_id: str): - KnowledgeVaultManager.delete_by_user_id() - MessageManager.delete_by_user_id() - BlockManager.delete_by_user_id() + - ConversationMessageManager.delete_by_user_id() (verbatim session transcripts) + - SkillExperienceManager.delete_by_user_id() (experiences distilled from them) 2. Each manager handles: - Bulk database deletion - Redis cache cleanup @@ -245,12 +247,14 @@ async def delete_memories_by_user_id(self, user_id: str): # Import managers from mirix.services.block_manager import BlockManager + from mirix.services.conversation_message_manager import ConversationMessageManager from mirix.services.episodic_memory_manager import EpisodicMemoryManager from mirix.services.knowledge_vault_manager import KnowledgeVaultManager from mirix.services.message_manager import MessageManager from mirix.services.procedural_memory_manager import ProceduralMemoryManager from mirix.services.resource_memory_manager import ResourceMemoryManager from mirix.services.semantic_memory_manager import SemanticMemoryManager + from mirix.services.skill_experience_manager import SkillExperienceManager # Initialize managers episodic_manager = EpisodicMemoryManager() @@ -260,6 +264,8 @@ async def delete_memories_by_user_id(self, user_id: str): knowledge_manager = KnowledgeVaultManager() message_manager = MessageManager() block_manager = BlockManager() + conversation_manager = ConversationMessageManager() + skill_experience_manager = SkillExperienceManager() # Use managers' bulk delete methods try: @@ -285,6 +291,15 @@ async def delete_memories_by_user_id(self, user_id: str): block_count = await block_manager.delete_by_user_id(user_id=user_id) logger.debug("Bulk deleted %d blocks", block_count) + # Verbatim session transcripts and the experiences distilled from + # them contain the user's raw conversation content โ€” erasure must + # cover them like every other memory table. + conversation_count = await conversation_manager.delete_by_user_id(user_id=user_id) + logger.debug("Bulk deleted %d conversation turns", conversation_count) + + experience_count = await skill_experience_manager.delete_by_user_id(user_id=user_id) + logger.debug("Bulk deleted %d skill experiences", experience_count) + # Clear message_ids from ALL agents in PostgreSQL (messages are user-scoped, agents are client-scoped) # IMPORTANT: Keep the first message (system message) as agents need it to function # We need to clear message_ids from all agents that might have cached this user's messages @@ -328,7 +343,8 @@ async def delete_memories_by_user_id(self, user_id: str): logger.info( "Bulk deleted all memories for user %s: " - "%d episodic, %d semantic, %d procedural, %d resource, %d knowledge_vault, %d messages, %d blocks " + "%d episodic, %d semantic, %d procedural, %d resource, %d knowledge_vault, %d messages, %d blocks, " + "%d conversation turns, %d skill experiences " "(user record preserved)", user_id, episodic_count, @@ -338,6 +354,8 @@ async def delete_memories_by_user_id(self, user_id: str): knowledge_count, message_count, block_count, + conversation_count, + experience_count, ) except Exception as e: logger.error("Failed to bulk delete memories for user %s: %s", user_id, e) diff --git a/samples/add_test_memory.py b/samples/add_test_memory.py index e797278cb..bfb945596 100755 --- a/samples/add_test_memory.py +++ b/samples/add_test_memory.py @@ -51,6 +51,7 @@ async def test_core_memory(client: MirixClient, user_id: str, filter_tags: Optio ], chaining=True, filter_tags=filter_tags, + session_id="sample-procedural-session", occurred_at="2025-11-16T10:30:00", ) logger.info("โœ… Core memory added successfully: %s", result.get("success", False)) @@ -128,7 +129,7 @@ async def test_procedural_memory(client: MirixClient, user_id: str, filter_tags: { "role": "assistant", "content": [ - {"type": "text", "text": "I've saved your deployment workflow procedure with 5 steps."} + {"type": "text", "text": "I've saved your deployment workflow."} ], }, ], diff --git a/samples/langgraph_integration.py b/samples/langgraph_integration.py index 1893edcae..07adfbba5 100644 --- a/samples/langgraph_integration.py +++ b/samples/langgraph_integration.py @@ -106,7 +106,7 @@ def format_memories_for_prompt(memories: dict) -> str: elif memory_type == "semantic": text = item.get("name", "") + ": " + item.get("summary", "") elif memory_type == "procedural": - text = item.get("summary", "") + text = item.get("name", "") + ": " + item.get("description", "") elif memory_type == "resource": text = item.get("title", "") + ": " + item.get("summary", "") elif memory_type == "knowledge": diff --git a/samples/mirix_memory_viewer.py b/samples/mirix_memory_viewer.py index 1303d55d7..7b9060e5f 100644 --- a/samples/mirix_memory_viewer.py +++ b/samples/mirix_memory_viewer.py @@ -124,7 +124,8 @@ async def main(): elif memory_type == "procedural": print(f" Name: {item.get('name', 'N/A')}") - print(f" Procedure: {item.get('procedure_text', 'N/A')[:100]}...") + print(f" Description: {item.get('description', 'N/A')[:100]}...") + print(f" Instructions: {item.get('instructions', 'N/A')[:100]}...") elif memory_type == "resources": print(f" Title: {item.get('title', 'N/A')}") diff --git a/samples/run_client.py b/samples/run_client.py index 739cf2c43..98dfbc7a2 100644 --- a/samples/run_client.py +++ b/samples/run_client.py @@ -109,21 +109,19 @@ def print_memories(memories): print() elif memory_type == "procedural": - # Procedural memory: procedure with summary (API doesn't return full details) + # Procedural memory: skill-shaped rows. for i, item in enumerate(items, 1): - # API returns 'summary' not 'description' - summary = item.get("summary", item.get("description", "N/A")) + name = item.get("name") or "N/A" + description = item.get("description") or "N/A" + instructions = item.get("instructions") or "" entry_type = item.get("entry_type", "N/A") - # API doesn't return 'steps' in retrieve endpoint (only id, entry_type, summary) - steps = item.get("steps", []) - - print(f" [{i}] [{entry_type}] {summary}") - if steps: - print(f" Steps ({len(steps)}):") - for step_num, step in enumerate(steps, 1): - print(f" {step_num}. {step}") - else: - print(" Steps: Not included in response (use search API for full details)") + version = item.get("version") + + suffix = f" v{version}" if version else "" + print(f" [{i}] [{entry_type}{suffix}] {name}") + print(f" Description: {description}") + if instructions: + print(f" Instructions: {instructions[:200]}...") print() elif memory_type == "semantic": diff --git a/scripts/migrate_add_agent_trigger_state.sql b/scripts/migrate_add_agent_trigger_state.sql new file mode 100644 index 000000000..e5f7bbcf4 --- /dev/null +++ b/scripts/migrate_add_agent_trigger_state.sql @@ -0,0 +1,59 @@ +-- Migration: create agent_trigger_state table. +-- +-- Run once on existing databases. New databases get this via SQLAlchemy +-- create_all at server startup, so there is no separate "phase 2" here. +-- +-- Purpose: persist per-(agent, user, trigger_type) cursors for interval- +-- driven memory triggers (e.g. "fire procedural extraction every N sessions"). +-- Only the last-fire cursor is stored; the "how many sessions since" counter +-- is derived from the messages table at read time. +-- +-- Entire migration is idempotent and safe to re-run. + +BEGIN; + +CREATE TABLE IF NOT EXISTS agent_trigger_state ( + id VARCHAR PRIMARY KEY, + organization_id VARCHAR REFERENCES organizations(id), + user_id VARCHAR NOT NULL REFERENCES users(id), + agent_id VARCHAR NOT NULL REFERENCES agents(id) ON DELETE CASCADE, + trigger_type VARCHAR(64) NOT NULL, + last_fired_at TIMESTAMPTZ NULL, + last_fired_session_id VARCHAR(64) NULL, + -- Session_ids tied at the watermark timestamp. Stored so the next window + -- can use `created_at >= last_fired_at` with session-level tie-break, + -- instead of `created_at > last_fired_at` which silently drops any row + -- that committed at the exact same microsecond as our SELECT. + last_fired_tied_session_ids JSON NULL, + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + is_deleted BOOLEAN NOT NULL DEFAULT FALSE, + _created_by_id VARCHAR NULL, + _last_updated_by_id VARCHAR NULL +); + +-- Back-compat: if an older deployment created this table before the tied +-- set was added, fill it in. Idempotent. +ALTER TABLE agent_trigger_state + ADD COLUMN IF NOT EXISTS last_fired_tied_session_ids JSON NULL; + +-- One live cursor per (agent, user, trigger_type). The app always addresses +-- a row by this triple, so the uniqueness closes an insert-race between +-- concurrent workers. +DO $$ +BEGIN + IF NOT EXISTS ( + SELECT 1 FROM pg_constraint + WHERE conname = 'uq_agent_trigger_state_agent_user_type' + AND conrelid = 'agent_trigger_state'::regclass + ) THEN + ALTER TABLE agent_trigger_state + ADD CONSTRAINT uq_agent_trigger_state_agent_user_type + UNIQUE (agent_id, user_id, trigger_type); + END IF; +END$$; + +CREATE INDEX IF NOT EXISTS ix_agent_trigger_state_agent_user_type + ON agent_trigger_state (agent_id, user_id, trigger_type); + +COMMIT; diff --git a/scripts/migrate_add_conversation_message.sql b/scripts/migrate_add_conversation_message.sql new file mode 100644 index 000000000..0e2a9b4b6 --- /dev/null +++ b/scripts/migrate_add_conversation_message.sql @@ -0,0 +1,72 @@ +-- Migration: create the conversation_message table + indexes. +-- Run once on existing databases. New databases get this via SQLAlchemy +-- create_all at server startup (the ORM model is registered in mirix/orm/__init__.py). +-- +-- The conversation_message store holds ONLY external conversation turns that +-- arrived through the memory-add API carrying a session_id, with their real +-- user/assistant roles preserved. It is the single source the procedural-memory +-- (skill) distiller reads โ€” separate from the agent-loop `messages` table. +-- +-- This migration is split into two phases because CREATE INDEX CONCURRENTLY +-- cannot run inside a transaction block. +-- +-- PHASE 1 (transactional): create the table with its columns, the CHECK +-- constraint on session_id (matching mirix.schemas.message), and the +-- non-concurrent indexes implied by column-level `index=True`. +-- +-- PHASE 2 (must run OUTSIDE a transaction, not inside psql -1): build the +-- composite indexes CONCURRENTLY so they do not take a write-blocking ACCESS +-- EXCLUSIVE lock on a large table. +-- +-- Phase 1 is idempotent: safe to re-run (e.g. if phase 2 failed and the whole +-- migration is replayed). All object creation guards against existing state. + +BEGIN; + +CREATE TABLE IF NOT EXISTS conversation_message ( + id VARCHAR PRIMARY KEY, + session_id VARCHAR(64) NOT NULL, + role VARCHAR NOT NULL, + content TEXT NOT NULL DEFAULT '', + distilled_at TIMESTAMPTZ, + user_id VARCHAR NOT NULL REFERENCES users(id), + organization_id VARCHAR REFERENCES organizations(id), + -- CommonSqlalchemyMetaMixins columns, matching sibling tables. + created_at TIMESTAMPTZ DEFAULT now(), + updated_at TIMESTAMPTZ DEFAULT now(), + is_deleted BOOLEAN DEFAULT FALSE, + _created_by_id VARCHAR, + _last_updated_by_id VARCHAR +); + +-- CHECK constraint matches mirix.schemas.message._validate_session_id. Unlike +-- the `messages` table, session_id here is NOT NULL (this store only holds +-- session'd turns), so the NULL branch is omitted. Guarded so a re-run is a +-- no-op. +DO $$ +BEGIN + IF NOT EXISTS ( + SELECT 1 FROM pg_constraint + WHERE conname = 'ck_conversation_message_session_id_format' + AND conrelid = 'conversation_message'::regclass + ) THEN + ALTER TABLE conversation_message + ADD CONSTRAINT ck_conversation_message_session_id_format + CHECK (session_id ~ '^[A-Za-z0-9_-]{1,64}$'); + END IF; +END$$; + +-- Non-concurrent single-column index from the ORM column-level `index=True` on +-- session_id. Small/cheap; kept in phase 1. +CREATE INDEX IF NOT EXISTS ix_conversation_message_session_id + ON conversation_message (session_id); + +COMMIT; + +-- PHASE 2 โ€” the composite indexes are built CONCURRENTLY in a SEPARATE, +-- EXECUTABLE file: scripts/migrate_add_conversation_message_phase2.sql. They +-- live there (rather than commented out here) because CREATE INDEX CONCURRENTLY +-- cannot run inside a transaction block, so they must NOT share this file's +-- BEGIN/COMMIT. Run that file next, OUTSIDE a transaction (plain `psql -f`, NOT +-- `psql -1`/--single-transaction). New databases get these indexes via +-- SQLAlchemy create_all and need neither migration file. diff --git a/scripts/migrate_add_conversation_message_phase2.sql b/scripts/migrate_add_conversation_message_phase2.sql new file mode 100644 index 000000000..b7d76918d --- /dev/null +++ b/scripts/migrate_add_conversation_message_phase2.sql @@ -0,0 +1,30 @@ +-- Migration PHASE 2 for conversation_message: build the composite indexes +-- CONCURRENTLY. Run this AFTER scripts/migrate_add_conversation_message.sql. +-- +-- IMPORTANT: run OUTSIDE a transaction โ€” plain `psql -f`, NOT +-- `psql -1`/--single-transaction. CREATE INDEX CONCURRENTLY cannot run inside a +-- transaction block. CONCURRENTLY avoids taking a write-blocking ACCESS +-- EXCLUSIVE lock on the table while the index builds. +-- +-- New databases get these indexes via SQLAlchemy create_all at server startup +-- (the ORM model in mirix/orm/conversation_message.py declares them), so only +-- existing databases being migrated in place need this file. +-- +-- If a build fails midway, Postgres leaves an INVALID index behind: drop it +-- (DROP INDEX ) and re-run. + +-- Primary access pattern: list/seal/order this (org, user)'s sessions by +-- first-appearance time. +CREATE INDEX CONCURRENTLY IF NOT EXISTS + ix_conversation_message_org_user_session_created + ON conversation_message (organization_id, user_id, session_id, created_at); + +-- Accelerates the per-session ascending fetch in list_turns_for_session. +CREATE INDEX CONCURRENTLY IF NOT EXISTS + ix_conversation_message_session_created + ON conversation_message (session_id, created_at); + +-- Single-column organization index (mirrors the ORM's pg-only org index). +CREATE INDEX CONCURRENTLY IF NOT EXISTS + ix_conversation_message_organization_id + ON conversation_message (organization_id); diff --git a/scripts/migrate_add_message_session_id.sql b/scripts/migrate_add_message_session_id.sql new file mode 100644 index 000000000..346c5b64a --- /dev/null +++ b/scripts/migrate_add_message_session_id.sql @@ -0,0 +1,54 @@ +-- Migration: add top-level session_id column + index on messages. +-- Run once on existing databases. New databases get this via SQLAlchemy create_all. +-- +-- This migration is split into two phases because CREATE INDEX CONCURRENTLY +-- cannot run inside a transaction block. +-- +-- PHASE 1 (transactional): add the column and a CHECK constraint so any writer +-- is immediately bound by the format rules. Runs quickly; acquires a short +-- ACCESS EXCLUSIVE lock only for the DDL itself. +-- +-- PHASE 2 (must run OUTSIDE a transaction, not inside psql -1): +-- CREATE INDEX CONCURRENTLY IF NOT EXISTS ix_messages_agent_session_created_at +-- ON messages (agent_id, session_id, created_at); +-- +-- CONCURRENTLY avoids the write-blocking ACCESS EXCLUSIVE lock on large tables. +-- If the index build fails midway, Postgres leaves an INVALID index behind; drop +-- it and retry. + +-- Phase 1 is idempotent: safe to re-run (e.g. if phase 2 failed and the whole +-- migration is replayed). Column/constraint additions all guard against +-- existing state. + +BEGIN; + +ALTER TABLE messages + ADD COLUMN IF NOT EXISTS session_id VARCHAR(64); + +-- CHECK constraint matches mirix.schemas.message._validate_session_id. +-- The constraint is NOT VALID first so adding it on a large table does not +-- rewrite existing rows; legacy rows all have NULL session_id and satisfy it. +-- Guarded so a re-run (after a partial failure) is a no-op. +DO $$ +BEGIN + IF NOT EXISTS ( + SELECT 1 FROM pg_constraint + WHERE conname = 'ck_messages_session_id_format' + AND conrelid = 'messages'::regclass + ) THEN + ALTER TABLE messages + ADD CONSTRAINT ck_messages_session_id_format + CHECK (session_id IS NULL OR session_id ~ '^[A-Za-z0-9_-]{1,64}$') + NOT VALID; + END IF; +END$$; + +-- Validate the constraint in a separate step so it only takes SHARE UPDATE +-- EXCLUSIVE (reads + writes continue; no concurrent DDL only). +-- VALIDATE is a no-op if the constraint is already validated. +ALTER TABLE messages VALIDATE CONSTRAINT ck_messages_session_id_format; + +COMMIT; + +-- PHASE 2 โ€” run scripts/migrate_add_message_session_id_phase2.sql separately, +-- NOT inside a transaction (do not combine with this file via `psql -1`). diff --git a/scripts/migrate_add_message_session_id_phase2.sql b/scripts/migrate_add_message_session_id_phase2.sql new file mode 100644 index 000000000..4746ce771 --- /dev/null +++ b/scripts/migrate_add_message_session_id_phase2.sql @@ -0,0 +1,11 @@ +-- Phase 2 of the session_id migration โ€” run AFTER migrate_add_message_session_id.sql. +-- MUST run outside a transaction (no BEGIN/COMMIT here, and do NOT run via `psql -1`). +-- Use e.g.: +-- psql "$DATABASE_URL" -f scripts/migrate_add_message_session_id_phase2.sql +-- +-- CONCURRENTLY avoids the write-blocking ACCESS EXCLUSIVE lock on large tables. +-- If the build fails midway, Postgres leaves an INVALID index behind; drop it and retry: +-- DROP INDEX CONCURRENTLY IF EXISTS ix_messages_agent_session_created_at; + +CREATE INDEX CONCURRENTLY IF NOT EXISTS ix_messages_agent_session_created_at + ON messages (agent_id, session_id, created_at); diff --git a/scripts/migrate_add_skill_experience.sql b/scripts/migrate_add_skill_experience.sql new file mode 100644 index 000000000..0a2f83768 --- /dev/null +++ b/scripts/migrate_add_skill_experience.sql @@ -0,0 +1,53 @@ +-- Migration: create skill_experience table. +-- +-- Run once on existing databases. New databases get this via SQLAlchemy +-- create_all at server startup (the ORM class is imported in +-- mirix/orm/__init__.py), so there is no separate "phase 2" here. +-- +-- Purpose: durable, general store for transferable EXPERIENCES distilled from +-- a single work session's transcript. Each experience is either +-- 'worth_learning' or 'worth_avoiding', scored by importance/credibility in +-- [0,1]. Consumed every N sessions by the skill-evolution run, ordered +-- by importance*credibility. Experiences flow pending -> consumed | superseded. +-- +-- Entire migration is idempotent and safe to re-run. + +BEGIN; + +CREATE TABLE IF NOT EXISTS skill_experience ( + id VARCHAR PRIMARY KEY, + organization_id VARCHAR REFERENCES organizations(id), + user_id VARCHAR NOT NULL REFERENCES users(id), + agent_id VARCHAR NOT NULL REFERENCES agents(id) ON DELETE CASCADE, + session_id VARCHAR NOT NULL, + -- experience_type / status are plain strings (NOT pg ENUMs); their value + -- spaces are validated in mirix/schemas/skill_experience.py, so adding a + -- value never requires a DB migration. + experience_type VARCHAR NOT NULL, + title VARCHAR NOT NULL, + -- content/evidence/importance/credibility are NOT NULL with defaults to + -- match the ORM + the pydantic full schema, which treat them as required. + -- A NULL here would break to_pydantic(). + content TEXT NOT NULL DEFAULT '', + importance DOUBLE PRECISION NOT NULL DEFAULT 0.0, + credibility DOUBLE PRECISION NOT NULL DEFAULT 0.0, + evidence TEXT NOT NULL DEFAULT '', + status VARCHAR NOT NULL DEFAULT 'pending', + consumed_by VARCHAR NULL, + influenced_skill_ids JSON NULL, + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + is_deleted BOOLEAN NOT NULL DEFAULT FALSE, + _created_by_id VARCHAR NULL, + _last_updated_by_id VARCHAR NULL +); + +-- Primary access pattern: list_experiences(agent_id, status='pending'). +CREATE INDEX IF NOT EXISTS ix_skill_experience_agent_status + ON skill_experience (agent_id, status); + +-- Organization-level query optimization (mirrors procedural_memory). +CREATE INDEX IF NOT EXISTS ix_skill_experience_organization_id + ON skill_experience (organization_id); + +COMMIT; diff --git a/scripts/migrate_backfill_procedural_scope.sql b/scripts/migrate_backfill_procedural_scope.sql new file mode 100644 index 000000000..4906152c2 --- /dev/null +++ b/scripts/migrate_backfill_procedural_scope.sql @@ -0,0 +1,69 @@ +-- Migration: backfill procedural_memory.filter_tags 'scope' from the owning client. +-- +-- Run ONCE on existing Postgres databases, AFTER migrate_procedural_to_skill.sql. +-- Postgres-only (jsonb operators). SQLite deployments get their schema from +-- SQLAlchemy create_all and have no multi-client scope to backfill (see header note). +-- +-- Problem: scope-based read authorization filters procedural skills on +-- filter_tags->>'scope' (mirix/database/filter_tags_query.py). Rows written before +-- scope-tagging existed -- or created internally by skill-evolve without the REST +-- scope-injection -- have no 'scope' key, so they are invisible to every scoped +-- reader. This reconstructs the scope from the row's owning client +-- (procedural_memory.client_id -> clients.write_scope), which is exactly the value +-- the REST layer stamps into filter_tags['scope'] on create. +-- +-- Data shapes handled: the ORM column is generic JSON with SQLAlchemy's default +-- none_as_null=False, so a Python None is stored as the JSON scalar 'null' (NOT +-- SQL NULL). jsonb_set() raises "cannot set path in scalar" on such rows, so any +-- non-object value (SQL NULL, JSON null, or another scalar) is replaced with '{}' +-- before setting the key. Verified against live rows where every unscoped row was +-- the JSON scalar 'null'. +-- +-- Idempotent: only touches rows whose filter_tags has no 'scope' key, so re-running +-- is a no-op. Runs in one transaction (pure UPDATE, no DDL, no CONCURRENTLY). + +BEGIN; + +-- Fixable rows: filter_tags has no 'scope' key AND the owning client has a non-null +-- write_scope to copy. jsonb_set(..., create_if_missing => true) preserves any other +-- keys an object row already has; non-object rows are normalized to '{}' first, +-- casting json->jsonb->json around the mutation. +-- +-- Deliberately LEFT UNTOUCHED (reported below, never guessed): +-- * rows with client_id IS NULL -- no client to derive a scope from +-- * rows whose client.write_scope IS NULL -- read-only client, no scope to donate +-- Stamping a guessed scope would be a silent authorization change, so unresolved rows +-- stay unscoped (invisible) rather than mis-scoped (leaked into the wrong scope). +UPDATE procedural_memory pm +SET filter_tags = jsonb_set( + CASE + WHEN pm.filter_tags IS NULL THEN '{}'::jsonb + WHEN jsonb_typeof(pm.filter_tags::jsonb) <> 'object' THEN '{}'::jsonb + ELSE pm.filter_tags::jsonb + END, + '{scope}', + to_jsonb(c.write_scope), + true + )::json +FROM clients c +WHERE pm.client_id = c.id + AND c.write_scope IS NOT NULL + AND (pm.filter_tags IS NULL + OR jsonb_typeof(pm.filter_tags::jsonb) <> 'object' + OR NOT jsonb_exists(pm.filter_tags::jsonb, 'scope')); + +-- Report rows still lacking a scope after the backfill (NULL client_id or read-only +-- client). Operators must scope these by hand if they need to be readable. +DO $$ +DECLARE + unresolved INTEGER; +BEGIN + SELECT COUNT(*) INTO unresolved + FROM procedural_memory pm + WHERE pm.filter_tags IS NULL + OR jsonb_typeof(pm.filter_tags::jsonb) <> 'object' + OR NOT jsonb_exists(pm.filter_tags::jsonb, 'scope'); + RAISE NOTICE 'procedural_memory rows still without a scope tag after backfill: %', unresolved; +END $$; + +COMMIT; diff --git a/scripts/migrate_procedural_to_skill.sql b/scripts/migrate_procedural_to_skill.sql new file mode 100644 index 000000000..9662bcda9 --- /dev/null +++ b/scripts/migrate_procedural_to_skill.sql @@ -0,0 +1,120 @@ +-- Migration: Convert procedural memory from flat format to skill-based format +-- Run this ONCE on existing databases before deploying the skill-evolve branch. +-- For new databases, the schema is created automatically by SQLAlchemy create_all. + +BEGIN; + +-- Step 1: Add new columns with safe defaults +ALTER TABLE procedural_memory ADD COLUMN IF NOT EXISTS name VARCHAR NOT NULL DEFAULT ''; +ALTER TABLE procedural_memory ADD COLUMN IF NOT EXISTS triggers JSONB DEFAULT '[]'; +ALTER TABLE procedural_memory ADD COLUMN IF NOT EXISTS examples JSONB DEFAULT '[]'; +ALTER TABLE procedural_memory ADD COLUMN IF NOT EXISTS version VARCHAR NOT NULL DEFAULT '0.1.0'; + +-- Step 2: Generate name from summary (slugify: lowercase, strip punctuation, spaces to hyphens, truncate). +-- Guard against NULL/empty summary by falling back to "skill-" so the +-- NOT NULL column constraint is never violated. The slug and the fallback MUST +-- be one COALESCE in a single statement: a two-statement "set NULL, then patch +-- NULLs" approach aborts the transaction on rows whose slug is empty (e.g. +-- summary=' ' or summary='!!!'), because NOT NULL is checked per row at +-- assignment time โ€” the patch statement never gets to run. +UPDATE procedural_memory +SET name = COALESCE( + NULLIF( + LOWER( + REGEXP_REPLACE( + REGEXP_REPLACE( + LEFT(TRIM(COALESCE(summary, '')), 60), + '[^a-zA-Z0-9\s-]', '', 'g' + ), + '\s+', '-', 'g' + ) + ), + '' + ), + 'skill-' || RIGHT(id, 8) +) +WHERE name = ''; + +-- Step 2b: Resolve slug collisions deterministically by suffixing -N, using +-- created_at (then id) as the tie-breaker. The first occurrence keeps the +-- bare slug; subsequent occurrences get -2, -3, etc. This runs before the +-- unique index so the unique constraint creation cannot fail on legacy data. +-- +-- A suffixed name can itself collide with a PRE-EXISTING slug (two 'task' +-- rows produce 'task-2' while a 'task 2' summary already slugged to +-- 'task-2'), so a single pass is not enough: loop until no row needs a +-- rename. Terminates because renamed names strictly lengthen each pass +-- (two renamed rows can never collide with each other โ€” same partition + +-- same occurrence is impossible โ€” only with fixed-length pre-existing +-- names, which they eventually outgrow). +DO $$ +DECLARE + renamed INTEGER; +BEGIN + LOOP + WITH ranked AS ( + SELECT + id, + ROW_NUMBER() OVER ( + PARTITION BY organization_id, user_id, name + ORDER BY created_at, id + ) AS occurrence + FROM procedural_memory + ) + UPDATE procedural_memory pm + SET name = pm.name || '-' || ranked.occurrence + FROM ranked + WHERE pm.id = ranked.id + AND ranked.occurrence > 1; + GET DIAGNOSTICS renamed = ROW_COUNT; + EXIT WHEN renamed = 0; + END LOOP; +END $$; + +-- Step 3: Change steps column type from JSON to TEXT first. This MUST happen +-- before the textification UPDATE below: assigning a text expression to a +-- json column has no assignment cast in Postgres and fails at plan time +-- with SQLSTATE 42804, which would abort the whole transaction. +ALTER TABLE procedural_memory ALTER COLUMN steps TYPE TEXT USING steps::TEXT; + +-- Step 3b: Convert steps from a JSON-array text form to plain text +-- (newline-joined). The column is TEXT now, so ::jsonb re-parses the +-- preserved JSON text and the assignment is text-to-text. +UPDATE procedural_memory +SET steps = ARRAY_TO_STRING( + ARRAY(SELECT jsonb_array_elements_text(steps::jsonb)), + E'\n' +) +WHERE steps IS NOT NULL AND steps LIKE '[%'; + +-- Step 3c: Normalize legacy entry_type values into the closed skill set +-- {workflow, guide, script}. Historically entry_type was free-form LLM +-- output ("process", "how-to", ...); the application now validates against +-- the closed set, so unmigrated variants would be rejected. Mirrors +-- normalize_entry_type() in mirix/schemas/procedural_memory.py. +UPDATE procedural_memory +SET entry_type = CASE + WHEN LOWER(TRIM(entry_type)) IN ('workflow', 'guide', 'script') THEN LOWER(TRIM(entry_type)) + WHEN entry_type ILIKE '%script%' THEN 'script' + WHEN entry_type ILIKE '%guide%' OR entry_type ILIKE '%how%' THEN 'guide' + ELSE 'workflow' +END +WHERE entry_type IS NULL OR entry_type NOT IN ('workflow', 'guide', 'script'); + +-- Step 4: Rename columns +ALTER TABLE procedural_memory RENAME COLUMN summary TO description; +ALTER TABLE procedural_memory RENAME COLUMN steps TO instructions; +ALTER TABLE procedural_memory RENAME COLUMN summary_embedding TO description_embedding; +ALTER TABLE procedural_memory RENAME COLUMN steps_embedding TO instructions_embedding; + +-- Step 5: Drop legacy non-unique lookup index (replaced by unique constraint below). +DROP INDEX IF EXISTS ix_procedural_memory_org_user_name; + +-- Step 6: Enforce per-user skill-name uniqueness at the DB level. Matches the +-- UniqueConstraint on the ORM model; without it, concurrent skill_create +-- calls race past the application-level pre-check. +ALTER TABLE procedural_memory + ADD CONSTRAINT uq_procedural_memory_org_user_name + UNIQUE (organization_id, user_id, name); + +COMMIT; diff --git a/tests/test_agent_trigger_state.py b/tests/test_agent_trigger_state.py new file mode 100644 index 000000000..348fc7614 --- /dev/null +++ b/tests/test_agent_trigger_state.py @@ -0,0 +1,429 @@ +"""Unit tests for the session-based procedural-memory trigger. + +Covers: +- Pydantic schema validation for AgentTriggerState and its trigger_type rules. +- ORM shape: table name, columns, unique constraint, index. +- Manager surface: exposes the expected async methods. +- memory_tools.trigger_memory_update wires to the new manager/threshold. +- Constants: SKILL_TRIGGER_SESSION_THRESHOLD is defined and env-overridable. +- Migration SQL creates the expected table. + +These are DB-free unit tests. DB-level UPSERT behavior is covered indirectly by +reading the manager source; an integration test that requires a running Postgres +is intentionally omitted here to keep this file in the fast lane. +""" +from __future__ import annotations + +import inspect +import re +from pathlib import Path + +import pytest +from sqlalchemy import inspect as sa_inspect + +from mirix import constants +from mirix.orm.agent_trigger_state import AgentTriggerState as AgentTriggerStateORM +from mirix.schemas.agent_trigger_state import ( + KNOWN_TRIGGER_TYPES, + TRIGGER_TYPE_MAX_LEN, + TRIGGER_TYPE_PATTERN, + TRIGGER_TYPE_PROCEDURAL_SKILL, + AgentTriggerState as PydanticAgentTriggerState, + _validate_trigger_type, +) +from mirix.services.agent_trigger_state_manager import ( + AgentTriggerStateManager, + ClaimFireResult, +) + + +# ----------------------------- Schema validation ------------------------- + + +class TestTriggerTypeValidation: + def test_procedural_skill_is_registered(self): + assert TRIGGER_TYPE_PROCEDURAL_SKILL in KNOWN_TRIGGER_TYPES + + def test_accepts_known_trigger_type(self): + assert _validate_trigger_type(TRIGGER_TYPE_PROCEDURAL_SKILL) == TRIGGER_TYPE_PROCEDURAL_SKILL + + def test_rejects_empty_string(self): + with pytest.raises(ValueError): + _validate_trigger_type("") + + def test_rejects_too_long(self): + with pytest.raises(ValueError): + _validate_trigger_type("a" * (TRIGGER_TYPE_MAX_LEN + 1)) + + def test_rejects_uppercase(self): + # Pattern is lowercase + digits + underscores; reject uppercase to + # keep the namespace tidy and avoid dupes by case. + with pytest.raises(ValueError): + _validate_trigger_type("Procedural_Skill") + + def test_rejects_hyphens(self): + with pytest.raises(ValueError): + _validate_trigger_type("procedural-skill") + + def test_pattern_matches_spec(self): + # Sanity-check the regex itself rejects a leading digit. + assert re.fullmatch(TRIGGER_TYPE_PATTERN, "1bad") is None + assert re.fullmatch(TRIGGER_TYPE_PATTERN, "a") is not None + + def test_rejects_unregistered_but_syntactically_valid(self): + # KNOWN_TRIGGER_TYPES is the source of truth; a syntactically valid + # but unregistered name must be rejected so typos cannot create a + # parallel bookkeeping row that no one watches. + assert "future_trigger" not in KNOWN_TRIGGER_TYPES + assert re.fullmatch(TRIGGER_TYPE_PATTERN, "future_trigger") is not None + with pytest.raises(ValueError): + _validate_trigger_type("future_trigger") + + +class TestPydanticAgentTriggerState: + def test_accepts_minimal_valid(self): + s = PydanticAgentTriggerState( + agent_id="agent-1", + user_id="user-1", + trigger_type=TRIGGER_TYPE_PROCEDURAL_SKILL, + ) + assert s.agent_id == "agent-1" + assert s.last_fired_at is None + assert s.last_fired_session_id is None + + def test_optional_cursor_fields(self): + from datetime import datetime, timezone + + ts = datetime(2026, 1, 1, tzinfo=timezone.utc) + s = PydanticAgentTriggerState( + agent_id="agent-1", + user_id="user-1", + trigger_type=TRIGGER_TYPE_PROCEDURAL_SKILL, + last_fired_at=ts, + last_fired_session_id="sess-abc", + ) + assert s.last_fired_at == ts + assert s.last_fired_session_id == "sess-abc" + + def test_rejects_invalid_trigger_type(self): + with pytest.raises(ValueError): + PydanticAgentTriggerState( + agent_id="agent-1", + user_id="user-1", + trigger_type="BAD", + ) + + +# ----------------------------- ORM shape --------------------------------- + + +class TestAgentTriggerStateORM: + def _column_names(self): + mapper = sa_inspect(AgentTriggerStateORM) + return {col.key for col in mapper.columns} + + def test_table_name(self): + assert AgentTriggerStateORM.__tablename__ == "agent_trigger_state" + + def test_has_required_columns(self): + cols = self._column_names() + expected = { + "id", + "organization_id", + "user_id", + "agent_id", + "trigger_type", + "last_fired_at", + "last_fired_session_id", + "last_fired_tied_session_ids", + "created_at", + "updated_at", + "is_deleted", + } + assert expected.issubset(cols), f"missing columns: {expected - cols}" + + def test_trigger_type_column_is_bounded(self): + col = AgentTriggerStateORM.__table__.c.trigger_type + assert col.type.length == TRIGGER_TYPE_MAX_LEN + + def test_unique_constraint_covers_triple(self): + uqs = [ + c for c in AgentTriggerStateORM.__table__.constraints + if c.__class__.__name__ == "UniqueConstraint" + ] + names = {u.name for u in uqs} + assert "uq_agent_trigger_state_agent_user_type" in names + target = next(u for u in uqs if u.name == "uq_agent_trigger_state_agent_user_type") + covered = {c.name for c in target.columns} + assert covered == {"agent_id", "user_id", "trigger_type"} + + def test_has_composite_index(self): + idx_names = {i.name for i in AgentTriggerStateORM.__table__.indexes} + assert "ix_agent_trigger_state_agent_user_type" in idx_names + + def test_is_registered_in_orm_package(self): + import mirix.orm as orm_pkg + + assert "AgentTriggerState" in orm_pkg.__all__ + assert orm_pkg.AgentTriggerState is AgentTriggerStateORM + + def test_pydantic_model_link(self): + # SqlalchemyBase.to_pydantic() dispatches via __pydantic_model__, so + # this link must point at the pydantic class to avoid runtime errors. + assert AgentTriggerStateORM.__pydantic_model__ is PydanticAgentTriggerState + + +# ----------------------------- Manager surface --------------------------- + + +class TestAgentTriggerStateManagerSurface: + """Ensure the manager exposes the contract memory_tools depends on. + + Wire-level correctness (UPSERT etc.) needs a real DB and is covered by + integration tests; here we just guard the public shape. + """ + + def test_exposes_async_methods(self): + # The manager methods are decorated by @enforce_types, whose wrapper + # is sync; unwrap to inspect the underlying async def. + mgr = AgentTriggerStateManager + for name in ( + "count_sealed_undistilled_sessions", + "check_and_claim_fire", + ): + assert hasattr(mgr, name), f"manager missing {name}" + method = getattr(mgr, name) + unwrapped = inspect.unwrap(method) + assert inspect.iscoroutinefunction(unwrapped), ( + f"{name} must be async under its decorators" + ) + + def test_check_and_claim_fire_signature(self): + sig = inspect.signature(AgentTriggerStateManager.check_and_claim_fire) + params = sig.parameters + for name in ( + "agent_id", + "user_id", + "trigger_type", + "threshold", + "organization_id", + "current_session_id", + ): + assert name in params, f"check_and_claim_fire missing kwarg {name}" + + def test_claim_result_shape(self): + # ClaimFireResult is the contract with memory_tools: if a field + # renames silently, the fire-path branches will read stale data. + fields = {f.name for f in __import__("dataclasses").fields(ClaimFireResult)} + assert fields == {"fired", "sessions_since", "just_installed", "state"} + + def test_count_signature_accepts_expected_kwargs(self): + # The cadence now counts SEALED, not-yet-distilled sessions in the + # Conversation Message Store, scoped per (user, org). The old + # cursor/watermark kwargs (since, tied_session_ids, exclude_session_id) + # are gone โ€” distilled_at is the durable dedup, not a timestamp window. + sig = inspect.signature( + AgentTriggerStateManager.count_sealed_undistilled_sessions + ) + params = sig.parameters + for name in ("user_id", "organization_id"): + assert name in params, ( + f"count_sealed_undistilled_sessions missing kwarg {name}" + ) + for gone in ("since", "tied_session_ids", "exclude_session_id", "agent_id"): + assert gone not in params, ( + f"obsolete windowing kwarg {gone} should be gone" + ) + + def test_uses_select_for_update(self): + # The whole point of check_and_claim_fire is serializing two workers + # on the same cursor row โ€” regress-guard against the lock hint being + # dropped during refactors. + src = inspect.getsource(AgentTriggerStateManager.check_and_claim_fire) + assert ".with_for_update(" in src, ( + "check_and_claim_fire must lock the cursor row with SELECT FOR UPDATE" + ) + + def test_excludes_soft_deleted_turns(self): + # Conversation turns carry an is_deleted flag from the common mixins. + # Leaving them in the count lets a user's wiped conversation keep + # firing procedural extraction โ€” mirror the manager's filter. + agg_src = inspect.getsource( + AgentTriggerStateManager._aggregate_sealed_undistilled + ) + assert "ConversationMessageModel.is_deleted.is_(False)" in agg_src, ( + "aggregate must exclude soft-deleted conversation turns" + ) + + def test_seals_by_min_created_at_and_dedups_by_max_distilled_at(self): + # Sealing/order uses MIN(created_at) per session (first-appearance order, + # immutable once the first turn lands). Dedup uses MAX(distilled_at) IS + # NULL โ€” a session counts iff none of its turns has been distilled. + agg_src = inspect.getsource( + AgentTriggerStateManager._aggregate_sealed_undistilled + ) + assert "func.min(ConversationMessageModel.created_at)" in agg_src + assert "func.max(ConversationMessageModel.distilled_at)" in agg_src + # created_at must NOT be aggregated with MAX (that would misorder sealing). + assert "func.max(ConversationMessageModel.created_at)" not in agg_src + # The open head (newest MIN) is dropped structurally so an in-progress + # session is never counted โ€” sealing is not a HAVING timestamp window. + assert "grouped[:-1]" in agg_src + + def test_sealed_set_from_single_group_by_query(self): + # The sealed-undistilled set comes from ONE GROUP BY session_id aggregate + # (order + distilled marker together), never separate queries. + agg_src = inspect.getsource( + AgentTriggerStateManager._aggregate_sealed_undistilled + ) + assert "group_by(ConversationMessageModel.session_id)" in agg_src, ( + "sealed set must be derived from one GROUP BY query" + ) + # No leftover machinery from the old message-table windowing scheme. + mgr_src = inspect.getsource(AgentTriggerStateManager) + assert "_tied_session_ids_at" not in mgr_src + assert "_aggregate_window" not in mgr_src + + def test_dedup_is_distilled_at_not_timestamp_tiebreaker(self): + # The old MIN>cursor + tied-session tie-breaker windowing is gone. + # Dedup is now durable: a session drops out of the sealed-undistilled + # set once its turns are stamped distilled_at, so there is no timestamp + # watermark to tie-break and no risk of double-firing a session. + agg_src = inspect.getsource( + AgentTriggerStateManager._aggregate_sealed_undistilled + ) + # A session is undistilled iff MAX(distilled_at) IS NULL. + assert "grp.last_distilled is None" in agg_src + # No resurrected watermark / tie-breaker machinery. + for gone in ( + "per_session_min > since", + "session_id.notin_(tied_session_ids)", + "per_session_min == since", + ".having(", + ): + assert gone not in agg_src, f"obsolete windowing fragment {gone!r} present" + + def test_fire_records_batch_for_diagnostics_only(self): + # On fire, the cursor records the oldest `threshold` sealed sessions it + # claimed โ€” for diagnostics ONLY. It must NOT be subtracted from future + # counts (distilled_at is the real dedup; subtracting would strand a + # session that failed to distill from its own retry). + src = inspect.getsource(AgentTriggerStateManager.check_and_claim_fire) + assert "last_fired_tied_session_ids = fired_batch" in src + assert "fired_batch = all_sealed[:threshold]" in src + # The fire counts the full sealed set, not a tied-subtracted remainder. + assert "count = len(all_sealed)" in src + + +# ----------------------------- memory_tools wiring ----------------------- + + +class TestMemoryToolsWiring: + """Guard the integration point inside trigger_memory_update. + + We inspect source text so the test stays DB-free. If the wiring ever + shifts to a different symbol, these assertions will flag it. + """ + + @staticmethod + def _source() -> str: + from mirix.functions.function_sets import memory_tools + + return inspect.getsource(memory_tools.trigger_memory_update) + + def test_uses_session_threshold(self): + src = self._source() + assert "SKILL_TRIGGER_SESSION_THRESHOLD" in src + # The old per-chunk message-count heuristic must be gone. + assert "SKILL_TRIGGER_MESSAGE_THRESHOLD" not in src + assert "[USER]" not in src and "[ASSISTANT]" not in src + + def test_uses_new_manager_and_trigger_type(self): + src = self._source() + assert "AgentTriggerStateManager" in src + assert "TRIGGER_TYPE_PROCEDURAL_SKILL" in src + + def test_uses_atomic_check_and_claim_fire(self): + # The fire path must go through the single atomic entry point โ€” not + # a separate read-then-write, which reintroduces the double-fire race. + src = self._source() + assert "check_and_claim_fire" in src + assert "record_fire" not in src + assert "count_distinct_sessions_since" not in src + + def test_reads_current_session_id_from_agent(self): + # The fire event must record the session in progress so the next + # count excludes it (no double counting of the open session). + src = self._source() + assert "_current_step_session_id" in src + + def test_skips_when_no_user(self): + # Counter is keyed per-user; without a user we must not bookkeep a + # fire event, or users would share a single cursor unexpectedly. + src = self._source() + assert "Skipping session-based procedural trigger" in src + + +# ----------------------------- Constants --------------------------------- + + +class TestThresholdConstant: + def test_default_is_five(self, monkeypatch): + # The env var may be set on the developer's machine; reimport under + # a pristine environment to check the default. + monkeypatch.delenv("SKILL_TRIGGER_SESSION_THRESHOLD", raising=False) + import importlib + + mod = importlib.reload(constants) + assert mod.SKILL_TRIGGER_SESSION_THRESHOLD == 5 + + def test_env_override(self, monkeypatch): + monkeypatch.setenv("SKILL_TRIGGER_SESSION_THRESHOLD", "12") + import importlib + + mod = importlib.reload(constants) + try: + assert mod.SKILL_TRIGGER_SESSION_THRESHOLD == 12 + finally: + monkeypatch.delenv("SKILL_TRIGGER_SESSION_THRESHOLD", raising=False) + importlib.reload(constants) + + +# ----------------------------- Migration SQL ----------------------------- + + +class TestMigrationSql: + SQL_PATH = Path("scripts/migrate_add_agent_trigger_state.sql") + + def _sql(self) -> str: + return self.SQL_PATH.read_text() + + def test_creates_table(self): + sql = self._sql() + assert "CREATE TABLE IF NOT EXISTS agent_trigger_state" in sql + + def test_declares_unique_and_index(self): + sql = self._sql() + assert "uq_agent_trigger_state_agent_user_type" in sql + assert "UNIQUE (agent_id, user_id, trigger_type)" in sql + assert "ix_agent_trigger_state_agent_user_type" in sql + + def test_has_tied_session_ids_column(self): + sql = self._sql() + # Present both in the CREATE TABLE ... and the back-compat ALTER + # so deployments that created the table before this fix get it too. + assert "last_fired_tied_session_ids" in sql + assert "ADD COLUMN IF NOT EXISTS last_fired_tied_session_ids" in sql + + def test_trigger_type_column_length_matches_schema(self): + sql = self._sql() + # Tolerate varying column alignment; only check that the SQL + # declares trigger_type with the shared max length. + assert re.search( + rf"\btrigger_type\s+VARCHAR\({TRIGGER_TYPE_MAX_LEN}\)", sql + ), "SQL column length must mirror TRIGGER_TYPE_MAX_LEN" + + def test_agent_fk_cascades(self): + # Deleting an agent must not leave orphan trigger-state rows behind. + sql = self._sql() + assert "REFERENCES agents(id) ON DELETE CASCADE" in sql diff --git a/tests/test_agent_trigger_state_integration.py b/tests/test_agent_trigger_state_integration.py new file mode 100644 index 000000000..8854173ff --- /dev/null +++ b/tests/test_agent_trigger_state_integration.py @@ -0,0 +1,744 @@ +"""DB-backed integration tests for the session-based procedural trigger. + +These tests require a real Postgres. `SELECT ... FOR UPDATE` is a no-op on +SQLite, so running these there would produce false-positive green on the +concurrency tests. The whole file is marked `integration` and skipped +unless MIRIX_PG_URI (or pg_user/pg_password/...) resolves to Postgres. + +Run with: + docker-compose up -d postgres + pytest tests/test_agent_trigger_state_integration.py -v -m integration + +The tests insert rows directly via SQLAlchemy ORM so `created_at` can be +controlled to the microsecond. This is essential for exercising the +tie-at-watermark path and the MIN-based "first-appearance" semantics. +""" +from __future__ import annotations + +import asyncio +import datetime as dt +import uuid +from datetime import datetime, timedelta, timezone +from typing import List, Optional + +import pytest +import pytest_asyncio + +from mirix.schemas.agent_trigger_state import TRIGGER_TYPE_PROCEDURAL_SKILL +from mirix.services.agent_trigger_state_manager import AgentTriggerStateManager +from mirix.settings import settings + + +pytestmark = [ + pytest.mark.integration, + pytest.mark.asyncio(loop_scope="module"), + pytest.mark.skipif( + not settings.mirix_pg_uri_no_default, + reason=( + "needs Postgres; SELECT FOR UPDATE is a no-op on SQLite so the " + "concurrency test would false-positive" + ), + ), +] + + +# ----------------------------- infrastructure ---------------------------- + + +@pytest_asyncio.fixture(scope="module", loop_scope="module") +async def server(): + """One AsyncServer instance per test module; owns the DB connection pool. + + `AsyncServer()` on its own does not run DDL โ€” that normally happens via + FastAPI's lifespan hook. We call `ensure_tables_created()` explicitly + so the ORM tables (including the new agent_trigger_state) exist before + any test runs. + """ + from mirix.server.server import AsyncServer, ensure_tables_created + + await ensure_tables_created() + return AsyncServer() + + +@pytest_asyncio.fixture(scope="module", loop_scope="module") +async def org(server): + """Dedicated org so we don't collide with other test suites.""" + from mirix.schemas.organization import Organization + + org_id = f"ats-org-{uuid.uuid4().hex[:8]}" + return await server.organization_manager.create_organization( + pydantic_org=Organization(id=org_id, name=org_id) + ) + + +@pytest_asyncio.fixture(scope="module", loop_scope="module") +async def user_a(server, org): + """Primary test user.""" + from mirix.schemas.user import User + + uid = f"ats-user-a-{uuid.uuid4().hex[:8]}" + return await server.user_manager.create_user( + pydantic_user=User(id=uid, name=uid, organization_id=org.id, timezone="UTC") + ) + + +@pytest_asyncio.fixture(scope="module", loop_scope="module") +async def user_b(server, org): + """Second user for cross-user isolation tests.""" + from mirix.schemas.user import User + + uid = f"ats-user-b-{uuid.uuid4().hex[:8]}" + return await server.user_manager.create_user( + pydantic_user=User(id=uid, name=uid, organization_id=org.id, timezone="UTC") + ) + + +async def _insert_minimal_agent(server, org_id: str) -> str: + """Insert a bare-bones agent row. All descriptive fields on the Agent + ORM are nullable, so a fresh UUID and the org_id satisfy the schema. + We bypass AgentManager to keep this test suite free of LLMConfig / + embedding-model requirements that don't matter for the trigger path. + """ + from mirix.orm.agent import Agent as AgentORM + from mirix.server.server import db_context + + agent_id = f"agent-{uuid.uuid4()}" + async with db_context() as session: + session.add(AgentORM(id=agent_id, organization_id=org_id)) + await session.commit() + return agent_id + + +@pytest_asyncio.fixture(loop_scope="module") +async def agent_id(server, org): + """Each test gets a fresh agent so the (agent_id, user_id, trigger_type) + unique key isolates trigger-state rows across tests. + """ + return await _insert_minimal_agent(server, org.id) + + +@pytest_asyncio.fixture(loop_scope="module") +async def other_agent_id(server, org): + """Second agent for cross-agent isolation tests.""" + return await _insert_minimal_agent(server, org.id) + + +async def _insert_message( + server, + *, + agent_id: str, + user_id: str, + org_id: str, + session_id: str, + created_at: datetime, + is_deleted: bool = False, +) -> str: + """Insert a message row with a caller-specified `created_at`. + + Going through MessageManager.create_message would stamp created_at via + the server default. For these tests we need microsecond-level control + so the tie-at-watermark and MIN-semantics cases can be set up + deterministically. + """ + from mirix.orm.message import Message as MessageORM + + from mirix.server.server import db_context + + msg_id = f"message-{uuid.uuid4()}" + async with db_context() as session: + msg = MessageORM( + id=msg_id, + agent_id=agent_id, + user_id=user_id, + organization_id=org_id, + role="user", + text=f"stub-{msg_id}", + tool_calls=[], + tool_returns=[], + session_id=session_id, + created_at=created_at, + is_deleted=is_deleted, + ) + session.add(msg) + await session.commit() + return msg_id + + +# ----------------------------- core behaviors ---------------------------- + + +class TestFirstInstall: + async def test_install_when_no_cursor(self, server, agent_id, user_a): + """First ever call must install the cursor at `now` and NOT fire, + even if there are pre-existing messages. Otherwise enabling the + feature would sweep up all legacy sessions on the first tick. + """ + mgr = AgentTriggerStateManager() + + # Seed some legacy messages BEFORE first install โ€” these must be + # ignored on install. + legacy_ts = datetime.now(timezone.utc) - timedelta(hours=1) + for i in range(5): + await _insert_message( + server, + agent_id=agent_id, + user_id=user_a.id, + org_id=user_a.organization_id, + session_id=f"legacy-{i}", + created_at=legacy_ts + timedelta(seconds=i), + ) + + claim = await mgr.check_and_claim_fire( + agent_id=agent_id, + user_id=user_a.id, + trigger_type=TRIGGER_TYPE_PROCEDURAL_SKILL, + threshold=5, + organization_id=user_a.organization_id, + current_session_id="bootstrap", + ) + + assert claim.just_installed is True + assert claim.fired is False + assert claim.sessions_since == 0 + assert claim.state.last_fired_at is not None + + +class TestBasicFire: + async def test_fires_when_threshold_met(self, server, agent_id, user_a): + mgr = AgentTriggerStateManager() + # Install cursor first. + await mgr.check_and_claim_fire( + agent_id=agent_id, + user_id=user_a.id, + trigger_type=TRIGGER_TYPE_PROCEDURAL_SKILL, + threshold=5, + organization_id=user_a.organization_id, + current_session_id="seed", + ) + # Insert 5 distinct NEW sessions strictly after the install cursor. + base = datetime.now(timezone.utc) + timedelta(seconds=1) + for i in range(5): + await _insert_message( + server, + agent_id=agent_id, + user_id=user_a.id, + org_id=user_a.organization_id, + session_id=f"sess-{i}", + created_at=base + timedelta(seconds=i), + ) + + claim = await mgr.check_and_claim_fire( + agent_id=agent_id, + user_id=user_a.id, + trigger_type=TRIGGER_TYPE_PROCEDURAL_SKILL, + threshold=5, + organization_id=user_a.organization_id, + current_session_id="sess-4", + ) + assert claim.fired is True + assert claim.sessions_since == 5 + # Cursor advances to the observed watermark, not wall clock. + assert claim.state.last_fired_at == base + timedelta(seconds=4) + + async def test_below_threshold_does_not_fire(self, server, agent_id, user_a): + mgr = AgentTriggerStateManager() + await mgr.check_and_claim_fire( + agent_id=agent_id, + user_id=user_a.id, + trigger_type=TRIGGER_TYPE_PROCEDURAL_SKILL, + threshold=5, + organization_id=user_a.organization_id, + ) + base = datetime.now(timezone.utc) + timedelta(seconds=1) + for i in range(4): # one short of threshold + await _insert_message( + server, + agent_id=agent_id, + user_id=user_a.id, + org_id=user_a.organization_id, + session_id=f"sess-{i}", + created_at=base + timedelta(seconds=i), + ) + + claim = await mgr.check_and_claim_fire( + agent_id=agent_id, + user_id=user_a.id, + trigger_type=TRIGGER_TYPE_PROCEDURAL_SKILL, + threshold=5, + organization_id=user_a.organization_id, + ) + assert claim.fired is False + assert claim.sessions_since == 4 + + +class TestNoDoubleCount: + """The whole point of switching to MIN semantics. Without it, a session + that was counted in window 1 and later sends more messages would get + re-counted in window 2 โ€” permanent over-fire. + """ + + async def test_old_session_continuing_does_not_recount( + self, server, agent_id, user_a + ): + mgr = AgentTriggerStateManager() + await mgr.check_and_claim_fire( + agent_id=agent_id, + user_id=user_a.id, + trigger_type=TRIGGER_TYPE_PROCEDURAL_SKILL, + threshold=5, + organization_id=user_a.organization_id, + ) + + base = datetime.now(timezone.utc) + timedelta(seconds=1) + for i in range(5): + await _insert_message( + server, + agent_id=agent_id, + user_id=user_a.id, + org_id=user_a.organization_id, + session_id=f"sess-{i}", + created_at=base + timedelta(seconds=i), + ) + + # First fire. + claim1 = await mgr.check_and_claim_fire( + agent_id=agent_id, + user_id=user_a.id, + trigger_type=TRIGGER_TYPE_PROCEDURAL_SKILL, + threshold=5, + organization_id=user_a.organization_id, + ) + assert claim1.fired is True + + # Now add more messages to the SAME sessions at a later timestamp. + # Under MAX semantics these would re-qualify; under MIN they don't. + later = base + timedelta(minutes=1) + for i in range(5): + await _insert_message( + server, + agent_id=agent_id, + user_id=user_a.id, + org_id=user_a.organization_id, + session_id=f"sess-{i}", # same session_ids + created_at=later + timedelta(seconds=i), + ) + + claim2 = await mgr.check_and_claim_fire( + agent_id=agent_id, + user_id=user_a.id, + trigger_type=TRIGGER_TYPE_PROCEDURAL_SKILL, + threshold=5, + organization_id=user_a.organization_id, + ) + assert claim2.fired is False, ( + "continuing old sessions must NOT re-fire โ€” each session " + "contributes to exactly one window under MIN semantics" + ) + assert claim2.sessions_since == 0 + + async def test_new_sessions_after_fire_do_refire( + self, server, agent_id, user_a + ): + """Sanity: genuinely new sessions after a fire DO re-qualify.""" + mgr = AgentTriggerStateManager() + await mgr.check_and_claim_fire( + agent_id=agent_id, + user_id=user_a.id, + trigger_type=TRIGGER_TYPE_PROCEDURAL_SKILL, + threshold=5, + organization_id=user_a.organization_id, + ) + base = datetime.now(timezone.utc) + timedelta(seconds=1) + for i in range(5): + await _insert_message( + server, + agent_id=agent_id, + user_id=user_a.id, + org_id=user_a.organization_id, + session_id=f"round1-{i}", + created_at=base + timedelta(seconds=i), + ) + claim1 = await mgr.check_and_claim_fire( + agent_id=agent_id, + user_id=user_a.id, + trigger_type=TRIGGER_TYPE_PROCEDURAL_SKILL, + threshold=5, + organization_id=user_a.organization_id, + ) + assert claim1.fired is True + + # Five brand-new session_ids at a later time. + later = base + timedelta(minutes=1) + for i in range(5): + await _insert_message( + server, + agent_id=agent_id, + user_id=user_a.id, + org_id=user_a.organization_id, + session_id=f"round2-{i}", + created_at=later + timedelta(seconds=i), + ) + claim2 = await mgr.check_and_claim_fire( + agent_id=agent_id, + user_id=user_a.id, + trigger_type=TRIGGER_TYPE_PROCEDURAL_SKILL, + threshold=5, + organization_id=user_a.organization_id, + ) + assert claim2.fired is True + assert claim2.sessions_since == 5 + + +class TestSoftDeleteExclusion: + async def test_deleted_sessions_do_not_count(self, server, agent_id, user_a): + """A user wiping history should not leave the procedural trigger + armed by the deleted sessions.""" + mgr = AgentTriggerStateManager() + await mgr.check_and_claim_fire( + agent_id=agent_id, + user_id=user_a.id, + trigger_type=TRIGGER_TYPE_PROCEDURAL_SKILL, + threshold=5, + organization_id=user_a.organization_id, + ) + base = datetime.now(timezone.utc) + timedelta(seconds=1) + # 3 live sessions... + for i in range(3): + await _insert_message( + server, + agent_id=agent_id, + user_id=user_a.id, + org_id=user_a.organization_id, + session_id=f"live-{i}", + created_at=base + timedelta(seconds=i), + ) + # ...plus 3 sessions whose only messages are soft-deleted. + for i in range(3): + await _insert_message( + server, + agent_id=agent_id, + user_id=user_a.id, + org_id=user_a.organization_id, + session_id=f"ghost-{i}", + created_at=base + timedelta(seconds=10 + i), + is_deleted=True, + ) + + claim = await mgr.check_and_claim_fire( + agent_id=agent_id, + user_id=user_a.id, + trigger_type=TRIGGER_TYPE_PROCEDURAL_SKILL, + threshold=5, + organization_id=user_a.organization_id, + ) + assert claim.fired is False + assert claim.sessions_since == 3, ( + "only the 3 live sessions should count; deleted sessions are " + "invisible to the trigger" + ) + + +class TestUserIsolation: + async def test_one_users_sessions_do_not_fire_another( + self, server, agent_id, user_a, user_b + ): + """Cursor is keyed per (agent, user, trigger_type). Five sessions + for user A must not trip user B's threshold.""" + mgr = AgentTriggerStateManager() + # Install cursors for both users. + await mgr.check_and_claim_fire( + agent_id=agent_id, + user_id=user_a.id, + trigger_type=TRIGGER_TYPE_PROCEDURAL_SKILL, + threshold=5, + organization_id=user_a.organization_id, + ) + await mgr.check_and_claim_fire( + agent_id=agent_id, + user_id=user_b.id, + trigger_type=TRIGGER_TYPE_PROCEDURAL_SKILL, + threshold=5, + organization_id=user_b.organization_id, + ) + + base = datetime.now(timezone.utc) + timedelta(seconds=1) + # Five new sessions for user A ONLY. + for i in range(5): + await _insert_message( + server, + agent_id=agent_id, + user_id=user_a.id, + org_id=user_a.organization_id, + session_id=f"a-sess-{i}", + created_at=base + timedelta(seconds=i), + ) + + claim_a = await mgr.check_and_claim_fire( + agent_id=agent_id, + user_id=user_a.id, + trigger_type=TRIGGER_TYPE_PROCEDURAL_SKILL, + threshold=5, + organization_id=user_a.organization_id, + ) + claim_b = await mgr.check_and_claim_fire( + agent_id=agent_id, + user_id=user_b.id, + trigger_type=TRIGGER_TYPE_PROCEDURAL_SKILL, + threshold=5, + organization_id=user_b.organization_id, + ) + assert claim_a.fired is True + assert claim_b.fired is False + assert claim_b.sessions_since == 0 + + +class TestTiedAtWatermark: + async def test_tied_sessions_recorded_and_not_double_counted( + self, server, agent_id, user_a + ): + """When multiple sessions share the exact watermark timestamp, + they should all be recorded in last_fired_tied_session_ids so the + next window does not re-count them via the `MIN == cursor` branch + of the HAVING filter. + """ + mgr = AgentTriggerStateManager() + await mgr.check_and_claim_fire( + agent_id=agent_id, + user_id=user_a.id, + trigger_type=TRIGGER_TYPE_PROCEDURAL_SKILL, + threshold=5, + organization_id=user_a.organization_id, + ) + + # 3 sessions with distinct timestamps, then 2 sharing the exact + # same watermark microsecond. The aggregate should see all 5, and + # the tied set should capture the 2 sharing the max. + base = datetime.now(timezone.utc) + timedelta(seconds=1) + for i in range(3): + await _insert_message( + server, + agent_id=agent_id, + user_id=user_a.id, + org_id=user_a.organization_id, + session_id=f"early-{i}", + created_at=base + timedelta(seconds=i), + ) + watermark = base + timedelta(seconds=10) + for sid in ("tied-a", "tied-b"): + await _insert_message( + server, + agent_id=agent_id, + user_id=user_a.id, + org_id=user_a.organization_id, + session_id=sid, + created_at=watermark, + ) + + claim = await mgr.check_and_claim_fire( + agent_id=agent_id, + user_id=user_a.id, + trigger_type=TRIGGER_TYPE_PROCEDURAL_SKILL, + threshold=5, + organization_id=user_a.organization_id, + ) + assert claim.fired is True + assert claim.sessions_since == 5 + assert claim.state.last_fired_at == watermark + assert set(claim.state.last_fired_tied_session_ids or []) == {"tied-a", "tied-b"} + + # Running again with NO new messages must not fire and must not + # re-count either tied-a or tied-b. + claim2 = await mgr.check_and_claim_fire( + agent_id=agent_id, + user_id=user_a.id, + trigger_type=TRIGGER_TYPE_PROCEDURAL_SKILL, + threshold=5, + organization_id=user_a.organization_id, + ) + assert claim2.fired is False + assert claim2.sessions_since == 0 + + async def test_late_arriving_session_at_watermark_picked_up_next_window( + self, server, agent_id, user_a + ): + """If a session's first message commits at the exact watermark but + was invisible to our SELECT, the tie-breaker must allow it into + the next window โ€” i.e. `MIN == cursor AND session_id NOT IN tied`. + We simulate the invisibility by inserting AFTER the first fire. + """ + mgr = AgentTriggerStateManager() + await mgr.check_and_claim_fire( + agent_id=agent_id, + user_id=user_a.id, + trigger_type=TRIGGER_TYPE_PROCEDURAL_SKILL, + threshold=5, + organization_id=user_a.organization_id, + ) + + base = datetime.now(timezone.utc) + timedelta(seconds=1) + watermark = base + timedelta(seconds=10) + # 4 "seen" sessions at strictly increasing ts, plus 1 at watermark. + for i in range(4): + await _insert_message( + server, + agent_id=agent_id, + user_id=user_a.id, + org_id=user_a.organization_id, + session_id=f"seen-{i}", + created_at=base + timedelta(seconds=i), + ) + await _insert_message( + server, + agent_id=agent_id, + user_id=user_a.id, + org_id=user_a.organization_id, + session_id="seen-at-watermark", + created_at=watermark, + ) + + claim1 = await mgr.check_and_claim_fire( + agent_id=agent_id, + user_id=user_a.id, + trigger_type=TRIGGER_TYPE_PROCEDURAL_SKILL, + threshold=5, + organization_id=user_a.organization_id, + ) + assert claim1.fired is True + + # Now insert a session whose first message is AT the watermark but + # wasn't visible during claim1 โ€” a delayed commit. Under a naive + # `MIN > cursor` filter this would be lost forever. The tie-breaker + # (`MIN = cursor AND session_id NOT IN tied_ids`) must rescue it. + await _insert_message( + server, + agent_id=agent_id, + user_id=user_a.id, + org_id=user_a.organization_id, + session_id="late-commit", + created_at=watermark, + ) + # Four more new sessions strictly after watermark โ€” together with + # the rescued late-commit we have 5 โ†’ fire. + for i in range(4): + await _insert_message( + server, + agent_id=agent_id, + user_id=user_a.id, + org_id=user_a.organization_id, + session_id=f"after-{i}", + created_at=watermark + timedelta(seconds=1 + i), + ) + + claim2 = await mgr.check_and_claim_fire( + agent_id=agent_id, + user_id=user_a.id, + trigger_type=TRIGGER_TYPE_PROCEDURAL_SKILL, + threshold=5, + organization_id=user_a.organization_id, + ) + assert claim2.fired is True, "late-commit session must be rescued" + assert claim2.sessions_since == 5 + # Sanity: the late-commit session is one of the counted, proving + # the tie-breaker really did include the `MIN == cursor` branch. + # (We can't assert the ids directly from the result, but the count + # of 5 only checks out if it was included.) + + +class TestConcurrency: + async def test_concurrent_fire_serialized_to_one_winner( + self, server, agent_id, user_a + ): + """Two concurrent check_and_claim_fire coroutines against the same + (agent, user, trigger_type) must serialize on SELECT FOR UPDATE. + Exactly one gets fired=True; the other sees the advanced cursor + and fired=False. + """ + mgr = AgentTriggerStateManager() + await mgr.check_and_claim_fire( + agent_id=agent_id, + user_id=user_a.id, + trigger_type=TRIGGER_TYPE_PROCEDURAL_SKILL, + threshold=5, + organization_id=user_a.organization_id, + ) + base = datetime.now(timezone.utc) + timedelta(seconds=1) + for i in range(5): + await _insert_message( + server, + agent_id=agent_id, + user_id=user_a.id, + org_id=user_a.organization_id, + session_id=f"conc-{i}", + created_at=base + timedelta(seconds=i), + ) + + async def claim(): + return await mgr.check_and_claim_fire( + agent_id=agent_id, + user_id=user_a.id, + trigger_type=TRIGGER_TYPE_PROCEDURAL_SKILL, + threshold=5, + organization_id=user_a.organization_id, + ) + + # Fire two coroutines at the same time. asyncio.gather schedules + # them concurrently; the real serialization happens at the + # Postgres row lock. + result_1, result_2 = await asyncio.gather(claim(), claim()) + fired_flags = sorted([result_1.fired, result_2.fired]) + assert fired_flags == [False, True], ( + "exactly one concurrent call must win the fire; got " + f"{fired_flags}" + ) + + +class TestAgentIsolation: + async def test_sessions_on_one_agent_do_not_count_for_another( + self, server, agent_id, other_agent_id, user_a + ): + """Cursor is per-agent. Sessions on agent X must not contribute to + the threshold on agent Y even for the same user. + """ + mgr = AgentTriggerStateManager() + await mgr.check_and_claim_fire( + agent_id=agent_id, + user_id=user_a.id, + trigger_type=TRIGGER_TYPE_PROCEDURAL_SKILL, + threshold=5, + organization_id=user_a.organization_id, + ) + await mgr.check_and_claim_fire( + agent_id=other_agent_id, + user_id=user_a.id, + trigger_type=TRIGGER_TYPE_PROCEDURAL_SKILL, + threshold=5, + organization_id=user_a.organization_id, + ) + base = datetime.now(timezone.utc) + timedelta(seconds=1) + for i in range(5): + await _insert_message( + server, + agent_id=agent_id, + user_id=user_a.id, + org_id=user_a.organization_id, + session_id=f"x-sess-{i}", + created_at=base + timedelta(seconds=i), + ) + + fire_x = await mgr.check_and_claim_fire( + agent_id=agent_id, + user_id=user_a.id, + trigger_type=TRIGGER_TYPE_PROCEDURAL_SKILL, + threshold=5, + organization_id=user_a.organization_id, + ) + fire_y = await mgr.check_and_claim_fire( + agent_id=other_agent_id, + user_id=user_a.id, + trigger_type=TRIGGER_TYPE_PROCEDURAL_SKILL, + threshold=5, + organization_id=user_a.organization_id, + ) + assert fire_x.fired is True + assert fire_y.fired is False + assert fire_y.sessions_since == 0 diff --git a/tests/test_auto_dream_response_changes.py b/tests/test_auto_dream_response_changes.py new file mode 100644 index 000000000..a9d39e238 --- /dev/null +++ b/tests/test_auto_dream_response_changes.py @@ -0,0 +1,184 @@ +"""Unit tests for the backward-compatible AutoDreamResponse schema extension. + +The generic-arm driver health-gates off structured evolution counts rather than +parsing the human-readable `message`, so AutoDreamResponse gained `skills_changed` +and `changes`. These tests pin the contract: + + * both new fields are OPTIONAL with sane defaults (0 / {}), so every existing + caller (and non-procedural modes) keeps working unchanged; + * when populated they round-trip through model_dump. + +No live server / API key needed. +""" + +import datetime as dt +from types import SimpleNamespace + +import pytest +from fastapi import HTTPException + +from mirix.schemas.agent import AgentType +from mirix.schemas.auto_dream import AutoDreamRequest +from mirix.schemas.auto_dream import AutoDreamResponse, MemoryTypeStats + + +def _now() -> dt.datetime: + return dt.datetime(2026, 6, 23, 12, 0, 0) + + +def test_defaults_are_zero_and_empty(): + resp = AutoDreamResponse( + start_date=None, + end_date=None, + processed={}, + last_dream_at=_now(), + dry_run=False, + ) + assert resp.skills_changed == 0 + assert resp.changes == {} + + +def test_populated_fields_round_trip(): + resp = AutoDreamResponse( + start_date=None, + end_date=None, + processed={"procedural": MemoryTypeStats(total=4)}, + last_dream_at=_now(), + dry_run=False, + skills_changed=2, + changes={"created": ["s1"], "edited": ["s2"], "deleted": []}, + ) + assert resp.skills_changed == 2 + assert resp.changes == {"created": ["s1"], "edited": ["s2"], "deleted": []} + + dumped = resp.model_dump() + assert dumped["skills_changed"] == 2 + assert dumped["changes"]["created"] == ["s1"] + assert dumped["processed"]["procedural"]["total"] == 4 + + # Re-parse the dump โ†’ identical structured counts (wire-stable for the driver). + reparsed = AutoDreamResponse.model_validate(dumped) + assert reparsed.skills_changed == 2 + assert reparsed.changes == resp.changes + + +def test_existing_caller_without_new_fields_unaffected(): + # Mirrors the exact shape pre-existing callers build (no skills_changed/changes). + resp = AutoDreamResponse( + start_date=None, + end_date=None, + processed={"procedural": MemoryTypeStats(total=0)}, + last_dream_at=_now(), + dry_run=True, + message="Dry run โ€” distilled 0 experience(s).", + ) + assert resp.message.startswith("Dry run") + assert resp.skills_changed == 0 + assert resp.changes == {} + + +def test_auto_dream_request_accepts_meta_agent_id(): + req = AutoDreamRequest(mode="procedural", meta_agent_id="agent-meta-1") + assert req.meta_agent_id == "agent-meta-1" + + +@pytest.mark.asyncio +async def test_auto_dream_handler_uses_explicit_meta_agent(monkeypatch): + from mirix.server import rest_api + + client = SimpleNamespace(id="client-1") + user = SimpleNamespace(id="user-1") + explicit_meta = SimpleNamespace( + id="agent-meta-2", + agent_type=AgentType.meta_memory_agent, + ) + captured = {} + + async def fake_auth(authorization=None, http_request=None): + return client, "api_key" + + class UserManager: + async def get_user_by_id(self, user_id=None): + assert user_id == "user-1" + return user + + class AgentManager: + async def get_agent_by_id(self, agent_id, actor): + assert agent_id == "agent-meta-2" + assert actor is client + return explicit_meta + + async def list_agents(self, actor): + raise AssertionError("explicit meta_agent_id should skip list_agents") + + class FakeAutoDreamManager: + async def run(self, *, request, user, actor, meta_agent_state): + captured["request"] = request + captured["user"] = user + captured["actor"] = actor + captured["meta_agent_state"] = meta_agent_state + return {"ok": True} + + fake_server = SimpleNamespace( + user_manager=UserManager(), + agent_manager=AgentManager(), + ) + + monkeypatch.setattr(rest_api, "get_client_from_jwt_or_api_key", fake_auth) + monkeypatch.setattr(rest_api, "get_server", lambda: fake_server) + + import mirix.services.auto_dream_manager as auto_dream_manager + + monkeypatch.setattr(auto_dream_manager, "AutoDreamManager", FakeAutoDreamManager) + + resp = await rest_api.auto_dream_handler( + AutoDreamRequest(mode="procedural", meta_agent_id="agent-meta-2"), + user_id="user-1", + authorization="Bearer test", + http_request=None, + ) + + assert resp == {"ok": True} + assert captured["meta_agent_state"] is explicit_meta + assert captured["user"] is user + assert captured["actor"] is client + assert captured["request"].meta_agent_id == "agent-meta-2" + + +@pytest.mark.asyncio +async def test_auto_dream_handler_rejects_non_meta_agent(monkeypatch): + from mirix.server import rest_api + + client = SimpleNamespace(id="client-1") + user = SimpleNamespace(id="user-1") + chat_agent = SimpleNamespace(id="agent-chat-1", agent_type=AgentType.chat_agent) + + async def fake_auth(authorization=None, http_request=None): + return client, "api_key" + + class UserManager: + async def get_user_by_id(self, user_id=None): + return user + + class AgentManager: + async def get_agent_by_id(self, agent_id, actor): + return chat_agent + + fake_server = SimpleNamespace( + user_manager=UserManager(), + agent_manager=AgentManager(), + ) + + monkeypatch.setattr(rest_api, "get_client_from_jwt_or_api_key", fake_auth) + monkeypatch.setattr(rest_api, "get_server", lambda: fake_server) + + with pytest.raises(HTTPException) as exc_info: + await rest_api.auto_dream_handler( + AutoDreamRequest(mode="procedural", meta_agent_id="agent-chat-1"), + user_id="user-1", + authorization="Bearer test", + http_request=None, + ) + + assert exc_info.value.status_code == 400 + assert "not a meta_memory_agent" in exc_info.value.detail diff --git a/tests/test_conversation_message_manager.py b/tests/test_conversation_message_manager.py new file mode 100644 index 000000000..b37bba073 --- /dev/null +++ b/tests/test_conversation_message_manager.py @@ -0,0 +1,1069 @@ +""" +Tests for the Conversation Message Store โ€” the single source the procedural +memory (skill) distiller reads. + +Layered exactly like the session_id tests: + + * DB-free shape/contract tests (mirroring tests/test_session_id.py): the ORM + column set, the Pydantic schema shape, role/session_id validation, and the + Postgres-only CHECK-constraint dialect gating. No DB, no running server. + * Postgres round-trip tests (marked @pytest.mark.integration, mirroring + tests/test_session_id_integration.py): record_turns ordering + accumulation, + count_distinct_sessions, the OLDEST-first sealed-only selection, + mark_sessions_distilled idempotency, and per-(user, org) scoping isolation. + +The DB-free run (`pytest -m "not integration"`) executes only the shape tests; +everything that needs Postgres carries the integration marker. +""" +from __future__ import annotations + +import pytest + + +# ============================ DB-free: ORM shape ============================ + + +class TestConversationMessageOrm: + """The table/column contract the migration + create_all must satisfy.""" + + def test_table_name(self): + from mirix.orm.conversation_message import ConversationMessage as CM + + assert CM.__tablename__ == "conversation_message" + + def test_key_columns_present(self): + from mirix.orm.conversation_message import ConversationMessage as CM + + cols = {c.name for c in CM.__table__.columns} + # The learnable-turn columns plus the owner scoping + barrier marker. + for expected in { + "id", + "session_id", + "role", + "content", + "distilled_at", + "user_id", + "organization_id", + "created_at", + }: + assert expected in cols, f"missing column {expected!r}" + + def test_session_id_is_not_null(self): + from mirix.orm.conversation_message import ConversationMessage as CM + + # This store only holds session'd turns โ€” session_id is NOT NULL. + assert CM.__table__.c.session_id.nullable is False + + def test_role_and_content_not_null(self): + from mirix.orm.conversation_message import ConversationMessage as CM + + assert CM.__table__.c.role.nullable is False + assert CM.__table__.c.content.nullable is False + + def test_distilled_at_is_nullable(self): + from mirix.orm.conversation_message import ConversationMessage as CM + + # NULL distilled_at == not yet consumed by a distill round. + assert CM.__table__.c.distilled_at.nullable is True + + def test_session_id_column_length_matches_constant(self): + from mirix.orm.conversation_message import ConversationMessage as CM + from mirix.schemas.message import SESSION_ID_MAX_LEN + + # One source of truth for the session_id length across schema + ORM. + assert CM.__table__.c.session_id.type.length == SESSION_ID_MAX_LEN + + def test_session_id_is_indexed(self): + from mirix.orm.conversation_message import ConversationMessage as CM + + covered = any( + any(c.name == "session_id" for c in idx.columns) + for idx in CM.__table__.indexes + ) + assert covered, "session_id should be indexed" + + def test_id_default_carries_convmsg_prefix(self): + from mirix.orm.conversation_message import ConversationMessage as CM + + default = CM.__table__.c.id.default + # The PK uses a `convmsg-` prefixed callable default (sexp-/proc- convention). + assert default is not None and default.is_callable + assert default.arg(None).startswith("convmsg-") + + +class TestConversationMessageCheckConstraintGating: + """The session_id CHECK uses Postgres' `~` regex operator, which SQLite + cannot parse. It must be emitted ONLY for Postgres so SQLite create_all works. + Mirrors tests/test_session_id.py::TestCheckConstraintDialectGating. + """ + + def test_check_compiles_for_postgresql(self): + from sqlalchemy.dialects import postgresql + from sqlalchemy.schema import CreateTable + + from mirix.orm.conversation_message import ConversationMessage as CM + + ddl = str(CreateTable(CM.__table__).compile(dialect=postgresql.dialect())) + assert "ck_conversation_message_session_id_format" in ddl + assert "~" in ddl # Postgres regex operator + + def test_check_is_suppressed_for_sqlite(self): + from sqlalchemy.dialects import sqlite + from sqlalchemy.schema import CreateTable + + from mirix.orm.conversation_message import ConversationMessage as CM + + ddl = str(CreateTable(CM.__table__).compile(dialect=sqlite.dialect())) + assert "ck_conversation_message_session_id_format" not in ddl + assert "session_id ~" not in ddl + + def test_check_constraint_uses_shared_pattern(self): + from mirix.orm.conversation_message import ConversationMessage as CM + from mirix.schemas.message import SESSION_ID_SQL_PATTERN + + ck_texts = [ + str(c.sqltext) + for c in CM.__table__.constraints + if getattr(c, "name", None) == "ck_conversation_message_session_id_format" + ] + assert ck_texts, "ck_conversation_message_session_id_format not found" + assert SESSION_ID_SQL_PATTERN in ck_texts[0] + + +# ============================ DB-free: Pydantic shape ====================== + + +class TestConversationMessageSchema: + def test_full_schema_field_shape(self): + from mirix.schemas.conversation_message import ConversationMessage + + fields = ConversationMessage.model_fields + for expected in { + "id", + "session_id", + "role", + "content", + "user_id", + "organization_id", + "distilled_at", + "created_at", + "updated_at", + }: + assert expected in fields, f"missing schema field {expected!r}" + + def test_create_carries_owner_ids(self): + from mirix.schemas.conversation_message import ConversationMessageCreate + + c = ConversationMessageCreate( + session_id="sess-1", + role="user", + content="hi", + user_id="user-1", + organization_id="org-1", + ) + assert c.role == "user" + assert c.user_id == "user-1" + assert c.organization_id == "org-1" + + def test_response_is_full_schema_alias(self): + from mirix.schemas.conversation_message import ( + ConversationMessage, + ConversationMessageResponse, + ) + + # Response must expose the same fields as the full schema. + assert ( + set(ConversationMessageResponse.model_fields) + == set(ConversationMessage.model_fields) + ) + + def test_role_accepts_user_assistant_and_tool(self): + from mirix.schemas.conversation_message import ConversationMessageCreate + + # 'tool' turns carry the work-process signal (tool errors/retries) the + # distiller's tool_error signal_type exists for. + for role in ("user", "assistant", "tool"): + c = ConversationMessageCreate( + session_id="sess-1", + role=role, + content="x", + user_id="u", + organization_id="o", + ) + assert c.role == role + + def test_role_rejects_other_values(self): + from mirix.schemas.conversation_message import ConversationMessageCreate + + # 'system' is meta-agent scaffolding, never a learnable turn; roles are + # also case-sensitive and non-empty. + for bad in ("system", "", "User", "TOOL"): + with pytest.raises(ValueError): + ConversationMessageCreate( + session_id="sess-1", + role=bad, + content="x", + user_id="u", + organization_id="o", + ) + + def test_session_id_required_rejects_none_and_empty(self): + from mirix.schemas.conversation_message import ConversationMessageCreate + + # This store never holds a session-less turn. + with pytest.raises(ValueError): + ConversationMessageCreate( + session_id="", + role="user", + content="x", + user_id="u", + organization_id="o", + ) + with pytest.raises(ValueError): + ConversationMessageCreate( + session_id=None, + role="user", + content="x", + user_id="u", + organization_id="o", + ) + + def test_session_id_charset_and_length_enforced(self): + from mirix.schemas.conversation_message import ConversationMessageCreate + from mirix.schemas.message import SESSION_ID_MAX_LEN + + with pytest.raises(ValueError): + ConversationMessageCreate( + session_id="bad/chars", + role="user", + content="x", + user_id="u", + organization_id="o", + ) + with pytest.raises(ValueError): + ConversationMessageCreate( + session_id="a" * (SESSION_ID_MAX_LEN + 1), + role="user", + content="x", + user_id="u", + organization_id="o", + ) + # Allowed charset passes. + ok = ConversationMessageCreate( + session_id="sess_Abc-123", + role="assistant", + content="x", + user_id="u", + organization_id="o", + ) + assert ok.session_id == "sess_Abc-123" + + def test_content_length_cap_enforced(self): + from mirix.schemas.conversation_message import ( + CONVERSATION_MESSAGE_MAX_CONTENT_LEN, + ConversationMessageCreate, + ) + + with pytest.raises(ValueError): + ConversationMessageCreate( + session_id="sess-1", + role="user", + content="c" * (CONVERSATION_MESSAGE_MAX_CONTENT_LEN + 1), + user_id="u", + organization_id="o", + ) + + def test_content_defaults_to_empty(self): + from mirix.schemas.conversation_message import ConversationMessageCreate + + c = ConversationMessageCreate( + session_id="sess-1", + role="user", + user_id="u", + organization_id="o", + ) + assert c.content == "" + + +class TestConversationMessageManagerContract: + """DB-free: the async method contract downstream agents depend on. + + No DB here โ€” we only assert the methods exist, are coroutines, and expose + the keyword-only signature the ingestion seam / trigger / distiller and + the data-lifecycle (retention + erasure) paths call. + """ + + def test_manager_exposes_contract_async_methods(self): + import inspect + + from mirix.services.conversation_message_manager import ( + ConversationMessageManager, + ) + + for name in ( + "record_turns", + "count_distinct_sessions", + "list_sealed_undistilled_sessions", + "list_turns_for_session", + "mark_sessions_distilled", + "prune_distilled_turns", + "delete_by_user_id", + "delete_by_client_id", + ): + fn = getattr(ConversationMessageManager, name, None) + assert fn is not None, f"manager missing {name}" + # Unwrap the @enforce_types decorator (a plain-def wrapper that + # returns the coroutine) to inspect the real `async def` underneath. + assert inspect.iscoroutinefunction( + inspect.unwrap(fn) + ), f"{name} must be async" + + def test_record_turns_signature_is_keyword_only(self): + import inspect + + from mirix.services.conversation_message_manager import ( + ConversationMessageManager, + ) + + sig = inspect.signature(ConversationMessageManager.record_turns) + params = sig.parameters + for expected in ( + "session_id", + "user_id", + "organization_id", + "turns", + "actor", + ): + assert expected in params, f"record_turns missing {expected}" + assert ( + params[expected].kind is inspect.Parameter.KEYWORD_ONLY + ), f"{expected} must be keyword-only" + + def test_count_distinct_sessions_has_only_undistilled_flag(self): + import inspect + + from mirix.services.conversation_message_manager import ( + ConversationMessageManager, + ) + + sig = inspect.signature( + ConversationMessageManager.count_distinct_sessions + ) + assert "only_undistilled" in sig.parameters + assert sig.parameters["only_undistilled"].default is False + + +# ============================ Integration: Postgres ======================== +# +# Everything below needs a live Postgres (docker-compose). It is marked with the +# integration marker so the default `-m "not integration"` CI subset skips it. + +import asyncio # noqa: E402 +import sys # noqa: E402 +import uuid # noqa: E402 +from pathlib import Path # noqa: E402 + +import pytest_asyncio # noqa: E402 + +pytestmark = [] # module-level markers applied per-class below instead + +project_root = Path(__file__).parent.parent +sys.path.insert(0, str(project_root)) + + +@pytest_asyncio.fixture(scope="module") +def event_loop(): + loop = asyncio.new_event_loop() + yield loop + loop.close() + + +async def _make_actor(): + """Create a fresh org + client (a distinct owner). Returns the client/actor. + + Each call mints its OWN organization so two actors are guaranteed to live in + different organizations โ€” exactly what the per-(user, org) scoping tests need. + """ + from mirix.schemas.client import Client as PydanticClient + from mirix.schemas.organization import Organization as PydanticOrganization + from mirix.services.client_manager import ClientManager + from mirix.services.organization_manager import OrganizationManager + + org_mgr = OrganizationManager() + client_mgr = ClientManager() + + org_id = f"test-convmsg-org-{uuid.uuid4().hex[:8]}" + try: + await org_mgr.get_organization_by_id(org_id) + except Exception: + await org_mgr.create_organization( + PydanticOrganization(id=org_id, name="ConvMsg Test Org") + ) + + client_id = f"test-convmsg-client-{uuid.uuid4().hex[:8]}" + try: + return await client_mgr.get_client_by_id(client_id) + except Exception: + return await client_mgr.create_client( + PydanticClient( + id=client_id, + organization_id=org_id, + name="ConvMsg Test Client", + write_scope="test-convmsg", + read_scopes=["test-convmsg"], + ) + ) + + +@pytest_asyncio.fixture(scope="module") +async def cm_actor(): + return await _make_actor() + + +@pytest_asyncio.fixture(scope="module") +async def cm_actor_org_b(): + """A second actor in a DIFFERENT organization, for cross-org isolation.""" + return await _make_actor() + + +async def _make_user(org_id: str): + from mirix.schemas.user import User as PydanticUser + from mirix.services.user_manager import UserManager + + user_mgr = UserManager() + user_id = f"test-convmsg-user-{uuid.uuid4().hex[:8]}" + try: + return await user_mgr.get_user_by_id(user_id) + except Exception: + return await user_mgr.create_user( + PydanticUser( + id=user_id, + name="ConvMsg Test User", + organization_id=org_id, + timezone="UTC", + ) + ) + + +@pytest_asyncio.fixture(scope="module") +async def cm_user(cm_actor): + return await _make_user(cm_actor.organization_id) + + +@pytest.fixture +def manager(): + from mirix.services.conversation_message_manager import ( + ConversationMessageManager, + ) + + return ConversationMessageManager() + + +@pytest.mark.integration +@pytest.mark.asyncio(loop_scope="module") +class TestRecordTurns: + async def test_records_turns_in_order_with_real_roles( + self, manager, cm_actor, cm_user + ): + sid = f"rt-{uuid.uuid4().hex[:8]}" + out = await manager.record_turns( + session_id=sid, + user_id=cm_user.id, + organization_id=cm_actor.organization_id, + turns=[ + {"role": "user", "content": "Q: 2+2?"}, + {"role": "assistant", "content": "4"}, + ], + actor=cm_actor, + ) + assert [t.role for t in out] == ["user", "assistant"] + assert [t.content for t in out] == ["Q: 2+2?", "4"] + # The STRICTLY increasing created_at is the documented contract guarantee + # (turns share one transaction, so func.now() would collapse them to one + # timestamp and destroy order). It is what later sealing/ordering relies + # on, so we pin it as behavior โ€” not merely the already-checked ascending + # retrieval below. + assert out[0].created_at < out[1].created_at + + got = await manager.list_turns_for_session( + session_id=sid, + user_id=cm_user.id, + organization_id=cm_actor.organization_id, + actor=cm_actor, + ) + assert [t.content for t in got] == ["Q: 2+2?", "4"] + + async def test_multiple_calls_same_session_accumulate_in_order( + self, manager, cm_actor, cm_user + ): + sid = f"acc-{uuid.uuid4().hex[:8]}" + await manager.record_turns( + session_id=sid, + user_id=cm_user.id, + organization_id=cm_actor.organization_id, + turns=[{"role": "user", "content": "first"}], + actor=cm_actor, + ) + await manager.record_turns( + session_id=sid, + user_id=cm_user.id, + organization_id=cm_actor.organization_id, + turns=[{"role": "assistant", "content": "second"}], + actor=cm_actor, + ) + got = await manager.list_turns_for_session( + session_id=sid, + user_id=cm_user.id, + organization_id=cm_actor.organization_id, + actor=cm_actor, + ) + # A later call's turns sort AFTER an earlier call's โ€” one ordered unit. + assert [t.content for t in got] == ["first", "second"] + + async def test_empty_turns_is_noop(self, manager, cm_actor, cm_user): + sid = f"empty-{uuid.uuid4().hex[:8]}" + out = await manager.record_turns( + session_id=sid, + user_id=cm_user.id, + organization_id=cm_actor.organization_id, + turns=[], + actor=cm_actor, + ) + assert out == [] + + async def test_bad_role_rejected_before_any_write( + self, manager, cm_actor, cm_user + ): + sid = f"badrole-{uuid.uuid4().hex[:8]}" + with pytest.raises(ValueError): + await manager.record_turns( + session_id=sid, + user_id=cm_user.id, + organization_id=cm_actor.organization_id, + turns=[ + {"role": "user", "content": "ok"}, + {"role": "system", "content": "bad"}, # invalid role + ], + actor=cm_actor, + ) + # The whole batch is validated up front: NOTHING was written. + got = await manager.list_turns_for_session( + session_id=sid, + user_id=cm_user.id, + organization_id=cm_actor.organization_id, + actor=cm_actor, + ) + assert got == [] + + +@pytest.mark.integration +@pytest.mark.asyncio(loop_scope="module") +class TestCountDistinctSessions: + async def test_counts_distinct_sessions_for_owner( + self, manager, cm_actor, cm_user + ): + user = await _make_user(cm_actor.organization_id) + for sid in ("c1", "c2", "c3"): + await manager.record_turns( + session_id=f"{sid}-{uuid.uuid4().hex[:6]}", + user_id=user.id, + organization_id=cm_actor.organization_id, + turns=[{"role": "user", "content": "x"}], + actor=cm_actor, + ) + n = await manager.count_distinct_sessions( + user_id=user.id, + organization_id=cm_actor.organization_id, + actor=cm_actor, + ) + assert n == 3 + + async def test_only_undistilled_excludes_distilled_sessions( + self, manager, cm_actor, cm_user + ): + user = await _make_user(cm_actor.organization_id) + sids = [] + for _ in range(3): + sid = f"u-{uuid.uuid4().hex[:8]}" + sids.append(sid) + await manager.record_turns( + session_id=sid, + user_id=user.id, + organization_id=cm_actor.organization_id, + turns=[{"role": "user", "content": "x"}], + actor=cm_actor, + ) + # Distill one session. + await manager.mark_sessions_distilled( + session_ids=[sids[0]], + user_id=user.id, + organization_id=cm_actor.organization_id, + actor=cm_actor, + ) + total = await manager.count_distinct_sessions( + user_id=user.id, + organization_id=cm_actor.organization_id, + actor=cm_actor, + ) + undistilled = await manager.count_distinct_sessions( + user_id=user.id, + organization_id=cm_actor.organization_id, + actor=cm_actor, + only_undistilled=True, + ) + assert total == 3 + assert undistilled == 2 + + +@pytest.mark.integration +@pytest.mark.asyncio(loop_scope="module") +class TestSealedUndistilledSelection: + async def test_returns_oldest_first_sealed_only( + self, manager, cm_actor, cm_user + ): + user = await _make_user(cm_actor.organization_id) + # Create 6 sessions, each a clearly-later MIN(created_at) than the last. + ordered = [] + for i in range(6): + sid = f"seal-{i}-{uuid.uuid4().hex[:6]}" + ordered.append(sid) + await manager.record_turns( + session_id=sid, + user_id=user.id, + organization_id=cm_actor.organization_id, + turns=[{"role": "user", "content": f"turn {i}"}], + actor=cm_actor, + ) + sealed = await manager.list_sealed_undistilled_sessions( + user_id=user.id, + organization_id=cm_actor.organization_id, + actor=cm_actor, + limit=5, + ) + # 6 sessions โ†’ 5 oldest sealed, oldest-FIRST; the newest is the open head. + assert sealed == ordered[:5] + assert ordered[5] not in sealed # open head never returned + + async def test_single_session_is_never_sealed( + self, manager, cm_actor, cm_user + ): + user = await _make_user(cm_actor.organization_id) + await manager.record_turns( + session_id=f"solo-{uuid.uuid4().hex[:8]}", + user_id=user.id, + organization_id=cm_actor.organization_id, + turns=[{"role": "user", "content": "alone"}], + actor=cm_actor, + ) + sealed = await manager.list_sealed_undistilled_sessions( + user_id=user.id, + organization_id=cm_actor.organization_id, + actor=cm_actor, + limit=5, + ) + assert sealed == [] + + async def test_distilled_sessions_drop_out_of_sealed_window( + self, manager, cm_actor, cm_user + ): + user = await _make_user(cm_actor.organization_id) + ordered = [] + for i in range(4): + sid = f"drop-{i}-{uuid.uuid4().hex[:6]}" + ordered.append(sid) + await manager.record_turns( + session_id=sid, + user_id=user.id, + organization_id=cm_actor.organization_id, + turns=[{"role": "user", "content": f"t{i}"}], + actor=cm_actor, + ) + # Distill the two oldest; they must drop out of the sealed window. + await manager.mark_sessions_distilled( + session_ids=ordered[:2], + user_id=user.id, + organization_id=cm_actor.organization_id, + actor=cm_actor, + ) + sealed = await manager.list_sealed_undistilled_sessions( + user_id=user.id, + organization_id=cm_actor.organization_id, + actor=cm_actor, + limit=5, + ) + # 4 sessions, newest is open head (ordered[3]); ordered[0..1] distilled โ†’ + # only ordered[2] remains sealed-and-undistilled. + assert sealed == [ordered[2]] + + +@pytest.mark.integration +@pytest.mark.asyncio(loop_scope="module") +class TestMarkSessionsDistilled: + async def test_idempotent_rerun_marks_zero(self, manager, cm_actor, cm_user): + user = await _make_user(cm_actor.organization_id) + sid = f"idem-{uuid.uuid4().hex[:8]}" + await manager.record_turns( + session_id=sid, + user_id=user.id, + organization_id=cm_actor.organization_id, + turns=[ + {"role": "user", "content": "a"}, + {"role": "assistant", "content": "b"}, + ], + actor=cm_actor, + ) + first = await manager.mark_sessions_distilled( + session_ids=[sid], + user_id=user.id, + organization_id=cm_actor.organization_id, + actor=cm_actor, + ) + second = await manager.mark_sessions_distilled( + session_ids=[sid], + user_id=user.id, + organization_id=cm_actor.organization_id, + actor=cm_actor, + ) + assert first == 2 # both turns stamped + assert second == 0 # re-run is a no-op + + async def test_empty_session_ids_is_noop(self, manager, cm_actor, cm_user): + n = await manager.mark_sessions_distilled( + session_ids=[], + user_id=cm_user.id, + organization_id=cm_actor.organization_id, + actor=cm_actor, + ) + assert n == 0 + + +@pytest.mark.integration +@pytest.mark.asyncio(loop_scope="module") +class TestPerUserOrgScoping: + async def test_sessions_isolated_per_user(self, manager, cm_actor): + user_a = await _make_user(cm_actor.organization_id) + user_b = await _make_user(cm_actor.organization_id) + sid_a = f"iso-a-{uuid.uuid4().hex[:6]}" + sid_b = f"iso-b-{uuid.uuid4().hex[:6]}" + + await manager.record_turns( + session_id=sid_a, + user_id=user_a.id, + organization_id=cm_actor.organization_id, + turns=[{"role": "user", "content": "from A"}], + actor=cm_actor, + ) + await manager.record_turns( + session_id=sid_b, + user_id=user_b.id, + organization_id=cm_actor.organization_id, + turns=[{"role": "user", "content": "from B"}], + actor=cm_actor, + ) + + # Each user counts only their own session. + assert ( + await manager.count_distinct_sessions( + user_id=user_a.id, + organization_id=cm_actor.organization_id, + actor=cm_actor, + ) + == 1 + ) + assert ( + await manager.count_distinct_sessions( + user_id=user_b.id, + organization_id=cm_actor.organization_id, + actor=cm_actor, + ) + == 1 + ) + # User A cannot see B's turns. + a_turns = await manager.list_turns_for_session( + session_id=sid_b, + user_id=user_a.id, + organization_id=cm_actor.organization_id, + actor=cm_actor, + ) + assert a_turns == [] + + async def test_mark_distilled_cannot_cross_user(self, manager, cm_actor): + user_a = await _make_user(cm_actor.organization_id) + user_b = await _make_user(cm_actor.organization_id) + sid = f"xuser-{uuid.uuid4().hex[:6]}" + + await manager.record_turns( + session_id=sid, + user_id=user_a.id, + organization_id=cm_actor.organization_id, + turns=[{"role": "user", "content": "A owns this"}], + actor=cm_actor, + ) + # User B attempting to mark A's session distills nothing. + crossed = await manager.mark_sessions_distilled( + session_ids=[sid], + user_id=user_b.id, + organization_id=cm_actor.organization_id, + actor=cm_actor, + ) + assert crossed == 0 + # A's session is still undistilled. + assert ( + await manager.count_distinct_sessions( + user_id=user_a.id, + organization_id=cm_actor.organization_id, + actor=cm_actor, + only_undistilled=True, + ) + == 1 + ) + + async def test_sessions_isolated_per_organization( + self, manager, cm_actor, cm_actor_org_b + ): + """A session recorded in org A must be invisible to org B's queries. + + The decisive move is the CROSS query: we ask for org A's user under org + B's id. A manager that filtered only by user_id (ignoring + organization_id) would still return org A's session here; correct + per-(user, org) scoping returns nothing. (A real `users` row is tied to + exactly one org by FK, so two orgs cannot share a user PK โ€” the cross + query, not a shared id, is what proves the org dimension is enforced.) + """ + org_a = cm_actor.organization_id + org_b = cm_actor_org_b.organization_id + assert org_a != org_b # the fixtures mint distinct orgs + + user_a = await _make_user(org_a) + sid_a = f"org-a-{uuid.uuid4().hex[:6]}" + await manager.record_turns( + session_id=sid_a, + user_id=user_a.id, + organization_id=org_a, + turns=[{"role": "user", "content": "belongs to org A"}], + actor=cm_actor, + ) + + # Correct, same-org query sees the sessionโ€ฆ + assert ( + await manager.count_distinct_sessions( + user_id=user_a.id, + organization_id=org_a, + actor=cm_actor, + ) + == 1 + ) + # โ€ฆbut the SAME user_id under org B's organization_id sees NOTHING. + assert ( + await manager.count_distinct_sessions( + user_id=user_a.id, + organization_id=org_b, + actor=cm_actor_org_b, + ) + == 0 + ) + # And org B cannot read org A's turns even with A's user_id. + assert ( + await manager.list_turns_for_session( + session_id=sid_a, + user_id=user_a.id, + organization_id=org_b, + actor=cm_actor_org_b, + ) + == [] + ) + + async def test_mark_distilled_cannot_cross_organization( + self, manager, cm_actor, cm_actor_org_b + ): + """mark_sessions_distilled is org-scoped: passing org B's + organization_id (with org A's user_id and session) distills nothing.""" + org_a = cm_actor.organization_id + org_b = cm_actor_org_b.organization_id + user_a = await _make_user(org_a) + sid = f"xorg-{uuid.uuid4().hex[:6]}" + await manager.record_turns( + session_id=sid, + user_id=user_a.id, + organization_id=org_a, + turns=[{"role": "user", "content": "A owns this"}], + actor=cm_actor, + ) + # Marking under org B's organization_id must be a no-op. + crossed = await manager.mark_sessions_distilled( + session_ids=[sid], + user_id=user_a.id, + organization_id=org_b, + actor=cm_actor_org_b, + ) + assert crossed == 0 + # A's session remains undistilled in its own org. + assert ( + await manager.count_distinct_sessions( + user_id=user_a.id, + organization_id=org_a, + actor=cm_actor, + only_undistilled=True, + ) + == 1 + ) + + +@pytest.mark.integration +@pytest.mark.asyncio(loop_scope="module") +class TestDataLifecycle: + """Retention pruning and the user/client erasure paths. + + The conversation store holds verbatim transcripts, so these are the + GDPR-facing guarantees: distilled turns age out after the retention + window, and the purge paths remove everything for a user/client. + """ + + async def _count_all_rows(self, manager, user_id: str, organization_id: str) -> int: + from sqlalchemy import func, select + + from mirix.orm.conversation_message import ( + ConversationMessage as ConversationMessageModel, + ) + + async with manager.session_maker() as session: + stmt = select(func.count()).where( + ConversationMessageModel.user_id == user_id, + ConversationMessageModel.organization_id == organization_id, + ) + result = await session.execute(stmt) + return int(result.scalar_one() or 0) + + async def _backdate_distilled(self, manager, session_id: str, days: int) -> None: + """Force distilled_at into the past so the retention cutoff applies.""" + from datetime import timedelta + + from sqlalchemy import update + + from mirix.client.utils import get_utc_time + from mirix.orm.conversation_message import ( + ConversationMessage as ConversationMessageModel, + ) + + async with manager.session_maker() as session: + await session.execute( + update(ConversationMessageModel) + .where(ConversationMessageModel.session_id == session_id) + .values(distilled_at=get_utc_time() - timedelta(days=days)) + ) + await session.commit() + + async def test_prune_zero_retention_is_noop(self, manager, cm_actor, cm_user): + pruned = await manager.prune_distilled_turns( + user_id=cm_user.id, + organization_id=cm_actor.organization_id, + actor=cm_actor, + retention_days=0, + ) + assert pruned == 0 + + async def test_prune_removes_only_distilled_turns_past_cutoff( + self, manager, cm_actor, cm_user + ): + org = cm_actor.organization_id + old_sid = f"prune-old-{uuid.uuid4().hex[:6]}" + fresh_sid = f"prune-fresh-{uuid.uuid4().hex[:6]}" + open_sid = f"prune-open-{uuid.uuid4().hex[:6]}" + for sid in (old_sid, fresh_sid, open_sid): + await manager.record_turns( + session_id=sid, + user_id=cm_user.id, + organization_id=org, + turns=[{"role": "user", "content": f"turn in {sid}"}], + actor=cm_actor, + ) + # old + fresh get distilled; old is then backdated past the window. + await manager.mark_sessions_distilled( + session_ids=[old_sid, fresh_sid], + user_id=cm_user.id, + organization_id=org, + actor=cm_actor, + ) + await self._backdate_distilled(manager, old_sid, days=60) + + pruned = await manager.prune_distilled_turns( + user_id=cm_user.id, + organization_id=org, + actor=cm_actor, + retention_days=30, + ) + + assert pruned == 1 # only the backdated session's turn + remaining = [ + turn.session_id + for sid in (old_sid, fresh_sid, open_sid) + for turn in await manager.list_turns_for_session( + session_id=sid, + user_id=cm_user.id, + organization_id=org, + actor=cm_actor, + ) + ] + assert old_sid not in remaining + # Freshly distilled and never-distilled turns both survive. + assert fresh_sid in remaining + assert open_sid in remaining + + async def test_delete_by_user_id_erases_all_turns_for_that_user_only( + self, manager, cm_actor + ): + org = cm_actor.organization_id + user_a = await _make_user(org) + user_b = await _make_user(org) + for user in (user_a, user_b): + await manager.record_turns( + session_id=f"erase-{uuid.uuid4().hex[:6]}", + user_id=user.id, + organization_id=org, + turns=[ + {"role": "user", "content": "sensitive"}, + {"role": "assistant", "content": "reply"}, + ], + actor=cm_actor, + ) + + deleted = await manager.delete_by_user_id(user_id=user_a.id) + + assert deleted == 2 + assert await self._count_all_rows(manager, user_a.id, org) == 0 + assert await self._count_all_rows(manager, user_b.id, org) == 2 + + async def test_delete_by_client_id_erases_only_that_clients_turns( + self, manager, cm_actor, cm_actor_org_b + ): + user_a = await _make_user(cm_actor.organization_id) + user_b = await _make_user(cm_actor_org_b.organization_id) + await manager.record_turns( + session_id=f"clienterase-{uuid.uuid4().hex[:6]}", + user_id=user_a.id, + organization_id=cm_actor.organization_id, + turns=[{"role": "user", "content": "ingested by client A"}], + actor=cm_actor, + ) + await manager.record_turns( + session_id=f"clienterase-{uuid.uuid4().hex[:6]}", + user_id=user_b.id, + organization_id=cm_actor_org_b.organization_id, + turns=[{"role": "user", "content": "ingested by client B"}], + actor=cm_actor_org_b, + ) + + deleted = await manager.delete_by_client_id(actor=cm_actor) + + assert deleted >= 1 + assert ( + await self._count_all_rows(manager, user_a.id, cm_actor.organization_id) + == 0 + ) + # Client B's ingested turns are untouched. + assert ( + await self._count_all_rows( + manager, user_b.id, cm_actor_org_b.organization_id + ) + == 1 + ) diff --git a/tests/test_conversation_turn_extraction.py b/tests/test_conversation_turn_extraction.py new file mode 100644 index 000000000..a7539ef4f --- /dev/null +++ b/tests/test_conversation_turn_extraction.py @@ -0,0 +1,338 @@ +"""Tests for conversation-turn extraction into the Conversation Message Store. + +The distiller's prompt has always treated tool errors/retries as strong +learning signals (`signal_type: tool_error`), but the ingestion seam used to +drop every non-user/assistant message โ€” blinding the distiller to exactly that +signal. These tests pin the repaired pipeline: tool results and assistant +tool_calls survive extraction, content parts are flattened to clean text, and +foreign roles are still dropped. +""" + +from types import SimpleNamespace + +import pytest + +from mirix.server.rest_api import ( + _extract_conversation_turns, + _ingest_session_turns, + _serialize_tool_calls, +) + + +class TestExtractConversationTurns: + def test_user_and_assistant_kept(self): + turns = _extract_conversation_turns( + [ + {"role": "user", "content": "deploy the app"}, + {"role": "assistant", "content": "done"}, + ] + ) + assert turns == [ + {"role": "user", "content": "deploy the app"}, + {"role": "assistant", "content": "done"}, + ] + + def test_tool_role_is_kept_with_name_prefix(self): + turns = _extract_conversation_turns( + [ + {"role": "user", "content": "run the tests"}, + {"role": "tool", "name": "run_tests", "content": "2 failed, 10 passed"}, + ] + ) + assert turns[1] == { + "role": "tool", + "content": "[run_tests] 2 failed, 10 passed", + } + + def test_function_role_maps_to_tool(self): + """OpenAI's legacy 'function' role is the same concept as 'tool'.""" + turns = _extract_conversation_turns( + [ + {"role": "user", "content": "q"}, + {"role": "function", "name": "search", "content": "no results"}, + ] + ) + assert turns[1]["role"] == "tool" + assert "no results" in turns[1]["content"] + + def test_assistant_tool_calls_are_serialized(self): + """A pure tool-call assistant turn (empty content) must still yield a + turn โ€” the call itself is the work-process signal.""" + turns = _extract_conversation_turns( + [ + {"role": "user", "content": "deploy"}, + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call_1", + "function": { + "name": "kubectl_apply", + "arguments": '{"file": "deploy.yaml"}', + }, + } + ], + }, + { + "role": "tool", + "name": "kubectl_apply", + "content": "error: forbidden", + }, + ] + ) + assert len(turns) == 3 + assert turns[1]["role"] == "assistant" + assert 'kubectl_apply({"file": "deploy.yaml"})' in turns[1]["content"] + assert turns[2] == { + "role": "tool", + "content": "[kubectl_apply] error: forbidden", + } + + def test_assistant_content_and_tool_calls_both_kept(self): + turns = _extract_conversation_turns( + [ + {"role": "user", "content": "q"}, + { + "role": "assistant", + "content": "let me check", + "tool_calls": [{"name": "search", "arguments": "query"}], + }, + ] + ) + assert turns[1]["content"].startswith("let me check") + assert "[tool_call] search(query)" in turns[1]["content"] + + def test_content_parts_flatten_to_text(self): + """The SDK's [{"type":"text","text":...}] parts must store clean text, + not dict reprs.""" + turns = _extract_conversation_turns( + [ + { + "role": "user", + "content": [ + {"type": "text", "text": "hello"}, + {"type": "text", "text": "world"}, + ], + }, + {"role": "assistant", "content": "hi"}, + ] + ) + assert turns[0]["content"] == "hello\nworld" + + def test_foreign_roles_are_dropped(self): + turns = _extract_conversation_turns( + [ + {"role": "user", "content": "q"}, + {"role": "system", "content": "scaffolding โ€” not a learnable turn"}, + {"role": "assistant", "content": "a"}, + ] + ) + assert [t["role"] for t in turns] == ["user", "assistant"] + + def test_non_role_bearing_payload_yields_nothing(self): + assert ( + _extract_conversation_turns([{"type": "text", "text": "screenshot"}]) == [] + ) + assert _extract_conversation_turns([]) == [] + + def test_oversized_turn_is_truncated_to_store_cap(self): + """One huge tool result must fail softly (truncate), not blow the whole + batch's Pydantic validation in record_turns.""" + from mirix.schemas.conversation_message import ( + CONVERSATION_MESSAGE_MAX_CONTENT_LEN, + ConversationMessageCreate, + ) + + huge = "x" * (CONVERSATION_MESSAGE_MAX_CONTENT_LEN + 1000) + turns = _extract_conversation_turns( + [ + {"role": "user", "content": "q"}, + {"role": "tool", "name": "dump", "content": huge}, + ] + ) + assert len(turns[1]["content"]) <= CONVERSATION_MESSAGE_MAX_CONTENT_LEN + assert turns[1]["content"].endswith("โ€ฆ[truncated]") + # And the truncated turn passes the store schema it will be validated by. + ConversationMessageCreate( + session_id="sess-1", + user_id="u", + organization_id="o", + role=turns[1]["role"], + content=turns[1]["content"], + ) + + def test_malformed_tool_calls_do_not_raise(self): + """Extraction must never abort the add path on odd shapes.""" + turns = _extract_conversation_turns( + [ + {"role": "user", "content": "q"}, + {"role": "assistant", "content": None, "tool_calls": "not-a-list"}, + { + "role": "assistant", + "content": "", + "tool_calls": [42, {"function": "not-a-dict"}], + }, + {"role": "tool", "content": {"weird": "dict"}}, + ] + ) + assert [t["role"] for t in turns] == ["user", "assistant", "assistant", "tool"] + + +class TestSerializeToolCalls: + def test_openai_shape(self): + rendered = _serialize_tool_calls( + [{"function": {"name": "grep", "arguments": '{"q": "x"}'}}] + ) + assert rendered == '[tool_call] grep({"q": "x"})' + + def test_flat_shape_and_dict_arguments(self): + rendered = _serialize_tool_calls( + [{"name": "read", "arguments": {"path": "a.py"}}] + ) + assert rendered == '[tool_call] read({"path": "a.py"})' + + def test_unrecognized_falls_back_to_str(self): + assert "[tool_call]" in _serialize_tool_calls(["opaque"]) + + +@pytest.mark.asyncio +class TestIngestSessionTurnsIsolation: + """The store write is ADDITIVE: a store or extraction failure must degrade + to a missed session for skill distillation, never abort the primary + /memory/add(_sync) ingestion.""" + + @staticmethod + def _request(session_id, messages): + return SimpleNamespace(session_id=session_id, messages=messages) + + @staticmethod + def _client(): + return SimpleNamespace(organization_id="org-1") + + async def test_record_turns_failure_does_not_propagate(self, monkeypatch): + from mirix.services.conversation_message_manager import ( + ConversationMessageManager, + ) + + async def boom(self, **kwargs): + raise RuntimeError("store down") + + monkeypatch.setattr(ConversationMessageManager, "record_turns", boom) + + request = self._request( + "sess-iso-1", + [ + {"role": "user", "content": "q"}, + {"role": "assistant", "content": "a"}, + ], + ) + # Must not raise โ€” the primary memory add would proceed. + await _ingest_session_turns(request, [], self._client(), "user-1") + + async def test_extraction_failure_does_not_propagate(self, monkeypatch): + """Extraction runs INSIDE the guard: a message whose content cannot be + rendered must not reach the store, and must not raise.""" + from mirix.services.conversation_message_manager import ( + ConversationMessageManager, + ) + + calls = [] + + async def recorder(self, **kwargs): + calls.append(kwargs) + + monkeypatch.setattr(ConversationMessageManager, "record_turns", recorder) + + class Unrenderable: + def __str__(self): + raise ValueError("unrenderable content part") + + request = self._request( + "sess-iso-2", [{"role": "user", "content": [Unrenderable()]}] + ) + await _ingest_session_turns(request, [], self._client(), "user-1") + assert calls == [] + + async def test_no_session_id_skips_store_entirely(self, monkeypatch): + from mirix.services.conversation_message_manager import ( + ConversationMessageManager, + ) + + calls = [] + + async def recorder(self, **kwargs): + calls.append(kwargs) + + monkeypatch.setattr(ConversationMessageManager, "record_turns", recorder) + + input_msg = SimpleNamespace(session_id=None) + request = self._request(None, [{"role": "user", "content": "q"}]) + await _ingest_session_turns(request, [input_msg], self._client(), "user-1") + assert calls == [] + assert input_msg.session_id is None # no stamping without a session + + async def test_session_id_stamped_and_turns_recorded(self, monkeypatch): + from mirix.services.conversation_message_manager import ( + ConversationMessageManager, + ) + + calls = [] + + async def recorder(self, **kwargs): + calls.append(kwargs) + + monkeypatch.setattr(ConversationMessageManager, "record_turns", recorder) + + unstamped = SimpleNamespace(session_id=None) + prestamped = SimpleNamespace(session_id="sess-own") + request = self._request( + "sess-batch", + [ + {"role": "user", "content": "q"}, + {"role": "tool", "name": "grep", "content": "no match"}, + {"role": "assistant", "content": "a"}, + ], + ) + await _ingest_session_turns( + request, [unstamped, prestamped], self._client(), "user-1" + ) + assert unstamped.session_id == "sess-batch" + assert prestamped.session_id == "sess-own" # own id wins over the batch id + assert len(calls) == 1 + assert calls[0]["session_id"] == "sess-batch" + assert calls[0]["user_id"] == "user-1" + assert calls[0]["organization_id"] == "org-1" + assert [t["role"] for t in calls[0]["turns"]] == ["user", "tool", "assistant"] + + +class TestToolTurnsRoundTripThroughSchema: + """The store schema must accept what extraction now produces.""" + + def test_tool_role_valid_in_schema(self): + from mirix.schemas.conversation_message import ConversationMessageCreate + + msg = ConversationMessageCreate( + session_id="sess-tool-1", + user_id="user-1", + organization_id="org-1", + role="tool", + content="[run_tests] 2 failed", + ) + assert msg.role == "tool" + + def test_distiller_renders_tool_turns(self): + from types import SimpleNamespace + + from mirix.services.session_experience_distiller import ( + SessionExperienceDistiller, + ) + + transcript = SessionExperienceDistiller._render_transcript( + [ + SimpleNamespace(role="user", content="deploy"), + SimpleNamespace(role="tool", content="[kubectl] error: forbidden"), + SimpleNamespace(role="assistant", content="fixed the RBAC and retried"), + ] + ) + assert "tool: [kubectl] error: forbidden" in transcript diff --git a/tests/test_deletion_apis.py b/tests/test_deletion_apis.py index 083169147..e3c161715 100644 --- a/tests/test_deletion_apis.py +++ b/tests/test_deletion_apis.py @@ -14,6 +14,7 @@ import asyncio import logging +import os import time import uuid from pathlib import Path @@ -340,8 +341,13 @@ async def test_3_delete_user_memories(client): for memory_type, count in counts_before.items(): logger.info(" %s: %d", memory_type, count) - # Delete memories via API - response = requests.delete(f"{BASE_URL}/users/{TEST_USER_ID}/memories") + # Delete memories via API. These endpoints are authenticated + tenant-scoped, + # so pass the test client's API key (the middleware resolves it to the + # X-Client-Id/X-Org-Id the tenant guard checks). + response = requests.delete( + f"{BASE_URL}/users/{TEST_USER_ID}/memories", + headers={"X-API-Key": os.environ["MIRIX_API_KEY"]}, + ) response.raise_for_status() result = response.json() @@ -401,8 +407,11 @@ async def test_5_delete_client_memories(client): for memory_type, count in counts_before.items(): logger.info(" %s: %d", memory_type, count) - # Delete memories via API - response = requests.delete(f"{BASE_URL}/clients/{TEST_CLIENT_ID}/memories") + # Delete memories via API (authenticated + tenant-scoped โ€” see the user case). + response = requests.delete( + f"{BASE_URL}/clients/{TEST_CLIENT_ID}/memories", + headers={"X-API-Key": os.environ["MIRIX_API_KEY"]}, + ) response.raise_for_status() result = response.json() diff --git a/tests/test_experience_to_skill_evolution.py b/tests/test_experience_to_skill_evolution.py new file mode 100644 index 000000000..a9ff5f8d0 --- /dev/null +++ b/tests/test_experience_to_skill_evolution.py @@ -0,0 +1,494 @@ +"""Goal-3 tests โ€” experiences -> procedural-skill self-evolution (DB-free). + +Drives the injectable core `_run_experience_evolution_core` with a fake +experience manager, fake snapshot/run_step collaborators, and a fake agent โ€” no +server, no DB, no LLM. Asserts the load-bearing contract: + +* budget mapping: worth_avoiding -> n_high_fail, worth_learning -> n_high_succ, + so the existing C4 formula (avoid weighted heavier than learn) applies; avoid + yields a budget >= the same count of learns. +* B_min=0 early-exit: no pending experiences -> skipped, NO step, NO consume. +* ordering invariant: snapshot BEFORE -> step -> snapshot AFTER -> diff -> + mark_consumed runs AFTER the diff (lineage can't be lost to a reset). +* no-delete gate: the agent gets a zero delete budget + empty auth set + soft + preference set at step time. +* lineage: influenced_skill_ids = created|edited|deleted from the diff, passed + to mark_consumed. + +Also a DB-free budget-mapping check and a compact-payload rendering check. +""" + +from __future__ import annotations + +import inspect +from types import SimpleNamespace + +import pytest + +from mirix.functions.function_sets.memory_tools import compute_edit_budget +from mirix.services.skill_experience_curator import ( + _run_experience_evolution_core, + build_experience_payload, +) + + +# ============================ Budget mapping (DB-free) ====================== + + +class TestBudgetMapping: + def test_avoid_weighted_heavier_than_learn(self): + # Same count: all-avoid budget >= all-learn budget (alpha_fail >= alpha_succ). + b_avoid = compute_edit_budget({"n_high_fail": 3, "n_high_succ": 0}) + b_learn = compute_edit_budget({"n_high_fail": 0, "n_high_succ": 3}) + assert b_avoid >= b_learn + + def test_zero_experiences_is_min_budget(self): + assert compute_edit_budget({"n_high_fail": 0, "n_high_succ": 0}) == 0 + + +# ============================ Payload rendering (DB-free) =================== + + +class TestPayloadRendering: + def test_compact_blocks_tagged_learn_avoid(self): + import json + + exps = [ + SimpleNamespace( + experience_type="worth_avoiding", title="Avoid X", + content="do not X", importance=0.9, credibility=0.9, + evidence=json.dumps({"quote": "no such column", "signal_type": "tool_error"}), + ), + SimpleNamespace( + experience_type="worth_learning", title="Do Y", + content="prefer Y", importance=0.5, credibility=0.5, + evidence=json.dumps({"quote": "great", "signal_type": "user_confirmation"}), + ), + ] + payload = build_experience_payload(exps) + assert "[AVOID] Avoid X" in payload + assert "[LEARN] Do Y" in payload + assert "no such column" in payload # evidence quote surfaced + assert "importance=0.90" in payload + + def test_empty_list_renders_placeholder(self): + payload = build_experience_payload([]) + assert "(no experiences)" in payload + + +# ============================ Fakes for the core =========================== + + +class _FakeExperienceManager: + """In-memory experience manager exposing the 3 methods the core calls.""" + + def __init__(self, experiences): + self._experiences = experiences + self.consumed_calls = [] # (ids, run_id, influenced) + self.superseded_calls = [] # [ids] + + async def list_experiences(self, *, agent_id, user_id, status, limit, ids=None): + assert status == "pending" + # Mirror the real manager: ids=None -> no filter; ids=[] -> nothing; + # otherwise scope to the given id set (this round's batch). + if ids is not None: + if not ids: + return [] + wanted = set(ids) + return [e for e in self._experiences if e.id in wanted][:limit] + return list(self._experiences)[:limit] + + async def aggregate(self, *, ids): + rows = [e for e in self._experiences if e.id in set(ids)] + n_avoid = sum(1 for e in rows if e.experience_type == "worth_avoiding") + n_learn = sum(1 for e in rows if e.experience_type == "worth_learning") + return { + "n": len(rows), + "n_worth_learning": n_learn, + "n_worth_avoiding": n_avoid, + "sum_priority": sum(e.importance * e.credibility for e in rows), + } + + async def mark_consumed(self, *, ids, run_id, influenced_skill_ids=None): + self.consumed_calls.append((list(ids), run_id, influenced_skill_ids)) + return len(ids) + + async def mark_superseded(self, *, ids, agent_id=None, user_id=None): + self.superseded_calls.append(list(ids)) + return len(ids) + + +class _Skill: + def __init__(self, sid, version=1, instructions="i", description="d"): + self.id = sid + self.version = version + self.instructions = instructions + self.description = description + + +def _exp(eid, etype, imp, cred): + return SimpleNamespace( + id=eid, experience_type=etype, title=f"title-{eid}", + content=f"content-{eid}", importance=imp, credibility=cred, + evidence='{"quote":"q","signal_type":"inferred"}', + ) + + +class _OrderTracker: + """Records the call order across snapshot/step/consume to assert ordering.""" + + def __init__(self): + self.events = [] + + +@pytest.mark.asyncio +class TestEvolutionCore: + async def test_b_min_skip_no_step_no_consume(self): + mgr = _FakeExperienceManager([]) # nothing pending + agent = SimpleNamespace() + tracker = _OrderTracker() + + async def snapshot(): + tracker.events.append("snapshot") + return [] + + async def run_step(a, payload, budget): + tracker.events.append("step") + + result = await _run_experience_evolution_core( + experience_manager=mgr, agent=agent, + agent_id="proc-1", meta_agent_id="meta-1", user_id="user-1", + snapshot_skills=snapshot, run_step=run_step, + ) + assert result["skipped"] is True + assert result["budget"] == 0 + assert "step" not in tracker.events # no step + assert mgr.consumed_calls == [] # no consume + + async def test_ordering_snapshot_step_snapshot_diff_then_consume(self): + exps = [ + _exp("sexp-a", "worth_avoiding", 0.9, 0.9), + _exp("sexp-b", "worth_learning", 0.7, 0.8), + ] + mgr = _FakeExperienceManager(exps) + agent = SimpleNamespace() + tracker = _OrderTracker() + + # before snapshot: skill s1 v1. after: s1 v2 (edited) + s2 (created). + snapshots = [ + [_Skill("s1", version=1)], + [_Skill("s1", version=2), _Skill("s2", version=1)], + ] + snap_iter = iter(snapshots) + + async def snapshot(): + tracker.events.append("snapshot") + return next(snap_iter) + + async def run_step(a, payload, budget): + tracker.events.append("step") + # The budget must be set on the agent BEFORE the step. + assert a._edit_budget_remaining == budget + + # mark_consumed records its order via a tracker wrapper. + orig_consume = mgr.mark_consumed + + async def tracked_consume(**kwargs): + tracker.events.append("consume") + return await orig_consume(**kwargs) + + mgr.mark_consumed = tracked_consume + + result = await _run_experience_evolution_core( + experience_manager=mgr, agent=agent, + agent_id="proc-1", meta_agent_id="meta-1", user_id="user-1", + snapshot_skills=snapshot, run_step=run_step, + ) + + # Ordering invariant: snapshot, step, snapshot, then consume LAST. + assert tracker.events == ["snapshot", "step", "snapshot", "consume"] + # Diff computed BEFORE consume. + assert set(result["changes"]["created"]) == {"s2"} + assert set(result["changes"]["edited"]) == {"s1"} + assert result["skills_changed"] == 2 + # Lineage = union of created|edited|deleted, passed to mark_consumed. + ids, run_id, influenced = mgr.consumed_calls[0] + assert set(ids) == {"sexp-a", "sexp-b"} + assert run_id == result["run_id"] + assert set(influenced) == {"s1", "s2"} + assert set(result["influenced_skill_ids"]) == {"s1", "s2"} + + async def test_step_failure_leaves_experiences_pending(self): + # If the step raises, the core must NOT consume the experiences (the + # consume runs AFTER a successful step), so a retry can re-process them. + exps = [_exp("sexp-a", "worth_avoiding", 0.9, 0.9)] + mgr = _FakeExperienceManager(exps) + agent = SimpleNamespace() + + async def snapshot(): + return [] + + async def run_step(a, payload, budget): + raise RuntimeError("step blew up") + + with pytest.raises(RuntimeError): + await _run_experience_evolution_core( + experience_manager=mgr, agent=agent, + agent_id="proc-1", meta_agent_id="meta-1", user_id="user-1", + snapshot_skills=snapshot, run_step=run_step, + ) + # mark_consumed must NOT have been called. + assert mgr.consumed_calls == [] + + async def test_no_delete_gate_set_at_step_time(self): + exps = [_exp("sexp-a", "worth_avoiding", 0.9, 0.9)] + mgr = _FakeExperienceManager(exps) + agent = SimpleNamespace() + observed = {} + + async def snapshot(): + return [] + + async def run_step(a, payload, budget): + observed["edit"] = a._edit_budget_remaining + observed["delete"] = a._delete_budget_remaining + observed["auth"] = a._delete_authorized_skill_ids + observed["soft"] = a._prefer_soft_delete + + await _run_experience_evolution_core( + experience_manager=mgr, agent=agent, + agent_id="proc-1", meta_agent_id="meta-1", user_id="user-1", + snapshot_skills=snapshot, run_step=run_step, + ) + assert observed["edit"] >= 1 + assert observed["delete"] == 0 + assert observed["auth"] == set() + assert observed["soft"] is True + + async def _budget_for(self, experiences): + """Drive the core with a no-op step; return (result, observed budget).""" + mgr = _FakeExperienceManager(experiences) + agent = SimpleNamespace() + observed = {} + + async def snapshot(): + return [] + + async def run_step(a, payload, budget): + observed["budget"] = budget + + result = await _run_experience_evolution_core( + experience_manager=mgr, agent=agent, + agent_id="proc-1", meta_agent_id="meta-1", user_id="user-1", + snapshot_skills=snapshot, run_step=run_step, + ) + return result, observed["budget"], mgr + + async def test_budget_reflects_avoid_count(self): + # 3 worth_avoiding -> budget mapped via n_high_fail. + exps = [_exp(f"sexp-{i}", "worth_avoiding", 0.9, 0.9) for i in range(3)] + result, budget, mgr = await self._budget_for(exps) + expected = compute_edit_budget({"n_high_fail": 3, "n_high_succ": 0}) + assert result["budget"] == expected + assert budget == expected + assert result["consumed_count"] == 3 + + async def test_budget_maps_worth_learning_to_n_high_succ(self): + # The core must map worth_learning -> n_high_succ (NOT n_high_fail and + # NOT ignored). All-learn budget must equal the formula's all-succ value. + exps = [_exp(f"sexp-{i}", "worth_learning", 0.8, 0.8) for i in range(3)] + result, budget, _ = await self._budget_for(exps) + expected_succ = compute_edit_budget({"n_high_fail": 0, "n_high_succ": 3}) + assert result["budget"] == expected_succ + # And it must NOT be mistakenly mapped to the avoid weight. + wrong_as_fail = compute_edit_budget({"n_high_fail": 3, "n_high_succ": 0}) + if wrong_as_fail != expected_succ: # only meaningful when weights differ + assert result["budget"] != wrong_as_fail + # Learns must not be silently ignored: 3 learns >= budget for 0 learns. + assert result["budget"] >= compute_edit_budget( + {"n_high_fail": 0, "n_high_succ": 0} + ) + + async def test_budget_mixed_avoid_and_learn(self): + # 2 avoid + 2 learn must map to the exact mixed formula value, proving + # both counts flow through (not just one). + exps = [ + _exp("sexp-a1", "worth_avoiding", 0.9, 0.9), + _exp("sexp-a2", "worth_avoiding", 0.7, 0.7), + _exp("sexp-l1", "worth_learning", 0.8, 0.8), + _exp("sexp-l2", "worth_learning", 0.6, 0.6), + ] + result, budget, _ = await self._budget_for(exps) + expected = compute_edit_budget({"n_high_fail": 2, "n_high_succ": 2}) + assert result["budget"] == expected + assert budget == expected + assert result["consumed_count"] == 4 + + +# ===================== Round scoping (this window only) ==================== + + +@pytest.mark.asyncio +class TestRoundScoping: + """A scoped evolve sees ONLY the round's freshly-distilled experiences. + + The auto-dream procedural path passes the ids it just distilled so an evolve + never re-surfaces earlier rounds' leftover pending experiences (which would + dilute the curator prompt and hurt skill quality). + """ + + @staticmethod + async def _run(mgr, *, experience_ids): + agent = SimpleNamespace() + seen = {} + + async def snapshot(): + return [] + + async def run_step(a, payload, budget): + seen["payload"] = payload + seen["budget"] = budget + + result = await _run_experience_evolution_core( + experience_manager=mgr, agent=agent, + agent_id="proc-1", meta_agent_id="meta-1", user_id="user-1", + snapshot_skills=snapshot, run_step=run_step, + experience_ids=experience_ids, + ) + return result, seen + + async def test_scopes_to_provided_ids_only(self): + # Pool has 2 "old" + 2 "new" pending; scope to the 2 new ones only. + old = [ + _exp("sexp-old1", "worth_avoiding", 0.9, 0.9), + _exp("sexp-old2", "worth_learning", 0.9, 0.9), + ] + new = [ + _exp("sexp-new1", "worth_avoiding", 0.6, 0.6), + _exp("sexp-new2", "worth_learning", 0.5, 0.5), + ] + mgr = _FakeExperienceManager(old + new) + new_ids = [e.id for e in new] + + result, seen = await self._run(mgr, experience_ids=new_ids) + + assert result["skipped"] is False + # Only the new batch is consumed โ€” the old pending stays untouched. + consumed_ids, _run, _infl = mgr.consumed_calls[0] + assert set(consumed_ids) == {"sexp-new1", "sexp-new2"} + assert result["consumed_count"] == 2 + # Budget reflects the scoped counts (1 avoid + 1 learn), NOT all four. + assert result["budget"] == compute_edit_budget( + {"n_high_fail": 1, "n_high_succ": 1} + ) + # The old experiences' titles must NOT appear in the curator payload. + assert "title-sexp-old1" not in seen["payload"] + assert "title-sexp-old2" not in seen["payload"] + assert "title-sexp-new1" in seen["payload"] + # Whole scoped batch fit under the cap -> no overflow to supersede. + assert mgr.superseded_calls == [] + + async def test_overflow_beyond_cap_is_superseded_not_stranded(self): + # A single round distilling MORE than the per-run cap must not strand its + # lowest-priority tail as pending forever (future rounds pass only THEIR + # fresh ids). The top _MAX_EXPERIENCES_PER_RUN are consumed; the rest of + # THIS round's scoped ids are superseded (out of sight, still auditable). + from mirix.services.skill_experience_curator import _MAX_EXPERIENCES_PER_RUN + + n = _MAX_EXPERIENCES_PER_RUN + 1 + exps = [_exp(f"sexp-{i:03d}", "worth_avoiding", 0.9, 0.9) for i in range(n)] + mgr = _FakeExperienceManager(exps) + all_ids = [e.id for e in exps] + + result, _seen = await self._run(mgr, experience_ids=all_ids) + + consumed_ids, _run, _infl = mgr.consumed_calls[0] + assert len(consumed_ids) == _MAX_EXPERIENCES_PER_RUN + assert result["consumed_count"] == _MAX_EXPERIENCES_PER_RUN + # Exactly the overflow tail (those not consumed) is superseded โ€” nothing + # from this round is left pending/stranded. + overflow = set(all_ids) - set(consumed_ids) + assert len(overflow) == 1 + assert set(mgr.superseded_calls[0]) == overflow + assert result["superseded_count"] == 1 + + async def test_empty_ids_skips_even_with_pending_pool(self): + # A round that distilled nothing (ids=[]) must NOT sweep the pending pool. + mgr = _FakeExperienceManager([ + _exp("sexp-old1", "worth_avoiding", 0.9, 0.9), + ]) + result, seen = await self._run(mgr, experience_ids=[]) + assert result["skipped"] is True + assert result["budget"] == 0 + assert "payload" not in seen # no step ran + assert mgr.consumed_calls == [] # nothing consumed + + async def test_none_ids_evolves_whole_pool(self): + # Backward-compat: ids=None means "no scope" -> the whole pending pool. + mgr = _FakeExperienceManager([ + _exp("sexp-a", "worth_avoiding", 0.9, 0.9), + _exp("sexp-b", "worth_learning", 0.7, 0.7), + ]) + result, _seen = await self._run(mgr, experience_ids=None) + consumed_ids, _run, _infl = mgr.consumed_calls[0] + assert set(consumed_ids) == {"sexp-a", "sexp-b"} + assert result["consumed_count"] == 2 + # Unscoped/global mode must NOT supersede leftovers โ€” a later global + # sweep is meant to pick them up. + assert mgr.superseded_calls == [] + + +# ============================ Module hygiene =============================== + + +class TestCuratorHygiene: + def test_core_is_async(self): + assert inspect.iscoroutinefunction(_run_experience_evolution_core) + + def test_no_asyncio_run(self): + import mirix.services.skill_experience_curator as mod + + src = inspect.getsource(mod) + assert "asyncio.run(" not in src + + def test_does_not_couple_to_metaclaw_record_store(self): + # The general curator must NOT drive benchmark-specific record stores. + # Its docs may contrast against other paths, but no executable line may + # reference benchmark manager/record types. Assert on code, stripping + # comments/docstrings. + import mirix.services.skill_experience_curator as mod + + src = inspect.getsource(mod) + code_only = _strip_comments_and_docstrings(src).lower() + for banned in [ + "skillevolutionrecord", + "run_records_evolution", + "record_round_result", + "round_index", + "quality_score", + ]: + assert banned not in code_only, f"benchmark coupling leaked into code: {banned!r}" + + +def _strip_comments_and_docstrings(src: str) -> str: + """Return ``src`` with comments and string/docstring literals removed. + + Tokenize-based so contrastive docs/comments do not trip the coupling + assertion โ€” only executable code is checked. + """ + import io + import tokenize + + out = [] + prev_type = tokenize.INDENT + tokens = tokenize.generate_tokens(io.StringIO(src).readline) + for tok_type, tok_str, _start, _end, _line in tokens: + if tok_type in (tokenize.COMMENT, tokenize.STRING): + continue + if tok_type == tokenize.NL or tok_type == tokenize.NEWLINE: + out.append("\n") + else: + out.append(tok_str + " ") + prev_type = tok_type + _ = prev_type + return "".join(out) diff --git a/tests/test_keyed_locks.py b/tests/test_keyed_locks.py new file mode 100644 index 000000000..7844992c8 --- /dev/null +++ b/tests/test_keyed_locks.py @@ -0,0 +1,110 @@ +"""Tests for the shared per-key lock registry. + +The registry replaces three hand-rolled `dict[key, asyncio.Lock]` registries +(raw memory updates, per-agent skill evolution, per-user procedural dreams) +whose entries were never evicted โ€” a slow leak on long-lived servers. These +tests pin the two properties the call sites rely on: per-key mutual exclusion, +and eviction back to empty once nobody holds or waits. +""" + +import asyncio + +import pytest + +from mirix.helpers.keyed_locks import KeyedLocks + + +@pytest.mark.asyncio +async def test_mutual_exclusion_per_key(): + locks = KeyedLocks() + order = [] + + async def worker(tag: str): + async with locks.acquire("k"): + order.append(f"{tag}-in") + await asyncio.sleep(0) # yield while holding โ€” the other must wait + order.append(f"{tag}-out") + + await asyncio.gather(worker("a"), worker("b")) + + # Critical sections must not interleave: each -in is followed by its -out. + assert order in (["a-in", "a-out", "b-in", "b-out"], + ["b-in", "b-out", "a-in", "a-out"]) + + +@pytest.mark.asyncio +async def test_different_keys_do_not_block_each_other(): + locks = KeyedLocks() + entered = asyncio.Event() + release = asyncio.Event() + + async def holder(): + async with locks.acquire("k1"): + entered.set() + await release.wait() + + task = asyncio.create_task(holder()) + await entered.wait() + # k2 is acquirable immediately even while k1 is held. + async with locks.acquire("k2"): + pass + release.set() + await task + + +@pytest.mark.asyncio +async def test_entries_are_evicted_when_idle(): + """The leak-fix property: the registry returns to empty after use.""" + locks = KeyedLocks() + + for i in range(100): + async with locks.acquire(f"key-{i}"): + pass + + assert len(locks) == 0 + + +@pytest.mark.asyncio +async def test_entry_survives_while_a_waiter_is_queued(): + """Eviction must not drop a lock another coroutine is still waiting on โ€” + that would hand out a NEW lock for the same key and break exclusion.""" + locks = KeyedLocks() + holding = asyncio.Event() + proceed = asyncio.Event() + order = [] + + async def first(): + async with locks.acquire("k"): + holding.set() + await proceed.wait() + order.append("first-out") + + async def second(): + await holding.wait() + async with locks.acquire("k"): + order.append("second-in") + + t1 = asyncio.create_task(first()) + t2 = asyncio.create_task(second()) + await holding.wait() + await asyncio.sleep(0) # let `second` queue up on the same entry + assert len(locks) == 1 # one live entry, refcounted by holder + waiter + proceed.set() + await asyncio.gather(t1, t2) + + assert order == ["first-out", "second-in"] + assert len(locks) == 0 + + +@pytest.mark.asyncio +async def test_exception_inside_critical_section_still_evicts(): + locks = KeyedLocks() + + with pytest.raises(RuntimeError): + async with locks.acquire("k"): + raise RuntimeError("boom") + + assert len(locks) == 0 + # And the key is usable again afterwards. + async with locks.acquire("k"): + pass diff --git a/tests/test_memory_agent_toolcall_truncation.py b/tests/test_memory_agent_toolcall_truncation.py deleted file mode 100644 index b56fe9364..000000000 --- a/tests/test_memory_agent_toolcall_truncation.py +++ /dev/null @@ -1,118 +0,0 @@ -from datetime import datetime -from unittest.mock import MagicMock - -import pytest - -from mirix.agent.agent import Agent -from mirix.schemas.agent import AgentState, AgentType -from mirix.schemas.client import Client -from mirix.schemas.enums import ToolType -from mirix.schemas.message import Message -from mirix.schemas.openai.chat_completion_response import ( - FunctionCall, -) -from mirix.schemas.openai.chat_completion_response import Message as ChatCompletionMessage -from mirix.schemas.openai.chat_completion_response import ( - ToolCall, -) -from mirix.schemas.tool import Tool - - -def make_client(id="client-1", org_id="org-1"): - return Client( - id=id, - organization_id=org_id, - name="Test Client", - status="active", - write_scope="test", - read_scopes=["test"], - created_at=datetime.now(), - updated_at=datetime.now(), - is_deleted=False, - ) - - -@pytest.mark.asyncio -async def test_memory_agent_truncates_extra_tool_calls_and_executes_only_first(): - """ - Regression test for memory-agent multi-tool-call bug: - If the LLM returns multiple tool calls in one response, memory agents must - truncate to only the first tool call to avoid duplicate state mutations. - """ - - # Create a minimal Agent instance without running __init__ - agent = Agent.__new__(Agent) - agent.logger = MagicMock() - agent.interface = MagicMock() - agent.model = "test-model" - agent.last_function_response = None - - # ToolRulesSolver is referenced later in the agent loop; keep it simple here. - agent.tool_rules_solver = MagicMock() - agent.tool_rules_solver.update_tool_usage = MagicMock() - agent.tool_rules_solver.has_children_tools = MagicMock(return_value=False) - agent.tool_rules_solver.is_terminal_tool = MagicMock(return_value=False) - - # Provide an AgentState-like object - agent_state = MagicMock(spec=AgentState) - agent_state.id = "agent-123" - agent_state.name = "semantic_memory_agent" - agent_state.agent_type = AgentType.semantic_memory_agent - - # Provide tools that match the function names in tool calls - tool_1 = Tool( - tool_type=ToolType.MIRIX_MEMORY_CORE, - name="semantic_memory_update", - json_schema={"name": "semantic_memory_update", "description": "test", "parameters": {}}, - return_char_limit=10000, - ) - tool_2 = Tool( - tool_type=ToolType.MIRIX_MEMORY_CORE, - name="semantic_memory_insert", - json_schema={"name": "semantic_memory_insert", "description": "test", "parameters": {}}, - return_char_limit=10000, - ) - agent_state.tools = [tool_1, tool_2] - agent.agent_state = agent_state - - executed = [] - - async def _fake_exec(function_name, function_args, *args, **kwargs): - executed.append(function_name) - return "ok" - - agent.execute_tool_and_persist_state = _fake_exec - - # Build input/response messages - input_message = Message.dict_to_message( - id="message-12345678", - agent_id=agent_state.id, - model=agent.model, - openai_message_dict={"role": "user", "content": "hi"}, - ) - - response_message = ChatCompletionMessage( - role="assistant", - content=None, - tool_calls=[ - ToolCall( - id="call-1", - function=FunctionCall(name="semantic_memory_update", arguments="{}"), - ), - ToolCall( - id="call-2", - function=FunctionCall(name="semantic_memory_insert", arguments="{}"), - ), - ], - ) - - await agent._handle_ai_response( - input_message=input_message, - response_message=response_message, - existing_file_uris=[], - response_message_id="message-87654321", - retrieved_memories=None, - chaining=False, - ) - - assert executed == ["semantic_memory_update"] diff --git a/tests/test_memory_api_surface.py b/tests/test_memory_api_surface.py new file mode 100644 index 000000000..5cc26f90f --- /dev/null +++ b/tests/test_memory_api_surface.py @@ -0,0 +1,101 @@ +"""Production REST surface for memory APIs. + +Procedural memory is learned by the automatic conversation-store distillation +pipeline. The public REST surface exposes generic memory ingestion and search, +plus DELETE for parity with every other memory type โ€” but no direct skill +WRITE controls (create/edit go through the evolution flow only). +""" + +from mirix.server.rest_api import router +from mirix.server.rest_api import _procedural_memory_response + + +def _routes() -> list: + return [ + route + for route in router.routes + if hasattr(route, "path") and getattr(route, "include_in_schema", True) + ] + + +def _paths() -> set[str]: + return {route.path for route in _routes()} + + +def test_public_memory_surface_excludes_direct_skill_write_routes(): + paths = _paths() + + assert "/memory/add" in paths + assert "/memory/add_sync" in paths + assert "/memory/search" in paths + assert "/memory/search_all_users" in paths + + assert not any(path.startswith("/v1/skills") for path in paths) + + # DELETE stays (uniform with episodic/semantic/resource/knowledge_vault); + # any other method on /memory/procedural* is a skill write and must not + # be public. + procedural_methods = { + method + for route in _routes() + if route.path.startswith("/memory/procedural") + for method in (getattr(route, "methods", None) or set()) + } + assert procedural_methods == {"DELETE"} + + +def test_procedural_delete_route_is_uniform_with_other_memory_types(): + paths = _paths() + for memory_type in ("episodic", "semantic", "procedural", "resource"): + assert f"/memory/{memory_type}/{{memory_id}}" in paths + + +def test_bulk_erasure_endpoints_require_authentication(): + """DELETE /users/{id}/memories and /clients/{id}/memories are irreversible + cross-table purges (including verbatim conversation transcripts); they must + resolve the caller. Pin that both endpoint signatures take the + authorization header + request the auth helper needs.""" + import inspect + + from mirix.server.rest_api import delete_client_memories, delete_user_memories + + for endpoint in (delete_user_memories, delete_client_memories): + params = inspect.signature(endpoint).parameters + assert "authorization" in params, f"{endpoint.__name__} missing auth header" + assert "http_request" in params, f"{endpoint.__name__} missing request param" + + +def test_procedural_read_shape_is_full_skill_schema(): + from datetime import datetime + from types import SimpleNamespace + + item = SimpleNamespace( + id="proc-1", + entry_type="workflow", + name="deploy-production", + description="Deploy the service", + instructions="Run tests, merge, and monitor.", + triggers=["user asks to deploy"], + examples=[{"input": "deploy", "output": "checklist"}], + version="0.2.0", + created_at=datetime(2026, 1, 1, 12, 0, 0), + updated_at=datetime(2026, 1, 2, 12, 0, 0), + user_id="user-1", + ) + + row = _procedural_memory_response(item, include_user_id=True) + + assert row == { + "memory_type": "procedural", + "id": "proc-1", + "entry_type": "workflow", + "name": "deploy-production", + "description": "Deploy the service", + "instructions": "Run tests, merge, and monitor.", + "triggers": ["user asks to deploy"], + "examples": [{"input": "deploy", "output": "checklist"}], + "version": "0.2.0", + "created_at": "2026-01-01T12:00:00", + "updated_at": "2026-01-02T12:00:00", + "user_id": "user-1", + } diff --git a/tests/test_message_retention.py b/tests/test_message_retention.py new file mode 100644 index 000000000..24a26189c --- /dev/null +++ b/tests/test_message_retention.py @@ -0,0 +1,354 @@ +""" +Tests for the configurable last-N-session raw-message retention feature. + +GOAL 1: stop hard-deleting the canonical conversation messages immediately after +memory extraction; retain at least the last N sessions of raw messages so a later +distiller / auto-dream can read them. Retention != in-context: retained rows stay +DETACHED from message_ids (never re-added). + +This file has two parts: + * TestSelectRetainedSessionIds โ€” DB-FREE unit tests of the pure selection helper + MessageManager._select_retained_session_ids. These run anywhere (no Postgres). + * TestDeleteDetachedRetention โ€” integration tests (marked `integration`) that seed + N+2 sessions for one (agent, user), call delete_detached_messages_for_agent with + retain_last_n_sessions=N, and assert exactly the last N distinct session_ids + survive while older ones are deleted, and a second user's sessions are + unaffected. Requires the docker-compose Postgres. Run: + pytest tests/test_message_retention.py -v -m integration +""" +from __future__ import annotations + +import asyncio +import sys +import uuid +from datetime import datetime, timedelta, timezone +from pathlib import Path + +import pytest + +project_root = Path(__file__).parent.parent +sys.path.insert(0, str(project_root)) + +from mirix.services.message_manager import MessageManager + + +# --------------------------------------------------------------------------- +# DB-free unit tests of the pure selection helper. +# --------------------------------------------------------------------------- +class TestSelectRetainedSessionIds: + def _ts(self, minute): + return datetime(2026, 1, 1, 0, minute, 0, tzinfo=timezone.utc) + + def test_zero_retains_nothing(self): + mapping = {"s1": self._ts(1), "s2": self._ts(2)} + assert MessageManager._select_retained_session_ids(mapping, 0) == set() + + def test_negative_retains_nothing(self): + mapping = {"s1": self._ts(1)} + assert MessageManager._select_retained_session_ids(mapping, -1) == set() + + def test_empty_mapping_retains_nothing(self): + assert MessageManager._select_retained_session_ids({}, 5) == set() + + def test_selects_most_recent_n_by_min_created_at(self): + # s1 oldest .. s5 newest by first-seen timestamp. + mapping = { + "s1": self._ts(1), + "s2": self._ts(2), + "s3": self._ts(3), + "s4": self._ts(4), + "s5": self._ts(5), + } + # Retain last 3 -> the 3 most-recent first-seen sessions. + assert MessageManager._select_retained_session_ids(mapping, 3) == { + "s3", + "s4", + "s5", + } + + def test_retain_n_larger_than_population_keeps_all(self): + mapping = {"s1": self._ts(1), "s2": self._ts(2)} + assert MessageManager._select_retained_session_ids(mapping, 10) == {"s1", "s2"} + + def test_default_five_keeps_last_five_of_seven(self): + mapping = {f"s{i}": self._ts(i) for i in range(1, 8)} # s1..s7 + retained = MessageManager._select_retained_session_ids(mapping, 5) + assert retained == {"s3", "s4", "s5", "s6", "s7"} + # The two oldest age out and become eligible for deletion. + assert "s1" not in retained + assert "s2" not in retained + + def test_tie_break_is_deterministic_on_session_id(self): + # Two sessions share the same MIN(created_at); tie broken by session_id desc. + ts = self._ts(1) + mapping = {"aaa": ts, "bbb": ts, "ccc": self._ts(2)} + # Retain 2 -> the newest (ccc) plus the higher session_id among the tie (bbb). + assert MessageManager._select_retained_session_ids(mapping, 2) == {"ccc", "bbb"} + + def test_none_timestamp_sessions_treated_as_oldest(self): + # A session whose rows all lack created_at (None ts) is treated as OLDEST: + # excluded when N < population, and only included (in the oldest slot) when + # N covers it. Guards against comparing None to datetime. + mapping = {"s1": self._ts(1), "s2": self._ts(2), "snull": None} + # Retain 2 -> the two timestamped sessions; the None-ts session is evicted. + assert MessageManager._select_retained_session_ids(mapping, 2) == {"s1", "s2"} + # Retain 3 -> the None-ts session is included as the oldest slot. + assert MessageManager._select_retained_session_ids(mapping, 3) == {"s1", "s2", "snull"} + + +# --------------------------------------------------------------------------- +# Integration tests against real Postgres. +# --------------------------------------------------------------------------- +pytest_asyncio = pytest.importorskip("pytest_asyncio") + +from mirix.schemas.client import Client as PydanticClient # noqa: E402 +from mirix.schemas.enums import MessageRole # noqa: E402 +from mirix.schemas.message import Message as PydanticMessage # noqa: E402 +from mirix.schemas.mirix_message_content import TextContent # noqa: E402 +from mirix.schemas.user import User as PydanticUser # noqa: E402 + + +@pytest_asyncio.fixture(scope="module") +def event_loop(): + loop = asyncio.new_event_loop() + yield loop + loop.close() + + +@pytest.fixture +def message_manager(): + return MessageManager() + + +@pytest_asyncio.fixture(scope="module") +async def test_actor(): + from mirix.schemas.organization import Organization as PydanticOrganization + from mirix.services.client_manager import ClientManager + from mirix.services.organization_manager import OrganizationManager + + org_mgr = OrganizationManager() + client_mgr = ClientManager() + + org_id = f"test-retain-org-{uuid.uuid4().hex[:8]}" + try: + await org_mgr.get_organization_by_id(org_id) + except Exception: + await org_mgr.create_organization( + PydanticOrganization(id=org_id, name="Retention Test Org") + ) + + client_id = f"test-retain-client-{uuid.uuid4().hex[:8]}" + try: + return await client_mgr.get_client_by_id(client_id) + except Exception: + return await client_mgr.create_client( + PydanticClient( + id=client_id, + organization_id=org_id, + name="Retention Test Client", + write_scope="test-retain", + read_scopes=["test-retain"], + ) + ) + + +async def _make_user(test_actor, label): + from mirix.services.user_manager import UserManager + + user_mgr = UserManager() + user_id = f"test-retain-user-{label}-{uuid.uuid4().hex[:8]}" + return await user_mgr.create_user( + PydanticUser( + id=user_id, + name=f"Retention Test User {label}", + organization_id=test_actor.organization_id, + timezone="UTC", + ) + ) + + +@pytest_asyncio.fixture(scope="module") +async def user_a(test_actor): + return await _make_user(test_actor, "A") + + +@pytest_asyncio.fixture(scope="module") +async def user_b(test_actor): + return await _make_user(test_actor, "B") + + +@pytest_asyncio.fixture(scope="module") +async def meta_agent(test_actor): + """A meta_memory_agent-shaped agent (canonical conversation recipient).""" + from mirix.schemas.agent import AgentType, CreateAgent + from mirix.services.agent_manager import AgentManager + + agent_mgr = AgentManager() + return await agent_mgr.create_agent( + agent_create=CreateAgent( + name=f"test-retain-meta-{uuid.uuid4().hex[:8]}", + agent_type=AgentType.meta_memory_agent, + description="Test meta agent for retention", + system=None, + llm_config=None, + embedding_config=None, + ), + actor=test_actor, + ) + + +async def _seed_session(mgr, actor, agent, user_id, session_id, n_msgs=2): + """Create n_msgs detached messages for (agent, user, session).""" + created = [] + for i in range(n_msgs): + msg = await mgr.create_message( + pydantic_msg=PydanticMessage( + agent_id=agent.id, + role=MessageRole.user, + content=[TextContent(text=f"{session_id}-{i}")], + session_id=session_id, + ), + actor=actor, + user_id=user_id, + use_cache=False, + ) + created.append(msg) + return created + + +pytestmark_integration = pytest.mark.integration + + +class TestDeleteDetachedRetention: + @pytest.mark.integration + @pytest.mark.asyncio + async def test_retains_last_n_sessions_and_scopes_per_user( + self, message_manager, test_actor, meta_agent, user_a, user_b + ): + N = 3 + # User A: seed N+2 = 5 sessions, ordered oldest->newest by creation. + a_sessions = [f"a-sess-{i}-{uuid.uuid4().hex[:6]}" for i in range(N + 2)] + for sid in a_sessions: + await _seed_session( + message_manager, test_actor, meta_agent, user_a.id, sid + ) + + # User B: seed 2 sessions on the SAME agent โ€” must be unaffected by A's window. + b_sessions = [f"b-sess-{i}-{uuid.uuid4().hex[:6]}" for i in range(2)] + for sid in b_sessions: + await _seed_session( + message_manager, test_actor, meta_agent, user_b.id, sid + ) + + # All messages are detached (agent has empty message_ids), so without + # retention they would all be deleted. Call delete scoped to user A, + # retaining the last N sessions. + await message_manager.delete_detached_messages_for_agent( + agent_id=meta_agent.id, + actor=test_actor, + retain_last_n_sessions=N, + user_id=user_a.id, + ) + + # The last N session_ids for user A survive; the 2 oldest are deleted. + surviving_a = a_sessions[-N:] + deleted_a = a_sessions[:-N] + + for sid in surviving_a: + got = await message_manager.list_messages_for_agent( + agent_id=meta_agent.id, + actor=test_actor, + session_id=sid, + limit=100, + use_cache=False, + ) + assert len(got) == 2, f"expected retained session {sid} to survive" + + for sid in deleted_a: + got = await message_manager.list_messages_for_agent( + agent_id=meta_agent.id, + actor=test_actor, + session_id=sid, + limit=100, + use_cache=False, + ) + assert got == [], f"expected aged-out session {sid} to be deleted" + + # User B's sessions were NOT in user A's retained set, but because we + # scoped the delete to user_id=user_a.id, user B's detached rows were not + # even considered for the retained-set computation. They are still detached + # though โ€” they DID get deleted because they are not in message_ids and not + # in A's retained set. Re-seed-and-verify the per-user retained SET instead: + # rebuild B fresh and retain scoped to B. + b_sessions2 = [f"b2-sess-{i}-{uuid.uuid4().hex[:6]}" for i in range(N + 1)] + for sid in b_sessions2: + await _seed_session( + message_manager, test_actor, meta_agent, user_b.id, sid + ) + + # Retain last N for user B โ€” user A's surviving sessions must NOT count + # toward or evict B's window (per-user scoping). + await message_manager.delete_detached_messages_for_agent( + agent_id=meta_agent.id, + actor=test_actor, + retain_last_n_sessions=N, + user_id=user_b.id, + ) + + # B keeps its last N; the oldest B session ages out. + for sid in b_sessions2[-N:]: + got = await message_manager.list_messages_for_agent( + agent_id=meta_agent.id, + actor=test_actor, + session_id=sid, + limit=100, + use_cache=False, + ) + assert len(got) == 2, f"expected retained B session {sid} to survive" + + got_old_b = await message_manager.list_messages_for_agent( + agent_id=meta_agent.id, + actor=test_actor, + session_id=b_sessions2[0], + limit=100, + use_cache=False, + ) + assert got_old_b == [], "expected oldest B session to be deleted" + + # CRITICAL per-user isolation: deleting/retaining for user B must NOT have + # touched user A's surviving sessions. + for sid in surviving_a: + got = await message_manager.list_messages_for_agent( + agent_id=meta_agent.id, + actor=test_actor, + session_id=sid, + limit=100, + use_cache=False, + ) + assert len(got) == 2, ( + f"user A retained session {sid} must be unaffected by user B's delete" + ) + + @pytest.mark.integration + @pytest.mark.asyncio + async def test_retain_zero_deletes_all_detached( + self, message_manager, test_actor, meta_agent, user_a + ): + sid = f"zero-sess-{uuid.uuid4().hex[:6]}" + await _seed_session(message_manager, test_actor, meta_agent, user_a.id, sid) + + await message_manager.delete_detached_messages_for_agent( + agent_id=meta_agent.id, + actor=test_actor, + retain_last_n_sessions=0, + user_id=user_a.id, + ) + + got = await message_manager.list_messages_for_agent( + agent_id=meta_agent.id, + actor=test_actor, + session_id=sid, + limit=100, + use_cache=False, + ) + assert got == [], "retain=0 must preserve legacy full-delete behavior" diff --git a/tests/test_procedural_hybrid_search.py b/tests/test_procedural_hybrid_search.py new file mode 100644 index 000000000..ba0c089d5 --- /dev/null +++ b/tests/test_procedural_hybrid_search.py @@ -0,0 +1,190 @@ +"""Unit tests for the EverOS-aligned procedural HYBRID retrieval path. + +Covers the two pieces most prone to silent divergence from EverOS/everalgo: + + 1. `_rrf_fuse` โ€” the Reciprocal Rank Fusion math. Must be rank-based (raw + lane scores discarded), 1-based, unweighted, k=60, dedup-by-id, and the + cross-lane-agreement property (an item ranked in BOTH lanes outranks an + item that tops only one lane). + 2. `_hybrid_search_for_user` orchestration โ€” runs exactly the bm25 + embedding + lanes, over-fetches `limit * SKILL_HYBRID_RECALL_MULTIPLIER` per lane, + applies the similarity_threshold quality floor to the DENSE lane only, and + falls the dense lane's search_field back to "description" when the caller's + field is not embeddable. + +Both are exercised WITHOUT a database: `_rrf_fuse` is pure, and the +orchestration test monkeypatches the recursive `list_procedures` lane calls, so +the manager's session_maker is never touched. +""" + +from types import SimpleNamespace + +import pytest + +from mirix.constants import SKILL_HYBRID_RECALL_MULTIPLIER, SKILL_HYBRID_RRF_K +from mirix.services.procedural_memory_manager import ( + ProceduralMemoryManager, + _rrf_fuse, +) + + +def _item(item_id): + """Minimal stand-in: _rrf_fuse only reads `.id`.""" + return SimpleNamespace(id=item_id) + + +class TestRrfFuse: + def test_k_is_canonical_60(self): + assert SKILL_HYBRID_RRF_K == 60 + + def test_rank_based_unweighted_order(self): + # k=60, 1-based: score(d) = sum 1/(60+rank). + sparse = [_item("a"), _item("b"), _item("c")] + dense = [_item("b"), _item("a"), _item("d")] + fused = _rrf_fuse([sparse, dense], k=60, limit=10) + # a: 1/61 + 1/62 ; b: 1/62 + 1/61 (tie with a); c: 1/63 ; d: 1/63. + # a,b (in both) >> c,d (one lane each). Ties broken by first-seen order. + assert [i.id for i in fused] == ["a", "b", "c", "d"] + + def test_cross_lane_agreement_beats_single_lane_top(self): + # `shared` is rank 2 in BOTH lanes; x/y are rank 1 in one lane only. + # Agreement (2/62) must outrank a single rank-1 hit (1/61). + sparse = [_item("x"), _item("shared")] + dense = [_item("y"), _item("shared")] + fused = _rrf_fuse([sparse, dense], k=60, limit=10) + assert fused[0].id == "shared" + assert set(i.id for i in fused) == {"shared", "x", "y"} + + def test_scores_are_rank_based_not_value_based(self): + # Even if one lane "wanted" to dominate, only POSITION matters: an item + # at rank 1 in lane A and absent from lane B scores exactly 1/(k+1). + only_a = _rrf_fuse([[_item("solo")], []], k=60, limit=10) + assert [i.id for i in only_a] == ["solo"] + + def test_limit_slices_after_fusion(self): + sparse = [_item("a"), _item("b"), _item("c")] + dense = [_item("a"), _item("b"), _item("c")] + fused = _rrf_fuse([sparse, dense], k=60, limit=2) + assert [i.id for i in fused] == ["a", "b"] + + def test_falsy_limit_returns_all(self): + sparse = [_item("a"), _item("b")] + fused = _rrf_fuse([sparse, []], k=60, limit=0) + assert {i.id for i in fused} == {"a", "b"} + + def test_none_id_items_are_skipped(self): + # Items without a stable id cannot be deduped across lanes -> skipped. + fused = _rrf_fuse([[_item(None), _item("a")], [_item("a")]], k=60, limit=10) + assert [i.id for i in fused] == ["a"] + + def test_empty_lanes(self): + assert _rrf_fuse([[], []], k=60, limit=10) == [] + + def test_dedup_keeps_first_seen_object(self): + # The object returned for a shared id is the one first encountered. + first = SimpleNamespace(id="dup", lane="sparse") + second = SimpleNamespace(id="dup", lane="dense") + fused = _rrf_fuse([[first], [second]], k=60, limit=10) + assert len(fused) == 1 + assert fused[0].lane == "sparse" + + +class TestHybridOrchestration: + @pytest.mark.asyncio + async def test_runs_both_lanes_overfetches_and_fuses(self, monkeypatch): + mgr = ProceduralMemoryManager() + calls = [] + + async def fake_list_procedures(**kwargs): + calls.append(kwargs) + if kwargs["search_method"] == "bm25": + return [_item("a"), _item("b")] + return [_item("b"), _item("c")] # embedding lane + + monkeypatch.setattr(mgr, "list_procedures", fake_list_procedures) + + out = await mgr._hybrid_search_for_user( + agent_state=None, + user=None, + query="fix dates", + embedded_text=None, + search_field="description", + limit=5, + timezone_str=None, + filter_tags=None, + scopes=None, + use_cache=True, + similarity_threshold=None, + ) + + # Exactly the two lanes ran. + assert sorted(c["search_method"] for c in calls) == ["bm25", "embedding"] + # Each lane over-fetched limit * multiplier. + for c in calls: + assert c["limit"] == 5 * SKILL_HYBRID_RECALL_MULTIPLIER + # `b` appears in both lanes -> fused to the top. + assert out[0].id == "b" + assert {i.id for i in out} == {"a", "b", "c"} + # Final result sliced back to the requested limit. + assert len(out) <= 5 + + @pytest.mark.asyncio + async def test_similarity_threshold_is_dense_only(self, monkeypatch): + mgr = ProceduralMemoryManager() + by_method = {} + + async def fake_list_procedures(**kwargs): + by_method[kwargs["search_method"]] = kwargs + return [] + + monkeypatch.setattr(mgr, "list_procedures", fake_list_procedures) + + await mgr._hybrid_search_for_user( + agent_state=None, + user=None, + query="q", + embedded_text=None, + search_field="description", + limit=3, + timezone_str=None, + filter_tags=None, + scopes=None, + use_cache=True, + similarity_threshold=0.4, + ) + + # The BM25 lane must NOT receive the cosine quality floor... + assert "similarity_threshold" not in by_method["bm25"] + # ...but the dense lane must, applied BEFORE fusion. + assert by_method["embedding"]["similarity_threshold"] == 0.4 + + @pytest.mark.asyncio + async def test_dense_lane_field_falls_back_when_not_embeddable(self, monkeypatch): + mgr = ProceduralMemoryManager() + by_method = {} + + async def fake_list_procedures(**kwargs): + by_method[kwargs["search_method"]] = kwargs + return [] + + monkeypatch.setattr(mgr, "list_procedures", fake_list_procedures) + + # "entry_type" is searchable lexically but has no embedding column. + await mgr._hybrid_search_for_user( + agent_state=None, + user=None, + query="q", + embedded_text=None, + search_field="entry_type", + limit=3, + timezone_str=None, + filter_tags=None, + scopes=None, + use_cache=True, + similarity_threshold=None, + ) + + # Lexical lane keeps the requested field; dense lane falls back to + # "description" (the only sensible embedding column). + assert by_method["bm25"]["search_field"] == "entry_type" + assert by_method["embedding"]["search_field"] == "description" diff --git a/tests/test_redis_index_upgrade.py b/tests/test_redis_index_upgrade.py new file mode 100644 index 000000000..82f5caf23 --- /dev/null +++ b/tests/test_redis_index_upgrade.py @@ -0,0 +1,118 @@ +"""DB-free tests for the procedural Redis index schema upgrade. + +An FT index created before the skill schema (summary / summary_embedding / +steps_embedding) persists across deploys. Because index creation used to +early-return whenever FT.INFO succeeded, that stale index was never rebuilt and +every skill-field query (description/instructions, description_embedding KNN) +missed Redis forever. These tests pin the upgrade behavior with a mocked +client: fresh โ†’ create; current schema โ†’ leave alone; stale schema โ†’ drop the +index (keep documents) and recreate. +""" + +import pytest + +from mirix.database.redis_client import RedisMemoryClient + + +class _FakeFT: + def __init__(self, info_result=None, info_raises=False, drop_raises=False): + self._info_result = info_result + self._info_raises = info_raises + self._drop_raises = drop_raises + self.drop_calls = [] + self.create_calls = [] + + async def info(self): + if self._info_raises: + raise Exception("Unknown index name") + return self._info_result + + async def dropindex(self, delete_documents=False): + if self._drop_raises: + raise Exception("cannot drop") + self.drop_calls.append({"delete_documents": delete_documents}) + + async def create_index(self, schema, definition=None): + self.create_calls.append({"schema": schema, "definition": definition}) + + +class _FakeRedis: + def __init__(self, ft): + self._ft = ft + + def ft(self, index_name): + return self._ft + + +def _client_with(ft) -> RedisMemoryClient: + client = RedisMemoryClient.__new__(RedisMemoryClient) + client.client = _FakeRedis(ft) + return client + + +def _skill_schema_info(): + return {"attributes": [[b"identifier", b"$.description_embedding", + b"attribute", b"description_embedding", b"type", b"VECTOR"]]} + + +def _pre_skill_info(): + return {"attributes": [ + [b"identifier", b"$.summary", b"attribute", b"summary", b"type", b"TEXT"], + [b"identifier", b"$.summary_embedding", b"attribute", b"summary_embedding", b"type", b"VECTOR"], + [b"identifier", b"$.steps_embedding", b"attribute", b"steps_embedding", b"type", b"VECTOR"], + ]} + + +class TestIndexHasAttribute: + def test_finds_attribute_in_bytes_tokens(self): + assert RedisMemoryClient._index_has_attribute( + _skill_schema_info(), "description_embedding" + ) + + def test_missing_attribute(self): + assert not RedisMemoryClient._index_has_attribute( + _pre_skill_info(), "description_embedding" + ) + + def test_str_tokens_and_str_key(self): + info = {"attributes": [["attribute", "description_embedding"]]} + assert RedisMemoryClient._index_has_attribute(info, "description_embedding") + + def test_bytes_top_level_key(self): + info = {b"attributes": [[b"attribute", b"description_embedding"]]} + assert RedisMemoryClient._index_has_attribute(info, "description_embedding") + + def test_non_dict_info_is_false(self): + assert not RedisMemoryClient._index_has_attribute(None, "x") + assert not RedisMemoryClient._index_has_attribute([], "x") + + +@pytest.mark.asyncio +class TestCreateProceduralIndexUpgrade: + async def test_fresh_index_is_created(self): + ft = _FakeFT(info_raises=True) + await _client_with(ft)._create_procedural_index() + + assert len(ft.create_calls) == 1 + assert ft.drop_calls == [] + + async def test_current_schema_left_alone(self): + ft = _FakeFT(info_result=_skill_schema_info()) + await _client_with(ft)._create_procedural_index() + + assert ft.create_calls == [] + assert ft.drop_calls == [] + + async def test_stale_pre_skill_index_dropped_and_recreated(self): + ft = _FakeFT(info_result=_pre_skill_info()) + await _client_with(ft)._create_procedural_index() + + # The documents must be preserved: index-only drop. + assert ft.drop_calls == [{"delete_documents": False}] + assert len(ft.create_calls) == 1 + + async def test_drop_failure_does_not_recreate_over_live_index(self): + ft = _FakeFT(info_result=_pre_skill_info(), drop_raises=True) + await _client_with(ft)._create_procedural_index() + + assert ft.create_calls == [] diff --git a/tests/test_remote_client_request_params.py b/tests/test_remote_client_request_params.py new file mode 100644 index 000000000..cd84a0a1b --- /dev/null +++ b/tests/test_remote_client_request_params.py @@ -0,0 +1,191 @@ +from types import SimpleNamespace + +import pytest + +from mirix.client.remote_client import MirixClient + + +def _messages(): + return [ + {"role": "user", "content": [{"type": "text", "text": "remember this"}]}, + {"role": "assistant", "content": [{"type": "text", "text": "noted"}]}, + ] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("async_add", "endpoint"), + [(True, "/memory/add"), (False, "/memory/add_sync")], +) +async def test_add_sends_top_level_session_id(monkeypatch, async_add, endpoint): + client = MirixClient(api_key="test-key", base_url="http://test") + client._meta_agent = SimpleNamespace(id="agent-meta-1") + calls = [] + + async def fake_ensure_user_exists(user_id=None, headers=None): + return None + + async def fake_request(method, endpoint, json=None, params=None, headers=None): + calls.append( + { + "method": method, + "endpoint": endpoint, + "json": json, + "params": params, + } + ) + return {"success": True} + + monkeypatch.setattr(client, "_ensure_user_exists", fake_ensure_user_exists) + monkeypatch.setattr(client, "_request", fake_request) + + try: + await client.add( + user_id="user-1", + messages=_messages(), + session_id="sess-1", + async_add=async_add, + ) + finally: + await client.close() + + assert calls == [ + { + "method": "POST", + "endpoint": endpoint, + "json": { + "user_id": "user-1", + "meta_agent_id": "agent-meta-1", + "messages": _messages(), + "chaining": True, + "verbose": False, + "session_id": "sess-1", + }, + "params": None, + } + ] + + +@pytest.mark.asyncio +async def test_add_rejects_mismatched_session_id_filter_tag(monkeypatch): + client = MirixClient(api_key="test-key", base_url="http://test") + client._meta_agent = SimpleNamespace(id="agent-meta-1") + + async def fail_request(*args, **kwargs): + raise AssertionError("_request should not be called") + + monkeypatch.setattr(client, "_request", fail_request) + + try: + with pytest.raises(ValueError, match="must agree"): + await client.add( + user_id="user-1", + messages=_messages(), + session_id="sess-1", + filter_tags={"session_id": "sess-2"}, + ) + finally: + await client.close() + + +@pytest.mark.asyncio +async def test_auto_dream_sends_meta_agent_and_last_n_sessions(monkeypatch): + client = MirixClient(api_key="test-key", base_url="http://test") + calls = [] + + async def fake_ensure_user_exists(user_id=None, headers=None): + return None + + async def fake_request(method, endpoint, json=None, params=None, headers=None): + calls.append( + { + "method": method, + "endpoint": endpoint, + "json": json, + "params": params, + } + ) + return {"ok": True} + + monkeypatch.setattr(client, "_ensure_user_exists", fake_ensure_user_exists) + monkeypatch.setattr(client, "_request", fake_request) + + try: + await client.auto_dream( + user_id="user-1", + mode="procedural", + meta_agent_id="agent-meta-1", + last_n_sessions=3, + ) + finally: + await client.close() + + assert calls == [ + { + "method": "POST", + "endpoint": "/memory/auto_dream", + "json": { + "mode": "procedural", + "dry_run": False, + "meta_agent_id": "agent-meta-1", + "last_n_sessions": 3, + }, + "params": {"user_id": "user-1"}, + } + ] + + +def _search_client(monkeypatch): + """A MirixClient whose _request records params instead of hitting HTTP.""" + client = MirixClient(api_key="test-key", base_url="http://test") + client._meta_agent = SimpleNamespace(id="agent-meta-1") + calls = [] + + async def fake_ensure_user_exists(user_id=None, headers=None): + return None + + async def fake_request(method, endpoint, json=None, params=None, headers=None): + calls.append({"endpoint": endpoint, "params": params}) + return {"success": True, "results": [], "count": 0} + + monkeypatch.setattr(client, "_ensure_user_exists", fake_ensure_user_exists) + monkeypatch.setattr(client, "_request", fake_request) + return client, calls + + +@pytest.mark.asyncio +@pytest.mark.parametrize("method_name", ["search", "search_all_users"]) +async def test_search_omits_search_method_by_default(monkeypatch, method_name): + """Omitting search_method must leave it OUT of the query params entirely, + so the server resolves its per-type default (procedural -> hybrid, + everything else -> embedding). Sending a value would override that.""" + client, calls = _search_client(monkeypatch) + kwargs = {"query": "q", "memory_type": "procedural"} + if method_name == "search": + kwargs["user_id"] = "user-1" + + try: + await getattr(client, method_name)(**kwargs) + finally: + await client.close() + + assert len(calls) == 1 + assert "search_method" not in calls[0]["params"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("method_name", ["search", "search_all_users"]) +@pytest.mark.parametrize("explicit", ["embedding", "bm25", "hybrid"]) +async def test_search_sends_explicit_search_method(monkeypatch, method_name, explicit): + client, calls = _search_client(monkeypatch) + kwargs = {"query": "q", "memory_type": "procedural", "search_method": explicit} + if method_name == "search": + kwargs["user_id"] = "user-1" + + try: + await getattr(client, method_name)(**kwargs) + finally: + await client.close() + + assert len(calls) == 1 + assert calls[0]["params"]["search_method"] == explicit diff --git a/tests/test_schema_drift_guard.py b/tests/test_schema_drift_guard.py new file mode 100644 index 000000000..47d81f814 --- /dev/null +++ b/tests/test_schema_drift_guard.py @@ -0,0 +1,66 @@ +"""Tests for the startup schema-drift guard. + +create_all only creates missing tables โ€” it never ALTERs existing ones โ€” so an +operator who upgrades without running the manual migration scripts would get +UndefinedColumn errors on every request. ensure_tables_created() detects that +drift at startup and aborts with the exact scripts to run. +""" + +import pytest +from sqlalchemy import create_engine, text + +from mirix.server.server import _MIGRATION_HINTS, _find_missing_columns + + +@pytest.fixture() +def sqlite_conn(): + engine = create_engine("sqlite://") + with engine.connect() as conn: + yield conn + engine.dispose() + + +def test_no_drift_on_empty_database(sqlite_conn): + """Tables that don't exist yet are create_all's job, not drift.""" + assert _find_missing_columns(sqlite_conn) == {} + + +def test_detects_missing_session_id_on_legacy_messages_table(sqlite_conn): + """A pre-upgrade messages table (no session_id) must be reported.""" + sqlite_conn.execute(text("CREATE TABLE messages (id VARCHAR PRIMARY KEY)")) + sqlite_conn.commit() + + missing = _find_missing_columns(sqlite_conn) + + assert "messages" in missing + assert "session_id" in missing["messages"] + + +def test_detects_missing_skill_columns_on_legacy_procedural_table(sqlite_conn): + """A pre-skill procedural_memory table (summary/steps era) must be reported.""" + sqlite_conn.execute( + text( + "CREATE TABLE procedural_memory (" + "id VARCHAR PRIMARY KEY, entry_type VARCHAR, summary VARCHAR, steps JSON)" + ) + ) + sqlite_conn.commit() + + missing = _find_missing_columns(sqlite_conn) + + assert "procedural_memory" in missing + for column in ("name", "description", "instructions", "version"): + assert column in missing["procedural_memory"] + + +def test_migration_hints_cover_all_migration_scripted_tables(): + """Every table with a manual migration script has an actionable hint.""" + for table_name in ( + "messages", + "procedural_memory", + "conversation_message", + "skill_experience", + "agent_trigger_state", + ): + assert table_name in _MIGRATION_HINTS + assert "scripts/migrate_" in _MIGRATION_HINTS[table_name] diff --git a/tests/test_scope_read_consistency.py b/tests/test_scope_read_consistency.py new file mode 100644 index 000000000..12bc35c91 --- /dev/null +++ b/tests/test_scope_read_consistency.py @@ -0,0 +1,607 @@ +"""Hermetic-SQLite regression tests for scope-consistent memory reads. + +These lock in the read-side half of the procedural-scope fix: the SQLite +in-memory ``bm25`` fallback and the ``fuzzy_match`` path in every document-memory +manager must apply the SAME scope filtering as that manager's scoped base_query. +Before the fix these candidate loads selected on ``user_id`` only, silently +leaking NULL-scope rows to scoped readers (and masking the Postgres bug on dev +machines). + +The suite forces the SQLite branch by patching ``settings.mirix_pg_uri_no_default`` +to ``None`` on the shared settings singleton's class (the property every manager +consults via ``from mirix.settings import settings``), so the tests are +deterministic regardless of the runner's PG env. + +Template: the hermetic throwaway-SQLite pattern from tests/test_skill_experience.py +(module-scoped sqlite+aiosqlite engine, Base.metadata.create_all, an +``@asynccontextmanager`` session_maker matching ``db_context``'s shape, per-manager +``session_maker`` override, and ``@pytest.mark.asyncio(loop_scope="module")``). +""" + +import datetime as dt +import uuid +from contextlib import asynccontextmanager +from datetime import datetime + +import pytest +import pytest_asyncio +from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine + +import mirix.orm # noqa: F401 -- register all ORM classes before Base is used +from mirix.orm.agent import Agent as AgentORM +from mirix.orm.base import Base +from mirix.orm.episodic_memory import EpisodicEvent +from mirix.orm.knowledge_vault import KnowledgeVaultItem +from mirix.orm.organization import Organization as OrganizationORM +from mirix.orm.procedural_memory import ProceduralMemoryItem +from mirix.orm.resource_memory import ResourceMemoryItem +from mirix.orm.semantic_memory import SemanticMemoryItem +from mirix.orm.user import User as UserORM +from mirix.schemas.agent import AgentState, AgentType +from mirix.schemas.client import Client as PydanticClient +from mirix.schemas.embedding_config import EmbeddingConfig +from mirix.schemas.llm_config import LLMConfig +from mirix.schemas.user import User as PydanticUser +from mirix.services.episodic_memory_manager import EpisodicMemoryManager +from mirix.services.knowledge_vault_manager import KnowledgeVaultManager +from mirix.services.procedural_memory_manager import ProceduralMemoryManager +from mirix.services.resource_memory_manager import ResourceMemoryManager +from mirix.services.semantic_memory_manager import SemanticMemoryManager +from mirix.settings import settings + +SCOPE = "test" + + +# --------------------------------------------------------------------------- # +# Force the SQLite in-memory bm25/fuzzy branch for the whole module. +# Every manager reads ``settings.mirix_pg_uri_no_default`` on the shared singleton +# to choose PG-native fulltext vs the SQLite fallback; returning None routes them +# all through the fallback (the exact path this fix hardens). +# --------------------------------------------------------------------------- # +@pytest.fixture(autouse=True) +def _force_sqlite_branch(monkeypatch): + monkeypatch.setattr( + type(settings), + "mirix_pg_uri_no_default", + property(lambda self: None), + ) + + +# --------------------------------------------------------------------------- # +# Hermetic sqlite engine + org/user/agent seed (module scoped). +# --------------------------------------------------------------------------- # +@pytest_asyncio.fixture(scope="module", loop_scope="module") +async def session_maker(tmp_path_factory): + db_path = tmp_path_factory.mktemp("scope_read") / "test.db" + engine = create_async_engine(f"sqlite+aiosqlite:///{db_path}") + async with engine.begin() as conn: + await conn.run_sync(Base.metadata.create_all) + + local = async_sessionmaker(engine, expire_on_commit=False) + + @asynccontextmanager + async def _ctx(): + async with local() as session: + try: + yield session + finally: + await session.close() + + yield _ctx + await engine.dispose() + + +@pytest_asyncio.fixture(scope="module", loop_scope="module") +async def seed(session_maker): + org_id = f"scope-org-{uuid.uuid4().hex[:8]}" + user_id = f"scope-user-{uuid.uuid4().hex[:8]}" + agent_id = f"agent-{uuid.uuid4()}" + async with session_maker() as session: + session.add(OrganizationORM(id=org_id, name=org_id)) + session.add( + UserORM( + id=user_id, + name=user_id, + organization_id=org_id, + status="active", + timezone="UTC", + ) + ) + session.add(AgentORM(id=agent_id, organization_id=org_id)) + await session.commit() + return {"org_id": org_id, "user_id": user_id, "agent_id": agent_id} + + +@pytest.fixture(scope="module") +def user(seed): + return PydanticUser( + id=seed["user_id"], + name="scope-user", + organization_id=seed["org_id"], + timezone="UTC", + ) + + +@pytest.fixture(scope="module") +def actor(seed): + return PydanticClient( + id=f"client-{uuid.uuid4().hex[:8]}", + organization_id=seed["org_id"], + name="scope-client", + write_scope=SCOPE, + read_scopes=[SCOPE], + ) + + +@pytest.fixture(scope="module") +def agent_state(): + return AgentState( + id="agent-scope", + name="scope-agent", + system="s", + agent_type=AgentType.procedural_memory_agent, + llm_config=LLMConfig( + model="x", model_endpoint_type="openai", context_window=8000 + ), + embedding_config=EmbeddingConfig( + embedding_endpoint_type="openai", embedding_model="x", embedding_dim=8 + ), + tools=[], + ) + + +# --------------------------------------------------------------------------- # +# Per-manager specs: how to build the two rows and how to call the list fn. +# Each spec seeds one row with filter_tags={"scope": SCOPE} and one with +# filter_tags=None, using a shared search token so bm25/fuzzy match both. +# --------------------------------------------------------------------------- # +_TOKEN = "alphaunique" + + +def _make_procedural_rows(seed, scoped_id, unscoped_id): + def _row(row_id, filter_tags, suffix): + return ProceduralMemoryItem( + id=row_id, + organization_id=seed["org_id"], + user_id=seed["user_id"], + agent_id=seed["agent_id"], + # Unique name per row: (org, user, name) is a unique constraint, so + # the bm25 and fuzzy specs must not collide on reused suffixes. + name=f"skill-{row_id}", + entry_type="workflow", + description=f"{_TOKEN} deploy workflow {suffix}", + instructions="do the thing", + filter_tags=filter_tags, + ) + + return [ + _row(scoped_id, {"scope": SCOPE}, "scoped"), + _row(unscoped_id, None, "unscoped"), + ] + + +def _make_episodic_rows(seed, scoped_id, unscoped_id): + def _row(row_id, filter_tags, suffix): + return EpisodicEvent( + id=row_id, + organization_id=seed["org_id"], + user_id=seed["user_id"], + agent_id=seed["agent_id"], + occurred_at=datetime.now(dt.timezone.utc), + actor="user", + event_type="note", + summary=f"{_TOKEN} event {suffix}", + details="details here", + filter_tags=filter_tags, + ) + + return [ + _row(scoped_id, {"scope": SCOPE}, "scoped"), + _row(unscoped_id, None, "unscoped"), + ] + + +def _make_semantic_rows(seed, scoped_id, unscoped_id): + def _row(row_id, filter_tags, suffix): + return SemanticMemoryItem( + id=row_id, + organization_id=seed["org_id"], + user_id=seed["user_id"], + agent_id=seed["agent_id"], + name=f"{_TOKEN} concept {suffix}", + summary="a summary", + details="the details", + source="manual", + filter_tags=filter_tags, + ) + + return [ + _row(scoped_id, {"scope": SCOPE}, "scoped"), + _row(unscoped_id, None, "unscoped"), + ] + + +def _make_resource_rows(seed, scoped_id, unscoped_id): + def _row(row_id, filter_tags, suffix): + return ResourceMemoryItem( + id=row_id, + organization_id=seed["org_id"], + user_id=seed["user_id"], + agent_id=seed["agent_id"], + title=f"resource {suffix}", + summary="a summary", + content=f"{_TOKEN} resource body {suffix}", + resource_type="doc", + filter_tags=filter_tags, + ) + + return [ + _row(scoped_id, {"scope": SCOPE}, "scoped"), + _row(unscoped_id, None, "unscoped"), + ] + + +def _make_knowledge_rows(seed, scoped_id, unscoped_id): + def _row(row_id, filter_tags, suffix): + return KnowledgeVaultItem( + id=row_id, + organization_id=seed["org_id"], + user_id=seed["user_id"], + agent_id=seed["agent_id"], + entry_type="credential", + source="manual", + sensitivity="low", + secret_value="x", + caption=f"{_TOKEN} secret note {suffix}", + filter_tags=filter_tags, + ) + + return [ + _row(scoped_id, {"scope": SCOPE}, "scoped"), + _row(unscoped_id, None, "unscoped"), + ] + + +def _procedural_call(session_maker): + mgr = ProceduralMemoryManager() + mgr.session_maker = session_maker + + async def _call(agent_state, user, method, scopes): + return await mgr.list_procedures( + agent_state=agent_state, + user=user, + query=_TOKEN, + search_field="description", + search_method=method, + scopes=scopes, + use_cache=False, + ) + + return _call + + +def _episodic_call(session_maker): + mgr = EpisodicMemoryManager() + mgr.session_maker = session_maker + + async def _call(agent_state, user, method, scopes): + return await mgr.list_episodic_memory( + agent_state=agent_state, + user=user, + query=_TOKEN, + search_field="summary", + search_method=method, + scopes=scopes, + use_cache=False, + ) + + return _call + + +def _semantic_call(session_maker): + mgr = SemanticMemoryManager() + mgr.session_maker = session_maker + + async def _call(agent_state, user, method, scopes): + return await mgr.list_semantic_items( + agent_state=agent_state, + user=user, + query=_TOKEN, + search_field="name", + search_method=method, + scopes=scopes, + use_cache=False, + ) + + return _call + + +def _resource_call(session_maker): + mgr = ResourceMemoryManager() + mgr.session_maker = session_maker + + async def _call(agent_state, user, method, scopes): + return await mgr.list_resources( + agent_state=agent_state, + user=user, + query=_TOKEN, + search_field="content", + search_method=method, + scopes=scopes, + use_cache=False, + ) + + return _call + + +def _knowledge_call(session_maker): + mgr = KnowledgeVaultManager() + mgr.session_maker = session_maker + + async def _call(agent_state, user, method, scopes): + return await mgr.list_knowledge( + agent_state=agent_state, + user=user, + query=_TOKEN, + search_field="caption", + search_method=method, + scopes=scopes, + use_cache=False, + ) + + return _call + + +# (id, make_rows, make_call, supports_fuzzy) +_SPECS = [ + ("procedural", _make_procedural_rows, _procedural_call, True), + ("episodic", _make_episodic_rows, _episodic_call, True), + ("semantic", _make_semantic_rows, _semantic_call, True), + ("resource", _make_resource_rows, _resource_call, False), + ("knowledge_vault", _make_knowledge_rows, _knowledge_call, True), +] + + +async def _seed_rows(session_maker, rows): + async with session_maker() as session: + for row in rows: + session.add(row) + await session.commit() + + +@pytest.mark.asyncio(loop_scope="module") +@pytest.mark.parametrize( + "manager_id,make_rows,make_call,supports_fuzzy", + _SPECS, + ids=[s[0] for s in _SPECS], +) +async def test_bm25_scope_filtering( + session_maker, + seed, + user, + agent_state, + manager_id, + make_rows, + make_call, + supports_fuzzy, +): + scoped_id = f"{manager_id}-scoped-{uuid.uuid4().hex[:6]}" + unscoped_id = f"{manager_id}-unscoped-{uuid.uuid4().hex[:6]}" + await _seed_rows(session_maker, make_rows(seed, scoped_id, unscoped_id)) + call = make_call(session_maker) + + # scopes=[SCOPE] -> only the scoped row is visible (NULL scope hidden). + got = await call(agent_state, user, "bm25", [SCOPE]) + ids = {r.id for r in got} + assert scoped_id in ids + assert unscoped_id not in ids, ( + f"{manager_id} bm25 SQLite fallback leaked a NULL-scope row to a scoped reader" + ) + + # scopes=None -> unscoped read returns both rows. + got = await call(agent_state, user, "bm25", None) + ids = {r.id for r in got} + assert {scoped_id, unscoped_id} <= ids + + # scopes=[] -> empty-scope read returns nothing (WHERE 1=0). + got = await call(agent_state, user, "bm25", []) + assert got == [] + + +@pytest.mark.asyncio(loop_scope="module") +@pytest.mark.parametrize( + "manager_id,make_rows,make_call,supports_fuzzy", + [s for s in _SPECS if s[3]], + ids=[s[0] for s in _SPECS if s[3]], +) +async def test_fuzzy_scope_filtering( + session_maker, + seed, + user, + agent_state, + manager_id, + make_rows, + make_call, + supports_fuzzy, +): + scoped_id = f"{manager_id}-fz-scoped-{uuid.uuid4().hex[:6]}" + unscoped_id = f"{manager_id}-fz-unscoped-{uuid.uuid4().hex[:6]}" + await _seed_rows(session_maker, make_rows(seed, scoped_id, unscoped_id)) + call = make_call(session_maker) + + got = await call(agent_state, user, "fuzzy_match", [SCOPE]) + ids = {r.id for r in got} + assert scoped_id in ids + assert unscoped_id not in ids, ( + f"{manager_id} fuzzy_match candidate load leaked a NULL-scope row" + ) + + got = await call(agent_state, user, "fuzzy_match", None) + ids = {r.id for r in got} + assert {scoped_id, unscoped_id} <= ids + + got = await call(agent_state, user, "fuzzy_match", []) + assert got == [] + + +# --------------------------------------------------------------------------- # +# Write-side round trip: a stamped skill is retrievable by a scoped reader. +# --------------------------------------------------------------------------- # +@pytest.mark.asyncio(loop_scope="module") +async def test_procedural_roundtrip_scoped_read( + session_maker, seed, user, actor, agent_state, monkeypatch +): + import mirix.services.procedural_memory_manager as pmm + from mirix.orm.client import Client as ClientORM + + # Avoid real embedding calls in the hermetic test. + monkeypatch.setattr(pmm, "BUILD_EMBEDDINGS_FOR_MEMORY", False) + + # insert_procedure stamps client_id + audit FKs (_created_by_id) from the + # actor, so the owning clients row must exist for the FK constraint. + async with session_maker() as session: + session.add( + ClientORM( + id=actor.id, + organization_id=seed["org_id"], + name="scope-client", + status="active", + write_scope=SCOPE, + read_scopes=[SCOPE], + ) + ) + await session.commit() + + mgr = pmm.ProceduralMemoryManager() + mgr.session_maker = session_maker + + inserted = await mgr.insert_procedure( + agent_state=agent_state, + agent_id=seed["agent_id"], + name=f"roundtrip-{uuid.uuid4().hex[:6]}", + description=f"{_TOKEN} roundtrip skill", + instructions="steps", + entry_type="workflow", + actor=actor, + organization_id=seed["org_id"], + filter_tags={"scope": SCOPE}, + use_cache=False, + user_id=user.id, + ) + assert inserted.filter_tags == {"scope": SCOPE} + + got = await mgr.list_procedures( + agent_state=agent_state, + user=user, + query=_TOKEN, + search_field="description", + search_method="bm25", + scopes=[SCOPE], + use_cache=False, + ) + assert inserted.id in {r.id for r in got}, ( + "a scope-stamped skill must be retrievable by a scoped reader" + ) + + +# --------------------------------------------------------------------------- # +# Unit: scope_filter_tags(actor) +# --------------------------------------------------------------------------- # +class TestScopeFilterTags: + def test_with_write_scope(self): + from mirix.services.skill_experience_curator import scope_filter_tags + + actor = type("A", (), {"write_scope": "admin"})() + assert scope_filter_tags(actor) == {"scope": "admin"} + + def test_without_write_scope(self): + from mirix.services.skill_experience_curator import scope_filter_tags + + actor = type("A", (), {"write_scope": None})() + assert scope_filter_tags(actor) is None + + def test_missing_attr(self): + from mirix.services.skill_experience_curator import scope_filter_tags + + actor = type("A", (), {})() + assert scope_filter_tags(actor) is None + + +# --------------------------------------------------------------------------- # +# Migration text-assertion (TestMigrationSql convention). +# --------------------------------------------------------------------------- # +class TestBackfillMigrationSql: + from pathlib import Path + + SQL_PATH = Path("scripts/migrate_backfill_procedural_scope.sql") + + def _sql(self) -> str: + return self.SQL_PATH.read_text() + + def test_contains_required_constructs(self): + sql = self._sql() + for token in ( + "BEGIN;", + "COMMIT;", + "jsonb_set", + # Non-object filter_tags (SQL NULL / JSON scalar 'null') must be + # normalized before jsonb_set, which errors on scalars. + "jsonb_typeof", + "FROM clients", + "write_scope IS NOT NULL", + "jsonb_exists", + "'{scope}'", + "RAISE NOTICE", + ): + assert token in sql, f"backfill migration missing {token!r}" + + +# --------------------------------------------------------------------------- # +# Production wiring: run_experience_evolution must construct the procedural +# agent WITH the actor-derived scope stamp. The round-trip test above supplies +# filter_tags by hand, so without this spy the stamping line in +# skill_experience_curator could be deleted and the suite would stay green. +# --------------------------------------------------------------------------- # +class _StampSentinel(Exception): + """Raised by the spy to stop the run right after agent construction.""" + + +@pytest.mark.asyncio +async def test_run_experience_evolution_stamps_actor_scope( + monkeypatch, actor, user, agent_state +): + import mirix.agent as agent_pkg + import mirix.server.server as server_module + import mirix.services.skill_experience_curator as curator + + captured: dict = {} + + class _SpyProceduralAgent: + def __init__(self, **kwargs): + captured.update(kwargs) + raise _StampSentinel + + class _StubAgentManager: + async def update_agent_tools_and_system_prompts(self, *a, **k): + raise RuntimeError("refresh skipped in spy test") + + class _StubServer: + agent_manager = _StubAgentManager() + + def default_interface_factory(self): + return None + + async def _fake_resolve(server, resolve_actor, meta_agent_state): + return agent_state + + monkeypatch.setattr(agent_pkg, "ProceduralMemoryAgent", _SpyProceduralAgent) + monkeypatch.setattr(server_module, "get_server", lambda: _StubServer()) + monkeypatch.setattr(curator, "_resolve_procedural_agent_state", _fake_resolve) + + with pytest.raises(_StampSentinel): + await curator.run_experience_evolution( + user=user, actor=actor, meta_agent_state=agent_state + ) + + assert captured["filter_tags"] == {"scope": SCOPE} + assert captured["actor"] is actor diff --git a/tests/test_session_experience_distillation.py b/tests/test_session_experience_distillation.py new file mode 100644 index 000000000..46fc9414e --- /dev/null +++ b/tests/test_session_experience_distillation.py @@ -0,0 +1,746 @@ +"""Tests for general per-session experience distillation (DB-free). + +Layered for the general Experience Distiller: no benchmark oracle and no +success/failure label. + +1. JSON-array parse robustness โ€” bare array, fenced ```json, prose-wrapped, + single-object fallback, empty `[]`, garbage -> []. +2. Parsed-experience -> SkillExperience mapping (_persist_experiences via a stub + manager, no DB): importance/credibility clamping, evidence normalization, + bad-type / empty-title skips, length caps, status='pending'. +3. Prioritization ordering invariant โ€” importance*credibility is the priority + the curator payload is built around. +4. Single session yields MULTIPLE mixed worth_learning + worth_avoiding from one + fixed fake LLM reply (no external label needed). +5. Benchmark-decoupling assertion โ€” the distiller prompt/flow carries no + benchmark vocabulary and needs no external success/failure oracle. + +The LLM client is a hand-rolled fake (no network). No Postgres: the experience +manager is replaced with an in-memory recorder so the per-session persistence +path is exercised without a DB. +""" + +from __future__ import annotations + +import inspect +import json +from pathlib import Path +from types import SimpleNamespace + +import pytest + +from mirix.helpers.json_parsing import parse_distiller_json_array +from mirix.schemas.skill_experience import ( + SKILL_EXPERIENCE_MAX_CONTENT_LEN, + SKILL_EXPERIENCE_MAX_TITLE_LEN, +) +from mirix.services.session_experience_distiller import SessionExperienceDistiller + +# ============================ JSON-array parse robustness =================== + + +class TestParseDistillerJsonArray: + def test_bare_array(self): + raw = '[{"experience_type":"worth_learning","title":"t"}]' + out = parse_distiller_json_array(raw) + assert len(out) == 1 + assert out[0]["title"] == "t" + + def test_empty_array_is_nothing_worth_remembering(self): + assert parse_distiller_json_array("[]") == [] + + def test_fenced_json_block(self): + raw = '```json\n[{"experience_type":"worth_avoiding","title":"x"}]\n```' + out = parse_distiller_json_array(raw) + assert len(out) == 1 + assert out[0]["experience_type"] == "worth_avoiding" + + def test_plain_fence(self): + raw = '```\n[{"title":"a"},{"title":"b"}]\n```' + out = parse_distiller_json_array(raw) + assert [o["title"] for o in out] == ["a", "b"] + + def test_prose_wrapped_array(self): + raw = 'Here are the experiences:\n[{"title":"only"}]\nThanks!' + out = parse_distiller_json_array(raw) + assert len(out) == 1 and out[0]["title"] == "only" + + def test_single_object_fallback_is_wrapped(self): + # A model that emits ONE experience without the enclosing array must not + # be silently dropped. + raw = '{"experience_type":"worth_learning","title":"solo"}' + out = parse_distiller_json_array(raw) + assert len(out) == 1 and out[0]["title"] == "solo" + + def test_non_dict_elements_dropped(self): + raw = '[{"title":"keep"}, 7, "noise", null]' + out = parse_distiller_json_array(raw) + assert out == [{"title": "keep"}] + + def test_garbage_yields_empty(self): + assert parse_distiller_json_array("not json at all") == [] + assert parse_distiller_json_array("") == [] + assert parse_distiller_json_array(None) == [] + + +# ===================== Parsed-experience -> SkillExperience mapping ========= + + +class _RecordingManager: + """In-memory stand-in for SkillExperienceManager.create_experience. + + Validates exactly like the real manager (via SkillExperienceCreate) so the + clamp / enum / length contracts are exercised, but never touches a DB. + """ + + def __init__(self): + self.created = [] + + async def create_experience(self, **kwargs): + from mirix.schemas.skill_experience import SkillExperienceCreate + + # Mirror the real manager: created_by_id feeds the _created_by_id audit + # column (client attribution for erasure), not SkillExperienceCreate. + created_by_id = kwargs.pop("created_by_id", None) + validated = SkillExperienceCreate(**kwargs) + rec = SimpleNamespace(**validated.model_dump(), created_by_id=created_by_id) + self.created.append(rec) + return rec + + +def _distiller_with(manager): + # llm_config=None / llm_client=None: we never call the LLM in these mapping + # tests, we drive _persist_experiences directly. + return SessionExperienceDistiller(experience_manager=manager) + + +def _meta_user_actor(): + meta = SimpleNamespace(id="agent-meta-1") + user = SimpleNamespace(id="user-1") + actor = SimpleNamespace(organization_id="org-1") + return meta, user, actor + + +@pytest.mark.asyncio +class TestPersistMapping: + async def test_basic_mapping_status_pending(self): + mgr = _RecordingManager() + d = _distiller_with(mgr) + meta, user, actor = _meta_user_actor() + parsed = [ + { + "experience_type": "worth_learning", + "title": "Cache the resolved path", + "content": "When repeatedly resolving X, cache it.", + "importance": 0.8, + "credibility": 0.9, + "evidence": { + "quote": "great, that worked", + "signal_type": "user_confirmation", + }, + } + ] + out = await d._persist_experiences( + meta_agent_state=meta, + user=user, + actor=actor, + session_id="sess-1", + parsed=parsed, + ) + assert len(out) == 1 + rec = mgr.created[0] + assert rec.experience_type == "worth_learning" + assert rec.status == "pending" + assert rec.session_id == "sess-1" + assert rec.agent_id == "agent-meta-1" + assert rec.user_id == "user-1" + assert rec.organization_id == "org-1" + ev = json.loads(rec.evidence) + assert ev["signal_type"] == "user_confirmation" + assert ev["quote"] == "great, that worked" + + async def test_clamps_importance_and_credibility(self): + mgr = _RecordingManager() + d = _distiller_with(mgr) + meta, user, actor = _meta_user_actor() + parsed = [ + { + "experience_type": "worth_avoiding", + "title": "t1", + "importance": 5.0, + "credibility": -3.0, + }, + { + "experience_type": "worth_learning", + "title": "t2", + "importance": "garbage", + "credibility": None, + }, + ] + await d._persist_experiences( + meta_agent_state=meta, + user=user, + actor=actor, + session_id="s", + parsed=parsed, + ) + assert mgr.created[0].importance == 1.0 + assert mgr.created[0].credibility == 0.0 + assert mgr.created[1].importance == 0.0 # garbage -> 0.0 + assert mgr.created[1].credibility == 0.0 + + async def test_bad_type_is_skipped(self): + mgr = _RecordingManager() + d = _distiller_with(mgr) + meta, user, actor = _meta_user_actor() + parsed = [ + {"experience_type": "partial", "title": "bad"}, # bad enum + {"experience_type": "worth_learning", "title": "good"}, # kept + {"experience_type": None, "title": "alsobad"}, # missing + ] + out = await d._persist_experiences( + meta_agent_state=meta, + user=user, + actor=actor, + session_id="s", + parsed=parsed, + ) + assert len(out) == 1 + assert mgr.created[0].title == "good" + + async def test_empty_title_is_skipped(self): + mgr = _RecordingManager() + d = _distiller_with(mgr) + meta, user, actor = _meta_user_actor() + parsed = [ + {"experience_type": "worth_learning", "title": " "}, + {"experience_type": "worth_learning", "title": ""}, + ] + out = await d._persist_experiences( + meta_agent_state=meta, + user=user, + actor=actor, + session_id="s", + parsed=parsed, + ) + assert out == [] + assert mgr.created == [] + + async def test_length_caps_enforced(self): + mgr = _RecordingManager() + d = _distiller_with(mgr) + meta, user, actor = _meta_user_actor() + parsed = [ + { + "experience_type": "worth_avoiding", + "title": "T" * (SKILL_EXPERIENCE_MAX_TITLE_LEN + 50), + "content": "C" * (SKILL_EXPERIENCE_MAX_CONTENT_LEN + 100), + } + ] + out = await d._persist_experiences( + meta_agent_state=meta, + user=user, + actor=actor, + session_id="s", + parsed=parsed, + ) + assert len(out) == 1 + assert len(mgr.created[0].title) <= SKILL_EXPERIENCE_MAX_TITLE_LEN + assert len(mgr.created[0].content) <= SKILL_EXPERIENCE_MAX_CONTENT_LEN + + async def test_evidence_normalized_when_missing_or_bad(self): + mgr = _RecordingManager() + d = _distiller_with(mgr) + meta, user, actor = _meta_user_actor() + parsed = [ + {"experience_type": "worth_learning", "title": "a"}, # no evidence + { + "experience_type": "worth_learning", + "title": "b", + "evidence": {"quote": "q", "signal_type": "bogus_signal"}, + }, + { + "experience_type": "worth_learning", + "title": "c", + "evidence": "a raw string", + }, + ] + await d._persist_experiences( + meta_agent_state=meta, + user=user, + actor=actor, + session_id="s", + parsed=parsed, + ) + for rec in mgr.created: + ev = json.loads(rec.evidence) + assert "quote" in ev and "signal_type" in ev + # missing -> inferred + assert json.loads(mgr.created[0].evidence)["signal_type"] == "inferred" + # bad signal_type -> inferred + assert json.loads(mgr.created[1].evidence)["signal_type"] == "inferred" + # raw string -> inferred, quote preserved + ev2 = json.loads(mgr.created[2].evidence) + assert ev2["signal_type"] == "inferred" + assert ev2["quote"] == "a raw string" + + async def test_one_bad_row_does_not_drop_the_rest(self): + # A manager that raises on a specific title must not abort the batch. + class _PartlyFailing(_RecordingManager): + async def create_experience(self, **kwargs): + if kwargs.get("title") == "boom": + raise RuntimeError("db blew up") + return await super().create_experience(**kwargs) + + mgr = _PartlyFailing() + d = _distiller_with(mgr) + meta, user, actor = _meta_user_actor() + parsed = [ + {"experience_type": "worth_learning", "title": "boom"}, + {"experience_type": "worth_learning", "title": "survivor"}, + ] + out = await d._persist_experiences( + meta_agent_state=meta, + user=user, + actor=actor, + session_id="s", + parsed=parsed, + ) + assert [r.title for r in out] == ["survivor"] + + +# ============== Single session -> MULTIPLE mixed experiences (fake LLM) ===== + + +class _FakeMessage: + def __init__(self, content): + self.content = content + + +class _FakeChoice: + def __init__(self, content): + self.message = _FakeMessage(content) + + +class _FakeResponse: + def __init__(self, content): + self.choices = [_FakeChoice(content)] + + +class _FakeLLMClient: + """Returns a FIXED reply regardless of input โ€” no network.""" + + def __init__(self, reply: str): + self._reply = reply + self.calls = 0 + + async def send_llm_request(self, *, messages): + self.calls += 1 + return _FakeResponse(self._reply) + + +MIXED_REPLY = json.dumps( + [ + { + "experience_type": "worth_learning", + "title": "Batch independent tool calls", + "content": "When calls are independent, issue them together.", + "importance": 0.7, + "credibility": 0.85, + "evidence": { + "quote": "perfect, much faster", + "signal_type": "user_confirmation", + }, + }, + { + "experience_type": "worth_avoiding", + "title": "Do not assume the column exists", + "content": "A query failed on a missing column; check schema first.", + "importance": 0.9, + "credibility": 0.95, + "evidence": {"quote": "no such column: foo", "signal_type": "tool_error"}, + }, + { + "experience_type": "worth_avoiding", + "title": "Avoid wholesale rewrites", + "content": "User said the rewrite lost context; prefer deltas.", + "importance": 0.6, + "credibility": 0.8, + "evidence": { + "quote": "this part isn't good enough", + "signal_type": "user_critique", + }, + }, + ] +) + + +@pytest.mark.asyncio +class TestSingleSessionMultipleMixed: + async def test_one_session_yields_multiple_mixed_kinds(self): + mgr = _RecordingManager() + fake = _FakeLLMClient(MIXED_REPLY) + d = SessionExperienceDistiller(llm_client=fake, experience_manager=mgr) + meta, user, actor = _meta_user_actor() + + parsed = await d._call_llm( + agent_id=meta.id, + session_id="sess-x", + transcript="user: ...\nassistant: ...", + skills_block="(none)", + ) + out = await d._persist_experiences( + meta_agent_state=meta, + user=user, + actor=actor, + session_id="sess-x", + parsed=parsed, + ) + assert fake.calls == 1 + types = sorted(r.experience_type for r in out) + # ONE session -> MULTIPLE experiences, a MIX of both kinds. + assert types == ["worth_avoiding", "worth_avoiding", "worth_learning"] + # No external success/failure label was needed or used. + assert len(out) == 3 + + +# ===================== Prioritization ordering (priority = imp*cred) ======== + + +class TestPriorityOrdering: + """The importance*credibility DESC ordering is the load-bearing priority. + + The DB-backed proof that the MANAGER emits this ordering lives in + test_skill_experience.py::TestCreateAndList::test_list_ordered_by_priority_desc + (against the real SQL `ORDER BY importance*credibility DESC`). Here we assert + the downstream guarantee the CURATOR relies on: build_experience_payload + PRESERVES the order it is handed (so the highest-priority experience leads + the prompt). This is non-tautological โ€” it pins real curator behavior. + """ + + def test_payload_preserves_handed_priority_order(self): + import json as _json + + from mirix.services.skill_experience_curator import build_experience_payload + + # Hand the payload builder experiences already in priority order (as the + # manager would). The first block in the rendered prompt must be the + # highest-priority one, the last the lowest. The titles are deliberately + # chosen so the handed (priority) order is the REVERSE of lexical order + # (Zeta > Mu > Alpha by priority, but Alpha < Mu < Zeta lexically) โ€” so a + # builder that accidentally sorted by title would FLIP them and fail. + ordered = [ + SimpleNamespace( + experience_type="worth_avoiding", + title="Zeta", # highest priority + content="c", + importance=0.8, + credibility=0.9, + evidence=_json.dumps({"quote": "", "signal_type": "inferred"}), + ), + SimpleNamespace( + experience_type="worth_learning", + title="Mu", # middle priority + content="c", + importance=0.5, + credibility=0.6, + evidence=_json.dumps({"quote": "", "signal_type": "inferred"}), + ), + SimpleNamespace( + experience_type="worth_avoiding", + title="Alpha", # lowest priority + content="c", + importance=0.9, + credibility=0.1, + evidence=_json.dumps({"quote": "", "signal_type": "inferred"}), + ), + ] + payload = build_experience_payload(ordered) + assert payload.index("Zeta") < payload.index("Mu") < payload.index("Alpha") + + +# ===================== Benchmark-decoupling assertion ======================= + + +class TestNoBenchmarkVocabulary: + PROMPT = Path("mirix/prompts/system/base/auto_dream_agent/procedural.txt") + + def _prompt_text(self) -> str: + return self.PROMPT.read_text(encoding="utf-8") + + def test_prompt_has_no_metaclaw_vocabulary(self): + text = self._prompt_text().lower() + for banned in [ + "metaclaw", + "previous feedback", + "\\bbox", + "bbox", + "round_id", + "round_index", + "oracle", + "quality_score", + "openclaw", + ]: + assert banned not in text, f"benchmark vocab leaked into prompt: {banned!r}" + + def test_prompt_states_no_external_grader(self): + text = self._prompt_text().lower() + assert "no external grader" in text or "no external" in text + + def test_prompt_does_not_label_session_success_or_failure(self): + text = self._prompt_text().lower() + # The prompt must instruct NOT to label the whole session success/failure. + assert "do not label" in text and "success or a failure" in text + + def test_prompt_uses_our_experience_kinds(self): + text = self._prompt_text() + assert "worth_learning" in text + assert "worth_avoiding" in text + + def test_distiller_module_has_no_metaclaw_terms(self): + src = inspect.getsource(SessionExperienceDistiller) + low = src.lower() + for banned in ["metaclaw", "round_index", "quality_score", "oracle"]: + assert banned not in low, ( + f"benchmark term leaked into distiller: {banned!r}" + ) + + +# ===== Source isolation โ€” the distiller reads ONLY the Conversation Message Store = +# The scaffolding-filter heuristic (_is_mirix_scaffolding) is DELETED. Correctness +# now comes from STRUCTURE: the distiller's transcript source is the dedicated +# Conversation Message Store โ€” external turns with their REAL user/assistant roles +# โ€” never the meta agent's own `messages` thread. So MIRIX's memory-management +# scaffolding (the "[System Message] As the meta memory managerโ€ฆ" instruction, its +# trigger_memory_update tool calls + results, the continue_chaining control +# replies) can NEVER reach the distiller: those rows simply do not live in this +# store. These tests pin that structural isolation rather than a string heuristic. + + +def _conv_turn(role, content, *, session_id="sess-1"): + """A ConversationMessage-shaped row (real role + plain-text content), the + exact shape ConversationMessageManager.list_turns_for_session returns.""" + return SimpleNamespace( + role=role, content=content, session_id=session_id, distilled_at=None + ) + + +class _RecordingConversationManager: + """In-memory stand-in for ConversationMessageManager. + + Records every call so a test can prove the distiller pulled transcripts from + THIS store (and, by construction, never from MessageManager). Returns the + canned turns for a session; everything else is a no-op the distiller needs. + """ + + def __init__(self, turns_by_session=None, sealed=None): + self._turns_by_session = turns_by_session or {} + self._sealed = sealed or [] + self.list_turns_calls = [] + self.sealed_calls = [] + + async def list_turns_for_session( + self, *, session_id, user_id, organization_id, actor + ): + self.list_turns_calls.append(session_id) + return list(self._turns_by_session.get(session_id, [])) + + async def list_sealed_undistilled_sessions( + self, *, user_id, organization_id, actor, limit + ): + self.sealed_calls.append({"limit": limit}) + return list(self._sealed[:limit]) + + +class TestRenderTranscriptFromConversationStore: + """_render_transcript operates on ConversationMessage rows with REAL roles + and plain-text content โ€” no scaffolding shapes, no tool-call flattening.""" + + def test_real_roles_are_preserved(self): + turns = [ + _conv_turn("user", "Q: which HTTP status means 'Not Found'?"), + _conv_turn("assistant", "404"), + ] + out = SessionExperienceDistiller._render_transcript(turns) + lines = [ln for ln in out.splitlines() if ln.strip()] + # Exactly the two real turns, with their real roles โ€” not [USER]/[ASSISTANT]. + assert lines == [ + "user: Q: which HTTP status means 'Not Found'?", + "assistant: 404", + ] + + def test_no_scaffolding_filtering_remains(self): + # The store never holds scaffolding, so the distiller does NOT strip it. + # A row whose content merely RESEMBLES old scaffolding is rendered verbatim + # (proving the heuristic is gone โ€” the source is trusted by structure). + turns = [ + _conv_turn( + "user", + "[System Message] As the meta memory manager, analyze the content.", + ), + ] + out = SessionExperienceDistiller._render_transcript(turns) + assert "As the meta memory manager" in out + + def test_empty_content_turn_is_dropped(self): + turns = [_conv_turn("user", ""), _conv_turn("assistant", "real answer")] + out = SessionExperienceDistiller._render_transcript(turns) + assert out == "assistant: real answer" + + def test_oversized_turn_is_capped_with_marker(self): + from mirix.services import session_experience_distiller as sed + + huge = "x" * (sed._MAX_MESSAGE_CHARS + 500) + out = SessionExperienceDistiller._render_transcript([_conv_turn("tool", huge)]) + line = out.splitlines()[0] + assert line.endswith("โ€ฆ[truncated]") + assert len(line) <= len("tool: ") + sed._MAX_MESSAGE_CHARS + len( + " โ€ฆ[truncated]" + ) + + def test_overlong_transcript_keeps_head_and_tail_elides_middle(self): + # Head (task framing) and tail (final verdict) must survive; only the + # middle is elided, and the result stays within the transcript budget. + from mirix.services import session_experience_distiller as sed + + per_turn = sed._MAX_MESSAGE_CHARS - 100 # below the per-turn cap + n = (sed._MAX_TRANSCRIPT_CHARS // per_turn) + 3 + turns = [ + _conv_turn("user", f"turn-{i:04d} " + "y" * per_turn) for i in range(n) + ] + out = SessionExperienceDistiller._render_transcript(turns) + assert "โ€ฆ[elided middle of session]โ€ฆ" in out + assert "turn-0000" in out + assert f"turn-{n - 1:04d}" in out + assert len(out) <= sed._MAX_TRANSCRIPT_CHARS + len( + "\nโ€ฆ[elided middle of session]โ€ฆ\n" + ) + + def test_transcript_within_budget_is_untouched(self): + out = SessionExperienceDistiller._render_transcript( + [_conv_turn("user", "short"), _conv_turn("assistant", "also short")] + ) + assert "โ€ฆ[elided middle of session]โ€ฆ" not in out + assert "โ€ฆ[truncated]" not in out + + def test_scaffolding_heuristic_is_deleted(self): + # The fragile string heuristic must be GONE โ€” correctness is structural. + assert not hasattr(SessionExperienceDistiller, "_is_mirix_scaffolding") + + def test_distiller_module_does_not_reference_scaffolding_terms(self): + src = inspect.getsource(SessionExperienceDistiller) + for banned in ( + "_is_mirix_scaffolding", + "trigger_memory_update", + "META_MEMORY_TOOLS", + ): + assert banned not in src, ( + f"dead scaffolding reference left behind: {banned!r}" + ) + + +@pytest.mark.asyncio +class TestDistillerReadsConversationStoreNotMetaMessages: + """The load-bearing isolation: the distiller's per-session transcript fetch + goes through the injected ConversationMessageManager, NEVER through the meta + agent's `messages` store (MessageManager).""" + + async def test_distill_one_pulls_turns_from_conversation_store(self, monkeypatch): + # Hard guard: if the distiller ever touches MessageManager, blow up. + import mirix.services.message_manager as mm + + def _boom(*a, **k): # pragma: no cover - only fires on regression + raise AssertionError( + "distiller must NOT read meta agent messages via MessageManager" + ) + + monkeypatch.setattr(mm.MessageManager, "list_messages_for_agent", _boom) + + conv = _RecordingConversationManager( + turns_by_session={ + "sess-x": [ + _conv_turn("user", "great, that worked", session_id="sess-x"), + _conv_turn("assistant", "glad to help", session_id="sess-x"), + ] + } + ) + exp_mgr = _RecordingManager() + fake = _FakeLLMClient(MIXED_REPLY) + d = SessionExperienceDistiller( + llm_client=fake, experience_manager=exp_mgr, conversation_manager=conv + ) + meta, user, actor = _meta_user_actor() + + out = await d._distill_one( + meta_agent_state=meta, + user=user, + actor=actor, + session_id="sess-x", + skills_block="(none)", + ) + # The transcript was fetched from the Conversation Message Storeโ€ฆ + assert conv.list_turns_calls == ["sess-x"] + # โ€ฆthe LLM was called once, and experiences were produced. + assert fake.calls == 1 + assert len(out) == 3 + + async def test_empty_session_in_store_skips_llm(self): + conv = _RecordingConversationManager(turns_by_session={"empty": []}) + fake = _FakeLLMClient(MIXED_REPLY) + d = SessionExperienceDistiller( + llm_client=fake, + experience_manager=_RecordingManager(), + conversation_manager=conv, + ) + meta, user, actor = _meta_user_actor() + out = await d._distill_one( + meta_agent_state=meta, + user=user, + actor=actor, + session_id="empty", + skills_block="(none)", + ) + assert out == [] + # No turns โ†’ no LLM call (the source store is the only thing consulted). + assert conv.list_turns_calls == ["empty"] + assert fake.calls == 0 + + async def test_enumerate_sealed_sessions_delegates_to_store(self): + conv = _RecordingConversationManager(sealed=["s1", "s2", "s3"]) + d = SessionExperienceDistiller(conversation_manager=conv) + _, user, actor = _meta_user_actor() + out = await d.enumerate_sealed_sessions( + user_id=user.id, + organization_id=actor.organization_id, + actor=actor, + limit=5, + ) + # Enumeration comes from the store's sealed-session query, oldest-first. + assert out == ["s1", "s2", "s3"] + assert conv.sealed_calls == [{"limit": 5}] + + +class TestDistillerSourceIsConversationStore: + """Static guard that the distiller's source-of-truth is the Conversation + Message Store and that it no longer enumerates the meta agent's messages.""" + + def test_distiller_fetches_via_conversation_manager(self): + src = inspect.getsource(SessionExperienceDistiller._distill_one) + assert "list_turns_for_session" in src + # Must NOT read the meta agent's message thread for distillation. + assert "list_messages_for_agent" not in src + + def test_distiller_constructor_injects_conversation_manager(self): + sig = inspect.signature(SessionExperienceDistiller.__init__) + assert "conversation_manager" in sig.parameters + + def test_distiller_does_not_import_message_manager(self): + # The distiller's source must not pull in the meta-agent MessageManager, + # nor query the meta `messages` store โ€” that store is structurally out of + # reach. (`ConversationMessageManager` legitimately contains the substring + # "MessageManager", so we match the precise import / query forms instead.) + src = inspect.getsource(inspect.getmodule(SessionExperienceDistiller)) + assert "from mirix.services.message_manager import" not in src + assert "list_messages_for_agent" not in src + # โ€ฆand it DOES use the Conversation Message Store. + assert "ConversationMessageManager" in src diff --git a/tests/test_session_id.py b/tests/test_session_id.py new file mode 100644 index 000000000..b3bc49fa8 --- /dev/null +++ b/tests/test_session_id.py @@ -0,0 +1,612 @@ +""" +Unit tests for the top-level `session_id` field on messages. + +Scope: +- Pydantic schema: MessageCreate / MessageUpdate / Message accept optional session_id. +- Validation: length/charset rules on session_id. +- Helper: prepare_input_message_create propagates session_id to the Message. +- REST schemas: SendMessageRequest / AddMemoryRequest accept session_id. +- Queue serialization: put_messages path sets session_id on proto, worker restores it. +- ORM column exists and is indexed. + +These are unit tests only; no DB, no running server. +""" +from __future__ import annotations + +import pytest + +from mirix.helpers.message_helpers import prepare_input_message_create +from mirix.schemas.enums import MessageRole +from mirix.schemas.message import Message, MessageCreate, MessageUpdate + + +# ----------------------------- Schema ------------------------------------ + + +class TestMessageCreateSessionId: + def test_accepts_session_id(self): + m = MessageCreate(role=MessageRole.user, content="hi", session_id="sess-abc") + assert m.session_id == "sess-abc" + + def test_session_id_optional(self): + m = MessageCreate(role=MessageRole.user, content="hi") + assert m.session_id is None + + def test_rejects_empty_string_session_id(self): + # Empty string is ambiguous; require None or a non-empty string. + with pytest.raises(ValueError): + MessageCreate(role=MessageRole.user, content="hi", session_id="") + + def test_rejects_too_long_session_id(self): + with pytest.raises(ValueError): + MessageCreate( + role=MessageRole.user, content="hi", session_id="a" * 65 + ) + + def test_rejects_invalid_chars(self): + with pytest.raises(ValueError): + MessageCreate( + role=MessageRole.user, content="hi", session_id="sess/abc" + ) + + def test_accepts_allowed_chars(self): + m = MessageCreate( + role=MessageRole.user, + content="hi", + session_id="sess_Abc-123", + ) + assert m.session_id == "sess_Abc-123" + + +class TestMessageSchemaSessionId: + def test_accepts_session_id(self): + m = Message( + agent_id="agent-1", + role=MessageRole.user, + session_id="sess-xyz", + ) + assert m.session_id == "sess-xyz" + + def test_defaults_to_none(self): + m = Message(agent_id="agent-1", role=MessageRole.user) + assert m.session_id is None + + +class TestMessageUpdateSessionId: + def test_accepts_session_id(self): + u = MessageUpdate(session_id="sess-upd") + assert u.session_id == "sess-upd" + + def test_defaults_to_none(self): + u = MessageUpdate() + assert u.session_id is None + + +# ----------------------------- Helper ------------------------------------ + + +class TestPrepareInputMessageCreate: + def test_propagates_session_id(self): + create = MessageCreate( + role=MessageRole.user, content="hi", session_id="sess-1" + ) + msg = prepare_input_message_create(create, agent_id="agent-1") + assert msg.session_id == "sess-1" + + def test_missing_session_id_is_none(self): + create = MessageCreate(role=MessageRole.user, content="hi") + msg = prepare_input_message_create(create, agent_id="agent-1") + assert msg.session_id is None + + +class TestDictToMessageSessionId: + """dict_to_message must FAITHFULLY pass through whatever session_id it is + given (the chat_agent path supplies one; non-chat callers pass None). The + decoupling is enforced at the CALLER (see TestAgentStepPropagation, which + pins that non-chat callers pass None) โ€” here we only pin the mechanism: the + kwarg round-trips and defaults to None when omitted.""" + + def test_passes_session_id_through_kwarg(self): + msg = Message.dict_to_message( + agent_id="agent-1", + openai_message_dict={"role": "user", "content": "hi"}, + session_id="sess-internal", + ) + assert msg.session_id == "sess-internal" + + def test_defaults_to_none(self): + msg = Message.dict_to_message( + agent_id="agent-1", + openai_message_dict={"role": "user", "content": "hi"}, + ) + assert msg.session_id is None + + def test_tool_role_message_carries_session_id(self): + msg = Message.dict_to_message( + agent_id="agent-1", + openai_message_dict={ + "role": "tool", + "content": "ok", + "tool_call_id": "tc-1", + "name": "do_stuff", + }, + session_id="sess-internal", + ) + assert msg.session_id == "sess-internal" + + +# ----------------------------- REST request schemas ---------------------- + + +class TestRestRequestSchemasSessionId: + def test_send_message_request_accepts_session_id(self): + from mirix.server.rest_api import SendMessageRequest + + req = SendMessageRequest( + message="hi", role="user", session_id="sess-req" + ) + assert req.session_id == "sess-req" + + def test_send_message_request_session_id_optional(self): + from mirix.server.rest_api import SendMessageRequest + + req = SendMessageRequest(message="hi", role="user") + assert req.session_id is None + + def test_add_memory_request_accepts_session_id(self): + from mirix.server.rest_api import AddMemoryRequest + + req = AddMemoryRequest( + meta_agent_id="meta-1", + messages=[{"role": "user", "content": "hi"}], + session_id="sess-mem", + ) + assert req.session_id == "sess-mem" + + def test_add_memory_request_session_id_optional(self): + from mirix.server.rest_api import AddMemoryRequest + + req = AddMemoryRequest( + meta_agent_id="meta-1", + messages=[{"role": "user", "content": "hi"}], + ) + assert req.session_id is None + + # Validation at the REST boundary (Codex review: C2 + I5). + # Invalid input should raise at model construction, not be queued. + def test_send_message_request_rejects_invalid_session_id(self): + from mirix.server.rest_api import SendMessageRequest + + with pytest.raises(ValueError): + SendMessageRequest( + message="hi", role="user", session_id="bad/chars" + ) + with pytest.raises(ValueError): + SendMessageRequest(message="hi", role="user", session_id="") + with pytest.raises(ValueError): + SendMessageRequest( + message="hi", role="user", session_id="a" * 65 + ) + + def test_add_memory_request_rejects_invalid_session_id(self): + from mirix.server.rest_api import AddMemoryRequest + + with pytest.raises(ValueError): + AddMemoryRequest( + meta_agent_id="meta-1", + messages=[{"role": "user", "content": "hi"}], + session_id="bad chars", + ) + with pytest.raises(ValueError): + AddMemoryRequest( + meta_agent_id="meta-1", + messages=[{"role": "user", "content": "hi"}], + session_id="", + ) + + +# ----------------------------- Queue proto ------------------------------- + + +class TestAddMemoryMismatchRejection: + """Codex review I1: top-level session_id and filter_tags.session_id must agree.""" + + def test_matching_values_are_allowed(self): + from mirix.server.rest_api import AddMemoryRequest + + # Same value in both places is fine. + req = AddMemoryRequest( + meta_agent_id="meta-1", + messages=[{"role": "user", "content": "hi"}], + session_id="sess-1", + filter_tags={"session_id": "sess-1"}, + ) + assert req.session_id == "sess-1" + assert req.filter_tags["session_id"] == "sess-1" + + def test_mismatch_is_rejected_at_request_model(self): + from mirix.server.rest_api import AddMemoryRequest + + with pytest.raises(ValueError): + AddMemoryRequest( + meta_agent_id="meta-1", + messages=[{"role": "user", "content": "hi"}], + session_id="sess-1", + filter_tags={"session_id": "sess-2"}, + ) + + +class TestQueueProtoSessionId: + """Ensure the proto round-trips session_id through MessageCreate.""" + + def test_proto_message_create_has_session_id_field(self): + from mirix.queue.message_pb2 import MessageCreate as ProtoMessageCreate + + field_names = {f.name for f in ProtoMessageCreate.DESCRIPTOR.fields} + assert "session_id" in field_names + + def test_put_messages_serializes_session_id(self, monkeypatch): + """put_messages should copy MessageCreate.session_id onto proto.""" + import asyncio + + from mirix.queue import queue_util + from mirix.schemas.client import Client + + saved = {} + + class FakeQueue: + async def save(self, msg): + saved["msg"] = msg + + monkeypatch.setattr(queue_util, "queue", FakeQueue()) + + async def run(): + await queue_util.put_messages( + actor=Client( + id="client-1", + organization_id="org-1", + name="c", + write_scope="w", + read_scopes=["w"], + ), + agent_id="agent-1", + input_messages=[ + MessageCreate( + role=MessageRole.user, + content="hi", + session_id="sess-q", + ) + ], + ) + + asyncio.run(run()) + + msg = saved["msg"] + assert len(msg.input_messages) == 1 + proto = msg.input_messages[0] + assert proto.HasField("session_id") + assert proto.session_id == "sess-q" + + def test_worker_restores_session_id_from_proto(self): + from mirix.queue.message_pb2 import MessageCreate as ProtoMessageCreate + from mirix.queue.worker import QueueWorker + + proto = ProtoMessageCreate() + proto.role = ProtoMessageCreate.ROLE_USER + proto.text_content = "hi" + proto.session_id = "sess-w" + + worker = QueueWorker.__new__(QueueWorker) + out = worker._convert_proto_message_to_pydantic(proto) + assert out.session_id == "sess-w" + + +# ----------------------------- ORM -------------------------------------- + + +class TestAgentStepPropagation: + """Decoupling lint (PRD: procedural-memory-distillation-decoupling). + + session_id identifies an EXTERNAL conversation, never the memory-production + machinery. The decoupling INVERTS the old rule: a non-chat agent + (meta_memory_agent + the memory sub-agents) must NOT stamp session_id on the + messages it synthesizes inside Agent.step โ€” those are transient bookkeeping + that gets discarded after each extraction. Only the chat_agent path may keep + inheriting it. + + The synthesized-message sites bind session_id to two derived locals โ€” + `input_session_id` (in the response-processing path) and `step_session_id` + (in the step loop / meta-bootstrap / summary path). The decoupling lives in + how those two locals are DERIVED: each is gated `... if else + None`. We lint that derivation rather than every call site, so the test is + robust to whitespace and to call-site churn while still pinning the + behavior: for a non-chat agent the value is None, so nothing it synthesizes + carries a session_id. + """ + + # Tolerant whitespace between tokens so the lint never breaks on reformatting. + _WS = r"\s*" + + def _gated_to_none_for_non_chat(self, src: str, local: str, attr_obj: str) -> bool: + """True iff ` = getattr(, "session_id", None) if + else None` appears in `src`, whitespace-insensitively. + + This is the load-bearing decoupling: the value is the input's session_id + ONLY for the chat_agent, and explicitly `None` for every other agent. + """ + import re + + ws = self._WS + # The assignment is wrapped in parens spanning several lines: + # = ( + # getattr(, "session_id", None) + # if self.agent_state.is_type(AgentType.chat_agent) + # else None + # ) + # `\s*` (in self._WS) spans newlines, so allow an optional opening paren + # after `=` and tolerate arbitrary whitespace between every token. + pattern = ( + rf"{re.escape(local)}{ws}={ws}\(?{ws}" + rf"getattr\({ws}{re.escape(attr_obj)}{ws},{ws}" + rf"[\"']session_id[\"']{ws},{ws}None{ws}\){ws}" + rf"if{ws}self\.agent_state\.is_type\({ws}AgentType\.chat_agent{ws}\){ws}" + rf"else{ws}None" + ) + return re.search(pattern, src) is not None + + def test_input_session_id_is_none_for_non_chat_agents(self): + """The response-path local `input_session_id` must resolve to None for + every non-chat agent, so the assistant/tool messages it stamps carry no + session_id.""" + from pathlib import Path + + src = Path("mirix/agent/agent.py").read_text() + assert self._gated_to_none_for_non_chat( + src, "input_session_id", "input_message" + ), ( + "input_session_id must be gated `... if chat_agent else None` so " + "non-chat synthesized messages carry NO session_id" + ) + + def test_step_session_id_is_none_for_non_chat_agents(self): + """The step-loop local `step_session_id` (heartbeats, meta-bootstrap, + summaries) must resolve to None for every non-chat agent.""" + from pathlib import Path + + src = Path("mirix/agent/agent.py").read_text() + assert self._gated_to_none_for_non_chat( + src, "step_session_id", "first_input_message" + ), ( + "step_session_id must be gated `... if chat_agent else None` so " + "non-chat synthesized messages carry NO session_id" + ) + + def test_synthesized_sites_never_bind_raw_input_session_id(self): + """Defense-in-depth: no synthesized message may read session_id straight + off the RAW triggering input, which would bypass the chat-agent gate and + leak the conversation id onto a non-chat agent's bookkeeping. The raw + inputs are `input_message` (response path) and `first_input_message` + (step path); their `.session_id` must only ever be read INSIDE the gated + derivation, never as a `session_id=` kwarg on a synthesized message. + """ + import re + from pathlib import Path + + src = Path("mirix/agent/agent.py").read_text() + ws = self._WS + for raw in ("input_message", "first_input_message"): + # A `session_id=` kwarg that reads the raw input's session_id โ€” in + # EITHER form, attribute access `.session_id` or + # `getattr(, "session_id", ...)` โ€” would bypass the chat-agent + # gate and leak the conversation id. The legitimate gated derivation + # assigns getattr(...) to the LOCAL (`input_session_id = ...`), never + # to a `session_id=` kwarg, so it is not matched here. + attr_leak = re.search( + r"session_id" + ws + r"=" + ws + re.escape(raw) + r"\.session_id", + src, + ) + getattr_leak = re.search( + r"session_id" + ws + r"=" + ws + + r"getattr\(" + ws + re.escape(raw) + ws + r",", + src, + ) + assert attr_leak is None and getattr_leak is None, ( + f"synthesized session_id must never bind {raw}'s session_id " + "directly (attribute or getattr) โ€” it must go through the " + "chat-agent-gated local" + ) + + def test_meta_bootstrap_does_not_inherit_session_id_for_non_chat(self): + """The meta-memory bootstrap MessageCreate is persisted (via + prepare_input_message_create) on the meta_memory_agent โ€” a NON-chat + agent. It must bind the gated `step_session_id` (which is None for the + meta agent), NOT a raw input attribute, so the bootstrap carries no + session_id. + """ + import re + from pathlib import Path + + src = Path("mirix/agent/agent.py").read_text() + ws = self._WS + meta_block_match = re.search( + r"meta_message" + ws + r"=" + ws + r"prepare_input_message_create\(" + r"[\s\S]*?MessageCreate\([\s\S]*?\)" + ws + r",", + src, + ) + assert meta_block_match, "meta-bootstrap MessageCreate block not found" + block = meta_block_match.group(0) + # It binds the gated local (whitespace/newline tolerant)โ€ฆ + assert re.search( + r"session_id" + ws + r"=" + ws + r"step_session_id", block + ), ( + "meta-memory bootstrap must bind session_id=step_session_id " + f"(gated to None for the meta agent), got:\n{block}" + ) + # โ€ฆand never a raw input attribute (attribute or getattr form). + for raw in ("input_message", "first_input_message"): + assert not re.search(re.escape(raw) + r"\.session_id", block) + assert not re.search(r"getattr\(" + ws + re.escape(raw) + ws + r",", block) + + +class TestOrmSessionId: + def test_orm_has_session_id_column(self): + from mirix.orm.message import Message as MessageORM + + cols = {c.name for c in MessageORM.__table__.columns} + assert "session_id" in cols + + def test_orm_session_id_is_indexed(self): + from mirix.orm.message import Message as MessageORM + + # At least one composite or single index covers session_id. + covered = any( + any(c.name == "session_id" for c in idx.columns) + for idx in MessageORM.__table__.indexes + ) + assert covered, "session_id should be indexed" + + +class TestStepUserMessageSessionContext: + """step_user_message() bypasses step() but is the SAME decoupling site: it + synthesizes a user Message and seeds _current_step_session_id for any + pre-persist summary, then restores the prior value. Per the + procedural-memory-distillation-decoupling PRD this site is also GATED โ€” a + non-chat agent must NOT inherit the conversation's session_id here either, so + the seeded value is gated `... if chat_agent else None` (not the raw + session_id argument). We inspect the source rather than run the coroutine + (full step needs LLM + DB).""" + + def test_step_user_message_gates_session_id_for_non_chat(self): + import inspect + import re + + from mirix.agent import agent as agent_mod + + src = inspect.getsource(agent_mod.Agent.step_user_message) + ws = r"\s*" + # The gate must bind to `effective_session_id` SPECIFICALLY โ€” i.e. + # effective_session_id = (session_id if chat_agent else None) + # Binding the regex to the actual local that flows into the synthesized + # message defeats a decoy like `unused = session_id if chat_agent else + # None; effective_session_id = session_id`, which an unbound `... if + # chat_agent else None` match would wrongly accept. + gate = re.search( + r"effective_session_id" + ws + r"=" + ws + r"\(?" + ws + + r"session_id" + ws + + r"if" + ws + r"self\.agent_state\.is_type\(" + ws + + r"AgentType\.chat_agent" + ws + r"\)" + ws + r"else" + ws + r"None", + src, + ) + assert gate is not None, ( + "step_user_message must gate `effective_session_id = (session_id if " + "chat_agent else None)` so a non-chat agent's synthesized message " + "carries NO session_id" + ) + # The synthesized message binds the GATED value, never the raw argument. + assert re.search(r"session_id" + ws + r"=" + ws + r"effective_session_id", src), ( + "step_user_message's dict_to_message must bind the gated " + "effective_session_id, not the raw session_id argument" + ) + # And it must NOT bind the raw `session_id` argument to the message. + assert not re.search( + r"dict_to_message\([\s\S]*?session_id" + ws + r"=" + ws + r"session_id\b", + src, + ), "step_user_message must not bind the RAW session_id argument" + + def test_step_user_message_seeds_gated_value_and_restores_stash(self): + import inspect + import re + + from mirix.agent import agent as agent_mod + + src = inspect.getsource(agent_mod.Agent.step_user_message) + ws = r"\s*" + # Seeds the GATED session context for the summarizer (not the raw arg)โ€ฆ + assert re.search( + r"self\._current_step_session_id" + ws + r"=" + ws + r"effective_session_id", + src, + ), "must seed _current_step_session_id with the gated effective_session_id" + # โ€ฆand uses try/finally to restore the prior value โ€” no leaks across calls. + assert "prev_session_id" in src + assert "finally" in src + assert "self._current_step_session_id = prev_session_id" in src + + def test_summarize_messages_inplace_accepts_explicit_session_id(self): + import inspect + + from mirix.agent import agent as agent_mod + + sig = inspect.signature(agent_mod.Agent.summarize_messages_inplace) + assert "session_id" in sig.parameters + # Default is optional None so existing callers don't break. + assert sig.parameters["session_id"].default is None + + +class TestCheckConstraintDialectGating: + """The CHECK uses PG's `~` regex operator, which SQLite doesn't understand. + Verify the constraint is only emitted for Postgres so SQLite create_all works.""" + + def test_check_compiles_for_postgresql(self): + from sqlalchemy.dialects import postgresql + from sqlalchemy.schema import CreateTable + + from mirix.orm.message import Message as MessageORM + + ddl = str( + CreateTable(MessageORM.__table__).compile(dialect=postgresql.dialect()) + ) + assert "ck_messages_session_id_format" in ddl + assert "~" in ddl # Postgres regex operator + + def test_check_is_suppressed_for_sqlite(self): + from sqlalchemy.dialects import sqlite + from sqlalchemy.schema import CreateTable + + from mirix.orm.message import Message as MessageORM + + ddl = str( + CreateTable(MessageORM.__table__).compile(dialect=sqlite.dialect()) + ) + # Constraint must not appear on SQLite โ€” its `~` regex would break CREATE TABLE. + assert "ck_messages_session_id_format" not in ddl + assert "session_id ~" not in ddl + + +class TestSessionIdConstantsInSync: + """Codex v2 nit: one source of truth for pattern/length across Python, ORM CheckConstraint, and SQL.""" + + def test_orm_column_length_matches_constant(self): + from mirix.orm.message import Message as MessageORM + from mirix.schemas.message import SESSION_ID_MAX_LEN + + col = MessageORM.__table__.c.session_id + assert col.type.length == SESSION_ID_MAX_LEN + + def test_orm_check_constraint_uses_shared_pattern(self): + from mirix.orm.message import Message as MessageORM + from mirix.schemas.message import SESSION_ID_SQL_PATTERN + + ck_texts = [ + str(c.sqltext) for c in MessageORM.__table__.constraints + if getattr(c, "name", None) == "ck_messages_session_id_format" + ] + assert ck_texts, "ck_messages_session_id_format not found" + assert SESSION_ID_SQL_PATTERN in ck_texts[0] + + def test_migration_sql_uses_shared_pattern_and_length(self): + from pathlib import Path + + from mirix.schemas.message import ( + SESSION_ID_MAX_LEN, + SESSION_ID_SQL_PATTERN, + ) + + phase1 = Path("scripts/migrate_add_message_session_id.sql").read_text() + phase2 = Path("scripts/migrate_add_message_session_id_phase2.sql").read_text() + + # Column length must match. + assert f"VARCHAR({SESSION_ID_MAX_LEN})" in phase1 + # CHECK regex must match the shared SQL pattern exactly. + assert SESSION_ID_SQL_PATTERN in phase1 + # Phase 2 must use CONCURRENTLY (online, write-safe). + assert "CREATE INDEX CONCURRENTLY" in phase2 diff --git a/tests/test_session_id_integration.py b/tests/test_session_id_integration.py new file mode 100644 index 000000000..a0351e01d --- /dev/null +++ b/tests/test_session_id_integration.py @@ -0,0 +1,189 @@ +""" +Integration test for the top-level session_id field: end-to-end round-trip +through MessageManager against real Postgres. + +Verifies: +- A message created with session_id persists that column. +- list_messages_for_agent(session_id=X) returns only messages for session X. +- list_messages_for_agent(session_id=None) returns all messages for that agent. +- Different session_ids correctly isolate messages. + +Requires the docker-compose Postgres (port 5433). Run: + pytest tests/test_session_id_integration.py -v -m integration +""" +from __future__ import annotations + +import asyncio +import sys +import uuid +from pathlib import Path + +import pytest +import pytest_asyncio + +pytestmark = [ + pytest.mark.integration, + pytest.mark.asyncio(loop_scope="module"), +] + +project_root = Path(__file__).parent.parent +sys.path.insert(0, str(project_root)) + + +@pytest_asyncio.fixture(scope="module") +def event_loop(): + loop = asyncio.new_event_loop() + yield loop + loop.close() + + +from mirix.schemas.client import Client as PydanticClient +from mirix.schemas.enums import MessageRole +from mirix.schemas.message import Message as PydanticMessage +from mirix.schemas.mirix_message_content import TextContent +from mirix.schemas.user import User as PydanticUser +from mirix.services.message_manager import MessageManager + + +@pytest.fixture +def message_manager(): + return MessageManager() + + +@pytest_asyncio.fixture(scope="module") +async def test_actor(): + from mirix.schemas.organization import Organization as PydanticOrganization + from mirix.services.client_manager import ClientManager + from mirix.services.organization_manager import OrganizationManager + + org_mgr = OrganizationManager() + client_mgr = ClientManager() + + org_id = f"test-session-id-org-{uuid.uuid4().hex[:8]}" + try: + await org_mgr.get_organization_by_id(org_id) + except Exception: + await org_mgr.create_organization( + PydanticOrganization(id=org_id, name="Session ID Test Org") + ) + + client_id = f"test-session-id-client-{uuid.uuid4().hex[:8]}" + try: + return await client_mgr.get_client_by_id(client_id) + except Exception: + return await client_mgr.create_client( + PydanticClient( + id=client_id, + organization_id=org_id, + name="Session ID Test Client", + write_scope="test-sid", + read_scopes=["test-sid"], + ) + ) + + +@pytest_asyncio.fixture(scope="module") +async def test_user(test_actor): + from mirix.services.user_manager import UserManager + + user_mgr = UserManager() + user_id = f"test-session-id-user-{uuid.uuid4().hex[:8]}" + try: + return await user_mgr.get_user_by_id(user_id) + except Exception: + return await user_mgr.create_user( + PydanticUser( + id=user_id, + name="Session ID Test User", + organization_id=test_actor.organization_id, + timezone="UTC", + ) + ) + + +@pytest_asyncio.fixture(scope="module") +async def test_agent(test_actor): + """Create a minimal agent we can attach messages to.""" + from mirix.schemas.agent import AgentType, CreateAgent + from mirix.services.agent_manager import AgentManager + + agent_mgr = AgentManager() + agent = await agent_mgr.create_agent( + agent_create=CreateAgent( + name=f"test-session-agent-{uuid.uuid4().hex[:8]}", + agent_type=AgentType.chat_agent, + description="Test agent for session_id round-trip", + system=None, + llm_config=None, + embedding_config=None, + ), + actor=test_actor, + ) + return agent + + +async def _create_msg(mgr, actor, agent, text, session_id): + return await mgr.create_message( + pydantic_msg=PydanticMessage( + agent_id=agent.id, + role=MessageRole.user, + content=[TextContent(text=text)], + session_id=session_id, + ), + actor=actor, + use_cache=False, + ) + + +class TestSessionIdRoundTrip: + async def test_persists_and_filters_by_session_id( + self, message_manager, test_actor, test_agent + ): + sid_a = f"sess-a-{uuid.uuid4().hex[:6]}" + sid_b = f"sess-b-{uuid.uuid4().hex[:6]}" + + a1 = await _create_msg(message_manager, test_actor, test_agent, "A1", sid_a) + a2 = await _create_msg(message_manager, test_actor, test_agent, "A2", sid_a) + b1 = await _create_msg(message_manager, test_actor, test_agent, "B1", sid_b) + null_msg = await _create_msg( + message_manager, test_actor, test_agent, "no-session", None + ) + + # Session A returns only A messages. + got_a = await message_manager.list_messages_for_agent( + agent_id=test_agent.id, + actor=test_actor, + session_id=sid_a, + limit=100, + use_cache=False, + ) + ids_a = {m.id for m in got_a} + assert a1.id in ids_a + assert a2.id in ids_a + assert b1.id not in ids_a + assert null_msg.id not in ids_a + + # Session B returns only B. + got_b = await message_manager.list_messages_for_agent( + agent_id=test_agent.id, + actor=test_actor, + session_id=sid_b, + limit=100, + use_cache=False, + ) + ids_b = {m.id for m in got_b} + assert b1.id in ids_b + assert a1.id not in ids_b + assert null_msg.id not in ids_b + + # No filter returns everything we created (and the session_id column round-trips). + all_msgs = await message_manager.list_messages_for_agent( + agent_id=test_agent.id, + actor=test_actor, + limit=100, + use_cache=False, + ) + by_id = {m.id: m for m in all_msgs} + assert by_id[a1.id].session_id == sid_a + assert by_id[b1.id].session_id == sid_b + assert by_id[null_msg.id].session_id is None diff --git a/tests/test_skill_cli_tools.py b/tests/test_skill_cli_tools.py new file mode 100644 index 000000000..3fcbb1fcd --- /dev/null +++ b/tests/test_skill_cli_tools.py @@ -0,0 +1,105 @@ +"""Tests for CLI-style skill tools, validators, and constants.""" +import pytest + + +def test_bump_patch_version(): + """Test version bumping.""" + # Import via importlib to avoid heavy dependency chain (rapidfuzz, etc.) + import importlib.util, types, sys + + spec = importlib.util.spec_from_file_location( + "_memory_tools_isolated", + "mirix/functions/function_sets/memory_tools.py", + ) + # We only need to extract the pure function source, not execute the whole module. + # Read the function directly to avoid import side-effects. + import ast, textwrap + + with open("mirix/functions/function_sets/memory_tools.py") as f: + source = f.read() + + tree = ast.parse(source) + for node in ast.walk(tree): + if isinstance(node, ast.FunctionDef) and node.name == "_bump_patch_version": + func_source = ast.get_source_segment(source, node) + break + else: + pytest.fail("_bump_patch_version not found in memory_tools.py") + + ns = {} + exec(func_source, ns) + _bump_patch_version = ns["_bump_patch_version"] + + assert _bump_patch_version("0.1.0") == "0.1.1" + assert _bump_patch_version("0.1.9") == "0.1.10" + assert _bump_patch_version("1.2.3") == "1.2.4" + assert _bump_patch_version("invalid") == "0.1.1" + assert _bump_patch_version("") == "0.1.1" + + +def test_skill_tools_constant(): + """SKILL_TOOLS contains all 5 CLI tools.""" + from mirix.constants import SKILL_TOOLS + + assert SKILL_TOOLS == ["skill_list", "skill_read", "skill_create", "skill_edit", "skill_delete"] + + +def test_trigger_threshold_constant(): + """SKILL_TRIGGER_MESSAGE_THRESHOLD exists and defaults to 10.""" + from mirix.constants import SKILL_TRIGGER_MESSAGE_THRESHOLD + + assert isinstance(SKILL_TRIGGER_MESSAGE_THRESHOLD, int) + assert SKILL_TRIGGER_MESSAGE_THRESHOLD == 10 + + +def test_skill_validators_registered(): + """Validators are registered for skill_create, skill_edit, skill_delete.""" + # Import validator module directly to avoid heavy dependency chain + import importlib.util + spec = importlib.util.spec_from_file_location( + "tool_validators", + "mirix/agent/tool_validators.py" + ) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + + # skill_create: missing name + result = mod.validate_tool_args( + "skill_create", + {"name": "", "description": "x", "instructions": "x", "entry_type": "workflow"}, + ) + assert result is not None + assert "name" in result + + # skill_create: valid + result = mod.validate_tool_args( + "skill_create", + {"name": "x", "description": "x", "instructions": "x", "entry_type": "workflow"}, + ) + assert result is None + + # skill_edit: missing field + result = mod.validate_tool_args("skill_edit", {"skill_id": "proc-1", "field": ""}) + assert result is not None + + # skill_edit: text field without old_text + result = mod.validate_tool_args("skill_edit", {"skill_id": "proc-1", "field": "instructions"}) + assert result is not None + assert "old_text" in result + + # skill_edit: value field valid + result = mod.validate_tool_args("skill_edit", {"skill_id": "proc-1", "field": "triggers", "value": ["a"]}) + assert result is None + + # skill_edit: invalid field + result = mod.validate_tool_args("skill_edit", {"skill_id": "proc-1", "field": "bogus", "value": "x"}) + assert result is not None + assert "must be one of" in result + + # skill_delete: valid + result = mod.validate_tool_args("skill_delete", {"skill_id": "proc-1"}) + assert result is None + + # skill_delete: empty + result = mod.validate_tool_args("skill_delete", {"skill_id": ""}) + assert result is not None diff --git a/tests/test_skill_edit_budget.py b/tests/test_skill_edit_budget.py new file mode 100644 index 000000000..170dde5c7 --- /dev/null +++ b/tests/test_skill_edit_budget.py @@ -0,0 +1,607 @@ +"""Tests for C4 โ€” bounded, count-driven, quality-aware edit budget + size/delete +gates on the skill tools. + +Three layers, all DB-free / network-free: + +1. Pure budget formula (`compute_edit_budget`): count-driven, clamped, B_min=0 + skip, hybrid `min(formula, autonomous)`. quality_score never drives it. +2. Pure size gate (`_edit_exceeds_size_gate`, `_instructions_over_ceiling`): + char-delta + change-ratio boundaries, instructions ceiling. +3. Budget/size/delete GATE behavior on `skill_create` / `skill_edit` / + `skill_delete`, exercised against a lightweight fake Agent whose + `procedural_memory_manager` is a stub. The counter is the per-instance + attribute `self._edit_budget_remaining`; we assert it is per-instance (no + cross-agent / cross-user leak) and that an unset attr means "no limit". + +The constants live in `mirix.constants` near the existing SKILL_* ones; the +gate logic lives in `mirix.functions.function_sets.memory_tools`. +""" + +from __future__ import annotations + +import asyncio +import difflib + +from mirix.constants import ( + SKILL_DELETE_BUDGET_MAX, + SKILL_EDIT_BUDGET_ALPHA_FAIL, + SKILL_EDIT_BUDGET_ALPHA_SUCC, + SKILL_EDIT_BUDGET_B0, + SKILL_EDIT_BUDGET_MAX, + SKILL_EDIT_BUDGET_MIN, + SKILL_EDIT_MAJOR_RATIO, + SKILL_MAX_EDIT_CHAR_DELTA, + SKILL_MAX_INSTRUCTIONS_CHARS, +) +from mirix.functions.function_sets import memory_tools as mt + + +# ============================ Budget formula ============================== + + +class TestBudgetConstants: + def test_constant_values_match_design(self): + # DESIGN ยงC4: B0=1, alpha_f=1.0, alpha_s=0.5, B_min=0, B_max=6. + assert SKILL_EDIT_BUDGET_B0 == 1 + assert SKILL_EDIT_BUDGET_ALPHA_FAIL == 1.0 + assert SKILL_EDIT_BUDGET_ALPHA_SUCC == 0.5 + assert SKILL_EDIT_BUDGET_MIN == 0 + assert SKILL_EDIT_BUDGET_MAX == 6 + assert SKILL_MAX_EDIT_CHAR_DELTA == 800 + assert SKILL_EDIT_MAJOR_RATIO == 0.4 + assert SKILL_MAX_INSTRUCTIONS_CHARS == 12000 + assert SKILL_DELETE_BUDGET_MAX == 1 + + +class TestComputeEditBudget: + def _agg(self, n_high_fail=0, n_high_succ=0, n=None, mean_q=0.0): + if n is None: + n = n_high_fail + n_high_succ + return { + "n": n, + "n_high_fail": n_high_fail, + "n_high_succ": n_high_succ, + "mean_q": mean_q, + } + + def test_all_failures_clamps_to_b_max(self): + # raw = 1 + 1.0*5 + 0.5*0 = 6 -> clamp(6, 0, 6) = 6 + agg = self._agg(n_high_fail=5) + assert mt.compute_edit_budget(agg) == 6 + + def test_more_failures_still_clamps_to_b_max(self): + # raw = 1 + 1.0*8 = 9 -> clamp to 6. + agg = self._agg(n_high_fail=8) + assert mt.compute_edit_budget(agg) == SKILL_EDIT_BUDGET_MAX + + def test_mixed_window_lands_mid(self): + # raw = 1 + 1.0*2 + 0.5*2 = 4 -> 4 + agg = self._agg(n_high_fail=2, n_high_succ=2) + assert mt.compute_edit_budget(agg) == 4 + + def test_single_failure(self): + # raw = 1 + 1.0*1 = 2 + agg = self._agg(n_high_fail=1) + assert mt.compute_edit_budget(agg) == 2 + + def test_successes_only_half_weight(self): + # raw = 1 + 0.5*3 = 2.5 -> round(2.5) banker's? we want round-half-up=3? -3 + # DESIGN says clamp(round(raw)). Python round(2.5)=2 (banker's). The + # implementation must use round-half-up so 2.5 -> 3 is deterministic and + # matches "mixed => 2-4" intuition. Assert the documented value. + agg = self._agg(n_high_succ=3) + assert mt.compute_edit_budget(agg) == 3 + + def test_noise_window_is_zero(self): + # No structurally-gated records: n_high_* == 0. Even though B0=1, the + # B_min=0 *skip* is handled by the curator (no aggregate -> no budget), + # but compute_edit_budget on an all-noise aggregate must itself yield 0 + # so the curator's "skip when 0" decision keys off the same number. + agg = self._agg(n_high_fail=0, n_high_succ=0, n=4) + assert mt.compute_edit_budget(agg) == 0 + + def test_quality_score_never_drives_budget(self): + # Two aggregates with identical counts but wildly different mean_q must + # produce the SAME budget: quality_score is ranking-only. + low_q = self._agg(n_high_fail=2, n_high_succ=1, mean_q=0.01) + high_q = self._agg(n_high_fail=2, n_high_succ=1, mean_q=0.99) + assert mt.compute_edit_budget(low_q) == mt.compute_edit_budget(high_q) + + def test_hybrid_min_formula_and_autonomous(self): + agg = self._agg(n_high_fail=5) # formula = 6 + # autonomous may only REDUCE: final = min(6, 2) = 2. + assert mt.compute_edit_budget(agg, autonomous=2) == 2 + # autonomous larger than formula cannot raise it: min(6, 10) = 6. + assert mt.compute_edit_budget(agg, autonomous=10) == 6 + + def test_hybrid_autonomous_clamped_to_band(self): + agg = self._agg(n_high_fail=5) # formula = 6 + # A negative autonomous is clamped to B_min then min'd: min(6, 0) = 0. + assert mt.compute_edit_budget(agg, autonomous=-3) == 0 + + +# ============================ Size gate (pure) ============================ + + +class TestEditSizeGate: + BASE = "x" * 1000 + + def test_char_delta_799_allowed(self): + # |len(new) - len(old)| = 799 < 800 -> allowed (no reason). + old = self.BASE + new = self.BASE + ("y" * 799) + assert abs(len(new) - len(old)) == 799 + # keep ratio low so only the char-delta dimension is under test + assert mt._edit_exceeds_size_gate(old, new) is None + + def test_char_delta_801_rejected(self): + old = self.BASE + new = self.BASE + ("y" * 801) + assert abs(len(new) - len(old)) == 801 + reason = mt._edit_exceeds_size_gate(old, new) + assert reason is not None + assert "too large" in reason.lower() + + def test_char_delta_exactly_800_allowed(self): + # Boundary: gate rejects only when delta > 800, so 800 is allowed. + old = self.BASE + new = self.BASE + ("y" * 800) + assert abs(len(new) - len(old)) == 800 + assert mt._edit_exceeds_size_gate(old, new) is None + + def test_ratio_039_allowed(self): + # Same-length old/new (char-delta 0) so ONLY the change-ratio gate is + # under test. For old="a"*L, new="a"*k+"b"*(L-k), SequenceMatcher finds + # the common run of `a`s (length k), so ratio = k/L and the change-ratio + # is 1 - k/L. L=1000, k=610 -> change-ratio = 0.39 (< 0.40). + old = "a" * 1000 + new = ("a" * 610) + ("b" * 390) + ratio_change = 1 - difflib.SequenceMatcher(None, old, new).ratio() + assert ratio_change < SKILL_EDIT_MAJOR_RATIO, ratio_change + assert mt._edit_exceeds_size_gate(old, new) is None + + def test_ratio_041_rejected(self): + # k=590 -> change-ratio = 0.41 (>= 0.40), so the gate trips. + old = "a" * 1000 + new = ("a" * 590) + ("b" * 410) + ratio_change = 1 - difflib.SequenceMatcher(None, old, new).ratio() + assert ratio_change >= SKILL_EDIT_MAJOR_RATIO, ratio_change + reason = mt._edit_exceeds_size_gate(old, new) + assert reason is not None + assert "too large" in reason.lower() + + def test_instructions_ceiling(self): + under = "z" * (SKILL_MAX_INSTRUCTIONS_CHARS - 1) + over = "z" * (SKILL_MAX_INSTRUCTIONS_CHARS + 1) + assert mt._instructions_over_ceiling(under) is False + assert mt._instructions_over_ceiling(over) is True + + +# ============================ Fake Agent harness ========================== + + +class _FakeSkill: + def __init__( + self, + skill_id, + name, + instructions="short instructions", + version="0.1.0", + description="d", + entry_type="guide", + ): + self.id = skill_id + self.name = name + self.instructions = instructions + self.version = version + self.description = description + self.entry_type = entry_type + self.triggers = [] + self.examples = [] + + +class _FakeProcManager: + """Stub manager: records mutations, never touches a DB or embeddings.""" + + def __init__(self, skills=None): + self._skills = {s.id: s for s in (skills or [])} + self.created = [] + self.updated = [] + self.deleted = [] + + async def list_procedures(self, **kwargs): + # Used by skill_create's name-dedup pre-check. Return empty (no dup). + return [] + + async def insert_procedure(self, **kwargs): + self.created.append(kwargs) + sid = f"proc-new-{len(self.created)}" + return _FakeSkill(sid, kwargs["name"], kwargs.get("instructions", "")) + + async def get_item_by_id(self, item_id, **kwargs): + if item_id not in self._skills: + raise KeyError(item_id) + return self._skills[item_id] + + async def update_item(self, item_update, **kwargs): + data = item_update.model_dump(exclude_unset=True) + self.updated.append(data) + skill = self._skills[data["id"]] + for k, v in data.items(): + if k != "id": + setattr(skill, k, v) + return skill + + async def delete_procedure_by_id(self, procedure_id, **kwargs): + self.deleted.append(procedure_id) + self._skills.pop(procedure_id, None) + + +class _FakeUser: + id = "user-1" + organization_id = "org-1" + + +class _FakeActor: + id = "client-1" + organization_id = "org-1" + + +class _FakeAgentState: + id = "agent-1" + parent_id = None + + +class _FakeAgent: + """Minimal `self` for the skill tools. Only the attributes the tools read.""" + + def __init__(self, manager): + self.procedural_memory_manager = manager + self.user = _FakeUser() + self.actor = _FakeActor() + self.agent_state = _FakeAgentState() + self.filter_tags = None + self.use_cache = True + self.user_id = "user-1" + + +def _run(coro): + return asyncio.run(coro) + + +# ============================ Budget gate ================================ + + +class TestBudgetGateExhaustion: + def test_create_decrements_and_exhausts(self): + agent = _FakeAgent(_FakeProcManager()) + agent._edit_budget_remaining = 2 + + r1 = _run(mt.skill_create(agent, "skill-a", "d", "i", "guide")) + assert "created" in r1.lower() + assert agent._edit_budget_remaining == 1 + + r2 = _run(mt.skill_create(agent, "skill-b", "d", "i", "guide")) + assert "created" in r2.lower() + assert agent._edit_budget_remaining == 0 + + # Third create: budget exhausted -> no mutation, advisory message. + r3 = _run(mt.skill_create(agent, "skill-c", "d", "i", "guide")) + assert "exhausted" in r3.lower() + assert "finish_memory_update" in r3 + assert len(agent.procedural_memory_manager.created) == 2 # only 2 created + assert agent._edit_budget_remaining == 0 + + def test_edit_decrements_same_counter_as_create(self): + skill = _FakeSkill("proc-1", "skill-x", "the quick brown fox") + agent = _FakeAgent(_FakeProcManager([skill])) + agent._edit_budget_remaining = 1 + + # One edit consumes the only unit. + r1 = _run( + mt.skill_edit( + agent, "proc-1", "instructions", old_text="quick", new_text="slow" + ) + ) + assert "updated" in r1.lower() + assert agent._edit_budget_remaining == 0 + + # Now a CREATE must also see the shared budget exhausted. + r2 = _run(mt.skill_create(agent, "skill-y", "d", "i", "guide")) + assert "exhausted" in r2.lower() + assert len(agent.procedural_memory_manager.created) == 0 + + def test_exhausted_edit_does_not_mutate(self): + skill = _FakeSkill("proc-1", "skill-x", "the quick brown fox") + agent = _FakeAgent(_FakeProcManager([skill])) + agent._edit_budget_remaining = 0 + + r = _run( + mt.skill_edit( + agent, "proc-1", "instructions", old_text="quick", new_text="slow" + ) + ) + assert "exhausted" in r.lower() + assert agent.procedural_memory_manager.updated == [] + + def test_unset_budget_means_no_limit(self): + agent = _FakeAgent(_FakeProcManager()) + # No _edit_budget_remaining attribute set at all. + assert not hasattr(agent, "_edit_budget_remaining") + for i in range(10): + r = _run(mt.skill_create(agent, f"skill-{i}", "d", "i", "guide")) + assert "created" in r.lower() + assert len(agent.procedural_memory_manager.created) == 10 + + +class TestBudgetCounterIsPerInstance: + def test_two_agents_do_not_share_budget(self): + a = _FakeAgent(_FakeProcManager()) + b = _FakeAgent(_FakeProcManager()) + a._edit_budget_remaining = 1 + b._edit_budget_remaining = 1 + + # Exhaust agent a. + _run(mt.skill_create(a, "skill-a1", "d", "i", "guide")) + assert a._edit_budget_remaining == 0 + r = _run(mt.skill_create(a, "skill-a2", "d", "i", "guide")) + assert "exhausted" in r.lower() + + # agent b is untouched: still has its own full budget. + assert b._edit_budget_remaining == 1 + r = _run(mt.skill_create(b, "skill-b1", "d", "i", "guide")) + assert "created" in r.lower() + assert b._edit_budget_remaining == 0 + + +# ============================ Size gate on edit ========================= + + +class TestEditSizeGateOnTool: + def test_oversized_edit_rejected_without_consuming_budget(self): + big = "a" * 1000 + skill = _FakeSkill("proc-1", "skill-x", big) + agent = _FakeAgent(_FakeProcManager([skill])) + agent._edit_budget_remaining = 3 + + # new_text differs by > 800 chars from old_text. + new = big + ("b" * 900) + r = _run( + mt.skill_edit(agent, "proc-1", "instructions", old_text=big, new_text=new) + ) + assert "too large" in r.lower() + # No mutation, and budget NOT consumed (rejection is free). + assert agent.procedural_memory_manager.updated == [] + assert agent._edit_budget_remaining == 3 + + def test_small_edit_passes_size_gate_and_consumes_budget(self): + skill = _FakeSkill("proc-1", "skill-x", "the quick brown fox jumps") + agent = _FakeAgent(_FakeProcManager([skill])) + agent._edit_budget_remaining = 3 + r = _run( + mt.skill_edit( + agent, "proc-1", "instructions", old_text="quick", new_text="speedy" + ) + ) + assert "updated" in r.lower() + assert agent._edit_budget_remaining == 2 + + def test_over_ceiling_instructions_routes_to_create(self): + # Editing instructions so the *resulting* text exceeds the hard ceiling + # must be refused with a "use skill_create" style message and not mutate. + base = "a" * (SKILL_MAX_INSTRUCTIONS_CHARS - 100) + skill = _FakeSkill("proc-1", "skill-x", base) + agent = _FakeAgent(_FakeProcManager([skill])) + agent._edit_budget_remaining = 3 + # Replace a short anchor with a huge block that pushes over the ceiling. + anchor = base[:50] + huge = anchor + ("b" * 500) # net +450, but total > ceiling + # Make the total exceed the ceiling regardless of size-gate by appending. + skill.instructions = base + ("a" * 200) # now total ~ ceiling+100 + r = _run( + mt.skill_edit( + agent, "proc-1", "instructions", old_text=anchor, new_text=huge + ) + ) + reason = r.lower() + assert ( + ("ceiling" in reason) + or ("skill_create" in reason) + or ("too large" in reason) + ) + assert agent.procedural_memory_manager.updated == [] + assert agent._edit_budget_remaining == 3 + + def test_size_gate_only_applies_to_text_fields(self): + # A triggers (non-text) edit must not be size-gated. + skill = _FakeSkill("proc-1", "skill-x", "i") + agent = _FakeAgent(_FakeProcManager([skill])) + agent._edit_budget_remaining = 3 + r = _run(mt.skill_edit(agent, "proc-1", "triggers", value='["a", "b"]')) + assert "updated" in r.lower() + assert agent._edit_budget_remaining == 2 + + +# ============================ Delete gate =============================== + + +class _FakeSoftDeleteManager(_FakeProcManager): + """Adds a soft-delete (exclude-from-retrieval) path the gate can prefer.""" + + def __init__(self, skills=None): + super().__init__(skills) + self.soft_deleted = [] + + async def soft_delete_procedure_by_id(self, procedure_id, **kwargs): + # Soft delete = excluded from retrieval but row retained. + self.soft_deleted.append(procedure_id) + skill = self._skills.get(procedure_id) + if skill is not None: + skill.is_deleted = True + + async def list_procedures(self, **kwargs): + # Retrieval excludes soft-deleted skills. + return [s for s in self._skills.values() if not getattr(s, "is_deleted", False)] + + +class TestDeleteAuthorizationHelper: + def test_record_names_skill_harmful_by_name(self): + rec = { + "detail": "root_cause: the deploy-prod skill is actively harmful, " + "it deletes the wrong namespace", + "record_type": "failure", + } + assert ( + mt._record_authorizes_delete( + rec, skill_name="deploy-prod", skill_id="proc-9" + ) + is True + ) + + def test_record_names_skill_harmful_by_id(self): + rec = { + "detail": "proc-9 caused the failure; it is harmful and should be removed", + "record_type": "failure", + } + assert ( + mt._record_authorizes_delete( + rec, skill_name="deploy-prod", skill_id="proc-9" + ) + is True + ) + + def test_record_merely_redundant_does_not_authorize(self): + rec = { + "detail": "the deploy-prod skill is redundant with deploy-staging", + "record_type": "failure", + } + # "redundant" is NOT "actively harmful" -> no delete authorization. + assert ( + mt._record_authorizes_delete( + rec, skill_name="deploy-prod", skill_id="proc-9" + ) + is False + ) + + def test_success_record_never_authorizes_delete(self): + rec = {"detail": "the deploy-prod skill is harmful", "record_type": "success"} + assert ( + mt._record_authorizes_delete( + rec, skill_name="deploy-prod", skill_id="proc-9" + ) + is False + ) + + def test_harmful_marker_must_co_occur_with_named_skill(self): + # The harmful marker and the skill mention are in DIFFERENT clauses: + # "deploy-prod is redundant; deploy-staging was harmful". This must NOT + # authorize deleting deploy-prod (the harmful marker refers to a DIFFERENT + # skill). Proximity gate: marker + name must share a clause/sentence. + rec = { + "detail": "the deploy-prod skill is redundant; " + "the deploy-staging skill was harmful and caused the failure", + "record_type": "failure", + } + assert ( + mt._record_authorizes_delete( + rec, skill_name="deploy-prod", skill_id="proc-9" + ) + is False + ) + # But deleting deploy-staging (the one the marker DOES refer to) is ok. + assert ( + mt._record_authorizes_delete( + rec, skill_name="deploy-staging", skill_id="proc-10" + ) + is True + ) + + def test_id_match_is_whole_token_not_prefix(self): + # A record naming 'proc-90' harmful must NOT authorize deleting 'proc-9' + # (substring/prefix cross-authorization guard). + rec = { + "detail": "proc-90 is actively harmful and caused the failure", + "record_type": "failure", + } + assert ( + mt._record_authorizes_delete(rec, skill_name="other", skill_id="proc-9") + is False + ) + # The exact id still authorizes. + assert ( + mt._record_authorizes_delete(rec, skill_name="other", skill_id="proc-90") + is True + ) + + def test_record_not_naming_skill_does_not_authorize(self): + rec = {"detail": "some other skill is harmful", "record_type": "failure"} + assert ( + mt._record_authorizes_delete( + rec, skill_name="deploy-prod", skill_id="proc-9" + ) + is False + ) + + +class TestDeleteGateOnTool: + def test_delete_rejected_without_authorization(self): + skill = _FakeSkill("proc-1", "skill-x", "i") + agent = _FakeAgent(_FakeProcManager([skill])) + # No _delete_authorized_skill_ids set -> nothing is authorized. + agent._delete_budget_remaining = 1 + r = _run(mt.skill_delete(agent, "proc-1")) + assert "not authorized" in r.lower() or "no failure record" in r.lower() + assert agent.procedural_memory_manager.deleted == [] + + def test_delete_d_max_enforced(self): + s1 = _FakeSkill("proc-1", "skill-1", "i") + s2 = _FakeSkill("proc-2", "skill-2", "i") + agent = _FakeAgent(_FakeProcManager([s1, s2])) + agent._delete_budget_remaining = 1 + agent._delete_authorized_skill_ids = {"proc-1", "proc-2"} + # Prefer hard delete here to keep the assertion on the hard path simple. + agent._prefer_soft_delete = False + + r1 = _run(mt.skill_delete(agent, "proc-1")) + assert "deleted" in r1.lower() + assert agent._delete_budget_remaining == 0 + + # Second authorized delete must be refused: D_max=1. + r2 = _run(mt.skill_delete(agent, "proc-2")) + assert ( + "budget" in r2.lower() or "d_max" in r2.lower() or "exhausted" in r2.lower() + ) + assert agent.procedural_memory_manager.deleted == ["proc-1"] + + def test_soft_delete_preferred_and_excludes_from_retrieval(self): + skill = _FakeSkill("proc-1", "skill-x", "i") + mgr = _FakeSoftDeleteManager([skill]) + agent = _FakeAgent(mgr) + agent._delete_budget_remaining = 1 + agent._delete_authorized_skill_ids = {"proc-1"} + agent._prefer_soft_delete = True + + r = _run(mt.skill_delete(agent, "proc-1")) + assert ( + "deleted" in r.lower() + or "superseded" in r.lower() + or "excluded" in r.lower() + ) + # Soft-delete path taken (not a hard delete). + assert mgr.soft_deleted == ["proc-1"] + assert mgr.deleted == [] + # Excluded from retrieval afterwards. + remaining = _run(mgr.list_procedures()) + assert all(s.id != "proc-1" for s in remaining) + + def test_delete_unset_budget_means_no_limit(self): + # When no delete-budget attr is set, deletes are NOT capped, but they + # still require authorization (the harmful-naming gate is independent of + # the budget counter). + s1 = _FakeSkill("proc-1", "skill-1", "i") + agent = _FakeAgent(_FakeProcManager([s1])) + agent._delete_authorized_skill_ids = {"proc-1"} + agent._prefer_soft_delete = False + r = _run(mt.skill_delete(agent, "proc-1")) + assert "deleted" in r.lower() + assert agent.procedural_memory_manager.deleted == ["proc-1"] diff --git a/tests/test_skill_experience.py b/tests/test_skill_experience.py new file mode 100644 index 000000000..238594cdb --- /dev/null +++ b/tests/test_skill_experience.py @@ -0,0 +1,718 @@ +"""Tests for the general session-experience store (Goal 2 storage layer). + +Three layers, mirroring test_skill_evolution_record.py: + +1. Pydantic schema validation (DB-free) โ€” rejects bad experience_type/status, + enforces length caps, defaults status='pending', CLAMPS importance/credibility + into [0,1] (garbage -> 0.0). +2. ORM shape (DB-free) โ€” table name, mixin + explicit columns, sexp- id prefix, + agent+status index, ORM package registration, migration SQL. +3. Manager behavior (DB-backed, integration) โ€” create, list ordering by + importance*credibility DESC, [0,1] enforcement at persistence, mark_consumed + idempotency + lineage, mark_superseded, aggregate, per-agent isolation. + +The DB-backed layer runs on a hermetic throwaway-SQLite sessionmaker (no +Postgres) and is marked `integration` per the TESTS-phase spec. +""" + +from __future__ import annotations + +import inspect +import uuid +from pathlib import Path + +import pytest +import pytest_asyncio +from pydantic import ValidationError +from sqlalchemy import inspect as sa_inspect + +from mirix.orm.skill_experience import SkillExperience as SkillExperienceORM +from mirix.schemas.skill_experience import ( + SKILL_EXPERIENCE_MAX_CONTENT_LEN, + SKILL_EXPERIENCE_MAX_EVIDENCE_LEN, + SKILL_EXPERIENCE_MAX_TITLE_LEN, + SkillExperience as PydanticSkillExperience, + SkillExperienceBase, + SkillExperienceCreate, + SkillExperienceResponse, + SkillExperienceUpdate, + _clamp01, +) +from mirix.services.skill_experience_manager import SkillExperienceManager + + +# ============================ Schema validation =========================== + + +class TestSkillExperienceSchema: + def _kwargs(self, **overrides) -> dict: + base = dict( + session_id="sess-1", + experience_type="worth_learning", + title="batch independent calls", + content="when calls are independent, batch them", + importance=0.8, + credibility=0.9, + evidence='{"quote":"great","signal_type":"user_confirmation"}', + ) + base.update(overrides) + return base + + def test_valid_base(self): + rec = SkillExperienceBase(**self._kwargs()) + assert rec.experience_type == "worth_learning" + assert rec.status == "pending" + + def test_experience_type_accepts_avoiding(self): + rec = SkillExperienceBase(**self._kwargs(experience_type="worth_avoiding")) + assert rec.experience_type == "worth_avoiding" + + def test_experience_type_rejects_unknown(self): + with pytest.raises(ValidationError): + SkillExperienceBase(**self._kwargs(experience_type="partial")) + + def test_status_rejects_unknown(self): + with pytest.raises(ValidationError): + SkillExperienceBase(**self._kwargs(status="archived")) + + def test_status_accepts_known(self): + for s in ("pending", "consumed", "superseded"): + assert SkillExperienceBase(**self._kwargs(status=s)).status == s + + def test_importance_clamped_high(self): + rec = SkillExperienceBase(**self._kwargs(importance=5.0)) + assert rec.importance == 1.0 + + def test_credibility_clamped_low(self): + rec = SkillExperienceBase(**self._kwargs(credibility=-3.0)) + assert rec.credibility == 0.0 + + def test_garbage_score_becomes_zero(self): + rec = SkillExperienceBase(**self._kwargs(importance="garbage")) + assert rec.importance == 0.0 + + def test_clamp01_helper(self): + assert _clamp01(5.0) == 1.0 + assert _clamp01(-3) == 0.0 + assert _clamp01("nope") == 0.0 + assert _clamp01(0.5) == 0.5 + assert _clamp01(float("inf")) == 1.0 + + def test_clamp01_nan_must_become_zero(self): + # Regression for the NaN gap (codex P1-1, fixed in _clamp01): NaN slips past + # both `< 0` and `> 1` (all NaN comparisons are False) and would poison the + # importance*credibility ordering. The [0,1] invariant requires NaN -> 0.0. + assert _clamp01(float("nan")) == 0.0 + assert _clamp01("nan") == 0.0 + + def test_title_length_cap(self): + with pytest.raises(ValidationError): + SkillExperienceBase( + **self._kwargs(title="x" * (SKILL_EXPERIENCE_MAX_TITLE_LEN + 1)) + ) + + def test_content_length_cap(self): + with pytest.raises(ValidationError): + SkillExperienceBase( + **self._kwargs(content="x" * (SKILL_EXPERIENCE_MAX_CONTENT_LEN + 1)) + ) + + def test_evidence_length_cap(self): + with pytest.raises(ValidationError): + SkillExperienceBase( + **self._kwargs(evidence="x" * (SKILL_EXPERIENCE_MAX_EVIDENCE_LEN + 1)) + ) + + def test_create_requires_owners(self): + rec = SkillExperienceCreate( + **self._kwargs(), agent_id="a", user_id="u", organization_id="o" + ) + assert rec.agent_id == "a" + + def test_full_schema_carries_db_fields(self): + rec = PydanticSkillExperience( + **self._kwargs(), id="sexp-abc", agent_id="a", + user_id="u", organization_id="o", + ) + assert rec.id == "sexp-abc" + assert rec.consumed_by is None + assert rec.influenced_skill_ids is None + + def test_update_validates_status(self): + with pytest.raises(ValidationError): + SkillExperienceUpdate(id="sexp-1", status="bogus") + + def test_update_accepts_consumed(self): + upd = SkillExperienceUpdate(id="sexp-1", status="consumed", consumed_by="run-1") + assert upd.status == "consumed" + + def test_response_is_subclass(self): + assert issubclass(SkillExperienceResponse, PydanticSkillExperience) + + def test_id_prefix(self): + assert SkillExperienceBase.__id_prefix__ == "sexp" + + +# ================================ ORM shape =============================== + + +class TestSkillExperienceORM: + def _cols(self): + return {c.key for c in sa_inspect(SkillExperienceORM).columns} + + def test_table_name(self): + assert SkillExperienceORM.__tablename__ == "skill_experience" + + def test_inherits_mixin_columns(self): + cols = self._cols() + for c in ["id", "organization_id", "user_id", "agent_id", "created_at"]: + assert c in cols, f"missing inherited column: {c}" + + def test_has_explicit_columns(self): + cols = self._cols() + for c in [ + "session_id", "experience_type", "title", "content", + "importance", "credibility", "evidence", "status", + "consumed_by", "influenced_skill_ids", + ]: + assert c in cols, f"missing explicit column: {c}" + + def test_id_default_uses_sexp_prefix(self): + default = SkillExperienceORM.__table__.c.id.default + assert default is not None + assert default.arg(None).startswith("sexp-") + + def test_status_defaults_to_pending(self): + assert SkillExperienceORM.__table__.c.status.default.arg == "pending" + + def test_pydantic_model_link(self): + assert SkillExperienceORM.__pydantic_model__ is PydanticSkillExperience + + def test_has_agent_status_index(self): + idx = {i.name for i in SkillExperienceORM.__table__.indexes} + assert "ix_skill_experience_agent_status" in idx + + def test_registered_in_orm_package(self): + import mirix.orm as orm_pkg + + assert "SkillExperience" in orm_pkg.__all__ + assert orm_pkg.SkillExperience is SkillExperienceORM + + +# =============================== Migration SQL ============================= + + +class TestMigrationSql: + SQL_PATH = Path("scripts/migrate_add_skill_experience.sql") + + def _sql(self) -> str: + return self.SQL_PATH.read_text() + + def test_creates_table(self): + assert "CREATE TABLE IF NOT EXISTS skill_experience" in self._sql() + + def test_idempotent_block(self): + sql = self._sql() + assert "CREATE INDEX IF NOT EXISTS" in sql + + def test_creates_agent_status_index(self): + assert "ix_skill_experience_agent_status" in self._sql() + + +# ============================== Manager surface ============================ + + +class TestManagerSurface: + def test_exposes_async_methods(self): + for name in ( + "create_experience", "list_experiences", + "mark_consumed", "mark_superseded", "aggregate", + "delete_by_user_id", "delete_by_client_id", + ): + assert hasattr(SkillExperienceManager, name) + unwrapped = inspect.unwrap(getattr(SkillExperienceManager, name)) + assert inspect.iscoroutinefunction(unwrapped), f"{name} must be async" + + def test_no_asyncio_run(self): + src = inspect.getsource(SkillExperienceManager) + assert "asyncio.run(" not in src + + +# ============================ Manager (DB-backed) ========================== +# Hermetic throwaway-SQLite (no Postgres). Marked integration per the spec. + + +@pytest_asyncio.fixture(scope="module", loop_scope="module") +async def session_maker(tmp_path_factory): + from contextlib import asynccontextmanager + + from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine + + import mirix.orm # noqa: F401 -- register all ORM classes + from mirix.orm.base import Base + + db_path = tmp_path_factory.mktemp("sexp") / "test.db" + engine = create_async_engine(f"sqlite+aiosqlite:///{db_path}") + async with engine.begin() as conn: + await conn.run_sync(Base.metadata.create_all) + + local = async_sessionmaker(engine, expire_on_commit=False) + + @asynccontextmanager + async def _ctx(): + async with local() as session: + try: + yield session + finally: + await session.close() + + yield _ctx + await engine.dispose() + + +@pytest_asyncio.fixture(scope="module", loop_scope="module") +async def org(session_maker): + from mirix.orm.organization import Organization as OrganizationORM + + org_id = f"sexp-org-{uuid.uuid4().hex[:8]}" + async with session_maker() as session: + session.add(OrganizationORM(id=org_id, name=org_id)) + await session.commit() + return type("Org", (), {"id": org_id})() + + +@pytest_asyncio.fixture(scope="module", loop_scope="module") +async def user_a(session_maker, org): + from mirix.orm.user import User as UserORM + + uid = f"sexp-user-{uuid.uuid4().hex[:8]}" + async with session_maker() as session: + session.add( + UserORM(id=uid, name=uid, organization_id=org.id, + status="active", timezone="UTC") + ) + await session.commit() + return type("User", (), {"id": uid, "organization_id": org.id})() + + +async def _insert_agent(session_maker, org_id: str) -> str: + from mirix.orm.agent import Agent as AgentORM + + agent_id = f"agent-{uuid.uuid4()}" + async with session_maker() as session: + session.add(AgentORM(id=agent_id, organization_id=org_id)) + await session.commit() + return agent_id + + +def _manager(session_maker): + mgr = SkillExperienceManager() + mgr.session_maker = session_maker + return mgr + + +def _exp_kwargs(idx, etype="worth_learning", importance=0.6, credibility=0.6, **ov): + base = dict( + session_id=f"sess-{idx}", + experience_type=etype, + title=f"title-{idx}", + content=f"content-{idx}", + importance=importance, + credibility=credibility, + evidence='{"quote":"q","signal_type":"inferred"}', + ) + base.update(ov) + return base + + +@pytest.mark.integration +@pytest.mark.asyncio(loop_scope="module") +class TestCreateAndList: + async def test_create_and_read_back(self, session_maker, org, user_a): + mgr = _manager(session_maker) + agent_id = await _insert_agent(session_maker, org.id) + rec = await mgr.create_experience( + agent_id=agent_id, user_id=user_a.id, + organization_id=user_a.organization_id, **_exp_kwargs(1), + ) + assert rec.id.startswith("sexp-") + assert rec.status == "pending" + assert rec.consumed_by is None + assert rec.experience_type == "worth_learning" + + async def test_rejects_bad_experience_type(self, session_maker, org, user_a): + mgr = _manager(session_maker) + agent_id = await _insert_agent(session_maker, org.id) + with pytest.raises(ValidationError): + await mgr.create_experience( + agent_id=agent_id, user_id=user_a.id, + organization_id=user_a.organization_id, + **_exp_kwargs(1, etype="partial"), + ) + listed = await mgr.list_experiences(agent_id=agent_id) + assert listed == [] + + async def test_clamps_at_persistence(self, session_maker, org, user_a): + mgr = _manager(session_maker) + agent_id = await _insert_agent(session_maker, org.id) + rec = await mgr.create_experience( + agent_id=agent_id, user_id=user_a.id, + organization_id=user_a.organization_id, + **_exp_kwargs(1, importance=9.0, credibility=-2.0), + ) + assert rec.importance == 1.0 + assert rec.credibility == 0.0 + + async def test_list_ordered_by_priority_desc(self, session_maker, org, user_a): + mgr = _manager(session_maker) + agent_id = await _insert_agent(session_maker, org.id) + owner = dict(agent_id=agent_id, user_id=user_a.id, + organization_id=user_a.organization_id) + # priorities: low=0.09, high=0.72, mid=0.30 + await mgr.create_experience(**owner, **_exp_kwargs("low", importance=0.9, credibility=0.1)) + await mgr.create_experience(**owner, **_exp_kwargs("high", importance=0.8, credibility=0.9)) + await mgr.create_experience(**owner, **_exp_kwargs("mid", importance=0.5, credibility=0.6)) + listed = await mgr.list_experiences(agent_id=agent_id) + sessions = [e.session_id for e in listed] + assert sessions == ["sess-high", "sess-mid", "sess-low"] + + async def test_isolated_per_agent(self, session_maker, org, user_a): + mgr = _manager(session_maker) + a1 = await _insert_agent(session_maker, org.id) + a2 = await _insert_agent(session_maker, org.id) + await mgr.create_experience( + agent_id=a1, user_id=user_a.id, + organization_id=user_a.organization_id, **_exp_kwargs(1), + ) + await mgr.create_experience( + agent_id=a2, user_id=user_a.id, + organization_id=user_a.organization_id, **_exp_kwargs(2), + ) + listed = await mgr.list_experiences(agent_id=a2) + assert len(listed) == 1 + assert all(e.agent_id == a2 for e in listed) + + async def test_respects_limit(self, session_maker, org, user_a): + mgr = _manager(session_maker) + agent_id = await _insert_agent(session_maker, org.id) + owner = dict(agent_id=agent_id, user_id=user_a.id, + organization_id=user_a.organization_id) + for i in range(4): + await mgr.create_experience(**owner, **_exp_kwargs(i)) + listed = await mgr.list_experiences(agent_id=agent_id, limit=2) + assert len(listed) == 2 + + async def test_ids_scopes_to_given_set(self, session_maker, org, user_a): + # The Goal-3 curator passes THIS round's freshly-distilled ids so an + # evolve only sees its own batch, never the whole pending pool. + mgr = _manager(session_maker) + agent_id = await _insert_agent(session_maker, org.id) + owner = dict(agent_id=agent_id, user_id=user_a.id, + organization_id=user_a.organization_id) + a = await mgr.create_experience(**owner, **_exp_kwargs("a", importance=0.9, credibility=0.9)) + await mgr.create_experience(**owner, **_exp_kwargs("b", importance=0.8, credibility=0.8)) + c = await mgr.create_experience(**owner, **_exp_kwargs("c", importance=0.7, credibility=0.7)) + # Scope to a + c only; b must stay invisible. + scoped = await mgr.list_experiences(agent_id=agent_id, ids=[a.id, c.id]) + # Priority-ordered WITHIN the scope (a=0.81 before c=0.49). + assert [e.id for e in scoped] == [a.id, c.id] + + async def test_empty_ids_returns_nothing(self, session_maker, org, user_a): + mgr = _manager(session_maker) + agent_id = await _insert_agent(session_maker, org.id) + await mgr.create_experience( + agent_id=agent_id, user_id=user_a.id, + organization_id=user_a.organization_id, **_exp_kwargs(1), + ) + # ids=[] is an explicit "scope to nothing" (distinct from ids=None). + assert await mgr.list_experiences(agent_id=agent_id, ids=[]) == [] + # ids=None (default) still returns the row. + assert len(await mgr.list_experiences(agent_id=agent_id)) == 1 + + +@pytest.mark.integration +@pytest.mark.asyncio(loop_scope="module") +class TestMarkTransitions: + async def test_mark_consumed_with_lineage(self, session_maker, org, user_a): + mgr = _manager(session_maker) + agent_id = await _insert_agent(session_maker, org.id) + owner = dict(agent_id=agent_id, user_id=user_a.id, + organization_id=user_a.organization_id) + r1 = await mgr.create_experience(**owner, **_exp_kwargs(1)) + r2 = await mgr.create_experience(**owner, **_exp_kwargs(2)) + n = await mgr.mark_consumed( + ids=[r1.id, r2.id], run_id="xprun-1", + influenced_skill_ids=["skill-a", "skill-b"], + ) + assert n == 2 + # No longer pending. + assert await mgr.list_experiences(agent_id=agent_id) == [] + # Lineage + consumed_by persisted. + consumed = await mgr.list_experiences(agent_id=agent_id, status="consumed") + assert len(consumed) == 2 + assert all(c.consumed_by == "xprun-1" for c in consumed) + assert all(set(c.influenced_skill_ids) == {"skill-a", "skill-b"} for c in consumed) + + async def test_mark_consumed_idempotent(self, session_maker, org, user_a): + mgr = _manager(session_maker) + agent_id = await _insert_agent(session_maker, org.id) + owner = dict(agent_id=agent_id, user_id=user_a.id, + organization_id=user_a.organization_id) + r1 = await mgr.create_experience(**owner, **_exp_kwargs(1)) + first = await mgr.mark_consumed(ids=[r1.id], run_id="run-A") + assert first == 1 + second = await mgr.mark_consumed(ids=[r1.id], run_id="run-B") + assert second == 0 # already consumed; consumed_by not clobbered + consumed = await mgr.list_experiences(agent_id=agent_id, status="consumed") + assert consumed[0].consumed_by == "run-A" + + async def test_mark_consumed_empty_noop(self, session_maker): + mgr = _manager(session_maker) + assert await mgr.mark_consumed(ids=[], run_id="run-0") == 0 + + async def test_mark_superseded(self, session_maker, org, user_a): + mgr = _manager(session_maker) + agent_id = await _insert_agent(session_maker, org.id) + owner = dict(agent_id=agent_id, user_id=user_a.id, + organization_id=user_a.organization_id) + r1 = await mgr.create_experience(**owner, **_exp_kwargs(1)) + n = await mgr.mark_superseded(ids=[r1.id]) + assert n == 1 + assert await mgr.list_experiences(agent_id=agent_id) == [] + sup = await mgr.list_experiences(agent_id=agent_id, status="superseded") + assert len(sup) == 1 + + +@pytest.mark.integration +@pytest.mark.asyncio(loop_scope="module") +class TestAggregate: + async def test_counts_by_type_and_priority(self, session_maker, org, user_a): + mgr = _manager(session_maker) + agent_id = await _insert_agent(session_maker, org.id) + owner = dict(agent_id=agent_id, user_id=user_a.id, + organization_id=user_a.organization_id) + a = await mgr.create_experience( + **owner, **_exp_kwargs(1, etype="worth_avoiding", importance=1.0, credibility=1.0) + ) + b = await mgr.create_experience( + **owner, **_exp_kwargs(2, etype="worth_avoiding", importance=0.5, credibility=0.5) + ) + c = await mgr.create_experience( + **owner, **_exp_kwargs(3, etype="worth_learning", importance=0.4, credibility=0.5) + ) + agg = await mgr.aggregate(ids=[a.id, b.id, c.id]) + assert agg["n"] == 3 + assert agg["n_worth_avoiding"] == 2 + assert agg["n_worth_learning"] == 1 + # 1.0 + 0.25 + 0.20 = 1.45 + assert agg["sum_priority"] == pytest.approx(1.45) + + async def test_empty_ids(self, session_maker): + mgr = _manager(session_maker) + agg = await mgr.aggregate(ids=[]) + assert agg == {"n": 0, "n_worth_learning": 0, + "n_worth_avoiding": 0, "sum_priority": 0.0} + + +# ====================== End-to-end-ish distill -> consume ================== + + +@pytest.mark.integration +@pytest.mark.asyncio(loop_scope="module") +class TestDistillThenConsume: + async def test_distilled_experiences_are_consumable(self, session_maker, org, user_a): + """Persist via the distiller's _persist_experiences (stubbed LLM upstream), + then drive the Goal-3 core with mock snapshot/step to consume them.""" + import json as _json + + from mirix.services.session_experience_distiller import ( + SessionExperienceDistiller, + ) + from mirix.services.skill_experience_curator import ( + _run_experience_evolution_core, + ) + + mgr = _manager(session_maker) + agent_id = await _insert_agent(session_maker, org.id) + meta = type("M", (), {"id": agent_id})() + user = type("U", (), {"id": user_a.id})() + actor = type("A", (), {"organization_id": user_a.organization_id})() + + d = SessionExperienceDistiller(experience_manager=mgr) + parsed = [ + {"experience_type": "worth_avoiding", "title": "no missing column", + "importance": 0.9, "credibility": 0.9, + "evidence": {"quote": "no such column", "signal_type": "tool_error"}}, + {"experience_type": "worth_learning", "title": "batch calls", + "importance": 0.7, "credibility": 0.8, + "evidence": {"quote": "great", "signal_type": "user_confirmation"}}, + ] + created = await d._persist_experiences( + meta_agent_state=meta, user=user, actor=actor, + session_id="sess-e2e", parsed=parsed, + ) + assert len(created) == 2 + + # Now consume via the Goal-3 core (mock agent/step; real manager+DB). + agent = type("Ag", (), {})() + + async def snapshot(): + return [] + + stepped = {} + + async def run_step(a, payload, budget): + stepped["payload"] = payload + stepped["budget"] = budget + + result = await _run_experience_evolution_core( + experience_manager=mgr, agent=agent, agent_id="proc-e2e", + meta_agent_id=agent_id, user_id=user_a.id, + snapshot_skills=snapshot, run_step=run_step, + ) + assert result["skipped"] is False + assert result["consumed_count"] == 2 + # Payload carries both kinds, priority-ordered (avoid 0.81 before learn 0.56). + assert stepped["payload"].index("[AVOID]") < stepped["payload"].index("[LEARN]") + # All pending consumed. + assert await mgr.list_experiences(agent_id=agent_id, status="pending") == [] + consumed = await mgr.list_experiences(agent_id=agent_id, status="consumed") + assert {c.consumed_by for c in consumed} == {result["run_id"]} + _ = _json # silence unused if branch above changes + + async def test_overflow_beyond_cap_superseded_no_stranding_in_db( + self, session_maker, org, user_a + ): + """A round distilling MORE than the per-run cap leaves NOTHING pending in + the DB: the top _MAX_EXPERIENCES_PER_RUN are consumed and the lowest- + priority tail is superseded (real end-state, not just a fake call).""" + from mirix.services.skill_experience_curator import ( + _MAX_EXPERIENCES_PER_RUN, + _run_experience_evolution_core, + ) + + mgr = _manager(session_maker) + agent_id = await _insert_agent(session_maker, org.id) + owner = dict(agent_id=agent_id, user_id=user_a.id, + organization_id=user_a.organization_id) + n = _MAX_EXPERIENCES_PER_RUN + 1 + created = [] + for i in range(n): + # Strictly descending priority so the LAST one is the overflow tail. + rec = await mgr.create_experience( + **owner, + **_exp_kwargs(i, etype="worth_avoiding", + importance=1.0, credibility=(n - i) / n), + ) + created.append(rec) + all_ids = [r.id for r in created] + + agent = type("Ag", (), {})() + + async def snapshot(): + return [] + + async def run_step(a, payload, budget): + pass + + result = await _run_experience_evolution_core( + experience_manager=mgr, agent=agent, agent_id="proc-overflow", + meta_agent_id=agent_id, user_id=user_a.id, + snapshot_skills=snapshot, run_step=run_step, + experience_ids=all_ids, + ) + assert result["consumed_count"] == _MAX_EXPERIENCES_PER_RUN + assert result["superseded_count"] == 1 + # The load-bearing assertion: NOTHING from this round is left stranded. + assert await mgr.list_experiences( + agent_id=agent_id, status="pending", ids=all_ids, limit=n + ) == [] + consumed = await mgr.list_experiences( + agent_id=agent_id, status="consumed", ids=all_ids, limit=n + ) + superseded = await mgr.list_experiences( + agent_id=agent_id, status="superseded", ids=all_ids, limit=n + ) + assert len(consumed) == _MAX_EXPERIENCES_PER_RUN + assert len(superseded) == 1 + + +@pytest.mark.integration +@pytest.mark.asyncio(loop_scope="module") +class TestErasure: + """User/client erasure: experiences are distilled from verbatim user + conversations, so the irreversible purge paths must cover them.""" + + async def test_delete_by_user_id_removes_all_statuses_for_that_user_only( + self, session_maker, org, user_a + ): + from mirix.orm.user import User as UserORM + + mgr = _manager(session_maker) + agent_id = await _insert_agent(session_maker, org.id) + other_uid = f"sexp-user-{uuid.uuid4().hex[:8]}" + async with session_maker() as session: + session.add( + UserORM(id=other_uid, name=other_uid, organization_id=org.id, + status="active", timezone="UTC") + ) + await session.commit() + + kept = await mgr.create_experience( + agent_id=agent_id, user_id=other_uid, + organization_id=org.id, **_exp_kwargs("erase-keep"), + ) + target = await mgr.create_experience( + agent_id=agent_id, user_id=user_a.id, + organization_id=org.id, **_exp_kwargs("erase-pending"), + ) + consumed = await mgr.create_experience( + agent_id=agent_id, user_id=user_a.id, + organization_id=org.id, **_exp_kwargs("erase-consumed"), + ) + await mgr.mark_consumed(ids=[consumed.id], run_id="run-x") + + deleted = await mgr.delete_by_user_id(user_id=user_a.id) + + assert deleted >= 2 # pending AND consumed rows both erased + for status in ("pending", "consumed", "superseded"): + assert await mgr.list_experiences( + agent_id=agent_id, user_id=user_a.id, status=status, limit=50 + ) == [] + # The other user's experience is untouched. + assert [e.id for e in await mgr.list_experiences( + agent_id=agent_id, user_id=other_uid, status="pending", limit=50 + )] == [kept.id] + + async def test_delete_by_client_id_uses_creation_attribution( + self, session_maker, org, user_a + ): + from mirix.schemas.client import Client as PydanticClient + + mgr = _manager(session_maker) + agent_id = await _insert_agent(session_maker, org.id) + client_a = PydanticClient( + id=f"sexp-client-{uuid.uuid4().hex[:8]}", organization_id=org.id, + name="A", write_scope="t", read_scopes=["t"], + ) + client_b = PydanticClient( + id=f"sexp-client-{uuid.uuid4().hex[:8]}", organization_id=org.id, + name="B", write_scope="t", read_scopes=["t"], + ) + mine = await mgr.create_experience( + agent_id=agent_id, user_id=user_a.id, organization_id=org.id, + created_by_id=client_a.id, **_exp_kwargs("client-a"), + ) + theirs = await mgr.create_experience( + agent_id=agent_id, user_id=user_a.id, organization_id=org.id, + created_by_id=client_b.id, **_exp_kwargs("client-b"), + ) + + deleted = await mgr.delete_by_client_id(actor=client_a) + + assert deleted == 1 + remaining = {e.id for e in await mgr.list_experiences( + agent_id=agent_id, user_id=user_a.id, status="pending", limit=50 + )} + assert mine.id not in remaining + assert theirs.id in remaining diff --git a/tests/test_skill_orm.py b/tests/test_skill_orm.py new file mode 100644 index 000000000..1ae126865 --- /dev/null +++ b/tests/test_skill_orm.py @@ -0,0 +1,33 @@ +"""Tests for the skill-based ProceduralMemoryItem ORM model.""" + +import pytest +from sqlalchemy import inspect as sa_inspect + +from mirix.orm.procedural_memory import ProceduralMemoryItem + + +class TestProceduralMemoryORM: + """Verify the ORM model has all new skill-based columns and no old columns.""" + + def _column_names(self): + mapper = sa_inspect(ProceduralMemoryItem) + return {col.key for col in mapper.columns} + + def test_has_new_columns(self): + cols = self._column_names() + for expected in [ + "name", + "triggers", + "examples", + "version", + "description", + "instructions", + "description_embedding", + "instructions_embedding", + ]: + assert expected in cols, f"Missing expected column: {expected}" + + def test_no_old_columns(self): + cols = self._column_names() + for removed in ["summary", "steps", "summary_embedding", "steps_embedding"]: + assert removed not in cols, f"Old column still present: {removed}" diff --git a/tests/test_skill_schema.py b/tests/test_skill_schema.py new file mode 100644 index 000000000..5932fc104 --- /dev/null +++ b/tests/test_skill_schema.py @@ -0,0 +1,192 @@ +"""Tests for the skill-based procedural memory schema.""" + +import pytest +from pydantic import ValidationError + +from mirix.schemas.procedural_memory import ( + ProceduralMemoryItem, + ProceduralMemoryItemBase, + ProceduralMemoryItemUpdate, + ProceduralMemoryItemResponse, + normalize_entry_type, +) + + +class TestEntryTypeNormalization: + """Legacy free-form entry_type values must never crash reads. + + The schema validator runs on every to_pydantic() (from_attributes reads), + so rows written before the skill schema ("process", "How-To", ...) are + normalized into the closed set instead of raising. Strict validation of + new values lives in the write path (tool_validators). + """ + + @pytest.mark.parametrize( + "legacy_value,expected", + [ + ("workflow", "workflow"), + ("Guide", "guide"), + (" SCRIPT ", "script"), + ("process", "workflow"), + ("how-to", "guide"), + ("howto guide", "guide"), + ("shell script", "script"), + ("procedure", "workflow"), + ("", "workflow"), + (None, "workflow"), + ], + ) + def test_normalize_entry_type(self, legacy_value, expected): + assert normalize_entry_type(legacy_value) == expected + + def test_legacy_entry_type_does_not_crash_schema_read(self): + """Simulates reading a pre-skill row whose entry_type is free-form.""" + item = ProceduralMemoryItemBase( + name="legacy-skill", + entry_type="process", + description="Written before the closed entry_type set existed", + instructions="Step 1: do the thing", + ) + assert item.entry_type == "workflow" + + +class TestProceduralMemoryItemBase: + """Tests for ProceduralMemoryItemBase skill schema.""" + + def test_name_is_required(self): + """name field must be provided; omitting it raises a ValidationError.""" + with pytest.raises(ValidationError): + ProceduralMemoryItemBase( + entry_type="workflow", + description="Deploy the app", + instructions="Step 1: build\nStep 2: deploy", + ) + + def test_valid_skill_creation(self): + """All required fields produce a valid skill object.""" + item = ProceduralMemoryItemBase( + name="deploy-production", + entry_type="workflow", + description="Deploy the production application", + instructions="Step 1: build the image\nStep 2: push to registry\nStep 3: deploy", + ) + assert item.name == "deploy-production" + assert item.entry_type == "workflow" + assert item.description == "Deploy the production application" + assert isinstance(item.instructions, str) + + def test_triggers_defaults_to_empty_list(self): + """triggers should default to an empty list.""" + item = ProceduralMemoryItemBase( + name="deploy-production", + entry_type="workflow", + description="Deploy the production application", + instructions="Step 1: build", + ) + assert item.triggers == [] + + def test_examples_defaults_to_empty_list(self): + """examples should default to an empty list.""" + item = ProceduralMemoryItemBase( + name="deploy-production", + entry_type="workflow", + description="Deploy the production application", + instructions="Step 1: build", + ) + assert item.examples == [] + + def test_triggers_and_examples_can_be_set(self): + """triggers and examples can be explicitly provided.""" + item = ProceduralMemoryItemBase( + name="deploy-production", + entry_type="workflow", + description="Deploy the production application", + instructions="Step 1: build", + triggers=["user says deploy", "CI pipeline triggers"], + examples=[{"input": "deploy now", "output": "deployed v1.2.3"}], + ) + assert len(item.triggers) == 2 + assert len(item.examples) == 1 + + +class TestProceduralMemoryItem: + """Tests for the full ProceduralMemoryItem with DB fields.""" + + def test_version_defaults_to_0_1_0(self): + """version should default to '0.1.0'.""" + item = ProceduralMemoryItem( + name="deploy-production", + entry_type="workflow", + description="Deploy the app", + instructions="Step 1: build", + user_id="user-123", + organization_id="org-456", + ) + assert item.version == "0.1.0" + + def test_description_embedding_defaults_to_none(self): + """description_embedding should be Optional and default to None.""" + item = ProceduralMemoryItem( + name="deploy-production", + entry_type="workflow", + description="Deploy the app", + instructions="Step 1: build", + user_id="user-123", + organization_id="org-456", + ) + assert item.description_embedding is None + + def test_instructions_embedding_defaults_to_none(self): + """instructions_embedding should be Optional and default to None.""" + item = ProceduralMemoryItem( + name="deploy-production", + entry_type="workflow", + description="Deploy the app", + instructions="Step 1: build", + user_id="user-123", + organization_id="org-456", + ) + assert item.instructions_embedding is None + + def test_version_can_be_set(self): + """version can be explicitly provided.""" + item = ProceduralMemoryItem( + name="deploy-production", + entry_type="workflow", + description="Deploy the app", + instructions="Step 1: build", + user_id="user-123", + organization_id="org-456", + version="1.2.3", + ) + assert item.version == "1.2.3" + + +class TestProceduralMemoryItemUpdate: + """Tests for the update schema.""" + + def test_update_has_renamed_fields(self): + """Update schema uses description/instructions instead of summary/steps.""" + update = ProceduralMemoryItemUpdate( + id="proc_item-abc12345", + description="Updated description", + instructions="Updated instructions", + ) + assert update.description == "Updated description" + assert update.instructions == "Updated instructions" + + def test_update_has_name_field(self): + """Update schema includes optional name field.""" + update = ProceduralMemoryItemUpdate( + id="proc_item-abc12345", + name="new-skill-name", + ) + assert update.name == "new-skill-name" + + +class TestProceduralMemoryItemResponse: + """Tests for the response schema.""" + + def test_response_inherits_from_item(self): + """Response schema should be a subclass of ProceduralMemoryItem.""" + assert issubclass(ProceduralMemoryItemResponse, ProceduralMemoryItem)