arize-phoenix 10.0.4__py3-none-any.whl → 12.28.1__py3-none-any.whl

phoenix/db/bulk_inserter.py

@@ -1,34 +1,52 @@
 import asyncio
 import logging
 from asyncio import Queue, as_completed
-from collections.abc import AsyncIterator, Awaitable, Callable, Iterable
+from collections import deque
 from dataclasses import dataclass, field
 from functools import singledispatchmethod
-from itertools import islice
-from time import perf_counter
-from typing import Any, Optional, cast
+from time import perf_counter, time
+from typing import Any, AsyncIterator, Awaitable, Callable, Iterable, Optional, cast
 
+from openinference.semconv.trace import SpanAttributes
 from typing_extensions import TypeAlias
 
 import phoenix.trace.v1 as pb
+from phoenix.db import models
 from phoenix.db.insertion.constants import DEFAULT_RETRY_ALLOWANCE, DEFAULT_RETRY_DELAY_SEC
 from phoenix.db.insertion.document_annotation import DocumentAnnotationQueueInserter
 from phoenix.db.insertion.evaluation import (
     InsertEvaluationError,
     insert_evaluation,
 )
-from phoenix.db.insertion.helpers import DataManipulation, DataManipulationEvent
+from phoenix.db.insertion.helpers import (
+    DataManipulation,
+    DataManipulationEvent,
+    should_calculate_span_cost,
+)
+from phoenix.db.insertion.session_annotation import SessionAnnotationQueueInserter
 from phoenix.db.insertion.span import SpanInsertionEvent, insert_span
 from phoenix.db.insertion.span_annotation import SpanAnnotationQueueInserter
 from phoenix.db.insertion.trace_annotation import TraceAnnotationQueueInserter
 from phoenix.db.insertion.types import Insertables, Precursors
+from phoenix.server.daemons.span_cost_calculator import (
+    SpanCostCalculator,
+)
 from phoenix.server.dml_event import DmlEvent, SpanInsertEvent
+from phoenix.server.prometheus import (
+    BULK_LOADER_EVALUATION_INSERTIONS,
+    BULK_LOADER_EXCEPTIONS,
+    BULK_LOADER_LAST_ACTIVITY,
+    BULK_LOADER_SPAN_EXCEPTIONS,
+    BULK_LOADER_SPAN_INSERTION_TIME,
+    SPAN_QUEUE_SIZE,
+)
 from phoenix.server.types import CanPutItem, DbSessionFactory
 from phoenix.trace.schemas import Span
 
 logger = logging.getLogger(__name__)
 
 ProjectRowId: TypeAlias = int
+ProjectName: TypeAlias = str
 
 
 @dataclass(frozen=True)
@@ -42,12 +60,13 @@ class BulkInserter:
         db: DbSessionFactory,
         *,
         event_queue: CanPutItem[DmlEvent],
-        initial_batch_of_spans: Optional[Iterable[tuple[Span, str]]] = None,
-        initial_batch_of_evaluations: Optional[Iterable[pb.Evaluation]] = None,
+        span_cost_calculator: SpanCostCalculator,
+        initial_batch_of_spans: Iterable[tuple[Span, ProjectName]] = (),
+        initial_batch_of_evaluations: Iterable[pb.Evaluation] = (),
         sleep: float = 0.1,
         max_ops_per_transaction: int = 1000,
         max_queue_size: int = 1000,
-        enable_prometheus: bool = False,
+        max_spans_queue_size: Optional[int] = None,
         retry_delay_sec: float = DEFAULT_RETRY_DELAY_SEC,
         retry_allowance: int = DEFAULT_RETRY_ALLOWANCE,
     ) -> None:
@@ -58,7 +77,6 @@ class BulkInserter:
         :param max_ops_per_transaction: The maximum number of operations to dequeue from
         the operations queue for each transaction.
         :param max_queue_size: The maximum length of the operations queue.
