diff --git a/README.md b/README.md index b5bb52b..f0950ba 100644 --- a/README.md +++ b/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 diff --git a/conf/base.yaml b/conf/base.yaml index 4c83152..65728d3 100644 --- a/conf/base.yaml +++ b/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 diff --git a/src/simulated_web_agent/agent/gpt.py b/src/simulated_web_agent/agent/gpt.py index 8e7e23a..bfc24a3 100644 --- a/src/simulated_web_agent/agent/gpt.py +++ b/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)