Python: latency improvements (#3014)

* latency improvements

* fixed mypy, added coding standards and instructions

* slight logic improvement

python/packages/azure-ai/pyproject.toml

@@ -83,6 +83,13 @@ include = "../../shared_tasks.toml"
 mypy = "mypy --config-file $POE_ROOT/pyproject.toml agent_framework_azure_ai"
 test = "pytest --cov=agent_framework_azure_ai --cov-report=term-missing:skip-covered tests"
 
+[tool.poe.tasks.integration-tests]
+cmd = """
+pytest --import-mode=importlib
+-n logical --dist loadfile --dist worksteal
+tests
+"""
+
 [build-system]
 requires = ["flit-core >= 3.11,<4.0"]
 build-backend = "flit_core.buildapi"

python/packages/core/agent_framework/_tools.py

@@ -573,7 +573,7 @@ class AIFunction(BaseTool, Generic[ArgsT, ReturnT]):
     """
 
     INJECTABLE: ClassVar[set[str]] = {"func"}
-    DEFAULT_EXCLUDE: ClassVar[set[str]] = {"input_model", "_invocation_duration_histogram"}
+    DEFAULT_EXCLUDE: ClassVar[set[str]] = {"input_model", "_invocation_duration_histogram", "_cached_parameters"}
 
     def __init__(
         self,
@@ -615,6 +615,7 @@ class AIFunction(BaseTool, Generic[ArgsT, ReturnT]):
         self.func = func
         self._instance = None  # Store the instance for bound methods
         self.input_model = self._resolve_input_model(input_model)
+        self._cached_parameters: dict[str, Any] | None = None  # Cache for model_json_schema()
         self.approval_mode = approval_mode or "never_require"
         if max_invocations is not None and max_invocations < 1:
             raise ValueError("max_invocations must be at least 1 or None.")
@@ -802,8 +803,11 @@ class AIFunction(BaseTool, Generic[ArgsT, ReturnT]):
 
         Returns:
             A dictionary containing the JSON schema for the function's parameters.
+            The result is cached after the first call for performance.
         """
-        return self.input_model.model_json_schema()
+        if self._cached_parameters is None:
+            self._cached_parameters = self.input_model.model_json_schema()
+        return self._cached_parameters
 
     def to_json_schema_spec(self) -> dict[str, Any]:
         """Convert a AIFunction to the JSON Schema function specification format.
@@ -825,7 +829,7 @@ class AIFunction(BaseTool, Generic[ArgsT, ReturnT]):
         as_dict = super().to_dict(exclude=exclude, exclude_none=exclude_none)
         if (exclude and "input_model" in exclude) or not self.input_model:
             return as_dict
-        as_dict["input_model"] = self.input_model.model_json_schema()
+        as_dict["input_model"] = self.parameters()  # Use cached parameters()
         return as_dict
 
 

python/packages/core/agent_framework/_types.py

@@ -101,7 +101,7 @@ def _parse_content(content_data: MutableMapping[str, Any]) -> "Contents":
     Raises:
         ContentError if parsing fails
     """
-    content_type = str(content_data.get("type"))
+    content_type: str | None = content_data.get("type", None)
     match content_type:
         case "text":
             return TextContent.from_dict(content_data)
@@ -127,6 +127,8 @@ def _parse_content(content_data: MutableMapping[str, Any]) -> "Contents":
             return FunctionApprovalResponseContent.from_dict(content_data)
         case "text_reasoning":
             return TextReasoningContent.from_dict(content_data)
+        case None:
+            raise ContentError("Content type is missing")
         case _:
             raise ContentError(f"Unknown content type '{content_type}'")
 
@@ -2248,27 +2250,30 @@ def _process_update(
     if update.message_id:
         message.message_id = update.message_id
     for content in update.contents:
-        if (
-            isinstance(content, FunctionCallContent)
-            and len(message.contents) > 0
-            and isinstance(message.contents[-1], FunctionCallContent)
-        ):
+        # Fast path: get type attribute (most content will have it)
+        content_type = getattr(content, "type", None)
+        # Slow path: only check for dict if type is None
+        if content_type is None and isinstance(content, (dict, MutableMapping)):
             try:
-                message.contents[-1] += content
-            except AdditionItemMismatch:
-                message.contents.append(content)
-        elif isinstance(content, UsageContent):
-            if response.usage_details is None:
-                response.usage_details = UsageDetails()
-            response.usage_details += content.details
-        elif isinstance(content, (dict, MutableMapping)):
-            try:
-                cont = _parse_content(content)
-                message.contents.append(cont)
+                content = _parse_content(content)
+                content_type = content.type
             except ContentError as exc:
                 logger.warning(f"Skipping unknown content type or invalid content: {exc}")
-        else:
-            message.contents.append(content)
+                continue
+        match content_type:
+            # mypy doesn't narrow type based on match/case, but we know these are FunctionCallContents
+            case "function_call" if message.contents and message.contents[-1].type == "function_call":
+                try:
+                    message.contents[-1] += content  # type: ignore[operator]
+                except AdditionItemMismatch:
+                    message.contents.append(content)
+            case "usage":
+                if response.usage_details is None:
+                    response.usage_details = UsageDetails()
+                # mypy doesn't narrow type based on match/case, but we know this is UsageContent
+                response.usage_details += content.details  # type: ignore[union-attr, arg-type]
+            case _:
+                message.contents.append(content)
     # Incorporate the update's properties into the response.
     if update.response_id:
         response.response_id = update.response_id

python/packages/core/agent_framework/observability.py

@@ -1680,13 +1680,12 @@ def _capture_messages(
     prepped = prepare_messages(messages, system_instructions=system_instructions)
     otel_messages: list[dict[str, Any]] = []
     for index, message in enumerate(prepped):
-        otel_messages.append(_to_otel_message(message))
-        try:
-            message_data = message.to_dict(exclude_none=True)
-        except Exception:
-            message_data = {"role": message.role.value, "contents": message.contents}
+        # Reuse the otel message representation for logging instead of calling to_dict()
+        # to avoid expensive Pydantic serialization overhead
+        otel_message = _to_otel_message(message)
+        otel_messages.append(otel_message)
         logger.info(
-            message_data,
+            otel_message,
             extra={
                 OtelAttr.EVENT_NAME: OtelAttr.CHOICE if output else ROLE_EVENT_MAP.get(message.role.value),
                 OtelAttr.PROVIDER_NAME: provider_name,