Spaces:
Paused
Paused
| """Shared fixtures: an offline vector store, a stand-in G1 model, and an app wired to a scripted LLM.""" | |
| from __future__ import annotations | |
| import hashlib | |
| import re | |
| from dataclasses import replace | |
| from pathlib import Path | |
| import numpy as np | |
| import pytest | |
| from chromadb.api.types import Documents, EmbeddingFunction, Embeddings | |
| from fastapi.testclient import TestClient | |
| from api import keys | |
| from api.auth import Allowlist | |
| from api.config import get_settings | |
| from api.llm import ScriptedLLM | |
| from api.logstore import LogStore | |
| from api.main import Services, create_app | |
| from api.runpod_client import RunPodError | |
| from api.storage import LocalStorage | |
| from rag.chunking import Chunk | |
| from rag.store import RagStore | |
| FAKE_G1 = """<mujoco model="fake g1"> | |
| <compiler angle="radian"/> | |
| <worldbody> | |
| <body name="pelvis" pos="0 0 0.8"> | |
| <freejoint name="floating_base_joint"/> | |
| <geom name="torso" type="capsule" size="0.1 0.3" mass="10"/> | |
| <body name="left_knee_link" pos="0 0.1 -0.4"> | |
| <joint name="left_knee_joint" axis="0 1 0" range="0 2.8"/> | |
| <geom type="capsule" size="0.05 0.2"/> | |
| </body> | |
| </body> | |
| </worldbody> | |
| <actuator><motor name="left_knee_joint" joint="left_knee_joint"/></actuator> | |
| <keyframe><key name="stand" qpos="0 0 0.8 1 0 0 0 0"/></keyframe> | |
| </mujoco> | |
| """ | |
| GOOD_SCENE = """<mujoco model="g1 on slope"> | |
| <include file="g1.xml"/> | |
| <worldbody> | |
| <light name="sun" pos="0 0 4"/> | |
| <geom name="floor" type="plane" size="0 0 0.05" friction="1 0.005 0.0001"/> | |
| <geom name="slope" type="box" size="1 1 0.05" pos="2 0 0.2" euler="0 -0.26 0"/> | |
| <body name="crate" pos="1 1 0.3"><freejoint/><geom name="crate_geom" type="box" size="0.2 0.2 0.2" mass="3"/></body> | |
| </worldbody> | |
| </mujoco>""" | |
| class HashEmbedding(EmbeddingFunction[Documents]): | |
| """Bag-of-words hashing embedding: deterministic, offline, good enough to rank by word overlap.""" | |
| def __init__(self) -> None: | |
| pass | |
| def __call__(self, input: Documents) -> Embeddings: | |
| vectors = [] | |
| for text in input: | |
| v = np.zeros(256, dtype=np.float32) | |
| for word in re.findall(r"[a-z0-9_]+", text.lower()): | |
| v[int(hashlib.md5(word.encode()).hexdigest(), 16) % 256] += 1.0 | |
| n = np.linalg.norm(v) | |
| vectors.append(v / n if n else v) | |
| return vectors | |
| def name() -> str: | |
| return "test-hash" | |
| def get_config(self) -> dict: | |
| return {} | |
| def build_from_config(config: dict) -> HashEmbedding: | |
| return HashEmbedding() | |
| def store(tmp_path) -> RagStore: | |
| s = RagStore(tmp_path / "chroma", embedding_function=HashEmbedding()) | |
| s.add( | |
| "mujoco", | |
| [ | |
| Chunk( | |
| "[XML reference > geom] geom type plane box sphere friction size pos euler", "doc#geom", {"kind": "doc"} | |
| ), | |
| Chunk("[XML reference > light] light pos dir diffuse", "doc#light", {"kind": "doc"}), | |
| Chunk( | |
| '[MJCF example scene.xml] <include file="g1.xml"/> floor plane', | |
| "menagerie:scene.xml", | |
| {"kind": "example", "robot": "unitree_g1"}, | |
| ), | |
| ], | |
| ) | |
| s.add( | |
| "ros2", | |
| [ | |
| Chunk( | |
| "[Writing a simple py publisher] create_publisher String timer_callback rclpy", | |
| "docs:pubsub", | |
| {"kind": "doc", "distro": "humble"}, | |
| ), | |
| Chunk( | |
| "[Launch file formats] LaunchDescription Node remappings parameters", | |
| "docs:launch", | |
| {"kind": "doc", "distro": "humble"}, | |
| ), | |
| Chunk( | |
| "[ROS 2 humble Message sensor_msgs/msg/LaserScan]\nfloat32[] ranges", | |
| "ci:LaserScan.msg", | |
| {"kind": "interface", "distro": "humble", "interface": "sensor_msgs/msg/LaserScan"}, | |
| ), | |
| ], | |
| ) | |
| return s | |
| def menagerie(tmp_path) -> Path: | |
| root = tmp_path / "menagerie" | |
| (root / "unitree_g1").mkdir(parents=True) | |
| (root / "unitree_g1" / "g1.xml").write_text(FAKE_G1) | |
| return root | |
| def settings(tmp_path, menagerie): | |
| base = get_settings() | |
| return replace( | |
| base, | |
| llm_api_key="test", | |
| data_dir=tmp_path / "data", | |
| testers_file=tmp_path / "testers.yaml", | |
| log_db=tmp_path / "data" / "log.sqlite3", | |
| menagerie_dir=menagerie, | |
| mjcf_sim_seconds=0.2, | |
| codegen_backend="direct", | |
| auto_review=False, | |
| prompt_enhance=False, | |
| free_debug=False, | |
| ) | |
| # --------------------------------------------------------------------------- job backend | |
| class FakeBackend: | |
| name = "runpod" | |
| def __init__(self): | |
| self.jobs: dict[str, dict] = {} | |
| self.submitted: list[dict] = [] | |
| self.fail_submit = False | |
| self.functions: list[str] = [] | |
| def supports(self, function): | |
| return True | |
| def submit(self, job_input, execution_timeout_s, function="run_simulation"): | |
| self.functions.append(function) | |
| if self.fail_submit: | |
| raise RunPodError("RunPod /run returned HTTP 500: boom") | |
| job_id = f"job-{len(self.submitted) + 1}" | |
| self.submitted.append(job_input) | |
| self.jobs[job_id] = {"id": job_id, "status": "IN_QUEUE"} | |
| return dict(self.jobs[job_id]) | |
| def status(self, job_id): | |
| return dict(self.jobs[job_id]) | |
| def run_env(settings, store, tmp_path): | |
| alice = keys.add(settings.testers_file, "alice") | |
| bob = keys.add(settings.testers_file, "bob") | |
| llm = ScriptedLLM([]) | |
| backend = FakeBackend() | |
| storage = LocalStorage(tmp_path / "r2", base_url="https://r2.example/bucket") | |
| services = Services( | |
| settings=settings, | |
| store=store, | |
| logs=LogStore(settings.log_db), | |
| allowlist=Allowlist(settings.testers_file), | |
| llm_factory=lambda: llm, | |
| backend=backend, | |
| storage=storage, | |
| ) | |
| client = TestClient(create_app(services)) | |
| llm.responses[:] = [f"```xml\n{GOOD_SCENE}\n```"] | |
| gen = client.post( | |
| "/simulate/generate", json={"description": "G1 near a slope and a crate"}, headers={"X-API-Key": alice} | |
| ).json() | |
| return client, {"alice": alice, "bob": bob}, backend, services, gen["id"] | |
| def complete(backend, job_id, run_id, **summary): | |
| backend.jobs[job_id] = { | |
| "id": job_id, | |
| "status": "COMPLETED", | |
| "executionTime": 12500, | |
| "delayTime": 3000, | |
| "output": { | |
| "run_id": run_id, | |
| "summary": {"simulated_seconds": 10.0, "fell": False, **summary}, | |
| "files": { | |
| "mcap": f"runs/{run_id}/run.mcap", | |
| "video": f"runs/{run_id}/video.mp4", | |
| "video_webm": f"runs/{run_id}/video.webm", | |
| "summary": f"runs/{run_id}/summary.json", | |
| }, | |
| }, | |
| } | |