agent/agent-framework
2
0
Fork 0
mirror of https://github.com/microsoft/agent-framework.git synced 2026-06-16 21:04:09 +08:00
Code Issues Packages Projects Releases Wiki Activity

Python: Web search file search tools (#395)

Browse Source 
* 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
Download Patch File Download Diff File Expand all files Collapse all files
python/packages/main/agent_framework/_clients.py
+2 -2
View File
@@ -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
python/packages/main/agent_framework/_tools.py
+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.
python/packages/main/agent_framework/openai/_assistants_client.py
+10 -1
View File
@@ -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)
python/packages/main/agent_framework/openai/_chat_client.py
+28 -3
View File
@@ -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 (
python/packages/main/agent_framework/openai/_responses_client.py
+46 -1
View File
@@ -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: