arize-phoenix 0.0.32rc1__py3-none-any.whl → 0.0.33__py3-none-any.whl

phoenix/server/api/types/Span.py

@@ -0,0 +1,229 @@
+import json
+from collections import defaultdict
+from datetime import datetime
+from enum import Enum
+from typing import Any, DefaultDict, List, Mapping, Optional, cast
+
+import strawberry
+from strawberry import ID
+from strawberry.types import Info
+
+import phoenix.trace.schemas as trace_schema
+from phoenix.core.traces import (
+    CUMULATIVE_LLM_TOKEN_COUNT_COMPLETION,
+    CUMULATIVE_LLM_TOKEN_COUNT_PROMPT,
+    CUMULATIVE_LLM_TOKEN_COUNT_TOTAL,
+    LATENCY_MS,
+)
+from phoenix.server.api.context import Context
+from phoenix.server.api.types.MimeType import MimeType
+from phoenix.trace.schemas import SpanID
+from phoenix.trace.semantic_conventions import (
+    EXCEPTION_MESSAGE,
+    INPUT_MIME_TYPE,
+    INPUT_VALUE,
+    LLM_TOKEN_COUNT_COMPLETION,
+    LLM_TOKEN_COUNT_PROMPT,
+    LLM_TOKEN_COUNT_TOTAL,
+    OUTPUT_MIME_TYPE,
+    OUTPUT_VALUE,
+)
+
+
+@strawberry.enum
+class SpanKind(Enum):
+    """
+    The type of work that a Span encapsulates.
+
+    NB: this is actively under construction
+    """
+
+    chain = trace_schema.SpanKind.CHAIN
+    tool = trace_schema.SpanKind.TOOL
+    llm = trace_schema.SpanKind.LLM
+    retriever = trace_schema.SpanKind.RETRIEVER
+    embedding = trace_schema.SpanKind.EMBEDDING
+    agent = trace_schema.SpanKind.AGENT
+    unknown = trace_schema.SpanKind.UNKNOWN
+
+    @classmethod
+    def _missing_(cls, v: Any) -> Optional["SpanKind"]:
+        return None if v else cls.unknown
+
+
+@strawberry.type
+class SpanContext:
+    trace_id: ID
+    span_id: ID
+
+
+@strawberry.type
+class SpanIOValue:
+    mime_type: MimeType
+    value: str
+
+
+@strawberry.enum
+class SpanStatusCode(Enum):
+    OK = trace_schema.SpanStatusCode.OK
+    ERROR = trace_schema.SpanStatusCode.ERROR
+    UNSET = trace_schema.SpanStatusCode.UNSET
+
+    @classmethod
+    def _missing_(cls, v: Any) -> Optional["SpanStatusCode"]:
+        return None if v else cls.UNSET
+
+
+@strawberry.type
+class SpanEvent:
+    name: str
+    message: str
+    timestamp: datetime
+
+    @staticmethod
+    def from_event(
+        event: trace_schema.SpanEvent,
+    ) -> "SpanEvent":
+        return SpanEvent(
+            name=event.name,
+            message=cast(str, event.attributes.get(EXCEPTION_MESSAGE) or ""),
+            timestamp=event.timestamp,
+        )
+
+
+@strawberry.type
+class Span:
+    name: str
+    status_code: SpanStatusCode
+    start_time: datetime
+    end_time: Optional[datetime]
+    latency_ms: Optional[float]
+    parent_id: Optional[ID] = strawberry.field(
+        description="the parent span ID. If null, it is a root span"
+    )
+    span_kind: SpanKind
+    context: SpanContext
+    attributes: str = strawberry.field(
+        description="Span attributes as a JSON string",
+    )
+    token_count_total: Optional[int]
+    token_count_prompt: Optional[int]
+    token_count_completion: Optional[int]
+    input: Optional[SpanIOValue]
+    output: Optional[SpanIOValue]
+    events: List[SpanEvent]
+    cumulative_token_count_total: Optional[int] = strawberry.field(
+        description="Cumulative (prompt plus completion) token count from "
+        "self and all descendant spans (children, grandchildren, etc.)",
+    )
+    cumulative_token_count_prompt: Optional[int] = strawberry.field(
+        description="Cumulative (prompt) token count from self and all "
+        "descendant spans (children, grandchildren, etc.)",
+    )
+    cumulative_token_count_completion: Optional[int] = strawberry.field(
+        description="Cumulative (completion) token count from self and all "
+        "descendant spans (children, grandchildren, etc.)",
+    )
+
+    @strawberry.field(
+        description="All descendant spans (children, grandchildren, etc.)",
+    )  # type: ignore
+    def descendants(
+        self,
+        info: Info[Context, None],
+    ) -> List["Span"]:
+        if (traces := info.context.traces) is None:
+            return []
+        return [
+            to_gql_span(cast(trace_schema.Span, traces[span_id]))
+            for span_id in traces.get_descendant_span_ids(
+                cast(SpanID, self.context.span_id),
+            )
+        ]
+
+
+def to_gql_span(span: trace_schema.Span) -> "Span":
+    events: List[SpanEvent] = list(map(SpanEvent.from_event, span.events))
+    input_value = cast(Optional[str], span.attributes.get(INPUT_VALUE))
+    output_value = cast(Optional[str], span.attributes.get(OUTPUT_VALUE))
+    return Span(
+        name=span.name,
+        status_code=SpanStatusCode(span.status_code),
+        parent_id=cast(Optional[ID], span.parent_id),
+        span_kind=SpanKind(span.span_kind),
+        start_time=span.start_time,
+        end_time=span.end_time,
+        latency_ms=cast(Optional[float], span.attributes.get(LATENCY_MS)),
+        context=SpanContext(
+            trace_id=cast(ID, span.context.trace_id),
+            span_id=cast(ID, span.context.span_id),
+        ),
+        attributes=json.dumps(
+            _nested_attributes(span.attributes),
+            default=_json_encode,
+        ),
+        token_count_total=cast(
+            Optional[int],
+            span.attributes.get(LLM_TOKEN_COUNT_TOTAL),
+        ),
+        token_count_prompt=cast(
+            Optional[int],
+            span.attributes.get(LLM_TOKEN_COUNT_PROMPT),
+        ),
+        token_count_completion=cast(
+            Optional[int],
+            span.attributes.get(LLM_TOKEN_COUNT_COMPLETION),
+        ),
+        cumulative_token_count_total=cast(
+            Optional[int],
+            span.attributes.get(CUMULATIVE_LLM_TOKEN_COUNT_TOTAL),
+        ),
+        cumulative_token_count_prompt=cast(
+            Optional[int],
+            span.attributes.get(CUMULATIVE_LLM_TOKEN_COUNT_PROMPT),
+        ),
+        cumulative_token_count_completion=cast(
+            Optional[int],
+            span.attributes.get(CUMULATIVE_LLM_TOKEN_COUNT_COMPLETION),
+        ),
+        events=events,
+        input=(
+            SpanIOValue(
+                mime_type=MimeType(span.attributes.get(INPUT_MIME_TYPE)),
+                value=input_value,
+            )
+            if input_value is not None
+            else None
+        ),
+        output=(
+            SpanIOValue(
+                mime_type=MimeType(span.attributes.get(OUTPUT_MIME_TYPE)),
+                value=output_value,
+            )
+            if output_value is not None
+            else None
+        ),
+    )
+
+
+def _json_encode(v: Any) -> str:
+    if isinstance(v, datetime):
+        return v.isoformat()
+    return str(v)
+
+
+def _trie() -> DefaultDict[str, Any]:
+    return defaultdict(_trie)
+
+
+def _nested_attributes(
+    attributes: Mapping[str, Any],
+) -> DefaultDict[str, Any]:
+    nested_attributes = _trie()
+    for attribute_name, attribute_value in attributes.items():
+        trie = nested_attributes
+        keys = attribute_name.split(".")
+        for key in keys[:-1]:
+            trie = trie[key]
+        trie[keys[-1]] = attribute_value
+    return nested_attributes

