mirror of
https://github.com/ollama/ollama-python.git
synced 2026-06-16 21:24:52 +00:00
Compare commits
8 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 588e338b2b | |||
| 422724ee8c | |||
| 70dd0b7e63 | |||
| 1066246ab5 | |||
| 4b10dee2b2 | |||
| e956a331e8 | |||
| 12f7302d5f | |||
| 366180aa8f |
@@ -37,9 +37,6 @@ See [_types.py](ollama/_types.py) for more information on the response types.
|
||||
|
||||
Response streaming can be enabled by setting `stream=True`.
|
||||
|
||||
> [!NOTE]
|
||||
> Streaming Tool/Function calling is not yet supported.
|
||||
|
||||
```python
|
||||
from ollama import chat
|
||||
|
||||
|
||||
@@ -30,6 +30,16 @@ python3 examples/<example>.py
|
||||
- [multimodal_generate.py](multimodal_generate.py)
|
||||
|
||||
|
||||
### Structured Outputs - Generate structured outputs with a model
|
||||
- [structured-outputs.py](structured-outputs.py)
|
||||
- [async-structured-outputs.py](async-structured-outputs.py)
|
||||
- [structured-outputs-image.py](structured-outputs-image.py)
|
||||
|
||||
|
||||
### Tokenization - Tokenize and detokenize text with a model
|
||||
- [tokenization.py](tokenization.py)
|
||||
|
||||
|
||||
### Ollama List - List all downloaded models and their properties
|
||||
- [list.py](list.py)
|
||||
|
||||
|
||||
@@ -0,0 +1,32 @@
|
||||
from pydantic import BaseModel
|
||||
from ollama import AsyncClient
|
||||
import asyncio
|
||||
|
||||
|
||||
# Define the schema for the response
|
||||
class FriendInfo(BaseModel):
|
||||
name: str
|
||||
age: int
|
||||
is_available: bool
|
||||
|
||||
|
||||
class FriendList(BaseModel):
|
||||
friends: list[FriendInfo]
|
||||
|
||||
|
||||
async def main():
|
||||
client = AsyncClient()
|
||||
response = await client.chat(
|
||||
model='llama3.1:8b',
|
||||
messages=[{'role': 'user', 'content': 'I have two friends. The first is Ollama 22 years old busy saving the world, and the second is Alonso 23 years old and wants to hang out. Return a list of friends in JSON format'}],
|
||||
format=FriendList.model_json_schema(), # Use Pydantic to generate the schema
|
||||
options={'temperature': 0}, # Make responses more deterministic
|
||||
)
|
||||
|
||||
# Use Pydantic to validate the response
|
||||
friends_response = FriendList.model_validate_json(response.message.content)
|
||||
print(friends_response)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
asyncio.run(main())
|
||||
+24
-10
@@ -41,21 +41,21 @@ subtract_two_numbers_tool = {
|
||||
},
|
||||
}
|
||||
|
||||
messages = [{'role': 'user', 'content': 'What is three plus one?'}]
|
||||
print('Prompt:', messages[0]['content'])
|
||||
|
||||
available_functions = {
|
||||
'add_two_numbers': add_two_numbers,
|
||||
'subtract_two_numbers': subtract_two_numbers,
|
||||
}
|
||||
|
||||
|
||||
async def main():
|
||||
client = ollama.AsyncClient()
|
||||
|
||||
prompt = 'What is three plus one?'
|
||||
print('Prompt:', prompt)
|
||||
|
||||
available_functions = {
|
||||
'add_two_numbers': add_two_numbers,
|
||||
'subtract_two_numbers': subtract_two_numbers,
|
||||
}
|
||||
|
||||
response: ChatResponse = await client.chat(
|
||||
'llama3.1',
|
||||
messages=[{'role': 'user', 'content': prompt}],
|
||||
messages=messages,
|
||||
tools=[add_two_numbers, subtract_two_numbers_tool],
|
||||
)
|
||||
|
||||
@@ -66,10 +66,24 @@ async def main():
|
||||
if function_to_call := available_functions.get(tool.function.name):
|
||||
print('Calling function:', tool.function.name)
|
||||
print('Arguments:', tool.function.arguments)
|
||||
print('Function output:', function_to_call(**tool.function.arguments))
|
||||
output = function_to_call(**tool.function.arguments)
|
||||
print('Function output:', output)
|
||||
else:
|
||||
print('Function', tool.function.name, 'not found')
|
||||
|
||||
# Only needed to chat with the model using the tool call results
|
||||
if response.message.tool_calls:
|
||||
# Add the function response to messages for the model to use
|
||||
messages.append(response.message)
|
||||
messages.append({'role': 'tool', 'content': str(output), 'name': tool.function.name})
|
||||
|
||||
# Get final response from model with function outputs
|
||||
final_response = await client.chat('llama3.1', messages=messages)
|
||||
print('Final response:', final_response.message.content)
|
||||
|
||||
else:
|
||||
print('No tool calls returned from model')
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
try:
|
||||
|
||||
@@ -0,0 +1,49 @@
|
||||
from pathlib import Path
|
||||
from pydantic import BaseModel
|
||||
from typing import List, Optional, Literal
|
||||
from ollama import chat
|
||||
|
||||
|
||||
# Define the schema for image objects
|
||||
class Object(BaseModel):
|
||||
name: str
|
||||
confidence: float
|
||||
attributes: str
|
||||
|
||||
|
||||
class ImageDescription(BaseModel):
|
||||
summary: str
|
||||
objects: List[Object]
|
||||
scene: str
|
||||
colors: List[str]
|
||||
time_of_day: Literal['Morning', 'Afternoon', 'Evening', 'Night']
|
||||
setting: Literal['Indoor', 'Outdoor', 'Unknown']
|
||||
text_content: Optional[str] = None
|
||||
|
||||
|
||||
# Get path from user input
|
||||
path = input('Enter the path to your image: ')
|
||||
path = Path(path)
|
||||
|
||||
# Verify the file exists
|
||||
if not path.exists():
|
||||
raise FileNotFoundError(f'Image not found at: {path}')
|
||||
|
||||
# Set up chat as usual
|
||||
response = chat(
|
||||
model='llama3.2-vision',
|
||||
format=ImageDescription.model_json_schema(), # Pass in the schema for the response
|
||||
messages=[
|
||||
{
|
||||
'role': 'user',
|
||||
'content': 'Analyze this image and return a detailed JSON description including objects, scene, colors and any text detected. If you cannot determine certain details, leave those fields empty.',
|
||||
'images': [path],
|
||||
},
|
||||
],
|
||||
options={'temperature': 0}, # Set temperature to 0 for more deterministic output
|
||||
)
|
||||
|
||||
|
||||
# Convert received content to the schema
|
||||
image_analysis = ImageDescription.model_validate_json(response.message.content)
|
||||
print(image_analysis)
|
||||
@@ -0,0 +1,26 @@
|
||||
from ollama import chat
|
||||
from pydantic import BaseModel
|
||||
|
||||
|
||||
# Define the schema for the response
|
||||
class FriendInfo(BaseModel):
|
||||
name: str
|
||||
age: int
|
||||
is_available: bool
|
||||
|
||||
|
||||
class FriendList(BaseModel):
|
||||
friends: list[FriendInfo]
|
||||
|
||||
|
||||
# schema = {'type': 'object', 'properties': {'friends': {'type': 'array', 'items': {'type': 'object', 'properties': {'name': {'type': 'string'}, 'age': {'type': 'integer'}, 'is_available': {'type': 'boolean'}}, 'required': ['name', 'age', 'is_available']}}}, 'required': ['friends']}
|
||||
response = chat(
|
||||
model='llama3.1:8b',
|
||||
messages=[{'role': 'user', 'content': 'I have two friends. The first is Ollama 22 years old busy saving the world, and the second is Alonso 23 years old and wants to hang out. Return a list of friends in JSON format'}],
|
||||
format=FriendList.model_json_schema(), # Use Pydantic to generate the schema or format=schema
|
||||
options={'temperature': 0}, # Make responses more deterministic
|
||||
)
|
||||
|
||||
# Use Pydantic to validate the response
|
||||
friends_response = FriendList.model_validate_json(response.message.content)
|
||||
print(friends_response)
|
||||
@@ -0,0 +1,10 @@
|
||||
import ollama
|
||||
|
||||
# Get tokens from a model
|
||||
response = ollama.tokenize(model='llama3.2', text='Why the sky is blue?')
|
||||
tokens = response.tokens
|
||||
print('Tokens from model', tokens)
|
||||
|
||||
# Convert tokens back to text
|
||||
response = ollama.detokenize(model='llama3.2', tokens=tokens)
|
||||
print('Text from tokens', response.text) # Prints: Why the sky is blue?
|
||||
+18
-4
@@ -40,8 +40,8 @@ subtract_two_numbers_tool = {
|
||||
},
|
||||
}
|
||||
|
||||
prompt = 'What is three plus one?'
|
||||
print('Prompt:', prompt)
|
||||
messages = [{'role': 'user', 'content': 'What is three plus one?'}]
|
||||
print('Prompt:', messages[0]['content'])
|
||||
|
||||
available_functions = {
|
||||
'add_two_numbers': add_two_numbers,
|
||||
@@ -50,7 +50,7 @@ available_functions = {
|
||||
|
||||
response: ChatResponse = chat(
|
||||
'llama3.1',
|
||||
messages=[{'role': 'user', 'content': prompt}],
|
||||
messages=messages,
|
||||
tools=[add_two_numbers, subtract_two_numbers_tool],
|
||||
)
|
||||
|
||||
@@ -61,6 +61,20 @@ if response.message.tool_calls:
|
||||
if function_to_call := available_functions.get(tool.function.name):
|
||||
print('Calling function:', tool.function.name)
|
||||
print('Arguments:', tool.function.arguments)
|
||||
print('Function output:', function_to_call(**tool.function.arguments))
|
||||
output = function_to_call(**tool.function.arguments)
|
||||
print('Function output:', output)
|
||||
else:
|
||||
print('Function', tool.function.name, 'not found')
|
||||
|
||||
# Only needed to chat with the model using the tool call results
|
||||
if response.message.tool_calls:
|
||||
# Add the function response to messages for the model to use
|
||||
messages.append(response.message)
|
||||
messages.append({'role': 'tool', 'content': str(output), 'name': tool.function.name})
|
||||
|
||||
# Get final response from model with function outputs
|
||||
final_response = chat('llama3.1', messages=messages)
|
||||
print('Final response:', final_response.message.content)
|
||||
|
||||
else:
|
||||
print('No tool calls returned from model')
|
||||
|
||||
@@ -12,6 +12,8 @@ from ollama._types import (
|
||||
ListResponse,
|
||||
ShowResponse,
|
||||
ProcessResponse,
|
||||
TokenizeResponse,
|
||||
DetokenizeResponse,
|
||||
RequestError,
|
||||
ResponseError,
|
||||
)
|
||||
@@ -31,6 +33,8 @@ __all__ = [
|
||||
'ListResponse',
|
||||
'ShowResponse',
|
||||
'ProcessResponse',
|
||||
'TokenizeResponse',
|
||||
'DetokenizeResponse',
|
||||
'RequestError',
|
||||
'ResponseError',
|
||||
]
|
||||
@@ -49,3 +53,5 @@ list = _client.list
|
||||
copy = _client.copy
|
||||
show = _client.show
|
||||
ps = _client.ps
|
||||
tokenize = _client.tokenize
|
||||
detokenize = _client.detokenize
|
||||
|
||||
+62
-13
@@ -23,6 +23,8 @@ from typing import (
|
||||
|
||||
import sys
|
||||
|
||||
from pydantic.json_schema import JsonSchemaValue
|
||||
|
||||
|
||||
from ollama._utils import convert_function_to_tool
|
||||
|
||||
@@ -46,6 +48,8 @@ from ollama._types import (
|
||||
CreateRequest,
|
||||
CopyRequest,
|
||||
DeleteRequest,
|
||||
DetokenizeRequest,
|
||||
DetokenizeResponse,
|
||||
EmbedRequest,
|
||||
EmbedResponse,
|
||||
EmbeddingsRequest,
|
||||
@@ -65,6 +69,8 @@ from ollama._types import (
|
||||
ShowRequest,
|
||||
ShowResponse,
|
||||
StatusResponse,
|
||||
TokenizeRequest,
|
||||
TokenizeResponse,
|
||||
Tool,
|
||||
)
|
||||
|
||||
@@ -186,7 +192,7 @@ class Client(BaseClient):
|
||||
context: Optional[Sequence[int]] = None,
|
||||
stream: Literal[False] = False,
|
||||
raw: bool = False,
|
||||
format: Optional[Literal['', 'json']] = None,
|
||||
format: Optional[Union[Literal['', 'json'], JsonSchemaValue]] = None,
|
||||
images: Optional[Sequence[Union[str, bytes]]] = None,
|
||||
options: Optional[Union[Mapping[str, Any], Options]] = None,
|
||||
keep_alive: Optional[Union[float, str]] = None,
|
||||
@@ -204,7 +210,7 @@ class Client(BaseClient):
|
||||
context: Optional[Sequence[int]] = None,
|
||||
stream: Literal[True] = True,
|
||||
raw: bool = False,
|
||||
format: Optional[Literal['', 'json']] = None,
|
||||
format: Optional[Union[Literal['', 'json'], JsonSchemaValue]] = None,
|
||||
images: Optional[Sequence[Union[str, bytes]]] = None,
|
||||
options: Optional[Union[Mapping[str, Any], Options]] = None,
|
||||
keep_alive: Optional[Union[float, str]] = None,
|
||||
@@ -221,7 +227,7 @@ class Client(BaseClient):
|
||||
context: Optional[Sequence[int]] = None,
|
||||
stream: bool = False,
|
||||
raw: Optional[bool] = None,
|
||||
format: Optional[Literal['', 'json']] = None,
|
||||
format: Optional[Union[Literal['', 'json'], JsonSchemaValue]] = None,
|
||||
images: Optional[Sequence[Union[str, bytes]]] = None,
|
||||
options: Optional[Union[Mapping[str, Any], Options]] = None,
|
||||
keep_alive: Optional[Union[float, str]] = None,
|
||||
@@ -265,7 +271,7 @@ class Client(BaseClient):
|
||||
*,
|
||||
tools: Optional[Sequence[Union[Mapping[str, Any], Tool, Callable]]] = None,
|
||||
stream: Literal[False] = False,
|
||||
format: Optional[Literal['', 'json']] = None,
|
||||
format: Optional[Union[Literal['', 'json'], JsonSchemaValue]] = None,
|
||||
options: Optional[Union[Mapping[str, Any], Options]] = None,
|
||||
keep_alive: Optional[Union[float, str]] = None,
|
||||
) -> ChatResponse: ...
|
||||
@@ -278,7 +284,7 @@ class Client(BaseClient):
|
||||
*,
|
||||
tools: Optional[Sequence[Union[Mapping[str, Any], Tool, Callable]]] = None,
|
||||
stream: Literal[True] = True,
|
||||
format: Optional[Literal['', 'json']] = None,
|
||||
format: Optional[Union[Literal['', 'json'], JsonSchemaValue]] = None,
|
||||
options: Optional[Union[Mapping[str, Any], Options]] = None,
|
||||
keep_alive: Optional[Union[float, str]] = None,
|
||||
) -> Iterator[ChatResponse]: ...
|
||||
@@ -290,7 +296,7 @@ class Client(BaseClient):
|
||||
*,
|
||||
tools: Optional[Sequence[Union[Mapping[str, Any], Tool, Callable]]] = None,
|
||||
stream: bool = False,
|
||||
format: Optional[Literal['', 'json']] = None,
|
||||
format: Optional[Union[Literal['', 'json'], JsonSchemaValue]] = None,
|
||||
options: Optional[Union[Mapping[str, Any], Options]] = None,
|
||||
keep_alive: Optional[Union[float, str]] = None,
|
||||
) -> Union[ChatResponse, Iterator[ChatResponse]]:
|
||||
@@ -327,7 +333,6 @@ class Client(BaseClient):
|
||||
|
||||
Returns `ChatResponse` if `stream` is `False`, otherwise returns a `ChatResponse` generator.
|
||||
"""
|
||||
|
||||
return self._request(
|
||||
ChatResponse,
|
||||
'POST',
|
||||
@@ -610,6 +615,28 @@ class Client(BaseClient):
|
||||
'/api/ps',
|
||||
)
|
||||
|
||||
def tokenize(self, model: str, text: str) -> TokenizeResponse:
|
||||
return self._request(
|
||||
TokenizeResponse,
|
||||
'POST',
|
||||
'/api/tokenize',
|
||||
json=TokenizeRequest(
|
||||
model=model,
|
||||
text=text,
|
||||
).model_dump(exclude_none=True),
|
||||
)
|
||||
|
||||
def detokenize(self, model: str, tokens: Sequence[int]) -> DetokenizeResponse:
|
||||
return self._request(
|
||||
DetokenizeResponse,
|
||||
'POST',
|
||||
'/api/detokenize',
|
||||
json=DetokenizeRequest(
|
||||
model=model,
|
||||
tokens=tokens,
|
||||
).model_dump(exclude_none=True),
|
||||
)
|
||||
|
||||
|
||||
class AsyncClient(BaseClient):
|
||||
def __init__(self, host: Optional[str] = None, **kwargs) -> None:
|
||||
@@ -689,7 +716,7 @@ class AsyncClient(BaseClient):
|
||||
context: Optional[Sequence[int]] = None,
|
||||
stream: Literal[False] = False,
|
||||
raw: bool = False,
|
||||
format: Optional[Literal['', 'json']] = None,
|
||||
format: Optional[Union[Literal['', 'json'], JsonSchemaValue]] = None,
|
||||
images: Optional[Sequence[Union[str, bytes]]] = None,
|
||||
options: Optional[Union[Mapping[str, Any], Options]] = None,
|
||||
keep_alive: Optional[Union[float, str]] = None,
|
||||
@@ -707,7 +734,7 @@ class AsyncClient(BaseClient):
|
||||
context: Optional[Sequence[int]] = None,
|
||||
stream: Literal[True] = True,
|
||||
raw: bool = False,
|
||||
format: Optional[Literal['', 'json']] = None,
|
||||
format: Optional[Union[Literal['', 'json'], JsonSchemaValue]] = None,
|
||||
images: Optional[Sequence[Union[str, bytes]]] = None,
|
||||
options: Optional[Union[Mapping[str, Any], Options]] = None,
|
||||
keep_alive: Optional[Union[float, str]] = None,
|
||||
@@ -724,7 +751,7 @@ class AsyncClient(BaseClient):
|
||||
context: Optional[Sequence[int]] = None,
|
||||
stream: bool = False,
|
||||
raw: Optional[bool] = None,
|
||||
format: Optional[Literal['', 'json']] = None,
|
||||
format: Optional[Union[Literal['', 'json'], JsonSchemaValue]] = None,
|
||||
images: Optional[Sequence[Union[str, bytes]]] = None,
|
||||
options: Optional[Union[Mapping[str, Any], Options]] = None,
|
||||
keep_alive: Optional[Union[float, str]] = None,
|
||||
@@ -767,7 +794,7 @@ class AsyncClient(BaseClient):
|
||||
*,
|
||||
tools: Optional[Sequence[Union[Mapping[str, Any], Tool, Callable]]] = None,
|
||||
stream: Literal[False] = False,
|
||||
format: Optional[Literal['', 'json']] = None,
|
||||
format: Optional[Union[Literal['', 'json'], JsonSchemaValue]] = None,
|
||||
options: Optional[Union[Mapping[str, Any], Options]] = None,
|
||||
keep_alive: Optional[Union[float, str]] = None,
|
||||
) -> ChatResponse: ...
|
||||
@@ -780,7 +807,7 @@ class AsyncClient(BaseClient):
|
||||
*,
|
||||
tools: Optional[Sequence[Union[Mapping[str, Any], Tool, Callable]]] = None,
|
||||
stream: Literal[True] = True,
|
||||
format: Optional[Literal['', 'json']] = None,
|
||||
format: Optional[Union[Literal['', 'json'], JsonSchemaValue]] = None,
|
||||
options: Optional[Union[Mapping[str, Any], Options]] = None,
|
||||
keep_alive: Optional[Union[float, str]] = None,
|
||||
) -> AsyncIterator[ChatResponse]: ...
|
||||
@@ -792,7 +819,7 @@ class AsyncClient(BaseClient):
|
||||
*,
|
||||
tools: Optional[Sequence[Union[Mapping[str, Any], Tool, Callable]]] = None,
|
||||
stream: bool = False,
|
||||
format: Optional[Literal['', 'json']] = None,
|
||||
format: Optional[Union[Literal['', 'json'], JsonSchemaValue]] = None,
|
||||
options: Optional[Union[Mapping[str, Any], Options]] = None,
|
||||
keep_alive: Optional[Union[float, str]] = None,
|
||||
) -> Union[ChatResponse, AsyncIterator[ChatResponse]]:
|
||||
@@ -1119,6 +1146,28 @@ class AsyncClient(BaseClient):
|
||||
'/api/ps',
|
||||
)
|
||||
|
||||
async def tokenize(self, model: str, text: str) -> TokenizeResponse:
|
||||
return await self._request(
|
||||
TokenizeResponse,
|
||||
'POST',
|
||||
'/api/tokenize',
|
||||
json=TokenizeRequest(
|
||||
model=model,
|
||||
text=text,
|
||||
).model_dump(exclude_none=True),
|
||||
)
|
||||
|
||||
async def detokenize(self, model: str, tokens: Sequence[int]) -> DetokenizeResponse:
|
||||
return await self._request(
|
||||
DetokenizeResponse,
|
||||
'POST',
|
||||
'/api/detokenize',
|
||||
json=DetokenizeRequest(
|
||||
model=model,
|
||||
tokens=tokens,
|
||||
).model_dump(exclude_none=True),
|
||||
)
|
||||
|
||||
|
||||
def _copy_messages(messages: Optional[Sequence[Union[Mapping[str, Any], Message]]]) -> Iterator[Message]:
|
||||
for message in messages or []:
|
||||
|
||||
+20
-1
@@ -4,6 +4,7 @@ from pathlib import Path
|
||||
from datetime import datetime
|
||||
from typing import Any, Mapping, Optional, Union, Sequence
|
||||
|
||||
from pydantic.json_schema import JsonSchemaValue
|
||||
from typing_extensions import Annotated, Literal
|
||||
|
||||
from pydantic import (
|
||||
@@ -150,7 +151,7 @@ class BaseGenerateRequest(BaseStreamableRequest):
|
||||
options: Optional[Union[Mapping[str, Any], Options]] = None
|
||||
'Options to use for the request.'
|
||||
|
||||
format: Optional[Literal['', 'json']] = None
|
||||
format: Optional[Union[Literal['', 'json'], JsonSchemaValue]] = None
|
||||
'Format of the response.'
|
||||
|
||||
keep_alive: Optional[Union[float, str]] = None
|
||||
@@ -493,6 +494,24 @@ class ProcessResponse(SubscriptableBaseModel):
|
||||
models: Sequence[Model]
|
||||
|
||||
|
||||
class TokenizeRequest(BaseRequest):
|
||||
model: str
|
||||
text: str
|
||||
|
||||
|
||||
class TokenizeResponse(BaseGenerateResponse):
|
||||
tokens: Sequence[int]
|
||||
|
||||
|
||||
class DetokenizeRequest(BaseRequest):
|
||||
model: str
|
||||
tokens: Sequence[int]
|
||||
|
||||
|
||||
class DetokenizeResponse(BaseGenerateResponse):
|
||||
text: str
|
||||
|
||||
|
||||
class RequestError(Exception):
|
||||
"""
|
||||
Common class for request errors.
|
||||
|
||||
Generated
+3
-3
@@ -559,13 +559,13 @@ typing-extensions = ">=4.6.0,<4.7.0 || >4.7.0"
|
||||
|
||||
[[package]]
|
||||
name = "pytest"
|
||||
version = "8.3.3"
|
||||
version = "8.3.4"
|
||||
description = "pytest: simple powerful testing with Python"
|
||||
optional = false
|
||||
python-versions = ">=3.8"
|
||||
files = [
|
||||
{file = "pytest-8.3.3-py3-none-any.whl", hash = "sha256:a6853c7375b2663155079443d2e45de913a911a11d669df02a50814944db57b2"},
|
||||
{file = "pytest-8.3.3.tar.gz", hash = "sha256:70b98107bd648308a7952b06e6ca9a50bc660be218d53c257cc1fc94fda10181"},
|
||||
{file = "pytest-8.3.4-py3-none-any.whl", hash = "sha256:50e16d954148559c9a74109af1eaf0c945ba2d8f30f0a3d3335edde19788b6f6"},
|
||||
{file = "pytest-8.3.4.tar.gz", hash = "sha256:965370d062bce11e73868e0335abac31b4d3de0e82f4007408d242b4f8610761"},
|
||||
]
|
||||
|
||||
[package.dependencies]
|
||||
|
||||
+249
-1
@@ -1,7 +1,7 @@
|
||||
import os
|
||||
import io
|
||||
import json
|
||||
from pydantic import ValidationError
|
||||
from pydantic import ValidationError, BaseModel
|
||||
import pytest
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
@@ -122,6 +122,128 @@ def test_client_chat_images(httpserver: HTTPServer):
|
||||
assert response['message']['content'] == "I don't know."
|
||||
|
||||
|
||||
def test_client_chat_format_json(httpserver: HTTPServer):
|
||||
httpserver.expect_ordered_request(
|
||||
'/api/chat',
|
||||
method='POST',
|
||||
json={
|
||||
'model': 'dummy',
|
||||
'messages': [{'role': 'user', 'content': 'Why is the sky blue?'}],
|
||||
'tools': [],
|
||||
'format': 'json',
|
||||
'stream': False,
|
||||
},
|
||||
).respond_with_json(
|
||||
{
|
||||
'model': 'dummy',
|
||||
'message': {
|
||||
'role': 'assistant',
|
||||
'content': '{"answer": "Because of Rayleigh scattering"}',
|
||||
},
|
||||
}
|
||||
)
|
||||
|
||||
client = Client(httpserver.url_for('/'))
|
||||
response = client.chat('dummy', messages=[{'role': 'user', 'content': 'Why is the sky blue?'}], format='json')
|
||||
assert response['model'] == 'dummy'
|
||||
assert response['message']['role'] == 'assistant'
|
||||
assert response['message']['content'] == '{"answer": "Because of Rayleigh scattering"}'
|
||||
|
||||
|
||||
def test_client_chat_format_pydantic(httpserver: HTTPServer):
|
||||
class ResponseFormat(BaseModel):
|
||||
answer: str
|
||||
confidence: float
|
||||
|
||||
httpserver.expect_ordered_request(
|
||||
'/api/chat',
|
||||
method='POST',
|
||||
json={
|
||||
'model': 'dummy',
|
||||
'messages': [{'role': 'user', 'content': 'Why is the sky blue?'}],
|
||||
'tools': [],
|
||||
'format': {'title': 'ResponseFormat', 'type': 'object', 'properties': {'answer': {'title': 'Answer', 'type': 'string'}, 'confidence': {'title': 'Confidence', 'type': 'number'}}, 'required': ['answer', 'confidence']},
|
||||
'stream': False,
|
||||
},
|
||||
).respond_with_json(
|
||||
{
|
||||
'model': 'dummy',
|
||||
'message': {
|
||||
'role': 'assistant',
|
||||
'content': '{"answer": "Because of Rayleigh scattering", "confidence": 0.95}',
|
||||
},
|
||||
}
|
||||
)
|
||||
|
||||
client = Client(httpserver.url_for('/'))
|
||||
response = client.chat('dummy', messages=[{'role': 'user', 'content': 'Why is the sky blue?'}], format=ResponseFormat.model_json_schema())
|
||||
assert response['model'] == 'dummy'
|
||||
assert response['message']['role'] == 'assistant'
|
||||
assert response['message']['content'] == '{"answer": "Because of Rayleigh scattering", "confidence": 0.95}'
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_client_chat_format_json(httpserver: HTTPServer):
|
||||
httpserver.expect_ordered_request(
|
||||
'/api/chat',
|
||||
method='POST',
|
||||
json={
|
||||
'model': 'dummy',
|
||||
'messages': [{'role': 'user', 'content': 'Why is the sky blue?'}],
|
||||
'tools': [],
|
||||
'format': 'json',
|
||||
'stream': False,
|
||||
},
|
||||
).respond_with_json(
|
||||
{
|
||||
'model': 'dummy',
|
||||
'message': {
|
||||
'role': 'assistant',
|
||||
'content': '{"answer": "Because of Rayleigh scattering"}',
|
||||
},
|
||||
}
|
||||
)
|
||||
|
||||
client = AsyncClient(httpserver.url_for('/'))
|
||||
response = await client.chat('dummy', messages=[{'role': 'user', 'content': 'Why is the sky blue?'}], format='json')
|
||||
assert response['model'] == 'dummy'
|
||||
assert response['message']['role'] == 'assistant'
|
||||
assert response['message']['content'] == '{"answer": "Because of Rayleigh scattering"}'
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_client_chat_format_pydantic(httpserver: HTTPServer):
|
||||
class ResponseFormat(BaseModel):
|
||||
answer: str
|
||||
confidence: float
|
||||
|
||||
httpserver.expect_ordered_request(
|
||||
'/api/chat',
|
||||
method='POST',
|
||||
json={
|
||||
'model': 'dummy',
|
||||
'messages': [{'role': 'user', 'content': 'Why is the sky blue?'}],
|
||||
'tools': [],
|
||||
'format': {'title': 'ResponseFormat', 'type': 'object', 'properties': {'answer': {'title': 'Answer', 'type': 'string'}, 'confidence': {'title': 'Confidence', 'type': 'number'}}, 'required': ['answer', 'confidence']},
|
||||
'stream': False,
|
||||
},
|
||||
).respond_with_json(
|
||||
{
|
||||
'model': 'dummy',
|
||||
'message': {
|
||||
'role': 'assistant',
|
||||
'content': '{"answer": "Because of Rayleigh scattering", "confidence": 0.95}',
|
||||
},
|
||||
}
|
||||
)
|
||||
|
||||
client = AsyncClient(httpserver.url_for('/'))
|
||||
response = await client.chat('dummy', messages=[{'role': 'user', 'content': 'Why is the sky blue?'}], format=ResponseFormat.model_json_schema())
|
||||
assert response['model'] == 'dummy'
|
||||
assert response['message']['role'] == 'assistant'
|
||||
assert response['message']['content'] == '{"answer": "Because of Rayleigh scattering", "confidence": 0.95}'
|
||||
|
||||
|
||||
def test_client_generate(httpserver: HTTPServer):
|
||||
httpserver.expect_ordered_request(
|
||||
'/api/generate',
|
||||
@@ -205,6 +327,108 @@ def test_client_generate_images(httpserver: HTTPServer):
|
||||
assert response['response'] == 'Because it is.'
|
||||
|
||||
|
||||
def test_client_generate_format_json(httpserver: HTTPServer):
|
||||
httpserver.expect_ordered_request(
|
||||
'/api/generate',
|
||||
method='POST',
|
||||
json={
|
||||
'model': 'dummy',
|
||||
'prompt': 'Why is the sky blue?',
|
||||
'format': 'json',
|
||||
'stream': False,
|
||||
},
|
||||
).respond_with_json(
|
||||
{
|
||||
'model': 'dummy',
|
||||
'response': '{"answer": "Because of Rayleigh scattering"}',
|
||||
}
|
||||
)
|
||||
|
||||
client = Client(httpserver.url_for('/'))
|
||||
response = client.generate('dummy', 'Why is the sky blue?', format='json')
|
||||
assert response['model'] == 'dummy'
|
||||
assert response['response'] == '{"answer": "Because of Rayleigh scattering"}'
|
||||
|
||||
|
||||
def test_client_generate_format_pydantic(httpserver: HTTPServer):
|
||||
class ResponseFormat(BaseModel):
|
||||
answer: str
|
||||
confidence: float
|
||||
|
||||
httpserver.expect_ordered_request(
|
||||
'/api/generate',
|
||||
method='POST',
|
||||
json={
|
||||
'model': 'dummy',
|
||||
'prompt': 'Why is the sky blue?',
|
||||
'format': {'title': 'ResponseFormat', 'type': 'object', 'properties': {'answer': {'title': 'Answer', 'type': 'string'}, 'confidence': {'title': 'Confidence', 'type': 'number'}}, 'required': ['answer', 'confidence']},
|
||||
'stream': False,
|
||||
},
|
||||
).respond_with_json(
|
||||
{
|
||||
'model': 'dummy',
|
||||
'response': '{"answer": "Because of Rayleigh scattering", "confidence": 0.95}',
|
||||
}
|
||||
)
|
||||
|
||||
client = Client(httpserver.url_for('/'))
|
||||
response = client.generate('dummy', 'Why is the sky blue?', format=ResponseFormat.model_json_schema())
|
||||
assert response['model'] == 'dummy'
|
||||
assert response['response'] == '{"answer": "Because of Rayleigh scattering", "confidence": 0.95}'
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_client_generate_format_json(httpserver: HTTPServer):
|
||||
httpserver.expect_ordered_request(
|
||||
'/api/generate',
|
||||
method='POST',
|
||||
json={
|
||||
'model': 'dummy',
|
||||
'prompt': 'Why is the sky blue?',
|
||||
'format': 'json',
|
||||
'stream': False,
|
||||
},
|
||||
).respond_with_json(
|
||||
{
|
||||
'model': 'dummy',
|
||||
'response': '{"answer": "Because of Rayleigh scattering"}',
|
||||
}
|
||||
)
|
||||
|
||||
client = AsyncClient(httpserver.url_for('/'))
|
||||
response = await client.generate('dummy', 'Why is the sky blue?', format='json')
|
||||
assert response['model'] == 'dummy'
|
||||
assert response['response'] == '{"answer": "Because of Rayleigh scattering"}'
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_client_generate_format_pydantic(httpserver: HTTPServer):
|
||||
class ResponseFormat(BaseModel):
|
||||
answer: str
|
||||
confidence: float
|
||||
|
||||
httpserver.expect_ordered_request(
|
||||
'/api/generate',
|
||||
method='POST',
|
||||
json={
|
||||
'model': 'dummy',
|
||||
'prompt': 'Why is the sky blue?',
|
||||
'format': {'title': 'ResponseFormat', 'type': 'object', 'properties': {'answer': {'title': 'Answer', 'type': 'string'}, 'confidence': {'title': 'Confidence', 'type': 'number'}}, 'required': ['answer', 'confidence']},
|
||||
'stream': False,
|
||||
},
|
||||
).respond_with_json(
|
||||
{
|
||||
'model': 'dummy',
|
||||
'response': '{"answer": "Because of Rayleigh scattering", "confidence": 0.95}',
|
||||
}
|
||||
)
|
||||
|
||||
client = AsyncClient(httpserver.url_for('/'))
|
||||
response = await client.generate('dummy', 'Why is the sky blue?', format=ResponseFormat.model_json_schema())
|
||||
assert response['model'] == 'dummy'
|
||||
assert response['response'] == '{"answer": "Because of Rayleigh scattering", "confidence": 0.95}'
|
||||
|
||||
|
||||
def test_client_pull(httpserver: HTTPServer):
|
||||
httpserver.expect_ordered_request(
|
||||
'/api/pull',
|
||||
@@ -1036,3 +1260,27 @@ def test_tool_validation():
|
||||
with pytest.raises(ValidationError):
|
||||
invalid_tool = {'type': 'invalid_type', 'function': {'name': 'test'}}
|
||||
list(_copy_tools([invalid_tool]))
|
||||
|
||||
|
||||
def test_client_tokenize(httpserver: HTTPServer):
|
||||
httpserver.expect_ordered_request(
|
||||
'/api/tokenize',
|
||||
method='POST',
|
||||
json={'model': 'dummy', 'text': 'Hello world!'},
|
||||
).respond_with_json({'tokens': [1, 2, 3]})
|
||||
|
||||
client = Client(httpserver.url_for('/'))
|
||||
response = client.tokenize('dummy', 'Hello world!')
|
||||
assert response.tokens == [1, 2, 3]
|
||||
|
||||
|
||||
def test_client_detokenize(httpserver: HTTPServer):
|
||||
httpserver.expect_ordered_request(
|
||||
'/api/detokenize',
|
||||
method='POST',
|
||||
json={'model': 'dummy', 'tokens': [1, 2, 3]},
|
||||
).respond_with_json({'text': 'Hello world!'})
|
||||
|
||||
client = Client(httpserver.url_for('/'))
|
||||
response = client.detokenize('dummy', [1, 2, 3])
|
||||
assert response.text == 'Hello world!'
|
||||
|
||||
Reference in New Issue
Block a user