Update app.py

app.py

@@ -12,6 +12,18 @@ 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")
@@ -26,6 +38,104 @@ 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"]
+
+
+
 
 
 
@@ -72,22 +182,19 @@ def process_pdf(pdf_file, uploaded_documents):
 
 # Gradio Interface
 with gr.Blocks() as demo:
-    gr.Markdown("# PDF Uploader to Pinecone with Logs")
+    gr.Markdown("# PDF Uploader to Pinecone with Query Response")
     
-    # File upload component
     with gr.Column():
         file_input = gr.File(label="Upload PDF", file_types=[".pdf"])
-        # Button to trigger processing
         process_button = gr.Button("Process PDF and Upload")
-
-        # Dataframe to display uploaded document records
         document_table = gr.Dataframe(headers=["Document Name", "Upload Time", "Chunks", "Pinecone Index"], interactive=False)
-        
-        # Output textbox for results
         output_textbox = gr.Textbox(label="Result")
-
-        # Define button click action
+        
+        query_input = gr.Textbox(label="Enter your query:")
+        query_button = gr.Button("Get Response")
+        response_output = gr.Textbox(label="Response:")
+        
         process_button.click(fn=process_pdf, inputs=[file_input, gr.State([])], outputs=[document_table, output_textbox])
+        query_button.click(fn=run_graph, inputs=query_input, outputs=response_output)
 
-demo.queue()
 demo.launch(show_error=True)