diff --git a/blueprints/conversation/__init__.py b/blueprints/conversation/__init__.py index 563336e..5171176 100644 --- a/blueprints/conversation/__init__.py +++ b/blueprints/conversation/__init__.py @@ -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, diff --git a/blueprints/conversation/logic.py b/blueprints/conversation/logic.py index 52f15d7..437d45e 100644 --- a/blueprints/conversation/logic.py +++ b/blueprints/conversation/logic.py @@ -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) diff --git a/blueprints/conversation/models.py b/blueprints/conversation/models.py index 60b8cee..0b1fd15 100644 --- a/blueprints/conversation/models.py +++ b/blueprints/conversation/models.py @@ -11,6 +11,7 @@ from tortoise.models import Model class Speaker(enum.Enum): USER = "user" SIMBA = "simba" + TOOL = "tool" class Conversation(Model): diff --git a/blueprints/email/__init__.py b/blueprints/email/__init__.py index 121a740..1a4b7af 100644 --- a/blueprints/email/__init__.py +++ b/blueprints/email/__init__.py @@ -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." diff --git a/blueprints/imessage/__init__.py b/blueprints/imessage/__init__.py index a2d2277..f1008c5 100644 --- a/blueprints/imessage/__init__.py +++ b/blueprints/imessage/__init__.py @@ -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}") diff --git a/blueprints/whatsapp/__init__.py b/blueprints/whatsapp/__init__.py index 6824d2d..30faeea 100644 --- a/blueprints/whatsapp/__init__.py +++ b/blueprints/whatsapp/__init__.py @@ -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"): diff --git a/raggr-frontend/src/api/conversationService.ts b/raggr-frontend/src/api/conversationService.ts index 7baab48..543ee79 100644 --- a/raggr-frontend/src/api/conversationService.ts +++ b/raggr-frontend/src/api/conversationService.ts @@ -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; } diff --git a/raggr-frontend/src/components/ChatScreen.tsx b/raggr-frontend/src/components/ChatScreen.tsx index d74e054..c5642ce 100644 --- a/raggr-frontend/src/components/ChatScreen.tsx +++ b/raggr-frontend/src/components/ChatScreen.tsx @@ -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 ; })} - {((isLoading && !streaming) || messagesLoading) && ( + {thinking && !streaming && ( + + )} + + {((isLoading && !streaming && !thinking) || messagesLoading) && ( )} diff --git a/raggr-frontend/src/components/ToolBubble.tsx b/raggr-frontend/src/components/ToolBubble.tsx index 22b0e85..213ef47 100644 --- a/raggr-frontend/src/components/ToolBubble.tsx +++ b/raggr-frontend/src/components/ToolBubble.tsx @@ -1,16 +1,19 @@ import React from "react"; import { cn } from "../lib/utils"; -export const ToolBubble = React.memo(({ text }: { text: string }) => ( -
-
- {text} +export const ToolBubble = React.memo( + ({ text, pulse = false }: { text: string; pulse?: boolean }) => ( +
+
+ {text} +
-
-)); + ), +); diff --git a/raggr-frontend/src/hooks/useChat.ts b/raggr-frontend/src/hooks/useChat.ts index d0a6d71..cc00670 100644 --- a/raggr-frontend/src/hooks/useChat.ts +++ b/raggr-frontend/src/hooks/useChat.ts @@ -52,6 +52,7 @@ export function useChat({ const [pendingImage, setPendingImage] = useState(null); const [error, setError] = useState(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, diff --git a/utils/image_process.py b/utils/image_process.py index 534f566..bb5cafa 100644 --- a/utils/image_process.py +++ b/utils/image_process.py @@ -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__":