Skip to content

Commit e115165

Browse files
nfcamposnfcampos
committed
Fix several issues
- combineCallOptions didn't respect the merge behavior defined for runnable config fields) - withStructuedOutput output type was too broad (eg included AIMessage) - withOutputParser type for message input was too broad, it is always an AIMessage/Chunk - bindTools was accidentally made required through addition of abstract keyword - combineCallOptions was called for stream but not for invoke - fake chat model classes used for testing updated
1 parent 16ee7d2 commit e115165

File tree

5 files changed

+97
-169
lines changed

5 files changed

+97
-169
lines changed

‎libs/langchain-core/src/language_models/base.ts‎

Lines changed: 12 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,4 @@
11
import type { Tiktoken, TiktokenModel } from "js-tiktoken/lite";
2-
import type { ZodType as ZodTypeV3 } from "zod/v3";
3-
import type { $ZodType as ZodTypeV4 } from "zod/v4/core";
42

53
import { type BaseCache, InMemoryCache } from "../caches/base.js";
64
import {
@@ -26,6 +24,7 @@ import {
2624
InteropZodObject,
2725
InteropZodType,
2826
} from "../utils/types/zod.js";
27+
import { type AIMessage, type AIMessageChunk } from "../messages/ai.js";
2928

3029
// https://www.npmjs.com/package/js-tiktoken
3130

@@ -280,6 +279,8 @@ export type BaseLanguageModelInput =
280279
| string
281280
| BaseMessageLike[];
282281

282+
export type AnyAIMessage = AIMessage | AIMessageChunk;
283+
283284
export type StructuredOutputType = InferInteropZodOutput<InteropZodObject>;
284285

285286
export type StructuredOutputMethodOptions<IncludeRaw extends boolean = false> =
@@ -560,15 +561,15 @@ export abstract class BaseLanguageModel<
560561
throw new Error("Use .toJSON() instead");
561562
}
562563

563-
withStructuredOutput?<Output>(
564+
withStructuredOutput?<Output extends Record<string, unknown>>(
564565
schema: InteropZodType<Output> | JSONSchema,
565566
config?: StructuredOutputMethodOptions<false>
566-
): BaseLanguageModel<BaseMessage, CallOptions>;
567+
): BaseLanguageModel<Output, CallOptions>;
567568

568-
withStructuredOutput?<Output>(
569+
withStructuredOutput?<Output extends Record<string, unknown>>(
569570
schema: InteropZodType<Output> | JSONSchema,
570-
config?: StructuredOutputMethodOptions<true>
571-
): BaseLanguageModel<{ raw: BaseMessage; parsed: Output }, CallOptions>;
571+
config: StructuredOutputMethodOptions<true>
572+
): BaseLanguageModel<{ raw: AnyAIMessage; parsed: Output }, CallOptions>;
572573

573574
/**
574575
* Model wrapper that returns outputs formatted to match the given schema.
@@ -581,14 +582,14 @@ export abstract class BaseLanguageModel<
581582
* @param {string} name The name of the function to call.
582583
* @param {"functionCalling" | "jsonMode"} [method=functionCalling] The method to use for getting the structured output. Defaults to "functionCalling".
583584
* @param {boolean | undefined} [includeRaw=false] Whether to include the raw output in the result. Defaults to false.
584-
* @returns {Runnable<RunInput, RunOutput> | Runnable<RunInput, { raw: BaseMessage; parsed: RunOutput }>} A new runnable that calls the LLM with structured output.
585+
* @returns {Runnable<RunInput, RunOutput> | Runnable<RunInput, { raw: AnyAIMessage; parsed: RunOutput }>} A new runnable that calls the LLM with structured output.
585586
*/
586-
withStructuredOutput?<Output>(
587+
withStructuredOutput?<Output extends Record<string, unknown>>(
587588
schema: InteropZodType<Output> | JSONSchema,
588589
config?: StructuredOutputMethodOptions<boolean>
589590
):
590-
| BaseLanguageModel<BaseMessage, CallOptions>
591-
| BaseLanguageModel<{ raw: BaseMessage; parsed: Output }, CallOptions>;
591+
| BaseLanguageModel<Output, CallOptions>
592+
| BaseLanguageModel<{ raw: AnyAIMessage; parsed: Output }, CallOptions>;
592593
}
593594

