chore: import upstream snapshot with attribution

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"],
+        )