Adds Kay and YDC (#13)

* Fix overly large payloads, retry on tab switch * Adds Kay and YDC
langchain-ai · Oct 18, 2023 · 483399b · 483399b
1 parent 2e8494f
commit 483399b
Show file tree

Hide file tree

Showing 7 changed files with 114 additions and 79 deletions.
diff --git a/README.md b/README.md
@@ -16,10 +16,16 @@ The code includes a simple backup that uses the Google Custom Search Engine for
 export OPENAI_API_KEY=
 export TAVILY_API_KEY=
 
-# if you'd like to use the backup retriever
+# if you'd like to use the You.com retriever
+export YDC_API_KEY=
+
+# if you'd like to use the Google retriever
 export GOOGLE_CSE_ID=
 export GOOGLE_API_KEY=
 
+# if you'd like to use the Kay.ai retriever
+export KAY_API_KEY=
+
 # for tracing
 export LANGCHAIN_TRACING_V2=true
 export LANGCHAIN_ENDPOINT="https://api.smith.langchain.com"

diff --git a/main.py b/main.py
@@ -18,6 +18,7 @@
                                   TavilySearchAPIRetriever)
 from langchain.retrievers.document_compressors import (
     DocumentCompressorPipeline, EmbeddingsFilter)
+from langchain.retrievers.kay import KayAiRetriever
 from langchain.retrievers.you import YouRetriever
 from langchain.schema import Document
 from langchain.schema.document import Document
@@ -173,13 +174,22 @@ def get_retriever():
     you_retriever = ContextualCompressionRetriever(
         base_compressor=pipeline_compressor, base_retriever=base_you_retriever
     )
+    base_kay_retriever = KayAiRetriever.create(
+        dataset_id="company",
+        data_types=["10-K", "10-Q", "PressRelease"],
+        num_contexts=3,
+    )
+    kay_retriever = ContextualCompressionRetriever(
+        base_compressor=pipeline_compressor, base_retriever=base_kay_retriever
+    )
     return tavily_retriever.configurable_alternatives(
         # This gives this field an id
         # When configuring the end runnable, we can then use this id to configure this field
         ConfigurableField(id="retriever"),
         default_key="tavily",
         google=google_retriever,
         you=you_retriever,
+        kay=kay_retriever,
     ).with_config(run_name="FinalSourceRetriever")
 
 

diff --git a/nextjs/app/components/ChatWindow.tsx b/nextjs/app/components/ChatWindow.tsx
@@ -24,6 +24,9 @@ import {
 } from "@chakra-ui/react";
 import { ArrowUpIcon } from "@chakra-ui/icons";
 import { Source } from "./SourceBubble";
