# Ai Agent Tool Augmentation

> Multi-agent LLM system that augments AI agent tools using code generators (ReAct-based code generation, security review, Docker sandbox execution)

- **Type:** MCP server
- **Install:** `agentstack add mcp-schneider-cherry-ai-agent-tool-augmentation`
- **Verified:** Yes — security-reviewed for prompt injection and unsafe behavior
- **Seller:** [Schneider-cherry](https://agentstack.voostack.com/s/schneider-cherry)
- **Installs:** 0
- **Category:** [Cloud & Infrastructure](https://agentstack.voostack.com/c/cloud-infrastructure)
- **Latest version:** 0.1.0
- **License:** MIT
- **Upstream author:** [Schneider-cherry](https://github.com/Schneider-cherry)
- **Source:** https://github.com/Schneider-cherry/ai-agent-tool-augmentation

## Install

```sh
agentstack add mcp-schneider-cherry-ai-agent-tool-augmentation
```

Requires the [AgentStack CLI](https://agentstack.voostack.com/docs/cli). Works with Claude Code, Cursor, and any MCP-compatible agent.

## About

# Tool Augmentation for AI Agents - agent and tool platform

Bachelor thesis project from TUKE FEI, 2026.

This repository contains an agent for code generation, review, and Docker execution, a tool registry, an MCP server, a Streamlit UI, and experimental pipelines. The main logic is in `src/`, the UI is in `UI/`, and experiments are in `experiments/`.

## Key Result

Routing generated code through the review -> execute pipeline improved Llama-3.1 task success rate from ~27% to ~54% on the benchmark set.

## Quick Start

1. Install dependencies:

```
python -m venv .venv
source .venv/bin/activate
pip install -r requirements.txt
```

2. Start Docker, which is required for code execution.

3. Start the UI:

```
streamlit run UI/StreamLitLangChain.py
```

## Architecture Diagram

```mermaid
sequenceDiagram
  actor User
  participant UI as Streamlit UI
  participant Client as MCPClient
  participant Server as MCPToolServer
  participant Orchestrator as AgentOrchestrator
  participant Dev as DeveloperAgent
  participant Conv as ConversationAgent
  participant Review as ReviewerAgent
  participant Exec as ExecutorAgent
  participant Registry as ToolRegistry
  participant Tool as GeneratedScriptTool
  participant Docker as DockerExecutor
  participant LLM as UnifiedLLMClient

  User->>UI: Send prompt

  alt Conversational prompt
    UI->>Client: chat(message, reset?)
    Client->>Server: chat(message, reset?)
    Server->>Orchestrator: chat(message, reset?)
    Orchestrator->>Conv: Build reply from history + user message
    Conv->>LLM: chat(messages)
    LLM-->>Conv: assistant text
    Conv-->>Orchestrator: response
    Orchestrator-->>Server: response
    Server-->>Client: response
    Client-->>UI: stream/render text
  else Code / tool / external-fact prompt
    UI->>Client: run_pipeline(prompt, mode="chat")
    Client->>Server: run_agent_pipeline(prompt, mode="chat")
    Server->>Orchestrator: run_pipeline(prompt, mode="chat")

    Orchestrator->>Registry: load_tool_instances()
    Registry-->>Orchestrator: enabled saved tools
    Orchestrator->>Dev: run(prompt, tools_info)
    Dev->>LLM: chat(system+user messages)
    LLM-->>Dev: ReAct output or Python code
    Dev-->>Orchestrator: raw output

    alt ReAct selected a saved tool
      Orchestrator->>Tool: run(params)
      Tool->>Docker: execute(tool script with TOOL_INPUT)
      Docker-->>Tool: tool output
      Tool-->>Orchestrator: observation
      Orchestrator->>Conv: final answer from observation
      Conv->>LLM: chat(messages)
      LLM-->>Conv: final answer
      Conv-->>Orchestrator: response
    else No tool selected, generate code
      Orchestrator->>Review: run(code)
      Review->>LLM: chat(review prompt)
      LLM-->>Review: fixed code
      Review-->>Orchestrator: reviewed code
      Orchestrator->>Exec: validate(code)
      Exec->>LLM: chat(validation prompt)
      LLM-->>Exec: final code
      Exec-->>Orchestrator: validated code
      Orchestrator->>Docker: execute(validated code)
      Docker-->>Orchestrator: runtime output
    end

    opt UI can save reusable generated code
      UI->>Client: save_tool(code, task_description)
      Client->>Server: save_tool(code, task_description)
      Server->>Registry: save_tool(code, task_description)
      Registry-->>Server: saved file path + enabled flag
      Server-->>Client: result
      Client-->>UI: success / error
    end
  end
```

This is a more accurate map of the real execution flow: the UI does not always go through the pipeline. For normal dialogue, it can stream the reply directly through `chat()`. The full pipeline runs only for technical, code-related, or tool-related requests.

## State Diagram (UML State Machine)

```mermaid
stateDiagram-v2
  [*] --> Idle
  Idle --> ClassifyIntent: user prompt

  ClassifyIntent --> ConversationResponse: no code/tool/external intent
  ClassifyIntent --> DeveloperStep: pipeline required

  DeveloperStep --> ToolExecution: ReAct tool selected
  DeveloperStep --> CodePipeline: no tool selected

  ConversationResponse --> Finalize
  ToolExecution --> Finalize
  CodePipeline --> Finalize

  Finalize --> Idle: ready for next prompt
```

## Main Components

| Component | Path | Purpose |
| --- | --- | --- |
| MCP stdio server | `src/mcp/mcp_stdio_server.py` | MCP server over stdio with tools and prompts |
| MCP server (in-process) | `src/mcp/server.py` | Server for UI/CLI calls through the Python API |
| MCP client | `src/mcp/client.py` | Client for tool management and pipeline execution |
| Agent orchestrator | `src/langchain_module/lc_agent.py` | Tool selection, ReAct, code generation, execution |
| Agents | `src/langchain_module/agents/*.py` | Developer, Reviewer, Executor, Conversation |
| Tool registry | `src/tools/registry.py` | Saving, enabling, and indexing tools |
| Generated tools | `src/generated_tools/` | Generated code files grouped by category |
| Docker executor | `src/executors/docker_executor.py` | Safe code execution |
| UI | `UI/StreamLitLangChain.py` | Streamlit interface |

## How Auto Agent Chat Works

The "Auto Agent Chat" mode in the UI uses `MCPClient.run_pipeline(..., mode="chat")`. Behavior is controlled by `AgentOrchestrator`:

1. If the request looks like normal conversation, the UI can go directly to `client.stream_chat()` or `client.chat()` without the pipeline.
2. If the request is technical, contains coding intent, requires external facts, or matches an existing tool, `run_pipeline(..., mode="chat")` is started.
3. `DeveloperAgent` either returns a ReAct `Action` or generates a Python code block.
4. If an `Action` is found, the selected saved tool is executed through `GeneratedScriptTool` and `DockerExecutor`.
5. If no tool is selected, the code goes through `ReviewerAgent` -> `ExecutorAgent.validate()` -> `DockerExecutor.execute()`.
6. If execution succeeds and the code looks like a reusable `def` or `class`, the UI shows a form to save a new tool.

## ReAct Format and Parsing

The Developer agent must respond in one of these formats:

```
Thought: ...
Action: tool_name
Action Input: {"key": "value"}
```

or with code blocks in ` ```python ``` `. Parsing is implemented in `src/utils/react_parsing.py`.

## Tools and Registry

### Where Tools Are Stored

Tools are stored in `src/generated_tools/` with subcategories such as `Date/`, `Math/`, `Logic/`, and `General/`. The enablement state is stored in `src/generated_tools/.enabled_tools.json`.

### Tool Format

A tool file is a Python file with a docstring at the top. The registry reads the docstring and available functions to build metadata.

### Tool Input Parameters

`GeneratedScriptTool` passes input through the `TOOL_INPUT` variable. The input can be:

- a string for a simple call,
- a dictionary for ReAct with parameters.

### How to Enable or Disable a Tool

1. Through the UI (the "Saved tools" section).
2. Through the MCP client:

```
from src.mcp.client import create_client

client = create_client(provider="openrouter", model="google/gemma-3-27b-it:free")
client.enable_tool("Date/tool_current_time_and_data")
```

## MCP Servers

### MCP Stdio Server

`src/mcp/mcp_stdio_server.py` is built on FastMCP and exposes:

| Tool | Purpose |
| --- | --- |
| `generate_code` | Code generation through the Developer agent |
| `review_code` | Code review |
| `execute_code` | Validation and execution in Docker |
| `run_full_pipeline` | Full generate -> review -> execute chain |
| `chat_with_agent` | Context-aware conversation |
| `list_tools` | List available tools |
| `toggle_tool` | Enable or disable a tool |
| `save_new_tool` | Save a new tool |

MCP prompts:

- `solve_task`: step-by-step scenario
- `full_pipeline`: one-shot execution of the full pipeline

### MCP In-Process Server

`src/mcp/server.py` is a server for use from Python code (UI/CLI). It provides the same capabilities, but without stdio transport.

### MCP Configuration

The `mcp_config.json` file describes how to run the MCP server:

```
{
  "mcpServers": {
    "zp-agents": {
      "command": "/path/to/.venv/bin/python",
      "args": ["/path/to/src/mcp/mcp_stdio_server.py"],
      "env": {
        "LLM_PROVIDER": "openrouter",
        "LLM_MODEL": "google/gemma-3-27b-it:free"
      }
    }
  }
}
```

## Providers and Models

Supported providers: `azure`, `openai`, `groq`, `openrouter`, `gemini`, `ollama`, `huggingface`.

| Provider | Key Environment Variables |
| --- | --- |
| Azure | `AZURE_OPENAI_KEY`, `AZURE_OPENAI_ENDPOINT`, `AZURE_OPENAI_DEPLOYMENT`, `AZURE_OPENAI_API_VERSION` |
| OpenAI | `OPENAI_API_KEY` |
| Groq | `GROQ_API_KEY`, `GROQ_BASE_URL` (optional) |
| OpenRouter | `OPENROUTER_API_KEY`, `OPENROUTER_BASE_URL`, `OPENROUTER_REFERER`, `OPENROUTER_TITLE` |
| Gemini | `GEMINI_API_KEY`, `GEMINI_BASE_URL` |
| Ollama | `OLLAMA_BASE_URL` |
| HuggingFace | `HUGGINGFACE_TOKEN` (for gated models) |

OpenRouter has a fallback: on a 429 error, it can switch to Groq or Ollama.

## HuggingFace Direct

The "HuggingFace Direct" mode in the UI bypasses the agent pipeline and calls local HuggingFace models directly:

- `bigcode/tiny_starcoder_py` (~200 MB)
- `microsoft/CodeGPT-small-py` (~500 MB)
- `Salesforce/codegen-350M-mono` (~700 MB)

This mode returns only generated code, without execution and without saving to the registry.

## UI (Streamlit)

The UI is located in `UI/StreamLitLangChain.py`.

UI features:

- chat with the agent,
- pipeline execution,
- switching provider and model,
- browsing and enabling tools,
- prompting to save a new tool after successful execution.

Interaction modes:

| Mode | What Happens |
| --- | --- |
| Auto Agent Chat | Conversations are streamed, technical requests go through the pipeline |
| HuggingFace Direct | Immediate code generation through HF |

## Docker Execution

Code execution happens through `DockerExecutor`:

- network is disabled,
- file access is read-only,
- memory limit is 256 MB,
- CPU limit is 1 core,
- default timeout is 30 seconds.

The Docker image is built from `Dockerfile` as `python-executor:latest`.

## CLI Launch

### LangChain + Ollama

```
python src/runners/lc_main.py
```

### AutoGPT (experiment)

```
PYTHONPATH=experiments python src/runners/autogpt_main.py
```

### CrewAI (experiment)

```
PYTHONPATH=experiments python src/runners/cw_main.py
```

### Example of Direct MCP Client Usage

```
python run_task.py
```

## Experiments

`experiments/` contains pipeline prototypes:

- `autoGpt_module` - planner, coder, critic, and TaskManager
- `crewai_module` - Agent + Task + Crew

These modules are not connected to the main UI by default.

## Directory Structure

```
UI/                     Streamlit UI
src/                    Main logic
src/mcp/                MCP server and client
src/langchain_module/   Agents, orchestrator, pipeline
src/tools/              Tool registry and wrappers
src/generated_tools/    Saved tools
src/executors/          Docker execution
experiments/            Prototypes
test/                   Tests
test/experiments_archive/  Experiment archive - model testing results
```

## Tests

```
make test
make test-unit
make test-integration
```

Integration tests require real API keys and network access.

## Common Issues

| Symptom | Cause | Fix |
| --- | --- | --- |
| `Execution unavailable: Docker executor is not initialized` | Docker is not running | Start Docker |
| `Missing ... API key` | The key is missing in `.env` | Add the required environment variable |
| 429 from OpenRouter | Rate limit or quota | Check the key or wait; fallback is available |
| HF model fails to load | Missing token or disk space | Add `HUGGINGFACE_TOKEN` and check free disk space |

## Source & license

This open-source MCP server is cataloged on AgentStack and links to its original source — we do not rehost the code.

- **Author:** [Schneider-cherry](https://github.com/Schneider-cherry)
- **Source:** [Schneider-cherry/ai-agent-tool-augmentation](https://github.com/Schneider-cherry/ai-agent-tool-augmentation)
- **License:** MIT

Install and usage instructions live in the source repository linked above.

## Pricing

- **Free** — Free

## Security capabilities

Automated source analysis of v0.1.0 — what this tool can access:

- **Network access:** no
- **Filesystem access:** no
- **Shell / process execution:** no
- **Environment & secrets:** yes
- **Dynamic code execution:** no

*"Yes" means the capability is present in the source — more access means more to trust, not that it is unsafe.*


## Versions

- **0.1.0** — security scan: passed — Imported from the upstream source.

## Links

- Listing page: https://agentstack.voostack.com/l/mcp-schneider-cherry-ai-agent-tool-augmentation
- Seller: https://agentstack.voostack.com/s/schneider-cherry
- Browse the marketplace: https://agentstack.voostack.com/browse

---
Listed on AgentStack — the marketplace for AI agent skills and MCP servers. Every listing is security-reviewed. Creators keep 70%.