forked from CopilotKit/CopilotKit
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathlanggraph_agent.py
More file actions
840 lines (702 loc) · 31.1 KB
/
Copy pathlanggraph_agent.py
File metadata and controls
840 lines (702 loc) · 31.1 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
"""LangGraph agent for CopilotKit"""
import uuid
import json
import math
from typing import Optional, List, Callable, Any, cast, Union, TypedDict, Literal
from langgraph.graph.state import CompiledStateGraph
from typing_extensions import NotRequired
from langgraph.types import Command
try:
from langchain.load.dump import dumps as langchain_dumps
from langchain.schema import BaseMessage, SystemMessage
except ImportError:
# Langchain >= 1.0.0
from langchain_core.load import dumps as langchain_dumps
from langchain_core.messages import BaseMessage, SystemMessage
from langchain_core.runnables import RunnableConfig, ensure_config
def _serialize_state(state):
"""Recursively convert Pydantic BaseModel instances to dicts for serialization."""
try:
from pydantic import BaseModel as PydanticBaseModel
except ImportError:
return state
if isinstance(state, PydanticBaseModel):
return state.model_dump()
elif isinstance(state, dict):
return {k: _serialize_state(v) for k, v in state.items()}
elif isinstance(state, (list, tuple)):
return type(state)(_serialize_state(item) for item in state)
return state
from langchain_core.messages import HumanMessage
from partialjson.json_parser import JSONParser
from .types import Message, MetaEvent
from .utils import filter_by_schema_keys
from .langgraph import copilotkit_messages_to_langchain, langchain_messages_to_copilotkit
from .action import ActionDict
from .agent import Agent
from .logging import get_logger
logger = get_logger(__name__)
def _sanitize_for_json(obj):
"""Replace NaN and Infinity float values with None for valid JSON serialization."""
if isinstance(obj, float):
if math.isnan(obj) or math.isinf(obj):
return None
return obj
if isinstance(obj, dict):
return {k: _sanitize_for_json(v) for k, v in obj.items()}
if isinstance(obj, (list, tuple)):
return [_sanitize_for_json(item) for item in obj]
return obj
class CopilotKitConfig(TypedDict):
"""
CopilotKit config for LangGraphAgent
This is used for advanced cases where you want to customize how CopilotKit interacts with
LangGraph.
```python
# Function signatures:
def merge_state(
*,
state: dict,
messages: List[BaseMessage],
actions: List[Any],
agent_name: str
):
# ...implementation...
def convert_messages(messages: List[Message]):
# ...implementation...
```
Parameters
----------
merge_state : Callable
This function lets you customize how CopilotKit merges the agent state.
convert_messages : Callable
Use this function to customize how CopilotKit converts its messages to LangChain messages.`
"""
merge_state: NotRequired[Callable]
convert_messages: NotRequired[Callable]
def langgraph_default_merge_state( # pylint: disable=unused-argument
*,
state: dict,
messages: List[BaseMessage],
actions: List[Any],
agent_name: str
):
"""Default merge state for LangGraph"""
if len(messages) > 0 and isinstance(messages[0], SystemMessage):
# remove system message
messages = messages[1:]
existing_messages = state.get("messages", [])
existing_message_ids = {message.id for message in existing_messages}
new_messages = [message for message in messages if message.id not in existing_message_ids]
return {
**state,
"messages": new_messages,
"copilotkit": {
"actions": actions
}
}
class LangGraphAgent(Agent):
"""
LangGraphAgent lets you define your agent for use with CopilotKit.
To install, run:
```bash
pip install copilotkit
```
### Examples
Every agent must have the `name` and `graph` properties defined. An optional `description`
can also be provided. This is used when CopilotKit is dynamically routing requests to the
agent.
```python
from copilotkit import LangGraphAgent
LangGraphAgent(
name="email_agent",
description="This agent sends emails",
graph=graph,
)
```
If you have a custom LangGraph/LangChain config that you want to use with the agent, you can
pass it in as the `langgraph_config` parameter.
```python
LangGraphAgent(
...
langgraph_config=config,
)
```
Parameters
----------
name : str
The name of the agent.
graph : CompiledStateGraph
The LangGraph graph to use with the agent.
description : Optional[str]
The description of the agent.
langgraph_config : Optional[RunnableConfig]
The LangGraph/LangChain config to use with the agent.
copilotkit_config : Optional[CopilotKitConfig]
The CopilotKit config to use with the agent.
"""
def __init__(
self,
*,
name: str,
graph: Optional[CompiledStateGraph] = None,
description: Optional[str] = None,
langgraph_config: Union[Optional[RunnableConfig], dict] = None,
copilotkit_config: Optional[CopilotKitConfig] = None,
# deprecated - use langgraph_config instead
config: Union[Optional[RunnableConfig], dict] = None,
# deprecated - use graph instead
agent: Optional[CompiledStateGraph] = None,
# deprecated - use copilotkit_config instead
merge_state: Optional[Callable] = None,
):
if config is not None:
logger.warning("Warning: config is deprecated, use langgraph_config instead")
if agent is not None:
logger.warning("Warning: agent is deprecated, use graph instead")
if merge_state is not None:
logger.warning("Warning: merge_state is deprecated, use copilotkit_config instead")
if graph is None and agent is None:
raise ValueError("graph must be provided")
super().__init__(
name=name,
description=description,
)
self.merge_state = None
self.thread_state = {}
if copilotkit_config is not None:
self.merge_state = copilotkit_config.get("merge_state")
if not self.merge_state and merge_state is not None:
self.merge_state = merge_state
if not self.merge_state:
self.merge_state = langgraph_default_merge_state
self.convert_messages = (
copilotkit_config.get("convert_messages")
if copilotkit_config
else None
) or copilotkit_messages_to_langchain(use_function_call=False)
self.langgraph_config = langgraph_config or config
self.graph = cast(CompiledStateGraph, graph or agent)
self.active_interrupt_event = False
def execute( # pylint: disable=too-many-arguments
self,
*,
state: dict,
config: Optional[dict] = None,
messages: List[Message],
thread_id: str,
actions: Optional[List[ActionDict]] = None,
meta_events: Optional[List[MetaEvent]] = None,
**kwargs
):
node_name = kwargs.get("node_name")
return self._stream_events(
state=state,
config=config,
messages=messages,
actions=actions,
thread_id=thread_id,
node_name=node_name,
meta_events=meta_events
)
async def prepare_stream( # pylint: disable=too-many-arguments
self,
*,
state_input: Any,
agent_state: Any,
config: Optional[dict] = None,
messages: List[Message],
thread_id: str,
actions: Optional[List[ActionDict]] = None,
node_name: Optional[str] = None,
meta_events: Optional[List[MetaEvent]] = None,
):
active_interrupts = agent_state.tasks[0].interrupts if agent_state.tasks and agent_state.tasks[0].interrupts else None
state_input["messages"] = agent_state.values.get("messages", [])
current_graph_state = agent_state.values
langchain_messages = self.convert_messages(messages)
state = cast(Callable, self.merge_state)(
state=state_input,
messages=langchain_messages,
actions=actions,
agent_name=self.name
)
# Only update graph state with keys that merge_state explicitly produced,
# not keys that were simply passed through from state_input unchanged.
# This preserves graph-owned state keys that the frontend may have sent stale values for.
for key, value in state.items():
if key not in state_input or value is not state_input.get(key):
current_graph_state[key] = value
lg_interrupt_meta_event = next((ev for ev in (meta_events or []) if ev.get("name") == "LangGraphInterruptEvent"), None)
has_active_interrupts = active_interrupts is not None and len(active_interrupts) > 0
resume_input = None
# An active interrupt event that runs through messages. Use latest message as response
if has_active_interrupts and lg_interrupt_meta_event is None:
# state["messages"] only includes the messages we need to add at this point, tool call+result if applicable, and user text
resume_input = Command(resume=state["messages"])
if lg_interrupt_meta_event and "response" in lg_interrupt_meta_event:
resume_input = Command(resume=lg_interrupt_meta_event["response"])
mode = "continue" if thread_id and node_name != "__end__" and node_name is not None else "start"
thread_id = thread_id or str(uuid.uuid4())
config["configurable"]["thread_id"] = thread_id
if mode == "continue" and not has_active_interrupts:
await self.graph.aupdate_state(config, state, as_node=node_name)
initial_state = state if mode == "start" else None
# Use provided resume_input or fallback to initial_state
stream_input = resume_input if resume_input else initial_state
# Get the output and input schema keys the user has allowed for this graph
input_keys, output_keys, config_keys = self.get_schema_keys(config)
self.output_schema_keys = output_keys
self.input_schema_keys = input_keys
stream_input = self.filter_state_on_schema_keys(stream_input, 'input')
config["configurable"] = filter_by_schema_keys(config["configurable"], config_keys)
if has_active_interrupts and (not resume_input):
value = active_interrupts[0].value
return {
"stream": None,
"state": None,
"config": None,
"interrupt_event": self.get_interrupt_event(value),
}
return {
"stream": self.graph.astream_events(stream_input, config, version="v2"),
"state": current_graph_state,
"config": config
}
async def prepare_regenerate_stream( # pylint: disable=too-many-arguments
self,
*,
state: Any,
config: Optional[dict] = None,
actions: Optional[List[ActionDict]] = None,
message_checkpoint: HumanMessage
):
thread_id = config.get("configurable", {}).get("thread_id")
time_travel_checkpoint = await self.get_checkpoint_before_message(message_checkpoint.id, thread_id)
if time_travel_checkpoint is None:
return None
fork = await self.graph.aupdate_state(
time_travel_checkpoint.config,
time_travel_checkpoint.values,
as_node=time_travel_checkpoint.next[0] if time_travel_checkpoint.next else "__start__"
)
stream_input = cast(Callable, self.merge_state)(
state=time_travel_checkpoint.values,
messages=[message_checkpoint],
actions=actions,
agent_name=self.name
)
stream = self.graph.astream_events(stream_input, fork, version="v2")
return {
"stream": stream,
"state": state,
"config": config
}
async def _stream_events( # pylint: disable=too-many-locals
self,
*,
state: Any,
config: Optional[dict] = None,
messages: List[Message],
thread_id: str,
actions: Optional[List[ActionDict]] = None,
node_name: Optional[str] = None,
meta_events: Optional[List[MetaEvent]] = None,
):
default_config = ensure_config(cast(Any, self.langgraph_config.copy()) if self.langgraph_config else {}) # pylint: disable=line-too-long
config = {**default_config, **(self.graph.config or {}), **(config or {})}
config["configurable"] = {**config.get("configurable", {}), **(config["configurable"] or {})}
config["configurable"]["thread_id"] = thread_id
streaming_state_extractor = _StreamingStateExtractor([])
prev_node_name = None
emit_intermediate_state_until_end = None
should_exit = False
manually_emitted_state = None
thread_id = cast(Any, config)["configurable"]["thread_id"]
agent_state = await self.graph.aget_state(config)
prepared_stream_response = await self.prepare_stream(
state_input=state,
agent_state=agent_state,
config=config,
messages=messages,
actions=actions,
thread_id=thread_id,
node_name=node_name,
meta_events=meta_events
)
langchain_messages = self.convert_messages(messages)
non_system_messages = [msg for msg in langchain_messages if not isinstance(msg, SystemMessage)]
if len(agent_state.values.get("messages", [])) > len(non_system_messages):
# Find the last user message by working backwards from the last message
last_user_message = None
for i in range(len(langchain_messages) - 1, -1, -1):
if isinstance(langchain_messages[i], HumanMessage):
last_user_message = langchain_messages[i]
break
if last_user_message:
prepared_stream_response = await self.prepare_regenerate_stream(
state=state,
config=config,
message_checkpoint=last_user_message,
actions=actions,
)
state = prepared_stream_response["state"]
current_graph_state = prepared_stream_response["state"]
stream = prepared_stream_response["stream"]
config = prepared_stream_response["config"]
interrupt_event = prepared_stream_response.get('interrupt_event', None)
if interrupt_event:
yield interrupt_event
return
try:
async for event in stream:
current_node_name = event.get("name")
event_type = event.get("event")
run_id = event.get("run_id")
metadata = event.get("metadata", {})
interrupt_event = (
event["data"].get("chunk", {}).get("__interrupt__", None)
if (
isinstance(event.get("data"), dict) and
isinstance(event["data"].get("chunk"), dict)
)
else None
)
if interrupt_event:
value = interrupt_event[0].value
yield self.get_interrupt_event(value)
continue
should_exit = should_exit or (
event_type == "on_custom_event" and
event["name"] == "copilotkit_exit"
)
# OPTIMIZATION: Update local state from chain_end events to avoid checkpointer calls
if event_type == "on_chain_end" and isinstance(
event.get("data", {}).get("output"), dict
):
current_graph_state.update(event["data"]["output"])
emit_intermediate_state = metadata.get("copilotkit:emit-intermediate-state")
manually_emit_intermediate_state = (
event_type == "on_custom_event" and
event["name"] == "copilotkit_manually_emit_intermediate_state"
)
# we only want to update the node name under certain conditions
# since we don't need any internal node names to be sent to the frontend
if current_node_name in self.graph.nodes.keys():
node_name = current_node_name
# we don't have a node name yet, so we can't update the state
if node_name is None:
continue
exiting_node = node_name == current_node_name and event_type == "on_chain_end"
if exiting_node:
manually_emitted_state = None
if manually_emit_intermediate_state:
manually_emitted_state = _merge_emit_state(current_graph_state, cast(Any, event["data"]))
if isinstance(manually_emitted_state, dict):
current_graph_state.update(manually_emitted_state)
yield self._emit_state_sync_event(
thread_id=thread_id,
run_id=run_id,
node_name=node_name,
state=manually_emitted_state,
running=True,
active=True
) + "\n"
continue
if emit_intermediate_state and emit_intermediate_state_until_end is None:
emit_intermediate_state_until_end = node_name
if emit_intermediate_state and event_type == "on_chat_model_start":
# reset the streaming state extractor
streaming_state_extractor = _StreamingStateExtractor(emit_intermediate_state)
# OPTIMIZATION: Use locally maintained state instead of hitting checkpointer repeatedly
updated_state = manually_emitted_state or current_graph_state
if emit_intermediate_state and event_type == "on_chat_model_stream":
streaming_state_extractor.buffer_tool_calls(event)
if emit_intermediate_state_until_end is not None:
updated_state = {
**updated_state,
**streaming_state_extractor.extract_state()
}
if (not emit_intermediate_state and
current_node_name == emit_intermediate_state_until_end and
event_type == "on_chain_end"):
# stop emitting function call state
emit_intermediate_state_until_end = None
# we send state sync events when:
# a) the state has changed
# b) the node has changed
# c) the node is ending
if updated_state != state or prev_node_name != node_name or exiting_node:
state = updated_state
prev_node_name = node_name
current_graph_state.update(updated_state)
yield self._emit_state_sync_event(
thread_id=thread_id,
run_id=run_id,
node_name=node_name,
state=state,
running=True,
active=not exiting_node
) + "\n"
yield langchain_dumps(_sanitize_for_json(event)) + "\n"
except Exception as error:
# Emit error information through streaming protocol before terminating
# This preserves the semantic error details that would otherwise be lost
error_message = str(error)
error_type = type(error).__name__
# Extract additional error details for common error types
error_details = {
"message": error_message,
"type": error_type,
"agent_name": self.name,
}
# Add specific details for OpenAI errors
if hasattr(error, 'status_code'):
error_details["status_code"] = error.status_code
if hasattr(error, 'response') and hasattr(error.response, 'json'):
try:
error_details["response_data"] = error.response.json()
except:
pass
# Emit error events in both formats to support both LangGraph Platform and direct LangGraph modes
# Format for LangGraph Platform (remote-lg-action.ts)
yield langchain_dumps({
"event": "error",
"data": {
"message": f"{error_type}: {error_message}",
"error_details": error_details,
"thread_id": thread_id,
"agent_name": self.name,
"node_name": node_name or "unknown"
}
}) + "\n"
# Format for direct LangGraph mode (event-source.ts)
yield langchain_dumps({
"event": "on_copilotkit_error",
"data": {
"error": error_details,
"thread_id": thread_id,
"agent_name": self.name,
"node_name": node_name or "unknown"
}
}) + "\n"
# Re-raise the exception to maintain normal error handling flow
raise
state = await self.graph.aget_state(config)
tasks = state.tasks
interrupts = tasks[0].interrupts if tasks and len(tasks) > 0 else None
if interrupts:
# node_name is already set earlier from the interrupt origin
pass
elif "writes" in state.metadata and state.metadata["writes"]:
node_name = list(state.metadata["writes"].keys())[0]
elif hasattr(state, "next") and state.next and state.next[0]:
node_name = state.next[0]
else:
node_name = "__end__"
is_end_node = state.next == () and not interrupts
yield self._emit_state_sync_event(
thread_id=thread_id,
run_id=run_id,
node_name=cast(str, node_name) if not is_end_node else "__end__",
state=state.values,
running=not should_exit,
# at this point, the node is ending so we set active to false
active=False,
# sync messages at the end of the run
include_messages=True
) + "\n"
def _emit_state_sync_event(
self,
*,
thread_id: str,
run_id: str,
node_name: str,
state: dict,
running: bool,
active: bool,
include_messages: bool = False
):
# First handle messages as before
if not include_messages:
state = {
k: v for k, v in state.items() if k != "messages"
}
else:
state = {
**state,
"messages": langchain_messages_to_copilotkit(state.get("messages", []))
}
# Filter by schema keys if available
state = self.filter_state_on_schema_keys(state, 'output')
# Convert Pydantic BaseModel instances to dicts for serialization
state = _serialize_state(state)
return langchain_dumps({
"event": "on_copilotkit_state_sync",
"thread_id": thread_id,
"run_id": run_id,
"agent_name": self.name,
"node_name": node_name,
"active": active,
"state": _sanitize_for_json(state),
"running": running,
"role": "assistant"
})
async def get_state(
self,
*,
thread_id: str,
):
if not thread_id:
return {
"threadId": "",
"threadExists": False,
"state": {},
"messages": []
}
config = ensure_config(cast(Any, self.langgraph_config.copy()) if self.langgraph_config else {}) # pylint: disable=line-too-long
config["configurable"] = config.get("configurable", {})
config["configurable"]["thread_id"] = thread_id
if self.thread_state.get(thread_id, None) is None:
self.thread_state[thread_id] = {**(await self.graph.aget_state(config)).values}
state = self.thread_state[thread_id]
if state == {}:
return {
"threadId": thread_id or "",
"threadExists": False,
"state": {},
"messages": []
}
messages = langchain_messages_to_copilotkit(state.get("messages", []))
state_copy = state.copy()
state_copy.pop("messages", None)
# Convert Pydantic BaseModel instances to dicts for serialization
state_copy = _serialize_state(state_copy)
return {
"threadId": thread_id,
"threadExists": True,
"state": _sanitize_for_json(state_copy),
"messages": messages
}
def dict_repr(self):
super_repr = super().dict_repr()
return {
**super_repr,
'type': 'langgraph'
}
def get_schema_keys(self, config):
CONSTANT_KEYS = ['copilotkit', 'messages']
CONSTANT_CONFIG_KEYS = ['checkpoint_id', 'checkpoint_ns', 'thread_id']
try:
input_schema = self.graph.get_input_jsonschema(config)
output_schema = self.graph.get_output_jsonschema(config)
input_schema_keys = list(input_schema["properties"].keys())
output_schema_keys = list(output_schema["properties"].keys())
try:
schema_dict = self.graph.config_schema().schema()
configurable_schema = schema_dict["$defs"]["Configurable"]
config_schema_keys = list(configurable_schema["properties"].keys())
# If only constant keys are present, it means no schema was passed, we allow everything
if set(config_schema_keys) == set(CONSTANT_CONFIG_KEYS):
config_schema_keys = None
except:
config_schema_keys = None
# We add "copilotkit" and "messages" as they are always sent and received.
for key in CONSTANT_KEYS:
if key not in input_schema_keys:
input_schema_keys.append(key)
if key not in output_schema_keys:
output_schema_keys.append(key)
return input_schema_keys, output_schema_keys, config_schema_keys
except Exception:
return None
def filter_state_on_schema_keys(self, state, schema_type: Literal["input", "output"]):
try:
schema_keys_name = f"{schema_type}_schema_keys"
if hasattr(self, schema_keys_name) and getattr(self, schema_keys_name):
return filter_by_schema_keys(state, getattr(self, schema_keys_name))
except Exception:
pass
return state
def get_interrupt_event(self, value):
if not isinstance(value, str) and "__copilotkit_interrupt_value__" in value:
ev_value = value["__copilotkit_interrupt_value__"]
return langchain_dumps({
"event": "on_copilotkit_interrupt",
"data": { "value": ev_value if isinstance(ev_value, str) else json.dumps(ev_value), "messages": langchain_messages_to_copilotkit(value["__copilotkit_messages__"]) }
}) + "\n"
else:
return langchain_dumps({
"event": "on_interrupt",
"value": value if isinstance(value, str) else json.dumps(value)
}) + "\n"
async def get_checkpoint_before_message(self, message_id: str, thread_id: str):
if not thread_id:
raise ValueError("Missing thread_id in config")
history_list = []
async for snapshot in self.graph.aget_state_history({"configurable": {"thread_id": thread_id}}):
history_list.append(snapshot)
history_list.reverse()
for idx, snapshot in enumerate(history_list):
messages = snapshot.values.get("messages", [])
if any(getattr(m, "id", None) == message_id for m in messages):
if idx == 0:
# No snapshot before this
# Return synthetic "empty before" version
empty_snapshot = snapshot
empty_snapshot.values["messages"] = []
return empty_snapshot
return history_list[idx - 1] # return one snapshot *before* the one that includes the message
raise ValueError("Message ID not found in history")
def _merge_emit_state(current_state: dict, emitted_state: Any) -> dict:
"""Merge emitted state on top of current graph state instead of replacing it."""
if isinstance(emitted_state, dict):
return {**current_state, **emitted_state}
return cast(Any, emitted_state)
class _StreamingStateExtractor:
def __init__(self, emit_intermediate_state: List[dict]):
self.emit_intermediate_state = emit_intermediate_state
self.tool_call_buffer = {}
self.current_tool_call = None
self.previously_parsable_state = {}
def buffer_tool_calls(self, event: Any):
"""Buffer the tool calls"""
if len(event["data"]["chunk"].tool_call_chunks) > 0:
chunk = event["data"]["chunk"].tool_call_chunks[0]
if chunk["name"] is not None:
self.current_tool_call = chunk["name"]
self.tool_call_buffer[self.current_tool_call] = chunk["args"]
elif self.current_tool_call is not None:
self.tool_call_buffer[self.current_tool_call] = (
self.tool_call_buffer[self.current_tool_call] + chunk["args"]
)
def get_emit_state_config(self, current_tool_name):
"""Get the emit state config"""
for config in self.emit_intermediate_state:
state_key = config.get("state_key")
tool = config.get("tool")
tool_argument = config.get("tool_argument")
if current_tool_name == tool:
return (tool_argument, state_key)
return (None, None)
def extract_state(self):
"""Extract the streaming state"""
parser = JSONParser()
state = {}
for key, value in self.tool_call_buffer.items():
argument_name, state_key = self.get_emit_state_config(key)
if state_key is None:
continue
try:
parsed_value = parser.parse(value)
except Exception as _exc: # pylint: disable=broad-except
if key in self.previously_parsable_state:
parsed_value = self.previously_parsable_state[key]
else:
continue
self.previously_parsable_state[key] = parsed_value
if argument_name is None:
state[state_key] = parsed_value
else:
state[state_key] = parsed_value.get(argument_name)
return state