tau-sim 0.1.0__py3-none-any.whl

tau_sim/api/routes.py

@@ -0,0 +1,148 @@
+from fastapi import APIRouter, HTTPException, Request, Response
+
+from tau_sim.llm import manual
+from tau_sim.llm.agent import parse_manual_response, run_chat
+from tau_sim.middleware.session import session_id_from_request
+from tau_sim.schemas import (
+    ChatRequest,
+    ChatResponse,
+    FileWrite,
+    ProjectCreate,
+    ProjectInfo,
+)
+from tau_sim.sim import menagerie
+from tau_sim.storage import projects
+
+router = APIRouter(prefix="/api")
+
+
+def _sid(request: Request) -> str:
+    return session_id_from_request(request)
+
+
+# ---- projects ----------------------------------------------------------
+
+@router.get("/projects", response_model=list[ProjectInfo])
+def list_projects(request: Request) -> list[ProjectInfo]:
+    return projects.list_projects(_sid(request))
+
+
+@router.post("/projects", response_model=ProjectInfo)
+def create_project(body: ProjectCreate, request: Request) -> ProjectInfo:
+    return projects.create_project(_sid(request), body.name, body.template)
+
+
+@router.get("/projects/{project_id}", response_model=ProjectInfo)
+def get_project(project_id: str, request: Request) -> ProjectInfo:
+    try:
+        return projects.get_project(_sid(request), project_id)
+    except FileNotFoundError as e:
+        raise HTTPException(status_code=404, detail=str(e))
+
+
+@router.delete("/projects/{project_id}")
+def delete_project(project_id: str, request: Request) -> dict:
+    projects.delete_project(_sid(request), project_id)
+    return {"ok": True}
+
+
+# ---- menagerie import --------------------------------------------------
+
+@router.get("/menagerie/models")
+def menagerie_models() -> list[str]:
+    try:
+        return menagerie.list_models()
+    except Exception as e:  # noqa: BLE001
+        raise HTTPException(status_code=502, detail=f"menagerie list failed: {e}")
+
+
+@router.post("/projects/{project_id}/import/menagerie")
+def import_menagerie(project_id: str, body: dict, request: Request) -> dict:
+    model = (body or {}).get("model", "").strip()
+    if not model:
+        raise HTTPException(status_code=400, detail="`model` is required")
+    try:
+        root = projects.project_root(_sid(request), project_id)
+    except (FileNotFoundError, ValueError) as e:
+        raise HTTPException(status_code=404, detail=str(e))
+    try:
+        return menagerie.import_model(model, root)
+    except (FileNotFoundError, ValueError) as e:
+        raise HTTPException(status_code=400, detail=str(e))
+    except PermissionError as e:
+        raise HTTPException(status_code=429, detail=str(e))
+    except Exception as e:  # noqa: BLE001
+        raise HTTPException(status_code=502, detail=f"import failed: {e}")
+
+
+# ---- files -------------------------------------------------------------
+
+@router.get("/projects/{project_id}/file")
+def read_file(project_id: str, path: str, request: Request) -> dict:
+    try:
+        return {
+            "path": path,
+            "content": projects.read_file(_sid(request), project_id, path),
+        }
+    except FileNotFoundError:
+        raise HTTPException(status_code=404, detail="file not found")
+    except ValueError as e:
+        raise HTTPException(status_code=400, detail=str(e))
+
+
+@router.put("/projects/{project_id}/file")
+def write_file(project_id: str, body: FileWrite, request: Request) -> dict:
+    try:
+        projects.write_file(_sid(request), project_id, body.path, body.content)
+        return {"ok": True}
+    except ValueError as e:
+        raise HTTPException(status_code=400, detail=str(e))
+
+
+# ---- chat --------------------------------------------------------------
+
+@router.post("/chat", response_model=ChatResponse)
+def chat(body: ChatRequest, request: Request) -> ChatResponse:
+    try:
+        return run_chat(_sid(request), body)
+    except RuntimeError as e:
+        raise HTTPException(status_code=400, detail=str(e))
+    except Exception as e:  # noqa: BLE001
+        raise HTTPException(status_code=500, detail=f"LLM error: {e!r}")
+
+
+# ---- manual / human-in-the-loop ----------------------------------------
+
+@router.get("/chat/result/{pending_id}")
+def chat_result(pending_id: str, response: Response) -> dict:
+    try:
+        payload = manual.read_response(pending_id)
+    except ValueError as e:
+        raise HTTPException(status_code=400, detail=str(e))
+    if payload is None:
+        response.status_code = 202
+        return {"pending": True}
+    chat_resp = parse_manual_response(payload)
+    return chat_resp.model_dump()
+
+
+@router.get("/manual/pending")
+def manual_pending() -> list[dict]:
+    return manual.list_pending()
+
+
+@router.get("/manual/pending/{pending_id}")
+def manual_pending_one(pending_id: str) -> dict:
+    req = manual.get_request(pending_id)
+    if req is None:
+        raise HTTPException(status_code=404, detail="not found")
+    return req
+
+
+@router.post("/manual/pending/{pending_id}/respond")
+def manual_respond(pending_id: str, payload: dict) -> dict:
+    try:
+        manual.write_response(pending_id, payload)
+    except FileNotFoundError:
+        raise HTTPException(status_code=404, detail="not found")
+    return {"ok": True}

