feat(ollama): Add support for Ollama built-in JSON schema with withStructuredOutput (#7672)

Author: jacoblee93

libs/langchain-ollama/package.json

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

libs/langchain-ollama/src/chat_models.ts

@@ -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,
@@ -21,9 +24,20 @@ 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 {
+  JsonOutputParser,
+  StructuredOutputParser,
+} from "@langchain/core/output_parsers";
+import { isZodSchema } from "@langchain/core/utils/types";
+import { z } from "zod";
+import { zodToJsonSchema } from "zod-to-json-schema";
 import {
   convertOllamaMessagesToLangChain,
   convertToOllamaMessages,
@@ -36,6 +50,8 @@ export interface ChatOllamaCallOptions extends BaseChatModelCallOptions {
    */
   stop?: string[];
   tools?: BindToolsInput[];
+  // eslint-disable-next-line @typescript-eslint/no-explicit-any
+  format?: string | Record<string, any>;
 }
 
 export interface PullModelOptions {
@@ -82,7 +98,8 @@ export interface ChatOllamaInput
    */
   checkOrPullModel?: boolean;
   streaming?: boolean;
-  format?: string;
+  // eslint-disable-next-line @typescript-eslint/no-explicit-any
+  format?: string | Record<string, any>;
 }
 
 /**
@@ -453,7 +470,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;
 
@@ -575,7 +593,7 @@ export class ChatOllama
 
     return {
       model: this.model,
-      format: this.format,
+      format: options?.format ?? this.format,
       keep_alive: this.keepAlive,
       options: {
         numa: this.numa,
@@ -763,4 +781,109 @@ 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;
+        }
+      > {
+    // TODO: Make this method the default in a minor bump
+    if (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 {
+      // TODO: Fix this type in core
+      // eslint-disable-next-line @typescript-eslint/no-explicit-any
+      return super.withStructuredOutput<RunOutput>(outputSchema, config as any);
+    }
+  }
 }

libs/langchain-ollama/src/tests/chat_models-tools.int.test.ts renamed to libs/langchain-ollama/src/tests/chat_models_structured_output.int.test.ts

@@ -89,7 +89,24 @@ test("Ollama can call withStructuredOutput", async () => {
   expect(result.location).not.toBe("");
 });
 
-test("Ollama can call withStructuredOutput includeRaw", async () => {
+test("Ollama can call withStructuredOutput includeRaw JSON Schema", async () => {
+  const model = new ChatOllama({
+    model: "llama3-groq-tool-use",
+    maxRetries: 1,
+  }).withStructuredOutput(weatherTool.schema, {
+    name: weatherTool.name,
+    includeRaw: true,
+    method: "jsonSchema",
+  });
+
+  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,

yarn.lock

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