mirror of
https://github.com/microsoft/agent-framework.git
synced 2026-06-16 21:04:09 +08:00
Python: added ChatClientBase with function calling (#147)
* added ChatClientBase with function calling * streaming update * fixed typing * test setup * small update * src setup * removed src, updated test naming * fixed test command * alolow args * updated test run * added unit test folder to azure * added init and unit test to azure * added other cross tests * restructured * reset test run * fix name * removed always * updated test * extend pytest.xml locations * run surface always * added decorators for FC and marked tests * fixed mypy settings and added tests * fix override import * removed import
This commit is contained in:
Eduard van Valkenburg
committed by
GitHub
Unverified
GitHub
parent
daf4788868
commit
3449902b03
@@ -0,0 +1,57 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import importlib
|
||||
import importlib.metadata
|
||||
|
||||
try:
|
||||
__version__ = importlib.metadata.version(__name__)
|
||||
except importlib.metadata.PackageNotFoundError:
|
||||
__version__ = "0.0.0" # Fallback for development mode
|
||||
|
||||
_IMPORTS = {
|
||||
"get_logger": "._logging",
|
||||
"Agent": "._agents",
|
||||
"AgentThread": "._agents",
|
||||
"AITool": "._tools",
|
||||
"ai_function": "._tools",
|
||||
"AIContent": "._types",
|
||||
"AIContents": "._types",
|
||||
"TextContent": "._types",
|
||||
"TextReasoningContent": "._types",
|
||||
"DataContent": "._types",
|
||||
"UriContent": "._types",
|
||||
"UsageContent": "._types",
|
||||
"UsageDetails": "._types",
|
||||
"FunctionCallContent": "._types",
|
||||
"FunctionResultContent": "._types",
|
||||
"ChatFinishReason": "._types",
|
||||
"ChatMessage": "._types",
|
||||
"ChatResponse": "._types",
|
||||
"StructuredResponse": "._types",
|
||||
"ChatResponseUpdate": "._types",
|
||||
"ChatRole": "._types",
|
||||
"ErrorContent": "._types",
|
||||
"GeneratedEmbeddings": "._types",
|
||||
"ChatOptions": "._types",
|
||||
"ChatToolMode": "._types",
|
||||
"ChatClient": "._clients",
|
||||
"ChatClientBase": "._clients",
|
||||
"use_tool_calling": "._clients",
|
||||
"EmbeddingGenerator": "._clients",
|
||||
"InputGuardrail": ".guard_rails",
|
||||
"OutputGuardrail": ".guard_rails",
|
||||
}
|
||||
|
||||
|
||||
def __getattr__(name: str):
|
||||
if name == "__version__":
|
||||
return __version__
|
||||
if name in _IMPORTS:
|
||||
submod_name = _IMPORTS[name]
|
||||
module = importlib.import_module(submod_name, package=__name__)
|
||||
return getattr(module, name)
|
||||
raise AttributeError(f"module {__name__} has no attribute {name}")
|
||||
|
||||
|
||||
def __dir__():
|
||||
return [*list(_IMPORTS.keys()), "__version__"]
|
||||
@@ -0,0 +1,65 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from . import __version__ # type: ignore[attr-defined]
|
||||
from ._agents import Agent, AgentThread
|
||||
from ._clients import ChatClient, ChatClientBase, EmbeddingGenerator, use_tool_calling
|
||||
from ._logging import get_logger
|
||||
from ._tools import AITool, ai_function
|
||||
from ._types import (
|
||||
AIContent,
|
||||
AIContents,
|
||||
ChatFinishReason,
|
||||
ChatMessage,
|
||||
ChatOptions,
|
||||
ChatResponse,
|
||||
ChatResponseUpdate,
|
||||
ChatRole,
|
||||
ChatToolMode,
|
||||
DataContent,
|
||||
ErrorContent,
|
||||
FunctionCallContent,
|
||||
FunctionResultContent,
|
||||
GeneratedEmbeddings,
|
||||
StructuredResponse,
|
||||
TextContent,
|
||||
TextReasoningContent,
|
||||
UriContent,
|
||||
UsageContent,
|
||||
UsageDetails,
|
||||
)
|
||||
from .guard_rails import InputGuardrail, OutputGuardrail
|
||||
|
||||
__all__ = [
|
||||
"AIContent",
|
||||
"AIContents",
|
||||
"AITool",
|
||||
"Agent",
|
||||
"AgentThread",
|
||||
"ChatClient",
|
||||
"ChatClientBase",
|
||||
"ChatFinishReason",
|
||||
"ChatMessage",
|
||||
"ChatOptions",
|
||||
"ChatResponse",
|
||||
"ChatResponseUpdate",
|
||||
"ChatRole",
|
||||
"ChatToolMode",
|
||||
"DataContent",
|
||||
"EmbeddingGenerator",
|
||||
"ErrorContent",
|
||||
"FunctionCallContent",
|
||||
"FunctionResultContent",
|
||||
"GeneratedEmbeddings",
|
||||
"InputGuardrail",
|
||||
"OutputGuardrail",
|
||||
"StructuredResponse",
|
||||
"TextContent",
|
||||
"TextReasoningContent",
|
||||
"UriContent",
|
||||
"UsageContent",
|
||||
"UsageDetails",
|
||||
"__version__",
|
||||
"ai_function",
|
||||
"get_logger",
|
||||
"use_tool_calling",
|
||||
]
|
||||
@@ -0,0 +1,148 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from abc import abstractmethod
|
||||
from collections.abc import AsyncIterable, Sequence
|
||||
from typing import Any, Protocol, runtime_checkable
|
||||
|
||||
from ._pydantic import AFBaseModel
|
||||
from ._types import ChatMessage, ChatResponse, ChatResponseUpdate
|
||||
|
||||
# region AgentThread
|
||||
|
||||
|
||||
class AgentThread(AFBaseModel):
|
||||
"""Base class for agent threads."""
|
||||
|
||||
id: str | None = None
|
||||
|
||||
async def create(self) -> str | None:
|
||||
"""Starts the thread and returns the thread ID."""
|
||||
# If the thread ID is already set, we're done, just return the Id.
|
||||
if self.id is not None:
|
||||
return self.id
|
||||
|
||||
# Otherwise, create the thread.
|
||||
self.id = await self._create()
|
||||
return self.id
|
||||
|
||||
async def delete(self) -> None:
|
||||
"""Ends the current thread."""
|
||||
await self._delete()
|
||||
self.id = None
|
||||
|
||||
async def on_new_message(
|
||||
self,
|
||||
new_messages: ChatMessage | Sequence[ChatMessage],
|
||||
) -> None:
|
||||
"""Invoked when a new message has been contributed to the chat by any participant."""
|
||||
# If the thread is not created yet, create it.
|
||||
if self.id is None:
|
||||
await self.create()
|
||||
|
||||
await self._on_new_message(new_messages=new_messages)
|
||||
|
||||
@abstractmethod
|
||||
async def _create(self) -> str:
|
||||
"""Starts the thread and returns the thread ID."""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def _delete(self) -> None:
|
||||
"""Ends the current thread."""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def _on_new_message(
|
||||
self,
|
||||
new_messages: ChatMessage | Sequence[ChatMessage],
|
||||
) -> None:
|
||||
"""Invoked when a new message has been contributed to the chat by any participant."""
|
||||
...
|
||||
|
||||
|
||||
# region Agent Protocol
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class Agent(Protocol):
|
||||
"""A protocol for an agent that can be invoked."""
|
||||
|
||||
@property
|
||||
def id(self) -> str:
|
||||
"""Returns the ID of the agent."""
|
||||
...
|
||||
|
||||
@property
|
||||
def name(self) -> str | None:
|
||||
"""Returns the name of the agent."""
|
||||
...
|
||||
|
||||
@property
|
||||
def description(self) -> str | None:
|
||||
"""Returns the description of the agent."""
|
||||
...
|
||||
|
||||
@property
|
||||
def instructions(self) -> str | None:
|
||||
"""Returns the instructions for the agent."""
|
||||
...
|
||||
|
||||
async def run(
|
||||
self,
|
||||
messages: str | ChatMessage | list[ChatMessage] | None = None,
|
||||
*,
|
||||
thread: AgentThread | None = None,
|
||||
**kwargs: Any,
|
||||
) -> ChatResponse:
|
||||
"""Get a response from the agent.
|
||||
|
||||
This method returns the final result of the agent's execution
|
||||
as a single ChatResponse object. The caller is blocked until
|
||||
the final result is available.
|
||||
|
||||
Note: For streaming responses, use the run_stream method, which returns
|
||||
intermediate steps and the final result as a stream of ChatResponseUpdate
|
||||
objects. Streaming only the final result is not feasible because the timing of
|
||||
the final result's availability is unknown, and blocking the caller until then
|
||||
is undesirable in streaming scenarios.
|
||||
|
||||
Args:
|
||||
messages: The message(s) to send to the agent.
|
||||
thread: The conversation thread associated with the message(s).
|
||||
kwargs: Additional keyword arguments.
|
||||
|
||||
Returns:
|
||||
An agent response item.
|
||||
"""
|
||||
...
|
||||
|
||||
def run_stream(
|
||||
self,
|
||||
messages: str | ChatMessage | list[ChatMessage] | None = None,
|
||||
*,
|
||||
thread: AgentThread | None = None,
|
||||
**kwargs: Any,
|
||||
) -> AsyncIterable[ChatResponseUpdate]:
|
||||
"""Run the agent as a stream.
|
||||
|
||||
This method will return the intermediate steps and final results of the
|
||||
agent's execution as a stream of ChatResponseUpdate objects to the caller.
|
||||
|
||||
To get the intermediate steps of the agent's execution as fully formed messages,
|
||||
use the on_intermediate_message callback.
|
||||
|
||||
Note: A ChatResponseUpdate object contains a chunk of a message.
|
||||
|
||||
Args:
|
||||
messages: The message(s) to send to the agent.
|
||||
thread: The conversation thread associated with the message(s).
|
||||
kwargs: Additional keyword arguments.
|
||||
|
||||
Yields:
|
||||
An agent response item.
|
||||
"""
|
||||
...
|
||||
|
||||
def get_new_thread(self) -> AgentThread:
|
||||
"""Creates a new conversation thread for the agent."""
|
||||
...
|
||||
@@ -0,0 +1,50 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import threading
|
||||
from asyncio import Future
|
||||
from collections.abc import Callable
|
||||
from typing import Any
|
||||
|
||||
|
||||
# from https://github.com/microsoft/autogen/blob/main/python/packages/autogen-core/src/autogen_core/_cancellation_token.py
|
||||
class CancellationToken:
|
||||
"""A token used to cancel pending async calls."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._cancelled: bool = False
|
||||
self._lock: threading.Lock = threading.Lock()
|
||||
self._callbacks: list[Callable[[], None]] = []
|
||||
|
||||
def cancel(self) -> None:
|
||||
"""Cancel pending async calls linked to this cancellation token."""
|
||||
with self._lock:
|
||||
if not self._cancelled:
|
||||
self._cancelled = True
|
||||
for callback in self._callbacks:
|
||||
callback()
|
||||
|
||||
def is_cancelled(self) -> bool:
|
||||
"""Check if the CancellationToken has been used."""
|
||||
with self._lock:
|
||||
return self._cancelled
|
||||
|
||||
def add_callback(self, callback: Callable[[], None]) -> None:
|
||||
"""Attach a callback that will be called when cancel is invoked."""
|
||||
with self._lock:
|
||||
if self._cancelled:
|
||||
callback()
|
||||
else:
|
||||
self._callbacks.append(callback)
|
||||
|
||||
def link_future(self, future: Future[Any]) -> Future[Any]:
|
||||
"""Link a pending async call to a token to allow its cancellation."""
|
||||
with self._lock:
|
||||
if self._cancelled:
|
||||
future.cancel()
|
||||
else:
|
||||
|
||||
def _cancel() -> None:
|
||||
future.cancel()
|
||||
|
||||
self._callbacks.append(_cancel)
|
||||
return future
|
||||
@@ -0,0 +1,530 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import asyncio
|
||||
from abc import ABC, abstractmethod
|
||||
from collections.abc import AsyncIterable, Awaitable, Callable, MutableSequence, Sequence
|
||||
from functools import wraps
|
||||
from typing import Annotated, Any, Generic, Literal, Protocol, TypeVar, runtime_checkable
|
||||
|
||||
from pydantic import BaseModel, PrivateAttr, StringConstraints
|
||||
|
||||
from ._logging import get_logger
|
||||
from ._pydantic import AFBaseModel
|
||||
from ._tools import AIFunction, AITool
|
||||
from ._types import (
|
||||
AIContents,
|
||||
ChatMessage,
|
||||
ChatOptions,
|
||||
ChatResponse,
|
||||
ChatResponseUpdate,
|
||||
ChatToolMode,
|
||||
FunctionCallContent,
|
||||
FunctionResultContent,
|
||||
GeneratedEmbeddings,
|
||||
)
|
||||
|
||||
TInput = TypeVar("TInput", contravariant=True)
|
||||
TEmbedding = TypeVar("TEmbedding")
|
||||
TInnerGetResponse = TypeVar("TInnerGetResponse", bound=Callable[..., Awaitable[ChatResponse]])
|
||||
TInnerGetStreamingResponse = TypeVar(
|
||||
"TInnerGetStreamingResponse", bound=Callable[..., AsyncIterable[ChatResponseUpdate]]
|
||||
)
|
||||
|
||||
TChatClientBase = TypeVar("TChatClientBase", bound="ChatClientBase")
|
||||
|
||||
logger = get_logger()
|
||||
|
||||
# region: Tool Calling Functions and Decorators
|
||||
|
||||
|
||||
def _merge_function_results(
|
||||
messages: list[ChatMessage],
|
||||
) -> ChatMessage:
|
||||
"""Combine multiple function result content types to one chat message content type.
|
||||
|
||||
This method combines the FunctionResultContent items from separate ChatMessageContent messages,
|
||||
and is used in the event that the `context.terminate = True` condition is met.
|
||||
"""
|
||||
contents: list[Any] = []
|
||||
for message in messages:
|
||||
contents.extend([item for item in message.contents if isinstance(item, FunctionResultContent)])
|
||||
|
||||
return ChatMessage(
|
||||
role="tool",
|
||||
contents=contents,
|
||||
)
|
||||
|
||||
|
||||
async def _auto_invoke_function(
|
||||
function_call_content: FunctionCallContent,
|
||||
custom_args: dict[str, Any] | None = None,
|
||||
*,
|
||||
tool_map: dict[str, AIFunction[BaseModel, Any]],
|
||||
sequence_index: int | None = None,
|
||||
request_index: int | None = None,
|
||||
) -> AIContents:
|
||||
"""Invoke a function call requested by the agent, applying filters that are defined in the agent."""
|
||||
tool: AIFunction[BaseModel, Any] | None = tool_map.get(function_call_content.name)
|
||||
if tool is None:
|
||||
raise KeyError(f"No tool or function named '{function_call_content.name}'")
|
||||
|
||||
parsed_args: dict[str, Any] = dict(function_call_content.parse_arguments() or {})
|
||||
|
||||
# Merge with user-supplied args; right-hand side dominates, so parsed args win on conflicts.
|
||||
merged_args: dict[str, Any] = (custom_args or {}) | parsed_args
|
||||
args = tool.input_model.model_validate(merged_args)
|
||||
exception = None
|
||||
try:
|
||||
function_result = await tool.invoke(arguments=args)
|
||||
except Exception as ex:
|
||||
exception = ex
|
||||
function_result = None
|
||||
return FunctionResultContent(
|
||||
call_id=function_call_content.call_id,
|
||||
exception=exception,
|
||||
result=function_result,
|
||||
)
|
||||
|
||||
|
||||
def _tool_to_json_schema_spec(tool: AITool) -> dict[str, Any]:
|
||||
"""Convert a AITool to the JSON Schema function specification format."""
|
||||
return {
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": tool.name,
|
||||
"description": tool.description,
|
||||
"parameters": tool.parameters(),
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def _prepare_tools_and_tool_choice(chat_options: ChatOptions) -> None:
|
||||
"""Prepare the tools and tool choice for the chat options."""
|
||||
chat_tool_mode: ChatToolMode | None = chat_options.tool_choice # type: ignore
|
||||
if chat_tool_mode is None or chat_tool_mode == ChatToolMode.NONE:
|
||||
chat_options.tools = None
|
||||
chat_options.tool_choice = ChatToolMode.NONE.mode
|
||||
return
|
||||
chat_options.tools = [
|
||||
(_tool_to_json_schema_spec(t) if isinstance(t, AITool) else t) for t in chat_options.tools or []
|
||||
]
|
||||
chat_options.tool_choice = chat_tool_mode.mode
|
||||
|
||||
|
||||
def _tool_call_non_streaming(func: TInnerGetResponse) -> TInnerGetResponse:
|
||||
"""Decorate the internal _inner_get_response method to enable tool calls.
|
||||
|
||||
Remarks:
|
||||
Relies on a class that has the _tool_map attribute for the executable tools to call.
|
||||
"""
|
||||
|
||||
@wraps(func)
|
||||
async def wrapper(
|
||||
self: "ChatClientBase",
|
||||
*,
|
||||
messages: MutableSequence[ChatMessage],
|
||||
chat_options: ChatOptions,
|
||||
**kwargs: Any,
|
||||
) -> ChatResponse:
|
||||
response: ChatResponse | None = None
|
||||
fcc_messages: list[ChatMessage] = []
|
||||
for attempt_idx in range(self.maximum_iterations_per_request):
|
||||
response = await func(self, messages=messages, chat_options=chat_options)
|
||||
# if there are function calls, we will handle them first
|
||||
function_calls = [it for it in response.messages[0].contents if isinstance(it, FunctionCallContent)]
|
||||
if function_calls:
|
||||
# Run all function calls concurrently
|
||||
results = await asyncio.gather(*[
|
||||
_auto_invoke_function(
|
||||
function_call,
|
||||
custom_args=kwargs,
|
||||
tool_map=self._tool_map,
|
||||
sequence_index=seq_idx,
|
||||
request_index=attempt_idx,
|
||||
)
|
||||
for seq_idx, function_call in enumerate(function_calls)
|
||||
])
|
||||
# add a single ChatMessage to the response with the results
|
||||
response.messages.append(ChatMessage(role="tool", contents=results))
|
||||
# response should contain 2 messages after this,
|
||||
# one with function call contents
|
||||
# and one with function result contents
|
||||
# the amount and call_id's should match
|
||||
# this runs in every but the first run
|
||||
# we need to keep track of all function call messages
|
||||
fcc_messages.extend(response.messages)
|
||||
# and add them as additional context to the messages
|
||||
messages.extend(response.messages)
|
||||
continue
|
||||
# If we reach this point, it means there were no function calls to handle,
|
||||
# we'll add the previous function call and responses
|
||||
# to the front of the list, so that the final response is the last one
|
||||
# TODO (eavanvalkenburg): control this behavior?
|
||||
if fcc_messages:
|
||||
for msg in reversed(fcc_messages):
|
||||
response.messages.insert(0, msg)
|
||||
return response
|
||||
|
||||
# Failsafe: give up on tools, ask model for plain answer
|
||||
chat_options.tool_choice = "none"
|
||||
_prepare_tools_and_tool_choice(chat_options=chat_options)
|
||||
response = await func(self, messages=messages, chat_options=chat_options)
|
||||
if fcc_messages:
|
||||
for msg in reversed(fcc_messages):
|
||||
response.messages.insert(0, msg)
|
||||
return response
|
||||
|
||||
return wrapper # type: ignore[reportReturnType, return-value]
|
||||
|
||||
|
||||
def _tool_call_streaming(func: TInnerGetStreamingResponse) -> TInnerGetStreamingResponse:
|
||||
"""Decorate the internal _inner_get_response method to enable tool calls.
|
||||
|
||||
Remarks:
|
||||
Relies on a class that has the _tool_map attribute for the executable tools to call.
|
||||
"""
|
||||
|
||||
@wraps(func)
|
||||
async def wrapper(
|
||||
self: "ChatClientBase",
|
||||
*,
|
||||
messages: MutableSequence[ChatMessage],
|
||||
chat_options: ChatOptions,
|
||||
**kwargs: Any,
|
||||
) -> AsyncIterable[ChatResponseUpdate]:
|
||||
for attempt_idx in range(self.maximum_iterations_per_request):
|
||||
function_call_returned = False
|
||||
all_messages: list[ChatResponseUpdate] = []
|
||||
async for update in func(self, messages=messages, chat_options=chat_options):
|
||||
if update.contents and any(isinstance(item, FunctionCallContent) for item in update.contents):
|
||||
all_messages.append(update)
|
||||
function_call_returned = True
|
||||
yield update
|
||||
|
||||
if not function_call_returned:
|
||||
return
|
||||
|
||||
# There is one FunctionCallContent response stream in the messages, combining now to create
|
||||
# the full completion depending on the prompt, the message may contain both function call
|
||||
# content and others
|
||||
response: ChatResponse = ChatResponse.from_chat_response_updates(all_messages)
|
||||
function_calls = [item for item in response.messages[0].contents if isinstance(item, FunctionCallContent)]
|
||||
messages.append(response.messages[0])
|
||||
|
||||
if function_calls:
|
||||
# Run all function calls concurrently
|
||||
results = await asyncio.gather(*[
|
||||
_auto_invoke_function(
|
||||
function_call,
|
||||
custom_args=kwargs,
|
||||
tool_map=self._tool_map,
|
||||
sequence_index=seq_idx,
|
||||
request_index=attempt_idx,
|
||||
)
|
||||
for seq_idx, function_call in enumerate(function_calls)
|
||||
])
|
||||
yield ChatResponseUpdate(contents=results, role="tool")
|
||||
response.messages.append(ChatMessage(role="tool", contents=results))
|
||||
messages.extend(response.messages)
|
||||
continue
|
||||
|
||||
# Failsafe: give up on tools, ask model for plain answer
|
||||
chat_options.tool_choice = "none"
|
||||
_prepare_tools_and_tool_choice(chat_options=chat_options)
|
||||
async for update in func(self, messages=messages, chat_options=chat_options, **kwargs):
|
||||
yield update
|
||||
|
||||
return wrapper # type: ignore[reportReturnType, return-value]
|
||||
|
||||
|
||||
def use_tool_calling(cls: type[TChatClientBase]) -> type[TChatClientBase]:
|
||||
inner_response = getattr(cls, "_inner_get_response", None)
|
||||
if inner_response is not None:
|
||||
cls._inner_get_response = _tool_call_non_streaming(inner_response) # type: ignore
|
||||
inner_streaming_response = getattr(cls, "_inner_get_streaming_response", None)
|
||||
if inner_streaming_response is not None:
|
||||
cls._inner_get_streaming_response = _tool_call_streaming(inner_streaming_response) # type: ignore
|
||||
return cls
|
||||
|
||||
|
||||
# region: ChatClient Protocol
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class ChatClient(Protocol):
|
||||
"""A protocol for a chat client that can generate responses."""
|
||||
|
||||
async def get_response(
|
||||
self,
|
||||
messages: str | ChatMessage | Sequence[ChatMessage],
|
||||
**kwargs: Any,
|
||||
) -> ChatResponse:
|
||||
"""Sends input and returns the response.
|
||||
|
||||
Args:
|
||||
messages: The sequence of input messages to send.
|
||||
**kwargs: Additional options for the request, such as ai_model_id, temperature, etc.
|
||||
See `ChatOptions` for more details.
|
||||
|
||||
Returns:
|
||||
The response messages generated by the client.
|
||||
|
||||
Raises:
|
||||
ValueError: If the input message sequence is `None`.
|
||||
"""
|
||||
...
|
||||
|
||||
async def get_streaming_response(
|
||||
self,
|
||||
messages: str | ChatMessage | Sequence[ChatMessage],
|
||||
**kwargs: Any,
|
||||
) -> AsyncIterable[ChatResponseUpdate]:
|
||||
"""Sends input messages and streams the response.
|
||||
|
||||
Args:
|
||||
messages: The sequence of input messages to send.
|
||||
**kwargs: Additional options for the request, such as ai_model_id, temperature, etc.
|
||||
See `ChatOptions` for more details.
|
||||
|
||||
Yields:
|
||||
An async iterable of chat response updates containing the content of the response messages
|
||||
generated by the client.
|
||||
|
||||
Raises:
|
||||
ValueError: If the input message sequence is `None`.
|
||||
"""
|
||||
...
|
||||
|
||||
|
||||
class ChatClientBase(AFBaseModel, ABC):
|
||||
"""Base class for chat clients."""
|
||||
|
||||
ai_model_id: Annotated[str, StringConstraints(strip_whitespace=True, min_length=1)]
|
||||
maximum_iterations_per_request: int = 10
|
||||
_tool_map: dict[str, AIFunction[BaseModel, Any]] = PrivateAttr(default_factory=dict) # type: ignore
|
||||
|
||||
# region Internal methods to be implemented by the derived classes
|
||||
|
||||
@abstractmethod
|
||||
async def _inner_get_response(
|
||||
self,
|
||||
*,
|
||||
messages: MutableSequence[ChatMessage],
|
||||
chat_options: ChatOptions,
|
||||
**kwargs: Any,
|
||||
) -> ChatResponse:
|
||||
"""Send a chat request to the AI service.
|
||||
|
||||
Args:
|
||||
messages: The chat messages to send.
|
||||
chat_options: The options for the request.
|
||||
kwargs: Any additional keyword arguments.
|
||||
|
||||
Returns:
|
||||
The chat response contents representing the response(s).
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
async def _inner_get_streaming_response(
|
||||
self,
|
||||
*,
|
||||
messages: MutableSequence[ChatMessage],
|
||||
chat_options: ChatOptions,
|
||||
**kwargs: Any,
|
||||
) -> AsyncIterable[ChatResponseUpdate]:
|
||||
"""Send a streaming chat request to the AI service.
|
||||
|
||||
Args:
|
||||
messages: The chat messages to send.
|
||||
chat_options: The chat_options for the request.
|
||||
kwargs: Any additional keyword arguments.
|
||||
|
||||
Yields:
|
||||
ChatResponseUpdate: The streaming chat message contents.
|
||||
"""
|
||||
# Below is needed for mypy: https://mypy.readthedocs.io/en/stable/more_types.html#asynchronous-iterators
|
||||
if False:
|
||||
yield
|
||||
await asyncio.sleep(0) # pragma: no cover
|
||||
# This is a no-op, but it allows the method to be async and return an AsyncIterable.
|
||||
# The actual implementation should yield ChatResponseUpdate instances as needed.
|
||||
|
||||
# endregion
|
||||
|
||||
# region Public method
|
||||
|
||||
async def get_response(
|
||||
self,
|
||||
messages: str | ChatMessage | list[ChatMessage],
|
||||
*,
|
||||
model: str | None = None,
|
||||
max_tokens: int | None = None,
|
||||
temperature: float | None = None,
|
||||
top_p: float | None = None,
|
||||
tool_choice: ChatToolMode | Literal["auto", "required", "none"] | dict[str, Any] | None = None,
|
||||
tools: Sequence[AITool] | None = None,
|
||||
response_format: type[BaseModel] | None = None,
|
||||
user: str | None = None,
|
||||
stop: str | Sequence[str] | None = None,
|
||||
frequency_penalty: float | None = None,
|
||||
logit_bias: dict[str | int, float] | None = None,
|
||||
presence_penalty: float | None = None,
|
||||
seed: int | None = None,
|
||||
store: bool | None = None,
|
||||
metadata: dict[str, Any] | None = None,
|
||||
additional_properties: dict[str, Any] | None = None,
|
||||
**kwargs: Any,
|
||||
) -> ChatResponse:
|
||||
"""Get a response from a chat client.
|
||||
|
||||
Args:
|
||||
messages: the message or messages to send to the model
|
||||
model: the model to use for the request
|
||||
max_tokens: the maximum number of tokens to generate
|
||||
temperature: the sampling temperature to use
|
||||
top_p: the nucleus sampling probability to use
|
||||
tool_choice: the tool choice for the request
|
||||
tools: the tools to use for the request
|
||||
response_format: the format of the response
|
||||
user: the user to associate with the request
|
||||
stop: the stop sequence(s) for the request
|
||||
frequency_penalty: the frequency penalty to use
|
||||
logit_bias: the logit bias to use
|
||||
presence_penalty: the presence penalty to use
|
||||
seed: the random seed to use
|
||||
store: whether to store the response
|
||||
metadata: additional metadata to include in the request
|
||||
additional_properties: additional properties to include in the request
|
||||
kwargs: any additional keyword arguments,
|
||||
will only be passed to functions that are called.
|
||||
|
||||
Returns:
|
||||
A chat response from the model.
|
||||
"""
|
||||
if tools is not None:
|
||||
self._tool_map = {tool.name: tool for tool in tools if isinstance(tool, AIFunction)}
|
||||
chat_options = ChatOptions(
|
||||
ai_model_id=model,
|
||||
max_tokens=max_tokens,
|
||||
temperature=temperature,
|
||||
top_p=top_p,
|
||||
tool_choice=tool_choice,
|
||||
tools=tools,
|
||||
response_format=response_format,
|
||||
user=user,
|
||||
stop=stop,
|
||||
frequency_penalty=frequency_penalty,
|
||||
logit_bias=logit_bias,
|
||||
presence_penalty=presence_penalty,
|
||||
seed=seed,
|
||||
store=store,
|
||||
metadata=metadata,
|
||||
additional_properties=additional_properties or {},
|
||||
)
|
||||
if isinstance(messages, str):
|
||||
messages = [ChatMessage(role="user", text=messages)]
|
||||
if isinstance(messages, ChatMessage):
|
||||
messages = [messages]
|
||||
_prepare_tools_and_tool_choice(chat_options=chat_options)
|
||||
return await self._inner_get_response(messages=messages, chat_options=chat_options, **kwargs)
|
||||
|
||||
async def get_streaming_response(
|
||||
self,
|
||||
messages: str | ChatMessage | list[ChatMessage],
|
||||
*,
|
||||
model: str | None = None,
|
||||
max_tokens: int | None = None,
|
||||
temperature: float | None = None,
|
||||
top_p: float | None = None,
|
||||
tool_choice: ChatToolMode | Literal["auto", "required", "none"] | dict[str, Any] | None = None,
|
||||
tools: Sequence[AITool] | None = None,
|
||||
response_format: type[BaseModel] | None = None,
|
||||
user: str | None = None,
|
||||
stop: str | Sequence[str] | None = None,
|
||||
frequency_penalty: float | None = None,
|
||||
logit_bias: dict[str | int, float] | None = None,
|
||||
presence_penalty: float | None = None,
|
||||
seed: int | None = None,
|
||||
store: bool | None = None,
|
||||
metadata: dict[str, Any] | None = None,
|
||||
additional_properties: dict[str, Any] | None = None,
|
||||
**kwargs: Any,
|
||||
) -> AsyncIterable[ChatResponseUpdate]:
|
||||
"""Get a streaming response from a chat client.
|
||||
|
||||
Args:
|
||||
messages: the message or messages to send to the model
|
||||
model: the model to use for the request
|
||||
max_tokens: the maximum number of tokens to generate
|
||||
temperature: the sampling temperature to use
|
||||
top_p: the nucleus sampling probability to use
|
||||
tool_choice: the tool choice for the request
|
||||
tools: the tools to use for the request
|
||||
response_format: the format of the response
|
||||
user: the user to associate with the request
|
||||
stop: the stop sequence(s) for the request
|
||||
frequency_penalty: the frequency penalty to use
|
||||
logit_bias: the logit bias to use
|
||||
presence_penalty: the presence penalty to use
|
||||
seed: the random seed to use
|
||||
store: whether to store the response
|
||||
metadata: additional metadata to include in the request
|
||||
additional_properties: additional properties to include in the request
|
||||
kwargs: any additional keyword arguments
|
||||
|
||||
Yields:
|
||||
A stream representing the response(s) from the LLM.
|
||||
"""
|
||||
if tools is not None:
|
||||
self._tool_map = {tool.name: tool for tool in tools if isinstance(tool, AIFunction)}
|
||||
chat_options = ChatOptions(
|
||||
ai_model_id=model,
|
||||
max_tokens=max_tokens,
|
||||
temperature=temperature,
|
||||
top_p=top_p,
|
||||
tool_choice=tool_choice,
|
||||
tools=tools,
|
||||
response_format=response_format,
|
||||
user=user,
|
||||
stop=stop,
|
||||
frequency_penalty=frequency_penalty,
|
||||
logit_bias=logit_bias,
|
||||
presence_penalty=presence_penalty,
|
||||
seed=seed,
|
||||
store=store,
|
||||
metadata=metadata,
|
||||
additional_properties=additional_properties or {},
|
||||
**kwargs,
|
||||
)
|
||||
if isinstance(messages, str):
|
||||
messages = [ChatMessage(role="user", text=messages)]
|
||||
if isinstance(messages, ChatMessage):
|
||||
messages = [messages]
|
||||
_prepare_tools_and_tool_choice(chat_options=chat_options)
|
||||
async for update in self._inner_get_streaming_response(messages=messages, chat_options=chat_options, **kwargs):
|
||||
yield update
|
||||
|
||||
|
||||
# region: Embedding Client
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class EmbeddingGenerator(Protocol, Generic[TInput, TEmbedding]):
|
||||
"""A protocol for an embedding generator that can create embeddings from input data."""
|
||||
|
||||
async def generate(
|
||||
self,
|
||||
input_data: Sequence[TInput],
|
||||
**kwargs: Any,
|
||||
) -> GeneratedEmbeddings[TEmbedding]:
|
||||
"""Generates an embedding for the given input data.
|
||||
|
||||
Args:
|
||||
input_data: The input data to generate an embedding for.
|
||||
**kwargs: Additional options for the request.
|
||||
|
||||
Returns:
|
||||
The generated embedding, this acts like a list, but has additional metadata and usage details.
|
||||
|
||||
"""
|
||||
...
|
||||
@@ -0,0 +1,24 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import logging
|
||||
|
||||
from .exceptions import AgentFrameworkException
|
||||
|
||||
logging.basicConfig(
|
||||
format="[%(asctime)s - %(pathname)s:%(lineno)d - %(levelname)s] %(message)s",
|
||||
datefmt="%Y-%m-%d %H:%M:%S",
|
||||
)
|
||||
|
||||
|
||||
def get_logger(name: str = "agent_framework") -> logging.Logger:
|
||||
"""Get a logger with the specified name, defaulting to 'agent_framework'.
|
||||
|
||||
Args:
|
||||
name (str): The name of the logger. Defaults to 'agent_framework'.
|
||||
|
||||
Returns:
|
||||
logging.Logger: The configured logger instance.
|
||||
"""
|
||||
if not name.startswith("agent_framework"):
|
||||
raise AgentFrameworkException("Logger name must start with 'agent_framework'.")
|
||||
return logging.getLogger(name)
|
||||
@@ -0,0 +1,10 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
|
||||
from pydantic import BaseModel, ConfigDict
|
||||
|
||||
|
||||
class AFBaseModel(BaseModel):
|
||||
"""Base class for all pydantic models in the Agent Framework."""
|
||||
|
||||
model_config = ConfigDict(populate_by_name=True, arbitrary_types_allowed=True, validate_assignment=True)
|
||||
@@ -0,0 +1,139 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import functools
|
||||
import inspect
|
||||
from collections.abc import Awaitable, Callable, Mapping
|
||||
from typing import Any, Generic, Protocol, TypeVar, runtime_checkable
|
||||
|
||||
from pydantic import BaseModel, create_model
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class AITool(Protocol):
|
||||
"""Represents a tool that can be specified to an AI service."""
|
||||
|
||||
name: str
|
||||
"""The name of the tool."""
|
||||
description: str | None = None
|
||||
"""A description of the tool, suitable for use in describing the purpose to a model."""
|
||||
additional_properties: dict[str, Any] | None = None
|
||||
"""Additional properties associated with the tool."""
|
||||
|
||||
def __str__(self) -> str:
|
||||
"""Return a string representation of the tool."""
|
||||
...
|
||||
|
||||
def parameters(self) -> Mapping[str, Any]:
|
||||
"""Return the parameters of the tool as a JSON schema."""
|
||||
...
|
||||
|
||||
|
||||
ArgsT = TypeVar("ArgsT", bound=BaseModel)
|
||||
ReturnT = TypeVar("ReturnT")
|
||||
|
||||
|
||||
class AIFunction(AITool, Generic[ArgsT, ReturnT]):
|
||||
"""A tool that represents a function that can be called by an AI service."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
func: Callable[..., Awaitable[ReturnT] | ReturnT],
|
||||
name: str,
|
||||
description: str,
|
||||
input_model: type[ArgsT],
|
||||
**kwargs: Any,
|
||||
):
|
||||
"""Initialize a FunctionTool.
|
||||
|
||||
Args:
|
||||
func: The function to wrap.
|
||||
name: The name of the tool.
|
||||
description: A description of the tool.
|
||||
input_model: A Pydantic model that defines the input parameters for the function.
|
||||
**kwargs: Additional properties to set on the tool.
|
||||
stored in additional_properties.
|
||||
"""
|
||||
self.name = name
|
||||
self.description = description
|
||||
self.input_model = input_model
|
||||
self.additional_properties: dict[str, Any] | None = kwargs
|
||||
self._func = func
|
||||
|
||||
def parameters(self) -> dict[str, Any]:
|
||||
"""Return the parameter json schemas of the input model."""
|
||||
return self.input_model.model_json_schema()
|
||||
|
||||
def __call__(self, *args: Any, **kwargs: Any) -> ReturnT | Awaitable[ReturnT]:
|
||||
"""Call the wrapped function with the provided arguments."""
|
||||
return self._func(*args, **kwargs)
|
||||
|
||||
def __str__(self) -> str:
|
||||
return f"AIFunction(name={self.name}, description={self.description})"
|
||||
|
||||
async def invoke(
|
||||
self,
|
||||
*,
|
||||
arguments: ArgsT | None = None,
|
||||
**kwargs: Any,
|
||||
) -> ReturnT:
|
||||
"""Run the AI function with the provided arguments as a Pydantic model.
|
||||
|
||||
Args:
|
||||
arguments: A Pydantic model instance containing the arguments for the function.
|
||||
kwargs: keyword arguments to pass to the function, will not be used if `args` is provided.
|
||||
"""
|
||||
if arguments is not None:
|
||||
if not isinstance(arguments, self.input_model):
|
||||
raise TypeError(f"Expected {self.input_model.__name__}, got {type(arguments).__name__}")
|
||||
kwargs = arguments.model_dump(exclude_none=True)
|
||||
res = self.__call__(**kwargs)
|
||||
if inspect.isawaitable(res):
|
||||
return await res
|
||||
return res
|
||||
|
||||
|
||||
def ai_function(
|
||||
func: Callable[..., ReturnT | Awaitable[ReturnT]] | None = None,
|
||||
*,
|
||||
name: str | None = None,
|
||||
description: str | None = None,
|
||||
additional_properties: dict[str, Any] | None = None,
|
||||
) -> AIFunction[Any, ReturnT] | Callable[[Callable[..., ReturnT | Awaitable[ReturnT]]], AIFunction[Any, ReturnT]]:
|
||||
"""Decorate a function to turn it into a AIFunction that can be passed to models.
|
||||
|
||||
Args:
|
||||
func: The function to wrap. If None, returns a decorator.
|
||||
name: The name of the tool. Defaults to the function's name.
|
||||
description: A description of the tool. Defaults to the function's docstring.
|
||||
additional_properties: Additional properties to set on the tool.
|
||||
|
||||
"""
|
||||
|
||||
def wrapper(f: Callable[..., ReturnT | Awaitable[ReturnT]]) -> AIFunction[Any, ReturnT]:
|
||||
tool_name: str = name or getattr(f, "__name__", "unknown_function") # type: ignore[assignment]
|
||||
tool_desc: str = description or (f.__doc__ or "")
|
||||
sig = inspect.signature(f)
|
||||
fields = {
|
||||
pname: (
|
||||
param.annotation if param.annotation is not inspect.Parameter.empty else str,
|
||||
param.default if param.default is not inspect.Parameter.empty else ...,
|
||||
)
|
||||
for pname, param in sig.parameters.items()
|
||||
if pname not in {"self", "cls"}
|
||||
}
|
||||
input_model: Any = create_model(f"{tool_name}_input", **fields) # type: ignore[call-overload]
|
||||
if not issubclass(input_model, BaseModel):
|
||||
raise TypeError(f"Input model for {tool_name} must be a subclass of BaseModel, got {input_model}")
|
||||
|
||||
return functools.update_wrapper( # type: ignore[return-value]
|
||||
AIFunction[Any, ReturnT](
|
||||
func=f,
|
||||
name=tool_name,
|
||||
description=tool_desc,
|
||||
input_model=input_model,
|
||||
**(additional_properties if additional_properties is not None else {}),
|
||||
),
|
||||
f,
|
||||
)
|
||||
|
||||
return wrapper(func) if func else wrapper
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,19 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
|
||||
class AgentFrameworkException(Exception):
|
||||
"""Base class for exceptions in the Agent Framework."""
|
||||
|
||||
pass
|
||||
|
||||
|
||||
class AgentException(AgentFrameworkException):
|
||||
"""Base class for all agent exceptions."""
|
||||
|
||||
pass
|
||||
|
||||
|
||||
class AgentExecutionException(AgentException):
|
||||
"""An error occurred while executing the agent."""
|
||||
|
||||
pass
|
||||
@@ -0,0 +1,25 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
|
||||
from typing import Generic, Protocol, TypeVar, runtime_checkable
|
||||
|
||||
TInput = TypeVar("TInput")
|
||||
TResponse = TypeVar("TResponse")
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class InputGuardrail(Protocol, Generic[TInput]):
|
||||
"""A protocol for input guardrails that can validate and transform input messages."""
|
||||
|
||||
def __call__(self, message: TInput) -> TInput:
|
||||
"""Validate and possibly transform the input message."""
|
||||
...
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class OutputGuardrail(Protocol, Generic[TResponse]):
|
||||
"""A protocol for output guardrails that can validate and transform output messages."""
|
||||
|
||||
def __call__(self, message: TResponse) -> TResponse:
|
||||
"""Validate and possibly transform the output message."""
|
||||
...
|
||||
Reference in New Issue
Block a user