Add local LLM support via Ollama and vLLM providers

README.md

@@ -78,6 +78,48 @@ https://github.com/user-attachments/assets/8f4b352b-1c36-4b16-9d83-b39046357c40
 
 ---
 
+## Using local models (Ollama / vLLM)
+
+UXAgent can run fully locally without any cloud API key, using [Ollama](https://ollama.com/) or a [vLLM](https://docs.vllm.ai/) OpenAI-compatible server. Select the provider with `llm_provider` in `conf/base.yaml` (or `llm_provider=ollama` on the command line) and configure models via environment variables.
+
+### Ollama
+
+```bash
+ollama pull llama3.1            # chat model
+ollama pull nomic-embed-text    # embedding model (required for agent memory)
+
+export OLLAMA_API_BASE=http://localhost:11434   # default
+export OLLAMA_CHAT_MODEL=llama3.1               # fast model ("small")
+export OLLAMA_SLOW_CHAT_MODEL=llama3.1          # deep-reasoning model
+export OLLAMA_EMBEDDING_MODEL=nomic-embed-text
+
+uv run -m src.simulated_web_agent.main --intent "..." --start-url "..." llm_provider=ollama
+```
+
+> [!IMPORTANT]
+> The agent feeds the full simplified web page into each prompt, so raise Ollama's
+> context window (default is 4096 tokens), e.g. `export OLLAMA_CONTEXT_LENGTH=32768`
+> before starting `ollama serve`, or bake `num_ctx` into a Modelfile.
+
+### vLLM
+
+```bash
+export VLLM_API_BASE=http://localhost:8000/v1
+export VLLM_API_KEY=dummy-key                        # only if your server requires one
+export VLLM_CHAT_MODEL=Qwen/Qwen2.5-32B-Instruct     # must match the served model name
+export VLLM_SLOW_CHAT_MODEL=Qwen/Qwen2.5-32B-Instruct
+export LLM_MAX_TOKENS=8192            # cap output tokens to fit the model context
+export EMBEDDING_PROVIDER=ollama      # vLLM usually serves one model; use Ollama for embeddings
+
+uv run -m src.simulated_web_agent.main --intent "..." --start-url "..." llm_provider=vllm
+```
+
+If your vLLM server also serves an embedding model, set `VLLM_EMBEDDING_MODEL` instead of `EMBEDDING_PROVIDER`.
+
+**Notes on model choice:** every agent step requires reading a long simplified-HTML page and emitting structured JSON. Small models (< 14B) often produce invalid JSON and trigger repeated retries; instruction-tuned models of 32B or larger are recommended for usable simulations.
+
+---
+
 ## Quick Start
 
 ### Running a single agent from the command line

conf/base.yaml

@@ -92,7 +92,9 @@ survey:
 # LLM CONFIGURATION
 # ============================================================================
 # Large Language Model provider settings
-llm_provider: "aws"     # Provider: "openai" or "aws" or "anthropic"
+# Cloud providers: "openai", "aws", "anthropic"
+# Local providers:  "ollama", "vllm" (see "Using local models" in README.md)
+llm_provider: "aws"
 
 # ============================================================================
 # GENERAL SETTINGS

src/simulated_web_agent/agent/gpt.py

@@ -1,5 +1,6 @@
 import asyncio
 import json
+import os
 import time
 from pathlib import Path
 from typing import Any, Dict, cast
@@ -22,10 +23,27 @@
 
 from . import context
 
-provider = "openai"  # "openai" or "aws" or "anthropic"
+load_dotenv()  # load API keys and local LLM settings from .env
+
+provider = "openai"  # "openai" or "aws" or "anthropic" or "ollama" or "vllm"
 
 prompt_dir = Path(__file__).parent.absolute() / "shop_prompts"
 
+# Local LLM settings (used when provider is "ollama" or "vllm").
+# Override via environment variables or a .env file.
+OLLAMA_API_BASE = os.getenv("OLLAMA_API_BASE", "http://localhost:11434")
+OLLAMA_CHAT_MODEL = os.getenv("OLLAMA_CHAT_MODEL", "llama3.1")
+OLLAMA_SLOW_CHAT_MODEL = os.getenv("OLLAMA_SLOW_CHAT_MODEL", OLLAMA_CHAT_MODEL)
+OLLAMA_EMBEDDING_MODEL = os.getenv("OLLAMA_EMBEDDING_MODEL", "nomic-embed-text")
+
+VLLM_API_BASE = os.getenv("VLLM_API_BASE", "http://localhost:8000/v1")
+VLLM_API_KEY = os.getenv("VLLM_API_KEY", "dummy-key")
+VLLM_CHAT_MODEL = os.getenv("VLLM_CHAT_MODEL", "Qwen/Qwen2.5-32B-Instruct")
+VLLM_SLOW_CHAT_MODEL = os.getenv("VLLM_SLOW_CHAT_MODEL", VLLM_CHAT_MODEL)
+VLLM_EMBEDDING_MODEL = os.getenv(
+    "VLLM_EMBEDDING_MODEL", "intfloat/e5-mistral-7b-instruct"
+)
+
 chat_router = Router(
     model_list=[
         {
@@ -77,6 +95,36 @@
                 },
             },
         },
+        {
+            "model_name": "ollama",
+            "litellm_params": {
+                "model": f"ollama_chat/{OLLAMA_CHAT_MODEL}",
+                "api_base": OLLAMA_API_BASE,
+            },
+        },
+        {
+            "model_name": "ollama_thinking",
+            "litellm_params": {
+                "model": f"ollama_chat/{OLLAMA_CHAT_MODEL}",
+                "api_base": OLLAMA_API_BASE,
+            },
+        },
+        {
+            "model_name": "vllm",
+            "litellm_params": {
+                "model": f"hosted_vllm/{VLLM_CHAT_MODEL}",
+                "api_base": VLLM_API_BASE,
+                "api_key": VLLM_API_KEY,
+            },
+        },
+        {
+            "model_name": "vllm_thinking",
+            "litellm_params": {
+                "model": f"hosted_vllm/{VLLM_CHAT_MODEL}",
+                "api_base": VLLM_API_BASE,
+                "api_key": VLLM_API_KEY,
+            },
+        },
     ]
 )
 
