Update chatbot and RAG assistant to use `StateGraph` in the backend #305

andrewnguonly · 2024-04-16T05:44:38Z

Summary

As a demonstration, get_chatbot_executor() and get_retrieval_executor() are updated to use StateGraph. Moving forward, when using the RAG assistant (assistantType === "chat_retrieval"), the POST /runs and /runs/stream API must be called with the follow request body:

{
    // input is an object instead of a list
    "input": {
        // messages key must be present
        "messages": [
            {
                "content": "hello!", // content key must be present
                "role": "human",     // role key must be present
                ...
            }
        ],
        ...
    },
    ...
}

All other assistant types require the existing request body format (e.g. "input": [{...}]).

Implementation

The MessageGraph in get_chatbot_executor() is migrated to StateGraph.
The MessageGraph in get_retrieval_executor() is migrated to StateGraph. The graph now accepts a TypedDict for the state. The interfaces for the corresponding nodes are updated accordingly.
API GET /threads/<tid>/state (storage layer) is updated to retrieve the graph state based on the assistant type.
Frontend is updated to call the API POST /runs/stream with the correct request body format based on the assistant type.

To Do

Update API documentation (if necessary, details TBD).
Remove TODOs in code. Update langchain_core dependency and implement new API. See core: forward config params to default langchain#20402 for details.
Add new keys to the get_retrieval_executor() StateGraph state.
The code that's causing the broken unit tests will be removed when the TODOs are resolved.
Retest everything once all the To Do's are resolved.

… type.

andrewnguonly · 2024-04-16T06:03:55Z

frontend/src/types.ts

