Python: Fix Python pyright package scoping and typing remediation (#4426)

* Fix Python pyright package scoping and typing remediation

Implements issue #4407 by removing the root pyright include, adding package-level pyright includes, and resolving pyright/mypy typing issues across Python packages. Also cleans unnecessary casts and applies line-level, rule-specific ignores where external libraries are too dynamic.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Reduce pyright cost in handoff cloning

Simplify cloned_options construction in HandoffAgentExecutor to avoid expensive TypedDict narrowing/inference in _handoff.py, which was causing pyright to spend a long time in orchestrations.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* fix types

* Fix lint and type-check regressions

Resolve current Python package check failures across lint, pyright, and mypy after recent code changes, including purview/declarative pyright issues and multiple ruff simplification findings.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* fixed hooks

* Stabilize package tests and test tasks

Resolve cross-package non-integration test failures, simplify streaming type flow, harden locale/culture handling, and standardize package test poe tasks to exclude integration tests where applicable.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* lots of small fixes

* Fix current Python test regressions

Address current failing unit tests in azure-ai, bedrock, and azure-cosmos while keeping Bedrock parsing logic inline (no new static helper methods).

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* small fixes

* small fixes

* removed pydantic from json

* final updates

* fix core

* fix tests

* fix obser

---------

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

python/packages/core/agent_framework/__init__.py

@@ -205,9 +205,6 @@ __all__ = [
     "AgentResponseUpdate",
     "AgentRunInputs",
     "AgentSession",
-    "Skill",
-    "SkillResource",
-    "SkillsProvider",
     "Annotation",
     "BaseAgent",
     "BaseChatClient",
@@ -272,6 +269,9 @@ __all__ = [
     "SecretString",
     "SessionContext",
     "SingleEdgeGroup",
+    "Skill",
+    "SkillResource",
+    "SkillsProvider",
     "SubWorkflowRequestMessage",
     "SubWorkflowResponseMessage",
     "SupportsAgentRun",

python/packages/core/agent_framework/_agents.py

@@ -83,10 +83,13 @@ OptionsCoT = TypeVar(
 
 def _get_tool_name(tool: Any) -> str | None:
     """Extract a tool's name from either an object with a .name attribute or a dict tool definition."""
-    if isinstance(tool, dict):
-        func = tool.get("function")
-        if isinstance(func, dict):
-            return func.get("name")
+    if isinstance(tool, Mapping):
+        tool_mapping = cast(Mapping[str, Any], tool)
+        func = tool_mapping.get("function")
+        if isinstance(func, Mapping):
+            func_mapping = cast(Mapping[str, Any], func)
+            name = func_mapping.get("name")
+            return name if isinstance(name, str) else None
         return None
     return getattr(tool, "name", None)
 
@@ -164,12 +167,12 @@ def _sanitize_agent_name(agent_name: str | None) -> str | None:
 class _RunContext(TypedDict):
     session: AgentSession | None
     session_context: SessionContext
-    input_messages: list[Message]
-    session_messages: list[Message]
+    input_messages: Sequence[Message]
+    session_messages: Sequence[Message]
     agent_name: str
-    chat_options: dict[str, Any]
-    filtered_kwargs: dict[str, Any]
-    finalize_kwargs: dict[str, Any]
+    chat_options: MutableMapping[str, Any]
+    filtered_kwargs: Mapping[str, Any]
+    finalize_kwargs: Mapping[str, Any]
 
 
 # region Agent Protocol
@@ -770,10 +773,9 @@ class RawAgent(BaseAgent, Generic[OptionsCoT]):  # type: ignore[misc]
         should check if there is already an agent name defined, and if not
         set it to this value.
         """
-        if hasattr(self.client, "_update_agent_name_and_description") and callable(
-            self.client._update_agent_name_and_description
-        ):  # type: ignore[reportAttributeAccessIssue, attr-defined]
-            self.client._update_agent_name_and_description(self.name, self.description)  # type: ignore[reportAttributeAccessIssue, attr-defined]
+        update_fn = getattr(self.client, "_update_agent_name_and_description", None)
+        if callable(update_fn):
+            update_fn(self.name, self.description)
 
     @overload
     def run(
@@ -860,11 +862,14 @@ class RawAgent(BaseAgent, Generic[OptionsCoT]):  # type: ignore[misc]
                     options=options,
                     kwargs=kwargs,
                 )
-                response = await self.client.get_response(  # type: ignore[call-overload]
-                    messages=ctx["session_messages"],
-                    stream=False,
-                    options=ctx["chat_options"],
-                    **ctx["filtered_kwargs"],
+                response = cast(
+                    ChatResponse[Any],
+                    await self.client.get_response(  # type: ignore
+                        messages=ctx["session_messages"],
+                        stream=False,
+                        options=ctx["chat_options"],  # type: ignore[reportArgumentType]
+                        **ctx["filtered_kwargs"],
+                    ),
                 )
 
                 if not response:
@@ -930,7 +935,7 @@ class RawAgent(BaseAgent, Generic[OptionsCoT]):  # type: ignore[misc]
             )
             await self._run_after_providers(session=ctx["session"], context=session_context)
 
-        async def _get_stream() -> ResponseStream[ChatResponseUpdate, ChatResponse]:
+        async def _get_stream() -> ResponseStream[ChatResponseUpdate, ChatResponse[Any]]:
             ctx_holder["ctx"] = await self._prepare_run_context(
                 messages=messages,
                 session=session,
@@ -942,7 +947,7 @@ class RawAgent(BaseAgent, Generic[OptionsCoT]):  # type: ignore[misc]
             return self.client.get_response(  # type: ignore[call-overload, no-any-return]
                 messages=ctx["session_messages"],
                 stream=True,
-                options=ctx["chat_options"],
+                options=ctx["chat_options"],  # type: ignore[reportArgumentType]
                 **ctx["filtered_kwargs"],
             )
 
@@ -965,12 +970,12 @@ class RawAgent(BaseAgent, Generic[OptionsCoT]):  # type: ignore[misc]
             rf = (
                 ctx.get("chat_options", {}).get("response_format")
                 if ctx
-                else (options.get("response_format") if options else None)
+                else (options.get("response_format") if options else None)  # type: ignore[union-attr]
             )
             return self._finalize_response_updates(updates, response_format=rf)
 
         return (
-            ResponseStream
+            ResponseStream  # type: ignore[reportUnknownMemberType]
             .from_awaitable(_get_stream())
             .map(
                 transform=partial(
@@ -988,10 +993,13 @@ class RawAgent(BaseAgent, Generic[OptionsCoT]):  # type: ignore[misc]
         updates: Sequence[AgentResponseUpdate],
         *,
         response_format: Any | None = None,
-    ) -> AgentResponse:
+    ) -> AgentResponse[Any]:
         """Finalize response updates into a single AgentResponse."""
         output_format_type = response_format if isinstance(response_format, type) else None
-        return AgentResponse.from_updates(updates, output_format_type=output_format_type)
+        return AgentResponse.from_updates(  # pyright: ignore[reportUnknownVariableType]
+            updates,
+            output_format_type=output_format_type,
+        )
 
     @staticmethod
     def _extract_conversation_id_from_streaming_response(response: AgentResponse[Any]) -> str | None:
@@ -1000,10 +1008,11 @@ class RawAgent(BaseAgent, Generic[OptionsCoT]):  # type: ignore[misc]
         if raw is None:
             return None
 
-        raw_items: list[Any] = raw if isinstance(raw, list) else [raw]
+        raw_items: list[Any] = list(cast(Any, raw)) if isinstance(raw, list) else [raw]
         for item in reversed(raw_items):
             if isinstance(item, Mapping):
-                value = item.get("conversation_id")
+                mapped_item = cast(Mapping[str, Any], item)
+                value = mapped_item.get("conversation_id")
                 if isinstance(value, str) and value:
                     return value
                 continue
@@ -1074,7 +1083,7 @@ class RawAgent(BaseAgent, Generic[OptionsCoT]):  # type: ignore[misc]
 
         # Merge runtime kwargs into additional_function_arguments so they're available
         # in function middleware context and tool invocation.
-        existing_additional_args = opts.pop("additional_function_arguments", None) or {}
+        existing_additional_args: dict[str, Any] = opts.pop("additional_function_arguments", None) or {}
         additional_function_arguments = {**kwargs, **existing_additional_args}
         # Include session so as_tool() wrappers with propagate_session=True can access it.
         if active_session is not None:

python/packages/core/agent_framework/_clients.py

@@ -317,10 +317,13 @@ class BaseChatClient(SerializationMixin, ABC, Generic[OptionsCoT]):
         updates: Sequence[ChatResponseUpdate],
         *,
         response_format: Any | None = None,
-    ) -> ChatResponse:
+    ) -> ChatResponse[Any]:
         """Finalize response updates into a single ChatResponse."""
         output_format_type = response_format if isinstance(response_format, type) else None
-        return ChatResponse.from_updates(updates, output_format_type=output_format_type)
+        return ChatResponse.from_updates(  # pyright: ignore[reportUnknownVariableType]
+            updates,
+            output_format_type=output_format_type,
+        )
 
     def _build_response_stream(
         self,
@@ -782,7 +785,7 @@ class BaseEmbeddingClient(SerializationMixin, ABC, Generic[EmbeddingInputT, Embe
         values: Sequence[EmbeddingInputT],
         *,
         options: EmbeddingOptionsT | None = None,
-    ) -> GeneratedEmbeddings[EmbeddingT]:
+    ) -> GeneratedEmbeddings[EmbeddingT, EmbeddingOptionsT]:
         """Generate embeddings for the given values.
 
         Args:

python/packages/core/agent_framework/_middleware.py

@@ -8,7 +8,7 @@ import sys
 from abc import ABC, abstractmethod
 from collections.abc import AsyncIterable, Awaitable, Callable, Mapping, Sequence
 from enum import Enum
-from typing import TYPE_CHECKING, Any, Generic, Literal, TypeAlias, overload
+from typing import TYPE_CHECKING, Any, Generic, Literal, TypeAlias, cast, overload
 
 from ._clients import SupportsChatGetResponse
 from ._types import (
@@ -170,9 +170,9 @@ class AgentContext:
         self.session = session
         self.options = options
         self.stream = stream
-        self.metadata = metadata if metadata is not None else {}
+        self.metadata: dict[str, Any] = dict(metadata) if metadata is not None else {}
         self.result = result
-        self.kwargs = kwargs if kwargs is not None else {}
+        self.kwargs: dict[str, Any] = dict(kwargs) if kwargs is not None else {}
         self.stream_transform_hooks = list(stream_transform_hooks or [])
         self.stream_result_hooks = list(stream_result_hooks or [])
         self.stream_cleanup_hooks = list(stream_cleanup_hooks or [])
@@ -231,9 +231,9 @@ class FunctionInvocationContext:
         """
         self.function = function
         self.arguments = arguments
-        self.metadata = metadata if metadata is not None else {}
+        self.metadata: dict[str, Any] = dict(metadata) if metadata is not None else {}
         self.result = result
-        self.kwargs = kwargs if kwargs is not None else {}
+        self.kwargs: dict[str, Any] = dict(kwargs) if kwargs is not None else {}
 
 
 class ChatContext:
@@ -314,9 +314,9 @@ class ChatContext:
         self.messages = messages
         self.options = options
         self.stream = stream
-        self.metadata = metadata if metadata is not None else {}
+        self.metadata: dict[str, Any] = dict(metadata) if metadata is not None else {}
         self.result = result
-        self.kwargs = kwargs if kwargs is not None else {}
+        self.kwargs: dict[str, Any] = dict(kwargs) if kwargs is not None else {}
         self.stream_transform_hooks = list(stream_transform_hooks or [])
         self.stream_result_hooks = list(stream_result_hooks or [])
         self.stream_cleanup_hooks = list(stream_cleanup_hooks or [])
@@ -754,9 +754,11 @@ class AgentMiddlewarePipeline(BaseMiddlewarePipeline):
             if index >= len(self._middleware):
 
                 async def final_wrapper() -> None:
-                    context.result = final_handler(context)  # type: ignore[assignment]
-                    if inspect.isawaitable(context.result):
-                        context.result = await context.result
+                    result = final_handler(context)
+                    if inspect.isawaitable(result):
+                        context.result = await cast(Awaitable[AgentResponse], result)
+                    else:
+                        context.result = result
 
                 return final_wrapper
 
@@ -893,12 +895,17 @@ class ChatMiddlewarePipeline(BaseMiddlewarePipeline):
             The chat response after processing through all middleware.
         """
         if not self._middleware:
-            context.result = final_handler(context)  # type: ignore[assignment]
-            if isinstance(context.result, Awaitable):
-                context.result = await context.result
-            if context.stream and not isinstance(context.result, ResponseStream):
+            result = final_handler(context)
+            if inspect.isawaitable(result):
+                resolved_result: ChatResponse | ResponseStream[ChatResponseUpdate, ChatResponse] = await cast(
+                    Awaitable[ChatResponse], result
+                )
+            else:
+                resolved_result = result
+            context.result = resolved_result
+            if context.stream and not isinstance(resolved_result, ResponseStream):
                 raise ValueError("Streaming agent middleware requires a ResponseStream result.")
-            return context.result
+            return resolved_result
 
         def create_next_handler(index: int) -> Callable[[], Awaitable[None]]:
             if index >= len(self._middleware):
@@ -1038,7 +1045,10 @@ class ChatMiddlewareLayer(Generic[OptionsCoT]):
                 # If result is ChatResponse (shouldn't happen for streaming), raise error
                 raise ValueError("Expected ResponseStream for streaming, got ChatResponse")
 
-            return ResponseStream.from_awaitable(_execute_stream())
+            return cast(
+                ResponseStream[ChatResponseUpdate, ChatResponse[Any]],
+                cast(Any, ResponseStream).from_awaitable(_execute_stream()),
+            )
 
         # For non-streaming, return the coroutine directly
         return _execute()  # type: ignore[return-value]
@@ -1120,7 +1130,10 @@ class AgentMiddlewareLayer:
     ) -> Awaitable[AgentResponse[Any]] | ResponseStream[AgentResponseUpdate, AgentResponse[Any]]:
         """MiddlewareTypes-enabled unified run method."""
         # Re-categorize self.middleware at runtime to support dynamic changes
-        base_middleware = getattr(self, "middleware", None) or []
+        base_middleware_attr = getattr(self, "middleware", None)
+        base_middleware: Sequence[MiddlewareTypes] = (
+            cast(Sequence[MiddlewareTypes], base_middleware_attr) if isinstance(base_middleware_attr, Sequence) else []
+        )
         base_middleware_list = categorize_middleware(base_middleware)
         run_middleware_list = categorize_middleware(middleware)
         pipeline = AgentMiddlewarePipeline(*base_middleware_list["agent"], *run_middleware_list["agent"])
@@ -1166,7 +1179,10 @@ class AgentMiddlewareLayer:
                 # If result is AgentResponse (shouldn't happen for streaming), convert to stream
                 raise ValueError("Expected ResponseStream for streaming, got AgentResponse")
 
-            return ResponseStream.from_awaitable(_execute_stream())
+            return cast(
+                ResponseStream[AgentResponseUpdate, AgentResponse[Any]],
+                cast(Any, ResponseStream).from_awaitable(_execute_stream()),
+            )
 
         # For non-streaming, return the coroutine directly
         return _execute()  # type: ignore[return-value]

python/packages/core/agent_framework/_serialization.py

@@ -303,7 +303,7 @@ class SerializationMixin:
                 # Handle lists containing SerializationProtocol objects
                 if isinstance(value, list):
                     value_as_list: list[Any] = []
-                    for item in value:
+                    for item in value:  # pyright: ignore[reportUnknownVariableType]
                         if isinstance(item, SerializationProtocol):
                             value_as_list.append(item.to_dict(exclude=exclude, exclude_none=exclude_none))
                             continue
@@ -311,7 +311,7 @@ class SerializationMixin:
                             value_as_list.append(item)
                             continue
                         logger.debug(
-                            f"Skipping non-serializable item in list attribute '{key}' of type {type(item).__name__}"
+                            f"Skipping non-serializable item in list attribute '{key}' of type {type(item).__name__}"  # pyright: ignore[reportUnknownArgumentType]
                         )
                     result[key] = value_as_list
                     continue
@@ -320,21 +320,22 @@ class SerializationMixin:
                     from datetime import date, datetime, time
 
                     serialized_dict: dict[str, Any] = {}
-                    for k, v in value.items():
+                    for raw_key, v in value.items():  # pyright: ignore[reportUnknownVariableType]
+                        dict_key = str(raw_key)  # pyright: ignore[reportUnknownArgumentType]
                         if isinstance(v, SerializationProtocol):
-                            serialized_dict[k] = v.to_dict(exclude=exclude, exclude_none=exclude_none)
+                            serialized_dict[dict_key] = v.to_dict(exclude=exclude, exclude_none=exclude_none)
                             continue
                         # Convert datetime objects to strings
                         if isinstance(v, (datetime, date, time)):
-                            serialized_dict[k] = str(v)
+                            serialized_dict[dict_key] = str(v)
                             continue
                         # Check if the value is JSON serializable
                         if is_serializable(v):
-                            serialized_dict[k] = v
+                            serialized_dict[dict_key] = v
                             continue
                         logger.debug(
-                            f"Skipping non-serializable value for key '{k}' in dict attribute '{key}' "
-                            f"of type {type(v).__name__}"
+                            f"Skipping non-serializable value for key '{dict_key}' in dict attribute '{key}' "
+                            f"of type {type(v).__name__}"  # pyright: ignore[reportUnknownArgumentType]
                         )
                     result[key] = serialized_dict
                     continue
@@ -505,7 +506,8 @@ class SerializationMixin:
                 # Only apply if the instance matches
                 if kwargs.get(field) == name and isinstance(dep_value, dict):
                     # Apply instance-specific dependencies
-                    for param_name, param_value in dep_value.items():
+                    for raw_param_name, param_value in dep_value.items():  # pyright: ignore[reportUnknownVariableType]
+                        param_name = str(raw_param_name)  # pyright: ignore[reportUnknownArgumentType]
                         if param_name not in cls.INJECTABLE:
                             logger.debug(
                                 f"Dependency '{param_name}' for type '{type_id}' is not in INJECTABLE set. "

python/packages/core/agent_framework/_sessions.py

@@ -16,7 +16,7 @@ import copy
 import uuid
 from abc import abstractmethod
 from collections.abc import Sequence
-from typing import TYPE_CHECKING, Any, ClassVar
+from typing import TYPE_CHECKING, Any, ClassVar, cast
 
 from ._types import AgentResponse, Message
 
@@ -92,7 +92,7 @@ def _deserialize_value(value: Any) -> Any:
                 from pydantic import BaseModel
 
                 if issubclass(cls, BaseModel):
-                    data = {k: v for k, v in value.items() if k != "type"}
+                    data: dict[str, Any] = {str(k): v for k, v in value.items() if k != "type"}  # pyright: ignore[reportUnknownVariableType, reportUnknownArgumentType]
                     return cls.model_validate(data)
             except ImportError:
                 pass
@@ -229,8 +229,11 @@ class SessionContext:
             tools: The tools to add.
         """
         for tool in tools:
-            if hasattr(tool, "additional_properties") and isinstance(tool.additional_properties, dict):
-                tool.additional_properties["context_source"] = source_id
+            if hasattr(tool, "additional_properties"):
+                additional_properties_obj = tool.additional_properties
+                if isinstance(additional_properties_obj, dict):
+                    additional_properties = cast(dict[str, Any], additional_properties_obj)
+                    additional_properties["context_source"] = source_id
         self.tools.extend(tools)
 
     def get_messages(

python/packages/core/agent_framework/_settings.py

@@ -215,9 +215,7 @@ def load_settings(
             raise FileNotFoundError(env_file_path)
 
         raw_dotenv_values = dotenv_values(dotenv_path=env_file_path, encoding=encoding)
-        loaded_dotenv_values = {
-            key: value for key, value in raw_dotenv_values.items() if key is not None and value is not None
-        }
+        loaded_dotenv_values = {key: value for key, value in raw_dotenv_values.items() if value is not None}
 
     # Filter out None overrides so defaults / env vars are preserved
     overrides = {k: v for k, v in overrides.items() if v is not None}

python/packages/core/agent_framework/_skills.py

@@ -151,6 +151,7 @@ class Skill:
                 content="Use this skill for DB tasks.",
             )
 
+
             @skill.resource
             def get_schema() -> str:
                 return "CREATE TABLE ..."
@@ -972,9 +973,7 @@ def _load_skills(
 
     if skills:
         for code_skill in skills:
-            error = _validate_skill_metadata(
-                code_skill.name, code_skill.description, "code skill"
-            )
+            error = _validate_skill_metadata(code_skill.name, code_skill.description, "code skill")
             if error:
                 logger.warning(error)
                 continue

python/packages/core/agent_framework/_tools.py

@@ -27,7 +27,7 @@ from typing import (
     Literal,
     TypeAlias,
     TypedDict,
-    Union,
+    cast,
     get_args,
     get_origin,
     overload,
@@ -77,6 +77,7 @@ else:
 
 logger = logging.getLogger("agent_framework")
 
+
 DEFAULT_MAX_ITERATIONS: Final[int] = 40
 DEFAULT_MAX_CONSECUTIVE_ERRORS_PER_REQUEST: Final[int] = 3
 SHELL_TOOL_KIND_VALUE: Final[str] = "shell"
@@ -84,7 +85,7 @@ ChatClientT = TypeVar("ChatClientT", bound="SupportsChatGetResponse[Any]")
 # region Helpers
 
 
-def _parse_inputs(
+def _parse_inputs(  # pyright: ignore[reportUnusedFunction]
     inputs: Content | dict[str, Any] | str | list[Content | dict[str, Any] | str] | None,
 ) -> list[Content]:
     """Parse the inputs for a tool, ensuring they are of type Content.
@@ -352,7 +353,8 @@ class FunctionTool(SerializationMixin):
     def declaration_only(self) -> bool:
         """Indicate whether the function is declaration only (i.e., has no implementation)."""
         # Check for explicit _declaration_only attribute first (used in tests)
-        if hasattr(self, "_declaration_only") and self._declaration_only:
+        declaration_flag = getattr(self, "_declaration_only", False)
+        if isinstance(declaration_flag, bool) and declaration_flag:
             return True
         return self.func is None
 
@@ -430,10 +432,13 @@ class FunctionTool(SerializationMixin):
             )
         self.invocation_count += 1
         try:
+            func = self.func
+            if func is None:
+                raise ToolException(f"Function '{self.name}' has no implementation.")
             # If we have a bound instance, call the function with self
             if self._instance is not None:
-                return self.func(self._instance, *args, **kwargs)
-            return self.func(*args, **kwargs)  # type:ignore[misc]
+                return func(self._instance, *args, **kwargs)
+            return func(*args, **kwargs)
         except Exception:
             self.invocation_exception_count += 1
             raise
@@ -600,9 +605,11 @@ class FunctionTool(SerializationMixin):
         from ._types import Content
 
         if isinstance(value, list):
-            return [FunctionTool._make_dumpable(item) for item in value]
+            list_value = cast(list[object], value)
+            return [FunctionTool._make_dumpable(item) for item in list_value]
         if isinstance(value, dict):
-            return {k: FunctionTool._make_dumpable(v) for k, v in value.items()}
+            dict_value = cast(dict[object, object], value)
+            return {key: FunctionTool._make_dumpable(item) for key, item in dict_value.items()}
         if isinstance(value, Content):
             return value.to_dict(exclude={"raw_representation", "additional_properties"})
         if isinstance(value, BaseModel):
@@ -661,7 +668,7 @@ class FunctionTool(SerializationMixin):
         return as_dict
 
 
-ToolTypes: TypeAlias = FunctionTool | MCPTool | Mapping[str, Any] | Any
+ToolTypes: TypeAlias = FunctionTool | MCPTool | Mapping[str, Any] | object
 
 
 def normalize_tools(
@@ -679,27 +686,31 @@ def normalize_tools(
     if not tools:
         return []
 
-    tool_items = (
-        list(tools)
-        if isinstance(tools, Sequence) and not isinstance(tools, (str, bytes, bytearray, Mapping))
-        else [tools]
-    )
+    if isinstance(tools, (str, bytes, bytearray, Mapping)) or not isinstance(tools, Sequence):
+        tools = cast(list[ToolTypes | Callable[..., Any]], [tools])
+
     from ._mcp import MCPTool
 
     normalized: list[ToolTypes] = []
-    for tool_item in tool_items:
+    for tool_item in tools:  # type: ignore[reportUnknownVariableType]
         # check known types, these are also callable, so we need to do that first
-        if isinstance(tool_item, (FunctionTool, Mapping, MCPTool)):
+        if isinstance(tool_item, FunctionTool):
             normalized.append(tool_item)
             continue
-        if callable(tool_item):
+        if isinstance(tool_item, dict):
+            normalized.append(tool_item)  # type: ignore[reportUnknownArgumentType]
+            continue
+        if isinstance(tool_item, MCPTool):
+            normalized.append(tool_item)
+            continue
+        if callable(tool_item):  # type: ignore[reportUnknownArgumentType]
             normalized.append(tool(tool_item))
             continue
-        normalized.append(tool_item)
+        normalized.append(tool_item)  # type: ignore[reportUnknownArgumentType]
     return normalized
 
 
-def _tools_to_dict(
+def _tools_to_dict(  # pyright: ignore[reportUnusedFunction]
     tools: ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None,
 ) -> list[str | dict[str, Any]] | None:
     """Parse the tools to a dict.
@@ -722,8 +733,8 @@ def _tools_to_dict(
         if isinstance(tool_item, SerializationMixin):
             results.append(tool_item.to_dict())
             continue
-        if isinstance(tool_item, Mapping):
-            results.append(dict(tool_item))
+        if isinstance(tool_item, dict):
+            results.append(tool_item)  # type: ignore[reportUnknownArgumentType]
             continue
         logger.warning("Can't parse tool.")
     return results
@@ -795,32 +806,28 @@ def _validate_arguments_against_schema(
     """Run lightweight argument checks for schema-supplied tools."""
     parsed_arguments = dict(arguments)
 
-    required_raw = schema.get("required", [])
-    required_fields = [field for field in required_raw if isinstance(field, str)]
+    required_fields = [field for field in schema.get("required", []) if isinstance(field, str)]
     missing_fields = [field for field in required_fields if field not in parsed_arguments]
     if missing_fields:
         raise TypeError(f"Missing required argument(s) for '{tool_name}': {', '.join(sorted(missing_fields))}")
 
-    properties_raw = schema.get("properties")
-    properties = properties_raw if isinstance(properties_raw, Mapping) else {}
-
+    properties: Mapping[str, Any] = schema.get("properties", {})
     if schema.get("additionalProperties") is False:
         unexpected_fields = sorted(field for field in parsed_arguments if field not in properties)
         if unexpected_fields:
             raise TypeError(f"Unexpected argument(s) for '{tool_name}': {', '.join(unexpected_fields)}")
 
     for field_name, field_value in parsed_arguments.items():
-        field_schema = properties.get(field_name)
-        if not isinstance(field_schema, Mapping):
+        if not isinstance(properties.get(field_name), dict):
             continue
 
-        enum_values = field_schema.get("enum")
+        enum_values = properties.get(field_name, {}).get("enum")  # type: ignore
         if isinstance(enum_values, list) and enum_values and field_value not in enum_values:
             raise TypeError(
                 f"Invalid value for '{field_name}' in '{tool_name}': {field_value!r} is not in {enum_values!r}"
             )
 
-        schema_type = field_schema.get("type")
+        schema_type = properties.get(field_name, {}).get("type")  # type: ignore
         if isinstance(schema_type, str):
             if not _matches_json_schema_type(field_value, schema_type):
                 raise TypeError(
@@ -830,7 +837,7 @@ def _validate_arguments_against_schema(
             continue
 
         if isinstance(schema_type, list):
-            allowed_types = [item for item in schema_type if isinstance(item, str)]
+            allowed_types: list[str] = [item for item in schema_type if isinstance(item, str)]  # type: ignore[reportUnknownVariableType]
             if allowed_types and not any(_matches_json_schema_type(field_value, item) for item in allowed_types):
                 raise TypeError(
                     f"Invalid type for '{field_name}' in '{tool_name}': expected one of "
@@ -840,240 +847,6 @@ def _validate_arguments_against_schema(
     return parsed_arguments
 
 
-# Map JSON Schema types to Pydantic types
-TYPE_MAPPING = {
-    "string": str,
-    "integer": int,
-    "number": float,
-    "boolean": bool,
-    "array": list,
-    "object": dict,
-    "null": type(None),
-}
-
-
-def _build_pydantic_model_from_json_schema(
-    model_name: str,
-    schema: Mapping[str, Any],
-) -> type[BaseModel]:
-    """Creates a Pydantic model from JSON Schema with support for $refs, nested objects, and typed arrays.
-
-    Args:
-        model_name: The name of the model to be created.
-        schema: The JSON Schema definition (should contain 'properties', 'required', '$defs', etc.).
-
-    Returns:
-        The dynamically created Pydantic model class.
-    """
-    properties = schema.get("properties")
-    required = schema.get("required", [])
-    definitions = schema.get("$defs", {})
-
-    # Check if 'properties' is missing or not a dictionary
-    if not properties:
-        return create_model(f"{model_name}_input")
-
-    def _resolve_literal_type(prop_details: dict[str, Any]) -> type | None:
-        """Check if property should be a Literal type (const or enum).
-
-        Args:
-            prop_details: The JSON Schema property details
-
-        Returns:
-            Literal type if const or enum is present, None otherwise
-        """
-        # const → Literal["value"]
-        if "const" in prop_details:
-            return Literal[prop_details["const"]]  # type: ignore
-
-        # enum → Literal["a", "b", ...]
-        if "enum" in prop_details and isinstance(prop_details["enum"], list):
-            enum_values = prop_details["enum"]
-            if enum_values:
-                return Literal[tuple(enum_values)]  # type: ignore
-
-        return None
-
-    def _resolve_type(prop_details: dict[str, Any], parent_name: str = "") -> type:
-        """Resolve JSON Schema type to Python type, handling $ref, nested objects, and typed arrays.
-
-        Args:
-            prop_details: The JSON Schema property details
-            parent_name: Name to use for creating nested models (for uniqueness)
-
-        Returns:
-            Python type annotation (could be int, str, list[str], or a nested Pydantic model)
-        """
-        # Handle oneOf + discriminator (polymorphic objects)
-        if "oneOf" in prop_details and "discriminator" in prop_details:
-            discriminator = prop_details["discriminator"]
-            disc_field = discriminator.get("propertyName")
-
-            variants = []
-            for variant in prop_details["oneOf"]:
-                if "$ref" in variant:
-                    ref = variant["$ref"]
-                    if ref.startswith("#/$defs/"):
-                        def_name = ref.split("/")[-1]
-                        resolved = definitions.get(def_name)
-                        if resolved:
-                            variant_model = _resolve_type(
-                                resolved,
-                                parent_name=f"{parent_name}_{def_name}",
-                            )
-                            variants.append(variant_model)
-
-            if variants and disc_field:
-                return Annotated[
-                    Union[tuple(variants)],  # type: ignore
-                    Field(discriminator=disc_field),
-                ]
-
-        # Handle $ref by resolving the reference
-        if "$ref" in prop_details:
-            ref = prop_details["$ref"]
-            # Extract the reference path (e.g., "#/$defs/CustomerIdParam" -> "CustomerIdParam")
-            if ref.startswith("#/$defs/"):
-                def_name = ref.split("/")[-1]
-                if def_name in definitions:
-                    # Resolve the reference and use its type
-                    resolved = definitions[def_name]
-                    return _resolve_type(resolved, def_name)
-            # If we can't resolve the ref, default to dict for safety
-            return dict
-
-        # Map JSON Schema types to Python types
-        json_type = prop_details.get("type", "string")
-        match json_type:
-            case "integer":
-                return int
-            case "number":
-                return float
-            case "boolean":
-                return bool
-            case "array":
-                # Handle typed arrays
-                items_schema = prop_details.get("items")
-                if items_schema and isinstance(items_schema, dict):
-                    # Recursively resolve the item type
-                    item_type = _resolve_type(items_schema, f"{parent_name}_item")
-                    # Return list[ItemType] instead of bare list
-                    return list[item_type]  # type: ignore
-                # If no items schema or invalid, return bare list
-                return list
-            case "object":
-                # Handle nested objects by creating a nested Pydantic model
-                nested_properties = prop_details.get("properties")
-                nested_required = prop_details.get("required", [])
-
-                if nested_properties and isinstance(nested_properties, dict):
-                    # Create the name for the nested model
-                    nested_model_name = f"{parent_name}_nested" if parent_name else "NestedModel"
-
-                    # Recursively build field definitions for the nested model
-                    nested_field_definitions: dict[str, Any] = {}
-                    for nested_prop_name, nested_prop_details in nested_properties.items():
-                        nested_prop_details = (
-                            json.loads(nested_prop_details)
-                            if isinstance(nested_prop_details, str)
-                            else nested_prop_details
-                        )
-
-                        # Check for Literal types first (const/enum)
-                        literal_type = _resolve_literal_type(nested_prop_details)
-                        if literal_type is not None:
-                            nested_python_type = literal_type
-                        else:
-                            nested_python_type = _resolve_type(
-                                nested_prop_details,
-                                f"{nested_model_name}_{nested_prop_name}",
-                            )
-                        nested_description = nested_prop_details.get("description", "")
-
-                        # Build field kwargs for nested property
-                        nested_field_kwargs: dict[str, Any] = {}
-                        if nested_description:
-                            nested_field_kwargs["description"] = nested_description
-
-                        # Create field definition
-                        if nested_prop_name in nested_required:
-                            nested_field_definitions[nested_prop_name] = (
-                                (
-                                    nested_python_type,
-                                    Field(**nested_field_kwargs),
-                                )
-                                if nested_field_kwargs
-                                else (nested_python_type, ...)
-                            )
-                        else:
-                            nested_field_kwargs["default"] = nested_prop_details.get("default", None)
-                            nested_field_definitions[nested_prop_name] = (
-                                nested_python_type,
-                                Field(**nested_field_kwargs),
-                            )
-
-                    # Create and return the nested Pydantic model
-                    return create_model(nested_model_name, **nested_field_definitions)  # type: ignore
-
-                # If no properties defined, return bare dict
-                return dict
-            case _:
-                return str  # default
-
-    field_definitions: dict[str, Any] = {}
-    for prop_name, prop_details in properties.items():
-        prop_details = json.loads(prop_details) if isinstance(prop_details, str) else prop_details
-
-        # Check for Literal types first (const/enum)
-        literal_type = _resolve_literal_type(prop_details)
-        if literal_type is not None:
-            python_type = literal_type
-        else:
-            python_type = _resolve_type(prop_details, f"{model_name}_{prop_name}")
-        description = prop_details.get("description", "")
-
-        # Build field kwargs (description, etc.)
-        field_kwargs: dict[str, Any] = {}
-        if description:
-            field_kwargs["description"] = description
-
-        # Create field definition for create_model
-        if prop_name in required:
-            if field_kwargs:
-                field_definitions[prop_name] = (python_type, Field(**field_kwargs))
-            else:
-                field_definitions[prop_name] = (python_type, ...)
-        else:
-            default_value = prop_details.get("default", None)
-            field_kwargs["default"] = default_value
-            if field_kwargs and any(k != "default" for k in field_kwargs):
-                field_definitions[prop_name] = (python_type, Field(**field_kwargs))
-            else:
-                field_definitions[prop_name] = (python_type, default_value)
-
-    return create_model(f"{model_name}_input", **field_definitions)
-
-
-def _create_model_from_json_schema(tool_name: str, schema_json: Mapping[str, Any]) -> type[BaseModel]:
-    """Creates a Pydantic model from a given JSON Schema.
-
-    Args:
-      tool_name: The name of the model to be created.
-      schema_json: The JSON Schema definition.
-
-    Returns:
-      The dynamically created Pydantic model class.
-    """
-    # Validate that 'properties' exists and is a dict
-    if "properties" not in schema_json or not isinstance(schema_json["properties"], dict):
-        raise ValueError(
-            f"JSON schema for tool '{tool_name}' must contain a 'properties' key of type dict. "
-            f"Got: {schema_json.get('properties', None)}"
-        )
-
-    return _build_pydantic_model_from_json_schema(tool_name, schema_json)
-
-
 @overload
 def tool(
     func: Callable[..., Any],
@@ -1348,8 +1121,6 @@ def normalize_function_invocation_configuration(
         raise ValueError("max_function_calls must be at least 1 or None.")
     if normalized["max_consecutive_errors_per_request"] < 0:
         raise ValueError("max_consecutive_errors_per_request must be 0 or more.")
-    if normalized["additional_tools"] is None:
-        normalized["additional_tools"] = []
     return normalized
 
 
@@ -1424,7 +1195,7 @@ async def _auto_invoke_function(
         if key not in {"_function_middleware_pipeline", "middleware", "conversation_id"}
     }
     try:
-        if not tool._schema_supplied and tool.input_model is not None:
+        if not cast(bool, getattr(tool, "_schema_supplied", False)) and tool.input_model is not None:
             args = tool.input_model.model_validate(parsed_args).model_dump(exclude_none=True)
         else:
             args = dict(parsed_args)
@@ -1435,7 +1206,7 @@ async def _auto_invoke_function(
         )
     except (TypeError, ValidationError) as exc:
         message = "Error: Argument parsing failed."
-        if config["include_detailed_errors"]:
+        if config.get("include_detailed_errors", False):
             message = f"{message} Exception: {exc}"
         return Content.from_function_result(
             call_id=function_call_content.call_id,  # type: ignore[arg-type]
@@ -1459,7 +1230,7 @@ async def _auto_invoke_function(
             )
         except Exception as exc:
             message = "Error: Function failed."
-            if config["include_detailed_errors"]:
+            if config.get("include_detailed_errors", False):
                 message = f"{message} Exception: {exc}"
             return Content.from_function_result(
                 call_id=function_call_content.call_id,  # type: ignore[arg-type]
@@ -1505,7 +1276,7 @@ async def _auto_invoke_function(
         raise
     except Exception as exc:
         message = "Error: Function failed."
-        if config["include_detailed_errors"]:
+        if config.get("include_detailed_errors", False):
             message = f"{message} Exception: {exc}"
         return Content.from_function_result(
             call_id=function_call_content.call_id,  # type: ignore[arg-type]
@@ -1560,7 +1331,8 @@ async def _try_execute_function_calls(
         approval_tools,
     )
     declaration_only = [tool_name for tool_name, tool in tool_map.items() if tool.declaration_only]
-    additional_tool_names = [tool.name for tool in config["additional_tools"]] if config["additional_tools"] else []
+    configured_additional_tools = config.get("additional_tools") or []
+    additional_tool_names = [tool.name for tool in configured_additional_tools]
     # check if any are calling functions that need approval
     # if so, we return approval request for all
     approval_needed = False
@@ -1581,7 +1353,7 @@ async def _try_execute_function_calls(
             declaration_only_flag = True
             break
         if (
-            config["terminate_on_unknown_calls"] and fcc.type == "function_call" and fcc.name not in tool_map  # type: ignore[attr-defined]
+            config.get("terminate_on_unknown_calls", False) and fcc.type == "function_call" and fcc.name not in tool_map  # type: ignore[attr-defined]
         ):
             raise KeyError(f'Error: Requested function "{fcc.name}" not found.')  # type: ignore[attr-defined]
     if approval_needed:
@@ -1598,7 +1370,7 @@ async def _try_execute_function_calls(
     if declaration_only_flag:
         # return the declaration only tools to the user, since we cannot execute them.
         # Mark as user_input_request so AgentExecutor emits request_info events and pauses the workflow.
-        declaration_only_calls = []
+        declaration_only_calls: list[Content] = []
         for fcc in function_calls:
             if fcc.type == "function_call":
                 fcc.user_input_request = True
@@ -1695,19 +1467,6 @@ def _update_conversation_id(
         options["conversation_id"] = conversation_id
 
 
-async def _ensure_response_stream(
-    stream_like: ResponseStream[Any, Any] | Awaitable[ResponseStream[Any, Any]],
-) -> ResponseStream[Any, Any]:
-    from ._types import ResponseStream
-
-    stream = await stream_like if isinstance(stream_like, Awaitable) else stream_like
-    if not isinstance(stream, ResponseStream):
-        raise ValueError("Streaming function invocation requires a ResponseStream result.")
-    if getattr(stream, "_stream", None) is None:
-        await stream
-    return stream
-
-
 def _extract_tools(
     options: dict[str, Any] | None,
 ) -> ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None:
@@ -1776,7 +1535,7 @@ def _replace_approval_contents_with_results(
         }
 
         # Track approval requests that should be removed (duplicates)
-        contents_to_remove = []
+        contents_to_remove: list[int] = []
 
         for content_idx, content in enumerate(msg.contents):
             if content.type == "function_approval_request":
@@ -2097,7 +1856,9 @@ class FunctionInvocationLayer(Generic[OptionsCoT]):
         function_middleware_pipeline = FunctionMiddlewarePipeline(
             *(self.function_middleware), *(function_middleware or [])
         )
-        max_errors: int = self.function_invocation_configuration["max_consecutive_errors_per_request"]  # type: ignore[assignment]
+        max_errors = self.function_invocation_configuration.get(
+            "max_consecutive_errors_per_request", DEFAULT_MAX_CONSECUTIVE_ERRORS_PER_REQUEST
+        )
         additional_function_arguments: dict[str, Any] = {}
         if options and (additional_opts := options.get("additional_function_arguments")):  # type: ignore[attr-defined]
             additional_function_arguments = additional_opts  # type: ignore
@@ -2122,7 +1883,7 @@ class FunctionInvocationLayer(Generic[OptionsCoT]):
 
         if not stream:
 
-            async def _get_response() -> ChatResponse:
+            async def _get_response() -> ChatResponse[Any]:
                 nonlocal mutable_options
                 nonlocal filtered_kwargs
                 errors_in_a_row: int = 0
@@ -2130,13 +1891,11 @@ class FunctionInvocationLayer(Generic[OptionsCoT]):
                 max_function_calls: int | None = self.function_invocation_configuration.get("max_function_calls")
                 prepped_messages = list(messages)
                 fcc_messages: list[Message] = []
-                response: ChatResponse | None = None
+                response: ChatResponse[Any] | None = None
 
-                for attempt_idx in range(
-                    self.function_invocation_configuration["max_iterations"]
-                    if self.function_invocation_configuration["enabled"]
-                    else 0
-                ):
+                loop_enabled = self.function_invocation_configuration.get("enabled", True)
+                max_iterations = self.function_invocation_configuration.get("max_iterations", DEFAULT_MAX_ITERATIONS)
+                for attempt_idx in range(max_iterations if loop_enabled else 0):
                     approval_result = await _process_function_requests(
                         response=None,
                         prepped_messages=prepped_messages,
@@ -2147,17 +1906,20 @@ class FunctionInvocationLayer(Generic[OptionsCoT]):
                         max_errors=max_errors,
                         execute_function_calls=execute_function_calls,
                     )
-                    if approval_result["action"] == "stop":
+                    if approval_result.get("action") == "stop":
                         response = ChatResponse(messages=prepped_messages)
                         break
-                    errors_in_a_row = approval_result["errors_in_a_row"]
+                    errors_in_a_row = approval_result.get("errors_in_a_row", errors_in_a_row)
                     total_function_calls += approval_result.get("function_call_count", 0)
 
-                    response = await super_get_response(
-                        messages=prepped_messages,
-                        stream=False,
-                        options=mutable_options,
-                        **filtered_kwargs,
+                    response = cast(
+                        ChatResponse[Any],
+                        await super_get_response(
+                            messages=prepped_messages,
+                            stream=False,
+                            options=mutable_options,
+                            **filtered_kwargs,
+                        ),
                     )
 
                     if response.conversation_id is not None:
@@ -2174,10 +1936,10 @@ class FunctionInvocationLayer(Generic[OptionsCoT]):
                         max_errors=max_errors,
                         execute_function_calls=execute_function_calls,
                     )
-                    if result["action"] == "return":
+                    if result.get("action") == "return":
                         return response
                     total_function_calls += result.get("function_call_count", 0)
-                    if result["action"] == "stop":
+                    if result.get("action") == "stop":
                         # Error threshold reached: force a final non-tool turn so
                         # function_call_output items are submitted before exit.
                         mutable_options["tool_choice"] = "none"
@@ -2190,7 +1952,7 @@ class FunctionInvocationLayer(Generic[OptionsCoT]):
                             max_function_calls,
                         )
                         mutable_options["tool_choice"] = "none"
-                    errors_in_a_row = result["errors_in_a_row"]
+                    errors_in_a_row = result.get("errors_in_a_row", errors_in_a_row)
 
                     # When tool_choice is 'required', reset tool_choice after one iteration to avoid infinite loops
                     if mutable_options.get("tool_choice") == "required" or (
@@ -2213,17 +1975,20 @@ class FunctionInvocationLayer(Generic[OptionsCoT]):
                 # Make a final model call with tool_choice="none" so the model
                 # produces a plain text answer instead of leaving orphaned
                 # function_call items without matching results.
-                if response is not None and self.function_invocation_configuration["enabled"]:
+                if response is not None and self.function_invocation_configuration.get("enabled", True):
                     logger.info(
                         "Maximum iterations reached (%d). Requesting final response without tools.",
-                        self.function_invocation_configuration["max_iterations"],
+                        self.function_invocation_configuration.get("max_iterations", DEFAULT_MAX_ITERATIONS),
                     )
                 mutable_options["tool_choice"] = "none"
-                response = await super_get_response(
-                    messages=prepped_messages,
-                    stream=False,
-                    options=mutable_options,
-                    **filtered_kwargs,
+                response = cast(
+                    ChatResponse[Any],
+                    await super_get_response(
+                        messages=prepped_messages,
+                        stream=False,
+                        options=mutable_options,
+                        **filtered_kwargs,
+                    ),
                 )
                 if fcc_messages:
                     for msg in reversed(fcc_messages):
@@ -2233,7 +1998,7 @@ class FunctionInvocationLayer(Generic[OptionsCoT]):
             return _get_response()
 
         response_format = mutable_options.get("response_format") if mutable_options else None
-        output_format_type = response_format if isinstance(response_format, type) else None
+        output_format_type: type[BaseModel] | None = response_format if isinstance(response_format, type) else None
         stream_result_hooks: list[Callable[[ChatResponse], Any]] = []
 
         async def _stream() -> AsyncIterable[ChatResponseUpdate]:
@@ -2245,13 +2010,11 @@ class FunctionInvocationLayer(Generic[OptionsCoT]):
             max_function_calls: int | None = self.function_invocation_configuration.get("max_function_calls")
             prepped_messages = list(messages)
             fcc_messages: list[Message] = []
-            response: ChatResponse | None = None
+            response: ChatResponse[Any] | None = None
 
-            for attempt_idx in range(
-                self.function_invocation_configuration["max_iterations"]
-                if self.function_invocation_configuration["enabled"]
-                else 0
-            ):
+            loop_enabled = self.function_invocation_configuration.get("enabled", True)
+            max_iterations = self.function_invocation_configuration.get("max_iterations", DEFAULT_MAX_ITERATIONS)
+            for attempt_idx in range(max_iterations if loop_enabled else 0):
                 approval_result = await _process_function_requests(
                     response=None,
                     prepped_messages=prepped_messages,
@@ -2262,20 +2025,22 @@ class FunctionInvocationLayer(Generic[OptionsCoT]):
                     max_errors=max_errors,
                     execute_function_calls=execute_function_calls,
                 )
-                errors_in_a_row = approval_result["errors_in_a_row"]
+                errors_in_a_row = approval_result.get("errors_in_a_row", errors_in_a_row)
                 total_function_calls += approval_result.get("function_call_count", 0)
-                if approval_result["action"] == "stop":
+                if approval_result.get("action") == "stop":
                     mutable_options["tool_choice"] = "none"
                     return
 
-                inner_stream = await _ensure_response_stream(
+                inner_stream = cast(
+                    ResponseStream[ChatResponseUpdate, ChatResponse[Any]],
                     super_get_response(
                         messages=prepped_messages,
                         stream=True,
                         options=mutable_options,
                         **filtered_kwargs,
-                    )
+                    ),
                 )
+                await inner_stream
                 # Collect result hooks from the inner stream to run later
                 stream_result_hooks[:] = _get_result_hooks_from_stream(inner_stream)
 
@@ -2308,18 +2073,18 @@ class FunctionInvocationLayer(Generic[OptionsCoT]):
                     max_errors=max_errors,
                     execute_function_calls=execute_function_calls,
                 )
-                errors_in_a_row = result["errors_in_a_row"]
+                errors_in_a_row = result.get("errors_in_a_row", errors_in_a_row)
                 total_function_calls += result.get("function_call_count", 0)
-                if role := result["update_role"]:
+                if role := result.get("update_role"):
                     yield ChatResponseUpdate(
-                        contents=result["function_call_results"] or [],
+                        contents=result.get("function_call_results") or [],
                         role=role,
                     )
-                if result["action"] == "stop":
+                if result.get("action") == "stop":
                     # Error threshold reached: submit collected function_call_output
                     # items once more with tools disabled.
                     mutable_options["tool_choice"] = "none"
-                elif result["action"] != "continue":
+                elif result.get("action") != "continue":
                     return
                 elif max_function_calls is not None and total_function_calls >= max_function_calls:
                     # Best-effort limit: checked after each batch of parallel calls completes,
@@ -2352,26 +2117,28 @@ class FunctionInvocationLayer(Generic[OptionsCoT]):
             # Make a final model call with tool_choice="none" so the model
             # produces a plain text answer instead of leaving orphaned
             # function_call items without matching results.
-            if response is not None and self.function_invocation_configuration["enabled"]:
+            if response is not None and self.function_invocation_configuration.get("enabled", True):
                 logger.info(
                     "Maximum iterations reached (%d). Requesting final response without tools.",
-                    self.function_invocation_configuration["max_iterations"],
+                    self.function_invocation_configuration.get("max_iterations", DEFAULT_MAX_ITERATIONS),
                 )
             mutable_options["tool_choice"] = "none"
-            inner_stream = await _ensure_response_stream(
+            final_inner_stream = cast(
+                ResponseStream[ChatResponseUpdate, ChatResponse[Any]],
                 super_get_response(
                     messages=prepped_messages,
                     stream=True,
                     options=mutable_options,
                     **filtered_kwargs,
-                )
+                ),
             )
-            async for update in inner_stream:
+            await final_inner_stream
+            async for update in final_inner_stream:
                 yield update
             # Finalize the inner stream to trigger its hooks
-            await inner_stream.get_final_response()
+            await final_inner_stream.get_final_response()
 
-        def _finalize(updates: Sequence[ChatResponseUpdate]) -> ChatResponse:
+        def _finalize(updates: Sequence[ChatResponseUpdate]) -> ChatResponse[Any]:
             # Note: stream_result_hooks are already run via inner stream's get_final_response()
             # We don't need to run them again here
             return ChatResponse.from_updates(updates, output_format_type=output_format_type)

python/packages/core/agent_framework/_types.py

@@ -17,12 +17,15 @@ from collections.abc import (
     Mapping,
     MutableMapping,
     Sequence,
+    Sized,
 )
 from copy import deepcopy
 from datetime import datetime
+from inspect import isawaitable
 from typing import TYPE_CHECKING, Any, ClassVar, Final, Generic, Literal, NewType, cast, overload
 
 from pydantic import BaseModel
+from typing_extensions import TypedDict
 
 from ._serialization import SerializationMixin
 from ._tools import ToolTypes
@@ -33,10 +36,6 @@ if sys.version_info >= (3, 13):
     from typing import TypeVar  # pragma: no cover
 else:
     from typing_extensions import TypeVar  # pragma: no cover
-if sys.version_info >= (3, 11):
-    from typing import TypedDict  # type: ignore # pragma: no cover
-else:
-    from typing_extensions import TypedDict  # type: ignore # pragma: no cover
 
 logger = logging.getLogger("agent_framework")
 
@@ -194,7 +193,7 @@ def _get_data_bytes_as_str(content: Content) -> str | None:
     return data  # type: ignore[return-value, no-any-return]
 
 
-def _get_data_bytes(content: Content) -> bytes | None:
+def _get_data_bytes(content: Content) -> bytes | None:  # pyright: ignore[reportUnusedFunction]
     """Extract and decode binary data from data URI.
 
     Args:
@@ -270,9 +269,9 @@ def _serialize_value(value: Any, exclude_none: bool) -> Any:
     if isinstance(value, Content):
         return value.to_dict(exclude_none=exclude_none)
     if isinstance(value, Sequence) and not isinstance(value, (str, bytes, bytearray)):
-        return [_serialize_value(item, exclude_none) for item in value]
+        return [_serialize_value(item, exclude_none) for item in cast(Iterable[Any], value)]
     if isinstance(value, Mapping):
-        return {k: _serialize_value(v, exclude_none) for k, v in value.items()}
+        return {k: _serialize_value(v, exclude_none) for k, v in value.items()}  # type: ignore[reportUnknownVariableType]
     if hasattr(value, "to_dict"):
         return value.to_dict()  # type: ignore[call-arg]
     return value
@@ -376,7 +375,7 @@ ContentT = TypeVar("ContentT", bound="Content")
 # endregion
 
 
-class UsageDetails(TypedDict, total=False):
+class UsageDetails(TypedDict, total=False, extra_items=int):  # type: ignore[call-arg]
     """A dictionary representing usage details.
 
     This is a non-closed dictionary, so any specific provider fields can be added as needed.
@@ -397,6 +396,9 @@ class UsageDetails(TypedDict, total=False):
 def add_usage_details(usage1: UsageDetails | None, usage2: UsageDetails | None) -> UsageDetails:
     """Add two UsageDetails dictionaries by summing all numeric values.
 
+    If any of the two usage details contains a key with a non-int value, it will be skipped,
+    even if the other contains a int-value on that key.
+
     Args:
         usage1: First usage details dictionary.
         usage2: Second usage details dictionary.
@@ -420,22 +422,15 @@ def add_usage_details(usage1: UsageDetails | None, usage2: UsageDetails | None)
         return usage1
 
     result = UsageDetails()
-
     # Combine all keys from both dictionaries
     all_keys = set(usage1.keys()) | set(usage2.keys())
-
     for key in all_keys:
-        val1 = usage1.get(key)
-        val2 = usage2.get(key)
-
-        # Sum if both present, otherwise use the non-None value
-        if val1 is not None and val2 is not None:
-            result[key] = val1 + val2  # type: ignore[literal-required, operator]
-        elif val1 is not None:
-            result[key] = val1  # type: ignore[literal-required]
-        elif val2 is not None:
-            result[key] = val2  # type: ignore[literal-required]
-
+        if not isinstance((val1 := usage1.get(key, 0)), (int | None)) or not isinstance(
+            (val2 := usage2.get(key, 0)), (int | None)
+        ):
+            logger.warning("Non `int` value found in usage details, skipping.")
+            continue
+        result[key] = (val1 or 0) + (val2 or 0)  # type: ignore[literal-required]
     return result
 
 
@@ -465,7 +460,7 @@ class Content:
         error_code: str | None = None,
         error_details: str | None = None,
         # Usage content fields
-        usage_details: dict[str, Any] | UsageDetails | None = None,
+        usage_details: UsageDetails | None = None,
         # Function call/result fields
         call_id: str | None = None,
         name: str | None = None,
@@ -1264,19 +1259,14 @@ class Content:
             return cls.from_data(remaining["data"], remaining["media_type"])
 
         # Handle nested Content objects (e.g., function_call in function_approval_request)
-        if "function_call" in remaining and isinstance(remaining["function_call"], dict):
-            remaining["function_call"] = cls.from_dict(remaining["function_call"])
+        if (function_call := remaining.get("function_call")) and isinstance(function_call, dict):
+            remaining["function_call"] = cls.from_dict(function_call)  # type: ignore[reportUnknownArgumentType]
 
         # Handle list of Content objects (e.g., inputs in code_interpreter_tool_call)
-        if "inputs" in remaining and isinstance(remaining["inputs"], list):
-            remaining["inputs"] = [
-                cls.from_dict(item) if isinstance(item, dict) else item for item in remaining["inputs"]
-            ]
-
-        if "outputs" in remaining and isinstance(remaining["outputs"], list):
-            remaining["outputs"] = [
-                cls.from_dict(item) if isinstance(item, dict) else item for item in remaining["outputs"]
-            ]
+        if (input_items := remaining.get("inputs")) and isinstance(input_items, list):
+            remaining["inputs"] = [cls.from_dict(item) if isinstance(item, dict) else item for item in input_items]  # type: ignore[reportUnknownVariableType]
+        if (output_items := remaining.get("outputs")) and isinstance(output_items, list):
+            remaining["outputs"] = [cls.from_dict(item) if isinstance(item, dict) else item for item in output_items]  # type: ignore[reportUnknownVariableType]
 
         return cls(
             type=content_type,
@@ -1306,55 +1296,16 @@ class Content:
 
     def _add_text_content(self, other: Content) -> Content:
         """Add two TextContent instances."""
-        # Merge raw representations
-        if self.raw_representation is None:
-            raw_representation = other.raw_representation
-        elif other.raw_representation is None:
-            raw_representation = self.raw_representation
-        else:
-            raw_representation = (
-                self.raw_representation if isinstance(self.raw_representation, list) else [self.raw_representation]
-            ) + (other.raw_representation if isinstance(other.raw_representation, list) else [other.raw_representation])
-
-        # Merge annotations
-        if self.annotations is None:
-            annotations = other.annotations
-        elif other.annotations is None:
-            annotations = self.annotations
-        else:
-            annotations = self.annotations + other.annotations  # type: ignore[operator]
-
         return Content(
             "text",
             text=self.text + other.text,  # type: ignore[attr-defined, operator]
-            annotations=annotations,
-            additional_properties={
-                **(other.additional_properties or {}),
-                **(self.additional_properties or {}),
-            },
-            raw_representation=raw_representation,
+            annotations=_combine_annotations(self.annotations, other.annotations),
+            additional_properties=_combine_additional_props(self.additional_properties, other.additional_properties),
+            raw_representation=_combine_raw_representations(self.raw_representation, other.raw_representation),
         )
 
     def _add_text_reasoning_content(self, other: Content) -> Content:
         """Add two TextReasoningContent instances."""
-        # Merge raw representations
-        if self.raw_representation is None:
-            raw_representation = other.raw_representation
-        elif other.raw_representation is None:
-            raw_representation = self.raw_representation
-        else:
-            raw_representation = (
-                self.raw_representation if isinstance(self.raw_representation, list) else [self.raw_representation]
-            ) + (other.raw_representation if isinstance(other.raw_representation, list) else [other.raw_representation])
-
-        # Merge annotations
-        if self.annotations is None:
-            annotations = other.annotations
-        elif other.annotations is None:
-            annotations = self.annotations
-        else:
-            annotations = self.annotations + other.annotations  # type: ignore[operator]
-
         # Concatenate text, handling None values
         self_text = self.text or ""  # type: ignore[attr-defined]
         other_text = other.text or ""  # type: ignore[attr-defined]
@@ -1367,12 +1318,9 @@ class Content:
             "text_reasoning",
             text=combined_text,
             protected_data=protected_data,
-            annotations=annotations,
-            additional_properties={
-                **(other.additional_properties or {}),
-                **(self.additional_properties or {}),
-            },
-            raw_representation=raw_representation,
+            annotations=_combine_annotations(self.annotations, other.annotations),
+            additional_properties=_combine_additional_props(self.additional_properties, other.additional_properties),
+            raw_representation=_combine_raw_representations(self.raw_representation, other.raw_representation),
         )
 
     def _add_function_call_content(self, other: Content) -> Content:
@@ -1396,64 +1344,23 @@ class Content:
         else:
             raise TypeError("Incompatible argument types")
 
-        # Merge raw representations
-        if self.raw_representation is None:
-            raw_representation: Any = other.raw_representation
-        elif other.raw_representation is None:
-            raw_representation = self.raw_representation
-        else:
-            raw_representation = (
-                self.raw_representation if isinstance(self.raw_representation, list) else [self.raw_representation]
-            ) + (other.raw_representation if isinstance(other.raw_representation, list) else [other.raw_representation])
-
         return Content(
             "function_call",
             call_id=self_call_id,
             name=getattr(self, "name", getattr(other, "name", None)),
             arguments=arguments,
             exception=getattr(self, "exception", None) or getattr(other, "exception", None),
-            additional_properties={
-                **(self.additional_properties or {}),
-                **(other.additional_properties or {}),
-            },
-            raw_representation=raw_representation,
+            additional_properties=_combine_additional_props(self.additional_properties, other.additional_properties),
+            raw_representation=_combine_raw_representations(self.raw_representation, other.raw_representation),
         )
 
     def _add_usage_content(self, other: Content) -> Content:
         """Add two UsageContent instances by combining their usage details."""
-        self_details = getattr(self, "usage_details", {})
-        other_details = getattr(other, "usage_details", {})
-
-        # Combine token counts
-        combined_details: dict[str, Any] = {}
-        for key in set(list(self_details.keys()) + list(other_details.keys())):
-            self_val = self_details.get(key)
-            other_val = other_details.get(key)
-            if isinstance(self_val, int) and isinstance(other_val, int):
-                combined_details[key] = self_val + other_val
-            elif self_val is not None:
-                combined_details[key] = self_val
-            elif other_val is not None:
-                combined_details[key] = other_val
-
-        # Merge raw representations
-        if self.raw_representation is None:
-            raw_representation = other.raw_representation
-        elif other.raw_representation is None:
-            raw_representation = self.raw_representation
-        else:
-            raw_representation = (
-                self.raw_representation if isinstance(self.raw_representation, list) else [self.raw_representation]
-            ) + (other.raw_representation if isinstance(other.raw_representation, list) else [other.raw_representation])
-
         return Content(
             "usage",
-            usage_details=combined_details,
-            additional_properties={
-                **(self.additional_properties or {}),
-                **(other.additional_properties or {}),
-            },
-            raw_representation=raw_representation,
+            usage_details=add_usage_details(self.usage_details, other.usage_details),
+            additional_properties=_combine_additional_props(self.additional_properties, other.additional_properties),
+            raw_representation=_combine_raw_representations(self.raw_representation, other.raw_representation),
         )
 
     def has_top_level_media_type(self, top_level_media_type: Literal["application", "audio", "image", "text"]) -> bool:
@@ -1530,6 +1437,42 @@ class Content:
         return self.arguments  # type: ignore[return-value]
 
 
+def _combine_additional_props(
+    self_additional_properties: dict[str, Any], other_additional_properties: dict[str, Any]
+) -> dict[str, Any]:
+    """Combine additional properties for addition operations."""
+    return {
+        **other_additional_properties,
+        **self_additional_properties,
+    }
+
+
+def _combine_raw_representations(
+    self_repr: Any,
+    other_repr: Any,
+) -> Any:
+    """Combine raw representations for addition operations."""
+    if self_repr is None:
+        return other_repr
+    if other_repr is None:
+        return self_repr
+    self_list = self_repr if isinstance(self_repr, list) else [self_repr]  # type: ignore[reportUnknownVariableType]
+    other_list = other_repr if isinstance(other_repr, list) else [other_repr]  # type: ignore[reportUnknownVariableType]
+    return self_list + other_list  # type: ignore[reportUnknownVariableType]
+
+
+def _combine_annotations(
+    self_annotations: Sequence[Annotation] | None,
+    other_annotations: Sequence[Annotation] | None,
+) -> Sequence[Annotation] | None:
+    """Combine annotations for addition operations."""
+    if self_annotations is None:
+        return other_annotations
+    if other_annotations is None:
+        return self_annotations
+    return [*self_annotations, *other_annotations]
+
+
 # endregion
 
 
@@ -1665,10 +1608,6 @@ class Message(SerializationMixin):
                 Additional properties are used within Agent Framework, they are not sent to services.
             raw_representation: Optional raw representation of the chat message.
         """
-        # Handle role conversion from legacy dict format
-        if isinstance(role, dict) and "value" in role:
-            role = role["value"]
-
         # Handle contents conversion
         parsed_contents = [] if contents is None else _parse_content_list(contents)
 
@@ -1836,14 +1775,14 @@ def _process_update(response: ChatResponse | AgentResponse, update: ChatResponse
     if update.created_at is not None:
         response.created_at = update.created_at
     if update.additional_properties is not None:
-        if response.additional_properties is None:
-            response.additional_properties = {}
         response.additional_properties.update(update.additional_properties)
     if response.raw_representation is None:
         response.raw_representation = []
     if not isinstance(response.raw_representation, list):
         response.raw_representation = [response.raw_representation]
-    response.raw_representation.append(update.raw_representation)
+    raw_representation_value = cast(Any, getattr(response, "raw_representation", None))
+    raw_representation_list = cast(list[Any], raw_representation_value)
+    raw_representation_list.append(update.raw_representation)
     if isinstance(response, ChatResponse) and isinstance(update, ChatResponseUpdate):
         if update.conversation_id is not None:
             response.conversation_id = update.conversation_id
@@ -2026,9 +1965,6 @@ class ChatResponse(SerializationMixin, Generic[ResponseModelT]):
         self.conversation_id = conversation_id
         self.model_id = model_id
         self.created_at = created_at
-        # Handle legacy dict format for finish_reason
-        if isinstance(finish_reason, dict) and "value" in finish_reason:
-            finish_reason = finish_reason["value"]
         self.finish_reason = finish_reason
         self.usage_details = usage_details
         self._value: ResponseModelT | None = value
@@ -2620,10 +2556,6 @@ class AgentResponseUpdate(SerializationMixin):
                     processed_contents.append(c)
             self.contents = processed_contents
 
-        # Handle legacy dict format for role
-        if isinstance(role, dict) and "value" in role:
-            role = role["value"]
-
         self.role: str | None = role
         self.author_name = author_name
         self.agent_id = agent_id
@@ -2717,7 +2649,7 @@ class ResponseStream(AsyncIterable[UpdateT], Generic[UpdateT, FinalT]):
         self._inner_stream: ResponseStream[Any, Any] | None = None
         self._inner_stream_source: ResponseStream[Any, Any] | Awaitable[ResponseStream[Any, Any]] | None = None
         self._wrap_inner: bool = False
-        self._map_update: Callable[[Any], Any | Awaitable[Any]] | None = None
+        self._map_update: Callable[[Any], UpdateT | Awaitable[UpdateT]] | None = None
 
     def map(
         self,
@@ -2757,11 +2689,11 @@ class ResponseStream(AsyncIterable[UpdateT], Generic[UpdateT, FinalT]):
             ...     AgentResponse.from_updates,
             ... )
         """
-        stream: ResponseStream[Any, Any] = ResponseStream(self, finalizer=finalizer)
+        stream: ResponseStream[OuterUpdateT, OuterFinalT] = ResponseStream(self, finalizer=finalizer)
         stream._inner_stream_source = self
         stream._wrap_inner = True
         stream._map_update = transform
-        return stream  # type: ignore[return-value]
+        return stream
 
     def with_finalizer(
         self,
@@ -2785,10 +2717,10 @@ class ResponseStream(AsyncIterable[UpdateT], Generic[UpdateT, FinalT]):
         Example:
             >>> stream.with_finalizer(AgentResponse.from_updates)
         """
-        stream: ResponseStream[Any, Any] = ResponseStream(self, finalizer=finalizer)
+        stream: ResponseStream[UpdateT, OuterFinalT] = ResponseStream(self, finalizer=finalizer)
         stream._inner_stream_source = self
         stream._wrap_inner = True
-        return stream  # type: ignore[return-value]
+        return stream
 
     @classmethod
     def from_awaitable(
@@ -2813,10 +2745,10 @@ class ResponseStream(AsyncIterable[UpdateT], Generic[UpdateT, FinalT]):
             >>> async def get_stream() -> ResponseStream[Update, Response]: ...
             >>> stream = ResponseStream.from_awaitable(get_stream())
         """
-        stream: ResponseStream[Any, Any] = cls(awaitable)  # type: ignore[arg-type]
-        stream._inner_stream_source = awaitable  # type: ignore[assignment]
+        stream: ResponseStream[UpdateT, FinalT] = cls(cast(Awaitable[AsyncIterable[UpdateT]], awaitable))
+        stream._inner_stream_source = awaitable
         stream._wrap_inner = True
-        return stream  # type: ignore[return-value]
+        return stream
 
     async def _get_stream(self) -> AsyncIterable[UpdateT]:
         if self._stream is None:
@@ -2826,10 +2758,10 @@ class ResponseStream(AsyncIterable[UpdateT], Generic[UpdateT, FinalT]):
                 if not iscoroutine(self._stream_source):
                     self._stream = self._stream_source  # type: ignore[assignment]
                 else:
-                    self._stream = await self._stream_source  # type: ignore[assignment]
+                    self._stream = await self._stream_source
             if isinstance(self._stream, ResponseStream) and self._wrap_inner:
-                self._inner_stream = self._stream
-                return self._stream
+                self._inner_stream = self._stream  # type: ignore[assignment]
+                return self._inner_stream
         return self._stream  # type: ignore[return-value]
 
     def __aiter__(self) -> ResponseStream[UpdateT, FinalT]:
@@ -2840,7 +2772,7 @@ class ResponseStream(AsyncIterable[UpdateT], Generic[UpdateT, FinalT]):
             stream = await self._get_stream()
             self._iterator = stream.__aiter__()
         try:
-            update = await self._iterator.__anext__()
+            update: UpdateT = await self._iterator.__anext__()
         except StopAsyncIteration:
             self._consumed = True
             await self._run_cleanup_hooks()
@@ -2849,18 +2781,16 @@ class ResponseStream(AsyncIterable[UpdateT], Generic[UpdateT, FinalT]):
             await self._run_cleanup_hooks()
             raise
         if self._map_update is not None:
-            mapped = self._map_update(update)
-            if isinstance(mapped, Awaitable):
-                update = await mapped
-            else:
-                update = mapped  # type: ignore[assignment]
+            update = self._map_update(update)  # type: ignore[assignment]
+            if isawaitable(update):
+                update = await update
         self._updates.append(update)
         for hook in self._transform_hooks:
             hooked = hook(update)
-            if isinstance(hooked, Awaitable):
-                update = await hooked
-            elif hooked is not None:
-                update = hooked  # type: ignore[assignment]
+            if isawaitable(hooked):
+                hooked = await hooked
+            if hooked is not None:
+                update = hooked
         return update
 
     def __await__(self) -> Any:
@@ -2903,58 +2833,71 @@ class ResponseStream(AsyncIterable[UpdateT], Generic[UpdateT, FinalT]):
 
                 # First, finalize the inner stream and run its result hooks
                 # This ensures inner post-processing (e.g., context provider notifications) runs
-                if self._inner_stream._finalizer is not None:
-                    inner_result: Any = self._inner_stream._finalizer(self._inner_stream._updates)
-                    if isinstance(inner_result, Awaitable):
+                inner_stream = self._inner_stream
+                inner_result: Any
+                if inner_stream._finalizer is not None:
+                    inner_finalizer = inner_stream._finalizer
+                    inner_result = inner_finalizer(inner_stream._updates)
+                    if isawaitable(inner_result):
                         inner_result = await inner_result
                 else:
-                    inner_result = self._inner_stream._updates
+                    inner_result = list(inner_stream._updates)
+
                 # Run inner stream's result hooks
-                for hook in self._inner_stream._result_hooks:
-                    hooked = hook(inner_result)
-                    if isinstance(hooked, Awaitable):
-                        hooked = await hooked
-                    if hooked is not None:
-                        inner_result = hooked
-                self._inner_stream._final_result = inner_result
-                self._inner_stream._finalized = True
+                inner_hooks = cast(list[Callable[[Any], Any | Awaitable[Any] | None]], inner_stream._result_hooks)
+                for hook in inner_hooks:
+                    hooked_result = hook(inner_result)
+                    if isawaitable(hooked_result):
+                        hooked_result = await hooked_result
+                    if hooked_result is not None:
+                        inner_result = hooked_result
+                inner_stream._final_result = inner_result
+                inner_stream._finalized = True
 
                 # Now finalize the outer stream with its own finalizer
                 # If outer has no finalizer, use inner's result (preserves from_awaitable behavior)
+                outer_result: Any
                 if self._finalizer is not None:
-                    result: Any = self._finalizer(self._updates)
-                    if isinstance(result, Awaitable):
-                        result = await result
+                    outer_result = self._finalizer(self._updates)
+                    if isawaitable(outer_result):
+                        outer_result = await outer_result
                 else:
                     # No outer finalizer - use inner's finalized result
-                    result = inner_result
+                    outer_result = inner_result
+
                 # Apply outer's result_hooks
-                for hook in self._result_hooks:
-                    hooked = hook(result)
-                    if isinstance(hooked, Awaitable):
-                        hooked = await hooked
-                    if hooked is not None:
-                        result = hooked
-                self._final_result = result
+                outer_hooks = cast(list[Callable[[Any], Any | Awaitable[Any] | None]], self._result_hooks)
+                for hook in outer_hooks:
+                    outer_hook_result = hook(outer_result)
+                    if isawaitable(outer_hook_result):
+                        outer_hook_result = await outer_hook_result
+                    if outer_hook_result is not None:
+                        outer_result = outer_hook_result
+                self._final_result = outer_result
                 self._finalized = True
             return self._final_result  # type: ignore[return-value]
+
         if not self._finalized:
             if not self._consumed:
                 async for _ in self:
                     pass
+
             # Use finalizer if configured, otherwise return collected updates
+            result: Any
             if self._finalizer is not None:
                 result = self._finalizer(self._updates)
-                if isinstance(result, Awaitable):
+                if isawaitable(result):
                     result = await result
             else:
-                result = self._updates
-            for hook in self._result_hooks:
-                hooked = hook(result)
-                if isinstance(hooked, Awaitable):
-                    hooked = await hooked
-                if hooked is not None:
-                    result = hooked
+                result = list(self._updates)
+
+            final_hooks = cast(list[Callable[[Any], Any | Awaitable[Any] | None]], self._result_hooks)
+            for hook in final_hooks:
+                final_hook_result = hook(result)
+                if isawaitable(final_hook_result):
+                    final_hook_result = await final_hook_result
+                if final_hook_result is not None:
+                    result = final_hook_result
             self._final_result = result
             self._finalized = True
         return self._final_result  # type: ignore[return-value]
@@ -2991,7 +2934,7 @@ class ResponseStream(AsyncIterable[UpdateT], Generic[UpdateT, FinalT]):
         self._cleanup_run = True
         for hook in self._cleanup_hooks:
             result = hook()
-            if isinstance(result, Awaitable):
+            if isawaitable(result):
                 await result
 
     @property
@@ -3302,9 +3245,9 @@ def merge_chat_options(
     # Copy base values (shallow copy for simple values, dict copy for dicts)
     for key, value in base.items():
         if isinstance(value, dict):
-            result[key] = dict(value)
+            result[key] = dict(value)  # type: ignore[reportUnknownArgumentType]
         elif isinstance(value, list):
-            result[key] = list(value)
+            result[key] = list(value)  # type: ignore[reportUnknownArgumentType]
         else:
             result[key] = value
 
@@ -3326,19 +3269,19 @@ def merge_chat_options(
             if base_tools and value:
                 # Add tools that aren't already present
                 merged_tools = list(base_tools)
-                for tool in value if isinstance(value, list) else [value]:
+                for tool in value if isinstance(value, Iterable) else [value]:  # type: ignore[reportUnknownVariableType]
                     if tool not in merged_tools:
                         merged_tools.append(tool)
                 result["tools"] = merged_tools
             elif value:
-                result["tools"] = list(value) if isinstance(value, list) else [value]
+                result["tools"] = value if isinstance(value, list) else [value]
         elif key in ("logit_bias", "metadata", "additional_properties"):
             # Merge dicts
             base_dict = result.get(key)
-            if base_dict and isinstance(value, dict):
+            if base_dict and isinstance(base_dict, dict) and isinstance(value, dict):
                 result[key] = {**base_dict, **value}
             elif value:
-                result[key] = dict(value) if isinstance(value, dict) else value
+                result[key] = dict(cast(Mapping[Any, Any], value)) if isinstance(value, dict) else value
         elif key == "tool_choice":
             # tool_choice from override takes precedence
             result["tool_choice"] = value if value else result.get("tool_choice")
@@ -3424,8 +3367,8 @@ class Embedding(Generic[EmbeddingT]):
         """
         if self._dimensions is not None:
             return self._dimensions
-        if isinstance(self.vector, (list, tuple, bytes)):
-            return len(self.vector)
+        if isinstance(self.vector, Sized) and not isinstance(self.vector, str):
+            return len(cast(Sized, self.vector))
         return None
 
 

python/packages/core/agent_framework/_workflows/_agent_executor.py

@@ -450,9 +450,9 @@ class AgentExecutor(Executor):
         options: dict[str, Any] = {}
         if options_from_workflow is not None:
             if isinstance(options_from_workflow, Mapping):
-                for key, value in options_from_workflow.items():
-                    if isinstance(key, str):
-                        options[key] = value
+                options_from_workflow_map = cast(Mapping[str, Any], options_from_workflow)
+                for key, value in options_from_workflow_map.items():
+                    options[key] = value
             else:
                 logger.warning(
                     "Ignoring non-mapping workflow 'options' kwarg of type %s for AgentExecutor %s.",
@@ -461,16 +461,17 @@ class AgentExecutor(Executor):
                 )
 
         existing_additional_args = options.get("additional_function_arguments")
+        additional_args: dict[str, Any]
         if isinstance(existing_additional_args, Mapping):
-            additional_args = {key: value for key, value in existing_additional_args.items() if isinstance(key, str)}
+            existing_additional_args_map = cast(Mapping[str, Any], existing_additional_args)
+            additional_args = {key: value for key, value in existing_additional_args_map.items()}
         else:
             additional_args = {}
 
         if workflow_additional_args is not None:
             if isinstance(workflow_additional_args, Mapping):
-                additional_args.update({
-                    key: value for key, value in workflow_additional_args.items() if isinstance(key, str)
-                })
+                workflow_additional_args_map = cast(Mapping[str, Any], workflow_additional_args)
+                additional_args.update({key: value for key, value in workflow_additional_args_map.items()})
             else:
                 logger.warning(
                     "Ignoring non-mapping workflow 'additional_function_arguments' kwarg of type %s for AgentExecutor %s.",  # noqa: E501

python/packages/core/agent_framework/_workflows/_function_executor.py

@@ -119,7 +119,7 @@ class FunctionExecutor(Executor):
         # Determine if function has WorkflowContext parameter
         self._has_context = ctx_annotation is not None
         # Determine if the function is an async function
-        self._is_async = asyncio.iscoroutinefunction(func)
+        self._is_async = inspect.iscoroutinefunction(func)
 
         # Initialize parent WITHOUT calling _discover_handlers yet
         # We'll manually set up the attributes first

python/packages/core/agent_framework/_workflows/_runner_context.py

@@ -99,11 +99,11 @@ class RunnerContext(Protocol):
     If checkpoint storage is not configured, checkpoint methods may raise.
     """
 
-    async def send_message(self, WorkflowMessage: WorkflowMessage) -> None:
+    async def send_message(self, message: WorkflowMessage) -> None:
         """Send a WorkflowMessage from the executor to the context.
 
         Args:
-            WorkflowMessage: The WorkflowMessage to be sent.
+            message: The WorkflowMessage to be sent.
         """
         ...
 
@@ -288,9 +288,9 @@ class InProcRunnerContext:
         self._streaming: bool = False
 
     # region Messaging and Events
-    async def send_message(self, WorkflowMessage: WorkflowMessage) -> None:
-        self._messages.setdefault(WorkflowMessage.source_id, [])
-        self._messages[WorkflowMessage.source_id].append(WorkflowMessage)
+    async def send_message(self, message: WorkflowMessage) -> None:
+        self._messages.setdefault(message.source_id, [])
+        self._messages[message.source_id].append(message)
 
     async def drain_messages(self) -> dict[str, list[WorkflowMessage]]:
         messages = copy(self._messages)

python/packages/core/agent_framework/_workflows/_typing_utils.py

@@ -193,36 +193,40 @@ def try_coerce_to_type(data: Any, target_type: type | UnionType | Any) -> Any:
     Returns:
         The coerced value, or the original value if coercion fails.
     """
+    original_data = data
+
     # If already the right type, return as-is
     if is_instance_of(data, target_type):
         return data
 
     # Can't coerce to non-concrete targets (Union, generic, etc.)
     if not isinstance(target_type, type):
-        return data
+        return original_data
+
+    target_cls: type[Any] = target_type
 
     # int -> float (JSON integers for float fields)
-    if isinstance(data, int) and target_type is float:
+    if isinstance(data, int) and target_cls is float:
         return float(data)
 
-    # dict -> dataclass
+    # dict -> dataclass or pydantic model
     if isinstance(data, dict):
         from dataclasses import is_dataclass
 
-        if is_dataclass(target_type):
+        if is_dataclass(target_cls):
             try:
-                return target_type(**data)
+                return target_cls(**data)
             except (TypeError, ValueError):
-                return data
+                return original_data
 
-        # dict -> Pydantic model
-        if hasattr(target_type, "model_validate"):
+        model_validate = getattr(target_cls, "model_validate", None)
+        if callable(model_validate):
             try:
-                return target_type.model_validate(data)
+                return model_validate(data)
             except Exception:
-                return data
+                return original_data
 
-    return data
+    return original_data
 
 
 def serialize_type(t: type) -> str:

python/packages/core/agent_framework/azure/_assistants_client.py

@@ -12,7 +12,7 @@ from .._settings import load_settings
 from ..openai import OpenAIAssistantsClient
 from ..openai._assistants_client import OpenAIAssistantsOptions
 from ._entra_id_authentication import AzureCredentialTypes, AzureTokenProvider, resolve_credential_to_token_provider
-from ._shared import AzureOpenAISettings, _apply_azure_defaults
+from ._shared import AzureOpenAISettings, _apply_azure_defaults  # pyright: ignore[reportPrivateUsage]
 
 if sys.version_info >= (3, 13):
     from typing import TypeVar  # type: ignore # pragma: no cover
@@ -145,43 +145,46 @@ class AzureOpenAIAssistantsClient(
         )
         _apply_azure_defaults(azure_openai_settings, default_api_version=self.DEFAULT_AZURE_API_VERSION)
 
-        if not azure_openai_settings["chat_deployment_name"]:
+        chat_deployment_name = azure_openai_settings.get("chat_deployment_name")
+        if not chat_deployment_name:
             raise ValueError(
                 "Azure OpenAI deployment name is required. Set via 'deployment_name' parameter "
                 "or 'AZURE_OPENAI_CHAT_DEPLOYMENT_NAME' environment variable."
             )
 
+        api_key_secret = azure_openai_settings.get("api_key")
+        token_scope = azure_openai_settings.get("token_endpoint")
+
         # Resolve credential to token provider
         ad_token_provider = None
-        if not async_client and not azure_openai_settings["api_key"] and credential:
-            ad_token_provider = resolve_credential_to_token_provider(
-                credential, azure_openai_settings["token_endpoint"]
-            )
+        if not async_client and not api_key_secret and credential:
+            ad_token_provider = resolve_credential_to_token_provider(credential, token_scope)
 
-        if not async_client and not azure_openai_settings["api_key"] and not ad_token_provider:
+        if not async_client and not api_key_secret and not ad_token_provider:
             raise ValueError("Please provide either api_key, credential, or a client.")
 
         # Create Azure client if not provided
         if not async_client:
             client_params: dict[str, Any] = {
-                "api_version": azure_openai_settings["api_version"],
                 "default_headers": default_headers,
             }
+            if resolved_api_version := azure_openai_settings.get("api_version"):
+                client_params["api_version"] = resolved_api_version
 
-            if azure_openai_settings["api_key"]:
-                client_params["api_key"] = azure_openai_settings["api_key"].get_secret_value()
+            if api_key_secret:
+                client_params["api_key"] = api_key_secret.get_secret_value()
             elif ad_token_provider:
                 client_params["azure_ad_token_provider"] = ad_token_provider
 
-            if azure_openai_settings["base_url"]:
-                client_params["base_url"] = str(azure_openai_settings["base_url"])
-            elif azure_openai_settings["endpoint"]:
-                client_params["azure_endpoint"] = str(azure_openai_settings["endpoint"])
+            if resolved_base_url := azure_openai_settings.get("base_url"):
+                client_params["base_url"] = str(resolved_base_url)
+            elif resolved_endpoint := azure_openai_settings.get("endpoint"):
+                client_params["azure_endpoint"] = str(resolved_endpoint)
 
             async_client = AsyncAzureOpenAI(**client_params)
 
         super().__init__(
-            model_id=azure_openai_settings["chat_deployment_name"],
+            model_id=chat_deployment_name,
             assistant_id=assistant_id,
             assistant_name=assistant_name,
             assistant_description=assistant_description,

python/packages/core/agent_framework/azure/_chat_client.py

@@ -6,7 +6,7 @@ import json
 import logging
 import sys
 from collections.abc import Mapping, Sequence
-from typing import TYPE_CHECKING, Any, Generic
+from typing import TYPE_CHECKING, Any, Generic, cast
 
 from openai.lib.azure import AsyncAzureOpenAI
 from openai.types.chat.chat_completion import Choice
@@ -31,7 +31,7 @@ from ._entra_id_authentication import AzureCredentialTypes, AzureTokenProvider
 from ._shared import (
     AzureOpenAIConfigMixin,
     AzureOpenAISettings,
-    _apply_azure_defaults,
+    _apply_azure_defaults,  # pyright: ignore[reportPrivateUsage]
 )
 
 if sys.version_info >= (3, 13):
@@ -260,19 +260,26 @@ class AzureOpenAIChatClient(  # type: ignore[misc]
         )
         _apply_azure_defaults(azure_openai_settings)
 
-        if not azure_openai_settings["chat_deployment_name"]:
+        chat_deployment_name = azure_openai_settings.get("chat_deployment_name")
+        if not chat_deployment_name:
             raise ValueError(
                 "Azure OpenAI deployment name is required. Set via 'deployment_name' parameter "
                 "or 'AZURE_OPENAI_CHAT_DEPLOYMENT_NAME' environment variable."
             )
 
+        endpoint_value = azure_openai_settings.get("endpoint")
+        base_url_value = azure_openai_settings.get("base_url")
+        api_version_value = cast(str, azure_openai_settings.get("api_version"))
+        api_key_value = azure_openai_settings.get("api_key")
+        token_endpoint_value = azure_openai_settings.get("token_endpoint")
+
         super().__init__(
-            deployment_name=azure_openai_settings["chat_deployment_name"],
-            endpoint=azure_openai_settings["endpoint"],
-            base_url=azure_openai_settings["base_url"],
-            api_version=azure_openai_settings["api_version"],  # type: ignore
-            api_key=azure_openai_settings["api_key"].get_secret_value() if azure_openai_settings["api_key"] else None,
-            token_endpoint=azure_openai_settings["token_endpoint"],
+            deployment_name=chat_deployment_name,
+            endpoint=endpoint_value,
+            base_url=base_url_value,
+            api_version=api_version_value,
+            api_key=api_key_value.get_secret_value() if api_key_value else None,
+            token_endpoint=token_endpoint_value,
             credential=credential,
             default_headers=default_headers,
             client=async_client,
@@ -302,24 +309,29 @@ class AzureOpenAIChatClient(  # type: ignore[misc]
         if not message.model_extra or "context" not in message.model_extra:
             return text_content
 
-        context: dict[str, Any] | str = message.context  # type: ignore[assignment, union-attr]
-        if isinstance(context, str):
+        context_raw: object = cast(object, message.context)  # type: ignore[union-attr]
+        if isinstance(context_raw, str):
             try:
-                context = json.loads(context)
+                context_raw = json.loads(context_raw)
             except json.JSONDecodeError:
                 logger.warning("Context is not a valid JSON string, ignoring context.")
                 return text_content
-        if not isinstance(context, dict):
+        if not isinstance(context_raw, dict):
             logger.warning("Context is not a valid dictionary, ignoring context.")
             return text_content
+        context = cast(dict[str, Any], context_raw)
         # `all_retrieved_documents` is currently not used, but can be retrieved
         # through the raw_representation in the text content.
         if intent := context.get("intent"):
             text_content.additional_properties = {"intent": intent}
-        if citations := context.get("citations"):
-            text_content.annotations = []
-            for citation in citations:
-                text_content.annotations.append(
+        citations = context.get("citations")
+        if isinstance(citations, list) and citations:
+            annotations: list[Annotation] = []
+            for citation_raw in cast(list[object], citations):
+                if not isinstance(citation_raw, dict):
+                    continue
+                citation = cast(dict[str, Any], citation_raw)
+                annotations.append(
                     Annotation(
                         type="citation",
                         title=citation.get("title", ""),
@@ -331,4 +343,5 @@ class AzureOpenAIChatClient(  # type: ignore[misc]
                         raw_representation=citation,
                     )
                 )
+            text_content.annotations = annotations
         return text_content

python/packages/core/agent_framework/azure/_embedding_client.py

@@ -17,7 +17,7 @@ from ._entra_id_authentication import AzureCredentialTypes, AzureTokenProvider
 from ._shared import (
     AzureOpenAIConfigMixin,
     AzureOpenAISettings,
-    _apply_azure_defaults,
+    _apply_azure_defaults,  # pyright: ignore[reportPrivateUsage]
 )
 
 if sys.version_info >= (3, 13):
@@ -118,19 +118,22 @@ class AzureOpenAIEmbeddingClient(
         )
         _apply_azure_defaults(azure_openai_settings)
 
-        if not azure_openai_settings.get("embedding_deployment_name"):
+        embedding_deployment_name = azure_openai_settings.get("embedding_deployment_name")
+        if not embedding_deployment_name:
             raise ValueError(
                 "Azure OpenAI embedding deployment name is required. Set via 'deployment_name' parameter "
                 "or 'AZURE_OPENAI_EMBEDDING_DEPLOYMENT_NAME' environment variable."
             )
 
+        api_key_secret = azure_openai_settings.get("api_key")
+
         super().__init__(
-            deployment_name=azure_openai_settings["embedding_deployment_name"],  # type: ignore[arg-type]
-            endpoint=azure_openai_settings["endpoint"],
-            base_url=azure_openai_settings["base_url"],
-            api_version=azure_openai_settings["api_version"],  # type: ignore
-            api_key=azure_openai_settings["api_key"].get_secret_value() if azure_openai_settings["api_key"] else None,
-            token_endpoint=azure_openai_settings["token_endpoint"],
+            deployment_name=embedding_deployment_name,
+            endpoint=azure_openai_settings.get("endpoint"),
+            base_url=azure_openai_settings.get("base_url"),
+            api_version=azure_openai_settings.get("api_version") or "",
+            api_key=api_key_secret.get_secret_value() if api_key_secret else None,
+            token_endpoint=azure_openai_settings.get("token_endpoint"),
             credential=credential,
             default_headers=default_headers,
             client=async_client,

python/packages/core/agent_framework/azure/_responses_client.py

@@ -20,7 +20,7 @@ from ._entra_id_authentication import AzureCredentialTypes, AzureTokenProvider
 from ._shared import (
     AzureOpenAIConfigMixin,
     AzureOpenAISettings,
-    _apply_azure_defaults,
+    _apply_azure_defaults,  # pyright: ignore[reportPrivateUsage]
 )
 
 if sys.version_info >= (3, 13):
@@ -207,27 +207,31 @@ class AzureOpenAIResponsesClient(  # type: ignore[misc]
         # TODO(peterychang): This is a temporary hack to ensure that the base_url is set correctly
         # while this feature is in preview.
         # But we should only do this if we're on azure. Private deployments may not need this.
+        endpoint_value = azure_openai_settings.get("endpoint")
         if (
             not azure_openai_settings.get("base_url")
-            and azure_openai_settings.get("endpoint")
-            and (hostname := urlparse(str(azure_openai_settings["endpoint"])).hostname)
+            and endpoint_value
+            and (hostname := urlparse(str(endpoint_value)).hostname)
             and hostname.endswith(".openai.azure.com")
         ):
-            azure_openai_settings["base_url"] = urljoin(str(azure_openai_settings["endpoint"]), "/openai/v1/")
+            azure_openai_settings["base_url"] = urljoin(str(endpoint_value), "/openai/v1/")
 
-        if not azure_openai_settings["responses_deployment_name"]:
+        responses_deployment_name = azure_openai_settings.get("responses_deployment_name")
+        if not responses_deployment_name:
             raise ValueError(
                 "Azure OpenAI deployment name is required. Set via 'deployment_name' parameter "
                 "or 'AZURE_OPENAI_RESPONSES_DEPLOYMENT_NAME' environment variable."
             )
 
+        api_key_secret = azure_openai_settings.get("api_key")
+
         super().__init__(
-            deployment_name=azure_openai_settings["responses_deployment_name"],
-            endpoint=azure_openai_settings["endpoint"],
-            base_url=azure_openai_settings["base_url"],
-            api_version=azure_openai_settings["api_version"],  # type: ignore
-            api_key=azure_openai_settings["api_key"].get_secret_value() if azure_openai_settings["api_key"] else None,
-            token_endpoint=azure_openai_settings["token_endpoint"],
+            deployment_name=responses_deployment_name,
+            endpoint=azure_openai_settings.get("endpoint"),
+            base_url=azure_openai_settings.get("base_url"),
+            api_version=azure_openai_settings.get("api_version") or "",
+            api_key=api_key_secret.get_secret_value() if api_key_secret else None,
+            token_endpoint=azure_openai_settings.get("token_endpoint"),
             credential=credential,
             default_headers=default_headers,
             client=async_client,

python/packages/core/agent_framework/azure/_shared.py

@@ -123,6 +123,9 @@ def _apply_azure_defaults(
         settings["token_endpoint"] = default_token_endpoint
 
 
+_AZURE_DEFAULTS_APPLIER = _apply_azure_defaults
+
+
 class AzureOpenAIConfigMixin(OpenAIBase):
     """Internal class for configuring a connection to an Azure OpenAI service."""
 

python/packages/core/agent_framework/declarative/__init__.pyi

@@ -4,7 +4,6 @@ from agent_framework_declarative import (
     AgentExternalInputRequest,
     AgentExternalInputResponse,
     AgentFactory,
-    AgentInvocationError,
     DeclarativeLoaderError,
     DeclarativeWorkflowError,
     ExternalInputRequest,
@@ -19,7 +18,6 @@ __all__ = [
     "AgentExternalInputRequest",
     "AgentExternalInputResponse",
     "AgentFactory",
-    "AgentInvocationError",
     "DeclarativeLoaderError",
     "DeclarativeWorkflowError",
     "ExternalInputRequest",

python/packages/core/agent_framework/observability.py

@@ -22,7 +22,7 @@ import weakref
 from collections.abc import Awaitable, Callable, Generator, Mapping, Sequence
 from enum import Enum
 from time import perf_counter, time_ns
-from typing import TYPE_CHECKING, Any, ClassVar, Final, Generic, Literal, TypedDict, overload
+from typing import TYPE_CHECKING, Any, ClassVar, Final, Generic, Literal, TypedDict, cast, overload
 
 from dotenv import load_dotenv
 from opentelemetry import metrics, trace
@@ -199,6 +199,7 @@ class OtelAttr(str, Enum):
     T_TYPE_INPUT = "input"
     T_TYPE_OUTPUT = "output"
     DURATION_UNIT = "s"
+
     # Agent attributes
     AGENT_NAME = "gen_ai.agent.name"
     AGENT_DESCRIPTION = "gen_ai.agent.description"
@@ -894,7 +895,6 @@ def get_meter(
         return metrics.get_meter(name=name, version=version, schema_url=schema_url)
 
 
-global OBSERVABILITY_SETTINGS
 OBSERVABILITY_SETTINGS: ObservabilitySettings = ObservabilitySettings()
 
 
@@ -1053,7 +1053,15 @@ def configure_otel_providers(
         if vs_code_extension_port is not None:
             settings_kwargs["vs_code_extension_port"] = vs_code_extension_port
 
-        OBSERVABILITY_SETTINGS = ObservabilitySettings(**settings_kwargs)
+        updated_settings = ObservabilitySettings(**settings_kwargs)
+        OBSERVABILITY_SETTINGS.enable_instrumentation = updated_settings.enable_instrumentation
+        OBSERVABILITY_SETTINGS.enable_sensitive_data = updated_settings.enable_sensitive_data
+        OBSERVABILITY_SETTINGS.enable_console_exporters = updated_settings.enable_console_exporters
+        OBSERVABILITY_SETTINGS.vs_code_extension_port = updated_settings.vs_code_extension_port
+        OBSERVABILITY_SETTINGS.env_file_path = updated_settings.env_file_path
+        OBSERVABILITY_SETTINGS.env_file_encoding = updated_settings.env_file_encoding
+        OBSERVABILITY_SETTINGS._resource = updated_settings._resource  # type: ignore[reportPrivateUsage]
+        OBSERVABILITY_SETTINGS._executed_setup = False  # type: ignore[reportPrivateUsage]
     else:
         # Update the observability settings with the provided values
         OBSERVABILITY_SETTINGS.enable_instrumentation = True
@@ -1146,6 +1154,8 @@ class ChatTelemetryLayer(Generic[OptionsCoT]):
         **kwargs: Any,
     ) -> Awaitable[ChatResponse[Any]] | ResponseStream[ChatResponseUpdate, ChatResponse[Any]]:
         """Trace chat responses with OpenTelemetry spans and metrics."""
+        from ._types import ChatResponse, ChatResponseUpdate, ResponseStream  # type: ignore[reportUnusedImport]
+
         global OBSERVABILITY_SETTINGS
         super_get_response = super().get_response  # type: ignore[misc]
 
@@ -1153,7 +1163,7 @@ class ChatTelemetryLayer(Generic[OptionsCoT]):
             return super_get_response(messages=messages, stream=stream, options=options, **kwargs)  # type: ignore[no-any-return]
 
         opts: dict[str, Any] = options or {}  # type: ignore[assignment]
-        provider_name = str(self.otel_provider_name)
+        provider_name = str(getattr(self, "otel_provider_name", "unknown"))
         model_id = kwargs.get("model_id") or opts.get("model_id") or getattr(self, "model_id", None) or "unknown"
         service_url_func = getattr(self, "service_url", None)
         service_url = str(service_url_func() if callable(service_url_func) else "unknown")
@@ -1166,15 +1176,10 @@ class ChatTelemetryLayer(Generic[OptionsCoT]):
         )
 
         if stream:
-            from ._types import ResponseStream
-
-            stream_result = super_get_response(messages=messages, stream=True, options=opts, **kwargs)
-            if isinstance(stream_result, ResponseStream):
-                result_stream = stream_result
-            elif isinstance(stream_result, Awaitable):
-                result_stream = ResponseStream.from_awaitable(stream_result)
-            else:
-                raise RuntimeError("Streaming telemetry requires a ResponseStream result.")
+            result_stream = cast(
+                ResponseStream[ChatResponseUpdate, ChatResponse[Any]],
+                super_get_response(messages=messages, stream=True, options=opts, **kwargs),
+            )
 
             # Create span directly without trace.use_span() context attachment.
             # Streaming spans are closed asynchronously in cleanup hooks, which run
@@ -1209,14 +1214,14 @@ class ChatTelemetryLayer(Generic[OptionsCoT]):
                 from ._types import ChatResponse
 
                 try:
-                    response = await result_stream.get_final_response()
+                    response: ChatResponse[Any] = await result_stream.get_final_response()
                     duration = duration_state.get("duration")
                     response_attributes = _get_response_attributes(attributes, response)
                     _capture_response(
                         span=span,
                         attributes=response_attributes,
-                        token_usage_histogram=self.token_usage_histogram,
-                        operation_duration_histogram=self.duration_histogram,
+                        token_usage_histogram=getattr(self, "token_usage_histogram", None),
+                        operation_duration_histogram=getattr(self, "duration_histogram", None),
                         duration=duration,
                     )
                     if (
@@ -1238,7 +1243,9 @@ class ChatTelemetryLayer(Generic[OptionsCoT]):
 
             # Register a weak reference callback to close the span if stream is garbage collected
             # without being consumed. This ensures spans don't leak if users don't consume streams.
-            wrapped_stream = result_stream.with_cleanup_hook(_record_duration).with_cleanup_hook(_finalize_stream)
+            wrapped_stream: ResponseStream[ChatResponseUpdate, ChatResponse[Any]] = result_stream.with_cleanup_hook(
+                _record_duration
+            ).with_cleanup_hook(_finalize_stream)
             weakref.finalize(wrapped_stream, _close_span)
             return wrapped_stream
 
@@ -1253,7 +1260,15 @@ class ChatTelemetryLayer(Generic[OptionsCoT]):
                     )
                 start_time_stamp = perf_counter()
                 try:
-                    response = await super_get_response(messages=messages, stream=False, options=opts, **kwargs)
+                    response = cast(
+                        ChatResponse[Any],
+                        await super_get_response(
+                            messages=messages,
+                            stream=False,
+                            options=opts,
+                            **kwargs,
+                        ),
+                    )
                 except Exception as exception:
                     capture_exception(span=span, exception=exception, timestamp=time_ns())
                     raise
@@ -1262,16 +1277,20 @@ class ChatTelemetryLayer(Generic[OptionsCoT]):
                 _capture_response(
                     span=span,
                     attributes=response_attributes,
-                    token_usage_histogram=self.token_usage_histogram,
-                    operation_duration_histogram=self.duration_histogram,
+                    token_usage_histogram=getattr(self, "token_usage_histogram", None),
+                    operation_duration_histogram=getattr(self, "duration_histogram", None),
                     duration=duration,
                 )
                 if OBSERVABILITY_SETTINGS.SENSITIVE_DATA_ENABLED and response.messages:
+                    finish_reason = cast(
+                        "FinishReason | None",
+                        response.finish_reason if response.finish_reason in FINISH_REASON_MAP else None,
+                    )
                     _capture_messages(
                         span=span,
                         provider_name=provider_name,
                         messages=response.messages,
-                        finish_reason=response.finish_reason,
+                        finish_reason=finish_reason,
                         output=True,
                     )
                 return response  # type: ignore[return-value,no-any-return]
@@ -1302,8 +1321,10 @@ class EmbeddingTelemetryLayer(Generic[EmbeddingInputT, EmbeddingT, EmbeddingOpti
         values: Sequence[EmbeddingInputT],
         *,
         options: EmbeddingOptionsT | None = None,
-    ) -> GeneratedEmbeddings[EmbeddingT]:
+    ) -> GeneratedEmbeddings[EmbeddingT, EmbeddingOptionsT]:
         """Trace embedding generation with OpenTelemetry spans and metrics."""
+        from ._types import GeneratedEmbeddings  # type: ignore[reportUnusedImport]
+
         global OBSERVABILITY_SETTINGS
         super_get_embeddings = super().get_embeddings  # type: ignore[misc]
 
@@ -1311,7 +1332,7 @@ class EmbeddingTelemetryLayer(Generic[EmbeddingInputT, EmbeddingT, EmbeddingOpti
             return await super_get_embeddings(values, options=options)  # type: ignore[no-any-return]
 
         opts: dict[str, Any] = options or {}  # type: ignore[assignment]
-        provider_name = str(self.otel_provider_name)
+        provider_name = str(getattr(self, "otel_provider_name", "unknown"))
         model_id = opts.get("model_id") or getattr(self, "model_id", None) or "unknown"
         service_url_func = getattr(self, "service_url", None)
         service_url = str(service_url_func() if callable(service_url_func) else "unknown")
@@ -1325,14 +1346,18 @@ class EmbeddingTelemetryLayer(Generic[EmbeddingInputT, EmbeddingT, EmbeddingOpti
         with _get_span(attributes=attributes, span_name_attribute=OtelAttr.REQUEST_MODEL) as span:
             start_time_stamp = perf_counter()
             try:
-                result = await super_get_embeddings(values, options=options)
+                result = cast(
+                    GeneratedEmbeddings[EmbeddingT, EmbeddingOptionsT],
+                    await super_get_embeddings(values, options=options),
+                )
             except Exception as exception:
                 capture_exception(span=span, exception=exception, timestamp=time_ns())
                 raise
             duration = perf_counter() - start_time_stamp
             response_attributes: dict[str, Any] = {**attributes}
-            if result.usage and "prompt_tokens" in result.usage:
-                response_attributes[OtelAttr.INPUT_TOKENS] = result.usage["prompt_tokens"]
+            usage = result.usage or {}
+            if (input_tokens := usage.get("input_token_count")) is not None:
+                response_attributes[OtelAttr.INPUT_TOKENS] = input_tokens
             _capture_response(
                 span=span,
                 attributes=response_attributes,
@@ -1391,7 +1416,12 @@ class AgentTelemetryLayer:
     ) -> Awaitable[AgentResponse[Any]] | ResponseStream[AgentResponseUpdate, AgentResponse[Any]]:
         """Trace agent runs with OpenTelemetry spans and metrics."""
         global OBSERVABILITY_SETTINGS
-        super_run = super().run  # type: ignore[misc]
+        from ._types import ResponseStream, merge_chat_options
+
+        super_run = cast(
+            "Callable[..., Awaitable[AgentResponse[Any]] | ResponseStream[AgentResponseUpdate, AgentResponse[Any]]]",
+            super().run,  # type: ignore[misc]
+        )
         provider_name = str(self.otel_provider_name)
         capture_usage = bool(getattr(self, "_otel_capture_usage", True))
 
@@ -1403,8 +1433,6 @@ class AgentTelemetryLayer:
                 **kwargs,
             )
 
-        from ._types import ResponseStream, merge_chat_options
-
         default_options = getattr(self, "default_options", {})
         options = kwargs.get("options")
         merged_options: dict[str, Any] = merge_chat_options(default_options, options or {})
@@ -1420,16 +1448,16 @@ class AgentTelemetryLayer:
         )
 
         if stream:
-            run_result = super_run(
+            run_result: object = super_run(
                 messages=messages,
                 stream=True,
                 session=session,
                 **kwargs,
             )
             if isinstance(run_result, ResponseStream):
-                result_stream = run_result
+                result_stream: ResponseStream[AgentResponseUpdate, AgentResponse[Any]] = run_result  # pyright: ignore[reportUnknownVariableType]
             elif isinstance(run_result, Awaitable):
-                result_stream = ResponseStream.from_awaitable(run_result)
+                result_stream = ResponseStream.from_awaitable(run_result)  # type: ignore[arg-type]  # pyright: ignore[reportArgumentType]
             else:
                 raise RuntimeError("Streaming telemetry requires a ResponseStream result.")
 
@@ -1466,7 +1494,7 @@ class AgentTelemetryLayer:
                 from ._types import AgentResponse
 
                 try:
-                    response = await result_stream.get_final_response()
+                    response: AgentResponse[Any] = await result_stream.get_final_response()
                     duration = duration_state.get("duration")
                     response_attributes = _get_response_attributes(
                         attributes,
@@ -1492,7 +1520,9 @@ class AgentTelemetryLayer:
 
             # Register a weak reference callback to close the span if stream is garbage collected
             # without being consumed. This ensures spans don't leak if users don't consume streams.
-            wrapped_stream = result_stream.with_cleanup_hook(_record_duration).with_cleanup_hook(_finalize_stream)
+            wrapped_stream: ResponseStream[AgentResponseUpdate, AgentResponse[Any]] = result_stream.with_cleanup_hook(
+                _record_duration
+            ).with_cleanup_hook(_finalize_stream)
             weakref.finalize(wrapped_stream, _close_span)
             return wrapped_stream
 
@@ -1507,7 +1537,7 @@ class AgentTelemetryLayer:
                     )
                 start_time_stamp = perf_counter()
                 try:
-                    response = await super_run(
+                    response: AgentResponse[Any] = await super_run(
                         messages=messages,
                         stream=False,
                         session=session,
@@ -1598,12 +1628,17 @@ def _get_span(
         yield current_span
 
 
-def _get_instructions_from_options(options: Any) -> str | None:
+def _get_instructions_from_options(options: Any) -> str | list[str] | None:
     """Extract instructions from options dict."""
     if options is None:
         return None
-    if isinstance(options, dict):
-        return options.get("instructions")
+    if isinstance(options, Mapping):
+        instructions = cast(Mapping[str, Any], options).get("instructions")
+        if isinstance(instructions, str):
+            return instructions
+        if isinstance(instructions, list) and all(isinstance(item, str) for item in instructions):  # type: ignore[reportUnknownVariableType]
+            return instructions  # type: ignore[reportUnknownVariableType]
+        return None
     return None
 
 
@@ -1662,8 +1697,7 @@ def _get_span_attributes(**kwargs: Any) -> dict[str, Any]:
     """Get the span attributes from a kwargs dictionary."""
     attributes: dict[str, Any] = {}
     options = kwargs.get("all_options", kwargs.get("options"))
-    if options is not None and not isinstance(options, dict):
-        options = None
+    options_mapping = cast(Mapping[str, Any], options) if isinstance(options, Mapping) else None
 
     for source_keys, (otel_key, transform_func, check_options, default_value) in OTEL_ATTR_MAP.items():
         # Normalize to tuple of keys
@@ -1671,8 +1705,8 @@ def _get_span_attributes(**kwargs: Any) -> dict[str, Any]:
 
         value = None
         for key in keys:
-            if check_options and options is not None:
-                value = options.get(key)
+            if check_options and options_mapping is not None:
+                value = options_mapping.get(key)
             if value is None:
                 value = kwargs.get(key)
             if value is not None:
@@ -1743,7 +1777,7 @@ def _to_otel_message(message: Message) -> dict[str, Any]:
 
 def _to_otel_part(content: Content) -> dict[str, Any] | None:
     """Create a otel representation of a Content."""
-    from ._types import _get_data_bytes_as_str
+    from ._types import _get_data_bytes_as_str  # pyright: ignore[reportPrivateUsage]
 
     match content.type:
         case "text":
@@ -1798,10 +1832,12 @@ def _get_response_attributes(
     if model_id := getattr(response, "model_id", None):
         attributes[OtelAttr.RESPONSE_MODEL] = model_id
     if capture_usage and (usage := response.usage_details):
-        if usage.get("input_token_count"):
-            attributes[OtelAttr.INPUT_TOKENS] = usage["input_token_count"]
-        if usage.get("output_token_count"):
-            attributes[OtelAttr.OUTPUT_TOKENS] = usage["output_token_count"]
+        input_tokens = usage.get("input_token_count")
+        if input_tokens:
+            attributes[OtelAttr.INPUT_TOKENS] = input_tokens
+        output_tokens = usage.get("output_token_count")
+        if output_tokens:
+            attributes[OtelAttr.OUTPUT_TOKENS] = output_tokens
     return attributes
 
 

python/packages/core/agent_framework/openai/_assistant_provider.py

@@ -3,7 +3,7 @@
 from __future__ import annotations
 
 import sys
-from collections.abc import Awaitable, Callable, MutableMapping, Sequence
+from collections.abc import Awaitable, Callable, Mapping, MutableMapping, Sequence
 from typing import TYPE_CHECKING, Any, Generic, cast
 
 from openai import AsyncOpenAI
@@ -149,24 +149,25 @@ class OpenAIAssistantProvider(Generic[OptionsCoT]):
                 env_file_encoding=env_file_encoding,
             )
 
-            if not settings["api_key"]:
+            api_key_setting = settings.get("api_key")
+            if not api_key_setting:
                 raise ValueError(
                     "OpenAI API key is required. Set via 'api_key' parameter or 'OPENAI_API_KEY' environment variable."
                 )
 
             # Get API key value
-            api_key_value: str | Callable[[], str | Awaitable[str]] | None
-            if isinstance(settings["api_key"], SecretString):
-                api_key_value = settings["api_key"].get_secret_value()
+            api_key_value: str | Callable[[], str | Awaitable[str]]
+            if isinstance(api_key_setting, SecretString):
+                api_key_value = api_key_setting.get_secret_value()
             else:
-                api_key_value = settings["api_key"]
+                api_key_value = api_key_setting
 
             # Create client
             client_args: dict[str, Any] = {"api_key": api_key_value}
-            if settings["org_id"]:
-                client_args["organization"] = settings["org_id"]
-            if settings["base_url"]:
-                client_args["base_url"] = settings["base_url"]
+            if org_id_value := settings.get("org_id"):
+                client_args["organization"] = org_id_value
+            if base_url_value := settings.get("base_url"):
+                client_args["base_url"] = base_url_value
 
             self._client = AsyncOpenAI(**client_args)
 
@@ -250,7 +251,9 @@ class OpenAIAssistantProvider(Generic[OptionsCoT]):
         """
         # Normalize tools
         normalized_tools = normalize_tools(tools)
-        assistant_tools = [tool for tool in normalized_tools if isinstance(tool, (FunctionTool, MutableMapping))]
+        assistant_tools: list[FunctionTool | MutableMapping[str, Any]] = [
+            tool for tool in normalized_tools if isinstance(tool, (FunctionTool, MutableMapping))
+        ]
         api_tools = to_assistant_tools(assistant_tools) if assistant_tools else []
 
         # Extract response_format from default_options if present
@@ -287,7 +290,7 @@ class OpenAIAssistantProvider(Generic[OptionsCoT]):
         if not self._client:
             raise RuntimeError("OpenAI client is not initialized.")
 
-        assistant = await self._client.beta.assistants.create(**create_params)
+        assistant = await self._client.beta.assistants.create(**create_params)  # type: ignore[reportDeprecated]
 
         # Create Agent - pass default_options which contains response_format
         return self._create_chat_agent_from_assistant(
@@ -353,7 +356,7 @@ class OpenAIAssistantProvider(Generic[OptionsCoT]):
         if not self._client:
             raise RuntimeError("OpenAI client is not initialized.")
 
-        assistant = await self._client.beta.assistants.retrieve(assistant_id)
+        assistant = await self._client.beta.assistants.retrieve(assistant_id)  # type: ignore[reportDeprecated]
 
         # Use as_agent to wrap it
         return self.as_agent(
@@ -466,12 +469,14 @@ class OpenAIAssistantProvider(Generic[OptionsCoT]):
             for tool in normalized:
                 if isinstance(tool, FunctionTool):
                     provided_functions.add(tool.name)
-                elif isinstance(tool, MutableMapping) and "function" in tool:
-                    func_spec = tool.get("function", {})
-                    if isinstance(func_spec, dict):
-                        func_dict = cast(dict[str, Any], func_spec)
-                        if "name" in func_dict:
-                            provided_functions.add(str(func_dict["name"]))
+                elif isinstance(tool, Mapping):
+                    typed_tool = cast(Mapping[str, Any], tool)
+                    raw_func_spec = typed_tool.get("function")
+                    if isinstance(raw_func_spec, Mapping):
+                        typed_func_spec = cast(Mapping[str, Any], raw_func_spec)
+                        raw_name = typed_func_spec.get("name")
+                        if isinstance(raw_name, str) and raw_name:
+                            provided_functions.add(raw_name)
 
         # Check for missing functions
         missing = required_functions - provided_functions

python/packages/core/agent_framework/openai/_assistants_client.py

@@ -360,23 +360,26 @@ class OpenAIAssistantsClient(  # type: ignore[misc]
             env_file_encoding=env_file_encoding,
         )
 
-        if not async_client and not openai_settings["api_key"]:
+        api_key_value = openai_settings.get("api_key")
+        if not async_client and not api_key_value:
             raise ValueError(
                 "OpenAI API key is required. Set via 'api_key' parameter or 'OPENAI_API_KEY' environment variable."
             )
-        if not openai_settings["chat_model_id"]:
+
+        chat_model_id = openai_settings.get("chat_model_id")
+        if not chat_model_id:
             raise ValueError(
                 "OpenAI model ID is required. "
                 "Set via 'model_id' parameter or 'OPENAI_CHAT_MODEL_ID' environment variable."
             )
 
         super().__init__(
-            model_id=openai_settings["chat_model_id"],
-            api_key=self._get_api_key(openai_settings["api_key"]),
-            org_id=openai_settings["org_id"],
+            model_id=chat_model_id,
+            api_key=self._get_api_key(api_key_value),
+            org_id=openai_settings.get("org_id"),
             default_headers=default_headers,
             client=async_client,
-            base_url=openai_settings["base_url"],
+            base_url=openai_settings.get("base_url"),
             middleware=middleware,
             function_invocation_configuration=function_invocation_configuration,
         )
@@ -403,7 +406,7 @@ class OpenAIAssistantsClient(  # type: ignore[misc]
         """Clean up any assistants we created."""
         if self._should_delete_assistant and self.assistant_id is not None:
             client = await self._ensure_client()
-            await client.beta.assistants.delete(self.assistant_id)
+            await client.beta.assistants.delete(self.assistant_id)  # type: ignore[reportDeprecated]
             object.__setattr__(self, "assistant_id", None)
             object.__setattr__(self, "_should_delete_assistant", False)
 
@@ -466,7 +469,7 @@ class OpenAIAssistantsClient(  # type: ignore[misc]
                 raise ValueError("Parameter 'model_id' is required for assistant creation.")
 
             client = await self._ensure_client()
-            created_assistant = await client.beta.assistants.create(
+            created_assistant = await client.beta.assistants.create(  # type: ignore[reportDeprecated]
                 model=self.model_id,
                 description=self.assistant_description,
                 name=self.assistant_name,
@@ -568,7 +571,8 @@ class OpenAIAssistantsClient(  # type: ignore[misc]
                         if isinstance(delta_block, TextDeltaBlock) and delta_block.text and delta_block.text.value:
                             text_content = Content.from_text(delta_block.text.value)
                             if delta_block.text.annotations:
-                                text_content.annotations = []
+                                annotations: list[Annotation] = []
+                                text_content.annotations = annotations
                                 for annotation in delta_block.text.annotations:
                                     if isinstance(annotation, FileCitationDeltaAnnotation):
                                         ann: Annotation = Annotation(
@@ -589,7 +593,7 @@ class OpenAIAssistantsClient(  # type: ignore[misc]
                                                     end_index=annotation.end_index,
                                                 )
                                             ]
-                                        text_content.annotations.append(ann)
+                                        annotations.append(ann)
                                     elif isinstance(annotation, FilePathDeltaAnnotation):
                                         ann = Annotation(
                                             type="citation",
@@ -609,7 +613,7 @@ class OpenAIAssistantsClient(  # type: ignore[misc]
                                                     end_index=annotation.end_index,
                                                 )
                                             ]
-                                        text_content.annotations.append(ann)
+                                        annotations.append(ann)
                             yield ChatResponseUpdate(
                                 role=role,  # type: ignore[arg-type]
                                 contents=[text_content],
@@ -628,7 +632,8 @@ class OpenAIAssistantsClient(  # type: ignore[misc]
                             continue
                         text_content = Content.from_text(block.text.value)
                         if block.text.annotations:
-                            text_content.annotations = []
+                            completed_annotations: list[Annotation] = []
+                            text_content.annotations = completed_annotations
                             for completed_annotation in block.text.annotations:
                                 if isinstance(completed_annotation, FileCitationAnnotation):
                                     props: dict[str, Any] = {
@@ -644,17 +649,13 @@ class OpenAIAssistantsClient(  # type: ignore[misc]
                                         and completed_annotation.file_citation.file_id
                                     ):
                                         ann["file_id"] = completed_annotation.file_citation.file_id
-                                    if (
-                                        completed_annotation.start_index is not None
-                                        and completed_annotation.end_index is not None
-                                    ):
-                                        ann["annotated_regions"] = [
-                                            TextSpanRegion(
-                                                type="text_span",
-                                                start_index=completed_annotation.start_index,
-                                                end_index=completed_annotation.end_index,
-                                            )
-                                        ]
+                                    ann["annotated_regions"] = [
+                                        TextSpanRegion(
+                                            type="text_span",
+                                            start_index=completed_annotation.start_index,
+                                            end_index=completed_annotation.end_index,
+                                        )
+                                    ]
                                     text_content.annotations.append(ann)
                                 elif isinstance(completed_annotation, FilePathAnnotation):
                                     ann = Annotation(
@@ -666,17 +667,13 @@ class OpenAIAssistantsClient(  # type: ignore[misc]
                                     )
                                     if completed_annotation.file_path and completed_annotation.file_path.file_id:
                                         ann["file_id"] = completed_annotation.file_path.file_id
-                                    if (
-                                        completed_annotation.start_index is not None
-                                        and completed_annotation.end_index is not None
-                                    ):
-                                        ann["annotated_regions"] = [
-                                            TextSpanRegion(
-                                                type="text_span",
-                                                start_index=completed_annotation.start_index,
-                                                end_index=completed_annotation.end_index,
-                                            )
-                                        ]
+                                    ann["annotated_regions"] = [
+                                        TextSpanRegion(
+                                            type="text_span",
+                                            start_index=completed_annotation.start_index,
+                                            end_index=completed_annotation.end_index,
+                                        )
+                                    ]
                                     text_content.annotations.append(ann)
                                 else:
                                     logger.debug("Unparsed annotation type: %s", completed_annotation.type)
@@ -823,15 +820,16 @@ class OpenAIAssistantsClient(  # type: ignore[misc]
                 tool_definitions.append(tool.to_json_schema_spec())  # type: ignore[reportUnknownArgumentType]
             elif isinstance(tool, MutableMapping):
                 # Pass through dict-based tools directly (from static factory methods)
-                tool_definitions.append(tool)
+                tool_definitions.append(cast(MutableMapping[str, Any], tool))
 
         if len(tool_definitions) > 0:
             run_options["tools"] = tool_definitions
 
         if tool_mode is not None:
-            if (mode := tool_mode["mode"]) == "required" and (
-                func_name := tool_mode.get("required_function_name")
-            ) is not None:
+            mode = tool_mode.get("mode")
+            if mode is None:
+                raise ValueError("tool_choice mode is required")
+            if mode == "required" and (func_name := tool_mode.get("required_function_name")) is not None:
                 run_options["tool_choice"] = {
                     "type": "function",
                     "function": {"name": func_name},

python/packages/core/agent_framework/openai/_chat_client.py

@@ -15,7 +15,7 @@ from collections.abc import (
 )
 from datetime import datetime, timezone
 from itertools import chain
-from typing import Any, Generic, Literal
+from typing import Any, Generic, Literal, cast
 
 from openai import AsyncOpenAI, BadRequestError
 from openai.lib._parsing._completions import type_to_response_format_param
@@ -301,11 +301,16 @@ class RawOpenAIChatClient(  # type: ignore[misc]
         for tool in normalize_tools(tools):
             if isinstance(tool, FunctionTool):
                 chat_tools.append(tool.to_json_schema_spec())
-            elif isinstance(tool, MutableMapping) and tool.get("type") == "web_search":
-                # Web search is handled via web_search_options, not tools array
-                web_search_options = {k: v for k, v in tool.items() if k != "type"}
+            elif isinstance(tool, MutableMapping):
+                typed_tool = cast(MutableMapping[str, Any], tool)
+                if typed_tool.get("type") == "web_search":
+                    # Web search is handled via web_search_options, not tools array
+                    web_search_options = {k: v for k, v in typed_tool.items() if k != "type"}
+                else:
+                    # Pass through all other dict-based tools unchanged
+                    chat_tools.append(typed_tool)
             else:
-                # Pass through all other tools (dicts, SDK types) unchanged
+                # Pass through all other tools (SDK types) unchanged
                 chat_tools.append(tool)
         result: dict[str, Any] = {}
         if chat_tools:
@@ -608,10 +613,21 @@ class RawOpenAIChatClient(  # type: ignore[misc]
         # See https://github.com/microsoft/agent-framework/issues/4084
         for msg in all_messages:
             msg_content: Any = msg.get("content")
-            if isinstance(msg_content, list) and all(
-                isinstance(c, dict) and c.get("type") == "text" for c in msg_content
-            ):
-                msg["content"] = "\n".join(c.get("text", "") for c in msg_content)
+            if isinstance(msg_content, list):
+                typed_msg_content = cast(list[object], msg_content)
+                text_items: list[Mapping[str, Any]] = []
+                for item in typed_msg_content:
+                    if not isinstance(item, Mapping):
+                        break
+                    text_item = cast(Mapping[str, Any], item)
+                    if text_item.get("type") != "text":
+                        break
+                    text_items.append(text_item)
+                else:
+                    msg["content"] = "\n".join(
+                        text_item.get("text", "") if isinstance(text_item.get("text", ""), str) else ""
+                        for text_item in text_items
+                    )
 
         return all_messages
 
@@ -775,21 +791,26 @@ class OpenAIChatClient(  # type: ignore[misc]
             env_file_encoding=env_file_encoding,
         )
 
-        if not async_client and not openai_settings["api_key"]:
+        api_key_value = openai_settings.get("api_key")
+        if not async_client and not api_key_value:
             raise ValueError(
                 "OpenAI API key is required. Set via 'api_key' parameter or 'OPENAI_API_KEY' environment variable."
             )
-        if not openai_settings["chat_model_id"]:
+
+        chat_model_id = openai_settings.get("chat_model_id")
+        if not chat_model_id:
             raise ValueError(
                 "OpenAI model ID is required. "
                 "Set via 'model_id' parameter or 'OPENAI_CHAT_MODEL_ID' environment variable."
             )
 
+        base_url_value = openai_settings.get("base_url")
+
         super().__init__(
-            model_id=openai_settings["chat_model_id"],
-            api_key=self._get_api_key(openai_settings["api_key"]),
-            base_url=openai_settings["base_url"] if openai_settings["base_url"] else None,
-            org_id=openai_settings["org_id"],
+            model_id=chat_model_id,
+            api_key=self._get_api_key(api_key_value),
+            base_url=base_url_value if base_url_value else None,
+            org_id=openai_settings.get("org_id"),
             default_headers=default_headers,
             client=async_client,
             instruction_role=instruction_role,

python/packages/core/agent_framework/openai/_embedding_client.py

@@ -67,7 +67,7 @@ class RawOpenAIEmbeddingClient(
         values: Sequence[str],
         *,
         options: OpenAIEmbeddingOptionsT | None = None,
-    ) -> GeneratedEmbeddings[list[float]]:
+    ) -> GeneratedEmbeddings[list[float], OpenAIEmbeddingOptionsT]:
         """Call the OpenAI embeddings API.
 
         Args:
@@ -81,9 +81,9 @@ class RawOpenAIEmbeddingClient(
             ValueError: If model_id is not provided or values is empty.
         """
         if not values:
-            return GeneratedEmbeddings([], options=options)
+            return GeneratedEmbeddings([], options=options)  # type: ignore
 
-        opts: dict[str, Any] = dict(options) if options else {}
+        opts: dict[str, Any] = options or {}  # type: ignore
         model = opts.get("model_id") or self.model_id
         if not model:
             raise ValueError("model_id is required")
@@ -193,21 +193,26 @@ class OpenAIEmbeddingClient(
             env_file_encoding=env_file_encoding,
         )
 
-        if not async_client and not openai_settings["api_key"]:
+        api_key_value = openai_settings.get("api_key")
+        if not async_client and not api_key_value:
             raise ValueError(
                 "OpenAI API key is required. Set via 'api_key' parameter or 'OPENAI_API_KEY' environment variable."
             )
-        if not openai_settings["embedding_model_id"]:
+
+        embedding_model_id = openai_settings.get("embedding_model_id")
+        if not embedding_model_id:
             raise ValueError(
                 "OpenAI embedding model ID is required. "
                 "Set via 'model_id' parameter or 'OPENAI_EMBEDDING_MODEL_ID' environment variable."
             )
 
+        base_url_value = openai_settings.get("base_url")
+
         super().__init__(
-            model_id=openai_settings["embedding_model_id"],
-            api_key=self._get_api_key(openai_settings["api_key"]),
-            base_url=openai_settings["base_url"] if openai_settings["base_url"] else None,
-            org_id=openai_settings["org_id"],
+            model_id=embedding_model_id,
+            api_key=self._get_api_key(api_key_value),
+            base_url=base_url_value if base_url_value else None,
+            org_id=openai_settings.get("org_id"),
             default_headers=default_headers,
             client=async_client,
             otel_provider_name=otel_provider_name,

python/packages/core/agent_framework/openai/_responses_client.py

@@ -460,14 +460,13 @@ class RawOpenAIResponsesClient(  # type: ignore[misc]
         for tool_item in tools_list:
             if isinstance(tool_item, FunctionTool) and tool_item.kind == SHELL_TOOL_KIND_VALUE:
                 shell_env = (tool_item.additional_properties or {}).get(OPENAI_SHELL_ENVIRONMENT_KEY)
-                if isinstance(shell_env, Mapping):
-                    response_tools.append(
-                        FunctionShellTool(
-                            type="shell",
-                            environment=dict(shell_env),
-                        )
+                response_tools.append(
+                    FunctionShellTool(
+                        type="shell",
+                        environment=shell_env,  # type: ignore[typeddict-item]
                     )
-                    continue
+                )
+                continue
             if isinstance(tool_item, FunctionTool):
                 params = tool_item.parameters()
                 params["additionalProperties"] = False
@@ -496,7 +495,7 @@ class RawOpenAIResponsesClient(  # type: ignore[misc]
             if tool_item.kind != SHELL_TOOL_KIND_VALUE:
                 continue
             shell_env = (tool_item.additional_properties or {}).get(OPENAI_SHELL_ENVIRONMENT_KEY)
-            if isinstance(shell_env, Mapping) and shell_env.get("type") == "local":
+            if isinstance(shell_env, Mapping) and shell_env.get("type") == "local":  # type: ignore[typeddict-item]
                 return tool_item.name
         return None
 
@@ -714,7 +713,7 @@ class RawOpenAIResponsesClient(  # type: ignore[misc]
             )
             if env_config.get("type") == "local":
                 raise ValueError("Local shell requires func. Provide func for local execution.")
-            return FunctionShellTool(type="shell", environment=env_config)
+            return FunctionShellTool(type="shell", environment=env_config)  # type: ignore[typeddict-item]
 
         if isinstance(environment, dict):
             raise ValueError("When func is provided, environment config is not supported.")
@@ -1226,7 +1225,7 @@ class RawOpenAIResponsesClient(  # type: ignore[misc]
         """Convert function tool output to the local shell JSON payload format."""
         payload: dict[str, Any]
         if isinstance(content.result, Mapping):
-            payload = dict(content.result)
+            payload = dict(content.result)  # type: ignore[assignment]
         else:
             payload = {
                 "stdout": "" if content.result is None else str(content.result),
@@ -1242,7 +1241,7 @@ class RawOpenAIResponsesClient(  # type: ignore[misc]
         """Convert function tool output to shell_call_output payload format."""
         payload: dict[str, Any]
         if isinstance(content.result, Mapping):
-            payload = dict(content.result)
+            payload = dict(content.result)  # type: ignore[assignment]
         else:
             payload = {
                 "stdout": "" if content.result is None else str(content.result),
@@ -1252,8 +1251,8 @@ class RawOpenAIResponsesClient(  # type: ignore[misc]
 
         # Pass through native payload shape when tool already returns shell output entries.
         direct_output = payload.get("output")
-        if isinstance(direct_output, list) and all(isinstance(item, Mapping) for item in direct_output):
-            return [dict(item) for item in direct_output]
+        if isinstance(direct_output, list) and all(isinstance(item, Mapping) for item in direct_output):  # type: ignore[reportUnknownMemberType]
+            return [dict(item) for item in direct_output]  # type: ignore[reportUnknownMemberType]
 
         stdout = str(payload.get("stdout", ""))
         stderr = str(payload.get("stderr", ""))
@@ -2293,24 +2292,26 @@ class OpenAIResponsesClient(  # type: ignore[misc]
             env_file_encoding=env_file_encoding,
         )
 
-        if not async_client and not openai_settings["api_key"]:
+        api_key_setting = openai_settings.get("api_key")
+        if not async_client and not api_key_setting:
             raise ValueError(
                 "OpenAI API key is required. Set via 'api_key' parameter or 'OPENAI_API_KEY' environment variable."
             )
-        if not openai_settings["responses_model_id"]:
+        responses_model_id = openai_settings.get("responses_model_id")
+        if not responses_model_id:
             raise ValueError(
                 "OpenAI model ID is required. "
                 "Set via 'model_id' parameter or 'OPENAI_RESPONSES_MODEL_ID' environment variable."
             )
 
         super().__init__(
-            model_id=openai_settings["responses_model_id"],
-            api_key=self._get_api_key(openai_settings["api_key"]),
-            org_id=openai_settings["org_id"],
+            model_id=responses_model_id,
+            api_key=self._get_api_key(api_key_setting),
+            org_id=openai_settings.get("org_id"),
             default_headers=default_headers,
             client=async_client,
             instruction_role=instruction_role,
-            base_url=openai_settings["base_url"],
+            base_url=openai_settings.get("base_url"),
             middleware=middleware,
             function_invocation_configuration=function_invocation_configuration,
             **kwargs,

python/packages/core/agent_framework/openai/_shared.py

@@ -6,7 +6,7 @@ import logging
 import sys
 from collections.abc import Awaitable, Callable, Mapping, MutableMapping, Sequence
 from copy import copy
-from typing import Any, ClassVar, Union
+from typing import Any, ClassVar, Union, cast
 
 import openai
 from openai import (
@@ -332,8 +332,10 @@ def from_assistant_tools(
     for tool in assistant_tools:
         if hasattr(tool, "type"):
             tool_type = tool.type
-        elif isinstance(tool, dict):
-            tool_type = tool.get("type")
+        elif isinstance(tool, Mapping):
+            typed_tool = cast(Mapping[str, Any], tool)
+            tool_type_value: Any = typed_tool.get("type")
+            tool_type = tool_type_value if isinstance(tool_type_value, str) else None
         else:
             tool_type = None
 

python/packages/core/pyproject.toml

@@ -104,11 +104,12 @@ extend = "../../pyproject.toml"
 
 [tool.pyright]
 extends = "../../pyproject.toml"
-include = ["tests/workflow"]
+include = ["agent_framework", "tests/workflow"]
 
 [tool.mypy]
 plugins = ['pydantic.mypy']
 strict = true
+incremental = false
 python_version = "3.10"
 ignore_missing_imports = true
 disallow_untyped_defs = true
@@ -130,7 +131,7 @@ include = "../../shared_tasks.toml"
 
 [tool.poe.tasks]
 mypy = "mypy --config-file $POE_ROOT/pyproject.toml agent_framework"
-test = "pytest --cov=agent_framework --cov-report=term-missing:skip-covered -n auto --dist worksteal tests"
+test = "pytest -m \"not integration\" --cov=agent_framework --cov-report=term-missing:skip-covered -n auto --dist worksteal tests"
 
 [tool.flit.module]
 name = "agent_framework"

python/packages/core/tests/core/test_skills.py

@@ -10,7 +10,7 @@ from unittest.mock import AsyncMock
 
 import pytest
 
-from agent_framework import Skill, SkillResource, SkillsProvider, SessionContext
+from agent_framework import SessionContext, Skill, SkillResource, SkillsProvider
 from agent_framework._skills import (
     DEFAULT_RESOURCE_EXTENSIONS,
     _create_instructions,
@@ -1348,9 +1348,7 @@ class TestReadAndParseSkillFile:
     def test_valid_file(self, tmp_path: Path) -> None:
         skill_dir = tmp_path / "my-skill"
         skill_dir.mkdir()
-        (skill_dir / "SKILL.md").write_text(
-            "---\nname: my-skill\ndescription: A skill.\n---\nBody.", encoding="utf-8"
-        )
+        (skill_dir / "SKILL.md").write_text("---\nname: my-skill\ndescription: A skill.\n---\nBody.", encoding="utf-8")
         result = _read_and_parse_skill_file(str(skill_dir))
         assert result is not None
         name, desc, content = result
@@ -1393,7 +1391,7 @@ class TestCreateResourceElement:
     def test_xml_escapes_name(self) -> None:
         r = SkillResource(name='ref"special', content="data")
         elem = _create_resource_element(r)
-        assert '"' in elem
+        assert """ in elem
 
     def test_xml_escapes_description(self) -> None:
         r = SkillResource(name="ref", description='Uses <tags> & "quotes"', content="data")

python/packages/core/tests/core/test_tools.py

@@ -5,7 +5,7 @@ from unittest.mock import Mock
 import pytest
 from opentelemetry import trace
 from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
-from pydantic import BaseModel, ValidationError
+from pydantic import BaseModel
 
 from agent_framework import (
     Content,
@@ -13,7 +13,6 @@ from agent_framework import (
     tool,
 )
 from agent_framework._tools import (
-    _build_pydantic_model_from_json_schema,
     _parse_annotation,
     _parse_inputs,
 )
@@ -1001,467 +1000,4 @@ def test_parse_annotation_with_annotated_and_literal():
     assert get_args(literal_type) == ("A", "B", "C")
 
 
-def test_build_pydantic_model_from_json_schema_array_of_objects_issue():
-    """Test for Tools with complex input schema (array of objects).
-
-    This test verifies that JSON schemas with array properties containing nested objects
-    are properly parsed, ensuring that the nested object schema is preserved
-    and not reduced to a bare dict.
-
-    Example from issue:
-    ```
-    const SalesOrderItemSchema = z.object({
-        customerMaterialNumber: z.string().optional(),
-        quantity: z.number(),
-        unitOfMeasure: z.string()
-    });
-
-    const CreateSalesOrderInputSchema = z.object({
-        contract: z.string(),
-        items: z.array(SalesOrderItemSchema)
-    });
-    ```
-
-    The issue was that agents only saw:
-    ```
-    {"contract": "str", "items": "list[dict]"}
-    ```
-
-    Instead of the proper nested schema with all fields.
-    """
-    # Schema matching the issue description
-    schema = {
-        "type": "object",
-        "properties": {
-            "contract": {"type": "string", "description": "Reference contract number"},
-            "items": {
-                "type": "array",
-                "description": "Sales order line items",
-                "items": {
-                    "type": "object",
-                    "properties": {
-                        "customerMaterialNumber": {
-                            "type": "string",
-                            "description": "Customer's material number",
-                        },
-                        "quantity": {"type": "number", "description": "Order quantity"},
-                        "unitOfMeasure": {
-                            "type": "string",
-                            "description": "Unit of measure (e.g., 'ST', 'KG', 'TO')",
-                        },
-                    },
-                    "required": ["quantity", "unitOfMeasure"],
-                },
-            },
-        },
-        "required": ["contract", "items"],
-    }
-
-    model = _build_pydantic_model_from_json_schema("create_sales_order", schema)
-
-    # Test valid data
-    valid_data = {
-        "contract": "CONTRACT-123",
-        "items": [
-            {
-                "customerMaterialNumber": "MAT-001",
-                "quantity": 10,
-                "unitOfMeasure": "ST",
-            },
-            {"quantity": 5.5, "unitOfMeasure": "KG"},
-        ],
-    }
-
-    instance = model(**valid_data)
-
-    # Verify the data was parsed correctly
-    assert instance.contract == "CONTRACT-123"
-    assert len(instance.items) == 2
-
-    # Verify first item
-    assert instance.items[0].customerMaterialNumber == "MAT-001"
-    assert instance.items[0].quantity == 10
-    assert instance.items[0].unitOfMeasure == "ST"
-
-    # Verify second item (optional field not provided)
-    assert instance.items[1].quantity == 5.5
-    assert instance.items[1].unitOfMeasure == "KG"
-
-    # Verify that items are proper BaseModel instances, not bare dicts
-    assert isinstance(instance.items[0], BaseModel)
-    assert isinstance(instance.items[1], BaseModel)
-
-    # Verify that the nested object has the expected fields
-    assert hasattr(instance.items[0], "customerMaterialNumber")
-    assert hasattr(instance.items[0], "quantity")
-    assert hasattr(instance.items[0], "unitOfMeasure")
-
-    # CRITICAL: Validate using the same methods that actual chat clients use
-    # This is what would actually be sent to the LLM
-
-    # Create a FunctionTool wrapper to access the client-facing APIs
-    def dummy_func(**kwargs):
-        return kwargs
-
-    test_func = FunctionTool(
-        func=dummy_func,
-        name="create_sales_order",
-        description="Create a sales order",
-        input_model=model,
-    )
-
-    # Test 1: Anthropic client uses tool.parameters() directly
-    anthropic_schema = test_func.parameters()
-
-    # Verify contract property
-    assert "contract" in anthropic_schema["properties"]
-    assert anthropic_schema["properties"]["contract"]["type"] == "string"
-
-    # Verify items array property exists
-    assert "items" in anthropic_schema["properties"]
-    items_prop = anthropic_schema["properties"]["items"]
-    assert items_prop["type"] == "array"
-
-    # THE KEY TEST for Anthropic: array items must have proper object schema
-    assert "items" in items_prop, "Array should have 'items' schema definition"
-    array_items_schema = items_prop["items"]
-
-    # Resolve schema if using $ref
-    if "$ref" in array_items_schema:
-        ref_path = array_items_schema["$ref"]
-        assert ref_path.startswith("#/$defs/") or ref_path.startswith("#/definitions/")
-        ref_name = ref_path.split("/")[-1]
-        defs = anthropic_schema.get("$defs", anthropic_schema.get("definitions", {}))
-        assert ref_name in defs, f"Referenced schema '{ref_name}' should exist"
-        item_schema = defs[ref_name]
-    else:
-        item_schema = array_items_schema
-
-    # Verify the nested object has all properties defined
-    assert "properties" in item_schema, "Array items should have properties (not bare dict)"
-    item_properties = item_schema["properties"]
-
-    # All three fields must be present in schema sent to LLM
-    assert "customerMaterialNumber" in item_properties, "customerMaterialNumber missing from LLM schema"
-    assert "quantity" in item_properties, "quantity missing from LLM schema"
-    assert "unitOfMeasure" in item_properties, "unitOfMeasure missing from LLM schema"
-
-    # Verify types are correct
-    assert item_properties["customerMaterialNumber"]["type"] == "string"
-    assert item_properties["quantity"]["type"] in ["number", "integer"]
-    assert item_properties["unitOfMeasure"]["type"] == "string"
-
-    # Test 2: OpenAI client uses tool.to_json_schema_spec()
-    openai_spec = test_func.to_json_schema_spec()
-
-    assert openai_spec["type"] == "function"
-    assert "function" in openai_spec
-    openai_schema = openai_spec["function"]["parameters"]
-
-    # Verify the same structure is present in OpenAI format
-    assert "items" in openai_schema["properties"]
-    openai_items_prop = openai_schema["properties"]["items"]
-    assert openai_items_prop["type"] == "array"
-    assert "items" in openai_items_prop
-
-    openai_array_items = openai_items_prop["items"]
-    if "$ref" in openai_array_items:
-        ref_path = openai_array_items["$ref"]
-        ref_name = ref_path.split("/")[-1]
-        defs = openai_schema.get("$defs", openai_schema.get("definitions", {}))
-        openai_item_schema = defs[ref_name]
-    else:
-        openai_item_schema = openai_array_items
-
-    assert "properties" in openai_item_schema
-    openai_props = openai_item_schema["properties"]
-    assert "customerMaterialNumber" in openai_props
-    assert "quantity" in openai_props
-    assert "unitOfMeasure" in openai_props
-
-    # Test validation - missing required quantity
-    with pytest.raises(ValidationError):
-        model(
-            contract="CONTRACT-456",
-            items=[
-                {
-                    "customerMaterialNumber": "MAT-002",
-                    "unitOfMeasure": "TO",
-                    # Missing required 'quantity'
-                }
-            ],
-        )
-
-    # Test validation - missing required unitOfMeasure
-    with pytest.raises(ValidationError):
-        model(
-            contract="CONTRACT-789",
-            items=[
-                {
-                    "quantity": 20
-                    # Missing required 'unitOfMeasure'
-                }
-            ],
-        )
-
-
-def test_one_of_discriminator_polymorphism():
-    """Test that oneOf with discriminator creates proper polymorphic union types.
-
-    Tests that oneOf + discriminator patterns are properly converted to Pydantic discriminated unions.
-    """
-    schema = {
-        "$defs": {
-            "CreateProject": {
-                "description": "Action: Create an Azure DevOps project.",
-                "properties": {
-                    "name": {
-                        "const": "create_project",
-                        "default": "create_project",
-                        "type": "string",
-                    },
-                    "params": {"$ref": "#/$defs/CreateProjectParams"},
-                },
-                "required": ["params"],
-                "type": "object",
-            },
-            "CreateProjectParams": {
-                "description": "Parameters for the create_project action.",
-                "properties": {
-                    "orgUrl": {"minLength": 1, "type": "string"},
-                    "projectName": {"minLength": 1, "type": "string"},
-                    "description": {"default": "", "type": "string"},
-                    "template": {"default": "Agile", "type": "string"},
-                    "sourceControl": {
-                        "default": "Git",
-                        "enum": ["Git", "Tfvc"],
-                        "type": "string",
-                    },
-                    "visibility": {"default": "private", "type": "string"},
-                },
-                "required": ["orgUrl", "projectName"],
-                "type": "object",
-            },
-            "DeployRequest": {
-                "description": "Request to deploy Azure DevOps resources.",
-                "properties": {
-                    "projectName": {"minLength": 1, "type": "string"},
-                    "organization": {"minLength": 1, "type": "string"},
-                    "actions": {
-                        "items": {
-                            "discriminator": {
-                                "mapping": {
-                                    "create_project": "#/$defs/CreateProject",
-                                    "hello_world": "#/$defs/HelloWorld",
-                                },
-                                "propertyName": "name",
-                            },
-                            "oneOf": [
-                                {"$ref": "#/$defs/HelloWorld"},
-                                {"$ref": "#/$defs/CreateProject"},
-                            ],
-                        },
-                        "type": "array",
-                    },
-                },
-                "required": ["projectName", "organization"],
-                "type": "object",
-            },
-            "HelloWorld": {
-                "description": "Action: Prints a greeting message.",
-                "properties": {
-                    "name": {
-                        "const": "hello_world",
-                        "default": "hello_world",
-                        "type": "string",
-                    },
-                    "params": {"$ref": "#/$defs/HelloWorldParams"},
-                },
-                "required": ["params"],
-                "type": "object",
-            },
-            "HelloWorldParams": {
-                "description": "Parameters for the hello_world action.",
-                "properties": {
-                    "name": {
-                        "description": "Name to greet",
-                        "minLength": 1,
-                        "type": "string",
-                    }
-                },
-                "required": ["name"],
-                "type": "object",
-            },
-        },
-        "properties": {"params": {"$ref": "#/$defs/DeployRequest"}},
-        "required": ["params"],
-        "type": "object",
-    }
-
-    # Build the model
-    model = _build_pydantic_model_from_json_schema("deploy_tool", schema)
-
-    # Verify the model structure
-    assert model is not None
-    assert issubclass(model, BaseModel)
-
-    # Test with HelloWorld action
-    hello_world_data = {
-        "params": {
-            "projectName": "MyProject",
-            "organization": "MyOrg",
-            "actions": [
-                {
-                    "name": "hello_world",
-                    "params": {"name": "Alice"},
-                }
-            ],
-        }
-    }
-
-    instance = model(**hello_world_data)
-    assert instance.params.projectName == "MyProject"
-    assert instance.params.organization == "MyOrg"
-    assert len(instance.params.actions) == 1
-    assert instance.params.actions[0].name == "hello_world"
-    assert instance.params.actions[0].params.name == "Alice"
-
-    # Test with CreateProject action
-    create_project_data = {
-        "params": {
-            "projectName": "MyProject",
-            "organization": "MyOrg",
-            "actions": [
-                {
-                    "name": "create_project",
-                    "params": {
-                        "orgUrl": "https://dev.azure.com/myorg",
-                        "projectName": "NewProject",
-                        "sourceControl": "Git",
-                    },
-                }
-            ],
-        }
-    }
-
-    instance2 = model(**create_project_data)
-    assert instance2.params.actions[0].name == "create_project"
-    assert instance2.params.actions[0].params.projectName == "NewProject"
-    assert instance2.params.actions[0].params.sourceControl == "Git"
-
-    # Test with mixed actions
-    mixed_data = {
-        "params": {
-            "projectName": "MyProject",
-            "organization": "MyOrg",
-            "actions": [
-                {"name": "hello_world", "params": {"name": "Bob"}},
-                {
-                    "name": "create_project",
-                    "params": {
-                        "orgUrl": "https://dev.azure.com/myorg",
-                        "projectName": "AnotherProject",
-                    },
-                },
-            ],
-        }
-    }
-
-    instance3 = model(**mixed_data)
-    assert len(instance3.params.actions) == 2
-    assert instance3.params.actions[0].name == "hello_world"
-    assert instance3.params.actions[1].name == "create_project"
-
-
-def test_const_creates_literal():
-    """Test that const in JSON Schema creates Literal type."""
-    schema = {
-        "properties": {
-            "action": {
-                "const": "create",
-                "type": "string",
-                "description": "Action type",
-            },
-            "value": {"type": "integer"},
-        },
-        "required": ["action", "value"],
-    }
-
-    model = _build_pydantic_model_from_json_schema("test_const", schema)
-
-    # Verify valid const value works
-    instance = model(action="create", value=42)
-    assert instance.action == "create"
-    assert instance.value == 42
-
-    # Verify incorrect const value fails
-    with pytest.raises(ValidationError):
-        model(action="delete", value=42)
-
-
-def test_enum_creates_literal():
-    """Test that enum in JSON Schema creates Literal type."""
-    schema = {
-        "properties": {
-            "status": {
-                "enum": ["pending", "approved", "rejected"],
-                "type": "string",
-                "description": "Status",
-            },
-            "priority": {"enum": [1, 2, 3], "type": "integer"},
-        },
-        "required": ["status"],
-    }
-
-    model = _build_pydantic_model_from_json_schema("test_enum", schema)
-
-    # Verify valid enum values work
-    instance = model(status="approved", priority=2)
-    assert instance.status == "approved"
-    assert instance.priority == 2
-
-    # Verify invalid enum value fails
-    with pytest.raises(ValidationError):
-        model(status="unknown")
-
-    with pytest.raises(ValidationError):
-        model(status="pending", priority=5)
-
-
-def test_nested_object_with_const_and_enum():
-    """Test that const and enum work in nested objects."""
-    schema = {
-        "properties": {
-            "config": {
-                "type": "object",
-                "properties": {
-                    "type": {
-                        "const": "production",
-                        "default": "production",
-                        "type": "string",
-                    },
-                    "level": {"enum": ["low", "medium", "high"], "type": "string"},
-                },
-                "required": ["level"],
-            }
-        },
-        "required": ["config"],
-    }
-
-    model = _build_pydantic_model_from_json_schema("test_nested", schema)
-
-    # Valid data
-    instance = model(config={"type": "production", "level": "high"})
-    assert instance.config.type == "production"
-    assert instance.config.level == "high"
-
-    # Invalid const in nested object
-    with pytest.raises(ValidationError):
-        model(config={"type": "development", "level": "low"})
-
-    # Invalid enum in nested object
-    with pytest.raises(ValidationError):
-        model(config={"type": "production", "level": "critical"})
-
-
 # endregion

python/packages/core/tests/core/test_types.py

@@ -550,7 +550,6 @@ def test_usage_details():
     assert usage["input_token_count"] == 5
     assert usage["output_token_count"] == 10
     assert usage["total_token_count"] == 15
-    assert usage.get("additional_counts", {}) == {}
 
 
 def test_usage_details_addition():
@@ -581,8 +580,8 @@ def test_usage_details_addition():
 def test_usage_details_fail():
     # TypedDict doesn't validate types at runtime, so this test no longer applies
     # Creating UsageDetails with wrong types won't raise ValueError
-    usage = UsageDetails(input_token_count=5, output_token_count=10, total_token_count=15, wrong_type="42.923")  # type: ignore[typeddict-item]
-    assert usage["wrong_type"] == "42.923"  # type: ignore[typeddict-item]
+    usage = UsageDetails(input_token_count=5, output_token_count=10, total_token_count=15, wrong_type="42.923")
+    assert usage["wrong_type"] == "42.923"
 
 
 def test_usage_details_additional_counts():
@@ -601,6 +600,15 @@ def test_usage_details_add_with_none_and_type_errors():
     # TypedDict doesn't support + operator, use add_usage_details
 
 
+def test_usage_details_add_skips_non_int():
+    u1 = UsageDetails(input_token_count=10, other="test")
+    u2 = UsageDetails(input_token_count=10, another="test")
+    u3 = add_usage_details(u1, u2)
+    assert len(u3.keys()) == 1
+    assert "input_token_count" in u3
+    assert u3["input_token_count"] == 20
+
+
 # region UserInputRequest and Response
 
 
@@ -1705,7 +1713,7 @@ def test_chat_response_complex_serialization():
             {"role": "user", "contents": [{"type": "text", "text": "Hello"}]},
             {"role": "assistant", "contents": [{"type": "text", "text": "Hi there"}]},
         ],
-        "finish_reason": {"value": "stop"},
+        "finish_reason": "stop",
         "usage_details": {
             "type": "usage_details",
             "input_token_count": 5,
@@ -1831,7 +1839,7 @@ def test_agent_run_response_update_all_content_types():
             },
             {"type": "text_reasoning", "text": "reasoning"},
         ],
-        "role": {"value": "assistant"},  # Test role as dict
+        "role": "assistant",  # Test role as dict
     }
 
     update = AgentResponseUpdate.from_dict(update_data)
@@ -2394,7 +2402,7 @@ def test_content_add_usage_content_non_integer_values():
     result = usage1 + usage2
 
     # Non-integer "model" should take first non-None value
-    assert result.usage_details["model"] == "gpt-4"
+    assert "model" not in result.usage_details
     # Integer "count" should be summed
     assert result.usage_details["count"] == 30
 

python/packages/core/tests/openai/test_openai_embedding_client.py

@@ -212,7 +212,8 @@ def test_azure_construction_with_existing_client() -> None:
     assert client.client is mock_client
 
 
-def test_azure_construction_missing_deployment_name_raises() -> None:
+def test_azure_construction_missing_deployment_name_raises(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.delenv("AZURE_OPENAI_EMBEDDING_DEPLOYMENT_NAME", raising=False)
     with pytest.raises(ValueError, match="deployment name is required"):
         AzureOpenAIEmbeddingClient(
             api_key="test-key",
@@ -272,6 +273,7 @@ skip_if_azure_openai_integration_tests_disabled = pytest.mark.skipif(
 
 @skip_if_openai_integration_tests_disabled
 @pytest.mark.flaky
+@pytest.mark.integration
 async def test_integration_openai_get_embeddings() -> None:
     """End-to-end test of OpenAI embedding generation."""
     client = OpenAIEmbeddingClient(model_id="text-embedding-3-small")
@@ -289,6 +291,7 @@ async def test_integration_openai_get_embeddings() -> None:
 
 @skip_if_openai_integration_tests_disabled
 @pytest.mark.flaky
+@pytest.mark.integration
 async def test_integration_openai_get_embeddings_multiple() -> None:
     """Test embedding generation for multiple inputs."""
     client = OpenAIEmbeddingClient(model_id="text-embedding-3-small")
@@ -302,6 +305,7 @@ async def test_integration_openai_get_embeddings_multiple() -> None:
 
 @skip_if_openai_integration_tests_disabled
 @pytest.mark.flaky
+@pytest.mark.integration
 async def test_integration_openai_get_embeddings_with_dimensions() -> None:
     """Test embedding generation with custom dimensions."""
     client = OpenAIEmbeddingClient(model_id="text-embedding-3-small")
@@ -315,6 +319,7 @@ async def test_integration_openai_get_embeddings_with_dimensions() -> None:
 
 @skip_if_azure_openai_integration_tests_disabled
 @pytest.mark.flaky
+@pytest.mark.integration
 async def test_integration_azure_openai_get_embeddings() -> None:
     """End-to-end test of Azure OpenAI embedding generation."""
     client = AzureOpenAIEmbeddingClient()
@@ -332,6 +337,7 @@ async def test_integration_azure_openai_get_embeddings() -> None:
 
 @skip_if_azure_openai_integration_tests_disabled
 @pytest.mark.flaky
+@pytest.mark.integration
 async def test_integration_azure_openai_get_embeddings_multiple() -> None:
     """Test Azure OpenAI embedding generation for multiple inputs."""
     client = AzureOpenAIEmbeddingClient()
@@ -345,6 +351,7 @@ async def test_integration_azure_openai_get_embeddings_multiple() -> None:
 
 @skip_if_azure_openai_integration_tests_disabled
 @pytest.mark.flaky
+@pytest.mark.integration
 async def test_integration_azure_openai_get_embeddings_with_dimensions() -> None:
     """Test Azure OpenAI embedding generation with custom dimensions."""
     client = AzureOpenAIEmbeddingClient()

python/packages/core/tests/workflow/test_agent_executor.py

@@ -5,6 +5,7 @@ from collections.abc import AsyncIterable, Awaitable
 from typing import TYPE_CHECKING, Any, Literal, overload
 
 import pytest
+
 from agent_framework import (
     AgentExecutor,
     AgentResponse,
@@ -59,30 +60,19 @@ class _CountingAgent(BaseAgent):
         stream: bool = False,
         session: AgentSession | None = None,
         **kwargs: Any,
-    ) -> (
-        Awaitable[AgentResponse[Any]]
-        | ResponseStream[AgentResponseUpdate, AgentResponse[Any]]
-    ):
+    ) -> Awaitable[AgentResponse[Any]] | ResponseStream[AgentResponseUpdate, AgentResponse[Any]]:
         self.call_count += 1
         if stream:
 
             async def _stream() -> AsyncIterable[AgentResponseUpdate]:
                 yield AgentResponseUpdate(
-                    contents=[
-                        Content.from_text(
-                            text=f"Response #{self.call_count}: {self.name}"
-                        )
-                    ]
+                    contents=[Content.from_text(text=f"Response #{self.call_count}: {self.name}")]
                 )
 
             return ResponseStream(_stream(), finalizer=AgentResponse.from_updates)
 
         async def _run() -> AgentResponse:
-            return AgentResponse(
-                messages=[
-                    Message("assistant", [f"Response #{self.call_count}: {self.name}"])
-                ]
-            )
+            return AgentResponse(messages=[Message("assistant", [f"Response #{self.call_count}: {self.name}"])])
 
         return _run()
 
@@ -120,10 +110,7 @@ class _StreamingHookAgent(BaseAgent):
         stream: bool = False,
         session: AgentSession | None = None,
         **kwargs: Any,
-    ) -> (
-        Awaitable[AgentResponse[Any]]
-        | ResponseStream[AgentResponseUpdate, AgentResponse[Any]]
-    ):
+    ) -> Awaitable[AgentResponse[Any]] | ResponseStream[AgentResponseUpdate, AgentResponse[Any]]:
         if stream:
 
             async def _stream() -> AsyncIterable[AgentResponseUpdate]:
@@ -138,9 +125,9 @@ class _StreamingHookAgent(BaseAgent):
                 self.result_hook_called = True
                 return response
 
-            return ResponseStream(
-                _stream(), finalizer=AgentResponse.from_updates
-            ).with_result_hook(_mark_result_hook_called)
+            return ResponseStream(_stream(), finalizer=AgentResponse.from_updates).with_result_hook(
+                _mark_result_hook_called
+            )
 
         async def _run() -> AgentResponse:
             return AgentResponse(messages=[Message("assistant", ["hook test"])])
@@ -148,9 +135,7 @@ class _StreamingHookAgent(BaseAgent):
         return _run()
 
 
-async def test_agent_executor_streaming_finalizes_stream_and_runs_result_hooks() -> (
-    None
-):
+async def test_agent_executor_streaming_finalizes_stream_and_runs_result_hooks() -> None:
     """AgentExecutor should call get_final_response() so stream result hooks execute."""
     agent = _StreamingHookAgent(id="hook_agent", name="HookAgent")
     executor = AgentExecutor(agent, id="hook_exec")
@@ -217,9 +202,7 @@ async def test_agent_executor_checkpoint_stores_and_restores_state() -> None:
 
     executor_state = executor_states[executor.id]  # type: ignore[index]
     assert "cache" in executor_state, "Checkpoint should store executor cache state"
-    assert "agent_session" in executor_state, (
-        "Checkpoint should store executor session state"
-    )
+    assert "agent_session" in executor_state, "Checkpoint should store executor session state"
 
     # Verify session state structure
     session_state = executor_state["agent_session"]  # type: ignore[index]
@@ -240,15 +223,11 @@ async def test_agent_executor_checkpoint_stores_and_restores_state() -> None:
     assert restored_agent.call_count == 0
 
     # Build new workflow with the restored executor
-    wf_resume = SequentialBuilder(
-        participants=[restored_executor], checkpoint_storage=storage
-    ).build()
+    wf_resume = SequentialBuilder(participants=[restored_executor], checkpoint_storage=storage).build()
 
     # Resume from checkpoint
     resumed_output: AgentExecutorResponse | None = None
-    async for ev in wf_resume.run(
-        checkpoint_id=restore_checkpoint.checkpoint_id, stream=True
-    ):
+    async for ev in wf_resume.run(checkpoint_id=restore_checkpoint.checkpoint_id, stream=True):
         if ev.type == "output":
             resumed_output = ev.data  # type: ignore[assignment]
         if ev.type == "status" and ev.state in (
@@ -391,11 +370,7 @@ async def test_prepare_agent_run_args_strips_all_reserved_kwargs_at_once(
     assert options is not None
     assert options["additional_function_arguments"]["custom"] == 1
 
-    warned_keys = {
-        r.message.split("'")[1]
-        for r in caplog.records
-        if "reserved" in r.message.lower()
-    }
+    warned_keys = {r.message.split("'")[1] for r in caplog.records if "reserved" in r.message.lower()}
     assert warned_keys == {"session", "stream", "messages"}
 
 

python/packages/core/tests/workflow/test_agent_utils.py

@@ -16,10 +16,31 @@ class MockAgent:
         self.description: str | None = None
 
     @overload
-    def run(self, messages: AgentRunInputs | None = ..., *, stream: Literal[False] = ..., session: AgentSession | None = ..., **kwargs: Any) -> Awaitable[AgentResponse[Any]]: ...
+    def run(
+        self,
+        messages: AgentRunInputs | None = ...,
+        *,
+        stream: Literal[False] = ...,
+        session: AgentSession | None = ...,
+        **kwargs: Any,
+    ) -> Awaitable[AgentResponse[Any]]: ...
     @overload
-    def run(self, messages: AgentRunInputs | None = ..., *, stream: Literal[True], session: AgentSession | None = ..., **kwargs: Any) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: ...
-    def run(self, messages: AgentRunInputs | None = None, *, stream: bool = False, session: AgentSession | None = None, **kwargs: Any) -> Awaitable[AgentResponse[Any]] | ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: ...
+    def run(
+        self,
+        messages: AgentRunInputs | None = ...,
+        *,
+        stream: Literal[True],
+        session: AgentSession | None = ...,
+        **kwargs: Any,
+    ) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: ...
+    def run(
+        self,
+        messages: AgentRunInputs | None = None,
+        *,
+        stream: bool = False,
+        session: AgentSession | None = None,
+        **kwargs: Any,
+    ) -> Awaitable[AgentResponse[Any]] | ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: ...
 
     def create_session(self, **kwargs: Any) -> AgentSession:
         """Creates a new conversation session for the agent."""

python/packages/core/tests/workflow/test_edge.py

@@ -4,9 +4,8 @@ from dataclasses import dataclass
 from typing import Any
 from unittest.mock import patch
 
-from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
-
 import pytest
+from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
 
 from agent_framework import (
     Executor,

python/packages/core/tests/workflow/test_executor.py

@@ -3,6 +3,8 @@
 from dataclasses import dataclass
 
 import pytest
+from typing_extensions import Never
+
 from agent_framework import (
     Executor,
     Message,
@@ -14,7 +16,6 @@ from agent_framework import (
     handler,
     response_handler,
 )
-from typing_extensions import Never
 
 
 # Module-level types for string forward reference tests
@@ -155,11 +156,7 @@ async def test_executor_invoked_event_contains_input_data():
     workflow = WorkflowBuilder(start_executor=upper).add_edge(upper, collector).build()
 
     events = await workflow.run("hello world")
-    invoked_events = [
-        e
-        for e in events
-        if isinstance(e, WorkflowEvent) and e.type == "executor_invoked"
-    ]
+    invoked_events = [e for e in events if isinstance(e, WorkflowEvent) and e.type == "executor_invoked"]
 
     assert len(invoked_events) == 2
 
@@ -193,16 +190,10 @@ async def test_executor_completed_event_contains_sent_messages():
     sender = MultiSenderExecutor(id="sender")
     collector = CollectorExecutor(id="collector")
 
-    workflow = (
-        WorkflowBuilder(start_executor=sender).add_edge(sender, collector).build()
-    )
+    workflow = WorkflowBuilder(start_executor=sender).add_edge(sender, collector).build()
 
     events = await workflow.run("hello")
-    completed_events = [
-        e
-        for e in events
-        if isinstance(e, WorkflowEvent) and e.type == "executor_completed"
-    ]
+    completed_events = [e for e in events if isinstance(e, WorkflowEvent) and e.type == "executor_completed"]
 
     # Sender should have completed with the sent messages
     sender_completed = next(e for e in completed_events if e.executor_id == "sender")
@@ -210,9 +201,7 @@ async def test_executor_completed_event_contains_sent_messages():
     assert sender_completed.data == ["hello-first", "hello-second"]
 
     # Collector should have completed with no sent messages (None)
-    collector_completed_events = [
-        e for e in completed_events if e.executor_id == "collector"
-    ]
+    collector_completed_events = [e for e in completed_events if e.executor_id == "collector"]
     # Collector is called twice (once per message from sender)
     assert len(collector_completed_events) == 2
     for collector_completed in collector_completed_events:
@@ -231,11 +220,7 @@ async def test_executor_completed_event_includes_yielded_outputs():
     workflow = WorkflowBuilder(start_executor=executor).build()
 
     events = await workflow.run("test")
-    completed_events = [
-        e
-        for e in events
-        if isinstance(e, WorkflowEvent) and e.type == "executor_completed"
-    ]
+    completed_events = [e for e in events if isinstance(e, WorkflowEvent) and e.type == "executor_completed"]
 
     assert len(completed_events) == 1
     assert completed_events[0].executor_id == "yielder"
@@ -263,9 +248,7 @@ async def test_executor_events_with_complex_message_types():
 
     class ProcessorExecutor(Executor):
         @handler
-        async def handle(
-            self, request: Request, ctx: WorkflowContext[Response]
-        ) -> None:
+        async def handle(self, request: Request, ctx: WorkflowContext[Response]) -> None:
             response = Response(results=[request.query.upper()] * request.limit)
             await ctx.send_message(response)
 
@@ -277,23 +260,13 @@ async def test_executor_events_with_complex_message_types():
     processor = ProcessorExecutor(id="processor")
     collector = CollectorExecutor(id="collector")
 
-    workflow = (
-        WorkflowBuilder(start_executor=processor).add_edge(processor, collector).build()
-    )
+    workflow = WorkflowBuilder(start_executor=processor).add_edge(processor, collector).build()
 
     input_request = Request(query="hello", limit=3)
     events = await workflow.run(input_request)
 
-    invoked_events = [
-        e
-        for e in events
-        if isinstance(e, WorkflowEvent) and e.type == "executor_invoked"
-    ]
-    completed_events = [
-        e
-        for e in events
-        if isinstance(e, WorkflowEvent) and e.type == "executor_completed"
-    ]
+    invoked_events = [e for e in events if isinstance(e, WorkflowEvent) and e.type == "executor_invoked"]
+    completed_events = [e for e in events if isinstance(e, WorkflowEvent) and e.type == "executor_completed"]
 
     # Check processor invoked event has the Request object
     processor_invoked = next(e for e in invoked_events if e.executor_id == "processor")
@@ -302,9 +275,7 @@ async def test_executor_events_with_complex_message_types():
     assert processor_invoked.data.limit == 3
 
     # Check processor completed event has the Response object
-    processor_completed = next(
-        e for e in completed_events if e.executor_id == "processor"
-    )
+    processor_completed = next(e for e in completed_events if e.executor_id == "processor")
     assert processor_completed.data is not None
     assert len(processor_completed.data) == 1
     assert isinstance(processor_completed.data[0], Response)
@@ -390,9 +361,7 @@ def test_executor_workflow_output_types_property():
     # Test executor with union workflow output types
     class UnionWorkflowOutputExecutor(Executor):
         @handler
-        async def handle(
-            self, text: str, ctx: WorkflowContext[int, str | bool]
-        ) -> None:
+        async def handle(self, text: str, ctx: WorkflowContext[int, str | bool]) -> None:
             pass
 
     executor = UnionWorkflowOutputExecutor(id="union_workflow_output")
@@ -403,15 +372,11 @@ def test_executor_workflow_output_types_property():
     # Test executor with multiple handlers having different workflow output types
     class MultiHandlerWorkflowExecutor(Executor):
         @handler
-        async def handle_string(
-            self, text: str, ctx: WorkflowContext[int, str]
-        ) -> None:
+        async def handle_string(self, text: str, ctx: WorkflowContext[int, str]) -> None:
             pass
 
         @handler
-        async def handle_number(
-            self, num: int, ctx: WorkflowContext[bool, float]
-        ) -> None:
+        async def handle_number(self, num: int, ctx: WorkflowContext[bool, float]) -> None:
             pass
 
     executor = MultiHandlerWorkflowExecutor(id="multi_workflow")
@@ -465,9 +430,7 @@ def test_executor_output_types_includes_response_handlers():
             pass
 
         @response_handler
-        async def handle_response(
-            self, original_request: str, response: bool, ctx: WorkflowContext[float]
-        ) -> None:
+        async def handle_response(self, original_request: str, response: bool, ctx: WorkflowContext[float]) -> None:
             pass
 
     executor = RequestResponseExecutor(id="request_response")
@@ -574,9 +537,7 @@ async def test_executor_invoked_event_data_not_mutated_by_handler():
     """Test that executor_invoked event (type='executor_invoked').data captures original input, not mutated input."""
 
     @executor(id="Mutator")
-    async def mutator(
-        messages: list[Message], ctx: WorkflowContext[list[Message]]
-    ) -> None:
+    async def mutator(messages: list[Message], ctx: WorkflowContext[list[Message]]) -> None:
         # The handler mutates the input list by appending new messages
         original_len = len(messages)
         messages.append(Message(role="assistant", text="Added by executor"))
@@ -591,11 +552,7 @@ async def test_executor_invoked_event_data_not_mutated_by_handler():
     events = await workflow.run(input_messages)
 
     # Find the invoked event for the Mutator executor
-    invoked_events = [
-        e
-        for e in events
-        if isinstance(e, WorkflowEvent) and e.type == "executor_invoked"
-    ]
+    invoked_events = [e for e in events if isinstance(e, WorkflowEvent) and e.type == "executor_invoked"]
     assert len(invoked_events) == 1
     mutator_invoked = invoked_events[0]
 
@@ -672,12 +629,8 @@ class TestHandlerExplicitTypes:
         assert handler_func._handler_spec["output_types"] == [list]  # pyright: ignore[reportFunctionMemberAccess]
 
         # Verify can_handle
-        assert exec_instance.can_handle(
-            WorkflowMessage(data={"key": "value"}, source_id="mock")
-        )
-        assert not exec_instance.can_handle(
-            WorkflowMessage(data="string", source_id="mock")
-        )
+        assert exec_instance.can_handle(WorkflowMessage(data={"key": "value"}, source_id="mock"))
+        assert not exec_instance.can_handle(WorkflowMessage(data="string", source_id="mock"))
 
     def test_handler_with_explicit_union_input_type(self):
         """Test that explicit union input_type is handled correctly."""
@@ -698,9 +651,7 @@ class TestHandlerExplicitTypes:
         assert exec_instance.can_handle(WorkflowMessage(data="hello", source_id="mock"))
         assert exec_instance.can_handle(WorkflowMessage(data=42, source_id="mock"))
         # Cannot handle float
-        assert not exec_instance.can_handle(
-            WorkflowMessage(data=3.14, source_id="mock")
-        )
+        assert not exec_instance.can_handle(WorkflowMessage(data=3.14, source_id="mock"))
 
     def test_handler_with_explicit_union_output_type(self):
         """Test that explicit union output is normalized to a list."""
@@ -776,9 +727,7 @@ class TestHandlerExplicitTypes:
 
             class OnlyWorkflowOutputExecutor(Executor):  # pyright: ignore[reportUnusedClass]
                 @handler(workflow_output=bool)
-                async def handle(
-                    self, message: str, ctx: WorkflowContext[int, str]
-                ) -> None:
+                async def handle(self, message: str, ctx: WorkflowContext[int, str]) -> None:
                     pass
 
     def test_handler_explicit_input_type_allows_no_message_annotation(self):
@@ -803,9 +752,7 @@ class TestHandlerExplicitTypes:
                 pass
 
             @handler
-            async def handle_introspected(
-                self, message: float, ctx: WorkflowContext[bool]
-            ) -> None:
+            async def handle_introspected(self, message: float, ctx: WorkflowContext[bool]) -> None:
                 pass
 
         exec_instance = MixedExecutor(id="mixed")
@@ -831,9 +778,7 @@ class TestHandlerExplicitTypes:
 
         # Should resolve the string to the actual type
         assert ForwardRefMessage in exec_instance._handlers  # pyright: ignore[reportPrivateUsage]
-        assert exec_instance.can_handle(
-            WorkflowMessage(data=ForwardRefMessage("hello"), source_id="mock")
-        )
+        assert exec_instance.can_handle(WorkflowMessage(data=ForwardRefMessage("hello"), source_id="mock"))
 
     def test_handler_with_string_forward_reference_union(self):
         """Test that string forward references work with union types."""
@@ -846,12 +791,8 @@ class TestHandlerExplicitTypes:
         exec_instance = StringUnionExecutor(id="string_union")
 
         # Should handle both types
-        assert exec_instance.can_handle(
-            WorkflowMessage(data=ForwardRefTypeA("hello"), source_id="mock")
-        )
-        assert exec_instance.can_handle(
-            WorkflowMessage(data=ForwardRefTypeB(42), source_id="mock")
-        )
+        assert exec_instance.can_handle(WorkflowMessage(data=ForwardRefTypeA("hello"), source_id="mock"))
+        assert exec_instance.can_handle(WorkflowMessage(data=ForwardRefTypeB(42), source_id="mock"))
 
     def test_handler_with_string_forward_reference_output_type(self):
         """Test that string forward references work for output_type."""
@@ -890,9 +831,7 @@ class TestHandlerExplicitTypes:
 
         class PrecedenceExecutor(Executor):
             @handler(input=int, output=float, workflow_output=str)
-            async def handle(
-                self, message: int, ctx: WorkflowContext[int, bool]
-            ) -> None:
+            async def handle(self, message: int, ctx: WorkflowContext[int, bool]) -> None:
                 pass
 
         exec_instance = PrecedenceExecutor(id="precedence")
@@ -958,9 +897,7 @@ class TestHandlerExplicitTypes:
             async def handle(self, message, ctx: WorkflowContext) -> None:  # type: ignore[no-untyped-def]
                 pass
 
-        exec_instance = StringUnionWorkflowOutputExecutor(
-            id="string_union_workflow_output"
-        )
+        exec_instance = StringUnionWorkflowOutputExecutor(id="string_union_workflow_output")
 
         # Should resolve both types from string union
         assert ForwardRefTypeA in exec_instance.workflow_output_types
@@ -971,14 +908,10 @@ class TestHandlerExplicitTypes:
 
         class IntrospectedWorkflowOutputExecutor(Executor):
             @handler
-            async def handle(
-                self, message: str, ctx: WorkflowContext[int, bool]
-            ) -> None:
+            async def handle(self, message: str, ctx: WorkflowContext[int, bool]) -> None:
                 pass
 
-        exec_instance = IntrospectedWorkflowOutputExecutor(
-            id="introspected_workflow_output"
-        )
+        exec_instance = IntrospectedWorkflowOutputExecutor(id="introspected_workflow_output")
 
         # Should use introspected types from WorkflowContext[int, bool]
         assert int in exec_instance.output_types

python/packages/core/tests/workflow/test_workflow_agent.py

@@ -717,9 +717,23 @@ class TestWorkflowAgent:
                 return AgentSession()
 
             @overload
-            def run(self, messages: str | Content | Message | Sequence[str | Content | Message] | None = ..., *, stream: Literal[False] = ..., session: AgentSession | None = ..., **kwargs: Any) -> Awaitable[AgentResponse[Any]]: ...
+            def run(
+                self,
+                messages: str | Content | Message | Sequence[str | Content | Message] | None = ...,
+                *,
+                stream: Literal[False] = ...,
+                session: AgentSession | None = ...,
+                **kwargs: Any,
+            ) -> Awaitable[AgentResponse[Any]]: ...
             @overload
-            def run(self, messages: str | Content | Message | Sequence[str | Content | Message] | None = ..., *, stream: Literal[True], session: AgentSession | None = ..., **kwargs: Any) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: ...
+            def run(
+                self,
+                messages: str | Content | Message | Sequence[str | Content | Message] | None = ...,
+                *,
+                stream: Literal[True],
+                session: AgentSession | None = ...,
+                **kwargs: Any,
+            ) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: ...
 
             def run(
                 self,
@@ -813,9 +827,23 @@ class TestWorkflowAgent:
                 return AgentSession()
 
             @overload
-            def run(self, messages: str | Content | Message | Sequence[str | Content | Message] | None = ..., *, stream: Literal[False] = ..., session: AgentSession | None = ..., **kwargs: Any) -> Awaitable[AgentResponse[Any]]: ...
+            def run(
+                self,
+                messages: str | Content | Message | Sequence[str | Content | Message] | None = ...,
+                *,
+                stream: Literal[False] = ...,
+                session: AgentSession | None = ...,
+                **kwargs: Any,
+            ) -> Awaitable[AgentResponse[Any]]: ...
             @overload
-            def run(self, messages: str | Content | Message | Sequence[str | Content | Message] | None = ..., *, stream: Literal[True], session: AgentSession | None = ..., **kwargs: Any) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: ...
+            def run(
+                self,
+                messages: str | Content | Message | Sequence[str | Content | Message] | None = ...,
+                *,
+                stream: Literal[True],
+                session: AgentSession | None = ...,
+                **kwargs: Any,
+            ) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: ...
 
             def run(
                 self,

python/packages/core/tests/workflow/test_workflow_kwargs.py

@@ -52,9 +52,23 @@ class _KwargsCapturingAgent(BaseAgent):
         self.captured_kwargs = []
 
     @overload
-    def run(self, messages: AgentRunInputs | None = ..., *, stream: Literal[False] = ..., session: AgentSession | None = ..., **kwargs: Any) -> Awaitable[AgentResponse[Any]]: ...
+    def run(
+        self,
+        messages: AgentRunInputs | None = ...,
+        *,
+        stream: Literal[False] = ...,
+        session: AgentSession | None = ...,
+        **kwargs: Any,
+    ) -> Awaitable[AgentResponse[Any]]: ...
     @overload
-    def run(self, messages: AgentRunInputs | None = ..., *, stream: Literal[True], session: AgentSession | None = ..., **kwargs: Any) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: ...
+    def run(
+        self,
+        messages: AgentRunInputs | None = ...,
+        *,
+        stream: Literal[True],
+        session: AgentSession | None = ...,
+        **kwargs: Any,
+    ) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: ...
 
     def run(
         self,
@@ -90,9 +104,23 @@ class _OptionsAwareAgent(BaseAgent):
         self.captured_kwargs = []
 
     @overload
-    def run(self, messages: AgentRunInputs | None = ..., *, stream: Literal[False] = ..., session: AgentSession | None = ..., **kwargs: Any) -> Awaitable[AgentResponse[Any]]: ...
+    def run(
+        self,
+        messages: AgentRunInputs | None = ...,
+        *,
+        stream: Literal[False] = ...,
+        session: AgentSession | None = ...,
+        **kwargs: Any,
+    ) -> Awaitable[AgentResponse[Any]]: ...
     @overload
-    def run(self, messages: AgentRunInputs | None = ..., *, stream: Literal[True], session: AgentSession | None = ..., **kwargs: Any) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: ...
+    def run(
+        self,
+        messages: AgentRunInputs | None = ...,
+        *,
+        stream: Literal[True],
+        session: AgentSession | None = ...,
+        **kwargs: Any,
+    ) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: ...
 
     def run(
         self,
@@ -475,9 +503,23 @@ async def test_kwargs_preserved_on_response_continuation() -> None:
             self._asked = False
 
         @overload
-        def run(self, messages: AgentRunInputs | None = ..., *, stream: Literal[False] = ..., session: AgentSession | None = ..., **kwargs: Any) -> Awaitable[AgentResponse[Any]]: ...
+        def run(
+            self,
+            messages: AgentRunInputs | None = ...,
+            *,
+            stream: Literal[False] = ...,
+            session: AgentSession | None = ...,
+            **kwargs: Any,
+        ) -> Awaitable[AgentResponse[Any]]: ...
         @overload
-        def run(self, messages: AgentRunInputs | None = ..., *, stream: Literal[True], session: AgentSession | None = ..., **kwargs: Any) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: ...
+        def run(
+            self,
+            messages: AgentRunInputs | None = ...,
+            *,
+            stream: Literal[True],
+            session: AgentSession | None = ...,
+            **kwargs: Any,
+        ) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: ...
 
         def run(
             self,
@@ -538,9 +580,23 @@ async def test_kwargs_overridden_on_response_continuation() -> None:
             self._asked = False
 
         @overload
-        def run(self, messages: AgentRunInputs | None = ..., *, stream: Literal[False] = ..., session: AgentSession | None = ..., **kwargs: Any) -> Awaitable[AgentResponse[Any]]: ...
+        def run(
+            self,
+            messages: AgentRunInputs | None = ...,
+            *,
+            stream: Literal[False] = ...,
+            session: AgentSession | None = ...,
+            **kwargs: Any,
+        ) -> Awaitable[AgentResponse[Any]]: ...
         @overload
-        def run(self, messages: AgentRunInputs | None = ..., *, stream: Literal[True], session: AgentSession | None = ..., **kwargs: Any) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: ...
+        def run(
+            self,
+            messages: AgentRunInputs | None = ...,
+            *,
+            stream: Literal[True],
+            session: AgentSession | None = ...,
+            **kwargs: Any,
+        ) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: ...
 
         def run(
             self,
@@ -605,9 +661,23 @@ async def test_kwargs_empty_value_passed_on_continuation() -> None:
             self._asked = False
 
         @overload
-        def run(self, messages: AgentRunInputs | None = ..., *, stream: Literal[False] = ..., session: AgentSession | None = ..., **kwargs: Any) -> Awaitable[AgentResponse[Any]]: ...
+        def run(
+            self,
+            messages: AgentRunInputs | None = ...,
+            *,
+            stream: Literal[False] = ...,
+            session: AgentSession | None = ...,
+            **kwargs: Any,
+        ) -> Awaitable[AgentResponse[Any]]: ...
         @overload
-        def run(self, messages: AgentRunInputs | None = ..., *, stream: Literal[True], session: AgentSession | None = ..., **kwargs: Any) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: ...
+        def run(
+            self,
+            messages: AgentRunInputs | None = ...,
+            *,
+            stream: Literal[True],
+            session: AgentSession | None = ...,
+            **kwargs: Any,
+        ) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: ...
 
         def run(
             self,

python/packages/core/tests/workflow/test_workflow_states.py

@@ -38,7 +38,9 @@ async def test_executor_failed_and_workflow_failed_events_streaming():
             events.append(ev)
 
     # executor_failed event (type='executor_failed') should be emitted before workflow failed event
-    executor_failed_events: list[WorkflowEvent[Any]] = [e for e in events if isinstance(e, WorkflowEvent) and e.type == "executor_failed"]
+    executor_failed_events: list[WorkflowEvent[Any]] = [
+        e for e in events if isinstance(e, WorkflowEvent) and e.type == "executor_failed"
+    ]
     assert executor_failed_events, "executor_failed event should be emitted when start executor fails"
     assert executor_failed_events[0].executor_id == "f"
     assert executor_failed_events[0].origin is WorkflowEventSource.FRAMEWORK
@@ -96,7 +98,9 @@ async def test_executor_failed_event_from_second_executor_in_chain():
             events.append(ev)
 
     # executor_failed event should be emitted for the failing executor
-    executor_failed_events: list[WorkflowEvent[Any]] = [e for e in events if isinstance(e, WorkflowEvent) and e.type == "executor_failed"]
+    executor_failed_events: list[WorkflowEvent[Any]] = [
+        e for e in events if isinstance(e, WorkflowEvent) and e.type == "executor_failed"
+    ]
     assert executor_failed_events, "executor_failed event should be emitted when second executor fails"
     assert executor_failed_events[0].executor_id == "failing"
     assert executor_failed_events[0].origin is WorkflowEventSource.FRAMEWORK