Update README for new tool call flow

Author: scotttrinh

packages/ai/README.md

@@ -31,21 +31,21 @@ const gpt4Rag = createRAGClient(client, {
 const astronomyRag = gpt4Rag.withContext({ query: "Astronomy" });
 
 console.time("gpt-4 Time");
-console.log(await astronomyRag.queryRag({ prompt: "What color is the sky on Mars?" }));
+console.log((await astronomyRag.queryRag({ prompt: "What color is the sky on Mars?" })).content);
 console.timeEnd("gpt-4 Time");
 
 const fastAstronomyRag = astronomyRag.withConfig({
   model: "gpt-4o",
 });
 
 console.time("gpt-4o Time");
-console.log(await fastAstronomyRag.queryRag({ prompt: "What color is the sky on Mars?" }));
+console.log((await fastAstronomyRag.queryRag({ prompt: "What color is the sky on Mars?" })).content);
 console.timeEnd("gpt-4o Time");
 
 const fastChemistryRag = fastAstronomyRag.withContext({ query: "Chemistry" });
 
 console.log(
-  await fastChemistryRag.queryRag({ prompt: "What is the atomic number of gold?" }),
+  (await fastChemistryRag.queryRag({ prompt: "What is the atomic number of gold?" })).content,
 );
 
 // handle the Response object
@@ -74,63 +74,60 @@ console.log(
 
 The `@gel/ai` package supports tool calls, allowing you to extend the capabilities of the AI model with your own functions. Here's how to use them:
 
-1.  **Define the tool:** Create a `SystemMessage` that describes the tool, its parameters, and when it should be used.
-2.  **Send the request:** Send a request to the model using `queryRag`, including the user's prompt and the tool definition.
-3.  **Handle the tool call:** If the model decides to use the tool, it will return an `AssistantMessage` with a `tool_calls` array. Your code needs to:
+1.  **Define your tools**: Create an array of `ToolDefinition` objects that describe your functions, their parameters, and what they do.
+2.  **Send the request**: Call `queryRag` or `streamRag` with the user's prompt and the `tools` array. You can also use the `tool_choice` parameter to control how the model uses your tools.
+3.  **Handle the tool call**: If the model decides to use a tool, it will return an `AssistantMessage` with a `tool_calls` array. Your code needs to:
     a.  Parse the `tool_calls` array to identify the tool and its arguments.
-    b.  Execute the tool.
+    b.  Execute the tool and get the result.
     c.  Create a `ToolMessage` with the result.
-    d.  Send the `ToolMessage` back to the model.
-4.  **Receive the final response:** The model will use the tool's output to generate a final response.
+    d.  Send the `ToolMessage` back to the model in a new request.
+4.  **Receive the final response**: The model will use the tool's output to generate a final response.
 
 ### Example
 
 ```typescript
 import type {
   Message,
-  SystemMessage,
+  ToolDefinition,
   UserMessage,
   ToolMessage,
   AssistantMessage,
 } from "@gel/ai";
 
-// 1. Define the tool in a system message
-const systemMessage: SystemMessage = {
-  role: "system",
-  content: `
-    You have access to a tool called "get_weather" that takes a city as a parameter.
-    Use this tool to answer questions about the weather.
-    The tool definition is:
-    {
-      "name": "get_weather",
-      "description": "Get the current weather for a given city.",
-      "parameters": {
-        "type": "object",
-        "properties": {
-          "city": {
-            "type": "string",
-            "description": "The city to get the weather for."
-          }
+// 1. Define your tools
+const tools: ToolDefinition[] = [
+  {
+    type: "function",
+    name: "get_weather",
+    description: "Get the current weather for a given city.",
+    parameters: {
+      type: "object",
+      properties: {
+        city: {
+          type: "string",
+          description: "The city to get the weather for.",
         },
-        "required": ["city"]
-      }
-    }
-  `,
-};
+      },
+      required: ["city"],
+    },
+  },
+];
 
 // 2. Send the request
 const userMessage: UserMessage = {
   role: "user",
   content: [{ type: "text", text: "What's the weather like in London?" }],
 };
 
-const messages: Message[] = [systemMessage, userMessage];
+const messages: Message[] = [userMessage];
 
 const response = await ragClient.queryRag({
   messages,
+  tools,
+  tool_choice: "auto",
 });
 
-// 3. Handle the tool call (this is a simplified example)
+// 3. Handle the tool call
 if (response.tool_calls) {
   const toolCall = response.tool_calls[0];
   if (toolCall.function.name === "get_weather") {
@@ -150,12 +147,13 @@ if (response.tool_calls) {
     // 4. Send the tool result back to the model
     const finalResponse = await ragClient.queryRag({
       messages,
+      tools,
     });
 
-    console.log(finalResponse.text);
+    console.log(finalResponse.content);
   }
 } else {
-  console.log(response.text);
+  console.log(response.content);
 }
 
 // Dummy function for the example
@@ -170,64 +168,72 @@ When using `streamRag`, you can handle tool calls as they arrive in the stream.
 
 ```typescript
 // Function to handle the streaming response
-async function handleStreamingResponse() {
+async function handleStreamingResponse(initialMessages: Message[]) {
   const stream = ragClient.streamRag({
-    messages,
+    messages: initialMessages,
+    tools,
+    tool_choice: "auto",
   });
 
-  let toolCallId: string | null = null;
-  let functionName: string | null = null;
-  let functionArguments = "";
-  let assistantResponse: AssistantMessage | null = null;
+  let toolCalls: { id: string; name: string; arguments: string }[] = [];
+  let currentToolCall: { id: string; name: string; arguments: string } | null =
+    null;
 
   for await (const chunk of stream) {
     if (
       chunk.type === "content_block_start" &&
       chunk.content_block.type === "tool_use"
     ) {
-      toolCallId = chunk.content_block.id;
-      functionName = chunk.content_block.name;
+      currentToolCall = {
+        id: chunk.content_block.id!,
+        name: chunk.content_block.name,
+        arguments: "",
+      };
     } else if (
       chunk.type === "content_block_delta" &&
-      chunk.delta.type === "tool_call_delta"
+      chunk.delta.type === "input_json_delta"
     ) {
-      functionArguments += chunk.delta.args;
+      if (currentToolCall) {
+        currentToolCall.arguments += chunk.delta.partial_json;
+      }
+    } else if (chunk.type === "content_block_stop") {
+      if (currentToolCall) {
+        toolCalls.push(currentToolCall);
+        currentToolCall = null;
+      }
     } else if (chunk.type === "message_stop") {
       // The model has finished its turn
-      if (functionName && toolCallId) {
-        // We have a tool call to execute
-        const args = JSON.parse(functionArguments);
-        const weather = await getWeather(args.city); // Your function to get the weather
-
-        const toolMessage: ToolMessage = {
-          role: "tool",
-          tool_call_id: toolCallId,
-          content: JSON.stringify({ weather }),
-        };
-
-        // Add the assistant's response and the tool message to the history
-        // A complete assistant message would be constructed from the stream
-        assistantResponse = {
+      if (toolCalls.length > 0) {
+        const assistantMessage: AssistantMessage = {
           role: "assistant",
-          content: "",
-          tool_calls: [
-            {
-              id: toolCallId,
-              type: "function",
-              function: { name: functionName, arguments: functionArguments },
-            },
-          ],
+          content: null,
+          tool_calls: toolCalls.map((tc) => ({
+            id: tc.id,
+            type: "function",
+            function: { name: tc.name, arguments: tc.arguments },
+          })),
         };
-        messages.push(assistantResponse);
-        messages.push(toolMessage);
 
-        // Reset for the next turn
-        toolCallId = null;
-        functionName = null;
-        functionArguments = "";
+        const toolMessages: ToolMessage[] = await Promise.all(
+          toolCalls.map(async (tc) => {
+            const args = JSON.parse(tc.arguments);
+            const weather = await getWeather(args.city); // Your function to get the weather
+            return {
+              role: "tool",
+              tool_call_id: tc.id,
+              content: JSON.stringify({ weather }),
+            };
+          }),
+        );
+
+        const newMessages: Message[] = [
+          ...initialMessages,
+          assistantMessage,
+          ...toolMessages,
+        ];
 
         // Call the function again to get the final response
-        await handleStreamingResponse();
+        await handleStreamingResponse(newMessages);
       }
     } else if (
       chunk.type === "content_block_delta" &&
@@ -239,5 +245,5 @@ async function handleStreamingResponse() {
   }
 }
 
-handleStreamingResponse();
+handleStreamingResponse(messages);
 ```