tau_sim/cli.py

@@ -0,0 +1,62 @@
+"""CLI entrypoint: ``tau`` launches the backend + opens a browser.
+
+Bundles the SPA if `backend/static/` exists (production build); otherwise
+expects Vite dev server on port 5173.
+"""
+from __future__ import annotations
+
+import argparse
+import os
+import sys
+import threading
+import time
+import webbrowser
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(prog="tau")
+    parser.add_argument("--host", default="127.0.0.1")
+    parser.add_argument("--port", default=int(os.environ.get("PORT", 7860)), type=int)
+    parser.add_argument("--no-browser", action="store_true",
+                        help="don't auto-open the browser")
+    parser.add_argument("--mujoco-gl",
+                        default=os.environ.get("MUJOCO_GL", "egl"),
+                        choices=["egl", "osmesa", "glfw"],
+                        help="MuJoCo rendering backend")
+    parser.add_argument("--projects-dir",
+                        default=os.environ.get("PROJECTS_DIR"),
+                        help="where to store project data "
+                             "(default: ~/.tau/projects)")
+    args = parser.parse_args()
+
+    from pathlib import Path
+
+    projects_dir = (
+        Path(args.projects_dir).expanduser()
+        if args.projects_dir
+        else Path.home() / ".tau" / "projects"
+    )
+    projects_dir.mkdir(parents=True, exist_ok=True)
+    os.environ["PROJECTS_DIR"] = str(projects_dir)
+    os.environ["MUJOCO_GL"] = args.mujoco_gl
+    # Local CLI = single user, no sandbox by default.
+    os.environ.setdefault("TAU_SINGLE_USER", "1")
+
+    print(f"\n  tau-sim  -  http://{args.host}:{args.port}", file=sys.stderr)
+    print(f"  data:     {projects_dir}", file=sys.stderr)
+    print(f"  render:   MUJOCO_GL={args.mujoco_gl}\n", file=sys.stderr)
+
+    if not args.no_browser:
+        url = f"http://{args.host}:{args.port}/"
+        threading.Thread(
+            target=lambda: (time.sleep(0.8), webbrowser.open(url)),
+            daemon=True,
+        ).start()
+
+    # Import uvicorn lazily so --help is fast.
+    import uvicorn
+    uvicorn.run("tau_sim.main:app", host=args.host, port=args.port, log_level="warning")
+
+
+if __name__ == "__main__":
+    main()

tau_sim/config.py

@@ -0,0 +1,23 @@
+from pathlib import Path
+from pydantic_settings import BaseSettings, SettingsConfigDict
+
+
+class Settings(BaseSettings):
+    model_config = SettingsConfigDict(env_file=".env", extra="ignore")
+
+    openai_api_key: str = ""
+    gemini_api_key: str = ""
+    openrouter_api_key: str = ""
+
+    llm_default_provider: str = "openai"
+    llm_default_model_openai: str = "gpt-4o-mini"
+    llm_default_model_gemini: str = "gemini-1.5-flash"
+    llm_default_model_openrouter: str = "deepseek/deepseek-v4-flash"
+    openrouter_base_url: str = "https://openrouter.ai/api/v1"
+
+    projects_dir: Path = Path("./projects_data")
+    mujoco_gl: str = "egl"
+
+
+settings = Settings()
+settings.projects_dir.mkdir(parents=True, exist_ok=True)

tau_sim/llm/agent.py

