@sentrial/sdk 0.3.3 → 0.4.0

package/dist/index.js

@@ -1073,8 +1073,13 @@ function wrapStreamText(originalFn, client, config) {
         return null;
       }
     })();
+    const wrappedParams = {
+      ...params,
+      tools: params.tools ? wrapToolsAsync(params.tools, sessionPromise, client) : void 0
+    };
+    const result = originalFn(wrappedParams);
     let tracked = false;
-    async function trackCompletion(fullText, usageInfo, error) {
+    async function trackCompletion(fullText, error) {
       if (tracked) return;
       tracked = true;
       const durationMs = Date.now() - startTime;
@@ -1094,10 +1099,26 @@ function wrapStreamText(originalFn, client, config) {
         });
         return;
       }
-      const promptTokens = usageInfo?.promptTokens || 0;
-      const completionTokens = usageInfo?.completionTokens || 0;
-      const totalTokens = usageInfo?.totalTokens || promptTokens + completionTokens;
+      let usage;
+      try {
+        usage = result.usage ? await result.usage : void 0;
+      } catch {
+      }
+      const promptTokens = usage?.promptTokens || 0;
+      const completionTokens = usage?.completionTokens || 0;
+      const totalTokens = usage?.totalTokens || promptTokens + completionTokens;
       const cost = calculateCostForCall(provider, modelId, promptTokens, completionTokens);
+      await client.trackEvent({
+        sessionId: sid,
+        eventType: "llm_call",
+        eventData: {
+          model: modelId,
+          provider,
+          prompt_tokens: promptTokens,
+          completion_tokens: completionTokens,
+          total_tokens: totalTokens
+        }
+      });
       await client.completeSession({
         sessionId: sid,
         success: true,
@@ -1109,28 +1130,39 @@ function wrapStreamText(originalFn, client, config) {
         totalTokens
       });
     }
-    const userOnFinish = params.onFinish;
-    const wrappedParams = {
-      ...params,
-      tools: params.tools ? wrapToolsAsync(params.tools, sessionPromise, client) : void 0,
-      onFinish: async (event) => {
-        try {
-          if (userOnFinish) await userOnFinish(event);
-        } catch {
-        }
-        await trackCompletion(event.text, event.usage);
-      }
-    };
-    const result = originalFn(wrappedParams);
     const textProp = result.text;
-    if (textProp != null && typeof textProp.then === "function") {
-      textProp.then((text) => {
+    if (typeof textProp === "string") {
+      trackCompletion(textProp).catch(() => {
+      });
+    } else if (textProp != null && typeof textProp.then === "function") {
+      const originalTextPromise = textProp;
+      result.text = originalTextPromise.then((text) => {
         trackCompletion(text).catch(() => {
         });
+        return text;
       }).catch((err) => {
-        trackCompletion("", void 0, err instanceof Error ? err : new Error(String(err))).catch(() => {
+        trackCompletion("", err instanceof Error ? err : new Error(String(err))).catch(() => {
         });
+        throw err;
       });
+    } else {
+      const originalTextStream = result.textStream;
+      let fullText = "";
+      result.textStream = (async function* () {
+        try {
+          for await (const chunk of originalTextStream) {
+            fullText += chunk;
+            yield chunk;
+          }
+          await trackCompletion(fullText);
+        } catch (error) {
+          await trackCompletion(
+            "",
+            error instanceof Error ? error : new Error(String(error))
+          );
+          throw error;
+        }
+      })();
     }
     return result;
   };
@@ -1218,52 +1250,54 @@ function wrapStreamObject(originalFn, client, config) {
         return null;
       }
     })();
-    let tracked = false;
-    const userOnFinish = params.onFinish;
-    const wrappedParams = {
-      ...params,
-      onFinish: async (event) => {
-        try {
-          if (userOnFinish) await userOnFinish(event);
-        } catch {
+    const result = originalFn(params);
+    if (result.object) {
+      const originalObjectPromise = result.object;
+      result.object = originalObjectPromise.then(async (obj) => {
+        const durationMs = Date.now() - startTime;
+        const sid = await sessionPromise;
+        if (sid) {
+          let usage;
+          try {
+            usage = result.usage ? await result.usage : void 0;
+          } catch {
+          }
+          const promptTokens = usage?.promptTokens || 0;
+          const completionTokens = usage?.completionTokens || 0;
+          const totalTokens = usage?.totalTokens || promptTokens + completionTokens;
+          const cost = calculateCostForCall(provider, modelId, promptTokens, completionTokens);
+          await client.completeSession({
+            sessionId: sid,
+            success: true,
+            output: JSON.stringify(obj),
+            durationMs,
+            estimatedCost: cost,
+            promptTokens,
+            completionTokens,
+            totalTokens
+          });
         }
-        if (tracked) return;
-        tracked = true;
+        return obj;
+      }).catch(async (error) => {
         const durationMs = Date.now() - startTime;
         const sid = await sessionPromise;
-        if (!sid) return;
-        if (event.error) {
-          const err = event.error instanceof Error ? event.error : new Error(String(event.error));
+        if (sid) {
           await client.trackError({
             sessionId: sid,
-            errorType: err.name || "Error",
-            errorMessage: err.message || "Unknown error"
+            errorType: error instanceof Error ? error.name : "Error",
+            errorMessage: error instanceof Error ? error.message : "Unknown error"
           });
           await client.completeSession({
             sessionId: sid,
             success: false,
-            failureReason: err.message || "Unknown error",
+            failureReason: error instanceof Error ? error.message : "Unknown error",
             durationMs
           });
-          return;
         }
-        const promptTokens = event.usage?.promptTokens || 0;
-        const completionTokens = event.usage?.completionTokens || 0;
-        const totalTokens = event.usage?.totalTokens || promptTokens + completionTokens;
-        const cost = calculateCostForCall(provider, modelId, promptTokens, completionTokens);
-        await client.completeSession({
-          sessionId: sid,
-          success: true,
-          output: event.object != null ? JSON.stringify(event.object) : "",
-          durationMs,
-          estimatedCost: cost,
-          promptTokens,
-          completionTokens,
-          totalTokens
-        });
-      }
-    };
-    return originalFn(wrappedParams);
+        throw error;
+      });
+    }
+    return result;
   };
 }
 function wrapAISDK(ai, options) {