Reverting braintrust implementation, will spin off a new branch (#2067)

fern-api · Jan 24, 2025 · 1016465 · 1016465
1 parent 9132dc6
commit 1016465
Show file tree

Hide file tree

Showing 4 changed files with 49 additions and 521 deletions.
diff --git a/packages/fern-docs/bundle/package.json b/packages/fern-docs/bundle/package.json
@@ -62,7 +62,6 @@
     "@workos-inc/node": "^7.31.0",
     "ai": "^4.0.18",
     "algoliasearch": "^5.13.0",
-    "braintrust": "^0.0.182",
     "cssnano": "^6.0.3",
     "es-toolkit": "^1.27.0",
     "esbuild": "0.20.2",

diff --git a/packages/fern-docs/bundle/src/app/api/fern-docs/search/v2/chat/route.ts b/packages/fern-docs/bundle/src/app/api/fern-docs/search/v2/chat/route.ts
@@ -13,7 +13,6 @@ import {
 } from "@fern-docs/search-server/turbopuffer";
 import { COOKIE_FERN_TOKEN, withoutStaging } from "@fern-docs/utils";
 import { embed, EmbeddingModel, streamText, tool } from "ai";
-import { initLogger, traced, wrapAISDKModel } from "braintrust";
 import { cookies } from "next/headers";
 import { NextRequest, NextResponse } from "next/server";
 import { z } from "zod";
@@ -22,18 +21,12 @@ export const maxDuration = 60;
 export const revalidate = 0;
 
 export async function POST(req: NextRequest) {
-  const _logger = initLogger({
-    projectName: "Braintrust Evaluation",
-    apiKey: process.env.BRAINTRUST_API_KEY,
-  });
   const bedrock = createAmazonBedrock({
     region: "us-west-2",
     accessKeyId: process.env.AWS_ACCESS_KEY_ID,
     secretAccessKey: process.env.AWS_SECRET_ACCESS_KEY,
   });
-  const languageModel = wrapAISDKModel(
-    bedrock("anthropic.claude-3-5-sonnet-20241022-v2:0")
-  );
+  const languageModel = bedrock("anthropic.claude-3-5-sonnet-20241022-v2:0");
 
   const openai = createOpenAI({ apiKey: openaiApiKey() });
   const embeddingModel = openai.embedding("text-embedding-3-small");
@@ -85,53 +78,51 @@ export async function POST(req: NextRequest) {
     date: new Date().toDateString(),
     documents,
   });
-  const result = 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 };
-            });
-          },
+  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(),
         }),
-      },
-      onFinish: async (e) => {
-        const end = Date.now();
-        await track("ask_ai", {
-          languageModel: languageModel.modelId,
-          embeddingModel: embeddingModel.modelId,
-          durationMs: end - start,
-          domain,
-          namespace,
-          numToolCalls: e.toolCalls.length,
-          finishReason: e.finishReason,
-          ...e.usage,
-        });
-        e.warnings?.forEach((warning) => {
-          console.warn(warning);
-        });
-      },
-    })
-  );
+        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 };
+          });
+        },
+      }),
+    },
+    onFinish: async (e) => {
+      const end = Date.now();
+      await track("ask_ai", {
+        languageModel: languageModel.modelId,
+        embeddingModel: embeddingModel.modelId,
+        durationMs: end - start,
+        domain,
+        namespace,
+        numToolCalls: e.toolCalls.length,
+        finishReason: e.finishReason,
+        ...e.usage,
+      });
+      e.warnings?.forEach((warning) => {
+        console.warn(warning);
+      });
+    },
+  });
 
   const response = result.toDataStreamResponse();
   response.headers.set("Access-Control-Allow-Origin", "*");

diff --git a/packages/fern-docs/bundle/src/server/loadDocsDefinitionFromS3.ts b/packages/fern-docs/bundle/src/server/loadDocsDefinitionFromS3.ts
@@ -28,6 +28,8 @@ export async function loadDocsDefinitionFromS3({
       dateLessThan: new Date(Date.now() + 1000 * 60 * 60 * 24 * 30).toString(),
     });
     const response = await fetch(signedUrl);
+    console.log("signedUrl", signedUrl);
+    console.log("response.ok", response.ok);
     if (response.ok) {
       const json = await response.json();
       return json as FdrAPI.docs.v2.read.LoadDocsForUrlResponse;