feat: add custom Swagger docs endpoint and import get_swagger_ui_html

api.py

@@ -1,9 +1,12 @@
+from fastapi.openapi.docs import get_swagger_ui_html
 from fastapi import FastAPI
 from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel
+from typing import Optional
+from dataclasses import asdict
 from environment import CustomerSupportEnv
 
-app = FastAPI(title="Customer Support AI OpenEnv")
+app = FastAPI(title="Customer Support AI Environment")
 
 app.add_middleware(
     CORSMiddleware,
@@ -23,59 +26,69 @@ class StepRequest(BaseModel):
     action: str
 
 
-def obs_to_dict(obs) -> dict:
-    step = obs.current_step
-    return {
-        "task_id": obs.task_id,
-        "task_description": obs.task_description,
-        "current_customer_message": obs.current_customer_message,
-        "current_step": {
-            "step_id": str(step.step_id),
-            "label": str(step.label),
-            "description": str(step.description),
-        },
-        "progress": obs.progress,
-        "step_number": int(obs.step_number),
-        "total_steps": int(obs.total_steps),
-        "max_steps": int(obs.max_steps),
-        "metadata": obs.metadata,
-        "conversation_history": obs.conversation_history,
-        "context_memory": obs.context_memory,
-        "steps_completed": list(obs.steps_completed),
-        "consecutive_failures": int(obs.consecutive_failures),
-        "episode_status": str(obs.episode_status),
-        "fail_conditions": list(obs.fail_conditions),
-    }
-
-
 @app.get("/")
-def root():
-    return {"status": "ok", "message": "Customer Support AI OpenEnv API"}
+def health_check():
+    return {"status": "ok", "message": "Customer Support AI Environment"}
 
 
 @app.post("/reset")
 def reset_env(body: ResetRequest = None):
     global _env
-    task = (body.task if body else "easy").lower()
-    if task not in ["easy", "medium", "hard"]:
-        task = "easy"
-    _env = CustomerSupportEnv(task=task)
+    task = body.task if body else "easy"
+    _env = CustomerSupportEnv(task=task.lower())
     obs = _env.reset()
-    return {"observation": obs_to_dict(obs)}
+    return {
+        "observation": {
+            "task_id": obs.task_id,
+            "task_description": obs.task_description,
+            "current_customer_message": obs.current_customer_message,
+            "current_step": {
+                "step_id": obs.current_step.step_id,
+                "label": obs.current_step.label,
+                "description": obs.current_step.description,
+            },
+            "progress": obs.progress,
+            "step_number": obs.step_number,
+            "total_steps": obs.total_steps,
+            "max_steps": obs.max_steps,
+            "metadata": obs.metadata,
+            "conversation_history": obs.conversation_history,
+            "context_memory": obs.context_memory,
+            "episode_status": obs.episode_status,
+        }
+    }
 
 
 @app.post("/step")
 def step_env(body: StepRequest):
     global _env
     if _env is None:
-        return {"error": "Call /reset first"}
+        return {"error": "Environment not initialized. Call /reset first."}
     result = _env.step(body.action)
-    info = result.info.copy()
+    obs = result.observation
     return {
-        "observation": obs_to_dict(result.observation),
-        "reward": float(result.reward),
-        "done": bool(result.done),
-        "info": info,
+        "observation": {
+            "task_id": obs.task_id,
+            "task_description": obs.task_description,
+            "current_customer_message": obs.current_customer_message,
+            "current_step": {
+                "step_id": obs.current_step.step_id,
+                "label": obs.current_step.label,
+                "description": obs.current_step.description,
+            },
+            "progress": obs.progress,
+            "step_number": obs.step_number,
+            "total_steps": obs.total_steps,
+            "max_steps": obs.max_steps,
+            "metadata": obs.metadata,
+            "conversation_history": obs.conversation_history,
+            "context_memory": obs.context_memory,
+            "episode_status": obs.episode_status,
+            "steps_completed": obs.steps_completed,
+        },
+        "reward": result.reward,
+        "done": result.done,
+        "info": result.info,
     }
 
 
@@ -83,7 +96,7 @@ def step_env(body: StepRequest):
 def get_state():
     global _env
     if _env is None:
-        return {"error": "Call /reset first"}
+        return {"error": "Environment not initialized. Call /reset first."}
     return _env.state()
 
 
@@ -96,6 +109,7 @@ def observation_space():
             "current_customer_message": "str",
             "current_step": "StepInfo",
             "progress": "str",
+            "conversation_history": "List[Dict]",
             "episode_status": "str",
         }
     }
@@ -105,6 +119,9 @@ def observation_space():
 def action_space():
     return {
         "type": "Text",
-        "description": "Agent free-text reply",
+        "description": "Agent free-text reply to customer",
         "min_words": 6,
     }
+@app.get("/docs", include_in_schema=False)
+def custom_docs():
+    return get_swagger_ui_html(openapi_url="/openapi.json", title="API Docs")