core[lint]: Fix typing in test_async_callbacks (#30788)

This commit is contained in:
Christophe Bornet 2025-04-11 13:26:38 +02:00 committed by GitHub
parent 8c6734325b
commit 89f28a24d3
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -1,40 +1,44 @@
# ruff: noqa: ARG002
import asyncio import asyncio
from itertools import cycle from itertools import cycle
from typing import Any from typing import Any, Optional, Union
from uuid import UUID
import pytest import pytest
from pytest_benchmark.fixture import BenchmarkFixture # type: ignore from pytest_benchmark.fixture import BenchmarkFixture # type: ignore[import-untyped]
from typing_extensions import override
from langchain_core.callbacks.base import AsyncCallbackHandler from langchain_core.callbacks.base import AsyncCallbackHandler
from langchain_core.language_models import GenericFakeChatModel from langchain_core.language_models import GenericFakeChatModel
from langchain_core.messages import AIMessage from langchain_core.messages import AIMessage, BaseMessage
from langchain_core.outputs import ChatGenerationChunk, GenerationChunk
class MyCustomAsyncHandler(AsyncCallbackHandler): class MyCustomAsyncHandler(AsyncCallbackHandler):
@override
async def on_chat_model_start( async def on_chat_model_start(
self, self,
serialized: Any, serialized: dict[str, Any],
messages: Any, messages: list[list[BaseMessage]],
*, *,
run_id: Any, run_id: UUID,
parent_run_id: Any = None, parent_run_id: Optional[UUID] = None,
tags: Any = None, tags: Optional[list[str]] = None,
metadata: Any = None, metadata: Optional[dict[str, Any]] = None,
**kwargs: Any, **kwargs: Any,
) -> Any: ) -> Any:
# Do nothing # Do nothing
# Required to implement since this is an abstract method # Required to implement since this is an abstract method
pass pass
@override
async def on_llm_new_token( async def on_llm_new_token(
self, self,
token: str, token: str,
*, *,
chunk: Any = None, chunk: Optional[Union[GenerationChunk, ChatGenerationChunk]] = None,
run_id: Any, run_id: UUID,
parent_run_id: Any = None, parent_run_id: Optional[UUID] = None,
tags: Any = None, tags: Optional[list[str]] = None,
**kwargs: Any, **kwargs: Any,
) -> None: ) -> None:
await asyncio.sleep(0) await asyncio.sleep(0)