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

phoenix/db/insertion/types.py

@@ -0,0 +1,256 @@
+from __future__ import annotations
+
+import asyncio
+import logging
+from abc import ABC, abstractmethod
+from collections.abc import Mapping, Sequence
+from copy import copy
+from dataclasses import dataclass, field
+from datetime import datetime, timezone
+from typing import Any, Generic, Optional, Protocol, TypeVar, cast
+
+from sqlalchemy.ext.asyncio import AsyncSession
+from sqlalchemy.sql.dml import Insert
+
+from phoenix.db import models
+from phoenix.db.insertion.constants import DEFAULT_RETRY_ALLOWANCE, DEFAULT_RETRY_DELAY_SEC
+from phoenix.db.insertion.helpers import insert_on_conflict
+from phoenix.server.dml_event import DmlEvent
+from phoenix.server.types import DbSessionFactory
+
+logger = logging.getLogger(__name__)
+
+
+class Insertable(Protocol):
+    @property
+    def row(self) -> models.Base: ...
+
+
+_AnyT = TypeVar("_AnyT")
+_PrecursorT = TypeVar("_PrecursorT")
+_InsertableT = TypeVar("_InsertableT", bound=Insertable)
+_RowT = TypeVar("_RowT", bound=models.Base)
+_DmlEventT = TypeVar("_DmlEventT", bound=DmlEvent)
+
+
+@dataclass(frozen=True)
+class Received(Generic[_AnyT]):
+    item: _AnyT
+    received_at: datetime = field(default_factory=lambda: datetime.now(timezone.utc))
+
+    def postpone(self, retries_left: int = DEFAULT_RETRY_ALLOWANCE) -> Postponed[_AnyT]:
+        return Postponed(item=self.item, received_at=self.received_at, retries_left=retries_left)
+
+
+@dataclass(frozen=True)
+class Postponed(Received[_AnyT]):
+    retries_left: int = field(default=DEFAULT_RETRY_ALLOWANCE)
+
+
+class QueueInserter(ABC, Generic[_PrecursorT, _InsertableT, _RowT, _DmlEventT]):
+    table: type[_RowT]
+    unique_by: Sequence[str]
+
+    def __init_subclass__(
+        cls,
+        table: type[_RowT],
+        unique_by: Sequence[str],
+    ) -> None:
+        cls.table = table
+        cls.unique_by = unique_by
+
+    def __init__(
+        self,
+        db: DbSessionFactory,
+        retry_delay_sec: float = DEFAULT_RETRY_DELAY_SEC,
+        retry_allowance: int = DEFAULT_RETRY_ALLOWANCE,
+    ) -> None:
+        self._queue: list[Received[_PrecursorT]] = []
+        self._db = db
+        self._retry_delay_sec = retry_delay_sec
+        self._retry_allowance = retry_allowance
+
+    @property
+    def empty(self) -> bool:
+        return not bool(self._queue)
+
+    async def enqueue(self, *items: _PrecursorT) -> None:
+        self._queue.extend([Received(item) for item in items])
+
+    @abstractmethod
+    async def _partition(
+        self,
+        session: AsyncSession,
+        *parcels: Received[_PrecursorT],
+    ) -> tuple[
+        list[Received[_InsertableT]],
+        list[Postponed[_PrecursorT]],
+        list[Received[_PrecursorT]],
+    ]: ...
+
+    async def insert(self) -> Optional[list[_DmlEventT]]:
+        if not self._queue:
+            return None
+        self._queue, parcels = [], self._queue
+        events: list[_DmlEventT] = []
+        async with self._db() as session:
+            to_insert, to_postpone, _ = await self._partition(session, *parcels)
+            if to_insert:
+                events, to_retry, _ = await self._insert(session, *to_insert)
+                if to_retry:
+                    to_postpone.extend(to_retry)
+        if to_postpone:
+            loop = asyncio.get_running_loop()
+            loop.call_later(self._retry_delay_sec, self._queue.extend, to_postpone)
+        return events
+
+    def _insert_on_conflict(self, *records: Mapping[str, Any]) -> Insert:
+        return insert_on_conflict(
+            *records,
+            table=self.table,
+            unique_by=self.unique_by,
+            dialect=self._db.dialect,
+        )
+
+    @abstractmethod
+    async def _events(
+        self,
+        session: AsyncSession,
+        *insertions: _InsertableT,
+    ) -> list[_DmlEventT]: ...
+
+    async def _insert(
+        self,
+        session: AsyncSession,
+        *parcels: Received[_InsertableT],
+    ) -> tuple[
+        list[_DmlEventT],
+        list[Postponed[_PrecursorT]],
+        list[Received[_InsertableT]],
+    ]:
+        to_retry: list[Postponed[_PrecursorT]] = []
+        failures: list[Received[_InsertableT]] = []
+        events: list[_DmlEventT] = []
+        try:
+            async with session.begin_nested():
+                events.extend(await self._events(session, *(p.item for p in parcels)))
+        except BaseException:
+            logger.exception(
+                f"Failed to bulk insert for {self.table.__name__}. "
+                f"Will try to insert ({len(parcels)} records) individually instead."
+            )
+            for p in parcels:
+                try:
+                    async with session.begin_nested():
+                        events.extend(await self._events(session, p.item))
+                except BaseException:
+                    logger.exception(f"Failed to insert for {self.table.__name__}.")
+                    if isinstance(p, Postponed) and p.retries_left == 1:
+                        failures.append(p)
+                    else:
+                        to_retry.append(
+                            Postponed(
+                                item=cast(_PrecursorT, p.item),
+                                received_at=p.received_at,
+                                retries_left=(p.retries_left - 1)
+                                if isinstance(p, Postponed)
+                                else self._retry_allowance,
+                            )
+                        )
+        return events, to_retry, failures
+
+
+class Precursors(ABC):
+    @dataclass(frozen=True)
+    class SpanAnnotation:
+        span_id: str
+        obj: models.SpanAnnotation
+
+        def as_insertable(
+            self,
+            span_rowid: int,
+            id_: Optional[int] = None,
+        ) -> Insertables.SpanAnnotation:
+            return Insertables.SpanAnnotation(
+                span_id=self.span_id,
+                obj=self.obj,
+                span_rowid=span_rowid,
+                id_=id_,
+            )
+
+    @dataclass(frozen=True)
+    class TraceAnnotation:
+        trace_id: str
+        obj: models.TraceAnnotation
+
+        def as_insertable(
+            self,
+            trace_rowid: int,
+            id_: Optional[int] = None,
+        ) -> Insertables.TraceAnnotation:
+            return Insertables.TraceAnnotation(
+                trace_id=self.trace_id,
+                obj=self.obj,
+                trace_rowid=trace_rowid,
+                id_=id_,
+            )
+
+    @dataclass(frozen=True)
+    class DocumentAnnotation:
+        span_id: str
+        document_position: int
+        obj: models.DocumentAnnotation
+
+        def as_insertable(
+            self,
+            span_rowid: int,
+            id_: Optional[int] = None,
+        ) -> Insertables.DocumentAnnotation:
+            return Insertables.DocumentAnnotation(
+                span_id=self.span_id,
+                document_position=self.document_position,
+                obj=self.obj,
+                span_rowid=span_rowid,
+                id_=id_,
+            )
+
+
+class Insertables(ABC):
+    @dataclass(frozen=True)
+    class SpanAnnotation(Precursors.SpanAnnotation):
+        span_rowid: int
+        id_: Optional[int] = None
+
+        @property
+        def row(self) -> models.SpanAnnotation:
+            obj = copy(self.obj)
+            obj.span_rowid = self.span_rowid
+            if self.id_ is not None:
+                obj.id = self.id_
+            return obj
+
+    @dataclass(frozen=True)
+    class TraceAnnotation(Precursors.TraceAnnotation):
+        trace_rowid: int
+        id_: Optional[int] = None
+
+        @property
+        def row(self) -> models.TraceAnnotation:
+            obj = copy(self.obj)
+            obj.trace_rowid = self.trace_rowid
+            if self.id_ is not None:
+                obj.id = self.id_
+            return obj
+
+    @dataclass(frozen=True)
+    class DocumentAnnotation(Precursors.DocumentAnnotation):
+        span_rowid: int
+        id_: Optional[int] = None
+
+        @property
+        def row(self) -> models.DocumentAnnotation:
+            obj = copy(self.obj)
+            obj.span_rowid = self.span_rowid
+            if self.id_ is not None:
+                obj.id = self.id_
+            return obj

