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

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

# %% ../../nbs/llama_index.obs.ipynb 3
from io import TextIOBase
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/llama_index.obs.ipynb 4
class TraceRecorder:
    def __init__(self, f: TextIOBase | Path):
        if isinstance(f, Path):
            f.parent.mkdir(parents=True, exist_ok=True)
            self.f = open(f, "a")
        else:
            self.f = f

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


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