Spaces:

ProthamD
/

prothamd-adaptive-world-env

Running

prothamd-adaptive-world-env / server /app.py

PROTHAM

Serve SPA with no-cache headers so HF UI picks up static/index.html after deploy

0db7f46 6 months ago

10.1 kB

import os
import threading
import uuid
from typing import List, Optional

from fastapi import FastAPI
from fastapi.responses import FileResponse
from fastapi.staticfiles import StaticFiles
from openenv.core.env_server import create_app
from pydantic import BaseModel

from models import AdaptiveAction, AdaptiveObservation
from server.adaptive_world_environment import AdaptiveWorldEnvironment
from server.mock_api import router as mock_router

_HERE   = os.path.dirname(os.path.abspath(__file__))   # server/
_ROOT   = os.path.dirname(_HERE)                        # repo root
_STATIC = os.path.join(_ROOT, "static")                # repo root/static/

# In-memory store for stateful 2-phase episodes
_sessions: dict = {}
_sessions_lock = threading.Lock()
import asyncio as _asyncio
_episode_lock = _asyncio.Lock()

app = create_app(AdaptiveWorldEnvironment, AdaptiveAction, AdaptiveObservation,
                 env_name="adaptive_world_env")

app.include_router(mock_router)

# Serve static assets (SPA + chart images)
app.mount("/static", StaticFiles(directory=_STATIC), name="static")


@app.get("/", include_in_schema=False)
async def serve_spa():
    # HF/proxy/browser caches often keep old index.html after git push; force revalidation.
    path = os.path.join(_STATIC, "index.html")
    return FileResponse(
        path,
        headers={
            "Cache-Control": "no-cache, no-store, must-revalidate",
            "Pragma": "no-cache",
            "Expires": "0",
        },
    )


@app.get("/health")
async def health():
    return {"status": "ok", "environment": "adaptive-world-env", "version": "2.3"}


# ── /run_episode  ─────────────────────────────────────────────────────────────
# OpenEnv's /reset and /step are stateless (fresh env per request).
# This endpoint runs a full multi-step episode inside a SINGLE env instance
# so state (step_count, drift_injected, world_truth) accumulates correctly.

class EpisodeRequest(BaseModel):
    scenario_id: str = "auto"
    difficulty: str = "easy"
    actions: List[dict]


class StepResult(BaseModel):
    step: int
    action_type: str
    status_code: Optional[int] = None
    response_body: Optional[str] = None
    feedback: Optional[str] = None
    done: bool = False
    task_reward: Optional[float] = None
    belief_accuracy: Optional[float] = None
    reward: Optional[float] = None


class EpisodeResponse(BaseModel):
    task_reward: float
    belief_accuracy: float
    reward: float
    steps_taken: int
    task_completed: bool
    steps: List[StepResult]


