@xsai-ext/telemetry 0.4.4 → 0.5.0-beta.2

package/dist/index.js

@@ -1,8 +1,9 @@
-import { embed as embed$1, clean, embedMany as embedMany$1, trampoline, chat, responseJSON, determineStepType, executeTool, DelayedPromise, objCamelToSnake } from 'xsai';
+import { embed as embed$1, clean, embedMany as embedMany$1, trampoline, resolveStepOptions, chat, responseJSON, InvalidResponseError, stepCountAtLeast, executeTool, shouldStop, determineStepType, DelayedPromise, objCamelToSnake } from 'xsai';
 export * from 'xsai';
 import { trace, SpanStatusCode } from '@opentelemetry/api';
+import { createControlledStream, errorControllers, closeControllers, EventSourceParserStream, JsonMessageTransformStream } from '@xsai/shared-stream';
 
-var version = "0.4.4";
+var version = "0.5.0-beta.2";
 var pkg = {
 	version: version};
 
@@ -131,127 +132,123 @@ const wrapTool = (tool, tracer) => ({
 
 const generateText = async (options) => {
   const tracer = getTracer();
-  const rawGenerateText = async (options2) => recordSpan(chatSpan(options2, tracer), async (span) => chat({
-    ...options2,
-    maxSteps: void 0,
-    steps: void 0,
-    stream: false
-  }).then(responseJSON).then(async (res) => {
-    const { choices, usage } = res;
-    if (!choices?.length)
-      throw new Error(`No choices returned, response body: ${JSON.stringify(res)}`);
-    const messages = structuredClone(options2.messages);
-    const steps = options2.steps ? structuredClone(options2.steps) : [];
-    const toolCalls = [];
-    const toolResults = [];
-    const { finish_reason: finishReason, message } = choices[0];
-    const msgToolCalls = message?.tool_calls ?? [];
-    const stepType = determineStepType({
-      finishReason,
-      maxSteps: options2.maxSteps ?? 1,
-      stepsLength: steps.length,
-      toolCallsLength: msgToolCalls.length
+  const rawGenerateText = async (options2) => {
+    const messages = options2.steps == null ? structuredClone(options2.messages) : options2.messages;
+    const steps = options2.steps ?? [];
+    const stepOptions = await resolveStepOptions({
+      messages,
+      model: options2.model,
+      prepareStep: options2.prepareStep,
+      stepNumber: steps.length,
+      steps,
+      toolChoice: options2.toolChoice
     });
-    messages.push(message);
-    span.setAttribute("gen_ai.output.messages", JSON.stringify([message]));
-    if (finishReason !== "stop" && stepType !== "done") {
-      for (const toolCall of msgToolCalls) {
-        const { completionToolCall, completionToolResult, message: message2 } = await executeTool({
-          abortSignal: options2.abortSignal,
-          messages,
-          toolCall,
-          tools: options2.tools
+    return recordSpan(chatSpan({
+      ...options2,
+      messages: stepOptions.messages,
+      model: stepOptions.model,
+      toolChoice: stepOptions.toolChoice
+    }, tracer), async (span) => chat({
+      ...options2,
+      maxSteps: void 0,
+      messages: stepOptions.messages,
+      model: stepOptions.model,
+      steps: void 0,
+      stopWhen: void 0,
+      stream: false,
+      toolChoice: stepOptions.toolChoice
+    }).then(responseJSON).then(async (res) => {
+      const { choices, usage } = res;
+      if (!choices?.length) {
+        const responseBody = JSON.stringify(res);
+        throw new InvalidResponseError(`No choices returned, response body: ${responseBody}`, {
+          reason: "no_choices",
+          responseBody
         });
-        toolCalls.push(completionToolCall);
-        toolResults.push(completionToolResult);
-        messages.push(message2);
       }
-    }
-    const step = {
-      finishReason,
-      stepType,
-      text: Array.isArray(message.content) ? message.content.filter((m) => m.type === "text").map((m) => m.text).join("\n") : message.content,
-      toolCalls,
-      toolResults,
-      usage
-    };
-    steps.push(step);
-    span.setAttributes({
-      "gen_ai.response.finish_reasons": [step.finishReason],
-      "gen_ai.usage.input_tokens": step.usage.prompt_tokens,
-      "gen_ai.usage.output_tokens": step.usage.completion_tokens
-    });
-    if (options2.onStepFinish)
-      await options2.onStepFinish(step);
-    if (step.finishReason === "stop" || step.stepType === "done") {
-      return {
-        finishReason: step.finishReason,
-        messages,
-        reasoningText: message.reasoning_content,
-        steps,
-        text: step.text,
-        toolCalls: step.toolCalls,
-        toolResults: step.toolResults,
-        usage: step.usage
+      const toolCalls = [];
+      const toolResults = [];
+      const { finish_reason: finishReason, message } = choices[0];
+      const msgToolCalls = message?.tool_calls ?? [];
+      const stopWhen = options2.stopWhen ?? stepCountAtLeast(1);
+      messages.push(message);
+      span.setAttribute("gen_ai.output.messages", JSON.stringify([message]));
+      if (msgToolCalls.length > 0) {
+        const results = await Promise.all(
+          msgToolCalls.map(async (toolCall) => executeTool({
+            abortSignal: options2.abortSignal,
+            messages,
+            toolCall,
+            tools: options2.tools
+          }))
+        );
+        for (const { completionToolCall, completionToolResult, message: message2 } of results) {
+          toolCalls.push(completionToolCall);
+          toolResults.push(completionToolResult);
+          messages.push(message2);
+        }
+      }
+      const stopStep = {
+        finishReason,
+        text: Array.isArray(message.content) ? message.content.filter((m) => m.type === "text").map((m) => m.text).join("\n") : message.content,
+        toolCalls,
+        toolResults,
+        usage
       };
-    } else {
-      return async () => rawGenerateText({
-        ...options2,
+      const stop = shouldStop(stopWhen, {
         messages,
-        steps
+        step: stopStep,
+        steps: [...steps, stopStep]
       });
-    }
-  }));
+      const willContinue = toolCalls.length > 0 && !stop;
+      const step = {
+        ...stopStep,
+        stepType: determineStepType({
+          finishReason,
+          stepsLength: steps.length,
+          toolCallsLength: toolCalls.length,
+          willContinue
+        })
+      };
+      steps.push(step);
+      span.setAttributes({
+        "gen_ai.response.finish_reasons": [step.finishReason],
+        "gen_ai.usage.input_tokens": step.usage.prompt_tokens,
+        "gen_ai.usage.output_tokens": step.usage.completion_tokens
+      });
+      if (options2.onStepFinish)
+        await options2.onStepFinish(step);
+      if (!willContinue) {
+        return {
+          finishReason: step.finishReason,
+          messages,
+          reasoningText: message.reasoning ?? message.reasoning_content,
+          steps,
+          text: step.text,
+          toolCalls: step.toolCalls,
+          toolResults: step.toolResults,
+          usage: step.usage
+        };
+      } else {
+        return async () => rawGenerateText({
+          ...options2,
+          messages,
+          steps
+        });
+      }
+    }));
+  };
   return trampoline(async () => rawGenerateText({
     ...options,
     tools: options.tools?.map((tool) => wrapTool(tool, tracer))
   }));
 };
 
-const parseChunk = (text) => {
-  if (!text || !text.startsWith("data:"))
-    return [void 0, false];
-  const content = text.slice("data:".length);
-  const data = content.startsWith(" ") ? content.slice(1) : content;
-  if (data === "[DONE]") {
-    return [void 0, true];
-  }
-  if (data.startsWith("{") && data.includes('"error":')) {
-    throw new Error(`Error from server: ${data}`);
-  }
-  const chunk = JSON.parse(data);
-  return [chunk, false];
-};
-const transformChunk = () => {
-  const decoder = new TextDecoder();
-  let buffer = "";
-  return new TransformStream({
-    transform: async (chunk, controller) => {
-      const text = decoder.decode(chunk, { stream: true });
-      buffer += text;
-      const lines = buffer.split("\n");
-      buffer = lines.pop() ?? "";
-      for (const line of lines) {
-        try {
-          const [chunk2, isEnd] = parseChunk(line);
-          if (isEnd)
-            break;
-          if (chunk2) {
-            controller.enqueue(chunk2);
-          }
-        } catch (error) {
-          controller.error(error);
-        }
-      }
-    }
-  });
-};
-
 const streamText = (options) => {
   const tracer = getTracer();
   const steps = [];
   const messages = structuredClone(options.messages);
-  const maxSteps = options.maxSteps ?? 1;
+  const stopWhen = options.stopWhen ?? stepCountAtLeast(1);
   let usage;
   let totalUsage;
   let reasoningField;
@@ -259,14 +256,11 @@ const streamText = (options) => {
   const resultMessages = new DelayedPromise();
   const resultUsage = new DelayedPromise();
   const resultTotalUsage = new DelayedPromise();
-  let eventCtrl;
-  let textCtrl;
-  let reasoningTextCtrl;
-  const eventStream = new ReadableStream({ start: (controller) => eventCtrl = controller });
-  const textStream = new ReadableStream({ start: (controller) => textCtrl = controller });
-  const reasoningTextStream = new ReadableStream({ start: (controller) => reasoningTextCtrl = controller });
+  const [eventStream, eventCtrl] = createControlledStream();
+  const [textStream, textCtrl] = createControlledStream();
+  const [reasoningTextStream, reasoningTextCtrl] = createControlledStream();
   const pushEvent = (stepEvent) => {
-    eventCtrl?.enqueue(stepEvent);
+    eventCtrl.current?.enqueue(stepEvent);
     void options.onEvent?.(stepEvent);
   };
   const pushStep = (step) => {
@@ -274,176 +268,211 @@ const streamText = (options) => {
     void options.onStepFinish?.(step);
   };
   const tools = options.tools != null && options.tools.length > 0 ? options.tools.map((tool) => wrapTool(tool, tracer)) : void 0;
-  const doStream = async () => recordSpan(chatSpan({ ...options, messages }, tracer), async (span) => {
-    const { body: stream } = await chat({
-      ...options,
-      maxSteps: void 0,
+  const doStream = async () => {
+    const stepOptions = await resolveStepOptions({
       messages,
-      stream: true,
-      streamOptions: options.streamOptions != null ? objCamelToSnake(options.streamOptions) : void 0,
-      tools
+      model: options.model,
+      prepareStep: options.prepareStep,
+      stepNumber: steps.length,
+      steps,
+      toolChoice: options.toolChoice
     });
-    const pushUsage = (u) => {
-      usage = u;
-      totalUsage = totalUsage ? {
-        completion_tokens: totalUsage.completion_tokens + u.completion_tokens,
-        prompt_tokens: totalUsage.prompt_tokens + u.prompt_tokens,
-        total_tokens: totalUsage.total_tokens + u.total_tokens
-      } : u;
-    };
-    let text = "";
-    let reasoningText;
-    const pushText = (content) => {
-      textCtrl?.enqueue(content);
-      text += content;
-    };
-    const pushReasoningText = (reasoningContent) => {
-      if (reasoningText == null)
-        reasoningText = "";
-      reasoningTextCtrl?.enqueue(reasoningContent);
-      reasoningText += reasoningContent;
-    };
-    const tool_calls = [];
-    const toolCalls = [];
-    const toolResults = [];
-    let finishReason = "other";
-    await stream.pipeThrough(transformChunk()).pipeTo(new WritableStream({
-      abort: (reason) => {
-        eventCtrl?.error(reason);
-        textCtrl?.error(reason);
-      },
-      close: () => {
-      },
-      // eslint-disable-next-line sonarjs/cognitive-complexity
-      write: (chunk) => {
-        if (chunk.usage)
-          pushUsage(chunk.usage);
-        if (chunk.choices == null || chunk.choices.length === 0)
-          return;
-        const choice = chunk.choices[0];
-        if (choice.delta.reasoning != null) {
-          if (reasoningField !== "reasoning")
-            reasoningField = "reasoning";
-          pushEvent({ text: choice.delta.reasoning, type: "reasoning-delta" });
-          pushReasoningText(choice.delta.reasoning);
-        } else if (choice.delta.reasoning_content != null) {
-          if (reasoningField !== "reasoning_content")
-            reasoningField = "reasoning_content";
-          pushEvent({ text: choice.delta.reasoning_content, type: "reasoning-delta" });
-          pushReasoningText(choice.delta.reasoning_content);
-        }
-        if (choice.finish_reason != null)
-          finishReason = choice.finish_reason;
-        if (choice.delta.tool_calls?.length === 0 || choice.delta.tool_calls == null) {
-          if (choice.delta.content != null) {
-            pushEvent({ text: choice.delta.content, type: "text-delta" });
-            pushText(choice.delta.content);
-          } else if (choice.delta.refusal != null) {
-            pushEvent({ error: choice.delta.refusal, type: "error" });
-          } else if (choice.finish_reason != null) {
-            pushEvent({ finishReason: choice.finish_reason, type: "finish", usage });
+    return recordSpan(chatSpan({
+      ...options,
+      messages: stepOptions.messages,
+      model: stepOptions.model,
+      toolChoice: stepOptions.toolChoice
+    }, tracer), async (span) => {
+      const { body: stream } = await chat({
+        ...options,
+        maxSteps: void 0,
+        messages: stepOptions.messages,
+        model: stepOptions.model,
+        stopWhen: void 0,
+        stream: true,
+        streamOptions: options.streamOptions != null ? objCamelToSnake(options.streamOptions) : void 0,
+        toolChoice: stepOptions.toolChoice,
+        tools
+      });
+      const pushUsage = (u) => {
+        usage = u;
+        totalUsage = totalUsage ? {
+          completion_tokens: totalUsage.completion_tokens + u.completion_tokens,
+          prompt_tokens: totalUsage.prompt_tokens + u.prompt_tokens,
+          total_tokens: totalUsage.total_tokens + u.total_tokens
+        } : u;
+      };
+      let text = "";
+      let reasoningText;
+      const pushText = (content) => {
+        textCtrl.current?.enqueue(content);
+        text += content;
+      };
+      const pushReasoningText = (reasoningContent) => {
+        if (reasoningText == null)
+          reasoningText = "";
+        reasoningTextCtrl.current?.enqueue(reasoningContent);
+        reasoningText += reasoningContent;
+      };
+      const tool_calls = [];
+      const toolCalls = [];
+      const toolResults = [];
+      let finishReason = "other";
+      await stream.pipeThrough(new TextDecoderStream()).pipeThrough(new EventSourceParserStream()).pipeThrough(new JsonMessageTransformStream()).pipeTo(new WritableStream({
+        abort: (reason) => {
+          errorControllers(reason, eventCtrl, textCtrl, reasoningTextCtrl);
+        },
+        close: () => {
+        },
+        // eslint-disable-next-line sonarjs/cognitive-complexity
+        write: (chunk) => {
+          if (chunk.usage)
+            pushUsage(chunk.usage);
+          if (chunk.choices == null || chunk.choices.length === 0)
+            return;
+          const choice = chunk.choices[0];
+          if (choice.delta.reasoning != null) {
+            if (reasoningField !== "reasoning")
+              reasoningField = "reasoning";
+            pushEvent({ text: choice.delta.reasoning, type: "reasoning-delta" });
+            pushReasoningText(choice.delta.reasoning);
+          } else if (choice.delta.reasoning_content != null) {
+            if (reasoningField !== "reasoning_content")
+              reasoningField = "reasoning_content";
+            pushEvent({ text: choice.delta.reasoning_content, type: "reasoning-delta" });
+            pushReasoningText(choice.delta.reasoning_content);
           }
-        } else {
-          for (const toolCall of choice.delta.tool_calls) {
-            const { index } = toolCall;
-            if (!tool_calls.at(index)) {
-              tool_calls[index] = {
-                ...toolCall,
-                function: {
-                  ...toolCall.function,
-                  arguments: toolCall.function.arguments ?? ""
-                }
-              };
-              pushEvent({
-                toolCallId: toolCall.id,
-                toolName: toolCall.function.name,
-                type: "tool-call-streaming-start"
-              });
-            } else {
-              tool_calls[index].function.arguments += toolCall.function.arguments;
-              pushEvent({
-                argsTextDelta: toolCall.function.arguments,
-                toolCallId: toolCall.id,
-                toolName: toolCall.function.name ?? tool_calls[index].function.name,
-                type: "tool-call-delta"
-              });
+          if (choice.finish_reason != null)
+            finishReason = choice.finish_reason;
+          if (choice.delta.tool_calls?.length === 0 || choice.delta.tool_calls == null) {
+            if (choice.delta.content != null) {
+              pushEvent({ text: choice.delta.content, type: "text-delta" });
+              pushText(choice.delta.content);
+            } else if (choice.delta.refusal != null) {
+              pushEvent({ error: choice.delta.refusal, type: "error" });
+            } else if (choice.finish_reason != null) {
+              pushEvent({ finishReason: choice.finish_reason, type: "finish", usage });
+            }
+          } else {
+            for (const toolCall of choice.delta.tool_calls) {
+              const { index } = toolCall;
+              if (!tool_calls.at(index)) {
+                tool_calls[index] = {
+                  ...toolCall,
+                  function: {
+                    ...toolCall.function,
+                    arguments: toolCall.function.arguments ?? ""
+                  }
+                };
+                pushEvent({
+                  toolCallId: toolCall.id,
+                  toolName: toolCall.function.name,
+                  type: "tool-call-streaming-start"
+                });
+              } else {
+                tool_calls[index].function.arguments += toolCall.function.arguments;
+                pushEvent({
+                  argsTextDelta: toolCall.function.arguments,
+                  toolCallId: toolCall.id,
+                  toolName: toolCall.function.name ?? tool_calls[index].function.name,
+                  type: "tool-call-delta"
+                });
+              }
             }
           }
         }
-      }
-    }));
-    const message = {
-      ...reasoningField != null ? { [reasoningField]: reasoningText } : {},
-      content: text,
-      role: "assistant",
-      tool_calls: tool_calls.length > 0 ? tool_calls : void 0
-    };
-    messages.push(message);
-    span.setAttribute("gen_ai.output.messages", JSON.stringify([message]));
-    if (tool_calls.length !== 0) {
-      for (const toolCall of tool_calls) {
-        if (toolCall == null)
-          continue;
-        const { completionToolCall, completionToolResult, message: message2 } = await executeTool({
-          abortSignal: options.abortSignal,
-          messages,
-          toolCall,
-          tools
+      }));
+      const message = {
+        ...reasoningField != null ? { [reasoningField]: reasoningText } : {},
+        content: text,
+        role: "assistant",
+        tool_calls: tool_calls.length > 0 ? tool_calls : void 0
+      };
+      messages.push(message);
+      span.setAttribute("gen_ai.output.messages", JSON.stringify([message]));
+      if (tool_calls.length !== 0) {
+        const validToolCalls = tool_calls.filter((tc) => tc != null);
+        const results = await Promise.all(
+          validToolCalls.map(async (toolCall) => executeTool({
+            abortSignal: options.abortSignal,
+            messages,
+            toolCall,
+            tools
+          }))
+        );
+        for (const { completionToolCall, completionToolResult, message: message2 } of results) {
+          toolCalls.push(completionToolCall);
+          toolResults.push(completionToolResult);
+          messages.push(message2);
+          pushEvent({ ...completionToolCall, type: "tool-call" });
+          pushEvent({ ...completionToolResult, type: "tool-result" });
+        }
+      } else {
+        pushEvent({
+          finishReason,
+          type: "finish",
+          usage
         });
-        toolCalls.push(completionToolCall);
-        toolResults.push(completionToolResult);
-        messages.push(message2);
-        pushEvent({ ...completionToolCall, type: "tool-call" });
-        pushEvent({ ...completionToolResult, type: "tool-result" });
       }
-    } else {
-      pushEvent({
+      const stopStep = {
         finishReason,
-        type: "finish",
+        text,
+        toolCalls,
+        toolResults,
         usage
+      };
+      const stop = shouldStop(stopWhen, {
+        messages,
+        step: stopStep,
+        steps: [...steps, stopStep]
       });
-    }
-    const step = {
-      finishReason,
-      stepType: determineStepType({ finishReason, maxSteps, stepsLength: steps.length, toolCallsLength: toolCalls.length }),
-      text,
-      toolCalls,
-      toolResults,
-      usage
-    };
-    pushStep(step);
-    span.setAttributes({
-      "gen_ai.response.finish_reasons": [step.finishReason],
-      ...step.usage && {
-        "gen_ai.usage.input_tokens": step.usage.prompt_tokens,
-        "gen_ai.usage.output_tokens": step.usage.completion_tokens
-      }
+      const willContinue = toolCalls.length > 0 && !stop;
+      const step = {
+        ...stopStep,
+        stepType: determineStepType({
+          finishReason,
+          stepsLength: steps.length,
+          toolCallsLength: toolCalls.length,
+          willContinue
+        })
+      };
+      pushStep(step);
+      span.setAttributes({
+        "gen_ai.response.finish_reasons": [step.finishReason],
+        ...step.usage && {
+          "gen_ai.usage.input_tokens": step.usage.prompt_tokens,
+          "gen_ai.usage.output_tokens": step.usage.completion_tokens
+        }
+      });
+      if (willContinue)
+        return async () => doStream();
     });
-    if (toolCalls.length !== 0 && steps.length < maxSteps)
-      return async () => doStream();
-  });
+  };
   void (async () => {
+    let finalError;
     try {
       await trampoline(async () => doStream());
-      eventCtrl?.close();
-      textCtrl?.close();
-      reasoningTextCtrl?.close();
     } catch (err) {
-      eventCtrl?.error(err);
-      textCtrl?.error(err);
-      reasoningTextCtrl?.error(err);
-      resultSteps.reject(err);
-      resultMessages.reject(err);
-      resultUsage.reject(err);
-      resultTotalUsage.reject(err);
-    } finally {
-      resultSteps.resolve(steps);
-      resultMessages.resolve(messages);
-      resultUsage.resolve(usage);
-      resultTotalUsage.resolve(totalUsage);
-      const finishStep = steps.at(-1);
-      void options.onFinish?.(finishStep);
+      finalError = err;
+    }
+    try {
+      await options.onFinish?.(steps.at(-1));
+    } catch (err) {
+      finalError ??= err;
+    }
+    if (finalError != null) {
+      errorControllers(finalError, eventCtrl, textCtrl, reasoningTextCtrl);
+      resultSteps.reject(finalError);
+      resultMessages.reject(finalError);
+      resultUsage.reject(finalError);
+      resultTotalUsage.reject(finalError);
+      return;
     }
+    closeControllers(eventCtrl, textCtrl, reasoningTextCtrl);
+    resultSteps.resolve(steps);
+    resultMessages.resolve(messages);
+    resultUsage.resolve(usage);
+    resultTotalUsage.resolve(totalUsage);
   })();
   return {
     fullStream: eventStream,

package/package.json

@@ -1,7 +1,7 @@
 {
   "name": "@xsai-ext/telemetry",
   "type": "module",
-  "version": "0.4.4",
+  "version": "0.5.0-beta.2",
   "description": "extra-small AI SDK.",
   "author": "Moeru AI",
   "license": "MIT",
@@ -30,7 +30,8 @@
   ],
   "dependencies": {
     "@opentelemetry/api": "^1.9.0",
-    "xsai": "~0.4.4"
+    "@xsai/shared-stream": "~0.5.0-beta.2",
+    "xsai": "~0.5.0-beta.2"
   },
   "devDependencies": {
     "@langfuse/otel": "^4.5.1",