arize-phoenix 1.6.0__py3-none-any.whl → 1.7.1__py3-none-any.whl

phoenix/session/evaluation.py

@@ -6,11 +6,9 @@ A set of **highly experimental** helper functions to
   - ingest evaluation results into Phoenix via HttpExporter
 """
 import math
+from time import sleep
 from typing import (
     Any,
-    Iterable,
-    List,
-    Mapping,
     Optional,
     Sequence,
     Tuple,
@@ -20,56 +18,19 @@ from typing import (
 
 import pandas as pd
 from google.protobuf.wrappers_pb2 import DoubleValue, StringValue
+from tqdm import tqdm
 
 import phoenix.trace.v1 as pb
-from phoenix.core.traces import TRACE_ID
-from phoenix.session.session import Session
+from phoenix.trace.dsl.helpers import get_qa_with_reference, get_retrieved_documents
 from phoenix.trace.exporter import HttpExporter
-from phoenix.trace.schemas import ATTRIBUTE_PREFIX
-from phoenix.trace.semantic_conventions import (
-    DOCUMENT_CONTENT,
-    DOCUMENT_SCORE,
-    INPUT_VALUE,
-    RETRIEVAL_DOCUMENTS,
-)
 
+__all__ = [
+    "get_retrieved_documents",
+    "get_qa_with_reference",
+    "add_evaluations",
+]
 
-def get_retrieved_documents(session: Session) -> pd.DataFrame:
-    data: List[Mapping[str, Any]] = []
-    if (df := session.get_spans_dataframe("span_kind == 'RETRIEVER'")) is not None:
-        for span_id, query, documents, trace_id in df.loc[
-            :,
-            [
-                ATTRIBUTE_PREFIX + INPUT_VALUE,
-                ATTRIBUTE_PREFIX + RETRIEVAL_DOCUMENTS,
-                TRACE_ID,
-            ],
-        ].itertuples():
-            if not isinstance(documents, Iterable):
-                continue
-            for position, document in enumerate(documents):
-                if not hasattr(document, "get"):
-                    continue
-                data.append(
-                    {
-                        "context.trace_id": trace_id,
-                        "context.span_id": span_id,
-                        "input": query,
-                        "document_position": position,
-                        "reference": document.get(DOCUMENT_CONTENT),
-                        "document_score": document.get(DOCUMENT_SCORE),
-                    }
-                )
-    index = ["context.span_id", "document_position"]
-    columns = [
-        "context.span_id",
-        "document_position",
-        "input",
-        "reference",
-        "document_score",
-        "context.trace_id",
-    ]
-    return pd.DataFrame(data=data, columns=columns).set_index(index)
+from phoenix.trace.span_evaluations import Evaluations
 
 
 def add_evaluations(
@@ -155,3 +116,22 @@ def _extract_result(row: "pd.Series[Any]") -> Optional[pb.Evaluation.Result]:
         label=StringValue(value=label) if label else None,
         explanation=StringValue(value=explanation) if explanation else None,
     )
+
+
+def log_evaluations(
+    *evals: Evaluations,
+    endpoint: Optional[str] = None,
+    host: Optional[str] = None,
+    port: Optional[int] = None,
+) -> None:
+    if not (n := sum(map(len, evals))):
+        return
+    exporter = HttpExporter(endpoint=endpoint, host=host, port=port)
+    for eval in filter(bool, evals):
+        add_evaluations(exporter, eval.dataframe, eval.eval_name)
+    with tqdm(total=n, desc="Sending Evaluations") as pbar:
+        while n:
+            sleep(0.1)
+            n_left = exporter._queue.qsize()
+            n, diff = n_left, n - n_left
+            pbar.update(diff)

phoenix/session/session.py

@@ -30,6 +30,7 @@ from phoenix.server.app import create_app
 from phoenix.server.thread_server import ThreadServer
 from phoenix.services import AppService
 from phoenix.trace.dsl import SpanFilter
+from phoenix.trace.dsl.query import SpanQuery
 from phoenix.trace.span_json_encoder import span_to_json
 from phoenix.trace.trace_dataset import TraceDataset
 
@@ -167,6 +168,25 @@ class Session(ABC):
         """Returns the url for the phoenix app"""
         return _get_url(self.host, self.port, self.notebook_env)
 
+    def query_spans(
+        self,
+        *queries: SpanQuery,
+        start_time: Optional[datetime] = None,
+        stop_time: Optional[datetime] = None,
+    ) -> Optional[Union[pd.DataFrame, List[pd.DataFrame]]]:
+        if len(queries) == 0 or (traces := self.traces) is None:
+            return None
+        spans = tuple(
+            traces.get_spans(
+                start_time=start_time,
+                stop_time=stop_time,
+            )
+        )
+        dataframes = [query(spans) for query in queries]
+        if len(dataframes) == 1:
+            return dataframes[0]
+        return dataframes
+
     def get_spans_dataframe(
         self,
         filter_condition: Optional[str] = None,

phoenix/trace/__init__.py

@@ -1,4 +1,4 @@
-from .span_evaluations import SpanEvaluations
+from .span_evaluations import DocumentEvaluations, SpanEvaluations, TraceEvaluations
 from .trace_dataset import TraceDataset
 
-__all__ = ["TraceDataset", "SpanEvaluations"]
+__all__ = ["TraceDataset", "SpanEvaluations", "DocumentEvaluations", "TraceEvaluations"]

phoenix/trace/dsl/__init__.py

@@ -1,5 +1,7 @@
 from phoenix.trace.dsl.filter import SpanFilter
+from phoenix.trace.dsl.query import SpanQuery
 
 __all__ = [
     "SpanFilter",
+    "SpanQuery",
 ]

phoenix/trace/dsl/helpers.py

@@ -0,0 +1,61 @@
+from typing import List, Optional, Protocol, Union, cast
+
+import pandas as pd
+
+from phoenix.trace.dsl import SpanQuery
+from phoenix.trace.semantic_conventions import (
+    DOCUMENT_CONTENT,
+    DOCUMENT_SCORE,
+    INPUT_VALUE,
+    OUTPUT_VALUE,
+    RETRIEVAL_DOCUMENTS,
+)
+
+INPUT = {"input": INPUT_VALUE}
+OUTPUT = {"output": OUTPUT_VALUE}
+IO = {**INPUT, **OUTPUT}
+
+IS_ROOT = "parent_id is None"
+IS_LLM = "span_kind == 'LLM'"
+IS_RETRIEVER = "span_kind == 'RETRIEVER'"
+
+
+class Session(Protocol):
+    def query_spans(self, *query: SpanQuery) -> Optional[Union[pd.DataFrame, List[pd.DataFrame]]]:
+        ...
+
+
+def get_retrieved_documents(session: Session) -> pd.DataFrame:
+    return cast(
+        pd.DataFrame,
+        session.query_spans(
+            SpanQuery()
+            .where(IS_RETRIEVER)
+            .select("trace_id", **INPUT)
+            .explode(
+                RETRIEVAL_DOCUMENTS,
+                reference=DOCUMENT_CONTENT,
+                document_score=DOCUMENT_SCORE,
+            )
+        ),
+    )
+
+
+def get_qa_with_reference(session: Session) -> pd.DataFrame:
+    return pd.concat(
+        cast(
+            List[pd.DataFrame],
+            session.query_spans(
+                SpanQuery().select(**IO).where(IS_ROOT),
+                SpanQuery()
+                .where(IS_RETRIEVER)
+                .select(span_id="parent_id")
+                .concat(
+                    RETRIEVAL_DOCUMENTS,
+                    reference=DOCUMENT_CONTENT,
+                ),
+            ),
+        ),
+        axis=1,
+        join="inner",
+    )

phoenix/trace/dsl/query.py

@@ -0,0 +1,261 @@
+from collections import defaultdict
+from dataclasses import dataclass, field, fields, replace
+from functools import cached_property, partial
+from types import MappingProxyType
+from typing import Any, Callable, ClassVar, Dict, Iterable, Iterator, List, Mapping, Sequence, Tuple
+
+import pandas as pd
+
+from phoenix.trace.dsl import SpanFilter
+from phoenix.trace.schemas import ATTRIBUTE_PREFIX, CONTEXT_PREFIX, Span
+from phoenix.trace.semantic_conventions import RETRIEVAL_DOCUMENTS
+
+_SPAN_ID = "context.span_id"
+_PRESCRIBED_POSITION_PREFIXES = {
+    RETRIEVAL_DOCUMENTS: "document_",
+    ATTRIBUTE_PREFIX + RETRIEVAL_DOCUMENTS: "document_",
+}
+_ALIASES = {
+    "span_id": "context.span_id",
+    "trace_id": "context.trace_id",
+}
+
+# Because UUIDs is not convertible to Parquet,
+# they need to be converted to string.
+_CONVERT_TO_STRING = (
+    "context.span_id",
+    "context.trace_id",
+    "parent_id",
+)
+
+
+def _unalias(key: str) -> str:
+    return _ALIASES.get(key, key)
+
+
+@dataclass(frozen=True)
+class Projection:
+    key: str = ""
+    value: Callable[[Span], Any] = field(init=False, repr=False)
+    span_fields: ClassVar[Tuple[str, ...]] = tuple(f.name for f in fields(Span))
+
+    def __bool__(self) -> bool:
+        return bool(self.key)
+
+    def __post_init__(self) -> None:
+        key = _unalias(self.key)
+        object.__setattr__(self, "key", key)
+        if key.startswith(CONTEXT_PREFIX):
+            key = key[len(CONTEXT_PREFIX) :]
+            value = partial(self._from_context, key=key)
+        elif key.startswith(ATTRIBUTE_PREFIX):
+            key = self.key[len(ATTRIBUTE_PREFIX) :]
+            value = partial(self._from_attributes, key=key)
+        elif key in self.span_fields:
+            value = partial(self._from_span, key=key)
+        else:
+            value = partial(self._from_attributes, key=key)
+        if self.key in _CONVERT_TO_STRING:
+            object.__setattr__(
+                self,
+                "value",
+                lambda span: None if (v := value(span)) is None else str(v),
+            )
+        else:
+            object.__setattr__(self, "value", value)
+
+    def __call__(self, span: Span) -> Any:
+        return self.value(span)
+
+    @staticmethod
+    def _from_attributes(span: Span, key: str) -> Any:
+        return span.attributes.get(key)
+
+    @staticmethod
+    def _from_context(span: Span, key: str) -> Any:
+        return getattr(span.context, key, None)
+
+    @staticmethod
+    def _from_span(span: Span, key: str) -> Any:
+        return getattr(span, key, None)
+
+
+@dataclass(frozen=True)
+class Explosion(Projection):
+    kwargs: Mapping[str, str] = field(default_factory=lambda: MappingProxyType({}))
+    primary_index_key: str = "context.span_id"
+
+    position_prefix: str = field(init=False, repr=False)
+    primary_index: Projection = field(init=False, repr=False)
+
+    def __post_init__(self) -> None:
+        super().__post_init__()
+        position_prefix = _PRESCRIBED_POSITION_PREFIXES.get(self.key, "")
+        object.__setattr__(self, "position_prefix", position_prefix)
+        object.__setattr__(self, "primary_index", Projection(self.primary_index_key))
+
+    @cached_property
+    def index_keys(self) -> Tuple[str, str]:
+        return (self.primary_index.key, f"{self.position_prefix}position")
+
+    def with_primary_index_key(self, primary_index_key: str) -> "Explosion":
+        return replace(self, primary_index_key=primary_index_key)
+
+    def __call__(self, span: Span) -> Iterator[Dict[str, Any]]:
+        if not isinstance(seq := self.value(span), Sequence):
+            return
+        has_mapping = False
+        for item in seq:
+            if isinstance(item, Mapping):
+                has_mapping = True
+                break
+        if not has_mapping:
+            for i, item in enumerate(seq):
+                if item is not None:
+                    yield {
+                        self.key: item,
+                        self.primary_index.key: self.primary_index(span),
+                        f"{self.position_prefix}position": i,
+                    }
+            return
+        for i, item in enumerate(seq):
+            if not isinstance(item, Mapping):
+                continue
+            record = (
+                {name: item.get(key) for name, key in self.kwargs.items()}
+                if self.kwargs
+                else dict(item)
+            )
+            for v in record.values():
+                if v is not None:
+                    break
+            else:
+                record = {}
+            if not record:
+                continue
+            record[self.primary_index.key] = self.primary_index(span)
+            record[f"{self.position_prefix}position"] = i
+            yield record
+
+
+@dataclass(frozen=True)
+class Concatenation(Projection):
+    kwargs: Mapping[str, str] = field(default_factory=lambda: MappingProxyType({}))
+    separator: str = "\n\n"
+
+    def with_separator(self, separator: str = "\n\n") -> "Concatenation":
+        return replace(self, separator=separator)
+
+    def __call__(self, span: Span) -> Iterator[Tuple[str, str]]:
+        if not isinstance(seq := self.value(span), Sequence):
+            return
+        if not self.kwargs:
+            yield self.key, self.separator.join(map(str, seq))
+        record = defaultdict(list)
+        for item in seq:
+            if not isinstance(item, Mapping):
+                continue
+            for k, v in self.kwargs.items():
+                if value := item.get(v):
+                    record[k].append(value)
+        for name, values in record.items():
+            yield name, self.separator.join(map(str, values))
+
+
+@dataclass(frozen=True)
+class SpanQuery:
+    _select: Mapping[str, Projection] = field(default_factory=lambda: MappingProxyType({}))
+    _concat: Concatenation = field(default_factory=Concatenation)
+    _explode: Explosion = field(default_factory=Explosion)
+    _filter: SpanFilter = field(default_factory=SpanFilter)
+    _rename: Mapping[str, str] = field(default_factory=lambda: MappingProxyType({}))
+    _index: Projection = field(default_factory=lambda: Projection("context.span_id"))
+
+    def __bool__(self) -> bool:
+        return bool(self._select) or bool(self._filter) or bool(self._explode) or bool(self._concat)
+
+    def select(self, *args: str, **kwargs: str) -> "SpanQuery":
+        _select = {
+            _unalias(name): Projection(key) for name, key in (*zip(args, args), *kwargs.items())
+        }
+        return replace(self, _select=MappingProxyType(_select))
+
+    def where(self, condition: str) -> "SpanQuery":
+        _filter = SpanFilter(condition)
+        return replace(self, _filter=_filter)
+
+    def explode(self, key: str, **kwargs: str) -> "SpanQuery":
+        _explode = Explosion(key=key, kwargs=kwargs, primary_index_key=self._index.key)
+        return replace(self, _explode=_explode)
+
+    def concat(self, key: str, **kwargs: str) -> "SpanQuery":
+        _concat = Concatenation(key=key, kwargs=kwargs)
+        return replace(self, _concat=_concat)
+
+    def rename(self, **kwargs: str) -> "SpanQuery":
+        _rename = MappingProxyType(kwargs)
+        return replace(self, _rename=_rename)
+
+    def with_index(self, key: str = "context.span_id") -> "SpanQuery":
+        _index = Projection(key=key)
+        return replace(self, _index=_index)
+
+    def with_concat_separator(self, separator: str = "\n\n") -> "SpanQuery":
+        _concat = self._concat.with_separator(separator)
+        return replace(self, _concat=_concat)
+
+    def with_explode_primary_index_key(self, primary_index_key: str) -> "SpanQuery":
+        _explode = self._explode.with_primary_index_key(primary_index_key)
+        return replace(self, _explode=_explode)
+
+    def __call__(self, spans: Iterable[Span]) -> pd.DataFrame:
+        if self._filter:
+            spans = filter(self._filter, spans)
+        if self._explode:
+            spans = filter(
+                lambda span: (isinstance(seq := self._explode.value(span), Sequence) and len(seq)),
+                spans,
+            )
+        if self._concat:
+            spans = filter(
+                lambda span: (isinstance(seq := self._concat.value(span), Sequence) and len(seq)),
+                spans,
+            )
+        _selected: List[Dict[str, Any]] = []
+        _exploded: List[Dict[str, Any]] = []
+        for span in spans:
+            if self._select:
+                record = {name: proj(span) for name, proj in self._select.items()}
+                for v in record.values():
+                    if v is not None:
+                        break
+                else:
+                    record = {}
+                if self._concat:
+                    record.update(self._concat(span))
+                if record:
+                    if self._index.key not in record:
+                        record[self._index.key] = self._index(span)
+                    _selected.append(record)
+            elif self._concat:
+                record = {self._index.key: self._index(span)}
+                record.update(self._concat(span))
+                if record:
+                    _selected.append(record)
+            if self._explode:
+                _exploded.extend(self._explode(span))
+        if _selected:
+            select_df = pd.DataFrame(_selected)
+        else:
+            select_df = pd.DataFrame(columns=[self._index.key])
+        select_df = select_df.set_index(self._index.key)
+        if self._explode:
+            if _exploded:
+                explode_df = pd.DataFrame(_exploded)
+            else:
+                explode_df = pd.DataFrame(columns=self._explode.index_keys)
+            explode_df = explode_df.set_index(list(self._explode.index_keys))
+            if not self._select:
+                return explode_df.rename(self._rename, axis=1, errors="ignore")
+            select_df = select_df.join(explode_df, how="outer")
+        return select_df.rename(self._rename, axis=1, errors="ignore")

phoenix/trace/llama_index/callback.py

@@ -27,7 +27,7 @@ from typing import (
     Union,
     cast,
 )
-from uuid import uuid4
+from uuid import UUID, uuid4
 
 import llama_index
 from llama_index.callbacks.base_handler import BaseCallbackHandler
@@ -40,8 +40,12 @@ from llama_index.callbacks.schema import (
 from llama_index.llms.base import ChatMessage, ChatResponse
 from llama_index.response.schema import Response, StreamingResponse
 from llama_index.tools import ToolMetadata
+from typing_extensions import TypeGuard
 
 from phoenix.trace.exporter import HttpExporter
+from phoenix.trace.llama_index.streaming import (
+    instrument_streaming_response as _instrument_streaming_response,
+)
 from phoenix.trace.schemas import (
     Span,
     SpanEvent,
@@ -49,6 +53,7 @@ from phoenix.trace.schemas import (
     SpanID,
     SpanKind,
     SpanStatusCode,
+    TraceID,
 )
 from phoenix.trace.semantic_conventions import (
     DOCUMENT_CONTENT,
@@ -109,6 +114,10 @@ class CBEventData:
     start_event: Optional[CBEvent] = field(default=None)
     end_event: Optional[CBEvent] = field(default=None)
     attributes: Dict[str, Any] = field(default_factory=dict)
+    span_id: Optional[CBEventID] = field(default=None)
+    parent_id: Optional[CBEventID] = field(default=None)
+    trace_id: Optional[TraceID] = field(default=None)
+    streaming_event: bool = field(default=False)
 
     def set_if_unset(self, key: str, value: Any) -> None:
         if not getattr(self, key):
@@ -270,9 +279,16 @@ class OpenInferenceTraceCallbackHandler(BaseCallbackHandler):
         **kwargs: Any,
     ) -> CBEventID:
         event_id = event_id or str(uuid4())
+        if parent_data := self._event_id_to_event_data.get(parent_id):
+            trace_id = parent_data.trace_id
+        else:
+            trace_id = uuid4()
         event_data = self._event_id_to_event_data[event_id]
         event_data.name = event_type.value
         event_data.event_type = event_type
+        event_data.parent_id = None if parent_id == "root" else parent_id
+        event_data.span_id = event_id
+        event_data.trace_id = trace_id
         event_data.start_event = CBEvent(
             event_type=event_type,
             payload=payload,
@@ -308,6 +324,10 @@ class OpenInferenceTraceCallbackHandler(BaseCallbackHandler):
             event_data.attributes.update(
                 payload_to_semantic_attributes(event_type, payload, is_event_end=True),
             )
+            response = payload.get(EventPayload.RESPONSE)
+            if _is_streaming_response(response):
+                event_data.streaming_event = True
+                response = _instrument_streaming_response(response, self._tracer, event_data)
 
     @graceful_fallback(_null_fallback)
     def start_trace(self, trace_id: Optional[str] = None) -> None:
@@ -353,7 +373,6 @@ def _add_spans_to_tracer(
         tracer (Tracer): The tracer that stores spans.
     """
 
