mirror of
https://github.com/microsoft/agent-framework.git
synced 2026-06-16 21:04:09 +08:00
Python: Web search file search tools (#395)
* Add web and file search tools * add tests * PR comments * Add tools support for chat and assistants clients * fix code checks * add tests for assistants client * Add samples * fix fn descriptions * Add openai responses model id to environment variables --------- Co-authored-by: Dmytro Struk <13853051+dmytrostruk@users.noreply.github.com>
This commit is contained in:
committed by
GitHub
Unverified
parent
ed86baa6cb
commit
0410f51777
@@ -141,7 +141,7 @@ def _tool_call_non_streaming(
|
||||
# Failsafe: give up on tools, ask model for plain answer
|
||||
chat_options.tool_choice = "none"
|
||||
self._prepare_tool_choice(chat_options=chat_options) # type: ignore[reportPrivateUsage]
|
||||
response = await func(self, messages=messages, chat_options=chat_options)
|
||||
response = await func(self, messages=messages, chat_options=chat_options, **kwargs)
|
||||
if fcc_messages:
|
||||
for msg in reversed(fcc_messages):
|
||||
response.messages.insert(0, msg)
|
||||
@@ -167,7 +167,7 @@ def _tool_call_streaming(
|
||||
for attempt_idx in range(getattr(self, "__maximum_iterations_per_request", 10)):
|
||||
function_call_returned = False
|
||||
all_messages: list[ChatResponseUpdate] = []
|
||||
async for update in func(self, messages=messages, chat_options=chat_options):
|
||||
async for update in func(self, messages=messages, chat_options=chat_options, **kwargs):
|
||||
if update.contents and any(isinstance(item, FunctionCallContent) for item in update.contents):
|
||||
all_messages.append(update)
|
||||
function_call_returned = True
|
||||
|
||||
@@ -4,7 +4,17 @@ import inspect
|
||||
from collections.abc import Awaitable, Callable
|
||||
from functools import wraps
|
||||
from time import perf_counter
|
||||
from typing import TYPE_CHECKING, Annotated, Any, Generic, Protocol, TypeVar, get_args, get_origin, runtime_checkable
|
||||
from typing import (
|
||||
TYPE_CHECKING,
|
||||
Annotated,
|
||||
Any,
|
||||
Generic,
|
||||
Protocol,
|
||||
TypeVar,
|
||||
get_args,
|
||||
get_origin,
|
||||
runtime_checkable,
|
||||
)
|
||||
|
||||
from opentelemetry import metrics, trace
|
||||
from pydantic import BaseModel, Field, PrivateAttr, create_model
|
||||
@@ -20,7 +30,14 @@ tracer: trace.Tracer = trace.get_tracer("agent_framework")
|
||||
meter: metrics.Meter = metrics.get_meter_provider().get_meter("agent_framework")
|
||||
logger = get_logger()
|
||||
|
||||
__all__ = ["AIFunction", "AITool", "HostedCodeInterpreterTool", "ai_function"]
|
||||
__all__ = [
|
||||
"AIFunction",
|
||||
"AITool",
|
||||
"HostedCodeInterpreterTool",
|
||||
"HostedFileSearchTool",
|
||||
"HostedWebSearchTool",
|
||||
"ai_function",
|
||||
]
|
||||
|
||||
|
||||
def _parse_inputs(
|
||||
@@ -160,6 +177,81 @@ class HostedCodeInterpreterTool(AIToolBase):
|
||||
super().__init__(**args, **kwargs)
|
||||
|
||||
|
||||
class HostedWebSearchTool(AIToolBase):
|
||||
"""Represents a web search tool that can be specified to an AI service to enable it to perform web searches."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
description: str | None = None,
|
||||
additional_properties: dict[str, Any] | None = None,
|
||||
**kwargs: Any,
|
||||
):
|
||||
"""Initialize a HostedWebSearchTool.
|
||||
|
||||
Args:
|
||||
description: A description of the tool.
|
||||
additional_properties: Additional properties associated with the tool
|
||||
(e.g., {"user_location": {"city": "Seattle", "country": "US"}}).
|
||||
**kwargs: Additional keyword arguments to pass to the base class.
|
||||
"""
|
||||
args: dict[str, Any] = {
|
||||
"name": "web_search",
|
||||
}
|
||||
if description is not None:
|
||||
args["description"] = description
|
||||
if additional_properties is not None:
|
||||
args["additional_properties"] = additional_properties
|
||||
if "name" in kwargs:
|
||||
raise ValueError("The 'name' argument is reserved for the HostedFileSearchTool and cannot be set.")
|
||||
super().__init__(**args, **kwargs)
|
||||
|
||||
|
||||
class HostedFileSearchTool(AIToolBase):
|
||||
"""Represents a file search tool that can be specified to an AI service to enable it to perform file searches."""
|
||||
|
||||
inputs: list[Any] | None = None
|
||||
max_results: int | None = None
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
inputs: "AIContents | dict[str, Any] | str | list[AIContents | dict[str, Any] | str] | None" = None,
|
||||
max_results: int | None = None,
|
||||
description: str | None = None,
|
||||
additional_properties: dict[str, Any] | None = None,
|
||||
**kwargs: Any,
|
||||
):
|
||||
"""Initialize a FileSearchTool.
|
||||
|
||||
Args:
|
||||
inputs: A list of contents that the tool can accept as input. Defaults to None.
|
||||
This should be one or more HostedVectorStoreContents.
|
||||
When supplying a list, it can contain:
|
||||
- AIContents instances
|
||||
- dicts with properties for AIContents (e.g., {"uri": "http://example.com", "media_type": "text/html"})
|
||||
- strings (which will be converted to UriContent with media_type "text/plain").
|
||||
If None, defaults to an empty list.
|
||||
max_results: The maximum number of results to return from the file search.
|
||||
If None, max limit is applied.
|
||||
description: A description of the tool.
|
||||
additional_properties: Additional properties associated with the tool.
|
||||
**kwargs: Additional keyword arguments to pass to the base class.
|
||||
"""
|
||||
args: dict[str, Any] = {
|
||||
"name": "file_search",
|
||||
}
|
||||
if inputs:
|
||||
args["inputs"] = _parse_inputs(inputs)
|
||||
if max_results:
|
||||
args["max_results"] = max_results
|
||||
if description is not None:
|
||||
args["description"] = description
|
||||
if additional_properties is not None:
|
||||
args["additional_properties"] = additional_properties
|
||||
if "name" in kwargs:
|
||||
raise ValueError("The 'name' argument is reserved for the HostedFileSearchTool and cannot be set.")
|
||||
super().__init__(**args, **kwargs)
|
||||
|
||||
|
||||
class AIFunction(AIToolBase, Generic[ArgsT, ReturnT]):
|
||||
"""A AITool that is callable as code.
|
||||
|
||||
|
||||
@@ -21,7 +21,7 @@ from openai.types.beta.threads.runs import RunStep
|
||||
from pydantic import Field, PrivateAttr, SecretStr, ValidationError
|
||||
|
||||
from .._clients import ChatClientBase, use_tool_calling
|
||||
from .._tools import AIFunction, HostedCodeInterpreterTool
|
||||
from .._tools import AIFunction, HostedCodeInterpreterTool, HostedFileSearchTool
|
||||
from .._types import (
|
||||
AIContents,
|
||||
ChatMessage,
|
||||
@@ -46,6 +46,7 @@ if sys.version_info >= (3, 11):
|
||||
else:
|
||||
from typing_extensions import Self # pragma: no cover
|
||||
|
||||
|
||||
__all__ = ["OpenAIAssistantsClient"]
|
||||
|
||||
|
||||
@@ -247,6 +248,7 @@ class OpenAIAssistantsClient(OpenAIConfigBase, ChatClientBase):
|
||||
metadata=run_options.get("metadata"),
|
||||
)
|
||||
run_options["additional_messages"] = []
|
||||
run_options.pop("tool_resources", None)
|
||||
return thread.id
|
||||
|
||||
if thread_run is not None:
|
||||
@@ -365,6 +367,13 @@ class OpenAIAssistantsClient(OpenAIConfigBase, ChatClientBase):
|
||||
tool_definitions.append(tool.to_json_schema_spec()) # type: ignore[reportUnknownArgumentType]
|
||||
elif isinstance(tool, HostedCodeInterpreterTool):
|
||||
tool_definitions.append({"type": "code_interpreter"})
|
||||
elif isinstance(tool, HostedFileSearchTool):
|
||||
params: dict[str, Any] = {
|
||||
"type": "file_search",
|
||||
}
|
||||
if tool.max_results is not None:
|
||||
params["max_num_results"] = tool.max_results
|
||||
tool_definitions.append(params)
|
||||
elif isinstance(tool, MutableMapping):
|
||||
tool_definitions.append(tool)
|
||||
|
||||
|
||||
@@ -19,6 +19,7 @@ from agent_framework import AIFunction, AITool, UsageContent
|
||||
|
||||
from .._clients import ChatClientBase, use_tool_calling
|
||||
from .._logging import get_logger
|
||||
from .._tools import HostedWebSearchTool
|
||||
from .._types import (
|
||||
AIContents,
|
||||
ChatFinishReason,
|
||||
@@ -125,16 +126,40 @@ class OpenAIChatClientBase(OpenAIHandler, ChatClientBase):
|
||||
chat_tools.append(tool if isinstance(tool, dict) else dict(tool))
|
||||
return chat_tools
|
||||
|
||||
def _process_web_search_tool(self, tools: list[AITool | MutableMapping[str, Any]]) -> dict[str, Any] | None:
|
||||
for tool in tools:
|
||||
if isinstance(tool, HostedWebSearchTool):
|
||||
# Web search tool requires special handling
|
||||
return (
|
||||
{
|
||||
"user_location": {
|
||||
"approximate": tool.additional_properties.get("user_location", None),
|
||||
"type": "approximate",
|
||||
}
|
||||
}
|
||||
if tool.additional_properties and "user_location" in tool.additional_properties
|
||||
else {}
|
||||
)
|
||||
|
||||
return None
|
||||
|
||||
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_provider_settings()
|
||||
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")
|
||||
if chat_options.tools is None:
|
||||
options_dict.pop("parallel_tool_calls", None)
|
||||
else:
|
||||
if chat_options.tools is not None:
|
||||
web_search_options = self._process_web_search_tool(chat_options.tools)
|
||||
if web_search_options:
|
||||
options_dict["web_search_options"] = web_search_options
|
||||
options_dict["tools"] = self._chat_to_tool_spec(chat_options.tools)
|
||||
if not options_dict.get("tools", None):
|
||||
options_dict.pop("tools", None)
|
||||
options_dict.pop("parallel_tool_calls", None)
|
||||
options_dict.pop("tool_choice", None)
|
||||
|
||||
if "model" not in options_dict:
|
||||
options_dict["model"] = self.ai_model_id
|
||||
if (
|
||||
|
||||
@@ -7,6 +7,7 @@ from itertools import chain
|
||||
from typing import TYPE_CHECKING, Any, Literal, TypeVar
|
||||
|
||||
from openai import AsyncOpenAI, BadRequestError
|
||||
from openai.types.responses.file_search_tool_param import FileSearchToolParam
|
||||
from openai.types.responses.function_tool_param import FunctionToolParam
|
||||
from openai.types.responses.parsed_response import (
|
||||
ParsedResponse,
|
||||
@@ -26,13 +27,15 @@ from openai.types.responses.tool_param import (
|
||||
CodeInterpreterContainerCodeInterpreterToolAuto,
|
||||
ToolParam,
|
||||
)
|
||||
from openai.types.responses.web_search_tool_param import UserLocation as WebSearchUserLocation
|
||||
from openai.types.responses.web_search_tool_param import WebSearchToolParam
|
||||
from pydantic import BaseModel, SecretStr, ValidationError
|
||||
|
||||
from agent_framework import DataContent, TextReasoningContent, UriContent, UsageContent
|
||||
|
||||
from .._clients import ChatClientBase, use_tool_calling
|
||||
from .._logging import get_logger
|
||||
from .._tools import AIFunction, AITool, HostedCodeInterpreterTool
|
||||
from .._tools import AIFunction, AITool, HostedCodeInterpreterTool, HostedFileSearchTool, HostedWebSearchTool
|
||||
from .._types import (
|
||||
AIContents,
|
||||
ChatMessage,
|
||||
@@ -44,6 +47,7 @@ from .._types import (
|
||||
FunctionCallContent,
|
||||
FunctionResultContent,
|
||||
HostedFileContent,
|
||||
HostedVectorStoreContent,
|
||||
TextContent,
|
||||
TextSpanRegion,
|
||||
UsageDetails,
|
||||
@@ -78,6 +82,8 @@ __all__ = ["OpenAIResponsesClient"]
|
||||
class OpenAIResponsesClientBase(OpenAIHandler, ChatClientBase):
|
||||
"""Base class for all OpenAI Responses based API's."""
|
||||
|
||||
FILE_SEARCH_MAX_RESULTS: int = 50
|
||||
|
||||
@override
|
||||
async def get_response(
|
||||
self,
|
||||
@@ -376,6 +382,45 @@ class OpenAIResponsesClientBase(OpenAIHandler, ChatClientBase):
|
||||
description=tool.description,
|
||||
)
|
||||
)
|
||||
case HostedFileSearchTool():
|
||||
if not tool.inputs:
|
||||
raise ValueError("HostedFileSearchTool requires inputs to be specified.")
|
||||
inputs: list[str] = [
|
||||
inp.vector_store_id for inp in tool.inputs if isinstance(inp, HostedVectorStoreContent)
|
||||
]
|
||||
if not inputs:
|
||||
raise ValueError(
|
||||
"HostedFileSearchTool requires inputs to be of type `HostedVectorStoreContent`."
|
||||
)
|
||||
|
||||
response_tools.append(
|
||||
FileSearchToolParam(
|
||||
type="file_search",
|
||||
vector_store_ids=inputs,
|
||||
max_num_results=tool.max_results
|
||||
or self.FILE_SEARCH_MAX_RESULTS, # default to max results if not specified
|
||||
)
|
||||
)
|
||||
case HostedWebSearchTool():
|
||||
location: dict[str, str] | None = (
|
||||
tool.additional_properties.get("user_location", None)
|
||||
if tool.additional_properties
|
||||
else None
|
||||
)
|
||||
response_tools.append(
|
||||
WebSearchToolParam(
|
||||
type="web_search_preview",
|
||||
user_location=WebSearchUserLocation(
|
||||
type="approximate",
|
||||
city=location.get("city", None),
|
||||
country=location.get("country", None),
|
||||
region=location.get("region", None),
|
||||
timezone=location.get("timezone", None),
|
||||
)
|
||||
if location
|
||||
else None,
|
||||
)
|
||||
)
|
||||
case _:
|
||||
logger.debug("Unsupported tool passed (type: %s)", type(tool))
|
||||
else:
|
||||
|
||||
@@ -12,6 +12,8 @@ from agent_framework import (
|
||||
ChatMessage,
|
||||
ChatResponse,
|
||||
ChatResponseUpdate,
|
||||
HostedFileSearchTool,
|
||||
HostedVectorStoreContent,
|
||||
TextContent,
|
||||
)
|
||||
from agent_framework.exceptions import ServiceInitializationError
|
||||
@@ -47,6 +49,29 @@ def create_test_openai_assistants_client(
|
||||
)
|
||||
|
||||
|
||||
async def create_vector_store(client: OpenAIAssistantsClient) -> tuple[str, HostedVectorStoreContent]:
|
||||
"""Create a vector store with sample documents for testing."""
|
||||
file = await client.client.files.create(
|
||||
file=("todays_weather.txt", b"The weather today is sunny with a high of 25C."), purpose="user_data"
|
||||
)
|
||||
vector_store = await client.client.vector_stores.create(
|
||||
name="knowledge_base",
|
||||
expires_after={"anchor": "last_active_at", "days": 1},
|
||||
)
|
||||
result = await client.client.vector_stores.files.create_and_poll(vector_store_id=vector_store.id, file_id=file.id)
|
||||
if result.last_error is not None:
|
||||
raise Exception(f"Vector store file processing failed with status: {result.last_error.message}")
|
||||
|
||||
return file.id, HostedVectorStoreContent(vector_store_id=vector_store.id)
|
||||
|
||||
|
||||
async def delete_vector_store(client: OpenAIAssistantsClient, file_id: str, vector_store_id: str) -> None:
|
||||
"""Delete the vector store after tests."""
|
||||
|
||||
await client.client.vector_stores.delete(vector_store_id=vector_store_id)
|
||||
await client.client.files.delete(file_id=file_id)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_async_openai() -> MagicMock:
|
||||
"""Mock AsyncOpenAI client."""
|
||||
@@ -386,3 +411,54 @@ async def test_openai_assistants_client_with_existing_assistant() -> None:
|
||||
assert response is not None
|
||||
assert isinstance(response, ChatResponse)
|
||||
assert len(response.text) > 0
|
||||
|
||||
|
||||
@skip_if_openai_integration_tests_disabled
|
||||
async def test_openai_assistants_client_file_search() -> None:
|
||||
"""Test OpenAI Assistants Client response."""
|
||||
async with OpenAIAssistantsClient() as openai_assistants_client:
|
||||
assert isinstance(openai_assistants_client, ChatClient)
|
||||
|
||||
messages: list[ChatMessage] = []
|
||||
messages.append(ChatMessage(role="user", text="What's the weather like today?"))
|
||||
|
||||
file_id, vector_store = await create_vector_store(openai_assistants_client)
|
||||
response = await openai_assistants_client.get_response(
|
||||
messages=messages,
|
||||
tools=[HostedFileSearchTool()],
|
||||
tool_resources={"file_search": {"vector_store_ids": [vector_store.vector_store_id]}},
|
||||
)
|
||||
await delete_vector_store(openai_assistants_client, file_id, vector_store.vector_store_id)
|
||||
|
||||
assert response is not None
|
||||
assert isinstance(response, ChatResponse)
|
||||
assert any(word in response.text.lower() for word in ["sunny", "25", "weather"])
|
||||
|
||||
|
||||
@skip_if_openai_integration_tests_disabled
|
||||
async def test_openai_assistants_client_file_search_streaming() -> None:
|
||||
"""Test OpenAI Assistants Client response."""
|
||||
async with OpenAIAssistantsClient() as openai_assistants_client:
|
||||
assert isinstance(openai_assistants_client, ChatClient)
|
||||
|
||||
messages: list[ChatMessage] = []
|
||||
messages.append(ChatMessage(role="user", text="What's the weather like today?"))
|
||||
|
||||
file_id, vector_store = await create_vector_store(openai_assistants_client)
|
||||
response = openai_assistants_client.get_streaming_response(
|
||||
messages=messages,
|
||||
tools=[HostedFileSearchTool()],
|
||||
tool_resources={"file_search": {"vector_store_ids": [vector_store.vector_store_id]}},
|
||||
)
|
||||
|
||||
assert response is not None
|
||||
full_message: str = ""
|
||||
async for chunk in response:
|
||||
assert chunk is not None
|
||||
assert isinstance(chunk, ChatResponseUpdate)
|
||||
for content in chunk.contents:
|
||||
if isinstance(content, TextContent) and content.text:
|
||||
full_message += content.text
|
||||
await delete_vector_store(openai_assistants_client, file_id, vector_store.vector_store_id)
|
||||
|
||||
assert any(word in full_message.lower() for word in ["sunny", "25", "weather"])
|
||||
|
||||
@@ -4,7 +4,15 @@ import os
|
||||
|
||||
import pytest
|
||||
|
||||
from agent_framework import ChatClient, ChatMessage, ChatResponse, ChatResponseUpdate, TextContent, ai_function
|
||||
from agent_framework import (
|
||||
ChatClient,
|
||||
ChatMessage,
|
||||
ChatResponse,
|
||||
ChatResponseUpdate,
|
||||
HostedWebSearchTool,
|
||||
TextContent,
|
||||
ai_function,
|
||||
)
|
||||
from agent_framework.exceptions import ServiceInitializationError
|
||||
from agent_framework.openai import OpenAIChatClient
|
||||
|
||||
@@ -231,3 +239,96 @@ async def test_openai_chat_client_streaming_tools() -> None:
|
||||
full_message += content.text
|
||||
|
||||
assert "scientists" in full_message
|
||||
|
||||
|
||||
@skip_if_openai_integration_tests_disabled
|
||||
async def test_openai_chat_client_web_search() -> None:
|
||||
# Currently only a select few models support web search tool calls
|
||||
openai_chat_client = OpenAIChatClient(ai_model_id="gpt-4o-search-preview")
|
||||
|
||||
assert isinstance(openai_chat_client, ChatClient)
|
||||
|
||||
# Test that the client will use the web search tool
|
||||
response = await openai_chat_client.get_response(
|
||||
messages=[
|
||||
ChatMessage(
|
||||
role="user",
|
||||
text="Who are the main characters of Kpop Demon Hunters? Do a web search to find the answer.",
|
||||
)
|
||||
],
|
||||
tools=[HostedWebSearchTool()],
|
||||
tool_choice="auto",
|
||||
)
|
||||
|
||||
assert response is not None
|
||||
assert isinstance(response, ChatResponse)
|
||||
assert "Rumi" in response.text
|
||||
assert "Mira" in response.text
|
||||
assert "Zoey" in response.text
|
||||
|
||||
# Test that the client will use the web search tool with location
|
||||
additional_properties = {
|
||||
"user_location": {
|
||||
"country": "US",
|
||||
"city": "Seattle",
|
||||
}
|
||||
}
|
||||
response = await openai_chat_client.get_response(
|
||||
messages=[ChatMessage(role="user", text="What is the current weather? Do not ask for my current location.")],
|
||||
tools=[HostedWebSearchTool(additional_properties=additional_properties)],
|
||||
tool_choice="auto",
|
||||
)
|
||||
assert "Seattle" in response.text
|
||||
|
||||
|
||||
@skip_if_openai_integration_tests_disabled
|
||||
async def test_openai_chat_client_web_search_streaming() -> None:
|
||||
openai_chat_client = OpenAIChatClient(ai_model_id="gpt-4o-search-preview")
|
||||
|
||||
assert isinstance(openai_chat_client, ChatClient)
|
||||
|
||||
# Test that the client will use the web search tool
|
||||
response = openai_chat_client.get_streaming_response(
|
||||
messages=[
|
||||
ChatMessage(
|
||||
role="user",
|
||||
text="Who are the main characters of Kpop Demon Hunters? Do a web search to find the answer.",
|
||||
)
|
||||
],
|
||||
tools=[HostedWebSearchTool()],
|
||||
tool_choice="auto",
|
||||
)
|
||||
|
||||
assert response is not None
|
||||
full_message: str = ""
|
||||
async for chunk in response:
|
||||
assert chunk is not None
|
||||
assert isinstance(chunk, ChatResponseUpdate)
|
||||
for content in chunk.contents:
|
||||
if isinstance(content, TextContent) and content.text:
|
||||
full_message += content.text
|
||||
assert "Rumi" in full_message
|
||||
assert "Mira" in full_message
|
||||
assert "Zoey" in full_message
|
||||
|
||||
# Test that the client will use the web search tool with location
|
||||
additional_properties = {
|
||||
"user_location": {
|
||||
"country": "US",
|
||||
"city": "Seattle",
|
||||
}
|
||||
}
|
||||
response = openai_chat_client.get_streaming_response(
|
||||
messages=[ChatMessage(role="user", text="What is the current weather? Do not ask for my current location.")],
|
||||
tools=[HostedWebSearchTool(additional_properties=additional_properties)],
|
||||
tool_choice="auto",
|
||||
)
|
||||
assert response is not None
|
||||
full_message: str = ""
|
||||
async for chunk in response:
|
||||
assert chunk is not None
|
||||
assert isinstance(chunk, ChatResponseUpdate)
|
||||
for content in chunk.contents:
|
||||
if isinstance(content, TextContent) and content.text:
|
||||
full_message += content.text
|
||||
assert "Seattle" in full_message
|
||||
|
||||
@@ -6,7 +6,17 @@ from typing import Annotated
|
||||
import pytest
|
||||
from pydantic import BaseModel
|
||||
|
||||
from agent_framework import ChatClient, ChatMessage, ChatResponse, ChatResponseUpdate, TextContent, ai_function
|
||||
from agent_framework import (
|
||||
ChatClient,
|
||||
ChatMessage,
|
||||
ChatResponse,
|
||||
ChatResponseUpdate,
|
||||
HostedFileSearchTool,
|
||||
HostedVectorStoreContent,
|
||||
HostedWebSearchTool,
|
||||
TextContent,
|
||||
ai_function,
|
||||
)
|
||||
from agent_framework.exceptions import ServiceInitializationError
|
||||
from agent_framework.openai import OpenAIResponsesClient
|
||||
|
||||
@@ -26,6 +36,29 @@ class OutputStruct(BaseModel):
|
||||
weather: str
|
||||
|
||||
|
||||
async def create_vector_store(client: OpenAIResponsesClient) -> tuple[str, HostedVectorStoreContent]:
|
||||
"""Create a vector store with sample documents for testing."""
|
||||
file = await client.client.files.create(
|
||||
file=("todays_weather.txt", b"The weather today is sunny with a high of 75F."), purpose="user_data"
|
||||
)
|
||||
vector_store = await client.client.vector_stores.create(
|
||||
name="knowledge_base",
|
||||
expires_after={"anchor": "last_active_at", "days": 1},
|
||||
)
|
||||
result = await client.client.vector_stores.files.create_and_poll(vector_store_id=vector_store.id, file_id=file.id)
|
||||
if result.last_error is not None:
|
||||
raise Exception(f"Vector store file processing failed with status: {result.last_error.message}")
|
||||
|
||||
return file.id, HostedVectorStoreContent(vector_store_id=vector_store.id)
|
||||
|
||||
|
||||
async def delete_vector_store(client: OpenAIResponsesClient, file_id: str, vector_store_id: str) -> None:
|
||||
"""Delete the vector store after tests."""
|
||||
|
||||
await client.client.vector_stores.delete(vector_store_id=vector_store_id)
|
||||
await client.client.files.delete(file_id=file_id)
|
||||
|
||||
|
||||
@ai_function
|
||||
async def get_weather(location: Annotated[str, "The location as a city name"]) -> str:
|
||||
"""Get the current weather in a given location."""
|
||||
@@ -132,7 +165,7 @@ def test_serialize_with_org_id(openai_unit_test_env: dict[str, str]) -> None:
|
||||
@skip_if_openai_integration_tests_disabled
|
||||
async def test_openai_responses_client_response() -> None:
|
||||
"""Test OpenAI chat completion responses."""
|
||||
openai_responses_client = OpenAIResponsesClient(ai_model_id="gpt-4.1-mini")
|
||||
openai_responses_client = OpenAIResponsesClient()
|
||||
|
||||
assert isinstance(openai_responses_client, ChatClient)
|
||||
|
||||
@@ -156,8 +189,8 @@ async def test_openai_responses_client_response() -> None:
|
||||
assert "scientists" in response.text
|
||||
|
||||
messages.clear()
|
||||
messages.append(ChatMessage(role="user", text="The weather in New York is sunny"))
|
||||
messages.append(ChatMessage(role="user", text="What is the weather in New York?"))
|
||||
messages.append(ChatMessage(role="user", text="The weather in Seattle is sunny"))
|
||||
messages.append(ChatMessage(role="user", text="What is the weather in Seattle?"))
|
||||
|
||||
# Test that the client can be used to get a response
|
||||
response = await openai_responses_client.get_response(
|
||||
@@ -168,14 +201,14 @@ async def test_openai_responses_client_response() -> None:
|
||||
assert response is not None
|
||||
assert isinstance(response, ChatResponse)
|
||||
output = OutputStruct.model_validate_json(response.text)
|
||||
assert output.location == "New York"
|
||||
assert "sunny" in output.weather
|
||||
assert "seattle" in output.location.lower()
|
||||
assert "sunny" in output.weather.lower()
|
||||
|
||||
|
||||
@skip_if_openai_integration_tests_disabled
|
||||
async def test_openai_responses_client_response_tools() -> None:
|
||||
"""Test OpenAI chat completion responses."""
|
||||
openai_responses_client = OpenAIResponsesClient(ai_model_id="gpt-4o-mini")
|
||||
openai_responses_client = OpenAIResponsesClient()
|
||||
|
||||
assert isinstance(openai_responses_client, ChatClient)
|
||||
|
||||
@@ -191,7 +224,7 @@ async def test_openai_responses_client_response_tools() -> None:
|
||||
|
||||
assert response is not None
|
||||
assert isinstance(response, ChatResponse)
|
||||
assert "sunny" in response.text
|
||||
assert "sunny" in response.text.lower()
|
||||
|
||||
messages.clear()
|
||||
messages.append(ChatMessage(role="user", text="What is the weather in Seattle?"))
|
||||
@@ -207,14 +240,14 @@ async def test_openai_responses_client_response_tools() -> None:
|
||||
assert response is not None
|
||||
assert isinstance(response, ChatResponse)
|
||||
output = OutputStruct.model_validate_json(response.text)
|
||||
assert "Seattle" in output.location
|
||||
assert "sunny" in output.weather
|
||||
assert "seattle" in output.location.lower()
|
||||
assert "sunny" in output.weather.lower()
|
||||
|
||||
|
||||
@skip_if_openai_integration_tests_disabled
|
||||
async def test_openai_responses_client_streaming() -> None:
|
||||
"""Test Azure OpenAI chat completion responses."""
|
||||
openai_responses_client = OpenAIResponsesClient(ai_model_id="gpt-4.1-mini")
|
||||
openai_responses_client = OpenAIResponsesClient()
|
||||
|
||||
assert isinstance(openai_responses_client, ChatClient)
|
||||
|
||||
@@ -260,14 +293,14 @@ async def test_openai_responses_client_streaming() -> None:
|
||||
full_message += content.text
|
||||
|
||||
output = OutputStruct.model_validate_json(full_message)
|
||||
assert "Seattle" in output.location
|
||||
assert "sunny" in output.weather
|
||||
assert "seattle" in output.location.lower()
|
||||
assert "sunny" in output.weather.lower()
|
||||
|
||||
|
||||
@skip_if_openai_integration_tests_disabled
|
||||
async def test_openai_responses_client_streaming_tools() -> None:
|
||||
"""Test OpenAI chat completion responses."""
|
||||
openai_responses_client = OpenAIResponsesClient(ai_model_id="gpt-4o-mini")
|
||||
openai_responses_client = OpenAIResponsesClient()
|
||||
|
||||
assert isinstance(openai_responses_client, ChatClient)
|
||||
|
||||
@@ -287,7 +320,7 @@ async def test_openai_responses_client_streaming_tools() -> None:
|
||||
if isinstance(content, TextContent) and content.text:
|
||||
full_message += content.text
|
||||
|
||||
assert "sunny" in full_message
|
||||
assert "sunny" in full_message.lower()
|
||||
|
||||
messages.clear()
|
||||
messages.append(ChatMessage(role="user", text="What is the weather in Seattle?"))
|
||||
@@ -307,5 +340,155 @@ async def test_openai_responses_client_streaming_tools() -> None:
|
||||
full_message += content.text
|
||||
|
||||
output = OutputStruct.model_validate_json(full_message)
|
||||
assert "Seattle" in output.location
|
||||
assert "sunny" in output.weather
|
||||
assert "seattle" in output.location.lower()
|
||||
assert "sunny" in output.weather.lower()
|
||||
|
||||
|
||||
@skip_if_openai_integration_tests_disabled
|
||||
async def test_openai_responses_client_web_search() -> None:
|
||||
openai_responses_client = OpenAIResponsesClient()
|
||||
|
||||
assert isinstance(openai_responses_client, ChatClient)
|
||||
|
||||
# Test that the client will use the web search tool
|
||||
response = await openai_responses_client.get_response(
|
||||
messages=[
|
||||
ChatMessage(
|
||||
role="user",
|
||||
text="Who are the main characters of Kpop Demon Hunters? Do a web search to find the answer.",
|
||||
)
|
||||
],
|
||||
tools=[HostedWebSearchTool()],
|
||||
tool_choice="auto",
|
||||
)
|
||||
|
||||
assert response is not None
|
||||
assert isinstance(response, ChatResponse)
|
||||
assert "Rumi" in response.text
|
||||
assert "Mira" in response.text
|
||||
assert "Zoey" in response.text
|
||||
|
||||
# Test that the client will use the web search tool with location
|
||||
additional_properties = {
|
||||
"user_location": {
|
||||
"country": "US",
|
||||
"city": "Seattle",
|
||||
}
|
||||
}
|
||||
response = await openai_responses_client.get_response(
|
||||
messages=[ChatMessage(role="user", text="What is the current weather? Do not ask for my current location.")],
|
||||
tools=[HostedWebSearchTool(additional_properties=additional_properties)],
|
||||
tool_choice="auto",
|
||||
)
|
||||
assert "Seattle" in response.text
|
||||
|
||||
|
||||
@skip_if_openai_integration_tests_disabled
|
||||
async def test_openai_responses_client_web_search_streaming() -> None:
|
||||
openai_responses_client = OpenAIResponsesClient()
|
||||
|
||||
assert isinstance(openai_responses_client, ChatClient)
|
||||
|
||||
# Test that the client will use the web search tool
|
||||
response = openai_responses_client.get_streaming_response(
|
||||
messages=[
|
||||
ChatMessage(
|
||||
role="user",
|
||||
text="Who are the main characters of Kpop Demon Hunters? Do a web search to find the answer.",
|
||||
)
|
||||
],
|
||||
tools=[HostedWebSearchTool()],
|
||||
tool_choice="auto",
|
||||
)
|
||||
|
||||
assert response is not None
|
||||
full_message: str = ""
|
||||
async for chunk in response:
|
||||
assert chunk is not None
|
||||
assert isinstance(chunk, ChatResponseUpdate)
|
||||
for content in chunk.contents:
|
||||
if isinstance(content, TextContent) and content.text:
|
||||
full_message += content.text
|
||||
assert "Rumi" in full_message
|
||||
assert "Mira" in full_message
|
||||
assert "Zoey" in full_message
|
||||
|
||||
# Test that the client will use the web search tool with location
|
||||
additional_properties = {
|
||||
"user_location": {
|
||||
"country": "US",
|
||||
"city": "Seattle",
|
||||
}
|
||||
}
|
||||
response = openai_responses_client.get_streaming_response(
|
||||
messages=[ChatMessage(role="user", text="What is the current weather? Do not ask for my current location.")],
|
||||
tools=[HostedWebSearchTool(additional_properties=additional_properties)],
|
||||
tool_choice="auto",
|
||||
)
|
||||
assert response is not None
|
||||
full_message: str = ""
|
||||
async for chunk in response:
|
||||
assert chunk is not None
|
||||
assert isinstance(chunk, ChatResponseUpdate)
|
||||
for content in chunk.contents:
|
||||
if isinstance(content, TextContent) and content.text:
|
||||
full_message += content.text
|
||||
assert "Seattle" in full_message
|
||||
|
||||
|
||||
@skip_if_openai_integration_tests_disabled
|
||||
async def test_openai_responses_client_file_search() -> None:
|
||||
openai_responses_client = OpenAIResponsesClient()
|
||||
|
||||
assert isinstance(openai_responses_client, ChatClient)
|
||||
|
||||
file_id, vector_store = await create_vector_store(openai_responses_client)
|
||||
# Test that the client will use the web search tool
|
||||
response = await openai_responses_client.get_response(
|
||||
messages=[
|
||||
ChatMessage(
|
||||
role="user",
|
||||
text="What is the weather today? Do a file search to find the answer.",
|
||||
)
|
||||
],
|
||||
tools=[HostedFileSearchTool(inputs=vector_store)],
|
||||
tool_choice="auto",
|
||||
)
|
||||
|
||||
await delete_vector_store(openai_responses_client, file_id, vector_store.vector_store_id)
|
||||
assert "sunny" in response.text.lower()
|
||||
assert "75" in response.text
|
||||
|
||||
|
||||
@skip_if_openai_integration_tests_disabled
|
||||
async def test_openai_responses_client_streaming_file_search() -> None:
|
||||
openai_responses_client = OpenAIResponsesClient()
|
||||
|
||||
assert isinstance(openai_responses_client, ChatClient)
|
||||
|
||||
file_id, vector_store = await create_vector_store(openai_responses_client)
|
||||
# Test that the client will use the web search tool
|
||||
response = openai_responses_client.get_streaming_response(
|
||||
messages=[
|
||||
ChatMessage(
|
||||
role="user",
|
||||
text="What is the weather today? Do a file search to find the answer.",
|
||||
)
|
||||
],
|
||||
tools=[HostedFileSearchTool(inputs=vector_store)],
|
||||
tool_choice="auto",
|
||||
)
|
||||
|
||||
assert response is not None
|
||||
full_message: str = ""
|
||||
async for chunk in response:
|
||||
assert chunk is not None
|
||||
assert isinstance(chunk, ChatResponseUpdate)
|
||||
for content in chunk.contents:
|
||||
if isinstance(content, TextContent) and content.text:
|
||||
full_message += content.text
|
||||
|
||||
await delete_vector_store(openai_responses_client, file_id, vector_store.vector_store_id)
|
||||
|
||||
assert "sunny" in full_message.lower()
|
||||
assert "75" in full_message
|
||||
|
||||
+53
@@ -0,0 +1,53 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import asyncio
|
||||
from agent_framework import ChatClientAgent, HostedFileSearchTool, HostedVectorStoreContent
|
||||
from agent_framework.openai import OpenAIAssistantsClient
|
||||
|
||||
# Helper functions
|
||||
|
||||
async def create_vector_store(client: OpenAIAssistantsClient) -> tuple[str, HostedVectorStoreContent]:
|
||||
"""Create a vector store with sample documents."""
|
||||
file = await client.client.files.create(
|
||||
file=("todays_weather.txt", b"The weather today is sunny with a high of 75F."), purpose="user_data"
|
||||
)
|
||||
vector_store = await client.client.vector_stores.create(
|
||||
name="knowledge_base",
|
||||
expires_after={"anchor": "last_active_at", "days": 1},
|
||||
)
|
||||
result = await client.client.vector_stores.files.create_and_poll(vector_store_id=vector_store.id, file_id=file.id)
|
||||
if result.last_error is not None:
|
||||
raise Exception(f"Vector store file processing failed with status: {result.last_error.message}")
|
||||
|
||||
return file.id, HostedVectorStoreContent(vector_store_id=vector_store.id)
|
||||
|
||||
|
||||
async def delete_vector_store(client: OpenAIAssistantsClient, file_id: str, vector_store_id: str) -> None:
|
||||
"""Delete the vector store after using it."""
|
||||
|
||||
await client.client.vector_stores.delete(vector_store_id=vector_store_id)
|
||||
await client.client.files.delete(file_id=file_id)
|
||||
|
||||
|
||||
async def main() -> None:
|
||||
client = OpenAIAssistantsClient()
|
||||
async with ChatClientAgent(
|
||||
chat_client=client,
|
||||
instructions="You are a helpful assistant that searches files in a knowledge base.",
|
||||
tools=HostedFileSearchTool(),
|
||||
) as agent:
|
||||
query = "What is the weather today? Do a file search to find the answer."
|
||||
file_id, vector_store = await create_vector_store(client)
|
||||
|
||||
print(f"User: {query}")
|
||||
print("Agent: ", end="", flush=True)
|
||||
async for chunk in agent.run_streaming(
|
||||
query,
|
||||
tool_resources={"file_search": {"vector_store_ids": [vector_store.vector_store_id]}}
|
||||
):
|
||||
if chunk.text:
|
||||
print(chunk.text, end="", flush=True)
|
||||
await delete_vector_store(client, file_id, vector_store.vector_store_id)
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
+39
@@ -0,0 +1,39 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import asyncio
|
||||
from agent_framework import HostedWebSearchTool
|
||||
from agent_framework.openai import OpenAIChatClient
|
||||
|
||||
async def main() -> None:
|
||||
client = OpenAIChatClient(ai_model_id="gpt-4o-search-preview")
|
||||
|
||||
message = "What is the current weather? Do not ask for my current location."
|
||||
# Test that the client will use the web search tool with location
|
||||
additional_properties = {
|
||||
"user_location": {
|
||||
"country": "US",
|
||||
"city": "Seattle",
|
||||
}
|
||||
}
|
||||
stream = False
|
||||
print(f"User: {message}")
|
||||
if stream:
|
||||
print("Assistant: ", end="")
|
||||
async for chunk in client.get_streaming_response(
|
||||
message,
|
||||
tools=[HostedWebSearchTool(additional_properties=additional_properties)],
|
||||
tool_choice="auto",
|
||||
):
|
||||
if chunk.text:
|
||||
print(chunk.text, end="")
|
||||
print("")
|
||||
else:
|
||||
response = await client.get_response(
|
||||
message,
|
||||
tools=[HostedWebSearchTool(additional_properties=additional_properties)],
|
||||
tool_choice="auto",
|
||||
)
|
||||
print(f"Assistant: {response}")
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
+60
@@ -0,0 +1,60 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import asyncio
|
||||
from agent_framework import HostedFileSearchTool, HostedVectorStoreContent
|
||||
from agent_framework.openai import OpenAIResponsesClient
|
||||
|
||||
# Helper functions
|
||||
|
||||
async def create_vector_store(client: OpenAIResponsesClient) -> tuple[str, HostedVectorStoreContent]:
|
||||
"""Create a vector store with sample documents."""
|
||||
file = await client.client.files.create(
|
||||
file=("todays_weather.txt", b"The weather today is sunny with a high of 75F."), purpose="user_data"
|
||||
)
|
||||
vector_store = await client.client.vector_stores.create(
|
||||
name="knowledge_base",
|
||||
expires_after={"anchor": "last_active_at", "days": 1},
|
||||
)
|
||||
result = await client.client.vector_stores.files.create_and_poll(vector_store_id=vector_store.id, file_id=file.id)
|
||||
if result.last_error is not None:
|
||||
raise Exception(f"Vector store file processing failed with status: {result.last_error.message}")
|
||||
|
||||
return file.id, HostedVectorStoreContent(vector_store_id=vector_store.id)
|
||||
|
||||
|
||||
async def delete_vector_store(client: OpenAIResponsesClient, file_id: str, vector_store_id: str) -> None:
|
||||
"""Delete the vector store after using it."""
|
||||
|
||||
await client.client.vector_stores.delete(vector_store_id=vector_store_id)
|
||||
await client.client.files.delete(file_id=file_id)
|
||||
|
||||
|
||||
async def main() -> None:
|
||||
client = OpenAIResponsesClient()
|
||||
|
||||
message = "What is the weather today? Do a file search to find the answer."
|
||||
|
||||
stream = False
|
||||
print(f"User: {message}")
|
||||
file_id, vector_store = await create_vector_store(client)
|
||||
if stream:
|
||||
print("Assistant: ", end="")
|
||||
async for chunk in client.get_streaming_response(
|
||||
message,
|
||||
tools=[HostedFileSearchTool(inputs=vector_store)],
|
||||
tool_choice="auto",
|
||||
):
|
||||
if chunk.text:
|
||||
print(chunk.text, end="")
|
||||
print("")
|
||||
else:
|
||||
response = await client.get_response(
|
||||
message,
|
||||
tools=[HostedFileSearchTool(inputs=vector_store)],
|
||||
tool_choice="auto",
|
||||
)
|
||||
print(f"Assistant: {response}")
|
||||
await delete_vector_store(client, file_id, vector_store.vector_store_id)
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
+39
@@ -0,0 +1,39 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import asyncio
|
||||
from agent_framework import HostedWebSearchTool
|
||||
from agent_framework.openai import OpenAIResponsesClient
|
||||
|
||||
async def main() -> None:
|
||||
client = OpenAIResponsesClient()
|
||||
|
||||
message = "What is the current weather? Do not ask for my current location."
|
||||
# Test that the client will use the web search tool with location
|
||||
additional_properties = {
|
||||
"user_location": {
|
||||
"country": "US",
|
||||
"city": "Seattle",
|
||||
}
|
||||
}
|
||||
stream = False
|
||||
print(f"User: {message}")
|
||||
if stream:
|
||||
print("Assistant: ", end="")
|
||||
async for chunk in client.get_streaming_response(
|
||||
message,
|
||||
tools=[HostedWebSearchTool(additional_properties=additional_properties)],
|
||||
tool_choice="auto",
|
||||
):
|
||||
if chunk.text:
|
||||
print(chunk.text, end="")
|
||||
print("")
|
||||
else:
|
||||
response = await client.get_response(
|
||||
message,
|
||||
tools=[HostedWebSearchTool(additional_properties=additional_properties)],
|
||||
tool_choice="auto",
|
||||
)
|
||||
print(f"Assistant: {response}")
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
Reference in New Issue
Block a user