Files
Artemis-Eapil/apps/admin_console/services/favorites.py
T

467 lines
20 KiB
Python

"""Independent favorite snapshots dispatched through the existing task services."""
from __future__ import annotations
from copy import deepcopy
import asyncio
import hashlib
import json
import time
from typing import Literal
from uuid import NAMESPACE_URL, uuid5
from weakref import WeakValueDictionary
from pydantic import BaseModel, ConfigDict, Field
from apps.admin_console.database.repositories.session_repository import SessionRepository
from apps.admin_console.services.automation_store import AutomationStore, identifier, store
from apps.admin_console.services.execution_config import execution_config
from apps.admin_console.services.model_service import model_service
from apps.admin_console.services.task_import import TaskCase, TaskOptions, execution_title
from artemis.config.exploration import ExplorationConfig
from artemis.exploration.case_prompt import require_resolved_case
from artemis.exploration.store import ExplorationStore
class FavoriteCreate(BaseModel):
model_config = ConfigDict(extra="forbid")
source_kind: Literal["task", "batch", "exploration", "exploration_case"] = "task"
source_id: str = Field(pattern=r"^[A-Za-z0-9_-]{1,100}$")
version: int | None = Field(default=None, ge=1)
name: str | None = Field(default=None, min_length=1, max_length=200)
class FavoriteUpdate(BaseModel):
model_config = ConfigDict(extra="forbid")
name: str | None = Field(default=None, min_length=1, max_length=200)
task_config: TaskOptions | None = None
case_configs: dict[int, TaskOptions] = Field(default_factory=dict)
class FavoriteRun(FavoriteUpdate):
request_id: str = Field(pattern=r"^[A-Za-z0-9_-]{1,100}$")
def _merge_options(target: dict, updates: dict) -> None:
if isinstance(updates.get("recording"), dict):
updates = {
**updates,
"recording": {**(target.get("recording") or {}), **updates["recording"]},
}
target.update(updates)
class FavoriteService:
def __init__(self, storage: AutomationStore = store):
self.store = storage
self._run_locks: WeakValueDictionary = WeakValueDictionary()
def get(self, key: str) -> dict:
value = self.store.get("favorite", identifier(key))
if value is None:
raise KeyError("收藏不存在")
if value.get("config_version") != 1:
# Complete legacy snapshots once and preserve their list position.
# Re-read under the write lock so a concurrent edit cannot be lost.
with self.store.connection() as conn:
conn.execute("BEGIN IMMEDIATE")
row = conn.execute(
"SELECT payload FROM automation_records WHERE kind='favorite' AND id=?", (key,)
).fetchone()
if row is None:
raise KeyError("收藏不存在")
value = self._complete(json.loads(row[0]))
conn.execute(
"UPDATE automation_records SET payload=? WHERE kind='favorite' AND id=?",
(json.dumps(value, ensure_ascii=False), key),
)
conn.commit()
return value
def list(self) -> list[dict]:
return [
self.get(item["id"]) if item.get("config_version") != 1 else item
for item in self.store.records("favorite", 10000)
]
def _complete(self, favorite: dict) -> dict:
result = deepcopy(favorite)
if result.get("config_version") == 1:
return result
snapshot = result["snapshot"]
if result["kind"] == "batch":
for case in snapshot["cases"]:
case["task_config"] = execution_config(
{
"profile": snapshot.get("profile"),
"device_serial": snapshot.get("device_serial"),
"runtime_overrides": {
"timeout_seconds": snapshot.get("timeout_seconds"),
**(case.get("task_config") or {}),
},
}
)["runtime_overrides"]
elif result["kind"] == "exploration":
config = ExplorationConfig.model_validate(snapshot.get("config") or {}).model_dump()
snapshot["config"] = config
snapshot["task_config"] = execution_config(
{
"runtime_overrides": {
**(snapshot.get("task_config") or {}),
**{
key: config[key]
for key in ("model", "reasoning_effort", "timeout_seconds")
},
}
}
)["runtime_overrides"]
else:
snapshot.update(execution_config(snapshot))
result["config_version"] = 1
return result
def _task(self, session_id: str, *, fallback: dict | None = None) -> tuple[dict, list[str]]:
job = self.store.job(session_id) or {}
if snapshot := job.get("execution_config"):
return execution_config(snapshot), []
row = SessionRepository(self.store.path).get_session_by_id(session_id) or {}
submission = job.get("submission") or fallback or {}
options = {
**(job.get("case") or {}).get("task_config", {}),
**(submission.get("runtime_overrides") or {}),
}
options = {key: value for key, value in options.items() if value is not None}
try:
device = row.get("device_info") or {}
device = json.loads(device) if isinstance(device, str) else device
except (ValueError, TypeError):
device = {}
device = device if isinstance(device, dict) else {}
profile = (
options.get("profile")
or submission.get("profile")
or model_service.resolve_session_profile(row)
or "flash"
)
options["profile"] = profile
serial = (
options.get("device_serial")
or submission.get("device_serial")
or job.get("device_serial")
or device.get("serial")
)
if serial:
options["device_serial"] = serial
if serial == "computer:local" or device.get("execution_target") == "computer":
options["execution_target"] = "computer"
# Historical diagnostics describe the actual model, unlike display metadata
# which may reflect today's defaults. Read only execution model roles.
with self.store.connection() as conn:
rows = conn.execute(
"SELECT payload FROM traces WHERE session_id=? AND type='llm_diagnostic' ORDER BY timestamp DESC",
(session_id,),
).fetchall()
for trace in rows:
try:
data = json.loads(trace[0])
except (ValueError, TypeError):
continue
if str(data.get("role", "")).lower() in {
"operator",
"flashrunner",
"reactiverunner",
"planner",
}:
for key in ("model", "reasoning_effort"):
if data.get(key):
options.setdefault(key, data[key])
if options.get("model"):
break
missing = [
key
for key in (
"model",
"reasoning_effort",
"record_video",
"recording",
"analyze_video",
"request_timeout_seconds",
"max_retries",
"retry_interval_seconds",
"timeout_seconds",
)
if key not in options
]
snapshot = execution_config(
{**submission, "profile": profile, "runtime_overrides": options}
)
return snapshot, missing
def create(self, value: FavoriteCreate) -> dict:
token = f"{value.source_kind}:{value.source_id}:{value.version or ''}"
key = hashlib.sha256(token.encode()).hexdigest()[:32]
if self.store.get("favorite", key):
return self.get(key)
snapshot, missing, kind, version = {}, [], value.source_kind, None
source_id = value.source_id
repo = ExplorationStore(self.store)
if kind == "task":
run = repo.run(source_id)
if run and run["kind"] == "exploration":
source_id = run["exploration_id"]
kind = "exploration_case" if repo.case(source_id) else "exploration"
if kind in {"exploration", "exploration_case"}:
session = repo.get(source_id)
options = deepcopy(session["task_config"])
if kind == "exploration":
snapshot = {
"goal": session["goal"],
"config": session["config"],
"task_config": options,
}
name = session["goal"][:120]
else:
case = repo.case(source_id, value.version)
if not case:
raise KeyError("确认用例不存在")
version = case["version"]
validations = [
r
for r in repo.runs(source_id)
if r["kind"] == "validation" and r.get("case_version") == version
]
if validations:
config, missing = self._task(
validations[-1]["id"],
fallback={"runtime_overrides": validations[-1]["task_config"]},
)
else:
options = {
k: v
for k, v in options.items()
if k not in {"profile", "model", "reasoning_effort", "timeout_seconds"}
}
config = execution_config({"profile": "flash", "runtime_overrides": options})
missing = ["首次用例复跑,默认 Flash"]
snapshot = {"goal": require_resolved_case(case), **config}
name = case["name"]
elif kind == "batch":
batch = self.store.get("batch", source_id)
if not batch or batch.get("status") == "deleting":
raise KeyError("批次不存在")
cases, seen = [], set()
for job in self.store.jobs(source_id):
position = job.get("case_position", job["position"])
if position in seen:
continue
seen.add(position)
case = deepcopy(job["case"])
config, inherited = self._task(
job["session_id"],
fallback={
"profile": batch["profile"],
"device_serial": batch.get("device_serial"),
"runtime_overrides": {
"timeout_seconds": batch.get("timeout_seconds"),
**case.get("task_config", {}),
},
},
)
case["task_config"] = config["runtime_overrides"]
cases.append(TaskCase.model_validate(case).model_dump(exclude_none=True))
missing.extend(inherited)
if not cases:
raise ValueError("批次没有可收藏的用例")
snapshot = {
key: batch[key]
for key in (
"name",
"formats",
"profile",
"device_serial",
"timeout_seconds",
"stop_on_failure",
"repeat_count",
)
if key in batch
}
snapshot["cases"] = cases
name = batch["name"]
else:
row = SessionRepository(self.store.path).get_session_by_id(source_id) or {}
job = self.store.job(source_id) or {}
goal = (
(job.get("case") or {}).get("goal")
or row.get("initial_goal")
or (job.get("submission") or {}).get("goal")
)
if not goal:
raise KeyError("任务不存在或没有可复用的提示词")
config, missing = self._task(source_id)
snapshot = {"goal": goal, **config}
name = execution_title(job) or goal[:120]
favorite = {
"id": key,
"name": (value.name or name).strip(),
"kind": kind,
"source_kind": value.source_kind,
"source_id": value.source_id,
"case_version": version,
"snapshot": snapshot,
"inherited_fields": sorted(set(missing)),
"created_at": time.time(),
"updated_at": time.time(),
}
if not favorite["name"]:
raise ValueError("收藏名称不能为空")
return self.store.put("favorite", key, self._complete(favorite))
def updated(self, favorite: dict, value: FavoriteUpdate) -> dict:
result = self._complete(favorite)
if value.name is not None:
if not value.name.strip():
raise ValueError("收藏名称不能为空")
result["name"] = value.name.strip()
snapshot = result["snapshot"]
options = value.task_config.model_dump(exclude_none=True) if value.task_config else {}
if result["kind"] == "batch":
for index, case in enumerate(snapshot["cases"]):
if override := value.case_configs.get(index):
_merge_options(case["task_config"], override.model_dump(exclude_none=True))
_merge_options(case["task_config"], options)
if set(value.case_configs) - set(range(len(snapshot["cases"]))):
raise ValueError("批次用例下标不存在")
if options.get("device_serial"):
snapshot["device_serial"] = options["device_serial"]
elif result["kind"] == "exploration":
config = snapshot["config"]
config.update(
{
k: options[k]
for k in ("model", "reasoning_effort", "timeout_seconds")
if k in options
}
)
snapshot["config"] = ExplorationConfig.model_validate(config).model_dump()
_merge_options(snapshot["task_config"], options)
else:
_merge_options(snapshot["runtime_overrides"], options)
return result
def update(self, key: str, value: FavoriteUpdate) -> dict:
favorite = self.updated(self.get(key), value)
favorite["updated_at"] = time.time()
return self.store.put("favorite", key, favorite)
def delete(self, key: str) -> None:
with self.store.connection() as conn:
conn.execute(
"DELETE FROM automation_records WHERE kind='favorite' AND id=?", (identifier(key),)
)
conn.commit()
async def run(self, key: str, value: FavoriteRun) -> dict:
lock = self._run_locks.setdefault(value.request_id, asyncio.Lock())
async with lock:
return await self._run(key, value)
async def _run(self, key: str, value: FavoriteRun) -> dict:
key = identifier(key)
request = {"favorite_id": key, **value.model_dump(exclude_none=True)}
signature = hashlib.sha256(json.dumps(request, sort_keys=True).encode()).hexdigest()
# Reserve a durable snapshot before admission. Retries survive restarts,
# edits/deletions of the favorite, and lost HTTP responses.
with self.store.connection() as conn:
conn.execute("BEGIN IMMEDIATE")
row = conn.execute(
"SELECT payload FROM automation_records WHERE kind='favorite_run' AND id=?",
(value.request_id,),
).fetchone()
record = json.loads(row[0]) if row else None
if record and record["signature"] != signature:
raise ValueError("request_id 已用于其他内容")
if record and record.get("result"):
return record["result"]
if record is None:
row = conn.execute(
"SELECT payload FROM automation_records WHERE kind='favorite' AND id=?", (key,)
).fetchone()
if row is None:
raise KeyError("收藏不存在")
record = {
"signature": signature,
"favorite": self.updated(json.loads(row[0]), value),
}
conn.execute(
"INSERT INTO automation_records VALUES (?,?,?,?)",
(
"favorite_run",
value.request_id,
json.dumps(record, ensure_ascii=False),
time.time(),
),
)
conn.commit()
favorite = record["favorite"]
snapshot = favorite["snapshot"]
run_id = str(uuid5(NAMESPACE_URL, "artemis-favorite:" + value.request_id))
if favorite["kind"] == "batch":
from apps.admin_console.services.automation_service import create_batch, batch_action
batch = create_batch(
{**snapshot, "name": favorite["name"], "request_id": "favorite-" + run_id}
)
if batch["status"] == "draft":
batch = batch_action(batch["id"], "resume")
result = {"kind": "batch", "batch_id": batch["id"], "status": batch["status"]}
elif favorite["kind"] == "exploration":
from apps.admin_console.services.exploration_service import exploration_service
from artemis.exploration.models import CreateExploration
session = exploration_service.create(CreateExploration(request_id=run_id, **snapshot))
admission = await exploration_service.start(session["id"], run_id)
run = admission.get("run", admission)
rejection = admission.get("admission", {})
outcome = run.get("outcome") or {}
if (
rejection.get("status") in {"error", "rejected"}
or outcome.get("status") == "rejected"
):
raise ValueError(
rejection.get("error")
or rejection.get("message")
or outcome.get("error")
or "探索未入队"
)
result = {
"kind": "exploration",
"session_id": run["id"],
"exploration_id": session["id"],
"admission": admission,
}
else:
from apps.admin_console.routers.tasks import run_task
from apps.admin_console.schemas.task_schema import RunRequest
admission = await run_task(
RunRequest(
goal=snapshot["goal"],
profile=snapshot["runtime_overrides"].get("profile", "flash"),
runtime_overrides=TaskOptions.model_validate(snapshot["runtime_overrides"]),
session_id=run_id,
ingress="favorite",
expected_output=snapshot.get("expected_output"),
enable_outputter=snapshot.get("enable_outputter"),
app_path=snapshot.get("app_path"),
)
)
if not admission.get("tasks"):
raise ValueError(admission.get("error") or admission.get("message") or "任务未入队")
result = {"kind": "task", "session_id": run_id, "admission": admission}
record["result"] = result
self.store.put("favorite_run", value.request_id, record)
return result
favorites = FavoriteService()