Update src/streamlit_app.py

src/streamlit_app.py

@@ -8,6 +8,12 @@ from datetime import datetime
 from typing import Generator, Dict, Any
 from requests.exceptions import RequestException
 import re
+import json
+from pydantic import BaseModel
+
+class AnswerFormat(BaseModel):
+    think: str
+    answer: str
 
 def parse_and_render_streamed_response(streamed):
     """
@@ -96,15 +102,20 @@ def call_perplexity_api(user_query: str) -> str:
         "messages": [{"role": "user", "content": user_query}],
         "stream": False,
         "search_mode": search_mode,
-        "search_after_date_filter": search_after_date_filter
+        "search_after_date_filter": search_after_date_filter,
+        "response_format": {
+            "type": "json_schema",
+            "json_schema": {
+                "schema": AnswerFormat.model_json_schema()
+            }
+        }
     }
-
     retries = 3
     for attempt in range(retries):
         try:
             response = requests.post(url, headers=headers, json=payload, timeout=120)
             response.raise_for_status()
-            return response.json()["choices"][0]["message"]["content"]
+            return response
         except RequestException as err:
             if attempt < retries - 1:
                 wait_time = 2 ** attempt
@@ -126,8 +137,17 @@ if user_input:
         with st.spinner("Thinking..."):
             try:
                 full_response = call_perplexity_api(user_input)
-                streamed = st.write_stream(stream_response(full_response))
-                parse_and_render_streamed_response(streamed)
+                full_response_json = json.loads(full_response["choices"][0]["message"]["content"])
+                st.subheader("🤔 Reasoning")
+                if full_response_json.think:
+                    st.markdown(full_response_json.think)
+                else:
+                    st.markdown("... Reasoning not found ...")
+                st.subheader("💡 Answer")
+                if full_response_json.answer:
+                    st.markdown(full_response_json.answer)
+                else:
+                    st.markdown("... Answer not found ...")
                 st.session_state.messages.append({"role": "assistant", "content": full_response})
             except requests.RequestException as err:
                 st.error(f"Error: {err}")