47 lines
1.2 KiB
Python
47 lines
1.2 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["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]
|
||
|
|
index = {item.get("id"): position for position, item in enumerate(merged) if item.get("id")}
|
||
|
|
for item in new:
|
||
|
|
run = dict(item)
|
||
|
|
run_id = run.get("id")
|
||
|
|
if run_id and run_id in index:
|
||
|
|
merged[index[run_id]] = {**merged[index[run_id]], **run}
|
||
|
|
else:
|
||
|
|
if run_id:
|
||
|
|
index[run_id] = len(merged)
|
||
|
|
merged.append(run)
|
||
|
|
return merged
|
||
|
|
|
||
|
|
|
||
|
|
class ChatBotState(BaseState):
|
||
|
|
subagent_runs: Annotated[list[SubAgentRunState], merge_subagent_runs]
|