@@ -125,6 +173,36 @@
                 },
             },
         },
+        {
+            "model_name": "ollama",
+            "litellm_params": {
+                "model": f"ollama_chat/{OLLAMA_SLOW_CHAT_MODEL}",
+                "api_base": OLLAMA_API_BASE,
+            },
+        },
+        {
+            "model_name": "ollama_thinking",
+            "litellm_params": {
+                "model": f"ollama_chat/{OLLAMA_SLOW_CHAT_MODEL}",
+                "api_base": OLLAMA_API_BASE,
+            },
+        },
+        {
+            "model_name": "vllm",
+            "litellm_params": {
+                "model": f"hosted_vllm/{VLLM_SLOW_CHAT_MODEL}",
+                "api_base": VLLM_API_BASE,
+                "api_key": VLLM_API_KEY,
+            },
+        },
+        {
+            "model_name": "vllm_thinking",
+            "litellm_params": {
+                "model": f"hosted_vllm/{VLLM_SLOW_CHAT_MODEL}",
+                "api_base": VLLM_API_BASE,
+                "api_key": VLLM_API_KEY,
+            },
+        },
     ]
 )
 
@@ -142,11 +220,25 @@
                 "truncate": "END",
             },
         },
+        {
+            "model_name": "ollama",
+            "litellm_params": {
+                "model": f"ollama/{OLLAMA_EMBEDDING_MODEL}",
+                "api_base": OLLAMA_API_BASE,
+            },
+        },
+        {
+            "model_name": "vllm",
+            "litellm_params": {
+                "model": f"hosted_vllm/{VLLM_EMBEDDING_MODEL}",
+                "api_base": VLLM_API_BASE,
+                "api_key": VLLM_API_KEY,
+            },
+        },
     ]
 )
 
 
-load_dotenv()  # load anthropic api key from .env
 anthropic_client = anthropic.Anthropic()
 anthropic_model = "claude-sonnet-4-20250514"
 
@@ -250,8 +342,13 @@ async def async_chat(
         router_model = provider + "_thinking"
     else:
         router_model = provider
-    if json_mode and provider == "openai":
+    if json_mode and provider in ("openai", "ollama", "vllm"):
         call_kwargs["response_format"] = {"type": "json_object"}
+    # Cap max_tokens for backends with smaller context windows (e.g. local
+    # models served by vLLM, which reject requests exceeding the model limit).
+    max_tokens_cap = os.getenv("LLM_MAX_TOKENS")
+    if max_tokens_cap:
+        max_tokens = min(max_tokens, int(max_tokens_cap))
     response = await router.acompletion(
         model=router_model,
         messages=messages,
@@ -312,7 +409,7 @@ def chat(
             if isinstance(enable_thinking, int)
             else 1024,
         }
-    if json_mode and provider == "openai":
+    if json_mode and provider in ("openai", "ollama", "vllm"):
         call_kwargs["response_format"] = {"type": "json_object"}
 
     try:
@@ -333,11 +430,15 @@ async def embed_text(texts: list[str]) -> list[list[float]]:
     """
     Embed a list of texts using the provider configured in /src/simulated_web_agent/agent/gpt.py
 
+    Set the EMBEDDING_PROVIDER environment variable to embed with a different
+    provider than the chat provider (e.g. chat on vLLM, embeddings on Ollama).
+
     Returns:
         List of list[float] representing each of the embedded texts
     """
     try:
-        response = await embed_router.aembedding(model=provider, input=texts)
+        embed_provider = os.getenv("EMBEDDING_PROVIDER") or provider
+        response = await embed_router.aembedding(model=embed_provider, input=texts)
         return [e["embedding"] for e in response.data]
     except Exception as e:
         print(texts)