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

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

# %% ../../nbs/llama_index.obs.ipynb 3
from uuid import uuid4
from datetime import datetime
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 phoenix.trace.tracer import SpanExporter

# %% ../../nbs/llama_index.obs.ipynb 4
class FileSpanExporter(SpanExporter):
    def __init__(self, filepath: Path | str | None = None):
        if filepath is None:
            filepath = Path("/tmp/llm-traces") / self._generate_filename()
        self.filepath = Path(filepath)
        self.filepath.parent.mkdir(exist_ok=True, parents=True)

    def export(self, span: Span) -> None:
        self.filepath.parent.mkdir(exist_ok=True, parents=True)
        with open(self.filepath, "a") as f:
            f.write("\n")
            f.write(spans_to_jsonl([span]))
            f.write("\n")

    def _generate_filename(self):
        return f"{datetime.utcnow().strftime('%Y_%m_%dT%H_%M_%SZ')}-{uuid4()}.jsonl"


# %% ../../nbs/llama_index.obs.ipynb 5
def make_phoenix_trace_callback_handler(output_filepath: Path | str | None = None):
    return OpenInferenceTraceCallbackHandler(exporter=FileSpanExporter(output_filepath))
