arc_agi_rs/

python.rs

// Copyright 2026 Mahmoud Harmouch.
//
// Licensed under the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! # Python Bindings
//!
//! Exposes the `arc-agi-rs` library to Python via [`pyo3`].
//! Every type and function is gated behind the `python` cargo feature.
//!
//! The bindings provide **synchronous** wrappers around the async Rust API by
//! driving a temporary, single-threaded [`tokio`] runtime inside each call,
//! making the API feel native and straightforward for Python callers.
//!
//! # See Also
//!
//! - [ARC-AGI-3 Reference](https://arcprize.org/arc-agi/3)

use std::fmt::Display;
use std::future::Future;

use pyo3::exceptions::PyRuntimeError;
use pyo3::prelude::*;
use pyo3::types::{PyAny, PyDict, PyList};
use serde_json::{Map, Value};
use tokio::runtime::Builder as RuntimeBuilder;

use crate::client::Client;
use crate::params::{MakeParams, ScorecardParams, StepParams};

/// Convert a Python dict/object to a [`Value`] without the `pythonize` crate.
///
/// Only `dict`, `list`, `str`, `int`, `float`, `bool`, and `None` are
/// supported - sufficient for the `data` and `reasoning` step payloads.
fn pyany_to_json(obj: &Bound<'_, PyAny>) -> PyResult<Value> {
    if obj.is_none() {
        return Ok(Value::Null);
    }
    if let Ok(b) = obj.extract::<bool>() {
        return Ok(Value::Bool(b));
    }
    if let Ok(i) = obj.extract::<i64>() {
        return Ok(Value::from(i));
    }
    if let Ok(f) = obj.extract::<f64>() {
        return Ok(Value::from(f));
    }
    if let Ok(s) = obj.extract::<String>() {
        return Ok(Value::String(s));
    }
    if let Ok(list) = obj.extract::<Bound<'_, PyList>>() {
        let arr: PyResult<Vec<_>> = list.iter().map(|v| pyany_to_json(&v)).collect();
        return Ok(Value::Array(arr?));
    }
    if let Ok(dict) = obj.extract::<Bound<'_, PyDict>>() {
        let mut map = Map::new();
        for (k, v) in dict.iter() {
            let key: String = k.extract()?;
            map.insert(key, pyany_to_json(&v)?);
        }
        return Ok(Value::Object(map));
    }
    Err(PyRuntimeError::new_err(format!(
        "cannot convert Python value of type {} to JSON",
        obj.get_type().name()?
    )))
}

/// Converts a Python `dict` to a [`Value`].
fn pydict_to_json(d: &Bound<'_, PyDict>) -> PyResult<Value> {
    pyany_to_json(d.as_any())
}

/// Drives an async future on a fresh single-threaded Tokio runtime.
///
/// This function is adapted from the `duckduckgo` crate's Python binding helper:
/// <https://github.com/wiseaidotdev/duckduckgo/blob/main/src/python.rs>
fn block_on<F, T, E>(future: F) -> PyResult<T>
where
    F: Future<Output = Result<T, E>>,
    E: Display,
{
    RuntimeBuilder::new_current_thread()
        .enable_all()
        .build()
        .map_err(|e| PyRuntimeError::new_err(e.to_string()))?
        .block_on(future)
        .map_err(|e| PyRuntimeError::new_err(e.to_string()))
}

/// Metadata for a single ARC-AGI-3 game environment.
///
/// All fields are read-only once the object is constructed by a
/// :meth:`ArcAgiClient.list_environments` or
/// :meth:`ArcAgiClient.get_environment` call.
///
/// # See Also
///
/// - [ARC-AGI-3 Reference](https://arcprize.org/arc-agi/3)
#[pyclass(name = "EnvironmentInfo", frozen, skip_from_py_object)]
#[derive(Debug, Clone)]
pub struct PyEnvironmentInfo {
    /// Unique game identifier (e.g. ``"ls20"``).
    #[pyo3(get)]
    pub game_id: String,
    /// Human-readable title.
    #[pyo3(get)]
    pub title: Option<String>,
    /// Default frames-per-second.
    #[pyo3(get)]
    pub default_fps: Option<u32>,
    /// Classification tags.
    #[pyo3(get)]
    pub tags: Option<Vec<String>>,
}

#[pymethods]
impl PyEnvironmentInfo {
    pub fn __repr__(&self) -> String {
        format!(
            "EnvironmentInfo(game_id={:?}, title={:?})",
            self.game_id, self.title
        )
    }
}

