Skip to content

datarobot_genai.dragent.frontends.tool_call_registry

tool_call_registry

Per-run handoff of LLM-issued tool_call_id from converter to adaptor.

Two-phase correlation tolerates parallel same-name calls completing out of order: register (converter) appends to a per-name FIFO; bind (adaptor, FUNCTION_START) pops the head into a UUID-keyed map; pop (adaptor, FUNCTION_END) drains by UUID.

Deferred end events prevent ToolCallEndEvent from racing ahead of ToolCallArgsEvent chunks still being streamed by the converter.

mark_args_done 

mark_args_done(tool_call_id: str) -> list[Any]

Mark argument streaming complete for tool_call_id.

Returns (and removes) any end/result events that the step adaptor deferred while arguments were still in flight.

Source code in datarobot_genai/dragent/frontends/tool_call_registry.py 
69
70
71
72
73
74
75
76
77
def mark_args_done(tool_call_id: str) -> list[Any]:
    """Mark argument streaming complete for *tool_call_id*.

    Returns (and removes) any end/result events that the step adaptor
    deferred while arguments were still in flight.
    """
    _, _, args_done, deferred = _get()
    args_done.add(tool_call_id)
    return deferred.pop(tool_call_id, [])

is_args_done 

is_args_done(tool_call_id: str) -> bool

Check whether the stream converter has finished emitting args for this call.

Source code in datarobot_genai/dragent/frontends/tool_call_registry.py 
80
81
82
def is_args_done(tool_call_id: str) -> bool:
    """Check whether the stream converter has finished emitting args for this call."""
    return tool_call_id in _get()[2]

defer_tool_end 

defer_tool_end(tool_call_id: str, events: list[Any]) -> None

Stash end/result events until argument streaming finishes.

Source code in datarobot_genai/dragent/frontends/tool_call_registry.py 
85
86
87
def defer_tool_end(tool_call_id: str, events: list[Any]) -> None:
    """Stash end/result events until argument streaming finishes."""
    _get()[3][tool_call_id] = events