arize-phoenix 3.16.0__py3-none-any.whl → 7.7.0__py3-none-any.whl

phoenix/db/engines.py

@@ -0,0 +1,208 @@
+import asyncio
+import json
+from collections.abc import Callable
+from datetime import datetime
+from enum import Enum
+from sqlite3 import Connection
+from typing import Any
+
+import aiosqlite
+import numpy as np
+import sqlalchemy
+import sqlean
+from sqlalchemy import URL, StaticPool, event, make_url
+from sqlalchemy.ext.asyncio import AsyncEngine, create_async_engine
+from typing_extensions import assert_never
+
+from phoenix.config import LoggingMode, get_env_database_schema
+from phoenix.db.helpers import SupportedSQLDialect
+from phoenix.db.migrate import migrate_in_thread
+from phoenix.db.models import init_models
+from phoenix.settings import Settings
+
+sqlean.extensions.enable("text", "stats")
+
+
+def set_sqlite_pragma(connection: Connection, _: Any) -> None:
+    cursor = connection.cursor()
+    cursor.execute("PRAGMA foreign_keys = ON;")
+    cursor.execute("PRAGMA journal_mode = WAL;")
+    cursor.execute("PRAGMA synchronous = OFF;")
+    cursor.execute("PRAGMA cache_size = -32000;")
+    cursor.execute("PRAGMA busy_timeout = 10000;")
+    cursor.close()
+
+
+def get_printable_db_url(connection_str: str) -> str:
+    return make_url(connection_str).render_as_string(hide_password=True)
+
+
+def get_async_db_url(connection_str: str) -> URL:
+    """
+    Parses the database URL string and returns a URL object that is async
+    """
+    url = make_url(connection_str)
+    if not url.database:
+        raise ValueError("Failed to parse database from connection string")
+    backend = SupportedSQLDialect(url.get_backend_name())
+    if backend is SupportedSQLDialect.SQLITE:
+        return url.set(drivername="sqlite+aiosqlite")
+    elif backend is SupportedSQLDialect.POSTGRESQL:
+        url = url.set(drivername="postgresql+asyncpg")
+        # For some reason username and password cannot be parsed from the typical slot
+        # So we need to parse them out manually
+        if url.username and url.password:
+            url = url.set(
+                query={**url.query, "user": url.username, "password": url.password},
+                password=None,
+                username=None,
+            )
+        return url
+    else:
+        assert_never(backend)
+
+
+def create_engine(
+    connection_str: str,
+    migrate: bool = True,
+    log_to_stdout: bool = False,
+) -> AsyncEngine:
+    """
+    Factory to create a SQLAlchemy engine from a URL string.
+    """
+    url = make_url(connection_str)
+    if not url.database:
+        raise ValueError("Failed to parse database from connection string")
+    backend = SupportedSQLDialect(url.get_backend_name())
+    url = get_async_db_url(url.render_as_string(hide_password=False))
+    # If Phoenix is run as an application, we want to pass log_migrations_to_stdout=False
+    # and let the configured sqlalchemy logger handle the migration logs
+    log_migrations_to_stdout = (
+        Settings.log_migrations and Settings.logging_mode != LoggingMode.STRUCTURED
+    )
+    if backend is SupportedSQLDialect.SQLITE:
+        return aio_sqlite_engine(
+            url=url,
+            migrate=migrate,
+            log_to_stdout=log_to_stdout,
+            log_migrations_to_stdout=log_migrations_to_stdout,
+        )
+    elif backend is SupportedSQLDialect.POSTGRESQL:
+        return aio_postgresql_engine(
+            url=url,
+            migrate=migrate,
+            log_to_stdout=log_to_stdout,
+            log_migrations_to_stdout=log_migrations_to_stdout,
+        )
+    else:
+        assert_never(backend)
+
+
+def aio_sqlite_engine(
+    url: URL,
+    migrate: bool = True,
+    shared_cache: bool = True,
+    log_to_stdout: bool = False,
+    log_migrations_to_stdout: bool = True,
+) -> AsyncEngine:
+    database = url.database or ":memory:"
+    if database.startswith("file:"):
+        database = database[5:]
+    if database.startswith(":memory:") and shared_cache:
+        url = url.set(query={**url.query, "cache": "shared"}, database=":memory:")
+    database = url.render_as_string().partition("///")[-1]
+
+    def async_creator() -> aiosqlite.Connection:
+        conn = aiosqlite.Connection(
+            lambda: sqlean.connect(f"file:{database}", uri=True),
+            iter_chunk_size=64,
+        )
+        conn.daemon = True
+        return conn
+
+    engine = create_async_engine(
+        url=url,
+        echo=log_to_stdout,
+        json_serializer=_dumps,
+        async_creator=async_creator,
+        poolclass=StaticPool,
+    )
+    event.listen(engine.sync_engine, "connect", set_sqlite_pragma)
+    if not migrate:
+        return engine
+    if database.startswith(":memory:"):
+        try:
+            asyncio.get_running_loop()
+        except RuntimeError:
+            asyncio.run(init_models(engine))
+        else:
+            asyncio.create_task(init_models(engine))
+    else:
+        sync_engine = sqlalchemy.create_engine(
+            url=url.set(drivername="sqlite"),
+            echo=log_migrations_to_stdout,
+            json_serializer=_dumps,
+            creator=lambda: sqlean.connect(f"file:{database}", uri=True),
+        )
+        migrate_in_thread(sync_engine)
+    return engine
+
+
+def set_postgresql_search_path(schema: str) -> Callable[[Connection, Any], None]:
+    def _(connection: Connection, _: Any) -> None:
+        cursor = connection.cursor()
+        cursor.execute(f"CREATE SCHEMA IF NOT EXISTS {schema};")
+        cursor.execute(f"SET search_path TO {schema};")
+
+    return _
+
+
+def aio_postgresql_engine(
+    url: URL,
+    migrate: bool = True,
+    log_to_stdout: bool = False,
+    log_migrations_to_stdout: bool = True,
+) -> AsyncEngine:
+    url_query = dict(url.query)
+    sslmode = url_query.pop("sslmode", None) or url_query.pop("ssl", None)
+    engine = create_async_engine(
+        url=url.set(
+            # https://github.com/MagicStack/asyncpg/issues/737
+            query={**url_query, "ssl": sslmode} if sslmode else url_query,
+        ),
+        echo=log_to_stdout,
+        json_serializer=_dumps,
+    )
+    if not migrate:
+        return engine
+    sync_engine = sqlalchemy.create_engine(
+        url=url.set(
+            drivername="postgresql+psycopg",
+            query={**url_query, "sslmode": sslmode} if sslmode else url_query,
+        ),
+        echo=log_migrations_to_stdout,
+        json_serializer=_dumps,
+    )
+    if schema := get_env_database_schema():
+        event.listen(sync_engine, "connect", set_postgresql_search_path(schema))
+    migrate_in_thread(sync_engine)
+    return engine
+
+
+def _dumps(obj: Any) -> str:
+    return json.dumps(obj, cls=_Encoder)
+
+
+class _Encoder(json.JSONEncoder):
+    def default(self, obj: Any) -> Any:
+        if isinstance(obj, datetime):
+            return obj.isoformat()
+        elif isinstance(obj, Enum):
+            return obj.value
+        elif isinstance(obj, np.ndarray):
+            return list(obj)
+        elif isinstance(obj, np.integer):
+            return int(obj)
+        elif isinstance(obj, np.floating):
+            return float(obj)
+        return super().default(obj)