/// The current state of a game run returned by reset and step calls.
///
/// All fields are read-only once the object is constructed by a
/// :meth:`ArcAgiClient.reset` or :meth:`ArcAgiClient.step` call.
///
/// # See Also
///
/// - [ARC-AGI-3 Reference](https://arcprize.org/arc-agi/3)
#[pyclass(name = "FrameData", frozen, skip_from_py_object)]
#[derive(Debug, Clone)]
pub struct PyFrameData {
    /// Game identifier.
    #[pyo3(get)]
    pub game_id: String,
    /// Unique run identifier assigned by the server.
    #[pyo3(get)]
    pub guid: Option<String>,
    /// Current lifecycle state (e.g. ``"NOT_FINISHED"``, ``"WIN"``).
    #[pyo3(get)]
    pub state: String,
    /// Number of levels completed in this run.
    #[pyo3(get)]
    pub levels_completed: u32,
    /// Total levels that must be completed to win.
    #[pyo3(get)]
    pub win_levels: u32,
    /// Action IDs the agent may send on the next step.
    #[pyo3(get)]
    pub available_actions: Vec<u32>,
    /// Whether this response corresponds to a full game reset.
    #[pyo3(get)]
    pub full_reset: bool,
}

#[pymethods]
impl PyFrameData {
    pub fn __repr__(&self) -> String {
        format!(
            "FrameData(game_id={:?}, guid={:?}, state={:?})",
            self.game_id, self.guid, self.state
        )
    }
}

/// A per-game score entry inside an :class:`EnvironmentScorecard`.
///
/// All fields are read-only.
#[pyclass(name = "EnvironmentScore", frozen, skip_from_py_object)]
#[derive(Debug, Clone)]
pub struct PyEnvironmentScore {
    /// Game identifier.
    #[pyo3(get)]
    pub id: Option<String>,
    /// Aggregate score (0.0–115.0).
    #[pyo3(get)]
    pub score: f64,
    /// Number of levels completed.
    #[pyo3(get)]
    pub levels_completed: u32,
    /// Total actions taken.
    #[pyo3(get)]
    pub actions: u32,
    /// Whether the environment was fully completed.
    #[pyo3(get)]
    pub completed: Option<bool>,
}

#[pymethods]
impl PyEnvironmentScore {
    pub fn __repr__(&self) -> String {
        format!(
            "EnvironmentScore(id={:?}, score={:.2})",
            self.id, self.score
        )
    }
}

/// The full scored scorecard returned by scorecard endpoints.
///
/// All fields are read-only once the object is constructed by a
/// :meth:`ArcAgiClient.get_scorecard` or :meth:`ArcAgiClient.close_scorecard`
/// call.
///
/// # See Also
///
/// - [ARC-AGI-3 Reference](https://arcprize.org/arc-agi/3)
#[pyclass(name = "EnvironmentScorecard", frozen, skip_from_py_object)]
#[derive(Debug, Clone)]
pub struct PyEnvironmentScorecard {
    /// Unique scorecard identifier.
    #[pyo3(get)]
    pub card_id: String,
    /// Aggregate score across all environments (0.0–115.0).
    #[pyo3(get)]
    pub score: f64,
    /// Whether the scorecard was created in competition mode.
    #[pyo3(get)]
    pub competition_mode: Option<bool>,
    /// Total environments completed.
    #[pyo3(get)]
    pub total_environments_completed: Option<u32>,
    /// Total environments.
    #[pyo3(get)]
    pub total_environments: Option<u32>,
    /// Total levels completed.
    #[pyo3(get)]
    pub total_levels_completed: Option<u32>,
    /// Total levels.
    #[pyo3(get)]
    pub total_levels: Option<u32>,
    /// Total actions taken.
    #[pyo3(get)]
    pub total_actions: Option<u32>,
}

#[pymethods]
impl PyEnvironmentScorecard {
    pub fn __repr__(&self) -> String {
        format!(
            "EnvironmentScorecard(card_id={:?}, score={:.2})",
            self.card_id, self.score
        )
    }
}

/// An HTTP client for interacting with the ARC-AGI-3 REST API.
///
/// Construct with ``ArcAgiClient()`` for a zero-configuration default that
/// reads credentials from the ``ARC_API_KEY`` and ``ARC_BASE_URL`` environment
/// variables, or supply keyword arguments to configure them explicitly. All
/// network methods are **synchronous** from Python's perspective; they drive
/// an internal single-threaded Tokio runtime for each call.
///
/// # See Also
///
/// - [ARC-AGI-3 Reference](https://arcprize.org/arc-agi/3)
#[pyclass(name = "ArcAgiClient")]
pub struct PyArcAgiClient {
    inner: Client,
}

