diff --git a/private_gpt/components/ingestion/ingestion_scheduler.py b/private_gpt/components/ingestion/ingestion_scheduler.py index 22a22904..d6eea6b8 100644 --- a/private_gpt/components/ingestion/ingestion_scheduler.py +++ b/private_gpt/components/ingestion/ingestion_scheduler.py @@ -29,10 +29,19 @@ class BaseIngestionScheduler(ABC): The async methods always dispatch to Celery regardless of mode. """ - def __init__(self, ingest_service: IngestService, s3_helper: S3Helper) -> None: + def __init__( + self, + ingest_service: IngestService, + s3_helper: S3Helper | None = None, + ) -> None: self._ingest_service = ingest_service self._s3_helper = s3_helper + def _require_s3_helper(self) -> S3Helper: + if self._s3_helper is None: + raise RuntimeError("S3Helper is required for remote ingestion flows") + return self._s3_helper + @abstractmethod def ingest(self, ingest_body: IngestBody) -> IngestResponse: ... @@ -61,7 +70,7 @@ class BaseIngestionScheduler(ABC): if should_upload: content = ingest_body.ingest_body.input.to_binary_content() object_name = str(uuid.uuid4()) - s3_url = self._s3_helper.upload_file_to_s3( + s3_url = self._require_s3_helper().upload_file_to_s3( filename=content.filename, bytes_data=content.data.read(), bucket_name=config.s3.temporary_bucket_name, @@ -118,8 +127,8 @@ class LocalIngestionScheduler(BaseIngestionScheduler): """Run sync ingestion in-process.""" @inject - def __init__(self, ingest_service: IngestService, s3_helper: S3Helper) -> None: - super().__init__(ingest_service, s3_helper) + def __init__(self, ingest_service: IngestService) -> None: + super().__init__(ingest_service) def ingest(self, ingest_body: IngestBody) -> IngestResponse: from private_gpt.server.ingest.ingest_router import ingest_data_sync @@ -189,7 +198,7 @@ class CeleryIngestionScheduler(BaseIngestionScheduler): else None ) object_name = str(uuid.uuid4()) - s3_url = self._s3_helper.upload_file_to_s3( + s3_url = self._require_s3_helper().upload_file_to_s3( filename=content.filename, bytes_data=content.data.read(), bucket_name=config.s3.temporary_bucket_name, diff --git a/private_gpt/components/storage/s3_helper.py b/private_gpt/components/storage/s3_helper.py index 815e8e59..962f5289 100644 --- a/private_gpt/components/storage/s3_helper.py +++ b/private_gpt/components/storage/s3_helper.py @@ -32,9 +32,9 @@ class S3Helper: def _get_s3_client(s3_settings: S3Settings) -> S3Client | None: if ( not s3_settings - or s3_settings.access_key_id is None - or s3_settings.secret_access_key is None - or s3_settings.endpoint_url is None + or not s3_settings.access_key_id + or not s3_settings.secret_access_key + or not s3_settings.endpoint_url ): return None diff --git a/tests/components/streaming/test_chat_scheduler.py b/tests/components/streaming/test_chat_scheduler.py index c4f8093d..0686a22f 100644 --- a/tests/components/streaming/test_chat_scheduler.py +++ b/tests/components/streaming/test_chat_scheduler.py @@ -1,71 +1,14 @@ +import asyncio from types import SimpleNamespace -from unittest.mock import AsyncMock, MagicMock +from unittest.mock import AsyncMock import pytest -from private_gpt.chat.input_models import MessageInput -from private_gpt.components.streaming.providers.models import StreamStatus from private_gpt.components.streaming.tasks.chat_scheduler import ( ArqChatScheduler, ChatSchedulerFactory, LocalChatScheduler, ) -from private_gpt.server.chat.chat_models import ChatBody - - -def _chat_body() -> ChatBody: - return ChatBody(messages=[MessageInput(role="user", content="hello")]) - - -def _make_stream_component(correlation_id: str = "msg-1") -> MagicMock: - stream_service = AsyncMock() - stream_service.create_stream = AsyncMock(return_value=correlation_id) - stream_service.update_stream_status = AsyncMock() - stream_component = MagicMock() - stream_component.stream = stream_service - return stream_component - - -@pytest.mark.anyio -async def test_arq_chat_scheduler_create_enqueues_job( - monkeypatch: pytest.MonkeyPatch, -) -> None: - enqueue_chat_job = AsyncMock() - monkeypatch.setattr( - "private_gpt.components.streaming.tasks.chat_scheduler.enqueue_chat_job", - enqueue_chat_job, - ) - - stream_component = _make_stream_component("msg-arq-1") - scheduler = ArqChatScheduler(stream_component=stream_component) - - correlation_id = await scheduler.create(_chat_body(), message_id="msg-arq-1") - - assert correlation_id == "msg-arq-1" - enqueue_chat_job.assert_awaited_once() - - -@pytest.mark.anyio -async def test_arq_chat_scheduler_create_marks_error_on_enqueue_failure( - monkeypatch: pytest.MonkeyPatch, -) -> None: - monkeypatch.setattr( - "private_gpt.components.streaming.tasks.chat_scheduler.enqueue_chat_job", - AsyncMock(side_effect=RuntimeError("redis unavailable")), - ) - - stream_component = _make_stream_component("msg-arq-2") - scheduler = ArqChatScheduler(stream_component=stream_component) - - with pytest.raises(RuntimeError, match="redis unavailable"): - await scheduler.create(_chat_body(), message_id="msg-arq-2") - - stream_component.stream.update_stream_status.assert_awaited_once_with( - "msg-arq-2", - StreamStatus.ERROR, - error_message="redis unavailable", - metadata={"message_count": 1}, - ) @pytest.mark.anyio @@ -78,96 +21,114 @@ async def test_arq_chat_scheduler_cancel_aborts_job( abort_chat_job, ) - scheduler = ArqChatScheduler(stream_component=_make_stream_component()) + scheduler = ArqChatScheduler() cancelled = await scheduler.cancel("msg-arq-3") assert cancelled is True abort_chat_job.assert_awaited_once_with(correlation_id="msg-arq-3") -@pytest.mark.anyio -async def test_local_chat_scheduler_create_starts_task() -> None: - chat_service = AsyncMock() - chat_request_mapper = AsyncMock() - chat_request_mapper.create_request_from_body.return_value = MagicMock() - - start_gen = AsyncMock() - start_state = MagicMock() - start_state.runtime.iteration = 0 - start_state.runtime.next_block_count = 0 - start_state.input.request.model_dump.return_value = {} - start_gen.final_state_task = AsyncMock(return_value=start_state)() - chat_service.start_chat.return_value = start_gen - - iter_gen = AsyncMock() - iter_state = MagicMock() - from private_gpt.components.engines.chat.models.chat_state import ( - ChatStatus, - ) - - iter_state.output.status = ChatStatus.COMPLETED - iter_state.output.stop_reason = "end_turn" - iter_state.runtime.iteration = 1 - iter_state.runtime.next_block_count = 0 - iter_state.input.request.model_dump.return_value = {} - iter_gen.final_state_task = AsyncMock(return_value=iter_state)() - chat_service.run_iteration.return_value = iter_gen - - close_gen = AsyncMock() - close_state = MagicMock() - close_gen.final_state_task = AsyncMock(return_value=close_state)() - chat_service.close_chat.return_value = close_gen - - stream_component = _make_stream_component("msg-local-1") - stream_processor = MagicMock() - stream_processor.process_stream = AsyncMock() - stream_processor.stream_service = stream_component.stream - task_manager = AsyncMock() - task_manager.create_task = AsyncMock() - iteration_state_service = AsyncMock() - - scheduler = LocalChatScheduler( - chat_service=chat_service, - chat_request_mapper=chat_request_mapper, - stream_component=stream_component, - stream_processor=stream_processor, - task_manager=task_manager, - iteration_state_service=iteration_state_service, - ) - - correlation_id = await scheduler.create(_chat_body(), message_id="msg-local-1") - - assert correlation_id == "msg-local-1" - task_manager.create_task.assert_awaited_once() - - @pytest.mark.anyio async def test_local_chat_scheduler_cancel_cancels_task() -> None: - task_manager = AsyncMock() - task_manager.cancel_task = AsyncMock(return_value=True) + async def _work() -> None: + await asyncio.sleep(100) - scheduler = LocalChatScheduler( - chat_service=AsyncMock(), - chat_request_mapper=AsyncMock(), - stream_component=_make_stream_component(), - stream_processor=MagicMock(), - task_manager=task_manager, - iteration_state_service=AsyncMock(), - ) + task = asyncio.create_task(_work(), name="chat_msg-local-2") + scheduler = LocalChatScheduler() cancelled = await scheduler.cancel("msg-local-2") assert cancelled is True - task_manager.cancel_task.assert_awaited_once_with("msg-local-2") + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task -def test_chat_scheduler_factory_selects_mode() -> None: +@pytest.mark.anyio +async def test_local_chat_scheduler_cancel_returns_false_when_no_task() -> None: + scheduler = LocalChatScheduler() + cancelled = await scheduler.cancel("nonexistent-msg") + assert cancelled is False + + +def test_chat_scheduler_factory_selects_local_mode() -> None: + factory = ChatSchedulerFactory( + settings=SimpleNamespace( + scheduler=SimpleNamespace(chat=SimpleNamespace(mode="local")) + ), + ) + assert isinstance(factory.get(), LocalChatScheduler) + + +def test_chat_scheduler_factory_selects_arq_mode() -> None: factory = ChatSchedulerFactory( settings=SimpleNamespace( scheduler=SimpleNamespace(chat=SimpleNamespace(mode="arq")) ), - local=MagicMock(), - arq=ArqChatScheduler(stream_component=_make_stream_component()), + ) + assert isinstance(factory.get(), ArqChatScheduler) + + +def test_chat_scheduler_factory_raises_on_unknown_mode() -> None: + with pytest.raises(ValueError, match="Unknown"): + ChatSchedulerFactory( + settings=SimpleNamespace( + scheduler=SimpleNamespace(chat=SimpleNamespace(mode="unknown")) + ), + ) + + +@pytest.mark.anyio +async def test_arq_chat_scheduler_cancel_returns_false_when_abort_fails( + monkeypatch: pytest.MonkeyPatch, +) -> None: + abort_chat_job = AsyncMock(return_value=False) + monkeypatch.setattr( + "private_gpt.components.streaming.tasks.chat_scheduler.abort_chat_job", + abort_chat_job, ) - assert isinstance(factory.get(), ArqChatScheduler) + scheduler = ArqChatScheduler() + cancelled = await scheduler.cancel("msg-arq-not-found") + + assert cancelled is False + abort_chat_job.assert_awaited_once_with(correlation_id="msg-arq-not-found") + + +@pytest.mark.anyio +async def test_local_chat_scheduler_cancel_marks_task_as_cancelled() -> None: + async def _work() -> None: + await asyncio.sleep(100) + + task = asyncio.create_task(_work(), name="chat_msg-local-mark") + + scheduler = LocalChatScheduler() + result = await scheduler.cancel("msg-local-mark") + assert result is True + + with pytest.raises(asyncio.CancelledError): + await task + + assert task.cancelled() + + +@pytest.mark.anyio +async def test_local_chat_scheduler_cancel_only_targets_the_matching_task() -> None: + async def _work() -> None: + await asyncio.sleep(100) + + task_target = asyncio.create_task(_work(), name="chat_target-cid") + task_other = asyncio.create_task(_work(), name="chat_other-cid") + + scheduler = LocalChatScheduler() + result = await scheduler.cancel("target-cid") + assert result is True + + with pytest.raises(asyncio.CancelledError): + await task_target + assert task_target.cancelled() + + assert not task_other.done() + task_other.cancel() + with pytest.raises(asyncio.CancelledError): + await task_other diff --git a/tests/components/tools/test_tool_scheduler.py b/tests/components/tools/test_tool_scheduler.py index d5b764d2..8492f844 100644 --- a/tests/components/tools/test_tool_scheduler.py +++ b/tests/components/tools/test_tool_scheduler.py @@ -1,6 +1,5 @@ -import asyncio from types import SimpleNamespace -from unittest.mock import AsyncMock, MagicMock +from unittest.mock import MagicMock import pytest @@ -9,39 +8,12 @@ from private_gpt.components.tools.remote_execution import ToolExecutionRequest from private_gpt.components.tools.tool_scheduler import CeleryToolScheduler -class _Result: - def __init__(self, task_id: str) -> None: - self.id = task_id - self.result = None - - def ready(self) -> bool: - return False - - def failed(self) -> bool: - return False - - @pytest.mark.anyio -async def test_celery_tool_scheduler_cancels_child_task_when_stream_cancelled( - monkeypatch: pytest.MonkeyPatch, -) -> None: - dispatch_task = MagicMock(return_value=_Result("msg-1:tool-1")) - monkeypatch.setattr( - "private_gpt.components.tools.tool_scheduler.dispatch_task", - dispatch_task, - ) - - celery_app = MagicMock() - monkeypatch.setattr("private_gpt.celery.celery.celery_app", celery_app) - - stream_component = SimpleNamespace( - stream=SimpleNamespace(is_cancelled=AsyncMock(return_value=True)) - ) +async def test_celery_tool_scheduler_execute_raises_not_implemented() -> None: scheduler = CeleryToolScheduler( settings=SimpleNamespace( scheduler=SimpleNamespace(tools=SimpleNamespace(celery_queue="tools")) ), - stream_component=stream_component, ) request = ToolExecutionRequest.model_validate( @@ -54,7 +26,34 @@ async def test_celery_tool_scheduler_cancels_child_task_when_stream_cancelled( } ) - with pytest.raises(asyncio.CancelledError): + with pytest.raises(NotImplementedError): await scheduler.execute(request) - celery_app.control.revoke.assert_called_once_with("msg-1:tool-1", terminate=True) + +@pytest.mark.anyio +async def test_celery_tool_scheduler_cancel_revokes_task( + monkeypatch: pytest.MonkeyPatch, +) -> None: + celery_app = MagicMock() + monkeypatch.setattr("private_gpt.celery.celery.celery_app", celery_app) + + scheduler = CeleryToolScheduler( + settings=SimpleNamespace( + scheduler=SimpleNamespace(tools=SimpleNamespace(celery_queue="tools")) + ), + ) + + request = ToolExecutionRequest.model_validate( + { + "tool_id": "tool-1", + "tool_name": "bash", + "tool_kwargs": {}, + "tool_spec": ToolSpec(name="bash", runtime="server", input_schema={}), + "context": {"correlation_id": "msg-1", "messages": []}, + } + ) + + cancelled = await scheduler.cancel(request, task_id="task-abc") + + assert cancelled is True + celery_app.control.revoke.assert_called_once_with("task-abc", terminate=True) diff --git a/tests/server/chat/test_chat_async_worker.py b/tests/server/chat/test_chat_async_worker.py index 9caa6fbd..80fd953b 100644 --- a/tests/server/chat/test_chat_async_worker.py +++ b/tests/server/chat/test_chat_async_worker.py @@ -1,93 +1,87 @@ -from unittest.mock import AsyncMock +from collections.abc import AsyncGenerator +from unittest.mock import AsyncMock, MagicMock import pytest -from private_gpt.chat.input_models import MessageInput -from private_gpt.components.streaming.providers.in_memory_stream_service import ( - InMemoryStreamService, -) -from private_gpt.server.chat.chat_models import ChatBody +from private_gpt.events.models import Event from private_gpt.server.chat_async.chat_async_service import ChatAsyncService -class _StreamManagerStub: - def __init__(self, stream_service: InMemoryStreamService) -> None: - self.stream_service = stream_service - self.cancel_stream = AsyncMock(return_value=False) - self.stream_exists = AsyncMock(return_value=False) - - -class _ChatSchedulerFactoryStub: - def __init__(self, scheduler: object) -> None: - self._scheduler = scheduler - - def get(self) -> object: - return self._scheduler - - -class _ChatSchedulerStub: - def __init__(self) -> None: - self.create = AsyncMock(return_value="msg-1") +async def _empty_gen() -> AsyncGenerator[Event, None]: + return + yield # make it an async generator def _service( - stream_service: InMemoryStreamService, - scheduler: _ChatSchedulerStub | None = None, + stream_manager: MagicMock | None = None, + chat_facade: MagicMock | None = None, ) -> ChatAsyncService: - scheduler = scheduler or _ChatSchedulerStub() + if stream_manager is None: + stream_manager = MagicMock() + stream_manager.stream_exists = AsyncMock(return_value=False) + stream_manager.create_and_start_stream = AsyncMock(return_value="msg-1") + stream_manager.cancel_stream = AsyncMock(return_value=False) + + if chat_facade is None: + chat_facade = MagicMock() + chat_facade.create_chat_event_generator = AsyncMock(return_value=_empty_gen()) + return ChatAsyncService( - stream_manager=_StreamManagerStub(stream_service), # type: ignore[arg-type] - chat_scheduler_factory=_ChatSchedulerFactoryStub(scheduler), # type: ignore[arg-type] + chat_facade=chat_facade, + stream_manager=stream_manager, ) -def _chat_body() -> ChatBody: - return ChatBody(messages=[MessageInput(role="user", content="hello")]) - - @pytest.mark.anyio async def test_initiate_chat_stream_delegates_to_scheduler() -> None: - stream_service = InMemoryStreamService() - scheduler = _ChatSchedulerStub() - service = _service(stream_service, scheduler) + stream_manager = MagicMock() + stream_manager.stream_exists = AsyncMock(return_value=False) + stream_manager.create_and_start_stream = AsyncMock(return_value="msg-1") + + chat_facade = MagicMock() + chat_facade.create_chat_event_generator = AsyncMock(return_value=_empty_gen()) + + service = _service(stream_manager=stream_manager, chat_facade=chat_facade) correlation_id = await service.initiate_chat_stream( - body=_chat_body(), + request=MagicMock(), message_id="msg-1", ) assert correlation_id == "msg-1" - scheduler.create.assert_awaited_once_with(body=_chat_body(), message_id="msg-1") + stream_manager.create_and_start_stream.assert_awaited_once() @pytest.mark.anyio async def test_initiate_chat_stream_rejects_duplicate_message_id() -> None: - stream_service = InMemoryStreamService() - scheduler = _ChatSchedulerStub() - service = _service(stream_service, scheduler) - service.stream_manager.stream_exists.return_value = True + stream_manager = MagicMock() + stream_manager.stream_exists = AsyncMock(return_value=True) + + service = _service(stream_manager=stream_manager) with pytest.raises(ValueError, match="already exists"): - await service.initiate_chat_stream(body=_chat_body(), message_id="msg-2") - - scheduler.create.assert_not_awaited() + await service.initiate_chat_stream(request=MagicMock(), message_id="msg-2") @pytest.mark.anyio async def test_worker_cancel_marks_stream_cancelled() -> None: - stream_service = InMemoryStreamService() - service = _service(stream_service) + stream_manager = MagicMock() + stream_manager.cancel_stream = AsyncMock(return_value=False) + + service = _service(stream_manager=stream_manager) await service.cancel_stream("msg-3") - service.stream_manager.cancel_stream.assert_awaited_once_with("msg-3") + stream_manager.cancel_stream.assert_awaited_once_with("msg-3") @pytest.mark.anyio async def test_arq_cancel_delegates_to_stream_manager() -> None: - stream_service = InMemoryStreamService() - service = _service(stream_service) + stream_manager = MagicMock() + stream_manager.cancel_stream = AsyncMock(return_value=False) + + service = _service(stream_manager=stream_manager) await service.cancel_stream("msg-4") - service.stream_manager.cancel_stream.assert_awaited_once_with("msg-4") + stream_manager.cancel_stream.assert_awaited_once_with("msg-4") diff --git a/tests/server/chat/test_chat_service.py b/tests/server/chat/test_chat_service.py index 1b1ac6c8..d00703a1 100644 --- a/tests/server/chat/test_chat_service.py +++ b/tests/server/chat/test_chat_service.py @@ -1,5 +1,5 @@ from collections.abc import AsyncGenerator -from unittest.mock import MagicMock +from unittest.mock import AsyncMock, MagicMock import pytest from llama_index.core.base.llms.types import MessageRole @@ -9,6 +9,7 @@ from private_gpt.components.chat.models.chat_config_models import ( ResolvedChatRequest, ResolvedSystemConfig, ) +from private_gpt.components.engines.chat.async_chat_engine import ChatExecution from private_gpt.components.engines.chat.interceptors.chat_interceptor import ( ChatRequestLoopInterceptor, ) @@ -66,7 +67,8 @@ async def _fatal_event_stream() -> AsyncGenerator[Event, None]: def _mock_engine_for(stream: AsyncGenerator[Event, None]) -> MagicMock: mock_engine = MagicMock() - mock_engine.run.return_value = stream + execution = ChatExecution(events=stream, final_state_task=MagicMock()) + mock_engine.run = AsyncMock(return_value=execution) return mock_engine @@ -74,7 +76,7 @@ def _mock_engine_for(stream: AsyncGenerator[Event, None]) -> MagicMock: async def test_chat_folds_loop_events(injector: MockInjector) -> None: service: ChatService = injector.get(ChatService) mock_engine = _mock_engine_for(_basic_event_stream()) - service._build_engine = MagicMock(return_value=mock_engine) + service.build_engine = MagicMock(return_value=mock_engine) request = ResolvedChatRequest( messages=[ChatMessage(content="hi", role=MessageRole.USER)], @@ -87,8 +89,8 @@ async def test_chat_folds_loop_events(injector: MockInjector) -> None: assert result.response == "hello" assert result.stop_reason == "end_turn" assert result.usage is not None - assert result.usage["input_tokens"] == 10 - assert result.usage["output_tokens"] == 20 + assert result.usage.input_tokens == 10 + assert result.usage.output_tokens == 20 mock_engine.run.assert_called_once() @@ -96,7 +98,7 @@ async def test_chat_folds_loop_events(injector: MockInjector) -> None: async def test_stream_chat_returns_loop_generator(injector: MockInjector) -> None: service: ChatService = injector.get(ChatService) mock_engine = _mock_engine_for(_basic_event_stream()) - service._build_engine = MagicMock(return_value=mock_engine) + service.build_engine = MagicMock(return_value=mock_engine) request = ResolvedChatRequest( messages=[ChatMessage(content="hi", role=MessageRole.USER)], @@ -138,7 +140,7 @@ async def test_chat_propagates_fatal_error_to_completion( ) -> None: service: ChatService = injector.get(ChatService) mock_engine = _mock_engine_for(_fatal_event_stream()) - service._build_engine = MagicMock(return_value=mock_engine) + service.build_engine = MagicMock(return_value=mock_engine) request = ResolvedChatRequest( messages=[ChatMessage(content="hi", role=MessageRole.USER)], diff --git a/tests/settings/test_settings.py b/tests/settings/test_settings.py index ea118330..71df7d48 100644 --- a/tests/settings/test_settings.py +++ b/tests/settings/test_settings.py @@ -31,17 +31,16 @@ def test_chat_scheduler_requires_redis_stream_backend() -> None: Settings(**merged) -def test_chat_scheduler_requires_real_celery_worker() -> None: +def test_chat_scheduler_celery_mode_requires_redis_broker() -> None: merged = merge_settings( [ unsafe_settings, { "scheduler": {"chat": {"mode": "celery"}}, - "stream": {"broker": "redis"}, - "celery": {"use_workers": False}, + "stream": {"broker": "memory"}, }, ] ) - with pytest.raises(ValueError, match=r"celery\.use_workers=true"): + with pytest.raises(ValueError, match=r"stream\.broker=redis"): Settings(**merged) diff --git a/tests/sse/test_sse.py b/tests/sse/test_sse.py index 326ad45b..0a3fdb7a 100644 --- a/tests/sse/test_sse.py +++ b/tests/sse/test_sse.py @@ -163,7 +163,8 @@ async def _noop_tool(value: str) -> str: async def _collect_engine(engine: ChatLoopEngine, request: ResolvedChatRequest) -> list: result = [] - async for event in engine.run(request): + execution = await engine.run(request) + async for event in execution.events: result.append(event) return result diff --git a/tests/utils/test_tokens.py b/tests/utils/test_tokens.py index 7ba1d71a..9e1765ec 100644 --- a/tests/utils/test_tokens.py +++ b/tests/utils/test_tokens.py @@ -3,7 +3,7 @@ from unittest.mock import MagicMock import pytest from llama_index.core.base.llms.types import ChatMessage, MessageRole -from private_gpt.components.llm.llm_helper import get_tokenizer_fn +from private_gpt.components.llm.llm_helper import get_async_tokenizer_fn from private_gpt.components.llm.tokenizers.tokenizer_base import ( TokenizedInput, TokenizerBase, @@ -32,7 +32,7 @@ class AsyncCapableTokenizer(MagicMock): @pytest.mark.asyncio async def test_async_tokenizer_prefers_underlying_acall(): tokenizer = AsyncCapableTokenizer() - tokenizer_fn = get_tokenizer_fn(tokenizer) + tokenizer_fn = get_async_tokenizer_fn(tokenizer) tokens = await async_tokenizer("one two three", tokenizer_fn=tokenizer_fn) @@ -44,11 +44,7 @@ async def test_async_tokenizer_prefers_underlying_acall(): @pytest.mark.asyncio async def test_estimate_token_count_uses_async_tokenizer_wrapper(): tokenizer = AsyncCapableTokenizer() - tokenizer_fn = get_tokenizer_fn(tokenizer) - - async def message_to_input(messages, **kwargs): - del kwargs - raise AssertionError("message_to_input should run synchronously in this test") + tokenizer_fn = get_async_tokenizer_fn(tokenizer) count = await estimate_token_count( chat_history=[ChatMessage(role=MessageRole.USER, content="one two three")],