diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 1173b0ef..69bc108f 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -53,6 +53,23 @@ jobs: python tests/test_ci_quick.py python -m pytest tests/test_plugins.py -v --tb=short || python tests/test_plugins.py python -m pytest tests/test_proxy_plugin.py -v --tb=short + python -m pytest tests/test_streaming_hints.py -v --tb=short + python -m pytest tests/test_plugin_gating.py -v --tb=short + python -m pytest tests/test_max_completion_tokens.py -v --tb=short + python -m pytest tests/test_bon_candidate_filtering.py -v --tb=short + python -m pytest tests/test_compact_plugin.py -v --tb=short + python -m pytest tests/test_completions_endpoint.py -v --tb=short + python -m pytest tests/test_entropy_decoding_sampling.py -v --tb=short + python -m pytest tests/test_mars_agent_routing.py -v --tb=short + python -m pytest tests/test_mars_aggregation_sampling.py -v --tb=short + python -m pytest tests/test_mars_imo25.py -v --tb=short + python -m pytest tests/test_mars_workspace_verification.py -v --tb=short + python -m pytest tests/test_mcts_selection.py -v --tb=short + python -m pytest tests/test_privacy_plugin_performance.py -v --tb=short + python -m pytest tests/test_readurls_ssrf.py -v --tb=short + python -m pytest tests/test_rstar_backpropagation.py -v --tb=short + python -m pytest tests/test_rstar_numeric_answers.py -v --tb=short + python -m pytest tests/test_z3_solver_guard.py -v --tb=short python tests/test_approaches.py python tests/test_reasoning_simple.py python tests/test_batching.py @@ -125,7 +142,7 @@ jobs: OPTILLM_API_KEY=optillm python tests/test_reasoning_integration.py OPTILLM_API_KEY=optillm python tests/test_json_plugin.py OPTILLM_API_KEY=optillm python tests/test_n_parameter.py - OPTILLM_API_KEY=optillm python -m pytest tests/test_api_compatibility.py -v --tb=short || echo "API compatibility tests require pytest" + OPTILLM_API_KEY=optillm python -m pytest tests/test_api_compatibility.py -v --tb=short OPTILLM_API_KEY=optillm python tests/test.py --approaches none --single-test "Simple Math Problem" || echo "Main test completed" # Run SSL config tests (no server needed but requires proper env setup) @@ -256,7 +273,6 @@ jobs: mcp-tests: runs-on: ubuntu-latest needs: unit-tests - if: github.event_name == 'push' && github.ref == 'refs/heads/main' # Only run on main branch pushes (secrets available) strategy: matrix: python-version: ['3.12'] @@ -296,5 +312,6 @@ jobs: echo "MCP tests completed successfully!" env: OPTILLM_API_KEY: optillm + # Not available on fork PRs, where the one test that needs it skips itself GITHUB_TOKEN: ${{ secrets.GH_TOKEN }} HF_TOKEN: ${{ secrets.HF_TOKEN }} \ No newline at end of file diff --git a/README.md b/README.md index 40c4f20e..e6c2a6a5 100644 --- a/README.md +++ b/README.md @@ -1,11 +1,11 @@ # OptiLLM

- OptiLLM Logo + OptiLLM Logo

- ๐Ÿš€ 2-10x accuracy improvements on reasoning tasks with zero training + 2-10x accuracy improvements on reasoning tasks with zero training

@@ -15,110 +15,85 @@

- ๐Ÿค— HuggingFace Space โ€ข - ๐Ÿ““ Colab Demo โ€ข - ๐Ÿ’ฌ Discussions + HuggingFace Space | + Colab Demo | + Discussions

