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/lab/gaia/agent_framework_lab_gaia/gaia.py

@@ -13,7 +13,7 @@ import time
 from collections.abc import Iterable
 from datetime import datetime
 from pathlib import Path
-from typing import Any
+from typing import Any, cast
 
 from opentelemetry.trace import NoOpTracer, SpanKind, get_tracer
 from tqdm import tqdm
@@ -163,7 +163,7 @@ def _normalize_str(s: str, remove_punct: bool = True) -> str:
     return no_spaces.lower()
 
 
-def gaia_scorer(model_answer: str, ground_truth: str) -> bool:
+def gaia_scorer(model_answer: str | None, ground_truth: str) -> bool:
     """Official GAIA scoring function.
 
     Args:
@@ -193,7 +193,7 @@ def gaia_scorer(model_answer: str, ground_truth: str) -> bool:
         ma_elems = _split_string(model_answer)
         if len(gt_elems) != len(ma_elems):
             return False
-        comparisons = []
+        comparisons: list[bool] = []
         for ma, gt in zip(ma_elems, gt_elems, strict=False):
             if is_float(gt):
                 comparisons.append(abs(_normalize_number_str(ma) - float(gt)) < 1e-6)
@@ -204,18 +204,39 @@ def gaia_scorer(model_answer: str, ground_truth: str) -> bool:
     return _normalize_str(model_answer) == _normalize_str(ground_truth)
 
 
+def _coerce_record(raw: object) -> dict[str, Any] | None:
+    if isinstance(raw, dict):
+        raw_dict = cast(dict[object, Any], raw)
+        if all(isinstance(key, str) for key in raw_dict):
+            return cast(dict[str, Any], raw_dict)
+    return None
+
+
+def _parse_level(level: object) -> int | None:
+    if isinstance(level, int):
+        return level
+    if isinstance(level, str) and level.isdigit():
+        return int(level)
+    return None
+
+
 def _read_jsonl(path: Path) -> Iterable[dict[str, Any]]:
     """Read JSONL file and yield parsed records."""
     with path.open("rb") as f:
         for line in f:
             if not line.strip():
                 continue
+            parsed: object
             try:
                 import orjson
 
-                yield orjson.loads(line)
+                parsed = orjson.loads(line)
             except Exception:
-                yield json.loads(line)
+                parsed = json.loads(line)
+
+            record = _coerce_record(parsed)
+            if record is not None:
+                yield record
 
 
 def _load_gaia_local(repo_dir: Path, wanted_levels: list[int] | None = None, max_n: int | None = None) -> list[Task]:
@@ -232,41 +253,43 @@ def _load_gaia_local(repo_dir: Path, wanted_levels: list[int] | None = None, max
         try:
             import pyarrow.parquet as pq
 
-            table = pq.read_table(p)
-            for row in table.to_pylist():
+            pq_any = cast(Any, pq)
+            table: Any = pq_any.read_table(p)
+            rows = cast(list[object], table.to_pylist())
+            for row in rows:
+                record = _coerce_record(row)
+                if record is None:
+                    continue
+
                 # Robustly extract fields used across variants
-                q = row.get("Question") or row.get("question") or row.get("query") or row.get("prompt")
-                ans = row.get("Final answer") or row.get("answer") or row.get("final_answer")
+                q_obj = record.get("Question") or record.get("question") or record.get("query") or record.get("prompt")
+                ans = record.get("Final answer") or record.get("answer") or record.get("final_answer")
+                if not isinstance(q_obj, str):
+                    continue
+                q = q_obj
+
                 qid = str(
-                    row.get("task_id")
-                    or row.get("question_id")
-                    or row.get("id")
-                    or row.get("uuid")
+                    record.get("task_id")
+                    or record.get("question_id")
+                    or record.get("id")
+                    or record.get("uuid")
                     or f"{p.stem}:{len(tasks)}"
                 )
-                lvl = row.get("Level") or row.get("level")
-
-                # Convert level to int if it's a string
-                def _parse_level(lvl: Any) -> int | None:
-                    """Parse level value to integer if possible."""
-                    if isinstance(lvl, int):
-                        return lvl
-                    if isinstance(lvl, str) and lvl.isdigit():
-                        return int(lvl)
-                    return None
-
-                lvl = _parse_level(lvl)
-                fname = row.get("file_name") or row.get("filename") or None
+                lvl = _parse_level(record.get("Level") or record.get("level"))
+                fname_obj = record.get("file_name") or record.get("filename")
+                fname = fname_obj if isinstance(fname_obj, str) else None
 
                 # Only evaluate examples with public answers (dev/validation split)
                 # Skip if no question, no answer, or answer is placeholder like "?"
-                if not q or ans is None or str(ans).strip() in ["?", ""]:
+                if ans is None or str(ans).strip() in ["?", ""]:
                     continue
 
                 if wanted_levels and (lvl not in wanted_levels):
                     continue
 
-                tasks.append(Task(task_id=qid, question=q, answer=str(ans), level=lvl, file_name=fname, metadata=row))
+                tasks.append(
+                    Task(task_id=qid, question=q, answer=str(ans), level=lvl, file_name=fname, metadata=record)
+                )
         except ImportError:
             print("Warning: pyarrow not installed. Install with: pip install pyarrow")
             continue
@@ -279,8 +302,12 @@ def _load_gaia_local(repo_dir: Path, wanted_levels: list[int] | None = None, max
         for p in repo_dir.rglob("metadata.jsonl"):
             for rec in _read_jsonl(p):
                 # Robustly extract fields used across variants
-                q = rec.get("Question") or rec.get("question") or rec.get("query") or rec.get("prompt")
+                q_obj = rec.get("Question") or rec.get("question") or rec.get("query") or rec.get("prompt")
                 ans = rec.get("Final answer") or rec.get("answer") or rec.get("final_answer")
+                if not isinstance(q_obj, str):
+                    continue
+                q = q_obj
+
                 qid = str(
                     rec.get("task_id")
                     or rec.get("question_id")
@@ -288,15 +315,13 @@ def _load_gaia_local(repo_dir: Path, wanted_levels: list[int] | None = None, max
                     or rec.get("uuid")
                     or f"{p.stem}:{len(tasks)}"
                 )
-                lvl = rec.get("Level") or rec.get("level")
-                # Convert level to int if it's a string
-                if isinstance(lvl, str) and lvl.isdigit():
-                    lvl = int(lvl)
-                fname = rec.get("file_name") or rec.get("filename") or None
+                lvl = _parse_level(rec.get("Level") or rec.get("level"))
+                fname_obj = rec.get("file_name") or rec.get("filename")
+                fname = fname_obj if isinstance(fname_obj, str) else None
 
                 # Only evaluate examples with public answers (dev/validation split)
                 # Skip if no question, no answer, or answer is placeholder like "?"
-                if not q or ans is None or str(ans).strip() in ["?", ""]:
+                if ans is None or str(ans).strip() in ["?", ""]:
                     continue
 
                 if wanted_levels and (lvl not in wanted_levels):
@@ -366,9 +391,10 @@ class GAIA:
                 "with access to gaia-benchmark/GAIA."
             )
 
-        from huggingface_hub import snapshot_download
+        import huggingface_hub
 
-        local_dir = snapshot_download(  # type: ignore
+        hf_hub = cast(Any, huggingface_hub)
+        local_dir = hf_hub.snapshot_download(
             repo_id="gaia-benchmark/GAIA",
             repo_type="dataset",
             revision="682dd723ee1e1697e00360edccf2366dc8418dd9",
@@ -376,6 +402,8 @@ class GAIA:
             local_dir=str(self.data_dir),
             force_download=False,
         )
+        if not isinstance(local_dir, str):
+            raise TypeError("snapshot_download returned unexpected non-string path")
         return Path(local_dir)
 
     async def _run_single_task(
@@ -522,7 +550,7 @@ class GAIA:
 
             # Run tasks
             semaphore = asyncio.Semaphore(parallel)
-            results = []
+            results: list[TaskResult] = []
 
             tasks_coroutines = [self._run_single_task(task, task_runner, semaphore, timeout) for task in tasks]
 
@@ -561,7 +589,7 @@ class GAIA:
         with open(output_path, "w", encoding="utf-8") as f:
             for result in results:
                 # Convert messages to serializable format
-                serializable_messages = []
+                serializable_messages: list[dict[str, Any] | str] = []
                 if result.prediction.messages:
                     for msg in result.prediction.messages:
                         if hasattr(msg, "model_dump"):
@@ -569,7 +597,7 @@ class GAIA:
                             serializable_messages.append(msg.model_dump())
                         elif hasattr(msg, "__dict__"):
                             # Regular object with attributes
-                            serializable_messages.append(vars(msg))
+                            serializable_messages.append(cast(dict[str, Any], getattr(msg, "__dict__", {})))
                         else:
                             # Fallback to string representation
                             serializable_messages.append(str(msg))
@@ -614,16 +642,20 @@ def viewer_main() -> None:
     args = parser.parse_args()
 
     # Load results
-    results = []
+    results: list[dict[str, Any]] = []
     with open(args.results_file, encoding="utf-8") as f:
         for line in f:
             if line.strip():
                 try:
                     import orjson
 
-                    results.append(orjson.loads(line))
+                    parsed: object = orjson.loads(line)
                 except ImportError:
-                    results.append(json.loads(line))
+                    parsed = json.loads(line)
+
+                record = _coerce_record(parsed)
+                if record is not None:
+                    results.append(record)
 
     # Apply filters
     if args.level is not None:

python/packages/lab/pyproject.toml

@@ -122,6 +122,7 @@ omit = [
 
 [tool.pyright]
 extends = "../../pyproject.toml"
+include = ["gaia/agent_framework_lab_gaia", "lightning/agent_framework_lab_lightning", "tau2/agent_framework_lab_tau2"]
 exclude = ['gaia/tests', 'lightning/tests', 'tau2/tests', 'namespace', '**/samples']
 
 [tool.mypy]
@@ -151,10 +152,10 @@ mypy-gaia = "mypy --config-file $POE_ROOT/pyproject.toml gaia/agent_framework_la
 mypy-lightning = "mypy --config-file $POE_ROOT/pyproject.toml lightning/agent_framework_lab_lightning"
 mypy-tau2 = "mypy --config-file $POE_ROOT/pyproject.toml tau2/agent_framework_lab_tau2"
 mypy = ["mypy-gaia", "mypy-lightning", "mypy-tau2"]
-test = "pytest --cov-report=term-missing:skip-covered --junitxml=test-results.xml"
-test-gaia = "pytest gaia/tests --cov=agent_framework_lab_gaia --cov-report=term-missing:skip-covered"
-test-lightning = "pytest lightning/tests --cov=agent_framework_lab_lightning --cov-report=term-missing:skip-covered"
-test-tau2 = "pytest tau2/tests --cov=agent_framework_lab_tau2 --cov-report=term-missing:skip-covered"
+test = "pytest -m \"not integration\" --cov-report=term-missing:skip-covered --junitxml=test-results.xml"
+test-gaia = "pytest -m \"not integration\" gaia/tests --cov=agent_framework_lab_gaia --cov-report=term-missing:skip-covered"
+test-lightning = "pytest -m \"not integration\" lightning/tests --cov=agent_framework_lab_lightning --cov-report=term-missing:skip-covered"
+test-tau2 = "pytest -m \"not integration\" tau2/tests --cov=agent_framework_lab_tau2 --cov-report=term-missing:skip-covered"
 build = "echo 'Skipping build'"
 publish = "echo 'Skipping publish'"
 

python/packages/lab/tau2/agent_framework_lab_tau2/_message_utils.py

@@ -23,7 +23,7 @@ def flip_messages(messages: list[Message]) -> list[Message]:
         """Remove function call content from message contents."""
         return [content for content in messages if content.type != "function_call"]
 
-    flipped_messages = []
+    flipped_messages: list[Message] = []
     for msg in messages:
         role_value = _get_role_value(msg.role)
         if role_value == "assistant":

python/packages/lab/tau2/agent_framework_lab_tau2/_tau2_utils.py

@@ -3,7 +3,7 @@
 import json
 from collections.abc import Mapping
 from copy import deepcopy
-from typing import Any
+from typing import Any, TypeGuard, cast
 
 import numpy as np
 from agent_framework._tools import FunctionTool
@@ -27,6 +27,26 @@ from tau2.environment.tool import Tool  # type: ignore[import-untyped]
 _original_set_state = Environment.set_state
 
 
+def _to_str(value: object, default: str = "") -> str:
+    if isinstance(value, str):
+        return value
+    if value is None:
+        return default
+    return str(value)
+
+
+def _is_any_list(value: Any) -> TypeGuard[list[Any]]:
+    return isinstance(value, list)
+
+
+def _is_any_mapping(value: Any) -> TypeGuard[Mapping[Any, Any]]:
+    return isinstance(value, Mapping)
+
+
+def _is_any_sequence(value: Any) -> TypeGuard[list[Any] | tuple[Any, ...] | set[Any]]:
+    return isinstance(value, (list, tuple, set))
+
+
 def convert_tau2_tool_to_function_tool(tau2_tool: Tool) -> FunctionTool:
     """Convert a tau2 Tool to a FunctionTool for agent framework compatibility.
 
