423 lines
17 KiB
Python
423 lines
17 KiB
Python
import json
|
||
import xml.etree.ElementTree as ET
|
||
from typing import AsyncGenerator, Dict, List, Tuple, Union
|
||
|
||
import openai
|
||
|
||
from letta.agents.base_agent import BaseAgent
|
||
from letta.schemas.agent import AgentState
|
||
from letta.schemas.block import BlockUpdate
|
||
from letta.schemas.enums import MessageStreamStatus
|
||
from letta.schemas.letta_message import LegacyLettaMessage, LettaMessage
|
||
from letta.schemas.letta_message_content import TextContent
|
||
from letta.schemas.letta_response import LettaResponse
|
||
from letta.schemas.message import MessageCreate
|
||
from letta.schemas.openai.chat_completion_request import ChatCompletionRequest, SystemMessage, Tool, UserMessage
|
||
from letta.schemas.usage import LettaUsageStatistics
|
||
from letta.schemas.user import User
|
||
from letta.server.rest_api.utils import convert_in_context_letta_messages_to_openai, create_input_messages
|
||
from letta.services.agent_manager import AgentManager
|
||
from letta.services.block_manager import BlockManager
|
||
from letta.services.message_manager import MessageManager
|
||
|
||
|
||
class EphemeralMemoryAgent(BaseAgent):
|
||
"""
|
||
A stateless agent that helps with offline memory computations.
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
agent_id: str,
|
||
openai_client: openai.AsyncClient,
|
||
message_manager: MessageManager,
|
||
agent_manager: AgentManager,
|
||
block_manager: BlockManager,
|
||
target_block_label: str,
|
||
message_transcripts: List[str],
|
||
actor: User,
|
||
):
|
||
super().__init__(
|
||
agent_id=agent_id,
|
||
openai_client=openai_client,
|
||
message_manager=message_manager,
|
||
agent_manager=agent_manager,
|
||
actor=actor,
|
||
)
|
||
|
||
self.block_manager = block_manager
|
||
self.target_block_label = target_block_label
|
||
self.message_transcripts = message_transcripts
|
||
|
||
def update_message_transcript(self, message_transcripts: List[str]):
|
||
self.message_transcripts = message_transcripts
|
||
|
||
async def step(self, input_messages: List[MessageCreate], max_steps: int = 10) -> LettaResponse:
|
||
"""
|
||
Process the user's input message, allowing the model to call memory-related tools
|
||
until it decides to stop and provide a final response.
|
||
"""
|
||
agent_state = self.agent_manager.get_agent_by_id(agent_id=self.agent_id, actor=self.actor)
|
||
in_context_messages = create_input_messages(input_messages=input_messages, agent_id=self.agent_id, actor=self.actor)
|
||
openai_messages = convert_in_context_letta_messages_to_openai(in_context_messages, exclude_system_messages=True)
|
||
|
||
# 1. Store memories
|
||
request = self._build_openai_request(
|
||
openai_messages, agent_state, tools=self._build_store_memory_tool_schemas(), system=self._get_memory_store_system_prompt()
|
||
)
|
||
|
||
chat_completion = await self.openai_client.chat.completions.create(**request.model_dump(exclude_unset=True))
|
||
assistant_message = chat_completion.choices[0].message
|
||
|
||
# Process tool calls
|
||
tool_call = assistant_message.tool_calls[0]
|
||
function_name = tool_call.function.name
|
||
function_args = json.loads(tool_call.function.arguments)
|
||
|
||
if function_name == "store_memories":
|
||
print("Called store_memories")
|
||
print(function_args)
|
||
for chunk_args in function_args.get("chunks"):
|
||
self.store_memory(agent_state=agent_state, **chunk_args)
|
||
result = "Successfully stored memories"
|
||
else:
|
||
raise ValueError("Error: Unknown tool function '{function_name}'")
|
||
|
||
openai_messages.append(
|
||
{
|
||
"role": "assistant",
|
||
"content": assistant_message.content,
|
||
"tool_calls": [
|
||
{
|
||
"id": tool_call.id,
|
||
"type": "function",
|
||
"function": {"name": function_name, "arguments": tool_call.function.arguments},
|
||
}
|
||
],
|
||
}
|
||
)
|
||
openai_messages.append({"role": "tool", "tool_call_id": tool_call.id, "content": str(result)})
|
||
|
||
# 2. Execute rethink block memory loop
|
||
human_block_content = self.agent_manager.get_block_with_label(
|
||
agent_id=self.agent_id, block_label=self.target_block_label, actor=self.actor
|
||
)
|
||
rethink_command = f"""
|
||
Here is the current memory block created earlier:
|
||
|
||
### CURRENT MEMORY
|
||
{human_block_content}
|
||
### END CURRENT MEMORY
|
||
|
||
Please refine this block:
|
||
|
||
- Merge in any new facts and remove outdated or contradictory details.
|
||
- Organize related information together (e.g., preferences, background, ongoing goals).
|
||
- Add any light, supportable inferences that deepen understanding—but do not invent unsupported details.
|
||
|
||
Use `rethink_user_memor(new_memory)` as many times as you need to iteratively improve the text. When it’s fully polished and complete, call `finish_rethinking_memory()`.
|
||
"""
|
||
rethink_command = UserMessage(content=rethink_command)
|
||
openai_messages.append(rethink_command.model_dump())
|
||
|
||
for _ in range(max_steps):
|
||
request = self._build_openai_request(
|
||
openai_messages, agent_state, tools=self._build_sleeptime_tools(), system=self._get_rethink_memory_system_prompt()
|
||
)
|
||
chat_completion = await self.openai_client.chat.completions.create(**request.model_dump(exclude_unset=True))
|
||
assistant_message = chat_completion.choices[0].message
|
||
|
||
# Process tool calls
|
||
tool_call = assistant_message.tool_calls[0]
|
||
function_name = tool_call.function.name
|
||
function_args = json.loads(tool_call.function.arguments)
|
||
|
||
if function_name == "rethink_user_memor":
|
||
print("Called rethink_user_memor")
|
||
print(function_args)
|
||
result = self.rethink_user_memory(agent_state=agent_state, **function_args)
|
||
elif function_name == "finish_rethinking_memory":
|
||
print("Called finish_rethinking_memory")
|
||
break
|
||
else:
|
||
result = f"Error: Unknown tool function '{function_name}'"
|
||
openai_messages.append(
|
||
{
|
||
"role": "assistant",
|
||
"content": assistant_message.content,
|
||
"tool_calls": [
|
||
{
|
||
"id": tool_call.id,
|
||
"type": "function",
|
||
"function": {"name": function_name, "arguments": tool_call.function.arguments},
|
||
}
|
||
],
|
||
}
|
||
)
|
||
openai_messages.append({"role": "tool", "tool_call_id": tool_call.id, "content": str(result)})
|
||
|
||
# Actually save the memory:
|
||
target_block = agent_state.memory.get_block(self.target_block_label)
|
||
self.block_manager.update_block(block_id=target_block.id, block_update=BlockUpdate(value=target_block.value), actor=self.actor)
|
||
|
||
return LettaResponse(messages=[], usage=LettaUsageStatistics())
|
||
|
||
def _format_messages_llm_friendly(self):
|
||
messages = self.message_manager.list_messages_for_agent(agent_id=self.agent_id, actor=self.actor)
|
||
|
||
llm_friendly_messages = [f"{m.role}: {m.content[0].text}" for m in messages if m.content and isinstance(m.content[0], TextContent)]
|
||
return "\n".join(llm_friendly_messages)
|
||
|
||
def _build_openai_request(
|
||
self, openai_messages: List[Dict], agent_state: AgentState, tools: List[Tool], system: str
|
||
) -> ChatCompletionRequest:
|
||
system_message = SystemMessage(role="system", content=system)
|
||
openai_request = ChatCompletionRequest(
|
||
model="gpt-4o", # agent_state.llm_config.model, # TODO: Separate config for summarizer?
|
||
messages=[system_message] + openai_messages,
|
||
tools=tools,
|
||
tool_choice="required",
|
||
user=self.actor.id,
|
||
max_completion_tokens=agent_state.llm_config.max_tokens,
|
||
temperature=agent_state.llm_config.temperature,
|
||
stream=False,
|
||
)
|
||
return openai_request
|
||
|
||
def _build_store_memory_tool_schemas(self) -> List[Tool]:
|
||
"""
|
||
Build the schemas for the three memory-related tools.
|
||
"""
|
||
tools = [
|
||
Tool(
|
||
type="function",
|
||
function={
|
||
"name": "store_memories",
|
||
"description": "Archive coherent chunks of dialogue that will be evicted, preserving raw lines and a brief contextual description.",
|
||
"parameters": {
|
||
"type": "object",
|
||
"properties": {
|
||
"chunks": {
|
||
"type": "array",
|
||
"items": {
|
||
"type": "object",
|
||
"properties": {
|
||
"start_index": {"type": "integer", "description": "Index of first line in original history."},
|
||
"end_index": {"type": "integer", "description": "Index of last line in original history."},
|
||
"context": {
|
||
"type": "string",
|
||
"description": "A high-level description providing context for why this chunk matters.",
|
||
},
|
||
},
|
||
"required": ["start_index", "end_index", "context"],
|
||
},
|
||
}
|
||
},
|
||
"required": ["chunks"],
|
||
"additionalProperties": False,
|
||
},
|
||
},
|
||
),
|
||
]
|
||
|
||
return tools
|
||
|
||
def _build_sleeptime_tools(self) -> List[Tool]:
|
||
tools = [
|
||
Tool(
|
||
type="function",
|
||
function={
|
||
"name": "rethink_user_memory",
|
||
"description": (
|
||
"Rewrite memory block for the main agent, new_memory should contain all current "
|
||
"information from the block that is not outdated or inconsistent, integrating any "
|
||
"new information, resulting in a new memory block that is organized, readable, and "
|
||
"comprehensive."
|
||
),
|
||
"parameters": {
|
||
"type": "object",
|
||
"properties": {
|
||
"new_memory": {
|
||
"type": "string",
|
||
"description": (
|
||
"The new memory with information integrated from the memory block. "
|
||
"If there is no new information, then this should be the same as the "
|
||
"content in the source block."
|
||
),
|
||
},
|
||
},
|
||
"required": ["new_memory"],
|
||
"additionalProperties": False,
|
||
},
|
||
},
|
||
),
|
||
Tool(
|
||
type="function",
|
||
function={
|
||
"name": "finish_rethinking_memory",
|
||
"description": ("This function is called when the agent is done rethinking the memory."),
|
||
"parameters": {
|
||
"type": "object",
|
||
"properties": {},
|
||
"required": [],
|
||
"additionalProperties": False,
|
||
},
|
||
},
|
||
),
|
||
]
|
||
|
||
return tools
|
||
|
||
def rethink_user_memory(self, new_memory: str, agent_state: AgentState) -> str:
|
||
if agent_state.memory.get_block(self.target_block_label) is None:
|
||
agent_state.memory.create_block(label=self.target_block_label, value=new_memory)
|
||
|
||
agent_state.memory.update_block_value(label=self.target_block_label, value=new_memory)
|
||
return "Successfully updated memory"
|
||
|
||
def store_memory(self, start_index: int, end_index: int, context: str, agent_state: AgentState) -> str:
|
||
"""
|
||
Store a memory.
|
||
"""
|
||
try:
|
||
messages = self.message_transcripts[start_index : end_index + 1]
|
||
memory = self.serialize(messages, context)
|
||
self.agent_manager.passage_manager.insert_passage(
|
||
agent_state=agent_state,
|
||
agent_id=agent_state.id,
|
||
text=memory,
|
||
actor=self.actor,
|
||
)
|
||
self.agent_manager.rebuild_system_prompt(agent_id=agent_state.id, actor=self.actor, force=True)
|
||
|
||
return "Sucessfully stored memory"
|
||
except Exception as e:
|
||
return f"Failed to store memory given start_index {start_index} and end_index {end_index}: {e}"
|
||
|
||
def serialize(self, messages: List[str], context: str) -> str:
|
||
"""
|
||
Produce an XML document like:
|
||
|
||
<memory>
|
||
<messages>
|
||
<message>…</message>
|
||
<message>…</message>
|
||
…
|
||
</messages>
|
||
<context>…</context>
|
||
</memory>
|
||
"""
|
||
root = ET.Element("memory")
|
||
|
||
msgs_el = ET.SubElement(root, "messages")
|
||
for msg in messages:
|
||
m = ET.SubElement(msgs_el, "message")
|
||
m.text = msg
|
||
|
||
sum_el = ET.SubElement(root, "context")
|
||
sum_el.text = context
|
||
|
||
# ET.tostring will escape reserved chars for you
|
||
return ET.tostring(root, encoding="unicode")
|
||
|
||
def deserialize(self, xml_str: str) -> Tuple[List[str], str]:
|
||
"""
|
||
Parse the XML back into (messages, context). Raises ValueError if tags are missing.
|
||
"""
|
||
try:
|
||
root = ET.fromstring(xml_str)
|
||
except ET.ParseError as e:
|
||
raise ValueError(f"Invalid XML: {e}")
|
||
|
||
msgs_el = root.find("messages")
|
||
if msgs_el is None:
|
||
raise ValueError("Missing <messages> section")
|
||
|
||
messages = []
|
||
for m in msgs_el.findall("message"):
|
||
# .text may be None if empty, so coerce to empty string
|
||
messages.append(m.text or "")
|
||
|
||
sum_el = root.find("context")
|
||
if sum_el is None:
|
||
raise ValueError("Missing <context> section")
|
||
context = sum_el.text or ""
|
||
|
||
return messages, context
|
||
|
||
async def step_stream(
|
||
self, input_messages: List[MessageCreate], max_steps: int = 10
|
||
) -> AsyncGenerator[Union[LettaMessage, LegacyLettaMessage, MessageStreamStatus], None]:
|
||
"""
|
||
This agent is synchronous-only. If called in an async context, raise an error.
|
||
"""
|
||
raise NotImplementedError("EphemeralMemoryAgent does not support async step.")
|
||
|
||
# TODO: Move these to independent text files
|
||
def _get_memory_store_system_prompt(self) -> str:
|
||
return """
|
||
You are a memory-recall assistant working asynchronously alongside a main chat agent that retains only a portion of the message history in its context window.
|
||
|
||
When given a full transcript with lines marked (Older) or (Newer), you should:
|
||
1. Segment the (Older) portion into coherent chunks by topic, instruction, or preference.
|
||
2. For each chunk, produce only:
|
||
- start_index: the first line’s index
|
||
- end_index: the last line’s index
|
||
- context: a blurb explaining why this chunk matters
|
||
|
||
Return exactly one JSON tool call to `store_memories`, consider this miniature example:
|
||
|
||
---
|
||
|
||
(Older)
|
||
0. user: Okay. Got it. Keep your answers shorter, please.
|
||
1. assistant: Sure thing! I’ll keep it brief. What would you like to know?
|
||
2. user: I like basketball.
|
||
3. assistant: That's great! Do you have a favorite team or player?
|
||
|
||
(Newer)
|
||
4. user: Yeah. I like basketball.
|
||
5. assistant: Awesome! What do you enjoy most about basketball?
|
||
|
||
---
|
||
|
||
Example output:
|
||
|
||
```json
|
||
{
|
||
"name": "store_memories",
|
||
"arguments": {
|
||
"chunks": [
|
||
{
|
||
"start_index": 0,
|
||
"end_index": 1,
|
||
"context": "User explicitly asked the assistant to keep responses concise."
|
||
},
|
||
{
|
||
"start_index": 2,
|
||
"end_index": 3,
|
||
"context": "User enjoys basketball and prompted follow-up about their favorite team or player."
|
||
}
|
||
]
|
||
}
|
||
}
|
||
```
|
||
"""
|
||
|
||
def _get_rethink_memory_system_prompt(self) -> str:
|
||
return """
|
||
SYSTEM
|
||
You are a Memory-Updater agent. Your job is to iteratively refine the given memory block until it’s concise, organized, and complete.
|
||
|
||
Instructions:
|
||
- Call `rethink_user_memor(new_memory: string)` as many times as you like. Each call should submit a fully revised version of the block so far.
|
||
- When you’re fully satisfied, call `finish_rethinking_memory()`.
|
||
- Don’t output anything else—only the JSON for these tool calls.
|
||
|
||
Goals:
|
||
- Merge in new facts and remove contradictions.
|
||
- Group related details (preferences, biography, goals).
|
||
- Draw light, supportable inferences without inventing facts.
|
||
- Preserve every critical piece of information.
|
||
"""
|