Repository: pydantic/pydantic-ai
File: pydantic_ai_slim/pydantic_ai/ag_ui.py (Python)

Code before:

            deps = replace(deps, state=state)
        elif raw_state:
            raise UserError(
                f'AG-UI state is provided but `deps` of type `{type(deps).__name__}` does not implement the `StateHandler` protocol: it needs to be a dataclass with a non-optional `state` field.'
            )
        else:
            # `deps` not being a `StateHandler` is OK if there is no state.
            pass

        messages = _messages_from_ag_ui(run_input.messages)

        async with agent.iter(
            user_prompt=None,
            output_type=[output_type or agent.output_type, DeferredToolCalls],
            message_history=messages,
            model=model,
            deps=deps,
            model_settings=model_settings,
            usage_limits=usage_limits,
            usage=usage,
            infer_name=infer_name,
            toolsets=toolsets,
        ) as run:
            async for event in _agent_stream(run):
                yield encoder.encode(event)
            if on_complete is not None:
                if _utils.is_async_callable(on_complete):
                    await on_complete(run)
                else:
                    await _utils.run_in_executor(on_complete, run)
    except _RunError as e:
        yield encoder.encode(
            RunErrorEvent(message=e.message, code=e.code),
        )
    except Exception as e:
        yield encoder.encode(
            RunErrorEvent(message=str(e)),
        )
        raise e
    else:
        yield encoder.encode(
            RunFinishedEvent(
                thread_id=run_input.thread_id,
                run_id=run_input.run_id,
            ),
        )


async def _agent_stream(run: AgentRun[AgentDepsT, Any]) -> AsyncIterator[BaseEvent]:
    """Run the agent streaming responses using AG-UI protocol events.

Review comment (suggestion):
@DouweM not sure this is what you meant...

```py
if _utils.is_async_callable(on_complete):
    await on_complete(run)
 ```  
 
 >"None" is not awaitable
  "None" is incompatible with protocol "Awaitable[_T_co@Awaitable]"
    "__await__" is not present
    
  And 
  
  ```py
  else:
       await _utils.run_in_executor(on_complete, run)
  ```
  
Is saying it's unreachable.  

Code after:
        else:
            # `deps` not being a `StateHandler` is OK if there is no state.
            pass

        messages = _messages_from_ag_ui(run_input.messages)

        async with agent.iter(
            user_prompt=None,
            output_type=[output_type or agent.output_type, DeferredToolRequests],
            message_history=messages,
            model=model,
            deps=deps,
            model_settings=model_settings,
            usage_limits=usage_limits,
            usage=usage,
            infer_name=infer_name,
            toolsets=toolsets,
        ) as run:
            async for event in _agent_stream(run):
                yield encoder.encode(event)

        if on_complete is not None and run.result is not None:
            if _utils.is_async_callable(on_complete):
                await on_complete(run.result)
            else:
                await _utils.run_in_executor(on_complete, run.result)
    except _RunError as e:
        yield encoder.encode(
            RunErrorEvent(message=e.message, code=e.code),
        )
    except Exception as e:
        yield encoder.encode(
            RunErrorEvent(message=str(e)),
        )
        raise e
    else:
        yield encoder.encode(
            RunFinishedEvent(
                thread_id=run_input.thread_id,
                run_id=run_input.run_id,
            ),
        )


async def _agent_stream(run: AgentRun[AgentDepsT, Any]) -> AsyncIterator[BaseEvent]:
    """Run the agent streaming responses using AG-UI protocol events.

    Args:
        run: The agent run to process.

    Yields: