Add stock analysis tools to app.py: implement sentiment analysis, stock price prediction, and stock comparison features; update requirements.txt to include statsmodels

app.py

@@ -11,6 +11,7 @@ from langchain_huggingface import HuggingFacePipeline
 from langchain import LLMChain
 from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 from langchain.memory import ConversationBufferWindowMemory
+from statsmodels.tsa.arima.model import ARIMA
 import torch
 import re
 from typing import List, Union
@@ -64,6 +65,37 @@ def calculate_moving_average(ticker, window=5):
     hist[f"{window}-day MA"] = hist["Close"].rolling(window=window).mean()
     return hist[["Close", f"{window}-day MA"]].tail(5)
 
+def analyze_sentiment(news_articles):
+    sentiment_pipeline = pipeline("sentiment-analysis")
+    results = [{"title": article["title"],
+                "sentiment": sentiment_pipeline(article["description"] or article["title"])[0]}
+               for article in news_articles]
+    return results
+
+def predict_stock_price(ticker, days=5):
+    stock = yf.Ticker(ticker)
+    hist = stock.history(period="6mo")
+    if hist.empty:
+        return {"error": f"No data found for ticker {ticker}"}
+
+    
+    model = ARIMA(hist["Close"], order=(5, 1, 0))
+    model_fit = model.fit()
+    forecast = model_fit.forecast(steps=days)
+    return forecast.tolist()
+
+def compare_stocks(ticker1, ticker2):
+    data1 = fetch_stock_data(ticker1)
+    data2 = fetch_stock_data(ticker2)
+    if "error" in data1 or "error" in data2:
+        return {"error": "Could not fetch stock data for comparison."}
+    comparison = {
+        ticker1: {"recent_close": data1["Close"][-1]},
+        ticker2: {"recent_close": data2["Close"][-1]},
+    }
+    return comparison
+
+
 # Define LangChain tools
 stock_data_tool = Tool(
     name="Stock Data Fetcher",
@@ -83,7 +115,32 @@ moving_average_tool = Tool(
     description="Calculate the moving average of a stock over a 5-day window."
 )
 
-tools = [stock_data_tool, stock_news_tool, moving_average_tool]
+sentiment_tool = Tool(
+    name="News Sentiment Analyzer",
+    func=lambda ticker: analyze_sentiment(fetch_stock_news(ticker, NEWSAPI_KEY)),
+    description="Analyze the sentiment of recent news articles about a stock ticker."
+)
+
+stock_prediction_tool = Tool(
+    name="Stock Price Predictor",
+    func=predict_stock_price,
+    description="Predict future stock prices for a given ticker based on historical data."
+)
+
+stock_comparator_tool = Tool(
+    name="Stock Comparator",
+    func=lambda tickers: compare_stocks(*tickers.split(',')),
+    description="Compare the recent performance of two stocks given their tickers, e.g., 'AAPL,MSFT'."
+)
+
+tools = [
+    stock_data_tool, 
+    stock_news_tool, 
+    moving_average_tool, 
+    sentiment_tool, 
+    stock_prediction_tool, 
+    stock_comparator_tool
+]
 
 # Set up a prompt template with history
 template_with_history = """You are SearchGPT, a professional search engine who provides informative answers to users. Answer the following questions as best you can. You have access to the following tools:

requirements.txt

@@ -4,4 +4,5 @@ requests
 pandas
 langchain
 langchain_huggingface
-python-dotenv
+python-dotenv
+statsmodels