phoenix/server/api/types/TimeSeries.py

@@ -7,7 +7,14 @@ import pandas as pd
 import strawberry
 from strawberry import UNSET
 
-from phoenix.core.model_schema import CONTINUOUS, PRIMARY, REFERENCE, Column, Dataset, Dimension
+from phoenix.core.model_schema import (
+    CONTINUOUS,
+    PRIMARY,
+    REFERENCE,
+    Column,
+    Dataset,
+    Dimension,
+)
 from phoenix.metrics import Metric, binning
 from phoenix.metrics.mixins import UnaryOperator
 from phoenix.metrics.timeseries import timeseries
@@ -31,7 +38,7 @@ class TimeSeriesDataPoint:
     """The value of the data point"""
     value: Optional[float] = strawberry.field(default=GqlValueMediator())
 
-    def __lt__(self, other: "TimeSeriesDataPoint") -> bool:
+    def __lt__(self, other: "TimeSeriesDataPoint") -> bool:  # type: ignore
         return self.timestamp < other.timestamp
 
 

phoenix/server/api/types/pagination.py

@@ -35,6 +35,7 @@ class PageInfo:
     has_previous_page: bool
     start_cursor: Optional[str]
     end_cursor: Optional[str]
+    total_count: int
 
 
 # A type alias for the connection cursor implementation
