Update pipeline.py

pipeline.py

@@ -16,9 +16,6 @@ from refusal_chain import get_refusal_chain
 from tailor_chain import get_tailor_chain
 from cleaner_chain import get_cleaner_chain, CleanerChain
 
-# We also import the relevant RAG logic here or define it directly
-# (We define build_rag_chain in this file for clarity)
-
 # 1) Environment: set up keys if missing
 if not os.environ.get("GEMINI_API_KEY"):
     os.environ["GEMINI_API_KEY"] = getpass.getpass("Enter your Gemini API Key: ")
@@ -55,8 +52,52 @@ def extract_main_topic(query: str) -> str:
     return main_topic if main_topic else "this topic"
 
 # 3) build_or_load_vectorstore (no changes)
+def build_or_load_vectorstore(csv_path: str, store_dir: str) -> FAISS:
+    if os.path.exists(store_dir):
+        print(f"DEBUG: Found existing FAISS store at '{store_dir}'. Loading...")
+        embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/multi-qa-mpnet-base-dot-v1")
+        vectorstore = FAISS.load_local(store_dir, embeddings)
+        return vectorstore
+    else:
+        print(f"DEBUG: Building new store from CSV: {csv_path}")
+        df = pd.read_csv(csv_path)
+        df = df.loc[:, ~df.columns.str.contains('^Unnamed')]
+        df.columns = df.columns.str.strip()
+        if "Answer" in df.columns:
+            df.rename(columns={"Answer": "Answers"}, inplace=True)
+        if "Question" not in df.columns and "Question " in df.columns:
+            df.rename(columns={"Question ": "Question"}, inplace=True)
+        if "Question" not in df.columns or "Answers" not in df.columns:
+            raise ValueError("CSV must have 'Question' and 'Answers' columns.")
+        docs = []
+        for _, row in df.iterrows():
+            q = str(row["Question"])
+            ans = str(row["Answers"])
+            doc = Document(page_content=ans, metadata={"question": q})
+            docs.append(doc)
+        embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/multi-qa-mpnet-base-dot-v1")
+        vectorstore = FAISS.from_documents(docs, embedding=embeddings)
+        vectorstore.save_local(store_dir)
+        return vectorstore
 
 # 4) Build RAG chain for Gemini (no changes)
+def build_rag_chain(llm_model: LiteLLMModel, vectorstore: FAISS) -> RetrievalQA:
+    class GeminiLangChainLLM(LLM):
+        def _call(self, prompt: str, stop: Optional[list] = None, **kwargs) -> str:
+            messages = [{"role": "user", "content": prompt}]
+            return llm_model(messages, stop_sequences=stop)
+        @property
+        def _llm_type(self) -> str:
+            return "custom_gemini"
+    retriever = vectorstore.as_retriever(search_type="similarity", search_kwargs={"k": 3})
+    gemini_as_llm = GeminiLangChainLLM()
+    rag_chain = RetrievalQA.from_chain_type(
+        llm=gemini_as_llm,
+        chain_type="stuff",
+        retriever=retriever,
+        return_source_documents=True
+    )
+    return rag_chain
 
 # 5) Initialize all the separate chains
 classification_chain = get_classification_chain()
@@ -78,6 +119,16 @@ wellness_rag_chain = build_rag_chain(gemini_llm, wellness_vectorstore)
 brand_rag_chain = build_rag_chain(gemini_llm, brand_vectorstore)
 
 # 7) Tools / Agents for web search (no changes)
+search_tool = DuckDuckGoSearchTool()
+web_agent = CodeAgent(tools=[search_tool], model=gemini_llm)
+managed_web_agent = ManagedAgent(agent=web_agent, name="web_search", description="Runs web search for you.")
+manager_agent = CodeAgent(tools=[], model=gemini_llm, managed_agents=[managed_web_agent])
+
+def do_web_search(query: str) -> str:
+    print("DEBUG: Attempting web search for more info...")
+    search_query = f"Give me relevant info: {query}"
+    response = manager_agent.run(search_query)
+    return response
 
 # 8) Orchestrator: run_with_chain
 def run_with_chain(query: str) -> str: