""" Root-level Gradio app for WorldSmithAI. This file is intended to live directly beside callbacks.py in the repository root. It does not assume an app/ package or folder. The UI is intentionally thin: - app.py defines the Gradio interface. - callbacks.py runs generation, parsing, validation, world construction, simulation, visualization, metrics, and narration. - The simulation engine remains generic and DSL-driven. Environment variables: WORLDSMITHAI_MODEL_ID: Optional Hugging Face model id for model-generated DSL. HF_TOKEN: Optional Hugging Face token for private or gated models. WORLDSMITHAI_OUTPUT_DIR: Optional directory for generated GIFs, MP4s, charts, and images. WORLDSMITHAI_DEFAULT_STEPS: Optional default simulation step count. WORLDSMITHAI_MAX_FRAMES: Optional maximum animation frame count. WORLDSMITHAI_LOG_LEVEL: Optional logging level. Defaults to INFO. """ from __future__ import annotations import json import logging import os import traceback from collections.abc import Mapping from dataclasses import dataclass from pathlib import Path from typing import Any import gradio as gr from callbacks import ( CallbackConfig, CallbackResult, apply_step_override, callback_failure_result, generate_world_from_prompt, load_example_callback, parse_constraints, run_worldsmith_pipeline, simulate_dsl_pipeline, validate_dsl_callback, ) from dsl.validator import validate_world_spec APP_TITLE = "WorldSmithAI" APP_SUBTITLE = "Production-Level Agent-Based World Simulation Framework" APP_VERSION = "0.1.0" DEFAULT_PROMPT = ( "Create a compact research ecosystem where scientists, engineers, reviewers, " "and funding agents exchange knowledge, compete for attention, collaborate, " "and adapt their goals over time." ) PROMPT_EXAMPLES = [ [ "A medieval civilization with merchants, rulers, artisans, scholars, guards, " "taxes, trade, regulations, and resource pressure." ], [ "A startup economy with founders, investors, customers, operators, competitors, " "market bidding, adoption dynamics, and shifting goals." ], [ "A fantasy world with mages, dragons, healers, merchants, mana resources, " "alliances, communication, negotiation, and magical infrastructure." ], [ "A transport network with hubs, carriers, dispatchers, chargers, passengers, " "queues, routes, deliveries, and congestion." ], [ "A power grid world with generators, storage nodes, consumers, regulators, " "energy resources, subsidies, enforcement, and rebalancing." ], [ "A space colony with colonists, engineers, botanists, navigators, miners, " "oxygen resources, habitat expansion, and survival goals." ], ] CUSTOM_CSS = """ #worldsmith-header { padding: 1.1rem 1.25rem; border-radius: 1rem; background: linear-gradient(135deg, rgba(80, 80, 120, 0.12), rgba(120, 120, 160, 0.08)); border: 1px solid rgba(120, 120, 160, 0.18); } .worldsmith-small { font-size: 0.92rem; opacity: 0.82; } .worldsmith-status textarea { font-family: ui-monospace, SFMono-Regular, Menlo, Monaco, Consolas, "Liberation Mono", monospace; } """ logging.basicConfig( level=os.getenv("WORLDSMITHAI_LOG_LEVEL", "INFO").upper(), format="%(asctime)s | %(levelname)s | %(name)s | %(message)s", ) logger = logging.getLogger(__name__) @dataclass(frozen=True) class AppModelStatus: """Status of the optional model client used by the UI.""" enabled: bool model_id: str | None message: str def to_markdown(self) -> str: """Return display Markdown for model status.""" if self.enabled: return f"**Model client:** configured with `{self.model_id}`." return f"**Model client:** not configured. {self.message}" class OptionalHFInferenceClient: """Small adapter around huggingface_hub.InferenceClient. This adapter tries both chat-completion and text-generation APIs and raises detailed errors instead of hiding provider/model failures. """ def __init__( self, model_id: str, token: str | None = None, provider: str | None = None, ) -> None: """Initialize the optional Hugging Face inference adapter.""" from huggingface_hub import InferenceClient self.model_id = model_id.strip().strip('"').strip("'").strip() self.provider = None if provider is None or not str(provider).strip() else str(provider).strip() api_key = ( token or os.getenv("HF_TOKEN") or os.getenv("HF_HUB_TOKEN") or os.getenv("HUGGINGFACEHUB_API_TOKEN") ) if not api_key: raise RuntimeError( "HF_TOKEN is required for Hugging Face Inference Providers. " "Create a Hugging Face token with Inference Providers permission " "and add it as a Space secret named HF_TOKEN." ) if self.provider: self.client = InferenceClient( model=self.model_id, provider=self.provider, api_key=api_key, ) else: self.client = InferenceClient( model=self.model_id, api_key=api_key, ) def generate( self, prompt: str | None = None, messages: list[dict[str, str]] | None = None, max_new_tokens: int = 2048, temperature: float = 0.2, **_: Any, ) -> str: """Generate text using Hugging Face chat or text-generation APIs.""" errors: list[str] = [] if messages: text = self._try_chat_completion( messages, max_new_tokens=max_new_tokens, temperature=temperature, errors=errors, ) if text: return text text = self._try_openai_compatible_chat( messages, max_new_tokens=max_new_tokens, temperature=temperature, errors=errors, ) if text: return text if prompt: text = self._try_text_generation( prompt, max_new_tokens=max_new_tokens, temperature=temperature, errors=errors, ) if text: return text error_text = "\n".join(errors) if errors else "No inference method was attempted." raise RuntimeError( "InferenceClient did not return generated text.\n\n" f"Model id: {self.model_id}\n\n" f"Attempted methods and errors:\n{error_text}" ) def _try_chat_completion( self, messages: list[dict[str, str]], *, max_new_tokens: int, temperature: float, errors: list[str], ) -> str | None: """Try InferenceClient.chat_completion().""" chat_completion = getattr(self.client, "chat_completion", None) if not callable(chat_completion): errors.append("chat_completion: method not available on InferenceClient") return None try: response = chat_completion( messages=messages, max_tokens=max_new_tokens, temperature=temperature, ) text = _extract_model_text(response) if text: return text errors.append( "chat_completion: response had no extractable text. " f"response_type={response.__class__.__name__}, response={_short_repr(response)}" ) return None except Exception as exc: errors.append(f"chat_completion: {exc.__class__.__name__}: {exc}") return None def _try_openai_compatible_chat( self, messages: list[dict[str, str]], *, max_new_tokens: int, temperature: float, errors: list[str], ) -> str | None: """Try InferenceClient.chat.completions.create().""" chat = getattr(self.client, "chat", None) completions = getattr(chat, "completions", None) create = getattr(completions, "create", None) if not callable(create): errors.append("chat.completions.create: method not available on InferenceClient") return None call_variants = ( { "model": self.model_id, "messages": messages, "max_tokens": max_new_tokens, "temperature": temperature, }, { "messages": messages, "max_tokens": max_new_tokens, "temperature": temperature, }, ) for kwargs in call_variants: try: response = create(**kwargs) text = _extract_model_text(response) if text: return text errors.append( "chat.completions.create: response had no extractable text. " f"kwargs_keys={list(kwargs.keys())}, " f"response_type={response.__class__.__name__}, " f"response={_short_repr(response)}" ) except Exception as exc: errors.append( "chat.completions.create: " f"kwargs_keys={list(kwargs.keys())}, " f"{exc.__class__.__name__}: {exc}" ) return None def _try_text_generation( self, prompt: str, *, max_new_tokens: int, temperature: float, errors: list[str], ) -> str | None: """Try InferenceClient.text_generation().""" text_generation = getattr(self.client, "text_generation", None) if not callable(text_generation): errors.append("text_generation: method not available on InferenceClient") return None call_variants = ( { "prompt": prompt, "max_new_tokens": max_new_tokens, "temperature": temperature, "return_full_text": False, }, { "prompt": prompt, "max_new_tokens": max_new_tokens, "return_full_text": False, }, { "prompt": prompt, "max_new_tokens": max_new_tokens, }, ) for kwargs in call_variants: try: response = text_generation(**kwargs) text = _extract_model_text(response) if text: return text errors.append( "text_generation: response had no extractable text. " f"kwargs_keys={list(kwargs.keys())}, " f"response_type={response.__class__.__name__}, " f"response={_short_repr(response)}" ) except Exception as exc: errors.append( "text_generation: " f"kwargs_keys={list(kwargs.keys())}, " f"{exc.__class__.__name__}: {exc}" ) return None def build_optional_model_client() -> tuple[Any | None, AppModelStatus]: """Build the optional model client from environment variables.""" raw_model_id = ( os.getenv("WORLDSMITHAI_MODEL_ID") or os.getenv("WORLD_SMITH_MODEL_ID") or os.getenv("HF_MODEL_ID") ) token = ( os.getenv("HF_TOKEN") or os.getenv("HF_HUB_TOKEN") or os.getenv("HUGGINGFACEHUB_API_TOKEN") ) provider = ( os.getenv("WORLDSMITHAI_PROVIDER") or os.getenv("HF_INFERENCE_PROVIDER") or None ) if not raw_model_id: return None, AppModelStatus( enabled=False, model_id=None, message="Set WORLDSMITHAI_MODEL_ID to enable model-generated DSL.", ) model_id = raw_model_id.strip().strip('"').strip("'").strip() logger.info( "WorldSmithAI model configuration: raw_model_id=%r sanitized_model_id=%r", raw_model_id, model_id, ) if not model_id: return None, AppModelStatus( enabled=False, model_id=None, message="WORLDSMITHAI_MODEL_ID was set but empty after sanitization.", ) try: client = OptionalHFInferenceClient( model_id=model_id, token=token, provider=provider, ) return client, AppModelStatus( enabled=True, model_id=model_id, message="Model client configured.", ) except Exception as exc: logger.warning("Optional model client could not be configured: %s", exc) return None, AppModelStatus( enabled=False, model_id=model_id, message=( f"WORLDSMITHAI_MODEL_ID={model_id} was found, but the client " f"could not be initialized: {exc}. The app will use deterministic fallback." ), ) MODEL_CLIENT, MODEL_STATUS = build_optional_model_client() def model_health_check_ui() -> str: """Return a quick model-client health check for the Gradio UI.""" if MODEL_CLIENT is None: return ( "Model client is NOT configured.\n\n" f"Runtime status: {MODEL_STATUS.message}\n\n" "Check that WORLDSMITHAI_MODEL_ID is set in Space settings." ) provider_label = ( getattr(MODEL_CLIENT, "provider", None) or os.getenv("WORLDSMITHAI_PROVIDER") or os.getenv("HF_INFERENCE_PROVIDER") or "auto" ) try: response = MODEL_CLIENT.generate( prompt='Return exactly this JSON object and nothing else: {"ok": true}', messages=[ { "role": "system", "content": "You return compact JSON only.", }, { "role": "user", "content": 'Return exactly this JSON object and nothing else: {"ok": true}', }, ], max_new_tokens=64, temperature=0.0, ) return ( "Model client is configured and returned text.\n\n" f"Model id: {MODEL_STATUS.model_id}\n\n" f"Provider: {provider_label}\n\n" f"Response:\n{response}" ) except Exception as exc: logger.exception("Model health check failed") return ( "Model client is configured but inference failed.\n\n" f"Model id: {MODEL_STATUS.model_id}\n\n" f"Provider: {provider_label}\n\n" f"Error type: {exc.__class__.__name__}\n" f"Error:\n{exc}" ) def app_callback_config( *, steps: int | float | str | None, animation_format: str, ) -> CallbackConfig: """Create callback configuration from UI values and environment defaults.""" default_steps = _env_int("WORLDSMITHAI_DEFAULT_STEPS", 60) max_frames = _env_int("WORLDSMITHAI_MAX_FRAMES", 80) output_dir = os.getenv("WORLDSMITHAI_OUTPUT_DIR") or None return CallbackConfig( output_dir=output_dir, steps=_coerce_optional_int(steps) if steps is not None else default_steps, max_animation_frames=max(1, max_frames), animation_format=animation_format, ) def run_from_prompt_ui( prompt: str, steps: int | float, constraints: str, animation_format: str, use_configured_model: bool, ) -> tuple[Any, ...]: """Run the full prompt-to-simulation pipeline for the main UI tab.""" client = MODEL_CLIENT if use_configured_model else None config = app_callback_config(steps=steps, animation_format=animation_format) try: result = run_worldsmith_pipeline( prompt=prompt, constraints=parse_constraints(constraints), client=client, config=config, ) return result_to_ui_tuple(result) except Exception as exc: logger.exception("Prompt simulation failed") return failure_to_ui_tuple(exc) def generate_dsl_only_ui( prompt: str, constraints: str, use_configured_model: bool, steps: int | float, ) -> tuple[str, str, str]: """Generate DSL only, without running simulation.""" client = MODEL_CLIENT if use_configured_model else None config = app_callback_config(steps=steps, animation_format="gif") try: generation = generate_world_from_prompt( prompt=prompt, constraints=parse_constraints(constraints), client=client, config=config, ) spec = apply_step_override(generation.spec, _coerce_optional_int(steps)) report = validate_world_spec( spec, config=config.resolved_validation_config(), ) status = ( f"Generated `{spec.id}` with {len(spec.agents)} agent(s), " f"{len(spec.resources)} resource(s), and " f"{len(spec.behavior_names)} behavior reference(s). " f"Validation: {len(report.errors)} error(s), {len(report.warnings)} warning(s)." ) return ( spec.to_json_string(indent=2, exclude_none=True), _safe_json_dumps(report.to_dict()), status, ) except Exception as exc: logger.exception("DSL generation failed") return "", "{}", f"DSL generation failed: {exc}" def simulate_dsl_ui( dsl_json: str, steps: int | float, animation_format: str, use_configured_model_for_narration: bool, ) -> tuple[Any, ...]: """Run simulation from user-supplied DSL JSON.""" client = MODEL_CLIENT if use_configured_model_for_narration else None config = app_callback_config(steps=steps, animation_format=animation_format) try: result = simulate_dsl_pipeline( dsl_json=dsl_json, client=client, config=config, ) return result_to_ui_tuple(result) except Exception as exc: logger.exception("DSL simulation failed") return failure_to_ui_tuple(exc, world_spec_json=dsl_json) def validate_dsl_ui(dsl_json: str) -> tuple[str, str]: """Validate DSL JSON in the validation tab.""" return validate_dsl_callback(dsl_json) def load_example_ui(example_name: str | None) -> tuple[str, str]: """Load an example JSON file into the DSL editor.""" if not example_name: return "", "Choose an example first." examples = discover_examples() path = examples.get(example_name) if path is None: return "", f"Unknown example: {example_name}" return load_example_callback(path) def clear_main_outputs() -> tuple[Any, ...]: """Clear main-tab outputs.""" return ( "", "", gr.update(value=None, visible=False), gr.update(value=None, visible=False), None, None, None, "", "", "", ) def result_to_ui_tuple(result: CallbackResult) -> tuple[Any, ...]: """Map CallbackResult to the app output component order.""" animation_image, animation_video = animation_path_updates(result.animation_path) return ( result.world_spec_json, result.validation_json, animation_image, animation_video, result.population_chart_path, result.resource_chart_path, result.final_image_path, _markdown_narrative(result.narrative), result.metrics_json, result.status_message, ) def failure_to_ui_tuple( error: BaseException, *, world_spec_json: str = "", ) -> tuple[Any, ...]: """Map an exception to a safe UI output tuple.""" result = callback_failure_result(error, world_spec_json=world_spec_json) details = { "success": False, "error_type": error.__class__.__name__, "error": str(error), "traceback": traceback.format_exc(), "callback_result": result.to_dict(), } return ( result.world_spec_json, result.validation_json, gr.update(value=None, visible=False), gr.update(value=None, visible=False), None, None, None, _markdown_narrative(result.narrative), _safe_json_dumps(details), result.status_message, ) def animation_path_updates(path: str | None) -> tuple[Any, Any]: """Return Gradio updates for GIF/image and MP4/video animation components.""" if not path: return gr.update(value=None, visible=False), gr.update(value=None, visible=False) suffix = Path(path).suffix.lower() if suffix in {".mp4", ".webm", ".mov", ".m4v"}: return gr.update(value=None, visible=False), gr.update(value=path, visible=True) return gr.update(value=path, visible=True), gr.update(value=None, visible=False) def discover_examples() -> dict[str, Path]: """Discover JSON example files in the root-level examples directory.""" examples_dir = Path(__file__).resolve().parent / "examples" if not examples_dir.exists(): return {} return { path.stem.replace("_", " ").title(): path for path in sorted(examples_dir.glob("*.json")) } def build_demo() -> gr.Blocks: """Build and return the Gradio Blocks demo.""" examples = discover_examples() example_names = list(examples.keys()) default_steps = _env_int("WORLDSMITHAI_DEFAULT_STEPS", 60) with gr.Blocks( title=f"{APP_TITLE} — Agent-Based World Simulation", css=CUSTOM_CSS, analytics_enabled=False, ) as demo: gr.Markdown( f"""