Enhance Content Moderation Environment: Update Dockerfile, add client, improve README, and refine environment handling

.gitignore

@@ -40,3 +40,7 @@ MANIFEST
 
 # uv
 uv.lock
+
+# llms
+.claude/
+.kilo/

Dockerfile

@@ -1,18 +1,42 @@
 FROM python:3.11-slim
 
-# Install git and other build deps just in case
-RUN apt-get update && apt-get install -y git && rm -rf /var/lib/apt/lists/*
+# Install system dependencies and cleanup
+RUN apt-get update && apt-get install -y --no-install-recommends \
+    git \
+    curl \
+    && rm -rf /var/lib/apt/lists/* \
+    && apt-get clean
 
 # Set working directory to the standard Hugging Face directory
 WORKDIR /code
 
-# Copy all the current files from the repo root to /code
-COPY . /code
+# Copy requirements first for better caching
+COPY pyproject.toml uv.lock ./
 
-# Install the project locally along with OpenEnv
-RUN pip install --no-cache-dir -e .
+# Install dependencies with uv for faster builds
+RUN pip install --no-cache-dir uv && \
+    uv sync --frozen && \
+    uv pip install --no-cache-dir -e .
+
+# Copy the rest of the application
+COPY . .
+
+# Create non-root user for security
+RUN useradd --create-home --shell /bin/bash app && \
+    chown -R app:app /code
+USER app
+
+# Health check endpoint
+HEALTHCHECK --interval=30s --timeout=10s --start-period=5s --retries=3 \
+    CMD curl -f http://localhost:8000/health || exit 1
 
 EXPOSE 8000
 
-# Run Uvicorn with factory mode
-CMD ["uvicorn", "server.app:app", "--host", "0.0.0.0", "--port", "8000"]
+# Run Uvicorn with optimized settings for production
+CMD ["uvicorn", "server.app:app", \
+     "--host", "0.0.0.0", \
+     "--port", "8000", \
+     "--workers", "1", \
+     "--timeout", "300", \
+     "--timeout-keep-alive", "65", \
+     "--log-level", "info"]

README.md

@@ -26,6 +26,35 @@ A small, deterministic OpenEnv for evaluating AI agents on user-generated conten
 
 ---
 
+## Environment Specification
+
+**OpenEnv Configuration**: Minimal specification in `openenv.yaml` matching reference environments:
+- `spec_version: 1`
+- `name: content_moderation_env` 
+- `type: space`
+- `runtime: fastapi`
+- `app: server.app:app`
+- `port: 8000`
+
+**Technical Details**:
+- **Version**: 0.1.0
+- **Python**: >=3.10
+- **License**: MIT
+- **Authors**: Saheb Motiani ([@saheb](https://github.com/saheb))
+- **Environment Type**: Sequential decision making
+- **Observation Space**: Structured text with Pydantic models
+- **Action Space**: Discrete moderation actions
+- **Max Episode Length**: 150 steps (very_hard)
+- **Deterministic**: Yes - fully reproducible task packs
+- **Reward Shaping**: Yes - dense rewards with policy bonuses
+- **Stateful**: Yes - maintains episode state
+
+**Tags**: content-moderation, reasoning, policy-grounded, multi-step, deterministic
+
+**Keywords**: content moderation, AI safety, policy reasoning, context awareness, reward shaping, thread context, multi-step reasoning
+
+---
+
 ## Why content moderation?
 
 Content moderation is one of the few domains where:

client.py

@@ -0,0 +1,131 @@
+# Copyright (c) Saheb Motiani
+# All rights reserved.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""Content Moderation Environment Client."""
+
+from typing import Dict
+
+from openenv.core import EnvClient
+from openenv.core.client_types import StepResult
+from openenv.core.env_server.types import State
+
+from models import ModerationAction, ModerationObservation
+
+
+class ContentModerationEnv(EnvClient[ModerationAction, ModerationObservation, State]):
+    """
+    Client for the Content Moderation Environment.
+
+    This client maintains a persistent WebSocket connection to the environment server,
+    enabling efficient multi-step interactions with lower latency.
+    Each client instance has its own dedicated environment session on the server.
+
+    Example:
+        >>> # Connect to a running server
+        >>> from content_moderation_env import ContentModerationEnv, ModerationAction
+        >>>
+        >>> env = ContentModerationEnv(base_url="http://localhost:8000")
+        >>> # Start with easy task
+        >>> result = env.reset(task_name="easy")
+        >>> print(f"Active post: {result.observation.active_post_summary}")
+        >>>
+        >>> # View the full post
+        >>> result = env.step(ModerationAction(type="view_post", post_id="p1"))
+        >>> print(f"Post content: {result.observation.current_post}")
+        >>>
+        >>> # View thread context if needed
+        >>> result = env.step(ModerationAction(type="view_thread", post_id="p1"))
+        >>> print(f"Thread: {result.observation.thread_context}")
+        >>>
+        >>> # Make moderation decision
+        >>> result = env.step(ModerationAction(
+        ...     type="moderate", 
+        ...     post_id="p1", 
+        ...     decision="keep",
+        ...     rationale="P4"
+        ... ))
+        >>> print(f"Score: {result.observation.grader_score}")
+        >>> print(f"Done: {result.observation.done}")
+        >>> env.close()
+
+    Example with Docker:
+        >>> # Automatically start container and connect
+        >>> client = ContentModerationEnv.from_docker_image("content-moderation-env:latest")
+        >>> try:
+        ...     result = client.reset(task_name="medium")
+        ...     # Process posts...
+        ... finally:
+        ...     client.close()
+    """
+
+    def _step_payload(self, action: ModerationAction) -> Dict:
+        """
+        Convert ModerationAction to JSON payload for step message.
+
+        Args:
+            action: ModerationAction instance
+
+        Returns:
+            Dictionary representation suitable for JSON encoding
+        """
+        payload = {"type": action.type}
+        
+        if action.post_id:
+            payload["post_id"] = action.post_id
+        if action.decision:
+            payload["decision"] = action.decision
+        if action.rationale:
+            payload["rationale"] = action.rationale
+        if action.policy_id:
+            payload["policy_id"] = action.policy_id
+        if action.category:
+            payload["category"] = action.category
+            
+        return payload
+
+    def _parse_result(self, payload: Dict) -> StepResult[ModerationObservation]:
+        """
+        Parse server response into StepResult[ModerationObservation].
+
+        Args:
+            payload: JSON response data from server
+
+        Returns:
+            StepResult with ModerationObservation
+        """
+        obs_data = payload.get("observation", {})
+        observation = ModerationObservation(
+            active_post_summary=obs_data.get("active_post_summary"),
+            failed_attempts=obs_data.get("failed_attempts", []),
+            current_post=obs_data.get("current_post"),
+            thread_context=obs_data.get("thread_context"),
+            last_action_result=obs_data.get("last_action_result", ""),
+            grader_score=obs_data.get("grader_score"),
+            reward=payload.get("reward"),
+            done=payload.get("done", False),
+            metadata=obs_data.get("metadata", {}),
+        )
+
+        return StepResult(
+            observation=observation,
+            reward=payload.get("reward"),
+            done=payload.get("done", False),
+        )
+
+    def _parse_state(self, payload: Dict) -> State:
+        """
+        Parse server response into State object.
+
+        Args:
+            payload: JSON response from state request
+
+        Returns:
+            State object with episode_id and step_count
+        """
+        return State(
+            episode_id=payload.get("episode_id"),
+            step_count=payload.get("step_count", 0),
+        )

openenv.yaml

@@ -1,14 +1,6 @@
-name: content-moderation-env
-version: 0.1.0
-description: "OpenEnv package for UGC content moderation (VibeNet mock)"
-authors:
-  - name: "Saheb Motiani"
-    github: "https://github.com/saheb"
-entrypoint: server.app:app
-python: ">=3.10"
-dependencies:
-  - openenv-core>=0.2.0
-  - pydantic>=2.0
-  - fastapi>=0.110
-  - uvicorn
-license: MIT
+spec_version: 1
+name: content_moderation_env
+type: space
+runtime: fastapi
+app: server.app:app
+port: 8000

requirements.txt

@@ -0,0 +1,9 @@
+# Core dependencies for content-moderation-env
+openenv-core>=0.2.0
+pydantic>=2.0
+fastapi>=0.110
+uvicorn[standard]>=0.25.0
+python-dotenv>=1.0.0
+
+# Optional dependencies for baseline inference
+openai>=1.0

server/app.py

@@ -1,6 +1,8 @@
 import os
 import uuid
 import logging
+import sys
+from typing import Optional
 
 import uvicorn
 from fastapi import HTTPException
@@ -11,11 +13,22 @@ from env_loader import load_environment
 from .environment import ContentModerationEnvironment
 from models import ModerationAction, ModerationObservation
 
+# Configure logging
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
+    handlers=[logging.StreamHandler(sys.stdout)]
+)
 logger = logging.getLogger(__name__)
 
-_loaded_env_file = load_environment()
-if _loaded_env_file:
-    logger.info("Loaded environment variables from %s", _loaded_env_file)
+# Load environment variables safely
+_loaded_env_file: Optional[str] = None
+try:
+    _loaded_env_file = load_environment()
+    if _loaded_env_file:
+        logger.info("Loaded environment variables from %s", _loaded_env_file)
+except Exception as e:
+    logger.warning("Failed to load environment variables: %s", e)
 
 # Demo session store — keyed by session UUID, capped at 50 concurrent sessions
 _demo_sessions: dict = {}
@@ -25,23 +38,56 @@ _MAX_DEMO_SESSIONS = 50
 def create_app():
     """Factory function for creating the FastAPI app instance."""
     logger.info("Creating FastAPI app...")
-    app = create_fastapi_app(
-        lambda: ContentModerationEnvironment(),
-        action_cls=ModerationAction,
-        observation_cls=ModerationObservation,
-    )
+    
+    try:
+        app = create_fastapi_app(
+            lambda: ContentModerationEnvironment(),
+            action_cls=ModerationAction,
+            observation_cls=ModerationObservation,
+        )
+        logger.info("OpenEnv app created successfully")
+    except Exception as e:
+        logger.error("Failed to create OpenEnv app: %s", e)
+        raise
 
     @app.get("/", response_class=HTMLResponse, tags=["UI"])
     def root():
         """Root endpoint serving the interactive demo UI."""
-        demo_html_path = os.path.join(os.path.dirname(__file__), "..", "demo.html")
-        with open(demo_html_path, encoding="utf-8") as f:
-            return HTMLResponse(content=f.read())
+        try:
+            demo_html_path = os.path.join(os.path.dirname(__file__), "..", "demo.html")
+            with open(demo_html_path, encoding="utf-8") as f:
+                return HTMLResponse(content=f.read())
+        except FileNotFoundError:
+            logger.warning("demo.html not found, serving simple response")
+            return HTMLResponse(content="<h1>Content Moderation Environment</h1><p>Interactive demo not available</p>")
+        except Exception as e:
+            logger.error("Error serving demo: %s", e)
+            return HTMLResponse(content="<h1>Content Moderation Environment</h1><p>Error loading demo</p>", status_code=500)
 
     @app.get("/health", tags=["Health"])
     def health():
         """Standard health check endpoint for monitoring."""
-        return {"status": "healthy"}
+        try:
+            # Test environment creation to ensure it's working
+            test_env = ContentModerationEnvironment()
+            test_obs = test_env.reset(task_name="easy")
+            return {
+                "status": "healthy",
+                "service": "content-moderation-env",
+                "environment_test": "passed"
+            }
+        except Exception as e:
+            logger.error("Health check failed: %s", e)
+            return {
+                "status": "unhealthy",
+                "service": "content-moderation-env", 
+                "error": str(e)
+            }, 500
+
+    @app.get("/ready", tags=["Health"])
+    def ready():
+        """Readiness endpoint often used by cloud platforms for startup probes."""
+        return {"status": "ready"}
 
     logger.info("FastAPI app created successfully")
     return app
@@ -57,62 +103,86 @@ if __name__ == "__main__":
     main()
 
 # Module-level app for runners that expect `server.app:app` (no --factory)
-app = create_app()
-
-
-@app.get("/ready", tags=["Health"])
-def ready():
-    """Readiness endpoint often used by cloud platforms for startup probes."""
-    return {"status": "ready"}
-
-
-# ── Demo endpoints ────────────────────────────────────────────────────────────
-
-@app.get("/demo", response_class=HTMLResponse, tags=["Demo"])
-def demo_page():
-    """Serve the interactive human-playable demo UI."""
-    demo_html_path = os.path.join(os.path.dirname(__file__), "..", "demo.html")
-    with open(demo_html_path, encoding="utf-8") as f:
-        return HTMLResponse(content=f.read())
-
-
-@app.get("/benchmark.html", response_class=HTMLResponse, tags=["UI"])
-def benchmark_page():
-    """Serve the interactive benchmark visualization UI."""
-    benchmark_html_path = os.path.join(os.path.dirname(__file__), "..", "benchmark.html")
-    with open(benchmark_html_path, encoding="utf-8") as f:
-        return HTMLResponse(content=f.read())
-
-
-@app.post("/demo/reset", tags=["Demo"])
-def demo_reset(task: str = "easy"):
-    """Start a new demo session. Returns session_id and initial observation."""
-    if task not in ("easy", "medium", "hard", "very_hard"):
-        raise HTTPException(status_code=400, detail=f"Unknown task: {task}")
-
-    # Evict oldest session if at capacity
-    if len(_demo_sessions) >= _MAX_DEMO_SESSIONS:
-        oldest_key = next(iter(_demo_sessions))
-        del _demo_sessions[oldest_key]
-
-    session_id = str(uuid.uuid4())
-    env = ContentModerationEnvironment()
-    obs = env.reset(task_name=task)
-    _demo_sessions[session_id] = env
-
-    return {"session_id": session_id, "observation": obs.model_dump()}
-
-
-@app.post("/demo/step/{session_id}", tags=["Demo"])
-def demo_step(session_id: str, action: ModerationAction):
-    """Submit one action in an existing demo session."""
-    env = _demo_sessions.get(session_id)
-    if env is None:
-        raise HTTPException(status_code=404, detail="Session not found or expired. Start a new session via /demo/reset.")
-
-    obs = env.step(action)
-
-    if obs.done:
-        del _demo_sessions[session_id]
-
-    return {"observation": obs.model_dump()}
+# Create lazily to avoid import-time errors
+_app = None
+
+def get_app():
+    global _app
+    if _app is None:
+        _app = create_app()
+    return _app
+
+# For direct imports
+app = get_app()
+
+# Demo endpoints - add them to the app with error handling
+try:
+    current_app = get_app()
+    
+    @current_app.get("/demo", response_class=HTMLResponse, tags=["Demo"])
+    def demo_page():
+        """Serve the interactive human-playable demo UI."""
+        try:
+            demo_html_path = os.path.join(os.path.dirname(__file__), "..", "demo.html")
+            with open(demo_html_path, encoding="utf-8") as f:
+                return HTMLResponse(content=f.read())
+        except FileNotFoundError:
+            logger.warning("demo.html not found")
+            return HTMLResponse(content="<h1>Demo Not Available</h1><p>demo.html file not found</p>")
+        except Exception as e:
+            logger.error("Error serving demo page: %s", e)
+            return HTMLResponse(content="<h1>Demo Error</h1><p>Failed to load demo</p>", status_code=500)
+
+    @current_app.get("/benchmark.html", response_class=HTMLResponse, tags=["UI"])
+    def benchmark_page():
+        """Serve the interactive benchmark visualization UI."""
+        try:
+            benchmark_html_path = os.path.join(os.path.dirname(__file__), "..", "benchmark.html")
+            with open(benchmark_html_path, encoding="utf-8") as f:
+                return HTMLResponse(content=f.read())
+        except FileNotFoundError:
+            logger.warning("benchmark.html not found")
+            return HTMLResponse(content="<h1>Benchmark Not Available</h1><p>benchmark.html file not found</p>")
+        except Exception as e:
+            logger.error("Error serving benchmark page: %s", e)
+            return HTMLResponse(content="<h1>Benchmark Error</h1><p>Failed to load benchmark</p>", status_code=500)
+
+    @current_app.post("/demo/reset", tags=["Demo"])
+    def demo_reset(task: str = "easy"):
+        """Start a new demo session. Returns session_id and initial observation."""
+        if task not in ("easy", "medium", "hard", "very_hard"):
+            raise HTTPException(status_code=400, detail=f"Unknown task: {task}")
+
+        # Evict oldest session if at capacity
+        if len(_demo_sessions) >= _MAX_DEMO_SESSIONS:
+            oldest_key = next(iter(_demo_sessions))
+            del _demo_sessions[oldest_key]
+
+        session_id = str(uuid.uuid4())
+        try:
+            env = ContentModerationEnvironment()
+            obs = env.reset(task_name=task)
+            _demo_sessions[session_id] = env
+            return {"session_id": session_id, "observation": obs.model_dump()}
+        except Exception as e:
+            logger.error("Error creating demo session: %s", e)
+            raise HTTPException(status_code=500, detail=f"Failed to create demo session: {e}")
+
+    @current_app.post("/demo/step/{session_id}", tags=["Demo"])
+    def demo_step(session_id: str, action: ModerationAction):
+        """Submit one action in an existing demo session."""
+        env = _demo_sessions.get(session_id)
+        if env is None:
+            raise HTTPException(status_code=404, detail="Session not found or expired. Start a new session via /demo/reset.")
+
+        try:
+            obs = env.step(action)
+            if obs.done:
+                del _demo_sessions[session_id]
+            return {"observation": obs.model_dump()}
+        except Exception as e:
+            logger.error("Error in demo step: %s", e)
+            raise HTTPException(status_code=500, detail=f"Failed to process action: {e}")
+
+except Exception as e:
+    logger.error("Failed to register demo endpoints: %s", e)

server/environment.py

@@ -61,12 +61,16 @@ class ContentModerationEnvironment(Environment):
 
         active_post = self._get_active_post()
         
+        # Always compute grader score (never None) - required for OpenEnv compatibility
+        initial_grader_score = self._compute_grader_score()
+        
         return ModerationObservation(
             active_post_summary=active_post,
             failed_attempts=self.post_failures.get(active_post["id"], []) if active_post else [],
             last_action_result="Episode started — VibeNet moderation pipeline ready",
             current_post=None,
             thread_context=None,
+            grader_score=initial_grader_score,
             done=False,
             reward=0.0,
             metadata={"task": task_name, "total_posts": len(self.task_data["items"])}
@@ -209,7 +213,8 @@ class ContentModerationEnvironment(Environment):
         # Episode termination
         done = len(self.moderated) == len(self.task_data["items"]) or self.step_count >= self.task_data["max_steps"] or loop_exceeded
 
-        grader = self._compute_grader_score() if done else None
+        # Always compute grader score (never None) - required for OpenEnv compatibility
+        grader = self._compute_grader_score()
 
         active_post = self._get_active_post()