@app.post("/run_episode", response_model=EpisodeResponse)
async def run_episode(req: EpisodeRequest) -> EpisodeResponse:
    """
    Run a full episode in a single env instance.
    Each env.step() is run in a thread-pool executor so the asyncio event
    loop stays free to process the internal localhost HTTP calls that
    _execute_api_call makes to /mock_api/*.  Without this, the async
    handler blocks the event loop and those self-calls deadlock (status=0).
    """
    import asyncio
    from concurrent.futures import ThreadPoolExecutor

    loop = asyncio.get_event_loop()
    # Dedicated single-thread executor keeps env calls sequential
    executor = ThreadPoolExecutor(max_workers=1)

    def _reset():
        env = AdaptiveWorldEnvironment()
        env.reset(scenario_id=req.scenario_id, difficulty=req.difficulty)
        return env

    env: AdaptiveWorldEnvironment = await loop.run_in_executor(executor, _reset)

    step_results: List[StepResult] = []
    final_task_reward = 0.0
    final_belief_accuracy = 0.0
    final_reward = 0.001
    final_done = False

    for raw_action in req.actions:
        try:
            action = AdaptiveAction(**raw_action)
        except Exception:
            continue

        obs: AdaptiveObservation = await loop.run_in_executor(
            executor, lambda a=action: env.step(a)
        )

        sr = StepResult(
            step=env.state.step_count,
            action_type=action.action_type,
            status_code=obs.last_status_code,
            response_body=(obs.last_response_body or "")[:500],
            feedback=(obs.step_feedback or "")[:300],
            done=obs.done,
        )

        if obs.done:
            sr.task_reward = obs.task_reward
            sr.belief_accuracy = obs.belief_accuracy
            sr.reward = obs.reward
            final_task_reward = float(obs.task_reward or 0.0)
            final_belief_accuracy = float(obs.belief_accuracy or 0.0)
            final_reward = float(obs.reward or 0.001)
            final_done = True

        step_results.append(sr)

        if final_done:
            break

    # If episode didn't end via submit_result, force-submit with no belief
    if not final_done:
        obs = await loop.run_in_executor(
            executor,
            lambda: env.step(AdaptiveAction(action_type="submit_result", belief_state={}))
        )
        final_task_reward = float(obs.task_reward or 0.0)
        final_belief_accuracy = float(obs.belief_accuracy or 0.0)
        final_reward = float(obs.reward or 0.001)
        step_results.append(StepResult(
            step=env.state.step_count,
            action_type="submit_result",
            done=True,
            task_reward=final_task_reward,
            belief_accuracy=final_belief_accuracy,
            reward=final_reward,
        ))

    executor.shutdown(wait=False)

    return EpisodeResponse(
        task_reward=final_task_reward,
        belief_accuracy=final_belief_accuracy,
        reward=final_reward,
        steps_taken=env.state.step_count,
        task_completed=env.state.task_completed,
        steps=step_results,
    )


# ── /start_episode + /finish_episode  ─────────────────────────────────────────
# 3-phase belief training:
#   Phase 1 (notebook): model generates task action
#   Phase 2 (server):   /start_episode runs pre-drift + query_history + probe_schema
#                       returns session_id + probe evidence to notebook
#   Phase 3 (notebook): model generates belief FROM probe evidence
#                       calls /finish_episode with belief → gets final scores

class StartEpisodeRequest(BaseModel):
    scenario_id: str = "auto"
    difficulty:  str = "easy"
    task_action: dict


class StartEpisodeResponse(BaseModel):
    session_id:       str
    probe_response:   str
    history_response: str
    pre_drift_ok:     bool


class FinishEpisodeRequest(BaseModel):
    session_id:    str
    task_action:   dict
    belief_state:  dict


@app.post("/start_episode", response_model=StartEpisodeResponse)
async def start_episode(req: StartEpisodeRequest) -> StartEpisodeResponse:
    """
    Phase 1 of 2-phase episode.
    Runs: N × task_action → query_history → probe_schema
    Stores the live env in memory; returns probe evidence to the notebook.
    """
    import asyncio
    from concurrent.futures import ThreadPoolExecutor

    loop     = asyncio.get_event_loop()
    executor = ThreadPoolExecutor(max_workers=1)
    pre_n    = {"easy": 3, "medium": 4, "hard": 4}.get(req.difficulty, 3)

    def _run():
        env = AdaptiveWorldEnvironment()
        env.reset(scenario_id=req.scenario_id, difficulty=req.difficulty)

        task_action   = AdaptiveAction(**req.task_action)
        pre_drift_ok  = False

        for _ in range(pre_n):
            obs = env.step(task_action)
            if obs.last_status_code and obs.last_status_code < 300:
                pre_drift_ok = True

        history_obs      = env.step(AdaptiveAction(action_type="query_history", history_steps=3))
        probe_obs        = env.step(AdaptiveAction(action_type="probe_schema"))
        history_response = (history_obs.last_response_body or "")[:500]
        probe_response   = (probe_obs.last_response_body   or "")[:500]
        return env, probe_response, history_response, pre_drift_ok

    async with _episode_lock:
        env, probe_response, history_response, pre_drift_ok = await loop.run_in_executor(executor, _run)

    sid = str(uuid.uuid4())[:8]
    with _sessions_lock:
        _sessions[sid] = {"env": env, "executor": executor}

    return StartEpisodeResponse(
        session_id       = sid,
        probe_response   = probe_response,
        history_response = history_response,
        pre_drift_ok     = pre_drift_ok,
    )


