fixed dict issue

app.py

@@ -10,13 +10,14 @@ from langchain_community.vectorstores import FAISS
 from langchain_huggingface import HuggingFaceEmbeddings
 from threading import Thread
 from dotenv import load_dotenv
+import json
 
 load_dotenv()
 
 # === CONFIG ===
 STORAGE_DIR = "storage"
-CLEANUP_INTERVAL = 600  # 10 min
-SESSION_TTL = 1000      # 30 min
+CLEANUP_INTERVAL = 600
+SESSION_TTL = 1000
 OPENROUTER_API_KEY = os.getenv("OPENROUTER_API_KEY")
 OPENROUTER_MODEL = "z-ai/glm-4.5-air:free"
 
@@ -39,106 +40,111 @@ Thread(target=cleanup_old_sessions, daemon=True).start()
 def process_pdf(pdf_file):
     if pdf_file is None:
         return "No file uploaded.", "", []
+
     session_id = str(uuid.uuid4())
     reader = PdfReader(pdf_file.name)
 
-    # Extract text
-    text = "".join([page.extract_text() for page in reader.pages if page.extract_text()])
+    metadata = reader.metadata or {}
+    num_pages = len(reader.pages)
+    title = metadata.get("/Title", "Unknown Title")
+    author = metadata.get("/Author", "Unknown Author")
 
-    # Metadata
-    page_count = len(reader.pages)
-    first_page_text = reader.pages[0].extract_text() if page_count > 0 else ""
-    guessed_title = first_page_text.split("\n")[0] if first_page_text else "Unknown Title"
+    text = "".join([page.extract_text() for page in reader.pages if page.extract_text()])
 
-    # Split text
     splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
     chunks = splitter.split_text(text)
 
-    # Embeddings + FAISS
     embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
     session_path = os.path.join(STORAGE_DIR, session_id)
     os.makedirs(session_path, exist_ok=True)
+
     db = FAISS.from_texts(chunks, embeddings)
     db.save_local(session_path)
 
-    # Save metadata
-    metadata_path = os.path.join(session_path, "metadata.txt")
-    with open(metadata_path, "w", encoding="utf-8") as f:
-        f.write(f"title={guessed_title}\n")
-        f.write(f"pages={page_count}\n")
+    meta_data = {
+        "title": title,
+        "author": author,
+        "pages": num_pages,
+        "session_id": session_id,
+        "created_at": time.ctime()
+    }
+    with open(os.path.join(session_path, "metadata.json"), "w") as f:
+        json.dump(meta_data, f)
 
-    chat_history = [("System", f"Paper uploaded. Title: {guessed_title}, Pages: {page_count}. You can now ask questions.")]
+    chat_history = [
+        {"role": "system", "content": f"📄 Paper uploaded.\nTitle: {title}\nAuthor: {author}\nPages: {num_pages}"}
+    ]
     return f"Paper uploaded successfully. Session ID: {session_id}", session_id, chat_history
 
 # === QUERY FUNCTION ===
 def query_paper(session_id, user_message, chat_history):
     if not session_id or not os.path.exists(os.path.join(STORAGE_DIR, session_id)):
         chat_history = chat_history or []
-        chat_history.append(("System", "Session expired or not found. Upload the paper again."))
+        chat_history.append({"role": "system", "content": "Session expired or not found. Upload the paper again."})
         return chat_history, ""
 
     if not user_message.strip():
         return chat_history, ""
 
     session_path = os.path.join(STORAGE_DIR, session_id)
-
-    # Load metadata
-    metadata_path = os.path.join(session_path, "metadata.txt")
-    metadata = {}
-    if os.path.exists(metadata_path):
-        with open(metadata_path, "r", encoding="utf-8") as f:
-            for line in f:
-                k, v = line.strip().split("=", 1)
-                metadata[k] = v
-
-    # Load retriever
     embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
     db = FAISS.load_local(session_path, embeddings, allow_dangerous_deserialization=True)
     retriever = db.as_retriever(search_kwargs={"k": 3})
 
