revert and lock

Author: avivhalfon

packages/traceloop-sdk/test/ai-sdk-integration.test.ts

@@ -35,7 +35,7 @@ Polly.register(NodeHttpAdapter);
 Polly.register(FetchAdapter);
 Polly.register(FSPersister);
 
-describe("Test AI SDK Integration with Recording", function() {
+describe("Test AI SDK Integration with Recording", function () {
   this.timeout(10000); // Increase timeout to 10 seconds for API calls
   setupPolly({
     adapters: ["node-http", "fetch"],
@@ -77,8 +77,10 @@ describe("Test AI SDK Integration with Recording", function() {
     const { server } = this.polly as Polly;
     server.any().on("beforePersist", (_req, recording) => {
       recording.request.headers = recording.request.headers.filter(
-        ({ name }: { name: string }) => 
-          !["authorization", "x-api-key", "x-goog-api-key"].includes(name.toLowerCase())
+        ({ name }: { name: string }) =>
+          !["authorization", "x-api-key", "x-goog-api-key"].includes(
+            name.toLowerCase(),
+          ),
       );
     });
   });
@@ -93,20 +95,24 @@ describe("Test AI SDK Integration with Recording", function() {
       { name: "test_openai_workflow" },
       async () => {
         return await generateText({
-          messages: [{ role: "user", content: "What is 2+2? Give a brief answer." }],
+          messages: [
+            { role: "user", content: "What is 2+2? Give a brief answer." },
+          ],
           model: vercel_openai("gpt-3.5-turbo"),
           experimental_telemetry: { isEnabled: true },
         });
-      }
+      },
     );
 
     // Force flush to ensure all spans are exported
     await traceloop.forceFlush();
-    
+
     const spans = memoryExporter.getFinishedSpans();
-    
+
     const generateTextSpan = spans.find(
-      (span) => span.name === "ai.generateText.generate" || span.name === "ai.generateText.doGenerate",
+      (span) =>
+        span.name === "ai.generateText.generate" ||
+        span.name === "ai.generateText.doGenerate",
     );
 
     assert.ok(result);
@@ -117,10 +123,7 @@ describe("Test AI SDK Integration with Recording", function() {
     assert.strictEqual(generateTextSpan.name, "ai.generateText.generate");
 
     // Verify vendor
-    assert.strictEqual(
-      generateTextSpan.attributes["gen_ai.system"],
-      "OpenAI",
-    );
+    assert.strictEqual(generateTextSpan.attributes["gen_ai.system"], "OpenAI");
 
     // Verify model information
     assert.strictEqual(
@@ -133,9 +136,7 @@ describe("Test AI SDK Integration with Recording", function() {
       generateTextSpan.attributes["gen_ai.prompt.0.role"],
       "user",
     );
-    assert.ok(
-      generateTextSpan.attributes["gen_ai.prompt.0.content"],
-    );
+    assert.ok(generateTextSpan.attributes["gen_ai.prompt.0.content"]);
 
     // Verify response
     assert.strictEqual(
@@ -148,41 +149,39 @@ describe("Test AI SDK Integration with Recording", function() {
     );
 
     // Verify token usage
-    assert.ok(
-      generateTextSpan.attributes["gen_ai.usage.prompt_tokens"],
-    );
-    assert.ok(
-      generateTextSpan.attributes["gen_ai.usage.completion_tokens"],
-    );
-    assert.ok(
-      generateTextSpan.attributes["llm.usage.total_tokens"],
-    );
+    assert.ok(generateTextSpan.attributes["gen_ai.usage.prompt_tokens"]);
+    assert.ok(generateTextSpan.attributes["gen_ai.usage.completion_tokens"]);
+    assert.ok(generateTextSpan.attributes["llm.usage.total_tokens"]);
   });
 
   it("should capture Google Gemini provider spans correctly with recording", async () => {
     // Clear any leftover spans from previous tests
     memoryExporter.reset();
-    
+
     const result = await traceloop.withWorkflow(
       { name: "test_google_workflow" },
       async () => {
         return await generateText({
-          messages: [{ role: "user", content: "What is 2+2? Give a brief answer." }],
+          messages: [
+            { role: "user", content: "What is 2+2? Give a brief answer." },
+          ],
           model: vercel_google("gemini-1.5-flash"),
           experimental_telemetry: { isEnabled: true },
         });
-      }
+      },
     );
 
     // Force flush to ensure all spans are exported
     await traceloop.forceFlush();
-    
+
     const spans = memoryExporter.getFinishedSpans();
-    
+
     // Find the Google span specifically (should have workflow name test_google_workflow)
     const generateTextSpan = spans.find(
-      (span) => (span.name === "ai.generateText.generate" || span.name === "ai.generateText.doGenerate") 
-        && span.attributes["traceloop.workflow.name"] === "test_google_workflow",
+      (span) =>
+        (span.name === "ai.generateText.generate" ||
+          span.name === "ai.generateText.doGenerate") &&
+        span.attributes["traceloop.workflow.name"] === "test_google_workflow",
     );
 
     assert.ok(result);
@@ -193,10 +192,7 @@ describe("Test AI SDK Integration with Recording", function() {
     assert.strictEqual(generateTextSpan.name, "ai.generateText.generate");
 
     // Verify vendor
-    assert.strictEqual(
-      generateTextSpan.attributes["gen_ai.system"],
-      "Google",
-    );
+    assert.strictEqual(generateTextSpan.attributes["gen_ai.system"], "Google");
 
     // Verify model information
     assert.strictEqual(
@@ -209,9 +205,7 @@ describe("Test AI SDK Integration with Recording", function() {
       generateTextSpan.attributes["gen_ai.prompt.0.role"],
       "user",
     );
-    assert.ok(
-      generateTextSpan.attributes["gen_ai.prompt.0.content"],
-    );
+    assert.ok(generateTextSpan.attributes["gen_ai.prompt.0.content"]);
 
     // Verify response
     assert.strictEqual(
@@ -224,15 +218,8 @@ describe("Test AI SDK Integration with Recording", function() {
     );
 
     // Verify token usage
-    assert.ok(
-      generateTextSpan.attributes["gen_ai.usage.prompt_tokens"],
-    );
-    assert.ok(
-      generateTextSpan.attributes["gen_ai.usage.completion_tokens"],
-    );
-    assert.ok(
-      generateTextSpan.attributes["llm.usage.total_tokens"],
-    );
+    assert.ok(generateTextSpan.attributes["gen_ai.usage.prompt_tokens"]);
+    assert.ok(generateTextSpan.attributes["gen_ai.usage.completion_tokens"]);
+    assert.ok(generateTextSpan.attributes["llm.usage.total_tokens"]);
   });
-
-});
+});