@app.post("/finish_episode", response_model=EpisodeResponse)
async def finish_episode(req: FinishEpisodeRequest) -> EpisodeResponse:
    """
    Phase 2 of 2-phase episode.
    Runs: corrected task_action → submit_result with model's belief_state.
    Returns final task_reward + belief_accuracy.
    """
    import asyncio

    with _sessions_lock:
        session = _sessions.pop(req.session_id, None)

    if session is None:
        return EpisodeResponse(
            task_reward=0.0, belief_accuracy=0.0, reward=0.0,
            steps_taken=0, task_completed=False, steps=[],
        )

    env      = session["env"]
    executor = session["executor"]
    loop     = asyncio.get_event_loop()

    def _finish():
        # One corrected call after probing
        env.step(AdaptiveAction(**req.task_action))
        # Submit belief
        obs = env.step(AdaptiveAction(
            action_type  = "submit_result",
            belief_state = req.belief_state,
        ))
        return obs

    obs      = await loop.run_in_executor(executor, _finish)
    executor.shutdown(wait=False)

    return EpisodeResponse(
        task_reward     = float(obs.task_reward     or 0.0),
        belief_accuracy = float(obs.belief_accuracy or 0.0),
        reward          = float(obs.reward          or 0.0),
        steps_taken     = env.state.step_count,
        task_completed  = env.state.task_completed,
        steps           = [],
    )


def main():
    import uvicorn
    uvicorn.run("server.app:app", host="0.0.0.0", port=7860)


