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)