kosong

 18async def generate(
 19    chat_provider: ChatProvider,
 20    system_prompt: str,
 21    tools: Sequence[Tool],
 22    history: Sequence[Message],
 23    *,
 24    on_message_part: Callback[[StreamedMessagePart], None] | None = None,
 25    on_tool_call: Callback[[ToolCall], None] | None = None,
 26    on_trace_id: Callback[[str | None], None] | None = None,
 27) -> "GenerateResult":
 28    """
 29    Generate one message based on the given context.
 30    Parts of the message will be streamed to the specified callbacks if provided.
 31
 32    Args:
 33        chat_provider: The chat provider to use for generation.
 34        system_prompt: The system prompt to use for generation.
 35        tools: The tools available for the model to call.
 36        history: The message history to use for generation.
 37        on_message_part: An optional callback to be called for each raw message part.
 38        on_tool_call: An optional callback to be called for each complete tool call.
 39        on_trace_id: An optional callback fired with the request's ``x-trace-id``
 40            response header as soon as it is available (before streaming starts).
 41
 42    Returns:
 43        A tuple of the generated message and the token usage (if available).
 44        All parts in the message are guaranteed to be complete and merged as much as possible.
 45
 46    Raises:
 47        APIConnectionError: If the API connection fails.
 48        APITimeoutError: If the API request times out.
 49        APIStatusError: If the API returns a status code of 4xx or 5xx.
 50        APIEmptyResponseError: If the API returns an empty response.
 51        ChatProviderError: If any other recognized chat provider error occurs.
 52    """
 53    message = Message(role="assistant", content=[])
 54    pending_part: StreamedMessagePart | None = None  # message part that is currently incomplete
 55
 56    logger.trace("Generating with history: {history}", history=history)
 57    stream = await chat_provider.generate(system_prompt, tools, history)
 58    if on_trace_id:
 59        # getattr for robustness against third-party StreamedMessage
 60        # implementations that predate the trace_id property.
 61        await callback(on_trace_id, getattr(stream, "trace_id", None))
 62    async for part in stream:
 63        logger.trace("Received part: {part}", part=part)
 64        if on_message_part:
 65            await callback(on_message_part, part.model_copy(deep=True))
 66
 67        if pending_part is None:
 68            pending_part = part
 69        elif not pending_part.merge_in_place(part):  # try merge into the pending part
 70            # unmergeable part must push the pending part to the buffer
 71            _message_append(message, pending_part)
 72            if isinstance(pending_part, ToolCall) and on_tool_call:
 73                await callback(on_tool_call, pending_part)
 74            pending_part = part
 75
 76    # end of message
 77    if pending_part is not None:
 78        _message_append(message, pending_part)
 79        if isinstance(pending_part, ToolCall) and on_tool_call:
 80            await callback(on_tool_call, pending_part)
 81
 82    if not message.content and not message.tool_calls:
 83        raise APIEmptyResponseError("The API returned an empty response.")
 84
 85    # A response with only ThinkPart (no TextPart, no tool calls) indicates an
 86    # abnormal termination — typically a stream interruption or max_tokens
 87    # exhaustion during reasoning.  The model should always produce visible
 88    # output after thinking; a think-only response is never intentional.
 89    has_think = any(isinstance(p, ThinkPart) for p in message.content)
 90    has_text = any(isinstance(p, TextPart) and p.text.strip() for p in message.content)
 91    if has_think and not has_text and not message.tool_calls:
 92        raise APIEmptyResponseError(
 93            "The API returned a response containing only thinking content "
 94            "without any text or tool calls. This usually indicates the "
 95            "stream was interrupted or the output token budget was exhausted "
 96            "during reasoning."
 97        )
 98
 99    return GenerateResult(
100        id=stream.id,
101        message=message,
102        usage=stream.usage,
103        trace_id=getattr(stream, "trace_id", None),
104    )

107@dataclass(frozen=True, slots=True)
108class GenerateResult:
109    """The result of a generation."""
110
111    id: str | None
112    """The ID of the generated message."""
113    message: Message
114    """The generated message."""
115    usage: TokenUsage | None
116    """The token usage of the generated message."""
117    trace_id: str | None = None
118    """The ``x-trace-id`` response header of the request, if the provider exposes it."""

