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

phoenix/db/insertion/span.py

@@ -0,0 +1,193 @@
+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_on_conflict
+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)
+    if (
+        project_rowid := await session.scalar(
+            select(models.Project.id).filter_by(name=project_name)
+        )
+    ) is None:
+        project_rowid = await session.scalar(
+            insert(models.Project).values(name=project_name).returning(models.Project.id)
+        )
+    assert project_rowid is not None
+
+    trace_id = span.context.trace_id
+    trace: models.Trace = await session.scalar(
+        select(models.Trace).filter_by(trace_id=trace_id)
+    ) or models.Trace(trace_id=trace_id)
+
+    if trace.id is not None:
+        # Trace record may need to be updated.
+        if trace.end_time < span.end_time:
+            trace.end_time = span.end_time
+            trace.project_rowid = project_rowid
+        if span.start_time < trace.start_time:
+            trace.start_time = span.start_time
+    else:
+        # Trace record needs to be persisted for the first time.
+        trace.start_time = span.start_time
+        trace.end_time = span.end_time
+        trace.project_rowid = project_rowid
+        session.add(trace)
+
+    session_id = get_attribute_value(span.attributes, SpanAttributes.SESSION_ID)
+    session_id = str(session_id).strip() if session_id is not None else ""
+    assert isinstance(session_id, str)
+
+    project_session: Optional[models.ProjectSession] = None
+    if trace.project_session_rowid is not None:
+        # ProjectSession record already exists in database for this Trace record, so we fetch
+        # it because it may need to be updated. However, the session_id on the span, if exists,
+        # will be ignored at this point. Otherwise, if session_id is different, we will need
+        # to create a new ProjectSession record, as well as to determine whether the old record
+        # needs to be deleted if this is the last Trace associated with it.
+        project_session = await session.scalar(
+            select(models.ProjectSession).filter_by(id=trace.project_session_rowid)
+        )
+    elif session_id:
+        project_session = await session.scalar(
+            select(models.ProjectSession).filter_by(session_id=session_id)
+        ) or models.ProjectSession(session_id=session_id)
+
+    if project_session is not None:
+        if project_session.id is None:
+            # ProjectSession record needs to be persisted for the first time.
+            project_session.start_time = trace.start_time
+            project_session.end_time = trace.end_time
+            project_session.project_id = project_rowid
+            session.add(project_session)
+            await session.flush()
+            assert project_session.id is not None
+            trace.project_session_rowid = project_session.id
+        else:
+            # ProjectSession record may need to be updated.
+            if trace.project_session_rowid is None:
+                trace.project_session_rowid = project_session.id
+            if trace.start_time < project_session.start_time:
+                project_session.start_time = trace.start_time
+            if project_session.end_time < trace.end_time:
+                project_session.end_time = trace.end_time
+
+    await session.flush()
+    assert trace.id is not None
+    assert project_session is None or (
+        project_session.id is not None and project_session.id == trace.project_session_rowid
+    )
+
+    cumulative_error_count = int(span.status_code is SpanStatusCode.ERROR)
+    try:
+        cumulative_llm_token_count_prompt = int(
+            get_attribute_value(span.attributes, SpanAttributes.LLM_TOKEN_COUNT_PROMPT) or 0
+        )
+    except BaseException:
+        cumulative_llm_token_count_prompt = 0
+    try:
+        cumulative_llm_token_count_completion = int(
+            get_attribute_value(span.attributes, SpanAttributes.LLM_TOKEN_COUNT_COMPLETION) or 0
+        )
+    except BaseException:
+        cumulative_llm_token_count_completion = 0
+    try:
+        llm_token_count_prompt = int(
+            get_attribute_value(span.attributes, SpanAttributes.LLM_TOKEN_COUNT_PROMPT) or 0
+        )
+    except BaseException:
+        llm_token_count_prompt = 0
+    try:
+        llm_token_count_completion = int(
+            get_attribute_value(span.attributes, SpanAttributes.LLM_TOKEN_COUNT_COMPLETION) or 0
+        )
+    except BaseException:
+        llm_token_count_completion = 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_on_conflict(
+            dict(
+                span_id=span.context.span_id,
+                trace_rowid=trace.id,
+                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,
+                llm_token_count_prompt=llm_token_count_prompt,
+                llm_token_count_completion=llm_token_count_completion,
+            ),
+            dialect=dialect,
+            table=models.Span,
+            unique_by=("span_id",),
+            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/insertion/span_annotation.py

@@ -0,0 +1,158 @@
+from collections.abc import Mapping
+from datetime import datetime
+from typing import Any, NamedTuple, Optional
+
+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
+from phoenix.db.insertion.helpers import as_kv
+from phoenix.db.insertion.types import (
+    Insertables,
+    Postponed,
+    Precursors,
+    QueueInserter,
+    Received,
+)
+from phoenix.server.dml_event import SpanAnnotationDmlEvent
+
+_Name: TypeAlias = str
+_SpanId: TypeAlias = str
+_SpanRowId: TypeAlias = int
+_AnnoRowId: TypeAlias = int
+
+_Key: TypeAlias = tuple[_Name, _SpanId]
+_UniqueBy: TypeAlias = tuple[_Name, _SpanRowId]
+_Existing: TypeAlias = tuple[
+    _SpanRowId,
+    _SpanId,
+    Optional[_AnnoRowId],
+    Optional[_Name],
+    Optional[datetime],
+]
+
+
+class SpanAnnotationQueueInserter(
+    QueueInserter[
+        Precursors.SpanAnnotation,
+        Insertables.SpanAnnotation,
+        models.SpanAnnotation,
+        SpanAnnotationDmlEvent,
+    ],
+    table=models.SpanAnnotation,
+    unique_by=("name", "span_rowid"),
+):
+    async def _events(
+        self,
+        session: AsyncSession,
+        *insertions: Insertables.SpanAnnotation,
+    ) -> list[SpanAnnotationDmlEvent]:
+        records = [dict(as_kv(ins.row)) for ins in insertions]
+        stmt = self._insert_on_conflict(*records).returning(self.table.id)
+        ids = tuple([_ async for _ in await session.stream_scalars(stmt)])
+        return [SpanAnnotationDmlEvent(ids)]
+
+    async def _partition(
+        self,
+        session: AsyncSession,
+        *parcels: Received[Precursors.SpanAnnotation],
+    ) -> tuple[
+        list[Received[Insertables.SpanAnnotation]],
+        list[Postponed[Precursors.SpanAnnotation]],
+        list[Received[Precursors.SpanAnnotation]],
+    ]:
+        to_insert: list[Received[Insertables.SpanAnnotation]] = []
+        to_postpone: list[Postponed[Precursors.SpanAnnotation]] = []
+        to_discard: list[Received[Precursors.SpanAnnotation]] = []
+
+        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) for e in existing
+        }
+        existing_annos: Mapping[_Key, _AnnoAttr] = {
+            (e.name, e.span_id): _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.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:
+                to_insert.append(
+                    Received(
+                        received_at=p.received_at,
+                        item=p.item.as_insertable(
+                            span_rowid=span.span_rowid,
+                        ),
+                    )
+                )
+            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)
+            .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).in_(keys),
+        )
+        return select(
+            span.c.id.label("span_rowid"),
+            span.c.span_id,
+            anno.id,
+            anno.name,
+            anno.updated_at,
+        ).outerjoin_from(span, anno, onclause)
+
+
+class _SpanAttr(NamedTuple):
+    span_rowid: _SpanRowId
+
+
+class _AnnoAttr(NamedTuple):
+    span_rowid: _SpanRowId
+    id_: _AnnoRowId
+    updated_at: datetime
+
+
+def _key(p: Received[Precursors.SpanAnnotation]) -> _Key:
+    return p.item.obj.name, p.item.span_id
+
+
+def _unique_by(p: Received[Insertables.SpanAnnotation]) -> _UniqueBy:
+    return p.item.obj.name, p.item.span_rowid
+
+
+def _time(p: Received[Any]) -> datetime:
+    return p.received_at

