diff --git a/docs/examples/openai_compat_server/README.md b/docs/examples/openai_compat_server/README.md new file mode 100644 index 0000000000..a06e388c93 --- /dev/null +++ b/docs/examples/openai_compat_server/README.md @@ -0,0 +1,137 @@ +# OpenAI-Compatible Server Examples + +This directory contains examples for using mellea's OpenAI-compatible HTTP server. + +## Overview + +The OpenAI-compatible server allows you to wrap any mellea backend and expose it through OpenAI's API format. This means you can use the official OpenAI Python SDK or any other OpenAI-compatible client to interact with local models, Ollama, or other backends supported by mellea. + +## Files + +- **`basic_server.py`**: Shows how to start the server with basic configuration +- **`client_example.py`**: Demonstrates using the OpenAI Python SDK to interact with the server +- **`curl_examples.sh`**: Command-line examples using curl + +## Quick Start + +### 1. Start the Server + +```python +from mellea.integrations.openai_compat.server import ServerConfig, run_server + +config = ServerConfig( + backend_name="ollama", + model_id="granite4:micro", +) + +run_server(host="0.0.0.0", port=8000, config=config) +``` + +Or run the example: + +```bash +uv run python docs/examples/openai_compat_server/basic_server.py +``` + +### 2. Use the OpenAI SDK + +```python +import openai + +client = openai.OpenAI( + base_url="http://localhost:8000/v1", + api_key="not-needed", +) + +response = client.chat.completions.create( + model="granite4:micro", + messages=[ + {"role": "user", "content": "Hello!"} + ], +) + +print(response.choices[0].message.content) +``` + +### 3. Or Use curl + +```bash +curl http://localhost:8000/v1/chat/completions \ + -H "Content-Type: application/json" \ + -d '{ + "model": "granite4:micro", + "messages": [{"role": "user", "content": "Hello!"}] + }' +``` + +## Supported Endpoints + +- **`POST /v1/chat/completions`**: Chat completions (with streaming support) +- **`GET /v1/models`**: List available models +- **`GET /health`**: Health check + +## Interactive API Documentation + +FastAPI automatically provides interactive API documentation when the server is running: + +- **Swagger UI**: http://localhost:8000/docs +- **ReDoc**: http://localhost:8000/redoc +- **OpenAPI JSON**: http://localhost:8000/openapi.json + +Visit the Swagger UI to explore all endpoints, view request/response schemas, and test the API directly from your browser. + +## Configuration Options + +The `ServerConfig` class accepts the following parameters: + +- **`backend_name`**: Name of the mellea backend (e.g., "ollama", "openai", "hf") +- **`model_id`**: Model identifier to use +- **`base_url`**: Base URL for the backend API (if applicable) +- **`api_key`**: API key for the backend (if applicable) +- **`default_model_options`**: Default model options for all requests +- **`**backend_kwargs`**: Additional arguments passed to the backend + +## Supported Features + +✅ Chat completions +✅ Streaming responses +✅ System messages +✅ Temperature, top_p, max_tokens +✅ Seed for reproducibility +✅ Tool calling (when supported by backend) +✅ Token usage tracking + +## Backend Support + +The server works with any mellea backend: + +- **Ollama**: Local models via Ollama +- **OpenAI**: OpenAI API models +- **HuggingFace**: Local HuggingFace models +- **Watsonx**: IBM Watsonx models +- **LiteLLM**: Any LiteLLM-supported provider + +## Deployment + +The server can be deployed as a standalone service: + +```bash +# Using uvicorn directly +uvicorn mellea.integrations.openai_compat.server:app --host 0.0.0.0 --port 8000 + +# Or using the run_server function +python -c "from mellea.integrations.openai_compat.server import run_server; run_server()" +``` + +## Testing + +Run the tests: + +```bash +uv run pytest test/integrations/test_openai_compat_server.py +``` + +Run the client example (requires server to be running): + +```bash +uv run python docs/examples/openai_compat_server/client_example.py \ No newline at end of file diff --git a/docs/examples/openai_compat_server/basic_server.py b/docs/examples/openai_compat_server/basic_server.py new file mode 100644 index 0000000000..72d9d94420 --- /dev/null +++ b/docs/examples/openai_compat_server/basic_server.py @@ -0,0 +1,61 @@ +# pytest: skip_always +"""Basic example of running the OpenAI-compatible server with mellea. + +This example demonstrates how to start a server that wraps a mellea backend +and exposes it via OpenAI's API format. +""" + +import socket + +from mellea.integrations.openai_compat.server import ServerConfig, run_server + + +def find_unused_port(start_port: int = 8000, max_attempts: int = 100) -> int: + """Find an unused port starting from start_port. + + Args: + start_port: Port to start searching from. + max_attempts: Maximum number of ports to try. + + Returns: + An available port number. + + Raises: + RuntimeError: If no available port is found. + """ + for port in range(start_port, start_port + max_attempts): + try: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + s.bind(("", port)) + return port + except OSError: + continue + raise RuntimeError( + f"Could not find an unused port in range {start_port}-{start_port + max_attempts}" + ) + + +def main(): + """Run the OpenAI-compatible server with Ollama backend.""" + # Configure the server to use Ollama with granite4:micro + config = ServerConfig( + backend_name="ollama", + model_id="granite4:micro", + default_model_options={"temperature": 0.7}, + ) + + # Find an available port starting from 8000 + port = find_unused_port(start_port=8000) + + # Start the server + print(f"Starting OpenAI-compatible server on http://localhost:{port}") + print(f"API endpoint: http://localhost:{port}/v1/chat/completions") + print(f"Models endpoint: http://localhost:{port}/v1/models") + print(f"Health check: http://localhost:{port}/health") + print(f"API docs: http://localhost:{port}/docs") + + run_server(host="0.0.0.0", port=port, config=config) + + +if __name__ == "__main__": + main() diff --git a/docs/examples/openai_compat_server/client_example.py b/docs/examples/openai_compat_server/client_example.py new file mode 100644 index 0000000000..4113411232 --- /dev/null +++ b/docs/examples/openai_compat_server/client_example.py @@ -0,0 +1,68 @@ +# pytest: skip_always +"""Example of using the OpenAI Python SDK with the mellea server. + +This example shows how to use the official OpenAI Python SDK to interact +with a mellea backend through the OpenAI-compatible server. +""" + +import openai + + +def test_openai_sdk_compatibility(): + """Test that the OpenAI SDK works with mellea server.""" + # Configure the OpenAI client to point to our mellea server + client = openai.OpenAI( + base_url="http://localhost:8000/v1", + api_key="not-needed", # mellea server doesn't require auth by default + ) + + # Make a simple chat completion request + response = client.chat.completions.create( + model="granite4:micro", + messages=[ + {"role": "system", "content": "You are a helpful assistant."}, + {"role": "user", "content": "What is the capital of France?"}, + ], + temperature=0.7, + max_tokens=50, + ) + + # Access the response + print(f"Response: {response.choices[0].message.content}") + assert response.choices[0].message.content is not None + assert len(response.choices[0].message.content) > 0 + + +def test_streaming_with_openai_sdk(): + """Test streaming responses with the OpenAI SDK.""" + client = openai.OpenAI(base_url="http://localhost:8000/v1", api_key="not-needed") + + # Make a streaming request + stream = client.chat.completions.create( + model="granite4:micro", + messages=[{"role": "user", "content": "Count from 1 to 5"}], + stream=True, + max_tokens=50, + ) + + # Collect streamed chunks + full_response = "" + for chunk in stream: + if chunk.choices[0].delta.content is not None: + content = chunk.choices[0].delta.content + full_response += content + print(content, end="", flush=True) + + print() # New line after streaming + assert len(full_response) > 0 + + +if __name__ == "__main__": + print("Testing OpenAI SDK compatibility...") + print("\n1. Basic completion:") + test_openai_sdk_compatibility() + + print("\n2. Streaming completion:") + test_streaming_with_openai_sdk() + + print("\nAll tests passed!") diff --git a/docs/examples/openai_compat_server/curl_examples.sh b/docs/examples/openai_compat_server/curl_examples.sh new file mode 100755 index 0000000000..4dd277ad34 --- /dev/null +++ b/docs/examples/openai_compat_server/curl_examples.sh @@ -0,0 +1,93 @@ +#!/bin/bash +# Example curl commands for the OpenAI-compatible server + +# Base URL for the server +BASE_URL="http://localhost:8000" + +echo "=== OpenAI-Compatible Server Examples ===" +echo "" + +# 1. Health check +echo "1. Health Check:" +curl -s "${BASE_URL}/health" | jq . +echo "" + +# 2. List models +echo "2. List Models:" +curl -s "${BASE_URL}/v1/models" | jq . +echo "" + +# 3. Basic chat completion +echo "3. Basic Chat Completion:" +curl -s "${BASE_URL}/v1/chat/completions" \ + -H "Content-Type: application/json" \ + -d '{ + "model": "granite4:micro", + "messages": [ + {"role": "user", "content": "What is 2+2?"} + ], + "max_tokens": 50 + }' | jq . +echo "" + +# 4. Chat completion with system message +echo "4. Chat Completion with System Message:" +curl -s "${BASE_URL}/v1/chat/completions" \ + -H "Content-Type: application/json" \ + -d '{ + "model": "granite4:micro", + "messages": [ + {"role": "system", "content": "You are a helpful math tutor."}, + {"role": "user", "content": "Explain what a prime number is."} + ], + "temperature": 0.7, + "max_tokens": 100 + }' | jq . +echo "" + +# 5. Streaming chat completion +echo "5. Streaming Chat Completion:" +curl -s "${BASE_URL}/v1/chat/completions" \ + -H "Content-Type: application/json" \ + -d '{ + "model": "granite4:micro", + "messages": [ + {"role": "user", "content": "Count from 1 to 5"} + ], + "stream": true, + "max_tokens": 50 + }' +echo "" + +# 6. Chat completion with seed for reproducibility +echo "6. Chat Completion with Seed:" +curl -s "${BASE_URL}/v1/chat/completions" \ + -H "Content-Type: application/json" \ + -d '{ + "model": "granite4:micro", + "messages": [ + {"role": "user", "content": "Pick a random number"} + ], + "seed": 42, + "max_tokens": 20 + }' | jq . +echo "" + +# 7. Chat completion with custom parameters +echo "7. Chat Completion with Custom Parameters:" +curl -s "${BASE_URL}/v1/chat/completions" \ + -H "Content-Type: application/json" \ + -d '{ + "model": "granite4:micro", + "messages": [ + {"role": "user", "content": "Write a haiku about coding"} + ], + "temperature": 0.9, + "top_p": 0.95, + "max_tokens": 100, + "presence_penalty": 0.5, + "frequency_penalty": 0.5 + }' | jq . +echo "" + +echo "=== All examples completed ===" diff --git a/mellea/integrations/openai_compat/__init__.py b/mellea/integrations/openai_compat/__init__.py new file mode 100644 index 0000000000..54e9263814 --- /dev/null +++ b/mellea/integrations/openai_compat/__init__.py @@ -0,0 +1,5 @@ +"""OpenAI API-compatible server for mellea.""" + +from .server import create_app, run_server + +__all__ = ["create_app", "run_server"] diff --git a/mellea/integrations/openai_compat/server.py b/mellea/integrations/openai_compat/server.py new file mode 100644 index 0000000000..b766b109f7 --- /dev/null +++ b/mellea/integrations/openai_compat/server.py @@ -0,0 +1,498 @@ +"""OpenAI API-compatible HTTP server that wraps mellea backends. + +This module provides a FastAPI server implementing OpenAI's chat completions API, +allowing any OpenAI-compatible client to use mellea backends. +""" + +import json +import time +import uuid +from collections.abc import AsyncIterator +from contextlib import asynccontextmanager +from typing import Any, Literal + +from fastapi import FastAPI, HTTPException +from fastapi.responses import StreamingResponse +from pydantic import BaseModel, Field + +from mellea.core import FancyLogger, ModelOutputThunk +from mellea.stdlib.components import Message +from mellea.stdlib.context import ChatContext +from mellea.stdlib.session import MelleaSession, start_session + + +class ChatMessage(BaseModel): + """A single message in a chat conversation.""" + + role: Literal["system", "user", "assistant", "tool"] + content: str | None = None + name: str | None = None + tool_call_id: str | None = None + + +class FunctionDefinition(BaseModel): + """Definition of a function that can be called by the model.""" + + name: str + description: str | None = None + parameters: dict[str, Any] = Field(default_factory=dict) + + +class ToolDefinition(BaseModel): + """Definition of a tool that can be called by the model.""" + + type: Literal["function"] = "function" + function: FunctionDefinition + + +class ChatCompletionRequest(BaseModel): + """Request body for chat completions endpoint.""" + + model: str + messages: list[ChatMessage] + temperature: float | None = Field(default=1.0, ge=0, le=2) + top_p: float | None = Field(default=1.0, ge=0, le=1) + n: int | None = Field(default=1, ge=1) + stream: bool = False + stop: str | list[str] | None = None + max_tokens: int | None = None + presence_penalty: float | None = Field(default=0, ge=-2, le=2) + frequency_penalty: float | None = Field(default=0, ge=-2, le=2) + logit_bias: dict[str, float] | None = None + user: str | None = None + seed: int | None = None + tools: list[ToolDefinition] | None = None + tool_choice: Literal["none", "auto"] | dict[str, Any] | None = None + + model_config = {"extra": "allow"} + + +class ChatCompletionMessage(BaseModel): + """A message in a chat completion response.""" + + role: Literal["assistant"] + content: str | None = None + tool_calls: list[dict[str, Any]] | None = None + + +class Choice(BaseModel): + """A single completion choice.""" + + index: int + message: ChatCompletionMessage + finish_reason: str | None = None + + +class Usage(BaseModel): + """Token usage information.""" + + prompt_tokens: int + completion_tokens: int + total_tokens: int + + +class ChatCompletion(BaseModel): + """Response from chat completions endpoint.""" + + id: str + object: Literal["chat.completion"] = "chat.completion" + created: int + model: str + choices: list[Choice] + usage: Usage | None = None + + +class ChatCompletionChunk(BaseModel): + """A chunk in a streaming chat completion response.""" + + id: str + object: Literal["chat.completion.chunk"] = "chat.completion.chunk" + created: int + model: str + choices: list[dict[str, Any]] + + +class Model(BaseModel): + """Information about an available model.""" + + id: str + object: Literal["model"] = "model" + created: int + owned_by: str + + +class ModelList(BaseModel): + """List of available models.""" + + object: Literal["list"] = "list" + data: list[Model] + + +class ServerConfig: + """Configuration for the OpenAI-compatible server. + + Args: + backend_name: Name of the mellea backend to use (e.g., "ollama", "openai"). + model_id: Model identifier to use with the backend. + base_url: Base URL for the backend API (if applicable). + api_key: API key for the backend (if applicable). + default_model_options: Default model options to apply to all requests. + backend_kwargs: Additional keyword arguments to pass to the backend. + """ + + def __init__( + self, + backend_name: str = "ollama", + model_id: str | None = None, + base_url: str | None = None, + api_key: str | None = None, + default_model_options: dict[str, Any] | None = None, + **backend_kwargs: Any, + ): + """Initialize server configuration.""" + self.backend_name = backend_name + self.model_id = model_id + self.base_url = base_url + self.api_key = api_key + self.default_model_options = default_model_options or {} + self.backend_kwargs = backend_kwargs + + +def create_app(config: ServerConfig | None = None) -> FastAPI: + """Create a FastAPI application with OpenAI-compatible endpoints. + + Args: + config: Server configuration. If None, uses default configuration. + + Returns: + FastAPI application instance. + """ + if config is None: + config = ServerConfig() + + app = FastAPI( + title="Mellea OpenAI-Compatible API", + description="OpenAI API-compatible server powered by mellea", + version="0.1.0", + ) + + # Initialize app state immediately (not in lifespan) for TestClient compatibility + app.state.config = config + app.state.sessions = {} + + @asynccontextmanager + async def lifespan(app: FastAPI): + """Manage application lifecycle.""" + yield + # Cleanup sessions on shutdown + for session in app.state.sessions.values(): + session.cleanup() + + # Set lifespan after state initialization + app.router.lifespan_context = lifespan + + def get_or_create_session(model: str) -> MelleaSession: + """Get or create a mellea session for the given model. + + Args: + model: Model identifier from the request. + + Returns: + MelleaSession instance. + """ + config = app.state.config + session_key = f"{config.backend_name}:{model}" + + if session_key not in app.state.sessions: + # Build backend kwargs + backend_kwargs = dict(config.backend_kwargs) + if config.base_url: + backend_kwargs["base_url"] = config.base_url + if config.api_key: + backend_kwargs["api_key"] = config.api_key + + # Start a new session + # Use the client-requested model, falling back to config default only if not provided + session = start_session( + backend_name=config.backend_name, + model_id=model or config.model_id, + model_options=config.default_model_options, + **backend_kwargs, + ) + app.state.sessions[session_key] = session + + return app.state.sessions[session_key] + + def build_model_options(request: ChatCompletionRequest) -> dict[str, Any]: + """Build model options from request parameters. + + Args: + request: Chat completion request. + + Returns: + Dictionary of model options. + """ + model_opts: dict[str, Any] = {} + + if request.temperature is not None: + model_opts["temperature"] = request.temperature + if request.top_p is not None: + model_opts["top_p"] = request.top_p + if request.max_tokens is not None: + model_opts["max_tokens"] = request.max_tokens + if request.seed is not None: + model_opts["seed"] = request.seed + if request.stop is not None: + model_opts["stop"] = request.stop + if request.presence_penalty is not None: + model_opts["presence_penalty"] = request.presence_penalty + if request.frequency_penalty is not None: + model_opts["frequency_penalty"] = request.frequency_penalty + if request.stream: + model_opts["stream"] = True + + return model_opts + + async def stream_response( + mot: ModelOutputThunk, completion_id: str, model: str, created: int + ) -> AsyncIterator[str]: + """Stream chat completion chunks. + + Args: + mot: Model output thunk to stream from. + completion_id: Unique completion ID. + model: Model identifier. + created: Creation timestamp. + + Yields: + Server-sent event formatted strings. + """ + try: + prev_length = 0 + while not mot.is_computed(): + # Get the next chunk + current_value = await mot.astream() + + # Extract only the new content since last iteration + new_content = current_value[prev_length:] if current_value else "" + prev_length = len(current_value) if current_value else 0 + + if new_content: + chunk_data = ChatCompletionChunk( + id=completion_id, + created=created, + model=model, + choices=[ + { + "index": 0, + "delta": {"role": "assistant", "content": new_content}, + "finish_reason": None, + } + ], + ) + yield f"data: {chunk_data.model_dump_json()}\n\n" + + # Send final chunk with finish_reason + final_chunk = ChatCompletionChunk( + id=completion_id, + created=created, + model=model, + choices=[{"index": 0, "delta": {}, "finish_reason": "stop"}], + ) + yield f"data: {final_chunk.model_dump_json()}\n\n" + yield "data: [DONE]\n\n" + + except Exception as e: + FancyLogger.get_logger().error(f"Error during streaming: {e}") + error_chunk = {"error": {"message": str(e), "type": "server_error"}} + yield f"data: {json.dumps(error_chunk)}\n\n" + + @app.post("/v1/chat/completions", response_model=None) + async def chat_completions(request: ChatCompletionRequest): + """Handle chat completion requests. + + Args: + request: Chat completion request body. + + Returns: + Chat completion response or streaming response. + + Raises: + HTTPException: If the request is invalid or processing fails. + """ + try: + # Get or create session + session = get_or_create_session(request.model) + + # Build model options + model_opts = build_model_options(request) + + # Extract the last user message content + last_message = request.messages[-1] + if not last_message.content: + raise HTTPException( + status_code=400, detail="Last message must have content" + ) + + # Determine if we should use tool calls + tool_calls = request.tools is not None and len(request.tools) > 0 + + # Build conversation history - add all messages to session context + for msg in request.messages[:-1]: + if msg.content: + session.ctx = session.ctx.add( + Message(role=msg.role, content=msg.content) + ) + + # Create the user message for the current request + user_message = Message(role=last_message.role, content=last_message.content) + + # Generate completion using the backend + # We use backend directly to get access to the ModelOutputThunk for streaming + mot, new_ctx = await session.backend._generate_from_context( + action=user_message, + ctx=session.ctx, + model_options=model_opts, + tool_calls=tool_calls, + ) + + # Update session context with the new context + session.ctx = new_ctx + + # Generate unique IDs + completion_id = f"chatcmpl-{uuid.uuid4().hex[:29]}" + created = int(time.time()) + + # Handle streaming + if request.stream: + return StreamingResponse( + stream_response(mot, completion_id, request.model, created), + media_type="text/event-stream", + ) + + # Wait for completion + await mot.avalue() + + # Extract usage information + usage = None + if mot.usage: + usage = Usage( + prompt_tokens=mot.usage.get("prompt_tokens", 0), + completion_tokens=mot.usage.get("completion_tokens", 0), + total_tokens=mot.usage.get("total_tokens", 0), + ) + + # Build response + response = ChatCompletion( + id=completion_id, + created=created, + model=request.model, + choices=[ + Choice( + index=0, + message=ChatCompletionMessage( + role="assistant", + content=mot.value, + tool_calls=( + [ + { + "id": f"call_{uuid.uuid4().hex[:24]}", + "type": "function", + "function": { + "name": name, + "arguments": json.dumps(call.args), + }, + } + for name, call in mot.tool_calls.items() + ] + if mot.tool_calls + else None + ), + ), + finish_reason="stop", + ) + ], + usage=usage, + ) + + return response + + except Exception as e: + FancyLogger.get_logger().error(f"Error processing request: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + @app.get("/v1/models") + async def list_models() -> ModelList: + """List available models. + + Returns: + List of available models. + """ + config = app.state.config + # Return the configured model + model_id = config.model_id or "default" + return ModelList( + data=[ + Model( + id=model_id, created=int(time.time()), owned_by=config.backend_name + ) + ] + ) + + @app.get("/health") + async def health() -> dict[str, str]: + """Health check endpoint. + + Returns: + Health status. + """ + return {"status": "healthy"} + + return app + + +def run_server( + host: str = "0.0.0.0", + port: int = 8000, + config: ServerConfig | None = None, + **uvicorn_kwargs: Any, +) -> None: + """Run the OpenAI-compatible server. + + Args: + host: Host to bind to. + port: Port to bind to. + config: Server configuration. + uvicorn_kwargs: Additional keyword arguments for uvicorn. + """ + try: + import uvicorn + except ImportError as e: + raise ImportError( + "uvicorn is required to run the server. " + "Please install it with: pip install mellea[server]" + ) from e + + app = create_app(config) + uvicorn.run(app, host=host, port=port, **uvicorn_kwargs) + + +if __name__ == "__main__": + """Run the server with default configuration when executed directly.""" + print("Starting mellea OpenAI-compatible server with default configuration...") + print("Backend: Ollama") + print("Model: granite4:micro") + print("") + print("Server will be available at:") + print(" - API: http://localhost:8000/v1/chat/completions") + print(" - Docs: http://localhost:8000/docs") + print(" - Health: http://localhost:8000/health") + print("") + print("Note: Requires Ollama to be running with granite4:micro model available") + print(" Install with: ollama pull granite4:micro") + print("") + + # Use default configuration with Ollama + default_config = ServerConfig(backend_name="ollama", model_id="granite4:micro") + + run_server(host="0.0.0.0", port=8000, config=default_config) diff --git a/test/integrations/test_openai_compat_server.py b/test/integrations/test_openai_compat_server.py new file mode 100644 index 0000000000..32d7b082fe --- /dev/null +++ b/test/integrations/test_openai_compat_server.py @@ -0,0 +1,633 @@ +"""Tests for the OpenAI-compatible server.""" + +import json + +import pytest +from fastapi.testclient import TestClient + +from mellea.integrations.openai_compat.server import ServerConfig, create_app + + +@pytest.mark.ollama +@pytest.mark.llm +def test_create_app(): + """Test that the app can be created.""" + config = ServerConfig(backend_name="ollama", model_id="granite4:micro") + app = create_app(config) + assert app is not None + + +@pytest.mark.ollama +@pytest.mark.llm +def test_health_endpoint(): + """Test the health check endpoint.""" + config = ServerConfig(backend_name="ollama", model_id="granite4:micro") + app = create_app(config) + client = TestClient(app) + + response = client.get("/health") + assert response.status_code == 200 + assert response.json() == {"status": "healthy"} + + +@pytest.mark.ollama +@pytest.mark.llm +def test_list_models(): + """Test the models listing endpoint.""" + config = ServerConfig(backend_name="ollama", model_id="granite4:micro") + app = create_app(config) + client = TestClient(app) + + response = client.get("/v1/models") + assert response.status_code == 200 + data = response.json() + assert data["object"] == "list" + assert len(data["data"]) > 0 + assert data["data"][0]["id"] == "granite4:micro" + + +@pytest.mark.ollama +@pytest.mark.llm +@pytest.mark.qualitative +def test_chat_completions_basic(): + """Test basic chat completion without streaming.""" + config = ServerConfig(backend_name="ollama", model_id="granite4:micro") + app = create_app(config) + client = TestClient(app) + + request_data = { + "model": "granite4:micro", + "messages": [{"role": "user", "content": "Say hello in one word"}], + "temperature": 0.7, + "max_tokens": 10, + } + + response = client.post("/v1/chat/completions", json=request_data) + assert response.status_code == 200 + + data = response.json() + assert data["object"] == "chat.completion" + assert data["model"] == "granite4:micro" + assert len(data["choices"]) == 1 + assert data["choices"][0]["message"]["role"] == "assistant" + assert data["choices"][0]["message"]["content"] is not None + assert len(data["choices"][0]["message"]["content"]) > 0 + + +@pytest.mark.ollama +@pytest.mark.llm +@pytest.mark.qualitative +def test_chat_completions_streaming(): + """Test streaming chat completion.""" + config = ServerConfig(backend_name="ollama", model_id="granite4:micro") + app = create_app(config) + client = TestClient(app) + + request_data = { + "model": "granite4:micro", + "messages": [{"role": "user", "content": "Count to 3"}], + "stream": True, + "max_tokens": 20, + } + + with client.stream("POST", "/v1/chat/completions", json=request_data) as response: + assert response.status_code == 200 + assert response.headers["content-type"] == "text/event-stream; charset=utf-8" + + chunks = [] + for line in response.iter_lines(): + if line.startswith("data: "): + data_str = line[6:] # Remove "data: " prefix + if data_str == "[DONE]": + break + chunk_data = json.loads(data_str) + chunks.append(chunk_data) + + # Verify we got multiple chunks + assert len(chunks) > 0 + # Verify chunk structure + assert chunks[0]["object"] == "chat.completion.chunk" + assert chunks[0]["model"] == "granite4:micro" + + +@pytest.mark.ollama +@pytest.mark.llm +@pytest.mark.qualitative +def test_chat_completions_with_system_message(): + """Test chat completion with system message.""" + config = ServerConfig(backend_name="ollama", model_id="granite4:micro") + app = create_app(config) + client = TestClient(app) + + request_data = { + "model": "granite4:micro", + "messages": [ + {"role": "system", "content": "You are a helpful assistant."}, + {"role": "user", "content": "What is 2+2?"}, + ], + "max_tokens": 20, + } + + response = client.post("/v1/chat/completions", json=request_data) + assert response.status_code == 200 + + data = response.json() + assert data["choices"][0]["message"]["content"] is not None + + +@pytest.mark.ollama +@pytest.mark.llm +def test_chat_completions_with_seed(): + """Test that seed parameter is accepted.""" + config = ServerConfig(backend_name="ollama", model_id="granite4:micro") + app = create_app(config) + client = TestClient(app) + + request_data = { + "model": "granite4:micro", + "messages": [{"role": "user", "content": "Hello"}], + "seed": 42, + "max_tokens": 10, + } + + response = client.post("/v1/chat/completions", json=request_data) + assert response.status_code == 200 + + +def test_server_config_defaults(): + """Test ServerConfig default values.""" + config = ServerConfig() + assert config.backend_name == "ollama" + assert config.model_id is None + assert config.base_url is None + assert config.api_key is None + assert config.default_model_options == {} + + +def test_server_config_custom(): + """Test ServerConfig with custom values.""" + config = ServerConfig( + backend_name="openai", + model_id="gpt-4", + base_url="https://api.openai.com/v1", + api_key="test-key", + default_model_options={"temperature": 0.5}, + timeout=30, + ) + assert config.backend_name == "openai" + assert config.model_id == "gpt-4" + assert config.base_url == "https://api.openai.com/v1" + assert config.api_key == "test-key" + assert config.default_model_options == {"temperature": 0.5} + assert config.backend_kwargs == {"timeout": 30} + + +@pytest.mark.ollama +@pytest.mark.llm +@pytest.mark.qualitative +def test_chat_completions_with_tools(): + """Test chat completion with tool calling.""" + config = ServerConfig(backend_name="ollama", model_id="granite4:micro") + app = create_app(config) + client = TestClient(app) + + # Define a simple tool + tool_definition = { + "type": "function", + "function": { + "name": "get_temperature", + "description": "Returns today's temperature of the given city in Celsius.", + "parameters": { + "type": "object", + "properties": { + "location": {"type": "string", "description": "A city name"} + }, + "required": ["location"], + }, + }, + } + + request_data = { + "model": "granite4:micro", + "messages": [ + { + "role": "user", + "content": "What is today's temperature in Boston? Use the get_temperature tool.", + } + ], + "tools": [tool_definition], + "tool_choice": "auto", + "max_tokens": 100, + } + + response = client.post("/v1/chat/completions", json=request_data) + assert response.status_code == 200 + + data = response.json() + assert data["object"] == "chat.completion" + assert len(data["choices"]) == 1 + + message = data["choices"][0]["message"] + assert message["role"] == "assistant" + + # Check if tool was called + if message.get("tool_calls"): + tool_call = message["tool_calls"][0] + assert tool_call["function"]["name"] == "get_temperature" + # Parse arguments + args = json.loads(tool_call["function"]["arguments"]) + assert "location" in args + assert "boston" in args["location"].lower() + + +@pytest.mark.ollama +@pytest.mark.llm +def test_chat_completions_empty_message(): + """Test that empty message content is rejected.""" + config = ServerConfig(backend_name="ollama", model_id="granite4:micro") + app = create_app(config) + client = TestClient(app) + + request_data = { + "model": "granite4:micro", + "messages": [{"role": "user", "content": ""}], + } + + response = client.post("/v1/chat/completions", json=request_data) + # Server returns 500 with HTTPException detail about empty content + assert response.status_code == 500 + assert "Last message must have content" in response.json()["detail"] + + +@pytest.mark.ollama +@pytest.mark.llm +def test_chat_completions_with_stop_sequences(): + """Test chat completion with stop sequences.""" + config = ServerConfig(backend_name="ollama", model_id="granite4:micro") + app = create_app(config) + client = TestClient(app) + + request_data = { + "model": "granite4:micro", + "messages": [{"role": "user", "content": "Count: 1, 2, 3"}], + "stop": [",", "."], + "max_tokens": 20, + } + + response = client.post("/v1/chat/completions", json=request_data) + assert response.status_code == 200 + + +@pytest.mark.ollama +@pytest.mark.llm +def test_chat_completions_with_penalties(): + """Test chat completion with presence and frequency penalties.""" + config = ServerConfig(backend_name="ollama", model_id="granite4:micro") + app = create_app(config) + client = TestClient(app) + + request_data = { + "model": "granite4:micro", + "messages": [{"role": "user", "content": "Hello"}], + "presence_penalty": 0.5, + "frequency_penalty": 0.3, + "max_tokens": 10, + } + + response = client.post("/v1/chat/completions", json=request_data) + assert response.status_code == 200 + + +@pytest.mark.ollama +@pytest.mark.llm +@pytest.mark.qualitative +def test_chat_completions_multi_turn(): + """Test multi-turn conversation.""" + config = ServerConfig(backend_name="ollama", model_id="granite4:micro") + app = create_app(config) + client = TestClient(app) + + # First turn + request_data = { + "model": "granite4:micro", + "messages": [{"role": "user", "content": "My name is Alice."}], + "max_tokens": 20, + } + + response = client.post("/v1/chat/completions", json=request_data) + assert response.status_code == 200 + first_response = response.json()["choices"][0]["message"]["content"] + + # Second turn with history + request_data = { + "model": "granite4:micro", + "messages": [ + {"role": "user", "content": "My name is Alice."}, + {"role": "assistant", "content": first_response}, + {"role": "user", "content": "What is my name?"}, + ], + "max_tokens": 20, + } + + response = client.post("/v1/chat/completions", json=request_data) + assert response.status_code == 200 + + +@pytest.mark.ollama +@pytest.mark.llm +def test_chat_completions_with_top_p(): + """Test chat completion with top_p parameter.""" + config = ServerConfig(backend_name="ollama", model_id="granite4:micro") + app = create_app(config) + client = TestClient(app) + + request_data = { + "model": "granite4:micro", + "messages": [{"role": "user", "content": "Hello"}], + "top_p": 0.9, + "max_tokens": 10, + } + + response = client.post("/v1/chat/completions", json=request_data) + assert response.status_code == 200 + + +def test_create_app_with_no_config(): + """Test that app can be created without explicit config.""" + app = create_app() + assert app is not None + assert app.state.config.backend_name == "ollama" + + +def test_server_config_with_backend_kwargs(): + """Test ServerConfig with additional backend kwargs.""" + config = ServerConfig( + backend_name="ollama", model_id="granite4:micro", timeout=60, max_retries=3 + ) + assert config.backend_kwargs["timeout"] == 60 + assert config.backend_kwargs["max_retries"] == 3 + + +@pytest.mark.ollama +@pytest.mark.llm +def test_session_reuse(): + """Test that sessions are reused for the same model.""" + config = ServerConfig(backend_name="ollama", model_id="granite4:micro") + app = create_app(config) + client = TestClient(app) + + request_data = { + "model": "granite4:micro", + "messages": [{"role": "user", "content": "Hello"}], + "max_tokens": 10, + } + + # Make two requests + response1 = client.post("/v1/chat/completions", json=request_data) + assert response1.status_code == 200 + + response2 = client.post("/v1/chat/completions", json=request_data) + assert response2.status_code == 200 + + # Verify session was reused (same session key) + session_key = f"{config.backend_name}:{request_data['model']}" + assert session_key in app.state.sessions + + +@pytest.mark.ollama +@pytest.mark.llm +def test_usage_information(): + """Test that usage information is returned.""" + config = ServerConfig(backend_name="ollama", model_id="granite4:micro") + app = create_app(config) + client = TestClient(app) + + request_data = { + "model": "granite4:micro", + "messages": [{"role": "user", "content": "Hi"}], + "max_tokens": 10, + } + + response = client.post("/v1/chat/completions", json=request_data) + assert response.status_code == 200 + + data = response.json() + # Usage may or may not be present depending on backend + if data.get("usage"): + assert "prompt_tokens" in data["usage"] + assert "completion_tokens" in data["usage"] + assert "total_tokens" in data["usage"] + + +@pytest.mark.ollama +@pytest.mark.llm +def test_model_list_with_custom_model(): + """Test models endpoint returns configured model.""" + config = ServerConfig(backend_name="ollama", model_id="custom-model") + app = create_app(config) + client = TestClient(app) + + response = client.get("/v1/models") + assert response.status_code == 200 + data = response.json() + assert data["data"][0]["id"] == "custom-model" + assert data["data"][0]["owned_by"] == "ollama" + + +def test_model_list_without_model_id(): + """Test models endpoint with no model_id configured.""" + config = ServerConfig(backend_name="ollama") + app = create_app(config) + client = TestClient(app) + + response = client.get("/v1/models") + assert response.status_code == 200 + data = response.json() + assert data["data"][0]["id"] == "default" + + +@pytest.mark.ollama +@pytest.mark.llm +def test_backend_config_with_base_url(): + """Test that base_url config is accepted and session is created.""" + config = ServerConfig( + backend_name="ollama", + model_id="granite4:micro", + base_url="http://localhost:11434", + ) + app = create_app(config) + client = TestClient(app) + + request_data = { + "model": "granite4:micro", + "messages": [{"role": "user", "content": "Hello"}], + "max_tokens": 10, + } + + response = client.post("/v1/chat/completions", json=request_data) + assert response.status_code == 200 + + # Verify session was created successfully with the base_url config + session_key = f"{config.backend_name}:{request_data['model']}" + assert session_key in app.state.sessions + assert app.state.sessions[session_key] is not None + + +@pytest.mark.ollama +@pytest.mark.llm +def test_session_with_different_models(): + """Test that different models create separate sessions.""" + config = ServerConfig(backend_name="ollama", model_id="granite4:micro") + app = create_app(config) + client = TestClient(app) + + # Request with first model + request_data_1 = { + "model": "granite4:micro", + "messages": [{"role": "user", "content": "Hello"}], + "max_tokens": 10, + } + response1 = client.post("/v1/chat/completions", json=request_data_1) + assert response1.status_code == 200 + + # Request with different model (simulated) + request_data_2 = { + "model": "llama3.2:1b", + "messages": [{"role": "user", "content": "Hello"}], + "max_tokens": 10, + } + response2 = client.post("/v1/chat/completions", json=request_data_2) + assert response2.status_code == 200 + + # Verify separate sessions were created + session_key_1 = f"{config.backend_name}:granite4:micro" + session_key_2 = f"{config.backend_name}:llama3.2:1b" + assert session_key_1 in app.state.sessions + assert session_key_2 in app.state.sessions + assert app.state.sessions[session_key_1] != app.state.sessions[session_key_2] + + +@pytest.mark.ollama +@pytest.mark.llm +@pytest.mark.qualitative +def test_chat_completions_streaming_with_tools(): + """Test streaming chat completion with tool calling.""" + config = ServerConfig(backend_name="ollama", model_id="granite4:micro") + app = create_app(config) + client = TestClient(app) + + tool_definition = { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the weather for a location", + "parameters": { + "type": "object", + "properties": { + "location": {"type": "string", "description": "City name"} + }, + "required": ["location"], + }, + }, + } + + request_data = { + "model": "granite4:micro", + "messages": [ + { + "role": "user", + "content": "What's the weather in Paris? Use the get_weather tool.", + } + ], + "tools": [tool_definition], + "stream": True, + "max_tokens": 100, + } + + with client.stream("POST", "/v1/chat/completions", json=request_data) as response: + assert response.status_code == 200 + assert response.headers["content-type"] == "text/event-stream; charset=utf-8" + + chunks = [] + for line in response.iter_lines(): + if line.startswith("data: "): + data_str = line[6:] + if data_str == "[DONE]": + break + chunk_data = json.loads(data_str) + chunks.append(chunk_data) + + # Verify we got chunks + assert len(chunks) > 0 + assert chunks[0]["object"] == "chat.completion.chunk" + + +@pytest.mark.ollama +@pytest.mark.llm +@pytest.mark.qualitative +def test_chat_completions_with_multiple_tools(): + """Test chat completion with multiple tool definitions.""" + config = ServerConfig(backend_name="ollama", model_id="granite4:micro") + app = create_app(config) + client = TestClient(app) + + tools = [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get weather for a location", + "parameters": { + "type": "object", + "properties": { + "location": {"type": "string", "description": "City name"} + }, + "required": ["location"], + }, + }, + }, + { + "type": "function", + "function": { + "name": "get_time", + "description": "Get current time for a timezone", + "parameters": { + "type": "object", + "properties": { + "timezone": {"type": "string", "description": "Timezone name"} + }, + "required": ["timezone"], + }, + }, + }, + ] + + request_data = { + "model": "granite4:micro", + "messages": [ + { + "role": "user", + "content": "What's the weather in Tokyo? Use the appropriate tool.", + } + ], + "tools": tools, + "tool_choice": "auto", + "max_tokens": 100, + } + + response = client.post("/v1/chat/completions", json=request_data) + assert response.status_code == 200 + + data = response.json() + assert data["object"] == "chat.completion" + + # If tools were called, verify the structure + message = data["choices"][0]["message"] + if message.get("tool_calls"): + for tool_call in message["tool_calls"]: + assert "id" in tool_call + assert "type" in tool_call + assert tool_call["type"] == "function" + assert "function" in tool_call + assert "name" in tool_call["function"] + assert tool_call["function"]["name"] in ["get_weather", "get_time"]