import abc import base64 import datetime import json from collections import Counter, defaultdict from typing import Any, Generic, Optional, Sequence, TypeVar, Union, overload import grpc from pydantic import BaseModel from typing_extensions import Self from .cost import cost_usd_from_usage from .meta import ProtoDecorator from .proto import chat_pb2, chat_pb2_grpc, image_pb2, sample_pb2, usage_pb2 from .search import SearchParameters from .service_tier import service_tier_from_proto, service_tier_to_proto from .telemetry import should_disable_sensitive_attributes from .types import ( AgentCount, AgentCountMap, ChatModel, Content, ImageDetail, IncludeOption, IncludeOptionMap, ReasoningEffort, ResponseFormat, ServiceTier, ToolMode, ) T = TypeVar("T") class BaseClient(abc.ABC, Generic[T]): """Base Client for interacting with the `Chat` API.""" _stub: chat_pb2_grpc.ChatStub def __init__(self, channel: Union[grpc.Channel, grpc.aio.Channel]): """Creates a new client based on a gRPC channel.""" self._stub = chat_pb2_grpc.ChatStub(channel) def create( self, model: Union[ChatModel, str], *, conversation_id: Optional[str] = None, messages: Optional[Sequence[chat_pb2.Message]] = None, user: Optional[str] = None, max_tokens: Optional[int] = None, seed: Optional[int] = None, stop: Optional[Sequence[str]] = None, temperature: Optional[float] = None, top_p: Optional[float] = None, logprobs: Optional[bool] = None, top_logprobs: Optional[int] = None, tools: Optional[Sequence[chat_pb2.Tool]] = None, tool_choice: Optional[Union[ToolMode, chat_pb2.ToolChoice]] = None, parallel_tool_calls: Optional[bool] = None, response_format: Optional[Union[ResponseFormat, chat_pb2.ResponseFormat, type[BaseModel]]] = None, frequency_penalty: Optional[float] = None, presence_penalty: Optional[float] = None, reasoning_effort: Optional[Union[ReasoningEffort, "chat_pb2.ReasoningEffort"]] = None, search_parameters: Optional[Union[SearchParameters, chat_pb2.SearchParameters]] = None, store_messages: Optional[bool] = None, previous_response_id: Optional[str] = None, use_encrypted_content: Optional[bool] = None, max_turns: Optional[int] = None, include: Optional[Sequence[Union[IncludeOption, "chat_pb2.IncludeOption"]]] = None, agent_count: Optional[Union[AgentCount, "chat_pb2.AgentCount"]] = None, batch_request_id: Optional[str] = None, service_tier: Optional[Union[ServiceTier, "usage_pb2.ServiceTier"]] = None, ) -> T: """Creates a new chat conversation. This function does not immediately perform an RPC. It only initializes a mutable request instance. Examples: ``` chat = client.chat.create( model="grok-4.20-non-reasoning", messages=[ system("You are a pirate"), user("How are you?"), ] ) response = chat.sample() chat.append(response) print(response) chat.append(user("Tell me a joke")) response = chat.sample() print(response) ``` Args: model: Model to use, e.g. "grok-4.20-non-reasoning". conversation_id: Optional ID to group all messages in this chat instance. When provided, this ID is added as a span attribute ("gen_ai.conversation.id") to all OpenTelemetry spans generated by this chat instance (e.g., calls to `sample`). This enables easy identification and grouping of spans from the same conversation for monitoring and debugging. messages: A list of messages that make up the the chat conversation. Different models support different message types, such as image and text. user: A unique identifier representing your end-user, which can help xAI to monitor and detect abuse. max_tokens: The maximum number of tokens that can be generated in the chat completion. seed: If specified, our system will make a best effort to sample deterministically, such that repeated requests with the same seed and parameters should return the same result. stop: Up to 4 sequences where the API will stop generating further tokens. temperature: What sampling temperature to use, between 0 and 2. Higher values like 0.8 will make the output more random, while lower values like 0.2 will make it more focused and deterministic. top_p: An alternative to sampling with temperature, called nucleus sampling, where the model considers the results of the tokens with `top_p` probability mass. logprobs: Whether to return log probabilities of the output tokens or not. If true, returns the log probabilities of each output token returned in the content of message. top_logprobs: An integer between 0 and 8 specifying the number of most likely tokens to return at each token position, each with an associated log probability. logprobs must be set to true if this parameter is used. tools: A list of tools the model may call in JSON-schema. Currently, only functions are supported as a tool. Use this to provide a list of functions the model may generate JSON inputs for. A max of 128 functions are supported. tool_choice: Controls which (if any) tool is called by the model. `none` means the model will not call any tool and instead generates a message. `auto` means the model can pick between generating a message or calling one or more tools. `required` means the model must call one or more tools. `auto` is the default if tools are present. To force the model to always invoke a specific tool, use the `required_tool(name)` function to create a `ToolChoice` with mode `required` and specify the tool's name. For example: ``` tool_choice=required_tool("get_weather") ``` This ensures the model must call the specified tool from the list provided in `tools`. parallel_tool_calls: If set to false, the model can perform maximum one tool call per response. Defaults to true. response_format: An object specifying the format that the model must output. `json_object` means the model must output a JSON object (although the shape is arbitrary). `text` means the model must output a text string. To force the model to output a JSON object adhering to a specific JSON Schema use the client's `parse` method instead. frequency_penalty: Positive values penalize new tokens based on their existing frequency in the text so far, decreasing the model's likelihood to repeat the same line verbatim. presence_penalty: Positive values penalize new tokens based on whether they appear in the text so far, increasing the model's likelihood to talk about new topics. reasoning_effort: Constrains how hard a reasoning model thinks before responding. Possible values are `none` (no reasoning), `low` (uses fewer reasoning tokens), `medium`, `high` (uses more reasoning tokens), and `xhigh` (maximum reasoning depth; only supported by some models, e.g. `grok-4.6`). The default varies by model (e.g. `grok-4.5` and `grok-4.6` default to `high`). search_parameters: The parameters that control search behavior. This includes settings like search mode, date range, sources (e.g., web, news, or X), and whether to return citations. See `SearchParameters` for detailed configuration options. store_messages: Whether to store responses generated by this chat instance on the xAI backend. When set to True, each response is persisted with a unique response ID (accessible via `response.id`), enabling retrieval via `client.chat.get_stored_completion()` and deletion via `client.chat.delete_stored_completion()`. Stored responses can be referenced by `previous_response_id` in new chat instances to branch conversations from specific points. Defaults to False. Note: Teams with ZDR enabled will not be able to use this feature. previous_response_id: The ID of a previously stored response (i.e., `response.id`) to use as the starting point for this conversation. When provided, the entire conversation history up to that response is prepended to the current chat instance, enabling conversation branching and continuation from any stored point. This prepending happens automatically on the server side and is opaque to the user. Note: Adding `previous_response_id` to a chat instance will not update the local message list of the chat instance (i.e., `chat.messages`). use_encrypted_content: Whether to return encrypted reasoning content from the model response. When enabled, encrypted reasoning content is included in responses, enabling optimal hydration of reasoning traces in follow-up conversations. Calling `append(response)` automatically passes this encrypted content back for subsequent requests. This is particularly useful for users with zero data retention (ZDR) enabled who cannot use `store_messages` and `previous_response_id` for conversation continuity. Defaults to False. max_turns: The maximum number of agentic turns the model can take. When set, the model will automatically iterate up to this many turns, calling tools and processing their results until it reaches a final answer or hits the turn limit. Defaults to server-side maximum. Note: This parameter has no effect on non-agentic requests (i.e. requests that do not use server-side tools). With parallel tool calls enabled, multiple tool calls can occur within a single turn, so max_turns does not necessarily equal the total number of tool calls. include: A list of output options to include in the response. Check the `IncludeOption` enum for all possible values. Defaults to None. agent_count: The number of agents to use for multi-agent models. Possible values are `4` and `16`. Defaults to None (server-side default). batch_request_id: An optional user-provided identifier for the batch request. **If provided, it must be unique within the batch.**Used to identify the corresponding result when the response is returned to the user. service_tier: The processing tier for this request. Set to `"priority"` for higher scheduling priority at a higher token price. Valid values: `"default"`, `"priority"`. The response includes a `service_tier` field indicating which tier was actually used. Returns: A new chat request bound to a client. """ tool_choice_pb: Optional[chat_pb2.ToolChoice] = None if isinstance(tool_choice, str): tool_choice_pb = chat_pb2.ToolChoice(mode=_tool_mode_to_proto(tool_choice)) else: tool_choice_pb = tool_choice response_format_pb: Optional[chat_pb2.ResponseFormat] = None if isinstance(response_format, str): response_format_pb = chat_pb2.ResponseFormat(format_type=_format_type_to_proto(response_format)) elif isinstance(response_format, type) and issubclass(response_format, BaseModel): response_format_pb = chat_pb2.ResponseFormat( format_type=chat_pb2.FORMAT_TYPE_JSON_SCHEMA, schema=json.dumps(response_format.model_json_schema()), ) else: response_format_pb = response_format reasoning_effort_pb: Optional[chat_pb2.ReasoningEffort] = None if isinstance(reasoning_effort, str): reasoning_effort_pb = _reasoning_effort_to_proto(reasoning_effort) else: reasoning_effort_pb = reasoning_effort search_parameters_pb: Optional[chat_pb2.SearchParameters] = None if isinstance(search_parameters, SearchParameters): search_parameters_pb = search_parameters._to_proto() else: search_parameters_pb = search_parameters include_pb: Optional[Sequence[chat_pb2.IncludeOption]] = None if include is not None: include_pb = [ _include_option_to_proto(include_option) if isinstance(include_option, str) else include_option for include_option in include ] agent_count_pb: Optional[chat_pb2.AgentCount] = None if isinstance(agent_count, int): agent_count_pb = _agent_count_to_proto(agent_count) else: agent_count_pb = agent_count service_tier_pb = service_tier_to_proto(service_tier) if isinstance(service_tier, str) else service_tier return self._make_chat( conversation_id=conversation_id, batch_request_id=batch_request_id, model=model, messages=messages, user=user, max_tokens=max_tokens, seed=seed, stop=stop, temperature=temperature, top_p=top_p, logprobs=logprobs, top_logprobs=top_logprobs, tools=tools, tool_choice=tool_choice_pb, parallel_tool_calls=parallel_tool_calls, response_format=response_format_pb, frequency_penalty=frequency_penalty, presence_penalty=presence_penalty, reasoning_effort=reasoning_effort_pb, search_parameters=search_parameters_pb, store_messages=store_messages, previous_response_id=previous_response_id, use_encrypted_content=use_encrypted_content, max_turns=max_turns, include=include_pb, agent_count=agent_count_pb, service_tier=service_tier_pb, ) @abc.abstractmethod def _make_chat(self, conversation_id: Optional[str], batch_request_id: Optional[str], **settings) -> T: """Creates the proto wrapper for chat requests.""" def _replay_tool_call_id(output_message: chat_pb2.CompletionMessage) -> Optional[str]: """Recovers the `tool_call_id` for replaying a completion output as an input message. `chat_pb2.CompletionMessage` does not carry a `tool_call_id` field, but server-side tool outputs (ROLE_TOOL) echo the originating call in `tool_calls`. The server pairs replayed tool turns with their calls via `Message.tool_call_id` (e.g. to re-hydrate binary tool results such as generated images), so it must be set when appending a response back onto the conversation. """ if output_message.role == chat_pb2.MessageRole.ROLE_TOOL and output_message.tool_calls: return output_message.tool_calls[0].id return None class BaseChat(ProtoDecorator[chat_pb2.GetCompletionsRequest]): """Utility class for simplifying the interaction with Chat requests and responses.""" _stub: chat_pb2_grpc.ChatStub def __init__( self, stub: chat_pb2_grpc.ChatStub, conversation_id: Optional[str], batch_request_id: Optional[str], **settings, ) -> None: """Prepares a new chat request. Args: stub: gRPC stub used to connect to the server. conversation_id: The ID of the conversation. batch_request_id: The ID of the batch request, should only be set when creating a chat completion for a batch request. **settings: See `chat_pb2.GetCompletionsRequest`. """ super().__init__(chat_pb2.GetCompletionsRequest(**settings)) self._stub = stub self._conversation_id = conversation_id self._batch_request_id = batch_request_id def append(self, message: Union[chat_pb2.Message, "Response", "CompactContextResponse"]) -> Self: """Adds a new message to the conversation history, enabling multi-turn interactions. This method appends a message to the chat's message sequence, which can be a user input, system prompt, assistant response, or tool result. It supports both `chat_pb2.Message` objects (created via helper functions like `user`, `system`, `assistant`, or `tool_result`) and `Response` objects from previous chat interactions. The method returns the chat object itself, allowing for method chaining. Examples: ``` # Adding a simple user message chat = client.chat.create(model="grok-4.20-non-reasoning") chat.append(user("How are you?")) ``` ``` # Adding a system prompt with an image chat = client.chat.create(model="grok-4.20-non-reasoning") chat.append(system( "Analyze the following image: ", image("https://example.com/image.jpg", detail="high"), "Provide a detailed description." )) ``` ``` # Appending an assistant's response in a multi-turn conversation chat = client.chat.create(model="grok-4.20-non-reasoning") chat.append(user("Tell me a pirate joke.")) response = chat.sample() # Get assistant's response print(f"Grok: {response.content}") chat.append(response) # Add assistant's response to history chat.append(user("Another one, please!")) ``` ``` # Multi-turn chat loop (inspired by sync/chat.py) chat = client.chat.create( model="grok-4.20-non-reasoning", messages=[system("You talk like a pirate.")] ) while True: prompt = input("You: ") if prompt.lower() == "exit": break chat.append(user(prompt)) response = chat.sample() print(f"Grok: {response.content}") chat.append(response) ``` Args: message: The message to append. Can be a ``chat_pb2.Message`` (e.g., created by ``user``, ``system``, ``assistant``, or ``tool_result``), a ``Response`` from a previous chat interaction, or a ``CompactContextResponse`` from ``client.chat.compact_context()``. **Note:** appending a ``CompactContextResponse`` drops all previous messages and replaces them with the compacted context. Returns: Self: The chat object, enabling method chaining. """ if isinstance(message, CompactContextResponse): del self._proto.messages[:] self._proto.messages.append( chat_pb2.Message( role=chat_pb2.MessageRole.ROLE_USER, encrypted_content=message.encrypted_content, ) ) elif isinstance(message, chat_pb2.Message): self._proto.messages.append(message) elif isinstance(message, Response): if message._index is None: # Every single output should be appended for agentic tool call responses. for output in message.proto.outputs: self._proto.messages.append( chat_pb2.Message( role=output.message.role, content=[text(output.message.content)], reasoning_content=output.message.reasoning_content, encrypted_content=output.message.encrypted_content, tool_calls=output.message.tool_calls, tool_call_id=_replay_tool_call_id(output.message), ) ) else: output_message = message._get_output().message self._proto.messages.append( chat_pb2.Message( role=output_message.role, content=[text(message.content)], reasoning_content=message.reasoning_content, encrypted_content=message.encrypted_content, tool_calls=message.tool_calls, tool_call_id=_replay_tool_call_id(output_message), ) ) else: raise ValueError("Unrecognized message type.") return self def _make_request(self, n: int) -> chat_pb2.GetCompletionsRequest: """Creates a request proto. Args: n: Number of completions to generate. """ request = chat_pb2.GetCompletionsRequest() # prevent requests with no messages. if not self._proto.messages: raise ValueError( "Cannot create a completion request: No messages provided. Please include at least one " "message (e.g., using user(), system(), or assistant()) in the request." ) request.CopyFrom(self._proto) request.n = n return request def _make_span_request_attributes(self) -> dict[str, Any]: # noqa: C901, PLR0912 """Creates a dictionary with all relevant request attributes to be set on the span as it is created.""" attributes: dict[str, Any] = { "gen_ai.operation.name": "chat", "gen_ai.provider.name": "xai", "gen_ai.output.type": "text", "gen_ai.request.model": self._proto.model, "server.port": 443, "server.address": "api.x.ai", } if should_disable_sensitive_attributes(): return attributes # Initialize optional fields to their default values (as set server-side). # Override with user-set values only if they are provided by the user. attributes["gen_ai.request.frequency_penalty"] = 0.0 attributes["gen_ai.request.presence_penalty"] = 0.0 attributes["gen_ai.request.temperature"] = 1.0 attributes["gen_ai.request.parallel_tool_calls"] = True attributes["gen_ai.request.store_messages"] = False attributes["gen_ai.request.use_encrypted_content"] = False attributes["gen_ai.request.logprobs"] = self._proto.logprobs # Float fields that need rounding float_fields = [ ("frequency_penalty", "gen_ai.request.frequency_penalty"), ("presence_penalty", "gen_ai.request.presence_penalty"), ("temperature", "gen_ai.request.temperature"), ("top_p", "gen_ai.request.top_p"), ] # Integer fields int_fields = [ ("n", "gen_ai.request.choice.count"), ("seed", "gen_ai.request.seed"), ("max_tokens", "gen_ai.request.max_tokens"), ("top_logprobs", "gen_ai.request.top_logprobs"), ] # Set float fields with rounding for proto_field, attr_name in float_fields: if self._proto.HasField(proto_field): attributes[attr_name] = round(getattr(self._proto, proto_field), 6) # Set integer fields for proto_field, attr_name in int_fields: if self._proto.HasField(proto_field): attributes[attr_name] = getattr(self._proto, proto_field) # Special cases if self._conversation_id: attributes["gen_ai.conversation.id"] = self._conversation_id if len(self._proto.stop) > 0: attributes["gen_ai.request.stop_sequences"] = list(self._proto.stop) if self._proto.HasField("response_format"): attributes["gen_ai.output.type"] = ( chat_pb2.FormatType.Name(self._proto.response_format.format_type).removeprefix("FORMAT_TYPE_").lower() ) if self._proto.HasField("reasoning_effort"): attributes["gen_ai.request.reasoning_effort"] = ( chat_pb2.ReasoningEffort.Name(self._proto.reasoning_effort).removeprefix("EFFORT_").lower() ) if self._proto.user: attributes["user_id"] = self._proto.user if self._proto.HasField("parallel_tool_calls"): attributes["gen_ai.request.parallel_tool_calls"] = self._proto.parallel_tool_calls if self._proto.store_messages: attributes["gen_ai.request.store_messages"] = self._proto.store_messages if self._proto.previous_response_id: attributes["gen_ai.request.previous_response_id"] = self._proto.previous_response_id if self._proto.use_encrypted_content: attributes["gen_ai.request.use_encrypted_content"] = self._proto.use_encrypted_content prompt_attributes = self._get_span_prompt_attributes() attributes.update(prompt_attributes) return attributes def _get_span_prompt_attributes(self) -> dict[str, Any]: """Creates a dictionary with prompt message attributes for span telemetry.""" prompt_attributes: dict[str, Any] = {} # Skip collecting sensitive attributes if disabled if should_disable_sensitive_attributes(): return prompt_attributes # Only text content is included in span attributes. for index, message in enumerate(self._proto.messages): if message.role == chat_pb2.MessageRole.ROLE_USER: prompt_attributes[f"gen_ai.prompt.{index}.role"] = "user" prompt_attributes[f"gen_ai.prompt.{index}.content"] = "".join([c.text for c in message.content]) elif message.role == chat_pb2.MessageRole.ROLE_ASSISTANT: prompt_attributes[f"gen_ai.prompt.{index}.role"] = "assistant" prompt_attributes[f"gen_ai.prompt.{index}.content"] = "".join([c.text for c in message.content]) if message.tool_calls: prompt_attributes[f"gen_ai.prompt.{index}.tool_calls"] = json.dumps( [ { "id": tool_call.id, "type": "function", "function": { "name": tool_call.function.name, "arguments": json.loads(tool_call.function.arguments), }, } for tool_call in message.tool_calls ] ) elif message.role == chat_pb2.MessageRole.ROLE_SYSTEM: prompt_attributes[f"gen_ai.prompt.{index}.role"] = "system" prompt_attributes[f"gen_ai.prompt.{index}.content"] = "".join([c.text for c in message.content]) elif message.role == chat_pb2.MessageRole.ROLE_DEVELOPER: prompt_attributes[f"gen_ai.prompt.{index}.role"] = "developer" prompt_attributes[f"gen_ai.prompt.{index}.content"] = "".join([c.text for c in message.content]) elif message.role == chat_pb2.MessageRole.ROLE_TOOL: prompt_attributes[f"gen_ai.prompt.{index}.role"] = "tool" prompt_attributes[f"gen_ai.prompt.{index}.content"] = "".join([c.text for c in message.content]) return prompt_attributes def _make_span_response_attributes(self, responses: Sequence["Response"]) -> dict[str, Any]: """Creates a dictionary with response metadata and completion attributes for span telemetry.""" attributes: dict[str, Any] = {} if should_disable_sensitive_attributes(): return attributes # All of these attributes are the same for all responses, so we can just use the first response to access them. response = responses[0] attributes["gen_ai.response.id"] = response.id attributes["gen_ai.response.model"] = response._proto.model attributes["gen_ai.usage.input_tokens"] = response.usage.prompt_tokens attributes["gen_ai.usage.output_tokens"] = response.usage.completion_tokens attributes["gen_ai.usage.total_tokens"] = response.usage.total_tokens attributes["gen_ai.usage.reasoning_tokens"] = response.usage.reasoning_tokens attributes["gen_ai.usage.cached_prompt_text_tokens"] = response.usage.cached_prompt_text_tokens attributes["gen_ai.usage.prompt_text_tokens"] = response.usage.prompt_text_tokens attributes["gen_ai.usage.prompt_image_tokens"] = response.usage.prompt_image_tokens if response.usage.HasField("cost_in_usd_ticks"): attributes["gen_ai.usage.cost_in_usd_ticks"] = response.usage.cost_in_usd_ticks attributes["gen_ai.response.system_fingerprint"] = response.system_fingerprint # Only finish reasons are different for each response. attributes["gen_ai.response.finish_reasons"] = [response.finish_reason for response in responses] completion_attributes = self._get_span_completion_attributes(responses) attributes.update(completion_attributes) return attributes def _get_span_completion_attributes(self, responses: Sequence["Response"]) -> dict[str, Any]: """Creates a dictionary with completion content attributes for span telemetry.""" completion_attributes: dict[str, Any] = {} # Skip collecting sensitive attributes if disabled if should_disable_sensitive_attributes(): return completion_attributes for index, response in enumerate(responses): completion_attributes[f"gen_ai.completion.{index}.role"] = response.role.removeprefix("ROLE_").lower() completion_attributes[f"gen_ai.completion.{index}.content"] = response.content if response.reasoning_content: completion_attributes[f"gen_ai.completion.{index}.reasoning_content"] = response.reasoning_content if response.tool_calls: completion_attributes[f"gen_ai.completion.{index}.tool_calls"] = json.dumps( [ { "id": tool_call.id, "type": "function", "function": { "name": tool_call.function.name, "arguments": json.loads(tool_call.function.arguments), }, } for tool_call in response.tool_calls ] ) return completion_attributes def _uses_server_side_tools(self) -> bool: """Returns True if any server-side tools are specified in the completions request.""" return any(tool.WhichOneof("tool") != "function" for tool in self._proto.tools) def _auto_detect_multi_output_mode( self, index: Optional[int], outputs: Sequence[Union[chat_pb2.CompletionOutput, chat_pb2.CompletionOutputChunk]] ) -> Optional[int]: """Auto-detects if the server is using multi-output mode and updates the index accordingly. When we expect single-output mode (index=0) but the server returns multiple outputs (likely because it added tools implicitly), this method switches to multi-output mode (index=None) to properly handle all outputs. Args: index: The current index value (0 for single-output, None for multi-output). outputs: The outputs from the response or chunk to check. Returns: The potentially updated index value (None if multi-output mode is detected, otherwise unchanged). """ if index == 0 and outputs: max_output_index = max(output.index for output in outputs) if max_output_index > 0: # Server is using multi-output mode (likely added tools implicitly) return None return index @property def messages(self) -> Sequence[chat_pb2.Message]: """Returns the messages in the conversation.""" return self._proto.messages def _make_compact_request(self) -> chat_pb2.CompactContextRequest: """Creates a CompactContextRequest from the chat's current model and messages.""" return chat_pb2.CompactContextRequest( model=self._proto.model, input=self._proto.messages, ) def _apply_compaction(self, response: "CompactContextResponse") -> None: """Replaces all messages with the compaction message to reduce payload size.""" self.append(response) def user(*args: Content) -> chat_pb2.Message: """Creates a new message of role "user".""" return chat_pb2.Message(role=chat_pb2.MessageRole.ROLE_USER, content=[_process_content(c) for c in args]) def assistant(*args: Content) -> chat_pb2.Message: """Creates a new message of role "assistant".""" return chat_pb2.Message(role=chat_pb2.MessageRole.ROLE_ASSISTANT, content=[_process_content(c) for c in args]) def system(*args: Content) -> chat_pb2.Message: """Creates a new message of role "system".""" return chat_pb2.Message(role=chat_pb2.MessageRole.ROLE_SYSTEM, content=[_process_content(c) for c in args]) def developer(*args: Content) -> chat_pb2.Message: """Creates a new message of role "developer". Note: This role is only supported by model versions higher than `grok-4.1`/`grok-4-1` (not included). Using the `developer` role in `grok-4.1` or below will be converted to `system` message by the backend. """ return chat_pb2.Message(role=chat_pb2.MessageRole.ROLE_DEVELOPER, content=[_process_content(c) for c in args]) def tool_result(result: str, tool_call_id: Optional[str] = None) -> chat_pb2.Message: """Creates a new message of role "tool". Use this to provide the result of a client-side tool execution back to the model in the conversation history. This enables multi-turn tool use and agentic workflows: the model calls a tool, you execute it, then append the result. Args: result: The string output/result from your tool's execution. This will be sent to the model as content. tool_call_id: Optional ID linking this result to a specific tool call (should match `tool_call.id` from the assistant's tool_calls list). Essential for parallel_tool_calls or multiple tools to associate results correctly. If omitted (for single tool calls), the association may still work but is less explicit. Examples: Basic function calling loop: ```python from xai_sdk.chat import tool_result import json # ... after chat.sample() or in stream if response.tool_calls: for tool_call in response.tool_calls: # Parse and execute args = json.loads(tool_call.function.arguments) tool_output = get_weather(args["city"]) # Your tool function # Append with tool_call_id for proper linking chat.append(tool_result(tool_output, tool_call_id=tool_call.id)) # Continue conversation response = chat.sample() ``` See `examples/sync/function_calling.py` and `examples/sync/server_side_tools.py` (for mixed client/server tools) for complete patterns. Returns: A `chat_pb2.Message` object with ROLE_TOOL, ready to append to chat. """ return chat_pb2.Message(role=chat_pb2.MessageRole.ROLE_TOOL, content=[text(result)], tool_call_id=tool_call_id) def tool(name: str, description: str, parameters: dict[str, Any]) -> chat_pb2.Tool: """Creates a new tool for function calling in chat conversations. This function defines a tool that the model can call to perform specific tasks, such as executing a function with provided arguments. The tool is represented as a `chat_pb2.Tool` object, which includes a function specification with a name, description, and JSON schema for the parameters. The model uses this schema to generate valid JSON inputs for the function when it decides to call the tool. Tools are typically passed to the `create` method of a chat client to enable function calling in a conversation. Examples: Using a Pydantic model to define the parameter schema: ```python from pydantic import BaseModel, Field from xai_sdk import Client from xai_sdk.chat import system, tool class GetWeatherRequest(BaseModel): city: str = Field(description="The name of the city to get the weather for.") units: Literal["C", "F"] = Field(description="The units to use for the temperature.") client = Client() weather_tool = tool( name="get_weather", description="Get the weather for a given city.", parameters=GetWeatherRequest.model_json_schema(), ) conversation = client.chat.create( model="grok-4.20-non-reasoning", messages=[system("You are a helpful assistant.")], tools=[weather_tool], ) ``` Using an explicit JSON schema definition: ```python from xai_sdk import Client from xai_sdk.chat import system, tool client = Client() weather_tool = tool( name="get_weather", description="Get the weather for a given city.", parameters={ "type": "object", "properties": { "city": {"type": "string", "description": "The name of the city to get the weather for."}, "units": { "type": "string", "description": "The units to use for the temperature.", "enum": ["C", "F"], }, }, "required": ["city", "units"], }, ) conversation = client.chat.create( model="grok-4.20-non-reasoning", messages=[system("You are a helpful assistant.")], tools=[weather_tool], ) ``` Args: name: The name of the function that the model can call. This should be unique and descriptive (e.g., "get_weather"). description: A brief description of what the function does, helping the model understand when to call it (e.g., "Get the weather for a given city."). parameters: A JSON schema dictionary or a dictionary derived from a Pydantic model's `model_json_schema()` that defines the structure and types of the function's input parameters. Returns: A `chat_pb2.Tool` object representing the function, which can be passed to the `tools` parameter of a chat client's `create` method. Note: - The `parameters` dictionary is serialized to a JSON string internally, so it must be JSON-serializable. - A maximum of 128 tools can be provided to a chat conversation. - The model decides whether to call the tool based on the conversation context and the tool's description. """ return chat_pb2.Tool( function=chat_pb2.Function( name=name, description=description, parameters=json.dumps(parameters), ) ) def required_tool(name: str) -> chat_pb2.ToolChoice: """Creates a new tool choice with function name `name`. Use this to force the model to always invoke a specific tool. `name` must be the name of a tool that has been provided in the `tools` parameter of a chat client's `create` method. """ return chat_pb2.ToolChoice(function_name=name) def text(content: str) -> chat_pb2.Content: """Returns a new content object of type text.""" return chat_pb2.Content(text=content) def image(image_url: str, *, detail: Optional[ImageDetail] = "auto") -> chat_pb2.Content: """Creates a new content object of type image for use in chat messages. Args: image_url: The URL or base64-encoded string of the image. Supported formats are PNG and JPG. If a URL is provided, the image is fetched for each API request without caching. Fetching uses the "XaiImageApiFetch/1.0" user agent with a 5-second timeout. The maximum image size is 10 MiB; larger images or failed fetches will cause the API request to fail. detail: Specifies the image resolution for model processing. One of: - `"auto"`: The system selects an appropriate resolution (default). - `"low"`: Uses a low-resolution image, reducing token usage and increasing speed. - `"high"`: Uses a high-resolution image, increasing token usage and processing time but capturing more detail. Returns: A `chat_pb2.Content` object representing the image content. """ pb_detail = image_pb2.ImageDetail.DETAIL_AUTO if detail == "low": pb_detail = image_pb2.ImageDetail.DETAIL_LOW elif detail == "high": pb_detail = image_pb2.ImageDetail.DETAIL_HIGH return chat_pb2.Content(image_url=image_pb2.ImageUrlContent(image_url=image_url, detail=pb_detail)) @overload def file(file_id: str) -> chat_pb2.Content: ... @overload def file( *, data: bytes, filename: Optional[str] = ..., mime_type: Optional[str] = ..., ) -> chat_pb2.Content: ... @overload def file( *, url: str, filename: Optional[str] = ..., mime_type: Optional[str] = ..., ) -> chat_pb2.Content: ... def file( file_id: Optional[str] = None, *, data: Optional[bytes] = None, url: Optional[str] = None, filename: Optional[str] = None, mime_type: Optional[str] = None, ) -> chat_pb2.Content: """Creates a new content object of type file for use in chat messages. This supports three modes: - Referencing a previously uploaded file via the Files API (`file_id`) - Providing inline file bytes directly in the chat request (`data`) - Providing a public URL to a file attachment (`url`) Exactly one of `file_id`, `data`, or `url` must be provided. Args: file_id: The ID of a previously uploaded file. You can obtain this ID by uploading a file using the Files API (`client.files.upload(...)`). data: Inline file bytes to attach directly in the chat request (no Files API upload required). url: A public URL pointing to a file attachment. The file will be fetched from this URL by the server. filename: Recommended when `data` or `url` is set. Used for display and may help downstream type inference. mime_type: Optional MIME type (e.g. "application/pdf"). Supported for `data` and `url` modes. Returns: A `chat_pb2.Content` object representing a file attachment. """ provided = sum(x is not None for x in (file_id, data, url)) if provided != 1: raise ValueError("Exactly one of `file_id`, `data`, or `url` must be set.") if file_id is not None: if filename is not None or mime_type is not None: raise ValueError("`filename`/`mime_type` are only supported for inline uploads (`data`) or URL (`url`).") return chat_pb2.Content(file=chat_pb2.FileContent(file_id=file_id)) if url is not None: file_content = chat_pb2.FileContent(url=url) else: file_content = chat_pb2.FileContent(data=data) if filename is not None: file_content.filename = filename if mime_type is not None: file_content.mime_type = mime_type return chat_pb2.Content(file=file_content) def _process_content(content: Content) -> chat_pb2.Content: """Converts a `Content` type to a proto.""" if isinstance(content, str): return text(content) else: return content def _reasoning_effort_to_proto(effort: ReasoningEffort) -> chat_pb2.ReasoningEffort: """Converts a `ReasoningEffort` literal to a proto.""" match effort: case "none": return chat_pb2.ReasoningEffort.EFFORT_NONE case "low": return chat_pb2.ReasoningEffort.EFFORT_LOW case "medium": return chat_pb2.ReasoningEffort.EFFORT_MEDIUM case "high": return chat_pb2.ReasoningEffort.EFFORT_HIGH case "xhigh": return chat_pb2.ReasoningEffort.EFFORT_XHIGH case _: raise ValueError(f"Invalid reasoning effort: {effort}. Must be one of: {ReasoningEffort.__args__}") def _include_option_to_proto(include_option: IncludeOption) -> chat_pb2.IncludeOption: """Converts a `IncludeOption` literal to a proto.""" if include_option in IncludeOptionMap: return IncludeOptionMap[include_option] raise ValueError(f"Invalid include option: {include_option}. Must be one of: {IncludeOptionMap.keys()}") def _agent_count_to_proto(agent_count: int) -> chat_pb2.AgentCount: """Converts an `AgentCount` literal to a proto.""" if agent_count in AgentCountMap: return AgentCountMap[agent_count] raise ValueError(f"Invalid agent count: {agent_count}. Must be one of: {list(AgentCountMap.keys())}") def _tool_mode_to_proto(mode: ToolMode) -> chat_pb2.ToolMode: """Converts a `ToolMode` literal to a proto.""" match mode: case "auto": return chat_pb2.ToolMode.TOOL_MODE_AUTO case "none": return chat_pb2.ToolMode.TOOL_MODE_NONE case "required": return chat_pb2.ToolMode.TOOL_MODE_REQUIRED case _: raise ValueError(f"Invalid tool mode: {mode}. Must be one of: {ToolMode.__args__}") def _format_type_to_proto(format_type: ResponseFormat) -> chat_pb2.FormatType: """Converts a `FormatType` literal to a proto.""" match format_type: case "text": return chat_pb2.FORMAT_TYPE_TEXT case "json_object": return chat_pb2.FORMAT_TYPE_JSON_OBJECT case "json_schema": return chat_pb2.FORMAT_TYPE_JSON_SCHEMA case _: raise ValueError(f"Invalid response format: {format_type}. Must be one of: {ResponseFormat.__args__}") class Chunk(ProtoDecorator[chat_pb2.GetChatCompletionChunk]): """Adds convenience functions to the chunk proto.""" _index: int | None def __init__(self, proto: chat_pb2.GetChatCompletionChunk, index: int | None): """Creates a new decorator instance. Args: proto: Chunk proto to wrap. index: Index of the response to track. If set to None, the chunk will expose all assistant outputs. """ super().__init__(proto) self._index = index @property def choices(self) -> Sequence["CompletionOutputChunk"]: """Returns the completion output chunks belonging to this index.""" return [ CompletionOutputChunk(output) for output in self.proto.outputs if output.delta.role == chat_pb2.MessageRole.ROLE_ASSISTANT and (output.index == self._index or self._index is None) ] @property def created(self) -> datetime.datetime: """Returns the creation timestamp of this chunk.""" return self.proto.created.ToDatetime() @property def output(self) -> str: """Concatenates all chunks into a single string.""" return "".join(c.content + c.reasoning_content for c in self.choices) @property def content(self) -> str: """Concatenates all content chunks into a single string.""" return "".join(c.content for c in self.choices) @property def reasoning_content(self) -> str: """Concatenates all reasoning chunks into a single string.""" return "".join(c.reasoning_content for c in self.choices) @property def tool_calls(self) -> Sequence[chat_pb2.ToolCall]: """Returns the tool calls of this chunk.""" tool_calls = [] for c in self.choices: tool_calls.extend(c.tool_calls) return tool_calls @property def server_side_tool_usage(self) -> dict[str, int]: """Returns the server side tools used for this chunk.""" tools_used = [usage_pb2.ServerSideTool.Name(tool) for tool in self.proto.usage.server_side_tools_used] return dict(Counter(tools_used)) @property def citations(self) -> Sequence[str]: """Returns the citations of this chunk.""" return self.proto.citations @property def inline_citations(self) -> Sequence[chat_pb2.InlineCitation]: """Returns the inline citations of this chunk. Inline citations provide structured citation metadata with position information, enabling you to know exactly where in the response text each citation appears. Each InlineCitation contains: - id: Display number as a string (e.g., "1", "2") - start_index: Character position where the citation starts in the response text - end_index: Character position where the citation ends (exclusive) - web_citation: Present if the citation is from a web source - x_citation: Present if the citation is from an X/Twitter source - collections_citation: Present if the citation is from a collections search Note: Inline citations are only populated when `include=["inline_citations"]` is passed when creating the chat. """ inline_citations = [] for c in self.choices: inline_citations.extend(c.proto.delta.citations) return inline_citations @property def tool_outputs(self) -> Sequence["CompletionOutputChunk"]: """Returns the completion output chunks that contain the tool outputs.""" return [ CompletionOutputChunk(output) for output in self.proto.outputs if output.delta.role == chat_pb2.MessageRole.ROLE_TOOL and (output.index == self._index or self._index is None) ] @property def debug_output(self) -> chat_pb2.DebugOutput: """Returns the debug output of this chunk.""" return self.proto.debug_output def __str__(self): """Concatenates all chunks into a single string.""" return "".join(c.content + c.reasoning_content for c in self.choices) class CompletionOutputChunk(ProtoDecorator[chat_pb2.CompletionOutputChunk]): """Adds convenience functions to the completion output chunk proto.""" @property def content(self) -> str: """Returns the main content/answer of this completion output chunk.""" return self.proto.delta.content @property def reasoning_content(self) -> str: """Returns the reasoning content of this completion output chunk.""" return self.proto.delta.reasoning_content @property def role(self) -> str: """Returns the role of this completion output chunk.""" return chat_pb2.MessageRole.Name(self.proto.delta.role) @property def tool_calls(self) -> Sequence[chat_pb2.ToolCall]: """Returns the tool calls of this completion output chunk.""" return self.proto.delta.tool_calls @property def finish_reason(self) -> sample_pb2.FinishReason: """Returns the finish reason of this completion output chunk.""" return self.proto.finish_reason class ImageGenerationOutput(ProtoDecorator[chat_pb2.CompletionOutput]): """A single image produced by the server-side `image_generation` tool. Exposes the generated image as decoded `image` bytes (with `mime_type` and `data_url` describing the payload), the `image_uuid` used to reference the image in follow-up edit requests, and the completed `tool_call` that produced it. """ _tool_call: chat_pb2.ToolCall _data_url: str _image_uuid: str _image: bytes def __init__(self, proto: chat_pb2.CompletionOutput, tool_call: chat_pb2.ToolCall) -> None: """Initializes a new instance of the `ImageGenerationOutput` class. Args: proto: The `ROLE_TOOL` completion output carrying the image generation result. tool_call: The completed image generation tool call that produced this image. Raises: ValueError: If the output content is not a well-formed image generation result envelope. Completed calls always carry the full envelope in a single chunk, so this indicates a corrupt payload, not a normal stream state. """ super().__init__(proto) self._tool_call = tool_call content = proto.message.content error_message = f"Output content is not a well-formed image generation result envelope: {content[:64]!r}" try: envelope = json.loads(content) except json.JSONDecodeError as e: raise ValueError(error_message) from e if not isinstance(envelope, dict) or envelope.get("__type") != "image_generation_result": raise ValueError(error_message) data_url = envelope.get("result") if not isinstance(data_url, str) or not data_url.startswith("data:image/") or "base64," not in data_url: raise ValueError(error_message) self._data_url = data_url image_uuid = envelope.get("image_uuid", "") self._image_uuid = image_uuid if isinstance(image_uuid, str) else "" _, encoded = data_url.split("base64,", 1) self._image = base64.b64decode(encoded) @property def tool_call(self) -> chat_pb2.ToolCall: """The completed image generation tool call that produced this image. `tool_call.function.name` distinguishes text-to-image generations (`imagine_text_to_image`) from image edits (`imagine_image_to_image`). """ return self._tool_call @property def data_url(self) -> str: """The generated image as a `data:;base64,` URL.""" return self._data_url @property def mime_type(self) -> str: """The MIME type of the generated image (e.g. `image/jpeg`).""" return self._data_url.removeprefix("data:").split(";", 1)[0] @property def image_uuid(self) -> str: """The short reference ID the model uses to address this image in follow-up tool calls. Empty if the server did not report one. """ return self._image_uuid @property def image(self) -> bytes: """The generated image as raw bytes.""" return self._image class _ResponseProtoDecorator(ProtoDecorator[chat_pb2.GetChatCompletionResponse]): def __init__(self, proto: chat_pb2.GetChatCompletionResponse) -> None: """Initialize with proto and content buffers for efficient accumulation.""" super().__init__(proto) # Buffers for efficient string accumulation: dict[output_index, [chunk_ones_content, chunk_twos_content, ...]] # for chunks with the same index self._content_buffers: dict[int, list[str]] = defaultdict(list) self._reasoning_content_buffers: dict[int, list[str]] = defaultdict(list) self._encrypted_content_buffers: dict[int, list[str]] = defaultdict(list) self._proto_in_sync = True def _sync_buffers_to_proto(self) -> None: """Materialize buffered content into proto messages.""" if self._proto_in_sync: return for index, buffer in self._content_buffers.items(): if buffer and index < len(self._proto.outputs): self._proto.outputs[index].message.content = "".join(buffer) self._content_buffers[index] = [self._proto.outputs[index].message.content] for index, buffer in self._reasoning_content_buffers.items(): if buffer and index < len(self._proto.outputs): self._proto.outputs[index].message.reasoning_content = "".join(buffer) self._reasoning_content_buffers[index] = [self._proto.outputs[index].message.reasoning_content] for index, buffer in self._encrypted_content_buffers.items(): if buffer and index < len(self._proto.outputs): self._proto.outputs[index].message.encrypted_content = "".join(buffer) self._encrypted_content_buffers[index] = [self._proto.outputs[index].message.encrypted_content] self._proto_in_sync = True @property def proto(self) -> chat_pb2.GetChatCompletionResponse: """Ensure buffers are synced before returning proto.""" self._sync_buffers_to_proto() return self._proto def process_chunk(self, chunk: chat_pb2.GetChatCompletionChunk): # Consolidate the response. self._proto.usage.CopyFrom(chunk.usage) self._proto.created.CopyFrom(chunk.created) self._proto.id = chunk.id self._proto.model = chunk.model self._proto.system_fingerprint = chunk.system_fingerprint if chunk.service_tier: self._proto.service_tier = chunk.service_tier self._proto.citations.extend(chunk.citations) # Make sure all chunk outputs has corresponding response outputs. if chunk.outputs: max_index = max(c.index for c in chunk.outputs) if max_index >= len(self._proto.outputs): self._proto.outputs.extend( [chat_pb2.CompletionOutput() for _ in range(max_index + 1 - len(self._proto.outputs))] ) for c in chunk.outputs: choice = self._proto.outputs[c.index] choice.index = c.index choice.message.role = c.delta.role choice.message.tool_calls.extend(c.delta.tool_calls) # c.delta.citations represents the inline citations on this chunk choice.message.citations.extend(c.delta.citations) choice.finish_reason = c.finish_reason # Accumulate content in buffers instead of concatenating strings if c.delta.content: self._content_buffers[c.index].append(c.delta.content) self._proto_in_sync = False if c.delta.reasoning_content: self._reasoning_content_buffers[c.index].append(c.delta.reasoning_content) self._proto_in_sync = False if c.delta.encrypted_content: self._encrypted_content_buffers[c.index].append(c.delta.encrypted_content) self._proto_in_sync = False class Response(_ResponseProtoDecorator): """Response of a chat request.""" # A single request can produce multiple responses. This index is used to retrieve the content of # a single answer from the response proto. _index: int | None def __init__(self, response: chat_pb2.GetChatCompletionResponse, index: int | None) -> None: """Initializes a new instance of the `Response` class. Args: response: The response proto, which can hold multiple answers. index: The index of the answer this class exposes via its convenience methods. If set to None, the response will expose all answers, the content and reasoning content will be only from the assistant response. """ super().__init__(response) self._index = index def _get_output(self, *, sync: bool = False) -> chat_pb2.CompletionOutput: # Sync buffers to proto only when content is needed if sync: self._sync_buffers_to_proto() outputs = [ output for output in self._proto.outputs if output.message.role == chat_pb2.MessageRole.ROLE_ASSISTANT and (output.index == self._index or self._index is None) ] if not outputs: return chat_pb2.CompletionOutput() return outputs[-1] @property def id(self) -> str: """Returns the id of this response.""" return self._proto.id @property def created(self) -> datetime.datetime: """Returns the creation timestamp of this response.""" return self._proto.created.ToDatetime() @property def content(self) -> str: """Returns the answer content of this response.""" return self._get_output(sync=True).message.content @property def encrypted_content(self) -> str: """Returns the encrypted reasoning content from the model response.""" return self._get_output(sync=True).message.encrypted_content @property def role(self) -> str: """Returns the role of this response.""" return chat_pb2.MessageRole.Name(self._get_output(sync=False).message.role) @property def usage(self) -> usage_pb2.SamplingUsage: """Returns the usage of this response.""" return self._proto.usage @property def cost_usd(self) -> Optional[float]: """Cost of this single request in USD, or None if the server did not report it. This is the cost of the gRPC request that produced this `Response` only -- not the running total of a multi-turn conversation. When the same `Chat` object is sampled across multiple turns (each `chat.sample()` / `chat.stream()` call is a separate request), accumulate per-response values yourself, e.g.: total_cost_usd = 0.0 while True: prompt = input("You: ") if prompt.lower() == "exit": break chat.append(user(prompt)) response = chat.sample() chat.append(response) # `or 0.0` skips turns where the server did not report a cost. total_cost_usd += response.cost_usd or 0.0 print(f"Grok: {response.content}") print(f"Total cost: ${total_cost_usd:.4f}") Server-side tools (`web_search`, `x_search`, `code_execution`, etc.) are an exception: the entire model + tool agentic loop runs inside one gRPC request, so `response.cost_usd` already covers all internal model decodes and tool calls. No accumulation is needed for those -- print `response.cost_usd` directly. """ return cost_usd_from_usage(self._proto.usage) @property def reasoning_content(self) -> str: """Returns the reasoning trace generated by the model. This is only available for models that support reasoning. """ return self._get_output(sync=True).message.reasoning_content @property def finish_reason(self) -> str: """Returns the finish reason of this response.""" return sample_pb2.FinishReason.Name(self._get_output(sync=False).finish_reason) @property def logprobs(self) -> chat_pb2.LogProbs: """Returns the logprobs of this response.""" return self._get_output(sync=False).logprobs @property def system_fingerprint(self) -> str: """Returns the system fingerprint of this response.""" return self.proto.system_fingerprint @property def service_tier(self) -> ServiceTier: """Returns the processing tier used for this request. Returns ``"priority"`` if the request was served at the priority tier, or ``"default"`` otherwise. """ return service_tier_from_proto(self.proto.service_tier) @property def tool_calls(self) -> Sequence[chat_pb2.ToolCall]: """Returns the all tool calls of this response.""" return [ tc for output in self.proto.outputs if output.message.role == chat_pb2.MessageRole.ROLE_ASSISTANT for tc in output.message.tool_calls ] @property def citations(self) -> Sequence[str]: """Returns the citations of this response.""" return self.proto.citations @property def inline_citations(self) -> Sequence[chat_pb2.InlineCitation]: """Returns the inline citations of this response. Inline citations provide structured citation metadata with position information, enabling you to know exactly where in the response text each citation appears. Each InlineCitation contains: - id: Display number as a string (e.g., "1", "2") - start_index: Character position where the citation starts in the response text - end_index: Character position where the citation ends (exclusive) - web_citation: Present if the citation is from a web source - x_citation: Present if the citation is from an X/Twitter source - collections_citation: Present if the citation is from a collections search Note: Inline citations are only populated when `include=["inline_citations"]` is passed when creating the chat. """ return [ citation for output in self.proto.outputs if output.message.role == chat_pb2.MessageRole.ROLE_ASSISTANT for citation in output.message.citations ] @property def tool_outputs(self) -> Sequence[chat_pb2.CompletionOutput]: """Returns the output entries that contain the tool outputs.""" return [output for output in self.proto.outputs if output.message.role == chat_pb2.MessageRole.ROLE_TOOL] @property def image_outputs(self) -> Sequence[ImageGenerationOutput]: """Returns the images generated by the server-side `image_generation` tool. One entry per completed image generation/edit call, in response order. Failed calls carry an error payload instead of an image and are omitted. Safe to access while streaming: the server delivers a completed call and its full result envelope in the same chunk, so entries appear fully formed as soon as their tool call completes. A `ValueError` here means a completed call carried a corrupt envelope (a server bug), not a normal stream state. Example: ``` for i, output in enumerate(response.image_outputs): extension = output.mime_type.removeprefix("image/") with open(f"image_{i}.{extension}", "wb") as f: f.write(output.image) ``` """ image_outputs = [] for output in self.tool_outputs: completed_calls = [ tool_call for tool_call in output.message.tool_calls if tool_call.type == chat_pb2.TOOL_CALL_TYPE_IMAGE_GENERATION_TOOL and tool_call.status == chat_pb2.TOOL_CALL_STATUS_COMPLETED ] if completed_calls: image_outputs.append(ImageGenerationOutput(output, completed_calls[0])) return image_outputs @property def server_side_tool_usage(self) -> dict[str, int]: """Returns the server side tools used for this response.""" tools_used = [usage_pb2.ServerSideTool.Name(tool) for tool in self.proto.usage.server_side_tools_used] return dict(Counter(tools_used)) @property def request_settings(self) -> chat_pb2.RequestSettings: """Returns the request settings, i.e. the model parameters set on the request used to generate this response.""" return self.proto.settings @property def debug_output(self) -> chat_pb2.DebugOutput: """Returns the debug output of this response. Only available to trusted testers.""" return self.proto.debug_output class CompactContextResponse(ProtoDecorator[chat_pb2.CompactContextResponse]): """Response from a context compaction request. Wraps a ``CompactContextResponse`` proto with convenience accessors. The ``encrypted_content`` string is an opaque, compact representation of the input messages produced by the model. It can be passed back as an ``encrypted_content`` field on an assistant message in a follow-up chat request to hydrate the compacted context. """ @property def id(self) -> str: """Unique identifier for this compaction response.""" return self._proto.id @property def encrypted_content(self) -> str: """Opaque compacted representation of the input messages.""" return self._proto.encrypted_content @property def dropped_message_count(self) -> int: """Number of input messages that were dropped during compaction.""" return self._proto.dropped_message_count @property def usage(self) -> usage_pb2.SamplingUsage: """Token usage for the compaction request.""" return self._proto.usage