nvidia_bot.py

import os
import time
from datetime import datetime

import schedule
import telebot
from dotenv import load_dotenv
from langchain_chroma import Chroma
from langchain_community.tools import DuckDuckGoSearchRun
from langchain_core.messages import HumanMessage, SystemMessage
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_openai import ChatOpenAI

from news_price_correlation import build_correlation_analysis
from nvidia_agent_core import forecast_markdown, predict_nvidia_stock_payload, route_query


load_dotenv()

BOT_TOKEN = os.getenv("TELEGRAM_BOT_TOKEN")
CHAT_ID = os.getenv("TELEGRAM_CHAT_ID")
USE_WEBHOOK = os.getenv("USE_WEBHOOK", "False").lower() == "true"
WEBHOOK_URL = os.getenv("RENDER_WEBHOOK_URL")

if not BOT_TOKEN or not CHAT_ID:
    raise ValueError("TELEGRAM_BOT_TOKEN or TELEGRAM_CHAT_ID not found in .env")

bot = telebot.TeleBot(BOT_TOKEN, parse_mode="Markdown")
llm = ChatOpenAI(model="gpt-5.4", temperature=0.3, max_tokens=1024)
ddg_search = DuckDuckGoSearchRun()

vectorstore = None
try:
    import chromadb
    from chromadb.config import Settings

    embeddings = HuggingFaceEmbeddings(model_name="all-mpnet-base-v2")
    client = chromadb.PersistentClient(path="./chroma_db_v2", settings=Settings(allow_reset=True))
    vectorstore = Chroma(
        client=client,
        collection_name="nvidia_annual_reports_2014_2025",
        embedding_function=embeddings,
    )
    print("RAG loaded successfully")
except Exception as exc:
    print(f"RAG load failed; continuing without it: {exc}")


def truncate_telegram(text: str, limit: int = 3800) -> str:
    return text[:limit] + "\n\n... (truncated)" if len(text) > limit else text


def trader_forecast() -> str:
    try:
        payload = predict_nvidia_stock_payload(periods=7)
        return truncate_telegram(forecast_markdown(payload))
    except Exception as exc:
        return f"Forecast error: {exc}"


def correlation_answer(query: str) -> str:
    try:
        answer = build_correlation_analysis(
            query=query,
            search=ddg_search.run,
            llm=llm,
            csv_path="nvda_2014_to_2026.csv",
        )
        return truncate_telegram(answer)
    except Exception as exc:
        return f"Correlation analysis error: {exc}"


def researcher_answer(query: str) -> str:
    context = ""
    if vectorstore:
        docs = vectorstore.similarity_search(query, k=4)
        context = "\n\n".join(doc.page_content[:600] for doc in docs)

    news = ddg_search.run(f"NVIDIA {query} latest news OR Blackwell OR Huawei")[:900]
    response = llm.invoke(
        [
            SystemMessage(content="You are a senior NVIDIA strategy analyst. Give concise, evidence-aware answers."),
            HumanMessage(content=f"Question: {query}\n\nAnnual-report context:\n{context}\n\nNews:\n{news}"),
        ]
    ).content
    return truncate_telegram(response)


def answer_query(query: str) -> str:
    route = route_query(query)
    if route == "correlation_agent":
        return correlation_answer(query)
    if route == "ML_agent":
        return trader_forecast()
    return researcher_answer(query)


@bot.message_handler(commands=["start", "help"])
def send_welcome(message):
    bot.reply_to(
        message,
        "NVIDIA Bot is online.\n\n"
        "/forecast - 7-business-day Prophet v3 residual ensemble forecast\n"
        "/news - latest NVIDIA news\n"
        "/correlate <event> - news/event to price-impact analysis or geopolitical shock stress test",
    )


@bot.message_handler(commands=["forecast"])
def send_forecast(message):
    bot.reply_to(message, trader_forecast())


@bot.message_handler(commands=["news"])
def send_news(message):
    bot.reply_to(message, truncate_telegram(ddg_search.run("NVIDIA latest news")))


@bot.message_handler(commands=["correlate"])
def send_correlation(message):
    query = message.text.replace("/correlate", "", 1).strip()
    if not query:
        query = "Nvidia China Huawei export controls selloff May 2026"
    bot.reply_to(message, correlation_answer(query))


@bot.message_handler(func=lambda message: True)
def handle_message(message):
    bot.send_chat_action(message.chat.id, "typing")
    bot.reply_to(message, answer_query(message.text.strip()))


def send_daily_alert():
    forecast = trader_forecast()
    news = truncate_telegram(ddg_search.run("NVIDIA latest news OR Blackwell OR earnings"), limit=800)
    alert = f"NVIDIA DAILY ALERT {datetime.now().strftime('%Y-%m-%d %H:%M')}\n\n{forecast}\n\nNews:\n{news}"
    try:
        bot.send_message(CHAT_ID, truncate_telegram(alert))
        print(f"Daily alert sent at {datetime.now()}")
    except Exception as exc:
        print(f"Alert error: {exc}")


if __name__ == "__main__":
    print("Nvidia_bot starting...")
    send_daily_alert()
    schedule.every(60).minutes.do(send_daily_alert)

    if USE_WEBHOOK and WEBHOOK_URL:
        bot.remove_webhook()
        time.sleep(1)
        bot.set_webhook(url=WEBHOOK_URL)
        print(f"Webhook set to: {WEBHOOK_URL}")
        while True:
            schedule.run_pending()
            time.sleep(60)
    else:
        print("Using polling mode...")
        while True:
            try:
                schedule.run_pending()
                bot.polling(none_stop=True, interval=1, timeout=30)
            except Exception as exc:
                print(f"Polling error: {exc}")
                time.sleep(10)