phoenix/db/migrate.py

@@ -0,0 +1,86 @@
+import codecs
+import logging
+import sys
+from pathlib import Path
+from queue import Empty, SimpleQueue
+from threading import Thread
+from time import perf_counter
+from typing import Optional
+
+from alembic import command
+from alembic.config import Config
+from sqlalchemy import Engine
+
+from phoenix.exceptions import PhoenixMigrationError
+from phoenix.settings import Settings
+
+logger = logging.getLogger(__name__)
+
+
+def printif(condition: bool, text: str) -> None:
+    if not condition:
+        return
+    if sys.platform.startswith("win"):
+        text = codecs.encode(text, "ascii", errors="ignore").decode("ascii").strip()
+    print(text)
+
+
+def migrate(
+    engine: Engine,
+    error_queue: Optional["SimpleQueue[BaseException]"] = None,
+) -> None:
+    """
+    Runs migrations on the database.
+    NB: Migrate only works on non-memory databases.
+
+    Args:
+        url: The database URL.
+    """
+    try:
+        log_migrations = Settings.log_migrations
+        printif(log_migrations, "🏃‍♀️‍➡️ Running migrations on the database.")
+        printif(log_migrations, "---------------------------")
+        config_path = str(Path(__file__).parent.resolve() / "alembic.ini")
+        alembic_cfg = Config(config_path)
+
+        # Explicitly set the migration directory
+        scripts_location = str(Path(__file__).parent.resolve() / "migrations")
+        alembic_cfg.set_main_option("script_location", scripts_location)
+        url = str(engine.url).replace("%", "%%")
+        alembic_cfg.set_main_option("sqlalchemy.url", url)
+        start_time = perf_counter()
+        with engine.connect() as conn:
+            alembic_cfg.attributes["connection"] = conn
+            command.upgrade(alembic_cfg, "head")
+        elapsed_time = perf_counter() - start_time
+        engine.dispose()
+        printif(log_migrations, "---------------------------")
+        printif(log_migrations, f"✅ Migrations completed in {elapsed_time:.3f} seconds.")
+    except BaseException as e:
+        if error_queue:
+            error_queue.put(e)
+            raise e
+
+
+def migrate_in_thread(engine: Engine) -> None:
+    """
+    Runs migrations on the database in a separate thread.
+    This is needed because depending on the context (notebook)
+    the migration process can fail to execute in the main thread.
+    """
+    error_queue: SimpleQueue[BaseException] = SimpleQueue()
+    t = Thread(target=migrate, args=(engine, error_queue))
+    t.start()
+    t.join()
+
+    try:
+        result = error_queue.get_nowait()
+    except Empty:
+        return
+
+    if result is not None:
+        error_message = (
+            "\n\nUnable to migrate configured Phoenix DB. Original error:\n"
+            f"{type(result).__name__}: {str(result)}"
+        )
+        raise PhoenixMigrationError(error_message) from result

