inspect-ai 0.3.102__py3-none-any.whl → 0.3.104__py3-none-any.whl

inspect_ai/model/__init__.py

@@ -1,8 +1,16 @@
 # ruff: noqa: F401 F403 F405
 
+from inspect_ai._util.citation import (
+    Citation,
+    CitationBase,
+    ContentCitation,
+    DocumentCitation,
+    UrlCitation,
+)
 from inspect_ai._util.content import (
     Content,
     ContentAudio,
+    ContentData,
     ContentImage,
     ContentReasoning,
     ContentText,
@@ -59,6 +67,7 @@ __all__ = [
     "ResponseSchema",
     "CachePolicy",
     "ContentAudio",
+    "ContentData",
     "ContentImage",
     "ContentReasoning",
     "ContentText",
@@ -93,6 +102,11 @@ __all__ = [
     "cache_size",
     "get_model",
     "modelapi",
+    "Citation",
+    "CitationBase",
+    "DocumentCitation",
+    "ContentCitation",
+    "UrlCitation",
 ]
 
 _TOOL_MODULE_VERSION = "0.3.18"

inspect_ai/model/_call_tools.py

@@ -1,7 +1,7 @@
 import inspect
 import json
 import types
-from copy import copy
+from copy import copy, deepcopy
 from dataclasses import is_dataclass
 from datetime import date, datetime, time
 from enum import EnumMeta
@@ -36,6 +36,7 @@ from pydantic import BaseModel
 from inspect_ai._util.content import (
     Content,
     ContentAudio,
+    ContentData,
     ContentImage,
     ContentText,
     ContentVideo,
@@ -188,13 +189,19 @@ async def execute_tools(
             # types to string as that is what the model APIs accept
             truncated: tuple[int, int] | None = None
             if isinstance(
-                result, ContentText | ContentImage | ContentAudio | ContentVideo
+                result,
+                ContentText | ContentImage | ContentAudio | ContentVideo | ContentData,
             ):
                 content: str | list[Content] = [result]
             elif isinstance(result, list) and (
                 len(result) == 0
                 or isinstance(
-                    result[0], ContentText | ContentImage | ContentAudio | ContentVideo
+                    result[0],
+                    ContentText
+                    | ContentImage
+                    | ContentAudio
+                    | ContentVideo
+                    | ContentData,
                 )
             ):
                 content = result
@@ -471,7 +478,9 @@ async def agent_handoff(
     limit_error: LimitExceededError | None = None
     agent_state = AgentState(messages=copy(agent_conversation))
     try:
-        with apply_limits(agent_tool.limits):
+        # The agent_tool's limits will be applied multiple times if the agent is handed
+        # off to multiple times which is not supported, so create a copy of each limit.
+        with apply_limits(deepcopy(agent_tool.limits)):
             async with span(name=agent_name, type="agent"):
                 agent_state = await agent_tool.agent(agent_state, **arguments)
     except LimitExceededError as ex:

inspect_ai/model/_chat_message.py

@@ -26,6 +26,9 @@ class ChatMessageBase(BaseModel):
     source: Literal["input", "generate"] | None = Field(default=None)
     """Source of message."""
 
+    metadata: dict[str, Any] | None = Field(default=None)
+    """Additional message metadata."""
+
     internal: JsonValue | None = Field(default=None)
     """Model provider specific payload - typically used to aid transformation back to model types."""
 

inspect_ai/model/_openai_responses.py

@@ -31,9 +31,16 @@ from openai.types.responses.response_create_params import (
     ToolChoice as ResponsesToolChoice,
 )
 from openai.types.responses.response_input_item_param import FunctionCallOutput, Message
+from openai.types.responses.response_output_text import (
+    Annotation,
+    AnnotationFileCitation,
+    AnnotationFilePath,
+    AnnotationURLCitation,
+)
 from openai.types.responses.response_reasoning_item_param import Summary
 from pydantic import JsonValue
 
+from inspect_ai._util.citation import Citation, DocumentCitation, UrlCitation
 from inspect_ai._util.content import (
     Content,
     ContentImage,
@@ -47,29 +54,30 @@ from inspect_ai.model._chat_message import ChatMessage, ChatMessageAssistant
 from inspect_ai.model._generate_config import GenerateConfig
 from inspect_ai.model._model_output import ChatCompletionChoice, StopReason
 from inspect_ai.model._openai import is_o_series
-from inspect_ai.model._openai_computer_use import (
+from inspect_ai.tool._tool_call import ToolCall
+from inspect_ai.tool._tool_choice import ToolChoice
+from inspect_ai.tool._tool_info import ToolInfo
+
+from ._providers._openai_computer_use import (
     computer_call_output,
     maybe_computer_use_preview_tool,
     tool_call_from_openai_computer_tool_call,
 )
-from inspect_ai.model._openai_web_search import maybe_web_search_tool
-from inspect_ai.tool._tool_call import ToolCall
-from inspect_ai.tool._tool_choice import ToolChoice
-from inspect_ai.tool._tool_info import ToolInfo
+from ._providers._openai_web_search import maybe_web_search_tool
 
 
 async def openai_responses_inputs(
-    messages: list[ChatMessage], model: str, store: bool
+    messages: list[ChatMessage], model: str
 ) -> list[ResponseInputItemParam]:
     return [
         item
         for message in messages
-        for item in await _openai_input_item_from_chat_message(message, model, store)
+        for item in await _openai_input_item_from_chat_message(message, model)
     ]
 
 
 async def _openai_input_item_from_chat_message(
-    message: ChatMessage, model: str, store: bool
+    message: ChatMessage, model: str
 ) -> list[ResponseInputItemParam]:
     if message.role == "system":
         content = await _openai_responses_content_list_param(message.content)
@@ -87,7 +95,7 @@ async def _openai_input_item_from_chat_message(
             )
         ]
     elif message.role == "assistant":
-        return _openai_input_items_from_chat_message_assistant(message, store)
+        return _openai_input_items_from_chat_message_assistant(message)
     elif message.role == "tool":
         if message.internal:
             internal = _model_tool_call_for_internal(message.internal)
@@ -252,7 +260,18 @@ def _chat_message_assistant_from_openai_response(
             case ResponseOutputMessage(content=content, id=id):
                 message_content.extend(
                     [
-                        ContentText(text=c.text, internal={"id": id})
+                        ContentText(
+                            text=c.text,
+                            internal={"id": id},
+                            citations=(
+                                [
+                                    _to_inspect_citation(annotation)
+                                    for annotation in c.annotations
+                                ]
+                                if c.annotations
+                                else None
+                            ),
+                        )
                         if isinstance(c, ResponseOutputText)
                         else ContentText(
                             text=c.refusal, refusal=True, internal={"id": id}
@@ -310,7 +329,7 @@ def _chat_message_assistant_from_openai_response(
 
 
 def _openai_input_items_from_chat_message_assistant(
-    message: ChatMessageAssistant, store: bool
+    message: ChatMessageAssistant,
 ) -> list[ResponseInputItemParam]:
     """
     Transform a `ChatMessageAssistant` into OpenAI `ResponseInputItem`'s for playback to the model.
@@ -343,10 +362,6 @@ def _openai_input_items_from_chat_message_assistant(
     )
     suppress_output_message = message.internal is not None and not has_content_with_ids
 
-    # if we are not storing messages on the server then blank these out
-    if not store:
-        tool_message_ids = {}
-
     # items to return
     items: list[ResponseInputItemParam] = []
     # group content by message ID
@@ -354,30 +369,21 @@ def _openai_input_items_from_chat_message_assistant(
         str | None, list[ResponseOutputTextParam | ResponseOutputRefusalParam]
     ] = {}
 
-    for content in (
-        list[ContentText | ContentReasoning]([ContentText(text=message.content)])
-        if isinstance(message.content, str)
-        else [
-            c for c in message.content if isinstance(c, ContentText | ContentReasoning)
-        ]
-    ):
+    for content in _filter_consecutive_reasoning_blocks(content_items):
         match content:
             case ContentReasoning(reasoning=reasoning):
                 assert content.signature is not None, (
                     "reasoning_id must be saved in signature"
                 )
-                # if items are not stored on the server then there is no
-                # sense appending the reasoning item as its just a pointer
-                if store:
-                    items.append(
-                        ResponseReasoningItemParam(
-                            type="reasoning",
-                            id=content.signature,
-                            summary=[Summary(type="summary_text", text=reasoning)]
-                            if reasoning
-                            else [],
-                        )
+                items.append(
+                    ResponseReasoningItemParam(
+                        type="reasoning",
+                        id=content.signature,
+                        summary=[Summary(type="summary_text", text=reasoning)]
+                        if reasoning
+                        else [],
                     )
+                )
             case ContentText(text=text, refusal=refusal):
                 if suppress_output_message:
                     continue
@@ -409,7 +415,7 @@ def _openai_input_items_from_chat_message_assistant(
             role="assistant",
             # this actually can be `None`, and it will in fact be `None` when the
             # assistant message is synthesized by the scaffold as opposed to being
-            # replayed from the model (or when store=False)
+            # replayed from the model
             id=msg_id,  # type: ignore[typeddict-item]
             content=content_list,
             status="completed",
@@ -531,3 +537,43 @@ def _responses_tool_alias(name: str) -> str:
 
 def _from_responses_tool_alias(name: str) -> str:
     return next((k for k, v in _responses_tool_aliases.items() if v == name), name)
+
+
+def _to_inspect_citation(input: Annotation) -> Citation:
+    match input:
+        case AnnotationURLCitation(
+            end_index=end_index, start_index=start_index, title=title, url=url
+        ):
+            return UrlCitation(
+                cited_text=(start_index, end_index), title=title, url=url
+            )
+
+        case (
+            AnnotationFileCitation(file_id=file_id, index=index)
+            | AnnotationFilePath(file_id=file_id, index=index)
+        ):
+            return DocumentCitation(internal={"file_id": file_id, "index": index})
+    assert False, f"Unexpected citation type: {input.type}"
+
+
+def _filter_consecutive_reasoning_blocks(
+    content_list: list[ContentText | ContentReasoning],
+) -> list[ContentText | ContentReasoning]:
+    return [
+        content
+        for i, content in enumerate(content_list)
+        if _should_keep_content(i, content, content_list)
+    ]
+
+
+def _should_keep_content(
+    i: int,
+    content: ContentText | ContentReasoning,
+    content_list: list[ContentText | ContentReasoning],
+) -> bool:
+    return (
+        True
+        if not isinstance(content, ContentReasoning)
+        else i == len(content_list) - 1
+        or not isinstance(content_list[i + 1], ContentReasoning)
+    )

inspect_ai/model/_providers/_anthropic_citations.py

@@ -0,0 +1,158 @@
+from anthropic.types import (
+    CitationCharLocation,
+    CitationCharLocationParam,
+    CitationContentBlockLocation,
+    CitationContentBlockLocationParam,
+    CitationPageLocation,
+    CitationPageLocationParam,
+    CitationsWebSearchResultLocation,
+    CitationWebSearchResultLocationParam,
+    TextCitation,
+    TextCitationParam,
+)
+
+from inspect_ai._util.citation import (
+    Citation,
+    DocumentCitation,
+    DocumentRange,
+    UrlCitation,
+)
+
+
+def to_inspect_citation(input: TextCitation) -> Citation:
+    match input:
+        case CitationsWebSearchResultLocation(
+            cited_text=cited_text,
+            title=title,
+            url=url,
+            encrypted_index=encrypted_index,
+        ):
+            # Sanitize a citation to work around https://github.com/anthropics/anthropic-sdk-python/issues/965.
+            return UrlCitation(
+                cited_text=cited_text,
+                title=title
+                if title is None or len(title) <= 255
+                else title[:254] + "…",
+                url=url,
+                internal={"encrypted_index": encrypted_index},
+            )
+
+        case CitationCharLocation(
+            cited_text=cited_text,
+            document_index=document_index,
+            document_title=title,
+            end_char_index=end_char_index,
+            start_char_index=start_char_index,
+        ):
+            return DocumentCitation(
+                cited_text=cited_text,
+                title=title,
+                range=DocumentRange(
+                    type="char", start_index=start_char_index, end_index=end_char_index
+                ),
+                internal={"document_index": document_index},
+            )
+
+        case CitationContentBlockLocation(
+            cited_text=cited_text,
+            document_index=document_index,
+            document_title=title,
+            end_block_index=end_block_index,
+            start_block_index=start_block_index,
+        ):
+            return DocumentCitation(
+                cited_text=cited_text,
+                title=title,
+                range=DocumentRange(
+                    type="block",
+                    start_index=start_block_index,
+                    end_index=end_block_index,
+                ),
+                internal={"document_index": document_index},
+            )
+
+        case CitationPageLocation(
+            cited_text=cited_text,
+            document_index=document_index,
+            document_title=title,
+            end_page_number=end_page_number,
+            start_page_number=start_page_number,
+        ):
+            return DocumentCitation(
+                cited_text=cited_text,
+                title=title,
+                range=DocumentRange(
+                    type="page",
+                    start_index=start_page_number - 1,
+                    end_index=end_page_number - 1,
+                ),
+                internal={"document_index": document_index},
+            )
+
+    assert False, f"Unexpected citation type: {input.type}"
+
+
+def to_anthropic_citation(input: Citation) -> TextCitationParam:
+    cited_text = input.cited_text
+    assert isinstance(cited_text, str), (
+        "anthropic citations must have a string cited_text"
+    )
+
+    match input:
+        case UrlCitation(title=title, url=url, internal=internal):
+            assert internal, "UrlCitation must have internal field"
+            encrypted_index = internal.get("encrypted_index", None)
+            assert isinstance(encrypted_index, str), (
+                "URL citations require encrypted_index in internal field"
+            )
+
+            return CitationWebSearchResultLocationParam(
+                type="web_search_result_location",
+                cited_text=cited_text,
+                title=title,
+                url=url,
+                encrypted_index=encrypted_index,
+            )
+
+        case DocumentCitation(title=title, range=range, internal=internal):
+            assert internal, "DocumentCharCitation must have internal field"
+            document_index = internal.get("document_index", None)
+            assert isinstance(document_index, int), (
+                "DocumentCharCitation require encrypted_index in internal field"
+            )
+            assert range, "DocumentCitation must have a range"
+
+            start_index = range.start_index
+            end_index = range.end_index
+
+            match range.type:
+                case "char":
+                    return CitationCharLocationParam(
+                        type="char_location",
+                        cited_text=cited_text,
+                        document_title=title,
+                        document_index=document_index,
+                        start_char_index=start_index,
+                        end_char_index=end_index,
+                    )
+                case "block":
+                    return CitationContentBlockLocationParam(
+                        type="content_block_location",
+                        cited_text=cited_text,
+                        document_title=title,
+                        document_index=document_index,
+                        start_block_index=start_index,
+                        end_block_index=end_index,
+                    )
+                case "page":
+                    return CitationPageLocationParam(
+                        type="page_location",
+                        cited_text=cited_text,
+                        document_title=title,
+                        document_index=document_index,
+                        start_page_number=start_index + 1,
+                        end_page_number=end_index + 1,
+                    )
+
+    # If we can't handle this citation type, raise an error
+    raise ValueError(f"Unsupported citation type: {input.type}")

inspect_ai/model/_providers/_google_citations.py

@@ -0,0 +1,100 @@
+from typing import Sequence
+
+from google.genai.types import (
+    Candidate,
+    GroundingChunk,
+    GroundingSupport,
+    Segment,
+)
+
+from inspect_ai._util.citation import Citation, UrlCitation
+
+
+def get_candidate_citations(candidate: Candidate) -> list[Citation]:
+    """Extract citations from Google AI candidate grounding metadata.
+
+    Understanding Google API Grounding Citations: `GroundingChunk`'s, `GroundingSupport`, and `Segment`'s
+
+    1. Grounding Chunks (`GroundingChunk`)
+       What they are: The raw source material that the AI retrieved to support its response.
+       Structure:
+         - Web (`GroundingChunkWeb`): Content from web searches
+           - domain: The website domain
+           - title: Page title
+           - uri: Web page URL
+       Think of chunks as: The library books or web pages that contain the information.
+
+    2. Segments (`Segment`)
+       What they are: Specific portions of the AI's generated response text.
+       Structure:
+         - start_index & end_index: Byte positions in the response text
+         - text: The actual text from the response that this segment represents
+       Think of segments as: Specific sentences or paragraphs in the AI's response that need citations.
+
+    3. Grounding Support (`GroundingSupport`)
+       What they are: The bridge that connects segments of the AI's response to the chunks that support them.
+       Structure:
+         - grounding_chunk_indices: Array of integers pointing to specific chunks (e.g., [1,3,4] means chunks 1, 3, and 4 support this claim)
+         - segment: Which part of the response this support applies to
+       Think of support as: The footnotes that say "this claim in my response is backed up by these specific sources."
+
+    Args:
+        candidate: The Google AI candidate response containing grounding metadata
+
+    Returns:
+        A list of `Citation` objects linking response segments to their web sources.
+        Currently only handles `GroundingChunkWeb` sources.
+    """
+    return (
+        []
+        if (
+            not candidate.content
+            or not candidate.content.parts
+            or not (metadata := candidate.grounding_metadata)
+            or not (chunks := metadata.grounding_chunks)
+            or not (supports := metadata.grounding_supports)
+        )
+        else [
+            citation
+            for support in supports
+            for citation in _citations_from_support(support, chunks)
+        ]
+    )
+
+
+def _create_citation_from_chunk_and_segment(
+    chunk: GroundingChunk, segment: Segment
+) -> UrlCitation | None:
+    """Create a citation from a chunk and segment, returning None if chunk is not web-based."""
+    return (
+        UrlCitation(
+            url=chunk.web.uri,
+            title=chunk.web.title,
+            cited_text=(
+                (segment.start_index or 0, segment.end_index)
+                if segment.end_index is not None
+                else None
+            ),
+        )
+        if (chunk.web and chunk.web.uri)
+        else None
+    )
+
+
+def _citations_from_support(
+    support: GroundingSupport, chunks: Sequence[GroundingChunk]
+) -> list[Citation]:
+    return (
+        []
+        if support.segment is None or support.grounding_chunk_indices is None
+        else [
+            citation
+            for chunk_index in support.grounding_chunk_indices
+            if chunk_index < len(chunks)
+            if (
+                citation := _create_citation_from_chunk_and_segment(
+                    chunks[chunk_index], support.segment
+                )
+            )
+        ]
+    )