Update app.py

app.py

@@ -1,5 +1,6 @@
 import gradio as gr
 import os
+import json
 from pageindex.core.tree_index import TreeIndex
 from llm_config import get_llm_client, get_model_name
 
@@ -12,41 +13,67 @@ from llm_config import get_llm_client, get_model_name
 # User provided specific token to use.
 REQUIRED_TOKEN = os.getenv("APP_TOKEN", "849ejdkf2Audjo2Jf3jdoirfjh")
 
-def process_docling_and_chat(markdown_text, user_query, token):
+def process_docling_and_chat(markdown_text, user_query, token, chat_history_json=None):
     if token != REQUIRED_TOKEN:
-        return "Error: Invalid Authentication Token."
+        yield "Error: Invalid Authentication Token."
+        return
 
     if not markdown_text:
-        return "Please provide document markdown text."
+        yield "Please provide document markdown text."
+        return
     if not user_query:
-        return "Please provide a query."
+        yield "Please provide a query."
+        return
 
     try:
+        # History parsing
+        chat_history = []
+        if chat_history_json:
+            try:
+                chat_history = json.loads(chat_history_json)
+            except:
+                pass
+
+        reasoning_log = ""
+        yield "<<<STATUS: Initializing PageIndex...>>>"
+
         # 1. Build the PageIndex Tree locally in the Space
+        reasoning_log += "<<<STATUS: Building Index from Markdown...>>>\n"
+        yield reasoning_log
         tree = TreeIndex()
         tree.build_from_markdown(markdown_text)
         
         # 2. Initialize the Navigator (The "Brain")
-        # Try Nvidia first, then Mistral
         try:
             client = get_llm_client(provider="nvidia")
             model = get_model_name(provider="nvidia")
-            # Test connection simply or just proceed
         except Exception as e:
             print(f"Nvidia client failed: {e}. Falling back to Mistral.")
             client = get_llm_client(provider="mistral")
             model = get_model_name(provider="mistral")
         
-        # 3. Perform Reasoning Search
-        # This uses the internal logic of the repo to navigate the tree
-        context = tree.reasoning_search(query=user_query, llm_client=client, model=model)
+        # 3. Perform Reasoning Search (Streamed)
+        context = ""
+        # Use stream method if available, else fallback
+        if hasattr(tree, 'reasoning_search_stream'):
+            for update in tree.reasoning_search_stream(user_query=user_query, llm_client=client, model=model):
+                if update.startswith("<<<STATUS:"):
+                    reasoning_log += update + "\n"
+                    yield reasoning_log
+                else:
+                    context = update # The last item is the full context
+        else:
+             reasoning_log += "<<<STATUS: Standard Reasoning Search...>>>\n"
+             yield reasoning_log
+             context = tree.reasoning_search(query=user_query, llm_client=client, model=model)
         
-        # 4. Final Answer Extraction
-        # Using the same client for consistency
-        response = client.chat.completions.create(
-            model=model, 
-            messages=[
-                {"role": "system", "content": """You are a Senior Petroleum Engineer assistant.
+        # 4. Final Answer Generation
+        reasoning_log += "<<<STATUS: Generating Final Answer...>>>\n"
+        yield reasoning_log
+
+        # Construct messages with history
+        messages = [
+            {"role": "system", "content": """You are a Senior Petroleum Engineer assistant.
 Your goal is to extract precise technical data from the provided document context.
 
 **Guidelines:**
@@ -64,14 +91,33 @@ Your goal is to extract precise technical data from the provided document contex
   "data": [{"x_label": 0, "y_label": 10}, ...]
 }
 ```
-"""},
-                {"role": "user", "content": f"Context:\n{context}\n\nQuery: {user_query}\n\nIf requesting data, provide a Markdown Table."}
-            ]
+"""}
+        ]
+        
+        # Add history
+        for msg in chat_history:
+            role = msg.get("role", "user")
+            content = msg.get("content", "")
+            messages.append({"role": role, "content": content})
+            
+        messages.append({"role": "user", "content": f"Context:\n{context}\n\nQuery: {user_query}\n\nIf requesting data, provide a Markdown Table."})
+
+        response_stream = client.chat.completions.create(
+            model=model, 
+            messages=messages,
+            stream=True,
+            max_tokens=8192
         )
-        return response.choices[0].message.content
+        
+        full_response_text = ""
+        for chunk in response_stream:
+            if chunk.choices[0].delta.content:
+                delta = chunk.choices[0].delta.content
+                full_response_text += delta
+                yield reasoning_log + "\n" + full_response_text
         
     except Exception as e:
-        return f"An error occurred: {str(e)}"
+        yield f"An error occurred: {str(e)}"
 
 # Gradio UI setup
 with gr.Blocks(title="Petromind AI - PageIndex RAG") as demo:
@@ -84,10 +130,11 @@ with gr.Blocks(title="Petromind AI - PageIndex RAG") as demo:
         with gr.Column(scale=1):
             query = gr.Textbox(label="What do you want to extract?", placeholder="e.g., What is the casing size?")
             token_input = gr.Textbox(label="API Token", placeholder="Enter access token", type="password")
+            history_json = gr.Textbox(visible=False, label="History JSON")
             btn = gr.Button("Analyze", variant="primary")
             output = gr.Textbox(label="Result", lines=10, interactive=False)
     
-    btn.click(fn=process_docling_and_chat, inputs=[input_md, query, token_input], outputs=output, api_name="process_docling_and_chat")
+    btn.click(fn=process_docling_and_chat, inputs=[input_md, query, token_input, history_json], outputs=output, api_name="process_docling_and_chat")
 
 if __name__ == "__main__":
     # Enable queue for concurrency