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:
peterychang
2025-08-15 15:35:23 -04:00
committed by GitHub
Unverified
parent ed86baa6cb
commit 0410f51777
13 changed files with 750 additions and 27 deletions
@@ -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
+94 -2
View File
@@ -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
@@ -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())
@@ -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())
@@ -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())
@@ -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())