Braintrust traced returns promise instead of raw object, requires await

packages/fern-docs/bundle/src/app/api/fern-docs/search/v2/chat/route.ts

@@ -11,6 +11,7 @@ import {
   queryTurbopuffer,
   toDocuments,
 } from "@fern-docs/search-server/turbopuffer";
+import { initLogger, traced, wrapAISDKModel } from "braintrust";
 import { COOKIE_FERN_TOKEN, withoutStaging } from "@fern-docs/utils";
 import { embed, EmbeddingModel, streamText, tool } from "ai";
 import { cookies } from "next/headers";
@@ -78,18 +79,33 @@ export async function POST(req: NextRequest) {
     date: new Date().toDateString(),
     documents,
   });
-  const result = streamText({
-    model: languageModel,
-    system,
-    messages,
-    maxSteps: 10,
-    maxRetries: 3,
-    tools: {
-      search: tool({
-        description:
-          "Search the knowledge base for the user's query. Semantic search is enabled.",
-        parameters: z.object({
-          query: z.string(),
+  const result = await traced(() =>
+    streamText({
+      model: languageModel,
+      system,
+      messages,
+      maxSteps: 10,
+      maxRetries: 3,
+      tools: {
+        search: tool({
+          description:
+            "Search the knowledge base for the user's query. Semantic search is enabled.",
+          parameters: z.object({
+            query: z.string(),
+          }),
+          async execute({ query }) {
+            const response = await runQueryTurbopuffer(query, {
+              embeddingModel,
+              namespace,
+              authed: user != null,
+              roles: user?.roles ?? [],
+            });
+            return response.map((hit) => {
+              const { domain, pathname, hash } = hit.attributes;
+              const url = `https://${domain}${pathname}${hash ?? ""}`;
+              return { url, ...hit.attributes };
+            });
+          },
         }),
         async execute({ query }) {
           const response = await runQueryTurbopuffer(query, {