@agentscope-ai/agentscope 0.0.2

package/src/model/ollama-model.test.ts

@@ -0,0 +1,307 @@
+import { OllamaChatModel } from './ollama-model';
+import { ChatResponse } from './response';
+import { createMsg } from '../message';
+
+// Mock fetch for streaming responses
+global.fetch = jest.fn();
+
+describe('OllamaChatModel', () => {
+    beforeEach(() => {
+        jest.clearAllMocks();
+    });
+
+    test('Test stream generation with delta output', async () => {
+        // Mock streaming response with multiple chunks (NDJSON format)
+        // Ollama returns newline-delimited JSON
+        const mockStreamChunks = [
+            '{"model":"qwen3:1.7b","created_at":"2024-01-01T00:00:00Z","message":{"role":"assistant","content":"Let me"},"done":false}\n',
+            '{"model":"qwen3:1.7b","created_at":"2024-01-01T00:00:00Z","message":{"role":"assistant","content":" check the weather"},"done":false}\n',
+            '{"model":"qwen3:1.7b","created_at":"2024-01-01T00:00:00Z","message":{"role":"assistant","content":"","tool_calls":[{"function":{"name":"get_current_weather","arguments":{"location":"Beijing"}}}]},"done":false}\n',
+            '{"model":"qwen3:1.7b","created_at":"2024-01-01T00:00:00Z","message":{"role":"assistant","content":""},"done":true,"prompt_eval_count":100,"eval_count":50}\n',
+        ];
+
+        const mockReadableStream = new ReadableStream({
+            start(controller) {
+                mockStreamChunks.forEach(chunk =>
+                    controller.enqueue(new TextEncoder().encode(chunk))
+                );
+                controller.close();
+            },
+        });
+
+        (global.fetch as jest.Mock).mockResolvedValue({
+            ok: true,
+            body: mockReadableStream,
+        });
+
+        const model = new OllamaChatModel({
+            modelName: 'qwen3:1.7b',
+            stream: true,
+            host: 'http://localhost:11434',
+        });
+
+        const res = await model.call({
+            messages: [
+                createMsg({
+                    name: 'user',
+                    role: 'user',
+                    content: [
+                        { type: 'text', text: "How's the weather today?", id: crypto.randomUUID() },
+                    ],
+                }),
+            ],
+            tools: [
+                {
+                    type: 'function',
+                    function: {
+                        name: 'get_current_weather',
+                        description: 'Get the current weather in a given location',
+                        parameters: {
+                            type: 'object',
+                            properties: {
+                                location: {
+                                    type: 'string',
+                                    description: 'The city and state, e.g. San Francisco, CA',
+                                },
+                            },
+                            required: ['location'],
+                        },
+                    },
+                },
+            ],
+        });
+
+        const generator = res as AsyncGenerator<ChatResponse, ChatResponse>;
+        let completeResponse: ChatResponse | undefined;
+        const yieldedChunks: ChatResponse[] = [];
+
+        // Manual iteration to capture both yielded and returned values
+        while (true) {
+            const result = await generator.next();
+            if (result.done) {
+                completeResponse = result.value;
+                break;
+            }
+            yieldedChunks.push(result.value);
+        }
+
+        // Verify we received multiple yielded chunks
+        expect(yieldedChunks.length).toBeGreaterThan(0);
+
+        // Verify the final complete response has correct structure
+        expect(completeResponse.content.length).toBe(2);
+
+        // Check text block - should be complete after accumulation
+        const textBlock = completeResponse.content.find(b => b.type === 'text');
+        expect(textBlock).toBeDefined();
+        expect(textBlock).toMatchObject({
+            type: 'text',
+            text: 'Let me check the weather',
+        });
+
+        // Check tool_call block
+        const toolCallBlock = completeResponse.content.find(b => b.type === 'tool_call');
+        expect(toolCallBlock).toBeDefined();
+        expect(toolCallBlock).toMatchObject({
+            type: 'tool_call',
+            name: 'get_current_weather',
+            input: '{"location":"Beijing"}',
+        });
+
+        // Verify usage
+        expect(completeResponse.usage).toBeDefined();
+        expect(completeResponse.usage?.inputTokens).toBe(100);
+        expect(completeResponse.usage?.outputTokens).toBe(50);
+    }, 10000);
+
+    test('Test non-streaming generation', async () => {
+        // Mock non-streaming response
+        const mockResponse = {
+            model: 'qwen3:8b',
+            created_at: '2024-01-01T00:00:00Z',
+            message: {
+                role: 'assistant',
+                content: '你好！我是一个AI助手。',
+            },
+            done: true,
+            prompt_eval_count: 50,
+            eval_count: 30,
+        };
+
+        (global.fetch as jest.Mock).mockResolvedValue({
+            ok: true,
+            json: async () => mockResponse,
+        });
+
+        const model = new OllamaChatModel({
+            modelName: 'qwen3:8b',
+            stream: false,
+            host: 'http://localhost:11434',
+        });
+
+        const res = await model.call({
+            messages: [
+                createMsg({
+                    name: 'user',
+                    role: 'user',
+                    content: [
+                        { type: 'text', text: '你好，请简单介绍一下自己', id: crypto.randomUUID() },
+                    ],
+                }),
+            ],
+        });
+
+        const completeResponse = res as ChatResponse;
+
+        // Verify complete response structure
+        expect(completeResponse.content.length).toBe(1);
+
+        // Check text block
+        const textBlock = completeResponse.content.find(b => b.type === 'text');
+        expect(textBlock).toBeDefined();
+        expect(textBlock).toMatchObject({
+            type: 'text',
+            text: '你好！我是一个AI助手。',
+        });
+
+        // Verify usage
+        expect(completeResponse.usage).toBeDefined();
+        expect(completeResponse.usage?.inputTokens).toBe(50);
+        expect(completeResponse.usage?.outputTokens).toBe(30);
+    }, 10000);
+
+    test('Test with thinking enabled', async () => {
+        // Mock streaming response with thinking
+        const mockStreamChunks = [
+            '{"model":"qwen3:8b","created_at":"2024-01-01T00:00:00Z","message":{"role":"assistant","thinking":"计算"},"done":false}\n',
+            '{"model":"qwen3:8b","created_at":"2024-01-01T00:00:00Z","message":{"role":"assistant","thinking":" 123 * 456"},"done":false}\n',
+            '{"model":"qwen3:8b","created_at":"2024-01-01T00:00:00Z","message":{"role":"assistant","content":"答案是 56088"},"done":false}\n',
+            '{"model":"qwen3:8b","created_at":"2024-01-01T00:00:00Z","message":{"role":"assistant","content":""},"done":true,"prompt_eval_count":80,"eval_count":40}\n',
+        ];
+
+        const mockReadableStream = new ReadableStream({
+            start(controller) {
+                mockStreamChunks.forEach(chunk =>
+                    controller.enqueue(new TextEncoder().encode(chunk))
+                );
+                controller.close();
+            },
+        });
+
+        (global.fetch as jest.Mock).mockResolvedValue({
+            ok: true,
+            body: mockReadableStream,
+        });
+
+        const model = new OllamaChatModel({
+            modelName: 'qwen3:8b',
+            stream: true,
+            thinkingConfig: {
+                enableThinking: true,
+            },
+            host: 'http://localhost:11434',
+        });
+
+        const res = await model.call({
+            messages: [
+                createMsg({
+                    name: 'user',
+                    role: 'user',
+                    content: [
+                        {
+                            type: 'text',
+                            text: '计算 123 * 456 等于多少？',
+                            id: crypto.randomUUID(),
+                        },
+                    ],
+                }),
+            ],
+        });
+
+        const generator = res as AsyncGenerator<ChatResponse, ChatResponse>;
+        let completeResponse: ChatResponse | undefined;
+        const yieldedChunks: ChatResponse[] = [];
+
+        while (true) {
+            const result = await generator.next();
+            if (result.done) {
+                completeResponse = result.value;
+                break;
+            }
+            yieldedChunks.push(result.value);
+        }
+
+        // Verify we received multiple yielded chunks
+        expect(yieldedChunks.length).toBeGreaterThan(0);
+
+        // Verify the final complete response has correct structure
+        expect(completeResponse.content.length).toBe(2);
+
+        // Check thinking block
+        const thinkingBlock = completeResponse.content.find(b => b.type === 'thinking');
+        expect(thinkingBlock).toBeDefined();
+        expect(thinkingBlock).toMatchObject({
+            type: 'thinking',
+            thinking: '计算 123 * 456',
+        });
+
+        // Check text block
+        const textBlock = completeResponse.content.find(b => b.type === 'text');
+        expect(textBlock).toBeDefined();
+        expect(textBlock).toMatchObject({
+            type: 'text',
+            text: '答案是 56088',
+        });
+
+        // Verify usage
+        expect(completeResponse.usage).toBeDefined();
+        expect(completeResponse.usage?.inputTokens).toBe(80);
+        expect(completeResponse.usage?.outputTokens).toBe(40);
+    }, 10000);
+
+    test('Test formatToolChoice function', () => {
+        const model = new OllamaChatModel({
+            modelName: 'qwen3:1.7b',
+            host: 'http://localhost:11434',
+        });
+
+        // Ollama's _formatToolChoice always returns undefined
+        expect(model._formatToolChoice('auto')).toBeUndefined();
+        expect(model._formatToolChoice('none')).toBeUndefined();
+        expect(model._formatToolChoice('my_function')).toBeUndefined();
+        expect(model._formatToolChoice(undefined)).toBeUndefined();
+    });
+
+    test('Test formatToolSchemas function', () => {
+        const model = new OllamaChatModel({
+            modelName: 'qwen3:1.7b',
+            host: 'http://localhost:11434',
+        });
+
+        const toolSchemas = [
+            {
+                type: 'function' as const,
+                function: {
+                    name: 'get_current_weather',
+                    description: 'Get the current weather in a given location',
+                    parameters: {
+                        type: 'object' as const,
+                        properties: {
+                            location: {
+                                type: 'string',
+                                description: 'The city and state, e.g. San Francisco, CA',
+                            },
+                        },
+                        required: ['location'],
+                    },
+                },
+            },
+        ];
+
+        // Test with tool schemas
+        expect(model._formatToolSchemas(toolSchemas)).toEqual(toolSchemas);
+
+        // Test with undefined (should return empty array)
+        expect(model._formatToolSchemas(undefined)).toEqual([]);
+    });
+});

