Context manager to automatically run span evaluators

[omrihar]

I have seen the examples of how to run span evaluations in a cron job, or alternatively from a notebook and these seem to be a bit less robust (e.g. if a job fails to run in a cron-job, the time-dependent logic might break). That's why I thought of creating a context manager that runs the evaluations in a thread. I'm sharing it here to get feedback and in case anyone finds it useful :)

class automated_evals(ContextDecorator):
    """Run automated evals after the model exited."""

    def __init__(self, phoenix_project: str = None):
        self.phoenix_project = phoenix_project or PHOENIX_PROJECT

    def __enter__(self):
        """Log the start time of the LLM execution thread."""
        self.start_time = datetime.now()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        """Run automated evals after the LLM is done."""
        eval_thread = threading.Thread(target=self.run_evals)
        eval_thread.start()

    def run_evals(self):
        """Run the automated evaluations."""
        print(f"PHOENIX PROJECT: {self.phoenix_project}")
        time.sleep(5)
        phoenix_client = px.Client()
        qa_spans_df = get_qa_with_reference(
            phoenix_client,
            project_name=self.phoenix_project,
            start_time=self.start_time,
        )

        retriever_spans_df = get_retrieved_documents(
            phoenix_client,
            project_name=self.phoenix_project,
            start_time=self.start_time,
        )

        eval_model = BedrockModel(
            model_id=SONNET_35,
            temperature=0.0,
        )

        hallucination_evaluator = HallucinationEvaluator(eval_model)
        qa_correctness_evaluator = QAEvaluator(eval_model)
        relevance_evaluator = RelevanceEvaluator(eval_model)
        [hallucination_evals_df, qa_correctness_evals_df] = run_evals(
            qa_spans_df,
            [hallucination_evaluator, qa_correctness_evaluator],
        )

        relevance_evals_df = run_evals(
            retriever_spans_df,
            [relevance_evaluator],
        )[0]

        phoenix_client.log_evaluations(
            SpanEvaluations(
                eval_name="Hallucination", dataframe=hallucination_evals_df
            ),
            SpanEvaluations(
                eval_name="QA Correctness", dataframe=qa_correctness_evals_df
            ),
            DocumentEvaluations(eval_name="Relevance", dataframe=relevance_evals_df),
        )

All you need to do to use it in your code is:

@automated_evals
def call_model(input: str) -> str:
    # logic to call your model

Any comments or suggestions would be very welcome!

[axiomofjoy]

Thanks @omrihar!