From e9deeab37f132fcfda688d13224d2ff004c080ee Mon Sep 17 00:00:00 2001 From: Eugene Yurtsev Date: Mon, 7 Aug 2023 10:32:53 -0400 Subject: [PATCH] x --- .../langchain/automaton/well_known_states.py | 16 ++++++++++++---- 1 file changed, 12 insertions(+), 4 deletions(-) diff --git a/libs/langchain/langchain/automaton/well_known_states.py b/libs/langchain/langchain/automaton/well_known_states.py index a1a9d54461a..d8fbab1b1b9 100644 --- a/libs/langchain/langchain/automaton/well_known_states.py +++ b/libs/langchain/langchain/automaton/well_known_states.py @@ -38,16 +38,24 @@ class LLMProgram(State): memory.add_message(message) if infer_message_type(message) == MessageType.AI_INVOKE: - data = result["data"] + function_call = result["function_call"] function_message = FunctionMessage( - name=data["function_call"]["name"], - content=data["function_call"]["result"], + name=function_call["name"], + content=function_call["result"], ) memory.add_message(function_message) + routing_message = function_message + else: + routing_message = message # What information should the state return in this case. # Does it matter, folks can use it or not... - return {"id": "llm_program", "data": result} + return { + "id": "llm_program", + "data": { + "message": routing_message, # Last message + }, + } @dataclasses.dataclass