import gradio as gr import pdfplumber import os from langchain.schema import Document from langchain.text_splitter import RecursiveCharacterTextSplitter from langchain_community.embeddings import OpenAIEmbeddings from langchain.vectorstores import Pinecone import pinecone import pandas as pd import time from pinecone.grpc import PineconeGRPC as Pinecone from pinecone import ServerlessSpec from langchain_pinecone import PineconeVectorStore from datetime import datetime import os from langchain.document_loaders import PyPDFLoader from langchain.text_splitter import RecursiveCharacterTextSplitter from langchain.embeddings.openai import OpenAIEmbeddings from langchain.vectorstores import Pinecone from typing import TypedDict,List from langgraph.graph import StateGraph from langgraph.prebuilt import ToolNode from langchain.schema import Document from langchain.prompts import PromptTemplate from langchain.tools import Tool from langchain.llms import OpenAI # OpenAI API key openai_api_key = os.getenv("OPENAI_API_KEY") # Embedding using OpenAI embeddings = OpenAIEmbeddings(api_key=openai_api_key) # Initialize Pinecone with PineconeGRPC from pinecone import Pinecone pc = Pinecone(api_key=os.environ['PINECONE_API_KEY']) # Define index name and parameters index_name = "italy-kg" vectorstore = PineconeVectorStore(index_name=index_name, embedding=embeddings) # Tool functions def search_vector_db(query: str, k: int = 3) -> List[Document]: docs = vectorstore.similarity_search(query, k=k) return docs def expand_query(query: str) -> str: return query def summarize_context(context: str) -> str: prompt = PromptTemplate(template="""Summarize the following Context to provide a concise overview: {context}""") summary = llm(prompt.format(context=context)) return summary.strip() def generate_response(context: str, query: str) -> str: prompt = PromptTemplate(template="""Question: {question}\nContext: {context}\nAnswer:""") formatted_prompt = prompt.format(context=context, question=query) response = llm(formatted_prompt) return response.strip() # Tool objects expand_tool = Tool( name="Expand Query", func=expand_query, description="Enhance the query with additional terms or context" ) summarize_tool = Tool( name="Summarize Context", func=summarize_context, description="Summarize the context to provide a concise overview" ) search_tool = Tool( name="Search Vector Database", func=search_vector_db, description="Search the vector database for relevant information" ) generate_tool = Tool( name="Generate Response", func=generate_response, description="Generate a response based on the context and query" ) # State for the graph class State(TypedDict): question: str context: List[Document] response: str expanded_query: str summarized_context: str # Workflow node definitions def expand(state: State) -> State: state["expanded_query"] = expand_tool.func(state["question"]) # Expand the query return state def search(state: State) -> State: results = search_tool.func(state["expanded_query"]) # Search using the expanded query state["context"] = results print(f"Retrieved Documents: {[doc.page_content[:100] for doc in results]}") return state def summarize(state: State) -> State: context = " ".join(doc.page_content for doc in state["context"]) if state["context"] else "" state["summarized_context"] = summarize_tool.func(context) print(f"Summarized Context: {state['summarized_context']}") return state def generate(state: State) -> State: response = generate_tool.func(state["summarized_context"], state["question"]) state["response"] = response print(f"Generated Response: {state['response']}") return state # Workflow graph workflow = StateGraph(State) workflow.add_node("expand", expand) workflow.add_node("search", search) workflow.add_node("summarize", summarize) workflow.add_node("generate", generate) workflow.set_entry_point("expand") workflow.add_edge("expand", "search") workflow.add_edge("search", "summarize") workflow.add_edge("summarize", "generate") workflow.set_finish_point("generate") graph = workflow.compile() # Function to run the graph def run_graph(question: str): result = graph.invoke({"question": question}) return result["response"] # Create a global list to store uploaded document records uploaded_documents = [] # Function to process PDF, extract text, split it into chunks, and upload to the vector DB def process_pdf(pdf_file, uploaded_documents): if pdf_file is None: return uploaded_documents, "No PDF file uploaded." with pdfplumber.open(pdf_file.name) as pdf: all_text = "" for page in pdf.pages: all_text += page.extract_text() # Split the text into chunks text_splitter = RecursiveCharacterTextSplitter(chunk_size=300, chunk_overlap=50) chunks = text_splitter.split_text(all_text) # Embed and upload the chunks into the vector database chunk_ids = [] for chunk in chunks: document = Document(page_content=chunk) chunk_id = vectorstore.add_documents([document]) chunk_ids.append(chunk_id) # Update the upload history document_record = { "Document Name": pdf_file.name, "Upload Time": datetime.now().strftime("%Y-%m-%d %H:%M:%S"), "Chunks": len(chunks), "Pinecone Index": index_name } # Add the record to the global list uploaded_documents.append(document_record) # Convert the list of dictionaries into a list of lists for the dataframe table_data = [[doc["Document Name"], doc["Upload Time"], doc["Chunks"], doc["Pinecone Index"]] for doc in uploaded_documents] return table_data, f"Uploaded {len(chunks)} chunks to the vector database." # Gradio Interface with gr.Blocks() as demo: gr.Markdown("# PDF Uploader to Pinecone with Query Response") with gr.Column(): with gr.Row(): response_output = gr.Textbox(label="Response:", lines=10, max_lines=10) query_input = gr.Textbox(label="Enter your query:") query_button = gr.Button("Get Response") query_button.click(fn=run_graph, inputs=query_input, outputs=response_output) with gr.Row(): file_input = gr.File(label="Upload PDF", file_types=[".pdf"]) document_table = gr.Dataframe(headers=["Document Name", "Upload Time", "Chunks", "Pinecone Index"], interactive=False) output_textbox = gr.Textbox(label="Result") process_button = gr.Button("Process PDF and Upload") process_button.click(fn=process_pdf, inputs=[file_input, gr.State([])], outputs=[document_table, output_textbox]) demo.launch(show_error=True)