chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,346 @@
|
||||
"""
|
||||
CopilotKit Run Loop
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import contextvars
|
||||
import json
|
||||
import traceback
|
||||
from typing import Callable
|
||||
from pydantic import BaseModel
|
||||
from typing_extensions import Any, Dict, Optional, List, TypedDict, cast
|
||||
from partialjson.json_parser import JSONParser as PartialJSONParser
|
||||
|
||||
from .protocol import (
|
||||
RuntimeEvent,
|
||||
RuntimeEventTypes,
|
||||
RuntimeMetaEventName,
|
||||
emit_runtime_event,
|
||||
emit_runtime_events,
|
||||
agent_state_message,
|
||||
AgentStateMessage,
|
||||
PredictStateConfig,
|
||||
RuntimeProtocolEvent,
|
||||
)
|
||||
|
||||
|
||||
async def yield_control():
|
||||
"""
|
||||
Yield control to the event loop.
|
||||
"""
|
||||
loop = asyncio.get_running_loop()
|
||||
future = loop.create_future()
|
||||
loop.call_soon(future.set_result, None)
|
||||
await future
|
||||
|
||||
|
||||
class CopilotKitRunExecution(TypedDict):
|
||||
"""
|
||||
CopilotKit Run Execution
|
||||
"""
|
||||
|
||||
thread_id: str
|
||||
agent_name: str
|
||||
run_id: str
|
||||
should_exit: bool
|
||||
node_name: str
|
||||
is_finished: bool
|
||||
predict_state_configuration: Dict[str, PredictStateConfig]
|
||||
predicted_state: Dict[str, Any]
|
||||
argument_buffer: str
|
||||
current_tool_call: Optional[str]
|
||||
state: Dict[str, Any]
|
||||
|
||||
|
||||
_CONTEXT_QUEUE = contextvars.ContextVar("queue", default=None)
|
||||
_CONTEXT_EXECUTION = contextvars.ContextVar("execution", default=None)
|
||||
|
||||
|
||||
def get_context_queue() -> asyncio.Queue:
|
||||
"""
|
||||
Retrieve the queue from this task's context.
|
||||
"""
|
||||
q = _CONTEXT_QUEUE.get()
|
||||
if q is None:
|
||||
raise RuntimeError("No context queue is set!")
|
||||
return q
|
||||
|
||||
|
||||
def set_context_queue(q: asyncio.Queue) -> contextvars.Token:
|
||||
"""
|
||||
Set the queue in this task's context.
|
||||
"""
|
||||
token = _CONTEXT_QUEUE.set(cast(Any, q))
|
||||
return token
|
||||
|
||||
|
||||
def reset_context_queue(token: contextvars.Token):
|
||||
"""
|
||||
Reset the queue in this task's context.
|
||||
"""
|
||||
_CONTEXT_QUEUE.reset(token)
|
||||
|
||||
|
||||
def get_context_execution() -> CopilotKitRunExecution:
|
||||
"""
|
||||
Get the execution from this task's context.
|
||||
"""
|
||||
return cast(CopilotKitRunExecution, _CONTEXT_EXECUTION.get())
|
||||
|
||||
|
||||
def set_context_execution(execution: CopilotKitRunExecution) -> contextvars.Token:
|
||||
"""
|
||||
Set the execution in this task's context.
|
||||
"""
|
||||
token = _CONTEXT_EXECUTION.set(cast(Any, execution))
|
||||
return token
|
||||
|
||||
|
||||
def reset_context_execution(token: contextvars.Token):
|
||||
"""
|
||||
Reset the execution in this task's context.
|
||||
"""
|
||||
_CONTEXT_EXECUTION.reset(token)
|
||||
|
||||
|
||||
async def queue_put(*events: RuntimeEvent, priority: bool = False):
|
||||
"""
|
||||
Put an event in the queue.
|
||||
"""
|
||||
if not priority:
|
||||
# yield control so that priority events can be processed first
|
||||
await yield_control()
|
||||
|
||||
q = get_context_queue()
|
||||
for event in events:
|
||||
await q.put(event)
|
||||
|
||||
# yield control so that the reader can process the event
|
||||
await yield_control()
|
||||
|
||||
|
||||
def _to_dict_if_pydantic(obj):
|
||||
if isinstance(obj, BaseModel):
|
||||
return obj.model_dump()
|
||||
return obj
|
||||
|
||||
|
||||
def _filter_state(
|
||||
*, state: Dict[str, Any], exclude_keys: Optional[List[str]] = None
|
||||
) -> Dict[str, Any]:
|
||||
"""Filter out messages and id from the state"""
|
||||
state = _to_dict_if_pydantic(state)
|
||||
exclude_keys = exclude_keys or ["messages", "id"]
|
||||
return {k: v for k, v in state.items() if k not in exclude_keys}
|
||||
|
||||
|
||||
async def copilotkit_run(fn: Callable, *, execution: CopilotKitRunExecution):
|
||||
"""
|
||||
Run a task with a local queue.
|
||||
"""
|
||||
local_queue = asyncio.Queue()
|
||||
token_queue = set_context_queue(local_queue)
|
||||
token_execution = set_context_execution(execution)
|
||||
|
||||
task = asyncio.create_task(fn())
|
||||
try:
|
||||
while True:
|
||||
event = await local_queue.get()
|
||||
local_queue.task_done()
|
||||
|
||||
json_lines = handle_runtime_event(event=event, execution=execution)
|
||||
|
||||
if json_lines is not None:
|
||||
yield json_lines
|
||||
|
||||
if execution["is_finished"]:
|
||||
break
|
||||
|
||||
# return control to the containing run loop to send events
|
||||
await yield_control()
|
||||
|
||||
await task
|
||||
|
||||
finally:
|
||||
reset_context_queue(token_queue)
|
||||
reset_context_execution(token_execution)
|
||||
|
||||
|
||||
def handle_runtime_event(
|
||||
*, event: RuntimeEvent, execution: CopilotKitRunExecution
|
||||
) -> Optional[str]:
|
||||
"""
|
||||
Handle a runtime event.
|
||||
"""
|
||||
|
||||
if event["type"] in [
|
||||
RuntimeEventTypes.TEXT_MESSAGE_START,
|
||||
RuntimeEventTypes.TEXT_MESSAGE_CONTENT,
|
||||
RuntimeEventTypes.TEXT_MESSAGE_END,
|
||||
RuntimeEventTypes.ACTION_EXECUTION_START,
|
||||
RuntimeEventTypes.ACTION_EXECUTION_ARGS,
|
||||
RuntimeEventTypes.ACTION_EXECUTION_END,
|
||||
RuntimeEventTypes.ACTION_EXECUTION_RESULT,
|
||||
RuntimeEventTypes.AGENT_STATE_MESSAGE,
|
||||
]:
|
||||
events: List[RuntimeProtocolEvent] = [cast(RuntimeProtocolEvent, event)]
|
||||
if event["type"] in [
|
||||
RuntimeEventTypes.ACTION_EXECUTION_START,
|
||||
RuntimeEventTypes.ACTION_EXECUTION_ARGS,
|
||||
]:
|
||||
message = predict_state(
|
||||
thread_id=execution["thread_id"],
|
||||
agent_name=execution["agent_name"],
|
||||
run_id=execution["run_id"],
|
||||
event=event,
|
||||
execution=execution,
|
||||
)
|
||||
if message is not None:
|
||||
events.append(message)
|
||||
return emit_runtime_events(*events)
|
||||
|
||||
if event["type"] == RuntimeEventTypes.META_EVENT:
|
||||
if event["name"] == RuntimeMetaEventName.PREDICT_STATE:
|
||||
execution["predict_state_configuration"] = event["value"]
|
||||
return None
|
||||
if event["name"] == RuntimeMetaEventName.EXIT:
|
||||
execution["should_exit"] = event["value"]
|
||||
return None
|
||||
return None
|
||||
|
||||
if event["type"] == RuntimeEventTypes.RUN_STARTED:
|
||||
execution["state"] = event["state"]
|
||||
return None
|
||||
|
||||
if event["type"] == RuntimeEventTypes.NODE_STARTED:
|
||||
execution["node_name"] = event["node_name"]
|
||||
execution["state"] = event["state"]
|
||||
|
||||
return emit_runtime_event(
|
||||
agent_state_message(
|
||||
thread_id=execution["thread_id"],
|
||||
agent_name=execution["agent_name"],
|
||||
node_name=execution["node_name"],
|
||||
run_id=execution["run_id"],
|
||||
active=True,
|
||||
role="assistant",
|
||||
state=json.dumps(_filter_state(state=execution["state"])),
|
||||
running=True,
|
||||
)
|
||||
)
|
||||
|
||||
if event["type"] == RuntimeEventTypes.NODE_FINISHED:
|
||||
# reset the predict state configuration at the end of the method execution
|
||||
execution["predict_state_configuration"] = {}
|
||||
execution["current_tool_call"] = None
|
||||
execution["argument_buffer"] = ""
|
||||
execution["predicted_state"] = {}
|
||||
execution["state"] = event["state"]
|
||||
|
||||
return emit_runtime_event(
|
||||
agent_state_message(
|
||||
thread_id=execution["thread_id"],
|
||||
agent_name=execution["agent_name"],
|
||||
node_name=execution["node_name"],
|
||||
run_id=execution["run_id"],
|
||||
active=False,
|
||||
role="assistant",
|
||||
state=json.dumps(_filter_state(state=execution["state"])),
|
||||
running=True,
|
||||
)
|
||||
)
|
||||
|
||||
if event["type"] == RuntimeEventTypes.RUN_FINISHED:
|
||||
execution["is_finished"] = True
|
||||
return None
|
||||
|
||||
if event["type"] == RuntimeEventTypes.RUN_ERROR:
|
||||
print("Flow execution error", flush=True)
|
||||
error_info = event["error"]
|
||||
|
||||
if isinstance(error_info, Exception):
|
||||
# If it's an exception, print the traceback
|
||||
print("Exception occurred:", flush=True)
|
||||
print(
|
||||
"".join(
|
||||
traceback.format_exception(
|
||||
None, error_info, error_info.__traceback__
|
||||
)
|
||||
),
|
||||
flush=True,
|
||||
)
|
||||
else:
|
||||
# Otherwise, assume it's a string and print it
|
||||
print(error_info, flush=True)
|
||||
|
||||
execution["is_finished"] = True
|
||||
return None
|
||||
|
||||
|
||||
def predict_state(
|
||||
*,
|
||||
thread_id: str,
|
||||
agent_name: str,
|
||||
run_id: str,
|
||||
event: Any,
|
||||
execution: CopilotKitRunExecution,
|
||||
) -> Optional[AgentStateMessage]:
|
||||
"""Predict the state"""
|
||||
|
||||
if event["type"] == RuntimeEventTypes.ACTION_EXECUTION_START:
|
||||
execution["current_tool_call"] = event["actionName"]
|
||||
execution["argument_buffer"] = ""
|
||||
elif event["type"] == RuntimeEventTypes.ACTION_EXECUTION_ARGS:
|
||||
execution["argument_buffer"] += event["args"]
|
||||
|
||||
tool_names = [
|
||||
config.get("tool_name")
|
||||
for config in execution["predict_state_configuration"].values()
|
||||
]
|
||||
|
||||
if execution["current_tool_call"] not in tool_names:
|
||||
return None
|
||||
|
||||
current_arguments = {}
|
||||
try:
|
||||
current_arguments = PartialJSONParser().parse(execution["argument_buffer"])
|
||||
except: # pylint: disable=bare-except
|
||||
return None
|
||||
|
||||
emit_update = False
|
||||
for k, v in execution["predict_state_configuration"].items():
|
||||
if v["tool_name"] == execution["current_tool_call"]:
|
||||
tool_argument = v.get("tool_argument")
|
||||
if tool_argument is not None:
|
||||
argument_value = current_arguments.get(tool_argument)
|
||||
if argument_value is not None:
|
||||
execution["predicted_state"][k] = argument_value
|
||||
emit_update = True
|
||||
else:
|
||||
execution["predicted_state"][k] = current_arguments
|
||||
emit_update = True
|
||||
|
||||
if emit_update:
|
||||
return agent_state_message(
|
||||
thread_id=thread_id,
|
||||
agent_name=agent_name,
|
||||
node_name=execution["node_name"],
|
||||
run_id=run_id,
|
||||
active=True,
|
||||
role="assistant",
|
||||
state=json.dumps(
|
||||
_filter_state(
|
||||
state={
|
||||
**(
|
||||
execution["state"].model_dump()
|
||||
if isinstance(execution["state"], BaseModel)
|
||||
else execution["state"]
|
||||
),
|
||||
**execution["predicted_state"],
|
||||
}
|
||||
)
|
||||
),
|
||||
running=True,
|
||||
)
|
||||
|
||||
return None
|
||||
Reference in New Issue
Block a user