feat: add internet search functionality to bot

n4ze3m · Aug 24, 2024 · 8eb1c66 · 8eb1c66
1 parent 7b9793e
commit 8eb1c66
Show file tree

Hide file tree

Showing 6 changed files with 87 additions and 36 deletions.
diff --git a/app/ui/src/@types/bot.ts b/app/ui/src/@types/bot.ts
@@ -22,6 +22,7 @@ export type BotSettings = {
     inactivityTimeout: number;
     autoResetSession: boolean;
     autoSyncDataSources: boolean;
+    internetSearchEnabled: boolean;
   };
   chatModel: {
     label: string;

diff --git a/app/ui/src/components/Bot/Playground/Message.tsx b/app/ui/src/components/Bot/Playground/Message.tsx
@@ -10,6 +10,7 @@ import React from "react";
 import { removeUUID } from "../../../utils/filename";
 import { useSpeechSynthesis } from "../../../hooks/useSpeechSynthesis";
 import { useElevenLabsTTS } from "../../../hooks/useElevenLabsTTS";
+import { Collapse } from "antd";
 
 type Props = Message & {
   onSourceClick(source: any): void;
@@ -56,23 +57,40 @@ export const PlaygroundMessage = (props: Props) => {
             </div>
 
             {props.isBot && (
-              <div className="mt-3 flex flex-wrap gap-2">
-                {props?.sources?.map((source, index) => (
-                  <button
-                    key={index}
-                    onClick={props.onSourceClick.bind(null, source)}
-                    className="inline-flex cursor-pointer transition-shadow duration-300 ease-in-out hover:shadow-lg  items-center rounded-md bg-gray-100 p-1 text-xs text-gray-800 border border-gray-300 dark:bg-gray-800 dark:border-gray-700 dark:text-gray-100 opacity-80 hover:opacity-100"
-                  >
-                    <span className="text-xs">
-                      {removeUUID(
-                        `${
-                          source?.metadata?.path || source?.metadata?.source
-                        }`.replace("./uploads/", "")
-                      )}
-                    </span>
-                  </button>
-                ))}
-              </div>
+              <Collapse
+                className="mt-6"
+                ghost
+                items={[
+                  {
+                    key: "1",
+                    label: (
+                      <div className="italic text-gray-500 dark:text-gray-400">
+                        Bot sources
+                      </div>
+                    ),
+                    children: (
+                      <div className="mt-3 flex flex-wrap gap-2">
+                        {props?.sources?.map((source, index) => (
+                          <button
+                            key={index}
+                            onClick={props.onSourceClick.bind(null, source)}
+                            className="inline-flex cursor-pointer transition-shadow !line-clamp-1 duration-300 ease-in-out hover:shadow-lg  items-center rounded-md bg-gray-200 p-1 text-xs text-gray-800 border border-gray-300 dark:bg-gray-800 dark:border-gray-700 dark:text-gray-100 opacity-80 hover:opacity-100 "
+                          >
+                            <span className="text-xs">
+                              {removeUUID(
+                                `${
+                                  source?.metadata?.path ||
+                                  source?.metadata?.source
+                                }`.replace("./uploads/", "")
+                              )}
+                            </span>
+                          </button>
+                        ))}
+                      </div>
+                    ),
+                  },
+                ]}
+              />
             )}
           </div>
 

diff --git a/app/ui/src/components/Bot/Settings/SettingsBody.tsx b/app/ui/src/components/Bot/Settings/SettingsBody.tsx
@@ -169,6 +169,7 @@ export const SettingsBody: React.FC<BotSettings> = ({
             autoResetSession: data.autoResetSession,
             inactivityTimeout: data.inactivityTimeout,
             autoSyncDataSources: data.autoSyncDataSources,
+            internetSearchEnabled: data.internetSearchEnabled,
           }}
           form={form}
           requiredMark={false}
@@ -226,7 +227,24 @@ export const SettingsBody: React.FC<BotSettings> = ({
                     options={chatModel}
                   />
                 </Form.Item>
-
+                <Form.Item
+                  label={"Embedding Method"}
+                  name="embedding"
+                  help={
+                    <>
+                      <p className="text-xs text-gray-500 dark:text-gray-400">
+                        If you change the embedding method, make sure to
+                        re-fetch the data source or choose a model with the same
+                        dimensions
+                      </p>
+                    </>
+                  }
+                >
+                  <Select
+                    placeholder="Select an embedding method"
+                    options={embeddingModel}
+                  />
+                </Form.Item>
                 <Form.Item
                   hasFeedback={!isStreamingSupported(currentModel)}
                   help={
@@ -266,23 +284,14 @@ export const SettingsBody: React.FC<BotSettings> = ({
                   />
                 </Form.Item>
 
+
                 <Form.Item
-                  label={"Embedding Method"}
-                  name="embedding"
-                  help={
-                    <>
-                      <p className="text-xs text-gray-500 dark:text-gray-400">
-                        If you change the embedding method, make sure to
-                        re-fetch the data source or choose a model with the same
-                        dimensions
-                      </p>
-                    </>
-                  }
+                  name="internetSearchEnabled"
+                  label="Enable internet search"
+                  valuePropName="checked"
+                  help="This is experimental and may not work as expected."
                 >
-                  <Select
-                    placeholder="Select an embedding method"
-                    options={embeddingModel}
-                  />
+                  <Switch />
                 </Form.Item>
 
                 <Form.Item
@@ -455,7 +464,7 @@ export const SettingsBody: React.FC<BotSettings> = ({
                   />
                 </Form.Item>
 
-              <Form.Item
+                <Form.Item
                   name="autoSyncDataSources"
                   label="Auto Sync Data Source(s)"
                   tooltip="This will automatically re-fetch the URL-based data sources at a certain interval."

diff --git a/app/ui/src/index.css b/app/ui/src/index.css
@@ -59,3 +59,7 @@ markdown-processing {
   border: none !important;
   box-shadow: none !important;
 }
+
+.ant-collapse-header {
+  padding:  0 !important;
+}
diff --git a/server/src/internet/index.ts b/server/src/internet/index.ts
@@ -116,6 +116,11 @@ const searchProviders = {
 };
 
 export const searchInternet = async (embedding: Embeddings, { query }: { query: string }) => {
+
+    if(process.env.DISABLE_INTERNET_SEARCH == "true") {
+        return [];
+    }
+
     const searchProvider = searchProviders[SERACH_PROVIDER];
     if (!searchProvider) {
         throw new Error(`Search provider ${SERACH_PROVIDER} not found`);

diff --git a/server/src/utils/hybrid.ts b/server/src/utils/hybrid.ts
@@ -3,6 +3,7 @@ import { PrismaClient } from "@prisma/client";
 import { Embeddings } from "langchain/embeddings/base";
 import { BaseRetriever, BaseRetrieverInput } from "@langchain/core/retrievers";
 import { CallbackManagerForRetrieverRun, Callbacks } from "langchain/callbacks";
+import { searchInternet } from "../internet";
 const prisma = new PrismaClient();
 export interface DialoqbaseLibArgs extends BaseRetrieverInput {
   botId: string;
@@ -47,7 +48,11 @@ export class DialoqbaseHybridRetrival extends BaseRetriever {
 
       const vector = `[${embeddedQuery.join(",")}]`;
       const bot_id = this.botId;
-
+      const botInfo = await prisma.bot.findFirst({
+        where: {
+          id: bot_id,
+        },
+      });
       const data = await prisma.$queryRaw`
      SELECT * FROM "similarity_search_v2"(query_embedding := ${vector}::vector, botId := ${bot_id}::text,match_count := ${k}::int)
     `;
@@ -62,8 +67,17 @@ export class DialoqbaseHybridRetrival extends BaseRetriever {
         resp.similarity * 10,
         resp.id,
       ]);
+      let internetSearchResults = [];
+      if (botInfo.internetSearchEnabled) {
+        internetSearchResults = await searchInternet(this.embeddings, {
+          query: query,
+        });
+      }
 
-      return result;
+      const combinedResults = [...result, ...internetSearchResults];
+      combinedResults.sort((a, b) => b[1] - a[1]);
+      const topResults = combinedResults.slice(0, k);
+      return topResults;
     } catch (e) {
       console.log(e);
       return [];