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>

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"])

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

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