From d6bfc8790fb64cf9c0dd32e154635b556011d69a Mon Sep 17 00:00:00 2001 From: Javier Martinez Date: Fri, 17 Jul 2026 20:34:22 +0200 Subject: [PATCH] fix: tools & params --- .../engines/chat/async_chat_engine.py | 4 + .../engines/chat/models/chat_state.py | 1 + .../engines/chat/resumable_runner.py | 1 + .../components/tools/processors/base.py | 41 ++ .../server/chat/chat_request_mapper.py | 1 + .../interceptors/chat_interceptor_service.py | 10 +- .../interceptors/runtime_model_interceptor.py | 62 +++ .../validator_request_interceptor.py | 37 +- .../tools/test_processor_builder_contracts.py | 416 ++++++++++++++++++ tests/components/tools/test_tool_pipeline.py | 79 +++- tests/engines/test_async_chat_engine.py | 138 +++++- tests/engines/test_resumable_runtime.py | 2 + .../test_runtime_model_interceptor.py | 120 +++++ .../test_validation_request_interceptor.py | 4 + 14 files changed, 878 insertions(+), 38 deletions(-) create mode 100644 private_gpt/server/chat/interceptors/runtime_model_interceptor.py create mode 100644 tests/components/tools/test_processor_builder_contracts.py create mode 100644 tests/server/chat/interceptors/test_runtime_model_interceptor.py diff --git a/private_gpt/components/engines/chat/async_chat_engine.py b/private_gpt/components/engines/chat/async_chat_engine.py index 85befcc5..584ddeb1 100644 --- a/private_gpt/components/engines/chat/async_chat_engine.py +++ b/private_gpt/components/engines/chat/async_chat_engine.py @@ -154,6 +154,7 @@ class _IterationCheckpoint: class IterationCheckpointPayload(BaseModel): + model_id: str | None = None pending_async_tools: dict[str, str] = Field(default_factory=dict) tool_responses: list[ToolExecutionResponse] = Field(default_factory=list) pending_external_tool_calls: list[ToolSelection] = Field(default_factory=list) @@ -394,6 +395,7 @@ class AsyncChatEngine: context, ) new_payload = IterationCheckpointPayload( + model_id=state.runtime.model_id, total_input_tokens=state.runtime.total_input_tokens, total_output_tokens=state.runtime.total_output_tokens, has_input_usage=state.runtime.has_input_usage, @@ -580,6 +582,7 @@ class AsyncChatEngine: request, iteration, next_block_count, payload, hooks, channel ) new_payload = IterationCheckpointPayload( + model_id=state.runtime.model_id, total_input_tokens=state.runtime.total_input_tokens, total_output_tokens=state.runtime.total_output_tokens, has_input_usage=state.runtime.has_input_usage, @@ -1522,6 +1525,7 @@ class AsyncChatEngine: def _apply_payload_usage( self, run: _Run, payload: IterationCheckpointPayload ) -> None: + run.state.runtime.model_id = payload.model_id if payload.has_input_usage: run.total_input_tokens = payload.total_input_tokens run.has_input_usage = True diff --git a/private_gpt/components/engines/chat/models/chat_state.py b/private_gpt/components/engines/chat/models/chat_state.py index 46641f3c..15ea22d6 100644 --- a/private_gpt/components/engines/chat/models/chat_state.py +++ b/private_gpt/components/engines/chat/models/chat_state.py @@ -39,6 +39,7 @@ class ChatInputState(BaseModel): class ChatRuntimeState(BaseModel): """Store runtime counters.""" + model_id: str | None = None effective_token_limit: int | None = None tokenizer_fn: TokenizerFn | AsyncTokenizerFn | None = None diff --git a/private_gpt/components/engines/chat/resumable_runner.py b/private_gpt/components/engines/chat/resumable_runner.py index 9217cf7a..bc578682 100644 --- a/private_gpt/components/engines/chat/resumable_runner.py +++ b/private_gpt/components/engines/chat/resumable_runner.py @@ -399,6 +399,7 @@ class ResumableChatRunner: ) return IterationCheckpointPayload( + model_id=state.runtime.model_id, pending_async_tools=state.output.pending_async_tools, pending_external_tool_calls=state.output.pending_external_tool_calls, total_input_tokens=state.runtime.total_input_tokens, diff --git a/private_gpt/components/tools/processors/base.py b/private_gpt/components/tools/processors/base.py index 2959712e..e8cd13e0 100644 --- a/private_gpt/components/tools/processors/base.py +++ b/private_gpt/components/tools/processors/base.py @@ -65,6 +65,7 @@ def _replace_tool( original: ToolSpec, replacements: list[ToolSpec], ) -> bool: + replacements = _inherit_tool_properties(original, replacements) tools = request.tool_config.tools for index, candidate in enumerate(tools): if candidate is original: @@ -75,3 +76,43 @@ def _replace_tool( ] return True return False + + +def _inherit_tool_properties( + original: ToolSpec, + replacements: list[ToolSpec], +) -> list[ToolSpec]: + if not replacements: + return replacements + + inherited = [ + replacement.model_copy( + update={ + "context": original.context + if original.context is not None + else replacement.context, + "defer_loading": original.defer_loading or replacement.defer_loading, + "instructions": original.instructions + if original.instructions is not None + else replacement.instructions, + "requirements": list( + dict.fromkeys([*replacement.requirements, *original.requirements]) + ), + } + ) + for replacement in replacements + ] + + if len(inherited) == 1: + inherited[0] = inherited[0].model_copy( + update={ + "description": original.description + if original.description is not None + else inherited[0].description, + "partial_params": original.partial_params + if original.partial_params is not None + else inherited[0].partial_params, + } + ) + + return inherited diff --git a/private_gpt/server/chat/chat_request_mapper.py b/private_gpt/server/chat/chat_request_mapper.py index 0eab39f0..00ef3552 100644 --- a/private_gpt/server/chat/chat_request_mapper.py +++ b/private_gpt/server/chat/chat_request_mapper.py @@ -215,6 +215,7 @@ class ChatRequestMapper: if body.metadata and body.metadata.user_id else str(uuid.uuid4()), container=body.container, + maximum_context_length=self._settings.chat.maximum_context_length, maximum_loaded_skills=( body.maximum_loaded_skills if body.maximum_loaded_skills is not None diff --git a/private_gpt/server/chat/interceptors/chat_interceptor_service.py b/private_gpt/server/chat/interceptors/chat_interceptor_service.py index df8ecf58..e50027ec 100644 --- a/private_gpt/server/chat/interceptors/chat_interceptor_service.py +++ b/private_gpt/server/chat/interceptors/chat_interceptor_service.py @@ -44,6 +44,9 @@ from private_gpt.server.chat.interceptors.multimodal_interceptor import ( from private_gpt.server.chat.interceptors.platform_guidelines_interceptor import ( PlatformGuidelinesInterceptor, ) +from private_gpt.server.chat.interceptors.runtime_model_interceptor import ( + RuntimeModelRequestInterceptor, +) from private_gpt.server.chat.interceptors.skill_tool_visibility_interceptor import ( SkillToolVisibilityInterceptor, ) @@ -81,6 +84,7 @@ class ChatInterceptorService: settings: Settings, prompt_builder_service: PromptBuilderService, # --- request interceptors (run once, order matters) --- + runtime_model_interceptor: RuntimeModelRequestInterceptor, validation_request_interceptor: ValidatorRequestInterceptor, default_values_interceptor: DefaultValuesRequestInterceptor, mcp_interceptor: McpRequestInterceptor, @@ -109,7 +113,11 @@ class ChatInterceptorService: # Init interceptors .add_range( "init", - requests=[validation_request_interceptor, default_values_interceptor], + requests=[ + runtime_model_interceptor, + validation_request_interceptor, + default_values_interceptor, + ], ) # Init tools, internal tools & platform skills .add_range( diff --git a/private_gpt/server/chat/interceptors/runtime_model_interceptor.py b/private_gpt/server/chat/interceptors/runtime_model_interceptor.py new file mode 100644 index 00000000..e14eaff4 --- /dev/null +++ b/private_gpt/server/chat/interceptors/runtime_model_interceptor.py @@ -0,0 +1,62 @@ +from injector import inject, singleton + +from private_gpt.components.engines.chat.interceptors.chat_interceptor import ( + ChatRequestLoopInterceptor, +) +from private_gpt.components.engines.chat.models.chat_interceptor_context import ( + ChatInterceptorContext, +) +from private_gpt.components.engines.chat.models.chat_phase import InterceptorPhase +from private_gpt.components.llm.custom.base import ZylonLLM +from private_gpt.components.llm.llm_component import LLMComponent + + +@singleton +class RuntimeModelRequestInterceptor(ChatRequestLoopInterceptor): + @inject + def __init__(self, llm_component: LLMComponent) -> None: + self._llm_component = llm_component + + async def intercept(self, context: ChatInterceptorContext) -> None: + if context.phase not in { + InterceptorPhase.VALIDATION, + InterceptorPhase.BEFORE_ITERATION, + }: + return + + runtime = context.state.runtime + if runtime.model_id is None: + runtime.model_id = context.state.input.request.system.model + + if ( + runtime.effective_token_limit is not None + and runtime.tokenizer_fn is not None + ): + return + + llm = self._llm_component.get_llm(runtime.model_id) + metadata = ( + llm.get_metadata(**context.state.input.llm_kwargs) + if isinstance(llm, ZylonLLM) + else llm.metadata + ) + runtime.effective_token_limit = self._token_limit( + llm.metadata.context_window, + metadata.num_output, + ) + try: + runtime.tokenizer_fn = self._llm_component.get_tokenizer(runtime.model_id) + except ValueError: + runtime.tokenizer_fn = None + context.set_state(context.state) + + @staticmethod + def _token_limit( + context_window: int | None, + num_output: int | None, + ) -> int | None: + if context_window is None or context_window <= 0: + return None + reserved_output = num_output or 0 + effective = context_window - reserved_output - 256 + return effective if effective > 0 else context_window diff --git a/private_gpt/server/chat/interceptors/validator_request_interceptor.py b/private_gpt/server/chat/interceptors/validator_request_interceptor.py index ccedf730..692ea851 100644 --- a/private_gpt/server/chat/interceptors/validator_request_interceptor.py +++ b/private_gpt/server/chat/interceptors/validator_request_interceptor.py @@ -16,7 +16,6 @@ from private_gpt.components.engines.chat.models.chat_interceptor_context import from private_gpt.components.engines.chat.models.chat_phase import ( InterceptorPhase, ) -from private_gpt.components.llm.custom.base import ZylonLLM from private_gpt.components.llm.llm_component import LLMComponent from private_gpt.components.llm.llm_helper import ( max_audios_supported, @@ -106,19 +105,9 @@ class ValidatorRequestInterceptor(ChatRequestLoopInterceptor): Errors.Codes.INVALID_REQUEST_AUDIO_MAX_NUM_ERROR, ) - metadata = ( - llm.get_metadata(**context.state.input.llm_kwargs) - if isinstance(llm, ZylonLLM) - else llm.metadata - ) - token_limit = self._token_limit( - llm.metadata.context_window, metadata.num_output - ) - if token_limit is None: - return - - tokenize = self._llm_component.get_tokenizer(model_id) - if tokenize is None: + token_limit = context.state.runtime.effective_token_limit + tokenize = context.state.runtime.tokenizer_fn + if token_limit is None or tokenize is None: return user_message_tokens = len(tokenize(user_text)) @@ -166,28 +155,8 @@ class ValidatorRequestInterceptor(ChatRequestLoopInterceptor): f"exceed the maximum token limit {token_limit}." ) - # Update state with effective token limit for downstream components - state = context.state - state.runtime.effective_token_limit = token_limit - state.runtime.tokenizer_fn = tokenize - context.set_state(state) - return - @staticmethod - def _token_limit( - context_window: int | None, - num_output: int | None, - ) -> int | None: - """Compute effective token limit from model metadata.""" - if context_window is None or context_window <= 0: - return None - reserved_output = num_output or 0 - effective = context_window - reserved_output - 256 - if effective > 0: - return effective - return context_window - @staticmethod def _extract_text(message: ChatMessage) -> str: """Extract normalized text from message blocks.""" diff --git a/tests/components/tools/test_processor_builder_contracts.py b/tests/components/tools/test_processor_builder_contracts.py new file mode 100644 index 00000000..416483c9 --- /dev/null +++ b/tests/components/tools/test_processor_builder_contracts.py @@ -0,0 +1,416 @@ +import inspect +from types import SimpleNamespace +from unittest.mock import AsyncMock, Mock + +import pytest +from llama_index.core.base.llms.types import ChatMessage, MessageRole + +from private_gpt.chat.extensions.context_filter import ContextFilter +from private_gpt.chat.input_models import BlobVisibilityMode +from private_gpt.components.chat.models.chat_config_models import ( + ResolvedChatRequest, + ResolvedContextConfig, + ResolvedSystemConfig, + ResolvedToolConfig, + ToolSpec, +) +from private_gpt.components.sandbox.content_bundle import ContentBundle +from private_gpt.components.tools.builders.bash_tool_builder import BashToolBuilder +from private_gpt.components.tools.builders.database_query_builder import ( + DatabaseQueryToolBuilder, +) +from private_gpt.components.tools.builders.present_files_tool_builder import ( + PresentFilesToolBuilder, +) +from private_gpt.components.tools.builders.present_server_tool_builder import ( + PresentServerToolBuilder, +) +from private_gpt.components.tools.builders.semantic_search_builder import ( + SemanticSearchToolBuilder, +) +from private_gpt.components.tools.builders.tabular_data_builder import ( + TabularDataToolBuilder, +) +from private_gpt.components.tools.builders.text_editor_tool_builder import ( + TextEditorToolBuilder, +) +from private_gpt.components.tools.builders.web_fetch_builder import WebFetchToolBuilder +from private_gpt.components.tools.builders.web_search_builder import ( + WebSearchToolBuilder, +) +from private_gpt.components.tools.processors.bash_processor import BashProcessor +from private_gpt.components.tools.processors.database_query_processor import ( + DatabaseQueryProcessor, +) +from private_gpt.components.tools.processors.present_files_processor import ( + PresentFilesProcessor, +) +from private_gpt.components.tools.processors.present_server_processor import ( + PresentServerProcessor, +) +from private_gpt.components.tools.processors.semantic_search_processor import ( + SemanticSearchProcessor, +) +from private_gpt.components.tools.processors.tabular_data_processor import ( + TabularDataProcessor, +) +from private_gpt.components.tools.processors.text_editor_processor import ( + TextEditorProcessor, +) +from private_gpt.components.tools.processors.web_fetch_processor import ( + WebFetchProcessor, +) +from private_gpt.components.tools.processors.web_search_processor import ( + WebSearchProcessor, +) +from private_gpt.components.tools.types import ToolValidationMode +from private_gpt.server.utils.artifact_input import ( + IngestedArtifact, + SqlDatabaseArtifact, +) + + +def _tool(name: str) -> ToolSpec: + return ToolSpec(name=name, type=f"{name}_v1") + + +def _resolved(name: str) -> ToolSpec: + return ToolSpec.from_defaults( + name=name, + type=f"{name}_v1", + runtime="server", + async_fn=AsyncMock(return_value=[]), + ) + + +def _request( + tool: ToolSpec, + *, + tool_context: list[object] | None = None, + content_bundles: list[ContentBundle] | None = None, + bundles_to_remove: list[str] | None = None, +) -> ResolvedChatRequest: + return ResolvedChatRequest( + messages=[ChatMessage(role=MessageRole.USER, content="hello")], + system=ResolvedSystemConfig( + model="contract-model", + prompt="Contract system prompt", + blob_visibility=BlobVisibilityMode.INTERNAL, + ), + tool_config=ResolvedToolConfig( + tools=[tool], + validation_mode=ToolValidationMode.EAGER, + ), + tool_context=tool_context or [], + context=ResolvedContextConfig( + correlation_id="contract-correlation", + maximum_context_length=98_765, + content_bundles=content_bundles or [], + bundles_to_remove=bundles_to_remove or [], + ), + ) + + +@pytest.mark.parametrize( + ("builder_method", "expected_parameters"), + [ + ( + SemanticSearchToolBuilder.build_tool, + { + "context_filter", + "model_id", + "embed_model_id", + "name", + "type", + "description", + "validate", + "runtime", + "kwargs", + }, + ), + ( + TabularDataToolBuilder.build_tool, + { + "context_filter", + "model_id", + "embed_model_id", + "llm", + "name", + "type", + "description", + "validate", + "runtime", + "blob_visibility", + "kwargs", + }, + ), + ( + DatabaseQueryToolBuilder.build_tool, + { + "sql_artifacts", + "chat_history", + "name", + "type", + "description", + "validate", + "runtime", + "blob_visibility", + }, + ), + ( + WebSearchToolBuilder.build_tool, + {"model_id", "name", "type", "description", "validate", "runtime"}, + ), + ( + WebFetchToolBuilder.build_tool, + {"name", "type", "description", "runtime"}, + ), + ( + BashToolBuilder.build_tool, + {"config", "name", "type", "description"}, + ), + ( + TextEditorToolBuilder.build_view_tool, + {"config", "name", "type", "description"}, + ), + ( + TextEditorToolBuilder.build_str_replace_tool, + {"config", "name", "type", "description"}, + ), + ( + TextEditorToolBuilder.build_create_tool, + {"config", "name", "type", "description"}, + ), + ( + TextEditorToolBuilder.build_insert_tool, + {"config", "name", "type", "description"}, + ), + ( + PresentFilesToolBuilder.build_tool, + {"session_id", "bundles", "name", "type", "description"}, + ), + ( + PresentServerToolBuilder.build_tool, + {"session_id", "name", "type", "description"}, + ), + ], +) +def test_processor_builder_contract_tracks_signature_changes( + builder_method: object, + expected_parameters: set[str], +) -> None: + parameters = set(inspect.signature(builder_method).parameters) - {"self"} + assert parameters == expected_parameters + + +@pytest.mark.asyncio +async def test_semantic_search_builder_receives_complete_request_contract() -> None: + context_filter = ContextFilter(collection="knowledge") + builder = SimpleNamespace( + build_tool=AsyncMock(return_value=_resolved("semantic_search")) + ) + request = _request( + _tool("semantic_search"), + tool_context=[IngestedArtifact(context_filter=context_filter)], + ) + request.citation.enabled = True + + assert await SemanticSearchProcessor(builder).intercept(request) + + builder.build_tool.assert_awaited_once_with( + model_id="contract-model", + name="semantic_search", + type="semantic_search_v1", + context_filter=context_filter, + generate_citations=True, + validate=ToolValidationMode.EAGER, + token_limit=98_765, + ) + + +@pytest.mark.asyncio +async def test_tabular_builder_receives_complete_request_contract() -> None: + context_filter = ContextFilter(collection="tables") + builder = SimpleNamespace( + build_tool=AsyncMock(return_value=_resolved("tabular_analysis")) + ) + request = _request( + _tool("tabular_analysis"), + tool_context=[IngestedArtifact(context_filter=context_filter)], + ) + + assert await TabularDataProcessor(builder).intercept(request) + + builder.build_tool.assert_awaited_once_with( + model_id="contract-model", + name="tabular_analysis", + type="tabular_analysis_v1", + context_filter=context_filter, + validate=ToolValidationMode.EAGER, + blob_visibility=BlobVisibilityMode.INTERNAL, + ) + + +@pytest.mark.asyncio +async def test_database_builder_receives_complete_request_contract() -> None: + artifact = SqlDatabaseArtifact( + connection_string="sqlite:///contract.db", + schemas=["main"], + ) + builder = SimpleNamespace( + build_tool=AsyncMock(return_value=_resolved("database_query")) + ) + request = _request(_tool("database_query"), tool_context=[artifact]) + + assert await DatabaseQueryProcessor(builder).intercept(request) + + kwargs = builder.build_tool.await_args.kwargs + assert kwargs == { + "name": "database_query", + "type": "database_query_v1", + "sql_artifacts": [artifact], + "chat_history": kwargs["chat_history"], + "validate": ToolValidationMode.EAGER, + "blob_visibility": BlobVisibilityMode.INTERNAL, + } + assert [message.role for message in kwargs["chat_history"]] == [ + MessageRole.SYSTEM, + MessageRole.USER, + ] + + +@pytest.mark.asyncio +async def test_web_search_builder_receives_complete_request_contract() -> None: + builder = SimpleNamespace( + build_tool=AsyncMock(return_value=_resolved("web_search")) + ) + + assert await WebSearchProcessor(builder).intercept(_request(_tool("web_search"))) + + builder.build_tool.assert_awaited_once_with( + model_id="contract-model", + name="web_search", + type="web_search_v1", + ) + + +@pytest.mark.asyncio +async def test_web_fetch_builder_receives_complete_request_contract() -> None: + builder = SimpleNamespace(build_tool=Mock(return_value=_resolved("web_fetch"))) + + assert await WebFetchProcessor(builder).intercept(_request(_tool("web_fetch"))) + + builder.build_tool.assert_called_once_with( + name="web_fetch", + type="web_fetch_v1", + ) + + +@pytest.mark.asyncio +async def test_bash_builder_receives_complete_session_contract() -> None: + bundle = ContentBundle(canonical_path="/mnt/skills/contract/") + builder = SimpleNamespace(build_tool=AsyncMock(return_value=_resolved("bash"))) + request = _request( + _tool("bash"), + content_bundles=[bundle], + bundles_to_remove=["/mnt/skills/old/"], + ) + + assert await BashProcessor(builder).intercept(request) + + config = builder.build_tool.await_args.args[0] + assert config.session_id == "contract-correlation" + assert config.extra_bundles == [bundle] + assert config.bundles_to_remove == ["/mnt/skills/old/"] + builder.build_tool.assert_awaited_once_with( + config, + name="bash", + type="bash_v1", + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("tool_name", "builder_method"), + [ + ("view", "build_view_tool"), + ("str_replace", "build_str_replace_tool"), + ("create", "build_create_tool"), + ("insert", "build_insert_tool"), + ], +) +async def test_text_editor_builders_receive_complete_session_contract( + tool_name: str, + builder_method: str, +) -> None: + bundle = ContentBundle(canonical_path="/mnt/skills/editor/") + builder = SimpleNamespace( + build_view_tool=AsyncMock(return_value=_resolved("view")), + build_str_replace_tool=AsyncMock(return_value=_resolved("str_replace")), + build_create_tool=AsyncMock(return_value=_resolved("create")), + build_insert_tool=AsyncMock(return_value=_resolved("insert")), + ) + request = _request( + _tool(tool_name), + content_bundles=[bundle], + bundles_to_remove=["/mnt/skills/removed/"], + ) + + assert await TextEditorProcessor(builder).intercept(request) + + method = getattr(builder, builder_method) + config = method.await_args.args[0] + assert config.session_id == "contract-correlation" + assert config.extra_bundles == [bundle] + assert config.bundles_to_remove == ["/mnt/skills/removed/"] + method.assert_awaited_once_with( + config, + name=tool_name, + type=f"{tool_name}_v1", + ) + + +@pytest.mark.asyncio +async def test_present_files_builder_receives_complete_request_contract() -> None: + bundle = ContentBundle(canonical_path="/mnt/skills/present/") + builder = SimpleNamespace( + build_tool=AsyncMock(return_value=_resolved("present_files")) + ) + settings = SimpleNamespace( + code_execution=SimpleNamespace( + tools=SimpleNamespace(present_files_enabled=True) + ) + ) + + assert await PresentFilesProcessor(builder, settings).intercept( + _request(_tool("present_files"), content_bundles=[bundle]) + ) + + builder.build_tool.assert_awaited_once_with( + "contract-correlation", + bundles=[bundle], + name="present_files", + type="present_files_v1", + ) + + +@pytest.mark.asyncio +async def test_present_server_builder_receives_complete_request_contract() -> None: + builder = SimpleNamespace( + build_tool=AsyncMock(return_value=_resolved("present_server")) + ) + settings = SimpleNamespace( + code_execution=SimpleNamespace( + tools=SimpleNamespace(present_server_enabled=True) + ) + ) + + assert await PresentServerProcessor(builder, settings).intercept( + _request(_tool("present_server")) + ) + + builder.build_tool.assert_awaited_once_with( + "contract-correlation", + name="present_server", + type="present_server_v1", + ) diff --git a/tests/components/tools/test_tool_pipeline.py b/tests/components/tools/test_tool_pipeline.py index 18123230..a98734eb 100644 --- a/tests/components/tools/test_tool_pipeline.py +++ b/tests/components/tools/test_tool_pipeline.py @@ -9,6 +9,7 @@ from private_gpt.components.chat.models.chat_config_models import ( ResolvedChatRequest, ResolvedContextConfig, ResolvedToolConfig, + ToolRequirements, ToolSpec, ) from private_gpt.components.skills.models.skill_entities import ( @@ -16,7 +17,7 @@ from private_gpt.components.skills.models.skill_entities import ( SkillFrontmatter, SkillVersionEntity, ) -from private_gpt.components.tools.processors.base import _session_id +from private_gpt.components.tools.processors.base import _replace_tool, _session_id from private_gpt.components.tools.processors.bash_processor import BashProcessor from private_gpt.components.tools.processors.code_execution_processor import ( CodeExecutionProcessor, @@ -40,6 +41,82 @@ def _request(tools: list[ToolSpec]) -> ResolvedChatRequest: ) +def test_replace_tool_preserves_single_replacement_properties() -> None: + original = ToolSpec( + name="semantic_search", + type="semantic_search_v1", + description="Custom search description", + defer_loading=True, + partial_params={"scope": "project"}, + instructions="Use the project knowledge base.", + requirements=[ToolRequirements.SANDBOX], + ) + replacement = ToolSpec.from_defaults( + name="semantic_search", + type="semantic_search_v1", + runtime="server", + description="Default search description", + async_fn=AsyncMock(return_value=[]), + ) + request = _request([original]) + + assert _replace_tool(request, original, [replacement]) + + resolved = request.tool_config.tools[0] + assert resolved.description == "Custom search description" + assert resolved.defer_loading is True + assert resolved.partial_params == {"scope": "project"} + assert resolved.instructions == "Use the project knowledge base." + assert resolved.requirements == [ToolRequirements.SANDBOX] + assert resolved.runtime == "server" + assert resolved.async_fn is replacement.async_fn + + +def test_replace_tool_preserves_shared_properties_across_expansion() -> None: + original = ToolSpec( + name="code_execution", + type="code_execution_v1", + description="Wrapper description", + defer_loading=True, + partial_params={"unsafe_for_children": True}, + instructions="Use the shared sandbox carefully.", + requirements=[ToolRequirements.SANDBOX], + ) + replacements = [ + ToolSpec.from_defaults( + name="bash", + type="bash_v1", + runtime="server", + description="Bash description", + async_fn=AsyncMock(return_value=[]), + ), + ToolSpec.from_defaults( + name="text_editor", + type="text_editor_v1", + runtime="server", + description="Editor description", + async_fn=AsyncMock(return_value=[]), + ), + ] + request = _request([original]) + + assert _replace_tool(request, original, replacements) + + bash, editor = request.tool_config.tools + assert bash.description == "Bash description" + assert editor.description == "Editor description" + assert bash.partial_params is None + assert editor.partial_params is None + assert all(tool.defer_loading for tool in (bash, editor)) + assert all( + tool.instructions == "Use the shared sandbox carefully." + for tool in (bash, editor) + ) + assert all( + tool.requirements == [ToolRequirements.SANDBOX] for tool in (bash, editor) + ) + + @pytest.mark.asyncio async def test_tool_pipeline_recursively_expands_code_execution_wrapper() -> None: bash_builder = SimpleNamespace( diff --git a/tests/engines/test_async_chat_engine.py b/tests/engines/test_async_chat_engine.py index 3df376c6..a27ef3c5 100644 --- a/tests/engines/test_async_chat_engine.py +++ b/tests/engines/test_async_chat_engine.py @@ -21,6 +21,13 @@ from private_gpt.components.engines.chat.async_chat_engine import ( LocalEventChannel, ) from private_gpt.components.engines.chat.chat_engine import ChatLoopEngine +from private_gpt.components.engines.chat.interceptors.chat_interceptor import ( + ChatRequestLoopInterceptor, +) +from private_gpt.components.engines.chat.models.chat_interceptor_context import ( + ChatInterceptorContext, +) +from private_gpt.components.engines.chat.models.chat_phase import InterceptorPhase from private_gpt.components.engines.chat.models.chat_state import ( ChatInputState, ChatState, @@ -42,6 +49,9 @@ from private_gpt.events.models import ( TextBlock, ToolResultBlock, ) +from private_gpt.server.chat.interceptors.runtime_model_interceptor import ( + RuntimeModelRequestInterceptor, +) from tests.fixtures.mock_function_llm import get_mock_function_calling_llm @@ -148,6 +158,28 @@ class _FakeChatScheduler: return True +@dataclass +class _RuntimeObservation: + phase: InterceptorPhase + model_id: str | None + effective_token_limit: int | None + has_tokenizer: bool + + +class _RuntimeRecordingInterceptor(ChatRequestLoopInterceptor): + observations: list[_RuntimeObservation] + + async def intercept(self, context: ChatInterceptorContext) -> None: + self.observations.append( + _RuntimeObservation( + phase=context.phase, + model_id=context.state.runtime.model_id, + effective_token_limit=context.state.runtime.effective_token_limit, + has_tokenizer=context.state.runtime.tokenizer_fn is not None, + ) + ) + + @dataclass class _AsyncRunResult: events: list[Any] @@ -235,10 +267,13 @@ async def _run_async_engine( request: ResolvedChatRequest, mock_llm: Any, tool_scheduler: BaseToolScheduler, + request_interceptors: list[ChatRequestLoopInterceptor] | None = None, + llm_component: LLMComponent | None = None, ) -> _AsyncRunResult: + resolved_llm_component = llm_component or _make_llm_component(mock_llm) engine = AsyncChatEngine( - llm_component=_make_llm_component(mock_llm), - request_interceptors=[], + llm_component=resolved_llm_component, + request_interceptors=request_interceptors or [], response_interceptors=[], max_iterations=6, tool_scheduler=tool_scheduler, @@ -273,6 +308,7 @@ async def _run_async_engine( iteration=state.runtime.iteration, next_block_count=state.runtime.next_block_count, payload=IterationCheckpointPayload( + model_id=state.runtime.model_id, pending_async_tools=state.output.pending_async_tools, tool_responses=responses, pending_external_tool_calls=state.output.pending_external_tool_calls, @@ -291,6 +327,18 @@ async def _run_async_engine( return _AsyncRunResult(events=all_events, states=states) +class _RecordingRequestInterceptor(ChatRequestLoopInterceptor): + observations: list[tuple[InterceptorPhase, list[MessageRole]]] + + async def intercept(self, context: ChatInterceptorContext) -> None: + self.observations.append( + ( + context.phase, + [message.role for message in context.state.input.request.messages], + ) + ) + + @pytest.mark.asyncio async def test_async_engine_matches_sync_simple_message( base_request: ResolvedChatRequest, @@ -397,6 +445,92 @@ async def test_async_engine_matches_sync_one_server_tool_and_resumes_same_point( assert _normalize_events(async_result.events) == _normalize_events(sync_events) +@pytest.mark.asyncio +async def test_async_engine_reruns_before_iteration_with_resumed_tool_results( + base_request: ResolvedChatRequest, +) -> None: + request = base_request.model_copy(deep=True) + request.tool_config = ResolvedToolConfig(tools=[_server_tool("echo")]) + recorder = _RecordingRequestInterceptor(observations=[]) + + await _run_async_engine( + request, + get_mock_function_calling_llm( + [ + [ + ToolSelection( + tool_id="tool_1", + tool_name="echo", + tool_kwargs={"value": "x"}, + ) + ], + ["done"], + ] + ), + tool_scheduler=_FakeAsyncToolScheduler(), + request_interceptors=[recorder], + ) + + before_iteration_roles = [ + roles + for phase, roles in recorder.observations + if phase == InterceptorPhase.BEFORE_ITERATION + ] + + assert before_iteration_roles == [ + [MessageRole.USER], + [MessageRole.USER, MessageRole.ASSISTANT, MessageRole.TOOL], + ] + + +@pytest.mark.asyncio +async def test_async_engine_rebuilds_runtime_before_condensation_after_tool_resume( + base_request: ResolvedChatRequest, +) -> None: + request = base_request.model_copy(deep=True) + request.system.model = "model-a" + request.tool_config = ResolvedToolConfig(tools=[_server_tool("echo")]) + mock_llm = get_mock_function_calling_llm( + [ + [ + ToolSelection( + tool_id="tool_1", + tool_name="echo", + tool_kwargs={"value": "x"}, + ) + ], + ["done"], + ] + ) + llm_component = _make_llm_component(mock_llm) + llm_component.get_tokenizer.return_value = lambda text: list(text) + runtime_interceptor = RuntimeModelRequestInterceptor(llm_component) + condensation_observer = _RuntimeRecordingInterceptor(observations=[]) + + await _run_async_engine( + request, + mock_llm, + tool_scheduler=_FakeAsyncToolScheduler(), + request_interceptors=[runtime_interceptor, condensation_observer], + llm_component=llm_component, + ) + + before_iteration = [ + observation + for observation in condensation_observer.observations + if observation.phase == InterceptorPhase.BEFORE_ITERATION + ] + + assert len(before_iteration) == 2 + assert all(observation.model_id == "model-a" for observation in before_iteration) + assert all( + observation.effective_token_limit is not None + for observation in before_iteration + ) + assert all(observation.has_tokenizer for observation in before_iteration) + assert llm_component.get_tokenizer.call_count == 3 + + @pytest.mark.asyncio async def test_async_engine_matches_sync_two_server_tools_plus_one_client_tool( base_request: ResolvedChatRequest, diff --git a/tests/engines/test_resumable_runtime.py b/tests/engines/test_resumable_runtime.py index 3b48736e..0a1a7351 100644 --- a/tests/engines/test_resumable_runtime.py +++ b/tests/engines/test_resumable_runtime.py @@ -371,6 +371,7 @@ async def test_runner_schedules_one_timeout_timer_per_pending_tool() -> None: state.input.context_stack.checkpoint_dump.return_value = {} state.runtime.iteration = 1 state.runtime.next_block_count = 3 + state.runtime.model_id = "default" state.runtime.total_input_tokens = 0 state.runtime.total_output_tokens = 0 state.runtime.has_input_usage = False @@ -575,6 +576,7 @@ async def test_cancelled_chat_cannot_resurrect_waiting_checkpoint() -> None: state.input.context_stack.checkpoint_dump.return_value = {} state.runtime.iteration = 1 state.runtime.next_block_count = 0 + state.runtime.model_id = "default" state.runtime.total_input_tokens = 0 state.runtime.total_output_tokens = 0 state.runtime.has_input_usage = False diff --git a/tests/server/chat/interceptors/test_runtime_model_interceptor.py b/tests/server/chat/interceptors/test_runtime_model_interceptor.py new file mode 100644 index 00000000..7615f5d4 --- /dev/null +++ b/tests/server/chat/interceptors/test_runtime_model_interceptor.py @@ -0,0 +1,120 @@ +from unittest.mock import MagicMock + +import pytest +from llama_index.core.base.llms.types import ChatMessage, LLMMetadata, MessageRole + +from private_gpt.components.chat.models.chat_config_models import ( + ResolvedChatRequest, + ResolvedSystemConfig, +) +from private_gpt.components.engines.chat.models.chat_interceptor_context import ( + ChatInterceptorContext, +) +from private_gpt.components.engines.chat.models.chat_phase import InterceptorPhase +from private_gpt.components.engines.chat.models.chat_state import ( + ChatInputState, + ChatOutputState, + ChatRuntimeState, + ChatState, +) +from private_gpt.components.engines.chat.resumable_runner import ResumableChatRunner +from private_gpt.server.chat.interceptors.runtime_model_interceptor import ( + RuntimeModelRequestInterceptor, +) +from tests.fixtures.mock_function_llm import get_mock_function_calling_llm + + +def _context( + *, + phase: InterceptorPhase, + model_id: str | None = "model-a", +) -> tuple[ChatInterceptorContext, MagicMock, MagicMock]: + tokenizer = MagicMock() + context_llm = get_mock_function_calling_llm(["ok"]) + llm = MagicMock() + llm.metadata = LLMMetadata(context_window=131_072, num_output=4_096) + llm_component = MagicMock() + llm_component.get_llm.return_value = llm + llm_component.get_tokenizer.return_value = tokenizer + request = ResolvedChatRequest( + messages=[ChatMessage(role=MessageRole.USER, content="hello")], + system=ResolvedSystemConfig(model=model_id), + ) + state = ChatState( + input=ChatInputState(request=request), + runtime=ChatRuntimeState(), + output=ChatOutputState(), + ) + return ( + ChatInterceptorContext( + state=state, + llm=context_llm, + phase=phase, + emit_fn=lambda _event: None, + ), + llm_component, + tokenizer, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "phase", + [InterceptorPhase.VALIDATION, InterceptorPhase.BEFORE_ITERATION], +) +async def test_runtime_model_interceptor_hydrates_model_runtime( + phase: InterceptorPhase, +) -> None: + context, llm_component, tokenizer = _context(phase=phase) + + await RuntimeModelRequestInterceptor(llm_component).intercept(context) + + assert context.state.runtime.model_id == "model-a" + assert context.state.runtime.effective_token_limit == 126_720 + assert context.state.runtime.tokenizer_fn is tokenizer + llm_component.get_tokenizer.assert_called_once_with("model-a") + + +@pytest.mark.asyncio +async def test_runtime_model_interceptor_rebuilds_process_local_fields_from_model_id() -> ( + None +): + context, llm_component, tokenizer = _context( + phase=InterceptorPhase.BEFORE_ITERATION + ) + context.state.runtime.model_id = "persisted-model" + + await RuntimeModelRequestInterceptor(llm_component).intercept(context) + + assert context.state.runtime.effective_token_limit == 126_720 + assert context.state.runtime.tokenizer_fn is tokenizer + llm_component.get_tokenizer.assert_called_once_with("persisted-model") + + +@pytest.mark.asyncio +async def test_runtime_model_interceptor_skips_already_hydrated_runtime() -> None: + context, llm_component, tokenizer = _context( + phase=InterceptorPhase.BEFORE_ITERATION + ) + context.state.runtime.model_id = "model-a" + context.state.runtime.effective_token_limit = 100_000 + context.state.runtime.tokenizer_fn = tokenizer + + await RuntimeModelRequestInterceptor(llm_component).intercept(context) + + assert context.state.runtime.effective_token_limit == 100_000 + llm_component.get_tokenizer.assert_not_called() + + +def test_checkpoint_payload_persists_only_model_identity() -> None: + context, _, tokenizer = _context(phase=InterceptorPhase.BEFORE_ITERATION) + context.state.runtime.model_id = "persisted-model" + context.state.runtime.effective_token_limit = 100_000 + context.state.runtime.tokenizer_fn = tokenizer + + payload = ResumableChatRunner._checkpoint_payload(context.state) + serialized = payload.model_dump(mode="json") + + assert serialized["model_id"] == "persisted-model" + assert "tokenizer_fn" not in serialized + assert "effective_token_limit" not in serialized diff --git a/tests/server/chat/interceptors/test_validation_request_interceptor.py b/tests/server/chat/interceptors/test_validation_request_interceptor.py index 6a4a64ae..0eaafea3 100644 --- a/tests/server/chat/interceptors/test_validation_request_interceptor.py +++ b/tests/server/chat/interceptors/test_validation_request_interceptor.py @@ -34,6 +34,9 @@ from private_gpt.components.engines.chat.utils.request_builder import ( build_initial_context_stack, ) from private_gpt.events.event_errors import Errors +from private_gpt.server.chat.interceptors.runtime_model_interceptor import ( + RuntimeModelRequestInterceptor, +) from private_gpt.server.chat.interceptors.validator_request_interceptor import ( ValidatorRequestInterceptor, ) @@ -141,6 +144,7 @@ async def _run_interceptor( phase=InterceptorPhase.VALIDATION, emit_fn=lambda _event: None, ) + await RuntimeModelRequestInterceptor(interceptor._llm_component).intercept(context) await interceptor.intercept(context)