594595
/**

‎libs/langchain-core/src/language_models/chat_models.ts‎

Lines changed: 29 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,6 @@ import {
55
type BaseMessageLike,
66
HumanMessage,
77
coerceMessageLikeToMessage,
8-
AIMessageChunk,
98
isAIMessageChunk,
109
isBaseMessage,
1110
isAIMessage,
@@ -29,13 +28,14 @@ import {
2928
type BaseLanguageModelCallOptions,
3029
type BaseLanguageModelInput,
3130
type BaseLanguageModelParams,
31+
type AnyAIMessage,
3232
} from "./base.js";
3333
import {
3434
CallbackManager,
3535
type CallbackManagerForLLMRun,
3636
type Callbacks,
3737
} from "../callbacks/manager.js";
38-
import type { RunnableConfig } from "../runnables/config.js";
38+
import { mergeConfigs, type RunnableConfig } from "../runnables/config.js";
3939
import type { BaseCache } from "../caches/base.js";
4040
import {
4141
StructuredToolInterface,
@@ -180,18 +180,18 @@ export type BindToolsInput =
180180
| RunnableToolLike
181181
| StructuredToolParams;
182182

183-
export type ChatModelOutputParser<T = unknown> = Runnable<BaseMessage, T>;
184-
export type InferChatModelOutputParser<TOutput> = TOutput extends BaseMessage
185-
? ChatModelOutputParser<TOutput>
186-
: undefined;
183+
export type ChatModelOutputParser<T = unknown> = Runnable<AnyAIMessage, T>;
184+
export type InferChatModelOutputParser<TOutput> = TOutput extends AnyAIMessage
185+
? undefined
186+
: ChatModelOutputParser<TOutput>;
187187

188188
/**
189189
* Base class for chat models. It extends the BaseLanguageModel class and
190190
* provides methods for generating chat based on input messages.
191191
*/
192192
export abstract class BaseChatModel<
193193
CallOptions extends BaseChatModelCallOptions = BaseChatModelCallOptions,
194-
TOutput = BaseMessage
194+
TOutput = AnyAIMessage
195195
> extends BaseLanguageModel<TOutput, CallOptions> {
196196
// Backwards compatibility since fields have been moved to RunnableConfig
197197
declare ParsedCallOptions: Omit<
@@ -214,11 +214,8 @@ export abstract class BaseChatModel<
214214

215215
protected _combineCallOptions(
216216
additionalOptions?: Partial<CallOptions>
217-
): CallOptions {
218-
return {
219-
...this.defaultOptions,
220-
...additionalOptions,
221-
};
217+
): Partial<CallOptions> {
218+
return mergeConfigs(this.defaultOptions, additionalOptions);
222219
}
223220

224221
protected async _parseOutput(output: BaseMessage): Promise<TOutput> {
@@ -237,7 +234,9 @@ export abstract class BaseChatModel<
237234
): [RunnableConfig, this["ParsedCallOptions"]] {
238235
// For backwards compat, keep `signal` in both runnableConfig and callOptions
239236
const [runnableConfig, callOptions] =
240-
super._separateRunnableConfigFromCallOptions(options);
237+
super._separateRunnableConfigFromCallOptions(
238+
this._combineCallOptions(options)
239+
);
241240
(callOptions as this["ParsedCallOptions"]).signal = runnableConfig.signal;
242241
return [runnableConfig, callOptions as this["ParsedCallOptions"]];
243242
}
@@ -250,10 +249,7 @@ export abstract class BaseChatModel<
250249
* matching the provider's specific tool schema.
251250
* @param kwargs Any additional parameters to bind.
252251
*/
253-
abstract bindTools?(
254-
tools: BindToolsInput[],
255-
options?: Partial<CallOptions>
256-
): BaseChatModel<CallOptions, TOutput>;
252+
bindTools?(tools: BindToolsInput[], options?: Partial<CallOptions>): this;
257253

258254
/**
259255
* Invokes the chat model with a single input.
@@ -298,9 +294,8 @@ export abstract class BaseChatModel<
298294
} else {
299295
const prompt = BaseChatModel._convertInputToPromptValue(input);
300296
const messages = prompt.toChatMessages();
301-
const combinedOptions = this._combineCallOptions(options);
302297
const [runnableConfig, callOptions] =
303-
this._separateRunnableConfigFromCallOptionsCompat(combinedOptions);
298+
this._separateRunnableConfigFromCallOptionsCompat(options);
304299

305300
const inheritableMetadata = {
306301
...runnableConfig.metadata,
@@ -936,7 +931,7 @@ export abstract class BaseChatModel<
936931
}
937932

938933
/** @internal */
939-
protected withOutputParser<TOutput>(
934+
protected withOutputParser<TOutput extends Record<string, unknown>>(
940935
outputParser: ChatModelOutputParser<TOutput>
941936
): BaseChatModel<CallOptions, TOutput> {
942937
const Cls = this.constructor as Constructor<
@@ -948,29 +943,29 @@ export abstract class BaseChatModel<
948943
return instance;
949944
}
950945

951-
withStructuredOutput<Output>(
946+
withStructuredOutput<Output extends Record<string, unknown>>(
952947
schema: InteropZodType<Output> | JSONSchema,
953948
config?: StructuredOutputMethodOptions<false>
954949
): BaseChatModel<CallOptions, Output>;
955950

956-
withStructuredOutput<Output>(
951+
withStructuredOutput<Output extends Record<string, unknown>>(
957952
schema: InteropZodType<Output> | JSONSchema,
958-
config?: StructuredOutputMethodOptions<true>
959-
): BaseChatModel<CallOptions, { raw: BaseMessage; parsed: Output }>;
953+
config: StructuredOutputMethodOptions<true>
954+
): BaseChatModel<CallOptions, { raw: AnyAIMessage; parsed: Output }>;
960955

961-
withStructuredOutput<Output>(
956+
withStructuredOutput<Output extends Record<string, unknown>>(
962957
schema: InteropZodType<Output> | JSONSchema,
963958
config?: StructuredOutputMethodOptions<boolean>
964959
):
965960
| BaseChatModel<CallOptions, Output>
966-
| BaseChatModel<CallOptions, { raw: BaseMessage; parsed: Output }>;
961+
| BaseChatModel<CallOptions, { raw: AnyAIMessage; parsed: Output }>;
967962

968-
withStructuredOutput<Output>(
963+
withStructuredOutput<Output extends Record<string, unknown>>(
969964
schema: InteropZodType<Output> | JSONSchema,
970965
config?: StructuredOutputMethodOptions<boolean>
971966
):
972967
| BaseChatModel<CallOptions, Output>
973-
| BaseChatModel<CallOptions, { raw: BaseMessage; parsed: Output }> {
968+
| BaseChatModel<CallOptions, { raw: AnyAIMessage; parsed: Output }> {
974969
if (typeof this.bindTools !== "function") {
975970
throw new Error(
976971
`Chat model must implement ".bindTools()" to use withStructuredOutput.`
@@ -1022,8 +1017,8 @@ export abstract class BaseChatModel<
10221017
}
10231018

10241019
const llm = this.bindTools(tools);
1025-
const toolMessageParser = RunnableLambda.from<AIMessageChunk, Output>(
1026-
(input: AIMessageChunk): Output => {
1020+
const toolMessageParser = RunnableLambda.from<AnyAIMessage, Output>(
1021+
(input: AnyAIMessage): Output => {
10271022
if (!input.tool_calls || input.tool_calls.length === 0) {
10281023
throw new Error("No tool calls found in the response.");
10291024
}
@@ -1046,12 +1041,12 @@ export abstract class BaseChatModel<
10461041
}
10471042

10481043
const rawOutputParser = RunnableLambda.from<
1049-
AIMessageChunk,
1050-
{ raw: BaseMessage; parsed: Output }
1044+
AnyAIMessage,
1045+
{ raw: AnyAIMessage; parsed: Output }
10511046
>(
10521047
async (
1053-
input: AIMessageChunk
1054-
): Promise<{ raw: BaseMessage; parsed: Output }> => {
1048+
input: AnyAIMessage
1049+
): Promise<{ raw: AnyAIMessage; parsed: Output }> => {
10551050
return {
10561051
raw: input,
10571052
parsed: await toolMessageParser.invoke(input),

‎libs/langchain-core/src/language_models/tests/v1.test.ts‎

Lines changed: 6 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -60,27 +60,23 @@ class SimpleChatModel<
6060
}
6161
}
6262

63-
withConfig(config: SimpleChatModelCallOptions) {
64-
return super.withConfig(config);
65-
}
66-
6763
bindTools(tools: BindToolsInput[]) {
6864
return this.withConfig({
69-
tools: [...(this.defaultOptions.tools ?? []), ...tools],
65+
tools: [...(this.defaultOptions?.tools ?? []), ...tools],
7066
});
7167
}
7268

73-
withStructuredOutput<Output>(
69+
withStructuredOutput<Output extends Record<string, unknown>>(
7470
schema: InteropZodType<Output> | JSONSchema,
7571
config?: StructuredOutputMethodOptions<false>
7672
): SimpleChatModel<CallOptions, Output>;
7773

78-
withStructuredOutput<Output>(
74+
withStructuredOutput<Output extends Record<string, unknown>>(
7975
schema: InteropZodType<Output> | JSONSchema,
80-
config?: StructuredOutputMethodOptions<true>
76+
config: StructuredOutputMethodOptions<true>
8177
): SimpleChatModel<CallOptions, { raw: BaseMessage; parsed: Output }>;
8278

83-
withStructuredOutput<Output>(
79+
withStructuredOutput<Output extends Record<string, unknown>>(
8480
schema: InteropZodType<Output> | JSONSchema,
8581
config?: StructuredOutputMethodOptions<boolean>
8682
):
@@ -114,7 +110,7 @@ describe("SimpleChatModel", () => {
114110
bar: z.string(),
115111
})
116112
);
117-
expect(model.defaultOptions.tools?.length).toBe(2);
113+
expect(model.defaultOptions.tools?.length).toBe(3); // + 1 structured output
118114
expect(model.outputParser).toBeDefined();
119115
});
120116

‎libs/langchain-core/src/prompts/tests/structured.test.ts‎

Lines changed: 9 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -2,13 +2,11 @@
22
import { ZodType, ZodTypeDef } from "zod";
33
import { test, expect } from "vitest";
44
import {
5-
StructuredOutputMethodParams,
65
StructuredOutputMethodOptions,
76
BaseLanguageModelInput,
7+
AnyAIMessage,
88
} from "../../language_models/base.js";
9-
import { BaseMessage } from "../../messages/index.js";
10-
import { Runnable, RunnableLambda } from "../../runnables/base.js";
11-
import { RunnableConfig } from "../../runnables/config.js";
9+
import { RunnableLambda } from "../../runnables/base.js";
1210
import { FakeListChatModel } from "../../utils/testing/index.js";
1311
import { StructuredPrompt } from "../structured.js";
1412
import { load } from "../../load/index.js";
@@ -17,42 +15,25 @@ class FakeStructuredChatModel extends FakeListChatModel {
1715
withStructuredOutput<
1816
RunOutput extends Record<string, any> = Record<string, any>
1917
>(
20-
_params:
21-
| Record<string, any>
22-
| StructuredOutputMethodParams<RunOutput, false>
23-
| ZodType<RunOutput, ZodTypeDef, RunOutput>,
18+
_params: Record<string, any> | ZodType<RunOutput, ZodTypeDef, RunOutput>,
2419
config?: StructuredOutputMethodOptions<false> | undefined
25-
): Runnable<BaseLanguageModelInput, RunOutput, RunnableConfig>;
20+
): FakeListChatModel<RunOutput>;
2621

2722
withStructuredOutput<
2823
RunOutput extends Record<string, any> = Record<string, any>
2924
>(
30-
_params:
31-
| Record<string, any>
32-
| StructuredOutputMethodParams<RunOutput, true>
33-
| ZodType<RunOutput, ZodTypeDef, RunOutput>,
25+
_params: Record<string, any> | ZodType<RunOutput, ZodTypeDef, RunOutput>,
3426
config?: StructuredOutputMethodOptions<true> | undefined
35-
): Runnable<
36-
BaseLanguageModelInput,
37-
{ raw: BaseMessage; parsed: RunOutput },
38-
RunnableConfig
39-
>;
27+
): FakeListChatModel<{ raw: AnyAIMessage; parsed: RunOutput }>;
4028

4129
withStructuredOutput<
4230
RunOutput extends Record<string, any> = Record<string, any>
4331
>(
44-
_params:
45-
| Record<string, any>
46-
| StructuredOutputMethodParams<RunOutput, boolean>
47-
| ZodType<RunOutput, ZodTypeDef, RunOutput>,
32+
_params: Record<string, any> | ZodType<RunOutput, ZodTypeDef, RunOutput>,
4833
_config?: StructuredOutputMethodOptions<boolean> | undefined
4934
):
50-
| Runnable<BaseLanguageModelInput, RunOutput, RunnableConfig>
51-
| Runnable<
52-
BaseLanguageModelInput,
53-
{ raw: BaseMessage; parsed: RunOutput },
54-
RunnableConfig
55-
> {
35+
| FakeListChatModel<RunOutput>
36+
| FakeListChatModel<{ raw: AnyAIMessage; parsed: RunOutput }> {
5637
if (!_config?.includeRaw) {
5738
if (typeof _params === "object") {
5839
const func = RunnableLambda.from(

0 commit comments

Comments
 (0)