UltraRAG

# AGENTS.md

This document is the primary engineering guide for autonomous coding agents working in the `UltraRAG` repository.

Use this file as the source of truth for architecture, conventions, workflows, and safe change patterns.
If `CLAUDE.md` exists, it should only point to this file.

---

## 1) Project Identity

`UltraRAG` is a lightweight RAG framework built around the Model Context Protocol (MCP).
The key design choice is strict modularization: retrieval, prompting, generation, routing, memory, and evaluation are implemented as independent MCP servers orchestrated by YAML pipelines.

Current core metadata:

- Package: `ultrarag`
- Version: `0.3.0`
- Python: `>=3.11, <3.13`
- CLI entrypoint: `ultrarag = ultrarag.client:main`
- Package manager: `uv` (`[tool.uv] package = true`)

---

## 2) Repository Map (What Matters Most)

```text
UltraRAG/
├── src/ultrarag/                    # Installable core package
│   ├── client.py                    # CLI + pipeline engine + run/build orchestration
│   ├── server.py                    # UltraRAG_MCP_Server (FastMCP extension)
│   ├── api.py                       # Python API wrappers (ToolCall, PipelineCall)
│   ├── cli.py                       # Rich banner and CLI visuals
│   ├── mcp_logging.py               # Central logging setup
│   ├── mcp_exceptions.py            # Node.js checks for remote MCP
│   └── utils.py                     # Subprocess lifecycle helpers
│
├── servers/                         # MCP microservices (each server is independent)
│   ├── retriever/
│   ├── generation/
│   ├── prompt/
│   ├── reranker/
│   ├── benchmark/
│   ├── evaluation/
│   ├── corpus/
│   ├── memory/
│   ├── router/
│   ├── custom/
│   ├── pageindex/
│   └── sayhello/
│
├── examples/
│   ├── demos/                       # UI-ready demo pipelines
│   └── experiments/                 # Experiment/research pipelines
│
├── ui/
│   ├── backend/                     # Flask backend + pipeline manager
│   └── frontend/                    # Vite + React + TypeScript
│
├── docs/                            # Docs and assets
├── script/                          # Utility scripts (deploy, case study, etc.)
├── pyproject.toml                   # Dependencies + package metadata
├── uv.lock                          # Locked dependency graph
├── Dockerfile*                      # Container variants
└── .gitignore
```

Important generated/derived files:

- `servers/*/server.yaml` (generated by each server `build` tool)
- `examples/**/parameter/*_parameter.yaml` (pipeline-merged parameters)
- `examples/**/server/*_server.yaml` (pipeline-merged server config)
- `output/memory_*.json` (per-run memory snapshots)

---

## 3) Mental Model of the System

Think of UltraRAG as a three-layer system:

1. **Interface layer**: CLI (`ultrarag ...`), UI (`ultrarag show ui`), and Python API (`ToolCall`, `PipelineCall`)
2. **Orchestration layer**: `src/ultrarag/client.py` (`build`, `load_pipeline_context`, `execute_pipeline`)
3. **Execution layer**: MCP servers in `servers/*`, each exposing tools/prompts over stdio (or remote MCP proxy)

The runtime contract is:

- A pipeline YAML declares **which servers** to use and **which steps** to execute.
- The client resolves I/O dependencies between steps.
- Each step calls exactly one MCP tool or prompt.
- Outputs are saved to a shared variable pool and can feed downstream steps.

---

## 4) Two-Phase Execution Lifecycle

### Phase A: Build

Command:

```bash
ultrarag build <pipeline.yaml>
```

What happens:

- Reads `servers:` from the pipeline YAML.
- For each referenced server, calls the server's `build` tool.
- Produces:
  - `<pipeline_dir>/parameter/<pipeline_name>_parameter.yaml`
  - `<pipeline_dir>/server/<pipeline_name>_server.yaml`

Why this matters:

- `build` materializes exact tool/prompt I/O metadata before runtime.
- UI and runner rely on these generated artifacts for deterministic execution.

### Phase B: Run

Command:

```bash
ultrarag run <pipeline.yaml> [--param path] [--is_demo]
```

What happens:

- Loads generated server config + parameter config.
- Creates `fastmcp.Client` transport config for each server.
- Executes pipeline steps in order, including `loop` and `branch`.
- Saves intermediate memory snapshots and writes `output/memory_*.json`.
- Invokes cleanup tools (e.g., tools ending with `vllm_shutdown`) if present.

