Add Hybrid Search

A new search endpoint and ui section that combines full
text search with vector search using Reciprocal rank fusion

Author: APIEngineering

builder/partnerproduct/src/config.yaml

@@ -14,6 +14,7 @@ vector_store:
     numCandidates: 150
     minScore: 0.1 
     vectorSearchIndexName: 'vector_index'
+    textSearchIndexName: 'text_index'
 llms:
     class_name: Fireworks
     model_name: 'accounts/fireworks/models/mixtral-8x22b-instruct'

builder/partnerproduct/src/loader.ts

@@ -26,10 +26,11 @@ try {
             chunksAdded += chunks.entriesAdded;
         });
     }
-    
+
     if (chunksAdded > 0) {
         console.log(`\n Total documents added : ${chunksAdded} `)
         await llmApplication.createVectorIndex();
+        await llmApplication.createTextIndex();
     }
     else {
         console.log("\n-- Data not inserted, please retry --")

builder/partnerproduct/src/semantic-search.ts

@@ -14,7 +14,7 @@ const app = express();
 const port = 9001;
 
 app.use(express.json());
-app.use(cors()); 
+app.use(cors());
 
 const llmApplication = await new RAGApplicationBuilder()
   .setModel(getModelClass())
@@ -42,6 +42,30 @@ app.get('/semantic-search', async (req: Request, res: Response) => {
   }
 });
 
+app.get('/hybrid-search', async (req: Request, res: Response) => {
+  try {
+    const userQuery = asString(req.query.query);
+    const vectorWeight = asFloat(req.query.vectorWeight ?? 0.5);
+    const fullTextWeight = asFloat(req.query.fullTextWeight ?? 0.5);
+
+    if (!userQuery) {
+      return res.status(400).send('Query is required');
+    }
+
+    llmApplication.hybridQuery(userQuery, vectorWeight, fullTextWeight).then((result) => {
+      console.log('Result:', result);
+      res.send(result);
+    });
+
+  } catch (error) {
+    console.error('Error during hybrid vector search:', error);
+    res.status(500).send('An error occurred while processing your request.');
+  }
+});
+
 app.listen(port, () => {
   console.log(`Server is running on http://localhost:${port}`);
-});
+});
+
+function asString(value: any): string { return typeof value !== 'undefined' ? value.toString() : ''; }
+function asFloat(value: any): number { return typeof value !== 'undefined' ? parseFloat(value.toString()) || 0 : 0; }

builder/partnerproduct/ui/package.json

@@ -6,6 +6,7 @@
     "@leafygreen-ui/card": "^11.0.0",
     "@leafygreen-ui/icon": "^12.6.0",
     "@leafygreen-ui/loading-indicator": "^2.0.12",
+    "@leafygreen-ui/number-input": "^2.2.1",
     "@leafygreen-ui/search-input": "^3.1.2",
     "@leafygreen-ui/side-nav": "^14.1.3",
     "@leafygreen-ui/tabs": "^13.0.1",

builder/partnerproduct/ui/src/App.js

@@ -49,13 +49,19 @@ function App() {
                   Vector Search
                 </SideNavItem>
               </Link>
+              <Link onClick={handleTabClick(3)} className="link" to="/hybrid-search">
+                <SideNavItem active={activeTab === 3} glyph={<Icon glyph="Wizard" />} >
+                  Hybrid Search
+                </SideNavItem>
+              </Link>
             </SideNavGroup>
           </SideNav>
           <div className="main-content">
             <Routes>
               <Route path="/" element={<Hero />} />
               <Route path="/rag-chatbot" element={<ChatModule />} />
               <Route path="/search" element={<Search />} />
+              <Route path="/hybrid-search" element={<Search hybrid="true" />} />
             </Routes>
           </div>
         </div>
@@ -64,4 +70,4 @@ function App() {
   );
 }
 
-export default App;
+export default App;

builder/partnerproduct/ui/src/modules/search/search.js

@@ -6,21 +6,27 @@ import Card from '@leafygreen-ui/card';
 import { PageLoader, Spinner } from "@leafygreen-ui/loading-indicator";
 import './search.css';
 import { H1, H2 } from '@leafygreen-ui/typography';
+import { NumberInput } from '@leafygreen-ui/number-input';
 
