Обработка вызовов инструментов в потоковых ответах

# In a non-streaming response, tool call is complete:
# choice.message.tool_calls[0].function.arguments = '{"city": "Paris"}'

# In a streaming response, arguments arrive in pieces:
# chunk 1: delta.tool_calls[0].function.arguments = '{'
# chunk 2: delta.tool_calls[0].function.arguments = '"city"'
# chunk 3: delta.tool_calls[0].function.arguments = ': "'
# chunk 4: delta.tool_calls[0].function.arguments = 'Paris'
# chunk 5: delta.tool_calls[0].function.arguments = '"}'
# You must concatenate these before JSON.parse can work

async def detect_stream_type(messages, tools):
    stream = await async_client.chat.completions.create(
        model='gpt-4o-mini',
        messages=messages,
        tools=tools,
        stream=True,
    )
    response_type = 'text'
    async for chunk in stream:
        choice = chunk.choices[0]
        if choice.delta.tool_calls:  # tool call fragment arriving
            response_type = 'tool_call'
        if choice.finish_reason == 'tool_calls':
            print('Model wants to call a function')
        elif choice.finish_reason == 'stop':
            print('Normal text response')
    return response_type

from collections import defaultdict

async def collect_streamed_tool_calls(messages, tools):
    stream = await async_client.chat.completions.create(
        model='gpt-4o-mini',
        messages=messages,
        tools=tools,
        stream=True,
    )

    tool_call_buffers = defaultdict(lambda: {'name': '', 'id': '', 'arguments': ''})
    text_buffer = ''

    async for chunk in stream:
        delta = chunk.choices[0].delta

        if delta.content:  # text content
            text_buffer += delta.content

        if delta.tool_calls:
            for tc in delta.tool_calls:
                idx = tc.index
                if tc.id:
                    tool_call_buffers[idx]['id'] = tc.id
                if tc.function.name:
                    tool_call_buffers[idx]['name'] += tc.function.name
                if tc.function.arguments:
                    tool_call_buffers[idx]['arguments'] += tc.function.arguments

    return text_buffer, dict(tool_call_buffers)

import json

# Example tool registry
tools_registry = {
    'get_weather': lambda city, unit='celsius': {'temp': 22, 'desc': 'sunny', 'city': city},
    'search_docs': lambda query, top_k=3: [{'title': 'Doc 1', 'snippet': 'Relevant info...'}],
}

def execute_tool_calls(tool_call_buffers: dict) -> list[dict]:
    tool_messages = []
    for idx in sorted(tool_call_buffers.keys()):
        tc = tool_call_buffers[idx]
        func_name = tc['name']
        args = json.loads(tc['arguments'])

        if func_name in tools_registry:
            result = tools_registry[func_name](**args)
        else:
            result = {'error': f'Unknown function: {func_name}'}

        tool_messages.append({
            'role': 'tool',
            'tool_call_id': tc['id'],
            'content': json.dumps(result),
        })
    return tool_messages

async def streaming_agent_loop(initial_messages, tools):
    messages = list(initial_messages)
    max_iterations = 5

    for iteration in range(max_iterations):
        text, tool_calls = await collect_streamed_tool_calls(messages, tools)

        if tool_calls:
            # Append assistant message with tool calls
            assistant_msg = {
                'role': 'assistant',
                'content': text or None,
                'tool_calls': [
                    {'id': tc['id'], 'type': 'function',
                     'function': {'name': tc['name'], 'arguments': tc['arguments']}}
                    for tc in tool_calls.values()
                ]
            }
            messages.append(assistant_msg)

            # Execute tools and append results
            tool_results = execute_tool_calls(tool_calls)
            messages.extend(tool_results)
        else:
            # No more tool calls — final text response
            print('Final answer:', text)
            return text

    return 'Max iterations reached'

async def stream_with_tools(messages, tools):
    stream = await async_client.chat.completions.create(
        model='gpt-4o-mini', messages=messages, tools=tools, stream=True
    )
    tool_buffers = defaultdict(lambda: {'name': '', 'id': '', 'arguments': ''})
    text_parts = []

    async for chunk in stream:
        delta = chunk.choices[0].delta
        finish = chunk.choices[0].finish_reason

        if delta.content:
            text_parts.append(delta.content)
            yield ('text', delta.content)   # stream to client immediately

        if delta.tool_calls:
            for tc in delta.tool_calls:
                if tc.id:    tool_buffers[tc.index]['id'] = tc.id
                if tc.function.name: tool_buffers[tc.index]['name'] += tc.function.name
                if tc.function.arguments: tool_buffers[tc.index]['arguments'] += tc.function.arguments

        if finish == 'tool_calls':
            yield ('tool_calls', dict(tool_buffers))  # signal tool execution needed

import asyncio

async def execute_tool_calls_parallel(tool_call_buffers: dict) -> list[dict]:
    async def execute_one(idx, tc):
        func_name = tc['name']
        args = json.loads(tc['arguments'])
        if func_name in async_tools_registry:
            result = await async_tools_registry[func_name](**args)
        else:
            result = {'error': f'Unknown function: {func_name}'}
        return {
            'role': 'tool',
            'tool_call_id': tc['id'],
            'content': json.dumps(result),
        }

    tasks = [execute_one(idx, tc) for idx, tc in sorted(tool_call_buffers.items())]
    return await asyncio.gather(*tasks)

async def full_tool_calling_stream(question: str, tools: list):
    messages = [{'role': 'user', 'content': question}]

    # First request: may produce tool calls
    tool_buffers = {}
    text1 = ''
    async for event_type, data in stream_with_tools(messages, tools):
        if event_type == 'text':
            text1 += data
            yield data  # stream partial text if any
        elif event_type == 'tool_calls':
            tool_buffers = data

    if tool_buffers:
        # Execute tools, then get final streaming answer
        tool_results = await execute_tool_calls_parallel(tool_buffers)
        messages += [{  # assistant tool call message
            'role': 'assistant',
            'tool_calls': [
                {'id': tc['id'], 'type': 'function',
                 'function': {'name': tc['name'], 'arguments': tc['arguments']}}
                for tc in tool_buffers.values()
            ]
        }] + tool_results

        # Second request: final answer streams directly
        async for token in token_stream(messages):  # from earlier lesson
            yield token

import json

async def stream_with_progress(question, tools):
    messages = [{'role': 'user', 'content': question}]
    tool_buffers = {}

    async for event_type, data in stream_with_tools(messages, tools):
        if event_type == 'tool_calls':
            tool_buffers = data

    for tc in tool_buffers.values():
        args = json.loads(tc['arguments'])
        yield f'data: {json.dumps({"type": "tool_start", "function": tc["name"], "args": args})}\n\n'

        result = tools_registry.get(tc['name'], lambda **kw: {})(** args)

        yield f'data: {json.dumps({"type": "tool_done", "function": tc["name"]})}\n\n'

    # Then stream final answer...

def safe_execute_tool(func_name: str, args: dict) -> str:
    try:
        if func_name not in tools_registry:
            return json.dumps({'error': f'Function {func_name!r} not found'})
        result = tools_registry[func_name](**args)
        return json.dumps(result)
    except TypeError as e:
        return json.dumps({'error': f'Invalid arguments: {str(e)}'})
    except Exception as e:
        return json.dumps({'error': f'Execution failed: {str(e)}'})

# Tool result message with error handled
tool_message = {
    'role': 'tool',
    'tool_call_id': tc['id'],
    'content': safe_execute_tool(tc['name'], json.loads(tc['arguments'])),
}