Structured Output JSON Not Compatible with Expected Message Format in Reflexion Agent

[braisdev]

Checked other resources

• I added a very descriptive title to this question.
• I searched the LangChain documentation with the integrated search.
• I used the GitHub search to find a similar question and didn't find it.

Commit to Help

• I commit to help with one of those options 👆

Example Code

from dotenv import load_dotenv
from langgraph.graph import MessageGraph
from chains import first_responder
from typing import List
from langchain_core.pydantic_v1 import BaseModel, Field


class Reflection(BaseModel):
    """Your reflection on the initial answer."""
    missing: str = Field(description="Critique of what is missing.")
    superfluous: str = Field(description="Critique of what is superfluous.")

    class Config:
        extra = "forbid"  # Disallow additional properties


class AnswerQuestion(BaseModel):
    """Answer the question."""

    answer: str = Field(description="~250 word detailed answer to the question.")
    reflection: Reflection
    search_queries: List[str] = Field(
        description="1-3 search queries for researching improvements to address the critique of your current answer."
    )


class ReviseAnswer(AnswerQuestion):
    """Revise your original answer to your question"""

    references: List[str] = Field(description="Citations that motivated your updated answer.")


load_dotenv()

llm = ChatOpenAI(model='gpt-4o-mini')

responder_llm = llm.with_structured_output(
    AnswerQuestion.schema(), method="json_schema"
)

actor_prompt_template = ChatPromptTemplate(
    [
        (
            "system",
            """You are an expert researcher.
            Current time: {time}

            1. {first_instruction}
            2. Reflect and critique your answer. Be severe to maximize improvement.
            3. Recommend search queries to research information and improve your answer"""
        ),
        MessagesPlaceholder(variable_name="messages"),
        (
            "system",
            "Answer the user's question above using the required format."
        ),
    ]
).partial(
    time=lambda: datetime.datetime.now().isoformat()
)

first_responder_prompt_template = actor_prompt_template.partial(
    first_instruction="Provide a detailed ~250 words per answer"
)

first_responder = first_responder_prompt_template | responder_llm

builder = MessageGraph()

builder.add_node("draft", first_responder)
builder.set_entry_point("draft")
builder.set_finish_point("draft")
graph = builder.compile()


res = graph.invoke("""Write about e-sports analytics problem domain,list startups that do that and raised capital.""")

Description

Hello,

I’m currently working on building a Reflexion Agent using LangGraph and Langchain.

The agent is designed to iteratively refine its response through a cycle between an initial responder and a reviser, utilizing Tavily as a real-time search engine as a tool to return a polished and accurate answer.

I'm leveraging OpenAI's structured output feature, recently supported by Langchain using json_schema with strict=true, and attempting to integrate this into my graph.

Issue:

The initial responder successfully generates a structured JSON output as expected. However, the process fails immediately afterward, leading to an error that I haven’t been able to resolve.:

ValueError: Message dict must contain 'role' and 'content' keys, got {'answer': 'E-sports analytics refers to the collection, analysis, and interp

Error Context:

The error appears to occur because the output from the responder LLM, which is in a structured JSON format, does not include the expected role and content keys that the next node in the graph expects.

I attached a minimum reproducible code with just "draft" node, but here's the code ready to be deployed: https://github.com/braisdev/basic-reflexion-agent

System Info

System Information

• Python Version: 3.11
• Platform: macOS

Library	Version
langchain	0.2.14
langchain-community	0.2.12
langchain-core	0.2.32
langchain-openai	0.1.21
langchain-text-splitters	0.2.2
langgraph	0.2.4
langgraph-checkpoint	1.0.3
langsmith	0.1.99

[dosubot[bot]]

Hello, @braisdev! I'm here to help you with your coding questions and bug fixes. Let's work together to resolve the issue you're experiencing.

To resolve the issue where the structured JSON output from the responder LLM does not include the expected role and content keys, you need to ensure that the JSON response from the LLM includes these keys. The _convert_dict_to_message function relies on these keys to convert the dictionary to a BaseMessage object.

Here is the relevant part of the code that expects the role and content keys:

def _convert_dict_to_message(_dict: Mapping[str, Any]) -> BaseMessage:
    role = _dict["role"]
    if role == "user":
        return HumanMessage(content=_dict["content"])
    elif role == "assistant":
        content = _dict.get("content") or ""
        if _dict.get("function_call"):
            _dict["function_call"]["arguments"] = json.dumps(
                _dict["function_call"]["arguments"]
            )
            additional_kwargs = {"function_call": dict(_dict["function_call"])}
        else:
            additional_kwargs = {}
        return AIMessage(content=content, additional_kwargs=additional_kwargs)
    elif role == "system":
        return SystemMessage(content=_dict["content"])
    elif role == "function":
        return FunctionMessage(content=_dict["content"], name=_dict["name"])
    else:
        return ChatMessage(content=_dict["content"], role=role)

To fix the issue, you need to ensure that the JSON response from the LLM includes the role and content keys. If the LLM response is missing these keys, you should modify the LLM or the client code to include them. If modifying the LLM response is not possible, you can add a check in the _create_chat_result method to handle missing keys gracefully.

Here is an example of how you can modify the _create_chat_result method to handle missing keys:

def _create_chat_result(self, response: Mapping[str, Any]) -> ChatResult:
    generations = []
    for res in response["choices"]:
        message_dict = res.get("message", {})
        if "role" not in message_dict or "content" not in message_dict:
            logger.error("Missing 'role' or 'content' in LLM response: %s", message_dict)
            continue  # Skip this message if required keys are missing
        message = _convert_dict_to_message(message_dict)
        gen = ChatGeneration(
            message=message,
            generation_info=dict(finish_reason=res.get("finish_reason")),
        )
        generations.append(gen)
    return ChatResult(generations=generations)

Additionally, you can use the _parse_lc_message function from the LLMonitorCallbackHandler to ensure that the role and content keys are correctly parsed from the LLM response:

def _parse_lc_message(message: BaseMessage) -> Dict[str, Any]:
    keys = ["function_call", "tool_calls", "tool_call_id", "name"]
    parsed = {"text": message.content, "role": _parse_lc_role(message.type)}
    parsed.update(
        {
            key: cast(Any, message.additional_kwargs.get(key))
            for key in keys
            if message.additional_kwargs.get(key) is not None
        }
    )
    return parsed

This function ensures that the role and content keys are included in the parsed message, which can then be used to create the ChatResult object.

Moreover, you can leverage the StructuredOutputParser from LangChain to define and enforce the expected schema for the LLM response. This parser can help ensure that the response adheres to the required structure, including the role and content keys.

Here is an example of how to define the schema and use the StructuredOutputParser:

from langchain.output_parsers.structured import StructuredOutputParser, ResponseSchema

response_schemas = [
    ResponseSchema(name="role", description="The role of the message", type="string"),
    ResponseSchema(name="content", description="The content of the message", type="string"),
]

parser = StructuredOutputParser.from_response_schemas(response_schemas)

# Get format instructions to ensure the LLM response adheres to the schema
format_instructions = parser.get_format_instructions()

# Use the parser to parse the LLM response
parsed_response = parser.parse(llm_response_text)

By using the StructuredOutputParser, you can ensure that the LLM response includes the necessary keys and adheres to the expected structure, thereby preventing the process from failing due to missing keys ^[1][2][3][4].

To continue talking to Dosu, mention @dosu.

[dosubot[bot]]

To address the issue where the structured JSON output from the responder LLM does not include the expected role and content keys, you can modify the _convert_dict_to_message function to include a check for these keys and handle the ValueError appropriately. Here is an updated version of the function:

def _convert_dict_to_message(_dict: Mapping[str, Any]) -> BaseMessage:
    if "role" not in _dict or "content" not in _dict:
        raise ValueError("The dictionary does not contain the required 'role' and 'content' keys.")
    
    role = _dict["role"]
    if role == "user":
        return HumanMessage(content=_dict["content"])
    elif role == "assistant":
        content = _dict.get("content") or ""
        if _dict.get("function_call"):
            _dict["function_call"]["arguments"] = json.dumps(
                _dict["function_call"]["arguments"]
            )
            additional_kwargs = {"function_call": dict(_dict["function_call"])}
        else:
            additional_kwargs = {}
        return AIMessage(content=content, additional_kwargs=additional_kwargs)
    elif role == "system":
        return SystemMessage(content=_dict["content"])
    elif role == "function":
        return FunctionMessage(content=_dict["content"], name=_dict["name"])
    else:
        return ChatMessage(content=_dict["content"], role=role)

Additionally, you can use the _parse_and_clean_json function from the llm.py file to ensure that the message is parsed correctly and includes the necessary keys. Here is an example of how you can integrate this:

def _convert_dict_to_message(_dict: Mapping[str, Any]) -> BaseMessage:
    if "role" not in _dict or "content" not in _dict:
        raise ValueError("The dictionary does not contain the required 'role' and 'content' keys.")
    
    role = _dict["role"]
    if role == "user":
        return HumanMessage(content=_dict["content"])
    elif role == "assistant":
        content = _dict.get("content") or ""
        if _dict.get("function_call"):
            _dict["function_call"]["arguments"] = json.dumps(
                _dict["function_call"]["arguments"]
            )
            additional_kwargs = {"function_call": dict(_dict["function_call"])}
        else:
            additional_kwargs = {}
        return AIMessage(content=content, additional_kwargs=additional_kwargs)
    elif role == "system":
        return SystemMessage(content=_dict["content"])
    elif role == "function":
        return FunctionMessage(content=_dict["content"], name=_dict["name"])
    else:
        return ChatMessage(content=_dict["content"], role=role)

def _parse_and_clean_json(argument_json: Dict[str, Any]) -> Tuple[List[Node], List[Relationship]]:
    nodes = []
    for node in argument_json["nodes"]:
        if not node.get("id"):  # Id is mandatory, skip this node
            continue
        node_properties = {}
        if "properties" in node and node["properties"]:
            for p in node["properties"]:
                node_properties[format_property_key(p["key"])] = p["value"]
        nodes.append(
            Node(
                id=node["id"],
                type=node.get("type"),
                properties=node_properties,
            )
        )
    relationships = []
    for rel in argument_json["relationships"]:
        # Mandatory props
        if (
            not rel.get("source_node_id")
            or not rel.get("target_node_id")
            or not rel.get("type")
        ):
            continue

        # Node type copying if needed from node list
        if not rel.get("source_node_type"):
            try:
                rel["source_node_type"] = [
                    el.get("type")
                    for el in argument_json["nodes"]
                    if el["id"] == rel["source_node_id"]
                ][0]
            except IndexError:
                rel["source_node_type"] = None
        if not rel.get("target_node_type"):
            try:
                rel["target_node_type"] = [
                    el.get("type")
                    for el in argument_json["nodes"]
                    if el["id"] == rel["target_node_id"]
                ][0]
            except IndexError:
                rel["target_node_type"] = None

        rel_properties = {}
        if "properties" in rel and rel["properties"]:
            for p in rel["properties"]:
                rel_properties[format_property_key(p["key"])] = p["value"]
        source_node = Node(
            id=rel["source_node_id"],
            type=rel["source_node_type"],
        )
        target_node = Node(
            id=rel["target_node_id"],
            type=rel["target_node_type"],
        )
        relationships.append(
            Relationship(
                source=source_node,
                target=target_node,
                type=rel["type"],
                properties=rel_properties,
            )
        )
    return nodes, relationships

This modification ensures that the function checks for the presence of the role and content keys and raises a ValueError if they are missing, which will help in identifying and handling the issue appropriately. Additionally, it leverages the _parse_and_clean_json function to ensure the message is correctly parsed and includes all necessary keys ^[1][2][3].

To continue talking to Dosu, mention @dosu.

[braisdev]

Hello, @dosu! Thanks for your response and the detailed explanation. However, I would prefer not to change any built-in Python functions from LangChain or Langraph. I'm looking for a way to resolve the error with the current setup.

Could you advise if I might be missing something in my current implementation, or suggest a workaround, perhaps using a parser or some other method? If you're unable to provide a direct solution, I'd appreciate it if you could reflect on other potential approaches before responding.

Thanks again for your help!

[dosubot[bot]]

To resolve the ValueError related to missing 'role' and 'content' keys without modifying built-in functions, you can ensure that the output from the LLM is correctly formatted before it is passed to the next node in the graph. Here are a few approaches you can consider:

1. Use a Custom Wrapper to Format the Output

Create a custom wrapper that formats the LLM output to include the necessary 'role' and 'content' keys before passing it to the next node.

from langchain_core.messages import AIMessage, HumanMessage, SystemMessage

def format_llm_output(output: dict) -> List[BaseMessage]:
    formatted_messages = []
    for item in output:
        if 'role' not in item or 'content' not in item:
            raise ValueError("The dictionary does not contain the required 'role' and 'content' keys.")
        if item['role'] == 'user':
            formatted_messages.append(HumanMessage(content=item['content']))
        elif item['role'] == 'assistant':
            formatted_messages.append(AIMessage(content=item['content']))
        elif item['role'] == 'system':
            formatted_messages.append(SystemMessage(content=item['content']))
    return formatted_messages

# Example usage
llm_output = responder_llm.invoke(query)
formatted_output = format_llm_output(llm_output)

2. Modify the Prompt Template to Include Role and Content

Ensure that the prompt template used in the ChatPromptTemplate includes the 'role' and 'content' keys explicitly.

actor_prompt_template = ChatPromptTemplate(
    [
        (
            "system",
            """You are an expert researcher.
            Current time: {time}

            1. {first_instruction}
            2. Reflect and critique your answer. Be severe to maximize improvement.
            3. Recommend search queries to research information and improve your answer"""
        ),
        MessagesPlaceholder(variable_name="messages"),
        (
            "system",
            "Answer the user's question above using the required format."
        ),
    ]
).partial(
    time=lambda: datetime.datetime.now().isoformat()
)

first_responder_prompt_template = actor_prompt_template.partial(
    first_instruction="Provide a detailed ~250 words per answer"
)

first_responder = first_responder_prompt_template | responder_llm

3. Use a Structured Output Parser

Leverage the StructuredOutputParser to enforce the expected schema for the LLM response, ensuring it includes the 'role' and 'content' keys.

from langchain.output_parsers.structured import StructuredOutputParser, ResponseSchema

response_schemas = [
    ResponseSchema(name="role", description="The role of the message", type="string"),
    ResponseSchema(name="content", description="The content of the message", type="string"),
]

parser = StructuredOutputParser.from_response_schemas(response_schemas)

# Get format instructions to ensure the LLM response adheres to the schema
format_instructions = parser.get_format_instructions()

# Use the parser to parse the LLM response
llm_response_text = responder_llm.invoke(query)
parsed_response = parser.parse(llm_response_text)

4. Use a Middleware to Process the Output

Implement a middleware function that processes the LLM output to ensure it includes the necessary keys before passing it to the next node.

def middleware_process_output(output: dict) -> dict:
    if 'role' not in output or 'content' not in output:
        output['role'] = 'assistant'  # Default role
        output['content'] = output.get('answer', '')  # Default content
    return output

# Example usage
llm_output = responder_llm.invoke(query)
processed_output = middleware_process_output(llm_output)

By using one of these approaches, you can ensure that the LLM output is correctly formatted with the required 'role' and 'content' keys, thereby avoiding the ValueError and ensuring smooth transitions in your LangChain and LangGraph setup.

To continue talking to Dosu, mention @dosu.

[braisdev]

Hi @dosu ! this is not what I'm looking for. Give me an example of a graph with one unique node, start -> node1 -> end, that works with structured output with method="json_schema"

[dosubot[bot]]

Here's an example of a graph with one unique node (node1) between the start and end nodes, using structured output with method="json_schema":

from langchain_core.runnables.graph import Graph
from langchain_core.output_parsers.string import StrOutputParser
from langchain_core.pydantic_v1 import BaseModel

class Schema(BaseModel):
    a: int

graph = Graph()
start = graph.add_node(Schema, id="__start__")
node1 = graph.add_node(StrOutputParser(), id="node1")
end = graph.add_node(Schema, id="__end__")
graph.add_edge(start, node1)
graph.add_edge(node1, end)

print(graph.to_json(with_schemas=True))

This code sets up a simple graph with a start node, a single intermediate node (node1), and an end node. The graph is configured to output in a structured format using method="json_schema" ^[1].

To continue talking to Dosu, mention @dosu.

[BharathKumarRavichandran]

Hi @braisdev , were you able to resolve this issue?