diff --git a/models.py b/models.py index fa890a3..5d35da3 100644 --- a/models.py +++ b/models.py @@ -20,3 +20,12 @@ class Track: id: str description: str tickets: List[Ticket] = field(default_factory=list) + +@dataclass +class WorkerContext: + """ + Represents the context provided to a Tier 3 Worker for a specific ticket. + """ + ticket_id: str + model_name: str + messages: List[dict] diff --git a/tests/test_mma_models.py b/tests/test_mma_models.py index c84389a..2b531b8 100644 --- a/tests/test_mma_models.py +++ b/tests/test_mma_models.py @@ -1,5 +1,5 @@ import pytest -from models import Ticket, Track +from models import Ticket, Track, WorkerContext def test_ticket_instantiation(): """ @@ -67,3 +67,25 @@ def test_track_can_handle_empty_tickets(): """ track = Track(id="TRACK-2", description="Empty Track", tickets=[]) assert track.tickets == [] + +def test_worker_context_instantiation(): + """ + Verifies that a WorkerContext can be instantiated with ticket_id, + model_name, and messages. + """ + ticket_id = "T1" + model_name = "gemini-2.0-flash-lite" + messages = [ + {"role": "user", "content": "Hello"}, + {"role": "assistant", "content": "Hi there!"} + ] + + context = WorkerContext( + ticket_id=ticket_id, + model_name=model_name, + messages=messages + ) + + assert context.ticket_id == ticket_id + assert context.model_name == model_name + assert context.messages == messages