code/ai_finance_agent/backend/app/infrastructure/ai.py

"""真实 HTTP 接入与熔断，未配置或解析失败明确降级，不构造虚假 AI 输出。"""
import asyncio
import hashlib
import json
import math
import time
from datetime import timedelta
from typing import TypeVar
import httpx
from pydantic import BaseModel, ValidationError
from sqlalchemy import select
from sqlalchemy.dialects.sqlite import insert
from app.domain.models import CacheEntry
from app.infrastructure.database import Session, transaction
from app.shared.common import BusinessError, utcnow
from app.shared.config import settings

T = TypeVar("T", bound=BaseModel)


class AIUnavailable(Exception):
    """分类为外部服务失败，上层可回规则，但不能标为成功。"""


class CircuitBreaker:
    """半开只允许一个探测请求，避免冷却结束时全部等待请求涌向故障服务。"""
    def __init__(self):
        self.failures = 0
        self.opened_at = None
        self.probing = False
        self.lock = asyncio.Lock()

    async def enter(self):
        async with self.lock:
            if self.opened_at is not None:
                if time.monotonic() - self.opened_at < settings.ai_recovery_seconds or self.probing:
                    raise AIUnavailable("服务熔断中")
                self.probing = True

    async def finish(self, success: bool):
        async with self.lock:
            self.probing = False
            if success:
                self.failures, self.opened_at = 0, None
            else:
                self.failures += 1
                if self.opened_at is not None or self.failures >= settings.ai_failure_threshold:
                    self.opened_at = time.monotonic()

    def status(self):
        return "half_open" if self.probing else "open" if self.opened_at is not None else "closed"


def parse_structured(raw: str, schema: type[T]) -> T:
    """用 JSON 解码器识别嵌套对象，而非贪婪正则吞掉两个独立 JSON。"""
    candidates = [raw.strip()]
    decoder = json.JSONDecoder()
    for position, char in enumerate(raw):
        if char == "{":
            try:
                item, end = decoder.raw_decode(raw[position:])
                if isinstance(item, dict):
                    candidates.append(json.dumps(item, ensure_ascii=False))
            except json.JSONDecodeError:
                continue
    for candidate in candidates:
        try:
            return schema.model_validate_json(candidate)
        except (ValidationError, ValueError):
            continue
    raise AIUnavailable("模型结构化输出不符合契约")


