diff --git a/frontend/interactive_server.py b/frontend/interactive_server.py index ac03d50..346c12d 100644 --- a/frontend/interactive_server.py +++ b/frontend/interactive_server.py @@ -2,7 +2,9 @@ import sys import os import json import asyncio -from llm_panda.metrics_logger import LANGSMITH_AVAILABLE, log_system_metrics, log_interaction +import uuid +import psutil +from llm_panda.metrics_logger import LANGSMITH_AVAILABLE, log_system_metrics, log_interaction, log_episode_metrics import math import time @@ -33,7 +35,7 @@ from llm_panda import config from llm_panda.execution.cancellation import ExecutionInterrupted from llm_panda import logger as _logger_setup _logger_setup.setup() -from frontend.ctx.web_runtime import WebRuntime +from frontend.web_runtime import WebRuntime import pybullet_data from pybullet_object_models import ycb_objects from fastapi.middleware.cors import CORSMiddleware @@ -115,12 +117,13 @@ class SessionContext: def __init__(self, session_id: str): self.session_id = session_id self.sim = None - self.ctx.executor_lock = asyncio.Lock() - self.ctx.pybullet_lock = threading.RLock() - self.ctx.pipeline_executing = False - self.ctx.web_runtime = WebRuntime() - self.ctx.command_worker_task = None - self.ctx.physics_settle_task = None + self.executor_lock = asyncio.Lock() + self.pybullet_lock = threading.RLock() + self.pipeline_executing = False + self.web_runtime = WebRuntime() + self.command_worker_task = None + self.physics_settle_task = None + self.disconnected = False self.drag_state = {"constraint_id": None, "active": False} self.last_active = time.time() @@ -345,12 +348,17 @@ def get_simulation_state_payload(ctx: SessionContext) -> dict: # Read objects pos/ori objects_state = {} for oid in sim.scene.object_ids: - name = sim.scene.get_name(oid) - pos, ori = pb.getBasePositionAndOrientation(oid, physicsClientId=sim.client) - objects_state[name] = { - "pos": list(pos), - "ori": list(ori) - } + if oid < 0: + continue + try: + name = sim.scene.get_name(oid) + pos, ori = pb.getBasePositionAndOrientation(oid, physicsClientId=sim.client) + objects_state[name] = { + "pos": list(pos), + "ori": list(ori) + } + except Exception as e: + print(f"Failed to get pose for object {oid}: {e}", flush=True) return { "type": "state", @@ -477,8 +485,10 @@ def run_pipeline_threadsafe( def publish(payload: dict, journal: bool = True): payload = {**payload, "task_id": task_id, "generation": generation} + if getattr(ctx, "episode_id", None) is not None: + payload["run_id"] = ctx.episode_id if getattr(ctx, "episode_run_tree", None) is not None: - payload["run_id"] = str(ctx.episode_run_tree.id) + payload["langsmith_run_id"] = str(ctx.episode_run_tree.id) ctx.web_runtime.publish_threadsafe(payload, loop, journal=journal) def log(message: str, step: Optional[int] = None): @@ -530,7 +540,7 @@ def run_pipeline_threadsafe( curr_time = time.time() if curr_time - last_stream_time[0] >= throttle_interval: last_stream_time[0] = curr_time - state_payload = get_simulation_state_payload(current_sim) + state_payload = get_simulation_state_payload(ctx) if state_payload: publish(state_payload, journal=False) @@ -560,7 +570,9 @@ def run_pipeline_threadsafe( state_text = json.dumps(scene_state, indent=2) if attempt == 0: + ctx.episode_id = str(uuid.uuid4()) ctx.episode_run_tree = None + ctx.llm_latency_total = 0 if LANGSMITH_AVAILABLE: from langsmith.run_trees import RunTree ctx.episode_run_tree = RunTree( @@ -608,6 +620,7 @@ def run_pipeline_threadsafe( # LangSmith automatically handles children if wrapped, but we can also store explicitly pass # Handled by LLMAgent wrapping if we did that, or we'll just end the tree at the end. llm_time = time.time() - llm_start_time + ctx.llm_latency_total += llm_time cancellation_token.checkpoint() # Discard an LLM result that arrived after the user pressed stop. @@ -683,7 +696,7 @@ def run_pipeline_threadsafe( "score_value": 0, "success": False, }) - state_payload = get_simulation_state_payload(current_sim) + state_payload = get_simulation_state_payload(ctx) if state_payload: publish(state_payload, journal=False) ctx.pipeline_executing = False @@ -692,13 +705,13 @@ def run_pipeline_threadsafe( if not plan: log("ÔÜá´©Å [1] Kß║┐ hoß║ích rß╗ùng.") log("ÔÅ│ [5/5] ─Éang t├¡nh to├ín ─æiß╗âm sß║»p xß║┐p...") - score_val, score_md = compute_and_format_score(current_sim, [], True) + score_val, score_md = compute_and_format_score(ctx, [], True) publish({ "type": "result", "score_display": score_md, "score_value": score_val, }) - state_payload = get_simulation_state_payload(current_sim) + state_payload = get_simulation_state_payload(ctx) if state_payload: publish(state_payload, journal=False) ctx.pipeline_executing = False @@ -740,7 +753,7 @@ def run_pipeline_threadsafe( try: # Trigger an initial simulation state transmission - state_payload = get_simulation_state_payload(current_sim) + state_payload = get_simulation_state_payload(ctx) if state_payload: publish(state_payload, journal=False) @@ -821,7 +834,7 @@ def run_pipeline_threadsafe( pb.stepSimulation = original_step # Send a final state to ensure visual sync - state_payload = get_simulation_state_payload(current_sim) + state_payload = get_simulation_state_payload(ctx) if state_payload: publish(state_payload, journal=False) @@ -886,7 +899,7 @@ def run_pipeline_threadsafe( # Step 5: Scoring cancellation_token.checkpoint() - score_val, score_md = compute_and_format_score(current_sim, plan, executed_successfully and goal_completed) + score_val, score_md = compute_and_format_score(ctx, plan, executed_successfully and goal_completed) error_cause = None if not plan: @@ -898,6 +911,43 @@ def run_pipeline_threadsafe( is_success = executed_successfully and goal_completed and score_val >= 50 + # Calculate PostHog metrics + token_usage = {"prompt_tokens": 0, "completion_tokens": 0} + if hasattr(ctx.sim.pipeline.agent, "planner") and hasattr(ctx.sim.pipeline.agent.planner, "last_token_usage"): + token_usage = getattr(ctx.sim.pipeline.agent.planner, "last_token_usage") + + system_ram_mb = psutil.virtual_memory().used / (1024 * 1024) + timeout_network_error = False + plan_validity_first_try = False + collision_free = True + kinematic_error = False + + if last_failed_sig: + if "Rate limit" in last_failed_sig or "Timeout" in last_failed_sig: + timeout_network_error = True + if "Kinematic" in last_failed_sig or "IK" in last_failed_sig or "outside workspace" in last_failed_sig.lower(): + kinematic_error = True + if "collision" in last_failed_sig.lower(): + collision_free = False + + if attempt == 0 and not last_failed_sig: + plan_validity_first_try = True + + try: + log_episode_metrics( + session_id=getattr(ctx, "episode_id", "unknown"), + llm_latency_sec=getattr(ctx, "llm_latency_total", 0), + token_usage=token_usage, + system_ram_mb=system_ram_mb, + timeout_network_error=timeout_network_error, + plan_validity_first_try=plan_validity_first_try, + collision_free=collision_free, + kinematic_error=kinematic_error, + replan_count=attempt + ) + except Exception as e: + log(f"Failed to send PostHog metrics: {e}") + publish({ "type": "result", "score_display": score_md, @@ -914,19 +964,19 @@ async def command_worker_loop(loop, ctx: SessionContext): try: while True: - command = ctx.web_ctx.web_runtime.take_pending() + command = ctx.web_runtime.take_pending() if command is None: return generation = command["generation"] task_id = command["task_id"] interrupted = False - await ctx.web_ctx.web_runtime.publish({ + await ctx.web_runtime.publish({ "type": "command_accepted", "prompt": command["prompt"], "task_id": task_id, "generation": generation, }) - await ctx.web_ctx.web_runtime.publish({ + await ctx.web_runtime.publish({ "type": "execution_status", "state": "planning", "task_id": task_id, @@ -938,7 +988,7 @@ async def command_worker_loop(loop, ctx: SessionContext): command["prompt"], loop, command["use_mock"], - ctx.web_runtime, + ctx, command["token"], generation, task_id, @@ -950,19 +1000,19 @@ async def command_worker_loop(loop, ctx: SessionContext): interrupted = True ctx.pipeline_executing = False with ctx.pybullet_lock: - if current_sim is not None and ctx.sim.pipeline is not None: + if ctx.sim is not None and ctx.sim.pipeline is not None: ctx.sim.pipeline.interrupt_hold() - ctx.web_ctx.web_runtime.set_status("interrupted", generation=generation) + ctx.web_runtime.set_status("interrupted", generation=generation) state_payload = await asyncio.to_thread( get_simulation_state_payload, - current_sim, + ctx.sim, ) if state_payload: - await ctx.web_ctx.web_runtime.publish(state_payload, journal=False) - ctx.web_ctx.web_runtime.conversation_manager.record_interruption( + await ctx.web_runtime.publish(state_payload, journal=False) + ctx.web_runtime.conversation_manager.record_interruption( command["prompt"] ) - await ctx.web_ctx.web_runtime.publish({ + await ctx.web_runtime.publish({ "type": "interrupted", "reason": str(exc), "task_id": task_id, @@ -970,7 +1020,7 @@ async def command_worker_loop(loop, ctx: SessionContext): }) except Exception as exc: ctx.pipeline_executing = False - await ctx.web_ctx.web_runtime.publish({ + await ctx.web_runtime.publish({ "type": "result", "success": False, "error_cause": "runtime", @@ -979,24 +1029,24 @@ async def command_worker_loop(loop, ctx: SessionContext): "generation": generation, }) finally: - if current_sim is not None and ctx.sim.pipeline is not None: + if ctx.sim is not None and ctx.sim.pipeline is not None: ctx.sim.pipeline.set_cancellation_token(None) - ctx.web_ctx.web_runtime.complete( + ctx.web_runtime.complete( generation, state="interrupted" if interrupted else "completed", ) finally: ctx.command_worker_task = None - if ctx.web_ctx.web_runtime.has_pending: + if ctx.web_runtime.has_pending: ctx.command_worker_task = asyncio.create_task(command_worker_loop(loop)) async def wait_for_active_task_stop(ctx: SessionContext, timeout=5.0): - ctx.web_ctx.web_runtime.cancel_all("runtime_mutation") + ctx.web_runtime.cancel_all("runtime_mutation") deadline = time.monotonic() + timeout - while ctx.web_ctx.web_runtime.active_token is not None and time.monotonic() < deadline: + while ctx.web_runtime.active_token is not None and time.monotonic() < deadline: await asyncio.sleep(0.02) - return ctx.web_ctx.web_runtime.active_token is None + return ctx.web_runtime.active_token is None @@ -1093,6 +1143,8 @@ async def websocket_endpoint(websocket: WebSocket): async with ctx.executor_lock: def init_sim(): with ctx.pybullet_lock: + if ctx.disconnected: + return if ctx.sim is None: print(f"[{session_id}] Initializing PandaApp lazily...", flush=True) env_path = os.path.join(PROJECT_SRC, "envs", "env1.json") @@ -1104,6 +1156,8 @@ async def websocket_endpoint(websocket: WebSocket): agent=None, viz=False ) + if ctx.disconnected: + pb.disconnect(physicsClientId=ctx.sim.client) except Exception as e: print(f"Failed to initialize lazy PandaApp: {e}", flush=True) raise e @@ -1603,6 +1657,7 @@ async def websocket_endpoint(websocket: WebSocket): print(f"[{session_id}] WS Exception: {e}") finally: if 'ctx' in locals(): + ctx.disconnected = True ctx.drag_state["active"] = False if ctx.drag_state.get("constraint_id") is not None: try: @@ -1632,4 +1687,5 @@ async def websocket_endpoint(websocket: WebSocket): if __name__ == "__main__": port = int(os.environ.get("PORT", 8000)) + import uvicorn uvicorn.run("interactive_server:app", host="0.0.0.0", port=port, reload=False)