File size: 10,587 Bytes
628936c
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
25c6c9c
628936c
 
 
 
 
 
 
 
 
 
 
 
 
25c6c9c
628936c
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
import os
import gradio as gr
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.document_loaders import PyPDFLoader, WebBaseLoader
from langchain_community.tools.tavily_search import TavilySearchResults
from langchain_community.vectorstores import SKLearnVectorStore
from langchain_openai import ChatOpenAI
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_pinecone import PineconeVectorStore
from langchain.prompts import PromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from pydantic import BaseModel, Field
from typing import List, TypedDict, Optional
from langchain.schema import Document
from langgraph.graph import START, END, StateGraph
from dotenv import load_dotenv

load_dotenv()

url = [
    "https://www.investopedia.com/",
    "https://www.fool.com/",
    "https://www.morningstar.com/",
    "https://www.kiplinger.com/",
    "https://www.nerdwallet.com/"
]

# Initialize Embedding and Vector DB
embedding_model = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")

# Initialize Pinecone connection
try:
    pc = PineconeVectorStore(
        pinecone_api_key=os.environ.get('PINCE_CONE_LIGHT'),
        embedding=embedding_model,
        index_name='rag-rubic',
        namespace='vectors_lightmodel'
    )
    retriever = pc.as_retriever(search_kwargs={"k": 10})
except Exception as e:
    print(f"Pinecone connection error: {e}")
    # Fallback to SKLearn vector store if Pinecone fails
    retriever = None

# Initialize the LLM
llm = ChatOpenAI(
    model='gpt-4o-mini',
    api_key=os.environ.get('OPEN_AI_KEY'),
    temperature=0.2
)

# Schema for grading documents
class GradeDocuments(BaseModel):
    binary_score: str = Field(description="Documents are relevant to the question, 'yes' or 'no'")

structured_llm_grader = llm.with_structured_output(GradeDocuments)

# Define System and Grading prompt
system = """You are a grader assessing relevance of a retrieved document to a user question. 
    If the document contains keyword(s) or semantic meaning related to the question, grade it as relevant. 
    Give a binary score 'yes' or 'no' score to indicate whether the document is relevant to the question."""

grade_prompt = ChatPromptTemplate.from_messages([
    ("system", system),
    ("human", "Retrieved document: \n\n {documents} \n\n User question: {question}")
])

retrieval_grader = grade_prompt | structured_llm_grader

# RAG Prompt template
prompt = PromptTemplate(
    template='''
    You are a Registered Investment Advisor with expertise in Indian financial markets and client relations.
    You must understand what the user is asking about their financial investments and respond to their queries based on the information in the documents only.

    Use the following documents to answer the question. If you do not know the answer, say you don't know.

    Query: {question}
    Documents: {context}
    ''',
    input_variables=['question', 'context']
)

rag_chain = prompt | llm | StrOutputParser()

# Web search tool for adding data from websites
web_search_tool = TavilySearchResults(api_key=os.environ.get('TAVILY_API_KEY'), k=5)

# Load website data
try:
    print("Loading web data...")
    docs = []
    for i in url:
        try:
            docs.append(WebBaseLoader(i).load())
        except Exception as e:
            print(f"Error loading {i}: {e}")
    
    docs_list = [item for sublist in docs for item in sublist]

    # Split documents into chunks
    text_splitter = RecursiveCharacterTextSplitter.from_tiktoken_encoder(
        chunk_size=1000,
        chunk_overlap=100
    )
    doc_splits = text_splitter.split_documents(docs_list)

    # VectorStore from the web-scraped documents
    vectorstore = SKLearnVectorStore.from_documents(
        documents=doc_splits,
        embedding=embedding_model
    )
    retriever_web = vectorstore.as_retriever(search_kwargs={"k": 5})
    print(f"Loaded {len(doc_splits)} document chunks from web sources")
except Exception as e:
    print(f"Error in web data processing: {e}")
    # Create a simple retriever that returns empty results if web loading fails
    retriever_web = lambda x: []

# Define Graph states and transitions
class GraphState(TypedDict):
    question: str
    generation: Optional[str]
    need_web_search: Optional[str]  # Changed from 'web_search' to 'need_web_search'
    documents: List

def retrieve_db(state):
    """Gather data for the query."""
    question = state['question']
    if retriever:
        try:
            results = retriever.invoke(question)
            return {'documents': results, 'question': question}
        except Exception as e:
            print(f"Retriever error: {e}")
    
    # If retriever fails or doesn't exist, return empty documents
    return {'documents': [], 'question': question, 'need_web_search': 'yes'}