-    trace_id = uuid4()
     parent_child_id_stack: List[Tuple[Optional[SpanID], CBEventID]] = [
         (None, root_event_id) for root_event_id in trace_map["root"]
     ]
@@ -383,7 +402,11 @@ def _add_spans_to_tracer(
 
         start_time = _timestamp_to_tz_aware_datetime(start_event.time)
         span_exceptions = _get_span_exceptions(event_data, start_time)
-        end_time = _get_end_time(event_data, span_exceptions)
+        if event_data.streaming_event:
+            # Do not set the end time for streaming events so we can update the event later
+            end_time = None
+        else:
+            end_time = _get_end_time(event_data, span_exceptions)
         start_time = start_time or end_time or datetime.now(timezone.utc)
 
         name = event_name if (event_name := event_data.name) is not None else "unknown"
@@ -391,7 +414,7 @@ def _add_spans_to_tracer(
         span = tracer.create_span(
             name=name,
             span_kind=span_kind,
-            trace_id=trace_id,
+            trace_id=event_data.trace_id,
             start_time=start_time,
             end_time=end_time,
             status_code=SpanStatusCode.ERROR if span_exceptions else SpanStatusCode.OK,
@@ -400,6 +423,7 @@ def _add_spans_to_tracer(
             attributes=attributes,
             events=sorted(span_exceptions, key=lambda event: event.timestamp) or None,
             conversation=None,
+            span_id=UUID(event_data.span_id),
         )
         new_parent_span_id = span.context.span_id
         for new_child_event_id in trace_map.get(event_id, []):
@@ -640,3 +664,7 @@ def _get_tool_call(tool_call: object) -> Iterator[Tuple[str, Any]]:
         if arguments := getattr(function, "arguments", None):
             assert isinstance(arguments, str), f"arguments must be str, found {type(arguments)}"
             yield TOOL_CALL_FUNCTION_ARGUMENTS_JSON, arguments
+
+
+def _is_streaming_response(response: Any) -> TypeGuard[StreamingResponse]:
+    return isinstance(response, StreamingResponse)

phoenix/trace/llama_index/streaming.py

@@ -0,0 +1,93 @@
+from datetime import datetime, timezone
+from typing import TYPE_CHECKING, Generator, List
+from uuid import UUID
+
+from llama_index.callbacks.schema import TIMESTAMP_FORMAT
+from llama_index.response.schema import StreamingResponse
+
+from phoenix.trace.schemas import SpanKind, SpanStatusCode
+from phoenix.trace.semantic_conventions import OUTPUT_VALUE
+from phoenix.trace.tracer import Tracer
+
+if TYPE_CHECKING:
+    from phoenix.trace.llama_index.callback import CBEventData
+
+_LOCAL_TZINFO = datetime.now().astimezone().tzinfo
+
+
+class TokenGenInstrumentor:
+    def __init__(
+        self, stream: Generator[str, None, None], tracer: Tracer, event_data: "CBEventData"
+    ):
+        self._stream = stream
+        self._token_stream: List[str] = []
+        self._finished = False
+        self._tracer = tracer
+        self._event_data = event_data
+
+    def __iter__(self) -> "TokenGenInstrumentor":
+        return self
+
+    def __next__(self) -> str:
+        if self._finished:
+            raise StopIteration
+
+        try:
+            value = next(self._stream)
+            self._token_stream.append(value)
+            return value
+        except StopIteration:
+            self._finished = True
+            self._handle_end_of_stream()
+            raise
+
+    def _handle_end_of_stream(self) -> None:
+        # Handle the end-of-stream logic here
+        parent_id = self._event_data.parent_id
+        output = "".join(self._token_stream)
+        start_event = self._event_data.start_event
+        if start_event:
+            start_time = _timestamp_to_tz_aware_datetime(start_event.time)
+        else:
+            start_time = datetime.now(timezone.utc)
+        attributes = self._event_data.attributes
+        attributes.update({OUTPUT_VALUE: output})
+        self._tracer.create_span(
+            name=self._event_data.name if self._event_data.name else "",
+            span_kind=SpanKind.LLM,
+            trace_id=self._event_data.trace_id,
+            start_time=start_time,
+            end_time=datetime.now(timezone.utc),
+            status_code=SpanStatusCode.OK,
+            status_message="",
+            parent_id=UUID(parent_id) if parent_id else None,
+            attributes=self._event_data.attributes,
+            events=[],
+            conversation=None,
+            span_id=UUID(self._event_data.span_id),
+        )
+
+
+def instrument_streaming_response(
+    response: StreamingResponse,
+    tracer: Tracer,
+    event_data: "CBEventData",
+) -> StreamingResponse:
+    if response.response_gen is not None:
+        response.response_gen = TokenGenInstrumentor(response.response_gen, tracer, event_data)  # type: ignore
+    return response
+
+
+def _timestamp_to_tz_aware_datetime(timestamp: str) -> datetime:
+    """Converts a timestamp string to a timezone-aware datetime."""
+    return _tz_naive_to_tz_aware_datetime(_timestamp_to_tz_naive_datetime(timestamp))
+
+
+def _timestamp_to_tz_naive_datetime(timestamp: str) -> datetime:
+    """Converts a timestamp string to a timezone-naive datetime."""
+    return datetime.strptime(timestamp, TIMESTAMP_FORMAT)
+
+
+def _tz_naive_to_tz_aware_datetime(timestamp: datetime) -> datetime:
+    """Converts a timezone-naive datetime to a timezone-aware datetime."""
+    return timestamp.replace(tzinfo=_LOCAL_TZINFO)