Python: Implemented FoundryChatClient (#193)

* Initial version of FoundryChatClient

* Updates to the tool call streaming wrapper

* Small fixes

* Small updates and addressed PR feedback

* Handle automatic client creation

* Small improvement

* Added credential parameter

* Small improvements

* Made FoundryChatClient disposable

* Small fixes

* Added unit tests

* Refactored samples

* Small improvements

* Small fix

* Addressed PR feedback

* Small fixes

* Small updates

* Small fix

* Addressed PR feedback

python/packages/main/agent_framework/_agents.py

@@ -1,10 +1,16 @@
 # Copyright (c) Microsoft. All rights reserved.
 
+import sys
 from collections.abc import AsyncIterable, Callable, MutableMapping, Sequence
 from enum import Enum
 from typing import Any, Literal, Protocol, TypeVar, runtime_checkable
 from uuid import uuid4
 
+if sys.version_info >= (3, 11):
+    from typing import Self  # pragma: no cover
+else:
+    from typing_extensions import Self  # pragma: no cover
+
 from pydantic import BaseModel, Field
 
 from ._clients import ChatClient
@@ -359,6 +365,23 @@ class ChatClientAgent(AgentBase):
 
         super().__init__(**args)
 
+    async def __aenter__(self) -> "Self":
+        """Async context manager entry.
+
+        If the chat_client supports async context management, enter its context.
+        """
+        if hasattr(self.chat_client, "__aenter__") and hasattr(self.chat_client, "__aexit__"):
+            await self.chat_client.__aenter__()  # type: ignore[reportUnknownMemberType]
+        return self
+
+    async def __aexit__(self, exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: Any) -> None:
+        """Async context manager exit.
+
+        If the chat_client supports async context management, exit its context.
+        """
+        if hasattr(self.chat_client, "__aenter__") and hasattr(self.chat_client, "__aexit__"):
+            await self.chat_client.__aexit__(exc_type, exc_val, exc_tb)  # type: ignore[reportUnknownMemberType]
+
     async def run(
         self,
         messages: str | ChatMessage | list[str] | list[ChatMessage] | None = None,

python/packages/main/agent_framework/_clients.py

@@ -4,9 +4,9 @@ import asyncio
 from abc import ABC, abstractmethod
 from collections.abc import AsyncIterable, Awaitable, Callable, MutableMapping, MutableSequence, Sequence
 from functools import wraps
-from typing import Annotated, Any, Generic, Literal, Protocol, TypeVar, runtime_checkable
+from typing import Any, Generic, Literal, Protocol, TypeVar, runtime_checkable
 
-from pydantic import BaseModel, StringConstraints
+from pydantic import BaseModel
 
 from ._logging import get_logger
 from ._pydantic import AFBaseModel
@@ -75,7 +75,7 @@ async def _auto_invoke_function(
     )
 
 
-def _tool_to_json_schema_spec(tool: AITool) -> dict[str, Any]:
+def tool_to_json_schema_spec(tool: AITool) -> dict[str, Any]:
     """Convert a AITool to the JSON Schema function specification format."""
     return {
         "type": "function",
@@ -95,7 +95,7 @@ def _prepare_tools_and_tool_choice(chat_options: ChatOptions) -> None:
         chat_options.tool_choice = ChatToolMode.NONE.mode
         return
     chat_options.tools = [
-        (_tool_to_json_schema_spec(t) if isinstance(t, AITool) else t)
+        (tool_to_json_schema_spec(t) if isinstance(t, AITool) else t)
         for t in chat_options._ai_tools or []  # type: ignore[reportPrivateUsage]
     ]
     if not chat_options.tools:
@@ -205,6 +205,12 @@ def _tool_call_streaming(func: TInnerGetStreamingResponse) -> TInnerGetStreaming
             messages.append(response.messages[0])
             function_calls = [item for item in response.messages[0].contents if isinstance(item, FunctionCallContent)]
 
+            # When conversation id is present, it means that messages are hosted on the server.
+            # In this case, we need to update ChatOptions with conversation id and also clear messages
+            if response.conversation_id is not None:
+                chat_options.conversation_id = response.conversation_id
+                messages = []
+
             if function_calls:
                 # Run all function calls concurrently
                 results = await asyncio.gather(*[
@@ -393,8 +399,6 @@ class ChatClient(Protocol):
 class ChatClientBase(AFBaseModel, ABC):
     """Base class for chat clients."""
 
-    ai_model_id: Annotated[str, StringConstraints(strip_whitespace=True, min_length=1)]
-
     def _prepare_messages(
         self, messages: str | ChatMessage | list[str] | list[ChatMessage]
     ) -> MutableSequence[ChatMessage]:

python/packages/main/agent_framework/_types.py

@@ -269,7 +269,7 @@ def _coalesce_text_content(
                 first_new_content = i
         else:
             if first_new_content is not None:
-                new_content = type_(text=" ".join(current_texts))
+                new_content = type_(text="".join(current_texts))
                 new_content.raw_representation = contents[first_new_content].raw_representation
                 new_content.additional_properties = contents[first_new_content].additional_properties
                 # Store the replacement node. We inherit the properties of the first text node. We don't
@@ -280,7 +280,7 @@ def _coalesce_text_content(
                 first_new_content = None
             coalesced_contents.append(content)
     if current_texts:
-        coalesced_contents.append(type_(text=" ".join(current_texts)))
+        coalesced_contents.append(type_(text="".join(current_texts)))
     contents.clear()
     contents.extend(coalesced_contents)
 
@@ -398,6 +398,7 @@ class DataContent(AIContent):
         uri: The URI of the data represented by this instance, typically in the form of a data URI.
             Should be in the form: "data:{media_type};base64,{base64_data}".
         type: The type of content, which is always "data" for this class.
+        media_type: The media type of the data.
         additional_properties: Optional additional properties associated with the content.
         raw_representation: Optional raw representation of the content.
 
@@ -405,6 +406,7 @@ class DataContent(AIContent):
 
     type: Literal["data"] = "data"  # type: ignore[assignment]
     uri: str
+    media_type: str | None = None
 
     @overload
     def __init__(
@@ -486,6 +488,7 @@ class DataContent(AIContent):
             uri = f"data:{media_type};base64,{base64.b64encode(data).decode('utf-8')}"
         super().__init__(
             uri=uri,  # type: ignore[reportCallIssue]
+            media_type=media_type,  # type: ignore[reportCallIssue]
             raw_representation=raw_representation,
             additional_properties=additional_properties,
             **kwargs,
@@ -506,6 +509,9 @@ class DataContent(AIContent):
             raise ValueError(f"Unknown media type: {media_type}")
         return uri
 
+    def has_top_level_media_type(self, top_level_media_type: str) -> bool:
+        return _has_top_level_media_type(self.media_type, top_level_media_type)
+
 
 class UriContent(AIContent):
     """Represents a URI content.
@@ -559,6 +565,19 @@ class UriContent(AIContent):
             **kwargs,
         )
 
+    def has_top_level_media_type(self, top_level_media_type: str) -> bool:
+        return _has_top_level_media_type(self.media_type, top_level_media_type)
+
+
+def _has_top_level_media_type(media_type: str | None, top_level_media_type: str) -> bool:
+    if media_type is None:
+        return False
+
+    slash_index = media_type.find("/")
+    span = media_type[:slash_index] if slash_index >= 0 else media_type
+    span = span.strip()
+    return span.lower() == top_level_media_type.lower()
+
 
 class ErrorContent(AIContent):
     """Represents an error.
@@ -1449,7 +1468,12 @@ ChatToolMode.NONE = ChatToolMode(mode="none")  # type: ignore[assignment]
 class ChatOptions(AFBaseModel):
     """Common request settings for AI services."""
 
+    additional_properties: MutableMapping[str, Any] = Field(
+        default_factory=dict, description="Provider-specific additional properties."
+    )
     ai_model_id: Annotated[str | None, Field(serialization_alias="model")] = None
+    allow_multiple_tool_calls: bool | None = None
+    conversation_id: str | None = None
     frequency_penalty: Annotated[float | None, Field(ge=-2.0, le=2.0)] = None
     logit_bias: MutableMapping[str | int, float] | None = None
     max_tokens: Annotated[int | None, Field(gt=0)] = None
@@ -1464,12 +1488,9 @@ class ChatOptions(AFBaseModel):
     temperature: Annotated[float | None, Field(ge=0.0, le=2.0)] = None
     tool_choice: ChatToolMode | Literal["auto", "required", "none"] | Mapping[str, Any] | None = None
     tools: list[AITool | MutableMapping[str, Any]] | None = None
-    _ai_tools: list[AITool | MutableMapping[str, Any]] | None = PrivateAttr(default=None)
     top_p: Annotated[float | None, Field(ge=0.0, le=1.0)] = None
     user: str | None = None
-    additional_properties: MutableMapping[str, Any] = Field(
-        default_factory=dict, description="Provider-specific additional properties."
-    )
+    _ai_tools: list[AITool | MutableMapping[str, Any]] | None = PrivateAttr(default=None)
 
     @model_validator(mode="after")
     def _copy_to_ai_tools(self) -> Self:

python/packages/main/agent_framework/foundry/__init__.py

@@ -0,0 +1,24 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import importlib
+from typing import Any
+
+PACKAGE_NAME = "agent_framework_foundry"
+PACKAGE_EXTRA = "foundry"
+_IMPORTS = ["__version__", "FoundryChatClient"]
+
+
+def __getattr__(name: str) -> Any:
+    if name in _IMPORTS:
+        try:
+            return getattr(importlib.import_module(PACKAGE_NAME), name)
+        except ModuleNotFoundError as exc:
+            raise ModuleNotFoundError(
+                f"The '{PACKAGE_EXTRA}' extra is not installed, "
+                f"please do `pip install agent-framework[{PACKAGE_EXTRA}]`"
+            ) from exc
+    raise AttributeError(f"Module {PACKAGE_NAME} has no attribute {name}.")
+
+
+def __dir__() -> list[str]:
+    return _IMPORTS

python/packages/main/agent_framework/foundry/__init__.pyi

@@ -0,0 +1,5 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from agent_framework_foundry import FoundryChatClient, __version__
+
+__all__ = ["FoundryChatClient", "__version__"]