#[pymethods]
impl PyArcAgiClient {
    /// Create a new ``ArcAgiClient``.
    ///
    /// Args:
    ///     api_key:      Optional API key string. Falls back to ``ARC_API_KEY``
    ///                   environment variable and then to an empty string.
    ///     base_url:     Optional server base URL. Falls back to ``ARC_BASE_URL``
    ///                   environment variable and then to
    ///                   ``"https://three.arcprize.org"``.
    ///     cookie_store: Enable cookie persistence across requests
    ///                   (default ``False``).
    ///     proxy:        Optional proxy URL, e.g. ``"socks5://127.0.0.1:9050"``.
    ///
    /// Raises:
    ///     RuntimeError: If the proxy URL is invalid or the HTTP client cannot
    ///                   be constructed.
    #[new]
    #[pyo3(signature = (api_key=None, base_url=None, cookie_store=false, proxy=None))]
    pub fn new(
        api_key: Option<String>,
        base_url: Option<String>,
        cookie_store: bool,
        proxy: Option<String>,
    ) -> PyResult<Self> {
        let mut builder = Client::builder();
        if let Some(key) = api_key {
            builder = builder.api_key(key);
        }
        if let Some(url) = base_url {
            builder = builder.base_url(url);
        }
        if cookie_store {
            builder = builder.cookie_store(true);
        }
        if let Some(proxy_url) = proxy {
            builder = builder.proxy(proxy_url);
        }
        let inner = builder
            .build()
            .map_err(|e| PyRuntimeError::new_err(e.to_string()))?;
        Ok(Self { inner })
    }

    /// Retrieve an anonymous API key from the server.
    ///
    /// Returns:
    ///     The anonymous API key string.
    ///
    /// Raises:
    ///     RuntimeError: On network or HTTP failure.
    pub fn get_anonymous_key(&self) -> PyResult<String> {
        block_on(self.inner.get_anonymous_key())
    }

    /// Return the list of all available game environments.
    ///
    /// Returns:
    ///     A list of :class:`EnvironmentInfo` objects.
    ///
    /// Raises:
    ///     RuntimeError: On network or HTTP failure.
    pub fn list_environments(&self) -> PyResult<Vec<PyEnvironmentInfo>> {
        let envs = block_on(self.inner.list_environments())?;
        Ok(envs
            .into_iter()
            .map(|e| PyEnvironmentInfo {
                game_id: e.game_id,
                title: e.title,
                default_fps: e.default_fps,
                tags: e.tags,
            })
            .collect())
    }

    /// Return metadata for a single game environment.
    ///
    /// Args:
    ///     game_id: The game identifier (e.g. ``"ls20"``).
    ///
    /// Returns:
    ///     An :class:`EnvironmentInfo` object.
    ///
    /// Raises:
    ///     RuntimeError: If the game is not found or on network failure.
    pub fn get_environment(&self, game_id: String) -> PyResult<PyEnvironmentInfo> {
        let info = block_on(self.inner.get_environment(&game_id))?;
        Ok(PyEnvironmentInfo {
            game_id: info.game_id,
            title: info.title,
            default_fps: info.default_fps,
            tags: info.tags,
        })
    }

    /// Create a new scorecard and return its ID.
    ///
    /// Args:
    ///     source_url:       Optional URL linking to the agent being evaluated.
    ///     tags:             Optional list of classification tag strings.
    ///     competition_mode: When ``True``, enables one-way competition semantics.
    ///
    /// Returns:
    ///     The ``card_id`` string of the newly created scorecard.
    ///
    /// Raises:
    ///     RuntimeError: On network or HTTP failure.
    #[pyo3(signature = (source_url=None, tags=None, competition_mode=None))]
    pub fn open_scorecard(
        &self,
        source_url: Option<String>,
        tags: Option<Vec<String>>,
        competition_mode: Option<bool>,
    ) -> PyResult<String> {
        let mut params = ScorecardParams::new();
        if let Some(url) = source_url {
            params = params.source_url(url);
        }
        if let Some(t) = tags {
            params = params.tags(t);
        }
        if let Some(cm) = competition_mode {
            params = params.competition_mode(cm);
        }
        block_on(self.inner.open_scorecard(Some(params)))
    }

    /// Retrieve an existing scorecard by its ID.
    ///
    /// Args:
    ///     card_id: The scorecard identifier.
    ///
    /// Returns:
    ///     An :class:`EnvironmentScorecard` object.
    ///
    /// Raises:
    ///     RuntimeError: If the scorecard is not found or on network failure.
    pub fn get_scorecard(&self, card_id: String) -> PyResult<PyEnvironmentScorecard> {
        let card = block_on(self.inner.get_scorecard(&card_id))?;
        Ok(PyEnvironmentScorecard {
            card_id: card.card_id,
            score: card.score,
            competition_mode: card.competition_mode,
            total_environments_completed: card.total_environments_completed,
            total_environments: card.total_environments,
            total_levels_completed: card.total_levels_completed,
            total_levels: card.total_levels,
            total_actions: card.total_actions,
        })
    }