@@ -0,0 +1,195 @@
+"""Chat agent that asks the LLM to return a JSON object with a reply
+plus an optional list of full-file replacement proposals."""
+from __future__ import annotations
+
+import json
+from typing import Any
+
+from tau_sim.llm import manual
+from tau_sim.llm.providers import get_provider
+from tau_sim.schemas import ChatMessage, ChatRequest, ChatResponse, DiffProposal
+from tau_sim.storage import projects
+
+_SYSTEM = """You are an expert coding assistant embedded in an IDE for
+building MuJoCo robotic-simulation environments. You collaborate with
+the user on ONE project at a time. The project contains MJCF XML scenes
+and a Python entry module that exposes `make_env()`.
+
+# Response format
+
+Respond with a single JSON object of EXACTLY this shape and nothing else
+(no prose, no markdown fences):
+
+{
+  "reply": "<short natural-language explanation for the user>",
+  "proposals": [
+    {
+      "path": "<project-relative file path>",
+      "new_content": "<the COMPLETE new contents of that file>",
+      "rationale": "<one-sentence why>"
+    }
+  ]
+}
+
+Rules for `proposals`:
+- `new_content` is the FULL file (never a diff/patch).
+- Include only files you actually want to change. May be empty.
+- Paths are relative to the project root. Use forward slashes.
+- Creating a new file inside the current project is allowed when the
+  user's request implies it (e.g. a helper module, a new asset).
+- NEVER include text outside the JSON object.
+
+# Project scope (very important)
+
+Each project represents ONE simulated system (one robot, one task).
+When the user asks for something that does not fit the current project:
+
+- If the request is a *modification* of the existing system (change
+  geometry, tune a reward, add a sensor, fix a bug) -> emit proposals.
+- If the request is a *different system* (e.g. project is a cartpole and
+  the user asks for a quadruped, a manipulator, a drone, etc.) -> do
+  NOT replace the existing files. Instead return an empty `proposals`
+  array and use `reply` to tell the user to click the “+ New” button in
+  the top bar to create a new project for that system, and offer to
+  generate the env once they switch to it. Mention the suggested
+  project name.
+- Never silently replace the entry env, the scene XML, or any file with
+  code for an unrelated robot or task.
+
+When in doubt, prefer to ASK in `reply` rather than guess.
+
+# Importing real robot models
+
+You CANNOT fetch URLs or external files yourself. If the user asks to
+import a model from `mujoco_menagerie` (or any other repo of pre-built
+MuJoCo assets), do NOT try to inline the XML/meshes. Instead, in
+`reply`, tell them to:
+
+1. Click **+ New** in the top bar.
+2. In the modal, choose template **"Import from MuJoCo Menagerie"**.
+3. Pick the model from the dropdown (e.g. `unitree_h1`, `franka_emika_panda`).
+
+The importer downloads the XML + meshes and writes a stub `env.py`. You
+can then help them edit the env's reward / policy in that new project.
+
+# Env API contract (the runner calls these exact methods)
+
+The entry module must define `make_env()` returning an object with:
+
+    class MyEnv:
+        xml_path = "scene.xml"   # any project-relative .xml file
+
+        def reset(self, data):              # may also accept (self, model, data)
+            data.qpos[:] = ...
+            data.qvel[:] = 0.0
+            return self._obs(data)          # numpy array of observations
+
+        def step(self, model, data, action):
+            data.ctrl[:] = action           # set actuator commands
+            obs = self._obs(data)
+            reward = float(...)
+            done = bool(...)
+            return obs, reward, done
+
+        def policy(self, obs):              # called every control step
+            return np.array([...])          # shape must match model.nu
+
+The runner advances physics (`mujoco.mj_step`) between `step()` calls;
+do NOT call `mj_step` yourself. Render setup is also handled for you.
+
+# Style
+
+- Keep MJCF valid (closing tags, valid attribute values).
+- Keep Python importable. `import numpy as np` if you use it.
+- Prefer small, focused diffs. Don't reformat unrelated code.
+- Don't add docstrings or comments unless the user asked for them.
+"""
+
+
+def _build_context(session_id: str, project_id: str) -> str:
+    info = projects.get_project(session_id, project_id)
+    files = projects.read_all_text(session_id, project_id)
+    entry = _entry_file(files)
+    header_lines = [
+        f"### Project: {info.name!r} (id: {info.id})",
+        f"Entry env module: {entry or '(none yet)'}",
+        "The runner loads this entry module, calls `make_env()`, then",
+        "drives `env.reset` / `env.step` / `env.policy` per the contract.",
+        "All file edits must stay within the scope of this single project.",
+        "",
+        f"### Current project files ({len(files)} files):",
+        "",
+    ]
+    parts = ["\n".join(header_lines)]
+    for path, content in files.items():
+        parts.append(f"--- {path} ---\n{content}\n")
+    return "\n".join(parts)
+
+
+def _entry_file(files: dict[str, str]) -> str | None:
+    """Mirror the runner's entry-file rule for the LLM's benefit."""
+    if "env.py" in files:
+        return "env.py"
+    candidates = sorted(p for p in files if p.endswith("_env.py"))
+    return candidates[0] if candidates else None
+
+
+def _parse(raw: str) -> dict[str, Any]:
+    raw = raw.strip()
+    # Strip stray markdown fences if a model ignores response_format.
+    if raw.startswith("```"):
+        raw = raw.strip("`")
+        if raw.lower().startswith("json"):
+            raw = raw[4:]
+        raw = raw.strip()
+    try:
+        return json.loads(raw)
+    except json.JSONDecodeError:
+        return {"reply": raw, "proposals": []}
+
+
+def run_chat(session_id: str, req: ChatRequest) -> ChatResponse:
+    context = _build_context(session_id, req.project_id)
+    # Inject project context as a leading user message.
+    messages = [ChatMessage(role="user", content=context), *req.messages]
+
+    if (req.provider or "").lower() == "manual":
+        pid = manual.enqueue(
+            req.project_id,
+            _SYSTEM,
+            [m.model_dump() for m in messages],
+        )
+        return ChatResponse(
+            reply="(waiting for human-in-the-loop response…)",
+            pending_id=pid,
+        )
+
+    provider = get_provider(req.provider, req.api_key)
+    raw = provider.complete(_SYSTEM, messages, req.model)
+    parsed = _parse(raw)
+
+    proposals: list[DiffProposal] = []
+    for p in parsed.get("proposals", []) or []:
+        try:
+            proposals.append(DiffProposal(**p))
+        except Exception:
+            continue
+
+    return ChatResponse(
+        reply=str(parsed.get("reply", "")).strip() or "(no reply)",
+        proposals=proposals,
+    )
+
+
+def parse_manual_response(payload: dict[str, Any]) -> ChatResponse:
+    """Validate a manually-written response and return a ChatResponse."""
+    proposals: list[DiffProposal] = []
+    for p in payload.get("proposals", []) or []:
+        try:
+            proposals.append(DiffProposal(**p))
+        except Exception:
+            continue
+    return ChatResponse(
+        reply=str(payload.get("reply", "")).strip() or "(no reply)",
+        proposals=proposals,
+    )

