Update rag/RAG.py

rag/RAG.py

@@ -1,6 +1,7 @@
 from google import genai
 from google.genai import types
 import numpy as np
+from concurrent.futures import ThreadPoolExecutor, as_completed
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 import os
 from dotenv import load_dotenv
@@ -26,23 +27,32 @@ class RAG:
         except Exception as e: 
             raise ValueError(f"an error occured: {e}")
     
-    def generate_embedding(self,text,task_type=None):
+    def generate_embedding(self, text, task_type=None):
         try:
-            if(not task_type):
+            if not task_type:
                 task_type = self.TASK_TYPE
-            embeddings = []
+    
             chunks = self.split_text(text)
-            for i in range(0,len(chunks),self.MAX_BATCH_SIZE):
+            batches = [chunks[i:i + self.MAX_BATCH_SIZE] for i in range(0, len(chunks), self.MAX_BATCH_SIZE)]
+    
+            def embed_batch(batch):
                 response = client.models.embed_content(
                     model=self.MODEL,
-                    contents=chunks[i:i + self.MAX_BATCH_SIZE],
+                    contents=batch,
                     config=types.EmbedContentConfig(task_type=task_type)
                 )
-                for chunk_embedding in response.embeddings:
-                    embeddings.append(chunk_embedding.values)
+                return [embedding.values for embedding in response.embeddings]
+    
+            embeddings = []
+            with ThreadPoolExecutor(max_workers=100) as executor:
+                futures = [executor.submit(embed_batch, batch) for batch in batches]
+                for future in as_completed(futures):
+                    embeddings.extend(future.result())
+    
             return {"embeddings": embeddings, "chunks": chunks}, 200
-        except Exception as e: 
-            return {"an error occured": f"{e}"}, 500
+    
+        except Exception as e:
+            return {"an error occurred": str(e)}, 500