phoenix/db/insertion/trace_annotation.py

@@ -0,0 +1,158 @@
+from collections.abc import Mapping
+from datetime import datetime
+from typing import Any, NamedTuple, Optional
+
+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
+from phoenix.db.insertion.helpers import as_kv
+from phoenix.db.insertion.types import (
+    Insertables,
+    Postponed,
+    Precursors,
+    QueueInserter,
+    Received,
+)
+from phoenix.server.dml_event import TraceAnnotationDmlEvent
+
+_Name: TypeAlias = str
+_TraceId: TypeAlias = str
+_TraceRowId: TypeAlias = int
+_AnnoRowId: TypeAlias = int
+
+_Key: TypeAlias = tuple[_Name, _TraceId]
+_UniqueBy: TypeAlias = tuple[_Name, _TraceRowId]
+_Existing: TypeAlias = tuple[
+    _TraceRowId,
+    _TraceId,
+    Optional[_AnnoRowId],
+    Optional[_Name],
+    Optional[datetime],
+]
+
+
+class TraceAnnotationQueueInserter(
+    QueueInserter[
+        Precursors.TraceAnnotation,
+        Insertables.TraceAnnotation,
+        models.TraceAnnotation,
+        TraceAnnotationDmlEvent,
+    ],
+    table=models.TraceAnnotation,
+    unique_by=("name", "trace_rowid"),
+):
+    async def _events(
+        self,
+        session: AsyncSession,
+        *insertions: Insertables.TraceAnnotation,
+    ) -> list[TraceAnnotationDmlEvent]:
+        records = [dict(as_kv(ins.row)) for ins in insertions]
+        stmt = self._insert_on_conflict(*records).returning(self.table.id)
+        ids = tuple([_ async for _ in await session.stream_scalars(stmt)])
+        return [TraceAnnotationDmlEvent(ids)]
+
+    async def _partition(
+        self,
+        session: AsyncSession,
+        *parcels: Received[Precursors.TraceAnnotation],
+    ) -> tuple[
+        list[Received[Insertables.TraceAnnotation]],
+        list[Postponed[Precursors.TraceAnnotation]],
+        list[Received[Precursors.TraceAnnotation]],
+    ]:
+        to_insert: list[Received[Insertables.TraceAnnotation]] = []
+        to_postpone: list[Postponed[Precursors.TraceAnnotation]] = []
+        to_discard: list[Received[Precursors.TraceAnnotation]] = []
+
+        stmt = self._select_existing(*map(_key, parcels))
+        existing: list[Row[_Existing]] = [_ async for _ in await session.stream(stmt)]
+        existing_traces: Mapping[str, _TraceAttr] = {
+            e.trace_id: _TraceAttr(e.trace_rowid) for e in existing
+        }
+        existing_annos: Mapping[_Key, _AnnoAttr] = {
+            (e.name, e.trace_id): _AnnoAttr(e.trace_rowid, e.id, e.updated_at)
+            for e in existing
+            if e.id is not None and e.name 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(
+                                trace_rowid=anno.trace_rowid,
+                                id_=anno.id_,
+                            ),
+                        )
+                    )
+            elif (trace := existing_traces.get(p.item.trace_id)) is not None:
+                to_insert.append(
+                    Received(
+                        received_at=p.received_at,
+                        item=p.item.as_insertable(
+                            trace_rowid=trace.trace_rowid,
+                        ),
+                    )
+                )
+            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
+        trace = (
+            select(models.Trace.id, models.Trace.trace_id)
+            .where(models.Trace.trace_id.in_({trace_id for _, trace_id in keys}))
+            .cte()
+        )
+        onclause = and_(
+            trace.c.id == anno.trace_rowid,
+            anno.name.in_({name for name, _ in keys}),
+            tuple_(anno.name, trace.c.trace_id).in_(keys),
+        )
+        return select(
+            trace.c.id.label("trace_rowid"),
+            trace.c.trace_id,
+            anno.id,
+            anno.name,
+            anno.updated_at,
+        ).outerjoin_from(trace, anno, onclause)
+
+
+class _TraceAttr(NamedTuple):
+    trace_rowid: _TraceRowId
+
+
+class _AnnoAttr(NamedTuple):
+    trace_rowid: _TraceRowId
+    id_: _AnnoRowId
+    updated_at: datetime
+
+
+def _key(p: Received[Precursors.TraceAnnotation]) -> _Key:
+    return p.item.obj.name, p.item.trace_id
+
+
+def _unique_by(p: Received[Insertables.TraceAnnotation]) -> _UniqueBy:
+    return p.item.obj.name, p.item.trace_rowid
+
+
+def _time(p: Received[Any]) -> datetime:
+    return p.received_at