---

## 5) Pipeline DSL Reference

UltraRAG accepts mixed step forms:

### 5.1 Plain step

```yaml
- retriever.retriever_search
```

### 5.2 Step with input/output remapping

```yaml
- generation.generate:
    input:
      prompt_ls: custom_prompt_ls
    output:
      ans_ls: final_answer_ls
```

### 5.3 Loop block

```yaml
- loop:
    times: 3
    steps:
    - retriever.retriever_search
    - generation.generate
```

### 5.4 Branch block

```yaml
- branch:
    router:
    - router.route_query
    branches:
      need_retrieval:
      - retriever.retriever_search
      direct_answer:
      - generation.generate
```

### 5.5 Prompt vs Tool step semantics

- Steps under the `prompt` server call `client.get_prompt(...)`.
- Non-prompt steps call `client.call_tool(...)`.

---

## 6) Variable Resolution and Data Flow Rules

In `UltraData`, each tool input value is interpreted by convention:

- `"$foo"` -> load from server-local params (`parameter.yaml`)
- `"bar"`  -> read from global variable pool (`global_vars["bar"]`)
- `"memory_xxx"` -> read/write memory lists

Output handling:

- Tool returns JSON payload -> keys mapped to declared outputs
- Output remapping (`output:` in pipeline step) is applied at save time
- Prompt outputs usually produce `prompt_ls`

Branch handling:

- Branch-aware values use wrapped list records with branch-state keys.
- Internal sentinel (`UNSET`) is used to distinguish "not filled yet" from `None`.

Memory handling:

- The engine tracks `memory_*` histories automatically.
- Final snapshots are serialized to `output/memory_<...>.json`.
- If a memory server is detected, turn-level memory auto-save can be triggered.

---

## 7) Core Python Modules (Authoritative Guide)

### `src/ultrarag/client.py`

This is the orchestration heart of the project.

Key responsibilities:

- CLI entrypoint (`main`)
- UI launch (`launch_ui`) and case-study launch (`launch_case_study`)
- Config loading (`Configuration`)
- Pipeline data graph and state (`UltraData`)
- Build pipeline configs (`build`)
- Load runtime context (`load_pipeline_context`)
- Execute step engine (`execute_pipeline`)
- Run full pipeline (`run`)

Important runtime behaviors:

- Supports both local python MCP servers (`path.endswith(".py")`) and remote MCP endpoints (`http(s)`).
- For remote MCP, requires Node.js >= 20 and uses `npx -y mcp-remote <url>`.
- Keeps loop-termination state in `ContextVar` for coroutine safety.
- Emits structured stream events in demo/UI flows (`step_start`, `step_end`, `token`, `sources`).

### `src/ultrarag/server.py`

Defines `UltraRAG_MCP_Server`, a compatibility wrapper over FastMCP.

Key responsibilities:

- Enhanced `tool()` and `prompt()` registration with `output` metadata support
- Metadata capture for automatic config generation
- `build(parameter_file)` to generate per-server `server.yaml`
- Compatibility filtering for FastMCP signature differences

### `src/ultrarag/api.py`

Provides ergonomic Python-side wrappers:

- `initialize(servers, server_root, log_level)`
- `ToolCall.server_name.tool_name(...)`
- `PipelineCall(pipeline_file, parameter_file, log_level)`

### `src/ultrarag/mcp_logging.py`

- Initializes root logger `UltraRAG`
- Rich console logging + file logging (`logs/<timestamp>.log`)
- Log level controlled by `log_level` argument and environment

### `src/ultrarag/mcp_exceptions.py`

- Validates local Node.js availability/version
- Raises `NodeNotInstalledError` / `NodeVersionTooLowError`

### `src/ultrarag/utils.py`

- Subprocess lifecycle helpers
- POSIX parent-death signal support
- Windows job object support for child-process cleanup

---

## 8) MCP Server Authoring Contract

Each server follows this shape:

```text
servers/<name>/
├── parameter.yaml
├── server.yaml            # generated
└── src/<name>.py
```

### 8.1 Registration styles

Use either:

1. Decorator style
2. Class-bound method registration style

Both are valid in this codebase.

### 8.2 `output=` grammar

Canonical form:

```text
input1,input2,$param_a -> output1,output2
```

Rules:

- Left side maps function args to pipeline inputs.
- Right side defines expected output keys from returned dict/JSON.
- `-> None` means no output variables.
- `$param` means value comes from server parameter file.