if __name__ == "__main__":
    main()

	import os
	import threading
	import uuid
	from typing import List, Optional

	from fastapi import FastAPI
	from fastapi.responses import FileResponse
	from fastapi.staticfiles import StaticFiles
	from openenv.core.env_server import create_app
	from pydantic import BaseModel

	from models import AdaptiveAction, AdaptiveObservation
	from server.adaptive_world_environment import AdaptiveWorldEnvironment
	from server.mock_api import router as mock_router

	_HERE = os.path.dirname(os.path.abspath(__file__)) # server/
	_ROOT = os.path.dirname(_HERE) # repo root
	_STATIC = os.path.join(_ROOT, "static") # repo root/static/

	# In-memory store for stateful 2-phase episodes
	_sessions: dict = {}
	_sessions_lock = threading.Lock()
	import asyncio as _asyncio
	_episode_lock = _asyncio.Lock()

	app = create_app(AdaptiveWorldEnvironment, AdaptiveAction, AdaptiveObservation,
	env_name="adaptive_world_env")

	app.include_router(mock_router)

	# Serve static assets (SPA + chart images)
	app.mount("/static", StaticFiles(directory=_STATIC), name="static")


	@app.get("/", include_in_schema=False)
	async def serve_spa():
	# HF/proxy/browser caches often keep old index.html after git push; force revalidation.
	path = os.path.join(_STATIC, "index.html")
	return FileResponse(
	path,
	headers={
	"Cache-Control": "no-cache, no-store, must-revalidate",
	"Pragma": "no-cache",
	"Expires": "0",
	},
	)


	@app.get("/health")
	async def health():
	return {"status": "ok", "environment": "adaptive-world-env", "version": "2.3"}


	# ── /run_episode ─────────────────────────────────────────────────────────────
	# OpenEnv's /reset and /step are stateless (fresh env per request).
	# This endpoint runs a full multi-step episode inside a SINGLE env instance
	# so state (step_count, drift_injected, world_truth) accumulates correctly.

	class EpisodeRequest(BaseModel):
	scenario_id: str = "auto"
	difficulty: str = "easy"
	actions: List[dict]


	class StepResult(BaseModel):
	step: int
	action_type: str
	status_code: Optional[int] = None
	response_body: Optional[str] = None
	feedback: Optional[str] = None
	done: bool = False
	task_reward: Optional[float] = None
	belief_accuracy: Optional[float] = None
	reward: Optional[float] = None


	class EpisodeResponse(BaseModel):
	task_reward: float
	belief_accuracy: float
	reward: float
	steps_taken: int
	task_completed: bool
	steps: List[StepResult]


	@app.post("/run_episode", response_model=EpisodeResponse)
	async def run_episode(req: EpisodeRequest) -> EpisodeResponse:
	"""
	Run a full episode in a single env instance.

	Each env.step() is run in a thread-pool executor so the asyncio event
	loop stays free to process the internal localhost HTTP calls that
	_execute_api_call makes to /mock_api/*. Without this, the async
	handler blocks the event loop and those self-calls deadlock (status=0).
	"""
	import asyncio
	from concurrent.futures import ThreadPoolExecutor

	loop = asyncio.get_event_loop()
	# Dedicated single-thread executor keeps env calls sequential
	executor = ThreadPoolExecutor(max_workers=1)

	def _reset():
	env = AdaptiveWorldEnvironment()
	env.reset(scenario_id=req.scenario_id, difficulty=req.difficulty)
	return env

	env: AdaptiveWorldEnvironment = await loop.run_in_executor(executor, _reset)

	step_results: List[StepResult] = []
	final_task_reward = 0.0
	final_belief_accuracy = 0.0
	final_reward = 0.001
	final_done = False

	for raw_action in req.actions:
	try:
	action = AdaptiveAction(**raw_action)
	except Exception:
	continue

	obs: AdaptiveObservation = await loop.run_in_executor(
	executor, lambda a=action: env.step(a)
	)

	sr = StepResult(
	step=env.state.step_count,
	action_type=action.action_type,
	status_code=obs.last_status_code,
	response_body=(obs.last_response_body or "")[:500],
	feedback=(obs.step_feedback or "")[:300],
	done=obs.done,
	)

	if obs.done:
	sr.task_reward = obs.task_reward
	sr.belief_accuracy = obs.belief_accuracy
	sr.reward = obs.reward
	final_task_reward = float(obs.task_reward or 0.0)
	final_belief_accuracy = float(obs.belief_accuracy or 0.0)
	final_reward = float(obs.reward or 0.001)
	final_done = True

	step_results.append(sr)

	if final_done:
	break

	# If episode didn't end via submit_result, force-submit with no belief
	if not final_done:
	obs = await loop.run_in_executor(
	executor,
	lambda: env.step(AdaptiveAction(action_type="submit_result", belief_state={}))
	)
	final_task_reward = float(obs.task_reward or 0.0)
	final_belief_accuracy = float(obs.belief_accuracy or 0.0)
	final_reward = float(obs.reward or 0.001)
	step_results.append(StepResult(
	step=env.state.step_count,
	action_type="submit_result",
	done=True,
	task_reward=final_task_reward,
	belief_accuracy=final_belief_accuracy,
	reward=final_reward,
	))

	executor.shutdown(wait=False)

	return EpisodeResponse(
	task_reward=final_task_reward,
	belief_accuracy=final_belief_accuracy,
	reward=final_reward,
	steps_taken=env.state.step_count,
	task_completed=env.state.task_completed,
	steps=step_results,
	)


	# ── /start_episode + /finish_episode ─────────────────────────────────────────
	# 3-phase belief training:
	# Phase 1 (notebook): model generates task action
	# Phase 2 (server): /start_episode runs pre-drift + query_history + probe_schema
	# returns session_id + probe evidence to notebook
	# Phase 3 (notebook): model generates belief FROM probe evidence
	# calls /finish_episode with belief → gets final scores

	class StartEpisodeRequest(BaseModel):
	scenario_id: str = "auto"
	difficulty: str = "easy"
	task_action: dict


	class StartEpisodeResponse(BaseModel):
	session_id: str
	probe_response: str
	history_response: str
	pre_drift_ok: bool


	class FinishEpisodeRequest(BaseModel):
	session_id: str
	task_action: dict
	belief_state: dict


	@app.post("/start_episode", response_model=StartEpisodeResponse)
	async def start_episode(req: StartEpisodeRequest) -> StartEpisodeResponse:
	"""
	Phase 1 of 2-phase episode.
	Runs: N × task_action → query_history → probe_schema
	Stores the live env in memory; returns probe evidence to the notebook.
	"""
	import asyncio
	from concurrent.futures import ThreadPoolExecutor

	loop = asyncio.get_event_loop()
	executor = ThreadPoolExecutor(max_workers=1)
	pre_n = {"easy": 3, "medium": 4, "hard": 4}.get(req.difficulty, 3)

	def _run():
	env = AdaptiveWorldEnvironment()
	env.reset(scenario_id=req.scenario_id, difficulty=req.difficulty)

	task_action = AdaptiveAction(**req.task_action)
	pre_drift_ok = False

	for _ in range(pre_n):
	obs = env.step(task_action)
	if obs.last_status_code and obs.last_status_code < 300:
	pre_drift_ok = True

	history_obs = env.step(AdaptiveAction(action_type="query_history", history_steps=3))
	probe_obs = env.step(AdaptiveAction(action_type="probe_schema"))
	history_response = (history_obs.last_response_body or "")[:500]
	probe_response = (probe_obs.last_response_body or "")[:500]
	return env, probe_response, history_response, pre_drift_ok

	async with _episode_lock:
	env, probe_response, history_response, pre_drift_ok = await loop.run_in_executor(executor, _run)

	sid = str(uuid.uuid4())[:8]
	with _sessions_lock:
	_sessions[sid] = {"env": env, "executor": executor}

	return StartEpisodeResponse(
	session_id = sid,
	probe_response = probe_response,
	history_response = history_response,
	pre_drift_ok = pre_drift_ok,
	)


	@app.post("/finish_episode", response_model=EpisodeResponse)
	async def finish_episode(req: FinishEpisodeRequest) -> EpisodeResponse:
	"""
	Phase 2 of 2-phase episode.
	Runs: corrected task_action → submit_result with model's belief_state.
	Returns final task_reward + belief_accuracy.
	"""
	import asyncio

	with _sessions_lock:
	session = _sessions.pop(req.session_id, None)

	if session is None:
	return EpisodeResponse(
	task_reward=0.0, belief_accuracy=0.0, reward=0.0,
	steps_taken=0, task_completed=False, steps=[],
	)

	env = session["env"]
	executor = session["executor"]
	loop = asyncio.get_event_loop()

	def _finish():
	# One corrected call after probing
	env.step(AdaptiveAction(**req.task_action))
	# Submit belief
	obs = env.step(AdaptiveAction(
	action_type = "submit_result",
	belief_state = req.belief_state,
	))
	return obs

	obs = await loop.run_in_executor(executor, _finish)
	executor.shutdown(wait=False)

	return EpisodeResponse(
	task_reward = float(obs.task_reward or 0.0),
	belief_accuracy = float(obs.belief_accuracy or 0.0),
	reward = float(obs.reward or 0.0),
	steps_taken = env.state.step_count,
	task_completed = env.state.task_completed,
	steps = [],
	)


	def main():
	import uvicorn
	uvicorn.run("server.app:app", host="0.0.0.0", port=7860)


	if __name__ == "__main__":
	main()