Repository: jerryjliu/gpt_index
File: llama-index-core/llama_index/core/agent/utils.py (Python)

Code before:

from llama_index.core.llms import ChatMessage, TextBlock
from typing import List, Type, Dict, Any, Optional, cast
from llama_index.core.bridge.pydantic import BaseModel
from llama_index.core.agent.types import TaskStep
from llama_index.core.base.llms.types import ChatMessage, MessageRole
from llama_index.core.llms import LLM
from llama_index.core.memory import BaseMemory


def add_user_step_to_memory(
    step: TaskStep, memory: BaseMemory, verbose: bool = False
) -> None:
    """Add user step to memory."""
    user_message = ChatMessage(content=step.input, role=MessageRole.USER)
    memory.put(user_message)
    if verbose:
        print(f"Added user message to memory: {step.input}")


def messages_to_xml_format(messages: List[ChatMessage]) -> List[ChatMessage]:
    blocks = [TextBlock(text="<current_conversation>\n")]
    system_msg: Optional[str] = None
    for message in messages:
        if message.role.value == "system":
            system_msg = message.content
        blocks.append(TextBlock(text=f"\t<{message.role.value}>\n"))
        for block in message.blocks:
            if isinstance(block, TextBlock):
                blocks.append(TextBlock(text=f"\t\t<message>{block.text}</message>\n"))
        blocks.append(TextBlock(text=f"\t</{message.role.value}>\n"))
    blocks.append(TextBlock(text="</current_conversation>\n\n"))
    blocks.append(
        TextBlock(
            text="Given the conversation, format the output according to the provided schema."
        )
    )
    messages = [ChatMessage(role="user", blocks=blocks)]
    if system_msg:
        messages.insert(
            0, ChatMessage(role=MessageRole.SYSTEM, blocks=[TextBlock(text=system_msg)])
        )
    return messages


async def generate_structured_response(
    messages: List[ChatMessage], llm: LLM, output_cls: Type[BaseModel]
) -> Dict[str, Any]:
    xml_message = messages_to_xml_format(messages)
    structured_response = await llm.as_structured_llm(
        output_cls,

Review comment (bug):
why not just copy the entire `message` object vs the content? This way multi-modal messages are copied over, along with any other edge cases

Code after:

from llama_index.core.llms import ChatMessage, TextBlock
from typing import List, Type, Dict, Any, Optional, cast
from llama_index.core.bridge.pydantic import BaseModel
from llama_index.core.agent.types import TaskStep
from llama_index.core.base.llms.types import ChatMessage, MessageRole
from llama_index.core.llms import LLM
from llama_index.core.memory import BaseMemory


def add_user_step_to_memory(
    step: TaskStep, memory: BaseMemory, verbose: bool = False
) -> None:
    """Add user step to memory."""
    user_message = ChatMessage(content=step.input, role=MessageRole.USER)
    memory.put(user_message)
    if verbose:
        print(f"Added user message to memory: {step.input}")


def messages_to_xml_format(messages: List[ChatMessage]) -> List[ChatMessage]:
    blocks = [TextBlock(text="<current_conversation>\n")]
    system_msg: Optional[ChatMessage] = None
    for message in messages:
        if message.role.value == "system":
            system_msg = message
        blocks.append(TextBlock(text=f"\t<{message.role.value}>\n"))
        for block in message.blocks:
            if isinstance(block, TextBlock):
                blocks.append(TextBlock(text=f"\t\t<message>{block.text}</message>\n"))
        blocks.append(TextBlock(text=f"\t</{message.role.value}>\n"))
    blocks.append(TextBlock(text="</current_conversation>\n\n"))
    blocks.append(
        TextBlock(
            text="Given the conversation, format the output according to the provided schema."
        )
    )
    messages = [ChatMessage(role="user", blocks=blocks)]
    if system_msg:
        messages.insert(0, system_msg)
    return messages


async def generate_structured_response(
    messages: List[ChatMessage], llm: LLM, output_cls: Type[BaseModel]
) -> Dict[str, Any]:
    xml_message = messages_to_xml_format(messages)
    structured_response = await llm.as_structured_llm(
        output_cls,
    ).achat(messages=xml_message)
    return cast(Dict[str, Any], json.loads(structured_response.message.content))