First_agent_template / tools /extract_market_data.py
egmaminta2's picture
Update tools/extract_market_data.py
2a12c09 verified
from smolagents.tools import Tool
# Creating Tool class
# Here we are creating a market data tool using `yfinance`
class ExtractMarketDataTool(Tool):
name = "extract_market_data"
description = ("Extracts market data (via `yfinance` or Yahoo! Finance) and "
"returns the information about a stock such as the stock's info, "
"address, city-state-zip-country, phone, website, industry, sector, "
"calendar data such as dividend rate, highs, lows, average, revenue, etc. "
"Use this when someone asks information about a market stock. "
"Ensure that a digestible or summarized version is the output (let the LLM summarize).")
inputs = {'stock': {'type': 'string', 'description': 'The market/stock code that the User wants to inspect.'}}
output_type = "string"
def forward(self, stock: str) -> str:
try:
import yfinance as yf
except ImportError as e:
raise ImportError(
"You must install packages `yfinance` to run this tool: for instance run `pip install yfinance`"
) from e
try:
data = yf.Ticker(f"{stock}")
info_data = data.info
# INFO
address = info_data['address1']
city_state_zip_country = f"{info_data['city']}, {info_data['state']}, {info_data['zip']}, {info_data['country']}"
phone = f"{info_data['phone']}"
website = f"{info_data['website']}"
industry = f"{info_data['industry']}"
sector = f"{info_data['sector']}"
calendar_data = data.calendar
# CALENDAR
dividend_date = calendar_data['Dividend Date']
ex_dividend_date = calendar_data['Ex-Dividend Date']
earnings_date = calendar_data['Earnings Date']
earnings_high = calendar_data['Earnings High']
earnings_low = calendar_data['Earnings Low']
earnings_avg = calendar_data['Earnings Average']
rev_high = calendar_data['Revenue High']
rev_low = calendar_data['Revenue Low']
rev_avg = calendar_data['Revenue Average']
output = (f"Here's information about {stock}:\n"
f"address: {address}\n"
f"city, state, zip, country: {city_state_zip_country}\n"
f"phone: {phone}\n"
f"website: {website}\n"
f"industry: {industry}\n"
f"sector: {sector}\n"
f"dividend date: {dividend_date}\n"
f"ex-dividend date: {ex_dividend_date}\n"
f"earnings date: {earnings_date}\n"
f"earnings high: {earnings_high}\n"
f"earnings low: {earnings_low}\n"
f"earnings average: {earnings_avg}\n"
f"revenue high: {rev_high}\n"
f"revenue low: {rev_low}\n"
f"revenue average: {rev_avg}\n")
return output
except Exception as e:
return f"An unexpected error occurred: {str(e)}"
def __init__(self, *args, **kwargs) -> None:
self.is_initialized = False