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

phoenix/db/README.md

@@ -0,0 +1,25 @@
+# Database
+
+This module is responsible for the database connection and the migrations.
+
+## Migrations
+
+All migrations are managed by Alembic. Migrations are applied to the database automatically when the application starts.
+
+### Applying migrations
+
+To manually apply the migrations, run the following command:
+
+```bash
+alembic upgrade head
+```
+
+### Creating a migration
+
+All migrations are stored in the `migrations` folder. To create a new migration, run the following command:
+
+```bash
+alembic revision -m "your_revision_name"
+```
+
+Then fill the migration file with the necessary changes.

phoenix/db/__init__.py

@@ -0,0 +1,4 @@
+from .engines import get_printable_db_url
+from .migrate import migrate
+
+__all__ = ["migrate", "get_printable_db_url"]

phoenix/db/alembic.ini

@@ -0,0 +1,119 @@
+# A generic, single database configuration.
+
+[alembic]
+# path to migration scripts
+# Note this is overridden in .migrate during programatic migrations
+script_location = migrations
+
+# template used to generate migration file names; The default value is %%(rev)s_%%(slug)s
+# Uncomment the line below if you want the files to be prepended with date and time
+# see https://alembic.sqlalchemy.org/en/latest/tutorial.html#editing-the-ini-file
+# for all available tokens
+# file_template = %%(year)d_%%(month).2d_%%(day).2d_%%(hour).2d%%(minute).2d-%%(rev)s_%%(slug)s
+
+# sys.path path, will be prepended to sys.path if present.
+# defaults to the current working directory.
+prepend_sys_path = .
+
+# timezone to use when rendering the date within the migration file
+# as well as the filename.
+# If specified, requires the python>=3.9 or backports.zoneinfo library.
+# Any required deps can installed by adding `alembic[tz]` to the pip requirements
+# string value is passed to ZoneInfo()
+# leave blank for localtime
+# timezone =
+
+# max length of characters to apply to the
+# "slug" field
+# truncate_slug_length = 40
+
+# set to 'true' to run the environment during
+# the 'revision' command, regardless of autogenerate
+# revision_environment = false
+
+# set to 'true' to allow .pyc and .pyo files without
+# a source .py file to be detected as revisions in the
+# versions/ directory
+# sourceless = false
+
+# version location specification; This defaults
+# to migrations/versions.  When using multiple version
+# directories, initial revisions must be specified with --version-path.
+# The path separator used here should be the separator specified by "version_path_separator" below.
+# version_locations = %(here)s/bar:%(here)s/bat:migrations/versions
+
+# version path separator; As mentioned above, this is the character used to split
+# version_locations. The default within new alembic.ini files is "os", which uses os.pathsep.
+# If this key is omitted entirely, it falls back to the legacy behavior of splitting on spaces and/or commas.
+# Valid values for version_path_separator are:
+#
+# version_path_separator = :
+# version_path_separator = ;
+# version_path_separator = space
+version_path_separator = os  # Use os.pathsep. Default configuration used for new projects.
+
+# set to 'true' to search source files recursively
+# in each "version_locations" directory
+# new in Alembic version 1.10
+# recursive_version_locations = false
+
+# the output encoding used when revision files
+# are written from script.py.mako
+# output_encoding = utf-8
+
+# NB: This is commented out intentionally as it is dynamic
+# See migrations/env.py
+# sqlalchemy.url = driver://user:pass@localhost/dbname
+
+
+[post_write_hooks]
+# post_write_hooks defines scripts or Python functions that are run
+# on newly generated revision scripts.  See the documentation for further
+# detail and examples
+
+# format using "black" - use the console_scripts runner, against the "black" entrypoint
+# hooks = black
+# black.type = console_scripts
+# black.entrypoint = black
+# black.options = -l 79 REVISION_SCRIPT_FILENAME
+
+# lint with attempts to fix using "ruff" - use the exec runner, execute a binary
+# hooks = ruff
+# ruff.type = exec
+# ruff.executable = %(here)s/.venv/bin/ruff
+# ruff.options = --fix REVISION_SCRIPT_FILENAME
+
+# Logging configuration
+[loggers]
+keys = root,sqlalchemy,alembic
+
+[handlers]
+keys = console
+
+[formatters]
+keys = generic
+
+[logger_root]
+level = WARN
+handlers = console
+qualname =
+
+[logger_sqlalchemy]
+level = WARN
+handlers =
+qualname = sqlalchemy.engine
+
+[logger_alembic]
+level = WARN
+handlers =
+qualname = alembic
+
+[handler_console]
+class = StreamHandler
+args = (sys.stderr,)
+level = NOTSET
+formatter = generic
+
+[formatter_generic]
+format = %(levelname)-5.5s [%(name)s] %(message)s
+datefmt = %H:%M:%S

