Ai Agent Tool Augmentation

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

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

1. Start the UI:

streamlit run UI/StreamLitLangChain.py

Architecture Diagram

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)

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
• Source: Schneider-cherry/ai-agent-tool-augmentation
• License: MIT

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