From 9ad2d3e534a5fa588e03544befbfdeeed115c152 Mon Sep 17 00:00:00 2001 From: Javier Martinez Date: Wed, 8 Jul 2026 15:51:07 +0200 Subject: [PATCH] fix: tools --- .../chat_loop/models/chat_loop_phase.py | 2 + .../components/tools/remote_execution.py | 107 ++++++++++++++---- .../configure_tool_execution_interceptor.py | 32 ++++++ .../null_tool_values_interceptor.py | 66 +++++------ .../schema_coercing_tool_interceptor.py | 86 ++++++-------- 5 files changed, 188 insertions(+), 105 deletions(-) create mode 100644 private_gpt/server/chat/interceptors/configure_tool_execution_interceptor.py diff --git a/private_gpt/components/engines/chat_loop/models/chat_loop_phase.py b/private_gpt/components/engines/chat_loop/models/chat_loop_phase.py index bd1c7fcb..71662654 100644 --- a/private_gpt/components/engines/chat_loop/models/chat_loop_phase.py +++ b/private_gpt/components/engines/chat_loop/models/chat_loop_phase.py @@ -6,7 +6,9 @@ class InterceptorPhase(StrEnum): VALIDATION = "validation" BEFORE_ITERATION = "before_iteration" + BEFORE_TOOL = "before_tool" STREAMING = "streaming" + AFTER_TOOL = "after_tool" AFTER_ITERATION = "after_iteration" diff --git a/private_gpt/components/tools/remote_execution.py b/private_gpt/components/tools/remote_execution.py index 2e3bc977..c6f72a75 100644 --- a/private_gpt/components/tools/remote_execution.py +++ b/private_gpt/components/tools/remote_execution.py @@ -2,8 +2,10 @@ from __future__ import annotations import importlib import inspect +from abc import ABC, abstractmethod from typing import TYPE_CHECKING, Any +from injector import inject, singleton from llama_index.core.base.llms.types import ChatMessage from llama_index.core.tools import adapt_to_async_tool from pydantic import BaseModel, Field @@ -12,6 +14,9 @@ from private_gpt.components.chat.models.chat_config_models import ( ToolExecutionMetadata, ToolSpec, ) +from private_gpt.components.engines.chat_loop.models.chat_loop_phase import ( + InterceptorPhase, +) from private_gpt.components.engines.chat_loop.utils.tool_utils import execute_tool_call from private_gpt.events.models import ( ResultContentBlockType, @@ -25,6 +30,9 @@ if TYPE_CHECKING: from private_gpt.components.engines.chat_loop.models.chat_loop_state import ( ChatLoopState, ) + from private_gpt.server.chat.interceptors.configure_tool_execution_interceptor import ( + ConfigureToolExecutionInterceptor, + ) class ToolExecutionRequest(BaseModel): @@ -43,6 +51,81 @@ class ToolExecutionResponse(BaseModel): tool_message: ChatMessage +class ToolExecutionInterceptorContext(BaseModel): + phase: InterceptorPhase + request: ToolExecutionRequest + tool_kwargs: dict[str, Any] + response: ToolExecutionResponse | None = None + + def set_tool_kwargs(self, tool_kwargs: dict[str, Any]) -> None: + self.tool_kwargs = tool_kwargs + + def set_response(self, response: ToolExecutionResponse) -> None: + self.response = response + + +class ToolExecutionInterceptor(ABC): + @abstractmethod + async def intercept(self, context: ToolExecutionInterceptorContext) -> None: + """Mutate tool execution context before/after tool invocation.""" + + +@singleton +class ToolExecutor: + @inject + def __init__( + self, + configure_tool_execution_interceptor: "ConfigureToolExecutionInterceptor", + ) -> None: + self._configure_tool_execution_interceptor = ( + configure_tool_execution_interceptor + ) + + async def execute( + self, + request: ToolExecutionRequest, + state_ctx: ChatLoopState | None = None, + ) -> ToolExecutionResponse: + tool = await rebuild_tool_from_spec(request.tool_spec) + + before_context = ToolExecutionInterceptorContext( + phase=InterceptorPhase.BEFORE_TOOL, + request=request, + tool_kwargs=dict(request.tool_kwargs), + ) + await self._configure_tool_execution_interceptor.intercept(before_context) + + result, tool_message = await execute_tool_call( + tool=tool, + tool_name=request.tool_name, + tool_id=request.tool_id, + tool_kwargs=before_context.tool_kwargs, + state_ctx=state_ctx, + ) + response = ToolExecutionResponse( + tool_name=request.tool_name, + tool_id=request.tool_id, + result_content=( + from_tool_output(result.tool_output.raw_output) + if result.tool_output.raw_output is not None + else [TextBlock(text=result.tool_output.content or "")] + ), + is_error=result.tool_output.is_error, + tool_message=tool_message, + ) + + after_context = ToolExecutionInterceptorContext( + phase=InterceptorPhase.AFTER_TOOL, + request=request, + tool_kwargs=before_context.tool_kwargs, + response=response, + ) + await self._configure_tool_execution_interceptor.intercept(after_context) + + assert after_context.response is not None + return after_context.response + + def build_rebuild_metadata( rebuild_callable: Any, rebuild_kwargs: dict[str, Any] | None = None, @@ -66,26 +149,10 @@ async def execute_tool_request( request: ToolExecutionRequest, state_ctx: ChatLoopState | None = None, ) -> ToolExecutionResponse: - tool = await rebuild_tool_from_spec(request.tool_spec) - result, tool_message = await execute_tool_call( - tool=tool, - tool_name=request.tool_name, - tool_id=request.tool_id, - tool_kwargs=request.tool_kwargs, - state_ctx=state_ctx, - ) - result_content = ( - from_tool_output(result.tool_output.raw_output) - if result.tool_output.raw_output is not None - else [TextBlock(text=result.tool_output.content or "")] - ) - return ToolExecutionResponse( - tool_name=request.tool_name, - tool_id=request.tool_id, - result_content=result_content, - is_error=result.tool_output.is_error, - tool_message=tool_message, - ) + from private_gpt.di import get_global_injector + + executor = get_global_injector().get(ToolExecutor) + return await executor.execute(request, state_ctx=state_ctx) def build_tool_execution_context(state: ChatLoopState) -> dict[str, Any]: diff --git a/private_gpt/server/chat/interceptors/configure_tool_execution_interceptor.py b/private_gpt/server/chat/interceptors/configure_tool_execution_interceptor.py new file mode 100644 index 00000000..be3d72fb --- /dev/null +++ b/private_gpt/server/chat/interceptors/configure_tool_execution_interceptor.py @@ -0,0 +1,32 @@ +from injector import inject, singleton + +from private_gpt.components.tools.remote_execution import ( + ToolExecutionInterceptor, + ToolExecutionInterceptorContext, +) +from private_gpt.server.chat.interceptors.null_tool_values_interceptor import ( + NullToolValuesRequestInterceptor, +) +from private_gpt.server.chat.interceptors.schema_coercing_tool_interceptor import ( + SchemaCoercingToolInterceptor, +) + + +@singleton +class ConfigureToolExecutionInterceptor(ToolExecutionInterceptor): + """Aggregate tool-execution sub-interceptors into a single step.""" + + @inject + def __init__( + self, + null_tool_values_interceptor: NullToolValuesRequestInterceptor, + schema_coercing_interceptor: SchemaCoercingToolInterceptor, + ) -> None: + self._interceptors: list[ToolExecutionInterceptor] = [ + null_tool_values_interceptor, + schema_coercing_interceptor, + ] + + async def intercept(self, context: ToolExecutionInterceptorContext) -> None: + for interceptor in self._interceptors: + await interceptor.intercept(context) diff --git a/private_gpt/server/chat/interceptors/null_tool_values_interceptor.py b/private_gpt/server/chat/interceptors/null_tool_values_interceptor.py index 30ce1600..7497ee03 100644 --- a/private_gpt/server/chat/interceptors/null_tool_values_interceptor.py +++ b/private_gpt/server/chat/interceptors/null_tool_values_interceptor.py @@ -1,50 +1,46 @@ -from typing import Any +from __future__ import annotations + +from typing import TYPE_CHECKING from injector import singleton -from private_gpt.components.chat.models.chat_config_models import ToolSpec -from private_gpt.components.context.models.context_layer import ToolDefinitionsLayer -from private_gpt.components.context.models.layer_type import LayerType from private_gpt.components.engines.chat_loop.interceptors.chat_loop_interceptor import ( ChatRequestLoopInterceptor, ) -from private_gpt.components.engines.chat_loop.models.chat_loop_interceptor_context import ( - ChatLoopInterceptorContext, -) from private_gpt.components.engines.chat_loop.models.chat_loop_phase import ( InterceptorPhase, ) +from private_gpt.components.tools.remote_execution import ( + ToolExecutionInterceptor, + ToolExecutionInterceptorContext, +) + +if TYPE_CHECKING: + from private_gpt.components.engines.chat_loop.models.chat_loop_interceptor_context import ( + ChatLoopInterceptorContext, + ) @singleton -class NullToolValuesRequestInterceptor(ChatRequestLoopInterceptor): - """Patch tool specs to strip None kwargs before async invocation.""" +class NullToolValuesRequestInterceptor( + ChatRequestLoopInterceptor, + ToolExecutionInterceptor, +): + """Strip ``None``-valued kwargs before tool execution.""" - async def intercept(self, context: ChatLoopInterceptorContext) -> None: - if context.phase != InterceptorPhase.BEFORE_ITERATION: + async def intercept( + self, + context: ChatLoopInterceptorContext | ToolExecutionInterceptorContext, + ) -> None: + if not isinstance(context, ToolExecutionInterceptorContext): + return + if context.phase != InterceptorPhase.BEFORE_TOOL: return - state = context.state - tools = [ - self._patch_tool(tool) for tool in state.input.context_stack.all_tools() - ] - - stack = state.input.context_stack - stack = stack.remove_layers_of_type(LayerType.TOOL_DEFINITIONS) - if tools: - stack = stack.append_layer( - ToolDefinitionsLayer(tools=tools, source="tool_patch") - ) - state.input.context_stack = stack - context.set_state(state) - - def _patch_tool(self, tool: ToolSpec) -> ToolSpec: - """Patch the tool to avoid to call using optional None params.""" - - async def async_patched_tool(*args: Any, **kwargs: Any) -> Any: - new_kwargs = {k: v for k, v in kwargs.items() if v is not None} - return await tool.async_fn(*args, **new_kwargs) - - tool_copy = tool.model_copy(deep=True) - tool_copy.async_fn = async_patched_tool - return tool_copy + context.set_tool_kwargs( + { + key: value + for key, value in context.tool_kwargs.items() + if value is not None + } + ) diff --git a/private_gpt/server/chat/interceptors/schema_coercing_tool_interceptor.py b/private_gpt/server/chat/interceptors/schema_coercing_tool_interceptor.py index f51994f7..432a9a77 100644 --- a/private_gpt/server/chat/interceptors/schema_coercing_tool_interceptor.py +++ b/private_gpt/server/chat/interceptors/schema_coercing_tool_interceptor.py @@ -1,24 +1,29 @@ +from __future__ import annotations + import ast import contextlib import json import logging import math -from typing import Any +from typing import TYPE_CHECKING, Any from injector import singleton -from private_gpt.components.chat.models.chat_config_models import ToolSpec -from private_gpt.components.context.models.context_layer import ToolDefinitionsLayer -from private_gpt.components.context.models.layer_type import LayerType from private_gpt.components.engines.chat_loop.interceptors.chat_loop_interceptor import ( ChatRequestLoopInterceptor, ) -from private_gpt.components.engines.chat_loop.models.chat_loop_interceptor_context import ( - ChatLoopInterceptorContext, -) from private_gpt.components.engines.chat_loop.models.chat_loop_phase import ( InterceptorPhase, ) +from private_gpt.components.tools.remote_execution import ( + ToolExecutionInterceptor, + ToolExecutionInterceptorContext, +) + +if TYPE_CHECKING: + from private_gpt.components.engines.chat_loop.models.chat_loop_interceptor_context import ( + ChatLoopInterceptorContext, + ) logger = logging.getLogger(__name__) @@ -86,10 +91,6 @@ def _parse_literal_string( raw: str, expected: type | tuple[type, ...], ) -> Any: - """Parse a string as JSON, falling back to ast.literal_eval for Python repr. - - Returns the parsed value if it is an instance of `expected`, else None. - """ with contextlib.suppress(json.JSONDecodeError, ValueError): parsed = json.loads(raw) if isinstance(parsed, expected): @@ -290,46 +291,31 @@ def _coerce_kwargs( @singleton -class SchemaCoercingToolInterceptor(ChatRequestLoopInterceptor): - """Coerce tool kwargs to match the declared input_schema before invocation.""" +class SchemaCoercingToolInterceptor( + ChatRequestLoopInterceptor, + ToolExecutionInterceptor, +): + """Coerce tool kwargs to the declared schema before execution.""" - async def intercept(self, context: ChatLoopInterceptorContext) -> None: - if context.phase != InterceptorPhase.BEFORE_ITERATION: + async def intercept( + self, + context: ChatLoopInterceptorContext | ToolExecutionInterceptorContext, + ) -> None: + if not isinstance(context, ToolExecutionInterceptorContext): + return + if context.phase != InterceptorPhase.BEFORE_TOOL: return - state = context.state - tools = [ - self._patch_tool(tool) for tool in state.input.context_stack.all_tools() - ] - - stack = state.input.context_stack.remove_layers_of_type( - LayerType.TOOL_DEFINITIONS - ) - if tools: - stack = stack.append_layer( - ToolDefinitionsLayer(tools=tools, source="schema_coercion_patch") + schema = context.request.tool_spec.input_schema or {} + try: + context.set_tool_kwargs( + _coerce_kwargs(context.tool_kwargs, input_schema=schema) + ) + except SchemaCoercionError: + raise + except Exception as e: + logger.exception( + "Schema coercion failed for tool '%s', invoking with original kwargs", + context.request.tool_spec.name, + exc_info=e, ) - state.input.context_stack = stack - context.set_state(state) - - def _patch_tool(self, tool: ToolSpec) -> ToolSpec: - schema = tool.input_schema or {} - original_fn = tool.async_fn - - async def _coerced_fn(*args: Any, **kwargs: Any) -> Any: - try: - fixed_kwargs = _coerce_kwargs(kwargs, input_schema=schema) - except SchemaCoercionError: - raise - except Exception as e: - logger.exception( - "Schema coercion failed for tool '%s', invoking with original kwargs", - tool.name, - exc_info=e, - ) - fixed_kwargs = kwargs - return await original_fn(*args, **fixed_kwargs) - - patched = tool.model_copy(deep=True) - patched.async_fn = _coerced_fn - return patched