tau_sim/llm/manual.py

@@ -0,0 +1,91 @@
+"""Human-in-the-loop \"LLM\" queue.
+
+When the user picks the `manual` provider, the chat request is written
+to ``manual_queue/<id>/request.json`` and the API returns a `pending_id`.
+A human (or Copilot) drops a JSON file at ``manual_queue/<id>/response.json``
+and the frontend, which is polling, picks it up.
+
+Response file shape (same as the regular LLM JSON contract):
+
+    {"reply": "...", "proposals": [{"path": "...", "new_content": "...",
+                                    "rationale": "..."}]}
+"""
+from __future__ import annotations
+
+import json
+import time
+import uuid
+from pathlib import Path
+from typing import Any
+
+from tau_sim.config import settings
+
+QUEUE_DIR: Path = settings.projects_dir.parent.resolve() / "manual_queue"
+QUEUE_DIR.mkdir(parents=True, exist_ok=True)
+
+
+def _slot(pending_id: str) -> Path:
+    p = (QUEUE_DIR / pending_id).resolve()
+    if not p.is_relative_to(QUEUE_DIR):
+        raise ValueError("invalid id")
+    return p
+
+
+def enqueue(project_id: str, system: str, messages: list[dict[str, str]]) -> str:
+    """Write a pending request and return its id."""
+    pid = uuid.uuid4().hex[:10]
+    slot = _slot(pid)
+    slot.mkdir(parents=True, exist_ok=False)
+    (slot / "request.json").write_text(
+        json.dumps(
+            {
+                "id": pid,
+                "project_id": project_id,
+                "created_at": time.time(),
+                "system": system,
+                "messages": messages,
+            },
+            indent=2,
+        )
+    )
+    return pid
+
+
+def list_pending() -> list[dict[str, Any]]:
+    out: list[dict[str, Any]] = []
+    for d in sorted(QUEUE_DIR.iterdir()):
+        if not d.is_dir():
+            continue
+        req = d / "request.json"
+        resp = d / "response.json"
+        if req.exists() and not resp.exists():
+            try:
+                out.append(json.loads(req.read_text()))
+            except json.JSONDecodeError:
+                continue
+    return out
+
+
+def get_request(pending_id: str) -> dict[str, Any] | None:
+    req = _slot(pending_id) / "request.json"
+    return json.loads(req.read_text()) if req.exists() else None
+
+
+def write_response(pending_id: str, payload: dict[str, Any]) -> None:
+    slot = _slot(pending_id)
+    if not slot.is_dir():
+        raise FileNotFoundError(pending_id)
+    (slot / "response.json").write_text(json.dumps(payload, indent=2))
+
+
+def read_response(pending_id: str) -> dict[str, Any] | None:
+    resp = _slot(pending_id) / "response.json"
+    return json.loads(resp.read_text()) if resp.exists() else None
+
+
+def delete(pending_id: str) -> None:
+    import shutil
+
+    slot = _slot(pending_id)
+    if slot.is_dir():
+        shutil.rmtree(slot)

