gate_scout_order/onchain_scout_gate/app/gemma_client.py

from __future__ import annotations

import json
import logging
import re
from typing import Any

import httpx

from .config import GemmaConfig

LOGGER = logging.getLogger("onchain_scout.gemma_client")


def _extract_json_object(text: str) -> dict[str, Any] | None:
    text = text.strip()
    m = re.search(r"\{[\s\S]*\}", text)
    if not m:
        return None
    raw = m.group(0)
    try:
        return json.loads(raw)
    except json.JSONDecodeError:
        return None


class OllamaGemmaClient:
    def __init__(self, conf: GemmaConfig) -> None:
        self.conf = conf
        self.timeout = httpx.Timeout(conf.timeout_seconds, read=conf.timeout_seconds + 30.0)

    async def rank_funnel(
        self,
        symbol: str,
        programmatic_text: str,
        ohlc_csv_block: str,
        image_base64: str | None,
    ) -> dict[str, Any]:
        """
        调用本地 Ollama，让 Gemma 按漏斗标准 JSON 回复。
        """
        system = (
            "你是加密货币永续合约的日线结构分析师。只输出一个 JSON 对象，不要 Markdown，不要代码围栏。"
            "字段必须全部存在且为英文枚举/数字："
            '{"daily_structure":"strong|ok|weak",'
            '"volume_view":"high|mid|low",'
            '"upside_space":"high|mid|low",'
            '"mid_resistance":"low|mid|high",'
            '"priority":1-10整数,'
            '"one_liner":"中文一句"}。'
            "priority 越高越值得优先关注：成交大、日线结构好、上方空间大、中间阻力小则给高分。"
        )
        user_body = (
            f"标的 {symbol} USDT 永续。\n"
            f"程序化摘要:\n{programmatic_text}\n\n"
            f"最近日线 OHLCV（时间正序最后一行为最新）:\n{ohlc_csv_block}\n"
        )
        url = f"{self.conf.ollama_base_url.rstrip('/')}/api/chat"
        message: dict[str, Any] = {"role": "user", "content": user_body}
        if image_base64 and self.conf.send_chart_image:
            message["images"] = [image_base64]

        payload: dict[str, Any] = {
            "model": self.conf.model,
            "messages": [{"role": "system", "content": system}, message],
            "stream": False,
            "options": {"temperature": self.conf.temperature},
        }
        if self.conf.json_mode:
            payload["format"] = "json"

        async with httpx.AsyncClient(timeout=self.timeout, trust_env=False) as client:
            resp = await client.post(url, json=payload)
            resp.raise_for_status()
            data = resp.json()

        msg = (data.get("message") or {}).get("content") or ""
        parsed = _extract_json_object(msg) if msg else None
        if parsed is None and isinstance(data.get("message"), dict):
            parsed = _extract_json_object(str(data["message"]))
        if parsed is None:
            LOGGER.warning("gemma_parse_failed symbol=%s raw_len=%s", symbol, len(msg))
            return {
                "error": "parse_failed",
                "raw": msg[:2000],
                "daily_structure": "weak",
                "volume_view": "low",
                "upside_space": "low",
                "mid_resistance": "high",
                "priority": 1,
                "one_liner": "模型输出无法解析为 JSON",
            }
        return _normalize_gemma_dict(parsed)

    async def generate_daily_report(self, report_day_cn: str, btc_snapshot: dict, stats: dict) -> dict[str, Any]:
        system = (
            "你是加密交易复盘助手。输出严格 JSON 对象，不要 Markdown。字段必须存在："
            '{"headline":"...","btc_explain":"...","summary":"...","risk_points":["..."],"action_hint":"..."}。'
            "用中文，简洁专业，不写投资建议免责声明。"
        )
        user_body = (
            f"请生成 {report_day_cn} 的晨报。\n"
            f"BTC 快照: {json.dumps(btc_snapshot, ensure_ascii=False)}\n"
            f"昨日统计: {json.dumps(stats, ensure_ascii=False)}\n"
            "要求：1) headline 一句话；2) btc_explain 解释方向；"
            "3) summary 覆盖 WATCH/TRIGGER/漏斗；4) risk_points 给1-3条；5) action_hint 给执行提示。"
        )
        url = f"{self.conf.ollama_base_url.rstrip('/')}/api/chat"
        payload: dict[str, Any] = {
            "model": self.conf.model,
            "messages": [{"role": "system", "content": system}, {"role": "user", "content": user_body}],
            "stream": False,
            "options": {"temperature": 0.1},
            "format": "json",
        }
        async with httpx.AsyncClient(timeout=self.timeout, trust_env=False) as client:
            resp = await client.post(url, json=payload)
            resp.raise_for_status()
            data = resp.json()

        msg = (data.get("message") or {}).get("content") or ""
        parsed = _extract_json_object(msg) if msg else None
        if parsed is None:
            return {"error": "parse_failed", "raw": msg[:1200]}
        risk = parsed.get("risk_points")
        if not isinstance(risk, list):
            risk = [str(risk or "")]
        risk = [str(x)[:120] for x in risk if str(x or "").strip()][:3] or ["注意高波动时的回撤风险。"]
        return {
            "headline": str(parsed.get("headline") or "")[:120],
            "btc_explain": str(parsed.get("btc_explain") or "")[:220],
            "summary": str(parsed.get("summary") or "")[:360],
            "risk_points": risk,
            "action_hint": str(parsed.get("action_hint") or "")[:220],
        }


def _normalize_gemma_dict(d: dict[str, Any]) -> dict[str, Any]:
    def _enum(v: Any, choices: set[str], default: str) -> str:
        s = str(v or "").strip().lower()
        return s if s in choices else default

    try:
        pr = int(float(d.get("priority", 1)))
    except (TypeError, ValueError):
        pr = 1
    pr = max(1, min(10, pr))
    return {
        "daily_structure": _enum(d.get("daily_structure"), {"strong", "ok", "weak"}, "weak"),
        "volume_view": _enum(d.get("volume_view"), {"high", "mid", "low"}, "low"),
        "upside_space": _enum(d.get("upside_space"), {"high", "mid", "low"}, "low"),
        "mid_resistance": _enum(d.get("mid_resistance"), {"low", "mid", "high"}, "high"),
        "priority": pr,
        "one_liner": str(d.get("one_liner") or "")[:280],
    }