mirror of
https://github.com/microsoft/agent-framework.git
synced 2026-06-16 21:04:09 +08:00
Python: [BREAKING] Python: Provider-leading client design & OpenAI package extraction (#4818)
* Python: Provider-leading client design & OpenAI package extraction Major refactoring of the Python Agent Framework client architecture: - Extract OpenAI clients into new `agent-framework-openai` package - Core package no longer depends on openai, azure-identity, azure-ai-projects - Rename clients for discoverability: OpenAIResponsesClient → OpenAIChatClient, OpenAIChatClient → OpenAIChatCompletionClient - Unify `model_id`/`deployment_name`/`model_deployment_name` → `model` param - New FoundryChatClient for Azure AI Foundry Responses API - New FoundryAgent/FoundryAgentClient for connecting to pre-configured Foundry agents - Remove OpenAIBase/OpenAIConfigMixin from non-deprecated client MRO - Deprecate AzureOpenAI* clients, AzureAIClient, OpenAIAssistantsClient - Reorganize samples: azure_openai+azure_ai+azure_ai_agent → azure/ - ADR-0020: Provider-Leading Client Design Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * fix: missing Agent imports in samples, .model_id → .model in foundry_local sample Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * fix: CI failures — mypy errors, coverage targets, sample imports - azure-ai mypy: add type ignores for TypedDict total=, model arg, forward ref - Coverage: replace core.azure/openai targets with openai package target - project_provider: add type annotation for opts dict Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * fix: populate openai .pyi stub, fix broken README links, coverage targets Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * fixes * updated observabilitty * reset azure init.pyi * fix errors * updated adr number * fix foundry local * fixed not renamed docstrings and comments, and added deprecated markers to old classes * fix tests and pyprojects * fix test vars * updated function tests * update durable * updated test setup for functions * Fix Foundry auth in workflow samples Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * Stabilize Python integration workflows Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * Update hosting samples for Foundry Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * Trigger full CI rerun Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * Trigger CI rerun again Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * trigger rerun * trigger rerun * fix for litellm * undo durabletask changes * Move Foundry APIs into foundry namespace Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * Fix Foundry pyproject formatting Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * Split provider samples by Foundry surface Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * Restore hosting sample requirements Also fix the Foundry Local sample link after the provider sample move. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * updated tests * udpated foundry integration tests * removed dist from azurefunctions tests * Use separate Foundry clients for concurrent agents Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * fix client setup in azfunc and durable * disabled two tests * updated setup for some function and durable tests * improved azure openai setup with new clients * ignore deprecated * fixes * skip 11 * remove openai assistants int tests --------- Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
This commit is contained in:
Eduard van Valkenburg
committed by
GitHub
Unverified
GitHub
parent
4b533608b6
commit
5e056b672e
@@ -0,0 +1,87 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""OpenAI integration for Microsoft Agent Framework.
|
||||
|
||||
This package provides OpenAI client implementations for the Agent Framework,
|
||||
including clients for the Responses API and Chat Completions API.
|
||||
"""
|
||||
|
||||
import importlib.metadata
|
||||
import sys
|
||||
|
||||
if sys.version_info >= (3, 13):
|
||||
from warnings import deprecated # type: ignore # pragma: no cover
|
||||
else:
|
||||
from typing_extensions import deprecated # type: ignore # pragma: no cover
|
||||
|
||||
from ._assistant_provider import OpenAIAssistantProvider
|
||||
from ._assistants_client import (
|
||||
AssistantToolResources,
|
||||
OpenAIAssistantsClient,
|
||||
OpenAIAssistantsOptions,
|
||||
)
|
||||
from ._chat_client import (
|
||||
OpenAIChatClient,
|
||||
OpenAIChatOptions,
|
||||
OpenAIContinuationToken,
|
||||
RawOpenAIChatClient,
|
||||
)
|
||||
from ._chat_completion_client import (
|
||||
OpenAIChatCompletionClient,
|
||||
OpenAIChatCompletionOptions,
|
||||
RawOpenAIChatCompletionClient,
|
||||
)
|
||||
from ._embedding_client import OpenAIEmbeddingClient, OpenAIEmbeddingOptions
|
||||
from ._exceptions import ContentFilterResultSeverity, OpenAIContentFilterException
|
||||
from ._shared import OpenAISettings
|
||||
|
||||
try:
|
||||
__version__ = importlib.metadata.version("agent-framework-openai")
|
||||
except importlib.metadata.PackageNotFoundError:
|
||||
__version__ = "0.0.0" # Fallback for development mode
|
||||
|
||||
# Deprecated aliases for old names — use subclasses so the warning only fires for the alias
|
||||
|
||||
|
||||
@deprecated(
|
||||
"OpenAIResponsesClient is deprecated, use OpenAIChatClient instead.",
|
||||
category=DeprecationWarning,
|
||||
)
|
||||
class OpenAIResponsesClient(OpenAIChatClient): # type: ignore[misc]
|
||||
"""Deprecated alias for :class:`OpenAIChatClient`."""
|
||||
|
||||
|
||||
@deprecated(
|
||||
"RawOpenAIResponsesClient is deprecated, use RawOpenAIChatClient instead.",
|
||||
category=DeprecationWarning,
|
||||
)
|
||||
class RawOpenAIResponsesClient(RawOpenAIChatClient): # type: ignore[misc]
|
||||
"""Deprecated alias for :class:`RawOpenAIChatClient`."""
|
||||
|
||||
|
||||
OpenAIResponsesOptions = OpenAIChatOptions
|
||||
"""Deprecated alias for :class:`OpenAIChatOptions`."""
|
||||
|
||||
|
||||
__all__ = [
|
||||
"AssistantToolResources",
|
||||
"ContentFilterResultSeverity",
|
||||
"OpenAIAssistantProvider",
|
||||
"OpenAIAssistantsClient",
|
||||
"OpenAIAssistantsOptions",
|
||||
"OpenAIChatClient",
|
||||
"OpenAIChatCompletionClient",
|
||||
"OpenAIChatCompletionOptions",
|
||||
"OpenAIChatOptions",
|
||||
"OpenAIContentFilterException",
|
||||
"OpenAIContinuationToken",
|
||||
"OpenAIEmbeddingClient",
|
||||
"OpenAIEmbeddingOptions",
|
||||
"OpenAIResponsesClient",
|
||||
"OpenAIResponsesOptions",
|
||||
"OpenAISettings",
|
||||
"RawOpenAIChatClient",
|
||||
"RawOpenAIChatCompletionClient",
|
||||
"RawOpenAIResponsesClient",
|
||||
"__version__",
|
||||
]
|
||||
@@ -0,0 +1,564 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import sys
|
||||
from collections.abc import Awaitable, Callable, Mapping, MutableMapping, Sequence
|
||||
from typing import TYPE_CHECKING, Any, Generic, cast
|
||||
|
||||
from agent_framework._agents import Agent
|
||||
from agent_framework._middleware import MiddlewareTypes
|
||||
from agent_framework._sessions import BaseContextProvider
|
||||
from agent_framework._settings import SecretString, load_settings
|
||||
from agent_framework._tools import FunctionTool, ToolTypes, normalize_tools
|
||||
from openai import AsyncOpenAI
|
||||
from openai.types.beta.assistant import Assistant
|
||||
from pydantic import BaseModel
|
||||
|
||||
from ._assistants_client import OpenAIAssistantsClient
|
||||
from ._shared import OpenAISettings, from_assistant_tools, to_assistant_tools
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ._assistants_client import OpenAIAssistantsOptions
|
||||
|
||||
if sys.version_info >= (3, 13):
|
||||
from typing import TypeVar # type:ignore # pragma: no cover
|
||||
else:
|
||||
from typing_extensions import TypeVar # type:ignore # pragma: no cover
|
||||
if sys.version_info >= (3, 11):
|
||||
from typing import Self, TypedDict # type:ignore # pragma: no cover
|
||||
else:
|
||||
from typing_extensions import Self, TypedDict # type:ignore # pragma: no cover
|
||||
|
||||
|
||||
# Type variable for options - allows typed OpenAIAssistantProvider[OptionsCoT] returns
|
||||
# Default matches OpenAIAssistantsClient's default options type
|
||||
OptionsCoT = TypeVar(
|
||||
"OptionsCoT",
|
||||
bound=TypedDict, # type: ignore[valid-type]
|
||||
default="OpenAIAssistantsOptions",
|
||||
covariant=True,
|
||||
)
|
||||
|
||||
|
||||
class OpenAIAssistantProvider(Generic[OptionsCoT]):
|
||||
"""Provider for creating Agent instances from OpenAI Assistants API.
|
||||
|
||||
This provider allows you to create, retrieve, and wrap OpenAI Assistants
|
||||
as Agent instances for use in the agent framework.
|
||||
|
||||
Examples:
|
||||
Basic usage with automatic client creation:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
from agent_framework.openai import OpenAIAssistantProvider
|
||||
|
||||
# Uses OPENAI_API_KEY environment variable
|
||||
provider = OpenAIAssistantProvider()
|
||||
|
||||
# Create a new assistant
|
||||
agent = await provider.create_agent(
|
||||
name="MyAssistant",
|
||||
model="gpt-4",
|
||||
instructions="You are a helpful assistant.",
|
||||
tools=[my_function],
|
||||
)
|
||||
|
||||
result = await agent.run("Hello!")
|
||||
|
||||
Using an existing client:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
from openai import AsyncOpenAI
|
||||
from agent_framework.openai import OpenAIAssistantProvider
|
||||
|
||||
client = AsyncOpenAI()
|
||||
provider = OpenAIAssistantProvider(client)
|
||||
|
||||
# Get an existing assistant by ID
|
||||
agent = await provider.get_agent(
|
||||
assistant_id="asst_123",
|
||||
tools=[my_function], # Provide implementations for function tools
|
||||
)
|
||||
|
||||
Wrapping an SDK Assistant object:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
# Fetch assistant directly via SDK
|
||||
assistant = await client.beta.assistants.retrieve("asst_123")
|
||||
|
||||
# Wrap without additional HTTP call
|
||||
agent = provider.as_agent(assistant, tools=[my_function])
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
client: AsyncOpenAI | None = None,
|
||||
*,
|
||||
api_key: str | SecretString | Callable[[], str | Awaitable[str]] | None = None,
|
||||
org_id: str | None = None,
|
||||
base_url: str | None = None,
|
||||
env_file_path: str | None = None,
|
||||
env_file_encoding: str | None = None,
|
||||
) -> None:
|
||||
"""Initialize the OpenAI Assistant Provider.
|
||||
|
||||
Args:
|
||||
client: An existing AsyncOpenAI client to use. If not provided,
|
||||
a new client will be created using the other parameters.
|
||||
|
||||
Keyword Args:
|
||||
api_key: OpenAI API key. Can also be set via OPENAI_API_KEY env var.
|
||||
org_id: OpenAI organization ID. Can also be set via OPENAI_ORG_ID env var.
|
||||
base_url: Base URL for the OpenAI API. Can also be set via OPENAI_BASE_URL env var.
|
||||
env_file_path: Path to .env file for configuration.
|
||||
env_file_encoding: Encoding of the .env file.
|
||||
|
||||
Raises:
|
||||
ValueError: If no client is provided and API key is missing.
|
||||
|
||||
Examples:
|
||||
.. code-block:: python
|
||||
|
||||
# Using environment variables
|
||||
provider = OpenAIAssistantProvider()
|
||||
|
||||
# Using explicit API key
|
||||
provider = OpenAIAssistantProvider(api_key="sk-...")
|
||||
|
||||
# Using existing client
|
||||
client = AsyncOpenAI()
|
||||
provider = OpenAIAssistantProvider(client)
|
||||
"""
|
||||
self._client: AsyncOpenAI | None = client
|
||||
self._should_close_client: bool = client is None
|
||||
|
||||
if client is None:
|
||||
# Load settings and create client
|
||||
settings = load_settings(
|
||||
OpenAISettings,
|
||||
env_prefix="OPENAI_",
|
||||
api_key=api_key,
|
||||
org_id=org_id,
|
||||
base_url=base_url,
|
||||
env_file_path=env_file_path,
|
||||
env_file_encoding=env_file_encoding,
|
||||
)
|
||||
|
||||
api_key_setting = settings.get("api_key")
|
||||
if not api_key_setting:
|
||||
raise ValueError(
|
||||
"OpenAI API key is required. Set via 'api_key' parameter or 'OPENAI_API_KEY' environment variable."
|
||||
)
|
||||
|
||||
# Get API key value
|
||||
api_key_value: str | Callable[[], str | Awaitable[str]]
|
||||
if isinstance(api_key_setting, SecretString):
|
||||
api_key_value = api_key_setting.get_secret_value()
|
||||
else:
|
||||
api_key_value = api_key_setting
|
||||
|
||||
# Create client
|
||||
client_args: dict[str, Any] = {"api_key": api_key_value}
|
||||
if org_id_value := settings.get("org_id"):
|
||||
client_args["organization"] = org_id_value
|
||||
if base_url_value := settings.get("base_url"):
|
||||
client_args["base_url"] = base_url_value
|
||||
|
||||
self._client = AsyncOpenAI(**client_args)
|
||||
|
||||
async def __aenter__(self) -> Self:
|
||||
"""Async context manager entry."""
|
||||
return self
|
||||
|
||||
async def __aexit__(self, exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: Any) -> None:
|
||||
"""Async context manager exit."""
|
||||
await self.close()
|
||||
|
||||
async def close(self) -> None:
|
||||
"""Close the provider and clean up resources.
|
||||
|
||||
If the provider created its own client, it will be closed.
|
||||
If an external client was provided, it will not be closed.
|
||||
"""
|
||||
if self._should_close_client and self._client is not None:
|
||||
await self._client.close()
|
||||
|
||||
async def create_agent(
|
||||
self,
|
||||
*,
|
||||
name: str,
|
||||
model: str,
|
||||
instructions: str | None = None,
|
||||
description: str | None = None,
|
||||
tools: ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None = None,
|
||||
metadata: dict[str, str] | None = None,
|
||||
default_options: OptionsCoT | None = None,
|
||||
middleware: Sequence[MiddlewareTypes] | None = None,
|
||||
context_providers: Sequence[BaseContextProvider] | None = None,
|
||||
) -> Agent[OptionsCoT]:
|
||||
"""Create a new assistant on OpenAI and return a Agent.
|
||||
|
||||
This method creates a new assistant on the OpenAI service and wraps it
|
||||
in a Agent instance. The assistant will persist on OpenAI until deleted.
|
||||
|
||||
Keyword Args:
|
||||
name: The name of the assistant (required).
|
||||
model: The model ID to use, e.g., "gpt-4", "gpt-4o" (required).
|
||||
instructions: System instructions for the assistant.
|
||||
description: A description of the assistant.
|
||||
tools: Tools available to the assistant. Can include:
|
||||
- FunctionTool instances or callables decorated with @tool
|
||||
- Dict-based tools from OpenAIAssistantsClient.get_code_interpreter_tool()
|
||||
- Dict-based tools from OpenAIAssistantsClient.get_file_search_tool()
|
||||
- Raw tool dictionaries
|
||||
metadata: Metadata to attach to the assistant (max 16 key-value pairs).
|
||||
default_options: A TypedDict containing default chat options for the agent.
|
||||
These options are applied to every run unless overridden.
|
||||
Include ``response_format`` here for structured output responses.
|
||||
middleware: MiddlewareTypes for the Agent.
|
||||
context_providers: Context providers for the Agent.
|
||||
|
||||
Returns:
|
||||
A Agent instance wrapping the created assistant.
|
||||
|
||||
Raises:
|
||||
ValueError: If assistant creation fails.
|
||||
|
||||
Examples:
|
||||
.. code-block:: python
|
||||
|
||||
provider = OpenAIAssistantProvider()
|
||||
|
||||
# Create with function tools
|
||||
agent = await provider.create_agent(
|
||||
name="WeatherBot",
|
||||
model="gpt-4",
|
||||
instructions="You are a helpful weather assistant.",
|
||||
tools=[get_weather],
|
||||
)
|
||||
|
||||
# Create with structured output
|
||||
agent = await provider.create_agent(
|
||||
name="StructuredBot",
|
||||
model="gpt-4",
|
||||
default_options={"response_format": MyPydanticModel},
|
||||
)
|
||||
"""
|
||||
# Normalize tools
|
||||
normalized_tools = normalize_tools(tools)
|
||||
assistant_tools: list[FunctionTool | MutableMapping[str, Any]] = [
|
||||
tool for tool in normalized_tools if isinstance(tool, (FunctionTool, MutableMapping))
|
||||
]
|
||||
api_tools = to_assistant_tools(assistant_tools) if assistant_tools else []
|
||||
|
||||
# Extract response_format from default_options if present
|
||||
opts = dict(default_options) if default_options else {}
|
||||
response_format = opts.get("response_format")
|
||||
|
||||
# Build assistant creation parameters
|
||||
create_params: dict[str, Any] = {
|
||||
"model": model,
|
||||
"name": name,
|
||||
}
|
||||
|
||||
if instructions is not None:
|
||||
create_params["instructions"] = instructions
|
||||
if description is not None:
|
||||
create_params["description"] = description
|
||||
if api_tools:
|
||||
create_params["tools"] = api_tools
|
||||
if metadata is not None:
|
||||
create_params["metadata"] = metadata
|
||||
|
||||
# Handle response format for OpenAI API
|
||||
if response_format is not None and isinstance(response_format, type) and issubclass(response_format, BaseModel):
|
||||
create_params["response_format"] = {
|
||||
"type": "json_schema",
|
||||
"json_schema": {
|
||||
"name": response_format.__name__,
|
||||
"schema": response_format.model_json_schema(),
|
||||
"strict": True,
|
||||
},
|
||||
}
|
||||
|
||||
# Create the assistant
|
||||
if not self._client:
|
||||
raise RuntimeError("OpenAI client is not initialized.")
|
||||
|
||||
assistant = await self._client.beta.assistants.create(**create_params) # type: ignore[reportDeprecated]
|
||||
|
||||
# Create Agent - pass default_options which contains response_format
|
||||
return self._create_chat_agent_from_assistant(
|
||||
assistant=assistant,
|
||||
tools=normalized_tools,
|
||||
instructions=instructions,
|
||||
middleware=middleware,
|
||||
context_providers=context_providers,
|
||||
default_options=default_options,
|
||||
)
|
||||
|
||||
async def get_agent(
|
||||
self,
|
||||
assistant_id: str,
|
||||
*,
|
||||
tools: ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None = None,
|
||||
instructions: str | None = None,
|
||||
default_options: OptionsCoT | None = None,
|
||||
middleware: Sequence[MiddlewareTypes] | None = None,
|
||||
context_providers: Sequence[BaseContextProvider] | None = None,
|
||||
) -> Agent[OptionsCoT]:
|
||||
"""Retrieve an existing assistant by ID and return a Agent.
|
||||
|
||||
This method fetches an existing assistant from OpenAI by its ID
|
||||
and wraps it in a Agent instance.
|
||||
|
||||
Args:
|
||||
assistant_id: The ID of the assistant to retrieve (e.g., "asst_123").
|
||||
|
||||
Keyword Args:
|
||||
tools: Function tools to make available. IMPORTANT: If the assistant
|
||||
was created with function tools, you MUST provide matching
|
||||
implementations here. Hosted tools (code_interpreter, file_search)
|
||||
are automatically included.
|
||||
instructions: Override the assistant's instructions (optional).
|
||||
default_options: A TypedDict containing default chat options for the agent.
|
||||
These options are applied to every run unless overridden.
|
||||
middleware: MiddlewareTypes for the Agent.
|
||||
context_providers: Context providers for the Agent.
|
||||
|
||||
Returns:
|
||||
A Agent instance wrapping the retrieved assistant.
|
||||
|
||||
Raises:
|
||||
RuntimeError: If the assistant cannot be retrieved.
|
||||
ValueError: If required function tools are missing.
|
||||
|
||||
Examples:
|
||||
.. code-block:: python
|
||||
|
||||
provider = OpenAIAssistantProvider()
|
||||
|
||||
# Get assistant without function tools
|
||||
agent = await provider.get_agent(assistant_id="asst_123")
|
||||
|
||||
# Get assistant with function tools
|
||||
agent = await provider.get_agent(
|
||||
assistant_id="asst_456",
|
||||
tools=[get_weather, search_database], # Implementations required!
|
||||
)
|
||||
"""
|
||||
# Fetch the assistant
|
||||
if not self._client:
|
||||
raise RuntimeError("OpenAI client is not initialized.")
|
||||
|
||||
assistant = await self._client.beta.assistants.retrieve(assistant_id) # type: ignore[reportDeprecated]
|
||||
|
||||
# Use as_agent to wrap it
|
||||
return self.as_agent(
|
||||
assistant=assistant,
|
||||
tools=tools,
|
||||
instructions=instructions,
|
||||
default_options=default_options,
|
||||
middleware=middleware,
|
||||
context_providers=context_providers,
|
||||
)
|
||||
|
||||
def as_agent(
|
||||
self,
|
||||
assistant: Assistant,
|
||||
*,
|
||||
tools: ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None = None,
|
||||
instructions: str | None = None,
|
||||
default_options: OptionsCoT | None = None,
|
||||
middleware: Sequence[MiddlewareTypes] | None = None,
|
||||
context_providers: Sequence[BaseContextProvider] | None = None,
|
||||
) -> Agent[OptionsCoT]:
|
||||
"""Wrap an existing SDK Assistant object as a Agent.
|
||||
|
||||
This method does NOT make any HTTP calls. It simply wraps an already-
|
||||
fetched Assistant object in a Agent.
|
||||
|
||||
Args:
|
||||
assistant: The OpenAI Assistant SDK object to wrap.
|
||||
|
||||
Keyword Args:
|
||||
tools: Function tools to make available. If the assistant has
|
||||
function tools defined, you MUST provide matching implementations.
|
||||
Hosted tools (code_interpreter, file_search) are automatically included.
|
||||
instructions: Override the assistant's instructions (optional).
|
||||
default_options: A TypedDict containing default chat options for the agent.
|
||||
These options are applied to every run unless overridden.
|
||||
middleware: MiddlewareTypes for the Agent.
|
||||
context_providers: Context providers for the Agent.
|
||||
|
||||
Returns:
|
||||
A Agent instance wrapping the assistant.
|
||||
|
||||
Raises:
|
||||
ValueError: If required function tools are missing.
|
||||
|
||||
Examples:
|
||||
.. code-block:: python
|
||||
|
||||
client = AsyncOpenAI()
|
||||
provider = OpenAIAssistantProvider(client)
|
||||
|
||||
# Fetch assistant via SDK
|
||||
assistant = await client.beta.assistants.retrieve("asst_123")
|
||||
|
||||
# Wrap without additional HTTP call
|
||||
agent = provider.as_agent(
|
||||
assistant,
|
||||
tools=[my_function],
|
||||
instructions="Custom instructions override",
|
||||
)
|
||||
"""
|
||||
# Validate that required function tools are provided
|
||||
self._validate_function_tools(assistant.tools or [], tools)
|
||||
|
||||
# Merge hosted tools with user-provided function tools
|
||||
merged_tools = self._merge_tools(assistant.tools or [], tools)
|
||||
|
||||
# Create Agent
|
||||
return self._create_chat_agent_from_assistant(
|
||||
assistant=assistant,
|
||||
tools=merged_tools,
|
||||
instructions=instructions,
|
||||
default_options=default_options,
|
||||
middleware=middleware,
|
||||
context_providers=context_providers,
|
||||
)
|
||||
|
||||
def _validate_function_tools(
|
||||
self,
|
||||
assistant_tools: list[Any],
|
||||
provided_tools: ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None,
|
||||
) -> None:
|
||||
"""Validate that required function tools are provided.
|
||||
|
||||
Args:
|
||||
assistant_tools: Tools defined on the assistant.
|
||||
provided_tools: Tools provided by the user.
|
||||
|
||||
Raises:
|
||||
ValueError: If a required function tool is missing.
|
||||
"""
|
||||
# Get function tool names from assistant
|
||||
required_functions: set[str] = set()
|
||||
for tool in assistant_tools:
|
||||
if (
|
||||
hasattr(tool, "type")
|
||||
and tool.type == "function"
|
||||
and hasattr(tool, "function")
|
||||
and hasattr(tool.function, "name")
|
||||
):
|
||||
required_functions.add(tool.function.name)
|
||||
|
||||
if not required_functions:
|
||||
return # No function tools required
|
||||
|
||||
# Get provided function names using normalize_tools
|
||||
provided_functions: set[str] = set()
|
||||
if provided_tools is not None:
|
||||
normalized = normalize_tools(provided_tools)
|
||||
for tool in normalized:
|
||||
if isinstance(tool, FunctionTool):
|
||||
provided_functions.add(tool.name)
|
||||
elif isinstance(tool, Mapping):
|
||||
typed_tool = cast(Mapping[str, Any], tool)
|
||||
raw_func_spec = typed_tool.get("function")
|
||||
if isinstance(raw_func_spec, Mapping):
|
||||
typed_func_spec = cast(Mapping[str, Any], raw_func_spec)
|
||||
raw_name = typed_func_spec.get("name")
|
||||
if isinstance(raw_name, str) and raw_name:
|
||||
provided_functions.add(raw_name)
|
||||
|
||||
# Check for missing functions
|
||||
missing = required_functions - provided_functions
|
||||
if missing:
|
||||
missing_list = ", ".join(sorted(missing))
|
||||
raise ValueError(
|
||||
f"Assistant requires function tool(s) '{missing_list}' but no implementation was provided. "
|
||||
f"Please pass the function implementation(s) in the 'tools' parameter."
|
||||
)
|
||||
|
||||
def _merge_tools(
|
||||
self,
|
||||
assistant_tools: list[Any],
|
||||
user_tools: ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None,
|
||||
) -> list[FunctionTool | MutableMapping[str, Any] | Any]:
|
||||
"""Merge hosted tools from assistant with user-provided function tools.
|
||||
|
||||
Args:
|
||||
assistant_tools: Tools defined on the assistant.
|
||||
user_tools: Tools provided by the user.
|
||||
|
||||
Returns:
|
||||
A list of all tools (hosted tools + user function implementations).
|
||||
"""
|
||||
merged: list[FunctionTool | MutableMapping[str, Any] | Any] = []
|
||||
|
||||
# Add hosted tools from assistant using shared conversion
|
||||
hosted_tools = from_assistant_tools(assistant_tools)
|
||||
merged.extend(hosted_tools)
|
||||
|
||||
# Add user-provided tools (normalized)
|
||||
if user_tools is not None:
|
||||
normalized_user_tools = normalize_tools(user_tools)
|
||||
merged.extend(normalized_user_tools)
|
||||
|
||||
return merged
|
||||
|
||||
def _create_chat_agent_from_assistant(
|
||||
self,
|
||||
assistant: Assistant,
|
||||
tools: list[FunctionTool | MutableMapping[str, Any] | Any] | None,
|
||||
instructions: str | None,
|
||||
middleware: Sequence[MiddlewareTypes] | None,
|
||||
context_providers: Sequence[BaseContextProvider] | None,
|
||||
default_options: OptionsCoT | None = None,
|
||||
**kwargs: Any,
|
||||
) -> Agent[OptionsCoT]:
|
||||
"""Create a Agent from an Assistant.
|
||||
|
||||
Args:
|
||||
assistant: The OpenAI Assistant object.
|
||||
tools: Tools for the agent.
|
||||
instructions: Instructions override.
|
||||
middleware: MiddlewareTypes for the agent.
|
||||
context_providers: Context providers for the agent.
|
||||
default_options: Default chat options for the agent (may include response_format).
|
||||
**kwargs: Additional arguments passed to Agent.
|
||||
|
||||
Returns:
|
||||
A configured Agent instance.
|
||||
"""
|
||||
# Create the chat client with the assistant
|
||||
client = OpenAIAssistantsClient(
|
||||
model=assistant.model,
|
||||
assistant_id=assistant.id,
|
||||
assistant_name=assistant.name,
|
||||
assistant_description=assistant.description,
|
||||
async_client=self._client,
|
||||
)
|
||||
|
||||
# Use instructions from assistant if not overridden
|
||||
final_instructions = instructions if instructions is not None else assistant.instructions
|
||||
|
||||
# Create and return Agent
|
||||
return Agent(
|
||||
client=client,
|
||||
id=assistant.id,
|
||||
name=assistant.name,
|
||||
description=assistant.description,
|
||||
instructions=final_instructions,
|
||||
tools=tools if tools else None,
|
||||
middleware=middleware,
|
||||
context_providers=context_providers,
|
||||
default_options=default_options, # type: ignore[arg-type]
|
||||
**kwargs,
|
||||
)
|
||||
@@ -0,0 +1,958 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
import sys
|
||||
from collections.abc import (
|
||||
AsyncIterable,
|
||||
Awaitable,
|
||||
Callable,
|
||||
Mapping,
|
||||
MutableMapping,
|
||||
Sequence,
|
||||
)
|
||||
from typing import TYPE_CHECKING, Any, Generic, Literal, TypedDict, cast
|
||||
|
||||
from agent_framework._clients import BaseChatClient
|
||||
from agent_framework._middleware import ChatMiddlewareLayer
|
||||
from agent_framework._settings import load_settings
|
||||
from agent_framework._tools import (
|
||||
FunctionInvocationConfiguration,
|
||||
FunctionInvocationLayer,
|
||||
FunctionTool,
|
||||
normalize_tools,
|
||||
)
|
||||
from agent_framework._types import (
|
||||
Annotation,
|
||||
ChatOptions,
|
||||
ChatResponse,
|
||||
ChatResponseUpdate,
|
||||
Content,
|
||||
Message,
|
||||
ResponseStream,
|
||||
TextSpanRegion,
|
||||
UsageDetails,
|
||||
)
|
||||
from agent_framework.observability import ChatTelemetryLayer
|
||||
from openai import AsyncOpenAI
|
||||
from openai.types.beta.threads import (
|
||||
FileCitationAnnotation,
|
||||
FileCitationDeltaAnnotation,
|
||||
FilePathAnnotation,
|
||||
FilePathDeltaAnnotation,
|
||||
ImageURLContentBlockParam,
|
||||
ImageURLParam,
|
||||
MessageContentPartParam,
|
||||
MessageDeltaEvent,
|
||||
Run,
|
||||
TextContentBlockParam,
|
||||
TextDeltaBlock,
|
||||
)
|
||||
from openai.types.beta.threads import (
|
||||
Message as ThreadMessage,
|
||||
)
|
||||
from openai.types.beta.threads.run_create_params import AdditionalMessage
|
||||
from openai.types.beta.threads.run_submit_tool_outputs_params import ToolOutput
|
||||
from openai.types.beta.threads.runs import RunStep
|
||||
from pydantic import BaseModel
|
||||
|
||||
from ._shared import OpenAIConfigMixin, OpenAISettings
|
||||
|
||||
if sys.version_info >= (3, 13):
|
||||
from typing import TypeVar # type: ignore # pragma: no cover
|
||||
else:
|
||||
from typing_extensions import TypeVar # type: ignore # pragma: no cover
|
||||
|
||||
if sys.version_info >= (3, 12):
|
||||
from typing import override # type: ignore # pragma: no cover
|
||||
else:
|
||||
from typing_extensions import override # type: ignore # pragma: no cover
|
||||
|
||||
if sys.version_info >= (3, 11):
|
||||
from typing import Self, TypedDict # type: ignore # pragma: no cover
|
||||
else:
|
||||
from typing_extensions import Self, TypedDict # type: ignore # pragma: no cover
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from agent_framework._middleware import MiddlewareTypes
|
||||
|
||||
logger = logging.getLogger("agent_framework.openai")
|
||||
|
||||
|
||||
# region OpenAI Assistants Options TypedDict
|
||||
|
||||
ResponseModelT = TypeVar("ResponseModelT", bound=BaseModel | None, default=None)
|
||||
|
||||
|
||||
class VectorStoreToolResource(TypedDict, total=False):
|
||||
"""Vector store configuration for file search tool resources."""
|
||||
|
||||
vector_store_ids: list[str]
|
||||
"""IDs of vector stores attached to this assistant."""
|
||||
|
||||
|
||||
class CodeInterpreterToolResource(TypedDict, total=False):
|
||||
"""Code interpreter tool resource configuration."""
|
||||
|
||||
file_ids: list[str]
|
||||
"""File IDs accessible by the code interpreter tool. Max 20 files per assistant."""
|
||||
|
||||
|
||||
class AssistantToolResources(TypedDict, total=False):
|
||||
"""Tool resources attached to the assistant.
|
||||
|
||||
See: https://platform.openai.com/docs/api-reference/assistants/createAssistant#assistants-createassistant-tool_resources
|
||||
"""
|
||||
|
||||
code_interpreter: CodeInterpreterToolResource
|
||||
"""Resources for code interpreter tool, including file IDs."""
|
||||
|
||||
file_search: VectorStoreToolResource
|
||||
"""Resources for file search tool, including vector store IDs."""
|
||||
|
||||
|
||||
class OpenAIAssistantsOptions(ChatOptions[ResponseModelT], Generic[ResponseModelT], total=False):
|
||||
"""OpenAI Assistants API-specific options dict.
|
||||
|
||||
Extends base ChatOptions with Assistants API-specific parameters
|
||||
for creating and running assistants.
|
||||
|
||||
See: https://platform.openai.com/docs/api-reference/assistants
|
||||
|
||||
Keys:
|
||||
# Inherited from ChatOptions:
|
||||
model_id: Deprecated. The model to use for the assistant,
|
||||
translates to ``model`` in OpenAI API.
|
||||
temperature: Sampling temperature between 0 and 2.
|
||||
top_p: Nucleus sampling parameter.
|
||||
max_tokens: Maximum number of tokens to generate,
|
||||
translates to ``max_completion_tokens`` in OpenAI API.
|
||||
tools: List of tools (functions, code_interpreter, file_search).
|
||||
tool_choice: How the model should use tools.
|
||||
allow_multiple_tool_calls: Whether to allow parallel tool calls,
|
||||
translates to ``parallel_tool_calls`` in OpenAI API.
|
||||
response_format: Structured output schema.
|
||||
metadata: Request metadata for tracking.
|
||||
|
||||
# Options not supported in Assistants API (inherited but unused):
|
||||
stop: Not supported.
|
||||
seed: Not supported (use assistant-level configuration instead).
|
||||
frequency_penalty: Not supported.
|
||||
presence_penalty: Not supported.
|
||||
user: Not supported.
|
||||
store: Not supported.
|
||||
|
||||
# Assistants-specific options:
|
||||
name: Name of the assistant.
|
||||
description: Description of the assistant.
|
||||
instructions: System instructions for the assistant.
|
||||
tool_resources: Resources for tools (file IDs, vector stores).
|
||||
reasoning_effort: Effort level for o-series reasoning models.
|
||||
conversation_id: Thread ID to continue conversation in.
|
||||
"""
|
||||
|
||||
# Assistants-specific options
|
||||
name: str
|
||||
"""Name of the assistant (max 256 characters)."""
|
||||
|
||||
description: str
|
||||
"""Description of the assistant (max 512 characters)."""
|
||||
|
||||
tool_resources: AssistantToolResources
|
||||
"""Tool-specific resources like file IDs and vector stores."""
|
||||
|
||||
reasoning_effort: Literal["low", "medium", "high"]
|
||||
"""Effort level for o-series reasoning models (o1, o3-mini).
|
||||
Higher effort = more reasoning time and potentially better results."""
|
||||
|
||||
conversation_id: str # type: ignore[misc]
|
||||
"""Thread ID to continue a conversation in an existing thread."""
|
||||
|
||||
# OpenAI/ChatOptions fields not supported in Assistants API
|
||||
stop: None # type: ignore[misc]
|
||||
"""Not supported in Assistants API."""
|
||||
|
||||
seed: None # type: ignore[misc]
|
||||
"""Not supported in Assistants API (use assistant-level configuration)."""
|
||||
|
||||
frequency_penalty: None # type: ignore[misc]
|
||||
"""Not supported in Assistants API."""
|
||||
|
||||
presence_penalty: None # type: ignore[misc]
|
||||
"""Not supported in Assistants API."""
|
||||
|
||||
user: None # type: ignore[misc]
|
||||
"""Not supported in Assistants API."""
|
||||
|
||||
store: None # type: ignore[misc]
|
||||
"""Not supported in Assistants API."""
|
||||
|
||||
|
||||
ASSISTANTS_OPTION_TRANSLATIONS: dict[str, str] = {
|
||||
"model_id": "model", # backward compat: accept model_id in options
|
||||
"max_tokens": "max_completion_tokens",
|
||||
"allow_multiple_tool_calls": "parallel_tool_calls",
|
||||
}
|
||||
"""Maps ChatOptions keys to OpenAI Assistants API parameter names."""
|
||||
|
||||
OpenAIAssistantsOptionsT = TypeVar(
|
||||
"OpenAIAssistantsOptionsT",
|
||||
bound=TypedDict, # type: ignore[valid-type]
|
||||
default="OpenAIAssistantsOptions",
|
||||
covariant=True,
|
||||
)
|
||||
|
||||
|
||||
# endregion
|
||||
|
||||
|
||||
class OpenAIAssistantsClient( # type: ignore[misc]
|
||||
OpenAIConfigMixin,
|
||||
FunctionInvocationLayer[OpenAIAssistantsOptionsT],
|
||||
ChatMiddlewareLayer[OpenAIAssistantsOptionsT],
|
||||
ChatTelemetryLayer[OpenAIAssistantsOptionsT],
|
||||
BaseChatClient[OpenAIAssistantsOptionsT],
|
||||
Generic[OpenAIAssistantsOptionsT],
|
||||
):
|
||||
"""OpenAI Assistants client with middleware, telemetry, and function invocation support."""
|
||||
|
||||
# region Hosted Tool Factory Methods
|
||||
|
||||
@staticmethod
|
||||
def get_code_interpreter_tool() -> dict[str, Any]:
|
||||
"""Create a code interpreter tool configuration for the Assistants API.
|
||||
|
||||
Returns:
|
||||
A dict tool configuration ready to pass to ChatAgent.
|
||||
|
||||
Examples:
|
||||
.. code-block:: python
|
||||
|
||||
from agent_framework.openai import OpenAIAssistantsClient
|
||||
|
||||
# Enable code interpreter
|
||||
tool = OpenAIAssistantsClient.get_code_interpreter_tool()
|
||||
|
||||
agent = ChatAgent(client, tools=[tool])
|
||||
"""
|
||||
return {"type": "code_interpreter"}
|
||||
|
||||
@staticmethod
|
||||
def get_file_search_tool(
|
||||
*,
|
||||
max_num_results: int | None = None,
|
||||
) -> dict[str, Any]:
|
||||
"""Create a file search tool configuration for the Assistants API.
|
||||
|
||||
Keyword Args:
|
||||
max_num_results: Maximum number of results to return from file search.
|
||||
|
||||
Returns:
|
||||
A dict tool configuration ready to pass to ChatAgent.
|
||||
|
||||
Examples:
|
||||
.. code-block:: python
|
||||
|
||||
from agent_framework.openai import OpenAIAssistantsClient
|
||||
|
||||
# Basic file search
|
||||
tool = OpenAIAssistantsClient.get_file_search_tool()
|
||||
|
||||
# With result limit
|
||||
tool = OpenAIAssistantsClient.get_file_search_tool(max_num_results=10)
|
||||
|
||||
agent = ChatAgent(client, tools=[tool])
|
||||
"""
|
||||
tool: dict[str, Any] = {"type": "file_search"}
|
||||
|
||||
if max_num_results is not None:
|
||||
tool["file_search"] = {"max_num_results": max_num_results}
|
||||
|
||||
return tool
|
||||
|
||||
# endregion
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
model: str | None = None,
|
||||
model_id: str | None = None,
|
||||
assistant_id: str | None = None,
|
||||
assistant_name: str | None = None,
|
||||
assistant_description: str | None = None,
|
||||
thread_id: str | None = None,
|
||||
api_key: str | Callable[[], str | Awaitable[str]] | None = None,
|
||||
org_id: str | None = None,
|
||||
base_url: str | None = None,
|
||||
default_headers: Mapping[str, str] | None = None,
|
||||
async_client: AsyncOpenAI | None = None,
|
||||
env_file_path: str | None = None,
|
||||
env_file_encoding: str | None = None,
|
||||
middleware: Sequence[MiddlewareTypes] | None = None,
|
||||
function_invocation_configuration: FunctionInvocationConfiguration | None = None,
|
||||
**kwargs: Any,
|
||||
) -> None:
|
||||
"""Initialize an OpenAI Assistants client.
|
||||
|
||||
Keyword Args:
|
||||
model: OpenAI model name, see https://platform.openai.com/docs/models.
|
||||
Can also be set via environment variable OPENAI_MODEL.
|
||||
model_id: Deprecated alias for ``model``.
|
||||
assistant_id: The ID of an OpenAI assistant to use.
|
||||
If not provided, a new assistant will be created (and deleted after the request).
|
||||
assistant_name: The name to use when creating new assistants.
|
||||
assistant_description: The description to use when creating new assistants.
|
||||
thread_id: Default thread ID to use for conversations. Can be overridden by
|
||||
conversation_id property when making a request.
|
||||
If not provided, a new thread will be created (and deleted after the request).
|
||||
api_key: The API key to use. If provided will override the env vars or .env file value.
|
||||
Can also be set via environment variable OPENAI_API_KEY.
|
||||
org_id: The org ID to use. If provided will override the env vars or .env file value.
|
||||
Can also be set via environment variable OPENAI_ORG_ID.
|
||||
base_url: The base URL to use. If provided will override the standard value.
|
||||
Can also be set via environment variable OPENAI_BASE_URL.
|
||||
default_headers: The default headers mapping of string keys to
|
||||
string values for HTTP requests.
|
||||
async_client: An existing client to use.
|
||||
env_file_path: Use the environment settings file as a fallback
|
||||
to environment variables.
|
||||
env_file_encoding: The encoding of the environment settings file.
|
||||
middleware: Optional sequence of middleware to apply to requests.
|
||||
function_invocation_configuration: Optional configuration for function invocation behavior.
|
||||
kwargs: Other keyword parameters.
|
||||
|
||||
Examples:
|
||||
.. code-block:: python
|
||||
|
||||
from agent_framework.openai import OpenAIAssistantsClient
|
||||
|
||||
# Using environment variables
|
||||
# Set OPENAI_API_KEY=sk-...
|
||||
# Set OPENAI_MODEL=gpt-4
|
||||
client = OpenAIAssistantsClient()
|
||||
|
||||
# Or passing parameters directly
|
||||
client = OpenAIAssistantsClient(model="gpt-4", api_key="sk-...")
|
||||
|
||||
# Or loading from a .env file
|
||||
client = OpenAIAssistantsClient(env_file_path="path/to/.env")
|
||||
|
||||
# Using custom ChatOptions with type safety:
|
||||
from typing import TypedDict
|
||||
from agent_framework.openai import OpenAIAssistantsOptions
|
||||
|
||||
|
||||
class MyOptions(OpenAIAssistantsOptions, total=False):
|
||||
my_custom_option: str
|
||||
|
||||
|
||||
client: OpenAIAssistantsClient[MyOptions] = OpenAIAssistantsClient(model="gpt-4")
|
||||
response = await client.get_response("Hello", options={"my_custom_option": "value"})
|
||||
"""
|
||||
if model_id is not None and model is None:
|
||||
import warnings
|
||||
|
||||
warnings.warn("model_id is deprecated, use model instead", DeprecationWarning, stacklevel=2)
|
||||
model = model_id
|
||||
openai_settings = load_settings(
|
||||
OpenAISettings,
|
||||
env_prefix="OPENAI_",
|
||||
api_key=api_key,
|
||||
base_url=base_url,
|
||||
org_id=org_id,
|
||||
model=model,
|
||||
env_file_path=env_file_path,
|
||||
env_file_encoding=env_file_encoding,
|
||||
)
|
||||
|
||||
api_key_value = openai_settings.get("api_key")
|
||||
if not async_client and not api_key_value:
|
||||
raise ValueError(
|
||||
"OpenAI API key is required. Set via 'api_key' parameter or 'OPENAI_API_KEY' environment variable."
|
||||
)
|
||||
|
||||
resolved_model = openai_settings.get("model")
|
||||
if not resolved_model:
|
||||
raise ValueError(
|
||||
"OpenAI model is required. Set via 'model' parameter or 'OPENAI_MODEL' environment variable."
|
||||
)
|
||||
|
||||
super().__init__(
|
||||
model=resolved_model,
|
||||
api_key=self._get_api_key(api_key_value),
|
||||
org_id=openai_settings.get("org_id"),
|
||||
default_headers=default_headers,
|
||||
client=async_client,
|
||||
base_url=openai_settings.get("base_url"),
|
||||
middleware=middleware,
|
||||
function_invocation_configuration=function_invocation_configuration,
|
||||
)
|
||||
self.assistant_id: str | None = assistant_id
|
||||
self.assistant_name: str | None = assistant_name
|
||||
self.assistant_description: str | None = assistant_description
|
||||
self.thread_id: str | None = thread_id
|
||||
self._should_delete_assistant: bool = False
|
||||
|
||||
async def __aenter__(self) -> Self:
|
||||
"""Async context manager entry."""
|
||||
return self
|
||||
|
||||
async def __aexit__(
|
||||
self,
|
||||
exc_type: type[BaseException] | None,
|
||||
exc_val: BaseException | None,
|
||||
exc_tb: Any,
|
||||
) -> None:
|
||||
"""Async context manager exit - clean up any assistants we created."""
|
||||
await self.close()
|
||||
|
||||
async def close(self) -> None:
|
||||
"""Clean up any assistants we created."""
|
||||
if self._should_delete_assistant and self.assistant_id is not None:
|
||||
client = await self._ensure_client()
|
||||
await client.beta.assistants.delete(self.assistant_id) # type: ignore[reportDeprecated]
|
||||
object.__setattr__(self, "assistant_id", None)
|
||||
object.__setattr__(self, "_should_delete_assistant", False)
|
||||
|
||||
@override
|
||||
def _inner_get_response(
|
||||
self,
|
||||
*,
|
||||
messages: Sequence[Message],
|
||||
options: Mapping[str, Any],
|
||||
stream: bool = False,
|
||||
**kwargs: Any,
|
||||
) -> Awaitable[ChatResponse] | ResponseStream[ChatResponseUpdate, ChatResponse]:
|
||||
if stream:
|
||||
# Streaming mode - return the async generator directly
|
||||
async def _stream() -> AsyncIterable[ChatResponseUpdate]:
|
||||
# prepare
|
||||
run_options, tool_results = self._prepare_options(messages, options, **kwargs)
|
||||
|
||||
# Get the thread ID
|
||||
thread_id: str | None = options.get(
|
||||
"conversation_id", run_options.get("conversation_id", self.thread_id)
|
||||
)
|
||||
|
||||
if thread_id is None and tool_results is not None:
|
||||
raise ValueError("No thread ID was provided, but chat messages includes tool results.")
|
||||
|
||||
# Determine which assistant to use and create if needed
|
||||
assistant_id = await self._get_assistant_id_or_create()
|
||||
|
||||
# execute
|
||||
stream_obj, thread_id = await self._create_assistant_stream(
|
||||
thread_id, assistant_id, run_options, tool_results
|
||||
)
|
||||
|
||||
# process
|
||||
async for update in self._process_stream_events(stream_obj, thread_id):
|
||||
yield update
|
||||
|
||||
return self._build_response_stream(_stream(), response_format=options.get("response_format"))
|
||||
|
||||
# Non-streaming mode - collect updates and convert to response
|
||||
async def _get_response() -> ChatResponse:
|
||||
stream_result = self._inner_get_response(messages=messages, options=options, stream=True, **kwargs)
|
||||
return await ChatResponse.from_update_generator(
|
||||
updates=stream_result, # type: ignore[arg-type]
|
||||
output_format_type=options.get("response_format"), # type: ignore[arg-type]
|
||||
)
|
||||
|
||||
return _get_response()
|
||||
|
||||
async def _get_assistant_id_or_create(self) -> str:
|
||||
"""Determine which assistant to use and create if needed.
|
||||
|
||||
Returns:
|
||||
str: The assistant_id to use.
|
||||
"""
|
||||
# If no assistant is provided, create a temporary assistant
|
||||
if self.assistant_id is None:
|
||||
if not self.model:
|
||||
raise ValueError("Parameter 'model' is required for assistant creation.")
|
||||
|
||||
client = await self._ensure_client()
|
||||
created_assistant = await client.beta.assistants.create( # type: ignore[reportDeprecated]
|
||||
model=self.model,
|
||||
description=self.assistant_description,
|
||||
name=self.assistant_name,
|
||||
)
|
||||
self.assistant_id = created_assistant.id
|
||||
self._should_delete_assistant = True
|
||||
|
||||
return self.assistant_id
|
||||
|
||||
async def _create_assistant_stream(
|
||||
self,
|
||||
thread_id: str | None,
|
||||
assistant_id: str,
|
||||
run_options: dict[str, Any],
|
||||
tool_results: list[Content] | None,
|
||||
) -> tuple[Any, str]:
|
||||
"""Create the assistant stream for processing.
|
||||
|
||||
Returns:
|
||||
tuple: (stream, final_thread_id)
|
||||
"""
|
||||
client = await self._ensure_client()
|
||||
# Get any active run for this thread
|
||||
thread_run = await self._get_active_thread_run(thread_id)
|
||||
|
||||
tool_run_id, tool_outputs = self._prepare_tool_outputs_for_assistants(tool_results)
|
||||
|
||||
if thread_run is not None and tool_run_id is not None and tool_run_id == thread_run.id and tool_outputs:
|
||||
# There's an active run and we have tool results to submit, so submit the results.
|
||||
stream = client.beta.threads.runs.submit_tool_outputs_stream( # type: ignore[reportDeprecated]
|
||||
run_id=tool_run_id,
|
||||
thread_id=thread_run.thread_id,
|
||||
tool_outputs=tool_outputs,
|
||||
)
|
||||
final_thread_id = thread_run.thread_id
|
||||
else:
|
||||
# Handle thread creation or cancellation
|
||||
final_thread_id = await self._prepare_thread(thread_id, thread_run, run_options)
|
||||
|
||||
# Now create a new run and stream the results.
|
||||
stream = client.beta.threads.runs.stream( # type: ignore[reportDeprecated]
|
||||
assistant_id=assistant_id, thread_id=final_thread_id, **run_options
|
||||
)
|
||||
|
||||
return stream, final_thread_id
|
||||
|
||||
async def _get_active_thread_run(self, thread_id: str | None) -> Run | None:
|
||||
"""Get any active run for the given thread."""
|
||||
client = await self._ensure_client()
|
||||
if thread_id is None:
|
||||
return None
|
||||
|
||||
async for run in client.beta.threads.runs.list(thread_id=thread_id, limit=1, order="desc"): # type: ignore[reportDeprecated]
|
||||
if run.status not in ["completed", "cancelled", "failed", "expired"]:
|
||||
return run
|
||||
return None
|
||||
|
||||
async def _prepare_thread(self, thread_id: str | None, thread_run: Run | None, run_options: dict[str, Any]) -> str:
|
||||
"""Prepare the thread for a new run, creating or cleaning up as needed."""
|
||||
client = await self._ensure_client()
|
||||
if thread_id is None:
|
||||
# No thread ID was provided, so create a new thread.
|
||||
thread = await client.beta.threads.create( # type: ignore[reportDeprecated]
|
||||
messages=run_options["additional_messages"],
|
||||
tool_resources=run_options.get("tool_resources"),
|
||||
metadata=run_options.get("metadata"),
|
||||
)
|
||||
run_options["additional_messages"] = []
|
||||
run_options.pop("tool_resources", None)
|
||||
return thread.id
|
||||
|
||||
if thread_run is not None:
|
||||
# There was an active run; we need to cancel it before starting a new run.
|
||||
await client.beta.threads.runs.cancel(run_id=thread_run.id, thread_id=thread_id) # type: ignore[reportDeprecated]
|
||||
|
||||
return thread_id
|
||||
|
||||
async def _process_stream_events(self, stream: Any, thread_id: str) -> AsyncIterable[ChatResponseUpdate]:
|
||||
response_id: str | None = None
|
||||
|
||||
async with stream as response_stream:
|
||||
async for response in response_stream:
|
||||
if response.event == "thread.run.created":
|
||||
yield ChatResponseUpdate(
|
||||
contents=[],
|
||||
conversation_id=thread_id,
|
||||
message_id=response_id,
|
||||
raw_representation=response.data,
|
||||
response_id=response_id,
|
||||
role="assistant",
|
||||
)
|
||||
elif response.event == "thread.run.step.created" and isinstance(response.data, RunStep):
|
||||
response_id = response.data.run_id
|
||||
elif response.event == "thread.message.delta" and isinstance(response.data, MessageDeltaEvent):
|
||||
delta = response.data.delta
|
||||
role = "user" if delta.role == "user" else "assistant"
|
||||
|
||||
for delta_block in delta.content or []:
|
||||
if isinstance(delta_block, TextDeltaBlock) and delta_block.text and delta_block.text.value:
|
||||
text_content = Content.from_text(delta_block.text.value)
|
||||
if delta_block.text.annotations:
|
||||
annotations: list[Annotation] = []
|
||||
text_content.annotations = annotations
|
||||
for annotation in delta_block.text.annotations:
|
||||
if isinstance(annotation, FileCitationDeltaAnnotation):
|
||||
ann: Annotation = Annotation(
|
||||
type="citation",
|
||||
additional_properties={
|
||||
"text": annotation.text,
|
||||
"index": annotation.index,
|
||||
},
|
||||
raw_representation=annotation,
|
||||
)
|
||||
if annotation.file_citation and annotation.file_citation.file_id:
|
||||
ann["file_id"] = annotation.file_citation.file_id
|
||||
if annotation.start_index is not None and annotation.end_index is not None:
|
||||
ann["annotated_regions"] = [
|
||||
TextSpanRegion(
|
||||
type="text_span",
|
||||
start_index=annotation.start_index,
|
||||
end_index=annotation.end_index,
|
||||
)
|
||||
]
|
||||
annotations.append(ann)
|
||||
elif isinstance(annotation, FilePathDeltaAnnotation):
|
||||
ann = Annotation(
|
||||
type="citation",
|
||||
additional_properties={
|
||||
"text": annotation.text,
|
||||
"index": annotation.index,
|
||||
},
|
||||
raw_representation=annotation,
|
||||
)
|
||||
if annotation.file_path and annotation.file_path.file_id:
|
||||
ann["file_id"] = annotation.file_path.file_id
|
||||
if annotation.start_index is not None and annotation.end_index is not None:
|
||||
ann["annotated_regions"] = [
|
||||
TextSpanRegion(
|
||||
type="text_span",
|
||||
start_index=annotation.start_index,
|
||||
end_index=annotation.end_index,
|
||||
)
|
||||
]
|
||||
annotations.append(ann)
|
||||
yield ChatResponseUpdate(
|
||||
role=role, # type: ignore[arg-type]
|
||||
contents=[text_content],
|
||||
conversation_id=thread_id,
|
||||
message_id=response_id,
|
||||
raw_representation=response.data,
|
||||
response_id=response_id,
|
||||
)
|
||||
elif response.event == "thread.message.completed" and isinstance(response.data, ThreadMessage):
|
||||
# Process completed message to extract fully resolved annotations.
|
||||
# Delta events may carry partial/empty annotation data; the completed
|
||||
# message contains the final text with all citation details populated.
|
||||
completed_contents: list[Content] = []
|
||||
for block in response.data.content:
|
||||
if block.type != "text":
|
||||
continue
|
||||
text_content = Content.from_text(block.text.value)
|
||||
if block.text.annotations:
|
||||
completed_annotations: list[Annotation] = []
|
||||
text_content.annotations = completed_annotations
|
||||
for completed_annotation in block.text.annotations:
|
||||
if isinstance(completed_annotation, FileCitationAnnotation):
|
||||
props: dict[str, Any] = {
|
||||
"text": completed_annotation.text,
|
||||
}
|
||||
ann = Annotation(
|
||||
type="citation",
|
||||
additional_properties=props,
|
||||
raw_representation=completed_annotation,
|
||||
)
|
||||
if (
|
||||
completed_annotation.file_citation
|
||||
and completed_annotation.file_citation.file_id
|
||||
):
|
||||
ann["file_id"] = completed_annotation.file_citation.file_id
|
||||
ann["annotated_regions"] = [
|
||||
TextSpanRegion(
|
||||
type="text_span",
|
||||
start_index=completed_annotation.start_index,
|
||||
end_index=completed_annotation.end_index,
|
||||
)
|
||||
]
|
||||
text_content.annotations.append(ann)
|
||||
elif isinstance(completed_annotation, FilePathAnnotation):
|
||||
ann = Annotation(
|
||||
type="citation",
|
||||
additional_properties={
|
||||
"text": completed_annotation.text,
|
||||
},
|
||||
raw_representation=completed_annotation,
|
||||
)
|
||||
if completed_annotation.file_path and completed_annotation.file_path.file_id:
|
||||
ann["file_id"] = completed_annotation.file_path.file_id
|
||||
ann["annotated_regions"] = [
|
||||
TextSpanRegion(
|
||||
type="text_span",
|
||||
start_index=completed_annotation.start_index,
|
||||
end_index=completed_annotation.end_index,
|
||||
)
|
||||
]
|
||||
text_content.annotations.append(ann)
|
||||
else:
|
||||
logger.debug("Unparsed annotation type: %s", completed_annotation.type)
|
||||
completed_contents.append(text_content)
|
||||
if completed_contents:
|
||||
yield ChatResponseUpdate(
|
||||
role="assistant",
|
||||
contents=completed_contents,
|
||||
conversation_id=thread_id,
|
||||
message_id=response_id,
|
||||
raw_representation=response.data,
|
||||
response_id=response_id,
|
||||
)
|
||||
elif response.event == "thread.run.requires_action" and isinstance(response.data, Run):
|
||||
contents = self._parse_function_calls_from_assistants(response.data, response_id)
|
||||
if contents:
|
||||
yield ChatResponseUpdate(
|
||||
role="assistant",
|
||||
contents=contents,
|
||||
conversation_id=thread_id,
|
||||
message_id=response_id,
|
||||
raw_representation=response.data,
|
||||
response_id=response_id,
|
||||
)
|
||||
elif (
|
||||
response.event == "thread.run.completed"
|
||||
and isinstance(response.data, Run)
|
||||
and response.data.usage is not None
|
||||
):
|
||||
usage = response.data.usage
|
||||
usage_content = Content.from_usage(
|
||||
UsageDetails(
|
||||
input_token_count=usage.prompt_tokens,
|
||||
output_token_count=usage.completion_tokens,
|
||||
total_token_count=usage.total_tokens,
|
||||
)
|
||||
)
|
||||
yield ChatResponseUpdate(
|
||||
role="assistant",
|
||||
contents=[usage_content],
|
||||
conversation_id=thread_id,
|
||||
message_id=response_id,
|
||||
raw_representation=response.data,
|
||||
response_id=response_id,
|
||||
)
|
||||
else:
|
||||
yield ChatResponseUpdate(
|
||||
contents=[],
|
||||
conversation_id=thread_id,
|
||||
message_id=response_id,
|
||||
raw_representation=response.data,
|
||||
response_id=response_id,
|
||||
role="assistant",
|
||||
)
|
||||
|
||||
def _parse_function_calls_from_assistants(self, event_data: Run, response_id: str | None) -> list[Content]:
|
||||
"""Parse function call contents from an assistants tool action event."""
|
||||
contents: list[Content] = []
|
||||
|
||||
if event_data.required_action is not None:
|
||||
for tool_call in event_data.required_action.submit_tool_outputs.tool_calls:
|
||||
tool_call_any = cast(Any, tool_call)
|
||||
call_id = json.dumps([response_id, tool_call.id])
|
||||
tool_type = getattr(tool_call, "type", None)
|
||||
if tool_type == "code_interpreter" and getattr(tool_call_any, "code_interpreter", None):
|
||||
code_input = getattr(tool_call_any.code_interpreter, "input", None)
|
||||
inputs = (
|
||||
[Content.from_text(text=code_input, raw_representation=tool_call)]
|
||||
if code_input is not None
|
||||
else None
|
||||
)
|
||||
contents.append(
|
||||
Content.from_code_interpreter_tool_call(
|
||||
call_id=call_id,
|
||||
inputs=inputs,
|
||||
raw_representation=tool_call,
|
||||
)
|
||||
)
|
||||
elif tool_type == "mcp":
|
||||
contents.append(
|
||||
Content.from_mcp_server_tool_call(
|
||||
call_id=call_id,
|
||||
tool_name=getattr(tool_call, "name", "") or "",
|
||||
server_name=getattr(tool_call, "server_label", None),
|
||||
arguments=getattr(tool_call, "args", None),
|
||||
raw_representation=tool_call,
|
||||
)
|
||||
)
|
||||
else:
|
||||
function_name = tool_call.function.name
|
||||
function_arguments = json.loads(tool_call.function.arguments)
|
||||
contents.append(
|
||||
Content.from_function_call(
|
||||
call_id=call_id,
|
||||
name=function_name,
|
||||
arguments=function_arguments,
|
||||
)
|
||||
)
|
||||
|
||||
return contents
|
||||
|
||||
def _prepare_options(
|
||||
self,
|
||||
messages: Sequence[Message],
|
||||
options: Mapping[str, Any],
|
||||
**kwargs: Any,
|
||||
) -> tuple[dict[str, Any], list[Content] | None]:
|
||||
from agent_framework._types import validate_tool_mode
|
||||
|
||||
run_options: dict[str, Any] = {**kwargs}
|
||||
|
||||
# Extract options from the dict
|
||||
max_tokens = options.get("max_tokens")
|
||||
model = options.get("model") or options.get("model_id") # backward compat
|
||||
top_p = options.get("top_p")
|
||||
temperature = options.get("temperature")
|
||||
allow_multiple_tool_calls = options.get("allow_multiple_tool_calls")
|
||||
tool_choice = options.get("tool_choice")
|
||||
tools = options.get("tools")
|
||||
response_format = options.get("response_format")
|
||||
tool_resources = options.get("tool_resources")
|
||||
|
||||
if max_tokens is not None:
|
||||
run_options["max_completion_tokens"] = max_tokens
|
||||
if model is not None:
|
||||
run_options["model"] = model
|
||||
if top_p is not None:
|
||||
run_options["top_p"] = top_p
|
||||
if temperature is not None:
|
||||
run_options["temperature"] = temperature
|
||||
|
||||
if allow_multiple_tool_calls is not None:
|
||||
run_options["parallel_tool_calls"] = allow_multiple_tool_calls
|
||||
|
||||
if tool_resources is not None:
|
||||
run_options["tool_resources"] = tool_resources
|
||||
|
||||
tool_mode = validate_tool_mode(tool_choice)
|
||||
tool_definitions: list[MutableMapping[str, Any]] = []
|
||||
# Always include tools if provided, regardless of tool_choice
|
||||
# tool_choice="none" means the model won't call tools, but tools should still be available
|
||||
for tool in normalize_tools(tools):
|
||||
if isinstance(tool, FunctionTool):
|
||||
tool_definitions.append(tool.to_json_schema_spec()) # type: ignore[reportUnknownArgumentType]
|
||||
elif isinstance(tool, MutableMapping):
|
||||
# Pass through dict-based tools directly (from static factory methods)
|
||||
tool_definitions.append(cast(MutableMapping[str, Any], tool))
|
||||
|
||||
if len(tool_definitions) > 0:
|
||||
run_options["tools"] = tool_definitions
|
||||
|
||||
if tool_mode is not None:
|
||||
mode = tool_mode.get("mode")
|
||||
if mode is None:
|
||||
raise ValueError("tool_choice mode is required")
|
||||
if mode == "required" and (func_name := tool_mode.get("required_function_name")) is not None:
|
||||
run_options["tool_choice"] = {
|
||||
"type": "function",
|
||||
"function": {"name": func_name},
|
||||
}
|
||||
else:
|
||||
run_options["tool_choice"] = mode
|
||||
|
||||
if response_format is not None:
|
||||
if isinstance(response_format, dict):
|
||||
run_options["response_format"] = response_format
|
||||
else:
|
||||
run_options["response_format"] = {
|
||||
"type": "json_schema",
|
||||
"json_schema": {
|
||||
"name": response_format.__name__,
|
||||
"schema": response_format.model_json_schema(),
|
||||
"strict": True,
|
||||
},
|
||||
}
|
||||
|
||||
instructions: list[str] = []
|
||||
tool_results: list[Content] | None = None
|
||||
|
||||
additional_messages: list[AdditionalMessage] | None = None
|
||||
|
||||
# System/developer messages are turned into instructions,
|
||||
# since there is no such message roles in OpenAI Assistants.
|
||||
# All other messages are added 1:1.
|
||||
for chat_message in messages:
|
||||
if chat_message.role in ["system", "developer"]:
|
||||
for text_content in [content for content in chat_message.contents if content.type == "text"]:
|
||||
text = getattr(text_content, "text", None)
|
||||
if text:
|
||||
instructions.append(text)
|
||||
|
||||
continue
|
||||
|
||||
message_contents: list[MessageContentPartParam] = []
|
||||
|
||||
for content in chat_message.contents:
|
||||
if content.type == "text":
|
||||
message_contents.append(TextContentBlockParam(type="text", text=content.text)) # type: ignore[attr-defined, typeddict-item]
|
||||
elif content.type == "uri" and content.has_top_level_media_type("image"):
|
||||
message_contents.append(
|
||||
ImageURLContentBlockParam(type="image_url", image_url=ImageURLParam(url=content.uri)) # type: ignore[attr-defined, typeddict-item]
|
||||
)
|
||||
elif content.type == "function_result":
|
||||
if tool_results is None:
|
||||
tool_results = []
|
||||
tool_results.append(content)
|
||||
|
||||
if len(message_contents) > 0:
|
||||
if additional_messages is None:
|
||||
additional_messages = []
|
||||
additional_messages.append(
|
||||
AdditionalMessage(
|
||||
role="assistant" if chat_message.role == "assistant" else "user",
|
||||
content=message_contents,
|
||||
)
|
||||
)
|
||||
|
||||
if additional_messages is not None:
|
||||
run_options["additional_messages"] = additional_messages
|
||||
|
||||
if len(instructions) > 0:
|
||||
run_options["instructions"] = "".join(instructions)
|
||||
|
||||
return run_options, tool_results
|
||||
|
||||
def _prepare_tool_outputs_for_assistants(
|
||||
self,
|
||||
tool_results: list[Content] | None,
|
||||
) -> tuple[str | None, list[ToolOutput] | None]:
|
||||
"""Prepare function results for submission to the assistants API."""
|
||||
run_id: str | None = None
|
||||
tool_outputs: list[ToolOutput] | None = None
|
||||
|
||||
if tool_results:
|
||||
for function_result_content in tool_results:
|
||||
# When creating the FunctionCallContent, we created it with a CallId == [runId, callId].
|
||||
# We need to extract the run ID and ensure that the ToolOutput we send back to Azure
|
||||
# is only the call ID.
|
||||
run_and_call_ids: list[str] = json.loads(function_result_content.call_id) # type: ignore[arg-type]
|
||||
|
||||
if (
|
||||
not run_and_call_ids
|
||||
or len(run_and_call_ids) != 2
|
||||
or not run_and_call_ids[0]
|
||||
or not run_and_call_ids[1]
|
||||
or (run_id is not None and run_id != run_and_call_ids[0])
|
||||
):
|
||||
continue
|
||||
|
||||
run_id = run_and_call_ids[0]
|
||||
call_id = run_and_call_ids[1]
|
||||
|
||||
if tool_outputs is None:
|
||||
tool_outputs = []
|
||||
output = (
|
||||
function_result_content.result
|
||||
if function_result_content.result is not None
|
||||
else "No output received."
|
||||
)
|
||||
tool_outputs.append(ToolOutput(tool_call_id=call_id, output=output))
|
||||
|
||||
return run_id, tool_outputs
|
||||
|
||||
def _update_agent_name_and_description(self, agent_name: str | None, description: str | None = None) -> None:
|
||||
"""Update the agent name in the chat client.
|
||||
|
||||
Args:
|
||||
agent_name: The new name for the agent.
|
||||
description: The new description for the agent.
|
||||
"""
|
||||
# This is a no-op in the base class, but can be overridden by subclasses
|
||||
# to update the agent name in the client.
|
||||
if agent_name and not self.assistant_name:
|
||||
self.assistant_name = agent_name
|
||||
if description and not self.assistant_description:
|
||||
self.assistant_description = description
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,303 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import struct
|
||||
import sys
|
||||
from collections.abc import Awaitable, Callable, Mapping, Sequence
|
||||
from copy import copy
|
||||
from typing import Any, ClassVar, Generic, Literal, TypedDict
|
||||
|
||||
from agent_framework._clients import BaseEmbeddingClient
|
||||
from agent_framework._settings import SecretString, load_settings
|
||||
from agent_framework._telemetry import APP_INFO, USER_AGENT_KEY, prepend_agent_framework_to_user_agent
|
||||
from agent_framework._types import Embedding, EmbeddingGenerationOptions, GeneratedEmbeddings, UsageDetails
|
||||
from agent_framework.observability import EmbeddingTelemetryLayer
|
||||
from openai import AsyncOpenAI
|
||||
|
||||
from ._shared import OpenAISettings, get_api_key
|
||||
|
||||
if sys.version_info >= (3, 13):
|
||||
from typing import TypeVar # type: ignore # pragma: no cover
|
||||
else:
|
||||
from typing_extensions import TypeVar # type: ignore # pragma: no cover
|
||||
|
||||
|
||||
class OpenAIEmbeddingOptions(EmbeddingGenerationOptions, total=False):
|
||||
"""OpenAI-specific embedding options.
|
||||
|
||||
Extends EmbeddingGenerationOptions with OpenAI-specific fields.
|
||||
|
||||
Examples:
|
||||
.. code-block:: python
|
||||
|
||||
from agent_framework.openai import OpenAIEmbeddingOptions
|
||||
|
||||
options: OpenAIEmbeddingOptions = {
|
||||
"model": "text-embedding-3-small",
|
||||
"dimensions": 1536,
|
||||
"encoding_format": "float",
|
||||
}
|
||||
"""
|
||||
|
||||
encoding_format: Literal["float", "base64"]
|
||||
user: str
|
||||
|
||||
|
||||
OpenAIEmbeddingOptionsT = TypeVar(
|
||||
"OpenAIEmbeddingOptionsT",
|
||||
bound=TypedDict, # type: ignore[valid-type]
|
||||
default="OpenAIEmbeddingOptions",
|
||||
covariant=True,
|
||||
)
|
||||
|
||||
|
||||
class RawOpenAIEmbeddingClient(
|
||||
BaseEmbeddingClient[str, list[float], OpenAIEmbeddingOptionsT],
|
||||
Generic[OpenAIEmbeddingOptionsT],
|
||||
):
|
||||
"""Raw OpenAI embedding client without telemetry."""
|
||||
|
||||
INJECTABLE: ClassVar[set[str]] = {"client"}
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
model: str | None = None,
|
||||
model_id: str | None = None,
|
||||
api_key: str | SecretString | Callable[[], str | Awaitable[str]] | None = None,
|
||||
org_id: str | None = None,
|
||||
base_url: str | None = None,
|
||||
default_headers: Mapping[str, str] | None = None,
|
||||
async_client: AsyncOpenAI | None = None,
|
||||
env_file_path: str | None = None,
|
||||
env_file_encoding: str | None = None,
|
||||
**kwargs: Any,
|
||||
) -> None:
|
||||
"""Initialize a raw OpenAI embedding client.
|
||||
|
||||
Keyword Args:
|
||||
model: OpenAI embedding model name.
|
||||
model_id: Deprecated alias for ``model``.
|
||||
api_key: OpenAI API key, SecretString, or callable returning a key.
|
||||
org_id: OpenAI organization ID.
|
||||
base_url: Custom API base URL.
|
||||
default_headers: Additional HTTP headers.
|
||||
async_client: Pre-configured AsyncOpenAI client (skips client creation).
|
||||
env_file_path: Path to .env file for settings.
|
||||
env_file_encoding: Encoding for .env file.
|
||||
kwargs: Additional keyword arguments forwarded to ``BaseEmbeddingClient``.
|
||||
"""
|
||||
if model_id is not None and model is None:
|
||||
import warnings
|
||||
|
||||
warnings.warn("model_id is deprecated, use model instead", DeprecationWarning, stacklevel=2)
|
||||
model = model_id
|
||||
|
||||
if not async_client:
|
||||
openai_settings = load_settings(
|
||||
OpenAISettings,
|
||||
env_prefix="OPENAI_",
|
||||
api_key=api_key,
|
||||
org_id=org_id,
|
||||
base_url=base_url,
|
||||
embedding_model=model,
|
||||
env_file_path=env_file_path,
|
||||
env_file_encoding=env_file_encoding,
|
||||
)
|
||||
|
||||
api_key_value = openai_settings.get("api_key")
|
||||
resolved_model = openai_settings.get("embedding_model") or model
|
||||
|
||||
# Only create a client when we have enough configuration.
|
||||
# Subclasses that manage their own client pass no args here
|
||||
if api_key_value:
|
||||
if not resolved_model:
|
||||
raise ValueError(
|
||||
"OpenAI embedding model is required. "
|
||||
"Set via 'model' parameter or 'OPENAI_EMBEDDING_MODEL' environment variable."
|
||||
)
|
||||
model = resolved_model
|
||||
|
||||
resolved_api_key = get_api_key(api_key_value)
|
||||
|
||||
# Merge APP_INFO into the headers
|
||||
merged_headers = dict(copy(default_headers)) if default_headers else {}
|
||||
if APP_INFO:
|
||||
merged_headers.update(APP_INFO)
|
||||
merged_headers = prepend_agent_framework_to_user_agent(merged_headers)
|
||||
|
||||
client_args: dict[str, Any] = {"api_key": resolved_api_key, "default_headers": merged_headers}
|
||||
if resolved_org_id := openai_settings.get("org_id"):
|
||||
client_args["organization"] = resolved_org_id
|
||||
if resolved_base_url := openai_settings.get("base_url"):
|
||||
client_args["base_url"] = resolved_base_url
|
||||
|
||||
async_client = AsyncOpenAI(**client_args)
|
||||
|
||||
self.client = async_client
|
||||
self.model: str | None = model.strip() if model else None
|
||||
|
||||
# Store configuration for serialization
|
||||
self.org_id = org_id
|
||||
self.base_url = str(base_url) if base_url else None
|
||||
if default_headers:
|
||||
self.default_headers: dict[str, Any] | None = {
|
||||
k: v for k, v in default_headers.items() if k != USER_AGENT_KEY
|
||||
}
|
||||
else:
|
||||
self.default_headers = None
|
||||
|
||||
super().__init__(**kwargs)
|
||||
|
||||
def service_url(self) -> str:
|
||||
"""Get the URL of the service."""
|
||||
return str(self.client.base_url) if self.client else "Unknown"
|
||||
|
||||
async def get_embeddings(
|
||||
self,
|
||||
values: Sequence[str],
|
||||
*,
|
||||
options: OpenAIEmbeddingOptionsT | None = None,
|
||||
) -> GeneratedEmbeddings[list[float], OpenAIEmbeddingOptionsT]:
|
||||
"""Call the OpenAI embeddings API.
|
||||
|
||||
Args:
|
||||
values: The text values to generate embeddings for.
|
||||
options: Optional embedding generation options.
|
||||
|
||||
Returns:
|
||||
Generated embeddings with usage metadata.
|
||||
|
||||
Raises:
|
||||
ValueError: If model is not provided or values is empty.
|
||||
"""
|
||||
if not values:
|
||||
return GeneratedEmbeddings([], options=options) # type: ignore
|
||||
|
||||
opts: dict[str, Any] = options or {} # type: ignore
|
||||
# backward compat: accept model_id in options
|
||||
model = opts.get("model") or opts.get("model_id") or self.model
|
||||
if not model:
|
||||
raise ValueError("model is required")
|
||||
|
||||
kwargs: dict[str, Any] = {"input": list(values), "model": model}
|
||||
if dimensions := opts.get("dimensions"):
|
||||
kwargs["dimensions"] = dimensions
|
||||
if encoding_format := opts.get("encoding_format"):
|
||||
kwargs["encoding_format"] = encoding_format
|
||||
if user := opts.get("user"):
|
||||
kwargs["user"] = user
|
||||
|
||||
response = await self.client.embeddings.create(**kwargs) # type: ignore[union-attr]
|
||||
|
||||
encoding = kwargs.get("encoding_format", "float")
|
||||
embeddings: list[Embedding[list[float]]] = []
|
||||
for item in response.data:
|
||||
vector: list[float]
|
||||
if encoding == "base64" and isinstance(item.embedding, str):
|
||||
# Decode base64-encoded floats (little-endian IEEE 754)
|
||||
raw = base64.b64decode(item.embedding)
|
||||
vector = list(struct.unpack(f"<{len(raw) // 4}f", raw))
|
||||
else:
|
||||
vector = item.embedding # type: ignore[assignment]
|
||||
embeddings.append(
|
||||
Embedding(
|
||||
vector=vector,
|
||||
dimensions=len(vector),
|
||||
model=response.model,
|
||||
)
|
||||
)
|
||||
|
||||
usage_dict: UsageDetails | None = None
|
||||
if response.usage:
|
||||
usage_dict = {
|
||||
"input_token_count": response.usage.prompt_tokens,
|
||||
"total_token_count": response.usage.total_tokens,
|
||||
}
|
||||
|
||||
return GeneratedEmbeddings(embeddings, options=options, usage=usage_dict)
|
||||
|
||||
|
||||
class OpenAIEmbeddingClient(
|
||||
EmbeddingTelemetryLayer[str, list[float], OpenAIEmbeddingOptionsT],
|
||||
RawOpenAIEmbeddingClient[OpenAIEmbeddingOptionsT],
|
||||
Generic[OpenAIEmbeddingOptionsT],
|
||||
):
|
||||
"""OpenAI embedding client with telemetry support.
|
||||
|
||||
Keyword Args:
|
||||
model: The embedding model (e.g. "text-embedding-3-small").
|
||||
Can also be set via environment variable OPENAI_EMBEDDING_MODEL.
|
||||
model_id: Deprecated alias for ``model``.
|
||||
api_key: OpenAI API key.
|
||||
Can also be set via environment variable OPENAI_API_KEY.
|
||||
org_id: OpenAI organization ID.
|
||||
default_headers: Additional HTTP headers.
|
||||
async_client: Pre-configured AsyncOpenAI client.
|
||||
base_url: Custom API base URL.
|
||||
otel_provider_name: Override the OpenTelemetry provider name for telemetry.
|
||||
env_file_path: Path to .env file for settings.
|
||||
env_file_encoding: Encoding for .env file.
|
||||
|
||||
Examples:
|
||||
.. code-block:: python
|
||||
|
||||
from agent_framework.openai import OpenAIEmbeddingClient
|
||||
|
||||
# Using environment variables
|
||||
# Set OPENAI_API_KEY=sk-...
|
||||
# Set OPENAI_EMBEDDING_MODEL=text-embedding-3-small
|
||||
client = OpenAIEmbeddingClient()
|
||||
|
||||
# Or passing parameters directly
|
||||
client = OpenAIEmbeddingClient(
|
||||
model="text-embedding-3-small",
|
||||
api_key="sk-...",
|
||||
)
|
||||
|
||||
# Generate embeddings
|
||||
result = await client.get_embeddings(["Hello, world!"])
|
||||
print(result[0].vector)
|
||||
"""
|
||||
|
||||
OTEL_PROVIDER_NAME: ClassVar[str] = "openai" # type: ignore[reportIncompatibleVariableOverride, misc]
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
model: str | None = None,
|
||||
api_key: str | Callable[[], str | Awaitable[str]] | None = None,
|
||||
org_id: str | None = None,
|
||||
default_headers: Mapping[str, str] | None = None,
|
||||
async_client: AsyncOpenAI | None = None,
|
||||
base_url: str | None = None,
|
||||
otel_provider_name: str | None = None,
|
||||
env_file_path: str | None = None,
|
||||
env_file_encoding: str | None = None,
|
||||
) -> None:
|
||||
"""Initialize an OpenAI embedding client."""
|
||||
super().__init__(
|
||||
model=model,
|
||||
api_key=api_key,
|
||||
org_id=org_id,
|
||||
base_url=base_url,
|
||||
default_headers=default_headers,
|
||||
async_client=async_client,
|
||||
env_file_path=env_file_path,
|
||||
env_file_encoding=env_file_encoding,
|
||||
)
|
||||
if otel_provider_name is not None:
|
||||
self.OTEL_PROVIDER_NAME = otel_provider_name # type: ignore[misc]
|
||||
|
||||
# Validate that the client was created successfully (from explicit args or env vars)
|
||||
if self.client is None:
|
||||
raise ValueError(
|
||||
"OpenAI API key is required. Set via 'api_key' parameter or 'OPENAI_API_KEY' environment variable."
|
||||
)
|
||||
if not self.model:
|
||||
raise ValueError(
|
||||
"OpenAI embedding model is required. "
|
||||
"Set via 'model' parameter or 'OPENAI_EMBEDDING_MODEL' environment variable."
|
||||
)
|
||||
@@ -0,0 +1,90 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from enum import Enum
|
||||
from typing import Any
|
||||
|
||||
from agent_framework.exceptions import ChatClientContentFilterException
|
||||
from openai import BadRequestError
|
||||
|
||||
|
||||
class ContentFilterResultSeverity(Enum):
|
||||
"""The severity of the content filter result."""
|
||||
|
||||
HIGH = "high"
|
||||
MEDIUM = "medium"
|
||||
SAFE = "safe"
|
||||
LOW = "low"
|
||||
|
||||
|
||||
@dataclass
|
||||
class ContentFilterResult:
|
||||
"""The result of a content filter check."""
|
||||
|
||||
filtered: bool = False
|
||||
detected: bool = False
|
||||
severity: ContentFilterResultSeverity = ContentFilterResultSeverity.SAFE
|
||||
|
||||
@classmethod
|
||||
def from_inner_error_result(cls, inner_error_results: dict[str, Any]) -> ContentFilterResult:
|
||||
"""Creates a ContentFilterResult from the inner error results.
|
||||
|
||||
Args:
|
||||
inner_error_results: The inner error results.
|
||||
|
||||
Returns:
|
||||
ContentFilterResult: The ContentFilterResult.
|
||||
"""
|
||||
return cls(
|
||||
filtered=inner_error_results.get("filtered", False),
|
||||
detected=inner_error_results.get("detected", False),
|
||||
severity=ContentFilterResultSeverity(
|
||||
inner_error_results.get("severity", ContentFilterResultSeverity.SAFE.value)
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
class ContentFilterCodes(Enum):
|
||||
"""Content filter codes."""
|
||||
|
||||
RESPONSIBLE_AI_POLICY_VIOLATION = "ResponsibleAIPolicyViolation"
|
||||
|
||||
|
||||
@dataclass
|
||||
class OpenAIContentFilterException(ChatClientContentFilterException):
|
||||
"""AI exception for an error from Azure OpenAI's content filter."""
|
||||
|
||||
# The parameter that caused the error.
|
||||
param: str | None
|
||||
|
||||
# The error code specific to the content filter.
|
||||
content_filter_code: ContentFilterCodes
|
||||
|
||||
# The results of the different content filter checks.
|
||||
content_filter_result: dict[str, ContentFilterResult]
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
message: str,
|
||||
inner_exception: BadRequestError,
|
||||
) -> None:
|
||||
"""Initializes a new instance of the ContentFilterAIException class.
|
||||
|
||||
Args:
|
||||
message: The error message.
|
||||
inner_exception: The inner exception.
|
||||
"""
|
||||
super().__init__(message)
|
||||
|
||||
self.param = inner_exception.param
|
||||
if inner_exception.body is not None and isinstance(inner_exception.body, dict):
|
||||
inner_error = inner_exception.body.get("innererror", {}) # type: ignore
|
||||
self.content_filter_code = ContentFilterCodes(
|
||||
inner_error.get("code", ContentFilterCodes.RESPONSIBLE_AI_POLICY_VIOLATION.value) # type: ignore
|
||||
)
|
||||
self.content_filter_result = {
|
||||
key: ContentFilterResult.from_inner_error_result(values) # type: ignore
|
||||
for key, values in inner_error.get("content_filter_result", {}).items() # type: ignore
|
||||
}
|
||||
@@ -0,0 +1,508 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import os
|
||||
import sys
|
||||
from collections.abc import Awaitable, Callable, Mapping, MutableMapping, Sequence
|
||||
from copy import copy
|
||||
from typing import Any, ClassVar, Union, cast
|
||||
|
||||
import openai
|
||||
from agent_framework._serialization import SerializationMixin
|
||||
from agent_framework._settings import SecretString, load_settings
|
||||
from agent_framework._telemetry import APP_INFO, USER_AGENT_KEY, prepend_agent_framework_to_user_agent
|
||||
from agent_framework._tools import FunctionTool
|
||||
from dotenv import dotenv_values
|
||||
from openai import AsyncOpenAI, AsyncStream, _legacy_response # type: ignore
|
||||
from openai.types import Completion
|
||||
from openai.types.audio import Transcription
|
||||
from openai.types.chat import ChatCompletion, ChatCompletionChunk
|
||||
from openai.types.images_response import ImagesResponse
|
||||
from openai.types.responses.response import Response
|
||||
from openai.types.responses.response_stream_event import ResponseStreamEvent
|
||||
from packaging.version import parse
|
||||
|
||||
logger: logging.Logger = logging.getLogger("agent_framework.openai")
|
||||
|
||||
DEFAULT_AZURE_OPENAI_CHAT_COMPLETION_API_VERSION = "2024-10-21"
|
||||
DEFAULT_AZURE_OPENAI_RESPONSES_API_VERSION = "preview"
|
||||
|
||||
|
||||
RESPONSE_TYPE = Union[
|
||||
ChatCompletion,
|
||||
Completion,
|
||||
AsyncStream[ChatCompletionChunk],
|
||||
AsyncStream[Completion],
|
||||
list[Any],
|
||||
ImagesResponse,
|
||||
Response,
|
||||
AsyncStream[ResponseStreamEvent],
|
||||
Transcription,
|
||||
_legacy_response.HttpxBinaryResponseContent,
|
||||
]
|
||||
|
||||
OPTION_TYPE = dict[str, Any]
|
||||
|
||||
if sys.version_info >= (3, 11):
|
||||
from typing import TypedDict # type: ignore # pragma: no cover
|
||||
else:
|
||||
from typing_extensions import TypedDict # type: ignore # pragma: no cover
|
||||
|
||||
|
||||
def _check_openai_version_for_callable_api_key() -> None:
|
||||
"""Check if OpenAI version supports callable API keys.
|
||||
|
||||
Callable API keys require OpenAI >= 1.106.0.
|
||||
If the version is too old, raise a ValueError with helpful message.
|
||||
"""
|
||||
try:
|
||||
current_version = parse(openai.__version__)
|
||||
min_required_version = parse("1.106.0")
|
||||
|
||||
if current_version < min_required_version:
|
||||
raise ValueError(
|
||||
f"Callable API keys require OpenAI SDK >= 1.106.0, but you have {openai.__version__}. "
|
||||
f"Please upgrade with 'pip install openai>=1.106.0' or provide a string API key instead. "
|
||||
f"Note: If you're using mem0ai, you may need to upgrade to mem0ai>=1.0.0 "
|
||||
f"to allow newer OpenAI versions."
|
||||
)
|
||||
except ValueError:
|
||||
raise # Re-raise our own exception
|
||||
except Exception as e:
|
||||
logger.warning(f"Could not check OpenAI version for callable API key support: {e}")
|
||||
|
||||
|
||||
class OpenAISettings(TypedDict, total=False):
|
||||
"""OpenAI environment settings.
|
||||
|
||||
Settings are resolved in this order: explicit keyword arguments, values from an
|
||||
explicitly provided .env file, then environment variables with the prefix
|
||||
'OPENAI_'. If settings are missing after resolution, validation will fail.
|
||||
|
||||
Keyword Args:
|
||||
api_key: OpenAI API key, see https://platform.openai.com/account/api-keys.
|
||||
Can be set via environment variable OPENAI_API_KEY.
|
||||
base_url: The base URL for the OpenAI API.
|
||||
Can be set via environment variable OPENAI_BASE_URL.
|
||||
org_id: This is usually optional unless your account belongs to multiple organizations.
|
||||
Can be set via environment variable OPENAI_ORG_ID.
|
||||
model: The OpenAI model to use, for example, gpt-4o or o1.
|
||||
Can be set via environment variable OPENAI_MODEL.
|
||||
embedding_model: The OpenAI embedding model to use, for example, text-embedding-3-small.
|
||||
Can be set via environment variable OPENAI_EMBEDDING_MODEL.
|
||||
|
||||
Examples:
|
||||
.. code-block:: python
|
||||
|
||||
from agent_framework.openai import OpenAISettings
|
||||
|
||||
# Using environment variables
|
||||
# Set OPENAI_API_KEY=sk-...
|
||||
# Set OPENAI_MODEL=gpt-4o
|
||||
settings = load_settings(OpenAISettings, env_prefix="OPENAI_")
|
||||
|
||||
# Or passing parameters directly
|
||||
settings = load_settings(OpenAISettings, env_prefix="OPENAI_", api_key="sk-...", model="gpt-4o")
|
||||
|
||||
# Or loading from a .env file
|
||||
settings = load_settings(OpenAISettings, env_prefix="OPENAI_", env_file_path="path/to/.env")
|
||||
"""
|
||||
|
||||
api_key: SecretString | Callable[[], str | Awaitable[str]] | None
|
||||
base_url: str | None
|
||||
org_id: str | None
|
||||
model: str | None
|
||||
embedding_model: str | None
|
||||
azure_endpoint: str | None
|
||||
api_version: str | None
|
||||
|
||||
|
||||
def _load_dotenv_values(*, env_file_path: str | None, env_file_encoding: str | None) -> dict[str, str]:
|
||||
"""Load dotenv values for non-standard environment variable aliases."""
|
||||
if env_file_path is None or not os.path.exists(env_file_path):
|
||||
return {}
|
||||
|
||||
raw_dotenv_values = dotenv_values(dotenv_path=env_file_path, encoding=env_file_encoding or "utf-8")
|
||||
return {key: value for key, value in raw_dotenv_values.items() if value is not None}
|
||||
|
||||
|
||||
def _get_setting_from_alias(
|
||||
name: str,
|
||||
*,
|
||||
dotenv_values_by_name: Mapping[str, str],
|
||||
) -> str | None:
|
||||
"""Resolve a setting from an explicit env-var alias."""
|
||||
if dotenv_value := dotenv_values_by_name.get(name):
|
||||
return dotenv_value
|
||||
return os.getenv(name)
|
||||
|
||||
|
||||
def load_openai_service_settings(
|
||||
*,
|
||||
model: str | None,
|
||||
api_key: str | SecretString | Callable[[], str | Awaitable[str]] | None,
|
||||
org_id: str | None,
|
||||
base_url: str | None,
|
||||
azure_endpoint: str | None,
|
||||
api_version: str | None,
|
||||
env_file_path: str | None,
|
||||
env_file_encoding: str | None,
|
||||
azure_model_env_vars: Sequence[str],
|
||||
default_azure_api_version: str,
|
||||
) -> tuple[OpenAISettings, bool]:
|
||||
"""Load OpenAI settings, including Azure OpenAI aliases.
|
||||
|
||||
The generic OpenAI clients primarily read from ``OPENAI_*`` variables. When an
|
||||
``AZURE_OPENAI_ENDPOINT`` (or ``AZURE_OPENAI_BASE_URL``) is available and no
|
||||
explicit OpenAI base URL is configured, this helper switches to Azure-specific
|
||||
environment variables for endpoint, API key, model deployment, and API version.
|
||||
"""
|
||||
openai_settings = load_settings(
|
||||
OpenAISettings,
|
||||
env_prefix="OPENAI_",
|
||||
api_key=api_key,
|
||||
org_id=org_id,
|
||||
base_url=base_url,
|
||||
model=model,
|
||||
azure_endpoint=azure_endpoint,
|
||||
api_version=api_version,
|
||||
env_file_path=env_file_path,
|
||||
env_file_encoding=env_file_encoding,
|
||||
)
|
||||
|
||||
dotenv_values_by_name = _load_dotenv_values(
|
||||
env_file_path=env_file_path,
|
||||
env_file_encoding=env_file_encoding,
|
||||
)
|
||||
|
||||
resolved_azure_endpoint = azure_endpoint
|
||||
resolved_azure_base_url: str | None = None
|
||||
if not openai_settings.get("base_url"):
|
||||
if resolved_azure_endpoint is None:
|
||||
resolved_azure_endpoint = _get_setting_from_alias(
|
||||
"AZURE_OPENAI_ENDPOINT",
|
||||
dotenv_values_by_name=dotenv_values_by_name,
|
||||
)
|
||||
if resolved_azure_endpoint is None:
|
||||
resolved_azure_base_url = _get_setting_from_alias(
|
||||
"AZURE_OPENAI_BASE_URL",
|
||||
dotenv_values_by_name=dotenv_values_by_name,
|
||||
)
|
||||
if resolved_azure_base_url is not None:
|
||||
openai_settings["base_url"] = resolved_azure_base_url
|
||||
|
||||
use_azure_client = resolved_azure_endpoint is not None or resolved_azure_base_url is not None
|
||||
if resolved_azure_endpoint is not None:
|
||||
openai_settings["azure_endpoint"] = resolved_azure_endpoint
|
||||
|
||||
if use_azure_client:
|
||||
if api_key is None:
|
||||
resolved_azure_api_key = _get_setting_from_alias(
|
||||
"AZURE_OPENAI_API_KEY",
|
||||
dotenv_values_by_name=dotenv_values_by_name,
|
||||
)
|
||||
if resolved_azure_api_key is not None:
|
||||
openai_settings["api_key"] = SecretString(resolved_azure_api_key)
|
||||
|
||||
if model is None:
|
||||
for env_var_name in azure_model_env_vars:
|
||||
resolved_model = _get_setting_from_alias(
|
||||
env_var_name,
|
||||
dotenv_values_by_name=dotenv_values_by_name,
|
||||
)
|
||||
if resolved_model is not None:
|
||||
openai_settings["model"] = resolved_model
|
||||
break
|
||||
|
||||
if not openai_settings.get("api_version"):
|
||||
resolved_api_version = _get_setting_from_alias(
|
||||
"AZURE_OPENAI_API_VERSION",
|
||||
dotenv_values_by_name=dotenv_values_by_name,
|
||||
)
|
||||
openai_settings["api_version"] = resolved_api_version or default_azure_api_version
|
||||
|
||||
return openai_settings, use_azure_client
|
||||
|
||||
|
||||
def maybe_append_azure_endpoint_guidance(message: str, *, azure_endpoint: str | None) -> str:
|
||||
"""Append Azure endpoint guidance only when the configured endpoint shape looks suspicious."""
|
||||
if not azure_endpoint or not azure_endpoint.rstrip("/").endswith("/openai/v1"):
|
||||
return message
|
||||
|
||||
return (
|
||||
f"{message} If you are using Azure OpenAI key auth, pass the resource endpoint without "
|
||||
"'/openai/v1' to 'azure_endpoint', or pass the full '/openai/v1' URL via 'base_url' instead."
|
||||
)
|
||||
|
||||
|
||||
def get_api_key(
|
||||
api_key: str | SecretString | Callable[[], str | Awaitable[str]] | None,
|
||||
) -> str | Callable[[], str | Awaitable[str]] | None:
|
||||
"""Get the appropriate API key value for client initialization.
|
||||
|
||||
Args:
|
||||
api_key: The API key parameter which can be a string, SecretString, callable, or None.
|
||||
|
||||
Returns:
|
||||
For callable API keys: returns the callable directly.
|
||||
For SecretString: returns the unwrapped secret value.
|
||||
For string/None API keys: returns as-is.
|
||||
"""
|
||||
if isinstance(api_key, SecretString):
|
||||
return api_key.get_secret_value()
|
||||
|
||||
# Check version compatibility for callable API keys
|
||||
if callable(api_key):
|
||||
_check_openai_version_for_callable_api_key()
|
||||
|
||||
return api_key # Pass callable, string, or None directly to OpenAI SDK
|
||||
|
||||
|
||||
class OpenAIBase(SerializationMixin):
|
||||
"""Base class for OpenAI Clients.
|
||||
|
||||
.. deprecated::
|
||||
``OpenAIBase`` is deprecated and only used by ``OpenAIAssistantsClient``
|
||||
and ``AzureOpenAIAssistantsClient``. New clients should manage ``client``
|
||||
and ``model`` directly in their own ``__init__``.
|
||||
"""
|
||||
|
||||
INJECTABLE: ClassVar[set[str]] = {"client"}
|
||||
|
||||
def __init__(
|
||||
self, *, model: str | None = None, model_id: str | None = None, client: AsyncOpenAI | None = None, **kwargs: Any
|
||||
) -> None:
|
||||
"""Initialize OpenAIBase.
|
||||
|
||||
Keyword Args:
|
||||
client: The AsyncOpenAI client instance.
|
||||
model: The AI model to use.
|
||||
model_id: Deprecated alias for ``model``.
|
||||
**kwargs: Additional keyword arguments.
|
||||
"""
|
||||
if model_id is not None and model is None:
|
||||
import warnings
|
||||
|
||||
warnings.warn("model_id is deprecated, use model instead", DeprecationWarning, stacklevel=2)
|
||||
model = model_id
|
||||
self.client = client
|
||||
self.model: str | None = None
|
||||
if model:
|
||||
self.model = model.strip()
|
||||
|
||||
# Call super().__init__() to continue MRO chain (e.g., RawChatClient)
|
||||
# Extract known kwargs that belong to other base classes
|
||||
additional_properties = kwargs.pop("additional_properties", None)
|
||||
middleware = kwargs.pop("middleware", None)
|
||||
instruction_role = kwargs.pop("instruction_role", None)
|
||||
function_invocation_configuration = kwargs.pop("function_invocation_configuration", None)
|
||||
|
||||
# Build super().__init__() args
|
||||
super_kwargs = {}
|
||||
if additional_properties is not None:
|
||||
super_kwargs["additional_properties"] = additional_properties
|
||||
if middleware is not None:
|
||||
super_kwargs["middleware"] = middleware
|
||||
if function_invocation_configuration is not None:
|
||||
super_kwargs["function_invocation_configuration"] = function_invocation_configuration
|
||||
|
||||
# Call super().__init__() with filtered kwargs
|
||||
super().__init__(**super_kwargs)
|
||||
|
||||
# Store instruction_role and any remaining kwargs as instance attributes
|
||||
if instruction_role is not None:
|
||||
self.instruction_role = instruction_role
|
||||
for key, value in kwargs.items():
|
||||
setattr(self, key, value)
|
||||
|
||||
async def _initialize_client(self) -> None:
|
||||
"""Initialize OpenAI client asynchronously.
|
||||
|
||||
Override in subclasses to initialize the OpenAI client asynchronously.
|
||||
"""
|
||||
pass
|
||||
|
||||
async def _ensure_client(self) -> AsyncOpenAI:
|
||||
"""Ensure OpenAI client is initialized."""
|
||||
await self._initialize_client()
|
||||
if self.client is None:
|
||||
raise RuntimeError("OpenAI client is not initialized")
|
||||
|
||||
return self.client
|
||||
|
||||
def _get_api_key(
|
||||
self, api_key: str | SecretString | Callable[[], str | Awaitable[str]] | None
|
||||
) -> str | Callable[[], str | Awaitable[str]] | None:
|
||||
"""Get the appropriate API key value for client initialization.
|
||||
|
||||
Args:
|
||||
api_key: The API key parameter which can be a string, SecretString, callable, or None.
|
||||
|
||||
Returns:
|
||||
For callable API keys: returns the callable directly.
|
||||
For SecretString/string/None API keys: returns as-is (SecretString is a str subclass).
|
||||
"""
|
||||
if isinstance(api_key, SecretString):
|
||||
return api_key.get_secret_value()
|
||||
|
||||
# Check version compatibility for callable API keys
|
||||
if callable(api_key):
|
||||
_check_openai_version_for_callable_api_key()
|
||||
|
||||
return api_key # Pass callable, string, or None directly to OpenAI SDK
|
||||
|
||||
|
||||
class OpenAIConfigMixin(OpenAIBase):
|
||||
"""Internal class for configuring a connection to an OpenAI service.
|
||||
|
||||
.. deprecated::
|
||||
``OpenAIConfigMixin`` is deprecated and only used by ``OpenAIAssistantsClient``
|
||||
and ``AzureOpenAIAssistantsClient``. New clients handle configuration
|
||||
directly in their own ``__init__``.
|
||||
"""
|
||||
|
||||
OTEL_PROVIDER_NAME: ClassVar[str] = "openai" # type: ignore[reportIncompatibleVariableOverride, misc]
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model: str,
|
||||
api_key: str | Callable[[], str | Awaitable[str]] | None = None,
|
||||
org_id: str | None = None,
|
||||
default_headers: Mapping[str, str] | None = None,
|
||||
client: AsyncOpenAI | None = None,
|
||||
instruction_role: str | None = None,
|
||||
base_url: str | None = None,
|
||||
**kwargs: Any,
|
||||
) -> None:
|
||||
"""Initialize a client for OpenAI services.
|
||||
|
||||
This constructor sets up a client to interact with OpenAI's API, allowing for
|
||||
different types of AI model interactions, like chat or text completion.
|
||||
|
||||
Args:
|
||||
model: OpenAI model identifier. Must be non-empty.
|
||||
Default to a preset value.
|
||||
api_key: OpenAI API key for authentication, or a callable that returns an API key.
|
||||
Must be non-empty. (Optional)
|
||||
org_id: OpenAI organization ID. This is optional
|
||||
unless the account belongs to multiple organizations.
|
||||
default_headers: Default headers
|
||||
for HTTP requests. (Optional)
|
||||
client: An existing OpenAI client, optional.
|
||||
instruction_role: The role to use for 'instruction'
|
||||
messages, for example, summarization prompts could use `developer` or `system`. (Optional)
|
||||
base_url: The optional base URL to use. If provided will override the standard value for a OpenAI connector.
|
||||
Will not be used when supplying a custom client.
|
||||
kwargs: Additional keyword arguments.
|
||||
|
||||
"""
|
||||
# Merge APP_INFO into the headers if it exists
|
||||
merged_headers = dict(copy(default_headers)) if default_headers else {}
|
||||
if APP_INFO:
|
||||
merged_headers.update(APP_INFO)
|
||||
merged_headers = prepend_agent_framework_to_user_agent(merged_headers)
|
||||
|
||||
# Handle callable API key using base class method
|
||||
api_key_value = self._get_api_key(api_key)
|
||||
|
||||
if not client:
|
||||
if not api_key:
|
||||
raise ValueError("Please provide an api_key")
|
||||
args: dict[str, Any] = {"api_key": api_key_value, "default_headers": merged_headers}
|
||||
if org_id:
|
||||
args["organization"] = org_id
|
||||
if base_url:
|
||||
args["base_url"] = base_url
|
||||
client = AsyncOpenAI(**args)
|
||||
|
||||
# Store configuration as instance attributes for serialization
|
||||
self.org_id = org_id
|
||||
self.base_url = str(base_url)
|
||||
# Store default_headers but filter out USER_AGENT_KEY for serialization
|
||||
if default_headers:
|
||||
self.default_headers: dict[str, Any] | None = {
|
||||
k: v for k, v in default_headers.items() if k != USER_AGENT_KEY
|
||||
}
|
||||
else:
|
||||
self.default_headers = None
|
||||
|
||||
args = {
|
||||
"model": model,
|
||||
"client": client,
|
||||
}
|
||||
if instruction_role:
|
||||
args["instruction_role"] = instruction_role
|
||||
|
||||
# Ensure additional_properties and middleware are passed through kwargs to RawChatClient
|
||||
# These are consumed by RawChatClient.__init__ via kwargs
|
||||
super().__init__(**args, **kwargs)
|
||||
|
||||
|
||||
def to_assistant_tools(
|
||||
tools: Sequence[FunctionTool | MutableMapping[str, Any]] | None,
|
||||
) -> list[dict[str, Any]]:
|
||||
"""Convert Agent Framework tools to OpenAI Assistants API format.
|
||||
|
||||
Handles FunctionTool instances and dict-based tools from static factory methods.
|
||||
|
||||
Args:
|
||||
tools: Sequence of Agent Framework tools.
|
||||
|
||||
Returns:
|
||||
List of tool definitions for OpenAI Assistants API.
|
||||
"""
|
||||
if not tools:
|
||||
return []
|
||||
|
||||
tool_definitions: list[dict[str, Any]] = []
|
||||
|
||||
for tool in tools:
|
||||
if isinstance(tool, FunctionTool):
|
||||
tool_definitions.append(tool.to_json_schema_spec())
|
||||
elif isinstance(tool, MutableMapping):
|
||||
# Pass through dict-based tools directly (from static factory methods)
|
||||
tool_definitions.append(dict(tool))
|
||||
|
||||
return tool_definitions
|
||||
|
||||
|
||||
def from_assistant_tools(
|
||||
assistant_tools: list[Any] | None,
|
||||
) -> list[dict[str, Any]]:
|
||||
"""Convert OpenAI Assistant tools to dict-based format.
|
||||
|
||||
This converts hosted tools (code_interpreter, file_search) from an OpenAI
|
||||
Assistant definition back to dict-based tool definitions.
|
||||
|
||||
Note: Function tools are skipped - user must provide implementations separately.
|
||||
|
||||
Args:
|
||||
assistant_tools: Tools from OpenAI Assistant object (assistant.tools).
|
||||
|
||||
Returns:
|
||||
List of dict-based tool definitions for hosted tools.
|
||||
"""
|
||||
if not assistant_tools:
|
||||
return []
|
||||
|
||||
tools: list[dict[str, Any]] = []
|
||||
|
||||
for tool in assistant_tools:
|
||||
if hasattr(tool, "type"):
|
||||
tool_type = tool.type
|
||||
elif isinstance(tool, Mapping):
|
||||
typed_tool = cast(Mapping[str, Any], tool)
|
||||
tool_type_value: Any = typed_tool.get("type")
|
||||
tool_type = tool_type_value if isinstance(tool_type_value, str) else None
|
||||
else:
|
||||
tool_type = None
|
||||
|
||||
if tool_type == "code_interpreter":
|
||||
tools.append({"type": "code_interpreter"})
|
||||
elif tool_type == "file_search":
|
||||
tools.append({"type": "file_search"})
|
||||
# Skip function tools - user must provide implementations
|
||||
|
||||
return tools
|
||||
Reference in New Issue
Block a user