用 AI 搭建个人投研系统：从数据管线到语义检索的四层架构

import akshare as ak
import pandas as pd

SCHEMA = ["date", "open", "high", "low", "close", "volume"]

def fetch_price(symbol: str, start: str, end: str) -> pd.DataFrame:
    """日线行情，前复权固定不变"""
    df = ak.stock_zh_a_hist(
        symbol=symbol, period="daily",
        start_date=start, end_date=end, adjust="qfq")
    df = df.rename(columns={
        "日期": "date", "开盘": "open", "收盘": "close",
        "最高": "high", "最低": "low", "成交量": "volume"})
    df["date"] = pd.to_datetime(df["date"])
    df = df.set_index("date").sort_index()
    return df[SCHEMA[1:]].astype(float)

def fetch_financial(symbol: str) -> pd.DataFrame:
    """财务指标，保留公告日用于 PIT 对齐"""
    df = ak.stock_financial_abstract(symbol=symbol)
    df["ann_date"] = pd.to_datetime(df["公告日期"])
    return df.dropna(subset=["ann_date"]).sort_values("ann_date")

import duckdb

con = duckdb.connect("research.db")

def init_schema():
    con.execute("""
    CREATE TABLE IF NOT EXISTS price (
        symbol VARCHAR, date DATE,
        open DOUBLE, high DOUBLE, low DOUBLE,
        close DOUBLE, volume DOUBLE,
        PRIMARY KEY (symbol, date)
    )""")
    con.execute("""
    CREATE TABLE IF NOT EXISTS financial (
        symbol VARCHAR, ann_date DATE, period VARCHAR,
        revenue DOUBLE, net_profit DOUBLE,
        PRIMARY KEY (symbol, ann_date, period)
    )""")

def upsert_price(symbol: str, df: pd.DataFrame):
    tmp = df.reset_index()
    tmp.insert(0, "symbol", symbol)
    con.execute("INSERT OR REPLACE INTO price SELECT * FROM tmp")

def query_pit(symbol: str, as_of: str) -> pd.DataFrame:
    """按公告日查询：只返回当时已知的财务数据"""
    return con.execute("""
        SELECT * FROM financial
        WHERE symbol = ? AND ann_date <= ?
        ORDER BY ann_date DESC LIMIT 1
    """, [symbol, as_of]).df()

import numpy as np

def compute_factors(df: pd.DataFrame,
                    windows=(5, 20, 60)) -> pd.DataFrame:
    f = pd.DataFrame(index=df.index)
    ret = df["close"].pct_change()
    for w in windows:
        f[f"mom_{w}"] = df["close"].pct_change(w)
        f[f"vol_{w}"] = ret.rolling(w).std()
        f[f"vr_{w}"] = df["volume"] / df["volume"].rolling(w).mean()
    f["ma20"] = df["close"].rolling(20).mean()
    f["ma60"] = df["close"].rolling(60).mean()
    f["trend"] = (f["ma20"] > f["ma60"]).astype(int)
    return f.dropna()

import chromadb
from openai import OpenAI

client = OpenAI(api_key="sk-xxx", base_url="https://api.deepseek.com/v1")
col = chromadb.PersistentClient(path="./ann_db") \
    .get_or_create_collection("announcements")

def index_announcement(doc_id: str, text: str, meta: dict):
    chunks = [text[i:i+500] for i in range(0, len(text), 450)]
    col.add(
        ids=[f"{doc_id}_{i}" for i in range(len(chunks))],
        documents=chunks,
        embeddings=[client.embeddings.create(
            model="text-embedding-3-small",
            input=c).data[0].embedding for c in chunks],
        metadatas=[meta] * len(chunks))

def search_ann(q: str, k: int = 3) -> list[dict]:
    emb = client.embeddings.create(
        model="text-embedding-3-small", input=q).data[0].embedding
    hits = col.query(query_embeddings=[emb], n_results=k)
    return [{"text": d, "meta": m}
            for d, m in zip(hits["documents"][0], hits["metadatas"][0])]

import json

def build_memo(symbol: str, as_of: str) -> str:
    px = con.execute(
        "SELECT * FROM price WHERE symbol=? AND date<=? "
        "ORDER BY date DESC LIMIT 120", [symbol, as_of]).df()
    factors = compute_factors(px.set_index("date"))
    fin = query_pit(symbol, as_of)
    anns = search_ann(f"{symbol} 产能 扩张", k=3)

    ctx = {
        "factors": factors.iloc[-1].round(4).to_dict(),
        "financial": fin.to_dict("records"),
        "announcements": [a["text"][:200] for a in anns],
    }
    prompt = f"""以下是某标的的公开数据。请用客观语言归纳当前状态。
严禁给出买卖建议，严禁预测价格，严禁评价估值高低。
数据：{json.dumps(ctx, ensure_ascii=False)}
输出：一段不超过 300 字的客观描述。"""
    resp = client.chat.completions.create(
        model="deepseek-chat",
        messages=[{"role": "user", "content": prompt}],
        temperature=0.2)
    return resp.choices[0].message.content