# AUTOGENERATED! DO NOT EDIT! File to edit: ../../../nbs/ml.llm.obs.ipynb.

# %% auto 0
__all__ = ['TracePersister', 'make_phoenix_trace_callback_handler']

# %% ../../../nbs/ml.llm.obs.ipynb 3
from pathlib import Path
from phoenix.trace.llama_index import OpenInferenceTraceCallbackHandler
from phoenix.trace.schemas import Span
from phoenix.trace.span_json_encoder import spans_to_jsonl
from typing import Iterator

# %% ../../../nbs/ml.llm.obs.ipynb 4
class TracePersister:
    def __init__(self, output_filepath: Path):
        self.output_filepath = output_filepath
        self.output_filepath.parent.mkdir(parents=True, exist_ok=True)

    def __call__(self, spans: Iterator[Span]) -> None:
        with open(self.output_filepath, "a") as f:
            f.write("\n")
            f.write(spans_to_jsonl(spans))


def make_phoenix_trace_callback_handler(output_filepath: Path):
    return OpenInferenceTraceCallbackHandler(callback=TracePersister(output_filepath))
