Update app.py

app.py

@@ -14,18 +14,30 @@ def query(payload):
         # Extract toolOutput and tool details
         tool_details = []
         if "agentReasoning" in data:
-            for reasoning in data["agentReasoning"]:
-                if "usedTools" in reasoning:
-                    for tool in reasoning["usedTools"]:
-                        tool_details.append({
-                            "toolInput": tool.get("toolInput", {}).get("input", ""),
-                            "toolOutput": tool.get("toolOutput", "")
-                        })
+            for reasoning in data.get("agentReasoning", []):
+                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}
     except requests.exceptions.RequestException as e:
         return {"error": str(e)}
+    except Exception as e:
+        return {"error": f"Unexpected error: {str(e)}"}
 
 # Streamlit app
 def main():
@@ -67,9 +79,13 @@ def main():
                         # Check if toolOutput is a list or dict and pretty print
                         if isinstance(tool_output, (list, dict)):
                             st.json(tool_output)
-                        else:
+                        elif isinstance(tool_output, str):
                             st.write(tool_output)
+                        else:
+                            st.write(str(tool_output))
                         
+                        # Optional: Add tool type information
+                        st.caption(f"Source: {tool.get('tool', 'Unknown')}")
             else:
                 st.write("No resources found.")