inspect-ai 0.3.81__py3-none-any.whl → 0.3.83__py3-none-any.whl

inspect_ai/agent/__init__.py

@@ -0,0 +1,36 @@
+from ._agent import Agent, AgentState, agent, agent_with
+from ._as_solver import as_solver
+from ._as_tool import as_tool
+from ._bridge.bridge import bridge
+from ._filter import MessageFilter, last_message, remove_tools
+from ._handoff import handoff
+from ._human.agent import human_cli
+from ._react import react
+from ._run import run
+from ._types import (
+    AgentAttempts,
+    AgentContinue,
+    AgentPrompt,
+    AgentSubmit,
+)
+
+__all__ = [
+    "react",
+    "bridge",
+    "human_cli",
+    "run",
+    "handoff",
+    "as_tool",
+    "as_solver",
+    "last_message",
+    "remove_tools",
+    "MessageFilter",
+    "Agent",
+    "AgentState",
+    "agent",
+    "agent_with",
+    "AgentPrompt",
+    "AgentAttempts",
+    "AgentContinue",
+    "AgentSubmit",
+]

inspect_ai/agent/_agent.py

@@ -0,0 +1,268 @@
+from copy import copy, deepcopy
+from functools import wraps
+from typing import (
+    Any,
+    Callable,
+    ParamSpec,
+    Protocol,
+    TypeGuard,
+    cast,
+    overload,
+    runtime_checkable,
+)
+
+from inspect_ai._util.registry import (
+    RegistryInfo,
+    is_registry_object,
+    registry_add,
+    registry_info,
+    registry_name,
+    registry_tag,
+    set_registry_info,
+)
+from inspect_ai.model._chat_message import (
+    ChatMessage,
+    ChatMessageAssistant,
+)
+from inspect_ai.model._model_output import ChatCompletionChoice, ModelOutput
+
+
+class AgentState:
+    """Agent state."""
+
+    def __init__(self, *, messages: list[ChatMessage]) -> None:
+        self._messages = messages
+        self._output: ModelOutput | None = None
+
+    @property
+    def messages(self) -> list[ChatMessage]:
+        """Conversation history."""
+        return self._messages
+
+    @messages.setter
+    def messages(self, messages: list[ChatMessage]) -> None:
+        """Set the conversation history."""
+        self._messages = messages
+
+    @property
+    def output(self) -> ModelOutput:
+        """Model output."""
+        # if there is no output yet then synthesize it from the last assistant message
+        if self._output is None:
+            # look for the last assistant message
+            for message in reversed(self.messages):
+                if isinstance(message, ChatMessageAssistant):
+                    self._output = ModelOutput(
+                        model=message.model or "",
+                        choices=[
+                            ChatCompletionChoice(
+                                message=message.model_copy(),
+                                stop_reason="stop",
+                            )
+                        ],
+                    )
+
+            # no assistant message, so generate an empty model output
+            if self._output is None:
+                self._output = ModelOutput()
+
+        return self._output
+
+    @output.setter
+    def output(self, output: ModelOutput) -> None:
+        """Set the model output."""
+        self._output = output
+
+    def __copy__(self) -> "AgentState":
+        state = AgentState(messages=copy(self.messages))
+        state.output = self.output.model_copy()
+        return state
+
+    def __deepcopy__(self, memo: dict[int, Any]) -> "AgentState":
+        state = AgentState(messages=deepcopy(self.messages, memo))
+        state.output = self.output.model_copy(deep=True)
+        return state
+
+
+@runtime_checkable
+class Agent(Protocol):
+    async def __call__(
+        self,
+        state: AgentState,
+        *args: Any,
+        **kwargs: Any,
+    ) -> AgentState:
+        """Agents perform tasks and participate in conversations.
+
+        Agents are similar to tools however they are participants
+        in conversation history and can optionally append messages
+        and model output to the current conversation state.
+
+        You can give the model a tool that enables handoff to
+        your agent using the `handoff()` function.
+
+        You can create a simple tool (that receives a string as
+        input) from an agent using `as_tool()`.
+
+        Args:
+            state: Agent state (conversation history and last model output)
+            *args: Arguments for the agent.
+            **kwargs: Keyword arguments for the agent.
+
+        Returns:
+            AgentState: Updated agent state.
+        """
+        ...
+
+
+P = ParamSpec("P")
+
+
+@overload
+def agent(func: Callable[P, Agent]) -> Callable[P, Agent]: ...
+
+
+@overload
+def agent() -> Callable[[Callable[P, Agent]], Callable[P, Agent]]: ...
+
+
+@overload
+def agent(
+    *,
+    name: str | None = None,
+    description: str | None = None,
+) -> Callable[[Callable[P, Agent]], Callable[P, Agent]]: ...
+
+
+def agent(
+    func: Callable[P, Agent] | None = None,
+    *,
+    name: str | None = None,
+    description: str | None = None,
+) -> Callable[P, Agent] | Callable[[Callable[P, Agent]], Callable[P, Agent]]:
+    r"""Decorator for registering agents.
+
+    Args:
+        func: Agent function
+        name: Optional name for agent. If the decorator has no name
+            argument then the name of the agent creation function
+            will be used as the name of the agent.
+        description: Description for the agent when used as
+            an ordinary tool or handoff tool.
+
+    Returns:
+        Agent with registry attributes.
+    """
+
+    def create_agent_wrapper(agent_type: Callable[P, Agent]) -> Callable[P, Agent]:
+        # determine the name (explicit or implicit from object)
+        agent_name = registry_name(
+            agent_type, name if name else getattr(agent_type, "__name__")
+        )
+
+        # wrap instantiations of agent so they carry registry info and metrics
+        @wraps(agent_type)
+        def agent_wrapper(*args: P.args, **kwargs: P.kwargs) -> Agent:
+            # create agent
+            agent = agent_type(*args, **kwargs)
+
+            # this might already have registry info, if so capture that
+            # and use it as default
+            if is_registry_object(agent):
+                info = registry_info(agent)
+                registry_name = info.name
+                registry_description = info.metadata.get(AGENT_DESCRIPTION, None)
+            else:
+                registry_name = None
+                registry_description = None
+
+            registry_tag(
+                agent_type,
+                agent,
+                RegistryInfo(
+                    type="agent",
+                    name=registry_name or agent_name,
+                    metadata={AGENT_DESCRIPTION: registry_description or description},
+                ),
+                *args,
+                **kwargs,
+            )
+            return agent
+
+        # register
+        return agent_register(cast(Callable[P, Agent], agent_wrapper), agent_name)
+
+    if func is not None:
+        return create_agent_wrapper(func)
+    else:
+        return create_agent_wrapper
+
+
+def agent_with(
+    agent: Agent,
+    *,
+    name: str | None = None,
+    description: str | None = None,
+) -> Agent:
+    """Agent with modifications to name and/or description
+
+    This function modifies the passed agent in place and
+    returns it. If you want to create multiple variations
+    of a single agent using `agent_with()` you should create
+    the underlying agent multiple times.
+
+    Args:
+       agent: Agent instance to modify.
+       name: Agent name (optional).
+       description: Agent description (optional).
+
+    Returns:
+       The passed agent with the requested modifications.
+    """
+    # resolve name and description
+    if is_registry_object(agent):
+        info = registry_info(agent)
+        name = name or info.name
+        description = description or info.metadata.get(AGENT_DESCRIPTION, None)
+
+    # if the name is null then raise
+    if name is None:
+        raise ValueError("You must provide a name to agent_with")
+
+    # now set registry info
+    set_registry_info(
+        agent,
+        RegistryInfo(
+            type="agent",
+            name=name,
+            metadata={AGENT_DESCRIPTION: description}
+            if description is not None
+            else {},
+        ),
+    )
+
+    return agent
+
+
+def agent_register(agent: Callable[P, Agent], name: str) -> Callable[P, Agent]:
+    r"""Register a function or class as an agent.
+
+    Args:
+        agent: Agent function or a class derived from Agent.
+        name (str): Name of agent (Optional, defaults to object name)
+
+    Returns:
+        Agent with registry attributes.
+    """
+    registry_add(
+        agent,
+        RegistryInfo(type="agent", name=name),
+    )
+    return agent
+
+
+def is_agent(obj: Any) -> TypeGuard[Agent]:
+    return is_registry_object(obj, type="agent")
+
+
+AGENT_DESCRIPTION = "description"