def grade_docs(state):
    """Grades the docs generated by the retriever_db"""
    question = state['question']
    docs = state['documents']
    
    if not docs:
        return {"documents": [], 'question': question, 'need_web_search': 'yes'}
    
    filtered_data = []
    web_search_needed = "no"
    
    try:
        for doc in docs:
            doc_content = doc.page_content if hasattr(doc, 'page_content') else str(doc)
            score = retrieval_grader.invoke({'question': question, 'documents': doc_content})
            grade = score.binary_score
            if grade == 'yes':
                filtered_data.append(doc)
    except Exception as e:
        print(f"Error in document grading: {e}")
        web_search_needed = "yes"
    
    # If no relevant documents were found, trigger web search
    if not filtered_data:
        web_search_needed = "yes"
        
    return {
        "documents": filtered_data,
        'question': question,
        'need_web_search': web_search_needed  # Updated key name
    }

def decide(state):
    """Decide if the generation should be based on DB or web search DATA"""
    web = state.get('need_web_search', 'no')  # Updated key name
    if web == 'yes':
        return 'web_search'
    else:
        return 'generate'

def web_search(state):
    """Based on the Grade, will proceed with WebSearch within the given URL's."""
    question = state['question']
    documents = state.get("documents", [])
    
    try:
        # First try website-specific retriever
        docs = retriever_web.invoke(question)
        if not docs:
            # If no results, try Tavily search
            search_results = web_search_tool.invoke(question)
            data = "\n".join(result["content"] for result in search_results)
            docs = [Document(page_content=data)]
    except Exception as e:
        print(f"Web search error: {e}")
        # Create a fallback document if search fails
        docs = [Document(page_content="Unable to retrieve additional information.")]
    
    # Combine existing documents with new ones
    all_docs = documents + docs
    
    return {'documents': all_docs, 'question': question}

def generate(state):
    """Generate response based on retrieved documents"""
    documents = state.get('documents', [])
    question = state['question']
    
    # Convert documents to text for the context
    if documents:
        try:
            context = "\n\n".join(
                doc.page_content if hasattr(doc, 'page_content') else str(doc) 
                for doc in documents
            )
        except Exception as e:
            print(f"Error processing documents: {e}")
            context = "Error retrieving relevant information."
    else:
        context = "No relevant information found."
    
    try:
        response = rag_chain.invoke({'context': context, 'question': question})
    except Exception as e:
        print(f"Generation error: {e}")
        response = "I apologize, but I encountered an error while generating a response. Please try asking your question again."

    return {
        'documents': documents,
        'question': question,
        'generation': response
    }

# Compile Workflow
workflow = StateGraph(GraphState)
workflow.add_node("retrieve", retrieve_db)
workflow.add_node("grader", grade_docs)
workflow.add_node("web_search", web_search)  # Now this won't conflict with the state key
workflow.add_node("generate", generate)

workflow.add_edge(START, "retrieve")
workflow.add_edge("retrieve", "grader")
workflow.add_conditional_edges(
    "grader", 
    decide,
    {
        'web_search': 'web_search',
        'generate': 'generate'
    },
)
workflow.add_edge("web_search", "generate")
workflow.add_edge("generate", END)

# Compile the graph
crag = workflow.compile()

# Define Gradio Interface with proper chat history management
def process_query(user_input, history):
    # Initialize history if it's None
    if history is None:
        history = []
    
    # Add user input to history
    history.append((user_input, ""))
    
    # Process the query
    inputs = {"question": user_input}
    response = ""
    
    try:
        # Execute the graph
        result = crag.invoke(inputs)
        if result and 'generation' in result:
            response = result['generation']
        else:
            response = "I couldn't find relevant information to answer your question."
    except Exception as e:
        print(f"Error in crag execution: {e}")
        response = "I encountered an error while processing your request. Please try again."
    
    # Update the last response in history
    history[-1] = (user_input, response)
    
    return history, ""

# Gradio Interface
with gr.Blocks() as demo:
    gr.Markdown("# 🤖 RAG-Powered Financial Advisor Chatbot")
    
    chatbot = gr.Chatbot(
        [],
        elem_id="chatbot",
        bubble_full_width=False,
        height=600,
        avatar_images=(None, "🤖")
    )
    
    with gr.Row():
        msg = gr.Textbox(
            placeholder="Ask me anything about Indian financial markets...",
            label="Your question:",
            scale=9
        )
        submit_btn = gr.Button("Send", scale=1)
    
    clear_btn = gr.Button("Clear Chat")
    
    # Set up event handlers
    submit_click_event = submit_btn.click(
        process_query,
        inputs=[msg, chatbot],
        outputs=[chatbot, msg]
    )
    
    msg.submit(
        process_query,
        inputs=[msg, chatbot],
        outputs=[chatbot, msg]
    )
    
    clear_btn.click(lambda: [], outputs=[chatbot])


if __name__ == "__main__":
    demo.launch()