diff --git a/docs/basic_usage/ollama_api.md b/docs/basic_usage/ollama_api.md new file mode 100644 index 000000000..c7b302c96 --- /dev/null +++ b/docs/basic_usage/ollama_api.md @@ -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 | diff --git a/docs/index.rst b/docs/index.rst index 59f15323f..bdb4a34ea 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -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 diff --git a/python/sglang/srt/entrypoints/http_server.py b/python/sglang/srt/entrypoints/http_server.py index 1be55d5ba..6eaf578a6 100644 --- a/python/sglang/srt/entrypoints/http_server.py +++ b/python/sglang/srt/entrypoints/http_server.py @@ -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: diff --git a/python/sglang/srt/entrypoints/ollama/README.md b/python/sglang/srt/entrypoints/ollama/README.md new file mode 100644 index 000000000..a3f87460a --- /dev/null +++ b/python/sglang/srt/entrypoints/ollama/README.md @@ -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 # 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 --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 diff --git a/python/sglang/srt/entrypoints/ollama/__init__.py b/python/sglang/srt/entrypoints/ollama/__init__.py new file mode 100644 index 000000000..cf572ac2f --- /dev/null +++ b/python/sglang/srt/entrypoints/ollama/__init__.py @@ -0,0 +1 @@ +# Ollama-compatible API for SGLang diff --git a/python/sglang/srt/entrypoints/ollama/protocol.py b/python/sglang/srt/entrypoints/ollama/protocol.py new file mode 100644 index 000000000..f7e18e2eb --- /dev/null +++ b/python/sglang/srt/entrypoints/ollama/protocol.py @@ -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) diff --git a/python/sglang/srt/entrypoints/ollama/serving.py b/python/sglang/srt/entrypoints/ollama/serving.py new file mode 100644 index 000000000..786ff43eb --- /dev/null +++ b/python/sglang/srt/entrypoints/ollama/serving.py @@ -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"], + ) diff --git a/python/sglang/srt/entrypoints/ollama/smart_router.py b/python/sglang/srt/entrypoints/ollama/smart_router.py new file mode 100644 index 000000000..de6fdddee --- /dev/null +++ b/python/sglang/srt/entrypoints/ollama/smart_router.py @@ -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()