inspect_ai/agent/_as_solver.py

@@ -0,0 +1,72 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from inspect_ai.solver._solver import Solver
+
+from inspect_ai._util.registry import (
+    is_registry_object,
+    registry_unqualified_name,
+)
+from inspect_ai.tool._tool_info import parse_tool_info
+
+from ._agent import Agent, AgentState
+
+
+def as_solver(agent: Agent, **agent_kwargs: Any) -> Solver:
+    """Convert an agent to a solver.
+
+    Note that agents used as solvers will only receive their first parameter
+    (`state`). Any other parameters must provide appropriate defaults
+    or be explicitly specified in `agent_kwargs`
+
+    Args:
+       agent: Agent to convert.
+       **agent_kwargs: Arguments to curry to Agent function (required
+          if the agent has parameters without default values).
+
+    Solver:
+       Solver from agent.
+    """
+    from inspect_ai.solver._solver import Generate, solver
+    from inspect_ai.solver._task_state import TaskState
+
+    # agent must be registered (so we can get its name)
+    if not is_registry_object(agent):
+        raise RuntimeError(
+            "Agent passed to as_solver was not created by an @agent decorated function"
+        )
+    agent_name = registry_unqualified_name(agent)
+
+    # check to make sure we have all the parameters we need to run the agent
+    agent_info = parse_tool_info(agent)
+    for name, param in list(agent_info.parameters.properties.items())[1:]:
+        if param.default is None and name not in agent_kwargs:
+            raise ValueError(
+                f"To use the {agent_name} agent as a solver "
+                + f"you must pass a value for the agent's required '{name}' "
+                + "parameter to the as_solver() function."
+            )
+
+    @solver(name=agent_name)
+    def agent_to_solver() -> Solver:
+        async def solve(state: TaskState, generate: Generate) -> TaskState:
+            # run agent
+            agent_state = await agent(
+                AgentState(messages=state.messages), **agent_kwargs
+            )
+
+            # update messages
+            state.messages = agent_state.messages
+
+            # update output if its not empty
+            if agent_state.output:
+                state.output = agent_state.output
+
+            return state
+
+        # return solver
+        return solve
+
+    return agent_to_solver()

inspect_ai/agent/_as_tool.py

@@ -0,0 +1,122 @@
+from typing import Any
+
+from inspect_ai._util.registry import (
+    is_registry_object,
+    registry_info,
+    registry_unqualified_name,
+)
+from inspect_ai.model._chat_message import ChatMessageAssistant, ChatMessageUser
+from inspect_ai.tool._tool import Tool, ToolResult, tool
+from inspect_ai.tool._tool_def import ToolDef, validate_tool_parameters
+from inspect_ai.tool._tool_info import ToolInfo, parse_tool_info
+from inspect_ai.tool._tool_params import ToolParam
+
+from ._agent import AGENT_DESCRIPTION, Agent, AgentState
+
+
+@tool
+def as_tool(agent: Agent, description: str | None = None, **agent_kwargs: Any) -> Tool:
+    """Convert an agent to a tool.
+
+    By default the model will see all of the agent's arguments as
+    tool arguments (save for `state` which is converted to an `input`
+    arguments of type `str`). Provide optional `agent_kwargs` to mask
+    out agent parameters with default values (these parameters will
+    not be presented to the model as part of the tool interface)
+
+    Args:
+       agent: Agent to convert.
+       description: Tool description (defaults to agent description)
+       **agent_kwargs: Arguments to curry to Agent function (arguments
+          provided here will not be presented to the model as part
+          of the tool interface).
+
+    Returns:
+        Tool from agent.
+    """
+    # agent must be registered (so we can get its name)
+    if not is_registry_object(agent):
+        raise RuntimeError(
+            "Agent passed to as_tool was not created by an @agent decorated function"
+        )
+
+    async def execute(input: str, *args: Any, **kwargs: Any) -> ToolResult:
+        # prepare state and call agent
+        state = AgentState(messages=[ChatMessageUser(content=input)])
+        state = await agent(state, *args, **(agent_kwargs | kwargs))
+
+        # find assistant message to read content from (prefer output)
+        if not state.output.empty:
+            return state.output.message.content
+        elif len(state.messages) > 0 and isinstance(
+            state.messages[-1], ChatMessageAssistant
+        ):
+            return state.messages[-1].content
+        else:
+            return ""
+
+    # get tool_info
+    tool_info = agent_tool_info(agent, description, **agent_kwargs)
+
+    # add "input" param
+    tool_info.parameters.properties = {
+        "input": ToolParam(type="string", description="Input message.")
+    } | tool_info.parameters.properties
+    tool_info.parameters.required.append("input")
+
+    # create tool
+    tool_def = ToolDef(
+        execute,
+        name=tool_info.name,
+        description=tool_info.description,
+        parameters=tool_info.parameters,
+    )
+    return tool_def.as_tool()
+
+
+def agent_tool_info(
+    agent: Agent, description: str | None, **agent_kwargs: Any
+) -> ToolInfo:
+    # get tool_info and name
+    tool_info = parse_tool_info(agent)
+    tool_info.name = registry_unqualified_name(agent)
+
+    # remove "state" param
+    def remove_param(param: str) -> None:
+        if param in tool_info.parameters.properties:
+            del tool_info.parameters.properties[param]
+        if param in tool_info.parameters.required:
+            tool_info.parameters.required.remove(param)
+
+    remove_param("state")
+
+    # validate and remove curried params
+    for agent_param in agent_kwargs.keys():
+        if agent_param in tool_info.parameters.properties:
+            remove_param(agent_param)
+        else:
+            raise ValueError(
+                f"Agent {tool_info.name} does not have a '{agent_param}' parameter."
+            )
+
+    # resolve and validate description. the description in the call takes
+    # precedence, then any @agent(description="<foo>"), and finally any
+    # doc comment on the agent's execute function
+    reg_info = registry_info(agent)
+    tool_info.description = (
+        description
+        or reg_info.metadata.get(AGENT_DESCRIPTION, None)
+        or tool_info.description
+    )
+    if len(tool_info.description) == 0:
+        raise ValueError(
+            f"Description not provided for agent function '{tool_info.name}'. Provide a "
+            + "description either via @agent(description='<description>'), the description "
+            + "argument to as_tool() or handoff(), or via a doc comment on the agent's "
+            + "execute function."
+        )
+
+    # validate parameter descriptions and types
+    validate_tool_parameters(tool_info.name, tool_info.parameters.properties)
+
+    return tool_info