@@ -12,6 +12,7 @@ export interface MessageDocument {
 export interface Message {
  id: string;
  type: string;
+  role?: string; // for chat_retrieval bot


I couldn't decide what to do about this...

Calling POST /runs/stream requires the role field in the message. Otherwise, the request body will not be serialized appropriately in the backend and an error is raised langgraph.graph.message.add_messages (state reducer function).

It seems more appropriate for the client (frontend) to handle this instead of the API (backend), but I'm not sure.

andrewnguonly · 2024-04-17T23:23:01Z

backend/app/checkpoint.py

@@ -19,8 +25,13 @@ def loads(value: bytes) -> Checkpoint:


 class PostgresCheckpoint(BaseCheckpointSaver):
-    class Config:
-        arbitrary_types_allowed = True
+    def __init__(


This change is required since the introduction of the serde API in BaseCheckpointSaver.

nfcampos · 2024-04-18T15:27:16Z

backend/app/agent.py

@@ -244,7 +246,10 @@ def __init__(
        llm=ConfigurableField(id="llm_type", name="LLM Type"),
        system_message=ConfigurableField(id="system_message", name="Instructions"),
    )
-    .with_types(input_type=Sequence[AnyMessage], output_type=Sequence[AnyMessage])
+    .with_types(
+        input_type=Union[Sequence[AnyMessage], Dict[str, Any]],


these types don't seem right, it either accepts list of messages or dict, not both?

Updated types to Messages/Sequence[AnyMessage] for input/output respectively.

nfcampos · 2024-04-18T15:29:03Z

backend/app/stream.py

+            state_chunk_msgs: Union[Sequence[AnyMessage], Dict[str, Any]] = event[
+                "data"
+            ]["chunk"]
+            if isinstance(state_chunk_msgs, Dict):


this should be lowercase dict

Good catch. I'll update all other instances of this error.

nfcampos · 2024-04-18T15:29:34Z

backend/app/chatbot.py

@@ -15,7 +18,7 @@ def _get_messages(messages):

    chatbot = _get_messages | llm

-    workflow = MessageGraph()
+    workflow = StateGraph(Annotated[Sequence[BaseMessage], add_messages])


we should be using list?

It didn't make a difference after this PR was merged: langchain-ai/langgraph#321. I'll update both instances to use List.

This is where the API feels ambiguous with respect to the underlying functionality. If there's a "preferred" or required type that should be used, then an end user isn't necessarily aware of it. Something to think about for later.

nfcampos · 2024-04-18T15:29:51Z

backend/app/retrieval.py

@@ -39,6 +42,10 @@ def get_retrieval_executor(
    system_message: str,
    checkpoint: BaseCheckpointSaver,
 ):
+    class AgentState(TypedDict):
+        messages: Annotated[Sequence[BaseMessage], add_messages]


we should be using list?

nfcampos · 2024-04-18T15:30:17Z

frontend/src/App.tsx

+          // Each message must contain a `role` field.
+          input = {
+            messages: input.map((msg: Message) => {
+              msg.role = "human";


The following error is raised:

opengpts-backend | Traceback (most recent call last): opengpts-backend | File "/backend/app/stream.py", line 63, in to_sse opengpts-backend | async for chunk in messages_stream: opengpts-backend | File "/backend/app/stream.py", line 23, in astream_state opengpts-backend | async for event in app.astream_events( opengpts-backend | File "/usr/local/lib/python3.11/site-packages/langchain_core/runnables/base.py", line 4711, in astream_events opengpts-backend | async for item in self.bound.astream_events( opengpts-backend | File "/usr/local/lib/python3.11/site-packages/langchain_core/runnables/base.py", line 1137, in astream_events opengpts-backend | async for log in _astream_log_implementation( # type: ignore[misc] opengpts-backend | File "/usr/local/lib/python3.11/site-packages/langchain_core/tracers/log_stream.py", line 616, in _astream_log_implementation opengpts-backend | await task opengpts-backend | File "/usr/local/lib/python3.11/site-packages/langchain_core/tracers/log_stream.py", line 570, in consume_astream opengpts-backend | async for chunk in runnable.astream(input, config, **kwargs): opengpts-backend | File "/usr/local/lib/python3.11/site-packages/langchain_core/runnables/configurable.py", line 221, in astream opengpts-backend | async for chunk in runnable.astream(input, config, **kwargs): opengpts-backend | File "/usr/local/lib/python3.11/site-packages/langchain_core/runnables/base.py", line 4698, in astream opengpts-backend | async for item in self.bound.astream( opengpts-backend | File "/usr/local/lib/python3.11/site-packages/langchain_core/runnables/configurable.py", line 221, in astream opengpts-backend | async for chunk in runnable.astream(input, config, **kwargs): opengpts-backend | File "/usr/local/lib/python3.11/site-packages/langchain_core/runnables/base.py", line 4698, in astream opengpts-backend | async for item in self.bound.astream( opengpts-backend | File "/usr/local/lib/python3.11/site-packages/langgraph/pregel/__init__.py", line 924, in astream opengpts-backend | _apply_writes(checkpoint, channels, pending_writes) opengpts-backend | File "/usr/local/lib/python3.11/site-packages/langgraph/pregel/__init__.py", line 1170, in _apply_writes opengpts-backend | channels[chan].update(vals) opengpts-backend | File "/usr/local/lib/python3.11/site-packages/langgraph/channels/binop.py", line 66, in update opengpts-backend | self.value = self.operator(self.value, value) opengpts-backend | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ opengpts-backend | File "/usr/local/lib/python3.11/site-packages/langgraph/graph/message.py", line 24, in add_messages opengpts-backend | right = [message_chunk_to_message(m) for m in convert_to_messages(right)] opengpts-backend | ^^^^^^^^^^^^^^^^^^^^^^^^^^ opengpts-backend | File "/usr/local/lib/python3.11/site-packages/langchain_core/messages/utils.py", line 234, in convert_to_messages opengpts-backend | return [_convert_to_message(m) for m in messages] opengpts-backend | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ opengpts-backend | File "/usr/local/lib/python3.11/site-packages/langchain_core/messages/utils.py", line 234, in <listcomp> opengpts-backend | return [_convert_to_message(m) for m in messages] opengpts-backend | ^^^^^^^^^^^^^^^^^^^^^^ opengpts-backend | File "/usr/local/lib/python3.11/site-packages/langchain_core/messages/utils.py", line 211, in _convert_to_message opengpts-backend | raise ValueError( opengpts-backend | ValueError: Message dict must contain 'role' and 'content' keys, got {'content': 'hello', 'additional_kwargs': {}, 'type': 'human', 'example': False, 'id': 'human-0.38348279892186277'}

FastAPI is not serializing the dict to an AnyMessage type (which contains role).

andrewnguonly · 2024-04-18T18:59:53Z

Thanks for the assist @nfcampos 😄

andrewnguonly added 9 commits April 11, 2024 17:41

Migrate MessageGraph to StateGraph in get_chatbot_executor().

0fe363d

Migrate MessageGraph in StateGraph in get_retrieval_executor().

0229a51

Merge branch 'main' into migrate-state-graph

ca91592

Support TypedDict agent state in retrieval chat bot.

2db5e0c

Update frontend to pass list of messages or object based on assistant…

9ac9364

… type.

Restore type field to Message and make role field optional.

40e9030

Remove commented out code.

e210dd5

Remove commented out code. Fix lint errors in frontend code.

024c224

Merge main branch. Resolve conflicts.

6966252

andrewnguonly requested a review from nfcampos April 16, 2024 05:57

andrewnguonly commented Apr 16, 2024

View reviewed changes

andrewnguonly added 2 commits April 16, 2024 12:25

Update README and API.md

6c1f922

Add msg_count field to retrieval bot state.

76e366f

ptgoetz added enhancement New feature or request documentation Improvements or additions to documentation backend Changes to the backend service labels Apr 16, 2024

andrewnguonly added 3 commits April 17, 2024 14:10

Update langchain-core. Remove TODOs. Refactor get_thread_state().

d4f6993

Fix broken unit test.

86c7e6c

Update langgraph. Fix broken unit test.

c154598

andrewnguonly commented Apr 17, 2024

View reviewed changes

Fix frontend lint errors.

595fdd7

andrewnguonly changed the title ~~Draft: Update chatbot and RAG assistant to use StateGraph in the backend~~ Update chatbot and RAG assistant to use StateGraph in the backend Apr 17, 2024

andrewnguonly added 2 commits April 17, 2024 21:09

Merge main branch.

c60978b

Fix lint errors in backend.

824d8ce

nfcampos reviewed Apr 18, 2024

View reviewed changes

Change Dict to dict. Change Sequence to List.

2ab65f8

andrewnguonly and others added 6 commits April 18, 2024 11:35

Change Dict to dict.

40d92f7

Fix up retrieval bot

e262abd

Fix null messages

dde4ec1

Remove unneeded role

0271758

Fix up chatbot

b1c7702

Update input/output types.

deab0ed

nfcampos merged commit f0c25df into langchain-ai:main Apr 18, 2024
6 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Update chatbot and RAG assistant to use `StateGraph` in the backend #305

Update chatbot and RAG assistant to use `StateGraph` in the backend #305

andrewnguonly commented Apr 16, 2024 •

edited

Loading

andrewnguonly Apr 16, 2024

andrewnguonly Apr 17, 2024

nfcampos Apr 18, 2024

andrewnguonly Apr 18, 2024

nfcampos Apr 18, 2024

andrewnguonly Apr 18, 2024

nfcampos Apr 18, 2024

andrewnguonly Apr 18, 2024

nfcampos Apr 18, 2024

nfcampos Apr 18, 2024

andrewnguonly Apr 18, 2024 •

edited

Loading

andrewnguonly commented Apr 18, 2024

Update chatbot and RAG assistant to use StateGraph in the backend #305

Update chatbot and RAG assistant to use StateGraph in the backend #305

Conversation

andrewnguonly commented Apr 16, 2024 • edited Loading

Summary

Implementation

To Do

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

andrewnguonly Apr 18, 2024 • edited Loading

Choose a reason for hiding this comment

andrewnguonly commented Apr 18, 2024

Update chatbot and RAG assistant to use `StateGraph` in the backend #305

Update chatbot and RAG assistant to use `StateGraph` in the backend #305

andrewnguonly commented Apr 16, 2024 •

edited

Loading

andrewnguonly Apr 18, 2024 •

edited

Loading