Agent improvements: image passthrough, context fixes, thinking indicator, observability #50

Open
ryan wants to merge 5 commits from agent-improvements into main
11 changed files with 223 additions and 99 deletions
+109 -24
View File
@@ -1,6 +1,7 @@
import asyncio
import json
import logging
import time
import uuid
from quart import Blueprint, jsonify, make_response, request
@@ -10,7 +11,7 @@ from quart_jwt_extended import (
)
import blueprints.users.models
from utils.image_process import analyze_user_image
from utils.image_process import image_content_block
from utils.image_upload import ImageValidationError, process_image
from utils.s3_client import generate_presigned_url as s3_presigned_url
from utils.s3_client import get_image as s3_get_image
@@ -19,12 +20,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
@@ -64,27 +68,29 @@ async def _build_system_prompt_with_memories(user_id: str) -> str:
def _build_messages_payload(
conversation,
query_text: str,
image_description: str | None = None,
image_block: dict | None = None,
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}"
messages_payload.append({"role": role, "content": text})
# Build the current user message with optional image description
if image_description:
content = f"[Image analysis: {image_description}]"
if query_text:
content = f"{query_text}\n\n{content}"
# Attach the image directly to the current message so the (multimodal)
# agent model sees it in one call, instead of a separate describe pass.
if image_block:
content = [
{"type": "text", "text": query_text or "The user sent this image."},
image_block,
]
else:
content = query_text
messages_payload.append({"role": "user", "content": content})
@@ -121,6 +127,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,
@@ -190,20 +197,19 @@ async def stream_query():
image_key=image_key,
)
# If an image was uploaded, analyze it with the vision model
image_description = None
# If an image was uploaded, attach it to the agent request directly
image_block = None
if image_key:
try:
image_bytes, _ = await s3_get_image(image_key)
image_description = await analyze_user_image(image_bytes)
logging.info(f"Image analysis complete for {image_key}")
image_block = image_content_block(image_bytes)
except Exception as e:
logging.error(f"Failed to analyze image: {e}")
image_description = "[Image could not be analyzed]"
logging.error(f"Failed to load image {image_key}: {e}")
query_text = f"{query_text or ''}\n\n[An image was sent but could not be loaded.]".strip()
system_prompt = await _build_system_prompt_with_memories(str(user.id))
messages_payload = _build_messages_payload(
conversation, query_text or "", image_description, system_prompt=system_prompt
conversation, query_text or "", image_block, system_prompt=system_prompt
)
payload = {"messages": messages_payload}
agent_config = {
@@ -213,6 +219,18 @@ async def stream_query():
async def event_generator():
final_message = None
streamed = ""
last_thinking = 0.0
turn_has_content = False
# Timing / observability
t0 = time.monotonic()
llm_start_at: float | None = None
first_content_at: float | None = None
tool_timers: dict[str, float] = {} # run_id -> start time
tool_spans: list[dict] = []
n_tokens = 0
model_name: str | None = None
cancelled = False
async def persist(text: str) -> None:
await add_message_to_conversation(
@@ -222,6 +240,35 @@ async def stream_query():
user=user,
)
def _log_summary(*, error: str | None = None) -> None:
now = time.monotonic()
total_ms = round((now - t0) * 1000)
thinking_ms = (
round((first_content_at - llm_start_at) * 1000)
if llm_start_at and first_content_at
else None
)
ttft_ms = (
round((first_content_at - t0) * 1000) if first_content_at else None
)
logging.info(
json.dumps(
{
"event": "stream_request",
"conv": str(conversation.id),
"user": user.username,
"total_ms": total_ms,
"thinking_ms": thinking_ms,
"ttft_ms": ttft_ms,
"tokens": n_tokens,
"tools": tool_spans,
"model": model_name,
"cancelled": cancelled,
**({"error": error} if error else {}),
}
)
)
try:
async for event in main_agent.astream_events(
payload, version="v2", config=agent_config
@@ -229,17 +276,54 @@ async def stream_query():
event_type = event.get("event")
if event_type == "on_tool_start":
tool_timers[event["run_id"]] = time.monotonic()
yield f"data: {json.dumps({'type': 'tool_start', 'tool': event['name']})}\n\n"
elif event_type == "on_tool_end":
start = tool_timers.pop(event["run_id"], None)
tool_spans.append(
{
"tool": event["name"],
"ms": round((time.monotonic() - start) * 1000)
if start
else None,
}
)
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_start":
turn_has_content = False
if llm_start_at is None:
llm_start_at = time.monotonic()
if model_name is None:
model_name = event.get("metadata", {}).get(
"ls_model_name"
) or event.get("name")
elif event_type == "on_chat_model_stream":
chunk = event.get("data", {}).get("chunk")
content = getattr(chunk, "content", None)
if isinstance(content, str) and content:
if first_content_at is None:
first_content_at = time.monotonic()
turn_has_content = True
n_tokens += 1
streamed += content
yield f"data: {json.dumps({'type': 'content', 'delta': content})}\n\n"
elif not turn_has_content:
now = time.monotonic()
if now - last_thinking >= 1.0:
last_thinking = now
yield f"data: {json.dumps({'type': 'thinking'})}\n\n"
elif event_type == "on_chain_end":
output = event.get("data", {}).get("output")
@@ -252,27 +336,25 @@ async def stream_query():
final_message = content
except asyncio.CancelledError:
# The client disconnected mid-stream (e.g. a route change aborts the
# fetch). CancelledError is not an Exception, so without this branch
# persistence below would be skipped and the answer lost. Save what
# was generated in a detached task, then honor the cancellation.
cancelled = True
_log_summary()
text = final_message or streamed
if text:
_persist_in_background(persist(text))
raise
except Exception as e:
_log_summary(error=str(e))
yield f"data: {json.dumps({'type': 'error', 'message': str(e)})}\n\n"
final_message = final_message or streamed
if final_message:
# Persist via a detached, shielded task so a disconnect at this
# point (after generation completes) still commits the answer.
task = _persist_in_background(persist(final_message))
await asyncio.shield(task)
yield f"data: {json.dumps({'type': 'response', 'message': final_message})}\n\n"
else:
yield f"data: {json.dumps({'type': 'error', 'message': 'No response generated'})}\n\n"
_log_summary()
yield "data: [DONE]\n\n"
return await make_response(
@@ -298,10 +380,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,
+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)
+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):
+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."
+18 -14
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
@@ -200,37 +202,36 @@ async def webhook():
user=user,
)
# Analyze image if present
image_description = None
# Attach image directly to the agent request if present
image_block = None
if media_url:
try:
from utils.image_process import analyze_user_image
from utils.image_process import image_content_block
async with httpx.AsyncClient() as client:
resp = await client.get(media_url, timeout=30)
resp.raise_for_status()
image_bytes = resp.content
image_description = await analyze_user_image(image_bytes)
logger.info(f"Image analysis complete for iMessage from {from_number}")
image_block = image_content_block(resp.content)
except Exception as e:
logger.error(f"Failed to analyze iMessage image: {e}")
image_description = "[Image could not be analyzed]"
logger.error(f"Failed to load iMessage image: {e}")
content = f"{content or ''}\n\n[An image was sent but could not be loaded.]".strip()
# 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_description:
user_content = f"[Image analysis: {image_description}]"
if content:
user_content = f"{content}\n\n{user_content}"
if image_block:
user_content = [
{"type": "text", "text": content or "The user sent this image."},
image_block,
]
else:
user_content = content
messages_payload.append({"role": "user", "content": user_content})
@@ -238,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}")
+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"):
@@ -1,6 +1,7 @@
import { userService } from "./userService";
export type SSEEvent =
| { type: "thinking" }
| { type: "tool_start"; tool: string }
| { type: "tool_end"; tool: string }
| { type: "content"; delta: string }
@@ -12,7 +13,7 @@ export type SSEEventCallback = (event: SSEEvent) => void;
interface Message {
id: string;
text: string;
speaker: "user" | "simba";
speaker: "user" | "simba" | "tool";
created_at: string;
image_key?: string | null;
}
+6 -1
View File
@@ -53,6 +53,7 @@ export const ChatScreen = ({ setAuthenticated, isAdmin }: ChatScreenProps) => {
messages,
isLoading,
streaming,
thinking,
messagesLoading,
pendingImage,
setPendingImage,
@@ -249,7 +250,11 @@ export const ChatScreen = ({ setAuthenticated, isAdmin }: ChatScreenProps) => {
return <QuestionBubble key={index} text={msg.text} image_key={msg.image_key} />;
})}
{((isLoading && !streaming) || messagesLoading) && (
{thinking && !streaming && (
<ToolBubble text="Thinking..." pulse />
)}
{((isLoading && !streaming && !thinking) || messagesLoading) && (
<AnswerBubble text="" loading={true} />
)}
+15 -12
View File
@@ -1,16 +1,19 @@
import React from "react";
import { cn } from "../lib/utils";
export const ToolBubble = React.memo(({ text }: { text: string }) => (
<div className="flex justify-center message-enter">
<div
className={cn(
"inline-flex items-center gap-1.5 px-3 py-1 rounded-full",
"bg-leaf-pale border border-leaf-light/50",
"text-xs text-leaf-dark italic",
)}
>
{text}
export const ToolBubble = React.memo(
({ text, pulse = false }: { text: string; pulse?: boolean }) => (
<div className="flex justify-center message-enter">
<div
className={cn(
"inline-flex items-center gap-1.5 px-3 py-1 rounded-full",
"bg-leaf-pale border border-leaf-light/50",
"text-xs text-leaf-dark italic",
pulse && "animate-pulse",
)}
>
{text}
</div>
</div>
</div>
));
),
);
+11 -2
View File
@@ -52,6 +52,7 @@ export function useChat({
const [pendingImage, setPendingImage] = useState<File | null>(null);
const [error, setError] = useState<string | null>(null);
const [streaming, setStreaming] = useState(false);
const [thinking, setThinking] = useState(false);
const isMountedRef = useRef(true);
// True once tokens for the current assistant reply have started arriving, so
@@ -66,7 +67,9 @@ export function useChat({
isMountedRef.current = true;
return () => {
isMountedRef.current = false;
abortControllerRef.current?.abort();
// Intentionally do NOT abort on unmount — navigation away lets the
// stream finish in the background so the response gets persisted.
// stopGeneration() still aborts on explicit user cancel.
};
}, []);
@@ -186,15 +189,19 @@ export function useChat({
activeId,
(event) => {
if (!isMountedRef.current) return;
if (event.type === "tool_start") {
if (event.type === "thinking") {
setThinking(true);
} else if (event.type === "tool_start") {
// Any text after a tool call belongs in a fresh bubble.
streamingActiveRef.current = false;
setThinking(false);
const friendly =
TOOL_MESSAGES[event.tool] ?? `Using ${event.tool}...`;
updateMessages((prev) =>
prev.concat([{ text: friendly, speaker: "tool" }]),
);
} else if (event.type === "content") {
setThinking(false);
if (streamingActiveRef.current) {
updateMessages((prev) => {
const updated = [...prev];
@@ -258,6 +265,7 @@ export function useChat({
if (isMountedRef.current) {
setIsLoading(false);
setStreaming(false);
setThinking(false);
// Refresh after every send so the just-used conversation floats to
// the top of the sidebar (sorted by last activity) and picks up the
// auto-generated title on the first message.
@@ -290,6 +298,7 @@ export function useChat({
setMessages: updateMessages,
isLoading,
streaming,
thinking,
messagesLoading,
pendingImage,
setPendingImage,
+10 -40
View File
@@ -87,50 +87,20 @@ def _to_jpeg_bytes(file_bytes: bytes) -> bytes:
return buf.getvalue()
async def analyze_user_image(file_bytes: bytes) -> str:
"""Analyze an image uploaded by a user and return a text description.
def image_content_block(file_bytes: bytes) -> dict:
"""Build an OpenAI-style image_url content part from raw image bytes.
Uses llama-server (OpenAI-compatible API) with vision support.
Falls back to OpenAI if llama-server is not configured.
Normalizes to JPEG so HEIC/HEIF and other formats work, then encodes as a
base64 data URL suitable for a multimodal chat message.
"""
import base64
from openai import AsyncOpenAI
file_bytes = _to_jpeg_bytes(file_bytes)
llama_url = os.getenv("LLAMA_SERVER_URL")
if llama_url:
aclient = AsyncOpenAI(base_url=llama_url, api_key="not-needed")
model = os.getenv("LLAMA_MODEL_NAME", "llama-3.1-8b-instruct")
else:
aclient = AsyncOpenAI()
model = "gpt-4o-mini"
b64 = base64.b64encode(file_bytes).decode("utf-8")
response = await aclient.chat.completions.create(
model=model,
messages=[
{
"role": "system",
"content": "You are a helpful image analyst. Describe what you see in the image in detail. Be thorough but concise.",
},
{
"role": "user",
"content": [
{"type": "text", "text": "Please describe this image in detail."},
{
"type": "image_url",
"image_url": {
"url": f"data:image/jpeg;base64,{b64}",
},
},
],
},
],
)
return response.choices[0].message.content
jpeg_bytes = _to_jpeg_bytes(file_bytes)
b64 = base64.b64encode(jpeg_bytes).decode("utf-8")
return {
"type": "image_url",
"image_url": {"url": f"data:image/jpeg;base64,{b64}"},
}
if __name__ == "__main__":