pipeline.py

# pipeline.py

import os
import getpass
import pandas as pd
from typing import Optional, Dict, Any

# (Optional) from langchain.schema import RunnableConfig
# If you have the latest "langchain_core", use from langchain_core.runnables.base import Runnable
# or from langchain.runnables.base import Runnable (depending on your version)
from langchain.runnables.base import Runnable

from langchain.docstore.document import Document
from langchain.embeddings import HuggingFaceEmbeddings
from langchain.vectorstores import FAISS
from langchain.chains import RetrievalQA

from smolagents import CodeAgent, DuckDuckGoSearchTool, ManagedAgent, LiteLLMModel
import litellm

# Classification/Refusal/Tailor/Cleaner
from classification_chain import get_classification_chain
from refusal_chain import get_refusal_chain
from tailor_chain import get_tailor_chain
from cleaner_chain import get_cleaner_chain

from langchain.llms.base import LLM

###############################################################################
# 1) Environment keys
###############################################################################
if not os.environ.get("GEMINI_API_KEY"):
    os.environ["GEMINI_API_KEY"] = getpass.getpass("Enter your Gemini API Key: ")
if not os.environ.get("GROQ_API_KEY"):
    os.environ["GROQ_API_KEY"] = getpass.getpass("Enter your GROQ API Key: ")

###############################################################################
# 2) Build or load VectorStore
###############################################################################
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 from disk.")
        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

###############################################################################
# 3) Build RAG chain
###############################################################################
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

###############################################################################
# 4) Initialize sub-chains
###############################################################################
classification_chain = get_classification_chain()
refusal_chain = get_refusal_chain()
tailor_chain = get_tailor_chain()
cleaner_chain = get_cleaner_chain()

###############################################################################
# 5) Build vectorstores & RAG
###############################################################################
gemini_llm = LiteLLMModel(model_id="gemini/gemini-pro", api_key=os.environ.get("GEMINI_API_KEY"))

wellness_csv = "AIChatbot.csv"
brand_csv = "BrandAI.csv"
wellness_store_dir = "faiss_wellness_store"
brand_store_dir = "faiss_brand_store"

wellness_vectorstore = build_or_load_vectorstore(wellness_csv, wellness_store_dir)
brand_vectorstore = build_or_load_vectorstore(brand_csv, brand_store_dir)

wellness_rag_chain = build_rag_chain(gemini_llm, wellness_vectorstore)
brand_rag_chain = build_rag_chain(gemini_llm, brand_vectorstore)

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

###############################################################################
# 6) Orchestrator function: returns a dict => {"answer": "..."}
###############################################################################
def run_with_chain_context(inputs: Dict[str, Any]) -> Dict[str, str]:
    """
    Called by the Runnable. 
    inputs: { "input": <user_query>, "chat_history": <list of messages> (optional) }
    Output: { "answer": <final string> }
    """
    user_query = inputs["input"]
    chat_history = inputs.get("chat_history", [])

    # 1) Classification
    class_result = classification_chain.invoke({"query": user_query})
    classification = class_result.get("text", "").strip()

    if classification == "OutOfScope":
        refusal_text = refusal_chain.run({})
        final_refusal = tailor_chain.run({"response": refusal_text})
        return {"answer": final_refusal.strip()}

    if classification == "Wellness":
        rag_result = wellness_rag_chain.invoke({"input": user_query, "chat_history": chat_history})
        csv_answer = rag_result["result"].strip()
        if not csv_answer:
            web_answer = do_web_search(user_query)
        else:
            lower_ans = csv_answer.lower()
            if any(phrase in lower_ans for phrase in ["i do not know", "not sure", "no context", "cannot answer"]):
                web_answer = do_web_search(user_query)
            else:
                web_answer = ""

        final_merged = cleaner_chain.merge(kb=csv_answer, web=web_answer)
        final_answer = tailor_chain.run({"response": final_merged}).strip()
        return {"answer": final_answer}

    if classification == "Brand":
        rag_result = brand_rag_chain.invoke({"input": user_query, "chat_history": chat_history})
        csv_answer = rag_result["result"].strip()
        final_merged = cleaner_chain.merge(kb=csv_answer, web="")
        final_answer = tailor_chain.run({"response": final_merged}).strip()
        return {"answer": final_answer}

    # fallback
    refusal_text = refusal_chain.run({})
    final_refusal = tailor_chain.run({"response": refusal_text}).strip()
    return {"answer": final_refusal}


###############################################################################
# 7) Build a "Runnable" wrapper so .with_listeners() works
###############################################################################
from langchain.runnables.base import Runnable

class PipelineRunnable(Runnable[Dict[str, Any], Dict[str, str]]):
    """
    Wraps run_with_chain_context(...) in a Runnable 
    so that RunnableWithMessageHistory can attach listeners.
    """
    def invoke(self, input: Dict[str, Any], config: Optional[Any] = None) -> Dict[str, str]:
        return run_with_chain_context(input)

# Export an instance of PipelineRunnable for use in my_memory_logic.py
pipeline_runnable = PipelineRunnable()