claude sequence 1. "in inspect_ai/ generate a class member diagram of TaskState" 2. "write this to a markdown file" 3. "include the fields of referenced classes" 4. "only include "public" properties of each class. I don't want private stuff (starts with underscore) nor non-property methods"

TaskState.md

TaskState Class Member Diagram

This document provides a visual representation of the TaskState class structure from inspect_ai/src/inspect_ai/solver/_task_state.py and all its referenced classes.

Complete Class Diagram

classDiagram
    class TaskState {
        +choices : Choices
        +tool_choice : ToolChoice|None
        +model : ModelName
        +sample_id : int|str
        +epoch : int
        +input : str|list~ChatMessage~
        +input_text : str
        +user_prompt : ChatMessageUser
        +metadata : dict~str,Any~
        +messages : list~ChatMessage~
        +output : ModelOutput
        +store : Store
        +tools : list~Tool~
        +max_messages : int|None
        +message_limit : int|None
        +token_limit : int|None
        +token_usage : int
        +completed : bool
        +target : Target
        +scores : dict~str,Score~|None
        +uuid : str
    }

    class Choice {
        +value : str
        +correct : bool|None
        +original_position : int
    }

    class Choices {
    }

    class ModelName {
        +api : str
        +name : str
    }

    class Target {
        +target : list~str~
        +text : str
    }

    class ChatMessageBase {
        +id : str|None
        +content : str|list~Content~
        +source : Literal~"input","generate"~|None
        +metadata : dict~str,Any~|None
        +text : str
    }

    class ChatMessageSystem {
        +role : Literal~"system"~
    }

    class ChatMessageUser {
        +role : Literal~"user"~
        +tool_call_id : list~str~|None
    }

    class ChatMessageAssistant {
        +role : Literal~"assistant"~
        +tool_calls : list~ToolCall~|None
        +model : str|None
    }

    class ChatMessageTool {
        +role : Literal~"tool"~
        +tool_call_id : str|None
        +function : str|None
        +error : ToolCallError|None
        +tool_error : str|None
    }

    class ModelOutput {
        +model : str
        +choices : list~ChatCompletionChoice~
        +completion : str
        +usage : ModelUsage|None
        +time : float|None
        +metadata : dict~str,Any~|None
        +error : str|None
        +empty : bool
        +stop_reason : StopReason
        +message : ChatMessageAssistant
    }

    class ModelUsage {
        +input_tokens : int
        +output_tokens : int
        +total_tokens : int
        +input_tokens_cache_write : int|None
        +input_tokens_cache_read : int|None
        +reasoning_tokens : int|None
    }

    class ChatCompletionChoice {
        +message : ChatMessageAssistant
        +stop_reason : StopReason
        +logprobs : Logprobs|None
    }

    class Store {
    }

    class Score {
        +value : Value
        +answer : str|None
        +explanation : str|None
        +metadata : dict~str,Any~|None
        +text : str
    }

    class Tool {
    }

    class ToolFunction {
        +name : str
    }

    TaskState *-- Choices : contains
    Choices *-- Choice : contains
    TaskState --> ModelName : uses
    TaskState --> Target : uses
    TaskState --> ChatMessageBase : uses
    TaskState --> Tool : uses
    TaskState --> ModelOutput : uses
    TaskState --> Store : uses
    TaskState --> Score : uses
    TaskState --> ToolChoice : uses

    ChatMessageSystem --|> ChatMessageBase : inherits
    ChatMessageUser --|> ChatMessageBase : inherits
    ChatMessageAssistant --|> ChatMessageBase : inherits
    ChatMessageTool --|> ChatMessageBase : inherits

    ModelOutput *-- ChatCompletionChoice : contains
    ModelOutput *-- ModelUsage : contains
    ChatCompletionChoice --> ChatMessageAssistant : uses

    ToolChoice --|> ToolFunction : can be

Loading

Class Structure Summary

Core Classes

TaskState Class

The main class representing the internal state of a Task being run for a single Sample.

Public Properties:

• choices: Choices - Multiple choice options (if applicable)
• tool_choice: ToolChoice|None - Tool choice directive
• model: ModelName - Name of model being evaluated
• sample_id: int|str - Unique id for sample
• epoch: int - Epoch number for sample
• input: str|list[ChatMessage] - Input from the Sample (immutable)
• input_text: str - Convenience property for input as text
• user_prompt: ChatMessageUser - User prompt for this state
• metadata: dict[str,Any] - Metadata from the Sample
• messages: list[ChatMessage] - Chat conversation history
• output: ModelOutput - The 'final' model output
• store: Store - Store for shared data
• tools: list[Tool] - Tools available to the model
• max_messages: int|None - Deprecated (use message_limit)
• message_limit: int|None - Limit on total messages allowed
• token_limit: int|None - Limit on total tokens allowed
• token_usage: int - Total tokens used for current sample
• completed: bool - Task completion status
• target: Target - The scoring target for this Sample
• scores: dict[str,Score]|None - Scores yielded by running task
• uuid: str - Globally unique identifier for sample run

Choice & Choices Classes

Handle multiple choice question functionality:

• Choice: Represents individual choices with value, correctness, and position
• Choices: Collection wrapper with shuffling and prompting capabilities

Referenced Classes

Model-Related Classes

• ModelName: Represents model API and name with pattern matching
• ModelOutput: Complete model response with choices, usage, and metadata
• ModelUsage: Token usage statistics including caching and reasoning tokens
• ChatCompletionChoice: Individual completion choice with message and stop reason

Message Classes

• ChatMessageBase: Base class for all chat messages with content and metadata
• ChatMessageSystem: System messages with "system" role
• ChatMessageUser: User messages with "user" role and optional tool call IDs
• ChatMessageAssistant: Assistant messages with "assistant" role and tool calls
• ChatMessageTool: Tool response messages with "tool" role and error handling

Utility Classes

• Target: Sequence-like class for scoring targets
• Store: Key-value store for shared state with change tracking
• Score: Scoring results with multiple access methods (str, int, float, bool, list, dict)
• Tool: Protocol for tools that agents can use
• ToolFunction: Specific tool function specification
• ToolChoice: Union type for tool selection ("auto", "any", "none", or specific function)

Key Relationships

• Composition: TaskState contains Choices, ModelOutput contains ChatCompletionChoice and ModelUsage
• Inheritance: All chat message types inherit from ChatMessageBase
• Usage: TaskState references all major classes for comprehensive state management
• Polymorphism: ToolChoice can be either a string literal or ToolFunction instance