Update app.py

app.py

@@ -1,99 +1,157 @@
 import streamlit as st
 import requests
+import json
 
 # API URL
 API_URL = "https://startrz-devi.hf.space/api/v1/prediction/e54adffc-ae77-42e5-9fc0-c4584e081093"
 
-# Function to query the API and extract toolOutput
+def parse_tool_details(tool):
+    """
+    Parse tool details with robust handling of different input types
+    
+    Args:
+        tool (dict): A single tool dictionary from the API response
+    
+    Returns:
+        dict: Parsed tool details with consistent structure
+    """
+    # Parse toolInput
+    input_value = ""
+    if isinstance(tool.get("toolInput"), dict):
+        input_value = tool["toolInput"].get("input", "")
+        # Fallback to full input dict as string if no 'input' key
+        if not input_value:
+            input_value = json.dumps(tool["toolInput"], indent=2)
+    elif isinstance(tool.get("toolInput"), str):
+        input_value = tool["toolInput"]
+    else:
+        input_value = str(tool.get("toolInput", "No input details"))
+
+    # Parse toolOutput
+    output_value = tool.get("toolOutput")
+    
+    # Flexible output handling
+    if output_value is None:
+        output_value = "No output available"
+    elif isinstance(output_value, (list, dict)):
+        # Convert to formatted JSON string for better readability
+        try:
+            output_value = json.dumps(output_value, indent=2)
+        except Exception:
+            output_value = str(output_value)
+    else:
+        # Convert to string for any other type
+        output_value = str(output_value)
+
+    return {
+        "tool": tool.get("tool", "Unknown Tool"),
+        "toolInput": input_value,
+        "toolOutput": output_value
+    }
+
 def query(payload):
+    """
+    Query the API and process the response
+    
+    Args:
+        payload (dict): Question payload to send to the API
+    
+    Returns:
+        dict: Processed response with tool details
+    """
     try:
+        # Send POST request to the API
         response = requests.post(API_URL, json=payload)
-        response.raise_for_status()  # Raise an error for bad HTTP responses
-        data = response.json()  # Parse the JSON response
-
-        # Extract toolOutput and tool details
+        response.raise_for_status()
+        
+        # Parse the JSON response
+        data = response.json()
+        
+        # Extract tool details
         tool_details = []
+        
+        # Handle different potential response structures
         if "agentReasoning" in data:
             for reasoning in data.get("agentReasoning", []):
+                # Extract used tools from each reasoning step
                 for tool in reasoning.get("usedTools", []):
-                    # Safely extract tool input
-                    tool_input = ""
-                    if isinstance(tool.get("toolInput"), dict):
-                        tool_input = tool["toolInput"].get("input", "")
-                    elif isinstance(tool.get("toolInput"), str):
-                        tool_input = tool["toolInput"]
-                    
-                    # Safely extract tool output
-                    tool_output = tool.get("toolOutput", "No output available")
-                    
-                    tool_details.append({
-                        "tool": tool.get("tool", "Unknown Tool"),
-                        "toolInput": tool_input,
-                        "toolOutput": tool_output
-                    })
-
-        # Return the full response and tool details
-        return {"raw_response": data, "tool_details": tool_details}
+                    parsed_tool = parse_tool_details(tool)
+                    tool_details.append(parsed_tool)
+        
+        return {
+            "raw_response": data,
+            "tool_details": tool_details
+        }
+    
     except requests.exceptions.RequestException as e:
-        return {"error": str(e)}
+        return {"error": f"API Request Error: {str(e)}"}
+    except json.JSONDecodeError as e:
+        return {"error": f"JSON Parsing Error: {str(e)}"}
     except Exception as e:
-        return {"error": f"Unexpected error: {str(e)}"}
+        return {"error": f"Unexpected Error: {str(e)}"}
 
-# Streamlit app
 def main():
-    st.title("DEVI RESEARCH")
-    st.write("Ask any question, and the API will provide an answer!")
+    """
+    Main Streamlit application function
+    """
+    st.set_page_config(
+        page_title="DEVI Research Assistant",
+        page_icon="🔍",
+        layout="wide"
+    )
+
+    st.title("🔬 DEVI RESEARCH ASSISTANT")
+    st.write("Explore insights by asking a research question!")
 
-    # Input box for user query
-    user_input = st.text_input("Enter your question below:")
+    # User input section
+    user_input = st.text_input(
+        "What would you like to research?", 
+        placeholder="Enter your research query here..."
+    )
 
-    # Add a button to submit the query
-    if st.button("Submit"):
+    # Submit button
+    if st.button("Explore Insights", type="primary"):
         if user_input:
-            # Send the user input to the API
-            with st.spinner("Fetching response..."):
-                response = query({
-                    "question": user_input,
-                })
+            # Progress spinner during API call
+            with st.spinner("Gathering research insights..."):
+                response = query({"question": user_input})
 
-            # Check for errors
+            # Error handling
             if "error" in response:
-                st.error(f"Error: {response['error']}")
+                st.error(response["error"])
                 return
 
-            # Display Online Resources section
-            st.subheader("Online Resources")
+            # Display Online Resources
+            st.header("🌐 Online Resources")
+            
             tool_details = response.get("tool_details", [])
             
             if tool_details:
-                # Create an expandable section for each resource
-                for idx, tool in enumerate(tool_details, 1):
-                    # Use toolInput as the resource name
-                    resource_name = tool.get("toolInput", f"Resource {idx}")
-                    
-                    # Create an expandable section
-                    with st.expander(resource_name):
-                        # Display the tool output
-                        tool_output = tool.get("toolOutput", "No output available")
-                        
-                        # Check if toolOutput is a list or dict and pretty print
-                        if isinstance(tool_output, (list, dict)):
-                            st.json(tool_output)
-                        elif isinstance(tool_output, str):
-                            st.write(tool_output)
-                        else:
-                            st.write(str(tool_output))
+                # Create tabs for each resource
+                tabs = st.tabs([
+                    f"{idx+1}. {tool['tool']}" 
+                    for idx, tool in enumerate(tool_details)
+                ])
+
+                # Populate each tab with resource details
+                for idx, (tool, tab) in enumerate(zip(tool_details, tabs)):
+                    with tab:
+                        st.subheader("Research Input")
+                        st.code(tool['toolInput'], language=None)
                         
-                        # Optional: Add tool type information
-                        st.caption(f"Source: {tool.get('tool', 'Unknown')}")
+                        st.subheader("Research Findings")
+                        # Use st.code for better formatting
+                        st.code(tool['toolOutput'], language=None)
+
             else:
-                st.write("No resources found.")
+                st.info("No resources found for this query.")
+
+            # Raw response in expander for advanced users
+            with st.expander("🔍 Advanced: Full API Response"):
+                st.json(response.get("raw_response", {}))
 
-            # Display the raw JSON response in the sidebar
-            st.sidebar.header("Raw JSON Response")
-            st.sidebar.json(response.get("raw_response", {}))
         else:
-            st.warning("Please enter a question before submitting.")
+            st.warning("Please enter a research question!")
 
 if __name__ == "__main__":
     main()