Improve/simplify tests; avoid user prompt duplication

Author: cpsievert

chatlas/_chat.py

@@ -23,7 +23,6 @@
     Optional,
     Sequence,
     TypeVar,
-    cast,
     overload,
 )
 
@@ -56,7 +55,6 @@
 
 if TYPE_CHECKING:
     from inspect_ai.model import ChatMessage as InspectChatMessage
-    from inspect_ai.model import ChatMessageAssistant as InspectChatMessageAssistant
     from inspect_ai.solver import TaskState as InspectTaskState
 
     from ._content import ToolAnnotations
@@ -888,6 +886,7 @@ def my_eval(grader_model: str = "openai/gpt-4o"):
             inspect_content_as_chatlas,
             inspect_messages_as_turns,
             try_import_inspect,
+            turn_as_inspect_messages,
         )
 
         (imodel, isolver, _) = try_import_inspect()
@@ -928,18 +927,23 @@ async def solve(state: InspectTaskState, generate):
                 # we translate the message state back to the chat instance.
                 # N.B., state.message can include non-trivial dataset of sample input
                 # (e.g., `Sample(input=[ChatMessage, ...])`)
-                system_prompts: list["InspectChatMessage"] = []
-                other_prompts: list["InspectChatMessage"] = []
-                for x in state.messages:
-                    if x.role == "system":
-                        system_prompts.append(x)
+                system_messages: list["InspectChatMessage"] = []
+                other_messages: list["InspectChatMessage"] = []
+                user_prompt: "InspectChatMessage | None" = None
+                for x in reversed(state.messages):
+                    if x.role == "user" and user_prompt is None:
+                        user_prompt = x
+                    elif x.role == "system":
+                        system_messages.append(x)
                     else:
-                        other_prompts.append(x)
+                        other_messages.append(x)
+
+                other_messages.reverse()
 
                 # Set the system prompt on the chat instance
-                if len(system_prompts) == 1:
-                    chat_instance.system_prompt = str(system_prompts[0])
-                elif len(system_prompts) > 1:
+                if len(system_messages) == 1:
+                    chat_instance.system_prompt = str(system_messages[0])
+                elif len(system_messages) > 1:
                     raise ValueError(
                         "Multiple system prompts detected in `.to_solver()`, but chatlas only "
                         "supports a single system prompt. This usually indicates that the system "
@@ -949,26 +953,21 @@ async def solve(state: InspectTaskState, generate):
                     )
 
                 # Now, set the other messages as turns on the chat instance
-                chat_instance.set_turns(inspect_messages_as_turns(other_prompts))
+                chat_instance.set_turns(inspect_messages_as_turns(other_messages))
 
-                # TODO: inspect docs mention this is always the _first_? user message??
-                user_content = state.user_prompt.content
-                if isinstance(user_content, str):
-                    input_content = [user_content]
-                else:
-                    input_content = [
-                        inspect_content_as_chatlas(x) for x in user_content
-                    ]
+                if user_prompt is None:
+                    raise ValueError("No user prompt found in InspectAI state messages")
+
+                input_content = [inspect_content_as_chatlas(x) for x in user_prompt.content]
 
                 await chat_instance.chat_async(*input_content, echo="none")
                 last_turn = chat_instance.get_last_turn(role="assistant")
                 if last_turn is None:
                     raise ValueError("No assistant turn found after chat completion")
 
-                last_turn_message = cast(
-                    "InspectChatMessageAssistant",
-                    last_turn.to_inspect_messages(model)[0],
-                )
+                last_turn_message = turn_as_inspect_messages(
+                    last_turn, "assistant", model
+                )[0]
                 state.messages.append(last_turn_message)
 
                 tokens = last_turn.tokens

tests/integration/__init__.py

@@ -2,6 +2,8 @@
 import tempfile
 
 import pytest
+from pydantic import BaseModel
+
 from chatlas import (
     ChatOpenAI,
     ContentToolRequest,
@@ -10,7 +12,6 @@
     Turn,
 )
 from chatlas._chat import ToolFailureWarning
-from pydantic import BaseModel
 
 
 def test_simple_batch_chat():
@@ -19,157 +20,6 @@ def test_simple_batch_chat():
     assert str(response) == "2"
 
 