tau_sim/llm/providers.py

@@ -0,0 +1,115 @@
+"""Pluggable LLM providers.
+
+Each provider exposes ``complete(system, messages, model) -> str``.
+The agent layer wraps these and parses structured diff proposals.
+"""
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+
+from tau_sim.config import settings
+from tau_sim.schemas import ChatMessage
+
+
+class LLMProvider(ABC):
+    name: str
+
+    @abstractmethod
+    def complete(self, system: str, messages: list[ChatMessage], model: str | None) -> str: ...
+
+
+class OpenAIProvider(LLMProvider):
+    name = "openai"
+
+    def __init__(self, api_key: str | None = None) -> None:
+        from openai import OpenAI
+
+        key = api_key or settings.openai_api_key
+        if not key:
+            raise RuntimeError("OPENAI_API_KEY is not set")
+        self._client = OpenAI(api_key=key)
+
+    def complete(self, system: str, messages: list[ChatMessage], model: str | None) -> str:
+        chat = [{"role": "system", "content": system}] + [
+            {"role": m.role, "content": m.content} for m in messages
+        ]
+        resp = self._client.chat.completions.create(
+            model=model or settings.llm_default_model_openai,
+            messages=chat,
+            response_format={"type": "json_object"},
+            temperature=0.2,
+        )
+        return resp.choices[0].message.content or ""
+
+
+class GeminiProvider(LLMProvider):
+    name = "gemini"
+
+    def __init__(self, api_key: str | None = None) -> None:
+        import google.generativeai as genai
+
+        key = api_key or settings.gemini_api_key
+        if not key:
+            raise RuntimeError("GEMINI_API_KEY is not set")
+        genai.configure(api_key=key)
+        self._genai = genai
+
+    def complete(self, system: str, messages: list[ChatMessage], model: str | None) -> str:
+        gm = self._genai.GenerativeModel(
+            model_name=model or settings.llm_default_model_gemini,
+            system_instruction=system,
+            generation_config={"response_mime_type": "application/json", "temperature": 0.2},
+        )
+        # Gemini expects alternating user/model turns.
+        history = []
+        for m in messages[:-1]:
+            history.append({"role": "user" if m.role == "user" else "model", "parts": [m.content]})
+        chat = gm.start_chat(history=history)
+        last = messages[-1].content if messages else ""
+        resp = chat.send_message(last)
+        return resp.text or ""
+
+
+class OpenRouterProvider(LLMProvider):
+    """OpenAI-compatible client pointed at OpenRouter."""
+
+    name = "openrouter"
+
+    def __init__(self, api_key: str | None = None) -> None:
+        from openai import OpenAI
+
+        key = api_key or settings.openrouter_api_key
+        if not key:
+            raise RuntimeError("OPENROUTER_API_KEY is not set")
+        self._client = OpenAI(
+            api_key=key,
+            base_url=settings.openrouter_base_url,
+            default_headers={
+                # Optional but recommended by OpenRouter for attribution.
+                "HTTP-Referer": "https://github.com/sim_llm",
+                "X-Title": "sim_llm",
+            },
+        )
+
+    def complete(self, system: str, messages: list[ChatMessage], model: str | None) -> str:
+        chat = [{"role": "system", "content": system}] + [
+            {"role": m.role, "content": m.content} for m in messages
+        ]
+        resp = self._client.chat.completions.create(
+            model=model or settings.llm_default_model_openrouter,
+            messages=chat,
+            response_format={"type": "json_object"},
+            temperature=0.2,
+        )
+        return resp.choices[0].message.content or ""
+
+
+def get_provider(name: str | None, api_key: str | None = None) -> LLMProvider:
+    name = (name or settings.llm_default_provider).lower()
+    if name == "openai":
+        return OpenAIProvider(api_key)
+    if name == "gemini":
+        return GeminiProvider(api_key)
+    if name == "openrouter":
+        return OpenRouterProvider(api_key)
+    raise ValueError(f"unknown provider: {name}")