diff --git a/.env.example b/.env.example deleted file mode 100644 index 7941a6c..0000000 --- a/.env.example +++ /dev/null @@ -1,32 +0,0 @@ -# Copy to .env and fill in your values. .env is gitignored; do not commit. -# -# Sourced by: -# - tests/smoke_test/run.sh / `dsagt smoke-test` -# - dsagt_config.yaml (per-project) reads these via ${VAR} interpolation -# when init writes the config - -# --- Upstream LLM gateway --- -# LLM_PROVIDER: LiteLLM provider prefix (selects request format + auth). -# Common: openai, anthropic, bedrock, vertex_ai, azure, gemini, -# ollama, mistral, groq, deepseek. -# Full list: https://docs.litellm.ai/docs/providers -LLM_PROVIDER=openai -LLM_API_KEY= -LLM_BASE_URL=https://your-gateway.example.com -LLM_MODEL=claude-sonnet-4-5 - -# --- Embedding endpoint (OPTIONAL) --- -# Embeddings default to ``backend: local`` in dsagt_config.yaml — sentence- -# transformers, CPU-side, no network or API key needed. The variables -# below are only consulted when you flip ``embedding.backend`` to ``api`` -# in a project's dsagt_config.yaml. Leave them blank if you stick with -# the local default. -# -# EMBEDDING_PROVIDER: LiteLLM provider prefix for embeddings. -# ``openai_like`` covers any OpenAI-wire-compatible /embeddings endpoint -# (including most lab gateways). Use ``bedrock`` for native AWS Bedrock -# InvokeModel (Titan, Cohere) or ``openai``, ``cohere``, ``voyage``, etc. -EMBEDDING_PROVIDER=openai_like -EMBEDDING_API_KEY= -EMBEDDING_BASE_URL=https://your-gateway.example.com -EMBEDDING_MODEL=text-embedding-3-small diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml index 0cae487..9ed5cb7 100644 --- a/.github/workflows/docs.yml +++ b/.github/workflows/docs.yml @@ -21,12 +21,15 @@ jobs: with: python-version: "3.12" + - name: Install uv + uses: astral-sh/setup-uv@v5 + - name: Install docs dependencies - run: pip install mkdocs-material + run: uv sync --group docs - name: Build docs - run: mkdocs build --strict + run: uv run mkdocs build --strict - name: Deploy to GitHub Pages if: github.ref == 'refs/heads/main' && github.event_name != 'pull_request' - run: mkdocs gh-deploy --force + run: uv run mkdocs gh-deploy --force diff --git a/.gitignore b/.gitignore index c37bb57..cb03964 100644 --- a/.gitignore +++ b/.gitignore @@ -105,6 +105,17 @@ test_registry.yaml test_runtime/ *.log provenance.log +# Serverless MLflow store + any stray file-store trace dumps — never commit +# a session's traces (they carry machine-local artifact paths). Guards +# against re-committing dumps like the pre-existing comb-flow-uni/**/mlruns/. +mlflow.db +mlruns/ +latex/ +design-notes/ +scratch/ +env_examples/ +notes.txt +uv.lock # Demo outputs demo/output/ @@ -151,3 +162,6 @@ cline_mcp.json *.bak *.swp .*.swp + +# Local scratch for removed/parked code (not tracked) +scratch/ diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 index 0000000..6f0f2f9 --- /dev/null +++ b/CHANGELOG.md @@ -0,0 +1,177 @@ +# Changelog + +All notable changes to DSAgt are documented here. The format is based on +[Keep a Changelog](https://keepachangelog.com/en/1.1.0/), and this project +adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). + +## [Unreleased] + +This cycle restores **observability and episodic memory without a proxy** — +both recovered from each agent's own on-disk transcript rather than intercepted +LLM traffic — renames registered executables to **codes**, and makes codes and +skills natively discoverable the moment they land in a project. + +### Added +- **Proxy-free trace pipeline.** An in-session heartbeat in `dsagt-server` + reads the agent's transcript and logs prompts, responses, tool calls, and + token usage to the serverless MLflow store — every supported agent (claude, + codex, goose, opencode, cline); no proxy, no OTel routing, no credentials. + DSAgt's own spans flow to the same store, tagged `dsagt.source` so the debug + view filters apart from agent traces. +- **`dsagt traces `** opens the MLflow viewer with nothing to + remember: runs trace catch-up first, deep-links straight to the project's + Traces tab, quiets the server noise. +- **Episodic memory (opt-in).** `dsagt init --episodic`: the heartbeat + mechanically chunks, keyword-tags, and embeds each completed turn into + `session_memory` (no LLM); retrieval is recency-weighted + (`episodic.recency_half_life_days`, default 14). +- New domain use cases: tokamak stability, AIDRIN data-readiness. + +### Changed +- **"Tools" are now "codes."** Registered CLI executables are *codes* + throughout — `/codes/`, `save_code_spec`, `dsagt-run --code`, the + `code_use` collection, `code.execute` spans — reserving "tool" for the + MCP/agent sense. +- **Codes share the skill envelope.** A code is a skill-standard dir + (`codes//SKILL.md`); bundled codes are copied into `/codes/` + at init — one place, one format. +- **Native discovery is immediate.** Installing/creating a skill or registering + a code mirrors it into the agent's native skills dir on the spot (previously + at the next `dsagt start`), so the next session auto-discovers it however the + agent is launched — and the agent no longer tells users a restart is needed. +- **Agents are pre-authenticated.** DSAgt never touches provider credentials; + all credential-hint machinery is gone. +- **Code-use indexing is incremental.** `dsagt-run` records embed into + `code_use` on the heartbeat and on demand before `reconstruct_pipeline`, + instead of waiting for the next session's startup catch-up. +- Docs split into per-capability pages; use-case walkthroughs restructured + (setup → prompt dialogue → post-conditions) and swept for staleness. + +### Fixed +- **Duplicate `code_use` entries** — indexing is now idempotent against a + persisted ack set shared by the heartbeat and startup catch-up. +- **cline:** per-project MCP config via `CLINE_MCP_SETTINGS_PATH`; `cline mcp + add` works on cline 3.x; global auth and settings are never touched. +- **codex:** the trace reader follows the per-project `CODEX_HOME`. + +### Removed +- The episodic **LLM-judge** distillation layer and **outlier-suggestion** + feature (incl. the `kb_get_suggestions` / `kb_dismiss_suggestion` MCP tools + and the `llama-cpp-python` dependency), plus their `dsagt init` prompts and + config keys. Episodic memory keeps the mechanical capture path so a Tier-0 + baseline can be measured first; design notes parked in + `design-notes/judge.md`. +- Dead `provenance.index_execution_record`. + +## [0.2.0] - 2026-06-24 + +This release adds an **external skill-catalog system** and consolidates the +agent-facing surface into a **single MCP server**. + +Agents can now discover and install skills from federated GitHub/GitLab catalogs +(Genesis, Anthropic, K-Dense, and more) and author their own — while installed +skills are picked up through each agent's *native* `SKILL.md` auto-discovery, so +`search_skills` is reserved for the one job native discovery can't do: browsing +the catalog of skills you haven't installed yet. + +In parallel, the registry and knowledge MCP servers — two processes that each +loaded their own embedder and opened their own ChromaDB (pure duplication, plus +a write-here/read-there hazard on the shared skill-catalog collections) — +collapse into one `dsagt-server`: one embedder, one Chroma owner, one connection +per agent, so startup is faster with fewer moving parts per project. + +**Upgrading from 0.1.0 (forwards compatibility).** There is no automatic +migration — adopting 0.2.0 is rebuild-not-migrate, and no project data changes: +- Re-run `dsagt start ` for each existing project; it regenerates the + per-agent MCP config to point at the single `dsagt-server`. +- For **cline** only, delete `/.cline-data` first — `cline mcp add` + has no remove, so the stale `dsagt-registry`/`dsagt-knowledge` entries would + otherwise linger next to the new one. +- Tools, skills, the KB index, traces, and memory all carry over untouched. + +### Added +- **External skill catalogs**: discover and install agent skills from GitHub / + GitLab sources via `add_skill_source`, `search_skills`, and `install_skill` + (plus the `dsagt skills sync/add/list/search` CLI), backed by per-source + ChromaDB collections. Curated sources ship out of the box (`scientific`, + `anthropic`, `antigravity`, `composio`, `genesis`); any git URL / `owner/repo` + also works. +- **Genesis catalog integration**: the curated `genesis` source (OSTI GitLab, + `gitlab.osti.gov/genesis/genesis-skills`) makes the BASE-Data / ModCon skills + — `datacard-generator` (frontmatter name `generating-datacards`), + `croissant-validator`, `hdmf-schema-builder` — pullable on demand + (`dsagt skills add genesis`, then `install_skill`) rather than + bundled in the package, alongside the rest of the Genesis catalog (HPC/Slurm, + HuggingFace, LangChain, and more). +- **Native skill discovery**: installed and bundled skills are mirrored into + each agent's native skill directory (`.claude/skills/`, `.agents/skills/`, …) + at init/start, so every supported agent auto-discovers them. +- **`skill-creator`** bundled skill for authoring new skills from the Anthropic + template. +- **Source-qualified catalog install**: when the same skill name exists in more + than one synced source, install a specific one with a `/` + name (via `install_skill` or `dsagt skills add /`) + instead of dead-ending on the ambiguity guard. +- **Keyword fallback** for `search_skills`: a zero-dependency token-overlap + scorer so catalog search works even when no embedding model is configured. +- **License / attribution provenance on install**: installing a catalog skill + preserves upstream `LICENSE` / `NOTICE` files and stamps a `PROVENANCE.txt` + recording the source repo and path into the installed skill directory. +- **`isaac_skills_demo` use case**: an end-to-end, skill-oriented walkthrough + (`use_cases/isaac_skills_demo/`) that drives a real agent through syncing a + catalog, installing a skill, and converting mock VASP output into an Isaac + record — with prompts and mock data included. +- **Install-from-GitHub instructions** for non-developers (`pip install + git+https://github.com/AI-ModCon/dsagt.git` into any Python 3.12/3.13 + environment) in the README and docs. + +### Changed +- **The two MCP servers are now one `dsagt-server`** — one shared + `KnowledgeBase`/embedder, one MCP entry per agent, one trace `service.name`. + The tool surface is organized by concern (registry / knowledge / memory / + skill) behind the single server. +- Skill discovery is now **catalog-only**: installed and bundled skills are + discovered natively by every supported agent, so `search_skills` covers only + the not-yet-installed external catalog. Catalogs are indexed on frontmatter + (name + description + tags) rather than the full SKILL.md body. +- `search_skills` now reports when no external catalog is synced instead of a + bare "no match", and `list_skill_sources` flags each known source as + `synced`/available with its indexed count. +- `install_skill` clarifies that an installed skill is usable in the current + session immediately — a restart is only needed for hands-free native + auto-invocation. +- The package version is single-sourced from `dsagt.__version__` (pyproject + reads it via setuptools dynamic metadata). +- Documentation home page (`docs/index.md`) pulls the supported-agents table + and install instructions directly from the README via the + `mkdocs-include-markdown` plugin, so the two no longer drift. + +### Removed +- **BREAKING:** the `dsagt-registry-server` and `dsagt-knowledge-server` console + scripts, replaced by `dsagt-server` (see **Upgrading** above). +- The bundled `datacard-generator` skill — it lives in the Genesis catalog and + is now installed on demand via `dsagt skills add genesis`. +- Dead indexing of installed/bundled skills into the `skills` ChromaDB + collection (nothing read it after the catalog-only search change). + +### Fixed +- CLI-added skill sources are now persisted to the project config. +- `dsagt --version` now works (it was documented but unimplemented — argparse + errored). Reports the version from `dsagt.__version__`. +- Catalog skills with technically-invalid YAML frontmatter (e.g. an unquoted + `description` containing a colon, like `…readiness levels: Level 1…`) are no + longer silently dropped from discovery. `_parse_frontmatter` falls back to a + lenient flat parse that recovers `name`/`description`/`tags`, so such skills — + including Genesis's `generating-datacards` (`datacard-generator`) — are + searchable and installable instead of skipped. + +## [0.1.0] - 2026-01-11 + +### Added +- Initial release: registry and knowledge MCP servers, BYOA per-agent config + generation, MLflow/OTel observability, the tool/skill registry, execution + provenance, and explicit + episodic memory. + +[Unreleased]: https://github.com/AI-ModCon/dsagt/compare/v0.2.0...HEAD +[0.2.0]: https://github.com/AI-ModCon/dsagt/compare/v0.1.0...v0.2.0 +[0.1.0]: https://github.com/AI-ModCon/dsagt/releases/tag/v0.1.0 diff --git a/CLAUDE.md b/CLAUDE.md new file mode 100644 index 0000000..eb7bb51 --- /dev/null +++ b/CLAUDE.md @@ -0,0 +1,142 @@ +# CLAUDE.md + +Guidance for Claude Code (claude.ai/code) when working in this repository. + +## Project Overview + +DSAGT (DataSmith Agent) is an AI-assisted data pipeline builder. It exposes an MCP (Model Context Protocol) server to agent platforms (Claude Code, Goose, Cline, Codex, opencode) that helps domain scientists build reproducible, auditable data-curation pipelines through iterative, knowledge-driven code generation. + +## Run model (BYOA, serverless) + +DSAGT is **bring-your-own-agent**: the agent talks to its own LLM provider directly — DSAGT never proxies that traffic. `dsagt init --agent ` writes a per-agent instructions file plus MCP config; nothing to paste, no launch shim, no OTel/env wiring. Two behaviorally-identical start flows: + +1. **Bare launch** — start the agent from the project dir (`cd && claude` / `goose` / …). +2. **`dsagt start `** — spawns the agent in the foreground and, on exit, runs post-session `run_extraction`. Its real job is owning a reliable session-end trigger. + +The MCP server is self-sufficient: it derives the project from its cwd (`.dsagt/config.yaml`), with the MCP-config env block as robustness, and behaves identically regardless of how it was launched. + +**Serverless store.** All self-logging goes to a `sqlite:////mlflow.db` MLflow store — no server to run (SQLite is MLflow's supported serverless backend). `observability.resolve_tracking_uri` always computes that path from the project dir (no env/config override) and never raises. The DB auto-creates on first span; DSAGT emits only traces, so no artifact dir materializes. View with `mlflow ui --backend-store-uri sqlite:////mlflow.db`. Agent LLM-call history is recovered post-hoc from the on-disk transcript (the trace pipeline), not by intercepting traffic. + +## Commands + +```bash +uv sync --all-groups # install +uv run --no-sync python -m pytest tests/test_.py -q # targeted tests +uv run black . # format +uv run ruff check . # lint +``` + +**`python -m pytest`, not bare `pytest`.** The bare binary picks up the wrong pytest on this machine and crashes with `ModuleNotFoundError: dsagt`. + +**Don't run the full suite by default** — ~50s for ~590 tests is too slow for an iteration loop. Run only the test file relevant to the change. `tests/test_config.py` covers session, init, agents, serverless-store resolution, and the no-launch-shim/no-OTel invariants. Skip `test_integration.py`, `test_*_integration.py`, `test_server_startup.py`, `test_dependency_integration.py` unless relevant — they hit the network or spawn subprocesses. + +## Code Organization + +The codebase separates **commands** (entry points with argparse, launched as CLI tools or subprocesses) from **modules** (importable logic). Commands live in `src/dsagt/commands/`, modules in `src/dsagt/`. + +**Commands** (`src/dsagt/commands/`): +- `cli.py` — `dsagt init / start / info / traces / list / mv / rm / smoke-test`. `dsagt start` launches the agent in the foreground and runs post-session extraction on exit. `dsagt traces ` opens the MLflow viewer over the project's store (runs catch-up first, deep-links to the Traces tab, quiets the mlflow noise; foreground, no managed daemon). +- `run_code.py` — `dsagt-run` (code execution wrapper). +- `setup_core_kb.py` — KB-asset build engine (`resolve_assets` / `ensure_assets`), called by `dsagt init` to provision the shared KB; not a CLI command of its own. +- `info.py` — `dsagt info` (project/config introspection + trace triage). + +(The MCP server — `dsagt-server` — lives in the `src/dsagt/mcp/` package, see below.) + +**Modules** (`src/dsagt/`): +- `session.py` — Project init, agent config generation, env-var resolution, config load/validate, session-id minting (`new_session_id`), end-of-session extraction (`run_extraction`), and startup **catch-up** (`catch_up_extraction`): code-use indexing + a chat-trace re-collect (`_catch_up_traces`) of the *previous* session (pinned to the `trace_source` token in `state.yaml`) so turns lost to an ungraceful shutdown still reach MLflow + episodic memory — uniform across all agents. +- `agents/` — Per-agent-platform setup (`base.py` ABC + `claude.py` / `goose.py` / `cline.py` / `codex.py` / `opencode.py`). Each subclass owns its `write_static`, `write_dynamic`, `runtime_env`, `vscode_hint`. Shared helpers (`_mcp_env_block`, `_build_mcp_servers_dict`) in `base.py`. DSAGT sets no telemetry/OTel env, writes no launch shim, and never touches provider credentials — agents are expected pre-authenticated (shell env / their own auth flows) before dsagt is pointed at them; dsagt prints no credential hints and never troubleshoots auth. +- `knowledge.py` — ChromaDB document retrieval, embedding backends, per-collection routing (the reference example of the house style). +- `registry.py` — `CodeRegistry` (CLI codes) + `SkillRegistry` (agent instruction skills), KB indexing. +- `provenance.py` — Code execution records (`run_and_record`), execution-record indexing into ChromaDB (`CodeUseIndexer` → `code_use` collection), pipeline reconstruction (`reconstruct_pipeline`, dependency graph). +- `observability.py` — first-party span emission over the serverless sqlite store via MLflow's native `mlflow.start_span` (no OTel `TracerProvider`). `resolve_tracking_uri` (never-raise), `init_tracing`, `@traced`/`obs`/`child_span` + typed span helpers. Each internal trace's root is tagged `dsagt.source` with the MCP tool *category* (`memory`/`skill`/`knowledge`/`registry`, or `execution` for dsagt-run) so the MLflow UI can filter the debug view apart from agent traces. `MLflowSink` (a `traces.Trace` consumer) replays finished transcripts via `start_span_no_context`. +- `memory.py` — Explicit memory (YAML, `ExplicitMemory`) + the episodic `MemoryExtractor` (a `traces.TraceCollector` consumer that mechanically chunks+tags+embeds every turn, no LLM). Turns carry `ts_epoch` for recency-weighted retrieval. `extract_session` is a no-op stub kept only for the deferred cross-session N+1 catch-up call site. +- `skills.py` — External skill-catalog data plane (`SkillsCatalog`: clone/sync/index/install), the `SkillRouter` render facade, and the Genesis-derived keyword scorer (`rank_skills`). +- `traces.py` — the whole trace pipeline in one module: the pure-data `Trace` (span dicts + compose/query/`to_exchanges`), the `Reader`/`Translator` ABCs with a per-agent subclass each (Claude bespoke; codex/goose/opencode/cline share the `Translator` turn-template; claude+codex share `JsonlReader`), and `TraceCollector` — the MCP-server heartbeat that reads→translates→hands the `Trace` to its consumers (MLflow logger, memory indexer), each with its own ack set for idempotency. Imports nothing heavy (mlflow is lazy, consumer-side). + +**MCP server** (`src/dsagt/mcp/`) — the single merged `dsagt-server`. `server.py` owns `main()`, the shared-KB startup (`_build_kb_from_config`), and the dispatch shell (`build_dispatch_server`). The 20-tool surface is split by concern: `registry_tools.py` (code registry + execution + provenance, 8), `knowledge_tools.py` (KB retrieval, 5), `memory_tools.py` (explicit memory, 2), `skill_tools.py` (skill search/install/sources, 5). Each `*_tools.py` exposes a `_*_tools_and_handlers()` factory (composed by `create_dsagt_server`) plus a `create_*_server` test wrapper. + +Entry points (`pyproject.toml` `[project.scripts]`): `dsagt` → `dsagt.commands.cli:main`, `dsagt-run` → `dsagt.commands.run_code:main`, `dsagt-server` → `dsagt.mcp.server:main`. + +**Bundled assets** (shipped as `package-data`): +- `src/dsagt/codes/` — built-in codes as skill-standard dirs (`/SKILL.md`), served from the package (never copied into projects). +- `src/dsagt/skills/` — built-in skills (e.g., `skill-creator`) the agent discovers via `search_skills`. +- `src/dsagt/dsagt_instructions.md` — agent-agnostic system instructions injected into per-agent files at init. + +**`use_cases/`** holds end-to-end domain walkthroughs. They are reference material for users, not part of the test suite. + +## Code style & conventions + +Distilled from working on this codebase; `knowledge.py` is the reference example of the house style. + +**No defensive swallowing.** Don't add guards that silently absorb empty/invalid input (`if not texts: return []`, empty-array short-circuits, disk-state "reconciliation" of can't-happen states). They convert a caller's bug into a silent success you'll never see. Empty/invalid input is out-of-contract — let it surface. Translating a *real, reachable* exception into an actionable message (e.g. a dim-mismatch hint) is different and welcome; swallowing is not. + +**YAGNI / no speculative generality.** Don't add a flag or option for a path never exercised in practice. Model real variation *structurally* (a subclass / distinct type), not with a runtime toggle nothing flips — e.g. the local store is unconditionally hybrid; "dense-only" is a future store *type*, not a `hybrid=False` flag. Don't extract a base class until a second concrete impl forces the seam. This is dev-stage: gut cleanly, no back-compat shims, aim net-minus LOC. + +**Explicit named arguments, never `**kwargs` config-splat.** A dict of kwargs threaded through layers and `**`-splatted into a constructor hides what's actually passed. Use explicit named params; unpack any config dict at the boundary (e.g. `Embedder.create(backend, *, model=, base_url=, ...)`, callers pass `model=cfg.get("model")`). + +**Put behavior where it belongs.** Factories live on the class they build (`Embedder.create` classmethod, not a free `_make_embedder`). Trivial field getters are unpythonic — expose the attribute; but a *method* is right when access does real work (lazy I/O + memoization like `_get_bm25`). Pure, stateless algorithms shared by multiple classes stay module-level functions (RRF: `_rrf_merge`/`_rrf_across`), not staticmethods nailed to one arbitrary owner. + +**Comments state the real reason, at the point they explain it.** A lazy import is justified *at the import site* with its actual cause, not as an "this is absent" note in the import block citing a stale rationale. If the reason changes, fix the comment. + +**Import hygiene on hot paths.** Modules on frequently-invoked paths (`dsagt-run` runs per tool call) must not transitively drag heavy modules in for *annotation-only* type hints. Use `from __future__ import annotations` + a `TYPE_CHECKING`-guarded import (verify the module doesn't introspect annotations at runtime first). Keep cold start lean; lazy-import the heavy leaf (llama_index) at its single use site. + +**Naming.** Prefer concise domain names (`APIEmbedder`/`LocalEmbedder`, not `…EmbeddingClient`). + +**Module docstrings (major modules).** Open with a title line + 3–5 sentences: what the module does, the capabilities it backs, the design motivations. Follow with an **ASCII-art UML class map** — one consistent notation throughout (`knowledge.py` uses `◇` holds · `◆` owns · `▷` inherits). Treat the class-map diagram as a deliverable of any **major module refactor** — refresh it whenever the class structure changes substantially. + +## BYOA artifacts + +`dsagt init --agent X --location ` writes, in the project dir: +- `.dsagt/config.yaml` — internal config (project name, agent, embedding/knowledge/extraction/skills settings). No mlflow port (the store is the serverless `sqlite:////mlflow.db`), no user-facing fields, no credentials. `.dsagt/state.yaml` (session log + memory cursor) and `.dsagt/explicit_memories.yaml` live alongside it, owned by the MCP server. +- Per-agent instructions file (e.g., `CLAUDE.md`, `.goosehints`, `AGENTS.md`). +- Per-agent MCP config artifact (`.mcp.json` for claude, `goose.yaml` for goose, `cline_mcp_settings.json` via `cline mcp add`, `.codex-data/config.toml`). The env block carries benign routing only (`DSAGT_PROJECT`, `DSAGT_PROJECT_DIR`, `DSAGT_SESSION_ID`, `MLFLOW_TRACKING_URI`, `EMBEDDING_*`) so MCP-server children of agents that don't inherit shell env (codex/cline) still log to the right store. No credentials, no OTel routing. + +No launch shim is written and `dsagt init` prints no env/OTel instructions — the user starts the agent directly or via `dsagt start`. DSAGT wires no MLflow autolog hook: Claude's traces (like every agent's) come from the heartbeat pipeline, not native autolog. + +## Architecture + +### MCP Server + +A single merged `dsagt-server` (`src/dsagt/mcp/`) exposes 20 tools across four concern modules under one `Server` + one shared `KnowledgeBase`: + +1. **Registry tools** (`mcp/registry_tools.py` + `registry.py` / `provenance.py`) — tool analysis, registration, dependency installation, command/file/http execution, pipeline reconstruction. Tools are saved as markdown specs with YAML frontmatter. +2. **Knowledge tools** (`mcp/knowledge_tools.py` + `knowledge.py`) — semantic search over document collections (ChromaDB, optional cross-encoder reranking); long ops run as background jobs. +3. **Memory tools** (`mcp/memory_tools.py` + `memory.py`) — explicit memory (`kb_remember` / `kb_get_memories`). +4. **Skill tools** (`mcp/skill_tools.py` + `skills.py`) — skill search/install + external catalog sources. + +### Observability + +- **Serverless MLflow store** — spans land in `sqlite:////mlflow.db` (no server). The tracking URI resolves via `observability.resolve_tracking_uri` (never raises). +- **dsagt-run** (`commands/run_code.py` + `provenance.py`) — wraps code commands; captures the execution layer (command, stdout/stderr, timing, file lists) into `trace_archive/` and emits `code.execute` spans. +- **MCP-server + tool spans (debug view)** — `dsagt-server` calls `init_tracing()` at startup; the dispatch shell opens one categorization-root span per tool call (subsystem `kb.*`/`registry.*` spans nest under it). Each root is tagged `dsagt.source` with its concern category so it filters apart as a debugging view. Session grouping via `DSAGT_SESSION_ID`. +- **Agent traces** — recovered post-hoc from the on-disk transcript by the MCP-server heartbeat's trace pipeline (`traces.py`), uniform across all five agents. No native OTel, no autolog. + +### Memory System + +- **Explicit memory** (`memory.py:ExplicitMemory`) — user-confirmed facts in YAML, loaded into agent context at session start via `kb_remember` / `kb_get_memories` (the vector mirror is optional — degrades to pure-YAML if the store is down). +- **Code-execution indexing** — `provenance.CodeUseIndexer` embeds `trace_archive/` records into the project's `code_use` collection incrementally on the heartbeat (idempotent via a persisted ack set), plus a startup catch-up and an on-demand tick before `reconstruct_pipeline`. No LLM. +- **Chat-trace catch-up** — the heartbeat logs the live transcript to MLflow (+ episodic memory) and a graceful shutdown flushes the deferred final turn; an ungraceful kill is backstopped at the *next* session's startup by `session._catch_up_traces`, which re-collects the previous session pinned to its recorded `trace_source` token. Idempotency rests on the collector's **session-qualified** ack keys (`:`). +- **Episodic memory** — live, **opt-in** (`episodic.enabled`, via `dsagt init --episodic`). The `memory.MemoryExtractor` consumer consumes `Trace.to_exchanges()` on the heartbeat and mechanically chunks+tags+embeds every turn into `session_memory` (no LLM). Retrieval is recency-weighted (`episodic.recency_half_life_days`). + +### Key Design Patterns + +- **Agent-agnostic**: DSAGT is infrastructure, not an agent. Capabilities are MCP services. +- **Session isolation**: each project gets its own directory with config, tools, skills, kb_index, trace_archive, and the `mlflow.db` sqlite store. +- **Codes vs Skills**: Codes are CLI executables in `/codes//` (skill-standard dirs whose SKILL.md frontmatter adds executable/parameters; wrapped by dsagt-run). Skills are agent instruction workflows in `/skills/` (SKILL.md + reference docs). Both share the skill envelope, so both mirror into the agent's native skills dir; both are also discoverable via ChromaDB-backed semantic search (`search_registry` / `search_skills`). + +## DSAGT Pipeline Builder Workflow + +When acting as a pipeline builder (using the MCP server), follow these constraints: + +1. **Never directly access data** — all data operations go through registered codes. +2. **Code preference hierarchy**: registered code → KB package code → custom implementation. +3. **Generate paired tools** — every data operation gets a check tool (pre/post audit) and an operation tool. +4. **Audit everything** — before/after JSON reports saved to `audit/`. +5. **One step at a time** — iterate with the user, confirming approach before execution. + +## Testing Patterns + +- pytest with `subprocess.run` mocking for command execution. +- MCP server tests invoke handlers directly (no stdio transport); async tests for server handlers. +- Temp directories for isolation; the `_use_tmp_registry` fixture in `tests/test_config.py` patches `DEFAULT_PROJECTS_BASE` and the project registry to `tmp_path`. +- Integration tests in `test_*_integration.py` require real `EMBEDDING_*` / `LLM_*` credentials. diff --git a/README.md b/README.md index c80e982..656aaf4 100644 --- a/README.md +++ b/README.md @@ -4,106 +4,127 @@ ![DSAgt architecture](latex/architecture.png) -DSAgt connects an MCP-compatible AI coding agent to tool registration, a semantic knowledge base, execution provenance, and observability infrastructure. DSAgt provides data-pipeline scaffolding around a user's existing agent CLI or VS Code extension (Claude Code, Goose, Codex, …); +DSAgt connects an MCP-compatible AI coding agent to code registration, a semantic knowledge base, skills discovery and creation, execution provenance, and observability infrastructure. DSAgt provides data-pipeline scaffolding around a user's existing agent CLI or VS Code extension (Claude Code, Goose, Codex, …); -**Prerequisites:** Python 3.10–3.13, [uv](https://github.com/astral-sh/uv), and one of the supported agent platforms below — already installed and authenticated against whatever LLM provider you intend to use. +**Prerequisites:** Python 3.12 or 3.13, and one of the supported agent platforms below — already installed and authenticated against whatever LLM provider you intend to use. + | Agent | Install | Verify | |-------|---------|--------| | [Claude Code](https://github.com/anthropics/claude-code) | `npm i -g @anthropic-ai/claude-code` | `claude --version` | | [Goose](https://github.com/block/goose) | See [Goose docs](https://github.com/block/goose#installation) | `goose --version` | | [Codex](https://github.com/openai/codex) | `npm i -g @openai/codex` (or `brew install --cask codex`) | `codex --version` | | [opencode](https://github.com/sst/opencode) | See [opencode docs](https://opencode.ai/docs/) | `opencode --version` | -| [Roo Code](https://github.com/RooCodeInc/Roo-Code) | `npm i -g @roo-code/cli` | `roo --version` | | [Cline](https://github.com/cline/cline) | `npm i -g cline` | `cline --version` | + + +## Installation + +### For use (no development) + + + +```bash +python3.12 -m venv ~/.venvs/dsagt # or: conda create -n dsagt python=3.12 && conda activate dsagt +source ~/.venvs/dsagt/bin/activate # (Windows venv: ~\.venvs\dsagt\Scripts\activate) +pip install "git+https://github.com/AI-ModCon/dsagt.git" +dsagt --version # 0.2.0 +``` + +This puts the `dsagt` CLI (and the `dsagt-run` / `dsagt-server` helpers) on your PATH. Create your first project — `dsagt init` is interactive (it walks you through the agent platform, project location, packaged knowledge collections, and skill-catalog sources) and sets up the knowledge base on first run: + +```bash +dsagt init # interactive; pick agent, collections, sources +``` + +Then start your agent from the project directory, use dsagt cli, or open a VS code : + +```bash +cd ~/dsagt-projects/my-project && claude # …or: dsagt start my-project +``` + +To upgrade later, reinstall — re-running `dsagt init` reconfigures an existing project in place: + +```bash +pip install --upgrade "git+https://github.com/AI-ModCon/dsagt.git" +``` + +> Pin to a specific release once tags are published, e.g. `pip install "git+https://github.com/AI-ModCon/dsagt.git@v0.2.0"`. + + +### For development + +Clone the repo and use `uv` (editable install with the full test suite) — see [Quick Start](#quick-start) below. ## Quick Start -Explore DSAgt knowledge ingest, tool registration, provenance, and explicit memory using the mock project in [`tests/smoke_test/`](tests/smoke_test/). Uses `claude`; substitute another agent (`goose` / `codex` / `opencode`) if you prefer — the prompts are agent-agnostic. +Explore DSAgt knowledge ingest, code registration, provenance, and explicit memory using the mock project in [`tests/smoke_test/`](tests/smoke_test/). Uses `claude`; substitute another agent (`goose` / `codex` / `opencode` / `cline`) if you prefer — the prompts are agent-agnostic. ```bash -# 0. Installation +# 0. Install git clone https://github.com/AI-ModCon/dsagt.git cd dsagt uv sync # add --all-groups for the test suite source .venv/bin/activate # so `dsagt` is on PATH -# Set convenience folder env variable for quickstart demo (not a normal dsagt step) +# A convenience variable for the demo paths below (not a normal dsagt step) export SMOKE_DIR="$(pwd)/tests/smoke_test" -# 1. Create a new project called quickstart -dsagt init quickstart --agent claude +# 1. Create a project. `dsagt init` is interactive — follow the menu to name it +# `quickstart`, pick your agent, and choose knowledge collections + skill sources. +# It sets up the knowledge base on first run (a ~130 MB local embedder downloads once). +dsagt init -# 2. Start MLflow in the background (writes /mlflow.log) and print -# the OTel routing exports for this session, including the resolved -# experiment id: -dsagt mlflow quickstart - -# 3. Paste the export block dsagt mlflow printed into THIS shell, then -# launch claude from the project directory: -cd ~/dsagt-projects/quickstart && claude +# 2. Launch the agent in the project: +dsagt start quickstart # …or: cd ~/dsagt-projects/quickstart && ``` Inside the agent, paste these prompts one at a time (substitute the absolute path you exported as `$SMOKE_DIR` — the chat doesn't expand env vars): 1. > Ingest the docs in `$SMOKE_DIR/knowledge/` into a collection named `knowledge`. -2. > Register the csvkit CLI tools `csvcut`, `csvgrep`, `csvstat`, and `csvlook`. -3. > Use the `scan_directory` tool from the registry to scan `$SMOKE_DIR/data/`. -4. > Summarize `samples.csv` — columns, row count, quality issues using csvkit tools from the registry. +2. > Register the CLI utility at `$SMOKE_DIR/csv_summary.py` as a code named `csv-summary` so we can reuse it. +3. > Use the `scan-directory` code from the registry to scan `$SMOKE_DIR/data/`. +4. > Run the `csv-summary` code on `$SMOKE_DIR/data/samples.csv` and tell me the columns, row count, and any columns with null values. 5. > Put this in explicit memory: samples.csv has null values in the status and timestamp columns. 6. > Tell me what you remember about the samples dataset. -Exit the agent (`Ctrl+C` or `/exit`), then distill the session into episodic memory and stop the MLflow daemon: - -```bash -# 4. After your session, distill traces into episodic memory: -dsagt memory --project quickstart +This exercised: -# 5. Stop the MLflow daemon dsagt mlflow started (writes a PID into -# /.runtime; this releases the port and the gunicorn workers): -dsagt stop quickstart -``` - -What this exercised: - -| Prompt | Layer | +| Prompt | Capability | |---|---| -| 1 | Knowledge MCP server (`kb_ingest`) — chunks and indexes docs into ChromaDB | -| 2 | Registry MCP server (`save_tool_spec`) — writes `tools/csvcut.md`, `tools/csvgrep.md`, etc. (one per registered tool) | -| 3 | `dsagt-run` provenance wrapper — records exec layer to `trace_archive/` | -| 4 | Explicit memory (`kb_remember` → `explicit_memories.yaml`) + KB recall | +| 1 | `dsagt-server` (`kb_ingest`) — chunks and indexes docs into ChromaDB | +| 2 | `dsagt-server` (`save_code_spec`) — writes `codes/csv-summary/SKILL.md` (a skill-standard dir), wrapping the executable with `dsagt-run` | +| 3–4 | `dsagt-run` provenance wrapper — records each execution to `trace_archive/` | +| 5–6 | Explicit memory (`kb_remember` → `.dsagt/explicit_memories.yaml`) + KB recall (`kb_get_memories`) | -Verify the artifacts and view traces in the MLflow UI (URL printed by `dsagt mlflow`): +(`csv_summary.py` is stdlib-only — no dependency to install — so registration and execution work out of the box. Step 4's null-column finding is the fact you store and recall in 5–6.) -```bash -dsagt info quickstart -ls ~/dsagt-projects/quickstart/{tools,trace_archive} -cat ~/dsagt-projects/quickstart/explicit_memories.yaml -``` +Exit the agent (`Ctrl+C` or `/exit`), then verify the artifacts and view traces: -The same flow runs non-interactively via `dsagt smoke-test --agent claude` (or `goose` / `codex` / `opencode`), which asserts each artifact is present. +```bash +dsagt info quickstart # config + a session/trace summary +ls ~/dsagt-projects/quickstart/{codes,trace_archive} +cat ~/dsagt-projects/quickstart/.dsagt/explicit_memories.yaml -### First-time knowledge base setup +# Traces land in a serverless SQLite store. Browse them with: +mlflow ui --backend-store-uri sqlite:///$HOME/dsagt-projects/quickstart/mlflow.db +``` -`dsagt setup-kb` builds the shared ChromaDB collections under `~/.dsagt/kb_index/` that every project on this machine reuses. Three of the six collections shown in the [architecture diagram](#architecture) are populated here — the other three are per-project and fill in automatically during use (see [Knowledge Base](#knowledge-base) below): +The same flow runs non-interactively via `dsagt smoke-test --agent claude` (or `goose` / `codex` / `opencode` / `cline`), which asserts each artifact is present. -- **Tool Specs** — DSAgt's bundled tool specs from `src/dsagt/tools/`, tagged with `source: bundled` so the agent finds them via `search_registry` from the very first session. -- **Skills** — DSAgt's bundled skill workflows from `src/dsagt/skills/` (e.g. `datacard-generator`), discovered via `search_skills`. -- **Domain Knowledge** — Reference corpora (NVIDIA NeMo Curator, AI Data Readiness Inspector) downloaded and embedded so the agent has data-curation domain knowledge out of the box. +### Knowledge base setup -The Tool Specs and Skills collections are wipe-and-rebuild on every run, so re-run `setup-kb` after upgrading DSAgt to pick up new bundled assets. +`dsagt init` sets up the project's knowledge base with three kinds of collection: -```bash -dsagt setup-kb # all collections (local embedder, no creds) -dsagt setup-kb --collection nemo_curator -dsagt setup-kb --embedding-backend api --embedding-base-url ... --embedding-api-key ... -``` +- **Code Specs** — DSAgt's bundled code specs, always set up so the agent finds them via `search_registry` from the first session. +- **Skill Catalogs** — the skill-catalog sources you pick at init (default `genesis`) are cloned and indexed so `search_skills` returns installable skills. The bundled `skill-creator` is discovered natively by the agent. +- **Knowledge Collections** — optional reference document sets you pick at init (`nemo_curator`, `aidrin`), downloaded and indexed for data-curation domain knowledge. -The default embedder is a local sentence-transformers model (~130 MB of weights downloaded on first run, CPU-side, no API key). Pass `--embedding-backend api` to route through a hosted embedder via LiteLLM (15–30 minutes typical for the reference corpora, depending on rate limits). +The default embedder is a local sentence-transformers model (~130 MB of weights downloaded on first run). ## Use Case Examples -End-to-end walkthroughs for representative scientific and data-readiness scenarios live in [`use_cases/`](use_cases/). Each one covers data acquisition, tool registration, pipeline construction, and agent-driven execution against a real dataset. +End-to-end walkthroughs for representative scientific and data-readiness scenarios live in [`use_cases/`](use_cases/). Each one covers data acquisition, code registration, pipeline construction, and agent-driven execution against a real dataset. | Use case | Domain | Guide | |----------|--------|-------| @@ -115,86 +136,103 @@ End-to-end walkthroughs for representative scientific and data-readiness scenari ## Project Directory -Default location: `~/dsagt-projects//`. Override with `--location`: - -```bash -dsagt init my-project --agent claude --location /data/runs # /data/runs/my-project/ -dsagt init my-project --agent claude --location . # ./my-project/ -``` +`dsagt init` prompts for the project location, defaulting to `~/dsagt-projects//` (e.g. enter `/data/runs` to place it at `/data/runs/my-project/`, or `.` for the current directory). -Projects are registered in `~/.dsagt/projects.yaml` so `dsagt mlflow ` and `dsagt info ` work from any directory. The data layer (knowledge base, MLflow store, registered tools, skills, audit records) is agent-agnostic, so re-running `dsagt init --agent ` switches platforms while preserving everything you've accumulated. +Projects are registered in `~/dsagt-projects/projects.yaml` so `dsagt info ` works from any directory. The project's data — knowledge base, trace store, registered codes, skills, audit records — is agent-agnostic, so re-running `dsagt init` for the same project and choosing a different agent switches platforms while preserving everything you've accumulated (it prompts before any destructive change). ``` ~/dsagt-projects/cheese-metagenome/ - dsagt_config.yaml # project configuration - tools/ # registered CLI tool specs (markdown + YAML frontmatter) - tools/code/ # agent-written tool scripts + .dsagt/ # dsagt-internal state (hidden) + config.yaml # project configuration (set by dsagt init) + state.yaml # session log + memory cursor (owned by the MCP server) + explicit_memories.yaml # user-confirmed facts + codes// # registered codes — skill-standard dirs (SKILL.md + scripts/) skills/ # agent skills (SKILL.md + reference docs) - trace_archive/ # tool execution records (JSON, from dsagt-run) - mlflow/ # MLflow traces, metrics, artifacts + trace_archive/ # code execution records (JSON, from dsagt-run) + mlflow.db # serverless MLflow SQLite trace store kb_index/ # knowledge base vector collections - explicit_memories.yaml # user-confirmed facts loaded at session start # Per-agent runtime config (one of, generated by dsagt init): # claude: CLAUDE.md, .mcp.json # goose: goose.yaml, .goosehints # codex: AGENTS.md, .codex-data/config.toml # opencode: AGENTS.md, opencode.json - # roo: .roomodes, .roo/mcp.json # cline: .clinerules/, cline_mcp_settings.json (managed via cline mcp add) ``` -### MCP Servers +### MCP Server -- **Registry** (`dsagt-registry-server`) — Tool registration and dependency installation. Tools are markdown files with YAML frontmatter under `/tools/`. Executables are wrapped with `dsagt-run` for provenance and `uv run --with` for Python dependencies. The agent discovers tools via `search_registry`. -- **Knowledge** (`dsagt-knowledge-server`) — Semantic search over indexed document collections (FAISS / ChromaDB). Background jobs handle long ingest operations. The agent searches via `kb_search`, ingests via `kb_ingest`, and saves user-confirmed facts via `kb_remember`. +DSAGT exposes a single MCP server, **`dsagt-server`**, that an agent connects to once. Its main tool groups are: -### Tools and Skills +- **Registry** — Code registration and dependency installation. Codes are markdown files with YAML frontmatter under `/codes/`. Executables are wrapped with `dsagt-run` for provenance and `uv run --with` for Python dependencies. The agent discovers codes via `search_registry`. +- **Knowledge** — Semantic search over indexed ChromaDB document collections. Background jobs handle long ingest operations. The agent searches via `kb_search`, ingests via `kb_ingest`, and saves user-confirmed facts via `kb_remember`. -**Tools** are CLI executables defined as markdown files with YAML frontmatter in `/tools/`. The agent registers new tools via the registry MCP server's `save_tool_spec`. +### Codes and Skills -**Skills** are instruction-based agent workflows in `/skills/`. Each skill is a directory containing a `SKILL.md` and optional reference docs. DSAgt ships with a bundled `datacard-generator` skill. The agent discovers skills via `search_skills`. +**Codes** are CLI executables defined as markdown files with YAML frontmatter in `/codes/`. The agent registers new codes via the MCP server's `save_code_spec`. + +**Skills** are instruction-based agent workflows — a directory with a `SKILL.md` and optional reference docs. They come in two tiers: + +- **Installed** skills live in `/skills/` (DSAgt ships a bundled `skill-creator`; domain skills like the MODCON datacard generator are installed from the `genesis` catalog). These are mirrored into the agent's native skill directory (e.g. `.claude/skills/`, `.agents/skills/`) at install time (and re-mirrored at `dsagt init`/`start`), where the agent auto-discovers and auto-invokes them — no `search_skills` needed (that covers only the catalog tier below). +- **Catalog** skills come from external Git repositories — GitHub *or* GitLab — indexed into a searchable catalog the agent browses with `search_skills` but that is **not** loaded into its context (so a catalog can hold thousands of skills). The agent enables a source with `add_skill_source(...)`, finds skills with `search_skills(...)`, then copies one into the project with `install_skill(...)`. + +The catalog is **opt-in**: a source must be synced before its skills are searchable. Curated named sources ship out of the box — `k-dense-ai`, `anthropic`, `antigravity`, `composio`, and `genesis` (the OSTI GENESIS catalog: HPC, HuggingFace, LangChain, OpenAI, plasma-sim, and more) — and any Git URL or `owner/repo` works too. Manage catalogs from the agent with `list_skill_sources` / `add_skill_source` / `search_skills` / `install_skill`. + +![DSAgt skills routing](latex/skills-routing.png) + +The diagram traces a skill's lifecycle: **discovery** — browse the catalog with `search_skills` for skills the agent doesn't yet have → **install** — `install_skill` copies one into the project → **use** — the agent auto-discovers installed skills natively and invokes them by relevance (and authors new ones with the bundled `skill-creator`). The diagram source is [`latex/skills-routing.tex`](latex/skills-routing.tex). ### Knowledge Base -Six independently-partitioned ChromaDB collections hold everything the agent searches semantically. The first three are global (under `~/.dsagt/kb_index/`, populated by `dsagt setup-kb`); the last three are per-project (under `/kb_index/`, populated automatically during use): +The agent searches these collections semantically: | Collection | Source | Populated by | |---|---|---| -| **Tool Specs** | Bundled CLI tool specs in `src/dsagt/tools/` | `dsagt setup-kb` | -| **Skills** | Bundled skill workflows in `src/dsagt/skills/` | `dsagt setup-kb` | -| **Domain Knowledge** | NeMo Curator + AIDRIN reference corpora; user-ingested docs | `dsagt setup-kb` + agent's `kb_ingest` | -| **Explicit Memory** | User-confirmed facts | Agent's `kb_remember` (also written to `/explicit_memories.yaml`); the agent fetches via `kb_get_memories` on demand — typically when you ask it to recall — not auto-loaded at session start | -| **Episodic Memory** | Distilled facts from MLflow traces | `dsagt memory --project ` (per-category outlier detection via embedding centroids) | -| **Tool Use Records** | `dsagt-run` execution traces | `dsagt-run` wrapper writes JSON to `/trace_archive/`; indexed into ChromaDB by `dsagt memory` | +| **Code Specs** | Bundled CLI code specs | `dsagt init` (always set up) | +| **Skill Catalogs** | Installable skills from external repos (one per source) | `dsagt init` (chosen sources) + `add_skill_source` | +| **Knowledge Collections** | NeMo Curator + AIDRIN reference collections; user-ingested docs | `dsagt init` (chosen collections) + agent's `kb_ingest` | +| **Explicit Memory** | User-confirmed facts | Agent's `kb_remember` (also written to `/.dsagt/explicit_memories.yaml`); the agent fetches via `kb_get_memories` on demand, not auto-loaded at session start | +| **Code Execution Records** | `dsagt-run` execution traces | `dsagt-run` writes JSON to `/trace_archive/`; indexed for search during the session, and before `reconstruct_pipeline` | +| **Episodic Memory** | Captured session turns | **Opt-in** (enabled in the `dsagt init` menu): DSAgt captures each completed turn into `session_memory` during the session (mechanical chunk + embed). Retrieval is recency-weighted. | + +The embedding backend is local (sentence-transformers, CPU-side, no API key). -The default embedding backend is local (sentence-transformers, CPU-side, no API needed). Switch to `embedding.backend: api` in `dsagt_config.yaml` to route through a hosted embedder via LiteLLM. Cross-encoder reranking is optional (`knowledge.rerank: true`). +The agent searches via `kb_search` and writes via `kb_ingest` / `kb_remember`. Registered codes have their own `search_registry` route over the same backend. Installed skills are discovered natively by the agent; enabling external skill catalogs adds one collection per source, which `search_skills` browses for installable skills. -The agent searches via `kb_search` (knowledge MCP server) and writes via `kb_ingest` / `kb_remember`. Tool Specs and Skills are queried through specialized routes (`search_registry`, `search_skills`) over the same backend. +### Memory + +DSAgt has two memory types, both retrievable via `kb_search` / `kb_get_memories`: + +- **Explicit memory** — user-confirmed facts the agent writes with `kb_remember` (mirrored to `/.dsagt/explicit_memories.yaml`). Always on; degrades to pure-YAML if the vector store is unavailable. +- **Episodic memory** — automatic session capture, **opt-in** (enabled in the `dsagt init` menu, off by default). As the session runs, DSAgt reads the agent's transcript and captures each completed turn into the `session_memory` collection — a fast, local chunk-and-embed pass. + + Retrieval is **recency-weighted** (`episodic.recency_half_life_days`): a newer turn edges out a stale one, but as a bounded boost — a strongly-relevant old turn is never buried. ### Observability -MLflow runs at `http://localhost:` (pinned at init time, listed by `dsagt info`). The trace view shows: +Self-logging goes to a serverless MLflow SQLite store at `/mlflow.db`. Browse it with `mlflow ui --backend-store-uri sqlite:////mlflow.db`. The trace view shows: - **Knowledge base operations** — `kb.search` / `kb.embed` / `kb.index_search` / `kb.rerank` span trees with per-phase timing. -- **Tool executions** — `tool.execute` spans with exit code, duration, file counts, truncated stderr. Full payload in `trace_archive/.json`. -- **Registry events** — `save_tool_spec`, `install_dependencies`, `reconstruct_pipeline` spans. -- **Native agent OTel** *(optional)* — when you export `MLFLOW_TRACKING_URI` and `OTEL_EXPORTER_OTLP_ENDPOINT` (printed by `dsagt init`), the agent's own LLM-call traces land in the same MLflow store. Trace coverage varies by agent: claude / goose emit full payloads, codex emits token counts + tool names, opencode emits nothing natively. +- **Code executions** — `code.execute` spans with exit code, duration, file counts, truncated stderr. Full payload in `trace_archive/.json`. +- **Registry events** — `save_code_spec`, `install_dependencies`, `reconstruct_pipeline` spans. +- **Agent traces** — recovered from the on-disk session transcript, so prompts, responses, and tool calls land in the store for every supported agent. -Every span carries the project's `session.id` for filtering. Tool execution records on disk provide the canonical provenance chain — the agent calls `reconstruct_pipeline` to render the trace archive as a reproducible bash script or Snakemake workflow. +Each launch gets a session id that every span carries, so you can filter the trace view by session. The code-execution records on disk are the provenance record — the agent calls `reconstruct_pipeline` to render them as a reproducible bash script or Snakemake workflow. ## CLI Reference | Command | Description | |---------|-------------| -| `dsagt init --agent [--location ] [--mlflow-port N]` | Create a project; write per-agent MCP config; print launch one-liner | -| `dsagt mlflow ` | Run MLflow in the foreground against a project's store (port pinned at init time) | -| `dsagt memory --project ` | Distill new traces from this project's MLflow into episodic memory | -| `dsagt info [--json]` | Resolved config (with source per value) and a session/error summary | -| `dsagt setup-kb [--collection ]` | Build the shared core knowledge base collections | -| `dsagt list` | List all projects with agent, status, and path | +| `dsagt init` | Create or reconfigure a project — interactive menu for name, location, agent, knowledge collections, skill sources, and the episodic-memory opt-in; sets up the KB and writes the per-agent MCP config | +| `dsagt start ` | Launch the agent in the project directory (equivalent to `cd && `) | +| `dsagt info [--json]` | Resolved config (with source per value) and a session/trace summary | +| `dsagt list` | List all projects with agent and path | | `dsagt mv ` | Move a project to a new location | | `dsagt rm [-y] [--keep-files]` | Unregister a project (and optionally delete its directory) | -| `dsagt smoke-test [--agent claude\|goose\|codex\|opencode]` | End-to-end install verification | +| `dsagt smoke-test [--agent claude\|goose\|codex\|opencode\|cline]` | End-to-end install verification | + +> **Deprecated `dsagt init` flags (backcompat).** The pre-menu flags still work for scripts/CI but are deprecated in favor of the interactive menu: `` (positional), `--agent `, `--location `, `--include … | --exclude …`, and `--episodic`. Passing any of them skips the corresponding prompt; new usage should prefer bare `dsagt init`. + +Skill catalogs are managed from the agent via the MCP tools (`add_skill_source` / `search_skills` / `install_skill`), and traces are viewed with `mlflow ui --backend-store-uri sqlite:////mlflow.db`. -For tests, proxy mode, troubleshooting, and other developer-facing material, see [developer.md](developer.md). +For tests, troubleshooting, and other developer-facing material, see [developer.md](developer.md). diff --git a/agent-card.md b/agent-card.md index 1e0eb0e..6bd7ca2 100644 --- a/agent-card.md +++ b/agent-card.md @@ -8,108 +8,107 @@ tags: - science:general - risk:general license: Apache-2.0 -base_model: N/A # DSAgt is agent-platform-agnostic; it wraps Claude Code, Goose, Codex, opencode, Roo Code, or Cline +base_model: N/A # DSAgt is agent-platform-agnostic; it wraps Claude Code, Goose, Codex, opencode, or Cline datasets: - - # NeMo Curator reference corpus (indexed into global KB by dsagt setup-kb) - - # AI Data Readiness Inspector (AIDRIN) reference corpus + - # NeMo Curator reference corpus (optional knowledge collection, indexed at dsagt init) + - # AI Data Readiness Inspector (AIDRIN) reference corpus (optional knowledge collection) metrics: - - # Tool registration success rate + - # Code registration success rate - # Knowledge base search precision - # Provenance trace completeness agent_card: name: "DSAgt (DataSmith Agent)" - description: "AI-assisted data pipeline builder developed under the DOE Genesis Mission. Wraps an MCP-compatible agent CLI with tool registration, a semantic knowledge base, execution provenance, and observability infrastructure to accelerate AI-ready scientific data preparation." + description: "AI-assisted data pipeline builder developed under the DOE Genesis Mission. Wraps an MCP-compatible agent CLI with code registration, a semantic knowledge base, skill discovery, execution provenance, and observability infrastructure to accelerate AI-ready scientific data preparation." provider: organization: "DOE AI ModCon Base Data Team (DOE Genesis Mission)" url: "https://github.com/AI-ModCon/dsagt" - version: "0.1.0" + version: "0.2.0" documentation_url: "https://ai-modcon.github.io/dsagt/" protocol_version: "N/A" # DSAgt extends existing agent CLIs via MCP; it is not itself an A2A service - preferred_transport: "HTTP" + preferred_transport: "stdio" capabilities: streaming: false push_notifications: false - state_transition_history: true # trace_archive + MLflow record full execution history + state_transition_history: true # trace_archive + the serverless MLflow store record full execution history authentication: - schemes: - - "Bearer" # used by the optional LiteLLM proxy and hosted embedding backends - credentials: "N/A" # public repo; agent CLI auth is handled by the underlying platform (Claude, Goose, etc.) + schemes: [] # BYOA — the agent platform owns its own LLM-provider auth; DSAgt never proxies or stores credentials + credentials: "N/A (optional: EMBEDDING_API_KEY in the shell for a hosted embedding backend — never written to disk)" default_input_modes: - "text/plain" default_output_modes: - "text/plain" - - "application/json" # tool execution records, pipeline reconstructions + - "application/json" # code execution records, pipeline reconstructions skills: - - id: "tool_registration" - name: "Tool Registration" - description: "Register CLI tools as markdown files with YAML frontmatter; install dependencies via uv; wrap executions with dsagt-run for provenance." - tags: [registry, provenance, mcp] + - id: "code_registration" + name: "Code Registration" + description: "Register CLI codes as skill-standard directories (codes//SKILL.md) with machine-readable parameters; install dependencies via uv; every execution is wrapped with dsagt-run for provenance. Registered codes are also mirrored into the agent's native skills directory for in-context discovery." + tags: [registry, provenance, mcp, skills] examples: - - "Register csvkit tools csvcut, csvgrep, csvstat, csvlook." + - "Register this analysis script as a reusable code." - "Install Python dependencies for a custom analysis script." input_modes: ["text/plain"] output_modes: ["text/plain", "application/json"] - id: "knowledge_base" name: "Knowledge Base" - description: "Hybrid dense+sparse semantic search over six ChromaDB collections: Tool Specs, Skills, Domain Knowledge, Explicit Memory, Episodic Memory, and Tool Use Records." + description: "Hybrid dense+sparse (sentence-transformers + BM25) semantic search over ChromaDB collections: code specs, skill catalogs, domain knowledge, code-use records, and session memory. Optional cross-encoder reranking and regex/substring document filters." tags: [knowledge, chromadb, semantic-search, mcp] examples: - "Ingest domain documentation into a named collection." - - "Search for tools matching 'CSV statistics'." - - "Save a user-confirmed fact to explicit memory." + - "Search for codes matching 'CSV statistics'." input_modes: ["text/plain"] output_modes: ["text/plain", "application/json"] - id: "provenance" name: "Execution Provenance" - description: "Record every tool invocation (command, args, exit code, duration, file counts, stderr) to trace_archive/ and emit OTLP spans to MLflow. Reconstruct the full execution history as a bash script or Snakemake workflow." - tags: [provenance, mlflow, otlp, reproducibility] + description: "Record every code invocation (command, stdout/stderr, exit code, timing, file I/O) to trace_archive/ and emit spans to the project's serverless MLflow store. Reconstruct the full execution history as a dependency-ordered pipeline." + tags: [provenance, mlflow, reproducibility] examples: - - "Reconstruct a reproducible pipeline from prior tool executions." - - "View tool execution traces in the MLflow UI." + - "Reconstruct a reproducible pipeline from prior code executions." + - "View execution traces in the MLflow UI." input_modes: ["text/plain"] - output_modes: ["text/plain", "application/x-sh", "text/x-snakemake"] + output_modes: ["text/plain", "application/json"] - - id: "episodic_memory" - name: "Episodic Memory Distillation" - description: "Distill new MLflow traces into per-category episodic memory using embedding-centroid outlier detection. Surfaces novel facts from prior sessions." - tags: [memory, mlflow, distillation] + - id: "memory" + name: "Explicit + Episodic Memory" + description: "Explicit memory stores user-confirmed facts (YAML + vector mirror) via kb_remember / kb_get_memories. Opt-in episodic memory mechanically chunks, tags, and embeds every session turn into a recency-weighted searchable collection — no LLM calls." + tags: [memory, chromadb] examples: - - "Distill today's session traces into episodic memory." + - "Put this in explicit memory: samples.csv has null values in the status column." + - "What do you remember about the samples dataset?" input_modes: ["text/plain"] output_modes: ["text/plain"] - - id: "datacard_generator" - name: "Datacard Generator" - description: "Bundled skill workflow for generating standardized dataset documentation (datacards) from indexed knowledge and agent-driven analysis." - tags: [datacard, documentation, skill] + - id: "skill_discovery" + name: "Skill Discovery and Installation" + description: "Search external skill catalogs (Genesis, Anthropic, K-Dense, and others cloned+indexed at init) and install skills into the project, where the agent auto-discovers them natively. Agents can also author and save their own skills." + tags: [skills, catalog, mcp] examples: - - "Generate a datacard for samples.csv." + - "Search the skill catalog for a literature-search skill and install it." input_modes: ["text/plain"] output_modes: ["text/plain", "text/markdown"] Extensions: agent_runtime: - framework: "MCP (Model Context Protocol) over stdio; supported agent platforms: Claude Code, Goose, Codex, opencode, Roo Code, Cline" - service_endpoint: "stdio (MCP servers launched as subprocesses by the agent platform)" - rate_limits: "Determined by the underlying LLM provider and optional LiteLLM proxy configuration." - logging: "OTLP HTTP spans emitted to local MLflow instance (http://localhost:/v1/traces); full tool execution records written to /trace_archive/" - memory: "Stateful per-project. Explicit memory: /explicit_memories.yaml + ChromaDB. Episodic memory: ChromaDB (distilled by dsagt memory). Tool use records: /trace_archive/ + ChromaDB." + framework: "MCP (Model Context Protocol) over stdio; supported agent platforms: Claude Code, Goose, Codex, opencode, Cline" + service_endpoint: "stdio (the single dsagt-server is launched as a subprocess by the agent platform)" + rate_limits: "Determined by the underlying LLM provider configured in the agent platform (BYOA — DSAgt never proxies LLM traffic)." + logging: "Serverless MLflow store at sqlite:////mlflow.db (no server to run); full code execution records written to /trace_archive/. Agent LLM-call history is recovered post-hoc from the agent's on-disk transcript." + memory: "Stateful per-project. Explicit memory: /.dsagt/explicit_memories.yaml + ChromaDB mirror. Episodic memory (opt-in): session_memory ChromaDB collection, embedded per turn. Code-use records: /trace_archive/ + code_use ChromaDB collection." --- # DSAgt (DataSmith Agent) -DSAgt is an AI-assisted data pipeline builder. It connects an MCP-compatible agent CLI (Claude Code, Goose, Codex, opencode, Roo Code, or Cline) to tool registration, a semantic knowledge base, execution provenance, and observability infrastructure — without modifying the agent itself. +DSAgt is an AI-assisted data pipeline builder. It connects an MCP-compatible agent CLI (Claude Code, Goose, Codex, opencode, or Cline) to code registration, a semantic knowledge base, skill discovery, execution provenance, and observability infrastructure — without modifying the agent itself. -*Last Updated*: **2026-06-30** +*Last Updated*: **2026-07-02** ## Developed by @@ -129,28 +128,29 @@ See https://github.com/AI-ModCon/dsagt/graphs/contributors for full list. ## Agent Changelog ++ **2026-07-02** v0.2.0 — single merged `dsagt-server` (20 tools); serverless SQLite MLflow store (no ports, no OTel, no proxy); external skill catalogs; agent-transcript trace pipeline + opt-in episodic memory; codes stored as skill-standard directories mirrored into the agent's native skills dir + **2026-06-30** initial public version (v0.1.0) ## Agent short description -Tool-using scaffolding layer that gives any MCP-compatible agent CLI persistent tool registration, semantic knowledge retrieval, execution provenance, and session memory — exposed via two MCP servers (Registry and Knowledge). +Scaffolding layer that gives any MCP-compatible agent CLI persistent code registration, semantic knowledge retrieval, skill discovery, execution provenance, and session memory — exposed as 20 tools on a single MCP server (`dsagt-server`). ## Agent description -DSAgt wraps an unmodified agent CLI with four independently-operable layers, each implemented as an MCP server the agent discovers through the standard MCP tool protocol: +DSAgt wraps an unmodified agent CLI with four independently-operable concerns, exposed by one MCP server the agent discovers through the standard MCP tool protocol: -1. **Tool Registry** — The agent registers CLI tools as markdown files with YAML frontmatter; the registry server handles dependency installation (`uv run --with`) and wraps each invocation with `dsagt-run` for provenance. Discovery is via `search_registry`. -2. **Knowledge Base** — Six independently-partitioned ChromaDB collections with hybrid dense (sentence-transformers) + sparse (BM25) search. Global collections (Tool Specs, Skills, Domain Knowledge) are populated by `dsagt setup-kb`; per-project collections (Explicit Memory, Episodic Memory, Tool Use Records) fill in during use. Discovery is via `kb_search`, `kb_ingest`, `kb_remember`, `search_skills`. -3. **Provenance** — `dsagt-run` captures every tool execution (command, args, exit code, duration, file I/O, stderr) to `trace_archive/` and emits OTLP spans to MLflow. `reconstruct_pipeline` renders the archive as a reproducible bash or Snakemake script. -4. **Observability** — Local MLflow instance with OTLP ingestion. All four layers emit spans. The agent's own LLM-call traces land in the same store when `OTEL_EXPORTER_OTLP_ENDPOINT` is set (printed by `dsagt init`). +1. **Code Registry** — The agent registers CLI codes as skill-standard directories (`codes//SKILL.md`, frontmatter carrying executable + parameters); the server handles dependency installation (`uv run --with`) and wraps each stored command with `dsagt-run` for provenance. Discovery is dual-path: semantic search via `search_registry`, plus a mirror into the agent's native skills directory so the exact runnable command is in context at invocation time. +2. **Knowledge Base** — ChromaDB collections with hybrid dense (sentence-transformers) + sparse (BM25) search and optional cross-encoder reranking. Code specs and selected skill catalogs are indexed at `dsagt init`; per-project collections (code-use records, session memory) fill in during use. Long ingests run as background jobs. +3. **Provenance** — `dsagt-run` captures every code execution (command, stdout/stderr, exit code, timing, file I/O) to `trace_archive/` and emits spans to the project's serverless MLflow store. `reconstruct_pipeline` renders the archive as a dependency-ordered execution history. +4. **Observability & Memory** — All self-logging lands in `sqlite:////mlflow.db` (MLflow's serverless backend — nothing to run). Agent LLM-call traces are recovered post-hoc from the agent's on-disk transcript, uniformly across all five platforms. Explicit memory stores user-confirmed facts; opt-in episodic memory embeds every session turn for recency-weighted recall. -The data layer is agent-platform-agnostic: switching platforms preserves all accumulated knowledge, tools, and traces. +The data layer is agent-platform-agnostic: switching platforms preserves all accumulated knowledge, codes, skills, and traces. ## Underlying model(s) -- Primary model(s): N/A — DSAgt is platform-agnostic and delegates LLM calls to the configured agent CLI (Claude, GPT-4o, etc.) -- Embedding model: `sentence-transformers` (local, CPU-side, default); optionally any LiteLLM-compatible hosted embedder -- Cross-encoder reranking: optional, enabled via `knowledge.rerank: true` in `dsagt_config.yaml` +- Primary model(s): N/A — DSAgt is platform-agnostic and delegates LLM calls to the configured agent CLI (BYOA; DSAgt never proxies LLM traffic) +- Embedding model: `sentence-transformers` (`bge-small-en-v1.5`, local, CPU-side, default); optionally any OpenAI-compatible hosted embedder (`embedding.backend: api`) +- Cross-encoder reranking: optional, enabled via `knowledge.rerank: true` in `.dsagt/config.yaml` ## Inputs and outputs @@ -159,121 +159,90 @@ The data layer is agent-platform-agnostic: switching platforms preserves all acc - defaultInputModes: `["text/plain"]` - defaultOutputModes: `["text/plain", "application/json"]` -The agent accepts natural-language instructions (text). Outputs include text responses, registered tool specs (markdown), execution trace records (JSON), pipeline reconstructions (bash / Snakemake), and datacard documents (markdown). +The agent accepts natural-language instructions (text). Outputs include text responses, registered code specs (skill-standard markdown), execution trace records (JSON), pipeline reconstructions, and installed skills (markdown). ### Skills -- **Skill ID**: `tool_registration` - **Name**: Tool Registration - **Description**: Register CLI tools as markdown+YAML specs; install dependencies; wrap executions with `dsagt-run` for provenance. - **Tags**: registry, provenance, mcp - **Examples**: "Register csvkit tools csvcut, csvgrep, csvstat, csvlook.", "Install Python dependencies for a custom analysis script." +- **Skill ID**: `code_registration` + **Name**: Code Registration + **Description**: Register CLI codes as skill-standard spec directories; install dependencies; wrap executions with `dsagt-run` for provenance; mirror specs into the agent's native skills dir. + **Tags**: registry, provenance, mcp, skills + **Examples**: "Register this analysis script as a reusable code.", "Install Python dependencies for a custom analysis script." **Input/Output Modes**: text/plain → text/plain, application/json - **Skill ID**: `knowledge_base` **Name**: Knowledge Base - **Description**: Hybrid semantic search (dense + BM25) over Tool Specs, Skills, Domain Knowledge, Explicit Memory, Episodic Memory, and Tool Use Records collections. + **Description**: Hybrid semantic search (dense + BM25, optional reranking, regex/substring filters) over code specs, skill catalogs, domain knowledge, code-use records, and session memory. **Tags**: knowledge, chromadb, semantic-search, mcp - **Examples**: "Ingest domain documentation into a named collection.", "Search for tools matching 'CSV statistics'.", "Save a user-confirmed fact to explicit memory." + **Examples**: "Ingest domain documentation into a named collection.", "Search for codes matching 'CSV statistics'." **Input/Output Modes**: text/plain → text/plain, application/json - **Skill ID**: `provenance` **Name**: Execution Provenance - **Description**: Record every tool invocation to `trace_archive/` + MLflow; reconstruct full execution history as a bash script or Snakemake workflow. - **Tags**: provenance, mlflow, otlp, reproducibility - **Examples**: "Reconstruct a reproducible pipeline from prior tool executions." - **Input/Output Modes**: text/plain → text/plain, application/x-sh - -- **Skill ID**: `episodic_memory` - **Name**: Episodic Memory Distillation - **Description**: Distill MLflow traces into per-category episodic memory using embedding-centroid outlier detection (`dsagt memory --project `). - **Tags**: memory, mlflow, distillation - **Examples**: "Distill today's session traces into episodic memory." + **Description**: Record every code invocation to `trace_archive/` + the serverless MLflow store; reconstruct the full execution history as a dependency-ordered pipeline. + **Tags**: provenance, mlflow, reproducibility + **Examples**: "Reconstruct a reproducible pipeline from prior code executions." + **Input/Output Modes**: text/plain → text/plain, application/json + +- **Skill ID**: `memory` + **Name**: Explicit + Episodic Memory + **Description**: User-confirmed facts via `kb_remember` / `kb_get_memories` (YAML + vector mirror); opt-in episodic memory embeds every turn for recency-weighted cross-session recall (mechanical — no LLM calls). + **Tags**: memory, chromadb + **Examples**: "Put this in explicit memory: samples.csv has null values in the status column." **Input/Output Modes**: text/plain → text/plain -- **Skill ID**: `datacard_generator` - **Name**: Datacard Generator - **Description**: Bundled skill workflow (`src/dsagt/skills/datacard-generator/`) for generating standardized dataset documentation. - **Tags**: datacard, documentation, skill - **Examples**: "Generate a datacard for samples.csv." - **Input/Output Modes**: text/plain → text/markdown +- **Skill ID**: `skill_discovery` + **Name**: Skill Discovery and Installation + **Description**: Search external skill catalogs cloned + indexed at init (Genesis, Anthropic, K-Dense, Composio, and others); install skills into the project for native auto-discovery; save agent-authored skills. + **Tags**: skills, catalog, mcp + **Examples**: "Search the skill catalog for a literature-search skill and install it." + **Input/Output Modes**: text/plain → text/plain, text/markdown ### Tools and permissions -- Tool: `search_registry` (Registry MCP server) - - Purpose: Semantic search over registered tool specs - - Inputs: query string, optional tag filters - - Outputs: ranked list of matching tool specs - - Side effects: reads data - - Required permissions: none - -- Tool: `save_tool_spec` (Registry MCP server) - - Purpose: Register a new CLI tool as a markdown file with YAML frontmatter - - Inputs: tool name, command, dependencies, tags, usage description - - Outputs: path to written tool spec file - - Side effects: writes data to `/tools/` - - Required permissions: filesystem write access to project directory - -- Tool: `install_dependencies` (Registry MCP server) - - Purpose: Install tool dependencies via `uv run --with` - - Inputs: list of package names - - Outputs: installation confirmation - - Side effects: executes jobs (uv), network calls (PyPI) - - Required permissions: network access, uv installed - -- Tool: `reconstruct_pipeline` (Registry MCP server) - - Purpose: Render the trace archive as a reproducible bash script or Snakemake workflow - - Inputs: optional time range or session filter - - Outputs: bash or Snakemake script - - Side effects: reads data from `/trace_archive/` - - Required permissions: none - -- Tool: `kb_search` (Knowledge MCP server) - - Purpose: Hybrid semantic search over one or more knowledge collections - - Inputs: query string, collection name(s), optional filters - - Outputs: ranked document chunks with metadata - - Side effects: reads data - - Required permissions: none - -- Tool: `kb_ingest` (Knowledge MCP server) - - Purpose: Index a file or directory into a named ChromaDB collection (runs in background for large corpora) - - Inputs: file path or directory, collection name - - Outputs: ingestion job ID; status queryable - - Side effects: reads data, writes data to `/kb_index/` - - Required permissions: filesystem read access to source files - -- Tool: `kb_remember` (Knowledge MCP server) - - Purpose: Save a user-confirmed fact to explicit memory - - Inputs: fact string - - Outputs: confirmation - - Side effects: writes data to `/explicit_memories.yaml` and ChromaDB - - Required permissions: none - -- Tool: `kb_get_memories` (Knowledge MCP server) - - Purpose: Retrieve explicit and episodic memories for the current project - - Inputs: optional query string - - Outputs: list of memory entries - - Side effects: reads data - - Required permissions: none - -- Tool: `search_skills` (Knowledge MCP server) - - Purpose: Discover agent skill workflows - - Inputs: query string - - Outputs: ranked list of matching skills with SKILL.md content - - Side effects: reads data - - Required permissions: none +All 20 tools live on the single `dsagt-server` (stdio), split across four concerns. + +**Registry (8):** + +- `search_registry` — semantic search over registered + bundled code specs. Side effects: reads data. +- `get_registry` — list every registered code with its MCP-compatible schema. Side effects: reads data. +- `save_code_spec` — register a code as `codes//SKILL.md` (executable auto-wrapped with `dsagt-run` + `uv run --with`). Side effects: writes to the project dir; indexes into ChromaDB. +- `install_dependencies` — install a code's Python dependencies via uv. Side effects: executes uv, network calls (PyPI). +- `run_command` — execute a shell command with a timeout. Side effects: executes subprocesses. +- `read_file` — read a file from disk. Side effects: reads data. +- `http_request` — issue an HTTP(S) request. Side effects: network calls. +- `reconstruct_pipeline` — render `trace_archive/` as a dependency-ordered execution history. Side effects: reads data. + +**Knowledge (5):** + +- `kb_search` — hybrid semantic search over one or more collections (optional metadata, regex, and substring filters). Side effects: reads data. +- `kb_ingest` — index a file or directory into a named collection (background job for large corpora). Side effects: reads sources, writes `/kb_index/`. +- `kb_append` — add documents to an existing collection (background job). Side effects: writes `/kb_index/`. +- `kb_list_collections` — list collections with document counts. Side effects: reads data. +- `kb_job_status` — poll a background ingest/append job. Side effects: none. + +**Memory (2):** + +- `kb_remember` — save a user-confirmed fact to explicit memory. Side effects: writes `/.dsagt/explicit_memories.yaml` + ChromaDB. +- `kb_get_memories` — retrieve explicit memories (optionally query-filtered). Side effects: reads data. + +**Skills (5):** + +- `search_skills` — rank installable skills across synced external catalogs. Side effects: reads data. +- `install_skill` — copy a catalog skill into `/skills/` (with upstream attribution). Side effects: writes to the project dir. +- `save_skill` — register an agent-authored skill into `/skills/`. Side effects: writes to the project dir. +- `add_skill_source` — clone + index a new external skill catalog. Side effects: network calls (git), writes `kb_index/`. +- `list_skill_sources` — list known/synced catalog sources. Side effects: reads data. ### Service endpoint and discovery - Base URL: `https://github.com/AI-ModCon/dsagt` -- MCP servers run as local subprocesses; there is no remote HTTP endpoint by default. -- Registry server: `dsagt-registry-server` (stdio) -- Knowledge server: `dsagt-knowledge-server` (stdio) -- Optional LiteLLM proxy: `dsagt-proxy` (HTTP, for routing embeddings/LLM calls to hosted providers) +- The MCP server runs as a local subprocess; there is no remote HTTP endpoint. +- Server: `dsagt-server` (stdio), self-sufficient — it derives the project from its working directory (`.dsagt/config.yaml`). ## Runtime Infrastructure -DSAgt runs locally as a CLI tool. Both MCP servers are launched as subprocesses by the configured agent platform. MLflow runs locally at a port pinned at `dsagt init` time. +DSAgt runs locally as a CLI tool. The MCP server is launched as a subprocess by the configured agent platform. There are no services to run: all self-logging goes to a serverless SQLite MLflow store (`sqlite:////mlflow.db`), browsable on demand with `mlflow ui --backend-store-uri sqlite:////mlflow.db`. ### Hardware @@ -281,25 +250,17 @@ Runs on any developer workstation or compute node with Python 3.12+. The default ### Software -Python 3.12+, `uv` package manager. Key dependencies: +Python 3.12 or 3.13, `uv` package manager. Key dependencies: - `mcp>=1.0.0` — MCP server framework +- `mlflow==3.11.1` — trace store and observability (serverless SQLite backend) - `chromadb>=1.5.1` — vector store -- `faiss-cpu>=1.8` — FAISS index backend -- `sentence-transformers==5.4.0` — local embedding model -- `llama-index-core>=0.11` — document chunking and indexing -- `mlflow>=3.11.1,<4.0` — observability and trace storage -- `opentelemetry-*>=1.27` — OTLP span emission -- `litellm[proxy]>=1.83.7` — optional hosted embedder / LLM proxy routing +- `sentence-transformers==5.4.0` — local embeddings and reranking +- `llama-index-core>=0.11` — document and code chunking - `rank-bm25>=0.2.2` — sparse keyword retrieval for hybrid search +- `questionary>=2.0` — interactive `dsagt init` menus -```txt -# To get full dependency list: -uv sync --all-groups -pip freeze -``` - -See `pyproject.toml` for the complete pinned dependency set. +See `pyproject.toml` for the complete dependency set. ## Papers and Scientific Outputs @@ -331,107 +292,106 @@ Tested Use cases include: - Microbial genomics pipelines (short-read QC and assembly) - Cryo-EM data curation (EMPIAR datasets) - Materials science DFT workflows (VASP via ISAAC) -- Tokamak stability analysis (fusion energy) -- Combustion flow simulations (BlastNet) +- Tokamak stability analysis (fusion energy, M3D-C1) +- AI data readiness assessment (AIDRIN) ## Out-of-Scope Use Cases -- Using DSAgt on controlled or proprietary data that should not be shared with configured LLM inference provider. +- Using DSAgt on controlled or proprietary data that should not be shared with the configured LLM inference provider. # How to use ## Install Instructions ```bash +# For use: +pip install "git+https://github.com/AI-ModCon/dsagt.git" + +# For development: git clone https://github.com/AI-ModCon/dsagt.git cd dsagt -uv sync +uv sync --all-groups source .venv/bin/activate -# Build shared knowledge base collections (Tool Specs, Skills, Domain Knowledge): -dsagt setup-kb - -# Create a project: -dsagt init my-project --agent claude # or goose / codex / opencode / roo / cline +# Create a project (interactive: pick agent platform, knowledge +# collections, and skill-catalog sources; the knowledge base is +# provisioned on first run): +dsagt init ``` ## Agent configuration -- **System prompt / instructions**: generated by `dsagt init` as `CLAUDE.md` (Claude Code), `AGENTS.md` (Codex/opencode), `.goosehints` (Goose), or `.roomodes` (Roo Code) -- **MCP server config**: generated by `dsagt init` as `.mcp.json` (Claude Code), `goose.yaml`, `.codex-data/config.toml`, `opencode.json`, or `.roo/mcp.json` -- **Embedding backend**: set `embedding.backend: api` in `/dsagt_config.yaml` to route through a hosted embedder via LiteLLM; provide `embedding.base_url` and `embedding.api_key` -- **Reranking**: set `knowledge.rerank: true` in `dsagt_config.yaml` -- **Observability**: `dsagt mlflow ` prints `OTEL_EXPORTER_OTLP_ENDPOINT` and `MLFLOW_TRACKING_URI` exports for the agent's own traces +- **System prompt / instructions**: generated by `dsagt init` as `CLAUDE.md` (Claude Code), `AGENTS.md` (Codex/opencode), `.goosehints` (Goose), or `.clinerules/dsagt_instructions.md` (Cline) +- **MCP server config**: generated by `dsagt init` as `.mcp.json` (Claude Code), `goose.yaml`, `.codex-data/config.toml`, `opencode.json`, or via `cline mcp add` +- **LLM provider auth**: owned entirely by the agent platform (BYOA) — configure the agent before pointing DSAgt at it; DSAgt never stores or proxies credentials +- **Embedding backend**: set `embedding.backend: api` in `.dsagt/config.yaml` to use an OpenAI-compatible hosted embedder; the key comes from `EMBEDDING_API_KEY` in the shell (never on disk) +- **Reranking**: set `knowledge.rerank: true` in `.dsagt/config.yaml` +- **Episodic memory**: opt in at init (`dsagt init --episodic` or the interactive prompt) ## Invocation / integration ```bash -# Start MLflow and get OTel exports: -dsagt mlflow my-project +# Launch the agent from the project directory (no env exports, no services): +cd ~/dsagt-projects/my-project && claude # …or goose / codex / opencode / cline -# Paste the printed export block, then launch the agent from the project directory: -cd ~/dsagt-projects/my-project && claude +# Or let dsagt own the launch + post-session extraction trigger: +dsagt start my-project ``` -The agent discovers DSAgt tools via MCP and can invoke `search_registry`, `kb_search`, `save_tool_spec`, etc. directly in the conversation. +The agent discovers DSAgt tools via MCP and can invoke `search_registry`, `kb_search`, `save_code_spec`, etc. directly in the conversation. Registered codes and installed skills also appear in the agent's native skills directory after the next `dsagt start`. # Code snippets of how to use the agent ```bash # Full quickstart (see README for step-by-step prompts): -export SMOKE_DIR="$(pwd)/tests/smoke_test" dsagt init quickstart --agent claude -dsagt mlflow quickstart -# paste printed exports, then: -cd ~/dsagt-projects/quickstart && claude +dsagt start quickstart # After the session: -dsagt memory --project quickstart -dsagt stop quickstart -dsagt info quickstart +dsagt info quickstart # config + session/trace summary +mlflow ui --backend-store-uri sqlite:///$HOME/dsagt-projects/quickstart/mlflow.db -# Non-interactive smoke test: +# Non-interactive smoke test (asserts every artifact): dsagt smoke-test --agent claude ``` ```python -# DSAgt MCP servers are invoked by the agent platform, not directly from Python. +# The DSAgt MCP server is invoked by the agent platform, not directly from Python. # To integrate programmatically, use the MCP Python SDK: # from mcp.client.stdio import StdioServerParameters -# See docs/mcp-servers.md for server command and tool schemas. +# Server command: dsagt-server (run from the project directory). ``` # Limitations ## Risks -DSAgt executes arbitrary CLI tools registered by the agent. The registry server wraps tool invocations with `dsagt-run`, but does not sandbox or restrict what commands the agent can register or execute. Users should review tool specs before registration and restrict filesystem access as appropriate. +DSAgt executes arbitrary CLI codes registered by the agent. The registry wraps code invocations with `dsagt-run` for provenance, but does not sandbox or restrict what commands the agent can register or execute. Users should review code specs before registration and restrict filesystem access as appropriate. ### Agent-specific risk notes (tool use) -- **Tool execution side effects**: Registered tools can read/write files, make network calls, and execute arbitrary subprocesses. The agent must be trusted to register only appropriate tools. -- **Prompt injection**: Knowledge base documents are retrieved and injected into the agent context; malicious content in indexed documents could influence agent behavior. -- **Secrets handling**: API keys for hosted embedding backends are stored in `dsagt_config.yaml`. Do not commit this file to version control if it contains secrets. The LiteLLM proxy credential file is similarly sensitive. +- **Code execution side effects**: Registered codes can read/write files, make network calls, and execute arbitrary subprocesses. The agent must be trusted to register only appropriate codes. +- **Prompt injection**: Knowledge base documents and installed catalog skills are retrieved and injected into the agent context; malicious content in indexed documents or third-party skill catalogs could influence agent behavior. +- **Secrets handling**: No credentials are written to disk by DSAgt. A hosted embedding backend reads `EMBEDDING_API_KEY` from the shell at runtime. - **Data exfiltration**: If a hosted embedding backend is configured, document chunks are sent to that external service during ingestion and search. ## Limitations - Local-first: designed for single-user local or HPC use; no multi-user access control - Embedding model quality: default local `sentence-transformers` model (~130 MB) is effective for general text but may underperform on highly domain-specific technical corpora -- Agent OTel coverage varies: Claude Code and Goose emit full LLM-call traces; Codex emits token counts and tool names; opencode emits nothing natively -- Memory distillation requires MLflow traces to be present; sessions without MLflow running will not generate episodic memories +- Agent LLM-call traces are recovered post-hoc from the agent's on-disk transcript (uniform across all five platforms) — recovery granularity follows what each platform records +- Cline batch mode is unsupported (cline's provider rewrites unrecognized model names); interactive cline use works - No GUI: all interaction is through the agent CLI or the MLflow web UI # Agent evaluation details -- **Smoke test**: `dsagt smoke-test --agent ` runs the full quickstart non-interactively and asserts all artifacts (tool specs, trace records, explicit memory) are produced -- **Unit and integration tests**: `pytest tests/` (requires `uv sync --all-groups`); integration tests marked with `@pytest.mark.integration` can be deselected with `-m 'not integration'` -- **Tool-call correctness**: verified by checking `trace_archive/` records for expected exit codes and file counts -- **Knowledge base precision**: evaluated informally via `kb_search` recall in the smoke test +- **Smoke test**: `dsagt smoke-test --agent ` runs two full agent sessions non-interactively and asserts 18 artifacts: code registration + execution provenance, knowledge ingest + retrieval, skill catalog install, native skill mirroring, explicit + episodic memory, cross-session recall, agent-trace recovery, and session state +- **Unit tests**: `uv run python -m pytest tests/` (~620 tests; integration tests requiring credentials live in `test_*_integration.py`) +- **Code-call correctness**: verified by checking `trace_archive/` records for expected exit codes and captured output +- **Knowledge base precision**: evaluated via retrieval assertions in the smoke test (the agent must answer from ingested docs) # More Information - Full documentation: https://ai-modcon.github.io/dsagt/ - Source code: https://github.com/AI-ModCon/dsagt -- Architecture diagram: `docs/assets/architecture.png` -- Use case walkthroughs: `use_cases/` (Microbial Isolates, Cryo-EM, ISAAC/VASP, Tokamak Stability, Combustion) +- Use case walkthroughs: `use_cases/` (Microbial Isolates, Cryo-EM, ISAAC/VASP, Tokamak Stability, AIDRIN) diff --git a/developer.md b/developer.md deleted file mode 100644 index 9be7849..0000000 --- a/developer.md +++ /dev/null @@ -1,47 +0,0 @@ -# DSAgt — Developer Notes - -Material that's useful once you start poking at internals or running modes beyond the default `dsagt init` → `dsagt mlflow` → agent flow. - -## Tests - -```bash -uv run python -m pytest -m "not integration" # unit tests, no creds required -uv run python -m pytest -m integration -v # integration tests (require .env) -``` - -Integration tests read endpoint and key values from `.env` at the repo root. Copy `.env.example` to `.env` and fill in your values. - -For per-flow hand-tests (CLI, proxy mode, VS Code extensions), see the scripts under [`tests/smoke_test/manual_runs/`](tests/smoke_test/manual_runs/). - -## Advanced: Proxy mode - -`dsagt init` followed by `dsagt start --enable-proxy` spawns a LiteLLM proxy in front of your agent's LLM calls. This adds: - -- Full LLM-call traces (request bodies, tool-use blocks, response payloads) in MLflow for agents whose native OTel doesn't emit those payloads (codex, opencode) or that don't emit OTel at all. -- Cache-breakpoint injection on outgoing requests (Anthropic prompt caching). -- Sidechannel detection for agent-internal title-generator / session-namer calls. -- Model-name aliasing — useful when an agent CLI hardcodes a model whitelist incompatible with your gateway's served names (cline, roo). - -Proxy mode reads upstream LLM credentials from `.env` (or the shell). See the prerequisites and the full setup walkthrough at [`tests/smoke_test/manual_runs/proxy_walkthrough.md`](tests/smoke_test/manual_runs/proxy_walkthrough.md). - -## Troubleshooting - -**Agent command not found.** The agent CLI isn't installed or isn't on PATH — see the install table in the [README](README.md#installation). - -**MCP servers not connecting.** Verify uv resolves the server commands: - -```bash -uv run which dsagt-registry-server -uv run which dsagt-knowledge-server -``` - -If missing, reinstall: `uv sync --reinstall`. - -**MLflow UI empty.** Confirm MLflow is running for the right project: - -```bash -dsagt info # shows the pinned port -curl http://localhost: -``` - -**Claude keychain conflict.** If `claude` won't authenticate against a non-default gateway, run `claude /logout` to clear the macOS Keychain OAuth, then re-export `ANTHROPIC_BASE_URL` / `ANTHROPIC_API_KEY` and re-launch. diff --git a/docs/architecture.md b/docs/architecture.md index 5724f82..f7adcaf 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -2,42 +2,48 @@ ![DSAgt architecture](assets/architecture.png) -DSAgt wraps an unmodified agent CLI with four independently-operable layers. Each layer exposes its own MCP server so the agent discovers and invokes capabilities through the standard MCP tool protocol. +DSAgt provides a preconfigured agent platform with augmented capabilities for AI-ready scientific data processing and curation. Its capabilities are designed to complement rather than compete with the fast-moving agent platforms it runs on — Claude Code, Codex, Cline, opencode, and Goose. Most are exposed to the agent through a central MCP server: skill discovery and installation, data-processing code execution with provenance, knowledge-base extension and retrieval, and explicit and cross-session memory. Others are spawned by the server and run in the background: observability through MLflow traces, episodic memory management, and vector-store indexing. DSAgt also ships a suite of agent skills for AI-ready data workflows. -## Layers +## Capabilities -**Tool Registry** (`dsagt-registry-server`) -The agent registers CLI tools as markdown files with YAML frontmatter under `/tools/`. The registry server handles dependency installation via `uv run --with` and wraps every execution with `dsagt-run` for provenance capture. The agent discovers tools via `search_registry`. +**Code Registry** (`dsagt-server`) +The agent registers CLI data processing codes as markdown files with YAML frontmatter under `/codes/`. DSAgt handles dependency installation via `uv run --with` and wraps every execution with `dsagt-run` for provenance capture. The agent discovers codes via `search_registry`. -**Knowledge Base** (`dsagt-knowledge-server`) -Semantic search over six independently-partitioned ChromaDB collections. Three are global (populated by `dsagt setup-kb`); three are per-project (filled automatically during use). Background jobs handle long ingest operations. The agent searches via `kb_search`, ingests via `kb_ingest`, and saves user-confirmed facts via `kb_remember`. +**Knowledge Base** (`dsagt-server`) +Hybrid semantic + BM25 search over ChromaDB collections partitioned by concern — code specs, skill catalogs, scientific documents, and per-project memory — served by the same process as the code registry (one shared embedder, one ChromaDB). A first `dsagt init` installs the bundled code specs and a default skill catalog; heavier scientific collections (NeMo Curator, AIDRIN) and additional catalogs are opt-in, and new collections can be added for the documents of a specific scientific endeavor. Long ingests run as background jobs. The agent searches via `kb_search`, ingests via `kb_ingest`, and saves user-confirmed facts via `kb_remember`. Opt-in episodic memory distills each session turn into the per-project `session_memory` collection. **Provenance** (`dsagt-run`) -A thin wrapper invoked by the registry server around every tool execution. Records the command, arguments, exit code, duration, file counts, and truncated stderr to `/trace_archive/.json` and emits an OTLP span to MLflow. The agent calls `reconstruct_pipeline` to render the trace archive as a reproducible bash script or Snakemake workflow. +A wrapper around every registered-code execution. Records the command, arguments, exit code, duration, file counts, and truncated stderr to `/trace_archive/.json` and emits a `code.execute` span to the trace store. The server incrementally embeds those records into a `code_use` collection so past executions are retrievable, and the agent calls `reconstruct_pipeline` to render the trace archive as a reproducible workflow. -**Observability** (MLflow + OTLP) -MLflow runs locally at a port pinned at `dsagt init` time. All four layers emit OTLP HTTP spans to MLflow's `/v1/traces` endpoint. The agent's own LLM-call traces land in the same store when you export the `OTEL_EXPORTER_OTLP_ENDPOINT` printed by `dsagt init`. +**Observability** (serverless MLflow) +Traces land in a serverless MLflow store — a SQLite file at `/mlflow.db`. DSAgt emits its own spans live; the agent's LLM-call traces are recovered post-hoc from the on-disk session transcript by the MCP server's in-session reader. View with `mlflow ui --backend-store-uri sqlite:////mlflow.db`. + +**Skills Discovery** +DSAgt exposes MCP tools to connect to external GitHub skill repositories and search them for skills that enhance scientific workflows. On top of the agent's own progressive disclosure of the skills already installed in its native skill folder, DSAgt maintains an extendable knowledge base of skills that can be searched and installed on demand — without flooding the agent's context window with skills it isn't using. The agent searches via `search_skills` and installs via `install_skill`. + +**Memory** +DSAgt adds two memory extensions that complement the host agent's own memory (which distills session facts into a managed set of Markdown files loaded into context). *Explicit memory* records user-confirmed facts as YAML (`kb_remember` / `kb_get_memories`). *Episodic memory* (opt-in) keeps a vector store of semantically chunked turn blocks and searches them by successive filtering — first to a session, then by regex over the query's key themes, then a final vector ranking of what remains — so long, multi-session history can augment agent context without the noisy retrieval that undifferentiated episodic accrual would otherwise produce. ## Project Layout ``` ~/dsagt-projects// - dsagt_config.yaml # project configuration - tools/ # registered CLI tool specs (markdown + YAML frontmatter) - tools/code/ # agent-written tool scripts + .dsagt/ # dsagt-internal state (hidden) + config.yaml # project configuration (set by dsagt init) + state.yaml # session log + memory cursor (owned by the MCP server) + explicit_memories.yaml # user-confirmed facts + codes// # registered codes — skill-standard dirs (SKILL.md + scripts/) skills/ # agent skills (SKILL.md + reference docs) - trace_archive/ # tool execution records (JSON, from dsagt-run) - mlflow/ # MLflow traces, metrics, artifacts + trace_archive/ # code execution records (JSON, from dsagt-run) + mlflow.db # serverless MLflow SQLite trace store kb_index/ # knowledge base vector collections - explicit_memories.yaml # user-confirmed facts # Per-agent runtime config (one of, generated by dsagt init): # claude: CLAUDE.md, .mcp.json # goose: goose.yaml, .goosehints # codex: AGENTS.md, .codex-data/config.toml # opencode: AGENTS.md, opencode.json - # roo: .roomodes, .roo/mcp.json # cline: .clinerules/, cline_mcp_settings.json ``` -Projects are registered in `~/.dsagt/projects.yaml` so `dsagt mlflow ` and `dsagt info ` work from any directory. The data layer is agent-agnostic — re-running `dsagt init --agent ` switches agent platforms while preserving all accumulated knowledge and traces. +Projects are registered in `~/dsagt-projects/projects.yaml` so `dsagt info ` works from any directory. The project's data is agent-agnostic — re-running `dsagt init` for the same project and choosing a different agent switches platforms while preserving all accumulated knowledge and traces. diff --git a/docs/assets/architecture.png b/docs/assets/architecture.png index 7dbdf46..013ca20 100644 Binary files a/docs/assets/architecture.png and b/docs/assets/architecture.png differ diff --git a/docs/assets/skills-routing.png b/docs/assets/skills-routing.png new file mode 100644 index 0000000..15c486d Binary files /dev/null and b/docs/assets/skills-routing.png differ diff --git a/docs/cli.md b/docs/cli.md index e4eb4d6..15da548 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -1,49 +1,50 @@ # CLI Reference -All commands are available after installing DSAgt. +All commands are available after [installation](index.md#installation) and activating your virtual environment. ## Project Management | Command | Description | |---------|-------------| -| `dsagt init --agent [--location ] [--mlflow-port N]` | Create a project; write per-agent MCP config; print the launch one-liner | -| `dsagt list` | List all projects with agent, status, and path | -| `dsagt info [--json]` | Resolved config (with source per value) and a session/error summary | -| `dsagt mv ` | Move a project to a new location | -| `dsagt rm [-y] [--keep-files]` | Unregister a project and optionally delete its directory | +| `dsagt init` | Create or reconfigure a project — interactive menu for name, location, agent platform, knowledge collections, skill sources, and the episodic-memory opt-in. Sets up the knowledge base and writes the per-agent instructions + MCP config. Re-runnable as a settings editor. | +| `dsagt start ` | Launch the agent in the project directory (equivalent to `cd && `); on exit, runs post-session catch-up. | +| `dsagt list` | List all projects with agent and path. | +| `dsagt info [--json]` | Resolved config (with the source of each value) and a session/trace summary read from the SQLite store. | +| `dsagt mv ` | Move a project to a new location. | +| `dsagt rm [-y] [--keep-files]` | Unregister a project and optionally delete its directory. | +| `dsagt smoke-test [--agent claude\|goose\|codex\|opencode\|cline]` | End-to-end install verification. | -## Session Lifecycle +### Deprecated `dsagt init` flags (backcompat) -| Command | Description | -|---------|-------------| -| `dsagt mlflow ` | Start MLflow for a project and print OTel routing exports | -| `dsagt stop ` | Stop the MLflow daemon | -| `dsagt memory --project ` | Distill new traces from MLflow into episodic memory | +The pre-menu flags still work — the automation/CI path — but are **deprecated** in favor of the interactive menu. Each one skips its corresponding prompt: -## Setup +| Flag | Prompt it replaces | +|------|--------------------| +| `` (positional) | Project name | +| `--agent ` | Agent platform | +| `--location ` | Project location | +| `--include … \| --exclude …` | Knowledge collections / skill sources | +| `--episodic` | "Enable episodic memory?" | -| Command | Description | -|---------|-------------| -| `dsagt setup-kb [--collection ]` | Build the shared core knowledge base collections | -| `dsagt smoke-test [--agent claude\|goose\|codex\|opencode]` | End-to-end install verification | +New usage should prefer bare `dsagt init` and the menu. ## Project Location -The default project location is `~/dsagt-projects//`. Override with `--location`: +The default project location is `~/dsagt-projects//`. + +## Viewing traces + +The trace store is a serverless SQLite file — browse it with MLflow's UI pointed at the file: ```bash -dsagt init my-project --agent claude --location /data/runs # /data/runs/my-project/ -dsagt init my-project --agent claude --location . # ./my-project/ +mlflow ui --backend-store-uri sqlite:////mlflow.db ``` ## Server Commands -These are launched automatically by `dsagt init` via the per-agent MCP config and are not typically run directly. +These are launched automatically by the per-agent MCP config and are not typically run directly. | Command | Description | |---------|-------------| -| `dsagt-registry-server` | Tool registry MCP server | -| `dsagt-knowledge-server` | Knowledge base MCP server | -| `dsagt-run` | Provenance-capturing tool execution wrapper | -| `dsagt-proxy` | LiteLLM proxy server (proxy mode only) | -| `dsagt-setup-kb` | Core knowledge base setup (called by `dsagt setup-kb`) | +| `dsagt-server` | The single MCP server — code registry, knowledge base, memory, and skills. Also runs the in-session heartbeat (trace capture + code-use/episodic indexing). | +| `dsagt-run` | Provenance-capturing code execution wrapper; writes execution records to `/trace_archive/`. | diff --git a/docs/developer.md b/docs/developer.md index 9ffbb2e..cb7afbe 100644 --- a/docs/developer.md +++ b/docs/developer.md @@ -1,47 +1,24 @@ # Developer Guide -Material for contributors and users who are working beyond the default `dsagt init` → `dsagt mlflow` → agent flow. +Material for contributors and users working beyond the default `dsagt init` → agent flow. ## Tests ```bash uv run python -m pytest -m "not integration" # unit tests, no creds required -uv run python -m pytest -m integration -v # integration tests (require .env) +uv run python -m pytest -m integration -v # integration tests (require real credentials / models) ``` -Integration tests read endpoint and key values from `.env` at the repo root. Copy `.env.example` to `.env` and fill in your values. - -For per-flow hand-tests (CLI, proxy mode, VS Code extensions), see the scripts under [`tests/smoke_test/manual_runs/`](https://github.com/AI-ModCon/dsagt/tree/main/tests/smoke_test/manual_runs/). - -## Proxy Mode - -`dsagt init` followed by `dsagt start --enable-proxy` spawns a LiteLLM proxy in front of your agent's LLM calls. This adds: - -- Full LLM-call traces (request bodies, tool-use blocks, response payloads) in MLflow for agents whose native OTel does not emit those payloads (codex, opencode). -- Cache-breakpoint injection on outgoing requests (Anthropic prompt caching). -- Sidechannel detection for agent-internal title-generator / session-namer calls. -- Model-name aliasing — useful when an agent CLI hardcodes a model whitelist incompatible with your gateway's served names (cline, roo). - -Proxy mode reads upstream LLM credentials from `.env` or the shell. See [`tests/smoke_test/manual_runs/proxy_walkthrough.md`](https://github.com/AI-ModCon/dsagt/blob/main/tests/smoke_test/manual_runs/proxy_walkthrough.md) for the full setup walkthrough. +For per-flow hand-tests (CLI, VS Code extensions), see the scripts under [`tests/manual_walkthroughs/`](https://github.com/AI-ModCon/dsagt/tree/main/tests/manual_walkthroughs/). ## Troubleshooting **Agent command not found.** The agent CLI is not installed or is not on PATH. See the [supported agents table](index.md#supported-agents). -**MCP servers not connecting.** Verify the server commands are on your PATH: - -```bash -which dsagt-registry-server -which dsagt-knowledge-server -``` - -If missing, reinstall: `pip install --force-reinstall https://github.com/AI-ModCon/dsagt/archive/refs/tags/0.1.0.zip`. - -**MLflow UI empty.** Confirm MLflow is running for the right project: +**MCP server not connecting.** Verify the server command resolves: ```bash -dsagt info # shows the pinned port -curl http://localhost: +uv run which dsagt-server ``` -**Claude keychain conflict.** If `claude` will not authenticate against a non-default gateway, run `claude /logout` to clear the macOS Keychain OAuth token, then re-export `ANTHROPIC_BASE_URL` / `ANTHROPIC_API_KEY` and re-launch. +If missing, reinstall: `pip install --force-reinstall "git+https://github.com/AI-ModCon/dsagt.git"`. \ No newline at end of file diff --git a/docs/index.md b/docs/index.md index 1b7fec3..c2f8c4e 100644 --- a/docs/index.md +++ b/docs/index.md @@ -2,39 +2,54 @@ **D**ata**S**mith **Ag**en**t** — AI-assisted data pipeline builder. -DSAgt connects an MCP-compatible AI coding agent to tool registration, a semantic knowledge base, execution provenance, and observability infrastructure. It provides data-pipeline scaffolding around your existing agent CLI or VS Code extension (Claude Code, Goose, Codex, and others). +DSAgt connects an MCP-compatible AI coding agent to code registration, a semantic knowledge base, execution provenance, and observability infrastructure. It provides data-pipeline scaffolding around your existing agent CLI or VS Code extension (Claude Code, Goose, Codex, and others). ## Supported Agents -| Agent | Install | Verify | -|-------|---------|--------| -| [Claude Code](https://github.com/anthropics/claude-code) | `npm i -g @anthropic-ai/claude-code` | `claude --version` | -| [Goose](https://github.com/block/goose) | See [Goose docs](https://github.com/block/goose#installation) | `goose --version` | -| [Codex](https://github.com/openai/codex) | `npm i -g @openai/codex` | `codex --version` | -| [opencode](https://github.com/sst/opencode) | See [opencode docs](https://opencode.ai/docs/) | `opencode --version` | -| [Roo Code](https://github.com/RooCodeInc/Roo-Code) | `npm i -g @roo-code/cli` | `roo --version` | -| [Cline](https://github.com/cline/cline) | `npm i -g cline` | `cline --version` | + +{% + include-markdown "../README.md" + start="" + end="" +%} ## Prerequisites -- Python 3.12–3.13 +- Python 3.12 or 3.13 - One of the supported agent platforms above, installed and authenticated against your LLM provider +- [uv](https://github.com/astral-sh/uv) — only for the development install ## Installation +### For use (no development) + + +{% + include-markdown "../README.md" + start="" + end="" +%} + +### For development + +Clone the repo and use `uv` (editable install; add `--all-groups` for the test suite): + ```bash -pip install https://github.com/AI-ModCon/dsagt/archive/refs/tags/0.1.0.zip +git clone https://github.com/AI-ModCon/dsagt.git +cd dsagt && uv sync --all-groups +source .venv/bin/activate ``` ## Key Capabilities -| Layer | What it does | +| Capability | What it does | |-------|-------------| -| **Tool Registry** | Register CLI tools as markdown specs; the agent discovers and runs them via `search_registry` | -| **Knowledge Base** | Semantic search over indexed document collections (ChromaDB + FAISS) | -| **Provenance** | `dsagt-run` wrapper records every tool execution to `trace_archive/` and MLflow | +| **Code Registry** | Register CLI codes as markdown specs; the agent discovers and runs them via `search_registry` | +| **Knowledge Base** | Hybrid semantic + keyword (BM25) search over indexed ChromaDB collections | +| **Skills Discovery** | Search external GitHub skill catalogs and install workflow skills on demand via `search_skills` / `install_skill`, without flooding the agent's context | +| **Provenance** | `dsagt-run` wrapper records every code execution to `trace_archive/`; `reconstruct_pipeline` renders it as a runnable script | | **Explicit Memory** | User-confirmed facts persisted to YAML and the knowledge base | -| **Episodic Memory** | Session distillation via outlier detection over MLflow traces | -| **Observability** | Full OTLP tracing to a local MLflow instance | +| **Episodic Memory** | Opt-in: the MCP server mechanically chunks and embeds each session turn into a searchable `session_memory` collection (recency-weighted retrieval) | +| **Observability** | Serverless MLflow tracing (a per-project SQLite file) — DSAgt's own spans plus agent traces recovered from the on-disk transcript | See the [Quick Start](quickstart.md) to try all of these in a single session. diff --git a/docs/knowledge-base.md b/docs/knowledge-base.md index 3370677..f6ef094 100644 --- a/docs/knowledge-base.md +++ b/docs/knowledge-base.md @@ -1,40 +1,47 @@ # Knowledge Base -DSAgt maintains six independently-partitioned ChromaDB collections. The first three are global (under `~/.dsagt/kb_index/`, populated by `dsagt setup-kb`); the last three are per-project (under `/kb_index/`, populated automatically during use). +The knowledge base is DSAgt's catalog of **domain knowledge** — reference corpora and your own documents — that the agent searches to ground its work on scientific data-processing and AI-readiness evaluation. Rather than let the agent guess at a library's API or a standard's requirements, it retrieves the relevant passages first, then acts. -## Collections +## Domain-knowledge collections | Collection | Source | Populated by | |---|---|---| -| **Tool Specs** | Bundled CLI tool specs in `src/dsagt/tools/` | `dsagt setup-kb` | -| **Skills** | Bundled skill workflows in `src/dsagt/skills/` | `dsagt setup-kb` | -| **Domain Knowledge** | NeMo Curator + AIDRIN reference corpora; user-ingested docs | `dsagt setup-kb` + agent's `kb_ingest` | -| **Explicit Memory** | User-confirmed facts | Agent's `kb_remember` (also written to `/explicit_memories.yaml`) | -| **Episodic Memory** | Distilled facts from MLflow traces | `dsagt memory --project ` | -| **Tool Use Records** | `dsagt-run` execution traces | `dsagt-run` wrapper writes JSON to `/trace_archive/`; indexed by `dsagt memory` | +| **Reference corpora** | NeMo Curator + AIDRIN (data-curation and AI-data-readiness references) | `dsagt init` (chosen collections) | +| **Your documents** | Papers, standards, protocols, schemas you ingest | Agent's `kb_ingest` | -## Explicit Memory +The agent works these with three tools: `kb_ingest` (index a file or directory into a named collection — long ingests run in the background), `kb_search` (retrieve across one or more collections), and `kb_list_collections` (see what's indexed). -Explicit memories are facts the user confirms during a session. The agent saves them via `kb_remember`, which writes to both the ChromaDB collection and `/explicit_memories.yaml`. The agent fetches them via `kb_get_memories` on demand (typically when you ask it to recall something) — they are not auto-loaded at session start. +## Why hybrid vector search -## Episodic Memory +Retrieval is **hybrid** — dense semantic embeddings fused with sparse BM25 keyword matching — and it's on by default per collection: -`dsagt memory --project ` distills new traces from the project's MLflow store into episodic memory using per-category outlier detection over embedding centroids. Run this after each session to accumulate cross-session memory. +- **Semantic embeddings** catch paraphrase and synonymy: a query about "missing values" finds a passage on "null rates" even with no shared words. +- **BM25 keyword matching** catches the exact terms embeddings tend to under-rank — identifiers, gene names, parameter flags, standard names — where a literal match matters. +- **Per-collection partitioning** scopes a search to a domain, so a materials-science query isn't diluted by genomics references. +- **Optional cross-encoder reranking** re-scores the top candidates for precision when it's worth the extra pass. -## Search +The default embedder is a local sentence-transformers model (~130 MB, CPU-side, no API key). +## One store, many concerns -The agent searches all collections via `kb_search` (knowledge MCP server) and writes via `kb_ingest` / `kb_remember`. Tool Specs and Skills are queried through specialized routes (`search_registry`, `search_skills`) over the same backend. - -Hybrid search (dense embeddings + sparse BM25 via Reciprocal Rank Fusion) is on by default per collection route. Cross-encoder reranking is optional. +The same vector store backs more than domain knowledge. DSAgt's [memory](memory.md) (explicit + episodic), [skills discovery](skills.md) (the installable-skill catalog), and [code execution tracking](provenance.md) (the `code_use` records) each live in it as their own partitioned collections, sharing one embedder and one ChromaDB. This page covers domain-knowledge cataloging; those pages carry the depth on their respective collections. ## Setup +`dsagt init` sets up the knowledge base from your choices in the interactive menu. Re-running `dsagt init` on an existing project reconfigures it in place. + +## Try it + ```bash -dsagt setup-kb # all global collections (local embedder) -dsagt setup-kb --collection nemo_curator -dsagt setup-kb --embedding-backend api \ - --embedding-base-url \ - --embedding-api-key +dsagt init # name it `demo`, and enable the AIDRIN collection in the prompts +dsagt start demo ``` -The Tool Specs and Skills collections are wiped and rebuilt on every `setup-kb` run — re-run after upgrading DSAgt to pick up new bundled assets. +Then, in the agent — substituting `` with any folder of your +own documents (papers, protocols, schemas): + +1. > Ingest the docs in `` into a collection named `domain`. +2. > Search the `domain` and `aidrin` collections for how to assess data completeness. + +## In practice + +See the [Use Cases](use-cases/index.md), where domain references and ingested protocols guide the agent through curating a real scientific dataset. diff --git a/docs/mcp-servers.md b/docs/mcp-servers.md index ee999cc..7edd9e4 100644 --- a/docs/mcp-servers.md +++ b/docs/mcp-servers.md @@ -1,25 +1,23 @@ -# MCP Servers +# MCP Server -DSAgt exposes its capabilities through two MCP servers. Both are launched automatically by `dsagt init` and configured in the per-agent runtime file (`.mcp.json` for Claude Code, `goose.yaml` for Goose, etc.). +DSAgt exposes its capabilities through a single MCP server, **`dsagt-server`**, configured in the per-agent runtime file (`.mcp.json` for Claude Code, `goose.yaml` for Goose, etc.) and launched automatically when the agent starts. It bundles four capabilities — a code registry, a knowledge base, explicit memory, and skill discovery — behind one process with one shared embedder and one ChromaDB. -## Registry Server +> Earlier versions ran two separate servers (`dsagt-registry-server` + `dsagt-knowledge-server`), merged in 0.2.0. Re-run `dsagt start ` on an existing project to regenerate its config against the single server (for cline, delete `/.cline-data` first). -**Command:** `dsagt-registry-server` +## Registry tools -Handles tool registration, dependency installation, and tool discovery. +Code registration, dependency installation, and code discovery. | Tool | Description | |------|-------------| -| `search_registry` | Semantic search over registered tool specs | -| `save_tool_spec` | Register a new CLI tool as a markdown file with YAML frontmatter | -| `install_dependencies` | Install tool dependencies via `uv run --with` | +| `search_registry` | Semantic search over registered code specs | +| `save_code_spec` | Register a new CLI code as a markdown file with YAML frontmatter | +| `install_dependencies` | Install code dependencies via `uv run --with` | | `reconstruct_pipeline` | Render the trace archive as a bash script or Snakemake workflow | -Tools are markdown files with YAML frontmatter under `/tools/`. Executables are wrapped with `dsagt-run` for provenance and `uv run --with` for Python dependencies. +Codes are markdown files with YAML frontmatter under `/codes/`. Executables are wrapped with `dsagt-run` for provenance and `uv run --with` for Python dependencies. -## Knowledge Server - -**Command:** `dsagt-knowledge-server` +## Knowledge tools Semantic search and ingestion over indexed document collections. @@ -29,10 +27,10 @@ Semantic search and ingestion over indexed document collections. | `kb_ingest` | Index a file or directory into a named collection (runs in background for large corpora) | | `kb_remember` | Save a user-confirmed fact to explicit memory | | `kb_get_memories` | Retrieve explicit memories for the current project | -| `search_skills` | Discover agent skill workflows | +| `search_skills` | Discover installable skills in the external catalog (installed skills are auto-discovered natively) | ### Backend -The default embedding backend is local (`sentence-transformers`, CPU-only, no API key needed). Switch to `embedding.backend: api` in `dsagt_config.yaml` to route through a hosted embedder via LiteLLM. Cross-encoder reranking is available via `knowledge.rerank: true`. +The default embedding backend is local (`sentence-transformers`, CPU-only, no API key needed). Switch to `embedding.backend: api` in `.dsagt/config.yaml` to route through a hosted, OpenAI-compatible `/v1/embeddings` endpoint (set `embedding.base_url` and export `EMBEDDING_API_KEY`). Cross-encoder reranking is available via `knowledge.rerank: true`. -Hybrid search (dense + sparse BM25) is on by default and controlled per-route via the `hybrid` flag. +Hybrid search (dense + sparse BM25, fused by Reciprocal Rank Fusion) is always on per collection — it is a property of the store, not a per-call flag. diff --git a/docs/memory.md b/docs/memory.md new file mode 100644 index 0000000..de68ddf --- /dev/null +++ b/docs/memory.md @@ -0,0 +1,39 @@ +# Memory + +DSAgt gives the agent two kinds of persistent memory backed by the project's vector store: **explicit memory** — facts the user confirms — and opt-in **episodic memory** — an automatic record of session turns. Both are retrievable by the agent with `kb_search` / `kb_get_memories` MCP tools. + +## Explicit memory + +Explicit memories are facts the user confirms during a session. The agent saves them via `kb_remember`, which writes to both the ChromaDB collection and `/.dsagt/explicit_memories.yaml`. It fetches them via `kb_get_memories` on demand — typically when you ask it to recall something — so they are not auto-loaded at session start. + +Explicit memory is always on. It is the reliable path: when you say "remember that…", the agent records the fact so a future session can retrieve it. + +## Episodic memory + +Episodic memory is **opt-in**. When enabled, DSAgt reads the agent's transcript as the session runs and captures each completed turn into the `session_memory` collection — a fast, local chunk-and-embed pass that reuses the same embedder as the rest of the knowledge base. + +Retrieval over `session_memory` filters first to a session, then by regex over the query's key terms, before a final **recency-weighted** semantic ranking: a newer turn edges out a stale one as a bounded boost, so a corrected fact wins by recency while a strongly-relevant old turn is never buried. + +## Try it + +Explicit memory needs no setup: + +```bash +dsagt init # name it `demo`; answer "yes" to "Enable episodic memory?" to also capture turns +dsagt start demo +``` + +Then, in the agent: + +1. > Remember that `samples.csv` has null values in the status and timestamp columns. +2. > *(later, or in a new session)* What do you remember about the samples dataset? + +Confirm it persisted to disk: + +```bash +cat ~/dsagt-projects/demo/.dsagt/explicit_memories.yaml +``` + +## In practice + +See the [Use Cases](use-cases/index.md), where confirmed facts about a dataset — its quirks, thresholds, and decisions — carry forward across a multi-step curation session. diff --git a/docs/observability.md b/docs/observability.md index 3344a9a..613a74b 100644 --- a/docs/observability.md +++ b/docs/observability.md @@ -1,45 +1,51 @@ # Observability -DSAgt provides end-to-end trace visibility through a local MLflow instance. All internal layers emit OTLP HTTP spans to MLflow's `/v1/traces` endpoint. - -## Starting MLflow +DSAgt logs traces to a **serverless MLflow store** — a SQLite file at `/mlflow.db`, created lazily on the first span. View it with MLflow's UI pointed at the file: ```bash -dsagt mlflow +mlflow ui --backend-store-uri sqlite:////mlflow.db ``` -Prints the MLflow UI URL and the `export` block for routing agent OTel output. The port is pinned at `dsagt init` time and listed by `dsagt info `. +`dsagt info ` prints the resolved tracking URI and a session/trace summary. The tracking URI resolves as `MLFLOW_TRACKING_URI` env → project config → the `sqlite:////mlflow.db` default. + +## Two feeds + +DSAgt is **BYOA**: the agent talks to its own LLM provider directly, and DSAgt reconstructs traces from what the agent writes to disk. Traces come from two places: + +1. **First-party spans (live).** DSAgt instruments its own code and emits spans directly to the store as it runs. +2. **Agent traces (post-hoc).** The MCP server's in-session heartbeat reads the agent's own on-disk session transcript, translates it to a canonical trace shape, and writes it to the same store via the MLflow sink — recovering prompts, responses, and tool calls. ## Trace Coverage | Source | Span type | Contents | |--------|-----------|----------| | Knowledge base | `kb.search`, `kb.embed`, `kb.index_search`, `kb.rerank` | Per-phase timing trees | -| Tool executions | `tool.execute` | Exit code, duration, file counts, truncated stderr. Full payload in `trace_archive/.json` | -| Registry events | `save_tool_spec`, `install_dependencies`, `reconstruct_pipeline` | Span metadata | -| Native agent OTel | LLM call spans | Coverage varies by agent (see below) | +| Code executions | `code.execute` | Exit code, duration, file counts, truncated stderr. Full payload in `trace_archive/.json` | +| Registry events | `save_code_spec`, `install_dependencies`, `reconstruct_pipeline` | Span metadata | +| Agent traces | one AGENT subtree per turn (`llm` / `tool_` children) | Prompts, responses, tool calls, and token usage where the transcript carries them | -### Agent OTel Coverage +### Agent trace coverage -Export the variables printed by `dsagt mlflow` before launching your agent: +Agent traces are reconstructed from each agent's on-disk session record. A per-agent reader + translator runs for every supported agent (claude, codex, goose, opencode, cline), uniformly. Fidelity is capped by what the transcript persisted (e.g. token counts and timing appear where the agent recorded them). -| Agent | Coverage | -|-------|----------| -| claude | Full request/response payloads | -| goose | Full request/response payloads | -| codex | Token counts and tool names | -| opencode | None natively | +Every span carries the project's session id (minted per launch into `/.dsagt/state.yaml`) for filtering in the MLflow trace view. -Every span carries the project's `session.id` for filtering in the MLflow trace view. +## The heartbeat -## Provenance and Reconstruction +The trace scan runs as a periodic heartbeat inside the long-lived MCP server — the one DSAgt process alive in every launch flow. Each tick reads new transcript records, translates completed turns, and fans out to subscribers (the MLflow sink always; the episodic-memory extractor when enabled). Correctness rests on idempotency: each subscriber keeps its own ack set, so a re-tick or a next-session catch-up can never double-log or lose a turn. The same heartbeat incrementally indexes `dsagt-run` code-execution records into the `code_use` collection. -Tool execution records on disk (`trace_archive/.json`) provide the canonical provenance chain. The agent calls `reconstruct_pipeline` to render the archive as a reproducible bash script or Snakemake workflow. +The same heartbeat also indexes `dsagt-run` execution records for search — the on-disk records and pipeline reconstruction are covered under [Provenance](provenance.md). -## Stopping MLflow +## Try it ```bash -dsagt stop +dsagt init # follow the prompts: name it `demo`, then pick your agent +dsagt start demo # …run a prompt or two, then exit the agent +mlflow ui --backend-store-uri sqlite:///$HOME/dsagt-projects/demo/mlflow.db ``` -Releases the port and stops the gunicorn workers. The PID is stored in `/.runtime`. +Open the MLflow UI to see both feeds in one store: DSAgt's own `kb.*` / `code.execute` spans and the per-turn agent traces recovered from the transcript. `dsagt info demo` prints the same session/trace summary from the command line. + +## In practice + +See the [Use Cases](use-cases/index.md) to watch a full curation session — every retrieval, code run, and agent turn — land in the trace store as it happens. diff --git a/docs/provenance.md b/docs/provenance.md new file mode 100644 index 0000000..52bbf54 --- /dev/null +++ b/docs/provenance.md @@ -0,0 +1,69 @@ +# Provenance + +DSAgt makes every data operation a reproducible, auditable step. The agent registers a **code** — a CLI executable — and every run of that code is wrapped for provenance capture, so the whole pipeline can later be reconstructed from the record. + +## Codes + +Codes are CLI executables defined as markdown files with YAML frontmatter under `/codes/`. The agent registers new codes via the MCP server's `save_code_spec` tool and finds existing ones via `search_registry`. + +A code spec includes: + +- A YAML frontmatter block describing the executable, parameters, dependencies, and tags. +- A markdown body with the exact runnable command, usage, and notes for the agent. + +Example code spec (`codes/csv-summary/SKILL.md`): + +```markdown +--- +name: csv-summary +description: Summarize a CSV — columns, row count, null counts, numeric stats. Use when profiling a tabular dataset. +executable: dsagt-run --code csv-summary -- python codes/csv-summary/scripts/csv_summary.py +parameters: + file: + type: string + required: true + cli: positional + description: Path to the CSV file +dependencies: [] +tags: [csv, profiling] +--- + +Run this registered code with the exact shell command below… +``` + +DSAgt wraps every registered code with `dsagt-run` for provenance capture and `uv run --with` for Python dependencies, so the agent can call any code without managing environments manually. It ships one bundled code, `scan-directory`, indexed for search by `dsagt init`. + +## Execution capture + +Every registered code runs through the `dsagt-run` wrapper. For each call it records the command, arguments, exit code, duration, input/output file counts, and truncated stderr to `/trace_archive/.json`, and emits a `code.execute` span to the trace store. The MCP server incrementally indexes those records into the `code_use` collection, so past executions are searchable. + +The wrapper is the whole point of code-mediated data access: a direct shell or editor call leaves no record and breaks reconstruction. + +## Pipeline reconstruction + +The on-disk execution records are the canonical provenance chain. The agent calls `reconstruct_pipeline` to render the trace archive as a reproducible **bash script** (`format="bash"`) or **Snakemake workflow** (`format="snakemake"`). It flushes the latest records into the searchable index first, then walks the dependency graph inferred from each step's input/output files to order the steps. + +## Try it + +```bash +export SMOKE_DIR="$(pwd)/tests/smoke_test" # a small bundled sample script + CSV +dsagt init # follow the prompts: name it `demo`, then pick your agent +dsagt start demo # launch the agent in the project +``` + +Then, in the agent (replace `$SMOKE_DIR` with the absolute path you exported): + +1. > Register the CLI utility at `$SMOKE_DIR/csv_summary.py` as a code named `csv-summary`. +2. > Run the `csv-summary` code from the registry on `$SMOKE_DIR/data/samples.csv` and summarize the columns. +3. > Reconstruct the pipeline as a bash script. + +Afterwards, inspect the trail: + +```bash +ls ~/dsagt-projects/demo/{codes,trace_archive} # the specs + execution records +mlflow ui --backend-store-uri sqlite:///$HOME/dsagt-projects/demo/mlflow.db # code.execute spans +``` + +## In practice + +See the [Use Cases](use-cases/index.md) for provenance-captured pipelines on real datasets — for example registering `fastp` and `megahit` as codes and reconstructing a genomics QC-and-assembly pipeline end to end. diff --git a/docs/quickstart.md b/docs/quickstart.md index 32d6702..f3020d3 100644 --- a/docs/quickstart.md +++ b/docs/quickstart.md @@ -1,24 +1,25 @@ # Quick Start -This guide walks through knowledge ingest, tool registration, provenance, and explicit memory using the mock project in [`tests/smoke_test/`](https://github.com/AI-ModCon/dsagt/tree/main/tests/smoke_test/). The examples use `claude`; substitute another agent (`goose`, `codex`, `opencode`) if you prefer — the prompts are agent-agnostic. +This guide walks through knowledge ingest, code registration, provenance, and explicit memory using the mock project in [`tests/smoke_test/`](https://github.com/AI-ModCon/dsagt/tree/main/tests/smoke_test/). The examples use `claude`; substitute another agent (`goose`, `codex`, `opencode`, `cline`) if you prefer — the prompts are agent-agnostic. + +DSAgt is **BYOA**: your agent talks to its own LLM provider directly, and the trace store is a serverless SQLite file per project. ## Setup ```bash -# Install -pip install https://github.com/AI-ModCon/dsagt/archive/refs/tags/0.1.0.zip +# Install (any Python 3.12/3.13 environment) +pip install "git+https://github.com/AI-ModCon/dsagt.git" # Set a convenience variable for the smoke test directory (not a normal dsagt step) export SMOKE_DIR="$(pwd)/tests/smoke_test" -# 1. Create a new project called quickstart -dsagt init quickstart --agent claude - -# 2. Start MLflow in the background and print the OTel routing exports -dsagt mlflow quickstart +# 1. Create a project. `dsagt init` is interactive — follow the menu to name it +# `quickstart`, pick your agent, and choose knowledge collections + skill sources. +# It sets up the knowledge base on first run (a ~130 MB local embedder downloads once). +dsagt init -# 3. Paste the export block from step 2 into this shell, then launch the agent -cd ~/dsagt-projects/quickstart && claude +# 2. Launch the agent in the project: +dsagt start quickstart # …or: cd ~/dsagt-projects/quickstart && ``` ## Agent Prompts @@ -26,43 +27,35 @@ cd ~/dsagt-projects/quickstart && claude Inside the agent, paste these prompts one at a time. Replace `$SMOKE_DIR` with the absolute path you exported — the chat does not expand shell variables. 1. > Ingest the docs in `$SMOKE_DIR/knowledge/` into a collection named `knowledge`. -2. > Register the csvkit CLI tools `csvcut`, `csvgrep`, `csvstat`, and `csvlook`. -3. > Use the `scan_directory` tool from the registry to scan `$SMOKE_DIR/data/`. -4. > Summarize `samples.csv` — columns, row count, quality issues using csvkit tools from the registry. +2. > Register the CLI utility at `$SMOKE_DIR/csv_summary.py` as a code named `csv-summary` so we can reuse it. +3. > Use the `scan-directory` code from the registry to scan `$SMOKE_DIR/data/`. +4. > Run the `csv-summary` code on `$SMOKE_DIR/data/samples.csv` and tell me the columns, row count, and any columns with null values. 5. > Put this in explicit memory: samples.csv has null values in the status and timestamp columns. 6. > Tell me what you remember about the samples dataset. -## Teardown - -After exiting the agent, distill the session into episodic memory and stop the MLflow daemon: - -```bash -# Distill traces into episodic memory -dsagt memory --project quickstart - -# Stop the MLflow daemon -dsagt stop quickstart -``` +`csv_summary.py` is stdlib-only, so there's no dependency to install — registration and execution work out of the box. Step 4's null-column finding is the fact you store and recall in 5–6. -## What Was Exercised +## Capabilities Covered -| Prompt | DSAgt layer | +| Prompt | DSAgt capability | |--------|-------------| -| 1 | Knowledge MCP server (`kb_ingest`) — chunks and indexes docs into ChromaDB | -| 2 | Registry MCP server (`save_tool_spec`) — writes `tools/csvcut.md`, etc. | -| 3 | `dsagt-run` provenance wrapper — records exec layer to `trace_archive/` | -| 4 | KB recall via `kb_search` and registered tool execution | -| 5–6 | Explicit memory (`kb_remember` → `explicit_memories.yaml`) + `kb_get_memories` | +| 1 | `dsagt-server` (`kb_ingest`) — chunks and indexes docs into ChromaDB | +| 2 | `dsagt-server` (`save_code_spec`) — writes `codes/csv-summary/SKILL.md` (a skill-standard dir), wrapping the executable with `dsagt-run` | +| 3–4 | `dsagt-run` provenance wrapper — records each registered-code execution to `trace_archive/` | +| 5–6 | Explicit memory (`kb_remember` → `.dsagt/explicit_memories.yaml`) + `kb_get_memories` | ## Verify the Artifacts +Exit the agent (`Ctrl+C` or `/exit`), then: + ```bash -dsagt info quickstart -ls ~/dsagt-projects/quickstart/{tools,trace_archive} -cat ~/dsagt-projects/quickstart/explicit_memories.yaml -``` +dsagt info quickstart # config + a session/trace summary +ls ~/dsagt-projects/quickstart/{codes,trace_archive} +cat ~/dsagt-projects/quickstart/.dsagt/explicit_memories.yaml -The MLflow UI URL is printed by `dsagt mlflow quickstart`. +# Traces land in a serverless SQLite store. Browse them with: +mlflow ui --backend-store-uri sqlite:///$HOME/dsagt-projects/quickstart/mlflow.db +``` ## Non-Interactive Smoke Test @@ -72,22 +65,16 @@ The same flow runs non-interactively and asserts each artifact is present: dsagt smoke-test --agent claude ``` -## First-Time Knowledge Base Setup +## Knowledge Base Setup -`dsagt setup-kb` builds shared ChromaDB collections under `~/.dsagt/kb_index/` that every project on this machine reuses. Run this once after installation. - -```bash -dsagt setup-kb # all collections (local embedder, no creds) -dsagt setup-kb --collection nemo_curator -dsagt setup-kb --embedding-backend api --embedding-base-url ... --embedding-api-key ... -``` +`dsagt init` sets up the project's knowledge base with three kinds of collection: -Three collections are populated: +- **Code Specs** — DSAgt's bundled code specs, always set up so the agent finds them via `search_registry` from the first session. +- **Skill Catalogs** — the skill-catalog sources you chose at init (default `genesis`), cloned and indexed so `search_skills` returns installable skills. +- **Knowledge Collections** — optional reference document sets you chose at init (`nemo_curator`, `aidrin`). -- **Tool Specs** — DSAgt's bundled tool specs from `src/dsagt/tools/`, tagged `source: bundled`. -- **Skills** — DSAgt's bundled skill workflows from `src/dsagt/skills/`. -- **Domain Knowledge** — NeMo Curator and AI Data Readiness Inspector reference corpora. +`--include` / `--exclude` (asset names, or `all`) select the set non-interactively. The default embedder is a local sentence-transformers model (~130 MB, CPU-side, no API key). -The Tool Specs and Skills collections are wiped and rebuilt on every run, so re-run `setup-kb` after upgrading DSAgt. +## Optional: Episodic Memory -The default embedder is a local sentence-transformers model (~130 MB, CPU-only, no API key). Pass `--embedding-backend api` to route through a hosted embedder via LiteLLM. +Answer **yes** to "Enable episodic memory?" in the `dsagt init` prompts to have the MCP server capture each session turn into a searchable `session_memory` collection. Capture is mechanical (chunk + embed) and reuses the local embedder, so there's nothing extra to download. See [Memory → Episodic Memory](memory.md#episodic-memory). diff --git a/docs/skills.md b/docs/skills.md new file mode 100644 index 0000000..23a6cd7 --- /dev/null +++ b/docs/skills.md @@ -0,0 +1,46 @@ +# Skills + +Skills are instruction-based agent workflows — a directory with a `SKILL.md` and optional reference docs that the agent reads and follows. DSAgt lets the agent discover and install skills from external catalogs on demand, without loading thousands of them into context. + +Skills live in `/skills/`. Each is a directory containing a `SKILL.md` file and optional reference documents. + +## Two tiers + +![DSAgt skill routing](assets/skills-routing.png) + +Skills live in **two tiers**, and a single MCP service — the **SkillRouter** — is the one entry point that routes every skill operation between them: + +- **Catalog tier** — skills that exist in external repositories but are *not yet installed*. DSAgt federates many sources (`k-dense-ai`, `anthropic`, `antigravity`, `composio`, `genesis`, or any git URL); each is cloned and indexed into its own collection. The agent browses this tier with `search_skills` and manages sources with `add_skill_source` / `list_skill_sources`. +- **Installed + created tier** — skills drawn into the project's Skill Directory (`/skills/`), either installed from the catalog (`install_skill`) or authored in place (e.g. with the bundled `skill-creator`). These are mirrored into each agent's *native* skill directory (`.claude/`, `.agents/`, `.cline/`) at install time (and re-mirrored at `dsagt init`/`start`), where the agent auto-discovers and auto-invokes them. + +The lifecycle runs **Discovery** (the router) → **Registration** (the searchable catalog) → **Progressive Exposure** (the native Skill Directory the agent loads on its own). + +## Design motivation + +- **Search is catalog-only.** Every supported agent (Claude, Codex, Goose, Cline, opencode) natively auto-discovers `SKILL.md` folders, so installed skills never need to be indexed or returned by a tool — the harness already loads them. So `search_skills` handles what native discovery can't reach: a catalog of potentially thousands of *un*installed skills, searchable without holding them all in context. Catalogs are indexed on name, description, and tags, which keeps those summaries compact and avoids diluting the embedding with full SKILL.md bodies. +- **Keyword fallback, no embedder required.** When no embedding model is configured, `search_skills` falls back to a keyword match over the local clones, so it still works (just less fuzzy) — no model or API key needed. +- **One router, not scattered policy.** Backend selection (semantic vs. keyword), the catalog/installed split, and source bookkeeping all live in the SkillRouter rather than being re-implemented at each MCP and CLI call site, so the behavior can't drift between them. +- **Federated and provenance-preserving.** Each source is an independent per-source collection, so re-syncing one never disturbs another; installing a catalog skill preserves its upstream `LICENSE`/`NOTICE` and stamps a `PROVENANCE.txt` into the installed directory. + +## Bundled and authored skills + +DSAgt ships a `skill-creator` skill (for scaffolding new `SKILL.md` skills). Bundled and installed skills are **not** indexed for search — the agent auto-discovers `SKILL.md` folders natively — so `search_skills` is reserved for the catalog tier. Domain skills, including the MODCON `datacard-generator`, are sourced from external catalogs rather than bundled, so they stay current upstream. + +To add one by hand, place a new directory under `/skills/` with a `SKILL.md` describing the workflow; the next `dsagt start` mirrors it into the agent's native skill directory, after which the agent auto-discovers and invokes it — no indexing step. + +## Try it + +```bash +dsagt init # follow the prompts: name it `demo`, then pick your agent +dsagt start demo # launch the agent in the project +``` + +Then, in the agent: + +1. > List the skill sources and their sync status. +2. > Sync the `genesis` catalog and search it for a data-card skill. +3. > Install the one that fits, then use it on this project. + +## In practice + +See the [Use Cases](use-cases/index.md), which draw on installed skills — such as the MODCON data-card generator — while working a real dataset end to end. diff --git a/docs/tools-skills.md b/docs/tools-skills.md deleted file mode 100644 index 6c057bd..0000000 --- a/docs/tools-skills.md +++ /dev/null @@ -1,43 +0,0 @@ -# Tools and Skills - -## Tools - -Tools are CLI executables defined as markdown files with YAML frontmatter under `/tools/`. The agent registers new tools via the registry MCP server's `save_tool_spec` tool. - -A tool spec includes: - -- A YAML frontmatter block describing the command, arguments, dependencies, and tags. -- A markdown body with usage examples and notes for the agent. - -Example tool spec structure: - -```markdown ---- -name: csvstat -command: csvstat -dependencies: [] -tags: [csv, statistics] ---- - -Prints descriptive statistics for all columns in a CSV file. - -Usage: csvstat [options] [FILE] -``` - -The registry server wraps every registered tool with `dsagt-run` for provenance capture and `uv run --with` for Python dependencies, so the agent can call any tool without managing environments manually. - -### Bundled Tools - -DSAgt ships a `scan_directory` tool in `src/dsagt/tools/` that is indexed into the global Tool Specs collection by `dsagt setup-kb`. - -## Skills - -Skills are instruction-based agent workflows in `/skills/`. Each skill is a directory containing a `SKILL.md` file and optional reference documents. The agent discovers skills via `search_skills`. - -### Bundled Skills - -DSAgt ships a `datacard-generator` skill in `src/dsagt/skills/` with reference templates for generating dataset documentation. It is indexed into the global Skills collection by `dsagt setup-kb`. - -### Adding Skills - -Place a new directory under `/skills/` with a `SKILL.md` describing the workflow. The knowledge server indexes it automatically on next startup, or trigger a re-index via `kb_ingest`. diff --git a/docs/use-cases/aidrin-cryoem.md b/docs/use-cases/aidrin-cryoem.md index 6853cdd..1114366 100644 --- a/docs/use-cases/aidrin-cryoem.md +++ b/docs/use-cases/aidrin-cryoem.md @@ -2,7 +2,7 @@ **Domain:** AI data readiness — quality assessment of a scientific pipeline -**Tools:** [`aidrin`](https://github.com/idtlab/AIDRIN) (AI Data Readiness Inspector CLI) +**Codes:** [`aidrin`](https://github.com/idtlab/AIDRIN) (AI Data Readiness Inspector CLI) **Dataset:** EMPIAR-10017 cryo-EM particle tables (via CryoPPP) diff --git a/docs/use-cases/aidrin-tour.md b/docs/use-cases/aidrin-tour.md index 75263cb..b5810ef 100644 --- a/docs/use-cases/aidrin-tour.md +++ b/docs/use-cases/aidrin-tour.md @@ -2,7 +2,7 @@ **Domain:** AI data readiness — quality, fairness, and privacy assessment -**Tools:** [`aidrin`](https://github.com/idtlab/AIDRIN) (AI Data Readiness Inspector CLI) +**Codes:** [`aidrin`](https://github.com/idtlab/AIDRIN) (AI Data Readiness Inspector CLI) **Dataset:** UCI Adult census (bundled with AIDRIN) diff --git a/docs/use-cases/cryoem.md b/docs/use-cases/cryoem.md index f07226e..8c97d2e 100644 --- a/docs/use-cases/cryoem.md +++ b/docs/use-cases/cryoem.md @@ -8,7 +8,7 @@ ## Overview -This use case demonstrates DSAgt-assisted curation of cryo-EM data from the EMPIAR public archive. The agent registers curation tools, ingests domain knowledge about cryo-EM data quality, and builds a pipeline for micrograph preprocessing. +This use case demonstrates DSAgt-assisted curation of cryo-EM data from the EMPIAR public archive. The agent registers curation codes, ingests domain knowledge about cryo-EM data quality, and builds a pipeline for micrograph preprocessing. ## Guides diff --git a/docs/use-cases/index.md b/docs/use-cases/index.md index e1d73a6..221b44c 100644 --- a/docs/use-cases/index.md +++ b/docs/use-cases/index.md @@ -1,6 +1,6 @@ # Use Cases -End-to-end walkthroughs for representative scientific and data-readiness scenarios live in [`use_cases/`](https://github.com/AI-ModCon/dsagt/tree/main/use_cases/). Each covers data acquisition, tool registration, pipeline construction, and agent-driven execution against a real dataset. +End-to-end walkthroughs for representative scientific and data-readiness scenarios live in [`use_cases/`](https://github.com/AI-ModCon/dsagt/tree/main/use_cases/). Each covers data acquisition, code registration, pipeline construction, and agent-driven execution against a real dataset. | Use case | Domain | Guide | |----------|--------|-------| diff --git a/docs/use-cases/microbial-isolates.md b/docs/use-cases/microbial-isolates.md index b6757dc..78effa5 100644 --- a/docs/use-cases/microbial-isolates.md +++ b/docs/use-cases/microbial-isolates.md @@ -2,7 +2,7 @@ **Domain:** Genomics — short-read QC and assembly -**Tools:** `fastp`, `megahit` +**Codes:** `fastp`, `megahit` **Source:** [`use_cases/microbial_isolates/`](https://github.com/AI-ModCon/dsagt/tree/main/use_cases/microbial_isolates/) diff --git a/docs/use-cases/vasp.md b/docs/use-cases/vasp.md index 9f4ce28..06dd4d5 100644 --- a/docs/use-cases/vasp.md +++ b/docs/use-cases/vasp.md @@ -2,13 +2,13 @@ **Domain:** Materials science — density functional theory -**Tools:** VASP, ISAAC +**Codes:** VASP, ISAAC **Source:** [`use_cases/isaac_vasp/`](https://github.com/AI-ModCon/dsagt/tree/main/use_cases/isaac_vasp/) ## Overview -This use case covers DFT input/output handling with VASP using DSAgt and the ISAAC workflow system. The agent registers VASP pre/post-processing tools and transfers NEB calculation results into the ISAAC database. +This use case covers DFT input/output handling with VASP using DSAgt and the ISAAC workflow system. The agent registers a VASP-to-ISAAC converter as a code and turns NEB calculation results into an ISAAC AI-Ready Record. ## Guides diff --git a/latex/architecture.png b/latex/architecture.png index 7dbdf46..013ca20 100644 Binary files a/latex/architecture.png and b/latex/architecture.png differ diff --git a/latex/dsagt.tex b/latex/dsagt.tex index c4c7a05..68cc740 100644 --- a/latex/dsagt.tex +++ b/latex/dsagt.tex @@ -343,19 +343,29 @@ \section{Architecture} % Knowledge centered above (Explicit, Domain); Registry above (Skills, % Tool Specs); MLflow directly above Episodic; dsagt-run directly above % Tool Records. All x-coordinates fall out of the row-3 grid below. -\node[svc] (knowledge) at (-0.9,-2) {Knowledge\\Server}; -\node[svc] (registry) at (3.5,-2) {Registry\\Server}; +\node[svc] (knowledge) at (-0.9,-2) {Knowledge}; +\node[svc] (registry) at (3.5,-2) {Registry}; \node[art] (mlflow) at (6.8,-2) {MLflow\\Traces}; % dsagt-run width matches the artifact boxes (18mm) so it visually % pairs with MLflow Traces and Tool Records on either side. \node[svc, minimum width=18mm] (dsagtrun) at (9.0,-2) {\texttt{dsagt-run}}; +% The Knowledge and Registry tool surfaces are now exposed by ONE merged +% MCP server (dsagt-server). A single container box spans + sits behind +% both (background layer, lighter fill) so the two inner boxes still read +% as distinct concern modules while the bridge conveys "one server". +\begin{scope}[on background layer] + \node[draw, rounded corners=3pt, thick, fill=green!6, + fit=(knowledge)(registry), inner sep=1.2mm] (mcpbox) {}; +\end{scope} + % Agent <-> MCP servers (symmetric forks; bidirectional because MCP is % request/response — the server replies travel back to the agent). \draw[bilink] (agent.south) -- ++(0,-0.3) -| (knowledge.north); \draw[bilink] (agent.south) -- ++(0,-0.3) -| (registry.north); -% MCP label centered below the agent, between the two MCP-server forks. -\node[lbl] at (1.3,-1.2) {MCP}; +% MCP label in the center of the merged-server box, in the gap between the +% Knowledge and Registry modules — the bridge that joins them into one server. +\node[lbl] at (1.3,-2) {MCP}; % Agent -> MLflow + dsagt-run (fork off the same trunk as the MCP servers, % drop vertically directly above each box). @@ -383,7 +393,7 @@ \section{Architecture} % on every registered-tool invocation. \node[art] (explicit) at (-2.0,-4) {Explicit\\Memory}; \node[art] (domain) at (0.2,-4) {Domain\\Knowledge}; -\node[art] (skills) at (2.4,-4) {Skills}; +\node[art] (skills) at (2.4,-4) {Skills\\Catalog}; \node[art] (tools) at (4.6,-4) {Tool\\Specs}; \node[art] (episodic) at (6.8,-4) {Episodic\\Memory}; \node[art] (history) at (9.0,-4) {Tool\\Records}; diff --git a/latex/skills-routing.png b/latex/skills-routing.png new file mode 100644 index 0000000..15c486d Binary files /dev/null and b/latex/skills-routing.png differ diff --git a/latex/skills-routing.tex b/latex/skills-routing.tex new file mode 100644 index 0000000..f4a9d89 --- /dev/null +++ b/latex/skills-routing.tex @@ -0,0 +1,89 @@ +% Skills routing diagram — companion to Figure 1 (architecture) in dsagt.tex. +% Standalone so it exports straight to latex/skills-routing.png: +% pdflatex skills-routing.tex +% pdftoppm -png -r 175 -singlefile skills-routing.pdf skills-routing +% +% Colour semantics (see legend): blue = external (skill repos, agent), +% green = DSAgt service (SkillRouter), orange = DSAgt store / skill (Skills +% Catalog, Skill Directory, the bundled skill-creator). Concern bands are +% uniform faint outlines. SkillRouter is the single skill-MCP entry point: +% every skill tool is labelled on the router operation it drives. +\documentclass[border=2mm]{standalone} +\usepackage{tikz} +\usepackage{lmodern} +\usetikzlibrary{positioning, arrows.meta, fit, backgrounds, calc} + +\begin{document} +\begin{tikzpicture}[ + >={Stealth[length=2mm]}, + box/.style={ + draw, rounded corners=2pt, thick, + minimum height=12mm, align=center, + font=\footnotesize, fill=white}, + ext/.style={box, fill=blue!12}, % external (repos, agent) + svc/.style={box, fill=green!12}, % DSAgt service + store/.style={box, fill=orange!14}, % DSAgt store / skill + skill/.style={draw, rounded corners=2pt, thick, fill=white, + font=\scriptsize, minimum height=5mm, align=center}, + link/.style={->, thick, blue!60!black}, + bilink/.style={<->, thick, blue!60!black}, + lbl/.style={font=\scriptsize\itshape, text=gray!55!black, align=center}, + concern/.style={font=\footnotesize\bfseries, text=gray!62!black}, + band/.style={draw=gray!55, dashed, rounded corners, fill=black!3}, +] + +% ============================ legend ==================================== +\node[ext, minimum width=5mm, minimum height=4mm] (lk1) at (-1.6,7.4) {}; +\node[right=1mm of lk1, font=\scriptsize, anchor=west] (lt1) {External}; +\node[svc, minimum width=5mm, minimum height=4mm, right=14mm of lk1] (lk2) {}; +\node[right=1mm of lk2, font=\scriptsize, anchor=west] {DSAgt service}; +\node[store, minimum width=5mm, minimum height=4mm, right=24mm of lk2] (lk3) {}; +\node[right=1mm of lk3, font=\scriptsize, anchor=west] {DSAgt store / skill}; + +% ============================ external (outside bands) ================== +\node[ext, minimum width=32mm] (extrepos) at (-0.2,0.6) + {External Skills Repos\\{\scriptsize scientific $\cdot$ anthropic $\cdot$ antigravity}\\{\scriptsize composio $\cdot$ genesis $\cdot$ any git URL}}; + +\node[ext, minimum width=22mm] (agent) at (16.8,3.8) {Agent}; + +% ============================ DSAgt stores ============================== +\node[store, minimum width=34mm] (catalog) at (5.6,0.6) + {\textbf{Skills Catalog}\\{\scriptsize federated across sources}\\{\scriptsize searchable $\cdot$ not yet installed}}; + +\node[store, minimum width=46mm, minimum height=26mm] (native) at (11.6,0.9) {}; +\node[anchor=north, align=center, font=\footnotesize] + at ([yshift=-1.6mm]native.north) + {Skill Directory\\{\scriptsize installed + created}\\{\scriptsize .claude $\cdot$ .agents $\cdot$ .cline $\cdot$ .roo /skills}}; +\node[skill, minimum width=26mm, anchor=south] at ([yshift=2mm]native.south) + {\texttt{skill-creator}}; + +% ============================ DSAgt service ============================= +\node[svc, minimum width=42mm] (router) at (5.6,5.8) + {\texttt{SkillRouter}\\{\scriptsize search $\cdot$ keyword fallback $\cdot$ list\_sources}}; + +% ============================ router operations (skill MCP tools) ======= +\draw[link] (router.west) -| (extrepos.north) + node[lbl, pos=0.25, above] {add\_skill\_source\\list\_skill\_sources}; +\draw[bilink] (router.south) -- (catalog.north) + node[lbl, pos=0.5, right] {search\_skills}; +\draw[link] (router.east) -| (native.north) + node[lbl, pos=0.25, above] {install\_skill}; +\draw[bilink] (router.north) -- ++(0,0.5) -| (agent.north) + node[lbl, pos=0.25, above] {MCP}; + +% ============================ agent <-> skill directory (two-way) ======= +\draw[bilink] (native.east) -| (agent.south) + node[lbl, pos=0.22, below, align=center] {auto-invoke\\author}; + +% ============================ concern bands ============================= +\begin{scope}[on background layer] + \node[band, fit=(catalog), inner sep=3mm] (regband) {}; + \node[band, fit=(native), inner sep=3mm] (expband) {}; + \node[band, fit=(router), inner sep=3mm] (disband) {}; +\end{scope} +\node[concern] at (disband.north) [above=0.6mm] {DISCOVERY}; +\node[concern] at (regband.south) [below=0.6mm] {REGISTRATION}; +\node[concern] at (expband.south) [below=0.6mm] {PROGRESSIVE EXPOSURE}; + +\end{tikzpicture} +\end{document} diff --git a/mkdocs.yml b/mkdocs.yml index 5745bf2..c94b656 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -31,6 +31,10 @@ theme: - content.code.copy - content.code.annotate +plugins: + - search + - include-markdown + markdown_extensions: - admonition - pymdownx.details @@ -48,10 +52,13 @@ nav: - Home: index.md - Quick Start: quickstart.md - Architecture: architecture.md + - Capabilities: + - Provenance: provenance.md + - Knowledge Base: knowledge-base.md + - Skills: skills.md + - Memory: memory.md + - Observability: observability.md - MCP Servers: mcp-servers.md - - Knowledge Base: knowledge-base.md - - Tools & Skills: tools-skills.md - - Observability: observability.md - CLI Reference: cli.md - Use Cases: - Overview: use-cases/index.md diff --git a/pyproject.toml b/pyproject.toml index d2db3a8..99e81fa 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,10 +1,11 @@ [project] name = "dsagt" -version = "0.1.0" +dynamic = ["version"] description = "DataSmith Agent - AI-assisted data pipeline builder" readme = "README.md" requires-python = ">=3.12,<3.14" -license = {text = "TBD"} +license = "Apache-2.0" +license-files = ["LICENSE"] authors = [ {name = "BASE-DATA Team", email = "aaron.tuor@pnnl.gov"} ] @@ -12,57 +13,26 @@ keywords = ["data", "pipeline", "agent", "mcp", "ai"] dependencies = [ # Core - "pyyaml>=6.0", - "httpx>=0.27.0", - "mcp>=1.0.0", - # LiteLLM is used by MCP servers (knowledge: embeddings), - # session.run_extraction (memory extraction LLM call), and Phase 2 - # ``dsagt start --enable-proxy`` (LiteLLM proxy server bootstrap). - # The ``[proxy]`` extra pulls in backoff / fastapi / uvicorn / ... - # required by ``litellm.proxy.proxy_cli.run_server``. - "litellm[proxy]>=1.83.7", # floor: closes GHSA-r75f / GHSA-v4p8 / GHSA-xqmj — don't downgrade - "python-dotenv>=1.0", # transitive via litellm, pinned directly so a fresh venv always has it - # MCP-server / dsagt-run observability emits OTLP HTTP to MLflow's - # /v1/traces endpoint; floor=3.11.1 so we ship with the latest GenAI - # semconv support and well past the 3.6.0 protobuf-response bug fixed - # in 3.7.0. The proxy-side (_DSAGTMlflowLogger in provenance.py) still - # uses MLflow's MlflowLogger SDK directly, so MLflow remains a hard dep. - "mlflow>=3.11.1,<4.0", - "opentelemetry-api>=1.27", - "opentelemetry-sdk>=1.27", - "opentelemetry-exporter-otlp-proto-http>=1.27", + "pyyaml>=6.0", # config & tool-spec parsing + "httpx>=0.27.0", # HTTP client; API embeddings + "questionary>=2.0", # interactive `dsagt init` select/checkbox menus + "mcp>=1.0.0", # MCP server protocol + "mlflow==3.11.1", # trace store & observability # Knowledge base - "numpy>=1.26", - "llama-index-core>=0.11", - "tree-sitter-language-pack>=0.2", - "chromadb>=1.5.1", - "faiss-cpu>=1.8", - "sentence-transformers==5.4.0", - "transformers==4.47.0", - "llama-index-readers-file>=0.1", - # Sparse keyword retrieval for hybrid search. rank_bm25 is pure Python - # (~10 KB), no native deps; the corpus and IDF stats live in memory and - # are pickled to disk per-collection. Used alongside dense embeddings - # via Reciprocal Rank Fusion in KnowledgeBase.search when the route's - # ``hybrid`` flag is set (default True). - "rank-bm25>=0.2.2", - # csvkit ships ~14 well-behaved CSV CLI tools (csvcut, csvgrep, - # csvstat, csvjson, csvlook, ...). Used by the README quickstart's - # CSV-summary prompts. Replaces the previous ``csvtool`` dep, which - # had inverted exit codes (returned True/False from main(), then - # sys.exit(main()) so every successful run exited 1) — that broke - # agent error-detection on every csvtool call. - "csvkit>=2.2.0", - "h5py>=3.16.0" + "numpy>=1.26", # embedding vector math + "llama-index-core>=0.11", # document & code chunking + "llama-index-readers-file>=0.1", # SimpleDirectoryReader file readers + "tree-sitter-language-pack>=0.2,<1.6", # code-aware chunking (1.6+ rewrite breaks CodeSplitter) + "chromadb>=1.5.1", # vector store + "sentence-transformers==5.4.0", # local embeddings & reranking + "transformers==4.47.0", # pin for sentence-transformers + "rank-bm25>=0.2.2", # sparse keyword retrieval ] [project.scripts] dsagt = "dsagt.commands.cli:main" -dsagt-run = "dsagt.commands.run_tool:main" -dsagt-proxy = "dsagt.commands.proxy_server:main" -dsagt-registry-server = "dsagt.commands.registry_server:main" -dsagt-knowledge-server = "dsagt.commands.knowledge_server:main" -dsagt-setup-kb = "dsagt.commands.setup_core_kb:main" +dsagt-run = "dsagt.commands.run_code:main" +dsagt-server = "dsagt.mcp.server:main" [dependency-groups] dev = [ @@ -72,23 +42,27 @@ dev = [ ] docs = [ "mkdocs-material>=9.5", + "mkdocs-include-markdown-plugin>=6.0", ] [project.optional-dependencies] fix-certificates = ["truststore>=0.10.4"] [build-system] -requires = ["setuptools>=61.0", "wheel"] +requires = ["setuptools>=77.0", "wheel"] build-backend = "setuptools.build_meta" [tool.setuptools] package-dir = {"" = "src"} +[tool.setuptools.dynamic] +version = {attr = "dsagt.__version__"} + [tool.setuptools.packages.find] where = ["src"] [tool.setuptools.package-data] -dsagt = ["py.typed", "tools/*.md", "tools/*.py", "skills/**/*", "dsagt_instructions.md"] +dsagt = ["py.typed", "codes/**/*", "skills/**/*", "dsagt_instructions.md"] [tool.pytest.ini_options] testpaths = ["tests"] diff --git a/src/dsagt/__init__.py b/src/dsagt/__init__.py index 6d1b160..df446b9 100644 --- a/src/dsagt/__init__.py +++ b/src/dsagt/__init__.py @@ -4,17 +4,20 @@ AI-assisted data pipeline builder for MCP-compatible agents. """ -__version__ = "0.1.0" +# Single source of truth for the package version: pyproject.toml reads this +# via `[tool.setuptools.dynamic] version = {attr = "dsagt.__version__"}`. +__version__ = "0.2.0" # Cap CPU thread count for embedding / tokenization libraries before any # heavy imports happen. Without this, PyTorch / sentence-transformers / # numpy+MKL default to using every available core, which pegs the # machine and causes visible system unresponsiveness during embed bursts -# (kb_ingest, kb_search bursts, setup-kb model load). Half the physical +# (kb_ingest, kb_search bursts, init's KB build). Half the physical # cores is a sensible default that leaves headroom for the OS, the # agent process, OneDrive sync, IDE, browser, etc. ``setdefault`` # preserves any value the user has already exported in their shell. import os as _os + _default_threads = str(max(1, (_os.cpu_count() or 4) // 2)) _os.environ.setdefault("OMP_NUM_THREADS", _default_threads) _os.environ.setdefault("MKL_NUM_THREADS", _default_threads) @@ -38,6 +41,6 @@ pass del _os, _default_threads -from dsagt.registry import ToolRegistry +from dsagt.registry import CodeRegistry -__all__ = ["ToolRegistry", "__version__"] +__all__ = ["CodeRegistry", "__version__"] diff --git a/src/dsagt/agents/__init__.py b/src/dsagt/agents/__init__.py index 2db06de..4899ee8 100644 --- a/src/dsagt/agents/__init__.py +++ b/src/dsagt/agents/__init__.py @@ -2,72 +2,21 @@ Agent platform configuration and launch. Generates platform-specific config files (MCP server entries, agent -instructions, env vars) from the single ``dsagt_config.yaml``. Launches +instructions, env vars) from the single ``.dsagt/config.yaml``. Launches the agent process in the foreground and blocks until it exits. -Post-proxy DSAGT: each agent talks directly to its own provider. We -inject MLflow + OTel env vars so every agent's native LLM-call telemetry -lands in the project's MLflow store, but model selection / API keys / -provider base URLs are the user's responsibility. - -Two orthogonal questions about each agent: - -1. **Native visibility.** Without ``--enable-proxy``, are the agent's - LLM calls visible in the MLflow UI (audit, drill-in, search)? -2. **Native extraction.** Without ``--enable-proxy``, does end-of- - session memory extraction work? - -Answer to #2 is uniformly **no for non-proxy paths**, on purpose. Each -agent that emits OTel does so in its own shape (Claude Code: span -events; Goose: domain spans; LiteLLM autolog: ``mlflow.spanInputs`` / -``mlflow.spanOutputs``). Memory extraction reads a single canonical -shape (the LiteLLM-autolog shape, emitted only by ``dsagt-proxy``) so -the parser stays small and we avoid per-agent maintenance forever. -Run with ``--enable-proxy`` to get extraction. - -Native visibility (``otel_payload_support`` class var, per-agent): - - ========= ============ ============================================= - Agent Support tier Native MLflow visibility (no proxy) - ========= ============ ============================================= - claude full yes — every turn lands as a trace with - messages, response, tool_use blocks - (gated by 4 ``CLAUDE_CODE_*`` / - ``OTEL_LOG_*`` flags we set) - goose full yes — ``dispatch_tool_call`` span carries - tool + arguments JSON - codex partial limited — Codex OTel spans carry only token - counts and tool names; full conversation - lives in ``$CODEX_HOME/sessions/rollout-*.jsonl`` - (side-channel reader not yet wired) - cline none no — Cline emits no OTel spans at all; - the agent is a black box from DSAgt's view - roo none no — Roo Code imports zero OTel SDKs; - PostHog telemetry is payload-free - ========= ============ ============================================= - -Pick the run mode by what you need: - -* **Visibility only** (Claude Code / Goose): default ``dsagt start`` - is enough. Real-time audit + drill-in works in the MLflow UI. - Memory extraction will produce nothing. -* **Visibility + extraction** (any agent): ``dsagt start - --enable-proxy``. Adds one subprocess; routes every LLM call - through it; lands traces in canonical shape that extraction reads. -* **Visibility for non-emitters** (Cline / Roo / Codex partial): - ``dsagt start --enable-proxy`` is the only way to see the agent's - LLM calls at all. Extraction is a free downstream consequence. - -Tool execution provenance (``dsagt-run`` ``tool.execute`` spans) and KB -observability (``kb.*`` / ``registry.*`` spans from MCP servers) always -work via OTLP regardless of run mode — the proxy decision affects only -the agent's own LLM-call traces. +BYOA: each agent talks directly to its own provider. DSAGT forces no +telemetry env on the agent — agent LLM-call history is recovered +post-hoc from the agent's on-disk session record, not by native OTel +emission. We set only ``MLFLOW_TRACKING_URI`` (so the MCP servers and +any MLflow client log to the project's store) and per-project state +dirs. Model selection / API keys / provider base URLs are the user's +responsibility. Each agent's quirks live in its own module — see ``base.py`` for the :class:`AgentSetup` ABC and one of the subclass modules -(``claude.py``, ``goose.py``, ``cline.py``, ``roo.py``, ``codex.py``) -for the platform-specific details and the per-agent investigation -behind its support tier. +(``claude.py``, ``goose.py``, ``cline.py``, ``codex.py``, ``opencode.py``) +for the platform-specific details. Public API exported here: @@ -76,8 +25,8 @@ - :func:`static_agent_record` — write instructions + state dirs. - :func:`static_agent_files_present` — has the static record been written? - :func:`dynamic_agent_record` — write runtime-dependent files. +- :func:`refresh_native_skills` — re-run the native-skills mirror on demand. - :func:`launch_agent` — fork the agent and block until exit. -- :func:`agent_otel_support` — query an agent's OTel payload tier. """ from __future__ import annotations @@ -92,14 +41,12 @@ _build_mcp_servers_dict, _mcp_env_block, _mcp_server_args, - _PROXY_FORWARDED_SENTINEL, ) from .claude import ClaudeSetup from .cline import ClineSetup from .codex import CodexSetup, _render_codex_config from .goose import GooseSetup from .opencode import OpenCodeSetup -from .roo import RooSetup logger = logging.getLogger(__name__) @@ -112,7 +59,6 @@ ClaudeSetup, GooseSetup, ClineSetup, - RooSetup, CodexSetup, OpenCodeSetup, ) @@ -120,26 +66,15 @@ AGENTS: dict[str, type[AgentSetup]] = {cls.name: cls for cls in _AGENT_CLASSES} -def _setup_for(agent_name: str, proxy_port: int | None = None) -> AgentSetup: +def _setup_for(agent_name: str) -> AgentSetup: """Return a fresh :class:`AgentSetup` instance for ``agent_name``. - *proxy_port*, when set, is forwarded to the constructor so the - instance binds its proxy-mode methods (``write_dynamic`` → - ``proxy_write_dynamic``, ``run_script`` → ``proxy_run_script``). - Callers don't have to know about the dispatch — they always call - ``setup.write_dynamic(...)`` / ``setup.run_script(...)``. - Raises ``KeyError`` with a helpful message if the agent isn't registered. """ cls = AGENTS.get(agent_name) if cls is None: raise KeyError(f"Unknown agent {agent_name!r}. Registered: {sorted(AGENTS)}.") - return cls(proxy_port=proxy_port) - - -def _proxy_port_from_config(config: dict) -> int | None: - """Extract the active proxy port from a config dict, if any.""" - return (config.get("proxy") or {}).get("port") + return cls() # --------------------------------------------------------------------------- @@ -154,18 +89,14 @@ def agent_env(config: dict) -> dict: 1. ``os.environ`` (user's shell env). 2. DSAGT-wide vars (``DSAGT_PROJECT``, ``DSAGT_PROJECT_DIR``, ``DSAGT_AGENT``, ``DSAGT_SESSION_ID``). - 3. MLflow + OTel telemetry env (``MLFLOW_TRACKING_URI``, - ``OTEL_EXPORTER_OTLP_ENDPOINT``, ``OTEL_EXPORTER_OTLP_HEADERS``, - ``OTEL_RESOURCE_ATTRIBUTES``) so the agent's native OTel SDK and - the MCP servers' ``init_tracing`` both land traces in the same - MLflow experiment. - 4. Per-agent overrides via :meth:`AgentSetup.env_overrides` — - each setup class owns its quirks: claude code's verbosity - flags, goose's ``GOOSE_PROVIDER``/``GOOSE_MODEL`` selectors, - translation of ``llm.{api_key,base_url}`` into the env-var - names that agent's runtime reads, codex/cline state-dir env, etc. - 5. Proxy overrides (when ``--enable-proxy`` set) — last so they - win over per-agent provider env. + 3. ``MLFLOW_TRACKING_URI`` so the MCP servers' ``init_tracing`` and + any MLflow client running under the agent log to the project's + store. No OTel routing env — DSAGT does not force native agent + telemetry; agent traces are recovered post-hoc from the on-disk + transcript. + 4. Per-agent dsagt-owned runtime env via + :meth:`AgentSetup.runtime_env` — per-project state dirs only + (``CLINE_DIR``, ``CODEX_HOME``). """ pdir = config["project_dir"] agent_name = config["agent"] @@ -178,129 +109,17 @@ def agent_env(config: dict) -> dict: if config.get("session_id"): env["DSAGT_SESSION_ID"] = config["session_id"] - mlflow_port = config.get("mlflow", {}).get("port") - proxy_port_for_otel = (config.get("proxy") or {}).get("port") - if mlflow_port: - mlflow_url = f"http://localhost:{mlflow_port}" - env["MLFLOW_TRACKING_URI"] = mlflow_url - # Claude in BYOA mode (no proxy) uses ``mlflow autolog claude`` - # for agent-side traces — its Stop hook produces richer - # transcript-based traces than native OTel. Skip OTel routing - # so we don't get duplicate (and inferior) trace shapes. MCP - # servers and dsagt-run still get tracing because their - # ``init_tracing`` reads MLflow URL from cwd's ``dsagt_config.yaml``, - # not these env vars. - skip_otel_for_claude_byoa = agent_name == "claude" and not proxy_port_for_otel - if not skip_otel_for_claude_byoa: - # OTel endpoint for the agent's native telemetry SDK. The - # x-mlflow-experiment-id header is mandatory for MLflow's OTLP - # receiver — resolve to the experiment's numeric id once at - # startup; if MLflow can't be reached yet we still write the env - # vars so init_tracing can resolve later. - env["OTEL_EXPORTER_OTLP_ENDPOINT"] = f"{mlflow_url}/v1/traces" - experiment_id = _resolve_experiment_id(mlflow_url, config["project"]) - if experiment_id: - env["OTEL_EXPORTER_OTLP_HEADERS"] = ( - f"x-mlflow-experiment-id={experiment_id}" - ) - # Resource attributes flow onto every span emitted by the agent's - # OTel SDK. ``session.id`` is what MLflow's OTLP receiver - # promotes to ``mlflow.trace.session`` trace_metadata — required - # so end-of-session memory extraction can find this run's traces. - resource_attrs = [f"service.name={agent_name}"] - if config.get("session_id"): - resource_attrs.append(f"session.id={config['session_id']}") - env["OTEL_RESOURCE_ATTRIBUTES"] = ",".join(resource_attrs) - - pre_runtime_env = dict(env) - # BYOA: only dsagt-owned env (telemetry capture flags, per-project - # state dirs). Provider credentials live in the user's shell; - # ``env_overrides`` is reserved for Phase 2 proxy mode. - env.update(setup.runtime_env(config)) - - # Transparency: surface what credential env vars the agent will - # actually pick up from the user's shell, so a missing or wrong - # value isn't silently consumed. - proxy_port = (config.get("proxy") or {}).get("port") - if not proxy_port: - _warn_on_preconfigured_creds(setup, env, pre_runtime_env) - - # Opt-in proxy routing (Phase 2). When ``--enable-proxy`` populated - # ``config["proxy"]["port"]``, delegate to the agent's proxy hooks. - # ``proxy_env_overrides`` sets the proxy URL; ``env_overrides`` is - # the per-agent llm-config translation hook the proxy path will use. - if proxy_port: - env.update(setup.env_overrides(config)) - env.update(setup.proxy_env_overrides(proxy_port)) - - return env - + from dsagt.observability import resolve_tracking_uri -def _warn_on_preconfigured_creds( - setup: AgentSetup, - env: dict, - pre_overrides_env: dict, -) -> None: - """Emit a one-line transparency warning when our env_overrides - didn't populate any of the agent's credential env vars but the - user's shell did. - - Lists the var NAMES (never values — keys are secrets). Helps the - user verify what the agent will actually pick up when the project - YAML is empty or has unresolved ``${VAR}`` placeholders. - - Skipped when the proxy is enabled (proxy_env_overrides plants - everything we need) and when the agent has no credential env vars - declared (IDE-only agents). - """ - if not setup.credential_env_vars: - return - - # What did our env_overrides actually inject? - injected = { - k - for k in setup.credential_env_vars - if env.get(k) and env.get(k) != pre_overrides_env.get(k) - } - if injected: - return # Project YAML supplied credentials; nothing to warn about. - - # Nothing injected — list the vars present from the user's shell. - from_shell = [k for k in setup.credential_env_vars if env.get(k)] - if from_shell: - logger.warning( - "%s: project config has no llm credentials — agent will use " - "preconfigured env vars: %s", - setup.name, - ", ".join(from_shell), - ) - else: - logger.warning( - "%s: project config has no llm credentials and none of %s are " - "set in the shell — agent may fall back to its own auth flow " - "(claude.ai subscription, codex login, etc.) or fail at first call.", - setup.name, - ", ".join(setup.credential_env_vars), - ) + env["MLFLOW_TRACKING_URI"] = resolve_tracking_uri(config) + # BYOA: only dsagt-owned env (per-project state dirs). DSAGT forces + # no telemetry env on the agent — agent traces are recovered post-hoc + # from the on-disk transcript. Provider credentials live in the + # user's shell. + env.update(setup.runtime_env(config)) -def _resolve_experiment_id(mlflow_url: str, project_name: str) -> str | None: - """Look up the MLflow experiment id for *project_name*; create if absent. - - Returns None when MLflow isn't reachable yet — agent_env writes the - rest of the OTel block anyway and ``init_tracing`` retries on the - server side. We keep this best-effort because dsagt-launch can race - against MLflow startup; observability.py's _resolve_experiment_id - runs again per-process and will succeed once MLflow is up. - """ - try: - import mlflow - - mlflow.set_tracking_uri(mlflow_url) - return str(mlflow.set_experiment(project_name).experiment_id) - except Exception as e: - logger.debug("could not resolve experiment id at agent_env time: %s", e) - return None + return env def agent_command(config: dict) -> list[str]: @@ -308,17 +127,6 @@ def agent_command(config: dict) -> list[str]: return _setup_for(config["agent"]).interactive_command(config) -def agent_otel_support(agent_name: str) -> str: - """Return the OTel-payload support tier for *agent_name*. - - See the module docstring matrix for the meaning of each tier. - Consumers (smoke-test, ``dsagt info``, future warning systems) use - this to decide whether to hard-fail or soft-warn when an agent's - LLM-call traces don't appear in MLflow. - """ - return _setup_for(agent_name).otel_payload_support - - def static_agent_record( config: dict, agent: str, @@ -344,32 +152,56 @@ def dynamic_agent_record( env: dict, working_dir: str | Path, ) -> list[str]: - """Write the agent's runtime-dependent files: MCP config + ``.dsagt_env`` - + ``dsagt-launch.sh``. + """Write the agent's runtime-dependent files: the per-agent MCP config. Caller must have already: - Resolved the agent and stored it in ``config["agent"]`` - - Started MLflow and updated ``config["mlflow"]["port"]`` to the - actually-bound port - Built ``env`` via :func:`agent_env` - """ - from .base import _write_launch_shim - setup = _setup_for(config["agent"], proxy_port=_proxy_port_from_config(config)) + No launch shim is written — ``dsagt init`` collapses to config + + instructions + MCP config; the user starts the agent directly in the + project dir or via ``dsagt start``. + """ + setup = _setup_for(config["agent"]) actions = setup.write_dynamic( config, env, Path(working_dir), Path(config["project_dir"]), ) - # Launch shim is BYOA-only. Skip when proxy mode is active — - # ``dsagt start --enable-proxy`` is the only sensible entry point - # in that mode (proxy URL must be plumbed through agent env). - if not _proxy_port_from_config(config): - actions.append(_write_launch_shim(setup, config, Path(working_dir))) + # Mirror installed skills into the agent's native skills dir (all agents). + # Central here so each agent only declares native_skills_dir. + actions += setup.setup_skills(Path(working_dir), config) return actions +def refresh_native_skills(working_dir: str | Path) -> list[str]: + """Re-run the native-skills mirror for the project's configured agent. + + Called by the MCP tools right after a skill is installed/created or a + code is registered, so the native skills dir is current the moment the + files land — the next session auto-discovers them no matter how the agent + is launched (bare or ``dsagt start``). An already-running session + enumerates its skills at startup (agent-side behavior), but can use a + fresh skill right away by reading its SKILL.md — which is all native + invocation does. Idempotent (manifest-tracked, the same mirror + :func:`dynamic_agent_record` runs at init/start). + + No-op when ``working_dir`` has no ``.dsagt/config.yaml`` agent — the dir + hasn't been ``dsagt init``-ed, so there is no native skills dir to mirror + into (the test-facing ``create_*_server`` wrappers over bare tmp dirs). + """ + # Lazy: session drags in knowledge/provenance at module level, which this + # package (imported by the CLI at cold start) must not pay for. + from dsagt.session import read_config_file + + working_dir = Path(working_dir) + config = read_config_file(working_dir) + if not config.get("agent"): + return [] + return _setup_for(config["agent"]).setup_skills(working_dir, config) + + def launch_agent( config: dict, env: dict, @@ -379,7 +211,7 @@ def launch_agent( ) -> int: """Launch the agent in the foreground. Blocks until it exits.""" working_dir = Path(working_dir) - setup = _setup_for(config["agent"], proxy_port=_proxy_port_from_config(config)) + setup = _setup_for(config["agent"]) if script_path is not None: return setup.run_script( @@ -406,7 +238,6 @@ def launch_agent( "AgentSetup", "agent_command", "agent_env", - "agent_otel_support", "dynamic_agent_record", "launch_agent", "static_agent_files_present", @@ -415,6 +246,5 @@ def launch_agent( "_build_mcp_servers_dict", "_mcp_env_block", "_mcp_server_args", - "_PROXY_FORWARDED_SENTINEL", "_render_codex_config", ] diff --git a/src/dsagt/agents/base.py b/src/dsagt/agents/base.py index 2835d80..21c2ef6 100644 --- a/src/dsagt/agents/base.py +++ b/src/dsagt/agents/base.py @@ -12,7 +12,7 @@ import json import logging -import shlex +import shutil import subprocess from abc import ABC, abstractmethod from pathlib import Path @@ -31,99 +31,34 @@ # files between init and start without losing edits on the next start. _DSAGT_MARKER = "DSAgt Pipeline Builder" -# Tools each dsagt MCP server exposes — listed in ``alwaysAllow`` so roo -# and cline auto-approve them without a human-in-the-loop prompt. Keep in -# sync with ``commands/registry_server.py`` and -# ``commands/knowledge_server.py`` tool registrations; a tool added there -# but not here means roo/cline will hang on its first call. -_DSAGT_MCP_ALWAYS_ALLOW = { - "registry": [ - "get_registry", - "http_request", - "install_dependencies", - "read_file", - "reconstruct_pipeline", - "run_command", - "save_skill", - "save_tool_spec", - "search_registry", - "search_skills", - ], - "knowledge": [ - "kb_add_vector_db", - "kb_append", - "kb_dismiss_suggestion", - "kb_get_memories", - "kb_get_suggestions", - "kb_ingest", - "kb_job_status", - "kb_list_collections", - "kb_remember", - "kb_search", - ], -} - - -# Sentinel API key planted in agent / MCP-child env when the optional -# proxy is enabled, so any direct call that bypasses the proxy returns -# 401 fast — fails loudly instead of silently leaking around our -# observability pipeline. -_PROXY_FORWARDED_SENTINEL = "dsagt-proxy-forwarded-disable-direct-calls" - - -def _real(value) -> str | None: - """Return *value* trimmed, or ``None`` if blank or an unresolved - ``${VAR}`` interpolation. Centralized so each agent's - ``env_overrides`` filters config values consistently — we never - propagate a placeholder to a downstream env var. - """ - s = (value or "").strip() if isinstance(value, str) else "" - if not s or s.startswith("${"): - return None - return s - - -def _anthropic_env(llm: dict) -> dict[str, str]: - """Anthropic-protocol env vars from ``llm.*`` config. - - Returns ``{}`` unless ``llm.provider == "anthropic"`` — never - propagates anthropic vars on top of an openai-shaped upstream, - where they'd be meaningless or actively misleading. - - Called only by agent setups whose native runtime speaks anthropic - (claude code), and by multi-protocol agents (goose, cline, roo) - when the configured provider matches. - """ - if (_real(llm.get("provider")) or "").lower() != "anthropic": - return {} - out: dict[str, str] = {} - if api_key := _real(llm.get("api_key")): - out["ANTHROPIC_API_KEY"] = api_key - if base_url := _real(llm.get("base_url")): - out["ANTHROPIC_BASE_URL"] = base_url - if model := _real(llm.get("model")): - out["ANTHROPIC_MODEL"] = model - return out - - -def _openai_env(llm: dict) -> dict[str, str]: - """OpenAI-protocol env vars from ``llm.*`` config. - - Returns ``{}`` unless ``llm.provider`` is one of the OpenAI-wire - aliases (``openai``, ``openai_like``, ``azure``). Only called by - agents whose native runtime speaks the OpenAI wire protocol (codex, - LiteLLM-backed paths) or by multi-protocol agents (goose, cline, - roo) when the configured provider matches. - """ - provider = (_real(llm.get("provider")) or "").lower() - if provider not in ("openai", "openai_like", "azure"): - return {} - out: dict[str, str] = {} - if api_key := _real(llm.get("api_key")): - out["OPENAI_API_KEY"] = api_key - if base_url := _real(llm.get("base_url")): - out["OPENAI_BASE_URL"] = base_url - return out +# Tools the dsagt MCP server exposes — listed in ``alwaysAllow`` so cline +# auto-approves them without a human-in-the-loop prompt. Keep in +# sync with the ``mcp/*_tools.py`` tool registrations (registry / knowledge / +# memory / skill); a tool added there but not here means cline will hang on +# its first call. All dsagt MCP tools live behind the single ``dsagt-server``, +# so the always-allow list is one flat union. +_DSAGT_MCP_ALWAYS_ALLOW = [ + "add_skill_source", + "get_registry", + "http_request", + "install_dependencies", + "install_skill", + "kb_append", + "kb_get_memories", + "kb_ingest", + "kb_job_status", + "kb_list_collections", + "kb_remember", + "kb_search", + "list_skill_sources", + "read_file", + "reconstruct_pipeline", + "run_command", + "save_skill", + "save_code_spec", + "search_registry", + "search_skills", +] # --------------------------------------------------------------------------- @@ -131,28 +66,39 @@ def _openai_env(llm: dict) -> dict[str, str]: # --------------------------------------------------------------------------- -def _mcp_server_args(server: str) -> list[str]: - """Build the argv tail for ``uv run dsagt--server``. +def _mcp_server_args() -> list[str]: + """Build the argv tail for ``uv run dsagt-server``. - Both servers read all configuration from ``DSAGT_PROJECT_DIR`` (env) - and dsagt_config.yaml — no CLI args needed. + The single merged server reads all configuration from the project's + ``.dsagt/config.yaml`` (located via cwd-walk) — no CLI args needed. """ - return ["run", f"dsagt-{server}-server"] + return ["run", "dsagt-server"] def _mcp_env_block(config: dict) -> dict[str, str]: """Env vars the dsagt MCP server children need at startup. - Project routing (project name, project_dir, mlflow port) lives in - ``dsagt_config.yaml`` and is read by services via cwd-walk — single - source of truth, no env duplication. This block carries only the - embedding-backend settings, which the embedding client still reads - from env (refactor TBD). For ``backend: api`` the user must set - ``EMBEDDING_API_KEY`` in their shell env (creds never on disk). + Benign routing only (no credentials, no provider redirection): the + project name + dir, the serverless ``MLFLOW_TRACKING_URI``, and the + embedding-backend settings. MCP children run with cwd == project_dir + and could read most of this from ``.dsagt/config.yaml``, but agents that + don't inherit the parent's shell env into their MCP children (codex / + cline) need it baked into the per-agent MCP config. For + ``backend: api`` the user still sets ``EMBEDDING_API_KEY`` in their shell + (creds never on disk). + + No session id here — the MCP server mints it at startup into + ``.dsagt/state.yaml`` (it owns the session lifecycle now), so there's + nothing to thread through the env. """ + from dsagt.observability import resolve_tracking_uri + emb = config.get("embedding") or {} block: dict[str, str] = {} for key, src in ( + ("DSAGT_PROJECT", config.get("project")), + ("DSAGT_PROJECT_DIR", config.get("project_dir")), + ("MLFLOW_TRACKING_URI", resolve_tracking_uri(config)), ("EMBEDDING_BACKEND", emb.get("backend")), ("EMBEDDING_MODEL", emb.get("model")), ("EMBEDDING_BASE_URL", emb.get("base_url")), @@ -170,29 +116,6 @@ def _load_master_instructions() -> str | None: return None -def _format_roomodes(instructions: str) -> str: - """Wrap master instructions as a Roo Code .roomodes JSON file.""" - parts = instructions.split("\n## ", 1) - role = parts[0].strip() - custom = ("## " + parts[1]).strip() if len(parts) > 1 else "" - - return json.dumps( - { - "customModes": [ - { - "slug": "dsagt", - "name": "DSAgt Pipeline Builder", - "roleDefinition": role, - "customInstructions": custom, - "groups": ["read", "edit", "browser", "command", "mcp"], - "source": "project", - } - ] - }, - indent=2, - ) - - def _append_or_write(path: Path, content: str, marker: str) -> str | None: """Idempotent write for instructions files. @@ -212,26 +135,100 @@ def _append_or_write(path: Path, content: str, marker: str) -> str | None: return f"Wrote {path}" +#: Claude Code caps a skill's frontmatter description (combined with +#: when_to_use) at this many characters; longer ones are rejected. We +#: truncate the *mirrored* copy only, never the project source. +_NATIVE_DESCRIPTION_CAP = 1536 + +#: Manifest filename inside a native skills dir listing the skill names +#: dsagt placed there, so the mirror can reap its own stale entries on +#: re-run without ever touching user-authored skills. +_SKILL_MANIFEST = ".dsagt-managed.json" + + +def _truncate_native_description(skill_md: Path) -> None: + """If the mirrored SKILL.md's description exceeds the native cap, trim it.""" + import yaml + + text = skill_md.read_text() + if not text.startswith("---"): + return + parts = text.split("---", 2) + if len(parts) < 3: + return + try: + front = yaml.safe_load(parts[1]) or {} + except yaml.YAMLError: + return + desc = front.get("description") + if isinstance(desc, str) and len(desc) > _NATIVE_DESCRIPTION_CAP: + front["description"] = desc[: _NATIVE_DESCRIPTION_CAP - 1].rstrip() + "…" + new_front = yaml.dump(front, default_flow_style=False, sort_keys=False) + skill_md.write_text(f"---\n{new_front}---{parts[2]}") + + +def _mirror_skills_to(target_dir: Path, skill_dirs: list[Path]) -> list[str]: + """Idempotently mirror *skill_dirs* into *target_dir* (e.g. .claude/skills). + + Copies each skill directory (SKILL.md + scripts/ + references/) under + ``target_dir//``. A manifest tracks the names dsagt owns so a + later run reaps skills that were removed upstream **without ever + touching user-authored skills** that dsagt didn't place. ``skill_dirs`` + should list bundled dirs before project dirs so a project skill wins a + name collision (copied last). + """ + actions: list[str] = [] + manifest_path = target_dir / _SKILL_MANIFEST + previously: list[str] = [] + if manifest_path.exists(): + try: + previously = json.loads(manifest_path.read_text()) + except (json.JSONDecodeError, OSError): + previously = [] + + target_dir.mkdir(parents=True, exist_ok=True) + managed: list[str] = [] + for src in skill_dirs: + if not (src / "SKILL.md").exists(): + continue + name = src.name + dest = target_dir / name + if dest.exists(): + shutil.rmtree(dest) + shutil.copytree(src, dest) + _truncate_native_description(dest / "SKILL.md") + if name not in managed: + managed.append(name) + + # Reap skills dsagt placed previously that are gone from the source set. + for stale in set(previously) - set(managed): + stale_dir = target_dir / stale + if stale_dir.is_dir(): + shutil.rmtree(stale_dir, ignore_errors=True) + + manifest_path.write_text(json.dumps(sorted(managed), indent=2) + "\n") + if managed: + actions.append(f"Mirrored {len(managed)} skill(s) into {target_dir}") + return actions + + def _build_mcp_servers_dict(env_block: dict | None) -> dict: - """Build the standard ``{"mcpServers": {...}}`` dict for dsagt servers. + """Build the standard ``{"mcpServers": {...}}`` dict for the dsagt server. - Used by agents that load MCP config from a JSON file (roo via - ``.roo/mcp.json``). Claude Code uses the same shape via ``.mcp.json`` + Used by agents that load MCP config from a JSON file. Claude Code + uses this shape via ``.mcp.json`` but builds it inline in :class:`ClaudeSetup.write_dynamic`. Cline - doesn't use this — it requires ``cline mcp add`` to register servers. + doesn't use this — it requires ``cline mcp add`` to register the server. """ - mcp_config: dict = {"mcpServers": {}} - for server in ("registry", "knowledge"): - entry: dict = { - "command": "uv", - "args": _mcp_server_args(server), - "disabled": False, - "alwaysAllow": _DSAGT_MCP_ALWAYS_ALLOW[server], - } - if env_block: - entry["env"] = env_block - mcp_config["mcpServers"][f"dsagt-{server}"] = entry - return mcp_config + entry: dict = { + "command": "uv", + "args": _mcp_server_args(), + "disabled": False, + "alwaysAllow": _DSAGT_MCP_ALWAYS_ALLOW, + } + if env_block: + entry["env"] = env_block + return {"mcpServers": {"dsagt": entry}} def _toml_quote(value: str) -> str: @@ -266,114 +263,6 @@ def _run_simple_script( return 0 -# --------------------------------------------------------------------------- -# Launch shim — dsagt-launch.sh -# --------------------------------------------------------------------------- - - -def _render_launch_shim(setup: AgentSetup, config: dict) -> str: - """Render the dsagt-launch.sh shell script for a project. - - The shim is the BYOA "transparency" entry point — users either run it - directly (``bash dsagt-launch.sh``), or read it and execute the lines - manually. It does NOT exec the agent — instead it prints how to - launch (CLI / VS Code), letting the user pick. - - Steps: - 1. Start MLflow in the background if not already running. - 2. Resolve the experiment-id for this project at run time - (curl MLflow's REST API). - 3. Export OTel routing env (skipped for Claude — mlflow autolog - claude's ``.claude/settings.json`` handles agent-side traces). - 4. Export per-agent telemetry verbosity flags. - 5. Print available agent-launch options. - """ - project = config["project"] - mlflow_port = (config.get("mlflow") or {}).get("port") - agent_name = setup.name - pdir = config.get("project_dir") or "." - - cli_cmd = " ".join(shlex.quote(p) for p in setup.interactive_command({})) - vscode_lines = setup.vscode_hint(Path(str(pdir))) or [] - - # Claude in BYOA gets agent-side traces via mlflow autolog claude's - # Stop hook, so we omit the OTel routing block (would create - # duplicate, lower-fidelity traces alongside the rich transcript). - skip_otel_routing = agent_name == "claude" - - lines: list[str] = [] - lines.append("#!/usr/bin/env bash") - lines.append("# dsagt-launch.sh — start MLflow, set env, show launch options.") - lines.append( - "# Generated by `dsagt init`. Re-running `dsagt init` overwrites this file." - ) - lines.append("set -euo pipefail") - lines.append('cd "$(dirname "$0")"') - lines.append("") - lines.append("# 1. Start MLflow in the background if not already running.") - lines.append(f"dsagt mlflow {shlex.quote(project)} --background-only") - lines.append("") - - if not skip_otel_routing and mlflow_port: - lines.append("# 2. Resolve experiment id for OTel routing.") - lines.append( - f"EXPERIMENT_ID=$(curl -s " - f'"http://localhost:{mlflow_port}/api/2.0/mlflow/experiments/get-by-name?experiment_name={project}" ' - '| python3 -c \'import json,sys; print(json.load(sys.stdin)["experiment"]["experiment_id"])\')' - ) - lines.append("") - lines.append( - "# 3. OTel routing env (agent's native OTel SDK ships traces to MLflow)." - ) - lines.append(f'export MLFLOW_TRACKING_URI="http://localhost:{mlflow_port}"') - lines.append("export OTEL_EXPORTER_OTLP_PROTOCOL=http/protobuf") - lines.append( - f'export OTEL_EXPORTER_OTLP_TRACES_ENDPOINT="http://localhost:{mlflow_port}/v1/traces"' - ) - lines.append( - 'export OTEL_EXPORTER_OTLP_HEADERS="x-mlflow-experiment-id=$EXPERIMENT_ID"' - ) - lines.append(f'export OTEL_RESOURCE_ATTRIBUTES="service.name={agent_name}"') - lines.append("") - elif mlflow_port: - lines.append( - "# Claude uses `mlflow autolog claude` (.claude/settings.json) for" - ) - lines.append("# agent-side traces — no OTel routing needed.") - lines.append(f'export MLFLOW_TRACKING_URI="http://localhost:{mlflow_port}"') - lines.append("") - - if setup.telemetry_env: - lines.append("# 4. Agent telemetry verbosity flags.") - for k, v in setup.telemetry_env.items(): - lines.append(f"export {k}={shlex.quote(v)}") - lines.append("") - - lines.append("# 5. Show how to launch the agent — pick one.") - lines.append('echo ""') - lines.append('echo "Environment ready. Launch the agent in one of these ways:"') - lines.append('echo ""') - lines.append(f'echo " CLI: {cli_cmd}"') - if vscode_lines: - for hint in vscode_lines: - lines.append(f'echo " VS Code: {hint}"') - lines.append('echo ""') - lines.append('echo "Run any of the above in this shell."') - lines.append("") - return "\n".join(lines) - - -def _write_launch_shim(setup: AgentSetup, config: dict, working_dir: Path) -> str: - """Write ``dsagt-launch.sh`` to working_dir and chmod it executable. - - Returns a one-line action description for the init output. - """ - shim_path = working_dir / "dsagt-launch.sh" - shim_path.write_text(_render_launch_shim(setup, config)) - shim_path.chmod(0o755) - return f"Wrote {shim_path}" - - # --------------------------------------------------------------------------- # AgentSetup ABC # --------------------------------------------------------------------------- @@ -389,7 +278,7 @@ class AgentSetup(ABC): Class attributes (set by every subclass): - - ``name`` — agent identifier as used in ``dsagt_config.yaml`` + - ``name`` — agent identifier as used in ``.dsagt/config.yaml`` (e.g. ``"claude"``, ``"goose"``). - ``base_command`` — argv list that launches the agent interactively. Subclasses may override :meth:`interactive_command` @@ -407,70 +296,15 @@ class AgentSetup(ABC): static_marker: ClassVar[str] install_hint: ClassVar[str] = "Install the agent CLI first." - def __init__(self, proxy_port: int | None = None): - """Rebind ``write_dynamic`` / ``run_script`` to their proxy-mode - analogs (``proxy_write_dynamic`` / ``proxy_run_script``) when - Phase 2 proxy mode is active AND the subclass overrides them. - - Encapsulates the dispatch so callers always say ``setup.write_dynamic(...)`` - without knowing whether the BYOA or proxy implementation runs. - - Why the override check: the base default ``proxy_write_dynamic`` - delegates to ``self.write_dynamic`` for fall-through. If we - rebind ``self.write_dynamic = self.proxy_write_dynamic`` when - the subclass DOESN'T override, the delegation infinite-loops - (proxy_write_dynamic → self.write_dynamic → proxy_write_dynamic). - Only rebind when the subclass has its own implementation. - - Net effect: agents whose proxy + BYOA setup is identical (claude, - goose, roo's write_dynamic) inherit the BYOA path under proxy - mode too — no override, no rebind, ``setup.write_dynamic`` stays - as ``write_dynamic``. Agents that genuinely differ (cline + - codex + opencode write_dynamic; cline + roo run_script) get - their proxy_* override bound in. - """ - self.proxy_port = proxy_port - if proxy_port: - cls = type(self) - if cls.proxy_write_dynamic is not AgentSetup.proxy_write_dynamic: - self.write_dynamic = self.proxy_write_dynamic # type: ignore[method-assign] - if cls.proxy_run_script is not AgentSetup.proxy_run_script: - self.run_script = self.proxy_run_script # type: ignore[method-assign] - - #: Env vars the agent's runtime reads for credentials / endpoint - #: routing — i.e. the vars our ``env_overrides`` translates ``llm.*`` - #: config into. Used by ``agent_env`` to surface a transparency - #: warning when the project YAML is empty: we list which of these - #: are actually present in the user's shell so the user can see - #: what the agent will pick up. Empty for IDE-extension agents - #: that never read env-var credentials. - credential_env_vars: ClassVar[tuple[str, ...]] = () - - #: Whether this agent makes its LLM calls visible in MLflow natively - #: (without DSAgt's optional ``dsagt-proxy``). Drives whether - #: ``dsagt info`` / live audit / memory extraction see the agent's - #: reasoning + tool calls or just see the agent as a black box. - #: - #: ``"full"`` — verified end-to-end (Claude Code, Goose). - #: Every agent turn lands in MLflow as a trace - #: with messages + response + tool_use blocks. - #: ``--enable-proxy`` is unnecessary. - #: ``"partial"`` — agent emits OTel but spans don't carry - #: message content (Codex: only token counts + - #: tool names). Conversation may be available - #: via a non-OTel side-channel (Codex's - #: ``~/.codex/sessions/*.jsonl``); not yet - #: wired in. ``--enable-proxy`` recommended. - #: ``"none"`` — agent emits no OTel traces, or emits only - #: metrics without payloads (Cline, Roo Code). - #: The agent is a black box from DSAgt's - #: perspective; ``--enable-proxy`` is the only - #: way to see what it's doing. Tool execution - #: + KB observability still work regardless via - #: dsagt-run / MCP-server spans. - #: - #: See agents/.py docstrings for the per-agent investigation. - otel_payload_support: ClassVar[str] = "full" + #: Directory (relative to the working dir) the agent natively auto-discovers + #: ``SKILL.md`` skill folders from. ``setup_skills`` mirrors installed + #: (bundled + project) skills AND registered codes here so the agent + #: discovers/auto-invokes them without an MCP round-trip. Every supported + #: agent has one — claude ``.claude/skills``, codex/goose/opencode + #: ``.agents/skills`` (the cross-agent standard), cline ``.cline/skills``. + #: ``None`` would mean the agent has no native skill discovery + #: (none currently). + native_skills_dir: ClassVar[str | None] = None @abstractmethod def write_static(self, working_dir: Path) -> list[str]: @@ -489,69 +323,76 @@ def write_dynamic( working_dir: Path, pdir: Path, ) -> list[str]: - """Write the agent's runtime-dependent files (MCP config, .dsagt_env). + """Write the agent's runtime-dependent files (the per-agent MCP config). - Caller must have already populated ``config["mlflow"]["port"]`` with - the actually-bound port and built ``env`` via :func:`agent_env`. - Returns a list of one-line action descriptions. + Caller must have built ``env`` via :func:`agent_env`. Returns a + list of one-line action descriptions. """ - def runtime_env(self, config: dict) -> dict[str, str]: - """Dsagt-owned env vars the agent process needs at runtime (BYOA). + def setup_skills(self, working_dir: Path, config: dict) -> list[str]: + """Mirror installed skills AND registered codes into the agent's + native skills dir so it auto-discovers/auto-invokes them. - Returned dict is layered into ``agent_env`` and the launch shim's - telemetry block. Default returns ``telemetry_env`` (claude's - OTEL_LOG_* / CLAUDE_CODE_* gates, etc.). Subclasses augment with - per-project state-dir env (``CLINE_DIR``, ``CODEX_HOME``). + Codes share the skill-standard envelope (``codes//SKILL.md``), + so the same copy serves both: native discovery puts a code's exact + dsagt-run command in context at invocation time — a second discovery + path alongside ``search_registry``, aimed at the from-memory + command-reconstruction failure mode. - This is the only method allowed to set agent-affecting env in - BYOA. Provider credentials (ANTHROPIC_*, OPENAI_*, GOOSE_*) are - the user's responsibility — exported in their shell, never - translated from ``config["llm"]``. + Idempotent — the manifest-tracked :func:`_mirror_skills_to` only + reaps skills dsagt placed, never user-authored ones. No-op when the + agent declares no ``native_skills_dir`` or ``skills.populate_native`` + is disabled. """ - del config - return dict(self.telemetry_env) + if not self.native_skills_dir: + return [] + if not (config.get("skills") or {}).get("populate_native", True): + return [] + from dsagt.registry import CodeRegistry, SkillRegistry + + codes = CodeRegistry(runtime_dir=working_dir, kb=None) + reg = SkillRegistry(runtime_dir=working_dir, kb=None) + # Later entries win name collisions: codes first, then bundled + # skills, then project skills — a deliberately installed instruction + # skill outranks a registered code of the same name. + src_dirs = ( + codes.code_dirs() + reg._bundled_skill_dirs() + reg._project_skill_dirs() + ) + target = working_dir + for part in self.native_skills_dir.split("/"): + target = target / part + return _mirror_skills_to(target, src_dirs) + + def owned_artifacts(self, working_dir: Path) -> list[Path]: + """Files/dirs this agent's setup writes, for cleanup when a project + re-inits onto a *different* agent platform. + + Lists the instruction file, the per-agent MCP-config file(s), and the + agent's private per-project state dir(s) — NOT the shared + ``.agents/`` skill-mirror dir (managed by the manifest reaper), and + never project data (``.dsagt/``, ``kb_index/``, ``trace_archive/``, + ``skills/``). Paths may not all exist; the caller filters. + + Default = just the static marker; subclasses extend. + """ + return [working_dir / self.static_marker] - def env_overrides(self, config: dict) -> dict[str, str]: - """Phase-2 proxy-mode hook. Not called in BYOA. + def runtime_env(self, config: dict) -> dict[str, str]: + """Dsagt-owned env vars the agent process needs at runtime (BYOA). + + Default is empty: DSAGT no longer forces any telemetry env on the + agent (agent traces are recovered post-hoc from the on-disk + transcript, not by native OTel emission). Subclasses override + only to set per-project state-dir env (``CLINE_DIR``, + ``CODEX_HOME``) that isolates their global config per project. - Reserved for the proxy/OTel-redirect path that ``observability.py`` - will own when Phase 2 lands — the hook each agent uses to translate - ``config["llm"].*`` into provider env vars (after the proxy has - already rewritten the upstream URL). Default implementation - returns ``{}``; agents override with their own translation logic - as needed. + Provider credentials (ANTHROPIC_*, OPENAI_*, GOOSE_*) are the + user's responsibility — exported in their shell, never translated + from ``config["llm"]``. """ del config return {} - def proxy_env_overrides(self, proxy_port: int) -> dict[str, str]: - """Env vars that route the agent's LLM calls through dsagt-proxy. - - Default implementation covers every agent we support today — - we set the well-known base URLs (``ANTHROPIC_BASE_URL``, - ``OPENAI_BASE_URL``) for both wire protocols so an agent that - speaks either still hits the proxy, plus the embedding base - URL so MCP-server children inherit proxy routing for - ``litellm.embedding`` calls. Real upstream credentials live - only in the proxy subprocess; we plant the sentinel API key in - every standard slot so a direct call (bypassing the proxy) - fails loudly with 401 rather than silently leaking around our - observability pipeline. - - Subclasses may override only if an agent reads non-standard - env names. None do today. - """ - proxy_url = f"http://localhost:{proxy_port}" - return { - "ANTHROPIC_BASE_URL": proxy_url, - "OPENAI_BASE_URL": proxy_url, - "EMBEDDING_BASE_URL": proxy_url, - "ANTHROPIC_API_KEY": _PROXY_FORWARDED_SENTINEL, - "OPENAI_API_KEY": _PROXY_FORWARDED_SENTINEL, - "EMBEDDING_API_KEY": _PROXY_FORWARDED_SENTINEL, - } - def interactive_command(self, config: dict) -> list[str]: """Return the argv list for interactive launch. @@ -561,44 +402,6 @@ def interactive_command(self, config: dict) -> list[str]: del config return list(self.base_command) - #: Telemetry env vars this agent emits via its native OTel SDK. - #: Set by subclasses with ``otel_payload_support`` of "full" or - #: "partial". Empty for agents that don't emit OTel (cline / roo). - #: Used by ``byoa_env_hints`` to surface what the user should - #: export in their shell to get full visibility. - telemetry_env: ClassVar[dict[str, str]] = {} - - #: Per-agent provider-credential hints surfaced by ``byoa_env_hints``. - #: List of ``(env_var_name, hint)`` tuples shown to the user with a - #: "skip if already configured" note. - credential_hints: ClassVar[tuple[tuple[str, str], ...]] = () - - def byoa_env_hints( - self, mlflow_port: int, project: str, project_dir: Path - ) -> list[tuple[str, str]]: - """Provider-credential env vars the user should set in their shell. - - BYOA design: dsagt-internal env (project routing, MLflow URL, - OTel endpoint + headers, telemetry verbosity flags) lives in - the per-project launch shim — not the user's shell. This - method returns only the credentials the user owns, with one - hint string each. - """ - del mlflow_port, project, project_dir - return list(self.credential_hints) - - def launch_oneliner(self, project: str, project_dir: Path) -> str: - """Shell command to launch the agent interactively. - - ``cd `` puts both the agent and any ``dsagt-run`` children - in the project directory; readers find ``dsagt_config.yaml`` - there and use it as the single source of truth for project / - agent / mlflow routing — no DSAGT_* env vars to manage. - """ - del project - cmd = " ".join(shlex.quote(part) for part in self.interactive_command({})) - return f"cd {shlex.quote(str(project_dir))} && {cmd}" - @abstractmethod def run_script( self, @@ -618,38 +421,10 @@ def vscode_hint(self, project_dir: Path) -> list[str] | None: """One-or-two-line hint for users who run this agent as a VS Code extension instead of via the CLI. Returns ``None`` for agents without a working VS Code extension (most of them — only claude - and roo currently have extensions that auto-discover dsagt's + currently have extensions that auto-discover dsagt's per-project files from the workspace root). ``dsagt init`` prints these lines under "Or with VS Code extension". """ del project_dir return None - - # --- Phase 2 proxy-mode analogs --------------------------------------- - # ``__init__`` rebinds ``write_dynamic`` / ``run_script`` to these when - # ``proxy_port`` is set, so callers don't branch. Defaults defer to - # the BYOA implementation — agents whose proxy + BYOA paths are the - # same (claude, goose, roo's write_dynamic) inherit this fallthrough. - # Override in subclasses where proxy mode genuinely differs (cline - # auth -b proxy, codex [model_providers.dsagt-proxy] block, opencode - # baseURL override, cline/roo run_script un-punt). - - def proxy_write_dynamic( - self, - config: dict, - env: dict, - working_dir: Path, - pdir: Path, - ) -> list[str]: - return self.write_dynamic(config, env, working_dir, pdir) - - def proxy_run_script( - self, - config: dict, - env: dict, - working_dir: Path, - script_path: Path, - max_turns: int, - ) -> int: - return self.run_script(config, env, working_dir, script_path, max_turns) diff --git a/src/dsagt/agents/claude.py b/src/dsagt/agents/claude.py index f2a4094..f13a44f 100644 --- a/src/dsagt/agents/claude.py +++ b/src/dsagt/agents/claude.py @@ -2,42 +2,20 @@ Claude Code agent setup. Install: ``npm i -g @anthropic-ai/claude-code``. -Generates: ``.mcp.json``, ``CLAUDE.md``, ``.dsagt_env``. - -Post-proxy: the user brings ``ANTHROPIC_API_KEY`` (and optionally -``ANTHROPIC_MODEL``, ``ANTHROPIC_BASE_URL``) themselves and Claude Code -talks directly to its provider. We only inject OTel telemetry env vars -so the agent's LLM-call traces land in the project's MLflow. - -OTel support: **full** for native MLflow visibility (verified). Tool -args land on the ``claude_code.tool_result`` event when -``OTEL_LOG_TOOL_DETAILS=1`` and the assistant's response (with -tool_use blocks) lands on the ``api_response_body`` event when -``OTEL_LOG_RAW_API_BODIES=1``; both are gated. -``CLAUDE_CODE_ENHANCED_TELEMETRY_BETA=1`` enables the trace hierarchy. -All four flags are set in ``_CLAUDE_TELEMETRY_ENV`` below. Cited from -https://code.claude.com/docs/en/monitoring-usage.md and verified -end-to-end by the smoke-test harness once ``session.id`` was added to -``OTEL_RESOURCE_ATTRIBUTES``. - -Memory extraction: **does NOT work without --enable-proxy**, even -though traces are visible in the MLflow UI. Claude Code emits its -conversation via OTel *log events* (``api_response_body``, -``tool_result``), which is a different shape from the LiteLLM-autolog -``mlflow.spanInputs`` / ``mlflow.spanOutputs`` shape that -``memory.drain_session_traces`` reads. Users who want both visibility -*and* extraction should run ``dsagt start --enable-proxy``. See -``agents/__init__.py`` module docstring for the design rationale. - -Truncation caveats: ``tool_input`` truncates per-value at 512 chars and -total at ~4 KB; ``api_response_body`` is capped at 60 KB. Large diffs -or huge tool outputs may be clipped. +Generates: ``CLAUDE.md`` (instructions) and ``.mcp.json`` (MCP config). + +The user brings ``ANTHROPIC_API_KEY`` (and optionally ``ANTHROPIC_MODEL``, +``ANTHROPIC_BASE_URL``) themselves and Claude Code talks directly to its +provider. DSAGT sets **no** telemetry env on the agent — agent-side traces +are recovered post-hoc from Claude's on-disk transcript by DSAGT's own +serverless pipeline (MCP-server heartbeat → ``ClaudeReader`` → +``ClaudeTranslator`` → ``MLflowSink``), uniformly with every other agent; not +by forcing native OTel emission or wiring MLflow's autolog hook. Cache-marker injection: Claude Code handles Anthropic prompt caching -natively against the Anthropic API, so the (deleted) proxy's -``_inject_cache_breakpoints`` is unnecessary here. Users on a custom +natively against the Anthropic API. Users on a custom ``ANTHROPIC_BASE_URL`` that proxies to a non-Anthropic provider lose -caching either way. +caching. """ from __future__ import annotations @@ -55,64 +33,20 @@ _run_simple_script, ) -# Env vars Claude Code reads to gate full LLM-call telemetry. Without -# these, the OTel spans only carry counts/cost/duration. See -# https://code.claude.com/docs/en/monitoring-usage.md. -# -# OTEL_LOG_RAW_API_BODIES is intentionally omitted here — its value is -# per-project (a path like ``file:/api_bodies``) and gets rendered -# dynamically by ``_cmd_mlflow``. The ``=1`` (inline) mode would drop -# bodies to ``/v1/logs`` which MLflow's OTLP receiver returns 404 for; -# ``file:`` writes bodies to disk and stamps a ``body_ref`` on the -# span event, which travels via ``/v1/traces`` (the path MLflow accepts). -# -# OTEL_LOGS_EXPORTER is also dropped — MLflow has no logs endpoint, so -# pointing the SDK at /v1/logs only generates 404s in mlflow.log. -_CLAUDE_TELEMETRY_ENV: dict[str, str] = { - "CLAUDE_CODE_ENABLE_TELEMETRY": "1", - "CLAUDE_CODE_ENHANCED_TELEMETRY_BETA": "1", - "OTEL_LOG_TOOL_DETAILS": "1", - # Without this, user prompts show as "[REDACTED]" in span attributes. - "OTEL_LOG_USER_PROMPTS": "1", - "OTEL_TRACES_EXPORTER": "otlp", -} - class ClaudeSetup(AgentSetup): name = "claude" base_command = ["claude"] static_marker = "CLAUDE.md" + native_skills_dir = ".claude/skills" install_hint = "Install with `npm i -g @anthropic-ai/claude-code`." - # Anthropic-protocol native; cross-protocol routing requires the proxy. - credential_env_vars = ( - "ANTHROPIC_API_KEY", - "ANTHROPIC_BASE_URL", - "ANTHROPIC_MODEL", - ) - otel_payload_support = "full" - telemetry_env = _CLAUDE_TELEMETRY_ENV - credential_hints = ( - ("ANTHROPIC_API_KEY", "your Anthropic API key (skip if subscription-authed)"), - ("ANTHROPIC_BASE_URL", "optional gateway / proxy URL"), - ("ANTHROPIC_MODEL", "optional model override"), - ) - - def env_overrides(self, config: dict) -> dict[str, str]: - """Phase-2 proxy-mode hook: pin ``ANTHROPIC_MODEL`` to the - upstream-served name so claude doesn't fall back to its - built-in default. Only fires when ``config["proxy"]["port"]`` - is set (gated by ``agents/__init__.py:agent_env``). - - ``ANTHROPIC_BASE_URL`` and ``ANTHROPIC_API_KEY`` are set by - :meth:`proxy_env_overrides` (base default) — point at the - localhost proxy with the sentinel key. Claude posts - ``/v1/messages`` to the proxy regardless of upstream protocol; - the proxy translates. - """ - model = (config.get("llm") or {}).get("model") - if model and not str(model).startswith("${"): - return {"ANTHROPIC_MODEL": model} - return {} + + def owned_artifacts(self, working_dir: Path) -> list[Path]: + return [ + working_dir / "CLAUDE.md", + working_dir / ".mcp.json", + working_dir / ".claude", + ] def vscode_hint(self, project_dir: Path) -> list[str]: return [f"Open {project_dir} in VS Code and start the Claude extension."] @@ -137,55 +71,35 @@ def write_dynamic( working_dir: Path, pdir: Path, ) -> list[str]: - """Write ``.mcp.json`` and configure ``mlflow autolog claude``. - - The env block in ``.mcp.json`` carries DSAGT/MLflow/embedding - routing for the MCP-server children — claude inherits parent - env into them, but baking it into the JSON is robust against - shells that don't have those vars set. - - Also wires MLflow's first-class Claude Code integration via - ``.claude/settings.json``: a Stop hook that processes Claude's - transcript at session end and creates a rich MLflow trace with - full prompts, responses, and tool_use blocks. This is the only - way to get high-fidelity agent-side traces in BYOA mode (without - the proxy) — Claude's native OTel emission carries only thin - ``api_response_body`` log events that don't roundtrip through - memory extraction. + """Write ``.mcp.json``. + + The env block carries DSAGT/MLflow/embedding routing for the MCP-server + children — claude inherits parent env into them, but baking it into the + JSON is robust against shells that don't have those vars set. + + No trace wiring here: DSAGT's own serverless pipeline (the MCP-server + heartbeat → ``ClaudeReader`` → ``ClaudeTranslator`` → ``MLflowSink``) + produces Claude's traces, uniformly with every other agent — so we do + NOT also wire MLflow's ``autolog claude`` Stop hook (which would + double-log the same turns, and only Claude can use it serverlessly). """ del env, pdir actions: list[str] = [] env_block = _mcp_env_block(config) - mcp_config: dict = {"mcpServers": {}} - for server in ("registry", "knowledge"): - entry: dict = {"command": "uv", "args": _mcp_server_args(server)} - if env_block: - entry["env"] = env_block - mcp_config["mcpServers"][f"dsagt-{server}"] = entry + entry: dict = {"command": "uv", "args": _mcp_server_args()} + if env_block: + entry["env"] = env_block + mcp_config: dict = {"mcpServers": {"dsagt": entry}} mcp_path = working_dir / ".mcp.json" mcp_path.write_text(json.dumps(mcp_config, indent=2) + "\n") actions.append(f"Wrote {mcp_path}") - # Configure mlflow autolog claude — writes .claude/settings.json - # with the MLflow Stop hook + tracking env vars. Idempotent and - # preserves any existing keys in settings.json (mlflow's setup - # functions do a load → update → save, not a replace). - mlflow_port = (config.get("mlflow") or {}).get("port") - project_name = config.get("project") - if mlflow_port and project_name: - from mlflow.claude_code.config import setup_environment_config - from mlflow.claude_code.hooks import setup_hooks_config - - settings_file = working_dir / ".claude" / "settings.json" - setup_hooks_config(settings_file) - setup_environment_config( - settings_file, - tracking_uri=f"http://localhost:{mlflow_port}", - experiment_name=project_name, - ) - actions.append(f"Wrote {settings_file} (mlflow autolog claude)") + # Skills are mirrored into .claude/skills/ centrally via + # AgentSetup.setup_skills (driven by native_skills_dir) in + # dynamic_agent_record — see base.py. Picked up on the next Claude + # start, which is fine: this runs at init/start, before launch. return actions def run_script( diff --git a/src/dsagt/agents/cline.py b/src/dsagt/agents/cline.py index e6fbac5..2bdff83 100644 --- a/src/dsagt/agents/cline.py +++ b/src/dsagt/agents/cline.py @@ -2,36 +2,46 @@ Cline agent setup. Install: ``npm i -g cline``. -Generates: ``.clinerules/dsagt_instructions.md``; -``cline auth`` + ``cline mcp add`` run per-init in -:meth:`ClineSetup.write_dynamic` and write to ``$CLINE_DIR/data/``. - -**BYOA Phase 1 status: PUNTED for batch / smoke-test.** Cline's bundled -``lib.mjs`` ships a hardcoded anthropic-provider model whitelist -(``claude-haiku-4-5-20251001``, ``claude-sonnet-4-5-20250929``, etc. — -all without the ``-v1-project`` suffix lab gateways like PNNL require). -Even when ``cline auth -m claude-haiku-4-5-20251001-v1-project`` stores -our model name correctly in ``globalState.json``, at request time -cline's anthropic provider validates against the whitelist, sees the -unrecognized suffixed name, and silently substitutes its hardcoded -default (``claude-sonnet-4-5-20250929``) — which the gateway then 401s -because PNNL only serves the suffixed variant. - -The openai-native path has the same problem (whitelisted ``gpt-5.5``, -``gpt-4o`` etc.; PNNL serves ``-project``-suffixed variants). Bedrock -is locked behind ``cline auth``'s interactive setup flow, not the CLI. - -Phase 2 reactivates batch mode: ``dsagt-proxy`` aliases the cline- -substituted name back to the upstream's served name (same trick old_code -used for roo). Until then, ``cline.run_script`` raises ``RuntimeError`` -and the smoke test short-circuits in ``tests/smoke_test/run.sh``. +Generates: ``.clinerules/dsagt_instructions.md`` and +``.cline-data/cline_mcp_settings.json`` (hand-written per-project MCP +config, loaded via the ``CLINE_MCP_SETTINGS_PATH`` env var at ``dsagt +start``). Provider auth is cline's own (subscription login / ``cline +auth`` / the VS Code extension), configured by the user before dsagt — +dsagt never writes cline auth state, since doing so can clobber an +existing provider integration. + +**Batch / smoke-test status: NOT SUPPORTED (verified against cline +3.0.34).** Batch itself works (``cline "prompt"`` runs act-mode with +auto-approve, honoring the user's pre-configured auth — the old +model-whitelist blocker is moot), and headless cline even SPAWNS +registered MCP servers correctly (cwd = the session dir, full shell +env — verified with a spawn probe). But it never bridges their tools +into the model's toolset: a session asked to enumerate every tool it +has lists only cline built-ins + team tools, zero MCP entries. With +no dsagt tools reachable there is nothing for a smoke run to exercise, +so ``cline.run_script`` raises ``RuntimeError`` and the smoke test +short-circuits in ``tests/smoke_test/run.sh``. Re-probe on cline +upgrades (the TUI / VS Code paths do bridge MCP): if a headless +session can call an MCP tool, batch support is one small run_script +away. + +Auth + config scoping, verified 3.0.34: cline stores provider +credentials in ``providers.json`` beside its MCP settings, and BOTH +follow a ``--config`` / ``CLINE_DIR`` redirect — so whole-directory +isolation loses a subscription/OAuth login (instant Unauthorized). +The escape hatch is ``CLINE_MCP_SETTINGS_PATH``: it relocates ONLY the +MCP settings file (``providers.json`` resolves independently via its +own ``CLINE_PROVIDER_SETTINGS_PATH`` / data-dir default). dsagt sets +it in :meth:`ClineSetup.runtime_env`, giving per-project dsagt MCP +config + the user's global auth + zero footprint in the global +settings, all at once (verified: per-project server spawns, +subscription answered, global mcpServers unchanged). Interactive use is unaffected — ``dsagt init --agent cline`` still -writes the project state, and users running cline via VS Code (where -the UI exposes ``awsBedrockEndpoint`` etc.) can drive the project -manually. +writes the project state, and users running cline via VS Code can +drive the project manually. -OTel support: **none** (verified, ``otel_payload_support = "none"``). +OTel support: **none** (verified). Cline ships ``@opentelemetry/*`` packages but installs only a ``MeterProvider`` + ``LoggerProvider`` — never a ``TracerProvider``; zero spans are ever created (``OpenTelemetryClientProvider.ts``). Its @@ -42,27 +52,18 @@ extraction will see no agent-conversation traces; tool execution and KB observability still work via dsagt-run / MCP-server spans. -Workaround: ``dsagt start --enable-proxy`` makes every cline LLM call -visible in MLflow — each agent turn becomes an inspectable trace -(messages + assistant response + tool_use blocks) you can audit live or -replay. The dsagt-proxy interposes on cline's LLM calls, forwards them -to the user's upstream, and emits an OTel trace on cline's behalf. -Memory extraction is a downstream consequence — but the primary value -of the flag is being able to see what cline is doing at all. Adds one -subprocess per project session. See ``commands/proxy_server.py``. - -MCP config: hand-writing ``cline_mcp_settings.json`` is silently ignored; -the only path cline loads is via ``cline mcp add``, which writes a -stripped schema (no ``env`` block). We patch in the env block ourselves -after ``cline mcp add`` runs so the dsagt MCP-server children inherit -``MLFLOW_TRACKING_URI``, ``EMBEDDING_*``, ``OTEL_*`` etc. +MCP config: cline 3.x loads whatever file ``CLINE_MCP_SETTINGS_PATH`` +names (hand-written files work — the old "only ``cline mcp add`` is +loaded" behavior is gone), with the schema +``{"mcpServers": {: {"transport": {type, command, args, env}}}}``. +The env block rides in ``transport.env`` so dsagt MCP-server children +get ``MLFLOW_TRACKING_URI``, ``DSAGT_PROJECT_DIR``, ``EMBEDDING_*``. """ from __future__ import annotations import json import logging -import subprocess from pathlib import Path from .base import ( @@ -71,7 +72,6 @@ _DSAGT_MARKER, _load_master_instructions, _mcp_env_block, - _run_simple_script, ) logger = logging.getLogger(__name__) @@ -81,29 +81,22 @@ class ClineSetup(AgentSetup): name = "cline" base_command = ["cline"] static_marker = ".clinerules/dsagt_instructions.md" + # Cline skills are opt-in (Settings → Features → Enable Skills); mirroring + # is harmless if unused, and search_skills covers the disabled case. + native_skills_dir = ".cline/skills" install_hint = "Install with `npm i -g cline`." - otel_payload_support = "none" - # Cline's CLI nominally supports openai-native + anthropic, but cline - # auth's ``-b/--baseurl`` flag is openai-only and the openai-native - # path needs a non-standard model env var. Anthropic is the only - # path with consistent env conventions: ``ANTHROPIC_BASE_URL`` is read - # by cline's anthropic SDK at runtime, covering api.anthropic.com and - # gateway endpoints alike. Match roo's policy. - credential_env_vars = ( - "ANTHROPIC_API_KEY", "ANTHROPIC_BASE_URL", "ANTHROPIC_MODEL", - ) - # Cline emits no OTel — agent-side telemetry only via --proxy_traces. - telemetry_env = {} - credential_hints = ( - ("ANTHROPIC_API_KEY", "your provider API key (works for openai-shape " - "gateways too — cline's anthropic SDK reaches them via " - "ANTHROPIC_BASE_URL)"), - ("ANTHROPIC_BASE_URL", "gateway / proxy URL " - "(cline auth's -b flag is openai-only; the anthropic SDK reads " - "this env var at runtime)"), - ("ANTHROPIC_MODEL", "model name your gateway serves " - "(e.g. claude-haiku-4-5-20251001-v1-project)"), - ) + # Cline owns its provider auth (subscription login, `cline auth`, or the + # VS Code extension's settings) — BYOA means the user configures cline + # BEFORE pointing dsagt at it, and dsagt never touches auth state: + # running `cline auth` from scavenged env vars can clobber an existing + # provider integration (e.g. an OpenAI subscription login). + + def owned_artifacts(self, working_dir: Path) -> list[Path]: + return [ + working_dir / ".clinerules", + working_dir / ".cline-data", + working_dir / ".cline", + ] def write_static(self, working_dir: Path) -> list[str]: actions: list[str] = [] @@ -128,238 +121,65 @@ def write_dynamic( working_dir: Path, pdir: Path, ) -> list[str]: - """Three side-effects, all idempotent: - - 1. ``cline auth`` — populate per-project auth state from - ``ANTHROPIC_API_KEY`` / ``ANTHROPIC_MODEL`` in the user's - shell. Cline's per-project state dir is empty until ``cline - auth`` writes to it; the user's global cline auth doesn't - propagate into a fresh ``--config /.cline-data``. - 2. ``cline mcp add`` per server (skipped if entry already exists - — cline's mcp subcommand has no remove, only add). - 3. Patch the JSON cline wrote with our env block (cline doesn't - inherit parent env into MCP children, so MLFLOW_TRACKING_URI, - DSAGT_PROJECT_DIR, EMBEDDING_* must live in the JSON). - - Requires the ``cline`` binary to be installed at init time. + """Write ``/.cline-data/cline_mcp_settings.json`` directly. + + Cline 3.x resolves its MCP settings file via the + ``CLINE_MCP_SETTINGS_PATH`` env var (set by :meth:`runtime_env`), + independent of ``providers.json`` (auth) — so a hand-written + per-project file gives project-scoped dsagt MCP config while the + user's global auth and global settings stay untouched (verified + 3.0.34: per-project server spawns, subscription auth honored, + global mcpServers list unchanged). The env block rides inside + ``transport.env`` per cline's ``McpStdioTransportConfig`` schema. + + Provider auth is cline's own (subscription login / ``cline auth`` / + the VS Code extension) and must be configured before using dsagt — + dsagt never writes cline auth state (see the class comment). + + Idempotent; preserves any non-dsagt entries the user added to the + per-project file. No cline binary needed at init. """ - del pdir - actions: list[str] = [] - cline_dir = str(working_dir / ".cline-data") - Path(cline_dir).mkdir(parents=True, exist_ok=True) + del env, pdir + settings_path = working_dir / ".cline-data" / "cline_mcp_settings.json" + settings_path.parent.mkdir(parents=True, exist_ok=True) - # 1. Configure per-project cline auth from shell env. - api_key = env.get("ANTHROPIC_API_KEY") - model = env.get("ANTHROPIC_MODEL") - if not api_key or not model: - raise RuntimeError( - "cline batch mode requires ANTHROPIC_API_KEY and " - "ANTHROPIC_MODEL in the shell env (and ANTHROPIC_BASE_URL " - "if you're on a gateway). Cline's anthropic SDK reads " - "ANTHROPIC_BASE_URL at runtime, so this single config " - "covers api.anthropic.com and lab gateways alike. For " - "openai-shape gateways like PNNL, alias " - "ANTHROPIC_API_KEY=$OPENAI_API_KEY — most lab gateways " - "serve both wire protocols on the same key." - ) - auth_cmd = [ - "cline", "auth", - "--config", cline_dir, - "-p", "anthropic", - "-k", api_key, - "-m", model, - ] - result = subprocess.run( - auth_cmd, cwd=str(working_dir), - capture_output=True, text=True, - ) - if result.returncode != 0: - detail = (result.stderr or result.stdout).strip() - raise RuntimeError( - f"cline auth failed (exit {result.returncode}): {detail}" - ) - actions.append(f"Configured cline auth at {cline_dir}") - - # Cline's mcp subcommand only has ``add`` (no ``remove``), and ``add`` - # errors if the server already exists. Detect pre-existing entries - # and skip the add so this method is idempotent (called by both - # ``dsagt init`` and ``_cmd_start`` → ``dynamic_agent_record``). - # The env-block patch below runs unconditionally to keep on-disk - # routing in sync if mlflow port / paths ever changed. - mcp_path = Path(cline_dir) / "data" / "settings" / "cline_mcp_settings.json" - existing: set[str] = set() - if mcp_path.exists(): + settings: dict = {"mcpServers": {}} + if settings_path.exists(): try: - existing = set(json.loads(mcp_path.read_text()) - .get("mcpServers", {}).keys()) + settings = json.loads(settings_path.read_text()) + settings.setdefault("mcpServers", {}) except (json.JSONDecodeError, OSError): - existing = set() - - for server in ("registry", "knowledge"): - name = f"dsagt-{server}" - if name in existing: - continue - add_cmd = [ - "cline", "mcp", "add", - "--config", cline_dir, - name, - "--", - "uv", "run", f"dsagt-{server}-server", - ] - result = subprocess.run( - add_cmd, cwd=str(working_dir), - capture_output=True, text=True, - ) - if result.returncode != 0: - detail = (result.stderr or result.stdout).strip() - raise RuntimeError( - f"cline mcp add {name} failed " - f"(exit {result.returncode}): {detail}" - ) + settings = {"mcpServers": {}} - mcp_path = Path(cline_dir) / "data" / "settings" / "cline_mcp_settings.json" - if not mcp_path.exists(): - raise RuntimeError( - f"cline mcp add succeeded but {mcp_path} was not created — " - "cline may have changed its config layout." - ) - settings = json.loads(mcp_path.read_text()) - mcp_env = _mcp_env_block(config) - for entry in settings.get("mcpServers", {}).values(): - entry["env"] = mcp_env - mcp_path.write_text(json.dumps(settings, indent=2) + "\n") - actions.append( - f"Registered MCP servers and patched {len(mcp_env)} env vars into {mcp_path}" - ) - - # ``launch_oneliner`` (used by ``dsagt init`` printout) tells - # the user to run ``cline -v -y -a --config /.cline-data`` - # — the per-project state dir we just populated above. - return actions - - # --- Phase 2 proxy-mode write_dynamic -------------------------------- - # Differs from BYOA only in the auth call: ``cline auth -p openai -b - # http://localhost: -k -m `` instead of - # ``-p anthropic``. The MCP-add + JSON-patch steps are identical, so - # we extract them into a helper. - - def _patch_mcp_servers( - self, config: dict, working_dir: Path, cline_dir: str, - ) -> str: - """Idempotent ``cline mcp add`` + JSON env-block patch. - Returns the action string for the printout. - """ - mcp_path = Path(cline_dir) / "data" / "settings" / "cline_mcp_settings.json" - existing: set[str] = set() - if mcp_path.exists(): - try: - existing = set(json.loads(mcp_path.read_text()) - .get("mcpServers", {}).keys()) - except (json.JSONDecodeError, OSError): - existing = set() - - for server in ("registry", "knowledge"): - name = f"dsagt-{server}" - if name in existing: - continue - add_cmd = [ - "cline", "mcp", "add", "--config", cline_dir, name, - "--", "uv", "run", f"dsagt-{server}-server", - ] - result = subprocess.run( - add_cmd, cwd=str(working_dir), - capture_output=True, text=True, - ) - if result.returncode != 0: - detail = (result.stderr or result.stdout).strip() - raise RuntimeError( - f"cline mcp add {name} failed " - f"(exit {result.returncode}): {detail}" - ) - - if not mcp_path.exists(): - raise RuntimeError( - f"cline mcp add succeeded but {mcp_path} was not created — " - "cline may have changed its config layout." - ) - settings = json.loads(mcp_path.read_text()) mcp_env = _mcp_env_block(config) - for entry in settings.get("mcpServers", {}).values(): - entry["env"] = mcp_env - mcp_path.write_text(json.dumps(settings, indent=2) + "\n") - return ( - f"Registered MCP servers and patched {len(mcp_env)} " - f"env vars into {mcp_path}" - ) - - def proxy_write_dynamic( - self, - config: dict, - env: dict, - working_dir: Path, - pdir: Path, - ) -> list[str]: - """Phase-2 proxy-mode setup. Cline auths against the - localhost dsagt-proxy with the sentinel API key (real upstream - creds live only in the proxy subprocess). Provider is forced - to ``openai`` because cline auth's ``-b/--baseurl`` flag is - openai-only — the proxy speaks /chat/completions on every - upstream regardless of what the agent emits, so this is fine. - """ - del pdir - from .base import _PROXY_FORWARDED_SENTINEL - actions: list[str] = [] - cline_dir = str(working_dir / ".cline-data") - Path(cline_dir).mkdir(parents=True, exist_ok=True) - - proxy_port = (config.get("proxy") or {}).get("port") - model = (config.get("llm") or {}).get("model") - if not proxy_port or not model: - raise RuntimeError( - "cline proxy-mode write_dynamic requires " - "config['proxy']['port'] and config['llm']['model']." - ) - auth_cmd = [ - "cline", "auth", - "--config", cline_dir, - "-p", "openai", - "-k", _PROXY_FORWARDED_SENTINEL, - "-m", model, - "-b", f"http://localhost:{proxy_port}", + transport: dict = { + "type": "stdio", + "command": "uv", + "args": ["run", "dsagt-server"], + } + if mcp_env: + transport["env"] = mcp_env + settings["mcpServers"]["dsagt"] = {"transport": transport} + settings_path.write_text(json.dumps(settings, indent=2) + "\n") + return [ + f"Wrote {settings_path} ({len(mcp_env)} MCP env vars; loaded via " + "CLINE_MCP_SETTINGS_PATH at dsagt start)" ] - result = subprocess.run( - auth_cmd, env=env, cwd=str(working_dir), - capture_output=True, text=True, - ) - if result.returncode != 0: - detail = (result.stderr or result.stdout).strip() - raise RuntimeError( - f"cline auth (proxy) failed (exit {result.returncode}): {detail}" - ) - actions.append(f"Configured cline auth at {cline_dir} (proxy mode)") - actions.append(self._patch_mcp_servers(config, working_dir, cline_dir)) - return actions def runtime_env(self, config: dict) -> dict[str, str]: - """BYOA infrastructure: per-project ``CLINE_DIR`` (state dir) - plus inherited telemetry flags. Isolates per-project MCP / - auth state from the global ``~/.cline-data``. + """BYOA infrastructure: point cline at the per-project MCP settings. + + ``CLINE_MCP_SETTINGS_PATH`` relocates ONLY the MCP settings file — + ``providers.json`` (auth) resolves independently, so the user's + subscription/OAuth login keeps working and no dsagt entry ever + lands in the global settings. """ env = super().runtime_env(config) - env["CLINE_DIR"] = str(Path(config["project_dir"]) / ".cline-data") + env["CLINE_MCP_SETTINGS_PATH"] = str( + Path(config["project_dir"]) / ".cline-data" / "cline_mcp_settings.json" + ) return env - def launch_oneliner(self, project: str, project_dir: Path) -> str: - """Cline needs ``--config /.cline-data`` so it picks up - the MCP-server registrations + auth state ``write_dynamic`` - wrote. Also ``-v -y -a`` for verbose / auto-yes / act-mode. - """ - del project - import shlex - pdir = shlex.quote(str(project_dir)) - cline_dir = shlex.quote(str(project_dir / ".cline-data")) - return f"cd {pdir} && cline -v -y -a --config {cline_dir}" - def run_script( self, config: dict, @@ -368,49 +188,16 @@ def run_script( script_path: Path, max_turns: int, ) -> int: - """Punted in BYOA Phase 1 — see module docstring. + """Not supported for cline — see module docstring. - Cline's bundled anthropic provider hardcodes a model-ID whitelist - that doesn't include lab-gateway aliases (``-v1-project`` etc.), - so the request silently substitutes a default cline thinks it - knows and the gateway 401s. Phase 2 ``dsagt-proxy`` fixes this - by aliasing names back to the upstream's served IDs. + Cline 3.x batch runs fine (act mode, pre-configured auth), but + its headless CLI never loads MCP servers, so a scripted session + has no dsagt tools to exercise. """ del config, env, working_dir, script_path, max_turns raise RuntimeError( - "cline batch mode is punted in Phase 1 BYOA — cline's anthropic " - "provider rewrites unrecognized model names (lab-gateway " - "aliases like ``-v1-project``) to its hardcoded default, " - "which the gateway then rejects. Phase 2 reactivates this " - "via dsagt-proxy aliasing. See agents/cline.py module " - "docstring." - ) - - def proxy_run_script( - self, - config: dict, - env: dict, - working_dir: Path, - script_path: Path, - max_turns: int, - ) -> int: - """Phase-2 proxy-mode batch run. Un-punts cline: - ``proxy_write_dynamic`` already configured cline auth against - the localhost proxy + registered our MCP servers, so we just - invoke cline pointing at the per-project state dir. The proxy - translates lab-gateway-aliased model names back to upstream- - served names, sidestepping cline's hardcoded whitelist. - """ - del max_turns - text = script_path.read_text().strip() - if not text: - return 1 - cline_dir = env.get("CLINE_DIR") or str( - Path(config["project_dir"]) / ".cline-data" + "cline batch mode is not supported — cline's headless CLI " + "(verified 3.0.34) does not load MCP servers, so a scripted " + "session has no dsagt tools. See agents/cline.py module " + "docstring; re-probe on cline upgrades." ) - # ``-a`` (act mode) — without it cline defaults to its task UI - # mode whose model fallback is ``gpt-5.5`` regardless of what - # ``cline auth`` configured for plan/act modes. Model comes - # from the per-project auth state ``proxy_write_dynamic`` populated. - cmd = ["cline", "-v", "-y", "-a", "--config", cline_dir, text] - return _run_simple_script(cmd, env, working_dir, self.install_hint) diff --git a/src/dsagt/agents/codex.py b/src/dsagt/agents/codex.py index 52ef5b3..5c293ab 100644 --- a/src/dsagt/agents/codex.py +++ b/src/dsagt/agents/codex.py @@ -10,14 +10,13 @@ servers). We point ``CODEX_HOME`` at ``/.codex-data`` to keep state isolated per project, and :meth:`CodexSetup.write_dynamic` writes ``[mcp_servers.*]`` blocks with explicit env (codex doesn't -inherit parent env into MCP children, same as cline/roo). +inherit parent env into MCP children, same as cline). -Post-proxy: the user owns model + provider config in -``$CODEX_HOME/config.toml`` (or via ``OPENAI_API_KEY`` / -``OPENAI_BASE_URL`` env). We don't write a model_providers block -anymore. +The user owns model + provider config in ``$CODEX_HOME/config.toml`` +(or via ``OPENAI_API_KEY`` / ``OPENAI_BASE_URL`` env). We don't write +a model_providers block. -OTel support: **partial** (verified, ``otel_payload_support = "partial"``). +OTel support: **partial** (verified). Codex's ``codex-otel`` Rust crate emits OTel spans/logs/metrics, BUT: * LLM-call spans (``stream_request``, ``handle_responses``) carry @@ -34,12 +33,10 @@ * Tool *results* go to ``codex.tool_result`` log events with full args + output (``session_telemetry.rs:962-1000``). -Conversation history *is* recoverable from +Conversation history is recovered from ``$CODEX_HOME/sessions/rollout--.jsonl`` (full assistant text -+ tool calls + responses). Reading that side-channel into -``extract_session`` is a TODO — not yet wired. Until then, memory -extraction will see only token counts and tool names from Codex -turns. ++ tool calls + responses) by the trace pipeline's Codex reader/translator +on the heartbeat — feeding MLflow and episodic memory like every other agent. Open Codex issues tracking richer OTel: openai/codex#12913, #10277, #6153, #16248. @@ -64,107 +61,54 @@ def _render_codex_config(mcp_env: dict) -> str: """Render the per-project ``$CODEX_HOME/config.toml`` body. - Emits ``[mcp_servers.*]`` sections plus an ``[otel]`` block. - - No top-level keys are emitted, so the output can be safely appended - to a copy of the user's ``~/.codex/config.toml`` without colliding - on top-level keys like ``model`` or ``approval_policy``. Batch-mode - approval / sandbox are set on the codex CLI directly + Emits only ``[mcp_servers.*]`` sections. No top-level keys are + emitted, so the output can be safely appended to a copy of the user's + ``~/.codex/config.toml`` without colliding on top-level keys like + ``model`` or ``approval_policy``. Batch-mode approval / sandbox are + set on the codex CLI directly (``--dangerously-bypass-approvals-and-sandbox``) instead of here. - The ``[otel]`` block opts codex's partial OTel into MLflow: - - * ``trace_exporter = "otlp-http"`` selects the OTLP-over-HTTP - exporter, which then reads ``OTEL_EXPORTER_OTLP_ENDPOINT`` / - ``OTEL_EXPORTER_OTLP_HEADERS`` from env (set by ``agent_env`` - in batch mode; the user sets them in their shell per the - ``dsagt init`` printout for interactive mode). Without this - key codex's traces never leave the process. - * ``log_user_prompt = true`` flips the ``codex.user_prompt`` log - event from REDACTED to full text, which is what memory - extraction needs to see what the user asked for. - - Codex's LLM-call spans still carry only token counts + tool names - (no request messages, no response bodies) — that's a codex - bundle limitation, not something we can override. Tool *results* - do land in ``codex.tool_result`` log events with full args + output. + No ``[otel]`` block: DSAGT no longer forces codex's native telemetry + (nor the ``log_user_prompt`` privacy override). Codex's conversation + history is recovered post-hoc from its on-disk session rollout. """ lines: list[str] = [] - for server in ("registry", "knowledge"): - lines.append(f"[mcp_servers.dsagt-{server}]") - lines.append('command = "uv"') - args = _mcp_server_args(server) - args_toml = ", ".join(_toml_quote(a) for a in args) - lines.append(f"args = [{args_toml}]") - if mcp_env: - lines.append(f"[mcp_servers.dsagt-{server}.env]") - for k, v in mcp_env.items(): - lines.append(f"{k} = {_toml_quote(v)}") - lines.append("") - lines.extend([ - "[otel]", - 'trace_exporter = "otlp-http"', - "log_user_prompt = true", - "", - ]) + lines.append("[mcp_servers.dsagt]") + lines.append('command = "uv"') + args = _mcp_server_args() + args_toml = ", ".join(_toml_quote(a) for a in args) + lines.append(f"args = [{args_toml}]") + if mcp_env: + lines.append("[mcp_servers.dsagt.env]") + for k, v in mcp_env.items(): + lines.append(f"{k} = {_toml_quote(v)}") + lines.append("") return "\n".join(lines) -def _render_codex_config_proxy(mcp_env: dict, proxy_port: int, model: str) -> str: - """Phase-2 proxy-mode config.toml body. - - Adds, on top of the BYOA MCP-server registrations: - * ``model_provider = "dsagt-proxy"`` (top-level, picks our provider) - * ``[model_providers.dsagt-proxy]`` block with ``base_url`` pointing - at the localhost proxy and ``wire_api = "chat"`` (the proxy - speaks /chat/completions on every upstream). - * Top-level ``model = `` so the user doesn't - need to pass ``-m`` on every ``codex exec`` call. - - These keys go ABOVE the ``[mcp_servers.*]`` sections that - ``_render_codex_config`` produces — top-level keys must precede - any table headers in TOML. - """ - base = _render_codex_config(mcp_env) - header = "\n".join([ - f'model = "{model}"', - 'model_provider = "dsagt-proxy"', - '', - '[model_providers.dsagt-proxy]', - 'name = "DSAGT Proxy"', - f'base_url = "http://localhost:{proxy_port}/v1"', - 'wire_api = "chat"', - '', - ]) - return header + base - - class CodexSetup(AgentSetup): name = "codex" base_command = ["codex"] static_marker = "AGENTS.md" + # Project-local .agents/skills (repo-root, codex-discovered) — never the + # global ~/.agents/skills or ~/.codex; manifest-tracked, user skills safe. + native_skills_dir = ".agents/skills" install_hint = ( - "Install with `npm i -g @openai/codex` or " - "`brew install --cask codex`." - ) - otel_payload_support = "partial" - # Codex is openai-protocol native. - credential_env_vars = ("OPENAI_API_KEY", "OPENAI_BASE_URL") - # Codex's OTel is config-file driven, not env; spans carry no - # message payload anyway, so we don't surface OTel hints here. - telemetry_env = {} - credential_hints = ( - ("OPENAI_API_KEY", "your OpenAI API key (skip if subscription-authed)"), - ("OPENAI_BASE_URL", "optional gateway / proxy URL"), + "Install with `npm i -g @openai/codex` or " "`brew install --cask codex`." ) + def owned_artifacts(self, working_dir: Path) -> list[Path]: + return [working_dir / "AGENTS.md", working_dir / ".codex-data"] + def write_static(self, working_dir: Path) -> list[str]: actions: list[str] = [] (working_dir / ".codex-data").mkdir(parents=True, exist_ok=True) instructions = _load_master_instructions() if instructions: action = _append_or_write( - working_dir / "AGENTS.md", instructions, _DSAGT_MARKER, + working_dir / "AGENTS.md", + instructions, + _DSAGT_MARKER, ) if action: actions.append(action) @@ -198,6 +142,7 @@ def write_dynamic( """ del env, pdir import shutil + actions: list[str] = [] codex_home = Path(working_dir) / ".codex-data" codex_home.mkdir(parents=True, exist_ok=True) @@ -220,79 +165,18 @@ def write_dynamic( base_toml = user_config.read_text() if user_config.exists() else "" mcp_env = _mcp_env_block(config) body = _render_codex_config(mcp_env) - merged = (base_toml.rstrip() + "\n\n" + body if base_toml else body) + merged = base_toml.rstrip() + "\n\n" + body if base_toml else body config_path.write_text(merged + "\n") actions.append(f"Wrote {config_path} ({len(mcp_env)} MCP env vars)") return actions - def launch_oneliner(self, project: str, project_dir: Path) -> str: - """Codex reads MCP servers from ``$CODEX_HOME/config.toml`` and - has no ``--config`` flag — must inline ``CODEX_HOME`` pointing - at our per-project ``.codex-data/``, otherwise codex reads - ``~/.codex`` and our MCP servers don't register. - """ - del project - import shlex - pdir = shlex.quote(str(project_dir)) - codex_home = shlex.quote(str(project_dir / ".codex-data")) - return f"cd {pdir} && CODEX_HOME={codex_home} codex" - - def proxy_write_dynamic( - self, - config: dict, - env: dict, - working_dir: Path, - pdir: Path, - ) -> list[str]: - """Phase-2 proxy-mode setup. Same MCP / auth / user-config - copy as BYOA, but renders config.toml with the - ``[model_providers.dsagt-proxy]`` block so codex routes its - LLM calls through our localhost proxy. No subscription auth - needed in proxy mode — the proxy authenticates upstream with - ``LLM_API_KEY``; we plant a sentinel API key in the codex env - so direct calls bypassing the proxy fail loudly. - """ - del env, pdir - import shutil - actions: list[str] = [] - codex_home = Path(working_dir) / ".codex-data" - codex_home.mkdir(parents=True, exist_ok=True) - user_codex = Path.home() / ".codex" - - # Still copy user's auth.json + config.toml as a base — they may - # have other prefs we want to preserve. Proxy routing layered on - # top via the _proxy renderer. - user_auth = user_codex / "auth.json" - if user_auth.exists(): - dest_auth = codex_home / "auth.json" - shutil.copy2(user_auth, dest_auth) - dest_auth.chmod(0o600) - actions.append(f"Copied {user_auth} → {dest_auth}") - - proxy_port = (config.get("proxy") or {}).get("port") - model = (config.get("llm") or {}).get("model") - if not proxy_port or not model: - raise RuntimeError( - "codex proxy_write_dynamic requires " - "config['proxy']['port'] and config['llm']['model']." - ) - - config_path = codex_home / "config.toml" - user_config = user_codex / "config.toml" - base_toml = user_config.read_text() if user_config.exists() else "" - mcp_env = _mcp_env_block(config) - body = _render_codex_config_proxy(mcp_env, proxy_port, model) - merged = (base_toml.rstrip() + "\n\n" + body if base_toml else body) - config_path.write_text(merged + "\n") - actions.append( - f"Wrote {config_path} ({len(mcp_env)} MCP env vars, proxy mode)" - ) - return actions - def runtime_env(self, config: dict) -> dict[str, str]: - """BYOA infrastructure: per-project ``CODEX_HOME`` (state dir) - plus inherited telemetry flags. Isolates per-project MCP / - config.toml from the global ``~/.codex``. + """BYOA infrastructure: per-project ``CODEX_HOME`` (state dir). + + Isolates per-project MCP / config.toml from the global + ``~/.codex``. Codex has no ``--config`` flag, so the agent must + be launched with this ``CODEX_HOME`` set for our MCP servers to + register. """ env = super().runtime_env(config) env["CODEX_HOME"] = str(Path(config["project_dir"]) / ".codex-data") @@ -312,10 +196,12 @@ def run_script( if not text: return 1 cmd = [ - "codex", "exec", + "codex", + "exec", "--dangerously-bypass-approvals-and-sandbox", "--skip-git-repo-check", - "-C", str(working_dir), + "-C", + str(working_dir), text, ] return _run_simple_script(cmd, env, working_dir, self.install_hint) diff --git a/src/dsagt/agents/goose.py b/src/dsagt/agents/goose.py index 71f5422..1317524 100644 --- a/src/dsagt/agents/goose.py +++ b/src/dsagt/agents/goose.py @@ -2,33 +2,25 @@ Goose agent setup. Install: see https://github.com/block/goose. -Generates: ``goose.yaml``, ``.goosehints``, ``.dsagt_env``. - -Post-proxy: Goose talks directly to the user's provider configured via -its own ``~/.config/goose/config.yaml`` or ``GOOSE_PROVIDER`` / -``GOOSE_MODEL`` env. We only inject the OTel endpoint env so Goose's -native ``dispatch_tool_call`` span (with full ``{tool, arguments}`` JSON -in ``input``) lands in the project's MLflow. - -OTel support: **full** for native MLflow visibility (verified by -source review of ``crates/goose/src/agents/agent.rs:572-585``). Goose -creates a ``dispatch_tool_call`` span with attribute -``input = {"tool": , "arguments": {...}}`` and ``session.id``, -activating automatically when ``OTEL_EXPORTER_OTLP_ENDPOINT`` is set -(handled by ``agents/__init__.py:agent_env``). No payload-stripping -flag exists — args are exported verbatim. - -Memory extraction: **does NOT work without --enable-proxy**, even -though Goose's traces are visible in the MLflow UI. The -``dispatch_tool_call`` span is in Goose's domain schema, not the -LiteLLM-autolog ``mlflow.spanInputs`` / ``mlflow.spanOutputs`` shape -that ``memory.drain_session_traces`` reads. Users who want both -visibility *and* extraction should run ``dsagt start --enable-proxy``. -See ``agents/__init__.py`` module docstring for the design rationale. - -Caveat: tool *outputs* are NOT in the span (the ``output`` field is -declared but never written). ``dsagt-run``'s ``tool.execute`` spans -cover the execution layer with stdout/stderr/exit-code. +Generates: ``goose.yaml``, ``.goosehints``. + +BYOA: Goose talks directly to the user's provider via its own +``~/.config/goose/config.yaml`` or ``GOOSE_PROVIDER`` / ``GOOSE_MODEL`` env; +DSAGT sets no telemetry env. + +Telemetry / episodic memory: Goose has **no Stop/turn hook and no MLflow +autolog integration**, so DSAGT does not offer the mlflow-autolog or +episodic-memory options for goose — those are gated on agents that have both +(claude / codex / opencode). Goose stays fully supported for the core, +agent-agnostic capabilities (KB retrieval, registered tools, skills, +tool-execution provenance via ``dsagt-run``). If goose gains a hook +mechanism the options can be enabled. + +Gateway note: goose's openai/anthropic providers read ``OPENAI_HOST`` / +``ANTHROPIC_HOST`` for the base URL (a goose-specific naming convention from +its Rust client), NOT the standard ``OPENAI_BASE_URL`` / ``ANTHROPIC_BASE_URL`` +everything else uses. Without HOST set, goose ignores BASE_URL and hits the +provider's default endpoint — silently, for users on a lab gateway. """ from __future__ import annotations @@ -51,40 +43,17 @@ class GooseSetup(AgentSetup): name = "goose" base_command = ["goose", "session"] static_marker = ".goosehints" + native_skills_dir = ".agents/skills" # cross-agent standard goose discovers install_hint = "See https://github.com/block/goose for installation." - otel_payload_support = "full" - # Multi-protocol; goose's runtime reads provider-specific creds plus - # its own GOOSE_PROVIDER / GOOSE_MODEL routing selectors. - # Goose's openai/anthropic providers read ``OPENAI_HOST`` / - # ``ANTHROPIC_HOST`` for the base URL (a goose-specific naming - # convention from its Rust client), NOT the standard - # ``OPENAI_BASE_URL`` / ``ANTHROPIC_BASE_URL`` everything else uses. - # Without HOST set, goose ignores BASE_URL and hits the provider's - # default endpoint — silently for users with a lab gateway. - credential_env_vars = ( - "GOOSE_PROVIDER", "GOOSE_MODEL", - "ANTHROPIC_API_KEY", "ANTHROPIC_BASE_URL", "ANTHROPIC_HOST", - "ANTHROPIC_MODEL", - "OPENAI_API_KEY", "OPENAI_BASE_URL", "OPENAI_HOST", - ) - # Goose's Rust client emits OTel automatically when - # OTEL_EXPORTER_OTLP_ENDPOINT is set — no per-platform flags needed. - telemetry_env = {} - credential_hints = ( - ("GOOSE_PROVIDER", "anthropic, openai, etc. (skip if global ~/.config/goose configured)"), - ("GOOSE_MODEL", "the model name your provider serves"), - ("ANTHROPIC_API_KEY", "if GOOSE_PROVIDER=anthropic"), - ("ANTHROPIC_HOST", "if GOOSE_PROVIDER=anthropic and on a gateway / proxy"), - ("OPENAI_API_KEY", "if GOOSE_PROVIDER=openai"), - ("OPENAI_HOST", "if GOOSE_PROVIDER=openai and on a gateway / proxy (NOT OPENAI_BASE_URL — goose ignores that)"), - ) def write_static(self, working_dir: Path) -> list[str]: actions: list[str] = [] instructions = _load_master_instructions() if instructions: action = _append_or_write( - working_dir / ".goosehints", instructions, _DSAGT_MARKER, + working_dir / ".goosehints", + instructions, + _DSAGT_MARKER, ) if action: actions.append(action) @@ -97,22 +66,23 @@ def write_dynamic( working_dir: Path, pdir: Path, ) -> list[str]: - """Write ``goose.yaml``. Goose inherits parent env into MCP - children, so extension entries don't need an explicit env list. - """ + """Write ``goose.yaml``. Goose inherits parent env into MCP children, + so extension entries don't need an explicit env list.""" del config, env, pdir actions: list[str] = [] - goose_config: dict = {"extensions": {}} - for server in ("registry", "knowledge"): - args = _mcp_server_args(server) - goose_config["extensions"][server] = { - "enabled": True, - "name": server, - "type": "stdio", - "cmd": "uv " + " ".join(args), - "timeout": 300, + args = _mcp_server_args() + goose_config: dict = { + "extensions": { + "dsagt": { + "enabled": True, + "name": "dsagt", + "type": "stdio", + "cmd": "uv " + " ".join(args), + "timeout": 300, + } } + } goose_path = working_dir / "goose.yaml" goose_path.write_text( @@ -121,47 +91,14 @@ def write_dynamic( actions.append(f"Wrote {goose_path}") return actions - def env_overrides(self, config: dict) -> dict[str, str]: - """Phase-2 proxy-mode hook: pin ``GOOSE_PROVIDER`` and - ``GOOSE_MODEL`` so a user's global ``~/.config/goose/config.yaml`` - doesn't override the project's configured model. - - Provider is forced to ``openai`` because the proxy speaks - /chat/completions (openai wire protocol) regardless of upstream; - goose's openai client posts there. Model is the upstream-served - name from ``config["llm"]["model"]``. - """ - model = (config.get("llm") or {}).get("model") - out: dict[str, str] = {"GOOSE_PROVIDER": "openai"} - if model and not str(model).startswith("${"): - out["GOOSE_MODEL"] = model - return out - - def proxy_env_overrides(self, proxy_port: int) -> dict[str, str]: - """Goose-specific proxy routing: add OPENAI_HOST / ANTHROPIC_HOST. - - The base-class default sets ``OPENAI_BASE_URL`` / ``ANTHROPIC_BASE_URL`` - which most agents read; goose ignores those and reads HOST. - Without this override, ``--enable-proxy`` would silently leave - goose talking to the upstream provider directly — same root - cause as the non-proxy path's HOST mapping. - """ - env = super().proxy_env_overrides(proxy_port) - proxy_url = f"http://localhost:{proxy_port}" - env["OPENAI_HOST"] = proxy_url - env["ANTHROPIC_HOST"] = proxy_url - return env - def interactive_command(self, config: dict) -> list[str]: - """Goose only reads ``~/.config/goose/config.yaml`` for extensions, not - a project-local file — so the MCP servers are passed via - ``--with-extension`` flags on the session command to guarantee they - attach for this project. + """Goose reads ``~/.config/goose/config.yaml`` for extensions, not a + project-local file — so the dsagt MCP server is passed via + ``--with-extension`` on the session command to attach for this project. """ del config cmd = list(self.base_command) - for server in ("registry", "knowledge"): - cmd.extend(["--with-extension", f"uv run dsagt-{server}-server"]) + cmd.extend(["--with-extension", "uv run dsagt-server"]) return cmd def run_script( @@ -175,8 +112,13 @@ def run_script( """Single ``goose run`` call — goose's instructions file IS multi-turn.""" del config env["GOOSE_MODE"] = "auto" - cmd = ["goose", "run", "--instructions", str(script_path), - "--max-turns", str(max_turns)] - for server in ("registry", "knowledge"): - cmd.extend(["--with-extension", f"uv run dsagt-{server}-server"]) + cmd = [ + "goose", + "run", + "--instructions", + str(script_path), + "--max-turns", + str(max_turns), + ] + cmd.extend(["--with-extension", "uv run dsagt-server"]) return _run_simple_script(cmd, env, working_dir, self.install_hint) diff --git a/src/dsagt/agents/opencode.py b/src/dsagt/agents/opencode.py index 34b6bf2..072e552 100644 --- a/src/dsagt/agents/opencode.py +++ b/src/dsagt/agents/opencode.py @@ -9,7 +9,7 @@ OTel support: **none** natively (one third-party plugin exists but isn't wired in by default). Agent transparency in MLflow is limited to MCP-server spans (kb.*, registry.*) and dsagt-run tool.execute spans; LLM-call payloads -require Phase 2 proxy mode. +are not captured. Auth model: opencode reads creds via ``{env:VAR}`` interpolation in its ``opencode.json`` provider block, so we can keep BYOA-pure — the file @@ -23,7 +23,7 @@ Model whitelist: pass-through for known providers (pulled from models.dev) and fully user-controlled for custom providers via ``provider..models``. -Unlike cline/roo, opencode does NOT rewrite gateway-aliased model names. +Unlike cline, opencode does NOT rewrite gateway-aliased model names. Batch mode: ``opencode run --dir --dangerously-skip-permissions -m ``. ``--dir`` is the cwd flag (not ``-C`` / @@ -72,15 +72,14 @@ def _render_opencode_config( "$schema": "https://opencode.ai/config.json", "mcp": {}, } - for server in ("registry", "knowledge"): - entry: dict = { - "type": "local", - "command": ["uv"] + _mcp_server_args(server), - "enabled": True, - } - if mcp_env: - entry["environment"] = dict(mcp_env) - config["mcp"][f"dsagt-{server}"] = entry + entry: dict = { + "type": "local", + "command": ["uv"] + _mcp_server_args(), + "enabled": True, + } + if mcp_env: + entry["environment"] = dict(mcp_env) + config["mcp"]["dsagt"] = entry providers: dict = {} if present_creds.get("OPENAI_API_KEY"): @@ -112,94 +111,25 @@ def _render_opencode_config( return json.dumps(config, indent=2) -def _render_opencode_config_proxy( - mcp_env: dict, proxy_port: int, opencode_model: str | None, -) -> str: - """Phase-2 proxy-mode opencode.json body. - - All provider routes hit ``http://localhost:`` regardless - of wire protocol — the proxy normalizes both /v1/messages and - /v1/chat/completions to the upstream's /chat/completions endpoint. - Sentinel API key plants a clear failure mode if a request bypasses - the proxy (the user sees 401, not silent leakage). - - Like the BYOA renderer, we register the user's chosen model under - ``provider..models`` so gateway-aliased names not in models.dev - don't trip ProviderModelNotFoundError. - """ - from .base import _PROXY_FORWARDED_SENTINEL - proxy_url = f"http://localhost:{proxy_port}" - config: dict = { - "$schema": "https://opencode.ai/config.json", - "mcp": {}, - } - for server in ("registry", "knowledge"): - entry: dict = { - "type": "local", - "command": ["uv"] + _mcp_server_args(server), - "enabled": True, - } - if mcp_env: - entry["environment"] = dict(mcp_env) - config["mcp"][f"dsagt-{server}"] = entry - - # Both providers point at the proxy — opencode picks via model prefix. - providers: dict = { - "openai": {"options": { - "apiKey": _PROXY_FORWARDED_SENTINEL, - "baseURL": proxy_url, - }}, - "anthropic": {"options": { - "apiKey": _PROXY_FORWARDED_SENTINEL, - "baseURL": proxy_url, - }}, - } - - if opencode_model and "/" in opencode_model: - provider_id, model_id = opencode_model.split("/", 1) - if provider_id in providers: - providers[provider_id]["models"] = { - model_id: {"name": model_id}, - } - config["model"] = opencode_model - - config["provider"] = providers - return json.dumps(config, indent=2) - - class OpenCodeSetup(AgentSetup): name = "opencode" base_command = ["opencode"] static_marker = "AGENTS.md" install_hint = "Install with `npm i -g opencode-ai`." - otel_payload_support = "none" - # OpenCode reads provider creds via ``{env:VAR}`` interpolation in - # its config — the file references these vars, opencode resolves - # them from the user's shell at runtime. Same shape as goose's - # multi-protocol story, no on-disk credential leakage. - credential_env_vars = ( - "OPENCODE_MODEL", - "OPENAI_API_KEY", "OPENAI_BASE_URL", - "ANTHROPIC_API_KEY", "ANTHROPIC_BASE_URL", - ) - telemetry_env = {} - credential_hints = ( - ("OPENCODE_MODEL", "model spec '/' " - "(e.g. 'openai/claude-haiku-4-5-20251001-v1-project' for a PNNL-shape " - "openai gateway, or 'anthropic/claude-sonnet-4-5')"), - ("OPENAI_API_KEY", "if your gateway speaks openai wire protocol"), - ("OPENAI_BASE_URL", "openai gateway URL " - "(referenced by opencode.json's provider.openai.options.baseURL)"), - ("ANTHROPIC_API_KEY", "if your gateway speaks anthropic wire protocol"), - ("ANTHROPIC_BASE_URL", "anthropic gateway URL"), - ) + # The AGENTS.md-convention skills dir codex/goose also use. + native_skills_dir = ".agents/skills" + + def owned_artifacts(self, working_dir: Path) -> list[Path]: + return [working_dir / "AGENTS.md", working_dir / "opencode.json"] def write_static(self, working_dir: Path) -> list[str]: actions: list[str] = [] instructions = _load_master_instructions() if instructions: action = _append_or_write( - working_dir / "AGENTS.md", instructions, _DSAGT_MARKER, + working_dir / "AGENTS.md", + instructions, + _DSAGT_MARKER, ) if action: actions.append(action) @@ -228,18 +158,21 @@ def write_dynamic( present = { name: bool(env.get(name)) for name in ( - "OPENAI_API_KEY", "OPENAI_BASE_URL", - "ANTHROPIC_API_KEY", "ANTHROPIC_BASE_URL", + "OPENAI_API_KEY", + "OPENAI_BASE_URL", + "ANTHROPIC_API_KEY", + "ANTHROPIC_BASE_URL", ) } body = _render_opencode_config( - mcp_env, present, opencode_model=env.get("OPENCODE_MODEL"), + mcp_env, + present, + opencode_model=env.get("OPENCODE_MODEL"), ) config_path = working_dir / "opencode.json" config_path.write_text(body + "\n") n_providers = sum( - 1 for k in ("OPENAI_API_KEY", "ANTHROPIC_API_KEY") - if present.get(k) + 1 for k in ("OPENAI_API_KEY", "ANTHROPIC_API_KEY") if present.get(k) ) actions.append( f"Wrote {config_path} ({len(mcp_env)} MCP env vars, " @@ -247,43 +180,6 @@ def write_dynamic( ) return actions - def proxy_write_dynamic( - self, - config: dict, - env: dict, - working_dir: Path, - pdir: Path, - ) -> list[str]: - """Phase-2 proxy-mode setup. All provider routes point at - ``http://localhost:``; the user's gateway creds live - in the proxy subprocess via ``LLM_API_KEY``. Same OPENCODE_MODEL - env var as BYOA — must be ``/``. - """ - del env, pdir - actions: list[str] = [] - proxy_port = (config.get("proxy") or {}).get("port") - opencode_model = (config.get("llm") or {}).get("model") - # opencode_model from config["llm"]["model"] is just a name; the - # user's OPENCODE_MODEL ``/`` form takes precedence - # if set, otherwise default to ``openai/`` since the - # proxy speaks /chat/completions on every upstream. - if opencode_model and "/" not in opencode_model: - opencode_model = f"openai/{opencode_model}" - if not proxy_port or not opencode_model: - raise RuntimeError( - "opencode proxy_write_dynamic requires " - "config['proxy']['port'] and config['llm']['model']." - ) - - mcp_env = _mcp_env_block(config) - body = _render_opencode_config_proxy(mcp_env, proxy_port, opencode_model) - config_path = working_dir / "opencode.json" - config_path.write_text(body + "\n") - actions.append( - f"Wrote {config_path} ({len(mcp_env)} MCP env vars, proxy mode)" - ) - return actions - def run_script( self, config: dict, @@ -309,15 +205,18 @@ def run_script( raise RuntimeError( "opencode batch mode requires OPENCODE_MODEL in the shell " "env, formatted as '/' (e.g. " - "'openai/claude-haiku-4-5-20251001-v1-project'). Plus the " - "matching {ANTHROPIC,OPENAI}_API_KEY / _BASE_URL. See " - "agents/opencode.py credential_hints." + "'openai/claude-haiku-4-5-20251001-v1-project'), plus the " + "matching {ANTHROPIC,OPENAI}_API_KEY / _BASE_URL — BYOA: " + "the agent must be pre-configured in the shell." ) cmd = [ - "opencode", "run", - "--dir", str(working_dir), + "opencode", + "run", + "--dir", + str(working_dir), "--dangerously-skip-permissions", - "-m", model, + "-m", + model, text, ] return _run_simple_script(cmd, env, working_dir, self.install_hint) diff --git a/src/dsagt/agents/roo.py b/src/dsagt/agents/roo.py deleted file mode 100644 index 6cbc403..0000000 --- a/src/dsagt/agents/roo.py +++ /dev/null @@ -1,244 +0,0 @@ -""" -Roo Code agent setup. - -Install: ``npm i -g @roo-code/cli`` (or curl install script — -binary lands at ``~/.local/bin/roo``). - -Generates: ``.roomodes``; per-init, ``.roo/mcp.json`` is written by -:meth:`RooSetup.write_dynamic` with the full env block (roo, like cline, -doesn't inherit parent env into MCP server children). - -**BYOA Phase 1 status: PUNTED for batch / smoke-test.** Roo CLI v0.1.17 -hardcodes endpoints per ``--provider``: - - * ``anthropic`` / ``openai-native`` etc. routes through Roo Code Cloud - (``ROO_CODE_PROVIDER_URL`` defaulting to ``api.roocode.com/proxy``). - With no ``--base-url`` flag, no `*_BASE_URL` env var honored, and - no native-SDK fallback, the CLI cannot be pointed at a self-hosted - gateway. Without cloud auth it falls back to the upstream provider's - default URL (``api.openai.com`` etc.) where the lab-gateway key 401s. - * ``bedrock`` is rejected by the CLI's ``--provider`` validator - despite the ``AwsBedrockHandler`` existing in the bundle and - ``awsBedrockEndpoint`` being a configurable schema field. - * ``vercel-ai-gateway`` has a literally hardcoded URL constant - (``AI_GATEWAY_BASE_URL = "https://ai-gateway.vercel.sh/v1"``); not - even an env var. - * The same model-whitelist substitution issue cline has would also - bite roo's anthropic path — old_code's docstring called it out: - "Roo rewrites our PNNL model name into its own default before - sending — the proxy aliases that name back to the upstream primary." - -Phase 2 reactivates batch mode via ``dsagt-proxy`` at ``localhost:`` -— roo's anthropic SDK happily talks to a localhost endpoint and the -proxy translates names back to upstream-served IDs. - -Interactive use is unaffected — ``dsagt init --agent roo`` still writes -the project state, and users running roo via VS Code (where the UI -exposes ``awsBedrockEndpoint``, ``openAiNativeBaseUrl`` etc.) can drive -the project manually. - -OTel support: **none** (verified, ``otel_payload_support = "none"``). -Roo Code imports zero ``@opentelemetry/*`` packages anywhere in the -agent runtime (``src/``, ``apps/cli/``, ``packages/core/``, -``packages/telemetry/``, ``packages/cloud/``). LLM-call sites in -``src/api/providers/*.ts`` have no span/tracer wrapping. Telemetry is -PostHog-only with payload-free events -(``LLM_COMPLETION``: token counts + cost; ``TASK_CONVERSATION_MESSAGE``: -``taskId`` + ``source``). Standard ``OTEL_EXPORTER_OTLP_ENDPOINT`` is -ignored. Memory extraction will see no agent-conversation traces; -tool execution and KB observability still work via dsagt-run / -MCP-server spans. - -Workaround: ``dsagt start --enable-proxy`` makes every roo LLM call -visible in MLflow — each agent turn becomes an inspectable trace you -can audit live or replay later. Roo's ``--provider anthropic`` honors -``ANTHROPIC_BASE_URL`` env (which ``agent_env`` overrides to the proxy -when the flag is set), so the proxy receives every LLM call and emits -an OTel trace on roo's behalf. Memory extraction is a downstream -consequence; the primary value is real-time transparency on the -agent's actions. See ``commands/proxy_server.py``. - -Batch mode: ``roo --print --oneshot --prompt-file FILE`` reads multi-line -prompts directly — no looping or argv encoding needed. -""" - -from __future__ import annotations - -import json -from pathlib import Path - -from .base import ( - AgentSetup, - _append_or_write, - _build_mcp_servers_dict, - _DSAGT_MARKER, - _format_roomodes, - _load_master_instructions, - _mcp_env_block, - _run_simple_script, -) - - -class RooSetup(AgentSetup): - name = "roo" - base_command = ["roo"] - static_marker = ".roomodes" - install_hint = ( - "Install via " - "https://github.com/RooCodeInc/Roo-Code/blob/main/apps/cli/install.sh" - ) - otel_payload_support = "none" - # Roo's CLI multi-provider story is misleading: openai-native has no - # ``--base-url`` flag and the SDK behind it doesn't read - # ``OPENAI_BASE_URL``, so reaching a non-default openai endpoint is - # impossible with that path. Only the anthropic provider works for - # gateways — the Anthropic SDK natively reads ``ANTHROPIC_BASE_URL``. - credential_env_vars = ( - "ANTHROPIC_API_KEY", "ANTHROPIC_BASE_URL", "ANTHROPIC_MODEL", - ) - # Roo emits no OTel — agent-side telemetry only via --proxy_traces. - telemetry_env = {} - credential_hints = ( - ("ANTHROPIC_API_KEY", "your provider API key (works for openai-shape " - "gateways too — the Anthropic SDK reaches them via ANTHROPIC_BASE_URL)"), - ("ANTHROPIC_BASE_URL", "gateway / proxy URL " - "(roo CLI has no --base-url flag; this env var is the only way)"), - ("ANTHROPIC_MODEL", "model name your gateway serves " - "(e.g. claude-haiku-4-5-20251001-v1-project)"), - ) - - def vscode_hint(self, project_dir: Path) -> list[str]: - return [ - f"Open {project_dir} in VS Code and start the Roo Code extension.", - "Pick 'DSAgt Pipeline Builder' from Roo's mode dropdown.", - ] - - def write_static(self, working_dir: Path) -> list[str]: - actions: list[str] = [] - (working_dir / ".roo").mkdir(exist_ok=True) - instructions = _load_master_instructions() - if instructions: - # Roo wraps the instructions in a customMode JSON envelope; the - # master marker text survives the wrap because it's part of the - # role definition body. - action = _append_or_write( - working_dir / ".roomodes", - _format_roomodes(instructions), - _DSAGT_MARKER, - ) - if action: - actions.append(action) - return actions - - def env_overrides(self, config: dict) -> dict[str, str]: - """Phase-2 proxy-mode hook: pin ``ANTHROPIC_MODEL`` so roo's - anthropic SDK posts the configured model to the proxy. - - ``proxy_env_overrides`` (base default) sets - ``ANTHROPIC_BASE_URL`` to the proxy URL and ``ANTHROPIC_API_KEY`` - to the sentinel. Roo's hardcoded model whitelist will rewrite - the model name to its current default (``claude-sonnet-4-5``) - before sending; the proxy aliases that rewrite back to the - upstream-served name (see ``_AGENT_PRIMARY_ALIASES`` in - proxy_server.py). - """ - model = (config.get("llm") or {}).get("model") - if model and not str(model).startswith("${"): - return {"ANTHROPIC_MODEL": model} - return {} - - def write_dynamic( - self, - config: dict, - env: dict, - working_dir: Path, - pdir: Path, - ) -> list[str]: - """Write ``.roo/mcp.json``. Roo doesn't inherit parent env into - MCP children — every var the dsagt servers need - (``MLFLOW_TRACKING_URI``, ``DSAGT_PROJECT_DIR``, ``EMBEDDING_*``) - must be explicit in the JSON env block. - """ - del env, pdir - actions: list[str] = [] - mcp_path = working_dir / ".roo" / "mcp.json" - mcp_path.parent.mkdir(parents=True, exist_ok=True) - mcp_env = _mcp_env_block(config) - mcp_config = _build_mcp_servers_dict(mcp_env) - mcp_path.write_text(json.dumps(mcp_config, indent=2) + "\n") - actions.append(f"Wrote {mcp_path} ({len(mcp_env)} env vars)") - return actions - - def launch_oneliner(self, project: str, project_dir: Path) -> str: - """Roo loads customModes from ``.roomodes``; without - ``--mode dsagt`` it runs in the default ``code`` mode and the - DSAgt customInstructions are dropped from the system prompt. - """ - del project - import shlex - pdir = shlex.quote(str(project_dir)) - return f"cd {pdir} && roo --mode dsagt" - - def run_script( - self, - config: dict, - env: dict, - working_dir: Path, - script_path: Path, - max_turns: int, - ) -> int: - """Punted in BYOA Phase 1 — see module docstring. - - Roo CLI v0.1.17 has no way to reach a self-hosted gateway (no - ``--base-url``, no honored env var, ``bedrock`` rejected at the - provider validator). Phase 2 ``dsagt-proxy`` provides a - localhost endpoint roo's anthropic SDK happily talks to. - """ - del config, env, working_dir, script_path, max_turns - raise RuntimeError( - "roo batch mode is punted in Phase 1 BYOA — roo CLI v0.1.17 " - "has no flag, env var, or non-cloud auth path that points it " - "at a self-hosted gateway. Phase 2 reactivates this via " - "dsagt-proxy. See agents/roo.py module docstring." - ) - - def proxy_run_script( - self, - config: dict, - env: dict, - working_dir: Path, - script_path: Path, - max_turns: int, - ) -> int: - """Phase-2 proxy-mode batch run. Un-punts roo: routes through - ``dsagt-proxy`` at ``ANTHROPIC_BASE_URL`` (set by - ``proxy_env_overrides`` in the base class), so roo's anthropic - SDK posts to localhost. The proxy aliases roo's hardcoded - ``claude-sonnet-4-5`` rewrite back to the upstream's served - name (see ``_AGENT_PRIMARY_ALIASES`` in proxy_server.py). - - Sentinel API key flows through ``--api-key`` so any direct - bypass 401s loudly. Model is the upstream-served name from - ``config["llm"]["model"]`` — roo will rewrite it to its - whitelist default before sending, but the proxy aliases that - rewrite back to the configured upstream. - """ - del max_turns - from .base import _PROXY_FORWARDED_SENTINEL - model = (config.get("llm") or {}).get("model") - if not model: - raise RuntimeError( - "roo proxy_run_script requires config['llm']['model']." - ) - cmd = [ - "roo", - "--print", "--oneshot", - "--mode", "dsagt", - "--prompt-file", str(script_path), - "--workspace", str(working_dir), - "--debug", - "--provider", "anthropic", - "--api-key", _PROXY_FORWARDED_SENTINEL, - "--model", model, - ] - return _run_simple_script(cmd, env, working_dir, self.install_hint) diff --git a/src/dsagt/tools/scan_directory.md b/src/dsagt/codes/scan-directory/SKILL.md similarity index 90% rename from src/dsagt/tools/scan_directory.md rename to src/dsagt/codes/scan-directory/SKILL.md index 2f4eb6e..b519101 100644 --- a/src/dsagt/tools/scan_directory.md +++ b/src/dsagt/codes/scan-directory/SKILL.md @@ -1,8 +1,8 @@ --- -name: scan_directory +name: scan-directory description: Scan a data directory and produce structured report with file counts, sizes, and directory tree -executable: dsagt-run --tool scan_directory -- python -m dsagt.tools.scan_directory +executable: dsagt-run --code scan-directory -- python codes/scan-directory/scripts/scan_directory.py parameters: directory: type: string @@ -23,7 +23,7 @@ parameters: description: Number of largest files to list --- -# scan_directory +# scan-directory Scan a data directory and produce a structured report with file counts, sizes, and directory tree. Use this as your first step when exploring a new dataset to understand its layout before deciding how to process it. diff --git a/src/dsagt/tools/scan_directory.py b/src/dsagt/codes/scan-directory/scripts/scan_directory.py similarity index 82% rename from src/dsagt/tools/scan_directory.py rename to src/dsagt/codes/scan-directory/scripts/scan_directory.py index a54742b..c4c4ebd 100644 --- a/src/dsagt/tools/scan_directory.py +++ b/src/dsagt/codes/scan-directory/scripts/scan_directory.py @@ -57,7 +57,9 @@ def du_total_bytes(directory: Path) -> int: return kb * 1024 -def iter_files_sizes_relative(directory: Path, max_depth: int) -> tuple[list[tuple[int, str]], list[dict]]: +def iter_files_sizes_relative( + directory: Path, max_depth: int +) -> tuple[list[tuple[int, str]], list[dict]]: """ Enumerate files under `directory` up to `max_depth`, INCLUDING hidden. Uses: find + stat (platform-specific flags) via -exec ... {} + @@ -102,10 +104,17 @@ def iter_files_sizes_relative(directory: Path, max_depth: int) -> tuple[list[tup continue items.append((size, rel)) except Exception: - errors.append({"path": str(directory), "error": f"Unparseable stat line: {line[:200]}"}) + errors.append( + { + "path": str(directory), + "error": f"Unparseable stat line: {line[:200]}", + } + ) if p.returncode != 0: - errors.append({"path": str(directory), "error": f"find/stat exited {p.returncode}"}) + errors.append( + {"path": str(directory), "error": f"find/stat exited {p.returncode}"} + ) return items, errors @@ -158,24 +167,32 @@ def scan(directory: Path, max_depth: int, top_n: int) -> dict: "total_size": v["total_size"], "total_size_human": format_size(v["total_size"]), } - for k, v in sorted(ext_summary.items(), key=lambda x: x[1]["total_size"], reverse=True) + for k, v in sorted( + ext_summary.items(), key=lambda x: x[1]["total_size"], reverse=True + ) ] - largest_files = [e for _, _, e in sorted(largest_heap, key=lambda t: (t[0], t[1]), reverse=True)] + largest_files = [ + e for _, _, e in sorted(largest_heap, key=lambda t: (t[0], t[1]), reverse=True) + ] - directory_tree = [{ - "directory": ".", - "size": total_size, - "file_count": total_files, - "size_human": format_size(total_size), - }] + directory_tree = [ + { + "directory": ".", + "size": total_size, + "file_count": total_files, + "size_human": format_size(total_size), + } + ] for d in sorted(dir_summary): - directory_tree.append({ - "directory": d, - "size": dir_summary[d]["size"], - "file_count": dir_summary[d]["file_count"], - "size_human": format_size(dir_summary[d]["size"]), - }) + directory_tree.append( + { + "directory": d, + "size": dir_summary[d]["size"], + "file_count": dir_summary[d]["file_count"], + "size_human": format_size(dir_summary[d]["size"]), + } + ) report = { "directory": str(directory.resolve()), @@ -195,7 +212,9 @@ def scan(directory: Path, max_depth: int, top_n: int) -> dict: def main(): - parser = argparse.ArgumentParser(description="Scan directory structure and file sizes") + parser = argparse.ArgumentParser( + description="Scan directory structure and file sizes" + ) parser.add_argument("directory", help="Path to directory to scan") # accept both hyphen and underscore variants (registry runners vary) parser.add_argument("--max-depth", "--max_depth", type=int, default=5) @@ -204,7 +223,9 @@ def main(): directory = Path(args.directory) if not directory.exists(): - print(json.dumps({"error": f"Directory not found: {directory}"}), file=sys.stderr) + print( + json.dumps({"error": f"Directory not found: {directory}"}), file=sys.stderr + ) sys.exit(1) if not directory.is_dir(): print(json.dumps({"error": f"Not a directory: {directory}"}), file=sys.stderr) diff --git a/src/dsagt/commands/cli.py b/src/dsagt/commands/cli.py index 40d30d4..e1fa39d 100644 --- a/src/dsagt/commands/cli.py +++ b/src/dsagt/commands/cli.py @@ -1,255 +1,444 @@ """ DSAgt CLI — project initialization and session management. -Two flows: - -1. **BYOA** (default): ``dsagt init --agent `` writes per-agent - MCP config artifacts. ``dsagt mlflow `` backgrounds MLflow - and prints the OTel routing exports the user pastes into the shell - that runs ``claude`` / ``goose``. Native-OTel traces appear in the - MLflow UI but use a shape (``api_response_body`` log events) that - ``dsagt memory`` cannot extract from — for episodic memory, use - proxy mode. -2. **Proxy mode**: ``dsagt start --enable-proxy `` interposes - a LiteLLM proxy between the agent and its provider, autologs every - LLM call into MLflow with ``mlflow.spanInputs`` / - ``mlflow.spanOutputs`` populated. This is the canonical shape that - ``dsagt memory --project X`` reads for episodic-memory extraction - and that the MLflow UI's request/response columns surface natively. +``dsagt init`` is the single, interactive, re-runnable place a user expresses +every choice; the prompts mirror ``.dsagt/config.yaml`` 1:1 and it writes the +per-agent instructions + MCP config. ``dsagt start `` refreshes the +dynamic agent record (MCP config + native-skills mirror, idempotent) and then +launches the agent — ``cd && `` plus the refresh (the MCP +server owns the session lifecycle, minting session ids into +``.dsagt/state.yaml`` and catching up post-session extraction in the +background at startup). + +The agent talks to its provider directly — DSAGT never interposes on its +traffic. Self-logging goes to a serverless ``sqlite:////mlflow.db`` +store (no server to run). Usage: - dsagt init --agent [--mlflow-port ] [--location ] - dsagt mlflow - dsagt memory --project - dsagt start [--agent ] [--mlflow-port ] + dsagt init [] # interactive; re-run to reconfigure + dsagt init --agent [--location ] + [--include ... | --exclude ...] # non-interactive + dsagt start dsagt info [--json] - dsagt stop + dsagt traces [--port ] dsagt smoke-test - dsagt setup-kb [--collection ] [--embedding-* flags] dsagt list dsagt mv dsagt rm [-y] [--keep-files] """ import argparse -import json import logging import os -import signal import subprocess import sys import time -from datetime import datetime, timezone from pathlib import Path from dsagt.agents import ( agent_env, dynamic_agent_record, launch_agent, - static_agent_files_present, static_agent_record, AGENTS, ) from dsagt.session import ( + DEFAULT_PROJECTS_BASE, VALID_AGENTS, list_projects, load_config, init_project, - mlflow_command, move_project, - persist_agent_choice, - pick_free_port, project_dir, + read_config_file, + remove_collection, remove_project, - run_extraction, - start_services, - stop_services, ) logger = logging.getLogger(__name__) +# --------------------------------------------------------------------------- +# Interactive prompt helpers +# +# Selection-style prompts (agent, KB collections, skill sources) use +# ``questionary`` for arrow-key navigation + space-to-toggle checkboxes — no +# typing required. Free-text (name / location) and y/N confirms stay plain. +# These run only on the interactive (TTY) path; automation drives init via +# flags and never reaches them. +# --------------------------------------------------------------------------- + + +def _prompt(text: str, default: str | None = None) -> str: + """Free-text prompt; empty input returns *default*.""" + suffix = f" [{default}]" if default not in (None, "") else "" + resp = input(f"{text}{suffix}: ").strip() + return resp or (default or "") + + +def _confirm(text: str, default: bool = False) -> bool: + """Yes/no prompt; empty input returns *default*.""" + hint = "Y/n" if default else "y/N" + resp = input(f"{text} [{hint}] ").strip().lower() + if not resp: + return default + return resp in ("y", "yes") + + +def _select(message: str, options: list[str], default: str) -> str: + """Single-select menu (arrow keys). Aborts init on cancel (Ctrl-C).""" + import questionary + + answer = questionary.select(message, choices=options, default=default).ask() + if answer is None: + raise SystemExit("dsagt init: cancelled.") + return answer + + +def _checkbox(message: str, choices: list[tuple[str, str, bool]]) -> list[str]: + """Multi-select checkbox menu (space toggles, enter confirms). + + *choices* is a list of ``(value, label, checked)``. Returns the selected + values. Aborts init on cancel (Ctrl-C). + """ + import questionary + + qchoices = [ + questionary.Choice(title=label, value=value, checked=checked) + for value, label, checked in choices + ] + # ``instruction=""`` suppresses questionary's own hint line — our message + # already spells out the controls. + answer = questionary.checkbox(message, choices=qchoices, instruction="").ask() + if answer is None: + raise SystemExit("dsagt init: cancelled.") + return answer + + +def _current_assets(pdir: Path) -> list[str]: + """Asset names whose collection already exists in the project's kb_index.""" + from dsagt.commands.setup_core_kb import all_assets, asset_collection_name + + present = [] + for asset in all_assets(): + try: + coll = asset_collection_name(asset) + except ValueError: + continue + if (pdir / "kb_index" / coll).exists(): + present.append(asset) + return present + + +def _skills_block_for(source_names: list[str]) -> dict: + """Build the ``skills`` config block from the selected skill-source names.""" + from dsagt.skills import KNOWN_SOURCES + + sources = [] + for name in source_names: + src = KNOWN_SOURCES[name] + sources.append( + { + "name": name, + "url": src["url"], + "branch": src.get("branch", "main"), + "subdir": src.get("subdir"), + } + ) + return {"sources": sources} + + +def _episodic_block(enabled: bool) -> dict | None: + """The ``episodic`` config block, or ``None`` when the user didn't opt in. + + Enabling captures each completed turn into ``session_memory`` (mechanical + chunk + tag + embed). ``None`` keeps a disabled project's config minimal + (``enabled: false`` is backfilled on read). + """ + if not enabled: + return None + return {"enabled": True} + + +def _collect_settings(args, interactive: bool, existing: dict, pdir: Path | None): + """Resolve the init choices (the 1:1 mirror of the config). + + Selection questions: agent platform, packaged KB document *collections*, + skill-catalog *sources*, and the episodic-memory opt-in. The bundled + ``tools`` collection is always provisioned and is NOT a per-project choice. + Project name + folder location are resolved by the caller. Embedding / + chunk_size / rerank are code defaults, not init choices. + + Interactive: questionary select/checkbox menus + y/N, pre-filled with the + project's current choices on re-init. Non-interactive (no TTY): drive from + ``--include`` / ``--exclude`` / ``--episodic`` flags — the automation/test + path. + """ + from dsagt.commands.setup_core_kb import COLLECTIONS, resolve_assets + from dsagt.skills import KNOWN_SOURCES + + coll_choices = list(COLLECTIONS) + skill_choices = list(KNOWN_SOURCES) + + if interactive: + agent = _select( + "Agent platform", + list(VALID_AGENTS), + default=existing.get("agent") or args.agent or "claude", + ) + + # Knowledge collections (heavy doc collections; default none). + # Labels are bare names — short enough to never wrap the terminal. + cur_colls = set(existing.get("knowledge", {}).get("collections", [])) + collections = _checkbox( + "Knowledge collections (space toggles, ↑/↓ to move, enter confirms)", + [(c, c, c in cur_colls) for c in coll_choices], + ) + + # Skill-catalog sources (default genesis on a fresh project). + cur_srcs = set( + s["name"] for s in existing.get("skills", {}).get("sources", []) + ) or {"genesis"} + skill_names = _checkbox( + "Skill catalog sources (space toggles, ↑/↓ to move, enter confirms)", + [(s, s, s in cur_srcs) for s in skill_choices], + ) + + # Episodic memory (opt-in): captures session turns into session_memory. + cur_epi = existing.get("episodic", {}) or {} + enable_epi = _confirm( + "Enable episodic memory? (captures session turns into searchable " + "memory)", + default=bool(cur_epi.get("enabled")), + ) + episodic = _episodic_block(enable_epi) + else: + agent = args.agent or existing.get("agent") + if not agent: + raise SystemExit("dsagt init: --agent is required (non-interactive).") + # --include / --exclude pick the full asset set; split it. + full = resolve_assets(include=args.include, exclude=args.exclude) + collections = [a for a in full if a in COLLECTIONS] + skill_names = [a for a in full if a in KNOWN_SOURCES] + # Episodic is flag-driven here (automation); omit --episodic to leave it + # off. Re-pass it on re-init — like --include/--exclude, flags are + # authoritative on the non-interactive path. + episodic = _episodic_block(getattr(args, "episodic", False)) + + return { + "agent": agent, + # The bundled ``tools`` collection is always provisioned. + "assets": ["codes", *collections, *skill_names], + "knowledge": {"collections": collections}, + "skills": _skills_block_for(skill_names), + "episodic": episodic, + } + + +def _handle_destructive( + existing: dict, settings: dict, pdir: Path, interactive: bool +) -> None: + """Detect destructive deltas on re-init and prompt delete-or-keep. + + Non-interactive (no TTY): never deletes — warns and keeps, so automation + can't lose data and ``input()`` is never called on a closed stdin. + + Never touches agent-populated data: ``tool_use`` / ``session_memory`` + collections, ``.dsagt/`` memory, ``trace_archive/``, ``skills/``. + """ + from dsagt.commands.setup_core_kb import asset_collection_name + + protected = {"code_use", "session_memory"} + + # Agent switch → old platform's files are now stale. + old_agent = existing.get("agent") + new_agent = settings["agent"] + if old_agent and old_agent != new_agent: + setup = AGENTS[old_agent]() + stale = [p for p in setup.owned_artifacts(pdir) if p.exists()] + if stale: + print(f"\n Switching agent {old_agent} → {new_agent} leaves stale files:") + for p in stale: + print(f" {p}") + if interactive and _confirm(" Delete these stale files?", default=True): + import shutil + + for p in stale: + if p.is_dir(): + shutil.rmtree(p) + else: + p.unlink() + print(" Deleted.") + else: + print(" Kept (remove them manually if you want them gone).") + + # Removed KB collections → their dirs are now orphaned. + new_colls = set() + for a in settings["assets"]: + try: + new_colls.add(asset_collection_name(a)) + except ValueError: + pass + for a in _current_assets(pdir): + try: + coll = asset_collection_name(a) + except ValueError: + continue + if coll in new_colls or coll in protected: + continue + if interactive and _confirm( + f"\n Collection '{coll}' was dropped from the asset set. Remove it?", + default=False, + ): + if remove_collection(pdir, coll): + print(f" Removed kb_index/{coll}.") + else: + print( + f"\n Collection '{coll}' was dropped from the asset set " + "(kept on disk)." + ) + + # Embedding backend/model change invalidates existing embeddings. + old_emb = existing.get("embedding") or {} + new_emb = settings["embedding"] + changed = old_emb.get("backend") != new_emb.get("backend") or old_emb.get( + "model" + ) != new_emb.get("model") + if changed and any((pdir / "kb_index").glob("*/")): + print( + "\n WARNING: embedding backend/model changed. Existing collections " + "were embedded with the old settings and KB search will degrade until " + "they are re-provisioned (delete + re-add the affected assets)." + ) + + def _cmd_init(args): - """Create a new BYOA project. + """Create or reconfigure a BYOA project — interactive and re-runnable. - Writes the agent's static files (instructions, state dirs) and the - runtime artifacts (MCP config: ``.mcp.json`` for claude, ``goose.yaml`` - for goose, etc.) populated with the MLflow port pinned at init time. - Then prints the env-var block + launch one-liner the user needs to - run their own agent. + ``dsagt init`` is the single place a user expresses every choice; the + prompts mirror ``.dsagt/config.yaml`` 1:1. On an existing project it + becomes a settings editor (prompts prefilled with current values) and + prompts before any destructive change (agent switch, removed collection). + Non-interactive (no TTY) drives from flags — the automation/test path. """ - location = Path(args.location).resolve() if args.location else None - pdir, mlflow_port = init_project( - args.project, - args.agent, - mlflow_port=args.mlflow_port, + interactive = sys.stdin.isatty() + + # Project name + name = args.project + if interactive and not name: + name = _prompt("Project name") + if not name: + raise SystemExit("dsagt init: project name required.") + + # Existing project? → re-init (settings editor). + try: + existing_pdir = project_dir(name) + except FileNotFoundError: + existing_pdir = None + existing = read_config_file(existing_pdir) if existing_pdir else {} + reinit = bool(existing) + + # Location (first init only; re-init keeps the registered path). The + # prompt collects the full project directory and defaults to one that + # already ends in the project name. If the user types a path that ends + # in the project name we take it as-is; otherwise we append the name — + # so both "~/proj/myproj" and "~/proj" land at "~/proj/myproj". + if reinit: + location = existing_pdir.parent + pdir_preview = existing_pdir + else: + if interactive: + default_full = str(DEFAULT_PROJECTS_BASE / name) + entered = Path(_prompt("Project location", default=default_full)).resolve() + proj_dir = entered if entered.name == name else entered / name + location = proj_dir.parent + else: + location = Path(args.location).resolve() if args.location else None + pdir_preview = (location or DEFAULT_PROJECTS_BASE) / name + + settings = _collect_settings(args, interactive, existing, pdir_preview) + + if reinit: + _handle_destructive(existing, settings, existing_pdir, interactive) + + include = settings["assets"] if settings["assets"] else None + exclude = ["all"] if not settings["assets"] else None + pdir = init_project( + name, + settings["agent"], location=location, + include=include, + exclude=exclude, + knowledge=settings["knowledge"], + skills=settings["skills"], + episodic=settings["episodic"], ) - print(f"Project created: {pdir}") - print(f" Agent: {args.agent}") - print(f" MLflow port: {mlflow_port}") - print() - config = load_config(args.project) - for action in static_agent_record(config, args.agent, pdir): - print(f" {action}") + agent = settings["agent"] + config = load_config(name) + + # 1. Actions first (Wrote … / Mirrored …). + print() + for action in static_agent_record(config, agent, pdir): + print(action) # Pass the user's shell env so per-agent ``write_dynamic`` can read - # provider creds (e.g., cline.write_dynamic invokes ``cline auth`` - # with ANTHROPIC_API_KEY / ANTHROPIC_MODEL from the shell). + # provider creds (e.g., cline.write_dynamic invokes ``cline auth``). for action in dynamic_agent_record(config, env=dict(os.environ), working_dir=pdir): - print(f" {action}") - - setup = AGENTS[args.agent]() + print(action) + # 2. Project summary. print() - cred_hints = setup.byoa_env_hints(mlflow_port, args.project, pdir) - if cred_hints: - print("Provider credentials (set in your shell; skip any your agent is") - print("already configured to handle for plain chat/coding):") - print() - for var, hint in cred_hints: - print(f" export {var}=... # {hint}") - print() + print(f"Project directory: {pdir}") + print(f"Agent: {agent}") + print(f"Trace store: sqlite:///{pdir}/mlflow.db") - print("Three ways to start:") - print() - print(f" 1. dsagt start {args.project} [--enable-proxy]") - print(" → DSAGT runs everything (MLflow + agent; proxy if requested).") - print() - print(f" 2. cat {pdir}/dsagt-launch.sh") - print(" → view & run the commands manually for full transparency.") - print() - print(f" 3. bash {pdir}/dsagt-launch.sh") - print(" → starts MLflow, sets env, then prints how to launch the agent.") + # 3. Startup instructions. print() - print("Options 2 & 3 are BYOA-only (no proxy).") - if args.agent == "claude": + print(f"Start {agent} in the project directory, or run:") + print(f" dsagt start {name}") + if AGENTS[agent]().vscode_hint(pdir): print() - print("Note: `mlflow autolog claude` was configured automatically.") print( - f" Traces appear at http://localhost:{mlflow_port} after each session." + f"Or open the project directory in VS Code and start the " + f"{agent} extension" ) - print() - print("After your session, extract memory:") - print(f" dsagt memory --project {args.project}") def _cmd_start(args): - """Start a project session: resolve agent → pick ports → start services → - write agent configs → launch. - - Order matters: services start before the agent's runtime configs are - written so the actually-bound ports flow into MCP env blocks. - Static files (instructions) are written here only if missing — when - the user ran ``dsagt init --agent X``, the static record was already - written at init time and we don't touch it. + """Refresh the dynamic agent record, then ``cd && ``. + + The refresh (``dynamic_agent_record``: per-agent MCP config + + native-skills mirror, all idempotent) backstops the install-time mirror + the MCP tools run (``agents.refresh_native_skills``) — e.g. a skill dir + dropped in by hand still mirrors here — and lets a credential-dependent + step skipped at init (cline auth) pick up once the vars are in the + shell. Session lifecycle (session-id + minting, post-session extraction catch-up) is owned by the MCP server + at startup. + + The per-project runtime env (e.g. ``CLINE_DIR`` / ``CODEX_HOME`` that + point an agent at its init-written config) is applied so the launched + agent finds what ``init`` set up. """ config = load_config(args.project) pdir = Path(config["project_dir"]) - # Step 1: resolve agent. CLI overrides YAML. If neither is set, - # fail with a one-line message naming both options. - yaml_agent = config.get("agent") - agent = args.agent or yaml_agent - if not agent: - raise ValueError( - "No agent specified. Pass --agent or set 'agent:' in " - f"{pdir / 'dsagt_config.yaml'}." - ) - if agent not in VALID_AGENTS: - raise ValueError(f"agent must be one of {VALID_AGENTS}, got '{agent}'") - config["agent"] = agent - - # Step 2: persist agent into YAML if this is the project's first - # encounter with one. CLI overrides on later runs are per-run only - # — they don't touch the YAML default. - if not yaml_agent: - persist_agent_choice(args.project, agent) - - # Step 3: write the static record on demand — when init didn't run - # with --agent, when the user switched agents, or when a marker - # file was deleted. Idempotent: skips when already present. - if not static_agent_files_present(agent, pdir): - for action in static_agent_record(config, agent, pdir): - print(f" {action}") - - # Step 4: synthesize session id (one per start, threaded through every - # subprocess via DSAGT_SESSION_ID so the MLflow UI can filter to one - # session). - config["session_id"] = ( - f"{config['project']}-" - f"{datetime.now(timezone.utc).strftime('%Y%m%dT%H%M%SZ')}" - ) - - # Step 5: start MLflow (and optionally dsagt-proxy). Picks free ports - # automatically and writes them to /.runtime; CLI overrides - # are honored if present. EVERYTHING from here on must be inside the - # try/finally — otherwise an exception between service start and - # ``launch_agent`` leaks subprocesses we just spawned. - if args.mlflow_port: - config.setdefault("mlflow", {})["port"] = args.mlflow_port - if args.enable_proxy: - # Marker that triggers proxy subprocess in start_services + URL - # injection in agent_env. Port is picked by start_services via - # pick_free_port — same path as MLflow. - config.setdefault("proxy", {}) - ports = start_services(config) - if "proxy" in ports: - print(f" Ports: mlflow={ports['mlflow']}, proxy={ports['proxy']}") - else: - print(f" Ports: mlflow={ports['mlflow']}") - - try: - # Step 7: build env, write dynamic agent configs (with the actual - # MLflow port baked in), launch. - env = agent_env(config) - for action in dynamic_agent_record(config, env, pdir): - print(f" {action}") - - print() - print(f" Project: {config['project']}") - print(f" Agent: {config['agent']}") - print(f" Dir: {pdir}") - print(f" MLflow: http://localhost:{config['mlflow']['port']}") - print() - - return launch_agent( - config, - env, - pdir, - script_path=args.script, - max_turns=args.max_turns, - ) - finally: - print() - # Extraction is best-effort — don't let its failures keep us from - # cleaning up services. - try: - result = run_extraction(args.project) - n_indexed = result.get("tool_use_indexed", 0) - if n_indexed: - print(f" Indexed {n_indexed} tool execution(s) into tool_use") - if result.get("status") == "ok": - print(f" Extracted {result['total_entries']} memories from session") - elif result.get("status") == "empty": - print(" No session exchanges to extract") - # ``tool_use_only`` is the BYOA default (no DSAGT_MEMORY_* - # configured); the indexed-count line above already covered it. - except Exception as e: - print(f" WARNING: extraction failed: {e}") - - # Phase 2 proxy mode: surface any sidechannel-call hits from this - # session so the user can spot a typo in their primary llm.model - # vs. a harmless agent title-gen / session-namer call. No-op - # when the proxy didn't run or no hits were logged. - try: - from dsagt.observability import print_sidechannel_warning + print(f" Project: {config['project']}") + print(f" Agent: {config['agent']}") + print(f" Dir: {pdir}") + print() - print_sidechannel_warning(pdir, config.get("session_id")) - except Exception as e: - logger.debug("sidechannel warning failed: %s", e) + for action in dynamic_agent_record(config, env=dict(os.environ), working_dir=pdir): + print(f" {action}") - _stop_one(args.project) + env = agent_env(config) + return launch_agent( + config, + env, + pdir, + script_path=args.script, + max_turns=args.max_turns, + ) def _cmd_list(args): @@ -263,28 +452,19 @@ def _cmd_list(args): for name, path in projects.items(): pdir = Path(path) - config_path = pdir / "dsagt_config.yaml" + cfg_file = pdir / ".dsagt" / "config.yaml" - # Best-effort: if the config is readable, show agent + service status. - # If the project dir is gone or config is broken, just show the path. + # Best-effort: if the config is readable, show the agent. If the + # project dir is gone or the config is broken, just show the path. agent = "" - status = "" - if config_path.exists(): + if cfg_file.exists(): try: config = load_config(name) agent = config.get("agent", "") except (FileNotFoundError, ValueError): pass - runtime_path = pdir / ".runtime" - if runtime_path.exists(): - state = json.loads(runtime_path.read_text()) - ports = state.get("ports", {}) - status = f"running (mlflow:{ports.get('mlflow','?')})" - else: - status = "stopped" - - print(f" {name:<20} {agent:<14} {status:<40} {path}") + print(f" {name:<20} {agent:<14} {path}") def _cmd_mv(args): @@ -297,9 +477,7 @@ def _cmd_mv(args): def _cmd_rm(args): """Unregister a project and (by default) delete its directory. - With ``--all``: bulk-remove every registered project. Reaps active - services per-project (via ``stop_services``) before removing so a - leftover ``.runtime`` doesn't block ``remove_project``. + With ``--all``: bulk-remove every registered project. """ if args.all and args.project: raise SystemExit("dsagt rm: pass either a project name or --all, not both.") @@ -354,12 +532,6 @@ def _cmd_rm_all(args) -> int: failures: list[tuple[str, str]] = [] for name in sorted(projects): - # Reap any active MLflow daemon so remove_project's .runtime - # safety check doesn't block bulk teardown. - try: - stop_services(name) - except Exception as e: - logger.debug("stop_services(%s) raised: %s", name, e) try: remove_project(name, keep_files=args.keep_files) verb_past = "Unregistered" if args.keep_files else "Removed" @@ -375,379 +547,6 @@ def _cmd_rm_all(args) -> int: return 0 -def _cmd_setup_kb(args): - """Build the core knowledge base collections.""" - from dsagt.commands.setup_core_kb import run_setup_kb - - run_setup_kb(args) - - -def _cmd_mlflow(args): - """Run MLflow in the foreground. - - Pins the port from the project's internal config so MCP servers - (which bake the URL into their artifacts at init time) agree on - where traces land. Reaps any prior MLflow we left behind on this - project before binding so a stale leftover doesn't block the new - one. If the port is still busy after reap, surfaces the offender - via ``lsof`` so the user knows what to kill. - - With ``--background-only``: idempotent fast-path used by the launch - shim. If MLflow is already running on this project's pinned port, - do nothing and exit 0; otherwise start it and return. - """ - config = load_config(args.project) - pdir = Path(config["project_dir"]) - background_only = getattr(args, "background_only", False) - - port = config.get("mlflow", {}).get("port") - if port is None: - port = pick_free_port() - - # --background-only: short-circuit if MLflow is already up on this port. - # Detect via a TCP probe — the .runtime PID may be stale across machines. - if background_only and _port_responds(port): - return 0 - - # Reap any prior dsagt-spawned MLflow on this project so it doesn't - # hold the pinned port. Same machinery dsagt start uses. - from dsagt.session import reap_runtime - - reaped = reap_runtime(pdir / ".runtime") - for msg in reaped: - print(f" {msg}") - - busy = _port_holder(port) - if busy: - line, pid = busy - print(f" Port {port} held by:") - print(f" {line}") - if pid: - ancestry = _process_ancestry(pid) - if ancestry: - print(f" Process tree: {ancestry}") - if pid and _free_port(port, pid): - print(f" Freed port {port}.") - else: - msg = [ - f"Error: could not free port {port}.", - ] - if pid: - msg.append(f"Try manually: kill -9 {pid}") - else: - msg.append( - "(could not parse PID from lsof output — kill the " - "process shown above by hand)" - ) - msg.append( - "Or re-init with a different port: " - "`dsagt init --mlflow-port `." - ) - print("\n".join(msg), file=sys.stderr) - return 1 - - cmd = mlflow_command(pdir, config.get("mlflow", {}), port=port) - - log_path = pdir / "mlflow.log" - log_fd = open(log_path, "wb") - proc = subprocess.Popen( - cmd, - start_new_session=True, - stdout=log_fd, - stderr=subprocess.STDOUT, - ) - - agent = config["agent"] - session_id = ( - f"{args.project}-" f"{datetime.now(timezone.utc).strftime('%Y%m%dT%H%M%SZ')}" - ) - - runtime_file = pdir / ".runtime" - runtime_file.write_text( - json.dumps( - { - "pids": {"mlflow": proc.pid}, - "ports": {"mlflow": port}, - "session_id": session_id, - "started_at": datetime.now(timezone.utc).isoformat(), - }, - indent=2, - ) - + "\n" - ) - - mlflow_url = f"http://localhost:{port}" - print(" Starting MLflow in background (gunicorn boot ~2-5s)...", flush=True) - experiment_id = _wait_and_resolve_experiment(port, args.project, timeout=20.0) - - if background_only: - # Shim is calling us; it handles the env exports itself. Print one - # confirmation line so the user sees what happened, then return. - print(f" MLflow ready at {mlflow_url} (pid {proc.pid})") - return 0 - - print() - print(f" Project: {args.project}") - print(f" PID: {proc.pid}") - print(f" UI: {mlflow_url}") - if experiment_id is not None: - print(f" Experiment id: {experiment_id}") - else: - print( - f" Experiment id: " - ) - print(f" Session id: {session_id}") - print(f" Logs: {log_path}") - print() - print(" OTel routing for the shell that runs your agent (these go") - print(" straight to the agent's external OTel SDK; project / agent /") - print(f" session_id are read from {pdir}/dsagt_config.yaml") - print(" + .runtime by dsagt's own services, no need to export them):") - print() - # Why http/protobuf + signal-specific TRACES_ENDPOINT: - # - OTel SDKs default to gRPC (port 4317); without this set, claude - # silently tries gRPC and drops every span. - # - Generic OTEL_EXPORTER_OTLP_ENDPOINT auto-appends /v1/traces; we - # use the signal-specific TRACES_ENDPOINT (used as-is) to avoid the - # double-append "...v1/traces/v1/traces" 404. - print(" export OTEL_EXPORTER_OTLP_PROTOCOL=http/protobuf") - print(f" export OTEL_EXPORTER_OTLP_TRACES_ENDPOINT={mlflow_url}/v1/traces") - if experiment_id is not None: - print( - f" export OTEL_EXPORTER_OTLP_HEADERS=" - f'"x-mlflow-experiment-id={experiment_id}"' - ) - print( - f" export OTEL_RESOURCE_ATTRIBUTES=" - f'"service.name={agent},session.id={session_id}"' - ) - - setup = AGENTS[agent]() - if setup.telemetry_env: - print() - print(f" Agent telemetry verbosity for {agent} — without these, OTel") - print(" spans carry only counts/cost/duration; tool_use payloads") - print(" (which memory extraction needs) are absent:") - print() - for k, v in sorted(setup.telemetry_env.items()): - print(f" export {k}={v}") - if agent == "claude": - # File-mode bodies — writes full request/response JSON to - # /api_bodies/, stamps body_ref on the span event so the - # trace links to the on-disk file. =1 (inline) mode would post - # to /v1/logs which MLflow's OTLP receiver returns 404 for, so - # the bodies vanish. See agents/claude.py for the full reasoning. - print(f" export OTEL_LOG_RAW_API_BODIES=file:{pdir}/api_bodies") - - print() - print(" Note: MLflow always creates a 'Default' experiment (id=0) on") - print(" init. It will stay empty; ignore it. Your traces land in the") - print(f" '{args.project}' experiment.") - print() - print(f" To stop MLflow: dsagt stop {args.project}") - print() - return 0 - - -def _process_ancestry(pid: str) -> str: - """Return ``pid (etime) cmd → ppid (etime) cmd → ...`` walking up to PID 1. - - Helps diagnose orphans: a zombie MLflow worker re-parented to PID 1 - after its dsagt parent died shows up as ``... → 1 systemd``, while a - genuinely-still-running dsagt parent shows up by name. Best-effort: - returns empty string if ps isn't available or the chain breaks. - """ - chain: list[str] = [] - cur = pid - seen: set[str] = set() - while cur and cur not in seen and cur != "0": - seen.add(cur) - try: - result = subprocess.run( - ["ps", "-p", cur, "-o", "pid=,ppid=,etime=,command="], - capture_output=True, - text=True, - check=False, - timeout=2.0, - ) - except (FileNotFoundError, subprocess.TimeoutExpired): - return "" - line = result.stdout.strip() - if not line: - break - parts = line.split(None, 3) - if len(parts) < 4: - break - this_pid, ppid, etime, cmd = parts - chain.append(f"{this_pid} ({etime}) {cmd[:60]}") - if cur == "1": - break - cur = ppid - return " → ".join(chain) - - -def _free_port( - port: int, pid: str, term_timeout: float = 3.0, kill_timeout: float = 1.5 -) -> bool: - """SIGTERM *pid*, wait for *port* to free, escalate to SIGKILL if not. - - Returns True if the port is free at the end. MLflow's gunicorn parent - usually releases the socket on SIGTERM after its workers wind down; - a stuck worker needs SIGKILL. We don't gate on process name — the - contract is "the pinned MLflow port is dsagt's; reclaim it" — but we - log what we kill so the user has a record. - """ - try: - target = int(pid) - except ValueError: - return False - - for sig, timeout in ( - (signal.SIGTERM, term_timeout), - (signal.SIGKILL, kill_timeout), - ): - try: - os.kill(target, sig) - except ProcessLookupError: - pass # already gone — still need to confirm port is free - except PermissionError: - return False - deadline = time.monotonic() + timeout - while time.monotonic() < deadline: - if _port_holder(port) is None: - return True - time.sleep(0.2) - return _port_holder(port) is None - - -def _port_responds(port: int) -> bool: - """Return True if something is accepting TCP connections on *port*. - - Used by ``--background-only`` to short-circuit when MLflow is - already up. socket.connect_ex returns 0 on success, errno otherwise. - """ - import socket - - with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: - s.settimeout(0.5) - return s.connect_ex(("127.0.0.1", port)) == 0 - - -def _port_holder(port: int) -> tuple[str, str | None] | None: - """Return ``(lsof_line, pid)`` for the process holding *port*, or None. - - Uses ``lsof`` (macOS + Linux). Returns None when the port is free - or when ``lsof`` isn't available (in which case the bind attempt - will surface the failure anyway). ``pid`` is the second whitespace - field of the lsof line; None if the line couldn't be parsed. - """ - try: - # Combine protocol + port into a single -i filter; multiple -i - # arguments are OR'd, which used to make us match any listening - # TCP socket (e.g., rapportd) regardless of the requested port. - result = subprocess.run( - ["lsof", f"-iTCP:{port}", "-sTCP:LISTEN", "-P", "-n"], - capture_output=True, - text=True, - check=False, - timeout=3.0, - ) - except (FileNotFoundError, subprocess.TimeoutExpired): - return None - lines = [l for l in result.stdout.splitlines() if l and not l.startswith("COMMAND")] - if not lines: - return None - line = lines[0] - parts = line.split() - pid = parts[1] if len(parts) >= 2 and parts[1].isdigit() else None - return (line, pid) - - -def _wait_and_resolve_experiment( - port: int, - project: str, - timeout: float, -) -> str | None: - """Poll MLflow until it answers, then look up / create the experiment. - - Returns the numeric experiment id on success, None on timeout or - error (caller still prints the URL — the user can find the id in - the UI). - """ - import socket - import time - - deadline = time.monotonic() + timeout - while time.monotonic() < deadline: - try: - with socket.create_connection(("127.0.0.1", port), timeout=0.5): - break - except (ConnectionRefusedError, OSError): - time.sleep(0.25) - else: - return None - try: - import mlflow - - mlflow.set_tracking_uri(f"http://localhost:{port}") - return str(mlflow.set_experiment(project).experiment_id) - except Exception as e: - logger.debug("could not resolve experiment id: %s", e) - return None - - -def _stop_one(project: str) -> int: - """Stop services for a single project via the .runtime state file. - - Returns the number of services killed (0 when nothing was running), - so callers can report "nothing to do" cleanly. Idempotent — safe - to call when no project state file exists. - """ - try: - load_config(project) # validates the project is registered + parsable - except (FileNotFoundError, ValueError) as e: - # Registered project whose config is missing or malformed — skip - # rather than abort a multi-project sweep. - print(f" [{project}] skipping: {e}") - return 0 - - msgs = stop_services(project) - for msg in msgs: - print(f" [{project}] {msg}") - if not msgs: - print(f" [{project}] no running services.") - return len(msgs) - - -def _cmd_stop(args): - """Stop running services. - - Without a project argument, sweeps every registered project so the - common "I just want to clean up whatever's running" case doesn't - require remembering which project name was active. With a project, - behaves as before — single-target. - - Output rules: every action gets a line, prefixed with the project - name; when nothing happened anywhere we still print a final summary. - """ - if args.project: - _stop_one(args.project) - return - - projects = list_projects() - if not projects: - print(" No projects registered.") - return - - total = 0 - for name in sorted(projects): - total += _stop_one(name) - - if total == 0: - print(f" Swept {len(projects)} project(s); nothing to stop.") - - def _cmd_info(args): """Triage summary of a project's MLflow traces.""" from dsagt.commands.info import run @@ -755,57 +554,11 @@ def _cmd_info(args): return run(args.project, as_json=args.json) -def _cmd_memory(args): - """Extract episodic memory from accumulated session traces. - - Tracks a high-water-mark timestamp in - ``/.dsagt/extracted_at.json``. Each invocation extracts - traces newer than the mark and updates it. In BYOA mode (no - ``DSAGT_SESSION_ID`` minted by ``dsagt start``), session boundaries - are fuzzy — we batch all unprocessed traces into one extraction. - """ - config = load_config(args.project) - pdir = Path(config["project_dir"]) - state_path = pdir / ".dsagt" / "extracted_at.json" - state_path.parent.mkdir(parents=True, exist_ok=True) +def _cmd_traces(args): + """Open the MLflow trace viewer over the project's store (catch-up first).""" + from dsagt.commands.traces import run - last_extracted = None - if state_path.exists(): - last_extracted = json.loads(state_path.read_text()).get("last_extracted_at") - print(f" Last extraction watermark: {last_extracted}") - else: - print(" No prior extraction recorded — processing all available traces.") - - now = datetime.now(timezone.utc).isoformat() - result = run_extraction(args.project) - status = result.get("status", "unknown") - n_indexed = result.get("tool_use_indexed", 0) - if n_indexed: - print(f" Indexed {n_indexed} tool execution(s) into tool_use") - if status == "ok": - print(f" Extracted {result.get('total_entries', 0)} memories") - state_path.write_text( - json.dumps( - { - "last_extracted_at": now, - "previous": last_extracted, - }, - indent=2, - ) - + "\n" - ) - elif status == "empty": - print(" No new traces to extract") - elif status == "tool_use_only": - print( - " LLM-based memory extraction skipped: set " - "DSAGT_MEMORY_API_KEY and DSAGT_MEMORY_MODEL in your shell " - "to enable. Optional: DSAGT_MEMORY_BASE_URL, " - "DSAGT_MEMORY_PROVIDER." - ) - else: - print(f" Extraction returned status={status}: {result}") - return 0 + return run(args.project, port=args.port) def _cmd_smoke_test(args): @@ -817,8 +570,9 @@ def _cmd_smoke_test(args): With ``--all``, run the harness in parallel for every agent in ``VALID_AGENTS``. Each agent has its own project name (``smoke-test-X``) - so they don't collide on MLflow ports, kb_index, or registry entries. - Output is per-agent log files; the summary prints in finish order. + so they don't collide on the sqlite MLflow store, kb_index, or registry + entries. Output is per-agent log files; the summary prints in finish + order. """ pkg_dir = Path(__file__).resolve().parent.parent.parent.parent script = pkg_dir / "tests" / "smoke_test" / "run.sh" @@ -842,7 +596,6 @@ def _run_smoke_all(script: Path) -> int: progress instead of waiting for the slowest agent before any output. """ import tempfile - import time from concurrent.futures import ThreadPoolExecutor, as_completed agents = list(VALID_AGENTS) @@ -890,29 +643,69 @@ def _run_one(agent: str) -> tuple[str, int, float]: def main(argv=None): + from dsagt import __version__ + + argv = list(sys.argv[1:] if argv is None else argv) + # `dsagt mlflow ` is an unlisted alias for `traces` — the word + # people reach for when they want the MLflow viewer. Rewritten before + # parsing (only the command slot: the first non-flag token) so argparse — + # and therefore --help — only ever knows `traces`. + for i, tok in enumerate(argv): + if tok.startswith("-"): + continue + if tok == "mlflow": + argv[i] = "traces" + break + parser = argparse.ArgumentParser( prog="dsagt", description="DSAgt project and session management." ) parser.add_argument("--verbose", action="store_true") + parser.add_argument("--version", action="version", version=f"dsagt {__version__}") sub = parser.add_subparsers(dest="command") - p_init = sub.add_parser("init", help="Create a new project") - p_init.add_argument("project", help="Project name (human-readable alias)") + p_init = sub.add_parser( + "init", + help="Create or reconfigure a project (interactive; re-runnable)", + ) p_init.add_argument( - "--agent", choices=VALID_AGENTS, required=True, help="Agent platform (required)" + "project", + nargs="?", + help="Project name (also the folder name). Prompted if omitted in a TTY.", ) p_init.add_argument( - "--mlflow-port", - type=int, + "--agent", + choices=VALID_AGENTS, default=None, - help="MLflow port to pin (default: pick a free one). Written to " - "the internal config so MCP servers + dsagt mlflow agree on it.", + help="Agent platform. Prompted interactively; required non-interactively.", ) p_init.add_argument( "--location", default=None, - help="Parent directory for the project (default: ~/dsagt-projects/)", + help="Parent directory for the project (default: ~/dsagt-projects/). " + "Non-interactive path; interactive prompts for the full project dir.", + ) + _kb_sel = p_init.add_mutually_exclusive_group() + _kb_sel.add_argument( + "--include", + nargs="+", + metavar="ASSET", + help="KB assets to provision into the project (or 'all' for " + "everything). Default: bundled tools + the genesis skill catalog.", + ) + _kb_sel.add_argument( + "--exclude", + nargs="+", + metavar="ASSET", + help="Provision the default KB set minus these assets ('all' to " + "create the project with no bundled KB content).", + ) + p_init.add_argument( + "--episodic", + action="store_true", + help="Enable episodic memory (captures session turns into searchable " + "memory). Off by default.", ) p_start = sub.add_parser("start", help="Start a project session") @@ -924,33 +717,13 @@ def main(argv=None): help="Agent platform. Required on first start if init didn't set one; " "thereafter, a per-run override (doesn't update the YAML default).", ) - p_start.add_argument( - "--mlflow-port", - type=int, - default=None, - help="Override the MLflow port from dsagt_config.yaml. Useful when " - "the configured port is permanently taken on your machine.", - ) - p_start.add_argument( - "--enable-proxy", - action="store_true", - help="Spawn dsagt-proxy and route the agent's LLM calls through it. " - "For agents that don't natively emit OTel traces with full " - "LLM-call payloads (cline, roo, codex partial), the proxy is " - "what makes their conversations visible in MLflow at all — " - "every agent turn becomes an inspectable trace (real-time " - "audit, replay, debugging) and memory extraction works as a " - "downstream consequence. Agents with " - "otel_payload_support='full' (claude, goose) emit their own " - "traces and don't need the flag. Port is auto-picked.", - ) p_start.add_argument( "--script", default=None, help="Path to a goose-run instructions file. When set, the agent runs " "non-interactively (GOOSE_MODE=auto) against this script — used by " "the smoke test to share the full dsagt start lifecycle (config " - "generation, services, memory extraction, cleanup) with manual runs.", + "generation, memory extraction) with manual runs.", ) p_start.add_argument( "--max-turns", @@ -959,24 +732,6 @@ def main(argv=None): help="Cap on agent turn count when --script is set (default: 30).", ) - p_mlflow = sub.add_parser( - "mlflow", help="Run MLflow in the foreground against a project's store" - ) - p_mlflow.add_argument("project", help="Project name") - p_mlflow.add_argument( - "--port", - type=int, - default=None, - help="Override the port from dsagt_config.yaml", - ) - p_mlflow.add_argument( - "--background-only", - action="store_true", - help="Start MLflow in background and exit. Skip the OTel-export " - "block (the launch shim handles env setup itself). " - "Idempotent: no-op if MLflow is already running on this project.", - ) - p_info = sub.add_parser( "info", help="Summarize a project's MLflow traces (tokens, errors, by session/source)", @@ -988,23 +743,17 @@ def main(argv=None): help="Emit the structured report as JSON instead of formatted text", ) - p_memory = sub.add_parser( - "memory", - help="Extract episodic memory from accumulated session traces " - "(BYOA: run after one or more agent sessions to populate the KB)", + p_traces = sub.add_parser( + "traces", + help="Open the MLflow trace viewer over a project's store (runs catch-up " + "first, deep-links to the Traces tab, quiets the mlflow noise)", ) - p_memory.add_argument("--project", required=True, help="Project name") - - p_stop = sub.add_parser( - "stop", - help="Stop project services (including orphans on configured ports). " - "Without a project argument, sweeps every registered project.", - ) - p_stop.add_argument( - "project", - nargs="?", - default=None, - help="Project name (omit to sweep all registered projects)", + p_traces.add_argument("project", help="Project name") + p_traces.add_argument( + "--port", + type=int, + default=5000, + help="Port for the local MLflow UI (default: 5000)", ) p_smoke = sub.add_parser( @@ -1025,13 +774,6 @@ def main(argv=None): "logs go to a temp dir; verdicts print in finish order.", ) - p_setup_kb = sub.add_parser( - "setup-kb", help="Build the core knowledge base collections" - ) - from dsagt.commands.setup_core_kb import add_setup_kb_args - - add_setup_kb_args(p_setup_kb) - sub.add_parser("list", help="List all registered projects and their status") p_mv = sub.add_parser("mv", help="Move a project to a new location") @@ -1058,8 +800,12 @@ def main(argv=None): args = parser.parse_args(argv) + # The CLI speaks to the user via ``print()``; library logs are diagnostic. + # Default the console to WARNING so INFO chatter (embedder load, route + # registration, catalog indexing) doesn't bury the init/start output. + # ``--verbose`` opts into the full DEBUG stream. logging.basicConfig( - level=logging.DEBUG if args.verbose else logging.INFO, + level=logging.DEBUG if args.verbose else logging.WARNING, format="%(asctime)s [dsagt] %(message)s", datefmt="%H:%M:%S", ) @@ -1071,12 +817,9 @@ def main(argv=None): cmds = { "init": _cmd_init, "start": _cmd_start, - "mlflow": _cmd_mlflow, - "memory": _cmd_memory, "info": _cmd_info, - "stop": _cmd_stop, + "traces": _cmd_traces, "smoke-test": _cmd_smoke_test, - "setup-kb": _cmd_setup_kb, "list": _cmd_list, "mv": _cmd_mv, "rm": _cmd_rm, diff --git a/src/dsagt/commands/info.py b/src/dsagt/commands/info.py index 63adf46..80f04fc 100644 --- a/src/dsagt/commands/info.py +++ b/src/dsagt/commands/info.py @@ -1,26 +1,32 @@ """ dsagt info — summary of MLflow traces for triage. -A terse, read-only snapshot of what the project's mlflow.db already knows: -counts and token totals grouped by session and by source (agent turn, -embedding, extraction). Errors surface inline so a user can see "which -session broke" without scrolling. Deep investigation still happens in -the MLflow UI (``dsagt mlflow ``); this command is the triage -layer that tells you *where* to look first. +A terse, read-only snapshot of what the project's serverless +``sqlite:////mlflow.db`` store already knows: counts and token +totals grouped by session and by source (agent turn, embedding, +extraction). Errors surface inline so a user can see "which session +broke" without scrolling. Deep investigation still happens in the MLflow +UI (``mlflow ui --backend-store-uri sqlite:////mlflow.db``); this +command is the triage layer that tells you *where* to look first. Aggregation reads ``trace_metadata`` for token totals + session id (MLflow stamps per-trace token usage as a JSON blob under -``mlflow.trace.tokenUsage``; the OTLP receiver promotes our -``session.id`` span attribute to ``mlflow.trace.session``). - -Source bucketing comes from the OTel ``service.name`` resource attribute -on each trace's root span (set per emitting process by ``init_tracing`` -and the agent's own OTel SDK). Possible values: - - ``claude-code`` / ``goose`` / ``cline`` / ``roo`` / ``codex`` — - agent-emitted LLM-call traces (the bulk of traffic) - - ``dsagt-knowledge-server`` / ``dsagt-registry-server`` — MCP server - spans (``kb.*``, ``registry.*``) - - ``dsagt-run`` — tool-execute spans +``mlflow.trace.tokenUsage``; the live tracer stamps ``mlflow.trace.session``). + +Source bucketing reads the metadata DSAGT itself stamps on each trace — no +span inspection: + - ``memory`` / ``skill`` / ``knowledge`` / ``registry`` — internal debug + traces, from the ``dsagt.source`` tag (the MCP tool category the agent + invoked, set on the trace root by the dispatch shell). + - ``execution`` — dsagt-run tool-execute traces (``dsagt.source``). + - ``episodic`` / ``code_use`` — background emitters (per-turn memory + extraction; trace-archive indexing), tagged at their off-thread call site + so their embedding writes don't orphan as untagged ``unknown`` roots. + - ``claude`` / ``goose`` / ``cline`` / ``codex`` — agent traces, from the + ``dsagt.agent`` metadata stamped by ``MLflowSink`` (the bulk of traffic). + +The ``Agent turns`` / ``Internal/debug`` headline splits the recovered agent +conversation from all of DSAGT's own bookkeeping traces (:data:`_INTERNAL_SOURCES`). """ from __future__ import annotations @@ -28,7 +34,6 @@ import json import os import re -import sys from pathlib import Path import yaml @@ -39,7 +44,7 @@ _SECRET_LEAF_KEYS = {"api_key"} # Internal/derived sections — irrelevant for "where does this credential # come from" triage and would just clutter the output. -_CONFIG_SOURCE_SKIP_PREFIXES = ("categories.", "extraction.", "knowledge.") +_CONFIG_SOURCE_SKIP_PREFIXES = ("knowledge.", "skills.") def _mask_secret(value: str) -> str: @@ -60,18 +65,18 @@ def _flatten(d: dict, prefix: str = ""): def _config_sources(project_name: str) -> list[dict]: - """Return per-leaf source info for the project's dsagt_config.yaml. + """Return per-leaf source info for the project's .dsagt/config.yaml. Walks the *raw* YAML (not the env-resolved version) so ${VAR} references are visible. For each leaf, reports where the resolved value came from: ``config`` (literal in YAML), ``shell`` (``${VAR}`` resolved against ``os.environ``), or ``unresolved`` (``${VAR}`` with no value anywhere). - No ``.env`` file is read — dsagt-internal env lives in - ``dsagt_config.yaml`` + ``.runtime``; user-provided shell exports are - the only other source. + No ``.env`` file is read — dsagt-internal config lives in + ``.dsagt/config.yaml``; user-provided shell exports are the only other + source. """ pdir = project_dir(project_name) - raw = yaml.safe_load((pdir / "dsagt_config.yaml").read_text()) or {} + raw = yaml.safe_load((pdir / ".dsagt" / "config.yaml").read_text()) or {} rows: list[dict] = [] for path, value in _flatten(raw): @@ -157,8 +162,9 @@ def _tokens(metadata: dict) -> tuple[int, int]: The value is a JSON string, not a dict — MLflow encodes structured metadata as strings so it round-trips through the same storage path - as arbitrary user tags. Missing key → (0, 0); some traces (e.g. the - hardcoded gpt-4o-mini session-namer mock) legitimately have no usage. + as arbitrary user tags. Missing key → (0, 0); some traces (e.g. an + agent's internal title-gen / session-namer call) legitimately have no + usage. """ raw = metadata.get("mlflow.trace.tokenUsage") if not raw: @@ -175,9 +181,9 @@ def _tokens_from_spans(spans) -> tuple[int, int]: Native claude OTel emission stamps these as span attributes on ``claude_code.llm_request`` spans (string-encoded — MLflow's OTLP - receiver JSON-encodes every attribute). ``mlflow.trace.tokenUsage`` - is only set by LiteLLM autolog (proxy mode); for BYOA we aggregate - from the spans themselves so the totals aren't always zero. + receiver JSON-encodes every attribute). When ``mlflow.trace.tokenUsage`` + isn't present on the trace metadata we aggregate from the spans + themselves so the totals aren't always zero. Returns ``(0, 0)`` for non-LLM traces (kb.search, tool.execute) — those don't carry token attributes. @@ -225,116 +231,34 @@ def _fmt_count(n: int) -> str: return f"{n / 1_000_000:.1f}M" -#: Map root-span name prefix → human-readable source bucket. Claude's -#: native OTel emission uses ``claude_code.*``; goose uses ``goose.*`` / -#: ``dispatch_tool_call``; our internal services use the prefixes their -#: ``init_tracing`` calls assign. Used as the canonical bucket because -#: MLflow's OTLP receiver doesn't surface the ``service.name`` resource -#: attribute on the span data ``mlflow.search_traces`` returns — the -#: name-prefix mapping is reliable across MLflow versions and gives the -#: same answer (which emitting process produced this trace). -_SPAN_NAME_TO_SOURCE: tuple[tuple[str, str], ...] = ( - ("claude_code.", "claude"), - ("goose.", "goose"), - # Goose's Rust runtime emits root spans named after agent/provider - # operations (no namespace prefix): ``reply`` (agent turn), - # ``complete_with_model`` (provider LLM call), ``dispatch_tool_call`` - # (tool invocation). Verified against goose 1.7+ spans in - # crates/goose/src/agents/agent.rs and providers/openai.rs. - ("dispatch_tool_call", "goose"), - ("complete_with_model", "goose"), - ("reply", "goose"), - # dsagt-proxy spans emitted via LiteLLM autolog. The proxy receives - # the agent's request as ``Received Proxy Server Request`` (LiteLLM's - # FastAPI handler span) and forwards via ``litellm_request`` (the - # actual upstream call). Both carry mlflow.spanInputs/Outputs, so - # the UI's request/response columns light up. - ("Received Proxy Server Request", "dsagt-proxy"), - ("litellm_request", "dsagt-proxy"), - # LiteLLM emits ``proxy_pre_call`` spans for pre-call hooks (our - # DSAGTCallback's cache-breakpoint injection runs here). These get - # their own trace_id but with parent_span_id pointing outside the - # trace, so the no-root-span fallback in _source_from_spans needs - # this name registered too. - ("proxy_pre_call", "dsagt-proxy"), - ("kb.", "dsagt-knowledge-server"), - ("registry.", "dsagt-registry-server"), - ("tool.execute", "dsagt-run"), +#: The ``dsagt.source`` categories — internal (debug) traces DSAGT emits, as +#: opposed to the recovered agent-conversation traces. MCP tool categories +#: (``memory`` / ``skill`` / ``knowledge`` / ``registry``), ``execution`` for +#: dsagt-run, plus the two background emitters: ``episodic`` (per-turn memory +#: extraction) and ``code_use`` (the trace-archive indexer). +_INTERNAL_SOURCES = frozenset( + {"memory", "skill", "knowledge", "registry", "execution", "episodic", "code_use"} ) -def _bucket_from_span_name(name: str | None) -> str | None: - if not name: - return None - for prefix, bucket in _SPAN_NAME_TO_SOURCE: - if name.startswith(prefix): - return bucket - return None - - -def _source_from_spans(spans) -> str | None: - """Bucket a trace by who emitted it. - - Tries ``service.name`` on span attributes / resource first (works for - proxy-mode traces and dsagt-internal services that stamp it - explicitly), then falls back to mapping the root span's NAME prefix - to a source bucket — necessary for native claude / goose OTel because - MLflow's OTLP receiver drops the ``service.name`` resource attribute - in the data ``mlflow.search_traces`` exposes. +def _row_source_for(tags: dict, metadata: dict) -> str: + """Bucket a trace by who emitted it, from the metadata DSAGT itself stamps. - ``mlflow.search_traces`` returns a ``spans`` column whose entries - vary in shape across MLflow versions (Span object, dict, or pandas - Series of dicts); we defend against both attribute and item access. + No span inspection: internal debug traces carry an explicit ``dsagt.source`` + tag (see :data:`_INTERNAL_SOURCES`), set on the trace root by the MCP + dispatch shell / ``code_execute_span`` / the background emitters. Agent + traces carry ``dsagt.agent`` metadata, stamped by ``MLflowSink``. Neither + overlaps, so the bucket is a direct lookup; anything else (a stray trace + with neither — which should no longer happen now background work is tagged) + is ``"unknown"``. """ - if spans is None: - return None - try: - for span in spans: - attrs = getattr(span, "attributes", None) - if attrs is None and isinstance(span, dict): - attrs = span.get("attributes") - if attrs and "service.name" in attrs: - return attrs["service.name"] - resource = getattr(span, "resource", None) or ( - span.get("resource") if isinstance(span, dict) else None - ) - if resource: - rattrs = getattr(resource, "attributes", None) or ( - resource.get("attributes") - if isinstance(resource, dict) else None - ) - if rattrs and "service.name" in rattrs: - return rattrs["service.name"] - # No service.name anywhere — fall back to span-name-prefix - # bucketing on the root span (the one without a parent) first. - for span in spans: - parent = ( - getattr(span, "parent_id", None) - or (span.get("parent_id") if isinstance(span, dict) else None) - or getattr(span, "parent_span_id", None) - or (span.get("parent_span_id") if isinstance(span, dict) else None) - ) - if parent: - continue - name = getattr(span, "name", None) or ( - span.get("name") if isinstance(span, dict) else None - ) - bucket = _bucket_from_span_name(name) - if bucket: - return bucket - # No root present (rare — LiteLLM's proxy_pre_call spans have - # parent_span_id pointing outside their own trace). Try any - # span whose name we recognize. - for span in spans: - name = getattr(span, "name", None) or ( - span.get("name") if isinstance(span, dict) else None - ) - bucket = _bucket_from_span_name(name) - if bucket: - return bucket - except (TypeError, AttributeError): - return None - return None + src = (tags or {}).get("dsagt.source") + if src: + return src + agent = (metadata or {}).get("dsagt.agent") + if agent and agent != "-": + return agent + return "unknown" def _is_error(state) -> bool: @@ -363,6 +287,7 @@ def _project_created(pdir: Path) -> str | None: if not ts: return None from datetime import datetime, timezone + return datetime.fromtimestamp(ts, tz=timezone.utc).strftime("%Y-%m-%d") @@ -393,14 +318,48 @@ def _kb_collections(pdir: Path) -> list[dict]: src = None if src: sources[src] = sources.get(src, 0) + 1 - rows.append({ - "collection": sub.name, - "chunks": n_chunks, - "by_source": sources, - }) + rows.append( + { + "collection": sub.name, + "chunks": n_chunks, + "by_source": sources, + } + ) return rows +def _skills(pdir: Path) -> list[dict]: + """Installed + bundled skills for the project. + + Reads the project's ``skills/`` plus the bundled skill dirs via + ``SkillRegistry`` (no embedder needed — this is a directory scan, not a + search). Returns ``[{"name", "description"}, ...]``; empty on any + failure so the report never crashes on a malformed skill. + """ + try: + from dsagt.registry import SkillRegistry + + skills = SkillRegistry(runtime_dir=pdir, kb=None).list_skills() + except Exception: + return [] + return [ + {"name": s.get("name", "?"), "description": s.get("description") or ""} + for s in skills + ] + + +def _print_skills(rows: list[dict]) -> None: + """Render the installed/bundled skill list (name — truncated description).""" + if not rows: + return + name_w = max(len(r["name"]) for r in rows) + print(f"Skills ({len(rows)}):") + for r in rows: + desc = r["description"][:80] + print(f" {r['name']:<{name_w}} {desc}") + print() + + def _kb_retrieval(traces) -> list[dict]: """Per-session ``kb.search`` activity pulled from MLflow trace spans. @@ -424,9 +383,11 @@ def _kb_retrieval(traces) -> list[dict]: ) if name != "kb.search": continue - attrs = getattr(span, "attributes", None) or ( - span.get("attributes") if isinstance(span, dict) else None - ) or {} + attrs = ( + getattr(span, "attributes", None) + or (span.get("attributes") if isinstance(span, dict) else None) + or {} + ) try: hits = int(attrs.get("hits", 0)) except (TypeError, ValueError): @@ -443,8 +404,10 @@ def _kb_retrieval(traces) -> list[dict]: def _load_traces(mlflow_db: Path, project_name: str): """Return (traces_df, experiment_id_or_none). - Separate from the main reporting logic so the caller can decide what to - print when the experiment doesn't exist yet (new project, never run). + Reads the serverless ``sqlite:////mlflow.db`` store directly — no + server required. Separate from the main reporting logic so the caller + can decide what to print when the experiment doesn't exist yet (new + project, never run). """ import mlflow @@ -453,7 +416,8 @@ def _load_traces(mlflow_db: Path, project_name: str): if exp is None: return None, None traces = mlflow.search_traces( - locations=[exp.experiment_id], max_results=5000, + locations=[exp.experiment_id], + max_results=5000, ) return traces, exp.experiment_id @@ -461,7 +425,9 @@ def _load_traces(mlflow_db: Path, project_name: str): def _report(project_name: str, config: dict, traces) -> dict: """Build the structured report dict. CLI formats it; --json prints it.""" agent_header = config.get("agent", "-") - model_header = config.get("llm", {}).get("model", "-") + # BYOA: dsagt no longer records the agent's LLM model (the agent talks to + # its provider directly). Surface the embedding model dsagt configures. + model_header = config.get("embedding", {}).get("model", "-") if traces is None or traces.empty: return { @@ -482,24 +448,25 @@ def _report(project_name: str, config: dict, traces) -> dict: # column once up front keeps pandas from re-parsing the dict on every # groupby. md = traces["trace_metadata"].apply(lambda m: m or {}) + tags = ( + traces["tags"].apply(lambda t: t or {}) + if "tags" in traces.columns + else md.apply(lambda _: {}) + ) session = md.apply(lambda m: m.get("mlflow.trace.session") or "(no-session)") agent = md.apply(lambda m: m.get("dsagt.agent") or "-") errored = traces["state"].apply(_is_error) - # Source + tokens both walk the spans column. Fall back to span data - # because MLflow's OTLP receiver doesn't surface ``service.name`` (so - # metadata-only bucketing returns "unknown" for everything) and only - # LiteLLM-autolog traces carry ``mlflow.trace.tokenUsage`` (so token - # totals are zero in BYOA mode without this fallback). + # Source comes from the metadata DSAGT stamps (dsagt.source tag / + # dsagt.agent), no span inspection. Tokens still walk the spans column as + # a fallback because not all traces carry ``mlflow.trace.tokenUsage``. spans_col = traces["spans"] if "spans" in traces.columns else None def _row_source(idx: int) -> str: - if spans_col is not None: - return _source_from_spans(spans_col.iloc[idx]) or "unknown" - return "unknown" + return _row_source_for(tags.iloc[idx], md.iloc[idx]) def _row_tokens(idx: int) -> tuple[int, int]: - # Trace metadata first (proxy / autolog shape), then aggregate + # Trace metadata first (when tokenUsage is present), then aggregate # span attrs (native claude OTel shape). m = md.iloc[idx] i, o = _tokens(m) @@ -537,7 +504,9 @@ def _group_rows(col: str, sort_by_recency: bool = False) -> list[dict]: "errors": int(g["_err"].sum()), } if col == "_session": - row["agent"] = g["_agent"].mode().iloc[0] if not g["_agent"].empty else "-" + row["agent"] = ( + g["_agent"].mode().iloc[0] if not g["_agent"].empty else "-" + ) row["latest_request_time"] = g["request_time"].max() rows.append(row) if sort_by_recency: @@ -552,12 +521,14 @@ def _group_rows(col: str, sort_by_recency: bool = False) -> list[dict]: # the request column is the display-friendly form. For an error we # just need "which session, which source, when" — the UI has the # payload. - errors.append({ - "session": row["_session"], - "source": row["_source"], - "request_time": row["request_time"], - "trace_id": row["trace_id"], - }) + errors.append( + { + "session": row["_session"], + "source": row["_source"], + "request_time": row["request_time"], + "trace_id": row["trace_id"], + } + ) return { "project": project_name, @@ -576,8 +547,8 @@ def _group_rows(col: str, sort_by_recency: bool = False) -> list[dict]: def _print_text(r: dict) -> None: print(f"Project: {r['project']}") - print(f" Agent: {r['agent']}") - print(f" Model: {r['model']}") + print(f" Agent: {r['agent']}") + print(f" Embedding: {r['model']}") if r.get("created"): print(f" Started: {r['created']}") print() @@ -587,6 +558,7 @@ def _print_text(r: dict) -> None: _print_config_sources(config_sources) _print_kb_collections(r.get("kb_collections") or []) + _print_skills(r.get("skills") or []) if r["total_traces"] == 0: print("No traces recorded yet (run `dsagt start` to create a session).") @@ -597,9 +569,23 @@ def _print_text(r: dict) -> None: f"Totals ({r['total_traces']} traces across {n_sessions} " f"session{'s' if n_sessions != 1 else ''}):" ) - print(f" Tokens: {_fmt_count(r['input_tokens'])} in / " - f"{_fmt_count(r['output_tokens'])} out") + print( + f" Tokens: {_fmt_count(r['input_tokens'])} in / " + f"{_fmt_count(r['output_tokens'])} out" + ) print(f" Errors: {r['total_errors']}") + # Split agent turns (the substance) from DSAGT's own internal/debug traces + # (embedding, indexing, tool spans) so the headline isn't dominated by + # bookkeeping. ``unknown`` counts as internal/debug — it's an orphaned + # DSAGT span, not an agent turn (and should be empty now background work is + # tagged); only genuine agent-conversation traces count as agent turns. + agent_traces = sum( + row["traces"] + for row in r["by_source"] + if row["source"] not in _INTERNAL_SOURCES and row["source"] != "unknown" + ) + internal_traces = r["total_traces"] - agent_traces + print(f" Agent turns: {agent_traces} Internal/debug: {internal_traces}") print() print("By source:") @@ -633,14 +619,15 @@ def _print_text(r: dict) -> None: def run(project: str, as_json: bool) -> int: - # Resolve ${ENV_VAR} references so the header shows the actual model - # name the proxy will route (not the placeholder from dsagt_config.yaml). + # Resolve ${ENV_VAR} references so the header shows resolved values + # (not ${VAR} placeholders from .dsagt/config.yaml). config = resolve_env_vars(load_config(project)) pdir = Path(config["project_dir"]) - mlflow_db = pdir / "mlflow" / "mlflow.db" + mlflow_db = pdir / "mlflow.db" sources = _config_sources(project) kb_collections = _kb_collections(pdir) + skills = _skills(pdir) created = _project_created(pdir) if not mlflow_db.exists(): @@ -650,7 +637,7 @@ def run(project: str, as_json: bool) -> int: r = { "project": project, "agent": config.get("agent", "-"), - "model": config.get("llm", {}).get("model", "-"), + "model": config.get("embedding", {}).get("model", "-"), "created": created, "total_traces": 0, "total_errors": 0, @@ -661,6 +648,7 @@ def run(project: str, as_json: bool) -> int: "errors": [], "kb_retrieval": [], "kb_collections": kb_collections, + "skills": skills, "config_sources": sources, } if as_json: @@ -673,6 +661,7 @@ def run(project: str, as_json: bool) -> int: r = _report(project, config, traces) r["created"] = created r["kb_collections"] = kb_collections + r["skills"] = skills r["config_sources"] = sources if as_json: diff --git a/src/dsagt/commands/knowledge_server.py b/src/dsagt/commands/knowledge_server.py deleted file mode 100644 index dd970cd..0000000 --- a/src/dsagt/commands/knowledge_server.py +++ /dev/null @@ -1,925 +0,0 @@ -""" -DSAgt Knowledge Base MCP Server. - -Provides semantic search over document collections for MCP-compatible agents. - -At startup, symlinks base indexes into a session-specific runtime directory. -All modifications (ingestion, append) happen in the runtime copy. - -Long-running operations (ingest, append) run in the background and return -immediately with a job_id. Use kb_job_status to poll for completion. - -Server configuration (chunk_size, vector_db, rerank) is read from the -project's dsagt_config.yaml. Embedding credentials flow through env vars -(LLM_API_KEY, OPENAI_BASE_URL, EMBEDDING_MODEL) set by dsagt start. - -Usage: - dsagt-knowledge-server --base-index-dir ./kb_index --runtime-dir ./runtime -""" - -import asyncio -import json -import logging -import os -import time -from dataclasses import dataclass, field -from functools import partial - -# Prevent fatal OpenMP crash when multiple libraries (FAISS, PyTorch/ -# sentence-transformers) each bundle their own libomp. -os.environ.setdefault("KMP_DUPLICATE_LIB_OK", "TRUE") - -import uuid -from pathlib import Path - -import mcp.server.stdio -import mcp.types as types -import yaml -from mcp.server.lowlevel import Server, NotificationOptions -from mcp.server.models import InitializationOptions - -from dsagt.knowledge import EMBEDDER_REGISTRY, VECTORINDEX_REGISTRY, CollectionRoute, KnowledgeBase -from dsagt.memory import SuggestionQueue -from dsagt.memory import ExplicitMemory -from dsagt.session import REGISTRY_DIR, _collection_exists, setup_runtime_kb # noqa: F401 - -logger = logging.getLogger(__name__) - - -# --------------------------------------------------------------------------- -# MCP server helpers -# --------------------------------------------------------------------------- - -async def _run_stdio(server: Server, name: str) -> None: - async with mcp.server.stdio.stdio_server() as (read_stream, write_stream): - await server.run( - read_stream, - write_stream, - InitializationOptions( - server_name=name, - server_version="0.1.0", - capabilities=server.get_capabilities( - notification_options=NotificationOptions(), - experimental_capabilities={}, - ), - ), - ) - - - - -# _collection_exists and setup_runtime_kb live in dsagt.session (imported above). - - -def _register_external_collection( - kb: KnowledgeBase, - collection_name: str, - vector_db: str, - connection_params: dict, - embedding_model: str, - description: str, -) -> None: - """Wire an already-built external vector store into the routing registry.""" - coll_dir = kb.index_dir / collection_name - coll_dir.mkdir(exist_ok=True) - - if description: - (coll_dir / "DESCRIPTION.md").write_text(description) - - if vector_db == "chroma": - index_kwargs = { - "collection_name": connection_params.get("collection", collection_name), - "persist_dir": None, - "host": connection_params.get("host", "localhost"), - "port": connection_params.get("port", 8000), - } - elif vector_db == "lancedb": - index_kwargs = { - "uri": connection_params["uri"], - "table": connection_params.get("table", collection_name), - } - elif vector_db == "qdrant": - index_kwargs = { - "url": connection_params["url"], - "collection": connection_params.get("collection", collection_name), - "api_key": connection_params.get("api_key"), - } - else: - raise ValueError( - f"Unsupported vector DB '{vector_db}'. " - f"Choose from: chroma, lancedb, qdrant" - ) - - route = CollectionRoute( - embedding_backend="api", - vector_db=vector_db, - embedder_kwargs={"model": embedding_model}, - index_kwargs=index_kwargs, - description=description, - ) - kb.register_route(collection_name, route) - - - - -# --------------------------------------------------------------------------- -# Background job tracker -# --------------------------------------------------------------------------- - -@dataclass -class _JobTracker: - """Tracks background ingest/append jobs and their completion state.""" - - jobs: dict[str, dict] = field(default_factory=dict) - active_collections: set[str] = field(default_factory=set) - - def start(self, coro, collection: str | None = None) -> str: - job_id = uuid.uuid4().hex[:8] - self.jobs[job_id] = { - "status": "running", - "result": None, - "error": None, - "collection": collection, - "started_at": time.monotonic(), - "message": "Starting -- embedding documents via API...", - } - if collection: - self.active_collections.add(collection) - - tracker = self # capture for the closure - - async def _run(): - try: - tracker.jobs[job_id]["message"] = "Embedding and indexing documents..." - result = await coro - tracker.jobs[job_id]["status"] = "complete" - tracker.jobs[job_id]["result"] = result - tracker.jobs[job_id]["message"] = "Done." - except Exception as e: - import traceback - tb = traceback.format_exc() - tracker.jobs[job_id]["status"] = "error" - tracker.jobs[job_id]["error"] = f"{type(e).__name__}: {e}" - tracker.jobs[job_id]["message"] = f"Failed: {type(e).__name__}: {e}" - tracker.jobs[job_id]["traceback"] = tb - logger.error("Job %s failed: %s\n%s", job_id, e, tb) - finally: - if collection: - tracker.active_collections.discard(collection) - - asyncio.get_event_loop().create_task(_run()) - return job_id - - -# --------------------------------------------------------------------------- -# Per-tool handlers (module-level, explicit dependencies) -# -# Each handler takes ``arguments: dict`` plus its dependencies as keyword -# args bound via functools.partial in create_knowledge_server(). Handlers -# return a result dict; the outer call_tool wrapper JSON-serializes it. -# --------------------------------------------------------------------------- - -async def _handle_kb_list_collections(arguments: dict, *, kb: KnowledgeBase) -> dict: - collections = await asyncio.to_thread(kb.list_collections) - return {"status": "ok", "collections": collections, "count": len(collections)} - - -async def _handle_kb_search( - arguments: dict, *, kb: KnowledgeBase, -) -> dict: - query = arguments["query"] - top_k = arguments.get("top_k", 5) - rerank = arguments.get("rerank") # None → kb.default_rerank - - collection_arg = arguments.get("collection") - collections_arg = arguments.get("collections") - - if not collection_arg and not collections_arg: - return {"status": "error", "error": "Provide 'collection' or 'collections'"} - - # Build ChromaDB where clause from the filter arguments. ChromaDB - # requires single-filter dicts or $and-wrapped lists; an empty dict - # would be invalid, so we only pass where when there are real filters. - where = { - key: arguments[key] - for key in ("category", "session_id", "source_type", "tool_name") - if arguments.get(key) is not None - } - return_code = arguments.get("return_code") - if return_code is not None: - where["return_code"] = int(return_code) - if len(where) > 1: - where = {"$and": [{k: v} for k, v in where.items()]} - - target_collections = collections_arg or [collection_arg] - all_results = [] - search_errors = [] - - for coll_name in target_collections: - try: - search_kwargs = dict(query=query, collection=coll_name, top_k=top_k, rerank=rerank) - if where: - search_kwargs["where"] = where - coll_results = await asyncio.to_thread(kb.search, **search_kwargs) - all_results.extend(coll_results) - except ValueError as e: - logger.warning("Search failed for '%s': %s", coll_name, e) - search_errors.append(str(e)) - - if search_errors and not all_results: - if len(target_collections) == 1: - return {"status": "error", "error": search_errors[0]} - return {"status": "error", "error": f"All collections failed: {'; '.join(search_errors)}"} - - score_key = "rerank_score" if rerank else "score" - all_results.sort(key=lambda r: r.get(score_key, r["score"]), reverse=True) - all_results = all_results[:top_k] - - result = { - "status": "ok", - "query": query, - "collection": collection_arg or ",".join(collections_arg), - "result_count": len(all_results), - "results": [ - { - "text": r["chunk"]["text"], - "score": r["score"], - "rerank_score": r.get("rerank_score"), - "source_file": r["chunk"]["metadata"].get("source_file", ""), - "chunk_index": r["chunk"]["metadata"].get("chunk_index", 0), - "metadata": { - k: v for k, v in r["chunk"]["metadata"].items() - if k not in ("source_file", "chunk_index", "collection", "file_type") - }, - } - for r in all_results - ], - } - if search_errors: - result["warnings"] = search_errors - return result - - -async def _handle_kb_ingest( - arguments: dict, *, kb: KnowledgeBase, job_tracker: _JobTracker, -) -> dict: - folder_path = Path(arguments["folder_path"]) - collection_name = arguments.get("collection_name") - file_types = arguments.get("file_types") - embedding_backend = arguments.get("embedding_backend") - embedding_model = arguments.get("embedding_model") - vector_db = arguments.get("vector_db") - - if not folder_path.exists(): - return {"status": "error", "error": f"Folder not found: {folder_path}"} - if not folder_path.is_dir(): - return {"status": "error", "error": f"Not a directory: {folder_path}"} - - target_name = collection_name or folder_path.name - warning = None - - if target_name in job_tracker.active_collections: - return { - "status": "error", - "error": ( - f"Collection '{target_name}' is already being ingested. " - f"Poll kb_job_status for progress." - ), - } - - if _collection_exists(kb.index_dir / target_name): - source_path = kb.index_dir / target_name / "source.txt" - existing_source = source_path.read_text().strip() if source_path.exists() else None - same_source = ( - existing_source is None - or Path(existing_source).resolve() == folder_path.resolve() - ) - if not same_source: - original_name = target_name - n = 1 - while ( - _collection_exists(kb.index_dir / target_name) - or target_name in job_tracker.active_collections - ): - target_name = f"{original_name}{n}" - n += 1 - warning = ( - f"Collection '{original_name}' already exists from a " - f"different folder; using '{target_name}'." - ) - - route = None - if embedding_backend or embedding_model or vector_db: - default = kb._default_route - inherited_model = embedding_model or default.embedder_kwargs.get("model") - route = CollectionRoute( - embedding_backend=embedding_backend or default.embedding_backend, - vector_db=vector_db or default.vector_db, - embedder_kwargs={"model": inherited_model} if inherited_model else {}, - ) - - ingest_kwargs: dict = {"collection_name": target_name} - if file_types: - ingest_kwargs["file_types"] = file_types - if route is not None: - ingest_kwargs["route"] = route - - async def _ingest_with_logging(): - import traceback as _tb - logger.info("Ingest starting: collection=%s folder=%s kwargs=%s", - target_name, folder_path, ingest_kwargs) - try: - result = await asyncio.to_thread(kb.ingest, folder_path, **ingest_kwargs) - logger.info("Ingest complete: %s", result) - return result - except Exception as _e: - logger.error("Ingest FAILED: %s\n%s", _e, _tb.format_exc()) - raise - - job_id = job_tracker.start(_ingest_with_logging(), collection=target_name) - result = { - "status": "started", - "job_id": job_id, - "collection": target_name, - "message": ( - f"Ingestion started. " - f"Poll kb_job_status(job_id='{job_id}') every 10 seconds. " - f"DO NOT call ingest again -- the job is running in the " - f"background. Large folders may take several minutes." - ), - } - if warning: - result["warning"] = warning - return result - - -async def _handle_kb_append( - arguments: dict, *, kb: KnowledgeBase, job_tracker: _JobTracker, -) -> dict: - collection = arguments["collection"] - paths = arguments["paths"] - if isinstance(paths, str): - paths = [paths] - file_types = arguments.get("file_types") - - if not _collection_exists(kb.index_dir / collection): - return {"status": "error", "error": f"Collection '{collection}' not found"} - - append_kwargs: dict = {} - if file_types: - append_kwargs["file_types"] = file_types - - job_id = job_tracker.start( - asyncio.to_thread(kb.append, collection, paths, **append_kwargs), - collection=collection, - ) - return { - "status": "started", - "job_id": job_id, - "collection": collection, - "message": f"Append started. Poll kb_job_status(job_id='{job_id}') for progress.", - } - - -async def _handle_kb_add_vector_db(arguments: dict, *, kb: KnowledgeBase) -> dict: - collection_name = arguments["collection_name"] - vector_db = arguments["vector_db"] - connection_params = arguments["connection_params"] - embedding_model = arguments["embedding_model"] - description = arguments.get("description", "") - - if (kb.index_dir / collection_name).exists(): - return { - "status": "error", - "error": ( - f"Collection '{collection_name}' already exists. " - "Choose a different name or delete the existing collection." - ), - } - - await asyncio.to_thread( - _register_external_collection, - kb, collection_name, vector_db, - connection_params, embedding_model, description, - ) - return { - "status": "ok", - "collection": collection_name, - "vector_db": vector_db, - "embedding_model": embedding_model, - "message": ( - f"External collection '{collection_name}' registered. " - "Use search to query it." - ), - } - - -async def _handle_kb_job_status(arguments: dict, *, job_tracker: _JobTracker) -> dict: - job_id = arguments["job_id"] - if job_id not in job_tracker.jobs: - return {"status": "error", "error": f"Unknown job: {job_id}"} - - job = job_tracker.jobs[job_id] - elapsed = int(time.monotonic() - job["started_at"]) - result = { - "status": job["status"], - "elapsed_seconds": elapsed, - "message": job.get("message", ""), - } - if job["status"] == "running": - result["instruction"] = ( - "Job is still running. DO NOT call ingest again. " - "Keep polling job_status every 10 seconds until " - "status is 'complete' or 'error'." - ) - if job["result"] is not None: - result["result"] = job["result"] - if job["error"] is not None: - result["error"] = job["error"] - if job.get("traceback") and job["status"] == "error": - result["traceback"] = job["traceback"] - return result - - -async def _handle_kb_remember( - arguments: dict, - *, - kb: KnowledgeBase, - memory: ExplicitMemory, - suggestions: SuggestionQueue, -) -> dict: - text = arguments["text"] - category = arguments.get("category", "") - session_id = arguments.get("session_id", "") - supersedes = arguments.get("supersedes") - promoted_from = arguments.get("promoted_from") - - store_result = await asyncio.to_thread( - memory.remember, - text=text, - category=category, - session_id=session_id, - supersedes=supersedes, - ) - - if not store_result.get("stored"): - return { - "status": "error", - "error": store_result.get("error", "Failed to store memory"), - } - - await asyncio.to_thread( - kb.add_entries, - texts=[text], - collection="session_memory", - metadatas=[{ - "source_type": "explicit_memory", - "category": category, - "session_id": session_id, - }], - ) - - if promoted_from: - suggestions.dismiss(promoted_from) - - return { - "status": "ok", - "entry_id": store_result["entry_id"], - "superseded_id": store_result.get("superseded_id"), - "promoted_from": promoted_from, - "total_memories": await asyncio.to_thread(memory.count), - } - - -async def _handle_kb_get_memories( - arguments: dict, *, memory: ExplicitMemory, suggestions: SuggestionQueue, -) -> dict: - entries = await asyncio.to_thread(memory.get_all) - pending = suggestions.get_all() - result = {"status": "ok", "count": len(entries), "memories": entries} - if pending: - result["suggestions"] = pending - result["suggestion_count"] = len(pending) - return result - - -async def _handle_kb_get_suggestions( - arguments: dict, *, suggestions: SuggestionQueue, -) -> dict: - pending = suggestions.get_all() - return {"status": "ok", "count": len(pending), "suggestions": pending} - - -async def _handle_kb_dismiss_suggestion( - arguments: dict, *, suggestions: SuggestionQueue, -) -> dict: - suggestion_id = arguments["suggestion_id"] - dismissed = suggestions.dismiss(suggestion_id) - if not dismissed: - return {"status": "error", "error": f"Suggestion not found: {suggestion_id}"} - return {"status": "ok", "dismissed": suggestion_id, "remaining": suggestions.count} - - -# --------------------------------------------------------------------------- -# Server factory (thin wiring — used by main() and tests) -# --------------------------------------------------------------------------- - -def create_knowledge_server( - kb: KnowledgeBase, - runtime_dir: str | Path | None = None, -): - """Create and configure the MCP knowledge server. - - This is the test-facing API: tests call it with a mock KB and get back - a server they can drive via call_tool_sync(). main() reads the project - config and constructs KB before calling this. - - The rerank default is on ``kb.default_rerank`` (set from - ``knowledge.rerank`` in dsagt_config.yaml). - """ - server = Server("knowledge") - - mem_dir = Path(runtime_dir) if runtime_dir else kb.index_dir.parent - memory = ExplicitMemory(runtime_dir=mem_dir) - suggestions = SuggestionQueue(mem_dir / "suggestions.json") - job_tracker = _JobTracker() - - handlers = { - "kb_list_collections": partial(_handle_kb_list_collections, kb=kb), - "kb_search": partial(_handle_kb_search, kb=kb), - "kb_ingest": partial(_handle_kb_ingest, kb=kb, job_tracker=job_tracker), - "kb_append": partial(_handle_kb_append, kb=kb, job_tracker=job_tracker), - "kb_add_vector_db": partial(_handle_kb_add_vector_db, kb=kb), - "kb_job_status": partial(_handle_kb_job_status, job_tracker=job_tracker), - "kb_remember": partial(_handle_kb_remember, kb=kb, memory=memory, suggestions=suggestions), - "kb_get_memories": partial(_handle_kb_get_memories, memory=memory, suggestions=suggestions), - "kb_get_suggestions": partial(_handle_kb_get_suggestions, suggestions=suggestions), - "kb_dismiss_suggestion": partial(_handle_kb_dismiss_suggestion, suggestions=suggestions), - } - - @server.list_tools() - async def list_tools() -> list[types.Tool]: - return [ - types.Tool( - name="kb_list_collections", - description=( - "List all available knowledge base collections with their " - "embedding model and vector DB. Use this to discover what " - "documentation is already indexed." - ), - inputSchema={"type": "object", "properties": {}}, - ), - types.Tool( - name="kb_search", - description=( - "Search knowledge base collections using semantic similarity. " - "Returns relevant chunks with source metadata. " - "Supports multi-collection search." - ), - inputSchema={ - "type": "object", - "properties": { - "query": { - "type": "string", - "description": "Natural language search query", - }, - "collection": { - "type": "string", - "description": "Name of a single collection to search", - }, - "collections": { - "type": "array", - "items": {"type": "string"}, - "description": "Search multiple collections and merge results (overrides 'collection')", - }, - "top_k": { - "type": "integer", - "description": "Number of results to return (default: 5)", - "default": 5, - }, - "rerank": { - "type": "boolean", - "description": "Use cross-encoder reranking (slower but more accurate). Default from config.", - "default": kb.default_rerank, - }, - "category": { - "type": "string", - "description": "Filter by category tag (ChromaDB collections only)", - }, - "session_id": { - "type": "string", - "description": "Filter by session ID (ChromaDB collections only)", - }, - "tool_name": { - "type": "string", - "description": "Filter by tool name (ChromaDB collections only)", - }, - "source_type": { - "type": "string", - "description": "Filter by source type (ChromaDB collections only)", - }, - "return_code": { - "type": "integer", - "description": "Filter by tool exit code (ChromaDB collections only)", - }, - }, - "required": ["query"], - }, - ), - types.Tool( - name="kb_ingest", - description=( - "Index a folder as a new knowledge base collection. " - "Returns immediately with a job_id. " - "IMPORTANT: poll kb_job_status every 10 seconds and wait for " - "status='complete'. DO NOT call ingest again for the same " - "folder while a job is running." - ), - inputSchema={ - "type": "object", - "properties": { - "folder_path": { - "type": "string", - "description": "Path to folder containing documents to index", - }, - "collection_name": { - "type": "string", - "description": "Name for the collection (default: folder name)", - }, - "file_types": { - "type": "array", - "items": {"type": "string"}, - "description": "File extensions to include, e.g. ['pdf', 'md', 'py']. Defaults to common types.", - }, - "embedding_backend": { - "type": "string", - "enum": list(EMBEDDER_REGISTRY.keys()), - "description": "Embedding backend override for this collection.", - }, - "embedding_model": { - "type": "string", - "description": "Embedding model override for this collection.", - }, - "vector_db": { - "type": "string", - "enum": list(VECTORINDEX_REGISTRY.keys()), - "description": "Vector database override for this collection.", - }, - }, - "required": ["folder_path"], - }, - ), - types.Tool( - name="kb_append", - description=( - "Add documents to an existing collection. Uses the same embedding " - "model and vector DB the collection was created with. " - "Returns immediately with a job_id -- poll kb_job_status for progress." - ), - inputSchema={ - "type": "object", - "properties": { - "collection": { - "type": "string", - "description": "Name of the existing collection to append to", - }, - "paths": { - "type": "array", - "items": {"type": "string"}, - "description": "List of file or folder paths to add", - }, - "file_types": { - "type": "array", - "items": {"type": "string"}, - "description": "File extensions to include when expanding folders.", - }, - }, - "required": ["collection", "paths"], - }, - ), - types.Tool( - name="kb_add_vector_db", - description=( - "Register an already-built external vector store as a collection. " - "Queries will be embedded via the API using the specified model." - ), - inputSchema={ - "type": "object", - "properties": { - "collection_name": {"type": "string", "description": "Unique name for this collection"}, - "vector_db": {"type": "string", "enum": ["chroma", "lancedb", "qdrant"], "description": "Vector store backend type"}, - "connection_params": {"type": "object", "description": "Backend-specific connection parameters."}, - "embedding_model": {"type": "string", "description": "The API model used to build this index"}, - "description": {"type": "string", "description": "Human-readable description for agent discovery"}, - }, - "required": ["collection_name", "vector_db", "connection_params", "embedding_model"], - }, - ), - types.Tool( - name="kb_job_status", - description="Check the status of a background ingest or append job.", - inputSchema={ - "type": "object", - "properties": { - "job_id": {"type": "string", "description": "Job ID returned by kb_ingest or kb_append"}, - }, - "required": ["job_id"], - }, - ), - types.Tool( - name="kb_remember", - description=( - "Store a user-confirmed fact as an explicit memory. " - "These persist across sessions. Use 'supersedes' to replace an outdated memory." - ), - inputSchema={ - "type": "object", - "properties": { - "text": {"type": "string", "description": "The fact to remember"}, - "category": {"type": "string", "description": "Classification tag"}, - "session_id": {"type": "string", "description": "Current session identifier"}, - "supersedes": {"type": "string", "description": "entry_id of an existing memory this replaces"}, - "promoted_from": {"type": "string", "description": "suggestion_id if promoted from outlier suggestion"}, - }, - "required": ["text"], - }, - ), - types.Tool( - name="kb_get_memories", - description=( - "Get all active explicit memories for this project. " - "Call at session start to load project context." - ), - inputSchema={"type": "object", "properties": {}}, - ), - types.Tool( - name="kb_get_suggestions", - description=( - "Get pending memory suggestions flagged by outlier detection. " - "Present to user for confirmation or dismissal." - ), - inputSchema={"type": "object", "properties": {}}, - ), - types.Tool( - name="kb_dismiss_suggestion", - description="Dismiss a pending memory suggestion.", - inputSchema={ - "type": "object", - "properties": { - "suggestion_id": {"type": "string", "description": "ID of the suggestion to dismiss"}, - }, - "required": ["suggestion_id"], - }, - ), - ] - - @server.call_tool() - async def call_tool(name: str, arguments: dict) -> list[types.TextContent]: - handler = handlers[name] - try: - result = await handler(arguments) - except ValueError as e: - result = {"status": "error", "error": str(e)} - except Exception as e: - logger.exception("Unexpected error in tool '%s'", name) - result = {"status": "error", "error": f"Unexpected error: {e}"} - return [types.TextContent(type="text", text=json.dumps(result, ensure_ascii=False))] - - return server - - -# --------------------------------------------------------------------------- -# Entry point -# --------------------------------------------------------------------------- - -def main(): - """Entry point for dsagt-knowledge-server. - - All configuration comes from the project directory: - - ``./dsagt_config.yaml`` → project path + non-secret settings - (chunk_size, vector_db, rerank) - - ``LLM_API_KEY``, ``OPENAI_BASE_URL`` env vars → embedding credentials - - No CLI arguments — the server derives everything from the YAML. By - contract the agent's launch one-liner is ``cd && ``, - so cwd is project_dir for the MCP children it spawns. - """ - from dsagt.observability import find_project_config - project_dir, _ = find_project_config() - if project_dir is None: - raise RuntimeError( - "dsagt-knowledge-server: no dsagt_config.yaml in cwd " - f"({Path.cwd()}). Launch the agent from the project " - "directory (`cd && `)." - ) - - log_file = project_dir / "dsagt_knowledge_server.log" - # Default INFO; users opt into DEBUG via DSAGT_LOG_LEVEL=DEBUG. See - # registry_server.py main() for rationale (httpcore/urllib3/llama_index - # at DEBUG floods agent debug output). - _level_name = os.environ.get("DSAGT_LOG_LEVEL", "INFO").upper() - _level = getattr(logging, _level_name, logging.INFO) - logging.basicConfig( - level=_level, - format="%(asctime)s %(levelname)s %(name)s: %(message)s", - handlers=[ - logging.FileHandler(log_file, mode="a"), - logging.StreamHandler(), - ], - ) - logger.info("Server starting — project_dir: %s, log: %s", project_dir, log_file) - - # Read project config. Required — this server runs inside a project - # created by dsagt init. Every section must be present with all fields - # filled. dsagt init generates complete defaults; if anything is missing, - # the config is broken and the server fails fast. - config_path = project_dir / "dsagt_config.yaml" - from dsagt.session import resolve_env_vars - config = resolve_env_vars(yaml.safe_load(config_path.read_text())) - - kb_config = config["knowledge"] - emb_config = config["embedding"] - - # Embedding backend selection. Default is "local" (sentence-transformers, - # CPU, no creds) so a fresh ``dsagt init`` works zero-config. Switching - # to "api" requires base_url + api_key — validate eagerly so a misconfig - # surfaces at MCP-server startup rather than at the first kb_search. - backend = (emb_config.get("backend") or "local").lower() - if backend not in ("local", "api"): - raise ValueError( - f"embedding.backend must be 'local' or 'api' (got {backend!r})" - ) - - # Only pass an explicit ``model`` when the user filled one in. Empty / - # ``${EMBEDDING_MODEL}`` placeholders mean "use the backend's default" - # — LocalEmbeddingClient has its own default; APIEmbeddingClient has - # no default and will raise downstream, which is what we want for a - # misconfigured api setup. - # - # Cross-backend leakage guard: HuggingFace identifiers ("org/repo") - # and OpenAI-style aliases ("text-embedding-3-small") share the same - # ``EMBEDDING_MODEL`` env var in most setups. When a user switches - # ``embedding.backend`` from api → local without also retargeting - # the env var, the api alias flows into LocalEmbeddingClient and - # produces a confusing 404 from HuggingFace at first embed. Drop - # the override when it's clearly mis-shaped for the active backend. - raw_model = (emb_config.get("model") or "").strip() - embedder_kwargs: dict = {} - if raw_model and not raw_model.startswith("${"): - looks_hf = "/" in raw_model - if backend == "local" and not looks_hf: - logger.warning( - "Ignoring embedding.model=%r for backend=local (does not " - "look like a HuggingFace identifier). Falling back to the " - "LocalEmbeddingClient default.", - raw_model, - ) - else: - embedder_kwargs["model"] = raw_model - if backend == "api": - base_url = emb_config.get("base_url") or "" - api_key = emb_config.get("api_key") or "" - if not base_url: - raise ValueError( - "embedding.backend='api' requires embedding.base_url in " - "dsagt_config.yaml. Either set it to your OpenAI-compatible " - "endpoint, or change backend to 'local'." - ) - if not api_key or api_key.startswith("${"): - raise ValueError( - "embedding.backend='api' requires embedding.api_key in " - "dsagt_config.yaml. Either fill it in (or export the " - "${EMBEDDING_API_KEY} env var), or change backend to 'local'." - ) - embedder_kwargs.update({"base_url": base_url, "api_key": api_key}) - - from dsagt.observability import init_tracing, configure_litellm_retries - init_tracing("dsagt-knowledge-server") # session_id picked up from DSAGT_SESSION_ID env - configure_litellm_retries() - - runtime_kb_dir = setup_runtime_kb(REGISTRY_DIR / "kb_index", project_dir) - - logger.info("Knowledge backend: %s", backend) - kb = KnowledgeBase( - index_dir=runtime_kb_dir, - chunk_size=kb_config["chunk_size"], - default_rerank=kb_config["rerank"], - default_embedder=backend, - default_index=kb_config["vector_db"], - embedder_kwargs=embedder_kwargs, - ) - # Background-load the embedder so the model is ready when the - # agent's first kb call lands. Without this, the first call pays - # the ~5–10s sentence-transformers import + SentenceTransformer - # construction cost, which looks like a hang to the operator. - kb.preload_default_embedder() - - server = create_knowledge_server(kb, runtime_dir=str(project_dir)) - try: - asyncio.run(_run_stdio(server, "knowledge")) - finally: - kb.close() - - -if __name__ == "__main__": - main() diff --git a/src/dsagt/commands/proxy_server.py b/src/dsagt/commands/proxy_server.py deleted file mode 100644 index 9e9b29c..0000000 --- a/src/dsagt/commands/proxy_server.py +++ /dev/null @@ -1,265 +0,0 @@ -""" -dsagt-proxy: opt-in LiteLLM forwarding proxy for agent transparency. - -DSAgt's normal observability path is the agent emitting OTel directly to -MLflow's ``/v1/traces`` (Claude Code, Goose). For agents that don't -emit OTel with full LLM-call payloads (Cline, Roo Code, Codex -partially), running with ``dsagt start --enable-proxy`` makes the -agent's actions visible at all — every LLM request the agent issues -becomes an MLflow trace you can inspect in real time and replay later: -which model, which messages, which tool_use blocks the assistant emits, -which tool_results came back, full token + cache stats. Without the -proxy, the same agent runs as a black box from DSAgt's perspective — -``dsagt info`` shows only embedding + tool-execute spans, MLflow shows -no agent turns, and end-of-session memory extraction has no -conversation to read. - -Real-time transparency is the primary value. Memory extraction works -as a downstream consequence because the data it needs (request + -response payloads tagged with the session id) lands in MLflow exactly -because the proxy autologged it. - -Architecture: ``init_proxy_tracing()`` installs MLflow's native tracer -provider as the OTel global and plants ``_DSAGTMlflowLogger`` (a -MlflowLogger subclass) into LiteLLM's logger cache. The subclass -stamps ``mlflow.trace.session``, ``dsagt.source=agent``, and -``dsagt.agent`` on every proxy-captured trace in the narrow window -between trace creation and export — giving rich -``mlflow.spanInputs``/``mlflow.spanOutputs`` traces with full request -and response payloads. - -Routing: requests come in on the local port and LiteLLM forwards them -to the user's configured upstream (LLM + embedding) using a minimal -two-route ``model_list`` config. No multi-provider abstraction beyond -what LiteLLM already provides. - -Activation: ``dsagt start --enable-proxy`` sets ``config["proxy"]`` and -``start_services`` spawns this command on a kernel-picked free port. -``agents/__init__.py`` sees the proxy port and overrides the agent's -``ANTHROPIC_BASE_URL`` / ``OPENAI_BASE_URL`` to point at it, plus -plants a sentinel API key so any direct call bypassing the proxy 401s -loudly instead of silently leaking. Without ``--enable-proxy``, this -command is never started — agents talk to their providers directly and -their visibility depends on whether they emit OTel themselves. -""" - -from __future__ import annotations - -import argparse -import logging -import os -import sys -import tempfile - -logger = logging.getLogger(__name__) - - -# Some agents rewrite the requested model name into one of their hardcoded -# "known" Anthropic IDs before sending to /v1/messages — they don't -# recognize lab-gateway-aliased names like -# ``claude-haiku-4-5-20251001-v1-project`` and silently substitute the -# agent's current default. Without aliasing, those primary-reasoning -# calls fall through to the sidechannel wildcard (mock) and the agent -# gets MODEL_NO_ASSISTANT_MESSAGES. -# - roo (v0.1.x): rewrites to ``claude-sonnet-4-5`` -# - cline (1.x): rewrites to ``claude-sonnet-4-5-20250929`` -# Each alias forwards to the configured upstream primary. Grow this list -# when new agents/versions surface their own defaults. -_AGENT_PRIMARY_ALIASES = ( - "claude-sonnet-4-5", - "claude-sonnet-4-5-20250929", -) - -# Agent-specific request fields that some upstreams reject. LiteLLM's -# global ``drop_params: true`` only drops fields it recognizes as -# "supported by some providers, not this one"; unknown fields pass -# through. ``additional_drop_params`` must be set per-model -# (``litellm_params`` level), not globally — verified empirically. -# - ``client_metadata``: Codex sends this; Bedrock Anthropic Messages -# adapter rejects it ("Extra inputs are not permitted"). -_DROP_PARAMS_YAML = ' additional_drop_params: ["client_metadata"]\n' - - -def _generate_config( - llm_model: str, - llm_base_url: str, - llm_provider: str, - embedding_model: str | None = None, - embedding_base_url: str | None = None, - embedding_provider: str | None = None, -) -> str: - """Render the LiteLLM proxy YAML and return the path to a tempfile. - - Layout: - * Primary route forwards the configured llm.model to the upstream. - * One alias route per ``_AGENT_PRIMARY_ALIASES`` entry, also - forwarded to the upstream primary — so cline/roo's hardcoded - model rewrites still reach the right model. - * Embedding route (only when ``embedding_*`` args provided) forwards - the configured embedding model. In ``local`` embedding mode the - knowledge MCP server uses sentence-transformers in-process and - bypasses the proxy entirely, so the embedding route is skipped. - * Sidechannel wildcard catches everything else (agent title-gen, - session-namer, etc.) and returns a canned mock response. - ``observability.SIDECHANNEL_WILDCARD_ROUTE_YAML`` provides this. - """ - from dsagt.observability import SIDECHANNEL_WILDCARD_ROUTE_YAML - - aliases_yaml = "".join(f""" - model_name: {alias} - litellm_params: - model: {llm_provider}/{llm_model} - api_base: {llm_base_url} - api_key: os.environ/LLM_API_KEY -{_DROP_PARAMS_YAML}""" for alias in _AGENT_PRIMARY_ALIASES) - - embedding_yaml = "" - if embedding_model and embedding_base_url and embedding_provider: - embedding_yaml = ( - f" - model_name: {embedding_model}\n" - f" litellm_params:\n" - f" model: {embedding_provider}/{embedding_model}\n" - f" api_base: {embedding_base_url}\n" - f" api_key: os.environ/EMBEDDING_API_KEY\n" - ) - - body = f"""\ -model_list: - - model_name: {llm_model} - litellm_params: - model: {llm_provider}/{llm_model} - api_base: {llm_base_url} - api_key: os.environ/LLM_API_KEY -{_DROP_PARAMS_YAML}{aliases_yaml}{embedding_yaml}{SIDECHANNEL_WILDCARD_ROUTE_YAML}\ -litellm_settings: - drop_params: true - num_retries: 5 - request_timeout: 300 -""" - tmp = tempfile.NamedTemporaryFile( - mode="w", - suffix=".yaml", - prefix="dsagt_litellm_", - delete=False, - ) - tmp.write(body) - tmp.close() - return tmp.name - - -def main() -> None: - parser = argparse.ArgumentParser(prog="dsagt-proxy") - parser.add_argument("--port", type=int, required=True) - parser.add_argument("--host", default="0.0.0.0") - parser.add_argument( - "--mlflow-url", - required=True, - help="MLflow server URL the OTLP exporter ships traces to.", - ) - parser.add_argument( - "--project", - required=True, - help="DSAGT project name (= MLflow experiment name).", - ) - parser.add_argument( - "--session", - required=True, - help="DSAGT session id stamped on every trace via OTel resource attrs.", - ) - parser.add_argument( - "--records-dir", - required=True, - help="Project's trace_archive/ — sidechannel.jsonl lands adjacent.", - ) - parser.add_argument("--model", required=True) - parser.add_argument("--base-url", required=True) - parser.add_argument("--provider", required=True) - # Embedding args optional — only needed when project's embedding - # backend is ``api``. In ``local`` mode the knowledge MCP server - # uses sentence-transformers in-process and never routes through - # the proxy, so we skip the embedding route entirely. - parser.add_argument("--embedding-model", default=None) - parser.add_argument("--embedding-base-url", default=None) - parser.add_argument("--embedding-provider", default=None) - parser.add_argument("--verbose", action="store_true") - args = parser.parse_args() - - logging.basicConfig( - level=logging.DEBUG if args.verbose else logging.INFO, - format="%(asctime)s [dsagt-proxy] %(levelname)s %(message)s", - datefmt="%H:%M:%S", - ) - - if not os.environ.get("LLM_API_KEY"): - logger.error("LLM_API_KEY not set; the proxy needs it to forward LLM requests.") - sys.exit(1) - embedding_routing = bool( - args.embedding_model and args.embedding_base_url and args.embedding_provider - ) - if embedding_routing and not os.environ.get("EMBEDDING_API_KEY"): - logger.error( - "EMBEDDING_API_KEY not set; the proxy needs it to forward embedding requests." - ) - sys.exit(1) - - # init_proxy_tracing installs MLflow's tracer provider as the OTel - # global, plants ``_DSAGTMlflowLogger`` (a MlflowLogger subclass) into - # LiteLLM's logger cache to stamp ``mlflow.trace.session`` / - # ``dsagt.source=agent`` / ``dsagt.agent`` on every proxy-captured - # trace, and registers DSAGTCallback for cache-breakpoint injection + - # sidechannel-call detection. - from dsagt.observability import init_proxy_tracing - - init_proxy_tracing( - mlflow_url=args.mlflow_url, - project=args.project, - session_id=args.session, - records_dir=args.records_dir, - ) - - # Claude Code sends native Anthropic-format requests (POST /v1/messages). - # Default LiteLLM behavior translates those to /responses, which most - # project gateways don't expose. Force the /chat/completions path so - # any openai-compatible upstream works. - import litellm - - litellm.use_chat_completions_url_for_anthropic_messages = True - - # Tell the DSAGT callback what the configured primary model is, so its - # sidechannel detector can distinguish "real upstream call" from - # "wildcard catchall hit". See observability.record_sidechannel_call. - os.environ["DSAGT_PRIMARY_MODEL"] = args.model - - config_path = _generate_config( - args.model, - args.base_url, - args.provider, - args.embedding_model, - args.embedding_base_url, - args.embedding_provider, - ) - logger.info("Generated LiteLLM config at %s", config_path) - logger.info("Starting LiteLLM proxy on %s:%d", args.host, args.port) - - # run_server is a Click command. standalone_mode=False makes Click - # raise on errors instead of sys.exit; we still catch SystemExit - # because Click can raise it on clean shutdown. - from litellm.proxy.proxy_cli import run_server - - try: - run_server.main( - args=[ - "--host", - args.host, - "--port", - str(args.port), - "--config", - config_path, - ], - standalone_mode=False, - ) - except SystemExit: - pass - - -if __name__ == "__main__": - main() diff --git a/src/dsagt/commands/registry_server.py b/src/dsagt/commands/registry_server.py deleted file mode 100644 index 4202811..0000000 --- a/src/dsagt/commands/registry_server.py +++ /dev/null @@ -1,777 +0,0 @@ -""" -DSAgt Registry MCP Server. - -Provides tools for building a tool registry by reading documentation, -fetching web resources, and running commands to extract tool specifications. - -Tool specs are saved as skill markdown files in the runtime skills directory -and indexed into a ChromaDB collection for semantic search. - -Server configuration (embedding credentials) flows through env vars -(LLM_API_KEY, OPENAI_BASE_URL, EMBEDDING_MODEL) set by dsagt start. - -Usage: - dsagt-registry-server --runtime-dir ./my_session -""" - -import asyncio -import json -import logging -import os -import subprocess -import sys -from functools import partial -from pathlib import Path - -import httpx -import yaml - -import mcp.server.stdio -import mcp.types as types -from mcp.server.lowlevel import Server, NotificationOptions -from mcp.server.models import InitializationOptions - -from dsagt.knowledge import KnowledgeBase -from dsagt.observability import ( - obs, - registry_install_deps_span, - registry_reconstruct_pipeline_span, - registry_save_tool_span, -) -from dsagt.provenance import reconstruct_pipeline -from dsagt.registry import ( - SKILLS_COLLECTION, - TOOLS_COLLECTION, - SkillRegistry, - ToolRegistry, -) - -os.environ["PYTHONUNBUFFERED"] = "1" - -logger = logging.getLogger(__name__) - - -# --------------------------------------------------------------------------- -# MCP server helpers -# --------------------------------------------------------------------------- - -async def _run_stdio(server: Server, name: str): - async with mcp.server.stdio.stdio_server() as (read_stream, write_stream): - await server.run( - read_stream, write_stream, - InitializationOptions( - server_name=name, server_version="0.1.0", - capabilities=server.get_capabilities( - notification_options=NotificationOptions(), - experimental_capabilities={}, - ), - ), - ) - - -def _install_dependencies(packages: list[str], timeout: int = 120) -> str: - """Install packages using uv pip install. Returns a status string.""" - cmd = ["uv", "pip", "install", "--python", sys.executable] + packages - try: - result = subprocess.run(cmd, capture_output=True, text=True, timeout=timeout) - if result.returncode == 0: - output = result.stdout.strip() - return f"Successfully installed: {', '.join(packages)}\n{output}" - else: - return ( - f"Installation failed (exit code {result.returncode}):\n" - f"{result.stderr.strip()}" - ) - except subprocess.TimeoutExpired: - return f"Installation timed out after {timeout}s for: {', '.join(packages)}" - except FileNotFoundError: - return "Error: 'uv' command not found. Install uv: https://github.com/astral-sh/uv" - - -# --------------------------------------------------------------------------- -# Per-tool handlers (module-level, explicit dependencies) -# --------------------------------------------------------------------------- - -async def _handle_read_file(arguments: dict) -> str: - path = Path(arguments["path"]) - try: - return path.read_text() - except (FileNotFoundError, PermissionError, IsADirectoryError, - OSError, UnicodeDecodeError) as e: - return f"Error reading file: {e}" - - -async def _handle_http_request(arguments: dict) -> str: - url = arguments["url"] - method = arguments.get("method", "GET") - headers = arguments.get("headers", {}) - try: - async with httpx.AsyncClient(follow_redirects=True) as client: - response = await client.request( - method=method, url=url, headers=headers, timeout=30.0, - ) - return f"Status: {response.status_code}\n\n{response.text}" - except (httpx.HTTPError, httpx.InvalidURL) as e: - return f"Error making request: {e}" - - -async def _handle_run_command(arguments: dict) -> str: - command = arguments["command"] - args = arguments.get("args", []) - timeout = arguments.get("timeout", 10) - try: - result = subprocess.run( - [command] + args, - capture_output=True, text=True, timeout=timeout, - ) - except subprocess.TimeoutExpired: - return f"Command timed out after {timeout} seconds" - except FileNotFoundError: - return f"Command '{command}' not found" - - output = "" - if result.stdout: - output += f"STDOUT:\n{result.stdout}\n" - if result.stderr: - output += f"STDERR:\n{result.stderr}\n" - output += f"\nReturn code: {result.returncode}" - return output - - -async def _handle_save_tool_spec( - arguments: dict, *, registry: ToolRegistry, -) -> str: - spec = arguments["spec"] - # Some MCP clients (notably Claude Sonnet/Haiku 4.x) serialize nested - # object args as JSON strings instead of objects. Accept both shapes. - if isinstance(spec, str): - try: - spec = json.loads(spec) - except json.JSONDecodeError as e: - return f"Error: spec must be a JSON object (or string-encoded JSON object): {e}" - with registry_save_tool_span(spec.get("name")): - obs.set("language", spec.get("language")) - obs.set("n_dependencies", len(spec.get("dependencies") or [])) - obs.set("n_tags", len(spec.get("tags") or [])) - try: - action = registry.save_tool(spec) - except (KeyError, ValueError, OSError) as e: - obs.event("save_tool_failed", error=str(e)[:256]) - return f"Error saving tool spec: {e}" - - tool_count = len(registry.list_tools_raw()) - obs.set("action", action) - obs.set("registry_size", tool_count) - message = ( - f"Tool '{spec['name']}' {action} successfully. " - f"Registry now contains {tool_count} tools." - ) - deps = spec.get("dependencies", []) - if deps: - with registry_install_deps_span(deps): - dep_result = _install_dependencies(deps) - if dep_result.startswith("Successfully installed:"): - obs.set("status", "ok") - else: - obs.set("status", "failed") - obs.event("install_failed", message=dep_result[:256]) - message += f"\n\nDependency installation:\n{dep_result}" - return message - - -async def _handle_save_skill( - arguments: dict, *, skill_registry: SkillRegistry, -) -> str: - """Register a skill (workflow / agent instructions) for later reuse. - - Symmetric with save_tool_spec — writes SKILL.md to - ``/skills//`` and indexes it into - ``registered_skills`` so future ``search_skills`` calls find it. - """ - spec = arguments["spec"] - if isinstance(spec, str): - try: - spec = json.loads(spec) - except json.JSONDecodeError as e: - return f"Error: spec must be a JSON object (or string-encoded JSON object): {e}" - body = arguments.get("body") - reference_files = arguments.get("reference_files") - if isinstance(reference_files, str): - try: - reference_files = json.loads(reference_files) - except json.JSONDecodeError as e: - return f"Error: reference_files must be a JSON object: {e}" - try: - action = skill_registry.save_skill(spec, body=body, reference_files=reference_files) - except (KeyError, ValueError, OSError) as e: - return f"Error saving skill: {e}" - skill_count = len(skill_registry.list_skills()) - return ( - f"Skill '{spec['name']}' {action} successfully. " - f"Registry now contains {skill_count} skills." - ) - - -async def _handle_get_registry( - arguments: dict, *, registry: ToolRegistry, -) -> str: - tools = registry.list_tools_raw() - if not tools: - return "Registry is empty. No tools registered yet." - return yaml.dump({"tools": tools}, default_flow_style=False, sort_keys=False) - - -async def _handle_search_registry( - arguments: dict, *, registry: ToolRegistry, kb: KnowledgeBase | None, -) -> str: - tool_name = arguments.get("tool_name") - query = arguments.get("query", "") - tag = arguments.get("tag") - top_k = arguments.get("top_k", 10) - - if tool_name: - tool = registry.get_tool(tool_name) - if tool: - return ( - f"Found tool '{tool_name}':\n\n" - + yaml.dump(tool, default_flow_style=False, sort_keys=False) - ) - return f"No tool named '{tool_name}'." - - if kb is None: - return ( - "search_registry requires a configured knowledge base " - "(set embedding.api_key + embedding.base_url + embedding.model " - "in dsagt_config.yaml). Use search_registry with an exact " - "tool_name for KB-free lookups." - ) - - # Single ``tools`` collection — bundled and registered entries - # coexist, distinguished by ``metadata.source`` if needed. - results = kb.search( - query=query or "tool", - collection=TOOLS_COLLECTION, - top_k=top_k * 3 if tag else top_k, - ) - if tag and results: - results = [ - r for r in results - if tag in r.get("chunk", {}).get("metadata", {}).get("tags", "") - ][:top_k] - if not results: - return "No tools found matching the query." - - summaries = [] - for r in results: - chunk = r.get("chunk", {}) - meta = chunk.get("metadata", {}) - summaries.append( - f"- **{meta.get('tool_name', 'unknown')}** " - f"(score: {r.get('score', 0):.2f})\n" - f" {chunk.get('text', '')[:200]}" - ) - return f"Found {len(results)} tool(s):\n\n" + "\n\n".join(summaries) - - -async def _handle_search_skills( - arguments: dict, - *, - kb: KnowledgeBase | None, - skill_registry: SkillRegistry | None, -) -> str: - skill_name = arguments.get("skill_name") - query = arguments.get("query", "") - tag = arguments.get("tag") - top_k = arguments.get("top_k", 10) - - if skill_name and skill_registry: - skill = skill_registry.get_skill(skill_name) - if skill: - return ( - f"Found skill '{skill_name}':\n\n" - + yaml.dump(skill, default_flow_style=False, sort_keys=False) - ) - return f"No skill named '{skill_name}'." - - if kb is None: - return ( - "search_skills requires a configured knowledge base " - "(set embedding.api_key + embedding.base_url + embedding.model " - "in dsagt_config.yaml). Use search_skills with an exact " - "skill_name for KB-free lookups." - ) - - # Single ``skills`` collection — bundled and registered entries. - results = kb.search( - query=query or "skill", - collection=SKILLS_COLLECTION, - top_k=top_k * 3 if tag else top_k, - ) - if tag and results: - results = [ - r for r in results - if tag in r.get("chunk", {}).get("metadata", {}).get("tags", "") - ][:top_k] - if not results: - return "No skills found matching the query." - - summaries = [] - for r in results: - chunk = r.get("chunk", {}) - meta = chunk.get("metadata", {}) - summaries.append( - f"- **{meta.get('skill_name', 'unknown')}** " - f"(score: {r.get('score', 0):.2f})\n" - f" {chunk.get('text', '')[:200]}" - ) - return f"Found {len(results)} skill(s):\n\n" + "\n\n".join(summaries) - - -async def _handle_reconstruct_pipeline( - arguments: dict, *, runtime_dir: Path, -) -> str: - fmt = arguments.get("format", "bash") - trace_dir = runtime_dir / "trace_archive" - with registry_reconstruct_pipeline_span(fmt): - try: - script = reconstruct_pipeline(trace_dir, fmt=fmt) - except (FileNotFoundError, ValueError, OSError) as e: - obs.event("reconstruct_failed", error=str(e)[:256]) - return f"Error reconstructing pipeline: {e}" - obs.set("output_chars", len(script)) - return script - - -async def _handle_install_dependencies( - arguments: dict, *, registry: ToolRegistry, -) -> str: - tool_name = arguments.get("tool_name") - tools = registry.list_tools_raw() - if not tools: - return "Registry is empty. No tools registered yet." - - all_deps = [] - tools_with_deps = [] - for tool in tools: - if tool_name and tool.get("name") != tool_name: - continue - tool_deps = tool.get("dependencies", []) - if tool_deps: - all_deps.extend(tool_deps) - tools_with_deps.append(tool["name"]) - - if not all_deps: - scope = f"tool '{tool_name}'" if tool_name else "registry" - return f"No dependencies declared in {scope}." - - seen = set() - unique_deps = [d for d in all_deps if not (d in seen or seen.add(d))] - - with registry_install_deps_span(unique_deps): - obs.set("scope_tool", tool_name) - obs.set("n_tools_with_deps", len(tools_with_deps)) - result = _install_dependencies(unique_deps) - if result.startswith("Successfully installed:"): - obs.set("status", "ok") - else: - obs.set("status", "failed") - obs.event("install_failed", message=result[:256]) - return f"Installing dependencies for: {', '.join(tools_with_deps)}\n\n{result}" - - -# --------------------------------------------------------------------------- -# Server factory (thin wiring — used by main() and tests) -# --------------------------------------------------------------------------- - -def create_registry_server( - registry: ToolRegistry, - kb: KnowledgeBase | None = None, - skill_registry: SkillRegistry | None = None, -): - """Create and configure the MCP registry server. - - Test-facing API: tests call with a mock registry and get back a server - they can drive via call_tool_sync(). main() constructs the registry - and KB from config before calling this. - """ - server = Server("registry") - runtime_dir = Path(registry.runtime_dir) - - # Dispatch table — maps MCP tool names to handler functions with - # dependencies bound via functools.partial. - handlers = { - "read_file": _handle_read_file, - "http_request": _handle_http_request, - "run_command": _handle_run_command, - "save_tool_spec": partial(_handle_save_tool_spec, registry=registry), - "save_skill": partial(_handle_save_skill, skill_registry=skill_registry), - "get_registry": partial(_handle_get_registry, registry=registry), - "search_registry": partial(_handle_search_registry, registry=registry, kb=kb), - "search_skills": partial(_handle_search_skills, kb=kb, skill_registry=skill_registry), - "reconstruct_pipeline": partial(_handle_reconstruct_pipeline, runtime_dir=runtime_dir), - "install_dependencies": partial(_handle_install_dependencies, registry=registry), - } - - @server.list_tools() - async def list_tools() -> list[types.Tool]: - return [ - types.Tool( - name="read_file", - description="Read contents of a text file", - inputSchema={ - "type": "object", - "properties": { - "path": {"type": "string", "description": "Path to the file to read"}, - }, - "required": ["path"], - }, - ), - types.Tool( - name="http_request", - description="Make an HTTP request to fetch documentation or API specs", - inputSchema={ - "type": "object", - "properties": { - "url": {"type": "string", "description": "URL to request"}, - "method": {"type": "string", "description": "HTTP method", "default": "GET"}, - "headers": {"type": "object", "description": "Optional headers"}, - }, - "required": ["url"], - }, - ), - types.Tool( - name="run_command", - description="Execute a command to get help/usage information", - inputSchema={ - "type": "object", - "properties": { - "command": {"type": "string", "description": "Command to execute"}, - "args": { - "type": "array", - "items": {"type": "string"}, - "description": "Arguments (e.g., ['--help'])", - "default": [], - }, - "timeout": {"type": "number", "default": 10}, - }, - "required": ["command"], - }, - ), - types.Tool( - name="save_tool_spec", - description="Save a tool specification to the registry as a skill file", - inputSchema={ - "type": "object", - "properties": { - # ``anyOf`` accepts both a structured object and a - # JSON-encoded string. Some MCP clients (notably - # Claude Sonnet/Haiku 4.x) serialize nested object - # arguments as JSON strings instead of objects; the - # handler unwraps either shape. - "spec": { - "description": "Tool specification (object or JSON-encoded string)", - "anyOf": [ - { - "type": "object", - "properties": { - "name": {"type": "string", "description": "Unique tool identifier"}, - "description": {"type": "string", "description": "What the tool does"}, - "executable": {"type": "string", "description": "Command to execute"}, - "parameters": { - "type": "object", - "description": "Parameter definitions", - "additionalProperties": { - "type": "object", - "properties": { - "type": {"type": "string", "description": "Parameter type"}, - "required": {"type": "boolean"}, - "description": {"type": "string"}, - "default": {"description": "Default value"}, - "cli": { - "type": "string", - "description": ( - "How to render this parameter on the command line: " - "'positional[:N]' for positional args, '--name' or '-n' " - "for spaced flags, '--name=' or '-n=' for glued flags, " - "'key=' for dd-style key=value. Defaults to '--' " - "if omitted." - ), - }, - }, - "required": ["type", "description"], - }, - }, - "dependencies": { - "type": "array", - "items": {"type": "string"}, - "description": "Python packages to install", - }, - "tags": { - "type": "array", - "items": {"type": "string"}, - "description": "Tags for categorizing the tool", - }, - }, - "required": ["name", "description", "executable", "parameters"], - }, - {"type": "string"}, - ], - }, - }, - "required": ["spec"], - }, - ), - types.Tool( - name="save_skill", - description=( - "Register a skill (agent workflow / instructions) into " - "/skills//SKILL.md and index it into the " - "registered_skills KB collection. Symmetric with " - "save_tool_spec — use this when you've designed a " - "reusable instruction set you want future sessions to " - "discover via search_skills." - ), - inputSchema={ - "type": "object", - "properties": { - # ``anyOf`` for spec mirrors save_tool_spec — accept - # both structured object and JSON-encoded string for - # MCP clients that serialize nested args. - "spec": { - "description": "Skill spec (object or JSON-encoded string)", - "anyOf": [ - { - "type": "object", - "properties": { - "name": {"type": "string", "description": "Unique skill identifier (becomes the directory name)"}, - "description": {"type": "string", "description": "What the skill does / when to use it"}, - "tags": { - "type": "array", - "items": {"type": "string"}, - "description": "Tags for categorizing the skill", - }, - }, - "required": ["name", "description"], - }, - {"type": "string"}, - ], - }, - "body": { - "type": "string", - "description": ( - "Markdown body of the SKILL.md (workflow / " - "instructions the agent will follow). When " - "updating an existing skill, omit to preserve " - "the existing body." - ), - }, - "reference_files": { - "description": ( - "Optional additional files to write into the " - "skill directory. Object mapping relative " - "path -> file contents, or JSON-encoded string." - ), - "anyOf": [ - {"type": "object", "additionalProperties": {"type": "string"}}, - {"type": "string"}, - ], - }, - }, - "required": ["spec"], - }, - ), - types.Tool( - name="get_registry", - description="Get all tools from the registry", - inputSchema={"type": "object", "properties": {}}, - ), - types.Tool( - name="search_registry", - description="Search for tools by name, tag, or description via semantic search.", - inputSchema={ - "type": "object", - "properties": { - "query": {"type": "string", "description": "Search query"}, - "tag": {"type": "string", "description": "Filter by tag"}, - "tool_name": {"type": "string", "description": "Exact tool name lookup"}, - "top_k": {"type": "integer", "default": 10}, - }, - }, - ), - types.Tool( - name="search_skills", - description="Search for agent skills (workflows, templates) by name, tag, or description.", - inputSchema={ - "type": "object", - "properties": { - "query": {"type": "string", "description": "Search query"}, - "tag": {"type": "string", "description": "Filter by tag"}, - "skill_name": {"type": "string", "description": "Exact skill name lookup"}, - "top_k": {"type": "integer", "default": 10}, - }, - }, - ), - types.Tool( - name="reconstruct_pipeline", - description="Reconstruct a reproducible pipeline script from tool execution records.", - inputSchema={ - "type": "object", - "properties": { - "format": { - "type": "string", - "enum": ["bash", "snakemake"], - "default": "bash", - }, - }, - }, - ), - types.Tool( - name="install_dependencies", - description="Install Python dependencies for one or all tools in the registry.", - inputSchema={ - "type": "object", - "properties": { - "tool_name": {"type": "string", "description": "Install deps for a specific tool (omit for all)"}, - }, - }, - ), - ] - - @server.call_tool() - async def call_tool(name: str, arguments: dict) -> list[types.TextContent]: - handler = handlers[name] # KeyError = bug in list_tools schema - text = await handler(arguments) - return [types.TextContent(type="text", text=text)] - - return server - - -# --------------------------------------------------------------------------- -# Entry point -# --------------------------------------------------------------------------- - -def main(): - """Entry point for dsagt-registry-server. - - All configuration comes from the project directory: - - ``./dsagt_config.yaml`` → project path + name - - ``LLM_API_KEY``, ``OPENAI_BASE_URL`` env vars → embedding credentials - - No CLI arguments. By contract the agent's launch one-liner is - ``cd && ``, so cwd is project_dir for the MCP - children it spawns. - """ - import logging as _logging - from dsagt.observability import find_project_config - - project_dir, _cfg = find_project_config() - if project_dir is None: - raise RuntimeError( - "dsagt-registry-server: no dsagt_config.yaml in cwd " - f"({Path.cwd()}). Launch the agent from the project " - "directory (`cd && `)." - ) - - log_file = project_dir / "dsagt_registry_server.log" - # Default INFO; users opt into DEBUG via DSAGT_LOG_LEVEL=DEBUG. At DEBUG, - # transitive libraries (httpcore, urllib3, llama_index, chromadb) flood - # stderr with one line per network operation — when an agent like roo - # pipes the MCP server's stderr into its own debug stream, the human - # output gets buried under thousands of low-value lines. - _level_name = os.environ.get("DSAGT_LOG_LEVEL", "INFO").upper() - _level = getattr(_logging, _level_name, _logging.INFO) - _logging.basicConfig( - level=_level, - format="%(asctime)s %(levelname)s %(name)s: %(message)s", - handlers=[ - _logging.FileHandler(log_file, mode="a"), - _logging.StreamHandler(), - ], - ) - log = _logging.getLogger(__name__) - log.info("Server starting — log file: %s", log_file) - - config_path = project_dir / "dsagt_config.yaml" - from dsagt.session import resolve_env_vars - config = resolve_env_vars(yaml.safe_load(config_path.read_text())) - - emb_config = config["embedding"] - - from dsagt.observability import init_tracing, configure_litellm_retries - init_tracing("dsagt-registry-server") # session_id picked up from DSAGT_SESSION_ID env - configure_litellm_retries() - - # The KB is optional for the registry server — most tools (save_tool_spec, - # get_registry, read_file, run_command, etc.) work without it. Only - # search_registry and search_skills need the KB for semantic search; - # they return a clear error if the KB is None. - backend = (emb_config.get("backend") or "local").lower() - # Cross-backend leakage guard: see the matching block in - # knowledge_server.py for the rationale. In short: when - # ``embedding.backend`` is ``local`` but the resolved model name is - # an OpenAI-style alias (no slash), drop the override so we fall - # back to the LocalEmbeddingClient default rather than 404 from HF. - raw_model = (emb_config.get("model") or "").strip() - embedder_kwargs: dict = {} - if raw_model and not raw_model.startswith("${"): - looks_hf = "/" in raw_model - if backend == "local" and not looks_hf: - log.warning( - "Ignoring embedding.model=%r for backend=local (does not " - "look like a HuggingFace identifier). Falling back to the " - "LocalEmbeddingClient default.", - raw_model, - ) - else: - embedder_kwargs["model"] = raw_model - if backend == "local": - kb_available = True - else: # backend == "api" - api_key = emb_config.get("api_key") or "" - kb_available = ( - api_key and not api_key.startswith("${") - and emb_config.get("base_url") - ) - embedder_kwargs.update({ - "base_url": emb_config.get("base_url") or "", - "api_key": api_key, - }) - - kb = None - if kb_available: - kb = KnowledgeBase( - index_dir=project_dir / "kb_index", - default_embedder=backend, - default_index=config["knowledge"]["vector_db"], - embedder_kwargs=embedder_kwargs, - ) - # Background-load the embedder so the model is ready when the - # agent's first search_registry / save_tool_spec call lands. - # See knowledge_server for the same rationale. - kb.preload_default_embedder() - - registry = ToolRegistry( - source_tools_dir=None, - runtime_dir=str(project_dir), - kb=kb, - ) - - skill_reg = SkillRegistry( - source_skills_dir=None, - runtime_dir=str(project_dir), - kb=kb, - ) - - # Bundled tools/skills are pre-embedded in the shared - # ~/dsagt-projects/kb_index/ by ``dsagt setup-kb`` (or by the auto-bootstrap - # in ``dsagt start``) and COPIED into the project's kb_index by - # ``setup_runtime_kb`` before either MCP server spawns. This server - # does no embedding work for bundled content at startup; agent's - # save_tool_spec / save_skill incur a single embed at save time. - - server = create_registry_server(registry, kb, skill_reg) - asyncio.run(_run_stdio(server, "registry")) - - -if __name__ == "__main__": - main() diff --git a/src/dsagt/commands/run_tool.py b/src/dsagt/commands/run_code.py similarity index 64% rename from src/dsagt/commands/run_tool.py rename to src/dsagt/commands/run_code.py index 4efc16c..757134f 100644 --- a/src/dsagt/commands/run_tool.py +++ b/src/dsagt/commands/run_code.py @@ -1,8 +1,8 @@ """ -dsagt-run: Tool execution wrapper for provenance capture. +dsagt-run: registered-code execution wrapper for provenance capture. Usage: - dsagt-run --tool fastp -- fastp -q 20 -l 50 --in1 reads.fq.gz + dsagt-run --code fastp -- fastp -q 20 -l 50 --in1 reads.fq.gz """ import argparse @@ -14,15 +14,26 @@ def _make_parser() -> argparse.ArgumentParser: parser = argparse.ArgumentParser( prog="dsagt-run", - description="Wrap a tool command and capture execution provenance.", + description="Wrap a code command and capture execution provenance.", + ) + parser.add_argument( + "--code", required=True, help="Name of the code being executed." + ) + parser.add_argument( + "--session", + default=None, + help="Session ID. Defaults to the DSAGT_SESSION_ID env var.", ) - parser.add_argument("--tool", required=True, help="Name of the tool being executed.") - parser.add_argument("--session", default=None, - help="Session ID. Defaults to session_id from /.runtime.") parser.add_argument("--record-id", default=None, help="Pre-assigned record ID.") - parser.add_argument("--records-dir", default=None, help="Directory for execution records.") - parser.add_argument("--input-files", default=None, help="Comma-separated input file paths.") - parser.add_argument("--output-files", default=None, help="Comma-separated output file paths.") + parser.add_argument( + "--records-dir", default=None, help="Directory for execution records." + ) + parser.add_argument( + "--input-files", default=None, help="Comma-separated input file paths." + ) + parser.add_argument( + "--output-files", default=None, help="Comma-separated output file paths." + ) return parser @@ -37,7 +48,7 @@ def _parse_args(argv: list[str] | None = None) -> tuple[argparse.Namespace, list sys.exit(1) wrapper_args = args_to_parse[:sep] - command_args = args_to_parse[sep + 1:] + command_args = args_to_parse[sep + 1 :] parsed = _make_parser().parse_args(wrapper_args) return parsed, command_args @@ -51,12 +62,13 @@ def main(argv: list[str] | None = None) -> int: return 1 from dsagt.observability import init_tracing + init_tracing("dsagt-run", session_id=args.session) records_dir = _resolve_records_dir(args.records_dir) return run_and_record( - tool_name=args.tool, + code_name=args.code, command=command, records_dir=records_dir, session_id=args.session, diff --git a/src/dsagt/commands/setup_core_kb.py b/src/dsagt/commands/setup_core_kb.py index 9cf977c..8c0858e 100644 --- a/src/dsagt/commands/setup_core_kb.py +++ b/src/dsagt/commands/setup_core_kb.py @@ -1,28 +1,23 @@ """ -dsagt-setup-kb: Build core knowledge base collections. - -Downloads and indexes: -- nemo_curator: NVIDIA NeMo Curator (code, docs, tutorials) -- aidrin: AI Data Readiness Inspector (code, papers) - -The embedding service is configured via CLI flags or environment variables -(``LLM_API_KEY``, ``OPENAI_BASE_URL``, ``EMBEDDING_MODEL``). API-backed -embedding of the full core KB typically takes 15-30 minutes. - -Usage: - dsagt-setup-kb - dsagt-setup-kb --index-dir ./my_kb_index - dsagt-setup-kb --collection nemo_curator - dsagt-setup-kb --embedding-base-url https://api.example.com/v1 \\ - --embedding-api-key sk-... \\ - --embedding-model text-embedding-3-small +Knowledge-base asset builder — the engine behind ``dsagt init``'s KB +provisioning. No longer a standalone command (``dsagt setup-kb`` was +retired); ``session._provision_kb`` calls :func:`resolve_assets` + +:func:`ensure_assets` to build the requested assets into the shared +``~/dsagt-projects/kb_index/`` once, then copies them per project. + +Asset namespace (the ``--include`` / ``--exclude`` selectors on ``dsagt init``): +- ``tools`` bundled tool specs (cheap, local) +- skill catalogs ``genesis`` (default), ``scientific``, ``composio``, … +- scientific collections ``nemo_curator``, ``aidrin`` (heavy; clone external repos) + +:data:`DEFAULT_ASSETS` (bundled tools + the genesis skill catalog) is the +cheap set installed automatically on a machine's first project. Embedding +config comes from the project's ``.dsagt/config.yaml`` (local backend by +default — no credentials needed). """ -import argparse -import os import shutil import subprocess -import sys import tarfile import tempfile from pathlib import Path @@ -135,7 +130,9 @@ } -def clone_github(url: str, dest: Path, branch: str = "main", include: list[str] | None = None): +def clone_github( + url: str, dest: Path, branch: str = "main", include: list[str] | None = None +): """Clone a GitHub repo, optionally keeping only specific directories. When *include* is set, the named subdirectories are copied AND any @@ -175,7 +172,7 @@ def clone_github(url: str, dest: Path, branch: str = "main", include: list[str] def download_arxiv(paper_id: str, dest: Path): """Download arXiv paper (source if available, else PDF).""" client = httpx.Client(timeout=60.0, follow_redirects=True) - + try: # Try source tarball first response = client.get(f"https://arxiv.org/e-print/{paper_id}") @@ -190,7 +187,7 @@ def download_arxiv(paper_id: str, dest: Path): return except tarfile.ReadError: tar_path.unlink() - + # Fall back to PDF response = client.get(f"https://arxiv.org/pdf/{paper_id}.pdf") response.raise_for_status() @@ -216,11 +213,10 @@ def setup_collection( compat with direct callers), a fresh KB is constructed and closed around this call. - ``run_setup_kb`` always passes a shared *kb* so a multi-collection - run pays the model-load cost once, not N times. + ``ensure_assets`` always passes a shared *kb* so a multi-asset build + pays the model-load cost once, not N times. It also prints the + user-facing progress line, so this function stays quiet. """ - print(f"Setting up {name}...", flush=True) - with tempfile.TemporaryDirectory() as tmp: download_dir = Path(tmp) / name download_dir.mkdir() @@ -248,16 +244,19 @@ def setup_collection( owned_kb = kb is None if owned_kb: from dsagt.knowledge import KnowledgeBase + kb = KnowledgeBase( index_dir=index_dir, default_embedder=embedding_backend, - default_index=vector_db, - embedder_kwargs=embedder_kwargs, + model=embedder_kwargs.get("model"), + base_url=embedder_kwargs.get("base_url"), + api_key=embedder_kwargs.get("api_key"), ) try: result = kb.ingest( download_dir, - exclude_patterns=config.get("exclude_patterns") or DEFAULT_EXCLUDE_PATTERNS, + exclude_patterns=config.get("exclude_patterns") + or DEFAULT_EXCLUDE_PATTERNS, ) skipped = result.get("skipped_files", 0) miss_msg = f", {skipped} file misses" if skipped else "" @@ -275,229 +274,250 @@ def _current_dsagt_version() -> str: """Return the installed dsagt package version, or ``"unknown"`` if absent.""" try: from importlib.metadata import version + return version("dsagt") except Exception: return "unknown" -def add_setup_kb_args(parser): - """Add setup-kb arguments to a parser or subparser. +# --------------------------------------------------------------------------- +# Installable KB assets — the namespace for ``dsagt init``'s +# ``--include`` / ``--exclude`` selectors. +# +# Three kinds, all built into the shared ``~/dsagt-projects/kb_index/`` and +# copied per-project: +# "codes" bundled tool specs (package data; cheap, fully local) +# a skill-catalog source from ``skills.KNOWN_SOURCES`` +# (e.g. "genesis", "k-dense-ai", "composio", "antigravity") +# a heavy scientific doc collection from ``COLLECTIONS`` +# (e.g. "nemo_curator", "aidrin" — clones external repos) +# +# DEFAULT_ASSETS is the cheap core a first-ever ``dsagt init`` installs +# automatically; everything else is opt-in via ``--include``. +# --------------------------------------------------------------------------- - Called from both the standalone ``dsagt-setup-kb`` entry point and the - ``dsagt setup-kb`` subcommand so the argument set is defined once. - """ - parser.add_argument( - "--index-dir", - type=Path, - default=DEFAULT_INDEX_DIR, - help=f"Index directory (default: {DEFAULT_INDEX_DIR})", - ) - parser.add_argument( - "--collection", - choices=list(COLLECTIONS.keys()), - help="Setup only this collection", - ) - parser.add_argument( - "--embedding-backend", - choices=["api", "local"], - default="local", - help=( - "Embedding backend (default: local — sentence-transformers, " - "no API credentials needed). Pass --embedding-backend api to " - "build collections against a hosted embedding endpoint; that " - "path requires --embedding-base-url and --embedding-api-key (or " - "the corresponding env vars)." - ), - ) - parser.add_argument( - "--embedding-model", default=None, - help="Embedding model name (falls back to EMBEDDING_MODEL env var)", - ) - parser.add_argument( - "--embedding-base-url", default=None, - help="Embedding API base URL (falls back to OPENAI_BASE_URL env var)", - ) - parser.add_argument( - "--embedding-api-key", default=None, - help="Embedding API key (falls back to LLM_API_KEY / OPENAI_API_KEY env var)", - ) - parser.add_argument( - "--vector-db", - choices=["chroma", "faiss"], - default="chroma", - help="Vector database backend (default: chroma)", - ) - parser.add_argument( - "--rebuild", - action="store_true", - help="Re-ingest collections that already exist in the index directory " - "(default: skip existing).", - ) +#: The default per-project / first-init asset set: bundled tools + the +#: genesis skill catalog. Kept deliberately cheap (one small local embed + +#: one git clone) so onboarding needs no manual step. +DEFAULT_ASSETS: tuple[str, ...] = ("codes", "genesis") + + +def all_assets() -> list[str]: + """Every installable asset name, in canonical install order (cheap → heavy).""" + from dsagt.skills import KNOWN_SOURCES + + return ["codes", *KNOWN_SOURCES, *COLLECTIONS] + + +def asset_collection_name(asset: str) -> str: + """The ``kb_index`` collection directory a given asset materializes as.""" + from dsagt.registry import CODES_COLLECTION, CATALOG_COLLECTION_PREFIX + from dsagt.skills import KNOWN_SOURCES, _repo_slug + if asset == "codes": + return CODES_COLLECTION + if asset in KNOWN_SOURCES: + return CATALOG_COLLECTION_PREFIX + _repo_slug(KNOWN_SOURCES[asset]["url"]) + if asset in COLLECTIONS: + return asset + raise ValueError(f"unknown KB asset: {asset!r}") -def run_setup_kb(args): - """Run the core knowledge base setup. - Accepts a parsed argparse.Namespace with the fields added by - ``add_setup_kb_args``. Called from both the ``dsagt setup-kb`` - subcommand and the standalone ``dsagt-setup-kb`` entry point. +def resolve_assets( + include: list[str] | None = None, + exclude: list[str] | None = None, +) -> list[str]: + """Resolve the requested asset set from ``--include`` / ``--exclude``. + + - ``include`` and ``exclude`` are mutually exclusive. + - The literal ``"all"`` expands to every installable asset. + - No selector → :data:`DEFAULT_ASSETS`. + - ``--exclude all`` → ``[]`` (empty stub; the project's KB is created but + holds no bundled content). + + Returns names in canonical install order so a build pays the cheap + (local) assets before the heavy (network) ones. """ - # Show only warnings/errors during setup-kb — the per-collection - # print() lines below are the user-visible progress story. Surfacing - # every "Found N files" / "Created M chunks" log line on top creates - # the noisy play-by-play we deliberately stripped out. - # - # ``force=True`` overrides cli.py's earlier basicConfig (which sets - # INFO + a "[dsagt]" format for the rest of the CLI). Without - # ``force``, the second basicConfig is a no-op because the root - # logger already has handlers, and the INFO-level chatter survives. - import logging as _logging - _logging.basicConfig( - level=_logging.WARNING, - format="%(levelname)s: %(message)s", - force=True, - ) + if include and exclude: + raise ValueError("--include and --exclude are mutually exclusive") + known = all_assets() - # Check git is available. - try: - subprocess.run(["git", "--version"], capture_output=True, check=True) - except (subprocess.CalledProcessError, FileNotFoundError): - raise RuntimeError("git is required but not found on PATH") - - # Resolve embedding config with env-var fallback so the user gets a - # clear error up front rather than 5 minutes into the first ingest. - embedder_kwargs: dict = {} - if args.embedding_backend == "api": - api_key = args.embedding_api_key or os.getenv("LLM_API_KEY") or os.getenv("OPENAI_API_KEY") - base_url = args.embedding_base_url or os.getenv("OPENAI_BASE_URL") - model = args.embedding_model or os.getenv("EMBEDDING_MODEL") - missing = [n for n, v in [("api key", api_key), ("base URL", base_url), ("model", model)] if not v] - if missing: + def _validate(names: list[str]) -> None: + bad = [n for n in names if n != "all" and n not in known] + if bad: raise ValueError( - "API embedding backend requires " - + ", ".join(missing) - + ". Pass via --embedding-* flags or set LLM_API_KEY, " - "OPENAI_BASE_URL, EMBEDDING_MODEL." + f"unknown KB asset(s): {', '.join(bad)}. " + f"Choose from: {', '.join(known)} (or 'all')." ) - embedder_kwargs = {"api_key": api_key, "base_url": base_url, "model": model} - elif args.embedding_model: - embedder_kwargs = {"model": args.embedding_model} - - # Configure LiteLLM retries before any embedding work. setup-kb is a - # one-shot bootstrap tool — no project exists yet, no MLflow to trace - # to, so we skip init_tracing entirely. @traced decorators inside - # KnowledgeBase see no backend and short-circuit cleanly. - from dsagt.observability import configure_litellm_retries - configure_litellm_retries() - - # One KnowledgeBase per setup-kb invocation. The embedder cache - # lives on the KB instance, so creating fresh KBs per collection - # would reload the local sentence-transformers model every time - # (~11s × N collections of pure waste). Threaded through to - # setup_collection (and used directly for the bundled tools+skills - # block below) so the model loads once. - args.index_dir.mkdir(parents=True, exist_ok=True) - from dsagt.knowledge import KnowledgeBase - from dsagt.registry import ( - SKILLS_COLLECTION, TOOLS_COLLECTION, - ToolRegistry, SkillRegistry, _parse_frontmatter, - ) - shared_kb = KnowledgeBase( - index_dir=args.index_dir, - default_embedder=args.embedding_backend, - default_index=args.vector_db, - embedder_kwargs=embedder_kwargs or {}, - ) + + if include is not None: + _validate(include) + if "all" in include: + return list(known) + sel = set(include) + return [a for a in known if a in sel] + if exclude is not None: + _validate(exclude) + if "all" in exclude: + return [] + excl = set(exclude) + return [a for a in DEFAULT_ASSETS if a not in excl] + return list(DEFAULT_ASSETS) + + +def _model_is_cached(model_id: str) -> bool: + """True if *model_id* is already in the local HuggingFace cache. + + Best-effort: on any error (offline, API change) returns True so we never + print a misleading "downloading" message for an already-present model. + """ try: - # Bundled tools + skills: each spec file is a single chunk with - # rich metadata. Wipe-and-rebuild every run — there's no - # version sentinel, so the user controls when this happens. - current_version = _current_dsagt_version() - - tool_paths = [ - p for p in sorted(ToolRegistry._PACKAGE_TOOLS_DIR.glob("*.md")) - if _parse_frontmatter(p).get("name") - ] - skill_dirs = [ - d for d in sorted(SkillRegistry._PACKAGE_SKILLS_DIR.iterdir()) - if d.is_dir() - and (d / "SKILL.md").exists() - and _parse_frontmatter(d / "SKILL.md").get("name") - ] - - for name in (TOOLS_COLLECTION, SKILLS_COLLECTION): - coll_dir = args.index_dir / name - if coll_dir.exists(): - shutil.rmtree(coll_dir) - - if tool_paths: - tool_specs = [_parse_frontmatter(p) for p in tool_paths] - shared_kb.add_entries( - texts=[p.read_text() for p in tool_paths], - collection=TOOLS_COLLECTION, - metadatas=[{ - "tool_name": s["name"], - "tags": ",".join(s.get("tags", [])), - "executable": s.get("executable", ""), - "has_dependencies": str(bool(s.get("dependencies"))), - "source": "bundled", - "dsagt_version": current_version, - } for s in tool_specs], - ) + from huggingface_hub import try_to_load_from_cache - if skill_dirs: - skill_specs = [_parse_frontmatter(d / "SKILL.md") for d in skill_dirs] - shared_kb.add_entries( - texts=[(d / "SKILL.md").read_text() for d in skill_dirs], - collection=SKILLS_COLLECTION, - metadatas=[{ - "skill_name": s["name"], - "tags": ",".join(s.get("tags", [])), - "source": "bundled", - "dsagt_version": current_version, - } for s in skill_specs], - ) + # A sentence-transformers model always ships a config.json; a str path + # back means the file is cached (None / sentinel ⇒ not cached). + return isinstance(try_to_load_from_cache(model_id, "config.json"), str) + except Exception: + return True - print(" bundled tools + skills: indexed", flush=True) - - collections = {args.collection: COLLECTIONS[args.collection]} if args.collection else COLLECTIONS - - for name, config in collections.items(): - target_dir = args.index_dir / name - if _collection_exists(target_dir): - if not args.rebuild: - print(f" {name}: already indexed (use --rebuild to force)", - flush=True) - continue - shutil.rmtree(target_dir) - setup_collection( - name, config, args.index_dir, - embedder_kwargs=embedder_kwargs, - embedding_backend=args.embedding_backend, - vector_db=args.vector_db, - kb=shared_kb, - ) - finally: - shared_kb.close() - - -def main(): - """Standalone entry point (``dsagt-setup-kb``).""" - parser = argparse.ArgumentParser( - description="Setup DSAGT core knowledge base", - formatter_class=argparse.RawDescriptionHelpFormatter, - epilog=( - "Embedding config is taken from CLI flags, then from these env vars:\n" - " LLM_API_KEY / OPENAI_API_KEY API key for the embedding endpoint\n" - " OPENAI_BASE_URL OpenAI-compatible base URL\n" - " EMBEDDING_MODEL Embedding model name\n\n" - "Full core KB embedding typically takes 15-30 minutes over an API.\n\n" - "This command can also be run as: dsagt setup-kb" - ), + +def _build_bundled_tools(kb, index_dir: Path) -> int: + """(Re)build the bundled-tools collection from the package's tool specs. + + Each spec file is one chunk with rich metadata. Wipe-and-rebuild so a + dsagt upgrade refreshes the bundled set. Returns the number indexed. + """ + from dsagt.registry import CODES_COLLECTION, CodeRegistry, _parse_frontmatter + + coll_dir = index_dir / CODES_COLLECTION + if coll_dir.exists(): + shutil.rmtree(coll_dir) + + code_paths = [ + p + for p in sorted(CodeRegistry._PACKAGE_CODES_DIR.glob("*/SKILL.md")) + if _parse_frontmatter(p).get("name") + ] + if not code_paths: + return 0 + + current_version = _current_dsagt_version() + code_specs = [_parse_frontmatter(p) for p in code_paths] + kb.add_entries( + texts=[p.read_text() for p in code_paths], + collection=CODES_COLLECTION, + metadatas=[ + { + "code_name": s["name"], + "tags": ",".join(s.get("tags", [])), + "executable": s.get("executable", ""), + "has_dependencies": str(bool(s.get("dependencies"))), + "source": "bundled", + "dsagt_version": current_version, + } + for s in code_specs + ], ) - add_setup_kb_args(parser) - run_setup_kb(parser.parse_args()) + return len(code_paths) + +def ensure_assets( + asset_names: list[str], + index_dir: Path, + *, + embedding_backend: str = "local", + embedder_kwargs: dict | None = None, + vector_db: str = "chroma", + rebuild: bool = False, + kb=None, +) -> dict: + """Build the named *asset_names* into *index_dir* (the shared KB cache). + + Idempotent: an asset whose collection already exists is skipped unless + *rebuild*, so a second ``dsagt init`` pays nothing. Reuses *kb* when + given (keeps the embedder model warm); otherwise constructs and closes + one. Best-effort per asset for catalogs (a clone failure warns and + continues); a heavy-collection failure propagates. + + Returns ``{"built": [...], "skipped": [...]}``. + """ + from dsagt.skills import KNOWN_SOURCES, sync_source + + index_dir.mkdir(parents=True, exist_ok=True) + + # Decide what actually needs building first, so we stay silent and cheap + # when a later init's requested set is already cached. + to_build = [ + a + for a in asset_names + if rebuild or not _collection_exists(index_dir / asset_collection_name(a)) + ] + skipped = [a for a in asset_names if a not in to_build] + if not to_build: + return {"built": [], "skipped": skipped} + + owned = kb is None + if owned: + from dsagt.knowledge import KnowledgeBase + + ek = embedder_kwargs or {} + kb = KnowledgeBase( + index_dir=index_dir, + default_embedder=embedding_backend, + model=ek.get("model"), + base_url=ek.get("base_url"), + api_key=ek.get("api_key"), + ) + + built: list[str] = [] + try: + # The model loads on the first embed below; announce it so the load + # (or one-time download) isn't a silent pause. Distinguish the two so + # we don't claim a download when the model is already cached. API + # backend has no local model to load. + if embedding_backend == "local": + from dsagt.knowledge import LocalEmbedder + + model_id = (embedder_kwargs or {}).get( + "model" + ) or LocalEmbedder.DEFAULT_MODEL + if _model_is_cached(model_id): + print(" Loading embedding model …", flush=True) + else: + print(" Downloading embedding model (one-time) …", flush=True) + + for asset in to_build: + if asset == "codes": + print(" Indexing bundled tools …", flush=True) + _build_bundled_tools(kb, index_dir) + built.append(asset) + elif asset in KNOWN_SOURCES: + print(f" Fetching skill catalog: {asset} …", flush=True) + try: + sync_source(asset, kb=kb, force=rebuild) + built.append(asset) + except Exception as e: # noqa: BLE001 — best-effort, keep going + print(f" skipped {asset} ({e})", flush=True) + else: # heavy scientific collection + print( + f" Fetching collection: {asset} (may take a few minutes) …", + flush=True, + ) + coll = index_dir / asset_collection_name(asset) + if coll.exists(): + shutil.rmtree(coll) + setup_collection( + asset, + COLLECTIONS[asset], + index_dir, + embedder_kwargs=embedder_kwargs or {}, + embedding_backend=embedding_backend, + vector_db=vector_db, + kb=kb, + ) + built.append(asset) + finally: + if owned: + kb.close() -if __name__ == "__main__": - main() + return {"built": built, "skipped": skipped} diff --git a/src/dsagt/commands/traces.py b/src/dsagt/commands/traces.py new file mode 100644 index 0000000..e5eb6fd --- /dev/null +++ b/src/dsagt/commands/traces.py @@ -0,0 +1,103 @@ +""" +dsagt traces — open the MLflow trace viewer over the project's +serverless store, frictionlessly. + +Three ergonomic wins over a raw +``mlflow ui --backend-store-uri sqlite:////mlflow.db``: + +1. **Catch-up first.** Runs :func:`dsagt.session.catch_up_extraction`, which + flushes the most-recent session's deferred final turn (the one an ungraceful + agent exit leaves unlogged) into the store — so "I don't see my last query" + is fixed before the viewer even opens. +2. **Deep link to the Traces tab.** DSAGT writes MLflow *traces*, not classic + runs, so the default Experiments/Runs view looks empty. We resolve the + project's experiment id and print the URL that lands directly on its Traces + tab. +3. **Quiet.** ``--workers 1`` and ``PYTHONWARNINGS=ignore`` drop the repeated + Starlette deprecation spam; the viewer runs in the foreground (Ctrl-C to + stop), so there is no background server to hunt down and kill later — the + store stays serverless. +""" + +from __future__ import annotations + +import logging +import os +import subprocess +from pathlib import Path + +from dsagt.session import catch_up_extraction, load_config + +logger = logging.getLogger(__name__) + +_DEFAULT_PORT = 5000 + + +def _resolve_experiment_id(tracking_uri: str, project: str) -> str | None: + """The MLflow experiment id for *project*, or None if not yet created.""" + try: + import mlflow + + mlflow.set_tracking_uri(tracking_uri) + exp = mlflow.get_experiment_by_name(project) + return exp.experiment_id if exp else None + except Exception as e: # noqa: BLE001 — a missing id only costs the deep link + logger.debug("Could not resolve experiment id for %s: %s", project, e) + return None + + +def run(project: str, port: int = _DEFAULT_PORT) -> int: + config = load_config(project) + pdir = Path(config["project_dir"]) + db = pdir / "mlflow.db" + if not db.exists(): + print( + f"No trace store yet for '{project}' ({db} not found). " + "Run a session first: dsagt start " + f"{project}" + ) + return 1 + + # 1. Catch-up: surface the most-recent session's deferred final turn before + # the viewer opens. Best-effort — a viewer must open even if catch-up + # hiccups. + try: + result = catch_up_extraction(pdir, config) + caught = result.get("traces_caught_up", 0) + if caught: + print(f"Caught up {caught} trailing trace(s) from the last session.") + except Exception as e: # noqa: BLE001 + logger.warning("Trace catch-up before viewer failed: %s", e) + + # 2. Deep-link to the project's Traces tab (DSAGT emits traces, not runs, so + # the default Runs view looks empty). + tracking_uri = f"sqlite:///{db}" + exp_id = _resolve_experiment_id(tracking_uri, project) + base = f"http://127.0.0.1:{port}" + url = f"{base}/#/experiments/{exp_id}/traces" if exp_id else base + + print(f"\nMLflow trace view for '{project}':\n {url}") + print("(Ctrl-C to stop the viewer — the store itself is serverless.)\n") + + # 3. Foreground, quiet: one worker + warnings off drops the Starlette noise. + env = {**os.environ, "PYTHONWARNINGS": "ignore"} + cmd = [ + "mlflow", + "ui", + "--backend-store-uri", + tracking_uri, + "--port", + str(port), + "--workers", + "1", + ] + try: + return subprocess.run(cmd, env=env).returncode + except FileNotFoundError: + print( + "mlflow not found on PATH. It ships with dsagt — activate the same " + "environment dsagt runs in." + ) + return 1 + except KeyboardInterrupt: + return 0 diff --git a/src/dsagt/dsagt_instructions.md b/src/dsagt/dsagt_instructions.md index d2bcfd6..acd634e 100644 --- a/src/dsagt/dsagt_instructions.md +++ b/src/dsagt/dsagt_instructions.md @@ -1,31 +1,44 @@ # DSAgt Pipeline Builder -You are an agentic data pipeline builder. You help domain scientists create **reproducible, auditable data curation pipelines** through iterative, knowledge-driven tool generation. +You are an agentic data pipeline builder. You help domain scientists create **reproducible, auditable data curation pipelines** through iterative, knowledge-driven code generation. ## CRITICAL CONSTRAINTS -### 1. Tool-Mediated Data Access +### 1. Code-Mediated Data Access **Never directly access, observe, assess, transform, or manipulate data.** This includes listing data directories, scanning files, and previewing samples — even when a built-in `shell` or `text_editor` tool would technically work. -All data operations must be performed by calling registered tools. The point is the execution record in `trace_archive/`, not just the result: a built-in shell or editor call leaves no record and breaks pipeline reconstruction. If a needed capability doesn't exist, generate and register it first, then call it. +All data operations must be performed by calling registered codes. The point is the execution record in `trace_archive/`, not just the result: a built-in shell or editor call leaves no record and breaks pipeline reconstruction. If a needed capability doesn't exist, generate and register it first, then call it. ### 1a. Memory: kb_remember / kb_get_memories Are Mandatory **Whenever the user says "remember", "note that", "keep in mind", "for future reference", or otherwise asks you to retain a fact, you MUST call `kb_remember(text=...)` in the same turn.** Mentioning the fact in your response or claiming you have "stored" or "noted" it without making the tool call is a hallucination — the fact is not persisted and a future session will not see it. End-of-session episodic extraction is automatic and unrelated; it is NOT a substitute for explicit memory. **Whenever the user says "what do you remember", "recall", or asks you to retrieve a previously-stored fact, you MUST call `kb_get_memories()` first** and answer based on its result, not from in-context message history. -### 1b. Registered-Tool Invocation: Use the `executable` String Verbatim -**When invoking a registered tool, copy the spec's `executable` field byte-for-byte, including any `dsagt-run --tool --` prefix.** The prefix is the wrapper that writes the execution record to `trace_archive/`; bypassing it (e.g. running `python tools/scan_directory.py` directly when the spec says `dsagt-run --tool scan_directory -- python tools/scan_directory.py`) loses provenance and breaks pipeline reconstruction. If `dsagt-run` errors with "command not found", surface the error rather than working around it. +### 1b. Registered-Code Invocation: Use the `executable` String Verbatim +**When invoking a registered code, copy the spec's `executable` field byte-for-byte, including any `dsagt-run --code --` prefix.** The prefix is the wrapper that writes the execution record to `trace_archive/`; bypassing it (e.g. running the bare script directly when the spec says `dsagt-run --code scan-directory -- python codes/scan-directory/scripts/scan_directory.py`) loses provenance and breaks pipeline reconstruction. If `dsagt-run` errors with "command not found", surface the error rather than working around it. -### 2. Tool and Skill Discovery +### 2. Code and Skill Discovery Before implementing anything, search for existing capabilities: -- `search_registry(query)` — find registered CLI tools by name, tag, or description (semantic search) +- `search_registry(query)` — find registered CLI codes by name, tag, or description (semantic search) - `search_skills(query)` — find agent skills (workflows, templates, procedures) -- `get_registry()` — list all registered tools +- `get_registry()` — list all registered codes -**When the user indicates they want a specific tool used** — phrasings like "use tool `foo`", "use `foo` from the registry", "run `foo`", or similar — look it up first (`search_registry(tool_name=...)` for exact match, `get_registry()` to browse). Read the returned spec's `executable` field and each parameter's `cli` field, then invoke via your shell. Do not substitute your own file/shell tools for a task a registered tool can do. (See section 1b for the verbatim-`executable` rule.) +**Skills come in two tiers.** *Installed* skills (in this project) are discovered **natively** by your platform — their names/descriptions are already in your context and you auto-invoke them; you do NOT need `search_skills` to find those. Separately there is a much larger *external catalog* of installable skills (entries marked `[catalog]`), NOT loaded into context. + +**Registered codes also appear among your native skills** (mirrored at registration and at `dsagt start`). A code's SKILL.md is an execution instruction: it opens with the exact shell command to run. When you invoke a code from its native skill entry, copy that command byte-for-byte — the same verbatim-`executable` rule as section 1b. + +**The external catalog is opt-in and starts empty — sources must be synced before `search_skills` can see them.** A blank/weak `search_skills` result usually means the relevant source isn't synced yet, NOT that no such skill exists. So before concluding the catalog has nothing, call `list_skill_sources()` — it reports each known source with its `synced` flag and `indexed` count. The flow: + +1. `list_skill_sources()` — see which sources are already synced vs only `available` (known name + URL, not yet indexed). For materials/chem/bio/DFT skills, the `k-dense-ai` source is the one to enable. +2. `add_skill_source(source=...)` — sync a source (a known name like `scientific`/`anthropic`, or a GitHub URL). Read-only indexing step; nothing is installed into the project. Only needed for sources whose `synced` is false. +3. `search_skills(query)` — now browse the synced catalog. Entries marked `[catalog]` are installable. +4. `install_skill(skill_name=...)` — copy a catalog skill into the project. The install is **complete and immediately usable**: its SKILL.md + scripts land in `skills//` and are mirrored into your native skills dir on the spot. To use it this session, read `skills//SKILL.md` and follow it — that is exactly what native invocation does. Future sessions auto-discover it hands-free. Never tell the user a restart or any other action is needed before an installed skill can be used. + +To author a brand-new skill instead of installing one, use the bundled `skill-creator` skill. + +**When the user indicates they want a specific code used** — phrasings like "use `foo`", "use `foo` from the registry", "run `foo`", or similar — look it up first (`search_registry(code_name=...)` for exact match, `get_registry()` to browse). Read the returned spec's `executable` field and each parameter's `cli` field, then invoke via your shell. Do not substitute your own file/shell tools for a task a registered code can do. (See section 1b for the verbatim-`executable` rule.) **Rendering parameters**: each parameter's `cli` field pins exactly how its value goes on the command line. Emit positional args first (in position order), then named args. Skip optional parameters whose value is absent; use the `default` when present. @@ -41,20 +54,20 @@ Before implementing anything, search for existing capabilities: Booleans render as a bare flag when truthy, nothing when falsy. -When registering a new tool via `save_tool_spec`, set the `cli` field on every parameter so the next invocation doesn't have to guess. +When registering a new code via `save_code_spec`, set the `cli` field on every parameter so the next invocation doesn't have to guess. Code names use lowercase letters, digits, and hyphens (e.g. `scan-directory`) — the skill-standard charset, since registered codes are mirrored into your native skills directory. -### 3. Tool Preference Hierarchy +### 3. Code Preference Hierarchy When implementing any data operation, follow this hierarchy: -1. **REGISTERED TOOL** — Use an existing tool (`search_registry`) -2. **KB PACKAGE TOOL** — Create a tool leveraging a package documented in the KB -3. **CUSTOM IMPLEMENTATION** — Write custom code to `tools/code/` and register it +1. **REGISTERED CODE** — Use an existing code (`search_registry`) +2. **KB PACKAGE CODE** — Create a code leveraging a package documented in the KB +3. **CUSTOM IMPLEMENTATION** — Write your script to `codes//scripts/` and register it Always exhaust higher-preference options before falling to lower ones. ### 4. Per-Operation Checks -Every filter/transform has an associated check tool. Run it before AND after: +Every filter/transform has an associated check code. Run it before AND after: ``` check_[X](input) → audit/step_N_pre.json operation(input, output) @@ -64,7 +77,7 @@ check_[X](output) → audit/step_N_post.json All check reports are saved to `audit/` for the audit trail. ### 5. File Organization -- All processing code goes in `tools/code/` +- Each registered code is a self-contained dir: spec at `codes//SKILL.md`, its scripts in `codes//scripts/` - All data output goes in a `data/` subdirectory - All audit reports go in `audit/` - All session artifacts stay within the project directory @@ -97,17 +110,17 @@ If yes: use `kb_ingest` to index them. Review what's available: `kb_list_collections()` -### 3. Register User's Custom Tools +### 3. Register User's Custom Codes -Ask: "Do you have existing scripts or tools you'd like to incorporate?" +Ask: "Do you have existing scripts or codes you'd like to incorporate?" -If yes, register them using `save_tool_spec`. +If yes, register them using `save_code_spec`. ### 4. Explore Available Resources Survey what's available before proceeding: -- `get_registry()` — list all tools -- `search_registry(query)` — semantic search for tools +- `get_registry()` — list all codes +- `search_registry(query)` — semantic search for codes - `search_skills(query)` — find available skills - `kb_list_collections()` — list knowledge base collections @@ -116,11 +129,11 @@ Survey what's available before proceeding: For **each data manipulation step**, cycle through: 1. **UNDERSTAND** what needs to happen at this step -2. **EXPLORE** knowledge base (`kb_search`) and tools (`search_registry`) -3. **APPLY** tool preference hierarchy +2. **EXPLORE** knowledge base (`kb_search`) and codes (`search_registry`) +3. **APPLY** code preference hierarchy 4. **DESIGN** the check and operation (confirm with user) -5. **GENERATE** code for check tool AND operation tool -6. **REGISTER** new tools via `save_tool_spec` +5. **GENERATE** the check code AND operation code +6. **REGISTER** new codes via `save_code_spec` 7. **EXECUTE** with before/after checks 8. **EVALUATE** results with user @@ -132,19 +145,19 @@ The knowledge base contains domain documentation, package references, implementa - `kb_search(query, collection, top_k)` — semantic search - `kb_ingest(folder_path)` — index new documents -## TOOL GENERATION PATTERN +## CODE GENERATION PATTERN -For each data operation, create TWO tools: +For each data operation, create TWO codes: -**Check Tool** — Quantifies the relevant metric +**Check Code** — Quantifies the relevant metric - Accepts: input data path, output report path, threshold parameters - Outputs: JSON report with counts, rates, distributions -**Operation Tool** — Performs the filter/transform +**Operation Code** — Performs the filter/transform - Accepts: input data path, output data path, parameters - Outputs: Transformed data -Write tool scripts to `tools/code/` and register them via `save_tool_spec`. Python dependencies declared in the spec are handled automatically via `uv run --with`. +Write each code's script to `codes//scripts/` and register it via `save_code_spec`. Python dependencies declared in the spec are handled automatically via `uv run --with`. ## PIPELINE RECONSTRUCTION @@ -154,12 +167,12 @@ At any point, you can reconstruct the pipeline from execution records: ## PRINCIPLES -1. **Setup first** — Extend KB and register user tools before iterating -2. **Follow the hierarchy** — Registered tool → KB package tool → Custom implementation +1. **Setup first** — Extend KB and register user codes before iterating +2. **Follow the hierarchy** — Registered code → KB package code → Custom implementation 3. **Explore first** — Search registry and KB before writing new code 4. **Iterate** — One manipulation step at a time; evaluate before proceeding -5. **Generate paired tools** — Both check and operation for each step +5. **Generate paired codes** — Both check and operation for each step 6. **Register everything** — Registry captures the complete pipeline 7. **Audit everything** — Before/after reports for every operation 8. **Confirm with user** — Domain scientist validates approach at each step -9. **No direct data access** — All operations through registered tools +9. **No direct data access** — All operations through registered codes diff --git a/src/dsagt/knowledge.py b/src/dsagt/knowledge.py index b6e9c83..ab3204b 100644 --- a/src/dsagt/knowledge.py +++ b/src/dsagt/knowledge.py @@ -1,8 +1,66 @@ """ -Knowledge base: semantic search over document collections. - -Pluggable embedding backends (local sentence-transformers or OpenAI-compatible API) -and vector-DB backends (FAISS or ChromaDB) with per-collection routing. +Knowledge base — hybrid semantic retrieval over document collections, and the +shared substrate every other DSAGT capability searches against. + +A :class:`KnowledgeBase` holds one or more vector stores (each a single embedder +over many collections) and fuses their results by Reciprocal Rank Fusion (RRF). +Fusion runs at two levels: *within* a collection, dense vector similarity is +blended with a BM25 sparse leg, so exact identifiers and error strings that dense +embeddings under-rank still surface; *across* collections and stores, per- +collection rankings are fused by rank alone — letting results from different +embedding spaces combine without any cross-space score normalization. The same +surface backs document/domain-knowledge retrieval, explicit and episodic memory, +tool-usage provenance, and skills discovery. Design-wise it keeps maintenance +bounded: one embedder per store, a bring-your-own external store is just another +:class:`VectorStore` subclass added to the list, and the heavy document parsers +are imported lazily so only actual ingestion pays for them, not server startup. + +The niche it fills +------------------ +General coding agents have settled retrieval into two modes this KB deliberately +does not compete with — agentic ``grep`` / tool search over *code* (Claude Code, +Cursor, Windsurf et al. dropped codebase vector indexing outright) and *web +search* for open-ended questions. Neither serves **bounded, domain-scoped, +vetted corpora of highly technical prose** — protocols, API references, domain +knowledge, tool-use provenance — which is exactly where hybrid dense+sparse +retrieval still wins and where general agents otherwise fall back to untrusted +web search. That curated, tagged, auditable, hit-it-*first* retrieval is the +niche DSAGT fills to supplement general agents on specialized, highly technical +task loads. The design above is the 2026 best-practice answer to the two +documented failure modes: hybrid search (BM25 recovers the exact identifiers +dense vectors miss) and per-collection domain scoping (the antidote to +"vector-search dilution", where pure-vector accuracy collapses on large +heterogeneous corpora). + +References: + - Code agents drop vector indexing for agentic grep — + https://www.mindstudio.ai/blog/is-rag-dead-what-ai-agents-use-instead ; + https://vadim.blog/claude-code-no-indexing/ + - Hybrid BM25 + dense + reranking, the highest-impact RAG upgrade — + https://www.digitalapplied.com/blog/hybrid-search-bm25-vector-reranking-reference-2026 + - "When More Documents Hurt RAG": vector-search dilution, fixed by + domain-scoped retrieval — https://arxiv.org/abs/2606.11350 + - 2026 enterprise knowledge management — trust / traceability favor curated + KBs over web search — + https://windowsforum.com/threads/2026-enterprise-ai-knowledge-management-from-search-to-governed-agent-workflows.410816/ + +Class map — every edge is ``─ Class``, where ```` is one of +``◇`` holds · ``◆`` owns · ``▷`` inherits (``*`` = one per collection):: + + KnowledgeBase federation infra: ingest pipeline + + │ cross-collection RRF (_rrf_across) + └─◇ VectorStore «abstract» 1..* BYO adapter port: one embedder, many + │ collections; single-collection add/search + └─▷ ChromaVectorStore local Chroma; hybrid dense+BM25 per + │ collection, fused by _rrf_merge + ├─◇ Embedder «abstract» 1 text → vectors; .create() factory + │ │ (one per store; may be shared) + │ ├─▷ LocalEmbedder sentence-transformers, offline + │ └─▷ APIEmbedder OpenAI /v1/embeddings + rate-limit retry + ├─◆ ChromaIndex * dense leg (add/search/save/load) + └─◆ BM25Index * sparse leg (build/search/save/load) + + free fns: _rrf_merge · _rrf_across (rank fusion, used by store + KB) """ from __future__ import annotations @@ -12,65 +70,41 @@ import logging import os import re +import threading import time from abc import ABC, abstractmethod -from dataclasses import asdict, dataclass, field from pathlib import Path from typing import Any, Iterator logger = logging.getLogger(__name__) -# Silence noisy parser libraries. pypdf logs "invalid pdf header", "EOF -# marker not found", and "Ignoring wrong pointing object" at WARNING for -# every malformed PDF it touches — usually image-derived "PDFs" or stripped -# arXiv source. We already count read/parse failures as ``skipped_files`` -# in the ingest result, so these per-file warnings add no signal and a lot -# of noise. Demote to ERROR so genuine PDF library bugs still surface. +# Silence noisy warnings, progress bars from a host of libraries +import warnings as _warnings + for _noisy in ("pypdf", "pypdf2", "PyPDF2", "pdfminer", "fontTools"): logging.getLogger(_noisy).setLevel(logging.ERROR) -# sentence-transformers / huggingface_hub / httpx emit a wall of INFO at -# embedding-model load time: every cached HEAD request to huggingface.co, -# the SentenceTransformer "Loading model from..." line, the cache redirects. -# Demote to WARNING so end-of-session extraction doesn't bury the actual -# extraction result under model-cache validation traffic. Errors still -# surface (genuine network failures, missing models, etc.). -for _noisy in ("httpx", "httpcore", "huggingface_hub", "sentence_transformers", - "transformers", "transformers_modules"): +for _noisy in ( + "httpx", + "httpcore", + "huggingface_hub", + "sentence_transformers", + "transformers", + "transformers_modules", +): logging.getLogger(_noisy).setLevel(logging.WARNING) -# huggingface_hub's "you are sending unauthenticated requests" warning is -# emitted both via ``logger.warning`` and via ``warnings.warn`` (one -# stderr line each). It's informational — anonymous downloads work -# fine, just at lower rate limits — but it surfaces under roo's -# ``--debug`` (which forwards MCP-server stderr) and confuses users -# who don't realise HF_TOKEN is unrelated to whether dsagt works. -# Demote the logger and filter the warning. logging.getLogger("huggingface_hub.utils._http").setLevel(logging.ERROR) -# ``transformers`` BertModel loader writes a "LOAD REPORT" table to stdout -# when missing/unexpected keys differ from the checkpoint — informational -# for ML researchers debugging finetunes, noise for everyone else. Their -# own env var silences it. -import os as _os -_os.environ.setdefault("TRANSFORMERS_VERBOSITY", "error") -# Disable huggingface_hub's tqdm progress bars (the "Loading weights: -# 100%|...| 199/199" line at model load). For cached models this loads -# in <1s — the bar is pure noise. -_os.environ.setdefault("HF_HUB_DISABLE_PROGRESS_BARS", "1") -# Silence the FutureWarning sentence-transformers prints about its own -# renamed method (``get_sentence_embedding_dimension`` → -# ``get_embedding_dimension``); upstream still calls the deprecated form. -import warnings as _warnings +os.environ.setdefault("TRANSFORMERS_VERBOSITY", "error") +os.environ.setdefault("HF_HUB_DISABLE_PROGRESS_BARS", "1") + _warnings.filterwarnings( "ignore", message="The `get_sentence_embedding_dimension` method has been renamed", category=FutureWarning, ) -# Silence huggingface_hub's "Warning: You are sending unauthenticated -# requests to the HF Hub" notice (UserWarning on the ``huggingface_hub`` -# module). Setting HF_TOKEN works too, but we never bake user creds -# into MCP env blocks on disk — see ``_mcp_env_block`` docstring. + _warnings.filterwarnings( "ignore", message=r".*unauthenticated requests to the HF Hub.*", @@ -78,11 +112,13 @@ import numpy as np -# llama_index is intentionally NOT imported at module top. -# It pulls in ~400 transitive submodules and adds ~8s to cold start, which -# matters because dsagt-run imports this module on every tool invocation -# but only ever needs the embedding/search code paths, never the parsers. -# llama_index is lazy-imported inside _chunk_file() and _get_parser(). +# llama_index is intentionally NOT imported at module top: it pulls in ~400 +# transitive submodules and adds ~8s to import. Only document *ingest* +# (kb_ingest / the `dsagt init` KB build) needs the parsers — search, memory, +# provenance indexing, and skills never do — so importing it at module top +# would add ~8s of blocking load to every MCP-server startup (and to +# post-session extraction) for a parser most sessions never invoke. +# Lazy-imported inside _chunk_file() / _get_parser() so only ingest pays for it. from dsagt.observability import ( kb_embed_span, @@ -92,61 +128,85 @@ traced, ) - CODE_LANGUAGES = { - ".py": "python", ".js": "javascript", ".ts": "typescript", - ".jsx": "javascript", ".tsx": "typescript", ".rs": "rust", - ".go": "go", ".java": "java", ".cpp": "cpp", ".c": "c", - ".rb": "ruby", ".php": "php", + ".py": "python", + ".js": "javascript", + ".ts": "typescript", + ".jsx": "javascript", + ".tsx": "typescript", + ".rs": "rust", + ".go": "go", + ".java": "java", + ".cpp": "cpp", + ".c": "c", + ".rb": "ruby", + ".php": "php", } -class BaseEmbeddingClient(ABC): - """Common interface for all embedding backends.""" - - @abstractmethod - def embed(self, texts: list[str]) -> np.ndarray: - """Return float32 array of shape (n_texts, dim), L2-normalized.""" +# =========================================================================== +# Embedder — one per VectorStore +# =========================================================================== - def close(self) -> None: - pass +class Embedder(ABC): + """Common interface for all embedding backends. -class BaseVectorIndex(ABC): - """Common interface for all vector-index backends.""" + A stateless-after-construction leaf: a single :class:`VectorStore` has-a one + ``Embedder``, and the *same* instance can be shared by several stores (e.g. a + Chroma store and a FAISS store on one embedding space). + """ - @abstractmethod - def add(self, embeddings: np.ndarray) -> None: - """Append pre-normalized embeddings.""" + #: Short backend tag used for tracing spans ("local" / "api"). + backend: str = "unknown" + #: Model identifier, for span labelling / collection listing. Subclasses + #: set this in ``__init__``. + model: str | None = None - @abstractmethod - def search(self, query_vec: np.ndarray, k: int) -> tuple[np.ndarray, np.ndarray]: - """Return (scores, int_indices) arrays of length k.""" + @classmethod + def create( + cls, + backend: str, + *, + model: str | None = None, + base_url: str | None = None, + api_key: str | None = None, + device: str | None = None, + ) -> "Embedder": + """Factory: construct the one embedder for a store, with explicit args. - @abstractmethod - def save(self, directory: Path) -> None: - """Persist index to *directory*.""" + The api-vs-local selector that stays after per-collection embedder + routing was removed: a store fixes a single embedder at construction. + Each backend pulls only the parameters it uses — no ``**kwargs`` splat. + """ + backend = (backend or "api").lower() + if backend == "local": + return LocalEmbedder(model=model, device=device) + if backend == "api": + return APIEmbedder(model=model, base_url=base_url, api_key=api_key) + raise ValueError( + f"Unknown embedding backend {backend!r}. Choose from: local, api" + ) - @classmethod @abstractmethod - def load(cls, directory: Path) -> "BaseVectorIndex": - """Load index from *directory*.""" + def embed(self, texts: list[str]) -> np.ndarray: + """Return float32 array of shape (n_texts, dim), L2-normalized.""" - @property - @abstractmethod - def size(self) -> int: - """Number of vectors stored.""" + def close(self) -> None: + pass -class LocalEmbeddingClient(BaseEmbeddingClient): +class LocalEmbedder(Embedder): """sentence-transformers, runs fully offline.""" + backend = "local" + #: Default local model. ``bge-small-en-v1.5`` (33M params, ~130 MB #: on disk, ~250 MB resident) is ~3× faster and ~3× smaller than the #: ``bge-base`` variant we used previously, with ~2 nDCG@10 points #: lower MTEB retrieval score — a hard-to-notice difference for #: typical DSAGT KB sizes (single-digit thousands of chunks). - #: Override via ``embedding.model`` in ``dsagt_config.yaml`` (e.g. + #: Override via ``embedding.model`` in ``.dsagt/config.yaml`` (e.g. #: ``BAAI/bge-large-en-v1.5`` for higher quality at the cost of #: ~10× memory and ~5× CPU). DEFAULT_MODEL = "BAAI/bge-small-en-v1.5" @@ -159,30 +219,36 @@ def __init__( ): import sys from sentence_transformers import SentenceTransformer + model = model or self.DEFAULT_MODEL + self.model = model self.batch_size = batch_size # Probe the HF cache for the model's config. If hit, load with # ``local_files_only=True`` so SentenceTransformer skips the # ETag-validation HEAD requests it would otherwise issue against # huggingface.co — those round-trips are anonymous (HF_TOKEN - # isn't propagated into MCP-server env blocks for cline / roo / + # isn't propagated into MCP-server env blocks for cline / # codex), so they trigger an "unauthenticated requests" warning - # surfaced under roo's ``--debug``. Cache miss path stays + # surfaced under the agent's debug stream. Cache miss path stays # online so first-run downloads still work. try: from huggingface_hub import try_to_load_from_cache + cache_hit = try_to_load_from_cache(model, "config.json") is not None except Exception: cache_hit = False if cache_hit: self._model = SentenceTransformer( - model, device=device, local_files_only=True, + model, + device=device, + local_files_only=True, ) else: print( f" Downloading {model} from HuggingFace " "(set HF_TOKEN for faster throughput)...", - file=sys.stderr, flush=True, + file=sys.stderr, + flush=True, ) self._model = SentenceTransformer(model, device=device) # Belt-and-suspenders: dsagt/__init__.py sets OMP_NUM_THREADS / @@ -191,20 +257,20 @@ def __init__( # vars in some configurations. Cap explicitly here. try: import torch + torch.set_num_threads(int(os.environ.get("OMP_NUM_THREADS", "4"))) except Exception: # noqa: BLE001 — best-effort cap, never fatal pass logger.info( "Loaded local embedding model: %s (dim=%d)", - model, self._model.get_sentence_embedding_dimension(), + model, + self._model.get_embedding_dimension(), ) def embed(self, texts: list[str]) -> np.ndarray: - if not texts: - return np.array([], dtype=np.float32) # Show the progress bar only for non-trivial inputs so single-query # search calls (kb.search → embed([query])) stay silent while large - # ingest / setup-kb runs surface tqdm progress. + # ingest runs surface tqdm progress. show_bar = len(texts) > self.batch_size return self._model.encode( texts, @@ -214,19 +280,18 @@ def embed(self, texts: list[str]) -> np.ndarray: ).astype(np.float32) -# --- rate-limit retry helpers (used by APIEmbeddingClient) ---------------- +# --- rate-limit retry helpers (used by APIEmbedder) ---------------- # -# We can't rely on litellm.num_retries alone for embedding rate limits. Lab -# LiteLLM proxies (notably PNNL's Azure-fronted instance) wrap upstream 429s -# as APIConnectionError carrying the rate-limit error in a JSON message body, -# which defeats litellm's exception-class-based retry detection. Even when -# litellm does retry, its default exponential backoff caps at ~30s total -# across 5 attempts, while Azure's per-minute quota window asks for 60s+ -# between retries. So we own this retry layer ourselves and disable -# litellm's by passing ``num_retries=0`` per call. +# The API embedder talks to an OpenAI-compatible ``/v1/embeddings`` endpoint +# over httpx. We own the retry layer because lab gateways (notably PNNL's +# Azure-fronted instance) return upstream 429s with the quota window in the +# body and ask for 60s+ between retries — longer than a generic exponential +# backoff would wait. Rate-limit / transient errors retry with a +# retry-after-aware wait; auth / bad-request errors fail fast. _RETRY_AFTER_RE = re.compile( - r"retry after (\d+(?:\.\d+)?)\s*seconds?", re.IGNORECASE, + r"retry after (\d+(?:\.\d+)?)\s*seconds?", + re.IGNORECASE, ) @@ -246,31 +311,22 @@ def _extract_retry_after_seconds(message: str, default: float = 60.0) -> float: def _is_retryable_embedding_error(exc: Exception) -> bool: """Decide whether an embedding exception is worth retrying. - Detects rate-limit and transient errors across the various shapes - LiteLLM exposes them in. Authentication / bad-request errors are - explicitly NOT retryable so we fail fast on misconfiguration. + Retries rate-limit (429) and transient server / network errors; + authentication and bad-request errors are explicitly NOT retryable so + we fail fast on misconfiguration. """ - # litellm is a hard dependency — broken install if this fails. - import litellm - - if isinstance(exc, ( - litellm.exceptions.AuthenticationError, - litellm.exceptions.BadRequestError, - litellm.exceptions.NotFoundError, - litellm.exceptions.PermissionDeniedError, - )): - return False - - if isinstance(exc, ( - litellm.exceptions.RateLimitError, - litellm.exceptions.APIConnectionError, - litellm.exceptions.Timeout, - litellm.exceptions.ServiceUnavailableError, - litellm.exceptions.InternalServerError, - )): + import httpx + + if isinstance(exc, httpx.HTTPStatusError): + code = exc.response.status_code + # 408 request timeout, 425 too early, 429 rate limit, 5xx server. + return code in (408, 425, 429, 500, 502, 503, 504) + + # Connect / read timeouts and other transport failures are transient. + if isinstance(exc, (httpx.TimeoutException, httpx.TransportError)): return True - # Some lab proxies wrap rate limits as plain Exceptions with the upstream + # Some gateways wrap rate limits in a plain Exception with the upstream # body in str(exc). Fall back to message inspection. msg = str(exc).lower() return any(k in msg for k in ("rate limit", "ratelimiterror", "429", "throttl")) @@ -279,39 +335,48 @@ def _is_retryable_embedding_error(exc: Exception) -> bool: def _retry_wait_seconds(exc: Exception, attempt: int) -> float: """How long to sleep before the next retry attempt. - Rate-limit errors get the upstream-suggested wait (or 60s default). - Other transient errors get exponential backoff capped at 30s. + Honors a ``Retry-After`` header (or a retry-after hint in the body) on + rate-limit responses; other transient errors get exponential backoff + capped at 30s. """ + import httpx + + if isinstance(exc, httpx.HTTPStatusError): + retry_after = exc.response.headers.get("retry-after") + if retry_after: + try: + return float(retry_after) + except ValueError: + pass + if exc.response.status_code == 429: + return _extract_retry_after_seconds(exc.response.text, default=60.0) msg = str(exc).lower() if "rate limit" in msg or "429" in msg or "throttl" in msg: return _extract_retry_after_seconds(str(exc), default=60.0) - return min(2.0 ** attempt, 30.0) + return min(2.0**attempt, 30.0) -class APIEmbeddingClient(BaseEmbeddingClient): - """OpenAI-compatible embedding client backed by LiteLLM. +class APIEmbedder(Embedder): + """Embedding client for an OpenAI-compatible ``/v1/embeddings`` endpoint. - LiteLLM normalizes a wide set of providers (OpenAI, Azure, Bedrock, - Vertex, Cohere, Voyage, Ollama, Together, etc.) behind a single - ``embedding(...)`` call. This client adds two things on top: + Talks to the endpoint directly over httpx — no provider-abstraction + layer. The model string is sent verbatim in the request body, so + gateway aliases (``text-embedding-3-small-project``) and HuggingFace- + style names with slashes (``nomic-ai/nomic-embed-text-v1``) pass + through unchanged. Two things on top of the raw call: 1. **Manual batching** of large inputs into ``batch_size``-sized requests so a single rate-limit hit only loses one batch and the user can see per-batch progress in the logs. 2. **Explicit rate-limit retry** with retry-after-aware backoff - (see ``_embed_batch_with_retry``). This is the layer that keeps - large ``kb_ingest`` and ``dsagt-setup-kb`` runs alive in the face - of TPM quotas — litellm's built-in retry isn't sufficient because - lab proxies wrap upstream 429s in a way that defeats its - exception-class detection, and its backoff is too short for the - 60-second quota windows Azure enforces. - - The model string passed to LiteLLM determines provider routing. Bare - model names get an ``openai_like/`` prefix so the lab proxy receives - the model alias unchanged. + (see ``_embed_batch_with_retry``) — what keeps large ``kb_ingest`` + and ``dsagt init`` KB builds alive against the 60-second quota + windows lab gateways enforce. """ + backend = "api" + def __init__( self, model: str | None = None, @@ -320,12 +385,17 @@ def __init__( timeout: float = 300.0, batch_size: int = 100, ): - self.model = model or os.getenv("EMBEDDING_MODEL", "text-embedding-3-small-project") - self.base_url = base_url or os.getenv("OPENAI_BASE_URL") + import httpx + + self.model = model or os.getenv( + "EMBEDDING_MODEL", "text-embedding-3-small-project" + ) + self.base_url = ( + base_url or os.getenv("EMBEDDING_BASE_URL") or os.getenv("OPENAI_BASE_URL") + ) # EMBEDDING_API_KEY is the canonical name; LLM_API_KEY/OPENAI_API_KEY # are legacy fallbacks from when the embedding endpoint shared the - # LLM endpoint's auth. When embeddings route through dsagt's proxy - # (the default), this is a sentinel — the proxy holds the real key. + # LLM endpoint's auth. self.api_key = ( api_key or os.getenv("EMBEDDING_API_KEY") @@ -336,30 +406,21 @@ def __init__( self.batch_size = batch_size if not self.base_url: - raise ValueError("Embedding API base URL required via argument or OPENAI_BASE_URL env var") + raise ValueError( + "Embedding API base URL required via argument or " + "EMBEDDING_BASE_URL / OPENAI_BASE_URL env var" + ) if not self.api_key: - raise ValueError("API key required via argument or LLM_API_KEY env var") - - # Always prefix with ``openai_like`` so LiteLLM dispatches to the - # OpenAI-wire-protocol client pointed at our ``base_url`` — the rest - # of DSAGT already assumes the embedding endpoint is OpenAI-compat, - # so there's no valid case for a different provider here. - # - # ``openai_like`` (not ``openai``) is deliberate on two counts: (a) - # ``openai`` matches LiteLLM's canonical model registry and silently - # normalizes aliases like ``text-embedding-3-small-project`` down to - # ``text-embedding-3-small``, which lab proxies then reject; (b) - # ``openai_like`` is the documented escape hatch for "endpoint - # speaks OpenAI but is not OpenAI" — it forwards the model string - # verbatim, preserving HuggingFace-style names (``lbl/nomic-embed-text``, - # ``nomic-ai/nomic-embed-text-v1``) whose slashes are part of the - # model identifier, not a provider prefix. - self._litellm_model = f"openai_like/{self.model}" + raise ValueError( + "API key required via argument or EMBEDDING_API_KEY env var" + ) - def embed(self, texts: list[str]) -> np.ndarray: - if not texts: - return np.array([], dtype=np.float32) + # ``base_url`` is the OpenAI-style root (typically ending in ``/v1``); + # the embeddings route hangs off it. + self._embeddings_url = self.base_url.rstrip("/") + "/embeddings" + self._client = httpx.Client(timeout=timeout) + def embed(self, texts: list[str]) -> np.ndarray: # Single small input: one call, no batch logging. if len(texts) <= self.batch_size: return self._embed_batch_with_retry(texts) @@ -369,11 +430,13 @@ def embed(self, texts: list[str]) -> np.ndarray: n_batches = (len(texts) + self.batch_size - 1) // self.batch_size batches: list[np.ndarray] = [] for i in range(0, len(texts), self.batch_size): - batch = texts[i:i + self.batch_size] + batch = texts[i : i + self.batch_size] batch_num = i // self.batch_size + 1 logger.info( "Embedding batch %d/%d (%d texts)", - batch_num, n_batches, len(batch), + batch_num, + n_batches, + len(batch), ) batches.append(self._embed_batch_with_retry(batch)) return np.vstack(batches) @@ -386,24 +449,22 @@ def _embed_batch_with_retry( """Embed one batch with explicit rate-limit and transient-error retry. See the module-level rate-limit retry helpers for the rationale. - We pass ``num_retries=0`` to litellm so its built-in retry doesn't - race with ours and we don't end up double-retrying with conflicting - backoff strategies. """ - import litellm + headers = { + "Authorization": f"Bearer {self.api_key}", + "Content-Type": "application/json", + } + payload = {"model": self.model, "input": texts} attempt = 0 while True: attempt += 1 try: - response = litellm.embedding( - model=self._litellm_model, - input=texts, - api_base=self.base_url, - api_key=self.api_key, - timeout=self.timeout, - num_retries=0, + resp = self._client.post( + self._embeddings_url, json=payload, headers=headers ) + resp.raise_for_status() + data = resp.json() break # success except Exception as exc: if not _is_retryable_embedding_error(exc) or attempt >= max_attempts: @@ -412,69 +473,35 @@ def _embed_batch_with_retry( logger.warning( "Embedding error (attempt %d/%d). " "Waiting %.0fs then retrying. Cause: %s", - attempt, max_attempts, wait, str(exc)[:200], + attempt, + max_attempts, + wait, + str(exc)[:200], ) time.sleep(wait) - # Response shape mirrors OpenAI: response.data is a list of objects - # with .embedding (or ['embedding'] for dict access). - sorted_data = sorted( - response.data, - key=lambda d: d["index"] if isinstance(d, dict) else d.index, - ) - vectors = [ - d["embedding"] if isinstance(d, dict) else d.embedding - for d in sorted_data - ] + # Response shape is OpenAI's: ``data`` is a list of ``{"index": i, + # "embedding": [...]}`` objects, not guaranteed to be in input order. + sorted_data = sorted(data["data"], key=lambda d: d["index"]) + vectors = [d["embedding"] for d in sorted_data] return np.array(vectors, dtype=np.float32) def close(self) -> None: - # No client object to close anymore — LiteLLM owns its own pool. - pass - -class FAISSIndex(BaseVectorIndex): - """Inner-product FAISS flat index (cosine similarity on L2-normed vecs).""" + self._client.close() - _FILENAME = "index.faiss" - def __init__(self, dim: int | None = None): - self._dim = dim - self._index = None # lazy-init on first add - - def _ensure_init(self, dim: int) -> None: - if self._index is None: - import faiss - self._dim = dim - self._index = faiss.IndexFlatIP(dim) - - def add(self, embeddings: np.ndarray) -> None: - self._ensure_init(embeddings.shape[1]) - self._index.add(embeddings) - - def search(self, query_vec: np.ndarray, k: int) -> tuple[np.ndarray, np.ndarray]: - k = min(k, self._index.ntotal) - scores, indices = self._index.search(query_vec.reshape(1, -1), k) - return scores[0], indices[0] - - def save(self, directory: Path) -> None: - import faiss - faiss.write_index(self._index, str(directory / self._FILENAME)) - - @classmethod - def load(cls, directory: Path) -> "FAISSIndex": - import faiss - obj = cls() - obj._index = faiss.read_index(str(directory / cls._FILENAME)) - obj._dim = obj._index.d - return obj +# =========================================================================== +# ChromaIndex — single-collection dense vector wrapper (ChromaVectorStore's leg) +# =========================================================================== - @property - def size(self) -> int: - return 0 if self._index is None else self._index.ntotal +class ChromaIndex: + """ChromaDB-backed dense index for ONE collection. Requires ``chromadb``. -class ChromaIndex(BaseVectorIndex): - """ChromaDB-backed index. Requires ``chromadb`` package.""" + A plain helper owned by :class:`ChromaVectorStore` — not a pluggable + "vector-DB backend". Stores cosine-space HNSW vectors plus a positional + id list so returned integer indices line up with the chunk list. + """ _META_FILE = "chroma_ids.json" # maps int position → chroma id @@ -495,7 +522,12 @@ def __init__(self, collection_name: str, persist_dir: Path | None = None): # ChromaDB stores ids internally; we maintain a positional list so that # returned integer indices are consistent with the chunk list. - def add(self, embeddings: np.ndarray, metadatas: list[dict] | None = None) -> None: + def add( + self, + embeddings: np.ndarray, + metadatas: list[dict] | None = None, + documents: list[str] | None = None, + ) -> None: start = len(self._ids) new_ids = [str(start + i) for i in range(len(embeddings))] self._ids.extend(new_ids) @@ -508,20 +540,32 @@ def add(self, embeddings: np.ndarray, metadatas: list[dict] | None = None) -> No batch_size = 5000 for i in range(0, len(new_ids), batch_size): kwargs: dict = { - "ids": new_ids[i:i + batch_size], - "embeddings": embeddings_list[i:i + batch_size], + "ids": new_ids[i : i + batch_size], + "embeddings": embeddings_list[i : i + batch_size], } if metadatas is not None: - kwargs["metadatas"] = metadatas[i:i + batch_size] + kwargs["metadatas"] = metadatas[i : i + batch_size] + # Store the chunk text as the Chroma *document* too — that's what + # ``where_document`` ($contains / $regex) matches against. + if documents is not None: + kwargs["documents"] = documents[i : i + batch_size] self._col.add(**kwargs) - def search(self, query_vec: np.ndarray, k: int, where: dict | None = None) -> tuple[np.ndarray, np.ndarray]: + def search( + self, + query_vec: np.ndarray, + k: int, + where: dict | None = None, + where_document: dict | None = None, + ) -> tuple[np.ndarray, np.ndarray]: k = min(k, len(self._ids)) if k == 0: return np.array([], dtype=np.float32), np.array([], dtype=np.int64) query_kwargs: dict = {"query_embeddings": [query_vec.tolist()], "n_results": k} if where is not None: query_kwargs["where"] = where + if where_document is not None: + query_kwargs["where_document"] = where_document results = self._col.query(**query_kwargs) chroma_ids = results["ids"][0] distances = results["distances"][0] # cosine distance (0=identical) @@ -547,6 +591,10 @@ def size(self) -> int: return len(self._ids) +# =========================================================================== +# BM25 — the sparse leg fused with dense vectors per collection +# =========================================================================== + # BM25 sparse-retrieval token splitter. Splits on every non-alphanumeric # character so ``snake_case`` and ``kebab-case`` identifiers fan out into # their parts; this matters because much of what an agent searches the KB @@ -563,15 +611,15 @@ def _bm25_tokenize(text: str) -> list[str]: class BM25Index: """Sparse BM25 keyword index over a collection's chunk texts. - Maintained alongside the dense vector index so :meth:`KnowledgeBase.search` - can fuse the two via Reciprocal Rank Fusion. Stored as a single - pickle file (``bm25.pkl``) per collection. + Maintained alongside the dense vector index so single-collection search + can fuse the two via Reciprocal Rank Fusion. Stored as a single pickle + file (``bm25.pkl``) per collection — its presence is also what marks a + collection as *hybrid*. - Rebuilt from scratch on every ``add_entries`` / ``append`` because BM25 - IDF stats are corpus-global — there is no cheap incremental update. - For DSAGT corpus sizes (single-digit thousands of chunks) the rebuild - is millisecond-scale, so this is fine. Watch the cost if a collection - grows past tens of thousands of entries. + Rebuilt from scratch on every write because BM25 IDF stats are + corpus-global — there is no cheap incremental update. For DSAGT corpus + sizes (single-digit thousands of chunks) the rebuild is millisecond-scale. + Watch the cost if a collection grows past tens of thousands of entries. """ _FILENAME = "bm25.pkl" @@ -582,6 +630,7 @@ def __init__(self): def build(self, texts: list[str]) -> None: from rank_bm25 import BM25Okapi + if not texts: self._bm25 = None self._n = 0 @@ -613,12 +662,14 @@ def search(self, query: str, k: int) -> tuple[np.ndarray, np.ndarray]: def save(self, directory: Path) -> None: import pickle + with open(directory / self._FILENAME, "wb") as f: pickle.dump({"bm25": self._bm25, "n": self._n}, f) @classmethod def load(cls, directory: Path) -> "BM25Index": import pickle + obj = cls() path = directory / cls._FILENAME if path.exists(): @@ -633,6 +684,11 @@ def size(self) -> int: return self._n +# =========================================================================== +# Reciprocal Rank Fusion — within a collection AND across collections/stores +# =========================================================================== + + def _rrf_merge( rankings: list[list[int]], k: int = 60, @@ -658,383 +714,685 @@ def _rrf_merge( return sorted(scores.items(), key=lambda kv: kv[1], reverse=True) -# map short names to embedding client constructors -EMBEDDER_REGISTRY: dict[str, type[BaseEmbeddingClient]] = { - "local": LocalEmbeddingClient, - "api": APIEmbeddingClient, -} +def _chunk_key(chunk: dict) -> tuple: + """Stable identity for a result chunk, for cross-collection RRF. -# map short names to vector-index constructors -VECTORINDEX_REGISTRY: dict[str, type[BaseVectorIndex]] = { - "faiss": FAISSIndex, - "chroma": ChromaIndex, -} + Different collections (and, later, different stores/embedders) live in + incomparable score spaces, so fusion is rank-only. We key on the chunk's + own metadata rather than score; ``id``-based fallback keeps distinct chunks + distinct even when metadata is sparse. + """ + meta = chunk.get("metadata", {}) + return ( + meta.get("collection"), + meta.get("source_file"), + meta.get("chunk_index"), + chunk.get("text", "")[:64], + ) -def _make_embedder(backend: str, **kwargs) -> BaseEmbeddingClient: - backend = backend.lower() - if backend not in EMBEDDER_REGISTRY: - raise ValueError(f"Unknown embedding backend '{backend}'. " - f"Choose from: {list(EMBEDDER_REGISTRY)}") - return EMBEDDER_REGISTRY[backend](**kwargs) +def _rrf_across(result_lists: list[list[dict]], k: int = 60) -> list[dict]: + """Rank-fuse per-collection result lists (each ``[{chunk, score}, ...]``). + Rank-only RRF across the lists — no cross-space score normalization, which + is exactly what lets results from different embedding spaces fuse correctly. + Returns a single list of result dicts (``score`` replaced by the RRF score), + sorted by descending fused score. + """ + scores: dict[tuple, float] = {} + rep: dict[tuple, dict] = {} + for results in result_lists: + for rank, r in enumerate(results): + key = _chunk_key(r["chunk"]) + scores[key] = scores.get(key, 0.0) + 1.0 / (k + rank + 1) + rep.setdefault(key, r) + ordered = sorted(scores.items(), key=lambda kv: kv[1], reverse=True) + return [{**rep[key], "score": float(score)} for key, score in ordered] + + +# --------------------------------------------------------------------------- +# Recency weighting (episodic session_memory only) +# --------------------------------------------------------------------------- + +#: The one collection recency applies to. Mirrors +#: ``memory.SESSION_MEMORY_COLLECTION`` as a literal to avoid a knowledge→memory +#: import (memory imports knowledge, not the reverse). +_RECENCY_COLLECTION = "session_memory" + +#: Max fractional boost a brand-new fact gets over its raw relevance. A *boost*, +#: never a penalty: recency only lifts recent facts, so a strongly-relevant old +#: fact (e.g. a day-1 threshold that never changed) is never buried — it keeps +#: its full relevance score while a same-relevance newer fact edges ahead. +_RECENCY_BOOST = 0.5 + + +def _apply_recency( + results: list[dict], + half_life_days: float, + now: float, + boost: float = _RECENCY_BOOST, +) -> list[dict]: + """Re-rank ``[{chunk, score}, ...]`` by relevance × a recency factor. + + ``factor = 1 + boost · 2^(-age / half_life)`` — newest facts get up to + ``1+boost``, decaying toward ``1`` (no change) with the given half-life. + Facts without a numeric ``ts_epoch`` get factor ``1.0`` (unweighted), never + dropped. Pure + side-effect-free so it's unit-testable without an embedder. + """ + hl_seconds = max(1.0, half_life_days * 86400.0) + weighted = [] + for r in results: + ts = r["chunk"].get("metadata", {}).get("ts_epoch") + if ts is None: + factor = 1.0 + else: + age = max(0.0, now - float(ts)) + factor = 1.0 + boost * (0.5 ** (age / hl_seconds)) + weighted.append({**r, "score": r["score"] * factor, "recency_factor": factor}) + weighted.sort(key=lambda r: r["score"], reverse=True) + return weighted + + +# =========================================================================== +# VectorStore — one embedder, many collections (the BYO adapter port) +# =========================================================================== + + +class VectorStore(ABC): + """One :class:`Embedder`, many collections: store + single-collection search. + + The adapter contract a BYO store implements: subclass this and wrap your + own backend. :class:`KnowledgeBase` holds a list of these and fuses their + collections. Heterogeneity (mixed embedding spaces) is expressed by having + *several* stores in the list, never by routing within one store. + + TODO (deferred to the second concrete store — BYO/FAISS, see + knowledge-base-plan.md): most of :class:`ChromaVectorStore` is backend- + agnostic — the BM25 sparse leg, the dense+sparse hybrid fusion, the + ``chunks.jsonl`` payload, ``_normalize``. Only the dense index + (add/query/persist over Chroma's HNSW) is Chroma-specific. When a second + store lands, lift that reusable machinery into a shared base that defers + only the dense primitives to subclasses, so a FAISS/BYO store reuses it + rather than copying it. Not extracted now: with one impl the seam can't be + validated. Dense-only stores (external adapters with no local sparse leg) + are modelled as a *separate store type*, not a per-instance flag. + """ -def _make_index(backend: str, **kwargs) -> BaseVectorIndex: - backend = backend.lower() - if backend not in VECTORINDEX_REGISTRY: - raise ValueError(f"Unknown vector-index backend '{backend}'. " - f"Choose from: {list(VECTORINDEX_REGISTRY)}") - return VECTORINDEX_REGISTRY[backend](**kwargs) + embedder: Embedder + @property + @abstractmethod + def collections(self) -> list[str]: + """Names of the collections this store hosts.""" -@dataclass -class CollectionRoute: - """ - Describes which embedding client and vector-index backend to use for a - named collection. - - Parameters - ---------- - embedding_backend : str - Key in EMBEDDER_REGISTRY, e.g. ``"local"`` or ``"api"``. - vector_db : str - Key in VECTORINDEX_REGISTRY, e.g. ``"faiss"`` or ``"chroma"``. - embedder_kwargs : dict - Extra kwargs forwarded to the embedding client constructor. - index_kwargs : dict - Extra kwargs forwarded to the vector-index constructor. - description : str - Human-readable note shown in ``list_collections()`` when no - DESCRIPTION.md file exists. - hybrid : bool - If True (default), maintain a BM25 sparse index alongside the - dense vector index and fuse results via Reciprocal Rank Fusion - on every search. Costs one pickle file per collection and a - rebuild on every write; gains are exact-token recall (API - names, error strings, identifiers) that cosine on dense - embeddings tends to under-rank. Set False to skip BM25 entirely - for a collection. - """ + @abstractmethod + def has_collection(self, name: str) -> bool: + """Whether *name* is hosted by this store.""" - embedding_backend: str = "api" - vector_db: str = "chroma" - embedder_kwargs: dict = field(default_factory=dict) - index_kwargs: dict = field(default_factory=dict) - description: str = "" - hybrid: bool = True + @abstractmethod + def collection_info(self, name: str) -> dict: + """Listing metadata for *name* (description, embedder, vector_db).""" - def to_dict(self) -> dict: - return asdict(self) + @abstractmethod + def embed(self, texts: list[str]) -> np.ndarray: + """Embed *texts* with this store's embedder (L2-normalized).""" - @classmethod - def from_dict(cls, d: dict) -> "CollectionRoute": - return cls(**d) + @abstractmethod + def add_chunks( + self, + collection: str, + chunks: list[dict], + return_embeddings: bool = False, + ) -> dict: + """Embed + store pre-built ``{text, metadata}`` chunks into *collection*.""" + @abstractmethod + def add_entries( + self, + texts: list[str], + collection: str, + metadatas: list[dict] | None = None, + return_embeddings: bool = False, + ) -> dict: + """Embed + store raw *texts* (with optional metadata) into *collection*.""" -class KnowledgeBase: - """ - Collection-based document retrieval with pluggable embedder & vector-DB routing. + @abstractmethod + def search( + self, + query: str, + collection: str, + top_k: int, + where: dict | None = None, + where_document: dict | None = None, + ) -> list[dict]: + """Single-collection hybrid search → ``[{chunk, score}, ...]``.""" - Each collection lives in ``//`` and may use a - *different* embedding model and vector-index backend. Routing is controlled - by a dictionary of :class:`CollectionRoute` objects keyed on collection name. + def close(self) -> None: + pass - Quick-start - ----------- - .. code-block:: python - from knowledge_base import KnowledgeBase, CollectionRoute - - kb = KnowledgeBase( - index_dir="./kb_store", - routes={ - # GPU-heavy research docs → local model + FAISS - "research": CollectionRoute( - embedding_backend="local", - vector_db="faiss", - embedder_kwargs={"model": "BAAI/bge-base-en-v1.5"}, - ), - # Fast live ingestion → API + Chroma - "support": CollectionRoute( - embedding_backend="api", - vector_db="chroma", - embedder_kwargs={"model": "text-embedding-3-small"}, - ), - }, - ) +class ChromaVectorStore(VectorStore): + """Local-ChromaDB store: one embedder, many collections under *index_dir*. - kb.ingest("./docs/research", "research") - results = kb.search("transformer architecture", "research") + Each collection lives in ``//`` as a Chroma persistent + collection (dense), a ``chunks.jsonl`` payload, and a ``bm25.pkl`` sparse + leg. The embedder is fixed for the whole store; it is built lazily (via + :meth:`Embedder.create`) so server startup can background-load it, or injected + directly for reuse across stores. - Embedding backends - ------------------ - ``"local"`` - sentence-transformers, no network required. - ``"api"`` - OpenAI-compatible REST API. Reads ``LLM_API_KEY`` env var. - - Vector-index backends - --------------------- - ``"chroma"`` - ChromaDB HNSW index (default). Supports metadata filtering - and incremental updates. - ``"faiss"`` - Flat inner-product index. Faster for small static collections - but no metadata filtering. + The local store is **unconditionally hybrid**: it writes a BM25 sparse leg on + every write and fuses dense + sparse on every (unfiltered) search. Dense-only + retrieval is not a toggle here — a metadata-``where`` filter falls back to + dense-only for that one query (BM25 has no filter equivalent), and dense-only + *stores* (external/BYO adapters with no local sparse leg) are a separate store + type, not a flag on this one. """ - FILE_TYPES = [ - "pdf", "md", "rst", "txt", "py", "docx", - "json", "yaml", "yml", - # Packaging metadata: agents reading this index need to know which - # version of a library to install when registering tools that - # depend on it. pyproject.toml is the modern standard; setup.cfg - # is still common in older codebases. - "toml", "cfg", - ] - _ROUTE_FILE = "route.json" - def __init__( self, index_dir: str | Path, - chunk_size: int = 1024, - chunk_overlap: int = 128, - rerank_model: str = "BAAI/bge-reranker-v2-m3", - default_rerank: bool = False, - # Global default route (used when collection has no specific route) - default_embedder: str | None = None, - default_index: str | None = None, - embedder_kwargs: dict | None = None, - # Per-collection routing registry - routes: dict[str, CollectionRoute] | None = None, + embedder: Embedder | None = None, + *, + backend: str = "api", + model: str | None = None, + base_url: str | None = None, + api_key: str | None = None, + device: str | None = None, ): self.index_dir = Path(index_dir) - self.chunk_size = chunk_size - self.chunk_overlap = chunk_overlap - self.rerank_model = rerank_model - self.default_rerank = default_rerank - self.index_dir.mkdir(parents=True, exist_ok=True) - # Build default route from explicit kwargs (config-driven). - # No env var reads — callers pass config values through. - self._default_route = CollectionRoute( - embedding_backend=default_embedder or "api", - vector_db=default_index or "chroma", - embedder_kwargs=embedder_kwargs or {}, - ) + # Embedder is either injected (shareable across stores) or built lazily + # from explicit args. Lazy build keeps the ~5-10s sentence-transformers + # load off the hot path (see preload via KnowledgeBase). backend/model + # are kept as lightweight labels so listing never forces a model load. + self._embedder = embedder + if embedder is not None: + self._backend, self._model = embedder.backend, embedder.model + else: + self._backend, self._model = backend, model + self._base_url, self._api_key, self._device = base_url, api_key, device + # Serializes embedder construction so a background preload and a + # foreground first-query call don't race and double-load the model. + self._embedder_lock = threading.Lock() + + # Collection runtime caches: name → (ChromaIndex, chunks) and name → BM25. + self._cache: dict[str, tuple[ChromaIndex, list[dict]]] = {} + self._bm25_cache: dict[str, BM25Index] = {} - # Per-collection route registry - self._routes: dict[str, CollectionRoute] = routes or {} + # -- embedder ------------------------------------------------------------ - # Shared embedder cache: embedder_key → client instance - # Key is "|" so identical configs share one client. - self._embedder_cache: dict[str, BaseEmbeddingClient] = {} - # Serializes embedder construction so a background preload (see - # ``preload_default_embedder``) and a foreground first-query call - # don't race and double-load the model. Construction is rare; - # the lock isn't a hot-path concern. - import threading as _threading - self._embedder_lock = _threading.Lock() + @property + def embedder(self) -> Embedder: # type: ignore[override] + with self._embedder_lock: + if self._embedder is None: + self._embedder = Embedder.create( + self._backend, + model=self._model, + base_url=self._base_url, + api_key=self._api_key, + device=self._device, + ) + return self._embedder - # Collection runtime cache: name → (BaseVectorIndex, list[dict]) - self._cache: dict[str, tuple[BaseVectorIndex, list[dict]]] = {} + @property + def backend(self) -> str: + return self._backend - # BM25 sparse-index cache for hybrid-enabled collections. Loaded - # from disk on first search, rebuilt and resaved on every write. - self._bm25_cache: dict[str, BM25Index] = {} + @property + def model(self) -> str | None: + return self._model - # Per-file-type parser cache. Constructing a CodeSplitter loads the - # tree-sitter language definition (~25ms each), so a 300-file ingest - # without this cache pays 7-8 seconds in parser construction alone. - # Parsers are stateless once built; safe to share across files. - self._parsers: dict[str, Any] = {} + def embed(self, texts: list[str]) -> np.ndarray: + with kb_embed_span(self.backend, self.model, len(texts)): + return self._normalize(self.embedder.embed(texts)) - # Per-ingest counter for files that _chunk_file skipped due to - # read/parse errors. ingest()/append() reset this before the loop - # and surface it in the result dict so users notice when a non-zero - # number of files are silently being dropped. - self._chunk_skip_count: int = 0 + @staticmethod + def _normalize(arr: np.ndarray) -> np.ndarray: + norms = np.linalg.norm(arr, axis=-1, keepdims=True) + return arr / np.where(norms > 0, norms, 1) - self._reranker = None + # -- collection inventory ------------------------------------------------ - @staticmethod - def _stale_index_message(collection: str, exc: Exception) -> str | None: - """Return a user-friendly hint when *exc* looks like a dim mismatch. + @property + def collections(self) -> list[str]: + return [ + p.name + for p in self.index_dir.iterdir() + if p.is_dir() and (p / "chroma_ids.json").exists() + ] - Chroma raises ``InvalidDimensionException`` and FAISS raises - ``RuntimeError``/``AssertionError`` with "dimension" in the - message when a query vector's dimension doesn't match the index's. - That happens when the embedder used to BUILD the index has been - swapped for one with a different output dim — typically because - the user changed ``embedding.backend`` or ``embedding.model`` in - their config after init. Detect on message-shape (no chroma - import here) and return the recovery instructions. - """ - msg = str(exc).lower() - if "dimension" in msg or "dim mismatch" in msg or "shape" in msg: - return ( - f"Embedding dimension mismatch in collection {collection!r}. " - f"The index was built with a different embedder than the one " - f"currently configured. Run `dsagt setup-kb` with the active " - f"embedding settings, then drop the stale collection: " - f"`rm -rf /kb_index/{collection}` and re-init/start " - f"the project." - ) - return None + def has_collection(self, name: str) -> bool: + return (self.index_dir / name / "chroma_ids.json").exists() + + def collection_info(self, name: str) -> dict: + coll_dir = self.index_dir / name + desc_path = coll_dir / "DESCRIPTION.md" + description = desc_path.read_text() if desc_path.exists() else "" + return { + "name": name, + "description": description, + "embedding_backend": self.backend, + "embedding_model": self.model or "default", + "vector_db": "chroma", + } - # route management + # -- writes -------------------------------------------------------------- - def register_route(self, collection: str, route: CollectionRoute) -> None: - """Register (or update) a routing rule for *collection*. + def add_chunks( + self, + collection: str, + chunks: list[dict], + return_embeddings: bool = False, + ) -> dict: + """Embed + store pre-built ``{text, metadata}`` chunks. - Updates the in-memory routing table AND persists ``route.json`` - to disk so the mapping survives server restarts. + The shared write path for both document ingest (chunks carry + ``source_file`` / ``file_type``) and structured entries. Rebuilds the + BM25 sparse leg on every write (the store is unconditionally hybrid). """ - self._routes[collection] = route coll_dir = self.index_dir / collection - coll_dir.mkdir(exist_ok=True) - (coll_dir / self._ROUTE_FILE).write_text(json.dumps(route.to_dict(), indent=2)) - logger.info("Registered route for '%s': embedder=%s index=%s", - collection, route.embedding_backend, route.vector_db) + coll_dir.mkdir(parents=True, exist_ok=True) - def _get_route(self, collection: str) -> CollectionRoute: - return self._routes.get(collection, self._default_route) + texts = [c["text"] for c in chunks] + metadatas = [c.get("metadata", {}) for c in chunks] + embeddings = self.embed(texts) - def _get_embedder(self, route: CollectionRoute) -> BaseEmbeddingClient: - """Return (possibly cached) embedder matching *route*. + if (coll_dir / "chunks.jsonl").exists(): + index, existing_chunks = self._load(collection) + else: + index = ChromaIndex(collection_name=coll_dir.name, persist_dir=coll_dir) + existing_chunks = [] - Routes carry per-collection overrides (e.g. a specific model name); - credentials (api_key, base_url) live on the kb's default route. - Merge so explicit routes inherit defaults unless they override. - """ - merged = {**self._default_route.embedder_kwargs, **route.embedder_kwargs} - key = f"{route.embedding_backend}|{sorted(merged.items())}" - with self._embedder_lock: - if key not in self._embedder_cache: - self._embedder_cache[key] = _make_embedder( - route.embedding_backend, **merged - ) - return self._embedder_cache[key] + try: + index.add(embeddings, metadatas=metadatas, documents=texts) + except Exception as e: + hint = self._stale_index_message(collection, e) + if hint: + raise RuntimeError(hint) from e + raise + index.save(coll_dir) - def preload_default_embedder(self) -> None: - """Kick off embedder construction in a daemon thread. - - Called at MCP server startup so the heavy load — sentence- - transformers import, ``SentenceTransformer(...)`` model load, - etc., ~5–10s for the default local backend — happens in - parallel with the rest of server bootstrap. By the time the - agent makes its first kb call, the embedder is cached and the - call completes immediately instead of paying the load cost. - - Failure inside the thread is swallowed: the model-load error - will surface again (with full traceback) on the first real - embedding call, where the agent can surface it usefully. The - background thread shouldn't crash the server before the agent - even gets a chance to query. - """ - import threading + with open(coll_dir / "chunks.jsonl", "a") as f: + for chunk in chunks: + f.write(json.dumps(chunk) + "\n") - def _load() -> None: - try: - self._get_embedder(self._default_route) - except Exception as e: - logger.warning("Background embedder preload failed: %s", e) + all_chunks = existing_chunks + chunks + self._rebuild_bm25(collection, all_chunks) + self._cache[collection] = (index, all_chunks) - threading.Thread( - target=_load, name="dsagt-embedder-preload", daemon=True, - ).start() + result = { + "collection": collection, + "entries_added": len(chunks), + "total_entries": len(all_chunks), + } + if return_embeddings: + result["embeddings"] = embeddings + return result - @property - def collections(self) -> list[str]: - return [p.name for p in self.index_dir.iterdir() if p.is_dir() and - ((p / "index.faiss").exists() or (p / "chroma_ids.json").exists())] + def add_entries( + self, + texts: list[str], + collection: str, + metadatas: list[dict] | None = None, + return_embeddings: bool = False, + ) -> dict: + """Add raw *texts* (no parsing/chunking) as entries in *collection*.""" + coll_dir = self.index_dir / collection + existing = 0 + if (coll_dir / "chunks.jsonl").exists(): + _, existing_chunks = self._load(collection) + existing = len(existing_chunks) - def list_collections(self) -> list[dict]: - result = [] - for name in self.collections: - coll_dir = self.index_dir / name - desc_path = coll_dir / "DESCRIPTION.md" - - # Ensure the persisted route is loaded into _routes so we always - # report the actual model/backend, not the server default. - if name not in self._routes: - persisted = self._load_route(name) - if persisted: - self._routes[name] = persisted - - route = self._routes.get(name, self._default_route) - description = ( - desc_path.read_text() if desc_path.exists() - else route.description + chunks = [] + for i, text in enumerate(texts): + entry_meta = metadatas[i] if metadatas else {} + chunks.append( + { + "text": text, + "metadata": { + "collection": collection, + "source_file": "entry", + "chunk_index": existing + i, + **entry_meta, + }, + } ) - result.append({ - "name": name, - "description": description, - "embedding_backend": route.embedding_backend, - "embedding_model": route.embedder_kwargs.get("model", "default"), - "vector_db": route.vector_db, - }) - return result + return self.add_chunks(collection, chunks, return_embeddings=return_embeddings) - @traced( - "kb.ingest", - capture=["collection_name"], - extract_return={ - "n_files": lambda r: r.get("files"), - "n_chunks": lambda r: r.get("chunks"), - }, - ) - def ingest( + # -- single-collection hybrid search ------------------------------------ + + def search( self, - folder: str | Path, - collection_name: str | None = None, + query: str, + collection: str, + top_k: int, + where: dict | None = None, + where_document: dict | None = None, + ) -> list[dict]: + index, chunks = self._load(collection) + query_emb = self.embed([query])[0] + + # A ``where`` (metadata) or ``where_document`` (text content) filter + # disables the BM25 leg — BM25 has no filter equivalent, so a filtered + # search is dense-only. + filtered = where is not None or where_document is not None + do_hybrid = not filtered + # Oversample the per-ranker pools so RRF has depth on small top_k. + candidate_k = min( + max(top_k * 10, 50) if do_hybrid else top_k, + len(chunks), + ) + + with kb_index_search_span("chroma", candidate_k, filtered): + try: + if filtered: + dense_scores, dense_indices = index.search( + query_emb, + candidate_k, + where=where, + where_document=where_document, + ) + else: + dense_scores, dense_indices = index.search(query_emb, candidate_k) + except Exception as e: + hint = self._stale_index_message(collection, e) + if hint: + raise RuntimeError(hint) from e + raise + + if do_hybrid: + bm25 = self._get_bm25(collection) + _, sparse_indices = bm25.search(query, candidate_k) + dense_ranking = [int(i) for i in dense_indices if i >= 0] + sparse_ranking = [int(i) for i in sparse_indices] + merged = _rrf_merge([dense_ranking, sparse_ranking]) + results = [ + {"chunk": chunks[idx], "score": float(score)} for idx, score in merged + ] + else: + results = [ + {"chunk": chunks[i], "score": float(dense_scores[j])} + for j, i in enumerate(dense_indices) + if i >= 0 + ] + return results[:top_k] + + # -- persistence helpers ------------------------------------------------- + + def _rebuild_bm25(self, collection: str, chunks: list[dict]) -> None: + """Build (or rebuild) the BM25 index for *collection* from *chunks*.""" + bm25 = BM25Index() + bm25.build([c["text"] for c in chunks]) + bm25.save(self.index_dir / collection) + self._bm25_cache[collection] = bm25 + + def _get_bm25(self, collection: str) -> BM25Index: + cached = self._bm25_cache.get(collection) + if cached is not None: + return cached + bm25 = BM25Index.load(self.index_dir / collection) + self._bm25_cache[collection] = bm25 + return bm25 + + def _load(self, name: str) -> tuple[ChromaIndex, list[dict]]: + """Load a collection's index + chunks (cached in memory).""" + if name in self._cache: + return self._cache[name] + + coll_dir = self.index_dir / name + if not coll_dir.exists(): + raise ValueError(f"Collection '{name}' not found") + + index = ChromaIndex.load(coll_dir) + with open(coll_dir / "chunks.jsonl") as f: + chunks = [json.loads(line) for line in f] + + self._cache[name] = (index, chunks) + return index, chunks + + @staticmethod + def _stale_index_message(collection: str, exc: Exception) -> str | None: + """Return a user-friendly hint when *exc* looks like a dim mismatch. + + Chroma raises ``InvalidDimensionException`` with "dimension" in the + message when a query vector's dimension doesn't match the index's — + i.e. the embedder that BUILT the index was swapped for one with a + different output dim (the user changed ``embedding.backend`` / + ``embedding.model`` after init). Detect on message-shape. + """ + msg = str(exc).lower() + if "dimension" in msg or "dim mismatch" in msg or "shape" in msg: + return ( + f"Embedding dimension mismatch in collection {collection!r}. " + f"The index was built with a different embedder than the one " + f"currently configured. Drop the stale collection " + f"(`rm -rf /kb_index/{collection}`) and re-create the " + f"project with `dsagt init` so it rebuilds with the active " + f"embedding settings." + ) + return None + + def close(self) -> None: + if self._embedder is not None: + self._embedder.close() + self._embedder = None + self._cache.clear() + self._bm25_cache.clear() + + +# =========================================================================== +# KnowledgeBase — the federation infra: a list of stores + ingest pipeline +# =========================================================================== + + +class KnowledgeBase: + """Collection-based document retrieval over one-or-more vector stores. + + Holds a **list** of :class:`VectorStore`s (today just the internal local + Chroma store) and its job is to **fuse their collections**: collection→store + routing plus rank-fusion across collections. It also owns the document + ingestion pipeline (collect / parse / chunk → ``VectorStore.add_chunks``) and + cross-collection reranking. This is the shared substrate every KB consumer + (retrieval, memory, provenance, skills) calls into. + + Quick-start + ----------- + .. code-block:: python + + kb = KnowledgeBase(index_dir="./kb_store", default_embedder="local") + kb.ingest("./docs/research", "research") + results = kb.search("transformer architecture", "research") + """ + + FILE_TYPES = [ + "pdf", + "md", + "rst", + "txt", + "py", + "docx", + "json", + "yaml", + "yml", + # Packaging metadata: agents reading this index need to know which + # version of a library to install when registering tools that + # depend on it. pyproject.toml is the modern standard; setup.cfg + # is still common in older codebases. + "toml", + "cfg", + ] + + def __init__( + self, + index_dir: str | Path, + chunk_size: int = 1024, + chunk_overlap: int = 128, + rerank_model: str = "BAAI/bge-reranker-v2-m3", + default_rerank: bool = False, + recency_half_life_days: float | None = None, + # Internal store's embedder (one per store, fixed at construction). + # Explicit args — callers unpack their config here, no kwargs dict. + default_embedder: str | None = None, + model: str | None = None, + base_url: str | None = None, + api_key: str | None = None, + device: str | None = None, + ): + self.index_dir = Path(index_dir) + self.chunk_size = chunk_size + self.chunk_overlap = chunk_overlap + self.rerank_model = rerank_model + self.default_rerank = default_rerank + # Episodic recency weighting (session_memory only); None = off. + self._recency_half_life_days = recency_half_life_days + + self.index_dir.mkdir(parents=True, exist_ok=True) + + # The internal store: local Chroma, one embedder, many collections. + # External BYO stores would be appended to ``self._stores`` (deferred). + self._store = ChromaVectorStore( + self.index_dir, + backend=default_embedder or "api", + model=model, + base_url=base_url, + api_key=api_key, + device=device, + ) + self._stores: list[VectorStore] = [self._store] + + # Per-file-type parser cache. Constructing a CodeSplitter loads the + # tree-sitter language definition (~25ms each), so a 300-file ingest + # without this cache pays 7-8 seconds in parser construction alone. + # Parsers are stateless once built; safe to share across files. + self._parsers: dict[str, Any] = {} + + # Per-ingest counter for files that _chunk_file skipped due to + # read/parse errors. ingest()/append() reset this before the loop + # and surface it in the result dict so users notice when a non-zero + # number of files are silently being dropped. + self._chunk_skip_count: int = 0 + + self._reranker = None + + # -- store routing ------------------------------------------------------- + + def _store_for(self, collection: str) -> VectorStore: + """Return the store hosting *collection*, else raise ``ValueError``.""" + for store in self._stores: + if store.has_collection(collection): + return store + raise ValueError(f"Collection '{collection}' not found") + + @property + def collections(self) -> list[str]: + """Union of collection names across all stores.""" + seen: list[str] = [] + for store in self._stores: + for name in store.collections: + if name not in seen: + seen.append(name) + return seen + + def list_collections(self) -> list[dict]: + return [ + store.collection_info(name) + for store in self._stores + for name in store.collections + ] + + def preload_default_embedder(self) -> None: + """Kick off internal-store embedder construction in a daemon thread. + + Called at MCP server startup so the heavy load (sentence-transformers + import + model load, ~5-10s) happens in parallel with the rest of + bootstrap. Failure is swallowed: it resurfaces with a full traceback + on the first real embedding call. + """ + + def _load() -> None: + try: + self._store.embedder # noqa: B018 — triggers lazy construction + except Exception as e: + logger.warning("Background embedder preload failed: %s", e) + + threading.Thread( + target=_load, name="dsagt-embedder-preload", daemon=True + ).start() + + def embed_texts( + self, texts: list[str], collection: str | None = None + ) -> np.ndarray: + """Embed *texts* with the internal store's embedder (L2-normalized).""" + return self._store.embed(texts) + + # -- structured entries (delegate to internal store) -------------------- + + @traced( + "kb.add_entries", + capture=["collection"], + extract_return={"n_entries": lambda r: r.get("entries_added")}, + ) + def add_entries( + self, + texts: list[str], + collection: str, + metadatas: list[dict] | None = None, + return_embeddings: bool = False, + ) -> dict: + """Add pre-formed text entries with optional metadata to a collection. + + Unlike ``ingest``/``append``, this skips document parsing and chunking. + Used by episodic memory, tool_executions, and other structured entry + types that produce their own text representations. Metadata is stored + as native Chroma metadata for ``where`` filtering. + """ + obs.set_inputs( + { + "collection": collection, + "n_entries": len(texts), + "texts_preview": [t[:200] for t in texts[:3]], + } + ) + return self._store.add_entries( + texts, + collection, + metadatas=metadatas, + return_embeddings=return_embeddings, + ) + + # -- document ingestion pipeline ---------------------------------------- + + @traced( + "kb.ingest", + capture=["collection_name"], + extract_return={ + "n_files": lambda r: r.get("files"), + "n_chunks": lambda r: r.get("chunks"), + }, + ) + def ingest( + self, + folder: str | Path, + collection_name: str | None = None, file_types: list[str] | None = None, - route: CollectionRoute | None = None, exclude_patterns: list[str] | None = None, ) -> dict: - """ - Ingest *folder* as a new collection. - - Parameters - ---------- - folder : path - Source directory. - collection_name : str, optional - Defaults to folder name. - file_types : list[str], optional - Extensions to include (without leading dot). - route : CollectionRoute, optional - Override routing for this collection. Persisted to disk so - subsequent ``search`` / ``append`` calls use the same backend. - exclude_patterns : list[str], optional - Glob-style patterns matched (via :func:`fnmatch.fnmatch`) - against each file's path *relative to* ``folder``. Any file - whose relative path matches any pattern is skipped. Patterns - are checked against both the full relative path and the - basename, so ``"tests/*"`` excludes a top-level tests dir and - ``"test_*.py"`` excludes test files anywhere in the tree. - Useful for skipping tests, build artifacts, and private modules - without inflating the embed cost of large libraries. - """ + """Ingest *folder* as a new collection in the internal store.""" folder = Path(folder) collection = collection_name or folder.name file_types = file_types or self.FILE_TYPES - obs.set_inputs({ - "folder": str(folder), - "collection": collection, - "file_types": file_types, - }) - - # Set route in memory so _get_embedder resolves during the build. - # Persisted to disk via register_route after the index is built. - if route is not None: - self._routes[collection] = route - active_route = self._get_route(collection) + obs.set_inputs( + { + "folder": str(folder), + "collection": collection, + "file_types": file_types, + } + ) coll_dir = self.index_dir / collection - coll_dir.mkdir(exist_ok=True) + coll_dir.mkdir(parents=True, exist_ok=True) # Record source folder so the MCP server can detect re-ingests vs. conflicts. (coll_dir / "source.txt").write_text(str(folder.resolve())) @@ -1046,16 +1404,10 @@ def ingest( files = self._collect_files(folder, file_types, exclude_patterns) logger.info("Found %d files to process", len(files)) - # Reset the per-ingest skip counter; _chunk_file bumps it on - # read/parse failure. We surface the result in the return dict - # so callers (and the user looking at dsagt-setup-kb output) - # notice if a non-trivial number of files were silently dropped. self._chunk_skip_count = 0 chunks = [chunk for f in files for chunk in self._chunk_file(f, collection)] n_skipped = self._chunk_skip_count - logger.info( - "Created %d chunks (skipped %d files)", len(chunks), n_skipped, - ) + logger.info("Created %d chunks (skipped %d files)", len(chunks), n_skipped) if not chunks: return { @@ -1065,28 +1417,7 @@ def ingest( "skipped_files": n_skipped, } - embedder = self._get_embedder(active_route) - with kb_embed_span(active_route.embedding_backend, - active_route.embedder_kwargs.get("model"), len(chunks)): - embeddings = self._normalize(embedder.embed([c["text"] for c in chunks])) - - index = _make_index(active_route.vector_db, - **self._index_init_kwargs(active_route, coll_dir)) - index.add(embeddings) - index.save(coll_dir) - - # Persist route AFTER index is built — guarantees route.json always - # reflects what was actually used, never overwritten by a racing job. - self.register_route(collection, active_route) - - with open(coll_dir / "chunks.jsonl", "w") as f: - for chunk in chunks: - f.write(json.dumps(chunk) + "\n") - - if active_route.hybrid: - self._rebuild_bm25(collection, chunks) - - self._cache[collection] = (index, chunks) + self._store.add_chunks(collection, chunks) return { "collection": collection, "files": len(files), @@ -1111,15 +1442,15 @@ def append( """Append documents to an existing collection.""" file_types = file_types or self.FILE_TYPES - obs.set_inputs({ - "collection": collection, - "n_paths": len(paths), - "paths_preview": [str(p) for p in paths[:5]], - }) + obs.set_inputs( + { + "collection": collection, + "n_paths": len(paths), + "paths_preview": [str(p) for p in paths[:5]], + } + ) - index, existing_chunks = self._load(collection) - active_route = self._get_route(collection) - coll_dir = self.index_dir / collection + store = self._store_for(collection) # raises ValueError if missing files: list[Path] = [] for p in paths: @@ -1132,365 +1463,137 @@ def append( logger.warning("Path not found, skipping: %s", p) if not files: - return {"collection": collection, "files": 0, "chunks_added": 0, - "total_chunks": len(existing_chunks), "skipped_files": 0} + return { + "collection": collection, + "files": 0, + "chunks_added": 0, + "skipped_files": 0, + } - # Reset and capture per-file skip count, same pattern as ingest(). self._chunk_skip_count = 0 new_chunks = [chunk for f in files for chunk in self._chunk_file(f, collection)] n_skipped = self._chunk_skip_count if not new_chunks: - return {"collection": collection, "files": len(files), - "chunks_added": 0, "total_chunks": len(existing_chunks), - "skipped_files": n_skipped} - - embedder = self._get_embedder(active_route) - with kb_embed_span(active_route.embedding_backend, - active_route.embedder_kwargs.get("model"), len(new_chunks)): - embeddings = self._normalize(embedder.embed([c["text"] for c in new_chunks])) - index.add(embeddings) - index.save(coll_dir) - - with open(coll_dir / "chunks.jsonl", "a") as f: - for chunk in new_chunks: - f.write(json.dumps(chunk) + "\n") + return { + "collection": collection, + "files": len(files), + "chunks_added": 0, + "skipped_files": n_skipped, + } - all_chunks = existing_chunks + new_chunks - if active_route.hybrid: - self._rebuild_bm25(collection, all_chunks) - self._cache[collection] = (index, all_chunks) + result = store.add_chunks(collection, new_chunks) return { "collection": collection, "files": len(files), "chunks_added": len(new_chunks), - "total_chunks": len(all_chunks), + "total_chunks": result.get("total_entries"), "skipped_files": n_skipped, } - @traced( - "kb.add_entries", - capture=["collection"], - extract_return={"n_entries": lambda r: r.get("entries_added")}, - ) - def add_entries( - self, - texts: list[str], - collection: str, - metadatas: list[dict] | None = None, - route: CollectionRoute | None = None, - return_embeddings: bool = False, - ) -> dict: - """Add pre-formed text entries with optional metadata to a collection. - - Unlike ``ingest``/``append``, this skips document parsing and chunking. - Each text is embedded and stored as-is. Used by episodic memory, - tool_executions, and other structured entry types that produce their - own text representations rather than ingesting raw documents. - - Parameters - ---------- - texts : list[str] - Text content for each entry (will be embedded). - collection : str - Target collection name (created if it doesn't exist). - metadatas : list[dict], optional - Per-entry metadata dicts. On ChromaDB-backed collections these - are stored as native metadata for ``where`` filtering. On all - backends they are merged into the chunk metadata in chunks.jsonl. - route : CollectionRoute, optional - Override routing for this collection (persisted to disk). - return_embeddings : bool, optional - If True, include the freshly-computed embeddings in the result - dict under the ``"embeddings"`` key. Callers that need the - embeddings for downstream work (e.g. centroid-based outlier - detection in memory extraction) should set this so they can - avoid a second round-trip to the embedding API. - - Returns - ------- - dict - ``{collection, entries_added, total_entries}`` plus - ``"embeddings"`` (numpy ndarray) if ``return_embeddings=True``. - """ - obs.set_inputs({ - "collection": collection, - "n_entries": len(texts), - "texts_preview": [t[:200] for t in texts[:3]], - }) - - if not texts: - result = {"collection": collection, "entries_added": 0, "total_entries": 0} - if return_embeddings: - result["embeddings"] = np.array([], dtype=np.float32) - return result - - coll_dir = self.index_dir / collection - coll_dir.mkdir(exist_ok=True) - - if route is not None: - self._routes[collection] = route - active_route = self._get_route(collection) - - embedder = self._get_embedder(active_route) - with kb_embed_span(active_route.embedding_backend, - active_route.embedder_kwargs.get("model"), len(texts)): - embeddings = self._normalize(embedder.embed(texts)) - - # Load existing or create new index - if (coll_dir / "chunks.jsonl").exists(): - index, existing_chunks = self._load(collection) - else: - index = _make_index( - active_route.vector_db, - **self._index_init_kwargs(active_route, coll_dir), - ) - existing_chunks = [] - - # Build chunk dicts (consistent with chunks.jsonl format) - new_chunks = [] - for i, text in enumerate(texts): - entry_meta = metadatas[i] if metadatas else {} - chunk = { - "text": text, - "metadata": { - "collection": collection, - "source_file": "entry", - "chunk_index": len(existing_chunks) + i, - **entry_meta, - }, - } - new_chunks.append(chunk) - - # Store embeddings (with metadata on ChromaDB) - try: - if active_route.vector_db == "chroma" and metadatas is not None: - index.add(embeddings, metadatas=metadatas) - else: - index.add(embeddings) - except Exception as e: - hint = self._stale_index_message(collection, e) - if hint: - raise RuntimeError(hint) from e - raise - - index.save(coll_dir) - self.register_route(collection, active_route) - - # Append to chunks.jsonl - with open(coll_dir / "chunks.jsonl", "a") as f: - for chunk in new_chunks: - f.write(json.dumps(chunk) + "\n") - - all_chunks = existing_chunks + new_chunks - if active_route.hybrid: - self._rebuild_bm25(collection, all_chunks) - self._cache[collection] = (index, all_chunks) - - result = { - "collection": collection, - "entries_added": len(texts), - "total_entries": len(all_chunks), - } - if return_embeddings: - result["embeddings"] = embeddings - return result - - def embed_texts(self, texts: list[str], collection: str) -> np.ndarray: - """Embed texts using the embedder configured for a collection. - - Returns L2-normalized float32 array of shape (n_texts, dim). - """ - route = self._get_route(collection) - embedder = self._get_embedder(route) - return self._normalize(embedder.embed(texts)) + # -- federated search ---------------------------------------------------- @traced("kb.search", capture=["collection", "top_k", "rerank"]) def search( self, query: str, - collection: str, + collection: str | None = None, + collections: list[str] | None = None, top_k: int = 5, rerank: bool | None = None, where: dict | None = None, + where_document: dict | None = None, ) -> list[dict]: - """Search *collection* with *query*. - - Parameters - ---------- - rerank : bool, optional - Use cross-encoder reranking. Defaults to ``self.default_rerank`` - (set from ``knowledge.rerank`` in dsagt_config.yaml). - where : dict, optional - ChromaDB ``where`` filter clause. Only effective on ChromaDB-backed - collections; silently ignored for FAISS collections. When set, - disables the BM25 sparse leg of hybrid search (BM25 has no - metadata-filter equivalent), so the result is dense-only. + """Search one or many collections, fusing across them by rank. + + A single collection routes straight to its store's hybrid search. + Multiple collections fan out — each store searched, then the per- + collection rankings fused by Reciprocal Rank Fusion (rank-only, so + different embedding spaces compose correctly). Optional cross-encoder + rerank runs over the fused candidates. + + Missing collections are skipped with a warning; the search fails only + when *every* requested collection is absent. """ if rerank is None: rerank = self.default_rerank - obs.set_inputs({"query": query, "collection": collection, "top_k": top_k}) + targets = collections or ([collection] if collection else []) + if not targets: + raise ValueError("Provide 'collection' or 'collections'") - index, chunks = self._load(collection) - active_route = self._get_route(collection) - embedder = self._get_embedder(active_route) - - with kb_embed_span(active_route.embedding_backend, - active_route.embedder_kwargs.get("model"), 1): - query_emb = self._normalize(embedder.embed([query]))[0] - - # Wider candidate pool when reranking (cross-encoder reorders top-N) - # OR when hybrid (RRF fuses two ranked lists; needs candidates from - # each ranker before merging). Floor of 50 keeps RRF meaningful on - # small top_k. - filtered = where is not None and active_route.vector_db == "chroma" - do_hybrid = active_route.hybrid and not filtered - oversample = (rerank or do_hybrid) - candidate_k = min( - max(top_k * 10, 50) if oversample else top_k, - len(chunks), + obs.set_inputs({"query": query, "collections": targets, "top_k": top_k}) + + # Oversample per-collection pools when reranking, fusing >1 collection, + # or recency-weighting (so a recent fact can be lifted from deep in the + # pool, not merely reordered within an already-cut top_k). + recency_target = bool( + self._recency_half_life_days and targets == [_RECENCY_COLLECTION] ) + oversample = rerank or len(targets) > 1 or recency_target + candidate_k = max(top_k * 10, 50) if oversample else top_k - with kb_index_search_span(active_route.vector_db, candidate_k, filtered): + per_coll: list[list[dict]] = [] + errors: list[str] = [] + for coll in targets: try: - if filtered: - dense_scores, dense_indices = index.search( - query_emb, candidate_k, where=where, + store = self._store_for(coll) + per_coll.append( + store.search( + query, + coll, + candidate_k, + where=where, + where_document=where_document, ) - else: - dense_scores, dense_indices = index.search( - query_emb, candidate_k, - ) - except Exception as e: - hint = self._stale_index_message(collection, e) - if hint: - raise RuntimeError(hint) from e - raise + ) + except (ValueError, FileNotFoundError, KeyError) as e: + logger.warning("Search failed for '%s': %s", coll, e) + errors.append(str(e)) + + if not per_coll: + if len(targets) == 1: + raise ValueError( + errors[0] if errors else f"Collection '{targets[0]}' not found" + ) + raise ValueError(f"All collections failed: {'; '.join(errors)}") - if do_hybrid: - bm25 = self._get_bm25(collection) - _, sparse_indices = bm25.search(query, candidate_k) - dense_ranking = [int(i) for i in dense_indices if i >= 0] - sparse_ranking = [int(i) for i in sparse_indices] - merged = _rrf_merge([dense_ranking, sparse_ranking]) - results = [ - {"chunk": chunks[idx], "score": float(score)} - for idx, score in merged - ] - else: - results = [ - {"chunk": chunks[i], "score": float(dense_scores[j])} - for j, i in enumerate(dense_indices) if i >= 0 - ] + fused = per_coll[0] if len(per_coll) == 1 else _rrf_across(per_coll) - if rerank and results: - with kb_rerank_span(self.rerank_model, len(results)): - final = self._rerank(query, results, top_k) + # Episodic recency: a recent corrected fact outranks a stale one without + # any contradiction detection — recency is the ranker for session_memory + # (mutually exclusive with the cross-encoder; the two are alternative + # rerankers and recency is the one that matters for a time-ordered log). + if recency_target and fused: + final = _apply_recency(fused, self._recency_half_life_days, time.time())[ + :top_k + ] + elif rerank and fused: + with kb_rerank_span(self.rerank_model, len(fused)): + final = self._rerank(query, fused, top_k) else: - final = results[:top_k] + final = fused[:top_k] obs.set("hits", len(final)) - obs.set_outputs({ - "hits": len(final), - "top_texts": [r["chunk"].get("text", "")[:200] for r in final[:3]], - }) + obs.set_outputs( + { + "hits": len(final), + "top_texts": [r["chunk"].get("text", "")[:200] for r in final[:3]], + } + ) return final - def _rebuild_bm25(self, collection: str, chunks: list[dict]) -> None: - """Build (or rebuild) the BM25 index for *collection* from *chunks*. - - Called from every write path of a hybrid-enabled collection. BM25 - IDF stats are corpus-global, so there is no incremental update — - we always rebuild from the full chunk list. Cached on the KB - instance and persisted to ``/bm25.pkl``. - """ - bm25 = BM25Index() - bm25.build([c["text"] for c in chunks]) - bm25.save(self.index_dir / collection) - self._bm25_cache[collection] = bm25 - - def _get_bm25(self, collection: str) -> BM25Index: - """Return the cached BM25 index, loading from disk on first hit. - - Raises ``FileNotFoundError`` if the collection's route has - ``hybrid=True`` but no ``bm25.pkl`` exists on disk — i.e. the - collection was built by a pre-hybrid dsagt. Migration is - ``delete the collection and re-ingest``; we deliberately do NOT - rebuild lazily on first search because that would silently turn - the next search into an N-document tokenization run. - """ - cached = self._bm25_cache.get(collection) - if cached is not None: - return cached - coll_dir = self.index_dir / collection - bm25_path = coll_dir / BM25Index._FILENAME - if not bm25_path.exists(): - raise FileNotFoundError( - f"Collection '{collection}' has hybrid=True in its route but " - f"no bm25.pkl on disk. The collection was built before hybrid " - f"retrieval was added. To fix: delete {coll_dir} and re-ingest, " - f"or set hybrid=False in route.json." - ) - bm25 = BM25Index.load(coll_dir) - self._bm25_cache[collection] = bm25 - return bm25 - - @staticmethod - def _normalize(arr: np.ndarray) -> np.ndarray: - norms = np.linalg.norm(arr, axis=-1, keepdims=True) - return arr / np.where(norms > 0, norms, 1) - - def _index_init_kwargs(self, route: CollectionRoute, coll_dir: Path) -> dict: - """Extra kwargs needed by some index constructors at init time.""" - if route.vector_db == "chroma": - return {"collection_name": coll_dir.name, - "persist_dir": coll_dir, **route.index_kwargs} - return dict(route.index_kwargs) - - def _load_route(self, collection: str) -> CollectionRoute | None: - route_path = self.index_dir / collection / self._ROUTE_FILE - if route_path.exists(): - return CollectionRoute.from_dict(json.loads(route_path.read_text())) - return None - - def _load(self, name: str) -> tuple[BaseVectorIndex, list[dict]]: - """Load collection index + chunks (cached in memory).""" - if name in self._cache: - return self._cache[name] - - coll_dir = self.index_dir / name - if not coll_dir.exists(): - raise ValueError(f"Collection '{name}' not found") - - # Restore persisted route if not already registered - if name not in self._routes: - persisted = self._load_route(name) - if persisted: - self._routes[name] = persisted - - active_route = self._get_route(name) - - # Pick correct loader - if active_route.vector_db == "chroma": - index = ChromaIndex.load(coll_dir) - else: - index = FAISSIndex.load(coll_dir) - - with open(coll_dir / "chunks.jsonl") as f: - chunks = [json.loads(line) for line in f] - - self._cache[name] = (index, chunks) - return index, chunks - def _rerank(self, query: str, results: list[dict], top_k: int) -> list[dict]: if self._reranker is None: from sentence_transformers import CrossEncoder + self._reranker = CrossEncoder(self.rerank_model, max_length=512) pairs = [[query, r["chunk"]["text"]] for r in results] scores = self._reranker.predict(pairs, show_progress_bar=False) ranked = sorted(zip(results, scores), key=lambda x: x[1], reverse=True) return [{**r, "rerank_score": float(s)} for r, s in ranked[:top_k]] + # -- file discovery + chunking ------------------------------------------ + def _collect_files( self, folder: Path, @@ -1501,19 +1604,13 @@ def _collect_files( glob exclusions. Pure function relative to filesystem state — no caching, no - side-effecting attributes. Extracted from ``ingest()`` so file - discovery and exclusion logic can be tested in isolation without - spinning up an embedder, an index, or a chunker. - - Patterns are checked against the relative path, the basename, and - each individual path segment, so ``"tests"`` excludes any file - whose path contains a ``tests/`` directory regardless of depth. + side-effecting attributes. Patterns are checked against the relative + path, the basename, and each individual path segment, so ``"tests"`` + excludes any file whose path contains a ``tests/`` directory. """ from fnmatch import fnmatch - all_files = [ - f for ext in file_types for f in folder.glob(f"**/*.{ext}") - ] + all_files = [f for ext in file_types for f in folder.glob(f"**/*.{ext}")] if not exclude_patterns: return all_files @@ -1522,7 +1619,8 @@ def _excluded(f: Path) -> bool: rel = str(f.relative_to(folder)) name = f.name return any( - fnmatch(rel, pat) or fnmatch(name, pat) + fnmatch(rel, pat) + or fnmatch(name, pat) or any(fnmatch(part, pat) for part in Path(rel).parts) for pat in exclude_patterns ) @@ -1532,7 +1630,9 @@ def _excluded(f: Path) -> bool: if n_excluded: logger.info( "Excluded %d/%d files via %d pattern(s)", - n_excluded, len(all_files), len(exclude_patterns), + n_excluded, + len(all_files), + len(exclude_patterns), ) return kept @@ -1544,17 +1644,7 @@ def _chunk_file(self, path: Path, collection: str) -> Iterator[dict]: # Per-file read/parse failures are kept as soft failures (count # surfaced in ingest()'s return dict) rather than aborting the - # whole ingest, because real-world directories like cloned - # upstream repos contain occasional unreadable files (binary - # blobs misnamed .txt, malformed UTF-8) that shouldn't kill an - # ingest of thousands of files. Callers see the skip count - # and can investigate if it's suspiciously high. - # - # llama_index's SimpleDirectoryReader prints "Failed to load file - # ... Skipping..." directly to stdout (literal print(), not - # logger) for every file it can't parse — there's no level switch - # for this. Capture it: the file is already counted as a miss - # below, so the per-file print is pure noise. + # whole ingest try: with contextlib.redirect_stdout(_io.StringIO()): docs = SimpleDirectoryReader(input_files=[str(path)]).load_data() @@ -1568,7 +1658,9 @@ def _chunk_file(self, path: Path, collection: str) -> Iterator[dict]: parser = self._get_parser(file_type) try: nodes = parser.get_nodes_from_documents(docs) - except (ValueError, RuntimeError) as e: + except Exception as e: + # A single malformed file (or a tree-sitter ABI mismatch in the + # code parser) must not abort the whole ingest — skip it. logger.warning("Could not parse %s: %s", path, e) self._chunk_skip_count += 1 return @@ -1577,7 +1669,9 @@ def _chunk_file(self, path: Path, collection: str) -> Iterator[dict]: if not text: continue yield { - "id": hashlib.sha256(f"{path}:{i}:{text[:100]}".encode()).hexdigest()[:16], + "id": hashlib.sha256(f"{path}:{i}:{text[:100]}".encode()).hexdigest()[ + :16 + ], "text": text, "metadata": { "source_file": str(path), @@ -1588,13 +1682,7 @@ def _chunk_file(self, path: Path, collection: str) -> Iterator[dict]: } def _get_parser(self, file_type: str): - """Return a cached parser for *file_type*, building it on first use. - - Parsers are stateless after construction. ``CodeSplitter`` in - particular loads a tree-sitter language definition on every - construction (~25ms), so a large ingest used to pay this cost per - file; now it pays once per file type. - """ + """Return a cached parser for *file_type*, building it on first use.""" cached = self._parsers.get(file_type) if cached is not None: return cached @@ -1625,11 +1713,9 @@ def _get_parser(self, file_type: str): return parser def close(self) -> None: - for client in self._embedder_cache.values(): - client.close() - self._embedder_cache.clear() + for store in self._stores: + store.close() self._parsers.clear() - self._bm25_cache.clear() def __enter__(self): return self diff --git a/src/dsagt/mcp/__init__.py b/src/dsagt/mcp/__init__.py new file mode 100644 index 0000000..2f8b1e6 --- /dev/null +++ b/src/dsagt/mcp/__init__.py @@ -0,0 +1,12 @@ +"""DSAGT MCP server package — the single merged ``dsagt-server``. + +The MCP tool surface is split by concern across sibling modules: + +* :mod:`dsagt.mcp.registry_tools` — tool registry + execution + provenance +* :mod:`dsagt.mcp.knowledge_tools` — knowledge-base retrieval +* :mod:`dsagt.mcp.memory_tools` — explicit memory +* :mod:`dsagt.mcp.skill_tools` — skill search / install / sources + +:mod:`dsagt.mcp.server` composes all four under one ``Server("dsagt")`` and +owns the ``dsagt-server`` entry point + shared-KB startup. +""" diff --git a/src/dsagt/mcp/knowledge_tools.py b/src/dsagt/mcp/knowledge_tools.py new file mode 100644 index 0000000..8b43b17 --- /dev/null +++ b/src/dsagt/mcp/knowledge_tools.py @@ -0,0 +1,541 @@ +"""MCP tools for knowledge-base retrieval. + +Semantic search over document collections + background ingest/append jobs. +Long-running operations (ingest, append) run in the background and return +immediately with a ``job_id``; poll ``kb_job_status`` for completion. + +Multi-collection search fans out and rank-fuses *below* this tool boundary, in +:meth:`dsagt.knowledge.KnowledgeBase.search` — the agent just names collection(s). +BYO external vector stores are deferred: ``kb_add_vector_db`` is intentionally +**not** registered (an external store is a ``VectorStore`` subclass appended to +the KB's store list, not a tool the agent calls). + +Server configuration (chunk_size, rerank) is read from the project's +.dsagt/config.yaml. Embedding credentials flow through env vars (EMBEDDING_API_KEY, +EMBEDDING_BASE_URL, EMBEDDING_MODEL) set by ``dsagt start``. + +These definitions + handlers run inside the merged ``dsagt-server`` (see +:mod:`dsagt.mcp.server`); ``create_knowledge_server`` is retained only as a +test-facing constructor. Explicit-memory tools (``kb_remember`` / etc.) live in +:mod:`dsagt.mcp.memory_tools`; skill-source tools in +:mod:`dsagt.mcp.skill_tools`. +""" + +import os + +# Prevent fatal OpenMP crash when multiple libraries (PyTorch / +# sentence-transformers) each bundle their own libomp. Must precede the +# ``dsagt.knowledge`` import below. +os.environ.setdefault("KMP_DUPLICATE_LIB_OK", "TRUE") + +import asyncio # noqa: E402 +import logging # noqa: E402 +import time # noqa: E402 +import uuid # noqa: E402 +from dataclasses import dataclass, field # noqa: E402 +from functools import partial # noqa: E402 +from pathlib import Path # noqa: E402 + +import mcp.types as types # noqa: E402 + +from dsagt.knowledge import KnowledgeBase # noqa: E402 +from dsagt.mcp.server import build_dispatch_server # noqa: E402 +from dsagt.session import _collection_exists # noqa: E402 +from dsagt.session import setup_runtime_kb # noqa: E402, F401 (re-exported for tests) + +logger = logging.getLogger(__name__) + + +# --------------------------------------------------------------------------- +# Background job tracker +# --------------------------------------------------------------------------- + + +@dataclass +class _JobTracker: + """Tracks background ingest/append jobs and their completion state.""" + + jobs: dict[str, dict] = field(default_factory=dict) + active_collections: set[str] = field(default_factory=set) + + def start(self, coro, collection: str | None = None) -> str: + job_id = uuid.uuid4().hex[:8] + self.jobs[job_id] = { + "status": "running", + "result": None, + "error": None, + "collection": collection, + "started_at": time.monotonic(), + "message": "Starting -- embedding documents via API...", + } + if collection: + self.active_collections.add(collection) + + tracker = self # capture for the closure + + async def _run(): + try: + tracker.jobs[job_id]["message"] = "Embedding and indexing documents..." + result = await coro + tracker.jobs[job_id]["status"] = "complete" + tracker.jobs[job_id]["result"] = result + tracker.jobs[job_id]["message"] = "Done." + except Exception as e: + import traceback + + tb = traceback.format_exc() + tracker.jobs[job_id]["status"] = "error" + tracker.jobs[job_id]["error"] = f"{type(e).__name__}: {e}" + tracker.jobs[job_id]["message"] = f"Failed: {type(e).__name__}: {e}" + tracker.jobs[job_id]["traceback"] = tb + logger.error("Job %s failed: %s\n%s", job_id, e, tb) + finally: + if collection: + tracker.active_collections.discard(collection) + + asyncio.get_event_loop().create_task(_run()) + return job_id + + +# --------------------------------------------------------------------------- +# Per-tool handlers (module-level, explicit dependencies) +# +# Each handler takes ``arguments: dict`` plus its dependencies as keyword +# args bound via functools.partial. Handlers return a result dict; the outer +# dispatch wrapper JSON-serializes it. +# --------------------------------------------------------------------------- + + +async def _handle_kb_list_collections(arguments: dict, *, kb: KnowledgeBase) -> dict: + collections = await asyncio.to_thread(kb.list_collections) + return {"status": "ok", "collections": collections, "count": len(collections)} + + +async def _handle_kb_search( + arguments: dict, + *, + kb: KnowledgeBase, +) -> dict: + query = arguments["query"] + top_k = arguments.get("top_k", 5) + rerank = arguments.get("rerank") # None → kb.default_rerank + + collection_arg = arguments.get("collection") + collections_arg = arguments.get("collections") + + if not collection_arg and not collections_arg: + return {"status": "error", "error": "Provide 'collection' or 'collections'"} + + # Build ChromaDB where clause from the filter arguments. ChromaDB + # requires single-filter dicts or $and-wrapped lists; an empty dict + # would be invalid, so we only pass where when there are real filters. + where = { + key: arguments[key] + for key in ("category", "session_id", "source_type", "code_name", "tool_name") + if arguments.get(key) is not None + } + return_code = arguments.get("return_code") + if return_code is not None: + where["return_code"] = int(return_code) + if len(where) > 1: + where = {"$and": [{k: v} for k, v in where.items()]} + + # Document-content filter (over the chunk text itself, complementary to the + # metadata ``where``). ``regex`` is the powerful leg — ChromaDB ``$regex``, + # with ``(?i)`` for case-insensitive; ``contains`` is a case-sensitive + # substring. Both narrow the candidate pool before vector ranking. + doc_filters = [] + if arguments.get("regex"): + doc_filters.append({"$regex": arguments["regex"]}) + if arguments.get("contains"): + doc_filters.append({"$contains": arguments["contains"]}) + where_document = ( + doc_filters[0] + if len(doc_filters) == 1 + else {"$and": doc_filters} if doc_filters else None + ) + + # Fan-out + rank-fusion across collections lives in KnowledgeBase.search; + # the tool just names collection(s). A single internal collection routes + # straight to its store; multiple/external targets federate by RRF. + try: + all_results = await asyncio.to_thread( + kb.search, + query=query, + collection=collection_arg, + collections=collections_arg, + top_k=top_k, + rerank=rerank, + where=where or None, + where_document=where_document, + ) + except ValueError as e: + return {"status": "error", "error": str(e)} + + return { + "status": "ok", + "query": query, + "collection": collection_arg or ",".join(collections_arg), + "result_count": len(all_results), + "results": [ + { + "text": r["chunk"]["text"], + "score": r["score"], + "rerank_score": r.get("rerank_score"), + "source_file": r["chunk"]["metadata"].get("source_file", ""), + "chunk_index": r["chunk"]["metadata"].get("chunk_index", 0), + "metadata": { + k: v + for k, v in r["chunk"]["metadata"].items() + if k + not in ("source_file", "chunk_index", "collection", "file_type") + }, + } + for r in all_results + ], + } + + +async def _handle_kb_ingest( + arguments: dict, + *, + kb: KnowledgeBase, + job_tracker: _JobTracker, +) -> dict: + folder_path = Path(arguments["folder_path"]) + collection_name = arguments.get("collection_name") + file_types = arguments.get("file_types") + + if not folder_path.exists(): + return {"status": "error", "error": f"Folder not found: {folder_path}"} + if not folder_path.is_dir(): + return {"status": "error", "error": f"Not a directory: {folder_path}"} + + target_name = collection_name or folder_path.name + warning = None + + if target_name in job_tracker.active_collections: + return { + "status": "error", + "error": ( + f"Collection '{target_name}' is already being ingested. " + f"Poll kb_job_status for progress." + ), + } + + if _collection_exists(kb.index_dir / target_name): + source_path = kb.index_dir / target_name / "source.txt" + existing_source = ( + source_path.read_text().strip() if source_path.exists() else None + ) + same_source = ( + existing_source is None + or Path(existing_source).resolve() == folder_path.resolve() + ) + if not same_source: + original_name = target_name + n = 1 + while ( + _collection_exists(kb.index_dir / target_name) + or target_name in job_tracker.active_collections + ): + target_name = f"{original_name}{n}" + n += 1 + warning = ( + f"Collection '{original_name}' already exists from a " + f"different folder; using '{target_name}'." + ) + + ingest_kwargs: dict = {"collection_name": target_name} + if file_types: + ingest_kwargs["file_types"] = file_types + + async def _ingest_with_logging(): + import traceback as _tb + + logger.info( + "Ingest starting: collection=%s folder=%s kwargs=%s", + target_name, + folder_path, + ingest_kwargs, + ) + try: + result = await asyncio.to_thread(kb.ingest, folder_path, **ingest_kwargs) + logger.info("Ingest complete: %s", result) + return result + except Exception as _e: + logger.error("Ingest FAILED: %s\n%s", _e, _tb.format_exc()) + raise + + job_id = job_tracker.start(_ingest_with_logging(), collection=target_name) + result = { + "status": "started", + "job_id": job_id, + "collection": target_name, + "message": ( + f"Ingestion started. " + f"Poll kb_job_status(job_id='{job_id}') every 10 seconds. " + f"DO NOT call ingest again -- the job is running in the " + f"background. Large folders may take several minutes." + ), + } + if warning: + result["warning"] = warning + return result + + +async def _handle_kb_append( + arguments: dict, + *, + kb: KnowledgeBase, + job_tracker: _JobTracker, +) -> dict: + collection = arguments["collection"] + paths = arguments["paths"] + if isinstance(paths, str): + paths = [paths] + file_types = arguments.get("file_types") + + if not _collection_exists(kb.index_dir / collection): + return {"status": "error", "error": f"Collection '{collection}' not found"} + + append_kwargs: dict = {} + if file_types: + append_kwargs["file_types"] = file_types + + job_id = job_tracker.start( + asyncio.to_thread(kb.append, collection, paths, **append_kwargs), + collection=collection, + ) + return { + "status": "started", + "job_id": job_id, + "collection": collection, + "message": f"Append started. Poll kb_job_status(job_id='{job_id}') for progress.", + } + + +async def _handle_kb_job_status(arguments: dict, *, job_tracker: _JobTracker) -> dict: + job_id = arguments["job_id"] + if job_id not in job_tracker.jobs: + return {"status": "error", "error": f"Unknown job: {job_id}"} + + job = job_tracker.jobs[job_id] + elapsed = int(time.monotonic() - job["started_at"]) + result = { + "status": job["status"], + "elapsed_seconds": elapsed, + "message": job.get("message", ""), + } + if job["status"] == "running": + result["instruction"] = ( + "Job is still running. DO NOT call ingest again. " + "Keep polling job_status every 10 seconds until " + "status is 'complete' or 'error'." + ) + if job["result"] is not None: + result["result"] = job["result"] + if job["error"] is not None: + result["error"] = job["error"] + if job.get("traceback") and job["status"] == "error": + result["traceback"] = job["traceback"] + return result + + +# --------------------------------------------------------------------------- +# Tool defs + handler map (used by the merged server and the test wrapper) +# --------------------------------------------------------------------------- + + +def _knowledge_tools_and_handlers(kb: KnowledgeBase): + """Build the knowledge-base ``(tool defs, handler map)``. + + Combined with the other concern modules' tools under one MCP ``Server`` by + :func:`dsagt.mcp.server.create_dsagt_server`. The rerank default is on + ``kb.default_rerank`` (set from ``knowledge.rerank`` in .dsagt/config.yaml). + """ + job_tracker = _JobTracker() + + handlers = { + "kb_list_collections": partial(_handle_kb_list_collections, kb=kb), + "kb_search": partial(_handle_kb_search, kb=kb), + "kb_ingest": partial(_handle_kb_ingest, kb=kb, job_tracker=job_tracker), + "kb_append": partial(_handle_kb_append, kb=kb, job_tracker=job_tracker), + "kb_job_status": partial(_handle_kb_job_status, job_tracker=job_tracker), + } + + tools = [ + types.Tool( + name="kb_list_collections", + description=( + "List all available knowledge base collections with their " + "embedding model and vector DB. Use this to discover what " + "documentation is already indexed." + ), + inputSchema={"type": "object", "properties": {}}, + ), + types.Tool( + name="kb_search", + description=( + "Search knowledge base collections using semantic similarity. " + "Returns relevant chunks with source metadata. " + "Supports multi-collection search. Narrow results with metadata " + "filters (session_id, code_name, tool_name, source_type, ...) and/or a " + "document-text filter ('regex' / 'contains') over the chunk text " + "itself — useful for pulling specific session-memory context." + ), + inputSchema={ + "type": "object", + "properties": { + "query": { + "type": "string", + "description": "Natural language search query", + }, + "collection": { + "type": "string", + "description": "Name of a single collection to search", + }, + "collections": { + "type": "array", + "items": {"type": "string"}, + "description": "Search multiple collections and merge results (overrides 'collection')", + }, + "top_k": { + "type": "integer", + "description": "Number of results to return (default: 5)", + "default": 5, + }, + "rerank": { + "type": "boolean", + "description": "Use cross-encoder reranking (slower but more accurate). Default from config.", + "default": kb.default_rerank, + }, + "category": { + "type": "string", + "description": "Filter by category tag (ChromaDB collections only)", + }, + "session_id": { + "type": "string", + "description": "Filter by session ID (ChromaDB collections only)", + }, + "code_name": { + "type": "string", + "description": "Filter by registered-code name — code execution records (ChromaDB collections only)", + }, + "tool_name": { + "type": "string", + "description": "Filter by agent tool-call name — session memory (ChromaDB collections only)", + }, + "source_type": { + "type": "string", + "description": "Filter by source type (ChromaDB collections only)", + }, + "return_code": { + "type": "integer", + "description": "Filter by tool exit code (ChromaDB collections only)", + }, + "regex": { + "type": "string", + "description": ( + "Require chunk text to match this regular expression " + "(ChromaDB $regex). Use (?i) for case-insensitive, " + "e.g. '(?i)parser'. Narrows before semantic ranking." + ), + }, + "contains": { + "type": "string", + "description": ( + "Require chunk text to contain this case-sensitive " + "substring (ChromaDB $contains)." + ), + }, + }, + "required": ["query"], + }, + ), + types.Tool( + name="kb_ingest", + description=( + "Index a folder as a new knowledge base collection. " + "Returns immediately with a job_id. " + "IMPORTANT: poll kb_job_status every 10 seconds and wait for " + "status='complete'. DO NOT call ingest again for the same " + "folder while a job is running." + ), + inputSchema={ + "type": "object", + "properties": { + "folder_path": { + "type": "string", + "description": "Path to folder containing documents to index", + }, + "collection_name": { + "type": "string", + "description": "Name for the collection (default: folder name)", + }, + "file_types": { + "type": "array", + "items": {"type": "string"}, + "description": "File extensions to include, e.g. ['pdf', 'md', 'py']. Defaults to common types.", + }, + }, + "required": ["folder_path"], + }, + ), + types.Tool( + name="kb_append", + description=( + "Add documents to an existing collection. Uses the same embedding " + "model and vector DB the collection was created with. " + "Returns immediately with a job_id -- poll kb_job_status for progress." + ), + inputSchema={ + "type": "object", + "properties": { + "collection": { + "type": "string", + "description": "Name of the existing collection to append to", + }, + "paths": { + "type": "array", + "items": {"type": "string"}, + "description": "List of file or folder paths to add", + }, + "file_types": { + "type": "array", + "items": {"type": "string"}, + "description": "File extensions to include when expanding folders.", + }, + }, + "required": ["collection", "paths"], + }, + ), + types.Tool( + name="kb_job_status", + description="Check the status of a background ingest or append job.", + inputSchema={ + "type": "object", + "properties": { + "job_id": { + "type": "string", + "description": "Job ID returned by kb_ingest or kb_append", + }, + }, + "required": ["job_id"], + }, + ), + ] + return tools, handlers + + +def create_knowledge_server(kb: KnowledgeBase): + """Create a standalone MCP server exposing only the knowledge-base tools. + + Test-facing API: tests call it with a mock KB and drive the server via + ``call_tool_sync()``. The merged ``dsagt-server`` uses + :func:`_knowledge_tools_and_handlers` directly instead of this wrapper. + """ + tools, handlers = _knowledge_tools_and_handlers(kb) + return build_dispatch_server( + "knowledge", tools, handlers, {t: "knowledge" for t in handlers} + ) diff --git a/src/dsagt/mcp/memory_tools.py b/src/dsagt/mcp/memory_tools.py new file mode 100644 index 0000000..1b515a2 --- /dev/null +++ b/src/dsagt/mcp/memory_tools.py @@ -0,0 +1,172 @@ +"""MCP tools for explicit memory. + +User-confirmed facts that persist across sessions (``kb_remember`` / +``kb_get_memories``). These front :mod:`dsagt.memory` (``ExplicitMemory``); the +``kb_`` tool-name prefix is historical (the tools were born in the knowledge +server) and is kept for agent-facing backward compatibility. + +These definitions + handlers run inside the merged ``dsagt-server`` (see +:mod:`dsagt.mcp.server`); ``create_memory_server`` is retained only as a +test-facing constructor. +""" + +import asyncio +import logging +from functools import partial +from pathlib import Path + +import mcp.types as types + +from dsagt.knowledge import KnowledgeBase +from dsagt.mcp.server import build_dispatch_server +from dsagt.memory import ExplicitMemory + +logger = logging.getLogger(__name__) + + +async def _handle_kb_remember( + arguments: dict, + *, + kb: KnowledgeBase, + memory: ExplicitMemory, +) -> dict: + text = arguments["text"] + category = arguments.get("category", "") + session_id = arguments.get("session_id", "") + supersedes = arguments.get("supersedes") + + store_result = await asyncio.to_thread( + memory.remember, + text=text, + category=category, + session_id=session_id, + supersedes=supersedes, + ) + + if not store_result.get("stored"): + return { + "status": "error", + "error": store_result.get("error", "Failed to store memory"), + } + + # Mirror into the VectorStore for semantic recall — optional infra. + # The durable YAML write above already succeeded, so a mirror failure + # degrades to pure-YAML explicit memory rather than failing the tool. + try: + await asyncio.to_thread( + kb.add_entries, + texts=[text], + collection="session_memory", + metadatas=[ + { + "source_type": "explicit_memory", + "category": category, + "session_id": session_id, + } + ], + ) + except Exception as e: + logger.warning( + "kb_remember: stored YAML memory %s but vector mirror failed: %s", + store_result["entry_id"], + e, + ) + + return { + "status": "ok", + "entry_id": store_result["entry_id"], + "superseded_id": store_result.get("superseded_id"), + "total_memories": await asyncio.to_thread(memory.count), + } + + +async def _handle_kb_get_memories( + arguments: dict, + *, + memory: ExplicitMemory, +) -> dict: + entries = await asyncio.to_thread(memory.get_all) + return {"status": "ok", "count": len(entries), "memories": entries} + + +# --------------------------------------------------------------------------- +# Tool defs + handler map (used by the merged server and the test wrapper) +# --------------------------------------------------------------------------- + + +def _memory_tools_and_handlers( + kb: KnowledgeBase, + runtime_dir: str | Path | None = None, +): + """Build the explicit-memory ``(tool defs, handler map)``. + + Combined with the other concern modules' tools under one MCP ``Server`` by + :func:`dsagt.mcp.server.create_dsagt_server`. ``ExplicitMemory`` lives in + ``/.dsagt/`` alongside config.yaml and state.yaml — the + server-owned internals — with ``runtime_dir`` (falling back to the KB + index's parent) as the project dir. + """ + project_dir = Path(runtime_dir) if runtime_dir else kb.index_dir.parent + memory = ExplicitMemory(runtime_dir=project_dir / ".dsagt") + + handlers = { + "kb_remember": partial(_handle_kb_remember, kb=kb, memory=memory), + "kb_get_memories": partial(_handle_kb_get_memories, memory=memory), + } + + tools = [ + types.Tool( + name="kb_remember", + description=( + "Store a user-confirmed fact as an explicit memory. " + "These persist across sessions. Use 'supersedes' to replace an outdated memory." + ), + inputSchema={ + "type": "object", + "properties": { + "text": { + "type": "string", + "description": "The fact to remember", + }, + "category": { + "type": "string", + "description": "Classification tag", + }, + "session_id": { + "type": "string", + "description": "Current session identifier", + }, + "supersedes": { + "type": "string", + "description": "entry_id of an existing memory this replaces", + }, + }, + "required": ["text"], + }, + ), + types.Tool( + name="kb_get_memories", + description=( + "Get all active explicit memories for this project. " + "Call at session start to load project context." + ), + inputSchema={"type": "object", "properties": {}}, + ), + ] + return tools, handlers + + +def create_memory_server( + kb: KnowledgeBase, + runtime_dir: str | Path | None = None, +): + """Create a standalone MCP server exposing only the explicit-memory tools. + + Test-facing API: tests call it with a mock KB and drive the server via + ``call_tool_sync()``. The merged ``dsagt-server`` uses + :func:`_memory_tools_and_handlers` directly instead of this wrapper. + """ + tools, handlers = _memory_tools_and_handlers(kb, runtime_dir) + return build_dispatch_server( + "memory", tools, handlers, {t: "memory" for t in handlers} + ) diff --git a/src/dsagt/mcp/registry_tools.py b/src/dsagt/mcp/registry_tools.py new file mode 100644 index 0000000..f26a9fb --- /dev/null +++ b/src/dsagt/mcp/registry_tools.py @@ -0,0 +1,556 @@ +"""MCP tools for the tool registry, execution, and provenance. + +The "tool lifecycle" surface of ``dsagt-server``: define a tool spec +(``save_code_spec``), discover tools (``get_registry`` / ``search_registry``), +execute / gather (``read_file`` / ``http_request`` / ``run_command`` / +``install_dependencies``), and reconstruct a reproducible pipeline from the +recorded executions (``reconstruct_pipeline``). + +Tool specs are saved as markdown files in the runtime tools directory and +indexed into a ChromaDB collection for semantic search. Server configuration +(embedding credentials) flows through env vars (LLM_API_KEY, OPENAI_BASE_URL, +EMBEDDING_MODEL) set by ``dsagt start``. + +These definitions + handlers run inside the merged ``dsagt-server`` (see +:mod:`dsagt.mcp.server`); ``create_registry_server`` is retained only as a +test-facing constructor. Skill tools (``save_skill`` / ``search_skills`` / +``install_skill``) live in :mod:`dsagt.mcp.skill_tools`. +""" + +import json +import logging +import subprocess +import sys +from functools import partial +from pathlib import Path + +import httpx +import yaml + +import mcp.types as types + +from dsagt.knowledge import KnowledgeBase +from dsagt.mcp.server import build_dispatch_server +from dsagt.observability import ( + obs, + registry_install_deps_span, + registry_reconstruct_pipeline_span, + registry_save_code_span, +) +from dsagt.provenance import CodeUseIndexer, reconstruct_pipeline +from dsagt.registry import CODES_COLLECTION, CodeRegistry + +logger = logging.getLogger(__name__) + + +def _install_dependencies(packages: list[str], timeout: int = 120) -> str: + """Install packages using uv pip install. Returns a status string.""" + cmd = ["uv", "pip", "install", "--python", sys.executable] + packages + try: + result = subprocess.run(cmd, capture_output=True, text=True, timeout=timeout) + if result.returncode == 0: + output = result.stdout.strip() + return f"Successfully installed: {', '.join(packages)}\n{output}" + else: + return ( + f"Installation failed (exit code {result.returncode}):\n" + f"{result.stderr.strip()}" + ) + except subprocess.TimeoutExpired: + return f"Installation timed out after {timeout}s for: {', '.join(packages)}" + except FileNotFoundError: + return ( + "Error: 'uv' command not found. Install uv: https://github.com/astral-sh/uv" + ) + + +# --------------------------------------------------------------------------- +# Per-tool handlers (module-level, explicit dependencies) +# --------------------------------------------------------------------------- + + +async def _handle_read_file(arguments: dict) -> str: + path = Path(arguments["path"]) + try: + return path.read_text() + except ( + FileNotFoundError, + PermissionError, + IsADirectoryError, + OSError, + UnicodeDecodeError, + ) as e: + return f"Error reading file: {e}" + + +async def _handle_http_request(arguments: dict) -> str: + url = arguments["url"] + method = arguments.get("method", "GET") + headers = arguments.get("headers", {}) + try: + async with httpx.AsyncClient(follow_redirects=True) as client: + response = await client.request( + method=method, + url=url, + headers=headers, + timeout=30.0, + ) + return f"Status: {response.status_code}\n\n{response.text}" + except (httpx.HTTPError, httpx.InvalidURL) as e: + return f"Error making request: {e}" + + +async def _handle_run_command(arguments: dict) -> str: + command = arguments["command"] + args = arguments.get("args", []) + timeout = arguments.get("timeout", 10) + try: + result = subprocess.run( + [command] + args, + capture_output=True, + text=True, + timeout=timeout, + ) + except subprocess.TimeoutExpired: + return f"Command timed out after {timeout} seconds" + except FileNotFoundError: + return f"Command '{command}' not found" + + output = "" + if result.stdout: + output += f"STDOUT:\n{result.stdout}\n" + if result.stderr: + output += f"STDERR:\n{result.stderr}\n" + output += f"\nReturn code: {result.returncode}" + return output + + +async def _handle_save_code_spec( + arguments: dict, + *, + registry: CodeRegistry, +) -> str: + spec = arguments["spec"] + # Some MCP clients (notably Claude Sonnet/Haiku 4.x) serialize nested + # object args as JSON strings instead of objects. Accept both shapes. + if isinstance(spec, str): + try: + spec = json.loads(spec) + except json.JSONDecodeError as e: + return f"Error: spec must be a JSON object (or string-encoded JSON object): {e}" + with registry_save_code_span(spec.get("name")): + obs.set("language", spec.get("language")) + obs.set("n_dependencies", len(spec.get("dependencies") or [])) + obs.set("n_tags", len(spec.get("tags") or [])) + try: + action = registry.save_tool(spec) + except (KeyError, ValueError, OSError) as e: + obs.event("save_tool_failed", error=str(e)[:256]) + return f"Error saving tool spec: {e}" + + # Codes share the skill envelope, so a fresh registration mirrors into + # the agent's native skills dir right away (next session discovers it). + from dsagt.agents import refresh_native_skills + + refresh_native_skills(registry.runtime_dir) + tool_count = len(registry.list_codes_raw()) + obs.set("action", action) + obs.set("registry_size", tool_count) + message = ( + f"Tool '{spec['name']}' {action} successfully. " + f"Registry now contains {tool_count} tools." + ) + deps = spec.get("dependencies", []) + if deps: + with registry_install_deps_span(deps): + dep_result = _install_dependencies(deps) + if dep_result.startswith("Successfully installed:"): + obs.set("status", "ok") + else: + obs.set("status", "failed") + obs.event("install_failed", message=dep_result[:256]) + message += f"\n\nDependency installation:\n{dep_result}" + return message + + +async def _handle_get_registry( + arguments: dict, + *, + registry: CodeRegistry, +) -> str: + tools = registry.list_codes_raw() + if not tools: + return "Registry is empty. No tools registered yet." + return yaml.dump({"codes": tools}, default_flow_style=False, sort_keys=False) + + +async def _handle_search_registry( + arguments: dict, + *, + registry: CodeRegistry, + kb: KnowledgeBase | None, +) -> str: + code_name = arguments.get("code_name") + query = arguments.get("query", "") + tag = arguments.get("tag") + top_k = arguments.get("top_k", 10) + + if code_name: + tool = registry.get_code(code_name) + if tool: + return f"Found tool '{code_name}':\n\n" + yaml.dump( + tool, default_flow_style=False, sort_keys=False + ) + return f"No tool named '{code_name}'." + + if kb is None: + return ( + "search_registry requires a configured knowledge base " + "(set embedding.api_key + embedding.base_url + embedding.model " + "in .dsagt/config.yaml). Use search_registry with an exact " + "code_name for KB-free lookups." + ) + + # Single ``tools`` collection — bundled and registered entries + # coexist, distinguished by ``metadata.source`` if needed. + results = kb.search( + query=query or "tool", + collection=CODES_COLLECTION, + top_k=top_k * 3 if tag else top_k, + ) + if tag and results: + results = [ + r + for r in results + if tag in r.get("chunk", {}).get("metadata", {}).get("tags", "") + ][:top_k] + if not results: + return "No tools found matching the query." + + summaries = [] + for r in results: + chunk = r.get("chunk", {}) + meta = chunk.get("metadata", {}) + summaries.append( + f"- **{meta.get('code_name', 'unknown')}** " + f"(score: {r.get('score', 0):.2f})\n" + f" {chunk.get('text', '')[:200]}" + ) + return f"Found {len(results)} tool(s):\n\n" + "\n\n".join(summaries) + + +async def _handle_reconstruct_pipeline( + arguments: dict, + *, + runtime_dir: Path, + kb: KnowledgeBase | None = None, +) -> str: + fmt = arguments.get("format", "bash") + trace_dir = runtime_dir / "trace_archive" + # Index the session's tool-use first: reconstruct is the moment the pipeline + # is "done enough" to review, so make the just-run executions searchable now + # rather than waiting on the heartbeat. Idempotent + file-locked, so this + # is safe to fire alongside the heartbeat's own CodeUseIndexer. + if kb is not None: + try: + CodeUseIndexer(kb, runtime_dir).tick() + except Exception as e: # noqa: BLE001 — indexing is best-effort here + logger.warning("tool_use indexing before reconstruct failed: %s", e) + with registry_reconstruct_pipeline_span(fmt): + try: + script = reconstruct_pipeline(trace_dir, fmt=fmt) + except (FileNotFoundError, ValueError, OSError) as e: + obs.event("reconstruct_failed", error=str(e)[:256]) + return f"Error reconstructing pipeline: {e}" + obs.set("output_chars", len(script)) + return script + + +async def _handle_install_dependencies( + arguments: dict, + *, + registry: CodeRegistry, +) -> str: + code_name = arguments.get("code_name") + tools = registry.list_codes_raw() + if not tools: + return "Registry is empty. No tools registered yet." + + all_deps = [] + codes_with_deps = [] + for tool in tools: + if code_name and tool.get("name") != code_name: + continue + code_deps = tool.get("dependencies", []) + if code_deps: + all_deps.extend(code_deps) + codes_with_deps.append(tool["name"]) + + if not all_deps: + scope = f"tool '{code_name}'" if code_name else "registry" + return f"No dependencies declared in {scope}." + + seen = set() + unique_deps = [d for d in all_deps if not (d in seen or seen.add(d))] + + with registry_install_deps_span(unique_deps): + obs.set("scope_code", code_name) + obs.set("n_tools_with_deps", len(codes_with_deps)) + result = _install_dependencies(unique_deps) + if result.startswith("Successfully installed:"): + obs.set("status", "ok") + else: + obs.set("status", "failed") + obs.event("install_failed", message=result[:256]) + return f"Installing dependencies for: {', '.join(codes_with_deps)}\n\n{result}" + + +# --------------------------------------------------------------------------- +# Tool defs + handler map (used by the merged server and the test wrapper) +# --------------------------------------------------------------------------- + + +def _registry_tools_and_handlers( + registry: CodeRegistry, + kb: KnowledgeBase | None = None, +): + """Build the registry/execution/provenance ``(tool defs, handler map)``. + + Combined with the other concern modules' tools under one MCP ``Server`` by + :func:`dsagt.mcp.server.create_dsagt_server`. + """ + runtime_dir = Path(registry.runtime_dir) + + handlers = { + "read_file": _handle_read_file, + "http_request": _handle_http_request, + "run_command": _handle_run_command, + "save_code_spec": partial(_handle_save_code_spec, registry=registry), + "get_registry": partial(_handle_get_registry, registry=registry), + "search_registry": partial(_handle_search_registry, registry=registry, kb=kb), + "reconstruct_pipeline": partial( + _handle_reconstruct_pipeline, runtime_dir=runtime_dir, kb=kb + ), + "install_dependencies": partial( + _handle_install_dependencies, registry=registry + ), + } + + tools = [ + types.Tool( + name="read_file", + description="Read contents of a text file", + inputSchema={ + "type": "object", + "properties": { + "path": { + "type": "string", + "description": "Path to the file to read", + }, + }, + "required": ["path"], + }, + ), + types.Tool( + name="http_request", + description="Make an HTTP request to fetch documentation or API specs", + inputSchema={ + "type": "object", + "properties": { + "url": {"type": "string", "description": "URL to request"}, + "method": { + "type": "string", + "description": "HTTP method", + "default": "GET", + }, + "headers": { + "type": "object", + "description": "Optional headers", + }, + }, + "required": ["url"], + }, + ), + types.Tool( + name="run_command", + description="Execute a command to get help/usage information", + inputSchema={ + "type": "object", + "properties": { + "command": { + "type": "string", + "description": "Command to execute", + }, + "args": { + "type": "array", + "items": {"type": "string"}, + "description": "Arguments (e.g., ['--help'])", + "default": [], + }, + "timeout": {"type": "number", "default": 10}, + }, + "required": ["command"], + }, + ), + types.Tool( + name="save_code_spec", + description=( + "Register a CLI code: writes a skill-standard spec dir " + "(codes//SKILL.md), mirrored into the agent's native " + "skills dir immediately so future sessions auto-discover it" + ), + inputSchema={ + "type": "object", + "properties": { + # ``anyOf`` accepts both a structured object and a + # JSON-encoded string. Some MCP clients (notably + # Claude Sonnet/Haiku 4.x) serialize nested object + # arguments as JSON strings instead of objects; the + # handler unwraps either shape. + "spec": { + "description": "Tool specification (object or JSON-encoded string)", + "anyOf": [ + { + "type": "object", + "properties": { + "name": { + "type": "string", + "description": ( + "Unique code name — lowercase " + "letters, digits, hyphens (e.g. " + "'scan-directory')" + ), + }, + "description": { + "type": "string", + "description": ( + "What the code does and when to " + "use it — phrase as 'Use when " + "…' so native skill routing can " + "match it" + ), + }, + "executable": { + "type": "string", + "description": "Command to execute", + }, + "parameters": { + "type": "object", + "description": "Parameter definitions", + "additionalProperties": { + "type": "object", + "properties": { + "type": { + "type": "string", + "description": "Parameter type", + }, + "required": {"type": "boolean"}, + "description": {"type": "string"}, + "default": { + "description": "Default value" + }, + "cli": { + "type": "string", + "description": ( + "How to render this parameter on the command line: " + "'positional[:N]' for positional args, '--name' or '-n' " + "for spaced flags, '--name=' or '-n=' for glued flags, " + "'key=' for dd-style key=value. Defaults to '--' " + "if omitted." + ), + }, + }, + "required": ["type", "description"], + }, + }, + "dependencies": { + "type": "array", + "items": {"type": "string"}, + "description": "Python packages to install", + }, + "tags": { + "type": "array", + "items": {"type": "string"}, + "description": "Tags for categorizing the tool", + }, + }, + "required": [ + "name", + "description", + "executable", + "parameters", + ], + }, + {"type": "string"}, + ], + }, + }, + "required": ["spec"], + }, + ), + types.Tool( + name="get_registry", + description="Get all tools from the registry", + inputSchema={"type": "object", "properties": {}}, + ), + types.Tool( + name="search_registry", + description="Search for tools by name, tag, or description via semantic search.", + inputSchema={ + "type": "object", + "properties": { + "query": {"type": "string", "description": "Search query"}, + "tag": {"type": "string", "description": "Filter by tag"}, + "code_name": { + "type": "string", + "description": "Exact tool name lookup", + }, + "top_k": {"type": "integer", "default": 10}, + }, + }, + ), + types.Tool( + name="reconstruct_pipeline", + description="Reconstruct a reproducible pipeline script from tool execution records.", + inputSchema={ + "type": "object", + "properties": { + "format": { + "type": "string", + "enum": ["bash", "snakemake"], + "default": "bash", + }, + }, + }, + ), + types.Tool( + name="install_dependencies", + description="Install Python dependencies for one or all tools in the registry.", + inputSchema={ + "type": "object", + "properties": { + "code_name": { + "type": "string", + "description": "Install deps for a specific tool (omit for all)", + }, + }, + }, + ), + ] + return tools, handlers + + +def create_registry_server( + registry: CodeRegistry, + kb: KnowledgeBase | None = None, +): + """Create a standalone MCP server exposing only the registry/exec/provenance tools. + + Test-facing API: tests call with a mock registry and drive the server via + ``call_tool_sync()``. The merged ``dsagt-server`` uses + :func:`_registry_tools_and_handlers` directly instead of this wrapper. + """ + tools, handlers = _registry_tools_and_handlers(registry, kb) + return build_dispatch_server( + "registry", tools, handlers, {t: "registry" for t in handlers} + ) diff --git a/src/dsagt/mcp/server.py b/src/dsagt/mcp/server.py new file mode 100644 index 0000000..e010140 --- /dev/null +++ b/src/dsagt/mcp/server.py @@ -0,0 +1,497 @@ +"""DSAGT MCP Server — the single merged registry + knowledge server. + +Supersedes the two former servers (``dsagt-registry-server`` + +``dsagt-knowledge-server``). Both previously constructed their own +:class:`~dsagt.knowledge.KnowledgeBase` — two embedders, two Chroma accesses, +and a write-here/read-there hazard on the ``skills_catalog__*`` collections +(synced by knowledge, searched by registry). Merging into one process gives one +embedder, one Chroma owner, one ``init_tracing``, and one MCP server per agent. + +The heavy/risky work is already offloaded out of the event loop (``run_command`` +→ ``dsagt-run`` subprocess; ``kb_ingest`` → background job thread), so collapsing +the two processes costs little isolation. + +Tool *definitions* and *handlers* live in their concern modules +(:mod:`~dsagt.mcp.registry_tools` / :mod:`~dsagt.mcp.knowledge_tools` / +:mod:`~dsagt.mcp.memory_tools` / :mod:`~dsagt.mcp.skill_tools`); this module only +composes their ``(tools, handlers)`` under one dispatch shell +(:func:`build_dispatch_server`) and owns the shared-KB startup. The factory +imports are *lazy* (inside :func:`create_dsagt_server` / :func:`main`) so the +concern modules can import :func:`build_dispatch_server` from here without a +cycle. + +See ``design-notes/skills-catalog-server-merge.md`` §2. + +Backward compatibility is **rebuild-not-migrate**: a project created against the +old two-server layout adopts this by re-running ``dsagt start`` (which +regenerates the per-agent MCP config to a single ``dsagt`` server). See the +upgrade note in the README. +""" + +import os + +# Set before any import that may pull in PyTorch / sentence-transformers +# (e.g. ``dsagt.knowledge`` below): prevents a fatal OpenMP crash when multiple +# libraries each bundle their own libomp. +os.environ["PYTHONUNBUFFERED"] = "1" +os.environ.setdefault("KMP_DUPLICATE_LIB_OK", "TRUE") + +import asyncio # noqa: E402 +import json # noqa: E402 +import logging # noqa: E402 +import threading # noqa: E402 +from pathlib import Path # noqa: E402 + +import yaml # noqa: E402 + +import mcp.server.stdio # noqa: E402 +import mcp.types as types # noqa: E402 +from mcp.server.lowlevel import Server, NotificationOptions # noqa: E402 +from mcp.server.models import InitializationOptions # noqa: E402 + +from dsagt.knowledge import KnowledgeBase # noqa: E402 +from dsagt.observability import open_span # noqa: E402 +from dsagt.registry import SkillRegistry, CodeRegistry # noqa: E402 + +logger = logging.getLogger(__name__) + + +# --------------------------------------------------------------------------- +# Shared dispatch shell (used by the merged server *and* the per-concern +# test-facing ``create_*_server`` wrappers) +# --------------------------------------------------------------------------- + + +def build_dispatch_server( + name: str, tools, handlers, tool_category: dict[str, str] | None = None +) -> Server: + """Wrap a ``(tools, handlers)`` pair in a configured MCP ``Server``. + + One dispatch contract for every concern module: catch + wrap errors, then + format by return type — a handler that returns ``str`` passes through, one + that returns ``dict`` is JSON-encoded. This is a superset of the old + per-server behavior (registry handlers returned ``str`` and never raised; + knowledge handlers returned ``dict`` and raised ``ValueError`` on bad + input), so it is behavior-preserving for both. + + ``tool_category`` maps tool name → concern (``memory`` / ``skill`` / + ``knowledge`` / ``registry``). Each call opens one categorization-root span + (named for the tool) tagged ``dsagt.source=``; the subsystem spans + the handler opens (``kb.*`` / ``registry.*``) nest under it — and inherit the + category, so the source reflects the tool the agent called, not whichever + subsystem did the work. Tracing no-ops outside a project, so this is inert + in the single-concern test servers / one-shot tools. + """ + tool_category = tool_category or {} + server = Server(name) + + @server.list_tools() + async def list_tools() -> list[types.Tool]: + return tools + + @server.call_tool() + async def call_tool(tool_name: str, arguments: dict) -> list[types.TextContent]: + handler = handlers[tool_name] # KeyError = bug in list_tools schema + with open_span(tool_name, source=tool_category.get(tool_name)): + try: + result = await handler(arguments) + except ValueError as e: + result = {"status": "error", "error": str(e)} + except Exception as e: + logger.exception("Unexpected error in tool '%s'", tool_name) + result = {"status": "error", "error": f"Unexpected error: {e}"} + text = ( + result + if isinstance(result, str) + else json.dumps(result, ensure_ascii=False) + ) + return [types.TextContent(type="text", text=text)] + + return server + + +HEARTBEAT_INTERVAL_S = 45.0 + + +async def _heartbeat(collector, tool_indexer, interval: float, project_dir) -> None: + """Periodically run the trace collector + tool-use indexer on wall-clock time. + + Runs regardless of tool traffic, so a quiet session (the agent thinking, + editing with its own tools, plain chat) is still captured. Both block on + disk (+ MLflow / embedding), so they run in a worker thread to keep handlers + responsive; a failure is logged, never fatal. + + It also records the live session's trace-source token into ``state.yaml`` + once resolved, so the *next* session's startup catch-up can pin this exact + session even if this one is killed ungracefully (the deferred final-turn + flush never runs). Uniform across agents — JSONL or SQLite. + """ + recorded = False + while True: + await asyncio.sleep(interval) + if collector is not None: + try: + n = await asyncio.to_thread(collector.collect) + if n: + logger.info("Trace heartbeat: logged %d trace(s)", n) + except Exception as e: # noqa: BLE001 + logger.warning("Trace heartbeat collect failed: %s", e) + if not recorded: + try: + source = collector.active_source() + if source is not None: + from dsagt.session import record_trace_source + + record_trace_source(project_dir, source) + recorded = True + except Exception as e: # noqa: BLE001 + logger.debug("Could not record trace source: %s", e) + if tool_indexer is not None: + try: + # tick_traced (not tick): opens a code_use categorization root on + # the worker thread so the indexer's kb.* writes nest under it + # instead of orphaning as untagged top-level traces. + n = await asyncio.to_thread(tool_indexer.tick_traced) + if n: + logger.info("Tool-use heartbeat: indexed %d record(s)", n) + except Exception as e: # noqa: BLE001 + logger.warning("Tool-use heartbeat tick failed: %s", e) + + +async def _run_stdio( + server: Server, name: str, collector=None, tool_indexer=None, project_dir=None +) -> None: + async with mcp.server.stdio.stdio_server() as (read_stream, write_stream): + hb = ( + asyncio.create_task( + _heartbeat( + collector, tool_indexer, HEARTBEAT_INTERVAL_S, project_dir + ) + ) + if (collector is not None or tool_indexer is not None) + else None + ) + try: + await server.run( + read_stream, + write_stream, + InitializationOptions( + server_name=name, + server_version="0.1.0", + capabilities=server.get_capabilities( + notification_options=NotificationOptions(), + experimental_capabilities={}, + ), + ), + ) + finally: + if hb is not None: + hb.cancel() + try: + await hb + except asyncio.CancelledError: + pass + # Best-effort end-of-session flush of the deferred final turn + + # any unindexed tool-use — covers the graceful-exit case. If + # this is killed before it runs, the next session's startup + # catch-up re-collects the previous session's transcript + # (session.catch_up_extraction → _catch_up_traces, pinned to the + # recorded transcript path); session-qualified acks make both + # paths idempotent. Tool-use likewise re-indexes via its ack set. + if collector is not None: + try: + await asyncio.to_thread(collector.collect, include_last=True) + except Exception as e: # noqa: BLE001 + logger.warning("Trace heartbeat final flush failed: %s", e) + if tool_indexer is not None: + try: + await asyncio.to_thread(tool_indexer.tick) + except Exception as e: # noqa: BLE001 + logger.warning("Tool-use final flush failed: %s", e) + + +# --------------------------------------------------------------------------- +# Composition — merge the four concern modules' tools under one Server +# --------------------------------------------------------------------------- + + +def create_dsagt_server( + registry: CodeRegistry, + kb: KnowledgeBase | None, + skill_registry: SkillRegistry | None, + runtime_dir: str | Path | None = None, +): + """Compose the registry + knowledge + memory + skill tools under one ``Server``. + + Test-facing API: build the registries + a (mock) KB, then drive the + returned server via ``call_tool_sync()``. ``main()`` constructs the real + deps from project config before calling this. Factory imports are lazy to + keep the concern modules' top-level import of :func:`build_dispatch_server` + cycle-free. + """ + from dsagt.mcp.knowledge_tools import _knowledge_tools_and_handlers + from dsagt.mcp.memory_tools import _memory_tools_and_handlers + from dsagt.mcp.registry_tools import _registry_tools_and_handlers + from dsagt.mcp.skill_tools import _skill_tools_and_handlers + + # (category, group) — the category is the dsagt.source bucket stamped on + # every trace rooted at one of that group's tools. + groups = [ + ("registry", _registry_tools_and_handlers(registry, kb)), + ("knowledge", _knowledge_tools_and_handlers(kb)), + ("memory", _memory_tools_and_handlers(kb, runtime_dir)), + ("skill", _skill_tools_and_handlers(skill_registry, kb, runtime_dir)), + ] + + tools: list[types.Tool] = [] + handlers: dict = {} + tool_category: dict[str, str] = {} + for category, (g_tools, g_handlers) in groups: + overlap = set(handlers) & set(g_handlers) + if overlap: + raise RuntimeError( + f"dsagt-server tool-name collision across modules: {overlap}" + ) + tools += g_tools + handlers.update(g_handlers) + for tool_name in g_handlers: + tool_category[tool_name] = category + + return build_dispatch_server("dsagt", tools, handlers, tool_category) + + +def _build_kb_from_config(config: dict, project_dir: Path) -> KnowledgeBase: + """Construct the one shared KnowledgeBase from project config. + + The single home for embedding-backend selection + the cross-backend + leakage guard that the two former server mains duplicated near-verbatim. + """ + from dsagt.session import REGISTRY_DIR, setup_runtime_kb + + # embedding is a backfilled code default (not a written config choice); + # chunk_size / rerank default in KnowledgeBase itself. + emb_config = config.get("embedding", {}) + + backend = (emb_config.get("backend") or "local").lower() + if backend not in ("local", "api"): + raise ValueError( + f"embedding.backend must be 'local' or 'api' (got {backend!r})" + ) + + # Cross-backend leakage guard: HuggingFace identifiers ("org/repo") and + # OpenAI-style aliases ("text-embedding-3-small") share the same + # EMBEDDING_MODEL env var in most setups. When backend=local but the + # resolved model is an OpenAI-style alias (no slash), drop the override so + # we fall back to the LocalEmbedder default rather than 404 from HF. + raw_model = (emb_config.get("model") or "").strip() + model: str | None = None + base_url: str | None = None + api_key: str | None = None + if raw_model and not raw_model.startswith("${"): + looks_hf = "/" in raw_model + if backend == "local" and not looks_hf: + logger.warning( + "Ignoring embedding.model=%r for backend=local (does not look " + "like a HuggingFace identifier). Falling back to the " + "LocalEmbedder default.", + raw_model, + ) + else: + model = raw_model + if backend == "api": + base_url = emb_config.get("base_url") or "" + # Credentials are never on disk: the api key comes from the shell env + # (EMBEDDING_API_KEY), threaded into MCP children via the env block. + api_key = os.environ.get("EMBEDDING_API_KEY") or emb_config.get("api_key") or "" + if not base_url: + raise ValueError( + "embedding.backend='api' requires embedding.base_url in " + ".dsagt/config.yaml. Either set it to your OpenAI-compatible " + "endpoint, or change backend to 'local'." + ) + if not api_key or api_key.startswith("${"): + raise ValueError( + "embedding.backend='api' requires the EMBEDDING_API_KEY env " + "var (export it in your shell), or change backend to 'local'." + ) + + from dsagt.session import _recency_half_life + + runtime_kb_dir = setup_runtime_kb(REGISTRY_DIR / "kb_index", project_dir) + logger.info("Knowledge backend: %s", backend) + kb = KnowledgeBase( + index_dir=runtime_kb_dir, + default_embedder=backend, + model=model, + base_url=base_url, + api_key=api_key, + recency_half_life_days=_recency_half_life(config), + ) + # Background-load the embedder so the model is ready when the agent's first + # search / kb call lands (otherwise the first call pays the ~5-10s + # sentence-transformers import + construction, which looks like a hang). + kb.preload_default_embedder() + return kb + + +def _spawn_catch_up(project_dir: Path, config: dict) -> None: + """Run :func:`dsagt.session.catch_up_extraction` in a daemon thread. + + Best-effort background catch-up of the previous session's post-session + work (tool-use indexing now; episodic stub later). Daemon so it never + holds the server open; exceptions are logged, never propagated. + """ + + def _run() -> None: + try: + from dsagt.session import catch_up_extraction + + result = catch_up_extraction(project_dir, config) + logger.info("Background catch-up complete: %s", result) + except Exception as e: # noqa: BLE001 + logger.warning("Background catch-up failed: %s", e) + + threading.Thread(target=_run, name="dsagt-catch-up", daemon=True).start() + + +def main(): + """Entry point for ``dsagt-server``. + + All configuration comes from the project directory: + - ``./.dsagt/config.yaml`` → project path + non-secret settings + - ``EMBEDDING_*`` env vars → embedding credentials + + No CLI arguments. By contract the agent's launch one-liner is + ``cd && ``, so cwd is project_dir for the MCP children it + spawns. + + The server owns the session lifecycle: it appends a new entry to + ``.dsagt/state.yaml`` (minting the session id) and spawns a background + thread that catches up post-session extraction for the *previous* + session — no reliable session-end trigger needed. + """ + from dsagt.observability import ( + find_project_config, + init_tracing, + ) + from dsagt.session import ( + DEFAULTS, + _deep_merge, + append_session, + resolve_env_vars, + session_tag, + ) + + project_dir, _cfg = find_project_config() + if project_dir is None: + raise RuntimeError( + "dsagt-server: no .dsagt/config.yaml in cwd " + f"({Path.cwd()}). Launch the agent from the project " + "directory (`cd && `)." + ) + + log_file = project_dir / "dsagt_server.log" + # Default INFO; users opt into DEBUG via DSAGT_LOG_LEVEL=DEBUG. At DEBUG, + # transitive libraries (httpcore, urllib3, llama_index, chromadb) flood + # stderr with one line per network op — when an agent pipes the MCP + # server's stderr into its own debug stream, human output gets buried. + _level_name = os.environ.get("DSAGT_LOG_LEVEL", "INFO").upper() + _level = getattr(logging, _level_name, logging.INFO) + logging.basicConfig( + level=_level, + format="%(asctime)s %(levelname)s %(name)s: %(message)s", + handlers=[ + logging.FileHandler(log_file, mode="a"), + logging.StreamHandler(), + ], + ) + logger.info("Server starting — project_dir: %s, log: %s", project_dir, log_file) + + cfg_file = project_dir / ".dsagt" / "config.yaml" + # Backfill code defaults (embedding, etc.) the same way ``load_config`` + # does — the written config carries only the user's init choices. + config = resolve_env_vars( + _deep_merge(DEFAULTS, yaml.safe_load(cfg_file.read_text()) or {}) + ) + + # Own the session lifecycle: mint this session's id into state.yaml and + # tag traces with it (replaces the DSAGT_SESSION_ID env minted by the old + # ``dsagt start``). Best-effort — never block startup on state I/O. + session_id = None + try: + entry = append_session(project_dir) + session_id = session_tag(config.get("project", ""), entry["id"]) + except Exception as e: # noqa: BLE001 + logger.warning("Could not mint session into state.yaml: %s", e) + + init_tracing("dsagt-server", session_id=session_id) + + # Catch up post-session extraction for the previous session in the + # background (tool-use indexing now; episodic stub later). Daemon thread: + # best-effort, never blocks or fails server startup. + _spawn_catch_up(project_dir, config) + + kb = _build_kb_from_config(config, project_dir) + + # Bundled tools are pre-embedded in the shared ~/dsagt-projects/kb_index/ + # by ``dsagt init`` (shared cache, one-time per machine) and + # copied into the project's kb_index by ``setup_runtime_kb`` above. No + # bundled embedding work happens here; save_code_spec incurs a single + # embed at save time. + registry = CodeRegistry( + runtime_dir=str(project_dir), + kb=kb, + ) + skill_reg = SkillRegistry( + source_skills_dir=None, + runtime_dir=str(project_dir), + kb=kb, + ) + + server = create_dsagt_server(registry, kb, skill_reg, runtime_dir=str(project_dir)) + + # The in-session trace heartbeat: read the live transcript → MLflow. The + # loop is agent-agnostic; ``make_trace_collector`` returns a collector for any + # agent with a registered (reader, translator) pair and ``None`` otherwise (so + # agents whose readers haven't landed yet simply run without it). + # Best-effort — a collector that can't be built never blocks the server. + collector = None + try: + from dsagt.memory import episodic_consumers + from dsagt.observability import resolve_tracking_uri + from dsagt.traces import make_trace_collector + + resolve_cfg = dict(config) + resolve_cfg["project_dir"] = str(project_dir) + collector = make_trace_collector( + config.get("agent"), + project_dir, + config.get("project", ""), + session_id or "", + resolve_tracking_uri(resolve_cfg), + extra_consumers=episodic_consumers(config, kb, project_dir, session_id), + ) + except Exception as e: # noqa: BLE001 + logger.warning("Could not start trace heartbeat: %s", e) + + # Tool-use indexer: incremental, idempotent embedding of dsagt-run records + # into the ``tool_use`` collection on the same heartbeat (no collector + # dependency — it reads trace_archive/, not the transcript). + tool_indexer = None + try: + from dsagt.provenance import CodeUseIndexer + + tool_indexer = CodeUseIndexer(kb, project_dir) + except Exception as e: # noqa: BLE001 + logger.warning("Could not start tool-use indexer: %s", e) + + try: + asyncio.run( + _run_stdio(server, "dsagt", collector, tool_indexer, project_dir) + ) + finally: + kb.close() + + +if __name__ == "__main__": + main() diff --git a/src/dsagt/mcp/skill_tools.py b/src/dsagt/mcp/skill_tools.py new file mode 100644 index 0000000..4a8bd7e --- /dev/null +++ b/src/dsagt/mcp/skill_tools.py @@ -0,0 +1,398 @@ +"""MCP tools for skill discovery, install, and catalog sources. + +The full skill surface of ``dsagt-server``, consolidated from the two former +servers: register a project skill (``save_skill``), enable + list external +catalog sources (``add_skill_source`` / ``list_skill_sources``), search the +catalog (``search_skills``), and install a catalog skill into the project +(``install_skill``). The catalog data plane + router live in +:mod:`dsagt.skills`; these handlers are the thin MCP wiring over it. + +Installed/created skills are natively auto-discovered by every supported agent, +so ``search_skills`` covers only the not-yet-installed *catalog* tier (plus the +no-embedder keyword fallback). + +These definitions + handlers run inside the merged ``dsagt-server`` (see +:mod:`dsagt.mcp.server`); ``create_skill_server`` is retained only as a +test-facing constructor. +""" + +import asyncio +import json +import logging +from functools import partial +from pathlib import Path + +import mcp.types as types + +from dsagt.knowledge import KnowledgeBase +from dsagt.mcp.server import build_dispatch_server +from dsagt.registry import SkillRegistry + +logger = logging.getLogger(__name__) + + +async def _handle_save_skill( + arguments: dict, + *, + skill_registry: SkillRegistry, +) -> str: + """Register a skill (workflow / agent instructions) for later reuse. + + Writes SKILL.md to ``/skills//`` and mirrors it into the + agent's native skills dir immediately, where every supported agent + auto-discovers it from its next session. No KB indexing — + ``search_skills`` covers only the not-yet-installed *catalog* tier, since + installed skills are already natively discoverable. + """ + spec = arguments["spec"] + if isinstance(spec, str): + try: + spec = json.loads(spec) + except json.JSONDecodeError as e: + return f"Error: spec must be a JSON object (or string-encoded JSON object): {e}" + body = arguments.get("body") + reference_files = arguments.get("reference_files") + if isinstance(reference_files, str): + try: + reference_files = json.loads(reference_files) + except json.JSONDecodeError as e: + return f"Error: reference_files must be a JSON object: {e}" + try: + action = skill_registry.save_skill( + spec, body=body, reference_files=reference_files + ) + except (KeyError, ValueError, OSError) as e: + return f"Error saving skill: {e}" + from dsagt.agents import refresh_native_skills + + refresh_native_skills(skill_registry.runtime_dir) + skill_count = len(skill_registry.list_skills()) + return ( + f"Skill '{spec['name']}' {action} successfully. " + f"Registry now contains {skill_count} skills." + ) + + +async def _handle_search_skills( + arguments: dict, + *, + kb: KnowledgeBase | None, + skill_registry: SkillRegistry | None, +) -> str: + if skill_registry is None: + return "search_skills is unavailable (no skill registry configured)." + + from dsagt.skills import SkillRouter + + router = SkillRouter(skill_registry=skill_registry, kb=kb) + return router.search( + arguments.get("query", ""), + top_k=arguments.get("top_k", 10), + tag=arguments.get("tag"), + skill_name=arguments.get("skill_name"), + ) + + +async def _handle_install_skill( + arguments: dict, + *, + runtime_dir: Path, +) -> str: + """Install a catalog skill into ``/skills//``. + + The skill's files land on disk and are mirrored into the agent's native + skills dir immediately — usable right away (the agent reads/follows its + SKILL.md, which is all native invocation does); hands-free auto-discovery + kicks in at the agent's next session, with no user action. + """ + from dsagt.skills import SkillRouter + + name = arguments.get("skill_name") + if not name: + return "install_skill requires 'skill_name'." + try: + info = SkillRouter().install(name, runtime_dir) + except LookupError as e: + return f"Error: {e}" + from dsagt.agents import refresh_native_skills + + refresh_native_skills(runtime_dir) + + # Bare confirmation by design: the install→use model and the + # license/PROVENANCE capture are already in the agent's instructions and on + # disk (PROVENANCE.txt), so repeating them on every install is just noise. + verb = "Updated" if info["action"] == "updated" else "Installed" + return f"{verb} '{info['name']}' → {info['dest_dir']}/" + + +async def _handle_add_skill_source( + arguments: dict, + *, + kb: KnowledgeBase, + runtime_dir: Path, +) -> dict: + """Enable a skill source (known name or GitHub URL): clone + index the catalog.""" + from dsagt.skills import ( + KNOWN_SOURCES, + SkillRouter, + persist_source_to_config, + resolve_source, + ) + + source = arguments.get("source") + if not source: + return { + "error": "add_skill_source requires 'source' (known name or GitHub URL)." + } + try: + spec = resolve_source(source) + if isinstance(source, str) and source in KNOWN_SOURCES: + spec.setdefault("name", source) + router = SkillRouter(kb=kb) + stats = await asyncio.to_thread(router.sync, source) + except (ValueError, RuntimeError) as e: + return {"error": str(e)} + persist_source_to_config( + runtime_dir, {"name": spec.get("name", stats["slug"]), **spec} + ) + return { + "source": spec["url"], + "slug": stats["slug"], + "skills_indexed": stats["indexed"], + "note": "Searchable via search_skills; install one with install_skill.", + } + + +async def _handle_list_skill_sources(arguments: dict, *, kb: KnowledgeBase) -> dict: + """List known skill sources, each flagged synced/available with its count. + + A source is ``synced`` (searchable via ``search_skills``) only after an + ``add_skill_source`` call has cloned + indexed it; otherwise it is + ``available`` (known name + URL, nothing indexed yet). Reporting the + flag + ``indexed`` count inline means the agent doesn't have to cross- + reference a separate ``synced_collections`` list to tell the difference. + """ + from dsagt.registry import CATALOG_COLLECTION_PREFIX, catalog_collection + from dsagt.skills import KNOWN_SOURCES, SkillRouter, _repo_slug + + synced = {c for c in kb.collections if c.startswith(CATALOG_COLLECTION_PREFIX)} + + # Single source of truth for the per-source synced/indexed view (shared + # with the CLI `skills list --catalog`). + sources = { + s["name"]: { + "url": s["url"], + "description": s["description"], + "synced": s["synced"], + "indexed": s["indexed"], + } + for s in SkillRouter(kb=kb).list_sources() + } + + # Surface any synced catalog whose source isn't in KNOWN_SOURCES (added + # by raw GitHub URL) so the count is never silently dropped. + known_colls = { + catalog_collection(_repo_slug(s["url"])) for s in KNOWN_SOURCES.values() + } + extra = sorted(synced - known_colls) + + any_synced = any(v["synced"] for v in sources.values()) or bool(extra) + return { + "sources": sources, + "other_synced_collections": extra, + "note": ( + "add_skill_source to sync a source whose synced=false; " + "then search_skills to browse. search_skills only sees synced sources." + if any_synced + else "No catalog synced yet — add_skill_source " + "(e.g. 'k-dense-ai') to enable one, then search_skills to browse." + ), + } + + +# --------------------------------------------------------------------------- +# Tool defs + handler map (used by the merged server and the test wrapper) +# --------------------------------------------------------------------------- + + +def _skill_tools_and_handlers( + skill_registry: SkillRegistry | None, + kb: KnowledgeBase | None = None, + runtime_dir: str | Path | None = None, +): + """Build the skill ``(tool defs, handler map)``. + + Combined with the other concern modules' tools under one MCP ``Server`` by + :func:`dsagt.mcp.server.create_dsagt_server`. ``runtime_dir`` (the project + dir, where skills install + sources persist) falls back to the skill + registry's ``runtime_dir`` then the KB index's parent. + """ + rt: Path | None = Path(runtime_dir) if runtime_dir else None + if rt is None and skill_registry is not None: + rt = Path(skill_registry.runtime_dir) + if rt is None and kb is not None: + rt = kb.index_dir.parent + + handlers = { + "save_skill": partial(_handle_save_skill, skill_registry=skill_registry), + "search_skills": partial( + _handle_search_skills, kb=kb, skill_registry=skill_registry + ), + "install_skill": partial(_handle_install_skill, runtime_dir=rt), + "add_skill_source": partial(_handle_add_skill_source, kb=kb, runtime_dir=rt), + "list_skill_sources": partial(_handle_list_skill_sources, kb=kb), + } + + tools = [ + types.Tool( + name="save_skill", + description=( + "Register a skill (agent workflow / instructions) into " + "/skills//SKILL.md, mirrored into the agent's " + "native skills dir immediately so future sessions " + "auto-discover it — no restart or user action needed. " + "Symmetric with save_code_spec — use this when you've " + "designed a reusable instruction set you want future " + "sessions to load automatically." + ), + inputSchema={ + "type": "object", + "properties": { + # ``anyOf`` for spec mirrors save_code_spec — accept + # both structured object and JSON-encoded string for + # MCP clients that serialize nested args. + "spec": { + "description": "Skill spec (object or JSON-encoded string)", + "anyOf": [ + { + "type": "object", + "properties": { + "name": { + "type": "string", + "description": "Unique skill identifier (becomes the directory name)", + }, + "description": { + "type": "string", + "description": "What the skill does / when to use it", + }, + "tags": { + "type": "array", + "items": {"type": "string"}, + "description": "Tags for categorizing the skill", + }, + }, + "required": ["name", "description"], + }, + {"type": "string"}, + ], + }, + "body": { + "type": "string", + "description": ( + "Markdown body of the SKILL.md (workflow / " + "instructions the agent will follow). When " + "updating an existing skill, omit to preserve " + "the existing body." + ), + }, + "reference_files": { + "description": ( + "Optional additional files to write into the " + "skill directory. Object mapping relative " + "path -> file contents, or JSON-encoded string." + ), + "anyOf": [ + { + "type": "object", + "additionalProperties": {"type": "string"}, + }, + {"type": "string"}, + ], + }, + }, + "required": ["spec"], + }, + ), + types.Tool( + name="add_skill_source", + description=( + "Enable an external agent-skill source (a known name like " + "'k-dense-ai'/'anthropic'/'antigravity'/'composio', or a GitHub URL). " + "Clones it and indexes its skills into the searchable catalog " + "(search_skills). Does NOT load them into context." + ), + inputSchema={ + "type": "object", + "properties": { + "source": { + "type": "string", + "description": "Known source name or GitHub repo URL / owner/repo", + }, + }, + "required": ["source"], + }, + ), + types.Tool( + name="list_skill_sources", + description="List known + synced external skill sources and their indexed catalogs.", + inputSchema={"type": "object", "properties": {}}, + ), + types.Tool( + name="search_skills", + description=( + "Search agent skills by name, tag, or description. Spans installed " + "skills and the external installable catalog. Catalog hits are marked " + "'[catalog]' — use install_skill to add one to this project." + ), + inputSchema={ + "type": "object", + "properties": { + "query": {"type": "string", "description": "Search query"}, + "tag": {"type": "string", "description": "Filter by tag"}, + "skill_name": { + "type": "string", + "description": "Exact skill name lookup", + }, + "top_k": {"type": "integer", "default": 10}, + }, + }, + ), + types.Tool( + name="install_skill", + description=( + "Install a skill from the external catalog (found via search_skills) " + "into this project. Copies SKILL.md + scripts/references and mirrors " + "it into the agent's native skills dir — usable immediately (read and " + "follow its SKILL.md); future sessions auto-discover it natively with " + "no user action." + ), + inputSchema={ + "type": "object", + "properties": { + "skill_name": { + "type": "string", + "description": "Catalog skill name to install", + }, + }, + "required": ["skill_name"], + }, + ), + ] + return tools, handlers + + +def create_skill_server( + skill_registry: SkillRegistry | None = None, + kb: KnowledgeBase | None = None, + runtime_dir: str | Path | None = None, +): + """Create a standalone MCP server exposing only the skill tools. + + Test-facing API: tests call it with mock deps and drive the server via + ``call_tool_sync()``. The merged ``dsagt-server`` uses + :func:`_skill_tools_and_handlers` directly instead of this wrapper. + """ + tools, handlers = _skill_tools_and_handlers(skill_registry, kb, runtime_dir) + return build_dispatch_server( + "skills", tools, handlers, {t: "skill" for t in handlers} + ) diff --git a/src/dsagt/memory.py b/src/dsagt/memory.py index 982cb07..e2d249f 100644 --- a/src/dsagt/memory.py +++ b/src/dsagt/memory.py @@ -10,53 +10,31 @@ demand via the ``kb_get_memories`` / ``kb_search`` MCP tools. Supports remember, supersede, remove, and retrieval. -**Episodic memory** (extract_session and friends): - End-of-session LLM extraction of facts, summaries, and insights. - Conversation history comes from MLflow traces with - ``service.name = "dsagt-proxy"`` — i.e., LLM calls forwarded - through ``dsagt-proxy`` and autologged via - ``mlflow.litellm.autolog()`` into a uniform shape - (``mlflow.spanInputs`` = request kwargs with ``messages``, - ``mlflow.spanOutputs`` = provider-specific response). - - Native agent OTel emission (Claude Code, Goose) does NOT feed - extraction, even though those traces are visible in the MLflow UI. - The reason is shape divergence: Claude Code puts conversation in - span events (``api_response_body``), Goose puts tool calls in - ``dispatch_tool_call`` spans with a domain-specific schema, and - LiteLLM autolog uses ``mlflow.spanInputs`` / ``mlflow.spanOutputs``. - Parsing all three would mean three parallel parsers in this module - and per-agent maintenance forever. Instead, extraction reads one - canonical shape (the one ``dsagt-proxy`` emits via autolog) and - users who want extraction run ``dsagt start --enable-proxy``. - - ``drain_session_traces`` queries proxy-shape traces in the session, - skips ones already tagged with ``dsagt.memory.extracted = "true"``, - formats each into the exchange shape the prompt expects, and tags - consumed traces so re-runs are idempotent. Stored in the - ``episodic_memory`` ChromaDB collection. Includes outlier - detection via per-category embedding centroids. +**Episodic memory** (MemoryExtractor): + A trace-pipeline *consumer* (``MemoryExtractor``) that consumes the + in-process ``Trace`` the heartbeat produces and writes per-block chunks into + the ``session_memory`` collection (producer/tool/turn_id metadata, no LLM). Files on disk (in project directory): explicit_memories.yaml — active user-confirmed facts explicit_memories_history.yaml — superseded/removed entries - centroids.json — per-category centroid vectors and counts - suggestions.json — queued outlier facts awaiting user review """ from __future__ import annotations import hashlib -import json import logging +import time from datetime import datetime, timezone from pathlib import Path - -import numpy as np +from typing import TYPE_CHECKING import yaml -from dsagt.knowledge import CollectionRoute, KnowledgeBase +from dsagt.knowledge import KnowledgeBase + +if TYPE_CHECKING: + from dsagt.traces import Trace logger = logging.getLogger(__name__) @@ -65,6 +43,7 @@ # Explicit memory (YAML store) # --------------------------------------------------------------------------- + def _now_iso() -> str: return datetime.now(timezone.utc).isoformat() @@ -97,7 +76,8 @@ def _save(self, entries: list[dict]) -> None: self._dir.mkdir(parents=True, exist_ok=True) self._path.write_text( yaml.dump(entries, default_flow_style=False, sort_keys=False) - if entries else "" + if entries + else "" ) def _append_history(self, entry: dict) -> None: @@ -199,249 +179,24 @@ def render_context(self) -> str: # Episodic memory: constants # --------------------------------------------------------------------------- -#: Project-local collection holding extracted-at-end-of-session facts, -#: insights, and summaries. Renamed from ``episodic_memory`` to match -#: the user-facing terminology in ``dsagt info`` ("session memory"). +#: Project-local collection holding the mechanically-indexed session turns. +#: Named ``session_memory`` to match the user-facing terminology in +#: ``dsagt info`` ("session memory"). SESSION_MEMORY_COLLECTION = "session_memory" -#: Backwards-compat alias. New code should use ``SESSION_MEMORY_COLLECTION``. -EPISODIC_COLLECTION = SESSION_MEMORY_COLLECTION - -# session_memory inherits the kb's default backend / vector_db — it -# used to hardcode ``embedding_backend="api"`` from when embedding was -# assumed to be API-only, but that forced ``kb_remember`` to retry -# against the (possibly invalid) embedding API even when the project -# was otherwise configured for local embeddings, hanging the agent -# for ~60s per call on the retry backoff. - -STOCK_CATEGORIES = { - "quality_control": "Assessment or filtering of data quality, QC metrics, thresholds, pass/fail rates", - "data_management": "File organization, data movement, format conversion, naming conventions", - "transformation": "Data processing steps, parameter choices, pipeline stage configuration", - "assembly": "Genome assembly, contig generation, scaffolding, assembly QC metrics", - "configuration": "Tool settings, environment setup, resource allocation decisions", - "performance": "Runtime, memory usage, throughput, resource consumption observations", - "tool_usage": "Tool selection rationale, parameter tuning, tool-specific behaviors or quirks", - "results": "Output summaries, key findings, deliverables produced", -} - -DEFAULT_SENSITIVITY = 0.35 - # --------------------------------------------------------------------------- -# Session-trace reading (MLflow) +# Trace chunking — split a turn's blocks into embeddable chunks +# +# The trace already carries mechanical boundaries (messages → content blocks), +# so a "chunk" is one block: a user/assistant text block or a tool_result, each +# embedded separately with a ``producer`` (user/llm/tool) label. tool_use blocks +# aren't embedded — they only supply the ``tool`` name resolved onto the matching +# tool_result. Over-long blocks split on newlines (cheap, boundary-meaningful). # --------------------------------------------------------------------------- -#: Trace tag we set after consuming a trace into an extraction run, so -#: re-runs of ``extract_session`` for the same session don't double-feed -#: the same exchange into the LLM. ``MlflowClient.set_trace_tag`` is the -#: idempotent equivalent of the old ``drain → unlink`` pattern. -DSAGT_MEMORY_PROCESSED_TAG = "dsagt.memory.extracted" - - -def _safe_parse_json(value): - """Parse value as JSON; return as-is if already a dict/list.""" - if isinstance(value, (dict, list)): - return value - if not isinstance(value, str) or not value: - return None - try: - return json.loads(value) - except (TypeError, ValueError): - return None - - -def _trace_to_exchange(row) -> dict | None: - """Format one MLflow ``search_traces`` row as an extraction-prompt exchange. - - The exchange shape mirrors what the proxy used to write to - ``session_log.jsonl`` so ``_render_conversation`` doesn't need to - change: - - {"timestamp": ..., "trace_id": ..., "model": ..., - "new_messages": [...], "response": [...content blocks...]} - - Returns None when the trace doesn't carry recognisable LLM-call - request/response shape (e.g. tool-execute spans, kb.* spans) so - callers can skip silently. - """ - request = _safe_parse_json(row.get("request")) - response = _safe_parse_json(row.get("response")) - - messages = request.get("messages") if isinstance(request, dict) else None - if not messages: - return None - - response_blocks: list[dict] = [] - if isinstance(response, dict): - # Anthropic shape: top-level ``content`` already a list of blocks. - if isinstance(response.get("content"), list): - response_blocks.extend(response["content"]) - # OpenAI shape: choices[].message.content (str or block list) plus - # tool_calls (translated to tool_use blocks for prompt consistency). - for choice in response.get("choices") or []: - msg = choice.get("message") or {} - content = msg.get("content") - if isinstance(content, str) and content: - response_blocks.append({"type": "text", "text": content}) - elif isinstance(content, list): - response_blocks.extend(content) - for tc in msg.get("tool_calls") or []: - func = tc.get("function") or {} - response_blocks.append({ - "type": "tool_use", - "id": tc.get("id"), - "name": func.get("name"), - "input": _safe_parse_json(func.get("arguments")) or {}, - }) - - return { - "timestamp": str(row.get("request_time") or ""), - "trace_id": row.get("trace_id") or "", - "model": (request.get("model") if isinstance(request, dict) else "") or "", - "new_messages": messages, - "response": response_blocks, - } - - -#: Service name extraction reads from. Restricting to ``dsagt-proxy`` -#: keeps the parser simple — every trace emitted by the proxy is -#: guaranteed LiteLLM-autolog shape (``mlflow.spanInputs`` / -#: ``mlflow.spanOutputs``), the same shape ``_trace_to_exchange`` parses. -#: Native agent traces (claude-code, goose) skip extraction by design; -#: see module docstring for why. -DSAGT_EXTRACTION_SOURCE_SERVICE_NAME = "dsagt-proxy" - - -def drain_session_traces( - project_name: str, session_id: str, mlflow_uri: str | None = None, -) -> list[dict]: - """Pull untagged proxy-shape session traces, format as exchanges, tag. - - Uses ``mlflow.search_traces`` to find traces in the project's - experiment whose ``mlflow.trace.session`` metadata matches - *session_id*, then filters to those emitted by ``dsagt-proxy`` - (the only canonical-shape source — see module docstring). Skips - ones already tagged with ``DSAGT_MEMORY_PROCESSED_TAG``; tags each - consumed trace so a subsequent extraction run on the same session - is a no-op. - - Returns a list of exchange dicts in chronological order. Empty - list when no experiment, no proxy-shape traces, or every trace was - already processed. Importantly, an empty result when the user ran - without ``--enable-proxy`` is *expected*: native-OTel agents - (Claude Code, Goose) emit traces in shapes this parser doesn't - handle, and the design choice is to require the proxy for - extraction rather than maintain N per-agent parsers. - """ - import os - import mlflow - from mlflow.tracking import MlflowClient - - uri = mlflow_uri or os.environ.get("MLFLOW_TRACKING_URI") - if not uri: - logger.warning("MLFLOW_TRACKING_URI not set; cannot drain session traces") - return [] - - mlflow.set_tracking_uri(uri) - client = MlflowClient(uri) - exp = client.get_experiment_by_name(project_name) - if exp is None: - return [] - - df = mlflow.search_traces( - locations=[exp.experiment_id], - filter_string=f"metadata.`mlflow.trace.session` = '{session_id}'", - max_results=10000, - order_by=["timestamp_ms ASC"], - ) - if df is None or df.empty: - return [] - - exchanges: list[dict] = [] - consumed_ids: list[str] = [] - for _, row in df.iterrows(): - tags = row.get("tags") or {} - if isinstance(tags, dict) and tags.get(DSAGT_MEMORY_PROCESSED_TAG) == "true": - continue - if not _trace_emitted_by(row, DSAGT_EXTRACTION_SOURCE_SERVICE_NAME): - continue - ex = _trace_to_exchange(row) - if ex: - exchanges.append(ex) - # Tag every proxy-shape trace we considered (parsed or not), so - # the next run doesn't re-inspect them. Non-proxy traces are - # left untagged on purpose — a future extraction run that - # broadens the source filter will pick them up. - trace_id = row.get("trace_id") - if trace_id: - consumed_ids.append(trace_id) - - for trace_id in consumed_ids: - try: - client.set_trace_tag(trace_id, DSAGT_MEMORY_PROCESSED_TAG, "true") - except Exception as e: - # Tag failure is non-fatal — worst case we re-extract that trace - # next session, which is annoying but not data-destroying. - logger.debug("set_trace_tag(%s) failed: %s", trace_id, e) - - return exchanges - - -def _trace_emitted_by(row, service_name: str) -> bool: - """Return True when *row*'s root span carries ``service.name == name``. - - MLflow's OTLP receiver flows the OTel resource attribute through to - each span's attributes, so every span in a trace from a given - process carries the same ``service.name``. We check the first span - we find — the shape varies (Span object / dict / pandas Series) - across MLflow versions, mirroring info.py:_source_from_spans. - """ - spans = row.get("spans") or [] - try: - for span in spans: - attrs = getattr(span, "attributes", None) - if attrs is None and isinstance(span, dict): - attrs = span.get("attributes") - if attrs and attrs.get("service.name") == service_name: - return True - except (TypeError, AttributeError): - pass - return False - - -# --------------------------------------------------------------------------- -# Extraction prompt construction -# --------------------------------------------------------------------------- - -def _render_conversation(exchanges: list[dict]) -> str: - lines = [] - for i, ex in enumerate(exchanges, 1): - lines.append(f"=== Exchange {i} ({ex.get('timestamp', '')}) ===") - - for msg in ex.get("new_messages", []): - role = msg.get("role", "unknown") - content = msg.get("content", "") - if isinstance(content, list): - parts = [] - for block in content: - if block.get("type") == "text": - parts.append(block.get("text", "")) - elif block.get("type") == "tool_result": - parts.append(f"[tool_result: {_extract_block_text(block)}]") - content = "\n".join(parts) - lines.append(f"[{role}] {content}") - - for block in ex.get("response", []): - if block.get("type") == "text": - lines.append(f"[assistant] {block.get('text', '')}") - elif block.get("type") == "tool_use": - name = block.get("name", "") - inp = json.dumps(block.get("input", {})) - lines.append(f"[assistant → tool_use: {name}({inp})]") - - lines.append("") - return "\n".join(lines) +#: Char budget above which a block is split on newlines before embedding. +_MAX_CHUNK_CHARS = 1200 def _extract_block_text(block: dict) -> str: @@ -450,385 +205,205 @@ def _extract_block_text(block: dict) -> str: return content if isinstance(content, list): return " ".join( - b.get("text", "") for b in content + b.get("text", "") + for b in content if isinstance(b, dict) and b.get("type") == "text" ) return str(content) -def _format_categories(categories: dict[str, str]) -> str: - lines = [] - for name, description in categories.items(): - lines.append(f"- {name}: {description}") - return "\n".join(lines) - - -def build_extraction_prompt( - exchanges: list[dict], - categories: dict[str, str] | None = None, -) -> str: - """Build the extraction prompt from session log exchanges.""" - cats = {**STOCK_CATEGORIES, **(categories or {})} - conversation = _render_conversation(exchanges) - category_list = _format_categories(cats) - - return f"""\ -You are a memory extraction system for a data pipeline building assistant. \ -Review the following session conversation to: - -1) Extract relevant facts — short declarative statements about what happened, \ -what was decided, what parameters were used, what succeeded or failed, \ -and any observations worth remembering for future sessions. - -2) Provide a complete summary of the session — the overall flow, what was \ -accomplished, key decisions made, and deliverables produced. - -3) Reflect across the facts and summary to synthesize insights — patterns, \ -lessons learned, generalizable observations that would help in future sessions \ -with similar data or tools. - -4) Classify each fact and insight into one of these categories: - -{category_list} - -Output ONLY a JSON object with this structure (no markdown, no preamble): - -{{ - "facts": [ - {{"text": "short declarative fact", "category": "category_name"}}, - ... - ], - "summary": "paragraph-length session summary", - "insights": [ - {{"text": "generalizable insight", "category": "category_name"}}, - ... - ] -}} - -Session conversation: - -{conversation}""" - +def _split_oversized(text: str) -> list[str]: + """Split an over-long block at meaningful boundaries — no tokenizer. -# --------------------------------------------------------------------------- -# Response parsing -# --------------------------------------------------------------------------- - -def parse_extraction_response(response_text: str) -> dict: - """Parse the LLM's JSON response into facts, summary, and insights.""" - text = response_text.strip() - if text.startswith("```"): - text = text.split("\n", 1)[1] if "\n" in text else text[3:] - if text.endswith("```"): - text = text[:-3] + Paragraph (``\\n\\n``) first, then line (``\\n``), greedily packing pieces up + to :data:`_MAX_CHUNK_CHARS`; a no-newline blob is hard-sliced as a last + resort. Short blocks pass through untouched (stripped). + """ text = text.strip() - - parsed = json.loads(text) - - return { - "facts": parsed.get("facts", []), - "summary": parsed.get("summary", ""), - "insights": parsed.get("insights", []), - } + if not text: + return [] + if len(text) <= _MAX_CHUNK_CHARS: + return [text] + sep = "\n\n" if "\n\n" in text else ("\n" if "\n" in text else None) + if sep is None: + return [ + text[i : i + _MAX_CHUNK_CHARS] + for i in range(0, len(text), _MAX_CHUNK_CHARS) + ] + chunks: list[str] = [] + cur = "" + for piece in (p.strip() for p in text.split(sep)): + if not piece: + continue + cur = f"{cur}{sep}{piece}" if cur else piece + while len(cur) > _MAX_CHUNK_CHARS: + chunks.append(cur[:_MAX_CHUNK_CHARS]) + cur = cur[_MAX_CHUNK_CHARS:].lstrip() + if cur: + chunks.append(cur) + return chunks -# --------------------------------------------------------------------------- -# LLM call -# --------------------------------------------------------------------------- +def _resolve_tool_names(exchanges: list[dict]) -> dict[str, str]: + """Map ``tool_use_id`` → tool name across a batch. -def call_extraction_llm( - prompt: str, - api_key: str, - model: str = "claude-sonnet-4-20250514", - base_url: str | None = None, - provider: str | None = None, -) -> str: - """Call the LLM with the extraction prompt. Returns raw response text. - - Uses LiteLLM so the call works against any OpenAI- or Anthropic-format - upstream the user's ``llm.base_url`` points at. Hand-rolling an - Anthropic-format request would 404 against an OpenAI-compatible gateway - (e.g. PNNL's ai-incubator-api), and vice versa. + A tool call and its result land in *different* turns (the result is the next + turn's input), so the lookup is built over the whole delivered batch, not a + single exchange. """ - import litellm - - # Mirror what the proxy does (commands/proxy_server.py): prefix the model - # with the configured provider so LiteLLM picks the right request format. - # Falls back to ``openai`` when provider is unset, preserving the historic - # behavior for callers that don't yet thread ``llm.provider`` through. - if base_url: - completion_model = f"{provider or 'openai'}/{model}" - else: - completion_model = model - - response = litellm.completion( - model=completion_model, - api_base=base_url, - api_key=api_key, - messages=[{"role": "user", "content": prompt}], - max_tokens=4096, - timeout=120.0, - ) - return response.choices[0].message.content or "" - - -# --------------------------------------------------------------------------- -# Outlier detection: category centroids -# --------------------------------------------------------------------------- - -class CategoryCentroids: - """Maintains running centroids per category.""" - - def __init__(self, path: Path): - self._path = path - self._data: dict[str, dict] = {} - if path.exists(): - self._data = json.loads(path.read_text()) - - def update(self, category: str, embedding: np.ndarray) -> float: - """Update centroid, return cosine distance before update.""" - vec = embedding.astype(np.float64) - - if category not in self._data: - self._data[category] = {"centroid": vec.tolist(), "count": 1} - return 0.0 - - entry = self._data[category] - old_centroid = np.array(entry["centroid"], dtype=np.float64) - count = entry["count"] - - distance = self._cosine_distance(vec, old_centroid) - - new_centroid = (old_centroid * count + vec) / (count + 1) - norm = np.linalg.norm(new_centroid) - if norm > 0: - new_centroid /= norm + names: dict[str, str] = {} + for ex in exchanges: + blocks = list(ex.get("response", [])) + for msg in ex.get("new_messages", []): + content = msg.get("content") + if isinstance(content, list): + blocks.extend(content) + for b in blocks: + if isinstance(b, dict) and b.get("type") == "tool_use": + names[b.get("id", "")] = b.get("name", "") + return names - entry["centroid"] = new_centroid.tolist() - entry["count"] = count + 1 - return float(distance) +def _turn_chunks(exchange: dict, tool_names: dict[str, str]) -> list[dict]: + """One turn's blocks → ``[{text, producer, tool}]`` chunks. - def save(self) -> None: - self._path.parent.mkdir(parents=True, exist_ok=True) - self._path.write_text(json.dumps(self._data, indent=2)) + text blocks carry the message's producer (``user``/``llm``); tool_result + blocks are ``tool`` with the name resolved via ``tool_use_id``; tool_use + blocks are skipped (captured in ``tool_names``, not embedded as prose). + """ + chunks: list[dict] = [] + + def emit(text: str, producer: str, tool_name: str) -> None: + for piece in _split_oversized(text or ""): + chunks.append({"text": piece, "producer": producer, "tool_name": tool_name}) + + def emit_block(block: dict, text_producer: str) -> None: + bt = block.get("type") + if bt == "text": + emit(block.get("text", ""), text_producer, "") + elif bt == "tool_result": + emit( + _extract_block_text(block), + "tool", + tool_names.get(block.get("tool_use_id", ""), ""), + ) - @staticmethod - def _cosine_distance(a: np.ndarray, b: np.ndarray) -> float: - dot = np.dot(a, b) - norm_a = np.linalg.norm(a) - norm_b = np.linalg.norm(b) - if norm_a == 0 or norm_b == 0: - return 1.0 - return float(1.0 - dot / (norm_a * norm_b)) + for msg in exchange.get("new_messages", []): + producer = "user" if msg.get("role", "user") == "user" else "llm" + content = msg.get("content", "") + if isinstance(content, str): + emit(content, producer, "") + else: + for block in content: + emit_block(block, producer) - @property - def categories(self) -> list[str]: - return list(self._data.keys()) + for block in exchange.get("response", []): + emit_block(block, "llm") - def count(self, category: str) -> int: - entry = self._data.get(category) - return entry["count"] if entry else 0 + return chunks # --------------------------------------------------------------------------- -# Outlier detection: suggestion queue +# Episodic memory: the trace-pipeline consumer # --------------------------------------------------------------------------- -class SuggestionQueue: - """Manages pending outlier suggestions on disk.""" - - def __init__(self, path: Path): - self._path = path - self._suggestions: list[dict] = [] - if path.exists(): - self._suggestions = json.loads(path.read_text()) - def add(self, text: str, category: str, distance: float, session_id: str = "") -> str: - suggestion_id = "sug_" + hashlib.sha256( - f"{text}:{category}:{session_id}".encode() - ).hexdigest()[:8] +def _epoch_or_now(ts: object) -> float: + """An exchange timestamp (epoch seconds) → float, else wall-clock now. - self._suggestions.append({ - "id": suggestion_id, - "text": text, - "category": category, - "distance": round(distance, 4), - "session_id": session_id, - "timestamp": datetime.now(timezone.utc).isoformat(), - }) - self._save() - return suggestion_id - - def dismiss(self, suggestion_id: str) -> bool: - before = len(self._suggestions) - self._suggestions = [s for s in self._suggestions if s["id"] != suggestion_id] - if len(self._suggestions) < before: - self._save() - return True - return False - - def get_all(self) -> list[dict]: - return list(self._suggestions) - - def clear(self) -> int: - count = len(self._suggestions) - self._suggestions = [] - self._save() - return count - - @property - def count(self) -> int: - return len(self._suggestions) - - def _save(self) -> None: - self._path.parent.mkdir(parents=True, exist_ok=True) - self._path.write_text(json.dumps(self._suggestions, indent=2)) + ``Trace.to_exchanges`` carries the span ``start_time`` (epoch + seconds) when the transcript recorded it, ``None`` otherwise — recency + weighting needs a number, so fall back to now for unstamped turns. + """ + return float(ts) if isinstance(ts, (int, float)) else time.time() -# --------------------------------------------------------------------------- -# Outlier detection: check and queue -# --------------------------------------------------------------------------- - -def check_and_queue_outliers( - texts: list[str], - categories: list[str], - embeddings: np.ndarray, - centroids: CategoryCentroids, - queue: SuggestionQueue, - threshold: float = DEFAULT_SENSITIVITY, - session_id: str = "", -) -> list[str]: - """Check facts against centroids, queue outliers. Returns suggestion IDs.""" - suggestion_ids = [] - - for text, category, embedding in zip(texts, categories, embeddings): - if not category: - continue +def episodic_consumers(config: dict, kb, runtime_dir, session_id) -> list: + """Build the episodic-memory consumer list from config (empty when off). - distance = centroids.update(category, embedding) + Shared by the live heartbeat (current session) and the startup catch-up + (previous session) so both feed the same :class:`MemoryExtractor` shape. + Episodic memory is a compute/storage opt-in (``episodic.enabled``). + Best-effort: a build failure returns ``[]`` rather than raising. + """ + epi = config.get("episodic", {}) or {} + if not epi.get("enabled"): + return [] + try: + return [ + MemoryExtractor( + kb, + runtime_dir=str(runtime_dir), + session_id=session_id or "", + ) + ] + except Exception as e: # noqa: BLE001 — memory is best-effort, never fatal + logger.warning("Could not build episodic-memory consumer: %s", e) + return [] - if centroids.count(category) <= 1: - continue - if distance > threshold: - sid = queue.add(text=text, category=category, distance=distance, session_id=session_id) - suggestion_ids.append(sid) - logger.info("Outlier flagged (distance=%.3f, threshold=%.3f): %s", distance, threshold, text[:80]) +class MemoryExtractor: + """Trace-pipeline consumer: ``Trace`` → ``session_memory`` chunks. - centroids.save() - return suggestion_ids + Plugged into :class:`~dsagt.traces.TraceCollector` alongside the MLflow sink; + each ``write`` receives the (subset of) just-completed turns and indexes them. + Idempotency is the heartbeat's job — it only delivers turns this consumer + hasn't acked — so ``write`` just does the work. + Chunks each turn per-block and embeds it — no LLM, nothing lost, the agent + is never blocked. + """ -# --------------------------------------------------------------------------- -# End-to-end extraction -# --------------------------------------------------------------------------- + #: Subscriber name → its own ack file (``.dsagt/trace_acks_memory.json``). + name = "memory" -def extract_session( - project_name: str, - kb: KnowledgeBase, - api_key: str, - model: str = "claude-sonnet-4-20250514", - base_url: str | None = None, - provider: str | None = None, - categories: dict[str, str] | None = None, - session_id: str | None = None, - runtime_dir: Path | None = None, - outlier_sensitivity: float = 0.0, - mlflow_uri: str | None = None, - exchanges: list[dict] | None = None, -) -> dict: - """Extract memories from a session's MLflow traces; store in episodic_memory. - - *project_name* is the MLflow experiment name (matches ``DSAGT_PROJECT``). - *session_id* selects which traces to consume; required when *exchanges* - is not supplied. Tests pass *exchanges* directly to bypass MLflow. - """ - if not session_id: - return {"status": "empty", "facts": 0, "insights": 0, - "reason": "no_session_id"} - - if exchanges is None: - exchanges = drain_session_traces(project_name, session_id, mlflow_uri) - if not exchanges: - return {"status": "empty", "facts": 0, "insights": 0} - - prompt = build_extraction_prompt(exchanges, categories) - response_text = call_extraction_llm(prompt, api_key, model, base_url, provider) - extracted = parse_extraction_response(response_text) - - timestamps = [ex.get("timestamp", "") for ex in exchanges if ex.get("timestamp")] - timestamp_start = min(timestamps) if timestamps else "" - timestamp_end = max(timestamps) if timestamps else "" - trace_ids = [ex.get("trace_id", "") for ex in exchanges if ex.get("trace_id")] - trace_refs = ",".join(trace_ids) if trace_ids else "" - - batch_meta = { - "session_id": session_id, - "timestamp_start": timestamp_start, - "timestamp_end": timestamp_end, - } - if trace_refs: - batch_meta["trace_refs"] = trace_refs - - fact_texts = [] - fact_metas = [] - fact_categories = [] - for fact in extracted["facts"]: - fact_texts.append(fact["text"]) - fact_categories.append(fact.get("category", "")) - fact_metas.append({**batch_meta, "source_type": "extraction", "category": fact.get("category", "")}) - - if extracted["summary"]: - fact_texts.append(extracted["summary"]) - fact_categories.append("results") - fact_metas.append({**batch_meta, "source_type": "summary", "category": "results"}) - - for insight in extracted["insights"]: - fact_texts.append(insight["text"]) - fact_categories.append(insight.get("category", "")) - fact_metas.append({**batch_meta, "source_type": "insight", "category": insight.get("category", "")}) - - stored = {"facts": 0, "insights": 0, "summary": 0, "suggestions": 0} - if fact_texts: - # Ask add_entries to hand back the freshly-computed embeddings so we - # don't pay a second embedding round-trip for outlier detection below. - # On API embedders this halves the wall time of memory extraction. - need_embeddings = outlier_sensitivity > 0 - add_result = kb.add_entries( - texts=fact_texts, - collection=EPISODIC_COLLECTION, - metadatas=fact_metas, - return_embeddings=need_embeddings, - ) - stored["facts"] = len(extracted["facts"]) - stored["insights"] = len(extracted["insights"]) - stored["summary"] = 1 if extracted["summary"] else 0 - - if need_embeddings: - if runtime_dir is None: - raise ValueError( - "extract_session: runtime_dir is required when " - "outlier_sensitivity > 0 (centroids/suggestions live there)" + def __init__( + self, + kb: KnowledgeBase, + *, + runtime_dir: str | Path, + session_id: str = "", + ): + self._kb = kb + self._runtime_dir = Path(runtime_dir) + self._session_id = session_id + + def write(self, trace: "Trace") -> None: + exchanges = trace.to_exchanges() + if not exchanges: + return + # Categorization root: this runs on the heartbeat, outside any MCP + # dispatch, so tag the whole extraction ``dsagt.source=episodic`` — the + # nested kb.* writes inherit it (otherwise they'd land uncategorized). + # ``episodic``, NOT ``memory``: this is per-turn internal embedding, and + # must filter apart from the user-facing memory tools (kb_remember / + # kb_get_memories) that carry ``dsagt.source=memory``. + # (This tags the *observability* span, not the memory chunks.) + from dsagt.observability import open_span + + with open_span("memory.extract", source="episodic"): + self._index_turns(exchanges) + + def _index_turns(self, exchanges: list[dict]) -> None: + """Chunk each turn per-block and embed into session_memory.""" + tool_names = _resolve_tool_names(exchanges) + texts, metas = [], [] + for ex in exchanges: + turn_id = ex.get("turn_id", "") + ts_epoch = _epoch_or_now(ex.get("timestamp")) + for chunk in _turn_chunks(ex, tool_names): + texts.append(chunk["text"]) + metas.append( + { + "session_id": self._session_id, + "source_type": "turn", + "turn_id": turn_id, + "producer": chunk["producer"], + "tool_name": chunk["tool_name"], + "ts_epoch": ts_epoch, + } ) - project_dir = runtime_dir - centroids_obj = CategoryCentroids(project_dir / "centroids.json") - queue = SuggestionQueue(project_dir / "suggestions.json") - - suggestion_ids = check_and_queue_outliers( - texts=fact_texts, - categories=fact_categories, - embeddings=add_result["embeddings"], - centroids=centroids_obj, - queue=queue, - threshold=outlier_sensitivity, - session_id=session_id, + if texts: + self._kb.add_entries( + texts=texts, collection=SESSION_MEMORY_COLLECTION, metadatas=metas ) - stored["suggestions"] = len(suggestion_ids) - - return { - "status": "ok", - **stored, - "total_entries": len(fact_texts), - "session_id": session_id, - } diff --git a/src/dsagt/observability.py b/src/dsagt/observability.py index 62405e3..8a8fa7b 100644 --- a/src/dsagt/observability.py +++ b/src/dsagt/observability.py @@ -1,50 +1,46 @@ """ -DSAgt observability — span emission via MLflow's native tracer provider. - -Business modules (knowledge.py, provenance.py, registry_server.py, run_tool.py) -import the small public surface defined here. ``init_tracing`` installs -MLflow's own ``TracerProvider`` as the OTel global, so every -``trace.get_tracer(...)`` call routes spans into MLflow's native trace store -with full ``mlflow.spanInputs`` / ``mlflow.spanOutputs`` integration and -direct access to ``InMemoryTraceManager`` for trace-metadata stamping. - -Public surface --------------- -init_tracing(service_name, *, mlflow_url=None, session_id=None) - Configure MLflow's tracer provider once per process. Reads - ``./dsagt_config.yaml`` for ``mlflow.port`` + ``project`` and - ``./.runtime`` for ``session_id``. Raises ``RuntimeError`` when no - MLflow URL is available — processes that legitimately run without - tracing (one-shot setup tools, tests) should not call this function. - -traced(span_name, *, capture=(), extract_return=None) - Decorator. Opens a span around a function call, captures named arguments - and (optionally) return-value-derived attributes, records exceptions, and - sets duration_ms. - -child_span(name, **attrs) - Context manager for nested spans inside a ``@traced`` method. - -obs - Process-wide proxy for the *current* span. ``obs.set("hits", 5)`` is a - no-op when no span is active, so business code can annotate without - branching on whether tracing is on. - -llm_source(source) / llm_call_context(source) - Marker decorator/context for LLM-call origin tagging. Stamps - ``dsagt.source`` (+ ``dsagt.agent``) on traces created inside the block - via MLflow's native ``tracing.context`` API, so the UI's session/source - filters surface every LLM call's origin (extraction, embedding, agent). +DSAgt observability — first-party span emission over the serverless MLflow store. + +DSAGT writes every trace to one ``sqlite:////mlflow.db`` store (no server). +TWO emission paths share that one store; they differ because MLflow's API forces +it, not by accident: + + * LIVE tracer (``mlflow.start_span``) — first-party DSAGT spans emitted *as the + MCP server / dsagt-run runs*. Uses MLflow's active-span context for + auto-nesting and the ``obs`` proxy. Each trace's root is tagged + ``dsagt.source`` with the MCP tool *category* the agent invoked + (memory / skill / knowledge / registry), or ``execution`` for dsagt-run — + set at the dispatch boundary, so the UI can filter this *debugging* view + apart from agent traces and bucket it by concern. + * REPLAY sink (``MLflowSink`` → ``mlflow.start_span_no_context``) — a finished + agent ``traces.Trace`` backfilled after the fact with the transcript's + original timestamps. ``start_span`` cannot backdate (it has no + ``start_time_ns`` param), so replay *must* use ``start_span_no_context``; + live *should* use ``start_span`` (no_context establishes no active span, + which would kill the ``obs`` proxy and auto-nesting). Hence two paths, one + store — neither is a historical leftover. + +Layout (top → bottom) +--------------------- + setup find_project_config · resolve_tracking_uri · init_tracing + live tracer open_span ─┬─ traced (decorate a function) + ├─ child_span (open a sub-span) + └─ obs (annotate the span you're inside) + tagging: open_span(source=…) → _attach_trace_metadata + (dsagt.source set on the trace's root only) + factories: kb_* · registry_* · code_execute_span + replay sink MLflowSink (Trace → backdated spans; a traces.TraceCollector consumer) + +``traced`` and ``child_span`` *open* a span; ``obs`` *annotates* whichever span +is currently open. All three no-op when tracing was never initialized, so +business code never imports MLflow and never branches on whether tracing is on. """ from __future__ import annotations -import asyncio -import atexit import functools import inspect import logging -import os import time from contextlib import contextmanager from pathlib import Path @@ -52,123 +48,22 @@ logger = logging.getLogger(__name__) -# Module-level state. _initialized guards against double-init in test runs -# and in subprocesses where init_tracing might be called more than once. -# -# NOTE on _default_session_id: -# This is intentionally a process-global rather than threaded through -# every span helper. DSAgt runs one process per project, so the session -# id is a process-wide constant set at startup by init_tracing(), and -# propagating it implicitly via this module-level value lets business -# code in knowledge.py / provenance.py / registry_server.py emit spans -# without ever knowing about session ids. The cost is that tests have -# to monkeypatch the global to isolate (see _reset_tracing fixture in -# test_observability.py). -# -# If DSAgt ever grows a multi-tenant mode (one server process serving -# many projects), the right replacement is OTel baggage: -# from opentelemetry import baggage -# sid = baggage.get_baggage("session.id") -# which gives per-request context propagation that works correctly -# across concurrent requests. Until then, the global is simpler. +# Module-level state. _initialized guards every span helper: without a tracking +# store + experiment set (init_tracing), mlflow.start_span would write to +# MLflow's default location, so the helpers no-op until init_tracing succeeds. +# DSAGT runs one process per project, so the session id is a process-wide +# constant set once at startup and tagged onto each internal trace for grouping. _initialized = False -_tracer_provider = None _default_session_id: str | None = None -# Strategy pointers bound at init_tracing time. Both exist to keep -# backend-specific behavior out of call sites: -# -# _metadata_stamper(dict) -# Write key/value metadata to the currently-active trace via MLflow's -# InMemoryTraceManager (so it lands in trace_metadata, queryable in -# the UI). -# -# _llm_context_factory(dict) → context manager -# Tag every trace created inside the returned context. Uses -# mlflow.tracing.context which stamps at trace-creation time — the only -# reliable window for LLM traces emitted via mlflow.litellm.autolog. -_metadata_stamper: "Callable[[dict], None] | None" = None -_llm_context_factory: "Callable[[dict], Any] | None" = None - - -def init_tracing( - service_name: str, - mlflow_url: str | None = None, - session_id: str | None = None, -) -> None: - """Install MLflow's tracer provider as the OTel global. - - Every ``trace.get_tracer(...)`` call (from ``@traced`` / ``child_span`` - in MCP servers, dsagt-run, and the proxy) routes spans into MLflow's - native trace store. This gives ``mlflow.spanInputs`` / ``mlflow.spanOutputs`` - full integration and direct access to ``InMemoryTraceManager`` for - trace-metadata stamping (session id, source, agent). - - Single source of truth: reads ``./dsagt_config.yaml`` for - ``mlflow.port`` + ``project``; reads ``./.runtime`` for ``session_id``. - No env-var fallback, no parent-walk — every dsagt service runs with - cwd == project_dir by contract. - - The ``mlflow_url`` and ``session_id`` keyword args are kept for tests - and proxy mode, where the caller plants known values directly. - - Raises ``RuntimeError`` when no MLflow URL is available. Processes - that legitimately run without tracing (one-shot setup tools, tests) - should not call this function; tests install a test provider directly. - """ - global _initialized, _default_session_id - global _metadata_stamper, _llm_context_factory - - if _initialized: - if session_id: - _default_session_id = session_id - return - cfg_pdir, cfg = find_project_config() - _default_session_id = session_id or _read_session_id_from_runtime(cfg_pdir) - mlflow_url = mlflow_url or _mlflow_url_from_config(cfg) - - if not mlflow_url: - raise RuntimeError( - f"{service_name}: no observability backend configured. " - f"Expected dsagt_config.yaml with mlflow.port set in cwd or a " - f"parent directory. Processes that legitimately run without " - f"tracing should not call init_tracing at all." - ) - project_name = (cfg or {}).get("project") - if not project_name: - raise RuntimeError( - "no 'project' in dsagt_config.yaml — this should never happen " - "for a project created via `dsagt init`." - ) - - _install_mlflow_provider(mlflow_url, project_name) - _metadata_stamper = _stamp_metadata_mlflow - _llm_context_factory = _mlflow_llm_context - - # Autolog every LiteLLM completion/embedding call from this process into - # MLflow. With MLflow's tracer provider installed, autolog stamps full - # request/response into ``mlflow.spanInputs`` / ``mlflow.spanOutputs`` - # natively. This feeds memory extraction (it queries MLflow for - # session traces) and ``dsagt info``. ``dsagt-run`` doesn't make - # litellm calls, so skip there. - if service_name != "dsagt-run": - import mlflow - - mlflow.litellm.autolog() - _initialized = True - atexit.register(_shutdown) - logger.info( - "init_tracing: service=%s mlflow=%s project=%s session=%s", - service_name, - mlflow_url, - project_name, - _default_session_id or "", - ) +# =========================================================================== +# Setup — where am I, where do I write, wire MLflow up +# =========================================================================== def find_project_config() -> tuple[Path | None, dict | None]: - """Read ``./dsagt_config.yaml`` from cwd. + """Read ``./.dsagt/config.yaml`` from cwd. Returns ``(cwd, parsed_config)`` or ``(None, None)`` if cwd isn't a project directory. No walking — services that need this info run @@ -181,7 +76,7 @@ def find_project_config() -> tuple[Path | None, dict | None]: helper is only for services running inside the project. """ cwd = Path.cwd().resolve() - candidate = cwd / "dsagt_config.yaml" + candidate = cwd / ".dsagt" / "config.yaml" if not candidate.exists(): return None, None try: @@ -193,299 +88,199 @@ def find_project_config() -> tuple[Path | None, dict | None]: return cwd, None -def _mlflow_url_from_config(cfg: dict | None) -> str | None: - """Return ``http://localhost:`` from a parsed dsagt_config.yaml.""" - if not cfg: - return None - port = (cfg.get("mlflow") or {}).get("port") - return f"http://localhost:{port}" if port else None +def resolve_tracking_uri(config: dict | None) -> str: + """Compute the serverless MLflow tracking URI for DSAGT self-logging. + Always ``sqlite:////mlflow.db`` — DSAGT knows where it writes; + there is no server to point at. ``project_dir`` comes from the resolved + config (injected by ``session.load_config``), falling back to cwd for + in-project callers. -def _read_session_id_from_runtime(project_dir: Path | None) -> str | None: - """Read the active session id from ``/.runtime``. - - ``dsagt mlflow`` writes ``session_id`` here at startup; every service - running for the duration of that MLflow daemon shares one session id. + The MLflow client honors a ``sqlite:`` URI directly (auto-creating + + migrating the DB on first use), so self-logging needs no listener and this + never has to fail. SQLite is MLflow's supported serverless backend — the + filesystem store (``file:`` / ``./mlruns``) is deprecated as of Feb 2026. + DSAGT emits only traces (no runs/models), so the experiment's default + artifact dir is never materialized. """ - if project_dir is None: - return None - runtime_file = project_dir / ".runtime" - if not runtime_file.exists(): - return None - try: - import json as _json + cfg = config or {} + pdir = cfg.get("project_dir") + base = Path(pdir).resolve() if pdir else Path.cwd().resolve() + return f"sqlite:///{base / 'mlflow.db'}" - return _json.loads(runtime_file.read_text()).get("session_id") - except Exception as e: - logger.debug("could not read session_id from %s: %s", runtime_file, e) - return None - - -def _install_mlflow_provider(mlflow_url: str, project_name: str) -> None: - """Wire MLflow's tracer provider in as the OTel global. - - Force MLflow's lazy provider to initialize via its private init hook - so we can hand the resulting TracerProvider to OTel below. The - underscore on the init function is MLflow-internal — pinning the - mlflow version range in pyproject.toml keeps that boundary stable. - """ - import mlflow - from mlflow.tracing import provider as mp - from opentelemetry import trace - from opentelemetry.util._once import Once - mlflow.set_tracking_uri(mlflow_url) - mlflow.set_experiment(project_name) - - mp._initialize_tracer_provider() - - # OTel guards set_tracer_provider with a one-shot Once flag — the first - # caller wins. Reset so installation always takes effect even if - # something accessed get_tracer earlier and locked in the no-op global. - trace._TRACER_PROVIDER = None # type: ignore[attr-defined] - trace._TRACER_PROVIDER_SET_ONCE = Once() # type: ignore[attr-defined] - trace.set_tracer_provider(mp.provider.get()) +def init_tracing( + service_name: str, + mlflow_url: str | None = None, + session_id: str | None = None, +) -> None: + """Point MLflow at the project's serverless store + experiment. + After this, every ``mlflow.start_span`` (from ``@traced`` / ``child_span`` + in the MCP server and dsagt-run) lands in ``sqlite:////mlflow.db``. + The tracking URI is the serverless ``sqlite:////mlflow.db`` computed by + :func:`resolve_tracking_uri`. The session id, passed by the MCP server at + startup, tags internal traces for grouping. -def _stamp_metadata_on_trace(request_id: str, metadata: dict) -> None: - """Write metadata to a specific MLflow trace by id. + The ``mlflow_url`` and ``session_id`` keyword args are kept for tests, where + the caller plants known values directly. - Used when the caller has the trace_id in hand (e.g. inside the proxy's - ``_DSAGTMlflowLogger`` right after ``start_trace``). ``stamp_metadata`` - is the higher-level version that looks up the current trace via the - active OTel span. + Never raises. When cwd isn't a dsagt project dir, logs and no-ops so + one-shot tools / tests outside a project simply run untraced. """ - try: - from mlflow.tracing.trace_manager import InMemoryTraceManager - - with InMemoryTraceManager.get_instance().get_trace(request_id) as t: - if t is not None: - t.info.trace_metadata.update({k: str(v) for k, v in metadata.items()}) - except Exception as e: - logger.debug("metadata stamp failed for %s: %s", request_id, e) - - -def _stamp_metadata_mlflow(metadata: dict) -> None: - """Write metadata to the currently-active trace's trace_metadata.""" - from opentelemetry import trace + global _initialized, _default_session_id - span = trace.get_current_span() - ctx = span.get_span_context() - if not ctx.is_valid: + if _initialized: + if session_id: + _default_session_id = session_id return - _stamp_metadata_on_trace(f"tr-{ctx.trace_id:032x}", metadata) + cfg_pdir, cfg = find_project_config() + project_name = (cfg or {}).get("project") + if not project_name: + logger.warning( + "%s: no .dsagt/config.yaml with a 'project' in cwd (%s) — tracing " + "disabled for this process.", + service_name, + Path.cwd(), + ) + return -def stamp_metadata(metadata: dict) -> None: - """Stamp arbitrary key/value metadata on the currently-active trace. + _default_session_id = session_id + if mlflow_url is None: + resolve_cfg = dict(cfg) + resolve_cfg["project_dir"] = str(cfg_pdir) + mlflow_url = resolve_tracking_uri(resolve_cfg) - No-op when no backend is configured (tests, standalone tools). - """ - if _metadata_stamper is not None and metadata: - _metadata_stamper(metadata) + import mlflow + mlflow.set_tracking_uri(mlflow_url) + mlflow.set_experiment(project_name) + _initialized = True + logger.info( + "init_tracing: service=%s mlflow=%s project=%s session=%s", + service_name, + mlflow_url, + project_name, + _default_session_id or "", + ) -def _mlflow_llm_context(metadata: dict): - """MLflow native tracing.context — stamps at trace-creation time. - Used by ``llm_call_context`` so every LLM trace emitted via - ``mlflow.litellm.autolog`` inside the block lands in MLflow with - ``dsagt.source`` / ``dsagt.agent`` / ``mlflow.trace.session`` already - set on its trace_metadata. - """ - import mlflow +# =========================================================================== +# Live tracer — first-party debug spans, emitted as code runs +# =========================================================================== - return mlflow.tracing.context(metadata=metadata) +# ----- the span primitive ----- @contextmanager -def llm_call_context(source: str): - """Stamp ``dsagt.source`` (+ ``dsagt.agent``, ``mlflow.trace.session``) - on traces created inside this block. - """ - metadata = {"dsagt.source": source} - if agent := os.environ.get("DSAGT_AGENT"): - metadata["dsagt.agent"] = agent - if _default_session_id: - metadata["mlflow.trace.session"] = _default_session_id - if _llm_context_factory is not None: - with _llm_context_factory(metadata): - yield - else: - yield - - -def llm_source(source: str): - """Decorator form of ``llm_call_context`` for tidy call sites. - - Handles both sync and async. Every LLM call made inside the decorated - function lands in MLflow with ``dsagt.source = `` metadata, - letting the UI distinguish extraction / embedding / agent-turn origins - at a glance. +def open_span(name: str, span_type: str | None = None, source: str | None = None): + """Open a span on the serverless MLflow store. + + Single tracing code path — ``mlflow.start_span`` auto-nests under whatever + span is already active (its own context model), so child spans Just Work. + Yields ``None`` when tracing was never initialized (one-shot tools / tests + outside a project), so the helpers below degrade to no-ops. + + ``source`` is set only on the *categorization root* of a trace — the MCP + dispatch span and ``tool.execute`` — and tags the whole trace + ``dsagt.source`` for the debug-view filter (see :func:`_attach_trace_metadata`). + Inner spans pass ``source=None`` and inherit the root's tag. """ + if not _initialized: + yield None + return + import mlflow + from mlflow.entities import SpanType - def dec(fn): - if asyncio.iscoroutinefunction(fn): - - @functools.wraps(fn) - async def aw(*a, **kw): - with llm_call_context(source): - return await fn(*a, **kw) - - return aw - - @functools.wraps(fn) - def sw(*a, **kw): - with llm_call_context(source): - return fn(*a, **kw) + with mlflow.start_span(name=name, span_type=span_type or SpanType.UNKNOWN) as span: + _attach_trace_metadata(source) + yield span - return sw - return dec +# ----- attribute-value helpers ----- -def extract_cache_stats(usage: dict) -> tuple[int, int]: - """Extract (cache_read, cache_write) tokens from a LiteLLM usage dict. +def _coerce_attr(value: Any) -> Any: + if isinstance(value, (str, bool, int, float)): + return value + if isinstance(value, (list, tuple)): + return [_coerce_attr(v) for v in value] + return str(value) - LiteLLM's ``Usage`` object doesn't backfill cache fields across - providers; whichever shape the upstream returned is the one populated. - The configured ``LLM_PROVIDER`` (passed to ``dsagt-proxy --provider X`` - and used as the LiteLLM provider prefix in the model_list) decides - which response format we get back. The four field-name conventions - we've seen in practice are all checked here so this function works - regardless of provider config: - Cache read (tokens served from cache at the lower rate): - - ``cache_read_input_tokens`` — Anthropic, Bedrock-Claude - - ``prompt_tokens_details.cached_tokens`` — OpenAI, Azure-OpenAI - - ``cached_content_token_count`` — Gemini - - ``prompt_cache_hit_tokens`` — DeepSeek +def truncate(value: str, limit: int = 256) -> str: + """Truncate a string for span attributes. - Cache write (only Anthropic-family bills cache writes separately at - the 1.25× premium; OpenAI/Gemini/DeepSeek auto-cache without a - distinct write fee, so the field doesn't exist): - - ``cache_creation_input_tokens`` — Anthropic, Bedrock-Claude + Span backends (and the MLflow trace UI) handle short attribute values + much better than multi-megabyte stdout/stderr blobs. The full payload + lives in ``trace_archive/.json``; the span just carries a + head/tail summary so a human glancing at the UI can tell what happened. """ - if not isinstance(usage, dict): - return (0, 0) - - details = usage.get("prompt_tokens_details") or {} - if not isinstance(details, dict): - details = {} - - read = ( - usage.get("cache_read_input_tokens") - or details.get("cached_tokens") - or usage.get("cached_content_token_count") - or usage.get("prompt_cache_hit_tokens") - or 0 - ) - write = usage.get("cache_creation_input_tokens") or 0 - return (int(read), int(write)) - + if value is None: + return "" + if len(value) <= limit: + return value + head = limit - 32 + return value[:head] + f"... [+{len(value) - head} chars]" -def _shutdown() -> None: - """Flush and shut down the tracer provider. Registered via atexit.""" - global _tracer_provider - if _tracer_provider is None: - return - try: - _tracer_provider.shutdown() - except Exception as e: - logger.debug("tracer shutdown raised: %s", e) - _tracer_provider = None +# ----- trace tagging — the dsagt.source debug filter + session grouping ----- +# +# dsagt.source names the MCP tool *category* that was invoked — one of +# {memory, skill, knowledge, registry} (the four concern modules of the merged +# dsagt-server), plus ``execution`` for dsagt-run's out-of-process data-tool +# runs. It is assigned at the *entry point*, not derived from the span name: +# the MCP dispatch shell knows which concern owns each tool and stamps the +# category on the trace's root span, so e.g. ``search_skills`` calling into +# ``kb.search`` is tagged ``skill`` (the tool the agent called), not +# ``knowledge`` (the subsystem that happened to do the work). Inner spans +# inherit the root's tag; agent traces carry no ``dsagt.source`` at all, which +# is what lets the MLflow UI filter the debug view in or out. -def get_tracer(name: str): - """Return an OTel tracer routed through the OTLP provider installed by - ``init_tracing``.""" - from opentelemetry import trace - return trace.get_tracer(name) +def _attach_trace_metadata(source: str | None) -> None: + """Tag the current trace as a DSAGT-internal (debug) trace. + Called from :func:`open_span` on a categorization root (``source`` set): -@contextmanager -def open_span(name: str): - """Open a span on the installed tracer provider. + - ``dsagt.source`` tag: the MCP tool category / ``execution`` — powers the + UI's debug-view filter. + - ``mlflow.trace.session`` metadata: groups this process's internal traces + under one session (reserved MLflow key, drives the native session filter). - Single OTel code path — provider selection happens once at - ``init_tracing`` time. Yields ``None`` if tracing was never initialized - (test paths that skip init_tracing monkeypatch module state directly). + No-op for inner spans (``source is None``) — they inherit the root's tag. """ - if not _initialized: - yield None + if not source: return - tracer = get_tracer("dsagt.observability") - with tracer.start_as_current_span(name) as span: - yield span + metadata = ( + {"mlflow.trace.session": _default_session_id} if _default_session_id else None + ) + import mlflow + mlflow.update_current_trace(tags={"dsagt.source": source}, metadata=metadata) -def configure_litellm_retries( - num_retries: int = 5, - request_timeout: float = 300.0, -) -> None: - """Configure LiteLLM module-level retry / timeout knobs and quiet its - stdout chatter. - - These are independent of tracing — even ``dsagt-setup-kb`` running before - any project exists (and therefore with no MLflow endpoint) gets the - rate-limit-resilient embedding path. - - LiteLLM's ``num_retries`` retries on transient failures including 429s - with exponential backoff. We also retry inside APIEmbeddingClient with - finer control, but the litellm-level setting acts as a backstop for - other code paths (e.g. the proxy). - - LiteLLM is also chatty by default: every error prints a "Give Feedback / - Get Help" footer to stdout, and every successful call duplicates a log - line through the ``LiteLLM`` logger. Both make ``dsagt-setup-kb`` output - nearly unreadable during long ingests with retries. We suppress them - here so DSAgt's own progress logs are visible. - """ - # litellm[proxy] is a hard dependency in pyproject.toml — failing this - # import means a broken install, not "tracing optional". - import litellm - - litellm.num_retries = num_retries - litellm.request_timeout = request_timeout - - # Kill the "Give Feedback / Get Help" + "Provider List" stdout footers - # that LiteLLM prints on every exception. These are aimed at first-time - # users; for a long-running embed job they make the output unreadable. - litellm.suppress_debug_info = True - - # Stop the duplicate INFO log lines. LiteLLM emits its own - # "Wrapper: Completed Call" line on every embedding call, and propagation - # to the root logger doubles each one (once via LiteLLM's namespace, once - # via root). Mute the namespace and stop propagation. - _llm_log = logging.getLogger("LiteLLM") - _llm_log.setLevel(logging.WARNING) - _llm_log.propagate = False - - logger.debug( - "configure_litellm_retries: num_retries=%d timeout=%.0fs", - num_retries, - request_timeout, - ) +# ----- annotate the active span ----- -# --------------------------------------------------------------------------- -# obs — process-wide proxy for the current span -# --------------------------------------------------------------------------- +class _ActiveSpanProxy: + """The *annotate* verb that complements the *open* verbs (traced/child_span). -class _Obs: - """No-op-safe proxy for the active span. + Where ``traced`` / ``child_span`` open a span, this annotates whichever span + is currently open — ``obs.set("hits", 5)`` from inside a ``@traced`` body + attaches to that body's span. It exists so business code (knowledge.py, + provenance.py, registry_tools.py) never imports MLflow and never branches on + whether tracing is on: when no span is active (tracing disabled, or call + site outside any traced block) every method silently does nothing. - Business code uses ``obs.set("hits", 5)`` instead of importing OTel. When - no span is active (tracing disabled, or call site outside any traced - block), every method silently does nothing. + The process-wide singleton is exported as ``obs``. """ def set(self, key: str, value: Any) -> None: span = self._current() if span is not None and value is not None: - span.set_attribute(key, value) + span.set_attribute(key, _coerce_attr(value)) def set_many(self, attrs: Mapping[str, Any]) -> None: span = self._current() @@ -493,53 +288,45 @@ def set_many(self, attrs: Mapping[str, Any]) -> None: return for k, v in attrs.items(): if v is not None: - span.set_attribute(k, v) + span.set_attribute(k, _coerce_attr(v)) def event(self, name: str, **attrs: Any) -> None: span = self._current() if span is None: return + from mlflow.entities import SpanEvent + clean_attrs = {k: v for k, v in attrs.items() if v is not None} - span.add_event(name, attributes=clean_attrs) + span.add_event(SpanEvent(name, attributes=clean_attrs)) def set_inputs(self, inputs: Any) -> None: - """Populate the trace's ``request`` field for the MLflow trace UI. - - Stamps ``mlflow.spanInputs`` as a JSON-serialized OTel attribute — - this is the same key MLflow's ``LiveSpan.set_inputs`` writes, so it - flows through to MLflow's ``request`` column in ``search_traces``. - For non-MLflow OTel backends (Jaeger, Tempo) it just shows up as a - regular span attribute, harmlessly. - """ + """Populate the trace's ``request`` field for the MLflow trace UI.""" span = self._current() if span is None or inputs is None: return - span.set_attribute("mlflow.spanInputs", _to_json(inputs)) + span.set_inputs(inputs) def set_outputs(self, outputs: Any) -> None: """Populate the trace's ``response`` field for the MLflow trace UI.""" span = self._current() if span is None or outputs is None: return - span.set_attribute("mlflow.spanOutputs", _to_json(outputs)) + span.set_outputs(outputs) @staticmethod def _current(): - """Return the currently-active OTel span, or ``None`` if none.""" - from opentelemetry import trace - - span = trace.get_current_span() - if span is None or not span.get_span_context().is_valid: + """Return the currently-active MLflow span, or ``None`` if none.""" + if not _initialized: return None - return span + import mlflow + return mlflow.get_current_active_span() -obs = _Obs() +obs = _ActiveSpanProxy() -# --------------------------------------------------------------------------- -# traced — decorator for top-level method instrumentation -# --------------------------------------------------------------------------- + +# ----- open a span — decorator + context manager ----- def traced( @@ -548,7 +335,7 @@ def traced( capture: Iterable[str] = (), extract_return: Mapping[str, Callable[[Any], Any]] | None = None, ) -> Callable: - """Wrap a function in an OTel span. + """Wrap a function in an MLflow span. Parameters ---------- @@ -560,16 +347,15 @@ def traced( work. extract_return Optional mapping of attribute name → function applied to the return - value to extract that attribute. Use this for things like - ``{"hits": lambda r: len(r)}``. + value to extract that attribute (e.g. ``{"hits": lambda r: len(r)}``). Behavior -------- * Captures the configured args as attributes (skipping ``None``). * Always sets ``duration_ms``. - * Always sets the DSAgt session id (if init_tracing was given one). - * Records exceptions via ``span.record_exception`` and sets ERROR status. - * Re-raises after recording. + * Tags the trace ``dsagt.source`` + session for the debug view. + * Records exceptions and sets ERROR status (MLflow auto-records the + exception on context exit); re-raises after recording. """ capture = tuple(capture) extract_return = dict(extract_return or {}) @@ -582,7 +368,6 @@ def wrapper(*args, **kwargs): with open_span(span_name) as span: if span is None: return fn(*args, **kwargs) - _attach_trace_metadata(span_name) # Capture configured arguments as span attributes. Looked up # by name against the function signature so positional and @@ -605,9 +390,8 @@ def wrapper(*args, **kwargs): start = time.perf_counter() try: result = fn(*args, **kwargs) - except Exception as exc: - span.record_exception(exc) - _set_error_status(span, str(exc)) + except Exception: + span.set_status("ERROR") raise finally: span.set_attribute( @@ -641,113 +425,31 @@ def wrapper(*args, **kwargs): return decorator -_SOURCE_BY_PREFIX = ( - ("tool.", "tool"), - ("kb.", "knowledge"), - ("registry.", "registry"), -) - - -def _derive_source(span_name: str | None) -> str | None: - """Map a span name to a ``dsagt.source`` value. - - Prefix-based dispatch so every span we open lands with a source tag - without touching the call site. LLM traces get their source from the - proxy's ``_DSAGTMlflowLogger`` (agent) or ``@llm_source`` (extraction, - embedding) instead — this helper covers the non-LLM spans our own - instrumentation emits. - """ - if not span_name: - return None - for prefix, source in _SOURCE_BY_PREFIX: - if span_name.startswith(prefix): - return source - return None - - -def _attach_trace_metadata(span_name: str | None) -> None: - """Stamp session + source on the currently-active trace. - - - ``mlflow.trace.session``: process-wide session id (reserved MLflow - key; powers the UI's native session filter). - - ``dsagt.source``: derived from the span name prefix — so every span - we open lands in ``dsagt info``'s "by source" bucket. - - No-op when no backend is configured or no span is active; the - metadata stamper handles both cases. - """ - md: dict[str, str] = {} - if _default_session_id: - md["mlflow.trace.session"] = _default_session_id - src = _derive_source(span_name) - if src: - md["dsagt.source"] = src - if md: - stamp_metadata(md) - - -def _to_json(value: Any) -> str: - """Serialize *value* to JSON for an OTel attribute. - - OTel attributes only accept primitives + sequences of primitives; - MLflow's ``mlflow.spanInputs`` / ``mlflow.spanOutputs`` keys expect - a JSON-encoded payload that the trace UI deserializes for display. - """ - import json - - try: - return json.dumps(value, default=str) - except (TypeError, ValueError): - return str(value) - - -def _coerce_attr(value: Any) -> Any: - if isinstance(value, (str, bool, int, float)): - return value - if isinstance(value, (list, tuple)): - return [_coerce_attr(v) for v in value] - return str(value) - - -def _set_error_status(span, message: str) -> None: - # OTel is mandatory; if Status/StatusCode disappears or moves, we want - # the resulting ImportError to surface immediately so the dep upgrade - # gets caught at the test level rather than silently disabling error - # status on every traced exception path. - from opentelemetry.trace import Status, StatusCode - - span.set_status(Status(StatusCode.ERROR, message)) - - -# --------------------------------------------------------------------------- -# Typed span helpers — populated incrementally as later stages need them. -# -# The convention is: every span name DSAgt emits has a helper here. Business -# modules call the helper, never tracer.start_as_current_span directly. This -# keeps span names, attribute schemas, and required fields in one file. -# --------------------------------------------------------------------------- - - @contextmanager -def child_span(name: str, **attrs: Any): +def child_span(name: str, *, span_type: str | None = None, **attrs: Any): """Open a child span with arbitrary attributes. - Use this from inside a ``@traced`` method when you need to break a method - into sub-phases (e.g. embed / index_search / rerank inside kb.search). - Prefer the typed helpers below when one exists for your operation. + Use this from inside a ``@traced`` method to break a method into sub-phases + (e.g. embed / index_search / rerank inside kb.search). Prefer the typed + factories below when one exists for your operation. """ - with open_span(name) as span: + with open_span(name, span_type=span_type) as span: if span is None: yield None return - _attach_trace_metadata(name) for k, v in attrs.items(): if v is not None: span.set_attribute(k, _coerce_attr(v)) yield span -# ----- Knowledge base spans (Stage 1) ----- +# ----- named span factories ----- +# +# Every span name DSAgt emits has a factory here. Business modules call the +# factory, never mlflow.start_span directly, so span names, attribute schemas, +# and span types stay in one file. + +# Knowledge base spans. def kb_embed_span(backend: str | None, model: str | None, n_texts: int): @@ -755,10 +457,13 @@ def kb_embed_span(backend: str | None, model: str | None, n_texts: int): Used for both query embedding (kb.search) and chunk embedding (kb.ingest, kb.append, kb.add_entries). Backend-agnostic: ``backend`` is ``"api"`` - for LiteLLM/HTTP embedders or ``"local"`` for sentence-transformers. + for the HTTP embedder or ``"local"`` for sentence-transformers. """ + from mlflow.entities import SpanType + return child_span( "kb.embed", + span_type=SpanType.EMBEDDING, backend=backend, model=model, n_texts=n_texts, @@ -767,8 +472,11 @@ def kb_embed_span(backend: str | None, model: str | None, n_texts: int): def kb_index_search_span(vector_db: str | None, k: int, filtered: bool): """Span around an underlying vector index search call.""" + from mlflow.entities import SpanType + return child_span( "kb.index_search", + span_type=SpanType.RETRIEVER, vector_db=vector_db, k=k, filtered=filtered, @@ -784,23 +492,28 @@ def kb_rerank_span(model: str | None, n_pairs: int): ) -# ----- Registry server spans (Stage 4) ----- -# -# Only the deliberate, infrequent registry operations are instrumented. -# search_registry / search_skills are intentionally NOT instrumented — they +# Registry spans. Only the deliberate, infrequent registry operations are +# instrumented. search_registry / search_skills are intentionally NOT — they # are high-frequency low-information per call, and the agent-side LLM trace # already records that they were invoked. -def registry_save_tool_span(tool_name: str | None): - """Span around ``save_tool_spec``.""" - return child_span("registry.save_tool_spec", tool_name=tool_name) +def registry_save_code_span(code_name: str | None): + """Span around ``save_code_spec``.""" + from mlflow.entities import SpanType + + return child_span( + "registry.save_code_spec", span_type=SpanType.TOOL, code_name=code_name + ) def registry_install_deps_span(packages: list[str] | None): """Span around an ``install_dependencies`` call.""" + from mlflow.entities import SpanType + return child_span( "registry.install_dependencies", + span_type=SpanType.TOOL, package_count=len(packages) if packages else 0, # First few package names are useful in the UI for at-a-glance # identification; full list is in the LLM call record if needed. @@ -810,414 +523,198 @@ def registry_install_deps_span(packages: list[str] | None): def registry_reconstruct_pipeline_span(fmt: str | None): """Span around a ``reconstruct_pipeline`` call.""" - return child_span("registry.reconstruct_pipeline", format=fmt or "bash") + from mlflow.entities import SpanType - -# ----- Tool execution spans (Stage 3) ----- + return child_span( + "registry.reconstruct_pipeline", span_type=SpanType.TOOL, format=fmt or "bash" + ) -def tool_execute_span(record_id: str, tool_name: str): - """Span around a single ``dsagt-run`` tool execution. +# Code execution spans (dsagt-run). - This is a *top-level* span, not a child of any LLM call span. The agent - CLI (Claude Code, Goose, etc.) spawns ``dsagt-run`` in its own process - tree without OTel context propagation, so the LLM-call-to-tool-execution - parent/child linkage cannot be expressed in the trace tree directly. - Instead, every ``tool.execute`` span carries: - * ``session.id`` (attached automatically via init_tracing's session_id) - * ``record_id`` — matches the ``tool_use_id`` in the proxy's intent record, - which is how trace_archive correlates execution to LLM intent today. - * ``tool_name`` — for filtering in the MLflow UI. +def code_execute_span(record_id: str, code_name: str): + """Span around a single ``dsagt-run`` code execution. - Filter by ``session.id`` in the MLflow trace view to see all tool - executions for a given project alongside the LLM calls that requested - them. Cross-reference via ``record_id`` for full intent → execution - linkage when needed. + A *top-level, categorization-root* span: the agent CLI spawns ``dsagt-run`` + in its own process tree, so this trace stands alone in the store rather than + nesting under any MCP-dispatch span. It carries ``record_id`` (correlates + to the ``trace_archive`` record) and ``code_name``, and is tagged + ``dsagt.source=execution`` — its own bucket, distinct from the four MCP tool + categories, since these are actual code runs rather than meta-ops. """ + from mlflow.entities import SpanType @contextmanager def _wrapper(): - with open_span("tool.execute") as span: + with open_span( + "code.execute", span_type=SpanType.TOOL, source="execution" + ) as span: if span is None: yield None return - _attach_trace_metadata("tool.execute") span.set_attribute("record_id", record_id) - span.set_attribute("tool_name", tool_name) + span.set_attribute("code_name", code_name) yield span return _wrapper() -def truncate(value: str, limit: int = 256) -> str: - """Truncate a string for span attributes. +# =========================================================================== +# Replay sink — finished agent Trace → backdated MLflow spans +# =========================================================================== - Span backends (and the MLflow trace UI) handle short attribute values - much better than multi-megabyte stdout/stderr blobs. The full payload - lives in ``trace_archive/.json``; the span just carries a - head/tail summary so a human glancing at the UI can tell what happened. - """ - if value is None: - return "" - if len(value) <= limit: - return value - head = limit - 32 - return value[:head] + f"... [+{len(value) - head} chars]" +_S_PER_NS = 1e9 -# --------------------------------------------------------------------------- -# Phase 2: dsagt-proxy callback (cache breakpoints + sidechannel detection) -# --------------------------------------------------------------------------- -# -# When ``dsagt start --enable-proxy`` is set, the proxy intercepts every -# LLM call and emits OTLP spans to MLflow's /v1/traces endpoint via -# LiteLLM's built-in "otel" callback. Two extra concerns can't be done -# by pure observation — they require request mutation / response logging -# at the proxy hot path: -# -# 1. Anthropic prompt-cache breakpoint injection (saves money on -# multi-turn sessions; anthropic + bedrock-anthropic honor it, -# others ignore the marker). -# 2. Sidechannel-call detection (agents' hardcoded title-gen / session- -# namer models would 400 against lab gateways; we mock them and -# surface a single warning at teardown so the user can spot a typo -# in their primary llm.model vs. a harmless sidechannel hit). -# -# These live here (not in provenance.py) because they're observation-time -# concerns, gated by proxy mode, and consumed by the proxy entry point -# in commands/proxy_server.py via :func:`init_proxy_tracing`. +def _to_ns(epoch_s: float | None) -> int | None: + return int(epoch_s * _S_PER_NS) if epoch_s is not None else None -_CACHE_MARKER = {"type": "ephemeral"} +class MLflowSink: + """Render a :class:`~dsagt.traces.Trace` into MLflow spans (a trace consumer). -def _inject_cache_breakpoints(messages: list, kwargs: dict) -> None: - """Mark the largest stable request prefix as cacheable. + The agent half of observability: where ``@traced`` / ``obs`` emit DSAGT's + own first-party debug spans live, this replays a finished transcript's + :class:`~dsagt.traces.Trace` after the fact into the *same* store. It uses + ``mlflow.start_span_no_context`` — the only API that accepts an explicit + ``parent_span`` and backdated ``start_time_ns`` — and mirrors the span + conventions of MLflow's own ``claude_code`` autolog so foreign traces render + identically in the Chat UI: an AGENT root, ``llm`` children carrying + ``message.format="anthropic"`` + ``mlflow.chat.tokenUsage``, and + ``tool_`` children. Agent traces carry no ``dsagt.source`` tag, so + they stay in the normal view, separate from the internal debug traces. - Anthropic prompt caching keys on the prefix UP TO each marked block, so - one marker at the end of the tools array caches "system + tools", and a - second on the system message itself ensures the system block is cached - even on requests with no tools. Subsequent turns within the 5-minute - TTL pay 10% on the cached prefix instead of 100%. + A session ``Trace`` carries one AGENT subtree per turn; the sink emits **one + MLflow trace per AGENT root**, matching the per-prompt granularity autolog's + Stop hook produces. MLflow mints its own trace/span ids, so each trace is + tagged ``dsagt.trace_id = :`` (a stable per-turn + idempotency key). - Mutates ``messages`` and ``kwargs`` in place — the proxy reads the same - objects on its way to the upstream call. - - Marker is ``{"type": "ephemeral"}``; LiteLLM forwards it as - ``cache_control`` to anthropic-family providers (Anthropic direct, - Bedrock-Claude). Providers without prompt caching (current OpenAI - automatic caching ignores explicit markers, Cohere/Mistral/etc just - drop them) treat it as a no-op, so this is safe to set unconditionally. - """ - # 1) Tools: stamp the last tool definition. In LiteLLM's OpenAI-shape - # the tool dict carries cache_control as a top-level key and the - # Anthropic translator picks it up. - tools = kwargs.get("tools") or [] - if tools and isinstance(tools[-1], dict): - tools[-1]["cache_control"] = _CACHE_MARKER - - # 2) System message: stamp the last text block. For string content we - # promote to block format because cache_control lives on the block, - # not the message itself. - for msg in messages: - if msg.get("role") != "system": - continue - content = msg.get("content") - if isinstance(content, str): - msg["content"] = [ - { - "type": "text", - "text": content, - "cache_control": _CACHE_MARKER, - } - ] - elif isinstance(content, list) and content: - last_block = content[-1] - if isinstance(last_block, dict): - last_block["cache_control"] = _CACHE_MARKER - break - - -# --------------------------------------------------------------------------- -# Sidechannel model-call handling -# --------------------------------------------------------------------------- -# -# Every agent platform hardcodes a small/fast model for internal features -# (goose → gpt-4o-mini session-namer; claude → claude-haiku-4-5... title -# generator; the next agent will pick its own). When the user's gateway -# doesn't carry that exact bare name — which is the norm for lab gateways -# that alias every model — those requests would 400 and clutter MLflow. -# -# Rather than maintain a per-vendor list of known hardcoded names (which -# rots as vendors rename their sidechannel models), DSAGT catches all of -# them with one wildcard LiteLLM route, records which names fired, and -# surfaces a single yellow warning at session teardown so the user can -# distinguish a harmless sidechannel from a typo in their own config. -# -# Callers import the specific thing they need: -# commands/proxy_server._generate_config → SIDECHANNEL_WILDCARD_ROUTE_YAML -# DSAGTCallback.log_success_event → record_sidechannel_call() -# commands/cli._cmd_start (teardown) → print_sidechannel_warning() - -import json as _json # local alias keeps the section self-contained -import sys as _sys -from datetime import datetime as _datetime, timezone as _timezone -from pathlib import Path as _Path - -#: Env var the parent process exports so the proxy's DSAGT callback can tell -#: "this is the configured primary model" from "this is a sidechannel hit". -SIDECHANNEL_PRIMARY_MODEL_ENV = "DSAGT_PRIMARY_MODEL" - -#: JSONL file (one entry per intercepted call) that the proxy subprocess -#: appends to and the parent reads at teardown. Lives at the project -#: directory root, adjacent to ``trace_archive/``. -SIDECHANNEL_LOG_FILENAME = "sidechannel.jsonl" - -#: Canned reply the wildcard returns. Short enough that goose's -#: session-namer (expects ≤4 words) and claude's title generator both -#: accept it without error. -_SIDECHANNEL_CANNED_RESPONSE = "session" - -#: Post-routing model name LiteLLM resolves the wildcard catchall to. Only -#: true mock hits end up here; explicit primary entries and aliases route -#: to the real upstream regardless of what name the client requested. The -#: sidechannel detector uses this to avoid flagging alias hits as -#: sidechannel — see ``record_sidechannel_call``. -SIDECHANNEL_CATCHALL_MODEL = "openai/dsagt-sidechannel-catchall" - -#: Where the user can read the longer explanation. Printed in the warning. -SIDECHANNEL_DOC_LOCATION = "README.md § Sidechannel model calls" - -#: YAML fragment appended to the proxy's ``model_list`` after the primary -#: route. LiteLLM prefers exact matches over wildcards, so the configured -#: model still routes normally; everything else falls through to the mock. -#: -#: ``api_base`` has to be a syntactically valid URL but is never dialed — -#: ``mock_response`` short-circuits upstream entirely. -SIDECHANNEL_WILDCARD_ROUTE_YAML = f"""\ - - model_name: "*" - litellm_params: - model: {SIDECHANNEL_CATCHALL_MODEL} - api_base: http://invalid.local - api_key: unused - mock_response: "{_SIDECHANNEL_CANNED_RESPONSE}" -""" - - -def _sidechannel_client_requested_model(kwargs: dict) -> str | None: - """Return the model name the client sent, not the post-routing target. - - LiteLLM mutates ``kwargs["model"]`` to the resolved route during - completion, so by the time callbacks fire the original name is gone - from there — a wildcard hit's ``kwargs["model"]`` is always the - wildcard's ``litellm_params.model`` target. - - ``standard_logging_object.model_group`` preserves the name from the - client's request body. For exact matches it equals the configured - primary; for wildcard hits it's the actual sidechannel name (e.g. - ``gpt-4o-mini``, ``claude-haiku-4-5-20251001``) — which is what the - warning needs to show. - """ - slo = kwargs.get("standard_logging_object") or {} - if isinstance(slo, dict): - grp = slo.get("model_group") - if grp: - return grp.split("/", 1)[-1] - # Fallback: whatever kwargs has. Worse (may be the catchall name) but - # still informative if standard_logging_object isn't populated. - m = kwargs.get("model") or "" - return m.split("/", 1)[-1] or None - - -def record_sidechannel_call(records_dir, kwargs: dict) -> None: - """Append a JSONL entry when a request hit the wildcard catchall. - - Detection rule: ``kwargs["model"]`` (the post-routing target LiteLLM - selected) equals ``SIDECHANNEL_CATCHALL_MODEL``. This is the only - reliable discriminator: name-based comparison against the primary - misclassifies alias hits (e.g. ``claude-sonnet-4-5`` aliased to a - longer lab-gateway model name) as sidechannel even though they - actually routed to the real upstream. - - Called from the DSAGT callback's success handler, so only successful - calls get logged — failures land in MLflow as errors regardless. - No-ops when ``SIDECHANNEL_PRIMARY_MODEL_ENV`` isn't set. - """ - primary = os.environ.get(SIDECHANNEL_PRIMARY_MODEL_ENV) - if not primary: - return - - routed_to = kwargs.get("model") or "" - if routed_to != SIDECHANNEL_CATCHALL_MODEL: - return # real upstream call (primary entry or alias) — not a sidechannel - - requested = _sidechannel_client_requested_model(kwargs) - if not requested: - return - - entry = { - "timestamp": _datetime.now(_timezone.utc).isoformat().replace("+00:00", "Z"), - "model": requested, - "agent": os.environ.get("DSAGT_AGENT", ""), - "session": os.environ.get("DSAGT_SESSION_ID", ""), - } - path = _Path(records_dir).parent / SIDECHANNEL_LOG_FILENAME - try: - path.parent.mkdir(parents=True, exist_ok=True) - with open(path, "a") as f: - f.write(_json.dumps(entry) + "\n") - except Exception as e: - logger.debug("sidechannel log append failed: %s", e) - - -def print_sidechannel_warning(project_dir, session_id: str | None) -> None: - """Read ``SIDECHANNEL_LOG_FILENAME`` under *project_dir*, dedup within - *session_id*, and print a yellow warning to stdout. - - No-op when nothing was logged for this session (common case). The - warning lists each unique model that hit the wildcard along with the - call count, and points the user at ``SIDECHANNEL_DOC_LOCATION`` for the - two possible causes (harmless sidechannel vs config typo). - - ANSI colors are only emitted when stdout is a TTY — CI logs stay clean. - """ - log_path = _Path(project_dir) / SIDECHANNEL_LOG_FILENAME - if not log_path.exists(): - return - - # Dedup by model name within the current session only. The file is - # append-only across runs, so older sessions' entries are still present - # but don't belong in this run's warning. - seen: dict[str, int] = {} - try: - for line in log_path.read_text().splitlines(): - if not line.strip(): - continue - entry = _json.loads(line) - if session_id and entry.get("session") != session_id: - continue - model = entry.get("model") or "" - seen[model] = seen.get(model, 0) + 1 - except (OSError, ValueError): - return - - if not seen: - return - - tty = _sys.stdout.isatty() - yellow = "\033[33m" if tty else "" - bold = "\033[1m" if tty else "" - reset = "\033[0m" if tty else "" - - print() - print(f"{yellow}{bold} ⚠ Sidechannel model calls intercepted:{reset}") - for model, count in sorted(seen.items(), key=lambda kv: -kv[1]): - s = "s" if count != 1 else "" - print(f"{yellow} {model} ({count} call{s}){reset}") - print(f"{yellow} Two possible causes:{reset}") - print( - f"{yellow} (1) agent sidechannel (e.g. title generator) — safe to ignore{reset}" - ) - print( - f"{yellow} (2) typo in dsagt_config.yaml llm.model — these replies are canned, not real{reset}" - ) - print(f"{yellow} See: {SIDECHANNEL_DOC_LOCATION}{reset}") - - -# --------------------------------------------------------------------------- -# DSAGTCallback — LiteLLM CustomLogger for proxy-mode cache + sidechannel -# --------------------------------------------------------------------------- - - -def _make_dsagt_callback(records_dir): - """Construct a LiteLLM CustomLogger with cache injection + sidechannel. - - Lazy-imports LiteLLM so the rest of observability.py is testable - without it on the import path. Trace transport is via - ``litellm.callbacks = ["otel"]``, configured by :func:`init_proxy_tracing` - — this callback is only for the two intercept-time concerns LiteLLM - autolog can't cover. + A *consumer* of :class:`~dsagt.traces.TraceCollector`: ``name`` keys its own + ack file (``.dsagt/trace_acks_mlflow.json``); ``write`` logs the trace. + Spans are plain dicts (see ``traces`` module docstring), so this reads them + directly — no per-object serialization. """ - from litellm.integrations.custom_logger import CustomLogger - - class DSAGTCallback(CustomLogger): - def __init__(self, records_dir): - self.records_dir = records_dir - - def log_pre_api_call(self, model, messages, kwargs): - _inject_cache_breakpoints(messages, kwargs) - - def log_success_event(self, kwargs, response_obj, start_time, end_time): - record_sidechannel_call(self.records_dir, kwargs) - - async def async_log_success_event( - self, kwargs, response_obj, start_time, end_time - ): - record_sidechannel_call(self.records_dir, kwargs) - - def log_failure_event(self, kwargs, response_obj, start_time, end_time): - logger.warning("LLM call failed: model=%s", kwargs.get("model")) - - async def async_log_failure_event( - self, kwargs, response_obj, start_time, end_time - ): - logger.warning("LLM call failed: model=%s", kwargs.get("model")) - return DSAGTCallback(records_dir) + name = "mlflow" + def __init__(self, tracking_uri: str, experiment: str): + self._uri = tracking_uri + self._experiment = experiment -def init_proxy_tracing( - mlflow_url: str, project: str, session_id: str | None, records_dir -) -> None: - """Configure LiteLLM proxy → MLflow native MlflowLogger transport, plus - install the DSAGT cache + sidechannel callback. - - Installs MLflow's tracer provider (so any span helpers in the proxy - process emit through MLflow), then plants ``_DSAGTMlflowLogger`` (a - subclass of LiteLLM's ``MlflowLogger``) into LiteLLM's logger cache so - the string ``"mlflow"`` in ``success_callback``/``failure_callback`` - routes through it. The subclass stamps ``mlflow.trace.session`` / - ``dsagt.source=agent`` / ``dsagt.agent`` on every proxy-captured trace - in the narrow window between trace creation and export. - - Caller is ``commands/proxy_server.py:main``. ``records_dir`` is the - project's ``trace_archive/`` so the sidechannel jsonl lands adjacent. - """ - global _initialized, _default_session_id - global _metadata_stamper, _llm_context_factory - - _install_mlflow_provider(mlflow_url, project) - _metadata_stamper = _stamp_metadata_mlflow - _llm_context_factory = _mlflow_llm_context - if session_id: - _default_session_id = session_id - _initialized = True + def write(self, trace) -> list[str]: + """Log every turn subtree; return the MLflow trace id of each.""" + import mlflow - # Install our MlflowLogger subclass so trace metadata stamping happens - # in the narrow window between trace creation and export. Lazy import - # to keep the import graph clean (provenance imports observability, - # not the other way round at top level). - from dsagt.provenance import install_mlflow_logger_with_session_tag + mlflow.set_tracking_uri(self._uri) + mlflow.set_experiment(trace.project or self._experiment) - install_mlflow_logger_with_session_tag() + children: dict[str, list] = {} + for span in trace.spans: + if span["parent_id"] is not None: + children.setdefault(span["parent_id"], []).append(span) - import litellm + trace_ids = [] + for root in trace.spans: + if root["parent_id"] is None: + trace_ids.append( + self._emit_subtree(root, children.get(root["span_id"], []), trace) + ) + return trace_ids - # Cache + sidechannel callback (intercept-time concerns). - callback = _make_dsagt_callback(records_dir) - litellm.callbacks = [callback] + def _emit_subtree(self, root, children, trace) -> str: + """Emit one MLflow trace for an AGENT ``root`` and its direct children.""" + import mlflow + from mlflow.entities import SpanType + from mlflow.tracing.constant import ( + SpanAttributeKey, + TokenUsageKey, + TraceMetadataKey, + ) + from mlflow.tracing.trace_manager import InMemoryTraceManager - # Register the built-in "mlflow" callback by NAME, not instance. LiteLLM - # async dispatch resolves names to logger classes via _in_memory_loggers - # at each call; passing a string ensures our pre-seeded subclass is used - # (it passes the isinstance(MlflowLogger) check). - if "mlflow" not in (litellm.success_callback or []): - litellm.success_callback = (litellm.success_callback or []) + ["mlflow"] - if "mlflow" not in (litellm.failure_callback or []): - litellm.failure_callback = (litellm.failure_callback or []) + ["mlflow"] + kind_to_type = { + "AGENT": SpanType.AGENT, + "LLM": SpanType.LLM, + "TOOL": SpanType.TOOL, + "OTHER": SpanType.UNKNOWN, + } + + ml_root = mlflow.start_span_no_context( + name=root["name"], + span_type=kind_to_type[root["kind"]], + inputs={"prompt": root["attributes"].get("prompt", "")}, + start_time_ns=_to_ns(root["start_time"]), + ) - logger.info( - "init_proxy_tracing: service=dsagt-proxy mlflow=%s session=%s", - mlflow_url, - session_id, - ) + for span in children: + if span["kind"] == "LLM": + child = mlflow.start_span_no_context( + name=span["name"], + parent_span=ml_root, + span_type=SpanType.LLM, + start_time_ns=_to_ns(span["start_time"]), + inputs={ + "model": span["model"] or "unknown", + "messages": span["request"], + }, + attributes={ + "model": span["model"] or "unknown", + SpanAttributeKey.MESSAGE_FORMAT: "anthropic", + }, + ) + if span["usage"]: + inp = span["usage"].get("input_tokens") or 0 + out = span["usage"].get("output_tokens") or 0 + child.set_attribute( + SpanAttributeKey.CHAT_USAGE, + { + TokenUsageKey.INPUT_TOKENS: inp, + TokenUsageKey.OUTPUT_TOKENS: out, + TokenUsageKey.TOTAL_TOKENS: inp + out, + }, + ) + child.set_outputs( + { + "type": "message", + "role": "assistant", + "content": span["response"], + } + ) + else: # TOOL / OTHER + child = mlflow.start_span_no_context( + name=span["name"], + parent_span=ml_root, + span_type=kind_to_type[span["kind"]], + start_time_ns=_to_ns(span["start_time"]), + inputs=span["attributes"].get("input", {}), + attributes={ + "tool_name": span["attributes"].get("tool_name"), + "tool_id": span["attributes"].get("tool_id"), + }, + ) + child.set_outputs({"result": span["attributes"].get("result", "")}) + child.end(end_time_ns=_to_ns(span["end_time"])) + + # Trace-level metadata: session correlation + the per-turn canonical id + # (idempotency key) + request/response previews for the trace list. + try: + mgr = InMemoryTraceManager.get_instance() + with mgr.get_trace(ml_root.trace_id) as in_mem: + meta = { + TraceMetadataKey.TRACE_SESSION: trace.session_id, + "dsagt.trace_id": f"{trace.trace_id}:{root['span_id']}", + "dsagt.agent": trace.agent, + } + in_mem.info.trace_metadata = {**in_mem.info.trace_metadata, **meta} + if prompt := root["attributes"].get("prompt"): + in_mem.info.request_preview = str(prompt)[:1000] + if response := root["attributes"].get("response"): + in_mem.info.response_preview = str(response)[:1000] + except Exception as e: # noqa: BLE001 + logger.warning("MLflowSink: could not stamp trace metadata: %s", e) + + ml_root.set_outputs({"response": root["attributes"].get("response", "")}) + ml_root.end(end_time_ns=_to_ns(root["end_time"])) + return ml_root.trace_id diff --git a/src/dsagt/provenance.py b/src/dsagt/provenance.py index 2fc1dc8..9abe20c 100644 --- a/src/dsagt/provenance.py +++ b/src/dsagt/provenance.py @@ -1,51 +1,46 @@ """ -Provenance for tool executions. +Provenance for code executions. **Execution capture** (dsagt-run wrapper): Wraps a shell command, captures exact execution data (command, exit code, stdout/stderr, input/output files), and writes a JSON record - to ``trace_archive/__.json``. + to ``trace_archive/__.json``. **Record indexing** (ChromaDB): - Indexes execution records into a ``tool_executions`` collection for + Indexes execution records into a ``code_use`` collection for semantic search and metadata filtering. **Pipeline reconstruction**: Reads execution records, builds a dependency graph from input/output file overlap, and renders as a bash script or Snakemake workflow. - -LLM-call provenance lives in MLflow (each MCP-server / agent process -autologs LiteLLM calls via ``init_tracing`` post-proxy-removal). """ from __future__ import annotations +import fcntl import json import logging -import os import subprocess import sys import time import uuid +from contextlib import contextmanager from pathlib import Path from datetime import datetime, timezone +from typing import TYPE_CHECKING -from dsagt.knowledge import CollectionRoute, KnowledgeBase +if TYPE_CHECKING: + # Annotation-only (this module has ``from __future__ import annotations``, so + # the hint is a string). Importing KnowledgeBase at runtime would drag the + # whole retrieval module into ``dsagt-run``, which only writes provenance + # records to disk and never touches a KB — the embedding of those records + # happens later, in ``session.run_extraction``. + from dsagt.knowledge import KnowledgeBase logger = logging.getLogger(__name__) -#: Project-local collection of indexed tool-execution records. -#: Renamed from ``tool_executions`` to match the user-facing -#: terminology ("tool_use index"). -TOOL_USE_COLLECTION = "tool_use" - -#: Backwards-compat alias. New code should use ``TOOL_USE_COLLECTION``. -TOOL_EXECUTIONS_COLLECTION = TOOL_USE_COLLECTION -TOOL_EXECUTIONS_ROUTE = CollectionRoute( - embedding_backend="api", - vector_db="chroma", - description="Indexed tool execution records from trace_archive.", -) +#: Project-local collection of indexed code-execution records. +CODE_USE_COLLECTION = "code_use" # --------------------------------------------------------------------------- @@ -57,7 +52,7 @@ def _resolve_records_dir(explicit: str | None) -> Path: """Determine the records directory. Priority: explicit ``--records-dir`` flag → ``/trace_archive``, - where cwd must contain ``dsagt_config.yaml`` (the project's + where cwd must contain ``.dsagt/config.yaml`` (the project's single-source-of-truth config written by ``dsagt init``). No env-var chain, no walking up the tree — if the agent's cwd isn't the project dir, that's the bug to fix, not something to recover from silently. @@ -65,14 +60,32 @@ def _resolve_records_dir(explicit: str | None) -> Path: if explicit: return Path(explicit) cwd = Path.cwd().resolve() - if not (cwd / "dsagt_config.yaml").exists(): + if not (cwd / ".dsagt" / "config.yaml").exists(): raise ValueError( - f"No dsagt_config.yaml in cwd ({cwd}); pass --records-dir or " + f"No .dsagt/config.yaml in cwd ({cwd}); pass --records-dir or " "run dsagt-run from a project directory." ) return cwd / "trace_archive" +def _current_session_tag_from_cwd() -> str | None: + """Read the current session tag from ``/.dsagt/state.yaml``. + + ``dsagt-run`` runs with cwd == project dir; the MCP server (also a child + of the agent) minted the session into ``state.yaml`` at startup. Lazy + import of ``session`` avoids a circular import (``session`` imports this + module for ``index_trace_archive``). + """ + from dsagt import session + + cwd = Path.cwd().resolve() + cfg = session.read_config_file(cwd) + project = cfg.get("project") + if not project: + return None + return session.current_session_tag(cwd, project) + + def _parse_file_list(raw: str | None) -> list[str]: """Split a comma-separated file list, stripping whitespace.""" if not raw: @@ -81,7 +94,7 @@ def _parse_file_list(raw: str | None) -> list[str]: def run_and_record( - tool_name: str, + code_name: str, command: list[str], records_dir: Path, session_id: str | None = None, @@ -90,19 +103,17 @@ def run_and_record( output_files: list[str] | None = None, ) -> int: """Execute a command, write an execution record, return the exit code.""" - from dsagt.observability import obs, tool_execute_span, truncate + from dsagt.observability import obs, code_execute_span, truncate record_id = record_id or uuid.uuid4().hex[:12] if session_id is None: - from dsagt.observability import ( - find_project_config, - _read_session_id_from_runtime, - ) - - pdir, _ = find_project_config() - session_id = _read_session_id_from_runtime(pdir) + # The MCP server mints the session at startup and records it in + # ``.dsagt/state.yaml``; read the current tag from there so this + # code span buckets with the rest of the session (cwd == project + # dir by contract). ``None`` if no session has been minted yet. + session_id = _current_session_tag_from_cwd() - with tool_execute_span(record_id, tool_name): + with code_execute_span(record_id, code_name): timestamp_start = datetime.now(timezone.utc).isoformat() start_perf = time.perf_counter() @@ -144,14 +155,14 @@ def run_and_record( if stderr.strip(): obs.set("stderr_truncated", truncate(stderr, 256)) if return_code != 0: - obs.event("tool_failed", exit_code=return_code) + obs.event("code_failed", exit_code=return_code) # Populate the MLflow trace UI's Input/Output tabs. Truncate to - # ~4KB per side so big tool results don't bloat the trace store + # ~4KB per side so big code results don't bloat the trace store # (the full payload is on disk in trace_archive/.json). obs.set_inputs( { - "tool": tool_name, + "code": code_name, "command": list(command), "input_files": input_files or [], } @@ -168,7 +179,7 @@ def run_and_record( record = { "record_id": record_id, - "tool_name": tool_name, + "code_name": code_name, "session_id": session_id, "execution": { "exact_command": command, @@ -197,7 +208,7 @@ def _write_record(record: dict, records_dir: Path) -> Path: records_dir.mkdir(parents=True, exist_ok=True) ts = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ") - filename = f"{record['tool_name']}_{ts}_{record['record_id']}.json" + filename = f"{record['code_name']}_{ts}_{record['record_id']}.json" path = records_dir / filename path.write_text(json.dumps(record, indent=2, ensure_ascii=False) + "\n") @@ -210,31 +221,22 @@ def _write_record(record: dict, records_dir: Path) -> Path: def render_execution_text(record: dict) -> str: - """Convert a tool execution record into embeddable natural-language text.""" - tool_name = record.get("tool_name", "unknown") - intent = record.get("intent") or {} + """Convert a code execution record into embeddable natural-language text.""" + code_name = record.get("code_name", "unknown") execution = record.get("execution") - report = record.get("report") or {} - parts = [f"Tool: {tool_name}"] + parts = [f"Code: {code_name}"] if execution and execution.get("exact_command"): cmd = execution["exact_command"] if isinstance(cmd, list): cmd = " ".join(cmd) parts.append(f"Command: {cmd}") - elif intent.get("parameters"): - parts.append(f"Parameters: {json.dumps(intent['parameters'])}") if execution and execution.get("return_code") is not None: rc = execution["return_code"] status = "succeeded" if rc == 0 else f"failed (exit code {rc})" parts.append(f"Outcome: {status}") - elif report.get("agent_output"): - output = report["agent_output"] - if len(output) > 500: - output = output[:500] + "..." - parts.append(f"Agent report: {output}") if execution: start = execution.get("timestamp_start", "") @@ -258,26 +260,18 @@ def render_execution_text(record: dict) -> str: def execution_metadata(record: dict) -> dict: - """Extract ChromaDB-filterable metadata from a tool execution record.""" + """Extract ChromaDB-filterable metadata from a code execution record.""" execution = record.get("execution") - intent = record.get("intent") or {} meta: dict = {} - meta["tool_name"] = record.get("tool_name", intent.get("command", "unknown")) - meta["session_id"] = record.get("session_id", intent.get("session_id", "unknown")) + meta["code_name"] = record.get("code_name", "unknown") + meta["session_id"] = record.get("session_id", "unknown") if execution and execution.get("return_code") is not None: meta["return_code"] = execution["return_code"] - meta["wrapper_used"] = 1 if execution is not None else 0 - - timestamp = None if execution and execution.get("timestamp_start"): - timestamp = execution["timestamp_start"] - elif intent.get("timestamp_requested"): - timestamp = intent["timestamp_requested"] - if timestamp: - meta["timestamp"] = timestamp + meta["timestamp"] = execution["timestamp_start"] record_id = record.get("record_id", "") if record_id: @@ -286,27 +280,12 @@ def execution_metadata(record: dict) -> dict: return meta -def index_execution_record(record: dict, kb: KnowledgeBase) -> dict: - """Render and store a single tool execution record in the knowledge base.""" - text = render_execution_text(record) - metadata = execution_metadata(record) - - # No ``route=`` — fall through to kb's default route so embedding - # backend follows the project's ``embedding.backend`` config (BYOA - # default = local sentence-transformers, no API call). - return kb.add_entries( - texts=[text], - collection=TOOL_EXECUTIONS_COLLECTION, - metadatas=[metadata], - ) - - def index_trace_archive( trace_dir: Path, kb: KnowledgeBase, indexed_ids: set[str] | None = None, ) -> dict: - """Batch-index all tool execution records in a trace archive directory.""" + """Batch-index all code execution records in a trace archive directory.""" if indexed_ids is None: indexed_ids = set() @@ -331,10 +310,8 @@ def index_trace_archive( skipped += 1 continue - has_intent = "intent" in record - has_execution = record.get("execution") is not None - if not has_intent and not has_execution: - logger.warning("Skipping %s: no intent or execution layer", path.name) + if record.get("execution") is None: + logger.warning("Skipping %s: no execution layer", path.name) errors += 1 continue @@ -348,7 +325,7 @@ def index_trace_archive( if texts: kb.add_entries( texts=texts, - collection=TOOL_EXECUTIONS_COLLECTION, + collection=CODE_USE_COLLECTION, metadatas=metadatas, ) @@ -360,6 +337,86 @@ def index_trace_archive( } +class CodeUseIndexer: + """Idempotent, incremental indexer of ``dsagt-run`` records into ``code_use``. + + The code-execution counterpart to :class:`~dsagt.trace_scan.TraceScan`: + ``dsagt-run`` writes one JSON record per call to ``trace_archive/``, and each + :meth:`tick` embeds only the records not already indexed — tracked by + ``record_id`` in a persisted ack set — so re-ticks and cross-session + re-reads can never duplicate (the bug the prior cursor-less batch had). + + One primitive, three triggers, all safe to overlap: the MCP-server heartbeat + (current-session freshness), startup catch-up (the previous session's tail), + and the ``reconstruct_pipeline`` code (index-then-reconstruct, so a pipeline + review reflects the calls just made). An OS file lock around + load→index→save serializes those callers — distinct instances in one + process, or a future cross-process ticker — against the shared ack file. + """ + + def __init__(self, kb: KnowledgeBase, project_dir: str | Path): + self._kb = kb + pdir = Path(project_dir) + self._trace_dir = pdir / "trace_archive" + self._acks_path = pdir / ".dsagt" / "code_use_acks.json" + + def _load_acks(self) -> set[str]: + try: + return set(json.loads(self._acks_path.read_text())) + except FileNotFoundError: + return set() + + def _save_acks(self, acks: set[str]) -> None: + self._acks_path.parent.mkdir(parents=True, exist_ok=True) + self._acks_path.write_text(json.dumps(sorted(acks))) + + @contextmanager + def _lock(self): + self._acks_path.parent.mkdir(parents=True, exist_ok=True) + lock_path = self._acks_path.with_suffix(".lock") + with open(lock_path, "w") as lf: + fcntl.flock(lf, fcntl.LOCK_EX) + try: + yield + finally: + fcntl.flock(lf, fcntl.LOCK_UN) + + def tick(self) -> int: + """Index newly-arrived records; return how many were indexed this tick. + + Emits no categorization root of its own: at the ``reconstruct_pipeline`` + call site this runs *inside* the registry tool's trace, so its ``kb.*`` + writes correctly inherit ``dsagt.source=registry``. The background + callers (heartbeat / startup catch-up) run outside any trace and use + :meth:`tick_traced` so their writes don't orphan as untagged roots. + """ + with self._lock(): + acks = self._load_acks() + before = len(acks) + # index_trace_archive skips record_ids already in ``acks`` and adds + # the newly-indexed ones to it (mutates the set we pass). + result = index_trace_archive(self._trace_dir, self._kb, indexed_ids=acks) + if len(acks) != before: + self._save_acks(acks) + return result.get("indexed", 0) + + def tick_traced(self) -> int: + """:meth:`tick` under a ``dsagt.source=code_use`` categorization root. + + For the background triggers (heartbeat, startup catch-up) that run off + any tool-call trace — otherwise the indexer's ``kb.add_entries`` / + ``kb.embed`` spans start their own untagged top-level traces, landing in + the ``unknown`` bucket and detached from the executions they index. + Must run on the same thread as the embedding (open the span inside the + worker), so callers dispatch *this* to the thread, not a wrapped + :meth:`tick`. + """ + from dsagt.observability import open_span + + with open_span("code_use.index", source="code_use"): + return self.tick() + + # --------------------------------------------------------------------------- # Pipeline reconstruction # --------------------------------------------------------------------------- @@ -416,20 +473,20 @@ def render_bash(records: list[dict], deps: dict[int, list[int]]) -> str: ] for i, record in enumerate(records): - tool = record["tool_name"] + code = record["code_name"] execution = record["execution"] cmd = execution["exact_command"] rc = execution.get("return_code", 0) inputs = execution.get("input_files", []) outputs = execution.get("output_files", []) - lines.append(f"# Step {i + 1}: {tool}") + lines.append(f"# Step {i + 1}: {code}") if inputs: lines.append(f"# inputs: {', '.join(inputs)}") if outputs: lines.append(f"# outputs: {', '.join(outputs)}") if deps[i]: - dep_names = [records[d]["tool_name"] for d in deps[i]] + dep_names = [records[d]["code_name"] for d in deps[i]] lines.append(f"# depends: {', '.join(dep_names)}") if rc != 0: lines.append(f"# WARNING: original run exited with code {rc}") @@ -450,8 +507,8 @@ def render_snakemake(records: list[dict], deps: dict[int, list[int]]) -> str: rule_names = [] for i, record in enumerate(records): - tool = record["tool_name"] - rule_name = f"{tool}_{i + 1}" + code = record["code_name"] + rule_name = f"{code}_{i + 1}" rule_names.append(rule_name) all_outputs = [] @@ -528,115 +585,3 @@ def reconstruct_pipeline( if fmt == "snakemake": return render_snakemake(records, deps) return render_bash(records, deps) - - -# --------------------------------------------------------------------------- -# Proxy-mode MLflow logger (LiteLLM callback subclass) -# --------------------------------------------------------------------------- - - -def install_mlflow_logger_with_session_tag() -> None: - """Inject a ``MlflowLogger`` subclass that stamps ``mlflow.trace.session``. - - Why a subclass instead of post-hoc tagging from another callback: - LiteLLM's MlflowLogger does the entire trace lifecycle (start_trace → - set_attributes → end_trace) inside one ``_handle_success`` call. By - the time any sibling success-callback fires, the trace is already - exported and ``mlflow.update_current_trace`` has nothing to update. - Subclassing lets us slot the metadata write between trace creation and - trace export, where the in-memory trace still exists and is mutable. - - Why register via LiteLLM's ``_in_memory_loggers`` cache: when the - proxy resolves the string ``"mlflow"`` from ``success_callback``, it - iterates ``_in_memory_loggers`` looking for any - ``isinstance(_, MlflowLogger)`` and returns it instead of constructing - a fresh one. A subclass passes the isinstance check, so pre-seeding - our subclass means the existing string-based registration in - ``success_callback``/``failure_callback`` automatically routes through - us — no sync-vs-async dispatch quirks to worry about. - - Idempotent: safe to call more than once. - """ - from litellm.integrations.mlflow import MlflowLogger - from litellm.litellm_core_utils import litellm_logging as _ll - - from dsagt.observability import _stamp_metadata_on_trace, extract_cache_stats - - class _DSAGTMlflowLogger(MlflowLogger): - def _start_span_or_trace(self, kwargs, start_time): - span = super()._start_span_or_trace(kwargs, start_time) - # span.request_id is MLflow's trace_id; the trace lives in - # InMemoryTraceManager until the parent _handle_success calls - # _end_span_or_trace, so we have a window here to mutate - # trace_metadata before export. - # - # Agent-turn LLM calls land here (proxy path). Stamp: - # mlflow.trace.session — session grouping in the UI - # dsagt.source=agent — distinguishes from extraction/embedding - # dsagt.agent — which platform (goose, claude, ...) - # Non-agent LLM calls (memory extraction, embeddings) go through - # llm_source(...) decorators and never touch this subclass, so - # hard-coding source="agent" here is safe. - if span is None: - return span - metadata: dict[str, str] = {"dsagt.source": "agent"} - if session_id := os.environ.get("DSAGT_SESSION_ID"): - metadata["mlflow.trace.session"] = session_id - if agent := os.environ.get("DSAGT_AGENT"): - metadata["dsagt.agent"] = agent - _stamp_metadata_on_trace(span.request_id, metadata) - return span - - def _extract_and_set_chat_attributes(self, span, kwargs, response_obj): - # Last window to stamp cache stats before _end_span_or_trace - # exports the trace. - super()._extract_and_set_chat_attributes(span, kwargs, response_obj) - if span is None: - return - usage = (_response_to_usage(response_obj) or {}).get("usage") or {} - read, write = extract_cache_stats(usage) - if not read and not write: - return - _stamp_metadata_on_trace( - span.request_id, - { - "dsagt.cache.read_tokens": str(read), - "dsagt.cache.write_tokens": str(write), - }, - ) - - # Already installed? Leave it. - for cb in _ll._in_memory_loggers: - if type(cb).__name__ == "_DSAGTMlflowLogger": - return - # Drop any vanilla MlflowLogger that beat us to the cache. - _ll._in_memory_loggers[:] = [ - cb - for cb in _ll._in_memory_loggers - if not ( - isinstance(cb, MlflowLogger) and type(cb).__name__ != "_DSAGTMlflowLogger" - ) - ] - _ll._in_memory_loggers.append(_DSAGTMlflowLogger()) - - -def _response_to_usage(response_obj) -> dict | None: - """Best-effort extraction of the ``usage`` dict from a litellm response. - - LiteLLM responses can be dataclass-like (``model_dump()``), pydantic - models (``dict()``), or already plain dicts. We try each, and fall - back to ``None`` if nothing yields a usable shape. - """ - if response_obj is None: - return None - if isinstance(response_obj, dict): - return response_obj - for method in ("model_dump", "dict"): - if hasattr(response_obj, method): - try: - d = getattr(response_obj, method)() - if isinstance(d, dict): - return d - except Exception: - continue - return None diff --git a/src/dsagt/registry.py b/src/dsagt/registry.py index 7e91aed..2278e6a 100644 --- a/src/dsagt/registry.py +++ b/src/dsagt/registry.py @@ -1,12 +1,26 @@ """ -Tool and Skill Registries. +Code and Skill Registries. Two parallel registries for agent capabilities: -**Tools** (CLI executables) — markdown files with YAML frontmatter specifying -name, description, executable, parameters, dependencies, tags. Stored in -`/tools/`. Agent-written scripts go in `/tools/code/`. +**Codes** (CLI executables) — skill-standard directories +(`/codes//SKILL.md`) whose frontmatter carries the machine +fields (name, description, executable, parameters, dependencies, tags) on +top of the skill-required name/description. Agent-written scripts live +beside their spec in `/codes//scripts/`, making each +registered code a self-contained, portable directory. The skill-standard +envelope means codes mirror into the agent's native skills dir unchanged +(see ``AgentSetup.setup_skills``) — native discovery puts the exact +runnable command in context at invocation time, alongside MCP discovery +via ``search_registry``. When registered, executables are wrapped with dsagt-run + uv run --with. +The wrapper lives *inside* the stored shell command by design: execution +used to be dispatched by MCP-server tools, but agents routinely sidestepped +those with their own bash tools, losing provenance. Baking dsagt-run into +the command the agent copies makes the bash path harmless — the residual +failure mode is an agent reconstructing the command from memory and +dropping the wrapper, which is why specs render the exact runnable command +and agent instructions say to copy it verbatim. **Skills** (agent instructions) — directories containing a SKILL.md with YAML frontmatter (name, description, tags) and optional reference docs. @@ -14,36 +28,64 @@ workflow instructions. Both registries support optional KB indexing for semantic search via -`search_registry` (tools) and `search_skills` (skills) MCP tools. +`search_registry` (codes) and `search_skills` (skills) MCP tools. """ +from __future__ import annotations + import logging +import re +import shutil from pathlib import Path +from typing import TYPE_CHECKING import yaml -from dsagt.knowledge import KnowledgeBase +if TYPE_CHECKING: + # Annotation-only. A runtime import would pull the whole retrieval module + # into anything that touches the registry — including ``dsagt-run`` via the + # package ``__init__`` — even though the registry only ever holds an + # injected KB instance, never references the class. + from dsagt.knowledge import KnowledgeBase logger = logging.getLogger(__name__) #: Single project-local collection holding both bundled (package-shipped) -#: and registered (agent-saved) tools. Bundled entries carry +#: and registered (agent-saved) codes. Bundled entries carry #: ``metadata.source = "bundled"`` and ``metadata.dsagt_version`` so #: they can be evicted and refreshed on dsagt upgrade without touching #: agent-registered entries. -TOOLS_COLLECTION = "tools" +CODES_COLLECTION = "codes" +#: Legacy installed-skills collection. No longer written or read: installed +#: skills are natively auto-discovered by every supported agent, so skill +#: search covers only the *catalog* tier below. Kept as a name for back-compat +#: and ``dsagt info`` display of any pre-existing index. SKILLS_COLLECTION = "skills" +#: External skill catalogs (fetched from GitHub repos) live in their own +#: per-source collections named ``skills_catalog__``. Keeping each +#: source in its own collection lets a re-sync drop+rebuild one source's +#: directory without disturbing other catalogs — no delete-by-metadata +#: primitive needed. +CATALOG_COLLECTION_PREFIX = "skills_catalog__" + + +def catalog_collection(slug: str) -> str: + """KB collection name holding the indexed catalog for source *slug*.""" + return f"{CATALOG_COLLECTION_PREFIX}{slug}" + + #: Backwards-compat aliases — kept so external code that imported the #: previous names still resolves. New code should use the names above. -TOOL_REGISTRY_COLLECTION = TOOLS_COLLECTION +TOOL_REGISTRY_COLLECTION = CODES_COLLECTION SKILL_REGISTRY_COLLECTION = SKILLS_COLLECTION # --------------------------------------------------------------------------- -# Helpers (tools only) +# Helpers (codes only) # --------------------------------------------------------------------------- + def _uv_run_prefix(deps: list[str]) -> str: """Build a 'uv run --with dep1,dep2 --' prefix for Python dependencies.""" if not deps: @@ -51,24 +93,38 @@ def _uv_run_prefix(deps: list[str]) -> str: return f"uv run --with {','.join(deps)} -- " +#: Skill-standard name charset — agent native skill loaders (claude et al.) +#: require lowercase-hyphen names, and codes mirror into those dirs. +_CODE_NAME_RE = re.compile(r"^[a-z0-9][a-z0-9-]*$") + + def _wrap_executable(name: str, executable: str, deps: list[str] | None = None) -> str: """Wrap an executable with uv run (for Python deps) and dsagt-run (for provenance). - Result: dsagt-run --tool -- [uv run --with deps --] + Result: dsagt-run --code -- [uv run --with deps --] """ if "dsagt-run" in executable: return executable inner = f"{_uv_run_prefix(deps or [])}{executable}" - return f"dsagt-run --tool {name} -- {inner}" + return f"dsagt-run --code {name} -- {inner}" -def _generate_tool_body(spec: dict) -> str: - """Generate a markdown body for a new tool file from its spec.""" +def _generate_code_body(spec: dict) -> str: + """Generate a markdown body for a new code file from its spec. + + The exact runnable command leads the body: native skill discovery + injects SKILL.md at invocation time, and the thing the agent must copy + verbatim (the dsagt-run-wrapped command) belongs at the top, not after + prose it may stop reading. + """ lines = [ - f"\n# {spec['name']}\n\n{spec['description']}\n\n", - "## Shell Command\n\n```bash\n", + f"\n# {spec['name']}\n\n", + "Run this registered code with the exact shell command below — copy " + "it byte-for-byte (the `dsagt-run` prefix writes the execution " + "record to `trace_archive/`):\n\n```bash\n", f"{spec['executable']} [options]\n", - "```\n\n## Parameters\n\n", + "```\n\n", + f"{spec['description']}\n\n## Parameters\n\n", ] params = spec.get("parameters", {}) if params: @@ -77,12 +133,24 @@ def _generate_tool_body(spec: dict) -> str: for name, p in params.items(): req = "yes" if p.get("required") else "no" default = p.get("default", "—") - lines.append(f"| `{name}` | {req} | {default} | {p.get('description', '')} |\n") + lines.append( + f"| `{name}` | {req} | {default} | {p.get('description', '')} |\n" + ) return "".join(lines) def _parse_frontmatter(path: Path) -> dict: - """Parse YAML frontmatter from a markdown file.""" + """Parse YAML frontmatter from a markdown file. + + Third-party skill catalogs (e.g. Genesis) ship SKILL.md files whose + frontmatter is *intended* as flat ``key: value`` but isn't strict YAML — + most commonly an unquoted ``description`` value that contains a colon + (``...readiness levels: Level 1...``), which PyYAML rejects as a nested + mapping. Rather than silently dropping such skills from discovery, fall back + to a best-effort flat parse (:func:`_lenient_frontmatter`) on YAML error so + ``name`` / ``description`` / ``tags`` are still recovered. dsagt-authored + code/skill specs are valid YAML, so the fallback never fires for them. + """ text = path.read_text() if not text.startswith("---"): return {} @@ -92,14 +160,63 @@ def _parse_frontmatter(path: Path) -> dict: try: return yaml.safe_load(parts[1]) or {} except yaml.YAMLError as e: - raise ValueError(f"Invalid YAML frontmatter in {path}: {e}") from e + # Benign: the frontmatter is flat ``key: value`` but not strict YAML; + # we recover the fields below. DEBUG, not WARNING — nothing is lost + # and it's pure noise during ``dsagt init`` catalog indexing. + logger.debug( + "Frontmatter in %s isn't strict YAML (%s); recovering flat fields.", + path, + str(e).splitlines()[0], + ) + return _lenient_frontmatter(parts[1]) + + +def _lenient_frontmatter(block: str) -> dict: + """Best-effort flat ``key: value`` parse for frontmatter that isn't strict YAML. + + Splits each top-level line on its **first** colon (so a value may itself + contain colons); indented ``- item`` lines extend the previous key into a + list, other indented lines continue the previous string value. Inline + ``[...]`` / ``{...}`` values are parsed as YAML when they can be. Lines + without a colon, and comments, are ignored. This recovers the discovery + fields (name/description/tags) from technically-invalid-but-obvious + frontmatter instead of dropping the skill. + """ + out: dict = {} + key: str | None = None + for raw in block.splitlines(): + stripped = raw.strip() + if not stripped or stripped.startswith("#"): + continue + if raw[:1].isspace() and key is not None: + # Continuation of the previous key. + if stripped.startswith("- "): + if not isinstance(out.get(key), list): + out[key] = [] + out[key].append(stripped[2:].strip()) + elif isinstance(out.get(key), str): + out[key] = (out[key] + " " + stripped).strip() + continue + if ":" not in stripped: + continue + k, _, v = stripped.partition(":") + key = k.strip() + v = v.strip() + if v.startswith(("[", "{")): + try: + out[key] = yaml.safe_load(v) + except yaml.YAMLError: + out[key] = v + else: + out[key] = v + return out # --------------------------------------------------------------------------- # CLI rendering # --------------------------------------------------------------------------- # -# Each parameter in a tool spec may declare a `cli` field that pins how its +# Each parameter in a code spec may declare a `cli` field that pins how its # value should be placed on the command line. Supported forms: # # positional — first positional slot @@ -115,6 +232,7 @@ def _parse_frontmatter(path: Path) -> dict: # Parameters with `type: boolean` render as a bare flag when truthy and emit # nothing when falsy; positional booleans are not supported. + def _parse_cli(cli: str, param_name: str) -> dict: """Classify a cli string into a rendering descriptor. Fails fast on invalid input.""" if cli == "positional": @@ -184,97 +302,87 @@ def render_arguments(parameters: dict, values: dict) -> list[str]: # --------------------------------------------------------------------------- -# Tool Registry +# Code Registry # --------------------------------------------------------------------------- -class ToolRegistry: + +class CodeRegistry: """ - Manages CLI tool spec files and optional KB indexing. - - Two layers: - - * **Bundled tools** ship with the dsagt package at - ``_PACKAGE_TOOLS_DIR``. They are read-only; their KB embeddings - live in the shared ``bundled_tools`` collection (built once per - machine per dsagt version by ``dsagt setup-kb``). Never copied - into projects, so package upgrades automatically reach all - existing projects. - * **Project tools** are agent-saved or user-edited specs in - ``/tools/``. Embeddings go into the project-local - ``registered_tools`` collection on save. - - Listing / lookup methods merge both layers (project wins on name - collision so agents can override a bundled tool). Search the - KB-side via ``search_registry`` which queries both collections. + Manages CLI code spec files and optional KB indexing. + + One layer: every code — bundled or agent-registered — is a + skill-standard directory in ``/codes//``. Bundled + codes ship with the package at ``_PACKAGE_CODES_DIR`` and are COPIED + into the project at ``dsagt init`` (:meth:`ensure_bundled_copies`), + so all available codes live in one place, in one format, fully + self-contained (spec + scripts). Re-running ``dsagt init`` after a + package upgrade refreshes unmodified copies; a user-edited copy is + never clobbered. KB-side search via ``search_registry``. """ - _PACKAGE_TOOLS_DIR = Path(__file__).parent / "tools" + _PACKAGE_CODES_DIR = Path(__file__).parent / "codes" def __init__( self, runtime_dir: str | Path, - source_tools_dir: str | None = None, kb: KnowledgeBase | None = None, ): self.runtime_dir = Path(runtime_dir) - self.tools_dir = self.runtime_dir / "tools" + self.codes_dir = self.runtime_dir / "codes" self._kb = kb self.runtime_dir.mkdir(parents=True, exist_ok=True) - # Optional override of the package-bundled directory (used by - # tests; production callers leave source_tools_dir=None and let - # the package dir stand). - self._bundled_dir = ( - Path(source_tools_dir) - if source_tools_dir and Path(source_tools_dir).exists() - else self._PACKAGE_TOOLS_DIR - ) - - # Project tool dir is always agent-writable. We no longer - # pre-populate it with bundled tools — they're served directly - # from the package via the merge in list_tools / get_tool. - self.tools_dir.mkdir(parents=True, exist_ok=True) - # Ensure code/ subdirectory exists for agent-written scripts. - (self.tools_dir / "code").mkdir(exist_ok=True) - - def _bundled_tool_paths(self) -> list[Path]: - """Return .md tool spec paths shipped with the package.""" - if not self._bundled_dir.exists(): - return [] - return sorted(self._bundled_dir.glob("*.md")) + # Each code is a self-contained skill-standard directory + # (``codes//SKILL.md`` + optional ``scripts/``), so there is + # no shared scripts/ dir to pre-create. + self.codes_dir.mkdir(parents=True, exist_ok=True) + + def ensure_bundled_copies(self) -> list[str]: + """Copy package-bundled code dirs into ``/codes/``. + + Called at ``dsagt init``. A dir whose name already exists in the + project is left untouched — user edits and agent overrides win; + delete the dir and re-init to restore the packaged version. + Returns one action line per copy made. + """ + actions: list[str] = [] + if not self._PACKAGE_CODES_DIR.exists(): + return actions + for spec in sorted(self._PACKAGE_CODES_DIR.glob("*/SKILL.md")): + dest = self.codes_dir / spec.parent.name + if dest.exists(): + continue + shutil.copytree(spec.parent, dest) + actions.append(f"Copied bundled code {spec.parent.name} into {dest}") + return actions - def _project_tool_paths(self) -> list[Path]: - """Return .md tool spec paths the agent has saved into this project.""" - return sorted(self.tools_dir.glob("*.md")) + def _project_code_paths(self) -> list[Path]: + """Return SKILL.md spec paths in this project's codes dir.""" + return sorted(self.codes_dir.glob("*/SKILL.md")) - def list_tools_raw(self) -> list[dict]: - """Return full frontmatter dicts for all tools. + def code_dirs(self) -> list[Path]: + """All code directories (for the native-skills mirror — see + ``AgentSetup.setup_skills``).""" + return [p.parent for p in self._project_code_paths()] - Merges bundled (package) + project (``/tools/``). - Project tools win on name collision so agents can override a - bundled tool with their own implementation. - """ + def list_codes_raw(self) -> list[dict]: + """Return full frontmatter dicts for all codes in the project.""" seen: dict[str, dict] = {} - for p in self._bundled_tool_paths(): + for p in self._project_code_paths(): spec = _parse_frontmatter(p) name = spec.get("name") if name: seen[name] = spec - for p in self._project_tool_paths(): - spec = _parse_frontmatter(p) - name = spec.get("name") - if name: - seen[name] = spec # project layer overrides bundled return [seen[name] for name in sorted(seen)] - def list_tools(self) -> list[dict]: - """List all tools with MCP-compatible schemas.""" - tools = [] - for tool in self.list_tools_raw(): - if not tool.get("name"): + def list_codes(self) -> list[dict]: + """List all codes with MCP-compatible schemas.""" + codes = [] + for code in self.list_codes_raw(): + if not code.get("name"): continue properties = {} required = [] - for param_name, param_def in tool.get("parameters", {}).items(): + for param_name, param_def in code.get("parameters", {}).items(): properties[param_name] = { "type": param_def.get("type", "string"), "description": param_def.get("description", ""), @@ -283,46 +391,55 @@ def list_tools(self) -> list[dict]: properties[param_name]["default"] = param_def["default"] if param_def.get("required", False): required.append(param_name) - tools.append({ - "name": tool["name"], - "description": tool["description"], - "inputSchema": { - "type": "object", - "properties": properties, - "required": required, - }, - }) - return tools - - def get_tool(self, name: str) -> dict | None: - """Look up a tool spec by name. Project layer overrides bundled.""" - project_path = self.tools_dir / f"{name}.md" - if project_path.exists(): - tool = _parse_frontmatter(project_path) - if tool.get("name") == name: - return tool - bundled_path = self._bundled_dir / f"{name}.md" - if bundled_path.exists(): - tool = _parse_frontmatter(bundled_path) - if tool.get("name") == name: - return tool + codes.append( + { + "name": code["name"], + "description": code["description"], + "inputSchema": { + "type": "object", + "properties": properties, + "required": required, + }, + } + ) + return codes + + def get_code(self, name: str) -> dict | None: + """Look up a code spec by name.""" + path = self.codes_dir / name / "SKILL.md" + if path.exists(): + code = _parse_frontmatter(path) + if code.get("name") == name: + return code return None def save_tool(self, spec: dict) -> str: - """Write or update a tool file. Returns 'added' or 'updated'. + """Write or update a code's SKILL.md. Returns 'added' or 'updated'. Automatically wraps the executable: - With `uv run --with ` if Python dependencies are specified - - With `dsagt-run --tool ` for provenance capture + - With `dsagt-run --code ` for provenance capture - If a KnowledgeBase is available, indexes the tool for semantic search. + If a KnowledgeBase is available, indexes the code for semantic search. """ - path = self.tools_dir / f"{spec['name']}.md" + # Codes share the skill-standard envelope so they mirror into agent + # native skills dirs, whose loaders require lowercase-hyphen names. + if not _CODE_NAME_RE.match(spec["name"]): + raise ValueError( + f"invalid code name {spec['name']!r}: use lowercase letters, " + "digits, and hyphens (the skill-standard charset agent native " + "skill loaders require), e.g. 'scan-directory'" + ) + code_dir = self.codes_dir / spec["name"] + path = code_dir / "SKILL.md" action = "updated" if path.exists() else "added" + code_dir.mkdir(parents=True, exist_ok=True) spec = dict(spec) spec["executable"] = _wrap_executable( - spec["name"], spec["executable"], spec.get("dependencies"), + spec["name"], + spec["executable"], + spec.get("dependencies"), ) # Preserve existing body when updating so hand-edited docs survive @@ -334,27 +451,27 @@ def save_tool(self, spec: dict) -> str: body = parts[2] if not body: - body = _generate_tool_body(spec) + body = _generate_code_body(spec) frontmatter = yaml.dump(spec, default_flow_style=False, sort_keys=False) path.write_text(f"---\n{frontmatter}---\n{body}") if self._kb: - self._index_tool(spec, path) + self._index_code(spec, path) return action - def _index_tool(self, spec: dict, tool_path: Path) -> None: - """Index a tool file into the ``tools`` KB collection. + def _index_code(self, spec: dict, tool_path: Path) -> None: + """Index a code file into the ``codes`` KB collection. - Errors propagate to the caller — a tool that lives on disk but + Errors propagate to the caller — a code that lives on disk but isn't searchable in the KB is a half-broken state that the agent cannot recover from (it would write a duplicate next time it searched). Atomic registration: in the index or not registered. """ text = tool_path.read_text() metadata = { - "tool_name": spec["name"], + "code_name": spec["name"], "tags": ",".join(spec.get("tags", [])), "executable": spec["executable"], "has_dependencies": str(bool(spec.get("dependencies"))), @@ -362,38 +479,21 @@ def _index_tool(self, spec: dict, tool_path: Path) -> None: } self._kb.add_entries( texts=[text], - collection=TOOLS_COLLECTION, + collection=CODES_COLLECTION, metadatas=[metadata], ) - def reindex_all(self) -> int: - """Reindex project-local tool files into the ``tools`` collection. - - Returns count indexed. Bundled tools are NOT indexed here — they - live in the shared ``tools`` collection built by ``dsagt setup-kb`` - and copied into the project at ``dsagt init`` time. Search via - ``search_registry`` queries the merged collection. - """ - if not self._kb: - return 0 - count = 0 - for path in self._project_tool_paths(): - spec = _parse_frontmatter(path) - if spec.get("name"): - self._index_tool(spec, path) - count += 1 - return count - # --------------------------------------------------------------------------- # Skill Registry # --------------------------------------------------------------------------- + class SkillRegistry: """ Manages instruction-based agent skills and optional KB indexing. - Two layers (mirroring ToolRegistry): + Two layers (mirroring CodeRegistry): * **Bundled skills** ship with the dsagt package at ``_PACKAGE_SKILLS_DIR``. Read-only; embeddings live in shared @@ -435,14 +535,16 @@ def _bundled_skill_dirs(self) -> list[Path]: if not self._bundled_dir.exists(): return [] return [ - d for d in sorted(self._bundled_dir.iterdir()) + d + for d in sorted(self._bundled_dir.iterdir()) if d.is_dir() and (d / "SKILL.md").exists() ] def _project_skill_dirs(self) -> list[Path]: """Return skill directories the agent has saved into this project.""" return [ - d for d in sorted(self.skills_dir.iterdir()) + d + for d in sorted(self.skills_dir.iterdir()) if d.is_dir() and (d / "SKILL.md").exists() ] @@ -474,9 +576,11 @@ def save_skill( ``{relative_path: contents}`` for additional files the skill wants in its directory (templates, schemas, etc.). - Returns "added" or "updated". Indexes the resulting SKILL.md - into ``registered_skills`` via ``_index_skill`` if a KB is - configured — symmetric with ``ToolRegistry.save_tool``. + Returns "added" or "updated". Does **not** index into a KB: + saved skills land in ``/skills/`` where every supported + agent natively auto-discovers them, so search only covers the + not-yet-installed *catalog* tier (see ``SkillRouter``). The old + ``skills`` collection is no longer read by anything. """ name = spec.get("name") if not name: @@ -488,7 +592,7 @@ def save_skill( skill_md = skill_dir / "SKILL.md" # Preserve hand-edited body when updating, unless caller passed - # an explicit replacement — same contract as ToolRegistry.save_tool. + # an explicit replacement — same contract as CodeRegistry.save_tool. if body is None and skill_md.exists(): existing = skill_md.read_text() parts = existing.split("---", 2) @@ -505,9 +609,6 @@ def save_skill( target.parent.mkdir(parents=True, exist_ok=True) target.write_text(contents) - if self._kb: - self._index_skill(spec, skill_md) - return action def _skill_md_path(self, name: str) -> Path | None: @@ -532,38 +633,3 @@ def get_skill_content(self, name: str) -> str | None: """Get the full SKILL.md content for a skill.""" path = self._skill_md_path(name) return path.read_text() if path is not None else None - - def _index_skill(self, spec: dict, skill_md: Path) -> None: - """Index a skill into the ``skills`` KB collection. - - Errors propagate to the caller — see _index_tool for the rationale. - """ - text = skill_md.read_text() - metadata = { - "skill_name": spec["name"], - "tags": ",".join(spec.get("tags", [])), - "source": "registered", # vs "bundled" - } - self._kb.add_entries( - texts=[text], - collection=SKILLS_COLLECTION, - metadatas=[metadata], - ) - - def reindex_all(self) -> int: - """Reindex project-local skills into ``registered_skills``. - - Bundled skills are NOT indexed here — they live in the shared - ``bundled_skills`` collection (see ToolRegistry.reindex_all - docstring for the same architecture). - """ - if not self._kb: - return 0 - count = 0 - for skill_dir in self._project_skill_dirs(): - skill_md = skill_dir / "SKILL.md" - spec = _parse_frontmatter(skill_md) - if spec.get("name"): - self._index_skill(spec, skill_md) - count += 1 - return count diff --git a/src/dsagt/session.py b/src/dsagt/session.py index 6ae19b3..e1d630d 100644 --- a/src/dsagt/session.py +++ b/src/dsagt/session.py @@ -3,38 +3,33 @@ Projects are registered in ~/dsagt-projects/projects.yaml (name → absolute path). Default project location is ~/dsagt-projects//. Shared bundled-content -KB lives alongside at ~/dsagt-projects/kb_index/ (built by dsagt setup-kb). +KB lives alongside at ~/dsagt-projects/kb_index/ (provisioned by dsagt init). Project directory layout:: / - dsagt_config.yaml # project configuration + .dsagt/ + config.yaml # project configuration (MCP-server object settings) + state.yaml # session log + memory cursor (owned by the MCP server) + explicit_memories.yaml, ... # explicit memory trace_archive/ # tool execution records - mlflow/ # MLflow data - tools/ # registered CLI tools - tools/code/ # agent-written tool scripts + mlflow.db # serverless MLflow SQLite trace store (created lazily) + codes// # registered codes (skill-standard dirs: SKILL.md + scripts/) skills/ # instruction-based agent skills kb_index/ # knowledge base collections """ -import json import logging import os import re import shutil -import signal -import socket -import subprocess -import sys -import time from datetime import datetime, timezone from pathlib import Path import yaml -from dsagt.memory import extract_session from dsagt.knowledge import KnowledgeBase -from dsagt.provenance import index_trace_archive +from dsagt.provenance import CodeUseIndexer logger = logging.getLogger(__name__) @@ -43,65 +38,72 @@ # Constants # --------------------------------------------------------------------------- -VALID_AGENTS = ("claude", "goose", "roo", "cline", "codex", "opencode") -VALID_MLFLOW_BACKENDS = ("sqlite", "flat-file") +VALID_AGENTS = ("claude", "goose", "cline", "codex", "opencode") DEFAULT_PROJECTS_BASE = Path.home() / "dsagt-projects" # Registry + shared KB live alongside projects under one visible tree — # ``~/dsagt-projects/projects.yaml`` (name → path) and -# ``~/dsagt-projects/kb_index/`` (shared bundled-content KB built by -# ``dsagt setup-kb``). Migrated from ``~/.dsagt/`` on 2026-05-07. +# ``~/dsagt-projects/kb_index/`` (shared bundled-content KB provisioned by +# ``dsagt init``). Migrated from ``~/.dsagt/`` on 2026-05-07. REGISTRY_DIR = DEFAULT_PROJECTS_BASE REGISTRY_FILE = REGISTRY_DIR / "projects.yaml" -RESERVED_PROJECT_NAMES = ("projects.yaml", "kb_index") +RESERVED_PROJECT_NAMES = ("projects.yaml", "kb_index", ".skill_sources") +# Per-project dsagt state lives under a hidden ``.dsagt/`` dir (alongside +# explicit memory): ``config.yaml`` (the MCP-server object settings the user +# chose at ``dsagt init``) and ``state.yaml`` (session log + memory cursor, +# owned by the MCP server). +CONFIG_DIRNAME = ".dsagt" +CONFIG_FILENAME = "config.yaml" +STATE_FILENAME = "state.yaml" + + +def config_path(pdir: Path) -> Path: + """Path to a project's ``.dsagt/config.yaml``.""" + return Path(pdir) / CONFIG_DIRNAME / CONFIG_FILENAME + + +def state_path(pdir: Path) -> Path: + """Path to a project's ``.dsagt/state.yaml``.""" + return Path(pdir) / CONFIG_DIRNAME / STATE_FILENAME + + +# Code defaults backfilled into a project's config on read (``_deep_merge``). +# These are NOT user choices and so are NOT written to ``.dsagt/config.yaml`` +# nor prompted at ``dsagt init`` — they live here as the single source of +# truth and are filled in for the MCP server / KB. Embedding is local-only +# for now (BYOA, no credentials); an ``api`` backend can be re-introduced as +# an init choice if it's ever requested. ``chunk_size`` / ``rerank`` default +# in :class:`~dsagt.knowledge.KnowledgeBase`; ``skills.populate_native`` in +# :meth:`AgentSetup.setup_skills`. DEFAULTS = { - # ``llm`` block uses ${VAR} placeholders so per-project config - # references the user's shell env at resolve time. In BYOA mode - # (the default), agent_env's proxy_port gate at agents/__init__.py - # short-circuits the env_overrides call, so this block sits dormant - # — no agent env-var leaks. In proxy mode (--enable-proxy), - # env_overrides reads these values to translate into per-agent env - # vars and proxy_server.py reads them to render its YAML. - "llm": { - "provider": "${LLM_PROVIDER}", - "model": "${LLM_MODEL}", - "base_url": "${LLM_BASE_URL}", - "api_key": "${LLM_API_KEY}", - }, "embedding": { - # Default backend: "local" — sentence-transformers, CPU-side, no - # API credentials needed. Switch to "api" to route through - # litellm to a hosted endpoint (then fill in base_url / api_key - # below and pick a model name your endpoint serves). - # Local default is the HuggingFace identifier ``LocalEmbeddingClient`` - # downloads; user can override with any sentence-transformers - # repo (e.g. ``BAAI/bge-large-en-v1.5`` for higher quality). "backend": "local", "model": "BAAI/bge-small-en-v1.5", "base_url": "", }, - "mlflow": { - "backend": "sqlite", + # External agent-skill catalogs. ``sources`` are GitHub repos whose + # SKILL.md skills get indexed into per-source catalog collections for + # ``search_skills``. This is the default when a project picks none. + "skills": { + "sources": [ + { + "name": "genesis", + "url": "https://gitlab.osti.gov/genesis/genesis-skills", + "branch": "main", + "subdir": "skills", + }, + ], }, - "categories": { - "quality_control": "Assessment or filtering of data quality, QC metrics, thresholds, pass/fail rates", - "data_management": "File organization, data movement, format conversion, naming conventions", - "transformation": "Data processing steps, parameter choices, pipeline stage configuration", - "assembly": "Genome assembly, contig generation, scaffolding, assembly QC metrics", - "configuration": "Tool settings, environment setup, resource allocation decisions", - "performance": "Runtime, memory usage, throughput, resource consumption observations", - "tool_usage": "Tool selection rationale, parameter tuning, tool-specific behaviors or quirks", - "results": "Output summaries, key findings, deliverables produced", - }, - "extraction": { - "threshold": 0, - "outlier_sensitivity": 0.0, - }, - "knowledge": { - "chunk_size": 1024, - "vector_db": "chroma", - "rerank": False, + # Episodic memory: the heartbeat's MemoryExtractor subscriber. ``enabled`` + # is a compute/storage opt-in that mechanically chunks/tags/embeds each + # completed turn into session_memory (no credentials). + "episodic": { + "enabled": False, + # Recency weighting for session_memory retrieval: a newer turn edges out + # a stale one without contradiction detection. Half-life in days (a + # *boost*, never a penalty — durable old turns keep full relevance). + "recency_half_life_days": 14, }, } @@ -112,6 +114,7 @@ # Config helpers # --------------------------------------------------------------------------- + def resolve_env_vars(value): """Replace ${VAR_NAME} references with environment variable values.""" if isinstance(value, str): @@ -134,30 +137,54 @@ def _deep_merge(defaults: dict, overrides: dict) -> dict: return result -def default_config_content( +def build_config( project_name: str, agent: str, - mlflow_port: int, -) -> str: - """Generate the internal dsagt_config.yaml for a new project. - - Internal-only: users don't edit this. Holds the project's - embedding / mlflow / knowledge / extraction settings plus the - pinned MLflow port so MCP servers (started fresh per agent run) - know where to log without relying on shell-env inheritance. - - User credentials are NOT here — the agent reads them from the - user's shell env directly (ANTHROPIC_API_KEY, OPENAI_API_KEY, etc). + *, + knowledge: dict | None = None, + skills: dict | None = None, + episodic: dict | None = None, +) -> dict: + """Assemble a project's ``.dsagt/config.yaml`` body. + + The schema is a strict 1:1 mirror of the choices ``dsagt init`` offers — + every key here corresponds to an init prompt and vice versa: + + - ``project`` / ``agent`` — identity (agent + name/location are prompted). + - ``knowledge.collections`` — the packaged document collections chosen + (default none; the bundled ``tools`` collection is always provisioned + and is not a per-project choice). + - ``skills.sources`` — the skill-catalog repos chosen. + - ``episodic`` — written *only when the user opted in* (it's an opt-in, so a + disabled project stays minimal and backfills ``enabled: false`` on read). + + Everything else (embedding backend, chunk_size, rerank, populate_native) + is a code default backfilled on read — NOT a written choice. Credentials + are never here (shell env only); no MLflow port (serverless sqlite store). """ - body: dict = { + body = { "project": project_name, "agent": agent, - "mlflow": {**DEFAULTS["mlflow"], "port": mlflow_port}, - "embedding": dict(DEFAULTS["embedding"]), - "knowledge": DEFAULTS["knowledge"], - "categories": DEFAULTS["categories"], - "extraction": DEFAULTS["extraction"], + "knowledge": knowledge or {"collections": []}, + "skills": skills or {"sources": list(DEFAULTS["skills"]["sources"])}, } + if episodic: + body["episodic"] = episodic + return body + + +def default_config_content( + project_name: str, + agent: str, + *, + knowledge: dict | None = None, + skills: dict | None = None, + episodic: dict | None = None, +) -> str: + """Serialize :func:`build_config` to YAML for ``.dsagt/config.yaml``.""" + body = build_config( + project_name, agent, knowledge=knowledge, skills=skills, episodic=episodic + ) return yaml.dump(body, default_flow_style=False, sort_keys=False) @@ -165,6 +192,7 @@ def default_config_content( # Project registry # --------------------------------------------------------------------------- + def _load_registry() -> dict[str, str]: """Load the project registry. Returns empty dict if no registry exists.""" if not REGISTRY_FILE.exists(): @@ -190,6 +218,47 @@ def list_projects() -> dict[str, str]: return _load_registry() +def kb_from_config(config: dict, index_dir: Path | None = None) -> "KnowledgeBase": + """Build a KnowledgeBase from a resolved project config. + + Resolves the project's embedding backend so callers (CLI ``skills`` group, + catalog sync) get a KB wired to the same embedder the project uses. + Defaults to ``/kb_index``. + """ + pdir = Path(config["project_dir"]) + emb = config.get("embedding", {}) + backend = emb.get("backend", "local") + recency = _recency_half_life(config) + if backend == "local": + model = emb.get("model") + if model and "/" not in str(model): + model = None + return KnowledgeBase( + index_dir=index_dir or (pdir / "kb_index"), + default_embedder=backend, + model=model, + recency_half_life_days=recency, + ) + return KnowledgeBase( + index_dir=index_dir or (pdir / "kb_index"), + default_embedder=backend, + model=emb.get("model"), + base_url=emb.get("base_url"), + api_key=os.environ.get("EMBEDDING_API_KEY", ""), + recency_half_life_days=recency, + ) + + +def _recency_half_life(config: dict) -> float | None: + """Episodic recency half-life (days) when enabled, else ``None`` (off). + + Recency weighting only matters for ``session_memory``, which only has + content when episodic memory is enabled — so it's gated on that opt-in. + """ + epi = config.get("episodic", {}) or {} + return epi.get("recency_half_life_days") if epi.get("enabled") else None + + def project_dir(name: str) -> Path: """Resolve a project name to its directory via the registry.""" registry = _load_registry() @@ -207,6 +276,7 @@ def project_dir(name: str) -> Path: # Config loading # --------------------------------------------------------------------------- + def load_config(project_name: str) -> dict: """Load and validate a project config by name. @@ -214,12 +284,12 @@ def load_config(project_name: str) -> dict: resolved config dict with defaults applied and 'project_dir' injected. """ pdir = project_dir(project_name) - config_path = pdir / "dsagt_config.yaml" + cfg_file = config_path(pdir) - if not config_path.exists(): - raise FileNotFoundError(f"Config not found: {config_path}") + if not cfg_file.exists(): + raise FileNotFoundError(f"Config not found: {cfg_file}") - raw = yaml.safe_load(config_path.read_text()) or {} + raw = yaml.safe_load(cfg_file.read_text()) or {} config = _deep_merge(DEFAULTS, raw) config = resolve_env_vars(config) @@ -229,45 +299,157 @@ def load_config(project_name: str) -> dict: return config +def read_config_file(pdir: Path) -> dict: + """Read a project's raw ``.dsagt/config.yaml`` by path (no registry, no + defaults merge). Returns ``{}`` if absent — used to prefill the + re-run ``dsagt init`` dialogue with current values. + """ + cfg_file = config_path(pdir) + if not cfg_file.exists(): + return {} + return yaml.safe_load(cfg_file.read_text()) or {} + + +def write_config_file(pdir: Path, config: dict) -> None: + """Write a project's ``.dsagt/config.yaml`` (creates ``.dsagt/`` if + needed). ``config`` is the trimmed schema from :func:`build_config`. + """ + cfg_file = config_path(pdir) + cfg_file.parent.mkdir(parents=True, exist_ok=True) + cfg_file.write_text(yaml.dump(config, default_flow_style=False, sort_keys=False)) + + def _validate(config: dict) -> None: """Validate required fields and values.""" if not config.get("project"): - raise ValueError("'project' is required in dsagt_config.yaml") + raise ValueError("'project' is required in .dsagt/config.yaml") agent = config.get("agent") if not agent: - raise ValueError("'agent' is required in dsagt_config.yaml") + raise ValueError("'agent' is required in .dsagt/config.yaml") if agent not in VALID_AGENTS: raise ValueError(f"'agent' must be one of {VALID_AGENTS}, got '{agent}'") - backend = config.get("mlflow", {}).get("backend") - if backend and backend not in VALID_MLFLOW_BACKENDS: - raise ValueError(f"'mlflow.backend' must be one of {VALID_MLFLOW_BACKENDS}, got '{backend}'") + +# --------------------------------------------------------------------------- +# Session state (`.dsagt/state.yaml`) — owned by the MCP server +# --------------------------------------------------------------------------- + + +def _empty_state() -> dict: + return {"sessions": [], "memory_cursor": {}} + + +def read_state(pdir: Path) -> dict: + """Read ``.dsagt/state.yaml``; return an empty skeleton if absent.""" + sp = state_path(pdir) + if not sp.exists(): + return _empty_state() + return yaml.safe_load(sp.read_text()) or _empty_state() + + +def write_state(pdir: Path, state: dict) -> None: + """Write ``.dsagt/state.yaml`` (creates ``.dsagt/`` if needed).""" + sp = state_path(pdir) + sp.parent.mkdir(parents=True, exist_ok=True) + sp.write_text(yaml.dump(state, default_flow_style=False, sort_keys=False)) + + +def append_session(pdir: Path) -> dict: + """Append a new session entry and return it. + + Called by the MCP server at startup — the server owns session-id + minting now (not ``dsagt start``), so a bare-launched agent gets a + session id too. Id is a monotonic per-project counter. + """ + state = read_state(pdir) + sessions = state.setdefault("sessions", []) + next_id = max((s.get("id", 0) for s in sessions), default=0) + 1 + entry = { + "id": next_id, + "started_at": datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ"), + } + sessions.append(entry) + state.setdefault("memory_cursor", {}) + write_state(pdir, state) + return entry + + +def current_session(pdir: Path) -> dict | None: + """The most recent session entry, or ``None`` if no session has run.""" + sessions = read_state(pdir).get("sessions") or [] + return sessions[-1] if sessions else None + + +def session_tag(project: str, session_id: int) -> str: + """The trace-correlation tag for a session: ``-``.""" + return f"{project}-{session_id}" + + +def current_session_tag(pdir: Path, project: str) -> str | None: + """The trace tag of the current session, or ``None`` if none exists. + + Used by ``dsagt-run`` to tag tool spans with the same session the MCP + server minted — read from ``state.yaml`` instead of a ``DSAGT_SESSION_ID`` + env var. + """ + cur = current_session(pdir) + if cur is None: + return None + return session_tag(project, cur["id"]) + + +def update_cursor(pdir: Path, **fields) -> None: + """Merge fields into ``state.yaml``'s ``memory_cursor`` map.""" + state = read_state(pdir) + state.setdefault("memory_cursor", {}).update(fields) + write_state(pdir, state) + + +def record_trace_source(pdir: Path, source) -> None: + """Stamp the current session's trace-source token into ``state.yaml``. + + ``source`` is the agent-shaped session token from + :meth:`dsagt.traces.Reader.active_source` (a transcript path, DB session id, + or session-dir name). The MCP server records it once the live collector + resolves a session, so the *next* session's catch-up can pin this exact + session — for *every* agent — when re-collecting turns an ungraceful + shutdown left unlogged. No-op if no session has been minted yet. + """ + state = read_state(pdir) + sessions = state.get("sessions") or [] + if not sessions: + return + if sessions[-1].get("trace_source") == source: + return # already recorded — avoid churning the file + sessions[-1]["trace_source"] = source + write_state(pdir, state) # --------------------------------------------------------------------------- # Project initialization # --------------------------------------------------------------------------- + def _collection_exists(path: Path) -> bool: """Return True if *path* looks like a persisted KB collection directory. - Accepts FAISS indexes, ChromaDB-in-a-dir layouts, routed external - collections, and bare ChromaDB sqlite files (as produced by - ``dsagt setup-kb`` for description-only collections). + Accepts ChromaDB-in-a-dir layouts, routed external collections, and + bare ChromaDB sqlite files (as produced by the KB asset build for + description-only collections). """ - return ( - path.is_dir() - and ( - (path / "index.faiss").exists() - or (path / "chroma_ids.json").exists() - or (path / "route.json").exists() - or (path / "chroma.sqlite3").exists() - ) + return path.is_dir() and ( + (path / "chroma_ids.json").exists() + or (path / "route.json").exists() + or (path / "chroma.sqlite3").exists() ) -def setup_runtime_kb(base_index_dir: Path, runtime_dir: Path) -> Path: +def setup_runtime_kb( + base_index_dir: Path, + runtime_dir: Path, + collections: list[str] | None = None, +) -> Path: """Copy base KB collections into a project's kb_index directory. Creates ``/kb_index`` if missing. For each collection @@ -275,9 +457,15 @@ def setup_runtime_kb(base_index_dir: Path, runtime_dir: Path) -> Path: twin doesn't already exist, **copies** (not symlinks) the entire collection directory into the project's kb_index. + *collections*, when given, is an allowlist of collection-directory + names to copy — so a project gets exactly its requested asset set even + when the shared KB holds more (e.g. heavy collections another project + installed). ``None`` copies every populated collection (legacy + copy-everything behavior). + Why copy instead of symlink: different projects on the same machine may run different dsagt versions, and a symlink would let one - project's ``dsagt setup-kb --rebuild`` mutate every project's view + project's KB rebuild mutate every project's view of bundled content. A copy pins each project to whatever bundled content was current when the project first ran. @@ -290,7 +478,10 @@ def setup_runtime_kb(base_index_dir: Path, runtime_dir: Path) -> Path: if not base_index_dir.exists(): return runtime_kb_dir + allow = set(collections) if collections is not None else None for collection_dir in base_index_dir.iterdir(): + if allow is not None and collection_dir.name not in allow: + continue if not _collection_exists(collection_dir): continue dest = runtime_kb_dir / collection_dir.name @@ -303,21 +494,119 @@ def setup_runtime_kb(base_index_dir: Path, runtime_dir: Path) -> Path: return runtime_kb_dir +def _provision_kb( + pdir: Path, + include: list[str] | None, + exclude: list[str] | None, + embedding: dict | None = None, +) -> None: + """Build the requested KB assets into the shared cache, then copy that + set into the project. + + The first project on a machine pays the one-time build (bundled tools + + genesis catalog by default); later projects just copy. The copy is + scoped to the requested set, so a project gets exactly what was asked + for regardless of what else the shared cache holds. + + Best-effort: a build failure (offline, no embedding model) degrades to + an empty-but-valid project KB with a warning rather than aborting + ``dsagt init`` — the build retries on a later ``dsagt init``. + """ + from dsagt.commands.setup_core_kb import ( + asset_collection_name, + ensure_assets, + resolve_assets, + ) + + assets = resolve_assets(include=include, exclude=exclude) + if not assets: + # ``--exclude all``: a valid project with an empty KB. + (pdir / "kb_index").mkdir(parents=True, exist_ok=True) + return + + shared = REGISTRY_DIR / "kb_index" + first_ever = not shared.exists() or not any( + _collection_exists(c) for c in shared.iterdir() + ) + # Which requested assets aren't in the shared cache yet — i.e. what this + # init will actually build (and narrate). Empty → silent fast path. + pending = [ + a for a in assets if not _collection_exists(shared / asset_collection_name(a)) + ] + if pending: + if first_ever: + print( + "Performing initial dsagt setup — first project on this " + "machine (one-time, may take a few minutes):", + flush=True, + ) + else: + print("Provisioning knowledge base assets:", flush=True) + + emb = embedding or DEFAULTS["embedding"] + backend = emb.get("backend", "local") + if backend == "local": + embedder_kwargs = {"model": emb.get("model")} + else: + embedder_kwargs = { + "model": emb.get("model"), + "base_url": emb.get("base_url"), + "api_key": os.environ.get("EMBEDDING_API_KEY", ""), + } + + try: + ensure_assets( + assets, + shared, + embedding_backend=backend, + embedder_kwargs=embedder_kwargs, + ) + except Exception as e: # noqa: BLE001 — never let a build failure block init + print( + f" Warning: could not build knowledge base ({e}). The project " + "works without it; re-run `dsagt init` for a project once a " + "network / embedding backend is available to install the shared " + "core KB.", + flush=True, + ) + + wanted = [asset_collection_name(a) for a in assets] + setup_runtime_kb(shared, pdir, collections=wanted) + + if pending: + print(" Knowledge base ready.", flush=True) + + def init_project( project_name: str, agent: str, - mlflow_port: int | None = None, location: Path | None = None, -) -> tuple[Path, int]: - """Create a new project directory with default config and subdirectories. - - BYOA model: at init we lay down everything the user needs to point - their own agent process at our MCP servers. Picks (or honors) an - MLflow port and writes it to the internal ``dsagt_config.yaml`` so - later ``dsagt mlflow `` and the MCP-server children all - agree on where traces land. - - Returns ``(project_dir, mlflow_port)``. + include: list[str] | None = None, + exclude: list[str] | None = None, + *, + embedding: dict | None = None, + knowledge: dict | None = None, + skills: dict | None = None, + episodic: dict | None = None, +) -> Path: + """Create or reconfigure a project — ``dsagt init`` is re-runnable. + + BYOA model: we lay down everything the user needs to point their own + agent process at our MCP server. The trace store is serverless + (``sqlite:////mlflow.db``), so there's no port to pick — the + MCP-server children resolve the store from the project dir. + + Idempotent: on a project that already has ``.dsagt/config.yaml`` this + overwrites the config with the new choices and provisions any + newly-requested KB assets. It never deletes agent-saved data; the + caller (``dsagt init`` in ``cli.py``) handles destructive deltas + (agent switch, removed collections) with explicit per-change prompts. + + Knowledge base: provisioned with a chosen set of KB assets + (``include`` / ``exclude``, default = bundled tools + genesis catalog), + built once into the shared ``~/dsagt-projects/kb_index/`` and copied in. + + Returns the project directory. """ if agent not in VALID_AGENTS: raise ValueError(f"agent must be one of {VALID_AGENTS}, got '{agent}'") @@ -330,67 +619,45 @@ def init_project( pdir = (location or DEFAULT_PROJECTS_BASE) / project_name - if (pdir / "dsagt_config.yaml").exists(): - raise FileExistsError(f"Project already exists: {pdir}") - pdir.mkdir(parents=True, exist_ok=True) - # `tools/` and `tools/code/` are created by ToolRegistry on first server - # startup so bundled tools get copied in (it short-circuits if tools/ - # already exists). - for subdir in ("trace_archive", "mlflow", "skills", ".dsagt"): + # ``mlflow.db`` is created lazily by the MLflow client on first span. + for subdir in ("trace_archive", "skills", CONFIG_DIRNAME): (pdir / subdir).mkdir(parents=True, exist_ok=True) - setup_runtime_kb(REGISTRY_DIR / "kb_index", pdir) + # Bundled codes are copied into /codes/ so every available + # code lives in one place, in one format (skill-standard dirs), fully + # self-contained. Re-init after a package upgrade refreshes copies + # the user hasn't touched; edited/overridden dirs are never clobbered. + from dsagt.registry import CodeRegistry - # If the shared KB hasn't been built yet, warn — the project's - # kb_index/ will be empty until ``dsagt setup-kb`` runs. Don't - # rebuild here: that conflicts with the contract that ``dsagt - # init`` does no embedding work. - shared_kb = REGISTRY_DIR / "kb_index" - if not shared_kb.exists() or not any( - _collection_exists(c) for c in shared_kb.iterdir() - ): - print( - f" Warning: shared KB at {shared_kb} is empty — " - "run `dsagt setup-kb` to build bundled tools and skills " - "before launching your agent.", - flush=True, - ) + CodeRegistry(runtime_dir=pdir).ensure_bundled_copies() - if mlflow_port is None: - mlflow_port = pick_free_port() + _provision_kb(pdir, include, exclude, embedding=embedding) - (pdir / "dsagt_config.yaml").write_text( - default_config_content(project_name, agent, mlflow_port) + write_config_file( + pdir, + build_config( + project_name, agent, knowledge=knowledge, skills=skills, episodic=episodic + ), ) register_project(project_name, pdir) - return pdir, mlflow_port + return pdir -def persist_agent_choice(project_name: str, agent: str) -> None: - """Add or update the ``agent:`` field in the project's YAML. +def remove_collection(pdir: Path, collection: str) -> bool: + """Delete a single KB collection directory from a project's ``kb_index``. - Called by ``dsagt start`` on the first run that resolves an agent - from ``--agent`` when the YAML didn't have one — so the next start - doesn't need the flag. Subsequent ``--agent`` overrides at start - are per-run only and don't touch the YAML. + Used by re-run ``dsagt init`` when the user opts to remove a collection + that was dropped from the asset set. Returns True if a directory was + removed. Caller must guard agent-populated collections (``tool_use``, + ``session_memory``) — this helper deletes whatever name it's given. """ - if agent not in VALID_AGENTS: - raise ValueError(f"agent must be one of {VALID_AGENTS}, got '{agent}'") - pdir = project_dir(project_name) - yaml_path = pdir / "dsagt_config.yaml" - raw = yaml.safe_load(yaml_path.read_text()) or {} - raw["agent"] = agent - # Preserve the comment header from default_config_content so - # readers still see the provider hint. Cheap to re-emit. - header = ( - "# llm.provider: LiteLLM provider prefix (selects request format + auth).\n" - "# Common: openai, anthropic, bedrock, vertex_ai, azure, gemini,\n" - "# ollama, mistral, groq, deepseek.\n" - "# Full list: https://docs.litellm.ai/docs/providers\n" - ) - yaml_path.write_text(header + yaml.dump(raw, default_flow_style=False, sort_keys=False)) + target = Path(pdir) / "kb_index" / collection + if target.exists(): + shutil.rmtree(target) + return True + return False def move_project(project_name: str, new_location: Path) -> Path: @@ -410,17 +677,12 @@ def move_project(project_name: str, new_location: Path) -> Path: def remove_project(project_name: str, keep_files: bool = False) -> Path: """Unregister a project. By default also deletes the project directory. - Raises RuntimeError if the project's .runtime file is present (services - still running) — stop the session first, or delete .runtime if stale. + Serverless: there are no background services to reap — ``dsagt start`` + runs the agent in the foreground and returns when it exits — so there + is nothing to stop before removing. """ pdir = project_dir(project_name) - if (pdir / ".runtime").exists(): - raise RuntimeError( - f"Project '{project_name}' has a .runtime file — services may still be " - f"running. Stop the session first, or remove {pdir / '.runtime'} if stale." - ) - if not keep_files and pdir.exists(): shutil.rmtree(pdir) @@ -434,6 +696,7 @@ def remove_project(project_name: str, keep_files: bool = False) -> Path: # Service start / stop # --------------------------------------------------------------------------- + def _embedding_provider(config: dict) -> str: """Resolve embedding provider with a fallback for two cases: @@ -448,435 +711,101 @@ def _embedding_provider(config: dict) -> str: # --------------------------------------------------------------------------- -# Service supervision -# -# Each ``dsagt start`` writes ``/.runtime`` containing the random -# port MLflow bound + its pid. The next start (or ``dsagt stop``) reads -# that file and reaps anything still alive whose command line still names -# ``mlflow`` — see ``reap_runtime``. Random ports + a project-local -# state file means we never have to ask "is this listener on port 5000 -# mine?" — the file IS the answer. +# Memory extraction orchestration # --------------------------------------------------------------------------- -#: Seconds to wait for SIGTERM-ed processes to exit before SIGKILL. Long -#: enough for uvicorn + mlflow graceful shutdown (a few seconds), short -#: enough that an unresponsive process doesn't drag teardown out forever. -_STOP_GRACE_SECONDS = 5 - -def mlflow_command(pdir: Path, mlflow_config: dict, port: int) -> list[str]: - """Build the argv for launching MLflow against a project's store. - - ``--workers 1``: dsagt is a single-user dev tool — the agent makes - serial LLM calls and MCP-server spans are low-volume. Default - workers=4 each spin up a fresh Python process re-importing MLflow's - full surface (fastapi, sqlalchemy, alembic), so dropping to 1 - shaves ~0.5s off startup with zero observable cost on this load. +def catch_up_extraction(pdir: Path, config: dict) -> dict: + """Background post-session catch-up — run by the MCP server at startup. + + The MCP server owns the session lifecycle now: each launch, it spawns + this against a snapshot taken at startup, so it processes the *previous* + session's trailing trace records, never the live one. This removes the + need for a reliable session-*end* trigger (``dsagt start`` no longer runs + any extraction) and gives bare-launched agents full parity. + + Two phases, both best-effort: + + 1. **Tool-execution indexing** (always): embed the previous session's + ``/trace_archive/`` records into the ``tool_use`` collection via + the shared :class:`~dsagt.provenance.CodeUseIndexer` — idempotent against + the same ``.dsagt/code_use_acks.json`` the live heartbeat uses, so the + startup catch-up and the heartbeat never double-index. No LLM, no + credentials (local-backend default). + 2. **Chat-trace catch-up** (:func:`_catch_up_traces`): re-collect the + previous session so any turns the heartbeat missed before an ungraceful + shutdown still reach MLflow (and episodic memory). Pinned to the + trace-source token recorded in ``state.yaml`` (uniform across agents); + session-qualified acks dedupe against the live pass, so only dangling + turns emit. """ - mlflow_dir = pdir / "mlflow" - mlflow_dir.mkdir(exist_ok=True) - backend_uri = ( - f"sqlite:///{mlflow_dir}/mlflow.db" - if mlflow_config.get("backend") == "sqlite" - else str(mlflow_dir) - ) - return [ - sys.executable, "-m", "mlflow", "server", - "--backend-store-uri", backend_uri, - "--default-artifact-root", str(mlflow_dir / "artifacts"), - "--host", "0.0.0.0", - "--port", str(port), - "--workers", "1", - ] + pdir = Path(pdir) + config = {**config, "project_dir": str(pdir)} + kb = kb_from_config(config) - -def pick_free_port() -> int: - """Bind ``("", 0)`` so the kernel assigns a free port, then release. - - There's a microsecond race between this returning and the subprocess - binding the same port — acceptable on a single-user dev machine. If - the subprocess fails to bind, the mlflow.log tail surfaces the error - via ``_wait_for_mlflow``. - """ - with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: - s.bind(("", 0)) - return s.getsockname()[1] - - -def _process_command(pid: int) -> str: - """Return the cmdline for *pid*, or ``""`` if dead/unreadable.""" try: - result = subprocess.run( - ["ps", "-p", str(pid), "-o", "command="], - capture_output=True, text=True, check=False, timeout=2.0, - ) - except (FileNotFoundError, subprocess.TimeoutExpired): - return "" - return result.stdout.strip() - - -def reap_runtime(runtime_file: Path) -> list[str]: - """SIGTERM, grace-wait, SIGKILL any live PIDs in *runtime_file*. - - PID-recycling guard: between writing the PID and reading it back, the - OS could have reassigned that PID to another process; we only signal - if its cmdline still names what we started (``mlflow`` / ``proxy``). - - Used by ``start_services`` to clear leftovers from a prior crashed - run, and by ``stop_services`` (user-invoked teardown). Idempotent — - safe when the file is missing or PIDs are already dead. - """ - if not runtime_file.exists(): - return [] - state = json.loads(runtime_file.read_text()) - pending: dict[str, tuple[int, int]] = {} # name -> (pid, pgid) - for name, pid in state.get("pids", {}).items(): - if name not in _process_command(pid): - continue - try: - pgid = os.getpgid(pid) - os.killpg(pgid, signal.SIGTERM) - pending[name] = (pid, pgid) - except (ProcessLookupError, PermissionError): - pass - - stopped: list[str] = [] - deadline = time.monotonic() + _STOP_GRACE_SECONDS - while pending and time.monotonic() < deadline: - time.sleep(0.2) - for name in list(pending): - pid, pgid = pending[name] - try: - os.killpg(pgid, 0) # liveness probe - except (ProcessLookupError, PermissionError): - stopped.append(f"Stopped {name} (pid {pid})") - del pending[name] - - for name, (pid, pgid) in pending.items(): - try: - os.killpg(pgid, signal.SIGKILL) - stopped.append(f"Stopped {name} (pid {pid}, SIGKILL after {_STOP_GRACE_SECONDS}s)") - except ProcessLookupError: - stopped.append(f"Stopped {name} (pid {pid})") - - runtime_file.unlink(missing_ok=True) - return stopped - - -def start_services(config: dict) -> dict[str, int]: - """Start MLflow, optionally start the dsagt-proxy too. - - Picks free ports (or honors pre-set ones), reaps any leftovers from a - prior crashed run, writes ``/.runtime`` (pids + ports + - start time), and waits for each service to accept connections. - - Returns ``{"mlflow": port, "proxy": port}`` when the proxy was - requested (``"proxy"`` key present in config), else just - ``{"mlflow": port}``. - - The proxy is opt-in via ``dsagt start --enable-proxy``. Used for - agents that don't natively emit OTel traces with full LLM-call - payloads (cline, roo, codex partial) — the proxy interposes on their - LLM calls and produces traces on their behalf via the same OTLP path - Claude Code uses natively. See ``commands/proxy_server.py``. - """ - pdir = Path(config["project_dir"]) - runtime_file = pdir / ".runtime" - - reap_runtime(runtime_file) # clear leftovers from any prior crashed run - - # KB bootstrap is intentionally NOT here. The contract: - # * ``dsagt setup-kb`` builds shared ~/dsagt-projects/kb_index/ (one-time - # per machine, the only place embedding work happens for bundled - # content) - # * ``dsagt init`` copies the shared collections into the project - # * ``dsagt start`` does no embedding work, no implicit rebuild, - # no sentinel checks — just spawns services - # If the project KB is empty or stale, the MCP servers surface a - # clear "run dsagt setup-kb" error rather than silently rebuilding. - - mlflow_port = config.get("mlflow", {}).get("port") or pick_free_port() - config.setdefault("mlflow", {})["port"] = mlflow_port - - proxy_requested = "proxy" in config - proxy_port = None - if proxy_requested: - proxy_port = config["proxy"].get("port") or pick_free_port() - config["proxy"]["port"] = proxy_port - - session_id = config.get( - "session_id", - f"{config['project']}-{datetime.now(timezone.utc).strftime('%Y%m%dT%H%M%SZ')}", - ) - config["session_id"] = session_id - - mlflow_log = pdir / "mlflow.log" - mlflow_proc = subprocess.Popen( - mlflow_command(pdir, config.get("mlflow", {}), port=mlflow_port), - stdout=open(mlflow_log, "w"), - stderr=subprocess.STDOUT, - start_new_session=True, - ) - logger.info( - "MLflow started (pid %d) → http://localhost:%d", - mlflow_proc.pid, mlflow_port, - ) - - pids = {"mlflow": mlflow_proc.pid} - ports = {"mlflow": mlflow_port} - - proxy_proc = None - if proxy_requested: - # MLflow has to be up before the proxy starts because the proxy's - # init_tracing() calls mlflow.set_experiment() during startup. - _wait_for_mlflow(mlflow_port, mlflow_proc, mlflow_log, timeout=30.0) - - proxy_proc = _start_proxy(config, pdir, mlflow_port, proxy_port, session_id) - pids["proxy"] = proxy_proc.pid - ports["proxy"] = proxy_port - - runtime_file.write_text(json.dumps({ - "pids": pids, - "ports": ports, - "started_at": datetime.now(timezone.utc).isoformat(), - }, indent=2) + "\n") - - if not proxy_requested: - _wait_for_mlflow(mlflow_port, mlflow_proc, mlflow_log, timeout=30.0) - if proxy_proc is not None: - _wait_for_proxy(proxy_port, proxy_proc, pdir / "proxy.log", timeout=45.0) - - return ports - - -def _start_proxy( - config: dict, pdir: Path, mlflow_port: int, proxy_port: int, session_id: str, -) -> subprocess.Popen: - """Spawn the dsagt-proxy subprocess. - - Forwards LLM + embedding requests using the user's configured - upstream credentials (LLM_API_KEY / EMBEDDING_API_KEY from - os.environ) and emits OTLP traces to MLflow via init_tracing. - - Crashes loudly if required config keys are missing — better than - spawning a half-configured proxy that 500s on the first agent - request. - """ - llm = config.get("llm") or {} - emb = config.get("embedding") or {} - for required in ("model", "base_url", "provider"): - if not llm.get(required): - raise RuntimeError( - f"--enable-proxy needs config.llm.{required} (got {llm.get(required)!r})" - ) - - cmd = [ - sys.executable, "-m", "dsagt.commands.proxy_server", - "--port", str(proxy_port), - "--mlflow-url", f"http://localhost:{mlflow_port}", - "--project", config["project"], - "--session", session_id, - "--records-dir", str(pdir / "trace_archive"), - "--model", llm["model"], - "--base-url", llm["base_url"], - "--provider", llm["provider"], - ] - # Embedding routing through the proxy is only relevant when the - # project's embedding backend is ``api`` — in ``local`` mode the - # knowledge MCP server uses sentence-transformers in-process and - # never makes HTTP embedding calls, so the proxy doesn't need an - # embedding route at all. - if (emb.get("backend") or "local").lower() == "api": - for required in ("model", "base_url", "provider"): - if not emb.get(required): - raise RuntimeError( - f"--enable-proxy with embedding.backend=api needs " - f"config.embedding.{required} (got {emb.get(required)!r})" - ) - cmd.extend([ - "--embedding-model", emb["model"], - "--embedding-base-url", emb["base_url"], - "--embedding-provider", emb["provider"], - ]) - proxy_log = pdir / "proxy.log" - # The proxy needs the *real* upstream credentials in env (not the - # sentinel agents see). os.environ already has them from the user's - # shell or .env file. We pass DSAGT_PROJECT explicitly so - # _resolve_experiment_id picks the right experiment. - proxy_proc = subprocess.Popen( - cmd, - stdout=open(proxy_log, "w"), - stderr=subprocess.STDOUT, - env={ - **os.environ, - "DSAGT_PROJECT": config["project"], - "DSAGT_PROJECT_DIR": str(pdir), - "DSAGT_SESSION_ID": session_id, - "DSAGT_AGENT": config.get("agent", ""), - "MLFLOW_TRACKING_URI": f"http://localhost:{mlflow_port}", - }, - start_new_session=True, - ) - logger.info( - "Proxy started (pid %d) → http://localhost:%d", - proxy_proc.pid, proxy_port, - ) - return proxy_proc - - -def _wait_for_proxy( - port: int, proc: subprocess.Popen, log_path: Path, timeout: float = 45.0, -) -> None: - """Poll *port* until the proxy answers, the subprocess dies, or we time out. - - Generous default (45s) because LiteLLM's transitive imports - (transformers/torch dependencies) can take 10-15s on warm cache, - longer cold. Raises with proxy.log tail attached so - ``dsagt start`` surfaces the failure instead of the agent's first - LLM call hitting ECONNREFUSED. - """ - deadline = time.monotonic() + timeout - while time.monotonic() < deadline: - if proc.poll() is not None: - tail = log_path.read_text().splitlines()[-20:] if log_path.exists() else [] - raise RuntimeError( - f"Proxy exited with code {proc.returncode} before becoming ready.\n " - + "\n ".join(tail) - ) + tool_use_indexed = 0 try: - with socket.create_connection(("127.0.0.1", port), timeout=0.5): - return - except (ConnectionRefusedError, OSError): - time.sleep(0.25) - raise RuntimeError( - f"Proxy did not accept connections on port {port} within {timeout:.0f}s. " - f"See {log_path} for details." - ) - - -def _wait_for_mlflow( - port: int, proc: subprocess.Popen, log_path: Path, timeout: float = 30.0, -) -> None: - """Poll *port* until MLflow answers, the subprocess dies, or we time out. - - Raises ``RuntimeError`` on failure with the mlflow.log tail attached, - so the failure surfaces at ``dsagt start`` rather than at the agent's - first OTLP export attempt. - """ - deadline = time.monotonic() + timeout - while time.monotonic() < deadline: - if proc.poll() is not None: - tail = log_path.read_text().splitlines()[-20:] if log_path.exists() else [] - raise RuntimeError( - f"MLflow exited with code {proc.returncode} before becoming ready.\n " - + "\n ".join(tail) - ) + # tick_traced: this catch-up runs off any tool-call trace, so tag + # the indexer's kb.* writes dsagt.source=code_use rather than let + # them orphan as untagged top-level traces. + tool_use_indexed = CodeUseIndexer(kb, pdir).tick_traced() + except Exception as e: # noqa: BLE001 — never let a background task crash + logger.warning("Tool execution indexing failed: %s", e) + + traces_caught_up = 0 try: - with socket.create_connection(("127.0.0.1", port), timeout=0.5): - return - except (ConnectionRefusedError, OSError): - time.sleep(0.25) - raise RuntimeError( - f"MLflow did not accept connections on port {port} within {timeout:.0f}s. " - f"See {log_path} for details." - ) - - -def stop_services(project_name: str) -> list[str]: - """User-invoked teardown. Returns ``[]`` when nothing was running.""" - return reap_runtime(project_dir(project_name) / ".runtime") + traces_caught_up = _catch_up_traces(pdir, config, kb) + except Exception as e: # noqa: BLE001 — never let a background task crash + logger.warning("Trace catch-up failed: %s", e) + + return { + "status": "ok", + "tool_use_indexed": tool_use_indexed, + "traces_caught_up": traces_caught_up, + } + finally: + kb.close() +def _catch_up_traces(pdir: Path, config: dict, kb) -> int: + """Re-collect the previous session's transcript; return turns newly emitted. -# --------------------------------------------------------------------------- -# Memory extraction orchestration -# --------------------------------------------------------------------------- + Builds a trace collector pinned to the previous session's recorded + trace-source token (and tagged with its session id), then runs one + ``collect(include_last=True)``. The collector's session-qualified ack files + are shared with the live pass, so already-logged turns are skipped and only + those lost to an ungraceful shutdown are emitted to MLflow + episodic memory. + Uniform across agents — JSONL or SQLite — since the pin is the agent's own + session token, not a transcript-file assumption. -def run_extraction(project_name: str) -> dict: - """Two-phase post-session work, both best-effort. - - 1. **Tool-execution indexing** (always): embed every JSON record in - ``/trace_archive/`` into the project's ``tool_use`` - collection so the agent can semantic-search prior tool runs in - future sessions. Pure embedding work — no LLM call, no - ``DSAGT_MEMORY_*`` needed. Local-backend embeddings (BYOA - default) need no credentials at all. - 2. **LLM-based memory extraction** (gated on ``DSAGT_MEMORY_*``): - summarises MLflow LLM-call traces into episodic memories. - Currently reads only ``service.name == "dsagt-proxy"`` traces - (see ``memory.DSAGT_EXTRACTION_SOURCE_SERVICE_NAME``); in BYOA - these don't exist yet, so this phase is effectively dormant - until Phase 2 / native-shape parsers land. + Returns 0 (a no-op) when there is no previous session, or it stamped no + trace-source (a session too short for the heartbeat to record one), where + guessing would risk reading the *new* session's records. """ - config = load_config(project_name) - pdir = Path(config["project_dir"]) - - emb_config = config.get("embedding", {}) - backend = emb_config.get("backend", "local") - # Local backend rejects base_url / api_key (no remote call); api - # backend reads creds from the shell. Local model must be a HF - # identifier (``org/repo``); if it isn't (legacy projects had - # Ollama-style ``nomic-embed-text`` here), fall through to - # LocalEmbeddingClient's default by passing model=None. - if backend == "local": - model = emb_config.get("model") - if model and "/" not in str(model): - model = None - embedder_kwargs = {"model": model} - else: - embedder_kwargs = { - "model": emb_config.get("model"), - "base_url": emb_config.get("base_url"), - "api_key": os.environ.get("EMBEDDING_API_KEY", ""), - } - kb = KnowledgeBase( - index_dir=pdir / "kb_index", - default_embedder=backend, - embedder_kwargs=embedder_kwargs, + from dsagt.memory import episodic_consumers + from dsagt.observability import resolve_tracking_uri + from dsagt.traces import make_trace_collector + + sessions = read_state(pdir).get("sessions") or [] + if len(sessions) < 2: + return 0 + prev = sessions[-2] + source = prev.get("trace_source") + if source is None: + return 0 + + project = config.get("project", "") + prev_tag = session_tag(project, prev["id"]) + collector = make_trace_collector( + config.get("agent"), + pdir, + project, + prev_tag, + resolve_tracking_uri(config), + extra_consumers=episodic_consumers(config, kb, pdir, prev_tag), + source=source, ) - - # Phase 1: index trace_archive into tool_use collection. - tool_use_indexed = 0 - try: - trace_result = index_trace_archive(pdir / "trace_archive", kb) - tool_use_indexed = trace_result.get("indexed", 0) - except Exception as e: - logger.warning("Tool execution indexing failed: %s", e) - - # Phase 2: LLM-based memory extraction. Skip silently if not configured. - api_key = os.environ.get("DSAGT_MEMORY_API_KEY", "") - model = os.environ.get("DSAGT_MEMORY_MODEL", "") - if not api_key or not model: - kb.close() - return {"status": "tool_use_only", "tool_use_indexed": tool_use_indexed} - - base_url = os.environ.get("DSAGT_MEMORY_BASE_URL", "") - provider = os.environ.get("DSAGT_MEMORY_PROVIDER") or None - session_id = config.get("session_id") or config.get("project", "") - categories = config.get("categories", {}) - - mlflow_port = config.get("mlflow", {}).get("port") - mlflow_uri = ( - f"http://localhost:{mlflow_port}" if mlflow_port - else os.environ.get("MLFLOW_TRACKING_URI") - ) - try: - result = extract_session( - project_name=project_name, - kb=kb, - api_key=api_key, - model=model, - base_url=base_url or None, - provider=provider, - session_id=session_id, - categories=categories if categories else None, - runtime_dir=pdir, - outlier_sensitivity=float( - config.get("extraction", {}).get("outlier_sensitivity", 0) - ), - mlflow_uri=mlflow_uri, - ) - result["tool_use_indexed"] = tool_use_indexed - return result - finally: - kb.close() + if collector is None: + return 0 + return collector.collect(include_last=True) diff --git a/src/dsagt/skills.py b/src/dsagt/skills.py new file mode 100644 index 0000000..fbd1ca7 --- /dev/null +++ b/src/dsagt/skills.py @@ -0,0 +1,838 @@ +"""Skill discovery — catalog data plane, keyword scorer, and the router facade. + +DSAGT fetches external Agent-Skills repos, indexes each per-source into a +``skills_catalog__`` KB collection, and searches/installs them into a +project. This is the one job native skill discovery can't do: a *catalog* skill +stays searchable without being copied locally or held in the agent's context +(you can't hold thousands of skill descriptions in context), while an +*installed* skill is copied into ``/skills//`` and mirrored into +the agent's native skills dir (``agents.base.setup_skills``). It backs the MCP +``search_skills`` tool and the ``dsagt skills`` CLI through the one +:class:`SkillRouter` facade, so search/install policy can't diverge between them. +Design-wise it stays cheap and degradable: :class:`SkillsCatalog` composes over +the host server's :class:`~dsagt.knowledge.KnowledgeBase` (shared embedder, no +second model load), falls back to a Genesis-derived keyword scorer +(:func:`rank_skills`) when no embedder/KB is configured, and indexes per-source +so re-sync is an idempotent drop-and-rebuild of just that source's collection. + +Class map — every edge is ``─ Class`` (``◇`` holds · ``◆`` owns):: + + SkillRouter render/MCP facade: the search_skills string, + │ the empty-result message, exact-name lookup + ├─◇ SkillsCatalog the catalog data plane (constructed here, or + │ │ shared in via catalog=) + │ └─◇ KnowledgeBase shared vector store + embedder; None selects + │ the keyword fallback over the clone cache + └─◇ SkillRegistry installed-skill registry, exact-name lookup only + + free fns: + keyword scorer score_skill · rank_skills (Genesis parity) + source resolve resolve_source · _repo_slug · persist_source_to_config + sync / index sync_source · _discover_skill_dirs · index_catalog + install find_catalog_skill · install_into_project · _capture_attribution + render _where_label + +Genesis Skills: Apache-2.0, gitlab.osti.gov/genesis/genesis-skills +(``skill_search/catalog.py``). +""" + +from __future__ import annotations + +import json +import logging +import re +import shutil +from pathlib import Path + +import yaml + +from dsagt.registry import ( + CATALOG_COLLECTION_PREFIX, + _parse_frontmatter, + catalog_collection, +) +from dsagt.session import REGISTRY_DIR + +logger = logging.getLogger(__name__) + + +# =========================================================================== +# Keyword scorer — Genesis-derived token-overlap fallback (stdlib only) +# =========================================================================== +# +# A faithful reimplementation (not an import) of the Genesis Skills +# ``skill-search`` engine (``skill_search/catalog.py``: ``_score_skill`` / +# ``rank_skills``). Used by :class:`SkillsCatalog` when no embedder / KB is +# configured: keyword overlap only, deterministic. +# +# Scoring (per skill, against a query) — matching Genesis exactly: +# +# * +2 for each query token that also appears in the skill **name** +# * +1 for each query token that also appears in the **description** +# * then **at most one** substring bonus (mutually exclusive, in priority +# order): +6 if the query equals the name, else +4 if it is a substring of +# the name, else +2 if it is a substring of the description +# +# Tokens are casefolded ``\w+`` runs with hyphens split, single-character +# tokens and stopwords dropped. Ties break by name (ascending); below +# ``min_score`` are dropped. + +#: Stopword set — kept identical to Genesis so ranking parity holds. +_STOPWORDS = frozenset( + { + "a", + "an", + "and", + "as", + "at", + "be", + "for", + "from", + "if", + "in", + "into", + "is", + "it", + "of", + "on", + "or", + "please", + "the", + "this", + "to", + "use", + "using", + "with", + } +) + +_TOKEN_RE = re.compile(r"\w+", flags=re.UNICODE) + + +def _tokens(text: str) -> set[str]: + """Casefolded word tokens (hyphens split), single-char + stopwords removed.""" + normalized = (text or "").casefold().replace("-", " ") + return { + t for t in _TOKEN_RE.findall(normalized) if len(t) > 1 and t not in _STOPWORDS + } + + +def score_skill(query: str, name: str, description: str) -> float: + """Token-overlap score of one skill against *query* (0.0 = no match).""" + qtokens = _tokens(query) + normalized_query = (query or "").casefold().strip() + score = 2 * len(qtokens & _tokens(name)) + len(qtokens & _tokens(description)) + + if normalized_query: + name_l = (name or "").casefold() + if normalized_query == name_l: + score += 6 + elif normalized_query in name_l: + score += 4 + elif normalized_query in (description or "").casefold(): + score += 2 + return float(score) + + +def rank_skills( + query: str, skills, top_k: int | None = 8, min_score: int = 1 +) -> list[tuple[dict, float]]: + """Rank *skills* (dicts with ``name`` + ``description``) against *query*. + + Returns ``[(skill, score), ...]`` for skills scoring at least *min_score*, + sorted by score descending then name ascending, truncated to *top_k* (all + when *top_k* is ``None``). + """ + scored: list[tuple[dict, float]] = [] + for s in skills: + sc = score_skill(query, s.get("name", ""), s.get("description", "")) + if sc >= min_score: + scored.append((s, sc)) + scored.sort(key=lambda kv: (-kv[1], (kv[0].get("name") or ""))) + return scored[:top_k] if top_k is not None else scored + + +# =========================================================================== +# Catalog data plane — sources, sync/index, lookup/install, SkillsCatalog +# =========================================================================== + +#: Default source enabled out of the box (matches .dsagt/config.yaml default). +DEFAULT_SOURCE = "k-dense-ai" + +#: Curated, named skill sources. ``subdir`` scopes the recursive SKILL.md +#: walk when set (cheaper clone); when omitted the whole repo is cloned and +#: walked, which is robust to category-nested layouts. +KNOWN_SOURCES: dict[str, dict] = { + "k-dense-ai": { + "url": "https://github.com/K-Dense-AI/scientific-agent-skills", + "branch": "main", + "subdir": "skills", + "description": "K-Dense scientific agent skills — chem/bio/medicine/materials (140+).", + }, + "anthropic": { + "url": "https://github.com/anthropics/skills", + "branch": "main", + "subdir": "skills", + "description": "Official Anthropic skills + document-editing examples.", + }, + "antigravity": { + "url": "https://github.com/sickn33/antigravity-awesome-skills", + "branch": "main", + "subdir": None, + "description": "Antigravity Awesome Skills — 1,500+ cross-platform agentic skills.", + }, + "composio": { + "url": "https://github.com/ComposioHQ/awesome-claude-skills", + "branch": "master", + "subdir": None, + "description": "Composio awesome-claude-skills — workflow skills for many SaaS apps.", + }, + "genesis": { + "url": "https://gitlab.osti.gov/genesis/genesis-skills", + "branch": "main", + "subdir": "skills", + "description": "GENESIS skills (OSTI GitLab) — aggregated agent-skill " + "catalog: HPC (Slurm/PBS, Perlmutter/Aurora/Frontier), HuggingFace, " + "LangChain, OpenAI, Anthropic, plasma-sim, ModCon, and more (70+).", + }, +} + +#: Shared, machine-global cache of cloned source repos (sibling of kb_index/). +SKILL_SOURCES_DIR = REGISTRY_DIR / ".skill_sources" + + +# --------------------------------------------------------------------------- +# Source resolution + slugging +# --------------------------------------------------------------------------- + + +def resolve_source(source: str | dict) -> dict: + """Resolve a known-source name, a git URL (any host), or a full spec dict. + + A full ``http(s)://`` / ``git@`` URL works for any host (GitHub, GitLab, + …); the bare ``owner/repo`` shorthand assumes GitHub. Returns a dict with + at least ``url``; optional ``branch`` / ``subdir``. + """ + if isinstance(source, dict): + if not source.get("url"): + raise ValueError("source dict must include a 'url'") + return source + if source in KNOWN_SOURCES: + return dict(KNOWN_SOURCES[source]) + if source.startswith(("http://", "https://", "git@")) or source.count("/") == 1: + # Full URL or ``owner/repo`` shorthand. + url = ( + source + if "://" in source or source.startswith("git@") + else f"https://github.com/{source}" + ) + return {"url": url, "branch": "main", "subdir": None} + raise ValueError( + f"Unknown skill source '{source}'. Use a known name " + f"({', '.join(sorted(KNOWN_SOURCES))}), a git URL (any host), " + f"or owner/repo (GitHub)." + ) + + +def persist_source_to_config(project_dir: str | Path, spec: dict) -> bool: + """Append a resolved source to ``skills.sources`` in the project config. + + Dedupes by URL. Returns True if the config was updated. No-op (returns + False) if the config file is missing — the catalog is still indexed + either way. Used by both the ``add_skill_source`` MCP tool and the + ``dsagt skills add`` CLI so a CLI-added source is re-synced by a later + config-driven ``dsagt skills sync``. + """ + cfg_path = Path(project_dir) / ".dsagt" / "config.yaml" + if not cfg_path.exists(): + return False + cfg = yaml.safe_load(cfg_path.read_text()) or {} + sources = cfg.setdefault("skills", {}).setdefault("sources", []) + if any(s.get("url") == spec.get("url") for s in sources): + return False + sources.append( + {k: spec[k] for k in ("name", "url", "branch", "subdir") if k in spec} + ) + cfg_path.write_text(yaml.dump(cfg, default_flow_style=False, sort_keys=False)) + return True + + +def _repo_slug(url: str) -> str: + """Stable, collection-name-safe slug from a repo URL (``owner-repo``). + + Host-agnostic: the scheme and host are stripped so github.com, gitlab.*, + etc. all reduce to the ``owner/repo`` path. GitHub URLs keep the slug + they had before this generalization, so existing catalog collections do + not need rebuilding. + """ + s = url.rstrip("/") + s = re.sub(r"^https?://", "", s) # drop scheme + s = re.sub(r"^git@", "", s) # ssh form: git@host:owner/repo + s = re.sub(r"\.git$", "", s).lower() + s = re.sub(r"^[^/:]+[/:]", "", s) # drop the host segment + s = re.sub(r"[^a-z0-9]+", "-", s).strip("-") + return s[:40] + + +# --------------------------------------------------------------------------- +# Discovery +# --------------------------------------------------------------------------- + + +def _discover_skill_dirs(root: Path) -> list[Path]: + """Recursively find skill directories (any dir holding a parseable SKILL.md). + + Recursive so both flat (``skills//SKILL.md``) and category-nested + (``skills///SKILL.md``) repo layouts work. A directory + qualifies only if its SKILL.md has YAML frontmatter with a ``name``. + """ + out: list[Path] = [] + if not root.exists(): + return out + for skill_md in sorted(root.rglob("SKILL.md")): + try: + spec = _parse_frontmatter(skill_md) + except ValueError as e: # malformed frontmatter — skip, don't abort + logger.warning("skipping %s: %s", skill_md, e) + continue + if spec.get("name"): + out.append(skill_md.parent) + return out + + +# --------------------------------------------------------------------------- +# Sync (clone + index) +# --------------------------------------------------------------------------- + + +def sync_source( + source: str | dict, + *, + kb=None, + cache_dir: Path = SKILL_SOURCES_DIR, + force: bool = False, +) -> dict: + """Clone *source* into the cache and (re)index its skills into the catalog. + + ``force`` re-clones from scratch. Indexing wipes and rebuilds only this + source's ``skills_catalog__`` collection, so other catalogs and the + installed/bundled ``skills`` collection are untouched. When *kb* is None + the clone still happens (so ``install`` works offline-of-KB) but nothing + is indexed. + """ + spec = resolve_source(source) + slug = _repo_slug(spec["url"]) + dest = cache_dir / slug + + if force and dest.exists(): + shutil.rmtree(dest) + if not dest.exists(): + from dsagt.commands.setup_core_kb import clone_github # lazy: break cycle + + dest.mkdir(parents=True, exist_ok=True) + subdir = spec.get("subdir") + include = [subdir] if subdir else None + clone_github( + spec["url"], dest, branch=spec.get("branch", "main"), include=include + ) + + walk_root = dest / spec["subdir"] if spec.get("subdir") else dest + skill_dirs = _discover_skill_dirs(walk_root) + indexed = index_catalog(skill_dirs, slug, spec["url"], kb) if kb is not None else 0 + if kb is not None and not skill_dirs: + logger.warning( + "source %s yielded no SKILL.md skills under %s", spec["url"], walk_root + ) + + return { + "slug": slug, + "url": spec["url"], + "discovered": len(skill_dirs), + "indexed": indexed, + "cache_dir": str(dest), + } + + +def _catalog_embed_text(spec: dict, fallback_name: str) -> str: + """Text embedded for catalog search: the frontmatter ``name`` + ``description`` + (+ ``tags``) only — *not* the SKILL.md body. + + Discovery is progressive-disclosure level 1: the description is authored to + say *what the skill does and when to use it*, which is exactly the match + target. Embedding the body would dilute that signal, and the embedder + truncates long input anyway (so a full SKILL.md is both incomplete and + misallocated). This also keeps the semantic backend ranking on the same + fields as the keyword fallback. + """ + name = spec.get("name") or fallback_name + desc = spec.get("description") or "" + tags = " ".join(spec.get("tags") or []) + return f"{name}: {desc} {tags}".strip() + + +def index_catalog(skill_dirs: list[Path], slug: str, url: str, kb) -> int: + """Wipe + rebuild source *slug*'s catalog collection from *skill_dirs*.""" + collection = catalog_collection(slug) + coll_dir = Path(kb.index_dir) / collection + if coll_dir.exists(): + shutil.rmtree(coll_dir) + + texts: list[str] = [] + metas: list[dict] = [] + for d in skill_dirs: + skill_md = d / "SKILL.md" + spec = _parse_frontmatter(skill_md) + name = spec.get("name") or d.name + texts.append(_catalog_embed_text(spec, d.name)) + metas.append( + { + "skill_name": name, + "description": (spec.get("description") or "")[:300], + "tags": ",".join(spec.get("tags") or []), + "source": f"catalog:{slug}", + "source_url": url, + "cache_path": str(d), + } + ) + if texts: + kb.add_entries(texts=texts, collection=collection, metadatas=metas) + return len(texts) + + +# --------------------------------------------------------------------------- +# Lookup + install +# --------------------------------------------------------------------------- + + +def find_catalog_skill(name: str, *, cache_dir: Path = SKILL_SOURCES_DIR) -> Path: + """Locate a cached catalog skill dir by name across all synced sources. + + Matches on frontmatter ``name`` first, then directory name. A bare name + must be unique across the machine-global clone cache; when the same name + exists in more than one synced source, pass a **source-qualified** + ``/`` (the slug is the per-source cache dir / catalog-collection + suffix, as shown by ``list_skill_sources`` / ``dsagt skills list + --catalog``) to pick one. Raises on no match or on a still-ambiguous bare + name. + """ + source_filter: str | None = None + skill = name + if "/" in name: + # Skill names never contain '/', so a slash means "/". + source_filter, skill = name.split("/", 1) + + matches: list[Path] = [] + if cache_dir.exists(): + for slug_dir in sorted(p for p in cache_dir.iterdir() if p.is_dir()): + if source_filter is not None and slug_dir.name != source_filter: + continue + for d in _discover_skill_dirs(slug_dir): + spec = _parse_frontmatter(d / "SKILL.md") + if spec.get("name") == skill or d.name == skill: + matches.append(d) + if not matches: + where = f" in source '{source_filter}'" if source_filter else "" + raise LookupError( + f"No catalog skill named '{skill}'{where}. Run 'dsagt skills sync' " + f"or add_skill_source first, then search_skills to find one." + ) + # Collapse matches that point at the same source repo (slug = first path + # part under cache_dir); ambiguity only matters across different sources. + by_source = {p.relative_to(cache_dir).parts[0]: p for p in matches} + if len(by_source) > 1: + sources = sorted(by_source) + raise LookupError( + f"Skill '{skill}' exists in multiple sources ({', '.join(sources)}); " + f"reinstall with a source-qualified name, e.g. '{sources[0]}/{skill}'." + ) + return next(iter(by_source.values())) + + +#: License / attribution files to preserve when installing a catalog skill. +_ATTRIBUTION_GLOBS = ( + "LICENSE*", + "NOTICE*", + "COPYING*", + "COPYRIGHT*", + "ATTRIBUTION*", +) + + +def _capture_attribution(src: Path, dest: Path, cache_dir: Path) -> list[str]: + """Preserve license/attribution when installing a (third-party) catalog skill. + + ``copytree`` already carries files *inside* the skill dir. A skill is often + governed by a per-subtree or repo-root ``LICENSE`` / ``NOTICE`` / + ``ATTRIBUTION`` that lives *outside* its own folder, so this also pulls those + from ancestor dirs up to the source repo root (which ``clone_github`` mirrors + into the cache root even for sparse ``subdir`` clones). Nearest ancestor + wins a filename collision; skill-local files (already in ``dest``) are never + overwritten. Always stamps a ``PROVENANCE.txt`` recording the source. + Returns the names of files captured from ancestors. + """ + src, dest, cache_dir = Path(src), Path(dest), Path(cache_dir) + try: + slug = src.relative_to(cache_dir).parts[0] + repo_root = cache_dir / slug + rel = src.relative_to(repo_root) + except ValueError: # src outside the cache (shouldn't happen) — degrade. + slug, repo_root, rel = src.parent.name, src.parent, Path(src.name) + + captured: list[str] = [] + node = src.parent + while True: + for pat in _ATTRIBUTION_GLOBS: + for f in sorted(node.glob(pat)): + if f.is_file() and not (dest / f.name).exists(): + shutil.copy2(f, dest / f.name) + captured.append(f.name) + if node == repo_root or node == node.parent: + break + node = node.parent + + (dest / "PROVENANCE.txt").write_text( + f"Installed by dsagt from catalog source: {slug}\n" + f"Source path in repo: {rel}\n" + ) + return captured + + +def install_into_project( + name: str, project_dir: str | Path, *, cache_dir: Path = SKILL_SOURCES_DIR +) -> dict: + """Copy a catalog skill into ``/skills//`` (with scripts/refs). + + The destination directory is named after the skill's frontmatter ``name`` + (falling back to its source dir name) so it matches the invocable name in + native discovery. Preserves upstream license/attribution (see + :func:`_capture_attribution`). Returns + ``{name, source_dir, dest_dir, action, attribution}``. + """ + src = find_catalog_skill(name, cache_dir=cache_dir) + spec = _parse_frontmatter(src / "SKILL.md") + skill_name = spec.get("name") or src.name + + dest = Path(project_dir) / "skills" / skill_name + action = "updated" if dest.exists() else "added" + if dest.exists(): + shutil.rmtree(dest) + shutil.copytree(src, dest) + attribution = _capture_attribution(src, dest, cache_dir) + + return { + "name": skill_name, + "source_dir": str(src), + "dest_dir": str(dest), + "action": action, + "attribution": attribution, + } + + +# --------------------------------------------------------------------------- +# SkillsCatalog — the catalog data plane (composition over KnowledgeBase) +# --------------------------------------------------------------------------- + + +class SkillsCatalog: + """The external-skill *catalog* behind one object. + + Composition over :class:`~dsagt.knowledge.KnowledgeBase`: it holds a KB + handle (the host server's existing instance → shared embedder, no second + model load) plus the clone-cache dir, and exposes the skill-domain ops — + ``sync`` / ``search`` / ``install`` / ``list_sources``. The skill-specific + behavior (frontmatter-indexed catalog collections, the no-embedder keyword + fallback over the clone cache) lives here; the vector store + embedder are + the shared KB. :class:`SkillRouter` is a thin render/MCP facade over this. + + Catalog tier only: installed/created skills are natively auto-discovered by + every supported agent, so they are never search candidates. ``search`` + covers the not-yet-installed catalog (which native discovery can't see) + plus the Cline-skills-disabled / no-embedder case via the keyword scorer. + """ + + def __init__(self, *, kb=None, cache_dir: Path | None = None): + """Compose a catalog over an existing KB + a clone-cache directory. + + ``kb`` is the host server's :class:`~dsagt.knowledge.KnowledgeBase` (so + the embedder/Chroma are shared, never a second model load); pass ``None`` + for the no-embedder keyword path. ``cache_dir`` overrides the default + machine-global clone cache (:data:`SKILL_SOURCES_DIR`) — handy for tests. + """ + self._kb = kb + self._cache_dir = cache_dir # default resolved lazily + + @property + def has_kb(self) -> bool: + """True when an embedder-backed KB is available (vs the keyword path).""" + return self._kb is not None + + def _resolved_cache_dir(self) -> Path: + """The clone-cache dir — the ``cache_dir`` override or the global default.""" + return Path(self._cache_dir or SKILL_SOURCES_DIR) + + # -- write ops (delegate to the module functions) ------------------------ + + def sync(self, source, *, force: bool = False) -> dict: + """Clone + frontmatter-index a source into its catalog collection.""" + return sync_source( + source, kb=self._kb, cache_dir=self._resolved_cache_dir(), force=force + ) + + def install(self, name: str, project_dir) -> dict: + """Copy a catalog skill into ``/skills//`` (+ attribution).""" + return install_into_project( + name, project_dir, cache_dir=self._resolved_cache_dir() + ) + + # -- backend selection --------------------------------------------------- + + def synced_collections(self) -> list[str]: + """The ``skills_catalog__*`` collections currently indexed in the KB.""" + if self._kb is None: + return [] + return [ + c for c in self._kb.collections if c.startswith(CATALOG_COLLECTION_PREFIX) + ] + + def search(self, query=None, *, top_k: int = 8, tag=None) -> list[dict]: + """Rank catalog candidates → normalized hit dicts (data, not rendered). + + ChromaDB semantic search when a KB exists, else the Genesis-derived + keyword scorer over the clone cache. + """ + if self._kb is not None: + return self._select_kb(query, top_k, tag) + return self._select_keyword(query, top_k, tag) + + def _select_kb(self, query, top_k: int, tag) -> list[dict]: + """Semantic backend: rank-fused search across every synced catalog + collection, normalized into hit dicts. + + Fan-out + RRF live in ``KnowledgeBase.search`` (the shared substrate); + catalog collections are homogeneous (one embedder) so the fusion is a + clean rank merge. When a ``tag`` filter is set we over-fetch + (``top_k * 3``) then post-filter so the tag doesn't starve the results. + """ + collections = self.synced_collections() + if not collections: + return [] + fetch_k = top_k * 3 if tag else top_k + # collections all come from synced_collections() (they exist), and + # KnowledgeBase.search already skips a missing collection with a warning + # and raises only when every target fails — so a raise here is a real + # failure worth surfacing, not a can't-happen state to swallow. + hits = self._kb.search( + query=query or "skill", collections=collections, top_k=fetch_k + ) + out = [] + for r in hits: + chunk = r.get("chunk", {}) + meta = chunk.get("metadata", {}) + out.append( + { + "name": meta.get("skill_name", "unknown"), + "summary": (meta.get("description") or chunk.get("text", "") or "")[ + :200 + ], + "source": meta.get("source", ""), + "tags": meta.get("tags", ""), + "score": r.get("score", 0.0), + } + ) + if tag: + out = [h for h in out if tag in (h["tags"] or "")] + out.sort(key=lambda h: h["score"], reverse=True) + return out[:top_k] + + def _candidate_skills(self) -> list[dict]: + """Cached-catalog skills as scorer-ready dicts (no KB needed).""" + cands: list[dict] = [] + cache = self._resolved_cache_dir() + if cache.exists(): + for slug_dir in sorted(p for p in cache.iterdir() if p.is_dir()): + for d in _discover_skill_dirs(slug_dir): + spec = _parse_frontmatter(d / "SKILL.md") + if spec.get("name"): + cands.append( + { + "name": spec["name"], + "description": spec.get("description", ""), + "tags": ",".join(spec.get("tags", []) or []), + "source": f"catalog:{slug_dir.name}", + } + ) + return cands + + def _select_keyword(self, query, top_k: int, tag) -> list[dict]: + """No-embedder backend: rank cached-catalog skills with the Genesis + token-overlap scorer (:func:`rank_skills`) into normalized hit dicts. + + With no ``query`` there's nothing to score, so it returns the first + ``top_k`` candidates (tag-filtered) at score 0.0 — a browse mode rather + than a search. + """ + cands = self._candidate_skills() + if tag: + cands = [c for c in cands if tag in (c["tags"] or "")] + if not query: + picks = cands[:top_k] + return [ + {**c, "summary": (c["description"] or "")[:200], "score": 0.0} + for c in picks + ] + ranked = rank_skills(query, cands, top_k=top_k) + return [ + {**c, "summary": (c["description"] or "")[:200], "score": sc} + for c, sc in ranked + ] + + # -- source view --------------------------------------------------------- + + def list_sources(self) -> list[dict]: + """Known sources + synced flag + indexed count (one source of truth).""" + synced = set(self.synced_collections()) + out = [] + for name, spec in KNOWN_SOURCES.items(): + coll = catalog_collection(_repo_slug(spec["url"])) + is_synced = coll in synced + out.append( + { + "name": name, + "url": spec["url"], + "description": spec.get("description", ""), + "synced": is_synced, + "indexed": self._indexed_count(coll) if is_synced else 0, + } + ) + return out + + def _indexed_count(self, collection: str) -> int: + """Number of skills indexed in a catalog *collection* (0 if absent/no KB). + + Reads the collection's persisted ``chroma_ids.json`` directly rather than + querying the store — cheap, and works without loading the embedder. + """ + if self._kb is None: + return 0 + ids = Path(self._kb.index_dir) / collection / "chroma_ids.json" + try: + return len(json.loads(ids.read_text())) + except (FileNotFoundError, ValueError, OSError): + return 0 + + +# =========================================================================== +# SkillRouter — the thin render/MCP facade over the catalog +# =========================================================================== +# +# ``SkillRouter`` adds only the presentation concerns that the MCP handlers and +# the CLI share: rendering a ranked hit list into the ``search_skills`` string, +# the empty-result message, and the exact-``skill_name`` lookup (which needs the +# installed-skill registry, not the catalog). +# +# Construct it from the same inputs at every call site (MCP ``search_skills``, +# CLI ``skills search/list``) so policy can't diverge between them — or hand it +# a prebuilt :class:`SkillsCatalog` via ``catalog=`` so a server that already +# owns a shared KB reuses one catalog instance. +# +# Skill *materialization* (mirroring installed skills into each agent's native +# skills directory) lives in the agent layer (``AgentSetup.setup_skills``), not +# here: every supported agent (claude/codex/goose/cline) natively +# auto-discovers ``SKILL.md`` folders, so there is no agent-facing disclosure +# tier for the router to own. ``search_skills`` exists for the *catalog* tier +# (skills not yet installed, which native discovery can't see) plus the +# no-embedder keyword fallback. + + +def _where_label(source: str) -> str: + """Human tag for a hit's origin, matching the legacy search output.""" + if source in ("bundled", "registered", "installed"): + return " [installed]" + if source.startswith("catalog:"): + return " [catalog · install_skill to add]" + return "" + + +class SkillRouter: + """Renders catalog discovery for the MCP ``search_skills`` tool + the CLI.""" + + def __init__(self, *, kb=None, skill_registry=None, cache_dir=None, catalog=None): + """Wire the router to a catalog + (optional) installed-skill registry. + + Pass a prebuilt ``catalog`` (a :class:`SkillsCatalog`) to share one + instance with the server; otherwise one is constructed from ``kb`` / + ``cache_dir``. ``skill_registry`` is only consulted for the exact + ``skill_name`` lookup in :meth:`search` (installed skills live there, + not in the catalog), so it may be ``None`` for catalog-only callers. + """ + self._catalog = ( + catalog + if catalog is not None + else SkillsCatalog(kb=kb, cache_dir=cache_dir) + ) + self._reg = skill_registry + + # -- rendering ----------------------------------------------------------- + + def _render_search(self, hits: list[dict]) -> str: + """Format ranked catalog hits into the ``search_skills`` markdown list. + + Each line carries the skill name, an origin tag (:func:`_where_label`), + the score, and the summary — the human-facing string the MCP tool and + CLI both return. + """ + lines = [] + for h in hits: + lines.append( + f"- **{h['name']}**{_where_label(h['source'])} " + f"(score: {h['score']:.2f})\n {h['summary']}" + ) + return f"Found {len(hits)} skill(s):\n\n" + "\n\n".join(lines) + + def _empty_message(self) -> str: + """The no-results string, tailored to *why* nothing matched. + + When a KB exists but no catalog source is synced yet, point the agent + at ``list_skill_sources`` / ``add_skill_source`` (the likely cause); + otherwise it's a genuine no-match for the query. + """ + if not self._catalog.synced_collections() and self._catalog.has_kb: + return ( + "No catalog skills found. No external skill catalog is synced " + "yet — search covers the catalog (skills you can install), since " + "installed skills are already natively discoverable. Call " + "list_skill_sources() to see available sources, then " + "add_skill_source(source=...) to sync one before searching again." + ) + return "No catalog skills found matching the query." + + # -- public API ---------------------------------------------------------- + + def search(self, query=None, *, top_k: int = 8, tag=None, skill_name=None) -> str: + """Stage B. Select + render. Stateless — no session/exposure tracking.""" + if skill_name: + if self._reg is None: + return f"No skill named '{skill_name}'." + spec = self._reg.get_skill(skill_name) + if spec: + return f"Found skill '{skill_name}':\n\n" + yaml.dump( + spec, default_flow_style=False, sort_keys=False + ) + return f"No skill named '{skill_name}'." + + hits = self._catalog.search(query, top_k=top_k, tag=tag) + if not hits: + return self._empty_message() + return self._render_search(hits) + + def sync(self, source, *, force: bool = False) -> dict: + """Stage A passthrough — see :meth:`SkillsCatalog.sync`.""" + return self._catalog.sync(source, force=force) + + def install(self, name: str, project_dir) -> dict: + """Stage C passthrough — see :meth:`SkillsCatalog.install`.""" + return self._catalog.install(name, project_dir) + + def list_sources(self) -> list[dict]: + """Stage A view passthrough — see :meth:`SkillsCatalog.list_sources`.""" + return self._catalog.list_sources() diff --git a/src/dsagt/skills/datacard-generator/SKILL.md b/src/dsagt/skills/datacard-generator/SKILL.md deleted file mode 100644 index 3a2c06a..0000000 --- a/src/dsagt/skills/datacard-generator/SKILL.md +++ /dev/null @@ -1,88 +0,0 @@ ---- -name: generating-datacards -description: "Generates MODCON Datacard v1 documentation for scientific datasets by introspecting a directory and filling a structured template. Use when the user asks to create a datacard, dataset card, dataset documentation, dataset metadata, document a dataset, or prepare a dataset for sharing. Supports three readiness levels: Level 1 (Discoverable), Level 2 (Interoperable & Reusable), Level 3 (Understandable & Trustworthy)." ---- - -# Generating Datacards - -Introspect a dataset directory to auto-populate a MODCON Datacard v1, then prompt the user for anything that couldn't be inferred. - -## Workflow - -Copy this checklist and check off steps as you go: - -``` -Progress: -- [ ] 1. Gather context (path + level) -- [ ] 2. Introspect dataset directory -- [ ] 3. Auto-fill the data card -- [ ] 4. Prompt for missing fields -- [ ] 5. Generate output file -- [ ] 6. Review summary -``` - -### 1. Gather Context - -Ask the user: -- **Dataset path** — directory to document -- **Readiness level** — choose one: - - **Level 1 — Discoverable**: identification, description, file structure, keywords - - **Level 2 — Interoperable & Reusable**: Level 1 + contacts, access, license, provenance, authors, related resources - - **Level 3 — Understandable & Trustworthy**: Level 2 + schema, quality, integrity, AI/ML, variable docs - -### 2. Introspect the Dataset Directory - -See [reference/introspection-commands.md](reference/introspection-commands.md) for exact commands to run for file structure, schema extraction, and existing metadata files. - -### 3. Auto-Fill the Data Card - -Read `references/datacard_template_v1.md`. Populate fields using this decision table: - -| Field | Auto-fill if… | Otherwise | -|-------|--------------|-----------| -| `datacard_creation.created_date` | Always | — | -| `datacard_creation.creation_method` | Always → `"hybrid"` | — | -| `dataset_info.data_formats` | File extensions found | Prompt | -| `dataset_info.modalities` | File types recognized | Prompt | -| `dataset_info.features` | Column headers extracted | Prompt at Level 2+ | -| `dataset_info.splits` | train/test/val dirs found | Leave empty | -| `dataset_counts` | File/record counts available | Prompt | -| `dataset_storage` | `du` output available | Prompt | -| `title` / `name` | README or metadata file found | Prompt | -| Description | README found | Prompt | -| Keywords | Metadata files found | Prompt at Level 1+ | -| License | LICENSE file found | Prompt at Level 2+ | -| Authors / contributors | CITATION.cff found | Prompt at Level 2+ | -| Citation | .bib or CITATION.cff found | Prompt at Level 2+ | - -Cross-check these fields for consistency between YAML frontmatter and markdown body: `title`/`name`, `authors`/`contributors`, `license`, `dataset_info.data_formats`, key dates. - -### 4. Prompt for Missing Fields - -Present auto-discovered values for confirmation, then ask for unfilled fields. Ask **3–5 at a time**. See [reference/field-prompts.md](reference/field-prompts.md) for the full per-level field list. - -### 5. Generate the Data Card - -- **Filename**: `modcon_datacard_.md` -- **Location**: save inside `/` by default; ask if the user prefers elsewhere -- Replace all `[!TODO]`, ``, ``, `` — none should appear in output -- Set `dataset_readiness.level` to chosen level -- Unprovided fields: `"N/A"` in YAML, brief "Not provided." in markdown body - -### 6. Review Summary - -Present: -- ✅ Auto-populated fields -- ✏️ User-provided fields -- ⬜ Empty/N/A fields (with reason) -- 💡 Suggestions for improvement - -Ask if the user wants to revise any sections. - -## References - -- **Template**: `references/datacard_template_v1.md` -- **Datasheet companion**: `references/datasheet_template.md` — consult if user mentions "datasheet" questions -- **Introspection commands**: [reference/introspection-commands.md](reference/introspection-commands.md) -- **Per-level field prompts**: [reference/field-prompts.md](reference/field-prompts.md) -- **Lookup tables** (OSTI codes, sensitivity tiers): [reference/lookup-tables.md](reference/lookup-tables.md) \ No newline at end of file diff --git a/src/dsagt/skills/datacard-generator/reference/datacard_template_v1.md b/src/dsagt/skills/datacard-generator/reference/datacard_template_v1.md deleted file mode 100644 index 3982f47..0000000 --- a/src/dsagt/skills/datacard-generator/reference/datacard_template_v1.md +++ /dev/null @@ -1,498 +0,0 @@ ---- -#YAML Metadata to support datacard metadata: (Instructions follow) -#----------------Datacard Information-- This section provides information about this datacard and its contents -datacard_info: - filename: modcon_datacard_${SNAKE_CASE_DATASET_NAME}.md # Follow this filename convention and enter the filename here - id: # PID assigned to this datacard if used. - type: #DOI | HANDLE | URL | LOCAL | OTHER - value: - datacard_access: #describe any access restrictions for this datacard, or "None - publicly accessible" - datacard_creation: - created_date: "${CREATED_DATE_YYYY_MM_DD}" # required - update_date: "${UPDATE_DATE_YYYY_MM_DD}" # add updates as needed - creation_method: #manual | automated | hybrid - created_by: [] # Contacts for this datacard. Include at least one (datacard_contact_person) or organizational contact (datacard_contact_organization) with email. - - datacard_contact_person: - name: - orcid: - email: - affiliation: - - datacard_contact_organization: - organization_name: - ror_id: - email: - datacard_template_version: 1.0 - language: en - -#-------------Data Identification-- Required. This section contains the information needed to link this datacard to the datasets it describes: -data_identifiers: - name: "${DATASET_NAME}" #provide a single, human readable name even if this datacard describes a collection of data. Ensure alignment with the datacard filename. - dataset_id: # Identify the dataset(s) described by this datacard. At least 1 identifier is required. Persistent, unique identifiers (PIDs) are preferred. - type: # DOI | HANDLE | URL | LOCAL | OTHER - value: - -#----- Level 1: Discoverable----- -#[Required Level1]: -# Identification -title: "${DATASET_NAME}" -id: # ID of the datacard if known. Otherwise this will be provided by the system. -project: # required - project name or identifier that the dataset is associated with - -#--- Data files/objects in dataset - -# for published datasets, include one or more dois. For unpublished datasets, provide a landing page url if available, or other identifier such as a local identifier or handle. -identifiers: [] # [required level1]. At least 1 identifier is required. - - type: # DOI | HANDLE | URL | LOCAL | OTHER - value: - - -#[Optional Level1]: -#--- Dataset Structure / Characteristics--- -dataset_info: - modalities: [] # e.g., ["tabular","image","time-series"] - data_formats: [] - features: [] # list of variables or features in the dataset. - splits: [] # e.g., ["train", "test", "validation"] - dataobject_type: dataset #Type of digital object (dataset, aiagent, eval, framework, model, etc.) - dataset_type: #See OSTI DOE Data Explorer Types - -#**OSTI DOE Data Explorer Dataset Types**: -#| Code | Type | Description | -#|------|------|-------------| -#| `GD` | Genome/Genetic Data | DNA/RNA sequences, genetic markers, genomic annotations | -#| `IM` | Image | Photographs, scans, visualizations, microscopy | -#| `ND` | Numeric Data | Measurements, time series, tabular data, sensor readings | -#| `SM` | Specialized Mix | Multiple data types combined | -#| `FP` | Figure/Plot | Charts, graphs, plots as primary deliverable | -#| `I` | Interactive Resource | Web applications, interactive visualizations, dashboards | -#| `MM` | Multimedia | Audio, video, combined media | -#| `MD` | Model | Computational models, simulations, trained ML models | -#| `AS` | Automated Software | Scripts, analysis pipelines, workflows | -#| `IP` | Instrumentation and Protocols | Experimental protocols, instrument specifications | -#| `IG` | Integrated Genomic Resources | Combined genomic databases and tools | - -dataset_readiness: - level: # 1 | 2 | 3 - evaluated_against: "Genesis Dataset Readiness Model v1.0" - evaluated_at: - evaluated_by: - confidence: # high | medium | low - - -#--- Lifecycle / Governance--- -release_status: # draft | approved | published | deprecated - -review_process: - review_purpose: # a short description of why is the data is being reviewed (e.g., "release", "partner sharing", "IRB",... ) - review_status: # "submitted" | "pending" | "approved" | "declined" | - review_institution: - name: # name of institution conducting the review - ror_id: # ROR ID of institution conducting the review - review_comments: - - -#------ Level 2: Interoperable and Reusable------- - -#[Required Level2]: -#--- Contacts -contact_point: - entity: - type: person # person - person: - given_name: - family_name: - orcid: - email: - affiliation: - entity: - type: organization # organization - organization: - name: - ror_id: - -additional_contacts: [] - -#--- Access, Permissions and Policy -access_policy: - # sensitivity_tier classification: - # tier0_openScience - # tier1_controlledResearch - # tier2_proprietary - # tier3_sensitive_exportControlled - # tier4_regulated_personal - # tier5_classified - sensitivity_tier: - access_level: # NOTE: currently only accepting "open", open | restricted | controlled - required NOTE: currently only accepting "open" - authorization: # none | accountRequired | userAgreement | dataUseAgreement | sponsorApproval | exportControlReview | irbApproval | other - policy_url: - policy_text: - -#[IF Applicable]: -cui_markings: # see CUI documentation -distribution_statement: #examples are: "Distribution A - Approved for public release"; "Distribution D - DoD only"; "Internal DOE use only" -handling_instructions: #examples are "No foreign dissemination" or "Export-controlled handling required" - -security_marking: - classification: # unclassified | cui | classified | confidential | Secret | TS | other - cui_marking: - distribution_statement: - handling_instructions: - declassification: - review_date: - authority: - -#--- Rights & License--- -license: - spdx_id: # SPDX ID is a standardized license identifier (e.g., CC-BY-4.0, MIT, Apache-2.0) - name: - link: - -additional_licenses: [] - -#[Optional Level2]: -# Domain and Purpose -categorization: - science_domain: "${SCIENCE_DOMAIN}" # required - high level scientific domain or discipline that the dataset is associated with, e.g., physics, chemistry, materials science, etc. - modalities: [] # e.g., ["tabular","image","time-series"] - task_category: [] # ML classification - task_subcategory: [] # ML subcategory classification - -# Resources -originating_research_organization: - entity: - type: organization # organization - organization: - name: - ror_id: - -facilities: [] -# list of facilities (as entities, type: organization) where the dataset was collected, processed, stored, or accessed. Examples include research labs, data centers, cloud platforms, supercomputing facilities, etc. - -fundings: -# list of funding sources that supported the creation, maintenance, or sharing of the dataset. Examples include grants, contracts, in-kind support, etc. - - funding: - award_number: - program: - funder: - entity: - type: organization # organization - organization: - name: - ror_id: - -#--- Dataset Responsibility & Credit -dataset_authors: - - entity: # at least 1 is required - type: person # person - person: - given_name: - family_name: - orcid: - email: - affiliation: - entity: - type: organization # organization - organization: - name: - ror_id: - role: creator - -dataset_contributors: [] # Acknowlege other contributors beyond dataset authors and include them as entities of type person or organization. - -#--- Stewardship & Maintenance--- -stewardship: - level: # projectManaged | repositoryManaged - -maintenance: - update_frequency: # none | ad_hoc | monthly | quarterly | annually | other - -# Provenance -#--- Dataset Provenance--- -dataset_provenance: - was_generated_by: - source_data: - processing_steps: - instrumentation: - simulation_details: - -#--- Related Resources (Optional but recommended) -related_resources: - related_datasets: [] - # - dataset_name: - # identifiers: - # - type: DOI | URL | LOCAL | OTHER - # value: - # relationship: isDerivedFrom | isBasedOn | isPartOf | hasPart | references | other - publications: - dois: [] - arxiv: [] - urls: [] - software: - # - name: - # version: - # identifiers: - # - type: DOI | URL | LOCAL | OTHER - # value: - # relationship: isDerivedFrom | isBasedOn | isPartOf | hasPart | references | other - aimodels: [] - # - name: #if available, provide the name of the AI model - # version: - # date accessed: - # identifiers: - # - type: DOI | URL | LOCAL | OTHER - # value: - # relationship: isDerivedFrom | isBasedOn | isPartOf | hasPart | references | other -#-- Methods -#--- Dates -dates: - data_collection_start: - data_collection_end: - issued: - modified: - - -#------ Level 3: Understandable & Trustworthy------- - -#[If Applicable Level3]: -#--- Semantic / Schema -semantic_layer: - # Describe formal schema / ontology alignment if available - # Leave blank if none exists (acceptable at Level 1 or Level 2) - schema_url: - ontology_alignment: [] - semantic_context: [] - controlled_vocabularies: [] - -#--- Data Quality -data_quality: - completeness: - known_issues: - validation_methods: - noise_characteristics: - uncertainty_notes: - -#--- Integrity / Fixity -integrity: - checksum_available: # true | false - checksum_type: - fixity_policy: - versioning_strategy: - -#--- AI / ML Usage -ai_usage: - ai_ready: # true | false | conditional - training_use_allowed: # true | false | conditional - inference_use_allowed: # true | false | conditional - restrictions: - bias_risks: - safety_considerations: - human_review_required: # true | false - -#---- System Level Info:---- - -#--- Dataset Scale -dataset_counts: - value: - category: # samples | files | records | timesteps | other - unit: - -dataset_storage: - compressed_bytes: - unpacked_bytes: - -#--- Repository-managed access endpoints--- -repository_access: - populated_by_repository: true - distributions: [] - # Example: - # - distribution: - # access_url: {} - # download_url: - # format: - # byte_size: - # checksum: - # type: [] - # items: [] - data_services: [] - # Example: - # - data_service: - # type: api | s3 | other - # additional_properties: false - # properties: - # endpoint: {} - # service_type: {} - # auth_hint: {} ---- - -### Instructions - - - - - - - - - - - -# Datacard for ${DATASET_NAME} -**Last Updated**: [!TODO] - -**Dataset Readiness Level:** - -### Machine Usability Snapshot -| Aspect | Status | -|--------|--------| -| AI Ready | Yes/No/Conditional| -| License Clarity | Yes/No| -| Machine Access | Yes/No| -| Checksum / Fixity | Yes/No| -| Semantic Context | Yes/No| - - -# ---- Level 1: Discoverable ---- ---- -## Identification - - -### Files & Structure -[!TODO] [required level1] - ---- -## Description - -### Dataset Description [required] -[!TODO] - -### Keywords [strongly recommended] -[!TODO] - -### Citation -[!TODO] - ---- - -# ---- Level 2: Interoperable and Reusable ---- - -### Sharing & Access -[!TODO] - -### Security / Marking Considerations -[!TODO] - ---- -### Access and Permissions -[!TODO] - -### Access conditions -[!TODO] - -### Release review process -[!TODO] - ---- -## Context - -### Domain and Purpose -[!TODO] - -### Resources used, including funding and facilities, to create the dataset -[!TODO] - ---- -## Provenance - -### Developed by -[!TODO] - -### Contributed by -[!TODO] - ---- -## Related Resources - -### Related datasets, standards, metadata, and ontologies -[!TODO] - -### Related publications -[!TODO] - -### Related software -[!TODO] - -### Related ai model -[!TODO] - ---- -## Methods - -### Dataset generation, collection, and procedures -[!TODO] - -### Maintenance & Updates -[!TODO] - - -# ---- Level 3: Understandable & Trustworthy ---- - -### Data Characteristics -[!TODO] - -### Data Quality & Limitations -[!TODO] - -### Related Schemas or Ontologies -[!TODO] - -### List of variable name(s), description(s), unit(s), and value labels for each variable in the dataset/file. -[!TODO] - -For example: -| Variable Name | Description | Unit | Value Labels | -|---------------|---------------------------|-----------|-----------------------------| -| temp | Temperature measurement | Celsius | N/A | -| status | Operational status | N/A | 0 = Off, 1 = On | - -### Codes used for missing data -[!TODO] - -For example: -| Code | Description | -|------|---------------------------| -| -999 | Data not collected | -| -888 | Measurement error | - -### Specialized formats or other abbreviations used -[!TODO] - -### Example of the contents -[!TODO] - -### Data Processing -[!TODO] - -### Software used to preprocess/ clean/ label the data -[!TODO] - -## Integrity & Versioning -[!TODO] - -## Semantic / Schema Information -[!TODO] - -## AI / Machine Learning Considerations -[!TODO] - ---- - -## Additional Information -[!TODO] - ---- \ No newline at end of file diff --git a/src/dsagt/skills/datacard-generator/reference/datasheet_template.md b/src/dsagt/skills/datacard-generator/reference/datasheet_template.md deleted file mode 100644 index c69afe9..0000000 --- a/src/dsagt/skills/datacard-generator/reference/datasheet_template.md +++ /dev/null @@ -1,320 +0,0 @@ ---- -language: -- en # ISO language tag -datasheet_version: 0.1.0 # version of the GENESIS datasheet template -name: Datasheet for {DATASET_NAME} # name of the datasheet -tags: -- project:genesis # include on all GENESIS project datasets -- science:lightsource # what kind of science is this for (e.g., materials, biology, lightsource, fusion, climate, etc.) -- keywords: [] # keywords associated with the dataset -- risk:general # indicates level of risk review {general, reviewed, restricted} ---- -# Datasheets for Datasets: Contextualizing Scientific Data - -Documenting a scientific dataset in a datasheet requires careful consideration of scope and resolution. -Scientific data often exists as part of a larger ecosystem of the scientific record: a single experiment may produce multiple datasets, or a dataset may be split into subsets for different analyses. -Deciding at what level to create a datasheet is not always straightforward—too granular, and you risk an unmanageable number of datasheets; too broad, and important details may be lost. -Consider how the dataset will be used and cited, who the intended consumers are (e.g., researchers, search engines, AI systems), and what level of documentation will provide meaningful transparency without unnecessary duplication. -Keep in mind that datasheets are not necessarily one-to-one with dataset DOIs: a datasheet may describe a single dataset, a collection of related datasets, or a versioned release. -Before you begin, reflect on these relationships and choose a resolution that balances clarity, usability, and sustainability. -This is not a prescriptive process; it deserves deliberate consideration. - -## Before you begin - -## Datasheet Structure - The datasheet is organized into the following sections: - - Motivation - - Composition - - Collection - - Preprocessing/Cleaning/Labeling - - Uses - - Distribution - - Maintenance - - Human Subject Research (if applicable) - -## Think about your dataset's modalities -Before filling out the datasheet, consider the different data modalities present in your dataset (e.g., sensor time series, microscopy images, simulation results, logbook text). -Each modality may have unique characteristics, collection methods, and intended uses that should be documented separately. -When addressing questions in the datasheet, specify which modality you are referring to, especially if the dataset contains multiple modalities. -This will help ensure clarity and provide a comprehensive understanding of the dataset's structure and content. - -## Think about your dataset's intended uses -Before filling out the datasheet, reflect on the intended uses of your dataset. -Consider the scientific tasks it is designed to support, the research gaps it aims to address, and the theories it seeks to test. -Understanding the intended applications will help you provide more accurate and relevant information in the datasheet, particularly in sections related to motivation, uses, and limitations. -This reflection will also guide you in identifying any potential risks or ethical considerations associated with the dataset's use. - -## Consider the scope and resolution of your datasheet -Before filling out the datasheet, carefully consider the scope and resolution at which you will document your dataset. -Decide whether the datasheet will cover a single dataset, a collection of related datasets, or a versioned release. -Think about the level of detail that will be most useful for your intended audience, balancing clarity and usability with the need to avoid unnecessary duplication. -This consideration will help ensure that the datasheet effectively communicates the essential information about your dataset while remaining manageable and relevant. - -## Consider the those filling out the datasheet -Before filling out the datasheet, think about who will be completing it. -Ensure that the individuals responsible for authoring the datasheet have a comprehensive understanding of the dataset, including its creation, composition, and intended uses. -Encourage collaboration among team members who contributed to the dataset to provide accurate and detailed responses. -This collaborative approach will help ensure that the datasheet reflects a well-rounded perspective and captures all relevant information about the dataset. - -### Guidelines -1. Recommend that there be a citable ``scientific record" that contains the Datasheet and Dataset and other elements of the scientific record, including facilities, authors, publications, software, datasets, schemas, and datasheets. -2. Recommend assigning a separate DOI to the Datasheet for independent versioning and authorship tracking, apart from related manuscripts and datasets. -3. Recommend providing a version of this datasheet with information that is suitable for the widest possible distribution and citation allowed under applicable laws. -If applicable and to the extent permissible in a document for such wide distribution, the datasheet should cite one or more controlled (at various classification levels, proprietary, or export controlled) versions of itself, and individual responses should indicate that the controlled version has different information. - Such different ``versions" of the document should have their own unique persistent identifiers. -4. Recommend a broad reading of the text of the questions, and provide partial answers if some parts cannot be provided. - For example, if a grant from a private foundation lacks an associated grant identifier, the entity's name and a description of the grant should be provided. - If providing the name of the grantor or the purpose of creation is not permissible at the wide distribution level, that portion can be left unanswered with an explicit marking of ``not available." -5. item If a combination of responses is not permitted (e.g., indicating the title and the motivation, or the motivation and the funding agency simultaneously may not be permitted) at the wide distribution level, then the topics describing the dataset and its allowed uses are to be given preference over the provenance. -6. In situations where there might be some uncertainty, make sure the datasheet is reviewed by a derivative classifier. - -**[[REMOVE "Datasheets for Datasets: Contextualizing Scientific Data" section prior to sharing the Datasheet]]** - -# Datasheet for ${DATASET_NAME} - -## Dataset Metadata - -### Dataset citation: - -Include DOI or URL. Use information that matches a data catalog entry if applicable. Citations in [bibtex format](https://www.bibtex.com/g/bibtex-format/). Please include either a `doi` or `url` field in the citation. - -### Name and contact information for the datasheet author(s): - -A person or group that was primarily responsible for authoring the datasheet. Provide the Name, [ORCID](https://orcid.org/), affiliation (ROR ID) and email address of the person or group responsible for the datasheet. - -## Motivation - -### What was the primary purpose for creating this dataset (e.g., to support a specific scientific task, address a research gap, or test a theory)? -Describe the intended use and, if applicable, any assumptions made about the underlying system (e.g., proxy variables, stationarity, symmetry, theoretical models). - -### Was the dataset created for or in the context of an AI application? -Please provide a description of the AI application, including the type of model(s) that will be trained or evaluated using this dataset. - -### Who created the dataset (e.g., which team, research group), what are their institutional affiliations, and on behalf of which research entity (e.g., company, institution, organization) was the dataset created? - -Provide the [ROR ID](https://ror.org/) for research entities if available. - -### What resources were used? For example, what funding, facility time, computing resources, and datasets were used to create the dataset? Provide answers in the sections below. - -#### Facilities: - -list the facilities. Provide the facility name and [ROR ID](https://ror.org/) if available. - -#### Funding: - -If there are associated grants, please provide the name of the grantor(s) (e.g., federal agency and program office name) and [ROR ID](https://ror.org/) , and the grant name(s) and number(s) (e.g., PAMS Award number and/or lab contract number). List the solicitation numbers and links, or citations. Citations in [bibtex format](https://www.bibtex.com/g/bibtex-format/). Please include either a `doi` or `url` field in the citation. - -#### Other Supporting Entities: - -If available, provide a [ROR ID](https://ror.org/), link, or citation for other supporting entities. Citations in [bibtex format](https://www.bibtex.com/g/bibtex-format/). Please include either a `doi` or `url` field in the citation. - -### Was the dataset created as part of a larger initiative? -If so, provide the name for the initiative, e.g., SciDAC, ModelTeam, Genesis. - -### Any other comments regarding the motivation for creating the dataset? (optional) -Provide any additional information about the motivation for creating the dataset. - -## Composition - -### Identify the data modalities contained in the dataset (e.g., sensor time series, microscopy images, simulation results, logbook text). -If multiple modalities exist, list each one. - -### What do the instances within each modality represent (e.g., detector hits, environmental monitoring readings, simulation runs, documents)? Are there multiple types of instances (e.g., beam pulses and cavity field maps; users and their proposals)? -Provide a description for each modality. - -### How many instances are there in total (of each type, if appropriate)? -Provide the number of instances for each modality. - -### Does the dataset include all possible instances for each modality or a subset (not necessarily random)? -If a subset, what is the larger collection (e.g., all experiments at a user facility during a given period)? Is the subset intended to be representative (e.g., energy ranges, geographic sites)? Describe the criteria and validation methods used to verify representativeness. -If it is not representative, please describe why not (e.g., to capture diversity in experimental conditions; due to data access restrictions; because data was intentionally excluded, and why). - -### What data does each instance within the dataset consist of (for each modality, if appropriate)? "Raw" data (e.g., unprocessed text or images) or features? -In either case, please provide a description. - -### Is any information missing from individual instances? -If so, please provide a description explaining why this information is missing (e.g., because it was unavailable). This does not include intentionally removed information but may include, e.g., redacted text or missing sensor data. - -### Are relationships between individual instances made explicit (e.g., parent-child relationships in a tree structure, protein–protein interaction networks)? -If so, please describe how these relationships are made explicit. - -### Have the data been split, or are there recommended data splits (e.g., training, development/validation, testing)? -If so, please provide a description of these splits, explaining the rationale behind them. - -### Are there any errors, sources of noise, or redundancies in the dataset? -If so, please provide a description. - -### Are external resources (e.g., file stores, calibration information, software packages, code, websites, other datasets) needed to interpret or reuse these data? -If not, respond "No." -Otherwise, for each external resource, provide a description, and if applicable, its access point (e.g., DOIs, link, [bibtex format](https://www.bibtex.com/g/bibtex-format/) citation). Please include either a `doi` or `url` field in the citation. - -#### If there are external resources, for each external resource, are there any restrictions (e.g., licenses, fees)? - -### Does the dataset contain data that might be considered Controlled Unclassified Information (CUI) or confidential (e.g., Personally Identifiable Information (PII), Confidential Business Information (CBI), pre-publication data, data protected by legal privilege, or data that includes the content of individuals' non-public communications)? -Refer to [Open Digital Rights Language (ODRL) Vocabulary & Expression](https://www.w3.org/TR/odrl-vocab/) and the [Information Model](https://www.w3.org/TR/odrl-model/) for a standard vocabulary to represent statements about the usage of data. - -### Does the dataset contain classified or export-controlled information? -If so, please provide an explanation. - -### Any other comments regarding the composition of the dataset? (optional) -Provide any additional information about the composition of the dataset. - -## Collection - -### How was the data for each instance obtained or generated (e.g., raw experimental measurements from user facilities, processed/physics-ready experimental data, outputs from computational simulations, or data derived from prior datasets)? -If the data were simulated or derived, provide details on how accuracy and validity were assessed. - -### Did the dataset incorporate data obtained through unconventional means (e.g., online crowdsourcing, volunteer-based citizen science)? -If yes, explain how these contributions were managed, including documentation, acknowledgment, and quality assurance measures. - -### For each instrument or facility used to generate and collect the data, what mechanisms or procedures were used to collect the data (e.g., hardware apparatuses or sensors, manual human curation, software programs, software APIs)? -Provide a description for each instrument or facility used, mechanisms and procedures used to collect the data. Provide relevant details on how these mechanisms or procedures were validated. - -### If the dataset is a sample from a larger set, what was the sampling strategy (e.g., deterministic, probabilistic with specific sampling probabilities)? -Provide a description of the sampling strategy and any methods used to assess sample representativeness. - -### Over what timeframe was the data collected or generated? Does this timeframe align with when the underlying phenomena or events occurred (e.g., recent simulation of historical accelerator configurations)? If not, describe the timeframe of the original events or conditions represented by the data. -Provide single date, range, or approximate date; suggested format YYYY-MM-DD, and, if appropriate, describe alignment with underlying phenomena or events - -### Any other comments regarding the dataset collection process? (optional) -Provide any additional information about the collection process for the dataset. - -## Preprocessing/cleaning/labeling - -### To create the final dataset, was any preprocessing/ cleaning/ labeling of raw data done? - -Preprocessing, cleaning, and labeling can include simple pipelines (e.g., discretization or bucketing, tokenization, part-of-speech tagging, SIFT feature extraction, removal of instances, processing of missing values), or more complex workflows including noise reduction and filtering, calibration, reconstruction and event building, simulation post-processing, meta-data enrichment (e.g., adding uncertainty estimates). -If so, please provide a description of the workflow. - -### Was the "raw" data saved in addition to the preprocessed/cleaned/labeled data (e.g., to support unanticipated future uses)? -If so, please provide a [bibtex format](https://www.bibtex.com/g/bibtex-format/) citation, link, or other access point to the "raw" data. Please include either a `doi` or `url` field in the citation. - -### Is the software that was used to preprocess/ clean/ label the data available? -If so, please provide a [bibtex format](https://www.bibtex.com/g/bibtex-format/) citation, PID, link, or other access point, along with descriptions of any required packages or libraries to run the scripts. Please include either a `doi` or `url` field in the citation. - -### Any other comments regarding the preprocessing, cleaning, and labeling workflow for the dataset? (optional) -Provide any additional information about the preprocessing, cleaning, and labeling workflow for the dataset. - -## Uses - -### Have AI-Readiness assessment tools been used? -If so, name the tool, and describe the results? - -### To what extent are the data prepared for AI? -Describe any steps taken to make the data suitable for AI applications, such as formatting, normalization, or feature extraction. - -### If the data has been used for AI, provide a description. -Highlight how the data was used for AI, including links to code or DOIs illustrating its use. - -### What (other) tasks could the dataset be used for? -Provide examples of tasks that the dataset could be used for, beyond its original purpose. - -### Are there limitations or characteristics, arising from the collection, preprocessing, or labeling, that may impact future uses of the dataset (e.g., incomplete coverage of parameter space, data quality variations, restrictions due to CUI or export control)? -If so, please provide a description. Describe mitigation strategies for the risks? - -### Are there tasks for which the dataset should not be used? -Describe tasks not recommended for this dataset's use. - -### Were any reviews for safety, cybersecurity, export control, or other considerations conducted? -If so, please provide a description of these review processes, including the outcomes. - -### Are there tutorials or other supporting information to assist a user of these data? -If so, please provide a description and [bibtex format](https://www.bibtex.com/g/bibtex-format/) citation, PID, link, or other access point. Please include either a `doi` or `url` field in the citation. - -### Any other comments on the use of the dataset? (optional) -Provide any additional information about the use of the dataset. - -## Distribution - -### Will all or parts of the dataset be shared outside the originating laboratory, facility, or site (e.g., with other DOE sites, collaborating institutions, or public repositories)? -If so, please provide a description of how the dataset will be shared, including any access mechanisms (e.g., data repositories, data catalogs, or other data sharing platforms). Provide relevant links or citations in [bibtex format](https://www.bibtex.com/g/bibtex-format/). Please include either a `doi` or `url` field in the citation. - -### Indicate the date or timeframe when the dataset was (or will be) made available. -Provide a single date or range; suggested format YYYY-MM-DD. - -### Access and reuse restrictions placed on the data: -Will the dataset be shared under a copyright or other intellectual property (IP) license, and/or under applicable terms of use (ToU)? -If your dataset is split into multiple parts (e.g., training and test sets), you may need to answer separately for each part. -If so, please describe this license and/or ToU, and provide a link or other access point to, or otherwise reproduce, any relevant licensing terms or ToU, as well as any fees associated with these restrictions. - -### Have any third parties imposed IP-based or other restrictions on the data associated with the instances? -If so, please describe these restrictions and provide a link or other access point to, or otherwise reproduce, any relevant licensing terms, as well as any fees associated with these restrictions. - -### Export control restrictions: -Do any export controls or other regulatory restrictions apply to the dataset or to individual instances? -If so, please describe these restrictions and provide a link or other access point to, or otherwise reproduce, any supporting documentation. - -### If the dataset or any individual instances have a classification or restriction, what is the specific level of classification or restriction (e.g., Classified, Unclassified, CUI)? -If there is no classification or restriction, respond "N/A." - -### Any other comments regarding the distribution of the dataset? (optional) -Provide any additional information about the distribution of the dataset. - -## Maintenance - -### Who will maintain and provide access to the dataset after its creation? -For example, DOE-supported repository, facility staff, or facility data steward. -Provide responsible parties’ names, roles, contact information, and persistent identifiers (e.g., [ORCID](https://orcid.org/)). -Indicate any DOE policies or agreements governing maintenance and retention. - -### Is there an erratum? -If so, please provide a link or other access point. - -### Will the dataset be updated (e.g., to correct labeling errors, add new instances, delete instances)? -If so, please describe how often, by whom, and how updates will be communicated to dataset consumers (e.g., mailing list, GitHub)? - -### Will older versions of the dataset continue to be supported/hosted/maintained? -If so, please describe how. If not, please describe how its obsolescence will be communicated to dataset consumers. - -### What is the expected lifespan of the dataset? -Provide a description of the expected lifespan, including any plans for archiving or decommissioning, including stakeholders, communities, and collaborators who will be consulted. If a retention policy applies, please provide a link or other access point to, or otherwise reproduce, the relevant documentation. - -### Will the dataset support controlled extensions or updates from collaborators? -If yes, describe the mechanism for submission, validation, and versioning, and how these changes will be communicated to users.Otherwise, explain why contributions are restricted. - -### Any other comments regarding the maintenance of the dataset? (optional) -Provide any additional information about the maintenance of the dataset. - -## Human Subject Research - -### Does the dataset include personal information (e.g., identifiable data, survey responses)? -If not, respond "No" and ignore the remainder of the Datasheet. -Otherwise, respond "Yes" and continue. - -### Human Subject Research: Dataset Composition - -#### Does the dataset identify any subpopulations (e.g., by age, gender, career stage)? -If not, answer "No." Otherwise, please describe how these subpopulations are identified and provide a description of their respective distributions within the dataset. - -#### Is it possible to identify individuals (i.e., one or more natural persons), either directly or indirectly (i.e., in combination with other data) from the dataset? -If not, answer "No." Otherwise, please describe how. - -#### Does the dataset contain data that might be considered sensitive in any way (e.g., data that reveals race or ethnic origins, sexual orientations, religious beliefs, political opinions or union memberships, or locations; financial or health data; biometric or genetic data; forms of government identification, such as social security numbers; criminal history)? -If not, answer "No." Otherwise, please provide a description. - -### Human Subject Research: Collection - -#### Were any ethical review processes conducted (e.g., by an institutional review board)? -If so, please provide a description of these review processes, including the outcomes, and a link or other access point to any supporting documentation. - -#### Were the individuals in question notified about the data collection? -If so, please describe (or show with screenshots or other information) how notice was provided, and provide a link or other access point to, or otherwise reproduce, the exact language of the notification itself. - -#### Did the individuals in question consent to the collection and use of their data? -If so, please describe (or show with screenshots or other information) how consent was requested and provided, and provide a link or other access point to, or otherwise reproduce, the exact language to which the individuals consented. - -#### If consent was obtained, were the consenting individuals provided with a mechanism to revoke their consent in the future or for certain uses? -If so, please provide a description and a link or other access point to the mechanism (if applicable). - -#### Has an analysis of the potential impact of the dataset and its use on data subjects (e.g., a data protection impact analysis) been conducted? -If so, please provide a description of this analysis, including the outcomes, as well as a link or other access point to any supporting documentation. - -### Human Subject Research: Maintenance - -#### Are there applicable limits on the retention of the data associated with the instances (e.g., were the individuals in question told that their data would be retained for a fixed period of time and then deleted)? -If so, please describe these limits and explain how they will be enforced. - -### Human Subject Research: Other - -#### Any other comments regarding the inclusion of human subject research in the dataset? (optional) -Provide any additional information about the inclusion of human subject research in the dataset. diff --git a/src/dsagt/skills/datacard-generator/reference/field-prompts.md b/src/dsagt/skills/datacard-generator/reference/field-prompts.md deleted file mode 100644 index 5be7460..0000000 --- a/src/dsagt/skills/datacard-generator/reference/field-prompts.md +++ /dev/null @@ -1,38 +0,0 @@ -# Per-Level Field Prompts - -Ask 3–5 fields at a time. Only prompt for fields not already auto-filled. - -## Level 1 — Discoverable - -- Dataset name and title -- Project name -- Dataset identifiers (DOI, URL) -- Dataset description -- Keywords -- Release status - -## Level 2 — Interoperable & Reusable (adds to Level 1) - -- Contact point (name, email, affiliation, ORCID) -- Access policy (sensitivity tier, access level, authorization) — see lookup-tables.md for sensitivity tiers -- License (SPDX ID, name, URL) -- Science domain -- Originating research organization -- Funding sources -- Dataset authors and contributors -- Provenance — how was the data generated or collected? -- Related resources (datasets, publications, software) -- Maintenance and stewardship plans -- Key dates (collection start/end, issued, modified) - -## Level 3 — Understandable & Trustworthy (adds to Level 2) - -- Semantic/schema info — ontologies, controlled vocabularies (must be user-provided) -- Data quality — completeness, known issues, validation methods, noise, uncertainty -- Integrity — checksums, fixity policy, versioning strategy -- AI/ML usage — AI-ready status, training/inference permissions, bias risks, safety notes -- Variable-level documentation — name, description, unit, value labels (one row per variable) -- Missing data codes -- Data processing steps - -**If the user selects Level 3 but provides minimal information:** generate the card with N/A for unprovided fields, list them in the review summary, and note which ones most improve trustworthiness. \ No newline at end of file diff --git a/src/dsagt/skills/datacard-generator/reference/introspection-commands.md b/src/dsagt/skills/datacard-generator/reference/introspection-commands.md deleted file mode 100644 index 0fd9517..0000000 --- a/src/dsagt/skills/datacard-generator/reference/introspection-commands.md +++ /dev/null @@ -1,80 +0,0 @@ -# Introspection Commands - -## Contents -- File structure and size -- Schema and data sampling (CSV, Parquet, JSON, HDF5/NetCDF, images) -- Existing metadata files -- Directory structure and splits -- Edge case handling - -## File Structure and Size - -```bash -# Full recursive file listing -find -type f | sort - -# Count files by extension -find -type f | sed 's/.*\.//' | sort | uniq -c | sort -rn - -# Total size (human-readable) -du -sh - -# Size of top-level items -du -sh /* -``` - -## Schema and Data Sampling - -```bash -# CSV/TSV — headers, row count, dtypes -head -n 3 -wc -l -python3 -c "import pandas as pd; df = pd.read_csv('', nrows=5); print(df.dtypes); print(df.shape)" - -# Parquet -python3 -c "import pandas as pd; df = pd.read_parquet(''); print(df.dtypes); print(df.shape)" - -# JSON — top-level structure -python3 -c "import json; d=json.load(open('')); print(type(d), list(d.keys()) if isinstance(d,dict) else f'{len(d)} items')" - -# HDF5 -python3 -c "import h5py; f=h5py.File('','r'); print(list(f.keys())); f.visititems(lambda n,o: print(n, type(o).__name__))" - -# NetCDF -python3 -c "import netCDF4 as nc; d=nc.Dataset(''); print(d.variables.keys()); print(d.dimensions)" - -# Images — count by format -find -type f \( -iname "*.jpg" -o -iname "*.png" -o -iname "*.tif" \) | wc -l -``` - -For files larger than 1GB, sample only (e.g., `nrows=1000`) and note this in the review summary. - -## Existing Metadata Files - -```bash -# Check for common metadata files -for f in README.md README.txt LICENSE CITATION.cff CITATION.bib .zenodo.json datacite.json croissant.json metadata.json; do - [ -f "/$f" ] && echo "FOUND: $f" -done - -# Print contents of found files -cat /README.md 2>/dev/null -cat /CITATION.cff 2>/dev/null -cat /LICENSE 2>/dev/null -``` - -## Directory Structure and Splits - -```bash -# Top-level listing and directory tree (2 levels) -ls / -find -maxdepth 2 -type d | sort -``` - -Look for `train/`, `test/`, `validation/`, or file naming patterns like `_train.csv`. - -## Edge Case Handling - -- **Empty directory**: Inform the user; ask whether to proceed with a skeleton card -- **Binary-only files**: Note under "Files & Structure" and mark schema fields N/A -- **Large files (>1GB)**: Sample only; note in the review summary \ No newline at end of file diff --git a/src/dsagt/skills/datacard-generator/reference/lookup-tables.md b/src/dsagt/skills/datacard-generator/reference/lookup-tables.md deleted file mode 100644 index e549645..0000000 --- a/src/dsagt/skills/datacard-generator/reference/lookup-tables.md +++ /dev/null @@ -1,32 +0,0 @@ -# Lookup Tables - -## OSTI Dataset Type Codes - -Use when setting `dataset_info.dataset_type`: - -| Code | Type | -|------|------| -| `GD` | Genome/Genetic Data | -| `IM` | Image | -| `ND` | Numeric Data | -| `SM` | Specialized Mix | -| `FP` | Figure/Plot | -| `I` | Interactive Resource | -| `MM` | Multimedia | -| `MD` | Model | -| `AS` | Automated Software | -| `IP` | Instrumentation and Protocols | -| `IG` | Integrated Genomic Resources | - -## Sensitivity Tiers - -Use when setting `access_policy.sensitivity_tier`: - -| Tier | Label | -|------|-------| -| `tier0` | Open Science | -| `tier1` | Controlled Research | -| `tier2` | Proprietary | -| `tier3` | Sensitive / Export Controlled | -| `tier4` | Regulated / Personal | -| `tier5` | Classified | \ No newline at end of file diff --git a/src/dsagt/skills/skill-creator/SKILL.md b/src/dsagt/skills/skill-creator/SKILL.md new file mode 100644 index 0000000..00f93f6 --- /dev/null +++ b/src/dsagt/skills/skill-creator/SKILL.md @@ -0,0 +1,65 @@ +--- +name: skill-creator +description: "Author a new agent Skill (a SKILL.md directory in the open Agent Skills format) from the Anthropic template. Use when the user wants to create a skill, scaffold a SKILL.md, package a repeatable workflow as a reusable skill, turn instructions into a skill, or capture a procedure so the agent can auto-invoke it later. Produces a valid SKILL.md (name + description frontmatter, optional scripts/ and references/) and saves it into the project's skills directory." +metadata: + version: "1.0" +--- + +# Skill Creator + +Scaffold a new, spec-valid agent Skill from the Anthropic template and save it so the agent can discover and auto-invoke it. + +A *skill* is a directory `/SKILL.md`: YAML frontmatter (`name`, `description`) plus markdown instructions the agent follows when the description matches the task. It can bundle `scripts/` and `references/`. See [references/agent_skills_spec.md](references/agent_skills_spec.md) for the full contract. + +## Workflow + +Copy this checklist and check off steps as you go: + +``` +Progress: +- [ ] 1. Gather skill intent (name, purpose, triggers) +- [ ] 2. Draft from the template +- [ ] 3. Write the body (instructions/workflow) +- [ ] 4. Add scripts/ and references/ if needed +- [ ] 5. Validate the frontmatter +- [ ] 6. Save into the project (save_skill) +- [ ] 7. Confirm + note how it activates +``` + +### 1. Gather Intent + +Ask the user (or infer from context): +- **name** — short, lowercase, hyphenated (e.g. `convert-vasp-outputs`). This becomes the directory name and the invocable name. +- **purpose** — one sentence on what the skill does. +- **triggers** — the user requests / phrasing that should make the agent reach for this skill. These become keywords in the `description`. + +### 2. Draft From the Template + +Start from [references/SKILL_template.md](references/SKILL_template.md). Fill the frontmatter: +- `name`: must equal the directory name. +- `description`: pack it with *what it does AND when to use it* (trigger phrases) — this is the only thing the agent sees when deciding to invoke. Keep it ≤ 1536 characters. + +### 3. Write the Body + +After the frontmatter, write the instructions the agent will follow. Prefer a copyable checklist (like this one) for multi-step workflows. Reference bundled files by relative path, e.g. `[reference](references/notes.md)`, or run a bundled script with `${CLAUDE_SKILL_DIR}/scripts/foo.py` so paths resolve regardless of cwd. + +### 4. Add Supporting Files (optional) + +- `scripts/` — runnable helpers the body invokes. +- `references/` — long docs/templates loaded on demand (keep them OUT of SKILL.md so they cost no tokens until used). + +### 5. Validate + +Confirm before saving: +- Frontmatter is valid YAML between `---` fences. +- `name` is present, lowercase-hyphenated, and equals the intended directory name. +- `description` is present and ≤ 1536 characters. +- Any `[link](references/...)` and `${CLAUDE_SKILL_DIR}/scripts/...` paths exist. + +### 6. Save + +Save via the **`save_skill`** MCP tool with the `spec` (frontmatter dict: `name`, `description`, optional `tags`), the `body` markdown, and any `reference_files` (a `{relative_path: contents}` map). This writes `/skills//` and mirrors it into the platform's native skill directory (e.g. `.claude/skills/`) immediately (no KB indexing — `search_skills` is for the not-yet-installed catalog). + +### 7. Confirm + +Tell the user the skill was saved and is ready to use. Future sessions auto-discover it natively; to apply it in the current session, read `skills//SKILL.md` and follow it — that is all native invocation does. Do not tell the user a restart or any other action is needed. diff --git a/src/dsagt/skills/skill-creator/references/SKILL_template.md b/src/dsagt/skills/skill-creator/references/SKILL_template.md new file mode 100644 index 0000000..99db025 --- /dev/null +++ b/src/dsagt/skills/skill-creator/references/SKILL_template.md @@ -0,0 +1,53 @@ +# SKILL.md template + +Copy the block below into `/SKILL.md` and fill it in. Only the +frontmatter `name` + `description` are required; everything else is +optional. (Based on the Anthropic skill template / open Agent Skills +standard — https://github.com/anthropics/skills/tree/main/template.) + +```markdown +--- +name: my-skill-name +description: A clear description of WHAT this skill does and WHEN to use it — include the user phrasings/triggers that should invoke it. (≤ 1536 chars; this is the only text the agent sees when deciding to invoke.) +# optional: +# tags: [domain, keyword] +# metadata: +# version: "1.0" +--- + +# My Skill Name + +One or two sentences framing the task this skill handles. + +## Workflow + +Copy this checklist and check off steps as you go: + +``` +Progress: +- [ ] 1. ... +- [ ] 2. ... +``` + +### 1. ... + +Step-by-step instructions. Reference bundled docs by relative path: +[details](references/details.md). Run bundled scripts with an absolute +skill-dir path so cwd doesn't matter: + + python ${CLAUDE_SKILL_DIR}/scripts/helper.py + +## Notes / Guidelines +- ... +``` + +## Optional bundled files + +``` +my-skill-name/ +├── SKILL.md (required) +├── references/ (long docs/templates, loaded on demand) +│ └── details.md +└── scripts/ (runnable helpers the body invokes) + └── helper.py +``` diff --git a/src/dsagt/skills/skill-creator/references/agent_skills_spec.md b/src/dsagt/skills/skill-creator/references/agent_skills_spec.md new file mode 100644 index 0000000..bd3f79a --- /dev/null +++ b/src/dsagt/skills/skill-creator/references/agent_skills_spec.md @@ -0,0 +1,48 @@ +# Agent Skills — condensed contract + +A *skill* packages instructions (and optionally code/docs) so an agent can +discover and follow a repeatable workflow. dsagt skills follow the open +Agent Skills standard, which is what Claude Code, Cursor, Codex, and +Antigravity all read — so one SKILL.md works across platforms. + +## Directory layout + +``` +/ +├── SKILL.md # required — frontmatter + instructions +├── references/ # optional — docs/templates loaded on demand +└── scripts/ # optional — runnable helpers +``` + +- The **directory name is the invocable name** (e.g. `.claude/skills/deploy/` → `/deploy`). Keep it lowercase, hyphenated. + +## Frontmatter + +YAML between `---` fences. Common fields: + +| Field | Required | Notes | +|-------|----------|-------| +| `name` | recommended | Should equal the directory name. Lowercase-hyphenated. | +| `description` | **yes (in practice)** | What it does AND when to use it (trigger phrases). The agent sees only this when deciding to invoke. **≤ 1536 characters.** | +| `tags` | no | List of keywords; dsagt uses these for `search_skills` tag filters. | +| `metadata` | no | Free-form (e.g. `version`). Ignored by the platform. | +| `license` | no | Free-form. Ignored by the platform. | + +Unknown/extra frontmatter fields are **silently ignored** by Claude Code, so dsagt-specific fields are safe to include. + +## How discovery works + +- At session start, each installed skill's `name` + `description` are loaded into the agent's context. The full SKILL.md body loads only when the skill is invoked (lazy — zero cost until used). +- The agent auto-invokes a skill when the `description` matches the task; the user can also invoke it directly (`/skill-name`). +- A **newly created** top-level skills directory is only picked up after the agent restarts. + +## Body conventions + +- Lead with a copyable progress checklist for multi-step workflows. +- Keep long material in `references/` (loaded on demand) rather than inline, to save context tokens. +- Reference bundled files by relative path, or run scripts via `${CLAUDE_SKILL_DIR}/scripts/...` so paths resolve regardless of working directory. + +## Two tiers in dsagt + +- **Catalog** — skills indexed from external GitHub source repos, searchable via `search_skills` but not installed. Not in context. +- **Installed** — skills in `/skills/` (saved via `save_skill` or installed via `install_skill`). Mirrored into the platform's native skill dir at `dsagt start`, then natively discovered. diff --git a/src/dsagt/traces.py b/src/dsagt/traces.py new file mode 100644 index 0000000..7818bf1 --- /dev/null +++ b/src/dsagt/traces.py @@ -0,0 +1,1395 @@ +""" +Trace pipeline — read each agent's on-disk session, normalize it to one common +``Trace``, and hand that to whatever consumes traces (the MLflow logger in +:mod:`dsagt.observability`, the episodic-memory indexer in :mod:`dsagt.memory`). + +While a session runs, the MCP server wakes on a timer and, for the running agent: +a **Reader** finds and reads the platform's session files on disk into raw +records; the matching **Translator** maps those records into one **Trace**; and +the **TraceCollector** hands the result to its consumers, skipping the turns it +already handled. + + A session ``Trace`` carries one AGENT subtree per turn + (an AGENT root with ``llm`` / ``tool_`` +children), matching the per-prompt granularity MLflow's own claude autolog +produces. Fidelity is capped by what the transcript persisted: every timestamp +and token count is ``None``-tolerant. + +The Claude grammar below is ported from MLflow's +``claude_code/tracing.py`` — © Databricks, Inc., Apache-2.0 — specifically its +turn-windowing, skill/command skips, and next-timestamp span durations. See +NOTICE. + +Class map — ``▷`` inherits · ``◆`` owns · ``◇`` holds (``*`` = many):: + + Trace one session: id fields + spans (list of dicts) + + compose / query / to_exchanges methods + + Reader «abstract» locate + read a platform's session → raw records + ├─▷ JsonlReader «abstract» shared whole-file line framing; active_file() hook + │ ├─▷ ClaudeReader ~/.claude/projects//*.jsonl (newest) + │ └─▷ CodexReader ~/.codex/sessions/**/rollout-*.jsonl (by cwd) + ├─▷ GooseReader goose sessions.db (sqlite, read-only) + ├─▷ OpenCodeReader opencode.db (sqlite, read-only) + └─▷ ClineReader ~/.cline/.../.messages.json (whole file) + + Translator «abstract» raw records → Trace (pure); shared turn template + ├─▷ ClaudeTranslator overrides translate() — bespoke grammar + ├─▷ CodexTranslator fills parse hooks (+ normalize / prompt-index) + ├─▷ GooseTranslator fills parse hooks + ├─▷ OpenCodeTranslator fills parse hooks + └─▷ ClineTranslator fills parse hooks + + TraceCollector the driver: read → translate → hand to consumers + (MLflow logger, memory indexer); a per-consumer + ack set makes repeated passes idempotent +""" + +from __future__ import annotations + +import fcntl +import json +import logging +import os +import re +import sqlite3 +import threading +from abc import ABC, abstractmethod +from contextlib import contextmanager +from datetime import datetime +from pathlib import Path + +logger = logging.getLogger(__name__) + +_DEFAULT_SPAN_SECONDS = ( + 1.0 # fallback turn-span duration when no next timestamp bounds it +) +_ROLE_ASSISTANT = "assistant" +_ROLE_USER = "user" +_TYPE_QUEUE_OP = "queue-operation" + + +def _parse_ts(ts: object) -> float | None: + """ISO-8601 string (or epoch number) → epoch seconds; ``None`` on failure.""" + if not ts: + return None + if isinstance(ts, (int, float)): + return float(ts) + if isinstance(ts, str): + try: + return datetime.fromisoformat(ts.replace("Z", "+00:00")).timestamp() + except ValueError: + return None + return None + + +# --------------------------------------------------------------------------- +# The block / message / usage shapes (the one place they're constructed) +# --------------------------------------------------------------------------- + + +def _text_block(text: str | None) -> dict: + return {"type": "text", "text": text or ""} + + +def _tool_use_block(name, tool_input, tool_call_id=None) -> dict: + return { + "type": "tool_use", + "id": tool_call_id, + "name": name, + "input": tool_input or {}, + } + + +def _tool_result_block(output, tool_call_id=None) -> dict: + return {"type": "tool_result", "tool_use_id": tool_call_id, "content": output} + + +def _message(role: str, blocks: list[dict]) -> dict: + return {"role": role, "content": blocks} + + +def _usage(raw: dict | None) -> dict | None: + """Token counts from a transcript ``usage`` dict; ``None`` when absent.""" + if not raw: + return None + return { + "input_tokens": raw.get("input_tokens"), + "output_tokens": raw.get("output_tokens"), + "cache_read_input_tokens": raw.get("cache_read_input_tokens"), + "cache_write_input_tokens": raw.get("cache_creation_input_tokens"), + } + + +# =========================================================================== +# Trace — the canonical form (nested data + composition / query methods) +# =========================================================================== + + +class Trace: + """One finished agent session: id fields plus a list of span dicts. + + The span/message/block shapes are documented in the module docstring and + built *only* by the ``add_*`` methods here, so the schema has one home. + Consumers read ``spans`` (and ``to_exchanges()``) directly — the data is + already in the dict shape both the MLflow logger and memory want. + """ + + def __init__(self, trace_id: str, session_id: str, agent: str, project: str): + self.trace_id = trace_id + self.session_id = session_id + self.agent = agent + self.project = project + self.spans: list[dict] = [] + + @property + def started_at(self) -> float | None: + return self.spans[0]["start_time"] if self.spans else None + + @property + def ended_at(self) -> float | None: + return self.spans[-1]["end_time"] if self.spans else None + + # -- composition (the only place a span dict is constructed) ------------- + + def add_agent_root(self, span_id, name, *, start_time, prompt) -> dict: + span = { + "span_id": span_id, + "name": name, + "kind": "AGENT", + "parent_id": None, + "start_time": start_time, + "end_time": None, + "status": "ok", + "request": [], + "response": [], + "model": None, + "usage": None, + "attributes": {"prompt": prompt}, + } + self.spans.append(span) + return span + + def add_llm_span( + self, + span_id, + *, + parent_id, + start_time, + end_time, + request, + response, + model=None, + usage=None, + ) -> dict: + span = { + "span_id": span_id, + "name": "llm", + "kind": "LLM", + "parent_id": parent_id, + "start_time": start_time, + "end_time": end_time, + "status": "ok", + "request": request, + "response": response, + "model": model, + "usage": usage, + "attributes": {}, + } + self.spans.append(span) + return span + + def add_tool_span( + self, + span_id, + *, + parent_id, + start_time, + end_time, + name, + tool_input, + result, + tool_id="", + ) -> dict: + span = { + "span_id": span_id, + "name": f"tool_{name}", + "kind": "TOOL", + "parent_id": parent_id, + "start_time": start_time, + "end_time": end_time, + "status": "ok", + "request": [], + "response": [], + "model": None, + "usage": None, + "attributes": { + "tool_name": name, + "tool_id": tool_id, + "input": tool_input, + "result": result, + }, + } + self.spans.append(span) + return span + + def add_turn(self, *, root_id, root_name, prompt, root_ts, events, last_ts) -> None: + """Append one AGENT subtree from a turn's ordered ``events``. + + Each event is a tuple — ``("llm", ts, text, model, usage)`` or + ``("tool", ts, name, input, result)`` — in transcript order. This is the + shared builder the four template translators use: it derives each span's + duration from the next event's timestamp (1s fallback for the last), and + threads the request "window" (the prompt, then each tool call+result) + into the following ``llm`` span's ``request`` — which is what memory's + ``to_exchanges`` reads. + """ + root = self.add_agent_root( + root_id, root_name, start_time=root_ts, prompt=prompt + ) + pending = [_message(_ROLE_USER, [_text_block(prompt)])] + ts_list = [e[1] for e in events] + final_response: str | None = None + for i, ev in enumerate(events): + ts = ev[1] + nxt = next((t for t in ts_list[i + 1 :] if t is not None), last_ts) + dur = ( + (nxt - ts) + if (ts is not None and nxt is not None and nxt > ts) + else _DEFAULT_SPAN_SECONDS + ) + end = (ts + dur) if ts is not None else None + if ev[0] == "llm": + _, _, text, model, usage = ev + final_response = text + self.add_llm_span( + f"{root_id}-{i}", + parent_id=root_id, + start_time=ts, + end_time=end, + request=list(pending), + response=[_text_block(text)], + model=model, + usage=usage, + ) + pending = [] + else: # "tool" + _, _, name, tin, result = ev + self.add_tool_span( + f"{root_id}-{i}", + parent_id=root_id, + start_time=ts, + end_time=end, + name=name, + tool_input=tin, + result=result, + ) + tool_input = tin if isinstance(tin, dict) else {"raw": tin} + pending.append( + _message(_ROLE_ASSISTANT, [_tool_use_block(name, tool_input)]) + ) + pending.append(_message(_ROLE_USER, [_tool_result_block(result)])) + root["end_time"] = last_ts if last_ts is not None else root_ts + if final_response is not None: + root["attributes"]["response"] = final_response + + # -- query / projection ------------------------------------------------- + + def roots(self) -> list[dict]: + return [s for s in self.spans if s["parent_id"] is None] + + def children(self, root_id) -> list[dict]: + return [s for s in self.spans if s["parent_id"] == root_id] + + def subset(self, root_ids: set[str]) -> "Trace": + """A copy carrying only the given AGENT roots and their direct children.""" + keep = [ + s + for s in self.spans + if (s["parent_id"] is None and s["span_id"] in root_ids) + or s["parent_id"] in root_ids + ] + out = Trace(self.trace_id, self.session_id, self.agent, self.project) + out.spans = keep + return out + + def to_exchanges(self) -> list[dict]: + """Project the ``llm`` spans onto memory's conversational shape. + + One ``llm`` span → one ``{turn_id, timestamp, new_messages, response}`` + exchange; ``turn_id`` is the span id (groups a turn's chunks back + together downstream), ``request`` is already the windowed message list, + ``response`` the output blocks — so this is a straight projection. + """ + return [ + { + "turn_id": s["span_id"], + "timestamp": s["start_time"], + "new_messages": s["request"], + "response": s["response"], + } + for s in self.spans + if s["kind"] == "LLM" + ] + + +# =========================================================================== +# Readers — locate + read a platform's session record into raw records +# =========================================================================== + + +class Reader(ABC): + """Find this project's active session for an agent and read its records. + + Each reader resolves *the latest session for this project* and reads it. A + reader can also be pinned (:meth:`pin`) to a specific session — the startup + catch-up does this to re-read the *previous* session rather than whatever is + newest now. :meth:`active_source` returns an opaque, agent-shaped token for + the session being read — a transcript path (claude/codex), a DB session id + (goose/opencode), or a session-dir name (cline) — which the server records + in ``state.yaml`` so the next session's catch-up can pin it back. The token + round-trips through YAML, so its native type (str/int) is preserved. + """ + + agent: str + _pinned = None + + def pin(self, source) -> None: + """Pin to a specific session token (as returned by :meth:`active_source`).""" + self._pinned = source + + def active_source(self): + """Token identifying the session being read, or ``None`` if none. + + Subclasses override to resolve the latest session when unpinned; the + base returns the pinned token (``None`` when neither pinned nor + overridden). + """ + return self._pinned + + @abstractmethod + def read(self) -> list[dict]: + """The whole active session's raw records (re-read each pass; the + collector's ack set dedupes, so reads need not be incremental).""" + + +class JsonlReader(Reader): + """Read the whole active ``*.jsonl`` file, framing only complete lines. + + A trailing half-written line is dropped (picked up next pass). Subclasses + supply :meth:`active_file`; the framing is identical for claude and codex. + A pinned reader reads that exact file instead of the newest. + """ + + @abstractmethod + def active_file(self) -> Path | None: ... + + def _file(self) -> Path | None: + if self._pinned is not None: + p = Path(self._pinned) + return p if p.is_file() else None + return self.active_file() + + def active_source(self) -> str | None: + f = self._file() + return str(f) if f else None + + def read(self) -> list[dict]: + f = self._file() + if f is None: + return [] + with open(f, "rb") as fh: + chunk = fh.read() + last_nl = chunk.rfind(b"\n") + if last_nl == -1: + return [] + return [ + json.loads(line) + for line in chunk[: last_nl + 1].splitlines() + if line.strip() + ] + + +def _transcript_dir(project_dir: str | Path, projects_root: Path | None = None) -> Path: + """``~/.claude/projects/`` for a project directory. + + Claude derives the dir name by replacing every non-alphanumeric character of + the launch cwd with ``-``. The MCP server launches as ``cd && + claude``, so cwd == project_dir. + """ + root = projects_root or (Path.home() / ".claude" / "projects") + mangled = re.sub(r"[^a-zA-Z0-9]", "-", os.path.abspath(project_dir)) + return root / mangled + + +class ClaudeReader(JsonlReader): + """The most-recently-modified transcript in the project's ``~/.claude`` dir.""" + + agent = "claude" + + def __init__(self, project_dir, *, projects_root: Path | None = None): + self._dir = _transcript_dir(project_dir, projects_root) + + def active_file(self) -> Path | None: + if not self._dir.is_dir(): + return None + files = list(self._dir.glob("*.jsonl")) + return max(files, key=lambda p: p.stat().st_mtime) if files else None + + +class CodexReader(JsonlReader): + """The newest ``rollout-*.jsonl`` whose ``session_meta.cwd`` is this project. + + DSAGT always runs codex with ``CODEX_HOME=/.codex-data`` (its MCP + config lives there — see agents/codex.py), so rollouts land under + ``/.codex-data/sessions/YYYY/MM/DD/``, not the global + ``~/.codex/sessions/``. Each rollout opens with a ``session_meta`` record + carrying the launch ``cwd``; the filter guards against stray files. + """ + + agent = "codex" + + def __init__(self, project_dir, *, sessions_root: Path | None = None): + self._project_dir = os.path.abspath(project_dir) + self._root = ( + Path(sessions_root) + if sessions_root + else Path(self._project_dir) / ".codex-data" / "sessions" + ) + + def _rollout_cwd(self, path: Path) -> str | None: + with open(path, encoding="utf-8") as fh: + first = fh.readline() + if not first.strip(): + return None + rec = json.loads(first) + if rec.get("type") != "session_meta": + return None + return (rec.get("payload") or {}).get("cwd") + + def active_file(self) -> Path | None: + if not self._root.is_dir(): + return None + files = sorted( + self._root.glob("**/rollout-*.jsonl"), + key=lambda p: p.stat().st_mtime, + reverse=True, + ) + for f in files: + if self._rollout_cwd(f) == self._project_dir: + return f + return None + + +_GOOSE_DB = Path.home() / ".local" / "share" / "goose" / "sessions" / "sessions.db" + + +class GooseReader(Reader): + """Read the project's active Goose session from its SQLite store (read-only).""" + + agent = "goose" + + def __init__(self, project_dir, *, db_path: Path | None = None): + self._project_dir = os.path.abspath(project_dir) + self._db = Path(db_path) if db_path else _GOOSE_DB + + def _latest_session(self, con): + row = con.execute( + "SELECT id FROM sessions WHERE working_dir = ? " + "ORDER BY updated_at DESC LIMIT 1", + (self._project_dir,), + ).fetchone() + return row[0] if row else None + + def active_source(self): + if self._pinned is not None: + return self._pinned + if not self._db.exists(): + return None + con = sqlite3.connect(f"file:{self._db}?mode=ro", uri=True) + try: + return self._latest_session(con) + finally: + con.close() + + def read(self) -> list[dict]: + if not self._db.exists(): + return [] + con = sqlite3.connect(f"file:{self._db}?mode=ro", uri=True) + try: + sid = ( + self._pinned if self._pinned is not None else self._latest_session(con) + ) + if sid is None: + return [] + return [ + {"role": role, "content": _loads_list(cj), "ts": ts} + for role, cj, ts in con.execute( + "SELECT role, content_json, created_timestamp FROM messages " + "WHERE session_id = ? ORDER BY id", + (sid,), + ) + ] + finally: + con.close() + + +_OPENCODE_DB = Path.home() / ".local" / "share" / "opencode" / "opencode.db" + + +class OpenCodeReader(Reader): + """Read + flatten the project's active opencode session (sqlite, read-only). + + opencode splits a session into ``message`` rows (role/model) and ``part`` + rows (the text/tool payloads); this joins them into one time-ordered list of + parts, each tagged with its message's role + model and a timestamp in + seconds (opencode stores milliseconds). + """ + + agent = "opencode" + + def __init__(self, project_dir, *, db_path: Path | None = None): + self._project_dir = os.path.abspath(project_dir) + self._db = Path(db_path) if db_path else _OPENCODE_DB + + def _latest_session(self, con): + row = con.execute( + "SELECT id FROM session WHERE directory = ? " + "ORDER BY time_updated DESC LIMIT 1", + (self._project_dir,), + ).fetchone() + return row[0] if row else None + + def active_source(self): + if self._pinned is not None: + return self._pinned + if not self._db.exists(): + return None + con = sqlite3.connect(f"file:{self._db}?mode=ro", uri=True) + try: + return self._latest_session(con) + finally: + con.close() + + def read(self) -> list[dict]: + if not self._db.exists(): + return [] + con = sqlite3.connect(f"file:{self._db}?mode=ro", uri=True) + try: + sid = ( + self._pinned if self._pinned is not None else self._latest_session(con) + ) + if sid is None: + return [] + messages = { + mid: _loads_dict(d) + for mid, d in con.execute( + "SELECT id, data FROM message WHERE session_id = ?", (sid,) + ) + } + parts = [] + for mid, t_created, data in con.execute( + "SELECT message_id, time_created, data FROM part " + "WHERE session_id = ? ORDER BY time_created", + (sid,), + ): + m = messages.get(mid, {}) + parts.append( + { + "role": m.get("role"), + "model": (m.get("model") or {}).get("modelID"), + "ts": t_created / 1000.0, + "data": _loads_dict(data), + } + ) + return parts + finally: + con.close() + + +_CLINE_SESSIONS_ROOT = Path.home() / ".cline" / "data" / "sessions" + + +class ClineReader(Reader): + """Locate the project's active Cline CLI session by its metadata ``cwd``. + + Each session is ``~/.cline/data/sessions//`` with ``.json`` (metadata + incl. ``cwd`` + ``model``) and ``.messages.json`` (the whole message list). + """ + + agent = "cline" + + def __init__(self, project_dir, *, sessions_root: Path | None = None): + self._project_dir = os.path.abspath(project_dir) + self._root = Path(sessions_root) if sessions_root else _CLINE_SESSIONS_ROOT + + def _active_dir(self) -> Path | None: + if not self._root.is_dir(): + return None + for d in sorted(self._root.iterdir(), key=lambda p: p.name, reverse=True): + meta = d / f"{d.name}.json" + if not meta.exists(): + continue + try: + cwd = json.loads(meta.read_text()).get("cwd", "") + except (json.JSONDecodeError, OSError): + continue + if os.path.abspath(cwd) == self._project_dir: + return d + return None + + def _dir(self) -> Path | None: + if self._pinned is not None: + d = self._root / str(self._pinned) + return d if d.is_dir() else None + return self._active_dir() + + def active_source(self) -> str | None: + d = self._dir() + return d.name if d else None + + def read(self) -> list[dict]: + d = self._dir() + if d is None: + return [] + msgs_path = d / f"{d.name}.messages.json" + if not msgs_path.exists(): + return [] + try: + data = json.loads(msgs_path.read_text()) + except (json.JSONDecodeError, OSError): + return [] + model = self._session_model(d) + messages = data.get("messages", []) if isinstance(data, dict) else data + for m in messages: + m["model"] = model + return messages + + def _session_model(self, session_dir: Path) -> str | None: + try: + meta = json.loads((session_dir / f"{session_dir.name}.json").read_text()) + return meta.get("model") + except (json.JSONDecodeError, OSError): + return None + + +def _loads_list(raw: str) -> list: + try: + out = json.loads(raw) + return out if isinstance(out, list) else [] + except (json.JSONDecodeError, TypeError): + return [] + + +def _loads_dict(raw: str) -> dict: + try: + out = json.loads(raw) + return out if isinstance(out, dict) else {} + except (json.JSONDecodeError, TypeError): + return {} + + +# =========================================================================== +# Translators — raw records → Trace (pure, no I/O) +# =========================================================================== + + +class Translator(ABC): + """Map one platform's records to a :class:`Trace`. + + The default ``translate`` is the shared template every platform but Claude + uses: build a tool-result map, find the turn-start (prompt) indices, and for + each turn lower its records into ordered ``("llm"|"tool", …)`` events that + ``Trace.add_turn`` turns into a span subtree. A subclass supplies the small + parse hooks (``_is_prompt`` / ``_prompt_text`` / ``_ts`` / ``_events`` and, + where needed, ``_tool_results`` / ``_normalize`` / ``_prompt_indices``). + Claude overrides ``translate`` outright — its grammar exceeds this shape. + """ + + agent: str + root_name: str + + def translate(self, records, *, trace_id, session_id, project) -> Trace | None: + records = self._normalize(records) + results = self._tool_results(records) + prompts = self._prompt_indices(records) + if not prompts: + return None + trace = Trace(trace_id, session_id, self.agent, project) + bounds = prompts + [len(records)] + for k, start in enumerate(prompts): + self._build_turn(trace, records, start, bounds[k + 1], results) + return trace if trace.spans else None + + def _build_turn(self, trace, records, start, end, results) -> None: + root_ts = self._ts(records[start]) + last_ts = root_ts + events = [] + for i in range(start + 1, end): + rec = records[i] + ts = self._ts(rec) + if ts is not None: + last_ts = ts + events.extend(self._events(rec, ts, results)) + trace.add_turn( + root_id=f"turn-{start}", + root_name=self.root_name, + prompt=self._prompt_text(records[start]), + root_ts=root_ts, + events=events, + last_ts=last_ts, + ) + + # -- hooks (defaults; concrete translators fill what they need) ---------- + + def _normalize(self, records): + return records + + def _prompt_indices(self, records) -> list[int]: + return [i for i, r in enumerate(records) if self._is_prompt(r)] + + def _tool_results(self, records) -> dict: + return {} + + def _is_prompt(self, rec) -> bool: + raise NotImplementedError + + def _prompt_text(self, rec) -> str: + raise NotImplementedError + + def _ts(self, rec) -> float | None: + raise NotImplementedError + + def _events(self, rec, ts, results) -> list: + raise NotImplementedError + + +class GooseTranslator(Translator): + """Goose message rows → Trace (content blocks: text / toolRequest / toolResponse).""" + + agent = "goose" + root_name = "goose_conversation" + + @staticmethod + def _blocks(msg) -> list[dict]: + c = msg.get("content") + return c if isinstance(c, list) else [] + + def _is_prompt(self, rec) -> bool: + return rec.get("role") == _ROLE_USER and any( + b.get("type") == "text" for b in self._blocks(rec) + ) + + def _prompt_text(self, rec) -> str: + return "".join( + b.get("text", "") for b in self._blocks(rec) if b.get("type") == "text" + ) + + def _ts(self, rec) -> float | None: + return _parse_ts(rec.get("ts")) + + def _tool_results(self, records) -> dict: + results: dict = {} + for msg in records: + for b in self._blocks(msg): + if b.get("type") == "toolResponse" and (tid := b.get("id")): + results[tid] = self._result_text(b) + return results + + @staticmethod + def _result_text(block) -> object: + value = (block.get("toolResult") or {}).get("value") or {} + content = value.get("content") + if isinstance(content, list): + return "".join( + c.get("text", "") for c in content if c.get("type") == "text" + ) + return content if content is not None else "" + + def _events(self, rec, ts, results) -> list: + if rec.get("role") != _ROLE_ASSISTANT: + return [] + out = [] + for b in self._blocks(rec): + if b.get("type") == "text" and b.get("text", "").strip(): + out.append(("llm", ts, b["text"], None, None)) + elif b.get("type") == "toolRequest": + call = (b.get("toolCall") or {}).get("value") or {} + out.append( + ( + "tool", + ts, + call.get("name", "unknown"), + call.get("arguments", {}), + results.get(b.get("id"), ""), + ) + ) + return out + + +class OpenCodeTranslator(Translator): + """opencode flattened parts → Trace (call + result live in one tool part).""" + + agent = "opencode" + root_name = "opencode_conversation" + + def _is_prompt(self, rec) -> bool: + d = rec.get("data") or {} + return ( + rec.get("role") == _ROLE_USER + and d.get("type") == "text" + and bool(d.get("text", "").strip()) + ) + + def _prompt_text(self, rec) -> str: + return (rec.get("data") or {}).get("text", "") + + def _ts(self, rec) -> float | None: + return rec.get("ts") + + def _events(self, rec, ts, results) -> list: + d = rec.get("data") or {} + if ( + d.get("type") == "text" + and rec.get("role") == _ROLE_ASSISTANT + and d.get("text", "").strip() + ): + return [("llm", ts, d["text"], rec.get("model"), None)] + if d.get("type") == "tool": + state = d.get("state") or {} + return [ + ( + "tool", + ts, + d.get("tool", "unknown"), + state.get("input", {}), + state.get("output", ""), + ) + ] + return [] + + +_CLINE_USER_INPUT_RE = re.compile(r"]*>(.*?)", re.DOTALL) + + +class ClineTranslator(Translator): + """Cline message records → Trace (Anthropic block list; ms timestamps).""" + + agent = "cline" + root_name = "cline_conversation" + + @staticmethod + def _content(msg) -> list[dict]: + c = msg.get("content") + return c if isinstance(c, list) else [] + + def _join_text(self, msg) -> str: + return "".join( + b.get("text", "") for b in self._content(msg) if b.get("type") == "text" + ) + + def _is_prompt(self, rec) -> bool: + return rec.get("role") == _ROLE_USER and any( + b.get("type") == "text" and b.get("text", "").strip() + for b in self._content(rec) + ) + + def _prompt_text(self, rec) -> str: + text = self._join_text(rec) + found = _CLINE_USER_INPUT_RE.findall(text) + return "".join(found).strip() if found else text.strip() + + def _ts(self, rec) -> float | None: + ts = rec.get("ts") + return ts / 1000.0 if isinstance(ts, (int, float)) else None + + def _tool_results(self, records) -> dict: + results: dict = {} + for msg in records: + for b in self._content(msg): + if b.get("type") == "tool_result" and (tid := b.get("tool_use_id")): + results[tid] = self._result_text(b) + return results + + @staticmethod + def _result_text(block) -> object: + content = block.get("content") + if isinstance(content, list): + return "".join( + c.get("text", "") + for c in content + if isinstance(c, dict) and c.get("type") == "text" + ) + return content if content is not None else "" + + def _events(self, rec, ts, results) -> list: + if rec.get("role") != _ROLE_ASSISTANT: + return [] + out = [] + text = self._join_text(rec) + if text.strip(): + out.append(("llm", ts, text, rec.get("model"), None)) + for b in self._content(rec): + if b.get("type") == "tool_use": + out.append( + ( + "tool", + ts, + b.get("name", "unknown"), + b.get("input", {}), + results.get(b.get("id"), ""), + ) + ) + return out + + +class CodexTranslator(Translator): + """Codex rollout records → Trace (OpenAI Responses format). + + Normalizes the rollout into ``{ts, p}`` conversation items, then fits the + template — except the prompt is the *last* user message in a consecutive run + (Codex injects an AGENTS.md context message earlier), so ``_prompt_indices`` + is overridden with that lookahead. + """ + + agent = "codex" + root_name = "codex_conversation" + + def _normalize(self, records) -> list[dict]: + return [ + {"ts": _parse_ts(r.get("timestamp")), "p": r["payload"]} + for r in records + if r.get("type") == "response_item" and isinstance(r.get("payload"), dict) + ] + + def _ts(self, rec) -> float | None: + return rec["ts"] + + def _tool_results(self, records) -> dict: + results: dict = {} + for rec in records: + p = rec["p"] + if p.get("type") in ("function_call_output", "custom_tool_call_output"): + if cid := p.get("call_id"): + results[cid] = p.get("output", "") + return results + + def _prompt_indices(self, records) -> list[int]: + msg_idxs = [i for i, r in enumerate(records) if r["p"].get("type") == "message"] + prompts = [] + for pos, i in enumerate(msg_idxs): + if records[i]["p"].get("role") != _ROLE_USER: + continue + nxt = msg_idxs[pos + 1] if pos + 1 < len(msg_idxs) else None + if nxt is None or records[nxt]["p"].get("role") == _ROLE_ASSISTANT: + prompts.append(i) + return prompts + + def _prompt_text(self, rec) -> str: + return self._msg_text(rec["p"]) + + @staticmethod + def _msg_text(item) -> str: + return "".join( + b.get("text", "") + for b in (item.get("content") or []) + if isinstance(b, dict) and b.get("type") in ("input_text", "output_text") + ) + + def _events(self, rec, ts, results) -> list: + p = rec["p"] + ptype = p.get("type") + if ptype == "message" and p.get("role") == _ROLE_ASSISTANT: + text = self._msg_text(p) + return [("llm", ts, text, None, None)] if text.strip() else [] + if ptype in ("function_call", "custom_tool_call"): + name, tool_input = self._tool_call(p) + return [ + ("tool", ts, name, tool_input, results.get(p.get("call_id", ""), "")) + ] + return [] + + @staticmethod + def _tool_call(item) -> tuple[str, object]: + name = item.get("name", "unknown") + if item.get("type") == "function_call": + raw = item.get("arguments", "") + try: + return name, json.loads(raw) + except (json.JSONDecodeError, TypeError): + return name, raw + return name, item.get("input", "") # custom_tool_call: raw input string + + +class ClaudeTranslator(Translator): + """Claude transcript → Trace — bespoke; overrides the template. + + Claude emits separate entries per thinking / text / tool_use, folds queued + "steer" messages into the request window, splits a turn's duration across + multiple tool calls, and carries token usage — richer than the event + template, so it builds spans directly via ``Trace``'s ``add_*`` methods. + """ + + agent = "claude" + + def translate(self, records, *, trace_id, session_id, project) -> Trace | None: + starts = self._prompt_indices(records) + if not starts: + return None + trace = Trace(trace_id, session_id, self.agent, project) + bounds = starts + [len(records)] + for k, user_idx in enumerate(starts): + self._build_turn(trace, records, user_idx, bounds[k + 1]) + return trace if trace.spans else None + + # _build_turn here takes no results arg (Claude looks results up per turn), + # so it intentionally shadows the template's signature. + def _build_turn(self, trace, records, user_idx, end_idx) -> None: # type: ignore[override] + user_rec = records[user_idx] + user_msg = (user_rec.get("message") or {}).get("content", "") + prompt = ( + user_msg if isinstance(user_msg, str) else self._text_and_tools(user_msg)[0] + ) + root_id = user_rec.get("uuid") or f"turn-{user_idx}" + root_ts = _parse_ts(user_rec.get("timestamp")) + root = trace.add_agent_root( + root_id, "claude_code_conversation", start_time=root_ts, prompt=prompt + ) + + counter = 0 + final_response: str | None = None + last_ts = root_ts + for i in range(user_idx + 1, end_idx): + rec = records[i] + if (t := _parse_ts(rec.get("timestamp"))) is not None: + last_ts = t + if rec.get("type") != _ROLE_ASSISTANT: + continue + msg = rec.get("message") or {} + ts = _parse_ts(rec.get("timestamp")) + text, tools = self._text_and_tools(msg.get("content")) + nxt = self._next_timestamp(records, i, stop=end_idx) + duration = ( + (nxt - ts) + if (ts is not None and nxt is not None and nxt > ts) + else _DEFAULT_SPAN_SECONDS + ) + + if text.strip() and not tools: + final_response = text + trace.add_llm_span( + f"{root_id}-{counter}", + parent_id=root_id, + start_time=ts, + end_time=(ts + duration) if ts is not None else None, + request=self._window_messages(records, i), + response=[_text_block(text)], + model=msg.get("model"), + usage=_usage(msg.get("usage")), + ) + counter += 1 + + if tools: + results = self._tool_results_after(records, i) + tool_duration = duration / len(tools) + for idx_t, tu in enumerate(tools): + tid = tu.get("id", "") + t_start = (ts + idx_t * tool_duration) if ts is not None else None + trace.add_tool_span( + f"{root_id}-{counter}", + parent_id=root_id, + start_time=t_start, + end_time=( + (t_start + tool_duration) if t_start is not None else None + ), + name=tu.get("name", "unknown"), + tool_input=tu.get("input", {}), + result=results.get(tid, ""), + tool_id=tid, + ) + counter += 1 + + root["end_time"] = last_ts + if final_response is not None: + root["attributes"]["response"] = final_response + + # -- Claude grammar (ported from mlflow claude_code; see NOTICE) --------- + + def _prompt_indices(self, records) -> list[int]: + return [i for i in range(len(records)) if self._is_user_prompt(records, i)] + + @staticmethod + def _is_user_prompt(records, i) -> bool: + rec = records[i] + if rec.get("type") != _ROLE_USER or rec.get("toolUseResult"): + return False + prev = records[i - 1] if i > 0 else None + prev_tur = prev.get("toolUseResult") if isinstance(prev, dict) else None + if isinstance(prev_tur, dict) and prev_tur.get("commandName"): + return False + content = (rec.get("message") or {}).get("content") + if isinstance(content, list) and content and isinstance(content[0], dict): + if content[0].get("type") == "tool_result": + return False + if isinstance(content, str): + if "" in content or not content.strip(): + return False + return bool(content) + + @staticmethod + def _text_and_tools(content) -> tuple[str, list[dict]]: + text, tools = "", [] + if isinstance(content, list): + for b in content: + if not isinstance(b, dict): + continue + if b.get("type") == "text": + text += b.get("text", "") + elif b.get("type") == "tool_use": + tools.append(b) + elif isinstance(content, str): + text = content + return text, tools + + @staticmethod + def _next_timestamp(records, idx, stop=None) -> float | None: + end = len(records) if stop is None else stop + for j in range(idx + 1, end): + if (t := _parse_ts(records[j].get("timestamp"))) is not None: + return t + return None + + def _block_from_raw(self, raw) -> dict | None: + btype = raw.get("type") + if btype == "text": + return _text_block(raw.get("text", "")) + if btype == "tool_use": + return _tool_use_block( + raw.get("name"), raw.get("input") or {}, raw.get("id") + ) + if btype == "tool_result": + return _tool_result_block(raw.get("content"), raw.get("tool_use_id")) + return None # thinking / unknown — no conversational payload + + def _message_from_record(self, record) -> dict | None: + msg = record.get("message") or {} + role = msg.get("role") + content = msg.get("content") + if not role or not content: + return None + if isinstance(content, str): + return _message(role, [_text_block(content)]) + blocks = [b for raw in content if (b := self._block_from_raw(raw))] + return _message(role, blocks) if blocks else None + + def _window_messages(self, records, current_idx) -> list[dict]: + """The user/tool entries since the previous text-bearing assistant turn.""" + out: list[dict] = [] + for i in range(current_idx - 1, -1, -1): + rec = records[i] + if rec.get("type") == _ROLE_ASSISTANT: + text, _ = self._text_and_tools( + (rec.get("message") or {}).get("content") + ) + if text.strip(): + break + if ( + rec.get("type") == _TYPE_QUEUE_OP + and rec.get("operation") == "enqueue" + and (steer := rec.get("content")) + ): + txt = steer if isinstance(steer, str) else str(steer) + out.append(_message(_ROLE_USER, [_text_block(txt)])) + continue + if (m := self._message_from_record(rec)) is not None: + out.append(m) + out.reverse() + return out + + @staticmethod + def _tool_results_after(records, start_idx) -> dict: + results: dict = {} + for i in range(start_idx + 1, len(records)): + rec = records[i] + if rec.get("type") != _ROLE_USER: + continue + content = (rec.get("message") or {}).get("content") + if isinstance(content, list): + for b in content: + if isinstance(b, dict) and b.get("type") == "tool_result": + if tid := b.get("tool_use_id"): + results[tid] = b.get("content", "") + return results + + +# =========================================================================== +# TraceCollector — the driver (read → translate → hand to consumers) +# =========================================================================== + +# An agent appears here once both its reader and translator exist; the collector +# runs for any agent in the table and is simply absent for the rest. +_PIPELINES = { + "claude": lambda pd, pr: (ClaudeReader(pd, projects_root=pr), ClaudeTranslator()), + "codex": lambda pd, pr: (CodexReader(pd), CodexTranslator()), + "goose": lambda pd, pr: (GooseReader(pd), GooseTranslator()), + "opencode": lambda pd, pr: (OpenCodeReader(pd), OpenCodeTranslator()), + "cline": lambda pd, pr: (ClineReader(pd), ClineTranslator()), +} + + +def make_trace_collector( + agent, + project_dir, + project, + session_id, + tracking_uri, + *, + projects_root: Path | None = None, + extra_consumers: list | None = None, + source=None, +) -> "TraceCollector | None": + """Build the collector for ``agent``, or ``None`` if no pipeline is registered. + + The MLflow logger is always the first consumer (observability is universal); + ``extra_consumers`` (e.g. a :class:`~dsagt.memory.MemoryExtractor` when + episodic memory is enabled) are appended, each acking independently. + + ``source`` pins the reader to a specific session (the startup catch-up passes + the *previous* session's recorded :meth:`Reader.active_source` token), so it + re-reads that exact session instead of whatever is newest now — uniformly + across all agents (transcript path, DB session id, or session-dir name). + """ + builder = _PIPELINES.get(agent) + if builder is None: + return None + reader, translator = builder(project_dir, projects_root) + if source is not None: + reader.pin(source) + # Imported here (not at module top) so traces stays a lean leaf — the MLflow + # logger drags in mlflow, the heaviest thing in the pipeline. + from dsagt.observability import MLflowSink + + consumers = [MLflowSink(tracking_uri, project), *(extra_consumers or [])] + return TraceCollector( + reader, + translator, + project=project, + session_id=session_id, + project_dir=project_dir, + consumers=consumers, + ) + + +class TraceCollector: + """Periodically read the session, translate it, and hand it to consumers. + + A *consumer* is anything with a ``name`` and a ``write(trace)`` — the MLflow + logger and the memory indexer both qualify (no shared base needed). Each + consumer keeps its own ack set (``.dsagt/trace_acks_.json``), keyed by + session-qualified turn id (``:``) so the per-transcript + ``turn-N`` indices can't collide across sessions in the shared file. A + re-pass or an N+1 catch-up can only waste work, never double-log or lose a + turn, and a failing consumer holds back only its own mark. + + Completeness watermark: a periodic pass emits only *completed* turns (all but + the still-open last one); the deferred final turn flushes when a later prompt + bounds it or at end-of-session (``include_last=True``). An OS file lock + serializes overlapping passes against the shared ack files. + """ + + def __init__( + self, reader, translator, *, project, session_id, project_dir, consumers + ): + self._reader = reader + self._translator = translator + self._project = project + self._session_id = session_id + self._project_dir = Path(project_dir) + self._consumers = list(consumers) + self._dsagt_dir = self._project_dir / ".dsagt" + self._lock = threading.Lock() + + def active_source(self): + """The reader's session token (see :meth:`Reader.active_source`), or + ``None``. Recorded in ``state.yaml`` so the next session's catch-up can + pin this exact session — uniform across all agents. + """ + try: + return self._reader.active_source() + except Exception: # noqa: BLE001 — best-effort; never break the heartbeat + return None + + def _acks_path(self, name: str) -> Path: + return self._dsagt_dir / f"trace_acks_{name}.json" + + def _load_acks(self, name: str) -> set[str]: + try: + return set(json.loads(self._acks_path(name).read_text())) + except FileNotFoundError: + return set() + + def _save_acks(self, name: str, acks: set[str]) -> None: + self._dsagt_dir.mkdir(parents=True, exist_ok=True) + self._acks_path(name).write_text(json.dumps(sorted(acks))) + + @contextmanager + def _lock_file(self): + self._dsagt_dir.mkdir(parents=True, exist_ok=True) + with open(self._dsagt_dir / "trace_acks.lock", "w") as lf: + fcntl.flock(lf, fcntl.LOCK_EX) + try: + yield + finally: + fcntl.flock(lf, fcntl.LOCK_UN) + + def collect(self, *, include_last: bool = False) -> int: + """Read → translate → hand completed turns to each consumer. + + Returns the number of turns newly delivered to at least one consumer. + ``include_last=True`` (end-of-session flush / per-turn hook) also emits + the otherwise-deferred final turn. Blocking — call via + ``asyncio.to_thread`` from the event loop. + """ + with self._lock, self._lock_file(): + records = self._reader.read() + if not records: + return 0 + trace = self._translator.translate( + records, + trace_id=self._session_id, + session_id=self._session_id, + project=self._project, + ) + if trace is None: + return 0 + + roots = trace.roots() + candidates = roots if include_last else roots[:-1] + # Ack keys are session-qualified. span_id is a per-transcript record + # index ("turn-N"), so the same ids recur in every session's + # transcript; a bare span_id would collide across sessions in the + # shared, never-reset ack file and suppress every turn after the + # first session. Qualifying by session id matches the key MLflowSink + # already uses for its own idempotency (``{trace_id}:{span_id}``). + key_by_span = { + r["span_id"]: f"{self._session_id}:{r['span_id']}" for r in candidates + } + if not key_by_span: + return 0 + + emitted: set[str] = set() + for consumer in self._consumers: + acks = self._load_acks(consumer.name) + emit_ids = {s for s, key in key_by_span.items() if key not in acks} + if not emit_ids: + continue + try: + consumer.write(trace.subset(emit_ids)) + except Exception as e: # noqa: BLE001 — per-consumer isolation + logger.warning("Trace consumer %r failed: %s", consumer.name, e) + continue + self._save_acks( + consumer.name, acks | {key_by_span[s] for s in emit_ids} + ) + emitted |= emit_ids + return len(emitted) diff --git a/src/session.py b/src/session.py deleted file mode 100755 index bfb9d64..0000000 --- a/src/session.py +++ /dev/null @@ -1,458 +0,0 @@ -""" -DSAGT session management. - -Handles project initialization, agent-specific config generation, -and service lifecycle (proxy + MLflow). - -Each agent platform gets its configs generated from the single -dsagt_config.yaml — MCP server entries, agent instructions, env vars, -and proxy routing are all derived mechanically. -""" - -import json -import logging -import os -import signal -import subprocess -import sys -from pathlib import Path - -import yaml - -from dsagt.config import ( - VALID_AGENTS, - default_config_content, - load_config, - project_dir_for, -) - -logger = logging.getLogger(__name__) - -# Where bundled agent instruction templates live (relative to dsagt package) -_AGENTS_DIR = Path(__file__).parent.parent.parent / "agents" - - -# --------------------------------------------------------------------------- -# Project initialization -# --------------------------------------------------------------------------- - -def init_project(project_name: str, agent: str, runtime_base: str | Path = "runtime") -> Path: - """Create a new project directory with default config and subdirectories. - - Returns the project directory path. - """ - if agent not in VALID_AGENTS: - raise ValueError(f"agent must be one of {VALID_AGENTS}, got '{agent}'") - - project_dir = project_dir_for(project_name, runtime_base) - - if (project_dir / "dsagt_config.yaml").exists(): - raise FileExistsError(f"Project already exists: {project_dir}") - - project_dir.mkdir(parents=True, exist_ok=True) - for subdir in ("trace_archive", "mlflow", "skills", "kb_index"): - (project_dir / subdir).mkdir(exist_ok=True) - - config_content = default_config_content(project_name, agent) - (project_dir / "dsagt_config.yaml").write_text(config_content) - - return project_dir - - -# --------------------------------------------------------------------------- -# Agent config generation -# --------------------------------------------------------------------------- - -def generate_agent_configs(config: dict, working_dir: str | Path) -> list[str]: - """Generate agent-platform-specific config files in the working directory. - - Reads the dsagt_config.yaml values and writes the files each agent - platform expects: MCP server configs, agent instructions, env setup. - - Returns a list of descriptions of what was written. - """ - agent = config["agent"] - working_dir = Path(working_dir) - project_dir = Path(config["project_dir"]) - proxy_port = config["proxy"]["port"] - - generators = { - "claude-code": _generate_claude_code, - "goose": _generate_goose, - "roo": _generate_roo, - "cline": _generate_cline, - } - - return generators[agent](config, working_dir, project_dir, proxy_port) - - -def _mcp_server_args(server: str, project_dir: Path) -> list[str]: - """Build the args list for an MCP server entry.""" - base = ["run", f"dsagt-{server}-server"] - if server == "registry": - base += ["--runtime-dir", str(project_dir)] - elif server == "knowledge": - base += [ - "--base-index-dir", str(project_dir / "kb_index"), - "--runtime-dir", str(project_dir), - ] - return base - - -def _mcp_env_block(config: dict) -> dict: - """Build the env block for MCP server entries.""" - env = {} - embedding_key = config.get("embedding", {}).get("api_key", "") - if embedding_key and not embedding_key.startswith("${"): - env["LLM_API_KEY"] = embedding_key - return env - - -def _generate_claude_code(config, working_dir, project_dir, proxy_port) -> list[str]: - actions = [] - env_block = _mcp_env_block(config) - - mcp_config = {"mcpServers": {}} - for server in ("registry", "knowledge"): - entry = { - "command": "uv", - "args": _mcp_server_args(server, project_dir), - } - if env_block: - entry["env"] = env_block - mcp_config["mcpServers"][f"dsagt-{server}"] = entry - - mcp_path = working_dir / ".mcp.json" - mcp_path.write_text(json.dumps(mcp_config, indent=2) + "\n") - actions.append(f"Wrote {mcp_path}") - - # Copy agent instructions - instructions = _load_instructions("claude-code", "dsagt_instructions.md") - if instructions: - claude_md = working_dir / "CLAUDE.md" - if claude_md.exists(): - existing = claude_md.read_text() - if "DSAGT Pipeline Builder" not in existing: - claude_md.write_text(existing + "\n\n" + instructions) - actions.append(f"Appended DSAGT instructions to {claude_md}") - else: - claude_md.write_text(instructions) - actions.append(f"Wrote {claude_md}") - - # Write env helper - env_path = working_dir / ".dsagt_env" - _write_env_file(env_path, { - "ANTHROPIC_BASE_URL": f"http://localhost:{proxy_port}", - "DSAGT_PROJECT": config["project"], - "DSAGT_PROJECT_DIR": str(project_dir), - }) - actions.append(f"Wrote {env_path} (source this or export the variables)") - - return actions - - -def _generate_goose(config, working_dir, project_dir, proxy_port) -> list[str]: - actions = [] - model = config["llm"]["model"] - - goose_config = { - "GOOSE_PROVIDER": "openai", - "GOOSE_MODEL": model, - "extensions": {}, - } - - for server in ("registry", "knowledge"): - args = _mcp_server_args(server, project_dir) - goose_config["extensions"][server] = { - "enabled": True, - "name": server, - "type": "stdio", - "cmd": "uv " + " ".join(args), - "timeout": 300, - } - - goose_path = working_dir / "goose.yaml" - goose_path.write_text(yaml.dump(goose_config, default_flow_style=False, sort_keys=False)) - actions.append(f"Wrote {goose_path}") - - # Copy agent instructions - instructions = _load_instructions("goose", ".goosehints") - if instructions: - hints_path = working_dir / ".goosehints" - hints_path.write_text(instructions) - actions.append(f"Wrote {hints_path}") - - env_path = working_dir / ".dsagt_env" - _write_env_file(env_path, { - "OPENAI_HOST": f"http://localhost:{proxy_port}", - "DSAGT_PROJECT": config["project"], - "DSAGT_PROJECT_DIR": str(project_dir), - }) - actions.append(f"Wrote {env_path}") - - return actions - - -def _generate_roo(config, working_dir, project_dir, proxy_port) -> list[str]: - actions = [] - env_block = _mcp_env_block(config) - - mcp_config = {"mcpServers": {}} - for server in ("registry", "knowledge"): - entry = { - "command": "uv", - "args": _mcp_server_args(server, project_dir), - "disabled": False, - } - if env_block: - entry["env"] = env_block - mcp_config["mcpServers"][f"dsagt-{server}"] = entry - - roo_dir = working_dir / ".roo" - roo_dir.mkdir(exist_ok=True) - mcp_path = roo_dir / "mcp.json" - mcp_path.write_text(json.dumps(mcp_config, indent=2) + "\n") - actions.append(f"Wrote {mcp_path}") - - # Copy roomodes - roomodes = _load_instructions("roo", "roomodes") - if roomodes: - roomodes_path = working_dir / ".roomodes" - roomodes_path.write_text(roomodes) - actions.append(f"Wrote {roomodes_path}") - - env_path = working_dir / ".dsagt_env" - _write_env_file(env_path, { - "DSAGT_PROJECT": config["project"], - "DSAGT_PROJECT_DIR": str(project_dir), - }) - actions.append(f"Wrote {env_path}") - - return actions - - -def _generate_cline(config, working_dir, project_dir, proxy_port) -> list[str]: - actions = [] - env_block = _mcp_env_block(config) - - mcp_config = {"mcpServers": {}} - for server in ("registry", "knowledge"): - entry = { - "command": "uv", - "args": _mcp_server_args(server, project_dir), - "disabled": False, - "alwaysAllow": [], - } - if env_block: - entry["env"] = env_block - mcp_config["mcpServers"][f"dsagt-{server}"] = entry - - # Write to project dir (user merges into global cline settings) - mcp_path = working_dir / "cline_mcp.json" - mcp_path.write_text(json.dumps(mcp_config, indent=2) + "\n") - actions.append(f"Wrote {mcp_path} (merge into Cline MCP settings)") - - # Copy agent instructions - instructions = _load_instructions("cline", "dsagt_instructions.md") - if instructions: - rules_dir = working_dir / ".clinerules" - rules_dir.mkdir(exist_ok=True) - instr_path = rules_dir / "dsagt_instructions.md" - instr_path.write_text(instructions) - actions.append(f"Wrote {instr_path}") - - env_path = working_dir / ".dsagt_env" - _write_env_file(env_path, { - "DSAGT_PROJECT": config["project"], - "DSAGT_PROJECT_DIR": str(project_dir), - }) - actions.append(f"Wrote {env_path}") - - return actions - - -def _load_instructions(agent_dir: str, filename: str) -> str | None: - """Load an instruction template from the agents directory.""" - path = _AGENTS_DIR / agent_dir / filename - if path.exists(): - return path.read_text() - logger.warning("Instruction template not found: %s", path) - return None - - -def _write_env_file(path: Path, env_vars: dict) -> None: - """Write a sourceable env file.""" - lines = [f'export {k}="{v}"' for k, v in env_vars.items()] - path.write_text("\n".join(lines) + "\n") - - -# --------------------------------------------------------------------------- -# Service start / stop -# --------------------------------------------------------------------------- - -def _pid_file(project_dir: Path) -> Path: - return project_dir / ".pids" - - -def start_services(config: dict) -> dict[str, int]: - """Start the proxy and MLflow for a project. Returns {name: pid}.""" - project_dir = Path(config["project_dir"]) - pids = {} - - # Start MLflow - mlflow_port = config["mlflow"]["port"] - mlflow_backend = config["mlflow"]["backend"] - mlflow_dir = project_dir / "mlflow" - mlflow_dir.mkdir(exist_ok=True) - - if mlflow_backend == "sqlite": - backend_uri = f"sqlite:///{mlflow_dir}/mlflow.db" - else: - backend_uri = str(mlflow_dir) - - mlflow_cmd = [ - sys.executable, "-m", "mlflow", "server", - "--backend-store-uri", backend_uri, - "--default-artifact-root", str(mlflow_dir / "artifacts"), - "--host", "0.0.0.0", - "--port", str(mlflow_port), - ] - - mlflow_log = project_dir / "mlflow.log" - mlflow_proc = subprocess.Popen( - mlflow_cmd, - stdout=open(mlflow_log, "w"), - stderr=subprocess.STDOUT, - start_new_session=True, - ) - pids["mlflow"] = mlflow_proc.pid - logger.info("MLflow started (pid %d) → http://localhost:%d", mlflow_proc.pid, mlflow_port) - - # Start proxy - proxy_port = config["proxy"]["port"] - otel_endpoint = f"http://localhost:{mlflow_port}" - trace_dir = str(project_dir / "trace_archive") - - proxy_cmd = [ - sys.executable, "-m", "dsagt.proxy", - "--port", str(proxy_port), - "--records-dir", trace_dir, - "--session", config["project"], - "--otel-endpoint", otel_endpoint, - "--model", config["llm"]["model"], - ] - - proxy_log = project_dir / "proxy.log" - proxy_proc = subprocess.Popen( - proxy_cmd, - stdout=open(proxy_log, "w"), - stderr=subprocess.STDOUT, - env={**os.environ, "DSAGT_PROJECT": config["project"]}, - start_new_session=True, - ) - pids["proxy"] = proxy_proc.pid - logger.info("Proxy started (pid %d) → http://localhost:%d", proxy_proc.pid, proxy_port) - - # Save PIDs - pid_path = _pid_file(project_dir) - pid_path.write_text(json.dumps(pids, indent=2) + "\n") - - return pids - - -def stop_services(project_dir: str | Path) -> list[str]: - """Stop running services for a project. Returns descriptions of what was stopped.""" - project_dir = Path(project_dir) - pid_path = _pid_file(project_dir) - stopped = [] - - if not pid_path.exists(): - return ["No running services found."] - - pids = json.loads(pid_path.read_text()) - - for name, pid in pids.items(): - try: - os.killpg(os.getpgid(pid), signal.SIGTERM) - stopped.append(f"Stopped {name} (pid {pid})") - except (ProcessLookupError, PermissionError): - stopped.append(f"{name} (pid {pid}) was not running") - - pid_path.unlink(missing_ok=True) - return stopped - - -# --------------------------------------------------------------------------- -# Agent launch -# --------------------------------------------------------------------------- - -def agent_env(config: dict) -> dict: - """Build the environment dict an agent process needs to inherit. - - Merges the current environment with DSAGT-specific variables so the - proxy intercept, project identity, and embedding key are all set. - """ - project_dir = config["project_dir"] - proxy_port = config["proxy"]["port"] - agent = config["agent"] - - env = dict(os.environ) - env["DSAGT_PROJECT"] = config["project"] - env["DSAGT_PROJECT_DIR"] = project_dir - - # Proxy routing — agent-specific env var - if agent in ("claude-code", "roo", "cline"): - env["ANTHROPIC_BASE_URL"] = f"http://localhost:{proxy_port}" - if agent == "goose": - env["OPENAI_HOST"] = f"http://localhost:{proxy_port}" - - # Embedding API key for MCP servers - embedding_key = config.get("embedding", {}).get("api_key", "") - if embedding_key and not embedding_key.startswith("${"): - env["LLM_API_KEY"] = embedding_key - - return env - - -def agent_command(config: dict) -> list[str] | None: - """Return the shell command to launch the agent, or None for VS Code agents.""" - commands = { - "claude-code": ["claude"], - "goose": ["goose", "session"], - } - return commands.get(config["agent"]) - - -def launch_agent(config: dict, working_dir: str | Path) -> int: - """Launch the agent process in the foreground with the correct environment. - - For CLI agents (Claude Code, Goose): runs the agent interactively and - returns its exit code. - - For VS Code agents (Roo, Cline): prints instructions and returns 0. - """ - working_dir = Path(working_dir) - env = agent_env(config) - cmd = agent_command(config) - - if cmd is None: - # VS Code agents can't be launched from the CLI - agent = config["agent"] - print(f"\n Open VS Code in: {working_dir}") - if agent == "roo": - print(" Switch to the DSAGT Pipeline Builder mode (Cmd+.)") - elif agent == "cline": - print(" Open the Cline panel and verify MCP servers are connected") - print() - return 0 - - logger.info("Launching: %s", " ".join(cmd)) - try: - result = subprocess.run(cmd, env=env, cwd=str(working_dir)) - return result.returncode - except FileNotFoundError: - logger.error("Command not found: %s", cmd[0]) - logger.error("Is %s installed?", config["agent"]) - return 1 - except KeyboardInterrupt: - return 0 diff --git a/tests/conftest.py b/tests/conftest.py index 3020f45..3f3a3fa 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -4,18 +4,16 @@ Integration-test credentials come from the project-root ``.env`` file — the same file smoke_test/run.sh sources. One source of truth beats two. -Tests that need real APIs use the ``embedding_config``, ``llm_config``, or -``integration_config`` fixtures. They FAIL loudly when ``.env`` is missing a -required variable, since an integration test with empty creds is worse than -a skipped one (silent false positive). +Tests that need real APIs use the ``embedding_config`` or ``llm_config`` +fixtures. They FAIL loudly when ``.env`` is missing a required variable, +since an integration test with empty creds is worse than a skipped one +(silent false positive). """ -import os import sys from pathlib import Path import pytest -import yaml # Ensure tests/ is on sys.path so mcp_helpers can be imported sys.path.insert(0, str(Path(__file__).parent)) @@ -23,26 +21,6 @@ _ENV_PATH = Path(__file__).parent.parent / ".env" -@pytest.fixture(autouse=True) -def _stub_mlflow_experiment_lookup(monkeypatch): - """Short-circuit ``_resolve_experiment_id`` in unit tests. - - Most tests pass ``mlflow.port=5001`` but don't actually run an MLflow - server. Real ``mlflow.set_experiment`` would block on the unreachable - URL until socket timeout (tens of seconds × every test that builds - ``agent_env``). Returning a fake numeric id keeps the env block - populated without paying for the lookup. - - Tests that DO run a real MLflow (the integration suite, smoke tests) - pass through unchanged because ``mlflow.set_experiment`` is the only - thing this stubs — the real network calls happen elsewhere. - """ - monkeypatch.setattr( - "dsagt.agents._resolve_experiment_id", - lambda mlflow_url, project_name: "1", - ) - - def _load_env() -> dict: """Parse .env as KEY=VALUE lines. Fails if the file is missing. @@ -74,7 +52,7 @@ def _require(env: dict, key: str) -> str: @pytest.fixture def embedding_config(monkeypatch): - """Embedding endpoint from .env. Sets env vars for APIEmbeddingClient.""" + """Embedding endpoint from .env. Sets env vars for APIEmbedder.""" env = _load_env() base_url = _require(env, "EMBEDDING_BASE_URL") api_key = _require(env, "EMBEDDING_API_KEY") @@ -97,48 +75,3 @@ def llm_config(): "base_url": _require(env, "LLM_BASE_URL"), "api_key": _require(env, "LLM_API_KEY"), } - - -@pytest.fixture -def integration_config(tmp_path, monkeypatch): - """Full DSAGT config as ``load_config()`` would return it. - - Builds a dsagt_config.yaml in a temp project dir from .env values, then - patches the registry so ``load_config(project_name)`` resolves to it. - """ - from dsagt.session import load_config - - env = _load_env() - project_name = "test-integration" - config_data = { - "project": project_name, - "agent": "claude", - "llm": { - "provider": _require(env, "LLM_PROVIDER"), - "model": _require(env, "LLM_MODEL"), - "base_url": _require(env, "LLM_BASE_URL"), - "api_key": _require(env, "LLM_API_KEY"), - }, - "embedding": { - "model": _require(env, "EMBEDDING_MODEL"), - "base_url": _require(env, "EMBEDDING_BASE_URL"), - "api_key": _require(env, "EMBEDDING_API_KEY"), - }, - "proxy": {"port": 14000}, - "mlflow": {"port": 15001, "backend": "sqlite"}, - } - - project_dir = tmp_path / project_name - project_dir.mkdir(parents=True) - for subdir in ("trace_archive", "mlflow", "tools", "tools/code", "skills", "kb_index"): - (project_dir / subdir).mkdir() - (project_dir / "dsagt_config.yaml").write_text( - yaml.dump(config_data, default_flow_style=False, sort_keys=False) - ) - - monkeypatch.setattr( - "dsagt.session._load_registry", - lambda: {project_name: str(project_dir)}, - ) - - return load_config(project_name) diff --git a/tests/smoke_test/manual_runs/cli_walkthrough.md b/tests/manual_walkthroughs/cli_walkthrough.md similarity index 91% rename from tests/smoke_test/manual_runs/cli_walkthrough.md rename to tests/manual_walkthroughs/cli_walkthrough.md index 72a82a6..18923f9 100644 --- a/tests/smoke_test/manual_runs/cli_walkthrough.md +++ b/tests/manual_walkthroughs/cli_walkthrough.md @@ -1,6 +1,6 @@ # BYOA CLI hand-test -Tests goose, codex, opencode, and claude in the BYOA-CLI flow (no proxy, no IDE extension). +Tests goose, codex, opencode, and claude in the BYOA-CLI flow (no IDE extension). Replace `` below with one of: `goose`, `codex`, `opencode`, `claude`. @@ -44,11 +44,11 @@ cd ~/dsagt-projects/smoke- && 2. > I have a CSV utility called `csvtool`. Its reference is at `$SMOKE_DIR/knowledge/api_reference.md` — register the `filter` subcommand. Use an underscore in the name. - **Expect:** `save_tool_spec` MCP call; `~/dsagt-projects/smoke-/tools/csvtool_filter.md` exists. + **Expect:** `save_code_spec` MCP call; `~/dsagt-projects/smoke-/codes/csvtool_filter.md` exists. -3. > Use the `scan_directory` tool from the registry to scan `$SMOKE_DIR/data/`. +3. > Use the `scan-directory` tool from the registry to scan `$SMOKE_DIR/data/`. - **Expect:** Agent invokes `dsagt-run --tool scan_directory ...`; one record in `~/dsagt-projects/smoke-/trace_archive/scan_directory_*.json`. + **Expect:** Agent invokes `dsagt-run --code scan-directory ...`; one record in `~/dsagt-projects/smoke-/trace_archive/scan-directory_*.json`. 4. > Look at `$SMOKE_DIR/data/samples.csv` and summarize — columns, row count, quality issues. @@ -68,7 +68,7 @@ Exit the agent (Ctrl+C / `/exit`). ```bash dsagt info smoke- -ls ~/dsagt-projects/smoke-/tools/ +ls ~/dsagt-projects/smoke-/codes/ ls ~/dsagt-projects/smoke-/trace_archive/ test -s ~/dsagt-projects/smoke-/explicit_memories.yaml && echo OK ``` diff --git a/tests/smoke_test/manual_runs/vscode_walkthrough.md b/tests/manual_walkthroughs/vscode_walkthrough.md similarity index 94% rename from tests/smoke_test/manual_runs/vscode_walkthrough.md rename to tests/manual_walkthroughs/vscode_walkthrough.md index 8349bb8..c81accc 100644 --- a/tests/smoke_test/manual_runs/vscode_walkthrough.md +++ b/tests/manual_walkthroughs/vscode_walkthrough.md @@ -40,7 +40,7 @@ Inside VS Code: 1. > Ingest the docs in `$SMOKE_DIR/knowledge/` into a collection named `knowledge`. 2. > I have a CSV utility called `csvtool`. Its reference is at `$SMOKE_DIR/knowledge/api_reference.md` — register the `filter` subcommand. Use an underscore in the name. -3. > Use the `scan_directory` tool from the registry to scan `$SMOKE_DIR/data/`. +3. > Use the `scan-directory` tool from the registry to scan `$SMOKE_DIR/data/`. 4. > Look at `$SMOKE_DIR/data/samples.csv` and summarize — columns, row count, quality issues. 5. > Put this in explicit memory: samples.csv has null values in the status and timestamp columns. 6. > What do you remember about the samples dataset? @@ -53,7 +53,7 @@ Close the extension chat / VS Code session when done. ```bash dsagt info smoke- -ls ~/dsagt-projects/smoke-/tools/ +ls ~/dsagt-projects/smoke-/codes/ ls ~/dsagt-projects/smoke-/trace_archive/ test -s ~/dsagt-projects/smoke-/explicit_memories.yaml && echo OK ``` @@ -61,7 +61,7 @@ test -s ~/dsagt-projects/smoke-/explicit_memories.yaml && echo OK In MLflow UI: - **claude extension:** same trace shape as claude CLI — full `api_response_body` + `tool_use` payloads in OTel log events (gated by `CLAUDE_CODE_*` / `OTEL_LOG_*` env vars; the extension reads them from the shell or VS Code's `terminal.integrated.env.osx` setting). -- **roo extension:** no agent LLM-call OTel (roo emits none natively). MCP-server spans (`kb.*`, `registry.*`) and `dsagt-run` `tool.execute` spans only. For LLM-call visibility with roo, switch to the proxy walkthrough. +- **roo extension:** no agent LLM-call OTel (roo emits none natively). MCP-server spans (`kb.*`, `registry.*`) and `dsagt-run` `tool.execute` spans only. ## Notes diff --git a/tests/mcp_helpers.py b/tests/mcp_helpers.py index 62c6a3a..007e371 100644 --- a/tests/mcp_helpers.py +++ b/tests/mcp_helpers.py @@ -14,11 +14,11 @@ import mcp.types as types - # --------------------------------------------------------------------------- # In-process MCP tool invocation (for unit tests) # --------------------------------------------------------------------------- + def call_tool_sync(server, name: str, arguments: dict) -> str: """Invoke a tool handler on an MCP server and return the response text.""" req = types.CallToolRequest( @@ -105,54 +105,72 @@ def read_mcp_message(proc, timeout: float = 10.0, expect_id=None) -> dict: def mcp_initialize(proc) -> dict: """Send MCP initialize handshake and return the server's response.""" - send_mcp_message(proc, { - "jsonrpc": "2.0", - "id": 1, - "method": "initialize", - "params": { - "protocolVersion": "2024-11-05", - "capabilities": {}, - "clientInfo": {"name": "test", "version": "1.0"}, + send_mcp_message( + proc, + { + "jsonrpc": "2.0", + "id": 1, + "method": "initialize", + "params": { + "protocolVersion": "2024-11-05", + "capabilities": {}, + "clientInfo": {"name": "test", "version": "1.0"}, + }, }, - }) + ) response = read_mcp_message(proc, expect_id=1) - send_mcp_message(proc, { - "jsonrpc": "2.0", - "method": "notifications/initialized", - }) + send_mcp_message( + proc, + { + "jsonrpc": "2.0", + "method": "notifications/initialized", + }, + ) return response def mcp_list_tools(proc) -> dict: """Request tools/list and return the response.""" - send_mcp_message(proc, { - "jsonrpc": "2.0", - "id": 2, - "method": "tools/list", - "params": {}, - }) + send_mcp_message( + proc, + { + "jsonrpc": "2.0", + "id": 2, + "method": "tools/list", + "params": {}, + }, + ) return read_mcp_message(proc, expect_id=2) -def mcp_call_tool(proc, tool_name: str, arguments: dict, - msg_id: int = 3, timeout: float = 30.0) -> dict: +def mcp_call_tool( + proc, tool_name: str, arguments: dict, msg_id: int = 3, timeout: float = 30.0 +) -> dict: """Call an MCP tool and return the JSON-RPC response.""" - send_mcp_message(proc, { - "jsonrpc": "2.0", - "id": msg_id, - "method": "tools/call", - "params": { - "name": tool_name, - "arguments": arguments, + send_mcp_message( + proc, + { + "jsonrpc": "2.0", + "id": msg_id, + "method": "tools/call", + "params": { + "name": tool_name, + "arguments": arguments, + }, }, - }) + ) return read_mcp_message(proc, timeout=timeout, expect_id=msg_id) -def start_server(cmd: list[str], env: dict = None) -> subprocess.Popen: - """Start a server subprocess with stdio pipes.""" +def start_server(cmd: list[str], env: dict = None, cwd: str = None) -> subprocess.Popen: + """Start a server subprocess with stdio pipes. + + ``cwd`` sets the working directory — ``dsagt-server`` discovers its project + config from cwd (see ``observability.find_project_config``), so startup tests + must run it from the project dir. + """ proc_env = os.environ.copy() if env: proc_env.update(env) @@ -164,4 +182,5 @@ def start_server(cmd: list[str], env: dict = None) -> subprocess.Popen: stderr=subprocess.PIPE, text=True, env=proc_env, + cwd=cwd, ) diff --git a/tests/smoke_test/csv_summary.py b/tests/smoke_test/csv_summary.py new file mode 100644 index 0000000..bcb8916 --- /dev/null +++ b/tests/smoke_test/csv_summary.py @@ -0,0 +1,92 @@ +"""csv_summary — the quickstart's registerable, executable fixture code. + +Stdlib-only (``csv`` + ``statistics``), so registering and running it via +dsagt-run needs no dependency install — the dependency-free replacement for +the csvkit demo. Summarizes a CSV: column list, row count, per-column null +count, and min/max/mean for numeric columns. On ``samples.csv`` the null +counts surface the empty ``status`` / ``timestamp`` cells, which is the fact +the quickstart's explicit-memory step stores and recalls. +""" + +import argparse +import csv +import json +import statistics +import sys + + +def _as_number(value): + if value is None or value.strip() == "": + return None + try: + return float(value) + except ValueError: + return None + + +def summarize(path): + with open(path, newline="") as fh: + reader = csv.DictReader(fh) + columns = reader.fieldnames or [] + rows = list(reader) + + nulls = {c: 0 for c in columns} + numeric = {c: [] for c in columns} + for row in rows: + for c in columns: + v = row.get(c) + if v is None or v.strip() == "": + nulls[c] += 1 + else: + n = _as_number(v) + if n is not None: + numeric[c].append(n) + + # A column is "numeric" only if every non-null value parsed as a number. + numeric_stats = {} + for c in columns: + vals = numeric[c] + non_null = len(rows) - nulls[c] + if vals and len(vals) == non_null: + numeric_stats[c] = { + "min": min(vals), + "max": max(vals), + "mean": round(statistics.fmean(vals), 4), + } + + return { + "file": path, + "columns": columns, + "row_count": len(rows), + "null_counts": nulls, + "columns_with_nulls": [c for c in columns if nulls[c] > 0], + "numeric_stats": numeric_stats, + } + + +def main(): + parser = argparse.ArgumentParser( + description="Summarize a CSV: columns, row count, null counts, numeric stats." + ) + parser.add_argument("file", help="Path to the CSV file") + args = parser.parse_args() + + try: + result = summarize(args.file) + except FileNotFoundError: + print( + json.dumps( + { + "status": "error", + "code": "CSV-404", + "error": f"no such file: {args.file}", + } + ) + ) + sys.exit(1) + + print(json.dumps(result, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/tests/smoke_test/greet.py b/tests/smoke_test/greet.py index 1ca5f46..0fa72cc 100755 --- a/tests/smoke_test/greet.py +++ b/tests/smoke_test/greet.py @@ -1,11 +1,26 @@ -"""Simple greeting script for smoke testing.""" +"""greet — the smoke test's registerable, executable fixture CLI. + +Stdlib-only so registration + execution via dsagt-run needs no +dependency install. The GRT-42 empty-name error code is documented in +knowledge/troubleshooting.md; the smoke script asks the agent to +retrieve it from the knowledge collection, so keep code and docs in +sync. +""" + import argparse import json +import sys parser = argparse.ArgumentParser(description="Generate a greeting") parser.add_argument("name", help="Name to greet") -parser.add_argument("--greeting", default="Hello", help="Greeting word (default: Hello)") +parser.add_argument( + "--greeting", default="Hello", help="Greeting word (default: Hello)" +) args = parser.parse_args() +if not args.name.strip(): + print(json.dumps({"status": "error", "code": "GRT-42", "error": "empty name"})) + sys.exit(1) + result = {"message": f"{args.greeting}, {args.name}!", "status": "ok"} print(json.dumps(result, indent=2)) diff --git a/tests/smoke_test/knowledge/DESCRIPTION.md b/tests/smoke_test/knowledge/DESCRIPTION.md index 1ec721d..08f904b 100644 --- a/tests/smoke_test/knowledge/DESCRIPTION.md +++ b/tests/smoke_test/knowledge/DESCRIPTION.md @@ -1,2 +1,5 @@ Sample documentation for smoke testing the DSAGT knowledge base. -Covers a fictional "csvtool" CLI for CSV data processing. +Covers the "greet" CLI fixture (tests/smoke_test/greet.py) — the same +utility the smoke script has the agent register and execute, so the +ingested docs, the registered code, and the provenance records all +describe one consistent tool. diff --git a/tests/smoke_test/knowledge/api_reference.md b/tests/smoke_test/knowledge/api_reference.md index 18cd7a9..d74b21e 100644 --- a/tests/smoke_test/knowledge/api_reference.md +++ b/tests/smoke_test/knowledge/api_reference.md @@ -1,35 +1,28 @@ -# csvtool CLI Reference +# greet CLI Reference -## csvtool filter - -Filter rows by column value. +Generate a JSON greeting for a name. ```bash -csvtool filter --input data.csv --column status --value active --output filtered.csv +python greet.py NAME [--greeting WORD] ``` Parameters: -- `--input` (required) — Input CSV file path -- `--column` (required) — Column name to filter on -- `--value` (required) — Value to match -- `--output` (required) — Output CSV file path +- `NAME` (positional, required) — Name to greet. +- `--greeting` (optional) — Greeting word. Default: `Hello`. -## csvtool validate +## Output -Check a CSV against a schema. +On success, prints a JSON object to stdout and exits 0: -```bash -csvtool validate --input data.csv --schema schema.json +```json +{ + "message": "Hello, DSAGT!", + "status": "ok" +} ``` -Outputs a JSON report with row count, null counts per column, type violations, and an overall pass/fail status. - -## csvtool summary - -Print column statistics. - -```bash -csvtool summary --input data.csv -``` +## Exit codes -Outputs: column names, types, null counts, unique value counts, and min/max for numeric columns. +- `0` — success. +- `1` — invalid input. The JSON output carries `"status": "error"` and an + error `code` field (see troubleshooting). diff --git a/tests/smoke_test/knowledge/installation.md b/tests/smoke_test/knowledge/installation.md index ab83d92..09f225e 100644 --- a/tests/smoke_test/knowledge/installation.md +++ b/tests/smoke_test/knowledge/installation.md @@ -1,11 +1,12 @@ -# csvtool Installation +# greet Installation -Install csvtool via pip: +greet ships as a single-file script (`greet.py`) with the smoke test +fixtures — there is nothing to install. -```bash -pip install csvtool -``` +Requirements: Python 3.10+, standard library only (argparse, json, sys). -Requirements: Python 3.10+, pandas >= 2.0. +Run it directly: -csvtool provides CLI commands for filtering, validating, and summarizing CSV files. +```bash +python greet.py World --greeting Hi +``` diff --git a/tests/smoke_test/knowledge/troubleshooting.md b/tests/smoke_test/knowledge/troubleshooting.md index 94c4af1..f40c62e 100644 --- a/tests/smoke_test/knowledge/troubleshooting.md +++ b/tests/smoke_test/knowledge/troubleshooting.md @@ -1,18 +1,25 @@ -# csvtool Troubleshooting +# greet Troubleshooting -## Encoding errors +## Error GRT-42: empty name -If csvtool raises a UnicodeDecodeError, the file may not be UTF-8. Use `--encoding latin-1` or detect encoding with `chardet`. +When the `NAME` argument is empty or whitespace-only, greet prints a JSON +error object with error code `GRT-42` and exits 1: -## Large file performance +```json +{"status": "error", "code": "GRT-42", "error": "empty name"} +``` -For files over 1 GB, csvtool streams rows in chunks. If you still run out of memory, increase the chunk size with `--chunk-size 100000` or filter to specific columns with `--columns col1,col2`. +Fix: pass a non-empty name. Shell quoting is the usual culprit — an unset +variable like `"$USER_NAME"` expands to an empty string. -## Schema validation failures +## Output is not valid JSON -The validate command checks column types against a JSON schema. Common issues: -- Mixed types in a column (e.g., "123" and "abc" in an integer column) -- Missing required columns -- Null values in non-nullable columns +greet writes only JSON to stdout. If you see extra text mixed in, another +tool in your pipeline is writing to the same stream — redirect its output +to stderr. -Fix the data or update the schema to match reality. +## Wrong greeting word + +`--greeting` must come after the name or be joined with `=` +(`--greeting=Ahoy`); a bare `--greeting` at the end consumes the name as +its value and fails with a missing-argument error. diff --git a/tests/smoke_test/manual_runs/proxy_walkthrough.md b/tests/smoke_test/manual_runs/proxy_walkthrough.md deleted file mode 100644 index 55340f9..0000000 --- a/tests/smoke_test/manual_runs/proxy_walkthrough.md +++ /dev/null @@ -1,109 +0,0 @@ -# Proxy mode hand-test - -Tests goose, claude, roo, cline, codex, opencode through `dsagt start --enable-proxy`. - -The proxy intercepts the agent's LLM calls, translates lab-gateway-aliased model names back to upstream-served names (`_AGENT_PRIMARY_ALIASES`), forwards via LiteLLM to the configured upstream, and emits OTLP spans to MLflow tagged with `service.name=dsagt-proxy`. Plus cache-breakpoint injection on outgoing requests + sidechannel detection on responses. - -Replace `` below with one of: `goose`, `claude`, `roo`, `cline`, `codex`, `opencode`. - -## Setup (once per agent) - -```bash -cd ~/dsagt -source .venv/bin/activate -export SMOKE_DIR="$(pwd)/tests/smoke_test" - -# Proxy mode reads upstream creds from .env (or shell env). These four -# are required: -export LLM_PROVIDER=openai -export LLM_MODEL=claude-haiku-4-5-20251001-v1-project -export LLM_BASE_URL=https://ai-incubator-api.pnnl.gov -export LLM_API_KEY=sk-... -# If `embedding.backend: api` in your project config, also set -# EMBEDDING_PROVIDER / EMBEDDING_MODEL / EMBEDDING_BASE_URL / -# EMBEDDING_API_KEY. Default `local` mode needs none of these. - -dsagt rm smoke- -y >/dev/null 2>&1 -dsagt init smoke- --agent -``` - -## Run - -```bash -# Terminal 1 -dsagt mlflow smoke- - -# Terminal 2 -dsagt start smoke- --enable-proxy -``` - -This spawns: -- MLflow on the pinned port (from `dsagt_config.yaml`) -- `dsagt-proxy` on a free port (its YAML config is regenerated each run) -- The agent, with its env / config files routed at the proxy URL + a sentinel API key (real upstream creds live only in the proxy subprocess) - -## Scripted prompts (paste one at a time) - -Same six prompts as `cli_walkthrough.md`: - -1. > Ingest the docs in `$SMOKE_DIR/knowledge/` into a collection named `knowledge`. -2. > I have a CSV utility called `csvtool`. Its reference is at `$SMOKE_DIR/knowledge/api_reference.md` — register the `filter` subcommand. Use an underscore in the name. -3. > Use the `scan_directory` tool from the registry to scan `$SMOKE_DIR/data/`. -4. > Look at `$SMOKE_DIR/data/samples.csv` and summarize — columns, row count, quality issues. -5. > Put this in explicit memory: samples.csv has null values in the status and timestamp columns. -6. > What do you remember about the samples dataset? - -Exit the agent. - -## Verify - -```bash -dsagt info smoke- -ls ~/dsagt-projects/smoke-/tools/ -ls ~/dsagt-projects/smoke-/trace_archive/ -test -s ~/dsagt-projects/smoke-/explicit_memories.yaml && echo OK -``` - -In MLflow UI: - -- **All six agents:** agent LLM-call traces with `service.name = "dsagt-proxy"`, full request `messages` + response payloads (LiteLLM autolog shape), token counts, latency. -- Plus the usual MCP / dsagt-run spans (separate `service.name`s). - -If the agent emits its own OTel too (claude/goose), you'll see those spans alongside the dsagt-proxy ones — same trace store, different service names. - -## Sidechannel warnings - -At end-of-session, if any agent-internal "title generator" / "session namer" calls fired (model names not in your `LLM_MODEL`), the teardown prints: - -``` - ⚠ Sidechannel model calls intercepted: - gpt-4o-mini (2 calls) - Two possible causes: - (1) agent sidechannel — safe to ignore - (2) typo in dsagt_config.yaml llm.model — these replies are canned, not real -``` - -Confirms the wildcard catchall + canned-response mock fired correctly. - -## Per-agent expected behavior - -| Agent | What proxy mode adds beyond BYOA | -|---|---| -| goose | Uniform trace shape; MLflow now has every LLM call from one parser-friendly source | -| claude | Sidesteps the macOS Keychain OAuth conflict — claude's auth token gets ignored by the proxy, which uses `LLM_API_KEY` for upstream | -| roo | **Un-punted.** Roo's anthropic SDK posts to `ANTHROPIC_BASE_URL` (the proxy URL); the proxy aliases roo's hardcoded `claude-sonnet-4-5` rewrite back to your upstream model | -| cline | **Un-punted.** `cline auth -p openai -b ` configures cline to talk openai-shape to the proxy; the proxy translates and aliases | -| codex | LLM-call message payloads now visible (codex's bundled OTel only emits token counts) | -| opencode | LLM-call traces visible (opencode emits no native OTel at all) | - -## Verify proxy logs - -```bash -tail -50 ~/dsagt-projects/smoke-/proxy.log -``` - -Look for: -- `init_proxy_tracing: service=dsagt-proxy mlflow=...` -- `Generated LiteLLM config at /tmp/dsagt_litellm_*.yaml` -- `Uvicorn running on http://0.0.0.0:` -- One `litellm` log line per agent LLM call diff --git a/tests/smoke_test/run.sh b/tests/smoke_test/run.sh index 2973567..9b9595a 100755 --- a/tests/smoke_test/run.sh +++ b/tests/smoke_test/run.sh @@ -2,11 +2,18 @@ # DSAGT smoke test — non-interactive end-to-end exercise. # # Drives the SAME `dsagt start` lifecycle as an interactive run (config -# generation → start_services → agent → run_extraction → stop_services). -# Only the agent-launch step swaps from `goose session` to `goose run -i`. +# generation → agent in the foreground → post-session run_extraction). +# Serverless: there are no services to start or stop — all self-logging +# lands in the project's sqlite MLflow store. Only the agent-launch +# step swaps from interactive to batch (`--script`). +# +# TWO sessions run back-to-back: session 1 exercises ingest, code +# registration + execution, provenance, KB retrieval, skill install, +# and explicit memory; session 2 exercises cross-session recall, +# registry persistence, and the startup catch-up path. +# # BYOA: the user's shell must already have the agent's provider creds -# (per `dsagt init` hints). No .env handling — that returns with the -# Phase 2 `--proxy_traces` flag. +# (per `dsagt init` hints). No .env handling. # # Run from anywhere: # bash tests/smoke_test/run.sh @@ -19,67 +26,40 @@ set -uo pipefail AGENT="${DSAGT_SMOKE_AGENT:-${1:-goose}}" # arg or env var, default goose # Per-agent project name so each agent's mlflow.db, trace_archive, and # kb_index/ survive across runs — crucial for cross-agent comparison -# (e.g., why does claude use 10x the tokens roo does?). Without this, +# (e.g., why does claude use 10x the tokens codex does?). Without this, # `dsagt rm` at the start of each run wipes the previous agent's state. PROJECT="smoke-test-${AGENT}" SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" DSAGT_ROOT="$(cd "${SCRIPT_DIR}/../.." && pwd)" -SCRIPT_FILE="${SCRIPT_DIR}/script.txt" # Project lives at the default ``dsagt init`` location so smoke-test # artifacts stay out of the dsagt source tree. PDIR mirrors # DEFAULT_PROJECTS_BASE in src/dsagt/session.py. PDIR="${HOME}/dsagt-projects/${PROJECT}" case "${AGENT}" in - goose|claude|cline|roo|codex|opencode) ;; + goose|claude|codex|opencode) ;; + cline) + # dsagt start --script hard-errors for cline: its headless CLI + # (verified 3.0.34) never loads MCP servers, so a scripted session + # has no dsagt tools to exercise — see agents/cline.py. Skip rather + # than report red checks; drop this arm when cline ships MCP in + # headless mode. + echo "[smoke] SKIP: cline headless CLI loads no MCP servers (see agents/cline.py) — hand-test via tests/manual_walkthroughs/ instead" + exit 0 + ;; *) - echo "ERROR: agent must be one of: goose, claude, cline, roo, codex, opencode (got '${AGENT}')" >&2 + echo "ERROR: agent must be one of: goose, claude, cline, codex, opencode (got '${AGENT}')" >&2 exit 2 ;; esac -# Proxy-mode opt-in. ``DSAGT_SMOKE_PROXY=1`` (or arg #2 == "proxy") -# routes the run through ``dsagt start --enable-proxy``, which spawns -# dsagt-proxy and forwards the agent's LLM calls through it — required -# for cline / roo (their CLIs hardcode model-ID whitelists incompatible -# with lab-gateway aliases; the proxy translates names back to the -# upstream's served IDs via _AGENT_PRIMARY_ALIASES). Default off. -PROXY_FLAG="" -if [[ "${DSAGT_SMOKE_PROXY:-${2:-}}" == "1" || "${DSAGT_SMOKE_PROXY:-${2:-}}" == "proxy" ]]; then - PROXY_FLAG="--enable-proxy" - echo "[smoke] Agent: ${AGENT} (proxy mode)" -else - echo "[smoke] Agent: ${AGENT}" -fi - -# Cline / roo only work in proxy mode (see agents/cline.py + roo.py -# module docstrings for the model-whitelist + endpoint-lockout reasons). -if [[ -z "${PROXY_FLAG}" && ( "${AGENT}" == "cline" || "${AGENT}" == "roo" ) ]]; then - echo - echo "[smoke] ${AGENT} requires proxy mode in BYOA — skipping." - echo "[smoke] Re-run with: DSAGT_SMOKE_PROXY=1 dsagt smoke-test --agent ${AGENT}" - echo "[smoke] (or pass 'proxy' as the second arg to this script)" - exit 0 -fi +echo "[smoke] Agent: ${AGENT}" cd "${DSAGT_ROOT}" -# Proxy mode needs LLM_*/EMBEDDING_* in env so the proxy can forward -# upstream. Source .env if present (BYOA runs without it; only proxy -# runs need these vars). Validation happens at proxy spawn time — -# session.py:_start_proxy raises if config["llm"] keys are missing. -if [[ -n "${PROXY_FLAG}" && -f .env ]]; then - set -a - # shellcheck disable=SC1091 - source .env - set +a - echo "[smoke] Sourced .env for proxy mode" -fi - # --------------------------------------------------------------------------- -# 2. Clean slate (idempotent — silent if nothing exists) +# 1. Clean slate (idempotent — silent if nothing exists) # --------------------------------------------------------------------------- -dsagt stop "${PROJECT}" >/dev/null 2>&1 || true dsagt rm "${PROJECT}" -y >/dev/null 2>&1 || true rm -rf "${PDIR}" @@ -95,64 +75,87 @@ if [[ "${AGENT}" == "claude" ]]; then fi # --------------------------------------------------------------------------- -# 3. Init at the default ``~/dsagt-projects/`` location so smoke artifacts -# don't pollute the dsagt source tree. The agent's cwd will be -# ``${PDIR}``; the script template uses ``{{SMOKE_DIR}}`` placeholders -# that we substitute below so prompt paths resolve regardless of where -# PDIR lives. +# 2. Init at the default ``~/dsagt-projects/`` location so smoke artifacts +# don't pollute the dsagt source tree. --episodic so session turns land +# in the ``session_memory`` collection (asserted below). The default KB +# set includes the genesis skill catalog, which the skill-install prompt +# relies on. # --------------------------------------------------------------------------- -dsagt init "${PROJECT}" --agent "${AGENT}" +# Force the non-interactive (flag-driven) init path regardless of TTY by +# closing stdin — `dsagt init` prompts only when stdin is a TTY. +dsagt init "${PROJECT}" --agent "${AGENT}" --episodic < /dev/null # Substitute {{SMOKE_DIR}} → absolute smoke_test/ path before the agent -# sees the script. Prompts can then reference data/knowledge files via +# sees the scripts. Prompts can then reference data/knowledge files via # absolute paths regardless of the agent's cwd. -RENDERED_SCRIPT=$(mktemp -t dsagt-smoke-script.XXXXXX) -trap 'rm -f "${RENDERED_SCRIPT}"' EXIT -sed "s|{{SMOKE_DIR}}|${SCRIPT_DIR}|g" "${SCRIPT_FILE}" > "${RENDERED_SCRIPT}" +RENDERED_SCRIPT_1=$(mktemp -t dsagt-smoke-script1.XXXXXX) +RENDERED_SCRIPT_2=$(mktemp -t dsagt-smoke-script2.XXXXXX) +SESSION_LOG_1=$(mktemp -t dsagt-smoke-log1.XXXXXX) +SESSION_LOG_2=$(mktemp -t dsagt-smoke-log2.XXXXXX) +trap 'rm -f "${RENDERED_SCRIPT_1}" "${RENDERED_SCRIPT_2}" "${SESSION_LOG_1}" "${SESSION_LOG_2}"' EXIT +sed "s|{{SMOKE_DIR}}|${SCRIPT_DIR}|g" "${SCRIPT_DIR}/script.txt" > "${RENDERED_SCRIPT_1}" +sed "s|{{SMOKE_DIR}}|${SCRIPT_DIR}|g" "${SCRIPT_DIR}/script2.txt" > "${RENDERED_SCRIPT_2}" # --------------------------------------------------------------------------- -# 4. Run the FULL `dsagt start` lifecycle, with the agent in batch mode. -# Wall-clock cap belt-and-suspenders the --max-turns inside. +# 3. Session runner: the FULL `dsagt start` lifecycle with the agent in +# batch mode, under a wall-clock watchdog. # # Pure-bash watcher pattern instead of GNU `timeout` so the smoke test # works on stock macOS without `brew install coreutils`. SIGTERM gives -# dsagt's finally-block a chance to stop services cleanly; the +# dsagt's finally-block a chance to run post-session extraction; the # follow-up SIGKILL after WALL_CLOCK_GRACE catches the agent if it # swallows the term signal. +# +# Output tees to a per-session log — the retrieval and recall +# assertions grep it for facts the agent can only have gotten from +# the KB / memory (process substitution keeps $! on dsagt itself). # --------------------------------------------------------------------------- -WALL_CLOCK_CAP=300 # seconds (5 minutes) WALL_CLOCK_GRACE=10 # extra seconds before SIGKILL +run_session() { + local script_file="$1" max_turns="$2" cap="$3" log_file="$4" + dsagt start "${PROJECT}" --script "${script_file}" --max-turns "${max_turns}" \ + > >(tee "${log_file}") 2>&1 & + local pid=$! + ( + sleep "${cap}" + kill -TERM "${pid}" 2>/dev/null && \ + echo "[smoke] WARN: ${cap}s cap exceeded — sent SIGTERM to dsagt start (pid ${pid})" + sleep "${WALL_CLOCK_GRACE}" + kill -KILL "${pid}" 2>/dev/null && \ + echo "[smoke] WARN: dsagt start did not exit on SIGTERM — sent SIGKILL" + ) & + local watcher=$! + wait "${pid}" + local rc=$? + # Tear down the watcher if dsagt exited on its own. + kill -TERM "${watcher}" 2>/dev/null + wait "${watcher}" 2>/dev/null + # Let the tee process-substitution drain before the log is grepped. + sleep 1 + return "${rc}" +} + echo -echo "[smoke] Running dsagt start --script ${PROXY_FLAG} (${WALL_CLOCK_CAP}s wall-clock cap)…" -dsagt start "${PROJECT}" --script "${RENDERED_SCRIPT}" --max-turns 30 ${PROXY_FLAG} & -DSAGT_PID=$! -( - sleep "${WALL_CLOCK_CAP}" - kill -TERM "${DSAGT_PID}" 2>/dev/null && \ - echo "[smoke] WARN: ${WALL_CLOCK_CAP}s cap exceeded — sent SIGTERM to dsagt start (pid ${DSAGT_PID})" - sleep "${WALL_CLOCK_GRACE}" - kill -KILL "${DSAGT_PID}" 2>/dev/null && \ - echo "[smoke] WARN: dsagt start did not exit on SIGTERM — sent SIGKILL" -) & -WATCHER_PID=$! -wait "${DSAGT_PID}" +echo "[smoke] Session 1: ingest / register / execute / provenance / skills / memory…" +run_session "${RENDERED_SCRIPT_1}" 40 420 "${SESSION_LOG_1}" START_EXIT=$? -# Tear down the watcher if dsagt exited on its own. -kill -TERM "${WATCHER_PID}" 2>/dev/null -wait "${WATCHER_PID}" 2>/dev/null - if [[ ${START_EXIT} -ne 0 ]]; then echo "WARN: dsagt start exited non-zero (${START_EXIT}) — continuing to artifact checks anyway" fi -# Defensive: ensure no stray services if the lifecycle's finally didn't run -# (timeout SIGKILL skips Python's finally blocks). Output kept visible — -# if any port is still in use after this, we want to see the warning so -# the next run doesn't race a half-shutdown orphan. +# --------------------------------------------------------------------------- +# 4. Session 2: cross-session recall + registry persistence. Its startup +# also runs the catch-up path over session 1 (code-use indexing + the +# pinned trace re-collect), so the post-session-2 assertions cover it. +# --------------------------------------------------------------------------- echo -echo "[smoke] Final cleanup…" -dsagt stop "${PROJECT}" || true +echo "[smoke] Session 2: cross-session recall + catch-up…" +run_session "${RENDERED_SCRIPT_2}" 15 240 "${SESSION_LOG_2}" +START_EXIT_2=$? +if [[ ${START_EXIT_2} -ne 0 ]]; then + echo "WARN: session 2 dsagt start exited non-zero (${START_EXIT_2}) — continuing to artifact checks anyway" +fi # --------------------------------------------------------------------------- # 5. Artifact checks @@ -170,42 +173,78 @@ check() { fi } -check "csvtool_filter spec written" "test -f '${PDIR}/tools/csvtool_filter.md'" -check "trace_archive has records" "ls '${PDIR}/trace_archive/'*.json | grep -q ." -check "scan_directory record" "ls '${PDIR}/trace_archive/'*scan_directory*.json | grep -q ." -# Both files are written by dsagt-knowledge-server's kb_ingest_directory MCP -# tool — chroma.sqlite3 is the actual vector DB, route.json is the collection -# manifest. Checking only `test -d kb_index/knowledge` is too weak: an agent +# -- registry + execution + provenance -------------------------------------- +check "greet spec written" "test -f '${PDIR}/codes/greet/SKILL.md'" +# Codes share the skill-standard envelope and mirror into the agent's +# native skills dir at dsagt start: the bundled scan-directory at +# session 1's start, greet (registered mid-session-1) at session 2's. +check "bundled code mirrored natively" "find '${PDIR}' -path '*skills/scan-directory/SKILL.md' | grep -q ." +check "greet mirrored natively" "find '${PDIR}' -path '*skills/greet/SKILL.md' | grep -q ." +# The execution went through dsagt-run iff the record captured greet's +# actual stdout — an agent that ran the script by hand can't fake the +# trace_archive record. Match only the greeting prefix: it proves our +# custom --greeting arg flowed through the registered code, while +# tolerating an agent flubbing which word goes in the name slot (goose +# produced "Ahoy, Ahoy!"). +check "greet executed via dsagt-run" "grep -l 'Ahoy,' '${PDIR}/trace_archive/'*greet*.json" +check "greet re-run in session 2" "test \$(ls '${PDIR}/trace_archive/'*greet*.json | wc -l) -ge 2" +check "scan-directory record" "ls '${PDIR}/trace_archive/'*scan-directory*.json" + +# -- knowledge base ---------------------------------------------------------- +# Both files are written by dsagt-server's kb_ingest MCP tool — chroma.sqlite3 +# is the actual vector DB, chroma_ids.json the internal-collection manifest +# (route.json marks routed *external* collections, which ingest never +# creates). Checking only `test -d kb_index/knowledge` is too weak: an agent # can satisfy it by hand-crafting an empty directory tree, masking a broken -# MCP wiring (which is exactly what we hit when cline's dsagt-knowledge -# server crashed silently and the LLM compensated by mkdir-ing the path). -check "knowledge ingested (route)" "test -f '${PDIR}/kb_index/knowledge/route.json'" +# MCP wiring (which is exactly what we hit when cline's dsagt server crashed +# silently and the LLM compensated by mkdir-ing the path). +check "knowledge ingested (ids)" "test -f '${PDIR}/kb_index/knowledge/chroma_ids.json'" check "knowledge ingested (vectors)" "test -f '${PDIR}/kb_index/knowledge/chroma.sqlite3'" -# Explicit memory writes to /explicit_memories.yaml (YAML at the -# project root), NOT to kb_index/. Only kb_remember (called deliberately -# by the agent in response to "Put this in explicit memory" / "remember -# this") populates the file. End-of-session episodic extraction writes -# elsewhere (kb_index/episodic_memory/...) and is independent. Checking -# the YAML's existence + non-empty catches the hallucination case where -# the agent claims it stored a fact but didn't actually call the tool. -check "explicit memory recorded" "test -s '${PDIR}/explicit_memories.yaml'" -check "mlflow has traces" "test -s '${PDIR}/mlflow/mlflow.db'" +# GRT-42 lives only in knowledge/troubleshooting.md — the agent answering +# with it proves retrieval reached the ingested docs. +check "kb retrieval answered (GRT-42)" "grep -q 'GRT-42' '${SESSION_LOG_1}'" + +# -- skills ------------------------------------------------------------------ +check "catalog skill installed" "ls '${PDIR}/skills/'*/SKILL.md" + +# -- memory ------------------------------------------------------------------ +# Explicit memory lives with the server-owned internals in .dsagt/. Only +# kb_remember (called deliberately by the agent in response to "Put this in +# explicit memory") populates the file; checking non-empty catches the +# hallucination case where the agent claims it stored a fact but didn't +# actually call the tool. +check "explicit memory recorded" "test -s '${PDIR}/.dsagt/explicit_memories.yaml'" +# Cross-session recall: session 2's answer must carry the stored fact's +# tokens, which only kb_get_memories (or episodic retrieval) can supply — +# session 2 never saw samples.csv. +check "cross-session recall" "grep -qi 'null' '${SESSION_LOG_2}' && grep -qi 'status' '${SESSION_LOG_2}'" +# Episodic memory (enabled via --episodic) chunks+embeds every turn into +# the session_memory collection on the heartbeat. +check "episodic memory indexed" "test -f '${PDIR}/kb_index/session_memory/chroma.sqlite3'" + +# -- observability + session state ------------------------------------------- +check "mlflow store has traces" "test -s '${PDIR}/mlflow.db'" +# The heartbeat indexes trace_archive/ execution records into the code_use +# collection (plus a startup catch-up in session 2). +check "code_use collection indexed" "test -f '${PDIR}/kb_index/code_use/chroma.sqlite3'" +# state.yaml is the anchor for crash catch-up: both sessions logged, and +# session 1 carries the trace_source token the session-2 catch-up pinned. +check "state.yaml logged 2 sessions" "uv run --quiet python -c \" +import yaml, sys +s = yaml.safe_load(open('${PDIR}/.dsagt/state.yaml')) +sessions = s.get('sessions') or [] +sys.exit(0 if len(sessions) >= 2 and sessions[0].get('trace_source') else 1)\"" +check "dsagt info runs" "dsagt info '${PROJECT}'" # --------------------------------------------------------------------------- -# 6. Agent LLM-call observability: every agent turn must produce at least -# one trace in MLflow with this session's id, tagged with the agent's -# service.name. Replaces the proxy-log parity check from before -# proxy removal — the agent now emits OTel directly to MLflow's OTLP -# receiver, so MLflow IS the source of truth. A zero count means -# either the agent didn't emit telemetry (e.g. CLAUDE_CODE_ENABLE_TELEMETRY=1 -# not honored) or the OTel endpoint was misconfigured. +# 6. Agent LLM-call transparency: the trace pipeline recovers every agent's +# turns from its on-disk transcript (heartbeat + graceful-shutdown flush, +# backstopped by session 2's startup catch-up), so agent traces in the +# store are a hard requirement for all five agents. # --------------------------------------------------------------------------- -AGENT_OTEL_SUPPORT=$(uv run --quiet python -c "from dsagt.agents import agent_otel_support; print(agent_otel_support('${AGENT}'))" 2>/dev/null) -AGENT_OTEL_SUPPORT="${AGENT_OTEL_SUPPORT:-unknown}" - AGENT_TRACES=$(uv run --quiet python </dev/null import mlflow -mlflow.set_tracking_uri("sqlite:///${PDIR}/mlflow/mlflow.db") +mlflow.set_tracking_uri("sqlite:///${PDIR}/mlflow.db") exp = mlflow.get_experiment_by_name("${PROJECT}") if exp is None: print(0); raise SystemExit @@ -213,49 +252,22 @@ df = mlflow.search_traces( locations=[exp.experiment_id], max_results=500, ) -n = 0 -for _, row in df.iterrows(): - spans = row.get("spans") or [] - # Match by service.name on root span — agent-emitted traces only. - # MCP-server traces (kb.*, registry.*, tool.execute) carry - # service.name = "dsagt-knowledge-server" / "dsagt-registry-server" / - # "dsagt-run" and shouldn't count toward agent turn parity. - for s in spans: - attrs = getattr(s, "attributes", None) or ( - s.get("attributes") if isinstance(s, dict) else None - ) - if attrs and not str(attrs.get("service.name", "")).startswith("dsagt-"): - n += 1 - break +# MLflowSink stamps every replayed agent trace with "dsagt.trace_id" in +# its trace metadata; DSAGT's internal MCP/dsagt-run debug traces carry a +# "dsagt.source" tag instead — the positive marker is the reliable +# filter. A service.name span heuristic previously counted internal +# spans lacking that attribute as agent traces, masking a codex reader +# that collected nothing. +n = sum( + 1 + for _, row in df.iterrows() + if "dsagt.trace_id" in (row.get("trace_metadata") or {}) +) print(n) PY ) AGENT_TRACES="${AGENT_TRACES:-0}" - -# Grade by the agent's verified support tier rather than fail-or-pass: -# agents we know don't emit OTel payloads (cline, roo) get SKIP, agents -# we know partial-emit (codex) get WARN, agents we know full-emit -# (claude, goose) get PASS-or-FAIL. See agents/__init__.py module -# docstring for the matrix. -case "${AGENT_OTEL_SUPPORT}" in - full) - if [[ "${AGENT_TRACES}" -eq 0 ]]; then - echo " FAIL agent transparency: 0 agent LLM-call traces in MLflow (agent supports full payload but emitted none — env vars not honored?)" - FAIL=1 - else - echo " PASS agent transparency: ${AGENT_TRACES} agent LLM-call trace(s) visible in MLflow" - fi - ;; - partial) - echo " WARN agent transparency: ${AGENT} emits only token counts + tool names natively (${AGENT_TRACES} agent trace(s)); use 'dsagt start --enable-proxy' to capture full LLM-call payloads" - ;; - none) - echo " SKIP agent transparency: ${AGENT} emits no payload-bearing OTel traces (${AGENT_TRACES} agent trace(s)); use 'dsagt start --enable-proxy' to make agent LLM calls visible in MLflow" - ;; - *) - echo " WARN agent transparency: support tier unknown for ${AGENT}; ${AGENT_TRACES} agent trace(s)" - ;; -esac +check "agent traces recovered (${AGENT_TRACES})" "test '${AGENT_TRACES}' -gt 0" echo if [[ ${FAIL} -eq 0 ]]; then @@ -263,5 +275,8 @@ if [[ ${FAIL} -eq 0 ]]; then exit 0 else echo "[smoke] FAIL" + echo "[smoke] session logs kept: ${SESSION_LOG_1} ${SESSION_LOG_2}" + trap - EXIT + rm -f "${RENDERED_SCRIPT_1}" "${RENDERED_SCRIPT_2}" exit 1 fi diff --git a/tests/smoke_test/script.txt b/tests/smoke_test/script.txt index b906147..85cacfc 100644 --- a/tests/smoke_test/script.txt +++ b/tests/smoke_test/script.txt @@ -1,11 +1,15 @@ Ingest the docs in {{SMOKE_DIR}}/knowledge/ into a collection named knowledge. -I have a CSV utility called csvtool. Its reference is at {{SMOKE_DIR}}/knowledge/api_reference.md — please register the filter subcommand so we can reuse it later. Use an underscore in the name, not a hyphen. +I have a small CLI utility at {{SMOKE_DIR}}/greet.py — its reference doc is in the knowledge collection you just ingested. Register it in the code registry as "greet" so we can reuse it later. -Use the scan_directory tool from the registry to scan the {{SMOKE_DIR}}/data/ directory so I can see what's in it. +Run the registered greet code with the name argument "DSAGT" and the greeting argument "Ahoy". Use the exact shell command from its registered spec (the dsagt-run wrapper) so the execution is recorded in the trace archive. -Look at {{SMOKE_DIR}}/data/samples.csv and give me a short summary — columns, row count, any obvious quality issues. +Use the scan-directory code from the registry to scan the {{SMOKE_DIR}}/data/ directory so I can see what's in it. -Put this fact in explicit memory: samples.csv has null values in the status and timestamp columns. +Search the knowledge collection: what error code does greet report when the name argument is empty? Reply with the code verbatim. + +Reconstruct the pipeline of code executions run in this project so far. -What do you remember about the samples dataset? +Search the external skill catalog for a skill that helps with literature search, and install the best match into this project. + +Put this fact in explicit memory: samples.csv has null values in the status and timestamp columns. diff --git a/tests/smoke_test/script2.txt b/tests/smoke_test/script2.txt new file mode 100644 index 0000000..b1f6d3d --- /dev/null +++ b/tests/smoke_test/script2.txt @@ -0,0 +1,3 @@ +What do you remember about the samples dataset? + +Run the registered greet code once more, greeting "again". diff --git a/tests/test_chroma_metadata.py b/tests/test_chroma_metadata.py index c468262..8edd0a3 100644 --- a/tests/test_chroma_metadata.py +++ b/tests/test_chroma_metadata.py @@ -7,7 +7,6 @@ """ import json -from pathlib import Path from unittest.mock import MagicMock, patch import numpy as np @@ -15,16 +14,14 @@ from dsagt.knowledge import ( ChromaIndex, - CollectionRoute, - FAISSIndex, KnowledgeBase, ) - # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- + def fake_embed(texts: list[str]) -> np.ndarray: """Deterministic fake embeddings: hash-based so different texts get different (but reproducible) vectors.""" @@ -44,6 +41,7 @@ def fake_embed(texts: list[str]) -> np.ndarray: # ChromaIndex — metadata on add # --------------------------------------------------------------------------- + class TestChromaIndexMetadataAdd: def test_add_without_metadata(self): @@ -90,6 +88,7 @@ def test_add_metadata_incremental(self): # ChromaIndex — where filter on search # --------------------------------------------------------------------------- + class TestChromaIndexWhereSearch: @pytest.fixture @@ -127,25 +126,23 @@ def test_search_with_compound_where(self, indexed): query = np.ones(8, dtype=np.float32) query /= np.linalg.norm(query) scores, indices = indexed.search( - query, k=4, + query, + k=4, where={"$and": [{"tool_name": "megahit"}, {"session_id": "s2"}]}, ) assert set(indices.tolist()) == {3} def test_search_where_no_matches(self, indexed): - """search() with where that matches nothing returns empty or raises.""" + """search() with a where that matches nothing returns empty arrays.""" query = np.ones(8, dtype=np.float32) query /= np.linalg.norm(query) - # ChromaDB behavior varies by version: either returns empty results - # or raises when where matches nothing. Both are acceptable. - try: - scores, indices = indexed.search( - query, k=4, where={"tool_name": "nonexistent"}, - ) - assert len(scores) == 0 - assert len(indices) == 0 - except (ValueError, RuntimeError): - pass # ChromaDB raised on no-match — acceptable + scores, indices = indexed.search( + query, + k=4, + where={"tool_name": "nonexistent"}, + ) + assert len(scores) == 0 + assert len(indices) == 0 def test_search_empty_index(self): """search() on empty index returns empty arrays.""" @@ -160,6 +157,7 @@ def test_search_empty_index(self): # ChromaIndex — persistence with metadata # --------------------------------------------------------------------------- + class TestChromaIndexPersistence: def test_save_load_preserves_metadata(self, tmp_path): @@ -181,29 +179,11 @@ def test_save_load_preserves_metadata(self, tmp_path): assert results["metadatas"][1]["tool"] == "b" -# --------------------------------------------------------------------------- -# FAISSIndex — unchanged behavior -# --------------------------------------------------------------------------- - -class TestFAISSUnchanged: - - def test_faiss_add_search_unchanged(self): - """FAISSIndex add/search signatures are unchanged.""" - idx = FAISSIndex() - emb = np.random.randn(3, 8).astype(np.float32) - # Normalize for inner-product search - emb /= np.linalg.norm(emb, axis=1, keepdims=True) - idx.add(emb) - assert idx.size == 3 - - scores, indices = idx.search(emb[0], k=2) - assert len(scores) == 2 - - # --------------------------------------------------------------------------- # KnowledgeBase.search — where parameter # --------------------------------------------------------------------------- + class TestKnowledgeBaseSearchWhere: @pytest.fixture @@ -212,17 +192,12 @@ def kb_with_chroma(self, tmp_path): entries with metadata.""" index_dir = tmp_path / "kb" - chroma_route = CollectionRoute( - embedding_backend="api", - vector_db="chroma", - ) - - with patch("dsagt.knowledge._make_embedder") as mock_make: + with patch("dsagt.knowledge.Embedder.create") as mock_make: mock_embedder = MagicMock() mock_embedder.embed = fake_embed mock_make.return_value = mock_embedder - kb = KnowledgeBase(index_dir=index_dir, routes={"mem": chroma_route}) + kb = KnowledgeBase(index_dir=index_dir) # Manually build the collection with metadata kb.add_entries( @@ -239,7 +214,6 @@ def kb_with_chroma(self, tmp_path): {"tool_name": "fastp", "session_id": "s2"}, {"tool_name": "megahit", "session_id": "s2"}, ], - route=chroma_route, ) yield kb @@ -248,14 +222,20 @@ def kb_with_chroma(self, tmp_path): def test_search_without_where(self, kb_with_chroma): """Search without where returns results from all entries.""" results = kb_with_chroma.search( - "fastp quality", collection="mem", top_k=4, rerank=False, + "fastp quality", + collection="mem", + top_k=4, + rerank=False, ) assert len(results) > 0 def test_search_with_where_filters(self, kb_with_chroma): """Search with where restricts to matching metadata.""" results = kb_with_chroma.search( - "quality filtering", collection="mem", top_k=4, rerank=False, + "quality filtering", + collection="mem", + top_k=4, + rerank=False, where={"tool_name": "fastp"}, ) # All results should be fastp entries @@ -265,49 +245,27 @@ def test_search_with_where_filters(self, kb_with_chroma): def test_search_with_session_filter(self, kb_with_chroma): """Search with session_id filter returns only that session.""" results = kb_with_chroma.search( - "assembly", collection="mem", top_k=4, rerank=False, + "assembly", + collection="mem", + top_k=4, + rerank=False, where={"session_id": "s2"}, ) for r in results: assert r["chunk"]["metadata"]["session_id"] == "s2" - def test_search_where_on_faiss_ignored(self, tmp_path): - """where parameter is silently ignored on FAISS collections.""" - index_dir = tmp_path / "kb_faiss" - - with patch("dsagt.knowledge._make_embedder") as mock_make: - mock_embedder = MagicMock() - mock_embedder.embed = fake_embed - mock_make.return_value = mock_embedder - - kb = KnowledgeBase(index_dir=index_dir, default_index="faiss") - - # Use add_entries on a FAISS collection - kb.add_entries( - texts=["hello world", "goodbye world"], - collection="faiss_coll", - ) - - # where should not cause an error - results = kb.search( - "hello", collection="faiss_coll", top_k=2, rerank=False, - where={"some_field": "value"}, - ) - # Should still return results (where was ignored) - assert len(results) > 0 - kb.close() - # --------------------------------------------------------------------------- # KnowledgeBase.add_entries # --------------------------------------------------------------------------- + class TestAddEntries: @pytest.fixture def mock_kb(self, tmp_path): """KnowledgeBase with mocked embedder.""" - with patch("dsagt.knowledge._make_embedder") as mock_make: + with patch("dsagt.knowledge.Embedder.create") as mock_make: mock_embedder = MagicMock() mock_embedder.embed = fake_embed mock_make.return_value = mock_embedder @@ -341,24 +299,19 @@ def test_add_entries_empty_list(self, mock_kb): def test_add_entries_with_metadata(self, tmp_path): """add_entries stores metadata in chunks.jsonl.""" - chroma_route = CollectionRoute( - embedding_backend="api", vector_db="chroma", - ) - with patch("dsagt.knowledge._make_embedder") as mock_make: + with patch("dsagt.knowledge.Embedder.create") as mock_make: mock_embedder = MagicMock() mock_embedder.embed = fake_embed mock_make.return_value = mock_embedder kb = KnowledgeBase( index_dir=tmp_path / "kb", - routes={"mem": chroma_route}, ) kb.add_entries( texts=["ran fastp with Q20"], collection="mem", metadatas=[{"tool_name": "fastp", "session_id": "s1"}], - route=chroma_route, ) # Verify chunks.jsonl includes metadata @@ -378,7 +331,10 @@ def test_add_entries_searchable(self, mock_kb): collection="test_search", ) results = mock_kb.search( - "quick fox", collection="test_search", top_k=2, rerank=False, + "quick fox", + collection="test_search", + top_k=2, + rerank=False, ) assert len(results) > 0 @@ -406,20 +362,9 @@ def test_add_entries_return_embeddings_on(self, mock_kb): norms = np.linalg.norm(embeddings, axis=1) assert np.allclose(norms, 1.0, atol=1e-5) - def test_add_entries_return_embeddings_empty(self, mock_kb): - """Empty input with return_embeddings=True returns an empty array.""" - result = mock_kb.add_entries( - texts=[], collection="ret_empty", return_embeddings=True, - ) - assert "embeddings" in result - assert result["embeddings"].shape == (0,) - - def test_add_entries_with_route(self, tmp_path): - """add_entries with explicit ChromaDB route uses ChromaDB.""" - chroma_route = CollectionRoute( - embedding_backend="api", vector_db="chroma", - ) - with patch("dsagt.knowledge._make_embedder") as mock_make: + def test_add_entries_uses_chroma(self, tmp_path): + """add_entries writes a ChromaDB-backed collection (chroma_ids.json).""" + with patch("dsagt.knowledge.Embedder.create") as mock_make: mock_embedder = MagicMock() mock_embedder.embed = fake_embed mock_make.return_value = mock_embedder @@ -428,37 +373,11 @@ def test_add_entries_with_route(self, tmp_path): kb.add_entries( texts=["test entry"], collection="chroma_coll", - route=chroma_route, ) - # Verify it used ChromaDB (chroma_ids.json exists, not index.faiss) + # Verify it used ChromaDB (chroma_ids.json exists) coll_dir = tmp_path / "kb" / "chroma_coll" assert (coll_dir / "chroma_ids.json").exists() - assert not (coll_dir / "index.faiss").exists() - - kb.close() - - def test_add_entries_preserves_route(self, tmp_path): - """Route is persisted so subsequent searches use the right backend.""" - chroma_route = CollectionRoute( - embedding_backend="api", vector_db="chroma", - ) - with patch("dsagt.knowledge._make_embedder") as mock_make: - mock_embedder = MagicMock() - mock_embedder.embed = fake_embed - mock_make.return_value = mock_embedder - - kb = KnowledgeBase(index_dir=tmp_path / "kb") - kb.add_entries( - texts=["test entry"], - collection="routed", - route=chroma_route, - ) - - route_path = tmp_path / "kb" / "routed" / "route.json" - assert route_path.exists() - route_data = json.loads(route_path.read_text()) - assert route_data["vector_db"] == "chroma" kb.close() @@ -467,21 +386,18 @@ def test_add_entries_preserves_route(self, tmp_path): # End-to-end: add_entries + filtered search # --------------------------------------------------------------------------- + class TestAddEntriesFilteredSearch: """Integration test: add entries with metadata, then search with where.""" def test_filtered_search_returns_correct_subset(self, tmp_path): - chroma_route = CollectionRoute( - embedding_backend="api", vector_db="chroma", - ) - with patch("dsagt.knowledge._make_embedder") as mock_make: + with patch("dsagt.knowledge.Embedder.create") as mock_make: mock_embedder = MagicMock() mock_embedder.embed = fake_embed mock_make.return_value = mock_embedder kb = KnowledgeBase( index_dir=tmp_path / "kb", - routes={"executions": chroma_route}, ) kb.add_entries( @@ -498,13 +414,14 @@ def test_filtered_search_returns_correct_subset(self, tmp_path): {"tool_name": "fastp", "session_id": "s2", "return_code": 0}, {"tool_name": "megahit", "session_id": "s2", "return_code": 1}, ], - route=chroma_route, ) # Filter by tool_name fastp_results = kb.search( - "quality filtering", collection="executions", - top_k=4, rerank=False, + "quality filtering", + collection="executions", + top_k=4, + rerank=False, where={"tool_name": "fastp"}, ) for r in fastp_results: @@ -512,8 +429,10 @@ def test_filtered_search_returns_correct_subset(self, tmp_path): # Filter by session s1_results = kb.search( - "assembly", collection="executions", - top_k=4, rerank=False, + "assembly", + collection="executions", + top_k=4, + rerank=False, where={"session_id": "s1"}, ) for r in s1_results: @@ -521,8 +440,10 @@ def test_filtered_search_returns_correct_subset(self, tmp_path): # Compound filter failed = kb.search( - "megahit", collection="executions", - top_k=4, rerank=False, + "megahit", + collection="executions", + top_k=4, + rerank=False, where={"$and": [{"tool_name": "megahit"}, {"return_code": 1}]}, ) for r in failed: diff --git a/tests/test_config.py b/tests/test_config.py index b616597..292c8cf 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -23,7 +23,7 @@ static_agent_record, static_agent_files_present, ) -from dsagt.session import init_project, persist_agent_choice +from dsagt.session import init_project @pytest.fixture(autouse=True) @@ -36,10 +36,10 @@ def _use_tmp_registry(tmp_path): registry = {} def fake_load(): - # Auto-discover: any subdir of tmp_path with dsagt_config.yaml counts + # Auto-discover: any subdir of tmp_path with .dsagt/config.yaml counts discovered = dict(registry) for child in tmp_path.iterdir(): - if child.is_dir() and (child / "dsagt_config.yaml").exists(): + if child.is_dir() and (child / ".dsagt" / "config.yaml").exists(): discovered.setdefault(child.name, str(child)) return discovered @@ -50,12 +50,21 @@ def fake_save(reg): def fake_register(name, path): registry[name] = str(Path(path).resolve()) + # Isolate the shared KB to tmp_path and stub the asset build so init + # tests stay fast and offline (no embedding-model load, no git clone). + def _noop_ensure_assets(*_a, **_k): + return {"built": [], "skipped": []} + with patch("dsagt.session._load_registry", fake_load): with patch("dsagt.session._save_registry", fake_save): with patch("dsagt.session.register_project", fake_register): with patch("dsagt.session.DEFAULT_PROJECTS_BASE", tmp_path): - with patch("dsagt.session.DEFAULT_PROJECTS_BASE", tmp_path): - yield + with patch("dsagt.session.REGISTRY_DIR", tmp_path): + with patch( + "dsagt.commands.setup_core_kb.ensure_assets", + _noop_ensure_assets, + ): + yield # --------------------------------------------------------------------------- @@ -111,8 +120,8 @@ class TestLoadConfig: def _write_config(self, tmp_path, name, content: dict): pdir = tmp_path / name - pdir.mkdir(exist_ok=True) - (pdir / "dsagt_config.yaml").write_text( + (pdir / ".dsagt").mkdir(parents=True, exist_ok=True) + (pdir / ".dsagt" / "config.yaml").write_text( yaml.dump(content, default_flow_style=False) ) return name @@ -132,37 +141,36 @@ def test_loads_minimal_config(self, tmp_path): assert config["project"] == "myproject" assert config["agent"] == "goose" - # Ports are no longer defaulted in YAML — start_services picks them - # at runtime via socket.bind(("", 0)) and writes them to .runtime. - assert "proxy" not in config or "port" not in config.get("proxy", {}) + # Serverless: no mlflow block at all — the store is a sqlite path + # resolved from the project dir, nothing to pin in config. + assert "proxy" not in config + assert "mlflow" not in config # User-supplied keys win on the merge; DEFAULTS fills in missing # llm.* fields with ``${VAR}`` placeholders (resolved by # ``resolve_env_vars`` against the user's shell, or filtered by - # ``_real()`` if env unset). In BYOA the agent_env gate prevents - # env_overrides from acting on these — they only matter when - # --enable-proxy populates ``config["proxy"]["port"]``. + # ``_real()`` if env unset). assert config["llm"]["provider"] == "openai" # user value preserved - assert config["mlflow"]["backend"] == "sqlite" # non-llm defaults kept - def test_overrides_defaults(self, tmp_path): + def test_missing_project_raises(self, tmp_path): + name = self._write_config(tmp_path, "myproject", {"agent": "goose"}) + with pytest.raises(ValueError, match="project"): + load_config(name) + + def test_skills_block_backfilled_for_old_config(self, tmp_path): + """A config with no skills block gets the default genesis source. + + ``populate_native`` is a code default (in ``AgentSetup.setup_skills``), + not a config key — so it isn't backfilled here.""" name = self._write_config( tmp_path, "myproject", - { - "project": "myproject", - "agent": "claude", - "llm": {"provider": "openai"}, - "proxy": {"port": 9000}, # explicit override stays honored - }, + {"project": "myproject", "agent": "claude"}, ) - config = load_config(name) - assert config["proxy"]["port"] == 9000 - - def test_missing_project_raises(self, tmp_path): - name = self._write_config(tmp_path, "myproject", {"agent": "goose"}) - with pytest.raises(ValueError, match="project"): - load_config(name) + sources = config["skills"]["sources"] + assert sources[0]["name"] == "genesis" + assert "genesis-skills" in sources[0]["url"] + assert "populate_native" not in config["skills"] def test_missing_agent_raises(self, tmp_path): name = self._write_config(tmp_path, "myproject", {"project": "myproject"}) @@ -194,6 +202,18 @@ def test_project_dir_injected(self, tmp_path): assert config["project_dir"] == str(tmp_path / "myproject") +class TestSkillsDefaults: + + def test_defaults_has_skills(self): + from dsagt.session import DEFAULTS + + assert DEFAULTS["skills"]["sources"][0]["name"] == "genesis" + + def test_default_config_content_includes_skills(self): + body = yaml.safe_load(default_config_content("p", "claude")) + assert body["skills"]["sources"][0]["name"] == "genesis" + + # --------------------------------------------------------------------------- # Config: helpers # --------------------------------------------------------------------------- @@ -216,7 +236,7 @@ def test_unregistered_project_raises(self): class TestDefaultConfigContent: def test_roundtrips_as_valid_yaml(self): - content = default_config_content("test", "goose", mlflow_port=5000) + content = default_config_content("test", "goose") parsed = yaml.safe_load(content) assert parsed["project"] == "test" assert parsed["agent"] == "goose" @@ -224,15 +244,162 @@ def test_roundtrips_as_valid_yaml(self): def test_no_user_facing_llm_block(self): """BYOA: project YAML is internal-only — no llm: block, no ${VAR} placeholders. User credentials live in their shell.""" - content = default_config_content("test", "claude", mlflow_port=5000) + content = default_config_content("test", "claude") parsed = yaml.safe_load(content) assert "llm" not in parsed assert "${" not in content - def test_pinned_mlflow_port_lands_in_yaml(self): - content = default_config_content("test", "claude", mlflow_port=12345) + def test_no_mlflow_port_pinned(self): + """Serverless store — nothing to pin. The config carries no + mlflow block at all.""" + content = default_config_content("test", "claude") parsed = yaml.safe_load(content) - assert parsed["mlflow"]["port"] == 12345 + assert "mlflow" not in parsed + + +# --------------------------------------------------------------------------- +# CLI: unlisted `mlflow` alias for `traces` +# --------------------------------------------------------------------------- + + +class TestMlflowAlias: + """``dsagt mlflow`` routes to ``traces`` without ever reaching argparse, + so it stays out of --help; only the command slot is rewritten (a project + may itself be named "mlflow").""" + + def _capture_traces(self, monkeypatch): + from dsagt.commands import cli + + seen = {} + + def fake(args): + seen.update(project=args.project, port=args.port) + return 0 + + monkeypatch.setattr(cli, "_cmd_traces", fake) + return cli, seen + + def test_mlflow_routes_to_traces(self, monkeypatch): + cli, seen = self._capture_traces(monkeypatch) + assert cli.main(["mlflow", "myproj", "--port", "5001"]) == 0 + assert seen == {"project": "myproj", "port": 5001} + + def test_project_named_mlflow_is_not_rewritten(self, monkeypatch): + cli, seen = self._capture_traces(monkeypatch) + assert cli.main(["traces", "mlflow"]) == 0 + assert seen == {"project": "mlflow", "port": 5000} + + def test_alias_absent_from_help_command_list(self, capsys): + from dsagt.commands import cli + + with pytest.raises(SystemExit): + cli.main(["--help"]) + out = capsys.readouterr().out + choices = out[out.index("{") + 1 : out.index("}")] + assert "mlflow" not in choices + + +# --------------------------------------------------------------------------- +# CLI: init choice resolution (_collect_settings) +# --------------------------------------------------------------------------- + + +class TestCollectSettings: + """``_collect_settings`` resolves the init choices into the config blocks. + + Two selection questions (KB collections, skill sources) + agent; the + bundled ``tools`` collection is always provisioned and not a choice. + """ + + def test_interactive_menus(self): + """Interactive path drives questionary select/checkbox; genesis is the + default-checked skill source, collections default to none, tools is + always in the provisioning set.""" + import types + import questionary + from unittest.mock import patch + from dsagt.commands import cli + + def fake_select(message, choices, default, **kw): + return _Ask(default) + + def fake_checkbox(message, choices, **kw): + return _Ask([c.value for c in choices if c.checked]) + + class _Ask: + def __init__(self, ret): + self.ret = ret + + def ask(self): + return self.ret + + args = types.SimpleNamespace(agent=None, include=None, exclude=None) + with ( + patch.object(questionary, "select", fake_select), + patch.object(questionary, "checkbox", fake_checkbox), + patch.object(cli, "_confirm", lambda *a, **k: False), # episodic off + ): + s = cli._collect_settings(args, interactive=True, existing={}, pdir=None) + + assert s["agent"] == "claude" + assert s["knowledge"] == {"collections": []} + assert [src["name"] for src in s["skills"]["sources"]] == ["genesis"] + assert s["assets"] == ["codes", "genesis"] # tools always provisioned + assert s["episodic"] is None # opt-in, off by default + + def test_interactive_episodic_enabled(self): + """Enabling episodic at the prompt yields the opt-in block.""" + import types + import questionary + from unittest.mock import patch + from dsagt.commands import cli + + class _Ask: + def __init__(self, ret): + self.ret = ret + + def ask(self): + return self.ret + + args = types.SimpleNamespace(agent=None, include=None, exclude=None) + with ( + patch.object(questionary, "select", lambda *a, **k: _Ask("claude")), + patch.object(questionary, "checkbox", lambda *a, **k: _Ask([])), + patch.object(cli, "_confirm", lambda *a, **k: True), + ): + s = cli._collect_settings(args, interactive=True, existing={}, pdir=None) + + assert s["episodic"] == {"enabled": True} + + def test_non_interactive_splits_assets(self): + """No-TTY path splits --include into collections vs skill sources; + tools is always provisioned. Episodic omitted → None.""" + import types + from dsagt.commands import cli + + args = types.SimpleNamespace( + agent="goose", include=["codes", "nemo_curator", "anthropic"], exclude=None + ) + s = cli._collect_settings(args, interactive=False, existing={}, pdir=None) + assert s["agent"] == "goose" + assert s["knowledge"]["collections"] == ["nemo_curator"] + assert [src["name"] for src in s["skills"]["sources"]] == ["anthropic"] + assert s["assets"] == ["codes", "nemo_curator", "anthropic"] + assert s["episodic"] is None + + def test_non_interactive_episodic_flag(self): + """--episodic builds the opt-in block on the no-TTY path.""" + import types + from dsagt.commands import cli + + args = types.SimpleNamespace( + agent="goose", + include=["codes"], + exclude=None, + episodic=True, + ) + s = cli._collect_settings(args, interactive=False, existing={}, pdir=None) + assert s["episodic"] == {"enabled": True} # --------------------------------------------------------------------------- @@ -243,18 +410,22 @@ def test_pinned_mlflow_port_lands_in_yaml(self): class TestInitProject: def test_creates_directory_structure(self): - pdir, _port = init_project("myproj", "goose") + pdir = init_project("myproj", "goose") assert pdir.exists() - assert (pdir / "dsagt_config.yaml").exists() + assert (pdir / ".dsagt" / "config.yaml").exists() assert (pdir / "trace_archive").is_dir() - assert (pdir / "mlflow").is_dir() assert (pdir / "skills").is_dir() assert (pdir / "kb_index").is_dir() assert (pdir / ".dsagt").is_dir() - # `tools/` is intentionally NOT created by init_project — ToolRegistry - # creates it on first server startup so bundled tools get copied in. - assert not (pdir / "tools").exists() + # Bundled codes are copied into codes/ at init — every available + # code in one place, one format (skill-standard dirs). + assert (pdir / "codes" / "scan-directory" / "SKILL.md").exists() + assert (pdir / "codes" / "scan-directory" / "scripts").is_dir() + # Serverless: no MLflow store is pre-created; ``mlflow.db`` is + # written lazily by the MLflow client on first span. + assert not (pdir / "mlflow.db").exists() + assert not (pdir / "mlflow").exists() def test_config_is_valid(self): init_project("myproj", "claude") @@ -262,30 +433,62 @@ def test_config_is_valid(self): assert config["project"] == "myproj" assert config["agent"] == "claude" - def test_returns_pdir_and_port(self): - """BYOA: init_project returns (pdir, mlflow_port).""" - pdir, port = init_project("myproj", "goose") - assert isinstance(port, int) and port > 0 - # Port is persisted in the YAML so MCP servers can read it. + def test_returns_pdir(self): + """Serverless: init_project returns just the project dir — no port.""" + pdir = init_project("myproj", "goose") + assert pdir.exists() config = load_config("myproj") - assert config["mlflow"]["port"] == port - - def test_explicit_mlflow_port_honored(self): - from dsagt.session import pick_free_port - - chosen = pick_free_port() - _pdir, port = init_project("myproj", "goose", mlflow_port=chosen) - assert port == chosen + assert "mlflow" not in config + + def test_exclude_all_creates_empty_kb_without_building(self): + """``--exclude all`` provisions a valid project with an empty KB and + never attempts an asset build.""" + from dsagt.commands import setup_core_kb + + with patch.object(setup_core_kb, "ensure_assets") as mock_ensure: + pdir = init_project("myproj", "goose", exclude=["all"]) + mock_ensure.assert_not_called() + kb_index = pdir / "kb_index" + assert kb_index.is_dir() + assert not any(kb_index.iterdir()) + + def test_default_init_provisions_default_asset_set(self): + """Default init builds exactly the default asset set into the shared + cache (here stubbed) — tools + genesis, nothing heavier.""" + from dsagt.commands import setup_core_kb + + with patch.object( + setup_core_kb, "ensure_assets", return_value={"built": [], "skipped": []} + ) as mock_ensure: + init_project("myproj", "goose") + mock_ensure.assert_called_once() + requested = mock_ensure.call_args.args[0] + assert requested == ["codes", "genesis"] - def test_duplicate_raises(self): + def test_reinit_is_idempotent_update(self): + """``dsagt init`` is re-runnable: a second init on the same project + updates settings in place rather than raising.""" init_project("myproj", "goose") - with pytest.raises(FileExistsError): - init_project("myproj", "goose") + init_project("myproj", "claude") # re-init, switch agent + config = load_config("myproj") + assert config["agent"] == "claude" def test_invalid_agent_raises(self): with pytest.raises(ValueError): init_project("myproj", "invalid-agent") + def test_episodic_block_round_trips_through_config(self): + """init_project writes the opted-in episodic block; load_config reads it + back. A project without it backfills ``enabled: False`` from DEFAULTS.""" + epi = {"enabled": True} + init_project("withmem", "goose", exclude=["all"], episodic=epi) + cfg = load_config("withmem") + assert cfg["episodic"]["enabled"] is True + + init_project("nomem", "goose", exclude=["all"]) + cfg2 = load_config("nomem") + assert cfg2["episodic"]["enabled"] is False # backfilled default + def test_static_record_written_eagerly(self, tmp_path): """BYOA flow writes static + dynamic records at init time so the user can edit instructions and inspect the MCP config artifact @@ -293,142 +496,68 @@ def test_static_record_written_eagerly(self, tmp_path): the CLI command, not init_project itself — but the agent dir layout exists post-init. """ - pdir, _ = init_project("myproj", "claude") + init_project("myproj", "claude") config = load_config("myproj") assert config["agent"] == "claude" # --------------------------------------------------------------------------- -# persist_agent_choice: first-start agent selection writes back to YAML +# Session state (.dsagt/state.yaml) — owned by the MCP server # --------------------------------------------------------------------------- -class TestPersistAgentChoice: - """``persist_agent_choice`` is called from ``_cmd_start`` when the - YAML doesn't already pin an agent — covers the case where the user - deferred ``--agent`` from init time and supplied it on first start. +class TestSessionState: + """The MCP server mints sessions into ``.dsagt/state.yaml`` (a monotonic + per-project counter) and ``dsagt-run`` reads the current tag from there. """ - def test_overwrites_existing_field(self, tmp_path): - init_project("myproj", "goose") - persist_agent_choice("myproj", "claude") - - config = load_config("myproj") - assert config["agent"] == "claude" - - def test_invalid_agent_raises(self, tmp_path): - init_project("myproj", "goose") - with pytest.raises(ValueError): - persist_agent_choice("myproj", "not-an-agent") - - -# --------------------------------------------------------------------------- -# pick_free_port: kernel-assigned via socket.bind(("", 0)) -# --------------------------------------------------------------------------- - - -class TestPickFreePort: - - def test_returns_a_usable_port(self): - """The kernel hands back a positive port number we can bind.""" - from dsagt.session import pick_free_port - import socket as _socket - - port = pick_free_port() - assert isinstance(port, int) - assert 1024 <= port <= 65535 - # And we can actually bind it (no leak, no half-stuck socket). - with _socket.socket(_socket.AF_INET, _socket.SOCK_STREAM) as s: - s.bind(("", port)) - - def test_consecutive_calls_pick_different_ports(self): - """No memoization, no preferred-port bias — consecutive calls just - reflect whatever the kernel hands out (almost always different on - a quiet box, very rarely the same; this test tolerates that).""" - from dsagt.session import pick_free_port - - ports = {pick_free_port() for _ in range(10)} - # Not strictly guaranteed all 10 differ, but >1 distinct is expected. - assert len(ports) > 1 - - -# --------------------------------------------------------------------------- -# reap_runtime: SIGTERM PIDs in /.runtime, grace, SIGKILL stragglers -# --------------------------------------------------------------------------- - - -class TestReapRuntime: - - def test_no_runtime_file_returns_empty(self, tmp_path): - from dsagt.session import reap_runtime - - assert reap_runtime(tmp_path / ".runtime") == [] - - def test_skips_pids_whose_cmdline_doesnt_match(self, tmp_path, monkeypatch): - """PID-recycling guard: cmdline doesn't name our service → leave alone.""" - from dsagt import session - - runtime = tmp_path / ".runtime" - runtime.write_text( - json.dumps( - { - "pids": {"mlflow": 99999, "proxy": 99998}, - "ports": {"mlflow": 12345, "proxy": 12346}, - } - ) - ) - # Pretend both PIDs are recycled — cmdline doesn't contain our name. - monkeypatch.setattr(session, "_process_command", lambda pid: "/bin/zsh") - killed = session.reap_runtime(runtime) - assert killed == [] - # File still gets cleaned up so a stale .runtime doesn't linger. - assert not runtime.exists() - - def test_signals_pids_whose_cmdline_matches(self, tmp_path, monkeypatch): - """When the cmdline still names our service, SIGTERM is sent.""" - from dsagt import session - - runtime = tmp_path / ".runtime" - runtime.write_text( - json.dumps( - { - "pids": {"mlflow": 11111, "proxy": 22222}, - "ports": {"mlflow": 12345, "proxy": 12346}, - } - ) + def test_append_session_increments(self, tmp_path): + from dsagt.session import append_session, current_session + + pdir = tmp_path / "proj" + pdir.mkdir() + e1 = append_session(pdir) + e2 = append_session(pdir) + assert e1["id"] == 1 + assert e2["id"] == 2 + assert e1["started_at"].endswith("Z") + assert current_session(pdir)["id"] == 2 + + def test_session_tag_shape(self): + from dsagt.session import session_tag + + assert session_tag("myproj", 3) == "myproj-3" + + def test_current_session_tag_from_state(self, tmp_path): + from dsagt.session import ( + append_session, + current_session_tag, + write_config_file, + build_config, ) - monkeypatch.setattr( - session, - "_process_command", - lambda pid: f"python -m {'mlflow' if pid == 11111 else 'dsagt.commands.proxy_server'}", - ) - signals_sent: list[tuple[int, int]] = [] - - def fake_killpg(pgid, sig): - signals_sent.append((pgid, sig)) - monkeypatch.setattr(session.os, "killpg", fake_killpg) - monkeypatch.setattr(session.os, "getpgid", lambda pid: pid) - # Make grace-wait short and the process appear to exit immediately. - monkeypatch.setattr(session, "_STOP_GRACE_SECONDS", 0) + pdir = tmp_path / "proj" + pdir.mkdir() + write_config_file(pdir, build_config("proj", "claude")) + assert current_session_tag(pdir, "proj") is None # no session yet + append_session(pdir) + assert current_session_tag(pdir, "proj") == "proj-1" - killed = session.reap_runtime(runtime) - assert len(killed) == 2 - # Both PIDs got SIGTERM. - import signal as _signal + def test_update_cursor_roundtrip(self, tmp_path): + from dsagt.session import read_state, update_cursor - sigterm_pgids = {pgid for pgid, sig in signals_sent if sig == _signal.SIGTERM} - assert sigterm_pgids == {11111, 22222} + pdir = tmp_path / "proj" + pdir.mkdir() + update_cursor(pdir, tool_use_indexed_through="2026-01-01T00:00:00Z") + cur = read_state(pdir)["memory_cursor"] + assert cur["tool_use_indexed_through"] == "2026-01-01T00:00:00Z" # --------------------------------------------------------------------------- # Per-agent record writers: static_agent_record + dynamic_agent_record # # Each test runs both writers against a project init'd with --agent, -# mirroring what `dsagt start` does in production. Cline's dynamic -# writer would shell out to `cline auth` here, so cline gets exercised -# by ``test_cline_mcp_config_shape`` (mocked subprocess) instead of -# the full dynamic writer. +# mirroring what `dsagt start` does in production. # --------------------------------------------------------------------------- @@ -439,13 +568,8 @@ def _init_and_load(self, agent): return load_config("testproj") def _write_both(self, config, working_dir): - """Run static then dynamic — what dsagt start does, minus services. - - ``start_services`` populates ``config["mlflow"]["port"]`` in - production; we stub it here so ``agent_env`` (which reads it to - build the OTLP endpoint) doesn't KeyError. - """ - config.setdefault("mlflow", {})["port"] = 5001 + """Run static then dynamic — what dsagt start does. Serverless: + no port to populate; the store resolves from the project dir.""" static_agent_record(config, config["agent"], working_dir) env = agent_env(config) dynamic_agent_record(config, env, working_dir) @@ -460,8 +584,8 @@ def test_claude_writes_mcp_json(self, tmp_path): mcp_path = working_dir / ".mcp.json" assert mcp_path.exists() mcp = json.loads(mcp_path.read_text()) - assert "dsagt-registry" in mcp["mcpServers"] - assert "dsagt-knowledge" in mcp["mcpServers"] + assert set(mcp["mcpServers"]) == {"dsagt"} + assert mcp["mcpServers"]["dsagt"]["args"] == ["run", "dsagt-server"] assert (working_dir / "CLAUDE.md").exists() # BYOA: .dsagt_env is no longer written; user manages shell env. assert not (working_dir / ".dsagt_env").exists() @@ -476,39 +600,62 @@ def test_goose_writes_goose_yaml(self, tmp_path): goose_path = working_dir / "goose.yaml" assert goose_path.exists() goose = yaml.safe_load(goose_path.read_text()) - assert "registry" in goose["extensions"] - assert "knowledge" in goose["extensions"] + assert set(goose["extensions"]) == {"dsagt"} + assert goose["extensions"]["dsagt"]["cmd"] == "uv run dsagt-server" assert (working_dir / ".goosehints").exists() - def test_roo_writes_static_and_dynamic(self, tmp_path): - config = self._init_and_load("roo") + def test_cline_writes_project_mcp_settings(self, tmp_path): + """Cline's dynamic writer hand-writes the per-project MCP settings + file (no cline binary needed); runtime_env points cline at it via + CLINE_MCP_SETTINGS_PATH, leaving global auth + settings untouched.""" + import json as _json + + from dsagt.agents import AGENTS + + config = self._init_and_load("cline") working_dir = tmp_path / "workdir" working_dir.mkdir() - self._write_both(config, working_dir) + static_agent_record(config, config["agent"], working_dir) + assert (working_dir / ".clinerules" / "dsagt_instructions.md").exists() + + dynamic_agent_record(config, env={}, working_dir=working_dir) + settings_path = working_dir / ".cline-data" / "cline_mcp_settings.json" + settings = _json.loads(settings_path.read_text()) + transport = settings["mcpServers"]["dsagt"]["transport"] + assert transport["type"] == "stdio" + assert [transport["command"], *transport["args"]] == [ + "uv", + "run", + "dsagt-server", + ] + assert "DSAGT_PROJECT_DIR" in transport["env"] + + env = AGENTS["cline"]().runtime_env(config) + assert env["CLINE_MCP_SETTINGS_PATH"].endswith( + ".cline-data/cline_mcp_settings.json" + ) + assert "CLINE_DIR" not in env + + def test_cline_dynamic_preserves_user_mcp_entries(self, tmp_path): + """Re-running the writer keeps non-dsagt servers the user added.""" + import json as _json - assert (working_dir / ".roo").is_dir() - assert (working_dir / ".roomodes").exists() - assert (working_dir / ".roo" / "mcp.json").exists() - # BYOA: project routing (project name, project_dir, mlflow port) - # comes from dsagt_config.yaml via cwd-walk; the MCP env block - # only carries EMBEDDING_* settings. - mcp = json.loads((working_dir / ".roo" / "mcp.json").read_text()) - assert "EMBEDDING_BACKEND" in mcp["mcpServers"]["dsagt-registry"]["env"] - - def test_cline_writes_static_only_in_split_test(self, tmp_path): - # Cline's dynamic writer shells out to `cline auth` and `cline mcp - # add`, which would fail without cline installed. Test only the - # static half here; ``test_cline_mcp_config_shape`` covers the - # dynamic half with mocked subprocess. config = self._init_and_load("cline") working_dir = tmp_path / "workdir" working_dir.mkdir() - static_agent_record(config, config["agent"], working_dir) + dynamic_agent_record(config, env={}, working_dir=working_dir) + settings_path = working_dir / ".cline-data" / "cline_mcp_settings.json" + settings = _json.loads(settings_path.read_text()) + settings["mcpServers"]["mytool"] = { + "transport": {"type": "stdio", "command": "mytool", "args": []} + } + settings_path.write_text(_json.dumps(settings)) - instructions = working_dir / ".clinerules" / "dsagt_instructions.md" - assert instructions.exists() + dynamic_agent_record(config, env={}, working_dir=working_dir) + settings = _json.loads(settings_path.read_text()) + assert set(settings["mcpServers"]) == {"dsagt", "mytool"} assert (working_dir / ".cline-data").is_dir() def test_codex_writes_static_and_dynamic(self, tmp_path): @@ -521,7 +668,7 @@ def test_codex_writes_static_and_dynamic(self, tmp_path): assert (working_dir / "AGENTS.md").exists() assert (working_dir / ".codex-data").is_dir() toml = (working_dir / ".codex-data" / "config.toml").read_text() - assert "[mcp_servers.dsagt-registry.env]" in toml + assert "[mcp_servers.dsagt.env]" in toml # Project routing comes from dsagt_config.yaml via cwd-walk; the # MCP env block only carries EMBEDDING_* settings. assert "EMBEDDING_BACKEND" in toml @@ -555,11 +702,11 @@ def test_static_files_present_check(self, tmp_path): assert static_agent_files_present("codex", working_dir) def test_codex_config_toml_shape(self, tmp_path): - """``_render_codex_config`` emits ``[mcp_servers.*]`` sections - plus an ``[otel]`` block (opting codex's partial-OTel emission - into MLflow via the OTLP exporter). No top-level keys — those - come from the user's ``~/.codex/config.toml`` which - ``write_dynamic`` copies as a base. + """``_render_codex_config`` emits only ``[mcp_servers.*]`` sections. + No ``[otel]`` block — DSAGT no longer forces codex's native + telemetry (nor the ``log_user_prompt`` privacy override). No + top-level keys — those come from the user's ``~/.codex/config.toml`` + which ``write_dynamic`` copies as a base. """ from dsagt.agents import _render_codex_config @@ -569,14 +716,12 @@ def test_codex_config_toml_shape(self, tmp_path): } toml = _render_codex_config(mcp_env) - assert "[mcp_servers.dsagt-registry]" in toml - assert "[mcp_servers.dsagt-knowledge]" in toml - assert "[mcp_servers.dsagt-registry.env]" in toml + assert "[mcp_servers.dsagt]" in toml + assert "[mcp_servers.dsagt.env]" in toml assert 'MLFLOW_TRACKING_URI = "http://localhost:5001"' in toml - # OTel opt-in: enables OTLP-HTTP export + un-redacts user prompt. - assert "[otel]" in toml - assert 'trace_exporter = "otlp-http"' in toml - assert "log_user_prompt = true" in toml + # No forced telemetry / privacy override. + assert "[otel]" not in toml + assert "log_user_prompt" not in toml # No top-level approval/sandbox keys — those come from user's # config.toml or the codex exec CLI flag. assert "approval_policy" not in toml @@ -606,10 +751,10 @@ def test_opencode_config_json_shape(self): parsed = json.loads(body) assert parsed["$schema"] == "https://opencode.ai/config.json" - assert set(parsed["mcp"]) == {"dsagt-registry", "dsagt-knowledge"} - reg = parsed["mcp"]["dsagt-registry"] + assert set(parsed["mcp"]) == {"dsagt"} + reg = parsed["mcp"]["dsagt"] assert reg["type"] == "local" - assert reg["command"] == ["uv", "run", "dsagt-registry-server"] + assert reg["command"] == ["uv", "run", "dsagt-server"] assert reg["environment"]["DSAGT_PROJECT_DIR"] == "/proj" # Provider block uses {env:VAR} reference, never the resolved value. assert ( @@ -681,19 +826,20 @@ def test_mcp_servers_dict_shape(self): } mcp = _build_mcp_servers_dict(env_block) - assert set(mcp["mcpServers"]) == {"dsagt-registry", "dsagt-knowledge"} - assert mcp["mcpServers"]["dsagt-knowledge"]["disabled"] is False + assert set(mcp["mcpServers"]) == {"dsagt"} + assert mcp["mcpServers"]["dsagt"]["disabled"] is False # Env block plumbs through so the MCP server children have what they need. assert ( - mcp["mcpServers"]["dsagt-registry"]["env"]["MLFLOW_TRACKING_URI"] + mcp["mcpServers"]["dsagt"]["env"]["MLFLOW_TRACKING_URI"] == "http://localhost:5001" ) - def test_mcp_config_omits_redundant_dsagt_env(self, tmp_path): - """Single source of truth: project routing (project, project_dir, - mlflow port) lives in dsagt_config.yaml and is read via cwd-walk - by every dsagt service. The MCP env block must NOT duplicate - those values — that's the contract being asserted here.""" + def test_mcp_config_carries_routing_env(self, tmp_path): + """Benign routing in the MCP env block: agents that don't inherit + the parent's shell env into their MCP children (codex / cline + — and claude's block is robust against shells that don't + export it) need project name + dir and the serverless + ``MLFLOW_TRACKING_URI`` baked in. No credentials, no OTel.""" config = self._init_and_load("claude") working_dir = tmp_path / "workdir" working_dir.mkdir() @@ -701,11 +847,13 @@ def test_mcp_config_omits_redundant_dsagt_env(self, tmp_path): self._write_both(config, working_dir) mcp = json.loads((working_dir / ".mcp.json").read_text()) - for server in ("dsagt-registry", "dsagt-knowledge"): - env = mcp["mcpServers"][server].get("env", {}) - assert "DSAGT_PROJECT" not in env - assert "DSAGT_PROJECT_DIR" not in env - assert "MLFLOW_TRACKING_URI" not in env + env = mcp["mcpServers"]["dsagt"].get("env", {}) + assert env["DSAGT_PROJECT"] == config["project"] + assert env["DSAGT_PROJECT_DIR"] == config["project_dir"] + assert env["MLFLOW_TRACKING_URI"].startswith("sqlite:///") + # No credentials or OTel routing leak into the MCP config. + assert "ANTHROPIC_API_KEY" not in env + assert "OTEL_EXPORTER_OTLP_ENDPOINT" not in env # --------------------------------------------------------------------------- @@ -714,7 +862,7 @@ def test_mcp_config_omits_redundant_dsagt_env(self, tmp_path): class TestResolveRecordsDirProjectAware: - """``_resolve_records_dir`` reads the project's ``dsagt_config.yaml`` + """``_resolve_records_dir`` reads the project's ``.dsagt/config.yaml`` from cwd as the single source of truth — no env-var chain.""" def test_explicit_overrides_cwd(self, tmp_path): @@ -724,12 +872,13 @@ def test_explicit_overrides_cwd(self, tmp_path): assert result == Path("/custom") def test_cwd_with_config(self, tmp_path, monkeypatch): - """When cwd contains ``dsagt_config.yaml``, records dir is + """When cwd contains ``.dsagt/config.yaml``, records dir is ``/trace_archive``. Even if env vars are set to point elsewhere, the config-in-cwd rule wins (env is ignored).""" from dsagt.provenance import _resolve_records_dir - (tmp_path / "dsagt_config.yaml").write_text("project: t\n") + (tmp_path / ".dsagt").mkdir() + (tmp_path / ".dsagt" / "config.yaml").write_text("project: t\n") monkeypatch.chdir(tmp_path) # Stale env vars must not be consulted. monkeypatch.setenv("DSAGT_PROJECT_DIR", "/stale/proj/dir") @@ -762,488 +911,54 @@ def test_dsagt_vars_set(self): assert env["DSAGT_AGENT"] == "claude" assert env["DSAGT_PROJECT_DIR"] == "/proj" - def test_claude_byoa_skips_otel_routing(self): - """Claude in BYOA mode (no proxy) uses ``mlflow autolog claude`` - for agent-side traces — its Stop hook produces richer - transcript-based traces than native OTel. agent_env should NOT - export OTEL_EXPORTER_OTLP_* vars for Claude in BYOA, to avoid - duplicate (and inferior) trace shapes. + def test_no_otel_routing_for_any_agent(self, monkeypatch): + """DSAGT forces no native OTel emission — agent traces are + recovered post-hoc from the on-disk transcript. ``agent_env`` + sets ``MLFLOW_TRACKING_URI`` (for MCP-server / MLflow-client + logging) but never the OTLP routing env, for any agent. """ from dsagt.agents import agent_env - env = agent_env(self._make_config("claude")) - assert env["MLFLOW_TRACKING_URI"] == "http://localhost:5001" - assert "OTEL_EXPORTER_OTLP_ENDPOINT" not in env - assert "OTEL_EXPORTER_OTLP_HEADERS" not in env - assert "OTEL_RESOURCE_ATTRIBUTES" not in env - - def test_non_claude_byoa_keeps_otel_routing(self): - """Goose has no ``mlflow autolog goose`` analog — it still gets - OTel routing so its native OTel emission lands in MLflow. - """ - from dsagt.agents import agent_env - - env = agent_env(self._make_config("goose")) - assert env["MLFLOW_TRACKING_URI"] == "http://localhost:5001" - assert env["OTEL_EXPORTER_OTLP_ENDPOINT"] == "http://localhost:5001/v1/traces" - assert "service.name=goose" in env["OTEL_RESOURCE_ATTRIBUTES"] - - def test_claude_telemetry_verbosity_flags(self): - """Claude Code needs OTEL_LOG_TOOL_DETAILS + OTEL_LOG_USER_PROMPTS - to emit unredacted tool_use payloads + user prompts memory - extraction depends on. OTEL_LOG_RAW_API_BODIES is intentionally - NOT in the static env block — its value is a per-project file - path rendered dynamically by ``_cmd_mlflow`` (see claude.py for - why ``=1`` mode would lose bodies to MLflow's missing /v1/logs - endpoint).""" - from dsagt.agents import agent_env - - env = agent_env(self._make_config("claude")) - assert env["CLAUDE_CODE_ENABLE_TELEMETRY"] == "1" - assert env["OTEL_LOG_TOOL_DETAILS"] == "1" - assert env["OTEL_LOG_USER_PROMPTS"] == "1" - assert "OTEL_LOG_RAW_API_BODIES" not in env - - -@pytest.mark.skip( - reason=( - "Old-code-shape env_overrides: assertions describe the pre-Phase-1 " - "design where env_overrides translated llm.* into ANTHROPIC_*/OPENAI_* " - "credential env vars across the board. Phase 2's env_overrides is " - "narrower — it pins the agent's MODEL env var only; provider creds " - "and base URLs are set by proxy_env_overrides (sentinel + proxy URL) " - "or live in agent config files (cline auth state, codex config.toml, " - "opencode.json). See ``TestPhase2EnvOverrides`` below for the new " - "contract. Keeping these tests around for the LoC reference; " - "delete when Phase 2 is stable." - ) -) -class TestProviderEnvInjection: - """``llm.{provider, model, api_key, base_url}`` from dsagt_config must - flow into the agent's expected env-var names — otherwise the agent - falls back to its own user-config and the project YAML's settings - don't actually drive the agent. This was a real bug for goose.""" - - def _config_with_llm(self, **llm) -> dict: - return { - "project": "test", - "agent": "goose", - "project_dir": "/proj", - "mlflow": {"port": 5001}, - "llm": llm, - "embedding": {}, - } - - def test_openai_provider_sets_openai_env(self): - from dsagt.agents import agent_env - - env = agent_env( - self._config_with_llm( - provider="openai", - model="gpt-4o", - api_key="sk-real", - base_url="https://gateway.example.com", - ) - ) - assert env["OPENAI_API_KEY"] == "sk-real" - assert env["OPENAI_BASE_URL"] == "https://gateway.example.com" - assert env["GOOSE_PROVIDER"] == "openai" - assert env["GOOSE_MODEL"] == "gpt-4o" - # Goose's Rust openai client reads OPENAI_HOST, not OPENAI_BASE_URL. - # Without this, goose silently hits api.openai.com regardless of - # the project's gateway. - assert env["OPENAI_HOST"] == "https://gateway.example.com" - - def test_goose_anthropic_sets_anthropic_host(self): - """Same HOST-vs-BASE_URL issue for goose's anthropic provider.""" - from dsagt.agents import agent_env - - env = agent_env( - self._config_with_llm( - provider="anthropic", - model="claude-sonnet", - api_key="anthropic-key", - base_url="https://api.anthropic.com", - ) - ) - assert env["ANTHROPIC_HOST"] == "https://api.anthropic.com" - - @pytest.mark.skip(reason="proxy mode deferred to Phase 2") - def test_goose_proxy_sets_host_at_proxy(self): - """Proxy mode: OPENAI_HOST / ANTHROPIC_HOST must point at the proxy - too, not just the standard BASE_URL slots.""" - from dsagt.agents import agent_env - - config = self._config_with_llm( - provider="openai", - api_key="sk-real", - base_url="https://upstream.example.com", - ) - config["proxy"] = {"port": 9999} - env = agent_env(config) - assert env["OPENAI_HOST"] == "http://localhost:9999" - assert env["ANTHROPIC_HOST"] == "http://localhost:9999" - - def test_openai_like_provider_treated_as_openai(self): - """openai_like (lab gateways speaking OpenAI wire protocol) maps - to OPENAI_* env vars too — that's what every OpenAI-compat client - reads.""" - from dsagt.agents import agent_env - - env = agent_env( - self._config_with_llm( - provider="openai_like", - api_key="sk-lab", - base_url="https://lab.example.com", - ) - ) - assert env["OPENAI_API_KEY"] == "sk-lab" - assert env["OPENAI_BASE_URL"] == "https://lab.example.com" - - def test_anthropic_provider_sets_anthropic_env(self): - from dsagt.agents import agent_env - - env = agent_env( - self._config_with_llm( - provider="anthropic", - model="claude-sonnet-4-5", - api_key="anthropic-key", - base_url="https://api.anthropic.com", - ) - ) - assert env["ANTHROPIC_API_KEY"] == "anthropic-key" - assert env["ANTHROPIC_BASE_URL"] == "https://api.anthropic.com" - # ANTHROPIC_MODEL pins the model so claude code doesn't fall back - # to its built-in default (which won't exist on lab gateways). - assert env["ANTHROPIC_MODEL"] == "claude-sonnet-4-5" - - def test_cline_gets_provider_env(self): - """Cline's CLI reads OPENAI_*/ANTHROPIC_* — must flow from llm.*.""" - from dsagt.agents import agent_env - - cfg = self._config_with_llm( - provider="openai", - api_key="sk-cline", - base_url="https://gateway.example.com", - ) - cfg["agent"] = "cline" - env = agent_env(cfg) - assert env["OPENAI_API_KEY"] == "sk-cline" - assert env["OPENAI_BASE_URL"] == "https://gateway.example.com" - - def test_roo_gets_provider_env(self): - """Roo's CLI reads ANTHROPIC_* — must flow from llm.*.""" - from dsagt.agents import agent_env - - cfg = self._config_with_llm( - provider="anthropic", - model="claude-haiku-4-5", - api_key="anthropic-roo-key", - base_url="https://api.anthropic.com", - ) - cfg["agent"] = "roo" - env = agent_env(cfg) - assert env["ANTHROPIC_API_KEY"] == "anthropic-roo-key" - assert env["ANTHROPIC_BASE_URL"] == "https://api.anthropic.com" - assert env["ANTHROPIC_MODEL"] == "claude-haiku-4-5" - - def test_unresolved_placeholder_skipped(self): - """``${VAR}`` (unresolved interpolation) must NOT be propagated — - we only inject real values.""" - from dsagt.agents import agent_env - - env = agent_env( - self._config_with_llm( - provider="openai", - api_key="${LLM_API_KEY}", - base_url="${LLM_BASE_URL}", - ) - ) - # Provider still set, but bogus key/url not injected. - assert env["GOOSE_PROVIDER"] == "openai" - assert env.get("OPENAI_API_KEY") != "${LLM_API_KEY}" - assert env.get("OPENAI_BASE_URL") != "${LLM_BASE_URL}" - - def test_blank_values_skipped(self): - from dsagt.agents import agent_env - - env = agent_env( - self._config_with_llm( - provider="openai", - api_key="", - base_url=" ", - ) - ) - # Blank values shouldn't override the user's shell-level env vars - # (which the os.environ copy in agent_env preserves). - assert env.get("OPENAI_API_KEY") != "" - assert env.get("OPENAI_BASE_URL") != " " - - def test_unknown_provider_still_sets_goose_vars(self): - """Goose can dispatch to unknown providers via its own config - fallback; we still set GOOSE_PROVIDER / GOOSE_MODEL so the agent - knows which provider to look up.""" - from dsagt.agents import agent_env - - env = agent_env( - self._config_with_llm( - provider="bedrock", - model="anthropic.claude-3-sonnet", - ) - ) - assert env["GOOSE_PROVIDER"] == "bedrock" - assert env["GOOSE_MODEL"] == "anthropic.claude-3-sonnet" - # No OPENAI_* / ANTHROPIC_* injection for unknown providers — user - # supplies those via their shell env (e.g. AWS_ACCESS_KEY_ID). - assert "OPENAI_API_KEY" not in env or env["OPENAI_API_KEY"] != "" - - @pytest.mark.skip(reason="proxy mode deferred to Phase 2") - def test_proxy_overrides_provider_injection(self): - """When --enable-proxy is set, the proxy block runs AFTER provider - injection and overrides OPENAI_BASE_URL / OPENAI_API_KEY with the - proxy URL + sentinel. Provider injection must not stomp the proxy.""" - from dsagt.agents import agent_env - - config = self._config_with_llm( - provider="openai", - api_key="sk-real", - base_url="https://upstream.example.com", - ) - config["proxy"] = {"port": 9999} - env = agent_env(config) - assert env["OPENAI_BASE_URL"] == "http://localhost:9999" - assert env["OPENAI_API_KEY"].startswith("dsagt-proxy-forwarded") - - def test_claude_does_not_get_openai_env(self): - """Protocol isolation: claude is anthropic-native. When the user - configures provider=openai, claude's env_overrides must NOT - propagate OPENAI_* — claude's runtime ignores them and the cross- - protocol mismatch is what the proxy is for. Setting them anyway - muddies the contract.""" - from dsagt.agents import agent_env - - cfg = self._config_with_llm( - provider="openai", - api_key="sk-real", - base_url="https://gateway.example.com", - ) - cfg["agent"] = "claude" - env = agent_env(cfg) - # Claude is openai-blind; project openai creds shouldn't show up - # in claude's process env (the user's shell may set them, in - # which case os.environ wins — that's a user-controlled choice). - # We assert ours specifically isn't injected. - assert env.get("OPENAI_BASE_URL") != "https://gateway.example.com" - - def test_codex_does_not_get_anthropic_env(self): - """Protocol isolation: codex is openai-native. Anthropic upstream - requires --enable-proxy for translation.""" - from dsagt.agents import agent_env - - cfg = self._config_with_llm( - provider="anthropic", - model="claude-sonnet", - api_key="anthropic-key", - base_url="https://api.anthropic.com", - ) - cfg["agent"] = "codex" - env = agent_env(cfg) - assert env.get("ANTHROPIC_BASE_URL") != "https://api.anthropic.com" - - -class TestPhase2EnvOverrides: - """Phase 2 ``env_overrides`` is narrower than old code: it pins the - agent's MODEL env var only. Provider creds / base URLs are set by - ``proxy_env_overrides`` (sentinel + proxy URL) or live in agent - config files (cline auth state, codex config.toml, opencode.json). - All gated by ``config["proxy"]["port"]`` — BYOA mode doesn't fire it. - """ - - def _proxy_config(self, agent: str, model: str = "test-model") -> dict: - return { - "project": "test", - "agent": agent, - "project_dir": "/proj", - "mlflow": {"port": 5001}, - "llm": { - "provider": "openai", - "model": model, - "base_url": "https://up", - "api_key": "sk-up", - }, - "embedding": {}, - "proxy": {"port": 9999}, - } - - def test_claude_pins_anthropic_model_in_proxy_mode(self): - from dsagt.agents import agent_env - - env = agent_env(self._proxy_config("claude", model="claude-test")) - assert env["ANTHROPIC_MODEL"] == "claude-test" - - def test_goose_pins_provider_and_model_in_proxy_mode(self): - from dsagt.agents import agent_env - - env = agent_env(self._proxy_config("goose", model="my-model")) - assert env["GOOSE_PROVIDER"] == "openai" - assert env["GOOSE_MODEL"] == "my-model" - - def test_roo_pins_anthropic_model_in_proxy_mode(self): - from dsagt.agents import agent_env - - env = agent_env(self._proxy_config("roo", model="my-roo-model")) - assert env["ANTHROPIC_MODEL"] == "my-roo-model" + for var in ( + "MLFLOW_TRACKING_URI", + "OTEL_EXPORTER_OTLP_ENDPOINT", + "OTEL_EXPORTER_OTLP_HEADERS", + "OTEL_RESOURCE_ATTRIBUTES", + ): + monkeypatch.delenv(var, raising=False) - def test_byoa_mode_does_not_fire_env_overrides(self): - """Without proxy_port in config, env_overrides is not called — - the user's shell ANTHROPIC_MODEL / GOOSE_MODEL etc. are not - overwritten by config["llm"] values. This is the gate that - prevents the Phase-1 GOOSE_MODEL bug from coming back. + for agent in ("claude", "goose", "codex", "cline", "opencode"): + env = agent_env(self._make_config(agent)) + # Serverless sqlite store derived from the project dir. + assert env["MLFLOW_TRACKING_URI"] == "sqlite:////proj/mlflow.db" + assert "OTEL_EXPORTER_OTLP_ENDPOINT" not in env + assert "OTEL_EXPORTER_OTLP_HEADERS" not in env + assert "OTEL_RESOURCE_ATTRIBUTES" not in env + + def test_no_telemetry_flags_for_claude(self, monkeypatch): + """The forced Claude telemetry/privacy flags are gone — DSAGT no + longer flips ``CLAUDE_CODE_ENABLE_TELEMETRY`` / ``OTEL_LOG_*`` + (the latter defeated Anthropic's off-by-default redaction). + + Cleared from the inherited shell env first so we test that DSAGT + doesn't *add* them — not whatever the test runner's own shell set. """ from dsagt.agents import agent_env - cfg = self._proxy_config("goose", model="config-model") - cfg.pop("proxy") # BYOA: no proxy block - env = agent_env(cfg) - # GOOSE_MODEL only set if user's shell has it; we don't override. - assert env.get("GOOSE_MODEL") != "config-model" - - def test_unresolved_placeholder_filtered(self): - """A ${VAR} that didn't resolve at load_config time must not - leak into env vars (would be an obvious bug surface).""" - from dsagt.agents import agent_env - - cfg = self._proxy_config("goose", model="${LLM_MODEL}") - env = agent_env(cfg) - assert "GOOSE_MODEL" not in env or env["GOOSE_MODEL"] != "${LLM_MODEL}" - - -class TestProxyEnvOverrides: - """``--enable-proxy`` plants the same proxy-routing env vars on - every agent. AgentSetup.proxy_env_overrides centralizes the - contract; agents inherit the default unless they have a non- - standard env-var convention (goose's OPENAI_HOST/ANTHROPIC_HOST).""" - - def _make_config(self, agent: str) -> dict: - return { - "project": "test", - "agent": agent, - "project_dir": "/proj", - "mlflow": {"port": 5001}, - "llm": {"provider": "openai", "api_key": "sk-up", "base_url": "https://up"}, - "embedding": {}, - "proxy": {"port": 9999}, - } - - def test_default_sets_proxy_url_for_both_protocols(self): - from dsagt.agents.base import AgentSetup - - # Spin up any concrete subclass to call the inherited default. - from dsagt.agents.goose import GooseSetup - - env = GooseSetup().proxy_env_overrides(9999) - assert env["ANTHROPIC_BASE_URL"] == "http://localhost:9999" - assert env["OPENAI_BASE_URL"] == "http://localhost:9999" - assert env["EMBEDDING_BASE_URL"] == "http://localhost:9999" - - def test_default_plants_sentinel_keys(self): - from dsagt.agents.goose import GooseSetup - - env = GooseSetup().proxy_env_overrides(9999) - assert env["ANTHROPIC_API_KEY"].startswith("dsagt-proxy-forwarded") - assert env["OPENAI_API_KEY"].startswith("dsagt-proxy-forwarded") - assert env["EMBEDDING_API_KEY"].startswith("dsagt-proxy-forwarded") - - def test_proxy_applied_uniformly_across_agents(self): - """Every agent gets identical proxy env from the inherited - default — that's the whole point of putting it on the base class.""" - from dsagt.agents import agent_env - - envs = { - a: agent_env(self._make_config(a)) - for a in ("claude", "goose", "cline", "roo", "codex") - } - for agent, env in envs.items(): - assert env["ANTHROPIC_BASE_URL"] == "http://localhost:9999", agent - assert env["OPENAI_BASE_URL"] == "http://localhost:9999", agent - assert env["OPENAI_API_KEY"].startswith("dsagt-proxy-forwarded"), agent - - -class TestPreconfiguredCredsWarning: - """When the project YAML has no llm credentials, ``agent_env`` warns - that the agent will fall back to the user's shell env, listing which - of the agent's credential env vars are actually present.""" - - def _config(self, agent: str = "claude", **llm) -> dict: - return { - "project": "test", - "agent": agent, - "project_dir": "/proj", - "mlflow": {"port": 5001}, - "llm": llm, - "embedding": {}, - } - - def test_warns_when_project_has_no_creds_but_shell_does(self, caplog, monkeypatch): - from dsagt.agents import agent_env - - monkeypatch.setenv("ANTHROPIC_API_KEY", "shell-key") - monkeypatch.setenv("ANTHROPIC_BASE_URL", "https://api.anthropic.com") - - with caplog.at_level("WARNING", logger="dsagt.agents"): - agent_env(self._config(agent="claude")) - - msgs = " ".join(r.getMessage() for r in caplog.records) - assert "preconfigured env vars" in msgs - assert "ANTHROPIC_API_KEY" in msgs - assert "ANTHROPIC_BASE_URL" in msgs - # Var names only, never values. - assert "shell-key" not in msgs - - def test_no_warning_when_project_supplies_creds(self, caplog): - from dsagt.agents import agent_env - - cfg = self._config( - agent="claude", - provider="anthropic", - api_key="project-key", - base_url="https://api.anthropic.com", - model="claude-haiku", + flags = ( + "CLAUDE_CODE_ENABLE_TELEMETRY", + "CLAUDE_CODE_ENHANCED_TELEMETRY_BETA", + "OTEL_LOG_TOOL_DETAILS", + "OTEL_LOG_USER_PROMPTS", + "OTEL_TRACES_EXPORTER", + "OTEL_LOG_RAW_API_BODIES", ) + for flag in flags: + monkeypatch.delenv(flag, raising=False) - with caplog.at_level("WARNING", logger="dsagt.agents"): - agent_env(cfg) - - msgs = " ".join(r.getMessage() for r in caplog.records) - assert "preconfigured env vars" not in msgs - - def test_no_warning_when_proxy_enabled(self, caplog): - """Proxy mode plants its own creds; transparency warning is moot.""" - from dsagt.agents import agent_env - - cfg = self._config(agent="claude") - cfg["proxy"] = {"port": 9999} - - with caplog.at_level("WARNING", logger="dsagt.agents"): - agent_env(cfg) - - msgs = " ".join(r.getMessage() for r in caplog.records) - assert "preconfigured env vars" not in msgs - - def test_warns_with_no_shell_either_lists_expected_vars(self, caplog, monkeypatch): - """When neither project nor shell has creds, surface the var names - the agent's runtime expects so the user can fix the gap.""" - from dsagt.agents import agent_env - - for var in ("ANTHROPIC_API_KEY", "ANTHROPIC_BASE_URL", "ANTHROPIC_MODEL"): - monkeypatch.delenv(var, raising=False) - - with caplog.at_level("WARNING", logger="dsagt.agents"): - agent_env(self._config(agent="claude")) - - msgs = " ".join(r.getMessage() for r in caplog.records) - assert "fall back to its own auth flow" in msgs - assert "ANTHROPIC_API_KEY" in msgs + env = agent_env(self._make_config("claude")) + for flag in flags: + assert flag not in env # --------------------------------------------------------------------------- @@ -1265,16 +980,9 @@ def test_goose(self): "goose", "session", "--with-extension", - "uv run dsagt-registry-server", - "--with-extension", - "uv run dsagt-knowledge-server", + "uv run dsagt-server", ] - def test_roo(self): - from dsagt.agents import agent_command - - assert agent_command({"agent": "roo"}) == ["roo"] - def test_cline(self): from dsagt.agents import agent_command @@ -1293,12 +1001,27 @@ def test_codex(self): class TestConfigFlow: - def test_default_config_has_embedding_block(self): - """default_config_content() includes embedding (KB needs it).""" - content = default_config_content("test", "claude", mlflow_port=5000) + def test_default_config_mirrors_init_choices(self): + """The written config holds only the init choices: project, agent, + knowledge.collections, skills.sources. The bundled ``tools`` collection + is always provisioned (not a choice); embedding / chunk_size / rerank + are code defaults backfilled on read, never written.""" + content = default_config_content("test", "claude") parsed = yaml.safe_load(content) - assert "embedding" in parsed - assert "backend" in parsed["embedding"] + assert set(parsed) == {"project", "agent", "knowledge", "skills"} + assert parsed["knowledge"] == {"collections": []} + assert parsed["skills"]["sources"][0]["name"] == "genesis" + assert "embedding" not in parsed + assert "episodic" not in parsed # opt-in, omitted unless enabled + + def test_episodic_written_only_when_enabled(self): + """An opted-in episodic block is written verbatim; absent otherwise + (and ``load_config`` backfills ``enabled: false`` for the absent case).""" + epi = {"enabled": True} + parsed = yaml.safe_load(default_config_content("t", "claude", episodic=epi)) + assert parsed["episodic"] == epi + # Omitted when None. + assert "episodic" not in yaml.safe_load(default_config_content("t", "claude")) def test_mcp_env_block_carries_embedding_routing(self): """_mcp_env_block plumbs embedding routing (model + base_url) @@ -1324,23 +1047,23 @@ def test_mcp_env_block_carries_embedding_routing(self): # Credentials never land in artifacts; user sets in shell. assert "EMBEDDING_API_KEY" not in env - def test_mcp_env_block_omits_project_routing(self): - """Single source of truth: MLflow URI / project name / project_dir - come from dsagt_config.yaml via cwd-walk, not from a duplicated - env block. The MCP env block carries only embedding-backend - settings (and EMBEDDING_API_KEY from the user's shell).""" + def test_mcp_env_block_carries_project_routing(self, monkeypatch): + """Benign routing: the MCP env block carries project name + dir and + the serverless ``MLFLOW_TRACKING_URI`` so MCP children of agents + that don't inherit the parent shell env still log to the right + store. No credentials, no OTel.""" from dsagt.agents import _mcp_env_block + monkeypatch.delenv("MLFLOW_TRACKING_URI", raising=False) config = { "project": "test", "project_dir": "/p", - "mlflow": {"port": 12345}, "embedding": {"backend": "local"}, } env = _mcp_env_block(config) - assert "MLFLOW_TRACKING_URI" not in env - assert "DSAGT_PROJECT" not in env - assert "DSAGT_PROJECT_DIR" not in env + assert env["DSAGT_PROJECT"] == "test" + assert env["DSAGT_PROJECT_DIR"] == "/p" + assert env["MLFLOW_TRACKING_URI"] == "sqlite:////p/mlflow.db" assert env["EMBEDDING_BACKEND"] == "local" def test_mcp_env_block_omits_empty_embedding_keys(self): @@ -1360,220 +1083,37 @@ def test_mcp_env_block_omits_empty_embedding_keys(self): assert "EMBEDDING_MODEL" not in env def test_mcp_server_args_are_just_command(self): - """MCP server args are just ["run", "dsagt--server"]. - All configuration flows through env vars and dsagt_config.yaml. + """MCP server args are just ["run", "dsagt-server"] — one merged + server. All configuration flows through .dsagt/config.yaml (cwd-walk). """ from dsagt.agents import _mcp_server_args - assert _mcp_server_args("knowledge") == ["run", "dsagt-knowledge-server"] - assert _mcp_server_args("registry") == ["run", "dsagt-registry-server"] + assert _mcp_server_args() == ["run", "dsagt-server"] - def test_mcp_env_block_carries_only_embedding_settings(self): + def test_mcp_env_block_carries_no_session_id(self): + """The MCP server owns the session lifecycle now (minted into + ``.dsagt/state.yaml`` at startup), so the env block never carries a + ``DSAGT_SESSION_ID`` — only project routing + embedding settings.""" from dsagt.agents import _mcp_env_block config = { "project": "test", "project_dir": "/home/user/dsagt-projects/test", - "mlflow": {"port": 5000}, "embedding": {"model": "m", "base_url": "u"}, } env = _mcp_env_block(config) - # Project routing comes from dsagt_config.yaml via cwd-walk — - # never duplicated into the MCP env block. - assert "DSAGT_PROJECT_DIR" not in env - assert "DSAGT_PROJECT" not in env + assert "DSAGT_SESSION_ID" not in env assert env["EMBEDDING_MODEL"] == "m" assert env["EMBEDDING_BASE_URL"] == "u" + # Even if a stray session_id is on the config dict, it isn't emitted. + env2 = _mcp_env_block({**config, "session_id": "test-1"}) + assert "DSAGT_SESSION_ID" not in env2 -# --------------------------------------------------------------------------- -# BYOA: per-agent env hints + launch one-liners surfaced by `dsagt init` -# --------------------------------------------------------------------------- - - -class TestByoaEnvHints: - """``dsagt init`` prints provider credentials only. Internal env - (DSAGT_*, MLFLOW_*, OTEL_*, telemetry verbosity flags) goes into - the per-project launch shim — the user's shell stays clean.""" - - @pytest.mark.parametrize( - "agent_name", ["claude", "goose", "cline", "roo", "codex", "opencode"] - ) - def test_returns_only_credential_hints(self, agent_name, tmp_path): - from dsagt.agents import AGENTS - - setup = AGENTS[agent_name]() - hints = setup.byoa_env_hints( - mlflow_port=5001, project="p", project_dir=tmp_path - ) - # Only credential hints come back; no DSAGT/MLflow/OTel routing. - names = [n for n, _ in hints] - assert "DSAGT_PROJECT" not in names - assert "MLFLOW_TRACKING_URI" not in names - assert "OTEL_EXPORTER_OTLP_ENDPOINT" not in names - - @pytest.mark.parametrize( - "agent_name", ["claude", "goose", "cline", "roo", "codex", "opencode"] - ) - def test_credentials_match_credential_hints(self, agent_name, tmp_path): - from dsagt.agents import AGENTS - - setup = AGENTS[agent_name]() - hints = setup.byoa_env_hints(5001, "p", tmp_path) - assert hints == list(setup.credential_hints) - - def test_goose_credentials_include_host_not_base_url(self, tmp_path): - """Goose's Rust client reads OPENAI_HOST / ANTHROPIC_HOST, NOT - OPENAI_BASE_URL — surfacing this gotcha to the user up front - avoids the silent 'agent hits api.openai.com regardless of - gateway' bug.""" - from dsagt.agents import AGENTS - - names = [n for n, _ in AGENTS["goose"]().byoa_env_hints(5001, "p", tmp_path)] - assert "OPENAI_HOST" in names - assert "ANTHROPIC_HOST" in names - - @pytest.mark.parametrize( - "agent_name,gateway_var", - [ - ("claude", "ANTHROPIC_BASE_URL"), - ("codex", "OPENAI_BASE_URL"), - # Cline / roo: ``cline auth -b`` is openai-only and openai-native - # path needs a non-standard model env var, so we standardize on - # the anthropic path for both — same env conventions as the rest - # of the BYOA story. See cline.py / roo.py credential_hints. - ("cline", "ANTHROPIC_BASE_URL"), - ("roo", "ANTHROPIC_BASE_URL"), - # opencode emits both — its provider config supports both wire - # protocols via {env:VAR} interpolation in opencode.json. - ("opencode", "OPENAI_BASE_URL"), - ("opencode", "ANTHROPIC_BASE_URL"), - ], - ) - def test_gateway_url_in_credential_hints(self, agent_name, gateway_var, tmp_path): - """Lab gateway / proxy URL hint is surfaced for every agent that - speaks the standard BASE_URL convention.""" - from dsagt.agents import AGENTS - - names = [n for n, _ in AGENTS[agent_name]().byoa_env_hints(5001, "p", tmp_path)] - assert gateway_var in names - - -class TestLaunchOneliner: - """``launch_oneliner`` returns the literal agent command (no shim). - The shim is a separate file (``dsagt-launch.sh``) written by - ``dynamic_agent_record`` — see TestLaunchShim below.""" - - @pytest.mark.parametrize( - "agent_name", ["claude", "goose", "cline", "roo", "codex", "opencode"] - ) - def test_oneliner_does_not_invoke_shim(self, agent_name, tmp_path): - from dsagt.agents import AGENTS - - cmd = AGENTS[agent_name]().launch_oneliner("myproj", tmp_path) - assert "dsagt-launch.sh" not in cmd - - @pytest.mark.parametrize( - "agent_name,expected_cmd", - [ - ("claude", "claude"), - ("goose", "goose session"), - ("roo", "roo"), - ("opencode", "opencode"), - ], - ) - def test_oneliner_runs_agent_directly(self, agent_name, expected_cmd, tmp_path): - from dsagt.agents import AGENTS - - cmd = AGENTS[agent_name]().launch_oneliner("myproj", tmp_path) - assert expected_cmd in cmd - assert f"cd {tmp_path}" in cmd - - def test_cline_oneliner_includes_config_flag(self, tmp_path): - """Cline's ``--config /.cline-data`` is the only way to - pick up the per-project MCP-server registrations.""" - from dsagt.agents import AGENTS - - cmd = AGENTS["cline"]().launch_oneliner("myproj", tmp_path) - assert "--config" in cmd - assert ".cline-data" in cmd - - def test_codex_oneliner_sets_codex_home(self, tmp_path): - """Codex has no ``--config`` flag — must export ``CODEX_HOME`` - before launching so codex finds the per-project config.toml.""" - from dsagt.agents import AGENTS - - cmd = AGENTS["codex"]().launch_oneliner("myproj", tmp_path) - assert "CODEX_HOME=" in cmd - assert ".codex-data" in cmd - - -# --------------------------------------------------------------------------- -# dsagt memory: high-water-mark extraction state -# --------------------------------------------------------------------------- - - -class TestMemoryWatermark: - """``dsagt memory --project X`` tracks which sessions have been - extracted in ``/.dsagt/extracted_at.json`` so re-runs only - process new traces.""" - - def test_first_run_creates_watermark_on_success(self, tmp_path, monkeypatch): - from argparse import Namespace - from dsagt.commands.cli import _cmd_memory - - pdir, _ = init_project("mem-test", "claude") - monkeypatch.setattr( - "dsagt.commands.cli.run_extraction", - lambda _project: {"status": "ok", "total_entries": 3}, - ) - _cmd_memory(Namespace(project="mem-test")) - - state_path = pdir / ".dsagt" / "extracted_at.json" - assert state_path.exists() - state = json.loads(state_path.read_text()) - assert "last_extracted_at" in state - assert state["previous"] is None - - def test_subsequent_run_records_previous_watermark(self, tmp_path, monkeypatch): - from argparse import Namespace - from dsagt.commands.cli import _cmd_memory - - pdir, _ = init_project("mem-test", "claude") - monkeypatch.setattr( - "dsagt.commands.cli.run_extraction", - lambda _project: {"status": "ok", "total_entries": 1}, - ) - _cmd_memory(Namespace(project="mem-test")) - first_state = json.loads((pdir / ".dsagt" / "extracted_at.json").read_text()) - - _cmd_memory(Namespace(project="mem-test")) - second_state = json.loads((pdir / ".dsagt" / "extracted_at.json").read_text()) - - # The previous run's mark moves into the "previous" slot. - assert second_state["previous"] == first_state["last_extracted_at"] - assert second_state["last_extracted_at"] != first_state["last_extracted_at"] - - def test_empty_extraction_does_not_advance_watermark(self, tmp_path, monkeypatch): - """When run_extraction returns status=empty, the state file is - not written — re-running later still picks up new traces.""" - from argparse import Namespace - from dsagt.commands.cli import _cmd_memory - - pdir, _ = init_project("mem-test", "claude") - monkeypatch.setattr( - "dsagt.commands.cli.run_extraction", - lambda _project: {"status": "empty"}, - ) - _cmd_memory(Namespace(project="mem-test")) - - assert not (pdir / ".dsagt" / "extracted_at.json").exists() - - -class TestLaunchShim: - """``dynamic_agent_record`` writes ``dsagt-launch.sh`` for BYOA-mode - projects. The shim starts MLflow, exports env, and prints how to - launch the agent (it does NOT exec the agent — user picks).""" +class TestNoLaunchShim: + """Phase 1 collapsed the launch surface: ``dynamic_agent_record`` + writes the MCP config but NO ``dsagt-launch.sh`` shim. The user + starts the agent directly in the project dir or via ``dsagt start``.""" def _make_config(self, agent_name: str, pdir): return { @@ -1586,103 +1126,38 @@ def _make_config(self, agent_name: str, pdir): "session_id": "sess-xyz", } - def test_shim_written_for_byoa(self, tmp_path): - """`dynamic_agent_record` writes dsagt-launch.sh in BYOA mode.""" + @pytest.mark.parametrize("agent_name", ["goose", "claude", "codex"]) + def test_no_shim_written(self, agent_name, tmp_path): from dsagt.agents import dynamic_agent_record - config = self._make_config("goose", tmp_path) + config = self._make_config(agent_name, tmp_path) dynamic_agent_record(config, env={}, working_dir=tmp_path) - shim = tmp_path / "dsagt-launch.sh" - assert shim.exists() - assert (shim.stat().st_mode & 0o111) != 0 # executable + assert not (tmp_path / "dsagt-launch.sh").exists() - def test_shim_does_not_exec_agent(self, tmp_path): - """Shim prints launch options instead of execing — user picks.""" - from dsagt.agents import dynamic_agent_record - - config = self._make_config("goose", tmp_path) - dynamic_agent_record(config, env={}, working_dir=tmp_path) - shim_text = (tmp_path / "dsagt-launch.sh").read_text() - assert "exec " not in shim_text - assert "Environment ready. Launch the agent" in shim_text - assert "goose session" in shim_text # CLI option printed - - def test_shim_starts_mlflow_in_background(self, tmp_path): - from dsagt.agents import dynamic_agent_record +class TestClaudeSetup: + """`dsagt init --agent claude` writes `.mcp.json` and does NOT wire MLflow's + autolog Stop hook — DSAGT's own serverless heartbeat pipeline (ClaudeReader → + ClaudeTranslator → MLflowSink) produces Claude's traces, uniformly with every + other agent, so wiring autolog too would double-log.""" - config = self._make_config("goose", tmp_path) - dynamic_agent_record(config, env={}, working_dir=tmp_path) - - shim_text = (tmp_path / "dsagt-launch.sh").read_text() - assert "dsagt mlflow test --background-only" in shim_text - - def test_shim_for_claude_skips_otel_routing(self, tmp_path): - """Claude shim omits OTEL_EXPORTER_OTLP_* exports — agent-side - traces come from `mlflow autolog claude`'s Stop hook.""" - from dsagt.agents import dynamic_agent_record - - config = self._make_config("claude", tmp_path) - dynamic_agent_record(config, env={}, working_dir=tmp_path) - - shim_text = (tmp_path / "dsagt-launch.sh").read_text() - assert "OTEL_EXPORTER_OTLP_TRACES_ENDPOINT" not in shim_text - assert "MLFLOW_TRACKING_URI" in shim_text # still exported - - def test_shim_for_goose_includes_otel_routing(self, tmp_path): - """Non-Claude agents need OTel routing (no autolog analog).""" - from dsagt.agents import dynamic_agent_record - - config = self._make_config("goose", tmp_path) - dynamic_agent_record(config, env={}, working_dir=tmp_path) - - shim_text = (tmp_path / "dsagt-launch.sh").read_text() - assert "OTEL_EXPORTER_OTLP_TRACES_ENDPOINT" in shim_text - assert "service.name=goose" in shim_text - - -class TestClaudeAutologSetup: - """`dsagt init --agent claude` configures `mlflow autolog claude` - by writing `.claude/settings.json` with the Stop hook + tracking - env vars.""" - - def test_settings_file_written(self, tmp_path): + def test_writes_mcp_json_no_autolog_hook(self, tmp_path, monkeypatch): from dsagt.agents.claude import ClaudeSetup + monkeypatch.delenv("MLFLOW_TRACKING_URI", raising=False) config = { "project": "myproj", "project_dir": str(tmp_path), "agent": "claude", - "mlflow": {"port": 5099}, "embedding": {}, "llm": {}, } actions = ClaudeSetup().write_dynamic( - config, - env={}, - working_dir=tmp_path, - pdir=tmp_path, + config, env={}, working_dir=tmp_path, pdir=tmp_path ) - settings_file = tmp_path / ".claude" / "settings.json" - assert settings_file.exists() - assert any("mlflow autolog claude" in a for a in actions) - - import json as _json - - settings = _json.loads(settings_file.read_text()) - env_block = settings.get("env") or {} - assert env_block.get("MLFLOW_TRACKING_URI") == "http://localhost:5099" - assert env_block.get("MLFLOW_EXPERIMENT_NAME") == "myproj" - assert env_block.get("MLFLOW_CLAUDE_TRACING_ENABLED") == "true" - - hooks = settings.get("hooks") or {} - stop_hooks = hooks.get("Stop") or [] - # Stop hook should reference mlflow autolog claude. - all_commands = [] - for group in stop_hooks: - for h in group.get("hooks") or []: - if h.get("command"): - all_commands.append(h["command"]) - assert any("mlflow autolog claude" in c for c in all_commands) + assert (tmp_path / ".mcp.json").exists() + # No autolog: no Stop hook, no .claude/settings.json. + assert not any("autolog" in a.lower() for a in actions) + assert not (tmp_path / ".claude" / "settings.json").exists() diff --git a/tests/test_csv_summary_fixture.py b/tests/test_csv_summary_fixture.py new file mode 100644 index 0000000..4b67238 --- /dev/null +++ b/tests/test_csv_summary_fixture.py @@ -0,0 +1,44 @@ +"""Guard the quickstart's bundled csv_summary fixture. + +The README/docs quickstart has the agent register and run +``tests/smoke_test/csv_summary.py`` on ``samples.csv``; step 4's +null-column finding is the fact the explicit-memory step (5-6) stores and +recalls. This pins that contract so the fixture can't silently drift out +from under the docs. +""" + +import importlib.util +from pathlib import Path + +_FIXTURE = Path(__file__).parent / "smoke_test" / "csv_summary.py" +_SAMPLES = Path(__file__).parent / "smoke_test" / "data" / "samples.csv" + + +def _load(): + spec = importlib.util.spec_from_file_location("csv_summary_fixture", _FIXTURE) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + return mod + + +def test_summary_reports_columns_and_row_count(): + result = _load().summarize(str(_SAMPLES)) + assert result["columns"] == ["id", "name", "status", "score", "timestamp"] + assert result["row_count"] == 8 + + +def test_summary_surfaces_the_null_columns_the_memory_step_stores(): + result = _load().summarize(str(_SAMPLES)) + # The quickstart's step 5 remembers "null values in the status and + # timestamp columns" — it must come from here, not thin air. + assert result["columns_with_nulls"] == ["status", "timestamp"] + assert result["null_counts"]["status"] == 1 + assert result["null_counts"]["timestamp"] == 1 + + +def test_summary_computes_numeric_stats_only_for_numeric_columns(): + result = _load().summarize(str(_SAMPLES)) + assert set(result["numeric_stats"]) == {"id", "score"} + assert result["numeric_stats"]["score"]["max"] == 92.1 + # timestamp is date-like, not numeric — must not appear. + assert "timestamp" not in result["numeric_stats"] diff --git a/tests/test_dependency_integration.py b/tests/test_dependency_integration.py index c0ad35b..53c4251 100644 --- a/tests/test_dependency_integration.py +++ b/tests/test_dependency_integration.py @@ -2,7 +2,7 @@ Integration test for dependency installation during tool registration. Registers a tool with a real dependency (cowsay), installs it via the -registry server, then executes the tool through the pipeline's ToolRegistry +registry server, then executes the tool through the pipeline's CodeRegistry to verify the package is usable. This test actually modifies the venv (installs and uninstalls cowsay). @@ -21,24 +21,18 @@ import subprocess import sys import textwrap -from pathlib import Path import pytest -pytestmark = pytest.mark.integration - -import pytest -import yaml - - -from dsagt.commands.registry_server import create_registry_server -from dsagt.registry import ToolRegistry +from dsagt.mcp.registry_tools import create_registry_server +from dsagt.registry import CodeRegistry # --------------------------------------------------------------------------- # Skip conditions # --------------------------------------------------------------------------- + def _uv_available() -> bool: return shutil.which("uv") is not None @@ -52,6 +46,7 @@ def _cowsay_installed() -> bool: pytestmark = [ + pytest.mark.integration, pytest.mark.skipif(not _uv_available(), reason="uv not available on PATH"), pytest.mark.skipif(_cowsay_installed(), reason="cowsay already installed"), ] @@ -61,6 +56,7 @@ def _cowsay_installed() -> bool: # Fixtures # --------------------------------------------------------------------------- + @pytest.fixture(scope="module", autouse=True) def uninstall_cowsay_after(): """Ensure cowsay is uninstalled after all tests in this module.""" @@ -74,11 +70,11 @@ def uninstall_cowsay_after(): from mcp_helpers import call_tool_sync as call_tool - # --------------------------------------------------------------------------- # Test # --------------------------------------------------------------------------- + def test_register_and_run_tool_with_dependency(tmp_path): """End-to-end: register a tool with a dependency, install it, run the tool.""" # 1. Write a test script that imports cowsay @@ -96,9 +92,8 @@ def test_register_and_run_tool_with_dependency(tmp_path): print(json.dumps({"cow_says": output, "status": "ok"})) """)) - # 2. Create a registry server with a fresh ToolRegistry - registry = ToolRegistry( - source_tools_dir=None, + # 2. Create a registry server with a fresh CodeRegistry + registry = CodeRegistry( runtime_dir=str(tmp_path / "runtime"), ) server = create_registry_server(registry) @@ -117,14 +112,14 @@ def test_register_and_run_tool_with_dependency(tmp_path): }, }, } - text = call_tool(server, "save_tool_spec", {"spec": spec}) + text = call_tool(server, "save_code_spec", {"spec": spec}) # Verify the tool was saved and deps were installed assert "added" in text assert "Successfully installed" in text # 4. Verify the spec is in the skill file with dsagt-run wrapping - tool = registry.get_tool("cowsay_tool") + tool = registry.get_code("cowsay_tool") assert tool is not None assert tool["dependencies"] == ["cowsay"] assert "dsagt-run" in tool["executable"] diff --git a/tests/test_dsagt_server.py b/tests/test_dsagt_server.py new file mode 100644 index 0000000..0174744 --- /dev/null +++ b/tests/test_dsagt_server.py @@ -0,0 +1,134 @@ +"""Tests for the merged ``dsagt-server`` (all four concern modules under one Server). + +These verify the *composition* contract: every tool from the registry / knowledge +/ memory / skill modules is exposed under one MCP ``Server``, and the single +``call_tool`` wrapper preserves both return-type contracts (registry + skill +handlers may return a plain string; knowledge / memory handlers return a dict +that gets JSON-encoded). Also covers ``_build_kb_from_config`` credential +validation in-process (the full subprocess boot needs a live MLflow — see +``test_server_startup.py``). +""" + +import asyncio +import json +from pathlib import Path +from unittest.mock import MagicMock + +import mcp.types as types +import pytest + +from dsagt.mcp.server import _build_kb_from_config, create_dsagt_server +from dsagt.registry import SkillRegistry, CodeRegistry + + +def _make_merged_server(tmp_path: Path): + kb = MagicMock() + kb.index_dir = tmp_path / "kb_index" + kb.index_dir.mkdir() + kb.default_rerank = True + kb.collections = [] + runtime = str(tmp_path / "runtime") + reg = CodeRegistry(runtime_dir=runtime, kb=None) + reg.ensure_bundled_copies() + sreg = SkillRegistry(source_skills_dir=None, runtime_dir=runtime, kb=None) + return create_dsagt_server(reg, kb, sreg, runtime_dir=runtime) + + +def _list_tools(server) -> list[str]: + handler = server.request_handlers[types.ListToolsRequest] + res = asyncio.run(handler(types.ListToolsRequest(method="tools/list"))) + return sorted(t.name for t in res.root.tools) + + +def _call(server, name: str, arguments: dict) -> str: + handler = server.request_handlers[types.CallToolRequest] + req = types.CallToolRequest( + method="tools/call", + params=types.CallToolRequestParams(name=name, arguments=arguments), + ) + res = asyncio.run(handler(req)) + return res.root.content[0].text + + +def test_merged_server_exposes_all_tools(tmp_path): + """Both concern modules' tools land under one server with no collision.""" + server = _make_merged_server(tmp_path) + names = _list_tools(server) + # 8 registry + 5 knowledge + 2 memory + 5 skill = 20 distinct tools. + assert set(names) == { + # registry / provenance (8) + "get_registry", + "search_registry", + "save_code_spec", + "install_dependencies", + "run_command", + "read_file", + "http_request", + "reconstruct_pipeline", + # knowledge (5) + "kb_search", + "kb_ingest", + "kb_list_collections", + "kb_job_status", + "kb_append", + # memory (2) + "kb_remember", + "kb_get_memories", + # skills (5) + "search_skills", + "install_skill", + "save_skill", + "add_skill_source", + "list_skill_sources", + } + assert len(set(names)) == len(names) # no name collision + + +def test_registry_tool_returns_plain_string(tmp_path): + """Registry handlers return a bare string — passed through unchanged.""" + server = _make_merged_server(tmp_path) + out = _call(server, "get_registry", {}) + # Not JSON — the registry contract is a human-readable string. + with pytest.raises(json.JSONDecodeError): + json.loads(out) + assert "codes:" in out + + +def test_dict_returning_handler_is_json_encoded(tmp_path): + """Dict-returning handlers (knowledge/memory/skill) are JSON-encoded by the wrapper.""" + server = _make_merged_server(tmp_path) + out = _call(server, "list_skill_sources", {}) + parsed = json.loads(out) + assert "sources" in parsed + + +class TestBuildKbFromConfig: + """``_build_kb_from_config`` validates embedding config before building a KB. + + These raise paths fire before any embedder / ChromaDB construction, so they + need no real backend. + """ + + def _cfg(self, **embedding): + return { + "embedding": embedding, + "knowledge": {"chunk_size": 1024, "rerank": False}, + } + + def test_invalid_backend_raises(self, tmp_path): + cfg = self._cfg(backend="not-a-backend") + with pytest.raises(ValueError, match="backend must be"): + _build_kb_from_config(cfg, tmp_path) + + def test_api_backend_without_base_url_raises(self, tmp_path, monkeypatch): + monkeypatch.setenv("EMBEDDING_API_KEY", "k") + cfg = self._cfg(backend="api", model="m", base_url="") + with pytest.raises(ValueError, match="requires embedding.base_url"): + _build_kb_from_config(cfg, tmp_path) + + def test_api_backend_without_api_key_raises(self, tmp_path, monkeypatch): + # Credentials come from the EMBEDDING_API_KEY env var, never on disk. + monkeypatch.delenv("EMBEDDING_API_KEY", raising=False) + cfg = self._cfg(backend="api", model="m", base_url="http://x") + with pytest.raises(ValueError, match="requires the EMBEDDING_API_KEY"): + _build_kb_from_config(cfg, tmp_path) diff --git a/tests/test_episodic_integration.py b/tests/test_episodic_integration.py new file mode 100644 index 0000000..bb7d785 --- /dev/null +++ b/tests/test_episodic_integration.py @@ -0,0 +1,149 @@ +"""End-to-end episodic-memory integration: heartbeat → session_memory. + +The full episodic chain with nothing faked: a real Claude transcript on +disk → ``TraceCollector`` (the heartbeat) → the ``MemoryExtractor`` consumer +built by the shared ``memory.episodic_consumers`` wiring → mechanically tagged +turns embedded into the ``session_memory`` collection of a real +``KnowledgeBase`` → retrieved by ``kb.search``. + +Marked ``integration``: loads the local embedder, so it's deselected from the +fast suite (``-m 'not integration'``). +""" + +import json + +import pytest + +from dsagt.knowledge import KnowledgeBase +from dsagt.memory import SESSION_MEMORY_COLLECTION +from dsagt.traces import _transcript_dir, make_trace_collector + + +def _asst(ts, text): + return { + "type": "assistant", + "timestamp": ts, + "message": { + "role": "assistant", + "model": "claude-x", + "content": [{"type": "text", "text": text}], + "usage": {"input_tokens": 5, "output_tokens": 2}, + }, + } + + +def _user(ts, text, uuid): + return { + "type": "user", + "timestamp": ts, + "uuid": uuid, + "message": {"role": "user", "content": text}, + } + + +@pytest.mark.integration +def test_heartbeat_indexes_turn_into_session_memory(tmp_path): + project_dir = tmp_path / "proj" + (project_dir / ".dsagt").mkdir(parents=True) + + # A real KB with the local embedder (the project's session_memory lives here). + kb = KnowledgeBase(index_dir=project_dir / "kb_index", default_embedder="local") + + # Episodic enabled — exactly what `dsagt init --episodic` writes — fed + # through the server's real subscriber builder. + from dsagt.memory import episodic_consumers + + config = {"episodic": {"enabled": True}} + subs = episodic_consumers(config, kb, project_dir, "proj:s") + assert subs and subs[0].name == "memory" + + # A real Claude transcript with one clearly fact-bearing completed turn, + # then an open turn so the first is a "completed" candidate for the tick. + proot = tmp_path / "projects" + tdir = _transcript_dir(project_dir, proot) + tdir.mkdir(parents=True) + transcript = tdir / "sess.jsonl" + records = [ + _user( + "2026-06-28T15:00:00.000Z", + "Run fastp on the reads at a Q30 quality threshold.", + "u1", + ), + _asst( + "2026-06-28T15:00:01.000Z", + "Done — fastp filtered the reads at Q30; 92% of reads passed and were written to clean.fq.gz.", + ), + _user("2026-06-28T15:00:02.000Z", "what next?", "u2"), # opens turn 2 + ] + with open(transcript, "w") as fh: + for r in records: + fh.write(json.dumps(r) + "\n") + + tracking_uri = f"sqlite:///{tmp_path / 'mlflow.db'}" + collector = make_trace_collector( + "claude", + project_dir, + "proj", + "proj:s", + tracking_uri, + projects_root=proot, + extra_consumers=subs, + ) + + # include_last=True flushes both turns; the memory consumer indexes them. + n = collector.collect(include_last=True) + assert n >= 1 + + # Both consumers advanced their own marks independently. + assert collector._load_acks("mlflow") # observability sink + assert collector._load_acks("memory") # episodic consumer + + # A mechanically-indexed block-chunk is retrievable from session_memory, + # tagged with its producer + turn from this session. + hits = kb.search( + "fastp quality filtering Q30", collection=SESSION_MEMORY_COLLECTION, top_k=5 + ) + assert hits, "expected at least one indexed chunk in session_memory" + chunk = hits[0]["chunk"] + text = chunk["text"].lower() + assert "fastp" in text or "q30" in text or "qc" in text + meta = chunk["metadata"] + assert meta.get("source_type") == "turn" + assert meta.get("session_id") == "proj:s" + assert meta.get("producer") in {"user", "llm", "tool"} + assert meta.get("turn_id") # groups chunks back into their turn + + kb.close() + + +@pytest.mark.integration +def test_where_document_regex_and_contains_filter_session_memory(tmp_path): + """The metadata-filter + regex strategy: a document-text filter narrows the + candidate pool before semantic ranking, end-to-end through a real KB.""" + kb = KnowledgeBase(index_dir=tmp_path / "kb_index", default_embedder="local") + kb.add_entries( + texts=["fastp filtered reads at Q30", "bowtie2 aligned the reads"], + collection=SESSION_MEMORY_COLLECTION, + metadatas=[{"producer": "llm"}, {"producer": "llm"}], + ) + + # $regex narrows to the fastp chunk (case-insensitive) despite both being + # about "reads"; the metadata-filter path is dense-only but still ranks. + hits = kb.search( + "reads", + collection=SESSION_MEMORY_COLLECTION, + top_k=5, + where_document={"$regex": "(?i)FASTP"}, + ) + assert [h["chunk"]["text"] for h in hits] == ["fastp filtered reads at Q30"] + + # $contains is a case-sensitive substring. + hits2 = kb.search( + "reads", + collection=SESSION_MEMORY_COLLECTION, + top_k=5, + where_document={"$contains": "bowtie2"}, + ) + assert [h["chunk"]["text"] for h in hits2] == ["bowtie2 aligned the reads"] + + kb.close() diff --git a/tests/test_extraction.py b/tests/test_extraction.py deleted file mode 100644 index 23177e9..0000000 --- a/tests/test_extraction.py +++ /dev/null @@ -1,299 +0,0 @@ -""" -Tests for memory extraction (post-proxy). - -Conversation history now comes from MLflow traces, not a local -``session_log.jsonl``. Tests inject ``exchanges=[...]`` directly into -``extract_session`` so they don't need to mock the MLflow SDK; the -MLflow-trace-to-exchange formatter is exercised by its own unit tests -on ``_trace_to_exchange``. -""" - -import json -from unittest.mock import MagicMock, patch - -import numpy as np -import pytest - -from dsagt.memory import ( - EPISODIC_COLLECTION as COLLECTION_NAME, - _trace_to_exchange, - build_extraction_prompt, - extract_session, - parse_extraction_response, -) -from dsagt.knowledge import KnowledgeBase - - -def fake_embed(texts: list[str]) -> np.ndarray: - dim = 8 - vecs = np.zeros((len(texts), dim), dtype=np.float32) - for i, t in enumerate(texts): - rng = np.random.RandomState(hash(t) & 0xFFFFFFFF) - vecs[i] = rng.randn(dim).astype(np.float32) - norms = np.linalg.norm(vecs, axis=1, keepdims=True) - norms[norms == 0] = 1 - return vecs / norms - - -# --------------------------------------------------------------------------- -# build_extraction_prompt -# --------------------------------------------------------------------------- - -class TestBuildExtractionPrompt: - - def _make_exchanges(self): - return [ - { - "timestamp": "2024-01-15T10:30:00Z", - "model": "claude-sonnet-4-20250514", - "new_messages": [{"role": "user", "content": "Run fastp on sample1.fq.gz"}], - "response": [ - {"type": "text", "text": "I'll run fastp with Q20 filtering."}, - {"type": "tool_use", "name": "bash", - "input": {"cmd": "fastp -q 20 --in1 sample1.fq.gz"}}, - ], - }, - { - "timestamp": "2024-01-15T10:31:00Z", - "model": "claude-sonnet-4-20250514", - "new_messages": [ - {"role": "user", "content": [ - {"type": "tool_result", "tool_use_id": "t1", - "content": "98% reads passed"}, - ]}, - ], - "response": [ - {"type": "text", "text": "Filtering complete. 98% of reads passed Q20."}, - ], - }, - ] - - def test_includes_conversation(self): - prompt = build_extraction_prompt(self._make_exchanges()) - assert "fastp" in prompt - assert "sample1.fq.gz" in prompt - assert "98% reads passed" in prompt - - def test_includes_categories(self): - prompt = build_extraction_prompt(self._make_exchanges()) - assert "quality_control" in prompt - assert "data_management" in prompt - - def test_custom_categories_merged(self): - custom = {"my_category": "custom stuff"} - prompt = build_extraction_prompt(self._make_exchanges(), categories=custom) - assert "my_category" in prompt - assert "custom stuff" in prompt - assert "quality_control" in prompt # stock still present - - def test_output_format_described(self): - prompt = build_extraction_prompt(self._make_exchanges()) - assert '"facts"' in prompt - assert '"summary"' in prompt - assert '"insights"' in prompt - - def test_empty_exchanges(self): - prompt = build_extraction_prompt([]) - assert "facts" in prompt # prompt structure still present - - -# --------------------------------------------------------------------------- -# parse_extraction_response -# --------------------------------------------------------------------------- - -class TestParseExtractionResponse: - - def test_valid_json(self): - response = json.dumps({ - "facts": [{"text": "fastp used Q20", "category": "quality_control"}], - "summary": "Ran fastp on sample1.", - "insights": [{"text": "Q20 is sufficient for isolates", - "category": "quality_control"}], - }) - result = parse_extraction_response(response) - assert len(result["facts"]) == 1 - assert result["summary"] == "Ran fastp on sample1." - assert len(result["insights"]) == 1 - - def test_strips_markdown_fences(self): - response = "```json\n" + json.dumps({ - "facts": [], "summary": "test", "insights": [], - }) + "\n```" - result = parse_extraction_response(response) - assert result["summary"] == "test" - - def test_missing_fields_default(self): - result = parse_extraction_response("{}") - assert result["facts"] == [] - assert result["summary"] == "" - assert result["insights"] == [] - - -# --------------------------------------------------------------------------- -# _trace_to_exchange — MLflow trace row → extraction exchange dict -# --------------------------------------------------------------------------- - -class TestTraceToExchange: - - def test_anthropic_response_shape(self): - """Anthropic-shape response: ``content`` is already a block list.""" - row = { - "request_time": "2024-05-01T12:00:00Z", - "trace_id": "tr-abc", - "request": json.dumps({ - "model": "claude-sonnet-4-5", - "messages": [{"role": "user", "content": "hi"}], - }), - "response": json.dumps({ - "content": [{"type": "text", "text": "hello"}], - }), - } - ex = _trace_to_exchange(row) - assert ex is not None - assert ex["new_messages"] == [{"role": "user", "content": "hi"}] - assert ex["response"] == [{"type": "text", "text": "hello"}] - assert ex["trace_id"] == "tr-abc" - assert ex["model"] == "claude-sonnet-4-5" - - def test_openai_response_shape_with_tool_calls(self): - """OpenAI-shape: choices[].message.content + tool_calls.""" - row = { - "request_time": "2024-05-01T12:00:00Z", - "trace_id": "tr-xyz", - "request": json.dumps({ - "model": "gpt-4o", - "messages": [{"role": "user", "content": "run it"}], - }), - "response": json.dumps({ - "choices": [{ - "message": { - "content": "Sure", - "tool_calls": [{ - "id": "call_1", - "function": { - "name": "bash", - "arguments": json.dumps({"cmd": "ls"}), - }, - }], - }, - }], - }), - } - ex = _trace_to_exchange(row) - assert ex is not None - types = [b["type"] for b in ex["response"]] - assert "text" in types - assert "tool_use" in types - tool_use = next(b for b in ex["response"] if b["type"] == "tool_use") - assert tool_use["name"] == "bash" - assert tool_use["input"] == {"cmd": "ls"} - - def test_unrecognised_shape_returns_none(self): - """Non-LLM-call traces (kb.* / tool.execute spans) get skipped.""" - row = { - "request_time": "2024-05-01T12:00:00Z", - "trace_id": "tr-kb", - "request": json.dumps({"query": "stuff"}), # no messages - "response": "{}", - } - assert _trace_to_exchange(row) is None - - def test_handles_missing_response(self): - row = { - "trace_id": "tr-empty", - "request": json.dumps({ - "messages": [{"role": "user", "content": "x"}], - }), - "response": None, - } - ex = _trace_to_exchange(row) - assert ex is not None - assert ex["response"] == [] - - -# --------------------------------------------------------------------------- -# extract_session (end-to-end with injected exchanges) -# --------------------------------------------------------------------------- - -class TestExtractSession: - - def _mock_llm_response(self): - return json.dumps({ - "facts": [ - {"text": "fastp was run with Q20 on sample1", - "category": "quality_control"}, - {"text": "98% of reads passed filtering", "category": "results"}, - ], - "summary": "Ran quality filtering on sample1 using fastp with Q20 threshold.", - "insights": [ - {"text": "Q20 filtering is sufficient for high-quality isolate data", - "category": "quality_control"}, - ], - }) - - def test_extracts_and_stores(self, tmp_path): - """End-to-end: injected exchanges → mocked LLM → KB write.""" - exchanges = [ - { - "timestamp": "2024-01-15T10:30:00Z", - "trace_id": "tr-1", - "model": "m", - "new_messages": [{"role": "user", "content": "run fastp"}], - "response": [{"type": "text", "text": "Running fastp."}], - }, - ] - with patch("dsagt.knowledge._make_embedder") as mock_make: - mock_embedder = MagicMock() - mock_embedder.embed = fake_embed - mock_make.return_value = mock_embedder - - kb = KnowledgeBase(index_dir=tmp_path / "kb") - with patch("dsagt.memory.call_extraction_llm") as mock_llm: - mock_llm.return_value = self._mock_llm_response() - result = extract_session( - project_name="proj", - kb=kb, - api_key="test-key", - session_id="test-session", - exchanges=exchanges, - ) - assert result["status"] == "ok" - assert result["facts"] == 2 - assert result["insights"] == 1 - assert result["summary"] == 1 - kb.close() - - def test_empty_exchanges_returns_status_empty(self, tmp_path): - """No exchanges → returns status=empty, no LLM call.""" - with patch("dsagt.knowledge._make_embedder") as mock_make: - mock_embedder = MagicMock() - mock_embedder.embed = fake_embed - mock_make.return_value = mock_embedder - - kb = KnowledgeBase(index_dir=tmp_path / "kb") - with patch("dsagt.memory.call_extraction_llm") as mock_llm: - result = extract_session( - project_name="proj", - kb=kb, - api_key="test-key", - session_id="test-session", - exchanges=[], - ) - mock_llm.assert_not_called() - assert result["status"] == "empty" - kb.close() - - def test_missing_session_id_returns_empty(self, tmp_path): - """``session_id`` is required when ``exchanges`` isn't supplied.""" - with patch("dsagt.knowledge._make_embedder") as mock_make: - mock_embedder = MagicMock() - mock_embedder.embed = fake_embed - mock_make.return_value = mock_embedder - - kb = KnowledgeBase(index_dir=tmp_path / "kb") - result = extract_session( - project_name="proj", kb=kb, api_key="test-key", - ) - assert result["status"] == "empty" - assert result.get("reason") == "no_session_id" - kb.close() diff --git a/tests/test_info.py b/tests/test_info.py index 6b1fd22..fba023b 100644 --- a/tests/test_info.py +++ b/tests/test_info.py @@ -22,34 +22,30 @@ ) -def _metadata(*, session: str, agent: str | None, - in_t: int, out_t: int) -> dict: +def _metadata(*, session: str, agent: str | None, in_t: int, out_t: int) -> dict: """Build a trace_metadata dict in MLflow's on-disk shape.""" md = {"mlflow.trace.session": session} if agent is not None: md["dsagt.agent"] = agent if in_t or out_t: - md["mlflow.trace.tokenUsage"] = json.dumps({ - "input_tokens": in_t, - "output_tokens": out_t, - "total_tokens": in_t + out_t, - }) + md["mlflow.trace.tokenUsage"] = json.dumps( + { + "input_tokens": in_t, + "output_tokens": out_t, + "total_tokens": in_t + out_t, + } + ) return md -def _spans_for(service_name: str | None) -> list: - """Build the ``spans`` column entry: one root span carrying service.name. +def _tags_for(source: str | None) -> dict: + """Build the ``tags`` column entry carrying the ``dsagt.source`` bucket. - ``_source_from_spans`` reads ``service.name`` off the root span's - attributes (MLflow's OTLP receiver flows the OTel resource attribute - through to span attributes), so a single SimpleNamespace is enough - for the report-side test. ``None`` → empty spans list, source falls - through to ``"unknown"``. + Internal debug traces are bucketed from this tag (the MCP tool category / + ``execution``); ``None`` → no tag, so the trace falls back to the + ``dsagt.agent`` metadata (agent traces) or ``"unknown"``. """ - from types import SimpleNamespace - if service_name is None: - return [] - return [SimpleNamespace(attributes={"service.name": service_name})] + return {"dsagt.source": source} if source else {} def _traces_df(rows: list[dict]) -> pd.DataFrame: @@ -66,6 +62,7 @@ def _traces_df(rows: list[dict]) -> pd.DataFrame: # _tokens / _fmt_count / _is_error primitives # --------------------------------------------------------------------------- + def test_tokens_missing_returns_zeros(): assert _tokens({}) == (0, 0) @@ -79,18 +76,29 @@ def test_tokens_extracts_input_output(): assert _tokens(md) == (123, 45) -@pytest.mark.parametrize("n,expected", [ - (0, "0"), (999, "999"), (1000, "1.0k"), - (12345, "12.3k"), (1_500_000, "1.5M"), -]) +@pytest.mark.parametrize( + "n,expected", + [ + (0, "0"), + (999, "999"), + (1000, "1.0k"), + (12345, "12.3k"), + (1_500_000, "1.5M"), + ], +) def test_fmt_count(n, expected): assert _fmt_count(n) == expected -@pytest.mark.parametrize("state,expected", [ - ("OK", False), ("ERROR", True), - ("TraceState.OK", False), ("TraceState.ERROR", True), -]) +@pytest.mark.parametrize( + "state,expected", + [ + ("OK", False), + ("ERROR", True), + ("TraceState.OK", False), + ("TraceState.ERROR", True), + ], +) def test_is_error_handles_enum_reprs(state, expected): assert _is_error(state) is expected @@ -99,11 +107,14 @@ def test_is_error_handles_enum_reprs(state, expected): # _report end-to-end # --------------------------------------------------------------------------- + @pytest.fixture def config(): return { "agent": "claude", - "llm": {"model": "claude-haiku-test"}, + # BYOA: dsagt no longer records the agent's LLM model; the info + # header surfaces the embedding model dsagt configures. + "embedding": {"model": "bge-test"}, } @@ -114,54 +125,68 @@ def test_report_empty_traces(config): assert r["by_session"] == [] assert r["errors"] == [] assert r["agent"] == "claude" - assert r["model"] == "claude-haiku-test" + assert r["model"] == "bge-test" def test_report_aggregates_by_source_and_session(config): - """Source bucketing comes from the root span's ``service.name``. + """Source bucketing: agent traces from ``dsagt.agent`` metadata, internal + traces from the ``dsagt.source`` tag. - Three claude-code agent traces (one of which errored) + one - knowledge-server trace, across two sessions. Sums + bucket counts - match expected totals. + Three claude agent traces (one of which errored) + one internal + knowledge-tool trace, across two sessions. Sums + bucket counts match. """ - df = _traces_df([ - { - "trace_id": "t1", - "state": "OK", - "request_time": 100, - "trace_metadata": _metadata( - session="sess-A", agent="claude", in_t=1000, out_t=100, - ), - "spans": _spans_for("claude-code"), - }, - { - "trace_id": "t2", - "state": "OK", - "request_time": 200, - "trace_metadata": _metadata( - session="sess-A", agent="claude", in_t=500, out_t=50, - ), - "spans": _spans_for("claude-code"), - }, - { - "trace_id": "t3", - "state": "OK", - "request_time": 150, - "trace_metadata": _metadata( - session="sess-A", agent="claude", in_t=200, out_t=0, - ), - "spans": _spans_for("dsagt-knowledge-server"), - }, - { - "trace_id": "t4", - "state": "ERROR", - "request_time": 300, - "trace_metadata": _metadata( - session="sess-B", agent="claude", in_t=800, out_t=20, - ), - "spans": _spans_for("claude-code"), - }, - ]) + df = _traces_df( + [ + { + "trace_id": "t1", + "state": "OK", + "request_time": 100, + "trace_metadata": _metadata( + session="sess-A", + agent="claude", + in_t=1000, + out_t=100, + ), + "tags": _tags_for(None), + }, + { + "trace_id": "t2", + "state": "OK", + "request_time": 200, + "trace_metadata": _metadata( + session="sess-A", + agent="claude", + in_t=500, + out_t=50, + ), + "tags": _tags_for(None), + }, + { + "trace_id": "t3", + "state": "OK", + "request_time": 150, + "trace_metadata": _metadata( + session="sess-A", + agent="claude", + in_t=200, + out_t=0, + ), + "tags": _tags_for("knowledge"), + }, + { + "trace_id": "t4", + "state": "ERROR", + "request_time": 300, + "trace_metadata": _metadata( + session="sess-B", + agent="claude", + in_t=800, + out_t=20, + ), + "tags": _tags_for(None), + }, + ] + ) r = _report("proj", config, df) @@ -171,12 +196,15 @@ def test_report_aggregates_by_source_and_session(config): assert r["output_tokens"] == 170 sources = {row["source"]: row for row in r["by_source"]} - assert sources["claude-code"]["traces"] == 3 - assert sources["claude-code"]["input_tokens"] == 2300 - assert sources["claude-code"]["output_tokens"] == 170 - assert sources["claude-code"]["errors"] == 1 - assert sources["dsagt-knowledge-server"]["traces"] == 1 - assert sources["dsagt-knowledge-server"]["errors"] == 0 + # t1/t2/t4 have no dsagt.source tag → bucket by dsagt.agent ("claude"). + assert sources["claude"]["traces"] == 3 + assert sources["claude"]["input_tokens"] == 2300 + assert sources["claude"]["output_tokens"] == 170 + assert sources["claude"]["errors"] == 1 + # t3 carries dsagt.source=knowledge → its own bucket even though it also + # has dsagt.agent (the tag wins for internal traces). + assert sources["knowledge"]["traces"] == 1 + assert sources["knowledge"]["errors"] == 0 assert [s["session"] for s in r["by_session"]] == ["sess-B", "sess-A"] sess_a = next(s for s in r["by_session"] if s["session"] == "sess-A") @@ -187,23 +215,110 @@ def test_report_aggregates_by_source_and_session(config): assert len(r["errors"]) == 1 err = r["errors"][0] assert err["session"] == "sess-B" - assert err["source"] == "claude-code" + assert err["source"] == "claude" assert err["trace_id"] == "t4" +def test_episodic_and_code_use_bucket_as_internal_sources(config): + """The background emitters carry their own dsagt.source and are counted as + internal/debug, not agent turns.""" + from dsagt.commands.info import _INTERNAL_SOURCES + + assert {"episodic", "code_use"} <= _INTERNAL_SOURCES + + df = _traces_df( + [ + { + "trace_id": "a1", + "state": "OK", + "request_time": 100, + "trace_metadata": _metadata( + session="s", agent="claude", in_t=900, out_t=90 + ), + "tags": _tags_for(None), + }, + { + "trace_id": "e1", + "state": "OK", + "request_time": 110, + "trace_metadata": _metadata(session="s", agent=None, in_t=0, out_t=0), + "tags": _tags_for("episodic"), + }, + { + "trace_id": "c1", + "state": "OK", + "request_time": 120, + "trace_metadata": _metadata(session="s", agent=None, in_t=0, out_t=0), + "tags": _tags_for("code_use"), + }, + ] + ) + r = _report("proj", config, df) + sources = {row["source"]: row for row in r["by_source"]} + assert sources["episodic"]["traces"] == 1 + assert sources["code_use"]["traces"] == 1 + # Agent-vs-internal split (the headline): 1 agent turn, 2 internal. + agent_traces = sum( + row["traces"] + for row in r["by_source"] + if row["source"] not in _INTERNAL_SOURCES and row["source"] != "unknown" + ) + assert agent_traces == 1 + assert r["total_traces"] - agent_traces == 2 + + +def test_agent_vs_internal_split_counts_unknown_as_internal(config): + """An orphaned/uncategorized trace (no source, no agent) is bookkeeping — + it counts as internal/debug, never as an agent turn.""" + from dsagt.commands.info import _INTERNAL_SOURCES + + df = _traces_df( + [ + { + "trace_id": "a1", + "state": "OK", + "request_time": 100, + "trace_metadata": _metadata( + session="s", agent="claude", in_t=10, out_t=1 + ), + "tags": _tags_for(None), + }, + { # neither dsagt.source nor dsagt.agent → "unknown" + "trace_id": "u1", + "state": "OK", + "request_time": 110, + "trace_metadata": _metadata(session="s", agent=None, in_t=0, out_t=0), + "tags": _tags_for(None), + }, + ] + ) + r = _report("proj", config, df) + agent_traces = sum( + row["traces"] + for row in r["by_source"] + if row["source"] not in _INTERNAL_SOURCES and row["source"] != "unknown" + ) + assert agent_traces == 1 # not 2 — the unknown is internal/debug + + def test_report_missing_source_falls_back_to_unknown(config): - """No spans column / no service.name → bucket is "unknown".""" - df = _traces_df([ - { - "trace_id": "t1", - "state": "OK", - "request_time": 100, - "trace_metadata": _metadata( - session="sess-A", agent=None, in_t=0, out_t=0, - ), - "spans": _spans_for(None), - }, - ]) + """No dsagt.source tag and no dsagt.agent → bucket is "unknown".""" + df = _traces_df( + [ + { + "trace_id": "t1", + "state": "OK", + "request_time": 100, + "trace_metadata": _metadata( + session="sess-A", + agent=None, + in_t=0, + out_t=0, + ), + "tags": _tags_for(None), + }, + ] + ) r = _report("proj", config, df) assert r["by_source"][0]["source"] == "unknown" assert r["by_session"][0]["agent"] == "-" @@ -213,6 +328,7 @@ def test_report_missing_source_falls_back_to_unknown(config): # Config source tracking (config / shell / unresolved) # --------------------------------------------------------------------------- + def test_mask_secret_short_value(): assert _mask_secret("short") == "***" assert _mask_secret("exactlytwelv") == "***" @@ -223,13 +339,12 @@ def test_mask_secret_long_value(): def _write_project(tmp_path, monkeypatch, raw_yaml: str): - """Register a temp project with the given dsagt_config.yaml content.""" - import yaml as _yaml + """Register a temp project with the given .dsagt/config.yaml content.""" from dsagt.session import register_project pdir = tmp_path / "proj" - pdir.mkdir() - (pdir / "dsagt_config.yaml").write_text(raw_yaml) + (pdir / ".dsagt").mkdir(parents=True) + (pdir / ".dsagt" / "config.yaml").write_text(raw_yaml) registry_dir = tmp_path / "registry" registry_dir.mkdir() @@ -242,20 +357,30 @@ def _write_project(tmp_path, monkeypatch, raw_yaml: str): def test_config_sources_classifies_shell(tmp_path, monkeypatch): """${VAR} resolves from os.environ — that's the only source for user-provided values now (no .env file is consulted).""" - _write_project(tmp_path, monkeypatch, - "project: proj\nagent: goose\nllm:\n provider: ${LLM_PROVIDER}\n model: ${LLM_MODEL}\n") + _write_project( + tmp_path, + monkeypatch, + "project: proj\nagent: goose\nllm:\n provider: ${LLM_PROVIDER}\n model: ${LLM_MODEL}\n", + ) monkeypatch.setenv("LLM_PROVIDER", "openai") monkeypatch.setenv("LLM_MODEL", "gpt-4") rows = {r["path"]: r for r in _config_sources("proj")} - assert rows["llm.provider"] == {"path": "llm.provider", "value": "openai", "source": "shell"} + assert rows["llm.provider"] == { + "path": "llm.provider", + "value": "openai", + "source": "shell", + } assert rows["llm.model"]["value"] == "gpt-4" assert rows["llm.model"]["source"] == "shell" def test_config_sources_classifies_unresolved(tmp_path, monkeypatch): - _write_project(tmp_path, monkeypatch, - "project: proj\nagent: goose\nllm:\n provider: ${LLM_PROVIDER}\n") + _write_project( + tmp_path, + monkeypatch, + "project: proj\nagent: goose\nllm:\n provider: ${LLM_PROVIDER}\n", + ) monkeypatch.delenv("LLM_PROVIDER", raising=False) rows = {r["path"]: r for r in _config_sources("proj")} @@ -264,16 +389,24 @@ def test_config_sources_classifies_unresolved(tmp_path, monkeypatch): def test_config_sources_classifies_literal(tmp_path, monkeypatch): - _write_project(tmp_path, monkeypatch, - "project: proj\nagent: goose\nproxy:\n port: 4000\n") + _write_project( + tmp_path, monkeypatch, "project: proj\nagent: goose\nmlflow:\n port: 4000\n" + ) rows = {r["path"]: r for r in _config_sources("proj")} - assert rows["proxy.port"] == {"path": "proxy.port", "value": "4000", "source": "config"} + assert rows["mlflow.port"] == { + "path": "mlflow.port", + "value": "4000", + "source": "config", + } def test_config_sources_masks_api_key(tmp_path, monkeypatch): - _write_project(tmp_path, monkeypatch, - "project: proj\nagent: goose\nllm:\n api_key: ${LLM_API_KEY}\n") + _write_project( + tmp_path, + monkeypatch, + "project: proj\nagent: goose\nllm:\n api_key: ${LLM_API_KEY}\n", + ) monkeypatch.setenv("LLM_API_KEY", "sk-1234567890abcdef") rows = {r["path"]: r for r in _config_sources("proj")} @@ -282,13 +415,15 @@ def test_config_sources_masks_api_key(tmp_path, monkeypatch): def test_config_sources_skips_internal_sections(tmp_path, monkeypatch): - _write_project(tmp_path, monkeypatch, - "project: proj\nagent: goose\nknowledge:\n chunk_size: 1024\nextraction:\n threshold: 0\ncategories:\n qc: stuff\n") + _write_project( + tmp_path, + monkeypatch, + "project: proj\nagent: goose\nknowledge:\n chunk_size: 1024\n rerank: false\nskills:\n populate_native: true\n", + ) paths = {r["path"] for r in _config_sources("proj")} assert "knowledge.chunk_size" not in paths - assert "extraction.threshold" not in paths - assert "categories.qc" not in paths + assert "skills.populate_native" not in paths assert "project" in paths @@ -296,6 +431,7 @@ def test_config_sources_skips_internal_sections(tmp_path, monkeypatch): # _kb_collections — read chunks.jsonl, count entries, break down by source # --------------------------------------------------------------------------- + def _write_chunk(path, text="x", source=None): """Append one chunks.jsonl line with optional metadata.source.""" md = {"chunk_index": 0} @@ -307,6 +443,7 @@ def _write_chunk(path, text="x", source=None): def test_kb_collections_empty_when_no_kb_index(tmp_path): from dsagt.commands.info import _kb_collections + assert _kb_collections(tmp_path) == [] @@ -314,14 +451,14 @@ def test_kb_collections_counts_chunks_and_breaks_down_sources(tmp_path): from dsagt.commands.info import _kb_collections kb = tmp_path / "kb_index" - (kb / "tools").mkdir(parents=True) + (kb / "codes").mkdir(parents=True) (kb / "skills").mkdir(parents=True) (kb / "research").mkdir(parents=True) # Tools: 2 bundled, 1 project for _ in range(2): - _write_chunk(kb / "tools" / "chunks.jsonl", source="bundled") - _write_chunk(kb / "tools" / "chunks.jsonl", source="project") + _write_chunk(kb / "codes" / "chunks.jsonl", source="bundled") + _write_chunk(kb / "codes" / "chunks.jsonl", source="project") # Skills: 1 bundled _write_chunk(kb / "skills" / "chunks.jsonl", source="bundled") @@ -331,8 +468,8 @@ def test_kb_collections_counts_chunks_and_breaks_down_sources(tmp_path): _write_chunk(kb / "research" / "chunks.jsonl") rows = {r["collection"]: r for r in _kb_collections(tmp_path)} - assert rows["tools"] == { - "collection": "tools", + assert rows["codes"] == { + "collection": "codes", "chunks": 3, "by_source": {"bundled": 2, "project": 1}, } @@ -355,9 +492,11 @@ def test_kb_collections_skips_dirs_without_chunks_jsonl(tmp_path): # _kb_retrieval — pull kb.search spans out of a synthetic traces frame # --------------------------------------------------------------------------- + def _trace_with_kb_spans(session: str, hits_per_search: list[int]) -> dict: """Build a single trace row carrying one kb.search span per hits entry.""" from types import SimpleNamespace + spans = [ SimpleNamespace(name="kb.search", attributes={"hits": h}) for h in hits_per_search @@ -375,6 +514,7 @@ def _trace_with_kb_spans(session: str, hits_per_search: list[int]) -> dict: def test_kb_retrieval_empty_when_no_traces(): from dsagt.commands.info import _kb_retrieval + assert _kb_retrieval(None) == [] assert _kb_retrieval(pd.DataFrame()) == [] @@ -382,11 +522,13 @@ def test_kb_retrieval_empty_when_no_traces(): def test_kb_retrieval_aggregates_searches_and_hits_by_session(): from dsagt.commands.info import _kb_retrieval - df = pd.DataFrame([ - _trace_with_kb_spans("sess-A", hits_per_search=[3, 5]), - _trace_with_kb_spans("sess-A", hits_per_search=[2]), - _trace_with_kb_spans("sess-B", hits_per_search=[7]), - ]) + df = pd.DataFrame( + [ + _trace_with_kb_spans("sess-A", hits_per_search=[3, 5]), + _trace_with_kb_spans("sess-A", hits_per_search=[2]), + _trace_with_kb_spans("sess-B", hits_per_search=[7]), + ] + ) rows = {r["session"]: r for r in _kb_retrieval(df)} assert rows["sess-A"] == {"session": "sess-A", "searches": 3, "hits": 10} assert rows["sess-B"] == {"session": "sess-B", "searches": 1, "hits": 7} @@ -396,13 +538,17 @@ def test_kb_retrieval_handles_missing_session(): from dsagt.commands.info import _kb_retrieval from types import SimpleNamespace - df = pd.DataFrame([{ - "trace_id": "t1", - "state": "OK", - "request_time": 1, - "trace_metadata": {}, # no mlflow.trace.session - "spans": [SimpleNamespace(name="kb.search", attributes={"hits": 4})], - }]) + df = pd.DataFrame( + [ + { + "trace_id": "t1", + "state": "OK", + "request_time": 1, + "trace_metadata": {}, # no mlflow.trace.session + "spans": [SimpleNamespace(name="kb.search", attributes={"hits": 4})], + } + ] + ) rows = _kb_retrieval(df) assert len(rows) == 1 assert rows[0]["session"] == "(no-session)" @@ -413,16 +559,20 @@ def test_kb_retrieval_ignores_non_kb_spans(): from dsagt.commands.info import _kb_retrieval from types import SimpleNamespace - df = pd.DataFrame([{ - "trace_id": "t1", - "state": "OK", - "request_time": 1, - "trace_metadata": {"mlflow.trace.session": "s"}, - "spans": [ - SimpleNamespace(name="kb.embed", attributes={}), - SimpleNamespace(name="registry.save_tool_spec", attributes={}), - ], - }]) + df = pd.DataFrame( + [ + { + "trace_id": "t1", + "state": "OK", + "request_time": 1, + "trace_metadata": {"mlflow.trace.session": "s"}, + "spans": [ + SimpleNamespace(name="kb.embed", attributes={}), + SimpleNamespace(name="registry.save_code_spec", attributes={}), + ], + } + ] + ) assert _kb_retrieval(df) == [] @@ -430,6 +580,7 @@ def test_kb_retrieval_ignores_non_kb_spans(): # _project_created — best-effort project-start date # --------------------------------------------------------------------------- + def test_project_created_returns_iso_date_for_existing_dir(tmp_path): from dsagt.commands.info import _project_created @@ -443,4 +594,5 @@ def test_project_created_returns_iso_date_for_existing_dir(tmp_path): def test_project_created_returns_none_when_dir_missing(tmp_path): from dsagt.commands.info import _project_created + assert _project_created(tmp_path / "does-not-exist") is None diff --git a/tests/test_integration.py b/tests/test_integration.py index 6cdfbe7..368cc04 100644 --- a/tests/test_integration.py +++ b/tests/test_integration.py @@ -9,8 +9,8 @@ Useful when debugging embedding failures, because the smoke test can't tell you whether a failure came from the network or the KB pipeline. - **TestLLMEndpoint** — the raw LLM API returns a completion. Same - reasoning: bisects "is the upstream reachable" from "is the proxy wired - correctly". + reasoning: bisects "is the upstream reachable" from "is the KB/agent + layer wired correctly". - **TestConfigChain** — dsagt init → static_agent_record → dynamic_agent_record → agent_env without touching the network, since config bugs can masquerade as runtime errors that are painful to diagnose via smoke. @@ -25,7 +25,7 @@ import numpy as np import pytest -from dsagt.knowledge import APIEmbeddingClient +from dsagt.knowledge import APIEmbedder pytestmark = pytest.mark.integration @@ -34,12 +34,13 @@ # Embedding endpoint # --------------------------------------------------------------------------- + class TestEmbeddingEndpoint: """Validates the embedding API is reachable and returns valid vectors.""" def test_embedding_returns_vector(self, embedding_config): """Single text produces a numpy array with positive dimension.""" - client = APIEmbeddingClient( + client = APIEmbedder( model=embedding_config["model"], base_url=embedding_config["base_url"], api_key=embedding_config["api_key"], @@ -55,7 +56,7 @@ def test_embedding_returns_vector(self, embedding_config): def test_embedding_batch(self, embedding_config): """Multiple texts produce correct batch shape with consistent dim.""" - client = APIEmbeddingClient( + client = APIEmbedder( model=embedding_config["model"], base_url=embedding_config["base_url"], api_key=embedding_config["api_key"], @@ -77,14 +78,13 @@ def test_embedding_batch(self, embedding_config): # LLM endpoint # --------------------------------------------------------------------------- + class TestLLMEndpoint: """Validates the LLM API is reachable and returns valid responses. Posts to ``/chat/completions`` (OpenAI-shape) because .env's LLM_BASE_URL is an OpenAI-compatible gateway — Anthropic-format ``/v1/messages`` would - 404 on it. The proxy (commands/proxy_server.py) sets - ``use_chat_completions_url_for_anthropic_messages = True`` for the same - reason. + 404 on it. """ def test_llm_returns_response(self, llm_config): @@ -92,7 +92,11 @@ def test_llm_returns_response(self, llm_config): # Some gateways expect /v1 already in base_url (RC Chat / Ollama # OpenAI-compat), others don't (ai-incubator-api). Don't double it. base = llm_config["base_url"].rstrip("/") - url = f"{base}/chat/completions" if base.endswith("/v1") else f"{base}/v1/chat/completions" + url = ( + f"{base}/chat/completions" + if base.endswith("/v1") + else f"{base}/v1/chat/completions" + ) with httpx.Client(timeout=30.0) as client: response = client.post( url, @@ -119,12 +123,13 @@ def test_llm_returns_response(self, llm_config): # Config chain (no network) # --------------------------------------------------------------------------- + class TestConfigChain: """Validates the init → config generation → env var chain (BYOA).""" def test_init_generates_valid_config(self, tmp_path, monkeypatch): - """init_project creates a valid dsagt_config.yaml and returns - ``(pdir, mlflow_port)``.""" + """init_project creates a valid .dsagt/config.yaml and returns the + project dir (serverless — no port).""" import yaml from dsagt.session import init_project @@ -132,71 +137,18 @@ def test_init_generates_valid_config(self, tmp_path, monkeypatch): reg = {} monkeypatch.setattr("dsagt.session._load_registry", lambda: dict(reg)) monkeypatch.setattr("dsagt.session._save_registry", lambda r: reg.update(r)) - monkeypatch.setattr("dsagt.session.register_project", lambda n, p: reg.update({n: str(p)})) + monkeypatch.setattr( + "dsagt.session.register_project", lambda n, p: reg.update({n: str(p)}) + ) - pdir, port = init_project("test-proj", "claude") + pdir = init_project("test-proj", "claude") - assert isinstance(port, int) and port > 0 - config_path = pdir / "dsagt_config.yaml" + config_path = pdir / ".dsagt" / "config.yaml" assert config_path.exists() config = yaml.safe_load(config_path.read_text()) assert config["project"] == "test-proj" assert config["agent"] == "claude" - assert config["mlflow"]["port"] == port + assert "mlflow" not in config assert "embedding" in config - # BYOA: no llm: block in the user-facing YAML. + # BYOA: no llm: block in the config. assert "llm" not in config - - -@pytest.mark.skip(reason="proxy mode deferred to Phase 2") -class TestProxyConfigChain: - """Proxy-mode env-routing matrix. Re-enable when --proxy_traces is - restored in Phase 2 — these tests are the regression net for the - 'agent talks to upstream instead of the proxy' bug class.""" - - def test_agent_configs_have_correct_ports(self, integration_config, tmp_path): - from dsagt.agents import agent_env, dynamic_agent_record, static_agent_record - - working_dir = tmp_path / "workdir" - working_dir.mkdir() - - static_agent_record(integration_config, integration_config["agent"], working_dir) - env = agent_env(integration_config) - dynamic_agent_record(integration_config, env, working_dir) - # Phase 2: assert proxy URL lands in the per-agent MCP-config artifact - # (cline_mcp_settings.json / .roo/mcp.json / codex config.toml). - - def test_env_vars_chain_complete(self, integration_config): - from dsagt.agents import agent_env - - env = agent_env(integration_config) - assert env["DSAGT_PROJECT"] == integration_config["project"] - assert env["DSAGT_AGENT"] == integration_config["agent"] - if integration_config["agent"] == "claude": - proxy_port = integration_config["proxy"]["port"] - assert str(proxy_port) in env["ANTHROPIC_BASE_URL"] - - @pytest.mark.parametrize("agent", ["goose", "claude", "roo", "cline", "codex"]) - def test_agent_env_matrix(self, integration_config, agent): - from dsagt.agents import agent_env - - config = dict(integration_config) - config["agent"] = agent - proxy_port = config["proxy"]["port"] - model = config["llm"]["model"] - sentinel = "dsagt-proxy-forwarded-disable-direct-calls" - - env = agent_env(config) - assert env["DSAGT_AGENT"] == agent - - if agent in ("claude", "roo"): - assert env["ANTHROPIC_BASE_URL"] == f"http://localhost:{proxy_port}" - assert env["ANTHROPIC_MODEL"] == model - assert env["ANTHROPIC_API_KEY"] == sentinel - elif agent == "goose": - assert env["OPENAI_HOST"] == f"http://localhost:{proxy_port}" - assert env["GOOSE_PROVIDER"] == "openai" - assert env["OPENAI_API_KEY"] == sentinel - elif agent == "codex": - assert env["CODEX_HOME"].endswith(".codex-data") - assert env["OPENAI_API_KEY"] == sentinel diff --git a/tests/test_kb_search_filters.py b/tests/test_kb_search_filters.py index 73919ca..c9e7b4c 100644 --- a/tests/test_kb_search_filters.py +++ b/tests/test_kb_search_filters.py @@ -11,7 +11,7 @@ import pytest import mcp.types as types -from dsagt.commands.knowledge_server import create_knowledge_server +from dsagt.mcp.knowledge_tools import create_knowledge_server from mcp_helpers import call_tool_json as call_tool @@ -47,121 +47,188 @@ def server(mock_kb): # Handler: filter threading # --------------------------------------------------------------------------- -class TestSearchFilterThreading: - - def test_no_filters_no_where(self, server, mock_kb): - """Search without filter params doesn't pass where to kb.search.""" - call_tool(server, "kb_search", { - "query": "test", - "collection": "docs", - }) - mock_kb.search.assert_called_once_with( - query="test", - collection="docs", - top_k=5, - rerank=None, - ) +class TestSearchFilterThreading: def test_tool_name_filter_passed(self, server, mock_kb): """tool_name filter is threaded through as where clause.""" - call_tool(server, "kb_search", { - "query": "quality filtering", - "collection": "tool_executions", - "tool_name": "fastp", - }) + call_tool( + server, + "kb_search", + { + "query": "quality filtering", + "collection": "tool_executions", + "tool_name": "fastp", + }, + ) mock_kb.search.assert_called_once_with( query="quality filtering", collection="tool_executions", + collections=None, top_k=5, rerank=None, where={"tool_name": "fastp"}, + where_document=None, ) def test_session_filter_passed(self, server, mock_kb): """session_id filter is threaded through.""" - call_tool(server, "kb_search", { - "query": "pipeline", - "collection": "tool_executions", - "session_id": "s3", - }) + call_tool( + server, + "kb_search", + { + "query": "pipeline", + "collection": "tool_executions", + "session_id": "s3", + }, + ) mock_kb.search.assert_called_once_with( query="pipeline", collection="tool_executions", + collections=None, top_k=5, rerank=None, where={"session_id": "s3"}, + where_document=None, ) def test_multiple_filters_combined(self, server, mock_kb): """Multiple filters produce a compound $and where clause.""" - call_tool(server, "kb_search", { - "query": "parameters", - "collection": "tool_executions", - "tool_name": "fastp", - "session_id": "s1", - }) - - call_kwargs = mock_kb.search.call_args[1] - assert "where" in call_kwargs - where = call_kwargs["where"] - assert "$and" in where + call_tool( + server, + "kb_search", + { + "query": "parameters", + "collection": "tool_executions", + "tool_name": "fastp", + "session_id": "s1", + }, + ) + + # The handler emits one single-key dict per filter, in the fixed + # source key order (category, session_id, source_type, tool_name) — + # so session_id precedes tool_name. Pin the exact payload: a dropped + # or duplicated filter must fail here. + where = mock_kb.search.call_args[1]["where"] + assert where == {"$and": [{"session_id": "s1"}, {"tool_name": "fastp"}]} def test_return_code_filter_passed(self, server, mock_kb): """return_code filter is threaded as integer.""" - call_tool(server, "kb_search", { - "query": "failures", - "collection": "tool_executions", - "return_code": 1, - }) + call_tool( + server, + "kb_search", + { + "query": "failures", + "collection": "tool_executions", + "return_code": 1, + }, + ) mock_kb.search.assert_called_once_with( query="failures", collection="tool_executions", + collections=None, top_k=5, rerank=None, where={"return_code": 1}, + where_document=None, ) def test_filters_with_multi_collection(self, server, mock_kb): - """Filters apply to each collection in a multi-collection search.""" + """Multi-collection search delegates to kb.search once with collections=.""" mock_kb.search.return_value = [ make_search_result("result", "/f.md"), ] - call_tool(server, "kb_search", { - "query": "test", - "collections": ["tool_executions", "episodic_memory"], - "tool_name": "fastp", - }) + call_tool( + server, + "kb_search", + { + "query": "test", + "collections": ["tool_executions", "episodic_memory"], + "tool_name": "fastp", + }, + ) + + # Fan-out lives in kb.search now; the handler makes a single call. + mock_kb.search.assert_called_once_with( + query="test", + collection=None, + collections=["tool_executions", "episodic_memory"], + top_k=5, + rerank=None, + where={"tool_name": "fastp"}, + where_document=None, + ) + - assert mock_kb.search.call_count == 2 - for call in mock_kb.search.call_args_list: - assert call[1]["where"] == {"tool_name": "fastp"} +class TestDocumentFilterThreading: + """The 'regex' / 'contains' args build a where_document passed to kb.search.""" + + def test_regex_threaded_as_where_document(self, server, mock_kb): + call_tool( + server, + "kb_search", + {"query": "q", "collection": "session_memory", "regex": "(?i)parser"}, + ) + assert mock_kb.search.call_args[1]["where_document"] == {"$regex": "(?i)parser"} + + def test_contains_threaded_as_where_document(self, server, mock_kb): + call_tool( + server, + "kb_search", + {"query": "q", "collection": "session_memory", "contains": "Q30"}, + ) + assert mock_kb.search.call_args[1]["where_document"] == {"$contains": "Q30"} + + def test_regex_and_contains_combined_with_and(self, server, mock_kb): + call_tool( + server, + "kb_search", + { + "query": "q", + "collection": "session_memory", + "regex": "fastp", + "contains": "Q30", + }, + ) + assert mock_kb.search.call_args[1]["where_document"] == { + "$and": [{"$regex": "fastp"}, {"$contains": "Q30"}] + } + + def test_no_document_filter_is_none(self, server, mock_kb): + call_tool(server, "kb_search", {"query": "q", "collection": "session_memory"}) + assert mock_kb.search.call_args[1]["where_document"] is None # --------------------------------------------------------------------------- # Handler: metadata in results # --------------------------------------------------------------------------- + class TestSearchResultMetadata: def test_metadata_included_in_results(self, mock_kb): """Search results include extra metadata fields.""" mock_kb.search.return_value = [ make_search_result( - "fastp ran", "/file.md", + "fastp ran", + "/file.md", extra_meta={"tool_name": "fastp", "session_id": "s1", "return_code": 0}, ), ] server = create_knowledge_server(mock_kb) - result = call_tool(server, "kb_search", { - "query": "test", - "collection": "tool_executions", - }) + result = call_tool( + server, + "kb_search", + { + "query": "test", + "collection": "tool_executions", + }, + ) hit = result["results"][0] assert hit["metadata"]["tool_name"] == "fastp" @@ -176,10 +243,14 @@ def test_metadata_excludes_standard_fields(self, mock_kb): ] server = create_knowledge_server(mock_kb) - result = call_tool(server, "kb_search", { - "query": "test", - "collection": "docs", - }) + result = call_tool( + server, + "kb_search", + { + "query": "test", + "collection": "docs", + }, + ) meta = result["results"][0]["metadata"] assert "source_file" not in meta @@ -194,10 +265,14 @@ def test_empty_metadata_for_reference_collections(self, mock_kb): ] server = create_knowledge_server(mock_kb) - result = call_tool(server, "kb_search", { - "query": "test", - "collection": "docs", - }) + result = call_tool( + server, + "kb_search", + { + "query": "test", + "collection": "docs", + }, + ) assert result["results"][0]["metadata"] == {} @@ -206,6 +281,7 @@ def test_empty_metadata_for_reference_collections(self, mock_kb): # Schema: filter params visible to agent # --------------------------------------------------------------------------- + class TestSearchSchemaFilters: def _get_kb_search_schema(self, server): @@ -221,7 +297,15 @@ def test_filter_params_in_schema(self, server): """All filter parameters are advertised in the kb_search schema.""" schema = self._get_kb_search_schema(server) props = schema["properties"] - for param in ("category", "session_id", "tool_name", "source_type", "return_code"): + for param in ( + "category", + "session_id", + "tool_name", + "source_type", + "return_code", + "regex", + "contains", + ): assert param in props, f"Missing filter param: {param}" def test_filter_params_not_required(self, server): @@ -234,6 +318,7 @@ def test_filter_params_not_required(self, server): # Handler: error behavior for missing collections # --------------------------------------------------------------------------- + class TestSearchCollectionErrors: def test_single_missing_collection_returns_error(self, mock_kb): @@ -241,55 +326,51 @@ def test_single_missing_collection_returns_error(self, mock_kb): mock_kb.search.side_effect = ValueError("Collection 'missing' not found") server = create_knowledge_server(mock_kb) - result = call_tool(server, "kb_search", { - "query": "test", - "collection": "missing", - }) + result = call_tool( + server, + "kb_search", + { + "query": "test", + "collection": "missing", + }, + ) assert result["status"] == "error" assert "not found" in result["error"] - def test_all_multi_collections_missing_returns_error(self, mock_kb): - """When every collection in a multi-search fails, return error.""" - mock_kb.search.side_effect = ValueError("not found") + def test_all_collections_missing_passes_through_error(self, mock_kb): + """kb.search owns the all-failed aggregation; the handler passes it through.""" + mock_kb.search.side_effect = ValueError( + "All collections failed: Collection 'missing1' not found" + ) server = create_knowledge_server(mock_kb) - result = call_tool(server, "kb_search", { - "query": "test", - "collections": ["missing1", "missing2"], - }) + result = call_tool( + server, + "kb_search", + { + "query": "test", + "collections": ["missing1", "missing2"], + }, + ) assert result["status"] == "error" assert "All collections failed" in result["error"] - def test_partial_multi_collection_returns_ok_with_warnings(self, mock_kb): - """When some collections fail, return ok with results and warnings.""" - def search_with_partial_error(**kwargs): - if kwargs["collection"] == "missing": - raise ValueError("Collection 'missing' not found") - return [make_search_result("found it", "/file.md")] - - mock_kb.search.side_effect = search_with_partial_error + def test_partial_multi_collection_returns_ok(self, mock_kb): + """Partial-skip happens inside kb.search; the handler just returns its + (already-fused) results as ok.""" + mock_kb.search.return_value = [make_search_result("found it", "/file.md")] server = create_knowledge_server(mock_kb) - result = call_tool(server, "kb_search", { - "query": "test", - "collections": ["docs", "missing"], - }) + result = call_tool( + server, + "kb_search", + { + "query": "test", + "collections": ["docs", "missing"], + }, + ) assert result["status"] == "ok" assert result["result_count"] == 1 - assert "warnings" in result - assert any("not found" in w for w in result["warnings"]) - - def test_successful_search_has_no_warnings(self, mock_kb): - """Successful search doesn't include a warnings field.""" - server = create_knowledge_server(mock_kb) - - result = call_tool(server, "kb_search", { - "query": "test", - "collection": "docs", - }) - - assert result["status"] == "ok" - assert "warnings" not in result diff --git a/tests/test_knowledge_base.py b/tests/test_knowledge_base.py index 6530dfc..a85be5c 100644 --- a/tests/test_knowledge_base.py +++ b/tests/test_knowledge_base.py @@ -1,9 +1,9 @@ """ -Tests for KnowledgeBase and APIEmbeddingClient. +Tests for KnowledgeBase and APIEmbedder. -APIEmbeddingClient tests mock litellm.embedding to avoid network calls. -KnowledgeBase tests mock _make_embedder with deterministic vectors -and use real FAISS indexes and llama-index chunking on temp files. +APIEmbedder tests mock the client's httpx POST to avoid network +calls. KnowledgeBase tests mock Embedder.create with deterministic vectors +and use real ChromaDB indexes and llama-index chunking on temp files. Reranking is mocked since sentence-transformers is a heavy dependency. """ @@ -13,11 +13,11 @@ from pathlib import Path from unittest.mock import patch, MagicMock -import faiss +import httpx import numpy as np import pytest -from dsagt.knowledge import APIEmbeddingClient, KnowledgeBase, CODE_LANGUAGES +from dsagt.knowledge import APIEmbedder, KnowledgeBase, CODE_LANGUAGES @pytest.fixture(autouse=True) @@ -48,21 +48,23 @@ def fake_embed(texts: list[str]) -> np.ndarray: return np.array(embeddings, dtype=np.float32) -def make_mock_litellm_response(texts: list[str], dim: int = EMBEDDING_DIM): - """Create a fake litellm.EmbeddingResponse for the given texts. +def make_http_response(texts: list[str], dim: int = EMBEDDING_DIM, status: int = 200): + """Build a real ``httpx.Response`` mimicking an OpenAI ``/v1/embeddings`` + reply for the given texts. - LiteLLM returns a Pydantic model whose ``.data`` field is a list of dicts - with ``index`` and ``embedding`` keys (matching the OpenAI API shape). - A MagicMock with the same attribute access is sufficient for tests. + Returning a real ``httpx.Response`` (not a MagicMock) lets the client's + own ``raise_for_status()`` / ``json()`` calls run for real, so tests + exercise the actual parsing + error-classification paths. """ rng = np.random.RandomState(0) data = [ - {"index": i, "embedding": rng.randn(dim).tolist()} - for i in range(len(texts)) + {"index": i, "embedding": rng.randn(dim).tolist()} for i in range(len(texts)) ] - resp = MagicMock() - resp.data = data - return resp + return httpx.Response( + status, + json={"data": data}, + request=httpx.Request("POST", "http://test/embeddings"), + ) def create_test_docs(folder: Path): @@ -90,10 +92,11 @@ def create_test_docs(folder: Path): # --------------------------------------------------------------------------- -# APIEmbeddingClient +# APIEmbedder # --------------------------------------------------------------------------- -class TestAPIEmbeddingClient: + +class TestAPIEmbedder: def test_missing_base_url_raises(self): """Constructor raises ValueError when no base URL is available.""" @@ -102,7 +105,7 @@ def test_missing_base_url_raises(self): env.pop("OPENAI_BASE_URL", None) with patch.dict(os.environ, env, clear=True): with pytest.raises(ValueError, match="base URL required"): - APIEmbeddingClient(api_key="test-key", base_url=None) + APIEmbedder(api_key="test-key", base_url=None) def test_missing_api_key_raises(self): """Constructor raises ValueError when no API key is available.""" @@ -112,342 +115,267 @@ def test_missing_api_key_raises(self): env.pop("OPENAI_API_KEY", None) with patch.dict(os.environ, env, clear=True): with pytest.raises(ValueError, match="API key required"): - APIEmbeddingClient(api_key=None, base_url="http://test") + APIEmbedder(api_key=None, base_url="http://test") def test_explicit_api_key(self): """Constructor accepts an explicit API key.""" - client = APIEmbeddingClient(api_key="explicit-key", base_url="http://test") + client = APIEmbedder(api_key="explicit-key", base_url="http://test") assert client.api_key == "explicit-key" client.close() - def test_bare_model_name_gets_openai_like_prefix(self): - """Bare model names should be routed via the openai_like/ prefix. - - ``openai_like`` is required (not ``openai``) so LiteLLM does not - normalize lab-specific suffixes like ``-project`` away — see the - comment in APIEmbeddingClient.__init__ for the full rationale. - """ - client = APIEmbeddingClient( - api_key="k", base_url="http://test", model="my-embed", - ) - assert client._litellm_model == "openai_like/my-embed" - client.close() - - def test_lab_suffixed_model_name_round_trips_verbatim(self): - """Regression: ``text-embedding-3-small-project`` must NOT be - normalized to ``text-embedding-3-small`` by the openai_like router. + def test_model_name_sent_verbatim(self): + """The model string is sent unchanged — no provider prefix. - Lab LiteLLM proxies route by alias. If the suffix gets stripped, the - request reaches the upstream as a name the proxy's ACL doesn't know - about, and we get a 401 ``team_model_access_denied``. + Gateways route by alias, so lab-specific suffixes + (``text-embedding-3-small-project``) and HuggingFace-style names + with slashes (``lbl/nomic-embed-text``) must reach the endpoint + exactly as configured. """ - client = APIEmbeddingClient( - api_key="k", base_url="http://test", - model="text-embedding-3-small-project", + for name in ( + "my-embed", + "text-embedding-3-small-project", + "lbl/nomic-embed-text", + ): + client = APIEmbedder(api_key="k", base_url="http://test", model=name) + assert client.model == name + client.close() + + def test_embeddings_url_built_from_base_url(self): + """The embeddings route hangs off the (``/v1``) base URL root.""" + client = APIEmbedder( + api_key="k", + base_url="https://gw.example.com/v1/", + model="m", ) - assert client._litellm_model == "openai_like/text-embedding-3-small-project" + assert client._embeddings_url == "https://gw.example.com/v1/embeddings" client.close() - def test_slash_in_model_name_still_gets_openai_like_prefix(self): - """HuggingFace-style names (``lbl/nomic-embed-text``) are model - identifiers, not LiteLLM provider prefixes — the whole thing needs - ``openai_like/`` in front so LiteLLM dispatches to the OpenAI-wire - client pointed at our base_url. The rest of DSAGT assumes an - OpenAI-compat endpoint, so there's no valid case for the user's - string to carry a LiteLLM provider prefix of its own. - """ - client = APIEmbeddingClient( - api_key="k", base_url="http://test", model="lbl/nomic-embed-text", - ) - assert client._litellm_model == "openai_like/lbl/nomic-embed-text" - client.close() - - def test_embed_empty_list(self): - """Embedding an empty list returns an empty array.""" - client = APIEmbeddingClient(api_key="test-key", base_url="http://test") - result = client.embed([]) - assert result.shape == (0,) - assert result.dtype == np.float32 - client.close() - - @patch("litellm.embedding") - def test_embed_calls_litellm_with_correct_args(self, mock_embedding): - """litellm.embedding receives the right model, input, api_base, api_key.""" - mock_embedding.return_value = make_mock_litellm_response(["hello"]) - - client = APIEmbeddingClient( + def test_embed_posts_correct_request(self): + """The POST carries the right URL, model/input body, and auth header.""" + client = APIEmbedder( api_key="my-key", model="test-model", - base_url="https://example.com", + base_url="https://example.com/v1", ) - result = client.embed(["hello"]) + with patch.object( + client._client, + "post", + return_value=make_http_response(["hello"]), + ) as mock_post: + result = client.embed(["hello"]) assert result.shape == (1, EMBEDDING_DIM) - assert mock_embedding.call_count == 1 - call_kwargs = mock_embedding.call_args.kwargs - assert call_kwargs["model"] == "openai_like/test-model" - assert call_kwargs["input"] == ["hello"] - assert call_kwargs["api_base"] == "https://example.com" - assert call_kwargs["api_key"] == "my-key" + assert mock_post.call_count == 1 + url = mock_post.call_args.args[0] + kwargs = mock_post.call_args.kwargs + assert url == "https://example.com/v1/embeddings" + assert kwargs["json"] == {"model": "test-model", "input": ["hello"]} + assert kwargs["headers"]["Authorization"] == "Bearer my-key" client.close() - @patch("litellm.embedding") - def test_embed_returns_vectors_in_index_order(self, mock_embedding): + def test_embed_returns_vectors_in_index_order(self): """Out-of-order response data is sorted back to input order.""" - # Construct a response where data is in reverse order to ensure - # the client sorts by 'index' field. rng = np.random.RandomState(7) out_of_order = [ {"index": 1, "embedding": rng.randn(EMBEDDING_DIM).tolist()}, {"index": 0, "embedding": rng.randn(EMBEDDING_DIM).tolist()}, ] - resp = MagicMock() - resp.data = out_of_order - mock_embedding.return_value = resp + resp = httpx.Response( + 200, + json={"data": out_of_order}, + request=httpx.Request("POST", "http://test/embeddings"), + ) - client = APIEmbeddingClient(api_key="k", base_url="http://test") - result = client.embed(["first", "second"]) + client = APIEmbedder(api_key="k", base_url="http://test") + with patch.object(client._client, "post", return_value=resp): + result = client.embed(["first", "second"]) assert result.shape == (2, EMBEDDING_DIM) # First-row vector matches the data entry with index=0 (the second list element) - assert np.allclose(result[0], np.array(out_of_order[1]["embedding"], dtype=np.float32)) - client.close() - - @patch("litellm.embedding") - def test_embed_handles_pydantic_style_data(self, mock_embedding): - """LiteLLM may return Pydantic objects with attribute access instead of dicts.""" - - class _Item: - def __init__(self, idx, vec): - self.index = idx - self.embedding = vec - - rng = np.random.RandomState(0) - items = [_Item(i, rng.randn(EMBEDDING_DIM).tolist()) for i in range(3)] - resp = MagicMock() - resp.data = items - mock_embedding.return_value = resp - - client = APIEmbeddingClient(api_key="k", base_url="http://test") - result = client.embed(["a", "b", "c"]) - assert result.shape == (3, EMBEDDING_DIM) + assert np.allclose( + result[0], np.array(out_of_order[1]["embedding"], dtype=np.float32) + ) client.close() # --------------------------------------------------------------------------- -# APIEmbeddingClient - retry and error propagation +# APIEmbedder - retry and error propagation # --------------------------------------------------------------------------- -class TestAPIEmbeddingClientErrors: - """Verify the explicit rate-limit retry layer in APIEmbeddingClient. - The retry layer exists because lab LiteLLM proxies wrap upstream 429s - in a way that defeats litellm's built-in retry classification, and - litellm's default backoff is too short for Azure-style 60s quota - windows. These tests pin the contract: +class TestAPIEmbedderErrors: + """Verify the explicit rate-limit retry layer in APIEmbedder. + + The retry layer exists because lab gateways enforce Azure-style 60s + quota windows that a generic exponential backoff would undershoot. + These tests pin the contract: * Authentication / bad-request errors are NOT retried (fail fast on misconfiguration). * Rate-limit and transient errors ARE retried up to max_attempts. - * The upstream "retry after N seconds" hint is honored. + * The ``Retry-After`` header / body hint is honored. """ - @patch("dsagt.knowledge.time.sleep") - @patch("litellm.embedding") - def test_authentication_error_propagates_immediately( - self, mock_embedding, mock_sleep, - ): - """A 401 must NOT be retried — this is a misconfiguration, not transient.""" - import litellm - - mock_embedding.side_effect = litellm.exceptions.AuthenticationError( - message="Invalid API key", llm_provider="openai", model="test", + def _err_response(self, status, headers=None, json_body=None): + return httpx.Response( + status, + headers=headers or {}, + json=json_body if json_body is not None else {"error": "x"}, + request=httpx.Request("POST", "http://test/embeddings"), ) - client = APIEmbeddingClient(api_key="bad-key", base_url="http://test") - with pytest.raises(litellm.exceptions.AuthenticationError): - client.embed(["test"]) + @patch("dsagt.knowledge.time.sleep") + def test_authentication_error_propagates_immediately(self, mock_sleep): + """A 401 must NOT be retried — this is a misconfiguration, not transient.""" + client = APIEmbedder(api_key="bad-key", base_url="http://test") + with patch.object( + client._client, + "post", + return_value=self._err_response(401), + ) as mock_post: + with pytest.raises(httpx.HTTPStatusError): + client.embed(["test"]) # No retries: one call, no sleeps. - assert mock_embedding.call_count == 1 + assert mock_post.call_count == 1 assert mock_sleep.call_count == 0 client.close() @patch("dsagt.knowledge.time.sleep") - @patch("litellm.embedding") - def test_rate_limit_retries_then_propagates(self, mock_embedding, mock_sleep): + def test_rate_limit_retries_then_propagates(self, mock_sleep): """A persistent rate limit retries up to max_attempts then raises.""" - import litellm - - mock_embedding.side_effect = litellm.exceptions.RateLimitError( - message="429 Please retry after 60 seconds", - llm_provider="openai", model="test", - ) - - client = APIEmbeddingClient(api_key="k", base_url="http://test") - with pytest.raises(litellm.exceptions.RateLimitError): - client.embed(["test"]) + resp = self._err_response(429, headers={"retry-after": "60"}) + client = APIEmbedder(api_key="k", base_url="http://test") + with patch.object(client._client, "post", return_value=resp) as mock_post: + with pytest.raises(httpx.HTTPStatusError): + client.embed(["test"]) # max_attempts is 6 — that's 6 calls and 5 sleeps between them. - assert mock_embedding.call_count == 6 + assert mock_post.call_count == 6 assert mock_sleep.call_count == 5 - # Each sleep should respect the upstream-suggested 60s wait. + # Each sleep should respect the Retry-After 60s hint. for call in mock_sleep.call_args_list: assert call.args[0] == 60.0 client.close() @patch("dsagt.knowledge.time.sleep") - @patch("litellm.embedding") - def test_rate_limit_retries_then_succeeds(self, mock_embedding, mock_sleep): + def test_rate_limit_retries_then_succeeds(self, mock_sleep): """If the rate limit clears on a retry, embed() returns successfully.""" - import litellm - - rng = np.random.RandomState(42) - success_resp = MagicMock() - success_resp.data = [ - {"index": 0, "embedding": rng.randn(EMBEDDING_DIM).tolist()}, - ] - + client = APIEmbedder(api_key="k", base_url="http://test") # Two rate-limit failures, then success. - mock_embedding.side_effect = [ - litellm.exceptions.RateLimitError( - message="429 Please retry after 60 seconds", - llm_provider="openai", model="test", - ), - litellm.exceptions.RateLimitError( - message="429 Please retry after 60 seconds", - llm_provider="openai", model="test", - ), - success_resp, + side = [ + self._err_response(429, headers={"retry-after": "60"}), + self._err_response(429, headers={"retry-after": "60"}), + make_http_response(["one"]), ] - - client = APIEmbeddingClient(api_key="k", base_url="http://test") - result = client.embed(["one"]) + with patch.object(client._client, "post", side_effect=side) as mock_post: + result = client.embed(["one"]) assert result.shape == (1, EMBEDDING_DIM) - assert mock_embedding.call_count == 3 + assert mock_post.call_count == 3 assert mock_sleep.call_count == 2 client.close() @patch("dsagt.knowledge.time.sleep") - @patch("litellm.embedding") - def test_transient_connection_error_retries(self, mock_embedding, mock_sleep): - """APIConnectionError is retryable (covers the lab-proxy 429 wrapping case).""" - import litellm - - mock_embedding.side_effect = litellm.exceptions.APIConnectionError( - message="connection reset", - llm_provider="openai_like", model="test", + def test_transient_connection_error_retries(self, mock_sleep): + """Transport errors (connection reset) are retryable.""" + client = APIEmbedder(api_key="k", base_url="http://test") + err = httpx.ConnectError( + "connection reset", + request=httpx.Request("POST", "http://test/embeddings"), ) - - client = APIEmbeddingClient(api_key="k", base_url="http://test") - with pytest.raises(litellm.exceptions.APIConnectionError): - client.embed(["test"]) - assert mock_embedding.call_count == 6 + with patch.object(client._client, "post", side_effect=err) as mock_post: + with pytest.raises(httpx.ConnectError): + client.embed(["test"]) + assert mock_post.call_count == 6 client.close() @patch("dsagt.knowledge.time.sleep") - @patch("litellm.embedding") - def test_lab_proxy_wrapped_429_retries(self, mock_embedding, mock_sleep): - """Real-world failure mode: openai_like wraps an upstream 429 as - APIConnectionError with the rate-limit body in the message. Our - retry layer must detect this via string matching, not just by - exception class, and must use the upstream-suggested wait time. - """ - import litellm - - # This is a lightly-paraphrased version of the actual lab error. - wrapped_429 = litellm.exceptions.APIConnectionError( - message=( - "Openai_likeException - " - '{"error":{"message":"litellm.RateLimitError: AzureException ' - "RateLimitError - rate limit exceeded. " - "Please retry after 90 seconds. " - 'To increase your default rate limit, visit ..."}}' - ), - llm_provider="openai_like", model="text-embedding-3-small-project", - ) - mock_embedding.side_effect = wrapped_429 - - client = APIEmbeddingClient(api_key="k", base_url="http://test") - with pytest.raises(litellm.exceptions.APIConnectionError): - client.embed(["test"]) - - assert mock_embedding.call_count == 6 - # The 90s hint from the upstream message must be honored, not the 60s default. + def test_rate_limit_body_hint_honored_without_header(self, mock_sleep): + """When there's no Retry-After header, a 429 body hint is parsed.""" + body = { + "error": { + "message": ( + "rate limit exceeded. Please retry after 90 seconds. " + "To increase your default rate limit, visit ..." + ) + } + } + resp = self._err_response(429, json_body=body) + client = APIEmbedder(api_key="k", base_url="http://test") + with patch.object(client._client, "post", return_value=resp) as mock_post: + with pytest.raises(httpx.HTTPStatusError): + client.embed(["test"]) + + assert mock_post.call_count == 6 + # The 90s hint from the body must be honored, not the 60s default. for call in mock_sleep.call_args_list: assert call.args[0] == 90.0 client.close() @patch("dsagt.knowledge.time.sleep") - @patch("litellm.embedding") - def test_timeout_retries(self, mock_embedding, mock_sleep): + def test_timeout_retries(self, mock_sleep): """Timeouts are transient and should be retried with exponential backoff.""" - import litellm - - mock_embedding.side_effect = litellm.exceptions.Timeout( - message="Request timed out", llm_provider="openai", model="test", + client = APIEmbedder(api_key="k", base_url="http://test") + err = httpx.ReadTimeout( + "Request timed out", + request=httpx.Request("POST", "http://test/embeddings"), ) - - client = APIEmbeddingClient(api_key="k", base_url="http://test") - with pytest.raises(litellm.exceptions.Timeout): - client.embed(["test"]) - assert mock_embedding.call_count == 6 + with patch.object(client._client, "post", side_effect=err) as mock_post: + with pytest.raises(httpx.ReadTimeout): + client.embed(["test"]) + assert mock_post.call_count == 6 # Exponential backoff capped at 30s: 2^1, 2^2, 2^3, 2^4, min(2^5, 30). sleeps = [c.args[0] for c in mock_sleep.call_args_list] assert sleeps == [2.0, 4.0, 8.0, 16.0, 30.0] client.close() -class TestAPIEmbeddingClientBatching: +class TestAPIEmbedderBatching: """Long inputs are split into batch_size chunks for the embedding API.""" - @patch("litellm.embedding") - def test_batches_when_over_batch_size(self, mock_embedding): + def test_batches_when_over_batch_size(self): """A 250-text input with batch_size=100 produces 3 API calls.""" - rng = np.random.RandomState(0) - - def make_response(input_texts, **_): - resp = MagicMock() - resp.data = [ - {"index": i, "embedding": rng.randn(EMBEDDING_DIM).tolist()} - for i in range(len(input_texts)) - ] - return resp + client = APIEmbedder( + api_key="k", + base_url="http://test", + batch_size=100, + ) - mock_embedding.side_effect = lambda **kwargs: make_response(kwargs["input"]) + def make_response(url, *, json, headers): + return make_http_response(json["input"]) - client = APIEmbeddingClient( - api_key="k", base_url="http://test", batch_size=100, - ) - texts = [f"chunk {i}" for i in range(250)] - result = client.embed(texts) + with patch.object( + client._client, + "post", + side_effect=make_response, + ) as mock_post: + texts = [f"chunk {i}" for i in range(250)] + result = client.embed(texts) assert result.shape == (250, EMBEDDING_DIM) - assert mock_embedding.call_count == 3 + assert mock_post.call_count == 3 # Verify the batch sizes were 100, 100, 50. batch_sizes = [ - len(call.kwargs["input"]) for call in mock_embedding.call_args_list + len(call.kwargs["json"]["input"]) for call in mock_post.call_args_list ] assert batch_sizes == [100, 100, 50] client.close() - @patch("litellm.embedding") - def test_single_call_when_under_batch_size(self, mock_embedding): + def test_single_call_when_under_batch_size(self): """Inputs within batch_size make a single call (no batching loop).""" - rng = np.random.RandomState(0) - resp = MagicMock() - resp.data = [ - {"index": i, "embedding": rng.randn(EMBEDDING_DIM).tolist()} - for i in range(5) - ] - mock_embedding.return_value = resp - - client = APIEmbeddingClient( - api_key="k", base_url="http://test", batch_size=100, + client = APIEmbedder( + api_key="k", + base_url="http://test", + batch_size=100, ) - result = client.embed(["a", "b", "c", "d", "e"]) + with patch.object( + client._client, + "post", + return_value=make_http_response(["a", "b", "c", "d", "e"]), + ) as mock_post: + result = client.embed(["a", "b", "c", "d", "e"]) assert result.shape == (5, EMBEDDING_DIM) - assert mock_embedding.call_count == 1 + assert mock_post.call_count == 1 client.close() @@ -456,23 +384,28 @@ class TestRetryAfterParsing: def test_extracts_seconds_from_message(self): from dsagt.knowledge import _extract_retry_after_seconds + msg = "Please retry after 60 seconds" assert _extract_retry_after_seconds(msg) == 60.0 def test_extracts_seconds_with_decimal(self): from dsagt.knowledge import _extract_retry_after_seconds + assert _extract_retry_after_seconds("retry after 12.5 seconds") == 12.5 def test_case_insensitive(self): from dsagt.knowledge import _extract_retry_after_seconds + assert _extract_retry_after_seconds("RETRY AFTER 30 SECONDS") == 30.0 def test_returns_default_when_no_hint(self): from dsagt.knowledge import _extract_retry_after_seconds + assert _extract_retry_after_seconds("rate limit", default=42.0) == 42.0 def test_finds_hint_in_long_message(self): from dsagt.knowledge import _extract_retry_after_seconds + # The real lab error nests the hint deep inside a JSON body. msg = ( 'Openai_likeException - {"error":{"message":' @@ -486,6 +419,7 @@ def test_finds_hint_in_long_message(self): # KnowledgeBase.ingest exclude_patterns # --------------------------------------------------------------------------- + class TestIngestExcludePatterns: """The exclude_patterns parameter on ingest() filters out files whose relative path matches any of the supplied glob patterns. Used by @@ -498,7 +432,7 @@ def kb(self, tmp_path): index_dir = tmp_path / "index" mock_client = MagicMock() mock_client.embed = fake_embed - with patch("dsagt.knowledge._make_embedder", return_value=mock_client): + with patch("dsagt.knowledge.Embedder.create", return_value=mock_client): kb = KnowledgeBase(index_dir=index_dir) yield kb kb.close() @@ -513,27 +447,27 @@ def repo_layout(self, tmp_path): # Public source (root / "mylib").mkdir() (root / "mylib" / "__init__.py").write_text('"""Mylib package."""\n') - (root / "mylib" / "core.py").write_text('def public_fn():\n return 1\n') - (root / "mylib" / "_internal.py").write_text('def _hidden():\n return 2\n') + (root / "mylib" / "core.py").write_text("def public_fn():\n return 1\n") + (root / "mylib" / "_internal.py").write_text("def _hidden():\n return 2\n") # Tests in a subdirectory (root / "mylib" / "tests").mkdir() (root / "mylib" / "tests" / "__init__.py").write_text("") (root / "mylib" / "tests" / "test_core.py").write_text( - 'def test_public_fn():\n assert True\n' + "def test_public_fn():\n assert True\n" ) # Top-level tests dir as well (root / "tests").mkdir() (root / "tests" / "test_integration.py").write_text( - 'def test_smoke():\n assert True\n' + "def test_smoke():\n assert True\n" ) (root / "tests" / "conftest.py").write_text("import pytest\n") # Examples (kept on purpose for the agent) (root / "examples").mkdir() (root / "examples" / "quickstart.py").write_text( - 'from mylib import public_fn\nprint(public_fn())\n' + "from mylib import public_fn\nprint(public_fn())\n" ) # Docs @@ -548,7 +482,8 @@ def repo_layout(self, tmp_path): def test_no_exclude_keeps_everything(self, kb, repo_layout): result = kb.ingest( - repo_layout, collection_name="full", + repo_layout, + collection_name="full", file_types=["py", "md"], ) # All py + md files except the .pyc which isn't in file_types. @@ -561,7 +496,8 @@ def test_no_exclude_keeps_everything(self, kb, repo_layout): def test_exclude_tests_directory(self, kb, repo_layout): """Pattern 'tests' should match the tests segment in any path.""" result = kb.ingest( - repo_layout, collection_name="no_tests", + repo_layout, + collection_name="no_tests", file_types=["py", "md"], exclude_patterns=["tests"], ) @@ -574,7 +510,8 @@ def test_exclude_tests_directory(self, kb, repo_layout): def test_exclude_test_files_by_basename(self, kb, repo_layout): """Pattern 'test_*.py' matches the basename anywhere in the tree.""" result = kb.ingest( - repo_layout, collection_name="no_test_files", + repo_layout, + collection_name="no_test_files", file_types=["py", "md"], exclude_patterns=["test_*.py"], ) @@ -591,7 +528,8 @@ def test_exclude_private_modules(self, kb, repo_layout): want to keep __init__.py should use a more specific pattern. """ result = kb.ingest( - repo_layout, collection_name="no_private", + repo_layout, + collection_name="no_private", file_types=["py", "md"], exclude_patterns=["_*.py"], ) @@ -607,11 +545,13 @@ def test_exclude_pycache_dir(self, kb, repo_layout): (repo_layout / "mylib" / "__pycache__" / "fake.py").write_text("x = 1\n") result_unfiltered = kb.ingest( - repo_layout, collection_name="with_cache", + repo_layout, + collection_name="with_cache", file_types=["py"], ) result_filtered = kb.ingest( - repo_layout, collection_name="no_cache", + repo_layout, + collection_name="no_cache", file_types=["py"], exclude_patterns=["__pycache__"], ) @@ -622,7 +562,8 @@ def test_combined_default_patterns(self, kb, repo_layout): from dsagt.commands.setup_core_kb import DEFAULT_EXCLUDE_PATTERNS result = kb.ingest( - repo_layout, collection_name="defaults", + repo_layout, + collection_name="defaults", file_types=["py", "md"], exclude_patterns=DEFAULT_EXCLUDE_PATTERNS, ) @@ -650,7 +591,8 @@ def test_default_patterns_keep_packaging_metadata(self, kb, tmp_path): (root / "mylib" / "core.py").write_text("def f():\n return 1\n") result = kb.ingest( - root, collection_name="pkg_meta", + root, + collection_name="pkg_meta", file_types=["py", "toml", "cfg"], exclude_patterns=DEFAULT_EXCLUDE_PATTERNS, ) @@ -665,7 +607,7 @@ class TestCollectFilesDirectly: The whole point of pulling this helper out of ingest() was to make file-discovery and exclude-pattern logic testable WITHOUT spinning up an embedder, an index, or a chunker. These tests exercise the helper - directly: no mocked _make_embedder context, no add_entries call, no + directly: no mocked Embedder.create context, no add_entries call, no cleanup of cached collections. If a regression in the file-walk or fnmatch logic ever lands, these tests fail in milliseconds and point at the exact problem instead of being buried under ingest() setup. @@ -678,7 +620,7 @@ def kb(self, tmp_path): return KnowledgeBase( index_dir=tmp_path / "index", default_embedder="local", - embedder_kwargs={"model": "unused"}, + model="unused", ) @pytest.fixture @@ -733,6 +675,7 @@ def test_returns_paths_not_strings(self, kb, repo): # KnowledgeBase - collections and ingest # --------------------------------------------------------------------------- + class TestKnowledgeBaseIngest: @pytest.fixture @@ -742,7 +685,7 @@ def kb(self, tmp_path): mock_client = MagicMock() mock_client.embed = fake_embed - with patch("dsagt.knowledge._make_embedder", return_value=mock_client): + with patch("dsagt.knowledge.Embedder.create", return_value=mock_client): kb = KnowledgeBase(index_dir=index_dir) yield kb kb.close() @@ -788,24 +731,6 @@ def test_list_collections_includes_description(self, kb, source_folder): assert collections[0]["name"] == "test_docs" assert "unit tests" in collections[0]["description"] - def test_ingest_creates_faiss_index(self, tmp_path, source_folder): - """Ingest into an explicitly FAISS-routed KB produces an index.faiss.""" - index_dir = tmp_path / "index_faiss" - mock_client = MagicMock() - mock_client.embed = fake_embed - - with patch("dsagt.knowledge._make_embedder", return_value=mock_client): - kb = KnowledgeBase(index_dir=index_dir, default_index="faiss") - try: - kb.ingest(source_folder) - index_path = kb.index_dir / "test_docs" / "index.faiss" - assert index_path.exists() - - index = faiss.read_index(str(index_path)) - assert index.ntotal > 0 - finally: - kb.close() - def test_ingest_creates_chunks_jsonl(self, kb, source_folder): """Ingest produces a chunks.jsonl with valid entries.""" result = kb.ingest(source_folder) @@ -864,6 +789,7 @@ def test_ingest_no_description(self, kb, tmp_path): # KnowledgeBase - search # --------------------------------------------------------------------------- + class TestKnowledgeBaseSearch: @pytest.fixture @@ -877,7 +803,7 @@ def kb_with_data(self, tmp_path): mock_client = MagicMock() mock_client.embed = fake_embed - with patch("dsagt.knowledge._make_embedder", return_value=mock_client): + with patch("dsagt.knowledge.Embedder.create", return_value=mock_client): kb = KnowledgeBase(index_dir=index_dir) kb.ingest(source_folder) yield kb @@ -924,7 +850,7 @@ def test_search_with_rerank(self, kb_with_data): mock_st = MagicMock() mock_st.CrossEncoder.return_value = mock_reranker - import sys + with patch.dict(sys.modules, {"sentence_transformers": mock_st}): # Ensure the lazy import triggers kb_with_data._reranker = None @@ -957,12 +883,14 @@ def test_search_collection_isolation(self, tmp_path): mock_client = MagicMock() mock_client.embed = fake_embed - with patch("dsagt.knowledge._make_embedder", return_value=mock_client): + with patch("dsagt.knowledge.Embedder.create", return_value=mock_client): kb = KnowledgeBase(index_dir=index_dir) kb.ingest(folder_a) kb.ingest(folder_b) - results = kb.search("rockets", collection="collection_a", top_k=5, rerank=False) + results = kb.search( + "rockets", collection="collection_a", top_k=5, rerank=False + ) sources = [r["chunk"]["metadata"]["collection"] for r in results] assert all(s == "collection_a" for s in sources) @@ -973,6 +901,7 @@ def test_search_collection_isolation(self, tmp_path): # KnowledgeBase - loading and caching # --------------------------------------------------------------------------- + class TestKnowledgeBaseLoad: def test_load_caches_collection(self, tmp_path): @@ -985,19 +914,20 @@ def test_load_caches_collection(self, tmp_path): mock_client = MagicMock() mock_client.embed = fake_embed - with patch("dsagt.knowledge._make_embedder", return_value=mock_client): + with patch("dsagt.knowledge.Embedder.create", return_value=mock_client): kb = KnowledgeBase(index_dir=index_dir) kb.ingest(source_folder) - # Clear the cache that ingest populated - kb._cache.clear() + # Clear the store cache that ingest populated + store = kb._store + store._cache.clear() # First load reads from disk - index1, chunks1 = kb._load("docs") - assert "docs" in kb._cache + index1, chunks1 = store._load("docs") + assert "docs" in store._cache # Second load returns cached - index2, chunks2 = kb._load("docs") + index2, chunks2 = store._load("docs") assert index1 is index2 assert chunks1 is chunks2 @@ -1008,28 +938,32 @@ def test_load_caches_collection(self, tmp_path): # KnowledgeBase - parser selection # --------------------------------------------------------------------------- + class TestGetParser: @pytest.fixture def kb(self, tmp_path): # New __init__ doesn't create an embedder, but mock to be safe - with patch("dsagt.knowledge._make_embedder"): + with patch("dsagt.knowledge.Embedder.create"): kb = KnowledgeBase(index_dir=tmp_path / "index") yield kb kb.close() def test_markdown_parser(self, kb): from llama_index.core.node_parser import MarkdownNodeParser + parser = kb._get_parser(".md") assert isinstance(parser, MarkdownNodeParser) def test_code_parser(self, kb): from llama_index.core.node_parser import CodeSplitter + parser = kb._get_parser(".py") assert isinstance(parser, CodeSplitter) def test_default_parser(self, kb): from llama_index.core.node_parser import SentenceSplitter + parser = kb._get_parser(".txt") assert isinstance(parser, SentenceSplitter) @@ -1044,108 +978,97 @@ def test_code_languages_coverage(self): # KnowledgeBase - context manager # --------------------------------------------------------------------------- + class TestContextManager: def test_context_manager_calls_close(self, tmp_path): """close() cleans up cached embedders created during the session.""" mock_client = MagicMock() - with patch("dsagt.knowledge._make_embedder", return_value=mock_client): + with patch("dsagt.knowledge.Embedder.create", return_value=mock_client): with KnowledgeBase(index_dir=tmp_path / "index") as kb: - # Trigger embedder creation so close() has something to clean up - route = kb._get_route("dummy") - kb._get_embedder(route) + # Trigger lazy embedder construction so close() has something + # to clean up. + kb._store.embedder mock_client.close.assert_called_once() -class TestEmbedderCredentialMerge: - """Routes carry per-collection overrides; credentials live on the kb's - default route. Explicit routes (e.g. EPISODIC_MEMORY_ROUTE) must inherit - api_key/base_url from the default kb config — otherwise memory extraction - falls through to env-var fallbacks and breaks under mixed-endpoint setups - where LLM_API_KEY != EMBEDDING_API_KEY. - """ +class TestStoreEmbedderConstruction: + """One embedder per store, built lazily from explicit args. - def test_explicit_route_inherits_default_credentials(self, tmp_path): - from dsagt.knowledge import CollectionRoute + Per-collection embedder routing was removed: the store fixes a single + embedder at construction, so the explicit args the KB was given flow + straight through to ``Embedder.create`` (named, no kwargs dict) on first use. + """ - with patch("dsagt.knowledge._make_embedder") as mock_make: + def test_store_builds_embedder_from_args(self, tmp_path): + with patch("dsagt.knowledge.Embedder.create") as mock_make: kb = KnowledgeBase( index_dir=tmp_path / "index", - embedder_kwargs={ - "api_key": "sk-real-key", - "base_url": "https://embed.example.com", - "model": "text-embedding-3-small", - }, + default_embedder="api", + api_key="sk-real-key", + base_url="https://embed.example.com", + model="text-embedding-3-small", ) - # Route with no embedder_kwargs (mimics EPISODIC_MEMORY_ROUTE). - route = CollectionRoute(embedding_backend="api", vector_db="chroma") - kb._get_embedder(route) + kb._store.embedder # trigger lazy construction mock_make.assert_called_once_with( "api", - api_key="sk-real-key", - base_url="https://embed.example.com", model="text-embedding-3-small", + base_url="https://embed.example.com", + api_key="sk-real-key", + device=None, ) - def test_route_kwargs_override_default(self, tmp_path): - """Per-route overrides win over kb defaults for matching keys.""" - from dsagt.knowledge import CollectionRoute - - with patch("dsagt.knowledge._make_embedder") as mock_make: + def test_embedder_built_once_and_cached(self, tmp_path): + """Repeated access builds the embedder exactly once.""" + with patch("dsagt.knowledge.Embedder.create") as mock_make: kb = KnowledgeBase( index_dir=tmp_path / "index", - embedder_kwargs={ - "api_key": "sk-real-key", - "base_url": "https://embed.example.com", - "model": "default-model", - }, - ) - # Route overrides model but not credentials. - route = CollectionRoute( - embedding_backend="api", - vector_db="chroma", - embedder_kwargs={"model": "BAAI/bge-base-en-v1.5"}, - ) - kb._get_embedder(route) - - mock_make.assert_called_once_with( - "api", + default_embedder="api", api_key="sk-real-key", base_url="https://embed.example.com", - model="BAAI/bge-base-en-v1.5", + model="default-model", ) + kb._store.embedder + kb._store.embedder + assert mock_make.call_count == 1 # --------------------------------------------------------------------------- # BM25 sparse index + RRF hybrid retrieval # --------------------------------------------------------------------------- + class TestBM25Tokenize: """The tokenizer drives recall — verify identifier-style splits land.""" def test_lowercases(self): from dsagt.knowledge import _bm25_tokenize + assert _bm25_tokenize("Hello WORLD") == ["hello", "world"] def test_splits_on_underscore(self): from dsagt.knowledge import _bm25_tokenize + # snake_case must fan out for BM25 to score "user_id" queries against # surrounding code. assert _bm25_tokenize("get_user_id") == ["get", "user", "id"] def test_splits_on_hyphen(self): from dsagt.knowledge import _bm25_tokenize + assert _bm25_tokenize("kb-ingest") == ["kb", "ingest"] def test_drops_punctuation(self): from dsagt.knowledge import _bm25_tokenize + assert _bm25_tokenize("foo.bar(baz)") == ["foo", "bar", "baz"] def test_keeps_numbers(self): from dsagt.knowledge import _bm25_tokenize + assert _bm25_tokenize("v1.2.3 model") == ["v1", "2", "3", "model"] @@ -1153,12 +1076,15 @@ class TestBM25Index: def test_build_then_search_finds_exact_match(self, tmp_path): from dsagt.knowledge import BM25Index + idx = BM25Index() - idx.build([ - "Alpha rocket fuel composition.", - "Beta submarine pressure hull.", - "Gamma rocket telemetry parser.", - ]) + idx.build( + [ + "Alpha rocket fuel composition.", + "Beta submarine pressure hull.", + "Gamma rocket telemetry parser.", + ] + ) scores, indices = idx.search("rocket", k=3) assert len(indices) == 3 # Docs 0 and 2 mention "rocket", should outrank doc 1. @@ -1167,6 +1093,7 @@ def test_build_then_search_finds_exact_match(self, tmp_path): def test_empty_corpus_returns_empty(self, tmp_path): from dsagt.knowledge import BM25Index + idx = BM25Index() idx.build([]) scores, indices = idx.search("anything", k=5) @@ -1174,6 +1101,7 @@ def test_empty_corpus_returns_empty(self, tmp_path): def test_empty_query_returns_empty(self, tmp_path): from dsagt.knowledge import BM25Index + idx = BM25Index() idx.build(["alpha", "beta"]) # Punctuation-only query tokenizes to []; should not crash. @@ -1182,6 +1110,7 @@ def test_empty_query_returns_empty(self, tmp_path): def test_save_load_roundtrip(self, tmp_path): from dsagt.knowledge import BM25Index + idx = BM25Index() idx.build(["foo bar", "baz qux", "foo qux"]) idx.save(tmp_path) @@ -1195,6 +1124,7 @@ def test_save_load_roundtrip(self, tmp_path): def test_load_missing_file_returns_empty(self, tmp_path): from dsagt.knowledge import BM25Index + loaded = BM25Index.load(tmp_path) assert loaded.size == 0 @@ -1203,11 +1133,13 @@ class TestRRFMerge: def test_single_ranker_passes_through_order(self): from dsagt.knowledge import _rrf_merge + merged = _rrf_merge([[5, 2, 7]]) assert [idx for idx, _ in merged] == [5, 2, 7] def test_two_rankers_promote_shared_top(self): from dsagt.knowledge import _rrf_merge + # Doc 1 is rank-1 in both rankings; doc 2 only in one. RRF must # rank doc 1 above all unique docs. merged = _rrf_merge([[1, 2, 3], [1, 4, 5]]) @@ -1216,12 +1148,71 @@ def test_two_rankers_promote_shared_top(self): def test_skips_negative_indices(self): from dsagt.knowledge import _rrf_merge - # FAISS pads with -1 when fewer than k results; must not pollute scores. + + # Backends pad with -1 when fewer than k results; must not pollute scores. merged = _rrf_merge([[0, -1, -1], [0, 1, 2]]) idxs = [idx for idx, _ in merged] assert -1 not in idxs +class TestFederatedSearch: + """KnowledgeBase.search owns collection→store routing + cross-collection RRF.""" + + @pytest.fixture + def kb(self, tmp_path): + mock_client = MagicMock() + mock_client.embed = fake_embed + with patch("dsagt.knowledge.Embedder.create", return_value=mock_client): + kb = KnowledgeBase(index_dir=tmp_path / "index") + kb.add_entries( + texts=["alpha quality filtering", "alpha assembly step"], + collection="coll_a", + ) + kb.add_entries( + texts=["beta quality filtering", "beta assembly step"], + collection="coll_b", + ) + yield kb + kb.close() + + def test_single_collection_routes_to_store(self, kb): + results = kb.search("quality", collection="coll_a", top_k=5, rerank=False) + assert results + assert all(r["chunk"]["metadata"]["collection"] == "coll_a" for r in results) + + def test_multi_collection_fuses_both(self, kb): + results = kb.search( + "quality filtering", + collections=["coll_a", "coll_b"], + top_k=10, + rerank=False, + ) + seen = {r["chunk"]["metadata"]["collection"] for r in results} + assert seen == {"coll_a", "coll_b"} + # Fused scores are descending RRF scores. + scores = [r["score"] for r in results] + assert scores == sorted(scores, reverse=True) + + def test_missing_single_collection_raises(self, kb): + with pytest.raises(ValueError, match="not found"): + kb.search("x", collection="nope", top_k=5) + + def test_partial_missing_skips_and_returns_found(self, kb): + results = kb.search( + "quality", collections=["coll_a", "nope"], top_k=5, rerank=False + ) + assert results + assert all(r["chunk"]["metadata"]["collection"] == "coll_a" for r in results) + + def test_all_missing_raises_all_failed(self, kb): + with pytest.raises(ValueError, match="All collections failed"): + kb.search("x", collections=["nope1", "nope2"], top_k=5) + + def test_no_target_raises(self, kb): + with pytest.raises(ValueError, match="Provide"): + kb.search("x", top_k=5) + + class TestHybridSearch: """End-to-end hybrid search behavior on a real KnowledgeBase.""" @@ -1235,8 +1226,8 @@ def kb_with_data(self, tmp_path): mock_client = MagicMock() mock_client.embed = fake_embed - with patch("dsagt.knowledge._make_embedder", return_value=mock_client): - kb = KnowledgeBase(index_dir=index_dir, default_index="faiss") + with patch("dsagt.knowledge.Embedder.create", return_value=mock_client): + kb = KnowledgeBase(index_dir=index_dir) kb.ingest(source_folder) yield kb kb.close() @@ -1256,46 +1247,6 @@ def test_search_succeeds_with_hybrid_default(self, kb_with_data): ) assert len(results) > 0 - def test_search_raises_when_hybrid_but_no_bm25(self, tmp_path): - """Hybrid=True with missing bm25.pkl is a loud failure.""" - from dsagt.knowledge import KnowledgeBase, CollectionRoute - - index_dir = tmp_path / "index" - source_folder = tmp_path / "docs" - source_folder.mkdir() - (source_folder / "doc.txt").write_text("Some content here.") - - mock_client = MagicMock() - mock_client.embed = fake_embed - - # Build with hybrid=False so no bm25.pkl is written, then flip the - # persisted route to hybrid=True to simulate a pre-hybrid collection - # in a now-hybrid-aware install. - with patch("dsagt.knowledge._make_embedder", return_value=mock_client): - kb = KnowledgeBase(index_dir=index_dir, default_index="faiss") - kb.ingest( - source_folder, - route=CollectionRoute( - embedding_backend="api", - vector_db="faiss", - hybrid=False, - ), - ) - kb.close() - - # Mutate the persisted route on disk. - import json as _json - route_path = index_dir / "docs" / "route.json" - route_data = _json.loads(route_path.read_text()) - route_data["hybrid"] = True - route_path.write_text(_json.dumps(route_data)) - - # Fresh KB instance: no in-memory cache, must read from disk. - kb2 = KnowledgeBase(index_dir=index_dir, default_index="faiss") - with pytest.raises(FileNotFoundError, match="bm25.pkl"): - kb2.search("anything", collection="docs", top_k=3, rerank=False) - kb2.close() - def test_add_entries_rebuilds_bm25(self, tmp_path): """add_entries must rebuild bm25.pkl with the new entry texts.""" from dsagt.knowledge import KnowledgeBase @@ -1303,8 +1254,8 @@ def test_add_entries_rebuilds_bm25(self, tmp_path): mock_client = MagicMock() mock_client.embed = fake_embed - with patch("dsagt.knowledge._make_embedder", return_value=mock_client): - kb = KnowledgeBase(index_dir=tmp_path / "index", default_index="faiss") + with patch("dsagt.knowledge.Embedder.create", return_value=mock_client): + kb = KnowledgeBase(index_dir=tmp_path / "index") kb.add_entries( texts=["alpha document", "beta document"], collection="memory", @@ -1314,32 +1265,6 @@ def test_add_entries_rebuilds_bm25(self, tmp_path): kb.add_entries(texts=["gamma document"], collection="memory") # BM25 must now know about all three docs. - bm25 = kb._get_bm25("memory") + bm25 = kb._store._get_bm25("memory") assert bm25.size == 3 kb.close() - - def test_route_persists_hybrid_flag(self, tmp_path): - """route.json must round-trip the hybrid field.""" - from dsagt.knowledge import KnowledgeBase, CollectionRoute - import json as _json - - mock_client = MagicMock() - mock_client.embed = fake_embed - - with patch("dsagt.knowledge._make_embedder", return_value=mock_client): - kb = KnowledgeBase(index_dir=tmp_path / "index", default_index="faiss") - kb.add_entries( - texts=["x"], - collection="c", - route=CollectionRoute( - embedding_backend="api", - vector_db="faiss", - hybrid=False, - ), - ) - kb.close() - - route_data = _json.loads( - (tmp_path / "index" / "c" / "route.json").read_text() - ) - assert route_data["hybrid"] is False diff --git a/tests/test_knowledge_integration.py b/tests/test_knowledge_integration.py index 5ff6fe5..e125c03 100755 --- a/tests/test_knowledge_integration.py +++ b/tests/test_knowledge_integration.py @@ -15,7 +15,6 @@ import asyncio import json -import os from pathlib import Path import pytest @@ -23,8 +22,11 @@ pytestmark = pytest.mark.integration from dsagt.knowledge import KnowledgeBase -from dsagt.commands.knowledge_server import create_knowledge_server, setup_runtime_kb -from mcp_helpers import call_tool_json as call_tool, call_tool_async as _call_tool_async_raw +from dsagt.mcp.knowledge_tools import create_knowledge_server, setup_runtime_kb +from mcp_helpers import ( + call_tool_json as call_tool, + call_tool_async as _call_tool_async_raw, +) async def _call_tool_async(server, name: str, arguments: dict) -> dict: @@ -32,7 +34,9 @@ async def _call_tool_async(server, name: str, arguments: dict) -> dict: return json.loads(await _call_tool_async_raw(server, name, arguments)) -async def call_tool_and_await_job(server, name: str, arguments: dict) -> tuple[dict, dict]: +async def call_tool_and_await_job( + server, name: str, arguments: dict +) -> tuple[dict, dict]: """Call a tool that starts a background job, wait for it, return (initial, final).""" initial = await _call_tool_async(server, name, arguments) assert initial["status"] == "started" @@ -51,6 +55,7 @@ async def call_tool_and_await_job(server, name: str, arguments: dict) -> tuple[d # Fixtures # --------------------------------------------------------------------------- + @pytest.fixture def smoke_test_dir(): """Path to the smoke test knowledge documents.""" @@ -85,7 +90,7 @@ def kb_server(tmp_path, smoke_test_dir, embedding_config): result = kb.ingest(smoke_test_dir) assert result["chunks"] > 0, "Ingest produced no chunks" - server = create_knowledge_server(kb, use_rerank=False) + server = create_knowledge_server(kb) yield server kb.close() @@ -94,12 +99,13 @@ def kb_server(tmp_path, smoke_test_dir, embedding_config): # Ingest via MCP handler # --------------------------------------------------------------------------- + class TestIngestIntegration: def test_ingest_smoke_test_docs(self, tmp_path, smoke_test_dir, embedding_config): """Ingest smoke test documents through the MCP handler.""" kb = _make_kb(tmp_path, embedding_config) - server = create_knowledge_server(kb, use_rerank=False) + server = create_knowledge_server(kb) async def run(): initial, final = await call_tool_and_await_job( @@ -121,14 +127,19 @@ async def run(): # Search via MCP handler # --------------------------------------------------------------------------- + class TestSearchIntegration: def test_search_returns_results(self, kb_server): """Search through MCP handler returns relevant chunks.""" - result = call_tool(kb_server, "kb_search", { - "query": "how to handle large files", - "collection": "knowledge", - }) + result = call_tool( + kb_server, + "kb_search", + { + "query": "how to handle large files", + "collection": "knowledge", + }, + ) assert result["status"] == "ok" assert result["result_count"] > 0 @@ -137,45 +148,61 @@ def test_search_returns_results(self, kb_server): def test_search_nonexistent_collection(self, kb_server): """Searching a nonexistent collection returns an error.""" - result = call_tool(kb_server, "kb_search", { - "query": "anything", - "collection": "nonexistent", - }) + result = call_tool( + kb_server, + "kb_search", + { + "query": "anything", + "collection": "nonexistent", + }, + ) assert result["status"] == "error" def test_search_with_top_k(self, kb_server): """top_k limits the number of results.""" - result = call_tool(kb_server, "kb_search", { - "query": "installation", - "collection": "knowledge", - "top_k": 2, - }) + result = call_tool( + kb_server, + "kb_search", + { + "query": "installation", + "collection": "knowledge", + "top_k": 2, + }, + ) assert result["status"] == "ok" assert result["result_count"] <= 2 def test_search_result_format(self, kb_server): """Each result has expected fields.""" - result = call_tool(kb_server, "kb_search", { - "query": "API reference", - "collection": "knowledge", - "top_k": 1, - }) + result = call_tool( + kb_server, + "kb_search", + { + "query": "API reference", + "collection": "knowledge", + "top_k": 1, + }, + ) assert result["status"] == "ok" - if result["result_count"] > 0: - hit = result["results"][0] - assert "text" in hit - assert "score" in hit - assert "source_file" in hit - assert "chunk_index" in hit + # The smoke-test corpus is fixed and owned by this test, so a known + # query must return at least one hit — then the shape is checked + # unconditionally (a zero-hit regression must not pass vacuously). + assert result["result_count"] > 0 + hit = result["results"][0] + assert "text" in hit + assert "score" in hit + assert "source_file" in hit + assert "chunk_index" in hit # --------------------------------------------------------------------------- # List collections via MCP handler # --------------------------------------------------------------------------- + class TestListCollectionsIntegration: def test_list_after_ingest(self, kb_server): @@ -189,32 +216,21 @@ def test_list_after_ingest(self, kb_server): # --------------------------------------------------------------------------- -# Setup runtime KB with symlinks +# Setup runtime KB (copies base collections into the project runtime) # --------------------------------------------------------------------------- -class TestSetupRuntimeKB: - def test_symlinks_base_collections(self, tmp_path, smoke_test_dir, embedding_config): - """setup_runtime_kb symlinks base collections into runtime.""" - base_dir = tmp_path / "base_kb" - base_dir.mkdir() - kb = _make_kb(tmp_path, embedding_config) - # Point KB at base_dir for this test - kb.index_dir = base_dir - kb.index_dir.mkdir(exist_ok=True) - kb.ingest(smoke_test_dir) - kb.close() - - runtime_dir = tmp_path / "runtime" - runtime_kb_dir = setup_runtime_kb(base_dir, runtime_dir) +class TestSetupRuntimeKB: - knowledge_link = runtime_kb_dir / "knowledge" - assert knowledge_link.exists() - assert knowledge_link.is_symlink() - assert (knowledge_link / "index.faiss").exists() + def test_runtime_search_after_setup(self, tmp_path, smoke_test_dir, embedding_config): + """End-to-end: a KB pointed at a runtime dir provisioned by + setup_runtime_kb can search the copied collections with real embeddings. - def test_runtime_search_via_symlink(self, tmp_path, smoke_test_dir, embedding_config): - """A KB pointing at runtime symlinks can search successfully.""" + Copy-vs-symlink structure is unit-tested network-free in + test_knowledge_server.py::TestSetupRuntimeKb; this is the only place + that exercises a real embedding search *through* the provisioned + runtime KB, so it stays an integration test. + """ base_dir = tmp_path / "base_kb" base_dir.mkdir() kb = _make_kb(tmp_path, embedding_config) diff --git a/tests/test_knowledge_server.py b/tests/test_knowledge_server.py index 1504490..0f35356 100644 --- a/tests/test_knowledge_server.py +++ b/tests/test_knowledge_server.py @@ -12,14 +12,13 @@ import asyncio import json -import time -from pathlib import Path -from unittest.mock import MagicMock, patch +import threading +from unittest.mock import MagicMock import pytest import mcp.types as types -from dsagt.commands.knowledge_server import create_knowledge_server, setup_runtime_kb +from dsagt.mcp.knowledge_tools import create_knowledge_server, setup_runtime_kb from mcp_helpers import call_tool_json as call_tool @@ -34,7 +33,9 @@ async def _call_tool_async(server, name: str, arguments: dict) -> dict: return json.loads(result.root.content[0].text) -async def call_tool_and_await_job(server, name: str, arguments: dict) -> tuple[dict, dict]: +async def call_tool_and_await_job( + server, name: str, arguments: dict +) -> tuple[dict, dict]: """Call a tool that starts a background job, wait for it, return (initial, final).""" initial = await _call_tool_async(server, name, arguments) assert initial["status"] == "started" @@ -50,7 +51,9 @@ async def call_tool_and_await_job(server, name: str, arguments: dict) -> tuple[d raise TimeoutError(f"Job {job_id} did not complete") -def make_search_result(text: str, source_file: str, chunk_index: int = 0, score: float = 0.9): +def make_search_result( + text: str, source_file: str, chunk_index: int = 0, score: float = 0.9 +): """Create a search result in the format KnowledgeBase.search returns.""" return { "chunk": { @@ -70,6 +73,7 @@ def make_search_result(text: str, source_file: str, chunk_index: int = 0, score: # Fixtures # --------------------------------------------------------------------------- + @pytest.fixture def mock_kb(tmp_path): """A mocked KnowledgeBase with default behaviors.""" @@ -86,7 +90,12 @@ def mock_kb(tmp_path): make_search_result("Second result text", "/path/to/file2.md", 1, 0.80), ] kb.ingest.return_value = {"collection": "new_docs", "files": 5, "chunks": 42} - kb.append.return_value = {"collection": "docs", "files": 2, "chunks_added": 10, "total_chunks": 50} + kb.append.return_value = { + "collection": "docs", + "files": 2, + "chunks_added": 10, + "total_chunks": 50, + } return kb @@ -100,6 +109,7 @@ def server(mock_kb): # kb_list_collections # --------------------------------------------------------------------------- + class TestListCollections: def test_returns_collections(self, server, mock_kb): @@ -128,14 +138,19 @@ def test_empty_collections(self, mock_kb): # kb_search # --------------------------------------------------------------------------- + class TestSearch: def test_search_success(self, server, mock_kb): """Successful search returns formatted results.""" - result = call_tool(server, "kb_search", { - "query": "how to install", - "collection": "docs", - }) + result = call_tool( + server, + "kb_search", + { + "query": "how to install", + "collection": "docs", + }, + ) assert result["status"] == "ok" assert result["query"] == "how to install" @@ -150,42 +165,60 @@ def test_search_success(self, server, mock_kb): def test_search_passes_parameters(self, server, mock_kb): """Search forwards top_k and rerank to the knowledge base.""" - call_tool(server, "kb_search", { - "query": "test", - "collection": "docs", - "top_k": 10, - "rerank": False, - }) + call_tool( + server, + "kb_search", + { + "query": "test", + "collection": "docs", + "top_k": 10, + "rerank": False, + }, + ) mock_kb.search.assert_called_once_with( query="test", collection="docs", + collections=None, top_k=10, rerank=False, + where=None, + where_document=None, ) def test_search_defaults(self, server, mock_kb): """Search uses default top_k=5 and server's use_rerank setting.""" - call_tool(server, "kb_search", { - "query": "test", - "collection": "docs", - }) + call_tool( + server, + "kb_search", + { + "query": "test", + "collection": "docs", + }, + ) mock_kb.search.assert_called_once_with( query="test", collection="docs", + collections=None, top_k=5, rerank=None, # agent didn't specify → kb.default_rerank resolves it + where=None, + where_document=None, ) def test_search_nonexistent_collection(self, server, mock_kb): """Searching a missing collection returns an error.""" mock_kb.search.side_effect = ValueError("Collection 'missing' not found") - result = call_tool(server, "kb_search", { - "query": "test", - "collection": "missing", - }) + result = call_tool( + server, + "kb_search", + { + "query": "test", + "collection": "missing", + }, + ) assert result["status"] == "error" assert "not found" in result["error"] @@ -196,10 +229,14 @@ def test_search_rerank_score_forwarded(self, server, mock_kb): {**make_search_result("text", "file.md"), "rerank_score": 0.99}, ] - result = call_tool(server, "kb_search", { - "query": "test", - "collection": "docs", - }) + result = call_tool( + server, + "kb_search", + { + "query": "test", + "collection": "docs", + }, + ) assert result["results"][0]["rerank_score"] == 0.99 @@ -208,6 +245,7 @@ def test_search_rerank_score_forwarded(self, server, mock_kb): # kb_ingest (background job pattern) # --------------------------------------------------------------------------- + class TestIngest: def test_ingest_returns_started(self, server, mock_kb, tmp_path): @@ -215,9 +253,13 @@ def test_ingest_returns_started(self, server, mock_kb, tmp_path): folder = tmp_path / "new_docs" folder.mkdir() - result = call_tool(server, "kb_ingest", { - "folder_path": str(folder), - }) + result = call_tool( + server, + "kb_ingest", + { + "folder_path": str(folder), + }, + ) assert result["status"] == "started" assert "job_id" in result @@ -245,23 +287,31 @@ def test_ingest_with_file_types(self, server, mock_kb, tmp_path): async def run(): await call_tool_and_await_job( - server, "kb_ingest", { + server, + "kb_ingest", + { "folder_path": str(folder), "file_types": ["md", "txt"], - } + }, ) # New server always passes collection_name to kb.ingest mock_kb.ingest.assert_called_once_with( - folder, collection_name="docs2", file_types=["md", "txt"], + folder, + collection_name="docs2", + file_types=["md", "txt"], ) asyncio.run(run()) def test_ingest_folder_not_found(self, server): """Ingesting a nonexistent folder returns an error immediately.""" - result = call_tool(server, "kb_ingest", { - "folder_path": "/nonexistent/folder", - }) + result = call_tool( + server, + "kb_ingest", + { + "folder_path": "/nonexistent/folder", + }, + ) assert result["status"] == "error" assert "not found" in result["error"].lower() @@ -271,9 +321,13 @@ def test_ingest_not_a_directory(self, server, tmp_path): file_path = tmp_path / "not_a_dir.txt" file_path.write_text("I'm a file") - result = call_tool(server, "kb_ingest", { - "folder_path": str(file_path), - }) + result = call_tool( + server, + "kb_ingest", + { + "folder_path": str(file_path), + }, + ) assert result["status"] == "error" assert "Not a directory" in result["error"] @@ -298,19 +352,23 @@ def test_ingest_deconflicts_existing_collection(self, server, mock_kb, tmp_path) folder = tmp_path / "docs" folder.mkdir() - # Simulate "docs" already exists with a FAISS index from a different source. + # Simulate "docs" already exists with a Chroma index from a different source. # _collection_exists() requires a marker file, and deconflict only triggers # when source.txt records a different folder than the one being ingested. existing = mock_kb.index_dir / "docs" existing.mkdir() - (existing / "index.faiss").write_bytes(b"fake") + (existing / "chroma_ids.json").write_bytes(b"fake") (existing / "source.txt").write_text("/some/other/folder") mock_kb.ingest.return_value = {"collection": "docs1", "files": 3, "chunks": 10} - result = call_tool(server, "kb_ingest", { - "folder_path": str(folder), - }) + result = call_tool( + server, + "kb_ingest", + { + "folder_path": str(folder), + }, + ) assert result["status"] == "started" assert result["collection"] == "docs1" @@ -325,15 +383,19 @@ def test_ingest_deconflicts_symlinked_collection(self, server, mock_kb, tmp_path # Simulate "docs" is a symlink to a base collection with index base_dir = tmp_path / "base_docs" base_dir.mkdir() - (base_dir / "index.faiss").write_bytes(b"fake") + (base_dir / "chroma_ids.json").write_bytes(b"fake") (base_dir / "source.txt").write_text("/some/other/folder") (mock_kb.index_dir / "docs").symlink_to(base_dir) mock_kb.ingest.return_value = {"collection": "docs1", "files": 3, "chunks": 10} - result = call_tool(server, "kb_ingest", { - "folder_path": str(folder), - }) + result = call_tool( + server, + "kb_ingest", + { + "folder_path": str(folder), + }, + ) assert result["status"] == "started" assert result["collection"] == "docs1" @@ -347,6 +409,7 @@ def test_ingest_deconflicts_symlinked_collection(self, server, mock_kb, tmp_path # kb_job_status # --------------------------------------------------------------------------- + class TestJobStatus: def test_unknown_job(self, server): @@ -361,22 +424,37 @@ def test_running_job(self, server, mock_kb, tmp_path): folder = tmp_path / "slow_docs" folder.mkdir() - # Make ingest block so the job stays in "running" + # Hold the job in "running" until the test has observed it. + # asyncio.run joins the worker thread at loop shutdown, so the + # event must be set before ``run()`` returns; the wait timeout + # only bounds a failing test. + release = threading.Event() + def blocking_ingest(*args, **kwargs): - time.sleep(10) + release.wait(timeout=10) return {"collection": "slow_docs", "files": 1, "chunks": 5} + mock_kb.ingest.side_effect = blocking_ingest async def run(): - initial = await _call_tool_async(server, "kb_ingest", { - "folder_path": str(folder), - }) - assert initial["status"] == "started" - job_id = initial["job_id"] - - # Immediately check — should still be running - status = await _call_tool_async(server, "kb_job_status", {"job_id": job_id}) - assert status["status"] == "running" + try: + initial = await _call_tool_async( + server, + "kb_ingest", + { + "folder_path": str(folder), + }, + ) + assert initial["status"] == "started" + job_id = initial["job_id"] + + # Immediately check — should still be running + status = await _call_tool_async( + server, "kb_job_status", {"job_id": job_id} + ) + assert status["status"] == "running" + finally: + release.set() asyncio.run(run()) @@ -385,6 +463,7 @@ async def run(): # kb_append (background job pattern) # --------------------------------------------------------------------------- + class TestAppend: def test_append_returns_started(self, server, mock_kb, tmp_path): @@ -392,12 +471,16 @@ def test_append_returns_started(self, server, mock_kb, tmp_path): # Create a fake existing collection coll_dir = mock_kb.index_dir / "docs" coll_dir.mkdir(exist_ok=True) - (coll_dir / "index.faiss").write_text("fake") - - result = call_tool(server, "kb_append", { - "collection": "docs", - "paths": [str(tmp_path)], - }) + (coll_dir / "chroma_ids.json").write_text("fake") + + result = call_tool( + server, + "kb_append", + { + "collection": "docs", + "paths": [str(tmp_path)], + }, + ) assert result["status"] == "started" assert "job_id" in result @@ -407,14 +490,16 @@ def test_append_job_completes(self, server, mock_kb, tmp_path): """Background append job completes successfully.""" coll_dir = mock_kb.index_dir / "docs" coll_dir.mkdir(exist_ok=True) - (coll_dir / "index.faiss").write_text("fake") + (coll_dir / "chroma_ids.json").write_text("fake") async def run(): initial, final = await call_tool_and_await_job( - server, "kb_append", { + server, + "kb_append", + { "collection": "docs", "paths": [str(tmp_path)], - } + }, ) assert final["status"] == "complete" assert final["result"]["chunks_added"] == 10 @@ -423,10 +508,14 @@ async def run(): def test_append_collection_not_found(self, server, mock_kb): """Appending to a nonexistent collection returns an error immediately.""" - result = call_tool(server, "kb_append", { - "collection": "nonexistent", - "paths": ["/some/path"], - }) + result = call_tool( + server, + "kb_append", + { + "collection": "nonexistent", + "paths": ["/some/path"], + }, + ) assert result["status"] == "error" assert "not found" in result["error"].lower() @@ -436,6 +525,7 @@ def test_append_collection_not_found(self, server, mock_kb): # kb_search — error handling (transport-closed diagnostics) # --------------------------------------------------------------------------- + class TestSearchErrorHandling: """Verify the server returns error responses (not crashes) for common failure modes that would otherwise cause 'transport closed'.""" @@ -443,12 +533,18 @@ class TestSearchErrorHandling: def test_search_httpx_connect_error(self, mock_kb): """Network unreachable during search returns error, not crash.""" import httpx + mock_kb.search.side_effect = httpx.ConnectError("Connection refused") server = create_knowledge_server(mock_kb) - result = call_tool(server, "kb_search", { - "query": "test", "collection": "docs", - }) + result = call_tool( + server, + "kb_search", + { + "query": "test", + "collection": "docs", + }, + ) assert result["status"] == "error" assert "Connection refused" in result["error"] @@ -456,12 +552,18 @@ def test_search_httpx_connect_error(self, mock_kb): def test_search_httpx_timeout(self, mock_kb): """Embedding API timeout during search returns error, not crash.""" import httpx + mock_kb.search.side_effect = httpx.ReadTimeout("Read timed out") server = create_knowledge_server(mock_kb) - result = call_tool(server, "kb_search", { - "query": "test", "collection": "docs", - }) + result = call_tool( + server, + "kb_search", + { + "query": "test", + "collection": "docs", + }, + ) assert result["status"] == "error" assert "timed out" in result["error"].lower() @@ -469,16 +571,24 @@ def test_search_httpx_timeout(self, mock_kb): def test_search_httpx_401(self, mock_kb): """Expired/invalid API key during search returns error, not crash.""" import httpx + mock_resp = MagicMock() mock_resp.status_code = 401 mock_kb.search.side_effect = httpx.HTTPStatusError( - "401 Unauthorized", request=MagicMock(), response=mock_resp, + "401 Unauthorized", + request=MagicMock(), + response=mock_resp, ) server = create_knowledge_server(mock_kb) - result = call_tool(server, "kb_search", { - "query": "test", "collection": "docs", - }) + result = call_tool( + server, + "kb_search", + { + "query": "test", + "collection": "docs", + }, + ) assert result["status"] == "error" assert "401" in result["error"] @@ -486,40 +596,58 @@ def test_search_httpx_401(self, mock_kb): def test_search_httpx_500(self, mock_kb): """Embedding API server error returns error, not crash.""" import httpx + mock_resp = MagicMock() mock_resp.status_code = 500 mock_kb.search.side_effect = httpx.HTTPStatusError( - "500 Internal Server Error", request=MagicMock(), response=mock_resp, + "500 Internal Server Error", + request=MagicMock(), + response=mock_resp, ) server = create_knowledge_server(mock_kb) - result = call_tool(server, "kb_search", { - "query": "test", "collection": "docs", - }) + result = call_tool( + server, + "kb_search", + { + "query": "test", + "collection": "docs", + }, + ) assert result["status"] == "error" assert "500" in result["error"] def test_search_runtime_error(self, mock_kb): """Unexpected RuntimeError during search returns error, not crash.""" - mock_kb.search.side_effect = RuntimeError("FAISS segfault simulation") + mock_kb.search.side_effect = RuntimeError("index segfault simulation") server = create_knowledge_server(mock_kb) - result = call_tool(server, "kb_search", { - "query": "test", "collection": "docs", - }) + result = call_tool( + server, + "kb_search", + { + "query": "test", + "collection": "docs", + }, + ) assert result["status"] == "error" - assert "FAISS segfault" in result["error"] + assert "index segfault" in result["error"] def test_search_os_error(self, mock_kb): """OS-level error (disk, permissions) returns error, not crash.""" - mock_kb.search.side_effect = OSError("Permission denied: index.faiss") + mock_kb.search.side_effect = OSError("Permission denied: chroma.sqlite3") server = create_knowledge_server(mock_kb) - result = call_tool(server, "kb_search", { - "query": "test", "collection": "docs", - }) + result = call_tool( + server, + "kb_search", + { + "query": "test", + "collection": "docs", + }, + ) assert result["status"] == "error" assert "Permission denied" in result["error"] @@ -530,6 +658,7 @@ def test_search_os_error(self, mock_kb): # setup_runtime_kb # --------------------------------------------------------------------------- + class TestSetupRuntimeKb: def test_copies_collections(self, tmp_path): @@ -543,7 +672,7 @@ def test_copies_collections(self, tmp_path): base = tmp_path / "base_index" coll_dir = base / "my_collection" coll_dir.mkdir(parents=True) - (coll_dir / "index.faiss").write_text("fake index") + (coll_dir / "chroma_ids.json").write_text("fake index") (coll_dir / "chunks.jsonl").write_text('{"id": "1"}\n') (coll_dir / "DESCRIPTION.md").write_text("Test collection") @@ -554,30 +683,30 @@ def test_copies_collections(self, tmp_path): copied = result / "my_collection" assert copied.exists() assert not copied.is_symlink() # copy not symlink - assert (copied / "index.faiss").exists() - assert (copied / "index.faiss").read_text() == "fake index" + assert (copied / "chroma_ids.json").exists() + assert (copied / "chroma_ids.json").read_text() == "fake index" assert (copied / "chunks.jsonl").exists() assert (copied / "DESCRIPTION.md").exists() def test_copy_is_independent(self, tmp_path): """Mutating the base after copy does not affect the project copy.""" base = tmp_path / "base_index" - coll = base / "tools" + coll = base / "codes" coll.mkdir(parents=True) - (coll / "index.faiss").write_text("v1") + (coll / "chroma_ids.json").write_text("v1") runtime = tmp_path / "runtime" setup_runtime_kb(base, runtime) # Mutate the base — simulating ``dsagt setup-kb --rebuild``. - (coll / "index.faiss").write_text("v2 newer") + (coll / "chroma_ids.json").write_text("v2 newer") # Project copy stays at v1. - project_copy = runtime / "kb_index" / "tools" / "index.faiss" + project_copy = runtime / "kb_index" / "codes" / "chroma_ids.json" assert project_copy.read_text() == "v1" def test_skips_non_collection_dirs(self, tmp_path): - """Directories without index.faiss are not copied.""" + """Directories without chroma_ids.json are not copied.""" base = tmp_path / "base_index" (base / "random_dir").mkdir(parents=True) (base / "random_dir" / "notes.txt").write_text("not a collection") @@ -600,34 +729,35 @@ def test_does_not_overwrite_existing(self, tmp_path): base = tmp_path / "base_index" coll = base / "docs" coll.mkdir(parents=True) - (coll / "index.faiss").write_text("base version") + (coll / "chroma_ids.json").write_text("base version") runtime = tmp_path / "runtime" runtime_coll = runtime / "kb_index" / "docs" runtime_coll.mkdir(parents=True) - (runtime_coll / "index.faiss").write_text("runtime version") + (runtime_coll / "chroma_ids.json").write_text("runtime version") setup_runtime_kb(base, runtime) - assert (runtime_coll / "index.faiss").read_text() == "runtime version" + assert (runtime_coll / "chroma_ids.json").read_text() == "runtime version" # --------------------------------------------------------------------------- # Regression: OpenMP duplicate library crash (transport closed) # --------------------------------------------------------------------------- + class TestOpenMPWorkaround: - """Importing knowledge_server must set KMP_DUPLICATE_LIB_OK to prevent - a fatal OpenMP crash when FAISS and sentence-transformers (PyTorch) - both bundle libomp. + """Importing the knowledge tools module must set KMP_DUPLICATE_LIB_OK to + prevent a fatal OpenMP crash when multiple native deps (e.g. ChromaDB and + sentence-transformers / PyTorch) both bundle libomp. Without this, kb_search with rerank=true kills the server process, producing 'transport closed' in MCP clients.""" def test_kmp_duplicate_lib_ok_is_set(self): - """KMP_DUPLICATE_LIB_OK is set after importing the knowledge server.""" + """KMP_DUPLICATE_LIB_OK is set after importing dsagt.mcp.knowledge_tools.""" import os - import dsagt.commands.knowledge_server # noqa: F401 + import dsagt.mcp.knowledge_tools # noqa: F401 assert os.environ.get("KMP_DUPLICATE_LIB_OK") == "TRUE" @@ -636,6 +766,7 @@ def test_kmp_duplicate_lib_ok_is_set(self): # Regression: rerank schema default must match server config # --------------------------------------------------------------------------- + class TestRerankSchemaDefault: """The kb_search schema previously hardcoded 'default': True for the rerank parameter, causing agents to request reranking even when the @@ -661,17 +792,93 @@ def test_rerank_default_from_kb(self, mock_kb): server = create_knowledge_server(mock_kb) assert self._get_rerank_default(server) is True - def test_search_omitted_rerank_passes_none(self, mock_kb): - """Omitting rerank passes None to kb.search, which resolves to - kb.default_rerank internally.""" - server = create_knowledge_server(mock_kb) - call_tool(server, "kb_search", { - "query": "test", - "collection": "docs", - }) + +# --------------------------------------------------------------------------- +# kb_search — multi-collection fan-out (moved from the former memory test file) +# --------------------------------------------------------------------------- + + +class TestKbSearchMultiCollection: + + def test_multi_collection_fanout(self, server, mock_kb): + """Multi-collection search delegates once to kb.search with collections=. + + Fan-out + fusion across collections is kb.search's job (covered by + TestFederatedSearch in test_knowledge_base.py); the handler just forwards. + """ + mock_kb.search.return_value = [ + make_search_result("result", "/file.md", 0, 0.9), + ] + + result = call_tool( + server, + "kb_search", + { + "query": "test", + "collections": ["docs", "papers"], + }, + ) + + assert result["status"] == "ok" mock_kb.search.assert_called_once_with( query="test", - collection="docs", + collection=None, + collections=["docs", "papers"], top_k=5, rerank=None, + where=None, + where_document=None, + ) + + def test_no_collection_returns_error(self, server): + result = call_tool( + server, + "kb_search", + { + "query": "test", + }, + ) + + assert result["status"] == "error" + + def test_multi_collection_merges_results(self, server, mock_kb): + """The handler returns kb.search's already-fused, sorted results.""" + # kb.search owns fusion now; it returns one merged, descending list. + mock_kb.search.return_value = [ + make_search_result("result_1", "/file_1.md", score=0.9), + make_search_result("result_2", "/file_2.md", score=0.7), + ] + + result = call_tool( + server, + "kb_search", + { + "query": "test", + "collections": ["docs", "papers"], + "top_k": 5, + }, ) + + assert result["result_count"] == 2 + scores = [r["score"] for r in result["results"]] + assert scores == sorted(scores, reverse=True) + + +class TestKbSearchSchema: + + def _get_tool(self, server, name): + req = types.ListToolsRequest(method="tools/list") + handler = server.request_handlers[types.ListToolsRequest] + result = asyncio.run(handler(req)) + for tool in result.root.tools: + if tool.name == name: + return tool + return None + + def test_kb_search_has_collections_param(self, server): + tool = self._get_tool(server, "kb_search") + assert "collections" in tool.inputSchema["properties"] + + def test_kb_search_query_is_only_required(self, server): + tool = self._get_tool(server, "kb_search") + assert tool.inputSchema["required"] == ["query"] diff --git a/tests/test_knowledge_server_memory.py b/tests/test_knowledge_server_memory.py deleted file mode 100644 index a4d1e9f..0000000 --- a/tests/test_knowledge_server_memory.py +++ /dev/null @@ -1,278 +0,0 @@ -""" -Tests for kb_remember, kb_get_memories, and extended kb_search handlers. - -Drop this file into tests/test_knowledge_server_memory.py - -These tests use a real ExplicitMemory (file-backed, no mocking needed) -and a mocked KnowledgeBase (same pattern as existing server tests). -""" - -import asyncio -from pathlib import Path -from unittest.mock import MagicMock - -import pytest -import mcp.types as types - -from dsagt.commands.knowledge_server import create_knowledge_server -from dsagt.memory import ExplicitMemory -from mcp_helpers import call_tool_json as call_tool - - -def make_search_result(text, source_file, chunk_index=0, score=0.9): - return { - "chunk": { - "text": text, - "metadata": { - "source_file": source_file, - "collection": "test_collection", - "chunk_index": chunk_index, - "file_type": ".md", - }, - }, - "score": score, - } - - -# --------------------------------------------------------------------------- -# Fixtures -# --------------------------------------------------------------------------- - - -@pytest.fixture -def mock_kb(tmp_path): - kb = MagicMock() - kb.index_dir = tmp_path / "kb_index" - kb.index_dir.mkdir() - kb.list_collections.return_value = [ - {"name": "docs", "description": "Documentation"}, - ] - kb.search.return_value = [ - make_search_result("Result one", "/path/to/file.md", 0, 0.95), - ] - kb.ingest.return_value = {"collection": "docs", "files": 1, "chunks": 10} - kb.append.return_value = {"collection": "docs", "files": 1, "chunks_added": 5, "total_chunks": 15} - return kb - - -@pytest.fixture -def server(mock_kb, tmp_path): - return create_knowledge_server(mock_kb, runtime_dir=tmp_path) - - -@pytest.fixture -def memory(tmp_path): - return ExplicitMemory(runtime_dir=tmp_path) - - -# --------------------------------------------------------------------------- -# kb_remember -# --------------------------------------------------------------------------- - - -class TestKbRemember: - - def test_stores_a_fact(self, server): - result = call_tool(server, "kb_remember", { - "text": "fastp quality threshold is Q20", - }) - - assert result["status"] == "ok" - assert result["entry_id"] - assert result["total_memories"] == 1 - - def test_stores_with_metadata(self, server): - result = call_tool(server, "kb_remember", { - "text": "some fact", - "category": "quality_control", - "session_id": "sess_01", - }) - - assert result["status"] == "ok" - - def test_supersede_existing(self, server): - r1 = call_tool(server, "kb_remember", { - "text": "old threshold Q20", - }) - r2 = call_tool(server, "kb_remember", { - "text": "new threshold Q30", - "supersedes": r1["entry_id"], - }) - - assert r2["status"] == "ok" - assert r2["superseded_id"] == r1["entry_id"] - assert r2["total_memories"] == 1 - - def test_supersede_nonexistent_returns_error(self, server): - result = call_tool(server, "kb_remember", { - "text": "new fact", - "supersedes": "bad_id", - }) - - assert result["status"] == "error" - assert "not found" in result["error"] - - def test_multiple_facts(self, server): - call_tool(server, "kb_remember", {"text": "fact one"}) - call_tool(server, "kb_remember", {"text": "fact two"}) - result = call_tool(server, "kb_remember", {"text": "fact three"}) - - assert result["total_memories"] == 3 - - -# --------------------------------------------------------------------------- -# kb_get_memories -# --------------------------------------------------------------------------- - - -class TestKbGetMemories: - - def test_empty_returns_zero(self, server): - result = call_tool(server, "kb_get_memories", {}) - - assert result["status"] == "ok" - assert result["count"] == 0 - assert result["memories"] == [] - - def test_returns_stored_memories(self, server): - call_tool(server, "kb_remember", {"text": "fact one"}) - call_tool(server, "kb_remember", {"text": "fact two"}) - - result = call_tool(server, "kb_get_memories", {}) - - assert result["count"] == 2 - texts = {m["text"] for m in result["memories"]} - assert texts == {"fact one", "fact two"} - - def test_excludes_superseded(self, server): - r1 = call_tool(server, "kb_remember", {"text": "old fact"}) - call_tool(server, "kb_remember", { - "text": "new fact", - "supersedes": r1["entry_id"], - }) - - result = call_tool(server, "kb_get_memories", {}) - - assert result["count"] == 1 - assert result["memories"][0]["text"] == "new fact" - - -# --------------------------------------------------------------------------- -# kb_search — multi-collection fan-out -# --------------------------------------------------------------------------- - - -class TestKbSearchMultiCollection: - - def test_single_collection_backward_compat(self, server, mock_kb): - """Plain search with collection still works.""" - result = call_tool(server, "kb_search", { - "query": "test", - "collection": "docs", - }) - - assert result["status"] == "ok" - mock_kb.search.assert_called_once() - - def test_multi_collection_fanout(self, server, mock_kb): - """Searching multiple collections calls search for each.""" - mock_kb.search.return_value = [ - make_search_result("result", "/file.md", 0, 0.9), - ] - - result = call_tool(server, "kb_search", { - "query": "test", - "collections": ["docs", "papers"], - }) - - assert result["status"] == "ok" - assert mock_kb.search.call_count == 2 - - def test_no_collection_returns_error(self, server): - result = call_tool(server, "kb_search", { - "query": "test", - }) - - assert result["status"] == "error" - - def test_multi_collection_merges_results(self, server, mock_kb): - """Results from multiple collections are merged and sorted.""" - call_count = [0] - - def varying_results(**kwargs): - call_count[0] += 1 - score = 0.9 if call_count[0] == 1 else 0.7 - return [make_search_result( - f"result_{call_count[0]}", - f"/file_{call_count[0]}.md", - score=score, - )] - - mock_kb.search.side_effect = varying_results - - result = call_tool(server, "kb_search", { - "query": "test", - "collections": ["docs", "papers"], - "top_k": 5, - }) - - assert result["result_count"] == 2 - scores = [r["score"] for r in result["results"]] - assert scores == sorted(scores, reverse=True) - - def test_missing_collection_skipped(self, server, mock_kb): - """A missing collection logs a warning but doesn't fail the search.""" - def search_with_error(**kwargs): - if kwargs["collection"] == "missing": - raise ValueError("Collection 'missing' not found") - return [make_search_result("result", "/file.md")] - - mock_kb.search.side_effect = search_with_error - - result = call_tool(server, "kb_search", { - "query": "test", - "collections": ["docs", "missing"], - }) - - assert result["status"] == "ok" - assert result["result_count"] == 1 - - -# --------------------------------------------------------------------------- -# Tool schemas -# --------------------------------------------------------------------------- - - -class TestToolSchemas: - - def _get_tool(self, server, name): - req = types.ListToolsRequest(method="tools/list") - handler = server.request_handlers[types.ListToolsRequest] - result = asyncio.run(handler(req)) - for tool in result.root.tools: - if tool.name == name: - return tool - return None - - def test_kb_remember_exists(self, server): - tool = self._get_tool(server, "kb_remember") - assert tool is not None - assert "text" in tool.inputSchema["properties"] - assert tool.inputSchema["required"] == ["text"] - - def test_kb_remember_has_optional_params(self, server): - tool = self._get_tool(server, "kb_remember") - for param in ("category", "session_id", "supersedes"): - assert param in tool.inputSchema["properties"] - - def test_kb_get_memories_exists(self, server): - tool = self._get_tool(server, "kb_get_memories") - assert tool is not None - - def test_kb_search_has_collections_param(self, server): - tool = self._get_tool(server, "kb_search") - assert "collections" in tool.inputSchema["properties"] - - def test_kb_search_query_is_only_required(self, server): - tool = self._get_tool(server, "kb_search") - assert tool.inputSchema["required"] == ["query"] diff --git a/tests/test_memory.py b/tests/test_memory.py index fdadf36..d09467d 100644 --- a/tests/test_memory.py +++ b/tests/test_memory.py @@ -232,17 +232,13 @@ def test_handles_empty_file(self, tmp_path): assert mem.count() == 0 def test_handles_corrupt_file(self, tmp_path): - """Corrupt YAML either returns empty list or raises YAMLError.""" + """Corrupt YAML fails fast — get_all surfaces the YAMLError, never + silently recovers to an empty list (which would hide data loss).""" mem = ExplicitMemory(runtime_dir=tmp_path) (tmp_path / ExplicitMemory.FILENAME).write_text("not: valid: yaml: [") - try: - result = mem.get_all() - # If it didn't raise, it should return an empty list (graceful recovery) - assert isinstance(result, list) - except yaml.YAMLError: - # Failing fast on corruption is also acceptable - pass + with pytest.raises(yaml.YAMLError): + mem.get_all() def test_file_is_human_readable(self, mem, tmp_path): mem.remember("readable fact", category="test") diff --git a/tests/test_memory_extractor.py b/tests/test_memory_extractor.py new file mode 100644 index 0000000..7015035 --- /dev/null +++ b/tests/test_memory_extractor.py @@ -0,0 +1,136 @@ +"""Tests for the episodic-memory subscriber (MemoryExtractor). + +The mechanical write path: per-block chunking with producer/tool/turn_id +metadata. The KB is faked so these stay unit tests with no embedding model. +""" + +from dsagt.memory import SESSION_MEMORY_COLLECTION, MemoryExtractor +from dsagt.traces import Trace + + +class FakeKB: + def __init__(self): + self.calls = [] + + def add_entries(self, *, texts, collection, metadatas): + self.calls.append( + {"texts": texts, "collection": collection, "metadatas": metadatas} + ) + return {} + + +def _one_turn_trace(): + """A trace whose ``to_exchanges`` yields one exchange.""" + trace = Trace("t", "proj:s", "claude", "proj") + trace.add_agent_root("r1", "conv", start_time=1.0, prompt="filter the reads") + trace.add_llm_span( + "r1-0", + parent_id="r1", + start_time=1.0, + end_time=None, + request=[ + {"role": "user", "content": [{"type": "text", "text": "filter the reads"}]} + ], + response=[{"type": "text", "text": "kept 90 percent after QC"}], + ) + return trace + + +def _tool_trace(): + """A two-turn trace: the assistant calls a tool, the result returns next turn.""" + trace = Trace("t", "proj:s", "claude", "proj") + trace.add_agent_root("r1", "conv", start_time=1.0, prompt="profile it") + trace.add_llm_span( + "r1-0", + parent_id="r1", + start_time=1.0, + end_time=None, + request=[ + {"role": "user", "content": [{"type": "text", "text": "profile sales.csv"}]} + ], + response=[{"type": "tool_use", "id": "t1", "name": "profile", "input": {}}], + ) + trace.add_llm_span( + "r1-1", + parent_id="r1", + start_time=2.0, + end_time=None, + request=[ + { + "role": "user", + "content": [ + {"type": "tool_result", "tool_use_id": "t1", "content": "1000 rows"} + ], + } + ], + response=[{"type": "text", "text": "The file has 1000 rows."}], + ) + return trace + + +def test_each_block_is_its_own_chunk_with_producer(tmp_path): + kb = FakeKB() + ext = MemoryExtractor(kb, runtime_dir=tmp_path, session_id="proj:s") + ext.write(_one_turn_trace()) + + assert len(kb.calls) == 1 + call = kb.calls[0] + assert call["collection"] == SESSION_MEMORY_COLLECTION + # The user question and the assistant answer are separate chunks/embeddings. + assert call["texts"] == ["filter the reads", "kept 90 percent after QC"] + metas = call["metadatas"] + assert [m["producer"] for m in metas] == ["user", "llm"] + assert all(m["turn_id"] == "r1-0" for m in metas) + assert all(m["source_type"] == "turn" for m in metas) + assert all(isinstance(m["ts_epoch"], float) for m in metas) # for recency + + +def test_tool_result_chunk_carries_producer_and_resolved_tool(tmp_path): + kb = FakeKB() + ext = MemoryExtractor(kb, runtime_dir=tmp_path, session_id="proj:s") + ext.write(_tool_trace()) + + texts = kb.calls[0]["texts"] + metas = kb.calls[0]["metadatas"] + # tool_use is not embedded as prose... + assert "profile" not in texts + # ...but the tool_result is a chunk tagged producer=tool with the name + # resolved across turns via tool_use_id. + i = texts.index("1000 rows") + assert metas[i]["producer"] == "tool" + assert metas[i]["tool_name"] == "profile" + assert metas[i]["turn_id"] == "r1-1" + + +def test_empty_trace_writes_nothing(tmp_path): + kb = FakeKB() + ext = MemoryExtractor(kb, runtime_dir=tmp_path, session_id="proj:s") + ext.write(Trace("t", "proj:s", "claude", "proj")) + assert kb.calls == [] + + +def test_extraction_span_tagged_episodic_not_memory(tmp_path): + """The per-turn embedding runs off the heartbeat, so it must carry + dsagt.source=episodic — filtering apart from the user-facing memory tools + (kb_remember / kb_get_memories) that carry dsagt.source=memory.""" + from unittest.mock import patch + + opened = {} + + class _Span: + def __enter__(self): + return self + + def __exit__(self, *a): + return False + + def _fake_open_span(name, span_type=None, source=None): + opened["source"] = source + return _Span() + + kb = FakeKB() + ext = MemoryExtractor(kb, runtime_dir=tmp_path, session_id="proj:s") + with patch("dsagt.observability.open_span", _fake_open_span): + ext.write(_one_turn_trace()) + + assert opened["source"] == "episodic" diff --git a/tests/test_memory_tools.py b/tests/test_memory_tools.py new file mode 100644 index 0000000..4b05111 --- /dev/null +++ b/tests/test_memory_tools.py @@ -0,0 +1,188 @@ +""" +Tests for the explicit-memory MCP tools (kb_remember, kb_get_memories). + +These tests use a real ExplicitMemory (file-backed, no mocking needed) and a +mocked KnowledgeBase (same pattern as the other server tests). The tools live +in :mod:`dsagt.mcp.memory_tools`; ``create_memory_server`` exposes just that +concern for driving via ``call_tool_sync``. +""" + +import asyncio +from unittest.mock import MagicMock + +import pytest +import mcp.types as types + +from dsagt.mcp.memory_tools import create_memory_server +from dsagt.memory import ExplicitMemory +from mcp_helpers import call_tool_json as call_tool + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + + +@pytest.fixture +def mock_kb(tmp_path): + kb = MagicMock() + kb.index_dir = tmp_path / "kb_index" + kb.index_dir.mkdir() + return kb + + +@pytest.fixture +def server(mock_kb, tmp_path): + return create_memory_server(mock_kb, runtime_dir=tmp_path) + + +@pytest.fixture +def memory(tmp_path): + return ExplicitMemory(runtime_dir=tmp_path) + + +# --------------------------------------------------------------------------- +# kb_remember +# --------------------------------------------------------------------------- + + +class TestKbRemember: + + def test_stores_a_fact(self, server): + result = call_tool( + server, + "kb_remember", + { + "text": "fastp quality threshold is Q20", + }, + ) + + assert result["status"] == "ok" + assert result["entry_id"] + assert result["total_memories"] == 1 + + def test_stores_with_metadata(self, server): + result = call_tool( + server, + "kb_remember", + { + "text": "some fact", + "category": "quality_control", + "session_id": "sess_01", + }, + ) + + assert result["status"] == "ok" + + def test_supersede_existing(self, server): + r1 = call_tool( + server, + "kb_remember", + { + "text": "old threshold Q20", + }, + ) + r2 = call_tool( + server, + "kb_remember", + { + "text": "new threshold Q30", + "supersedes": r1["entry_id"], + }, + ) + + assert r2["status"] == "ok" + assert r2["superseded_id"] == r1["entry_id"] + assert r2["total_memories"] == 1 + + def test_supersede_nonexistent_returns_error(self, server): + result = call_tool( + server, + "kb_remember", + { + "text": "new fact", + "supersedes": "bad_id", + }, + ) + + assert result["status"] == "error" + assert "not found" in result["error"] + + def test_multiple_facts(self, server): + call_tool(server, "kb_remember", {"text": "fact one"}) + call_tool(server, "kb_remember", {"text": "fact two"}) + result = call_tool(server, "kb_remember", {"text": "fact three"}) + + assert result["total_memories"] == 3 + + +# --------------------------------------------------------------------------- +# kb_get_memories +# --------------------------------------------------------------------------- + + +class TestKbGetMemories: + + def test_empty_returns_zero(self, server): + result = call_tool(server, "kb_get_memories", {}) + + assert result["status"] == "ok" + assert result["count"] == 0 + assert result["memories"] == [] + + def test_returns_stored_memories(self, server): + call_tool(server, "kb_remember", {"text": "fact one"}) + call_tool(server, "kb_remember", {"text": "fact two"}) + + result = call_tool(server, "kb_get_memories", {}) + + assert result["count"] == 2 + texts = {m["text"] for m in result["memories"]} + assert texts == {"fact one", "fact two"} + + def test_excludes_superseded(self, server): + r1 = call_tool(server, "kb_remember", {"text": "old fact"}) + call_tool( + server, + "kb_remember", + { + "text": "new fact", + "supersedes": r1["entry_id"], + }, + ) + + result = call_tool(server, "kb_get_memories", {}) + + assert result["count"] == 1 + assert result["memories"][0]["text"] == "new fact" + + +# --------------------------------------------------------------------------- +# Tool schemas +# --------------------------------------------------------------------------- + + +class TestToolSchemas: + + def _get_tool(self, server, name): + req = types.ListToolsRequest(method="tools/list") + handler = server.request_handlers[types.ListToolsRequest] + result = asyncio.run(handler(req)) + for tool in result.root.tools: + if tool.name == name: + return tool + return None + + def test_kb_remember_exists(self, server): + tool = self._get_tool(server, "kb_remember") + assert tool is not None + assert "text" in tool.inputSchema["properties"] + assert tool.inputSchema["required"] == ["text"] + + def test_kb_remember_has_optional_params(self, server): + tool = self._get_tool(server, "kb_remember") + for param in ("category", "session_id", "supersedes"): + assert param in tool.inputSchema["properties"] + + def test_kb_get_memories_exists(self, server): + tool = self._get_tool(server, "kb_get_memories") + assert tool is not None diff --git a/tests/test_observability.py b/tests/test_observability.py index dd1ff46..2809ac6 100644 --- a/tests/test_observability.py +++ b/tests/test_observability.py @@ -1,14 +1,20 @@ """ Tests for dsagt.observability — Stage 0. -These tests use OTel's InMemorySpanExporter so they don't require a real -collector. They cover: +These tests read spans back from the serverless MLflow trace store (a +per-test ``sqlite:////mlflow.db``) rather than from OTel's +InMemorySpanExporter — the live-span path now uses ``mlflow.start_span`` +directly and installs no OTel TracerProvider. They cover: -* init_tracing is a no-op without an endpoint -* init_tracing with an endpoint installs a tracer provider +* init_tracing is a no-op outside a dsagt project dir +* init_tracing points MLflow at the resolved store + experiment * @traced opens a span, captures args, sets duration_ms, records exceptions * obs.set / obs.event are no-ops outside a span and write attributes inside -* child_span nests under the active span +* child_span / typed helpers nest under the active span +* internal traces are tagged ``dsagt.source`` + ``mlflow.trace.session`` + +Each test gets its own fresh store, so the LAST trace in the store is the +operation under test — no clear-then-read dance is needed. """ from __future__ import annotations @@ -22,52 +28,48 @@ @pytest.fixture(autouse=True) -def _reset_tracing(monkeypatch): - """Reset module state and install an in-memory exporter for each test. +def _reset_tracing(monkeypatch, tmp_path): + """Point MLflow at a fresh per-test sqlite store and mark tracing live. - We bypass ``init_tracing`` (no MLflow endpoint to talk to in tests) and - install our own TracerProvider with an InMemorySpanExporter directly. + Each test gets its own store, so reading the LAST active trace always + yields the operation under test — no exporter to clear between calls. """ - from opentelemetry import trace - from opentelemetry.sdk.trace import TracerProvider - from opentelemetry.sdk.trace.export import SimpleSpanProcessor - from opentelemetry.sdk.trace.export.in_memory_span_exporter import ( - InMemorySpanExporter, - ) + import mlflow - monkeypatch.setattr(obs_module, "_initialized", False) - monkeypatch.setattr(obs_module, "_tracer_provider", None) + mlflow.set_tracking_uri(f"sqlite:///{tmp_path}/mlflow.db") + mlflow.set_experiment("test") + + monkeypatch.setattr(obs_module, "_initialized", True) monkeypatch.setattr(obs_module, "_default_session_id", None) - exporter = InMemorySpanExporter() - provider = TracerProvider() - provider.add_span_processor(SimpleSpanProcessor(exporter)) - from opentelemetry.util._once import Once + yield - trace._TRACER_PROVIDER = None # type: ignore[attr-defined] - trace._TRACER_PROVIDER_SET_ONCE = Once() # type: ignore[attr-defined] - trace.set_tracer_provider(provider) - monkeypatch.setattr(obs_module, "_initialized", True) - yield exporter +def _last_trace(): + import mlflow + from mlflow import MlflowClient - exporter.clear() + return MlflowClient().get_trace(mlflow.get_last_active_trace_id()) -def test_init_tracing_no_endpoint_raises(monkeypatch): - """init_tracing must fail loudly when no backend is configured — silent - no-op behavior would let a misconfigured subprocess run with tracing - silently dropped, which is exactly the kind of silent-fallback bug - DSAGT's design principles prohibit.""" +def _spans_by_name(_ignored=None): + return {s.name: s for s in _last_trace().data.spans} + + +def test_init_tracing_outside_project_is_noop(monkeypatch): + """Serverless + never-raise: when cwd isn't a dsagt project dir (no + ``.dsagt/config.yaml`` with a ``project``), ``init_tracing`` logs and + no-ops rather than raising — one-shot tools / tests outside a project + simply run untraced. The store itself never needs a server, so the + only reason to skip is "not in a project", which must not be fatal.""" monkeypatch.setattr(obs_module, "_initialized", False) monkeypatch.delenv("MLFLOW_TRACKING_URI", raising=False) - with pytest.raises(RuntimeError, match="no observability backend"): - init_tracing("test-service") + # Repo root has no .dsagt/config.yaml → find_project_config returns None. + init_tracing("test-service") # must not raise + assert obs_module._initialized is False def test_traced_emits_span_with_args(_reset_tracing): - exporter = _reset_tracing - @traced("test.op", capture=["a", "b"]) def f(a, b, c=10): return a + b + c @@ -75,10 +77,9 @@ def f(a, b, c=10): result = f(1, 2, c=3) assert result == 6 - spans = exporter.get_finished_spans() - assert len(spans) == 1 - span = spans[0] - assert span.name == "test.op" + spans = _spans_by_name() + assert "test.op" in spans + span = spans["test.op"] assert span.attributes["a"] == 1 assert span.attributes["b"] == 2 assert "c" not in span.attributes # not in capture list @@ -86,20 +87,16 @@ def f(a, b, c=10): def test_traced_extract_return(_reset_tracing): - exporter = _reset_tracing - @traced("test.op", extract_return={"hits": len}) def search(): return ["a", "b", "c"] search() - span = exporter.get_finished_spans()[0] + span = _spans_by_name()["test.op"] assert span.attributes["hits"] == 3 def test_traced_records_exception(_reset_tracing): - exporter = _reset_tracing - @traced("test.boom") def boom(): raise ValueError("kaboom") @@ -107,9 +104,9 @@ def boom(): with pytest.raises(ValueError): boom() - span = exporter.get_finished_spans()[0] + span = _spans_by_name()["test.boom"] assert span.status.status_code.name == "ERROR" - # exception is recorded as a span event + # exception is recorded as a span event (MLflow auto-records it) assert any(e.name == "exception" for e in span.events) # duration is still set even on error path assert "duration_ms" in span.attributes @@ -123,8 +120,6 @@ def test_obs_set_outside_span_is_noop(): def test_obs_set_inside_span_writes_attribute(_reset_tracing): - exporter = _reset_tracing - @traced("test.op") def f(): obs.set("hits", 7) @@ -133,7 +128,7 @@ def f(): return None f() - span = exporter.get_finished_spans()[0] + span = _spans_by_name()["test.op"] assert span.attributes["hits"] == 7 assert span.attributes["foo"] == "bar" assert "skipped" not in span.attributes @@ -141,43 +136,68 @@ def f(): def test_child_span_nests(_reset_tracing): - exporter = _reset_tracing - @traced("test.parent") def parent(): with child_span("test.child", phase="embed"): pass parent() - spans = exporter.get_finished_spans() - # Children finish first, then the parent. - by_name = {s.name: s for s in spans} - child = by_name["test.child"] - parent_span = by_name["test.parent"] - assert child.parent is not None - assert child.parent.span_id == parent_span.context.span_id + spans = _spans_by_name() + child = spans["test.child"] + parent_span = spans["test.parent"] + assert child.parent_id is not None + assert child.parent_id == parent_span.span_id assert child.attributes["phase"] == "embed" -def test_session_id_stamped_via_metadata(_reset_tracing, monkeypatch): - """Session id flows through ``_metadata_stamper`` → MLflow's - ``InMemoryTraceManager``, not as a span attribute. We verify the - stamper is invoked with the reserved ``mlflow.trace.session`` key. +def test_root_span_source_tags_trace_and_session(_reset_tracing, monkeypatch): + """A categorization root (``open_span(source=...)``) tags the trace's + ``dsagt.source`` and, when a session id is set, the reserved + ``mlflow.trace.session`` metadata key (the native session filter). """ - del _reset_tracing - captured: list[dict] = [] monkeypatch.setattr(obs_module, "_default_session_id", "proj-xyz") - monkeypatch.setattr(obs_module, "_metadata_stamper", captured.append) - @traced("test.op") + with obs_module.open_span("search_knowledge", source="knowledge"): + pass + + trace = _last_trace() + assert trace.info.tags["dsagt.source"] == "knowledge" + assert trace.info.trace_metadata["mlflow.trace.session"] == "proj-xyz" + + +def test_inner_spans_inherit_root_source(_reset_tracing): + """Source is set at the entry point, not derived from the span name: a + child span opened under a ``skill`` root inherits ``skill`` even when the + child is a ``kb.*`` (knowledge-subsystem) span. This is what makes + ``search_skills`` → ``kb.search`` tag as ``skill``, not ``knowledge``. + """ + with obs_module.open_span("search_skills", source="skill"): + + @traced("kb.search") + def inner(): + pass + + inner() + + trace = _last_trace() + assert trace.info.tags["dsagt.source"] == "skill" + # Both the root and the kb.search child live in the one trace. + names = {s.name for s in trace.data.spans} + assert {"search_skills", "kb.search"} <= names + + +def test_uncategorized_span_has_no_source(_reset_tracing): + """A span opened with no source (inner span outside any root, e.g. a + background ``kb.*`` write) carries no ``dsagt.source`` — it doesn't leak + into the debug-view filter as a miscategorized concern. + """ + + @traced("kb.add_entries") def f(): pass f() - # _attach_trace_metadata should have stamped session id (and source, - # if a prefix matched — "test.op" doesn't, so just the session). - found = [md for md in captured if "mlflow.trace.session" in md] - assert any(md["mlflow.trace.session"] == "proj-xyz" for md in found) + assert "dsagt.source" not in _last_trace().info.tags def test_init_tracing_double_call_only_updates_session(monkeypatch): @@ -188,81 +208,51 @@ def test_init_tracing_double_call_only_updates_session(monkeypatch): assert obs_module._default_session_id == "new" -def test_init_tracing_installs_mlflow_provider(monkeypatch): - """init_tracing installs MLflow's native tracer provider as the OTel - global, so every span flows into MLflow's trace store with full - ``mlflow.spanInputs`` / ``mlflow.spanOutputs`` integration. Pins the - set_experiment call (so the project's experiment exists) and verifies - the strategy pointers (_metadata_stamper, _llm_context_factory) get - bound. +def test_init_tracing_points_mlflow_at_store_and_experiment(monkeypatch): + """init_tracing resolves the project name from ``.dsagt/config.yaml``, + points MLflow's tracking URI at the passed store, sets the experiment to + the project name, and flips ``_initialized`` true. """ captured: dict = {} - class _FakeExperiment: - experiment_id = "42" - def _fake_set_experiment(name): - captured["experiment_name"] = name - return _FakeExperiment() + captured["experiment"] = name def _fake_set_tracking_uri(uri): captured["tracking_uri"] = uri - def _fake_install(mlflow_url, project_name): - captured["installed_url"] = mlflow_url - captured["installed_project"] = project_name - import mlflow monkeypatch.setattr(mlflow, "set_experiment", _fake_set_experiment) monkeypatch.setattr(mlflow, "set_tracking_uri", _fake_set_tracking_uri) - # Stub the actual provider install — it touches OTel internals that - # aren't safely mockable in a unit test (Once flag, lazy provider). - monkeypatch.setattr(obs_module, "_install_mlflow_provider", _fake_install) monkeypatch.setattr(obs_module, "_initialized", False) - monkeypatch.setattr(obs_module, "_metadata_stamper", None) - monkeypatch.setattr(obs_module, "_llm_context_factory", None) monkeypatch.setattr( obs_module, "find_project_config", lambda: (None, {"project": "my-project"}), ) - monkeypatch.setattr( - obs_module, - "_read_session_id_from_runtime", - lambda _pdir: None, - ) - # Skip autolog wiring — using service_name="dsagt-run" exits the - # autolog branch in init_tracing. The strategy pointers and provider - # install are what we actually care about here. try: - init_tracing("dsagt-run", mlflow_url="http://localhost:5000/") + init_tracing("dsagt-run", mlflow_url="sqlite:///x.db") + assert captured["tracking_uri"] == "sqlite:///x.db" + assert captured["experiment"] == "my-project" + assert obs_module._initialized is True finally: monkeypatch.setattr(obs_module, "_initialized", False) - assert captured["installed_url"] == "http://localhost:5000/" - assert captured["installed_project"] == "my-project" - # Strategy pointers should be bound to the MLflow-backed implementations. - assert obs_module._metadata_stamper is obs_module._stamp_metadata_mlflow - assert obs_module._llm_context_factory is obs_module._mlflow_llm_context - # --------------------------------------------------------------------------- # Safety nets for the remaining defensive catches in observability.py. # -# After the fallback-purge pass, three "soft" catches remain in the -# observability layer. Two of them now live inline inside traced()'s -# wrapper (previously _attach_captured_args and _attach_return_attrs): +# After the fallback-purge pass, two "soft" catches remain in the +# observability layer, both inline inside traced()'s wrapper (previously +# _attach_captured_args and _attach_return_attrs): # -# 1. _shutdown wraps tracer_provider.shutdown() in try/except so that -# a shutdown failure during process exit can't escape into the -# atexit chain and corrupt the parent process exit. -# 2. traced() wraps sig.bind_partial in except TypeError so that a +# 1. traced() wraps sig.bind_partial in except TypeError so that a # function whose signature was mangled by another decorator doesn't # crash on every traced call. -# 3. traced() wraps each user-supplied extractor lambda in except +# 2. traced() wraps each user-supplied extractor lambda in except # Exception so a buggy lambda doesn't crash the instrumented # function. # @@ -272,49 +262,6 @@ def _fake_install(mlflow_url, project_name): # --------------------------------------------------------------------------- -def test_shutdown_runs_cleanly_against_real_provider(monkeypatch): - """The _shutdown helper must successfully flush and shut down a real - BatchSpanProcessor-backed provider on the happy path. - - If the catch in _shutdown ever fires during dsagt-setup-kb's atexit - pass, the user loses any spans still buffered in the BatchSpanProcessor. - This test catches "shutdown raised an exception we silently swallowed" - by exercising the real shutdown path against a real (in-memory) provider. - """ - from opentelemetry.sdk.resources import Resource - from opentelemetry.sdk.trace import TracerProvider - from opentelemetry.sdk.trace.export import SimpleSpanProcessor - from opentelemetry.sdk.trace.export.in_memory_span_exporter import ( - InMemorySpanExporter, - ) - - exporter = InMemorySpanExporter() - provider = TracerProvider( - resource=Resource.create({"service.name": "shutdown-test"}) - ) - provider.add_span_processor(SimpleSpanProcessor(exporter)) - - # Plug it into the module so _shutdown reaches the real object. - monkeypatch.setattr(obs_module, "_tracer_provider", provider) - - # Emit a span first so there's actually something to flush. - tracer = provider.get_tracer("test") - with tracer.start_as_current_span("pre-shutdown"): - pass - - # Run shutdown — must not raise, and must clear the module-level handle. - obs_module._shutdown() - assert obs_module._tracer_provider is None - - # The exporter should still hold the span we emitted before shutdown. - spans = exporter.get_finished_spans() - assert len(spans) == 1 - assert spans[0].name == "pre-shutdown" - - # Calling shutdown again is a no-op (idempotent atexit safety). - obs_module._shutdown() - - def test_extract_return_failure_logs_at_debug_and_does_not_crash( _reset_tracing, caplog ): @@ -343,9 +290,7 @@ def f(): assert result == ["a", "b", "c"] # Span was emitted with the good attributes set, bad one missing. - spans = _reset_tracing.get_finished_spans() - assert len(spans) == 1 - span = spans[0] + span = _spans_by_name()["test.extractor_bug"] assert span.attributes["good"] == 3 assert span.attributes["another_good"] == "a" assert "bad" not in span.attributes @@ -354,16 +299,16 @@ def f(): # what makes the silent skip visible to a developer running with # --verbose / DEBUG logging. debug_messages = [r.message for r in caplog.records if r.levelno == _logging.DEBUG] - assert any("extract_return['bad']" in m or "'bad'" in m for m in debug_messages), ( + assert any("extract_return['bad']" in m for m in debug_messages), ( f"Expected a DEBUG log mentioning the broken 'bad' extractor, " f"got: {debug_messages}" ) def test_attach_captured_args_happy_path_protects_bind_partial_catch(_reset_tracing): - """Pin the happy path for _attach_captured_args so the silent - 'except TypeError: return' catch can't hide a regression where args - stop being captured due to a signature-introspection bug. + """Pin the happy path for arg capture so the silent 'except TypeError: + skip' catch can't hide a regression where args stop being captured due to + a signature-introspection bug. If sig.bind_partial silently failed for any reason, this test would fail because the captured 'a' and 'b' attributes would be missing @@ -375,50 +320,18 @@ def f(a, b, c=42, *, d=None): return None f(1, b=2, d="ignored") - span = _reset_tracing.get_finished_spans()[0] + span = _spans_by_name()["test.signature_capture"] assert span.attributes["a"] == 1 assert span.attributes["b"] == 2 assert span.attributes["c"] == 42 # default value still captured assert "d" not in span.attributes -def test_litellm_imports_at_observability_init_no_fallback(): - """Regression test for the deletion of `except ImportError: return` in - configure_litellm_retries. - - litellm is a hard dependency in pyproject.toml. If anyone re-introduces - a try/except around the import (turning litellm "optional" again), the - function would silently no-op and the rate-limit retry/backoff knobs - would never be configured — exactly the silent-degradation pattern - we're trying to eliminate. - - This test asserts the import succeeds AND the side-effects of - configure_litellm_retries actually happened. - """ - import litellm - from dsagt.observability import configure_litellm_retries - - # Mutate litellm state to known-bad values, then call configure and - # verify it stomped them. If configure ever silently no-ops on a - # caught ImportError, the asserts below would fail. - litellm.num_retries = -999 - litellm.request_timeout = -999.0 - - configure_litellm_retries(num_retries=7, request_timeout=42.0) - - assert litellm.num_retries == 7 - assert litellm.request_timeout == 42.0 - - # --------------------------------------------------------------------------- # Stage 1: KnowledgeBase instrumentation # --------------------------------------------------------------------------- -def _spans_by_name(exporter): - return {s.name: s for s in exporter.get_finished_spans()} - - @contextmanager def _kb_with_mocked_embedder(tmp_path, backend: str = "api", model: str = "test-model"): """Build a KnowledgeBase with a mocked embedder for the given backend. @@ -438,11 +351,11 @@ def fake_embed(texts): mock_client = MagicMock() mock_client.embed = fake_embed - with patch("dsagt.knowledge._make_embedder", return_value=mock_client): + with patch("dsagt.knowledge.Embedder.create", return_value=mock_client): kb = KnowledgeBase( index_dir=tmp_path / f"kb_{backend}", default_embedder=backend, - embedder_kwargs={"model": model}, + model=model, ) try: yield kb @@ -452,16 +365,15 @@ def fake_embed(texts): def test_kb_search_emits_three_child_spans(_reset_tracing, tmp_path): """kb.search should produce kb.search → {kb.embed, kb.index_search}.""" - exporter = _reset_tracing with _kb_with_mocked_embedder(tmp_path) as kb: # Seed a collection so search has something to load. kb.add_entries(texts=["hello world", "goodbye"], collection="tcoll") - exporter.clear() # ignore add_entries spans for this assertion results = kb.search("hello", collection="tcoll", top_k=2, rerank=False) assert isinstance(results, list) - spans = _spans_by_name(exporter) + # The last trace is the search (add_entries is an earlier trace). + spans = _spans_by_name() assert "kb.search" in spans assert "kb.embed" in spans assert "kb.index_search" in spans @@ -472,8 +384,8 @@ def test_kb_search_emits_three_child_spans(_reset_tracing, tmp_path): embed = spans["kb.embed"] index = spans["kb.index_search"] - assert embed.parent.span_id == parent.context.span_id - assert index.parent.span_id == parent.context.span_id + assert embed.parent_id == parent.span_id + assert index.parent_id == parent.span_id # Captured args + obs.set('hits', ...) on the parent. assert parent.attributes["collection"] == "tcoll" @@ -492,13 +404,11 @@ def test_kb_search_emits_three_child_spans(_reset_tracing, tmp_path): def test_kb_search_local_backend_same_span_shape(_reset_tracing, tmp_path): """The local embedding backend should emit the same span tree.""" - exporter = _reset_tracing with _kb_with_mocked_embedder(tmp_path, backend="local", model="bge-base") as kb: kb.add_entries(texts=["hello world"], collection="tcoll") - exporter.clear() kb.search("hello", collection="tcoll", top_k=1, rerank=False) - spans = _spans_by_name(exporter) + spans = _spans_by_name() assert "kb.search" in spans assert "kb.embed" in spans assert spans["kb.embed"].attributes["backend"] == "local" @@ -507,7 +417,6 @@ def test_kb_search_local_backend_same_span_shape(_reset_tracing, tmp_path): def test_kb_ingest_emits_embed_child(_reset_tracing, tmp_path): """kb.ingest opens an outer span and one child kb.embed span.""" - exporter = _reset_tracing src = tmp_path / "docs" src.mkdir() (src / "a.txt").write_text("alpha beta gamma") @@ -516,13 +425,13 @@ def test_kb_ingest_emits_embed_child(_reset_tracing, tmp_path): with _kb_with_mocked_embedder(tmp_path) as kb: kb.ingest(src) - spans = _spans_by_name(exporter) + spans = _spans_by_name() assert "kb.ingest" in spans assert "kb.embed" in spans ingest = spans["kb.ingest"] embed = spans["kb.embed"] - assert embed.parent.span_id == ingest.context.span_id + assert embed.parent_id == ingest.span_id assert ingest.attributes["n_files"] == 2 assert ingest.attributes["n_chunks"] >= 2 assert embed.attributes["n_texts"] == ingest.attributes["n_chunks"] @@ -530,73 +439,15 @@ def test_kb_ingest_emits_embed_child(_reset_tracing, tmp_path): def test_kb_add_entries_emits_span(_reset_tracing, tmp_path): """kb.add_entries should emit a top-level span with n_entries.""" - exporter = _reset_tracing with _kb_with_mocked_embedder(tmp_path) as kb: kb.add_entries(texts=["one", "two", "three"], collection="epis") - spans = _spans_by_name(exporter) + spans = _spans_by_name() assert "kb.add_entries" in spans assert spans["kb.add_entries"].attributes["collection"] == "epis" assert spans["kb.add_entries"].attributes["n_entries"] == 3 -# --------------------------------------------------------------------------- -# Stage 2: LiteLLM retry wiring -# --------------------------------------------------------------------------- - - -def test_configure_litellm_retries_sets_module_globals(monkeypatch): - """configure_litellm_retries must set litellm.num_retries / request_timeout.""" - import litellm - - from dsagt.observability import configure_litellm_retries - - # Save and clear so the test owns the values. - monkeypatch.setattr(litellm, "num_retries", None, raising=False) - monkeypatch.setattr(litellm, "request_timeout", 0.0, raising=False) - - configure_litellm_retries(num_retries=7, request_timeout=42.0) - - assert litellm.num_retries == 7 - assert litellm.request_timeout == 42.0 - - -def test_configure_litellm_retries_works_without_tracing(monkeypatch): - """The retry knobs must be applied even when tracing is not initialized. - - This is the dsagt-setup-kb path: the long-running embed job that has to - survive rate limits also runs before any MLflow endpoint exists. - """ - import litellm - - from dsagt.observability import configure_litellm_retries - - monkeypatch.setattr(obs_module, "_initialized", False) - monkeypatch.setattr(obs_module, "_tracer_provider", None) - monkeypatch.setattr(litellm, "num_retries", None, raising=False) - - configure_litellm_retries(num_retries=3, request_timeout=120.0) - - assert litellm.num_retries == 3 - assert litellm.request_timeout == 120.0 - - -def test_kb_search_does_not_call_real_litellm(_reset_tracing, tmp_path): - """Sanity check: with a mocked embedder, no litellm.embedding call escapes.""" - from unittest.mock import patch as _patch - - exporter = _reset_tracing - with _patch("litellm.embedding") as mock_embed: - with _kb_with_mocked_embedder(tmp_path) as kb: - kb.add_entries(texts=["hello"], collection="tcoll") - kb.search("hello", collection="tcoll", top_k=1) - - # _kb_with_mocked_embedder patches _make_embedder, so litellm.embedding - # should never be called even though APIEmbeddingClient now uses it. - assert mock_embed.call_count == 0 - assert "kb.search" in _spans_by_name(exporter) - - # --------------------------------------------------------------------------- # Stage 3: tool execution spans # --------------------------------------------------------------------------- @@ -624,34 +475,30 @@ def test_truncate_handles_none(): assert truncate(None, 256) == "" -def test_tool_execute_span_attributes(_reset_tracing): - """tool_execute_span sets record_id and tool_name on the span.""" - exporter = _reset_tracing +def test_code_execute_span_attributes(_reset_tracing): + """code_execute_span sets record_id and code_name on the span.""" + from dsagt.observability import obs, code_execute_span - from dsagt.observability import obs, tool_execute_span - - with tool_execute_span(record_id="abc123", tool_name="fastp"): + with code_execute_span(record_id="abc123", code_name="fastp"): obs.set("exit_code", 0) obs.set("duration_ms", 42.5) - spans = _spans_by_name(exporter) - assert "tool.execute" in spans - span = spans["tool.execute"] + spans = _spans_by_name() + assert "code.execute" in spans + span = spans["code.execute"] assert span.attributes["record_id"] == "abc123" - assert span.attributes["tool_name"] == "fastp" + assert span.attributes["code_name"] == "fastp" assert span.attributes["exit_code"] == 0 assert span.attributes["duration_ms"] == 42.5 -def test_run_and_record_emits_tool_execute_span(_reset_tracing, tmp_path): +def test_run_and_record_emits_code_execute_span(_reset_tracing, tmp_path): """run_and_record() should produce a tool.execute span with the expected execution attributes.""" - exporter = _reset_tracing - from dsagt.provenance import run_and_record rc = run_and_record( - tool_name="echo", + code_name="echo", command=["echo", "hello world"], records_dir=tmp_path / "records", session_id="test-session", @@ -661,12 +508,12 @@ def test_run_and_record_emits_tool_execute_span(_reset_tracing, tmp_path): ) assert rc == 0 - spans = _spans_by_name(exporter) - assert "tool.execute" in spans - span = spans["tool.execute"] + spans = _spans_by_name() + assert "code.execute" in spans + span = spans["code.execute"] assert span.attributes["record_id"] == "rec-001" - assert span.attributes["tool_name"] == "echo" + assert span.attributes["code_name"] == "echo" assert span.attributes["exit_code"] == 0 assert span.attributes["duration_ms"] >= 0 assert span.attributes["n_input_files"] == 1 @@ -676,48 +523,42 @@ def test_run_and_record_emits_tool_execute_span(_reset_tracing, tmp_path): def test_run_and_record_failed_tool_records_event_and_status(_reset_tracing, tmp_path): - """A failed tool call should emit a tool_failed event with the exit code + """A failed tool call should emit a code_failed event with the exit code and the truncated stderr should be attached.""" - exporter = _reset_tracing - from dsagt.provenance import run_and_record rc = run_and_record( - tool_name="missing_tool", + code_name="missing_tool", command=["this-binary-does-not-exist-anywhere"], records_dir=tmp_path / "records", session_id="test-session", ) assert rc == 127 - spans = _spans_by_name(exporter) - span = spans["tool.execute"] + span = _spans_by_name()["code.execute"] assert span.attributes["exit_code"] == 127 # Stderr was set by the FileNotFoundError branch — should be on the span. assert "stderr_truncated" in span.attributes - # tool_failed event was added. - assert any(e.name == "tool_failed" for e in span.events) + # code_failed event was added. + assert any(e.name == "code_failed" for e in span.events) def test_run_and_record_long_stderr_is_truncated(_reset_tracing, tmp_path): """Span attributes should never carry multi-megabyte stderr blobs.""" - exporter = _reset_tracing - from dsagt.provenance import run_and_record # Use python -c to emit a large stderr deterministically. big = "x" * 5000 rc = run_and_record( - tool_name="echo_err", + code_name="echo_err", command=["python", "-c", f"import sys; sys.stderr.write('{big}'); sys.exit(0)"], records_dir=tmp_path / "records", session_id="s", ) assert rc == 0 - spans = _spans_by_name(exporter) - span = spans["tool.execute"] + span = _spans_by_name()["code.execute"] truncated = span.attributes["stderr_truncated"] # Truncated to ~256 chars even though stderr was 5000. assert len(truncated) < 300 @@ -735,15 +576,10 @@ def _make_registry_server(tmp_path): Mirrors the pattern used by tests/test_registry_server.py so we exercise the real call_tool dispatcher rather than reaching into private state. """ - from dsagt.commands.registry_server import create_registry_server - from dsagt.registry import ToolRegistry - - source_dir = tmp_path / "source_skills" - source_dir.mkdir() - reg = ToolRegistry( - source_tools_dir=str(source_dir), - runtime_dir=str(tmp_path / "runtime"), - ) + from dsagt.mcp.registry_tools import create_registry_server + from dsagt.registry import CodeRegistry + + reg = CodeRegistry(runtime_dir=str(tmp_path / "runtime")) return create_registry_server(reg) @@ -758,37 +594,38 @@ def _minimal_spec(name: str, **extras) -> dict: return spec -def test_save_tool_spec_emits_registry_save_span(_reset_tracing, tmp_path): - """save_tool_spec should produce a registry.save_tool_spec span with - tool_name, language, n_dependencies, n_tags, action, and registry_size.""" +def test_save_code_spec_emits_registry_save_span(_reset_tracing, tmp_path): + """save_code_spec should produce a registry.save_code_spec span with + code_name, language, n_dependencies, n_tags, action, and registry_size.""" from mcp_helpers import call_tool_sync as call_tool - exporter = _reset_tracing server = _make_registry_server(tmp_path) spec = _minimal_spec("alpha", language="python", tags=["genomics", "qc"]) - call_tool(server, "save_tool_spec", {"spec": spec}) + call_tool(server, "save_code_spec", {"spec": spec}) - spans = _spans_by_name(exporter) - assert "registry.save_tool_spec" in spans - span = spans["registry.save_tool_spec"] - assert span.attributes["tool_name"] == "alpha" + spans = _spans_by_name() + assert "registry.save_code_spec" in spans + span = spans["registry.save_code_spec"] + assert span.attributes["code_name"] == "alpha" assert span.attributes["language"] == "python" assert span.attributes["n_dependencies"] == 0 assert span.attributes["n_tags"] == 2 assert span.attributes["action"] == "added" assert span.attributes["registry_size"] == 1 + # The dispatch root tags the whole trace with the tool's concern category. + assert _last_trace().info.tags["dsagt.source"] == "registry" -def test_save_tool_spec_with_deps_nests_install_span( +def test_save_code_spec_with_deps_nests_install_span( _reset_tracing, tmp_path, monkeypatch ): - """When a spec carries dependencies, save_tool_spec should open a + """When a spec carries dependencies, save_code_spec should open a nested registry.install_dependencies span as a child.""" from mcp_helpers import call_tool_sync as call_tool # Stub the actual uv install so the test doesn't hit the network. - import dsagt.commands.registry_server as rs_mod + import dsagt.mcp.registry_tools as rs_mod monkeypatch.setattr( rs_mod, @@ -796,17 +633,16 @@ def test_save_tool_spec_with_deps_nests_install_span( lambda packages, timeout=120: f"Successfully installed: {', '.join(packages)}", ) - exporter = _reset_tracing server = _make_registry_server(tmp_path) spec = _minimal_spec("beta", dependencies=["numpy", "pandas"]) - call_tool(server, "save_tool_spec", {"spec": spec}) + call_tool(server, "save_code_spec", {"spec": spec}) - spans = _spans_by_name(exporter) - save_span = spans["registry.save_tool_spec"] + spans = _spans_by_name() + save_span = spans["registry.save_code_spec"] install_span = spans["registry.install_dependencies"] - assert install_span.parent.span_id == save_span.context.span_id + assert install_span.parent_id == save_span.span_id assert install_span.attributes["package_count"] == 2 assert install_span.attributes["status"] == "ok" assert "numpy" in install_span.attributes["packages_preview"] @@ -819,7 +655,7 @@ def test_install_dependencies_failed_records_event( an install_failed event with the error message truncated.""" from mcp_helpers import call_tool_sync as call_tool - import dsagt.commands.registry_server as rs_mod + import dsagt.mcp.registry_tools as rs_mod monkeypatch.setattr( rs_mod, @@ -827,17 +663,16 @@ def test_install_dependencies_failed_records_event( lambda packages, timeout=120: "Installation failed (exit code 1):\nresolution failure", ) - exporter = _reset_tracing server = _make_registry_server(tmp_path) # First register a tool with deps so install_dependencies has something # to operate on, then call install_dependencies directly. spec = _minimal_spec("gamma", dependencies=["broken-package"]) - call_tool(server, "save_tool_spec", {"spec": spec}) - exporter.clear() + call_tool(server, "save_code_spec", {"spec": spec}) call_tool(server, "install_dependencies", {}) - spans = _spans_by_name(exporter) + # The last trace is the install_dependencies call. + spans = _spans_by_name() assert "registry.install_dependencies" in spans span = spans["registry.install_dependencies"] assert span.attributes["package_count"] == 1 @@ -849,7 +684,6 @@ def test_reconstruct_pipeline_emits_span(_reset_tracing, tmp_path): """reconstruct_pipeline should produce a span with format and output_chars.""" from mcp_helpers import call_tool_sync as call_tool - exporter = _reset_tracing server = _make_registry_server(tmp_path) # Empty trace_archive — reconstruct_pipeline should still emit a span, @@ -858,93 +692,27 @@ def test_reconstruct_pipeline_emits_span(_reset_tracing, tmp_path): call_tool(server, "reconstruct_pipeline", {"format": "bash"}) - spans = _spans_by_name(exporter) + spans = _spans_by_name() assert "registry.reconstruct_pipeline" in spans span = spans["registry.reconstruct_pipeline"] assert span.attributes["format"] == "bash" assert "output_chars" in span.attributes -def test_search_registry_does_not_emit_span(_reset_tracing, tmp_path): - """High-frequency search calls are deliberately not instrumented.""" +def test_search_registry_categorized_but_no_internal_span(_reset_tracing, tmp_path): + """Every MCP call gets a categorized dispatch root span (so the concern + shows up in the debug view), but high-frequency search still adds no + internal subsystem span — the trace is just the root, tagged ``registry``, + with no ``registry.*`` child.""" from mcp_helpers import call_tool_sync as call_tool - exporter = _reset_tracing server = _make_registry_server(tmp_path) call_tool(server, "search_registry", {"query": "anything"}) - spans = _spans_by_name(exporter) - assert "registry.search" not in spans - assert "registry.search_registry" not in spans - - -# --------------------------------------------------------------------------- -# extract_cache_stats — provider-agnostic cache-token reader -# --------------------------------------------------------------------------- -# -# Each provider returns cache stats under a different field name; LiteLLM -# doesn't backfill across them. Lock in the field-name handling so a -# regression can't silently hide cache hits in `dsagt info`. - - -def test_extract_cache_stats_anthropic_format(): - from dsagt.observability import extract_cache_stats - - usage = { - "prompt_tokens": 5000, - "completion_tokens": 100, - "cache_read_input_tokens": 3000, - "cache_creation_input_tokens": 2000, - } - assert extract_cache_stats(usage) == (3000, 2000) - - -def test_extract_cache_stats_openai_format(): - from dsagt.observability import extract_cache_stats - - # OpenAI/Azure: cached_tokens nested under prompt_tokens_details, no write field - usage = { - "prompt_tokens": 5000, - "completion_tokens": 100, - "prompt_tokens_details": {"cached_tokens": 2400, "audio_tokens": None}, - } - assert extract_cache_stats(usage) == (2400, 0) - - -def test_extract_cache_stats_gemini_format(): - from dsagt.observability import extract_cache_stats - - usage = {"prompt_tokens": 1000, "cached_content_token_count": 700} - assert extract_cache_stats(usage) == (700, 0) - - -def test_extract_cache_stats_deepseek_format(): - from dsagt.observability import extract_cache_stats - - # DeepSeek: prompt_cache_hit_tokens is the read; prompt_cache_miss_tokens - # is the COMPLEMENT (uncached prompt tokens), not a "write" — don't count it. - usage = { - "prompt_tokens": 1000, - "prompt_cache_hit_tokens": 600, - "prompt_cache_miss_tokens": 400, - } - assert extract_cache_stats(usage) == (600, 0) - - -def test_extract_cache_stats_no_cache_fields(): - from dsagt.observability import extract_cache_stats - - assert extract_cache_stats({"prompt_tokens": 100, "completion_tokens": 10}) == ( - 0, - 0, - ) - - -def test_extract_cache_stats_handles_garbage(): - from dsagt.observability import extract_cache_stats - - # Real responses occasionally have None or non-dict values where we expect dicts - assert extract_cache_stats({"prompt_tokens_details": None}) == (0, 0) - assert extract_cache_stats({"prompt_tokens_details": "garbage"}) == (0, 0) - assert extract_cache_stats(None) == (0, 0) - assert extract_cache_stats({}) == (0, 0) + trace = _last_trace() + names = {s.name for s in trace.data.spans} + # The dispatch root exists and is categorized... + assert "search_registry" in names + assert trace.info.tags["dsagt.source"] == "registry" + # ...but search opens no internal subsystem span. + assert not any(n.startswith("registry.") for n in names) diff --git a/tests/test_outlier.py b/tests/test_outlier.py deleted file mode 100644 index 9bebed8..0000000 --- a/tests/test_outlier.py +++ /dev/null @@ -1,228 +0,0 @@ -""" -Tests for outlier detection and suggestion queue. - -Tests CategoryCentroids (incremental update, cosine distance, persistence), -SuggestionQueue (add, dismiss, persistence), and check_and_queue_outliers -(threshold behavior, first-in-category skip). -""" - -import numpy as np -import pytest - -from dsagt.memory import CategoryCentroids, SuggestionQueue, check_and_queue_outliers - - -# --------------------------------------------------------------------------- -# CategoryCentroids -# --------------------------------------------------------------------------- - -class TestCategoryCentroids: - - def test_first_entry_returns_zero_distance(self, tmp_path): - """First entry in a category has no centroid to compare against.""" - centroids = CategoryCentroids(tmp_path / "centroids.json") - vec = np.array([1.0, 0.0, 0.0]) - distance = centroids.update("quality_control", vec) - assert distance == 0.0 - assert centroids.count("quality_control") == 1 - - def test_identical_vectors_zero_distance(self, tmp_path): - """Identical vectors have zero distance.""" - centroids = CategoryCentroids(tmp_path / "centroids.json") - vec = np.array([1.0, 0.0, 0.0]) - centroids.update("qc", vec) - distance = centroids.update("qc", vec) - assert distance < 0.01 - - def test_orthogonal_vectors_high_distance(self, tmp_path): - """Orthogonal vectors have distance ~1.0.""" - centroids = CategoryCentroids(tmp_path / "centroids.json") - centroids.update("qc", np.array([1.0, 0.0, 0.0])) - distance = centroids.update("qc", np.array([0.0, 1.0, 0.0])) - assert distance > 0.9 - - def test_incremental_update(self, tmp_path): - """Centroid shifts toward new vectors incrementally.""" - centroids = CategoryCentroids(tmp_path / "centroids.json") - centroids.update("qc", np.array([1.0, 0.0, 0.0])) - centroids.update("qc", np.array([1.0, 0.0, 0.0])) - centroids.update("qc", np.array([1.0, 0.0, 0.0])) - - # Outlier vector — mostly in the y direction - distance = centroids.update("qc", np.array([0.0, 1.0, 0.0])) - assert distance > 0.8 - assert centroids.count("qc") == 4 - - def test_save_and_reload(self, tmp_path): - """Centroids persist across reloads.""" - path = tmp_path / "centroids.json" - c1 = CategoryCentroids(path) - c1.update("qc", np.array([1.0, 0.0])) - c1.update("qc", np.array([0.9, 0.1])) - c1.save() - - c2 = CategoryCentroids(path) - assert c2.count("qc") == 2 - assert "qc" in c2.categories - - def test_separate_categories(self, tmp_path): - """Different categories maintain independent centroids.""" - centroids = CategoryCentroids(tmp_path / "centroids.json") - centroids.update("qc", np.array([1.0, 0.0])) - centroids.update("config", np.array([0.0, 1.0])) - - assert centroids.count("qc") == 1 - assert centroids.count("config") == 1 - - -# --------------------------------------------------------------------------- -# SuggestionQueue -# --------------------------------------------------------------------------- - -class TestSuggestionQueue: - - def test_add_returns_id(self, tmp_path): - queue = SuggestionQueue(tmp_path / "suggestions.json") - sid = queue.add("test fact", "qc", 0.45, "s1") - assert sid.startswith("sug_") - assert queue.count == 1 - - def test_get_all(self, tmp_path): - queue = SuggestionQueue(tmp_path / "suggestions.json") - queue.add("fact 1", "qc", 0.4) - queue.add("fact 2", "config", 0.5) - suggestions = queue.get_all() - assert len(suggestions) == 2 - assert suggestions[0]["text"] == "fact 1" - assert suggestions[1]["category"] == "config" - - def test_dismiss_removes(self, tmp_path): - queue = SuggestionQueue(tmp_path / "suggestions.json") - sid = queue.add("test fact", "qc", 0.45) - assert queue.dismiss(sid) is True - assert queue.count == 0 - - def test_dismiss_nonexistent(self, tmp_path): - queue = SuggestionQueue(tmp_path / "suggestions.json") - assert queue.dismiss("sug_nonexistent") is False - - def test_clear(self, tmp_path): - queue = SuggestionQueue(tmp_path / "suggestions.json") - queue.add("a", "qc", 0.3) - queue.add("b", "qc", 0.4) - removed = queue.clear() - assert removed == 2 - assert queue.count == 0 - - def test_persistence(self, tmp_path): - """Suggestions survive reload.""" - path = tmp_path / "suggestions.json" - q1 = SuggestionQueue(path) - q1.add("persisted fact", "qc", 0.5) - - q2 = SuggestionQueue(path) - assert q2.count == 1 - assert q2.get_all()[0]["text"] == "persisted fact" - - -# --------------------------------------------------------------------------- -# check_and_queue_outliers -# --------------------------------------------------------------------------- - -class TestCheckAndQueueOutliers: - - def _make_embeddings(self, vectors): - return np.array(vectors, dtype=np.float32) - - def test_flags_outlier(self, tmp_path): - """A vector far from its category centroid gets flagged.""" - centroids = CategoryCentroids(tmp_path / "centroids.json") - queue = SuggestionQueue(tmp_path / "suggestions.json") - - # Seed the centroid with a few similar vectors - for _ in range(5): - centroids.update("qc", np.array([1.0, 0.0, 0.0])) - centroids.save() - - texts = ["outlier observation"] - categories = ["qc"] - embeddings = self._make_embeddings([[0.0, 1.0, 0.0]]) # orthogonal - - ids = check_and_queue_outliers( - texts, categories, embeddings, centroids, queue, - threshold=0.3, session_id="s1", - ) - - assert len(ids) == 1 - assert queue.count == 1 - - def test_normal_fact_not_flagged(self, tmp_path): - """A vector close to its category centroid is not flagged.""" - centroids = CategoryCentroids(tmp_path / "centroids.json") - queue = SuggestionQueue(tmp_path / "suggestions.json") - - for _ in range(5): - centroids.update("qc", np.array([1.0, 0.0, 0.0])) - centroids.save() - - texts = ["normal observation"] - categories = ["qc"] - embeddings = self._make_embeddings([[0.98, 0.1, 0.0]]) # close - - ids = check_and_queue_outliers( - texts, categories, embeddings, centroids, queue, - threshold=0.3, session_id="s1", - ) - - assert len(ids) == 0 - assert queue.count == 0 - - def test_first_in_category_not_flagged(self, tmp_path): - """First entry in a category cannot be an outlier (no centroid yet).""" - centroids = CategoryCentroids(tmp_path / "centroids.json") - queue = SuggestionQueue(tmp_path / "suggestions.json") - - texts = ["first observation"] - categories = ["new_category"] - embeddings = self._make_embeddings([[0.5, 0.5, 0.5]]) - - ids = check_and_queue_outliers( - texts, categories, embeddings, centroids, queue, - threshold=0.1, session_id="s1", - ) - - assert len(ids) == 0 - - def test_empty_category_skipped(self, tmp_path): - """Facts with empty category are skipped.""" - centroids = CategoryCentroids(tmp_path / "centroids.json") - queue = SuggestionQueue(tmp_path / "suggestions.json") - - texts = ["uncategorized"] - categories = [""] - embeddings = self._make_embeddings([[1.0, 0.0, 0.0]]) - - ids = check_and_queue_outliers( - texts, categories, embeddings, centroids, queue, - threshold=0.1, - ) - - assert len(ids) == 0 - - def test_centroids_saved(self, tmp_path): - """Centroids are saved to disk after checking.""" - centroids = CategoryCentroids(tmp_path / "centroids.json") - queue = SuggestionQueue(tmp_path / "suggestions.json") - - texts = ["fact"] - categories = ["qc"] - embeddings = self._make_embeddings([[1.0, 0.0]]) - - check_and_queue_outliers( - texts, categories, embeddings, centroids, queue, - threshold=0.3, - ) - - assert (tmp_path / "centroids.json").exists() - - diff --git a/tests/test_pipeline.py b/tests/test_pipeline.py index 9af66c1..129f9f3 100644 --- a/tests/test_pipeline.py +++ b/tests/test_pipeline.py @@ -5,7 +5,6 @@ import json from pathlib import Path -import pytest from dsagt.provenance import ( build_dependency_graph, @@ -15,22 +14,22 @@ render_snakemake, ) - # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- + def _write_record(trace_dir: Path, record: dict) -> Path: trace_dir.mkdir(parents=True, exist_ok=True) rid = record.get("record_id", "r0") - tool = record.get("tool_name", "tool") + tool = record.get("code_name", "tool") path = trace_dir / f"{tool}_{rid}.json" path.write_text(json.dumps(record)) return path def _make_record( - tool_name: str, + code_name: str, command: list[str], input_files: list[str] | None = None, output_files: list[str] | None = None, @@ -41,7 +40,7 @@ def _make_record( ) -> dict: return { "record_id": record_id, - "tool_name": tool_name, + "code_name": code_name, "session_id": session_id, "execution": { "exact_command": command, @@ -60,6 +59,7 @@ def _make_record( # load_pipeline_records # --------------------------------------------------------------------------- + class TestLoadPipelineRecords: def test_loads_wrapper_records(self, tmp_path): @@ -68,13 +68,13 @@ def test_loads_wrapper_records(self, tmp_path): records = load_pipeline_records(tmp_path) assert len(records) == 1 - assert records[0]["tool_name"] == "fastp" + assert records[0]["code_name"] == "fastp" def test_skips_proxy_only_records(self, tmp_path): """Records without execution layer are skipped.""" proxy_record = { "record_id": "r1", - "tool_name": "fastp", + "code_name": "fastp", "session_id": "s1", "intent": {"command": "fastp", "parameters": {}}, "execution": None, @@ -86,20 +86,34 @@ def test_skips_proxy_only_records(self, tmp_path): assert len(records) == 0 def test_filters_by_session(self, tmp_path): - _write_record(tmp_path, _make_record("a", ["a"], session_id="s1", record_id="r1")) - _write_record(tmp_path, _make_record("b", ["b"], session_id="s2", record_id="r2")) + _write_record( + tmp_path, _make_record("a", ["a"], session_id="s1", record_id="r1") + ) + _write_record( + tmp_path, _make_record("b", ["b"], session_id="s2", record_id="r2") + ) records = load_pipeline_records(tmp_path, session_id="s1") assert len(records) == 1 - assert records[0]["tool_name"] == "a" + assert records[0]["code_name"] == "a" def test_sorted_by_timestamp(self, tmp_path): - _write_record(tmp_path, _make_record("late", ["late"], timestamp="2024-01-15T11:00:00Z", record_id="r2")) - _write_record(tmp_path, _make_record("early", ["early"], timestamp="2024-01-15T09:00:00Z", record_id="r1")) + _write_record( + tmp_path, + _make_record( + "late", ["late"], timestamp="2024-01-15T11:00:00Z", record_id="r2" + ), + ) + _write_record( + tmp_path, + _make_record( + "early", ["early"], timestamp="2024-01-15T09:00:00Z", record_id="r1" + ), + ) records = load_pipeline_records(tmp_path) - assert records[0]["tool_name"] == "early" - assert records[1]["tool_name"] == "late" + assert records[0]["code_name"] == "early" + assert records[1]["code_name"] == "late" def test_empty_directory(self, tmp_path): tmp_path.mkdir(exist_ok=True) @@ -113,6 +127,7 @@ def test_nonexistent_directory(self, tmp_path): # build_dependency_graph # --------------------------------------------------------------------------- + class TestBuildDependencyGraph: def test_linear_dependency(self): @@ -161,6 +176,7 @@ def test_self_dependency_excluded(self): # render_bash # --------------------------------------------------------------------------- + class TestRenderBash: def test_basic_script(self): @@ -175,7 +191,9 @@ def test_basic_script(self): def test_includes_file_comments(self): records = [ - _make_record("fastp", ["fastp"], input_files=["raw.fq"], output_files=["clean.fq"]), + _make_record( + "fastp", ["fastp"], input_files=["raw.fq"], output_files=["clean.fq"] + ), ] deps = build_dependency_graph(records) script = render_bash(records, deps) @@ -212,12 +230,17 @@ def test_quotes_special_characters(self): # render_snakemake # --------------------------------------------------------------------------- + class TestRenderSnakemake: def test_basic_workflow(self): records = [ - _make_record("fastp", ["fastp", "-q", "20"], - input_files=["raw.fq"], output_files=["clean.fq"]), + _make_record( + "fastp", + ["fastp", "-q", "20"], + input_files=["raw.fq"], + output_files=["clean.fq"], + ), ] deps = build_dependency_graph(records) workflow = render_snakemake(records, deps) @@ -230,7 +253,12 @@ def test_basic_workflow(self): def test_multi_step(self): records = [ _make_record("fastp", ["fastp"], output_files=["clean.fq"]), - _make_record("megahit", ["megahit"], input_files=["clean.fq"], output_files=["contigs.fa"]), + _make_record( + "megahit", + ["megahit"], + input_files=["clean.fq"], + output_files=["contigs.fa"], + ), ] deps = build_dependency_graph(records) workflow = render_snakemake(records, deps) @@ -245,21 +273,29 @@ def test_multi_step(self): # reconstruct_pipeline (end-to-end) # --------------------------------------------------------------------------- + class TestReconstructPipeline: def test_bash_format(self, tmp_path): - _write_record(tmp_path, _make_record("fastp", ["fastp", "-q", "20"], record_id="r1")) + _write_record( + tmp_path, _make_record("fastp", ["fastp", "-q", "20"], record_id="r1") + ) script = reconstruct_pipeline(tmp_path, fmt="bash") assert "#!/usr/bin/env bash" in script assert "fastp -q 20" in script def test_snakemake_format(self, tmp_path): - _write_record(tmp_path, _make_record( - "fastp", ["fastp"], - input_files=["raw.fq"], output_files=["clean.fq"], - record_id="r1", - )) + _write_record( + tmp_path, + _make_record( + "fastp", + ["fastp"], + input_files=["raw.fq"], + output_files=["clean.fq"], + record_id="r1", + ), + ) workflow = reconstruct_pipeline(tmp_path, fmt="snakemake") assert "rule fastp_1:" in workflow @@ -270,31 +306,52 @@ def test_empty_returns_comment(self, tmp_path): assert "No execution records found" in result def test_session_filter(self, tmp_path): - _write_record(tmp_path, _make_record("a", ["a"], session_id="s1", record_id="r1")) - _write_record(tmp_path, _make_record("b", ["b"], session_id="s2", record_id="r2")) + _write_record( + tmp_path, _make_record("a", ["a"], session_id="s1", record_id="r1") + ) + _write_record( + tmp_path, _make_record("b", ["b"], session_id="s2", record_id="r2") + ) script = reconstruct_pipeline(tmp_path, session_id="s1", fmt="bash") + # Only the s1 record survives the filter: exactly one step (tool "a"), + # the s2 record ("b") is excluded — so there is no second step. assert "Step 1: a" in script - assert "b" not in script.split("Step")[1] if "Step" in script else True + assert "Step 2:" not in script def test_full_pipeline_with_deps(self, tmp_path): """Three-step pipeline: fastp → megahit → quast.""" - _write_record(tmp_path, _make_record( - "fastp", ["fastp", "-q", "20", "--in1", "raw.fq.gz"], - output_files=["clean.fq.gz"], - timestamp="2024-01-15T10:00:00Z", record_id="r1", - )) - _write_record(tmp_path, _make_record( - "megahit", ["megahit", "-r", "clean.fq.gz", "-o", "assembly"], - input_files=["clean.fq.gz"], - output_files=["assembly/final.contigs.fa"], - timestamp="2024-01-15T10:10:00Z", record_id="r2", - )) - _write_record(tmp_path, _make_record( - "quast", ["quast", "assembly/final.contigs.fa", "-o", "quast_out"], - input_files=["assembly/final.contigs.fa"], - timestamp="2024-01-15T10:20:00Z", record_id="r3", - )) + _write_record( + tmp_path, + _make_record( + "fastp", + ["fastp", "-q", "20", "--in1", "raw.fq.gz"], + output_files=["clean.fq.gz"], + timestamp="2024-01-15T10:00:00Z", + record_id="r1", + ), + ) + _write_record( + tmp_path, + _make_record( + "megahit", + ["megahit", "-r", "clean.fq.gz", "-o", "assembly"], + input_files=["clean.fq.gz"], + output_files=["assembly/final.contigs.fa"], + timestamp="2024-01-15T10:10:00Z", + record_id="r2", + ), + ) + _write_record( + tmp_path, + _make_record( + "quast", + ["quast", "assembly/final.contigs.fa", "-o", "quast_out"], + input_files=["assembly/final.contigs.fa"], + timestamp="2024-01-15T10:20:00Z", + record_id="r3", + ), + ) script = reconstruct_pipeline(tmp_path, session_id="s1", fmt="bash") diff --git a/tests/test_proxy_callback.py b/tests/test_proxy_callback.py deleted file mode 100644 index 369d73c..0000000 --- a/tests/test_proxy_callback.py +++ /dev/null @@ -1,226 +0,0 @@ -""" -Tests for the Phase 2 proxy callback (cache breakpoints + sidechannel -detection) in ``observability.py``. The callback is the only piece -LiteLLM autolog can't cover — request mutation (cache markers) and -canned-response detection (sidechannel) need the proxy hot path. - -LiteLLM's standard ``"otel"`` callback handles trace transport; -``DSAGTCallback`` only intercepts for these two intercept-time concerns. -""" -from __future__ import annotations - -import json -from pathlib import Path -from unittest.mock import patch - -import pytest - - -# --------------------------------------------------------------------------- -# _inject_cache_breakpoints -# --------------------------------------------------------------------------- - - -class TestInjectCacheBreakpoints: - """Anthropic prompt caching keys on the prefix UP TO each marked - block. We mark the last system text block + last tool definition - so subsequent turns within the 5-min TTL pay 10% on the cached - prefix. Providers without caching ignore the marker as a no-op.""" - - def test_marks_last_tool_definition(self): - from dsagt.observability import _inject_cache_breakpoints - messages: list = [] - kwargs = { - "tools": [ - {"name": "tool_a", "description": "a"}, - {"name": "tool_b", "description": "b"}, - ], - } - _inject_cache_breakpoints(messages, kwargs) - assert "cache_control" not in kwargs["tools"][0] - assert kwargs["tools"][1]["cache_control"] == {"type": "ephemeral"} - - def test_promotes_system_string_to_block_with_marker(self): - from dsagt.observability import _inject_cache_breakpoints - messages = [{"role": "system", "content": "you are helpful"}] - kwargs = {} - _inject_cache_breakpoints(messages, kwargs) - content = messages[0]["content"] - assert isinstance(content, list) - assert content[0]["type"] == "text" - assert content[0]["text"] == "you are helpful" - assert content[0]["cache_control"] == {"type": "ephemeral"} - - def test_marks_last_system_block(self): - from dsagt.observability import _inject_cache_breakpoints - messages = [{"role": "system", "content": [ - {"type": "text", "text": "rule 1"}, - {"type": "text", "text": "rule 2"}, - ]}] - kwargs = {} - _inject_cache_breakpoints(messages, kwargs) - blocks = messages[0]["content"] - assert "cache_control" not in blocks[0] - assert blocks[1]["cache_control"] == {"type": "ephemeral"} - - def test_no_tools_no_system_is_noop(self): - from dsagt.observability import _inject_cache_breakpoints - messages = [{"role": "user", "content": "hi"}] - kwargs = {} - _inject_cache_breakpoints(messages, kwargs) - assert "tools" not in kwargs - assert messages[0]["content"] == "hi" - - def test_only_first_system_message_marked(self): - """If multiple system messages somehow exist, only the first - gets stamped (loop breaks after first match).""" - from dsagt.observability import _inject_cache_breakpoints - messages = [ - {"role": "system", "content": "first"}, - {"role": "system", "content": "second"}, - ] - _inject_cache_breakpoints(messages, {}) - assert isinstance(messages[0]["content"], list) - assert messages[1]["content"] == "second" # untouched - - -# --------------------------------------------------------------------------- -# record_sidechannel_call / print_sidechannel_warning -# --------------------------------------------------------------------------- - - -class TestSidechannelDetection: - """``record_sidechannel_call`` writes to ``sidechannel.jsonl`` only - when ``kwargs["model"]`` matches ``SIDECHANNEL_CATCHALL_MODEL`` — - that's the wildcard's post-routing target, the only reliable - discriminator from primary / alias hits.""" - - def _kwargs(self, routed_to: str, requested: str | None = "gpt-4o-mini"): - slo = {} - if requested: - slo["model_group"] = f"openai/{requested}" - return {"model": routed_to, "standard_logging_object": slo} - - def test_records_when_routed_to_catchall(self, tmp_path, monkeypatch): - from dsagt.observability import ( - record_sidechannel_call, SIDECHANNEL_CATCHALL_MODEL, - SIDECHANNEL_LOG_FILENAME, - ) - monkeypatch.setenv("DSAGT_PRIMARY_MODEL", "real-model") - records_dir = tmp_path / "trace_archive" - records_dir.mkdir() - record_sidechannel_call(records_dir, self._kwargs(SIDECHANNEL_CATCHALL_MODEL)) - log_path = tmp_path / SIDECHANNEL_LOG_FILENAME - assert log_path.exists() - entry = json.loads(log_path.read_text().strip()) - assert entry["model"] == "gpt-4o-mini" - - def test_skipped_when_routed_to_primary(self, tmp_path, monkeypatch): - """Real upstream calls (primary entry or alias) shouldn't log — - only true mock hits do.""" - from dsagt.observability import ( - record_sidechannel_call, SIDECHANNEL_LOG_FILENAME, - ) - monkeypatch.setenv("DSAGT_PRIMARY_MODEL", "real-model") - records_dir = tmp_path / "trace_archive" - records_dir.mkdir() - record_sidechannel_call(records_dir, self._kwargs("openai/real-model")) - log_path = tmp_path / SIDECHANNEL_LOG_FILENAME - assert not log_path.exists() - - def test_no_op_when_primary_env_unset(self, tmp_path, monkeypatch): - """If the proxy never set DSAGT_PRIMARY_MODEL we don't know what - to compare against — refuse to log rather than misclassify.""" - from dsagt.observability import ( - record_sidechannel_call, SIDECHANNEL_CATCHALL_MODEL, - SIDECHANNEL_LOG_FILENAME, - ) - monkeypatch.delenv("DSAGT_PRIMARY_MODEL", raising=False) - records_dir = tmp_path / "trace_archive" - records_dir.mkdir() - record_sidechannel_call(records_dir, self._kwargs(SIDECHANNEL_CATCHALL_MODEL)) - assert not (tmp_path / SIDECHANNEL_LOG_FILENAME).exists() - - -class TestPrintSidechannelWarning: - """End-of-session warning: dedups by model name within the current - session, only emits ANSI on a TTY, says nothing when no entries - matched (most common case).""" - - def _write_log(self, project_dir: Path, entries: list[dict]): - from dsagt.observability import SIDECHANNEL_LOG_FILENAME - path = project_dir / SIDECHANNEL_LOG_FILENAME - path.write_text("\n".join(json.dumps(e) for e in entries) + "\n") - - def test_silent_when_no_log(self, tmp_path, capsys): - from dsagt.observability import print_sidechannel_warning - print_sidechannel_warning(tmp_path, "any-session") - assert capsys.readouterr().out == "" - - def test_silent_when_session_doesnt_match(self, tmp_path, capsys): - from dsagt.observability import print_sidechannel_warning - self._write_log(tmp_path, [ - {"model": "gpt-4o-mini", "session": "OTHER", "agent": "goose"}, - ]) - print_sidechannel_warning(tmp_path, "MINE") - assert capsys.readouterr().out == "" - - def test_lists_unique_models_with_counts(self, tmp_path, capsys): - from dsagt.observability import print_sidechannel_warning - self._write_log(tmp_path, [ - {"model": "gpt-4o-mini", "session": "S", "agent": "goose"}, - {"model": "gpt-4o-mini", "session": "S", "agent": "goose"}, - {"model": "claude-haiku-4-5", "session": "S", "agent": "claude"}, - ]) - print_sidechannel_warning(tmp_path, "S") - out = capsys.readouterr().out - assert "Sidechannel model calls intercepted" in out - assert "gpt-4o-mini" in out - assert "(2 calls)" in out - assert "claude-haiku-4-5" in out - assert "(1 call)" in out - - -# --------------------------------------------------------------------------- -# DSAGTCallback (the LiteLLM CustomLogger) -# --------------------------------------------------------------------------- - - -class TestDSAGTCallback: - """Minimal callback that hooks two LiteLLM events: - log_pre_api_call → cache injection - log_success_event → sidechannel detection. - Trace transport is via ``litellm.callbacks = ["otel"]``, not here. - """ - - def test_pre_api_call_injects_cache(self, tmp_path): - from dsagt.observability import _make_dsagt_callback - # litellm.integrations.custom_logger may not be installed in some - # CI envs; if so, skip — the callback is purely a litellm wrapper. - pytest.importorskip("litellm") - cb = _make_dsagt_callback(records_dir=tmp_path / "trace_archive") - messages = [{"role": "system", "content": "hi"}] - kwargs = {} - cb.log_pre_api_call(model="m", messages=messages, kwargs=kwargs) - assert messages[0]["content"][0]["cache_control"] == {"type": "ephemeral"} - - def test_success_event_records_sidechannel(self, tmp_path, monkeypatch): - from dsagt.observability import ( - _make_dsagt_callback, SIDECHANNEL_CATCHALL_MODEL, - SIDECHANNEL_LOG_FILENAME, - ) - pytest.importorskip("litellm") - monkeypatch.setenv("DSAGT_PRIMARY_MODEL", "real-model") - records_dir = tmp_path / "trace_archive" - records_dir.mkdir() - cb = _make_dsagt_callback(records_dir) - kwargs = { - "model": SIDECHANNEL_CATCHALL_MODEL, - "standard_logging_object": {"model_group": "openai/gpt-4o-mini"}, - } - cb.log_success_event(kwargs=kwargs, response_obj=None, - start_time=0, end_time=0) - log_path = tmp_path / SIDECHANNEL_LOG_FILENAME - assert log_path.exists() - entry = json.loads(log_path.read_text().strip()) - assert entry["model"] == "gpt-4o-mini" diff --git a/tests/test_recency.py b/tests/test_recency.py new file mode 100644 index 0000000..5277603 --- /dev/null +++ b/tests/test_recency.py @@ -0,0 +1,69 @@ +"""Recency weighting for episodic (session_memory) retrieval. + +``_apply_recency`` is a pure re-ranker, so these test the *behavior that matters* +without an embedder: a recent fact edges out a same-relevance stale one, while a +strongly-relevant old fact still beats an irrelevant recent one (the guard that +recency is a bounded boost, never a filter). +""" + +from dsagt.knowledge import _apply_recency + +DAY = 86400.0 +NOW = 1_000_000_000.0 + + +def _r(score, age_days, label): + ts = NOW - age_days * DAY + return {"chunk": {"text": label, "metadata": {"ts_epoch": ts}}, "score": score} + + +def _labels(results): + return [r["chunk"]["text"] for r in results] + + +def test_recent_edges_out_same_relevance_stale_fact(): + # Equal base relevance → the newer one wins purely on recency. + out = _apply_recency( + [_r(1.0, age_days=30, label="old"), _r(1.0, age_days=0, label="new")], + half_life_days=14, + now=NOW, + ) + assert _labels(out)[0] == "new" + + +def test_strong_old_fact_still_beats_irrelevant_recent_one(): + # The guard: recency is a bounded boost (≤ +50%), so a much-more-relevant old + # fact is never buried by a barely-relevant recent one. + out = _apply_recency( + [ + _r(1.0, age_days=60, label="relevant_old"), + _r(0.4, age_days=0, label="recent_junk"), + ], + half_life_days=14, + now=NOW, + ) + assert _labels(out)[0] == "relevant_old" + + +def test_half_life_controls_decay_strength(): + # At one half-life the boost halves: factor = 1 + 0.5 * 0.5 = 1.25. + (one,) = _apply_recency( + [_r(1.0, age_days=14, label="x")], half_life_days=14, now=NOW + ) + assert abs(one["recency_factor"] - 1.25) < 1e-9 + # Brand-new → full boost 1.5; very old → ~1.0 (no penalty below relevance). + (fresh,) = _apply_recency( + [_r(1.0, age_days=0, label="x")], half_life_days=14, now=NOW + ) + assert abs(fresh["recency_factor"] - 1.5) < 1e-9 + (ancient,) = _apply_recency( + [_r(1.0, age_days=3650, label="x")], half_life_days=14, now=NOW + ) + assert 1.0 <= ancient["recency_factor"] < 1.001 + + +def test_missing_ts_epoch_is_unweighted_not_dropped(): + res = [{"chunk": {"text": "no_ts", "metadata": {}}, "score": 1.0}] + (out,) = _apply_recency(res, half_life_days=14, now=NOW) + assert out["recency_factor"] == 1.0 + assert out["chunk"]["text"] == "no_ts" diff --git a/tests/test_registry.py b/tests/test_registry.py index 2b1c74f..ace3af3 100644 --- a/tests/test_registry.py +++ b/tests/test_registry.py @@ -1,5 +1,5 @@ """ -Tests for ToolRegistry and SkillRegistry. +Tests for CodeRegistry and SkillRegistry. Covers tool listing, MCP schema conversion, tool lookup, tool file writing, dsagt-run wrapping, runtime isolation, and skill discovery. @@ -9,11 +9,50 @@ import yaml from dsagt.registry import ( - ToolRegistry, SkillRegistry, _wrap_executable, _uv_run_prefix, _parse_frontmatter, + CodeRegistry, + _parse_frontmatter, + _lenient_frontmatter, render_arguments, ) +class TestLenientFrontmatter: + """Frontmatter that isn't strict YAML must still yield discovery fields. + + Real third-party skill catalogs (e.g. Genesis) ship SKILL.md files whose + unquoted ``description`` contains a colon (``...readiness levels: Level + 1...``) — invalid YAML. These must be recovered, not dropped from discovery. + """ + + def test_unquoted_colon_in_description_is_recovered(self, tmp_path): + path = tmp_path / "SKILL.md" + path.write_text( + "---\n" + "name: generating-datacards\n" + "description: Generates a datacard. Supports levels: Level 1, Level 2.\n" + "---\n\n# Body\n" + ) + spec = _parse_frontmatter(path) # must NOT raise + assert spec["name"] == "generating-datacards" + assert spec["description"].startswith("Generates a datacard") + assert "Level 1" in spec["description"] # colon-bearing tail preserved + + def test_lenient_parses_inline_list_and_continuation(self): + spec = _lenient_frontmatter( + "\nname: x\ndescription: a: b: c\ntags: [one, two]\n" + ) + assert spec["name"] == "x" + assert spec["description"] == "a: b: c" # split on first colon only + assert spec["tags"] == ["one", "two"] + + def test_valid_yaml_still_uses_strict_path(self, tmp_path): + # Sanity: well-formed frontmatter is unchanged by the fallback. + path = tmp_path / "SKILL.md" + path.write_text("---\nname: ok\ndescription: clean\ntags:\n - a\n---\n") + spec = _parse_frontmatter(path) + assert spec == {"name": "ok", "description": "clean", "tags": ["a"]} + + # --------------------------------------------------------------------------- # Fixtures # --------------------------------------------------------------------------- @@ -52,23 +91,22 @@ } -def _write_tool(tools_dir, spec: dict) -> None: - """Write a minimal skill file for the given spec dict.""" - path = tools_dir / f"{spec['name']}.md" +def _write_tool(codes_dir, spec: dict) -> None: + """Write a minimal skill-standard code dir for the given spec dict.""" + code_dir = codes_dir / spec["name"] + code_dir.mkdir(parents=True, exist_ok=True) frontmatter = yaml.dump(spec, default_flow_style=False, sort_keys=False) - path.write_text(f"---\n{frontmatter}---\n\n# {spec['name']}\n") + (code_dir / "SKILL.md").write_text(f"---\n{frontmatter}---\n\n# {spec['name']}\n") -def make_registry(tmp_path, tools: list[dict]) -> ToolRegistry: - """Create a ToolRegistry with the given tool definitions.""" - tools_dir = tmp_path / "source_skills" - tools_dir.mkdir() +def make_registry(tmp_path, tools: list[dict]) -> CodeRegistry: + """Create a CodeRegistry with the given tool definitions.""" + runtime_dir = tmp_path / "runtime" + codes_dir = runtime_dir / "codes" + codes_dir.mkdir(parents=True, exist_ok=True) for tool in tools: - _write_tool(tools_dir, tool) - return ToolRegistry( - source_tools_dir=str(tools_dir), - runtime_dir=str(tmp_path / "runtime"), - ) + _write_tool(codes_dir, tool) + return CodeRegistry(runtime_dir=str(runtime_dir)) @pytest.fixture @@ -84,14 +122,15 @@ def empty_registry(tmp_path): # --------------------------------------------------------------------------- -# list_tools +# list_codes # --------------------------------------------------------------------------- + class TestListTools: def test_schema_structure(self, registry): """MCP schema has name, description, and well-formed inputSchema.""" - tools = registry.list_tools() + tools = registry.list_codes() assert len(tools) == 1 tool = tools[0] @@ -105,7 +144,7 @@ def test_schema_structure(self, registry): def test_required_vs_optional(self, registry): """Required params appear in 'required', optional ones don't.""" - tool = registry.list_tools()[0] + tool = registry.list_codes()[0] required = tool["inputSchema"]["required"] assert "input_file" in required @@ -114,7 +153,7 @@ def test_required_vs_optional(self, registry): def test_default_values_propagate(self, registry): """Default values from the skill file appear in the MCP schema.""" - tool = registry.list_tools()[0] + tool = registry.list_codes()[0] props = tool["inputSchema"]["properties"] assert props["threshold"]["default"] == 0.5 @@ -122,12 +161,12 @@ def test_default_values_propagate(self, registry): def test_empty_registry(self, empty_registry): """An empty skills directory gives an empty tool list.""" - assert empty_registry.list_tools() == [] + assert empty_registry.list_codes() == [] def test_multiple_tools(self, tmp_path): """Multiple tools are listed in alphabetical filename order.""" reg = make_registry(tmp_path, [TOOL_WITH_MIXED_PARAMS, TOOL_NO_PARAMS]) - tools = reg.list_tools() + tools = reg.list_codes() assert len(tools) == 2 names = {t["name"] for t in tools} @@ -136,80 +175,102 @@ def test_multiple_tools(self, tmp_path): def test_no_params_tool(self, tmp_path): """Tool with empty parameters gives empty properties and required.""" reg = make_registry(tmp_path, [TOOL_NO_PARAMS]) - tool = reg.list_tools()[0] + tool = reg.list_codes()[0] assert tool["inputSchema"]["properties"] == {} assert tool["inputSchema"]["required"] == [] # --------------------------------------------------------------------------- -# get_tool +# get_code # --------------------------------------------------------------------------- + class TestGetTool: def test_found(self, registry): """Returns the raw tool definition for an existing tool.""" - tool = registry.get_tool("process") + tool = registry.get_code("process") assert tool is not None assert tool["name"] == "process" assert tool["executable"] == "python process.py" def test_not_found(self, registry): """Returns None for a nonexistent tool.""" - assert registry.get_tool("nonexistent") is None + assert registry.get_code("nonexistent") is None # --------------------------------------------------------------------------- # save_tool # --------------------------------------------------------------------------- + class TestSaveTool: def test_add_new_tool(self, empty_registry): """Saving a new tool creates a skill file with dsagt-run wrapping.""" empty_registry.save_tool(TOOL_NO_PARAMS) - tool = empty_registry.get_tool("ping") + tool = empty_registry.get_code("ping") assert tool is not None assert tool["name"] == "ping" - assert tool["executable"] == "dsagt-run --tool ping -- echo pong" + assert tool["executable"] == "dsagt-run --code ping -- echo pong" def test_wraps_executable_with_dsagt_run(self, empty_registry): """save_tool automatically wraps the executable with dsagt-run.""" - empty_registry.save_tool({"name": "mytool", "description": "test", - "executable": "python mytool.py", "parameters": {}}) - tool = empty_registry.get_tool("mytool") - assert tool["executable"] == "dsagt-run --tool mytool -- python mytool.py" + empty_registry.save_tool( + { + "name": "mytool", + "description": "test", + "executable": "python mytool.py", + "parameters": {}, + } + ) + tool = empty_registry.get_code("mytool") + assert tool["executable"] == "dsagt-run --code mytool -- python mytool.py" def test_does_not_double_wrap(self, empty_registry): """If executable already has dsagt-run, don't wrap again.""" - empty_registry.save_tool({"name": "mytool", "description": "test", - "executable": "dsagt-run --tool mytool -- python mytool.py", - "parameters": {}}) - tool = empty_registry.get_tool("mytool") + empty_registry.save_tool( + { + "name": "mytool", + "description": "test", + "executable": "dsagt-run --code mytool -- python mytool.py", + "parameters": {}, + } + ) + tool = empty_registry.get_code("mytool") assert tool["executable"].count("dsagt-run") == 1 def test_python_deps_use_uv_run(self, empty_registry): """Python dependencies are wrapped with uv run --with.""" - empty_registry.save_tool({ - "name": "analyzer", "description": "test", - "executable": "python analyzer.py", - "parameters": {}, - "dependencies": ["pandas>=2.0", "numpy"], - }) - tool = empty_registry.get_tool("analyzer") + empty_registry.save_tool( + { + "name": "analyzer", + "description": "test", + "executable": "python analyzer.py", + "parameters": {}, + "dependencies": ["pandas>=2.0", "numpy"], + } + ) + tool = empty_registry.get_code("analyzer") assert tool["executable"] == ( - "dsagt-run --tool analyzer -- uv run --with pandas>=2.0,numpy -- python analyzer.py" + "dsagt-run --code analyzer -- uv run --with pandas>=2.0,numpy -- python analyzer.py" ) def test_no_deps_no_uv_run(self, empty_registry): """Tools without dependencies don't get uv run prefix.""" - empty_registry.save_tool({"name": "simple", "description": "test", - "executable": "echo hi", "parameters": {}}) - tool = empty_registry.get_tool("simple") + empty_registry.save_tool( + { + "name": "simple", + "description": "test", + "executable": "echo hi", + "parameters": {}, + } + ) + tool = empty_registry.get_code("simple") assert "uv run" not in tool["executable"] - assert tool["executable"] == "dsagt-run --tool simple -- echo hi" + assert tool["executable"] == "dsagt-run --code simple -- echo hi" def test_add_returns_added(self, empty_registry): """save_tool returns 'added' for new tools.""" @@ -222,7 +283,8 @@ def test_update_returns_updated(self, empty_registry): def test_update_preserves_body(self, empty_registry): """Updating a tool preserves any hand-edited markdown body.""" - skill_path = empty_registry.tools_dir / "ping.md" + skill_path = empty_registry.codes_dir / "ping" / "SKILL.md" + skill_path.parent.mkdir(parents=True) spec = TOOL_NO_PARAMS fm = __import__("yaml").dump(spec, default_flow_style=False, sort_keys=False) skill_path.write_text(f"---\n{fm}---\n\n# Custom docs written by hand.\n") @@ -239,7 +301,7 @@ def test_update_overwrites_frontmatter(self, empty_registry): updated = {**TOOL_NO_PARAMS, "description": "New description"} empty_registry.save_tool(updated) - tool = empty_registry.get_tool("ping") + tool = empty_registry.get_code("ping") assert tool["description"] == "New description" @@ -247,47 +309,60 @@ def test_update_overwrites_frontmatter(self, empty_registry): # Runtime isolation # --------------------------------------------------------------------------- -class TestRuntimeIsolation: - - def test_source_unchanged_after_init(self, tmp_path): - """Source skills directory is not modified; runtime copy is separate.""" - source_dir = tmp_path / "source_skills" - source_dir.mkdir() - _write_tool(source_dir, TOOL_NO_PARAMS) - runtime_dir = tmp_path / "runtime" - reg = ToolRegistry( - source_tools_dir=str(source_dir), - runtime_dir=str(runtime_dir), +class TestBundledCopies: + + def test_copies_bundled_codes_into_project(self, tmp_path): + """ensure_bundled_copies lands each packaged code dir in codes/.""" + reg = CodeRegistry(runtime_dir=str(tmp_path / "rt")) + actions = reg.ensure_bundled_copies() + assert any("scan-directory" in a for a in actions) + copied = tmp_path / "rt" / "codes" / "scan-directory" + assert (copied / "SKILL.md").exists() + # Fully self-contained: the implementation script rides along. + assert (copied / "scripts" / "scan_directory.py").exists() + assert reg.get_code("scan-directory") is not None + + def test_never_clobbers_existing_copy(self, tmp_path): + """A user-edited (or agent-overridden) code dir is left untouched.""" + reg = CodeRegistry(runtime_dir=str(tmp_path / "rt")) + reg.ensure_bundled_copies() + spec = tmp_path / "rt" / "codes" / "scan-directory" / "SKILL.md" + spec.write_text(spec.read_text() + "\nUser edit.\n") + actions = reg.ensure_bundled_copies() + assert actions == [] + assert "User edit." in spec.read_text() + + def test_package_source_unmodified(self, tmp_path): + """Copying never touches the packaged source dirs.""" + before = sorted( + p.relative_to(CodeRegistry._PACKAGE_CODES_DIR) + for p in CodeRegistry._PACKAGE_CODES_DIR.rglob("*") ) - - # Add a tool to runtime - reg.save_tool(TOOL_WITH_MIXED_PARAMS) - - # Source should be unchanged - source_files = list(source_dir.glob("*.md")) - assert len(source_files) == 1 - assert source_files[0].stem == "ping" - - # Runtime should have both tools - assert len(reg.list_tools()) == 2 + CodeRegistry(runtime_dir=str(tmp_path / "rt")).ensure_bundled_copies() + after = sorted( + p.relative_to(CodeRegistry._PACKAGE_CODES_DIR) + for p in CodeRegistry._PACKAGE_CODES_DIR.rglob("*") + ) + assert before == after # --------------------------------------------------------------------------- # Default skills # --------------------------------------------------------------------------- + class TestDefaultTools: """Validate the tool files that ship with the package.""" def test_tools_directory_exists(self): """The package ships a tools directory.""" - assert ToolRegistry._PACKAGE_TOOLS_DIR.exists() - assert ToolRegistry._PACKAGE_TOOLS_DIR.is_dir() + assert CodeRegistry._PACKAGE_CODES_DIR.exists() + assert CodeRegistry._PACKAGE_CODES_DIR.is_dir() def test_tools_are_valid(self): """Every tool file must parse cleanly and have required fields.""" - tool_files = list(ToolRegistry._PACKAGE_TOOLS_DIR.glob("*.md")) + tool_files = list(CodeRegistry._PACKAGE_CODES_DIR.glob("*/SKILL.md")) assert len(tool_files) > 0, "No tool files found in package" for path in tool_files: @@ -297,19 +372,20 @@ def test_tools_are_valid(self): assert tool.get("executable"), f"{path.name}: missing 'executable'" assert "parameters" in tool, f"{path.name}: missing 'parameters'" - def test_default_init_fallback(self, tmp_path): - """ToolRegistry with no source_tools_dir falls back to package skills.""" - reg = ToolRegistry(source_tools_dir=None, runtime_dir=str(tmp_path / "rt")) - tools = reg.list_tools() - assert len(tools) > 0 - names = [t["name"] for t in tools] - assert "scan_directory" in names + def test_bundled_codes_appear_after_copy(self, tmp_path): + """A fresh registry is empty until ensure_bundled_copies runs.""" + reg = CodeRegistry(runtime_dir=str(tmp_path / "rt")) + assert reg.list_codes() == [] + reg.ensure_bundled_copies() + names = [t["name"] for t in reg.list_codes()] + assert "scan-directory" in names # --------------------------------------------------------------------------- # render_arguments # --------------------------------------------------------------------------- + class TestRenderArguments: def test_default_cli_is_double_dash_name(self): @@ -353,7 +429,8 @@ def test_positionals_before_named(self): "path": {"type": "string", "cli": "positional:0"}, } assert render_arguments(params, {"path": "/x", "verbose": True}) == [ - "/x", "--verbose", + "/x", + "--verbose", ] def test_boolean_true_emits_flag(self): @@ -378,7 +455,9 @@ def test_optional_missing_skipped(self): assert render_arguments(params, {}) == [] def test_required_missing_raises(self): - params = {"directory": {"type": "string", "cli": "positional", "required": True}} + params = { + "directory": {"type": "string", "cli": "positional", "required": True} + } with pytest.raises(ValueError, match="directory"): render_arguments(params, {}) diff --git a/tests/test_registry_server.py b/tests/test_registry_server.py index a7e8cb2..cc9ef41 100644 --- a/tests/test_registry_server.py +++ b/tests/test_registry_server.py @@ -1,26 +1,30 @@ """ Tests for the registry MCP server. -Tests tool handlers: save_tool_spec, get_registry, search_registry, -read_file, run_command, install_dependencies. +Tests tool handlers: save_code_spec, get_registry, search_registry, +read_file, run_command, http_request, install_dependencies. """ -import json import subprocess import sys from pathlib import Path from unittest.mock import patch, MagicMock +import httpx import pytest import yaml -from dsagt.registry import SkillRegistry, ToolRegistry -from dsagt.commands.registry_server import create_registry_server +from dsagt.registry import CodeRegistry +from dsagt.mcp.registry_tools import create_registry_server from mcp_helpers import call_tool_sync as call_tool -def make_spec(name="test_tool", description="A test tool", executable="echo hello", - dependencies=None): +def make_spec( + name="test-tool", + description="A test tool", + executable="echo hello", + dependencies=None, +): """Create a minimal valid tool spec.""" spec = { "name": name, @@ -39,32 +43,25 @@ def make_spec(name="test_tool", description="A test tool", executable="echo hell return spec -def _write_tool(tools_dir: Path, spec: dict) -> None: - path = tools_dir / f"{spec['name']}.md" +def _write_tool(codes_dir: Path, spec: dict) -> None: + code_dir = codes_dir / spec["name"] + code_dir.mkdir(parents=True, exist_ok=True) fm = yaml.dump(spec, default_flow_style=False, sort_keys=False) - path.write_text(f"---\n{fm}---\n\n# {spec['name']}\n") + (code_dir / "SKILL.md").write_text(f"---\n{fm}---\n\n# {spec['name']}\n") def _make_server(tmp_path, tools=None): """Create (server, registry) with optional pre-populated tools. - Pre-populated tools are written into ``/tools/`` (the - project layer) so they exercise the agent-saved code path — - ``reindex_all`` and most lookups happen here. The ``source_dir`` - is still passed as the bundled layer override but left empty; - tests that need bundled-layer behavior populate it explicitly. + Pre-populated tools are written into ``/codes/`` — the + single project layer every lookup reads. """ - source_dir = tmp_path / "source_skills" - source_dir.mkdir() runtime_dir = tmp_path / "runtime" - project_tools_dir = runtime_dir / "tools" + project_tools_dir = runtime_dir / "codes" project_tools_dir.mkdir(parents=True, exist_ok=True) - for spec in (tools or []): + for spec in tools or []: _write_tool(project_tools_dir, spec) - reg = ToolRegistry( - source_tools_dir=str(source_dir), - runtime_dir=str(runtime_dir), - ) + reg = CodeRegistry(runtime_dir=str(runtime_dir)) return create_registry_server(reg), reg @@ -72,6 +69,7 @@ def _make_server(tmp_path, tools=None): # Fixtures # --------------------------------------------------------------------------- + @pytest.fixture def server_and_registry(tmp_path): return _make_server(tmp_path) @@ -89,10 +87,13 @@ def registry(server_and_registry): @pytest.fixture def populated(tmp_path): - server, reg = _make_server(tmp_path, tools=[ - make_spec("tool_alpha", "Alpha tool", "python alpha.py"), - make_spec("tool_beta", "Beta data processor", "python beta.py"), - ]) + server, reg = _make_server( + tmp_path, + tools=[ + make_spec("tool-alpha", "Alpha tool", "python alpha.py"), + make_spec("tool-beta", "Beta data processor", "python beta.py"), + ], + ) return server, reg @@ -102,134 +103,70 @@ def populated_server(populated): # --------------------------------------------------------------------------- -# save_tool_spec +# save_code_spec # --------------------------------------------------------------------------- + class TestSaveToolSpec: def test_add_new_tool(self, server, registry): """Saving a new spec creates a skill file and reports added.""" - spec = make_spec("my_tool") - text = call_tool(server, "save_tool_spec", {"spec": spec}) + spec = make_spec("my-tool") + text = call_tool(server, "save_code_spec", {"spec": spec}) assert "added" in text assert "1 tools" in text - assert registry.get_tool("my_tool") is not None + assert registry.get_code("my-tool") is not None def test_update_existing_tool(self, server, registry): """Saving a spec with the same name updates rather than duplicates.""" - call_tool(server, "save_tool_spec", {"spec": make_spec("my_tool", description="Version 1")}) - text = call_tool(server, "save_tool_spec", {"spec": make_spec("my_tool", description="Version 2")}) + call_tool( + server, + "save_code_spec", + {"spec": make_spec("my-tool", description="Version 1")}, + ) + text = call_tool( + server, + "save_code_spec", + {"spec": make_spec("my-tool", description="Version 2")}, + ) assert "updated" in text assert "1 tools" in text - assert registry.get_tool("my_tool")["description"] == "Version 2" + assert registry.get_code("my-tool")["description"] == "Version 2" def test_add_multiple_tools(self, server, registry): """Multiple distinct tools accumulate as separate skill files.""" - call_tool(server, "save_tool_spec", {"spec": make_spec("tool_a")}) - text = call_tool(server, "save_tool_spec", {"spec": make_spec("tool_b")}) + call_tool(server, "save_code_spec", {"spec": make_spec("tool-a")}) + text = call_tool(server, "save_code_spec", {"spec": make_spec("tool-b")}) assert "2 tools" in text - assert registry.get_tool("tool_a") is not None - assert registry.get_tool("tool_b") is not None + assert registry.get_code("tool-a") is not None + assert registry.get_code("tool-b") is not None def test_accepts_stringified_spec(self, server, registry): """Some MCP clients (Claude Sonnet/Haiku 4.x) send nested-object args as JSON strings. The handler must accept both shapes.""" import json - spec = make_spec("stringy_tool") - text = call_tool(server, "save_tool_spec", {"spec": json.dumps(spec)}) + + spec = make_spec("stringy-tool") + text = call_tool(server, "save_code_spec", {"spec": json.dumps(spec)}) assert "added" in text - assert registry.get_tool("stringy_tool") is not None + assert registry.get_code("stringy-tool") is not None def test_rejects_invalid_stringified_spec(self, server, registry): """Non-JSON strings produce a clear error rather than crashing.""" - text = call_tool(server, "save_tool_spec", {"spec": "not valid json {"}) + text = call_tool(server, "save_code_spec", {"spec": "not valid json {"}) assert "Error" in text assert "JSON object" in text # --------------------------------------------------------------------------- -# save_skill +# install_skill # --------------------------------------------------------------------------- -class TestSaveSkill: - - def test_add_new_skill_creates_files_and_indexes(self, tmp_path): - """save_skill writes SKILL.md and indexes when KB is configured. - - The skill count includes any bundled skills that ship in the - package (see SkillRegistry.list_skills which merges bundled + - project layers), so we assert the file was created and the - count went up by one rather than equality on a specific number. - """ - server, reg, kb = _make_server_with_kb(tmp_path) - from dsagt.registry import SkillRegistry as _SR - skill_reg = _SR(runtime_dir=str(tmp_path / "runtime"), kb=kb) - before = len(skill_reg.list_skills()) - - spec = { - "name": "csv_inspector", - "description": "Workflow for inspecting CSV columns and quality", - "tags": ["data_management", "quality_control"], - } - body = "# csv_inspector\n\nFirst, run head on the file. Then check nulls.\n" - text = call_tool(server, "save_skill", {"spec": spec, "body": body}) - - assert "added" in text - skill_md = tmp_path / "runtime" / "skills" / "csv_inspector" / "SKILL.md" - assert skill_md.exists() - content = skill_md.read_text() - assert "csv_inspector" in content - assert "First, run head" in content - after = len(skill_reg.list_skills()) - assert after == before + 1 - - def test_update_existing_skill_preserves_body_when_omitted(self, tmp_path): - """Saving a spec for an existing skill without body keeps the body.""" - server, reg, kb = _make_server_with_kb(tmp_path) - first_body = "# orig\n\nOriginal workflow body.\n" - call_tool(server, "save_skill", { - "spec": {"name": "wf", "description": "v1"}, - "body": first_body, - }) - # Update the description only — body should be preserved. - text = call_tool(server, "save_skill", { - "spec": {"name": "wf", "description": "v2 description"}, - }) - assert "updated" in text - skill_md = tmp_path / "runtime" / "skills" / "wf" / "SKILL.md" - content = skill_md.read_text() - assert "v2 description" in content - assert "Original workflow body" in content - - def test_save_skill_writes_reference_files(self, tmp_path): - """reference_files dict lands as additional files in the skill dir.""" - server, reg, kb = _make_server_with_kb(tmp_path) - text = call_tool(server, "save_skill", { - "spec": {"name": "with_template", "description": "Has a template"}, - "body": "# with_template\n\nUses template.json.\n", - "reference_files": {"template.json": '{"foo": "bar"}\n'}, - }) - assert "added" in text - skill_dir = tmp_path / "runtime" / "skills" / "with_template" - assert (skill_dir / "SKILL.md").exists() - assert (skill_dir / "template.json").read_text() == '{"foo": "bar"}\n' - - def test_save_skill_string_encoded_spec(self, tmp_path): - """MCP clients that JSON-encode nested object args still work.""" - server, reg, kb = _make_server_with_kb(tmp_path) - spec_json = json.dumps({"name": "s1", "description": "d"}) - text = call_tool(server, "save_skill", {"spec": spec_json, "body": "x"}) - assert "added" in text - - -# --------------------------------------------------------------------------- -# get_registry -# --------------------------------------------------------------------------- class TestGetRegistry: @@ -243,16 +180,17 @@ def test_populated_registry(self, populated_server, populated): text = call_tool(populated_server, "get_registry", {}) data = yaml.safe_load(text) - assert len(data["tools"]) == 2 - names = [t["name"] for t in data["tools"]] - assert "tool_alpha" in names - assert "tool_beta" in names + assert len(data["codes"]) == 2 + names = [t["name"] for t in data["codes"]] + assert "tool-alpha" in names + assert "tool-beta" in names # --------------------------------------------------------------------------- # search_registry # --------------------------------------------------------------------------- + class TestSearchRegistryNoKB: """search_registry with no KB configured. @@ -265,19 +203,29 @@ class TestSearchRegistryNoKB: """ def test_exact_name_lookup_works_without_kb(self, populated_server): - """tool_name lookup is KB-free and must keep working.""" - text = call_tool(populated_server, "search_registry", {"tool_name": "tool_alpha"}) - assert "tool_alpha" in text + """code_name lookup is KB-free and must keep working.""" + text = call_tool( + populated_server, "search_registry", {"code_name": "tool-alpha"} + ) + assert "tool-alpha" in text def test_exact_name_miss_without_kb(self, populated_server): - """tool_name with a non-existent name returns a clean 'no tool' message.""" - text = call_tool(populated_server, "search_registry", {"tool_name": "nonexistent"}) + """code_name with a non-existent name returns a clean 'no tool' message.""" + text = call_tool( + populated_server, "search_registry", {"code_name": "nonexistent"} + ) assert "No tool named 'nonexistent'" in text def test_query_search_without_kb_returns_helpful_error(self, populated_server): """A semantic search request when no KB is configured must surface - the missing-KB condition clearly, not silently degrade.""" + the missing-KB condition clearly, not silently degrade. + + The query "alpha" is a substring of the registered ``tool_alpha``; + the deleted string-matching fallback would have returned it, so the + ``not in`` assertion pins that the fallback stays gone. + """ text = call_tool(populated_server, "search_registry", {"query": "alpha"}) + assert "tool-alpha" not in text # no silent substring fallback assert "knowledge base" in text.lower() assert "embedding" in text.lower() @@ -291,6 +239,7 @@ def test_empty_query_without_kb_returns_helpful_error(self, populated_server): # read_file # --------------------------------------------------------------------------- + class TestReadFile: def test_read_success(self, server, tmp_path): @@ -311,105 +260,192 @@ def test_read_missing_file(self, server): # run_command # --------------------------------------------------------------------------- + class TestRunCommand: def test_success(self, server): """Running a valid command returns its output.""" - text = call_tool(server, "run_command", { - "command": "echo", - "args": ["hello"], - }) + text = call_tool( + server, + "run_command", + { + "command": "echo", + "args": ["hello"], + }, + ) assert "hello" in text assert "Return code: 0" in text def test_command_not_found(self, server): """Running a nonexistent command returns not found error.""" - text = call_tool(server, "run_command", { - "command": "nonexistent_command_xyz", - }) + text = call_tool( + server, + "run_command", + { + "command": "nonexistent_command_xyz", + }, + ) assert "not found" in text def test_timeout(self, server): """A command that exceeds the timeout reports timeout.""" - text = call_tool(server, "run_command", { - "command": "sleep", - "args": ["30"], - "timeout": 0.1, - }) + text = call_tool( + server, + "run_command", + { + "command": "sleep", + "args": ["30"], + "timeout": 0.1, + }, + ) assert "timed out" in text # --------------------------------------------------------------------------- -# save_tool_spec — dependency installation +# http_request # --------------------------------------------------------------------------- + +class TestHttpRequest: + + @patch("dsagt.mcp.registry_tools.httpx.AsyncClient") + def test_success_with_defaults(self, mock_client_cls, server): + """A plain URL issues a GET and returns status + body.""" + client = mock_client_cls.return_value.__aenter__.return_value + client.request.return_value = MagicMock(status_code=200, text="pong") + + text = call_tool( + server, + "http_request", + {"url": "https://example.test/ping"}, + ) + + assert text == "Status: 200\n\npong" + client.request.assert_awaited_once_with( + method="GET", + url="https://example.test/ping", + headers={}, + timeout=30.0, + ) + + @patch("dsagt.mcp.registry_tools.httpx.AsyncClient") + def test_method_and_headers_forwarded(self, mock_client_cls, server): + """Explicit method and headers reach the client unchanged.""" + client = mock_client_cls.return_value.__aenter__.return_value + client.request.return_value = MagicMock(status_code=201, text="created") + + text = call_tool( + server, + "http_request", + { + "url": "https://example.test/items", + "method": "POST", + "headers": {"Authorization": "Bearer tok"}, + }, + ) + + assert text.startswith("Status: 201") + client.request.assert_awaited_once_with( + method="POST", + url="https://example.test/items", + headers={"Authorization": "Bearer tok"}, + timeout=30.0, + ) + + @patch("dsagt.mcp.registry_tools.httpx.AsyncClient") + def test_transport_error_reported(self, mock_client_cls, server): + """httpx transport failures surface as a clean error string.""" + client = mock_client_cls.return_value.__aenter__.return_value + client.request.side_effect = httpx.ConnectError("Connection refused") + + text = call_tool( + server, + "http_request", + {"url": "https://example.test/down"}, + ) + + assert text == "Error making request: Connection refused" + + +# --------------------------------------------------------------------------- +# save_code_spec — dependency installation +# --------------------------------------------------------------------------- + + class TestSaveToolSpecDependencies: - @patch("dsagt.commands.registry_server.subprocess.run") + @patch("dsagt.mcp.registry_tools.subprocess.run") def test_deps_installed_on_save(self, mock_run, server, registry): """When dependencies are provided, uv pip install is called.""" mock_run.return_value = MagicMock( returncode=0, stdout="Successfully installed pandas-2.1.0", stderr="" ) - spec = make_spec("tool_with_deps", dependencies=["pandas>=2.0", "numpy"]) - text = call_tool(server, "save_tool_spec", {"spec": spec}) + spec = make_spec("tool-with-deps", dependencies=["pandas>=2.0", "numpy"]) + text = call_tool(server, "save_code_spec", {"spec": spec}) assert "added" in text assert "Successfully installed" in text mock_run.assert_called_once() cmd = mock_run.call_args[0][0] - assert cmd == ["uv", "pip", "install", "--python", sys.executable, - "pandas>=2.0", "numpy"] - - @patch("dsagt.commands.registry_server.subprocess.run") + assert cmd == [ + "uv", + "pip", + "install", + "--python", + sys.executable, + "pandas>=2.0", + "numpy", + ] + + @patch("dsagt.mcp.registry_tools.subprocess.run") def test_deps_failure_still_saves_spec(self, mock_run, server, registry): """Even if uv pip install fails, the spec is saved as a skill file.""" mock_run.return_value = MagicMock( returncode=1, stdout="", stderr="No matching distribution for bogus-pkg" ) - spec = make_spec("tool_bad_deps", dependencies=["bogus-pkg"]) - text = call_tool(server, "save_tool_spec", {"spec": spec}) + spec = make_spec("tool-bad-deps", dependencies=["bogus-pkg"]) + text = call_tool(server, "save_code_spec", {"spec": spec}) assert "added" in text assert "Installation failed" in text - tool = registry.get_tool("tool_bad_deps") + tool = registry.get_code("tool-bad-deps") assert tool is not None assert tool["dependencies"] == ["bogus-pkg"] - @patch("dsagt.commands.registry_server.subprocess.run") + @patch("dsagt.mcp.registry_tools.subprocess.run") def test_deps_timeout(self, mock_run, server): """Timeout during install is reported, spec is still saved.""" mock_run.side_effect = subprocess.TimeoutExpired("uv", 120) - spec = make_spec("tool_slow_deps", dependencies=["heavy-pkg"]) - text = call_tool(server, "save_tool_spec", {"spec": spec}) + spec = make_spec("tool-slow-deps", dependencies=["heavy-pkg"]) + text = call_tool(server, "save_code_spec", {"spec": spec}) assert "added" in text assert "timed out" in text def test_no_deps_no_install_message(self, server, registry): """When no dependencies are provided, no install message appears.""" - spec = make_spec("tool_no_deps") - text = call_tool(server, "save_tool_spec", {"spec": spec}) + spec = make_spec("tool-no-deps") + text = call_tool(server, "save_code_spec", {"spec": spec}) assert "added" in text assert "Dependency" not in text - @patch("dsagt.commands.registry_server.subprocess.run") + @patch("dsagt.mcp.registry_tools.subprocess.run") def test_deps_persisted_in_skill_file(self, mock_run, server, registry): """Dependencies are stored in the skill file frontmatter.""" mock_run.return_value = MagicMock(returncode=0, stdout="ok", stderr="") - spec = make_spec("dep_tool", dependencies=["requests>=2.28"]) - call_tool(server, "save_tool_spec", {"spec": spec}) + spec = make_spec("dep-tool", dependencies=["requests>=2.28"]) + call_tool(server, "save_code_spec", {"spec": spec}) - tool = registry.get_tool("dep_tool") + tool = registry.get_code("dep-tool") assert tool["dependencies"] == ["requests>=2.28"] - @patch("dsagt.commands.registry_server.subprocess.run") + @patch("dsagt.mcp.registry_tools.subprocess.run") def test_uv_not_found(self, mock_run, server): """FileNotFoundError from missing uv is reported gracefully.""" mock_run.side_effect = FileNotFoundError("uv") - spec = make_spec("tool_no_uv", dependencies=["pandas"]) - text = call_tool(server, "save_tool_spec", {"spec": spec}) + spec = make_spec("tool-no-uv", dependencies=["pandas"]) + text = call_tool(server, "save_code_spec", {"spec": spec}) assert "added" in text assert "'uv' command not found" in text @@ -419,40 +455,55 @@ def test_uv_not_found(self, mock_run, server): # install_dependencies # --------------------------------------------------------------------------- + class TestInstallDependencies: - @patch("dsagt.commands.registry_server.subprocess.run") + @patch("dsagt.mcp.registry_tools.subprocess.run") def test_install_all(self, mock_run, tmp_path): - """install_dependencies with no tool_name installs all unique deps.""" - server, reg = _make_server(tmp_path, tools=[ - make_spec("tool_a", dependencies=["pandas", "numpy"]), - make_spec("tool_b", dependencies=["numpy", "scipy"]), - ]) + """install_dependencies with no code_name installs all unique deps.""" + server, reg = _make_server( + tmp_path, + tools=[ + make_spec("tool-a", dependencies=["pandas", "numpy"]), + make_spec("tool-b", dependencies=["numpy", "scipy"]), + ], + ) mock_run.return_value = MagicMock(returncode=0, stdout="ok", stderr="") text = call_tool(server, "install_dependencies", {}) - assert "tool_a" in text - assert "tool_b" in text + assert "tool-a" in text + assert "tool-b" in text cmd = mock_run.call_args[0][0] - assert cmd == ["uv", "pip", "install", "--python", sys.executable, - "pandas", "numpy", "scipy"] - - @patch("dsagt.commands.registry_server.subprocess.run") + assert cmd == [ + "uv", + "pip", + "install", + "--python", + sys.executable, + "pandas", + "numpy", + "scipy", + ] + + @patch("dsagt.mcp.registry_tools.subprocess.run") def test_install_single_tool(self, mock_run, tmp_path): - """install_dependencies with tool_name targets only that tool.""" - server, reg = _make_server(tmp_path, tools=[ - make_spec("tool_a", dependencies=["pandas"]), - make_spec("tool_b", dependencies=["scipy"]), - ]) + """install_dependencies with code_name targets only that tool.""" + server, reg = _make_server( + tmp_path, + tools=[ + make_spec("tool-a", dependencies=["pandas"]), + make_spec("tool-b", dependencies=["scipy"]), + ], + ) mock_run.return_value = MagicMock(returncode=0, stdout="ok", stderr="") - text = call_tool(server, "install_dependencies", {"tool_name": "tool_b"}) + text = call_tool(server, "install_dependencies", {"code_name": "tool-b"}) cmd = mock_run.call_args[0][0] assert cmd == ["uv", "pip", "install", "--python", sys.executable, "scipy"] - assert "tool_b" in text - assert "tool_a" not in text + assert "tool-b" in text + assert "tool-a" not in text def test_no_deps_in_registry(self, server): """install_dependencies on empty registry reports no tools.""" @@ -471,38 +522,30 @@ def test_tools_without_deps(self, tmp_path): # KB-backed tool indexing and search # --------------------------------------------------------------------------- + def _make_server_with_kb(tmp_path, tools=None): """Create (server, registry, kb) with a real local-embedding KnowledgeBase. Pre-populated tools are written to ``/tools/`` so they - exercise the agent-saved code path that ``reindex_all`` operates on. + exercise the agent-saved code path. """ from dsagt.knowledge import KnowledgeBase - source_dir = tmp_path / "source_skills" - source_dir.mkdir() runtime_dir = tmp_path / "runtime" - project_tools_dir = runtime_dir / "tools" + project_tools_dir = runtime_dir / "codes" project_tools_dir.mkdir(parents=True, exist_ok=True) - for spec in (tools or []): + for spec in tools or []: _write_tool(project_tools_dir, spec) kb = KnowledgeBase( index_dir=tmp_path / "kb_index", default_embedder="local", - default_index="chroma", - ) - reg = ToolRegistry( - source_tools_dir=str(source_dir), - runtime_dir=str(runtime_dir), - kb=kb, ) - skill_reg = SkillRegistry( - source_skills_dir=None, # use package default (empty bundled is fine) + reg = CodeRegistry( runtime_dir=str(runtime_dir), kb=kb, ) - server = create_registry_server(reg, kb, skill_reg) + server = create_registry_server(reg, kb) return server, reg, kb @@ -515,40 +558,39 @@ def test_save_tool_indexes_into_kb(self, tmp_path): server, reg, kb = _make_server_with_kb(tmp_path) - call_tool(server, "save_tool_spec", {"spec": make_spec( - name="csv_filter", - description="Filter CSV rows by column value", - )}) + call_tool( + server, + "save_code_spec", + { + "spec": make_spec( + name="csv-filter", + description="Filter CSV rows by column value", + ) + }, + ) results = kb.search("filter", collection=TOOL_REGISTRY_COLLECTION) assert len(results) > 0 - assert any("csv_filter" in r["chunk"].get("text", "") for r in results) - - def test_search_registry_by_name(self, tmp_path): - """Exact tool_name lookup returns the tool.""" - server, reg, kb = _make_server_with_kb(tmp_path) - call_tool(server, "save_tool_spec", {"spec": make_spec(name="fastp")}) - - text = call_tool(server, "search_registry", {"tool_name": "fastp"}) - assert "fastp" in text - - def test_search_registry_by_name_not_found(self, tmp_path): - """Exact lookup for nonexistent tool returns not found.""" - server, reg, kb = _make_server_with_kb(tmp_path) - - text = call_tool(server, "search_registry", {"tool_name": "nonexistent"}) - assert "No tool" in text + assert any("csv-filter" in r["chunk"].get("text", "") for r in results) def test_search_registry_semantic(self, tmp_path): """Semantic search finds tools by description similarity.""" server, reg, kb = _make_server_with_kb(tmp_path) - call_tool(server, "save_tool_spec", {"spec": make_spec( - name="csv_filter", - description="Filter and remove rows from a CSV spreadsheet based on column values", - )}) + call_tool( + server, + "save_code_spec", + { + "spec": make_spec( + name="csv-filter", + description="Filter and remove rows from a CSV spreadsheet based on column values", + ) + }, + ) - text = call_tool(server, "search_registry", {"query": "delete rows from tabular data"}) - assert "csv_filter" in text + text = call_tool( + server, "search_registry", {"query": "delete rows from tabular data"} + ) + assert "csv-filter" in text def test_search_registry_by_tag(self, tmp_path): """Tag-based filtering returns only matching tools.""" @@ -556,46 +598,13 @@ def test_search_registry_by_tag(self, tmp_path): spec_genomics = make_spec(name="fastp", description="FASTQ preprocessor") spec_genomics["tags"] = ["genomics", "data_processing"] - call_tool(server, "save_tool_spec", {"spec": spec_genomics}) + call_tool(server, "save_code_spec", {"spec": spec_genomics}) spec_other = make_spec(name="csvtool", description="CSV processor") spec_other["tags"] = ["data_processing"] - call_tool(server, "save_tool_spec", {"spec": spec_other}) - - text = call_tool(server, "search_registry", {"query": "tool", "tag": "genomics"}) - assert "fastp" in text + call_tool(server, "save_code_spec", {"spec": spec_other}) - def test_reindex_all(self, tmp_path): - """reindex_all populates KB from existing skill files.""" - from dsagt.registry import TOOL_REGISTRY_COLLECTION - - server, reg, kb = _make_server_with_kb( - tmp_path, - tools=[make_spec(name="preexisting", description="Already registered tool")], + text = call_tool( + server, "search_registry", {"query": "tool", "tag": "genomics"} ) - - # Skills were copied to runtime on init but not indexed (KB was empty) - # reindex_all should pick them up - count = reg.reindex_all() - assert count >= 1 - - results = kb.search("registered", collection=TOOL_REGISTRY_COLLECTION) - assert len(results) > 0 - - def test_no_kb_query_search_returns_explicit_error(self, tmp_path): - """Without a configured KB, query-based search MUST NOT silently - fall back to substring matching. It must return an explicit - error so the user knows the KB is missing. - - Regression test for the deletion of the string-matching fallback - in search_registry. The old fallback hid embedding/KB failures - and produced dramatically worse search results without telling - anyone. - """ - server, reg = _make_server(tmp_path, tools=[ - make_spec(name="csv_filter", description="Filter CSV rows"), - ]) - - text = call_tool(server, "search_registry", {"query": "csv"}) - assert "csv_filter" not in text # the substring match must NOT happen - assert "knowledge base" in text.lower() + assert "fastp" in text diff --git a/tests/test_run.py b/tests/test_run.py index fe37f03..c0be23d 100644 --- a/tests/test_run.py +++ b/tests/test_run.py @@ -6,10 +6,7 @@ """ import json -import os -import sys from pathlib import Path -from unittest.mock import patch import pytest @@ -19,34 +16,45 @@ _write_record, run_and_record, ) -from dsagt.commands.run_tool import ( +from dsagt.commands.run_code import ( _parse_args, main, ) - # --------------------------------------------------------------------------- # Argument parsing # --------------------------------------------------------------------------- + class TestParseArgs: def test_basic(self): - args, command = _parse_args(["--tool", "fastp", "--", "fastp", "-q", "20"]) - assert args.tool == "fastp" + args, command = _parse_args(["--code", "fastp", "--", "fastp", "-q", "20"]) + assert args.code == "fastp" assert command == ["fastp", "-q", "20"] def test_all_flags(self): - args, command = _parse_args([ - "--tool", "megahit", - "--session", "sess-1", - "--record-id", "rec-42", - "--records-dir", "/tmp/records", - "--input-files", "a.fq,b.fq", - "--output-files", "out/contigs.fa", - "--", "megahit", "-1", "a.fq", - ]) - assert args.tool == "megahit" + args, command = _parse_args( + [ + "--code", + "megahit", + "--session", + "sess-1", + "--record-id", + "rec-42", + "--records-dir", + "/tmp/records", + "--input-files", + "a.fq,b.fq", + "--output-files", + "out/contigs.fa", + "--", + "megahit", + "-1", + "a.fq", + ] + ) + assert args.code == "megahit" assert args.session == "sess-1" assert args.record_id == "rec-42" assert args.records_dir == "/tmp/records" @@ -57,10 +65,10 @@ def test_all_flags(self): def test_no_separator_exits(self): """Missing '--' separator causes a SystemExit (from argparse --help).""" with pytest.raises(SystemExit): - _parse_args(["--tool", "fastp", "fastp", "-q", "20"]) + _parse_args(["--code", "fastp", "fastp", "-q", "20"]) def test_defaults(self): - args, _ = _parse_args(["--tool", "x", "--", "echo"]) + args, _ = _parse_args(["--code", "x", "--", "echo"]) assert args.session is None assert args.record_id is None assert args.records_dir is None @@ -72,6 +80,7 @@ def test_defaults(self): # File list parsing # --------------------------------------------------------------------------- + class TestParseFileList: def test_none(self): @@ -94,24 +103,26 @@ def test_trailing_comma(self): # Records directory resolution # --------------------------------------------------------------------------- + class TestResolveRecordsDir: def test_explicit_wins(self): assert _resolve_records_dir("/custom/dir") == Path("/custom/dir") def test_uses_cwd_dsagt_config(self, tmp_path, monkeypatch): - """No --records-dir → reads ``/dsagt_config.yaml`` and uses + """No --records-dir → reads ``/.dsagt/config.yaml`` and uses ``/trace_archive``. Env vars are not consulted; the project dir is the single source of truth.""" - (tmp_path / "dsagt_config.yaml").write_text("project: t\n") + (tmp_path / ".dsagt").mkdir() + (tmp_path / ".dsagt" / "config.yaml").write_text("project: t\n") monkeypatch.chdir(tmp_path) assert _resolve_records_dir(None) == tmp_path / "trace_archive" def test_no_config_in_cwd_raises(self, tmp_path, monkeypatch): - """If cwd has no dsagt_config.yaml, fail clearly — don't walk + """If cwd has no .dsagt/config.yaml, fail clearly — don't walk up the tree, don't fall back to env vars.""" monkeypatch.chdir(tmp_path) - with pytest.raises(ValueError, match="No dsagt_config.yaml"): + with pytest.raises(ValueError, match="No .dsagt/config.yaml"): _resolve_records_dir(None) @@ -119,13 +130,14 @@ def test_no_config_in_cwd_raises(self, tmp_path, monkeypatch): # Record writing # --------------------------------------------------------------------------- + class TestWriteRecord: def test_creates_directory_and_file(self, tmp_path): records_dir = tmp_path / "nested" / "records" record = { "record_id": "abc123", - "tool_name": "fastp", + "code_name": "fastp", "session_id": None, "execution": { "exact_command": ["fastp", "-q", "20"], @@ -146,13 +158,13 @@ def test_creates_directory_and_file(self, tmp_path): assert "abc123" in path.name data = json.loads(path.read_text()) - assert data["tool_name"] == "fastp" + assert data["code_name"] == "fastp" assert data["execution"]["return_code"] == 0 def test_record_is_valid_json(self, tmp_path): record = { "record_id": "x", - "tool_name": "t", + "code_name": "t", "session_id": None, "execution": { "exact_command": ["echo"], @@ -174,12 +186,13 @@ def test_record_is_valid_json(self, tmp_path): # run_and_record # --------------------------------------------------------------------------- + class TestRunAndRecord: def test_successful_command(self, tmp_path): """Runs echo, captures output, writes record, returns 0.""" exit_code = run_and_record( - tool_name="echo_test", + code_name="echo_test", command=["echo", "hello world"], records_dir=tmp_path, record_id="test-001", @@ -191,7 +204,7 @@ def test_successful_command(self, tmp_path): assert len(records) == 1 data = json.loads(records[0].read_text()) - assert data["tool_name"] == "echo_test" + assert data["code_name"] == "echo_test" assert data["record_id"] == "test-001" assert data["execution"]["return_code"] == 0 assert "hello world" in data["execution"]["stdout"] @@ -200,7 +213,7 @@ def test_successful_command(self, tmp_path): def test_failing_command(self, tmp_path): """A command that fails returns non-zero and captures stderr.""" exit_code = run_and_record( - tool_name="false_test", + code_name="false_test", command=["bash", "-c", "echo oops >&2; exit 42"], records_dir=tmp_path, record_id="test-002", @@ -215,7 +228,7 @@ def test_failing_command(self, tmp_path): def test_command_not_found(self, tmp_path): """A missing command returns exit code 127.""" exit_code = run_and_record( - tool_name="missing", + code_name="missing", command=["this_command_does_not_exist_xyz"], records_dir=tmp_path, record_id="test-003", @@ -227,40 +240,34 @@ def test_command_not_found(self, tmp_path): assert data["execution"]["return_code"] == 127 assert "command not found" in data["execution"]["stderr"] - def test_session_from_runtime(self, tmp_path, monkeypatch): - """Session ID falls back to /.runtime when not passed. + def test_session_from_state(self, tmp_path, monkeypatch): + """Session ID falls back to the current tag in ``.dsagt/state.yaml`` + when not passed — the MCP server mints it there at startup and + ``dsagt-run`` (cwd == project dir) reads it.""" + from dsagt.session import append_session, write_config_file, build_config - Single source of truth: dsagt mlflow writes session_id into - .runtime; services read from there, not from env vars. - """ - (tmp_path / "dsagt_config.yaml").write_text( - "project: t\nmlflow: {port: 5000}\n" - ) - (tmp_path / ".runtime").write_text( - json.dumps({"session_id": "runtime-session"}) - ) + write_config_file(tmp_path, build_config("t", "claude")) + append_session(tmp_path) # mints session id 1 → tag "t-1" monkeypatch.chdir(tmp_path) run_and_record( - tool_name="t", + code_name="t", command=["echo"], records_dir=tmp_path, record_id="test-004", ) data = json.loads(list(tmp_path.glob("*_test-004.json"))[0].read_text()) - assert data["session_id"] == "runtime-session" + assert data["session_id"] == "t-1" - def test_explicit_session_overrides_runtime(self, tmp_path, monkeypatch): - """Explicit --session takes precedence over .runtime.""" - (tmp_path / "dsagt_config.yaml").write_text( - "project: t\nmlflow: {port: 5000}\n" - ) - (tmp_path / ".runtime").write_text( - json.dumps({"session_id": "runtime-session"}) - ) + def test_explicit_session_overrides_state(self, tmp_path, monkeypatch): + """Explicit --session takes precedence over the state-file tag.""" + from dsagt.session import append_session, write_config_file, build_config + + write_config_file(tmp_path, build_config("t", "claude")) + append_session(tmp_path) monkeypatch.chdir(tmp_path) run_and_record( - tool_name="t", + code_name="t", command=["echo"], records_dir=tmp_path, session_id="explicit-session", @@ -273,7 +280,7 @@ def test_explicit_session_overrides_runtime(self, tmp_path, monkeypatch): def test_file_lists_recorded(self, tmp_path): """Input and output file lists appear in the record.""" run_and_record( - tool_name="t", + code_name="t", command=["echo"], records_dir=tmp_path, input_files=["in1.fq", "in2.fq"], @@ -288,7 +295,7 @@ def test_file_lists_recorded(self, tmp_path): def test_timestamps_are_populated(self, tmp_path): """Start and end timestamps are non-empty ISO strings.""" run_and_record( - tool_name="t", + code_name="t", command=["echo"], records_dir=tmp_path, record_id="test-007", @@ -297,12 +304,14 @@ def test_timestamps_are_populated(self, tmp_path): data = json.loads(list(tmp_path.glob("*.json"))[0].read_text()) assert data["execution"]["timestamp_start"] assert data["execution"]["timestamp_end"] - assert data["execution"]["timestamp_start"] <= data["execution"]["timestamp_end"] + assert ( + data["execution"]["timestamp_start"] <= data["execution"]["timestamp_end"] + ) def test_auto_generates_record_id(self, tmp_path): """Omitting record_id auto-generates one.""" run_and_record( - tool_name="t", + code_name="t", command=["echo"], records_dir=tmp_path, ) @@ -316,60 +325,48 @@ def test_auto_generates_record_id(self, tmp_path): # main() CLI entry point # --------------------------------------------------------------------------- + class TestMain: @pytest.fixture(autouse=True) def _mlflow_file_store(self, tmp_path, monkeypatch): """Point MLflow tracing at a scratch file-store so init_tracing has a real backend. In production dsagt-run runs with cwd inside the - project directory, where ``dsagt_config.yaml`` (project name + - mlflow port) and ``.runtime`` (session id) live; tests mirror - that by writing both files in tmp_path and chdir-ing into it. - - Also stub the OTLPSpanExporter — init_tracing now builds one pointed - at the resolved MLflow URL, and a file:// URL would otherwise emit - async export-failure warnings to stderr. + project directory, where ``.dsagt/config.yaml`` (project name) lives + and the session id comes from ``.dsagt/state.yaml``; + tests mirror that by chdir-ing into tmp_path. """ - # dsagt_config.yaml carries the resolved MLflow URL via a synthetic - # port; tests use a file:// store, so monkeypatch the resolver to - # return the file:// URI directly. + # Serverless: init_tracing resolves a sqlite store from the project + # dir via MLflow's native provider — no OTLP exporter. Stub the + # resolver to a known sqlite URI so a shell-set MLFLOW_TRACKING_URI + # can't redirect the test. from dsagt import observability as obs_module - cfg = {"project": "test", "mlflow": {"port": 5000}} + + cfg = {"project": "test"} monkeypatch.setattr( - obs_module, "find_project_config", + obs_module, + "find_project_config", lambda: (tmp_path, cfg), ) monkeypatch.setattr( - obs_module, "_mlflow_url_from_config", - lambda c: f"file://{tmp_path}/mlruns", - ) - - class _NoopExporter: - def __init__(self, *args, **kwargs): - pass - - def export(self, spans): - from opentelemetry.sdk.trace.export import SpanExportResult - return SpanExportResult.SUCCESS - - def shutdown(self): - pass - - def force_flush(self, timeout_millis=30000): - return True - - monkeypatch.setattr( - "opentelemetry.exporter.otlp.proto.http.trace_exporter.OTLPSpanExporter", - _NoopExporter, + obs_module, + "resolve_tracking_uri", + lambda c: f"sqlite:///{tmp_path}/mlflow.db", ) def test_basic_invocation(self, tmp_path): """main() runs a command and returns its exit code.""" - exit_code = main([ - "--tool", "echo_tool", - "--records-dir", str(tmp_path), - "--", "echo", "from main", - ]) + exit_code = main( + [ + "--code", + "echo_tool", + "--records-dir", + str(tmp_path), + "--", + "echo", + "from main", + ] + ) assert exit_code == 0 records = list(tmp_path.glob("*.json")) @@ -377,20 +374,29 @@ def test_basic_invocation(self, tmp_path): def test_empty_command_returns_1(self, tmp_path): """No command after '--' returns exit code 1.""" - exit_code = main([ - "--tool", "empty", - "--records-dir", str(tmp_path), - "--", - ]) + exit_code = main( + [ + "--code", + "empty", + "--records-dir", + str(tmp_path), + "--", + ] + ) assert exit_code == 1 def test_exit_code_propagation(self, tmp_path): """main() returns the wrapped command's exit code.""" - exit_code = main([ - "--tool", "fail", - "--records-dir", str(tmp_path), - "--", "bash", "-c", "exit 7", - ]) + exit_code = main( + [ + "--code", + "fail", + "--records-dir", + str(tmp_path), + "--", + "bash", + "-c", + "exit 7", + ] + ) assert exit_code == 7 - - diff --git a/tests/test_server_startup.py b/tests/test_server_startup.py index fdf5fd4..243d78d 100644 --- a/tests/test_server_startup.py +++ b/tests/test_server_startup.py @@ -1,185 +1,72 @@ """ -Process-level MCP server startup tests. +Process-level ``dsagt-server`` entry-point tests. -Each server is spawned as a subprocess and must complete the MCP handshake -without an API key or network access. End-to-end ingest/search flows -belong in smoke_test (which drives them through the real agent), not here. +The single merged server (``dsagt.mcp.server:main``) is spawned as a subprocess +to verify the entry point is wired and fails fast + clearly on a misconfigured +project — without needing a live MLflow backend or network access. + +The full boot (init_tracing → shared KB → 20-tool MCP handshake) requires a +running MLflow server, so it is exercised by ``dsagt smoke-test`` (real agent), +not here. The 20-tool composition + dispatch contract is unit-tested in-process +by ``test_dsagt_server.py``; ``_build_kb_from_config``'s credential validation by +``test_dsagt_server.py::TestBuildKbFromConfig``. """ -import os -import shutil import sys -from pathlib import Path - -import pytest - -from mcp_helpers import mcp_initialize, mcp_list_tools, start_server - - -# --------------------------------------------------------------------------- -# Skip conditions -# --------------------------------------------------------------------------- - -_uv_available = shutil.which("uv") is not None - -pytestmark = pytest.mark.skipif(not _uv_available, reason="uv not available") - - -def _write_minimal_config( - project_dir: Path, - project_name: str = "test", - backend: str = "api", -) -> None: - """Write the minimum dsagt_config.yaml the servers need to start. - - Default backend is ``"api"`` here (with fake credentials) so startup - doesn't trigger a sentence-transformers model load — keeps these - tests fast. Tests that need to exercise the local backend pass - ``backend="local"`` explicitly. - """ - import yaml - config = { - "project": project_name, - "agent": "claude", - "embedding": { - "backend": backend, - "model": "test-model", - "base_url": "http://localhost:9999", - "api_key": "test-fake-key", - }, - "knowledge": { - "chunk_size": 1024, - "vector_db": "chroma", - "rerank": False, - }, - } - (project_dir / "dsagt_config.yaml").write_text( - yaml.dump(config, default_flow_style=False) - ) - - -# --------------------------------------------------------------------------- -# Startup tests (no API key needed) -# --------------------------------------------------------------------------- - -class TestRegistryServerStartup: - - def test_starts_and_lists_tools(self, tmp_path): - """Registry server starts without embedding credentials (kb=None) - and completes the MCP handshake.""" - import yaml as _yaml - project = tmp_path / "runtime" - project.mkdir() - # backend="api" with empty credentials → kb_available=False → - # registry runs with kb=None. This is the fast path the test - # cares about (no ChromaDB / sentence-transformers init). With - # backend="local" (the default for new projects) registry would - # eagerly load the local embedder, which is the slow path - # exercised by other tests. - (project / "dsagt_config.yaml").write_text(_yaml.dump({ - "project": "test", - "agent": "claude", - "embedding": { - "backend": "api", - "model": "", "base_url": "", "api_key": "", - }, - "knowledge": {"chunk_size": 1024, "vector_db": "chroma", "rerank": False}, - })) - proc = start_server( - [sys.executable, "-m", "dsagt.commands.registry_server"], - env={ - "DSAGT_PROJECT_DIR": str(project), - # init_tracing raises without a backend; give it a local - # file store so the handshake proceeds without needing an - # MLflow server. - "MLFLOW_TRACKING_URI": f"file://{tmp_path}/mlruns", - }, - ) - try: - resp = mcp_initialize(proc) - assert "result" in resp, f"Init failed: {resp}" - tools_resp = mcp_list_tools(proc) - assert "result" in tools_resp, f"list_tools failed: {tools_resp}" +from mcp_helpers import start_server - tool_names = [t["name"] for t in tools_resp["result"]["tools"]] - assert "save_tool_spec" in tool_names - assert "install_dependencies" in tool_names - finally: - proc.terminate() - proc.wait(timeout=5) +_SERVER_CMD = [sys.executable, "-m", "dsagt.mcp.server"] -class TestKnowledgeServerStartup: +class TestServerEntryPoint: - def test_starts_and_lists_tools(self, tmp_path): - """Knowledge server starts, completes MCP handshake, and lists tools. + def test_fails_fast_without_project_config(self, tmp_path): + """Run from a dir with no .dsagt/config.yaml → clean fail-fast. - Uses a fake API key — server accepts it at init time since it only - validates the key is non-empty. Actual API calls would fail, but we - don't make any here. + ``dsagt-server`` discovers its project from cwd; launched anywhere else + it must say so rather than boot a half-configured server. """ - project = tmp_path / "runtime" - project.mkdir() - _write_minimal_config(project) - proc = start_server( - [sys.executable, "-m", "dsagt.commands.knowledge_server"], - env={ - "DSAGT_PROJECT_DIR": str(project), - "LLM_API_KEY": "test-fake-key-for-startup", - "MLFLOW_TRACKING_URI": f"file://{tmp_path}/mlruns", - }, - ) - try: - resp = mcp_initialize(proc) - assert "result" in resp, f"Init failed: {resp}" - - tools_resp = mcp_list_tools(proc) - assert "result" in tools_resp, f"list_tools failed: {tools_resp}" - - tool_names = [t["name"] for t in tools_resp["result"]["tools"]] - assert "kb_search" in tool_names - assert "kb_list_collections" in tool_names - assert "kb_ingest" in tool_names - finally: - proc.terminate() - proc.wait(timeout=5) - - def test_api_backend_fails_fast_without_api_key(self, tmp_path): - """When the user explicitly opts into backend='api', the knowledge - server must fail at startup if the api key is missing — better than - booting with a broken embedder that 401s on the first kb_search. - - With the post-refactor default of backend='local' (no creds - needed), this hard-fail only fires for users who have explicitly - said they want the api backend. + empty = tmp_path / "empty" + empty.mkdir() + proc = start_server(_SERVER_CMD, cwd=str(empty)) + rc = proc.wait(timeout=15) + assert rc != 0 + assert "no .dsagt/config.yaml in cwd" in proc.stderr.read() + + def test_mints_session_into_state_on_boot(self, tmp_path): + """The server owns the session lifecycle: on boot it appends a session + entry to ``.dsagt/state.yaml`` (serverless — no MLflow backend needed). + + Session minting happens before the (slow) KB build, so the state file + appears within a second; we poll for it, then terminate. """ + import time + import yaml + project = tmp_path / "runtime" - project.mkdir() + (project / ".dsagt").mkdir(parents=True) config = { "project": "test", "agent": "claude", - "embedding": { - "backend": "api", - "model": "test-model", - "base_url": "http://localhost:9999", - "api_key": "${LLM_API_KEY}", # unresolved placeholder - }, - "knowledge": {"chunk_size": 1024, "vector_db": "chroma", "rerank": False}, + "embedding": {"backend": "local", "model": "BAAI/bge-small-en-v1.5"}, + "knowledge": {"chunk_size": 1024, "rerank": False}, } - (project / "dsagt_config.yaml").write_text( + (project / ".dsagt" / "config.yaml").write_text( yaml.dump(config, default_flow_style=False) ) - - stripped = {"LLM_API_KEY", "OPENAI_API_KEY"} - clean_env = {k: v for k, v in os.environ.items() if k not in stripped} - clean_env["DSAGT_PROJECT_DIR"] = str(project) - clean_env["MLFLOW_TRACKING_URI"] = f"file://{tmp_path}/mlruns" - - proc = start_server( - [sys.executable, "-m", "dsagt.commands.knowledge_server"], - env=clean_env, - ) - rc = proc.wait(timeout=10) - assert rc != 0, "backend='api' without api_key must fail at startup" + state_file = project / ".dsagt" / "state.yaml" + proc = start_server(_SERVER_CMD, cwd=str(project)) + try: + for _ in range(60): # up to ~12s + if state_file.exists(): + break + time.sleep(0.2) + assert state_file.exists(), "server did not mint a session into state.yaml" + state = yaml.safe_load(state_file.read_text()) + assert state["sessions"][-1]["id"] == 1 + assert state["sessions"][-1]["started_at"].endswith("Z") + finally: + proc.terminate() + proc.wait(timeout=10) diff --git a/tests/test_setup_core_kb.py b/tests/test_setup_core_kb.py index ff5472d..0e008dd 100644 --- a/tests/test_setup_core_kb.py +++ b/tests/test_setup_core_kb.py @@ -1,10 +1,10 @@ """ -Tests for dsagt-setup-kb internals. +Tests for the knowledge-base asset builder (dsagt.commands.setup_core_kb), +the engine behind ``dsagt init``'s KB provisioning. -These cover the helpers in dsagt.commands.setup_core_kb that don't require -network access — the git clone subprocess is mocked so the tests can run -offline. The actual end-to-end behavior of dsagt-setup-kb against real -upstream repos is exercised manually. +These cover the helpers that don't require network access — the git clone +subprocess is mocked so the tests can run offline. The actual end-to-end +behavior against real upstream repos is exercised manually. """ from __future__ import annotations @@ -16,12 +16,18 @@ import pytest +import numpy as np + from dsagt.commands.setup_core_kb import ( + DEFAULT_ASSETS, DEFAULT_EXCLUDE_PATTERNS, + all_assets, + asset_collection_name, clone_github, + ensure_assets, + resolve_assets, ) - # --------------------------------------------------------------------------- # clone_github top-level-files behavior # --------------------------------------------------------------------------- @@ -135,3 +141,80 @@ def test_default_exclude_patterns_keeps_packaging_metadata(): f"DEFAULT_EXCLUDE_PATTERNS contains packaging metadata files " f"that the agent needs: {overlap}" ) + + +# --------------------------------------------------------------------------- +# Installable-asset selection (--include / --exclude namespace) +# --------------------------------------------------------------------------- + + +class TestResolveAssets: + + def test_default_is_tools_plus_genesis(self): + assert resolve_assets() == list(DEFAULT_ASSETS) == ["codes", "genesis"] + + def test_include_all_is_everything(self): + assert resolve_assets(include=["all"]) == all_assets() + + def test_include_subset_returns_canonical_order(self): + # input order shouldn't matter — cheap assets always built first. + assert resolve_assets(include=["aidrin", "codes"]) == ["codes", "aidrin"] + + def test_exclude_trims_the_default_set(self): + assert resolve_assets(exclude=["genesis"]) == ["codes"] + + def test_exclude_all_is_empty(self): + assert resolve_assets(exclude=["all"]) == [] + + def test_unknown_asset_raises(self): + with pytest.raises(ValueError, match="unknown KB asset"): + resolve_assets(include=["not-a-real-asset"]) + + def test_include_exclude_mutually_exclusive(self): + with pytest.raises(ValueError, match="mutually exclusive"): + resolve_assets(include=["codes"], exclude=["genesis"]) + + +class TestAssetCollectionName: + + def test_tools(self): + assert asset_collection_name("codes") == "codes" + + def test_catalog_uses_catalog_prefix(self): + name = asset_collection_name("genesis") + assert name.startswith("skills_catalog__") and "genesis" in name + + def test_scientific_collection_is_its_own_name(self): + assert asset_collection_name("nemo_curator") == "nemo_curator" + + def test_unknown_raises(self): + with pytest.raises(ValueError, match="unknown KB asset"): + asset_collection_name("bogus") + + +class TestEnsureAssetsTools: + """``ensure_assets`` for the bundled-tools asset, with a mocked embedder + so the test stays offline (no model download, no git clone).""" + + def _fake_embedder(self): + emb = MagicMock() + emb.embed = lambda texts: np.full((len(texts), 8), 0.1, dtype=np.float32) + return emb + + def test_builds_tools_collection(self, tmp_path): + with patch( + "dsagt.knowledge.Embedder.create", return_value=self._fake_embedder() + ): + result = ensure_assets(["codes"], tmp_path) + assert "codes" in result["built"] + # ChromaIndex.save writes chroma_ids.json — the collection marker. + assert (tmp_path / "codes" / "chroma_ids.json").exists() + + def test_is_idempotent(self, tmp_path): + with patch( + "dsagt.knowledge.Embedder.create", return_value=self._fake_embedder() + ): + ensure_assets(["codes"], tmp_path) + second = ensure_assets(["codes"], tmp_path) + assert second["skipped"] == ["codes"] + assert second["built"] == [] diff --git a/tests/test_site_config.yaml b/tests/test_site_config.yaml new file mode 100755 index 0000000..b53f049 --- /dev/null +++ b/tests/test_site_config.yaml @@ -0,0 +1,26 @@ +# Test site configuration — institution-specific values for integration tests. +# +# Copy this file to tests/test_site_config.yaml and fill in your values. +# Unit tests (mocked) run without this file. Integration tests that hit +# a real embedding API will skip if this file is missing. +# +# DO NOT commit test_site_config.yaml — it contains API endpoints and +# model names specific to your institution. + +project: test-integration +agent: claude + +llm: + model: claude-sonnet-4-5-20250929-v1-project # LLM model name at your endpoint + base_url: https://api.example.com # LLM API endpoint (Anthropic-compatible) + api_key: sk-QKnGITknkwJbBjhaPcj4QQ # API key for LLM service + +embedding: + model: text-embedding-3-small-project # embedding model at your endpoint + base_url: https://ai-incubator-api.pnnl.gov # OpenAI-compatible embeddings URL + api_key: sk-QKnGITknkwJbBjhaPcj4QQ # API key for embedding service + +mlflow: + port: 15001 # test MLflow port + backend: sqlite + diff --git a/tests/test_site_config.yaml.example b/tests/test_site_config.yaml.example index eb47c3c..58f3844 100755 --- a/tests/test_site_config.yaml.example +++ b/tests/test_site_config.yaml.example @@ -10,7 +10,7 @@ project: test-integration agent: claude llm: - model: claude-sonnet-4-20250514 # LLM model the proxy routes to + model: claude-sonnet-4-20250514 # LLM model name at your endpoint base_url: https://api.example.com # LLM API endpoint (Anthropic-compatible) api_key: your-llm-api-key # API key for LLM service @@ -19,9 +19,6 @@ embedding: base_url: https://api.example.com/v1 # OpenAI-compatible embeddings URL api_key: your-embedding-api-key # API key for embedding service -proxy: - port: 14000 # test proxy port (offset to avoid conflicts) - mlflow: port: 15001 # test MLflow port backend: sqlite diff --git a/tests/test_skill_discovery.py b/tests/test_skill_discovery.py new file mode 100644 index 0000000..b955d07 --- /dev/null +++ b/tests/test_skill_discovery.py @@ -0,0 +1,236 @@ +"""Unit tests for the keyword scorer and SkillRouter (no network, no embedder). + +The router is exercised in its KB-free keyword mode against a real +``SkillRegistry`` (bundled skills suppressed via an empty source dir) plus a +fake catalog cache, and in KB mode against a small fake KnowledgeBase. +""" + +import json + +from dsagt.registry import SkillRegistry +from dsagt.skills import SkillRouter, rank_skills, score_skill + + +def _mkskill(d, name, desc): + d.mkdir(parents=True, exist_ok=True) + (d / "SKILL.md").write_text( + f"---\nname: {name}\ndescription: {desc}\n---\n# {name}\nbody\n" + ) + return d + + +def _registry(tmp_path, skills=None): + """SkillRegistry with bundled skills suppressed and optional project skills.""" + empty_bundled = tmp_path / "no_bundled" + empty_bundled.mkdir() + reg = SkillRegistry( + runtime_dir=tmp_path / "proj", + source_skills_dir=str(empty_bundled), + kb=None, + ) + for name, desc in (skills or {}).items(): + _mkskill(reg.skills_dir / name, name, desc) + return reg + + +class FakeKB: + """Minimal duck-typed KnowledgeBase: collections + search + index_dir.""" + + def __init__(self, collections, hits_by_collection=None, index_dir="/tmp/none"): + self.collections = list(collections) + self._hits = hits_by_collection or {} + self.index_dir = index_dir + + def search(self, query, collection=None, collections=None, top_k=5, **kwargs): + # Mirror KnowledgeBase.search's fan-out surface; a simplified stand-in + # that merges by raw score (real cross-collection RRF is exercised in + # test_knowledge_base.TestFederatedSearch). + targets = collections or ([collection] if collection else []) + hits: list[dict] = [] + for c in targets: + hits.extend(self._hits.get(c, [])) + hits.sort(key=lambda h: h["score"], reverse=True) + return hits[:top_k] + + +def _hit(name, text, source, score, tags=""): + return { + "chunk": { + "metadata": {"skill_name": name, "source": source, "tags": tags}, + "text": text, + }, + "score": score, + } + + +# --------------------------------------------------------------------------- +# keyword scorer +# --------------------------------------------------------------------------- + + +def test_score_name_token_weighs_more_than_description(): + name_hit = score_skill("slurm", "slurm-submit", "unrelated text") + desc_hit = score_skill("slurm", "other", "submit a slurm job") + assert name_hit > desc_hit + + +def test_score_exact_name_beats_substring(): + exact = score_skill("datacard", "datacard", "x") + substr = score_skill("datacard", "datacard-generator", "x") + assert exact > substr > 0 + + +def test_score_substring_description_bonus(): + assert score_skill("batch job", "x", "submit a batch job") > 0 + + +def test_stopwords_do_not_score(): + # only stopwords overlap → no score + assert score_skill("the and of", "the skill", "and of the") == 0.0 + + +def test_empty_query_scores_zero(): + assert score_skill("", "anything", "anything") == 0.0 + + +def test_substring_bonuses_are_mutually_exclusive(): + # Genesis parity: at most ONE of the +6/+4/+2 substring bonuses fires. + # name-token 2 + desc-token 1 + exact-name 6 = 9; the desc-substring +2 is + # NOT also added (a stacking bug would give 11). + assert score_skill("alpha", "alpha", "alpha tool") == 9.0 + + +def test_single_char_tokens_dropped(): + # "x" is a single char → not a token, so no name-token overlap. + assert score_skill("x", "x", "y") == 6.0 # only the exact-name bonus + + +def test_rank_orders_and_breaks_ties_by_name(): + skills = [ + {"name": "zeta", "description": "submit jobs"}, + {"name": "alpha", "description": "submit jobs"}, + {"name": "unrelated", "description": "nothing here"}, + ] + ranked = rank_skills("submit", skills, top_k=5) + names = [s["name"] for s, _ in ranked] + assert names == ["alpha", "zeta"] # equal score → name asc; unrelated dropped + + +# --------------------------------------------------------------------------- +# keyword-mode search (kb is None) +# --------------------------------------------------------------------------- + + +def test_search_keyword_is_catalog_only(tmp_path): + # Installed skills are NOT search candidates (they're natively discovered); + # only the cached catalog is keyword-scored. + reg = _registry(tmp_path, {"slurm-submit": "submit a batch job to slurm"}) + cache = tmp_path / "cache" + _mkskill( + cache / "genesis" / "slurm-catalog", + "slurm-catalog", + "submit a batch job to slurm", + ) + r = SkillRouter(skill_registry=reg, cache_dir=cache) + out = r.search("slurm batch") + assert "slurm-catalog" in out # catalog skill found + assert "slurm-submit" not in out # installed skill NOT surfaced by search + assert "[catalog · install_skill to add]" in out + + +def test_search_keyword_includes_catalog_cache(tmp_path): + reg = _registry(tmp_path) + cache = tmp_path / "cache" + _mkskill( + cache / "genesis-skills" / "croissant", + "croissant-validator", + "validate a croissant metadata file", + ) + r = SkillRouter(skill_registry=reg, cache_dir=cache) + out = r.search("croissant") + assert "croissant-validator" in out + assert "[catalog · install_skill to add]" in out + + +def test_search_is_stateless(tmp_path): + # No recency queue: repeating a query yields the same result, no suppression. + reg = _registry(tmp_path) + cache = tmp_path / "cache" + _mkskill(cache / "src" / "slurm-x", "slurm-x", "submit a batch job to slurm") + r = SkillRouter(skill_registry=reg, cache_dir=cache) + first = r.search("slurm") + second = r.search("slurm") + assert first == second + assert "Found 1 skill" in second + + +def test_search_exact_name_is_kb_free(tmp_path): + reg = _registry(tmp_path, {"datacard-gen": "make a dataset card"}) + r = SkillRouter(skill_registry=reg) + out = r.search(skill_name="datacard-gen") + assert "datacard-gen" in out + assert r.search(skill_name="nope").startswith("No skill named") + + +# --------------------------------------------------------------------------- +# KB-mode search +# --------------------------------------------------------------------------- + + +def test_search_kb_merges_catalog_collections(tmp_path): + # Only skills_catalog__* collections are searched; the installed 'skills' + # collection is ignored even if present. + kb = FakeKB( + collections=["skills", "skills_catalog__a", "skills_catalog__b"], + hits_by_collection={ + "skills": [_hit("installed-one", "installed", "registered", 0.99)], + "skills_catalog__a": [_hit("cat-a", "catalog a", "catalog:a", 0.9)], + "skills_catalog__b": [_hit("cat-b", "catalog b", "catalog:b", 0.5)], + }, + ) + r = SkillRouter(skill_registry=_registry(tmp_path), kb=kb) + out = r.search("anything") + assert "installed-one" not in out # installed collection not searched + assert "cat-a" in out and "cat-b" in out + assert out.index("cat-a") < out.index("cat-b") # higher score first + + +def test_search_kb_tag_filter(tmp_path): + kb = FakeKB( + collections=["skills_catalog__x"], + hits_by_collection={ + "skills_catalog__x": [ + _hit("tagged", "x", "catalog:x", 0.9, tags="hpc,slurm"), + _hit("untagged", "y", "catalog:x", 0.8, tags=""), + ] + }, + ) + r = SkillRouter(skill_registry=_registry(tmp_path), kb=kb) + out = r.search("x", tag="slurm") + assert "tagged" in out and "untagged" not in out + + +# --------------------------------------------------------------------------- +# list_sources +# --------------------------------------------------------------------------- + + +def test_list_sources_flags_synced(tmp_path): + from dsagt.skills import KNOWN_SOURCES, _repo_slug + from dsagt.registry import catalog_collection + + genesis_coll = catalog_collection(_repo_slug(KNOWN_SOURCES["genesis"]["url"])) + index_dir = tmp_path / "idx" + (index_dir / genesis_coll).mkdir(parents=True) + (index_dir / genesis_coll / "chroma_ids.json").write_text( + json.dumps(["1", "2", "3"]) + ) + + kb = FakeKB(collections=[genesis_coll], index_dir=str(index_dir)) + # list_sources needs only a KB — no skill_registry required. + r = SkillRouter(kb=kb) + sources = {s["name"]: s for s in r.list_sources()} + assert sources["genesis"]["synced"] is True + assert sources["genesis"]["indexed"] == 3 + assert sources["anthropic"]["synced"] is False + assert sources["anthropic"]["indexed"] == 0 diff --git a/tests/test_skill_tools.py b/tests/test_skill_tools.py new file mode 100644 index 0000000..9bc0e29 --- /dev/null +++ b/tests/test_skill_tools.py @@ -0,0 +1,249 @@ +""" +Tests for the skill MCP tools (save_skill, search_skills, install_skill, +add_skill_source, list_skill_sources). + +The skill surface lives in :mod:`dsagt.mcp.skill_tools`; ``create_skill_server`` +exposes just that concern for driving via the MCP helpers. Handlers return a +mix of ``str`` (save/search/install) and ``dict`` (add/list sources), so the two +``call_tool`` helpers are both used. +""" + +import json +from unittest.mock import MagicMock + +import pytest + +from dsagt.mcp.skill_tools import create_skill_server +from dsagt.registry import SkillRegistry +from mcp_helpers import call_tool_json, call_tool_sync + + +def _make_skill_server(tmp_path): + """Create (server, skill_registry, kb) with a real local-embedding KB. + + The skill registry is rooted at ``/runtime`` so save_skill writes to + ``/runtime/skills//`` — the project layer the agent natively + discovers. + """ + from dsagt.knowledge import KnowledgeBase + + runtime_dir = tmp_path / "runtime" + runtime_dir.mkdir(parents=True, exist_ok=True) + kb = KnowledgeBase( + index_dir=tmp_path / "kb_index", + default_embedder="local", + ) + skill_reg = SkillRegistry( + source_skills_dir=None, # package default (empty bundled is fine) + runtime_dir=str(runtime_dir), + kb=kb, + ) + server = create_skill_server(skill_reg, kb, runtime_dir=str(runtime_dir)) + return server, skill_reg, kb + + +# --------------------------------------------------------------------------- +# save_skill +# --------------------------------------------------------------------------- + + +class TestSaveSkill: + + def test_add_new_skill_creates_files_and_indexes(self, tmp_path): + """save_skill writes SKILL.md and the skill count goes up by one. + + The count includes any bundled skills that ship in the package (see + SkillRegistry.list_skills which merges bundled + project layers), so we + assert the file was created and the count incremented rather than + equality on a specific number. + """ + server, skill_reg, kb = _make_skill_server(tmp_path) + before = len(skill_reg.list_skills()) + + spec = { + "name": "csv_inspector", + "description": "Workflow for inspecting CSV columns and quality", + "tags": ["data_management", "quality_control"], + } + body = "# csv_inspector\n\nFirst, run head on the file. Then check nulls.\n" + text = call_tool_sync(server, "save_skill", {"spec": spec, "body": body}) + + assert "added" in text + skill_md = tmp_path / "runtime" / "skills" / "csv_inspector" / "SKILL.md" + assert skill_md.exists() + content = skill_md.read_text() + assert "csv_inspector" in content + assert "First, run head" in content + after = len(skill_reg.list_skills()) + assert after == before + 1 + + def test_update_existing_skill_preserves_body_when_omitted(self, tmp_path): + """Saving a spec for an existing skill without body keeps the body.""" + server, skill_reg, kb = _make_skill_server(tmp_path) + first_body = "# orig\n\nOriginal workflow body.\n" + call_tool_sync( + server, + "save_skill", + { + "spec": {"name": "wf", "description": "v1"}, + "body": first_body, + }, + ) + # Update the description only — body should be preserved. + text = call_tool_sync( + server, + "save_skill", + { + "spec": {"name": "wf", "description": "v2 description"}, + }, + ) + assert "updated" in text + skill_md = tmp_path / "runtime" / "skills" / "wf" / "SKILL.md" + content = skill_md.read_text() + assert "v2 description" in content + assert "Original workflow body" in content + + def test_save_skill_writes_reference_files(self, tmp_path): + """reference_files dict lands as additional files in the skill dir.""" + server, skill_reg, kb = _make_skill_server(tmp_path) + text = call_tool_sync( + server, + "save_skill", + { + "spec": {"name": "with_template", "description": "Has a template"}, + "body": "# with_template\n\nUses template.json.\n", + "reference_files": {"template.json": '{"foo": "bar"}\n'}, + }, + ) + assert "added" in text + skill_dir = tmp_path / "runtime" / "skills" / "with_template" + assert (skill_dir / "SKILL.md").exists() + assert (skill_dir / "template.json").read_text() == '{"foo": "bar"}\n' + + def test_save_skill_string_encoded_spec(self, tmp_path): + """MCP clients that JSON-encode nested object args still work.""" + server, skill_reg, kb = _make_skill_server(tmp_path) + spec_json = json.dumps({"name": "s1", "description": "d"}) + text = call_tool_sync(server, "save_skill", {"spec": spec_json, "body": "x"}) + assert "added" in text + + +# --------------------------------------------------------------------------- +# search_skills +# --------------------------------------------------------------------------- + + +class TestSearchSkills: + + def test_search_skills_empty_catalog_hints_to_sync(self, tmp_path): + """With no catalog synced, search_skills explains how to enable one + instead of returning a bare 'no match' the agent reads as exhausted.""" + server, skill_reg, kb = _make_skill_server(tmp_path) + + text = call_tool_sync(server, "search_skills", {"query": "vasp pymatgen dft"}) + assert "No catalog skills found" in text + assert "no external skill catalog is synced" in text.lower() + assert "add_skill_source" in text + + +class TestNativeMirror: + + def test_save_skill_mirrors_into_native_skills_dir(self, tmp_path): + """A skill created mid-session mirrors into the agent's native skills + dir immediately (not at the next ``dsagt start``), so the next session + auto-discovers it however the agent is launched. install_skill and + save_code_spec run the same ``refresh_native_skills`` call.""" + server, skill_reg, kb = _make_skill_server(tmp_path) + runtime = tmp_path / "runtime" + (runtime / ".dsagt").mkdir() + (runtime / ".dsagt" / "config.yaml").write_text("project: p\nagent: claude\n") + + call_tool_sync( + server, + "save_skill", + { + "spec": {"name": "mirror-me", "description": "mirrors right away"}, + "body": "# mirror-me\n\nDo the thing.\n", + }, + ) + + native = runtime / ".claude" / "skills" + assert (native / "mirror-me" / "SKILL.md").exists() + manifest = json.loads((native / ".dsagt-managed.json").read_text()) + assert "mirror-me" in manifest + + def test_save_skill_without_project_config_skips_mirror(self, tmp_path): + """No ``.dsagt/config.yaml`` (not dsagt init-ed) → the skill still + saves; no native dir appears.""" + server, skill_reg, kb = _make_skill_server(tmp_path) + text = call_tool_sync( + server, + "save_skill", + {"spec": {"name": "bare", "description": "no config"}, "body": "# b\n"}, + ) + assert "added" in text + assert not (tmp_path / "runtime" / ".claude").exists() + + +# --------------------------------------------------------------------------- +# install_skill +# --------------------------------------------------------------------------- + + +class TestInstallSkill: + + def test_install_skill_routes_and_reports_missing(self, tmp_path, monkeypatch): + """install_skill is registered and reports a clean error when the + named skill isn't in any synced catalog. + + The handler builds its own SkillRouter over the machine-global clone + cache, so point that at an empty tmp dir — the test must not depend on + (or scan) whatever catalogs are synced on the developer's machine. No + KB is involved on the install path, so a mock keeps the test fast. + """ + monkeypatch.setattr( + "dsagt.skills.SKILL_SOURCES_DIR", tmp_path / "skill_sources" + ) + server = create_skill_server( + kb=MagicMock(), runtime_dir=str(tmp_path / "runtime") + ) + text = call_tool_sync( + server, + "install_skill", + {"skill_name": "zzz-definitely-not-a-real-skill-xyz"}, + ) + assert "No catalog skill" in text + + +# --------------------------------------------------------------------------- +# skill sources (add_skill_source / list_skill_sources) +# --------------------------------------------------------------------------- + + +@pytest.fixture +def mock_kb(tmp_path): + kb = MagicMock() + kb.index_dir = tmp_path / "kb_index" + kb.index_dir.mkdir() + kb.collections = [] + return kb + + +class TestSkillSources: + + def test_list_skill_sources_returns_known(self, mock_kb): + server = create_skill_server(kb=mock_kb) + result = call_tool_json(server, "list_skill_sources", {}) + assert "k-dense-ai" in result["sources"] + # Nothing synced → every known source flagged available, not synced. + assert result["sources"]["k-dense-ai"]["synced"] is False + assert result["sources"]["k-dense-ai"]["indexed"] == 0 + assert result["other_synced_collections"] == [] + assert "k-dense-ai" in result["note"] + + def test_add_skill_source_bad_source_errors(self, mock_kb): + server = create_skill_server(kb=mock_kb) + result = call_tool_json( + server, "add_skill_source", {"source": "not-a-real-known-name"} + ) + assert "error" in result diff --git a/tests/test_skills_catalog.py b/tests/test_skills_catalog.py new file mode 100644 index 0000000..1a3c097 --- /dev/null +++ b/tests/test_skills_catalog.py @@ -0,0 +1,423 @@ +"""Unit tests for the external skill catalog (fetch / index / install) and +the native-skill mirror. No network: ``clone_github`` is monkeypatched and +the KB is a lightweight fake that records ``add_entries`` calls.""" + +import json + +import pytest + +from dsagt.agents.base import ( + _NATIVE_DESCRIPTION_CAP, + _SKILL_MANIFEST, + _mirror_skills_to, +) +from dsagt import skills as sc +from dsagt.registry import CATALOG_COLLECTION_PREFIX, catalog_collection + + +def _mkskill(d, name, desc="a short description"): + d.mkdir(parents=True, exist_ok=True) + (d / "SKILL.md").write_text( + f"---\nname: {name}\ndescription: {desc}\n---\n# {name}\nbody\n" + ) + return d + + +# --------------------------------------------------------------------------- +# slug + source resolution +# --------------------------------------------------------------------------- + + +def test_repo_slug_is_collection_safe(): + slug = sc._repo_slug("https://github.com/K-Dense-AI/scientific-agent-skills") + assert slug == "k-dense-ai-scientific-agent-skills" + assert sc._repo_slug("git@github.com:Foo/Bar.git") == "foo-bar" + + +def test_repo_slug_is_host_agnostic(): + # Non-GitHub hosts (GitLab, etc.) reduce to owner-repo, scheme/host dropped. + assert sc._repo_slug("https://gitlab.osti.gov/genesis/genesis-skills") == ( + "genesis-genesis-skills" + ) + assert sc._repo_slug("git@gitlab.osti.gov:genesis/genesis-skills.git") == ( + "genesis-genesis-skills" + ) + + +def test_known_source_genesis_covers_whole_skills_tree(): + spec = sc.resolve_source("genesis") + assert spec["url"] == "https://gitlab.osti.gov/genesis/genesis-skills" + # subdir scopes the recursive SKILL.md walk to the whole skills/ tree so + # every category (hpc, huggingface, langchain, …) is discoverable. + assert spec["subdir"] == "skills" + assert spec["branch"] == "main" + + +def test_persist_source_to_config_appends_and_dedupes(tmp_path): + import yaml + + (tmp_path / ".dsagt").mkdir() + cfg = tmp_path / ".dsagt" / "config.yaml" + cfg.write_text(yaml.dump({"project": "p", "skills": {"sources": []}})) + spec = { + "name": "anthropic", + "url": "https://github.com/anthropics/skills", + "branch": "main", + } + assert sc.persist_source_to_config(tmp_path, spec) is True + sources = yaml.safe_load(cfg.read_text())["skills"]["sources"] + assert sources[-1]["name"] == "anthropic" + # Idempotent: same URL is not appended twice. + assert sc.persist_source_to_config(tmp_path, spec) is False + assert len(yaml.safe_load(cfg.read_text())["skills"]["sources"]) == 1 + # No config file → no-op, no crash. + assert sc.persist_source_to_config(tmp_path / "nope", spec) is False + + +def test_resolve_source_known_url_and_shorthand(): + assert ( + sc.resolve_source("k-dense-ai")["url"] == sc.KNOWN_SOURCES["k-dense-ai"]["url"] + ) + assert ( + sc.resolve_source("https://github.com/a/b")["url"] == "https://github.com/a/b" + ) + assert sc.resolve_source("a/b")["url"] == "https://github.com/a/b" + with pytest.raises(ValueError): + sc.resolve_source("not-a-known-name") + + +# --------------------------------------------------------------------------- +# discovery +# --------------------------------------------------------------------------- + + +def test_discover_skill_dirs_flat_and_nested(tmp_path): + root = tmp_path / "skills" + _mkskill(root / "flat", "flat") + _mkskill(root / "domain" / "nested", "nested") + # A dir whose SKILL.md has no name is ignored. + bad = root / "noname" + bad.mkdir(parents=True) + (bad / "SKILL.md").write_text("---\ndescription: x\n---\nbody") + names = sorted(p.name for p in sc._discover_skill_dirs(tmp_path)) + assert names == ["flat", "nested"] + + +# --------------------------------------------------------------------------- +# find + install +# --------------------------------------------------------------------------- + + +def test_find_catalog_skill_and_ambiguity(tmp_path): + cache = tmp_path / "cache" + _mkskill(cache / "srcA" / "skills" / "alpha", "alpha") + found = sc.find_catalog_skill("alpha", cache_dir=cache) + assert found.name == "alpha" + with pytest.raises(LookupError): + sc.find_catalog_skill("missing", cache_dir=cache) + # Same skill name in a second source → ambiguous. + _mkskill(cache / "srcB" / "skills" / "alpha", "alpha") + with pytest.raises(LookupError, match="multiple sources"): + sc.find_catalog_skill("alpha", cache_dir=cache) + + +def test_find_catalog_skill_source_qualified(tmp_path): + cache = tmp_path / "cache" + _mkskill(cache / "srcA" / "skills" / "alpha", "alpha") + _mkskill(cache / "srcB" / "skills" / "alpha", "alpha") + + # A "/" qualifier disambiguates which source to install from. + a = sc.find_catalog_skill("srcA/alpha", cache_dir=cache) + b = sc.find_catalog_skill("srcB/alpha", cache_dir=cache) + assert a.relative_to(cache).parts[0] == "srcA" + assert b.relative_to(cache).parts[0] == "srcB" + assert a.name == b.name == "alpha" + + # Qualifying with a source that lacks the skill is a clear, source-scoped miss. + with pytest.raises(LookupError, match="in source 'srcA'"): + sc.find_catalog_skill("srcA/missing", cache_dir=cache) + + +def test_install_into_project_source_qualified(tmp_path): + cache = tmp_path / "cache" + _mkskill(cache / "srcA" / "skills" / "dup", "dup", desc="from A") + _mkskill(cache / "srcB" / "skills" / "dup", "dup", desc="from B") + proj = tmp_path / "proj" + proj.mkdir() + + # Bare ambiguous name refuses; the source-qualified form installs srcB's copy. + with pytest.raises(LookupError, match="multiple sources"): + sc.install_into_project("dup", proj, cache_dir=cache) + info = sc.install_into_project("srcB/dup", proj, cache_dir=cache) + assert info["name"] == "dup" + assert (proj / "skills" / "dup" / "SKILL.md").read_text().count("from B") == 1 + + +def test_install_into_project_copies_subdirs(tmp_path): + cache = tmp_path / "cache" + skill = _mkskill(cache / "src" / "vasp-to-isaac", "vasp-to-isaac") + (skill / "scripts").mkdir() + (skill / "scripts" / "run.py").write_text("print(1)") + (skill / "references").mkdir() + (skill / "references" / "spec.md").write_text("# spec") + + proj = tmp_path / "proj" + proj.mkdir() + info = sc.install_into_project("vasp-to-isaac", proj, cache_dir=cache) + dest = proj / "skills" / "vasp-to-isaac" + assert info["action"] == "added" + assert (dest / "SKILL.md").exists() + assert (dest / "scripts" / "run.py").exists() + assert (dest / "references" / "spec.md").exists() + # Re-install reports "updated". + assert ( + sc.install_into_project("vasp-to-isaac", proj, cache_dir=cache)["action"] + == "updated" + ) + + +# --------------------------------------------------------------------------- +# sync_source (mocked clone + fake KB) +# --------------------------------------------------------------------------- + + +class _FakeKB: + def __init__(self, index_dir): + self.index_dir = index_dir + self.collections = [] + self.adds = [] # (collection, metadatas) + + def add_entries(self, texts, collection, metadatas=None): + self.adds.append((collection, metadatas)) + if collection not in self.collections: + self.collections.append(collection) + return {"collection": collection, "entries_added": len(texts)} + + +def test_sync_source_indexes_per_source_collection(tmp_path, monkeypatch): + # Fake clone: populate dest/ with two skills. + def fake_clone(url, dest, branch="main", include=None): + sub = include[0] if include else "" + base = dest / sub if sub else dest + _mkskill(base / "s1", "s1") + _mkskill(base / "s2", "s2") + + monkeypatch.setattr("dsagt.commands.setup_core_kb.clone_github", fake_clone) + + kb = _FakeKB(tmp_path / "kb_index") + cache = tmp_path / "cache" + stats = sc.sync_source( + {"url": "https://github.com/x/y", "branch": "main", "subdir": "skills"}, + kb=kb, + cache_dir=cache, + ) + slug = sc._repo_slug("https://github.com/x/y") + coll = catalog_collection(slug) + assert stats["discovered"] == 2 and stats["indexed"] == 2 + assert coll.startswith(CATALOG_COLLECTION_PREFIX) + added_coll, metas = kb.adds[-1] + assert added_coll == coll + assert all(m["source"] == f"catalog:{slug}" for m in metas) + assert {m["skill_name"] for m in metas} == {"s1", "s2"} + + +# --------------------------------------------------------------------------- +# native mirror +# --------------------------------------------------------------------------- + + +def test_mirror_manifest_preserves_user_skills_and_reaps(tmp_path): + target = tmp_path / ".claude" / "skills" + target.mkdir(parents=True) + # A user-authored skill dsagt must never touch. + _mkskill(target / "user-skill", "user-skill") + + bundled = _mkskill(tmp_path / "bundled" / "skill-creator", "skill-creator") + proj = _mkskill(tmp_path / "proj" / "alpha", "alpha") + + _mirror_skills_to(target, [bundled, proj]) + assert sorted(p.name for p in target.iterdir() if p.is_dir()) == [ + "alpha", + "skill-creator", + "user-skill", + ] + manifest = json.loads((target / _SKILL_MANIFEST).read_text()) + assert manifest == ["alpha", "skill-creator"] + assert "user-skill" not in manifest + + # Re-run with skill-creator gone → reaped; user-skill preserved. + _mirror_skills_to(target, [proj]) + assert sorted(p.name for p in target.iterdir() if p.is_dir()) == [ + "alpha", + "user-skill", + ] + + +def test_mirror_truncates_long_description(tmp_path): + long_desc = "x" * (_NATIVE_DESCRIPTION_CAP + 500) + src = _mkskill(tmp_path / "src" / "big", "big", desc=long_desc) + target = tmp_path / ".claude" / "skills" + _mirror_skills_to(target, [src]) + + import yaml + + mirrored = (target / "big" / "SKILL.md").read_text() + front = yaml.safe_load(mirrored.split("---", 2)[1]) + assert len(front["description"]) <= _NATIVE_DESCRIPTION_CAP + # Source untouched. + assert len((src / "SKILL.md").read_text()) > _NATIVE_DESCRIPTION_CAP + + +# --------------------------------------------------------------------------- +# AgentSetup.setup_skills — per-agent native-dir mirror +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "agent,subdir", + [ + ("claude", ".claude/skills"), + ("goose", ".agents/skills"), + ("cline", ".cline/skills"), + ("codex", ".agents/skills"), + ("opencode", ".agents/skills"), + ], +) +def test_setup_skills_mirrors_into_native_dir(tmp_path, agent, subdir): + from dsagt.agents import AGENTS + + _mkskill(tmp_path / "skills" / "myskill", "myskill") # a project skill + actions = AGENTS[agent]().setup_skills(tmp_path, {}) + target = tmp_path + for part in subdir.split("/"): + target = target / part + assert (target / "myskill" / "SKILL.md").exists() + assert any("kill" in a for a in actions) # reported a mirror action + + +def test_setup_skills_mirrors_registered_codes(tmp_path): + """Registered codes share the skill envelope, so they mirror natively too.""" + from dsagt.agents import AGENTS + from dsagt.registry import CodeRegistry + + CodeRegistry(runtime_dir=tmp_path).save_tool( + { + "name": "my-code", + "description": "Use when testing the native code mirror", + "executable": "echo hi", + "parameters": {}, + } + ) + AGENTS["claude"]().setup_skills(tmp_path, {}) + mirrored = tmp_path / ".claude" / "skills" / "my-code" / "SKILL.md" + assert mirrored.exists() + # The mirrored copy carries the exact dsagt-run command the agent must run. + assert "dsagt-run --code my-code -- echo hi" in mirrored.read_text() + + +def test_setup_skills_project_skill_wins_code_name_collision(tmp_path): + """A deliberately installed instruction skill outranks a same-named code.""" + from dsagt.agents import AGENTS + from dsagt.registry import CodeRegistry + + CodeRegistry(runtime_dir=tmp_path).save_tool( + { + "name": "clash", + "description": "the code", + "executable": "echo code", + "parameters": {}, + } + ) + _mkskill(tmp_path / "skills" / "clash", "clash") + AGENTS["claude"]().setup_skills(tmp_path, {}) + text = (tmp_path / ".claude" / "skills" / "clash" / "SKILL.md").read_text() + assert "dsagt-run" not in text # the skill copy, not the code copy + + +def test_setup_skills_respects_populate_native_false(tmp_path): + from dsagt.agents import AGENTS + + _mkskill(tmp_path / "skills" / "myskill", "myskill") + actions = AGENTS["claude"]().setup_skills( + tmp_path, {"skills": {"populate_native": False}} + ) + assert actions == [] + assert not (tmp_path / ".claude" / "skills").exists() + + +# --------------------------------------------------------------------------- +# install_into_project — license / attribution capture +# --------------------------------------------------------------------------- + + +def test_install_captures_ancestor_attribution(tmp_path): + cache = tmp_path / "cache" + repo = cache / "srcrepo" + repo.mkdir(parents=True) + (repo / "LICENSE").write_text("Apache-2.0") # repo-root license + cat = repo / "skills" / "modcon" + cat.mkdir(parents=True) + (cat / "ATTRIBUTION.md").write_text("upstream credits") # per-subtree + _mkskill(cat / "myskill", "myskill") + + proj = tmp_path / "proj" + proj.mkdir() + info = sc.install_into_project("myskill", proj, cache_dir=cache) + dest = proj / "skills" / "myskill" + assert (dest / "SKILL.md").exists() + assert (dest / "ATTRIBUTION.md").read_text() == "upstream credits" + assert (dest / "LICENSE").read_text() == "Apache-2.0" + prov = (dest / "PROVENANCE.txt").read_text() + assert "srcrepo" in prov and "skills/modcon/myskill" in prov + assert set(info["attribution"]) == {"ATTRIBUTION.md", "LICENSE"} + + +def test_install_skill_local_license_wins(tmp_path): + cache = tmp_path / "cache" + repo = cache / "srcrepo" + repo.mkdir(parents=True) + (repo / "LICENSE").write_text("ROOT") # repo-root license + skill = _mkskill(repo / "myskill", "myskill") + (skill / "LICENSE").write_text("SKILL-LOCAL") # skill bundles its own + + proj = tmp_path / "proj" + proj.mkdir() + info = sc.install_into_project("myskill", proj, cache_dir=cache) + dest = proj / "skills" / "myskill" + # The skill's own LICENSE (copied by copytree) must not be overwritten. + assert (dest / "LICENSE").read_text() == "SKILL-LOCAL" + assert "LICENSE" not in info["attribution"] + + +# --------------------------------------------------------------------------- +# index_catalog — frontmatter-only embedding (progressive disclosure) +# --------------------------------------------------------------------------- + + +def test_index_catalog_embeds_frontmatter_not_body(tmp_path): + captured = {} + + class _KB: + index_dir = tmp_path / "idx" + collections: list = [] + + def add_entries(self, texts, collection, metadatas=None): + captured["texts"] = texts + captured["metas"] = metadatas + return {} + + skill = tmp_path / "myskill" + skill.mkdir() + (skill / "SKILL.md").write_text( + "---\nname: myskill\ndescription: does a thing\ntags: [hpc, slurm]\n---\n" + "# Body\nSECRET_BODY_MARKER should not be embedded.\n" + ) + dirs = sc._discover_skill_dirs(tmp_path) + sc.index_catalog(dirs, "slug", "http://x", _KB()) + + joined = " ".join(captured["texts"]) + assert "myskill" in joined and "does a thing" in joined # frontmatter embedded + assert "hpc" in joined and "slurm" in joined # tags embedded + assert "SECRET_BODY_MARKER" not in joined # body NOT embedded + # description is also carried in metadata for the search summary. + assert captured["metas"][0]["description"] == "does a thing" diff --git a/tests/test_tool_executions.py b/tests/test_tool_executions.py index 775471c..2e60c67 100644 --- a/tests/test_tool_executions.py +++ b/tests/test_tool_executions.py @@ -4,9 +4,9 @@ Tests render_execution_text, execution_metadata, index_execution_record, and index_trace_archive. KnowledgeBase embedding is mocked. -Record formats match the actual output from: - - proxy_callback.py (intent + report, execution=None) - - run.py (execution only, no intent/report) +Records match the output of ``dsagt-run`` (run.py): an ``execution`` block +with no intent/report. (The pre-BYOA proxy_callback.py producer, which wrote +intent/report records, was removed — see scratch/excised_proxy_provenance.py.) """ import json @@ -14,22 +14,21 @@ from unittest.mock import MagicMock, patch import numpy as np -import pytest from dsagt.provenance import ( - TOOL_EXECUTIONS_COLLECTION as COLLECTION_NAME, + CODE_USE_COLLECTION as COLLECTION_NAME, + CodeUseIndexer, execution_metadata, - index_execution_record, index_trace_archive, render_execution_text, ) -from dsagt.knowledge import CollectionRoute, KnowledgeBase - +from dsagt.knowledge import KnowledgeBase # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- + def fake_embed(texts: list[str]) -> np.ndarray: dim = 8 vecs = np.zeros((len(texts), dim), dtype=np.float32) @@ -41,33 +40,6 @@ def fake_embed(texts: list[str]) -> np.ndarray: return vecs / norms -def make_proxy_record( - tool="fastp", - params=None, - session="s1", - record_id="toolu_001", - agent_output="Filtering complete.", -): - """Proxy callback record: intent + report, execution=None.""" - return { - "record_id": record_id, - "tool_name": tool, - "session_id": session, - "intent": { - "command": tool, - "parameters": params or {"q": 20, "in1": "sample.fq.gz"}, - "timestamp_requested": "2024-01-15T10:30:00+00:00", - "session_id": session, - }, - "execution": None, - "report": { - "agent_output": agent_output, - "timestamp_reported": "2024-01-15T10:30:13+00:00", - "wrapper_used": False, - }, - } - - def make_wrapper_record( tool="fastp", session="s1", @@ -81,7 +53,7 @@ def make_wrapper_record( """Wrapper (dsagt-run) record: execution only, no intent/report.""" return { "record_id": record_id, - "tool_name": tool, + "code_name": tool, "session_id": session, "execution": { "exact_command": [tool, "-q", "20", "--in1", "sample.fq.gz"], @@ -100,15 +72,8 @@ def make_wrapper_record( # render_execution_text # --------------------------------------------------------------------------- -class TestRenderExecutionText: - def test_proxy_record_uses_intent(self): - """Proxy record renders parameters from intent layer.""" - record = make_proxy_record() - text = render_execution_text(record) - assert "Tool: fastp" in text - assert "Parameters:" in text - assert "Agent report:" in text +class TestRenderExecutionText: def test_wrapper_record_uses_exact_command(self): """Wrapper record renders exact command from execution layer.""" @@ -137,12 +102,6 @@ def test_long_stderr_truncated(self): text = render_execution_text(record) assert "..." in text - def test_long_agent_output_truncated(self): - """Long agent output in proxy record is truncated.""" - record = make_proxy_record(agent_output="y" * 600) - text = render_execution_text(record) - assert "..." in text - def test_exact_command_list_joined(self): """exact_command as a list is joined with spaces.""" record = make_wrapper_record() @@ -157,9 +116,9 @@ def test_timing_from_wrapper(self): assert "Duration:" in text def test_minimal_record(self): - """Record with only tool_name renders.""" - text = render_execution_text({"tool_name": "ls"}) - assert "Tool: ls" in text + """Record with only code_name renders.""" + text = render_execution_text({"code_name": "ls"}) + assert "Code: ls" in text def test_empty_record(self): """Record with nothing renders tool as unknown.""" @@ -171,29 +130,17 @@ def test_empty_record(self): # execution_metadata # --------------------------------------------------------------------------- -class TestExecutionMetadata: - def test_proxy_record_metadata(self): - """Proxy record extracts metadata from top-level and intent fields.""" - record = make_proxy_record() - meta = execution_metadata(record) - - assert meta["tool_name"] == "fastp" - assert meta["session_id"] == "s1" - assert meta["wrapper_used"] == 0 - assert "return_code" not in meta - assert meta["timestamp"] == "2024-01-15T10:30:00+00:00" - assert meta["record_id"] == "toolu_001" +class TestExecutionMetadata: def test_wrapper_record_metadata(self): """Wrapper record extracts metadata from top-level and execution fields.""" record = make_wrapper_record() meta = execution_metadata(record) - assert meta["tool_name"] == "fastp" + assert meta["code_name"] == "fastp" assert meta["session_id"] == "s1" assert meta["return_code"] == 0 - assert meta["wrapper_used"] == 1 assert meta["timestamp"] == "2024-01-15T10:30:01+00:00" assert meta["record_id"] == "rec_001" @@ -205,86 +152,112 @@ def test_failed_execution_metadata(self): def test_missing_session_defaults_to_unknown(self): """Missing session_id defaults to 'unknown'.""" - record = {"tool_name": "ls"} + record = {"code_name": "ls"} meta = execution_metadata(record) assert meta["session_id"] == "unknown" - def test_top_level_fields_preferred(self): - """Top-level tool_name/session_id are used over intent fields.""" - record = make_proxy_record(tool="fastp", session="s1") - record["tool_name"] = "override_tool" - record["session_id"] = "override_session" - meta = execution_metadata(record) - assert meta["tool_name"] == "override_tool" - assert meta["session_id"] == "override_session" - # --------------------------------------------------------------------------- -# index_execution_record +# CodeUseIndexer — idempotent, incremental heartbeat indexing # --------------------------------------------------------------------------- -class TestIndexExecutionRecord: - def test_indexes_into_tool_executions_collection(self, tmp_path): - """Single record is indexed into the tool_executions collection.""" - with patch("dsagt.knowledge._make_embedder") as mock_make: +class TestCodeUseIndexer: + + def _write(self, trace_dir: Path, record: dict): + trace_dir.mkdir(parents=True, exist_ok=True) + rid = record["record_id"] + (trace_dir / f"{record.get('code_name', 'x')}_{rid}.json").write_text( + json.dumps(record) + ) + + def test_incremental_and_idempotent(self, tmp_path): + """Each tick indexes only new records; a re-tick with nothing new is a + no-op (the bug the cursor-less batch had — re-indexing everything).""" + with patch("dsagt.knowledge.Embedder.create") as mock_make: mock_embedder = MagicMock() mock_embedder.embed = fake_embed mock_make.return_value = mock_embedder - kb = KnowledgeBase(index_dir=tmp_path / "kb") - record = make_proxy_record() - result = index_execution_record(record, kb) + pdir = tmp_path / "proj" + (pdir / ".dsagt").mkdir(parents=True) + kb = KnowledgeBase(index_dir=pdir / "kb") + indexer = CodeUseIndexer(kb, pdir) + + r1 = make_wrapper_record() + r1["record_id"] = "r1" + self._write(pdir / "trace_archive", r1) + assert indexer.tick() == 1 # first record indexed + assert indexer.tick() == 0 # nothing new → no-op (idempotent) + + r2 = make_wrapper_record() + r2["record_id"] = "r2" + self._write(pdir / "trace_archive", r2) + assert indexer.tick() == 1 # only the new one + assert indexer.tick() == 0 + + # The ack set persists exactly the indexed record ids. + acks = json.loads((pdir / ".dsagt" / "code_use_acks.json").read_text()) + assert set(acks) == {"r1", "r2"} + kb.close() - assert result["collection"] == COLLECTION_NAME - assert result["entries_added"] == 1 + def test_tick_traced_wraps_in_code_use_source_span(self, tmp_path): + """The background triggers use tick_traced so the indexer's kb.* writes + nest under a dsagt.source=code_use root instead of orphaning.""" + pdir = tmp_path / "proj" + (pdir / ".dsagt").mkdir(parents=True) + kb = MagicMock() + indexer = CodeUseIndexer(kb, pdir) - results = kb.search( - "fastp quality", collection=COLLECTION_NAME, - top_k=5, rerank=False, - ) - assert len(results) > 0 - assert "fastp" in results[0]["chunk"]["text"] + opened = {} - kb.close() + class _Span: + def __enter__(self): + return self - def test_indexes_wrapper_record(self, tmp_path): - """Wrapper-only record is also indexable.""" - with patch("dsagt.knowledge._make_embedder") as mock_make: - mock_embedder = MagicMock() - mock_embedder.embed = fake_embed - mock_make.return_value = mock_embedder + def __exit__(self, *a): + return False - kb = KnowledgeBase(index_dir=tmp_path / "kb") - record = make_wrapper_record() - result = index_execution_record(record, kb) + def _fake_open_span(name, span_type=None, source=None): + opened["name"] = name + opened["source"] = source + return _Span() - assert result["entries_added"] == 1 - kb.close() + with patch("dsagt.observability.open_span", _fake_open_span): + # No records → tick returns 0, but the span must still be opened + # with the right source. + assert indexer.tick_traced() == 0 + + assert opened["source"] == "code_use" + assert opened["name"] == "code_use.index" # --------------------------------------------------------------------------- # index_trace_archive # --------------------------------------------------------------------------- + class TestIndexTraceArchive: def _write_records(self, trace_dir: Path, records: list[dict]): trace_dir.mkdir(parents=True, exist_ok=True) for i, record in enumerate(records): rid = record.get("record_id", f"rec_{i}") - path = trace_dir / f"{record.get('tool_name', 'unknown')}_{rid}.json" + path = trace_dir / f"{record.get('code_name', 'unknown')}_{rid}.json" path.write_text(json.dumps(record)) def test_indexes_all_records(self, tmp_path): """Indexes all records in trace_dir.""" trace_dir = tmp_path / "trace_archive" - self._write_records(trace_dir, [ - make_proxy_record(record_id="t1"), - make_wrapper_record(tool="megahit", record_id="t2"), - ]) - - with patch("dsagt.knowledge._make_embedder") as mock_make: + self._write_records( + trace_dir, + [ + make_wrapper_record(record_id="t1"), + make_wrapper_record(tool="megahit", record_id="t2"), + ], + ) + + with patch("dsagt.knowledge.Embedder.create") as mock_make: mock_embedder = MagicMock() mock_embedder.embed = fake_embed mock_make.return_value = mock_embedder @@ -300,12 +273,15 @@ def test_indexes_all_records(self, tmp_path): def test_skips_already_indexed(self, tmp_path): """Records with IDs in indexed_ids are skipped.""" trace_dir = tmp_path / "trace_archive" - self._write_records(trace_dir, [ - make_proxy_record(record_id="t1"), - make_proxy_record(record_id="t2"), - ]) - - with patch("dsagt.knowledge._make_embedder") as mock_make: + self._write_records( + trace_dir, + [ + make_wrapper_record(record_id="t1"), + make_wrapper_record(record_id="t2"), + ], + ) + + with patch("dsagt.knowledge.Embedder.create") as mock_make: mock_embedder = MagicMock() mock_embedder.embed = fake_embed mock_make.return_value = mock_embedder @@ -322,12 +298,15 @@ def test_skips_already_indexed(self, tmp_path): def test_updates_indexed_ids(self, tmp_path): """indexed_ids set is updated with newly indexed record IDs.""" trace_dir = tmp_path / "trace_archive" - self._write_records(trace_dir, [ - make_proxy_record(record_id="t1"), - make_wrapper_record(record_id="t2"), - ]) - - with patch("dsagt.knowledge._make_embedder") as mock_make: + self._write_records( + trace_dir, + [ + make_wrapper_record(record_id="t1"), + make_wrapper_record(record_id="t2"), + ], + ) + + with patch("dsagt.knowledge.Embedder.create") as mock_make: mock_embedder = MagicMock() mock_embedder.embed = fake_embed mock_make.return_value = mock_embedder @@ -344,7 +323,7 @@ def test_empty_directory(self, tmp_path): trace_dir = tmp_path / "trace_archive" trace_dir.mkdir() - with patch("dsagt.knowledge._make_embedder") as mock_make: + with patch("dsagt.knowledge.Embedder.create") as mock_make: mock_embedder = MagicMock() mock_embedder.embed = fake_embed mock_make.return_value = mock_embedder @@ -358,7 +337,7 @@ def test_empty_directory(self, tmp_path): def test_nonexistent_directory(self, tmp_path): """Non-existent trace_dir returns zeros.""" - with patch("dsagt.knowledge._make_embedder") as mock_make: + with patch("dsagt.knowledge.Embedder.create") as mock_make: mock_embedder = MagicMock() mock_embedder.embed = fake_embed mock_make.return_value = mock_embedder @@ -369,15 +348,18 @@ def test_nonexistent_directory(self, tmp_path): assert result["indexed"] == 0 kb.close() - def test_skips_records_without_intent_or_execution(self, tmp_path): - """Records missing both intent and execution are skipped.""" + def test_skips_records_without_execution(self, tmp_path): + """Records with no execution layer are skipped (counted as errors).""" trace_dir = tmp_path / "trace_archive" - self._write_records(trace_dir, [ - {"record_id": "bad", "tool_name": "unknown", "report": {"agent_output": "something"}}, - make_proxy_record(record_id="good"), - ]) - - with patch("dsagt.knowledge._make_embedder") as mock_make: + self._write_records( + trace_dir, + [ + {"record_id": "bad", "code_name": "unknown"}, + make_wrapper_record(record_id="good"), + ], + ) + + with patch("dsagt.knowledge.Embedder.create") as mock_make: mock_embedder = MagicMock() mock_embedder.embed = fake_embed mock_make.return_value = mock_embedder @@ -390,13 +372,16 @@ def test_skips_records_without_intent_or_execution(self, tmp_path): kb.close() def test_accepts_wrapper_only_records(self, tmp_path): - """Wrapper-only records (no intent) are valid and indexed.""" + """Wrapper records (execution only) are valid and indexed.""" trace_dir = tmp_path / "trace_archive" - self._write_records(trace_dir, [ - make_wrapper_record(record_id="w1"), - ]) - - with patch("dsagt.knowledge._make_embedder") as mock_make: + self._write_records( + trace_dir, + [ + make_wrapper_record(record_id="w1"), + ], + ) + + with patch("dsagt.knowledge.Embedder.create") as mock_make: mock_embedder = MagicMock() mock_embedder.embed = fake_embed mock_make.return_value = mock_embedder @@ -411,9 +396,9 @@ def test_accepts_wrapper_only_records(self, tmp_path): def test_idempotent_reindex(self, tmp_path): """Running index_trace_archive twice doesn't duplicate entries.""" trace_dir = tmp_path / "trace_archive" - self._write_records(trace_dir, [make_proxy_record(record_id="t1")]) + self._write_records(trace_dir, [make_wrapper_record(record_id="t1")]) - with patch("dsagt.knowledge._make_embedder") as mock_make: + with patch("dsagt.knowledge.Embedder.create") as mock_make: mock_embedder = MagicMock() mock_embedder.embed = fake_embed mock_make.return_value = mock_embedder @@ -434,6 +419,7 @@ def test_idempotent_reindex(self, tmp_path): # End-to-end: index + filtered search # --------------------------------------------------------------------------- + class TestIndexAndSearch: def _index_records(self, tmp_path, records): @@ -442,20 +428,20 @@ def _index_records(self, tmp_path, records): trace_dir.mkdir(exist_ok=True) for r in records: rid = r.get("record_id", "x") - path = trace_dir / f"{r.get('tool_name', 'unknown')}_{rid}.json" + path = trace_dir / f"{r.get('code_name', 'unknown')}_{rid}.json" path.write_text(json.dumps(r)) return trace_dir def test_search_by_tool_name(self, tmp_path): - """Index multiple records, search filtered by tool_name.""" + """Index multiple records, search filtered by code_name.""" records = [ - make_proxy_record(tool="fastp", session="s1", record_id="t1"), - make_proxy_record(tool="megahit", session="s1", record_id="t2"), - make_proxy_record(tool="fastp", session="s2", record_id="t3"), + make_wrapper_record(tool="fastp", session="s1", record_id="t1"), + make_wrapper_record(tool="megahit", session="s1", record_id="t2"), + make_wrapper_record(tool="fastp", session="s2", record_id="t3"), ] trace_dir = self._index_records(tmp_path, records) - with patch("dsagt.knowledge._make_embedder") as mock_make: + with patch("dsagt.knowledge.Embedder.create") as mock_make: mock_embedder = MagicMock() mock_embedder.embed = fake_embed mock_make.return_value = mock_embedder @@ -466,22 +452,23 @@ def test_search_by_tool_name(self, tmp_path): results = kb.search( "quality filtering parameters", collection=COLLECTION_NAME, - top_k=10, rerank=False, - where={"tool_name": "fastp"}, + top_k=10, + rerank=False, + where={"code_name": "fastp"}, ) for r in results: - assert r["chunk"]["metadata"]["tool_name"] == "fastp" + assert r["chunk"]["metadata"]["code_name"] == "fastp" kb.close() def test_search_by_session(self, tmp_path): """Index multiple records, search filtered by session_id.""" records = [ - make_proxy_record(tool="fastp", session="s1", record_id="t1"), - make_proxy_record(tool="fastp", session="s2", record_id="t2"), + make_wrapper_record(tool="fastp", session="s1", record_id="t1"), + make_wrapper_record(tool="fastp", session="s2", record_id="t2"), ] trace_dir = self._index_records(tmp_path, records) - with patch("dsagt.knowledge._make_embedder") as mock_make: + with patch("dsagt.knowledge.Embedder.create") as mock_make: mock_embedder = MagicMock() mock_embedder.embed = fake_embed mock_make.return_value = mock_embedder @@ -492,7 +479,8 @@ def test_search_by_session(self, tmp_path): results = kb.search( "fastp", collection=COLLECTION_NAME, - top_k=10, rerank=False, + top_k=10, + rerank=False, where={"session_id": "s1"}, ) for r in results: @@ -507,7 +495,7 @@ def test_search_wrapper_records_by_return_code(self, tmp_path): ] trace_dir = self._index_records(tmp_path, records) - with patch("dsagt.knowledge._make_embedder") as mock_make: + with patch("dsagt.knowledge.Embedder.create") as mock_make: mock_embedder = MagicMock() mock_embedder.embed = fake_embed mock_make.return_value = mock_embedder @@ -518,7 +506,8 @@ def test_search_wrapper_records_by_return_code(self, tmp_path): results = kb.search( "tool execution", collection=COLLECTION_NAME, - top_k=10, rerank=False, + top_k=10, + rerank=False, where={"return_code": 1}, ) for r in results: diff --git a/tests/test_trace_pipeline.py b/tests/test_trace_pipeline.py new file mode 100644 index 0000000..5d5fd5e --- /dev/null +++ b/tests/test_trace_pipeline.py @@ -0,0 +1,224 @@ +"""Parity test: our pipeline vs MLflow's own claude_code autolog. + +The same Claude transcript is logged two ways into one serverless sqlite store: + oracle: mlflow.claude_code.tracing.process_transcript(path) + ours: ClaudeTranslator → Trace → observability.MLflowSink + +then the two resulting MLflow traces are compared structurally. This is the +"how does our product match up" check — it pins our foreign-trace output to the +autolog look (AGENT root + llm/tool children, anthropic message format, token +usage) that makes those traces navigable. + +Real mlflow logging to a local sqlite file — no network. Skips cleanly if the +maintained autolog parser isn't importable. +""" + +import json + +import pytest + +from dsagt.observability import MLflowSink +from dsagt.traces import ClaudeTranslator + +process_transcript = pytest.importorskip( + "mlflow.claude_code.tracing" +).process_transcript + + +def _asst(ts, *blocks, model="claude-sonnet-4-5", usage=None): + return { + "type": "assistant", + "timestamp": ts, + "message": { + "role": "assistant", + "model": model, + "content": list(blocks), + "usage": usage or {"input_tokens": 100, "output_tokens": 20}, + }, + } + + +def _user(ts, content, tool_use_result=False): + rec = { + "type": "user", + "timestamp": ts, + "message": {"role": "user", "content": content}, + } + if tool_use_result: + rec["toolUseResult"] = {"stdout": "a.db\nb.db"} + return rec + + +TRANSCRIPT = [ + _user("2026-06-19T15:49:29.000Z", "list the db files"), + _asst("2026-06-19T15:49:30.000Z", {"type": "thinking", "thinking": "let me think"}), + _asst( + "2026-06-19T15:49:31.000Z", + {"type": "text", "text": "Let me look at the files."}, + ), + _asst( + "2026-06-19T15:49:32.000Z", + { + "type": "tool_use", + "id": "tu1", + "name": "Bash", + "input": {"command": "ls *.db"}, + }, + ), + _user( + "2026-06-19T15:49:33.000Z", + [{"type": "tool_result", "tool_use_id": "tu1", "content": "a.db\nb.db"}], + tool_use_result=True, + ), + _asst( + "2026-06-19T15:49:35.000Z", + {"type": "text", "text": "Found two databases: a.db and b.db."}, + ), +] + + +def _span_shape(trace): + """(span_type, name) for every span, root first then children in start order.""" + spans = sorted(trace.data.spans, key=lambda s: s.start_time_ns) + return [(str(s.span_type), s.name) for s in spans] + + +def _llm_spans(trace): + return [s for s in trace.data.spans if str(s.span_type) == "LLM"] + + +def _durations_ns(trace): + """Per-span (end - start) in ns, ordered by start time.""" + spans = sorted(trace.data.spans, key=lambda s: s.start_time_ns) + return [s.end_time_ns - s.start_time_ns for s in spans] + + +@pytest.fixture +def mlflow_sqlite(tmp_path, monkeypatch): + """Point MLflow at an isolated sqlite store under tmp; chdir so the autolog + parser's ``.claude/mlflow`` log dir also lands in tmp.""" + import mlflow + + monkeypatch.chdir(tmp_path) + uri = f"sqlite:///{tmp_path / 'mlflow.db'}" + mlflow.set_tracking_uri(uri) + mlflow.set_experiment("parity") + return uri + + +def test_our_pipeline_matches_autolog_span_layout(tmp_path, mlflow_sqlite): + import mlflow + + # --- oracle: mlflow's maintained parser --- + path = tmp_path / "transcript.jsonl" + path.write_text("\n".join(json.dumps(r) for r in TRANSCRIPT)) + oracle = process_transcript(str(path), session_id="proj:sess1") + assert oracle is not None, "autolog parser produced no trace" + + # --- ours: translator → sink --- + trace = ClaudeTranslator().translate( + TRANSCRIPT, trace_id="tr1", session_id="proj:sess1", project="parity" + ) + (ours_id,) = MLflowSink(mlflow_sqlite, "parity").write( + trace + ) # one turn → one trace + ours = mlflow.get_trace(ours_id) + + # Same span tree: one AGENT root, two llm turns, one tool span — thinking + # turn produces no span in either. + assert _span_shape(ours) == _span_shape(oracle) + assert _span_shape(ours) == [ + ("AGENT", "claude_code_conversation"), + ("LLM", "llm"), + ("TOOL", "tool_Bash"), + ("LLM", "llm"), + ] + + # Span durations match the autolog model (ported #1): both derive each + # span's end from the next entry's timestamp, so the per-span durations line + # up — and crucially none is the absurd "now minus backdated start". + ours_durs = _durations_ns(ours) + oracle_durs = _durations_ns(oracle) + assert all(d > 0 for d in ours_durs) + for d_ours, d_oracle in zip(ours_durs, oracle_durs): + assert abs(d_ours - d_oracle) <= 2_000_000 # within 2ms (int-ns rounding) + + +def test_our_llm_spans_carry_the_navigable_attributes(mlflow_sqlite): + import mlflow + from mlflow.tracing.constant import SpanAttributeKey + + trace = ClaudeTranslator().translate( + TRANSCRIPT, trace_id="tr2", session_id="proj:sess1", project="parity" + ) + (tid,) = MLflowSink(mlflow_sqlite, "parity").write(trace) + ours = mlflow.get_trace(tid) + + llm = _llm_spans(ours) + assert llm, "expected llm spans" + for span in llm: + attrs = span.attributes + # the anthropic message-format flag is what triggers Chat-UI rendering + assert attrs.get(SpanAttributeKey.MESSAGE_FORMAT) == "anthropic" + # outputs are an anthropic assistant message + assert span.outputs["role"] == "assistant" + assert span.outputs["type"] == "message" + # token usage present + assert attrs.get(SpanAttributeKey.CHAT_USAGE)["input_tokens"] == 100 + + +def test_multi_turn_each_subtree_matches_autolog_on_its_slice(tmp_path, mlflow_sqlite): + """A 2-turn session → 2 of our MLflow traces; each must match what autolog + produces from that turn's slice in isolation.""" + import mlflow + + turn1 = [ + _user("2026-06-19T16:00:00.000Z", "first question"), + _asst("2026-06-19T16:00:01.000Z", {"type": "text", "text": "first answer"}), + ] + turn2 = [ + _user("2026-06-19T16:00:02.000Z", "second question"), + _asst( + "2026-06-19T16:00:03.000Z", + {"type": "tool_use", "id": "tu9", "name": "Bash", "input": {"c": "ls"}}, + ), + _user( + "2026-06-19T16:00:04.000Z", + [{"type": "tool_result", "tool_use_id": "tu9", "content": "ok"}], + tool_use_result=True, + ), + _asst("2026-06-19T16:00:05.000Z", {"type": "text", "text": "second answer"}), + ] + session = turn1 + turn2 + + # ours: one translate of the whole session → two traces, in turn order + trace = ClaudeTranslator().translate( + session, trace_id="trM", session_id="proj:sessM", project="parity" + ) + our_ids = MLflowSink(mlflow_sqlite, "parity").write(trace) + assert len(our_ids) == 2 + + # oracle: autolog on each turn's slice independently + for k, slice_records in enumerate([turn1, turn2]): + path = tmp_path / f"turn{k}.jsonl" + path.write_text("\n".join(json.dumps(r) for r in slice_records)) + oracle = process_transcript(str(path), session_id="proj:sessM") + assert oracle is not None + ours = mlflow.get_trace(our_ids[k]) + assert _span_shape(ours) == _span_shape(oracle) + + +def test_session_tag_and_canonical_id_on_trace(mlflow_sqlite): + import mlflow + from mlflow.tracing.constant import TraceMetadataKey + + trace = ClaudeTranslator().translate( + TRANSCRIPT, trace_id="tr3", session_id="proj:sess1", project="parity" + ) + (tid,) = MLflowSink(mlflow_sqlite, "parity").write(trace) + ours = mlflow.get_trace(tid) + md = ours.info.trace_metadata + assert md.get(TraceMetadataKey.TRACE_SESSION) == "proj:sess1" + # Per-turn idempotency key: :. + assert md.get("dsagt.trace_id").startswith("tr3:") + assert md.get("dsagt.agent") == "claude" diff --git a/tests/test_trace_scan.py b/tests/test_trace_scan.py new file mode 100644 index 0000000..be4d2b8 --- /dev/null +++ b/tests/test_trace_scan.py @@ -0,0 +1,413 @@ +"""Tests for the in-session trace heartbeat (TraceCollector). + +Exercise the collect logic directly (no event loop): the completeness watermark +(defer the open last turn) and ack-set idempotency, end-to-end through the real +ClaudeReader + ClaudeTranslator + MLflowSink into a tmp sqlite store. +""" + +import json + +import pytest + +from dsagt.traces import ( + Trace, + TraceCollector, + _transcript_dir, + make_trace_collector, +) + + +def _asst(ts, *blocks): + return { + "type": "assistant", + "timestamp": ts, + "message": { + "role": "assistant", + "model": "claude-x", + "content": list(blocks), + "usage": {"input_tokens": 5, "output_tokens": 2}, + }, + } + + +def _user(ts, content, uuid, tool_use_result=False): + rec = { + "type": "user", + "timestamp": ts, + "uuid": uuid, + "message": {"role": "user", "content": content}, + } + if tool_use_result: + rec["toolUseResult"] = {"stdout": "ok"} + return rec + + +def _append(path, *records): + with open(path, "a") as fh: + for r in records: + fh.write(json.dumps(r) + "\n") + + +@pytest.fixture +def scan_env(tmp_path, monkeypatch): + monkeypatch.chdir(tmp_path) + proj = tmp_path / "proj" + (proj / ".dsagt").mkdir(parents=True) + proot = tmp_path / "projects" + tdir = _transcript_dir(proj, proot) + tdir.mkdir(parents=True) + f = tdir / "sess.jsonl" + uri = f"sqlite:///{tmp_path / 'mlflow.db'}" + collector = make_trace_collector( + "claude", proj, "proj", "proj:s", uri, projects_root=proot + ) + return collector, f, proj + + +def test_make_trace_collector_registered_agents(tmp_path): + # The heartbeat isn't Claude-special: it runs for any agent with a pipeline. + for agent in ("claude", "codex", "goose", "opencode", "cline"): + assert ( + make_trace_collector(agent, tmp_path, "p", "p:s", "sqlite:///x.db") + is not None + ) + # An agent with no pipeline registered simply gets no heartbeat. + assert ( + make_trace_collector("nonesuch", tmp_path, "p", "p:s", "sqlite:///x.db") is None + ) + + +def test_subset_keeps_only_named_roots_and_their_children(): + trace = Trace("t", "s", "claude", "p") + trace.add_agent_root("r1", "claude_code_conversation", start_time=None, prompt="") + trace.add_llm_span( + "r1-0", parent_id="r1", start_time=None, end_time=None, request=[], response=[] + ) + trace.add_agent_root("r2", "claude_code_conversation", start_time=None, prompt="") + trace.add_llm_span( + "r2-0", parent_id="r2", start_time=None, end_time=None, request=[], response=[] + ) + sub = trace.subset({"r1"}) + assert [s["span_id"] for s in sub.spans] == ["r1", "r1-0"] + + +def test_periodic_collect_emits_complete_turns_and_defers_the_open_last(scan_env): + collector, f, _ = scan_env + # Two complete turns, then an open turn (prompt only). + _append( + f, + _user("2026-06-19T15:00:00.000Z", "q1", "u1"), + _asst("2026-06-19T15:00:01.000Z", {"type": "text", "text": "a1"}), + _user("2026-06-19T15:00:02.000Z", "q2", "u2"), + _asst("2026-06-19T15:00:03.000Z", {"type": "text", "text": "a2"}), + _user("2026-06-19T15:00:04.000Z", "q3 (still working)", "u3"), + ) + # roots = [u1, u2, u3]; periodic emits roots[:-1] = u1, u2; u3 deferred. + assert collector.collect() == 2 + assert collector._load_acks("mlflow") == {"proj:s:u1", "proj:s:u2"} + + +def test_collect_is_idempotent(scan_env): + collector, f, _ = scan_env + _append( + f, + _user("2026-06-19T15:00:00.000Z", "q1", "u1"), + _asst("2026-06-19T15:00:01.000Z", {"type": "text", "text": "a1"}), + _user("2026-06-19T15:00:02.000Z", "q2", "u2"), + ) + assert collector.collect() == 1 # u1 emitted, u2 (last) deferred + assert collector.collect() == 0 # nothing new + assert collector.collect() == 0 + + +def test_deferred_turn_emits_once_a_later_prompt_bounds_it(scan_env): + collector, f, _ = scan_env + _append( + f, + _user("2026-06-19T15:00:00.000Z", "q1", "u1"), + _asst("2026-06-19T15:00:01.000Z", {"type": "text", "text": "a1"}), + _user("2026-06-19T15:00:02.000Z", "q2", "u2"), + ) + assert collector.collect() == 1 # u1 only + # A new prompt arrives → u2 is no longer the open turn. + _append( + f, + _asst("2026-06-19T15:00:03.000Z", {"type": "text", "text": "a2"}), + _user("2026-06-19T15:00:04.000Z", "q3", "u3"), + ) + assert collector.collect() == 1 # u2 now emitted; u3 deferred + assert collector._load_acks("mlflow") == {"proj:s:u1", "proj:s:u2"} + + +def test_final_flush_emits_the_deferred_last_turn(scan_env): + collector, f, _ = scan_env + _append( + f, + _user("2026-06-19T15:00:00.000Z", "q1", "u1"), + _asst("2026-06-19T15:00:01.000Z", {"type": "text", "text": "a1"}), + _user("2026-06-19T15:00:02.000Z", "q2", "u2"), + _asst("2026-06-19T15:00:03.000Z", {"type": "text", "text": "a2"}), + ) + assert collector.collect() == 1 # u1; u2 deferred + assert collector.collect(include_last=True) == 1 # end-of-session flush emits u2 + assert collector.collect(include_last=True) == 0 # idempotent + assert collector._load_acks("mlflow") == {"proj:s:u1", "proj:s:u2"} + + +def test_collect_on_empty_transcript_is_zero(scan_env): + collector, f, _ = scan_env + f.write_text("") + assert collector.collect() == 0 + assert collector.collect(include_last=True) == 0 + + +class _Recorder: + def __init__(self, name, fail=False): + self.name = name + self.fail = fail + self.seen = [] + + def write(self, trace): + if self.fail: + raise RuntimeError("consumer down") + self.seen.append({s["span_id"] for s in trace.spans if s["parent_id"] is None}) + + +class _FakeReader: + def read(self): + return [{"x": 1}] # non-empty so collect proceeds; translator ignores + + +class _FakeTranslator: + def __init__(self, trace): + self._trace = trace + + def translate(self, records, **kw): + return self._trace + + +def test_consumers_ack_independently(tmp_path): + # Two AGENT roots → two complete turns (include_last emits both). + trace = Trace("t", "p:s", "claude", "p") + trace.add_agent_root("r1", "c", start_time=None, prompt="") + trace.add_agent_root("r2", "c", start_time=None, prompt="") + good = _Recorder("good") + bad = _Recorder("bad", fail=True) + collector = TraceCollector( + _FakeReader(), + _FakeTranslator(trace), + project="p", + session_id="p:s", + project_dir=tmp_path, + consumers=[good, bad], + ) + + # One turn advanced for at least one consumer → returns 2 (both turns new + # to the good consumer); the failing one logs and holds its mark back. + assert collector.collect(include_last=True) == 2 + assert good.seen == [{"r1", "r2"}] + # Acks are session-qualified (:). + assert collector._load_acks("good") == {"p:s:r1", "p:s:r2"} + assert collector._load_acks("bad") == set() # wedged consumer didn't advance + + # Good is fully caught up; bad retries (and fails again) — good doesn't redo. + assert collector.collect(include_last=True) == 0 + assert good.seen == [{"r1", "r2"}] # not re-delivered + assert collector._load_acks("bad") == set() + + +def test_acks_are_session_qualified_no_cross_session_collision(tmp_path): + """Two sessions in one project share the ack file, but each turn is keyed by + ``:`` — so session B's turns (the same per-transcript + ``turn-N`` indices) are not suppressed by session A's acks. Regression for + the cross-session collision that silently dropped every session after the + first. + """ + + def _trace(session_id): + # Both sessions produce the SAME index-based span ids ("r1"/"r2"). + t = Trace("t", session_id, "claude", "p") + t.add_agent_root("r1", "c", start_time=None, prompt="") + t.add_agent_root("r2", "c", start_time=None, prompt="") + return t + + rec_a = _Recorder("mlflow") + collector_a = TraceCollector( + _FakeReader(), + _FakeTranslator(_trace("p:1")), + project="p", + session_id="p:1", + project_dir=tmp_path, + consumers=[rec_a], + ) + assert collector_a.collect(include_last=True) == 2 + + # Session B: new collector, SAME project_dir (shared ack file), new session + # id, identical span ids. Must still emit both turns. + rec_b = _Recorder("mlflow") + collector_b = TraceCollector( + _FakeReader(), + _FakeTranslator(_trace("p:2")), + project="p", + session_id="p:2", + project_dir=tmp_path, + consumers=[rec_b], + ) + assert collector_b.collect(include_last=True) == 2 + assert rec_b.seen == [{"r1", "r2"}] + # Both sessions' qualified keys coexist in the one shared ack file. + assert collector_b._load_acks("mlflow") == { + "p:1:r1", + "p:1:r2", + "p:2:r1", + "p:2:r2", + } + + +def test_make_trace_collector_pins_source(tmp_path): + """A pinned source is read regardless of the newest-file logic, and the + collector reports it via active_source().""" + transcript = tmp_path / "pinned.jsonl" + _append( + transcript, + _user("2026-06-19T15:00:00.000Z", "q1", "u1"), + _asst("2026-06-19T15:00:01.000Z", {"type": "text", "text": "a1"}), + ) + collector = make_trace_collector( + "claude", + tmp_path / "proj", + "p", + "p:1", + f"sqlite:///{tmp_path / 'm.db'}", + source=str(transcript), + ) + assert collector.active_source() == str(transcript) + assert collector.collect(include_last=True) == 1 # the one complete turn + + +def test_record_trace_source_updates_current_session(tmp_path): + from dsagt.session import read_state, record_trace_source, write_state + + proj = tmp_path / "proj" + (proj / ".dsagt").mkdir(parents=True) + write_state(proj, {"sessions": [{"id": 1}], "memory_cursor": {}}) + + record_trace_source(proj, "/path/to/t.jsonl") + assert read_state(proj)["sessions"][-1]["trace_source"] == "/path/to/t.jsonl" + # A non-str token (e.g. a SQLite session id) round-trips through YAML. + record_trace_source(proj, 42) + assert read_state(proj)["sessions"][-1]["trace_source"] == 42 + # No session minted yet → safe no-op. + write_state(tmp_path / "empty", {"sessions": [], "memory_cursor": {}}) + record_trace_source(tmp_path / "empty", "/x") # must not raise + + +def test_catch_up_traces_emits_previous_session_dangling_and_is_idempotent(tmp_path): + """The startup catch-up re-collects the previous session's pinned transcript, + emits the turns the live pass missed, and a second pass is a no-op (the + session-qualified ack files dedupe).""" + from unittest.mock import MagicMock + + from dsagt.session import _catch_up_traces, write_state + + proj = tmp_path / "proj" + (proj / ".dsagt").mkdir(parents=True) + transcript = tmp_path / "prev.jsonl" # pinned, so location is irrelevant + _append( + transcript, + _user("2026-06-19T15:00:00.000Z", "q1", "u1"), + _asst("2026-06-19T15:00:01.000Z", {"type": "text", "text": "a1"}), + _user("2026-06-19T15:00:02.000Z", "q2", "u2"), + _asst("2026-06-19T15:00:03.000Z", {"type": "text", "text": "a2"}), + ) + # Previous session (id=1) recorded its trace source; current session is id=2. + write_state( + proj, + { + "sessions": [{"id": 1, "trace_source": str(transcript)}, {"id": 2}], + "memory_cursor": {}, + }, + ) + config = {"project": "proj", "agent": "claude", "project_dir": str(proj)} + + n = _catch_up_traces(proj, config, MagicMock()) + assert n == 2 # both completed turns flushed via include_last + # Idempotent — acks (keyed proj-1:) suppress a re-collect. + assert _catch_up_traces(proj, config, MagicMock()) == 0 + + +def test_catch_up_traces_noop_without_previous_or_transcript(tmp_path): + from unittest.mock import MagicMock + + from dsagt.session import _catch_up_traces, write_state + + proj = tmp_path / "proj" + (proj / ".dsagt").mkdir(parents=True) + config = {"project": "proj", "agent": "claude", "project_dir": str(proj)} + + # Only one session → nothing to catch up. + write_state(proj, {"sessions": [{"id": 1}], "memory_cursor": {}}) + assert _catch_up_traces(proj, config, MagicMock()) == 0 + + # Previous session exists but recorded no transcript (SQLite agent / too + # short) → skip rather than risk reading the new session's transcript. + write_state( + proj, {"sessions": [{"id": 1}, {"id": 2}], "memory_cursor": {}} + ) + assert _catch_up_traces(proj, config, MagicMock()) == 0 + + +def test_sqlite_reader_pins_to_a_specific_session(tmp_path): + """SQLite agents pin uniformly: a pinned GooseReader reads the *specified* + session, not the latest — so the catch-up backstops goose/opencode/cline + too, not just the JSONL agents. (opencode/cline mirror this shape.) + """ + import os + import sqlite3 + + from dsagt.traces import GooseReader + + db = tmp_path / "sessions.db" + con = sqlite3.connect(db) + con.executescript( + "CREATE TABLE sessions (id TEXT, working_dir TEXT, updated_at INTEGER);" + "CREATE TABLE messages (id INTEGER PRIMARY KEY, session_id TEXT, role TEXT," + " content_json TEXT, created_timestamp INTEGER);" + ) + wd = os.path.abspath(tmp_path / "proj") + con.execute("INSERT INTO sessions VALUES ('A', ?, 100)", (wd,)) # older + con.execute("INSERT INTO sessions VALUES ('B', ?, 200)", (wd,)) # newer + con.execute("INSERT INTO messages VALUES (1, 'A', 'user', '[\"qA\"]', 1)") + con.execute("INSERT INTO messages VALUES (2, 'B', 'user', '[\"qB\"]', 2)") + con.commit() + con.close() + + reader = GooseReader(tmp_path / "proj", db_path=db) + # Unpinned → newest session (B), and reports its id. + assert reader.active_source() == "B" + assert reader.read()[0]["content"] == ["qB"] + # Pinned to the older session A → reads exactly A (the catch-up's job). + reader.pin("A") + assert reader.active_source() == "A" + assert reader.read()[0]["content"] == ["qA"] + + +def test_emitted_traces_land_in_the_store(scan_env): + import mlflow + + collector, f, _ = scan_env + _append( + f, + _user("2026-06-19T15:00:00.000Z", "q1", "u1"), + _asst("2026-06-19T15:00:01.000Z", {"type": "text", "text": "a1"}), + _user("2026-06-19T15:00:02.000Z", "q2", "u2"), + _asst("2026-06-19T15:00:03.000Z", {"type": "text", "text": "a2"}), + ) + collector.collect(include_last=True) # emit both turns + exp = mlflow.get_experiment_by_name("proj") + traces = mlflow.search_traces( + locations=[exp.experiment_id], return_type="list" + ) + sessions = {t.info.trace_metadata.get("mlflow.trace.session") for t in traces} + assert len(traces) == 2 + assert sessions == {"proj:s"} diff --git a/tests/test_traces.py b/tests/test_traces.py new file mode 100644 index 0000000..2b70cdf --- /dev/null +++ b/tests/test_traces.py @@ -0,0 +1,437 @@ +"""Fixture tests for the canonical trace waist + Claude translator (Phase 2). + +Pure, no disk / no network. Two things under test: +- ``to_exchanges`` projects LLM spans onto the conversational shape the episodic + indexer chunks + embeds (carrying ``turn_id`` + content) — the seam. +- ``ClaudeTranslator`` turns raw transcript records into a ``Trace`` + with the autolog-style AGENT/LLM/TOOL span layout. +""" + +import json + +from dsagt.traces import ( + ClaudeReader, + ClaudeTranslator, + Trace, + _transcript_dir, +) + +# --------------------------------------------------------------------------- +# to_exchanges projection (windowed model) +# --------------------------------------------------------------------------- + + +def _windowed_trace() -> Trace: + """Two LLM spans whose ``request`` already holds the per-turn window.""" + trace = Trace( + trace_id="tr1", session_id="proj:sess1", agent="claude", project="proj" + ) + trace.add_llm_span( + "s1", + parent_id="r1", + start_time=100.0, + end_time=None, + request=[ + {"role": "user", "content": [{"type": "text", "text": "Profile sales.csv"}]} + ], + response=[ + { + "type": "tool_use", + "id": None, + "name": "profile", + "input": {"path": "sales.csv"}, + } + ], + ) + trace.add_tool_span( + "s3", + parent_id="r1", + start_time=101.0, + end_time=None, + name="profile", + tool_input={"path": "sales.csv"}, + result="1000 rows", + ) + trace.add_llm_span( + "s2", + parent_id="r1", + start_time=102.0, + end_time=None, + request=[ + { + "role": "user", + "content": [ + {"type": "tool_result", "tool_use_id": "t1", "content": "1000 rows"} + ], + } + ], + response=[{"type": "text", "text": "The file has 1000 rows."}], + ) + return trace + + +def test_to_exchanges_one_per_llm_span_skipping_tool_spans(): + exchanges = _windowed_trace().to_exchanges() + assert len(exchanges) == 2 + assert [e["timestamp"] for e in exchanges] == [100.0, 102.0] + + +def test_to_exchanges_uses_request_window_directly(): + ex1, ex2 = _windowed_trace().to_exchanges() + assert [m["role"] for m in ex1["new_messages"]] == ["user"] + # The tool_result block round-trips in Anthropic shape memory can read. + block = ex2["new_messages"][0]["content"][0] + assert block == {"type": "tool_result", "tool_use_id": "t1", "content": "1000 rows"} + assert ex2["response"] == [{"type": "text", "text": "The file has 1000 rows."}] + + +def test_projection_carries_turn_id_and_content(): + """The waist's whole justification: to_exchanges carries ``turn_id`` (groups + a turn's chunks) and the real content the episodic indexer embeds.""" + exchanges = _windowed_trace().to_exchanges() + assert all(ex["turn_id"] for ex in exchanges) # span id, for chunk grouping + flat = json.dumps(exchanges) + assert "Profile sales.csv" in flat + assert "profile" in flat # tool_use name (in the assistant response) + assert "1000 rows" in flat # tool_result content + assert "The file has 1000 rows." in flat # final answer + + +def test_none_tolerant_usage_and_timing(): + trace = Trace(trace_id="t", session_id="s", agent="goose", project="p") + span = trace.add_llm_span( + "s1", parent_id="r1", start_time=None, end_time=None, request=[], response=[] + ) + assert span["usage"] is None and span["start_time"] is None + assert trace.to_exchanges() == [ + {"turn_id": "s1", "timestamp": None, "new_messages": [], "response": []} + ] + + +# --------------------------------------------------------------------------- +# ClaudeTranslator +# --------------------------------------------------------------------------- + + +def _asst(ts, *blocks, model="claude-x", usage=None): + return { + "type": "assistant", + "timestamp": ts, + "message": { + "role": "assistant", + "model": model, + "content": list(blocks), + "usage": usage or {"input_tokens": 10, "output_tokens": 3}, + }, + } + + +def _user(ts, content, tool_use_result=False): + rec = { + "type": "user", + "timestamp": ts, + "message": {"role": "user", "content": content}, + } + if tool_use_result: + rec["toolUseResult"] = {"stdout": "..."} + return rec + + +def _single_turn_records(): + """user → thinking → text → tool_use → tool_result → text (final).""" + return [ + _user("2026-06-19T15:49:29.000Z", "list the db files"), + _asst("2026-06-19T15:49:30.000Z", {"type": "thinking", "thinking": "hmm"}), + _asst("2026-06-19T15:49:31.000Z", {"type": "text", "text": "Let me look."}), + _asst( + "2026-06-19T15:49:32.000Z", + {"type": "tool_use", "id": "tu1", "name": "Bash", "input": {"cmd": "ls"}}, + ), + _user( + "2026-06-19T15:49:33.000Z", + [{"type": "tool_result", "tool_use_id": "tu1", "content": "a.db\nb.db"}], + tool_use_result=True, + ), + _asst( + "2026-06-19T15:49:34.000Z", {"type": "text", "text": "Found a.db and b.db."} + ), + ] + + +def _translate(records): + return ClaudeTranslator().translate( + records, trace_id="tr", session_id="proj:s", project="proj" + ) + + +def test_translate_builds_agent_root_with_prompt_and_response(): + trace = _translate(_single_turn_records()) + root = trace.spans[0] + assert root["kind"] == "AGENT" + assert root["name"] == "claude_code_conversation" + assert root["attributes"]["prompt"] == "list the db files" + assert root["attributes"]["response"] == "Found a.db and b.db." + assert trace.started_at is not None and trace.ended_at is not None + + +def test_translate_span_layout_matches_autolog_shape(): + trace = _translate(_single_turn_records()) + kinds = [(s["kind"], s["name"]) for s in trace.spans] + # thinking turn produces NO span; two text turns → 2 LLM; one tool_use → 1 TOOL. + assert kinds == [ + ("AGENT", "claude_code_conversation"), + ("LLM", "llm"), + ("TOOL", "tool_Bash"), + ("LLM", "llm"), + ] + + +def test_translate_llm_span_carries_model_usage_and_window(): + trace = _translate(_single_turn_records()) + llms = [s for s in trace.spans if s["kind"] == "LLM"] + first, second = llms + assert first["model"] == "claude-x" + assert first["usage"]["input_tokens"] == 10 and first["usage"]["output_tokens"] == 3 + # The final LLM turn's window is the tool_use + tool_result run since the + # previous text turn (the "Let me look." boundary is excluded). + roles = [m["role"] for m in second["request"]] + assert "user" in roles # the tool_result message + assert second["response"][0]["text"] == "Found a.db and b.db." + + +def test_translate_tool_span_has_input_and_result(): + trace = _translate(_single_turn_records()) + tool = next(s for s in trace.spans if s["kind"] == "TOOL") + assert tool["attributes"]["tool_name"] == "Bash" + assert tool["attributes"]["input"] == {"cmd": "ls"} + assert tool["attributes"]["result"] == "a.db\nb.db" + assert tool["parent_id"] == trace.spans[0]["span_id"] + + +def test_translate_empty_transcript_is_none(): + assert _translate([]) is None + + +# --------------------------------------------------------------------------- +# Ported edge cases (from mlflow's claude_code parser) +# --------------------------------------------------------------------------- + + +def test_spans_get_real_durations_not_a_now_end(): + """#1 — every span ends via the next-timestamp model, never an open end_time + (which the sink would otherwise stamp as wall-clock now → absurd duration).""" + trace = _translate(_single_turn_records()) + for span in trace.spans: + assert span["start_time"] is not None and span["end_time"] is not None + assert span["end_time"] > span["start_time"] + # Final LLM turn has no following entry → 1s fallback, not a huge span. + final_llm = [s for s in trace.spans if s["kind"] == "LLM"][-1] + assert final_llm["end_time"] - final_llm["start_time"] == 1.0 + + +def test_skill_injection_user_message_is_not_taken_as_the_prompt(): + """#2 — a user entry following a Skill tool result (commandName) is an + injection, not the human prompt; the real prompt is selected instead.""" + records = [ + _user("2026-06-19T15:00:00.000Z", "real prompt"), + _asst("2026-06-19T15:00:01.000Z", {"type": "text", "text": "answer"}), + { + "type": "user", + "timestamp": "2026-06-19T15:00:02.000Z", + "toolUseResult": {"commandName": "some-skill"}, + "message": {"role": "user", "content": "skill tool result"}, + }, + _user("2026-06-19T15:00:03.000Z", "injected skill content"), + ] + trace = _translate(records) + assert trace.spans[0]["attributes"]["prompt"] == "real prompt" + + +def test_local_command_stdout_is_not_taken_as_the_prompt(): + """#2 — slash-command stdout echoes are skipped when finding the prompt.""" + records = [ + _user("2026-06-19T15:00:00.000Z", "real prompt"), + _asst("2026-06-19T15:00:01.000Z", {"type": "text", "text": "answer"}), + _user( + "2026-06-19T15:00:02.000Z", + "ran /foo", + ), + ] + trace = _translate(records) + assert trace.spans[0]["attributes"]["prompt"] == "real prompt" + + +def test_steer_message_folds_into_the_llm_window(): + """#3 — a queue-operation/enqueue steer message becomes a user message in + the following text turn's input window.""" + records = [ + _user("2026-06-19T15:00:00.000Z", "start a long task"), + _asst( + "2026-06-19T15:00:01.000Z", + {"type": "tool_use", "id": "tu1", "name": "Bash", "input": {"c": "sleep"}}, + ), + { + "type": "queue-operation", + "operation": "enqueue", + "timestamp": "2026-06-19T15:00:02.000Z", + "content": "actually, stop and summarize", + }, + _user( + "2026-06-19T15:00:03.000Z", + [{"type": "tool_result", "tool_use_id": "tu1", "content": "done"}], + tool_use_result=True, + ), + _asst("2026-06-19T15:00:04.000Z", {"type": "text", "text": "Summary."}), + ] + trace = _translate(records) + final_llm = [s for s in trace.spans if s["kind"] == "LLM"][-1] + texts = [ + b["text"] + for m in final_llm["request"] + for b in m["content"] + if b["type"] == "text" + ] + assert "actually, stop and summarize" in texts + + +# --------------------------------------------------------------------------- +# Multi-turn segmentation +# --------------------------------------------------------------------------- + + +def _two_turn_records(): + return [ + _user("2026-06-19T15:00:00.000Z", "first question"), + _asst("2026-06-19T15:00:01.000Z", {"type": "text", "text": "first answer"}), + _user("2026-06-19T15:00:02.000Z", "second question"), + _asst( + "2026-06-19T15:00:03.000Z", + {"type": "tool_use", "id": "t", "name": "Bash", "input": {"c": "ls"}}, + ), + _user( + "2026-06-19T15:00:04.000Z", + [{"type": "tool_result", "tool_use_id": "t", "content": "ok"}], + tool_use_result=True, + ), + _asst("2026-06-19T15:00:05.000Z", {"type": "text", "text": "second answer"}), + ] + + +def test_translate_segments_one_agent_subtree_per_turn(): + trace = _translate(_two_turn_records()) + roots = [s for s in trace.spans if s["parent_id"] is None] + assert len(roots) == 2 + assert roots[0]["attributes"]["prompt"] == "first question" + assert roots[1]["attributes"]["prompt"] == "second question" + # Turn 2's tool + final-answer LLM hang off turn 2's root, not turn 1's. + t2 = [s for s in trace.spans if s["parent_id"] == roots[1]["span_id"]] + assert {s["kind"] for s in t2} == {"LLM", "TOOL"} + assert roots[1]["attributes"]["response"] == "second answer" + + +def test_each_turn_is_self_contained_no_cross_turn_duration_borrow(): + trace = _translate(_two_turn_records()) + roots = [s for s in trace.spans if s["parent_id"] is None] + # Turn 1's only LLM span is its last → 1s fallback, NOT the 2s gap to turn 2. + t1_llm = next( + s + for s in trace.spans + if s["parent_id"] == roots[0]["span_id"] and s["kind"] == "LLM" + ) + assert t1_llm["end_time"] - t1_llm["start_time"] == 1.0 + + +def test_translate_skips_a_leading_partial_turn(): + """Records before the first prompt in the batch belong to an already-processed + turn (incremental read) — they produce no subtree, which is what keeps + cursor-driven reads idempotent at turn granularity.""" + records = [ + _asst( + "2026-06-19T15:00:00.000Z", {"type": "text", "text": "tail of prev turn"} + ), + _user("2026-06-19T15:00:01.000Z", "new prompt"), + _asst("2026-06-19T15:00:02.000Z", {"type": "text", "text": "answer"}), + ] + trace = _translate(records) + roots = [s for s in trace.spans if s["parent_id"] is None] + assert len(roots) == 1 + assert roots[0]["attributes"]["prompt"] == "new prompt" + + +# --------------------------------------------------------------------------- +# ClaudeReader — whole-file read +# --------------------------------------------------------------------------- + + +def _write_lines(path, objs, mode="w"): + with open(path, mode) as fh: + for o in objs: + fh.write(json.dumps(o) + "\n") + + +def test_reader_reads_the_whole_active_file(tmp_path): + proj = tmp_path / "proj" + proj.mkdir() + root = tmp_path / "projects" + tdir = _transcript_dir(proj, root) + tdir.mkdir(parents=True) + f = tdir / "sess.jsonl" + _write_lines(f, [{"a": 1}, {"a": 2}]) + + reader = ClaudeReader(proj, projects_root=root) + recs = reader.read() + assert [r["a"] for r in recs] == [1, 2] + + _write_lines(f, [{"a": 3}], mode="a") + recs2 = reader.read() + assert [r["a"] for r in recs2] == [1, 2, 3] # re-reads the whole file + + +def test_reader_leaves_a_partial_trailing_line(tmp_path): + proj = tmp_path / "proj" + proj.mkdir() + root = tmp_path / "projects" + tdir = _transcript_dir(proj, root) + tdir.mkdir(parents=True) + f = tdir / "sess.jsonl" + _write_lines(f, [{"a": 1}]) + + reader = ClaudeReader(proj, projects_root=root) + assert [r["a"] for r in reader.read()] == [1] + + # A half-written append (no newline) must not be parsed. + with open(f, "a") as fh: + fh.write('{"a": 2') + assert [r["a"] for r in reader.read()] == [1] + + # Once the line completes, it's read. + with open(f, "a") as fh: + fh.write("}\n") + assert [r["a"] for r in reader.read()] == [1, 2] + + +def test_reader_picks_the_most_recent_session(tmp_path): + proj = tmp_path / "proj" + proj.mkdir() + root = tmp_path / "projects" + tdir = _transcript_dir(proj, root) + tdir.mkdir(parents=True) + old = tdir / "old.jsonl" + _write_lines(old, [{"a": 1}, {"a": 2}]) + + # A newer session file becomes active; the reader follows the newest mtime. + import os + + new = tdir / "new.jsonl" + _write_lines(new, [{"b": 1}]) + newer = old.stat().st_mtime + 10 + os.utime(new, (newer, newer)) # ensure new is the most-recent file + reader = ClaudeReader(proj, projects_root=root) + assert [r.get("b") for r in reader.read()] == [1] + + +def test_reader_no_transcripts_is_empty(tmp_path): + proj = tmp_path / "proj" + proj.mkdir() + reader = ClaudeReader(proj, projects_root=tmp_path / "projects") + assert reader.read() == [] diff --git a/tests/test_traces_cline.py b/tests/test_traces_cline.py new file mode 100644 index 0000000..c2e26d6 --- /dev/null +++ b/tests/test_traces_cline.py @@ -0,0 +1,169 @@ +"""Fixture tests for the Cline CLI translator + reader. + +The text path mirrors real ``~/.cline`` sessions (validated in discovery); the +tool path uses standard Anthropic blocks (what Cline's SDK emits) since the +captured sessions were text-only. +""" + +import json + +import pytest + +from dsagt.observability import MLflowSink +from dsagt.traces import ClineReader, ClineTranslator + + +def _user(ts, text): + return { + "id": f"u{ts}", + "role": "user", + "ts": ts, + "content": [{"type": "text", "text": text}], + } + + +def _asst(ts, text=None, tools=(), model=None): + content = [] + if text is not None: + content.append({"type": "text", "text": text}) + for name, inp, tid in tools: + content.append({"type": "tool_use", "id": tid, "name": name, "input": inp}) + return { + "id": f"a{ts}", + "role": "assistant", + "ts": ts, + "content": content, + "model": model, + } + + +def _toolresult(ts, tid, result): + return { + "id": f"t{ts}", + "role": "user", + "ts": ts, + "content": [{"type": "tool_result", "tool_use_id": tid, "content": result}], + } + + +def _session(): + return [ + _user(1_000_000, 'do the thing'), # turn 1 + _asst( + 1_001_000, + "Let me run a command.", + tools=[("shell", {"cmd": "ls"}, "t1")], + model="m", + ), + _toolresult(1_002_000, "t1", "a.txt"), # tool result (user msg, not a prompt) + _asst(1_003_000, "Done.", model="m"), + _user(1_004_000, "next"), # turn 2 + _asst(1_005_000, "OK.", model="m"), + ] + + +def _translate(records=None): + return ClineTranslator().translate( + records if records is not None else _session(), + trace_id="cl", + session_id="proj:s", + project="clineproj", + ) + + +def test_user_input_wrapper_stripped_and_turns_segmented(): + roots = [s for s in _translate().spans if s["parent_id"] is None] + assert len(roots) == 2 # the tool_result user message is not a turn + assert roots[0]["attributes"]["prompt"] == "do the thing" # unwrapped + assert roots[1]["attributes"]["prompt"] == "next" + + +def test_assistant_text_and_tool_use_in_one_message(): + trace = _translate() + roots = [s for s in trace.spans if s["parent_id"] is None] + t1 = [ + (s["kind"], s["name"]) + for s in trace.spans + if s["parent_id"] == roots[0]["span_id"] + ] + # message 1 has text + tool_use → an LLM span then a TOOL span; then "Done." LLM. + assert t1 == [ + ("LLM", "llm"), + ("TOOL", "tool_shell"), + ("LLM", "llm"), + ] + tool = next(s for s in trace.spans if s["kind"] == "TOOL") + assert tool["attributes"]["input"] == {"cmd": "ls"} + assert tool["attributes"]["result"] == "a.txt" + assert roots[0]["attributes"]["response"] == "Done." + assert trace.spans[1]["model"] == "m" # first LLM carries the session model + + +def test_tool_result_as_text_block_list(): + records = [ + _user(1, "go"), + _asst(2, "calling", tools=[("t", {}, "x1")]), + _toolresult(3, "x1", [{"type": "text", "text": "block result"}]), + _asst(4, "ok"), + ] + tool = next(s for s in _translate(records).spans if s["kind"] == "TOOL") + assert tool["attributes"]["result"] == "block result" + + +def test_empty_is_none(): + assert _translate([]) is None + assert _translate([_asst(1, "no user prompt")]) is None + + +# --- reader over a tiny ~/.cline/data/sessions tree --- + + +def _make_session(root, sid, cwd, messages, model="m"): + d = root / sid + d.mkdir(parents=True) + (d / f"{sid}.json").write_text(json.dumps({"cwd": cwd, "model": model})) + (d / f"{sid}.messages.json").write_text(json.dumps({"messages": messages})) + + +def test_reader_picks_newest_session_for_the_cwd(tmp_path): + root = tmp_path / "sessions" + _make_session(root, "200_p1", "/proj", [_user(1, "old")]) + _make_session(root, "300_p2", "/proj", [_user(2, "new")]) + _make_session( + root, "400_x", "/elsewhere", [_user(3, "other")] + ) # newest overall, wrong cwd + reader = ClineReader("/proj", sessions_root=root) + assert reader._active_dir().name == "300_p2" + records = reader.read() + assert records[0]["content"][0]["text"] == "new" + assert records[0]["model"] == "m" # attached from metadata + + +def test_reader_no_match_is_empty(tmp_path): + root = tmp_path / "sessions" + _make_session(root, "100_a", "/elsewhere", [_user(1, "x")]) + records = ClineReader("/proj", sessions_root=root).read() + assert records == [] + + +@pytest.fixture +def mlflow_sqlite(tmp_path, monkeypatch): + import mlflow + + monkeypatch.chdir(tmp_path) + uri = f"sqlite:///{tmp_path / 'mlflow.db'}" + mlflow.set_tracking_uri(uri) + mlflow.set_experiment("clineproj") + return uri + + +def test_end_to_end_through_the_sink(mlflow_sqlite): + import mlflow + + ids = MLflowSink(mlflow_sqlite, "clineproj").write(_translate()) + assert len(ids) == 2 + exp = mlflow.get_experiment_by_name("clineproj") + traces = mlflow.search_traces( + locations=[exp.experiment_id], return_type="list" + ) + assert len(traces) == 2 diff --git a/tests/test_traces_codex.py b/tests/test_traces_codex.py new file mode 100644 index 0000000..4cbef4a --- /dev/null +++ b/tests/test_traces_codex.py @@ -0,0 +1,214 @@ +"""Fixture tests for the Codex rollout translator. + +Synthetic but grammar-faithful records (validated against a real rollout in +discovery): developer/AGENTS context, the real prompt, reasoning, assistant +text, function_call + custom_tool_call with their outputs, across two turns. +""" + +import json + +import pytest + +from dsagt.observability import MLflowSink +from dsagt.traces import CodexReader, CodexTranslator + + +def _rec(ts, rtype, payload): + return {"timestamp": ts, "type": rtype, "payload": payload} + + +def _msg(ts, role, *texts): + bt = "output_text" if role == "assistant" else "input_text" + return _rec( + ts, + "response_item", + { + "type": "message", + "role": role, + "content": [{"type": bt, "text": t} for t in texts], + }, + ) + + +def _fcall(ts, name, args, cid): + return _rec( + ts, + "response_item", + { + "type": "function_call", + "name": name, + "arguments": json.dumps(args), + "call_id": cid, + }, + ) + + +def _fout(ts, cid, out): + return _rec( + ts, + "response_item", + {"type": "function_call_output", "call_id": cid, "output": out}, + ) + + +def _custom(ts, name, raw_input, cid): + return _rec( + ts, + "response_item", + {"type": "custom_tool_call", "name": name, "input": raw_input, "call_id": cid}, + ) + + +def _custom_out(ts, cid, out): + return _rec( + ts, + "response_item", + {"type": "custom_tool_call_output", "call_id": cid, "output": out}, + ) + + +def _reasoning(ts): + return _rec( + ts, + "response_item", + {"type": "reasoning", "summary": [], "encrypted_content": "x"}, + ) + + +def _session(): + T = "2026-06-18T21:30:%02d.000Z" + return [ + _rec(T % 0, "session_meta", {"id": "s", "cwd": "/p"}), + _msg(T % 1, "developer", "system instructions"), + _msg(T % 2, "user", "# AGENTS.md instructions for /p"), # injected context + _msg(T % 3, "user", "fix the hang"), # the real prompt (turn 1) + _reasoning(T % 4), # skipped + _msg(T % 5, "assistant", "Let me look."), # LLM + _fcall(T % 6, "exec_command", {"cmd": "ls"}, "c1"), # TOOL + _fout(T % 7, "c1", "a.py\nb.py"), + _custom(T % 8, "apply_patch", "*** Begin Patch ...", "c2"), # TOOL (raw input) + _custom_out(T % 9, "c2", "Success."), + _msg(T % 10, "assistant", "Fixed it."), # LLM (final of turn 1) + _msg(T % 11, "user", "now run the tests"), # turn 2 prompt + _msg(T % 12, "assistant", "Tests pass."), # LLM + ] + + +def _translate(records=None): + return CodexTranslator().translate( + records if records is not None else _session(), + trace_id="cx", + session_id="proj:s", + project="codexproj", + ) + + +def test_agents_md_injection_is_not_the_prompt(): + roots = [s for s in _translate().spans if s["parent_id"] is None] + assert len(roots) == 2 + assert roots[0]["attributes"]["prompt"] == "fix the hang" # not the AGENTS.md msg + assert roots[1]["attributes"]["prompt"] == "now run the tests" + + +def test_turn1_span_layout_skips_reasoning(): + trace = _translate() + roots = [s for s in trace.spans if s["parent_id"] is None] + t1 = [ + (s["kind"], s["name"]) + for s in trace.spans + if s["parent_id"] == roots[0]["span_id"] + ] + # reasoning → no span; 2 assistant texts → 2 LLM; 2 tool calls → 2 TOOL. + assert t1 == [ + ("LLM", "llm"), + ("TOOL", "tool_exec_command"), + ("TOOL", "tool_apply_patch"), + ("LLM", "llm"), + ] + assert roots[0]["attributes"]["response"] == "Fixed it." + + +def test_tool_input_parsed_and_result_matched_by_call_id(): + trace = _translate() + by_name = {s["name"]: s for s in trace.spans if s["kind"] == "TOOL"} + exec_span = by_name["tool_exec_command"] + assert exec_span["attributes"]["input"] == {"cmd": "ls"} # JSON args parsed + assert exec_span["attributes"]["result"] == "a.py\nb.py" + patch_span = by_name["tool_apply_patch"] + assert ( + patch_span["attributes"]["input"] == "*** Begin Patch ..." + ) # raw custom input + assert patch_span["attributes"]["result"] == "Success." + + +def test_durations_are_bounded_per_turn(): + trace = _translate() + for s in trace.spans: + if s["parent_id"] is not None: + assert s["end_time"] is not None and s["end_time"] > s["start_time"] + # turn 1's final LLM is bounded by turn 2's prompt? No — turns are + # independent, so its last span falls back to 1s, not the gap to turn 2. + roots = [s for s in trace.spans if s["parent_id"] is None] + t1_last_llm = [ + s + for s in trace.spans + if s["parent_id"] == roots[0]["span_id"] and s["kind"] == "LLM" + ][-1] + assert t1_last_llm["end_time"] - t1_last_llm["start_time"] == 1.0 + + +def test_empty_records_is_none(): + assert _translate([]) is None + assert _translate([{"type": "event_msg", "payload": {}}]) is None + + +def test_reader_picks_the_rollout_matching_the_project_cwd(tmp_path): + root = tmp_path / "sessions" / "2026" / "06" / "18" + root.mkdir(parents=True) + # Two rollouts; only one's session_meta.cwd matches the project dir. + mine = root / "rollout-mine.jsonl" + mine.write_text("\n".join(json.dumps(r) for r in _session())) + other = root / "rollout-other.jsonl" + other.write_text( + json.dumps(_rec("t", "session_meta", {"id": "o", "cwd": "/elsewhere"})) + "\n" + ) + reader = CodexReader("/p", sessions_root=tmp_path / "sessions") + # _session()'s session_meta cwd is "/p". + assert reader.active_file() == mine + records = reader.read() + assert any(r.get("type") == "response_item" for r in records) + + +def test_reader_no_match_is_empty(tmp_path): + (tmp_path / "sessions").mkdir() + records = CodexReader("/p", sessions_root=tmp_path / "sessions").read() + assert records == [] + + +@pytest.fixture +def mlflow_sqlite(tmp_path, monkeypatch): + import mlflow + + monkeypatch.chdir(tmp_path) + uri = f"sqlite:///{tmp_path / 'mlflow.db'}" + mlflow.set_tracking_uri(uri) + mlflow.set_experiment("codexproj") + return uri + + +def test_end_to_end_through_the_sink(mlflow_sqlite): + import mlflow + + trace = _translate() + ids = MLflowSink(mlflow_sqlite, "codexproj").write(trace) + assert len(ids) == 2 # one MLflow trace per turn + exp = mlflow.get_experiment_by_name("codexproj") + traces = mlflow.search_traces( + locations=[exp.experiment_id], return_type="list" + ) + assert len(traces) == 2 + # the tool-bearing turn rendered llm + tool spans under its agent root + shapes = {len(t.data.spans) for t in traces} + # tool-bearing turn: root + 2 llm + 2 tool = 5; second turn: root + 1 llm = 2. + # Pin both — `5 in shapes` alone would pass even if turn 2 collapsed. + assert shapes == {5, 2} diff --git a/tests/test_traces_command.py b/tests/test_traces_command.py new file mode 100644 index 0000000..c8e475c --- /dev/null +++ b/tests/test_traces_command.py @@ -0,0 +1,77 @@ +"""Tests for ``dsagt traces`` — the frictionless MLflow viewer launcher.""" + +from __future__ import annotations + +from unittest.mock import MagicMock, patch + +from dsagt.commands import traces as traces_cmd + + +def _config(pdir): + return {"project": "proj", "project_dir": str(pdir), "agent": "claude"} + + +def test_missing_store_returns_error(tmp_path): + pdir = tmp_path / "proj" + pdir.mkdir() + with patch.object(traces_cmd, "load_config", return_value=_config(pdir)): + rc = traces_cmd.run("proj") + assert rc == 1 # no mlflow.db yet + + +def test_runs_catchup_then_launches_deeplinked_quiet_viewer(tmp_path, capsys): + pdir = tmp_path / "proj" + pdir.mkdir() + (pdir / "mlflow.db").write_text("") # existence is all run() checks + + launched = {} + + def _fake_subprocess_run(cmd, env=None): + launched["cmd"] = cmd + launched["env"] = env + return MagicMock(returncode=0) + + with ( + patch.object(traces_cmd, "load_config", return_value=_config(pdir)), + patch.object( + traces_cmd, "catch_up_extraction", return_value={"traces_caught_up": 2} + ) as catchup, + patch.object(traces_cmd, "_resolve_experiment_id", return_value="42"), + patch.object(traces_cmd.subprocess, "run", _fake_subprocess_run), + ): + rc = traces_cmd.run("proj", port=5001) + + assert rc == 0 + catchup.assert_called_once() # catch-up ran before the viewer + + cmd = launched["cmd"] + assert cmd[:2] == ["mlflow", "ui"] + assert "--workers" in cmd and cmd[cmd.index("--workers") + 1] == "1" + assert cmd[cmd.index("--port") + 1] == "5001" + assert cmd[cmd.index("--backend-store-uri") + 1].startswith("sqlite:///") + # Quiet: warnings suppressed in the child env. + assert launched["env"]["PYTHONWARNINGS"] == "ignore" + + out = capsys.readouterr().out + # Deep-links to the project's Traces tab (not the empty default Runs view). + assert "http://127.0.0.1:5001/#/experiments/42/traces" in out + assert "Caught up 2" in out + + +def test_launch_survives_catchup_failure(tmp_path): + pdir = tmp_path / "proj" + pdir.mkdir() + (pdir / "mlflow.db").write_text("") + + with ( + patch.object(traces_cmd, "load_config", return_value=_config(pdir)), + patch.object( + traces_cmd, "catch_up_extraction", side_effect=RuntimeError("boom") + ), + patch.object(traces_cmd, "_resolve_experiment_id", return_value=None), + patch.object( + traces_cmd.subprocess, "run", return_value=MagicMock(returncode=0) + ), + ): + # A hiccup in catch-up must not stop the viewer from opening. + assert traces_cmd.run("proj") == 0 diff --git a/tests/test_traces_goose.py b/tests/test_traces_goose.py new file mode 100644 index 0000000..4cac896 --- /dev/null +++ b/tests/test_traces_goose.py @@ -0,0 +1,167 @@ +"""Fixture tests for the Goose SQLite translator + reader. + +Translator records mirror what ``GooseReader`` yields (``{role, content, ts}`` +with parsed content_json blocks); the reader test builds a tiny sessions.db. +""" + +import sqlite3 + +import pytest + +from dsagt.observability import MLflowSink +from dsagt.traces import GooseReader, GooseTranslator + + +def _utext(ts, text): + return {"role": "user", "ts": ts, "content": [{"type": "text", "text": text}]} + + +def _atext(ts, text): + return {"role": "assistant", "ts": ts, "content": [{"type": "text", "text": text}]} + + +def _atool(ts, name, args, tid): + return { + "role": "assistant", + "ts": ts, + "content": [ + { + "type": "toolRequest", + "id": tid, + "toolCall": {"value": {"name": name, "arguments": args}}, + } + ], + } + + +def _uresp(ts, tid, text): + return { + "role": "user", + "ts": ts, + "content": [ + { + "type": "toolResponse", + "id": tid, + "toolResult": {"value": {"content": [{"type": "text", "text": text}]}}, + } + ], + } + + +def _session(): + return [ + _utext(1000, "do the thing"), # turn 1 prompt + _atext(1001, "Let me check."), # LLM + _atool(1002, "shell", {"command": "ls"}, "t1"), # TOOL + _uresp(1003, "t1", "a.txt\nb.txt"), # result (user msg, not a prompt) + _atext(1004, "Done."), # LLM + _utext(1005, "now another"), # turn 2 prompt + _atext(1006, "OK."), # LLM + ] + + +def _translate(records=None): + return GooseTranslator().translate( + records if records is not None else _session(), + trace_id="g", + session_id="proj:s", + project="gooseproj", + ) + + +def test_turns_segment_on_user_text_not_tool_responses(): + roots = [s for s in _translate().spans if s["parent_id"] is None] + assert len(roots) == 2 # the toolResponse user message is NOT a turn + assert roots[0]["attributes"]["prompt"] == "do the thing" + assert roots[1]["attributes"]["prompt"] == "now another" + + +def test_span_layout_and_tool_extraction(): + trace = _translate() + roots = [s for s in trace.spans if s["parent_id"] is None] + t1 = [ + (s["kind"], s["name"]) + for s in trace.spans + if s["parent_id"] == roots[0]["span_id"] + ] + assert t1 == [ + ("LLM", "llm"), + ("TOOL", "tool_shell"), + ("LLM", "llm"), + ] + tool = next(s for s in trace.spans if s["kind"] == "TOOL") + assert tool["attributes"]["input"] == {"command": "ls"} + assert ( + tool["attributes"]["result"] == "a.txt\nb.txt" + ) # toolResult.value.content text + assert roots[0]["attributes"]["response"] == "Done." + + +def test_empty_is_none(): + assert _translate([]) is None + assert _translate([_atext(1, "no user prompt")]) is None + + +# --- reader over a tiny sessions.db --- + + +def _make_db(path): + con = sqlite3.connect(path) + con.execute("CREATE TABLE sessions(id TEXT, working_dir TEXT, updated_at INTEGER)") + con.execute( + "CREATE TABLE messages(id INTEGER PRIMARY KEY, session_id TEXT, " + "role TEXT, content_json TEXT, created_timestamp INTEGER)" + ) + con.execute("INSERT INTO sessions VALUES('s_old', '/proj', 100)") + con.execute( + "INSERT INTO sessions VALUES('s_new', '/proj', 200)" + ) # newest for /proj + con.execute("INSERT INTO sessions VALUES('s_other', '/elsewhere', 300)") + rows = [ + ("s_new", "user", '[{"type":"text","text":"hi"}]', 1000), + ("s_new", "assistant", '[{"type":"text","text":"hello"}]', 1001), + ("s_old", "user", '[{"type":"text","text":"old session"}]', 50), + ("s_other", "user", '[{"type":"text","text":"other dir"}]', 5), + ] + con.executemany( + "INSERT INTO messages(session_id, role, content_json, created_timestamp) " + "VALUES(?,?,?,?)", + rows, + ) + con.commit() + con.close() + + +def test_reader_picks_newest_session_for_the_project_dir(tmp_path): + db = tmp_path / "sessions.db" + _make_db(db) + records = GooseReader("/proj", db_path=db).read() + assert [r["content"][0]["text"] for r in records] == ["hi", "hello"] # s_new only + + +def test_reader_missing_db_is_empty(tmp_path): + records = GooseReader("/proj", db_path=tmp_path / "nope.db").read() + assert records == [] + + +@pytest.fixture +def mlflow_sqlite(tmp_path, monkeypatch): + import mlflow + + monkeypatch.chdir(tmp_path) + uri = f"sqlite:///{tmp_path / 'mlflow.db'}" + mlflow.set_tracking_uri(uri) + mlflow.set_experiment("gooseproj") + return uri + + +def test_end_to_end_through_the_sink(mlflow_sqlite): + import mlflow + + ids = MLflowSink(mlflow_sqlite, "gooseproj").write(_translate()) + assert len(ids) == 2 + exp = mlflow.get_experiment_by_name("gooseproj") + traces = mlflow.search_traces( + locations=[exp.experiment_id], return_type="list" + ) + assert len(traces) == 2 diff --git a/tests/test_traces_opencode.py b/tests/test_traces_opencode.py new file mode 100644 index 0000000..93e4eaa --- /dev/null +++ b/tests/test_traces_opencode.py @@ -0,0 +1,170 @@ +"""Fixture tests for the opencode SQLite translator + reader.""" + +import json +import sqlite3 + +import pytest + +from dsagt.observability import MLflowSink +from dsagt.traces import OpenCodeReader, OpenCodeTranslator + + +def _part(role, ts, data, model=None): + return {"role": role, "ts": ts, "model": model, "data": data} + + +def _session(): + return [ + _part("user", 1000.0, {"type": "text", "text": "do the thing"}), # turn 1 + _part("assistant", 1001.0, {"type": "step-start"}), # skip + _part("assistant", 1002.0, {"type": "text", "text": "Working."}, "m"), # LLM + _part( + "assistant", + 1003.0, + { + "type": "tool", + "tool": "shell", + "callID": "c1", + "state": {"input": {"cmd": "ls"}, "output": "a.txt"}, + }, + ), # TOOL (call + result in one part) + _part("assistant", 1004.0, {"type": "step-finish", "tokens": {}}), # skip + _part("assistant", 1005.0, {"type": "text", "text": "Done."}, "m"), # LLM + _part("user", 1006.0, {"type": "text", "text": "next"}), # turn 2 + _part("assistant", 1007.0, {"type": "text", "text": "OK."}, "m"), # LLM + ] + + +def _translate(records=None): + return OpenCodeTranslator().translate( + records if records is not None else _session(), + trace_id="oc", + session_id="proj:s", + project="ocproj", + ) + + +def test_two_turns_with_step_parts_skipped(): + trace = _translate() + roots = [s for s in trace.spans if s["parent_id"] is None] + assert len(roots) == 2 + assert roots[0]["attributes"]["prompt"] == "do the thing" + assert roots[1]["attributes"]["prompt"] == "next" + t1 = [ + (s["kind"], s["name"]) + for s in trace.spans + if s["parent_id"] == roots[0]["span_id"] + ] + assert t1 == [ + ("LLM", "llm"), + ("TOOL", "tool_shell"), + ("LLM", "llm"), + ] + assert roots[0]["attributes"]["response"] == "Done." + + +def test_tool_call_and_result_from_one_part(): + tool = next(s for s in _translate().spans if s["kind"] == "TOOL") + assert tool["attributes"]["input"] == {"cmd": "ls"} + assert tool["attributes"]["result"] == "a.txt" + + +def test_model_carried_onto_llm_spans(): + llm = next(s for s in _translate().spans if s["kind"] == "LLM") + assert llm["model"] == "m" + + +def test_empty_is_none(): + assert _translate([]) is None + assert ( + _translate([_part("assistant", 1.0, {"type": "text", "text": "no prompt"})]) + is None + ) + + +# --- reader over a tiny opencode.db --- + + +def _make_db(path): + con = sqlite3.connect(path) + con.execute("CREATE TABLE session(id TEXT, directory TEXT, time_updated INTEGER)") + con.execute("CREATE TABLE message(id TEXT, session_id TEXT, data TEXT)") + con.execute( + "CREATE TABLE part(id TEXT PRIMARY KEY, message_id TEXT, session_id TEXT, " + "time_created INTEGER, data TEXT)" + ) + con.execute("INSERT INTO session VALUES('s_new','/proj',200)") + con.execute("INSERT INTO session VALUES('s_old','/proj',100)") + con.execute("INSERT INTO session VALUES('s_other','/elsewhere',300)") + con.executemany( + "INSERT INTO message VALUES(?,?,?)", + [ + ("m1", "s_new", json.dumps({"role": "user", "model": {"modelID": "x"}})), + ( + "m2", + "s_new", + json.dumps({"role": "assistant", "model": {"modelID": "x"}}), + ), + ("mo", "s_old", json.dumps({"role": "user"})), + ], + ) + con.executemany( + "INSERT INTO part VALUES(?,?,?,?,?)", + [ + ( + "p1", + "m1", + "s_new", + 1777819990000, + json.dumps({"type": "text", "text": "hi"}), + ), + ( + "p2", + "m2", + "s_new", + 1777819991000, + json.dumps({"type": "text", "text": "hello"}), + ), + ("po", "mo", "s_old", 1, json.dumps({"type": "text", "text": "old"})), + ], + ) + con.commit() + con.close() + + +def test_reader_picks_newest_session_and_converts_ms(tmp_path): + db = tmp_path / "opencode.db" + _make_db(db) + parts = OpenCodeReader("/proj", db_path=db).read() + assert [p["role"] for p in parts] == ["user", "assistant"] # s_new only + assert [p["data"]["text"] for p in parts] == ["hi", "hello"] + assert parts[0]["ts"] == 1777819990.0 # ms → s + assert parts[0]["model"] == "x" + + +def test_reader_missing_db_is_empty(tmp_path): + parts = OpenCodeReader("/proj", db_path=tmp_path / "nope.db").read() + assert parts == [] + + +@pytest.fixture +def mlflow_sqlite(tmp_path, monkeypatch): + import mlflow + + monkeypatch.chdir(tmp_path) + uri = f"sqlite:///{tmp_path / 'mlflow.db'}" + mlflow.set_tracking_uri(uri) + mlflow.set_experiment("ocproj") + return uri + + +def test_end_to_end_through_the_sink(mlflow_sqlite): + import mlflow + + ids = MLflowSink(mlflow_sqlite, "ocproj").write(_translate()) + assert len(ids) == 2 + exp = mlflow.get_experiment_by_name("ocproj") + traces = mlflow.search_traces( + locations=[exp.experiment_id], return_type="list" + ) + assert len(traces) == 2 diff --git a/use_cases/aidrin_readiness/aidrin_full_tour_demo.md b/use_cases/aidrin_readiness/aidrin_full_tour_demo.md index 6557dc2..e8715d0 100644 --- a/use_cases/aidrin_readiness/aidrin_full_tour_demo.md +++ b/use_cases/aidrin_readiness/aidrin_full_tour_demo.md @@ -1,5 +1,10 @@ # DSAgt Demo: Full AIDRIN Feature Tour +> **Estimated time:** ~30–40 minutes — this is a long-form tour, not a quick +> demo. Most of the time is the one-time AIDRIN build (clone + `pip install -e` +> in a Python 3.10 venv) plus the agent issuing ~15 separate metric runs. The +> dataset ships with AIDRIN (no large download). + This guide drives **every** [AIDRIN](https://github.com/idtlab/AIDRIN) (AI Data Readiness Inspector) metric through DSAgt on a single tabular dataset — exercising all 15 metrics across all four categories, with full execution provenance. It is the companion to the @@ -22,9 +27,9 @@ prediction **target** (`income`). ## Prerequisites -- DSAgt installed (`uv sync --all-groups`) and an agent platform installed (e.g. `claude`). -- `dsagt_config.yaml` configured (the default local embedder needs no API key; the agent needs its - own LLM credentials). +- DSAgt installed (`uv sync --all-groups`) and an agent platform installed and **already + authenticated** (BYOA — your agent talks to its own LLM provider; dsagt writes no + credentials). The default local embedder needs no API key. - **AIDRIN** installed from its `develop` branch in its own Python 3.10 virtual environment. - Git installed. (No large download — the dataset ships with AIDRIN.) @@ -40,7 +45,6 @@ AIDRIN_BIN="$(pwd)/aidrin-venv/bin/aidrin"; echo "$AIDRIN_BIN" deactivate dsagt init aidrin-tour --agent claude -# Edit ~/dsagt-projects/aidrin-tour/dsagt_config.yaml — set the agent's LLM credentials. PROJ=~/dsagt-projects/aidrin-tour mkdir -p "$PROJ/data" cp AIDRIN/examples/sample_data/csv/adult.csv "$PROJ/data/" @@ -54,18 +58,18 @@ Paste these prompts one at a time (substitute the absolute `$AIDRIN_BIN` path). ### 1. Register the AIDRIN CLI ```text -Register a data-readiness CLI named aidrin into the tool registry. The executable is at +Register a data-readiness CLI named aidrin into the code registry. The executable is at . Run " --help" and " list" to discover its subcommands and the -15 metrics, then save a tool spec named aidrin describing the run/batch/data-quality subcommands +15 metrics, then save a code spec named aidrin describing the run/batch/data-quality subcommands and their positional arguments. ``` -**Verify:** `Search the registry for the aidrin data-readiness tool.` +**Verify:** `Search the registry for the aidrin data-readiness code.` ### 2. Run all 15 metrics through `dsagt-run` ```text -Using the registry aidrin tool, run AIDRIN's full readiness assessment on data/adult.csv, executing +Using the registry aidrin code, run AIDRIN's full readiness assessment on data/adult.csv, executing every metric through dsagt-run so each is recorded. Cover all four categories: (1) data-quality: completeness, duplicity, outliers; (2) impact-of-data-on-AI: correlations on "age,education.num,sex,race", and feature-relevance with @@ -129,7 +133,7 @@ quasi-identifiers — bin or suppress before sharing. ```text Write an aidrin batch config (YAML) that runs completeness, class-imbalance, statistical-rates, and representation-rate on data/adult.csv with target income and sensitive attribute sex, then run it -through the registry aidrin tool. +through the registry aidrin code. ``` Batch config keys: `file-path`, `file-type`, `metrics`, `target-column`, @@ -153,7 +157,7 @@ Reconstruct the full readiness assessment you just ran from the execution record ## Post-Conditions -1. Tool registry contains the `aidrin` spec (`tools/aidrin.md`). +1. Code registry contains the `aidrin` spec (`codes/aidrin/SKILL.md`). 2. `trace_archive/` holds one provenance record per metric run (15 from step 2). 3. Results span all four categories, with the gender-fairness gap and the `k = 1` / `l = 1` re-identification risks identified. @@ -165,19 +169,21 @@ Reconstruct the full readiness assessment you just ran from the execution record | DSAgt Capability | Steps | |------------------|-------| -| External-CLI registration (`save_tool_spec`) | 1 | +| External-CLI registration (`save_code_spec`) | 1 | | Registry search | 1 (Verify) | -| Tool execution with provenance (`dsagt-run` → `trace_archive/`) | 2 | +| Code execution with provenance (`dsagt-run` → `trace_archive/`) | 2 | | Full-suite (15-metric) orchestration | 2 | | Multi-metric / batch execution | 3 | | Skill discovery and use (datacard generation) | 4 | | Pipeline reconstruction from execution records | 5 | -| Observability (MLflow spans) | all | +| Observability (MLflow spans in the serverless `mlflow.db` store) | all | + +View the traces any time with +`mlflow ui --backend-store-uri sqlite:///$PROJ/mlflow.db`. ## Cleanup ```bash -dsagt stop aidrin-tour dsagt rm aidrin-tour -y rm -rf AIDRIN aidrin-venv ``` diff --git a/use_cases/aidrin_readiness/cryoem_readiness_demo.md b/use_cases/aidrin_readiness/cryoem_readiness_demo.md index c4eaf42..e135457 100644 --- a/use_cases/aidrin_readiness/cryoem_readiness_demo.md +++ b/use_cases/aidrin_readiness/cryoem_readiness_demo.md @@ -1,5 +1,10 @@ # DSAgt Demo: AIDRIN Readiness Gate on a Cryo-EM Pipeline +> **Estimated time:** ~30 minutes, dominated by a **~2 GB EMPIAR-10017 +> download** (from `calla.rnet.missouri.edu`) plus the one-time AIDRIN build. +> The agent portion (register + before/after assessment + datacard) is ~10 +> minutes once the data is staged. + This guide demonstrates DSAgt using [AIDRIN](https://github.com/idtlab/AIDRIN) (AI Data Readiness Inspector) as a **readiness gate** around a cryo-EM data-curation step. The agent registers the AIDRIN CLI, then runs the applicable readiness metrics **before and after** particle curation to @@ -31,9 +36,9 @@ tabular dataset where they do apply. ## Prerequisites -- DSAgt installed (`uv sync --all-groups`) and an agent platform installed (e.g. `claude`). -- `dsagt_config.yaml` configured (the default local embedder needs no API key; the agent needs its - own LLM credentials). +- DSAgt installed (`uv sync --all-groups`) and an agent platform installed and **already + authenticated** (BYOA — dsagt writes no credentials). The default local embedder needs no + API key. - **AIDRIN** installed from its `develop` branch in its own Python 3.10 virtual environment (see Setup). - ~2 GB disk for the EMPIAR-10017 cryo-EM data. @@ -85,7 +90,6 @@ deactivate ```bash dsagt init cryoem-readiness --agent claude -# Edit ~/dsagt-projects/cryoem-readiness/dsagt_config.yaml — set the agent's LLM credentials. PROJ=~/dsagt-projects/cryoem-readiness mkdir -p "$PROJ/data" cp demo_data/cryoem/cryoem_before.csv demo_data/cryoem/cryoem_after.csv "$PROJ/data/" @@ -99,19 +103,19 @@ Paste these prompts into the agent one at a time (substitute the absolute `$AIDR ### 1. Register the AIDRIN CLI ```text -Register a data-readiness CLI named aidrin into the tool registry. The executable is at +Register a data-readiness CLI named aidrin into the code registry. The executable is at . Run " --help" and " list" to discover its subcommands and the -15 metrics, then save a tool spec named aidrin describing the run/batch/data-quality subcommands +15 metrics, then save a code spec named aidrin describing the run/batch/data-quality subcommands and their positional arguments. ``` -**Verify:** `Search the registry for the aidrin data-readiness tool.` → -`~/dsagt-projects/cryoem-readiness/tools/aidrin.md` should exist. +**Verify:** `Search the registry for the aidrin data-readiness code.` → +`~/dsagt-projects/cryoem-readiness/codes/aidrin/SKILL.md` should exist. ### 2. Readiness assessment BEFORE curation ```text -Using the registry aidrin tool, assess the AI data readiness of data/cryoem_before.csv. Run these +Using the registry aidrin code, assess the AI data readiness of data/cryoem_before.csv. Run these metrics through dsagt-run: completeness; duplicity; outliers; correlations on the columns "Defocus U,Defocus V,Defocus Angle,CTF B Factor,Origin X (Ang),Origin Y (Ang)"; class-imbalance on the Class Number column; and feature-relevance with no categorical columns, those same numerical @@ -157,7 +161,7 @@ Reconstruct the readiness assessment you just ran from the execution records as ## Post-Conditions -1. Tool registry contains the `aidrin` spec (`tools/aidrin.md`). +1. Code registry contains the `aidrin` spec (`codes/aidrin/SKILL.md`). 2. `trace_archive/` holds one provenance record per metric run (before and after). 3. Before/after scores show curation reduced outliers (~0.041 → ~0.029) and class imbalance (~22.2 → ~11.1) while completeness and duplicity stayed clean. @@ -169,18 +173,20 @@ Reconstruct the readiness assessment you just ran from the execution records as | DSAgt Capability | Steps | |------------------|-------| -| External-CLI registration (`save_tool_spec`) | 1 | +| External-CLI registration (`save_code_spec`) | 1 | | Registry search | 1 (Verify) | -| Tool execution with provenance (`dsagt-run` → `trace_archive/`) | 2, 3 | +| Code execution with provenance (`dsagt-run` → `trace_archive/`) | 2, 3 | | Before/after comparison driven by the agent | 3 | | Skill discovery and use (datacard generation) | 4 | | Pipeline reconstruction from execution records | 5 | -| Observability (MLflow spans) | all | +| Observability (MLflow spans in the serverless `mlflow.db` store) | all | + +View the traces any time with +`mlflow ui --backend-store-uri sqlite:///$PROJ/mlflow.db`. ## Cleanup ```bash -dsagt stop cryoem-readiness dsagt rm cryoem-readiness -y rm -rf demo_data/cryoem AIDRIN aidrin-venv ``` diff --git a/use_cases/cryoem/cryoem_demo.md b/use_cases/cryoem/cryoem_demo.md index ec86c4e..173dc34 100644 --- a/use_cases/cryoem/cryoem_demo.md +++ b/use_cases/cryoem/cryoem_demo.md @@ -1,12 +1,18 @@ # DSAgt Demo: Cryo-EM Data Curation Pipeline -This guide documents a comprehensive DSAgt demonstration using cryo-electron microscopy (cryo-EM) data. It exercises knowledge ingestion, KB-guided pipeline design, tool registration from third-party code, cross-collection knowledge synthesis, and multi-stage pipeline execution with domain-specific evaluation. +> **Estimated time:** ~40 minutes — this is the broadest demo (7 stages) and +> the least quick. It pulls a **~2 GB EMPIAR-10017 download**, an arXiv PDF, +> and a full CryoPPP repo clone, then KB-ingests the whole repo (minutes on the +> local embedder) before any pipeline work. Best treated as an advanced, +> bring-time walkthrough rather than a 10-minute demo. + +This guide documents a comprehensive DSAgt demonstration using cryo-electron microscopy (cryo-EM) data. It exercises knowledge ingestion, KB-guided pipeline design, code registration from third-party scripts, cross-collection knowledge synthesis, and multi-stage pipeline execution with domain-specific evaluation. ## Prerequisites - DSAgt installed (`uv sync --all-groups`) -- An agent platform installed (e.g., `claude` for Claude Code) -- `test_site_config.yaml` configured with valid API keys and embedding endpoint +- An agent platform installed and **already authenticated** (e.g., `claude` for Claude Code) + — BYOA: dsagt writes no credentials. The default local embedder needs no API key. - ~2 GB disk space for the cryo-EM test data - Git installed @@ -44,7 +50,9 @@ git clone https://github.com/BioinfoMachineLearning/cryoppp.git demo_repos/cryop dsagt init cryoem-pipeline --agent claude ``` -Edit `~/dsagt-projects/cryoem-pipeline/dsagt_config.yaml` — set your API keys and embedding endpoint. +(The default local embedder needs no key. To use a hosted embedder instead, set +`embedding.backend: api` in `~/dsagt-projects/cryoem-pipeline/.dsagt/config.yaml` +and export `EMBEDDING_API_KEY` in your shell — never written to disk.) ### 5. Start the session @@ -82,19 +90,19 @@ Search the cryoppp collection for guidance on creating an AI-ready data processi The agent should return chunks describing quality metrics: CTF resolution, defocus ranges, ice thickness thresholds, and motion statistics. -### 3. Register CryoPPP processing tools +### 3. Register CryoPPP processing codes ```text -Look at the scripts in demo_repos/cryoppp/ and register any data processing or evaluation tools you find. Run --help on each script to discover its interface. +Look at the scripts in demo_repos/cryoppp/ and register any data-processing or evaluation codes you find. Run --help on each script to discover its interface. ``` **Verify:** ```text -Search the registry for cryo-EM tools. +Search the registry for cryo-EM codes. ``` -### 4. Create a quality scoring tool +### 4. Create a quality scoring code ```text Write a Python script that scores cryo-EM micrographs based on: @@ -103,10 +111,10 @@ Write a Python script that scores cryo-EM micrographs based on: - Ice thickness - Motion statistics -Use the CryoCRAB 0-7 scoring scheme described in the cryoppp collection. The script should read a metadata CSV and output a scored CSV with quality_score and quality_tier (high/medium/low) columns. Save the script under tools/code/ and register it as a tool. +Use the CryoCRAB 0-7 scoring scheme described in the cryoppp collection. The script should read a metadata CSV and output a scored CSV with quality_score and quality_tier (high/medium/low) columns. Save the script under codes//scripts/ and register it as a code. ``` -The agent should search the knowledge base, write the script, and register it via `save_tool_spec`. +The agent should search the knowledge base, write the script, and register it via `save_code_spec`. ### 5. Run the pipeline @@ -133,12 +141,12 @@ Reconstruct the pipeline from the execution records as a bash script. ## Post-Conditions 1. Knowledge base contains `cryoppp` collection with repo code, docs, and appended paper. -2. Tool registry includes CryoPPP processing tools and the quality scoring tool. +2. Code registry includes the CryoPPP processing codes and the quality-scoring code. 3. Quality-scored CSV exists with tier distribution. 4. A datacard exists for the processed dataset. 5. A reconstructed pipeline script is available. -6. Tool execution records in `trace_archive/` document the full provenance chain. -7. MLflow traces capture token usage, latency, and full request/response history. +6. Code execution records in `trace_archive/` document the full provenance chain. +7. MLflow traces (in the serverless `mlflow.db` store) capture token usage, latency, and full request/response history. ## What This Tests @@ -147,15 +155,16 @@ Reconstruct the pipeline from the execution records as a bash script. | Knowledge ingestion (folder) | 1 | | Knowledge append (single file) | 1 | | Semantic search | 2 | -| Tool discovery via registry | 3 | -| Tool registration | 3, 4 | +| Code discovery via registry | 3 | +| Code registration | 3, 4 | | KB-guided code generation | 4 | -| Tool execution with provenance | 5 | +| Code execution with provenance | 5 | | Skill discovery and use | 6 | | Pipeline reconstruction | 7 | ## Cleanup ```bash -rm -rf ~/dsagt-projects/cryoem-pipeline demo_data/cryoem demo_repos/cryoppp +dsagt rm cryoem-pipeline -y # unregisters the project and removes its dir +rm -rf demo_data/cryoem demo_repos/cryoppp ``` diff --git a/use_cases/fusion-fm/README.md b/use_cases/fusion-fm/README.md index 858a18a..14053ef 100644 --- a/use_cases/fusion-fm/README.md +++ b/use_cases/fusion-fm/README.md @@ -1,5 +1,13 @@ # Fusion Foundation Model — XGC AI Training Use Case +> **Estimated time:** ~30 minutes on an HPC login node — **not a 10-minute demo +> and data is not included.** This is an advanced, bring-your-own-data example: +> the XGC cases below are HPC-scale ADIOS2 BP5 output (up to ~1.3M mesh nodes) +> that must be supplied by the user, and the scripts require `adios2` + `torch` +> plus MATEY's `BaseCFDGraphDataset`. Point the paths below at your own XGC run. +> DSAgt drives this via the bundled skill (see **Skill** below), which registers +> and runs these scripts as codes with provenance. + **Domain:** Plasma physics — gyrokinetic turbulence simulation **Simulation code:** XGC (X-point Gyrokinetic Code) **Data format:** ADIOS2 BP5 (one directory per simulation run) @@ -7,7 +15,8 @@ ## Data -Three example simulation cases live in `example_xgc_data/`: +Three example simulation cases are referenced under `example_xgc_data/` (**not +shipped in-repo** — substitute your own XGC output directories): | Case | Machine | Nodes | nphi | Steps | Has f3d | |------|---------|-------|------|-------|---------| diff --git a/use_cases/genesis_skills/README.md b/use_cases/genesis_skills/README.md new file mode 100644 index 0000000..f52ab62 --- /dev/null +++ b/use_cases/genesis_skills/README.md @@ -0,0 +1,122 @@ +# DSAgt Demo: Genesis Skills for a Data-Curation Pipeline + +> **Estimated time:** ~10 minutes (all data is bundled and tiny; the one +> external dependency is a shallow clone of the Genesis catalog from OSTI +> GitLab — needs network access to `gitlab.osti.gov`). + +An end-to-end **data-preparation** walkthrough that flexes the skill catalog +against the **Genesis** source (OSTI GitLab). The agent pulls in the +BASE-Data/ModCon curation skills, grounds itself in domain context loaded into +the **knowledge base**, then prepares and **datacards a finished dataset**. + +The "finished product" is a small curated dataset — a CO2-methanation **catalyst +screen** (`mock_data/dataset/catalyst_screening.csv`, 8 rows) — plus the domain +docs that describe how it was produced. Everything is tiny, so the whole thing +runs in seconds with no real instruments or HPC. + +## What this demonstrates + +1. **Genesis skills, agent-facing.** `add_skill_source genesis` syncs the OSTI + GitLab catalog; `search_skills` finds the ModCon curation skills + (`generating-datacards`, `croissant-validator`); `install_skill` draws them + into the project where the agent **natively** auto-discovers them. +2. **Domain → knowledge base.** `kb_ingest` indexes the data dictionary + + measurement protocol into a project collection, so the agent describes + methods and provenance accurately (via `kb_search`) instead of guessing. +3. **Datacard for a finished product.** The installed `generating-datacards` + skill runs over the dataset — pulling field definitions, methodology, and + license from the KB-ingested domain docs — and emits a datacard, then + `croissant-validator` checks the Croissant metadata. + +## Setup + +Assumes DSAgt (`uv sync --all-groups`) and Claude Code +(`npm i -g @anthropic-ai/claude-code`) are installed, plus git **with network +access to `gitlab.osti.gov`** (the Genesis catalog clones from OSTI GitLab, not +GitHub). Embedding credentials are optional — `search_skills` / `kb_search` use +semantic search when `EMBEDDING_*` is set and fall back to a keyword scorer +otherwise (configure it for sharper relevance over the domain docs). + +```bash +dsagt init genesis-skills --agent claude # provisions the bundled codes/skills + core KB +cp -r use_cases/genesis_skills/mock_data ~/dsagt-projects/genesis-skills/mock_data +dsagt start genesis-skills +``` + +The Genesis catalog is not synced at `init` — the agent enables it in step 1. + +## Walkthrough + +Paste each prompt into Claude Code (running inside the project), one at a time. +Confirmation checks are consolidated in **Post-conditions** below. + +### 1 — Enable the Genesis source + +> Enable the "genesis" skill source so we have the GENESIS / ModCon data-curation skills available. Then tell me how many skills it indexed. + +*Expect:* `add_skill_source(source="genesis")` → a shallow clone of OSTI GitLab, +**74** skills indexed, source written to `.dsagt/config.yaml`. + +### 2 — Find and install the curation skills + +> Search the catalog for two skills — one that creates a datacard / dataset documentation for a dataset, and one that validates Croissant / JSON-LD dataset metadata — and install the best match for each into this project. + +*Expect:* `search_skills` surfaces **`generating-datacards`** and +**`croissant-validator`** → `install_skill` for each. The reply presents both +as installed and immediately usable (no restart caveat), each with a +`PROVENANCE.txt` crediting the Genesis source. + +### 3 — Ingest the domain docs into the KB + +> Ingest the domain docs under mock_data/domain/ into a new knowledge-base collection called "methanation_domain". Poll until it finishes, then tell me what's in it. + +*Expect:* `kb_ingest(folder_path="mock_data/domain", collection_name="methanation_domain")` +returns a `job_id`; the agent polls `kb_job_status` to completion, then +`kb_list_collections` shows `methanation_domain` (2 docs). + +### 4 — Retrieve domain grounding + +> Using the knowledge base, what reactor conditions were used for the CO2 conversion measurement, and what license applies to this dataset? + +*Expect:* `kb_search` over `methanation_domain` → **250 °C, 1 atm, H2:CO2 = 4:1, +GHSV 12,000**; license **CC-BY-4.0**. + +### 5 — Generate the datacard for the finished dataset + +> Use the generating-datacards skill to write a datacard for mock_data/dataset/catalyst_screening.csv. Pull the field definitions, measurement methodology, provenance, and license from the methanation_domain knowledge-base collection — don't invent them. Save it to audit/catalyst_screening_datacard.md. Then compare your sections against mock_data/expected_datacard.md and report anything missing. + +*Expect:* the agent reads the installed skill's `SKILL.md`, queries the KB, +computes basic stats from the 8-row CSV, and writes +`audit/catalyst_screening_datacard.md` covering summary / provenance / schema / +methodology / stats / limitations / license. + +### 6 — Validate the metadata + +> Use the croissant-validator skill to check the Croissant/JSON-LD metadata for this dataset (generate it from the datacard if needed), and report any schema errors. + +*Expect:* the validator skill runs and reports a clean pass or names specific +schema issues. + +## Post-conditions + +Confirm from a shell: + +```bash +dsagt info genesis-skills # KB lists skills_catalog__genesis-genesis-skills + methanation_domain +ls ~/dsagt-projects/genesis-skills/skills/ # generating-datacards croissant-validator +ls ~/dsagt-projects/genesis-skills/.claude/skills/ # both mirrored here at install time +cat ~/dsagt-projects/genesis-skills/skills/generating-datacards/PROVENANCE.txt +ls ~/dsagt-projects/genesis-skills/audit/ # catalyst_screening_datacard.md +``` + +1. The KB holds a `skills_catalog__genesis-genesis-skills` collection + (searchable via `search_skills`) **and** a `methanation_domain` document + collection (retrievable via `kb_search`). +2. `generating-datacards` and `croissant-validator` are installed into + `/skills/` and mirrored into `.claude/skills/`, each with a + `PROVENANCE.txt` crediting the Genesis source. The next Claude session + auto-invokes them natively; this session used them by reading their + `SKILL.md`. +3. `audit/catalyst_screening_datacard.md` was produced for the finished dataset, + grounded in the KB-ingested domain docs, covering the sections in + `mock_data/expected_datacard.md`. diff --git a/use_cases/genesis_skills/mock_data/dataset/catalyst_screening.csv b/use_cases/genesis_skills/mock_data/dataset/catalyst_screening.csv new file mode 100644 index 0000000..b18c4b9 --- /dev/null +++ b/use_cases/genesis_skills/mock_data/dataset/catalyst_screening.csv @@ -0,0 +1,9 @@ +sample_id,composition,calcination_temp_C,bet_surface_area_m2_g,co2_conversion_pct,ch4_selectivity_pct,test_date,operator +CAT-001,Ni/Al2O3,500,142.3,38.5,91.2,2026-03-02,jlee +CAT-002,Ni-Ce/Al2O3,500,138.7,46.1,93.8,2026-03-02,jlee +CAT-003,Ni-Ce/Al2O3,650,121.4,52.7,95.1,2026-03-03,jlee +CAT-004,Ru/TiO2,400,88.6,29.3,88.4,2026-03-03,akumar +CAT-005,Ru-K/TiO2,400,85.1,34.8,90.7,2026-03-04,akumar +CAT-006,Co/SiO2,550,210.9,41.2,84.6,2026-03-04,akumar +CAT-007,Co-Mn/SiO2,550,198.3,49.5,87.9,2026-03-05,jlee +CAT-008,Fe/ZrO2,600,76.2,33.1,79.5,2026-03-05,akumar diff --git a/use_cases/genesis_skills/mock_data/domain/data_dictionary.md b/use_cases/genesis_skills/mock_data/domain/data_dictionary.md new file mode 100644 index 0000000..db37622 --- /dev/null +++ b/use_cases/genesis_skills/mock_data/domain/data_dictionary.md @@ -0,0 +1,25 @@ +# Data Dictionary — CO2 Methanation Catalyst Screening + +One row per catalyst sample tested in the fixed-bed reactor screen. + +| Column | Type | Units | Definition | +|---|---|---|---| +| `sample_id` | string | — | Unique catalyst identifier (`CAT-NNN`). Primary key. | +| `composition` | string | — | Active metal / promoter / support, e.g. `Ni-Ce/Al2O3` = Ni + Ce promoter on alumina. | +| `calcination_temp_C` | integer | °C | Calcination temperature during catalyst preparation. | +| `bet_surface_area_m2_g` | float | m²/g | BET specific surface area (N2 physisorption, see protocol). | +| `co2_conversion_pct` | float | % | CO2 converted at steady state (250 °C, 1 atm, H2:CO2 = 4:1). | +| `ch4_selectivity_pct` | float | % | Carbon selectivity to CH4 (balance: CO + higher hydrocarbons). | +| `test_date` | date | ISO 8601 | Date the reactor run was performed. | +| `operator` | string | — | Lab notebook operator initials. | + +## Controlled vocabularies + +- **Supports:** `Al2O3`, `TiO2`, `SiO2`, `ZrO2`. +- **Promoters:** `Ce`, `K`, `Mn` (optional; absent for unpromoted catalysts). + +## Quality rules + +- `co2_conversion_pct` and `ch4_selectivity_pct` are in `[0, 100]`. +- `bet_surface_area_m2_g > 0`. +- Every `sample_id` is unique and matches `^CAT-\d{3}$`. diff --git a/use_cases/genesis_skills/mock_data/domain/measurement_protocol.md b/use_cases/genesis_skills/mock_data/domain/measurement_protocol.md new file mode 100644 index 0000000..7ce2251 --- /dev/null +++ b/use_cases/genesis_skills/mock_data/domain/measurement_protocol.md @@ -0,0 +1,40 @@ +# Measurement Protocol — Methanation Screen v2 + +This protocol describes how the values in `catalyst_screening.csv` were +produced. It is domain context for the curation agent: ingest it into the +knowledge base so the agent can describe methods + provenance accurately in the +datacard without re-deriving them. + +## Catalyst preparation + +Catalysts were prepared by incipient-wetness impregnation of the support with +aqueous metal-nitrate precursors, dried at 110 °C overnight, and calcined in +static air for 4 h at the temperature recorded in `calcination_temp_C`. +Promoted samples (Ce, K, Mn) were co-impregnated. + +## BET surface area + +Specific surface area (`bet_surface_area_m2_g`) was measured by N2 physisorption +at 77 K on a Micromeritics ASAP 2020, multipoint BET in the relative-pressure +range 0.05–0.30, after degassing at 200 °C for 6 h. + +## Reactor screen + +Activity was measured in a fixed-bed quartz reactor, 50 mg catalyst diluted with +SiC, at **250 °C, 1 atm, H2:CO2 = 4:1, GHSV = 12,000 mL·g⁻¹·h⁻¹**. Steady state +was reached after 60 min on stream. Effluent was analyzed by online GC-TCD/FID. + +- `co2_conversion_pct` = (CO2_in − CO2_out) / CO2_in × 100 +- `ch4_selectivity_pct` = CH4 carbon / (total converted carbon) × 100 + +## Known limitations + +- Single-run measurements (no replicate error bars in this screen). +- Selectivity excludes trace C2+ (< 0.5 %), lumped into the balance. +- Intended for **relative** ranking of formulations, not absolute kinetics. + +## Provenance + +- Instrument campaign: `methanation-screen-2026Q1` +- Raw GC traces + reactor logs archived in the lab LIMS under the same campaign id. +- License for the curated dataset: **CC-BY-4.0**. diff --git a/use_cases/genesis_skills/mock_data/expected_datacard.md b/use_cases/genesis_skills/mock_data/expected_datacard.md new file mode 100644 index 0000000..4c57c96 --- /dev/null +++ b/use_cases/genesis_skills/mock_data/expected_datacard.md @@ -0,0 +1,43 @@ +# Expected Datacard Shape (reference target) + +This is the *shape* the generated datacard should cover — not a byte-for-byte +answer. The installed `generating-datacards` skill owns the authoritative +template; this file just lists the sections the agent should populate from the +dataset + the KB-ingested domain docs, so you can spot anything missing. + +--- + +## Dataset: CO2 Methanation Catalyst Screen (2026 Q1) + +**Summary.** 8 catalyst formulations screened for CO2 methanation activity and +CH4 selectivity in a fixed-bed reactor. One row per sample. + +### Provenance +- Campaign: `methanation-screen-2026Q1` +- Curated from `catalyst_screening.csv`; methods per the lab protocol. +- License: **CC-BY-4.0** + +### Schema +A row per `sample_id` with: `composition`, `calcination_temp_C` (°C), +`bet_surface_area_m2_g` (m²/g), `co2_conversion_pct` (%), `ch4_selectivity_pct` +(%), `test_date`, `operator`. (Field units + definitions pulled from the data +dictionary.) + +### Collection methodology +- Catalysts: incipient-wetness impregnation, calcined 4 h in static air. +- BET: N2 physisorption, 77 K, multipoint (P/P0 0.05–0.30). +- Activity: 250 °C, 1 atm, H2:CO2 = 4:1, GHSV 12,000 mL·g⁻¹·h⁻¹, steady state at 60 min. + +### Descriptive statistics (computed from the data) +- Rows: 8; samples unique; no missing values. +- `co2_conversion_pct`: range ~29–53 %. +- `ch4_selectivity_pct`: range ~80–95 %. + +### Limitations / caveats +- Single-run (no replicate error bars). +- Relative ranking, not absolute kinetics. +- Selectivity excludes trace C2+ (< 0.5 %). + +### Quality checks +- `sample_id` matches `^CAT-\d{3}$`, unique. +- Percentages in `[0, 100]`; `bet_surface_area_m2_g > 0`. diff --git a/use_cases/isaac_skills_demo/README.md b/use_cases/isaac_skills_demo/README.md new file mode 100644 index 0000000..8daf6e5 --- /dev/null +++ b/use_cases/isaac_skills_demo/README.md @@ -0,0 +1,189 @@ +# DSAgt Demo: Skill-Driven VASP → ISAAC Conversion + +> **Estimated time:** ~15 minutes — the agent flow runs in seconds on the +> bundled few-KB mock data; the one real cost is a one-time +> `pip install pymatgen` (needed for step 6) plus a shallow clone of the +> K-Dense catalog from GitHub. + +A lightweight mock of the [`isaac_vasp`](../isaac_vasp/) workflow, built to **vet +the skill-management feature**. It follows the same arc as `isaac_vasp` — install +the **pymatgen** skill, author a `vasp-to-isaac` converter that parses VASP output +**with pymatgen**, and emit an ISAAC record — but the agent **discovers, syncs, +installs, and authors** those skills itself, surfaced through Claude Code's +*native* skill discovery. It uses tiny **mock VASP outputs** (`mock_data/`, a few +KB), so the whole thing runs in seconds with no DFT, no NERSC, and no 32 MB +OUTCAR files — yet the parsing is the real `pymatgen.io.vasp`, not a hand-rolled +stand-in. + +## What this demonstrates + +- **`list_skill_sources`** — the agent discovers what external sources it can + pull from (curated names + arbitrary git URLs) and which are synced. +- **`add_skill_source`** — the agent syncs a source (here: K-Dense + `k-dense-ai`, 140+ skills) into a searchable catalog that is + **not** loaded into context; with the single `dsagt-server` it's searchable + immediately, no restart. +- **`search_skills`** + **`install_skill`** — find a catalog skill (hits marked + `[catalog]`) and draw it into the project + Claude's native `.claude/skills/`. +- **`skill-creator`** — the bundled meta-skill scaffolds a new `vasp-to-isaac` + skill (from the Anthropic template) whose converter *uses the installed + pymatgen skill* to parse the VASP output — install-then-build, not install-and-ignore. +- **Native mirror** — installed + bundled skills appear under + `.claude/skills//` (tracked by `.dsagt-managed.json`), so Claude + auto-invokes them with no MCP round-trip. + +The two tiers in one sentence: **catalog = searchable but not in context; +installed = native and auto-invoked.** + +## Setup + +Assumes DSAgt (`uv sync --all-groups`) and Claude Code +(`npm i -g @anthropic-ai/claude-code`) are installed, plus git. Embedding +credentials are optional — `search_skills` uses semantic search when +`EMBEDDING_*` is set and falls back to a keyword scorer otherwise (configure it +for sharper relevance). + +```bash +dsagt init isaac-skills-demo --agent claude --exclude genesis # core KB + bundled codes/skills, but NO catalog synced +cp -r use_cases/isaac_skills_demo/mock_data ~/dsagt-projects/isaac-skills-demo/mock_data +dsagt start isaac-skills-demo # mirrors the bundled skill-creator into .claude/skills/ before launch +``` + +The project starts with **no external catalog synced** — that's deliberate (the +`--exclude genesis` drops the default catalog): the walkthrough has the agent +*discover, sync, and search* it from inside the session. The single +`dsagt-server` owns the KB, so a source the agent syncs mid-session is +**immediately searchable, no restart**. + +> To instead pre-sync a source at init, pass `--include` with a source name +> (e.g. `dsagt init … --include k-dense-ai`); `init` provisions it into the KB +> and step 3 below becomes an idempotent refresh. + +## Walkthrough + +Paste each prompt into Claude Code (running inside the project), one at a time. +The arc: **see what you have → find more → sync a source → install the relevant +skill → author a new one → run it.** + +### 1 — What do we have? (native discovery) +> Do you have a skill available for scaffolding new skills? Name it and give me a one-line summary of what it does. + +*Expect:* Claude names **`skill-creator`** and summarizes it — discovered +**natively, with no MCP call** and no file digging. That's the mirror working: +`dsagt start` copied the bundled `skill-creator` into `.claude/skills/`, so Claude +sees its name + description like any native skill (and loads the full `SKILL.md` +only when the skill is invoked — progressive disclosure). `search_skills` is for +the not-yet-installed *catalog* only, so it should not fire here. You confirm +*which* skills dsagt placed from a shell in step 7 (`cat .dsagt-managed.json`) — +that manifest is dsagt's internal mirror bookkeeping, not something the agent +reads. + +### 2 — Where can we find more skills? +> Where can I get more skills from? List the skill sources you can pull from and which are already synced. + +*Expect:* `list_skill_sources` → the known sources (`k-dense-ai`, `anthropic`, +`antigravity`, `composio`, `genesis`) with URLs, each flagged **available, not +synced** (nothing is synced yet on an `--exclude genesis` setup). + +### 3 — Sync skills from an external repo +> Sync the "k-dense-ai" source so we can search its catalog. + +*Expect:* `add_skill_source(source="k-dense-ai")` → a shallow clone of K-Dense +`scientific-agent-skills`, ~140 skills indexed into +`skills_catalog__k-dense-ai-scientific-agent-skills`, source persisted to +`.dsagt/config.yaml`. Because it's one `dsagt-server`, the catalog is searchable +**immediately** — the next prompt can hit it with no restart. + +### 4 — Add the relevant skill +> Search the catalog for a skill that helps parse VASP output with pymatgen, then install the most relevant one into this project. + +*Expect:* `search_skills` (catalog hits tagged `[catalog · install_skill to add]`, +`pymatgen` at/near the top) → `install_skill(skill_name="pymatgen")`; the reply +notes it'll be native after the next start. The installed `pymatgen` skill carries +the reference docs (`pymatgen.io.vasp.Incar` / `Poscar` / `Outcar`) the converter +uses next. **Verify** the skill dir (with any `scripts/`/`references/`) landed: + +```bash +ls ~/dsagt-projects/isaac-skills-demo/skills/ +``` + +### 5 — Create the converter skill with skill-creator +This mirrors `isaac_vasp`'s `vasp-to-isaac` skill — the lightweight version reads +the small mock directory, but does the parsing with **real pymatgen**, the same +way the full workflow does (just without the heavy `vasprun.xml`). + +> Use the skill-creator skill to author a new project skill named "vasp-to-isaac". Following the pymatgen skill you just installed, its converter should use `pymatgen.io.vasp` — `Incar.from_file` (ENCUT, NSW, ISPIN, LDAUU), `Poscar.from_file` (formula, atom counts), and `Outcar` (final energy, energy(sigma->0), total magnetization, max force) — to read a VASP slab calc directory and emit an ISAAC-style JSON record. The mock has no vasprun.xml, so take energy/forces from the OUTCAR. Target the shape in mock_data/expected_isaac_record.json. Save it with save_skill. + +*Expect:* the agent reads `skill-creator`'s template + the `pymatgen` skill's IO +docs, then `save_skill` writes `/skills/vasp-to-isaac/` whose script +imports `pymatgen.io.vasp` (not a hand-rolled regex parser). + +### 6 — Run the converter on the mock data +> Invoke the vasp-to-isaac skill on mock_data/mock_slab/ and write the result to audit/mock_slab_isaac.json. Then diff its structure and values against mock_data/expected_isaac_record.json and report any differences. + +*Expect:* pymatgen parses the mock dir and the agent writes +`audit/mock_slab_isaac.json` with the key fields **pymatgen extracted** — final +energy ≈ -132.8421 eV (`Outcar.final_energy`), 12 atoms (`Poscar`), ENCUT 520 / +NSW 50 (`Incar`), total mag ≈ 8.0123 (`Outcar.total_mag`) — matching the +reference. (`pymatgen` must be importable in the project env — see Notes.) + +### 7 — Inspect the tiers (run in a shell) +```bash +dsagt info isaac-skills-demo # KB shows the k-dense-ai catalog collection +ls ~/dsagt-projects/isaac-skills-demo/skills/ # installed: pymatgen + vasp-to-isaac +ls ~/dsagt-projects/isaac-skills-demo/.claude/skills/ +cat ~/dsagt-projects/isaac-skills-demo/.claude/skills/.dsagt-managed.json +``` + +The manifest lists only the skills **dsagt** placed; any skill you hand-create +under `.claude/skills/` is never touched. Restart Claude +(`dsagt start isaac-skills-demo` again) to pick up newly-mirrored skills as +native auto-invoked skills. + +## Post-Conditions + +1. The KB holds the `skills_catalog__k-dense-ai-scientific-agent-skills` + collection, **synced in-session by the agent** (step 3), searchable via + `search_skills` but absent from Claude's context. +2. The `pymatgen` catalog skill was installed into `/skills/` and + mirrored into `.claude/skills/`. +3. A new `vasp-to-isaac` skill, authored via `skill-creator` and parsing with + **pymatgen** (`pymatgen.io.vasp`), exists and is native-discoverable. +4. `audit/mock_slab_isaac.json` was produced from the mock VASP directory by + pymatgen and matches the ISAAC shape + values. +5. `.claude/skills/.dsagt-managed.json` tracks exactly the dsagt-placed skills. + +## Cleanup + +```bash +dsagt rm isaac-skills-demo # add -y to skip the prompt +``` + +The shared catalog cache lives at `~/dsagt-projects/.skill_sources/` and is +reused across projects; delete it to force a fresh clone. + +## Notes + +- **pymatgen must be importable** in the project env to run step 6 (the converter + uses `pymatgen.io.vasp`, exactly as `isaac_vasp` does). Install it once — + `pip install pymatgen` (or `uv pip install pymatgen`) into the same environment + `dsagt` runs in; the installed `pymatgen` skill's `references/` document this. + This is the demo's one real dependency — "lightweight" is about the data, not + avoiding pymatgen. +- `mock_data/` is intentionally tiny and **not** real DFT output, but it is *valid + VASP format*: the INCAR/POSCAR parse cleanly, and the OUTCAR stub keeps exactly + the lines pymatgen's `Outcar` reads (TOTEN, `energy(sigma->0)`, magnetization, + the force block) while omitting the ~250k-line SCF/eigenvalue blocks. No + `vasprun.xml` (it'd be large), so the converter takes energy/forces from OUTCAR. +- With the default **local** embedder (`bge-small`), absolute search scores are + low (~0.03) because short queries under-score long SKILL.md text — *ranking* is + still correct (pymatgen #1). Switch `embedding.backend` to `api` for sharper + relevance. With no embedder at all, `search_skills` falls back to keyword + scoring; `install_skill` and the native mirror are pure filesystem ops. +- Add **more** sources the same way — ask the agent to "enable the anthropic + source" (or `antigravity`, `composio`, `genesis`, or any + `https://github.com/owner/repo`), which fires `add_skill_source`. Each lands + in its own `skills_catalog__*` collection. +- Sister demo: [`genesis_skills`](../genesis_skills/) flexes the same catalog → + install → native loop plus KB domain ingest and datacard generation, against + the Genesis (OSTI GitLab) source. diff --git a/use_cases/isaac_skills_demo/mock_data/expected_isaac_record.json b/use_cases/isaac_skills_demo/mock_data/expected_isaac_record.json new file mode 100644 index 0000000..5d03192 --- /dev/null +++ b/use_cases/isaac_skills_demo/mock_data/expected_isaac_record.json @@ -0,0 +1,63 @@ +{ + "isaac_record_version": "1.05", + "record_id": "mock-iro2-110-slab-0001", + "record_type": "dft_calculation", + "record_domain": "materials", + "source_type": "MOCK", + "_note": "Reference SHAPE/values for the isaac_skills_demo. Derived from the tiny mock_slab/ stub, NOT a real DFT run.", + "timestamps": { + "created_utc": "2026-06-01T10:00:00Z" + }, + "sample": { + "material": { + "name": "IrO2 (110) slab", + "formula": "IrO2", + "provenance": "mock" + }, + "sample_form": "surface_slab" + }, + "system": { + "domain": "materials", + "technique": "DFT", + "instrument": { + "instrument_type": "compute", + "instrument_name": "VASP", + "vendor_or_project": "VASP" + }, + "configuration": { + "code_version": "6.3.2", + "compute_architecture": "cpu", + "cores": null + } + }, + "computation": { + "method": { + "family": "DFT", + "functional_class": "GGA", + "functional_name": "PBE", + "pseudopotential": "PAW", + "cutoff_eV": 520.0, + "spin_treatment": "collinear", + "hubbard_u": {"Ir": 4.0} + }, + "relaxation": { + "is_relaxation": true, + "nsw": 50, + "converged": true, + "ionic_steps": 50 + } + }, + "results": { + "total_energy_eV": -132.84210000, + "energy_sigma0_eV": -132.84210000, + "n_atoms": 12, + "n_species": {"Ir": 4, "O": 8}, + "total_magnetization_muB": 8.0123, + "max_residual_force_eV_per_A": 0.011 + }, + "assets": [ + {"name": "POSCAR", "role": "structure_input"}, + {"name": "INCAR", "role": "calc_parameters"}, + {"name": "OUTCAR", "role": "calc_output"} + ] +} diff --git a/use_cases/isaac_skills_demo/mock_data/mock_slab/INCAR b/use_cases/isaac_skills_demo/mock_data/mock_slab/INCAR new file mode 100644 index 0000000..4c80632 --- /dev/null +++ b/use_cases/isaac_skills_demo/mock_data/mock_slab/INCAR @@ -0,0 +1,18 @@ +# MOCK INCAR — slab ionic relaxation (NSW > 0 marks this as a slab calc) +SYSTEM = IrO2(110) mock slab +ISTART = 0 +ICHARG = 2 +ENCUT = 520 +ISMEAR = 0 +SIGMA = 0.05 +IBRION = 2 +NSW = 50 +ISIF = 2 +EDIFF = 1E-5 +EDIFFG = -0.02 +ISPIN = 2 +LDAU = .TRUE. +LDAUTYPE = 2 +LDAUL = 2 -1 +LDAUU = 4.0 0.0 +GGA = PE diff --git a/use_cases/isaac_skills_demo/mock_data/mock_slab/OUTCAR b/use_cases/isaac_skills_demo/mock_data/mock_slab/OUTCAR new file mode 100644 index 0000000..4640af5 --- /dev/null +++ b/use_cases/isaac_skills_demo/mock_data/mock_slab/OUTCAR @@ -0,0 +1,42 @@ + MOCK OUTCAR — heavily truncated stub for the skills demo. NOT real VASP output. + Only the few lines a converter typically greps for are kept; the SCF/eigenvalue + blocks that make a real OUTCAR ~250k lines are omitted on purpose. + + vasp.6.3.2 mock build + executed on LinuxIFC date 2026.06.01 10:00:00 + + INCAR: + ENCUT = 520.0 + ISPIN = 2 + NSW = 50 + LDAUU = 4.000 0.000 + + energy without entropy= -123.45678901 energy(sigma->0) = -123.40000000 + ... + FREE ENERGIE OF THE ION-ELECTRON SYSTEM (eV) + --------------------------------------------------- + free energy TOTEN = -132.10000000 eV (ionic step 1) + + energy without entropy= -131.98000000 energy(sigma->0) = -131.99000000 + + FREE ENERGIE OF THE ION-ELECTRON SYSTEM (eV) + --------------------------------------------------- + free energy TOTEN = -132.84210000 eV (ionic step 50, converged) + + energy without entropy= -132.80000000 energy(sigma->0) = -132.84210000 + + POSITION TOTAL-FORCE (eV/Angst) + ----------------------------------------------------------------------------------- + 0.00000 0.00000 7.04000 0.000000 0.000000 -0.004000 + 3.19250 3.24900 7.04000 0.000000 0.000000 0.003000 + 0.00000 0.00000 9.90000 0.000000 0.000000 0.011000 + ----------------------------------------------------------------------------------- + + magnetization (x) + number of electron 192.0000000 magnetization 8.0123000 + + General timing and accounting informations for this job: + ======================================================== + Total CPU time used (sec): 4210.123 + Elapsed time (sec): 1130.456 + reached required accuracy - stopping structural energy minimisation diff --git a/use_cases/isaac_skills_demo/mock_data/mock_slab/POSCAR b/use_cases/isaac_skills_demo/mock_data/mock_slab/POSCAR new file mode 100644 index 0000000..5101962 --- /dev/null +++ b/use_cases/isaac_skills_demo/mock_data/mock_slab/POSCAR @@ -0,0 +1,21 @@ +IrO2 (110) mock slab [MOCK — not real DFT input] +1.0 + 6.3850000000 0.0000000000 0.0000000000 + 0.0000000000 6.4980000000 0.0000000000 + 0.0000000000 0.0000000000 22.0000000000 + Ir O + 4 8 +Selective dynamics +Direct + 0.0000 0.0000 0.3200 F F F + 0.5000 0.5000 0.3200 F F F + 0.0000 0.0000 0.4500 T T T + 0.5000 0.5000 0.4500 T T T + 0.2500 0.2500 0.3850 T T T + 0.7500 0.7500 0.3850 T T T + 0.2500 0.7500 0.3850 T T T + 0.7500 0.2500 0.3850 T T T + 0.2500 0.2500 0.5100 T T T + 0.7500 0.7500 0.5100 T T T + 0.2500 0.7500 0.5100 T T T + 0.7500 0.2500 0.5100 T T T diff --git a/use_cases/isaac_vasp/README.md b/use_cases/isaac_vasp/README.md index 2199ade..02c141f 100644 --- a/use_cases/isaac_vasp/README.md +++ b/use_cases/isaac_vasp/README.md @@ -1,12 +1,103 @@ -# VASP -> ISAAC AI Ready Data conversion +# DSAgt Demo: VASP → ISAAC AI-Ready Record -**Goal:** Generate ISAAC AI-Ready Records (https://github.com/ISAAC-DOE/isaac-ai-ready-record) +> **Estimated time:** ~15 minutes — the NEB fixture data is bundled in this +> folder, so the only real cost is a one-time `pip install pymatgen`. No DFT run, +> no HPC. -See schema: https://github.com/ISAAC-DOE/isaac-ai-ready-record/blob/main/schema/isaac_record_v1.json +**Goal:** register a converter as a DSAgt **code**, then have the agent run it +through `dsagt-run` to turn a VASP nudged-elastic-band (NEB) calculation into an +[ISAAC AI-Ready Record](https://github.com/ISAAC-DOE/isaac-ai-ready-record) — +with the execution captured in `trace_archive/`. + +- Converter: [`vasp_neb_to_isaac.py`](vasp_neb_to_isaac.py) parses the NEB + `OUTCAR`s with `pymatgen.io.vasp` and emits a v1.05 record. +- Data: the `neb/00..04/` subdirs are a small NEB fixture (5 images) copied from + the pymatgen test suite — bundled here so the demo runs without a DFT code. +- Schema: [isaac_record_v1.json](https://github.com/ISAAC-DOE/isaac-ai-ready-record/blob/main/schema/isaac_record_v1.json). + +For the **skill-management** counterpart of this workflow (the agent discovers, +installs, and authors the converter as a skill on tiny mock data), see the +sister demo [`isaac_skills_demo`](../isaac_skills_demo/). + +## Prerequisites + +- DSAgt installed (`uv sync --all-groups`) and an agent platform installed and + **already authenticated** (BYOA — dsagt writes no credentials; the default + local embedder needs no API key). +- `pymatgen` importable in the environment `dsagt` runs in + (`uv pip install pymatgen`) — the converter uses `pymatgen.io.vasp`. + +## Setup + +```bash +uv pip install pymatgen # the converter's one real dependency +dsagt init isaac-vasp --agent claude +PROJ=~/dsagt-projects/isaac-vasp +mkdir -p "$PROJ/codes/scripts" +cp use_cases/isaac_vasp/vasp_neb_to_isaac.py "$PROJ/codes/scripts/" +cp -r use_cases/isaac_vasp/neb "$PROJ/data_neb" +dsagt start isaac-vasp +``` + +## Execution + +Paste each prompt into the agent, one at a time. + +### 1. Register the converter as a code + +```text +Register a code named vasp-neb-to-isaac. Its executable is +`python codes/scripts/vasp_neb_to_isaac.py`, which takes a positional NEB +directory argument (containing 00/, 01/, ... image subdirs) and an optional +`--output` path. Run it with `--help` first to confirm the interface, then save +the code spec with the positional `neb_dir` and the `--output` option. +``` + +**Verify:** `Search the registry for the vasp-neb-to-isaac code.` → +`$PROJ/codes/vasp-neb-to-isaac/SKILL.md` should exist. + +### 2. Run the conversion through dsagt-run + +```text +Using the registered vasp-neb-to-isaac code, convert the NEB calculation in +data_neb/ and write the ISAAC record to data_neb/isaac_neb_record.json. Use the +exact dsagt-run command from the spec so the execution is recorded. Then tell me +the record's reaction-energy / barrier fields and how many images it summarized. +``` + +**Expect:** the agent runs `dsagt-run --code vasp-neb-to-isaac -- python +codes/scripts/vasp_neb_to_isaac.py data_neb/ --output data_neb/isaac_neb_record.json`, +pymatgen parses the five OUTCARs, and a v1.05 ISAAC record lands with the +`computation` / `measurement` blocks populated (5 NEB images). Compare against the +reference [`isaac_neb_record.json`](isaac_neb_record.json) shipped here. + +### 3. Reconstruct the pipeline + +```text +Reconstruct the pipeline from the execution records as a script. +``` + +## Post-Conditions + +1. Code registry contains the `vasp-neb-to-isaac` spec (`codes/vasp-neb-to-isaac/SKILL.md`). +2. `trace_archive/` holds the conversion's provenance record. +3. `data_neb/isaac_neb_record.json` is a valid ISAAC v1.05 record matching the + shape of the bundled `isaac_neb_record.json`. +4. MLflow traces (in the serverless `mlflow.db` store) capture the run — + `mlflow ui --backend-store-uri sqlite:///$PROJ/mlflow.db`. + +## Cleanup + +```bash +dsagt rm isaac-vasp -y +``` ## Notes -- We convert example VASP outputs from: -https://github.com/materialsproject/pymatgen/tree/master/test-files/io/vasp/fixtures/neb_analysis/neb2 -- Uses "pymatgen" Skills from: https://github.com/K-Dense-AI/scientific-agent-skills - +- The `neb/` OUTCARs are public pymatgen test fixtures vendored here for a + self-contained demo. They are large (~32 MB total); a future revision may fetch + them on demand instead of shipping them in-repo. +- The bundled `skills/vasp-to-isaac/` skill is a broader **slab/bulk** converter + (a different pymatgen workflow that needs `vasprun.xml`-bearing slab/bulk data, + not the NEB fixture here). It's shipped as a reference skill; the NEB converter + above is what this folder's data actually exercises. diff --git a/use_cases/isaac_vasp/skills/vasp-to-isaac/references/cathub_organize.md b/use_cases/isaac_vasp/skills/vasp-to-isaac/references/cathub_organize.md index b39813e..be31a87 100644 --- a/use_cases/isaac_vasp/skills/vasp-to-isaac/references/cathub_organize.md +++ b/use_cases/isaac_vasp/skills/vasp-to-isaac/references/cathub_organize.md @@ -95,7 +95,7 @@ The `.db` file is a standard ASE database containing one row per structure, with Once the `.db` file is built, convert to ISAAC records: ```bash -python3 ase_slab_db_to_isaac.py /.db ./isaac_records/ --electrode-type anode +python3 ase_db_to_isaac.py /.db ./isaac_records/ --electrode-type anode ``` See `slab_workflow.md` for the full slab pipeline. @@ -107,6 +107,6 @@ See `slab_workflow.md` for the full slab pipeline. | Route | When to use | |---|---| | **cathub organize → folder2db** (this file) | Data is organized as surface reactions with gas references and adsorbate slabs. Produces CatHub-schema entries with reaction energies, `state`, `facet`, `species` key-value pairs. | -| **make_slab_ase_db.py** (direct from VASP) | Raw VASP slab directories without a CatHub reaction structure. Useful for datasets not organized around specific reactions, or when you want richer INCAR-level metadata (smearing, convergence, Hubbard-U) in the db. | +| **make_ase_db.py** (direct from VASP) | Raw VASP slab directories without a CatHub reaction structure. Useful for datasets not organized around specific reactions, or when you want richer INCAR-level metadata (smearing, convergence, Hubbard-U) in the db. | -Both routes produce an ASE SQLite `.db` file that `ase_slab_db_to_isaac.py` can read directly. +Both routes produce an ASE SQLite `.db` file that `ase_db_to_isaac.py` can read directly. diff --git a/use_cases/microbial_isolates/demoplan.md b/use_cases/microbial_isolates/demoplan.md deleted file mode 100644 index bb76fd1..0000000 --- a/use_cases/microbial_isolates/demoplan.md +++ /dev/null @@ -1,12 +0,0 @@ -# Demo Plan - -1. Initialize project: `dsagt init isolate-pipeline --agent claude` -2. Configure `dsagt_config.yaml` with API keys and embedding endpoint -3. Start session: `dsagt start isolate-pipeline` -4. Ask agent to load fastp and megahit docs into knowledge base -5. Ask agent to register fastp and megahit tools (conda bin paths) -6. Ask agent to design pipeline for preprocessing and assembling short read sequence data — give one file as example -7. After pipeline design, ask agent to run on all files in the directory -8. Ask agent to search for and use the datacard-generator skill -9. Ask agent to reconstruct the pipeline from execution records -10. Show artifacts: knowledge base collection, processed data, datacard, registered tools, pipeline script, MLflow traces diff --git a/use_cases/microbial_isolates/isolate_demo.md b/use_cases/microbial_isolates/isolate_demo.md index b5e6cee..178c5c5 100644 --- a/use_cases/microbial_isolates/isolate_demo.md +++ b/use_cases/microbial_isolates/isolate_demo.md @@ -1,12 +1,18 @@ # DSAgt Demo: Microbial Isolate Processing +> **Estimated time:** advanced / not a 10-minute demo. The isolate data is +> pulled from **NERSC** (requires an account + allocation — step 2), and +> `megahit` assembly runs minutes per sample across ~11 isolates. Treat this as +> a bring-your-own-HPC-data walkthrough; substitute your own FASTQ files for the +> NERSC path if you don't have NERSC access. + This guide documents a reproducible DSAgt demonstration for microbial isolate data processing using `fastp` and `megahit`. ## Prerequisites - DSAgt installed (`uv sync --all-groups`) -- An agent platform installed (e.g., `claude` for Claude Code, or `goose`) -- `test_site_config.yaml` configured with valid API keys and embedding endpoint +- An agent platform installed and **already authenticated** (e.g., `claude` for Claude Code, or + `goose`) — BYOA: dsagt writes no credentials. The default local embedder needs no API key. - Conda (for installing fastp and megahit) ## Setup @@ -21,7 +27,7 @@ conda run -n isolate fastp --version conda run -n isolate megahit --version ``` -Note the conda env prefix (e.g., `~/miniconda3/envs/isolate/bin/`) — you'll reference these paths when registering the tools. +Note the conda env prefix (e.g., `~/miniconda3/envs/isolate/bin/`) — you'll reference these paths when registering the codes. ### 2. Collect data @@ -45,7 +51,9 @@ git clone https://github.com/voutcn/megahit.git demo_repos/megahit dsagt init isolate-pipeline --agent claude ``` -Edit `~/dsagt-projects/isolate-pipeline/dsagt_config.yaml` — set your API keys and embedding endpoint. +(The default local embedder needs no key. To use a hosted embedder instead, set +`embedding.backend: api` in `~/dsagt-projects/isolate-pipeline/.dsagt/config.yaml` +and export `EMBEDDING_API_KEY` in your shell — never written to disk.) ### 5. Start the session @@ -53,7 +61,7 @@ Edit `~/dsagt-projects/isolate-pipeline/dsagt_config.yaml` — set your API keys dsagt start isolate-pipeline ``` -The agent launches from the project directory with MCP servers connected. Services clean up automatically when the agent exits. +The agent launches from the project directory with the MCP server connected. Serverless — there are no background services to clean up. ## Execution @@ -70,7 +78,7 @@ The collection will contain: 4) best practices for fastp and megahit: use_cases/microbial_isolates/fastp_megahit_best_practices.md ``` -### 2. Register tools +### 2. Register codes ```text Let's add /fastp to the registry @@ -113,7 +121,7 @@ The agent calls `reconstruct_pipeline` to generate a reproducible script from th ## Post-Conditions 1. Knowledge base contains collection `microbial_isolates` with all listed references indexed. -2. Tool registry includes `fastp` and `megahit` tool specs (wrapped with `dsagt-run`). +2. Code registry includes `fastp` and `megahit` code specs (wrapped with `dsagt-run`). 3. Processed output directories exist for each isolate sample. 4. For each completed sample: - Preprocessed FASTQ output exists @@ -121,8 +129,8 @@ The agent calls `reconstruct_pipeline` to generate a reproducible script from th - Assembly output exists, including `final.contigs.fa` 5. A Level 1 datacard exists for the processed dataset. 6. A reconstructed pipeline script (bash or Snakemake) is available. -7. Tool execution records in `trace_archive/` document the full provenance chain. -8. MLflow traces capture token usage, latency, and full request/response history. +7. Code execution records in `trace_archive/` document the full provenance chain. +8. MLflow traces (in the serverless `mlflow.db` store) capture token usage, latency, and full request/response history. View with `mlflow ui --backend-store-uri sqlite:///~/dsagt-projects/isolate-pipeline/mlflow.db`. ### Note diff --git a/use_cases/microbial_isolates/isolate_session.txt b/use_cases/microbial_isolates/isolate_session.txt deleted file mode 100644 index 810b53d..0000000 --- a/use_cases/microbial_isolates/isolate_session.txt +++ /dev/null @@ -1,27 +0,0 @@ -goose session \ - --with-extension 'uv run dsagt-registry-server' \ - --with-extension 'uv run dsagt-knowledge-server' - -I'd like to create a new collection in the knowledge base: microbial_isolates. -The collection will contain -1) the code package files for fastp: /Users/tuor472/Documents/dsagt_knowledge_complete/repos/fastp/ -2) the code package files for megahit: /Users/tuor472/Documents/dsagt_knowledge_complete/repos/megahit/ and -3) a short document describing a processing pipeline: /Users/tuor472/Documents/dsagt_knowledge_complete/DSAGT/demo/genomics.md -4) A document describing best practices for using fastp and -megahit: /Users/tuor472/Documents/dsagt_knowledge_complete/DSAGT/demo/fastp_megahit_best_practices.md - -Let's add /Users/tuor472/miniconda3/envs/isolate/bin/fastp to the registry -Let's add /Users/tuor472/miniconda3/envs/isolate/bin/megahit to the registry - -Okay good I have an isolate file located at -/Users/tuor472/Documents/dsagt_knowledge_complete/data/microbial_isolate/53162.2.609630.AAAGGCTAGA-GATTCAGTTA.filter-ISO.fastq.gz -Information about the dataset is contained in the Readme in that directory. I need to preprocess this file and assemble it. -fastp and megahit both have data assessment capability so we don't need to create additional tools. -megahit should be run with kmax=21 and memory=0.3 to avoid OOM on this laptop - Tell me your plan before proceeding. - -Let's run this same pipeline the rest of the fastq files at -/Users/tuor472/Documents/dsagt_knowledge_complete/data/microbial_isolate/ -We can process them one at a time. We don't need to create a script for batch processing - -Use the skill located at /Users/tuor472/.config/goose/skills/datacard-generator/ to create a datacard for the processed microbial isolate data diff --git a/use_cases/tokamak_stability/AGENTS.md b/use_cases/tokamak_stability/AGENTS.md index 1d44a78..f2f4ea1 100644 --- a/use_cases/tokamak_stability/AGENTS.md +++ b/use_cases/tokamak_stability/AGENTS.md @@ -1,8 +1,8 @@ # Agent Guide - tokamak_stability -## Tools +## Codes -This subdirectory provides python modules containing functions for interacting with HDF5 data files created by the M3D-C1 finite-element simulation code. Before using these functions, and in particular before creating CLI tools based on these functions, read the agent skill file at `skills/m3dc1-skill/SKILL.md` for calling conventions, input/output formats, and examples. +This subdirectory provides python modules containing functions for interacting with HDF5 data files created by the M3D-C1 finite-element simulation code. Before using these functions, and in particular before registering them as CLI codes, read the agent skill file at `skills/m3dc1-skill/SKILL.md` for calling conventions, input/output formats, and examples. The three python modules containing useful functions are: `m3dc1_tools.py`, `m3dc1_plots.py`, and `hdf5.py`. diff --git a/use_cases/tokamak_stability/README.md b/use_cases/tokamak_stability/README.md index e8c5c00..214b8dd 100644 --- a/use_cases/tokamak_stability/README.md +++ b/use_cases/tokamak_stability/README.md @@ -1,6 +1,12 @@ # Fusion energy use case -Here we're going to use dsagt to investigate the stability properties of a tokamak configuration. We'll be looking at linear MHD simulation data produced by the [M3D-C1](https://sites.google.com/pppl.gov/m3d-c1) unstructured-mesh finite-element code. The session below has been tested with Claude Code using Sonnet 4.6. +> **Estimated time:** advanced / not a 10-minute demo. Setup is the cost: +> building the **fusion-io** C/C++ library from source and downloading a +> Google-Drive-hosted **M3D-C1** dataset. Budget an hour+ for first-time setup; +> the agent session itself is ~15 minutes once dependencies and data are in +> place. + +Here we're going to use dsagt to investigate the stability properties of a tokamak configuration. We'll be looking at linear MHD simulation data produced by the [M3D-C1](https://sites.google.com/pppl.gov/m3d-c1) unstructured-mesh finite-element code. The session below has been tested with Claude Code. ## Dependencies @@ -40,26 +46,12 @@ Now create a new dsagt project and associated directory, here called `fusion-use dsagt init fusion-use-case --agent claude --location ~/data/ ``` -Here we're using Claude Code; see the dsagt [documentation](https://ai-modcon.github.io/dsagt/) for guidelines for using other agents. The created project directory will be a subdirectory of `~/data/`; omitting this with place your project directories in `~/dsagt-projects/`. - -We can start MLflow in the background: - -``` -dsagt mlflow fusion-use-case -``` - -This isn't strictly necessary but will allow us to more accurately recreate our work in a shell script later on. Copy and paste the MLflow export block printed by this command into this current shell. For example the block may be: +Here we're using Claude Code; see the dsagt [documentation](https://ai-modcon.github.io/dsagt/) for guidelines for using other agents. The created project directory will be a subdirectory of `~/data/`; omitting this will place your project directories in `~/dsagt-projects/`. -``` -export CLAUDE_CODE_ENABLE_TELEMETRY=1 -export CLAUDE_CODE_ENHANCED_TELEMETRY_BETA=1 -export OTEL_LOG_TOOL_DETAILS=1 -export OTEL_LOG_USER_PROMPTS=1 -export OTEL_TRACES_EXPORTER=otlp -export OTEL_LOG_RAW_API_BODIES=file:/path/to/data/fusion-use-case/api_bodies -``` - -Now enter the project directory and start the agent: +Observability is serverless — there is nothing to start. Every code execution +and agent turn is logged automatically to a SQLite MLflow store at +`~/data/fusion-use-case/mlflow.db`, which we'll use to reconstruct the session +later. Now enter the project directory and start the agent: ``` cd ~/data/fusion-use-case @@ -71,11 +63,11 @@ claude Inside the agent, enter these prompts one at a time, replacing the placeholder directory paths with the corresponding paths on your system: -- There are three python files in the `/path/to/your/dsagt/use_cases/tokamak_stability` directory containing functions for dealing with HDF5 datasets produced by the M3D-C1 code: `hdf5.py`, `m3dc1_tools.py`, and `m3dc1_plots.py`. Register these functions as tools and print a summary here. Read the `AGENTS.md` file in that directory first. +- There are three python files in the `/path/to/your/dsagt/use_cases/tokamak_stability` directory containing functions for dealing with HDF5 datasets produced by the M3D-C1 code: `hdf5.py`, `m3dc1_tools.py`, and `m3dc1_plots.py`. Register these functions as codes and print a summary here. Read the `AGENTS.md` file in that directory first. -Creating and registering the tools may take several minutes. +Creating and registering the codes may take several minutes. -- Using your tools, tell me about the data in the `/path/to/your/m3dc1_data` directory. +- Using your codes, tell me about the data in the `/path/to/your/m3dc1_data` directory. - What are the Miller parameters for this configuration? @@ -93,14 +85,13 @@ New data products will go to a `processed_data/` subdirectory of your project di - Extract the electron temperature and electron density data at t=1 and place them in an HDF5 file `electrons.h5`. -- Using the MLflow traces in the `mlflow/` directory, create a shell script that recreates this session's tool calls on a general data directory that is set at the top of the script. Save as `dsagt_session_script.sh`. +- Using the MLflow traces in the `mlflow.db` store, create a shell script that recreates this session's code executions on a general data directory that is set at the top of the script. Save as `dsagt_session_script.sh`. +(You can browse those traces any time with +`mlflow ui --backend-store-uri sqlite:///~/data/fusion-use-case/mlflow.db`.) -Now exit the agent session as usual and stop the MLflow server: - -``` -dsagt stop fusion-use-case -``` +Now exit the agent session as usual (there is no server to stop — the store is +serverless). The project name (here "fusion-use-case") is registered in `~/dsagt-projects/projects.yaml` (default location). You can unregister the project name with diff --git a/use_cases/tokamak_stability/skills/m3dc1-skill/SKILL.md b/use_cases/tokamak_stability/skills/m3dc1-skill/SKILL.md index dbd6e05..67cd089 100644 --- a/use_cases/tokamak_stability/skills/m3dc1-skill/SKILL.md +++ b/use_cases/tokamak_stability/skills/m3dc1-skill/SKILL.md @@ -1,17 +1,17 @@ --- name: m3dc1-skill -description: Use this skill when working with the python modules `m3dc1_tools.py`, `m3dc1_plots.py`, `hdf5.py` and the tools created from the functions within. Triggers when working with M3D-C1 simulation data or repackaging general HDF5 files. +description: Use this skill when working with the python modules `m3dc1_tools.py`, `m3dc1_plots.py`, `hdf5.py` and the codes created from the functions within. Triggers when working with M3D-C1 simulation data or repackaging general HDF5 files. --- # m3dc1-skill -## Creating CLI tools from M3D-C1 python modules +## Creating CLI codes from M3D-C1 python modules -Always read the relevant API guide before creating CLI tools --- see the `Reference material` section below. +Always read the relevant API guide before creating CLI codes --- see the `Reference material` section below. -IMPORTANT NOTE FOR AGENTS GENERATING NEW TOOLS FROM m3dc1_tools.py +IMPORTANT NOTE FOR AGENTS GENERATING NEW CODES FROM m3dc1_tools.py ------------------------------------------------------------------ The functions marked "requires m3dc1 + fpy" in the initial comment in `m3dc1_tools.py` call into compiled C/Fortran code that writes diagnostic @@ -19,14 +19,14 @@ messages (e.g. "deleting simulation object", "period = 6.28318") directly to the OS-level stdout file descriptor. This output bypasses Python's sys.stdout entirely and cannot be suppressed or redirected from Python. -Consequence: any CLI tool that (1) calls one of these functions and (2) prints +Consequence: any CLI code that (1) calls one of these functions and (2) prints its JSON result to stdout will produce contaminated output that cannot be parsed as JSON when captured via shell redirect. -Required pattern for any new CLI tool wrapping these functions: +Required pattern for any new CLI code wrapping these functions: - Accept an --output-json FILE argument. - Write the JSON result to that file instead of printing to stdout. - - Document in the tool spec that the agent MUST use --output-json and read + - Document in the code spec that the agent MUST use --output-json and read from the file rather than capturing stdout. Functions that ARE affected: @@ -40,18 +40,18 @@ Functions NOT affected (h5py / numpy only — no compiled stdout writes): `compute_ke_growth_trace`, `compute_growth_rate`, `compute_q95` -## When to use these tools +## When to use these codes -Always use the M3D-C1 tools when processing or investigating datasets created by this code. Only create custom tools when your needs cannot be met by any of the existing tools. Be watchful for synonyms of alternative expressions from those used in the tool descriptions; for example, use of `repackage_hdf5` should also be triggered by calls to repack or reorganize HDF5 data, or to create a new file to hold a new dataset etc. +Always use the M3D-C1 codes when processing or investigating datasets created by this code. Only create custom codes when your needs cannot be met by any of the existing codes. Be watchful for synonyms of alternative expressions from those used in the tool descriptions; for example, use of `repackage_hdf5` should also be triggered by calls to repack or reorganize HDF5 data, or to create a new file to hold a new dataset etc. -When using any tools that write a temporary JSON file (to circumvent the corruption of stdout by fpy as described above) that temporary file should be deleted as soon as it is no longer needed to connect tool calls. Do not leave unnecessary JSON files in the filesystem. +When using any codes that write a temporary JSON file (to circumvent the corruption of stdout by fpy as described above) that temporary file should be deleted as soon as it is no longer needed to connect code calls. Do not leave unnecessary JSON files in the filesystem. -If the user's meaning or intent is unclear always ask a clarifying question --- do not silently fail or say that a required tool does not exist. +If the user's meaning or intent is unclear always ask a clarifying question --- do not silently fail or say that a required code does not exist. ## Note on how M3D-C1 field data is stored -M3D-C1 stores fields using the coefficients of the basis functions rather than the pointwise field values themselves. The field values can be found using the `eval_field` function in the `m3dc1` python module; the `eval_field_on_grid` function in `m3dc1_tools.py` evaluates field values across a spatial grid and either returns a dict or writes an HDF5 file. Use these functions or tools derived from them to evaluate fields. If the user asks for field data (e.g. temperature, density, pressure, magnetic field or current density components) to be repackaged, repacked, extracted, moved, rewritten etc. assume that they are interested in these real evaluated field values rather than the raw coefficient values. Only repackage raw coefficients if the user specifically requests this. If the user's intent is unclear always ask a clarifying question. +M3D-C1 stores fields using the coefficients of the basis functions rather than the pointwise field values themselves. The field values can be found using the `eval_field` function in the `m3dc1` python module; the `eval_field_on_grid` function in `m3dc1_tools.py` evaluates field values across a spatial grid and either returns a dict or writes an HDF5 file. Use these functions or codes derived from them to evaluate fields. If the user asks for field data (e.g. temperature, density, pressure, magnetic field or current density components) to be repackaged, repacked, extracted, moved, rewritten etc. assume that they are interested in these real evaluated field values rather than the raw coefficient values. Only repackage raw coefficients if the user specifically requests this. If the user's intent is unclear always ask a clarifying question. ## Where to store products @@ -68,7 +68,7 @@ When asked to make a shell script always check the user's default shell. Create ## Reference material -ALWAYS read the relevant API guide before attempting to create CLI tools from one of the python modules. +ALWAYS read the relevant API guide before attempting to create CLI codes from one of the python modules. - `references/m3dc1_output_guide.md`: explains M3D-C1's native output format in detail, including how the field coefficients are stored and evaluated.