---- +**OptiLLM** is an OpenAI API-compatible optimizing inference proxy that implements 20+ state-of-the-art techniques to dramatically improve LLM accuracy and performance on reasoning tasks, without requiring any model training or fine-tuning. -**OptiLLM** is an OpenAI API-compatible optimizing inference proxy that implements 20+ state-of-the-art techniques to dramatically improve LLM accuracy and performance on reasoning tasks - without requiring any model training or fine-tuning. +It is possible to beat the frontier models using these techniques across diverse tasks by doing additional compute at inference time. A good example of how to combine such techniques is the [CePO approach](optillm/cepo) from Cerebras. -It is possible to beat the frontier models using these techniques across diverse tasks by doing additional compute at inference time. A good example of how to combine such techniques together is the [CePO approach](optillm/cepo) from Cerebras. +## Key features -## โœจ Key Features +- **Instant improvements**: 2-10x better accuracy on math, coding, and logical reasoning +- **Drop-in replacement**: Works with any OpenAI-compatible API endpoint +- **20+ optimization techniques**: From simple best-of-N to advanced MCTS and planning +- **Zero training required**: Just proxy your existing API calls through OptiLLM +- **Production ready**: Used in production by companies and researchers worldwide +- **Multi-provider**: Supports OpenAI, Anthropic, Google, Cerebras, and 100+ models via LiteLLM -- **๐ŸŽฏ Instant Improvements**: 2-10x better accuracy on math, coding, and logical reasoning -- **๐Ÿ”Œ Drop-in Replacement**: Works with any OpenAI-compatible API endpoint -- **๐Ÿง  20+ Optimization Techniques**: From simple best-of-N to advanced MCTS and planning -- **๐Ÿ“ฆ Zero Training Required**: Just proxy your existing API calls through OptiLLM -- **โšก Production Ready**: Used in production by companies and researchers worldwide -- **๐ŸŒ Multi-Provider**: Supports OpenAI, Anthropic, Google, Cerebras, and 100+ models via LiteLLM - -## ๐Ÿš€ Quick Start - -Get powerful reasoning improvements in 3 simple steps: +## Quick start ```bash -# 1. Install OptiLLM pip install optillm - -# 2. Start the server export OPENAI_API_KEY="your-key-here" optillm - -# 3. Use with any OpenAI client - just change the model name! ``` +Point any OpenAI client at the proxy and prefix the model name with a technique: + ```python from openai import OpenAI client = OpenAI(base_url="http://localhost:8000/v1") -# Add 'moa-' prefix for Mixture of Agents optimization response = client.chat.completions.create( - model="moa-gpt-4o-mini", # This gives you GPT-4o performance from GPT-4o-mini! - messages=[{"role": "user", "content": "Solve: If 2x + 3 = 7, what is x?"}] + model="moa-gpt-4o-mini", # This gives you GPT-4o performance from GPT-4o-mini + messages=[{"role": "user", "content": "Solve: If 2x + 3 = 7, what is x?"}], ) ``` -**Before OptiLLM**: "x = 1" โŒ -**After OptiLLM**: "Let me work through this step by step: 2x + 3 = 7, so 2x = 4, therefore x = 2" โœ… +**Before OptiLLM**: "x = 1" +**After OptiLLM**: "Let me work through this step by step: 2x + 3 = 7, so 2x = 4, therefore x = 2" -## ๐Ÿ“Š Proven Results +## Results -OptiLLM delivers measurable improvements across diverse benchmarks: - -| Technique | Base Model | Improvement | Benchmark | +| Technique | Base model | Improvement | Benchmark | |-----------|------------|-------------|-----------| -| **MARS** | Gemini 2.5 Flash Lite | **+30.0 points** | AIME 2025 (43.3โ†’73.3) | -| **CePO** | Llama 3.3 70B | **+18.6 points** | Math-L5 (51.0โ†’69.6) | -| **AutoThink** | DeepSeek-R1-1.5B | **+9.34 points** | GPQA-Diamond (21.72โ†’31.06) | -| **LongCePO** | Llama 3.3 70B | **+13.6 points** | InfiniteBench (58.0โ†’71.6) | -| **MOA** | GPT-4o-mini | **Matches GPT-4** | Arena-Hard-Auto | -| **PlanSearch** | GPT-4o-mini | **+20% pass@5** | LiveCodeBench | +| MARS | Gemini 2.5 Flash Lite | +30.0 points | AIME 2025 (43.3 to 73.3) | +| CePO | Llama 3.3 70B | +18.6 points | Math-L5 (51.0 to 69.6) | +| LongCePO | Llama 3.3 70B | +13.6 points | InfiniteBench (58.0 to 71.6) | +| AutoThink | DeepSeek-R1-1.5B | +9.34 points | GPQA-Diamond (21.72 to 31.06) | +| PlanSearch | GPT-4o-mini | +20% pass@5 | LiveCodeBench | +| MOA | GPT-4o-mini | Matches GPT-4 | Arena-Hard-Auto | -*Full benchmark results [below](#sota-results-on-benchmarks-with-optillm)* โฌ‡๏ธ +See [benchmark results](#benchmark-results) for the full tables. -## ๐Ÿ—๏ธ Installation +## Installation -### Using pip +**pip** ```bash pip install optillm -optillm -2024-10-22 07:45:05,612 - INFO - Loaded plugin: privacy -2024-10-22 07:45:06,293 - INFO - Loaded plugin: memory -2024-10-22 07:45:06,293 - INFO - Starting server with approach: auto ``` -### Using docker +**Docker** ```bash docker pull ghcr.io/algorithmicsuperintelligence/optillm:latest docker run -p 8000:8000 ghcr.io/algorithmicsuperintelligence/optillm:latest -2024-10-22 07:45:05,612 - INFO - Loaded plugin: privacy -2024-10-22 07:45:06,293 - INFO - Loaded plugin: memory -2024-10-22 07:45:06,293 - INFO - Starting server with approach: auto ``` -**Available Docker image variants:** +| Image tag | Contents | +|-----------|----------| +| `latest` | All dependencies for local inference and plugins | +| `latest-proxy` | Lightweight image without local inference | +| `latest-offline` | Self-contained image with pre-downloaded models (spaCy) for fully offline operation | -- **Full image** (`latest`): Includes all dependencies for local inference and plugins -- **Proxy-only** (`latest-proxy`): Lightweight image without local inference capabilities -- **Offline** (`latest-offline`): Self-contained image with pre-downloaded models (spaCy) for fully offline operation +To build and run from the repository, set your keys in `docker-compose.yaml` or a `.env` file and run `docker compose up -d`. -```bash -# Proxy-only (smallest) -docker pull ghcr.io/algorithmicsuperintelligence/optillm:latest-proxy - -# Offline (largest, includes pre-downloaded models) -docker pull ghcr.io/algorithmicsuperintelligence/optillm:latest-offline -``` - -### Install from source - -Clone the repository with `git` and use `pip install` to setup the dependencies. +**From source** ```bash git clone https://github.com/algorithmicsuperintelligence/optillm.git @@ -126,347 +101,187 @@ cd optillm python3 -m venv .venv source .venv/bin/activate pip install -r requirements.txt -``` - -## ๐Ÿ”’ SSL Configuration - -OptILLM supports SSL certificate verification configuration for working with self-signed certificates or corporate proxies. - -**Disable SSL verification (development only):** -```bash -# Command line -optillm --no-ssl-verify - -# Environment variable -export OPTILLM_SSL_VERIFY=false -optillm -``` - -**Use custom CA certificate:** -```bash -# Command line -optillm --ssl-cert-path /path/to/ca-bundle.crt - -# Environment variable -export OPTILLM_SSL_CERT_PATH=/path/to/ca-bundle.crt -optillm -``` - -โš ๏ธ **Security Note**: Disabling SSL verification is insecure and should only be used in development. For production environments with custom CAs, use `--ssl-cert-path` instead. See [SSL_CONFIGURATION.md](SSL_CONFIGURATION.md) for details. - -## ๐Ÿ”’ URL Fetch Security (readurls plugin) - -The `readurls` plugin fetches URLs found in prompts and blocks requests to non-public destinations (private, loopback, link-local, reserved addresses such as the cloud metadata IP `169.254.169.254`) to prevent SSRF. Every redirect hop is re-validated and unresolvable hosts are rejected. - -If your deployment intentionally needs to read internal documentation pages, you can explicitly opt out: - -```bash -# Command line -optillm --readurls-allow-internal - -# Environment variable -export OPTILLM_READURLS_ALLOW_INTERNAL=true -optillm -``` - -โš ๏ธ **Security Note**: Only enable this on trusted networks where internal URL fetching is actually required. - -**Known limitation (DNS rebinding)**: the guard validates DNS resolution before each request, but the underlying HTTP client performs a second resolution when it opens the connection, so a domain under attacker control could rebind to an internal address between the two lookups. For high-security deployments, additionally restrict the network egress of the optillm process (firewall or service-mesh policy) instead of relying on this check alone. - -## Implemented techniques - -| Approach | Slug | Description | -| ------------------------------------ | ------------------ | ---------------------------------------------------------------------------------------------- | -| [MARS (Multi-Agent Reasoning System)](optillm/mars) | `mars` | Multi-agent reasoning with diverse temperature exploration, cross-verification, and iterative improvement | -| [Cerebras Planning and Optimization](optillm/cepo) | `cepo` | Combines Best of N, Chain-of-Thought, Self-Reflection, Self-Improvement, and various prompting techniques | -| CoT with Reflection | `cot_reflection` | Implements chain-of-thought reasoning with \, \ and \ sections | -| PlanSearch | `plansearch` | Implements a search algorithm over candidate plans for solving a problem in natural language | -| ReRead | `re2` | Implements rereading to improve reasoning by processing queries twice | -| Self-Consistency | `self_consistency` | Implements an advanced self-consistency method | -| Z3 Solver | `z3` | Utilizes the Z3 theorem prover for logical reasoning | -| R* Algorithm | `rstar` | Implements the R* algorithm for problem-solving | -| LEAP | `leap` | Learns task-specific principles from few shot examples | -| Round Trip Optimization | `rto` | Optimizes responses through a round-trip process | -| Best of N Sampling | `bon` | Generates multiple responses and selects the best one | -| Mixture of Agents | `moa` | Combines responses from multiple critiques | -| Monte Carlo Tree Search | `mcts` | Uses MCTS for decision-making in chat responses | -| PV Game | `pvg` | Applies a prover-verifier game approach at inference time | -| [Deep Confidence](optillm/deepconf) | N/A for proxy | Implements confidence-guided reasoning with multiple intensity levels for enhanced accuracy | -| CoT Decoding | N/A for proxy | Implements chain-of-thought decoding to elicit reasoning without explicit prompting | -| Entropy Decoding | N/A for proxy | Implements adaptive sampling based on the uncertainty of tokens during generation | -| Thinkdeeper | N/A for proxy | Implements the `reasoning_effort` param from OpenAI for reasoning models like DeepSeek R1 | -| [AutoThink](optillm/autothink) | N/A for proxy | Combines query complexity classification with steering vectors to enhance reasoning | - -## Implemented plugins - -| Plugin | Slug | Description | -| ----------------------- | ------------------ | ---------------------------------------------------------------------------------------------- | -| [System Prompt Learning](optillm/plugins/spl) | `spl` | Implements what [Andrej Karpathy called the third paradigm](https://x.com/karpathy/status/1921368644069765486) for LLM learning, this enables the model to acquire program solving knowledge and strategies | -| [Deep Think](optillm/plugins/deepthink) | `deepthink` | Implements a Gemini-like Deep Think approach using inference time scaling for reasoning LLMs | -| [Long-Context Cerebras Planning and Optimization](optillm/plugins/longcepo) | `longcepo` | Combines planning and divide-and-conquer processing of long documents to enable infinite context | -| Majority Voting | `majority_voting` | Generates k candidate solutions and selects the most frequent answer through majority voting (default k=6) | -| MCP Client | `mcp` | Implements the model context protocol (MCP) client, enabling you to use any LLM with any MCP Server | -| Router | `router` | Uses the [optillm-modernbert-large](https://huggingface.co/codelion/optillm-modernbert-large) model to route requests to different approaches based on the user prompt | -| Chain-of-Code | `coc` | Implements a chain of code approach that combines CoT with code execution and LLM based code simulation | -| Memory | `memory` | Implements a short term memory layer, enables you to use unbounded context length with any LLM. Set `OPTILLM_MEMORY_FILE` to opt in to file-backed persistence so memories survive across requests | -| Privacy | `privacy` | Anonymize PII data in request and deanonymize it back to original value in response | -| Read URLs | `readurls` | Reads all URLs found in the request, fetches the content at the URL and adds it to the context | -| Execute Code | `executecode` | Enables use of code interpreter to execute python code in requests and LLM generated responses | -| JSON | `json` | Enables structured outputs using the outlines library, supports pydantic types and JSON schema | -| GenSelect | `genselect` | Generative Solution Selection - generates multiple candidates and selects the best based on quality criteria | -| Web Search | `web_search` | Performs Google searches using Chrome automation (Selenium) to gather search results and URLs | -| [Deep Research](optillm/plugins/deep_research) | `deep_research` | Implements Test-Time Diffusion Deep Researcher (TTD-DR) for comprehensive research reports using iterative refinement | -| [Proxy](optillm/plugins/proxy) | `proxy` | Load balancing and failover across multiple LLM providers with health monitoring and round-robin routing | - -We support all major LLM providers and models for inference. You need to set the correct environment variable and the proxy will pick the corresponding client. - -| Provider | Required Environment Variables | Additional Notes | -|----------|-------------------------------|------------------| -| OptiLLM | `OPTILLM_API_KEY` | Uses the inbuilt local server for inference, supports logprobs and decoding techniques like `cot_decoding` & `entropy_decoding` | -| OpenAI | `OPENAI_API_KEY` | You can use this with any OpenAI compatible endpoint (e.g. OpenRouter) by setting the `base_url` | -| Cerebras | `CEREBRAS_API_KEY` | You can use this for fast inference with supported models, see [docs for details](https://inference-docs.cerebras.ai/introduction) | -| Azure OpenAI | `AZURE_OPENAI_API_KEY`
`AZURE_API_VERSION`
`AZURE_API_BASE` | - | -| Azure OpenAI (Managed Identity) | `AZURE_API_VERSION`
`AZURE_API_BASE` | Login required using `az login`, see [docs for details](https://learn.microsoft.com/en-us/azure/ai-services/openai/how-to/managed-identity) | -| LiteLLM | depends on the model | See [docs for details](https://docs.litellm.ai/docs/providers) | - -You can then run the optillm proxy as follows. - -```bash python optillm.py -2024-09-06 07:57:14,191 - INFO - Starting server with approach: auto -2024-09-06 07:57:14,191 - INFO - Server configuration: {'approach': 'auto', 'mcts_simulations': 2, 'mcts_exploration': 0.2, 'mcts_depth': 1, 'best_of_n': 3, 'model': 'gpt-4o-mini', 'rstar_max_depth': 3, 'rstar_num_rollouts': 5, 'rstar_c': 1.4, 'base_url': '', 'host': '127.0.0.1'} - * Serving Flask app 'optillm' - * Debug mode: off -2024-09-06 07:57:14,212 - INFO - WARNING: This is a development server. Do not use it in a production deployment. Use a production WSGI server instead. - * Running on http://127.0.0.1:8000 -2024-09-06 07:57:14,212 - INFO - Press CTRL+C to quit ``` -> **Security Note**: By default, optillm binds to `127.0.0.1` (localhost only) for security. To allow external connections (e.g., for Docker or remote access), use `--host 0.0.0.0`. Only do this on trusted networks or with proper authentication configured via `--optillm-api-key`. -## Usage +## Providers -Once the proxy is running, you can use it as a drop in replacement for an OpenAI client by setting the `base_url` as `http://localhost:8000/v1`. +Set the environment variable for your provider and the proxy picks the matching client. -```python -import os -from openai import OpenAI +| Provider | Required environment variables | Notes | +|----------|--------------------------------|-------| +| OptiLLM | `OPTILLM_API_KEY` | Built-in local inference server; supports logprobs and decoding techniques such as `cot_decoding` and `entropy_decoding` | +| OpenAI | `OPENAI_API_KEY` | Also works with any OpenAI-compatible endpoint (e.g. OpenRouter) by setting `--base-url` | +| Cerebras | `CEREBRAS_API_KEY` | See the [Cerebras docs](https://inference-docs.cerebras.ai/introduction) | +| Azure OpenAI | `AZURE_OPENAI_API_KEY`
`AZURE_API_VERSION`
`AZURE_API_BASE` | | +| Azure OpenAI (Managed Identity) | `AZURE_API_VERSION`
`AZURE_API_BASE` | Requires `az login`; see the [Azure docs](https://learn.microsoft.com/en-us/azure/ai-services/openai/how-to/managed-identity) | +| LiteLLM | Depends on the model | See the [LiteLLM docs](https://docs.litellm.ai/docs/providers) | -OPENAI_KEY = os.environ.get("OPENAI_API_KEY") -OPENAI_BASE_URL = "http://localhost:8000/v1" -client = OpenAI(api_key=OPENAI_KEY, base_url=OPENAI_BASE_URL) +Other models work through [LiteLLM](https://docs.litellm.ai/docs/#litellm-python-sdk). For example, set `GEMINI_API_KEY` and call `moa-gemini/gemini-1.5-flash-002`. -response = client.chat.completions.create( - model="moa-gpt-4o", - messages=[ - { - "role": "user", - "content": "Write a Python program to build an RL model to recite text from any position that the user provides, using only numpy." - } - ], - temperature=0.2 -) - -print(response) -``` -The code above applies to both OpenAI and Azure OpenAI, just remember to populate the `OPENAI_API_KEY` env variable with the proper key. -There are multiple ways to control the optimization techniques, they are applied in the follow order of preference: - -- You can control the technique you use for optimization by prepending the slug to the model name `{slug}-model-name`. E.g. in the above code we are using `moa` or mixture of agents as the optimization approach. In the proxy logs you will see the following showing the `moa` is been used with the base model as `gpt-4o-mini`. +## Usage -```bash -2024-09-06 08:35:32,597 - INFO - Using approach moa, with gpt-4o-mini -2024-09-06 08:35:35,358 - INFO - HTTP Request: POST https://api.openai.com/v1/chat/completions "HTTP/1.1 200 OK" -2024-09-06 08:35:39,553 - INFO - HTTP Request: POST https://api.openai.com/v1/chat/completions "HTTP/1.1 200 OK" -2024-09-06 08:35:44,795 - INFO - HTTP Request: POST https://api.openai.com/v1/chat/completions "HTTP/1.1 200 OK" -2024-09-06 08:35:44,797 - INFO - 127.0.0.1 - - [06/Sep/2024 08:35:44] "POST /v1/chat/completions HTTP/1.1" 200 - -``` +### Selecting a technique -- Or, you can pass the slug in the `optillm_approach` field in the `extra_body`. +When the server runs with the default `auto` approach, a technique can be selected in three ways, in this order of preference: -```bash -response = client.chat.completions.create( - model="gpt-4o-mini", - messages=[{ "role": "user","content": "" }], - temperature=0.2, - extra_body={"optillm_approach": "bon|moa|mcts"} +```python +# 1. Model name prefix +client.chat.completions.create(model="moa-gpt-4o-mini", messages=messages) + +# 2. extra_body field +client.chat.completions.create( + model="gpt-4o-mini", + messages=messages, + extra_body={"optillm_approach": "bon|moa|mcts"}, ) -``` -- Or, you can just mention the approach in either your `system` or `user` prompt, within ` ` tags. -```bash -response = client.chat.completions.create( - model="gpt-4o-mini", - messages=[{ "role": "user","content": "re2 How many r's are there in strawberry?" }], - temperature=0.2 +# 3. Tag in the system or user prompt +client.chat.completions.create( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "re2 How many r's are there in strawberry?"}], ) ``` -> [!TIP] -> You can also combine different techniques either by using symbols `&` and `|`. When you use `&` the techniques are processed in the order from left to right in a pipeline -> with response from previous stage used as request to the next. While, with `|` we run all the requests in parallel and generate multiple responses that are returned as a list. +If the server is started with a specific `--approach`, set `model` to the model name only. -Please note that the convention described above works only when the optillm server has been started with inference approach set to `auto`. Otherwise, the `model` attribute in the client request must be set with the model name only. - -We now support all LLM providers (by wrapping around the [LiteLLM sdk](https://docs.litellm.ai/docs/#litellm-python-sdk)). E.g. you can use the Gemini Flash model with `moa` by setting passing the api key in the environment variable `os.environ['GEMINI_API_KEY']` and then calling the model `moa-gemini/gemini-1.5-flash-002`. In the output you will then see that LiteLLM is being used to call the base model. - -```bash -9:43:21 - LiteLLM:INFO: utils.py:2952 - -LiteLLM completion() model= gemini-1.5-flash-002; provider = gemini -2024-09-29 19:43:21,011 - INFO - -LiteLLM completion() model= gemini-1.5-flash-002; provider = gemini -2024-09-29 19:43:21,481 - INFO - HTTP Request: POST https://generativelanguage.googleapis.com/v1beta/models/gemini-1.5-flash-002:generateContent?key=[redacted] "HTTP/1.1 200 OK" -19:43:21 - LiteLLM:INFO: utils.py:988 - Wrapper: Completed Call, calling success_handler -2024-09-29 19:43:21,483 - INFO - Wrapper: Completed Call, calling success_handler -19:43:21 - LiteLLM:INFO: utils.py:2952 - -LiteLLM completion() model= gemini-1.5-flash-002; provider = gemini -``` +### Combining techniques -> [!TIP] -> optillm is a transparent proxy and will work with any LLM API or provider that has an OpenAI API compatible chat completions endpoint, and in turn, optillm also exposes -the same OpenAI API compatible chat completions endpoint. This should allow you to integrate it into any existing tools or frameworks easily. If the LLM you want to use -doesn't have an OpenAI API compatible endpoint (like Google or Anthropic) you can use [LiteLLM proxy server](https://docs.litellm.ai/docs/proxy/quick_start) that supports most LLMs. +- `&` runs techniques as a pipeline from left to right, using each response as the request to the next stage: `cot_reflection&moa` +- `|` runs techniques in parallel and returns the responses as a list: `bon|moa|mcts` -The following sequence diagram illustrates how the request and responses go through optillm. +### How requests flow -![Sequance diagram showing optillm in use](https://raw.githubusercontent.com/algorithmicsuperintelligence/optillm/main/optillm-sequence-diagram.png) +![Sequence diagram showing optillm in use](https://raw.githubusercontent.com/algorithmicsuperintelligence/optillm/main/optillm-sequence-diagram.png) -In the diagram: -- `A` is an existing tool (like [oobabooga](https://github.com/oobabooga/text-generation-webui/)), framework (like [patchwork](https://github.com/patched-codes/patchwork)) -or your own code where you want to use the results from optillm. You can use it directly using any OpenAI client sdk. -- `B` is the optillm service (running directly or in a docker container) that will send requests to the `base_url`. -- `C` is any service providing an OpenAI API compatible chat completions endpoint. +- `A` is an existing tool, framework, or your own code using any OpenAI client SDK. +- `B` is the optillm service, which sends requests to the configured `base_url`. +- `C` is any service providing an OpenAI-compatible chat completions endpoint. -### Local inference server +For an LLM without an OpenAI-compatible endpoint, you can use the [LiteLLM proxy server](https://docs.litellm.ai/docs/proxy/quick_start). -We support loading any HuggingFace model or LoRA directly in optillm. To use the built-in inference server set the `OPTILLM_API_KEY` to any value (e.g. `export OPTILLM_API_KEY="optillm"`) -and then use the same in your OpenAI client. You can pass any HuggingFace model in model field. If it is a private model make sure you set the `HF_TOKEN` environment variable -with your HuggingFace key. We also support adding any number of LoRAs on top of the model by using the `+` separator. +### Local inference -E.g. The following code loads the base model `meta-llama/Llama-3.2-1B-Instruct` and then adds two LoRAs on top - `patched-codes/Llama-3.2-1B-FixVulns` and `patched-codes/Llama-3.2-1B-FastApply`. -You can specify which LoRA to use using the `active_adapter` param in `extra_body` field of OpenAI SDK client. By default we will load the last specified adapter. +Set `OPTILLM_API_KEY` to any value to use the built-in inference server, and use the same value in your client. Any HuggingFace model can be passed as the model name; set `HF_TOKEN` for private models. Add LoRA adapters with the `+` separator and choose one with `active_adapter` (the last one is used by default). ```python -OPENAI_BASE_URL = "http://localhost:8000/v1" -OPENAI_KEY = "optillm" +client = OpenAI(api_key="optillm", base_url="http://localhost:8000/v1") + response = client.chat.completions.create( - model="meta-llama/Llama-3.2-1B-Instruct+patched-codes/Llama-3.2-1B-FastApply+patched-codes/Llama-3.2-1B-FixVulns", - messages=messages, - temperature=0.2, - logprobs = True, - top_logprobs = 3, - extra_body={"active_adapter": "patched-codes/Llama-3.2-1B-FastApply"}, + model="meta-llama/Llama-3.2-1B-Instruct+patched-codes/Llama-3.2-1B-FastApply+patched-codes/Llama-3.2-1B-FixVulns", + messages=messages, + temperature=0.2, + logprobs=True, + top_logprobs=3, + extra_body={"active_adapter": "patched-codes/Llama-3.2-1B-FastApply"}, ) ``` -You can also use the alternate decoding techniques like `cot_decoding` and `entropy_decoding` directly with the local inference server. +Alternate decoding techniques are available on the local server: ```python response = client.chat.completions.create( - model="meta-llama/Llama-3.2-1B-Instruct", - messages=messages, - temperature=0.2, - extra_body={ + model="meta-llama/Llama-3.2-1B-Instruct", + messages=messages, + temperature=0.2, + extra_body={ "decoding": "cot_decoding", # or "entropy_decoding" # CoT specific params "k": 10, "aggregate_paths": True, - # OR Entropy specific params + # OR entropy specific params "top_k": 27, "min_p": 0.03, - } + }, ) ``` -The server also exposes the legacy text completions endpoint `POST /v1/completions` (returning `object: "text_completion"` responses), so OpenAI-legacy clients work against the inbuilt inference server too: +The server also exposes the legacy `POST /v1/completions` endpoint. Streaming is not supported on it yet; use `/v1/chat/completions` for streaming. ```bash curl http://localhost:8000/v1/completions \ -H "Content-Type: application/json" \ -H "Authorization: Bearer $OPTILLM_API_KEY" \ - -d '{ - "model": "meta-llama/Llama-3.2-1B-Instruct", - "prompt": "Write a haiku about recursion.", - "max_tokens": 64 - }' + -d '{"model": "meta-llama/Llama-3.2-1B-Instruct", "prompt": "Write a haiku about recursion.", "max_tokens": 64}' ``` -Streaming (`stream: true`) is not supported on `/v1/completions` yet; use `/v1/chat/completions` for streaming responses. - -### Starting the optillm proxy with an external server (e.g. llama.cpp or ollama) +### External servers (llama.cpp, ollama) -- Set the `OPENAI_API_KEY` env variable to a placeholder value - - e.g. `export OPENAI_API_KEY="sk-no-key"` -- Run `./llama-server -c 4096 -m path_to_model` to start the server with the specified model and a context length of 4096 tokens -- Run `python3 optillm.py --base_url base_url` to start the proxy - - e.g. for llama.cpp, run `python3 optillm.py --base_url http://localhost:8080/v1` +```bash +export OPENAI_API_KEY="sk-no-key" # placeholder +./llama-server -c 4096 -m path_to_model +python3 optillm.py --base-url http://localhost:8080/v1 +``` > [!WARNING] -> The Anthropic API, llama.cpp-server, and ollama currently do not support sampling multiple responses from a model, which limits the available approaches to the following: -> `cot_reflection`, `leap`, `plansearch`, `rstar`, `rto`, `self_consistency`, `re2`, and `z3`. For models on HuggingFace, you can use the built-in local inference server as it supports multiple responses. - -### MCP Plugin - -The Model Context Protocol (MCP) plugin enables OptiLLM to connect with MCP servers, bringing external tools, resources, and prompts into the context of language models. This allows for powerful integrations with filesystem access, database queries, API connections, and more. +> The Anthropic API, llama.cpp-server, and ollama do not support sampling multiple responses from a model, which limits the available approaches to `cot_reflection`, `leap`, `plansearch`, `rstar`, `rto`, `self_consistency`, `re2`, and `z3`. For HuggingFace models, the built-in local inference server supports multiple responses. + +## Use as an LLM gateway + +The `proxy` plugin turns OptiLLM into a self-hosted gateway with load balancing, failover, and health checks across any OpenAI-compatible providers, with no gateway fees. It can also apply any of the techniques below to the requests it routes, combining routing and optimization in one service. See [Using OptiLLM as an LLM gateway](docs/llm-gateway.md) for setup. + +## Techniques + +| Approach | Slug | Description | +|----------|------|-------------| +| [MARS (Multi-Agent Reasoning System)](optillm/mars) | `mars` | Multi-agent reasoning with diverse temperature exploration, cross-verification, and iterative improvement | +| [Cerebras Planning and Optimization](optillm/cepo) | `cepo` | Combines Best of N, Chain-of-Thought, Self-Reflection, Self-Improvement, and various prompting techniques | +| CoT with Reflection | `cot_reflection` | Chain-of-thought reasoning with ``, `` and `` sections | +| PlanSearch | `plansearch` | Search over candidate plans for solving a problem in natural language | +| ReRead | `re2` | Improves reasoning by processing queries twice | +| Self-Consistency | `self_consistency` | An advanced self-consistency method | +| Z3 Solver | `z3` | Uses the Z3 theorem prover for logical reasoning | +| R* Algorithm | `rstar` | The R* algorithm for problem solving | +| LEAP | `leap` | Learns task-specific principles from few-shot examples | +| Round Trip Optimization | `rto` | Optimizes responses through a round-trip process | +| Best of N Sampling | `bon` | Generates multiple responses and selects the best one | +| Mixture of Agents | `moa` | Combines responses from multiple critiques | +| Monte Carlo Tree Search | `mcts` | Uses MCTS for decision-making in chat responses | +| PV Game | `pvg` | Applies a prover-verifier game at inference time | +| [Deep Confidence](optillm/deepconf) | Local only | Confidence-guided reasoning with multiple intensity levels | +| CoT Decoding | Local only | Chain-of-thought decoding to elicit reasoning without explicit prompting | +| Entropy Decoding | Local only | Adaptive sampling based on the uncertainty of tokens during generation | +| Thinkdeeper | Local only | The `reasoning_effort` parameter for reasoning models like DeepSeek R1 | +| [AutoThink](optillm/autothink) | Local only | Combines query complexity classification with steering vectors | + +## Plugins + +| Plugin | Slug | Description | +|--------|------|-------------| +| [System Prompt Learning](optillm/plugins/spl) | `spl` | Implements what [Andrej Karpathy called the third paradigm](https://x.com/karpathy/status/1921368644069765486) for LLM learning, letting the model acquire problem-solving knowledge and strategies | +| [Deep Think](optillm/plugins/deepthink) | `deepthink` | A Gemini-like Deep Think approach using inference-time scaling for reasoning LLMs | +| [Long-Context Cerebras Planning and Optimization](optillm/plugins/longcepo) | `longcepo` | Combines planning and divide-and-conquer processing of long documents | +| Majority Voting | `majority_voting` | Generates k candidate solutions and selects the most frequent answer (default k=6) | +| [MCP Client](#mcp-plugin) | `mcp` | Model Context Protocol client for using any LLM with any MCP server | +| Router | `router` | Uses the [optillm-modernbert-large](https://huggingface.co/codelion/optillm-modernbert-large) model to route requests to approaches based on the prompt | +| Chain-of-Code | `coc` | Combines CoT with code execution and LLM-based code simulation | +| Memory | `memory` | Short-term memory layer for unbounded context length with any LLM. Set `OPTILLM_MEMORY_FILE` to persist memories across requests | +| Compact | `compact` | Compresses conversation context that exceeds a token budget into a structured summary, keeping recent turns verbatim | +| Privacy | `privacy` | Anonymizes PII in the request and restores it in the response | +| Read URLs | `readurls` | Fetches URLs found in the request and adds their content to the context | +| Execute Code | `executecode` | Runs Python code from requests and LLM responses. **Disabled by default**; see [Security](#security) | +| JSON | `json` | Structured outputs using the outlines library; supports pydantic types and JSON schema | +| GenSelect | `genselect` | Generates multiple candidates and selects the best based on quality criteria | +| Web Search | `web_search` | Performs Google searches using Chrome automation (Selenium) | +| [Deep Research](optillm/plugins/deep_research) | `deep_research` | Test-Time Diffusion Deep Researcher (TTD-DR) for research reports using iterative refinement | +| [Proxy](optillm/plugins/proxy) | `proxy` | Load balancing and failover across LLM providers with health monitoring and round-robin routing | +| [Streaming Hints](optillm/plugins/streaming_hints) | `streaming_hints` | Watches reasoning as it is generated and injects relevant hints from a memory bank at moments of uncertainty | + +### MCP plugin + +The [Model Context Protocol](https://modelcontextprotocol.io/) plugin connects OptiLLM to MCP servers, making their tools, resources, and prompts available to the model. When the model calls a tool, the plugin executes it and returns the result to the model. + +Configure servers in `~/.optillm/mcp_config.json`: -OptiLLM supports both **local** and **remote** MCP servers through multiple transport methods: -- **stdio**: Local servers (traditional) -- **SSE**: Remote servers via Server-Sent Events -- **Streamable HTTP**: Remote servers via the MCP Streamable HTTP transport - -#### What is MCP? - -The [Model Context Protocol](https://modelcontextprotocol.io/) (MCP) is an open protocol standard that allows LLMs to securely access tools and data sources through a standardized interface. MCP servers can provide: - -- **Tools**: Callable functions that perform actions (like writing files, querying databases, etc.) -- **Resources**: Data sources for providing context (like file contents) -- **Prompts**: Reusable prompt templates for specific use cases - -#### Configuration - -##### Setting up MCP Config - -> **Note on Backwards Compatibility**: Existing MCP configurations will continue to work unchanged. The `transport` field defaults to "stdio" when not specified, maintaining full backwards compatibility with existing setups. - -1. Create a configuration file at `~/.optillm/mcp_config.json` with the following structure: - -**Local Server (stdio) - Traditional Method:** ```json { "mcpServers": { "filesystem": { "transport": "stdio", "command": "npx", - "args": [ - "-y", - "@modelcontextprotocol/server-filesystem", - "/path/to/allowed/directory1", - "/path/to/allowed/directory2" - ], - "env": {}, + "args": ["-y", "@modelcontextprotocol/server-filesystem", "/home/user/docs"], "description": "Local filesystem access" - } - }, - "log_level": "INFO" -} -``` - -**Legacy Format (still works):** -```json -{ - "mcpServers": { - "filesystem": { - "command": "npx", - "args": ["-y", "@modelcontextprotocol/server-filesystem", "/path/to/directory"], - "env": {} - } - } -} -``` - -**Remote Server (SSE) - New Feature:** -```json -{ - "mcpServers": { + }, "github": { "transport": "sse", "url": "https://api.githubcopilot.com/mcp", @@ -474,19 +289,8 @@ The [Model Context Protocol](https://modelcontextprotocol.io/) (MCP) is an open "Authorization": "Bearer ${GITHUB_TOKEN}", "Accept": "text/event-stream" }, - "timeout": 30.0, - "sse_read_timeout": 300.0, - "description": "GitHub MCP server for repository access" - } - }, - "log_level": "INFO" -} -``` - -**Remote Server (Streamable HTTP):** -```json -{ - "mcpServers": { + "description": "GitHub MCP server" + }, "remote-http": { "transport": "streamable_http", "url": "https://api.example.com/mcp", @@ -500,225 +304,44 @@ The [Model Context Protocol](https://modelcontextprotocol.io/) (MCP) is an open } ``` -**Mixed Configuration (Local + Remote):** -```json -{ - "mcpServers": { - "filesystem": { - "transport": "stdio", - "command": "npx", - "args": ["-y", "@modelcontextprotocol/server-filesystem", "/home/user/docs"], - "description": "Local filesystem access" - }, - "github": { - "transport": "sse", - "url": "https://api.githubcopilot.com/mcp", - "headers": { - "Authorization": "Bearer ${GITHUB_TOKEN}" - }, - "description": "GitHub MCP server" - }, - "remote-api": { - "transport": "streamable_http", - "url": "https://api.company.com/mcp", - "description": "Company internal MCP server" - } - }, - "log_level": "INFO" -} -``` - -##### Configuration Parameters - -**Common Parameters:** -- **Server name**: A unique identifier for the server (e.g., "filesystem", "github") -- **transport**: Transport method - "stdio" (default), "sse", or "streamable_http" -- **description** (optional): Description of the server's functionality -- **timeout** (optional): Connection timeout in seconds (default: 5.0) - -**stdio Transport (Local Servers):** -- **command**: The executable to run the server -- **args**: Command-line arguments for the server -- **env**: Environment variables for the server process +| Transport | Required | Optional | +|-----------|----------|----------| +| `stdio` (default) | `command` | `args`, `env` | +| `sse` | `url` | `headers`, `sse_read_timeout` (default 300.0) | +| `streamable_http` | `url` | `headers`, `sse_read_timeout` (default 300.0) | -**sse Transport (Server-Sent Events):** -- **url**: The SSE endpoint URL -- **headers** (optional): HTTP headers for authentication -- **sse_read_timeout** (optional): SSE read timeout in seconds (default: 300.0) +All transports accept `description` and `timeout` (connection timeout in seconds, default 5.0). The `websocket` transport was removed in MCP Python SDK 2.x; use `streamable_http` or `sse` instead. String values such as headers support `${VARIABLE_NAME}` expansion, and those variables must be set before starting OptiLLM. -**streamable_http Transport (Streamable HTTP):** -- **url**: The MCP endpoint URL -- **headers** (optional): HTTP headers for authentication -- **sse_read_timeout** (optional): Read timeout in seconds (default: 300.0) +Any of the [official MCP servers](https://modelcontextprotocol.io/examples) can be used. Logs are written to `~/.optillm/logs/mcp_plugin.log`; check there for connection and tool errors. -The `websocket` transport was removed in MCP Python SDK 2.x and is no longer supported; use `streamable_http` or `sse` instead. +## Configuration -**Environment Variable Expansion:** -Headers and other string values support environment variable expansion using `${VARIABLE_NAME}` syntax. This is especially useful for API keys: -```json -{ - "headers": { - "Authorization": "Bearer ${GITHUB_TOKEN}", - "X-API-Key": "${MY_API_KEY}" - } -} -``` +Parameters can be passed on the command line. Most can also be set as environment variables prefixed with `OPTILLM_`, for example `OPTILLM_APPROACH=mcts`. Run `optillm --help` for the full list. -#### Available MCP Servers - -OptiLLM supports both local and remote MCP servers: - -##### Local MCP Servers (stdio transport) - -You can use any of the [official MCP servers](https://modelcontextprotocol.io/examples) or third-party servers that run as local processes: - -- **Filesystem**: `@modelcontextprotocol/server-filesystem` - File operations -- **Git**: `mcp-server-git` - Git repository operations -- **SQLite**: `@modelcontextprotocol/server-sqlite` - SQLite database access -- **Brave Search**: `@modelcontextprotocol/server-brave-search` - Web search capabilities - -##### Remote MCP Servers (SSE/Streamable HTTP transport) - -Remote servers provide centralized access without requiring local installation: - -- **GitHub MCP Server**: `https://api.githubcopilot.com/mcp` - Repository management, issue tracking, and code analysis -- **Third-party servers**: Any MCP server that supports the SSE or Streamable HTTP transports - -##### Example: Comprehensive Configuration - -```json -{ - "mcpServers": { - "filesystem": { - "transport": "stdio", - "command": "npx", - "args": ["-y", "@modelcontextprotocol/server-filesystem", "/home/user/documents"], - "description": "Local file system access" - }, - "search": { - "transport": "stdio", - "command": "npx", - "args": ["-y", "@modelcontextprotocol/server-brave-search"], - "env": { - "BRAVE_API_KEY": "your-api-key-here" - }, - "description": "Web search capabilities" - }, - "github": { - "transport": "sse", - "url": "https://api.githubcopilot.com/mcp", - "headers": { - "Authorization": "Bearer ${GITHUB_TOKEN}", - "Accept": "text/event-stream" - }, - "description": "GitHub repository and issue management" - } - }, - "log_level": "INFO" -} -``` - -#### Using the MCP Plugin - -Once configured, the MCP plugin will automatically: - -1. Connect to all configured MCP servers -2. Discover available tools, resources, and prompts -3. Make these capabilities available to the language model -4. Handle tool calls and resource requests - -The plugin enhances the system prompt with MCP capabilities so the model knows which tools are available. When the model decides to use a tool, the plugin: - -1. Executes the tool with the provided arguments -2. Returns the results to the model -3. Allows the model to incorporate the results into its response - -#### Example Queries - -Here are some examples of queries that will engage MCP tools: - -**Local Server Examples:** -- "List all the Python files in my documents directory" (Filesystem) -- "What are the recent commits in my Git repository?" (Git) -- "Search for the latest information about renewable energy" (Search) -- "Query my database for all users who registered this month" (Database) - -**Remote Server Examples:** -- "Show me the open issues in my GitHub repository" (GitHub MCP) -- "Create a new branch for the feature I'm working on" (GitHub MCP) -- "What are the most recent pull requests that need review?" (GitHub MCP) -- "Get the file contents from my remote repository" (GitHub MCP) - -#### Troubleshooting - -##### Logs - -The MCP plugin logs detailed information to: -``` -~/.optillm/logs/mcp_plugin.log -``` - -Check this log file for connection issues, tool execution errors, and other diagnostic information. - -##### Common Issues - -**Local Server Issues (stdio transport):** - -1. **Command not found**: Make sure the server executable is available in your PATH, or use an absolute path in the configuration. - -2. **Access denied**: For filesystem operations, ensure the paths specified in the configuration are accessible to the process. - -**Remote Server Issues (SSE/Streamable HTTP transport):** - -3. **Connection timeout**: Remote servers may take longer to connect. Increase the `timeout` value in your configuration. - -4. **Authentication failed**: Verify your API keys and tokens are correct. For GitHub MCP server, ensure your `GITHUB_TOKEN` environment variable is set with appropriate permissions. - -5. **Network errors**: Check your internet connection and verify the server URL is accessible. - -6. **Environment variable not found**: If using `${VARIABLE_NAME}` syntax, ensure the environment variables are set before starting OptILLM. - -**General Issues:** - -7. **Method not found**: Some servers don't implement all MCP capabilities (tools, resources, prompts). Verify which capabilities the server supports. - -8. **Transport not supported**: Ensure you're using a supported transport: "stdio", "sse", or "streamable_http". - -**Example: Testing GitHub MCP Connection** - -To test if your GitHub MCP server configuration is working: - -1. Set your GitHub token: `export GITHUB_TOKEN="your-github-token"` -2. Start OptILLM and check the logs at `~/.optillm/logs/mcp_plugin.log` -3. Look for connection success messages and discovered capabilities - -## Available parameters - -optillm supports various command-line arguments for configuration. When using Docker, these can also be set as environment variables prefixed with `OPTILLM_`. - -| Parameter | Description | Default Value | -|--------------------------|-----------------------------------------------------------------|-----------------| -| `--approach` | Inference approach to use | `"auto"` | -| `--simulations` | Number of MCTS simulations | 2 | -| `--exploration` | Exploration weight for MCTS | 0.2 | -| `--depth` | Simulation depth for MCTS | 1 | -| `--best-of-n` | Number of samples for best_of_n approach | 3 | -| `--model` | OpenAI model to use | `"gpt-4o-mini"` | -| `--base-url` | Base URL for OpenAI compatible endpoint | `""` | -| `--rstar-max-depth` | Maximum depth for rStar algorithm | 3 | -| `--rstar-num-rollouts` | Number of rollouts for rStar algorithm | 5 | -| `--rstar-c` | Exploration constant for rStar algorithm | 1.4 | -| `--n` | Number of final responses to be returned | 1 | -| `--return-full-response` | Return the full response including the CoT with tags | `False` | -| `--port` | Specify the port to run the proxy | 8000 | -| `--optillm-api-key` | Optional API key for client authentication to optillm | `""` | -| `--cepo_*` | See CePO Parameters section below for detailed config options | Various | +| Parameter | Description | Default | +|-----------|-------------|---------| +| `--approach` | Inference approach to use | `auto` | +| `--model` | OpenAI model to use | `gpt-4o-mini` | +| `--base-url` | Base URL for an OpenAI-compatible endpoint | `""` | +| `--host` | Host address to bind the server | `127.0.0.1` | +| `--port` | Port to run the proxy on | 8000 | +| `--optillm-api-key` | Optional API key for client authentication to optillm | `""` | +| `--n` | Number of final responses to return | 1 | +| `--return-full-response` | Return the full response including the CoT with `` tags | `False` | +| `--mcts-simulations` | Number of MCTS simulations | 2 | +| `--mcts-exploration` | Exploration weight for MCTS | 0.2 | +| `--mcts-depth` | Simulation depth for MCTS | 1 | +| `--best-of-n` | Number of samples for the best_of_n approach | 3 | +| `--rstar-max-depth` | Maximum depth for the R* algorithm | 3 | +| `--rstar-num-rollouts` | Number of rollouts for the R* algorithm | 5 | +| `--rstar-c` | Exploration constant for the R* algorithm | 1.4 | +| `--cepo_*` | See CePO parameters below | Various |
-CePO Parameters +CePO parameters -| Parameter | Description | Default Value | -|-----------|-------------|---------------| +| Parameter | Description | Default | +|-----------|-------------|---------| | `--cepo_bestofn_n` | Number of responses to be generated in best of n stage | 3 | | `--cepo_bestofn_temperature` | Temperature for verifier in best of n stage | 0.1 | | `--cepo_bestofn_max_tokens` | Maximum number of tokens for verifier in best of n stage | 4096 | @@ -744,50 +367,34 @@ optillm supports various command-line arguments for configuration. When using Do
-## Running with Docker - -optillm can optionally be built and run using Docker and the provided [Dockerfile](https://github.com/algorithmicsuperintelligence/optillm/blob/main/Dockerfile). - -### Using Docker Compose - -1. Make sure you have Docker and Docker Compose installed on your system. +## Security -2. Either update the environment variables in the docker-compose.yaml file or create a `.env` file in the project root directory and add any environment variables you want to set. For example, to set the OpenAI API key, add the following line to the `.env` file: +**Network binding and authentication.** The server binds to `127.0.0.1` by default. Use `--host 0.0.0.0` to allow external connections, and only on trusted networks or with authentication enabled. Set `--optillm-api-key` (or `OPTILLM_API_KEY`) to require clients to send `Authorization: Bearer `. - ```bash - OPENAI_API_KEY=your_openai_api_key_here - ``` - -3. Run the following command to start optillm: - - ```bash - docker compose up -d - ``` - - This will build the Docker image if it doesn't exist and start the optillm service. - -4. optillm will be available at `http://localhost:8000`. - -When using Docker, you can set these parameters as environment variables. For example, to set the approach and model, you would use: +**Code execution.** The `executecode` plugin runs Python from requests and LLM responses in a live Jupyter kernel with no sandboxing, so it is disabled unless you opt in. Only enable it in single-user, air-gapped environments where whoever submits requests is trusted with full access to the host. ```bash -OPTILLM_APPROACH=mcts -OPTILLM_MODEL=gpt-4 +optillm --enable-unsafe-plugins=executecode +# or: export OPTILLM_ENABLE_UNSAFE_PLUGINS=executecode ``` -To secure the optillm proxy with an API key, set the `OPTILLM_API_KEY` environment variable: +**URL fetching.** The `readurls` plugin blocks requests to non-public destinations (private, loopback, link-local, and reserved addresses such as the cloud metadata IP `169.254.169.254`) to prevent SSRF. Every redirect hop is re-validated and unresolvable hosts are rejected. Only opt out on trusted networks that need internal URL fetching: ```bash -OPTILLM_API_KEY=your_secret_api_key +optillm --readurls-allow-internal +# or: export OPTILLM_READURLS_ALLOW_INTERNAL=true ``` -When the API key is set, clients must include it in their requests using the `Authorization` header: +The guard validates DNS resolution before each request, but the HTTP client resolves the host again when it connects, so a domain under attacker control could rebind to an internal address between the two lookups. For high-security deployments, also restrict the network egress of the optillm process with a firewall or service-mesh policy. + +**SSL verification.** Use a custom CA bundle for self-signed certificates or corporate proxies. Disabling verification is insecure and intended for development only. See [SSL_CONFIGURATION.md](SSL_CONFIGURATION.md) for details. -```plain -Authorization: Bearer your_secret_api_key +```bash +optillm --ssl-cert-path /path/to/ca-bundle.crt # or: export OPTILLM_SSL_CERT_PATH=... +optillm --no-ssl-verify # or: export OPTILLM_SSL_VERIFY=false ``` -## SOTA results on benchmarks with optillm +## Benchmark results ### MARS on AIME 2025, IMO 2025, and LiveCodeBench (Oct 2025) @@ -800,10 +407,22 @@ Authorization: Bearer your_secret_api_key | **LiveCodeBench v5/v6** | Baseline | 105 | 41 | 39.05% | - | | **LiveCodeBench v5/v6** | **MARS** | 105 | **53** | **50.48%** | **+11.43pp (+29.3%)** | -Model: google/gemini-2.5-flash-lite-preview-09-2025 via OpenRouter -Configuration: 3 agents, 2-pass verification, thinking tags disabled for proofs +Model: google/gemini-2.5-flash-lite-preview-09-2025 via OpenRouter. Configuration: 3 agents, 2-pass verification, thinking tags disabled for proofs. + +### CePO on math and code benchmarks (Sep 2025) + +| Method | AIME 2024 | AIME 2025 | GPQA | LiveCodeBench | +| ----------------------: | :-------: | :-------: | :----: | :-----------: | +| Qwen3 8B | 74.0 | 68.3 | 59.3 | 55.7 | +| CePO (using Qwen3 8B) | 86.7 | 80.0 | 62.5 | 60.5 | +| Qwen3 32B | 81.4 | 72.9 | 66.8 | 65.7 | +| CePO (using Qwen3 32B) | **90.7** | **83.3** | 70.0 | **71.9** | +| Qwen3 235B | 85.7 | 81.5 | 71.1 | 70.7 | +| DeepSeek R1 | 79.8 | 70.0 | 71.5 | 64.3 | +| OpenAI o3-mini | 79.6 | 74.8 | 76.8 | 66.3 | +| Grok3 Think | 83.9 | 77.3 |**80.2**| 70.6 | -### AutoThink on GPQA-Diamond & MMLU-Pro (May 2025) +### AutoThink on GPQA-Diamond and MMLU-Pro (May 2025) | **Model** | **GPQA-Diamond** | | **MMLU-Pro** | | |----------------|-----------------------------|--------------------------|----------------------------|--------------------------| @@ -812,6 +431,8 @@ Configuration: 3 agents, 2-pass verification, thinking tags disabled for proofs | with Fixed Budget | 28.47 | 3570.00 | 26.18 | 1815.67 | | **with AutoThink** | **31.06** | **3520.52** | **26.38** | **1792.50** | +
+Earlier results (2024 to April 2025) ### LongCePO on LongBench v2 (Apr 2025) @@ -824,9 +445,9 @@ Configuration: 3 agents, 2-pass verification, thinking tags disabled for proofs | Claude-3.5-Sonnet-20241022 | 200K | 46.1 (53.9) | 38.6 (41.9) | | Llama-4-Maverick-17B-128E-Instruct | 524K | 32.22 (50.56) | 28.84 (41.86) | - ยน Performance numbers reported by LongBench v2 authors, except for LongCePO and Llama-4-Maverick results. +ยน Performance numbers reported by LongBench v2 authors, except for LongCePO and Llama-4-Maverick results. - ยฒ Numbers in parentheses for LongCePO indicate accuracy of majority voting from 5 runs. +ยฒ Numbers in parentheses for LongCePO indicate accuracy of majority voting from 5 runs. ### LongCePO on HELMET - InfiniteBench En.MC, 128K length (Apr 2025) @@ -837,20 +458,7 @@ Configuration: 3 agents, 2-pass verification, thinking tags disabled for proofs | o1-mini-2024-09-12 (full context) | 58.0 | | gpt-4o-2024-08-06 (full context) | 74.0 | - ยน Numbers in parentheses for LongCePO indicate accuracy of majority voting from 5 runs. - -### CePO on math and code benchmarks (Sep 2025) - -| Method | AIME 2024 | AIME 2025 | GPQA | LiveCodeBench | -| ----------------------: | :-------: | :-------: | :----: | :-----------: | -| Qwen3 8B | 74.0 | 68.3 | 59.3 | 55.7 | -| CePO (using Qwen3 8B) | 86.7 | 80.0 | 62.5 | 60.5 | -| Qwen3 32B | 81.4 | 72.9 | 66.8 | 65.7 | -| CePO (using Qwen3 32B) | **90.7** | **83.3** | 70.0 | **71.9** | -| Qwen3 235B | 85.7 | 81.5 | 71.1 | 70.7 | -| DeepSeek R1 | 79.8 | 70.0 | 71.5 | 64.3 | -| OpenAI o3-mini | 79.6 | 74.8 | 76.8 | 66.3 | -| Grok3 Think | 83.9 | 77.3 |**80.2**| 70.6 | +ยน Numbers in parentheses for LongCePO indicate accuracy of majority voting from 5 runs. ### CePO on math and code benchmarks (Mar 2025) @@ -878,6 +486,7 @@ Configuration: 3 agents, 2-pass verification, thinking tags disabled for proofs | gemini-1.5-flash-002 | 16.67 | ### readurls&memory-gpt-4o-mini on Google FRAMES Benchmark (Oct 2024) + | Model | Accuracy | | ----- | -------- | | readurls&memory-gpt-4o-mini | 61.29 | @@ -904,78 +513,41 @@ Configuration: 3 agents, 2-pass verification, thinking tags disabled for proofs ### optillm with Patchwork (July 2024) -Since optillm is a drop-in replacement for OpenAI API you can easily integrate it with existing tools and frameworks using the OpenAI client. We used optillm with [patchwork](https://github.com/patched-codes/patchwork) which is an open-source framework that automates development gruntwork like PR reviews, bug fixing, security patching using workflows -called patchflows. We saw huge performance gains across all the supported patchflows as shown below when using the mixture of agents approach (moa). +We used optillm with [patchwork](https://github.com/patched-codes/patchwork), an open-source framework that automates development work such as PR reviews, bug fixing, and security patching through workflows called patchflows. The mixture of agents approach (moa) improved performance across all supported patchflows. ![Results showing optillm mixture of agents approach used with patchflows](https://raw.githubusercontent.com/algorithmicsuperintelligence/optillm/main/moa-patchwork-results.png) -## Testing - -OptiLLM includes a comprehensive test suite to ensure reliability and compatibility. +
-### Running Tests +## Testing -The main test suite can be run from the project root: ```bash -# Test all approaches with default test cases -python tests/test.py +pip install -r tests/requirements.txt -# Test specific approaches +# Approach tests +python tests/test.py python tests/test.py --approaches moa bon mcts - -# Run a single test python tests/test.py --single-test "Simple Math Problem" -``` -### Unit and Integration Tests - -Additional tests are available in the `tests/` directory: -```bash -# Run all tests (requires pytest) +# Unit and integration tests ./tests/run_tests.sh - -# Run specific test modules pytest tests/test_plugins.py -v -pytest tests/test_api_compatibility.py -v ``` -### CI/CD - -All tests are automatically run on pull requests via GitHub Actions. The workflow tests: -- Multiple Python versions (3.10, 3.11, 3.12) -- Unit tests for plugins and core functionality -- API compatibility tests -- Integration tests with various approaches - -See `tests/README.md` for more details on the test structure and how to write new tests. - -## ๐Ÿค Contributing +Tests run on pull requests through GitHub Actions. See [tests/README.md](tests/README.md) for the test structure and how to write new tests. -We โค๏ธ contributions! OptiLLM is built by the community, for the community. +## Contributing -- ๐Ÿ› **Found a bug?** [Open an issue](https://github.com/algorithmicsuperintelligence/optillm/issues/new) -- ๐Ÿ’ก **Have an idea?** [Start a discussion](https://github.com/algorithmicsuperintelligence/optillm/discussions) -- ๐Ÿ”ง **Want to code?** Check out [good first issues](https://github.com/algorithmicsuperintelligence/optillm/labels/good%20first%20issue) - -### Development Setup -```bash -git clone https://github.com/algorithmicsuperintelligence/optillm.git -cd optillm -python -m venv .venv -source .venv/bin/activate # or `.venv\Scripts\activate` on Windows -pip install -r requirements.txt -pip install -r tests/requirements.txt - -# Run tests -python -m pytest tests/ -``` +- Report bugs by [opening an issue](https://github.com/algorithmicsuperintelligence/optillm/issues/new). +- Share ideas by [starting a discussion](https://github.com/algorithmicsuperintelligence/optillm/discussions). +- Look for a place to start in the [good first issues](https://github.com/algorithmicsuperintelligence/optillm/labels/good%20first%20issue). ## References + - [Eliciting Fine-Tuned Transformer Capabilities via Inference-Time Techniques](https://arxiv.org/abs/2506.08060) - [AutoThink: efficient inference for reasoning LLMs](https://dx.doi.org/10.2139/ssrn.5253327) - [Implementation](optillm/autothink) - [Deep Think with Confidence: Confidence-guided reasoning and inference-time scaling](https://arxiv.org/abs/2508.15260) - [Implementation](optillm/deepconf) -- [Self-Discover: Large Language Models Self-Compose Reasoning Structures -](https://arxiv.org/abs/2402.03620) - [Implementation](optillm/plugings/deepthink) +- [Self-Discover: Large Language Models Self-Compose Reasoning Structures](https://arxiv.org/abs/2402.03620) - [Implementation](optillm/plugins/deepthink) - [CePO: Empowering Llama with Reasoning using Test-Time Compute](https://cerebras.ai/blog/cepo) - [Implementation](optillm/cepo) - [LongCePO: Empowering LLMs to efficiently leverage infinite context](https://cerebras.ai/blog/longcepo) - [Implementation](optillm/plugins/longcepo) - [Chain of Code: Reasoning with a Language Model-Augmented Code Emulator](https://arxiv.org/abs/2312.04474) - [Inspired the implementation of coc plugin](optillm/plugins/coc_plugin.py) @@ -1010,13 +582,3 @@ If you use this library in your research, please cite: url = {https://github.com/algorithmicsuperintelligence/optillm} } ``` - ---- - -

- Ready to optimize your LLMs? Install OptiLLM and see the difference! ๐Ÿš€ -

- -

- โญ Star us on GitHub if you find OptiLLM useful! -

diff --git a/docs/llm-gateway.md b/docs/llm-gateway.md new file mode 100644 index 00000000..cb864f2e --- /dev/null +++ b/docs/llm-gateway.md @@ -0,0 +1,74 @@ +# Using OptiLLM as an LLM gateway + +OptiLLM's proxy mode gives you one self-hosted, OpenAI-compatible endpoint in front of any number of providers, with load balancing, failover, and health checks. It is free and Apache-2.0 licensed, and you pay your providers directly. + +OptiLLM combines routing and inference optimization in the same service. The same request that is routed across providers can also run through [20+ state-of-the-art reasoning techniques](../README.md#techniques), so your gateway can improve answers as well as deliver them. + +## Quick start + +Create `~/.optillm/proxy_config.yaml`: + +```yaml +providers: + - name: openai + base_url: https://api.openai.com/v1 + api_key: ${OPENAI_API_KEY} + weight: 3 + + - name: together + base_url: https://api.together.xyz/v1 + api_key: ${TOGETHER_API_KEY} + weight: 1 + model_map: + gpt-4o-mini: meta-llama/Llama-3.3-70B-Instruct-Turbo + + - name: local + base_url: http://localhost:8080/v1 + api_key: none + max_concurrent: 2 + +routing: + strategy: weighted # or round_robin, failover + health_check: + enabled: true + interval: 30 +``` + +Start the server with the proxy as the default approach and protect it with an API key: + +```bash +optillm --approach proxy --optillm-api-key "$GATEWAY_KEY" +``` + +Point any OpenAI client at it: + +```python +from openai import OpenAI + +client = OpenAI(api_key=GATEWAY_KEY, base_url="http://localhost:8000/v1") + +response = client.chat.completions.create( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "Hello"}], +) +``` + +If a provider fails, the request moves to the next one, and then to the server's default client (`--base-url`) if one is available. + +## Adding inference optimization + +Start the server without `--approach proxy` and select the proxy per request, on its own or wrapped around a technique: + +```python +# Route only +client.chat.completions.create(model="proxy-gpt-4o-mini", messages=messages) + +# Run Mixture of Agents, sending its calls through the proxy +client.chat.completions.create( + model="gpt-4o-mini", + messages=messages, + extra_body={"optillm_approach": "proxy", "proxy_wrap": "moa"}, +) +``` + +See the [proxy plugin documentation](../optillm/plugins/proxy/README.md) for the full configuration reference. diff --git a/optillm/__init__.py b/optillm/__init__.py index bdb79053..6c0d6714 100644 --- a/optillm/__init__.py +++ b/optillm/__init__.py @@ -1,5 +1,5 @@ # Version information -__version__ = "0.4.0" +__version__ = "0.4.1" import os as _os diff --git a/optillm/inference.py b/optillm/inference.py index de4ebbb3..49d95d27 100644 --- a/optillm/inference.py +++ b/optillm/inference.py @@ -1872,6 +1872,10 @@ def create( ) -> ChatCompletion: """Create a chat completion with OpenAI-compatible parameters""" logger.info("Starting chat completion creation") + # max_completion_tokens is the preferred name for the same limit + max_completion_tokens = kwargs.pop("max_completion_tokens", None) + if max_completion_tokens is not None: + max_tokens = max_completion_tokens if stream: raise NotImplementedError("Streaming is not yet supported") diff --git a/optillm/plugins/mcp_plugin.py b/optillm/plugins/mcp_plugin.py index a3638d1c..a5414d86 100644 --- a/optillm/plugins/mcp_plugin.py +++ b/optillm/plugins/mcp_plugin.py @@ -192,7 +192,6 @@ def load_config(self) -> bool: with open(self.config_path, 'r') as f: config_data = f.read() - logger.debug(f"Raw config data: {config_data}") config = json.loads(config_data) # Set log level @@ -204,7 +203,7 @@ def load_config(self) -> bool: servers_config = config.get("mcpServers", {}) for server_name, server_config in servers_config.items(): self.servers[server_name] = ServerConfig.from_dict(server_config) - logger.debug(f"Loaded server config for {server_name}: {server_config}") + logger.debug(f"Loaded server config for {server_name} (transport: {self.servers[server_name].transport})") logger.info(f"Loaded configuration with {len(self.servers)} servers") return True @@ -398,7 +397,7 @@ async def connect_sse(self) -> bool: """Connect to server using SSE transport and discover capabilities""" logger.info(f"Connecting to SSE server: {self.server_name}") logger.debug(f"SSE URL: {self.config.url}") - logger.debug(f"Headers: {self.config.headers}") + logger.debug(f"Headers: {list(self.config.headers or {})}") if not self.config.url: logger.error(f"SSE transport requires URL for server {self.server_name}") @@ -435,7 +434,6 @@ async def connect_websocket(self) -> bool: async def connect_stdio_native(self) -> bool: """Connect using stdio transport with local executable""" - logger.debug(f"Server configuration: {vars(self.config)}") # Validate stdio configuration if not self.config.command: @@ -455,7 +453,7 @@ async def connect_stdio_native(self) -> bool: logger.debug(f"Using command: {full_command}") logger.debug(f"Arguments: {self.config.args}") - logger.debug(f"Environment: {self.config.env}") + logger.debug(f"Environment variables: {list(self.config.env or {})}") # Create server parameters server_params = StdioServerParameters( diff --git a/optillm/plugins/proxy_plugin.py b/optillm/plugins/proxy_plugin.py index 151b0421..10d0361c 100644 --- a/optillm/plugins/proxy_plugin.py +++ b/optillm/plugins/proxy_plugin.py @@ -37,6 +37,18 @@ def _request_messages(system_prompt: str, initial_query: str, messages: Optional {"role": "user", "content": initial_query} ] +def _passthrough_config(request_config: Optional[dict]) -> dict: + """ + Build the kwargs forwarded to a provider. Streaming is stripped to force a + complete response (server.py converts it to SSE if needed), and the max_tokens + copy of max_completion_tokens is dropped because o-series and gpt-5 models reject it. + """ + api_config = dict(request_config or {}) + api_config.pop('stream', None) + if api_config.get('max_completion_tokens') is not None: + api_config.pop('max_tokens', None) + return api_config + def run(system_prompt: str, initial_query: str, client, model: str, request_config: dict = None, messages: list = None) -> Tuple[str, int]: """ @@ -65,9 +77,7 @@ def run(system_prompt: str, initial_query: str, client, model: str, if not config.get('providers'): logger.warning(f"No providers configured in {ProxyConfig._config_path}, falling back to original client") - # Strip stream parameter to force complete response - api_config = dict(request_config or {}) - api_config.pop('stream', None) + api_config = _passthrough_config(request_config) response = client.chat.completions.create( model=model, @@ -143,10 +153,7 @@ def run(system_prompt: str, initial_query: str, client, model: str, # message support and merges the system prompt when needed. logger.info(f"Direct proxy routing for model: {model}") - # Strip stream parameter to force complete response - # server.py will handle converting to SSE streaming format if needed - api_config = dict(request_config or {}) - api_config.pop('stream', None) + api_config = _passthrough_config(request_config) response = proxy_client.chat.completions.create( model=model, @@ -162,9 +169,7 @@ def run(system_prompt: str, initial_query: str, client, model: str, logger.error(f"Proxy plugin error: {e}", exc_info=True) # Fallback to original client logger.info("Falling back to original client") - # Strip stream parameter to force complete response - api_config = dict(request_config or {}) - api_config.pop('stream', None) + api_config = _passthrough_config(request_config) response = client.chat.completions.create( model=model, diff --git a/optillm/plugins/streaming_hints/README.md b/optillm/plugins/streaming_hints/README.md new file mode 100644 index 00000000..d5438a2a --- /dev/null +++ b/optillm/plugins/streaming_hints/README.md @@ -0,0 +1,73 @@ +# Streaming Hints + +Injects hints from a memory bank into the model's reasoning while it is being generated, inspired by [this idea from Will Brown](https://x.com/willccbb/status/1940557166248972387). Instead of putting every piece of guidance in the prompt up front, the plugin watches the reasoning and adds one targeted hint at the moment the model hesitates or moves to a new step. + +This is the inference half of the approach. The retriever is rule-based for now; training the injection policy with RL is future work (see below). + +## How it works + +1. The model answers the query in segments of at most `segment_tokens` tokens. +2. Each segment is fed sentence by sentence to the `StreamingHintRetriever`. +3. A sentence is an **injection point** when at least `min_gap_chars` characters have passed since the last hint and the sentence contains an uncertainty signal ("wait", "hmm", "not sure", "let me check", ...) or ends a paragraph. +4. At an injection point the retriever queries the `HintMemoryBank` with the problem plus the most recent reasoning. If the best unused hint scores at least `hint_threshold`, the segment is cut after that sentence and the hint is appended as `...`. +5. Generation continues from the cut, with the reasoning so far (hints included) in the prompt. This repeats until the model finishes, `max_segments` is reached, or `max_tokens` is spent. + +Segments are used instead of a token stream so the plugin works with every client optillm supports, including the local inference server, which does not stream. The cost is that up to one segment of text after an injection point is generated and then discarded. + +The memory bank ranks hints by TF-IDF cosine similarity over the hint text and its keywords. It is pure Python, so the plugin has no extra dependencies. + +## Usage + +```python +from openai import OpenAI + +client = OpenAI(api_key="sk-...", base_url="http://localhost:8000/v1") + +response = client.chat.completions.create( + model="streaming_hints-gpt-4o-mini", + messages=[{"role": "user", "content": "A train travels 120 km in 1.5 hours. What is its speed in m/s?"}], + extra_body={"hint_threshold": 0.25, "show_hints": True}, +) +``` + +It combines with other approaches like any plugin, for example `streaming_hints&cot_reflection-gpt-4o-mini`. + +## Configuration + +All options are read from the request body (`extra_body`). + +| Option | Default | Description | +|---|---|---| +| `hint_threshold` | `0.2` | Minimum similarity score for a hint to be injected | +| `max_hints` | `3` | Maximum number of hints injected per request | +| `min_gap_chars` | `200` | Minimum characters of reasoning between two injections | +| `segment_tokens` | `512` | Maximum tokens generated per segment | +| `max_segments` | `8` | Maximum number of segments per request | +| `hints_path` | bundled `hints.json` | Path to a custom hint bank | +| `show_hints` | `false` | Return the reasoning with the `` tags instead of the model text only | + +`max_tokens`, if set, caps the completion tokens across all segments. + +## Custom hints + +A hint bank is a JSON list. `keywords` count twice as much as words in `text` when matching. + +```json +[ + { + "id": "math-units", + "domain": "math", + "text": "Write down the units of every quantity and check both sides match.", + "keywords": ["units", "speed", "rate", "convert"] + } +] +``` + +The bundled bank has 20 general hints for math, coding and logical reasoning. Domain-specific banks will usually do better. + +## Future work + +- **RL-trained retriever.** Learn when to inject and which hint to pick from downstream task accuracy, with and without hints as the reward signal. `Hint.uses` is tracked as a starting point for usage statistics. +- **Embedding retrieval.** A bank backed by sentence embeddings can replace TF-IDF behind the same `search` interface. +- **Token streaming.** For providers that stream, the retriever can be fed token deltas directly and the stream cancelled at the injection point, which avoids the discarded tail of a segment. +- **Evaluation.** Benchmark results (GSM8K, MATH, HumanEval, MMLU) against the plain model and against putting all hints in the prompt. diff --git a/optillm/plugins/streaming_hints/__init__.py b/optillm/plugins/streaming_hints/__init__.py new file mode 100644 index 00000000..208ec2ea --- /dev/null +++ b/optillm/plugins/streaming_hints/__init__.py @@ -0,0 +1,12 @@ +""" +Streaming Hints Plugin for OptiLLM + +Monitors the model's reasoning as it is generated and injects relevant hints +from a memory bank at moments of uncertainty or between reasoning steps. +""" + +from .memory_bank import Hint, HintMemoryBank +from .retriever import Injection, StreamingHintRetriever +from .streaming_hints import run_streaming_hints + +__all__ = ['Hint', 'HintMemoryBank', 'Injection', 'StreamingHintRetriever', 'run_streaming_hints'] diff --git a/optillm/plugins/streaming_hints/hints.json b/optillm/plugins/streaming_hints/hints.json new file mode 100644 index 00000000..4e805cba --- /dev/null +++ b/optillm/plugins/streaming_hints/hints.json @@ -0,0 +1,122 @@ +[ + { + "id": "math-units", + "domain": "math", + "text": "Write down the units of every quantity and check that both sides of each equation have the same units.", + "keywords": ["units", "speed", "rate", "distance", "time", "convert", "per", "hour", "minute", "meters"] + }, + { + "id": "math-substitute-back", + "domain": "math", + "text": "Substitute the candidate answer back into the original conditions to confirm it satisfies all of them.", + "keywords": ["solve", "equation", "solution", "root", "satisfy", "check", "verify", "answer", "x"] + }, + { + "id": "math-small-cases", + "domain": "math", + "text": "Try the smallest cases (n = 1, 2, 3) by hand and look for a pattern before generalising.", + "keywords": ["pattern", "sequence", "integer", "n", "general", "formula", "induction", "sum", "series"] + }, + { + "id": "math-casework", + "domain": "math", + "text": "Split the problem into exhaustive, non-overlapping cases and make sure no case is counted twice.", + "keywords": ["count", "ways", "cases", "arrangements", "combinations", "permutations", "how", "many", "choose"] + }, + { + "id": "math-complement", + "domain": "math", + "text": "When counting or computing a probability is hard directly, compute the complement and subtract from the total.", + "keywords": ["probability", "at", "least", "one", "complement", "chance", "random", "dice", "coin"] + }, + { + "id": "math-percent-base", + "domain": "math", + "text": "For percentage changes, be explicit about the base each percentage refers to; successive changes multiply rather than add.", + "keywords": ["percent", "percentage", "increase", "decrease", "discount", "interest", "price", "growth"] + }, + { + "id": "math-geometry-diagram", + "domain": "math", + "text": "Label every known length and angle on a diagram, then look for similar triangles, right angles or symmetry.", + "keywords": ["triangle", "angle", "circle", "area", "perimeter", "radius", "geometry", "length", "square"] + }, + { + "id": "math-modular", + "domain": "math", + "text": "For remainders and divisibility, work modulo the divisor throughout and reduce intermediate results early.", + "keywords": ["remainder", "divisible", "modulo", "mod", "divides", "digit", "last", "prime"] + }, + { + "id": "code-edge-cases", + "domain": "coding", + "text": "List the edge cases first: empty input, a single element, duplicates, negative numbers and the maximum size.", + "keywords": ["function", "array", "list", "input", "empty", "edge", "case", "string", "return"] + }, + { + "id": "code-off-by-one", + "domain": "coding", + "text": "Check loop bounds and slice indices for off-by-one errors by tracing the first and last iteration.", + "keywords": ["loop", "index", "range", "bounds", "slice", "iterate", "for", "while", "off", "by"] + }, + { + "id": "code-complexity", + "domain": "coding", + "text": "Estimate the time complexity against the input limits before writing code; nested loops over large inputs usually need a hash map, sorting or two pointers.", + "keywords": ["complexity", "efficient", "performance", "large", "constraints", "time", "optimize", "fast", "big", "o"] + }, + { + "id": "code-trace-example", + "domain": "coding", + "text": "Trace the code by hand on the given example and compare each intermediate value with what you expect.", + "keywords": ["example", "output", "expected", "trace", "debug", "test", "bug", "wrong", "result"] + }, + { + "id": "code-recursion-base", + "domain": "coding", + "text": "For recursion, state the base case and confirm every recursive call moves strictly towards it; consider memoisation for overlapping subproblems.", + "keywords": ["recursion", "recursive", "base", "case", "dynamic", "programming", "memoization", "fibonacci", "subproblem"] + }, + { + "id": "logic-restate", + "domain": "reasoning", + "text": "Restate exactly what the question asks for, including the form of the answer, before continuing.", + "keywords": ["question", "asks", "find", "determine", "which", "answer", "what"] + }, + { + "id": "logic-assumptions", + "domain": "reasoning", + "text": "Separate what the problem states from what you have assumed, and check whether any assumption is unjustified.", + "keywords": ["assume", "assumption", "given", "statement", "premise", "conclusion", "infer", "implies"] + }, + { + "id": "logic-contrapositive", + "domain": "reasoning", + "text": "An implication 'if A then B' is equivalent to 'if not B then not A', but not to 'if B then A'.", + "keywords": ["if", "then", "implies", "converse", "contrapositive", "necessary", "sufficient", "logic", "only"] + }, + { + "id": "logic-eliminate", + "domain": "reasoning", + "text": "Eliminate options that contradict any stated fact, then compare the remaining options directly.", + "keywords": ["option", "choice", "options", "eliminate", "multiple", "correct", "statement", "true", "false"] + }, + { + "id": "logic-table", + "domain": "reasoning", + "text": "For puzzles that match people to attributes, build a grid of possibilities and cross out cells as each clue is applied.", + "keywords": ["puzzle", "each", "person", "clue", "who", "house", "owns", "lives", "match"] + }, + { + "id": "general-estimate", + "domain": "general", + "text": "Make a rough estimate of the answer's magnitude first so an implausible final result stands out.", + "keywords": ["estimate", "approximately", "roughly", "magnitude", "total", "large", "number"] + }, + { + "id": "general-recheck-arithmetic", + "domain": "general", + "text": "Recompute the last arithmetic step another way, for example by reversing the operation.", + "keywords": ["calculate", "compute", "multiply", "divide", "add", "subtract", "arithmetic", "total", "sum"] + } +] diff --git a/optillm/plugins/streaming_hints/memory_bank.py b/optillm/plugins/streaming_hints/memory_bank.py new file mode 100644 index 00000000..f00fbed8 --- /dev/null +++ b/optillm/plugins/streaming_hints/memory_bank.py @@ -0,0 +1,97 @@ +""" +Hint memory bank for the streaming hints plugin. + +Stores curated hints and retrieves them by TF-IDF cosine similarity. This is +pure Python on purpose, so the plugin adds no dependencies and the bank can +later be swapped for an embedding-backed one with the same `search` interface. +""" + +import json +import math +import os +import re +from collections import Counter +from dataclasses import dataclass, field +from typing import Dict, Iterable, List, Optional, Tuple + +DEFAULT_HINTS_PATH = os.path.join(os.path.dirname(__file__), "hints.json") + +_STOPWORDS = frozenset( + "a an and are as at be by for from has have i if in is it its of on or so " + "that the then there this to was we what when which will with you your".split() +) + + +def tokenize(text: str) -> List[str]: + return [t for t in re.findall(r"[a-z0-9]+", text.lower()) if t not in _STOPWORDS] + + +@dataclass +class Hint: + id: str + text: str + keywords: List[str] = field(default_factory=list) + domain: str = "general" + uses: int = 0 + + def index_terms(self) -> List[str]: + # Keywords are counted twice so a curated keyword outweighs an + # incidental word in the hint text. + keyword_terms = tokenize(" ".join(self.keywords)) + return tokenize(self.text) + keyword_terms + keyword_terms + + +class HintMemoryBank: + def __init__(self, hints: Optional[Iterable[Hint]] = None): + self.hints: List[Hint] = [] + self._idf: Dict[str, float] = {} + self._vectors: List[Dict[str, float]] = [] + for hint in hints or []: + self.hints.append(hint) + self._reindex() + + @classmethod + def from_file(cls, path: Optional[str] = None) -> "HintMemoryBank": + with open(path or DEFAULT_HINTS_PATH, "r", encoding="utf-8") as f: + data = json.load(f) + return cls( + Hint( + id=str(item.get("id", i)), + text=item["text"], + keywords=list(item.get("keywords", [])), + domain=item.get("domain", "general"), + ) + for i, item in enumerate(data) + ) + + def add(self, hint: Hint) -> None: + self.hints.append(hint) + self._reindex() + + def _reindex(self) -> None: + docs = [Counter(h.index_terms()) for h in self.hints] + n = len(docs) + df = Counter(term for doc in docs for term in doc) + self._idf = {term: math.log((1 + n) / (1 + count)) + 1 for term, count in df.items()} + self._vectors = [self._weigh(doc) for doc in docs] + + def _weigh(self, counts: Counter) -> Dict[str, float]: + vector = {t: c * self._idf[t] for t, c in counts.items() if t in self._idf} + norm = math.sqrt(sum(w * w for w in vector.values())) + return {t: w / norm for t, w in vector.items()} if norm else {} + + def search(self, query: str, k: int = 1, exclude: Iterable[str] = ()) -> List[Tuple[Hint, float]]: + """Return up to k (hint, score) pairs, best first, skipping ids in exclude.""" + query_vector = self._weigh(Counter(tokenize(query))) + if not query_vector: + return [] + excluded = set(exclude) + scored = [] + for hint, vector in zip(self.hints, self._vectors): + if hint.id in excluded: + continue + score = sum(w * vector.get(t, 0.0) for t, w in query_vector.items()) + if score > 0: + scored.append((hint, score)) + scored.sort(key=lambda pair: pair[1], reverse=True) + return scored[:k] diff --git a/optillm/plugins/streaming_hints/retriever.py b/optillm/plugins/streaming_hints/retriever.py new file mode 100644 index 00000000..f62634d7 --- /dev/null +++ b/optillm/plugins/streaming_hints/retriever.py @@ -0,0 +1,103 @@ +""" +Streaming hint retriever. + +Consumes the model's reasoning incrementally, decides where a hint may be +injected (an uncertainty signal or a paragraph break), and retrieves the best +matching unused hint for the problem plus the most recent reasoning. +""" + +import re +from dataclasses import dataclass +from typing import List, Optional + +from .memory_bank import Hint, HintMemoryBank + +UNCERTAINTY_PATTERN = re.compile( + r"\b(wait|hmm+|hold on|not sure|unsure|maybe|perhaps|actually|alternatively|" + r"on second thought|let me (?:re)?(?:check|think|consider|verify|try)|" + r"(?:doesn'?t|does not) (?:seem|look) right|i (?:might|may) have made)\b", + re.IGNORECASE, +) + +# A sentence ends at ., ! or ? followed by whitespace, or at a newline. +_BOUNDARY = re.compile(r"[.!?](?=\s)|\n") + + +@dataclass +class Injection: + hint: Hint + score: float + # Offset into the text passed to the feed() call that produced this + # injection. The caller keeps text[:offset] and discards the rest. + offset: int + + +class StreamingHintRetriever: + def __init__( + self, + memory_bank: HintMemoryBank, + threshold: float = 0.2, + max_hints: int = 3, + min_gap_chars: int = 200, + window_chars: int = 600, + ): + self.memory_bank = memory_bank + self.threshold = threshold + self.max_hints = max_hints + self.min_gap_chars = min_gap_chars + self.window_chars = window_chars + self.reset("") + + def reset(self, problem: str) -> None: + self.problem = problem + self.text = "" + self.pending = "" + self.chars_since_injection = 0 + self.used: List[str] = [] + + def feed(self, delta: str) -> Optional[Injection]: + """ + Consume the next piece of reasoning. Returns the first injection found + in it, or None. On an injection the text after the offset is treated + as discarded and is not kept in the retriever's state. + """ + if len(self.used) >= self.max_hints: + self.text += delta + return None + + consumed = 0 + for match in _BOUNDARY.finditer(delta): + end = match.end() + sentence = self.pending + delta[consumed:end] + consumed = end + self.pending = "" + self.text += sentence + self.chars_since_injection += len(sentence) + + if not self._is_injection_point(sentence): + continue + injection = self._retrieve(end) + if injection: + self.chars_since_injection = 0 + return injection + + self.pending += delta[consumed:] + return None + + def _is_injection_point(self, sentence: str) -> bool: + if self.chars_since_injection < self.min_gap_chars: + return False + if UNCERTAINTY_PATTERN.search(sentence): + return True + # A blank line marks a transition between reasoning steps. + return self.text.endswith("\n\n") + + def _retrieve(self, offset: int) -> Optional[Injection]: + query = f"{self.problem}\n{self.text[-self.window_chars:]}" + results = self.memory_bank.search(query, k=1, exclude=self.used) + if not results or results[0][1] < self.threshold: + return None + hint, score = results[0] + hint.uses += 1 + self.used.append(hint.id) + return Injection(hint=hint, score=score, offset=offset) diff --git a/optillm/plugins/streaming_hints/streaming_hints.py b/optillm/plugins/streaming_hints/streaming_hints.py new file mode 100644 index 00000000..96e18942 --- /dev/null +++ b/optillm/plugins/streaming_hints/streaming_hints.py @@ -0,0 +1,117 @@ +""" +Streaming hint injection. + +Generates the answer in segments of at most `segment_tokens` tokens and feeds +each segment to the retriever as it arrives. When the retriever finds an +injection point with a relevant hint, the segment is cut there, the hint is +appended inside tags, and generation continues from the cut. +Segments are used instead of a token stream so the approach works with every +client optillm supports, including the local inference server, which does not +stream. +""" + +import logging +from typing import Any, Dict, Optional, Tuple + +from .memory_bank import HintMemoryBank +from .retriever import StreamingHintRetriever + +logger = logging.getLogger(__name__) + +DEFAULT_CONFIG = { + "hint_threshold": 0.2, + "max_hints": 3, + "min_gap_chars": 200, + "segment_tokens": 512, + "max_segments": 8, + "hints_path": None, + "show_hints": False, +} + +CONTINUE_TEMPLATE = ( + "{query}\n\n" + "Your reasoning so far is below. Continue it from exactly where it stops, " + "without repeating it. Text inside tags was retrieved to help " + "you; use it if it is relevant.\n\n" + "\n{reasoning}\n" +) + + +def _parse_config(request_config: Optional[Dict[str, Any]]) -> Dict[str, Any]: + request_config = request_config or {} + return {key: request_config.get(key, default) for key, default in DEFAULT_CONFIG.items()} + + +def run_streaming_hints( + system_prompt: str, + initial_query: str, + client, + model: str, + request_config: Optional[Dict[str, Any]] = None, + memory_bank: Optional[HintMemoryBank] = None, +) -> Tuple[str, int]: + config = _parse_config(request_config) + request_config = request_config or {} + max_tokens = request_config.get("max_tokens") + temperature = request_config.get("temperature") + + retriever = StreamingHintRetriever( + memory_bank or HintMemoryBank.from_file(config["hints_path"]), + threshold=config["hint_threshold"], + max_hints=config["max_hints"], + min_gap_chars=config["min_gap_chars"], + ) + retriever.reset(initial_query) + + output = "" # model text only + reasoning = "" # model text with injected hints + total_tokens = 0 + + for segment in range(config["max_segments"]): + segment_tokens = config["segment_tokens"] + if max_tokens: + remaining = max_tokens - total_tokens + if remaining <= 0: + break + segment_tokens = min(segment_tokens, remaining) + + user_content = initial_query if not reasoning else CONTINUE_TEMPLATE.format( + query=initial_query, reasoning=reasoning + ) + params = { + "model": model, + "messages": [ + {"role": "system", "content": system_prompt}, + {"role": "user", "content": user_content}, + ], + "max_tokens": segment_tokens, + } + if temperature is not None: + params["temperature"] = temperature + + response = client.chat.completions.create(**params) + choice = response.choices[0] + text = choice.message.content or "" + used_tokens = response.usage.completion_tokens if response.usage is not None else 0 + total_tokens += used_tokens + # The local inference server always reports "stop", so a segment that + # used its whole budget is treated as cut off too. + truncated = choice.finish_reason == "length" or used_tokens >= segment_tokens + + injection = retriever.feed(text) + if injection: + kept = text[: injection.offset] + output += kept + reasoning += f"{kept}\n{injection.hint.text}\n" + logger.info( + f"Segment {segment + 1}: injected hint {injection.hint.id} (score {injection.score:.3f})" + ) + continue + + output += text + reasoning += text + if not truncated: + break + + logger.info(f"Streaming hints done: {len(retriever.used)} hints injected, {total_tokens} completion tokens") + return (reasoning if config["show_hints"] else output), total_tokens diff --git a/optillm/plugins/streaming_hints_plugin.py b/optillm/plugins/streaming_hints_plugin.py new file mode 100644 index 00000000..73bcaef1 --- /dev/null +++ b/optillm/plugins/streaming_hints_plugin.py @@ -0,0 +1,27 @@ +""" +Streaming Hints Plugin for OptiLLM + +Injects hints from a memory bank into the model's reasoning while it is being +generated. See optillm/plugins/streaming_hints/README.md for details. +""" + +import logging +from typing import Any, Dict, Tuple + +from optillm.plugins.streaming_hints import run_streaming_hints + +# Plugin identifier for optillm +SLUG = "streaming_hints" + +logger = logging.getLogger(__name__) + + +def run( + system_prompt: str, + initial_query: str, + client, + model: str, + request_config: Dict[str, Any] = None +) -> Tuple[str, int]: + logger.info("Starting streaming hints") + return run_streaming_hints(system_prompt, initial_query, client, model, request_config) diff --git a/optillm/server.py b/optillm/server.py index ca8f4b25..6bc7b4e1 100644 --- a/optillm/server.py +++ b/optillm/server.py @@ -197,6 +197,7 @@ def count_reasoning_tokens(text: str, tokenizer=None) -> int: 'ssl_verify': True, 'ssl_cert_path': '', 'readurls_allow_internal': False, + 'enable_unsafe_plugins': '', } # List of known approaches @@ -286,10 +287,20 @@ def none_approach( logger.error(f"Error in none approach: {str(e)}") raise +# Plugins that execute arbitrary code are disabled by default and require an +# explicit opt-in via --enable-unsafe-plugins (env: OPTILLM_ENABLE_UNSAFE_PLUGINS). +UNSAFE_DEFAULT_DISABLED_PLUGINS = {"executecode"} + def load_plugins(): # Clear existing plugins first but modify the global dict in place plugin_approaches.clear() + # Slugs explicitly opted in via --enable-unsafe-plugins (env: + # OPTILLM_ENABLE_UNSAFE_PLUGINS). The env fallback covers the initial + # load in main(), which runs before parse_args() populates server_config. + enabled_arg = server_config.get('enable_unsafe_plugins', '') or os.environ.get('OPTILLM_ENABLE_UNSAFE_PLUGINS', '') + enabled_unsafe_plugins = {slug.strip() for slug in enabled_arg.split(',') if slug.strip()} + # Get installed package plugins directory import optillm package_plugin_dir = os.path.join(os.path.dirname(optillm.__file__), 'plugins') @@ -329,6 +340,12 @@ def load_plugins(): spec.loader.exec_module(module) if hasattr(module, 'SLUG') and hasattr(module, 'run'): + if module.SLUG in UNSAFE_DEFAULT_DISABLED_PLUGINS and module.SLUG not in enabled_unsafe_plugins: + logger.warning( + f"Plugin '{module.SLUG}' is disabled by default because it executes arbitrary code. " + f"Start optillm with --enable-unsafe-plugins={module.SLUG} " + f"(env: OPTILLM_ENABLE_UNSAFE_PLUGINS) to allow it.") + continue if module.SLUG in plugin_approaches: logger.info(f"Overriding {source} plugin: {module.SLUG}") plugin_approaches[module.SLUG] = module.run @@ -405,6 +422,11 @@ def execute_single_approach(approach, system_prompt, initial_query, client, mode # Note: 'n' is NOT removed - the none_approach passes it to the client which handles multiple completions kwargs.pop('stream', None) # stream is handled by proxy() + # _execute_completion copies max_completion_tokens to max_tokens for approaches; + # don't forward the copy, since o-series and gpt-5 models reject max_tokens. + if kwargs.get('max_completion_tokens') is not None: + kwargs.pop('max_tokens', None) + # Pass the original messages through when available so multi-turn # structure, tool_calls and tool results reach the provider intact. # Otherwise reconstruct them from system_prompt and initial_query. @@ -1260,6 +1282,7 @@ def parse_args(): ("--log-conversations", "OPTILLM_LOG_CONVERSATIONS", bool, False, "Enable conversation logging with full metadata"), ("--conversation-log-dir", "OPTILLM_CONVERSATION_LOG_DIR", str, str(Path.home() / ".optillm" / "conversations"), "Directory to save conversation logs"), ("--readurls-allow-internal", "OPTILLM_READURLS_ALLOW_INTERNAL", bool, False, "Allow the readurls plugin to fetch non-public/internal URLs (disables SSRF protection)"), + ("--enable-unsafe-plugins", "OPTILLM_ENABLE_UNSAFE_PLUGINS", str, "", "Comma-separated slugs of code-executing plugins (e.g. executecode) to enable; such plugins are disabled by default"), ] for arg, env, type_, default, help_text, *extra in args_env: @@ -1354,6 +1377,10 @@ def main(): # Update server_config with all argument values server_config.update(vars(args)) + # Reload plugins now that server_config reflects the CLI/env options, so + # --enable-unsafe-plugins (and --plugins-dir) take effect for the loaded set. + load_plugins() + port = server_config['port'] # Initialize request batcher if batch mode is enabled diff --git a/pyproject.toml b/pyproject.toml index 28f35b2b..16a5a5cb 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "optillm" -version = "0.4.0" +version = "0.4.1" description = "An optimizing inference proxy for LLMs." readme = "README.md" license = "Apache-2.0" @@ -86,6 +86,7 @@ optillm = [ "plugins/proxy/*.py", "plugins/proxy/*.yaml", "plugins/proxy/*.md", + "plugins/streaming_hints/*.json", "cepo/*.py", "cepo/configs/*.yaml", ] \ No newline at end of file diff --git a/tests/test_api_compatibility.py b/tests/test_api_compatibility.py index 5324acd1..d2585267 100644 --- a/tests/test_api_compatibility.py +++ b/tests/test_api_compatibility.py @@ -11,6 +11,8 @@ # Add parent directory to path for imports sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) +# pytest puts the repo root on sys.path because tests/ is a package, not tests/ itself +sys.path.append(os.path.dirname(os.path.abspath(__file__))) # Import test utilities from test_utils import setup_test_env, get_test_client, TEST_MODEL diff --git a/tests/test_entropy_decoding_sampling.py b/tests/test_entropy_decoding_sampling.py index 6483b5a5..b1ed8a1b 100644 --- a/tests/test_entropy_decoding_sampling.py +++ b/tests/test_entropy_decoding_sampling.py @@ -22,7 +22,8 @@ def test_top_p_keeps_the_most_likely_token(): assert 0 in tokens # The nucleus keeps tokens until those before them exceed 0.9: 36 of 40. - assert tokens == set(range(36)) + # The 39 tail tokens tie, so which of them survive depends on the platform. + assert len(tokens) == 36 def test_min_p_removes_unlikely_tokens(): diff --git a/tests/test_max_completion_tokens.py b/tests/test_max_completion_tokens.py new file mode 100644 index 00000000..c8d701ac --- /dev/null +++ b/tests/test_max_completion_tokens.py @@ -0,0 +1,129 @@ +#!/usr/bin/env python3 +""" +Regression tests for issue #350: pass-through requests must carry the +max_completion_tokens the client sent and not add a max_tokens copy. +""" + +import os +import sys +import unittest +from unittest.mock import MagicMock, patch + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +import optillm.server as server +from optillm.inference import InferenceClient +from optillm.plugins import proxy_plugin +from optillm.plugins.proxy.config import ProxyConfig + +COMPLETION = { + "id": "chatcmpl-1", "object": "chat.completion", "created": 1, "model": "o4-mini", + "choices": [{"index": 0, "finish_reason": "stop", "message": {"role": "assistant", "content": "8"}}], + "usage": {"prompt_tokens": 3, "completion_tokens": 1, "total_tokens": 4}, +} +MESSAGES = [{"role": "user", "content": "What is 3 + 5?"}] + + +class TestNonePassthrough(unittest.TestCase): + def post(self, **body): + upstream = MagicMock() + upstream.chat.completions.create.return_value = COMPLETION + with patch.object(server, "get_config", return_value=(upstream, "key")): + resp = server.app.test_client().post( + "/v1/chat/completions", json={"model": "none-o4-mini", "messages": MESSAGES, **body}) + self.assertEqual(resp.status_code, 200) + return upstream.chat.completions.create.call_args.kwargs + + def test_max_completion_tokens_is_not_copied_upstream(self): + kwargs = self.post(max_completion_tokens=1024) + self.assertEqual(kwargs["max_completion_tokens"], 1024) + self.assertNotIn("max_tokens", kwargs) + + def test_max_completion_tokens_wins_when_both_are_sent(self): + kwargs = self.post(max_tokens=10, max_completion_tokens=1024) + self.assertEqual(kwargs["max_completion_tokens"], 1024) + self.assertNotIn("max_tokens", kwargs) + + def test_max_tokens_alone_is_forwarded(self): + kwargs = self.post(max_tokens=10) + self.assertEqual(kwargs["max_tokens"], 10) + self.assertNotIn("max_completion_tokens", kwargs) + + +class TestApproachesStillSeeTheLimit(unittest.TestCase): + def test_non_passthrough_approach_receives_max_tokens(self): + """The copy stays in request_config for approaches that read max_tokens.""" + upstream = MagicMock() + seen = {} + + def fake_approach(approach, system_prompt, initial_query, client, model, request_config, *args, **kwargs): + seen.update(request_config) + return "8", 1 + + with patch.object(server, "get_config", return_value=(upstream, "key")), \ + patch.object(server, "execute_single_approach", side_effect=fake_approach): + resp = server.app.test_client().post("/v1/chat/completions", json={ + "model": "moa-o4-mini", "messages": MESSAGES, "max_completion_tokens": 1024}) + self.assertEqual(resp.status_code, 200) + self.assertEqual(seen["max_tokens"], 1024) + + +class TestProxyPassthrough(unittest.TestCase): + CONFIG = {"max_completion_tokens": 1024, "max_tokens": 1024, "stream": False, "n": 1} + + def run_plugin(self, **load_patch): + upstream = MagicMock() + upstream.chat.completions.create.return_value = COMPLETION + with patch.object(ProxyConfig, "load", **load_patch): + proxy_plugin.run("system", "hi", upstream, "o4-mini", dict(self.CONFIG), MESSAGES) + return upstream.chat.completions.create.call_args.kwargs + + def assert_limit_only(self, kwargs): + self.assertEqual(kwargs["max_completion_tokens"], 1024) + self.assertNotIn("max_tokens", kwargs) + self.assertNotIn("stream", kwargs) + + def test_no_providers_route(self): + self.assert_limit_only(self.run_plugin(return_value={"providers": []})) + + def test_error_fallback_route(self): + self.assert_limit_only(self.run_plugin(side_effect=RuntimeError("bad config"))) + + def test_provider_route(self): + config = {"providers": [{"name": "p"}], "routing": {}, "timeouts": {}, "queue": {}} + proxy_client = MagicMock() + proxy_client.chat.completions.create.return_value = COMPLETION + proxy_plugin._proxy_client_cache.clear() + self.addCleanup(proxy_plugin._proxy_client_cache.clear) + with patch.object(ProxyConfig, "load", return_value=config), \ + patch.object(proxy_plugin, "ProxyClient", return_value=proxy_client): + proxy_plugin.run("system", "hi", MagicMock(), "o4-mini", dict(self.CONFIG), MESSAGES) + self.assert_limit_only(proxy_client.chat.completions.create.call_args.kwargs) + + +class TestLocalInference(unittest.TestCase): + """The local client only used to read max_tokens, so it must honor the preferred name.""" + + def generated_limit(self, **limits): + client = InferenceClient.__new__(InferenceClient) + pipeline = MagicMock() + pipeline.generate.return_value = (["8"], [1], [None]) + pipeline.tokenizer.encode.return_value = [1, 2, 3] + client.get_pipeline = MagicMock(return_value=pipeline) + completions = InferenceClient.Chat.Completions(client) + with patch("optillm.inference.count_reasoning_tokens", return_value=0): + completions.create(messages=MESSAGES, model="m", **limits) + return pipeline.generate.call_args.kwargs["generation_params"]["max_new_tokens"] + + def test_max_completion_tokens_limits_generation(self): + self.assertEqual(self.generated_limit(max_completion_tokens=7), 7) + + def test_max_completion_tokens_wins_over_max_tokens(self): + self.assertEqual(self.generated_limit(max_completion_tokens=7, max_tokens=99), 7) + + def test_max_tokens_still_limits_generation(self): + self.assertEqual(self.generated_limit(max_tokens=99), 99) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_mcp_plugin.py b/tests/test_mcp_plugin.py index 621d6f1d..825dfe29 100644 --- a/tests/test_mcp_plugin.py +++ b/tests/test_mcp_plugin.py @@ -7,6 +7,7 @@ import os import asyncio import json +import logging import pytest from unittest.mock import Mock, AsyncMock, patch, MagicMock from pathlib import Path @@ -24,6 +25,7 @@ execute_tool, execute_tool_stdio, execute_tool_sse, execute_tool_websocket, execute_tool_streamable_http, execute_tool_with_session, LoggingClientSession, SLUG ) +from optillm.plugins.mcp_plugin import logger as mcp_logger class TestServerConfig: @@ -372,6 +374,49 @@ def test_get_capabilities_description_no_servers(self): assert "No MCP servers available" in description +class TestSecretsAreNotLogged: + """Env and header values in the MCP config hold API keys and must stay out of the log file""" + + SECRET = "s3cret-value-for-log-test" + + @pytest.fixture(autouse=True) + def restore_log_level(self): + level = mcp_logger.level + mcp_logger.setLevel(logging.DEBUG) + yield + mcp_logger.setLevel(level) + + def test_load_config(self, tmp_path, caplog): + path = tmp_path / "mcp_config.json" + path.write_text(json.dumps({"log_level": "DEBUG", "mcpServers": { + "local": {"command": "cmd", "env": {"API_KEY": self.SECRET}}, + "remote": {"transport": "sse", "url": "http://x", "headers": {"Authorization": self.SECRET}}, + }})) + with caplog.at_level(logging.DEBUG, logger=mcp_logger.name): + assert MCPConfigManager(str(path)).load_config() + assert caplog.records + assert self.SECRET not in caplog.text + + @pytest.mark.asyncio + async def test_connect_stdio(self, caplog): + config = ServerConfig(transport="stdio", command="definitely-not-a-command-xyz", env={"API_KEY": self.SECRET}) + with caplog.at_level(logging.DEBUG, logger=mcp_logger.name), \ + patch("optillm.plugins.mcp_plugin.find_executable", return_value="/bin/false"), \ + patch("optillm.plugins.mcp_plugin.stdio_client", side_effect=RuntimeError("stop")): + await MCPServer("local", config).connect_stdio_native() + assert "API_KEY" in caplog.text + assert self.SECRET not in caplog.text + + @pytest.mark.asyncio + async def test_connect_sse(self, caplog): + config = ServerConfig(transport="sse", url="http://127.0.0.1:9/mcp", headers={"Authorization": self.SECRET}) + with caplog.at_level(logging.DEBUG, logger=mcp_logger.name), \ + patch("optillm.plugins.mcp_plugin.sse_client", side_effect=RuntimeError("stop")): + await MCPServer("remote", config).connect_sse() + assert "Authorization" in caplog.text + assert self.SECRET not in caplog.text + + @pytest.mark.asyncio @pytest.mark.skipif(not os.getenv("GITHUB_TOKEN"), reason="GITHUB_TOKEN not set") class TestGitHubMCPServer: diff --git a/tests/test_plugin_gating.py b/tests/test_plugin_gating.py new file mode 100644 index 00000000..6450c8f7 --- /dev/null +++ b/tests/test_plugin_gating.py @@ -0,0 +1,104 @@ +#!/usr/bin/env python3 +""" +Tests for the plugin gating that disables code-executing plugins (executecode) +by default and requires an explicit opt-in (issue #302). + +All tests are offline: they exercise load_plugins() registration against the +packaged plugin directory and the pure helper functions of the executecode +plugin. No server, LLM client, or Jupyter kernel is started. +""" + +import sys +import os + +import pytest + +sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +from optillm import plugin_approaches, load_plugins +from optillm.server import UNSAFE_DEFAULT_DISABLED_PLUGINS, server_config +from optillm.plugins.executecode_plugin import ( + extract_python_code, + should_execute_request_code, +) + +ENV_VAR = "OPTILLM_ENABLE_UNSAFE_PLUGINS" + + +@pytest.fixture(autouse=True) +def restore_default_plugins(monkeypatch): + """Make every test start (and leave behind) the default plugin registry.""" + monkeypatch.setitem(server_config, 'enable_unsafe_plugins', '') + monkeypatch.delenv(ENV_VAR, raising=False) + yield + load_plugins() + + +def test_unsafe_plugin_set_contains_executecode(): + assert "executecode" in UNSAFE_DEFAULT_DISABLED_PLUGINS + + +def test_executecode_disabled_by_default(): + load_plugins() + assert "executecode" not in plugin_approaches, \ + "executecode executes arbitrary code and must be opt-in" + # Gating must not disable unrelated plugins. + assert "memory" in plugin_approaches + assert "readurls" in plugin_approaches + + +def test_executecode_opt_in_via_server_config(monkeypatch): + monkeypatch.setitem(server_config, 'enable_unsafe_plugins', 'executecode') + load_plugins() + assert "executecode" in plugin_approaches + + +def test_executecode_opt_in_via_env_var(monkeypatch): + monkeypatch.setenv(ENV_VAR, "executecode") + load_plugins() + assert "executecode" in plugin_approaches + + +def test_opt_in_accepts_comma_separated_list(monkeypatch): + monkeypatch.setitem(server_config, 'enable_unsafe_plugins', 'memory, executecode ,') + load_plugins() + assert "executecode" in plugin_approaches + assert "memory" in plugin_approaches + + +def test_disabled_plugin_logs_how_to_enable(caplog): + import logging + with caplog.at_level(logging.WARNING, logger="optillm.server"): + load_plugins() + warnings = [r.message for r in caplog.records if "executecode" in r.message] + assert warnings, "expected a warning explaining why executecode was skipped" + assert "--enable-unsafe-plugins=executecode" in warnings[0] + assert ENV_VAR in warnings[0] + + +def test_extract_python_code_multiple_blocks(): + text = "Before\n```python\nprint(1)\n```\nmiddle\n```python\nprint(2)\n```\nafter" + assert extract_python_code(text) == ["print(1)", "print(2)"] + + +def test_extract_python_code_no_blocks(): + assert extract_python_code("no code here at all") == [] + + +def test_extract_python_code_requires_python_tag(): + # A bare fenced block without the `python` tag must not be extracted. + assert extract_python_code("```\nprint(1)\n```") == [] + + +def test_should_execute_request_code_keyword_hits(): + for query in ["please run this", "execute it", "show the output", "give me the result"]: + assert should_execute_request_code(query) is True, query + + +def test_should_execute_request_code_keyword_misses(): + assert should_execute_request_code("what is the capital of France?") is False + assert should_execute_request_code("") is False + + +if __name__ == "__main__": + sys.exit(pytest.main([__file__, "-v"])) diff --git a/tests/test_plugins.py b/tests/test_plugins.py index ace9d2ef..bdee792d 100644 --- a/tests/test_plugins.py +++ b/tests/test_plugins.py @@ -33,7 +33,8 @@ def test_plugin_module_imports(): 'optillm.plugins.spl_plugin', 'optillm.plugins.proxy_plugin', 'optillm.plugins.mcp_plugin', - 'optillm.plugins.compact_plugin' + 'optillm.plugins.compact_plugin', + 'optillm.plugins.streaming_hints_plugin' ] for module_name in plugin_modules: @@ -54,7 +55,7 @@ def test_plugin_approach_detection(): load_plugins() # Check if known plugins are loaded - expected_plugins = ["memory", "readurls", "privacy", "web_search", "deep_research", "deepthink", "longcepo", "spl", "proxy", "mcp", "compact"] + expected_plugins = ["memory", "readurls", "privacy", "web_search", "deep_research", "deepthink", "longcepo", "spl", "proxy", "mcp", "compact", "streaming_hints"] for plugin_name in expected_plugins: assert plugin_name in plugin_approaches, f"Plugin {plugin_name} not loaded" diff --git a/tests/test_streaming_hints.py b/tests/test_streaming_hints.py new file mode 100644 index 00000000..f1c7b3a2 --- /dev/null +++ b/tests/test_streaming_hints.py @@ -0,0 +1,190 @@ +from types import SimpleNamespace + +from optillm.plugins.streaming_hints import ( + Hint, + HintMemoryBank, + StreamingHintRetriever, + run_streaming_hints, +) +from optillm.plugins import streaming_hints_plugin + +PROBLEM = "A train travels 120 km in 1.5 hours. What is its speed in meters per second?" + + +def response(text, reason="stop", tokens=10): + return SimpleNamespace( + choices=[SimpleNamespace(message=SimpleNamespace(content=text), finish_reason=reason)], + usage=SimpleNamespace(completion_tokens=tokens), + ) + + +def fake_client(*responses): + calls = [] + queue = list(responses) + + def create(**kwargs): + calls.append(kwargs) + return queue.pop(0) + + client = SimpleNamespace(chat=SimpleNamespace(completions=SimpleNamespace(create=create))) + return client, calls + + +def units_bank(): + return HintMemoryBank([ + Hint(id="units", text="Check the units of every quantity.", keywords=["units", "speed", "meters", "hours"]), + Hint(id="cases", text="Split into cases.", keywords=["count", "cases", "arrangements"]), + ]) + + +def test_default_bank_loads_and_ranks_relevant_hint_first(): + bank = HintMemoryBank.from_file() + assert len(bank.hints) >= 10 + assert len({h.id for h in bank.hints}) == len(bank.hints) + + (best, score), = bank.search("how many ways can 5 people be arranged in a row", k=1) + assert best.id == "math-casework" + assert score > 0 + + +def test_search_skips_excluded_and_unrelated_queries(): + bank = units_bank() + assert bank.search(PROBLEM, k=1)[0][0].id == "units" + assert bank.search(PROBLEM, k=1, exclude=["units"]) == [] + assert bank.search("zzz qqq", k=2) == [] + + +def test_retriever_waits_for_uncertainty_signal_and_gap(): + retriever = StreamingHintRetriever(units_bank(), threshold=0.1, min_gap_chars=20) + retriever.reset(PROBLEM) + + # A confident sentence is not an injection point. + assert retriever.feed("The distance is 120 km and the time is 1.5 hours. ") is None + + text = "So the speed is 80. Wait, the answer must be in meters per second. " + injection = retriever.feed(text) + assert injection is not None + assert injection.hint.id == "units" + assert text[: injection.offset].endswith("meters per second.") + assert injection.hint.uses == 1 + + # Each hint is injected at most once. + assert retriever.feed("Hmm, let me check that again. " * 3) is None + + +def test_retriever_handles_sentences_split_across_feeds(): + retriever = StreamingHintRetriever(units_bank(), threshold=0.1, min_gap_chars=0) + retriever.reset(PROBLEM) + + assert retriever.feed("Hmm, I am not sure about the un") is None + injection = retriever.feed("its here. More text that is discarded.") + assert injection is not None + assert injection.offset == len("its here.") + + +def test_retriever_respects_max_hints_and_threshold(): + retriever = StreamingHintRetriever(units_bank(), threshold=0.99, min_gap_chars=0) + retriever.reset(PROBLEM) + assert retriever.feed("Wait, what are the units of speed here? ") is None + + retriever = StreamingHintRetriever(units_bank(), threshold=0.1, min_gap_chars=0, max_hints=0) + retriever.reset(PROBLEM) + assert retriever.feed("Wait, what are the units of speed here? ") is None + + +def test_run_injects_hint_and_continues_from_cut(): + first = ( + "Speed is distance over time, so 120 / 1.5 = 80. " + "Wait, the question asks for meters per second. " + "This tail is discarded." + ) + client, calls = fake_client( + response(first, tokens=30), + response(" 80 km/h is 80000 m / 3600 s, about 22.2 m/s.", tokens=20), + ) + + answer, tokens = run_streaming_hints( + "system", PROBLEM, client, "model", + request_config={"hint_threshold": 0.1, "min_gap_chars": 20}, + memory_bank=units_bank(), + ) + + assert tokens == 50 + assert "discarded" not in answer + assert "" not in answer + assert answer.endswith("about 22.2 m/s.") + + assert len(calls) == 2 + assert calls[0]["messages"][1]["content"] == PROBLEM + continuation = calls[1]["messages"][1]["content"] + assert continuation.startswith(PROBLEM) + assert "Check the units of every quantity." in continuation + assert "meters per second.\n" in continuation + assert "discarded" not in continuation + + +def test_run_show_hints_returns_reasoning_with_hint_tags(): + client, _ = fake_client( + response("Speed is 80. Wait, the units are meters per second. Tail."), + response(" So 22.2 m/s."), + ) + answer, _ = run_streaming_hints( + "system", PROBLEM, client, "model", + request_config={"hint_threshold": 0.1, "min_gap_chars": 0, "show_hints": True}, + memory_bank=units_bank(), + ) + assert "Check the units of every quantity." in answer + assert answer.endswith("So 22.2 m/s.") + + +def test_run_continues_after_length_cutoff_and_respects_token_budget(): + client, calls = fake_client( + response("Part one", reason="length", tokens=6), + response(" part two", reason="length", tokens=4), + response(" never requested"), + ) + answer, tokens = run_streaming_hints( + "system", "question", client, "model", + request_config={"max_tokens": 10, "segment_tokens": 6}, + memory_bank=units_bank(), + ) + assert answer == "Part one part two" + assert tokens == 10 + assert [c["max_tokens"] for c in calls] == [6, 4] + + +def test_run_treats_full_budget_as_cut_off_when_finish_reason_is_stop(): + # The local inference server reports "stop" even when max_tokens was hit. + client, calls = fake_client( + response("Part one", reason="stop", tokens=6), + response(" part two", reason="stop", tokens=3), + ) + answer, tokens = run_streaming_hints( + "system", "question", client, "model", + request_config={"segment_tokens": 6}, + memory_bank=units_bank(), + ) + assert answer == "Part one part two" + assert tokens == 9 + assert len(calls) == 2 + + +def test_run_without_relevant_hints_is_a_single_call(): + client, calls = fake_client(response("The answer is 4. Wait, yes, 4.", tokens=7)) + answer, tokens = run_streaming_hints( + "system", "What is 2 + 2?", client, "model", + request_config={"min_gap_chars": 0, "temperature": 0.3}, + memory_bank=units_bank(), + ) + assert answer == "The answer is 4. Wait, yes, 4." + assert tokens == 7 + assert len(calls) == 1 + assert calls[0]["temperature"] == 0.3 + + +def test_plugin_entry_point(): + assert streaming_hints_plugin.SLUG == "streaming_hints" + client, _ = fake_client(response("Done.")) + answer, tokens = streaming_hints_plugin.run("system", "question", client, "model") + assert answer == "Done." + assert tokens == 10