app.py

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