tools.py

from pydantic import BaseModel
from datetime import datetime, timedelta
import yfinance as yf
from langchain.prompts import MessagesPlaceholder, ChatPromptTemplate, HumanMessagePromptTemplate, PromptTemplate, AIMessagePromptTemplate
from pydantic import BaseModel, Field
from langchain.tools import BaseTool
from typing import Optional, Type
from typing import List
from functools import lru_cache


@lru_cache()
def get_stock_price(symbol):
    ticker = yf.Ticker(symbol)
    todays_data = ticker.history(period='1d')
    price = round(todays_data['Close'][0], 2)
    currency = ticker.info['currency']
    return price, currency

@lru_cache()
def get_stock_data_yahoo(ticker):
    stock = yf.Ticker(ticker)
    data = stock.history(period="1y")
    return data

@lru_cache()
def get_company_profile_yahoo(ticker):
    stock = yf.Ticker(ticker)
    info = stock.info
    profile = {
        "name": info.get("shortName"),
        "sector": info.get("sector"),
        "industry": info.get("industry"),
        "marketCap": info.get("marketCap"),
        "website": info.get("website"),
        "description": info.get("longBusinessSummary"),
    }
    return profile

@lru_cache()
def get_company_news_yahoo(ticker):
    stock = yf.Ticker(ticker)
    news = stock.news
    return news

@lru_cache()
def get_price_change_percent(symbol, days_ago):
    ticker = yf.Ticker(symbol)
    end_date = datetime.now()
    start_date = end_date - timedelta(days=days_ago)

    # Convert dates to string format that yfinance can accept
    start_date = start_date.strftime('%Y-%m-%d')
    end_date = end_date.strftime('%Y-%m-%d')

    historical_data = ticker.history(start=start_date, end=end_date)

    old_price = historical_data['Close'].iloc[0]
    new_price = historical_data['Close'].iloc[-1]

    percent_change = ((new_price - old_price) / old_price) * 100
    return round(percent_change, 2)

@lru_cache()
def calculate_performance(symbol, days_ago):
    ticker = yf.Ticker(symbol)
    end_date = datetime.now()
    start_date = end_date - timedelta(days=days_ago)
    start_date = start_date.strftime('%Y-%m-%d')
    end_date = end_date.strftime('%Y-%m-%d')
    historical_data = ticker.history(start=start_date, end=end_date)
    old_price = historical_data['Close'].iloc[0]
    new_price = historical_data['Close'].iloc[-1]
    percent_change = ((new_price - old_price) / old_price) * 100
    return round(percent_change, 2)

@lru_cache()
def get_best_performing(stocks, days_ago):
    best_stock = None
    best_performance = None
    for stock in stocks:
        try:
            performance = calculate_performance(stock, days_ago)
            if best_performance is None or performance > best_performance:
                best_stock = stock
                best_performance = performance
        except Exception as e:
            print(f"Could not calculate performance for {stock}: {e}")
    return best_stock, best_performance

class StockPriceCheckInput(BaseModel):
    """Input for Stock price check."""

    stockticker: str = Field(..., description="Ticker symbol for stock or index")

class StockPriceTool(BaseTool):
    name = "get_stock_ticker_price"
    description = "Useful for when you need to find out the price of the stock today. You should input the stock ticker used on the yfinance API"

    def _run(self, stockticker: str):
        # print("i'm running")
        price_response, currency = get_stock_price(stockticker)

        return f"{currency} {price_response}"

    def _arun(self, stockticker: str):
        raise NotImplementedError("This tool does not support async")

    args_schema: Optional[Type[BaseModel]] = StockPriceCheckInput

class PrevYearStockTool(BaseTool):
    name = "get_past_year_stock_data"
    description = "Useful for when you need to find out the past 1 year performance of a stock. You should input the stock ticker used on the yfinance API"

    def _run(self, stockticker: str):
        price_response = get_stock_data_yahoo(stockticker)
        return price_response

    def _arun(self, stockticker: str):
        raise NotImplementedError("This tool does not support async")

    args_schema: Optional[Type[BaseModel]] = StockPriceCheckInput
    
class StockNewsTool(BaseTool):
    name = "get_news_about_stock"
    description = "Useful for when you need recent news related to a stock. You should input the stock ticker used on the yfinance API"

    def _run(self, stockticker: str):
        price_response = get_company_news_yahoo(stockticker)
        return price_response

    def _arun(self, stockticker: str):
        raise NotImplementedError("This tool does not support async")

    args_schema: Optional[Type[BaseModel]] = StockPriceCheckInput

class StockProfileTool(BaseTool):
    name = "get_profile_of_stock"
    description = "Useful for when you need details or profile of a stock. You should input the stock ticker used on the yfinance API"

    def _run(self, stockticker: str):
        price_response = get_company_profile_yahoo(stockticker)
        return price_response

    def _arun(self, stockticker: str):
        raise NotImplementedError("This tool does not support async")

    args_schema: Optional[Type[BaseModel]] = StockPriceCheckInput

class StockChangePercentageCheckInput(BaseModel):
    """Input for Stock ticker check. for percentage check"""

    stockticker: str = Field(..., description="Ticker symbol for stock or index")
    days_ago: int = Field(..., description="Int number of days to look back")

class StockPercentageChangeTool(BaseTool):
    name = "get_price_change_percent"
    description = "Useful for when you need to find out the performance or percentage change in a stock's value. You should input the stock ticker used on the yfinance API and also input the number of days to check the change over"

    def _run(self, stockticker: str, days_ago: int):
        price_change_response = get_price_change_percent(stockticker, days_ago)

        return price_change_response

    def _arun(self, stockticker: str, days_ago: int):
        raise NotImplementedError("This tool does not support async")

    args_schema: Optional[Type[BaseModel]] = StockChangePercentageCheckInput

class StockBestPerformingInput(BaseModel):
    """Input for Stock ticker check. for percentage check"""

    stocktickers: List[str] = Field(..., description="Ticker symbols for stocks or indices")
    days_ago: int = Field(..., description="Int number of days to look back")

class StockGetBestPerformingTool(BaseTool):
    name = "get_best_performing"
    description = "Useful for when you need to the performance of multiple stocks over a period. You should input a list of stock tickers used on the yfinance API and also input the number of days to check the change over"

    def _run(self, stocktickers: List[str], days_ago: int):
        price_change_response = get_best_performing(stocktickers, days_ago)

        return price_change_response

    def _arun(self, stockticker: List[str], days_ago: int):
        raise NotImplementedError("This tool does not support async")

    args_schema: Optional[Type[BaseModel]] = StockBestPerformingInput