Improving streaming entity serialization

dev/microsoft-agents-testing/microsoft_agents/testing/core/utils.py

@@ -90,4 +90,4 @@ def generate_token_from_config(sdk_config: dict, connection_name: str = "SERVICE
 
     if not client_id or not client_secret or not tenant_id:
         raise ValueError("Incorrect configuration provided for token generation.")
-    return generate_token(client_id, client_secret, tenant_id)
+    return generate_token(client_id, client_secret, tenant_id)

dev/microsoft-agents-testing/microsoft_agents/testing/utils.py

@@ -102,4 +102,4 @@ def resolve_scenario(scenario_or_str: Scenario | str ) -> Scenario:
         else:
             return scenario_registry.get(scenario_or_str)
     else:
-        raise TypeError("Input must be a Scenario instance or a string key.")
+        raise TypeError("Input must be a Scenario instance or a string key.")

dev/tests/sdk/test_streaming_response.py

@@ -0,0 +1,122 @@
+import pytest
+import asyncio
+
+from microsoft_agents.activity import (
+    Activity,
+    ActivityTypes,
+    Channels,
+    Entity
+)
+
+from microsoft_agents.hosting.core import (
+    TurnContext,
+    TurnState,
+)
+
+from microsoft_agents.testing import (
+    AgentClient,
+    AgentEnvironment,
+    AiohttpScenario,
+)
+
+FULL_TEXT = "This is a streaming response."
+CHUNKS = FULL_TEXT.split()
+
+def get_streaminfo(activity: Activity) -> Entity:
+    for entity in activity.entities:
+        if isinstance(entity, dict) and entity.get("type") == "streaminfo":
+            return Entity.model_validate(entity)
+        elif isinstance(entity, Entity) and entity.type == "streaminfo":
+            return entity
+    raise ValueError("No streaminfo entity found")
+
+async def init_agent(env: AgentEnvironment):
+
+    app = env.agent_application
+
+    @app.message("/stream")
+    async def stream_handler(context: TurnContext, state: TurnState):
+
+        assert context.streaming_response is not None
+
+        context.streaming_response.queue_informative_update("Starting stream...")
+        await asyncio.sleep(1.0)  # Simulate delay before starting stream
+
+        for chunk in CHUNKS[:-1]:
+            context.streaming_response.queue_text_chunk(chunk)
+            await asyncio.sleep(1.0)  # Simulate delay between chunks
+
+        context.streaming_response.queue_text_chunk(CHUNKS[-1])
+        await context.streaming_response.end_stream()
+
+_SCENARIO = AiohttpScenario(init_agent=init_agent, use_jwt_middleware=False)
+
+@pytest.mark.asyncio
+@pytest.mark.agent_test(_SCENARIO)
+async def test_basic_streaming_response_non_streaming_channel(agent_client: AgentClient):
+
+    expected_len = len(FULL_TEXT.split())
+
+    agent_client.template = agent_client.template.with_updates(channel_id=Channels.emulator)
+
+    # give enough time for all the activities to send
+    await agent_client.send("/stream", wait=expected_len * 2.0)
+
+    stream_activities = agent_client.select().where(
+        entities=lambda x: any(e["type"] == "streaminfo" for e in x)
+    ).get()
+
+    assert len(stream_activities) == 1
+
+    final_streaminfo = get_streaminfo(stream_activities[0])
+
+    assert final_streaminfo.stream_sequence == 1
+    assert final_streaminfo.stream_type == "final"
+    assert stream_activities[0].text == FULL_TEXT.replace(" ", "")
+
+
+
+@pytest.mark.asyncio
+@pytest.mark.agent_test(_SCENARIO)
+async def test_basic_streaming_response_streaming_channel(agent_client: AgentClient):
+
+    expected_len = len(FULL_TEXT.split())
+
+    agent_client.template = agent_client.template.with_updates(channel_id=Channels.webchat)
+
+    # give enough time for all the activities to send
+    await agent_client.send("/stream", wait=expected_len * 2.0)
+
+    stream_activities = agent_client.select().where(
+        entities=lambda x: any(e["type"] == "streaminfo" for e in x)
+    ).get()
+
+    assert len(stream_activities) == len(CHUNKS) + 1
+
+    informative = stream_activities[0]
+    informative_streaminfo = get_streaminfo(informative)
+
+    assert informative_streaminfo.stream_type == "informative"
+    assert informative_streaminfo.stream_sequence == 1
+    assert informative.text == "Starting stream..."
+    assert informative.type == ActivityTypes.typing
+
+    t = ""
+    for i, chunk in enumerate(CHUNKS[:-1]):
+        t += chunk
+
+        j = i + 1
+
+        streaminfo = get_streaminfo(stream_activities[j])
+
+        assert stream_activities[j].text == t
+        assert stream_activities[j].type == ActivityTypes.typing
+        assert streaminfo.stream_type == "streaming"
+        assert streaminfo.stream_sequence == j + 1
+
+    final_streaminfo = get_streaminfo(stream_activities[-1])
+
+    assert final_streaminfo.stream_sequence == len(stream_activities)
+    assert final_streaminfo.stream_type == "final"
+    assert stream_activities[-1].text == FULL_TEXT.replace(" ", "")
+

libraries/microsoft-agents-activity/microsoft_agents/activity/__init__.py

@@ -44,6 +44,7 @@
     Place,
     ProductInfo,
     Thing,
+    StreamInfo,
 )
 from .error import Error
 from .error_response import ErrorResponse