+import { DefaultQuestion } from "./DefaultQuestion";
+
+type RetrieverName = "tavily" | "kay" | "you" | "google";
 
 export function ChatWindow(props: {
   apiBaseUrl: string;
@@ -35,7 +38,7 @@ export function ChatWindow(props: {
   const [messages, setMessages] = useState<Array<Message>>([]);
   const [input, setInput] = useState("");
   const [isLoading, setIsLoading] = useState(false);
-  const [retriever, setRetriever] = useState("tavily");
+  const [retriever, setRetriever] = useState<RetrieverName>("tavily");
   const [llm, setLlm] = useState("openai");
 
   const [chatHistory, setChatHistory] = useState<
@@ -141,7 +144,8 @@ export function ChatWindow(props: {
               sources = streamedResponse.logs[
                 sourceStepName
               ].final_output.documents.map((doc: Record<string, any>) => ({
-                url: doc.metadata.source,
+                url: doc.metadata.source ?? doc.metadata.data_source_link,
+                defaultSourceUrl: retriever === "you" ? "https://you.com" : "",
                 title: doc.metadata.title,
                 images: doc.metadata.images,
               }));
@@ -183,6 +187,30 @@ export function ChatWindow(props: {
     }
   };
 
+  const defaultQuestions = [
+    "what is langchain?",
+    "history of mesopotamia",
+    "how to build a discord bot",
+    "leonardo dicaprio girlfriend",
+    "fun gift ideas for software engineers",
+    "how does a prism separate light",
+    "what bear is best",
+  ];
+
+  const DEFAULT_QUESTIONS: Record<RetrieverName, string[]> = {
+    tavily: defaultQuestions,
+    you: defaultQuestions,
+    google: defaultQuestions,
+    kay: [
+      "Is Johnson & Johnson increasing its marketing budget?",
+      "How is Lululemon adapting to new customer trends?",
+      "Which industries are growing in recent 10-Q reports?",
+      "Who are Etsy’s competitors?",
+      "Which companies reported data breaches?",
+      "What were the biggest strategy changes made by Roku in 2023?",
+    ],
+  };
+
   const sendInitialQuestion = async (question: string) => {
     await sendMessage(question);
   };
@@ -216,9 +244,12 @@ export function ChatWindow(props: {
         </Heading>
         <div className="text-white flex items-center mt-4">
           <span className="shrink-0 mr-2">Powered by</span>
-          <Select onChange={(e) => setRetriever(e.target.value)}>
+          <Select
+            onChange={(e) => setRetriever(e.target.value as RetrieverName)}
+          >
             <option value="tavily">Tavily</option>
-            {/* <option value="you">You.com</option> */}
+            <option value="kay">Kay.ai SEC Filings</option>
+            <option value="you">You.com</option>
             <option value="google">Google</option>
           </Select>
           <span className="shrink-0 ml-2 mr-2">and</span>
@@ -285,64 +316,34 @@ export function ChatWindow(props: {
       {messages.length === 0 ? (
         <div className="w-full text-center flex flex-col">
           <div className="flex grow justify-center w-full mt-4">
-            <div
-              onMouseUp={(e) =>
-                sendInitialQuestion((e.target as HTMLDivElement).innerText)
-              }
-              className="bg-stone-700 px-2 py-1 mx-2 rounded cursor-pointer justify-center text-stone-200 hover:bg-stone-500"
-            >
-              what is langchain?
-            </div>
-            <div
-              onMouseUp={(e) =>
-                sendInitialQuestion((e.target as HTMLDivElement).innerText)
-              }
-              className="bg-stone-700 px-2 py-1 mx-2 rounded cursor-pointer justify-center text-stone-200 hover:bg-stone-500"
-            >
-              history of mesopotamia
-            </div>
-            <div
-              onMouseUp={(e) =>
-                sendInitialQuestion((e.target as HTMLDivElement).innerText)
-              }
-              className="bg-stone-700 px-2 py-1 mx-2 rounded cursor-pointer justify-center text-stone-200 hover:bg-stone-500"
-            >
-              how to build a discord bot
-            </div>
-            <div
-              onMouseUp={(e) =>
-                sendInitialQuestion((e.target as HTMLDivElement).innerText)
-              }
-              className="bg-stone-700 px-2 py-1 mx-2 rounded cursor-pointer justify-center text-stone-200 hover:bg-stone-500"
-            >
-              leonardo dicaprio girlfriend
-            </div>
+            {DEFAULT_QUESTIONS[retriever]
+              .slice(0, 4)
+              .map((defaultQuestion, i) => {
+                return (
+                  <DefaultQuestion
+                    key={`defaultquestion:${i}`}
+                    question={defaultQuestion}
+                    onMouseUp={(e) =>
+                      sendInitialQuestion(
+                        (e.target as HTMLDivElement).innerText,
+                      )
+                    }
+                  ></DefaultQuestion>
+                );
+              })}
           </div>
           <div className="flex grow justify-center w-full mt-4">
-            <div
-              onMouseUp={(e) =>
-                sendInitialQuestion((e.target as HTMLDivElement).innerText)
-              }
-              className="bg-stone-700 px-2 py-1 mx-2 rounded cursor-pointer justify-center text-stone-200 hover:bg-stone-500"
-            >
-              fun gift ideas for software engineers
-            </div>
-            <div
-              onMouseUp={(e) =>
-                sendInitialQuestion((e.target as HTMLDivElement).innerText)
-              }
-              className="bg-stone-700 px-2 py-1 mx-2 rounded cursor-pointer justify-center text-stone-200 hover:bg-stone-500"
-            >
-              how does a prism separate light
-            </div>
-            <div
-              onMouseUp={(e) =>
-                sendInitialQuestion((e.target as HTMLDivElement).innerText)
-              }
-              className="bg-stone-700 px-2 py-1 mx-2 rounded cursor-pointer justify-center text-stone-200 hover:bg-stone-500"
-            >
-              what bear is best
-            </div>
+            {DEFAULT_QUESTIONS[retriever].slice(4).map((defaultQuestion, i) => {
+              return (
+                <DefaultQuestion
+                  key={`defaultquestion:${i + 4}`}
+                  question={defaultQuestion}
+                  onMouseUp={(e) =>
+                    sendInitialQuestion((e.target as HTMLDivElement).innerText)
+                  }
+                ></DefaultQuestion>
+              );
+            })}
           </div>
         </div>
       ) : (

diff --git a/nextjs/app/components/DefaultQuestion.tsx b/nextjs/app/components/DefaultQuestion.tsx
@@ -0,0 +1,15 @@
+import { MouseEventHandler } from "react";
+
+export function DefaultQuestion(props: {
+  question: string;
+  onMouseUp: MouseEventHandler;
+}) {
+  return (
+    <div
+      onMouseUp={props.onMouseUp}
+      className="bg-stone-700 px-2 py-1 mx-2 rounded cursor-pointer justify-center text-stone-200 hover:bg-stone-500"
+    >
+      {props.question}
+    </div>
+  );
+}
diff --git a/nextjs/app/components/SourceBubble.tsx b/nextjs/app/components/SourceBubble.tsx
@@ -5,6 +5,7 @@ export type Source = {
   url: string;
   title: string;
   images: string[];
+  defaultSourceUrl?: string;
 };
 
 export function SourceBubble(props: {
@@ -14,26 +15,13 @@ export function SourceBubble(props: {
   onMouseEnter: () => any;
   onMouseLeave: () => any;
 }) {
-  const cumulativeOffset = function (element: HTMLElement | null) {
-    var top = 0,
-      left = 0;
-    do {
-      top += element?.offsetTop || 0;
-      left += element?.offsetLeft || 0;
-      element = (element?.offsetParent as HTMLElement) || null;
-    } while (element);
-
-    return {
-      top: top,
-      left: left,
-    };
-  };
-
-  const hostname = new URL(props.source.url).hostname.replace("www.", "");
+  const hostname = new URL(
+    props.source.url ?? props.source.defaultSourceUrl,
+  ).hostname.replace("www.", "");
 
   return (
     <a
-      href={props.source.url}
+      href={props.source.url ?? props.source.defaultSourceUrl}
       target="_blank"
       onMouseEnter={props.onMouseEnter}
       onMouseLeave={props.onMouseLeave}

diff --git a/poetry.lock b/poetry.lock
diff --git a/pyproject.toml b/pyproject.toml
@@ -19,6 +19,7 @@ sse-starlette = "^1.6.5"
 google-api-python-client = "^2.102.0"
 html2text = "^2020.1.16"
 anthropic = "^0.4.1"
+kay = "^0.1.2"
 
 
 [tool.poetry.group.dev.dependencies]