Cannot inspect schema when working with special type annotations

[eyurtsev]

Privileged issue

• I am a LangChain maintainer, or was asked directly by a LangChain maintainer to create an issue here.

Issue Content

schema inspection + special type annotations

from langchain_core.tools import tool, InjectedToolArg

@tool
def user_specific_tool(input_data: str, user_id: InjectedToolArg) -> str:
    """Some tool"""
    return f"User {user_id} processed {input_data}"

user_specific_tool.args

[chloepan33]

Hey @eyurtsev! We are a group of students from UofT interested in this issue. Is it possible we could know more details of this issue and the desire output?

[eyurtsev]

A tool is an association between:

1. underlying function
2. a schema that describes the function name, description and its inputs

There are use cases, where some of the tools inputs should not be controllable by the chat model, but dictated by the application logic. (e.g., user_id -- should be associated with the user logged into the application and not chosen by the chat model as that would be a (very serious) security vulnerability)

We use special type annotations to mark arguments as "hidden" from the chat model. This allows us to use the @tool decorator to easily generate a schema that can be presented to a chat model without surfacing other parameters that the chat model should not know about (these parameters will be injected during run time).

FYI I haven't triaged this issue yet to determine whether it manifests itself when used with chat models (there might be another code path used) -- i don't know if this is a high priority issue yet or not, if it's high priority we will likely tackle within a few days time

[chloepan33]

hey @eyurtsev,

To confirm our understanding of the issue:

The goal is that when we use:

from langchain_core.tools import tool, InjectedToolArg

@tool
def user_specific_tool(input_data: str, user_id: InjectedToolArg) -> str:
    """Some tool"""
    return f"User {user_id} processed {input_data}"

The schema generated by user_specific_tool.args should only include the input_data argument, without any reference to user_id, since user_id is meant to be injected at runtime and not controlled by the chat model.

Additionally, if we define:

def call_tool(input_data: str, user_id: str):
    return user_specific_tool.run({"input_data": input_data}, user_id=user_id)

result = call_tool("some data", "12345")
print(result)

The expected output should be:

User 12345 processed some data

---

Loop in team members @Chloekyuu @JANERUBBISHTOEAT @XiaoConan

[JANERUBBISHTOEAT]

Our team has determined that modifying the base.py::_get_filtered_args() and base.py::_is_injected_arg_type methods, along with updating the base.py::InjectedToolArg class definition, would address this issue. Does this approach seem appropriate?

ccing team: @chloepan33 @Chloekyuu @XiaoConan

[XiaoConan]

Hi @eyurtsev,

Our team had identified that the issue arises because the InjectedToolArg type annotation is not compatible with Pydantic's JSON schema generation. This leads to two primary problems:

1. Schema Generation Failure: When InjectedToolArg needs to be included as an argument in a tool or model, attempts to generate its JSON schema result in errors (PydanticInvalidForJsonSchema), which makes tools or models with InjectedToolArg arguments uninspectable or imcompatible with workflows requireing schema validation.
2. Argument Filtering Logic: The current _is_injected_arg_type function fails to detect runtime-injected arguments like InjectedToolArg, which couldn't be filtering out during the create schema process when we need to exclude such arguments.

To address these issues, we propose the following solutions:

1. We create a function _check_injected_arg_type to replace the functionality of _is_injected_arg_type, which will correctly identify InjectedToolArg Arguments. This can be helpful in the filtering our runtime-injected argument process.
2. We defined InjectedToolArgSchema that wraps InjectedToolArg in Annotated with a custom schema defined via WithJsonSchema, which allows InjectedToolArg to be parsed into a JSON schema format without causing errors, enabling proper inspection and compatibility with Pydantic.

We created a PR here: #28435

Here's a testing file to observe the actual behavior:

from langchain_core.tools import tool, InjectedToolArgSchema

@tool
def user_specific_tool(input_data: str, user_id: InjectedToolArgSchema) -> str:
    """Some tool"""
    return f"User {user_id} processed {input_data}"

user_specific_tool.args

Loop in team members @Chloekyuu @JANERUBBISHTOEAT @chloepan33