Add streaming

web/src/app/page.tsx

@@ -1,11 +1,11 @@
-import { SearchSection } from "@/components/SearchBar";
+import { SearchSection } from "@/components/search/SearchSection";
 import { Header } from "@/components/Header";
 
 export default function Home() {
   return (
     <>
       <Header />
-      <div className="p-24 flex flex-col items-center min-h-screen bg-gray-900 text-gray-100">
+      <div className="px-24 pt-10 flex flex-col items-center min-h-screen bg-gray-900 text-gray-100">
         <div className="max-w-[800px] w-full">
           <SearchSection />
         </div>

web/src/components/SearchResultsDisplay.tsx

@@ -1,101 +0,0 @@
-import React from "react";
-import { Globe, SlackLogo, GoogleDriveLogo } from "@phosphor-icons/react";
-import "tailwindcss/tailwind.css";
-import { Quote, SearchResponse } from "./types";
-import { ThinkingAnimation } from "./Thinking";
-
-interface SearchResultsDisplayProps {
-  data: SearchResponse | undefined;
-  isFetching: boolean;
-}
-
-const ICON_SIZE = "20";
-const ICON_STYLE = "text-blue-600 my-auto mr-1 flex flex-shrink-0";
-
-const getSourceIcon = (sourceType: string) => {
-  switch (sourceType) {
-    case "web":
-      return <Globe size={ICON_SIZE} className={ICON_STYLE} />;
-    case "slack":
-      return <SlackLogo size={ICON_SIZE} className={ICON_STYLE} />;
-    case "google_drive":
-      return <GoogleDriveLogo size={ICON_SIZE} className={ICON_STYLE} />;
-    default:
-      return null;
-  }
-};
-
-export const SearchResultsDisplay: React.FC<SearchResultsDisplayProps> = ({
-  data,
-  isFetching,
-}) => {
-  if (isFetching) {
-    return <ThinkingAnimation />;
-  }
-
-  if (!data) {
-    return null;
-  }
-
-  const { answer, quotes } = data;
-  if (!answer || !quotes) {
-    return <div>Unable to find an answer</div>;
-  }
-
-  const dedupedQuotes: Quote[] = [];
-  const seen = new Set<string>();
-  Object.values(quotes).forEach((quote) => {
-    if (!seen.has(quote.document_id)) {
-      dedupedQuotes.push(quote);
-      seen.add(quote.document_id);
-    }
-  });
-
-  return (
-    <>
-      <div className="p-4 border rounded-md border-gray-700">
-        <h2 className="text font-bold mb-2">AI Answer</h2>
-        <p className="mb-4">{answer}</p>
-
-        <h2 className="text-sm font-bold mb-2">Sources</h2>
-        <div className="flex">
-          {dedupedQuotes.map((quoteInfo) => (
-            <a
-              key={quoteInfo.document_id}
-              className="p-2 border border-gray-800 rounded-lg text-sm flex max-w-[230px] hover:bg-gray-800"
-              href={quoteInfo.link}
-              target="_blank"
-              rel="noopener noreferrer"
-            >
-              {getSourceIcon(quoteInfo.source_type)}
-              <p className="truncate break-all">
-                {quoteInfo.semantic_identifier || quoteInfo.document_id}
-              </p>
-            </a>
-          ))}
-        </div>
-      </div>
-      <div className="mt-4">
-        <div className="font-bold border-b mb-4 pb-1 border-gray-800">
-          Results
-        </div>
-        {dedupedQuotes.map((quoteInfo) => (
-          <div key={quoteInfo.document_id} className="text-sm">
-            <a
-              className="rounded-lg flex font-bold"
-              href={quoteInfo.link}
-              target="_blank"
-              rel="noopener noreferrer"
-            >
-              {getSourceIcon(quoteInfo.source_type)}
-              <p className="truncate break-all">
-                {quoteInfo.semantic_identifier || quoteInfo.document_id}
-              </p>
-            </a>
-            <p className="p-2 mb-2 text-gray-200">{quoteInfo.blurb}</p>
-          </div>
-        ))}
-      </div>
-    </>
-  );
-};

web/src/components/search/SearchBar.tsx

@@ -1,63 +1,20 @@
-"use client";
-
 import React, { useState, KeyboardEvent, ChangeEvent } from "react";
 import { MagnifyingGlass } from "@phosphor-icons/react";
-import "tailwindcss/tailwind.css";
-import { SearchResultsDisplay } from "./SearchResultsDisplay";
-import { SearchResponse } from "./types";
-
-const searchRequest = async (query: string): Promise<SearchResponse> => {
-  const url = new URL("/api/direct-qa", window.location.origin);
-  const params = new URLSearchParams({
-    query,
-    collection: "semantic_search",
-  }).toString();
-  url.search = params;
-
-  const response = await fetch(url);
-  return response.json();
-};
-
-export const SearchSection: React.FC<{}> = () => {
-  const [answer, setAnswer] = useState<SearchResponse>();
-  const [isFetching, setIsFetching] = useState(false);
-
-  return (
-    <>
-      <SearchBar
-        onSearch={(query) => {
-          setIsFetching(true);
-          searchRequest(query).then((response) => {
-            setIsFetching(false);
-            setAnswer(response);
-          });
-        }}
-      />
-      <div className="mt-2">
-        <SearchResultsDisplay data={answer} isFetching={isFetching} />
-      </div>
-    </>
-  );
-};
 
 interface SearchBarProps {
   onSearch: (searchTerm: string) => void;
 }
 
-const SearchBar: React.FC<SearchBarProps> = ({ onSearch }) => {
+export const SearchBar: React.FC<SearchBarProps> = ({ onSearch }) => {
   const [searchTerm, setSearchTerm] = useState<string>("");
 
   const handleChange = (event: ChangeEvent<HTMLTextAreaElement>) => {
     const target = event.target;
     setSearchTerm(target.value);
 
-    // Reset the textarea height
+    // Resize the textarea to fit the content
     target.style.height = "24px";
-
-    // Calculate the new height based on scrollHeight
     const newHeight = target.scrollHeight;
-
-    // Apply the new height
     target.style.height = `${newHeight}px`;
   };
 

web/src/components/search/SearchResultsDisplay.tsx

@@ -0,0 +1,115 @@
+import React from "react";
+import { Globe, SlackLogo, GoogleDriveLogo } from "@phosphor-icons/react";
+import "tailwindcss/tailwind.css";
+import { Quote, Document } from "./types";
+import { ThinkingAnimation } from "../Thinking";
+
+interface SearchResultsDisplayProps {
+  answer: string | null;
+  quotes: Record<string, Quote> | null;
+  documents: Document[] | null;
+  isFetching: boolean;
+}
+
+const ICON_SIZE = "20";
+const ICON_STYLE = "text-blue-600 my-auto mr-1 flex flex-shrink-0";
+
+const getSourceIcon = (sourceType: string) => {
+  switch (sourceType) {
+    case "web":
+      return <Globe size={ICON_SIZE} className={ICON_STYLE} />;
+    case "slack":
+      return <SlackLogo size={ICON_SIZE} className={ICON_STYLE} />;
+    case "google_drive":
+      return <GoogleDriveLogo size={ICON_SIZE} className={ICON_STYLE} />;
+    default:
+      return null;
+  }
+};
+
+export const SearchResultsDisplay: React.FC<SearchResultsDisplayProps> = ({
+  answer,
+  quotes,
+  documents,
+  isFetching,
+}) => {
+  if (!answer) {
+    if (isFetching) {
+      return <ThinkingAnimation />;
+    }
+    return null;
+  }
+
+  if (answer === null) {
+    return <div>Unable to find an answer</div>;
+  }
+
+  const dedupedQuotes: Quote[] = [];
+  const seen = new Set<string>();
+  if (quotes) {
+    Object.values(quotes).forEach((quote) => {
+      if (!seen.has(quote.document_id)) {
+        dedupedQuotes.push(quote);
+        seen.add(quote.document_id);
+      }
+    });
+  }
+
+  return (
+    <>
+      <div className="p-4 border-2 rounded-md border-gray-700">
+        <h2 className="text font-bold mb-2">AI Answer</h2>
+        <p className="mb-4">{answer}</p>
+
+        {dedupedQuotes.length > 0 && (
+          <>
+            <h2 className="text-sm font-bold mb-2">Sources</h2>
+            <div className="flex">
+              {dedupedQuotes.map((quoteInfo) => (
+                <a
+                  key={quoteInfo.document_id}
+                  className="p-2 border border-gray-800 rounded-lg text-sm flex max-w-[230px] hover:bg-gray-800"
+                  href={quoteInfo.link}
+                  target="_blank"
+                  rel="noopener noreferrer"
+                >
+                  {getSourceIcon(quoteInfo.source_type)}
+                  <p className="truncate break-all">
+                    {quoteInfo.semantic_identifier || quoteInfo.document_id}
+                  </p>
+                </a>
+              ))}
+            </div>
+          </>
+        )}
+      </div>
+      {/* Only display docs once we're done fetching to avoid distracting from the AI answer*/}
+      {!isFetching && documents && documents.length > 0 && (
+        <div className="mt-4">
+          <div className="font-bold border-b mb-4 pb-1 border-gray-800">
+            Results
+          </div>
+          {documents.slice(0, 5).map((doc) => (
+            <div
+              key={doc.document_id}
+              className="text-sm border-b border-gray-800 mb-3"
+            >
+              <a
+                className="rounded-lg flex font-bold"
+                href={doc.link}
+                target="_blank"
+                rel="noopener noreferrer"
+              >
+                {getSourceIcon(doc.source_type)}
+                <p className="truncate break-all">
+                  {doc.semantic_name || doc.document_id}
+                </p>
+              </a>
+              <p className="pl-1 py-3 text-gray-200">{doc.blurb}</p>
+            </div>
+          ))}
+        </div>
+      )}
+    </>
+  );
+};

web/src/components/search/SearchSection.tsx

@@ -0,0 +1,93 @@
+"use client";
+
+import { useState } from "react";
+import { SearchBar } from "./SearchBar";
+import { SearchResultsDisplay } from "./SearchResultsDisplay";
+import { Quote, Document } from "./types";
+
+const searchRequestStreamed = async (
+  query: string,
+  updateCurrentAnswer: (val: string) => void,
+  updateQuotes: (quotes: Record<string, Quote>) => void,
+  updateDocs: (docs: Document[]) => void
+) => {
+  const url = new URL("/api/stream-direct-qa", window.location.origin);
+  const params = new URLSearchParams({
+    query,
+    collection: "semantic_search",
+  }).toString();
+  url.search = params;
+
+  let answer = "";
+  try {
+    const response = await fetch(url);
+    const reader = response.body?.getReader();
+    const decoder = new TextDecoder("utf-8");
+
+    while (true) {
+      const rawChunk = await reader?.read();
+      if (!rawChunk) {
+        throw new Error("Unable to process chunk");
+      }
+      const { done, value } = rawChunk;
+      if (done) {
+        break;
+      }
+
+      // Process each chunk as it arrives
+      const chunk = decoder.decode(value, { stream: true });
+      if (!chunk) {
+        break;
+      }
+      const chunkJson = JSON.parse(chunk);
+      const answerChunk = chunkJson.answer_data;
+      if (answerChunk) {
+        answer += answerChunk;
+        updateCurrentAnswer(answer);
+      } else {
+        const docs = chunkJson.top_documents as any[];
+        if (docs) {
+          updateDocs(docs.map((doc) => JSON.parse(doc)));
+        } else {
+          updateQuotes(chunkJson);
+        }
+      }
+    }
+  } catch (err) {
+    console.error("Fetch error:", err);
+  }
+  return answer;
+};
+
+export const SearchSection: React.FC<{}> = () => {
+  const [answer, setAnswer] = useState<string | null>("");
+  const [quotes, setQuotes] = useState<Record<string, Quote> | null>(null);
+  const [documents, setDocuments] = useState<Document[] | null>(null);
+  const [isFetching, setIsFetching] = useState(false);
+
+  return (
+    <>
+      <SearchBar
+        onSearch={(query) => {
+          setIsFetching(true);
+          setAnswer("");
+          setQuotes(null);
+          setDocuments(null);
+          searchRequestStreamed(query, setAnswer, setQuotes, setDocuments).then(
+            () => {
+              setIsFetching(false);
+            }
+          );
+        }}
+      />
+      <div className="mt-2">
+        <SearchResultsDisplay
+          answer={answer}
+          quotes={quotes}
+          documents={documents}
+          isFetching={isFetching}
+        />
+      </div>
+    </>
+  );
+};

web/src/components/search/types.tsx

@@ -6,6 +6,14 @@ export interface Quote {
   semantic_identifier: string | null;
 }
 
+export interface Document {
+  document_id: string;
+  link: string;
+  source_type: string;
+  blurb: string;
+  semantic_name: string | null;
+}
+
 export interface SearchResponse {
   answer: string;
   quotes: Record<string, Quote>;