arize-phoenix 4.20.1__py3-none-any.whl → 4.20.2__py3-none-any.whl

{arize_phoenix-4.20.1.dist-info → arize_phoenix-4.20.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: arize-phoenix
-Version: 4.20.1
+Version: 4.20.2
 Summary: AI Observability and Evaluation
 Project-URL: Documentation, https://docs.arize.com/phoenix/
 Project-URL: Issues, https://github.com/Arize-ai/phoenix/issues

{arize_phoenix-4.20.1.dist-info → arize_phoenix-4.20.2.dist-info}/RECORD

@@ -5,7 +5,7 @@ phoenix/exceptions.py,sha256=n2L2KKuecrdflB9MsCdAYCiSEvGJptIsfRkXMoJle7A,169
 phoenix/py.typed,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
 phoenix/services.py,sha256=aTxhcOA1pZHB6U-B3TEcp6fqDF5oT0xCUvEUNMZVTUQ,5175
 phoenix/settings.py,sha256=cO-qgis_S27nHirTobYI9hHPfZH18R--WMmxNdsVUwc,273
-phoenix/version.py,sha256=x2GXndd_stTekM40AdALx0Eq-Nbip56uWQ4qvw-qCgE,23
+phoenix/version.py,sha256=hh55rhGJPmETx024BaT0pa35ouf8-J2B1_CJUXY4XFg,23
 phoenix/core/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 phoenix/core/embedding_dimension.py,sha256=zKGbcvwOXgLf-yrJBpQyKtd-LEOPRKHnUToyAU8Owis,87
 phoenix/core/model.py,sha256=km_a--PBHOuA337ClRw9xqhOHhrUT6Rl9pz_zV0JYkQ,4843
@@ -18,23 +18,23 @@ phoenix/db/bulk_inserter.py,sha256=qgg8pt5k4VnHKOE0-KoReXVAfXRhLt-sMZihI-b4X9I,1
 phoenix/db/engines.py,sha256=R3btYTSOSd6BwRA59EmhhojL0HCQ7NnzFIXQrPYS0iU,4812
 phoenix/db/helpers.py,sha256=2zSc4n5IJfu-CaOFoBfqTB35M1nTFcAc8tqLsNtF2Jw,3488
 phoenix/db/migrate.py,sha256=MuhtNWnR24riROvarvKfbRb4_D5xuQi6P760vBUKl1E,2270
-phoenix/db/models.py,sha256=7DBWbxY3cx3ve2P1I0kkDKXzlt04zEFJuRPJWsVpH-I,20422
+phoenix/db/models.py,sha256=1fSwI9NjXcVv7PT40VEeKGGI13TWbqzrYfIemXV-DYY,20837
 phoenix/db/insertion/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 phoenix/db/insertion/constants.py,sha256=8wifm7X-1XvroZ__R2Gc96NsgLhTDn0zXl4lehlXtcA,70
 phoenix/db/insertion/dataset.py,sha256=_vxy5e6W5jEuvO2fMKbbNCn9JvHkwI4LRKk_10eKFVg,7171
 phoenix/db/insertion/document_annotation.py,sha256=qp8E63LzlthAScQg6opqln5Qg1d7UdtP3rkYL4riTgU,5983
 phoenix/db/insertion/evaluation.py,sha256=SoI85N3MYUSeNgjKa5WzFw14OfNjNTjExv-2m3sxaR8,6371
 phoenix/db/insertion/helpers.py,sha256=z_Wnckhdf-F7xadWgaAV5ScXnLft8EtaYJCuIkma2Vw,3486
-phoenix/db/insertion/span.py,sha256=T9jOW3lyWte-JGD7wlP2ZqtO0-V57d8z6U_TldnuGuk,5527
+phoenix/db/insertion/span.py,sha256=-CXn2LlEv2u7xbz7m8X5jALQ-BUNNzTxPJEzSmjhTpA,5958
 phoenix/db/insertion/span_annotation.py,sha256=eQK7fFjKjZGj25Qf_PTU9Q8DZiYJAw4lfcfdLKFsKe0,5259
 phoenix/db/insertion/trace_annotation.py,sha256=36CwcxSvDo2r-_y_GlmMXGnlH4BKVaMu5BWM9w-9l7A,5324
 phoenix/db/insertion/types.py,sha256=nQYYnpzcPxj2kdUoXfKE8ilOKlx1zpKLPc40OGuBlfk,8149
 phoenix/db/migrations/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 phoenix/db/migrations/env.py,sha256=QbzB5zrRs6XQQmrYeUpuzeilcMlM-MsbaAgHHYcIHTI,3626
 phoenix/db/migrations/script.py.mako,sha256=MEqL-2qATlST9TAOeYgscMn1uy6HUS9NFvDgl93dMj8,635
-phoenix/db/migrations/types.py,sha256=Frq1AKSyBKQQ0FLzON-EmgTqE4kNkOpHMsbWnI-WgCE,605
-phoenix/db/migrations/versions/10460e46d750_datasets.py,sha256=l69yZfScFrjfZZpY0gnqwhsDUEctLeo02qMgA_aOGDg,8155
-phoenix/db/migrations/versions/cf03bd6bae1d_init.py,sha256=CbWT3ZTR0CZqeT3zWLoTWhboFmnOy3Ju1z6Ztpq8WIM,8122
+phoenix/db/migrations/versions/10460e46d750_datasets.py,sha256=eZAyz720DmpOd7RnuxDN2dVNXVuMrdlCA7eAOxyMtfs,8695
+phoenix/db/migrations/versions/3be8647b87d8_add_token_columns_to_spans_table.py,sha256=x6oKFwn7Zmite4G0trDQPpMCn0I7jejuBcN3-ivEuDg,3938
+phoenix/db/migrations/versions/cf03bd6bae1d_init.py,sha256=09cpofqje8zi4eQFfUn-i21x7VcsUYOfLKKUlrtKrGc,8662
 phoenix/experiments/__init__.py,sha256=6JGwgUd7xCbGpuHqYZlsmErmYvVgv7N_j43bn3dUqsk,123
 phoenix/experiments/functions.py,sha256=4XaOLE1Co9sW_yjM1sypQClmOLtt9kwoxmhIEJ3f_rk,32209
 phoenix/experiments/tracing.py,sha256=wVpt8Ie9WNPoi1djJdcrkwCokHdTO0bicXViLg3O-1Y,2831
@@ -102,7 +102,7 @@ phoenix/server/api/dataloaders/span_dataset_examples.py,sha256=BtLZp11fyyeaWGGBP
 phoenix/server/api/dataloaders/span_descendants.py,sha256=b7jGTn0Hi22gv2yskloLnf3BG3upS9z5hnKLMT9Sxac,2094
 phoenix/server/api/dataloaders/span_evaluations.py,sha256=IfwXW23GQaWti8F49wSJocWf7Tklf2ZJ0F6aB4cSVHs,1248
 phoenix/server/api/dataloaders/span_projects.py,sha256=LbQWiboCFqq4CHS18OzvRUwL9yORqP26fh5p7JbpFdg,1244
-phoenix/server/api/dataloaders/token_counts.py,sha256=riJC3jGRJv2eyenk8Gkx7xwDaulb9SWla58wfcXZeCE,4947
+phoenix/server/api/dataloaders/token_counts.py,sha256=6gDVely8BYiCBdmiq1ECO0lMChUYPIlsZbB34rmL1xM,4684
 phoenix/server/api/dataloaders/trace_evaluations.py,sha256=vraPehNsausR4dbdvq-HudRVNARJUbep3T-Ud9jwWYY,1262
 phoenix/server/api/dataloaders/trace_row_ids.py,sha256=RODX4NULlBzMxHMrsq0dp1ij6ZlLH4ZzQBnafGdxOvU,1100
 phoenix/server/api/dataloaders/cache/__init__.py,sha256=SYoOM9n8FJaMdQarma5d1blu-jIg2GB8Shqg5ezSzZ8,106
@@ -133,7 +133,7 @@ phoenix/server/api/input_types/PatchDatasetExamplesInput.py,sha256=E86aBGXDBC83j
 phoenix/server/api/input_types/PatchDatasetInput.py,sha256=OURtTVY8Z_oFEDtKwT1LCMaOK5D4QYo5TVQ6mDrex-g,328
 phoenix/server/api/input_types/PerformanceMetricInput.py,sha256=fElsLTSEYYgGFGMYTEGcYid39tXUKFdV_JkdHavMcbA,591
 phoenix/server/api/input_types/SpanAnnotationSort.py,sha256=T5pAGzmh4MiJp9JMAzNDByFVTczfw02FH4WFWwFezyI,361
-phoenix/server/api/input_types/SpanSort.py,sha256=BqV0DjI4m3RKFbzsHN47XNhJfFmErFIYg64c8EJAVMo,6312
+phoenix/server/api/input_types/SpanSort.py,sha256=Dhvl8BIoV52yHoqntfOax_gUc15uH8ITI_00Ha7PvYc,5959
 phoenix/server/api/input_types/TimeRange.py,sha256=yzx-gxj8mDeGLft1FzU_x1MVEgIG5Pt6-f8PUVDgipQ,522
 phoenix/server/api/input_types/TraceAnnotationSort.py,sha256=BzwiUnMh2VsgQYnhDlbJ6ljHugqIS4YDUlYzvq_tl3o,365
 phoenix/server/api/input_types/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -205,7 +205,7 @@ phoenix/server/api/types/Retrieval.py,sha256=OhMK2ncjoyp5h1yjKhjlKpoTbQrMHuxmgSF
 phoenix/server/api/types/ScalarDriftMetricEnum.py,sha256=IUAcRPpgL41WdoIgK6cNk2Te38SspXGyEs-S1fY23_A,232
 phoenix/server/api/types/Segments.py,sha256=m2yoegrxA1Tn7ZAy1rMjjD1isc752MaAXMoffkBlvrM,2921
 phoenix/server/api/types/SortDir.py,sha256=OUpXhlCzCxPoXSDkJJygEs9Rw9pMymfaZUG5zPTrw4Y,152
-phoenix/server/api/types/Span.py,sha256=k3j8P06927jNK5P8VrqQLRhOu3pyiPlV5wdiGdXC1cE,15263
+phoenix/server/api/types/Span.py,sha256=xzJoRuzFf1S64jDuPmGLjSFZ4oPyKUyfK1CtmtZ4LY0,14801
 phoenix/server/api/types/SpanAnnotation.py,sha256=6b5G-b_OoRvDL2ayWk7MkbqarLK-F-pQMx21CpUuNGY,1168
 phoenix/server/api/types/TimeSeries.py,sha256=wjzuxHFqCey0O7Ys25qiXyuqXK8an-osyNWUE8A_8G4,5227
 phoenix/server/api/types/Trace.py,sha256=-nh3A-S_BlQK1VSSOTWqM85l-WwJsRHifxeDi0sFWZE,3246
@@ -227,13 +227,13 @@ phoenix/server/static/apple-touch-icon-76x76.png,sha256=CT_xT12I0u2i0WU8JzBZBuOQ
 phoenix/server/static/apple-touch-icon.png,sha256=fOfpjqGpWYbJ0eAurKsyoZP1EAs6ZVooBJ_SGk2ZkDs,3801
 phoenix/server/static/favicon.ico,sha256=bY0vvCKRftemZfPShwZtE93DiiQdaYaozkPGwNFr6H8,34494
 phoenix/server/static/modernizr.js,sha256=mvK-XtkNqjOral-QvzoqsyOMECXIMu5BQwSVN_wcU9c,2564
-phoenix/server/static/.vite/manifest.json,sha256=iHg8InkOETcckHJQlUfToMmMPaocrvCM68tSU6BEU1Q,1929
-phoenix/server/static/assets/components-CAummAJx.js,sha256=LvIzotUMWSCvu2IT9QHFGK_JJB_LYXyFJagHFyu6AVI,184229
-phoenix/server/static/assets/index-Cg5hdf3g.js,sha256=-mDNUU79jpnfsiTaJztdQ_WVKySUcyRbbZyRCztZKJI,7362
-phoenix/server/static/assets/pages-BU__X1UX.js,sha256=tNQZgWWC9ycS2CMKUm5qdRlnhpPfIECvuM8E5AUOAtY,445374
+phoenix/server/static/.vite/manifest.json,sha256=f_O0wfSXGRjeXvLiAE1oNla0ADk3X6PTSnCeHtd4xB8,1929
+phoenix/server/static/assets/components-BSw2e1Zr.js,sha256=nvfZM0lYZkeOklzR8VLMikEgIOTpb2jNpF8b9IMMcps,185927
+phoenix/server/static/assets/index-BYUFcdtx.js,sha256=bYcfXq3huGgFkCR0SX2680A1MbACjpdxd5MIbMcOIyY,7362
+phoenix/server/static/assets/pages-p_fuED5k.js,sha256=MA3xlc64hC6bRMJRcXZ1v1sWZdsngaU8bcejbho402A,449363
 phoenix/server/static/assets/vendor-BMWfu6zp.js,sha256=AAVTM5SjGUI_CmAWFUFmhpp5VDhvCD-MrEoh-pXXADY,1355423
 phoenix/server/static/assets/vendor-DxkFTwjz.css,sha256=nZrkr0u6NNElFGvpWHk9GTHeGoibCXCli1bE7mXZGZg,1816
-phoenix/server/static/assets/vendor-arizeai-CkyzG9Wl.js,sha256=m1xvoung59IMplKeV-0XmsSPRwu1AJDRxwHbq6BN-oc,296355
+phoenix/server/static/assets/vendor-arizeai-CIETbKDq.js,sha256=JYd0o3cIY0av7lzU5LT5GDMMhmMPE3FbbGjtlhKBhLo,304008
 phoenix/server/static/assets/vendor-codemirror-DO3VqEcD.js,sha256=M7t6xd6WpgKes25OOeGyxT1MU1dDrEKdmUBHgy5zslw,503031
 phoenix/server/static/assets/vendor-recharts-BGN0SxgJ.js,sha256=L9LAYSjuf0GHh1_PQh9bF4l9euWCDVQcnQN1RgMDMBw,282859
 phoenix/server/static/assets/vendor-three-DwGkEfCM.js,sha256=0D12ZgKzfKCTSdSTKJBFR2RZO_xxeMXrqDp0AszZqHY,620972
@@ -281,8 +281,8 @@ phoenix/utilities/logging.py,sha256=lDXd6EGaamBNcQxL4vP1au9-i_SXe0OraUDiJOcszSw,
 phoenix/utilities/project.py,sha256=8IJuMM4yUMoooPi37sictGj8Etu9rGmq6RFtc9848cQ,436
 phoenix/utilities/re.py,sha256=PDve_OLjRTM8yQQJHC8-n3HdIONi7aNils3ZKRZ5uBM,2045
 phoenix/utilities/span_store.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-arize_phoenix-4.20.1.dist-info/METADATA,sha256=m65-hbRGN_Pg_ajtvbWhoKoGdqvq4Y8o3_850jJTGKM,11902
-arize_phoenix-4.20.1.dist-info/WHEEL,sha256=1yFddiXMmvYK7QYTqtRNtX66WJ0Mz8PYEiEUoOUUxRY,87
-arize_phoenix-4.20.1.dist-info/licenses/IP_NOTICE,sha256=JBqyyCYYxGDfzQ0TtsQgjts41IJoa-hiwDrBjCb9gHM,469
-arize_phoenix-4.20.1.dist-info/licenses/LICENSE,sha256=HFkW9REuMOkvKRACuwLPT0hRydHb3zNg-fdFt94td18,3794
-arize_phoenix-4.20.1.dist-info/RECORD,,
+arize_phoenix-4.20.2.dist-info/METADATA,sha256=DAffUocQNJm0BONMNSzCucpDNIILl9e2JIdARb3As3c,11902
+arize_phoenix-4.20.2.dist-info/WHEEL,sha256=1yFddiXMmvYK7QYTqtRNtX66WJ0Mz8PYEiEUoOUUxRY,87
+arize_phoenix-4.20.2.dist-info/licenses/IP_NOTICE,sha256=JBqyyCYYxGDfzQ0TtsQgjts41IJoa-hiwDrBjCb9gHM,469
+arize_phoenix-4.20.2.dist-info/licenses/LICENSE,sha256=HFkW9REuMOkvKRACuwLPT0hRydHb3zNg-fdFt94td18,3794
+arize_phoenix-4.20.2.dist-info/RECORD,,

phoenix/db/insertion/span.py

@@ -71,6 +71,13 @@ async def insert_span(
     cumulative_llm_token_count_completion = cast(
         int, get_attribute_value(span.attributes, SpanAttributes.LLM_TOKEN_COUNT_COMPLETION) or 0
     )
+    llm_token_count_prompt = cast(
+        Optional[int], get_attribute_value(span.attributes, SpanAttributes.LLM_TOKEN_COUNT_PROMPT)
+    )
+    llm_token_count_completion = cast(
+        Optional[int],
+        get_attribute_value(span.attributes, SpanAttributes.LLM_TOKEN_COUNT_COMPLETION),
+    )
     if accumulation := (
         await session.execute(
             select(
@@ -100,6 +107,8 @@ async def insert_span(
                 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,

phoenix/db/migrations/versions/10460e46d750_datasets.py

@@ -6,11 +6,37 @@ Create Date: 2024-05-10 11:24:23.985834
 
 """
 
-from typing import Sequence, Union
+from typing import Any, Sequence, Union
 
 import sqlalchemy as sa
 from alembic import op
-from phoenix.db.migrations.types import JSON_
+from sqlalchemy import JSON
+from sqlalchemy.dialects import postgresql
+from sqlalchemy.ext.compiler import compiles
+
+
+class JSONB(JSON):
+    # See https://docs.sqlalchemy.org/en/20/core/custom_types.html
+    __visit_name__ = "JSONB"
+
+
+@compiles(JSONB, "sqlite")  # type: ignore
+def _(*args: Any, **kwargs: Any) -> str:
+    # See https://docs.sqlalchemy.org/en/20/core/custom_types.html
+    return "JSONB"
+
+
+JSON_ = (
+    JSON()
+    .with_variant(
+        postgresql.JSONB(),  # type: ignore
+        "postgresql",
+    )
+    .with_variant(
+        JSONB(),
+        "sqlite",
+    )
+)
 
 # revision identifiers, used by Alembic.
 revision: str = "10460e46d750"

phoenix/db/migrations/versions/3be8647b87d8_add_token_columns_to_spans_table.py

@@ -0,0 +1,134 @@
+"""add token columns to spans table
+
+Revision ID: 3be8647b87d8
+Revises: 10460e46d750
+Create Date: 2024-08-03 22:11:28.733133
+
+"""
+
+from typing import Any, Dict, List, Optional, Sequence, TypedDict, Union
+
+import sqlalchemy as sa
+from alembic import op
+from openinference.semconv.trace import SpanAttributes
+from sqlalchemy import (
+    JSON,
+    Dialect,
+    MetaData,
+    TypeDecorator,
+    update,
+)
+from sqlalchemy.dialects import postgresql
+from sqlalchemy.ext.asyncio.engine import AsyncConnection
+from sqlalchemy.ext.compiler import compiles
+from sqlalchemy.orm import (
+    DeclarativeBase,
+    Mapped,
+    mapped_column,
+)
+
+
+class JSONB(JSON):
+    # See https://docs.sqlalchemy.org/en/20/core/custom_types.html
+    __visit_name__ = "JSONB"
+
+
+@compiles(JSONB, "sqlite")  # type: ignore
+def _(*args: Any, **kwargs: Any) -> str:
+    # See https://docs.sqlalchemy.org/en/20/core/custom_types.html
+    return "JSONB"
+
+
+JSON_ = (
+    JSON()
+    .with_variant(
+        postgresql.JSONB(),  # type: ignore
+        "postgresql",
+    )
+    .with_variant(
+        JSONB(),
+        "sqlite",
+    )
+)
+
+
+class JsonDict(TypeDecorator[Dict[str, Any]]):
+    # See # See https://docs.sqlalchemy.org/en/20/core/custom_types.html
+    cache_ok = True
+    impl = JSON_
+
+    def process_bind_param(self, value: Optional[Dict[str, Any]], _: Dialect) -> Dict[str, Any]:
+        return value if isinstance(value, dict) else {}
+
+
+class JsonList(TypeDecorator[List[Any]]):
+    # See # See https://docs.sqlalchemy.org/en/20/core/custom_types.html
+    cache_ok = True
+    impl = JSON_
+
+    def process_bind_param(self, value: Optional[List[Any]], _: Dialect) -> List[Any]:
+        return value if isinstance(value, list) else []
+
+
+class ExperimentRunOutput(TypedDict, total=False):
+    task_output: Any
+
+
+class Base(DeclarativeBase):
+    # Enforce best practices for naming constraints
+    # https://alembic.sqlalchemy.org/en/latest/naming.html#integration-of-naming-conventions-into-operations-autogenerate
+    metadata = MetaData(
+        naming_convention={
+            "ix": "ix_%(table_name)s_%(column_0_N_name)s",
+            "uq": "uq_%(table_name)s_%(column_0_N_name)s",
+            "ck": "ck_%(table_name)s_`%(constraint_name)s`",
+            "fk": "fk_%(table_name)s_%(column_0_name)s_%(referred_table_name)s",
+            "pk": "pk_%(table_name)s",
+        }
+    )
+    type_annotation_map = {
+        Dict[str, Any]: JsonDict,
+        List[Dict[str, Any]]: JsonList,
+        ExperimentRunOutput: JsonDict,
+    }
+
+
+class Span(Base):
+    __tablename__ = "spans"
+    id: Mapped[int] = mapped_column(primary_key=True)
+    attributes: Mapped[Dict[str, Any]]
+    llm_token_count_prompt: Mapped[Optional[int]]
+    llm_token_count_completion: Mapped[Optional[int]]
+
+
+# revision identifiers, used by Alembic.
+revision: str = "3be8647b87d8"
+down_revision: Union[str, None] = "10460e46d750"
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+LLM_TOKEN_COUNT_PROMPT = SpanAttributes.LLM_TOKEN_COUNT_PROMPT.split(".")
+LLM_TOKEN_COUNT_COMPLETION = SpanAttributes.LLM_TOKEN_COUNT_COMPLETION.split(".")
+
+
+async def get_token_counts_from_attributes(connection: AsyncConnection) -> None:
+    """
+    Gets token counts from attributes if present.
+    """
+    await connection.execute(
+        update(Span).values(
+            llm_token_count_prompt=Span.attributes[LLM_TOKEN_COUNT_PROMPT].as_float(),
+            llm_token_count_completion=Span.attributes[LLM_TOKEN_COUNT_COMPLETION].as_float(),
+        )
+    )
+
+
+def upgrade() -> None:
+    op.add_column("spans", sa.Column("llm_token_count_prompt", sa.Integer, nullable=True))
+    op.add_column("spans", sa.Column("llm_token_count_completion", sa.Integer, nullable=True))
+    op.run_async(get_token_counts_from_attributes)
+
+
+def downgrade() -> None:
+    op.drop_column("spans", "llm_token_count_completion")
+    op.drop_column("spans", "llm_token_count_prompt")

phoenix/db/migrations/versions/cf03bd6bae1d_init.py

@@ -6,11 +6,37 @@ Create Date: 2024-04-03 19:41:48.871555
 
 """
 
-from typing import Sequence, Union
+from typing import Any, Sequence, Union
 
 import sqlalchemy as sa
 from alembic import op
-from phoenix.db.migrations.types import JSON_
+from sqlalchemy import JSON
+from sqlalchemy.dialects import postgresql
+from sqlalchemy.ext.compiler import compiles
+
+
+class JSONB(JSON):
+    # See https://docs.sqlalchemy.org/en/20/core/custom_types.html
+    __visit_name__ = "JSONB"
+
+
+@compiles(JSONB, "sqlite")  # type: ignore
+def _(*args: Any, **kwargs: Any) -> str:
+    # See https://docs.sqlalchemy.org/en/20/core/custom_types.html
+    return "JSONB"
+
+
+JSON_ = (
+    JSON()
+    .with_variant(
+        postgresql.JSONB(),  # type: ignore
+        "postgresql",
+    )
+    .with_variant(
+        JSONB(),
+        "sqlite",
+    )
+)
 
 # revision identifiers, used by Alembic.
 revision: str = "cf03bd6bae1d"

phoenix/db/models.py

@@ -197,7 +197,7 @@ class Span(Base):
         ForeignKey("traces.id", ondelete="CASCADE"),
         index=True,
     )
-    span_id: Mapped[str]
+    span_id: Mapped[str] = mapped_column(index=True)
     parent_id: Mapped[Optional[str]] = mapped_column(index=True)
     name: Mapped[str]
     span_kind: Mapped[str]
@@ -214,6 +214,8 @@ class Span(Base):
     cumulative_error_count: Mapped[int]
     cumulative_llm_token_count_prompt: Mapped[int]
     cumulative_llm_token_count_completion: Mapped[int]
+    llm_token_count_prompt: Mapped[Optional[int]]
+    llm_token_count_completion: Mapped[Optional[int]]
 
     @hybrid_property
     def latency_ms(self) -> float:
@@ -230,6 +232,12 @@ class Span(Base):
     def cumulative_llm_token_count_total(self) -> int:
         return self.cumulative_llm_token_count_prompt + self.cumulative_llm_token_count_completion
 
+    @hybrid_property
+    def llm_token_count_total(self) -> Optional[int]:
+        if self.llm_token_count_prompt is None and self.llm_token_count_completion is None:
+            return None
+        return (self.llm_token_count_prompt or 0) + (self.llm_token_count_completion or 0)
+
     trace: Mapped["Trace"] = relationship("Trace", back_populates="spans")
     document_annotations: Mapped[List["DocumentAnnotation"]] = relationship(back_populates="span")
     dataset_examples: Mapped[List["DatasetExample"]] = relationship(back_populates="span")

phoenix/server/api/dataloaders/token_counts.py

@@ -10,7 +10,6 @@ from typing import (
 )
 
 from cachetools import LFUCache, TTLCache
-from openinference.semconv.trace import SpanAttributes
 from sqlalchemy import Select, func, select
 from sqlalchemy.sql.functions import coalesce
 from strawberry.dataloader import AbstractCache, DataLoader
@@ -107,8 +106,8 @@ def _get_stmt(
     *params: Param,
 ) -> Select[Any]:
     (start_time, end_time), filter_condition = segment
-    prompt = func.sum(models.Span.attributes[_LLM_TOKEN_COUNT_PROMPT].as_float())
-    completion = func.sum(models.Span.attributes[_LLM_TOKEN_COUNT_COMPLETION].as_float())
+    prompt = func.sum(models.Span.llm_token_count_prompt)
+    completion = func.sum(models.Span.llm_token_count_completion)
     total = coalesce(prompt, 0) + coalesce(completion, 0)
     pid = models.Trace.project_rowid
     stmt: Select[Any] = (
@@ -130,7 +129,3 @@ def _get_stmt(
         stmt = sf(stmt)
     stmt = stmt.where(pid.in_([rowid for rowid, _ in params]))
     return stmt
-
-
-_LLM_TOKEN_COUNT_PROMPT = SpanAttributes.LLM_TOKEN_COUNT_PROMPT.split(".")
-_LLM_TOKEN_COUNT_COMPLETION = SpanAttributes.LLM_TOKEN_COUNT_COMPLETION.split(".")

phoenix/server/api/input_types/SpanSort.py

@@ -3,7 +3,6 @@ from enum import Enum, auto
 from typing import Any, Optional, Protocol
 
 import strawberry
-from openinference.semconv.trace import SpanAttributes
 from sqlalchemy import and_, desc, nulls_last
 from sqlalchemy.orm import InstrumentedAttribute
 from sqlalchemy.sql.expression import Select
@@ -16,10 +15,6 @@ from phoenix.server.api.types.pagination import CursorSortColumnDataType
 from phoenix.server.api.types.SortDir import SortDir
 from phoenix.trace.schemas import SpanID
 
-LLM_TOKEN_COUNT_PROMPT = SpanAttributes.LLM_TOKEN_COUNT_PROMPT.split(".")
-LLM_TOKEN_COUNT_COMPLETION = SpanAttributes.LLM_TOKEN_COUNT_COMPLETION.split(".")
-LLM_TOKEN_COUNT_TOTAL = SpanAttributes.LLM_TOKEN_COUNT_TOTAL.split(".")
-
 
 @strawberry.enum
 class SpanColumn(Enum):
@@ -47,11 +42,11 @@ class SpanColumn(Enum):
         elif self is SpanColumn.latencyMs:
             expr = models.Span.latency_ms
         elif self is SpanColumn.tokenCountTotal:
-            expr = models.Span.attributes[LLM_TOKEN_COUNT_TOTAL].as_float()
+            expr = models.Span.llm_token_count_total
         elif self is SpanColumn.tokenCountPrompt:
-            expr = models.Span.attributes[LLM_TOKEN_COUNT_PROMPT].as_float()
+            expr = models.Span.llm_token_count_prompt
         elif self is SpanColumn.tokenCountCompletion:
-            expr = models.Span.attributes[LLM_TOKEN_COUNT_COMPLETION].as_float()
+            expr = models.Span.llm_token_count_completion
         elif self is SpanColumn.cumulativeTokenCountTotal:
             expr = (
                 models.Span.cumulative_llm_token_count_prompt

phoenix/server/api/types/Span.py

@@ -40,9 +40,6 @@ EMBEDDING_EMBEDDINGS = SpanAttributes.EMBEDDING_EMBEDDINGS
 EMBEDDING_VECTOR = EmbeddingAttributes.EMBEDDING_VECTOR
 INPUT_MIME_TYPE = SpanAttributes.INPUT_MIME_TYPE
 INPUT_VALUE = SpanAttributes.INPUT_VALUE
-LLM_TOKEN_COUNT_COMPLETION = SpanAttributes.LLM_TOKEN_COUNT_COMPLETION
-LLM_TOKEN_COUNT_PROMPT = SpanAttributes.LLM_TOKEN_COUNT_PROMPT
-LLM_TOKEN_COUNT_TOTAL = SpanAttributes.LLM_TOKEN_COUNT_TOTAL
 LLM_PROMPT_TEMPLATE_VARIABLES = SpanAttributes.LLM_PROMPT_TEMPLATE_VARIABLES
 LLM_INPUT_MESSAGES = SpanAttributes.LLM_INPUT_MESSAGES
 LLM_OUTPUT_MESSAGES = SpanAttributes.LLM_OUTPUT_MESSAGES
@@ -308,18 +305,9 @@ def to_gql_span(span: models.Span) -> Span:
         attributes=json.dumps(_hide_embedding_vectors(span.attributes), cls=_JSONEncoder),
         metadata=_convert_metadata_to_string(get_attribute_value(span.attributes, METADATA)),
         num_documents=num_documents,
-        token_count_total=cast(
-            Optional[int],
-            get_attribute_value(span.attributes, LLM_TOKEN_COUNT_TOTAL),
-        ),
-        token_count_prompt=cast(
-            Optional[int],
-            get_attribute_value(span.attributes, LLM_TOKEN_COUNT_PROMPT),
-        ),
-        token_count_completion=cast(
-            Optional[int],
-            get_attribute_value(span.attributes, LLM_TOKEN_COUNT_COMPLETION),
-        ),
+        token_count_total=span.llm_token_count_total,
+        token_count_prompt=span.llm_token_count_prompt,
+        token_count_completion=span.llm_token_count_completion,
         cumulative_token_count_total=span.cumulative_llm_token_count_prompt
         + span.cumulative_llm_token_count_completion,
         cumulative_token_count_prompt=span.cumulative_llm_token_count_prompt,

phoenix/server/static/.vite/manifest.json

@@ -1,22 +1,22 @@
 {
-  "_components-CAummAJx.js": {
-    "file": "assets/components-CAummAJx.js",
+  "_components-BSw2e1Zr.js": {
+    "file": "assets/components-BSw2e1Zr.js",
     "name": "components",
     "imports": [
       "_vendor-BMWfu6zp.js",
-      "_vendor-arizeai-CkyzG9Wl.js",
-      "_pages-BU__X1UX.js",
+      "_vendor-arizeai-CIETbKDq.js",
+      "_pages-p_fuED5k.js",
       "_vendor-three-DwGkEfCM.js",
       "_vendor-codemirror-DO3VqEcD.js"
     ]
   },
-  "_pages-BU__X1UX.js": {
-    "file": "assets/pages-BU__X1UX.js",
+  "_pages-p_fuED5k.js": {
+    "file": "assets/pages-p_fuED5k.js",
     "name": "pages",
     "imports": [
       "_vendor-BMWfu6zp.js",
-      "_components-CAummAJx.js",
-      "_vendor-arizeai-CkyzG9Wl.js",
+      "_components-BSw2e1Zr.js",
+      "_vendor-arizeai-CIETbKDq.js",
       "_vendor-recharts-BGN0SxgJ.js",
       "_vendor-codemirror-DO3VqEcD.js"
     ]
@@ -35,8 +35,8 @@
       "assets/vendor-DxkFTwjz.css"
     ]
   },
-  "_vendor-arizeai-CkyzG9Wl.js": {
-    "file": "assets/vendor-arizeai-CkyzG9Wl.js",
+  "_vendor-arizeai-CIETbKDq.js": {
+    "file": "assets/vendor-arizeai-CIETbKDq.js",
     "name": "vendor-arizeai",
     "imports": [
       "_vendor-BMWfu6zp.js"
@@ -61,15 +61,15 @@
     "name": "vendor-three"
   },
   "index.tsx": {
-    "file": "assets/index-Cg5hdf3g.js",
+    "file": "assets/index-BYUFcdtx.js",
     "name": "index",
     "src": "index.tsx",
     "isEntry": true,
     "imports": [
       "_vendor-BMWfu6zp.js",
-      "_vendor-arizeai-CkyzG9Wl.js",
-      "_components-CAummAJx.js",
-      "_pages-BU__X1UX.js",
+      "_vendor-arizeai-CIETbKDq.js",
+      "_components-BSw2e1Zr.js",
+      "_pages-p_fuED5k.js",
       "_vendor-three-DwGkEfCM.js",
       "_vendor-codemirror-DO3VqEcD.js",
       "_vendor-recharts-BGN0SxgJ.js"