Minor fixes.

app.py

@@ -98,7 +98,10 @@ default_model = smolagents.HfApiModel()
 mistral_api_key = os.environ["MISTRAL_API_KEY"]
 mistral_model_id = "mistral/codestral-latest"
 mistral_model = smolagents.LiteLLMModel(
-    model_id=mistral_model_id, api_key=mistral_api_key)
+    model_id=mistral_model_id,
+    api_key=mistral_api_key,
+    temperature=0.0
+)
 
 #
 # Define the agent
@@ -108,17 +111,73 @@ agent = smolagents.CodeAgent(
     model=mistral_model
 )
 
+#
+# Handler to extract the response's content
+#
+from typing import Union, Any
+from dataclasses import is_dataclass
+import json
+
+class ResponseHandler:
+    @staticmethod
+    def extract_content(response: Any) -> str:
+        """
+        Extract content from various types of agent responses.
+        
+        Args:
+            response: The response from the agent, could be string, Message object, or dict
+            
+        Returns:
+            str: The extracted content
+        """
+        # If it's already a string, return it
+        if isinstance(response, str):
+            return response
+            
+        # If it's a Message object
+        if hasattr(response, 'content') and isinstance(response.content, str):
+            return response.content
+            
+        # If it's a dictionary (e.g., from json.loads())
+        if isinstance(response, dict) and 'content' in response:
+            return response['content']
+            
+        # If it's a dataclass
+        if is_dataclass(response):
+            if hasattr(response, 'content'):
+                return response.content
+                
+        # If it's JSON string
+        if isinstance(response, str):
+            try:
+                parsed = json.loads(response)
+                if isinstance(parsed, dict) and 'content' in parsed:
+                    return parsed['content']
+            except json.JSONDecodeError:
+                pass
+                
+        # If we can't determine the type, return the string representation
+        return str(response)
+
+handler = ResponseHandler()
+
+
 def generate_response(query: str) -> str:
     """Generate a response given query.
 
     Args:
+        - query: the question from the user
 
     Returns:
         - the response from the agent having access to a database over the ICIJ
           data and a large language model.
     """
     agent_output = agent.run(query)
-    return agent_output
+
+    # At times, the response appears to be a class instance with a 'content'
+    # part. Hence, we will pass the agent's response to some handler that will
+    # extract the response's content.
+    return handler.extract_content(agent_output)
 
 
 #
@@ -127,7 +186,7 @@ def generate_response(query: str) -> str:
 with gr.Blocks() as demo:
     gr.Markdown("""
         # SQL agent
-        Database: ICIJ data on offshore financial data leaks.
+        Database: ICIJ data on offshore financial data leaks. Very early "fast" prorotyping.
     """)
 
     # Inputs: question
@@ -157,6 +216,9 @@ with gr.Blocks() as demo:
                         "Please give the name of the entity an its  address."
                     ),
                 ],
+                [
+                    "Are there any entities located on Montreal, Canada?",
+                ]
             ],
             inputs=[question,],
             outputs=[response,],