diff --git a/.gitattributes b/.gitattributes new file mode 100644 index 0000000..1d3d194 --- /dev/null +++ b/.gitattributes @@ -0,0 +1,8 @@ +# Keep launcher and installer scripts executable in WSL and readable by +# PowerShell. Without this, a Windows checkout can pass CRLF into the bash +# shebang and WSL reports `/usr/bin/env: bash\r: No such file or directory`. +*.sh text eol=lf +*.ps1 text eol=crlf +*.py text eol=lf +*.json text eol=lf +*.toml text eol=lf diff --git a/datamind/agent/options.py b/datamind/agent/options.py index 806d403..67e08be 100644 --- a/datamind/agent/options.py +++ b/datamind/agent/options.py @@ -13,24 +13,21 @@ from pathlib import Path from typing import Any -from datamind.capabilities.db import DBService, build_db_service, build_db_tools +from datamind.capabilities.db import DBService, build_db_service from datamind.capabilities.embedding import build_embedding -from datamind.capabilities.graph import GraphService, build_graph_service, build_graph_tools +from datamind.capabilities.graph import GraphService, build_graph_service from datamind.capabilities.hooks import AuditLogHook, DestructiveSqlHook, PathAllowlistHook from datamind.capabilities.ingest import ( IngestLedger, IngestService, build_ingest_service, - build_ingest_tools, with_receipts, ) -from datamind.capabilities.kb import KBService, build_kb_service, build_kb_tools -from datamind.capabilities.memory import MemoryService, build_memory_service, build_memory_tools +from datamind.capabilities.kb import KBService, build_kb_service +from datamind.capabilities.memory import MemoryService, build_memory_service from datamind.capabilities.skills import ( SkillsService, build_skills_service, - build_skills_store_tools, - build_skills_tools, ) from datamind.config import Settings from datamind.core.contracts import ToolAccess @@ -40,7 +37,7 @@ from datamind.core.context import RequestContext from datamind.core.model_clients import build_model_client from datamind.core.protocols import EmbeddingProvider, TextModelClient, ToolCallingModelClient -from datamind.core.tools import ToolRegistry +from datamind.core.tools import ToolRegistry, tool_provider_registry from .base import AgentLoopConfig, AgentLoopProtocol from .loop_native import NativeAgentLoop @@ -63,6 +60,29 @@ class AgentServices: skills: SkillsService | None = None memory: MemoryService | None = None ingest: IngestService | None = None + _closed: bool = False + + async def aclose(self) -> None: + """Close shared runtime resources once, regardless of facade owner.""" + if self._closed: + return + self._closed = True + resources = [ + self.db, + self.graph, + self.kb, + self.embedding, + self.fallback_client, + self.client, + ] + seen: set[int] = set() + for resource in resources: + if resource is None or id(resource) in seen: + continue + seen.add(id(resource)) + close = getattr(resource, "aclose", None) + if callable(close): + await close() @dataclass @@ -140,6 +160,9 @@ async def query( user_message=message, history=history, final_contract=final_contract, ) + async def aclose(self) -> None: + await self.services.aclose() + @dataclass class StoreAgent: @@ -167,6 +190,9 @@ async def store( ) -> dict[str, Any]: return await self.loop.run_turn(user_message=message, history=history) + async def aclose(self) -> None: + await self.services.aclose() + @dataclass class DataMind: @@ -176,7 +202,6 @@ class DataMind: retrieve_agent: RetrieveAgent services: AgentServices profile: str = "default" - _closed: bool = False @property def store(self) -> StoreAgent: @@ -224,25 +249,7 @@ async def query( async def aclose(self) -> None: """Idempotently close clients, engines, stores, and providers.""" - if self._closed: - return - self._closed = True - resources = [ - self.services.db, - self.services.graph, - self.services.kb, - self.services.embedding, - self.services.fallback_client, - self.services.client, - ] - seen: set[int] = set() - for resource in resources: - if resource is None or id(resource) in seen: - continue - seen.add(id(resource)) - close = getattr(resource, "aclose", None) - if callable(close): - await close() + await self.services.aclose() def _build_hook_chain(settings: Settings) -> HookChain | None: @@ -367,27 +374,22 @@ async def build_datamind( ) catalogue = ToolRegistry() - ingest_tools = build_ingest_tools(ingest) if ingest is not None else [] - catalogue.extend([t for t in ingest_tools if t.surface is None]) - if "kb" in active: - assert kb is not None - catalogue.extend(build_kb_tools(kb)) - catalogue.extend([t for t in ingest_tools if t.surface and t.surface.value == "kb"]) - if "db" in active: - assert db is not None - catalogue.extend(build_db_tools(db)) - catalogue.extend([t for t in ingest_tools if t.surface and t.surface.value == "db"]) - if "graph" in active: - assert graph is not None - catalogue.extend(build_graph_tools(graph)) - catalogue.extend([t for t in ingest_tools if t.surface and t.surface.value == "graph"]) - if "skills" in active: - assert skills is not None - catalogue.extend(build_skills_tools(skills)) - catalogue.extend(build_skills_store_tools(skills)) - if "memory" in active: - assert memory is not None - catalogue.extend(build_memory_tools(memory)) + provider_services = { + "kb_service": kb, + "db_service": db, + "graph_service": graph, + "skills_service": skills, + "memory_service": memory, + "ingest_service": ingest, + } + if ingest is not None: + ingest_provider = tool_provider_registry.create("ingest") + for spec in ingest_provider.build(**provider_services): + if spec.surface is None or spec.surface.value in active: + catalogue.add(spec) + for provider_name in sorted(active): + provider = tool_provider_registry.create(provider_name) + catalogue.extend(provider.build(**provider_services)) retrieve_tools = catalogue.select(access={ToolAccess.READ, ToolAccess.UTILITY}) raw_store_tools = catalogue.select(access={ToolAccess.WRITE}) diff --git a/datamind/capabilities/ingest/ledger.py b/datamind/capabilities/ingest/ledger.py index 594d07a..b08cbe0 100644 --- a/datamind/capabilities/ingest/ledger.py +++ b/datamind/capabilities/ingest/ledger.py @@ -5,6 +5,8 @@ import hashlib import inspect import json +import os +import time from pathlib import Path from typing import Any, Awaitable, Callable @@ -116,9 +118,46 @@ def __init__(self, *, storage_dir: Path, profile: str) -> None: self._profile = profile self._state_path = self._storage_dir / "ingest_state.json" self._receipts_path = self._storage_dir / "ingest_receipts.jsonl" + self._process_lock_path = self._storage_dir / "ingest.lock" self._lock = asyncio.Lock() self._storage_dir.mkdir(parents=True, exist_ok=True) + async def _acquire_process_lock(self) -> None: + """Serialize ledger read/execute/write across MCP processes. + + The in-memory asyncio lock only protects one DataMind instance. An + atomic lock-file claim extends the same idempotency boundary to a + second MCP process or a separately-created Ledger object. + """ + while True: + try: + fd = os.open( + self._process_lock_path, + os.O_CREAT | os.O_EXCL | os.O_WRONLY, + ) + try: + os.write(fd, f"{os.getpid()} {time.time():.6f}\n".encode()) + finally: + os.close(fd) + return + except FileExistsError: + try: + age = time.time() - self._process_lock_path.stat().st_mtime + if age > 3600: + self._process_lock_path.unlink(missing_ok=True) + continue + except FileNotFoundError: + continue + await asyncio.sleep(0.02) + + async def _release_process_lock(self) -> None: + try: + self._process_lock_path.unlink(missing_ok=True) + except OSError: + # The lock may have been recovered as stale by another process; + # the receipt/state writes remain the source of truth. + pass + def _load_state(self) -> dict[str, Any]: if not self._state_path.is_file(): return {"revision": 0, "successful": {}} @@ -150,7 +189,21 @@ async def execute( args: dict[str, Any], invoke: Callable[..., Awaitable[Any]], ) -> dict[str, Any]: - """Deduplicate retryable imports, but always apply explicit DB replacement.""" + """Deduplicate retryable imports, including across MCP processes.""" + await self._acquire_process_lock() + try: + return await self._execute_locked(spec=spec, args=args, invoke=invoke) + finally: + await self._release_process_lock() + + async def _execute_locked( + self, + *, + spec: ToolSpec, + args: dict[str, Any], + invoke: Callable[..., Awaitable[Any]], + ) -> dict[str, Any]: + """Run one call while both ledger locks are held.""" source = _source_from_call(spec, args) fingerprint = _hash_text( _canonical( diff --git a/datamind/capabilities/ingest/service.py b/datamind/capabilities/ingest/service.py index b04387e..43d2d5f 100644 --- a/datamind/capabilities/ingest/service.py +++ b/datamind/capabilities/ingest/service.py @@ -820,6 +820,18 @@ async def _upsert_chunks( store = self._kb.vector_store texts = [c.text for c in chunks] vectors = await provider.embed_texts(texts) + # Keep track of IDs that existed before this call. If a third-party + # store fails after partially accepting a batch, only remove newly + # created IDs during rollback; an idempotent re-write must not erase + # the previous revision. + existing_ids: set[str] = set() + try: + existing_ids = { + str(chunk_id) + for chunk_id, _text, _metadata in await store.get_all_texts() + } + except Exception: # noqa: BLE001 - rollback is best effort + pass stale_ids: list[str] = [] if replace_sources: new_ids = {chunk.id for chunk in chunks} @@ -830,15 +842,28 @@ async def _upsert_chunks( } if recorded_sources & replace_sources and chunk_id not in new_ids: stale_ids.append(chunk_id) - await store.add( - ids=[c.id for c in chunks], - texts=texts, - embeddings=vectors, - metadatas=[ - {**(c.metadata or {}), "source": c.source or ""} - for c in chunks - ], - ) + try: + await store.add( + ids=[c.id for c in chunks], + texts=texts, + embeddings=vectors, + metadatas=[ + {**(c.metadata or {}), "source": c.source or ""} + for c in chunks + ], + ) + except Exception: + rollback_ids = [c.id for c in chunks if c.id not in existing_ids] + delete = getattr(store, "delete", None) + if rollback_ids and callable(delete): + try: + await delete(rollback_ids) + except Exception: # noqa: BLE001 - preserve original write error + _log.warning( + "kb_partial_write_rollback_failed", + extra={"chunk_ids": rollback_ids}, + ) + raise if stale_ids: await store.delete(stale_ids) # Incremental writes must leave the same compatibility metadata that diff --git a/datamind/capabilities/ingest/tools.py b/datamind/capabilities/ingest/tools.py index c9b4d26..b2784ae 100644 --- a/datamind/capabilities/ingest/tools.py +++ b/datamind/capabilities/ingest/tools.py @@ -15,7 +15,7 @@ """ from __future__ import annotations -from datamind.core.tools import ToolSpec +from datamind.core.tools import ToolSpec, tool_provider_registry from .service import IngestService @@ -448,4 +448,13 @@ async def _graph_add_path( ] +@tool_provider_registry.register("ingest") +class _IngestToolProvider: + def build(self, **services: object) -> list[ToolSpec]: + ingest = services.get("ingest_service") + if not isinstance(ingest, IngestService): + raise ValueError("ingest tool provider requires 'ingest_service'") + return build_ingest_tools(ingest) + + __all__ = ["build_ingest_tools"] diff --git a/datamind/capabilities/kb/indexer.py b/datamind/capabilities/kb/indexer.py index 07f754c..3b5df42 100644 --- a/datamind/capabilities/kb/indexer.py +++ b/datamind/capabilities/kb/indexer.py @@ -74,6 +74,10 @@ def _split_text( they fall inside the size window. Not fancy; deterministic; good enough for the default. """ + if chunk_size <= 0: + raise ValueError("chunk_size must be > 0") + if chunk_overlap < 0: + raise ValueError("chunk_overlap must be >= 0") if chunk_overlap >= chunk_size: raise ValueError("chunk_overlap must be < chunk_size") if len(text) <= chunk_size: diff --git a/datamind/capabilities/kb/providers/chroma_store.py b/datamind/capabilities/kb/providers/chroma_store.py index 689374b..a5a753b 100644 --- a/datamind/capabilities/kb/providers/chroma_store.py +++ b/datamind/capabilities/kb/providers/chroma_store.py @@ -6,6 +6,7 @@ from pathlib import Path from typing import Any, Sequence +from datamind.core.errors import CapabilityError from datamind.core.logging import get_logger from datamind.core.protocols import RetrievedChunk from datamind.core.registry import vector_store_registry @@ -24,26 +25,39 @@ def __init__( collection_name: str, dimension: int, ) -> None: - import chromadb # type: ignore + try: + import chromadb # type: ignore + except Exception as exc: # noqa: BLE001 + raise CapabilityError("kb", "Chroma dependency is unavailable", cause=exc) from exc self.dimension = dimension self._persist_dir = Path(persist_dir) self._persist_dir.mkdir(parents=True, exist_ok=True) self._collection_name = collection_name - self._client = chromadb.PersistentClient(path=str(self._persist_dir)) - # We supply our own embeddings — disable the default model download. - self._collection = self._client.get_or_create_collection( - name=collection_name, - embedding_function=None, # type: ignore[arg-type] - metadata={"hnsw:space": "cosine"}, - ) - self.existing_count = int(self._collection.count()) + try: + self._client = chromadb.PersistentClient(path=str(self._persist_dir)) + # We supply our own embeddings — disable the default model download. + self._collection = self._client.get_or_create_collection( + name=collection_name, + embedding_function=None, # type: ignore[arg-type] + metadata={"hnsw:space": "cosine"}, + ) + except Exception as exc: # noqa: BLE001 + raise CapabilityError( + "kb", f"cannot open Chroma vector index at {self._persist_dir}: {exc}", cause=exc, + ) from exc + try: + self.existing_count = int(self._collection.count()) + except Exception as exc: # noqa: BLE001 + raise CapabilityError( + "kb", f"cannot inspect Chroma vector index at {self._persist_dir}: {exc}", cause=exc, + ) from exc _log.info( "chroma_collection_ready", extra={ "collection": collection_name, "path": str(self._persist_dir), - "count": self._collection.count(), + "count": self.existing_count, }, ) @@ -77,13 +91,16 @@ async def query( top_k: int = 5, where: dict[str, Any] | None = None, ) -> list[RetrievedChunk]: - result = await asyncio.to_thread( - self._collection.query, - query_embeddings=[list(embedding)], - n_results=top_k, - where=where, - include=["documents", "metadatas", "distances"], - ) + try: + result = await asyncio.to_thread( + self._collection.query, + query_embeddings=[list(embedding)], + n_results=top_k, + where=where, + include=["documents", "metadatas", "distances"], + ) + except Exception as exc: # noqa: BLE001 + raise CapabilityError("kb", f"vector index query failed: {exc}", cause=exc) from exc ids = (result.get("ids") or [[]])[0] docs = (result.get("documents") or [[]])[0] metas = (result.get("metadatas") or [[]])[0] diff --git a/datamind/capabilities/kb/providers/hybrid_retriever.py b/datamind/capabilities/kb/providers/hybrid_retriever.py index b82447d..d4482d6 100644 --- a/datamind/capabilities/kb/providers/hybrid_retriever.py +++ b/datamind/capabilities/kb/providers/hybrid_retriever.py @@ -99,6 +99,10 @@ async def aretrieve( top_k: int = 5, filters: dict[str, Any] | None = None, ) -> list[RetrievedChunk]: + if top_k <= 0: + raise ValueError("top_k must be > 0") + if not query or not query.strip(): + raise ValueError("query must not be empty") validate_metadata_filter(filters) await self._ensure_lexical() k_inner = top_k * self._cm diff --git a/datamind/capabilities/kb/providers/simple_retriever.py b/datamind/capabilities/kb/providers/simple_retriever.py index 7e81757..3652459 100644 --- a/datamind/capabilities/kb/providers/simple_retriever.py +++ b/datamind/capabilities/kb/providers/simple_retriever.py @@ -34,6 +34,10 @@ async def aretrieve( top_k: int = 5, filters: dict[str, Any] | None = None, ) -> list[RetrievedChunk]: + if top_k <= 0: + raise ValueError("top_k must be > 0") + if not query or not query.strip(): + raise ValueError("query must not be empty") validate_metadata_filter(filters) vec = await self._embed.embed_query(query) chunks = await self._store.query(vec, top_k=top_k, where=filters) diff --git a/datamind/capabilities/kb/service.py b/datamind/capabilities/kb/service.py index 487ad86..ecfd8cb 100644 --- a/datamind/capabilities/kb/service.py +++ b/datamind/capabilities/kb/service.py @@ -82,7 +82,11 @@ async def search( ) -> list[dict[str, Any]]: if self._compatibility_error: raise ConfigError(self._compatibility_error) - k = top_k or self.retrieval_cfg.top_k + k = self.retrieval_cfg.top_k if top_k is None else top_k + if k <= 0: + raise ValueError("top_k must be > 0") + if not query or not query.strip(): + raise ValueError("query must not be empty") chunks = await self.retriever.aretrieve(query, top_k=k, filters=filters) return [c.model_dump() for c in chunks] diff --git a/datamind/capabilities/memory/providers/sqlite_store.py b/datamind/capabilities/memory/providers/sqlite_store.py index b9cd08e..a249d9f 100644 --- a/datamind/capabilities/memory/providers/sqlite_store.py +++ b/datamind/capabilities/memory/providers/sqlite_store.py @@ -37,6 +37,7 @@ import asyncio import json import math +import numbers import sqlite3 import struct import time @@ -79,6 +80,29 @@ def _cosine(a: list[float], b: list[float]) -> float: return dot / (na * nb) +def _validate_vector( + vector: Sequence[float], *, expected_dimension: int | None, label: str, +) -> list[float]: + """Reject malformed vectors before they can poison memory recall.""" + try: + values = list(vector) + except TypeError as exc: + raise CapabilityError("memory", f"{label} embedding is not a sequence", cause=exc) from exc + if expected_dimension and len(values) != expected_dimension: + raise CapabilityError( + "memory", + f"{label} embedding dimension mismatch: expected {expected_dimension}, got {len(values)}", + ) + if any( + isinstance(value, bool) + or not isinstance(value, numbers.Real) + or not math.isfinite(float(value)) + for value in values + ): + raise CapabilityError("memory", f"{label} embedding contains non-finite or non-numeric values") + return [float(value) for value in values] + + def _row_to_item(row: tuple, *, score: float = 0.0) -> MemoryItem: rid, scope, profile, session_id, kind, status, content, meta_json, _emb, created_at, _updated, _archived = row try: @@ -239,17 +263,39 @@ async def save( kind: Kind = "fact", metadata: dict[str, Any] | None = None, ) -> str: + if not isinstance(content, str) or not content.strip(): + raise CapabilityError("memory", "content must be a non-empty string") + if scope not in {"global", "profile", "session"}: + raise CapabilityError("memory", f"unsupported scope: {scope!r}") + if kind not in {"preference", "decision", "workflow", "summary", "skill", "fact"}: + raise CapabilityError("memory", f"unsupported kind: {kind!r}") if scope == "profile" and not profile: raise CapabilityError("memory", "scope='profile' requires profile= argument") if scope == "session" and not session_id: raise CapabilityError("memory", "scope='session' requires session_id= argument") + if scope == "global" and (profile or session_id): + raise CapabilityError("memory", "scope='global' cannot include profile or session_id") + if scope == "profile" and session_id: + raise CapabilityError("memory", "scope='profile' cannot include session_id") + if scope == "session" and profile: + raise CapabilityError("memory", "scope='session' cannot include profile") + if metadata is not None and not isinstance(metadata, dict): + raise CapabilityError("memory", "metadata must be an object") item_id = uuid.uuid4().hex ts = time.time() - meta_json = json.dumps(metadata or {}, ensure_ascii=False) + try: + meta_json = json.dumps(metadata or {}, ensure_ascii=False) + except (TypeError, ValueError) as exc: + raise CapabilityError("memory", "metadata must be JSON-serializable", cause=exc) from exc emb: bytes | None = None if self._embedding is not None and content.strip(): vec = await self._embedding.embed_query(content) + vec = _validate_vector( + vec, + expected_dimension=int(getattr(self._embedding, "dimension", 0) or 0) or None, + label="memory", + ) emb = _pack(vec) def _run() -> None: @@ -291,15 +337,37 @@ async def recall( # advanced — let callers tune the per-scope budget for ablation per_scope: dict[str, int] | None = None, ) -> list[MemoryItem]: + if not isinstance(query, str) or not query.strip(): + raise CapabilityError("memory", "query must be a non-empty string") + if isinstance(top_k, bool) or not isinstance(top_k, int) or top_k <= 0: + raise CapabilityError("memory", "top_k must be a positive integer") + allowed_kinds = {"preference", "decision", "workflow", "summary", "skill", "fact"} + if kinds is not None: + if not isinstance(kinds, (list, tuple)) or any(kind not in allowed_kinds for kind in kinds): + raise CapabilityError("memory", "kinds must contain only supported memory kinds") # Default per-scope budget: 2 (session) + 4 (profile) + 2 (global) = 8. budgets: dict[str, int] = {"session": 2, "profile": 4, "global": 2} if per_scope: - budgets.update({k: v for k, v in per_scope.items() if k in budgets}) + if not isinstance(per_scope, dict): + raise CapabilityError("memory", "per_scope must be an object") + for key, value in per_scope.items(): + if key not in budgets: + raise CapabilityError("memory", f"unsupported memory scope budget: {key!r}") + if isinstance(value, bool) or not isinstance(value, int) or value < 0: + raise CapabilityError("memory", f"per_scope[{key!r}] must be a non-negative integer") + budgets[key] = value # If no embedding, fall back to lexical scoring within the same # scope filters so the contract stays identical. scoring = self._score_with_embedding if self._embedding else self._score_lexical - qv = await self._embedding.embed_query(query) if self._embedding else None + qv = None + if self._embedding: + raw_qv = await self._embedding.embed_query(query) + qv = _validate_vector( + raw_qv, + expected_dimension=int(getattr(self._embedding, "dimension", 0) or 0) or None, + label="query", + ) merged: list[MemoryItem] = [] seen: set[str] = set() @@ -470,7 +538,16 @@ def _score_with_embedding( out: list[tuple[float, MemoryItem]] = [] for row in rows: emb_blob = row[8] - score = _cosine(qvec, _unpack(emb_blob)) if (qvec and emb_blob) else 0.0 + if qvec and emb_blob: + stored = _unpack(emb_blob) + if len(stored) != len(qvec): + raise CapabilityError( + "memory", + f"stored embedding dimension mismatch: expected {len(qvec)}, got {len(stored)}", + ) + score = _cosine(qvec, stored) + else: + score = 0.0 out.append((score, _row_to_item(row, score=score))) return out diff --git a/datamind/capabilities/memory/short_term.py b/datamind/capabilities/memory/short_term.py index 0d51f54..79b2396 100644 --- a/datamind/capabilities/memory/short_term.py +++ b/datamind/capabilities/memory/short_term.py @@ -38,6 +38,8 @@ class ShortTermMemory: """Session-keyed rolling window of Turns.""" def __init__(self, *, max_turns: int = 20) -> None: + if isinstance(max_turns, bool) or not isinstance(max_turns, int) or max_turns <= 0: + raise ValueError("max_turns must be a positive integer") self._max = max_turns self._buffers: dict[str, deque[Turn]] = {} self._locks: dict[str, asyncio.Lock] = {} @@ -50,16 +52,28 @@ def _lock(self, session_id: str) -> asyncio.Lock: return lock async def append(self, session_id: str, role: str, content: str, **metadata: Any) -> None: + if not isinstance(session_id, str) or not session_id.strip(): + raise ValueError("session_id must be a non-empty string") + if role not in {"user", "assistant", "system"}: + raise ValueError("role must be user, assistant, or system") + if not isinstance(content, str) or not content.strip(): + raise ValueError("content must be a non-empty string") async with self._lock(session_id): buf = self._buffers.setdefault(session_id, deque(maxlen=self._max)) buf.append(Turn(role=role, content=content, metadata=dict(metadata))) async def recent(self, session_id: str, *, limit: int | None = None) -> list[Turn]: + if not isinstance(session_id, str) or not session_id.strip(): + raise ValueError("session_id must be a non-empty string") + if limit is not None and (isinstance(limit, bool) or not isinstance(limit, int) or limit < 0): + raise ValueError("limit must be a non-negative integer") async with self._lock(session_id): buf = self._buffers.get(session_id) if not buf: return [] items = list(buf) + if limit == 0: + return [] if limit is not None: items = items[-limit:] return items diff --git a/datamind/capabilities/skills/loader.py b/datamind/capabilities/skills/loader.py index 159fbf3..d2f998f 100644 --- a/datamind/capabilities/skills/loader.py +++ b/datamind/capabilities/skills/loader.py @@ -82,10 +82,23 @@ def load_skill(path: Path) -> SkillManifest | None: except OSError as exc: _log.warning("skill_read_failed", extra={"path": str(path), "err": str(exc)}) return None + has_frontmatter = bool(_FRONTMATTER_RE.match(text)) + if text.lstrip().startswith("---") and not has_frontmatter: + _log.warning("skill_manifest_invalid", extra={"path": str(path), "reason": "malformed frontmatter"}) + return None meta, body = _parse_frontmatter(text) - name = str(meta.get("name") or path.parent.name).strip() - desc = str(meta.get("description") or "").strip() - if not desc: + if has_frontmatter: + name = str(meta.get("name") or "").strip() + desc = str(meta.get("description") or "").strip() + if not name or not desc: + _log.warning( + "skill_manifest_invalid", + extra={"path": str(path), "reason": "name and description are required"}, + ) + return None + else: + name = path.parent.name.strip() + desc = "" # Fall back to the first non-empty Markdown line so the skill is at # least searchable. for line in body.splitlines(): @@ -93,6 +106,12 @@ def load_skill(path: Path) -> SkillManifest | None: if s: desc = s break + if not re.fullmatch(r"[a-z0-9][a-z0-9_-]{0,63}", name): + _log.warning("skill_manifest_invalid", extra={"path": str(path), "reason": "invalid name"}) + return None + if not body.strip(): + _log.warning("skill_manifest_invalid", extra={"path": str(path), "reason": "empty body"}) + return None keywords_raw = meta.get("keywords") if isinstance(keywords_raw, list): keywords = tuple(str(k) for k in keywords_raw) diff --git a/datamind/capabilities/skills/service.py b/datamind/capabilities/skills/service.py index 5e8cea5..0b48a37 100644 --- a/datamind/capabilities/skills/service.py +++ b/datamind/capabilities/skills/service.py @@ -9,6 +9,8 @@ from __future__ import annotations import json +import math +import numbers import re from pathlib import Path from typing import Any @@ -71,6 +73,30 @@ async def load(self) -> dict[str, int]: for m in manifests ] vectors = await self._embedding.embed_texts(texts) + expected_dimension = int(getattr(self._embedding, "dimension", 0) or 0) + if len(vectors) != len(manifests): + raise CapabilityError( + "skills", + f"embedding count mismatch: expected {len(manifests)}, got {len(vectors)}", + ) + for index, vector in enumerate(vectors): + if not isinstance(vector, (list, tuple)) or not vector: + raise CapabilityError("skills", f"embedding {index} must be a non-empty vector") + if expected_dimension and len(vector) != expected_dimension: + raise CapabilityError( + "skills", + f"embedding dimension mismatch at item {index}: " + f"expected {expected_dimension}, got {len(vector)}", + ) + if any( + isinstance(value, bool) + or not isinstance(value, numbers.Real) + or not math.isfinite(float(value)) + for value in vector + ): + raise CapabilityError( + "skills", f"embedding {index} contains non-finite or non-numeric values", + ) await self._store.reset() await self._store.add(ids=ids, texts=texts, embeddings=vectors, metadatas=metas) indexed = len(ids) @@ -116,6 +142,10 @@ def get(self, name: str) -> dict[str, Any]: } async def search(self, query: str, *, top_k: int = 3) -> list[dict[str, Any]]: + if not isinstance(query, str) or not query.strip(): + raise CapabilityError("skills", "query must be a non-empty string") + if isinstance(top_k, bool) or not isinstance(top_k, int) or top_k <= 0: + raise CapabilityError("skills", "top_k must be a positive integer") if not self._embedding or not self._store: return [] vec = await self._embedding.embed_query(query) @@ -141,6 +171,14 @@ async def upsert( overwrite: bool = True, ) -> dict[str, Any]: """Write a profile-scoped SKILL.md and refresh the live index.""" + if not isinstance(name, str): + raise CapabilityError("skills", "skill name must be a string") + if not isinstance(description, str): + raise CapabilityError("skills", "skill description must be a string") + if not isinstance(body, str): + raise CapabilityError("skills", "skill body must be a string") + if keywords is not None and not isinstance(keywords, (list, tuple)): + raise CapabilityError("skills", "skill keywords must be an array") normalized = name.strip().lower() if not re.fullmatch(r"[a-z0-9][a-z0-9_-]{0,63}", normalized): raise CapabilityError( @@ -159,9 +197,9 @@ async def upsert( raise CapabilityError("skills", f"skill '{normalized}' already exists") target_dir.mkdir(parents=True, exist_ok=True) - safe_keywords = [ - str(k).strip() for k in (keywords or []) if str(k).strip() - ] + if any(not isinstance(keyword, str) for keyword in (keywords or [])): + raise CapabilityError("skills", "skill keywords must contain strings") + safe_keywords = [keyword.strip() for keyword in (keywords or []) if keyword.strip()] keyword_json = json.dumps(safe_keywords, ensure_ascii=False) description_json = json.dumps(description.strip(), ensure_ascii=False) text = ( diff --git a/datamind/capabilities/skills/tools.py b/datamind/capabilities/skills/tools.py index ed720f5..1e29d4a 100644 --- a/datamind/capabilities/skills/tools.py +++ b/datamind/capabilities/skills/tools.py @@ -115,7 +115,7 @@ def build(self, **services: Any) -> list[ToolSpec]: s = services.get("skills_service") if s is None: raise ValueError("skills tool provider requires 'skills_service'") - return build_skills_tools(s) + return build_skills_tools(s) + build_skills_store_tools(s) __all__ = ["build_skills_tools", "build_skills_store_tools"] diff --git a/datamind/tests/test_agent_loop.py b/datamind/tests/test_agent_loop.py index b4c5662..aa9d804 100644 --- a/datamind/tests/test_agent_loop.py +++ b/datamind/tests/test_agent_loop.py @@ -258,6 +258,32 @@ async def test_tool_error_is_surfaced_as_tool_result(): assert "RuntimeError" in error_blocks[0]["content"] +def test_tool_context_truncation_keeps_count_and_evidence_locators(): + registry = ToolRegistry() + registry.add(_tool_echo()) + loop = NativeAgentLoop( + client=_FakeClient([]), + tools=registry, + config=AgentLoopConfig(model="m", max_tool_result_rows=1, max_tool_result_chars=200), + ) + result = { + "results": [ + {"id": "c1", "source": "doc.md", "text": "第一段", "score": 0.9}, + {"id": "c2", "source": "doc.md", "text": "第二段", "score": 0.8}, + ] + } + block = loop._tool_result_block("tool-1", result, None) + trace, evidence = loop._trace_and_evidence( + name="echo", tool_input={"query": "q"}, result=result, error=None, + ) + + assert block["_datamind_truncated"] is True + assert block["_datamind_total_count"] == 2 + assert '"truncated": true' in block["content"] + assert [item["locator"]["chunk_id"] for item in evidence] == ["c1", "c2"] + assert trace["result_size_chars"] > 0 + + @pytest.mark.asyncio async def test_max_tool_turns_enforced(): # Script: infinite tool_use loops. diff --git a/datamind/tests/test_ingest_revisions.py b/datamind/tests/test_ingest_revisions.py index bcf4ed7..1bda3e0 100644 --- a/datamind/tests/test_ingest_revisions.py +++ b/datamind/tests/test_ingest_revisions.py @@ -3,6 +3,9 @@ import pytest from datamind.capabilities.ingest.service import IngestService +from datamind.capabilities.ingest.ledger import IngestLedger, with_receipts +from datamind.capabilities.ingest.tools import build_ingest_tools +from datamind.core.tools import ToolRegistry class _Embedding: @@ -29,6 +32,17 @@ async def get_all_texts(self): ] +class _PartiallyFailingStore(_VectorStore): + async def add(self, ids, texts, embeddings, metadatas=None): + first = True + for chunk_id, text, metadata in zip(ids, texts, metadatas or []): + if first: + self.items[chunk_id] = (text, dict(metadata)) + first = False + continue + raise RuntimeError("simulated chunk write failure") + + class _KB: def __init__(self): self.embedding = _Embedding() @@ -38,6 +52,12 @@ async def record_incremental_ingest(self): return None +class _FailingKB(_KB): + def __init__(self): + self.embedding = _Embedding() + self.vector_store = _PartiallyFailingStore() + + class _Model: async def generate_text(self, prompt, **kwargs): return "[]" @@ -73,3 +93,27 @@ async def test_reingesting_same_path_replaces_old_kb_chunks(tmp_path: Path): assert [item[0] for item in kb.vector_store.items.values()] == [ "负责人:周宁\n验收日期:2026年12月2日" ] + + +@pytest.mark.asyncio +async def test_partial_chunk_failure_returns_failed_receipt_without_new_chunks(tmp_path: Path): + profile = tmp_path / "profile" + profile.mkdir() + service = IngestService( + kb=_FailingKB(), db=None, graph=None, llm_client=_Model(), llm_model="test", + profile_data_dir=profile, chunk_size=4, chunk_overlap=0, + allowed_roots=[tmp_path], + ) + raw = ToolRegistry() + raw.extend(build_ingest_tools(service)) + tools = with_receipts( + raw, + IngestLedger(storage_dir=tmp_path / "ledger", profile="test"), + ) + + receipt = await tools.get("kb_add_text").handler( + text="abcdefgh", source="partial.txt", persist=False, + ) + + assert receipt["results"][0]["status"] == "failed" + assert service._kb.vector_store.items == {} diff --git a/datamind/tests/test_rag_failure_contracts.py b/datamind/tests/test_rag_failure_contracts.py new file mode 100644 index 0000000..a9f4f10 --- /dev/null +++ b/datamind/tests/test_rag_failure_contracts.py @@ -0,0 +1,25 @@ +from __future__ import annotations + +import pytest + +from datamind.capabilities.kb.providers.chroma_store import ChromaVectorStore +from datamind.core.errors import CapabilityError + + +@pytest.mark.asyncio +async def test_empty_vector_index_is_a_valid_empty_result(tmp_path): + store = ChromaVectorStore( + persist_dir=tmp_path, collection_name="empty-index", dimension=2, + ) + + assert await store.count() == 0 + assert await store.query([1.0, 0.0], top_k=5) == [] + + +def test_corrupt_vector_index_is_not_reported_as_empty(tmp_path): + (tmp_path / "chroma.sqlite3").write_bytes(b"not a sqlite database") + + with pytest.raises(CapabilityError, match="cannot open Chroma vector index"): + ChromaVectorStore( + persist_dir=tmp_path, collection_name="broken-index", dimension=2, + ) diff --git a/datamind/tests/test_replace_receipts.py b/datamind/tests/test_replace_receipts.py index ee4b2d6..c0f95c0 100644 --- a/datamind/tests/test_replace_receipts.py +++ b/datamind/tests/test_replace_receipts.py @@ -1,4 +1,5 @@ """A historical receipt cannot establish the current contents of a table.""" +import asyncio from pathlib import Path import pytest @@ -9,7 +10,7 @@ from datamind.capabilities.ingest.service import IngestService from datamind.capabilities.ingest.tools import build_ingest_tools from datamind.config import DBConfig -from datamind.core.tools import ToolRegistry +from datamind.core.tools import ToolRegistry, ToolSpec @pytest.fixture @@ -77,3 +78,33 @@ async def test_append_retry_remains_deduplicated(database, tmp_path): assert second["results"][0]["status"] == "unchanged" assert second["revision"] == first["revision"] assert (await db.query_sql("SELECT amount FROM sales")).rows == [["200"]] + + +@pytest.mark.asyncio +async def test_independent_ledgers_deduplicate_concurrent_same_call(tmp_path): + calls = 0 + + async def handler(*, text): + nonlocal calls + calls += 1 + await asyncio.sleep(0.05) + return {"chunks_added": 1} + + raw = ToolRegistry() + raw.add(ToolSpec( + name="kb_add_text", + description="test", + input_schema={"type": "object"}, + handler=handler, + metadata={"surface": "kb", "access": "write"}, + )) + first = with_receipts(raw, IngestLedger(storage_dir=tmp_path / "ledger", profile="test")) + second = with_receipts(raw, IngestLedger(storage_dir=tmp_path / "ledger", profile="test")) + + results = await asyncio.gather( + first.get("kb_add_text").handler(text="same payload"), + second.get("kb_add_text").handler(text="same payload"), + ) + + assert calls == 1 + assert sorted(item["results"][0]["status"] for item in results) == ["stored", "unchanged"] diff --git a/plugins/datamind-context/src/datamind_mcp.py b/plugins/datamind-context/src/datamind_mcp.py index 1f653bc..d63cabd 100755 --- a/plugins/datamind-context/src/datamind_mcp.py +++ b/plugins/datamind-context/src/datamind_mcp.py @@ -121,18 +121,40 @@ def repo_root() -> Path: # MCP hosts use these hints to distinguish harmless inspection from operations # that need confirmation. Keep write-capable tools unannotated so Codex retains # its normal approval gate for ingestion and mutation. -READ_ONLY_TOOLS = { - "datamind_raw_file_read", - "datamind_build_status", - "datamind_workspace_inspect", - "datamind_ask", - "datamind_build_verify", - "datamind_rag_query", - "datamind_graph_query", - "datamind_list_profiles", - "datamind_status", +class McpRoute: + __slots__ = ("surfaces", "read_only") + + def __init__(self, surfaces: frozenset[str] | None, read_only: bool = False) -> None: + self.surfaces = surfaces + self.read_only = read_only + + +_ALL_SURFACES: frozenset[str] | None = None +ROUTES: dict[str, McpRoute] = { + "datamind_raw_file_read": McpRoute(frozenset({"graph"}), True), + "datamind_build_status": McpRoute(frozenset({"graph"}), True), + "datamind_workspace_inspect": McpRoute(frozenset({"graph"}), True), + "datamind_ask": McpRoute(_ALL_SURFACES, True), + "datamind_store": McpRoute(_ALL_SURFACES), + "datamind_use_folder": McpRoute(frozenset({"kb", "graph"})), + "datamind_graph_ingest": McpRoute(frozenset({"graph"})), + "datamind_graph_build_lineage": McpRoute(frozenset({"graph"})), + "datamind_table_ingest": McpRoute(frozenset({"db"})), + "datamind_build_start": McpRoute(frozenset({"graph"})), + "datamind_build_freeze": McpRoute(frozenset({"graph"})), + "datamind_build_verify": McpRoute(frozenset({"graph"}), True), + "datamind_build_export": McpRoute(frozenset({"graph"})), + "datamind_surface_ingest": McpRoute(_ALL_SURFACES), + "datamind_rag_query": McpRoute(frozenset({"kb"}), True), + "datamind_graph_query": McpRoute(frozenset({"graph"}), True), + "datamind_remember": McpRoute(frozenset({"memory"})), + "datamind_list_profiles": McpRoute(frozenset(), True), + "datamind_status": McpRoute(_ALL_SURFACES, True), } +if set(ROUTES) != set(TOOLS): + raise RuntimeError("MCP route table and tool catalogue are out of sync") + def profile_name(args: dict[str, Any]) -> str: value = str(args.get("profile") or "default").strip() allowed = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789_.-" @@ -143,105 +165,149 @@ def profile_name(args: dict[str, Any]) -> str: def enabled_surfaces(name: str) -> set[str] | None: """Build only the capability services required by one MCP tool.""" - if name in {"datamind_raw_file_read", "datamind_workspace_inspect", - "datamind_build_status", "datamind_build_start", - "datamind_build_freeze", "datamind_build_verify", - "datamind_build_export"}: - # Generic ingest/build tools need an ingest service but no data - # surface. Graph is the lightest surface because it does not - # initialise an embedding provider. - return {"graph"} - if name == "datamind_rag_query": - return {"kb"} - if name == "datamind_use_folder": - return {"kb", "graph"} - if name == "datamind_table_ingest": - return {"db"} - if name in {"datamind_graph_query", "datamind_graph_ingest", - "datamind_graph_build_lineage"}: - return {"graph"} - if name == "datamind_remember": - return {"memory"} - if name == "datamind_surface_ingest": - # Routing may target any combination supplied at call time. - return None - # Agent-level ask/store and status intentionally expose/warm all services. - return None + try: + surfaces = ROUTES[name].surfaces + except KeyError as exc: + raise ValueError(f"unknown tool: {name}") from exc + return None if surfaces is None else set(surfaces) + + +class RuntimeEntry: + __slots__ = ("system", "settings", "warmup") + + def __init__(self, *, system: Any, settings: Any, warmup: dict[str, Any]) -> None: + self.system = system + self.settings = settings + self.warmup = warmup + + +class RuntimeFactory: + """Create one warmed DataMind runtime per profile and surface set.""" + + def __init__(self) -> None: + self._entries: dict[tuple[str, tuple[str, ...] | None], RuntimeEntry] = {} + self._locks: dict[tuple[str, tuple[str, ...] | None], asyncio.Lock] = {} + self._closed = False + @staticmethod + def _key(profile: str, surfaces: set[str] | None) -> tuple[str, tuple[str, ...] | None]: + return (profile, None if surfaces is None else tuple(sorted(surfaces))) -async def execute(name: str, args: dict[str, Any]) -> dict[str, Any]: - from datamind.agent import build_datamind - from datamind.config import Settings + async def get(self, profile: str, surfaces: set[str] | None) -> RuntimeEntry: + if self._closed: + raise RuntimeError("DataMind runtime factory is closed") + key = self._key(profile, surfaces) + existing = self._entries.get(key) + if existing is not None: + return existing + lock = self._locks.setdefault(key, asyncio.Lock()) + async with lock: + existing = self._entries.get(key) + if existing is not None: + return existing + from datamind.agent import build_datamind + from datamind.config import Settings + + settings = Settings() + settings.data.profile = profile + system = await build_datamind(settings, enable=surfaces) + try: + warmup = await system.warmup() + except BaseException: + await system.aclose() + raise + entry = RuntimeEntry(system=system, settings=settings, warmup=warmup) + self._entries[key] = entry + return entry + + async def aclose(self) -> None: + if self._closed: + return + self._closed = True + entries = list(self._entries.values()) + self._entries.clear() + self._locks.clear() + for entry in entries: + await entry.system.aclose() + + +_RUNTIME_FACTORY = RuntimeFactory() + + +async def execute( + name: str, + args: dict[str, Any], + *, + runtime_factory: RuntimeFactory | None = None, +) -> dict[str, Any]: from datamind.core.context import RequestContext from datamind.core.logging import bind_context profile = profile_name(args) - settings = Settings() - settings.data.profile = profile - system = await build_datamind(settings, enable=enabled_surfaces(name)) + factory = runtime_factory or _RUNTIME_FACTORY + entry = await factory.get(profile, enabled_surfaces(name)) + system = entry.system + settings = entry.settings context = RequestContext( session_id=str(args.get("session_id") or "codex"), profile=profile, user_id="codex", ) - try: - with bind_context(context): - if name == "datamind_raw_file_read": - return await system.retrieve.tools.get("raw_file_read").handler( - path=str(args["path"]), offset=int(args.get("offset", 0)), max_chars=int(args.get("max_chars", 20000))) - if name == "datamind_build_status": - return await system.retrieve.tools.get("build_status").handler(build_id=str(args["build_id"])) - if name == "datamind_workspace_inspect": - spec = system.retrieve.tools.get("workspace_inspect") - return await spec.handler( - path=str(args["path"]), - recursive=bool(args.get("recursive", True)), - include_hash=bool(args.get("include_hash", True)), - max_files=int(args.get("max_files", 2000)), - ) - if name == "datamind_ask": - return await system.query(str(args["question"])) - if name == "datamind_store": - return await system.ingest(str(args["message"])) - if name == "datamind_use_folder": - spec = system.store.tools.get("kb_add_path") - result = {"kb": await spec.handler(path=str(args["path"]), recursive=bool(args.get("recursive", True)))} - if args.get("build_graph", True): - graph_spec = system.store.tools.get("graph_add_path") - result["graph"] = await graph_spec.handler(path=str(args["path"]), recursive=bool(args.get("recursive", True))) - return result - if name == "datamind_graph_ingest": - spec = system.store.tools.get("graph_add_path") - return await spec.handler(path=str(args["path"]), recursive=bool(args.get("recursive", True)), max_triples_per_file=int(args.get("max_triples_per_file", 30))) - if name == "datamind_graph_build_lineage": - spec = system.store.tools.get("graph_build_lineage") - return await spec.handler(path=str(args["path"]), recursive=bool(args.get("recursive", True)), max_files=int(args.get("max_files", 2000)), dependencies=args.get("dependencies")) - if name == "datamind_table_ingest": - spec = system.store.tools.get("db_import_path") - return await spec.handler(path=str(args["path"]), table_prefix=args.get("table_prefix"), if_exists=str(args.get("if_exists", "append"))) - if name == "datamind_build_start": - return await system.store.tools.get("build_start").handler(path=str(args["path"])) - if name == "datamind_build_freeze": - return await system.store.tools.get("build_freeze").handler(build_id=str(args["build_id"])) - if name == "datamind_build_verify": - return await system.retrieve.tools.get("build_verify").handler(build_id=str(args["build_id"])) - if name == "datamind_build_export": - return await system.store.tools.get("build_export").handler(build_id=str(args["build_id"]), output_path=str(args["output_path"])) - if name == "datamind_surface_ingest": - return await system.store.tools.get("surface_ingest_path").handler(path=str(args["path"]), surfaces=args.get("surfaces"), recursive=bool(args.get("recursive", True))) - if name == "datamind_rag_query": - spec = system.retrieve.tools.get("kb_search") - return await spec.handler(query=str(args["query"]), top_k=int(args.get("top_k", 5))) - if name == "datamind_graph_query": - return await system.query("Answer using the knowledge graph and relationship tools where useful: " + str(args["query"])) - if name == "datamind_remember": - spec = system.store.tools.get("memory_save") - return await spec.handler(content=str(args["content"]), kind=str(args.get("kind", "fact")), scope=str(args.get("scope", "profile")), session_id=args.get("session_id")) - if name == "datamind_status": - return {"profile": profile, "data_dir": str(settings.data.data_dir), "storage_dir": str(settings.data.storage_dir), "warmup": await system.warmup()} - raise ValueError(f"unknown tool: {name}") - finally: - await system.aclose() + with bind_context(context): + if name == "datamind_raw_file_read": + return await system.retrieve.tools.get("raw_file_read").handler( + path=str(args["path"]), offset=int(args.get("offset", 0)), max_chars=int(args.get("max_chars", 20000))) + if name == "datamind_build_status": + return await system.retrieve.tools.get("build_status").handler(build_id=str(args["build_id"])) + if name == "datamind_workspace_inspect": + spec = system.retrieve.tools.get("workspace_inspect") + return await spec.handler( + path=str(args["path"]), + recursive=bool(args.get("recursive", True)), + include_hash=bool(args.get("include_hash", True)), + max_files=int(args.get("max_files", 2000)), + ) + if name == "datamind_ask": + return await system.query(str(args["question"])) + if name == "datamind_store": + return await system.ingest(str(args["message"])) + if name == "datamind_use_folder": + spec = system.store.tools.get("kb_add_path") + result = {"kb": await spec.handler(path=str(args["path"]), recursive=bool(args.get("recursive", True)))} + if args.get("build_graph", True): + graph_spec = system.store.tools.get("graph_add_path") + result["graph"] = await graph_spec.handler(path=str(args["path"]), recursive=bool(args.get("recursive", True))) + return result + if name == "datamind_graph_ingest": + spec = system.store.tools.get("graph_add_path") + return await spec.handler(path=str(args["path"]), recursive=bool(args.get("recursive", True)), max_triples_per_file=int(args.get("max_triples_per_file", 30))) + if name == "datamind_graph_build_lineage": + spec = system.store.tools.get("graph_build_lineage") + return await spec.handler(path=str(args["path"]), recursive=bool(args.get("recursive", True)), max_files=int(args.get("max_files", 2000)), dependencies=args.get("dependencies")) + if name == "datamind_table_ingest": + spec = system.store.tools.get("db_import_path") + return await spec.handler(path=str(args["path"]), table_prefix=args.get("table_prefix"), if_exists=str(args.get("if_exists", "append"))) + if name == "datamind_build_start": + return await system.store.tools.get("build_start").handler(path=str(args["path"])) + if name == "datamind_build_freeze": + return await system.store.tools.get("build_freeze").handler(build_id=str(args["build_id"])) + if name == "datamind_build_verify": + return await system.retrieve.tools.get("build_verify").handler(build_id=str(args["build_id"])) + if name == "datamind_build_export": + return await system.store.tools.get("build_export").handler(build_id=str(args["build_id"]), output_path=str(args["output_path"])) + if name == "datamind_surface_ingest": + return await system.store.tools.get("surface_ingest_path").handler(path=str(args["path"]), surfaces=args.get("surfaces"), recursive=bool(args.get("recursive", True))) + if name == "datamind_rag_query": + spec = system.retrieve.tools.get("kb_search") + return await spec.handler(query=str(args["query"]), top_k=int(args.get("top_k", 5))) + if name == "datamind_graph_query": + return await system.query("Answer using the knowledge graph and relationship tools where useful: " + str(args["query"])) + if name == "datamind_remember": + spec = system.store.tools.get("memory_save") + return await spec.handler(content=str(args["content"]), kind=str(args.get("kind", "fact")), scope=str(args.get("scope", "profile")), session_id=args.get("session_id")) + if name == "datamind_status": + return {"profile": profile, "data_dir": str(settings.data.data_dir), "storage_dir": str(settings.data.storage_dir), "warmup": entry.warmup} + raise ValueError(f"unknown tool: {name}") def list_profiles() -> dict[str, Any]: configured = os.environ.get("DATAMIND_DATA_ROOT", "").strip() @@ -259,7 +325,7 @@ def tool_result(request_id: Any, value: Any) -> dict[str, Any]: def error(request_id: Any, message: str) -> dict[str, Any]: return {"jsonrpc": "2.0", "id": request_id, "result": {"isError": True, "content": [{"type": "text", "text": message}]}} -def handle(message: dict[str, Any]) -> dict[str, Any] | None: +async def handle(message: dict[str, Any]) -> dict[str, Any] | None: request_id = message.get("id") if request_id is None: return None @@ -274,7 +340,7 @@ def handle(message: dict[str, Any]) -> dict[str, Any] | None: listed: list[dict[str, Any]] = [] for name, spec in TOOLS.items(): item = {"name": name, **spec} - if name in READ_ONLY_TOOLS: + if ROUTES[name].read_only: item["annotations"] = { "readOnlyHint": True, "destructiveHint": False, @@ -287,21 +353,31 @@ def handle(message: dict[str, Any]) -> dict[str, Any] | None: name = str(params.get("name") or "") if name == "datamind_list_profiles": return tool_result(request_id, list_profiles()) - return tool_result(request_id, asyncio.run(execute(name, params.get("arguments") or {}))) + return tool_result(request_id, await execute(name, params.get("arguments") or {})) return {"jsonrpc": "2.0", "id": request_id, "error": {"code": -32601, "message": f"Method not found: {method}"}} except Exception as exc: return error(request_id, f"{type(exc).__name__}: {exc}") +async def serve() -> None: + try: + while True: + line = await asyncio.to_thread(sys.stdin.readline) + if line == "": + break + if not line.strip(): + continue + try: + reply = await handle(json.loads(line)) + except json.JSONDecodeError as exc: + reply = {"jsonrpc": "2.0", "id": None, "error": {"code": -32700, "message": str(exc)}} + if reply is not None: + print(json.dumps(reply, ensure_ascii=False, separators=(",", ":")), flush=True) + finally: + await _RUNTIME_FACTORY.aclose() + + def main() -> None: - for line in sys.stdin: - if not line.strip(): - continue - try: - reply = handle(json.loads(line)) - except json.JSONDecodeError as exc: - reply = {"jsonrpc": "2.0", "id": None, "error": {"code": -32700, "message": str(exc)}} - if reply is not None: - print(json.dumps(reply, ensure_ascii=False, separators=(",", ":")), flush=True) + asyncio.run(serve()) if __name__ == "__main__": main()