class AIClient:
    def __init__(self):
        self.http = httpx.AsyncClient(timeout=settings.ai_timeout, follow_redirects=False)
        self.semaphore = asyncio.Semaphore(settings.ai_concurrency)
        self.breakers = {}

    async def close(self):
        await self.http.aclose()

    async def cached(self, namespace: str, model: str, payload: dict, ttl: int, operation):
        """模型名、端点和参数都参与键生成；只缓存成功响应，不让短暂故障长期固化。"""
        key = hashlib.sha256(json.dumps([namespace, model, payload], sort_keys=True, ensure_ascii=False).encode()).hexdigest()
        try:
            async with Session() as session:
                row = await session.get(CacheEntry, key)
                if row and row.expires_at > utcnow():
                    value = dict(row.value_json)
                    value["cache_hit"] = True
                    return value
        except Exception:
            # 缓存不是业务前提，读取失败仍调用真实服务。
            pass
        value = await operation()
        try:
            async with transaction() as session:
                await session.execute(insert(CacheEntry).values(key=key, namespace=namespace, model_version=model,
                    value_json=value, expires_at=utcnow() + timedelta(seconds=ttl)).on_conflict_do_update(
                    index_elements=["key"], set_={"value_json": value, "expires_at": utcnow() + timedelta(seconds=ttl)}))
        except Exception:
            pass
        return value

    async def request(self, url: str, payload: dict, timeout: float | None = None):
        """HTTP 层不自动重试；图级重试掌握总尝试预算，避免叠加指数放大。"""
        if not url.startswith(("http://", "https://")):
            raise AIUnavailable("AI 服务未配置")
        breaker = self.breakers.setdefault(url, CircuitBreaker())
        await breaker.enter()
        success = False
        try:
            async with self.semaphore:
                response = await self.http.post(url, json=payload, timeout=timeout or settings.ai_timeout,
                    headers={"Authorization": "Bearer " + settings.ai_api_key} if settings.ai_api_key else {})
                response.raise_for_status()
                result = response.json()
            success = True
            return result
        except (httpx.HTTPError, ValueError, KeyError):
            raise AIUnavailable("AI 服务不可用或响应格式错误")
        finally:
            await breaker.finish(success)

    async def structured(self, prompt: str, schema: type[T], vision: str | None = None, timeout: float | None = None):
        """OpenAI 兼容 JSON Schema 请求；不支持时明确上抛，图级预算控制降级。"""
        base = settings.vlm_base_url or settings.llm_base_url if vision else settings.llm_base_url
        model = settings.vlm_model or settings.llm_model if vision else settings.llm_model
        if not base or not model:
            raise AIUnavailable("模型未配置")
        content = [{"type": "text", "text": prompt}, {"type": "image_url", "image_url": {"url": vision}}] if vision else prompt
        payload = {"model": model, "messages": [{"role": "user", "content": content}], "temperature": 0.1,
            "response_format": {"type": "json_schema", "json_schema": {"name": schema.__name__, "schema": schema.model_json_schema()}}}

        async def invoke():
            response = await self.request(base.rstrip("/") + "/chat/completions", payload, timeout)
            try:
                data = parse_structured(response["choices"][0]["message"]["content"], schema)
            except (KeyError, IndexError, TypeError):
                raise AIUnavailable("模型未返回可解析的消息")
            return {"result": data.model_dump(mode="json"), "usage": response.get("usage", {})}

        value = await self.cached("llm", model, {"base": base, **payload}, 86400, invoke)
        usage = {"cache_hit": True} if value.get("cache_hit") else value.get("usage", {})
        return schema.model_validate(value["result"]), usage

    async def embeddings(self, texts: list[str]) -> list[list[float]]:
        """不混用不同模型的向量；服务失败可走同模型 OpenAI embeddings 协议。"""
        async def invoke():
            response = None
            if settings.embedding_base_url:
                try:
                    response = await self.request(settings.embedding_base_url.rstrip("/") + "/embeddings", {"texts": texts})
                except AIUnavailable:
                    response = None
            if response is None and settings.llm_base_url:
                response = await self.request(settings.llm_base_url.rstrip("/") + "/embeddings", {"model": settings.embedding_model, "input": texts})
            if response is None:
                raise AIUnavailable("Embedding 服务未配置")
            try:
                vectors = response.get("vectors")
                if vectors is None:
                    vectors = [item["embedding"] for item in sorted(response["data"], key=lambda item: item["index"])]
                if len(vectors) != len(texts):
                    raise ValueError()
                for vector in vectors:
                    if len(vector) != settings.embedding_dimension or not all(isinstance(x, (int, float)) and math.isfinite(x) for x in vector) or sum(x * x for x in vector) <= 0:
                        raise ValueError()
            except (KeyError, TypeError, ValueError):
                raise AIUnavailable("Embedding 数量、维度或数值不合法")
            return {"vectors": vectors}

        value = await self.cached("embed", settings.embedding_model, {"texts": texts, "base": settings.embedding_base_url,
            "fallback": settings.llm_base_url, "dim": settings.embedding_dimension}, 604800, invoke)
        return value["vectors"]

    async def rerank(self, query: str, candidates: list[dict]) -> list[dict]:
        """精排失败保留原 COSINE 排序，同时返回降级标记供界面解释。"""
        if not settings.rerank_base_url or not candidates:
            return candidates
        try:
            response = await self.request(settings.rerank_base_url.rstrip("/") + "/rerank",
                                          {"query": query, "documents": [item["text"] for item in candidates]})
            results = response["results"]
            indexes = [item["index"] for item in results]
            if len(indexes) != len(candidates) or set(indexes) != set(range(len(candidates))):
                raise ValueError()
            ordered = []
            for result in results:
                score = float(result["relevance_score"])
                if not math.isfinite(score):
                    raise ValueError()
                ordered.append({**candidates[result["index"]], "rerank_score": score})
            return sorted(ordered, key=lambda item: item["rerank_score"], reverse=True)
        except (AIUnavailable, KeyError, TypeError, ValueError):
            return [{**item, "rerank_degraded": True} for item in candidates]


ai = AIClient()