-def test_chat_to_solver_creates_solver():
-    pytest.importorskip("inspect_ai")
-
-    chat = ChatOpenAI()
-    solver = chat.to_solver()
-
-    assert callable(solver)
-
-
-def test_chat_to_solver_with_history():
-    pytest.importorskip("inspect_ai")
-
-    chat = ChatOpenAI(system_prompt="You are a helpful assistant.")
-    chat.set_turns(
-        [
-            Turn("user", "What is 2 + 2?"),
-            Turn("assistant", "4"),
-        ]
-    )
-
-    solver = chat.to_solver()
-
-    assert callable(solver)
-    assert len(chat.get_turns(include_system_prompt=False)) == 2
-    assert chat.system_prompt == "You are a helpful assistant."
-
-
-def test_chat_to_solver_with_rich_content():
-    pytest.importorskip("inspect_ai")
-    from chatlas import content_image_url
-
-    chat = ChatOpenAI(system_prompt="You analyze images.")
-
-    # Create a turn with mixed content (text + image)
-    image_content = content_image_url("https://example.com/image.jpg")
-    chat.set_turns(
-        [
-            Turn("user", ["Describe this image", image_content]),
-            Turn("assistant", "This is a test image."),
-        ]
-    )
-
-    solver = chat.to_solver()
-
-    assert callable(solver)
-    turns = chat.get_turns(include_system_prompt=False)
-    assert len(turns) == 2
-    assert len(turns[0].contents) == 2  # text + image
-    assert chat.system_prompt == "You analyze images."
-
-
-def test_chat_to_solver_with_tool_calls():
-    pytest.importorskip("inspect_ai")
-    from chatlas import ContentToolRequest, ContentToolResult
-
-    chat = ChatOpenAI()
-
-    # Create a turn with tool calls
-    tool_request = ContentToolRequest(
-        id="call_123", name="get_weather", arguments={"city": "NYC"}
-    )
-    tool_result = ContentToolResult(value="Sunny, 75°F", request=tool_request)
-
-    chat.set_turns(
-        [
-            Turn("user", "What's the weather in NYC?"),
-            Turn("assistant", [tool_request]),
-            Turn("user", [tool_result]),
-            Turn("assistant", "The weather in NYC is sunny and 75°F."),
-        ]
-    )
-
-    solver = chat.to_solver()
-
-    assert callable(solver)
-    turns = chat.get_turns(include_system_prompt=False)
-    assert len(turns) == 4
-    assert isinstance(turns[1].contents[0], ContentToolRequest)
-    assert isinstance(turns[2].contents[0], ContentToolResult)
-
-
-def test_chat_to_solver_without_inspect_ai():
-    import sys
-    from unittest.mock import patch
-
-    chat = ChatOpenAI()
-
-    # Mock inspect_ai as not installed in sys.modules
-    with patch.dict(sys.modules, {"inspect_ai": None, "inspect_ai.model": None}):
-        with pytest.raises(ImportError, match="pip install inspect-a"):
-            chat.to_solver()
-
-
-@pytest.mark.asyncio
-async def test_chat_to_solver_with_pdf_content():
-    pytest.importorskip("inspect_ai")
-    from chatlas._content import ContentPDF
-
-    chat = ChatOpenAI(system_prompt="You analyze documents.")
-
-    pdf_content = ContentPDF(data=b"Mock PDF data")
-    chat.set_turns(
-        [
-            Turn("user", ["Analyze this document", pdf_content]),
-            Turn("assistant", "This appears to be a PDF document."),
-        ]
-    )
-
-    solver = chat.to_solver()
-    assert callable(solver)
-
-    turns = chat.get_turns(include_system_prompt=False)
-    assert len(turns) == 2
-    assert any(isinstance(c, ContentPDF) for c in turns[0].contents)
-
-
-@pytest.mark.asyncio
-async def test_chat_to_solver_with_json_content():
-    pytest.importorskip("inspect_ai")
-    from chatlas._content import ContentJson
-
-    chat = ChatOpenAI()
-
-    json_content = ContentJson(value={"key": "value", "number": 42})
-    chat.set_turns(
-        [
-            Turn("user", "Process this data"),
-            Turn("assistant", [json_content]),
-        ]
-    )
-
-    solver = chat.to_solver()
-    assert callable(solver)
-
-    turns = chat.get_turns(include_system_prompt=False)
-    assert len(turns) == 2
-    assert any(isinstance(c, ContentJson) for c in turns[1].contents)
-
-
-def test_chat_to_solver_deepcopy_isolation():
-    pytest.importorskip("inspect_ai")
-
-    chat = ChatOpenAI()
-    initial_turns_count = len(chat.get_turns())
-
-    solver = chat.to_solver()
-
-    assert len(chat.get_turns()) == initial_turns_count
-    assert callable(solver)
-
-
 @pytest.mark.asyncio
 async def test_simple_async_batch_chat():
     chat = ChatOpenAI()