chore: import upstream snapshot with attribution

python/samples/02-agents/chat_client/custom_chat_client.py

@@ -0,0 +1,213 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import asyncio
+import random
+import sys
+from collections.abc import AsyncIterable, Awaitable, Mapping, Sequence
+from typing import Any, ClassVar, TypeAlias, TypedDict
+
+from agent_framework import (
+    Agent,
+    BaseChatClient,
+    ChatMiddlewareLayer,
+    ChatResponse,
+    ChatResponseUpdate,
+    Content,
+    FunctionInvocationLayer,
+    InMemoryHistoryProvider,
+    Message,
+    ResponseStream,
+)
+from agent_framework.observability import ChatTelemetryLayer
+
+if sys.version_info >= (3, 12):
+    from typing import override  # type: ignore # pragma: no cover
+else:
+    from typing_extensions import override  # type: ignore[import] # pragma: no cover
+
+
+"""
+Custom Chat Client Implementation Example
+
+This sample demonstrates implementing a custom chat client and optionally composing
+middleware, telemetry, and function invocation layers explicitly. The recommended
+layer order is `FunctionInvocationLayer -> ChatMiddlewareLayer -> ChatTelemetryLayer`
+so chat middleware runs within each tool-loop iteration while telemetry records
+per-call spans without middleware latency.
+"""
+
+
+class EchoingChatClientOptions(TypedDict, total=False):
+    """Custom options for EchoingChatClient."""
+
+    uppercase: bool
+    suffix: str
+    stream_delay_seconds: float
+
+
+OptionsT: TypeAlias = EchoingChatClientOptions
+
+
+class EchoingChatClient(BaseChatClient[OptionsT]):
+    """A custom chat client that echoes messages back with modifications.
+
+    This demonstrates how to implement a custom chat client by extending BaseChatClient
+    and implementing the required _inner_get_response() method.
+    """
+
+    OTEL_PROVIDER_NAME: ClassVar[str] = "EchoingChatClient"
+
+    def __init__(self, *, prefix: str = "Echo:", **kwargs: Any) -> None:
+        """Initialize the EchoingChatClient.
+
+        Args:
+            prefix: Prefix to add to echoed messages.
+            **kwargs: Additional keyword arguments passed to BaseChatClient.
+        """
+        super().__init__(**kwargs)
+        self.prefix = prefix
+
+    @override
+    def _inner_get_response(
+        self,
+        *,
+        messages: Sequence[Message],
+        stream: bool = False,
+        options: Mapping[str, Any],
+        **kwargs: Any,
+    ) -> Awaitable[ChatResponse] | ResponseStream[ChatResponseUpdate, ChatResponse]:
+        """Echo back the user's message with a prefix."""
+        if not messages:
+            response_text = "No messages to echo!"
+        else:
+            # Echo the last user message
+            last_user_message = None
+            for message in reversed(messages):
+                if message.role == "user":
+                    last_user_message = message
+                    break
+
+            if last_user_message and last_user_message.text:
+                response_text = f"{self.prefix} {last_user_message.text}"
+            else:
+                response_text = f"{self.prefix} [No text message found]"
+
+        if options.get("uppercase"):
+            response_text = response_text.upper()
+        if suffix := options.get("suffix"):
+            response_text = f"{response_text} {suffix}"
+        stream_delay_seconds = float(options.get("stream_delay_seconds", 0.05))
+
+        response_message = Message(role="assistant", contents=[response_text])
+
+        response = ChatResponse(
+            messages=[response_message],
+            model="echo-model-v1",
+            response_id=f"echo-resp-{random.randint(1000, 9999)}",
+        )
+
+        if not stream:
+
+            async def _get_response() -> ChatResponse:
+                return response
+
+            return _get_response()
+
+        async def _stream() -> AsyncIterable[ChatResponseUpdate]:
+            response_text_local = response_message.text or ""
+            for char in response_text_local:
+                yield ChatResponseUpdate(
+                    contents=[Content.from_text(char)],
+                    role="assistant",
+                    response_id=f"echo-stream-resp-{random.randint(1000, 9999)}",
+                    model="echo-model-v1",
+                )
+                await asyncio.sleep(stream_delay_seconds)
+
+        return ResponseStream(_stream(), finalizer=lambda updates: response)
+
+
+class EchoingChatClientWithLayers(  # type: ignore[misc]
+    FunctionInvocationLayer[OptionsT],
+    ChatMiddlewareLayer[OptionsT],
+    ChatTelemetryLayer[OptionsT],
+    EchoingChatClient,
+):
+    """Echoing chat client that explicitly composes middleware, telemetry, and function layers."""
+
+    OTEL_PROVIDER_NAME: ClassVar[str] = "EchoingChatClientWithLayers"
+
+
+async def main() -> None:
+    """Demonstrates how to implement and use a custom chat client with Agent."""
+    print("=== Custom Chat Client Example ===\n")
+
+    # Create the custom chat client
+    print("--- EchoingChatClient Example ---")
+
+    echo_client = EchoingChatClientWithLayers(prefix="🔊 Echo:")
+
+    # Use the chat client directly
+    print("Using chat client directly:")
+    direct_response = await echo_client.get_response(
+        [Message(role="user", contents=["Hello, custom chat client!"])],
+        options={
+            "uppercase": True,
+            "suffix": "(CUSTOM OPTIONS)",
+            "stream_delay_seconds": 0.02,
+        },
+    )
+    print(f"Direct response: {direct_response.messages[0].text}")
+
+    # Create an agent using the custom chat client
+    echo_agent = Agent(
+        client=echo_client,
+        name="EchoAgent",
+        instructions="You are a helpful assistant that echoes back what users say.",
+    )
+
+    print(f"\nAgent Name: {echo_agent.name}")
+
+    # Test non-streaming with agent
+    query = "This is a test message"
+    print(f"\nUser: {query}")
+    result = await echo_agent.run(query)
+    print(f"Agent: {result.messages[0].text}")
+
+    # Test streaming with agent
+    query2 = "Stream this message back to me"
+    print(f"\nUser: {query2}")
+    print("Agent: ", end="", flush=True)
+    async for chunk in echo_agent.run(query2, stream=True):
+        if chunk.text:
+            print(chunk.text, end="", flush=True)
+    print()
+
+    # Example: Using with sessions and conversation history
+    print("\n--- Using Custom Chat Client with Session ---")
+
+    session = echo_agent.create_session()
+
+    # Multiple messages in conversation
+    messages = [
+        "Hello, I'm starting a conversation",
+        "How are you doing?",
+        "Thanks for chatting!",
+    ]
+
+    for msg in messages:
+        result = await echo_agent.run(msg, session=session)
+        print(f"User: {msg}")
+        print(f"Agent: {result.messages[0].text}\n")
+
+    # Check conversation history
+    memory_state = session.state.get(InMemoryHistoryProvider.DEFAULT_SOURCE_ID, {})
+    session_messages = memory_state.get("messages", [])
+    if session_messages:
+        print(f"Session contains {len(session_messages)} messages")
+    else:
+        print("Session has no messages stored")
+
+
+if __name__ == "__main__":
+    asyncio.run(main())