arize-phoenix 3.24.0__py3-none-any.whl → 4.0.0__py3-none-any.whl

phoenix/db/insertion/evaluation.py

@@ -0,0 +1,209 @@
+from typing import NamedTuple, Optional
+
+from sqlalchemy import select
+from sqlalchemy.ext.asyncio import AsyncSession
+from typing_extensions import assert_never
+
+from phoenix.db import models
+from phoenix.db.helpers import SupportedSQLDialect, num_docs_col
+from phoenix.db.insertion.helpers import OnConflict, insert_stmt
+from phoenix.exceptions import PhoenixException
+from phoenix.trace import v1 as pb
+
+
+class InsertEvaluationError(PhoenixException):
+    pass
+
+
+class EvaluationInsertionResult(NamedTuple):
+    project_rowid: int
+    evaluation_name: str
+
+
+class SpanEvaluationInsertionEvent(EvaluationInsertionResult): ...
+
+
+class TraceEvaluationInsertionEvent(EvaluationInsertionResult): ...
+
+
+class DocumentEvaluationInsertionEvent(EvaluationInsertionResult): ...
+
+
+async def insert_evaluation(
+    session: AsyncSession,
+    evaluation: pb.Evaluation,
+) -> Optional[EvaluationInsertionResult]:
+    evaluation_name = evaluation.name
+    result = evaluation.result
+    label = result.label.value if result.HasField("label") else None
+    score = result.score.value if result.HasField("score") else None
+    explanation = result.explanation.value if result.HasField("explanation") else None
+    if (evaluation_kind := evaluation.subject_id.WhichOneof("kind")) is None:
+        raise InsertEvaluationError("Cannot insert an evaluation that has no evaluation kind")
+    elif evaluation_kind == "trace_id":
+        trace_id = evaluation.subject_id.trace_id
+        return await _insert_trace_evaluation(
+            session, trace_id, evaluation_name, label, score, explanation
+        )
+    elif evaluation_kind == "span_id":
+        span_id = evaluation.subject_id.span_id
+        return await _insert_span_evaluation(
+            session, span_id, evaluation_name, label, score, explanation
+        )
+    elif evaluation_kind == "document_retrieval_id":
+        span_id = evaluation.subject_id.document_retrieval_id.span_id
+        document_position = evaluation.subject_id.document_retrieval_id.document_position
+        return await _insert_document_evaluation(
+            session, span_id, document_position, evaluation_name, label, score, explanation
+        )
+    else:
+        assert_never(evaluation_kind)
+
+
+async def _insert_trace_evaluation(
+    session: AsyncSession,
+    trace_id: str,
+    evaluation_name: str,
+    label: Optional[str],
+    score: Optional[float],
+    explanation: Optional[str],
+) -> TraceEvaluationInsertionEvent:
+    stmt = select(
+        models.Trace.project_rowid,
+        models.Trace.id,
+    ).where(models.Trace.trace_id == trace_id)
+    if not (row := (await session.execute(stmt)).first()):
+        raise InsertEvaluationError(
+            f"Cannot insert a trace evaluation for a missing trace: {evaluation_name=}, {trace_id=}"
+        )
+    project_rowid, trace_rowid = row
+    dialect = SupportedSQLDialect(session.bind.dialect.name)
+    values = dict(
+        trace_rowid=trace_rowid,
+        name=evaluation_name,
+        label=label,
+        score=score,
+        explanation=explanation,
+        metadata_={},  # `metadata_` must match ORM
+        annotator_kind="LLM",
+    )
+    set_ = dict(values)
+    set_.pop("metadata_")
+    set_["metadata"] = values["metadata_"]  # `metadata` must match database
+    await session.execute(
+        insert_stmt(
+            dialect=dialect,
+            table=models.TraceAnnotation,
+            values=values,
+            constraint="uq_trace_annotations_name_trace_rowid",
+            column_names=("name", "trace_rowid"),
+            on_conflict=OnConflict.DO_UPDATE,
+            set_=set_,
+        )
+    )
+    return TraceEvaluationInsertionEvent(project_rowid, evaluation_name)
+
+
+async def _insert_span_evaluation(
+    session: AsyncSession,
+    span_id: str,
+    evaluation_name: str,
+    label: Optional[str],
+    score: Optional[float],
+    explanation: Optional[str],
+) -> SpanEvaluationInsertionEvent:
+    stmt = (
+        select(
+            models.Trace.project_rowid,
+            models.Span.id,
+        )
+        .join_from(models.Span, models.Trace)
+        .where(models.Span.span_id == span_id)
+    )
+    if not (row := (await session.execute(stmt)).first()):
+        raise InsertEvaluationError(
+            f"Cannot insert a span evaluation for a missing span: {evaluation_name=}, {span_id=}"
+        )
+    project_rowid, span_rowid = row
+    dialect = SupportedSQLDialect(session.bind.dialect.name)
+    values = dict(
+        span_rowid=span_rowid,
+        name=evaluation_name,
+        label=label,
+        score=score,
+        explanation=explanation,
+        metadata_={},  # `metadata_` must match ORM
+        annotator_kind="LLM",
+    )
+    set_ = dict(values)
+    set_.pop("metadata_")
+    set_["metadata"] = values["metadata_"]  # `metadata` must match database
+    await session.execute(
+        insert_stmt(
+            dialect=dialect,
+            table=models.SpanAnnotation,
+            values=values,
+            constraint="uq_span_annotations_name_span_rowid",
+            column_names=("name", "span_rowid"),
+            on_conflict=OnConflict.DO_UPDATE,
+            set_=set_,
+        )
+    )
+    return SpanEvaluationInsertionEvent(project_rowid, evaluation_name)
+
+
+async def _insert_document_evaluation(
+    session: AsyncSession,
+    span_id: str,
+    document_position: int,
+    evaluation_name: str,
+    label: Optional[str],
+    score: Optional[float],
+    explanation: Optional[str],
+) -> EvaluationInsertionResult:
+    dialect = SupportedSQLDialect(session.bind.dialect.name)
+    stmt = (
+        select(
+            models.Trace.project_rowid,
+            models.Span.id,
+            num_docs_col(dialect),
+        )
+        .join_from(models.Span, models.Trace)
+        .where(models.Span.span_id == span_id)
+    )
+    if not (row := (await session.execute(stmt)).first()):
+        raise InsertEvaluationError(
+            f"Cannot insert a document evaluation for a missing span: {span_id=}"
+        )
+    project_rowid, span_rowid, num_docs = row
+    if num_docs is None or num_docs <= document_position:
+        raise InsertEvaluationError(
+            f"Cannot insert a document evaluation for a non-existent "
+            f"document position: {evaluation_name=}, {span_id=}, {document_position=}"
+        )
+    dialect = SupportedSQLDialect(session.bind.dialect.name)
+    values = dict(
+        span_rowid=span_rowid,
+        document_position=document_position,
+        name=evaluation_name,
+        label=label,
+        score=score,
+        explanation=explanation,
+        metadata_={},  # `metadata_` must match ORM
+        annotator_kind="LLM",
+    )
+    set_ = dict(values)
+    set_.pop("metadata_")
+    set_["metadata"] = values["metadata_"]  # `metadata` must match database
+    await session.execute(
+        insert_stmt(
+            dialect=dialect,
+            table=models.DocumentAnnotation,
+            values=values,
+            constraint="uq_document_annotations_name_span_rowid_document_position",
+            column_names=("name", "span_rowid", "document_position"),
+            on_conflict=OnConflict.DO_UPDATE,
+            set_=set_,
+        )
+    )
+    return DocumentEvaluationInsertionEvent(project_rowid, evaluation_name)

