18 lines
525 B
Python
18 lines
525 B
Python
"""Agent utils."""
|
|
|
|
|
|
from llama_index.agent.types import TaskStep
|
|
from llama_index.core.llms.types import MessageRole
|
|
from llama_index.llms.base import ChatMessage
|
|
from llama_index.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}")
|