feat: add thought summary in the response

app/agent.py

@@ -15,6 +15,7 @@ CONFIG = types.GenerateContentConfig(
         literature.search_medical_literature,
     ],
     system_instruction=(Path(__file__).parent / "system_instruction.txt").read_text(),
+    thinking_config=types.ThinkingConfig(include_thoughts=True),
 )
 
 SOURCE_TOOL_NAMES = {
@@ -48,6 +49,13 @@ def hydrate_sources(
 
 
 def validate_response(response: types.GenerateContentResponse) -> models.Statements:
+    thoughts = []
+
+    for part in response.candidates[0].content.parts:  # type: ignore
+        if part.thought:
+            thoughts.append(part.text)
+
+    thoughts = " ".join(thoughts)
     text = (response.text or "").strip()
 
     # Extract content inside the first markdown code block (``` or ```json)
@@ -57,8 +65,12 @@ def validate_response(response: types.GenerateContentResponse) -> models.Stateme
 
     try:
         statements = models.Statements.model_validate_json(f'{{"statements":{text}}}')
+        statements.thoughts = thoughts
     except ValidationError:
-        statements = models.Statements(statements=[models.Statement(text=text)])
+        statements = models.Statements(
+            statements=[models.Statement(text=text)],
+            thoughts=thoughts,
+        )
 
     statements = hydrate_sources(
         statements, response.automatic_function_calling_history or []

app/models.py

@@ -25,3 +25,4 @@ class Statement(BaseModel):
 
 class Statements(BaseModel):
     statements: list[Statement]
+    thoughts: str | None = None

main.py

@@ -28,7 +28,9 @@ def format_output(statements: models.Statements) -> tuple[str, str]:
 
     answer = " ".join(sentences)
     footnotes = "\n".join(f"[^{id}]: {citation}" for citation, id in citations.items())
-    return answer, footnotes
+    thought = statements.thoughts or ""
+
+    return f"{answer}\n\n{footnotes}", thought
 
 
 def main():
@@ -38,14 +40,14 @@ def main():
     with st.form("search", border=False):
         query = st.text_input("Your medical question")
         submit = st.form_submit_button("Ask")
-        response = st.empty()
 
         if submit:
             with st.spinner("Thinking...", show_time=True):
                 output = agent.respond(gemini_client, query)
-
-            answer, footnotes = format_output(output)
-            response.markdown(f"{answer}\n\n{footnotes}")
+            answer, thoughts = format_output(output)
+            with st.expander("Thinking Process"):
+                st.markdown(thoughts)
+            st.markdown(answer)
 
 
 if __name__ == "__main__":