Skip to content

Commit

Permalink
Make default Ollama withStructuredOutput use built-in structured outp…
Browse files Browse the repository at this point in the history
…ut format
  • Loading branch information
jacoblee93 committed Feb 8, 2025
1 parent 6b7d2a5 commit af86324
Show file tree
Hide file tree
Showing 4 changed files with 161 additions and 17 deletions.
10 changes: 5 additions & 5 deletions libs/langchain-ollama/package.json
Original file line number Diff line number Diff line change
Expand Up @@ -32,8 +32,10 @@
"author": "LangChain",
"license": "MIT",
"dependencies": {
"ollama": "^0.5.9",
"uuid": "^10.0.0"
"ollama": "^0.5.12",
"uuid": "^10.0.0",
"zod": "^3.24.1",
"zod-to-json-schema": "^3.24.1"
},
"peerDependencies": {
"@langchain/core": ">=0.2.21 <0.4.0"
Expand Down Expand Up @@ -62,9 +64,7 @@
"release-it": "^17.6.0",
"rollup": "^4.5.2",
"ts-jest": "^29.1.0",
"typescript": "<5.2.0",
"zod": "^3.22.4",
"zod-to-json-schema": "^3.23.0"
"typescript": "<5.2.0"
},
"publishConfig": {
"access": "public"
Expand Down
128 changes: 123 additions & 5 deletions libs/langchain-ollama/src/chat_models.ts
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,10 @@ import {
UsageMetadata,
type BaseMessage,
} from "@langchain/core/messages";
import { BaseLanguageModelInput } from "@langchain/core/language_models/base";
import {
BaseLanguageModelInput,
StructuredOutputMethodOptions,
} from "@langchain/core/language_models/base";
import { CallbackManagerForLLMRun } from "@langchain/core/callbacks/manager";
import {
type BaseChatModelParams,
Expand All @@ -21,21 +24,32 @@ import type {
Message as OllamaMessage,
Tool as OllamaTool,
} from "ollama";
import { Runnable } from "@langchain/core/runnables";
import {
Runnable,
RunnablePassthrough,
RunnableSequence,
} from "@langchain/core/runnables";
import { convertToOpenAITool } from "@langchain/core/utils/function_calling";
import { concat } from "@langchain/core/utils/stream";
import { z } from "zod";
import { zodToJsonSchema } from "zod-to-json-schema";
import {

Check failure on line 36 in libs/langchain-ollama/src/chat_models.ts

View workflow job for this annotation

GitHub Actions / Check linting

`./utils.js` import should occur after import of `@langchain/core/output_parsers`
convertOllamaMessagesToLangChain,
convertToOllamaMessages,
} from "./utils.js";
import { OllamaCamelCaseOptions } from "./types.js";

Check failure on line 40 in libs/langchain-ollama/src/chat_models.ts

View workflow job for this annotation

GitHub Actions / Check linting

`./types.js` import should occur after import of `@langchain/core/output_parsers`
import { isZodSchema } from "@langchain/core/utils/types";
import { JsonOutputParser } from "@langchain/core/output_parsers";

Check failure on line 42 in libs/langchain-ollama/src/chat_models.ts

View workflow job for this annotation

GitHub Actions / Check linting

'@langchain/core/output_parsers' imported multiple times
import { StructuredOutputParser } from "@langchain/core/output_parsers";

Check failure on line 43 in libs/langchain-ollama/src/chat_models.ts

View workflow job for this annotation

GitHub Actions / Check linting

'@langchain/core/output_parsers' imported multiple times

export interface ChatOllamaCallOptions extends BaseChatModelCallOptions {
/**
* An array of strings to stop on.
*/
stop?: string[];
tools?: BindToolsInput[];
// eslint-disable-next-line @typescript-eslint/no-explicit-any
format?: string | Record<string, any>;
}

export interface PullModelOptions {
Expand Down Expand Up @@ -82,7 +96,8 @@ export interface ChatOllamaInput
*/
checkOrPullModel?: boolean;
streaming?: boolean;
format?: string;
// eslint-disable-next-line @typescript-eslint/no-explicit-any
format?: string | Record<string, any>;
}

/**
Expand Down Expand Up @@ -453,7 +468,8 @@ export class ChatOllama

streaming?: boolean;

format?: string;
// eslint-disable-next-line @typescript-eslint/no-explicit-any
format?: string | Record<string, any>;

keepAlive?: string | number;

Expand Down Expand Up @@ -575,7 +591,7 @@ export class ChatOllama

return {
model: this.model,
format: this.format,
format: options?.format ?? this.format,
keep_alive: this.keepAlive,
options: {
numa: this.numa,
Expand Down Expand Up @@ -763,4 +779,106 @@ export class ChatOllama
}),
});
}

withStructuredOutput<
// eslint-disable-next-line @typescript-eslint/no-explicit-any
RunOutput extends Record<string, any> = Record<string, any>
>(
outputSchema:
| z.ZodType<RunOutput>
// eslint-disable-next-line @typescript-eslint/no-explicit-any
| Record<string, any>,
config?: StructuredOutputMethodOptions<false>
): Runnable<BaseLanguageModelInput, RunOutput>;

withStructuredOutput<
// eslint-disable-next-line @typescript-eslint/no-explicit-any
RunOutput extends Record<string, any> = Record<string, any>
>(
outputSchema:
| z.ZodType<RunOutput>
// eslint-disable-next-line @typescript-eslint/no-explicit-any
| Record<string, any>,
config?: StructuredOutputMethodOptions<true>
): Runnable<BaseLanguageModelInput, { raw: BaseMessage; parsed: RunOutput }>;

withStructuredOutput<
// eslint-disable-next-line @typescript-eslint/no-explicit-any
RunOutput extends Record<string, any> = Record<string, any>
>(
outputSchema:
| z.ZodType<RunOutput>
// eslint-disable-next-line @typescript-eslint/no-explicit-any
| Record<string, any>,
config?: StructuredOutputMethodOptions<boolean>
):
| Runnable<BaseLanguageModelInput, RunOutput>
| Runnable<
BaseLanguageModelInput,
{
raw: BaseMessage;
parsed: RunOutput;
}
>;

withStructuredOutput<
// eslint-disable-next-line @typescript-eslint/no-explicit-any
RunOutput extends Record<string, any> = Record<string, any>
>(
outputSchema:
| z.ZodType<RunOutput>
// eslint-disable-next-line @typescript-eslint/no-explicit-any
| Record<string, any>,
config?: StructuredOutputMethodOptions<boolean>
):
| Runnable<BaseLanguageModelInput, RunOutput>
| Runnable<
BaseLanguageModelInput,
{
raw: BaseMessage;
parsed: RunOutput;
}
> {
if (config?.method === undefined || config?.method === "jsonSchema") {
const outputSchemaIsZod = isZodSchema(outputSchema);
const jsonSchema = outputSchemaIsZod
? zodToJsonSchema(outputSchema)
: outputSchema;
const llm = this.bind({
format: jsonSchema,
});
const outputParser = outputSchemaIsZod
? StructuredOutputParser.fromZodSchema(outputSchema)
: new JsonOutputParser<RunOutput>();

if (!config?.includeRaw) {
return llm.pipe(outputParser) as Runnable<
BaseLanguageModelInput,
RunOutput
>;
}

const parserAssign = RunnablePassthrough.assign({
// eslint-disable-next-line @typescript-eslint/no-explicit-any
parsed: (input: any, config) => outputParser.invoke(input.raw, config),
});
const parserNone = RunnablePassthrough.assign({
parsed: () => null,
});
const parsedWithFallback = parserAssign.withFallbacks({
fallbacks: [parserNone],
});
return RunnableSequence.from<
BaseLanguageModelInput,
{ raw: BaseMessage; parsed: RunOutput }
>([
{
raw: llm,
},
parsedWithFallback,
]);
} else {
return super.withStructuredOutput<RunOutput>(outputSchema, config);
}
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -98,6 +98,23 @@ test("Ollama can call withStructuredOutput includeRaw", async () => {
includeRaw: true,
});

const result = await model.invoke(messageHistory);
expect(result).toBeDefined();
expect(result.parsed.location).toBeDefined();
expect(result.parsed.location).not.toBe("");
expect((result.raw as AIMessage).tool_calls?.length).toBe(0);
});

test("Ollama can call withStructuredOutput includeRaw with tool calling", async () => {
const model = new ChatOllama({
model: "llama3-groq-tool-use",
maxRetries: 1,
}).withStructuredOutput(weatherTool.schema, {
name: weatherTool.name,
includeRaw: true,
method: "functionCalling",
});

const result = await model.invoke(messageHistory);
expect(result).toBeDefined();
expect(result.parsed.location).toBeDefined();
Expand Down
23 changes: 16 additions & 7 deletions yarn.lock
Original file line number Diff line number Diff line change
Expand Up @@ -12990,15 +12990,15 @@ __metadata:
eslint-plugin-prettier: ^4.2.1
jest: ^29.5.0
jest-environment-node: ^29.6.4
ollama: ^0.5.9
ollama: ^0.5.12
prettier: ^2.8.3
release-it: ^17.6.0
rollup: ^4.5.2
ts-jest: ^29.1.0
typescript: <5.2.0
uuid: ^10.0.0
zod: ^3.22.4
zod-to-json-schema: ^3.23.0
zod: ^3.24.1
zod-to-json-schema: ^3.24.1
peerDependencies:
"@langchain/core": ">=0.2.21 <0.4.0"
languageName: unknown
Expand Down Expand Up @@ -36216,12 +36216,12 @@ __metadata:
languageName: node
linkType: hard

"ollama@npm:^0.5.9":
version: 0.5.9
resolution: "ollama@npm:0.5.9"
"ollama@npm:^0.5.12":
version: 0.5.12
resolution: "ollama@npm:0.5.12"
dependencies:
whatwg-fetch: ^3.6.20
checksum: bfaadcec6273d86fcc7c94e5e9e571a7b6b84b852b407a473f3bac7dc69b7b11815a163ae549b5318267a00f192d39696225309812319d2edc8a98a079ace475
checksum: 0abc1151d2cfd02198829f706f8efca978c8562691e7502924166798f6a0cd7e1bf51e085d313ddf5a76507a36ffa12b48a66d4dd659b419474c2f33e3f03b44
languageName: node
linkType: hard

Expand Down Expand Up @@ -44942,6 +44942,15 @@ __metadata:
languageName: node
linkType: hard

"zod-to-json-schema@npm:^3.24.1":
version: 3.24.1
resolution: "zod-to-json-schema@npm:3.24.1"
peerDependencies:
zod: ^3.24.1
checksum: 7195563f611bc21ea7f44129b8e32780125a9bd98b2e6b8709ac98bd2645729fecd87b8aeeaa8789617ee3f38e6585bab23dd613e2a35c31c6c157908f7a1681
languageName: node
linkType: hard

"zod@npm:3.23.8":
version: 3.23.8
resolution: "zod@npm:3.23.8"
Expand Down

0 comments on commit af86324

Please sign in to comment.