phoenix/db/migrations/data_migration_scripts/populate_project_sessions.py

@@ -0,0 +1,199 @@
+# /// script
+# dependencies = [
+#   "arize-phoenix[pg]",
+# ]
+# ///
+"""
+Populate the `project_sessions` table with data from the traces and spans tables.
+
+Environment variables.
+
+- `PHOENIX_SQL_DATABASE_URL` must be set to the database connection string.
+- (optional) Postgresql schema can be set via `PHOENIX_SQL_DATABASE_SCHEMA`.
+"""
+
+import os
+from datetime import datetime
+from time import perf_counter
+from typing import Any, Optional, Union
+
+import sqlean
+from openinference.semconv.trace import SpanAttributes
+from sqlalchemy import (
+    JSON,
+    Engine,
+    NullPool,
+    create_engine,
+    event,
+    func,
+    insert,
+    make_url,
+    select,
+    update,
+)
+from sqlalchemy.dialects import postgresql
+from sqlalchemy.ext.compiler import compiles
+from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column, sessionmaker
+
+from phoenix.config import ENV_PHOENIX_SQL_DATABASE_SCHEMA, get_env_database_connection_str
+from phoenix.db.engines import set_postgresql_search_path
+
+
+class JSONB(JSON):
+    # See https://docs.sqlalchemy.org/en/20/core/custom_types.html
+    __visit_name__ = "JSONB"
+
+
+@compiles(JSONB, "sqlite")
+def _(*args: Any, **kwargs: Any) -> str:
+    # See https://docs.sqlalchemy.org/en/20/core/custom_types.html
+    return "JSONB"
+
+
+JSON_ = (
+    JSON()
+    .with_variant(
+        postgresql.JSONB(),  # type: ignore
+        "postgresql",
+    )
+    .with_variant(
+        JSONB(),
+        "sqlite",
+    )
+)
+
+
+class Base(DeclarativeBase): ...
+
+
+class ProjectSession(Base):
+    __tablename__ = "project_sessions"
+    id: Mapped[int] = mapped_column(primary_key=True)
+    session_id: Mapped[str]
+    project_id: Mapped[int]
+    start_time: Mapped[datetime]
+    end_time: Mapped[datetime]
+
+
+class Trace(Base):
+    __tablename__ = "traces"
+    id: Mapped[int] = mapped_column(primary_key=True)
+    project_session_rowid: Mapped[Union[int, None]]
+    project_rowid: Mapped[int]
+    start_time: Mapped[datetime]
+    end_time: Mapped[datetime]
+
+
+class Span(Base):
+    __tablename__ = "spans"
+    id: Mapped[int] = mapped_column(primary_key=True)
+    trace_rowid: Mapped[int]
+    parent_id: Mapped[Optional[str]]
+    attributes: Mapped[dict[str, Any]] = mapped_column(JSON_, nullable=False)
+
+
+SESSION_ID = SpanAttributes.SESSION_ID.split(".")
+USER_ID = SpanAttributes.USER_ID.split(".")
+
+
+def populate_project_sessions(
+    engine: Engine,
+) -> None:
+    sessions_from_span = (
+        select(
+            Span.attributes[SESSION_ID].as_string().label("session_id"),
+            Trace.project_rowid.label("project_id"),
+            Trace.start_time.label("start_time"),
+            func.row_number()
+            .over(
+                partition_by=Span.attributes[SESSION_ID],
+                order_by=[Trace.start_time, Trace.id, Span.id],
+            )
+            .label("rank"),
+            func.max(Trace.end_time)
+            .over(partition_by=Span.attributes[SESSION_ID])
+            .label("end_time"),
+        )
+        .join_from(Span, Trace, Span.trace_rowid == Trace.id)
+        .where(Span.parent_id.is_(None))
+        .where(Span.attributes[SESSION_ID].as_string() != "")
+        .subquery()
+    )
+    sessions_for_trace_id = (
+        select(
+            Span.trace_rowid,
+            ProjectSession.id.label("project_session_rowid"),
+        )
+        .join_from(
+            Span,
+            ProjectSession,
+            Span.attributes[SESSION_ID].as_string() == ProjectSession.session_id,
+        )
+        .where(Span.parent_id.is_(None))
+        .where(Span.attributes[SESSION_ID].as_string() != "")
+        .subquery()
+    )
+    start_time = perf_counter()
+    with sessionmaker(engine).begin() as session:
+        session.execute(
+            insert(ProjectSession).from_select(
+                [
+                    "session_id",
+                    "project_id",
+                    "start_time",
+                    "end_time",
+                ],
+                select(
+                    sessions_from_span.c.session_id,
+                    sessions_from_span.c.project_id,
+                    sessions_from_span.c.start_time,
+                    sessions_from_span.c.end_time,
+                ).where(sessions_from_span.c.rank == 1),
+            )
+        )
+        session.execute(
+            (
+                update(Trace)
+                .values(project_session_rowid=sessions_for_trace_id.c.project_session_rowid)
+                .where(Trace.id == sessions_for_trace_id.c.trace_rowid)
+            )
+        )
+    elapsed_time = perf_counter() - start_time
+    print(f"✅ Populated project_sessions in {elapsed_time:.3f} seconds.")
+
+
+if __name__ == "__main__":
+    sql_database_url = make_url(get_env_database_connection_str())
+    print(f"Using database URL: {sql_database_url}")
+    ans = input("Is that correct? [y]/n: ")
+    if ans.lower().startswith("n"):
+        url = input("Please enter the correct database URL: ")
+        sql_database_url = make_url(url)
+    backend = sql_database_url.get_backend_name()
+    if backend == "sqlite":
+        file = sql_database_url.database
+        engine = create_engine(
+            url=sql_database_url.set(drivername="sqlite"),
+            creator=lambda: sqlean.connect(f"file:///{file}", uri=True),
+            poolclass=NullPool,
+            echo=True,
+        )
+    elif backend == "postgresql":
+        schema = os.getenv(ENV_PHOENIX_SQL_DATABASE_SCHEMA)
+        if schema:
+            print(f"Using schema: {schema}")
+        else:
+            print("No PostgreSQL schema set. (This is the default.)")
+        ans = input("Is that correct? [y]/n: ")
+        if ans.lower().startswith("n"):
+            schema = input("Please enter the correct schema: ")
+        engine = create_engine(
+            url=sql_database_url.set(drivername="postgresql+psycopg"),
+            poolclass=NullPool,
+            echo=True,
+        )
+        if schema:
+            event.listen(engine, "connect", set_postgresql_search_path(schema))
+    else:
+        raise ValueError(f"Unknown database backend: {backend}")
+    populate_project_sessions(engine)

