feat(console): add Artemis Admin & Trace Console backend and management plane
Introduces the comprehensive developer/administrator control plane and trace
inspection console backend for the ARTEMIS mobile automation platform.
Key Features & Backend Architecture:
- Core Server & Modular Routing:
* FastAPI modular architecture with dedicated routers for sessions, steps, tasks, media, stream, and replay.
* Multi-layer repository pattern with SQLite persistence for sessions, traces, and step data.
- Task Control & Queue Management:
* FIFO asynchronous task queue service with support for single/batch goal execution and profile selection (Flash / Pro).
* Session lifecycle controls (run, pause, resume, stop, clear queue).
* Startup orphan cleanup and automatic older replay archiving.
- Live Telemetry & IPC Communication:
* High-performance IPC server for streaming real-time execution logs, agent thoughts, and tool invocations.
* Server-Sent Events (SSE) /api/stream/{session_id} endpoint for live UI updates.
- Real-Time Device Streaming:
* Live ADB screen capture and MJPEG frame generator (/api/stream/device-live).
* Device connectivity and state inspection endpoint (/api/stream/device-state).
- Trace Inspection & Step Replay Engine:
* Hierarchical call tree inspection, pre/post step screenshots, and action overlay metadata.
* Isolated single-step replay manager (/api/replay) for interactive debugging and validation.
This commit is contained in:
Weilin Liu committed
1 parent
2fd70465e6
commit
c1803fcb86
29 files changed
+8889
No files matched your search
@@ -0,0 +1 @@
|
||||
# Copyright 2026 Google LLC
|
||||
@@ -0,0 +1,43 @@
|
||||
# Copyright 2026 Google LLC
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
from contextlib import contextmanager
|
||||
from pathlib import Path
|
||||
import sqlite3
|
||||
from fastapi import HTTPException
|
||||
|
||||
try:
|
||||
from admin_console.core.config import DB_PATH
|
||||
except ImportError:
|
||||
from apps.admin_console.core.config import DB_PATH
|
||||
|
||||
|
||||
def get_db(db_path: Path | None = None) -> sqlite3.Connection:
|
||||
"""Returns a SQLite connection with timeout and Row row_factory."""
|
||||
path = db_path or DB_PATH
|
||||
if not path.exists():
|
||||
raise HTTPException(status_code=500, detail=f"Database not found at {path}")
|
||||
conn = sqlite3.connect(path, timeout=10.0)
|
||||
conn.row_factory = sqlite3.Row
|
||||
return conn
|
||||
|
||||
|
||||
@contextmanager
|
||||
def db_session(db_path: Path | None = None):
|
||||
"""Context manager for SQLite connections ensuring clean closure."""
|
||||
conn = get_db(db_path)
|
||||
try:
|
||||
yield conn
|
||||
finally:
|
||||
conn.close()
|
||||
@@ -0,0 +1 @@
|
||||
# Copyright 2026 Google LLC
|
||||
@@ -0,0 +1,202 @@
|
||||
# Copyright 2026 Google LLC
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
import time
|
||||
from typing import Any
|
||||
|
||||
try:
|
||||
from admin_console.database.connection import db_session
|
||||
except ImportError:
|
||||
from apps.admin_console.database.connection import db_session
|
||||
|
||||
|
||||
class SessionRepository:
|
||||
"""Repository handling all database queries and updates for Sessions."""
|
||||
|
||||
def __init__(self, db_path=None):
|
||||
self.db_path = db_path
|
||||
|
||||
def get_all_sessions(self) -> list[dict[str, Any]]:
|
||||
with db_session(self.db_path) as conn:
|
||||
cursor = conn.cursor()
|
||||
cursor.execute("SELECT * FROM sessions ORDER BY start_time DESC")
|
||||
return [dict(r) for r in cursor.fetchall()]
|
||||
|
||||
def get_session_by_id(self, session_id: str) -> dict[str, Any] | None:
|
||||
with db_session(self.db_path) as conn:
|
||||
cursor = conn.cursor()
|
||||
cursor.execute("SELECT * FROM sessions WHERE session_id = ?", (str(session_id),))
|
||||
row = cursor.fetchone()
|
||||
return dict(row) if row else None
|
||||
|
||||
def get_video_recordings_map(self) -> dict[str, str]:
|
||||
with db_session(self.db_path) as conn:
|
||||
cursor = conn.cursor()
|
||||
cursor.execute(
|
||||
"SELECT name FROM sqlite_master WHERE type='table' AND name='video_recordings'"
|
||||
)
|
||||
if cursor.fetchone() is None:
|
||||
return {}
|
||||
cursor.execute(
|
||||
"SELECT session_id, local_video_path FROM video_recordings WHERE session_id IS NOT NULL AND local_video_path IS NOT NULL ORDER BY start_time ASC"
|
||||
)
|
||||
return {str(r[0]): str(r[1]) for r in cursor.fetchall() if r[0] and r[1]}
|
||||
|
||||
def get_llm_traces_for_profile(self, session_id: str, limit: int = 3) -> list[str]:
|
||||
with db_session(self.db_path) as conn:
|
||||
cursor = conn.cursor()
|
||||
cursor.execute(
|
||||
"SELECT payload FROM traces WHERE session_id = ? AND type = 'llm_call' LIMIT ?",
|
||||
(session_id, limit),
|
||||
)
|
||||
rows = cursor.fetchall()
|
||||
return [r[0] for r in rows if r and r[0]]
|
||||
|
||||
def get_agent_trace_names(self, session_id: str) -> list[str]:
|
||||
try:
|
||||
with db_session(self.db_path) as conn:
|
||||
cursor = conn.cursor()
|
||||
cursor.execute(
|
||||
"SELECT DISTINCT name FROM traces WHERE session_id = ? AND type IN ('agent', 'llm_call')",
|
||||
(session_id,),
|
||||
)
|
||||
rows = cursor.fetchall()
|
||||
return [r[0] for r in rows if r and r[0]]
|
||||
except Exception:
|
||||
return []
|
||||
|
||||
def get_session_by_id(self, session_id: str) -> dict[str, Any] | None:
|
||||
try:
|
||||
with db_session(self.db_path) as conn:
|
||||
cursor = conn.cursor()
|
||||
cursor.execute(
|
||||
"SELECT * FROM sessions WHERE session_id = ?",
|
||||
(str(session_id),),
|
||||
)
|
||||
row = cursor.fetchone()
|
||||
return dict(row) if row else None
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
def harvest_orphaned_sessions(self, orphaned_ids: list[str]) -> int:
|
||||
if not orphaned_ids:
|
||||
return 0
|
||||
with db_session(self.db_path) as conn:
|
||||
cursor = conn.cursor()
|
||||
now = time.time()
|
||||
for s_id in orphaned_ids:
|
||||
cursor.execute(
|
||||
"UPDATE sessions SET status = ?, end_time = ? "
|
||||
"WHERE session_id = ? AND status = ?",
|
||||
("failed", now, s_id, "running"),
|
||||
)
|
||||
conn.commit()
|
||||
return len(orphaned_ids)
|
||||
|
||||
def cleanup_orphans_on_startup(self) -> int:
|
||||
try:
|
||||
with db_session(self.db_path) as conn:
|
||||
cursor = conn.cursor()
|
||||
cursor.execute(
|
||||
"UPDATE sessions SET status = ?, end_time = ? WHERE status = ?",
|
||||
("failed", time.time(), "running"),
|
||||
)
|
||||
count = cursor.rowcount
|
||||
conn.commit()
|
||||
return count
|
||||
except Exception:
|
||||
return 0
|
||||
|
||||
def get_latest_session(self) -> dict[str, Any] | None:
|
||||
try:
|
||||
with db_session(self.db_path) as conn:
|
||||
cursor = conn.cursor()
|
||||
cursor.execute(
|
||||
"SELECT session_id, status FROM sessions ORDER BY start_time DESC LIMIT 1"
|
||||
)
|
||||
row = cursor.fetchone()
|
||||
return dict(row) if row else None
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
def get_running_session_id(self) -> str | None:
|
||||
try:
|
||||
with db_session(self.db_path) as conn:
|
||||
cursor = conn.cursor()
|
||||
cursor.execute(
|
||||
"SELECT session_id FROM sessions WHERE status = 'running' "
|
||||
"ORDER BY start_time DESC LIMIT 1"
|
||||
)
|
||||
row = cursor.fetchone()
|
||||
return row["session_id"] if row else None
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
def get_session_status(self, session_id: str) -> str | None:
|
||||
try:
|
||||
with db_session(self.db_path) as conn:
|
||||
cursor = conn.cursor()
|
||||
cursor.execute(
|
||||
"SELECT status FROM sessions WHERE session_id = ?",
|
||||
(str(session_id),),
|
||||
)
|
||||
row = cursor.fetchone()
|
||||
return row["status"] if row else None
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
def update_session_status(
|
||||
self, session_id: str, status: str, end_time: float | None = None
|
||||
) -> bool:
|
||||
try:
|
||||
with db_session(self.db_path) as conn:
|
||||
cursor = conn.cursor()
|
||||
t = end_time or time.time()
|
||||
cursor.execute(
|
||||
"UPDATE sessions SET status = ?, end_time = ? WHERE session_id = ?",
|
||||
(status, t, str(session_id)),
|
||||
)
|
||||
conn.commit()
|
||||
return cursor.rowcount > 0
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
def mark_all_running_cancelled(self) -> int:
|
||||
try:
|
||||
with db_session(self.db_path) as conn:
|
||||
cursor = conn.cursor()
|
||||
cursor.execute(
|
||||
"UPDATE sessions SET status = ?, end_time = ? WHERE status = ?",
|
||||
("cancelled", time.time(), "running"),
|
||||
)
|
||||
count = cursor.rowcount
|
||||
conn.commit()
|
||||
return count
|
||||
except Exception:
|
||||
return 0
|
||||
|
||||
def get_background_tasks(self, session_id: str) -> list[dict[str, Any]]:
|
||||
try:
|
||||
with db_session(self.db_path) as conn:
|
||||
cursor = conn.cursor()
|
||||
cursor.execute(
|
||||
"SELECT * FROM background_tasks WHERE session_id = ? ORDER BY start_time ASC",
|
||||
(session_id,),
|
||||
)
|
||||
return [dict(r) for r in cursor.fetchall()]
|
||||
except Exception:
|
||||
return []
|
||||
|
||||
|
||||
session_repo = SessionRepository()
|
||||
@@ -0,0 +1,411 @@
|
||||
# Copyright 2026 Google LLC
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
import json
|
||||
from typing import Any
|
||||
|
||||
try:
|
||||
from admin_console.database.connection import db_session
|
||||
except ImportError:
|
||||
from apps.admin_console.database.connection import db_session
|
||||
|
||||
|
||||
class StepRepository:
|
||||
"""Repository handling querying and formatting of steps and associated traces."""
|
||||
|
||||
def __init__(self, db_path=None):
|
||||
self.db_path = db_path
|
||||
|
||||
def get_step_session_id(self, step_id: str) -> str | None:
|
||||
with db_session(self.db_path) as conn:
|
||||
cursor = conn.cursor()
|
||||
cursor.execute("SELECT session_id FROM steps WHERE step_id = ?", (step_id,))
|
||||
row = cursor.fetchone()
|
||||
return row["session_id"] if row else None
|
||||
|
||||
@staticmethod
|
||||
def _clean_value(val: Any) -> Any:
|
||||
"""Recursively cleans values to remove non-serializable and internal
|
||||
Python runtime representations.
|
||||
"""
|
||||
if val is None:
|
||||
return None
|
||||
if isinstance(val, (bool, int, float)):
|
||||
return val
|
||||
if isinstance(val, str):
|
||||
if (
|
||||
"object at 0x" in val
|
||||
or val.startswith("<artemis.")
|
||||
or val.startswith("<controller")
|
||||
):
|
||||
return None
|
||||
return val
|
||||
if isinstance(val, dict):
|
||||
cleaned = {}
|
||||
for k, v in val.items():
|
||||
cv = StepRepository._clean_value(v)
|
||||
if cv is not None:
|
||||
cleaned[k] = cv
|
||||
return cleaned
|
||||
if isinstance(val, (list, tuple)):
|
||||
cleaned_list = []
|
||||
for item in val:
|
||||
cv = StepRepository._clean_value(item)
|
||||
if cv is not None:
|
||||
cleaned_list.append(cv)
|
||||
return cleaned_list
|
||||
val_str = str(val)
|
||||
if "object at 0x" in val_str or val_str.startswith("<"):
|
||||
return None
|
||||
return val_str
|
||||
|
||||
@staticmethod
|
||||
def _extract_clean_text_and_images(content: Any) -> tuple[str, int]:
|
||||
"""Extracts readable text and count of image blocks from LLM content."""
|
||||
if not content:
|
||||
return "", 0
|
||||
if isinstance(content, str):
|
||||
if content.startswith("data:image/") or "base64," in content[:50]:
|
||||
return "", 1
|
||||
return content, 0
|
||||
if isinstance(content, dict):
|
||||
if content.get("type") == "image_url" or "image_url" in content:
|
||||
return "", 1
|
||||
if content.get("type") == "text":
|
||||
return str(content.get("text", "")), 0
|
||||
t = content.get("text") or content.get("thought") or ""
|
||||
return str(t), 0
|
||||
if isinstance(content, list):
|
||||
total_text = ""
|
||||
num_images = 0
|
||||
for item in content:
|
||||
txt, imgs = StepRepository._extract_clean_text_and_images(item)
|
||||
total_text += txt
|
||||
num_images += imgs
|
||||
return total_text, num_images
|
||||
return str(content), 0
|
||||
|
||||
@staticmethod
|
||||
def _extract_token_usage_from_trace(payload_raw: Any) -> tuple[int, int, int]:
|
||||
"""Extract prompt, completion, and total tokens from an LLM trace payload with multimodal awareness."""
|
||||
if not payload_raw:
|
||||
return 0, 0, 0
|
||||
payload_obj = payload_raw
|
||||
if isinstance(payload_raw, str):
|
||||
try:
|
||||
payload_obj = json.loads(payload_raw)
|
||||
except Exception:
|
||||
pass
|
||||
if not isinstance(payload_obj, dict):
|
||||
return 0, 0, 0
|
||||
|
||||
u = payload_obj.get("token_usage") or payload_obj.get("usage_metadata")
|
||||
if isinstance(u, dict):
|
||||
pr = u.get("prompt_tokens") or u.get("prompt_token_count") or u.get("input_tokens") or 0
|
||||
co = (
|
||||
u.get("completion_tokens")
|
||||
or u.get("candidates_token_count")
|
||||
or u.get("output_tokens")
|
||||
or 0
|
||||
)
|
||||
to = u.get("total_tokens") or u.get("total_token_count") or (pr + co)
|
||||
if to > 0:
|
||||
return int(pr), int(co), int(to)
|
||||
|
||||
# Fallback for historical traces: estimate tokens from text and images
|
||||
msgs = payload_obj.get("messages", [])
|
||||
resp = payload_obj.get("response", [])
|
||||
|
||||
prompt_text = ""
|
||||
prompt_images = 0
|
||||
if isinstance(msgs, list):
|
||||
for m in msgs:
|
||||
if isinstance(m, dict):
|
||||
t, imgs = StepRepository._extract_clean_text_and_images(m.get("content"))
|
||||
prompt_text += t
|
||||
prompt_images += imgs
|
||||
else:
|
||||
t, imgs = StepRepository._extract_clean_text_and_images(m)
|
||||
prompt_text += t
|
||||
prompt_images += imgs
|
||||
else:
|
||||
t, imgs = StepRepository._extract_clean_text_and_images(msgs)
|
||||
prompt_text += t
|
||||
prompt_images += imgs
|
||||
|
||||
resp_text = ""
|
||||
resp_images = 0
|
||||
if isinstance(resp, list):
|
||||
for r in resp:
|
||||
if isinstance(r, dict):
|
||||
t, imgs = StepRepository._extract_clean_text_and_images(r.get("content"))
|
||||
resp_text += t
|
||||
resp_images += imgs
|
||||
else:
|
||||
t, imgs = StepRepository._extract_clean_text_and_images(r)
|
||||
resp_text += t
|
||||
resp_images += imgs
|
||||
|
||||
# Standard Gemini multimodal token pricing: 258 tokens per image
|
||||
pr = (len(prompt_text) // 4) + (prompt_images * 258)
|
||||
co = (len(resp_text) // 4) + (resp_images * 258)
|
||||
return int(pr), int(co), int(pr + co)
|
||||
|
||||
def _clean_tool_payload(self, payload_raw: Any, trace_type: str | None = None) -> Any:
|
||||
if not payload_raw:
|
||||
return None
|
||||
payload_obj = payload_raw
|
||||
if isinstance(payload_raw, str):
|
||||
try:
|
||||
payload_obj = json.loads(payload_raw)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
if not isinstance(payload_obj, dict):
|
||||
return None
|
||||
|
||||
if trace_type == "llm_call":
|
||||
return (
|
||||
{"error": self._clean_value(payload_obj.get("error"))}
|
||||
if "error" in payload_obj
|
||||
else None
|
||||
)
|
||||
|
||||
args = payload_obj.get("args", payload_obj)
|
||||
if isinstance(args, dict):
|
||||
return {"args": self._clean_value(args)}
|
||||
return {"args": self._clean_value(args)}
|
||||
|
||||
def get_session_steps(self, session_id: str, client: str | None = None) -> list[dict[str, Any]]:
|
||||
with db_session(self.db_path) as conn:
|
||||
cursor = conn.cursor()
|
||||
cursor.execute(
|
||||
"SELECT * FROM steps WHERE session_id = ? ORDER BY step_number ASC",
|
||||
(session_id,),
|
||||
)
|
||||
rows = cursor.fetchall()
|
||||
|
||||
steps = []
|
||||
for r in rows:
|
||||
step_dict = dict(r)
|
||||
|
||||
# Strip raw binary bytes for API efficiency (clients fetch images by url/hash)
|
||||
step_dict["pre_screenshot_bytes"] = None
|
||||
step_dict["post_screenshot_bytes"] = None
|
||||
|
||||
# Parse JSON fields safely
|
||||
for json_col in [
|
||||
"action_taken",
|
||||
"last_execution_result",
|
||||
"extra_metadata",
|
||||
]:
|
||||
if step_dict.get(json_col):
|
||||
try:
|
||||
step_dict[json_col] = json.loads(step_dict[json_col])
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# If post_image_name is identical to pre_image_name, normalize to None
|
||||
pre_img = step_dict.get("pre_image_name")
|
||||
if pre_img and step_dict.get("post_image_name") == pre_img:
|
||||
step_dict["post_image_name"] = None
|
||||
|
||||
# Fetch generic tool traces for this step
|
||||
step_id = step_dict.get("step_id")
|
||||
if step_id:
|
||||
cursor.execute(
|
||||
"SELECT t1.trace_id, t1.parent_trace_id, t1.step_id,"
|
||||
" t1.type, t1.name, t1.timestamp, t1.duration, t1.status,"
|
||||
" t1.payload, t2.name as agent_name FROM traces t1 LEFT"
|
||||
" JOIN traces t2 ON t1.parent_trace_id = t2.trace_id WHERE"
|
||||
" t1.step_id = ? AND (t1.type = 'tool' OR (t1.type ="
|
||||
" 'llm_call' AND t1.status = 'failed')) ORDER BY"
|
||||
" t1.timestamp ASC",
|
||||
(step_id,),
|
||||
)
|
||||
trace_rows = cursor.fetchall()
|
||||
generic_tools = []
|
||||
for tr in trace_rows:
|
||||
trace_dict = dict(tr)
|
||||
trace_dict["payload"] = self._clean_tool_payload(
|
||||
trace_dict.get("payload"), trace_dict.get("type")
|
||||
)
|
||||
generic_tools.append(trace_dict)
|
||||
step_dict["generic_tools"] = generic_tools
|
||||
|
||||
# Compute token usage for this step from all LLM call traces
|
||||
cursor.execute(
|
||||
"SELECT payload FROM traces WHERE step_id = ? AND type = 'llm_call'",
|
||||
(step_id,),
|
||||
)
|
||||
llm_rows = cursor.fetchall()
|
||||
step_p, step_c, step_t = 0, 0, 0
|
||||
for lr in llm_rows:
|
||||
p, c, t = self._extract_token_usage_from_trace(lr["payload"])
|
||||
step_p += p
|
||||
step_c += c
|
||||
step_t += t
|
||||
|
||||
# If step_t is 0, also check extra_metadata for token_usage
|
||||
if step_t == 0 and step_dict.get("extra_metadata"):
|
||||
meta = (
|
||||
step_dict["extra_metadata"]
|
||||
if isinstance(step_dict["extra_metadata"], dict)
|
||||
else {}
|
||||
)
|
||||
u = meta.get("token_usage")
|
||||
if isinstance(u, dict):
|
||||
step_p = int(u.get("prompt_tokens") or u.get("input_tokens") or 0)
|
||||
step_c = int(u.get("completion_tokens") or u.get("output_tokens") or 0)
|
||||
step_t = int(u.get("total_tokens") or (step_p + step_c))
|
||||
|
||||
step_dict["token_usage"] = {
|
||||
"prompt_tokens": step_p,
|
||||
"completion_tokens": step_c,
|
||||
"total_tokens": step_t,
|
||||
}
|
||||
step_dict["total_tokens"] = step_t
|
||||
else:
|
||||
step_dict["generic_tools"] = []
|
||||
step_dict["token_usage"] = {
|
||||
"prompt_tokens": 0,
|
||||
"completion_tokens": 0,
|
||||
"total_tokens": 0,
|
||||
}
|
||||
step_dict["total_tokens"] = 0
|
||||
|
||||
try:
|
||||
from artemis.utils.coordinates import normalize_step_actions
|
||||
|
||||
step_dict = normalize_step_actions(step_dict)
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
steps.append(step_dict)
|
||||
|
||||
# Fetch step-less tool traces (e.g. planner tools)
|
||||
cursor.execute(
|
||||
"SELECT t1.trace_id, t1.parent_trace_id, t1.step_id, t1.type,"
|
||||
" t1.name, t1.timestamp, t1.duration, t1.status, t1.payload,"
|
||||
" t2.name as agent_name FROM traces t1 LEFT JOIN traces t2 ON"
|
||||
" t1.parent_trace_id = t2.trace_id WHERE t1.session_id = ? AND"
|
||||
" (t1.step_id IS NULL OR t1.step_id = '') AND (t1.type = 'tool' OR"
|
||||
" (t1.type = 'llm_call' AND t1.status = 'failed')) ORDER BY"
|
||||
" t1.timestamp ASC",
|
||||
(session_id,),
|
||||
)
|
||||
stepless_rows = cursor.fetchall()
|
||||
if stepless_rows:
|
||||
unassigned_tools = []
|
||||
sorted_steps = sorted(steps, key=lambda s: s.get("timestamp") or 0)
|
||||
|
||||
for tr in stepless_rows:
|
||||
trace_dict = dict(tr)
|
||||
trace_dict["payload"] = self._clean_tool_payload(
|
||||
trace_dict.get("payload"), trace_dict.get("type")
|
||||
)
|
||||
|
||||
tr_time = trace_dict.get("timestamp") or 0
|
||||
target_step = None
|
||||
for s in sorted_steps:
|
||||
s_time = s.get("timestamp") or 0
|
||||
if s_time <= tr_time:
|
||||
target_step = s
|
||||
else:
|
||||
break
|
||||
|
||||
if target_step:
|
||||
if "generic_tools" not in target_step or not target_step["generic_tools"]:
|
||||
target_step["generic_tools"] = []
|
||||
target_step["generic_tools"].append(trace_dict)
|
||||
target_step["generic_tools"].sort(key=lambda t: t.get("timestamp") or 0)
|
||||
else:
|
||||
unassigned_tools.append(trace_dict)
|
||||
|
||||
if unassigned_tools:
|
||||
unassigned_tools.sort(key=lambda t: t.get("timestamp") or 0)
|
||||
first_timestamp = unassigned_tools[0].get("timestamp") or 0
|
||||
|
||||
# Calculate token usage ONLY for pre-planning LLM calls (before first step)
|
||||
first_recorded_step_time = (
|
||||
sorted_steps[0].get("timestamp") if sorted_steps else float("inf")
|
||||
)
|
||||
cursor.execute(
|
||||
"SELECT payload FROM traces WHERE session_id = ? AND (step_id IS NULL OR step_id = '') AND timestamp < ? AND type = 'llm_call'",
|
||||
(session_id, first_recorded_step_time),
|
||||
)
|
||||
plan_llm_rows = cursor.fetchall()
|
||||
plan_p, plan_c, plan_t = 0, 0, 0
|
||||
for lr in plan_llm_rows:
|
||||
p, c, t = self._extract_token_usage_from_trace(lr["payload"])
|
||||
plan_p += p
|
||||
plan_c += c
|
||||
plan_t += t
|
||||
|
||||
# Retrieve planner thinking / thoughts
|
||||
cursor.execute(
|
||||
"SELECT type, payload FROM traces WHERE session_id = ? AND (step_id IS NULL OR step_id = '') AND timestamp < ? AND type IN ('raw_thinking', 'thinking') ORDER BY timestamp ASC",
|
||||
(session_id, first_recorded_step_time),
|
||||
)
|
||||
think_rows = cursor.fetchall()
|
||||
p_raw_th = []
|
||||
p_nat_th = []
|
||||
for thr in think_rows:
|
||||
if thr["payload"]:
|
||||
try:
|
||||
th_obj = (
|
||||
json.loads(thr["payload"])
|
||||
if isinstance(thr["payload"], str)
|
||||
else thr["payload"]
|
||||
)
|
||||
th_text = th_obj.get("thought") or th_obj.get("text") or ""
|
||||
if thr["type"] == "raw_thinking" and th_text:
|
||||
p_raw_th.append(str(th_text).strip())
|
||||
elif thr["type"] == "thinking" and th_text:
|
||||
p_nat_th.append(str(th_text).strip())
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
virtual_step = {
|
||||
"step_id": "pre-planning",
|
||||
"step_type": "planning",
|
||||
"session_id": session_id,
|
||||
"step_number": 0,
|
||||
"timestamp": first_timestamp,
|
||||
"operator_native_thinking": "\n\n".join(p_nat_th) if p_nat_th else "",
|
||||
"operator_raw_thinking": "\n\n".join(p_raw_th) if p_raw_th else "",
|
||||
"action_taken": None,
|
||||
"last_execution_result": None,
|
||||
"extra_metadata": None,
|
||||
"generic_tools": unassigned_tools,
|
||||
"token_usage": {
|
||||
"prompt_tokens": plan_p,
|
||||
"completion_tokens": plan_c,
|
||||
"total_tokens": plan_t,
|
||||
},
|
||||
"total_tokens": plan_t,
|
||||
}
|
||||
steps.append(virtual_step)
|
||||
|
||||
steps.sort(
|
||||
key=lambda x: (
|
||||
x.get("timestamp") or 0,
|
||||
x.get("step_number") or 0,
|
||||
)
|
||||
)
|
||||
|
||||
return steps
|
||||
|
||||
|
||||
step_repo = StepRepository()
|
||||
@@ -0,0 +1,131 @@
|
||||
# Copyright 2026 Google LLC
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
import json
|
||||
from typing import Any
|
||||
|
||||
try:
|
||||
from admin_console.database.connection import db_session
|
||||
except ImportError:
|
||||
from apps.admin_console.database.connection import db_session
|
||||
|
||||
|
||||
class TraceRepository:
|
||||
"""Repository handling querying traces and constructing hierarchical trace trees."""
|
||||
|
||||
def __init__(self, db_path=None):
|
||||
self.db_path = db_path
|
||||
|
||||
def get_trace_by_id(self, trace_id: str, db_path=None) -> dict[str, Any] | None:
|
||||
with db_session(db_path or self.db_path) as conn:
|
||||
cursor = conn.cursor()
|
||||
cursor.execute("SELECT * FROM traces WHERE trace_id = ?", (trace_id,))
|
||||
row = cursor.fetchone()
|
||||
return dict(row) if row else None
|
||||
|
||||
def get_trace_tree(self, session_id: str) -> list[dict[str, Any]]:
|
||||
with db_session(self.db_path) as conn:
|
||||
cursor = conn.cursor()
|
||||
cursor.execute(
|
||||
"SELECT trace_id, parent_trace_id, type, name, status, timestamp,"
|
||||
" duration, CASE WHEN type IN ('raw_thinking', 'thinking') THEN"
|
||||
" payload ELSE NULL END AS payload FROM traces WHERE session_id = ?"
|
||||
" ORDER BY timestamp ASC",
|
||||
(session_id,),
|
||||
)
|
||||
rows = cursor.fetchall()
|
||||
|
||||
nodes = {}
|
||||
for r in rows:
|
||||
row_dict = dict(r)
|
||||
if row_dict.get("payload"):
|
||||
try:
|
||||
row_dict["payload"] = json.loads(row_dict["payload"])
|
||||
except Exception:
|
||||
pass
|
||||
nodes[row_dict["trace_id"]] = {**row_dict, "children": []}
|
||||
|
||||
root_nodes = []
|
||||
for node_id, node in nodes.items():
|
||||
parent_id = node["parent_trace_id"]
|
||||
if parent_id and parent_id in nodes:
|
||||
nodes[parent_id]["children"].append(node)
|
||||
else:
|
||||
root_nodes.append(node)
|
||||
|
||||
def sanitize_tree(nodes_list):
|
||||
sanitized = []
|
||||
for node in nodes_list:
|
||||
node["children"] = sanitize_tree(node["children"])
|
||||
if node["name"] == "node":
|
||||
sanitized.extend(node["children"])
|
||||
else:
|
||||
sanitized.append(node)
|
||||
return sanitized
|
||||
|
||||
return sanitize_tree(root_nodes)
|
||||
|
||||
def get_step_traces_tree(self, session_id: str, step_id: str) -> list[dict[str, Any]]:
|
||||
with db_session(self.db_path) as conn:
|
||||
cursor = conn.cursor()
|
||||
cursor.execute(
|
||||
"SELECT trace_id, parent_trace_id, step_id, type, name, status,"
|
||||
" timestamp, duration, CASE WHEN type IN ('raw_thinking',"
|
||||
" 'thinking') THEN payload ELSE NULL END AS payload FROM traces"
|
||||
" WHERE session_id = ? ORDER BY timestamp ASC",
|
||||
(session_id,),
|
||||
)
|
||||
rows = cursor.fetchall()
|
||||
|
||||
nodes = {}
|
||||
for r in rows:
|
||||
row_dict = dict(r)
|
||||
if row_dict.get("payload"):
|
||||
try:
|
||||
row_dict["payload"] = json.loads(row_dict["payload"])
|
||||
except Exception:
|
||||
pass
|
||||
nodes[row_dict["trace_id"]] = {**row_dict, "children": []}
|
||||
|
||||
step_roots = []
|
||||
for n_id, node in nodes.items():
|
||||
if node.get("step_id") == step_id:
|
||||
step_roots.append(node)
|
||||
|
||||
for n_id, node in nodes.items():
|
||||
parent_id = node["parent_trace_id"]
|
||||
if parent_id and parent_id in nodes:
|
||||
nodes[parent_id]["children"].append(node)
|
||||
|
||||
def sanitize_tree(nodes_list, visited=None):
|
||||
if visited is None:
|
||||
visited = set()
|
||||
sanitized = []
|
||||
for node in nodes_list:
|
||||
node_id = node.get("trace_id")
|
||||
if node_id and node_id in visited:
|
||||
continue
|
||||
if node_id:
|
||||
visited.add(node_id)
|
||||
node["children"] = sanitize_tree(node["children"], visited.copy())
|
||||
if node["name"] == "node":
|
||||
sanitized.extend(node["children"])
|
||||
else:
|
||||
sanitized.append(node)
|
||||
return sanitized
|
||||
|
||||
return sanitize_tree(step_roots)
|
||||
|
||||
|
||||
trace_repo = TraceRepository()
|
||||
Reference in new issue
Block a user