-        :param enable_prometheus: Whether Prometheus is enabled.
         """
         self._db = db
         self._running = False
@@ -66,18 +84,19 @@ class BulkInserter:
         self._max_ops_per_transaction = max_ops_per_transaction
         self._operations: Optional[Queue[DataManipulation]] = None
         self._max_queue_size = max_queue_size
-        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._max_spans_queue_size = max_spans_queue_size
+        self._spans: deque[tuple[Span, ProjectName]] = deque(initial_batch_of_spans)
+        self._evaluations: deque[pb.Evaluation] = deque(initial_batch_of_evaluations)
         self._task: Optional[asyncio.Task[None]] = None
         self._event_queue = event_queue
-        self._enable_prometheus = enable_prometheus
         self._retry_delay_sec = retry_delay_sec
         self._retry_allowance = retry_allowance
         self._queue_inserters = _QueueInserters(db, self._retry_delay_sec, self._retry_allowance)
+        self._span_cost_calculator = span_cost_calculator
+
+    @property
+    def is_full(self) -> bool:
+        return bool(self._max_spans_queue_size and self._max_spans_queue_size <= len(self._spans))
 
     async def __aenter__(
         self,
@@ -91,9 +110,9 @@ class BulkInserter:
         self._operations = Queue(maxsize=self._max_queue_size)
         self._task = asyncio.create_task(self._bulk_insert())
         return (
-            self._enqueue,
-            self._queue_span,
-            self._queue_evaluation,
+            self._enqueue_annotations,
+            self._enqueue_span,
+            self._enqueue_evaluation,
             self._enqueue_operation,
         )
 
@@ -103,23 +122,22 @@ class BulkInserter:
             self._task.cancel()
             self._task = None
 
-    async def _enqueue(self, *items: Any) -> None:
+    async def _enqueue_annotations(self, *items: Any) -> None:
         await self._queue_inserters.enqueue(*items)
 
     def _enqueue_operation(self, operation: DataManipulation) -> None:
         cast("Queue[DataManipulation]", self._operations).put_nowait(operation)
 
-    async def _queue_span(self, span: Span, project_name: str) -> None:
+    async def _enqueue_span(self, span: Span, project_name: str) -> None:
         self._spans.append((span, project_name))
 
-    async def _queue_evaluation(self, evaluation: pb.Evaluation) -> None:
+    async def _enqueue_evaluation(self, evaluation: pb.Evaluation) -> None:
         self._evaluations.append(evaluation)
 
     async def _process_events(self, events: Iterable[Optional[DataManipulationEvent]]) -> None: ...
 
     async def _bulk_insert(self) -> None:
         assert isinstance(self._operations, Queue)
-        spans_buffer, evaluations_buffer = None, None
         # start first insert immediately if the inserter has not run recently
         while (
             self._running
@@ -128,6 +146,8 @@ class BulkInserter:
             or self._spans
             or self._evaluations
         ):
+            BULK_LOADER_LAST_ACTIVITY.set(time())
+            SPAN_QUEUE_SIZE.set(len(self._spans))
             if (
                 self._queue_inserters.empty
                 and self._operations.empty()
@@ -145,100 +165,100 @@ class BulkInserter:
                         async with session.begin_nested():
                             await op(session)
                     except Exception as e:
-                        if self._enable_prometheus:
-                            from phoenix.server.prometheus import BULK_LOADER_EXCEPTIONS
-
-                            BULK_LOADER_EXCEPTIONS.inc()
+                        BULK_LOADER_EXCEPTIONS.inc()
                         logger.exception(str(e))
             # 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 = []
+            num_spans_to_insert = min(self._max_ops_per_transaction, len(self._spans))
+            num_evals_to_insert = min(self._max_ops_per_transaction, len(self._evaluations))
             # Spans should be inserted before the evaluations, since an evaluation
             # insertion will fail if the span it references doesn't exist.
-            if spans_buffer:
-                await self._insert_spans(spans_buffer)
-                spans_buffer = None
-            if evaluations_buffer:
-                await self._insert_evaluations(evaluations_buffer)
-                evaluations_buffer = None
+            await self._insert_spans(num_spans_to_insert)
+            await self._insert_evaluations(num_evals_to_insert)
             async for event in self._queue_inserters.insert():
                 self._event_queue.put(event)
             await asyncio.sleep(self._sleep)
 
-    async def _insert_spans(self, spans: list[tuple[Span, str]]) -> None:
+    async def _insert_spans(self, num_spans_to_insert: int) -> None:
+        if not num_spans_to_insert or not self._spans:
+            return
         project_ids = set()
-        for i in range(0, len(spans), self._max_ops_per_transaction):
-            try:
-                start = perf_counter()
-                async with self._db() as session:
-                    for span, project_name in islice(spans, i, i + self._max_ops_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:
-                            project_ids.add(result.project_rowid)
-                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")
-        self._event_queue.put(SpanInsertEvent(tuple(project_ids)))
-
-    async def _insert_evaluations(self, evaluations: list[pb.Evaluation]) -> None:
-        for i in range(0, len(evaluations), self._max_ops_per_transaction):
-            try:
-                start = perf_counter()
-                async with self._db() as session:
-                    for evaluation in islice(evaluations, i, i + self._max_ops_per_transaction):
-                        if self._enable_prometheus:
-                            from phoenix.server.prometheus import BULK_LOADER_EVALUATION_INSERTIONS
-
-                            BULK_LOADER_EVALUATION_INSERTIONS.inc()
-                        try:
-                            async with session.begin_nested():
-                                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 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
+        span_costs: list[models.SpanCost] = []
+        try:
+            start = perf_counter()
+            async with self._db() as session:
+                while num_spans_to_insert > 0:
+                    num_spans_to_insert -= 1
+                    if not self._spans:
+                        break
+                    span, project_name = self._spans.popleft()
+                    result: Optional[SpanInsertionEvent] = None
+                    try:
+                        async with session.begin_nested():
+                            result = await insert_span(session, span, project_name)
+                    except Exception:
+                        BULK_LOADER_SPAN_EXCEPTIONS.inc()
+                        logger.exception(
+                            f"Failed to insert span with span_id={span.context.span_id}"
+                        )
+                    if result is None:
+                        continue
+                    project_ids.add(result.project_rowid)
+                    try:
+                        if not should_calculate_span_cost(span.attributes):
+                            continue
+                        span_cost = self._span_cost_calculator.calculate_cost(
+                            span.start_time,
+                            span.attributes,
+                        )
+                    except Exception:
+                        logger.exception(
+                            f"Failed to calculate span cost for span with "
+                            f"span_id={span.context.span_id}"
+                        )
+                    else:
+                        if span_cost is None:
+                            continue
+                        span_cost.span_rowid = result.span_rowid
+                        span_cost.trace_rowid = result.trace_rowid
+                        span_costs.append(span_cost)
+            BULK_LOADER_SPAN_INSERTION_TIME.observe(perf_counter() - start)
+        except Exception:
+            BULK_LOADER_SPAN_EXCEPTIONS.inc()
+            logger.exception("Failed to insert spans")
+        if project_ids:
+            self._event_queue.put(SpanInsertEvent(tuple(project_ids)))
+        if not span_costs:
+            return
+        try:
+            async with self._db() as session:
+                session.add_all(span_costs)
+        except Exception:
+            logger.exception("Failed to insert span costs")
 
-                    BULK_LOADER_EXCEPTIONS.inc()
-                logger.exception("Failed to insert evaluations")
+    async def _insert_evaluations(self, num_evals_to_insert: int) -> None:
+        if not num_evals_to_insert or not self._evaluations:
+            return
+        try:
+            async with self._db() as session:
+                while num_evals_to_insert > 0:
+                    num_evals_to_insert -= 1
+                    if not self._evaluations:
+                        break
+                    evaluation = self._evaluations.popleft()
+                    BULK_LOADER_EVALUATION_INSERTIONS.inc()
+                    try:
+                        async with session.begin_nested():
+                            await insert_evaluation(session, evaluation)
+                    except InsertEvaluationError as error:
+                        BULK_LOADER_EXCEPTIONS.inc()
+                        logger.exception(f"Failed to insert evaluation: {str(error)}")
+        except Exception:
+            BULK_LOADER_EXCEPTIONS.inc()
+            logger.exception("Failed to insert evaluations")
 
 
 class _QueueInserters:
@@ -253,10 +273,12 @@ class _QueueInserters:
         self._span_annotations = SpanAnnotationQueueInserter(*args)
         self._trace_annotations = TraceAnnotationQueueInserter(*args)
         self._document_annotations = DocumentAnnotationQueueInserter(*args)
+        self._session_annotations = SessionAnnotationQueueInserter(*args)
         self._queues = (
             self._span_annotations,
             self._trace_annotations,
             self._document_annotations,
+            self._session_annotations,
         )
 
     async def insert(self) -> AsyncIterator[DmlEvent]:
@@ -292,3 +314,23 @@ class _QueueInserters:
     @_enqueue.register(Insertables.DocumentAnnotation)
     async def _(self, item: Precursors.DocumentAnnotation) -> None:
         await self._document_annotations.enqueue(item)
+
+    @_enqueue.register(Precursors.SessionAnnotation)
+    @_enqueue.register(Insertables.SessionAnnotation)
+    async def _(self, item: Precursors.SessionAnnotation) -> None:
+        await self._session_annotations.enqueue(item)
+
+
+LLM_MODEL_NAME = SpanAttributes.LLM_MODEL_NAME
+LLM_PROVIDER = SpanAttributes.LLM_PROVIDER
+LLM_TOKEN_COUNT_COMPLETION = SpanAttributes.LLM_TOKEN_COUNT_COMPLETION
+LLM_TOKEN_COUNT_COMPLETION_DETAILS_AUDIO = SpanAttributes.LLM_TOKEN_COUNT_COMPLETION_DETAILS_AUDIO
+LLM_TOKEN_COUNT_COMPLETION_DETAILS_REASONING = (
+    SpanAttributes.LLM_TOKEN_COUNT_COMPLETION_DETAILS_REASONING
+)
+LLM_TOKEN_COUNT_PROMPT = SpanAttributes.LLM_TOKEN_COUNT_PROMPT
+LLM_TOKEN_COUNT_PROMPT_DETAILS_AUDIO = SpanAttributes.LLM_TOKEN_COUNT_PROMPT_DETAILS_AUDIO
+LLM_TOKEN_COUNT_PROMPT_DETAILS_CACHE_READ = SpanAttributes.LLM_TOKEN_COUNT_PROMPT_DETAILS_CACHE_READ
+LLM_TOKEN_COUNT_PROMPT_DETAILS_CACHE_WRITE = (
+    SpanAttributes.LLM_TOKEN_COUNT_PROMPT_DETAILS_CACHE_WRITE
+)

phoenix/db/engines.py

@@ -1,16 +1,15 @@
 from __future__ import annotations
 
 import asyncio
-import json
 import logging
 from collections.abc import Callable
-from datetime import datetime
 from enum import Enum
 from sqlite3 import Connection
-from typing import Any
+from typing import Any, Optional
 
 import aiosqlite
 import numpy as np
+import orjson
 import sqlalchemy
 import sqlean
 from sqlalchemy import URL, StaticPool, event, make_url
@@ -123,7 +122,7 @@ def aio_sqlite_engine(
             lambda: sqlean.connect(f"file:{database}", uri=True),
             iter_chunk_size=64,
         )
-        conn.daemon = True
+        conn.daemon = True  # type: ignore[attr-defined]
         return conn
 
     engine = create_async_engine(
@@ -169,43 +168,151 @@ def aio_postgresql_engine(
     log_to_stdout: bool = False,
     log_migrations_to_stdout: bool = True,
 ) -> AsyncEngine:
-    asyncpg_url, asyncpg_args = get_pg_config(url, "asyncpg")
-    engine = create_async_engine(
-        url=asyncpg_url,
-        connect_args=asyncpg_args,
-        echo=log_to_stdout,
-        json_serializer=_dumps,
+    from phoenix.config import (
+        get_env_postgres_iam_token_lifetime,
+        get_env_postgres_use_iam_auth,
     )
+
+    use_iam_auth = get_env_postgres_use_iam_auth()
+
+    asyncpg_url, asyncpg_args = get_pg_config(url, "asyncpg", enforce_ssl=use_iam_auth)
+
+    iam_config: Optional[dict[str, Any]] = None
+    token_lifetime: int = 0
+    if use_iam_auth:
+        iam_config = _extract_iam_config_from_url(url)
+        token_lifetime = get_env_postgres_iam_token_lifetime()
+
+        async def iam_async_creator() -> Any:
+            import asyncpg  # type: ignore
+
+            from phoenix.db.iam_auth import generate_aws_rds_token
+
+            assert iam_config is not None
+            token = generate_aws_rds_token(
+                host=iam_config["host"],
+                port=iam_config["port"],
+                user=iam_config["user"],
+            )
+
+            conn_kwargs = {
+                "host": iam_config["host"],
+                "port": iam_config["port"],
+                "user": iam_config["user"],
+                "password": token,
+                "database": iam_config["database"],
+            }
+
+            if asyncpg_args:
+                conn_kwargs.update(asyncpg_args)
+
+            return await asyncpg.connect(**conn_kwargs)
+
+        engine = create_async_engine(
+            url=asyncpg_url,
+            async_creator=iam_async_creator,
+            echo=log_to_stdout,
+            json_serializer=_dumps,
+            pool_recycle=token_lifetime,
+        )
+    else:
+        engine = create_async_engine(
+            url=asyncpg_url,
+            connect_args=asyncpg_args,
+            echo=log_to_stdout,
+            json_serializer=_dumps,
+        )
+
     if not migrate:
         return engine
 
-    psycopg_url, psycopg_args = get_pg_config(url, "psycopg")
-    sync_engine = sqlalchemy.create_engine(
-        url=psycopg_url,
-        connect_args=psycopg_args,
-        echo=log_migrations_to_stdout,
-        json_serializer=_dumps,
-    )
+    psycopg_url, psycopg_args = get_pg_config(url, "psycopg", enforce_ssl=use_iam_auth)
+
+    if use_iam_auth:
+        assert iam_config is not None
+
+        def iam_sync_creator() -> Any:
+            import psycopg
+
+            from phoenix.db.iam_auth import generate_aws_rds_token
+
+            token = generate_aws_rds_token(
+                host=iam_config["host"],
+                port=iam_config["port"],
+                user=iam_config["user"],
+            )
+
+            conn_kwargs = {
+                "host": iam_config["host"],
+                "port": iam_config["port"],
+                "user": iam_config["user"],
+                "password": token,
+                "dbname": iam_config["database"],
+            }
+
+            if psycopg_args:
+                conn_kwargs.update(psycopg_args)
+
+            return psycopg.connect(**conn_kwargs)
+
+        sync_engine = sqlalchemy.create_engine(
+            url=psycopg_url,
+            creator=iam_sync_creator,
+            echo=log_migrations_to_stdout,
+            json_serializer=_dumps,
+            pool_recycle=token_lifetime,
+        )
+    else:
+        sync_engine = sqlalchemy.create_engine(
+            url=psycopg_url,
+            connect_args=psycopg_args,
+            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 _extract_iam_config_from_url(url: URL) -> dict[str, Any]:
+    """Extract connection parameters needed for IAM authentication from a SQLAlchemy URL.
+
+    Args:
+        url: SQLAlchemy database URL
+
+    Returns:
+        Dictionary with host, port, user, and database
+    """
+    host = url.host
+    if not host:
+        raise ValueError("Database host is required for IAM authentication")
+
+    port = url.port or 5432
+    user = url.username
+    if not user:
+        raise ValueError("Database user is required for IAM authentication")
+
+    database = url.database or "postgres"
+
+    return {
+        "host": host,
+        "port": port,
+        "user": user,
+        "database": database,
+    }
+
+
 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)
+    return orjson.dumps(obj, default=_default).decode()
+
+
+def _default(obj: Any) -> Any:
+    if isinstance(obj, np.ndarray):
+        return obj.tolist()
+    if isinstance(obj, (np.integer, np.floating, np.bool_)):
+        return obj.item()
+    if isinstance(obj, Enum):
+        return obj.value
+    raise TypeError(f"Object of type {type(obj).__name__} is not serializable")

phoenix/db/enums.py

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