ryan/simbarag
1
0
Fork 0
Code Issues Pull Requests 4 Actions Packages Projects Releases Wiki Activity

Fix history roles and persist tool results for cross-turn context

Browse Source 
Two fixes for the agent losing track of conversation context:

- History replay compared msg.speaker (a Speaker enum) to the string
  "user", which is always false, so every prior message — including the
  user's — was sent to the LLM with role "assistant". Compare against
  Speaker.USER instead. Affected web, iMessage, WhatsApp, and email.

- Tool outputs (web searches, RAG chunks) were discarded after each
  turn, so follow-up questions about fetched content had nothing to
  reference. Persist them as speaker="tool" messages (capped at 4000
  chars) and replay them in history as assistant context. The
  conversation GET endpoint serves only their header line so the UI
  tool pill stays compact.

Also bumps the replayed history window from 10 to 20 messages since
tool messages now share it.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
Ryan Chen
2026-09-12 15:14:33 -04:00
parent 0389ec45c8
commit 395c97958f
6 changed files with 79 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files
blueprints/conversation/__init__.py
+20 -4
View File
@@ -19,12 +19,15 @@ from utils.s3_client import upload_image as s3_upload_image
from .agents import main_agent
from .logic import (
add_message_to_conversation,
persist_agent_tool_messages,
tool_result_text,
)
from .memory import get_memories_for_user
from .models import (
Conversation,
PydConversation,
PydListConversation,
Speaker,
)
from .prompts import SIMBA_SYSTEM_PROMPT
@@ -68,13 +71,13 @@ def _build_messages_payload(
system_prompt: str | None = None,
) -> list:
recent_messages = (
conversation.messages[-10:]
if len(conversation.messages) > 10
conversation.messages[-20:]
if len(conversation.messages) > 20
else conversation.messages
)
messages_payload = [{"role": "system", "content": system_prompt or _SYSTEM_PROMPT}]
for msg in recent_messages[:-1]: # Exclude the message we just added
role = "user" if msg.speaker == "user" else "assistant"
role = "user" if msg.speaker == Speaker.USER else "assistant"
text = msg.text
if msg.image_key and role == "user":
text = f"[User sent an image]\n{text}"
@@ -123,6 +126,7 @@ async def query():
response = await main_agent.ainvoke(payload, config=agent_config)
message = response.get("messages", [])[-1].content
await persist_agent_tool_messages(conversation, user, response.get("messages", []))
await add_message_to_conversation(
conversation=conversation,
message=message,
@@ -234,6 +238,15 @@ async def stream_query():
elif event_type == "on_tool_end":
yield f"data: {json.dumps({'type': 'tool_end', 'tool': event['name']})}\n\n"
# Persist tool output so later turns keep the fetched context
output = event.get("data", {}).get("output")
if output is not None:
await add_message_to_conversation(
conversation=conversation,
message=tool_result_text(event["name"], output),
speaker="tool",
user=user,
)
elif event_type == "on_chat_model_stream":
chunk = event.get("data", {}).get("chunk")
@@ -299,10 +312,13 @@ async def get_conversation(conversation_id: str):
# Manually serialize the conversation with messages
messages = []
for msg in conversation.messages:
# Tool messages hold raw fetched context for the LLM; only surface
# their header line (e.g. "[web_search result]") in the UI.
text = msg.text.split("\n", 1)[0] if msg.speaker == Speaker.TOOL else msg.text
messages.append(
{
"id": str(msg.id),
"text": msg.text,
"text": text,
"speaker": msg.speaker.value,
"created_at": msg.created_at.isoformat(),
"image_key": msg.image_key,
blueprints/conversation/logic.py
+36
View File
@@ -1,11 +1,47 @@
from datetime import datetime, timezone
import tortoise.exceptions
from langchain_core.messages import ToolMessage
import blueprints.users.models
from .models import Conversation, ConversationMessage
# Cap persisted tool output so replayed history stays within the LLM context
# budget (llama-server slots are 32k tokens; base prompt + tools already use ~6k).
TOOL_RESULT_MAX_CHARS = 4000
def tool_result_text(tool_name: str, output) -> str:
"""Format a tool's output for persistence as a conversation message."""
content = getattr(output, "content", output)
if isinstance(content, (list, tuple)):
content = "\n".join(str(item) for item in content)
text = str(content)
if len(text) > TOOL_RESULT_MAX_CHARS:
text = text[:TOOL_RESULT_MAX_CHARS] + "\n[truncated]"
return f"[{tool_name} result]\n{text}"
async def persist_agent_tool_messages(
conversation: Conversation,
user: blueprints.users.models.User,
response_messages: list,
) -> None:
"""Persist ToolMessages from an agent response as 'tool' conversation messages.
Keeps fetched context (web searches, RAG chunks, etc.) available to later
turns, which otherwise only see the final text answers.
"""
for msg in response_messages:
if isinstance(msg, ToolMessage):
await add_message_to_conversation(
conversation=conversation,
message=tool_result_text(msg.name or "tool", msg),
speaker="tool",
user=user,
)
async def create_conversation(name: str = "") -> Conversation:
conversation = await Conversation.create(name=name)
blueprints/conversation/models.py
+1
View File
@@ -11,6 +11,7 @@ from tortoise.models import Model
class Speaker(enum.Enum):
USER = "user"
SIMBA = "simba"
TOOL = "tool"
class Conversation(Model):
blueprints/email/__init__.py
+7 -2
View File
@@ -13,8 +13,10 @@ from blueprints.users.models import User
from blueprints.conversation.logic import (
get_conversation_for_channel,
add_message_to_conversation,
persist_agent_tool_messages,
)
from blueprints.conversation.agents import main_agent
from blueprints.conversation.models import Speaker
from blueprints.conversation.prompts import SIMBA_SYSTEM_PROMPT
from . import models # noqa: F401 — register Tortoise ORM models
from .helpers import generate_email_token, get_user_email_address # noqa: F401
@@ -193,17 +195,20 @@ async def webhook():
# Build messages payload
try:
messages = await conversation.messages.all()
recent_messages = list(messages)[-10:]
recent_messages = list(messages)[-20:]
messages_payload = [{"role": "system", "content": SIMBA_SYSTEM_PROMPT}]
for msg in recent_messages[:-1]:
role = "user" if msg.speaker == "user" else "assistant"
role = "user" if msg.speaker == Speaker.USER else "assistant"
messages_payload.append({"role": role, "content": msg.text})
messages_payload.append({"role": "user", "content": body})
logger.info(f"Invoking LangChain agent with {len(messages_payload)} messages")
response = await main_agent.ainvoke({"messages": messages_payload})
response_text = response.get("messages", [])[-1].content
await persist_agent_tool_messages(
conversation, user, response.get("messages", [])
)
except Exception as e:
logger.error(f"Error invoking agent for email: {e}")
response_text = "Sorry, I'm having trouble thinking right now."
blueprints/imessage/__init__.py
+7 -2
View File
@@ -12,7 +12,9 @@ from blueprints.users.models import User
from blueprints.conversation.logic import (
get_conversation_for_channel,
add_message_to_conversation,
persist_agent_tool_messages,
)
from blueprints.conversation.models import Speaker
from blueprints.conversation.agents import main_agent
from blueprints.conversation.prompts import SIMBA_SYSTEM_PROMPT
@@ -217,12 +219,12 @@ async def webhook():
# Build messages payload for LangChain agent
try:
messages = await conversation.messages.all()
recent_messages = list(messages)[-10:]
recent_messages = list(messages)[-20:]
messages_payload = [{"role": "system", "content": SIMBA_SYSTEM_PROMPT}]
for msg in recent_messages[:-1]:
role = "user" if msg.speaker == "user" else "assistant"
role = "user" if msg.speaker == Speaker.USER else "assistant"
messages_payload.append({"role": role, "content": msg.text})
if image_block:
@@ -237,6 +239,9 @@ async def webhook():
logger.info(f"Invoking LangChain agent with {len(messages_payload)} messages")
response = await main_agent.ainvoke({"messages": messages_payload})
response_text = response.get("messages", [])[-1].content
await persist_agent_tool_messages(
conversation, user, response.get("messages", [])
)
except Exception as e:
logger.error(f"Error invoking agent: {e}")
blueprints/whatsapp/__init__.py
+8 -3
View File
@@ -11,7 +11,9 @@ from blueprints.users.models import User
from blueprints.conversation.logic import (
get_conversation_for_channel,
add_message_to_conversation,
persist_agent_tool_messages,
)
from blueprints.conversation.models import Speaker
from blueprints.conversation.agents import main_agent
from blueprints.conversation.prompts import SIMBA_SYSTEM_PROMPT
@@ -178,16 +180,16 @@ async def webhook():
# Build messages payload for LangChain agent with system prompt and conversation history
try:
# Get last 10 messages for conversation history
# Get last 20 messages for conversation history
messages = await conversation.messages.all()
recent_messages = list(messages)[-10:]
recent_messages = list(messages)[-20:]
# Build messages payload
messages_payload = [{"role": "system", "content": SIMBA_SYSTEM_PROMPT}]
# Add recent conversation history (exclude the message we just added)
for msg in recent_messages[:-1]:
role = "user" if msg.speaker == "user" else "assistant"
role = "user" if msg.speaker == Speaker.USER else "assistant"
messages_payload.append({"role": role, "content": msg.text})
# Add current query
@@ -197,6 +199,9 @@ async def webhook():
logger.info(f"Invoking LangChain agent with {len(messages_payload)} messages")
response = await main_agent.ainvoke({"messages": messages_payload})
response_text = response.get("messages", [])[-1].content
await persist_agent_tool_messages(
conversation, user, response.get("messages", [])
)
# Log YNAB availability
if os.getenv("YNAB_ACCESS_TOKEN"):