@@ -168,5 +169,6 @@ def connection_from_list_slice(
             end_cursor=last_edge.cursor if last_edge else None,
             has_previous_page=start_offset > lower_bound if isinstance(args.last, int) else False,
             has_next_page=end_offset < upper_bound if isinstance(args.first, int) else False,
+            total_count=list_length,
         ),
     )

phoenix/server/app.py

@@ -19,9 +19,10 @@ from strawberry.schema import BaseSchema
 
 from phoenix.config import SERVER_DIR
 from phoenix.core.model_schema import Model
-
-from .api.context import Context
-from .api.schema import schema
+from phoenix.core.traces import Traces
+from phoenix.server.api.context import Context
+from phoenix.server.api.schema import schema
+from phoenix.server.span_handler import SpanHandler
 
 logger = logging.getLogger(__name__)
 
@@ -65,9 +66,11 @@ class GraphQLWithContext(GraphQL):  # type: ignore
         export_path: Path,
         graphiql: bool = False,
         corpus: Optional[Model] = None,
+        traces: Optional[Traces] = None,
     ) -> None:
         self.model = model
         self.corpus = corpus
+        self.traces = traces
         self.export_path = export_path
         super().__init__(schema, graphiql=graphiql)
 
@@ -81,6 +84,7 @@ class GraphQLWithContext(GraphQL):  # type: ignore
             response=response,
             model=self.model,
             corpus=self.corpus,
+            traces=self.traces,
             export_path=self.export_path,
         )
 
@@ -104,12 +108,14 @@ def create_app(
     export_path: Path,
     model: Model,
     corpus: Optional[Model] = None,
+    traces: Optional[Traces] = None,
     debug: bool = False,
 ) -> Starlette:
     graphql = GraphQLWithContext(
         schema=schema,
         model=model,
         corpus=corpus,
+        traces=traces,
         export_path=export_path,
         graphiql=True,
     )
@@ -118,7 +124,21 @@ def create_app(
             Middleware(HeadersMiddleware),
         ],
         debug=debug,
