demo: trieve

packages/fern-docs/components/.storybook/variables.css

@@ -64,7 +64,7 @@
   --accent-aaa: 173, 255, 140;
   --accent-tinted: 191, 255, 165;
   --background: 0, 0, 0;
-  --accent-contrast: 0, 0, 0;
+  --accent-contrast: #000;
   --bg-color-card: var(--grayscale-a2);
   --sidebar-background: transparent;
   --header-background: transparent;

packages/fern-docs/components/src/FernButtonV2.tsx

@@ -9,7 +9,7 @@ const buttonVariants = cva(
     variants: {
       variant: {
         default:
-          "bg-[var(--accent-10)] text-[var(--accent-1)] hover:bg-[var(--accent-9)]",
+          "bg-[var(--accent-10)] text-[var(--accent-contrast,var(--accent-1))] hover:bg-[var(--accent-9)]",
         destructive:
           "bg-[var(--red-10)] text-[var(--red-12)] hover:bg-[var(--red-11)]",
         outline:

packages/fern-docs/search-ui/package.json

@@ -20,6 +20,8 @@
     "@ai-sdk/cohere": "^1.0.5",
     "@ai-sdk/openai": "^1.0.8",
     "@ai-sdk/provider": "^1.0.2",
+    "@ai-sdk/provider-utils": "^2.0.4",
+    "@ai-sdk/ui-utils": "^1.0.5",
     "@algolia/client-search": "^5.15.0",
     "@fern-api/ui-core-utils": "workspace:*",
     "@fern-docs/components": "workspace:*",
@@ -59,6 +61,7 @@
     "remark-gfm": "^4.0.0",
     "server-only": "^0.0.1",
     "tailwind-merge": "^2.3.0",
+    "trieve-ts-sdk": "^0.0.49",
     "ts-essentials": "^10.0.1",
     "unist-util-visit": "^5.0.0",
     "zod": "^3.23.8",

packages/fern-docs/search-ui/src/app/(demo)/api/suggest/route.ts

@@ -1,67 +1,33 @@
-import { algoliaAppId } from "@/server/env-variables";
-import { models } from "@/server/models";
-import { searchClient } from "@algolia/client-search";
 import { SuggestionsSchema } from "@fern-docs/search-server";
-import {
-  SEARCH_INDEX,
-  type AlgoliaRecord,
-} from "@fern-docs/search-server/algolia";
-import { kv } from "@vercel/kv";
-import { streamObject } from "ai";
+import { TrieveSDK } from "trieve-ts-sdk";
 import { z } from "zod";
 
-// Allow streaming responses up to 30 seconds
-export const maxDuration = 30;
-
 const BodySchema = z.object({
-  algoliaSearchKey: z.string(),
-  model: z.string().optional(),
+  apiKey: z.string(),
+  datasetId: z.string(),
 });
 
 export async function POST(request: Request): Promise<Response> {
-  const { algoliaSearchKey, model: _model } = BodySchema.parse(
-    await request.json()
-  );
-
-  const model = models[(_model as keyof typeof models) ?? "gpt-4o-mini"];
+  const { apiKey, datasetId } = BodySchema.parse(await request.json());
 
-  if (!algoliaSearchKey || typeof algoliaSearchKey !== "string") {
-    return new Response("Missing search key", { status: 400 });
-  }
+  const client = new TrieveSDK({
+    apiKey,
+    datasetId,
+  });
+  const suggestions = SuggestionsSchema.parse({
+    suggestions: [],
+  });
 
-  const cacheKey = `suggestions:${algoliaSearchKey}`;
-  const cachedSuggestions = await kv.get<string>(cacheKey);
-  if (cachedSuggestions) {
-    return new Response(JSON.stringify(cachedSuggestions), {
-      headers: { "Content-Type": "text/plain; charset=utf-8" },
+  try {
+    const suggestedQueriesResult = await client.suggestedQueries({
+      suggestion_type: "question",
     });
+    suggestions.suggestions = suggestedQueriesResult.queries.slice(0, 5);
+  } catch (e) {
+    console.error("error getting suggestions from trieve", e);
   }
 
-  const client = searchClient(algoliaAppId(), algoliaSearchKey);
-  const response = await client.searchSingleIndex<AlgoliaRecord>({
-    indexName: SEARCH_INDEX,
-    searchParams: { query: "", hitsPerPage: 100, attributesToSnippet: [] },
-  });
-
-  const result = streamObject({
-    model,
-    system:
-      "You are a helpful assistant that suggestions of questions for the user to ask about the documentation. Generate 5 questions based on the following search results.",
-    prompt: response.hits
-      .map(
-        (hit) =>
-          `# ${hit.title}\n${hit.description ?? ""}\n${hit.type === "changelog" || hit.type === "markdown" ? (hit.content ?? "") : ""}`
-      )
-      .join("\n\n"),
-    maxRetries: 3,
-    schema: SuggestionsSchema,
-    onFinish: async ({ object }) => {
-      if (object) {
-        await kv.set(cacheKey, object);
-        await kv.expire(cacheKey, 60 * 60);
-      }
-    },
+  return new Response(JSON.stringify(suggestions), {
+    headers: { "Content-Type": "application/json" },
   });
-
-  return result.toTextStreamResponse();
 }

packages/fern-docs/search-ui/src/app/(demo)/api/trieve/route.ts

@@ -0,0 +1,98 @@
+import { formatDataStreamPart } from "@ai-sdk/ui-utils";
+import { TrieveSDK } from "trieve-ts-sdk";
+import { z } from "zod";
+// Allow streaming responses up to 30 seconds
+export const maxDuration = 60;
+
+const BodySchema = z.object({
+  apiKey: z.string(),
+  datasetId: z.string().optional(),
+  organizationId: z.string().optional(),
+  topicId: z.string().optional(),
+  messages: z.array(
+    z.object({
+      role: z.enum(["user", "assistant"]),
+      content: z.string(),
+    })
+  ),
+});
+
+export async function POST(request: Request): Promise<Response> {
+  const {
+    apiKey,
+    datasetId,
+    organizationId,
+    topicId: _topicId,
+    messages,
+  } = BodySchema.parse(await request.json());
+
+  const client = new TrieveSDK({
+    apiKey,
+    datasetId,
+    organizationId,
+  });
+
+  const first_user_message = messages.find(
+    (message) => message.role === "user"
+  )?.content;
+
+  const topic_id =
+    _topicId ??
+    (
+      await client.createTopic({
+        owner_id: "testing",
+        first_user_message,
+      })
+    ).id;
+
+  const new_message_content = messages[messages.length - 1].content;
+
+  const reader = await client.createMessageReader({
+    topic_id,
+    new_message_content,
+  });
+
+  const decoder = new TextDecoder("utf-8");
+
+  const stream = new ReadableStream({
+    async start(controller) {
+      controller.enqueue(formatDataStreamPart("data", [{ topic_id }]));
+
+      let is_data = true;
+      let data = "";
+
+      while (true) {
+        const { done, value } = await reader.read();
+        if (done) {
+          break;
+        }
+        const str = decoder.decode(value);
+        if (is_data) {
+          if (str.endsWith("||")) {
+            data += str.slice(0, -2);
+            is_data = false;
+            controller.enqueue(formatDataStreamPart("data", JSON.parse(data)));
+          } else {
+            data += str;
+          }
+        } else {
+          controller.enqueue(formatDataStreamPart("text", str));
+        }
+      }
+      controller.enqueue(
+        formatDataStreamPart("finish_message", {
+          finishReason: "stop",
+          usage: {
+            promptTokens: NaN,
+            completionTokens: NaN,
+          },
+        })
+      );
+      controller.close();
+    },
+  });
+
+  return new Response(stream, {
+    headers: { "Content-Type": "text/plain; charset=utf-8" },
+  });
+}

packages/fern-docs/search-ui/src/app/(demo)/search-demo.tsx

@@ -20,9 +20,11 @@ import {
 import { ChatbotModelSelect } from "@/components/chatbot/model-select";
 import { DesktopCommandWithAskAI } from "@/components/desktop/desktop-ask-ai";
 import { useIsMobile } from "@/hooks/use-mobile";
+import { isPlainObject } from "@fern-api/ui-core-utils";
 import { FacetsResponse, SEARCH_INDEX } from "@fern-docs/search-server/algolia";
 import { useAtom } from "jotai";
 import { atomWithStorage } from "jotai/utils";
+import { useSearchParams } from "next/navigation";
 
 const USER_TOKEN_KEY = "user-token";
 
@@ -41,12 +43,23 @@ export function DemoInstantSearchClient({
   appId: string;
   domain: string;
 }): ReactElement | false {
+  const searchParams = useSearchParams();
+  const isTrieve = searchParams.get("provider") === "trieve";
+  const trieveApiKey = isTrieve ? searchParams.get("apiKey") : undefined;
+  const trieveDatasetId = isTrieve ? searchParams.get("datasetId") : undefined;
+  const trieveOrganizationId = isTrieve
+    ? searchParams.get("organizationId")
+    : undefined;
   const [initialInput, setInitialInput] = useState("");
   const [askAi, setAskAi] = useState(false);
   const [open, setOpen] = useState(false);
   const { setTheme } = useTheme();
   const isMobile = useIsMobile();
 
+  const [trieveTopicId, setTrieveTopicId] = useState<string | undefined>(
+    undefined
+  );
+
   useEffect(() => {
     const handleKeyDown = (event: KeyboardEvent) => {
       if (isMobile) {
@@ -141,20 +154,46 @@ export function DemoInstantSearchClient({
             domain={domain}
             askAI={askAi}
             setAskAI={setAskAi}
+            api={isTrieve ? "/api/trieve" : "/api/chat"}
             onEscapeKeyDown={() => setOpen(false)}
-            api="/api/chat"
             suggestionsApi="/api/suggest"
             initialInput={initialInput}
             setInitialInput={setInitialInput}
-            body={{
-              algoliaSearchKey: apiKey,
-              domain,
-              model,
-            }}
+            body={
+              isTrieve
+                ? {
+                    apiKey: trieveApiKey,
+                    datasetId: trieveDatasetId,
+                    organizationId: trieveOrganizationId,
+                    topicId: trieveTopicId,
+                  }
+                : {
+                    algoliaSearchKey: apiKey,
+                    domain,
+                    model,
+                  }
+            }
             onSelectHit={handleSubmit}
             composerActions={
-              <ChatbotModelSelect value={model} onValueChange={setModel} />
+              isTrieve ? (
+                <span className="align-bottom text-xs text-[var(--grayscale-a9)]">
+                  Powered by{" "}
+                  <a href="https://trieve.ai" target="_blank">
+                    Trieve
+                  </a>
+                </span>
+              ) : (
+                <ChatbotModelSelect value={model} onValueChange={setModel} />
+              )
             }
+            onData={(data) => {
+              const topicId = data
+                .filter(isPlainObject)
+                .find((d) => d.type === "topic_id")?.topic_id;
+              if (topicId && typeof topicId === "string") {
+                setTrieveTopicId(topicId);
+              }
+            }}
           >
             <DefaultDesktopBackButton />
 

packages/fern-docs/search-ui/src/components/desktop/desktop-ask-ai.tsx

@@ -75,6 +75,7 @@ export const DesktopCommandWithAskAI = forwardRef<
     prefetch?: (path: string) => Promise<void>;
     composerActions?: ReactNode;
     domain: string;
+    onData?: (data: unknown[]) => void;
     renderActions?: (message: SqueezedMessage) => ReactNode;
     setInitialInput?: (initialInput: string) => void;
     children?: ReactNode;
@@ -220,6 +221,7 @@ const DesktopAskAIContent = (props: {
   prefetch?: (path: string) => Promise<void>;
   composerActions?: ReactNode;
   domain: string;
+  onData?: (data: unknown[]) => void;
   renderActions?: (message: SqueezedMessage) => ReactNode;
 }) => {
   return (
@@ -261,6 +263,7 @@ const DesktopAskAIChat = ({
   prefetch,
   composerActions,
   domain,
+  onData,
   renderActions,
 }: {
   onReturnToSearch?: () => void;
@@ -274,6 +277,7 @@ const DesktopAskAIChat = ({
   prefetch?: (path: string) => Promise<void>;
   composerActions?: ReactNode;
   domain: string;
+  onData?: (data: unknown[]) => void;
   renderActions?: (message: SqueezedMessage) => ReactNode;
 }) => {
   const inputRef = useRef<HTMLTextAreaElement>(null);
@@ -293,6 +297,13 @@ const DesktopAskAIChat = ({
     }),
   });
 
+  useEffect(() => {
+    if (chat.data) {
+      onData?.(chat.data);
+    }
+    // eslint-disable-next-line react-hooks/exhaustive-deps
+  }, [chat.data]);
+
   // Reset userScrolled when the chat is loading
   useIsomorphicLayoutEffect(() => {
     if (chat.isLoading) {
@@ -481,7 +492,7 @@ const AskAIComposer = forwardRef<
     const inputRef = useRef<HTMLTextAreaElement>(null);
     return (
       <div
-        className="cursor-text border-t border-[var(--grayscale-a6)]"
+        className="cursor-text border-t border-[var(--grayscale-a6)] p-2"
         onClick={() => inputRef.current?.focus()}
       >
         <DesktopCommandInput asChild>
@@ -538,12 +549,12 @@ const AskAIComposer = forwardRef<
             )}
           />
         </DesktopCommandInput>
-        <div className="flex items-center justify-between px-2 pb-2">
+        <div className="flex items-center justify-between">
           <div>{actions}</div>
           <Button
             size="icon"
             className="rounded-full"
-            variant="ghost"
+            variant="default"
             onClick={isLoading ? stop : () => onSend?.(value)}
             disabled={!isLoading && !canSubmit}
           >

packages/fern-docs/search-ui/src/server/with-search-api-key.ts

@@ -1,8 +1,8 @@
 import {
   DEFAULT_SEARCH_API_KEY_EXPIRATION_SECONDS,
   SEARCH_INDEX,
-  getSearchApiKey,
 } from "@fern-docs/search-server/algolia";
+import { getSearchApiKey } from "@fern-docs/search-server/algolia/edge";
 
 interface WithSearchApiKeyOptions {
   searchApiKey: string;

packages/fern-docs/tailwind.config.js

@@ -100,7 +100,7 @@ module.exports = {
         "accent-aa": withOpacity("--accent-aa"),
         "accent-aaa": withOpacity("--accent-aaa"),
         "accent-tinted": "var(--accent-10)",
-        "accent-contrast": withOpacity("--accent-contrast"),
+        "accent-contrast": "var(--accent-contrast)",
         "accent-muted": `var(--accent-6)`,
         "accent-highlight": "var(--accent-3)",
         "accent-highlight-faded": "var(--accent-2)",