Add choices to output state

Author: cpsievert

chatlas/_chat.py

@@ -4,6 +4,7 @@
 import inspect
 import os
 import sys
+import time
 import traceback
 import warnings
 from pathlib import Path
@@ -868,9 +869,10 @@ def my_eval(grader_model: str = "openai/gpt-4o"):
         @isolver.solver("chatlas_solver")
         def _solver():
             async def solve(state: InspectTaskState, generate):
+                start_time = time.perf_counter()
                 if not state.messages:
                     for turn in chat_instance.get_turns():
-                        state.messages.append(*turn_as_messages(turn, model=model))
+                        state.messages.extend(turn_as_messages(turn, turn.role, model))
 
                 user_content = state.user_prompt.content
                 if isinstance(user_content, str):
@@ -885,19 +887,27 @@ async def solve(state: InspectTaskState, generate):
                 last_turn = chat_instance.get_last_turn(role="assistant")
                 if last_turn is None:
                     raise ValueError("No assistant turn found after chat completion")
-                state.messages.append(*turn_as_messages(last_turn, model=model))
-                tokens = last_turn.tokens or (0, 0, 0)
-                state.output = imodel.ModelOutput(
-                    model=model,
-                    # TODO: add choices?
-                    # choices=<choices>,
-                    completion=last_turn.text,
-                    usage=imodel.ModelUsage(
+
+                last_turn_message = turn_as_messages(last_turn, "assistant", model)[0]
+                state.messages.append(last_turn_message)
+
+                tokens = last_turn.tokens
+                if tokens is None:
+                    usage = None
+                else:
+                    usage = imodel.ModelUsage(
                         input_tokens=tokens[0],
                         output_tokens=tokens[1],
                         total_tokens=tokens[0] + tokens[1],
                         input_tokens_cache_read=tokens[2],
-                    ),
+                    )
+
+                state.output = imodel.ModelOutput(
+                    model=model,
+                    choices=[imodel.ChatCompletionChoice(message=last_turn_message)],
+                    completion=last_turn.text,
+                    usage=usage,
+                    time=time.perf_counter() - start_time,
                 )
                 return state
 

chatlas/_inspect.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Literal, overload
 
 from ._content import (
     Content,
@@ -19,8 +19,34 @@
     import inspect_ai.solver as isolver
     import inspect_ai.tool as itool
 
+Role = Literal["system", "user", "assistant"]
 
-def turn_as_messages(turn: Turn, model: str | None = None) -> list:
+
+@overload
+def turn_as_messages(
+    turn: Turn, role: Literal["system"], model: str | None = None
+) -> list[imodel.ChatMessageSystem]: ...
+
+
+@overload
+def turn_as_messages(
+    turn: Turn, role: Literal["user"], model: str | None = None
+) -> list[imodel.ChatMessage]: ...
+
+
+@overload
+def turn_as_messages(
+    turn: Turn, role: Literal["assistant"], model: str | None = None
+) -> list[imodel.ChatMessageAssistant]: ...
+
+
+def turn_as_messages(
+    turn: Turn, role: Role, model: str | None = None
+) -> (
+    list[imodel.ChatMessageSystem]
+    | list[imodel.ChatMessage]
+    | list[imodel.ChatMessageAssistant]
+):
     """
     Translate a chatlas Turn into InspectAI ChatMessages.