From 0410f5177713b09fc64b323ff345021f3f04345b Mon Sep 17 00:00:00 2001 From: peterychang <49209570+peterychang@users.noreply.github.com> Date: Fri, 15 Aug 2025 15:35:23 -0400 Subject: [PATCH] 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> --- .github/workflows/python-merge-tests.yml | 1 + .../packages/main/agent_framework/_clients.py | 4 +- .../packages/main/agent_framework/_tools.py | 96 +++++++- .../openai/_assistants_client.py | 11 +- .../agent_framework/openai/_chat_client.py | 31 ++- .../openai/_responses_client.py | 47 +++- .../openai/test_openai_assistants_client.py | 76 ++++++ .../tests/openai/test_openai_chat_client.py | 103 ++++++++- .../openai/test_openai_responses_client.py | 217 ++++++++++++++++-- .../openai_assistants_with_file_search.py | 53 +++++ .../openai_chat_client_with_web_search.py | 39 ++++ ...penai_responses_client_with_file_search.py | 60 +++++ ...openai_responses_client_with_web_search.py | 39 ++++ 13 files changed, 750 insertions(+), 27 deletions(-) create mode 100644 python/samples/getting_started/agents/openai_assistants_client/openai_assistants_with_file_search.py create mode 100644 python/samples/getting_started/agents/openai_chat_client/openai_chat_client_with_web_search.py create mode 100644 python/samples/getting_started/agents/openai_responses_client/openai_responses_client_with_file_search.py create mode 100644 python/samples/getting_started/agents/openai_responses_client/openai_responses_client_with_web_search.py diff --git a/.github/workflows/python-merge-tests.yml b/.github/workflows/python-merge-tests.yml index 90746b9cd2..6378e4520e 100644 --- a/.github/workflows/python-merge-tests.yml +++ b/.github/workflows/python-merge-tests.yml @@ -57,6 +57,7 @@ jobs: env: UV_PYTHON: ${{ matrix.python-version }} OPENAI_CHAT_MODEL_ID: ${{ vars.OPENAI__CHATMODELID }} + OPENAI_RESPONSES_MODEL_ID: ${{ vars.OPENAI__RESPONSESMODELID }} OPENAI_API_KEY: ${{ secrets.OPENAI__APIKEY }} LOCAL_MCP_URL: ${{ vars.LOCAL_MCP__URL }} PACKAGE_NAME: "main" diff --git a/python/packages/main/agent_framework/_clients.py b/python/packages/main/agent_framework/_clients.py index 68cfbd8ee1..88f7fc7ca0 100644 --- a/python/packages/main/agent_framework/_clients.py +++ b/python/packages/main/agent_framework/_clients.py @@ -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 diff --git a/python/packages/main/agent_framework/_tools.py b/python/packages/main/agent_framework/_tools.py index 9059367471..f3930d82d1 100644 --- a/python/packages/main/agent_framework/_tools.py +++ b/python/packages/main/agent_framework/_tools.py @@ -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. diff --git a/python/packages/main/agent_framework/openai/_assistants_client.py b/python/packages/main/agent_framework/openai/_assistants_client.py index e26459d51e..2cde074875 100644 --- a/python/packages/main/agent_framework/openai/_assistants_client.py +++ b/python/packages/main/agent_framework/openai/_assistants_client.py @@ -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) diff --git a/python/packages/main/agent_framework/openai/_chat_client.py b/python/packages/main/agent_framework/openai/_chat_client.py index 5e7d126182..3a6f122d7f 100644 --- a/python/packages/main/agent_framework/openai/_chat_client.py +++ b/python/packages/main/agent_framework/openai/_chat_client.py @@ -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 ( diff --git a/python/packages/main/agent_framework/openai/_responses_client.py b/python/packages/main/agent_framework/openai/_responses_client.py index 4ad6ad7dbf..af17d513e6 100644 --- a/python/packages/main/agent_framework/openai/_responses_client.py +++ b/python/packages/main/agent_framework/openai/_responses_client.py @@ -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: diff --git a/python/packages/main/tests/openai/test_openai_assistants_client.py b/python/packages/main/tests/openai/test_openai_assistants_client.py index 6dabe8ac43..fae14234cb 100644 --- a/python/packages/main/tests/openai/test_openai_assistants_client.py +++ b/python/packages/main/tests/openai/test_openai_assistants_client.py @@ -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"]) diff --git a/python/packages/main/tests/openai/test_openai_chat_client.py b/python/packages/main/tests/openai/test_openai_chat_client.py index 8ba5efc440..f85a9c0dc3 100644 --- a/python/packages/main/tests/openai/test_openai_chat_client.py +++ b/python/packages/main/tests/openai/test_openai_chat_client.py @@ -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 diff --git a/python/packages/main/tests/openai/test_openai_responses_client.py b/python/packages/main/tests/openai/test_openai_responses_client.py index e610963966..0eafea37df 100644 --- a/python/packages/main/tests/openai/test_openai_responses_client.py +++ b/python/packages/main/tests/openai/test_openai_responses_client.py @@ -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 diff --git a/python/samples/getting_started/agents/openai_assistants_client/openai_assistants_with_file_search.py b/python/samples/getting_started/agents/openai_assistants_client/openai_assistants_with_file_search.py new file mode 100644 index 0000000000..0840b7e1eb --- /dev/null +++ b/python/samples/getting_started/agents/openai_assistants_client/openai_assistants_with_file_search.py @@ -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()) diff --git a/python/samples/getting_started/agents/openai_chat_client/openai_chat_client_with_web_search.py b/python/samples/getting_started/agents/openai_chat_client/openai_chat_client_with_web_search.py new file mode 100644 index 0000000000..e2ba52474d --- /dev/null +++ b/python/samples/getting_started/agents/openai_chat_client/openai_chat_client_with_web_search.py @@ -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()) diff --git a/python/samples/getting_started/agents/openai_responses_client/openai_responses_client_with_file_search.py b/python/samples/getting_started/agents/openai_responses_client/openai_responses_client_with_file_search.py new file mode 100644 index 0000000000..369aae1d43 --- /dev/null +++ b/python/samples/getting_started/agents/openai_responses_client/openai_responses_client_with_file_search.py @@ -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()) diff --git a/python/samples/getting_started/agents/openai_responses_client/openai_responses_client_with_web_search.py b/python/samples/getting_started/agents/openai_responses_client/openai_responses_client_with_web_search.py new file mode 100644 index 0000000000..b6c3511039 --- /dev/null +++ b/python/samples/getting_started/agents/openai_responses_client/openai_responses_client_with_web_search.py @@ -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())