-function Search() {
+function Search({ hybrid }) {
+  const isHybridSearch = hybrid
   const [searchValue, setSearchValue] = useState('');
+  const [vectorWeight, setVectorWeight] = useState(0.5);
+  const [textWeight, setTextWeight] = useState(0.5);
   const [response, setResponse] = useState([]);
   const [loading, setLoading] = useState(false);
 
   const handleInputChange = (event) => {
     setSearchValue(event.target.value);
   };
 
+  const searchEndpoint = `http://localhost:9001/${isHybridSearch ? 'hybrid-search' : 'semantic-search'}`;
+
   const searchResults = () => {
     console.log(searchValue);
     setLoading(true);
     const query = encodeURIComponent(searchValue);
-    fetch(`http://localhost:9001/semantic-search?query=${query}`, {
+    fetch(`${searchEndpoint}?query=${query}&vectorWeight=${vectorWeight}&fullTextWeight=${textWeight}`, {
       method: 'GET',
       headers: {
         'Content-Type': 'application/json'
@@ -38,20 +44,37 @@ function Search() {
 
   return (
     <div className="search-bar">
-      <H1>Vector Search</H1>  
+      <H1>{ isHybridSearch ? "Hybrid Search" : "Vector Search"}</H1>
       <header className="search-header">
         <SearchInput
           value={searchValue}
           onChange={handleInputChange}
           onSubmit={searchResults}
         />
+        { isHybridSearch && <>
+          <NumberInput
+            label='Vector Weight'
+            description='Affects the weighted reciprocal rank'
+            darkMode="true"
+            value={vectorWeight}
+            onChange={(e) => setVectorWeight(e.target.value)}
+          />
+          <NumberInput
+            label='Text Weight'
+            description='Affects the weighted reciprocal rank'
+            darkMode="true"
+            value={textWeight}
+            onChange={(e) => setTextWeight(e.target.value)}
+          />
+        </> }
       </header>
       <div className='results'>
         <>
           {loading ? <PageLoader /> : <div>
             {response.map((item, index) => (
               <Card key={index} className="card-styles" as="article">
                 <h3>Score: {item.score}</h3>
+                { isHybridSearch && <pre>Text Score: {item.fts_score}, Vector Score: {item.vs_score}</pre> }
                 <p><strong>Content:</strong> {item.pageContent}</p>
                 <div>
                   <h4>Metadata:</h4>
@@ -66,4 +89,4 @@ function Search() {
   );
 }
 
-export default Search;
+export default Search;

src/core/rag-application.ts

@@ -216,6 +216,12 @@ export class RAGApplication {
         return await this.vectorDb.similaritySearch(queryEmbedded, this.searchResultCount);
     }
 
+    public async hybridQuery(query: string, vectorWeight = 0.1, fullTextWeight = 0.9) {
+        const cleanQuery = cleanString(query);
+        const queryEmbedded = await RAGEmbedding.getEmbedding().embedQuery(cleanQuery);
+        return await this.vectorDb.hybridSearch(query, queryEmbedded, this.searchResultCount, vectorWeight, fullTextWeight);
+    }
+
     public async getContext(query: string) {
         const cleanQuery = cleanString(query);
         const rawContext = await this.getEmbeddings(cleanQuery);
@@ -224,7 +230,7 @@ export class RAGApplication {
 
 
     public async getQueryContext(cleanQuery: string, aggregatePipelineName: string) {
-        //TODO: Method override. Create a MQL query with user prompts using LLM. 
+        //TODO: Method override. Create a MQL query with user prompts using LLM.
         // Generate output query with the user prompt and the context.
         let mqlQuery = await this.dbLookup.get(aggregatePipelineName).aggregateQuery;
         mqlQuery = JSON.stringify(mqlQuery);
@@ -263,7 +269,7 @@ export class RAGApplication {
             // If there is not context lookup provided
             sources = [];
         }
-        
+
         return {
             sources,
             result: await this.model.query(this.queryTemplate, userQuery, context, conversationId),
@@ -274,11 +280,15 @@ export class RAGApplication {
         await this.vectorDb.createVectorIndex(RAGEmbedding.getEmbedding().getDimensions());
     }
 
+    public async createTextIndex() {
+        await this.vectorDb.createTextIndex();
+    }
+
     public async docsCount() : Promise<number> {
         return await this.vectorDb.docsCount();
     }
 
     public getDb(key: string): Map<string, any> {
         return this.dbLookup.get(key);
     }
-}
+}

src/interfaces/base-db.ts

@@ -4,8 +4,10 @@ export interface BaseDb {
     init({}: { dimensions: number }): Promise<void>;
     insertChunks(chunks: InsertChunkData[]): Promise<number>;
     similaritySearch(query: number[], k: number): Promise<ExtractChunkData[]>;
+    hybridSearch(textQuery: string, query: number[], k: number, vectorWeight, fullTextWeight): Promise<ExtractChunkData[]>
     getVectorCount(): Promise<number>;
     createVectorIndex(numDimensions: number): Promise<void>;
+    createTextIndex(): Promise<void>;
     docsCount(): Promise<number>;
 
     deleteKeys(uniqueLoaderId: string): Promise<boolean>;