feature: streaming LLM responses

devflowinc · Sep 5, 2024 · d90c242 · d90c242
1 parent e6bdac9
commit d90c242
Show file tree

Hide file tree

Showing 6 changed files with 106 additions and 28 deletions.
diff --git a/clients/search-component/example/index.html b/clients/search-component/example/index.html
@@ -4,7 +4,17 @@
     <meta charset="UTF-8" />
     <link rel="icon" type="image/svg+xml" href="/vite.svg" />
     <meta name="viewport" content="width=device-width, initial-scale=1.0" />
-    <title>Vite + React + TS</title>
+    <link rel="apple-touch-icon" sizes="180x180" href="https://cdn.trieve.ai/apple-touch-icon.png" />
+    <link rel="icon" type="image/png" sizes="32x32" href="https://cdn.trieve.ai/favicon-32x32.png" />
+    <link rel="icon" type="image/png" sizes="16x16" href="https://cdn.trieve.ai/favicon-16x16.png" />
+    <meta name="viewport" content="width=device-width, initial-scale=1.0" />
+    <meta property="og:title" content="Trieve Drop-In RAG and Search Component">
+    <meta property="og:site_name" content="advanced relevance semantic search and RAG API">
+    <meta property="og:url" content="https://docsearch.trieve.ai">
+    <meta property="og:description" content="Build better, faster, and more relevant search and RAG with our open source API. Date recency biasing, re-ranker models, dense vector search, sub-sentence highlighting, and more all on one endpoint that you can host yourself." >
+    <meta property="og:type" content="">
+    <meta property="og:image" content="https://cdn.trieve.ai/trieve-og.png">
+    <title>Trieve Drop-In RAG and Search Component</title>
   </head>
   <body>
     <div id="root"></div>

