File size: 8,118 Bytes
b0c64f6 99474e2 b0c64f6 99474e2 b0c64f6 99474e2 b0c64f6 99474e2 b0c64f6 99474e2 b0c64f6 99474e2 b0c64f6 99474e2 b0c64f6 99474e2 b0c64f6 99474e2 b0c64f6 99474e2 b0c64f6 99474e2 b0c64f6 99474e2 b0c64f6 99474e2 b0c64f6 99474e2 b0c64f6 99474e2 b0c64f6 99474e2 b0c64f6 99474e2 b0c64f6 99474e2 b0c64f6 99474e2 b0c64f6 99474e2 b0c64f6 99474e2 b0c64f6 99474e2 b0c64f6 99474e2 b0c64f6 99474e2 b0c64f6 99474e2 b0c64f6 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 |
# pipeline.py
import os
import getpass
import pandas as pd
from typing import Optional, List
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
# Import your classification/refusal/tailor/cleaner chains
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
# For RAG chain building
from langchain.llms.base import LLM
###############################################################################
# 1) Environment: set up 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...")
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 for Gemini
###############################################################################
def build_rag_chain(llm_model: LiteLLMModel, vectorstore: FAISS) -> RetrievalQA:
class GeminiLangChainLLM(LLM):
def _call(self, prompt: str, stop: Optional[list] = None, **kwargs) -> str:
# We'll treat the entire prompt as 'user' content
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 your 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 Chains
###############################################################################
wellness_csv = "AIChatbot.csv"
brand_csv = "BrandAI.csv"
wellness_store_dir = "faiss_wellness_store"
brand_store_dir = "faiss_brand_store"
gemini_llm = LiteLLMModel(model_id="gemini/gemini-pro", api_key=os.environ.get("GEMINI_API_KEY"))
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: run_with_chain_context
###############################################################################
def run_with_chain_context(query: str, chat_history: list) -> str:
"""
Like run_with_chain, but also references `chat_history`.
We'll do single-turn classification, but pass chat_history
to the RAG chain if needed.
Example usage:
chat_history = []
question = "What is Self-Reflection?"
resp1 = run_with_chain_context(question, chat_history)
# then chat_history.extend([...]) with HumanMessage/AIMessage
"""
print("DEBUG: Starting run_with_chain_context...")
# 1) Classification (no multi-turn, just single-turn classification)
class_result = classification_chain.invoke({"query": query})
classification = class_result.get("text", "").strip()
print("DEBUG: Classification =>", classification)
# 2) If OutOfScope => refusal => tailor => return
if classification == "OutOfScope":
refusal_text = refusal_chain.run({})
final_refusal = tailor_chain.run({"response": refusal_text})
return final_refusal.strip()
# 3) If Wellness => call wellness_rag_chain with chat_history
if classification == "Wellness":
# pass the conversation to .invoke(...) so it can see it if needed
rag_result = wellness_rag_chain.invoke({
"input": query,
"chat_history": chat_history # pass the entire list of prior messages
})
csv_answer = rag_result["result"].strip()
if not csv_answer:
web_answer = do_web_search(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(query)
else:
web_answer = ""
final_merged = cleaner_chain.merge(kb=csv_answer, web=web_answer)
final_answer = tailor_chain.run({"response": final_merged})
return final_answer.strip()
# 4) If Brand => brand_rag_chain with chat_history
if classification == "Brand":
rag_result = brand_rag_chain.invoke({
"input": 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})
return final_answer.strip()
# fallback => refusal
refusal_text = refusal_chain.run({})
final_refusal = tailor_chain.run({"response": refusal_text})
return final_refusal.strip()
|