phoenix/db/insertion/helpers.py

@@ -0,0 +1,54 @@
+from enum import Enum, auto
+from typing import Any, Mapping, Optional, Sequence
+
+from sqlalchemy import Insert, insert
+from sqlalchemy.dialects.postgresql import insert as insert_postgresql
+from sqlalchemy.dialects.sqlite import insert as insert_sqlite
+from typing_extensions import assert_never
+
+from phoenix.db.helpers import SupportedSQLDialect
+
+
+class OnConflict(Enum):
+    DO_NOTHING = auto()
+    DO_UPDATE = auto()
+
+
+def insert_stmt(
+    dialect: SupportedSQLDialect,
+    table: Any,
+    values: Mapping[str, Any],
+    constraint: Optional[str] = None,
+    column_names: Sequence[str] = (),
+    on_conflict: OnConflict = OnConflict.DO_NOTHING,
+    set_: Optional[Mapping[str, Any]] = None,
+) -> Insert:
+    """
+    Dialect specific insertion statement using ON CONFLICT DO syntax.
+    """
+    if bool(constraint) != bool(column_names):
+        raise ValueError(
+            "Both `constraint` and `column_names` must be provided or omitted at the same time."
+        )
+    if (
+        dialect is SupportedSQLDialect.POSTGRESQL
+        and constraint is None
+        or dialect is SupportedSQLDialect.SQLITE
+        and not column_names
+    ):
+        return insert(table).values(values)
+    if dialect is SupportedSQLDialect.POSTGRESQL:
+        stmt_postgresql = insert_postgresql(table).values(values)
+        if on_conflict is OnConflict.DO_NOTHING or not set_:
+            return stmt_postgresql.on_conflict_do_nothing(constraint=constraint)
+        if on_conflict is OnConflict.DO_UPDATE:
+            return stmt_postgresql.on_conflict_do_update(constraint=constraint, set_=set_)
+        assert_never(on_conflict)
+    if dialect is SupportedSQLDialect.SQLITE:
+        stmt_sqlite = insert_sqlite(table).values(values)
+        if on_conflict is OnConflict.DO_NOTHING or not set_:
+            return stmt_sqlite.on_conflict_do_nothing(column_names)
+        if on_conflict is OnConflict.DO_UPDATE:
+            return stmt_sqlite.on_conflict_do_update(column_names, set_=set_)
+        assert_never(on_conflict)
+    assert_never(dialect)