diff --git a/clients/search-component/src/TrieveModal/ChatMode.tsx b/clients/search-component/src/TrieveModal/ChatMode.tsx
@@ -44,7 +44,7 @@ export const ChatMode = ({
       setIsLoading(true);
       const fingerprint = await getFingerprint();
       const topic = await trieve.createTopic({
-        first_user_message: currentQuestion,
+        name: currentQuestion,
         owner_id: fingerprint.toString(),
       });
       setCurrentQuestion("");
@@ -53,25 +53,56 @@ export const ChatMode = ({
     }
   };
 
+  const handleReader = async (
+    reader: ReadableStreamDefaultReader<Uint8Array>
+  ) => {
+    setIsLoading(true);
+    let done = false;
+    let textInStream = "";
+
+    while (!done) {
+      const { value, done: doneReading } = await reader.read();
+      if (doneReading) {
+        done = doneReading;
+        setIsLoading(false);
+      } else if (value) {
+        const decoder = new TextDecoder();
+        const newText = decoder.decode(value);
+        textInStream += newText;
+        const [text, jsonData] = textInStream.split("||");
+        let json;
+        try {
+          json = JSON.parse(jsonData);
+          // eslint-disable-next-line @typescript-eslint/no-unused-vars
+        } catch (_) {
+          json = null;
+        }
+
+        setMessages((m) => [
+          [
+            m[0][0],
+            {
+              type: "system",
+              text: text,
+              additional: json ? json : null,
+            },
+          ],
+          ...m.slice(1),
+        ]);
+      }
+    }
+  };
+
   const createQuestion = async (id?: string) => {
     setIsLoading(true);
-    const answer = await trieve.createMessage({
+    const reader = await trieve.createMessageReader({
       topic_id: id || currentTopic,
       new_message_content: currentQuestion,
+      llm_options: {
+        completion_first: true,
+      },
     });
-    const [json, text] = answer.split("||");
-    setMessages((m) => [
-      [
-        m[0][0],
-        {
-          type: "system",
-          text: text,
-          additional: JSON.parse(json),
-        },
-      ],
-      ...m.slice(1),
-    ]);
-    setIsLoading(false);
+    handleReader(reader);
   };
 
   useEffect(() => {
@@ -105,17 +136,17 @@ export const ChatMode = ({
       </div>
       <div className="chat-modal-wrapper">
         {messages.map((chat, i) => (
-          <ul className="chat-ul">
+          <ul className="chat-ul" key={i}>
             {chat.map((message, idx) => {
               return (
-                <div>
+                <div key={idx}>
                   {message.type == "user" ? (
                     <div className={message.type}>
                       {isLoading && i === 0 ? <LoadingAIIcon /> : <CheckIcon />}
                       <span> {message.text}</span>
                     </div>
                   ) : null}
-                  {isLoading && i === 0 ? (
+                  {isLoading && i === 0 && !message.text ? (
                     <div className="system">
                       <LoadingIcon />
                     </div>

diff --git a/clients/search-component/src/TrieveModal/index.css b/clients/search-component/src/TrieveModal/index.css
@@ -7,7 +7,7 @@ body {
     @apply data-[state=open]:animate-contentShow fixed top-[50%] left-[50%] max-h-[80vh] w-[90vw] sm:max-w-[550px] -translate-x-[50%] -translate-y-[50%] rounded-lg bg-zinc-50 p-2 pt-0 shadow-2xl focus:outline-none text-zinc-950;
 
     .trieve-elements-search {
-      @apply pr-2
+      @apply pr-2;
     }
 
     .trieve-powered {
@@ -67,7 +67,7 @@ body {
     h2,
     h3,
     h4 {
-      @apply font-semibold py-2;
+      @apply font-semibold pb-2;
     }
 
     a {
@@ -94,7 +94,7 @@ body {
       @apply sticky top-0 bg-zinc-50 pt-2 flex items-center z-10 mt-1;
 
       &::after {
-        @apply absolute h-[17px] w-full z-50 block top-[56px] bg-gradient-to-b from-zinc-50 via-zinc-50  to-transparent;
+        @apply absolute h-[17px] w-full z-50 block top-[56px] bg-gradient-to-b from-zinc-50 via-zinc-50 to-transparent dark:from-zinc-900 dark:via-zinc-900 dark:to-transparent;
         content: "";
       }
 
@@ -257,7 +257,7 @@ body {
           display: none;
         }
       }
-      
+
       .mac {
         display: none;
       }

diff --git a/clients/ts-sdk/package.json b/clients/ts-sdk/package.json
@@ -6,7 +6,7 @@
   "files": [
     "dist"
   ],
-  "version": "0.0.3",
+  "version": "0.0.4",
   "license": "MIT",
   "scripts": {
     "lint": "eslint 'src/**/*.ts'",
@@ -19,7 +19,7 @@
     "build:docs": "typedoc"
   },
   "devDependencies": {
-    "@hey-api/openapi-ts": "^0.52.9",
+    "@hey-api/openapi-ts": "^0.53",
     "@vitest/ui": "^2.0.5",
     "esbuild": "^0.23.1",
     "esbuild-register": "^3.6.0",

diff --git a/clients/ts-sdk/src/fetch-client.ts b/clients/ts-sdk/src/fetch-client.ts
@@ -88,9 +88,9 @@ interface TrieveOpts {
 }
 
 export class TrieveFetchClient {
-  private apiKey?: string;
-  private baseUrl: string;
-  private debug: boolean = false;
+  apiKey?: string;
+  baseUrl: string;
+  debug: boolean = false;
 
   constructor(opts: TrieveOpts) {
     this.apiKey = opts.apiKey;

diff --git a/clients/ts-sdk/src/functions/message/index.ts b/clients/ts-sdk/src/functions/message/index.ts
@@ -39,6 +39,43 @@ export async function createMessage(
   );
 }
 
+/**
+ * Create a message. Messages are attached to topics in order to coordinate memory of gen-AI chat sessions.Auth’ed user or api key must have an admin or owner role for the specified dataset’s organization.
+ * 
+ * Example:
+ * ```js
+ *const reader = await trieve.createMessageReader({
+  topic_id: "3c90c3cc-1d76-27198-8888-8dd25736052a",
+  new_message_content: "a new message"
+});
+ * ```
+ */
+export async function createMessageReader(
+  /** @hidden */
+  this: TrieveSDK,
+  data: CreateMessageReqPayload,
+  signal?: AbortSignal
+) {
+  const response = await fetch(this.trieve.baseUrl + "/api/message", {
+    method: "post",
+    headers: {
+      "Content-Type": "application/json",
+      "TR-Dataset": this.datasetId,
+      Authorization: `Bearer ${this.trieve.apiKey}`,
+    },
+    body: JSON.stringify(data),
+    signal,
+  });
+
+  const reader = response.body?.getReader();
+
+  if (!reader) {
+    throw new Error("Failed to get reader from response body");
+  }
+
+  return reader;
+}
+
 /**
  * Edit message which exists within the topic’s chat history. This will delete the message and replace it with a new message. The new message will be generated by the AI based on the new content provided in the request body. The response will include Chunks first on the stream if the topic is using RAG. The structure will look like [chunks]||mesage. See docs.trieve.ai for more information. Auth’ed user or api key must have an admin or owner role for the specified dataset’s organization.
  *