chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

Add Ollama-compatible API endpoints + Smart Router (#14376)

Browse Source 
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Co-authored-by: Xinyuan Tong <115166877+JustinTong0323@users.noreply.github.com>
This commit is contained in:
Alison Shao
2025-12-16 20:43:38 -08:00
committed by GitHub
parent 0129c911e0
commit 31d48d7f6f
8 changed files with 1032 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

91
docs/basic_usage/ollama_api.md Normal file
View File

@@ -0,0 +1,91 @@
# Ollama-Compatible API
SGLang provides Ollama API compatibility, allowing you to use the Ollama CLI and Python library with SGLang as the inference backend.
## Prerequisites
```bash
# Install the Ollama Python library (for Python client usage)
pip install ollama
```
> **Note**: You don't need the Ollama server installed - SGLang acts as the backend. You only need the `ollama` CLI or Python library as the client.
## Endpoints
| Endpoint | Method | Description |
|----------|--------|-------------|
| `/` | GET, HEAD | Health check for Ollama CLI |
| `/api/tags` | GET | List available models |
| `/api/chat` | POST | Chat completions (streaming & non-streaming) |
| `/api/generate` | POST | Text generation (streaming & non-streaming) |
| `/api/show` | POST | Model information |
## Quick Start
### 1. Launch SGLang Server
```bash
python -m sglang.launch_server \
--model Qwen/Qwen2.5-1.5B-Instruct \
--port 30001 \
--host 0.0.0.0
```
> **Note**: The model name used with `ollama run` must match exactly what you passed to `--model`.
### 2. Use Ollama CLI
```bash
# List available models
OLLAMA_HOST=http://localhost:30001 ollama list
# Interactive chat
OLLAMA_HOST=http://localhost:30001 ollama run "Qwen/Qwen2.5-1.5B-Instruct"
```
If connecting to a remote server behind a firewall:
```bash
# SSH tunnel
ssh -L 30001:localhost:30001 user@gpu-server -N &
# Then use Ollama CLI as above
OLLAMA_HOST=http://localhost:30001 ollama list
```
### 3. Use Ollama Python Library
```python
import ollama
client = ollama.Client(host='http://localhost:30001')
# Non-streaming
response = client.chat(
model='Qwen/Qwen2.5-1.5B-Instruct',
messages=[{'role': 'user', 'content': 'Hello!'}]
)
print(response['message']['content'])
# Streaming
stream = client.chat(
model='Qwen/Qwen2.5-1.5B-Instruct',
messages=[{'role': 'user', 'content': 'Tell me a story'}],
stream=True
)
for chunk in stream:
print(chunk['message']['content'], end='', flush=True)
```
## Smart Router
For intelligent routing between local Ollama (fast) and remote SGLang (powerful) using an LLM judge, see the [Smart Router documentation](https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/entrypoints/ollama/README.md).
## Summary
| Component | Purpose |
|-----------|---------|
| **Ollama API** | Familiar CLI/API that developers already know |
| **SGLang Backend** | High-performance inference engine |
| **Smart Router** | Intelligent routing - fast local for simple tasks, powerful remote for complex tasks |

1
docs/index.rst
View File

@@ -23,6 +23,7 @@ Its core features include:
basic_usage/send_request.ipynb
basic_usage/openai_api.rst
basic_usage/ollama_api.md
basic_usage/offline_engine_api.ipynb
basic_usage/native_api.ipynb
basic_usage/sampling_params.md

45
python/sglang/srt/entrypoints/http_server.py
View File

@@ -54,6 +54,12 @@ from fastapi.responses import ORJSONResponse, Response, StreamingResponse
from sglang.srt.disaggregation.utils import FAKE_BOOTSTRAP_HOST, DisaggregationMode
from sglang.srt.entrypoints.engine import _launch_subprocesses
from sglang.srt.entrypoints.ollama.protocol import (
OllamaChatRequest,
OllamaGenerateRequest,
OllamaShowRequest,
)
from sglang.srt.entrypoints.ollama.serving import OllamaServing
from sglang.srt.entrypoints.openai.protocol import (
ChatCompletionRequest,
ClassifyRequest,
@@ -281,6 +287,9 @@ async def lifespan(fast_api_app: FastAPI):
_global_state.tokenizer_manager
)
# Initialize Ollama-compatible serving handler
fast_api_app.state.ollama_serving = OllamaServing(_global_state.tokenizer_manager)
# Launch tool server
tool_server = None
if server_args.tool_server == "demo":
@@ -1363,6 +1372,42 @@ async def v1_rerank_request(request: V1RerankReqInput, raw_request: Request):
)
##### Ollama-compatible API endpoints #####
@app.get(os.environ.get("SGLANG_OLLAMA_ROOT_ROUTE", "/"))
@app.head(os.environ.get("SGLANG_OLLAMA_ROOT_ROUTE", "/"))
async def ollama_root():
"""Ollama-compatible root endpoint for health check."""
return "Ollama is running"
@app.post(os.environ.get("SGLANG_OLLAMA_CHAT_ROUTE", "/api/chat"))
async def ollama_chat(request: OllamaChatRequest, raw_request: Request):
"""Ollama-compatible chat endpoint."""
return await raw_request.app.state.ollama_serving.handle_chat(request, raw_request)
@app.post(os.environ.get("SGLANG_OLLAMA_GENERATE_ROUTE", "/api/generate"))
async def ollama_generate(request: OllamaGenerateRequest, raw_request: Request):
"""Ollama-compatible generate endpoint."""
return await raw_request.app.state.ollama_serving.handle_generate(
request, raw_request
)
@app.get(os.environ.get("SGLANG_OLLAMA_TAGS_ROUTE", "/api/tags"))
async def ollama_tags(raw_request: Request):
"""Ollama-compatible list models endpoint."""
return raw_request.app.state.ollama_serving.get_tags()
@app.post(os.environ.get("SGLANG_OLLAMA_SHOW_ROUTE", "/api/show"))
async def ollama_show(request: OllamaShowRequest, raw_request: Request):
"""Ollama-compatible show model info endpoint."""
return raw_request.app.state.ollama_serving.get_show(request.model)
## SageMaker API
@app.get("/ping")
async def sagemaker_health() -> Response:

112
python/sglang/srt/entrypoints/ollama/README.md Normal file
View File

@@ -0,0 +1,112 @@
# SGLang Ollama Integration
Ollama API compatibility for SGLang, plus a Smart Router for intelligent routing between local and remote models.
## Features
1. **Ollama-compatible API** - Use Ollama CLI/library with SGLang backend
2. **Smart Router** - LLM-based routing between local and remote models
## Ollama API
For basic Ollama API usage with SGLang (CLI and Python examples), see the [Ollama API documentation](https://sgl-project.github.io/basic_usage/ollama_api.html).
## Smart Router
### Prerequisites
```bash
pip install ollama
```
Intelligently routes requests between local Ollama and remote SGLang using an LLM judge.
### How It Works
```
User Request
│
▼
┌─────────────────────┐
│ LLM Judge │ Classifies as SIMPLE or COMPLEX
│ (local model) │
└─────────────────────┘
│
▼
┌─────────────────────┐
│ SIMPLE → Local │ Fast response from local Ollama
│ COMPLEX → Remote │ Powerful response from SGLang
└─────────────────────┘
```
The LLM judge (running on local Ollama) analyzes each request and decides:
- **SIMPLE**: Quick responses, greetings, factual questions, definitions, basic Q&A
- **COMPLEX**: Deep reasoning, multi-step analysis, long explanations, creative writing
### Setup
**Terminal 1: Local Ollama**
```bash
ollama pull <LOCAL_MODEL> # e.g., llama3.2, mistral, phi3
ollama serve # This will block the terminal
```
**Terminal 2: Remote SGLang (GPU)**
```bash
ssh user@gpu-server
python -m sglang.launch_server --model <REMOTE_MODEL> --port 30001 --host 0.0.0.0
```
**Terminal 3: Smart Router**
```bash
ssh -L 30001:localhost:30001 user@gpu-server -N &
python python/sglang/srt/entrypoints/ollama/smart_router.py
```
### Configuration
```python
from sglang.srt.entrypoints.ollama.smart_router import SmartRouter
router = SmartRouter(
# Local Ollama
local_host="http://localhost:11434",
local_model="llama3.2", # or any Ollama model
# Remote SGLang
remote_host="http://localhost:30001",
remote_model="Qwen/Qwen2.5-1.5B-Instruct", # or any HuggingFace model
# LLM Judge (optional, defaults to local_model)
judge_model="llama3.2",
)
```
### Usage
```python
# Auto-routing via LLM judge
response = router.chat("Hello!", verbose=True)
# [Router] LLM Judge: SIMPLE
# [Router] -> Local Ollama | Model: llama3.2
response = router.chat("Explain quantum computing in detail", verbose=True)
# [Router] LLM Judge: COMPLEX
# [Router] -> Remote SGLang | Model: Qwen/Qwen2.5-1.5B-Instruct
# Force routing (skip LLM judge)
response = router.chat("question", force_local=True)
response = router.chat("question", force_remote=True)
# Streaming
for chunk in router.chat_stream("Tell me a story"):
print(chunk['message']['content'], end='')
```
---
## Value
- **Ollama**: Simple CLI/API developers already know
- **SGLang**: High-performance inference
- **Smart Router**: Intelligent routing - fast local for simple tasks, powerful remote for complex tasks

1
python/sglang/srt/entrypoints/ollama/__init__.py Normal file
View File

@@ -0,0 +1 @@
# Ollama-compatible API for SGLang

137
python/sglang/srt/entrypoints/ollama/protocol.py Normal file
View File

@@ -0,0 +1,137 @@
"""
Ollama-compatible API protocol definitions.
These models match the Ollama API format:
https://github.com/ollama/ollama/blob/main/docs/api.md
"""
from typing import Any, Dict, List, Literal, Optional, Union
from pydantic import BaseModel, Field
class OllamaMessage(BaseModel):
"""Ollama message format."""
role: str
content: str
images: Optional[List[str]] = None
class OllamaChatRequest(BaseModel):
"""Ollama /api/chat request format."""
model: str
messages: List[OllamaMessage]
stream: bool = True
format: Optional[Union[Literal["json"], Dict[str, Any]]] = None
options: Optional[Dict[str, Any]] = None
keep_alive: Optional[Union[float, str]] = None
think: Optional[Union[bool, Literal["low", "medium", "high"]]] = None
class OllamaChatResponse(BaseModel):
"""Ollama /api/chat response format (non-streaming)."""
model: str
created_at: str
message: OllamaMessage
done: bool = True
done_reason: Optional[str] = "stop"
total_duration: Optional[int] = None
load_duration: Optional[int] = None
prompt_eval_count: Optional[int] = None
prompt_eval_duration: Optional[int] = None
eval_count: Optional[int] = None
eval_duration: Optional[int] = None
class OllamaChatStreamResponse(BaseModel):
"""Ollama /api/chat streaming response chunk."""
model: str
created_at: str
message: OllamaMessage
done: bool = False
done_reason: Optional[str] = None
class OllamaGenerateRequest(BaseModel):
"""Ollama /api/generate request format."""
model: str
prompt: str
suffix: Optional[str] = None
system: Optional[str] = None
template: Optional[str] = None
context: Optional[List[int]] = None
stream: bool = True
raw: bool = False
format: Optional[Union[Literal["json"], Dict[str, Any]]] = None
options: Optional[Dict[str, Any]] = None
keep_alive: Optional[Union[float, str]] = None
images: Optional[List[str]] = None
think: Optional[bool] = None
class OllamaGenerateResponse(BaseModel):
"""Ollama /api/generate response format (non-streaming)."""
model: str
created_at: str
response: str
done: bool = True
done_reason: Optional[str] = "stop"
context: Optional[List[int]] = None
total_duration: Optional[int] = None
load_duration: Optional[int] = None
prompt_eval_count: Optional[int] = None
prompt_eval_duration: Optional[int] = None
eval_count: Optional[int] = None
eval_duration: Optional[int] = None
class OllamaGenerateStreamResponse(BaseModel):
"""Ollama /api/generate streaming response chunk."""
model: str
created_at: str
response: str
done: bool = False
done_reason: Optional[str] = None
class OllamaModelInfo(BaseModel):
"""Model information for /api/tags response."""
name: str
model: str
modified_at: str
size: int
digest: str
details: Optional[Dict[str, Any]] = None
class OllamaTagsResponse(BaseModel):
"""Ollama /api/tags response format."""
models: List[OllamaModelInfo]
class OllamaShowRequest(BaseModel):
"""Ollama /api/show request format."""
model: str
class OllamaShowResponse(BaseModel):
"""Ollama /api/show response format."""
license: str = ""
modelfile: str = ""
parameters: str = ""
template: str = ""
modified_at: str = ""
details: Dict[str, Any] = Field(default_factory=dict)
model_info: Dict[str, Any] = Field(default_factory=dict)
capabilities: List[str] = Field(default_factory=list)

349
python/sglang/srt/entrypoints/ollama/serving.py Normal file
View File

@@ -0,0 +1,349 @@
"""
Ollama-compatible API serving handlers.
This module provides handlers that convert Ollama API requests to SGLang's
internal format and return Ollama-compatible responses.
"""
import time
from datetime import datetime, timezone
from typing import AsyncIterator, Union
import orjson
from fastapi import Request
from fastapi.responses import StreamingResponse
from sglang.srt.entrypoints.ollama.protocol import (
OllamaChatRequest,
OllamaChatResponse,
OllamaChatStreamResponse,
OllamaGenerateRequest,
OllamaGenerateResponse,
OllamaGenerateStreamResponse,
OllamaMessage,
OllamaModelInfo,
OllamaShowResponse,
OllamaTagsResponse,
)
from sglang.srt.managers.io_struct import GenerateReqInput
class OllamaServing:
"""Handler for Ollama-compatible API endpoints."""
def __init__(self, tokenizer_manager):
self.tokenizer_manager = tokenizer_manager
def _get_timestamp(self) -> str:
"""Get current timestamp in Ollama format."""
return datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%S.%f")[:-3] + "Z"
def _convert_options_to_sampling_params(self, options: dict = None) -> dict:
"""Convert Ollama options to SGLang sampling params."""
sampling_params = {}
if options:
# Map Ollama options to SGLang params
param_mapping = {
"temperature": "temperature",
"top_p": "top_p",
"top_k": "top_k",
"num_predict": "max_new_tokens",
"stop": "stop",
"presence_penalty": "presence_penalty",
"frequency_penalty": "frequency_penalty",
"seed": "seed",
}
for ollama_param, sglang_param in param_mapping.items():
if ollama_param in options:
sampling_params[sglang_param] = options[ollama_param]
# Set a reasonable default for max_new_tokens if not specified
# Ollama users typically expect longer responses than SGLang's default (128)
if "max_new_tokens" not in sampling_params:
sampling_params["max_new_tokens"] = 2048
return sampling_params
async def handle_chat(
self, request: OllamaChatRequest, raw_request: Request
) -> Union[OllamaChatResponse, StreamingResponse]:
"""Handle /api/chat endpoint."""
model_name = self.tokenizer_manager.served_model_name
# Convert messages to SGLang format
messages = [
{"role": msg.role, "content": msg.content} for msg in request.messages
]
# Apply chat template using tokenizer
prompt_ids = self.tokenizer_manager.tokenizer.apply_chat_template(
messages,
tokenize=True,
add_generation_prompt=True,
)
# Convert options to sampling params
sampling_params = self._convert_options_to_sampling_params(request.options)
# Create SGLang request with input_ids
gen_request = GenerateReqInput(
input_ids=prompt_ids,
sampling_params=sampling_params,
stream=request.stream,
)
if request.stream:
return await self._stream_chat_response(
gen_request, raw_request, model_name
)
else:
return await self._generate_chat_response(
gen_request, raw_request, model_name
)
async def _generate_chat_response(
self, gen_request: GenerateReqInput, raw_request: Request, model_name: str
) -> OllamaChatResponse:
"""Generate non-streaming chat response."""
start_time = time.time_ns()
# Get response from tokenizer manager
response = await self.tokenizer_manager.generate_request(
gen_request, raw_request
).__anext__()
end_time = time.time_ns()
total_duration = end_time - start_time
output_text = response.get("text", "")
return OllamaChatResponse(
model=model_name,
created_at=self._get_timestamp(),
message=OllamaMessage(role="assistant", content=output_text),
done=True,
done_reason="stop",
total_duration=total_duration,
prompt_eval_count=response.get("meta_info", {}).get("prompt_tokens", None),
eval_count=response.get("meta_info", {}).get("completion_tokens", None),
)
async def _stream_chat_response(
self, gen_request: GenerateReqInput, raw_request: Request, model_name: str
) -> StreamingResponse:
"""Generate streaming chat response."""
async def generate_stream() -> AsyncIterator[bytes]:
previous_text = ""
async for chunk in self.tokenizer_manager.generate_request(
gen_request, raw_request
):
text = chunk.get("text", "")
is_done = chunk.get("meta_info", {}).get("finish_reason") is not None
# Calculate delta (new text since last chunk)
delta = text[len(previous_text) :]
previous_text = text
if is_done:
# Final chunk
response = OllamaChatStreamResponse(
model=model_name,
created_at=self._get_timestamp(),
message=OllamaMessage(role="assistant", content=""),
done=True,
done_reason="stop",
)
else:
response = OllamaChatStreamResponse(
model=model_name,
created_at=self._get_timestamp(),
message=OllamaMessage(role="assistant", content=delta),
done=False,
)
yield orjson.dumps(response.model_dump()) + b"\n"
return StreamingResponse(
generate_stream(),
media_type="application/x-ndjson",
)
async def handle_generate(
self, request: OllamaGenerateRequest, raw_request: Request
) -> Union[OllamaGenerateResponse, StreamingResponse]:
"""Handle /api/generate endpoint."""
model_name = self.tokenizer_manager.served_model_name
# Build prompt
prompt = request.prompt
if request.system:
prompt = f"{request.system}\n\n{prompt}"
# Handle empty prompt - Ollama CLI sends empty requests on initialization
if not prompt or not prompt.strip():
empty_response = OllamaGenerateResponse(
model=model_name,
created_at=self._get_timestamp(),
response="",
done=True,
done_reason="stop",
)
if request.stream:
# Return streaming response with done=True
async def empty_stream() -> AsyncIterator[bytes]:
yield orjson.dumps(empty_response.model_dump()) + b"\n"
return StreamingResponse(
empty_stream(),
media_type="application/x-ndjson",
)
return empty_response
# Convert options to sampling params
sampling_params = self._convert_options_to_sampling_params(request.options)
# Create SGLang request
gen_request = GenerateReqInput(
text=prompt,
sampling_params=sampling_params,
stream=request.stream,
)
if request.stream:
return await self._stream_generate_response(
gen_request, raw_request, model_name
)
else:
return await self._generate_generate_response(
gen_request, raw_request, model_name
)
async def _generate_generate_response(
self, gen_request: GenerateReqInput, raw_request: Request, model_name: str
) -> OllamaGenerateResponse:
"""Generate non-streaming generate response."""
start_time = time.time_ns()
response = await self.tokenizer_manager.generate_request(
gen_request, raw_request
).__anext__()
end_time = time.time_ns()
total_duration = end_time - start_time
output_text = response.get("text", "")
return OllamaGenerateResponse(
model=model_name,
created_at=self._get_timestamp(),
response=output_text,
done=True,
done_reason="stop",
total_duration=total_duration,
prompt_eval_count=response.get("meta_info", {}).get("prompt_tokens", None),
eval_count=response.get("meta_info", {}).get("completion_tokens", None),
)
async def _stream_generate_response(
self, gen_request: GenerateReqInput, raw_request: Request, model_name: str
) -> StreamingResponse:
"""Generate streaming generate response."""
async def generate_stream() -> AsyncIterator[bytes]:
previous_text = ""
async for chunk in self.tokenizer_manager.generate_request(
gen_request, raw_request
):
text = chunk.get("text", "")
is_done = chunk.get("meta_info", {}).get("finish_reason") is not None
# Calculate delta (new text since last chunk)
delta = text[len(previous_text) :]
previous_text = text
if is_done:
response = OllamaGenerateStreamResponse(
model=model_name,
created_at=self._get_timestamp(),
response="",
done=True,
done_reason="stop",
)
else:
response = OllamaGenerateStreamResponse(
model=model_name,
created_at=self._get_timestamp(),
response=delta,
done=False,
)
yield orjson.dumps(response.model_dump()) + b"\n"
return StreamingResponse(
generate_stream(),
media_type="application/x-ndjson",
)
def get_tags(self) -> OllamaTagsResponse:
"""Handle /api/tags endpoint - list available models."""
model_name = self.tokenizer_manager.served_model_name
model_info = OllamaModelInfo(
name=model_name,
model=model_name,
modified_at=self._get_timestamp(),
size=0, # We don't track model size
digest="sha256:sglang0000000000000000000000000000000000000000000000000000000000",
details={
"format": "sglang",
"family": (
model_name.split("/")[-1] if "/" in model_name else model_name
),
"parameter_size": "unknown",
},
)
return OllamaTagsResponse(models=[model_info])
def get_show(self, model: str) -> OllamaShowResponse:
"""Handle /api/show endpoint - show model information."""
model_config = self.tokenizer_manager.model_config
# Extract model family from model name
model_family = model.split("/")[-1] if "/" in model else model
# Remove common suffixes to get base family
for suffix in ["-Instruct", "-Chat", "-Base"]:
if model_family.endswith(suffix):
model_family = model_family[: -len(suffix)]
break
# Build context length info
context_len = model_config.context_len if model_config else 4096
return OllamaShowResponse(
license="", # License info not available from SGLang
modelfile=f"FROM {model}\nPARAMETER num_ctx {context_len}\n",
parameters=f"num_ctx {context_len}",
template="", # Template info not easily accessible
modified_at=self._get_timestamp(),
details={
"parent_model": "",
"format": "sglang",
"family": model_family,
"families": [model_family],
"parameter_size": "unknown",
"quantization_level": "",
},
model_info={
"general.architecture": model_family,
"general.name": model,
"general.parameter_count": 0,
f"{model_family}.context_length": context_len,
f"{model_family}.block_count": 0,
f"{model_family}.embedding_length": 0,
f"{model_family}.attention.head_count": 0,
},
capabilities=["completion"],
)

296
python/sglang/srt/entrypoints/ollama/smart_router.py Normal file
View File

@@ -0,0 +1,296 @@
"""
Smart Router: Automatically routes requests between local Ollama and remote SGLang.
Uses an LLM judge to classify tasks as simple or complex, then routes accordingly:
- Simple tasks → Local Ollama (fast response)
- Complex tasks → Remote SGLang (powerful model)
Usage:
from sglang.srt.entrypoints.ollama.smart_router import SmartRouter
router = SmartRouter(
local_host="http://localhost:11434",
remote_host="http://sglang-server:30001",
)
response = router.chat("Hello!")
"""
from typing import Optional
import ollama
class SmartRouter:
"""Routes requests between local Ollama and remote SGLang using LLM-based classification."""
# Classification prompt for LLM judge
CLASSIFICATION_PROMPT = """You are a task classifier. Classify the following user request into one of two categories.
Categories:
- SIMPLE: Quick responses, greetings, factual questions, definitions, translations, basic Q&A
- COMPLEX: Tasks requiring deep reasoning, multi-step analysis, long explanations, creative writing, detailed research
Reply with ONLY one word: either SIMPLE or COMPLEX.
User request: "{prompt}"
Category:"""
def __init__(
self,
local_host: str = "http://localhost:11434",
remote_host: str = "http://localhost:30001",
local_model: str = "llama3.2",
remote_model: str = "Qwen/Qwen2.5-1.5B-Instruct",
judge_model: Optional[str] = None,
judge_host: Optional[str] = None,
):
"""
Initialize the smart router.
Args:
local_host: URL of local Ollama server
remote_host: URL of remote SGLang server
local_model: Model name for local Ollama
remote_model: Model name for remote SGLang
judge_model: Model for LLM-based classification (default: same as local_model)
judge_host: Host for judge model (default: same as local_host)
"""
self.local_client = ollama.Client(host=local_host)
self.remote_client = ollama.Client(host=remote_host)
self.local_model = local_model
self.remote_model = remote_model
# Judge model configuration
self.judge_model = judge_model or local_model
self.judge_host = judge_host or local_host
self.judge_client = ollama.Client(host=self.judge_host)
def _classify_with_llm(
self, prompt: str, verbose: bool = False
) -> tuple[bool, str]:
"""
Use LLM to classify the prompt.
Returns:
Tuple of (use_remote, reason)
"""
try:
classification_prompt = self.CLASSIFICATION_PROMPT.format(
prompt=prompt[:500] # Limit prompt length for classification
)
response = self.judge_client.chat(
model=self.judge_model,
messages=[{"role": "user", "content": classification_prompt}],
options={"temperature": 0, "num_predict": 10},
)
result = response["message"]["content"].strip().upper()
if verbose:
print(f"[Router] LLM Judge: {result}")
if "COMPLEX" in result:
return True, "Complex task"
else:
return False, "Simple task"
except Exception as e:
if verbose:
print(f"[Router] LLM Judge failed: {e}, defaulting to local")
return False, "Judge failed, defaulting to local"
def should_use_remote(self, prompt: str, verbose: bool = False) -> tuple[bool, str]:
"""
Determine if the prompt should be routed to remote SGLang.
Args:
prompt: User's input prompt
verbose: Print debug information
Returns:
Tuple of (should_use_remote, reason)
"""
return self._classify_with_llm(prompt, verbose)
def chat(
self,
prompt: str,
messages: Optional[list] = None,
verbose: bool = False,
force_local: bool = False,
force_remote: bool = False,
) -> dict:
"""
Route the request and get response.
Args:
prompt: User's input (used if messages is None)
messages: Full message history (overrides prompt if provided)
verbose: Print routing decision
force_local: Force use of local model
force_remote: Force use of remote model
Returns:
Response dict with 'content', 'model', 'location', 'reason' keys
"""
# Build messages
if messages is None:
messages = [{"role": "user", "content": prompt}]
check_prompt = prompt
else:
# Use the last user message for routing decision
check_prompt = ""
for msg in reversed(messages):
if msg.get("role") == "user":
check_prompt = msg.get("content", "")
break
# Determine routing
if force_remote:
use_remote, reason = True, "Forced remote"
elif force_local:
use_remote, reason = False, "Forced local"
else:
use_remote, reason = self.should_use_remote(check_prompt, verbose)
if use_remote:
client = self.remote_client
model = self.remote_model
location = "Remote SGLang"
else:
client = self.local_client
model = self.local_model
location = "Local Ollama"
if verbose:
print(f"[Router] -> {location} | Model: {model}")
try:
response = client.chat(model=model, messages=messages)
return {
"content": response["message"]["content"],
"model": model,
"location": location,
"reason": reason,
}
except Exception as e:
# Fallback to the other option
if verbose:
print(f"[Router] {location} failed: {e}, falling back...")
fallback_client = (
self.remote_client if not use_remote else self.local_client
)
fallback_model = self.remote_model if not use_remote else self.local_model
fallback_location = "Remote SGLang" if not use_remote else "Local Ollama"
response = fallback_client.chat(model=fallback_model, messages=messages)
return {
"content": response["message"]["content"],
"model": fallback_model,
"location": fallback_location,
"reason": f"Fallback from {location}",
}
def chat_stream(
self,
prompt: str,
messages: Optional[list] = None,
verbose: bool = False,
force_local: bool = False,
force_remote: bool = False,
):
"""
Route the request and stream response.
Yields:
Response chunks
"""
if messages is None:
messages = [{"role": "user", "content": prompt}]
check_prompt = prompt
else:
check_prompt = ""
for msg in reversed(messages):
if msg.get("role") == "user":
check_prompt = msg.get("content", "")
break
if force_remote:
use_remote, reason = True, "Forced remote"
elif force_local:
use_remote, reason = False, "Forced local"
else:
use_remote, reason = self.should_use_remote(check_prompt, verbose)
if use_remote:
client = self.remote_client
model = self.remote_model
location = "Remote SGLang"
else:
client = self.local_client
model = self.local_model
location = "Local Ollama"
if verbose:
print(f"[Router] -> {location} | Model: {model}")
for chunk in client.chat(model=model, messages=messages, stream=True):
yield chunk
def main():
"""Interactive demo of the smart router."""
print("=" * 60)
print("Smart Router: Local Ollama <-> Remote SGLang")
print("=" * 60)
print("\nRouting strategy:")
print(" LLM Judge classifies each request as SIMPLE or COMPLEX")
print(" - SIMPLE tasks -> Local Ollama (fast)")
print(" - COMPLEX tasks -> Remote SGLang (powerful)")
print("\nType 'quit' to exit\n")
router = SmartRouter(
local_host="http://localhost:11434",
remote_host="http://localhost:30001",
local_model="llama3.2",
remote_model="Qwen/Qwen2.5-1.5B-Instruct",
)
messages = []
while True:
try:
user_input = input("You: ").strip()
if user_input.lower() in ["quit", "exit", "q"]:
print("Goodbye!")
break
if not user_input:
continue
messages.append({"role": "user", "content": user_input})
# Use streaming for real-time output
print("\nAssistant: ", end="", flush=True)
full_response = ""
for chunk in router.chat_stream(
prompt=user_input, messages=messages, verbose=True
):
content = chunk.get("message", {}).get("content", "")
if content:
print(content, end="", flush=True)
full_response += content
print("\n")
messages.append({"role": "assistant", "content": full_response})
except KeyboardInterrupt:
print("\nGoodbye!")
break
except Exception as e:
print(f"Error: {e}\n")
if __name__ == "__main__":
main()