arize-phoenix 4.14.1__tar.gz → 4.16.0__tar.gz

{arize_phoenix-4.14.1 → arize_phoenix-4.16.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: arize-phoenix
-Version: 4.14.1
+Version: 4.16.0
 Summary: AI Observability and Evaluation
 Project-URL: Documentation, https://docs.arize.com/phoenix/
 Project-URL: Issues, https://github.com/Arize-ai/phoenix/issues
@@ -22,6 +22,7 @@ Requires-Dist: aiosqlite
 Requires-Dist: alembic<2,>=1.3.0
 Requires-Dist: arize-phoenix-evals>=0.13.1
 Requires-Dist: cachetools
+Requires-Dist: fastapi
 Requires-Dist: grpcio
 Requires-Dist: hdbscan>=0.8.33
 Requires-Dist: httpx
@@ -40,8 +41,8 @@ Requires-Dist: pandas>=1.0
 Requires-Dist: protobuf<6.0,>=3.20
 Requires-Dist: psutil
 Requires-Dist: pyarrow
+Requires-Dist: pydantic!=2.0.*,<3,>=1.0
 Requires-Dist: python-multipart
-Requires-Dist: pyyaml
 Requires-Dist: scikit-learn
 Requires-Dist: scipy
 Requires-Dist: sqlalchemy[asyncio]<3,>=2.0.4
@@ -56,9 +57,9 @@ Requires-Dist: uvicorn
 Requires-Dist: wrapt
 Provides-Extra: container
 Requires-Dist: opentelemetry-exporter-otlp; extra == 'container'
+Requires-Dist: opentelemetry-instrumentation-fastapi; extra == 'container'
 Requires-Dist: opentelemetry-instrumentation-grpc; extra == 'container'
 Requires-Dist: opentelemetry-instrumentation-sqlalchemy; extra == 'container'
-Requires-Dist: opentelemetry-instrumentation-starlette; extra == 'container'
 Requires-Dist: opentelemetry-proto>=1.12.0; extra == 'container'
 Requires-Dist: opentelemetry-sdk; extra == 'container'
 Requires-Dist: opentelemetry-semantic-conventions; extra == 'container'
@@ -71,6 +72,7 @@ Requires-Dist: anthropic; extra == 'dev'
 Requires-Dist: arize[autoembeddings,llm-evaluation]; extra == 'dev'
 Requires-Dist: asgi-lifespan; extra == 'dev'
 Requires-Dist: asyncpg; extra == 'dev'
+Requires-Dist: faker>=26.0.0; extra == 'dev'
 Requires-Dist: gcsfs; extra == 'dev'
 Requires-Dist: google-cloud-aiplatform>=1.3; extra == 'dev'
 Requires-Dist: hatch; extra == 'dev'

{arize_phoenix-4.14.1 → arize_phoenix-4.16.0}/pyproject.toml

@@ -59,7 +59,8 @@ dependencies = [
   "cachetools",
   "python-multipart",  # see https://www.starlette.io/#dependencies
   "arize-phoenix-evals>=0.13.1",
-  "pyyaml", # for OpenAPI
+  "fastapi",
+  "pydantic>=1.0,!=2.0.*,<3,", # exclude 2.0.* since it does not support the `json_encoders` configuration setting
 ]
 dynamic = ["version"]
 
@@ -92,6 +93,7 @@ dev = [
   "anthropic",
   "prometheus_client",
   "asgi-lifespan",
+  "Faker>=26.0.0",
 ]
 evals = []
 experimental = []
@@ -111,7 +113,7 @@ container = [
   "opentelemetry-proto>=1.12.0",
   "opentelemetry-exporter-otlp",
   "opentelemetry-semantic-conventions",
-  "opentelemetry-instrumentation-starlette",
+  "opentelemetry-instrumentation-fastapi",
   "opentelemetry-instrumentation-sqlalchemy",
   "opentelemetry-instrumentation-grpc",
   "py-grpc-prometheus",
@@ -161,13 +163,16 @@ dependencies = [
   "protobuf==3.20",  # version minimum (for tests)
   "responses",
   "tiktoken",
-  "typing-extensions==4.5.0; python_version<'3.12'",
-  "typing-extensions==4.6.0; python_version>='3.12'",
+  "typing-extensions==4.5.0; python_version=='3.8'",
+  "typing-extensions==4.6.0; python_version=='3.9'",
+  "pydantic==1.9.0; python_version<='3.9'",  # minimum version of pydantic compatible with openai
+  "pydantic==2.8.2; python_version=='3.12'",
   "httpx", # For OpenAI testing
   "respx", # For OpenAI testing
   "nest-asyncio", # for executor testing
   "astunparse; python_version<'3.9'",  # `ast.unparse(...)` is only available starting with Python 3.9
   "asgi-lifespan",
+  "Faker>=26.0.0",
 ]
 
 [tool.hatch.envs.type]
@@ -190,12 +195,15 @@ dependencies = [
   "opentelemetry-proto>=1.12.0",
   "opentelemetry-exporter-otlp",
   "opentelemetry-semantic-conventions",
-  "opentelemetry-instrumentation-starlette",
+  "opentelemetry-instrumentation-fastapi",
   "opentelemetry-instrumentation-sqlalchemy",
   "opentelemetry-instrumentation-grpc",
   "py-grpc-prometheus",
   "strawberry-graphql[opentelemetry]==0.236.0",  # need to pin version because we're monkey-patching
   "requests",  # this is needed to type-check third-party packages
+  "pydantic==1.10.17; python_version=='3.8'",  # lower minor versions of pydantic break strawberry mypy plugin
+  "pydantic==1.10.17; python_version=='3.9'",  # lower minor versions of pydantic break strawberry mypy plugin
+  "pydantic==2.8.2; python_version=='3.12'",
 ]
 
 [[tool.hatch.envs.type.matrix]]
@@ -310,11 +318,12 @@ build = 'strawberry export-schema phoenix.server.api.schema:schema > app/schema.
 
 [tool.hatch.envs.openapi]
 dependencies = [
-  "pyyaml",
+  "pydantic==2.8.2",
+  "fastapi==0.111.0",
 ]
 
 [tool.hatch.envs.openapi.scripts]
-build = 'python -m phoenix.server.api.openapi.main > schemas/openapi.yaml'
+build = "python -m phoenix.server.api.openapi.main > schemas/openapi.json"
 
 [tool.hatch.envs.proto]
 detached = true
@@ -379,6 +388,7 @@ module = [
   "sqlean",
   "grpc.*",
   "py_grpc_prometheus.*",
+  "orjson",  # suppress fastapi internal type errors
 ]
 ignore_missing_imports = true
 

{arize_phoenix-4.14.1 → arize_phoenix-4.16.0}/src/phoenix/db/bulk_inserter.py

@@ -1,28 +1,36 @@
 import asyncio
 import logging
-from asyncio import Queue
+from asyncio import Queue, as_completed
+from collections import defaultdict
 from dataclasses import dataclass, field
 from datetime import datetime, timezone
+from functools import singledispatchmethod
 from itertools import islice
 from time import perf_counter
 from typing import (
     Any,
-    AsyncContextManager,
     Awaitable,
     Callable,
+    DefaultDict,
+    Dict,
     Iterable,
     List,
+    Mapping,
     Optional,
     Set,
     Tuple,
+    Type,
     cast,
 )
 
 from cachetools import LRUCache
-from sqlalchemy.ext.asyncio import AsyncSession
+from sqlalchemy import Select, select
 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 (
     EvaluationInsertionEvent,
     InsertEvaluationError,
@@ -30,7 +38,11 @@ from phoenix.db.insertion.evaluation import (
 )
 from phoenix.db.insertion.helpers import DataManipulation, DataManipulationEvent
 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.api.dataloaders import CacheForDataLoaders
+from phoenix.server.types import DbSessionFactory
 from phoenix.trace.schemas import Span
 
 logger = logging.getLogger(__name__)
@@ -46,7 +58,7 @@ class TransactionResult:
 class BulkInserter:
     def __init__(
         self,
-        db: Callable[[], AsyncContextManager[AsyncSession]],
+        db: DbSessionFactory,
         *,
         cache_for_dataloaders: Optional[CacheForDataLoaders] = None,
         initial_batch_of_operations: Iterable[DataManipulation] = (),
@@ -56,6 +68,8 @@ class BulkInserter:
         max_ops_per_transaction: int = 1000,
         max_queue_size: int = 1000,
         enable_prometheus: bool = False,
+        retry_delay_sec: float = DEFAULT_RETRY_DELAY_SEC,
+        retry_allowance: int = DEFAULT_RETRY_ALLOWANCE,
     ) -> None:
         """
         :param db: A function to initiate a new database session.
@@ -82,6 +96,9 @@ class BulkInserter:
         self._last_updated_at_by_project: LRUCache[ProjectRowId, datetime] = LRUCache(maxsize=100)
         self._cache_for_dataloaders = cache_for_dataloaders
         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)
 
     def last_updated_at(self, project_rowid: Optional[ProjectRowId] = None) -> Optional[datetime]:
         if isinstance(project_rowid, ProjectRowId):
@@ -91,6 +108,7 @@ class BulkInserter:
     async def __aenter__(
         self,
     ) -> Tuple[
+        Callable[[Any], Awaitable[None]],
         Callable[[Span, str], Awaitable[None]],
         Callable[[pb.Evaluation], Awaitable[None]],
         Callable[[DataManipulation], None],
@@ -99,6 +117,7 @@ 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_operation,
@@ -110,6 +129,9 @@ class BulkInserter:
             self._task.cancel()
             self._task = None
 
+    async def _enqueue(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)
 
@@ -125,7 +147,17 @@ class BulkInserter:
         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 or not self._operations.empty() or self._spans or self._evaluations:
+        while (
+            self._running
+            or not self._queue_inserters.empty
+            or not self._operations.empty()
+            or self._spans
+            or self._evaluations
+        ):
+            if not self._queue_inserters.empty:
+                if inserted_ids := await self._queue_inserters.insert():
+                    for project_rowid in await self._get_project_rowids(inserted_ids):
+                        self._last_updated_at_by_project[project_rowid] = datetime.now(timezone.utc)
             if self._operations.empty() and not (self._spans or self._evaluations):
                 await asyncio.sleep(self._sleep)
                 continue
@@ -245,3 +277,97 @@ class BulkInserter:
                     BULK_LOADER_EXCEPTIONS.inc()
                 logger.exception("Failed to insert evaluations")
         return transaction_result
+
+    async def _get_project_rowids(
+        self,
+        inserted_ids: Mapping[Type[models.Base], List[int]],
+    ) -> Set[int]:
+        ans: Set[int] = set()
+        if not inserted_ids:
+            return ans
+        stmt: Select[Tuple[int]]
+        for table, ids in inserted_ids.items():
+            if not ids:
+                continue
+            if issubclass(table, models.SpanAnnotation):
+                stmt = (
+                    select(models.Project.id)
+                    .join(models.Trace)
+                    .join_from(models.Trace, models.Span)
+                    .join_from(models.Span, models.SpanAnnotation)
+                    .where(models.SpanAnnotation.id.in_(ids))
+                )
+            elif issubclass(table, models.DocumentAnnotation):
+                stmt = (
+                    select(models.Project.id)
+                    .join(models.Trace)
+                    .join_from(models.Trace, models.Span)
+                    .join_from(models.Span, models.DocumentAnnotation)
+                    .where(models.DocumentAnnotation.id.in_(ids))
+                )
+            elif issubclass(table, models.TraceAnnotation):
+                stmt = (
+                    select(models.Project.id)
+                    .join(models.Trace)
+                    .join_from(models.Trace, models.TraceAnnotation)
+                    .where(models.TraceAnnotation.id.in_(ids))
+                )
+            else:
+                continue
+            async with self._db() as session:
+                project_rowids = [_ async for _ in await session.stream_scalars(stmt)]
+                ans.update(project_rowids)
+        return ans
+
+
+class _QueueInserters:
+    def __init__(
+        self,
+        db: DbSessionFactory,
+        retry_delay_sec: float = DEFAULT_RETRY_DELAY_SEC,
+        retry_allowance: int = DEFAULT_RETRY_ALLOWANCE,
+    ) -> None:
+        self._db = db
+        args = (db, retry_delay_sec, retry_allowance)
+        self._span_annotations = SpanAnnotationQueueInserter(*args)
+        self._trace_annotations = TraceAnnotationQueueInserter(*args)
+        self._document_annotations = DocumentAnnotationQueueInserter(*args)
+        self._queues = (
+            self._span_annotations,
+            self._trace_annotations,
+            self._document_annotations,
+        )
+
+    async def insert(self) -> Dict[Type[models.Base], List[int]]:
+        ans: DefaultDict[Type[models.Base], List[int]] = defaultdict(list)
+        for coro in as_completed([q.insert() for q in self._queues]):
+            table, inserted_ids = await coro
+            if inserted_ids:
+                ans[table].extend(inserted_ids)
+        return ans
+
+    @property
+    def empty(self) -> bool:
+        return all(q.empty for q in self._queues)
+
+    async def enqueue(self, *items: Any) -> None:
+        for item in items:
+            await self._enqueue(item)
+
+    @singledispatchmethod
+    async def _enqueue(self, item: Any) -> None: ...
+
+    @_enqueue.register(Precursors.SpanAnnotation)
+    @_enqueue.register(Insertables.SpanAnnotation)
+    async def _(self, item: Precursors.SpanAnnotation) -> None:
+        await self._span_annotations.enqueue(item)
+
+    @_enqueue.register(Precursors.TraceAnnotation)
+    @_enqueue.register(Insertables.TraceAnnotation)
+    async def _(self, item: Precursors.TraceAnnotation) -> None:
+        await self._trace_annotations.enqueue(item)
+
+    @_enqueue.register(Precursors.DocumentAnnotation)
+    @_enqueue.register(Insertables.DocumentAnnotation)
+    async def _(self, item: Precursors.DocumentAnnotation) -> None:
+        await self._document_annotations.enqueue(item)

{arize_phoenix-4.14.1 → arize_phoenix-4.16.0}/src/phoenix/db/engines.py

@@ -8,7 +8,7 @@ from typing import Any
 import aiosqlite
 import numpy as np
 import sqlean
-from sqlalchemy import URL, event, make_url
+from sqlalchemy import URL, StaticPool, event, make_url
 from sqlalchemy.ext.asyncio import AsyncEngine, create_async_engine
 from typing_extensions import assert_never
 
@@ -105,6 +105,7 @@ def aio_sqlite_engine(
         echo=echo,
         json_serializer=_dumps,
         async_creator=async_creator,
+        poolclass=StaticPool,
     )
     event.listen(engine.sync_engine, "connect", set_sqlite_pragma)
     if not migrate:

{arize_phoenix-4.14.1 → arize_phoenix-4.16.0}/src/phoenix/db/helpers.py

@@ -1,5 +1,5 @@
 from enum import Enum
-from typing import Any, Optional, Tuple
+from typing import Any, Callable, Hashable, Iterable, List, Optional, Set, Tuple, TypeVar
 
 from openinference.semconv.trace import (
     OpenInferenceSpanKindValues,
@@ -80,3 +80,25 @@ def get_project_names_for_experiments(*experiment_ids: int) -> Select[Tuple[Opti
         .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

arize_phoenix-4.16.0/src/phoenix/db/insertion/constants.py

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

arize_phoenix-4.16.0/src/phoenix/db/insertion/document_annotation.py

@@ -0,0 +1,157 @@
+from datetime import datetime
+from typing import Any, List, Mapping, NamedTuple, Optional, Tuple
+
+from sqlalchemy import Row, Select, and_, select, tuple_
+from sqlalchemy.ext.asyncio import AsyncSession
+from typing_extensions import TypeAlias
+
+from phoenix.db import models
+from phoenix.db.helpers import dedup, num_docs_col
+from phoenix.db.insertion.types import (
+    Insertables,
+    Postponed,
+    Precursors,
+    QueueInserter,
+    Received,
+)
+
+_Name: TypeAlias = str
+_SpanId: TypeAlias = str
+_SpanRowId: TypeAlias = int
+_DocumentPosition: TypeAlias = int
+_AnnoRowId: TypeAlias = int
+_NumDocs: TypeAlias = int
+
+_Key: TypeAlias = Tuple[_Name, _SpanId, _DocumentPosition]
+_UniqueBy: TypeAlias = Tuple[_Name, _SpanRowId, _DocumentPosition]
+_Existing: TypeAlias = Tuple[
+    _SpanRowId,
+    _SpanId,
+    _NumDocs,
+    Optional[_AnnoRowId],
+    Optional[_Name],
+    Optional[_DocumentPosition],
+    Optional[datetime],
+]
+
+
+class DocumentAnnotationQueueInserter(
+    QueueInserter[
+        Precursors.DocumentAnnotation,
+        Insertables.DocumentAnnotation,
+        models.DocumentAnnotation,
+    ],
+    table=models.DocumentAnnotation,
+    unique_by=("name", "span_rowid", "document_position"),
+):
+    async def _partition(
+        self,
+        session: AsyncSession,
+        *parcels: Received[Precursors.DocumentAnnotation],
+    ) -> Tuple[
+        List[Received[Insertables.DocumentAnnotation]],
+        List[Postponed[Precursors.DocumentAnnotation]],
+        List[Received[Precursors.DocumentAnnotation]],
+    ]:
+        to_insert: List[Received[Insertables.DocumentAnnotation]] = []
+        to_postpone: List[Postponed[Precursors.DocumentAnnotation]] = []
+        to_discard: List[Received[Precursors.DocumentAnnotation]] = []
+
+        stmt = self._select_existing(*map(_key, parcels))
+        existing: List[Row[_Existing]] = [_ async for _ in await session.stream(stmt)]
+        existing_spans: Mapping[str, _SpanAttr] = {
+            e.span_id: _SpanAttr(e.span_rowid, e.num_docs) for e in existing
+        }
+        existing_annos: Mapping[_Key, _AnnoAttr] = {
+            (e.name, e.span_id, e.document_position): _AnnoAttr(e.span_rowid, e.id, e.updated_at)
+            for e in existing
+            if e.id is not None
+            and e.name is not None
+            and e.document_position is not None
+            and e.updated_at is not None
+        }
+
+        for p in parcels:
+            if (anno := existing_annos.get(_key(p))) is not None:
+                if p.received_at <= anno.updated_at:
+                    to_discard.append(p)
+                else:
+                    to_insert.append(
+                        Received(
+                            received_at=p.received_at,
+                            item=p.item.as_insertable(
+                                span_rowid=anno.span_rowid,
+                                id_=anno.id_,
+                            ),
+                        )
+                    )
+            elif (span := existing_spans.get(p.item.span_id)) is not None:
+                if 0 <= p.item.document_position < span.num_docs:
+                    to_insert.append(
+                        Received(
+                            received_at=p.received_at,
+                            item=p.item.as_insertable(
+                                span_rowid=span.span_rowid,
+                            ),
+                        )
+                    )
+                else:
+                    to_discard.append(p)
+            elif isinstance(p, Postponed):
+                if p.retries_left > 1:
+                    to_postpone.append(p.postpone(p.retries_left - 1))
+                else:
+                    to_discard.append(p)
+            elif isinstance(p, Received):
+                to_postpone.append(p.postpone(self._retry_allowance))
+            else:
+                to_discard.append(p)
+
+        assert len(to_insert) + len(to_postpone) + len(to_discard) == len(parcels)
+        to_insert = dedup(sorted(to_insert, key=_time, reverse=True), _unique_by)[::-1]
+        return to_insert, to_postpone, to_discard
+
+    def _select_existing(self, *keys: _Key) -> Select[_Existing]:
+        anno = self.table
+        span = (
+            select(models.Span.id, models.Span.span_id, num_docs_col(self._db.dialect))
+            .where(models.Span.span_id.in_({span_id for _, span_id, *_ in keys}))
+            .cte()
+        )
+        onclause = and_(
+            span.c.id == anno.span_rowid,
+            anno.name.in_({name for name, *_ in keys}),
+            tuple_(anno.name, span.c.span_id, anno.document_position).in_(keys),
+        )
+        return select(
+            span.c.id.label("span_rowid"),
+            span.c.span_id,
+            span.c.num_docs,
+            anno.id,
+            anno.name,
+            anno.document_position,
+            anno.updated_at,
+        ).outerjoin_from(span, anno, onclause)
+
+
+class _SpanAttr(NamedTuple):
+    span_rowid: _SpanRowId
+    num_docs: _NumDocs
+
+
+class _AnnoAttr(NamedTuple):
+    span_rowid: _SpanRowId
+    id_: _AnnoRowId
+    updated_at: datetime
+
+
+def _key(p: Received[Precursors.DocumentAnnotation]) -> _Key:
+    return p.item.obj.name, p.item.span_id, p.item.document_position
+
+
+def _unique_by(p: Received[Insertables.DocumentAnnotation]) -> _UniqueBy:
+    return p.item.obj.name, p.item.span_rowid, p.item.document_position
+
+
+def _time(p: Received[Any]) -> datetime:
+    return p.received_at

{arize_phoenix-4.14.1 → arize_phoenix-4.16.0}/src/phoenix/db/insertion/helpers.py

@@ -20,6 +20,7 @@ from sqlalchemy.ext.asyncio import AsyncSession
 from sqlalchemy.sql.elements import KeyedColumnElement
 from typing_extensions import TypeAlias, assert_never
 
+from phoenix.db import models
 from phoenix.db.helpers import SupportedSQLDialect
 from phoenix.db.models import Base
 
@@ -93,3 +94,15 @@ def _clean(
             yield "metadata", v
         else:
             yield k, v
+
+
+def as_kv(obj: models.Base) -> Iterator[Tuple[str, Any]]:
+    for k, c in obj.__table__.c.items():
+        if k in ["created_at", "updated_at"]:
+            continue
+        k = "metadata_" if k == "metadata" else k
+        v = getattr(obj, k, None)
+        if c.primary_key and v is None:
+            # postgresql disallows None for primary key
+            continue
+        yield k, v