141 lines
4.4 KiB
Python
141 lines
4.4 KiB
Python
import json
|
||
import time
|
||
from dataclasses import dataclass
|
||
from typing import Literal
|
||
|
||
import httpx
|
||
from pydantic import BaseModel, Field, ValidationError
|
||
|
||
|
||
Sentiment = Literal["positive", "negative", "neutral", "unknown"]
|
||
|
||
|
||
@dataclass(frozen=True)
|
||
class AIAnalysisResult:
|
||
comment_id: str
|
||
sentiment: str
|
||
labels: list[str]
|
||
reason: str = ""
|
||
ai_analysis_status: str = "success"
|
||
|
||
|
||
class AIAnalysisItem(BaseModel):
|
||
comment_id: str
|
||
sentiment: Sentiment
|
||
labels: list[str] = Field(default_factory=list, max_length=3)
|
||
reason: str = ""
|
||
|
||
|
||
def build_comment_prompt(comments: list[dict[str, str]]) -> str:
|
||
payload = [
|
||
{
|
||
"comment_id": str(comment["comment_id"]),
|
||
"content": str(comment.get("content", ""))[:150],
|
||
}
|
||
for comment in comments
|
||
]
|
||
return (
|
||
"只返回 JSON Array，不输出 Markdown 或解释性自然语言。"
|
||
"请原样回填输入中的 comment_id，不得修改或生成新 ID。"
|
||
"sentiment 只能是 positive、negative、neutral、unknown；labels 最多 3 个简短中文短语。\n"
|
||
f"{json.dumps(payload, ensure_ascii=False)}"
|
||
)
|
||
|
||
|
||
def parse_ai_comment_response(response_text: str, *, expected_comment_ids: set[str]) -> list[AIAnalysisResult]:
|
||
try:
|
||
raw = json.loads(response_text)
|
||
except json.JSONDecodeError as exc:
|
||
raise ValueError("AI response is not valid JSON") from exc
|
||
if not isinstance(raw, list):
|
||
raise ValueError("AI response must be a JSON Array")
|
||
|
||
results = []
|
||
for item in raw:
|
||
try:
|
||
parsed = AIAnalysisItem.model_validate(item)
|
||
except ValidationError as exc:
|
||
raise ValueError("AI response item does not match schema") from exc
|
||
if parsed.comment_id not in expected_comment_ids:
|
||
raise ValueError("AI response comment_id does not match input")
|
||
results.append(
|
||
AIAnalysisResult(
|
||
comment_id=parsed.comment_id,
|
||
sentiment=parsed.sentiment,
|
||
labels=parsed.labels,
|
||
reason=parsed.reason,
|
||
)
|
||
)
|
||
return results
|
||
|
||
|
||
def calculate_analysis_status(*, success_count: int, total_count: int) -> tuple[float, str]:
|
||
if total_count <= 0:
|
||
return 0.0, "insufficient"
|
||
rate = success_count / total_count
|
||
return rate, "normal" if rate >= 0.8 else "insufficient"
|
||
|
||
|
||
def analyze_comments_with_retry(
|
||
comments: list[dict[str, str]],
|
||
*,
|
||
requester,
|
||
max_retries: int = 3,
|
||
) -> list[AIAnalysisResult]:
|
||
expected_ids = {str(comment["comment_id"]) for comment in comments}
|
||
prompt = build_comment_prompt(comments)
|
||
for attempt in range(max_retries):
|
||
try:
|
||
return parse_ai_comment_response(requester(prompt), expected_comment_ids=expected_ids)
|
||
except Exception:
|
||
if attempt >= max_retries - 1:
|
||
break
|
||
time.sleep(2**attempt)
|
||
return [
|
||
AIAnalysisResult(
|
||
comment_id=comment_id,
|
||
sentiment="unknown",
|
||
labels=[],
|
||
reason="ai_parse_failed",
|
||
ai_analysis_status="failed",
|
||
)
|
||
for comment_id in expected_ids
|
||
]
|
||
|
||
|
||
class OpenAICompatibleAIClient:
|
||
def __init__(
|
||
self,
|
||
*,
|
||
base_url: str,
|
||
api_key: str,
|
||
model: str,
|
||
timeout_seconds: int = 30,
|
||
http_client: httpx.Client | None = None,
|
||
) -> None:
|
||
self.base_url = base_url.rstrip("/")
|
||
self.api_key = api_key
|
||
self.model = model
|
||
self._http_client = http_client or httpx.Client(timeout=timeout_seconds)
|
||
|
||
def request(self, prompt: str) -> str:
|
||
endpoint = f"{self.base_url}/chat/completions" if self.base_url.endswith("/v1") else f"{self.base_url}/v1/chat/completions"
|
||
response = self._http_client.post(
|
||
endpoint,
|
||
headers={"Authorization": f"Bearer {self.api_key}"},
|
||
json={
|
||
"model": self.model,
|
||
"messages": [
|
||
{
|
||
"role": "system",
|
||
"content": "只返回 JSON Array，不输出 Markdown 或解释性自然语言。",
|
||
},
|
||
{"role": "user", "content": prompt},
|
||
],
|
||
"temperature": 0,
|
||
},
|
||
)
|
||
response.raise_for_status()
|
||
payload = response.json()
|
||
return str(payload["choices"][0]["message"]["content"])
|