ai-workspace-services/litellm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

[Feat] New LLM Provider - Docker Model Runner (#16948)

Browse Source 
* add DOCKER_MODEL_RUNNER

* add DockerModelRunnerChatConfig Transorm

* add docker_model_runner

* add docker_model_runner

* docs docker model runner

* add DockerModelRunnerChatConfig

* add docker_model_runner to providers

* test_completion_hits_correct_url_and_body

* fix sidebar

* TestDockerModelRunnerIntegration

* test_completion_with_custom_engine_and_host

* docs docker model runner

* docs fix
This commit is contained in:
Ishaan Jaff 2025-11-21 16:09:32 -08:00 committed by GitHub
parent d88580fa28
commit 8e318dd06c
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
10 changed files with 688 additions and 59 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

277
docs/my-website/docs/providers/docker_model_runner.md Normal file
View File

@ -0,0 +1,277 @@
import Tabs from '@theme/Tabs';
import TabItem from '@theme/TabItem';
# Docker Model Runner
## Overview
| Property | Details |
|-------|-------|
| Description | Docker Model Runner allows you to run large language models locally using Docker Desktop. |
| Provider Route on LiteLLM | `docker_model_runner/` |
| Link to Provider Doc | [Docker Model Runner ↗](https://docs.docker.com/ai/model-runner/) |
| Base URL | `http://localhost:22088` |
| Supported Operations | [`/chat/completions`](#sample-usage) |
<br />
<br />
https://docs.docker.com/ai/model-runner/
**We support ALL Docker Model Runner models, just set `docker_model_runner/` as a prefix when sending completion requests**
## Quick Start
Docker Model Runner is a Docker Desktop feature that lets you run AI models locally. It provides better performance than other local solutions while maintaining OpenAI compatibility.
### Installation
1. Install [Docker Desktop](https://www.docker.com/products/docker-desktop/)
2. Enable Docker Model Runner in Docker Desktop settings
3. Download your preferred model through Docker Desktop
## Environment Variables
```python showLineNumbers title="Environment Variables"
os.environ["DOCKER_MODEL_RUNNER_API_BASE"] = "http://localhost:22088/engines/llama.cpp" # Optional - defaults to this
os.environ["DOCKER_MODEL_RUNNER_API_KEY"] = "dummy-key" # Optional - Docker Model Runner may not require auth for local instances
```
**Note:**
- Docker Model Runner typically runs locally and may not require authentication. LiteLLM will use a dummy key by default if no key is provided.
- The API base should include the engine path (e.g., `/engines/llama.cpp`)
## API Base Structure
Docker Model Runner uses a unique URL structure:
```
http://model-runner.docker.internal/engines/{engine}/v1/chat/completions
```
Where `{engine}` is the engine you want to use (typically `llama.cpp`).
**Important:** Specify the engine in your `api_base` URL, not in the model name:
- ✅ Correct: `api_base="http://localhost:22088/engines/llama.cpp"`, `model="docker_model_runner/llama-3.1"`
- ❌ Incorrect: `api_base="http://localhost:22088"`, `model="docker_model_runner/llama.cpp/llama-3.1"`
## Usage - LiteLLM Python SDK
### Non-streaming
```python showLineNumbers title="Docker Model Runner Non-streaming Completion"
import os
import litellm
from litellm import completion
# Specify the engine in the api_base URL
os.environ["DOCKER_MODEL_RUNNER_API_BASE"] = "http://localhost:22088/engines/llama.cpp"
messages = [{"content": "Hello, how are you?", "role": "user"}]
# Docker Model Runner call
response = completion(
model="docker_model_runner/llama-3.1",
messages=messages
)
print(response)
```
### Streaming
```python showLineNumbers title="Docker Model Runner Streaming Completion"
import os
import litellm
from litellm import completion
# Specify the engine in the api_base URL
os.environ["DOCKER_MODEL_RUNNER_API_BASE"] = "http://localhost:22088/engines/llama.cpp"
messages = [{"content": "Hello, how are you?", "role": "user"}]
# Docker Model Runner call with streaming
response = completion(
model="docker_model_runner/llama-3.1",
messages=messages,
stream=True
)
for chunk in response:
print(chunk)
```
### Custom API Base and Engine
```python showLineNumbers title="Custom API Base with Different Engine"
import litellm
from litellm import completion
messages = [{"content": "Hello, how are you?", "role": "user"}]
# Specify the engine in the api_base URL
# Using a different host and engine
response = completion(
model="docker_model_runner/llama-3.1",
messages=messages,
api_base="http://model-runner.docker.internal/engines/llama.cpp"
)
print(response)
```
### Using Different Engines
```python showLineNumbers title="Using a Different Engine"
import litellm
from litellm import completion
messages = [{"content": "Hello, how are you?", "role": "user"}]
# To use a different engine, specify it in the api_base
# For example, if Docker Model Runner supports other engines:
response = completion(
model="docker_model_runner/mistral-7b",
messages=messages,
api_base="http://localhost:22088/engines/custom-engine"
)
print(response)
```
## Usage - LiteLLM Proxy
Add the following to your LiteLLM Proxy configuration file:
```yaml showLineNumbers title="config.yaml"
model_list:
- model_name: llama-3.1
litellm_params:
model: docker_model_runner/llama-3.1
api_base: http://localhost:22088/engines/llama.cpp
- model_name: mistral-7b
litellm_params:
model: docker_model_runner/mistral-7b
api_base: http://localhost:22088/engines/llama.cpp
```
Start your LiteLLM Proxy server:
```bash showLineNumbers title="Start LiteLLM Proxy"
litellm --config config.yaml
# RUNNING on http://0.0.0.0:4000
```
<Tabs>
<TabItem value="openai-sdk" label="OpenAI SDK">
```python showLineNumbers title="Docker Model Runner via Proxy - Non-streaming"
from openai import OpenAI
# Initialize client with your proxy URL
client = OpenAI(
base_url="http://localhost:4000", # Your proxy URL
api_key="your-proxy-api-key" # Your proxy API key
)
# Non-streaming response
response = client.chat.completions.create(
model="llama-3.1",
messages=[{"role": "user", "content": "hello from litellm"}]
)
print(response.choices[0].message.content)
```
```python showLineNumbers title="Docker Model Runner via Proxy - Streaming"
from openai import OpenAI
# Initialize client with your proxy URL
client = OpenAI(
base_url="http://localhost:4000", # Your proxy URL
api_key="your-proxy-api-key" # Your proxy API key
)
# Streaming response
response = client.chat.completions.create(
model="llama-3.1",
messages=[{"role": "user", "content": "hello from litellm"}],
stream=True
)
for chunk in response:
if chunk.choices[0].delta.content is not None:
print(chunk.choices[0].delta.content, end="")
```
</TabItem>
<TabItem value="litellm-sdk" label="LiteLLM SDK">
```python showLineNumbers title="Docker Model Runner via Proxy - LiteLLM SDK"
import litellm
# Configure LiteLLM to use your proxy
response = litellm.completion(
model="litellm_proxy/llama-3.1",
messages=[{"role": "user", "content": "hello from litellm"}],
api_base="http://localhost:4000",
api_key="your-proxy-api-key"
)
print(response.choices[0].message.content)
```
```python showLineNumbers title="Docker Model Runner via Proxy - LiteLLM SDK Streaming"
import litellm
# Configure LiteLLM to use your proxy with streaming
response = litellm.completion(
model="litellm_proxy/llama-3.1",
messages=[{"role": "user", "content": "hello from litellm"}],
api_base="http://localhost:4000",
api_key="your-proxy-api-key",
stream=True
)
for chunk in response:
if hasattr(chunk.choices[0], 'delta') and chunk.choices[0].delta.content is not None:
print(chunk.choices[0].delta.content, end="")
```
</TabItem>
<TabItem value="curl" label="cURL">
```bash showLineNumbers title="Docker Model Runner via Proxy - cURL"
curl http://localhost:4000/v1/chat/completions \
-H "Content-Type: application/json" \
-H "Authorization: Bearer your-proxy-api-key" \
-d '{
"model": "llama-3.1",
"messages": [{"role": "user", "content": "hello from litellm"}]
}'
```
```bash showLineNumbers title="Docker Model Runner via Proxy - cURL Streaming"
curl http://localhost:4000/v1/chat/completions \
-H "Content-Type: application/json" \
-H "Authorization: Bearer your-proxy-api-key" \
-d '{
"model": "llama-3.1",
"messages": [{"role": "user", "content": "hello from litellm"}],
"stream": true
}'
```
</TabItem>
</Tabs>
For more detailed information on using the LiteLLM Proxy, see the [LiteLLM Proxy documentation](../providers/litellm_proxy).
## API Reference
For detailed API information, see the [Docker Model Runner API Reference](https://docs.docker.com/ai/model-runner/api-reference/).

117
docs/my-website/sidebars.js
View File

@ -530,13 +530,39 @@ const sidebars = {
"providers/bedrock_vector_store",
]
},
"providers/milvus_vector_stores",
"providers/litellm_proxy",
"providers/meta_llama",
"providers/mistral",
"providers/ai21",
"providers/aiml",
"providers/aleph_alpha",
"providers/anyscale",
"providers/baseten",
"providers/bytez",
"providers/cerebras",
"providers/clarifai",
"providers/cloudflare_workers",
"providers/codestral",
"providers/cohere",
"providers/anyscale",
"providers/cometapi",
"providers/compactifai",
"providers/custom_llm_server",
"providers/dashscope",
"providers/databricks",
"providers/datarobot",
"providers/deepgram",
"providers/deepinfra",
"providers/deepseek",
"providers/docker_model_runner",
"providers/elevenlabs",
"providers/fal_ai",
"providers/featherless_ai",
"providers/fireworks_ai",
"providers/friendliai",
"providers/galadriel",
"providers/github",
"providers/github_copilot",
"providers/gradient_ai",
"providers/groq",
"providers/heroku",
{
type: "category",
label: "HuggingFace",
@ -546,10 +572,21 @@ const sidebars = {
]
},
"providers/hyperbolic",
"providers/databricks",
"providers/deepgram",
"providers/watsonx",
"providers/predibase",
"providers/infinity",
"providers/jina_ai",
"providers/lambda_ai",
"providers/lemonade",
"providers/llamafile",
"providers/lm_studio",
"providers/meta_llama",
"providers/milvus_vector_stores",
"providers/mistral",
"providers/moonshot",
"providers/morph",
"providers/nebius",
"providers/nlp_cloud",
"providers/novita",
{ type: "doc", id: "providers/nscale", label: "Nscale (EU Sovereign)" },
{
type: "category",
label: "Nvidia NIM",
@ -558,37 +595,13 @@ const sidebars = {
"providers/nvidia_nim_rerank",
]
},
{ type: "doc", id: "providers/nscale", label: "Nscale (EU Sovereign)" },
"providers/xai",
"providers/moonshot",
"providers/lm_studio",
"providers/cerebras",
"providers/volcano",
"providers/triton-inference-server",
"providers/oci",
"providers/ollama",
"providers/openrouter",
"providers/ovhcloud",
"providers/perplexity",
"providers/friendliai",
"providers/galadriel",
"providers/topaz",
"providers/groq",
"providers/deepseek",
"providers/elevenlabs",
"providers/fal_ai",
"providers/fireworks_ai",
"providers/clarifai",
"providers/compactifai",
"providers/lemonade",
"providers/vllm",
"providers/llamafile",
"providers/infinity",
"providers/xinference",
"providers/aiml",
"providers/cloudflare_workers",
"providers/deepinfra",
"providers/github",
"providers/github_copilot",
"providers/ai21",
"providers/nlp_cloud",
"providers/petals",
"providers/predibase",
"providers/recraft",
"providers/replicate",
{
@ -599,32 +612,20 @@ const sidebars = {
"providers/runwayml/videos",
]
},
"providers/sambanova",
"providers/snowflake",
"providers/togetherai",
"providers/topaz",
"providers/triton-inference-server",
"providers/v0",
"providers/vercel_ai_gateway",
"providers/morph",
"providers/lambda_ai",
"providers/novita",
"providers/vllm",
"providers/volcano",
"providers/voyage",
"providers/jina_ai",
"providers/aleph_alpha",
"providers/baseten",
"providers/openrouter",
"providers/sambanova",
"providers/custom_llm_server",
"providers/petals",
"providers/snowflake",
"providers/gradient_ai",
"providers/featherless_ai",
"providers/nebius",
"providers/dashscope",
"providers/bytez",
"providers/heroku",
"providers/oci",
"providers/datarobot",
"providers/ovhcloud",
"providers/wandb_inference",
"providers/cometapi",
"providers/watsonx",
"providers/xai",
"providers/xinference",
],
},
{

5
litellm/__init__.py
View File

@ -563,6 +563,7 @@ wandb_models: Set = set(WANDB_MODELS)
ovhcloud_models: Set = set()
ovhcloud_embedding_models: Set = set()
lemonade_models: Set = set()
docker_model_runner_models: Set = set()
def is_bedrock_pricing_only_model(key: str) -> bool:
@ -797,6 +798,8 @@ def add_known_models():
ovhcloud_embedding_models.add(key)
elif value.get("litellm_provider") == "lemonade":
lemonade_models.add(key)
elif value.get("litellm_provider") == "docker_model_runner":
docker_model_runner_models.add(key)
add_known_models()
@ -900,6 +903,7 @@ model_list = list(
| wandb_models
| ovhcloud_models
| lemonade_models
| docker_model_runner_models
| set(clarifai_models)
)
@ -1350,6 +1354,7 @@ from .llms.nebius.chat.transformation import NebiusConfig
from .llms.wandb.chat.transformation import WandbConfig
from .llms.dashscope.chat.transformation import DashScopeChatConfig
from .llms.moonshot.chat.transformation import MoonshotChatConfig
from .llms.docker_model_runner.chat.transformation import DockerModelRunnerChatConfig
from .llms.v0.chat.transformation import V0ChatConfig
from .llms.oci.chat.transformation import OCIChatConfig
from .llms.morph.chat.transformation import MorphChatConfig

2
litellm/constants.py
View File

@ -381,6 +381,7 @@ LITELLM_CHAT_PROVIDERS = [
"wandb",
"ovhcloud",
"lemonade",
"docker_model_runner",
]
LITELLM_EMBEDDING_PROVIDERS_SUPPORTING_INPUT_ARRAY_OF_TOKENS = [
@ -567,6 +568,7 @@ openai_compatible_providers: List = [
"wandb",
"cometapi",
"clarifai",
"docker_model_runner",
]
openai_text_completion_compatible_providers: List = (
[ # providers that support `/v1/completions`

7
litellm/litellm_core_utils/get_llm_provider_logic.py
View File

@ -741,6 +741,13 @@ def _get_openai_compatible_provider_info( # noqa: PLR0915
) = litellm.MoonshotChatConfig()._get_openai_compatible_provider_info(
api_base, api_key
)
elif custom_llm_provider == "docker_model_runner":
(
api_base,
dynamic_api_key,
) = litellm.DockerModelRunnerChatConfig()._get_openai_compatible_provider_info(
api_base, api_key
)
elif custom_llm_provider == "v0":
(
api_base,

144
litellm/llms/docker_model_runner/chat/transformation.py Normal file
View File

@ -0,0 +1,144 @@
"""
Translates from OpenAI's `/v1/chat/completions` to Docker Model Runner's `/engines/{engine}/v1/chat/completions`
Docker Model Runner API Reference: https://docs.docker.com/ai/model-runner/api-reference/
"""
from typing import Any, Coroutine, List, Literal, Optional, Tuple, Union, overload
from litellm.litellm_core_utils.prompt_templates.common_utils import (
handle_messages_with_content_list_to_str_conversion,
)
from litellm.secret_managers.main import get_secret_str
from litellm.types.llms.openai import AllMessageValues
from ...openai.chat.gpt_transformation import OpenAIGPTConfig
class DockerModelRunnerChatConfig(OpenAIGPTConfig):
"""
Configuration for Docker Model Runner API.
Docker Model Runner uses URLs in the format: /engines/{engine}/v1/chat/completions
The engine name (e.g., "llama.cpp") is part of the API endpoint path.
"""
@overload
def _transform_messages(
self, messages: List[AllMessageValues], model: str, is_async: Literal[True]
) -> Coroutine[Any, Any, List[AllMessageValues]]:
...
@overload
def _transform_messages(
self,
messages: List[AllMessageValues],
model: str,
is_async: Literal[False] = False,
) -> List[AllMessageValues]:
...
def _transform_messages(
self, messages: List[AllMessageValues], model: str, is_async: bool = False
) -> Union[List[AllMessageValues], Coroutine[Any, Any, List[AllMessageValues]]]:
"""
Docker Model Runner is OpenAI-compatible, so we use standard message transformation.
"""
messages = handle_messages_with_content_list_to_str_conversion(messages)
if is_async:
return super()._transform_messages(
messages=messages, model=model, is_async=True
)
else:
return super()._transform_messages(
messages=messages, model=model, is_async=False
)
def _get_openai_compatible_provider_info(
self, api_base: Optional[str], api_key: Optional[str]
) -> Tuple[Optional[str], Optional[str]]:
"""
Get API base and key for Docker Model Runner.
Default API base: http://localhost:22088/engines/llama.cpp
The engine path should be included in the api_base.
"""
api_base = (
api_base
or get_secret_str("DOCKER_MODEL_RUNNER_API_BASE")
or "http://localhost:22088/engines/llama.cpp"
) # type: ignore
# Docker Model Runner may not require authentication for local instances
dynamic_api_key = api_key or get_secret_str("DOCKER_MODEL_RUNNER_API_KEY") or "dummy-key"
return api_base, dynamic_api_key
def get_complete_url(
self,
api_base: Optional[str],
api_key: Optional[str],
model: str,
optional_params: dict,
litellm_params: dict,
stream: Optional[bool] = None,
) -> str:
"""
Build the complete URL for Docker Model Runner API.
Docker Model Runner uses URLs in the format: /engines/{engine}/v1/chat/completions
The engine name should be specified in the api_base:
- api_base="http://model-runner.docker.internal/engines/llama.cpp"
- Default: "http://localhost:22088/engines/llama.cpp"
Args:
api_base: Base URL for the Docker Model Runner instance including engine path
api_key: API key (may not be required for local instances)
model: Model name (e.g., "llama-3.1")
optional_params: Optional parameters
litellm_params: LiteLLM parameters
stream: Whether streaming is enabled
Returns:
Complete URL for the API call
"""
if not api_base:
api_base = "http://localhost:22088/engines/llama.cpp"
# Remove trailing slashes from api_base
api_base = api_base.rstrip("/")
# Build the URL: {api_base}/v1/chat/completions
# api_base is expected to already contain the engine path
complete_url = f"{api_base}/v1/chat/completions"
return complete_url
def get_supported_openai_params(self, model: str) -> list:
"""
Get the supported OpenAI params for Docker Model Runner.
Docker Model Runner is OpenAI-compatible and supports standard parameters.
"""
return super().get_supported_openai_params(model=model)
def map_openai_params(
self,
non_default_params: dict,
optional_params: dict,
model: str,
drop_params: bool,
) -> dict:
"""
Map OpenAI parameters to Docker Model Runner parameters.
Docker Model Runner is OpenAI-compatible, so most parameters map directly.
"""
supported_openai_params = self.get_supported_openai_params(model)
for param, value in non_default_params.items():
if param == "max_completion_tokens":
optional_params["max_tokens"] = value
elif param in supported_openai_params:
optional_params[param] = value
return optional_params

1
litellm/types/utils.py
View File

@ -2594,6 +2594,7 @@ class LlmProviders(str, Enum):
EMPOWER = "empower"
GITHUB = "github"
COMPACTIFAI = "compactifai"
DOCKER_MODEL_RUNNER = "docker_model_runner"
CUSTOM = "custom"
LITELLM_PROXY = "litellm_proxy"
HOSTED_VLLM = "hosted_vllm"

6
litellm/utils.py
View File

@ -7205,6 +7205,8 @@ class ProviderConfigManager:
return litellm.DashScopeChatConfig()
elif litellm.LlmProviders.MOONSHOT == provider:
return litellm.MoonshotChatConfig()
elif litellm.LlmProviders.DOCKER_MODEL_RUNNER == provider:
return litellm.DockerModelRunnerChatConfig()
elif litellm.LlmProviders.V0 == provider:
return litellm.V0ChatConfig()
elif litellm.LlmProviders.MORPH == provider:
@ -7758,7 +7760,9 @@ class ProviderConfigManager:
return LiteLLMProxyImageEditConfig()
elif LlmProviders.VERTEX_AI == provider:
from litellm.llms.vertex_ai.image_edit import get_vertex_ai_image_edit_config
from litellm.llms.vertex_ai.image_edit import (
get_vertex_ai_image_edit_config,
)
return get_vertex_ai_image_edit_config(model)
return None

16
provider_endpoints_support.json
View File

@ -1036,6 +1036,22 @@
"rerank": false
}
},
"docker_model_runner": {
"display_name": "Docker Model Runner (`docker_model_runner`)",
"url": "https://docs.litellm.ai/docs/providers/docker_model_runner",
"endpoints": {
"chat_completions": true,
"messages": true,
"responses": true,
"embeddings": false,
"image_generations": false,
"audio_transcriptions": false,
"audio_speech": false,
"moderations": false,
"batches": false,
"rerank": false
}
},
"morph": {
"display_name": "Morph (`morph`)",
"url": "https://docs.litellm.ai/docs/providers/morph",

172
tests/test_litellm/llms/docker_model_runner/test_docker_model_runner_chat_transformation.py Normal file
View File

@ -0,0 +1,172 @@
"""
Unit tests for Docker Model Runner configuration.
This test validates that litellm.completion correctly routes requests to Docker Model Runner
with the proper URL structure and request body.
"""
import os
import sys
sys.path.insert(
0, os.path.abspath("../../../../..")
)
import json
from unittest.mock import Mock, patch
import pytest
import litellm
from litellm import completion
class TestDockerModelRunnerIntegration:
"""Integration test for Docker Model Runner"""
@pytest.mark.asyncio
async def test_completion_hits_correct_url_and_body(self):
"""
Test that litellm.completion with docker_model_runner provider:
1. Hits the correct URL: {api_base}/v1/chat/completions where api_base includes engine path
2. Sends the correct request body with messages and parameters
"""
with patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.post") as mock_post:
# Mock the response
mock_response = Mock()
mock_response.json.return_value = {
"id": "chatcmpl-123",
"object": "chat.completion",
"created": 1677652288,
"model": "llama-3.1",
"choices": [{
"index": 0,
"message": {
"role": "assistant",
"content": "Hello! How can I help you today?"
},
"finish_reason": "stop"
}],
"usage": {
"prompt_tokens": 10,
"completion_tokens": 20,
"total_tokens": 30
}
}
mock_response.status_code = 200
mock_response.headers = {"content-type": "application/json"}
mock_post.return_value = mock_response
# Make the completion call with engine in api_base
response = completion(
model="docker_model_runner/llama-3.1",
messages=[{"role": "user", "content": "Hello, how are you?"}],
api_base="http://localhost:22088/engines/llama.cpp",
temperature=0.7,
max_tokens=100
)
# Verify the URL was correct
assert mock_post.called
call_args = mock_post.call_args
url = call_args[1]["url"]
print("URL For request", url)
print("request body for request", json.dumps(call_args[1]["data"], indent=4))
# Should hit {api_base}/v1/chat/completions where api_base includes engine
assert "/engines/llama.cpp/v1/chat/completions" in url
assert "http://localhost:22088" in url
# Verify the request body
request_data = call_args[1]["data"]
if isinstance(request_data, str):
request_data = json.loads(request_data)
# Check messages
assert "messages" in request_data
assert len(request_data["messages"]) == 1
assert request_data["messages"][0]["role"] == "user"
assert request_data["messages"][0]["content"] == "Hello, how are you?"
# Check parameters
assert request_data["temperature"] == 0.7
assert request_data["max_tokens"] == 100
# Verify response
assert response.choices[0].message.content == "Hello! How can I help you today?"
@pytest.mark.asyncio
async def test_completion_with_custom_engine_and_host(self):
"""
Test that litellm.completion works with custom engine and host:
1. Uses model-runner.docker.internal as host
2. Specifies a different engine in the api_base
3. Model name is sent in the request body
"""
with patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.post") as mock_post:
# Mock the response
mock_response = Mock()
mock_response.json.return_value = {
"id": "chatcmpl-456",
"object": "chat.completion",
"created": 1677652288,
"model": "mistral-7b",
"choices": [{
"index": 0,
"message": {
"role": "assistant",
"content": "Bonjour! How can I assist you?"
},
"finish_reason": "stop"
}],
"usage": {
"prompt_tokens": 15,
"completion_tokens": 25,
"total_tokens": 40
}
}
mock_response.status_code = 200
mock_response.headers = {"content-type": "application/json"}
mock_post.return_value = mock_response
# Make the completion call with custom engine and host
response = completion(
model="docker_model_runner/mistral-7b",
messages=[{"role": "user", "content": "Hello!"}],
api_base="http://model-runner.docker.internal/engines/custom-engine",
temperature=0.5,
max_tokens=200
)
# Verify the URL was correct
assert mock_post.called
call_args = mock_post.call_args
url = call_args[1]["url"]
print("URL For request", url)
print("request body for request", json.dumps(call_args[1]["data"], indent=4))
# Should hit the custom host and engine
assert "model-runner.docker.internal" in url
assert "/engines/custom-engine/v1/chat/completions" in url
# Verify the request body contains the model name
request_data = call_args[1]["data"]
if isinstance(request_data, str):
request_data = json.loads(request_data)
# Check that model name is in the request body
assert request_data["model"] == "mistral-7b"
# Check messages
assert "messages" in request_data
assert len(request_data["messages"]) == 1
assert request_data["messages"][0]["role"] == "user"
assert request_data["messages"][0]["content"] == "Hello!"
# Check parameters
assert request_data["temperature"] == 0.5
assert request_data["max_tokens"] == 200
# Verify response
assert response.choices[0].message.content == "Bonjour! How can I assist you?"