@@ -41,7 +61,7 @@ def convert_tau2_tool_to_function_tool(tau2_tool: Tool) -> FunctionTool:
 
     return FunctionTool(
         name=tau2_tool.name,
-        description=tau2_tool._get_description(),
+        description=tau2_tool._get_description(),  # pyright: ignore[reportPrivateUsage]
         func=wrapped_func,
         input_model=tau2_tool.params,
     )
@@ -53,27 +73,26 @@ def convert_agent_framework_messages_to_tau2_messages(messages: list[Message]) -
     Handles role mapping, text extraction, function calls, and function results.
     Function results are converted to separate ToolMessage instances.
     """
-    tau2_messages = []
+    tau2_messages: list[Tau2Message] = []
 
     for msg in messages:
         role_str = str(msg.role)
 
         # Extract text content from all text-type contents
-        text_content = None
         text_contents = [c for c in msg.contents if hasattr(c, "text") and hasattr(c, "type") and c.type == "text"]
-        if text_contents:
-            text_content = " ".join(c.text for c in text_contents)  # type: ignore[misc]
+        content_parts: list[str] = [_to_str(getattr(c, "text", "")) for c in text_contents]
+        content_value = " ".join(content_parts)
 
         # Extract function calls and convert to ToolCall objects
         function_calls = [c for c in msg.contents if hasattr(c, "type") and c.type == "function_call"]
-        tool_calls = None
+        tool_calls: list[ToolCall] | None = None
         if function_calls:
             tool_calls = []
             for fc in function_calls:
                 arguments = fc.parse_arguments() or {}
                 tool_call = ToolCall(
-                    id=fc.call_id,
-                    name=fc.name,
+                    id=_to_str(fc.call_id),
+                    name=_to_str(fc.name),
                     arguments=arguments,
                     requestor="assistant" if role_str == "assistant" else "user",
                 )
@@ -84,11 +103,11 @@ def convert_agent_framework_messages_to_tau2_messages(messages: list[Message]) -
 
         # Create main message based on role
         if role_str == "system":
-            tau2_messages.append(SystemMessage(role="system", content=text_content))
+            tau2_messages.append(SystemMessage(role="system", content=content_value))
         elif role_str == "user":
-            tau2_messages.append(UserMessage(role="user", content=text_content, tool_calls=tool_calls))
+            tau2_messages.append(UserMessage(role="user", content=content_value, tool_calls=tool_calls))
         elif role_str == "assistant":
-            tau2_messages.append(AssistantMessage(role="assistant", content=text_content, tool_calls=tool_calls))
+            tau2_messages.append(AssistantMessage(role="assistant", content=content_value, tool_calls=tool_calls))
         elif role_str == "tool":
             # Tool messages are handled as function results below
             pass
@@ -98,7 +117,7 @@ def convert_agent_framework_messages_to_tau2_messages(messages: list[Message]) -
             dumpable_content = _dump_function_result(fr.result)
             content = dumpable_content if isinstance(dumpable_content, str) else json.dumps(dumpable_content)
             tool_msg = ToolMessage(
-                id=fr.call_id,
+                id=_to_str(fr.call_id),
                 role="tool",
                 content=content,
                 requestor="assistant",  # Most tool calls originate from assistant
@@ -126,12 +145,10 @@ def patch_env_set_state() -> None:
         if self.solo_mode and any(isinstance(message, UserMessage) for message in message_history):
             raise ValueError("User messages are not allowed in solo mode")
 
-        def get_actions_from_messages(
-            messages: list[Tau2Message],
-        ) -> list[tuple[ToolCall, ToolMessage]]:
+        def get_actions_from_messages(messages: list[Tau2Message]) -> list[tuple[ToolCall, ToolMessage]]:
             """Get the actions from the messages."""
             messages = deepcopy(messages)[::-1]
-            actions = []
+            actions: list[tuple[ToolCall, ToolMessage]] = []
             while messages:
                 message = messages.pop()
                 if isinstance(message, ToolMessage):
@@ -153,10 +170,13 @@ def patch_env_set_state() -> None:
             return actions
 
         if initialization_data is not None:
-            if initialization_data.agent_data is not None:
-                self.tools.update_db(initialization_data.agent_data)
-            if initialization_data.user_data is not None:
-                self.user_tools.update_db(initialization_data.user_data)
+            agent_data = cast(object, getattr(initialization_data, "agent_data", None))
+            if isinstance(agent_data, dict):
+                self.tools.update_db(cast(dict[str, Any], agent_data))
+
+            user_data = cast(object, getattr(initialization_data, "user_data", None))
+            if isinstance(user_data, dict):
+                self.user_tools.update_db(cast(dict[str, Any], user_data))
 
         if initialization_actions is not None:
             for action in initialization_actions:
@@ -188,10 +208,11 @@ def unpatch_env_set_state() -> None:
 def _dump_function_result(result: Any) -> Any:
     if isinstance(result, BaseModel):
         return result.model_dump_json()
-    if isinstance(result, list):
+    if _is_any_list(result):
         return [_dump_function_result(item) for item in result]
     if isinstance(result, dict):
-        return {k: _dump_function_result(v) for k, v in result.items()}
+        result_dict = cast(dict[str, Any], result)
+        return {k: _dump_function_result(v) for k, v in result_dict.items()}
     if result is None:
         return None
     return result
@@ -208,11 +229,11 @@ def _to_native(obj: Any) -> Any:
         return _to_native(obj.item())
 
     # 3) Dict-like -> dict
-    if isinstance(obj, Mapping):
+    if _is_any_mapping(obj):
         return {_to_native(k): _to_native(v) for k, v in obj.items()}
 
     # 4) Lists/Tuples/Sets -> list
-    if isinstance(obj, (list, tuple, set)):
+    if _is_any_sequence(obj):
         return [_to_native(x) for x in obj]
 
     # 5) Anything else: leave as-is
@@ -227,9 +248,10 @@ def _recursive_json_deserialize(obj: Any) -> Any:
             return _recursive_json_deserialize(deserialized)
         except (json.JSONDecodeError, TypeError):
             return obj
-    elif isinstance(obj, list):
+    elif _is_any_list(obj):
         return [_recursive_json_deserialize(item) for item in obj]
     elif isinstance(obj, dict):
-        return {k: _recursive_json_deserialize(v) for k, v in obj.items()}
+        typed_obj = cast(dict[str, Any], obj)
+        return {k: _recursive_json_deserialize(v) for k, v in typed_obj.items()}
     else:
         return obj

python/packages/lab/tau2/agent_framework_lab_tau2/runner.py

@@ -3,7 +3,7 @@
 from __future__ import annotations
 
 import uuid
-from typing import Any
+from typing import Any, cast
 
 from agent_framework import (
     Agent,
@@ -38,6 +38,16 @@ from ._tau2_utils import convert_agent_framework_messages_to_tau2_messages, conv
 
 __all__ = ["ASSISTANT_AGENT_ID", "ORCHESTRATOR_ID", "USER_SIMULATOR_ID", "TaskRunner"]
 
+
+def _get_openai_schema(tool: Any) -> dict[str, Any]:
+    schema = getattr(tool, "openai_schema", None)
+    if isinstance(schema, dict):
+        schema_dict = cast(dict[object, Any], schema)
+        if all(isinstance(key, str) for key in schema_dict):
+            return cast(dict[str, Any], schema_dict)
+    raise TypeError(f"Tool {tool} does not expose a dict openai_schema")
+
+
 # Agent instructions matching tau2's LLMAgent
 ASSISTANT_AGENT_INSTRUCTION = """
 You are a customer service agent that helps the user according to the <policy> provided below.
@@ -205,7 +215,7 @@ class TaskRunner:
             context_providers=[
                 SlidingWindowHistoryProvider(
                     system_message=assistant_system_prompt,
-                    tool_definitions=[tool.openai_schema for tool in tools],
+                    tool_definitions=[_get_openai_schema(tool) for tool in tools],
                     max_tokens=self.assistant_window_size,
                 )
             ],