stock_analysis.py

import time
import logging
from datetime import datetime
import os

from crewai import LLM
from crewai_tools import ScrapeWebsiteTool, SerperDevTool

from src.config import (
    GEMINI_MODEL,
    LLM_TEMPERATURE,
    EST_MINUTES_PER_REQUEST,
    DEFAULT_INITIAL_CAPITAL,
    DEFAULT_RISK_TOLERANCE,
)
from src.stocks import is_valid_stock_code
from src.tools.calculator import CalculatorTool
from src.tools.langsearch import LangsearchTool
from src.crew.runner import build_crew
from src.storage.cache import get_stock_analysis, store_stock_analysis
from src.storage.inflight import insert_stock, search_stock, remove_stock
from src.storage.queue import add_to_queue, is_in_queue, get_queue_size

logger = logging.getLogger(__name__)


def analyse_stock(stock, initial_capital=None, risk_tolerance=None):
    """
    UI-facing entry point. Checks cache, in-progress, queue — enqueues if needed.
    Returns the analysis (str) on cache hit, or a status message on queue/enqueue.
    """
    if initial_capital is None:
        initial_capital = DEFAULT_INITIAL_CAPITAL
    if risk_tolerance is None:
        risk_tolerance = DEFAULT_RISK_TOLERANCE

    if is_in_queue(stock):
        queue_size = get_queue_size()
        estimated_time = queue_size * EST_MINUTES_PER_REQUEST
        return (
            f"Your request is already in the queue. "
            f"There are {queue_size} requests ahead of you. "
            f"Estimated completion time is {estimated_time} minutes."
        )

    if search_stock(stock):
        return "The analysis for this stock is already in progress. Please check back in a few minutes."

    cached = get_stock_analysis(stock)
    if cached:
        return cached

    add_to_queue(stock)
    queue_size = get_queue_size()
    estimated_time = queue_size * EST_MINUTES_PER_REQUEST
    return (
        f"Your request has been queued. "
        f"There are {queue_size - 1} requests ahead of you. "
        f"Estimated completion time is {estimated_time} minutes."
    )


def run_analysis(stock, initial_capital=None, risk_tolerance=None):
    """
    Worker-facing. Actually runs the crew and writes the result to cache.
    Called only by the background queue worker.
    """
    if initial_capital is None:
        initial_capital = DEFAULT_INITIAL_CAPITAL
    if risk_tolerance is None:
        risk_tolerance = DEFAULT_RISK_TOLERANCE

    if search_stock(stock):
        return None

    cached = get_stock_analysis(stock)
    if cached:
        return cached

    insert_stock(stock)

    try:
        os.environ["SERPER_API_KEY"] = os.getenv("SERPER_API_KEY")
        gemini_llm = LLM(
            model=GEMINI_MODEL,
            temperature=LLM_TEMPERATURE,
            api_key=os.getenv("GEMINI_API_KEY"),
        )
        tools = [
            LangsearchTool(),
            ScrapeWebsiteTool(),
            SerperDevTool(),
            CalculatorTool(),
        ]
        crew = build_crew(gemini_llm, tools)

        financial_trading_inputs = {
            "stock_selection": is_valid_stock_code(stock),
            "initial_capital": initial_capital,
            "risk_tolerance": risk_tolerance,
            "date": datetime.today().strftime("%Y-%m-%d"),
            "news_impact_consideration": True,
        }

        start_time = time.time()
        stock_analysis = crew.kickoff(inputs=financial_trading_inputs)
        result = stock_analysis.raw

        logger.info("Stock analysis completed for %s", stock)

        store_stock_analysis(stock, result)
        end_time = time.time()
        total_time = end_time - start_time
        with open("time.txt", "a") as f:
            f.write(f"\nstock: {stock}, Time taken: {total_time}, Query time: {datetime.fromtimestamp(start_time).isoformat()}\n")

        return result

    except Exception as e:
        logger.exception("Error processing stock %s", stock)
        return f"Error encountered, Please retry: {str(e)}"
    finally:
        remove_stock(stock)