first commit
This commit is contained in:
192
evaluator/engine.py
Normal file
192
evaluator/engine.py
Normal file
@@ -0,0 +1,192 @@
|
||||
from __future__ import annotations
|
||||
import asyncio, inspect, json, math, random
|
||||
from datetime import date, datetime
|
||||
from typing import Any, Awaitable, Callable
|
||||
from evaluator.collectors.base import ConversationCollector
|
||||
from evaluator.collectors.langfuse import LangfuseCollector
|
||||
from evaluator.collectors.agent_framework import AgentFrameworkCollector
|
||||
from evaluator.collectors.mock import MockCollector
|
||||
from evaluator.config.agents import AgentConfig
|
||||
from evaluator.config.settings import settings
|
||||
from evaluator.core.models import ConversationRecord, RunStatus
|
||||
from evaluator.judges.llm_judge import TIMStyleLLMJudge
|
||||
from evaluator.output.legacy_exporter import export_legacy_txt_gz
|
||||
from evaluator.persistence.repository import EvaluationRepository
|
||||
from evaluator.publishers.langfuse_scores import LangfuseScorePublisher
|
||||
from evaluator.publishers.langfuse_job_audit import LangfuseJobAuditPublisher
|
||||
from evaluator.publishers.gcs_export import upload_evaluation_file
|
||||
|
||||
ProgressCallback = Callable[[dict[str, Any]], Awaitable[None] | None]
|
||||
|
||||
|
||||
class EvaluationEngine:
|
||||
def __init__(self, repository: EvaluationRepository | None = None, progress_callback: ProgressCallback | None = None):
|
||||
self.repository = repository or EvaluationRepository(auto_init_schema=False)
|
||||
self.progress_callback = progress_callback
|
||||
self.judge = TIMStyleLLMJudge()
|
||||
self.langfuse_publisher = LangfuseScorePublisher()
|
||||
self.job_audit = LangfuseJobAuditPublisher()
|
||||
|
||||
async def _emit(self, progress_run_id: str, stage: str, message: str = "", **details):
|
||||
details.pop("run_id", None)
|
||||
await self.repository.arecord_progress(progress_run_id, stage, message, details)
|
||||
event = {"run_id": progress_run_id, "stage": stage, "message": message, "details": details}
|
||||
if self.progress_callback:
|
||||
r = self.progress_callback(event)
|
||||
if inspect.isawaitable(r):
|
||||
await r
|
||||
|
||||
def collector_for(self, source: str, agent_id: str | None = None) -> ConversationCollector:
|
||||
if source == 'langfuse': return LangfuseCollector(agent_id=agent_id)
|
||||
if source == 'agent_framework': return AgentFrameworkCollector()
|
||||
if source == 'mock': return MockCollector()
|
||||
raise ValueError('source must be langfuse, agent_framework or mock')
|
||||
|
||||
async def _audit(self, progress_run_id: str, audit_trace_id: str, agent_id: str, status: str, period_start: datetime, period_end: datetime, *, required: bool = True, **details) -> bool:
|
||||
try:
|
||||
result = await self.job_audit.publish(
|
||||
audit_trace_id, agent_id, status, period_start=period_start, period_end=period_end, details=details
|
||||
)
|
||||
if not result or not result.get("ok"):
|
||||
raise RuntimeError(f"Langfuse audit did not confirm persistence: {result!r}")
|
||||
await self._emit(progress_run_id, 'LANGFUSE_JOB_AUDIT', f'judge-job-{agent_id} status={status}', audit_status=status, audit_result=result)
|
||||
return True
|
||||
except Exception as exc:
|
||||
await self._emit(progress_run_id, 'LANGFUSE_JOB_AUDIT_FAILED', f'Langfuse job audit failed: {exc}', audit_status=status, error=str(exc))
|
||||
if required:
|
||||
raise
|
||||
return False
|
||||
|
||||
async def run_agent(self, agent: AgentConfig, period_start: datetime, period_end: datetime, source: str = 'langfuse', limit: int | None = None, reference_date: date | None = None) -> dict:
|
||||
run_id = await self.repository.acreate_run(period_start, period_end, source, agent.agent_id)
|
||||
reference_date = reference_date or period_start.date()
|
||||
audit_trace_id = self.job_audit.new_trace_id()
|
||||
try:
|
||||
# Emit RUN_CREATED immediately after persistence so CLI/monitor can expose
|
||||
# the RUN_ID before any external Langfuse audit/network call is performed.
|
||||
await self._emit(run_id, 'RUN_CREATED', f'Agent run created: {agent.agent_id}', agent_id=agent.agent_id, source=source, reference_date=reference_date.isoformat())
|
||||
await self._audit(run_id, audit_trace_id, agent.agent_id, 'created', period_start, period_end, run_id=run_id, source=source)
|
||||
collector = self.collector_for(source, agent.agent_id)
|
||||
await self._emit(run_id, 'COLLECTING', 'Collecting sessions and conversations')
|
||||
seed = f'{agent.agent_id}:{reference_date.isoformat()}'
|
||||
if isinstance(collector, LangfuseCollector):
|
||||
records, sampled_sessions, collection_stats = await collector.collect_sampled(
|
||||
period_start, period_end, agent_aliases=agent.aliases, percentage=agent.percentage, seed=seed, limit=limit
|
||||
)
|
||||
await self._emit(
|
||||
run_id, 'COLLECTED', f'Collected {len(records)} traces from {len(sampled_sessions)} sampled sessions',
|
||||
**collection_stats
|
||||
)
|
||||
else:
|
||||
records = await collector.collect(period_start, period_end, agent_aliases=agent.aliases, limit=limit)
|
||||
records, sampled_sessions = self._sample_sessions(records, agent.percentage, agent.agent_id, reference_date)
|
||||
collection_stats = {}
|
||||
await self._emit(run_id, 'COLLECTED', f'Collected {len(records)} traces before session sampling')
|
||||
await self._emit(
|
||||
run_id,
|
||||
'SAMPLED',
|
||||
f'Kept {len(sampled_sessions)} sessions / {len(records)} traces',
|
||||
percentage=agent.percentage,
|
||||
seed=seed,
|
||||
sampled_sessions=sampled_sessions,
|
||||
)
|
||||
inserted = await self.repository.ainsert_items(run_id, records)
|
||||
await self._emit(run_id, 'ITEMS_INSERTED', f'Inserted {inserted} items')
|
||||
summary = await self._process(run_id)
|
||||
output_path = export_legacy_txt_gz(self.repository, run_id, agent.agent_id)
|
||||
await self._emit(run_id, 'EXPORTED', f'Exported {output_path}', output_file=str(output_path))
|
||||
if not settings.enable_gcs_upload:
|
||||
raise RuntimeError("TIM specification requires GCS upload for run-agents; set ENABLE_GCS_UPLOAD=true")
|
||||
uploaded_to = await asyncio.to_thread(upload_evaluation_file, output_path, agent.gcs_prefix)
|
||||
await self._emit(run_id, 'GCS_UPLOADED', f'Uploaded {uploaded_to}', uploaded_to=uploaded_to)
|
||||
output_path.unlink()
|
||||
await self._emit(run_id, 'LOCAL_FILE_DELETED', 'Local export deleted after successful GCS upload')
|
||||
result = {**summary, 'agent_id': agent.agent_id, 'output_file': str(output_path), 'uploaded_to': uploaded_to}
|
||||
# _audit() receives status as an explicit positional argument. The run
|
||||
# summary also contains a "status" key (e.g. COMPLETED), so forwarding
|
||||
# **result unchanged would pass status twice and raise:
|
||||
# TypeError: _audit() got multiple values for argument 'status'
|
||||
audit_details = dict(result)
|
||||
# Never forward names that are explicit _audit() parameters.
|
||||
for reserved_key in ('progress_run_id', 'audit_trace_id', 'agent_id', 'status', 'period_start', 'period_end'):
|
||||
audit_details.pop(reserved_key, None)
|
||||
await self._audit(run_id, audit_trace_id, agent.agent_id, 'completed', period_start, period_end, **audit_details)
|
||||
return result
|
||||
except Exception as exc:
|
||||
await self.repository.amark_run_status(run_id, RunStatus.FAILED, str(exc))
|
||||
await self._emit(run_id, 'FAILED', f'Agent run failed: {exc}', error=str(exc))
|
||||
await self._audit(run_id, audit_trace_id, agent.agent_id, 'failed', period_start, period_end, required=False, run_id=run_id, error=str(exc))
|
||||
return {'status': 'FAILED', 'run_id': run_id, 'agent_id': agent.agent_id, 'error': str(exc)}
|
||||
|
||||
async def run(self, period_start: datetime, period_end: datetime, source: str = 'langfuse', limit: int | None = None) -> dict:
|
||||
run_id = await self.repository.acreate_run(period_start, period_end, source, None)
|
||||
try:
|
||||
collector = self.collector_for(source)
|
||||
await self._emit(run_id, 'COLLECTING', 'Collecting conversations')
|
||||
records = await collector.collect(period_start, period_end, limit=limit)
|
||||
await self._emit(run_id, 'COLLECTED', f'Collected {len(records)} records')
|
||||
await self.repository.ainsert_items(run_id, records)
|
||||
return await self._process(run_id)
|
||||
except Exception as exc:
|
||||
await self.repository.amark_run_status(run_id, RunStatus.PARTIAL, str(exc))
|
||||
await self._emit(run_id, 'PARTIAL', f'Run failed: {exc}', error=str(exc))
|
||||
return {'status': 'PARTIAL', 'run_id': run_id, 'error': str(exc)}
|
||||
|
||||
async def _process(self, run_id: str) -> dict:
|
||||
processed_records: list[ConversationRecord] = []
|
||||
while True:
|
||||
items = await self.repository.afetch_next_items(run_id, settings.batch_size)
|
||||
if not items: break
|
||||
await self._emit(run_id, 'BATCH_STARTED', f'Processing {len(items)} items')
|
||||
for item in items:
|
||||
item_id = item['item_id']
|
||||
await self.repository.amark_item_processing(item_id)
|
||||
try:
|
||||
raw = item['raw_json']
|
||||
if hasattr(raw, 'read'): raw = raw.read()
|
||||
record = ConversationRecord.model_validate(json.loads(raw))
|
||||
result = await self.judge.judge_trace(record)
|
||||
await self.repository.asave_trace_result(run_id, item_id, record, result)
|
||||
await self.langfuse_publisher.publish_trace_score(record, result)
|
||||
await self.repository.amark_item_completed(run_id, item_id)
|
||||
processed_records.append(record)
|
||||
await self._emit(
|
||||
run_id,
|
||||
"ITEM_COMPLETED",
|
||||
f"Item completed {item_id}",
|
||||
trace_id=record.trace_id,
|
||||
session_id=record.session_id,
|
||||
judgeScore=result.judgeScore,
|
||||
accuracyScore=result.accuracyScore,
|
||||
alucinationScore=result.alucinationScore,
|
||||
rationale=result.rationale,
|
||||
)
|
||||
except Exception as exc:
|
||||
await self.repository.amark_item_failed(run_id, item_id, str(exc))
|
||||
await self._emit(run_id, 'ITEM_FAILED', f'Item failed {item_id}', error=str(exc))
|
||||
if processed_records:
|
||||
sessions = await self.judge.judge_sessions(processed_records)
|
||||
for sid, result in sessions.items():
|
||||
agent_id = next((r.agent_id for r in processed_records if r.session_id == sid), None)
|
||||
await self.repository.asave_session_result(run_id, sid, agent_id, result)
|
||||
await self._emit(run_id, 'SESSION_JUDGE_COMPLETED', f'Evaluated {len(sessions)} sessions')
|
||||
await self.repository.amark_run_status(run_id, RunStatus.COMPLETED)
|
||||
summary = await self.repository.asummarize_run(run_id)
|
||||
await self._emit(run_id, 'COMPLETED', 'Run completed', **summary)
|
||||
return {'status': 'COMPLETED', **summary}
|
||||
|
||||
def _sample_sessions(self, records: list[ConversationRecord], percentage: float, agent_id: str, reference_date: date) -> tuple[list[ConversationRecord], list[str]]:
|
||||
"""Deterministic session sampling; all traces from selected sessions are retained."""
|
||||
if not records:
|
||||
return [], []
|
||||
session_ids = sorted({r.session_id for r in records})
|
||||
if percentage >= 1:
|
||||
selected = session_ids
|
||||
elif percentage <= 0:
|
||||
selected = []
|
||||
else:
|
||||
count = min(len(session_ids), max(1, math.ceil(len(session_ids) * percentage)))
|
||||
rng = random.Random(f"{agent_id}:{reference_date.isoformat()}")
|
||||
selected = sorted(rng.sample(session_ids, count))
|
||||
selected_set = set(selected)
|
||||
return [r for r in records if r.session_id in selected_set], selected
|
||||
Reference in New Issue
Block a user