phoenix/db/bulk_inserter.py

@@ -0,0 +1,206 @@
+import asyncio
+import logging
+from dataclasses import dataclass, field
+from datetime import datetime, timezone
+from itertools import islice
+from time import perf_counter, time
+from typing import (
+    Any,
+    AsyncContextManager,
+    Awaitable,
+    Callable,
+    Iterable,
+    List,
+    Optional,
+    Set,
+    Tuple,
+)
+
+from cachetools import LRUCache
+from sqlalchemy.ext.asyncio import AsyncSession
+from typing_extensions import TypeAlias
+
+import phoenix.trace.v1 as pb
+from phoenix.db.insertion.evaluation import (
+    EvaluationInsertionResult,
+    InsertEvaluationError,
+    insert_evaluation,
+)
+from phoenix.db.insertion.span import SpanInsertionEvent, insert_span
+from phoenix.server.api.dataloaders import CacheForDataLoaders
+from phoenix.trace.schemas import Span
+
+logger = logging.getLogger(__name__)
+
+ProjectRowId: TypeAlias = int
+
+
+@dataclass(frozen=True)
+class TransactionResult:
+    updated_project_rowids: Set[ProjectRowId] = field(default_factory=set)
+
+
+class BulkInserter:
+    def __init__(
+        self,
+        db: Callable[[], AsyncContextManager[AsyncSession]],
+        *,
+        cache_for_dataloaders: Optional[CacheForDataLoaders] = None,
+        initial_batch_of_spans: Optional[Iterable[Tuple[Span, str]]] = None,
+        initial_batch_of_evaluations: Optional[Iterable[pb.Evaluation]] = None,
+        run_interval_in_seconds: float = 2,
+        max_num_per_transaction: int = 1000,
+        enable_prometheus: bool = False,
+    ) -> None:
+        """
+        :param db: A function to initiate a new database session.
+        :param initial_batch_of_spans: Initial batch of spans to insert.
+        :param run_interval_in_seconds: The time interval between the starts of each
+        bulk insert. If there's nothing to insert, the inserter goes back to sleep.
+        :param max_num_per_transaction: The maximum number of items to insert in a single
+        transaction. Multiple transactions will be used if there are more items in the batch.
+        """
+        self._db = db
+        self._running = False
+        self._run_interval_seconds = run_interval_in_seconds
+        self._max_num_per_transaction = max_num_per_transaction
+        self._spans: List[Tuple[Span, str]] = (
+            [] if initial_batch_of_spans is None else list(initial_batch_of_spans)
+        )
+        self._evaluations: List[pb.Evaluation] = (
+            [] if initial_batch_of_evaluations is None else list(initial_batch_of_evaluations)
+        )
+        self._task: Optional[asyncio.Task[None]] = None
+        self._last_updated_at_by_project: LRUCache[ProjectRowId, datetime] = LRUCache(maxsize=100)
+        self._cache_for_dataloaders = cache_for_dataloaders
+        self._enable_prometheus = enable_prometheus
+
+    def last_updated_at(self, project_rowid: Optional[ProjectRowId] = None) -> Optional[datetime]:
+        if isinstance(project_rowid, ProjectRowId):
+            return self._last_updated_at_by_project.get(project_rowid)
+        return max(self._last_updated_at_by_project.values(), default=None)
+
+    async def __aenter__(
+        self,
+    ) -> Tuple[Callable[[Span, str], Awaitable[None]], Callable[[pb.Evaluation], Awaitable[None]]]:
+        self._running = True
+        self._task = asyncio.create_task(self._bulk_insert())
+        return self._queue_span, self._queue_evaluation
+
+    async def __aexit__(self, *args: Any) -> None:
+        self._running = False
+
+    async def _queue_span(self, span: Span, project_name: str) -> None:
+        self._spans.append((span, project_name))
+
+    async def _queue_evaluation(self, evaluation: pb.Evaluation) -> None:
+        self._evaluations.append(evaluation)
+
+    async def _bulk_insert(self) -> None:
+        spans_buffer, evaluations_buffer = None, None
+        next_run_at = time() + self._run_interval_seconds
+        while self._spans or self._evaluations or self._running:
+            await asyncio.sleep(next_run_at - time())
+            next_run_at = time() + self._run_interval_seconds
+            if not (self._spans or self._evaluations):
+                continue
+            # It's important to grab the buffers at the same time so there's
+            # no race condition, since an eval insertion will fail if the span
+            # it references doesn't exist. Grabbing the eval buffer later may
+            # include an eval whose span is in the queue but missed being
+            # included in the span buffer that was grabbed previously.
+            if self._spans:
+                spans_buffer = self._spans
+                self._spans = []
+            if self._evaluations:
+                evaluations_buffer = self._evaluations
+                self._evaluations = []
+            # Spans should be inserted before the evaluations, since an evaluation
+            # insertion will fail if the span it references doesn't exist.
+            transaction_result = TransactionResult()
+            if spans_buffer:
+                result = await self._insert_spans(spans_buffer)
+                transaction_result.updated_project_rowids.update(result.updated_project_rowids)
+                spans_buffer = None
+            if evaluations_buffer:
+                result = await self._insert_evaluations(evaluations_buffer)
+                transaction_result.updated_project_rowids.update(result.updated_project_rowids)
+                evaluations_buffer = None
+            for project_rowid in transaction_result.updated_project_rowids:
+                self._last_updated_at_by_project[project_rowid] = datetime.now(timezone.utc)
+
+    async def _insert_spans(self, spans: List[Tuple[Span, str]]) -> TransactionResult:
+        transaction_result = TransactionResult()
+        for i in range(0, len(spans), self._max_num_per_transaction):
+            try:
+                start = perf_counter()
+                async with self._db() as session:
+                    for span, project_name in islice(spans, i, i + self._max_num_per_transaction):
+                        if self._enable_prometheus:
+                            from phoenix.server.prometheus import BULK_LOADER_SPAN_INSERTIONS
+
+                            BULK_LOADER_SPAN_INSERTIONS.inc()
+                        result: Optional[SpanInsertionEvent] = None
+                        try:
+                            async with session.begin_nested():
+                                result = await insert_span(session, span, project_name)
+                        except Exception:
+                            if self._enable_prometheus:
+                                from phoenix.server.prometheus import BULK_LOADER_EXCEPTIONS
+
+                                BULK_LOADER_EXCEPTIONS.inc()
+                            logger.exception(
+                                f"Failed to insert span with span_id={span.context.span_id}"
+                            )
+                        if result is not None:
+                            transaction_result.updated_project_rowids.add(result.project_rowid)
+                            if (cache := self._cache_for_dataloaders) is not None:
+                                cache.invalidate(result)
+                if self._enable_prometheus:
+                    from phoenix.server.prometheus import BULK_LOADER_INSERTION_TIME
+
+                    BULK_LOADER_INSERTION_TIME.observe(perf_counter() - start)
+            except Exception:
+                if self._enable_prometheus:
+                    from phoenix.server.prometheus import BULK_LOADER_EXCEPTIONS
+
+                    BULK_LOADER_EXCEPTIONS.inc()
+                logger.exception("Failed to insert spans")
+        return transaction_result
+
+    async def _insert_evaluations(self, evaluations: List[pb.Evaluation]) -> TransactionResult:
+        transaction_result = TransactionResult()
+        for i in range(0, len(evaluations), self._max_num_per_transaction):
+            try:
+                start = perf_counter()
+                async with self._db() as session:
+                    for evaluation in islice(evaluations, i, i + self._max_num_per_transaction):
+                        if self._enable_prometheus:
+                            from phoenix.server.prometheus import BULK_LOADER_EVALUATION_INSERTIONS
+
+                            BULK_LOADER_EVALUATION_INSERTIONS.inc()
+                        result: Optional[EvaluationInsertionResult] = None
+                        try:
+                            async with session.begin_nested():
+                                result = await insert_evaluation(session, evaluation)
+                        except InsertEvaluationError as error:
+                            if self._enable_prometheus:
+                                from phoenix.server.prometheus import BULK_LOADER_EXCEPTIONS
+
+                                BULK_LOADER_EXCEPTIONS.inc()
+                            logger.exception(f"Failed to insert evaluation: {str(error)}")
+                        if result is not None:
+                            transaction_result.updated_project_rowids.add(result.project_rowid)
+                            if (cache := self._cache_for_dataloaders) is not None:
+                                cache.invalidate(result)
+                if self._enable_prometheus:
+                    from phoenix.server.prometheus import BULK_LOADER_INSERTION_TIME
+
+                    BULK_LOADER_INSERTION_TIME.observe(perf_counter() - start)
+            except Exception:
+                if self._enable_prometheus:
+                    from phoenix.server.prometheus import BULK_LOADER_EXCEPTIONS
+
+                    BULK_LOADER_EXCEPTIONS.inc()
+                logger.exception("Failed to insert evaluations")
+        return transaction_result

