app.py

import os, json, time, random
from collections import defaultdict
from datetime import date, datetime, timedelta
import gradio as gr

import pandas as pd
import finnhub
from openai import OpenAI

from io import StringIO
import requests


# ---------- 0  CONFIG ---------------------------------------------------------

OPENAI_MODEL  = os.getenv("OPENAI_MODEL", "gpt-4o-mini")
FINNHUB_KEY   = os.getenv("FINNHUB_API_KEY")
ALPHA_KEY = os.getenv("ALPHAVANTAGE_API_KEY")

if not FINNHUB_KEY:
    raise RuntimeError("FINNHUB_API_KEY not set")

client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
finnhub_client = finnhub.Client(api_key=FINNHUB_KEY)


SYSTEM_PROMPT = (
    "You are a seasoned stock-market analyst. "
    "Given recent company news and optional basic financials, "
    "return:\n"
    "[Positive Developments] – 2-4 bullets\n"
    "[Potential Concerns] – 2-4 bullets\n"
    "[Prediction & Analysis] – a one-week price outlook with rationale."
)


# ---------- 1  DATE / UTILITY HELPERS ----------------------------------------

def today() -> str:
    return date.today().strftime("%Y-%m-%d")

def n_weeks_before(date_string: str, n: int) -> str:
    return (datetime.strptime(date_string, "%Y-%m-%d") -
            timedelta(days=7 * n)).strftime("%Y-%m-%d")


# ---------- 2  DATA FETCHING --------------------------------------------------

def get_stock_data(symbol: str, steps: list[str]) -> pd.DataFrame:

    if not ALPHA_KEY:
        raise RuntimeError("ALPHAVANTAGE_API_KEY is Missing")

    # 免费端点：TIME_SERIES_DAILY :contentReference[oaicite:8]{index=8}
    url = (
        "https://www.alphavantage.co/query"
        "?function=TIME_SERIES_DAILY"
        f"&symbol={symbol}"
        f"&apikey={ALPHA_KEY}"
        "&datatype=csv"
        "&outputsize=full"
    )

    # 重试 3 次
    text = None
    for attempt in range(3):
        resp = requests.get(url, timeout=10)
        if not resp.ok:
            time.sleep(1)
            continue
        text = resp.text.strip()
        if text.startswith("{"):
            info = resp.json()
            msg = info.get("Note") or info.get("Error Message") or str(info)
            raise RuntimeError(f"Alpha Vantage Return Error：{msg}")
        break

    if not text:
        raise RuntimeError(f"Alpha Vantage Connection Error：{url}")

    df = pd.read_csv(StringIO(text))
    date_col = "timestamp" if "timestamp" in df.columns else df.columns[0]
    df[date_col] = pd.to_datetime(df[date_col])
    df = df.sort_values(date_col).set_index(date_col)

    data = {"Start Date": [], "End Date": [], "Start Price": [], "End Price": []}
    for i in range(len(steps) - 1):
        s_date = pd.to_datetime(steps[i])
        e_date = pd.to_datetime(steps[i+1])
        seg = df.loc[s_date:e_date]
        if seg.empty:
            raise RuntimeError(
                f"Alpha Vantage 无法获取 {symbol} 在 {steps[i]} – {steps[i+1]} 的数据"
            )
        data["Start Date"].append(seg.index[0])
        data["Start Price"].append(seg["close"].iloc[0])
        data["End Date"].append(seg.index[-1])
        data["End Price"].append(seg["close"].iloc[-1])
        # Limits：5 times/min
        time.sleep(12)

    return pd.DataFrame(data)

def current_basics(symbol: str, curday: str) -> dict:
    raw = finnhub_client.company_basic_financials(symbol, "all")
    if not raw["series"]:
        return {}
    merged = defaultdict(dict)
    for metric, vals in raw["series"]["quarterly"].items():
        for v in vals:
            merged[v["period"]][metric] = v["v"]

    latest = max((p for p in merged if p <= curday), default=None)
    if latest is None:
        return {}
    d = dict(merged[latest])
    d["period"] = latest
    return d

def attach_news(symbol: str, df: pd.DataFrame) -> pd.DataFrame:
    news_col = []
    for _, row in df.iterrows():
        start = row["Start Date"].strftime("%Y-%m-%d")
        end   = row["End Date"].strftime("%Y-%m-%d")
        time.sleep(1)                                        # Finnhub QPM guard
        weekly = finnhub_client.company_news(symbol, _from=start, to=end)
        weekly_fmt = [
            {
                "date"    : datetime.fromtimestamp(n["datetime"]).strftime("%Y%m%d%H%M%S"),
                "headline": n["headline"],
                "summary" : n["summary"],
            }
            for n in weekly
        ]
        weekly_fmt.sort(key=lambda x: x["date"])
        news_col.append(json.dumps(weekly_fmt))
    df["News"] = news_col
    return df


# ---------- 3  PROMPT CONSTRUCTION -------------------------------------------

def sample_news(news: list[str], k: int = 5) -> list[str]:
    if len(news) <= k: return news
    return [news[i] for i in sorted(random.sample(range(len(news)), k))]


def make_prompt(symbol: str, df: pd.DataFrame, curday: str, use_basics=False) -> str:
    # Company profile
    prof = finnhub_client.company_profile2(symbol=symbol)
    company_blurb = (
        f"[Company Introduction]:\n{prof['name']} operates in the "
        f"{prof['finnhubIndustry']} sector ({prof['country']}). "
        f"Founded {prof['ipo']}, market cap {prof['marketCapitalization']:.1f} "
        f"{prof['currency']}; ticker {symbol} on {prof['exchange']}.\n"
    )

    # Past weeks block
    past_block = ""
    for _, row in df.iterrows():
        term = "increased" if row["End Price"] > row["Start Price"] else "decreased"
        head = (f"From {row['Start Date']:%Y-%m-%d} to {row['End Date']:%Y-%m-%d}, "
                f"{symbol}'s stock price {term} from "
                f"{row['Start Price']:.2f} to {row['End Price']:.2f}.")
        news_items = json.loads(row["News"])
        summaries  = [
            f"[Headline] {n['headline']}\n[Summary] {n['summary']}\n"
            for n in news_items
            if not n["summary"].startswith("Looking for stock market analysis")
        ]
        past_block += "\n" + head + "\n" + "".join(sample_news(summaries, 5))

    # Optional basic financials
    if use_basics:
        basics = current_basics(symbol, curday)
        if basics:
            basics_txt = "\n".join(f"{k}: {v}" for k, v in basics.items() if k != "period")
            basics_block = (f"\n[Basic Financials] (reported {basics['period']}):\n{basics_txt}\n")
        else:
            basics_block = "\n[Basic Financials]: not available\n"
    else:
        basics_block = "\n[Basic Financials]: not requested\n"

    horizon = f"{curday} to {n_weeks_before(curday, -1)}"
    final_user_msg = (
        company_blurb
        + past_block
        + basics_block
        + f"\nBased on all information before {curday}, analyse positive "
          "developments and potential concerns for {symbol}, then predict its "
          f"price movement for next week ({horizon})."
    )
    return final_user_msg


# ---------- 4  LLM CALL -------------------------------------------------------

def chat_completion(prompt: str,
                    model: str = OPENAI_MODEL,
                    temperature: float = 0.3,
                    stream: bool = False) -> str:

    response = client.chat.completions.create(
        model=model,
        temperature=temperature,
        stream=stream,
        messages=[
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user",   "content": prompt}
        ],
    )

    if stream:
        collected = []
        for chunk in response:
            delta = chunk.choices[0].delta.content or ""
            print(delta, end="", flush=True)
            collected.append(delta)
        print()
        return "".join(collected)

    # without stream
    return response.choices[0].message.content


# ---------- 5  MAIN ENTRY (CLI test) -----------------------------------------

def predict(symbol: str = "AAPL",
            curday: str = today(),
            n_weeks: int = 3,
            use_basics: bool = False,
            stream: bool = False) -> tuple[str, str]:
    steps = [n_weeks_before(curday, n) for n in range(n_weeks + 1)][::-1]
    df    = get_stock_data(symbol, steps)
    df    = attach_news(symbol, df)

    prompt_info = make_prompt(symbol, df, curday, use_basics)
    answer      = chat_completion(prompt_info, stream=stream)

    return prompt_info, answer



# ---------- 6  SETUP HF -----------------------------------------


def hf_predict(symbol, n_weeks, use_basics):
    # 1. get curday
    curday = date.today().strftime("%Y-%m-%d")
    # 2. call predict
    prompt, answer = predict(
        symbol=symbol.upper(),
        curday=curday,
        n_weeks=int(n_weeks),
        use_basics=bool(use_basics),
        stream=False
    )
    return prompt, answer

with gr.Blocks() as demo:
    gr.Markdown("FinRobot_Forecaster")
    with gr.Row():
        symbol = gr.Textbox(label="Ticker（eg. AAPL）", value="AAPL")
        n_weeks = gr.Slider(1, 6, value=3, step=1, label="Trace Back Weeks")
        use_basics = gr.Checkbox(label="Add Basic Financials", value=False)
    output_prompt = gr.Textbox(label="Model Prompt", lines=8)
    output_answer = gr.Textbox(label="Model Output", lines=12)
    btn = gr.Button("Run Forecaster")
    btn.click(fn=hf_predict,
              inputs=[symbol, n_weeks, use_basics],
              outputs=[output_prompt, output_answer])

if __name__ == "__main__":
    demo.launch()