Update app.py

app.py

@@ -1,4 +1,4 @@
-from smolagents import CodeAgent,DuckDuckGoSearchTool, HfApiModel,load_tool,tool
+from smolagents import CodeAgent,DuckDuckGoSearchTool, HfApiModel, load_tool, tool
 import datetime
 import requests
 import pytz
@@ -6,8 +6,10 @@ import yaml
 from tools.final_answer import FinalAnswerTool
 import os
 from huggingface_hub import list_models
-
 from Gradio_UI import GradioUI
+import baostock as bs
+import pandas as pd
+from tools.house_drawing_tool import HouseDrawingTool
 
 # Below is an example of a tool that does nothing. Amaze us with your creativity !
 @tool
@@ -47,8 +49,6 @@ def get_current_time_in_timezone(timezone: str) -> str:
     except Exception as e:
         return f"Error fetching time for timezone '{timezone}': {str(e)}"
 
-import baostock as bs
-import pandas as pd
 @tool
 def get_stock_data_baostock(stock_code: str) -> str:
     """Fetch the latest stock data for a given stock code using Baostock.
@@ -70,40 +70,43 @@ def get_stock_data_baostock(stock_code: str) -> str:
         )
         if rs.error_code != '0':
             return f"Failed to fetch stock data: {rs.error_msg}"
-
         # Convert the data to a DataFrame
         data_list = []
         while (rs.error_code == '0') & rs.next():
             data_list.append(rs.get_row_data())
         df = pd.DataFrame(data_list, columns=rs.fields)
-
         # Log out of the Baostock system
         bs.logout()
-
         # Check if the data is empty
         if df.empty:
             return f"No data found for stock code '{stock_code}'."
-
         # Extract the latest trading day's data
         latest_data = df.iloc[-1]  # Get the last row of data (most recent trading day)
         date = latest_data['date']  # Date
         stock_name = latest_data['code']  # Stock code
         close_price = latest_data['close']  # Closing price
         volume = latest_data['volume']  # Trading volume
-
         # Construct the result string
         result = f"Stock Code: {stock_name}\nDate: {date}\nClosing Price: {close_price}\nVolume: {volume}"
         return result
-
     except Exception as e:
         return f"An error occurred while fetching stock data: {str(e)}"
-
 # Example call
 print(get_stock_data_baostock('sh.600941'))
 
+@tool
+final_answer = FinalAnswerTool()
+
+@tool
+# Import tool from Hub
+image_generation_tool = load_tool("agents-course/text-to-image", trust_remote_code=True)
+
+@tool
+house_draw = HouseDrawingTool()
 
+with open("prompts.yaml", 'r') as stream:
+    prompt_templates = yaml.safe_load(stream)
 
-final_answer = FinalAnswerTool()
 model = HfApiModel(
     max_tokens=2096,
     temperature=0.5,
@@ -112,19 +115,9 @@ model = HfApiModel(
     token=os.getenv('hf_token'),
 )
 
-
-# Import tool from Hub
-image_generation_tool = load_tool("agents-course/text-to-image", trust_remote_code=True)
-
-with open("prompts.yaml", 'r') as stream:
-    prompt_templates = yaml.safe_load(stream)
-    
 agent = CodeAgent(
     model=model,
-    tools=[
-        final_answer,
-        image_generation_tool,
-        get_stock_data_baostock], ## add your tools here (don't remove final answer)
+    tools=[final_answer, image_generation_tool, get_stock_data_baostock, house_draw], ## add your tools here (don't remove final answer)
     max_steps=6,
     verbosity_level=1,
     grammar=None,
@@ -135,5 +128,4 @@ agent = CodeAgent(
     add_base_tools=True,
 )
 
-
 GradioUI(agent).launch()