@@ -133,6 +134,7 @@
     "ExpectedReplies",
     "Entity",
     "AIEntity",
+    "EntityTypes",
     "ClientCitation",
     "ClientCitationAppearance",
     "ClientCitationImage",
@@ -154,6 +156,7 @@
     "OAuthCard",
     "PagedMembersResult",
     "Place",
+    "StreamInfo",
     "ProductInfo",
     "ReceiptCard",
     "ReceiptItem",

libraries/microsoft-agents-activity/microsoft_agents/activity/activity.py

@@ -754,9 +754,6 @@ def add_ai_metadata(
         """
         if citations:
             ai_entity = AIEntity(
-                type="https://schema.org/Message",
-                schema_type="Message",
-                context="https://schema.org",
                 id="",
                 additional_type=["AIGeneratedContent"],
                 citation=citations,

libraries/microsoft-agents-activity/microsoft_agents/activity/agents_model.py

@@ -8,7 +8,9 @@
 
 
 class AgentsModel(BaseModel):
-    model_config = ConfigDict(alias_generator=to_camel, populate_by_name=True)
+    model_config = ConfigDict(
+        alias_generator=to_camel, validate_by_name=True, validate_by_alias=True
+    )
 
     """
     @model_serializer

libraries/microsoft-agents-activity/microsoft_agents/activity/entity/__init__.py

@@ -16,6 +16,7 @@
 from .geo_coordinates import GeoCoordinates
 from .place import Place
 from .product_info import ProductInfo
+from .stream_info import StreamInfo
 from .thing import Thing
 
 __all__ = [
@@ -29,6 +30,7 @@
     "Mention",
     "SensitivityUsageInfo",
     "SensitivityPattern",
+    "StreamInfo",
     "GeoCoordinates",
     "Place",
     "ProductInfo",

libraries/microsoft-agents-activity/microsoft_agents/activity/entity/_schema_mixin.py

@@ -0,0 +1,55 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+from typing import Any
+from pydantic import (
+    model_serializer,
+    model_validator,
+    ModelWrapValidatorHandler,
+    SerializerFunctionWrapHandler,
+    SerializationInfo,
+    BaseModel,
+)
+
+
+def validate_schema_model(data: Any, handler: ModelWrapValidatorHandler):
+    """Custom validator to handle the aliases @type, @context, and @id if defined in the destination type."""
+    model = handler(data)
+    if isinstance(data, dict):
+        if "@type" in data:
+            setattr(model, "at_type", data["@type"])
+        if "@context" in data:
+            setattr(model, "at_context", data["@context"])
+        if "@id" in data:
+            setattr(model, "at_id", data["@id"])
+    return model
+
+
+def serialize_schema_model(
+    self, handler: SerializerFunctionWrapHandler, info: SerializationInfo
+) -> dict[str, Any]:
+    """Custom serializer to convert keys to force inclusion of @type, @context, and @id if defined."""
+    serialized = handler(self)
+    if info.by_alias:
+        if hasattr(self, "at_type"):
+            serialized["@type"] = getattr(self, "at_type")
+        if hasattr(self, "at_context"):
+            serialized["@context"] = getattr(self, "at_context")
+        if hasattr(self, "at_id"):
+            serialized["@id"] = getattr(self, "at_id")
+    return serialized
+
+
+class _SchemaMixin(BaseModel):
+    """Mixin class to force inclusion of @property fields when serializing."""
+
+    @model_validator(mode="wrap")
+    @classmethod
+    def _validate_model(cls, data: Any, handler: ModelWrapValidatorHandler):
+        return validate_schema_model(data, handler)
+
+    @model_serializer(mode="wrap")
+    def _serialize_model(
+        self, handler: SerializerFunctionWrapHandler, info: SerializationInfo
+    ) -> dict[str, Any]:
+        return serialize_schema_model(self, handler, info)

libraries/microsoft-agents-activity/microsoft_agents/activity/entity/ai_entity.py

@@ -2,10 +2,11 @@
 # Licensed under the MIT License.
 
 from enum import Enum
-from typing import List, Optional, Union, Literal
-from dataclasses import dataclass
+from typing import List, Optional, Literal
 
+from pydantic import Field
 from ..agents_model import AgentsModel
+from ._schema_mixin import _SchemaMixin
 from .entity import Entity
 
 
@@ -42,33 +43,36 @@ class ClientCitationImage(AgentsModel):
     name: str = ""
 
 
-class SensitivityPattern(AgentsModel):
+class SensitivityPattern(AgentsModel, _SchemaMixin):
     """Pattern information for sensitivity usage info."""
 
-    type: str = "DefinedTerm"
+    at_type: Literal["DefinedTerm"] = "DefinedTerm"
+
     in_defined_term_set: str = ""
     name: str = ""
     term_code: str = ""
 
 
-class SensitivityUsageInfo(AgentsModel):
+class SensitivityUsageInfo(AgentsModel, _SchemaMixin):
     """
     Sensitivity usage info for content sent to the user.
     This is used to provide information about the content to the user.
     """
 
     type: str = "https://schema.org/Message"
-    schema_type: str = "CreativeWork"
+    at_type: Literal["CreativeWork"] = "CreativeWork"
+
     description: Optional[str] = None
     name: str = ""
     position: Optional[int] = None
     pattern: Optional[SensitivityPattern] = None
 
 
-class ClientCitationAppearance(AgentsModel):
+class ClientCitationAppearance(AgentsModel, _SchemaMixin):
     """Appearance information for a client citation."""
 
-    type: str = "DigitalDocument"
+    at_type: Literal["DigitalDocument"] = "DigitalDocument"
+
     name: str = ""
     text: Optional[str] = None
     url: Optional[str] = None
@@ -79,33 +83,30 @@ class ClientCitationAppearance(AgentsModel):
     usage_info: Optional[SensitivityUsageInfo] = None
 
 
-class ClientCitation(AgentsModel):
+class ClientCitation(AgentsModel, _SchemaMixin):
     """
     Represents a Teams client citation to be included in a message.
     See Bot messages with AI-generated content for more details.
     https://learn.microsoft.com/en-us/microsoftteams/platform/bots/how-to/bot-messages-ai-generated-content?tabs=before%2Cbotmessage
     """
 
-    type: str = "Claim"
-    position: int = 0
-    appearance: Optional[ClientCitationAppearance] = None
+    at_type: Literal["Claim"] = "Claim"
 
-    def __post_init__(self):
-        if self.appearance is None:
-            self.appearance = ClientCitationAppearance()
+    position: int = 0
+    appearance: ClientCitationAppearance = Field(
+        default_factory=ClientCitationAppearance
+    )
 
 
 class AIEntity(Entity):
     """Entity indicating AI-generated content."""
 
+    at_type: Literal["Message"] = "Message"
+    at_context: Literal["https://schema.org"] = "https://schema.org"
+
     type: str = "https://schema.org/Message"
-    schema_type: str = "Message"
-    context: str = "https://schema.org"
     id: str = ""
-    additional_type: Optional[List[str]] = None
+
+    additional_type: List[str] = Field(default_factory=lambda: ["AIGeneratedContent"])
     citation: Optional[List[ClientCitation]] = None
     usage_info: Optional[SensitivityUsageInfo] = None
-
-    def __post_init__(self):
-        if self.additional_type is None:
-            self.additional_type = ["AIGeneratedContent"]