Agentic Swarm Intelligence

pip install tensorlake openai pydantic python-dotenv

from dotenv import load_dotenv
load_dotenv()  # Load environment variables from .env file

from tensorlake.sandbox import SandboxClient
from pydantic import BaseModel
from typing import List
from openai import OpenAI
from concurrent.futures import ThreadPoolExecutor

class ScoutReport(BaseModel):
    agent_id: int
    raw_data: str

class FinalInsight(BaseModel):
    summary: str

# 1. Worker Agent: LLM + Sandbox Execution
def scout_agent(task_id: int) -> ScoutReport:
    """Each scout analyzes a specific aspect of the mission."""
    perspectives = ["Scientific", "Economic", "Ethical", "Logistical", "Psychological"]
    perspective = perspectives[task_id % len(perspectives)]

    print(f"🕵️  Scout {task_id}: Analyzing {perspective} perspective...")
    client = OpenAI()
    # Step A: LLM decides what to do
    prompt = f"""
You are a {perspective} analyst for a Mars mission.
Write a Python script to perform a simple simulation using the 'numpy' library.
The simulation should model a key factor from your perspective (e.g., scientific sensor data, economic cost projection, logistical supply levels).

The script MUST print a single valid JSON string to standard output. This JSON should contain:
'perspective': '{perspective}',
'score': an integer from 0-100 derived from your simulation (higher is better),
'insight': a brief, unique risk or opportunity revealed by the simulation.
Do NOT use markdown blocks."""
    response = client.chat.completions.create(
        model="gpt-4o",
        messages=[{"role": "user", "content": prompt}]
    )
    
    # Clean up markdown formatting (remove ```python ... ```)
    generated_code = response.choices[0].message.content.replace("```python", "").replace("```", "").strip()
    print(f"🕵️  Scout {task_id}: Generated code -> {generated_code}")

    # Step B: Secure execution in a Sandbox
    sb_client = SandboxClient()
    with sb_client.create_and_connect() as sandbox:
        print(f"🕵️  Scout {task_id}: Installing dependencies in Sandbox...")
        sandbox.run("pip", ["install", "numpy", "--user", "--break-system-packages"])
        print(f"🕵️  Scout {task_id}: Running simulation in Sandbox...")
        execution = sandbox.run("python3", ["-c", generated_code])
        output = execution.stdout.strip()
        print(f"🕵️  Scout {task_id}: Execution complete. Output: {output}")
        return ScoutReport(agent_id=task_id, raw_data=output)

# 2. Lead Agent: LLM Aggregator (The "Reducer")
def lead_aggregator(reports: List[ScoutReport]) -> FinalInsight:
    """The Lead LLM reviews all sandbox outputs to find patterns."""
    print(f"👑 Lead Agent: Received {len(reports)} scout reports. Aggregating...")
    client = OpenAI()
    combined_reports = "\n".join([f"Report {r.agent_id}: {r.raw_data}" for r in reports])
    
    # The 'Intelligence' step: synthesizing multiple sources
    prompt = (
        f"You are the Mission Commander for Mars Colonization. Review these viability reports:\n{combined_reports}\n\n"
        "1. Calculate the average viability score.\n"
        "2. Synthesize a strategic Go/No-Go recommendation.\n"
        "3. Summarize key risks."
    )
    
    response = client.chat.completions.create(
        model="gpt-4o",
        messages=[{"role": "user", "content": prompt}]
    )
    
    return FinalInsight(summary=response.choices[0].message.content)

# 3. The Swarm Application
def intelligence_swarm(count: int) -> str:
    print(f"🚀 Launching a swarm of {count} scouts...")
    # Parallel Map: Launch multiple sandboxed scouts
    with ThreadPoolExecutor() as executor:
        reports = list(executor.map(scout_agent, range(count)))
    
    # Reduce: Use the Lead Agent to combine results
    final_insight = lead_aggregator(reports)
    
    return final_insight.summary

if __name__ == "__main__":
    # This runs 3 parallel LLMs, 3 parallel Sandboxes, and 1 Aggregator LLM
    result = intelligence_swarm(count=5)
    print(f"\n--- SWARM INTELLIGENCE REPORT ---\n{result}")

# 1. Create a snapshot ID (do this once)
# snapshot = client.snapshot_and_wait(base_sandbox_id)

# 2. Use it in your agent
with sb_client.create_and_connect(snapshot_id="snps_abc123") as sandbox:
    # Numpy is already installed!
    sandbox.run("python", ["-c", generated_code])

with sb_client.create_and_connect(allow_internet_access=False) as sandbox:
    # ...