mirror of
https://github.com/microsoft/agent-framework.git
synced 2026-06-16 21:04:09 +08:00
Python : Ollama Connector for Agent Framework (#1104)
* Initial Commit for Olama Connector * Added Olama Sample * Add Sample & Fixed Open Telemetry * Fixed Spelling from Olama to Ollama * remove"opentelemetry-semantic-conventions-ai ~=0.4.13" since its handled in a different pr * Added Tool Calling * Finalizing test cases * Adjust samples to be more reliable * Update python/packages/ollama/agent_framework_ollama/_chat_client.py Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> * Update python/packages/ollama/pyproject.toml Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> * Update python/packages/ollama/tests/test_ollama_chat_client.py Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> * Update python/packages/ollama/agent_framework_ollama/_chat_client.py Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> * Improved Docstrings & Sample * Update python/packages/ollama/agent_framework_ollama/_chat_client.py Co-authored-by: Eduard van Valkenburg <eavanvalkenburg@users.noreply.github.com> * Integrate PR Feedback - Divided Streaming and Non-Streaming into independent Methods - Catch Ollama Validation Error - Add OTEL Provider Name - Checked Ollama Messages - Add Usage Statistics * Revert setting, so it can be none * Validate Message formatting between AF and Ollama * Catch Ollama Error and raise a ServiceResponse Error * Fix mypy error * remove .vscode comma * Add Reasoning support & adjust to new structure * Add Ollama Multimodality and Reasoning * Add test cases for reasoning * Add Tests for Error Handling in Ollama Client * Update python/samples/getting_started/multimodal_input/ollama_chat_multimodal.py Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> * Integrated Copilot Feedback * Implement first PR Feedback * Adjust Readme files for examples * Adjust argument passing via additional chat options * Implemented PR Feedback * Removing Ollama Package from Core and moving samples * Fix Link & Adding Samples to Main Sample Readme * Fixing Links in Readme * Moved Multimodal and Chat Example * Fixed Link in ChatClient to Ollama * Fix AgentFramework Links in Ollama Project * Fix observability breaking change --------- Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> Co-authored-by: Eduard van Valkenburg <eavanvalkenburg@users.noreply.github.com>
This commit is contained in:
Nico Möller
committed by
GitHub
Unverified
GitHub
parent
1dbf3fd5cf
commit
2f06fe557a
@@ -0,0 +1,21 @@
|
||||
MIT License
|
||||
|
||||
Copyright (c) Microsoft Corporation.
|
||||
|
||||
Permission is hereby granted, free of charge, to any person obtaining a copy
|
||||
of this software and associated documentation files (the "Software"), to deal
|
||||
in the Software without restriction, including without limitation the rights
|
||||
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
||||
copies of the Software, and to permit persons to whom the Software is
|
||||
furnished to do so, subject to the following conditions:
|
||||
|
||||
The above copyright notice and this permission notice shall be included in all
|
||||
copies or substantial portions of the Software.
|
||||
|
||||
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
||||
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
||||
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
||||
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
||||
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
||||
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
||||
SOFTWARE
|
||||
@@ -0,0 +1,13 @@
|
||||
# Get Started with Microsoft Agent Framework Ollama
|
||||
|
||||
Please install this package as the extra for `agent-framework`:
|
||||
|
||||
```bash
|
||||
pip install agent-framework-ollama --pre
|
||||
```
|
||||
|
||||
and see the [README](https://github.com/microsoft/agent-framework/tree/main/python/README.md) for more information.
|
||||
|
||||
# Run samples with the Ollama Conector
|
||||
|
||||
You can find samples how to run the connector under the [Getting_started] (./getting_started/README.md) folder
|
||||
@@ -0,0 +1,16 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import importlib.metadata
|
||||
|
||||
from ._chat_client import OllamaChatClient, OllamaSettings
|
||||
|
||||
try:
|
||||
__version__ = importlib.metadata.version(__name__)
|
||||
except importlib.metadata.PackageNotFoundError:
|
||||
__version__ = "0.0.0" # Fallback for development mode
|
||||
|
||||
__all__ = [
|
||||
"OllamaChatClient",
|
||||
"OllamaSettings",
|
||||
"__version__",
|
||||
]
|
||||
@@ -0,0 +1,315 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import json
|
||||
from collections.abc import (
|
||||
AsyncIterable,
|
||||
Callable,
|
||||
Mapping,
|
||||
MutableMapping,
|
||||
MutableSequence,
|
||||
Sequence,
|
||||
)
|
||||
from itertools import chain
|
||||
from typing import Any, ClassVar
|
||||
|
||||
from agent_framework import (
|
||||
AIFunction,
|
||||
BaseChatClient,
|
||||
ChatMessage,
|
||||
ChatOptions,
|
||||
ChatResponse,
|
||||
ChatResponseUpdate,
|
||||
Contents,
|
||||
DataContent,
|
||||
FunctionCallContent,
|
||||
FunctionResultContent,
|
||||
Role,
|
||||
TextContent,
|
||||
TextReasoningContent,
|
||||
ToolProtocol,
|
||||
UsageDetails,
|
||||
get_logger,
|
||||
use_chat_middleware,
|
||||
use_function_invocation,
|
||||
)
|
||||
from agent_framework._pydantic import AFBaseSettings
|
||||
from agent_framework.exceptions import (
|
||||
ServiceInitializationError,
|
||||
ServiceInvalidRequestError,
|
||||
ServiceResponseException,
|
||||
)
|
||||
from agent_framework.observability import use_instrumentation
|
||||
from ollama import AsyncClient
|
||||
|
||||
# Rename imported types to avoid naming conflicts with Agent Framework types
|
||||
from ollama._types import ChatResponse as OllamaChatResponse
|
||||
from ollama._types import Message as OllamaMessage
|
||||
from pydantic import ValidationError
|
||||
|
||||
|
||||
class OllamaSettings(AFBaseSettings):
|
||||
"""Ollama settings."""
|
||||
|
||||
env_prefix: ClassVar[str] = "OLLAMA_"
|
||||
|
||||
host: str | None = None
|
||||
model_id: str | None = None
|
||||
|
||||
|
||||
logger = get_logger("agent_framework.ollama")
|
||||
|
||||
|
||||
@use_function_invocation
|
||||
@use_instrumentation
|
||||
@use_chat_middleware
|
||||
class OllamaChatClient(BaseChatClient):
|
||||
"""Ollama Chat completion class."""
|
||||
|
||||
OTEL_PROVIDER_NAME: ClassVar[str] = "ollama"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
host: str | None = None,
|
||||
client: AsyncClient | None = None,
|
||||
model_id: str | None = None,
|
||||
env_file_path: str | None = None,
|
||||
env_file_encoding: str | None = None,
|
||||
**kwargs: Any,
|
||||
) -> None:
|
||||
"""Initialize an Ollama Chat client.
|
||||
|
||||
Keyword Args:
|
||||
host: Ollama server URL, if none `http://localhost:11434` is used.
|
||||
Can be set via the OLLAMA_HOST env variable.
|
||||
client: An optional Ollama Client instance. If not provided, a new instance will be created.
|
||||
model_id: The Ollama chat model ID to use. Can be set via the OLLAMA_MODEL_ID env variable.
|
||||
env_file_path: An optional path to a dotenv (.env) file to load environment variables from.
|
||||
env_file_encoding: The encoding to use when reading the dotenv (.env) file. Defaults to 'utf-8'.
|
||||
**kwargs: Additional keyword arguments passed to BaseChatClient.
|
||||
"""
|
||||
try:
|
||||
ollama_settings = OllamaSettings(
|
||||
host=host,
|
||||
model_id=model_id,
|
||||
env_file_encoding=env_file_encoding,
|
||||
env_file_path=env_file_path,
|
||||
)
|
||||
except ValidationError as ex:
|
||||
raise ServiceInitializationError("Failed to create Ollama settings.", ex) from ex
|
||||
|
||||
if ollama_settings.model_id is None:
|
||||
raise ServiceInitializationError(
|
||||
"Ollama chat model ID must be provided via model_id or OLLAMA_MODEL_ID environment variable."
|
||||
)
|
||||
|
||||
self.model_id = ollama_settings.model_id
|
||||
self.client = client or AsyncClient(host=ollama_settings.host)
|
||||
# Save Host URL for serialization with to_dict()
|
||||
self.host = str(self.client._client.base_url)
|
||||
|
||||
super().__init__(**kwargs)
|
||||
|
||||
async def _inner_get_response(
|
||||
self,
|
||||
*,
|
||||
messages: MutableSequence[ChatMessage],
|
||||
chat_options: ChatOptions,
|
||||
**kwargs: Any,
|
||||
) -> ChatResponse:
|
||||
options_dict = self._prepare_options(messages, chat_options)
|
||||
|
||||
try:
|
||||
response: OllamaChatResponse = await self.client.chat( # type: ignore[misc]
|
||||
stream=False,
|
||||
**options_dict,
|
||||
**kwargs,
|
||||
)
|
||||
except Exception as ex:
|
||||
raise ServiceResponseException(f"Ollama chat request failed : {ex}", ex) from ex
|
||||
|
||||
return self._ollama_response_to_agent_framework_response(response)
|
||||
|
||||
async def _inner_get_streaming_response(
|
||||
self,
|
||||
*,
|
||||
messages: MutableSequence[ChatMessage],
|
||||
chat_options: ChatOptions,
|
||||
**kwargs: Any,
|
||||
) -> AsyncIterable[ChatResponseUpdate]:
|
||||
options_dict = self._prepare_options(messages, chat_options)
|
||||
|
||||
try:
|
||||
response_object: AsyncIterable[OllamaChatResponse] = await self.client.chat( # type: ignore[misc]
|
||||
stream=True,
|
||||
**options_dict,
|
||||
**kwargs,
|
||||
)
|
||||
except Exception as ex:
|
||||
raise ServiceResponseException(f"Ollama streaming chat request failed : {ex}", ex) from ex
|
||||
|
||||
async for part in response_object:
|
||||
yield self._ollama_streaming_response_to_agent_framework_response(part)
|
||||
|
||||
def _prepare_options(self, messages: MutableSequence[ChatMessage], chat_options: ChatOptions) -> dict[str, Any]:
|
||||
# Preprocess web search tool if it exists
|
||||
options_dict = chat_options.to_dict(exclude={"instructions", "type"})
|
||||
|
||||
# Promote additional_properties to the top level of options_dict
|
||||
additional_props = options_dict.pop("additional_properties", {})
|
||||
options_dict.update(additional_props)
|
||||
|
||||
# Prepare Messages from Agent Framework format to Ollama format
|
||||
if messages and "messages" not in options_dict:
|
||||
options_dict["messages"] = self._prepare_chat_history_for_request(messages)
|
||||
if "messages" not in options_dict:
|
||||
raise ServiceInvalidRequestError("Messages are required for chat completions")
|
||||
|
||||
# Prepare Tools from Agent Framework format to Json Schema format
|
||||
if chat_options.tools:
|
||||
options_dict["tools"] = self._chat_to_tool_spec(chat_options.tools)
|
||||
|
||||
# Currently Ollama only supports auto tool choice
|
||||
if chat_options.tool_choice == "required":
|
||||
raise ServiceInvalidRequestError("Ollama does not support required tool choice.")
|
||||
# Always auto: remove tool_choice since Ollama does not expose configuration to force or disable tools.
|
||||
if "tool_choice" in options_dict:
|
||||
del options_dict["tool_choice"]
|
||||
|
||||
# Rename model_id to model for Ollama API, if no model is provided use the one from client initialization
|
||||
if "model_id" in options_dict:
|
||||
options_dict["model"] = options_dict.pop("model_id")
|
||||
|
||||
if "model_id" not in options_dict:
|
||||
options_dict["model"] = self.model_id
|
||||
|
||||
return options_dict
|
||||
|
||||
def _prepare_chat_history_for_request(self, messages: MutableSequence[ChatMessage]) -> list[OllamaMessage]:
|
||||
ollama_messages = [self._agent_framework_message_to_ollama_message(msg) for msg in messages]
|
||||
# Flatten the list of lists into a single list
|
||||
return list(chain.from_iterable(ollama_messages))
|
||||
|
||||
def _agent_framework_message_to_ollama_message(self, message: ChatMessage) -> list[OllamaMessage]:
|
||||
message_converters: dict[str, Callable[[ChatMessage], list[OllamaMessage]]] = {
|
||||
Role.SYSTEM.value: self._format_system_message,
|
||||
Role.USER.value: self._format_user_message,
|
||||
Role.ASSISTANT.value: self._format_assistant_message,
|
||||
Role.TOOL.value: self._format_tool_message,
|
||||
}
|
||||
return message_converters[message.role.value](message)
|
||||
|
||||
def _format_system_message(self, message: ChatMessage) -> list[OllamaMessage]:
|
||||
return [OllamaMessage(role="system", content=message.text)]
|
||||
|
||||
def _format_user_message(self, message: ChatMessage) -> list[OllamaMessage]:
|
||||
if not any(isinstance(c, (DataContent, TextContent)) for c in message.contents) and not message.text:
|
||||
raise ServiceInvalidRequestError(
|
||||
"Ollama connector currently only supports user messages with TextContent or DataContent."
|
||||
)
|
||||
|
||||
if not any(isinstance(c, DataContent) for c in message.contents):
|
||||
return [OllamaMessage(role="user", content=message.text)]
|
||||
|
||||
user_message = OllamaMessage(role="user", content=message.text)
|
||||
data_contents = [c for c in message.contents if isinstance(c, DataContent)]
|
||||
if data_contents:
|
||||
if not any(c.has_top_level_media_type("image") for c in data_contents):
|
||||
raise ServiceInvalidRequestError("Only image data content is supported for user messages in Ollama.")
|
||||
# Ollama expects base64 strings without prefix
|
||||
user_message["images"] = [c.uri.split(",")[1] for c in data_contents]
|
||||
return [user_message]
|
||||
|
||||
def _format_assistant_message(self, message: ChatMessage) -> list[OllamaMessage]:
|
||||
text_content = message.text
|
||||
reasoning_contents = "".join(c.text for c in message.contents if isinstance(c, TextReasoningContent))
|
||||
|
||||
assistant_message = OllamaMessage(role="assistant", content=text_content, thinking=reasoning_contents)
|
||||
|
||||
tool_calls = [item for item in message.contents if isinstance(item, FunctionCallContent)]
|
||||
if tool_calls:
|
||||
assistant_message["tool_calls"] = [
|
||||
{
|
||||
"function": {
|
||||
"call_id": tool_call.call_id,
|
||||
"name": tool_call.name,
|
||||
"arguments": tool_call.arguments
|
||||
if isinstance(tool_call.arguments, Mapping)
|
||||
else json.loads(tool_call.arguments or "{}"),
|
||||
}
|
||||
}
|
||||
for tool_call in tool_calls
|
||||
]
|
||||
return [assistant_message]
|
||||
|
||||
def _format_tool_message(self, message: ChatMessage) -> list[OllamaMessage]:
|
||||
# Ollama does not support multiple tool results in a single message, so we create a separate
|
||||
return [
|
||||
OllamaMessage(role="tool", content=str(item.result), tool_name=item.call_id)
|
||||
for item in message.contents
|
||||
if isinstance(item, FunctionResultContent)
|
||||
]
|
||||
|
||||
def _ollama_response_to_agent_framework_content(self, response: OllamaChatResponse) -> list[Contents]:
|
||||
contents: list[Contents] = []
|
||||
if response.message.thinking:
|
||||
contents.append(TextReasoningContent(text=response.message.thinking))
|
||||
if response.message.content:
|
||||
contents.append(TextContent(text=response.message.content))
|
||||
if response.message.tool_calls:
|
||||
tool_calls = self._parse_ollama_tool_calls(response.message.tool_calls)
|
||||
contents.extend(tool_calls)
|
||||
return contents
|
||||
|
||||
def _ollama_streaming_response_to_agent_framework_response(
|
||||
self, response: OllamaChatResponse
|
||||
) -> ChatResponseUpdate:
|
||||
contents = self._ollama_response_to_agent_framework_content(response)
|
||||
return ChatResponseUpdate(
|
||||
contents=contents,
|
||||
role=Role.ASSISTANT,
|
||||
ai_model_id=response.model,
|
||||
created_at=response.created_at,
|
||||
)
|
||||
|
||||
def _ollama_response_to_agent_framework_response(self, response: OllamaChatResponse) -> ChatResponse:
|
||||
contents = self._ollama_response_to_agent_framework_content(response)
|
||||
|
||||
return ChatResponse(
|
||||
messages=[ChatMessage(role=Role.ASSISTANT, contents=contents)],
|
||||
model_id=response.model,
|
||||
created_at=response.created_at,
|
||||
usage_details=UsageDetails(
|
||||
input_token_count=response.prompt_eval_count,
|
||||
output_token_count=response.eval_count,
|
||||
),
|
||||
)
|
||||
|
||||
def _parse_ollama_tool_calls(self, tool_calls: Sequence[OllamaMessage.ToolCall]) -> list[Contents]:
|
||||
resp: list[Contents] = []
|
||||
for tool in tool_calls:
|
||||
fcc = FunctionCallContent(
|
||||
call_id=tool.function.name, # Use name of function as call ID since Ollama doesn't provide a call ID
|
||||
name=tool.function.name,
|
||||
arguments=tool.function.arguments if isinstance(tool.function.arguments, dict) else "",
|
||||
raw_representation=tool.function,
|
||||
)
|
||||
resp.append(fcc)
|
||||
return resp
|
||||
|
||||
def _chat_to_tool_spec(self, tools: list[ToolProtocol | MutableMapping[str, Any]]) -> list[dict[str, Any]]:
|
||||
chat_tools: list[dict[str, Any]] = []
|
||||
for tool in tools:
|
||||
if isinstance(tool, ToolProtocol):
|
||||
match tool:
|
||||
case AIFunction():
|
||||
chat_tools.append(tool.to_json_schema_spec())
|
||||
case _:
|
||||
raise ServiceInvalidRequestError(
|
||||
"Unsupported tool type '"
|
||||
f"{type(tool).__name__}"
|
||||
"' for Ollama client. Supported tool types: AIFunction."
|
||||
)
|
||||
else:
|
||||
chat_tools.append(tool if isinstance(tool, dict) else dict(tool))
|
||||
return chat_tools
|
||||
@@ -0,0 +1,38 @@
|
||||
# Ollama Examples
|
||||
|
||||
This folder contains examples demonstrating how to use Ollama models with the Agent Framework.
|
||||
|
||||
## Prerequisites
|
||||
|
||||
1. **Install Ollama**: Download and install Ollama from [ollama.com](https://ollama.com/)
|
||||
2. **Start Ollama**: Ensure Ollama is running on your local machine
|
||||
3. **Pull a model**: Run `ollama pull mistral` (or any other model you prefer)
|
||||
- For function calling examples, use models that support tool calling like `mistral` or `qwen2.5`
|
||||
- For reasoning examples, use models that support reasoning like `qwen2.5:8b`
|
||||
- For Multimodality you can use models like `gemma3:4b`
|
||||
|
||||
> **Note**: Not all models support all features. Function calling and reasoning capabilities depend on the specific model you're using.
|
||||
|
||||
## Examples
|
||||
|
||||
| File | Description |
|
||||
|------|-------------|
|
||||
| [`ollama_agent_basic.py`](ollama_agent_basic.py) | Demonstrates basic Ollama agent usage with the native Ollama Chat Client. Shows both streaming and non-streaming responses with tool calling capabilities. |
|
||||
| [`ollama_agent_reasoning.py`](ollama_agent_reasoning.py) | Demonstrates Ollama agent with reasoning capabilities using the native Ollama Chat Client. Shows how to enable thinking/reasoning mode. |
|
||||
| [`ollama_chat_client.py`](ollama_chat_client.py) | Ollama Chat Client with native Ollama Chat Client |
|
||||
| [`ollama_chat_multimodal.py`](ollama_chat_multimodal.py) | Ollama Chat with multimodal native Ollama Chat Client |
|
||||
|
||||
## Configuration
|
||||
|
||||
The examples use environment variables for configuration. Set the appropriate variables based on which example you're running:
|
||||
|
||||
### For Native Ollama Examples (`ollama_agent_basic.py`, `ollama_agent_reasoning.py`)
|
||||
|
||||
Set the following environment variables:
|
||||
|
||||
- `OLLAMA_HOST`: The base URL for your Ollama server (optional, defaults to `http://localhost:11434`)
|
||||
- Example: `export OLLAMA_HOST="http://localhost:11434"`
|
||||
|
||||
- `OLLAMA_CHAT_MODEL_ID`: The model name to use
|
||||
- Example: `export OLLAMA_CHAT_MODEL_ID="qwen2.5:8b"`
|
||||
- Must be a model you have pulled with Ollama
|
||||
@@ -0,0 +1,68 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import asyncio
|
||||
from datetime import datetime
|
||||
|
||||
from agent_framework_ollama import OllamaChatClient
|
||||
|
||||
"""
|
||||
Ollama Agent Basic Example
|
||||
|
||||
This sample demonstrates implementing a Ollama agent with basic tool usage.
|
||||
|
||||
Ensure to install Ollama and have a model running locally before running the sample
|
||||
Not all Models support function calling, to test function calling try llama3.2 or qwen3:4b
|
||||
Set the model to use via the OLLAMA_CHAT_MODEL_ID environment variable or modify the code below.
|
||||
https://ollama.com/
|
||||
|
||||
"""
|
||||
|
||||
|
||||
def get_time(location: str) -> str:
|
||||
"""Get the current time."""
|
||||
return f"The current time in {location} is {datetime.now().strftime('%I:%M %p')}."
|
||||
|
||||
|
||||
async def non_streaming_example() -> None:
|
||||
"""Example of non-streaming response (get the complete result at once)."""
|
||||
print("=== Non-streaming Response Example ===")
|
||||
|
||||
agent = OllamaChatClient().create_agent(
|
||||
name="TimeAgent",
|
||||
instructions="You are a helpful time agent answer in one sentence.",
|
||||
tools=get_time,
|
||||
)
|
||||
|
||||
query = "What time is it in Seattle? Use a tool call"
|
||||
print(f"User: {query}")
|
||||
result = await agent.run(query)
|
||||
print(f"Result: {result}\n")
|
||||
|
||||
|
||||
async def streaming_example() -> None:
|
||||
"""Example of streaming response (get results as they are generated)."""
|
||||
print("=== Streaming Response Example ===")
|
||||
|
||||
agent = OllamaChatClient().create_agent(
|
||||
name="TimeAgent",
|
||||
instructions="You are a helpful time agent answer in one sentence.",
|
||||
tools=get_time,
|
||||
)
|
||||
query = "What time is it in San Francisco? Use a tool call"
|
||||
print(f"User: {query}")
|
||||
print("Agent: ", end="", flush=True)
|
||||
async for chunk in agent.run_stream(query):
|
||||
if chunk.text:
|
||||
print(chunk.text, end="", flush=True)
|
||||
print("\n")
|
||||
|
||||
|
||||
async def main() -> None:
|
||||
print("=== Basic Ollama Chat Client Agent Example ===")
|
||||
|
||||
await non_streaming_example()
|
||||
await streaming_example()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
@@ -0,0 +1,46 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import asyncio
|
||||
|
||||
from agent_framework import TextReasoningContent
|
||||
|
||||
from agent_framework_ollama import OllamaChatClient
|
||||
|
||||
"""
|
||||
Ollama Agent Reasoning Example
|
||||
|
||||
This sample demonstrates implementing a Ollama agent with reasoning.
|
||||
|
||||
Ensure to install Ollama and have a model running locally before running the sample
|
||||
Not all Models support reasoning, to test reasoning try qwen3:8b
|
||||
Set the model to use via the OLLAMA_CHAT_MODEL_ID environment variable or modify the code below.
|
||||
https://ollama.com/
|
||||
|
||||
"""
|
||||
|
||||
|
||||
async def reasoning_example() -> None:
|
||||
print("=== Response Reasoning Example ===")
|
||||
|
||||
agent = OllamaChatClient().create_agent(
|
||||
name="TimeAgent",
|
||||
instructions="You are a helpful agent answer in one sentence.",
|
||||
additional_chat_options={"think": True}, # Enable Reasoning on agent level
|
||||
)
|
||||
query = "Hey what is 3+4? Can you explain how you got to that answer?"
|
||||
print(f"User: {query}")
|
||||
# Enable Reasoning on per request level
|
||||
result = await agent.run(query)
|
||||
reasoning = "".join(c.text for c in result.messages[-1].contents if isinstance(c, TextReasoningContent))
|
||||
print(f"Reasoning: {reasoning}")
|
||||
print(f"Answer: {result}\n")
|
||||
|
||||
|
||||
async def main() -> None:
|
||||
print("=== Basic Ollama Chat Client Agent Reasoning ===")
|
||||
|
||||
await reasoning_example()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
@@ -0,0 +1,36 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import asyncio
|
||||
from datetime import datetime
|
||||
|
||||
from agent_framework_ollama import OllamaChatClient
|
||||
|
||||
# Ensure to install Ollama and have a model running locally before running the sample
|
||||
# Not all Models support function calling, to test function calling try llama3.2
|
||||
# Set the model to use via the OLLAMA_CHAT_MODEL_ID environment variable or modify the code below.
|
||||
# https://ollama.com/
|
||||
|
||||
|
||||
def get_time():
|
||||
"""Get the current time."""
|
||||
return f"The current time is {datetime.now().strftime('%I:%M %p')}."
|
||||
|
||||
|
||||
async def main() -> None:
|
||||
client = OllamaChatClient()
|
||||
message = "What time is it? Use a tool call"
|
||||
stream = False
|
||||
print(f"User: {message}")
|
||||
if stream:
|
||||
print("Assistant: ", end="")
|
||||
async for chunk in client.get_streaming_response(message, tools=get_time):
|
||||
if str(chunk):
|
||||
print(str(chunk), end="")
|
||||
print("")
|
||||
else:
|
||||
response = await client.get_response(message, tools=get_time)
|
||||
print(f"Assistant: {response}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
@@ -0,0 +1,54 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import asyncio
|
||||
|
||||
from agent_framework import ChatMessage, DataContent, Role, TextContent
|
||||
|
||||
from agent_framework_ollama import OllamaChatClient
|
||||
|
||||
"""
|
||||
Ollama Agent Multimodal Example
|
||||
|
||||
This sample demonstrates implementing a Ollama agent with multimodal input capabilities.
|
||||
|
||||
Ensure to install Ollama and have a model running locally before running the sample
|
||||
Not all Models support multimodal input, to test multimodal input try gemma3:4b
|
||||
Set the model to use via the OLLAMA_CHAT_MODEL_ID environment variable or modify the code below.
|
||||
https://ollama.com/
|
||||
|
||||
"""
|
||||
|
||||
|
||||
def create_sample_image() -> str:
|
||||
"""Create a simple 1x1 pixel PNG image for testing."""
|
||||
# This is a tiny red pixel in PNG format
|
||||
png_data = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8/5+hHgAHggJ/PchI7wAAAABJRU5ErkJggg=="
|
||||
return f"data:image/png;base64,{png_data}"
|
||||
|
||||
|
||||
async def test_image() -> None:
|
||||
"""Test image analysis with Ollama."""
|
||||
|
||||
client = OllamaChatClient()
|
||||
|
||||
image_uri = create_sample_image()
|
||||
|
||||
message = ChatMessage(
|
||||
role=Role.USER,
|
||||
contents=[
|
||||
TextContent(text="What's in this image?"),
|
||||
DataContent(uri=image_uri, media_type="image/png"),
|
||||
],
|
||||
)
|
||||
|
||||
response = await client.get_response(message)
|
||||
print(f"Image Response: {response}")
|
||||
|
||||
|
||||
async def main() -> None:
|
||||
print("=== Testing Ollama Multimodal ===")
|
||||
await test_image()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
@@ -0,0 +1,95 @@
|
||||
[project]
|
||||
name = "agent-framework-ollama"
|
||||
description = "Ollama integration for Microsoft Agent Framework."
|
||||
authors = [{ name = "Microsoft", email = "af-support@microsoft.com"}]
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.10"
|
||||
version = "0.1.0b1"
|
||||
license-files = ["LICENSE"]
|
||||
urls.homepage = "https://learn.microsoft.com/en-us/agent-framework/"
|
||||
urls.source = "https://github.com/microsoft/agent-framework/tree/main/python"
|
||||
urls.release_notes = "https://github.com/microsoft/agent-framework/releases?q=tag%3Apython-1&expanded=true"
|
||||
urls.issues = "https://github.com/microsoft/agent-framework/issues"
|
||||
classifiers = [
|
||||
"License :: OSI Approved :: MIT License",
|
||||
"Development Status :: 4 - Beta",
|
||||
"Intended Audience :: Developers",
|
||||
"Programming Language :: Python :: 3",
|
||||
"Programming Language :: Python :: 3.10",
|
||||
"Programming Language :: Python :: 3.11",
|
||||
"Programming Language :: Python :: 3.12",
|
||||
"Programming Language :: Python :: 3.13",
|
||||
"Framework :: Pydantic :: 2",
|
||||
"Typing :: Typed",
|
||||
]
|
||||
dependencies = [
|
||||
"agent-framework-core",
|
||||
"ollama >= 0.5.3",
|
||||
]
|
||||
|
||||
[tool.uv]
|
||||
prerelease = "if-necessary-or-explicit"
|
||||
environments = [
|
||||
"sys_platform == 'darwin'",
|
||||
"sys_platform == 'linux'",
|
||||
"sys_platform == 'win32'"
|
||||
]
|
||||
|
||||
[tool.uv-dynamic-versioning]
|
||||
fallback-version = "0.0.0"
|
||||
[tool.pytest.ini_options]
|
||||
testpaths = 'tests'
|
||||
addopts = "-ra -q -r fEX"
|
||||
asyncio_mode = "auto"
|
||||
asyncio_default_fixture_loop_scope = "function"
|
||||
filterwarnings = []
|
||||
timeout = 120
|
||||
|
||||
[tool.ruff]
|
||||
line-length = 120
|
||||
|
||||
[tool.ruff.lint]
|
||||
select = ["E", "F", "I", "N", "W"]
|
||||
|
||||
[tool.coverage.run]
|
||||
omit = [
|
||||
"**/__init__.py"
|
||||
]
|
||||
|
||||
[tool.pyright]
|
||||
extend = "../../pyproject.toml"
|
||||
exclude = ['tests']
|
||||
|
||||
[tool.mypy]
|
||||
plugins = ['pydantic.mypy']
|
||||
strict = true
|
||||
python_version = "3.10"
|
||||
ignore_missing_imports = true
|
||||
disallow_untyped_defs = true
|
||||
no_implicit_optional = true
|
||||
check_untyped_defs = true
|
||||
warn_return_any = true
|
||||
show_error_codes = true
|
||||
warn_unused_ignores = false
|
||||
disallow_incomplete_defs = true
|
||||
disallow_untyped_decorators = true
|
||||
disallow_any_unimported = true
|
||||
|
||||
[tool.bandit]
|
||||
targets = ["agent_framework_ollama"]
|
||||
exclude_dirs = ["tests"]
|
||||
|
||||
[tool.poe]
|
||||
executor.type = "uv"
|
||||
include = "../../shared_tasks.toml"
|
||||
[tool.poe.tasks]
|
||||
mypy = "mypy --config-file $POE_ROOT/pyproject.toml agent_framework_ollama"
|
||||
test = "pytest --cov=agent_framework_ollama --cov-report=term-missing:skip-covered tests"
|
||||
|
||||
[tool.uv.build-backend]
|
||||
module-name = "agent_framework_ollama"
|
||||
module-root = ""
|
||||
|
||||
[build-system]
|
||||
requires = ["uv_build>=0.8.2,<0.9.0"]
|
||||
build-backend = "uv_build"
|
||||
@@ -0,0 +1,47 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
from typing import Any
|
||||
|
||||
from agent_framework import ChatMessage
|
||||
from pytest import fixture
|
||||
|
||||
|
||||
# region: Connector Settings fixtures
|
||||
@fixture
|
||||
def exclude_list(request: Any) -> list[str]:
|
||||
"""Fixture that returns a list of environment variables to exclude."""
|
||||
return request.param if hasattr(request, "param") else []
|
||||
|
||||
|
||||
@fixture
|
||||
def override_env_param_dict(request: Any) -> dict[str, str]:
|
||||
"""Fixture that returns a dict of environment variables to override."""
|
||||
return request.param if hasattr(request, "param") else {}
|
||||
|
||||
|
||||
# These two fixtures are used for multiple things, also non-connector tests
|
||||
@fixture()
|
||||
def ollama_unit_test_env(monkeypatch, exclude_list, override_env_param_dict): # type: ignore
|
||||
"""Fixture to set environment variables for OllamaSettings."""
|
||||
|
||||
if exclude_list is None:
|
||||
exclude_list = []
|
||||
|
||||
if override_env_param_dict is None:
|
||||
override_env_param_dict = {}
|
||||
|
||||
env_vars = {"OLLAMA_HOST": "http://localhost:12345", "OLLAMA_MODEL_ID": "test"}
|
||||
|
||||
env_vars.update(override_env_param_dict) # type: ignore
|
||||
|
||||
for key, value in env_vars.items():
|
||||
if key in exclude_list:
|
||||
monkeypatch.delenv(key, raising=False) # type: ignore
|
||||
continue
|
||||
monkeypatch.setenv(key, value) # type: ignore
|
||||
|
||||
return env_vars
|
||||
|
||||
|
||||
@fixture
|
||||
def chat_history() -> list[ChatMessage]:
|
||||
return []
|
||||
@@ -0,0 +1,511 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import os
|
||||
from collections.abc import AsyncIterable
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
from agent_framework import (
|
||||
BaseChatClient,
|
||||
ChatMessage,
|
||||
ChatResponseUpdate,
|
||||
DataContent,
|
||||
FunctionCallContent,
|
||||
FunctionResultContent,
|
||||
HostedWebSearchTool,
|
||||
TextContent,
|
||||
TextReasoningContent,
|
||||
UriContent,
|
||||
chat_middleware,
|
||||
)
|
||||
from agent_framework.exceptions import (
|
||||
ServiceInitializationError,
|
||||
ServiceInvalidRequestError,
|
||||
ServiceResponseException,
|
||||
)
|
||||
from ollama import AsyncClient
|
||||
from ollama._types import ChatResponse as OllamaChatResponse
|
||||
from ollama._types import Message as OllamaMessage
|
||||
from openai import AsyncStream
|
||||
|
||||
from agent_framework_ollama import OllamaChatClient
|
||||
|
||||
# region Service Setup
|
||||
|
||||
skip_if_azure_integration_tests_disabled = pytest.mark.skipif(
|
||||
os.getenv("RUN_INTEGRATION_TESTS", "false").lower() != "true"
|
||||
or os.getenv("OLLAMA_MODEL_ID", "") in ("", "test-model"),
|
||||
reason="No real Ollama chat model provided; skipping integration tests."
|
||||
if os.getenv("RUN_INTEGRATION_TESTS", "false").lower() == "true"
|
||||
else "Integration tests are disabled.",
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_streaming_chat_completion_response() -> AsyncStream[OllamaChatResponse]:
|
||||
response = OllamaChatResponse(
|
||||
message=OllamaMessage(content="test", role="assistant"),
|
||||
model="test",
|
||||
)
|
||||
stream = MagicMock(spec=AsyncStream)
|
||||
stream.__aiter__.return_value = [response]
|
||||
return stream
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_streaming_chat_completion_response_reasoning() -> AsyncStream[OllamaChatResponse]:
|
||||
response = OllamaChatResponse(
|
||||
message=OllamaMessage(thinking="test", role="assistant"),
|
||||
model="test",
|
||||
)
|
||||
stream = MagicMock(spec=AsyncStream)
|
||||
stream.__aiter__.return_value = [response]
|
||||
return stream
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_chat_completion_response() -> OllamaChatResponse:
|
||||
return OllamaChatResponse(
|
||||
message=OllamaMessage(content="test", role="assistant"),
|
||||
model="test",
|
||||
eval_count=1,
|
||||
prompt_eval_count=1,
|
||||
created_at="2024-01-01T00:00:00Z",
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_chat_completion_response_reasoning() -> OllamaChatResponse:
|
||||
return OllamaChatResponse(
|
||||
message=OllamaMessage(thinking="test", role="assistant"),
|
||||
model="test",
|
||||
eval_count=1,
|
||||
prompt_eval_count=1,
|
||||
created_at="2024-01-01T00:00:00Z",
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_streaming_chat_completion_tool_call() -> AsyncStream[OllamaChatResponse]:
|
||||
ollama_tool_call = OllamaChatResponse(
|
||||
message=OllamaMessage(
|
||||
content="",
|
||||
role="assistant",
|
||||
tool_calls=[{"function": {"name": "hello_world", "arguments": {"arg1": "value1"}}}],
|
||||
),
|
||||
model="test",
|
||||
)
|
||||
stream = MagicMock(spec=AsyncStream)
|
||||
stream.__aiter__.return_value = [ollama_tool_call]
|
||||
return stream
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_chat_completion_tool_call() -> OllamaChatResponse:
|
||||
return OllamaChatResponse(
|
||||
message=OllamaMessage(
|
||||
content="",
|
||||
role="assistant",
|
||||
tool_calls=[{"function": {"name": "hello_world", "arguments": {"arg1": "value1"}}}],
|
||||
),
|
||||
model="test",
|
||||
created_at="2024-01-01T00:00:00Z",
|
||||
)
|
||||
|
||||
|
||||
def hello_world(arg1: str) -> str:
|
||||
return "Hello World"
|
||||
|
||||
|
||||
def test_init(ollama_unit_test_env: dict[str, str]) -> None:
|
||||
# Test successful initialization
|
||||
ollama_chat_client = OllamaChatClient()
|
||||
|
||||
assert ollama_chat_client.client is not None
|
||||
assert isinstance(ollama_chat_client.client, AsyncClient)
|
||||
assert ollama_chat_client.model_id == ollama_unit_test_env["OLLAMA_MODEL_ID"]
|
||||
assert isinstance(ollama_chat_client, BaseChatClient)
|
||||
|
||||
|
||||
def test_init_client(ollama_unit_test_env: dict[str, str]) -> None:
|
||||
# Test successful initialization with provided client
|
||||
test_client = MagicMock(spec=AsyncClient)
|
||||
# Mock underlying HTTP client's base_url
|
||||
test_client._client = MagicMock()
|
||||
test_client._client.base_url = ollama_unit_test_env["OLLAMA_MODEL_ID"]
|
||||
ollama_chat_client = OllamaChatClient(client=test_client)
|
||||
|
||||
assert ollama_chat_client.client is test_client
|
||||
assert ollama_chat_client.model_id == ollama_unit_test_env["OLLAMA_MODEL_ID"]
|
||||
assert isinstance(ollama_chat_client, BaseChatClient)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("exclude_list", [["OLLAMA_MODEL_ID"]], indirect=True)
|
||||
def test_with_invalid_settings(ollama_unit_test_env: dict[str, str]) -> None:
|
||||
with pytest.raises(ServiceInitializationError):
|
||||
OllamaChatClient(
|
||||
host="http://localhost:12345",
|
||||
model_id=None,
|
||||
env_file_path="test.env",
|
||||
)
|
||||
|
||||
|
||||
def test_serialize(ollama_unit_test_env: dict[str, str]) -> None:
|
||||
settings = {
|
||||
"host": ollama_unit_test_env["OLLAMA_HOST"],
|
||||
"model_id": ollama_unit_test_env["OLLAMA_MODEL_ID"],
|
||||
}
|
||||
|
||||
ollama_chat_client = OllamaChatClient.from_dict(settings)
|
||||
serialized = ollama_chat_client.to_dict()
|
||||
|
||||
assert isinstance(serialized, dict)
|
||||
assert serialized["host"] == ollama_unit_test_env["OLLAMA_HOST"]
|
||||
assert serialized["model_id"] == ollama_unit_test_env["OLLAMA_MODEL_ID"]
|
||||
|
||||
|
||||
def test_chat_middleware(ollama_unit_test_env: dict[str, str]) -> None:
|
||||
@chat_middleware
|
||||
async def sample_middleware(context, next):
|
||||
await next(context)
|
||||
|
||||
ollama_chat_client = OllamaChatClient(middleware=[sample_middleware])
|
||||
assert len(ollama_chat_client.middleware) == 1
|
||||
assert ollama_chat_client.middleware[0] == sample_middleware
|
||||
|
||||
|
||||
def test_additional_properties(ollama_unit_test_env: dict[str, str]) -> None:
|
||||
additional_properties = {
|
||||
"user_location": {
|
||||
"country": "US",
|
||||
"city": "Seattle",
|
||||
}
|
||||
}
|
||||
ollama_chat_client = OllamaChatClient(
|
||||
additional_properties=additional_properties,
|
||||
)
|
||||
assert ollama_chat_client.additional_properties == additional_properties
|
||||
|
||||
|
||||
# region CMC
|
||||
|
||||
|
||||
async def test_empty_messages() -> None:
|
||||
ollama_chat_client = OllamaChatClient(
|
||||
host="http://localhost:12345",
|
||||
model_id="test-model",
|
||||
)
|
||||
with pytest.raises(ServiceInvalidRequestError):
|
||||
await ollama_chat_client.get_response(messages=[])
|
||||
|
||||
|
||||
async def test_function_choice_required_argument() -> None:
|
||||
ollama_chat_client = OllamaChatClient(
|
||||
host="http://localhost:12345",
|
||||
model_id="test-model",
|
||||
)
|
||||
with pytest.raises(ServiceInvalidRequestError):
|
||||
await ollama_chat_client.get_response(
|
||||
messages=[ChatMessage(text="hello world", role="user")],
|
||||
tool_choice="required",
|
||||
tools=[hello_world],
|
||||
)
|
||||
|
||||
|
||||
@patch.object(AsyncClient, "chat", new_callable=AsyncMock)
|
||||
async def test_cmc(
|
||||
mock_chat: AsyncMock,
|
||||
ollama_unit_test_env: dict[str, str],
|
||||
chat_history: list[ChatMessage],
|
||||
mock_chat_completion_response: AsyncStream[OllamaChatResponse],
|
||||
) -> None:
|
||||
mock_chat.return_value = mock_chat_completion_response
|
||||
chat_history.append(ChatMessage(text="hello world", role="system"))
|
||||
chat_history.append(ChatMessage(text="hello world", role="user"))
|
||||
|
||||
ollama_client = OllamaChatClient()
|
||||
result = await ollama_client.get_response(messages=chat_history)
|
||||
|
||||
assert result.text == "test"
|
||||
|
||||
|
||||
@patch.object(AsyncClient, "chat", new_callable=AsyncMock)
|
||||
async def test_cmc_reasoning(
|
||||
mock_chat: AsyncMock,
|
||||
ollama_unit_test_env: dict[str, str],
|
||||
chat_history: list[ChatMessage],
|
||||
mock_chat_completion_response_reasoning: AsyncStream[OllamaChatResponse],
|
||||
) -> None:
|
||||
mock_chat.return_value = mock_chat_completion_response_reasoning
|
||||
chat_history.append(ChatMessage(text="hello world", role="user"))
|
||||
|
||||
ollama_client = OllamaChatClient()
|
||||
result = await ollama_client.get_response(messages=chat_history)
|
||||
|
||||
reasoning = "".join(c.text for c in result.messages.pop().contents if isinstance(c, TextReasoningContent))
|
||||
assert reasoning == "test"
|
||||
|
||||
|
||||
@patch.object(AsyncClient, "chat", new_callable=AsyncMock)
|
||||
async def test_cmc_chat_failure(
|
||||
mock_chat: AsyncMock,
|
||||
ollama_unit_test_env: dict[str, str],
|
||||
chat_history: list[ChatMessage],
|
||||
) -> None:
|
||||
# Simulate a failure in the Ollama client
|
||||
mock_chat.side_effect = Exception("Connection error")
|
||||
chat_history.append(ChatMessage(text="hello world", role="user"))
|
||||
|
||||
ollama_client = OllamaChatClient()
|
||||
|
||||
with pytest.raises(ServiceResponseException) as exc_info:
|
||||
await ollama_client.get_response(messages=chat_history)
|
||||
|
||||
assert "Ollama chat request failed" in str(exc_info.value)
|
||||
assert "Connection error" in str(exc_info.value)
|
||||
|
||||
|
||||
@patch.object(AsyncClient, "chat", new_callable=AsyncMock)
|
||||
async def test_cmc_streaming(
|
||||
mock_chat: AsyncMock,
|
||||
ollama_unit_test_env: dict[str, str],
|
||||
chat_history: list[ChatMessage],
|
||||
mock_streaming_chat_completion_response: AsyncStream[OllamaChatResponse],
|
||||
) -> None:
|
||||
mock_chat.return_value = mock_streaming_chat_completion_response
|
||||
chat_history.append(ChatMessage(text="hello world", role="system"))
|
||||
chat_history.append(ChatMessage(text="hello world", role="user"))
|
||||
|
||||
ollama_client = OllamaChatClient()
|
||||
result = ollama_client.get_streaming_response(messages=chat_history)
|
||||
|
||||
async for chunk in result:
|
||||
assert chunk.text == "test"
|
||||
|
||||
|
||||
@patch.object(AsyncClient, "chat", new_callable=AsyncMock)
|
||||
async def test_cmc_streaming_reasoning(
|
||||
mock_chat: AsyncMock,
|
||||
ollama_unit_test_env: dict[str, str],
|
||||
chat_history: list[ChatMessage],
|
||||
mock_streaming_chat_completion_response_reasoning: AsyncStream[OllamaChatResponse],
|
||||
) -> None:
|
||||
mock_chat.return_value = mock_streaming_chat_completion_response_reasoning
|
||||
chat_history.append(ChatMessage(text="hello world", role="user"))
|
||||
|
||||
ollama_client = OllamaChatClient()
|
||||
result = ollama_client.get_streaming_response(messages=chat_history)
|
||||
|
||||
async for chunk in result:
|
||||
reasoning = "".join(c.text for c in chunk.contents if isinstance(c, TextReasoningContent))
|
||||
assert reasoning == "test"
|
||||
|
||||
|
||||
@patch.object(AsyncClient, "chat", new_callable=AsyncMock)
|
||||
async def test_cmc_streaming_chat_failure(
|
||||
mock_chat: AsyncMock,
|
||||
ollama_unit_test_env: dict[str, str],
|
||||
chat_history: list[ChatMessage],
|
||||
) -> None:
|
||||
# Simulate a failure in the Ollama client for streaming
|
||||
mock_chat.side_effect = Exception("Streaming connection error")
|
||||
chat_history.append(ChatMessage(text="hello world", role="user"))
|
||||
|
||||
ollama_client = OllamaChatClient()
|
||||
|
||||
with pytest.raises(ServiceResponseException) as exc_info:
|
||||
async for _ in ollama_client.get_streaming_response(messages=chat_history):
|
||||
pass
|
||||
|
||||
assert "Ollama streaming chat request failed" in str(exc_info.value)
|
||||
assert "Streaming connection error" in str(exc_info.value)
|
||||
|
||||
|
||||
@patch.object(AsyncClient, "chat", new_callable=AsyncMock)
|
||||
async def test_cmc_streaming_with_tool_call(
|
||||
mock_chat: AsyncMock,
|
||||
ollama_unit_test_env: dict[str, str],
|
||||
chat_history: list[ChatMessage],
|
||||
mock_streaming_chat_completion_response: AsyncStream[OllamaChatResponse],
|
||||
mock_streaming_chat_completion_tool_call: AsyncStream[OllamaChatResponse],
|
||||
) -> None:
|
||||
mock_chat.side_effect = [
|
||||
mock_streaming_chat_completion_tool_call,
|
||||
mock_streaming_chat_completion_response,
|
||||
]
|
||||
|
||||
chat_history.append(ChatMessage(text="hello world", role="user"))
|
||||
|
||||
ollama_client = OllamaChatClient()
|
||||
result = ollama_client.get_streaming_response(messages=chat_history, tools=[hello_world])
|
||||
|
||||
chunks: list[ChatResponseUpdate] = []
|
||||
async for chunk in result:
|
||||
chunks.append(chunk)
|
||||
|
||||
# Check parsed Toolcalls
|
||||
assert isinstance(chunks[0].contents[0], FunctionCallContent)
|
||||
tool_call = chunks[0].contents[0]
|
||||
assert tool_call.name == "hello_world"
|
||||
assert tool_call.arguments == {"arg1": "value1"}
|
||||
assert isinstance(chunks[1].contents[0], FunctionResultContent)
|
||||
tool_result = chunks[1].contents[0]
|
||||
assert tool_result.result == "Hello World"
|
||||
assert isinstance(chunks[2].contents[0], TextContent)
|
||||
text_result = chunks[2].contents[0]
|
||||
assert text_result.text == "test"
|
||||
|
||||
|
||||
async def test_cmc_with_hosted_tool_call(
|
||||
ollama_unit_test_env: dict[str, str],
|
||||
chat_history: list[ChatMessage],
|
||||
) -> None:
|
||||
with pytest.raises(ServiceInvalidRequestError):
|
||||
additional_properties = {
|
||||
"user_location": {
|
||||
"country": "US",
|
||||
"city": "Seattle",
|
||||
}
|
||||
}
|
||||
|
||||
chat_history.append(ChatMessage(text="hello world", role="user"))
|
||||
|
||||
ollama_client = OllamaChatClient()
|
||||
await ollama_client.get_response(
|
||||
messages=chat_history,
|
||||
tools=[HostedWebSearchTool(additional_properties=additional_properties)],
|
||||
)
|
||||
|
||||
|
||||
@patch.object(AsyncClient, "chat", new_callable=AsyncMock)
|
||||
async def test_cmc_with_data_content_type(
|
||||
mock_chat: AsyncMock,
|
||||
ollama_unit_test_env: dict[str, str],
|
||||
chat_history: list[ChatMessage],
|
||||
mock_chat_completion_response: OllamaChatResponse,
|
||||
) -> None:
|
||||
mock_chat.return_value = mock_chat_completion_response
|
||||
chat_history.append(
|
||||
ChatMessage(
|
||||
contents=[DataContent(uri="data:image/png;base64,xyz", media_type="image/png")],
|
||||
role="user",
|
||||
)
|
||||
)
|
||||
|
||||
ollama_client = OllamaChatClient()
|
||||
|
||||
result = await ollama_client.get_response(messages=chat_history)
|
||||
assert result.text == "test"
|
||||
|
||||
|
||||
@patch.object(AsyncClient, "chat", new_callable=AsyncMock)
|
||||
async def test_cmc_with_invalid_data_content_media_type(
|
||||
mock_chat: AsyncMock,
|
||||
ollama_unit_test_env: dict[str, str],
|
||||
chat_history: list[ChatMessage],
|
||||
mock_streaming_chat_completion_response: AsyncStream[OllamaChatResponse],
|
||||
) -> None:
|
||||
with pytest.raises(ServiceInvalidRequestError):
|
||||
mock_chat.return_value = mock_streaming_chat_completion_response
|
||||
# Remote Uris are not supported by Ollama client
|
||||
chat_history.append(
|
||||
ChatMessage(
|
||||
contents=[DataContent(uri="data:audio/mp3;base64,xyz", media_type="audio/mp3")],
|
||||
role="user",
|
||||
)
|
||||
)
|
||||
|
||||
ollama_client = OllamaChatClient()
|
||||
ollama_client.client.chat = AsyncMock(return_value=mock_streaming_chat_completion_response)
|
||||
|
||||
await ollama_client.get_response(messages=chat_history)
|
||||
|
||||
|
||||
@patch.object(AsyncClient, "chat", new_callable=AsyncMock)
|
||||
async def test_cmc_with_invalid_content_type(
|
||||
mock_chat: AsyncMock,
|
||||
ollama_unit_test_env: dict[str, str],
|
||||
chat_history: list[ChatMessage],
|
||||
mock_chat_completion_response: AsyncStream[OllamaChatResponse],
|
||||
) -> None:
|
||||
with pytest.raises(ServiceInvalidRequestError):
|
||||
mock_chat.return_value = mock_chat_completion_response
|
||||
# Remote Uris are not supported by Ollama client
|
||||
chat_history.append(
|
||||
ChatMessage(
|
||||
contents=[UriContent(uri="http://example.com/image.png", media_type="image/png")],
|
||||
role="user",
|
||||
)
|
||||
)
|
||||
|
||||
ollama_client = OllamaChatClient()
|
||||
|
||||
await ollama_client.get_response(messages=chat_history)
|
||||
|
||||
|
||||
@skip_if_azure_integration_tests_disabled
|
||||
async def test_cmc_integration_with_tool_call(
|
||||
chat_history: list[ChatMessage],
|
||||
) -> None:
|
||||
chat_history.append(ChatMessage(text="Call the hello world function and repeat what it says", role="user"))
|
||||
|
||||
ollama_client = OllamaChatClient()
|
||||
result = await ollama_client.get_response(messages=chat_history, tools=[hello_world])
|
||||
|
||||
assert "hello" in result.text.lower() and "world" in result.text.lower()
|
||||
assert isinstance(result.messages[-2].contents[0], FunctionResultContent)
|
||||
tool_result = result.messages[-2].contents[0]
|
||||
assert tool_result.result == "Hello World"
|
||||
|
||||
|
||||
@skip_if_azure_integration_tests_disabled
|
||||
async def test_cmc_integration_with_chat_completion(
|
||||
chat_history: list[ChatMessage],
|
||||
) -> None:
|
||||
chat_history.append(ChatMessage(text="Say Hello World", role="user"))
|
||||
|
||||
ollama_client = OllamaChatClient()
|
||||
result = await ollama_client.get_response(messages=chat_history)
|
||||
|
||||
assert "hello" in result.text.lower()
|
||||
|
||||
|
||||
@skip_if_azure_integration_tests_disabled
|
||||
async def test_cmc_streaming_integration_with_tool_call(
|
||||
chat_history: list[ChatMessage],
|
||||
) -> None:
|
||||
chat_history.append(ChatMessage(text="Call the hello world function and repeat what it says", role="user"))
|
||||
|
||||
ollama_client = OllamaChatClient()
|
||||
result: AsyncIterable[ChatResponseUpdate] = ollama_client.get_streaming_response(
|
||||
messages=chat_history, tools=[hello_world]
|
||||
)
|
||||
|
||||
chunks: list[ChatResponseUpdate] = []
|
||||
async for chunk in result:
|
||||
chunks.append(chunk)
|
||||
|
||||
for c in chunks:
|
||||
if len(c.contents) > 0:
|
||||
if isinstance(c.contents[0], FunctionResultContent):
|
||||
tool_result = c.contents[0]
|
||||
assert tool_result.result == "Hello World"
|
||||
if isinstance(c.contents[0], FunctionCallContent):
|
||||
tool_call = c.contents[0]
|
||||
assert tool_call.name == "hello_world"
|
||||
|
||||
|
||||
@skip_if_azure_integration_tests_disabled
|
||||
async def test_cmc_streaming_integration_with_chat_completion(
|
||||
chat_history: list[ChatMessage],
|
||||
) -> None:
|
||||
chat_history.append(ChatMessage(text="Say Hello World", role="user"))
|
||||
|
||||
ollama_client = OllamaChatClient()
|
||||
result: AsyncIterable[ChatResponseUpdate] = ollama_client.get_streaming_response(messages=chat_history)
|
||||
|
||||
full_text = ""
|
||||
async for chunk in result:
|
||||
full_text += chunk.text
|
||||
|
||||
assert "hello" in full_text.lower() and "world" in full_text.lower()
|
||||
Reference in New Issue
Block a user