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()