Update app.py

app.py

@@ -4,12 +4,22 @@ import requests
 # API URL
 API_URL = "https://startrz-devi.hf.space/api/v1/prediction/e54adffc-ae77-42e5-9fc0-c4584e081093"
 
-# Function to query the API
+# Function to query the API and extract toolOutput
 def query(payload):
     try:
         response = requests.post(API_URL, json=payload)
         response.raise_for_status()  # Raise an error for bad HTTP responses
-        return response.json()  # Return the JSON response
+        data = response.json()  # Parse the JSON response
+
+        # Extract toolOutput if present in the response
+        tool_outputs = []
+        if "agentReasoning" in data:
+            for reasoning in data["agentReasoning"]:
+                if "usedTools" in reasoning:
+                    tool_outputs.extend(reasoning["usedTools"])
+
+        # Return the full response and toolOutput
+        return {"raw_response": data, "tool_output": tool_outputs}
     except requests.exceptions.RequestException as e:
         return {"error": str(e)}
 
@@ -30,11 +40,19 @@ def main():
                     "question": user_input,
                 })
 
-            # Display the raw JSON response
-            st.subheader("Raw JSON Response:")
-            st.json(response)
+            # Display the extracted toolOutput on the main page
+            st.subheader("Extracted Tool Output:")
+            tool_output = response.get("tool_output", [])
+            if tool_output:
+                st.write(tool_output)
+            else:
+                st.write("No toolOutput data found in the 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.")
 
 if __name__ == "__main__":
-    main()
+    main()