Multi-Agent Code Llama Demo

A minimal proof-of-concept showing how two LLM-powered agents (a Developer and a Reviewer) can collaborate in natural language to solve a programming task end-to-end.
The demo loads CodeLlama-7B-Instruct to write code and CodeLlama-13B-Instruct to review it, runs unit tests automatically, and iterates until the solution is APPROVED.

---

What you get

• Multi-agent loop driven by plain Python – no hidden frameworks.
• Deterministic(ish) generation with low-temperature sampling for code.
• Offline execution (CPU-only is possible, ~16 GB RAM per model when loaded in 4-bit).
• Unit-test feedback injected into the Reviewer’s context.
• Docker + docker-compose setup that mounts your local Hugging Face cache.
• A ready-made tests.py example (Sieve of Eratosthenes).

---

Folder layout

├── download_models.py # robust HF snapshot downloader with resume 
├── multi_agent_codellama.py # core demo script 
├── tests.py # sample unit tests 
├── Dockerfile # slim Python image with CPU PyTorch + transformers 
└── docker-compose.yml # single-service stack mounting HF cache

---

Quick start (Docker)

# 1. clone the repo
git clone https://github.com/YOUR_USERNAME/multi-agent-codellama.git
cd multi-agent-codellama

# 2. (one-time) download the models to ~/.cache/huggingface
python download_models.py          # ≈ 13 GB total – resumable

# 3. build and launch the container (CPU image)
docker compose build
docker compose run --rm llama-agents \
  python multi_agent_codellama.py \
  "Write a function primes_up_to(n) that returns all primes ≤ n using the Sieve of Eratosthenes."

The compose file mounts your local HF cache into the container, so the heavy model files are reused across runs.

---

How it works — deep-dive

Below is the full control-flow executed by multi_agent_codellama.py.

1. Model bootstrap

1. make_agents() loads two Code Llama checkpoints via Hug Transformers:
   • Developer → codellama/CodeLlama-7b-Instruct-hf
   • Reviewer → codellama/CodeLlama-13b-Instruct-hf
     The models are instantiated with device_map="auto" and 4-bit GPTQ so they fit in CPU RAM.
2. A dedicated AutoTokenizer is created for each model; both are stored inside an Agent dataclass together with:
   • a system prompt (role instructions),
   • a private chat history list.

2. Conversation seed

1. Task ingestion from the CLI
   When you launch the script you pass a natural-language brief, e.g.:

   python multi_agent_codellama.py \
     "Write a function primes_up_to(n) ..."

   That string is stored in the local variable requirement.

2. Tagging the brief as a system message
   The script rewrites the raw text into the fixed pattern:

   "System: Requirement — " + requirement
   

   Why?

   • Using the System label tells the LLM “this is immutable context, not a turn you must answer.”
   • Prefixing with Requirement — makes the intent crystal-clear, even if the prompt later grows.

3. Injecting into the Developer’s memory

   developer.hear(f"System: Requirement — {requirement}")

   Under the hood this is a simple list.append() on developer.history, so the Developer now holds a single entry:

   ["System: Requirement — Write a function primes_up_to(n) ..."]

   The Reviewer’s history remains an empty list at this stage.

4. Zero tokens spent so far No model call happens during seeding. The first inference is triggered later by:

   dev_reply = developer.speak()

   At that moment the Developer’s prompt becomes:

   <system_prompt>
   
   System: Requirement — ...
   Developer:
   

   Only then does Code Llama start generating.

---

3. Dialogue loop (max --rounds)

Phase	What happens	Key implementation details
Developer speaks	1. The helper Agent._format_prompt() concatenates: • system prompt • full history • literal “`Developer:`” 2. The string is passed to llama_generate(), which calls the 7 B model with: • low temperature = 0.15 (makes answers stable) • top_p=0.9, repetition_penalty=1.05. 3. The raw completion is appended to history and printed to the terminal.	Developer’s prompt explicitly forbids commentary and demands only one fenced code block.
Code extraction	A simple regex pulls the first triple-fenced block (optionally tagged “python”).	Regex pattern lives in CODE_BLOCK_RE.
Unit-test harness	1. Code is written to a temp file. 2. The temp file is imported as module solution. 3. The chosen tests.py is exec’d with that module in scope. 4. Stdout + traceback get captured so the Reviewer can see failures.	Entire routine happens in run_tests() and returns a boolean + feedback string.
Reviewer context update	Two messages are appended to the Reviewer’s history: 1. A fenced code dump of what the Developer wrote. 2. Plain text: “Unit tests passed: <True/False>” plus captured output.	Keeps the review process reproducible and transparent.
Reviewer speaks	Reviewer’s formatted prompt (similar to Developer’s) is run through the 13 B model. The reply can only be: • APPROVED — tests passed and code looks clean. • REJECTED — followed by specific bullet-point suggestions.	If the Reviewer writes APPROVED, the script exits successfully. Otherwise the whole Reviewer reply is appended to the Developer’s history and the next round begins.

The loop terminates as soon as:

• Reviewer says APPROVED, or
• The maximum number of rounds (CLI flag --rounds, default 3) is reached.

---

4. Prompt engineering choices

• Few-shot memory – Every turn’s transcript is re-fed to the respective agent; this creates an implicit shared memory without an external vector DB.
• Role separation – Although both agents use Code Llama, divergent system prompts make them behave differently.
• Explicit test feedback – Piping the unit-test outcome into the Reviewer dramatically increases the chance of meaningful reviews compared to asking the model to “imagine” runtime errors.

---

5. Fail-safes & tweaks

• Resume downloads – download_models.py skips files already present and uses resume_download=True to survive network drops.
• Footgun guard – The Reviewer must literally start with “APPROVED”; accidental praise that lacks the keyword will keep the loop alive.
• Scalability – Swap in larger/smaller Llama variants, or add more agents (e.g. “Optimizer”) by instantiating extra Agent objects and adjusting the turn order.

That’s it ~200 lines of Python orchestrate two LLMs, tests, and an approval workflow without any heavyweight dependencies.