56c56a7b4a
All checks were successful
CI / test (pull_request) Successful in 4m9s
新增成本跟踪模块，为对话级和任务级成本计算提供基础。 - Turn 模型新增 prompt_tokens、completion_tokens、total_tokens 字段 - OpenAI 协议适配器新增 parse_usage() 提取 token 使用量 - 新增 evaluation/cost_tracking.py 模块： - ModelPricing：模型定价配置 - TokenUsage：token 使用量聚合 - CostBreakdown：成本明细 - calculate_cost()：根据 token 使用量和定价计算费用 - calculate_turn_cost()、calculate_case_cost()、calculate_run_cost() - 内置常见模型定价（GPT-4o、GPT-4o-mini、Claude 等） - 新增 11 项单元测试（677 tests passed） 注：引擎集成（实际捕获 API 调用的 token 使用量）留待后续实现。 Closes #25
73 lines
2.9 KiB
Python
73 lines
2.9 KiB
Python
"""OpenAI-compatible chat, embedding, and moderation protocol."""
|
|
|
|
from typing import Any
|
|
|
|
from agenteval.models import ModelCapability, ModelProtocol
|
|
from agenteval.utils.llm import extract_content_from_llm_response
|
|
|
|
from .base import ModelProtocolAdapter, ProtocolAdapterError
|
|
|
|
|
|
class OpenAICompatibleAdapter(ModelProtocolAdapter):
|
|
protocol = ModelProtocol.OPENAI_COMPATIBLE
|
|
supported_capabilities = frozenset(ModelCapability)
|
|
|
|
def headers(self, api_key: str | None) -> dict[str, str]:
|
|
headers = super().headers(api_key)
|
|
if api_key:
|
|
headers["Authorization"] = f"Bearer {api_key}"
|
|
return headers
|
|
|
|
def chat_payload(
|
|
self,
|
|
model_name: str | None,
|
|
messages: list[dict[str, str]],
|
|
temperature: float,
|
|
) -> dict[str, Any]:
|
|
return {"model": self.require_model(model_name), "messages": messages, "temperature": temperature}
|
|
|
|
def parse_chat(self, data: dict[str, Any]) -> str:
|
|
content = extract_content_from_llm_response(data)
|
|
if not content:
|
|
raise ProtocolAdapterError("模型接口返回内容为空")
|
|
return content
|
|
|
|
def parse_usage(self, data: dict[str, Any]) -> dict[str, int] | None:
|
|
"""Extract token usage from OpenAI-compatible response."""
|
|
usage = data.get("usage")
|
|
if not usage or not isinstance(usage, dict):
|
|
return None
|
|
return {
|
|
"prompt_tokens": usage.get("prompt_tokens", 0),
|
|
"completion_tokens": usage.get("completion_tokens", 0),
|
|
"total_tokens": usage.get("total_tokens", 0),
|
|
}
|
|
|
|
def embedding_payload(self, model_name: str | None, inputs: str | list[str]) -> dict[str, Any]:
|
|
return {"model": self.require_model(model_name), "input": inputs}
|
|
|
|
def parse_embeddings(self, data: dict[str, Any]) -> list[list[float]]:
|
|
try:
|
|
rows = sorted(data["data"], key=lambda item: item.get("index", 0))
|
|
vectors = [row["embedding"] for row in rows]
|
|
except (KeyError, TypeError) as exc:
|
|
raise ProtocolAdapterError("Embedding 接口返回格式不正确") from exc
|
|
if not all(isinstance(vector, list) for vector in vectors):
|
|
raise ProtocolAdapterError("Embedding 接口返回格式不正确")
|
|
return vectors
|
|
|
|
def moderation_payload(self, model_name: str | None, text: str) -> dict[str, Any]:
|
|
payload: dict[str, Any] = {"input": text}
|
|
if model_name:
|
|
payload["model"] = model_name
|
|
return payload
|
|
|
|
def parse_moderation(self, data: dict[str, Any]) -> dict[str, Any]:
|
|
try:
|
|
result = data["results"][0]
|
|
except (KeyError, IndexError, TypeError) as exc:
|
|
raise ProtocolAdapterError("Moderation 接口返回格式不正确") from exc
|
|
if not isinstance(result, dict):
|
|
raise ProtocolAdapterError("Moderation 接口返回格式不正确")
|
|
return result
|