phoenix/db/migrations/env.py

@@ -0,0 +1,114 @@
+import asyncio
+
+from alembic import context
+from sqlalchemy import Connection, engine_from_config, pool
+from sqlalchemy.ext.asyncio import AsyncEngine
+
+from phoenix.config import get_env_database_connection_str
+from phoenix.db.engines import get_async_db_url
+from phoenix.db.models import Base
+from phoenix.settings import Settings
+
+# this is the Alembic Config object, which provides
+# access to the values within the .ini file in use.
+config = context.config
+
+# add your model's MetaData object here
+# for 'autogenerate' support
+target_metadata = Base.metadata
+
+# other values from the config, defined by the needs of env.py,
+# can be acquired:
+# my_important_option = config.get_main_option("my_important_option")
+# ... etc.
+
+
+def run_migrations_offline() -> None:
+    """Run migrations in 'offline' mode.
+
+    This configures the context with just a URL
+    and not an Engine, though an Engine is acceptable
+    here as well.  By skipping the Engine creation
+    we don't even need a DBAPI to be available.
+
+    Calls to context.execute() here emit the given string to the
+    script output.
+
+    """
+    url = config.get_main_option("sqlalchemy.url")
+    context.configure(
+        url=url,
+        target_metadata=target_metadata,
+        literal_binds=True,
+        dialect_opts={"paramstyle": "named"},
+        transaction_per_migration=True,
+    )
+
+    with context.begin_transaction():
+        context.run_migrations()
+
+
+def run_migrations_online() -> None:
+    """Run migrations in 'online' mode.
+
+    In this scenario we need to create an Engine
+    and associate a connection with the context.
+
+    """
+    connectable = context.config.attributes.get("connection", None)
+    if connectable is None:
+        config = context.config.get_section(context.config.config_ini_section) or {}
+        if "sqlalchemy.url" not in config:
+            connection_str = get_env_database_connection_str()
+            config["sqlalchemy.url"] = get_async_db_url(connection_str).render_as_string(
+                hide_password=False
+            )
+        connectable = AsyncEngine(
+            engine_from_config(
+                config,
+                prefix="sqlalchemy.",
+                poolclass=pool.NullPool,
+                future=True,
+                echo=Settings.log_migrations,
+            )
+        )
+
+    if isinstance(connectable, AsyncEngine):
+        try:
+            asyncio.get_running_loop()
+        except RuntimeError:
+            asyncio.run(run_async_migrations(connectable))
+        else:
+            asyncio.create_task(run_async_migrations(connectable))
+    else:
+        run_migrations(connectable)
+
+
+async def run_async_migrations(connectable: AsyncEngine) -> None:
+    async with connectable.connect() as connection:
+        await connection.run_sync(run_migrations)
+
+
+def run_migrations(connection: Connection) -> None:
+    transaction = connection.begin()
+    try:
+        context.configure(
+            connection=connection,
+            target_metadata=target_metadata,
+            compare_type=True,
+            transactional_ddl=True,
+            transaction_per_migration=True,
+        )
+        context.run_migrations()
+        transaction.commit()
+    except Exception:
+        transaction.rollback()
+        raise
+    finally:
+        connection.close()
+
+
+if context.is_offline_mode():
+    run_migrations_offline()
+else:
+    run_migrations_online()

phoenix/db/migrations/script.py.mako

@@ -0,0 +1,26 @@
+"""${message}
+
+Revision ID: ${up_revision}
+Revises: ${down_revision | comma,n}
+Create Date: ${create_date}
+
+"""
+from typing import Sequence, Union
+
+from alembic import op
+import sqlalchemy as sa
+${imports if imports else ""}
+
+# revision identifiers, used by Alembic.
+revision: str = ${repr(up_revision)}
+down_revision: Union[str, None] = ${repr(down_revision)}
+branch_labels: Union[str, Sequence[str], None] = ${repr(branch_labels)}
+depends_on: Union[str, Sequence[str], None] = ${repr(depends_on)}
+
+
+def upgrade() -> None:
+    ${upgrades if upgrades else "pass"}
+
+
+def downgrade() -> None:
+    ${downgrades if downgrades else "pass"}