Skip to content

Commit c8fc585

Browse files
christian-bromannchristian-bromann
committed
feat(@langchain/core): support tools with custom state or context provided by ToolRuntime
1 parent dcffe06 commit c8fc585

File tree

2 files changed

+302
-25
lines changed

2 files changed

+302
-25
lines changed

‎libs/langchain-core/src/tools/index.ts‎

Lines changed: 205 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,6 @@ import {
1717
pickRunnableConfigKeys,
1818
type RunnableConfig,
1919
} from "../runnables/config.js";
20-
import type { RunnableFunc } from "../runnables/base.js";
2120
import { isDirectToolOutput, ToolCall, ToolMessage } from "../messages/tool.js";
2221
import { AsyncLocalStorageProviderSingleton } from "../singletons/index.js";
2322
import {
@@ -54,6 +53,7 @@ import type {
5453
StringInputToolSchema,
5554
ToolInterface,
5655
ToolOutputType,
56+
ToolRuntime,
5757
} from "./types.js";
5858
import { type JSONSchema, validatesOnlyStrings } from "../utils/json_schema.js";
5959

@@ -71,6 +71,7 @@ export type {
7171
ToolReturnType,
7272
ToolRunnableConfig,
7373
ToolInputSchemaBase as ToolSchemaBase,
74+
ToolRuntime,
7475
} from "./types.js";
7576

7677
export {
@@ -511,14 +512,64 @@ export abstract class BaseToolkit {
511512
}
512513
}
513514

515+
/**
516+
* Helper type to check if a schema is defined (not undefined).
517+
*/
518+
type IsSchemaDefined<T> = T extends undefined ? false : true;
519+
520+
/**
521+
* Helper type to determine if runtime should be passed to the function.
522+
*/
523+
type ShouldPassRuntime<
524+
StateSchema extends InteropZodObject | undefined,
525+
ContextSchema extends InteropZodObject | undefined
526+
> = IsSchemaDefined<StateSchema> extends true
527+
? true
528+
: IsSchemaDefined<ContextSchema> extends true
529+
? true
530+
: false;
531+
532+
/**
533+
* Helper type to create RunnableFunc with optional runtime parameter.
534+
*/
535+
type RunnableFuncWithRuntime<
536+
RunInput,
537+
RunOutput,
538+
StateSchema extends InteropZodObject | undefined,
539+
ContextSchema extends InteropZodObject | undefined,
540+
CallOptions extends RunnableConfig = RunnableConfig
541+
> = ShouldPassRuntime<StateSchema, ContextSchema> extends true
542+
? (
543+
input: RunInput,
544+
runtime: ToolRuntime<StateSchema, ContextSchema>,
545+
options?:
546+
| CallOptions
547+
// eslint-disable-next-line @typescript-eslint/no-explicit-any
548+
| Record<string, any>
549+
// eslint-disable-next-line @typescript-eslint/no-explicit-any
550+
| (Record<string, any> & CallOptions)
551+
) => RunOutput | Promise<RunOutput>
552+
: (
553+
input: RunInput,
554+
options?:
555+
| CallOptions
556+
// eslint-disable-next-line @typescript-eslint/no-explicit-any
557+
| Record<string, any>
558+
// eslint-disable-next-line @typescript-eslint/no-explicit-any
559+
| (Record<string, any> & CallOptions)
560+
) => RunOutput | Promise<RunOutput>;
561+
514562
/**
515563
* Parameters for the tool function.
516564
* Schema can be provided as Zod or JSON schema.
517565
* Both schema types will be validated.
518566
* @template {ToolInputSchemaBase} RunInput The input schema for the tool.
519567
*/
520-
interface ToolWrapperParams<RunInput = ToolInputSchemaBase | undefined>
521-
extends ToolParams {
568+
interface ToolWrapperParams<
569+
RunInput = ToolInputSchemaBase | undefined,
570+
StateSchema extends InteropZodObject | undefined = undefined,
571+
ContextSchema extends InteropZodObject | undefined = undefined
572+
> extends ToolParams {
522573
/**
523574
* The name of the tool. If using with an LLM, this
524575
* will be passed as the tool name.
@@ -552,6 +603,14 @@ interface ToolWrapperParams<RunInput = ToolInputSchemaBase | undefined>
552603
* an agent should stop looping.
553604
*/
554605
returnDirect?: boolean;
606+
/**
607+
* The state schema for the tool runtime.
608+
*/
609+
stateSchema?: StateSchema;
610+
/**
611+
* The context schema for the tool runtime.
612+
*/
613+
contextSchema?: ContextSchema;
555614
}
556615

557616
/**
@@ -562,6 +621,8 @@ interface ToolWrapperParams<RunInput = ToolInputSchemaBase | undefined>
562621
* @function
563622
* @template {ToolInputSchemaBase} SchemaT The input schema for the tool.
564623
* @template {ToolReturnType} ToolOutputT The output type of the tool.
624+
* @template {InteropZodObject | undefined} StateSchema The state schema for the tool runtime.
625+
* @template {InteropZodObject | undefined} ContextSchema The context schema for the tool runtime.
565626
*
566627
* @param {RunnableFunc<z.output<SchemaT>, ToolOutputT>} func - The function to invoke when the tool is called.
567628
* @param {ToolWrapperParams<SchemaT>} fields - An object containing the following properties:
@@ -571,56 +632,90 @@ interface ToolWrapperParams<RunInput = ToolInputSchemaBase | undefined>
571632
*
572633
* @returns {DynamicStructuredTool<SchemaT>} A new StructuredTool instance.
573634
*/
574-
export function tool<SchemaT extends ZodStringV3, ToolOutputT = ToolOutputType>(
575-
func: RunnableFunc<
635+
export function tool<
636+
SchemaT extends ZodStringV3,
637+
ToolOutputT = ToolOutputType,
638+
StateSchema extends InteropZodObject | undefined = undefined,
639+
ContextSchema extends InteropZodObject | undefined = undefined
640+
>(
641+
func: RunnableFuncWithRuntime<
576642
InferInteropZodOutput<SchemaT>,
577643
ToolOutputT,
644+
StateSchema,
645+
ContextSchema,
578646
ToolRunnableConfig
579647
>,
580-
fields: ToolWrapperParams<SchemaT>
648+
fields: ToolWrapperParams<SchemaT, StateSchema, ContextSchema>
581649
): DynamicTool<ToolOutputT>;
582650

583-
export function tool<SchemaT extends ZodStringV4, ToolOutputT = ToolOutputType>(
584-
func: RunnableFunc<
651+
export function tool<
652+
SchemaT extends ZodStringV4,
653+
ToolOutputT = ToolOutputType,
654+
StateSchema extends InteropZodObject | undefined = undefined,
655+
ContextSchema extends InteropZodObject | undefined = undefined
656+
>(
657+
func: RunnableFuncWithRuntime<
585658
InferInteropZodOutput<SchemaT>,
586659
ToolOutputT,
660+
StateSchema,
661+
ContextSchema,
587662
ToolRunnableConfig
588663
>,
589-
fields: ToolWrapperParams<SchemaT>
664+
fields: ToolWrapperParams<SchemaT, StateSchema, ContextSchema>
590665
): DynamicTool<ToolOutputT>;
591666

592667
export function tool<
593668
SchemaT extends ZodObjectV3,
594669
SchemaOutputT = InferInteropZodOutput<SchemaT>,
595670
SchemaInputT = InferInteropZodInput<SchemaT>,
596-
ToolOutputT = ToolOutputType
671+
ToolOutputT = ToolOutputType,
672+
StateSchema extends InteropZodObject | undefined = undefined,
673+
ContextSchema extends InteropZodObject | undefined = undefined
597674
>(
598-
func: RunnableFunc<SchemaOutputT, ToolOutputT, ToolRunnableConfig>,
599-
fields: ToolWrapperParams<SchemaT>
675+
func: RunnableFuncWithRuntime<
676+
SchemaOutputT,
677+
ToolOutputT,
678+
StateSchema,
679+
ContextSchema,
680+
ToolRunnableConfig
681+
>,
682+
fields: ToolWrapperParams<SchemaT, StateSchema, ContextSchema>
600683
): DynamicStructuredTool<SchemaT, SchemaOutputT, SchemaInputT, ToolOutputT>;
601684

602685
export function tool<
603686
SchemaT extends ZodObjectV4,
604687
SchemaOutputT = InferInteropZodOutput<SchemaT>,
605688
SchemaInputT = InferInteropZodInput<SchemaT>,
606-
ToolOutputT = ToolOutputType
689+
ToolOutputT = ToolOutputType,
690+
StateSchema extends InteropZodObject | undefined = undefined,
691+
ContextSchema extends InteropZodObject | undefined = undefined
607692
>(
608-
func: RunnableFunc<SchemaOutputT, ToolOutputT, ToolRunnableConfig>,
609-
fields: ToolWrapperParams<SchemaT>
693+
func: RunnableFuncWithRuntime<
694+
SchemaOutputT,
695+
ToolOutputT,
696+
StateSchema,
697+
ContextSchema,
698+
ToolRunnableConfig
699+
>,
700+
fields: ToolWrapperParams<SchemaT, StateSchema, ContextSchema>
610701
): DynamicStructuredTool<SchemaT, SchemaOutputT, SchemaInputT, ToolOutputT>;
611702

612703
export function tool<
613704
SchemaT extends JSONSchema,
614705
SchemaOutputT = ToolInputSchemaOutputType<SchemaT>,
615706
SchemaInputT = ToolInputSchemaInputType<SchemaT>,
616-
ToolOutputT = ToolOutputType
707+
ToolOutputT = ToolOutputType,
708+
StateSchema extends InteropZodObject | undefined = undefined,
709+
ContextSchema extends InteropZodObject | undefined = undefined
617710
>(
618-
func: RunnableFunc<
711+
func: RunnableFuncWithRuntime<
619712
Parameters<DynamicStructuredToolInput<SchemaT>["func"]>[0],
620713
ToolOutputT,
714+
StateSchema,
715+
ContextSchema,
621716
ToolRunnableConfig
622717
>,
623-
fields: ToolWrapperParams<SchemaT>
718+
fields: ToolWrapperParams<SchemaT, StateSchema, ContextSchema>
624719
): DynamicStructuredTool<SchemaT, SchemaOutputT, SchemaInputT, ToolOutputT>;
625720

626721
export function tool<
@@ -630,15 +725,26 @@ export function tool<
630725
| JSONSchema = InteropZodObject,
631726
SchemaOutputT = ToolInputSchemaOutputType<SchemaT>,
632727
SchemaInputT = ToolInputSchemaInputType<SchemaT>,
633-
ToolOutputT = ToolOutputType
728+
ToolOutputT = ToolOutputType,
729+
StateSchema extends InteropZodObject | undefined = undefined,
730+
ContextSchema extends InteropZodObject | undefined = undefined
634731
>(
635-
func: RunnableFunc<SchemaOutputT, ToolOutputT, ToolRunnableConfig>,
636-
fields: ToolWrapperParams<SchemaT>
732+
func: RunnableFuncWithRuntime<
733+
SchemaOutputT,
734+
ToolOutputT,
735+
StateSchema,
736+
ContextSchema,
737+
ToolRunnableConfig
738+
>,
739+
fields: ToolWrapperParams<SchemaT, StateSchema, ContextSchema>
637740
):
638741
| DynamicStructuredTool<SchemaT, SchemaOutputT, SchemaInputT, ToolOutputT>
639742
| DynamicTool<ToolOutputT> {
640743
const isSimpleStringSchema = isSimpleStringZodSchema(fields.schema);
641744
const isStringJSONSchema = validatesOnlyStrings(fields.schema);
745+
const hasStateSchema = fields.stateSchema !== undefined;
746+
const hasContextSchema = fields.contextSchema !== undefined;
747+
const shouldPassRuntime = hasStateSchema || hasContextSchema;
642748

643749
// If the schema is not provided, or it's a simple string schema, create a DynamicTool
644750
if (!fields.schema || isSimpleStringSchema || isStringJSONSchema) {
@@ -658,9 +764,50 @@ export function tool<
658764
pickRunnableConfigKeys(childConfig),
659765
async () => {
660766
try {
661-
// TS doesn't restrict the type here based on the guard above
662-
// eslint-disable-next-line @typescript-eslint/no-explicit-any
663-
resolve(func(input as any, childConfig));
767+
if (shouldPassRuntime) {
768+
// Construct runtime object from config
769+
// State will be provided by ToolNode, but we create a minimal runtime here
770+
// eslint-disable-next-line @typescript-eslint/no-explicit-any
771+
const lgConfig = config as any;
772+
const toolConfig = childConfig as ToolRunnableConfig;
773+
const runtime: ToolRuntime<StateSchema, ContextSchema> = {
774+
state: (lgConfig?.state ||
775+
{}) as StateSchema extends InteropZodObject
776+
? InferInteropZodOutput<StateSchema>
777+
: Record<string, unknown>,
778+
toolCallId: toolConfig?.toolCall?.id || "",
779+
config: toolConfig,
780+
context: (lgConfig?.context ||
781+
undefined) as ContextSchema extends InteropZodObject
782+
? InferInteropZodOutput<ContextSchema>
783+
: unknown,
784+
store: lgConfig?.store || null,
785+
writer: lgConfig?.writer || null,
786+
};
787+
const funcWithRuntime = func as (
788+
input: unknown,
789+
runtime: ToolRuntime<StateSchema, ContextSchema>,
790+
options?: unknown
791+
) => ToolOutputT | Promise<ToolOutputT>;
792+
resolve(
793+
await funcWithRuntime(
794+
input as InferInteropZodOutput<SchemaT>,
795+
runtime,
796+
childConfig
797+
)
798+
);
799+
} else {
800+
const funcWithoutRuntime = func as (
801+
input: unknown,
802+
options?: unknown
803+
) => ToolOutputT | Promise<ToolOutputT>;
804+
resolve(
805+
await funcWithoutRuntime(
806+
input as InferInteropZodOutput<SchemaT>,
807+
childConfig
808+
)
809+
);
810+
}
664811
} catch (e) {
665812
reject(e);
666813
}
@@ -703,7 +850,40 @@ export function tool<
703850
pickRunnableConfigKeys(childConfig),
704851
async () => {
705852
try {
706-
const result = await func(input, childConfig);
853+
let result: ToolOutputT;
854+
if (shouldPassRuntime) {
855+
// Construct runtime object from config
856+
// State will be provided by ToolNode, but we create a minimal runtime here
857+
// eslint-disable-next-line @typescript-eslint/no-explicit-any
858+
const lgConfig = config as any;
859+
const toolConfig = childConfig as ToolRunnableConfig;
860+
const runtime: ToolRuntime<StateSchema, ContextSchema> = {
861+
state: (lgConfig?.state ||
862+
{}) as StateSchema extends InteropZodObject
863+
? InferInteropZodOutput<StateSchema>
864+
: Record<string, unknown>,
865+
toolCallId: toolConfig?.toolCall?.id || "",
866+
config: toolConfig,
867+
context: (lgConfig?.context ||
868+
undefined) as ContextSchema extends InteropZodObject
869+
? InferInteropZodOutput<ContextSchema>
870+
: unknown,
871+
store: lgConfig?.store || null,
872+
writer: lgConfig?.writer || null,
873+
};
874+
const funcWithRuntime = func as (
875+
input: SchemaOutputT,
876+
runtime: ToolRuntime<StateSchema, ContextSchema>,
877+
options?: unknown
878+
) => ToolOutputT | Promise<ToolOutputT>;
879+
result = await funcWithRuntime(input, runtime, childConfig);
880+
} else {
881+
const funcWithoutRuntime = func as (
882+
input: SchemaOutputT,
883+
options?: unknown
884+
) => ToolOutputT | Promise<ToolOutputT>;
885+
result = await funcWithoutRuntime(input, childConfig);
886+
}
707887

708888
/**
709889
* If the signal is aborted, we don't want to resolve the promise

0 commit comments

Comments
 (0)