    /// Close and finalise a scorecard.
    ///
    /// Args:
    ///     card_id: The scorecard identifier.
    ///
    /// Returns:
    ///     An :class:`EnvironmentScorecard` with the final scores.
    ///
    /// Raises:
    ///     RuntimeError: If the scorecard is not found or on network failure.
    pub fn close_scorecard(&self, card_id: String) -> PyResult<PyEnvironmentScorecard> {
        let card = block_on(self.inner.close_scorecard(&card_id))?;
        Ok(PyEnvironmentScorecard {
            card_id: card.card_id,
            score: card.score,
            competition_mode: card.competition_mode,
            total_environments_completed: card.total_environments_completed,
            total_environments: card.total_environments,
            total_levels_completed: card.total_levels_completed,
            total_levels: card.total_levels,
            total_actions: card.total_actions,
        })
    }

    /// Reset (or start) a game environment.
    ///
    /// Args:
    ///     game_id:      The game identifier.
    ///     scorecard_id: The scorecard to record this run under.
    ///     guid:         Optional existing run GUID to re-use.
    ///     seed:         Random seed for reproducible level ordering (default ``0``).
    ///
    /// Returns:
    ///     A :class:`FrameData` object with the initial game state.
    ///
    /// Raises:
    ///     RuntimeError: On network or HTTP failure.
    #[pyo3(signature = (game_id, scorecard_id, guid=None, seed=0))]
    pub fn reset(
        &self,
        game_id: String,
        scorecard_id: String,
        guid: Option<String>,
        seed: u32,
    ) -> PyResult<PyFrameData> {
        let mut params = MakeParams::new(&game_id, &scorecard_id).seed(seed);
        if let Some(g) = guid {
            params = params.guid(g);
        }
        let frame = block_on(self.inner.reset(params))?;
        Ok(PyFrameData {
            game_id: frame.game_id,
            guid: frame.guid,
            state: frame.state.as_str().to_string(),
            levels_completed: frame.levels_completed,
            win_levels: frame.win_levels,
            available_actions: frame.available_actions,
            full_reset: frame.full_reset,
        })
    }

    /// Send one game action and receive the resulting frame.
    ///
    /// Args:
    ///     game_id:      The game identifier.
    ///     scorecard_id: The scorecard identifier.
    ///     guid:         The run GUID from a prior :meth:`reset` call.
    ///     action_id:    Numeric action ID (0 = RESET).
    ///     data:         Optional dict of action data (e.g. ``{"x": 3, "y": 4}``).
    ///     reasoning:    Optional dict of freeform reasoning.
    ///
    /// Returns:
    ///     A :class:`FrameData` object with the updated game state.
    ///
    /// Raises:
    ///     RuntimeError: On network or HTTP failure.
    #[pyo3(signature = (game_id, scorecard_id, guid, action_id, data=None, reasoning=None))]
    pub fn step(
        &self,
        game_id: String,
        scorecard_id: String,
        guid: String,
        action_id: u32,
        data: Option<&pyo3::Bound<'_, pyo3::types::PyDict>>,
        reasoning: Option<&pyo3::Bound<'_, pyo3::types::PyDict>>,
    ) -> PyResult<PyFrameData> {
        let data_json = data.map(pydict_to_json).transpose()?;
        let reasoning_json = reasoning.map(pydict_to_json).transpose()?;

        let mut params = StepParams::new(&game_id, &scorecard_id, &guid, action_id);
        if let Some(d) = data_json {
            params = params.data(d);
        }
        if let Some(r) = reasoning_json {
            params = params.reasoning(r);
        }
        let frame = block_on(self.inner.step(params))?;
        Ok(PyFrameData {
            game_id: frame.game_id,
            guid: frame.guid,
            state: frame.state.as_str().to_string(),
            levels_completed: frame.levels_completed,
            win_levels: frame.win_levels,
            available_actions: frame.available_actions,
            full_reset: frame.full_reset,
        })
    }

    pub fn __repr__(&self) -> String {
        format!("ArcAgiClient(base_url={:?})", self.inner.base_url())
    }
}

/// Register all Python-exposed types and functions into the ``arc_agi_rs`` module.
pub fn register_python_module(_py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> {
    m.add_class::<PyEnvironmentInfo>()?;
    m.add_class::<PyFrameData>()?;
    m.add_class::<PyEnvironmentScore>()?;
    m.add_class::<PyEnvironmentScorecard>()?;
    m.add_class::<PyArcAgiClient>()?;
    Ok(())
}

// Copyright 2026 Mahmoud Harmouch.
//
// Licensed under the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.