-        routes=[
+        routes=(
+            []
+            if traces is None
+            else [
+                Route(
+                    "/v1/spans",
+                    type(
+                        "SpanEndpoint",
+                        (SpanHandler,),
+                        {"queue": traces},
+                    ),
+                ),
+            ]
+        )
+        + [
             Route(
                 "/exports",
                 type(

phoenix/server/main.py

@@ -1,53 +1,67 @@
 import atexit
-import errno
 import logging
 import os
 from argparse import ArgumentParser
 from pathlib import Path
+from threading import Thread
+from time import sleep, time
 from typing import Optional
 
-import uvicorn
+from uvicorn import Config, Server
 
-import phoenix.config as config
+from phoenix.config import EXPORT_DIR, get_env_host, get_env_port, get_pids_path
 from phoenix.core.model_schema_adapter import create_model_from_datasets
-from phoenix.datasets.dataset import Dataset
+from phoenix.core.traces import Traces
+from phoenix.datasets.dataset import EMPTY_DATASET, Dataset
 from phoenix.datasets.fixtures import FIXTURES, get_datasets
 from phoenix.server.app import create_app
+from phoenix.trace.fixtures import (
+    TRACES_FIXTURES,
+    _download_traces_fixture,
+    _get_trace_fixture_by_name,
+)
+from phoenix.trace.span_json_decoder import json_string_to_span
 
 logger = logging.getLogger(__name__)
 
 
-def _write_pid_file() -> None:
-    with open(_get_pid_file(), "w"):
-        pass
+def _write_pid_file_when_ready(
+    server: Server,
+    wait_up_to_seconds: float = 5,
+) -> None:
+    """Write PID file after server is started (or when time is up)."""
+    time_limit = time() + wait_up_to_seconds
+    while time() < time_limit and not server.should_exit and not server.started:
+        sleep(1e-3)
+    if time() >= time_limit and not server.started:
+        server.should_exit = True
+    _get_pid_file().touch()
 
 
 def _remove_pid_file() -> None:
-    try:
-        os.unlink(_get_pid_file())
-    except OSError as e:
-        if e.errno == errno.ENOENT:
-            # If the pid file doesn't exist, ignore and continue on since
-            # we are already in the desired end state; This should not happen
-            pass
-        else:
-            raise
+    _get_pid_file().unlink(missing_ok=True)
 
 
-def _get_pid_file() -> str:
-    return os.path.join(config.get_pids_path(), "%d" % os.getpid())
+def _get_pid_file() -> Path:
+    return get_pids_path() / str(os.getpid())
 
 
 if __name__ == "__main__":
     primary_dataset_name: str
     reference_dataset_name: Optional[str]
+    trace_dataset_name: Optional[str] = None
+
+    primary_dataset: Dataset = EMPTY_DATASET
+    reference_dataset: Optional[Dataset] = None
+    corpus_dataset: Optional[Dataset] = None
+
     # automatically remove the pid file when the process is being gracefully terminated
     atexit.register(_remove_pid_file)
-    _write_pid_file()
 
     parser = ArgumentParser()
     parser.add_argument("--export_path")
-    parser.add_argument("--port", type=int, default=config.PORT)
+    parser.add_argument("--host", type=str, required=False)
+    parser.add_argument("--port", type=int, required=False)
     parser.add_argument("--no-internet", action="store_true")
     parser.add_argument("--debug", action="store_false")  # TODO: Disable before public launch
     subparsers = parser.add_subparsers(dest="command", required=True)
@@ -55,11 +69,16 @@ if __name__ == "__main__":
     datasets_parser.add_argument("--primary", type=str, required=True)
     datasets_parser.add_argument("--reference", type=str, required=False)
     datasets_parser.add_argument("--corpus", type=str, required=False)
+    datasets_parser.add_argument("--trace", type=str, required=False)
     fixture_parser = subparsers.add_parser("fixture")
     fixture_parser.add_argument("fixture", type=str, choices=[fixture.name for fixture in FIXTURES])
     fixture_parser.add_argument("--primary-only", type=bool)
+    trace_fixture_parser = subparsers.add_parser("trace-fixture")
+    trace_fixture_parser.add_argument(
+        "fixture", type=str, choices=[fixture.name for fixture in TRACES_FIXTURES]
+    )
     args = parser.parse_args()
-    export_path = Path(args.export_path) if args.export_path else config.EXPORT_DIR
+    export_path = Path(args.export_path) if args.export_path else EXPORT_DIR
     if args.command == "datasets":
         primary_dataset_name = args.primary
         reference_dataset_name = args.reference
@@ -73,7 +92,7 @@ if __name__ == "__main__":
         corpus_dataset = (
             None if corpus_dataset_name is None else Dataset.from_name(corpus_dataset_name)
         )
-    else:
+    elif args.command == "fixture":
         fixture_name = args.fixture
         primary_only = args.primary_only
         primary_dataset, reference_dataset, corpus_dataset = get_datasets(
@@ -83,16 +102,33 @@ if __name__ == "__main__":
         if primary_only:
             reference_dataset_name = None
             reference_dataset = None
+    elif args.command == "trace-fixture":
+        trace_dataset_name = args.fixture
 
     model = create_model_from_datasets(
         primary_dataset,
         reference_dataset,
     )
+    traces = Traces()
+    if trace_dataset_name is not None:
+        for span in map(
+            json_string_to_span,
+            _download_traces_fixture(
+                _get_trace_fixture_by_name(
+                    trace_dataset_name,
+                ),
+            ),
+        ):
+            traces.put(span)
     app = create_app(
         export_path=export_path,
         model=model,
+        traces=traces,
         corpus=None if corpus_dataset is None else create_model_from_datasets(corpus_dataset),
         debug=args.debug,
     )
-
-    uvicorn.run(app, port=args.port)
+    host = args.host or get_env_host()
+    port = args.port or get_env_port()
+    server = Server(config=Config(app, host=host, port=port))
+    Thread(target=_write_pid_file_when_ready, args=(server,), daemon=True).start()
+    server.run()

phoenix/server/span_handler.py

@@ -0,0 +1,39 @@
+import gzip
+from typing import Protocol
+
+from starlette.endpoints import HTTPEndpoint
+from starlette.requests import Request
+from starlette.responses import Response
+
+from phoenix.trace.schemas import Span
+from phoenix.trace.span_json_decoder import json_to_span
+from phoenix.trace.v1 import encode
+from phoenix.trace.v1 import trace_pb2 as pb
+
+
+class SupportsPutSpan(Protocol):
+    def put(self, span: pb.Span) -> None:
+        ...
+
+
+class SpanHandler(HTTPEndpoint):
+    queue: SupportsPutSpan
+
+    async def post(self, request: Request) -> Response:
+        try:
+            content_type = request.headers.get("content-type")
+            if content_type == "application/x-protobuf":
+                body = await request.body()
+                content_encoding = request.headers.get("content-encoding")
+                if content_encoding == "gzip":
+                    body = gzip.decompress(body)
+                pb_span = pb.Span()
+                pb_span.ParseFromString(body)
+            else:
+                span = json_to_span(await request.json())
+                assert isinstance(span, Span)
+                pb_span = encode(span)
+        except Exception:
+            return Response(status_code=422)
+        self.queue.put(pb_span)
+        return Response()