Update app.py

app.py

@@ -4,9 +4,110 @@ import requests
 import pytz
 import yaml
 from tools.final_answer import FinalAnswerTool
+import os
 
 from Gradio_UI import GradioUI
 
+from stocksymbol import StockSymbol
+from polygon import RESTClient
+import time
+import difflib
+from datetime import datetime
+from dateutil.relativedelta import relativedelta
+
+# API
+symbol_api_key = os.getenv("STOCK_SYMBOL")
+polygon_api = os.getenv('POLYGON_API')
+
+# Get the us market stock name
+ss = StockSymbol(symbol_api_key)
+symbol_list_us = ss.get_symbol_list(market="US")
+
+# Create stock price retriever
+client = RESTClient(api_key=polygon_api)
+
+# utility function
+def convert_timestamp_to_date(timestamp):
+    # Convert milliseconds to seconds
+    timestamp_seconds = timestamp / 1000.0
+    # Create a datetime object
+    date_time = datetime.fromtimestamp(timestamp_seconds)
+    # Format the date to a readable string
+    return date_time.strftime('%B %d, %Y')
+    
+def find_stock_symbol(company_name, stock_list):
+    """
+    Find the stock symbol based on the closest matching company name.
+
+    Args:
+        company_name (str): The name of the company to search for.
+        stock_list (list): A list of dictionaries containing stock information.
+
+    Returns:
+        str: The stock symbol of the closest matching company name or None if no match is found.
+    """
+    # Extract long names from the stock list
+    long_names = [stock['longName'] for stock in stock_list]
+    
+    # Find the closest match
+    closest_match = difflib.get_close_matches(company_name, long_names, n=1)
+
+    if closest_match:
+        # Get the index of the closest match
+        match_index = long_names.index(closest_match[0])
+        # Return the corresponding symbol
+        return stock_list[match_index]['symbol']
+    else:
+        return None
+
+def get_dates():
+    # Get current date
+    current_date = datetime.now()
+    # Get the date 5 months before
+    five_months_ago = current_date - relativedelta(months=6)
+    one_months_ago = current_date - relativedelta(months=1)
+    # Format dates as YYYY-MM-DD
+    one_months_ago_formatted = one_months_ago.strftime('%Y-%m-%d')
+    five_months_ago_formatted = five_months_ago.strftime('%Y-%m-%d')
+    
+    return one_months_ago_formatted, five_months_ago_formatted
+### Finally my tool
+@tool
+def get_stock_price(name_of_company: str):
+    """
+    Retrive the stock price of the company for the last 5 months.
+
+    Args:
+        name_of_company: the name of the company
+
+    Returns:
+        List: a list of aggregate stock price data. Each agg contains the following
+        Each Agg entry contains the following fields:
+            open: The price at which the stock opened at the beginning of the trading period.
+            high: The highest price reached by the stock during the trading period.
+            low: The lowest price reached by the stock during the trading period.
+            close: The price at which the stock closed at the end of the trading period.
+            volume: The total number of shares traded during the period.
+            vwap: The Volume Weighted Average Price, which gives an average price of the stock weighted by volume traded.
+            timestamp: The time at which the data point was recorded, represented as a Unix timestamp (milliseconds since January 1, 1970).
+            transactions: The total number of transactions that occurred during the trading period.
+            otc: Indicates if the trade was over-the-counter (OTC), which is typically not applicable for major stocks and is None here.
+    """
+    symbol = find_stock_symbol(name_of_company, symbol_list_us)
+    # The dates
+    current_date, five_months_ago = get_dates()
+    ticker = "AAPL"
+
+    # List Aggregates (Bars)
+    aggs = []
+    for a in client.list_aggs(ticker=ticker, multiplier=1, timespan="month", from_=five_months_ago, to=current_date, limit=5):
+        aggs.append(a)
+
+    for agg in aggs:
+        agg.timestamp = convert_timestamp_to_date(agg.timestamp)
+
+    return aggs
+
 # Below is an example of a tool that does nothing. Amaze us with your creativity !
 image_generation_tool = load_tool("agents-course/text-to-image", trust_remote_code=True)
 
@@ -73,7 +174,7 @@ with open("prompts.yaml", 'r') as stream:
     
 agent = CodeAgent(
     model=model,
-    tools=[generate_pictures_with_watermelon_in_it], ## add your tools here (don't remove final answer)
+    tools=[final_answer, get_stock_price], ## add your tools here (don't remove final answer)
     max_steps=6,
     verbosity_level=1,
     grammar=None,