📦 NEW: Stream support in tool call (#101)

* 📦 NEW: Stream support when tools are present

* 📦 NEW: Stream tool call example

* 👌 IMPROVE: Code

* 🐛 FIX: Anthropic tool call stream support

* 👌 IMPROVE: Code

* 👌 IMPROVE: Code

* 📦 NEW: Anthropic check

Author: msaaddev

examples/nextjs/app/api/langbase/pipes/run-tool-stream/route.ts

@@ -0,0 +1,23 @@
+import pipeWithToolsStream from '@/baseai/pipes/pipe-with-tool-stream';
+import {Pipe, RunResponseStream} from '@baseai/core';
+import {NextRequest} from 'next/server';
+
+export async function POST(req: NextRequest) {
+	const runOptions = await req.json();
+
+	// 1. Initiate the Pipe.
+	const pipe = new Pipe(pipeWithToolsStream());
+
+	// 2. Run the pipe with user messages and other run options.
+	let {stream, threadId} = (await pipe.run(
+		runOptions,
+	)) as unknown as RunResponseStream;
+
+	// 3. Stream the response.
+	return new Response(stream, {
+		status: 200,
+		headers: {
+			'lb-thread-id': threadId ?? '',
+		},
+	});
+}

examples/nextjs/app/demo/tool-calling-stream/page.tsx

@@ -0,0 +1,18 @@
+import PipeRunToolStreamExample from '@/components/pipe-run-with-tool-stream';
+import GoHome from '@/components/ui/go-home';
+
+export default function Page() {
+	return (
+		<div className="w-full max-w-md">
+			<GoHome />
+
+			<h1 className="text-2xl font-light text-gray-800 mb-1 text-center">
+				AI Agent Pipes: Tool Calling
+			</h1>
+			<p className="text-muted-foreground text-base font-light mb-20 text-center">
+				Run a pipe with tool calling.
+			</p>
+			<PipeRunToolStreamExample />
+		</div>
+	);
+}

examples/nextjs/app/page.tsx

@@ -2,16 +2,17 @@ import Link from 'next/link';
 
 export default function Page() {
 	const examples = [
-		{title: 'Pipe Run', href: '/demo/pipe-run'},
-		{title: 'Pipe Run Stream', href: '/demo/pipe-run-stream'},
-		{title: 'Chat Simple', href: '/demo/chat-simple'},
-		{title: 'Chat Advanced', href: '/demo/chat-advanced'},
-		{title: 'Tool Calling', href: '/demo/tool-calling'},
+		{ title: 'Pipe Run', href: '/demo/pipe-run' },
+		{ title: 'Pipe Run Stream', href: '/demo/pipe-run-stream' },
+		{ title: 'Chat Simple', href: '/demo/chat-simple' },
+		{ title: 'Chat Advanced', href: '/demo/chat-advanced' },
+		{ title: 'Tool Calling', href: '/demo/tool-calling' },
+		{ title: 'Tool Calling Stream', href: '/demo/tool-calling-stream' },
 		{
 			title: 'Tool Calling: Pipes as Tools',
 			href: '/demo/pipe-run-pipes-as-tools',
 		},
-		{title: 'Memory', href: '/demo/memory'},
+		{ title: 'Memory', href: '/demo/memory' },
 	];
 
 	return (

examples/nextjs/baseai/pipes/pipe-with-tool-stream.ts

@@ -0,0 +1,28 @@
+import {PipeI} from '@baseai/core';
+import toolCalculator from '../tools/calculator';
+import toolGetWeather from '../tools/weather';
+
+const pipeWithToolsStream = (): PipeI => ({
+	apiKey: process.env.LANGBASE_API_KEY!,
+	name: 'pipe-with-tool',
+	description: 'An AI agent pipe that can call tools',
+	status: 'public',
+	model: 'openai:gpt-4o-mini',
+	stream: true,
+	json: false,
+	store: true,
+	moderate: true,
+	top_p: 1,
+	max_tokens: 1000,
+	temperature: 0.7,
+	presence_penalty: 1,
+	frequency_penalty: 1,
+	stop: [],
+	tool_choice: 'auto',
+	parallel_tool_calls: true,
+	messages: [{role: 'system', content: `You are a helpful AI assistant.`}],
+	variables: [],
+	memory: [],
+	tools: [toolGetWeather(), toolCalculator()],
+});
+export default pipeWithToolsStream;

examples/nextjs/baseai/pipes/pipe-with-tool.ts

@@ -8,7 +8,7 @@ const pipeWithTools = (): PipeI => ({
 	description: 'An AI agent pipe that can call tools',
 	status: 'public',
 	model: 'openai:gpt-4o-mini',
-	stream: true,
+	stream: false,
 	json: false,
 	store: true,
 	moderate: true,

examples/nextjs/components/pipe-run-with-tool-stream.tsx

@@ -0,0 +1,75 @@
+'use client';
+
+import { Button } from '@/components/ui/button';
+import { Input } from '@/components/ui/input';
+import { getRunner, getTextContent } from '@baseai/core';
+import { useState } from 'react';
+
+export default function PipeRunToolStreamExample() {
+	const [prompt, setPrompt] = useState(
+		'What is the weather in SF. Square root of 9 and then add 7?',
+	);
+	const [completion, setCompletion] = useState('');
+	const [loading, setLoading] = useState(false);
+
+	const handleSubmit = async (e: any) => {
+		e.preventDefault();
+		if (!prompt.trim()) return;
+
+		setLoading(true);
+		try {
+			const response = await fetch('/api/langbase/pipes/run-tool-stream', {
+				method: 'POST',
+				headers: { 'Content-Type': 'application/json' },
+				body: JSON.stringify({
+					messages: [{ role: 'user', content: prompt }],
+				}),
+			});
+
+			if (!response.ok) {
+				throw new Error('Network response was not ok');
+			}
+
+			const runner = getRunner(response.body as ReadableStream<Uint8Array>);
+
+			let localCompletion = '';
+			for await (const chunk of runner) {
+				const textPart = getTextContent(chunk);
+				localCompletion += textPart;
+				setCompletion(localCompletion);
+			}
+		} catch (error) {
+			console.error('Error:', error);
+			setCompletion('An error occurred while generating the completion.');
+		} finally {
+			setLoading(false);
+		}
+	};
+
+	return (
+		<div className="bg-neutral-200 rounded-md p-2 flex flex-col gap-2 w-full">
+			<form
+				onSubmit={handleSubmit}
+				className="flex flex-col w-full items-center gap-2"
+			>
+				<Input
+					type="text"
+					placeholder="Enter prompt message here"
+					value={prompt}
+					onChange={e => setPrompt(e.target.value)}
+					required
+				/>
+
+				<Button type="submit" className="w-full" disabled={loading}>
+					{loading ? 'AI is thinking...' : 'Ask AI'}
+				</Button>
+			</form>
+
+			{!loading && completion && (
+				<p className="mt-4">
+					<strong>AI:</strong> {completion}
+				</p>
+			)}
+		</div>
+	);
+}

packages/core/src/helpers/stream.ts

@@ -1,6 +1,7 @@
 import {ChatCompletionStream} from 'openai/lib/ChatCompletionStream';
 import {ChunkStream} from 'src/pipes';
 import {Stream} from 'openai/streaming';
+import {ToolCall} from 'types/pipes';
 
 export interface Runner extends ChatCompletionStream<null> {}
 
@@ -81,3 +82,17 @@ export function handleResponseStream({
 	}
 	return result;
 }
+
+/**
+ * Retrieves tool calls from a given readable stream.
+ *
+ * @param stream - The readable stream from which to extract tool calls.
+ * @returns A promise that resolves to an array of `ToolCall` objects.
+ */
+export async function getToolsFromStream(
+	stream: ReadableStream<any>,
+): Promise<ToolCall[]> {
+	let run = getRunner(stream);
+	const {choices} = await run.finalChatCompletion();
+	return choices[0].message.tool_calls;
+}

packages/core/src/pipes/pipes.ts

@@ -5,6 +5,8 @@ import {getLLMApiKey} from '../utils/get-llm-api-key';
 import {getApiUrl, isProd} from '../utils/is-prod';
 import {toOldPipeFormat} from '../utils/to-old-pipe-format';
 import {isLocalServerRunning} from 'src/utils/local-server-running';
+import {getToolsFromStream} from 'src/helpers';
+import {ANTHROPIC} from 'src/data/models';
 
 export interface Variable {
 	name: string;
@@ -140,17 +142,89 @@ export class Pipe {
 	}
 
 	private isStreamRequested(options: RunOptions | RunOptionsStream): boolean {
-		return 'stream' in options && options.stream === true;
+		return (
+			('stream' in options && options.stream === true) ||
+			this.pipe.meta.stream
+		);
 	}
 
 	private warnIfToolsWithStream(requestedStream: boolean): void {
 		if (this.hasTools && requestedStream) {
 			console.warn(
-				'Warning: Streaming is not yet supported when tools are present in the pipe. Falling back to non-streaming mode.',
+				'Warning: Streaming is not yet supported in Anthropic models when tools are present in the pipe. Falling back to non-streaming mode.',
 			);
 		}
 	}
 
+	private async handleStreamResponse(
+		options: RunOptionsStream,
+		response: RunResponseStream,
+	): Promise<RunResponseStream> {
+		const endpoint = '/beta/pipes/run';
+		const stream = this.isStreamRequested(options);
+		const body = {...options, stream};
+
+		const [streamForToolCall, streamForReturn] = response.stream.tee();
+		const tools = await getToolsFromStream(streamForToolCall);
+
+		if (tools.length) {
+			let messages = options.messages || [];
+
+			let currentResponse: RunResponseStream = {
+				stream: streamForReturn,
+				threadId: response.threadId,
+				rawResponse: response.rawResponse,
+			};
+
+			let callCount = 0;
+
+			while (callCount < this.maxCalls) {
+				const [streamForToolCall, streamForReturn] =
+					currentResponse.stream.tee();
+
+				const tools = await getToolsFromStream(streamForToolCall);
+
+				if (tools.length === 0) {
+					return {
+						stream: streamForReturn,
+						threadId: currentResponse.threadId,
+						rawResponse: response.rawResponse,
+					};
+				}
+
+				const toolResults = await this.runTools(tools);
+
+				const responseMessage = {
+					role: 'assistant',
+					content: null,
+					tool_calls: tools,
+				} as Message;
+
+				messages = this.getMessagesToSend(
+					messages,
+					responseMessage,
+					toolResults,
+				);
+
+				currentResponse = await this.createRequest<RunResponseStream>(
+					endpoint,
+					{
+						...body,
+						messages,
+						threadId: currentResponse.threadId,
+					},
+				);
+
+				callCount++;
+			}
+		}
+
+		return {
+			...response,
+			stream: streamForReturn,
+		} as RunResponseStream;
+	}
+
 	public async run(options: RunOptionsStream): Promise<RunResponseStream>;
 	public async run(options: RunOptions): Promise<RunResponse>;
 	public async run(
@@ -163,9 +237,16 @@ export class Pipe {
 		// logger('pipe.run.options');
 		// logger(options, {depth: null, colors: true});
 
-		const requestedStream = this.isStreamRequested(options);
-		const stream = this.hasTools ? false : requestedStream;
-		this.warnIfToolsWithStream(requestedStream);
+		const isAnthropic = this.pipe.model.provider === ANTHROPIC;
+		const hasTools = this.pipe.tools.length > 0;
+
+		let stream = this.isStreamRequested(options);
+
+		// Anthropic models don't support streaming with tools.
+		if (isAnthropic && hasTools && stream) {
+			this.warnIfToolsWithStream(stream);
+			stream = false;
+		}
 
 		const runTools = options.runTools ?? true;
 		delete options.runTools;
@@ -179,10 +260,18 @@ export class Pipe {
 			return {} as RunResponse | RunResponseStream;
 		}
 
-		if (stream || !runTools) {
+		if (!runTools) {
 			return response as RunResponseStream;
 		}
 
+		if (stream) {
+			return await this.handleStreamResponse(
+				options as RunOptionsStream,
+				response as RunResponseStream,
+			);
+		}
+
+		// STREAM IS OFF
 		let messages = options.messages || [];
 		let currentResponse = response as RunResponse;
 		let callCount = 0;