fhaviary 0.1.0__py3-none-any.whl

aviary/api.py

@@ -0,0 +1,195 @@
+import os
+import secrets
+import uuid
+from typing import TYPE_CHECKING, Annotated, Any
+
+import httpx
+
+from aviary.db import (
+    EnvironmentBase,
+    EnvironmentDB,
+    EnvironmentDBBackend,
+    EnvironmentDBSchema,
+    FrameDB,
+    FrameDBSchema,
+)
+from aviary.render import Frame
+
+if TYPE_CHECKING:
+    from fastapi import FastAPI
+
+
+class EnvDBClient:
+    """Interact with an Environment DB via REST API."""
+
+    def __init__(
+        self,
+        server_url: str,
+        request_headers: httpx._types.HeaderTypes | None = None,
+        request_timeout: float | None = None,
+    ):
+        self._request_url = server_url
+        self._request_headers = request_headers
+        self._request_timeout = request_timeout
+
+    async def write_environment_instance(self, name: str) -> uuid.UUID:
+        async with httpx.AsyncClient() as client:
+            response = await client.post(
+                f"{self._request_url}/environment_instance",
+                params={
+                    "env_name": name,
+                },
+                headers=self._request_headers,
+                timeout=self._request_timeout,
+            )
+            response.raise_for_status()
+            return response.json()
+
+    async def write_environment_frame(
+        self,
+        environment_id: uuid.UUID,
+        frame: Frame,
+    ) -> uuid.UUID:
+        async with httpx.AsyncClient() as client:
+            frame_data = frame.model_dump()
+            response = await client.post(
+                f"{self._request_url}/environment_frame",
+                params={"environment_id": str(environment_id)},
+                json={
+                    "state": frame_data.get("state"),
+                    "metadata": frame_data.get("info"),
+                },
+                headers=self._request_headers,
+                timeout=self._request_timeout,
+            )
+            response.raise_for_status()
+            return response.json()
+
+    async def get_environment_instances(
+        self,
+        name: str | None = None,
+        environment_id: uuid.UUID | None = None,
+    ) -> list[EnvironmentDB]:
+        async with httpx.AsyncClient() as client:
+            params = (
+                {"env_name": name} if name else {"environment_id": str(environment_id)}
+            )
+            response = await client.get(
+                f"{self._request_url}/environment_instance",
+                params=params,
+                headers=self._request_headers,
+                timeout=self._request_timeout,
+            )
+            response.raise_for_status()
+            return [EnvironmentDB(**obj) for obj in response.json()]
+
+    async def get_environment_frames(
+        self,
+        environment_id: uuid.UUID | None = None,
+        frame_id: uuid.UUID | None = None,
+    ) -> list[FrameDB]:
+        async with httpx.AsyncClient() as client:
+            params = (
+                {"environment_id": str(environment_id)}
+                if environment_id
+                else {"frame_id": str(frame_id)}
+            )
+            response = await client.get(
+                f"{self._request_url}/environment_frame",
+                params=params,
+                headers=self._request_headers,
+                timeout=self._request_timeout,
+            )
+            response.raise_for_status()
+            return [FrameDB(**obj) for obj in response.json()]
+
+
+def make_environment_db_server(render_docs: bool = False) -> "FastAPI":  # noqa: C901
+    """Make a FastAPI app, an interface for the Environment DB."""
+    try:
+        from fastapi import Depends, FastAPI, HTTPException, status
+        from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
+    except ModuleNotFoundError as exc:
+        raise ImportError(
+            "Environment DB server requires the 'server' extra for 'fastapi'. Please:"
+            " `pip install aviary[server]`."
+        ) from exc
+
+    backend = EnvironmentDBBackend()
+
+    async def ensure_session():
+        if not backend.Session:
+            await backend.populate_session(
+                base=EnvironmentBase, uri=os.environ["ENV_DB_URI"]
+            )
+        yield
+
+    asgi_app = FastAPI(
+        title="Aviary Environment DB API",
+        description="CRUD operations for the Aviary Environment DB",
+        # Only render Swagger docs if local since we don't have a login here
+        docs_url="/docs" if render_docs else None,
+        redoc_url="/redoc" if render_docs else None,
+    )
+    auth_scheme = HTTPBearer()
+
+    async def validate_token(
+        token: Annotated[HTTPAuthorizationCredentials, Depends(auth_scheme)],
+    ) -> HTTPAuthorizationCredentials:
+        # NOTE: don't use os.environ.get() to avoid possible empty string matches, and
+        # to have clearer server failures if the AUTH_TOKEN env var isn't present
+        if not secrets.compare_digest(token.credentials, os.environ["AUTH_TOKEN"]):
+            raise HTTPException(
+                status_code=status.HTTP_401_UNAUTHORIZED,
+                detail="Incorrect bearer token",
+                headers={"WWW-Authenticate": "Bearer"},
+            )
+        return token
+
+    @asgi_app.post("/environment_instance")
+    async def add_environment_instance(
+        env_name: str,
+        _: Annotated[HTTPAuthorizationCredentials, Depends(validate_token)],
+        _db: Annotated[None, Depends(ensure_session)],
+    ) -> uuid.UUID:
+        return await backend.add_environment_instance(name=env_name)
+
+    @asgi_app.post("/environment_frame")
+    async def add_environment_frame(
+        environment_id: uuid.UUID,
+        state: dict[str, Any],
+        _: Annotated[HTTPAuthorizationCredentials, Depends(validate_token)],
+        _db: Annotated[None, Depends(ensure_session)],
+        metadata: dict[str, Any] | None = None,
+    ) -> uuid.UUID:
+        return await backend.add_environment_frame(
+            environment_id=environment_id, state=state, metadata=metadata
+        )
+
+    @asgi_app.get("/environment_instance")
+    async def get_environment_instances(
+        _: Annotated[HTTPAuthorizationCredentials, Depends(validate_token)],
+        _db: Annotated[None, Depends(ensure_session)],
+        env_name: str | None = None,
+        environment_id: uuid.UUID | None = None,
+    ) -> list[EnvironmentDBSchema]:
+        if not (env_name or environment_id):
+            raise HTTPException(400, "Must provide either a name or environment_id")
+        return await backend.get_environment_instances(
+            name=env_name, environment_id=environment_id
+        )
+
+    @asgi_app.get("/environment_frame")
+    async def get_environment_frames(
+        _: Annotated[HTTPAuthorizationCredentials, Depends(validate_token)],
+        _db: Annotated[None, Depends(ensure_session)],
+        environment_id: uuid.UUID | None = None,
+        frame_id: uuid.UUID | None = None,
+    ) -> list[FrameDBSchema]:
+        if not (frame_id or environment_id):
+            raise HTTPException(400, "Must provide either a frame_id or environment_id")
+        return await backend.get_environment_frames(
+            environment_id=environment_id, frame_id=frame_id
+        )
+
+    return asgi_app

