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) -> "GenerateResult":
27    """
28    Generate one message based on the given context.
29    Parts of the message will be streamed to the specified callbacks if provided.
30
31    Args:
32        chat_provider: The chat provider to use for generation.
33        system_prompt: The system prompt to use for generation.
34        tools: The tools available for the model to call.
35        history: The message history to use for generation.
36        on_message_part: An optional callback to be called for each raw message part.
37        on_tool_call: An optional callback to be called for each complete tool call.
38
39    Returns:
40        A tuple of the generated message and the token usage (if available).
41        All parts in the message are guaranteed to be complete and merged as much as possible.
42
43    Raises:
44        APIConnectionError: If the API connection fails.
45        APITimeoutError: If the API request times out.
46        APIStatusError: If the API returns a status code of 4xx or 5xx.
47        APIEmptyResponseError: If the API returns an empty response.
48        ChatProviderError: If any other recognized chat provider error occurs.
49    """
50    message = Message(role="assistant", content=[])
51    pending_part: StreamedMessagePart | None = None  # message part that is currently incomplete
52
53    logger.trace("Generating with history: {history}", history=history)
54    stream = await chat_provider.generate(system_prompt, tools, history)
55    async for part in stream:
56        logger.trace("Received part: {part}", part=part)
57        if on_message_part:
58            await callback(on_message_part, part.model_copy(deep=True))
59
60        if pending_part is None:
61            pending_part = part
62        elif not pending_part.merge_in_place(part):  # try merge into the pending part
63            # unmergeable part must push the pending part to the buffer
64            _message_append(message, pending_part)
65            if isinstance(pending_part, ToolCall) and on_tool_call:
66                await callback(on_tool_call, pending_part)
67            pending_part = part
68
69    # end of message
70    if pending_part is not None:
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
75    if not message.content and not message.tool_calls:
76        raise APIEmptyResponseError("The API returned an empty response.")
77
78    return GenerateResult(
79        id=stream.id,
80        message=message,
81        usage=stream.usage,
82    )

85@dataclass(frozen=True, slots=True)
86class GenerateResult:
87    """The result of a generation."""
88
89    id: str | None
90    """The ID of the generated message."""
91    message: Message
92    """The generated message."""
93    usage: TokenUsage | None
94    """The token usage of the generated message."""

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

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