app.py

import os
from typing import List, Dict, Any, Optional, TypedDict
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.document_loaders import PyPDFLoader
from langchain_community.vectorstores import Chroma
from langchain_google_genai import GoogleGenerativeAIEmbeddings, ChatGoogleGenerativeAI
from langchain.prompts import ChatPromptTemplate
from langgraph.graph import StateGraph, END
from langgraph.checkpoint.memory import MemorySaver
import chromadb
import gradio as gr

GOOGLE_API_KEY = "AIzaSyDERyKbxN9Da7eMfkO0zw3b4-qCH715h24"
os.environ["GOOGLE_API_KEY"] = "AIzaSyDERyKbxN9Da7eMfkO0zw3b4-qCH715h24"

class AgentState(TypedDict):
    query: str
    documents: List[Dict[str, Any]]
    context: str
    answer: str
    sources: List[str]
    error: Optional[str]

class ResearchAgent:
    def __init__(self, api_key):
        self.api_key = api_key
        self.chroma_client = chromadb.PersistentClient(path="/tmp/chroma_db")
        try:
            self.collection = self.chroma_client.get_collection(name="docs")
        except:
            self.collection = self.chroma_client.create_collection(name="docs")
        self.embeddings = GoogleGenerativeAIEmbeddings(model="models/text-embedding-004", google_api_key=self.api_key)
        self.llm = ChatGoogleGenerativeAI(model="gemini-1.5-flash", google_api_key=self.api_key, temperature=0.1)
        self.vector_store = Chroma(client=self.chroma_client, collection_name="docs", embedding_function=self.embeddings)
        self.documents = []
        self.setup_workflow()

    def setup_workflow(self):
        def retrieve_docs(state):
            query = state["query"]
            docs = self.vector_store.similarity_search(query, k=5)
            context_parts = []
            sources = []
            for doc in docs:
                context_parts.append(doc.page_content)
                source = f"📄 {doc.metadata.get('source', 'Unknown')} (Page {doc.metadata.get('page', '?')})"
                if source not in sources:
                    sources.append(source)
            context = "\n\n".join(context_parts)
            return {**state, "context": context, "sources": sources, "documents": [{"content": doc.page_content, "metadata": doc.metadata} for doc in docs]}

        def make_answer(state):
            query = state["query"]
            context = state["context"]
            if not context:
                return {**state, "answer": "No documents found. Upload some PDFs first."}
            prompt = ChatPromptTemplate.from_messages([("system", "Answer based on the context: {context}"), ("human", "{question}")])
            chain = prompt | self.llm
            response = chain.invoke({"context": context, "question": query})
            return {**state, "answer": response.content}

        def check_docs(state):
            if not self.documents:
                return {**state, "answer": "No documents uploaded yet. Please upload PDFs first.", "sources": []}
            return state

        def continue_or_end(state):
            if not self.documents:
                return "end"
            return "retrieve"

        workflow = StateGraph(AgentState)
        workflow.add_node("check", check_docs)
        workflow.add_node("retrieve", retrieve_docs)
        workflow.add_node("generate", make_answer)
        workflow.set_entry_point("check")
        workflow.add_conditional_edges("check", continue_or_end, {"end": END, "retrieve": "retrieve"})
        workflow.add_edge("retrieve", "generate")
        workflow.add_edge("generate", END)
        memory = MemorySaver()
        self.workflow = workflow.compile(checkpointer=memory)

    def add_pdf(self, pdf_path):
        try:
            loader = PyPDFLoader(pdf_path)
            pages = loader.load()
            text_splitter = RecursiveCharacterTextSplitter(chunk_size=1024, chunk_overlap=200)
            documents = []
            for page in pages:
                chunks = text_splitter.split_documents([page])
                for chunk in chunks:
                    chunk.metadata.update({"source": os.path.basename(pdf_path), "file_path": pdf_path})
                    documents.append(chunk)
            if documents:
                self.vector_store.add_documents(documents)
                self.documents.extend(documents)
                return f"✅ Added {len(documents)} chunks from {os.path.basename(pdf_path)}"
            else:
                return f"⚠️ No content found in {os.path.basename(pdf_path)}"
        except Exception as e:
            return f"Error: {str(e)}"

    def ask_question(self, question):
        try:
            state = {"query": question, "documents": [], "context": "", "answer": "", "sources": [], "error": None}
            config = {"configurable": {"thread_id": "main"}}
            result = self.workflow.invoke(state, config)
            return result["answer"], result.get("sources", [])
        except Exception as e:
            return f"Error: {str(e)}", []

    def get_stats(self):
        if not self.documents:
            return "No documents loaded."
        source_counts = {}
        for doc in self.documents:
            source = doc.metadata.get('source', 'Unknown')
            source_counts[source] = source_counts.get(source, 0) + 1
        stats = f"📊 Total chunks: {len(self.documents)}\n"
        stats += f"📁 Total files: {len(source_counts)}\n\n"
        for source, count in source_counts.items():
            stats += f"- {source}: {count} chunks\n"
        return stats

agent = ResearchAgent(GOOGLE_API_KEY)

def upload_file(file):
    if file is None:
        return "❌ No file selected.", ""
    result = agent.add_pdf(file.name)
    stats = agent.get_stats()
    return result, stats

def ask_question_ui(question, history):
    if not question.strip():
        return history
    answer, sources = agent.ask_question(question)
    formatted_answer = answer
    if sources:
        formatted_answer += "\n\n**Sources:**\n" + "\n".join(sources)
    history.append([question, formatted_answer])
    return history

def clear_history():
    return []

with gr.Blocks() as demo:
    gr.Markdown("# 🔬 PDF Research Assistant")
    with gr.Tab("📁 Upload Files"):
        with gr.Row():
            with gr.Column():
                file_upload = gr.File(label="Choose PDF File", file_types=[".pdf"], type="filepath")
                upload_btn = gr.Button("Upload PDF")
            with gr.Column():
                upload_status = gr.Textbox(label="Status", interactive=False, max_lines=3)
        doc_stats = gr.Markdown("No files uploaded yet.")
        upload_btn.click(upload_file, inputs=[file_upload], outputs=[upload_status, doc_stats])
    with gr.Tab("💬 Ask Questions"):
        chatbot = gr.Chatbot(label="Chat with your documents", height=400)
        with gr.Row():
            question_input = gr.Textbox(label="Your Question", placeholder="Ask something about your PDFs...", scale=4)
            ask_btn = gr.Button("Ask", variant="primary", scale=1)
        clear_btn = gr.Button("Clear Chat")
        ask_btn.click(ask_question_ui, inputs=[question_input, chatbot], outputs=[chatbot]).then(lambda: "", outputs=[question_input])
        question_input.submit(ask_question_ui, inputs=[question_input, chatbot], outputs=[chatbot]).then(lambda: "", outputs=[question_input])
        clear_btn.click(clear_history, outputs=[chatbot])

demo.launch(server_name="0.0.0.0", server_port=7860)