Update app.py

app.py

@@ -41,6 +41,19 @@ def extract_answer(text: str, original_question: str) -> str:
     # If it's a "how many" question, try to extract just the number
     if 'how many' in original_question.lower():
         # Look for numbers in the response
+
+        number_patterns = [
+            r'(?:is|are)\s*(\d+)',
+            r'^\s*(\d+)',
+            r'\D(\d+)\D*$',
+            r'(\d+)'
+        ]
+    
+        for pattern in number_patterns:
+            numbers = re.findall(pattern, cleaned)
+            if numbers:
+                return numbers[0]
+            
         numbers = re.findall(r'\d+', cleaned)
         if numbers:
             return numbers[0]  # Return the first number found
@@ -153,6 +166,24 @@ class FileAttachmentQueryTool(Tool):
         except Exception as e:
             return f"File download error: {e}"
 
+class WikipediaSearchTool(Tool):
+    name = "wikipedia_search"
+    description = "Search Wikipedia for information relevant to the user's query."
+    inputs = {"query": {"type": "string", "description": "The search query."}}
+    output_type = "string"
+
+    def forward(self, query: str) -> str:
+        try:
+            search_url = f"https://en.wikipedia.org/api/rest_v1/page/summary/{requests.utils.quote(query)}"
+            response = requests.get(search_url)
+            if response.status_code == 200:
+                data = response.json()
+                return data.get("extract", "No summary available.")
+            else:
+                return f"Wikipedia search error: {response.status_code}"
+        except Exception as e:
+            return f"Wikipedia search exception: {e}"
+            
 # --- Agent Implementation ---
 def select_model(provider="groq"):
     """Select and return a model based on the provider."""
@@ -187,6 +218,7 @@ class BasicAgent:
         tools = [
             MathSolver(),
             FileAttachmentQueryTool(),
+            WikipediaSearchTool()
         ]
         self.agent = CodeAgent(
             model=model,