Skip to content

LLM - Evaluation: Unit Test

Install packages

1
2
3
4
!uv pip install -q \
    litellm==1.78.5 \
    python-dotenv==1.1.1 \
    pydantic==2.12.3

Import packages

1
2
3
4
5
6
7
8
9
import io
import json
from unittest import TestCase, TestLoader, TestSuite, TextTestRunner

import litellm  # type: ignore
from dotenv import load_dotenv  # type: ignore
from pydantic import BaseModel  # type: ignore

load_dotenv()

True

Define structured output

1
2
3
4
class CalendarEvent(BaseModel):
    name: str
    date: str
    participants: list[str]

Define llm completion

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
def get_calendar_event(message):
    messages = [
        {"role": "system", "content": "Extract the event information."},
        {"role": "user", "content": message},
    ]

    completion = litellm.completion(
        model="gemini/gemini-2.0-flash",
        messages=messages,
        response_format=CalendarEvent,
    )

    content = completion.choices[0].message.content

    data = json.loads(content)

    return CalendarEvent.model_validate(data)

Create a test case

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
class TestGetCalendarEvent(TestCase):
    def test_calendar_event(self):
        # Arrange
        message = "Alice and Bob are going to a science fair on friday."

        # Act
        calendar_event = get_calendar_event(message)

        # Assert
        self.assertIsInstance(calendar_event, CalendarEvent)
        self.assertIsInstance(calendar_event.participants, list)

        self.assertEqual(calendar_event.name.lower(), "science fair")
        self.assertEqual(calendar_event.date.lower(), "friday")
        self.assertEqual(calendar_event.participants, ["Alice", "Bob"])
  • Create a test suit that will allow to run one or more test cases
  • Run the test cases
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
loader = TestLoader()
suite = TestSuite()

buffer = io.StringIO()
suite.addTest(loader.loadTestsFromTestCase(TestGetCalendarEvent))

TextTestRunner(stream=buffer, verbosity=2).run(suite)

output = buffer.getvalue()

print(output)

test_calendar_event (__main__.TestGetCalendarEvent.test_calendar_event) ... ok



----------------------------------------------------------------------

Ran 1 test in 1.775s



OK