Add Ollama support for free local LLM inference

New OllamaClient talks to locally-running open-source models (Llama 3,
Mistral, Qwen, etc.) via Ollama's API. Auto-detects provider based on
whether ANTHROPIC_API_KEY is set. Added --provider and --model CLI flags.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

.env.example

@@ -1 +1,10 @@
-ANTHROPIC_API_KEY=sk-ant-your-key-here
+# LLM Provider: "claude" or "ollama" (auto-detects if not set)
+# LLM_PROVIDER=ollama
+
+# For Claude (paid API):
+# ANTHROPIC_API_KEY=sk-ant-api03-your-key-here
+
+# For Ollama (free, local):
+# Install: https://ollama.com
+# Then: ollama pull llama3.1
+# No API key needed!

main.py

@@ -27,7 +27,7 @@ from rich.text import Text
 # Add src to path
 sys.path.insert(0, str(Path(__file__).parent / "src"))
 
-from soci.engine.llm import ClaudeClient
+from soci.engine.llm import create_llm_client
 from soci.engine.simulation import Simulation
 from soci.persistence.database import Database
 from soci.persistence.snapshots import save_simulation, load_simulation
@@ -118,16 +118,21 @@ async def run_simulation(
     max_agents: int = 20,
     tick_delay: float = 0.5,
     resume: bool = False,
+    provider: str = "",
+    model: str = "",
 ) -> None:
     """Run the simulation with a live Rich dashboard."""
     # Initialize
     console.print("[bold blue]Initializing Soci City Simulation...[/]")
 
     try:
-        llm = ClaudeClient()
-    except ValueError as e:
+        llm = create_llm_client(
+            provider=provider or None,
+            model=model or None,
+        )
+        console.print(f"[green]LLM provider: {llm.provider} (model: {llm.default_model})[/]")
+    except (ValueError, ConnectionError) as e:
         console.print(f"[bold red]Error: {e}[/]")
-        console.print("Copy .env.example to .env and add your ANTHROPIC_API_KEY.")
         return
 
     db = Database()
@@ -212,6 +217,10 @@ def main():
     parser.add_argument("--agents", type=int, default=20, help="Max number of agents (default: 20)")
     parser.add_argument("--speed", type=float, default=0.5, help="Delay between ticks in seconds (default: 0.5)")
     parser.add_argument("--resume", action="store_true", help="Resume from last save")
+    parser.add_argument("--provider", type=str, default="", choices=["", "claude", "ollama"],
+                        help="LLM provider: claude or ollama (default: auto-detect)")
+    parser.add_argument("--model", type=str, default="",
+                        help="Model name (e.g. llama3.1, mistral, qwen2.5)")
     args = parser.parse_args()
 
     Path("data").mkdir(exist_ok=True)
@@ -226,6 +235,8 @@ def main():
         max_agents=args.agents,
         tick_delay=args.speed,
         resume=args.resume,
+        provider=args.provider,
+        model=args.model,
     ))
 
 

src/soci/api/server.py

@@ -11,7 +11,7 @@ from typing import Optional
 from fastapi import FastAPI
 from fastapi.middleware.cors import CORSMiddleware
 
-from soci.engine.llm import ClaudeClient
+from soci.engine.llm import create_llm_client
 from soci.engine.simulation import Simulation
 from soci.persistence.database import Database
 from soci.persistence.snapshots import load_simulation, save_simulation
@@ -63,7 +63,7 @@ async def lifespan(app: FastAPI):
 
     # Start up
     logger.info("Starting Soci API server...")
-    llm = ClaudeClient()
+    llm = create_llm_client()
     db = Database()
     await db.connect()
     _database = db

src/soci/engine/llm.py

@@ -1,4 +1,4 @@
-"""LLM client — Claude API wrapper with model routing, cost tracking, and prompt templates."""
+"""LLM client — supports Claude API and Ollama (local LLMs) with model routing and cost tracking."""
 
 from __future__ import annotations
 
@@ -9,15 +9,26 @@ import time
 from dataclasses import dataclass, field
 from typing import Optional
 
-import anthropic
+import httpx
 
 logger = logging.getLogger(__name__)
 
-# Model IDs
+# --- Provider constants ---
+PROVIDER_CLAUDE = "claude"
+PROVIDER_OLLAMA = "ollama"
+
+# Claude model IDs
 MODEL_SONNET = "claude-sonnet-4-5-20250929"
 MODEL_HAIKU = "claude-haiku-4-5-20251001"
 