aviary/db.py

@@ -0,0 +1,155 @@
+from __future__ import annotations
+
+import logging
+import uuid
+from datetime import datetime
+from typing import Any, ClassVar
+
+from pydantic import BaseModel, ConfigDict, Field, JsonValue
+from sqlalchemy import Column, DateTime, ForeignKey, select
+from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column
+from sqlalchemy.orm.decl_api import _TypeAnnotationMapType
+from sqlalchemy.sql import func
+from sqlalchemy.types import JSON
+
+from aviary.utils import DBBackend
+from aviary.version import __version__
+
+logger = logging.getLogger(__name__)
+
+
+class EnvironmentBase(DeclarativeBase):
+    # Allowing dict to correspond with arbitrary JSON
+    type_annotation_map: ClassVar[_TypeAnnotationMapType] = {JsonValue: JSON}
+
+
+class EnvironmentDBBackend(DBBackend):
+    @classmethod
+    async def get_environment_instances(
+        cls, name: str | None = None, environment_id: uuid.UUID | None = None
+    ) -> list[EnvironmentDBSchema]:
+        """Given a name, pull associated environment instance data.
+
+        Args:
+            name: The name of the environment this instance is associated with
+            environment_id: The UUID of a specific environment instances
+
+        Returns:
+           A list of environment UUIDs.
+        """
+        async with cls.begin_session() as session:
+            if environment_id is not None:
+                query = select(EnvironmentDB).where(EnvironmentDB.id == environment_id)
+            elif name is not None:
+                query = select(EnvironmentDB).where(EnvironmentDB.name == name)
+            else:
+                raise ValueError("Must provide either a name or environment_id")
+
+            results = (await session.execute(query)).scalars().all()
+            return [EnvironmentDBSchema.model_validate(result) for result in results]
+
+    @classmethod
+    async def get_environment_frames(
+        cls, environment_id: uuid.UUID | None, frame_id: uuid.UUID | None
+    ) -> list[FrameDBSchema]:
+        """Given a environment_id and frame_id data, get environment frame data.
+
+        Args:
+            environment_id: The environment instance we are looking for
+            frame_id: The specific frame we are looking for
+
+        Returns:
+           A list of FrameDB.
+        """
+        async with cls.begin_session() as session:
+            if frame_id is not None:
+                query = select(FrameDB).where(FrameDB.id == frame_id)
+            elif environment_id is not None:
+                query = select(FrameDB).where(FrameDB.environment_id == environment_id)
+            else:
+                raise ValueError("Must provide either a environment_id or frame_id")
+            results = (await session.execute(query)).scalars().all()
+            return [FrameDBSchema.model_validate(result) for result in results]
+
+    @classmethod
+    async def add_environment_frame(
+        cls,
+        environment_id: uuid.UUID,
+        state: dict[str, Any],
+        metadata: dict[str, Any] | None = None,
+    ) -> uuid.UUID:
+        frame_id = uuid.uuid4()
+        async with cls.begin_session() as session:
+            session.add(
+                FrameDB(
+                    id=frame_id,
+                    environment_id=environment_id,
+                    state=state,
+                    supplemental_data=metadata,
+                )
+            )
+        return frame_id
+
+    @classmethod
+    async def add_environment_instance(
+        cls,
+        name: str,
+    ) -> uuid.UUID:
+        environment_id = uuid.uuid4()
+        async with cls.begin_session() as session:
+            session.add(
+                EnvironmentDB(id=environment_id, name=name, aviary_version=__version__)
+            )
+        return environment_id
+
+
+class EnvironmentDB(EnvironmentBase):
+    """Rows here correspond to a single environment trajectory, a single named set of frames."""
+
+    __tablename__ = "environments"
+
+    id: Mapped[uuid.UUID] = mapped_column(primary_key=True)
+    created_at: Mapped[datetime] = mapped_column(default=func.now())
+    name: Mapped[str] = mapped_column(index=True)
+    aviary_version: Mapped[str] = mapped_column(default=__version__)
+
+
+class FrameDB(EnvironmentBase):
+    """Rows are a snapshot of an environment after each step.
+
+    Note: corresponds to a frame in the renderer object
+
+    """
+
+    __tablename__ = "frames"
+
+    id: Mapped[uuid.UUID] = mapped_column(primary_key=True)
+    environment_id: Mapped[uuid.UUID] = mapped_column(
+        ForeignKey(EnvironmentDB.id), index=True
+    )
+    state: Mapped[JsonValue]  # Just for human readability
+    supplemental_data: Mapped[JsonValue]
+    created_at = Column(DateTime(timezone=True), server_default=func.now())
+
+
+class EnvironmentDBSchema(BaseModel):
+    """Sister model for EnvironmentDB."""
+
+    id: uuid.UUID
+    created_at: datetime = Field(default_factory=datetime.now)
+    name: str
+    aviary_version: str = __version__
+
+    model_config = ConfigDict(from_attributes=True)
+
+
+class FrameDBSchema(BaseModel):
+    """Sister model for FrameDB."""
+
+    id: uuid.UUID
+    environment_id: uuid.UUID
+    state: JsonValue
+    supplemental_data: JsonValue
+    created_at: datetime = Field(default_factory=datetime.now)
+
+    model_config = ConfigDict(from_attributes=True)