phoenix/db/enums.py

@@ -0,0 +1,20 @@
+from collections.abc import Mapping
+from enum import Enum
+
+from sqlalchemy.orm import InstrumentedAttribute
+
+from phoenix.db import models
+from phoenix.db.models import AuthMethod
+
+__all__ = ["AuthMethod", "UserRole", "COLUMN_ENUMS"]
+
+
+class UserRole(Enum):
+    SYSTEM = "SYSTEM"
+    ADMIN = "ADMIN"
+    MEMBER = "MEMBER"
+
+
+COLUMN_ENUMS: Mapping[InstrumentedAttribute[str], type[Enum]] = {
+    models.UserRole.name: UserRole,
+}

phoenix/db/facilitator.py

@@ -0,0 +1,113 @@
+from __future__ import annotations
+
+import asyncio
+import secrets
+from functools import partial
+
+from sqlalchemy import (
+    distinct,
+    insert,
+    select,
+)
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from phoenix.auth import (
+    DEFAULT_ADMIN_EMAIL,
+    DEFAULT_ADMIN_USERNAME,
+    DEFAULT_SECRET_LENGTH,
+    DEFAULT_SYSTEM_EMAIL,
+    DEFAULT_SYSTEM_USERNAME,
+    compute_password_hash,
+)
+from phoenix.config import get_env_default_admin_initial_password
+from phoenix.db import models
+from phoenix.db.enums import COLUMN_ENUMS, UserRole
+from phoenix.server.types import DbSessionFactory
+
+
+class Facilitator:
+    """
+    Facilitates the creation of database records necessary for Phoenix to function. This includes
+    ensuring that all enum values are present in their respective tables, ensuring that all user
+    roles are present, and ensuring that the admin user has a password hash. These tasks will be
+    carried out as callbacks at the very beginning of Starlette's lifespan process.
+    """
+
+    def __init__(self, *, db: DbSessionFactory) -> None:
+        self._db = db
+
+    async def __call__(self) -> None:
+        async with self._db() as session:
+            for fn in (
+                _ensure_enums,
+                _ensure_user_roles,
+            ):
+                async with session.begin_nested():
+                    await fn(session)
+
+
+async def _ensure_enums(session: AsyncSession) -> None:
+    """
+    Ensure that all enum values are present in their respective tables. If any values are missing,
+    they will be added. If any values are present in the database but not in the enum, an error will
+    be raised. This function is idempotent: it will not add duplicate values to the database.
+    """
+    for column, enum in COLUMN_ENUMS.items():
+        table = column.class_
+        existing = set([_ async for _ in await session.stream_scalars(select(distinct(column)))])
+        expected = set(e.value for e in enum)
+        if unexpected := existing - expected:
+            raise ValueError(f"Unexpected values in {table.name}.{column.key}: {unexpected}")
+        if not (missing := expected - existing):
+            continue
+        await session.execute(insert(table), [{column.key: v} for v in missing])
+
+
+async def _ensure_user_roles(session: AsyncSession) -> None:
+    """
+    Ensure that the system and admin roles are present in the database. If they are not, they will
+    be added. The system user will have the email "system@localhost" and the admin user will have
+    the email "admin@localhost".
+    """
+    role_ids = {
+        name: id_
+        async for name, id_ in await session.stream(
+            select(models.UserRole.name, models.UserRole.id)
+        )
+    }
+    existing_roles = [
+        name
+        async for name in await session.stream_scalars(
+            select(distinct(models.UserRole.name)).join_from(models.User, models.UserRole)
+        )
+    ]
+    if (system_role := UserRole.SYSTEM.value) not in existing_roles and (
+        system_role_id := role_ids.get(system_role)
+    ) is not None:
+        system_user = models.User(
+            user_role_id=system_role_id,
+            username=DEFAULT_SYSTEM_USERNAME,
+            email=DEFAULT_SYSTEM_EMAIL,
+            reset_password=False,
+            password_salt=secrets.token_bytes(DEFAULT_SECRET_LENGTH),
+            password_hash=secrets.token_bytes(DEFAULT_SECRET_LENGTH),
+        )
+        session.add(system_user)
+    if (admin_role := UserRole.ADMIN.value) not in existing_roles and (
+        admin_role_id := role_ids.get(admin_role)
+    ) is not None:
+        salt = secrets.token_bytes(DEFAULT_SECRET_LENGTH)
+        password = get_env_default_admin_initial_password()
+        compute = partial(compute_password_hash, password=password, salt=salt)
+        loop = asyncio.get_running_loop()
+        hash_ = await loop.run_in_executor(None, compute)
+        admin_user = models.User(
+            user_role_id=admin_role_id,
+            username=DEFAULT_ADMIN_USERNAME,
+            email=DEFAULT_ADMIN_EMAIL,
+            password_salt=salt,
+            password_hash=hash_,
+            reset_password=True,
+        )
+        session.add(admin_user)
+    await session.flush()