phoenix/db/insertion/span.py

@@ -0,0 +1,142 @@
+from dataclasses import asdict
+from typing import NamedTuple, Optional, cast
+
+from openinference.semconv.trace import SpanAttributes
+from sqlalchemy import func, insert, select, update
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from phoenix.db import models
+from phoenix.db.helpers import SupportedSQLDialect
+from phoenix.db.insertion.helpers import OnConflict, insert_stmt
+from phoenix.trace.attributes import get_attribute_value
+from phoenix.trace.schemas import Span, SpanStatusCode
+
+
+class SpanInsertionEvent(NamedTuple):
+    project_rowid: int
+
+
+class ClearProjectSpansEvent(NamedTuple):
+    project_rowid: int
+
+
+async def insert_span(
+    session: AsyncSession,
+    span: Span,
+    project_name: str,
+) -> Optional[SpanInsertionEvent]:
+    dialect = SupportedSQLDialect(session.bind.dialect.name)
+    project_rowid = await session.scalar(
+        insert_stmt(
+            dialect=dialect,
+            table=models.Project,
+            constraint="uq_projects_name",
+            column_names=("name",),
+            values=dict(name=project_name),
+            on_conflict=OnConflict.DO_UPDATE,
+            set_=dict(name=project_name),
+        ).returning(models.Project.id)
+    )
+    assert project_rowid is not None
+    if trace := await session.scalar(
+        select(models.Trace).where(models.Trace.trace_id == span.context.trace_id)
+    ):
+        trace_rowid = trace.id
+        if span.start_time < trace.start_time or trace.end_time < span.end_time:
+            trace_start_time = min(trace.start_time, span.start_time)
+            trace_end_time = max(trace.end_time, span.end_time)
+            await session.execute(
+                update(models.Trace)
+                .where(models.Trace.id == trace_rowid)
+                .values(
+                    start_time=trace_start_time,
+                    end_time=trace_end_time,
+                )
+            )
+    else:
+        trace_rowid = cast(
+            int,
+            await session.scalar(
+                insert(models.Trace)
+                .values(
+                    project_rowid=project_rowid,
+                    trace_id=span.context.trace_id,
+                    start_time=span.start_time,
+                    end_time=span.end_time,
+                )
+                .returning(models.Trace.id)
+            ),
+        )
+    cumulative_error_count = int(span.status_code is SpanStatusCode.ERROR)
+    cumulative_llm_token_count_prompt = cast(
+        int, get_attribute_value(span.attributes, SpanAttributes.LLM_TOKEN_COUNT_PROMPT) or 0
+    )
+    cumulative_llm_token_count_completion = cast(
+        int, get_attribute_value(span.attributes, SpanAttributes.LLM_TOKEN_COUNT_COMPLETION) or 0
+    )
+    if accumulation := (
+        await session.execute(
+            select(
+                func.sum(models.Span.cumulative_error_count),
+                func.sum(models.Span.cumulative_llm_token_count_prompt),
+                func.sum(models.Span.cumulative_llm_token_count_completion),
+            ).where(models.Span.parent_id == span.context.span_id)
+        )
+    ).first():
+        cumulative_error_count += cast(int, accumulation[0] or 0)
+        cumulative_llm_token_count_prompt += cast(int, accumulation[1] or 0)
+        cumulative_llm_token_count_completion += cast(int, accumulation[2] or 0)
+    span_rowid = await session.scalar(
+        insert_stmt(
+            dialect=dialect,
+            table=models.Span,
+            constraint="uq_spans_span_id",
+            column_names=("span_id",),
+            values=dict(
+                span_id=span.context.span_id,
+                trace_rowid=trace_rowid,
+                parent_id=span.parent_id,
+                span_kind=span.span_kind.value,
+                name=span.name,
+                start_time=span.start_time,
+                end_time=span.end_time,
+                attributes=span.attributes,
+                events=[asdict(event) for event in span.events],
+                status_code=span.status_code.value,
+                status_message=span.status_message,
+                cumulative_error_count=cumulative_error_count,
+                cumulative_llm_token_count_prompt=cumulative_llm_token_count_prompt,
+                cumulative_llm_token_count_completion=cumulative_llm_token_count_completion,
+            ),
+            on_conflict=OnConflict.DO_NOTHING,
+        ).returning(models.Span.id)
+    )
+    if span_rowid is None:
+        return None
+    # Propagate cumulative values to ancestors. This is usually a no-op, since
+    # the parent usually arrives after the child. But in the event that a
+    # child arrives after its parent, we need to make sure that all the
+    # ancestors' cumulative values are updated.
+    ancestors = (
+        select(models.Span.id, models.Span.parent_id)
+        .where(models.Span.span_id == span.parent_id)
+        .cte(recursive=True)
+    )
+    child = ancestors.alias()
+    ancestors = ancestors.union_all(
+        select(models.Span.id, models.Span.parent_id).join(
+            child, models.Span.span_id == child.c.parent_id
+        )
+    )
+    await session.execute(
+        update(models.Span)
+        .where(models.Span.id.in_(select(ancestors.c.id)))
+        .values(
+            cumulative_error_count=models.Span.cumulative_error_count + cumulative_error_count,
+            cumulative_llm_token_count_prompt=models.Span.cumulative_llm_token_count_prompt
+            + cumulative_llm_token_count_prompt,
+            cumulative_llm_token_count_completion=models.Span.cumulative_llm_token_count_completion
+            + cumulative_llm_token_count_completion,
+        )
+    )
+    return SpanInsertionEvent(project_rowid)

phoenix/db/migrate.py

@@ -0,0 +1,71 @@
+import logging
+from pathlib import Path
+from queue import Empty, Queue
+from threading import Thread
+from typing import Optional
+
+from alembic import command
+from alembic.config import Config
+from sqlalchemy import URL
+
+from phoenix.exceptions import PhoenixMigrationError
+from phoenix.settings import Settings
+
+logger = logging.getLogger(__name__)
+
+
+def printif(condition: bool, text: str) -> None:
+    if condition:
+        print(text)
+
+
+def migrate(url: URL, error_queue: Optional["Queue[Exception]"] = 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)
+        alembic_cfg.set_main_option("sqlalchemy.url", str(url))
+        command.upgrade(alembic_cfg, "head")
+        printif(log_migrations, "---------------------------")
+        printif(log_migrations, "✅ Migrations complete.")
+    except Exception as e:
+        if error_queue:
+            error_queue.put(e)
+            raise e
+
+
+def migrate_in_thread(url: URL) -> 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: Queue[Exception] = Queue()
+    t = Thread(target=migrate, args=(url, 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/env.py

@@ -0,0 +1,121 @@
+import asyncio
+from logging.config import fileConfig
+
+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
+
+# Interpret the config file for Python logging.
+# This line sets up loggers basically.
+if config.config_file_name is not None:
+    fileConfig(config.config_file_name, disable_existing_loggers=False)
+
+# 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"}