-# Approximate cost per 1M tokens (USD)
+# Ollama model IDs (popular open-source models)
+MODEL_LLAMA = "llama3.1"
+MODEL_LLAMA_SMALL = "llama3.2"
+MODEL_MISTRAL = "mistral"
+MODEL_QWEN = "qwen2.5"
+MODEL_GEMMA = "gemma2"
+
+# Approximate cost per 1M tokens (USD) — Ollama is free
 COST_PER_1M = {
     MODEL_SONNET: {"input": 3.0, "output": 15.0},
     MODEL_HAIKU: {"input": 0.80, "output": 4.0},
@@ -48,7 +59,7 @@ class LLMUsage:
     def estimated_cost_usd(self) -> float:
         total = 0.0
         for model, tokens in self.tokens_by_model.items():
-            costs = COST_PER_1M.get(model, {"input": 3.0, "output": 15.0})
+            costs = COST_PER_1M.get(model, {"input": 0.0, "output": 0.0})
             total += tokens["input"] / 1_000_000 * costs["input"]
             total += tokens["output"] / 1_000_000 * costs["output"]
         return total
@@ -65,8 +76,35 @@ class LLMUsage:
         return "\n".join(lines)
 
 
+def _parse_json_response(text: str) -> dict:
+    """Extract JSON from an LLM response, handling markdown blocks and extra text."""
+    text = text.strip()
+    # Handle markdown code blocks
+    if text.startswith("```"):
+        lines = text.split("\n")
+        text = "\n".join(lines[1:-1]) if len(lines) > 2 else text
+        text = text.strip()
+    try:
+        return json.loads(text)
+    except json.JSONDecodeError:
+        # Try to find JSON object in the response
+        start = text.find("{")
+        end = text.rfind("}") + 1
+        if start >= 0 and end > start:
+            try:
+                return json.loads(text[start:end])
+            except json.JSONDecodeError:
+                pass
+        logger.warning(f"Failed to parse JSON from LLM response: {text[:200]}")
+        return {}
+
+
+# ============================================================
+# Claude (Anthropic API) Client
+# ============================================================
+
 class ClaudeClient:
-    """Wrapper around the Anthropic Claude API with model routing and retries."""
+    """Wrapper around the Anthropic Claude API."""
 
     def __init__(
         self,
@@ -74,6 +112,7 @@ class ClaudeClient:
         default_model: str = MODEL_HAIKU,
         max_retries: int = 3,
     ) -> None:
+        import anthropic
         self.api_key = api_key or os.environ.get("ANTHROPIC_API_KEY", "")
         if not self.api_key:
             raise ValueError(
@@ -83,6 +122,7 @@ class ClaudeClient:
         self.default_model = default_model
         self.max_retries = max_retries
         self.usage = LLMUsage()
+        self.provider = PROVIDER_CLAUDE
 
     async def complete(
         self,
@@ -92,7 +132,7 @@ class ClaudeClient:
         temperature: float = 0.7,
         max_tokens: int = 1024,
     ) -> str:
-        """Send a message to Claude and return the text response."""
+        import anthropic
         model = model or self.default_model
 
         for attempt in range(self.max_retries):
@@ -104,7 +144,6 @@ class ClaudeClient:
                     system=system,
                     messages=[{"role": "user", "content": user_message}],
                 )
-                # Track usage
                 self.usage.record(
                     model=model,
                     input_tokens=response.usage.input_tokens,
@@ -121,7 +160,123 @@ class ClaudeClient:
                 if attempt == self.max_retries - 1:
                     raise
                 time.sleep(1)
+        return ""
+
+    async def complete_json(
+        self,
+        system: str,
+        user_message: str,
+        model: Optional[str] = None,
+        temperature: float = 0.7,
+        max_tokens: int = 1024,
+    ) -> dict:
+        json_instruction = (
+            "\n\nRespond ONLY with valid JSON. No markdown, no explanation, no extra text. "
+            "Just the JSON object."
+        )
+        text = await self.complete(
+            system=system,
+            user_message=user_message + json_instruction,
+            model=model,
+            temperature=temperature,
+            max_tokens=max_tokens,
+        )
+        return _parse_json_response(text)
+
 
+# ============================================================
+# Ollama (Local LLM) Client
+# ============================================================
+
+class OllamaClient:
+    """Wrapper around Ollama's local API for running open-source LLMs.
+
+    Ollama serves models locally at http://localhost:11434.
+    Install: https://ollama.com
+    Pull a model: ollama pull llama3.1
+    """
+
+    def __init__(
+        self,
+        base_url: str = "http://localhost:11434",
+        default_model: str = MODEL_LLAMA,
+        max_retries: int = 2,
+    ) -> None:
+        self.base_url = base_url.rstrip("/")
+        self.default_model = default_model
+        self.max_retries = max_retries
+        self.usage = LLMUsage()
+        self.provider = PROVIDER_OLLAMA
+        self._http = httpx.Client(timeout=120.0)
+
+    async def complete(
+        self,
+        system: str,
+        user_message: str,
+        model: Optional[str] = None,
+        temperature: float = 0.7,
+        max_tokens: int = 1024,
+    ) -> str:
+        """Send a message to the local Ollama model."""
+        model = model or self.default_model
+        # Map Claude model names to Ollama models
+        model = self._map_model(model)
+
+        payload = {
+            "model": model,
+            "messages": [
+                {"role": "system", "content": system},
+                {"role": "user", "content": user_message},
+            ],
+            "stream": False,
+            "options": {
+                "temperature": temperature,
+                "num_predict": max_tokens,
+            },
+        }
+
+        for attempt in range(self.max_retries):
+            try:
+                response = self._http.post(
+                    f"{self.base_url}/api/chat",
+                    json=payload,
+                )
+                response.raise_for_status()
+                data = response.json()
+
+                # Track usage
+                input_tokens = data.get("prompt_eval_count", 0)
+                output_tokens = data.get("eval_count", 0)
+                self.usage.record(model, input_tokens, output_tokens)
+
+                return data.get("message", {}).get("content", "")
+
+            except httpx.ConnectError:
+                msg = (
+                    f"Cannot connect to Ollama at {self.base_url}. "
+                    "Make sure Ollama is running: 'ollama serve'"
+                )
+                logger.error(msg)
+                if attempt == self.max_retries - 1:
+                    raise ConnectionError(msg)
+                time.sleep(1)
+            except httpx.HTTPStatusError as e:
+                if e.response.status_code == 404:
+                    msg = (
+                        f"Model '{model}' not found in Ollama. "
+                        f"Pull it first: 'ollama pull {model}'"
+                    )
+                    logger.error(msg)
+                    raise ValueError(msg)
+                logger.error(f"Ollama API error: {e}")
+                if attempt == self.max_retries - 1:
+                    raise
+                time.sleep(1)
+            except Exception as e:
+                logger.error(f"Ollama error: {e}")
+                if attempt == self.max_retries - 1:
+                    raise
+                time.sleep(1)
         return ""
 
     async def complete_json(
@@ -132,8 +287,6 @@ class ClaudeClient:
         temperature: float = 0.7,
         max_tokens: int = 1024,
     ) -> dict:
-        """Send a message and parse the response as JSON."""
-        # Add JSON instruction to the prompt
         json_instruction = (
             "\n\nRespond ONLY with valid JSON. No markdown, no explanation, no extra text. "
             "Just the JSON object."
@@ -145,25 +298,52 @@ class ClaudeClient:
             temperature=temperature,
             max_tokens=max_tokens,
         )
-        # Try to extract JSON from the response
-        text = text.strip()
-        # Handle markdown code blocks
-        if text.startswith("```"):
-            lines = text.split("\n")
-            text = "\n".join(lines[1:-1]) if len(lines) > 2 else text
-        try:
-            return json.loads(text)
-        except json.JSONDecodeError:
-            # Try to find JSON in the response
-            start = text.find("{")
-            end = text.rfind("}") + 1
-            if start >= 0 and end > start:
-                try:
-                    return json.loads(text[start:end])
-                except json.JSONDecodeError:
-                    pass
-            logger.warning(f"Failed to parse JSON from LLM response: {text[:200]}")
-            return {}
+        return _parse_json_response(text)
+
+    def _map_model(self, model: str) -> str:
+        """Map Claude model names to Ollama equivalents so existing code works."""
+        mapping = {
+            MODEL_SONNET: self.default_model,  # Use the main local model
+            MODEL_HAIKU: self.default_model,    # Same model for both (local is free)
+        }
+        return mapping.get(model, model)
+
+
+# ============================================================
+# Factory — create the right client based on config
+# ============================================================
+
+def create_llm_client(
+    provider: Optional[str] = None,
+    model: Optional[str] = None,
+    ollama_url: str = "http://localhost:11434",
+) -> ClaudeClient | OllamaClient:
+    """Create an LLM client based on environment or explicit config.
+
+    Provider detection order:
+    1. Explicit provider argument
+    2. LLM_PROVIDER env var
+    3. If ANTHROPIC_API_KEY is set → Claude
+    4. Default → Ollama (free, local)
+    """
+    if provider is None:
+        provider = os.environ.get("LLM_PROVIDER", "").lower()
+
+    if not provider:
+        # Auto-detect: use Claude if key is set, otherwise Ollama
+        if os.environ.get("ANTHROPIC_API_KEY"):
+            provider = PROVIDER_CLAUDE
+        else:
+            provider = PROVIDER_OLLAMA
+
+    if provider == PROVIDER_CLAUDE:
+        default_model = model or MODEL_HAIKU
+        return ClaudeClient(default_model=default_model)
+    elif provider == PROVIDER_OLLAMA:
+        default_model = model or MODEL_LLAMA
+        return OllamaClient(base_url=ollama_url, default_model=default_model)
+    else:
+        raise ValueError(f"Unknown LLM provider: {provider}. Use 'claude' or 'ollama'.")
 
 
 # --- Prompt Templates ---