### 8.3 Return payload expectations

- For tools: return JSON-serializable dict payloads matching declared output keys.
- For prompts: return prompt messages (typically list-like prompt payloads consumed by `get_prompt`).

### 8.4 Entrypoint requirement

Server modules should end with:

```python
if __name__ == "__main__":
    app.run(transport="stdio")
```

---

## 9) Retriever/Generation/Prompt Specific Notes

### Retriever (`servers/retriever`)

- Supports multiple retrieval modes:
  - Dense retrieval
  - BM25
  - Web search
  - Project-memory retrieval
- Index backends are pluggable via factory:
  - `faiss`
  - `milvus`
  - `qdrant`
- Web search backends are pluggable via factory:
  - `exa`
  - `tavily`
  - `zhipuai`

### Generation (`servers/generation`)

- Supports generation backends including `openai`, `vllm`, and hf workflows.
- Provides explicit cleanup tool `vllm_shutdown`.
- Demo mode can use local streaming generation service.

### Prompt (`servers/prompt`)

- Uses `SandboxedEnvironment` from Jinja2 for safer rendering.
- Validates template paths to reduce traversal risk.
- Escapes string inputs before template rendering.

---

## 10) UI Backend Architecture (`ui/backend`)

### `app.py`

- Flask app factory (`create_app`)
- Serves frontend static assets
- Exposes chat/pipeline/KB/auth endpoints
- Reads optional frontend override via `ULTRARAG_FRONTEND_DIR`

### `pipeline_manager.py`

This module is large and central to UI behavior.

Major responsibilities:

- Session lifecycle and streaming chat management
- Background chat task management
- Pipeline CRUD (list/load/save/rename/delete)
- Parameter load/save/build wrappers
- Knowledge base file ingest and pipeline triggering
- Memory synchronization to per-user KB collections
- Optional server introspection via AST stub generation if `server.yaml` is missing

Notable behavior:

- Applies defensive patches to suppress noisy closed-event-loop teardown logs.
- Uses a queue bridge for async-to-sync SSE event streaming.
- Supports automatic memory-to-KB sync for pipelines that include memory components.

---

## 11) Storage Model and Paths

Default UI storage root:

- `ui/storage`

Can be overridden by:

- `ULTRARAG_UI_STORAGE_ROOT`

Key subpaths:

- `db/users.sqlite3`
- `chat_sessions/`
- `knowledge_base/raw|corpus|chunks|index`
- `memory/`
- `knowledge_base/_memory_sync`

Runtime outputs:

- `output/memory_*.json`
- `logs/*.log`

---

## 12) Environment Variables You Should Know

- `ULTRARAG_UI_STORAGE_ROOT`: override UI storage root
- `ULTRARAG_FRONTEND_DIR`: override frontend static directory
- `ULTRARAG_SESSION_TIMEOUT`: foreground chat session timeout
- `ULTRARAG_BG_SESSION_TIMEOUT`: background session timeout
- `ULTRARAG_LOG_TS`: custom timestamp seed for log file naming
- `log_level`: consumed by core logger initialization

---

## 13) Dependency Model

Install tiers from `pyproject.toml`:

- Core install: no extras
- `retriever` extra
- `generation` extra
- `evaluation` extra
- `corpus` extra
- `all` extra (union)

Typical commands:

```bash
uv sync
uv sync --extra retriever
uv sync --extra generation
uv sync --all-extras
```

Development dependencies include:

- `ruff`
- `ipython`
- `jupyter`
- `pytest`

---

## 14) CLI Commands (Canonical)

```bash
ultrarag build <pipeline.yaml>
ultrarag run <pipeline.yaml> [--param <parameter.yaml>] [--log_level info|debug|warn|error] [--is_demo]
ultrarag show ui [--host 127.0.0.1] [--port 5050]
ultrarag show case [--config_path <memory.json>] [--host 127.0.0.1] [--port 8080]
```

Minimal smoke check:

```bash
ultrarag run examples/experiments/sayhello.yaml
```

---

## 15) Docker Variants

- `Dockerfile`: full image (builds frontend, installs all extras)
- `Dockerfile.base-cpu`: CPU base image
- `Dockerfile.base-gpu`: GPU base image

All variants start UI with:

```bash
ultrarag show ui --port 5050 --host 0.0.0.0
```

---

## 16) Development Playbooks

### 16.1 Add a new MCP server

