modelfusion 0.104.0 → 0.106.0

package/model-provider/ollama/OllamaCompletionModel.cjs

@@ -12,8 +12,7 @@ const PromptTemplateTextStreamingModel_js_1 = require("../../model-function/gene
 const TextGenerationModel_js_1 = require("../../model-function/generate-text/TextGenerationModel.cjs");
 const TextGenerationToolCallModel_js_1 = require("../../tool/generate-tool-call/TextGenerationToolCallModel.cjs");
 const TextGenerationToolCallsOrGenerateTextModel_js_1 = require("../../tool/generate-tool-calls-or-text/TextGenerationToolCallsOrGenerateTextModel.cjs");
-const AsyncQueue_js_1 = require("../../util/AsyncQueue.cjs");
-const parseJsonStream_js_1 = require("../../util/streaming/parseJsonStream.cjs");
+const createJsonStreamResponseHandler_js_1 = require("../../util/streaming/createJsonStreamResponseHandler.cjs");
 const OllamaApiConfiguration_js_1 = require("./OllamaApiConfiguration.cjs");
 const OllamaError_js_1 = require("./OllamaError.cjs");
 class OllamaCompletionModel extends AbstractModel_js_1.AbstractModel {
@@ -135,6 +134,10 @@ class OllamaCompletionModel extends AbstractModel_js_1.AbstractModel {
             responseFormat: exports.OllamaCompletionResponseFormat.deltaIterable,
         });
     }