105async def step(
106    chat_provider: ChatProvider,
107    system_prompt: str,
108    toolset: Toolset,
109    history: Sequence[Message],
110    *,
111    on_message_part: Callback[[StreamedMessagePart], None] | None = None,
112    on_tool_result: Callable[[ToolResult], None] | None = None,
113    on_trace_id: Callback[[str | None], None] | None = None,
114) -> "StepResult":
115    """
116    Run one agent "step". In one step, the function generates LLM response based on the given
117    context for exactly one time. All new message parts will be streamed to `on_message_part` in
118    real-time if provided. Tool calls will be handled by `toolset`. The generated message will be
119    returned in a `StepResult`. Depending on the toolset implementation, the tool calls may be
120    handled asynchronously and the results need to be fetched with `await result.tool_results()`.
121
122    The message history will NOT be modified in this function.
123
124    The token usage will be returned in the `StepResult` if available.
125
126    Raises:
127        APIConnectionError: If the API connection fails.
128        APITimeoutError: If the API request times out.
129        APIStatusError: If the API returns a status code of 4xx or 5xx.
130        APIEmptyResponseError: If the API returns an empty response.
131        ChatProviderError: If any other recognized chat provider error occurs.
132        asyncio.CancelledError: If the step is cancelled.
133    """
134
135    tool_calls: list[ToolCall] = []
136    tool_result_futures: dict[str, ToolResultFuture] = {}
137
138    def future_done_callback(future: ToolResultFuture):
139        if on_tool_result:
140            try:
141                result = future.result()
142                on_tool_result(result)
143            except asyncio.CancelledError:
144                return
145
146    async def on_tool_call(tool_call: ToolCall):
147        tool_calls.append(tool_call)
148        result = toolset.handle(tool_call)
149
150        if isinstance(result, ToolResult):
151            future = ToolResultFuture()
152            future.add_done_callback(future_done_callback)
153            future.set_result(result)
154            tool_result_futures[tool_call.id] = future
155        else:
156            result.add_done_callback(future_done_callback)
157            tool_result_futures[tool_call.id] = result
158
159    try:
160        result = await generate(
161            chat_provider,
162            system_prompt,
163            toolset.tools,
164            history,
165            on_message_part=on_message_part,
166            on_tool_call=on_tool_call,
167            on_trace_id=on_trace_id,
168        )
169    except (ChatProviderError, asyncio.CancelledError):
170        # cancel all the futures to avoid hanging tasks
171        for future in tool_result_futures.values():
172            future.remove_done_callback(future_done_callback)
173            future.cancel()
174        await asyncio.gather(*tool_result_futures.values(), return_exceptions=True)
175        raise
176
177    return StepResult(
178        result.id,
179        result.message,
180        result.usage,
181        tool_calls,
182        tool_result_futures,
183        trace_id=result.trace_id,
184    )

187@dataclass(frozen=True, slots=True)
188class StepResult:
189    id: str | None
190    """The ID of the generated message."""
191
192    message: Message
193    """The message generated in this step."""
194
195    usage: TokenUsage | None
196    """The token usage in this step."""
197
198    tool_calls: list[ToolCall]
199    """All the tool calls generated in this step."""
200
201    _tool_result_futures: dict[str, ToolResultFuture]
202    """@private The futures of the results of the spawned tool calls."""
203
204    trace_id: str | None = None
205    """The ``x-trace-id`` response header of the request, if the provider exposes it."""
206
207    async def tool_results(self) -> list[ToolResult]:
208        """All the tool results returned by corresponding tool calls."""
209        if not self._tool_result_futures:
210            return []
211
212        try:
213            results: list[ToolResult] = []
214            for tool_call in self.tool_calls:
215                future = self._tool_result_futures[tool_call.id]
216                result = await future
217                results.append(result)
218            return results
219        finally:
220            # one exception should cancel all the futures to avoid hanging tasks
221            for future in self._tool_result_futures.values():
222                future.cancel()
223            await asyncio.gather(*self._tool_result_futures.values(), return_exceptions=True)