1. Create `servers/<name>/parameter.yaml`
2. Implement `servers/<name>/src/<name>.py`
3. Register tools/prompts via `app.tool` / `app.prompt` (or class-bound registration)
4. Ensure `app.run(transport="stdio")` exists
5. Add server to `servers:` in a pipeline YAML
6. Run `ultrarag build <pipeline.yaml>`

### 16.2 Add a new tool to an existing server

1. Implement function/method
2. Register with explicit `output=...` contract
3. Ensure return payload keys match outputs
4. Update relevant parameter keys in `parameter.yaml` if using `$...`
5. Add the step in pipeline YAML and rebuild

### 16.3 Add a retriever index backend

1. Implement backend in `servers/retriever/src/index_backends/`
2. Follow `BaseIndexBackend` contract
3. Register in `_INDEX_BACKENDS` map in `index_backends/__init__.py`

### 16.4 Add a web-search backend

1. Implement backend in `servers/retriever/src/websearch_backends/`
2. Follow `BaseWebSearchBackend` contract
3. Register in `_WEBSEARCH_BACKENDS` map

### 16.5 Modify UI pipeline behavior

Primary files:

- `ui/backend/app.py`
- `ui/backend/pipeline_manager.py`

If touching build/runtime semantics, cross-check against:

- `src/ultrarag/client.py`

---

## 17) Coding Standards (Repository-Conformant)

- Use type hints for function signatures.
- Prefer `pathlib.Path` for filesystem paths.
- Use `yaml.safe_load` / `yaml.safe_dump`.
- Use project logger (`get_logger` or `app.logger`) instead of `print`.
- Keep tool outputs deterministic and JSON-serializable.
- Keep imports grouped: stdlib -> third-party -> local.
- Keep async boundaries explicit (`async`/`await`).

---

## 18) What Not To Edit Blindly

Treat these as generated or runtime artifacts unless intentionally regenerating:

- `servers/*/server.yaml`
- `examples/**/parameter/*_parameter.yaml`
- `examples/**/server/*_server.yaml`
- `output/*`
- `logs/*`
- `ui/storage/*` runtime data

Also avoid committing secrets:

- `.env`
- any credential-bearing local config

---

## 19) Common Failure Modes and Fixes

### Error: server file not found

- Check `servers.<name>` path in pipeline YAML.
- Ensure server entry script exists at `servers/<name>/src/<name>.py` (or update path in config).

### Error: missing variable in pipeline execution

- Verify output key names from upstream tool match downstream input names.
- Verify remapping under step-level `output:` is correct.
- Verify `$param` keys exist in that server's parameter config.

### Remote MCP server fails to start

- Ensure Node.js >= 20.
- Confirm `npx` is available.
- Confirm remote URL in server path is reachable.

### Build succeeds but UI cannot list tools

- Ensure `server.yaml` exists or can be inferred by AST parsing.
- Check for unusual dynamic registration patterns that static analysis cannot infer.

### No final answer in chat

- Inspect `output/memory_*.json`.
- Check whether generation step ran and produced `ans_ls`.
- Check stream events in UI path (`step_start`, `step_end`, `sources`, `token`, `final`).

---

## 20) Validation Checklist for Agents

Before finalizing any non-trivial change:

1. Build the affected pipeline:
   - `ultrarag build <pipeline.yaml>`
2. Run a relevant smoke case:
   - `ultrarag run <pipeline.yaml>`
3. If UI behavior changed:
   - run `ultrarag show ui` and verify route behavior
4. If retriever/generation backends changed:
   - validate parameter schema keys and output names
5. Keep generated artifacts intentional:
   - do not accidentally commit transient runtime outputs

---

## 21) Minimal Quickstart for New Agents

```bash
# 1) Install dependencies
uv sync --all-extras

# 2) Smoke test
ultrarag run examples/experiments/sayhello.yaml

# 3) Build and run a demo pipeline
ultrarag build examples/demos/LLM.yaml
ultrarag run examples/demos/LLM.yaml

# 4) Launch UI
ultrarag show ui --host 127.0.0.1 --port 5050
```

---

## 22) Final Notes

- This repository is orchestration-first: correctness depends heavily on I/O naming consistency across tools and pipeline steps.
- Most regressions come from mismatched variable names, stale generated configs, or incomplete parameter updates.
- When in doubt, inspect `src/ultrarag/client.py` first: it is the execution truth.

More agent context in OpenBMB/UltraRAG

2 other files this repository gives its agents.