phoenix/db/engines.py

@@ -0,0 +1,152 @@
+import asyncio
+import json
+from datetime import datetime
+from enum import Enum
+from sqlite3 import Connection
+from typing import Any
+
+import aiosqlite
+import numpy as np
+import sqlean
+from sqlalchemy import URL, event, make_url
+from sqlalchemy.ext.asyncio import AsyncEngine, create_async_engine
+from typing_extensions import assert_never
+
+from phoenix.db.helpers import SupportedSQLDialect
+from phoenix.db.migrate import migrate_in_thread
+from phoenix.db.models import init_models
+
+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={"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,
+    echo: 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 backend is SupportedSQLDialect.SQLITE:
+        return aio_sqlite_engine(url=url, migrate=migrate, echo=echo)
+    elif backend is SupportedSQLDialect.POSTGRESQL:
+        return aio_postgresql_engine(url=url, migrate=migrate, echo=echo)
+    else:
+        assert_never(backend)
+
+
+def aio_sqlite_engine(
+    url: URL,
+    migrate: bool = True,
+    echo: bool = False,
+    shared_cache: 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=echo,
+        json_serializer=_dumps,
+        async_creator=async_creator,
+    )
+    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:
+        migrate_in_thread(engine.url)
+    return engine
+
+
+def aio_postgresql_engine(
+    url: URL,
+    migrate: bool = True,
+    echo: bool = False,
+) -> AsyncEngine:
+    engine = create_async_engine(url=url, echo=echo, json_serializer=_dumps)
+    if not migrate:
+        return engine
+    migrate_in_thread(engine.url)
+    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/helpers.py

@@ -0,0 +1,47 @@
+from enum import Enum
+from typing import Any
+
+from openinference.semconv.trace import (
+    OpenInferenceSpanKindValues,
+    RerankerAttributes,
+    SpanAttributes,
+)
+from sqlalchemy import Integer, SQLColumnExpression, case, func
+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(".")