61 lines
1.8 KiB
Python
61 lines
1.8 KiB
Python
from __future__ import annotations
|
|
|
|
from typing import Annotated, Literal, TypedDict
|
|
|
|
from yuxi.agents.state import BaseState
|
|
|
|
|
|
class SubAgentRunState(TypedDict, total=False):
|
|
id: str
|
|
subagent_type: str
|
|
subagent_name: str
|
|
child_thread_id: str
|
|
description: str
|
|
status: Literal["running", "completed", "failed"]
|
|
created_at: str
|
|
completed_at: str
|
|
result_preview: str
|
|
error: str | None
|
|
artifacts: list[str]
|
|
|
|
|
|
def merge_subagent_runs(
|
|
existing: list[SubAgentRunState] | None,
|
|
new: list[SubAgentRunState] | None,
|
|
) -> list[SubAgentRunState]:
|
|
if existing is None:
|
|
return list(new or [])
|
|
if new is None:
|
|
return existing
|
|
|
|
merged = [dict(item) for item in existing]
|
|
child_thread_index = {
|
|
item.get("child_thread_id"): position for position, item in enumerate(merged) if item.get("child_thread_id")
|
|
}
|
|
id_index = {item.get("id"): position for position, item in enumerate(merged) if item.get("id")}
|
|
for item in new:
|
|
run = dict(item)
|
|
child_thread_id = run.get("child_thread_id")
|
|
run_id = run.get("id")
|
|
position = None
|
|
if child_thread_id and child_thread_id in child_thread_index:
|
|
position = child_thread_index[child_thread_id]
|
|
elif run_id and run_id in id_index:
|
|
position = id_index[run_id]
|
|
|
|
if position is None:
|
|
position = len(merged)
|
|
merged.append(run)
|
|
else:
|
|
merged[position] = {**merged[position], **run}
|
|
|
|
if child_thread_id:
|
|
child_thread_index[child_thread_id] = position
|
|
if run_id:
|
|
id_index[run_id] = position
|
|
return merged
|
|
|
|
|
|
class ChatBotState(BaseState):
|
|
subagent_runs: Annotated[list[SubAgentRunState], merge_subagent_runs]
|