+    extractTextDelta(delta) {
+        const chunk = delta;
+        return chunk.done === true ? undefined : chunk.response;
+    }
     asToolCallGenerationModel(promptTemplate) {
         return new TextGenerationToolCallModel_js_1.TextGenerationToolCallModel({
             model: this,
@@ -150,11 +153,25 @@ class OllamaCompletionModel extends AbstractModel_js_1.AbstractModel {
     withTextPrompt() {
         return this.withPromptTemplate({
             format(prompt) {
-                return { prompt: prompt };
+                return { prompt };
             },
             stopSequences: [],
         });
     }
+    /**
+     * Maps the prompt for a text version of the Ollama completion prompt template (without image support).
+     */
+    withTextPromptTemplate(promptTemplate) {
+        return new PromptTemplateTextStreamingModel_js_1.PromptTemplateTextStreamingModel({
+            model: this.withTextPrompt().withSettings({
+                stopSequences: [
+                    ...(this.settings.stopSequences ?? []),
+                    ...promptTemplate.stopSequences,
+                ],
+            }),
+            promptTemplate,
+        });
+    }
     withPromptTemplate(promptTemplate) {
         return new PromptTemplateTextStreamingModel_js_1.PromptTemplateTextStreamingModel({
             model: this.withSettings({
@@ -184,7 +201,7 @@ const ollamaCompletionResponseSchema = zod_1.z.object({
     eval_duration: zod_1.z.number(),
     context: zod_1.z.array(zod_1.z.number()).optional(),
 });
-const ollamaCompletionStreamSchema = new ZodSchema_js_1.ZodSchema(zod_1.z.discriminatedUnion("done", [
+const ollamaCompletionStreamChunkSchema = (0, ZodSchema_js_1.zodSchema)(zod_1.z.discriminatedUnion("done", [
     zod_1.z.object({
         done: zod_1.z.literal(false),
         model: zod_1.z.string(),
@@ -206,44 +223,6 @@ const ollamaCompletionStreamSchema = new ZodSchema_js_1.ZodSchema(zod_1.z.discri
         context: zod_1.z.array(zod_1.z.number()).optional(),
     }),
 ]));
-async function createOllamaFullDeltaIterableQueue(stream) {
-    const queue = new AsyncQueue_js_1.AsyncQueue();
-    let accumulatedText = "";
-    // process the stream asynchonously (no 'await' on purpose):
-    (0, parseJsonStream_js_1.parseJsonStream)({
-        stream,
-        schema: ollamaCompletionStreamSchema,
-        process(event) {
-            if (event.done === true) {
-                queue.push({
-                    type: "delta",
-                    fullDelta: {
-                        content: accumulatedText,
-                        isComplete: true,
-                        delta: "",
-                    },
-                    valueDelta: "",
-                });
-            }
-            else {
-                accumulatedText += event.response;
-                queue.push({
-                    type: "delta",
-                    fullDelta: {
-                        content: accumulatedText,
-                        isComplete: false,
-                        delta: event.response,
-                    },
-                    valueDelta: event.response,
-                });
-            }
-        },
-        onDone() {
-            queue.close();
-        },
-    });
-    return queue;
-}
 exports.OllamaCompletionResponseFormat = {
     /**
      * Returns the response as a JSON object.
@@ -293,6 +272,6 @@ exports.OllamaCompletionResponseFormat = {
      */
     deltaIterable: {
         stream: true,
-        handler: async ({ response }) => createOllamaFullDeltaIterableQueue(response.body),
+        handler: (0, createJsonStreamResponseHandler_js_1.createJsonStreamResponseHandler)(ollamaCompletionStreamChunkSchema),
     },
 };

package/model-provider/ollama/OllamaCompletionModel.d.ts

@@ -2,13 +2,14 @@ import { z } from "zod";
 import { FunctionOptions } from "../../core/FunctionOptions.js";
 import { ApiConfiguration } from "../../core/api/ApiConfiguration.js";
 import { ResponseHandler } from "../../core/api/postToApi.js";
+import { ZodSchema } from "../../core/schema/ZodSchema.js";
 import { AbstractModel } from "../../model-function/AbstractModel.js";
-import { Delta } from "../../model-function/Delta.js";
 import { PromptTemplateTextStreamingModel } from "../../model-function/generate-text/PromptTemplateTextStreamingModel.js";
 import { TextStreamingModel } from "../../model-function/generate-text/TextGenerationModel.js";
 import { TextGenerationPromptTemplate } from "../../model-function/generate-text/TextGenerationPromptTemplate.js";
 import { TextGenerationToolCallModel, ToolCallPromptTemplate } from "../../tool/generate-tool-call/TextGenerationToolCallModel.js";
-import { TextGenerationToolCallsOrGenerateTextModel, ToolCallsOrGenerateTextPromptTemplate } from "../../tool/generate-tool-calls-or-text/TextGenerationToolCallsOrGenerateTextModel.js";
+import { TextGenerationToolCallsOrGenerateTextModel } from "../../tool/generate-tool-calls-or-text/TextGenerationToolCallsOrGenerateTextModel.js";
+import { ToolCallsOrGenerateTextPromptTemplate } from "../../tool/generate-tool-calls-or-text/ToolCallsOrGenerateTextPromptTemplate.js";
 import { OllamaTextGenerationSettings } from "./OllamaTextGenerationSettings.js";
 export interface OllamaCompletionPrompt {
     /**
@@ -70,10 +71,33 @@ export declare class OllamaCompletionModel<CONTEXT_WINDOW_SIZE extends number |
             finishReason: "unknown";
         }[];
     }>;
-    doStreamText(prompt: OllamaCompletionPrompt, options?: FunctionOptions): Promise<AsyncIterable<Delta<string>>>;
+    doStreamText(prompt: OllamaCompletionPrompt, options?: FunctionOptions): Promise<AsyncIterable<import("../../index.js").Delta<{
+        response: string;
+        model: string;
+        done: false;
+        created_at: string;
+    } | {
+        model: string;
+        done: true;
+        created_at: string;
+        total_duration: number;
+        prompt_eval_count: number;
+        eval_count: number;
+        eval_duration: number;
+        load_duration?: number | undefined;
+        sample_count?: number | undefined;
+        sample_duration?: number | undefined;
+        prompt_eval_duration?: number | undefined;
+        context?: number[] | undefined;
+    }>>>;
+    extractTextDelta(delta: unknown): string | undefined;
     asToolCallGenerationModel<INPUT_PROMPT>(promptTemplate: ToolCallPromptTemplate<INPUT_PROMPT, OllamaCompletionPrompt>): TextGenerationToolCallModel<INPUT_PROMPT, OllamaCompletionPrompt, this>;
     asToolCallsOrTextGenerationModel<INPUT_PROMPT>(promptTemplate: ToolCallsOrGenerateTextPromptTemplate<INPUT_PROMPT, OllamaCompletionPrompt>): TextGenerationToolCallsOrGenerateTextModel<INPUT_PROMPT, OllamaCompletionPrompt, this>;
     withTextPrompt(): PromptTemplateTextStreamingModel<string, OllamaCompletionPrompt, OllamaCompletionModelSettings<CONTEXT_WINDOW_SIZE>, this>;
+    /**
+     * Maps the prompt for a text version of the Ollama completion prompt template (without image support).
+     */
+    withTextPromptTemplate<INPUT_PROMPT>(promptTemplate: TextGenerationPromptTemplate<INPUT_PROMPT, string>): PromptTemplateTextStreamingModel<INPUT_PROMPT, string, OllamaCompletionModelSettings<CONTEXT_WINDOW_SIZE>, PromptTemplateTextStreamingModel<string, OllamaCompletionPrompt, OllamaCompletionModelSettings<CONTEXT_WINDOW_SIZE>, this>>;
     withPromptTemplate<INPUT_PROMPT>(promptTemplate: TextGenerationPromptTemplate<INPUT_PROMPT, OllamaCompletionPrompt>): PromptTemplateTextStreamingModel<INPUT_PROMPT, OllamaCompletionPrompt, OllamaCompletionModelSettings<CONTEXT_WINDOW_SIZE>, this>;
     withSettings(additionalSettings: Partial<OllamaCompletionModelSettings<CONTEXT_WINDOW_SIZE>>): this;
 }
@@ -115,11 +139,26 @@ declare const ollamaCompletionResponseSchema: z.ZodObject<{
     context?: number[] | undefined;
 }>;
 export type OllamaCompletionResponse = z.infer<typeof ollamaCompletionResponseSchema>;
-export type OllamaCompletionDelta = {
-    content: string;
-    isComplete: boolean;
-    delta: string;
-};
+declare const ollamaCompletionStreamChunkSchema: ZodSchema<{
+    response: string;
+    model: string;
+    done: false;
+    created_at: string;
+} | {
+    model: string;
+    done: true;
+    created_at: string;
+    total_duration: number;
+    prompt_eval_count: number;
+    eval_count: number;
+    eval_duration: number;
+    load_duration?: number | undefined;
+    sample_count?: number | undefined;
+    sample_duration?: number | undefined;
+    prompt_eval_duration?: number | undefined;
+    context?: number[] | undefined;
+}>;
+export type OllamaCompletionStreamChunk = (typeof ollamaCompletionStreamChunkSchema)["_type"];
 export type OllamaCompletionResponseFormatType<T> = {
     stream: boolean;
     handler: ResponseHandler<T>;
@@ -153,10 +192,28 @@ export declare const OllamaCompletionResponseFormat: {
      * of the response stream.
      */
     deltaIterable: {
-        stream: true;
+        stream: boolean;
         handler: ({ response }: {
             response: Response;
-        }) => Promise<AsyncIterable<Delta<string>>>;
+        }) => Promise<AsyncIterable<import("../../index.js").Delta<{
+            response: string;
+            model: string;
+            done: false;
+            created_at: string;
+        } | {
+            model: string;
+            done: true;
+            created_at: string;
+            total_duration: number;
+            prompt_eval_count: number;
+            eval_count: number;
+            eval_duration: number;
+            load_duration?: number | undefined;
+            sample_count?: number | undefined;
+            sample_duration?: number | undefined;
+            prompt_eval_duration?: number | undefined;
+            context?: number[] | undefined;
+        }>>>;
     };
 };
 export {};

package/model-provider/ollama/OllamaCompletionModel.js

@@ -2,15 +2,14 @@ import { z } from "zod";
 import { ApiCallError } from "../../core/api/ApiCallError.js";
 import { callWithRetryAndThrottle } from "../../core/api/callWithRetryAndThrottle.js";
 import { postJsonToApi } from "../../core/api/postToApi.js";
-import { ZodSchema } from "../../core/schema/ZodSchema.js";
+import { ZodSchema, zodSchema } from "../../core/schema/ZodSchema.js";
 import { safeParseJSON } from "../../core/schema/parseJSON.js";
 import { AbstractModel } from "../../model-function/AbstractModel.js";
 import { PromptTemplateTextStreamingModel } from "../../model-function/generate-text/PromptTemplateTextStreamingModel.js";
 import { textGenerationModelProperties, } from "../../model-function/generate-text/TextGenerationModel.js";
 import { TextGenerationToolCallModel, } from "../../tool/generate-tool-call/TextGenerationToolCallModel.js";
-import { TextGenerationToolCallsOrGenerateTextModel, } from "../../tool/generate-tool-calls-or-text/TextGenerationToolCallsOrGenerateTextModel.js";
-import { AsyncQueue } from "../../util/AsyncQueue.js";
-import { parseJsonStream } from "../../util/streaming/parseJsonStream.js";
+import { TextGenerationToolCallsOrGenerateTextModel } from "../../tool/generate-tool-calls-or-text/TextGenerationToolCallsOrGenerateTextModel.js";
+import { createJsonStreamResponseHandler } from "../../util/streaming/createJsonStreamResponseHandler.js";
 import { OllamaApiConfiguration } from "./OllamaApiConfiguration.js";
 import { failedOllamaCallResponseHandler } from "./OllamaError.js";
 export class OllamaCompletionModel extends AbstractModel {
@@ -132,6 +131,10 @@ export class OllamaCompletionModel extends AbstractModel {
             responseFormat: OllamaCompletionResponseFormat.deltaIterable,
         });
     }
+    extractTextDelta(delta) {
+        const chunk = delta;
+        return chunk.done === true ? undefined : chunk.response;
+    }
     asToolCallGenerationModel(promptTemplate) {
         return new TextGenerationToolCallModel({
             model: this,
@@ -147,11 +150,25 @@ export class OllamaCompletionModel extends AbstractModel {
     withTextPrompt() {
         return this.withPromptTemplate({
             format(prompt) {
-                return { prompt: prompt };
+                return { prompt };
             },
             stopSequences: [],
         });
     }
+    /**
+     * Maps the prompt for a text version of the Ollama completion prompt template (without image support).
+     */
+    withTextPromptTemplate(promptTemplate) {
+        return new PromptTemplateTextStreamingModel({
+            model: this.withTextPrompt().withSettings({
+                stopSequences: [
+                    ...(this.settings.stopSequences ?? []),
+                    ...promptTemplate.stopSequences,
+                ],
+            }),
+            promptTemplate,
+        });
+    }
     withPromptTemplate(promptTemplate) {
         return new PromptTemplateTextStreamingModel({
             model: this.withSettings({
@@ -180,7 +197,7 @@ const ollamaCompletionResponseSchema = z.object({
     eval_duration: z.number(),
     context: z.array(z.number()).optional(),
 });
-const ollamaCompletionStreamSchema = new ZodSchema(z.discriminatedUnion("done", [
+const ollamaCompletionStreamChunkSchema = zodSchema(z.discriminatedUnion("done", [
     z.object({
         done: z.literal(false),
         model: z.string(),
@@ -202,44 +219,6 @@ const ollamaCompletionStreamSchema = new ZodSchema(z.discriminatedUnion("done",
         context: z.array(z.number()).optional(),
     }),
 ]));
-async function createOllamaFullDeltaIterableQueue(stream) {
-    const queue = new AsyncQueue();
-    let accumulatedText = "";
-    // process the stream asynchonously (no 'await' on purpose):
-    parseJsonStream({
-        stream,
-        schema: ollamaCompletionStreamSchema,
-        process(event) {
-            if (event.done === true) {
-                queue.push({
-                    type: "delta",
-                    fullDelta: {
-                        content: accumulatedText,
-                        isComplete: true,
-                        delta: "",
-                    },
-                    valueDelta: "",
-                });
-            }
-            else {
-                accumulatedText += event.response;
-                queue.push({
-                    type: "delta",
-                    fullDelta: {
-                        content: accumulatedText,
-                        isComplete: false,
-                        delta: event.response,
-                    },
-                    valueDelta: event.response,
-                });
-            }
-        },
-        onDone() {
-            queue.close();
-        },
-    });
-    return queue;
-}
 export const OllamaCompletionResponseFormat = {
     /**
      * Returns the response as a JSON object.
@@ -289,6 +268,6 @@ export const OllamaCompletionResponseFormat = {
      */
     deltaIterable: {
         stream: true,
-        handler: async ({ response }) => createOllamaFullDeltaIterableQueue(response.body),
+        handler: createJsonStreamResponseHandler(ollamaCompletionStreamChunkSchema),
     },
 };

package/model-provider/ollama/OllamaCompletionModel.test.cjs

@@ -1,25 +1,48 @@
 "use strict";
+var __createBinding = (this && this.__createBinding) || (Object.create ? (function(o, m, k, k2) {
+    if (k2 === undefined) k2 = k;
+    var desc = Object.getOwnPropertyDescriptor(m, k);
+    if (!desc || ("get" in desc ? !m.__esModule : desc.writable || desc.configurable)) {
+      desc = { enumerable: true, get: function() { return m[k]; } };
+    }
+    Object.defineProperty(o, k2, desc);
+}) : (function(o, m, k, k2) {
+    if (k2 === undefined) k2 = k;
+    o[k2] = m[k];
+}));
+var __setModuleDefault = (this && this.__setModuleDefault) || (Object.create ? (function(o, v) {
+    Object.defineProperty(o, "default", { enumerable: true, value: v });
+}) : function(o, v) {
+    o["default"] = v;
+});
+var __importStar = (this && this.__importStar) || function (mod) {
+    if (mod && mod.__esModule) return mod;
+    var result = {};
+    if (mod != null) for (var k in mod) if (k !== "default" && Object.prototype.hasOwnProperty.call(mod, k)) __createBinding(result, mod, k);
+    __setModuleDefault(result, mod);
+    return result;
+};
 Object.defineProperty(exports, "__esModule", { value: true });
 const assert_1 = require("assert");
-const msw_1 = require("msw");
-const node_1 = require("msw/node");
+const zod_1 = require("zod");
 const ApiCallError_js_1 = require("../../core/api/ApiCallError.cjs");
 const retryNever_js_1 = require("../../core/api/retryNever.cjs");
+const ZodSchema_js_1 = require("../../core/schema/ZodSchema.cjs");
+const jsonStructurePrompt_js_1 = require("../../model-function/generate-structure/jsonStructurePrompt.cjs");
+const streamStructure_js_1 = require("../../model-function/generate-structure/streamStructure.cjs");
 const generateText_js_1 = require("../../model-function/generate-text/generateText.cjs");
+const TextPrompt = __importStar(require("../../model-function/generate-text/prompt-template/TextPromptTemplate.cjs"));
+const streamText_js_1 = require("../../model-function/generate-text/streamText.cjs");
+const JsonTestServer_js_1 = require("../../test/JsonTestServer.cjs");
+const StreamingTestServer_js_1 = require("../../test/StreamingTestServer.cjs");
+const arrayFromAsync_js_1 = require("../../test/arrayFromAsync.cjs");
 const OllamaApiConfiguration_js_1 = require("./OllamaApiConfiguration.cjs");
 const OllamaCompletionModel_js_1 = require("./OllamaCompletionModel.cjs");
-// eslint-disable-next-line @typescript-eslint/no-explicit-any
-let responseBodyJson = {};
-const server = (0, node_1.setupServer)(msw_1.http.post("http://127.0.0.1:11434/api/generate", () => msw_1.HttpResponse.json(responseBodyJson)));
-beforeAll(() => server.listen());
-beforeEach(() => {
-    responseBodyJson = {};
-});
-afterEach(() => server.resetHandlers());
-afterAll(() => server.close());
 describe("generateText", () => {
+    const server = new JsonTestServer_js_1.JsonTestServer("http://127.0.0.1:11434/api/generate");
+    server.setupTestEnvironment();
     it("should return the generated text", async () => {
-        responseBodyJson = {
+        server.responseBodyJson = {
             model: "test-model",
             created_at: "2023-08-04T19:22:45.499127Z",
             response: "test response",
@@ -40,7 +63,7 @@ describe("generateText", () => {
         expect(result).toEqual("test response");
     });
     it("should throw retryable ApiCallError when Ollama is overloaded", async () => {
-        responseBodyJson = {
+        server.responseBodyJson = {
             model: "",
             created_at: "0001-01-01T00:00:00Z",
             response: "",
@@ -61,3 +84,62 @@ describe("generateText", () => {
         }
     });
 });
+describe("streamText", () => {
+    const server = new StreamingTestServer_js_1.StreamingTestServer("http://127.0.0.1:11434/api/generate");
+    server.setupTestEnvironment();
+    it("should return a text stream", async () => {
+        server.responseChunks = [
+            `{"model":"mistral:text","created_at":"2023-12-24T16:11:17.715003Z","response":"Hello","done":false}\n`,
+            `{"model":"mistral:text","created_at":"2023-12-24T16:11:17.715003Z","response":", ","done":false}\n`,
+            `{"model":"mistral:text","created_at":"2023-12-24T16:11:17.715003Z","response":"world!","done":false}\n`,
+            `{"model":"mistral:text","created_at":"2023-12-24T16:11:19.697067Z","response":"",` +
+                `"done":true,"context":[123,456,789],"total_duration":2165354041,"load_duration":1293958,` +
+                `"prompt_eval_count":5,"prompt_eval_duration":193273000,"eval_count":136,"eval_duration":1966852000}\n`,
+        ];
+        const stream = await (0, streamText_js_1.streamText)(new OllamaCompletionModel_js_1.OllamaCompletionModel({ model: "mistral:text" }).withTextPrompt(), "hello");
+        // note: space moved to last chunk bc of trimming
+        expect(await (0, arrayFromAsync_js_1.arrayFromAsync)(stream)).toStrictEqual([
+            "Hello",
+            ",",
+            " world!",
+        ]);
+    });
+});
+describe("streamStructure", () => {
+    const server = new StreamingTestServer_js_1.StreamingTestServer("http://127.0.0.1:11434/api/generate");
+    server.setupTestEnvironment();
+    it("should return a text stream", async () => {
+        server.responseChunks = [
+            `{"model":"mistral:text","created_at":"2023-12-25T11:48:02.253175Z","response":"{","done":false}\n`,
+            `{"model":"mistral:text","created_at":"2023-12-25T11:48:02.273505Z","response":"\\n","done":false}\n`,
+            `{"model":"mistral:text","created_at":"2023-12-25T11:48:02.293192Z","response":"   ","done":false}\n`,
+            `{"model":"mistral:text","created_at":"2023-12-25T11:48:02.312446Z","response":" \\"","done":false}\n`,
+            `{"model":"mistral:text","created_at":"2023-12-25T11:48:02.332021Z","response":"name","done":false}\n`,
+            `{"model":"mistral:text","created_at":"2023-12-25T11:48:02.351128Z","response":"\\":","done":false}\n`,
+            `{"model":"mistral:text","created_at":"2023-12-25T11:48:02.372082Z","response":" \\"","done":false}\n`,
+            `{"model":"mistral:text","created_at":"2023-12-25T11:48:02.391903Z","response":"M","done":false}\n`,
+            `{"model":"mistral:text","created_at":"2023-12-25T11:48:02.411056Z","response":"ike","done":false}\n`,
+            `{"model":"mistral:text","created_at":"2023-12-25T11:48:02.430789Z","response":"\\"","done":false}\n`,
+            `{"model":"mistral:text","created_at":"2023-12-25T11:48:02.450216Z","response":"\\n","done":false}\n`,
+            `{"model":"mistral:text","created_at":"2023-12-25T11:48:02.47009Z","response":"}","done":false}\n`,
+            `{"model":"mistral:text","created_at":"2023-12-25T11:48:02.48885Z","response":"","done":true,` +
+                `"total_duration":521893000,"load_duration":957666,"prompt_eval_count":74,"prompt_eval_duration":302508000,` +
+                `"eval_count":12,"eval_duration":215282000}\n`,
+        ];
+        const stream = await (0, streamStructure_js_1.streamStructure)(new OllamaCompletionModel_js_1.OllamaCompletionModel({
+            model: "mistral:text",
+            format: "json",
+            raw: true,
+        })
+            .withTextPromptTemplate(TextPrompt.instruction())
+            .asStructureGenerationModel(jsonStructurePrompt_js_1.jsonStructurePrompt.text()), (0, ZodSchema_js_1.zodSchema)(zod_1.z.object({ name: zod_1.z.string() })), "generate a name");
+        // note: space moved to last chunk bc of trimming
+        expect(await (0, arrayFromAsync_js_1.arrayFromAsync)(stream)).toStrictEqual([
+            { isComplete: false, value: {} },
+            { isComplete: false, value: { name: "" } },
+            { isComplete: false, value: { name: "M" } },
+            { isComplete: false, value: { name: "Mike" } },
+            { isComplete: true, value: { name: "Mike" } },
+        ]);
+    });
+});

package/model-provider/ollama/OllamaCompletionModel.test.js

@@ -1,23 +1,23 @@
 import { fail } from "assert";
-import { HttpResponse, http } from "msw";
-import { setupServer } from "msw/node";
+import { z } from "zod";
 import { ApiCallError } from "../../core/api/ApiCallError.js";
 import { retryNever } from "../../core/api/retryNever.js";
+import { zodSchema } from "../../core/schema/ZodSchema.js";
+import { jsonStructurePrompt } from "../../model-function/generate-structure/jsonStructurePrompt.js";
+import { streamStructure } from "../../model-function/generate-structure/streamStructure.js";
 import { generateText } from "../../model-function/generate-text/generateText.js";
+import * as TextPrompt from "../../model-function/generate-text/prompt-template/TextPromptTemplate.js";
+import { streamText } from "../../model-function/generate-text/streamText.js";
+import { JsonTestServer } from "../../test/JsonTestServer.js";
+import { StreamingTestServer } from "../../test/StreamingTestServer.js";
+import { arrayFromAsync } from "../../test/arrayFromAsync.js";
 import { OllamaApiConfiguration } from "./OllamaApiConfiguration.js";
 import { OllamaCompletionModel } from "./OllamaCompletionModel.js";
-// eslint-disable-next-line @typescript-eslint/no-explicit-any
-let responseBodyJson = {};
-const server = setupServer(http.post("http://127.0.0.1:11434/api/generate", () => HttpResponse.json(responseBodyJson)));
-beforeAll(() => server.listen());
-beforeEach(() => {
-    responseBodyJson = {};
-});
-afterEach(() => server.resetHandlers());
-afterAll(() => server.close());
 describe("generateText", () => {
+    const server = new JsonTestServer("http://127.0.0.1:11434/api/generate");
+    server.setupTestEnvironment();
     it("should return the generated text", async () => {
-        responseBodyJson = {
+        server.responseBodyJson = {
             model: "test-model",
             created_at: "2023-08-04T19:22:45.499127Z",
             response: "test response",
@@ -38,7 +38,7 @@ describe("generateText", () => {
         expect(result).toEqual("test response");
     });
     it("should throw retryable ApiCallError when Ollama is overloaded", async () => {
-        responseBodyJson = {
+        server.responseBodyJson = {
             model: "",
             created_at: "0001-01-01T00:00:00Z",
             response: "",
@@ -59,3 +59,62 @@ describe("generateText", () => {
         }
     });
 });
+describe("streamText", () => {
+    const server = new StreamingTestServer("http://127.0.0.1:11434/api/generate");
+    server.setupTestEnvironment();
+    it("should return a text stream", async () => {
+        server.responseChunks = [
+            `{"model":"mistral:text","created_at":"2023-12-24T16:11:17.715003Z","response":"Hello","done":false}\n`,
+            `{"model":"mistral:text","created_at":"2023-12-24T16:11:17.715003Z","response":", ","done":false}\n`,
+            `{"model":"mistral:text","created_at":"2023-12-24T16:11:17.715003Z","response":"world!","done":false}\n`,
+            `{"model":"mistral:text","created_at":"2023-12-24T16:11:19.697067Z","response":"",` +
+                `"done":true,"context":[123,456,789],"total_duration":2165354041,"load_duration":1293958,` +
+                `"prompt_eval_count":5,"prompt_eval_duration":193273000,"eval_count":136,"eval_duration":1966852000}\n`,
+        ];
+        const stream = await streamText(new OllamaCompletionModel({ model: "mistral:text" }).withTextPrompt(), "hello");
+        // note: space moved to last chunk bc of trimming
+        expect(await arrayFromAsync(stream)).toStrictEqual([
+            "Hello",
+            ",",
+            " world!",
+        ]);
+    });
+});
+describe("streamStructure", () => {
+    const server = new StreamingTestServer("http://127.0.0.1:11434/api/generate");
+    server.setupTestEnvironment();
+    it("should return a text stream", async () => {
+        server.responseChunks = [
+            `{"model":"mistral:text","created_at":"2023-12-25T11:48:02.253175Z","response":"{","done":false}\n`,
+            `{"model":"mistral:text","created_at":"2023-12-25T11:48:02.273505Z","response":"\\n","done":false}\n`,
+            `{"model":"mistral:text","created_at":"2023-12-25T11:48:02.293192Z","response":"   ","done":false}\n`,
+            `{"model":"mistral:text","created_at":"2023-12-25T11:48:02.312446Z","response":" \\"","done":false}\n`,
+            `{"model":"mistral:text","created_at":"2023-12-25T11:48:02.332021Z","response":"name","done":false}\n`,
+            `{"model":"mistral:text","created_at":"2023-12-25T11:48:02.351128Z","response":"\\":","done":false}\n`,
+            `{"model":"mistral:text","created_at":"2023-12-25T11:48:02.372082Z","response":" \\"","done":false}\n`,
+            `{"model":"mistral:text","created_at":"2023-12-25T11:48:02.391903Z","response":"M","done":false}\n`,
+            `{"model":"mistral:text","created_at":"2023-12-25T11:48:02.411056Z","response":"ike","done":false}\n`,
+            `{"model":"mistral:text","created_at":"2023-12-25T11:48:02.430789Z","response":"\\"","done":false}\n`,
+            `{"model":"mistral:text","created_at":"2023-12-25T11:48:02.450216Z","response":"\\n","done":false}\n`,
+            `{"model":"mistral:text","created_at":"2023-12-25T11:48:02.47009Z","response":"}","done":false}\n`,
+            `{"model":"mistral:text","created_at":"2023-12-25T11:48:02.48885Z","response":"","done":true,` +
+                `"total_duration":521893000,"load_duration":957666,"prompt_eval_count":74,"prompt_eval_duration":302508000,` +
+                `"eval_count":12,"eval_duration":215282000}\n`,
+        ];
+        const stream = await streamStructure(new OllamaCompletionModel({
+            model: "mistral:text",
+            format: "json",
+            raw: true,
+        })
+            .withTextPromptTemplate(TextPrompt.instruction())
+            .asStructureGenerationModel(jsonStructurePrompt.text()), zodSchema(z.object({ name: z.string() })), "generate a name");
+        // note: space moved to last chunk bc of trimming
+        expect(await arrayFromAsync(stream)).toStrictEqual([
+            { isComplete: false, value: {} },
+            { isComplete: false, value: { name: "" } },
+            { isComplete: false, value: { name: "M" } },
+            { isComplete: false, value: { name: "Mike" } },
+            { isComplete: true, value: { name: "Mike" } },
+        ]);
+    });
+});