fix(ag-ui): partially unknown types

Fix partially known types lint violations in the ag-ui tests due to
to conditional imports.

Author: stevenh

pydantic_ai_slim/pydantic_ai/models/test.py

@@ -6,10 +6,10 @@
 from contextlib import asynccontextmanager
 from dataclasses import InitVar, dataclass, field
 from datetime import date, datetime, timedelta
-from typing import Any, Literal, TypeAlias, Union
+from typing import Any, Literal, Union
 
 import pydantic_core
-from typing_extensions import assert_never
+from typing_extensions import TypeAlias, assert_never
 
 from .. import _utils
 from ..messages import (

tests/test_ag_ui.py

@@ -18,12 +18,6 @@
 from pydantic import BaseModel
 
 from pydantic_ai import Agent
-from pydantic_ai.ag_ui import (
-    SSE_CONTENT_TYPE,
-    Adapter,
-    Role,
-    StateDeps,
-)
 from pydantic_ai.models.test import TestModel, TestNode, TestToolCallPart
 
 has_ag_ui: bool = False
@@ -44,6 +38,13 @@
         UserMessage,
     )
 
+    from pydantic_ai.ag_ui import (
+        SSE_CONTENT_TYPE,
+        Adapter,
+        Role,
+        StateDeps,
+    )
+
     has_ag_ui = True
 
 
@@ -122,7 +123,7 @@ async def create_adapter(
                 call_tools=call_tools,
                 tool_call_deltas={'get_weather_parts', 'current_time'},
             ),
-            deps_type=StateDeps[StateInt],
+            deps_type=StateDeps[StateInt],  # type: ignore[reportUnknownArgumentType]
             tools=[send_snapshot, send_custom, current_time],
         )
     )
@@ -763,7 +764,7 @@ async def test_run_method(mock_uuid: _MockUUID, tc: AdapterRunTest) -> None:
     events: list[str] = []
     thread_id: str = f'{THREAD_ID_PREFIX}{mock_uuid()}'
     adapter: Adapter[StateDeps[StateInt], str] = await create_adapter(tc.call_tools)
-    deps: StateDeps[StateInt] = StateDeps[StateInt](state_type=StateInt)
+    deps: StateDeps[StateInt] = StateDeps[StateInt](state_type=StateInt)  # type: ignore[reportUnknownArgumentType]
     for run in tc.runs:
         if run.nodes is not None:
             assert isinstance(adapter.agent.model, TestModel), 'Agent model is not TestModel'
@@ -774,11 +775,11 @@ async def test_run_method(mock_uuid: _MockUUID, tc: AdapterRunTest) -> None:
             run_id=f'{RUN_ID_PREFIX}{mock_uuid()}',
         )
 
-        events.extend([event async for event in adapter.run(run_input, deps=deps)])
+        events.extend([event async for event in adapter.run(run_input, deps=deps)])  # type: ignore[reportUnknownArgumentType]
 
     assert_events(events, tc.expected_events)
     if tc.expected_state is not None:
-        assert deps.state.value == tc.expected_state
+        assert deps.state.value == tc.expected_state  # type: ignore[reportUnknownArgumentType]
 
 
 async def test_concurrent_runs(mock_uuid: _MockUUID, adapter: Adapter[None, str]) -> None: