Update app.py

app.py

@@ -17,9 +17,18 @@ class AIRegulationModel:
             model_id=model_id
         )
 
-    def get_regulation_info(self, query):
-        response = self.model.generate(prompt=query)  # Fix here!
-        return response if isinstance(response, str) else response.get("generated_text", "No response generated.")
+@tool
+def ai_regulation_query_tool(query: str) -> str:
+    """Queries an AI model for information on AI regulation.
+    
+    Args:
+        query (str): The input question related to AI regulations.
+
+    Returns:
+        str: A relevant response from the AI regulation model.
+    """
+    return ai_regulation_tool.get_regulation_info(query)
+
 
 # Initialize the AI Regulation Tool
 ai_regulation_tool = AIRegulationModel()