phoenix/db/helpers.py

@@ -0,0 +1,159 @@
+from collections.abc import Callable, Hashable, Iterable
+from enum import Enum
+from typing import Any, Optional, TypeVar
+
+from openinference.semconv.trace import (
+    OpenInferenceSpanKindValues,
+    RerankerAttributes,
+    SpanAttributes,
+)
+from sqlalchemy import (
+    Integer,
+    Select,
+    SQLColumnExpression,
+    and_,
+    case,
+    distinct,
+    func,
+    select,
+)
+from typing_extensions import assert_never
+
+from phoenix.db import models
+
+
+class SupportedSQLDialect(Enum):
+    SQLITE = "sqlite"
+    POSTGRESQL = "postgresql"
+
+    @classmethod
+    def _missing_(cls, v: Any) -> "SupportedSQLDialect":
+        if isinstance(v, str) and v and v.isascii() and not v.islower():
+            return cls(v.lower())
+        raise ValueError(f"`{v}` is not a supported SQL backend/dialect.")
+
+
+def num_docs_col(dialect: SupportedSQLDialect) -> SQLColumnExpression[Integer]:
+    if dialect is SupportedSQLDialect.POSTGRESQL:
+        array_length = func.jsonb_array_length
+    elif dialect is SupportedSQLDialect.SQLITE:
+        array_length = func.json_array_length
+    else:
+        assert_never(dialect)
+    retrieval_docs = models.Span.attributes[_RETRIEVAL_DOCUMENTS]
+    num_retrieval_docs = array_length(retrieval_docs)
+    reranker_docs = models.Span.attributes[_RERANKER_OUTPUT_DOCUMENTS]
+    num_reranker_docs = array_length(reranker_docs)
+    return case(
+        (
+            func.upper(models.Span.span_kind) == OpenInferenceSpanKindValues.RERANKER.value.upper(),
+            num_reranker_docs,
+        ),
+        else_=num_retrieval_docs,
+    ).label("num_docs")
+
+
+_RETRIEVAL_DOCUMENTS = SpanAttributes.RETRIEVAL_DOCUMENTS.split(".")
+_RERANKER_OUTPUT_DOCUMENTS = RerankerAttributes.RERANKER_OUTPUT_DOCUMENTS.split(".")
+
+
+def get_eval_trace_ids_for_datasets(*dataset_ids: int) -> Select[tuple[Optional[str]]]:
+    return (
+        select(distinct(models.ExperimentRunAnnotation.trace_id))
+        .join(models.ExperimentRun)
+        .join_from(models.ExperimentRun, models.Experiment)
+        .where(models.Experiment.dataset_id.in_(set(dataset_ids)))
+        .where(models.ExperimentRunAnnotation.trace_id.isnot(None))
+    )
+
+
+def get_project_names_for_datasets(*dataset_ids: int) -> Select[tuple[Optional[str]]]:
+    return (
+        select(distinct(models.Experiment.project_name))
+        .where(models.Experiment.dataset_id.in_(set(dataset_ids)))
+        .where(models.Experiment.project_name.isnot(None))
+    )
+
+
+def get_eval_trace_ids_for_experiments(*experiment_ids: int) -> Select[tuple[Optional[str]]]:
+    return (
+        select(distinct(models.ExperimentRunAnnotation.trace_id))
+        .join(models.ExperimentRun)
+        .where(models.ExperimentRun.experiment_id.in_(set(experiment_ids)))
+        .where(models.ExperimentRunAnnotation.trace_id.isnot(None))
+    )
+
+
+def get_project_names_for_experiments(*experiment_ids: int) -> Select[tuple[Optional[str]]]:
+    return (
+        select(distinct(models.Experiment.project_name))
+        .where(models.Experiment.id.in_(set(experiment_ids)))
+        .where(models.Experiment.project_name.isnot(None))
+    )
+
+
+_AnyT = TypeVar("_AnyT")
+_KeyT = TypeVar("_KeyT", bound=Hashable)
+
+
+def dedup(
+    items: Iterable[_AnyT],
+    key: Callable[[_AnyT], _KeyT],
+) -> list[_AnyT]:
+    """
+    Discard subsequent duplicates after the first appearance in `items`.
+    """
+    ans = []
+    seen: set[_KeyT] = set()
+    for item in items:
+        if (k := key(item)) in seen:
+            continue
+        else:
+            ans.append(item)
+            seen.add(k)
+    return ans
+
+
+def get_dataset_example_revisions(
+    dataset_version_id: int,
+) -> Select[tuple[models.DatasetExampleRevision]]:
+    version = (
+        select(
+            models.DatasetVersion.id,
+            models.DatasetVersion.dataset_id,
+        )
+        .filter_by(id=dataset_version_id)
+        .subquery()
+    )
+    table = models.DatasetExampleRevision
+    revision = (
+        select(
+            table.dataset_example_id,
+            func.max(table.dataset_version_id).label("dataset_version_id"),
+        )
+        .join_from(
+            table,
+            models.DatasetExample,
+            table.dataset_example_id == models.DatasetExample.id,
+        )
+        .join_from(
+            models.DatasetExample,
+            version,
+            models.DatasetExample.dataset_id == version.c.dataset_id,
+        )
+        .where(models.DatasetExample.dataset_id == version.c.dataset_id)
+        .where(table.dataset_version_id <= version.c.id)
+        .group_by(table.dataset_example_id)
+        .subquery()
+    )
+    return (
+        select(table)
+        .where(table.revision_kind != "DELETE")
+        .join(
+            revision,
+            onclause=and_(
+                revision.c.dataset_example_id == table.dataset_example_id,
+                revision.c.dataset_version_id == table.dataset_version_id,
+            ),
+        )
+    )

phoenix/db/insertion/constants.py

@@ -0,0 +1,2 @@
+DEFAULT_RETRY_DELAY_SEC: float = 10
+DEFAULT_RETRY_ALLOWANCE: int = 60