diff --git a/engine.py b/engine.py index 0f157d4e..03e3464f 100644 --- a/engine.py +++ b/engine.py @@ -64,9 +64,6 @@ def __init__(self, config: LCMConfig | None = None, # genuinely new messages (appended after compaction) get ingested. self._ingest_cursor: int = 0 - # Wire tool handlers - lcm_tools.set_engine(self) - # State required by ContextEngine ABC and run_agent.py compatibility self.model = "" self.base_url = "" @@ -263,7 +260,7 @@ def handle_tool_call(self, name: str, args: Dict[str, Any], **kwargs) -> str: } handler = handlers.get(name) if handler: - return handler(args) + return handler(args, engine=self) return json.dumps({"error": f"Unknown LCM tool: {name}"}) def get_status(self) -> Dict[str, Any]: diff --git a/tests/test_lcm_engine.py b/tests/test_lcm_engine.py index 86219a69..93d4fef1 100644 --- a/tests/test_lcm_engine.py +++ b/tests/test_lcm_engine.py @@ -5,6 +5,7 @@ from agent.context_engine import ContextEngine from hermes_lcm.config import LCMConfig +from hermes_lcm.dag import SummaryNode from hermes_lcm.engine import LCMEngine @@ -387,3 +388,62 @@ def test_handle_describe_overview(self, engine): def test_handle_unknown_tool(self, engine): result = json.loads(engine.handle_tool_call("unknown_tool", {})) assert "error" in result + + def test_tool_dispatch_is_bound_to_engine_instance(self, tmp_path): + config_a = LCMConfig(database_path=str(tmp_path / "a.db")) + config_b = LCMConfig(database_path=str(tmp_path / "b.db")) + + engine_a = LCMEngine(config=config_a) + engine_a._session_id = "session-a" + engine_b = LCMEngine(config=config_b) + engine_b._session_id = "session-b" + + engine_a._store.append("session-a", {"role": "user", "content": "alpha project"}) + engine_b._store.append("session-b", {"role": "user", "content": "beta project"}) + + result_a = json.loads(engine_a.handle_tool_call("lcm_grep", {"query": "alpha"})) + result_b = json.loads(engine_b.handle_tool_call("lcm_grep", {"query": "beta"})) + + assert result_a["total_results"] == 1 + assert result_b["total_results"] == 1 + assert "alpha" in result_a["results"][0]["snippet"] + assert "beta" in result_b["results"][0]["snippet"] + + def test_describe_and_expand_are_session_scoped(self, engine): + node_id = engine._dag.add_node( + SummaryNode( + session_id="session-a", + depth=0, + summary="secret summary", + token_count=10, + source_token_count=20, + source_ids=[], + source_type="messages", + created_at=0, + ) + ) + + engine._session_id = "session-b" + + describe = json.loads(engine.handle_tool_call("lcm_describe", {"node_id": node_id})) + expand = json.loads(engine.handle_tool_call("lcm_expand", {"node_id": node_id})) + + assert "error" in describe + assert "error" in expand + + def test_describe_overview_includes_sparse_high_depth_nodes(self, engine): + engine._dag.add_node( + SummaryNode( + session_id="test-session", + depth=2, + summary="durable summary", + token_count=100, + source_token_count=500, + source_ids=[], + source_type="messages", + created_at=0, + ) + ) + + overview = json.loads(engine.handle_tool_call("lcm_describe", {})) + assert "d2" in overview["depths"] diff --git a/tools.py b/tools.py index 19748486..1659f11d 100644 --- a/tools.py +++ b/tools.py @@ -1,5 +1,7 @@ """Tool handlers for LCM — the code that runs when the LLM calls each tool.""" +from __future__ import annotations + import json import logging from typing import Any, Dict, TYPE_CHECKING @@ -9,18 +11,23 @@ logger = logging.getLogger(__name__) -# The engine instance is set by the engine on init -_engine: "LCMEngine | None" = None + +def _require_engine(kwargs: Dict[str, Any]) -> "LCMEngine | None": + engine = kwargs.get("engine") + return engine if engine is not None else None -def set_engine(engine: "LCMEngine") -> None: - global _engine - _engine = engine +def _get_session_node(engine: "LCMEngine", node_id: int): + node = engine._dag.get_node(node_id) + if node is None or node.session_id != engine._session_id: + return None + return node def lcm_grep(args: Dict[str, Any], **kwargs) -> str: - """Search across the full DAG and raw messages.""" - if _engine is None: + """Search across the full DAG and raw messages for the current session.""" + engine = _require_engine(kwargs) + if engine is None: return json.dumps({"error": "LCM engine not initialized"}) query = args.get("query", "").strip() @@ -28,153 +35,156 @@ def lcm_grep(args: Dict[str, Any], **kwargs) -> str: return json.dumps({"error": "No query provided"}) limit = args.get("limit", 10) - session_id = _engine._session_id - + session_id = engine._session_id results = [] - # Search raw messages try: - msg_hits = _engine._store.search(query, session_id=session_id, limit=limit) + msg_hits = engine._store.search(query, session_id=session_id, limit=limit) for hit in msg_hits: - results.append({ - "type": "message", - "depth": "raw", - "store_id": hit["store_id"], - "role": hit["role"], - "snippet": hit.get("snippet", hit.get("content", "")[:200]), - }) - except Exception as e: - logger.debug("Message search failed: %s", e) - - # Search summary nodes + results.append( + { + "type": "message", + "depth": "raw", + "store_id": hit["store_id"], + "role": hit["role"], + "snippet": hit.get("snippet", hit.get("content", "")[:200]), + } + ) + except Exception as exc: + logger.debug("Message search failed: %s", exc) + try: - node_hits = _engine._dag.search(query, session_id=session_id, limit=limit) + node_hits = engine._dag.search(query, session_id=session_id, limit=limit) for node in node_hits: - results.append({ - "type": "summary", - "depth": f"d{node.depth}", - "node_id": node.node_id, - "snippet": node.summary[:300], - "token_count": node.token_count, - "expand_hint": node.expand_hint, - }) - except Exception as e: - logger.debug("Node search failed: %s", e) - - # Sort by relevance (raw messages first, then by depth ascending) - results.sort(key=lambda r: (0 if r["type"] == "message" else 1, r.get("depth", ""))) - - return json.dumps({ - "query": query, - "total_results": len(results), - "results": results[:limit], - }) + results.append( + { + "type": "summary", + "depth": f"d{node.depth}", + "node_id": node.node_id, + "snippet": node.summary[:300], + "token_count": node.token_count, + "expand_hint": node.expand_hint, + } + ) + except Exception as exc: + logger.debug("Node search failed: %s", exc) + + results.sort(key=lambda result: (0 if result["type"] == "message" else 1, result.get("depth", ""))) + return json.dumps({"query": query, "total_results": len(results), "results": results[:limit]}) def lcm_describe(args: Dict[str, Any], **kwargs) -> str: """Inspect a node's subtree or get session DAG overview.""" - if _engine is None: + engine = _require_engine(kwargs) + if engine is None: return json.dumps({"error": "LCM engine not initialized"}) node_id = args.get("node_id") - session_id = _engine._session_id + session_id = engine._session_id if node_id is not None: - info = _engine._dag.describe_subtree(node_id) + node = _get_session_node(engine, node_id) + if node is None: + return json.dumps({"error": f"Node {node_id} not found in current session"}) + info = engine._dag.describe_subtree(node_id) return json.dumps(info) - # Session overview: count nodes at each depth + all_nodes = engine._dag.get_session_nodes(session_id) overview = { "session_id": session_id, - "store_message_count": _engine._store.get_session_count(session_id), + "store_message_count": engine._store.get_session_count(session_id), "depths": {}, } - for depth in range(10): # check up to d9 - count = _engine._dag.count_at_depth(session_id, depth) - if count == 0 and depth > 0: - break - if count > 0: - nodes = _engine._dag.get_session_nodes(session_id, depth=depth, limit=100) - overview["depths"][f"d{depth}"] = { - "count": count, - "total_tokens": sum(n.token_count for n in nodes), - "total_source_tokens": sum(n.source_token_count for n in nodes), - "nodes": [ - { - "node_id": n.node_id, - "token_count": n.token_count, - "expand_hint": n.expand_hint, - } - for n in nodes[:20] # cap at 20 per depth for display - ], - } + for depth in sorted({node.depth for node in all_nodes}): + nodes = [node for node in all_nodes if node.depth == depth] + overview["depths"][f"d{depth}"] = { + "count": len(nodes), + "total_tokens": sum(node.token_count for node in nodes), + "total_source_tokens": sum(node.source_token_count for node in nodes), + "nodes": [ + { + "node_id": node.node_id, + "token_count": node.token_count, + "expand_hint": node.expand_hint, + } + for node in nodes[:20] + ], + } return json.dumps(overview) def lcm_expand(args: Dict[str, Any], **kwargs) -> str: """Expand a summary node to its source content.""" - if _engine is None: + engine = _require_engine(kwargs) + if engine is None: return json.dumps({"error": "LCM engine not initialized"}) node_id = args.get("node_id") if node_id is None: return json.dumps({"error": "node_id is required"}) - max_tokens = args.get("max_tokens", 4000) + node = _get_session_node(engine, node_id) + if node is None: + return json.dumps({"error": f"Node {node_id} not found in current session"}) - node = _engine._dag.get_node(node_id) - if not node: - return json.dumps({"error": f"Node {node_id} not found"}) + max_tokens = args.get("max_tokens", 4000) if node.source_type == "messages": - # Expand to raw messages - messages = [] from .tokens import count_tokens + + messages = [] budget_used = 0 - for sid in node.source_ids: - stored = _engine._store.get(sid) - if not stored: + for store_id in node.source_ids: + stored = engine._store.get(store_id) + if not stored or stored.get("session_id") != engine._session_id: continue content = stored.get("content", "") msg_tokens = count_tokens(content) if budget_used + msg_tokens > max_tokens and messages: - messages.append({ - "note": f"Truncated — {len(node.source_ids) - len(messages)} more messages available", - }) + messages.append( + { + "note": f"Truncated — {len(node.source_ids) - len(messages)} more messages available", + } + ) break - messages.append({ - "store_id": stored["store_id"], - "role": stored["role"], - "content": content[:2000] if len(content) > 2000 else content, - }) - budget_used += msg_tokens - - return json.dumps({ - "node_id": node_id, - "depth": node.depth, - "source_type": "messages", - "expanded": messages, - }) - - elif node.source_type == "nodes": - # Expand to child summaries - children = _engine._dag.get_source_nodes(node) - return json.dumps({ - "node_id": node_id, - "depth": node.depth, - "source_type": "nodes", - "expanded": [ + messages.append( { - "node_id": c.node_id, - "depth": c.depth, - "summary": c.summary[:1000], - "token_count": c.token_count, - "expand_hint": c.expand_hint, + "store_id": stored["store_id"], + "role": stored["role"], + "content": content[:2000] if len(content) > 2000 else content, } - for c in children - ], - }) + ) + budget_used += msg_tokens + + return json.dumps( + { + "node_id": node_id, + "depth": node.depth, + "source_type": "messages", + "expanded": messages, + } + ) + + if node.source_type == "nodes": + children = [child for child in engine._dag.get_source_nodes(node) if child.session_id == engine._session_id] + return json.dumps( + { + "node_id": node_id, + "depth": node.depth, + "source_type": "nodes", + "expanded": [ + { + "node_id": child.node_id, + "depth": child.depth, + "summary": child.summary[:1000], + "token_count": child.token_count, + "expand_hint": child.expand_hint, + } + for child in children + ], + } + ) return json.dumps({"error": f"Unknown source_type: {node.source_type}"})