github/privateGPT
1
0
Fork 1
mirror of https://github.com/imartinez/privateGPT.git synced 2026-07-18 14:04:29 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix: tools

Browse Source
This commit is contained in:
Javier Martinez
2026-07-08 15:51:07 +02:00
parent 9022d5fadf
commit 9ad2d3e534
5 changed files with 188 additions and 105 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
private_gpt/components/engines/chat_loop/models/chat_loop_phase.py
View File

@@ -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"

107
private_gpt/components/tools/remote_execution.py
View File

@@ -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]:

32
private_gpt/server/chat/interceptors/configure_tool_execution_interceptor.py Normal file
View File

@@ -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)

66
private_gpt/server/chat/interceptors/null_tool_values_interceptor.py
View File

@@ -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
}
)

86
private_gpt/server/chat/interceptors/schema_coercing_tool_interceptor.py
View File

@@ -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