18
18
from azure .ai .agents .models ._enums import (
19
19
AgentsResponseFormatMode ,
20
20
MessageRole ,
21
+ MessageStatus ,
21
22
RunStepStatus ,
23
+ RunStepType ,
22
24
)
23
25
from azure .ai .agents .models import (
24
26
MessageAttachment ,
28
30
RunStepDeltaChunk ,
29
31
RunStepError ,
30
32
RunStepFunctionToolCall ,
33
+ RunStepOpenAPIToolCall ,
31
34
RunStepToolCallDetails ,
32
35
RunStepCodeInterpreterToolCall ,
33
36
RunStepBingGroundingToolCall ,
@@ -439,6 +442,12 @@ def _process_tool_calls(self, step: RunStep) -> List[Dict[str, Any]]:
439
442
"type" : t .type ,
440
443
t .type : t .bing_grounding ,
441
444
}
445
+ elif isinstance (t , RunStepOpenAPIToolCall ):
446
+ tool_call = {
447
+ "id" : t .id ,
448
+ "type" : t .type ,
449
+ 'function' : t .as_dict ().get ('function' , {})
450
+ }
442
451
else :
443
452
tool_details = t .as_dict ()[t .type ]
444
453
@@ -2057,7 +2066,8 @@ def on_thread_message(self, message: "ThreadMessage") -> None: # type: ignore[f
2057
2066
else :
2058
2067
retval = super ().on_thread_message (message ) # pylint: disable=assignment-from-none # type: ignore
2059
2068
2060
- if message .status in {"completed" , "incomplete" }:
2069
+ # Message status may be in progress, even if the thread.message.completed event has arrived.
2070
+ if message .status in {MessageStatus .COMPLETED , MessageStatus .INCOMPLETE } or (message .status == MessageStatus .IN_PROGRESS and message .content ):
2061
2071
self .last_message = message
2062
2072
2063
2073
return retval # type: ignore
@@ -2081,14 +2091,14 @@ def on_run_step(self, step: "RunStep") -> None: # type: ignore[func-returns-val
2081
2091
retval = super ().on_run_step (step ) # pylint: disable=assignment-from-none # type: ignore
2082
2092
2083
2093
if (
2084
- step .type == "tool_calls"
2094
+ step .type == RunStepType . TOOL_CALLS
2085
2095
and isinstance (step .step_details , RunStepToolCallDetails )
2086
2096
and step .status == RunStepStatus .COMPLETED
2087
2097
):
2088
2098
self .instrumentor ._add_tool_assistant_message_event ( # pylint: disable=protected-access # pyright: ignore [reportFunctionMemberAccess]
2089
2099
self .span , step
2090
2100
)
2091
- elif step .type == "message_creation" and step .status == RunStepStatus .COMPLETED :
2101
+ elif step .type == RunStepType . MESSAGE_CREATION and step .status == RunStepStatus .COMPLETED :
2092
2102
self .instrumentor .add_thread_message_event (self .span , cast (ThreadMessage , self .last_message ), step .usage )
2093
2103
if (
2094
2104
self .span
@@ -2194,7 +2204,8 @@ async def on_thread_message(self, message: "ThreadMessage") -> None: # type: ig
2194
2204
else :
2195
2205
retval = await super ().on_thread_message (message ) # type: ignore
2196
2206
2197
- if message .status in {"completed" , "incomplete" }:
2207
+ # Message status may be in progress, even if the thread.message.completed event has arrived.
2208
+ if message .status in {MessageStatus .COMPLETED , MessageStatus .INCOMPLETE } or (message .status == MessageStatus .IN_PROGRESS and message .content ):
2198
2209
self .last_message = message
2199
2210
2200
2211
return retval # type: ignore
@@ -2218,14 +2229,14 @@ async def on_run_step(self, step: "RunStep") -> None: # type: ignore[func-retur
2218
2229
retval = await super ().on_run_step (step ) # type: ignore
2219
2230
2220
2231
if (
2221
- step .type == "tool_calls"
2232
+ step .type == RunStepType . TOOL_CALLS
2222
2233
and isinstance (step .step_details , RunStepToolCallDetails )
2223
2234
and step .status == RunStepStatus .COMPLETED
2224
2235
):
2225
2236
self .instrumentor ._add_tool_assistant_message_event ( # pylint: disable=protected-access # pyright: ignore [reportFunctionMemberAccess]
2226
2237
self .span , step
2227
2238
)
2228
- elif step .type == "message_creation" and step .status == RunStepStatus .COMPLETED :
2239
+ elif step .type == RunStepType . MESSAGE_CREATION and step .status == RunStepStatus .COMPLETED :
2229
2240
self .instrumentor .add_thread_message_event (self .span , cast (ThreadMessage , self .last_message ), step .usage )
2230
2241
if (
2231
2242
self .span
0 commit comments