98 lines
3.5 KiB
Python
98 lines
3.5 KiB
Python
"""Semantic similarity evaluation rule.
|
||
|
||
Uses an external embedding API to compute cosine similarity between the
|
||
agent reply and a reference answer. Requires an OpenAI-compatible
|
||
embeddings endpoint (POST /v1/embeddings or equivalent).
|
||
|
||
Configuration params:
|
||
api_url Embeddings API endpoint (required)
|
||
api_key Bearer token (optional)
|
||
model Embedding model name (default: text-embedding-3-small)
|
||
reference Reference text to compare against (required)
|
||
min_score Minimum cosine similarity to pass, 0-1 (default: 0.7)
|
||
"""
|
||
|
||
import asyncio
|
||
import math
|
||
|
||
import httpx
|
||
|
||
from agenteval.evaluation.rules.base import EvalRule, RuleResult, register_rule
|
||
from agenteval.models import Case, Turn
|
||
from agenteval.utils.llm import extract_reply_text
|
||
|
||
|
||
def _cosine(a: list[float], b: list[float]) -> float:
|
||
dot = sum(x * y for x, y in zip(a, b))
|
||
norm_a = math.sqrt(sum(x * x for x in a))
|
||
norm_b = math.sqrt(sum(x * x for x in b))
|
||
if norm_a == 0 or norm_b == 0:
|
||
return 0.0
|
||
return dot / (norm_a * norm_b)
|
||
|
||
|
||
async def _embed(
|
||
client: httpx.AsyncClient,
|
||
api_url: str,
|
||
api_key: str | None,
|
||
model: str,
|
||
text: str,
|
||
) -> list[float]:
|
||
headers = {"Content-Type": "application/json"}
|
||
if api_key:
|
||
headers["Authorization"] = f"Bearer {api_key}"
|
||
resp = await client.post(
|
||
api_url,
|
||
headers=headers,
|
||
json={"model": model, "input": text},
|
||
timeout=30,
|
||
)
|
||
resp.raise_for_status()
|
||
data = resp.json()
|
||
return data["data"][0]["embedding"]
|
||
|
||
|
||
@register_rule
|
||
class SemanticSimilarityRule(EvalRule):
|
||
"""Score reply by cosine similarity to a reference answer via embedding API."""
|
||
|
||
name = "semantic_similarity"
|
||
|
||
async def evaluate(self, case: Case, dialog: list[Turn]) -> RuleResult:
|
||
if not dialog:
|
||
return RuleResult(passed=False, reason="无回复记录")
|
||
|
||
reply_text = extract_reply_text(dialog[-1].reply)
|
||
if not reply_text:
|
||
return RuleResult(passed=False, reason="回复内容为空")
|
||
|
||
reference: str | None = self.params.get("reference")
|
||
min_score: float = float(self.params.get("min_score", 0.7))
|
||
|
||
if not reference:
|
||
return RuleResult(passed=False, reason="semantic_similarity 未配置 reference")
|
||
|
||
try:
|
||
if self.model_config and self.gateway:
|
||
reply_vec, ref_vec = await self.gateway.embed(self.model_config, [reply_text, reference])
|
||
else:
|
||
api_url: str | None = self.params.get("api_url")
|
||
api_key: str | None = self.params.get("api_key")
|
||
model: str = self.params.get("model", "text-embedding-3-small")
|
||
if not api_url:
|
||
return RuleResult(passed=False, reason="semantic_similarity 未绑定向量模型(兼容配置缺少 api_url)")
|
||
async with httpx.AsyncClient() as client:
|
||
reply_vec, ref_vec = await asyncio.gather(
|
||
_embed(client, api_url, api_key, model, reply_text),
|
||
_embed(client, api_url, api_key, model, reference),
|
||
)
|
||
similarity = _cosine(reply_vec, ref_vec)
|
||
passed = similarity >= min_score
|
||
return RuleResult(
|
||
passed=passed,
|
||
score=round(similarity, 4),
|
||
reason=f"语义相似度 {similarity:.3f}(阈值 {min_score})",
|
||
)
|
||
except Exception as exc:
|
||
return RuleResult(passed=False, reason=f"embedding 调用失败: {exc}")
|