package/src/model/ollama-model.ts

@@ -0,0 +1,356 @@
+import { Ollama, ChatResponse as OllamaChatResponse, AbortableAsyncIterator } from 'ollama';
+
+import { ChatModelBase, ChatModelOptions, ChatModelRequestOptions } from './base';
+import { ChatResponse } from './response';
+import { TextBlock, ThinkingBlock, ToolCallBlock } from '../message';
+import { ToolChoice, ToolSchema } from '../type';
+import { ChatUsage } from './usage';
+import { OllamaChatFormatter } from '../formatter';
+
+interface OllamaThinkingConfig {
+    /**
+     * Whether to enable thinking or not.
+     */
+    enableThinking: boolean;
+
+    /**
+     * Thinking level for Ollama models (high, medium, low).
+     * Only applicable when enableThinking is true.
+     */
+    thinkingLevel?: 'high' | 'medium' | 'low';
+}
+
+interface OllamaChatModelOptions extends ChatModelOptions {
+    /**
+     * Additional parameters to pass to the Ollama API (e.g., temperature).
+     */
+    options?: Record<string, unknown>;
+
+    /**
+     * Duration to keep the model loaded in memory (e.g., "5m", "1h").
+     */
+    keepAlive?: string;
+
+    /**
+     * Thinking configuration for Ollama models.
+     */
+    thinkingConfig?: OllamaThinkingConfig;
+
+    /**
+     * The host address of the Ollama server.
+     */
+    host?: string;
+
+    /**
+     * Extra keyword arguments to initialize the Ollama client.
+     */
+    clientKwargs?: Record<string, unknown>;
+
+    /**
+     * Extra keyword arguments used in Ollama API generation.
+     */
+    generateKwargs?: Record<string, unknown>;
+}
+
+/**
+ * The Ollama chat model class in AgentScope.
+ */
+export class OllamaChatModel extends ChatModelBase {
+    protected client: Ollama;
+    protected options?: Record<string, unknown>;
+    protected keepAlive: string;
+    protected thinkingConfig: OllamaThinkingConfig;
+    protected generateKwargs: Record<string, unknown>;
+
+    /**
+     * Initializes a new instance of the OllamaChatModel class.
+     * @param root0
+     * @param root0.modelName
+     * @param root0.stream
+     * @param root0.options
+     * @param root0.keepAlive
+     * @param root0.thinkingConfig
+     * @param root0.host
+     * @param root0.maxRetries
+     * @param root0.fallbackModelName
+     * @param root0.clientKwargs
+     * @param root0.generateKwargs
+     * @param root0.formatter
+     */
+    constructor({
+        modelName,
+        stream = true,
+        options,
+        keepAlive = '5m',
+        thinkingConfig,
+        host,
+        maxRetries = 0,
+        fallbackModelName,
+        clientKwargs,
+        generateKwargs,
+        formatter,
+    }: OllamaChatModelOptions) {
+        // If no formatter is provided, create a default OllamaChatFormatter
+        const defaultFormatter = formatter || new OllamaChatFormatter();
+        super({
+            modelName,
+            stream,
+            maxRetries,
+            fallbackModelName,
+            formatter: defaultFormatter,
+        } as ChatModelOptions);
+
+        this.options = options;
+        this.keepAlive = keepAlive;
+        this.thinkingConfig = thinkingConfig || {
+            enableThinking: false,
+        };
+        this.generateKwargs = generateKwargs || {};
+
+        // Initialize Ollama client
+        this.client = new Ollama({
+            host: host,
+            ...clientKwargs,
+        });
+    }
+
+    /**
+     * Calls the Ollama API with the given parameters.
+     * @param modelName
+     * @param options
+     * @returns A promise that resolves to either a ChatResponse or an AsyncGenerator of ChatResponses.
+     */
+    async _callAPI(
+        modelName: string,
+        options: ChatModelRequestOptions<Record<string, unknown>>
+    ): Promise<ChatResponse | AsyncGenerator<ChatResponse, ChatResponse>> {
+        const kwargs: Record<string, unknown> = {
+            model: modelName,
+            messages: options.messages,
+            stream: this.stream,
+            options: this.options,
+            keep_alive: this.keepAlive,
+            ...this.generateKwargs,
+        };
+
+        if (this.thinkingConfig.enableThinking) {
+            // If thinkingLevel is specified, use it; otherwise use true
+            kwargs.think = this.thinkingConfig.thinkingLevel || true;
+        } else {
+            kwargs.think = false;
+        }
+
+        if (options.tools) {
+            kwargs.tools = this._formatToolSchemas(options.tools);
+        }
+
+        if (options.toolChoice) {
+            console.warn('Ollama does not support tool_choice yet, ignored.');
+        }
+
+        const startTime = Date.now();
+
+        if (this.stream) {
+            const response = (await this.client.chat({
+                ...kwargs,
+                stream: true,
+            } as Parameters<
+                typeof this.client.chat
+            >[0])) as unknown as AbortableAsyncIterator<OllamaChatResponse>;
+            return this._parseOllamaStreamResponse(response, startTime);
+        }
+
+        const response = (await this.client.chat({
+            ...kwargs,
+            stream: false,
+        } as Parameters<typeof this.client.chat>[0])) as unknown as OllamaChatResponse;
+        return this._parseOllamaResponse(response, startTime);
+    }
+
+    /**
+     * Parse Ollama streaming response.
+     * @param stream
+     * @param startTime
+     * @returns An async generator that yields delta ChatResponse objects and returns the complete ChatResponse.
+     */
+    async *_parseOllamaStreamResponse(
+        stream: AbortableAsyncIterator<OllamaChatResponse>,
+        startTime: number
+    ): AsyncGenerator<ChatResponse, ChatResponse> {
+        let accText = '';
+        let accThinking = '';
+        const toolCalls: Map<string, ToolCallBlock> = new Map();
+        let lastUsage: ChatUsage | null = null;
+
+        for await (const chunk of stream) {
+            const msg = chunk.message;
+
+            // Delta data for this chunk
+            let deltaText = '';
+            let deltaThinking = '';
+            const deltaToolCalls: Map<string, ToolCallBlock> = new Map();
+
+            // Accumulate text and thinking
+            if (msg.thinking) {
+                deltaThinking = msg.thinking;
+                accThinking += msg.thinking;
+            }
+            if (msg.content) {
+                deltaText = msg.content;
+                accText += msg.content;
+            }
+
+            // Handle tool calls
+            if (msg.tool_calls && Array.isArray(msg.tool_calls)) {
+                for (let idx = 0; idx < msg.tool_calls.length; idx++) {
+                    const toolCall = msg.tool_calls[idx];
+                    const func = toolCall.function;
+                    const toolId = `${idx}_${func.name}`;
+
+                    const toolCallBlock = {
+                        type: 'tool_call' as const,
+                        id: toolId,
+                        name: func.name,
+                        input: JSON.stringify(func.arguments),
+                    };
+
+                    toolCalls.set(toolId, toolCallBlock);
+                    deltaToolCalls.set(toolId, toolCallBlock);
+                }
+            }
+
+            // Calculate usage
+            const currentTime = (Date.now() - startTime) / 1000;
+            lastUsage = {
+                type: 'chat_usage',
+                inputTokens: chunk.prompt_eval_count || 0,
+                outputTokens: chunk.eval_count || 0,
+                time: currentTime,
+            };
+
+            // Yield delta response
+            const deltaBlocks = this._buildContentBlocks(deltaText, deltaThinking, deltaToolCalls);
+            yield {
+                type: 'chat',
+                id: crypto.randomUUID(),
+                createdAt: new Date().toISOString(),
+                content: deltaBlocks,
+                usage: lastUsage,
+            } as ChatResponse;
+        }
+
+        // Return complete response
+        const blocks = this._buildContentBlocks(accText, accThinking, toolCalls);
+        return {
+            type: 'chat',
+            id: crypto.randomUUID(),
+            createdAt: new Date().toISOString(),
+            content: blocks,
+            usage: lastUsage,
+        } as ChatResponse;
+    }
+
+    /**
+     * Parse Ollama non-streaming response.
+     * @param response
+     * @param startTime
+     * @returns A ChatResponse object containing the content blocks and usage.
+     */
+    _parseOllamaResponse(response: OllamaChatResponse, startTime: number): ChatResponse {
+        const blocks: Array<TextBlock | ThinkingBlock | ToolCallBlock> = [];
+
+        if (response.message.thinking) {
+            blocks.push({
+                id: crypto.randomUUID(),
+                type: 'thinking',
+                thinking: response.message.thinking,
+            });
+        }
+
+        if (response.message.content) {
+            blocks.push({
+                id: crypto.randomUUID(),
+                type: 'text',
+                text: response.message.content,
+            });
+        }
+
+        // Handle tool calls
+        if (response.message.tool_calls && Array.isArray(response.message.tool_calls)) {
+            for (let idx = 0; idx < response.message.tool_calls.length; idx++) {
+                const toolCall = response.message.tool_calls[idx];
+                blocks.push({
+                    type: 'tool_call',
+                    id: `${idx}_${toolCall.function.name}`,
+                    name: toolCall.function.name,
+                    input: JSON.stringify(toolCall.function.arguments),
+                });
+            }
+        }
+
+        const usage =
+            response.prompt_eval_count !== undefined && response.eval_count !== undefined
+                ? {
+                      type: 'chat_usage',
+                      inputTokens: response.prompt_eval_count || 0,
+                      outputTokens: response.eval_count || 0,
+                      time: (Date.now() - startTime) / 1000,
+                  }
+                : undefined;
+
+        return {
+            type: 'chat',
+            id: crypto.randomUUID(),
+            createdAt: new Date().toISOString(),
+            content: blocks,
+            usage,
+        } as ChatResponse;
+    }
+
+    /**
+     * Build content blocks from accumulated data.
+     * @param text
+     * @param thinking
+     * @param toolCalls
+     * @returns An array of content blocks.
+     */
+    _buildContentBlocks(
+        text: string,
+        thinking: string,
+        toolCalls: Map<string, ToolCallBlock>
+    ): Array<TextBlock | ThinkingBlock | ToolCallBlock> {
+        const blocks: Array<TextBlock | ThinkingBlock | ToolCallBlock> = [];
+
+        if (thinking) {
+            blocks.push({ id: crypto.randomUUID(), type: 'thinking', thinking });
+        }
+
+        if (text) {
+            blocks.push({ id: crypto.randomUUID(), type: 'text', text });
+        }
+
+        toolCalls.forEach(toolCall => {
+            blocks.push(toolCall);
+        });
+
+        return blocks;
+    }
+
+    /**
+     * Format tool choice parameter (not supported by Ollama).
+     * @param _toolChoice
+     * @returns undefined as Ollama does not support tool choice.
+     */
+    _formatToolChoice(_toolChoice?: ToolChoice): unknown {
+        return undefined;
+    }
+
+    /**
+     * Format tool schemas for Ollama API (no special formatting needed).
+     * @param tools
+     * @returns The same array of tool schemas, or an empty array if undefined.
+     */
+    _formatToolSchemas(tools: ToolSchema[] | undefined): ToolSchema[] {
+        return tools || [];
+    }
+}