-    # Retrieve context
-    docs = retriever.invoke(user_message)
-    context = "\n\n".join([d.page_content for d in docs])
-
-    # Prompt
-    prompt = f"""
-You are an AI assistant that explains research papers in clear, structured, simple terms. 
-You can use BOTH metadata and the paper content.
-
-Metadata:
-- Title: {metadata.get('title','Unknown')}
-- Pages: {metadata.get('pages','Unknown')}
-
-Paper content (retrieved chunks):
+    metadata_path = os.path.join(session_path, "metadata.json")
+    if os.path.exists(metadata_path):
+        with open(metadata_path, "r") as f:
+            metadata = json.load(f)
+    else:
+        metadata = {"title": "Unknown", "author": "Unknown", "pages": "Unknown"}
+
+    lower_q = user_message.lower()
+    if "title" in lower_q or "name of this paper" in lower_q:
+        answer = f"The title of this paper is: **{metadata['title']}**."
+    elif "author" in lower_q or "who wrote" in lower_q:
+        answer = f"The author of this paper is: **{metadata['author']}**."
+    elif "pages" in lower_q or "how many pages" in lower_q:
+        answer = f"This paper has **{metadata['pages']} pages**."
+    else:
+        docs = retriever.invoke(user_message)
+        context = "\n\n".join([d.page_content for d in docs])
+
+        prompt = f"""
+You are an AI research assistant. Use the paper content and metadata to answer clearly.
+
+Paper Metadata:
+- Title: {metadata['title']}
+- Author: {metadata['author']}
+- Pages: {metadata['pages']}
+
+Context from paper:
 {context}
 
-User Question: {user_message}
-Answer in plain English with clarity.
+Question: {user_message}
+Answer:
 """
-
-    headers = {"Authorization": f"Bearer {OPENROUTER_API_KEY}", "Content-Type": "application/json"}
-    payload = {
-        "model": OPENROUTER_MODEL,
-        "messages": [
-            {"role": "system", "content": "You are a helpful research paper explainer. Use metadata if the user asks about title, authors, or page count. Otherwise, use the retrieved context."},
-            {"role": "user", "content": prompt}
-        ]
-    }
-
-    try:
-        response = requests.post("https://openrouter.ai/api/v1/chat/completions",
-                                 headers=headers, json=payload)
-
-        if response.status_code == 200:
-            answer = response.json()["choices"][0]["message"]["content"].strip()
-        else:
-            answer = f"Error: {response.status_code} - {response.text}"
-    except Exception as e:
-        answer = f"Error: {str(e)}"
+        headers = {"Authorization": f"Bearer {OPENROUTER_API_KEY}", "Content-Type": "application/json"}
+        payload = {
+            "model": OPENROUTER_MODEL,
+            "messages": [
+                {"role": "system", "content": "You are a helpful research explainer. Always use metadata if available."},
+                {"role": "user", "content": prompt}
+            ]
+        }
+
+        try:
+            response = requests.post("https://openrouter.ai/api/v1/chat/completions",
+                                     headers=headers, json=payload)
+            if response.status_code == 200:
+                answer = response.json()["choices"][0]["message"]["content"].strip()
+            else:
+                answer = f"Error: {response.status_code} - {response.text}"
+        except Exception as e:
+            answer = f"Error: {str(e)}"
 
     chat_history = chat_history or []
-    chat_history.append((user_message, answer))
-    
+    chat_history.append({"role": "user", "content": user_message})
+    chat_history.append({"role": "assistant", "content": answer})
     return chat_history, ""
 
 # === GRADIO UI ===
@@ -149,7 +155,7 @@ with gr.Blocks() as demo:
         pdf_input = gr.File(label="Upload Research Paper (PDF)", file_types=[".pdf"])
         session_box = gr.Textbox(label="Session ID", interactive=False)
 
-    chatbot = gr.Chatbot(label="Chat about your paper", height=400, type="messages")
+    chatbot = gr.Chatbot(label="Chat about your paper", type="messages", height=400)
     user_message = gr.Textbox(label="Ask a question", placeholder="What is this paper about?")
 
     with gr.Row():
@@ -157,11 +163,9 @@ with gr.Blocks() as demo:
         ask_btn = gr.Button("Send Question")
         clear_btn = gr.Button("Clear Chat")
 
-    # Store chat history and session
     state_chat = gr.State([])
     state_session = gr.State("")
 
-    # Upload
     def handle_upload(pdf_file):
         status, session_id, chat_history = process_pdf(pdf_file)
         return status, session_id, chat_history
@@ -172,7 +176,6 @@ with gr.Blocks() as demo:
         outputs=[session_box, state_session, state_chat]
     )
 
-    # Ask
     def handle_question(session_id, message, chat_history):
         updated_chat, _ = query_paper(session_id, message, chat_history)
         return updated_chat, ""
@@ -197,7 +200,6 @@ with gr.Blocks() as demo:
         outputs=[state_chat]
     )
 
-    # Clear
     def clear_chat():
         return [], []
 
@@ -212,4 +214,4 @@ with gr.Blocks() as demo:
         outputs=[chatbot]
     )
 
-demo.launch(debug=True)
+demo.launch(debug=True, server_name="0.0.0.0", server_port=7860)