inspect_ai/{solver → agent}/_bridge/bridge.py

@@ -5,17 +5,15 @@ from pydantic import BaseModel, Field, ValidationError
 from pydantic_core import to_json
 
 from inspect_ai._util._async import is_callable_coroutine
-from inspect_ai.model._chat_message import ChatMessage, ChatMessageUser
+from inspect_ai.agent._agent import Agent, AgentState, agent
+from inspect_ai.model._model import get_model
+from inspect_ai.model._model_output import ModelOutput
 from inspect_ai.model._providers.providers import validate_openai_client
-from inspect_ai.scorer._metric import Score
 
-from .._solver import Generate, Solver, solver
-from .._task_state import TaskState
 
-
-@solver
-def bridge(agent: Callable[[dict[str, Any]], Awaitable[dict[str, Any]]]) -> Solver:
-    """Bridge an external agent into an Inspect Solver.
+@agent
+def bridge(agent: Callable[[dict[str, Any]], Awaitable[dict[str, Any]]]) -> Agent:
+    """Bridge an external agent into an Inspect Agent.
 
     See documentation at <https://inspect.aisi.org.uk/agent-bridge.html>
 
@@ -25,7 +23,7 @@ def bridge(agent: Callable[[dict[str, Any]], Awaitable[dict[str, Any]]]) -> Solv
     Returns:
       Standard Inspect solver.
     """
-    validate_openai_client("Solver bridge()")
+    validate_openai_client("Agent bridge()")
 
     from openai.types.chat import ChatCompletionMessageParam
 
@@ -36,17 +34,15 @@ def bridge(agent: Callable[[dict[str, Any]], Awaitable[dict[str, Any]]]) -> Solv
 
     from .patch import openai_request_to_inspect_model
 
-    class BridgeSample(BaseModel):
-        sample_id: str
-        epoch: int
+    class BridgeInput(BaseModel):
+        messages: list[ChatCompletionMessageParam]
+
+        # temporarily here for backward compatibility w/ previous bridge
         input: list[ChatCompletionMessageParam]
-        metadata: dict[str, Any]
-        target: list[str]
 
     class BridgeResult(BaseModel):
         output: str
         messages: list[ChatCompletionMessageParam] | None = Field(default=None)
-        scores: dict[str, Score] | None = Field(default=None)
 
     result_schema = BridgeResult.model_json_schema()
     result_validator = Draft7Validator(result_schema)
@@ -55,27 +51,15 @@ def bridge(agent: Callable[[dict[str, Any]], Awaitable[dict[str, Any]]]) -> Solv
     if not is_callable_coroutine(agent):
         raise TypeError(f"'{agent.__name__}' is not declared as an async callable.")
 
-    async def solve(state: TaskState, generate: Generate) -> TaskState:
-        # resolve input to array
-        input: list[ChatMessage] = (
-            [ChatMessageUser(content=state.input)]
-            if isinstance(state.input, str)
-            else state.input
-        )
-
-        # create sample (use standard gpt-4 message encoding -- i.e. no 'developer' messages)
-        sample = BridgeSample(
-            sample_id=str(state.sample_id),
-            epoch=state.epoch,
-            input=await openai_chat_messages(input, model="gpt-4"),
-            metadata=state.metadata,
-            target=list(state.target),
-        )
+    async def execute(state: AgentState) -> AgentState:
+        # create input (use standard gpt-4 message encoding -- i.e. no 'developer' messages)
+        messages = await openai_chat_messages(state.messages, model="gpt-4")
+        input = BridgeInput(messages=messages, input=messages)
 
         # run target function
         async with openai_request_to_inspect_model():
             # call the function
-            result_dict = await agent(sample.model_dump())
+            result_dict = await agent(input.model_dump())
             try:
                 result = BridgeResult.model_validate(result_dict)
             except ValidationError:
@@ -89,12 +73,14 @@ def bridge(agent: Callable[[dict[str, Any]], Awaitable[dict[str, Any]]]) -> Solv
                 raise ValueError(message)
 
         # update and return state
-        state.output.completion = result.output
+        state.output = ModelOutput.from_content(
+            model=get_model().name, content=result.output
+        )
         if result.messages is not None:
-            state.messages = chat_messages_from_openai(result.messages)
-        if result.scores is not None:
-            state.scores = result.scores
+            state.messages = chat_messages_from_openai(
+                state.output.model, result.messages
+            )
 
         return state
 
-    return solve
+    return execute

inspect_ai/{solver → agent}/_bridge/patch.py

@@ -23,7 +23,6 @@ from inspect_ai.model._openai import (
     openai_chat_choices,
     openai_completion_usage,
 )
-from inspect_ai.solver._task_state import sample_state
 from inspect_ai.tool._tool_choice import ToolChoice, ToolFunction
 from inspect_ai.tool._tool_info import ToolInfo
 from inspect_ai.tool._tool_params import ToolParams
@@ -98,10 +97,18 @@ def init_openai_request_patch() -> None:
 async def inspect_model_request(
     model_name: str, options: FinalRequestOptions
 ) -> ChatCompletion:
+    from inspect_ai.solver._task_state import sample_state
+
+    # resolve model
+    if model_name == "inspect":
+        model = get_model()
+    else:
+        model = get_model(model_name.removeprefix("inspect/"))
+
     # convert openai messages to inspect messages
     json_data = cast(dict[str, Any], options.json_data)
     messages: list[ChatCompletionMessageParam] = json_data["messages"]
-    input = chat_messages_from_openai(messages)
+    input = chat_messages_from_openai(model.api.model_name, messages)
 
     # convert openai tools to inspect tools
     tools: list[ChatCompletionToolParam] = json_data.get("tools", [])
@@ -130,12 +137,6 @@ async def inspect_model_request(
             case _:
                 inspect_tool_choice = ToolFunction(name=tool_choice["function"]["name"])
 
-    # resolve model
-    if model_name == "inspect":
-        model = get_model()
-    else:
-        model = get_model(model_name.removeprefix("inspect/"))
-
     output = await model.generate(
         input=input,
         tools=inspect_tools,

inspect_ai/agent/_filter.py

@@ -0,0 +1,46 @@
+from typing import Awaitable, Callable
+
+from inspect_ai.model._chat_message import (
+    ChatMessage,
+    ChatMessageAssistant,
+    ChatMessageTool,
+)
+
+MessageFilter = Callable[[list[ChatMessage]], Awaitable[list[ChatMessage]]]
+"""Filter messages sent to or received from agent handoffs."""
+
+
+async def remove_tools(messages: list[ChatMessage]) -> list[ChatMessage]:
+    """Remove tool calls from messages.
+
+    Removes all instances of `ChatMessageTool` as well as the `tool_calls`
+    field from `ChatMessageAssistant`.
+
+    Args:
+       messages: Messages to remove tool calls from.
+
+    Returns:
+       Messages without tool calls.
+    """
+    filtered: list[ChatMessage] = []
+    for message in messages:
+        if isinstance(message, ChatMessageTool):
+            continue
+        if isinstance(message, ChatMessageAssistant):
+            message = message.model_copy(update=dict(tool_calls=None))
+        filtered.append(message)
+
+    return filtered
+
+
+async def last_message(messages: list[ChatMessage]) -> list[ChatMessage]:
+    """Remove all but the last message.
+
+    Args:
+       messages: Target messages.
+
+    Returns:
+       List containing only the last message from the input list.
+
+    """
+    return messages[-1:]