import mlflow
from mlflow.genai import make_judge

mlflow.set_tracking_uri("databricks")
mlflow.set_experiment(experiment_id="3011836326718646")


def test_litellm_adapter_openai():
    judge = make_judge(
        name="test_litellm_judge",
        instructions="Evaluate if the {{ inputs }} matches {{ outputs }}",
        model="openai:/gpt-4o-mini",
    )
    result = judge(inputs="What is 2+2?", outputs="4")
    print(f"LiteLLM: {result}")


def test_databricks_serving_endpoint_adapter():
    judge = make_judge(
        name="test_databricks_endpoint_judge",
        instructions="Evaluate if the {{ inputs }} matches {{ outputs }}",
        model="databricks:/databricks-claude-sonnet-4-5",
    )
    result = judge(inputs="What is 2+2?", outputs="4")
    print(f"Databricks Endpoint: {result}")


def test_databricks_managed_judge_adapter():
    judge = make_judge(
        name="test_databricks_default_judge",
        instructions="Evaluate if the {{ inputs }} matches {{ outputs }}",
        model="databricks",
    )
    result = judge(inputs="What is 2+2?", outputs="4")
    print(f"Databricks Managed: {result}")


if __name__ == "__main__":
    test_litellm_adapter_openai()
    test_databricks_serving_endpoint_adapter()
    test_databricks_managed_judge_adapter()

LiteLLM: Feedback(name='test_litellm_judge', source=AssessmentSource(source_type='LLM_JUDGE', source_id='openai:/gpt-4o-mini'), trace_id=None, run_id=None, rationale="The output correctly answers the query by providing the accurate result of the mathematical expression '2+2', which is indeed '4'. This matches the expected response for the given input.", metadata={'mlflow.assessment.judgeCost': 5.46e-05}, span_id=None, create_time_ms=1765291025597, last_update_time_ms=1765291025597, assessment_id=None, error=None, expectation=None, feedback=FeedbackValue(value='Pass', error=None), overrides=None, valid=True)
Databricks Endpoint: Feedback(name='test_databricks_endpoint_judge', source=AssessmentSource(source_type='LLM_JUDGE', source_id='databricks:/databricks-claude-sonnet-4-5'), trace_id=None, run_id=None, rationale='The input asks "What is 2+2?" and the output provides "4", which is the correct mathematical answer to this basic arithmetic question. The output accurately and directly answers the question posed in the input. This is a clear match between the question and its correct answer.', metadata={'mlflow.assessment.judgeCost': 0.00405599383}, span_id=None, create_time_ms=1765291029632, last_update_time_ms=1765291029632, assessment_id=None, error=None, expectation=None, feedback=FeedbackValue(value='Match', error=None), overrides=None, valid=True)
Databricks Managed: Feedback(name='test_databricks_default_judge', source=AssessmentSource(source_type='LLM_JUDGE', source_id='databricks'), trace_id=None, run_id=None, rationale="The output '4' is the correct answer to the mathematical query 'What is 2+2?'. The evaluation confirms that the AI agent accurately provided the expected result for the given input.", metadata=None, span_id=None, create_time_ms=1765291031518, last_update_time_ms=1765291031518, assessment_id=None, error=None, expectation=None, feedback=FeedbackValue(value='Correct', error=None), overrides=None, valid=True)