From c56ca3fda4dc36b6db989cbf178d2c01ff12c722 Mon Sep 17 00:00:00 2001 From: Sam Zhong Date: Mon, 1 Jun 2026 19:56:21 +0000 Subject: [PATCH 01/17] Add chat memory design spec Document the planned independent chat-session memory model with optional history-session attachment so implementation can follow the agreed LangGraph approach. Co-Authored-By: Claude Opus 4.7 (1M context) --- .../specs/2026-06-01-chat-memory-design.md | 108 ++++++++++++++++++ 1 file changed, 108 insertions(+) create mode 100644 docs/superpowers/specs/2026-06-01-chat-memory-design.md diff --git a/docs/superpowers/specs/2026-06-01-chat-memory-design.md b/docs/superpowers/specs/2026-06-01-chat-memory-design.md new file mode 100644 index 0000000..f81604d --- /dev/null +++ b/docs/superpowers/specs/2026-06-01-chat-memory-design.md @@ -0,0 +1,108 @@ +# Chat Short-Term Memory Design + +## Goal + +Add short-term memory to the chatbot using LangGraph so each chat conversation can remember prior turns. Chat memory should be independent from the existing document-processing history sessions, while still allowing a chat to attach a history session when the user wants to ask about processed files. + +## Current Context + +The current chat endpoint is stateless. The frontend chat page sends a user query and optional document markdown to `/api/chat/query`, and the backend forwards the prompt to the existing LLM service. The app already has a separate history/session system for processed files and results, but `ChatPage` does not currently receive or persist a session id. + +## Core Model + +Use two separate concepts: + +- `chat_session_id`: identifies one chatbot conversation and maps directly to LangGraph `thread_id`. +- `attached_session_id`: optionally points to an existing document/history session whose processed files should be available as request-time context. + +A chat session owns conversation memory. A document/history session owns processed files and extraction results. The chat session may reference one attached document session, but document state is not part of the chat identity. + +## User Experience + +From the chat page, the user can: + +1. Start a blank independent chat. +2. Select an existing history session in advanced mode and attach it to the current chat. +3. Upload/process new files from chat, creating or updating a document/history session, then attach that session to the chat. + +This supports multiple independent conversations about the same processed files without mixing chat histories. It also lets users reset chat memory without losing their document-processing history. + +## Backend Architecture + +Add a `ChatMemoryService` that owns LangGraph execution for chat. The FastAPI route remains responsible for request validation, model/provider parameters, session attachment lookup, and response formatting. + +`ChatMemoryService` responsibilities: + +- Build and invoke the LangGraph chat graph. +- Use `chat_session_id` as `configurable.thread_id`. +- Persist graph state with a Postgres checkpointer. +- Keep graph state small. +- Compose prompts from system instructions, remembered messages, and request-time document context. + +Graph state should include: + +- `messages`: prior user and assistant turns. +- `attachment`: lightweight metadata such as `attached_session_id`, selected document ids, hashes, timestamps, or titles. + +Graph state should not include full document markdown or large extraction payloads. Existing document services remain the source of truth for processed file content. + +## Data Flow + +1. The frontend creates or loads a `chat_session_id` for the active chat. +2. The frontend sends `/api/chat/query` with: + - `chat_session_id` + - user query + - optional `attached_session_id` + - existing model/provider fields +3. The backend validates the request. +4. If `attached_session_id` is present, the backend loads current document context from the existing history/session services. +5. The backend invokes `ChatMemoryService` with LangGraph config `thread_id = chat_session_id`. +6. LangGraph loads prior state from the Postgres checkpointer, appends the new user turn, calls the model, and checkpoints the updated state. +7. The backend returns the assistant response plus chat metadata. + +Document context can influence the current response, but it is not stored as full content in chat memory. + +## Attachment Behavior + +A chat can point at one attached document/history session at a time. Changing the attachment updates the chat's lightweight attachment metadata and affects future responses, but does not erase prior messages. + +If the user wants a clean conversation about different files, the frontend should make starting a new chat easy. The system should not silently merge unrelated document contexts into one conversation. + +## Error Handling + +- Missing `chat_session_id`: the frontend should create it; the backend should validate and return a clear error if absent. +- Invalid `attached_session_id`: return a user-friendly error and allow the chat to continue unattached. +- Oversized document context: do not checkpoint large content; summarize, truncate, or surface the existing context-window warning. +- Checkpointer failure: fail the chat request visibly instead of silently losing memory. +- Model failure: avoid creating ambiguous half-saved turns; retry behavior should be explicit in service tests. + +## Testing Strategy + +Backend tests: + +- Request validation for `chat_session_id` and optional `attached_session_id`. +- One `chat_session_id` remembers prior turns. +- Two different `chat_session_id`s do not share history. +- Attached document context influences a response without being stored as full markdown in graph state. +- Invalid or unavailable `attached_session_id` returns a clear error. + +Frontend tests: + +- Starting a blank chat creates an independent `chat_session_id`. +- Selecting an existing history session attaches it to the current chat. +- Starting a new chat does not reuse the prior chat's memory unless explicitly selected. + +Manual validation: + +- Process files in the workflow. +- Open Chat and attach the processed history session. +- Ask a document-specific question and a follow-up that depends on chat history. +- Start a blank chat and confirm prior chat memory does not leak. +- Reattach the same history session in a second chat and confirm document context is available without sharing conversation history. + +## Recommended Implementation Order + +1. Add backend request fields and `ChatMemoryService` with LangGraph checkpointing. +2. Add frontend `chat_session_id` lifecycle for blank independent chats. +3. Add history-session attachment UI and backend document-context lookup. +4. Add tests for memory isolation and attachment behavior. From 0fcbd5367e2d2cd11c41465753b113647bbc85eb Mon Sep 17 00:00:00 2001 From: Sam Zhong Date: Thu, 4 Jun 2026 16:14:39 +0000 Subject: [PATCH 02/17] feat: implement chat session management with history attachment - Added functions to create, retrieve, and reset chat session IDs using sessionStorage and in-memory storage. - Integrated chat session handling in ChatPage component, allowing users to attach history sessions and manage document uploads accordingly. - Updated UI to reflect changes in session management, including disabling document uploads when a history session is attached. - Enhanced error handling and user feedback for session-related actions. --- backend/api/chat/router.py | 100 +- backend/requirements.txt | 2 + backend/services/chat_memory/__init__.py | 3 + .../chat_memory/chat_memory_service.py | 347 +++++ .../services/chat_memory/document_context.py | 77 ++ backend/services/llm/anthropic.py | 5 +- backend/services/llm/llm_service.py | 1 + .../2026-06-01-chat-memory-implementation.md | 1123 +++++++++++++++++ frontend/components/ChatPage.tsx | 215 +++- frontend/utils/chatSession.ts | 79 ++ 10 files changed, 1874 insertions(+), 78 deletions(-) create mode 100644 backend/services/chat_memory/__init__.py create mode 100644 backend/services/chat_memory/chat_memory_service.py create mode 100644 backend/services/chat_memory/document_context.py create mode 100644 docs/superpowers/plans/2026-06-01-chat-memory-implementation.md create mode 100644 frontend/utils/chatSession.ts diff --git a/backend/api/chat/router.py b/backend/api/chat/router.py index 2652e34..bc3921f 100644 --- a/backend/api/chat/router.py +++ b/backend/api/chat/router.py @@ -1,73 +1,81 @@ """Chat API endpoint for support staff chatbot""" +from typing import Optional + from fastapi import APIRouter, Depends from fastapi.responses import JSONResponse -from pydantic import BaseModel -from typing import Optional +from pydantic import BaseModel, Field from core.auth import get_current_user -from services.llm.llm_service import LLMService +from services.chat_memory import ChatMemoryRequest, ChatMemoryService +from services.chat_memory.chat_memory_service import GENERIC_MODEL_ERROR_MESSAGE +from services.chat_memory.document_context import build_attached_session_context router = APIRouter(prefix="/api/chat", tags=["chat"]) -llm_service = LLMService() +chat_memory_service: Optional[ChatMemoryService] = None class ChatQueryRequest(BaseModel): - query: str + chat_session_id: str = Field(min_length=1) + query: str = Field(min_length=1) document_markdown: Optional[str] = None - model_type: str # "azure", "gemini", "anthropic", "llama", "azure-llama", "macbook" + attached_session_id: Optional[str] = None + model_type: str model_id: Optional[str] = None deployment: Optional[str] = None api_version: Optional[str] = None +def get_chat_memory_service() -> ChatMemoryService: + global chat_memory_service + if chat_memory_service is None: + chat_memory_service = ChatMemoryService() + return chat_memory_service + + @router.post("/query", dependencies=[Depends(get_current_user)]) -async def chat_query(request: ChatQueryRequest): - """ - Send a chat message with optional document context. +async def chat_query( + request: ChatQueryRequest, + current_user: dict = Depends(get_current_user), +): + document_context = request.document_markdown + attached_session_id = request.attached_session_id + attached_session_warning: Optional[str] = None - When document_markdown is provided, it is injected into the prompt so the - model can answer questions about the uploaded document. - """ - if request.document_markdown: - user_prompt = ( - "The following document has been uploaded by the user:\n\n" - f"\n{request.document_markdown}\n\n\n" - f"User question: {request.query}" + if request.attached_session_id: + attached_context = await build_attached_session_context( + user_id=current_user["id"], + attached_session_id=request.attached_session_id, ) - system_message = ( - "You are a helpful document assistant for Health Canada support staff. " - "Answer the user's question based on the provided document. " - "If the answer is not found in the document, say so clearly and offer " - "general guidance if possible." - ) - else: - user_prompt = request.query - system_message = ( - "You are a helpful assistant for Health Canada support staff. " - "Answer questions clearly and concisely." - ) - - result = await llm_service.generate_paragraph( - user_prompt=user_prompt, - model_type=request.model_type, - model_id=request.model_id, - deployment=request.deployment, - api_version=request.api_version, - max_tokens=4096, - temperature=0.3, - system_message=system_message, - ) + if attached_context is None: + attached_session_id = None + attached_session_warning = ( + "Attached session not found or has no processed document content. Continuing without attachment." + ) + else: + document_context = attached_context - if not result.get("success"): + try: + result = await get_chat_memory_service().invoke( + ChatMemoryRequest( + chat_session_id=request.chat_session_id, + query=request.query, + model_type=request.model_type, + model_id=request.model_id, + deployment=request.deployment, + api_version=request.api_version, + document_context=document_context, + attached_session_id=attached_session_id, + ) + ) + if attached_session_warning and isinstance(result, dict): + result["attached_session_warning"] = attached_session_warning + return result + except Exception: return JSONResponse( status_code=500, content={ "success": False, - "error": result.get( - "error", "The model call failed. Please try again." - ), + "error": GENERIC_MODEL_ERROR_MESSAGE, }, ) - - return {"success": True, "response": result.get("content", "")} diff --git a/backend/requirements.txt b/backend/requirements.txt index ac2de57..930e095 100644 --- a/backend/requirements.txt +++ b/backend/requirements.txt @@ -18,6 +18,8 @@ deepeval>=3.6.9 google-cloud-aiplatform>=1.122.0 langchain-openai>=1.0.1 langchain-google-vertexai>=3.0.1 +langgraph>=1.2.4 +langgraph-checkpoint-postgres>=3.1.0 pytest>=8.4.2 pytest-asyncio>=1.2.0 anthropic[vertex]>=0.39.0 diff --git a/backend/services/chat_memory/__init__.py b/backend/services/chat_memory/__init__.py new file mode 100644 index 0000000..202685b --- /dev/null +++ b/backend/services/chat_memory/__init__.py @@ -0,0 +1,3 @@ +from .chat_memory_service import ChatMemoryRequest, ChatMemoryService + +__all__ = ["ChatMemoryRequest", "ChatMemoryService"] diff --git a/backend/services/chat_memory/chat_memory_service.py b/backend/services/chat_memory/chat_memory_service.py new file mode 100644 index 0000000..2e76219 --- /dev/null +++ b/backend/services/chat_memory/chat_memory_service.py @@ -0,0 +1,347 @@ +from __future__ import annotations + +import asyncio +import atexit +import inspect +import weakref +from dataclasses import dataclass +from typing import Annotated, Any, Dict, List, Optional, TypedDict + +from services.llm.llm_service import LLMService + +LANGGRAPH_INSTALL_MESSAGE = ( + "LangGraph chat memory dependencies are unavailable. " + "Install backend requirements and retry." +) +GENERIC_MODEL_ERROR_MESSAGE = "The model call failed. Please try again." + + +class ModelProviderError(Exception): + pass + +try: + from langchain_core.messages import AIMessage, BaseMessage, HumanMessage + from langchain_core.runnables import RunnableConfig + + LANGCHAIN_MESSAGES_IMPORT_ERROR: Optional[ImportError] = None +except ImportError as exc: + AIMessage = BaseMessage = HumanMessage = RunnableConfig = Any # type: ignore[assignment] + LANGCHAIN_MESSAGES_IMPORT_ERROR = exc + +try: + from langgraph.graph import END, START, StateGraph + from langgraph.graph.message import add_messages + + LANGGRAPH_GRAPH_IMPORT_ERROR: Optional[ImportError] = None +except ImportError as exc: + END = "__end__" + START = "__start__" + StateGraph = None + LANGGRAPH_GRAPH_IMPORT_ERROR = exc + + def add_messages(existing: List[Any], new: List[Any]) -> List[Any]: + return [*(existing or []), *(new or [])] + + +try: + from langgraph.checkpoint.memory import InMemorySaver as LangGraphInMemorySaver + + LANGGRAPH_MEMORY_IMPORT_ERROR: Optional[ImportError] = None +except ImportError: + try: + from langgraph.checkpoint.memory import MemorySaver as LangGraphInMemorySaver + + LANGGRAPH_MEMORY_IMPORT_ERROR = None + except ImportError as exc: + LangGraphInMemorySaver = None + LANGGRAPH_MEMORY_IMPORT_ERROR = exc + + +@dataclass +class ChatMemoryRequest: + chat_session_id: str + query: str + model_type: str + model_id: Optional[str] = None + deployment: Optional[str] = None + api_version: Optional[str] = None + document_context: Optional[str] = None + attached_session_id: Optional[str] = None + + +class ChatState(TypedDict): + messages: Annotated[List[BaseMessage], add_messages] + attachment: Optional[Dict[str, Optional[str]]] + + +class ChatMemoryService: + def __init__( + self, + llm_service: Optional[LLMService] = None, + checkpointer: Optional[Any] = None, + use_memory_checkpointer: bool = False, + ): + self.llm_service = llm_service or LLMService() + self._use_memory_checkpointer = use_memory_checkpointer + self._async_checkpointer_context = None + self._did_register_atexit_close = False + self._initialization_lock = asyncio.Lock() + self._session_locks: weakref.WeakValueDictionary[str, asyncio.Lock] = weakref.WeakValueDictionary() + self._ensure_langgraph_dependencies(require_in_memory=use_memory_checkpointer) + self.checkpointer = checkpointer + self.graph = self._build_graph(checkpointer) if checkpointer is not None else None + + def _ensure_langgraph_dependencies(self, require_in_memory: bool = False) -> None: + import_errors: List[str] = [] + + if LANGCHAIN_MESSAGES_IMPORT_ERROR is not None: + import_errors.append(f"langchain_core.messages: {LANGCHAIN_MESSAGES_IMPORT_ERROR}") + if LANGGRAPH_GRAPH_IMPORT_ERROR is not None: + import_errors.append(f"langgraph.graph: {LANGGRAPH_GRAPH_IMPORT_ERROR}") + if require_in_memory and LANGGRAPH_MEMORY_IMPORT_ERROR is not None: + import_errors.append( + f"langgraph.checkpoint.memory: {LANGGRAPH_MEMORY_IMPORT_ERROR}" + ) + + if import_errors: + details = "; ".join(import_errors) + raise RuntimeError(f"{LANGGRAPH_INSTALL_MESSAGE} Import errors: {details}") + + def _get_async_postgres_saver_class(self) -> Any: + from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver + + return AsyncPostgresSaver + + async def _build_checkpointer(self) -> Any: + if self._use_memory_checkpointer: + return self._build_in_memory_checkpointer() + + try: + AsyncPostgresSaver = self._get_async_postgres_saver_class() + except Exception as exc: + raise RuntimeError( + f"Failed to import LangGraph Postgres checkpointer: {exc}" + ) from exc + + try: + from models.base import DATABASE_URL + + saver_context = AsyncPostgresSaver.from_conn_string(DATABASE_URL) + self._async_checkpointer_context = saver_context + saver = await saver_context.__aenter__() + self._register_checkpointer_cleanup() + if hasattr(saver, "setup"): + setup_result = saver.setup() + if inspect.isawaitable(setup_result): + await setup_result + return saver + except Exception as exc: + await self._close_async_checkpointer_context() + raise RuntimeError( + f"Failed to initialize LangGraph Postgres checkpointer: {exc}" + ) from exc + + def _register_checkpointer_cleanup(self) -> None: + if self._did_register_atexit_close: + return + atexit.register(self._close_checkpointer_context_at_exit) + self._did_register_atexit_close = True + + def _close_checkpointer_context_at_exit(self) -> None: + if self._async_checkpointer_context is None: + return + try: + asyncio.run(self._close_async_checkpointer_context()) + except Exception: + pass + + async def _close_async_checkpointer_context(self) -> None: + if self._async_checkpointer_context is None: + return + try: + await self._async_checkpointer_context.__aexit__(None, None, None) + except Exception: + pass + finally: + self._async_checkpointer_context = None + + def _build_in_memory_checkpointer(self) -> Any: + if LangGraphInMemorySaver is None: + detail = ( + str(LANGGRAPH_MEMORY_IMPORT_ERROR) + if LANGGRAPH_MEMORY_IMPORT_ERROR is not None + else "LangGraph in-memory saver is unavailable." + ) + raise RuntimeError( + f"{LANGGRAPH_INSTALL_MESSAGE} Unable to import LangGraph in-memory saver: {detail}" + ) from LANGGRAPH_MEMORY_IMPORT_ERROR + return LangGraphInMemorySaver() + + def _build_graph(self, checkpointer: Any) -> Any: + self._ensure_langgraph_dependencies() + + graph = StateGraph(ChatState) + graph.add_node("generate", self._generate_response) + graph.add_edge(START, "generate") + graph.add_edge("generate", END) + return graph.compile(checkpointer=checkpointer) + + async def _ensure_graph(self) -> Any: + if self.graph is not None: + return self.graph + + async with self._initialization_lock: + if self.graph is not None: + return self.graph + + checkpointer = self.checkpointer or await self._build_checkpointer() + self.checkpointer = checkpointer + self.graph = self._build_graph(checkpointer) + return self.graph + + def _get_session_lock(self, chat_session_id: str) -> asyncio.Lock: + lock = self._session_locks.get(chat_session_id) + if lock is None: + lock = asyncio.Lock() + self._session_locks[chat_session_id] = lock + return lock + + async def invoke(self, request: ChatMemoryRequest) -> Dict[str, Any]: + config = { + "configurable": { + "thread_id": request.chat_session_id, + "model_type": request.model_type, + "model_id": request.model_id, + "deployment": request.deployment, + "api_version": request.api_version, + "document_context": request.document_context, + } + } + attachment = self._build_attachment(request.attached_session_id) + lock = self._get_session_lock(request.chat_session_id) + graph = await self._ensure_graph() + + try: + async with lock: + starting_state = await graph.aget_state(config) + state_config = getattr(starting_state, "config", None) or config + prior_messages = list(starting_state.values.get("messages", [])) + next_messages = [*prior_messages, HumanMessage(content=request.query)] + result = await self._generate_response( + { + "messages": next_messages, + "attachment": attachment, + }, + config=config, + ) + ai_messages = result.get("messages", []) + persisted_config = await graph.aupdate_state( + state_config, + { + "messages": [*next_messages, *ai_messages], + "attachment": attachment, + }, + as_node="generate", + ) + persisted_state = await graph.aget_state(persisted_config) + response = self._extract_ai_response(persisted_state.values.get("messages", [])) + persisted_attachment = persisted_state.values.get("attachment") or {} + return { + "success": True, + "response": response, + "chat_session_id": request.chat_session_id, + "attached_session_id": persisted_attachment.get("attached_session_id"), + } + except ModelProviderError: + return { + "success": False, + "response": "", + "chat_session_id": request.chat_session_id, + "attached_session_id": request.attached_session_id, + "error": GENERIC_MODEL_ERROR_MESSAGE, + } + + async def _generate_response( + self, + state: ChatState, + config: Optional[RunnableConfig] = None, + ) -> Dict[str, Any]: + configurable = (config or {}).get("configurable", {}) + prompt = self._build_user_prompt( + messages=state.get("messages", []), + document_context=configurable.get("document_context"), + ) + system_message = self._build_system_message( + has_document_context=bool(configurable.get("document_context")) + ) + + try: + result = await self.llm_service.generate_paragraph( + user_prompt=prompt, + model_type=configurable["model_type"], + model_id=configurable.get("model_id"), + deployment=configurable.get("deployment"), + api_version=configurable.get("api_version"), + max_tokens=4096, + temperature=0.3, + system_message=system_message, + ) + except Exception as exc: + raise ModelProviderError(GENERIC_MODEL_ERROR_MESSAGE) from exc + if not result.get("success"): + raise ModelProviderError(str(result.get("error", GENERIC_MODEL_ERROR_MESSAGE))) + + return {"messages": [AIMessage(content=str(result.get("content", "")))]} + + def _build_user_prompt( + self, + messages: List[BaseMessage], + document_context: Optional[str], + ) -> str: + history_lines: List[str] = [] + current_query = "" + + for index, message in enumerate(messages): + is_last = index == len(messages) - 1 + if isinstance(message, HumanMessage): + if is_last: + current_query = str(message.content) + else: + history_lines.append(f"User: {message.content}") + elif isinstance(message, AIMessage): + history_lines.append(f"Assistant: {message.content}") + + sections: List[str] = [] + if history_lines: + sections.append("Conversation so far:\n" + "\n".join(history_lines)) + if document_context: + sections.append( + "Current document context for this request:\n" + f"{document_context}\n\n" + "Use this context when it is relevant. Do not assume it remains available in future turns unless it is provided again." + ) + sections.append(f"User question: {current_query}") + return "\n\n".join(sections) + + def _build_system_message(self, has_document_context: bool) -> str: + if has_document_context: + return ( + "You are a helpful document assistant for Health Canada support staff. " + "Answer the user's question using the current document context and remembered conversation. " + "If the answer is not found in the document context, say so clearly and offer general guidance if possible." + ) + return ( + "You are a helpful assistant for Health Canada support staff. " + "Answer questions clearly and concisely using remembered conversation when relevant." + ) + + def _build_attachment(self, attached_session_id: Optional[str]) -> Optional[Dict[str, Optional[str]]]: + if attached_session_id is None: + return None + return {"attached_session_id": attached_session_id} + + def _extract_ai_response(self, messages: List[BaseMessage]) -> str: + for message in reversed(messages): + if isinstance(message, AIMessage): + return str(message.content) + return "" diff --git a/backend/services/chat_memory/document_context.py b/backend/services/chat_memory/document_context.py new file mode 100644 index 0000000..a4bbaf8 --- /dev/null +++ b/backend/services/chat_memory/document_context.py @@ -0,0 +1,77 @@ +from typing import Optional + +from services.document.organized_file_service import get_organized_file_service +from services.session.session_service import get_session_service + +MAX_ATTACHED_SESSION_CONTEXT_CHARS = 200_000 + + +def _build_document_block(filename: str, file_hash: str, markdown: str) -> str: + return ( + f'\n' + f"{markdown}\n" + "" + ) + + +async def build_attached_session_context( + user_id: str, + attached_session_id: Optional[str], + session_service=None, + file_service=None, +) -> Optional[str]: + if not attached_session_id: + return None + + session_service = session_service or get_session_service() + file_service = file_service or get_organized_file_service() + + session = session_service.get_session(user_id, attached_session_id) + if session is None: + return None + + document_blocks = [] + current_length = 0 + separator = "\n\n" + + for document in session.documents: + processor_used = document.processor_used or "azure_doc_intelligence" + markdown = await file_service.get_processed_content( + document.file_hash, processor_used + ) + if not markdown: + continue + + block = _build_document_block(document.filename, document.file_hash, markdown) + separator_length = len(separator) if document_blocks else 0 + remaining_budget = MAX_ATTACHED_SESSION_CONTEXT_CHARS - current_length - separator_length + if remaining_budget <= 0: + break + + if len(block) > remaining_budget: + if document_blocks: + break + + opening_tag = f'\n' + closing_tag = "\n" + max_markdown_length = remaining_budget - len(opening_tag) - len(closing_tag) + if max_markdown_length <= 0: + break + block = _build_document_block( + document.filename, + document.file_hash, + markdown[:max_markdown_length], + ) + if len(block) > remaining_budget: + block = block[:remaining_budget] + if not block.endswith(closing_tag): + content_budget = max(0, remaining_budget - len(opening_tag) - len(closing_tag)) + block = opening_tag + markdown[:content_budget] + closing_tag + + document_blocks.append(block) + current_length += separator_length + len(block) + + if current_length >= MAX_ATTACHED_SESSION_CONTEXT_CHARS: + break + + return "\n\n".join(document_blocks) if document_blocks else None diff --git a/backend/services/llm/anthropic.py b/backend/services/llm/anthropic.py index 79f34c0..e872d0b 100644 --- a/backend/services/llm/anthropic.py +++ b/backend/services/llm/anthropic.py @@ -394,6 +394,7 @@ async def generate_paragraph_with_anthropic( model_id: Optional[str] = None, max_tokens: int = 2048, temperature: float = 0.0, + system_message: Optional[str] = None, project_id_override: Optional[str] = None, location_override: Optional[str] = None, service_account_path_override: Optional[Path] = None, @@ -402,7 +403,9 @@ async def generate_paragraph_with_anthropic( used_model_id = model_id or "claude-sonnet-4-5@20250929" # System message for paragraph generation - system = "You are a scientific writing assistant. Your task is to synthesize extracted information into a cohesive, well-structured paragraph while maintaining complete accuracy. Follow the instructions exactly and preserve all factual details from the provided entities." + system = system_message or ( + "You are a scientific writing assistant. Your task is to synthesize extracted information into a cohesive, well-structured paragraph while maintaining complete accuracy. Follow the instructions exactly and preserve all factual details from the provided entities." + ) messages = [{"role": "user", "content": user_prompt}] diff --git a/backend/services/llm/llm_service.py b/backend/services/llm/llm_service.py index aa62465..64bdb4e 100644 --- a/backend/services/llm/llm_service.py +++ b/backend/services/llm/llm_service.py @@ -419,6 +419,7 @@ async def generate_paragraph( model_id, max_tokens, temperature, + system_message=system_message, ) self._record_session_metrics(session_id, "gcp", result) return result diff --git a/docs/superpowers/plans/2026-06-01-chat-memory-implementation.md b/docs/superpowers/plans/2026-06-01-chat-memory-implementation.md new file mode 100644 index 0000000..3847782 --- /dev/null +++ b/docs/superpowers/plans/2026-06-01-chat-memory-implementation.md @@ -0,0 +1,1123 @@ +# Chat Memory Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** Add independent LangGraph-backed memory for each chatbot conversation, with optional attachment to an existing document/history session. + +**Architecture:** The frontend creates a durable `chat_session_id` for each chat. The backend uses that ID as LangGraph's `thread_id`, stores short-term message state through a checkpointer, and loads attached document context from existing session/document services at request time. Chat memory and document-processing sessions remain separate; `attached_session_id` is only a reference. + +**Tech Stack:** FastAPI, Pydantic, Python LangGraph, LangGraph Postgres checkpointer, existing `LLMService`, React/TypeScript, existing authenticated fetch/session APIs. + +--- + +## File Structure + +Backend files: + +- Modify: `backend/requirements.txt` + - Add `langgraph` and `langgraph-checkpoint-postgres`. +- Create: `backend/services/chat_memory/__init__.py` + - Export `ChatMemoryService`, `ChatMessage`, and `ChatMemoryRequest`. +- Create: `backend/services/chat_memory/chat_memory_service.py` + - Build the LangGraph chat graph, use `chat_session_id` as `thread_id`, and call `LLMService.generate_paragraph`. +- Create: `backend/services/chat_memory/document_context.py` + - Load markdown for an optional `attached_session_id` using existing `SessionService` and `OrganizedFileService` methods. +- Modify: `backend/api/chat/router.py` + - Accept `chat_session_id` and optional `attached_session_id`; delegate to `ChatMemoryService`. +- Create: `backend/tests/test_chat_memory_service.py` + - Unit tests for independent memory and prompt construction using an in-memory LangGraph checkpointer. +- Create: `backend/tests/test_chat_router.py` + - Unit tests for request validation and route delegation. + +Frontend files: + +- Create: `frontend/utils/chatSession.ts` + - Generate, store, reset, and retrieve independent chat session IDs. +- Modify: `frontend/components/ChatPage.tsx` + - Send `chat_session_id`, optionally send `attached_session_id`, support starting a new chat, and add a simple history-session attachment selector. +- Reuse: `frontend/types/session.ts` + - Use existing `SessionSummary` type for the selector. + +No database model table is needed for chat messages in this first implementation. LangGraph checkpoint tables are owned by the LangGraph checkpointer package. + +--- + +### Task 1: Add LangGraph dependencies + +**Files:** +- Modify: `backend/requirements.txt:19-29` + +- [ ] **Step 1: Add dependency lines** + +Edit `backend/requirements.txt` so the LangChain/LangGraph block includes these lines: + +```txt +langchain-openai>=1.0.1 +langchain-google-vertexai>=3.0.1 +langgraph>=1.0.0 +langgraph-checkpoint-postgres>=3.0.0 +``` + +- [ ] **Step 2: Install backend dependencies locally** + +Run: + +```bash +cd backend && python -m pip install -r requirements.txt +``` + +Expected: installation completes with `langgraph` and `langgraph-checkpoint-postgres` installed. + +- [ ] **Step 3: Commit** + +```bash +git add backend/requirements.txt +git commit -m "Add LangGraph chat memory dependencies" +``` + +--- + +### Task 2: Write failing ChatMemoryService tests + +**Files:** +- Create: `backend/tests/test_chat_memory_service.py` + +- [ ] **Step 1: Create the failing tests** + +Create `backend/tests/test_chat_memory_service.py` with: + +```python +import sys +from pathlib import Path + +import pytest + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) + +from services.chat_memory.chat_memory_service import ChatMemoryRequest, ChatMemoryService + + +class FakeLLMService: + def __init__(self): + self.prompts = [] + + async def generate_paragraph(self, **kwargs): + self.prompts.append(kwargs) + user_prompt = kwargs["user_prompt"] + if "What did I ask before?" in user_prompt: + return {"success": True, "content": "You asked about licensing."} + return {"success": True, "content": "First answer."} + + +@pytest.mark.asyncio +async def test_chat_session_remembers_prior_turns(): + fake_llm = FakeLLMService() + service = ChatMemoryService(llm_service=fake_llm, use_memory_checkpointer=True) + + first = await service.invoke( + ChatMemoryRequest( + chat_session_id="chat-a", + query="Tell me about licensing.", + model_type="azure", + ) + ) + second = await service.invoke( + ChatMemoryRequest( + chat_session_id="chat-a", + query="What did I ask before?", + model_type="azure", + ) + ) + + assert first["success"] is True + assert second["success"] is True + assert second["response"] == "You asked about licensing." + assert "Tell me about licensing." in fake_llm.prompts[1]["user_prompt"] + + +@pytest.mark.asyncio +async def test_chat_sessions_do_not_share_history(): + fake_llm = FakeLLMService() + service = ChatMemoryService(llm_service=fake_llm, use_memory_checkpointer=True) + + await service.invoke( + ChatMemoryRequest( + chat_session_id="chat-a", + query="Tell me about licensing.", + model_type="azure", + ) + ) + await service.invoke( + ChatMemoryRequest( + chat_session_id="chat-b", + query="What did I ask before?", + model_type="azure", + ) + ) + + chat_b_prompt = fake_llm.prompts[-1]["user_prompt"] + assert "Tell me about licensing." not in chat_b_prompt + + +@pytest.mark.asyncio +async def test_document_context_is_injected_but_not_saved_as_message(): + fake_llm = FakeLLMService() + service = ChatMemoryService(llm_service=fake_llm, use_memory_checkpointer=True) + + await service.invoke( + ChatMemoryRequest( + chat_session_id="chat-doc", + query="Summarize this.", + model_type="azure", + document_context="Large protocol markdown", + attached_session_id="session-1", + ) + ) + await service.invoke( + ChatMemoryRequest( + chat_session_id="chat-doc", + query="What did I ask?", + model_type="azure", + ) + ) + + first_prompt = fake_llm.prompts[0]["user_prompt"] + second_prompt = fake_llm.prompts[1]["user_prompt"] + assert "Large protocol markdown" in first_prompt + assert "Large protocol markdown" not in second_prompt + assert "Summarize this." in second_prompt +``` + +- [ ] **Step 2: Run tests to verify they fail** + +Run: + +```bash +cd backend && python -m pytest tests/test_chat_memory_service.py -v +``` + +Expected: FAIL with `ModuleNotFoundError: No module named 'services.chat_memory'`. + +--- + +### Task 3: Implement ChatMemoryService + +**Files:** +- Create: `backend/services/chat_memory/__init__.py` +- Create: `backend/services/chat_memory/chat_memory_service.py` +- Test: `backend/tests/test_chat_memory_service.py` + +- [ ] **Step 1: Create package export** + +Create `backend/services/chat_memory/__init__.py`: + +```python +from .chat_memory_service import ChatMemoryRequest, ChatMemoryService + +__all__ = ["ChatMemoryRequest", "ChatMemoryService"] +``` + +- [ ] **Step 2: Create the service implementation** + +Create `backend/services/chat_memory/chat_memory_service.py`: + +```python +from dataclasses import dataclass +from typing import Any, Dict, List, Optional, TypedDict + +from langchain_core.messages import AIMessage, BaseMessage, HumanMessage +from langgraph.checkpoint.memory import InMemorySaver +from langgraph.graph import END, START, StateGraph + +from services.llm.llm_service import LLMService + + +@dataclass +class ChatMemoryRequest: + chat_session_id: str + query: str + model_type: str + model_id: Optional[str] = None + deployment: Optional[str] = None + api_version: Optional[str] = None + document_context: Optional[str] = None + attached_session_id: Optional[str] = None + + +class ChatState(TypedDict): + messages: List[BaseMessage] + query: str + model_type: str + model_id: Optional[str] + deployment: Optional[str] + api_version: Optional[str] + document_context: Optional[str] + attached_session_id: Optional[str] + + +class ChatMemoryService: + def __init__( + self, + llm_service: Optional[LLMService] = None, + checkpointer: Optional[Any] = None, + use_memory_checkpointer: bool = False, + ): + self.llm_service = llm_service or LLMService() + self.checkpointer = checkpointer or self._build_checkpointer(use_memory_checkpointer) + self.graph = self._build_graph() + + def _build_checkpointer(self, use_memory_checkpointer: bool): + if use_memory_checkpointer: + return InMemorySaver() + + try: + from langgraph.checkpoint.postgres import PostgresSaver + from models.base import DATABASE_URL + + saver = PostgresSaver.from_conn_string(DATABASE_URL) + saver.setup() + return saver + except Exception as exc: + print(f"[CHAT_MEMORY] Falling back to in-memory checkpointer: {exc}") + return InMemorySaver() + + def _build_graph(self): + graph = StateGraph(ChatState) + graph.add_node("generate", self._generate_response) + graph.add_edge(START, "generate") + graph.add_edge("generate", END) + return graph.compile(checkpointer=self.checkpointer) + + async def invoke(self, request: ChatMemoryRequest) -> Dict[str, Any]: + state: ChatState = { + "messages": [HumanMessage(content=request.query)], + "query": request.query, + "model_type": request.model_type, + "model_id": request.model_id, + "deployment": request.deployment, + "api_version": request.api_version, + "document_context": request.document_context, + "attached_session_id": request.attached_session_id, + } + config = {"configurable": {"thread_id": request.chat_session_id}} + result = await self.graph.ainvoke(state, config=config) + messages = result.get("messages", []) + response = "" + for message in reversed(messages): + if isinstance(message, AIMessage): + response = str(message.content) + break + return { + "success": True, + "response": response, + "chat_session_id": request.chat_session_id, + "attached_session_id": request.attached_session_id, + } + + async def _generate_response(self, state: ChatState) -> Dict[str, Any]: + prompt = self._build_user_prompt( + messages=state.get("messages", []), + current_query=state["query"], + document_context=state.get("document_context"), + ) + system_message = self._build_system_message(bool(state.get("document_context"))) + result = await self.llm_service.generate_paragraph( + user_prompt=prompt, + model_type=state["model_type"], + model_id=state.get("model_id"), + deployment=state.get("deployment"), + api_version=state.get("api_version"), + max_tokens=4096, + temperature=0.3, + system_message=system_message, + ) + if not result.get("success"): + raise RuntimeError(result.get("error", "The model call failed. Please try again.")) + return {"messages": [AIMessage(content=result.get("content", ""))]} + + def _build_user_prompt( + self, + messages: List[BaseMessage], + current_query: str, + document_context: Optional[str], + ) -> str: + history_lines = [] + for message in messages[:-1]: + if isinstance(message, HumanMessage): + history_lines.append(f"User: {message.content}") + elif isinstance(message, AIMessage): + history_lines.append(f"Assistant: {message.content}") + + sections = [] + if history_lines: + sections.append("Conversation so far:\n" + "\n".join(history_lines)) + if document_context: + sections.append( + "Current document context for this request:\n" + f"{document_context}\n\n" + "Use this context when it is relevant. Do not assume it remains attached in future turns unless it is provided again." + ) + sections.append(f"User question: {current_query}") + return "\n\n".join(sections) + + def _build_system_message(self, has_document_context: bool) -> str: + if has_document_context: + return ( + "You are a helpful document assistant for Health Canada support staff. " + "Answer the user's question using the current document context and the remembered conversation. " + "If the answer is not found in the document context, say so clearly and offer general guidance if possible." + ) + return ( + "You are a helpful assistant for Health Canada support staff. " + "Answer questions clearly and concisely using the remembered conversation when relevant." + ) +``` + +- [ ] **Step 3: Run service tests** + +Run: + +```bash +cd backend && python -m pytest tests/test_chat_memory_service.py -v +``` + +Expected: PASS for all three tests. + +- [ ] **Step 4: Commit** + +```bash +git add backend/services/chat_memory/__init__.py backend/services/chat_memory/chat_memory_service.py backend/tests/test_chat_memory_service.py +git commit -m "Add LangGraph chat memory service" +``` + +--- + +### Task 4: Add document-session context loader tests + +**Files:** +- Create: `backend/services/chat_memory/document_context.py` +- Modify: `backend/tests/test_chat_memory_service.py` + +- [ ] **Step 1: Add tests for attached session context** + +Append to `backend/tests/test_chat_memory_service.py`: + +```python +from services.chat_memory.document_context import build_attached_session_context + + +class FakeDocument: + def __init__(self, file_hash, filename, processor_used="azure_doc_intelligence"): + self.file_hash = file_hash + self.filename = filename + self.processor_used = processor_used + + +class FakeSession: + def __init__(self): + self.documents = [ + FakeDocument("hash-1", "protocol.pdf"), + FakeDocument("hash-2", "consent.pdf"), + ] + + +class FakeSessionService: + def get_session(self, user_id, session_id): + if session_id == "missing-session": + return None + return FakeSession() + + +class FakeFileService: + async def get_processed_content(self, file_hash, processor): + return f"markdown for {file_hash} via {processor}" + + +@pytest.mark.asyncio +async def test_build_attached_session_context_loads_session_documents(): + context = await build_attached_session_context( + user_id="user-1", + attached_session_id="session-1", + session_service=FakeSessionService(), + file_service=FakeFileService(), + ) + + assert '' in context + assert "markdown for hash-1 via azure_doc_intelligence" in context + assert '' in context + + +@pytest.mark.asyncio +async def test_build_attached_session_context_returns_none_for_missing_session(): + context = await build_attached_session_context( + user_id="user-1", + attached_session_id="missing-session", + session_service=FakeSessionService(), + file_service=FakeFileService(), + ) + + assert context is None +``` + +- [ ] **Step 2: Run tests to verify they fail** + +Run: + +```bash +cd backend && python -m pytest tests/test_chat_memory_service.py -v +``` + +Expected: FAIL with `ModuleNotFoundError` or missing `build_attached_session_context`. + +--- + +### Task 5: Implement attached session context loading + +**Files:** +- Create: `backend/services/chat_memory/document_context.py` +- Test: `backend/tests/test_chat_memory_service.py` + +- [ ] **Step 1: Create context loader** + +Create `backend/services/chat_memory/document_context.py`: + +```python +from typing import Optional + +from services.document.organized_file_service import get_organized_file_service +from services.session.session_service import get_session_service + + +async def build_attached_session_context( + user_id: str, + attached_session_id: Optional[str], + session_service=None, + file_service=None, +) -> Optional[str]: + if not attached_session_id: + return None + + session_service = session_service or get_session_service() + file_service = file_service or get_organized_file_service() + session = session_service.get_session(user_id, attached_session_id) + if session is None: + return None + + document_blocks = [] + for doc in session.documents: + processor_used = doc.processor_used or "azure_doc_intelligence" + markdown = await file_service.get_processed_content(doc.file_hash, processor_used) + if not markdown: + continue + document_blocks.append( + f'\n' + f"{markdown}\n" + "" + ) + + return "\n\n".join(document_blocks) if document_blocks else None +``` + +- [ ] **Step 2: Run tests** + +Run: + +```bash +cd backend && python -m pytest tests/test_chat_memory_service.py -v +``` + +Expected: PASS. + +- [ ] **Step 3: Commit** + +```bash +git add backend/services/chat_memory/document_context.py backend/tests/test_chat_memory_service.py +git commit -m "Load attached session context for chat" +``` + +--- + +### Task 6: Write failing chat router tests + +**Files:** +- Create: `backend/tests/test_chat_router.py` + +- [ ] **Step 1: Create route tests** + +Create `backend/tests/test_chat_router.py`: + +```python +import sys +from pathlib import Path + +import pytest + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) + +from api.chat import router as chat_router +from api.chat.router import ChatQueryRequest, chat_query + + +class FakeChatMemoryService: + def __init__(self): + self.requests = [] + + async def invoke(self, request): + self.requests.append(request) + return { + "success": True, + "response": "remembered response", + "chat_session_id": request.chat_session_id, + "attached_session_id": request.attached_session_id, + } + + +@pytest.mark.asyncio +async def test_chat_query_delegates_to_memory_service(monkeypatch): + fake_service = FakeChatMemoryService() + monkeypatch.setattr(chat_router, "chat_memory_service", fake_service) + monkeypatch.setattr(chat_router, "build_attached_session_context", lambda **kwargs: None) + + response = await chat_query( + ChatQueryRequest( + chat_session_id="chat-1", + query="hello", + model_type="azure", + ), + current_user={"id": "user-1"}, + ) + + assert response["success"] is True + assert response["response"] == "remembered response" + assert fake_service.requests[0].chat_session_id == "chat-1" + assert fake_service.requests[0].query == "hello" + + +@pytest.mark.asyncio +async def test_chat_query_uses_attached_session_context(monkeypatch): + fake_service = FakeChatMemoryService() + + async def fake_context(**kwargs): + return 'protocol markdown' + + monkeypatch.setattr(chat_router, "chat_memory_service", fake_service) + monkeypatch.setattr(chat_router, "build_attached_session_context", fake_context) + + await chat_query( + ChatQueryRequest( + chat_session_id="chat-1", + attached_session_id="session-1", + query="summarize", + model_type="azure", + ), + current_user={"id": "user-1"}, + ) + + assert fake_service.requests[0].attached_session_id == "session-1" + assert "protocol markdown" in fake_service.requests[0].document_context +``` + +- [ ] **Step 2: Run route tests to verify they fail** + +Run: + +```bash +cd backend && python -m pytest tests/test_chat_router.py -v +``` + +Expected: FAIL because `ChatQueryRequest` does not accept `chat_session_id` and `chat_query` does not accept `current_user` directly. + +--- + +### Task 7: Update chat route to use ChatMemoryService + +**Files:** +- Modify: `backend/api/chat/router.py:3-73` +- Test: `backend/tests/test_chat_router.py` + +- [ ] **Step 1: Replace route implementation** + +Update `backend/api/chat/router.py` to: + +```python +"""Chat API endpoint for support staff chatbot""" + +from typing import Optional + +from fastapi import APIRouter, Depends +from fastapi.responses import JSONResponse +from pydantic import BaseModel, Field + +from core.auth import get_current_user +from services.chat_memory import ChatMemoryRequest, ChatMemoryService +from services.chat_memory.document_context import build_attached_session_context + +router = APIRouter(prefix="/api/chat", tags=["chat"]) +chat_memory_service = ChatMemoryService() + + +class ChatQueryRequest(BaseModel): + chat_session_id: str = Field(min_length=1) + query: str = Field(min_length=1) + document_markdown: Optional[str] = None + attached_session_id: Optional[str] = None + model_type: str + model_id: Optional[str] = None + deployment: Optional[str] = None + api_version: Optional[str] = None + + +@router.post("/query", dependencies=[Depends(get_current_user)]) +async def chat_query( + request: ChatQueryRequest, + current_user: dict = Depends(get_current_user), +): + document_context = request.document_markdown + if request.attached_session_id: + attached_context = await build_attached_session_context( + user_id=current_user["id"], + attached_session_id=request.attached_session_id, + ) + if attached_context is None: + return JSONResponse( + status_code=404, + content={ + "success": False, + "error": "Attached session not found or has no processed document content.", + }, + ) + document_context = attached_context + + try: + return await chat_memory_service.invoke( + ChatMemoryRequest( + chat_session_id=request.chat_session_id, + attached_session_id=request.attached_session_id, + query=request.query, + document_context=document_context, + model_type=request.model_type, + model_id=request.model_id, + deployment=request.deployment, + api_version=request.api_version, + ) + ) + except Exception as exc: + return JSONResponse( + status_code=500, + content={ + "success": False, + "error": str(exc) or "The model call failed. Please try again.", + }, + ) +``` + +- [ ] **Step 2: Run route tests** + +Run: + +```bash +cd backend && python -m pytest tests/test_chat_router.py -v +``` + +Expected: PASS. + +- [ ] **Step 3: Run all chat memory backend tests** + +Run: + +```bash +cd backend && python -m pytest tests/test_chat_memory_service.py tests/test_chat_router.py -v +``` + +Expected: PASS. + +- [ ] **Step 4: Commit** + +```bash +git add backend/api/chat/router.py backend/tests/test_chat_router.py +git commit -m "Route chat requests through memory service" +``` + +--- + +### Task 8: Add frontend chat session utility + +**Files:** +- Create: `frontend/utils/chatSession.ts` + +- [ ] **Step 1: Create utility** + +Create `frontend/utils/chatSession.ts`: + +```ts +const CHAT_SESSION_STORAGE_KEY = "science-gpt.chat_session_id"; + +export function createChatSessionId(): string { + return crypto.randomUUID(); +} + +export function getOrCreateChatSessionId(): string { + const existing = sessionStorage.getItem(CHAT_SESSION_STORAGE_KEY); + if (existing) return existing; + const next = createChatSessionId(); + sessionStorage.setItem(CHAT_SESSION_STORAGE_KEY, next); + return next; +} + +export function resetChatSessionId(): string { + const next = createChatSessionId(); + sessionStorage.setItem(CHAT_SESSION_STORAGE_KEY, next); + return next; +} +``` + +- [ ] **Step 2: Run frontend build** + +Run: + +```bash +cd frontend && npm run build +``` + +Expected: TypeScript compilation succeeds. + +- [ ] **Step 3: Commit** + +```bash +git add frontend/utils/chatSession.ts +git commit -m "Add frontend chat session id utility" +``` + +--- + +### Task 9: Send chat_session_id from ChatPage + +**Files:** +- Modify: `frontend/components/ChatPage.tsx:1-619` + +- [ ] **Step 1: Add imports** + +In `frontend/components/ChatPage.tsx`, add this import near the existing utility imports: + +```ts +import { + getOrCreateChatSessionId, + resetChatSessionId, +} from "../utils/chatSession"; +``` + +- [ ] **Step 2: Add chat session state** + +Inside `ChatPage`, after the existing `ratings` state, add: + +```ts +const [chatSessionId, setChatSessionId] = useState(() => + getOrCreateChatSessionId() +); +``` + +- [ ] **Step 3: Include chat_session_id in request body** + +In the `/api/chat/query` request body, add `chat_session_id` before `query`: + +```ts +body: JSON.stringify({ + chat_session_id: chatSessionId, + query, + document_markdown: documentMarkdown, + model_type: modelConfig.modelType, + model_id: modelConfig.modelId, + deployment: modelConfig.deployment ?? null, + api_version: modelConfig.apiVersion ?? null, +}), +``` + +- [ ] **Step 4: Update sendQuery dependencies** + +Change the `sendQuery` dependency array to include `chatSessionId`: + +```ts +[chatSessionId, docs, getModelConfig] +``` + +- [ ] **Step 5: Add a New Chat handler** + +Below `handleRegenerate`, add: + +```ts +const handleNewChat = useCallback(() => { + const nextSessionId = resetChatSessionId(); + setChatSessionId(nextSessionId); + setMessages([]); + setRatings({}); + setContextError(false); +}, []); +``` + +- [ ] **Step 6: Add a New Chat button** + +In the header right actions before `Simplified Mode`, add: + +```tsx + +``` + +- [ ] **Step 7: Run frontend build** + +Run: + +```bash +cd frontend && npm run build +``` + +Expected: TypeScript compilation succeeds. + +- [ ] **Step 8: Commit** + +```bash +git add frontend/components/ChatPage.tsx +git commit -m "Send independent chat session id" +``` + +--- + +### Task 10: Add optional history-session attachment UI + +**Files:** +- Modify: `frontend/components/ChatPage.tsx:1-740` +- Reuse: `frontend/types/session.ts` + +- [ ] **Step 1: Add imports** + +Add these imports to `frontend/components/ChatPage.tsx`: + +```ts +import { authenticatedFetch } from "../utils/authUtils"; +import { SessionSummary } from "../types/session"; +``` + +Keep the existing `getValidToken` import because upload/document calls still use it. + +- [ ] **Step 2: Add attachment state** + +Inside `ChatPage`, near document state, add: + +```ts +const [historySessions, setHistorySessions] = useState([]); +const [sessionsLoading, setSessionsLoading] = useState(false); +const [attachedSessionId, setAttachedSessionId] = useState(null); +``` + +- [ ] **Step 3: Fetch sessions for the selector** + +Add this effect below the model-loading effect: + +```ts +useEffect(() => { + let cancelled = false; + + async function fetchSessions() { + try { + setSessionsLoading(true); + const response = await authenticatedFetch("/api/sessions"); + if (!response.ok) throw new Error("Failed to fetch sessions"); + const data = await response.json(); + if (!cancelled) setHistorySessions(data.sessions ?? []); + } catch { + if (!cancelled) setHistorySessions([]); + } finally { + if (!cancelled) setSessionsLoading(false); + } + } + + fetchSessions(); + return () => { + cancelled = true; + }; +}, []); +``` + +- [ ] **Step 4: Send attached_session_id in request body** + +Add `attached_session_id` to the chat request body: + +```ts +body: JSON.stringify({ + chat_session_id: chatSessionId, + attached_session_id: attachedSessionId, + query, + document_markdown: documentMarkdown, + model_type: modelConfig.modelType, + model_id: modelConfig.modelId, + deployment: modelConfig.deployment ?? null, + api_version: modelConfig.apiVersion ?? null, +}), +``` + +- [ ] **Step 5: Update sendQuery dependencies** + +Change the `sendQuery` dependency array to: + +```ts +[attachedSessionId, chatSessionId, docs, getModelConfig] +``` + +- [ ] **Step 6: Add a session selector in the header** + +In the header left area, after the model selector ``, add: + +```tsx + +``` + +- [ ] **Step 7: Clear attachment when starting a new chat** + +Update `handleNewChat`: + +```ts +const handleNewChat = useCallback(() => { + const nextSessionId = resetChatSessionId(); + setChatSessionId(nextSessionId); + setMessages([]); + setRatings({}); + setContextError(false); + setAttachedSessionId(null); +}, []); +``` + +- [ ] **Step 8: Run frontend build** + +Run: + +```bash +cd frontend && npm run build +``` + +Expected: TypeScript compilation succeeds. + +- [ ] **Step 9: Commit** + +```bash +git add frontend/components/ChatPage.tsx +git commit -m "Attach history sessions to chat requests" +``` + +--- + +### Task 11: Verify backend and frontend together + +**Files:** +- Verify only; no expected file changes unless a prior task fails. + +- [ ] **Step 1: Run backend chat tests** + +Run: + +```bash +cd backend && python -m pytest tests/test_chat_memory_service.py tests/test_chat_router.py -v +``` + +Expected: PASS. + +- [ ] **Step 2: Run existing backend smoke test that does not require DB** + +Run: + +```bash +cd backend && python -m pytest tests/test_auth_migration.py::test_models_load -v +``` + +Expected: PASS and includes existing app tables. + +- [ ] **Step 3: Run frontend build** + +Run: + +```bash +cd frontend && npm run build +``` + +Expected: PASS. + +- [ ] **Step 4: Run frontend lint if the repo already passes lint on this branch** + +Run: + +```bash +cd frontend && npm run lint +``` + +Expected: PASS. If lint fails on pre-existing unrelated files, capture the failure and do not fix unrelated lint issues. + +- [ ] **Step 5: Manual browser validation** + +Run the app locally with the usual backend and frontend dev commands, then verify: + +1. Open Chat. +2. Send “Remember that my topic is licensing.” +3. Send “What is my topic?” and confirm the assistant remembers licensing. +4. Click New Chat. +5. Send “What is my topic?” and confirm the prior chat memory does not leak. +6. Attach a history session from the selector. +7. Ask a document-specific question and confirm the assistant uses attached document content. +8. Start another New Chat, attach the same history session, and confirm the document is available but prior conversation is not. + +- [ ] **Step 6: Final commit if verification required fixes** + +If verification required fixes, commit only those files: + +```bash +git add +git commit -m "Fix chat memory verification issues" +``` + +--- + +## Self-Review + +Spec coverage: + +- Independent `chat_session_id` memory: Tasks 2, 3, 8, 9, and 11. +- LangGraph checkpointing: Tasks 1 and 3. +- Optional `attached_session_id`: Tasks 4, 5, 7, 10, and 11. +- No full document markdown in persistent memory: Tasks 2, 3, and 5. +- Existing history/session source of truth: Tasks 4, 5, 7, and 10. +- Testing and manual validation: Tasks 2, 4, 6, 11. + +Placeholder scan: no `TBD`, `TODO`, “implement later”, or unspecified code steps remain. + +Type consistency: + +- Backend request field is `chat_session_id` everywhere. +- Optional attachment field is `attached_session_id` everywhere. +- Frontend utility names are `getOrCreateChatSessionId` and `resetChatSessionId` everywhere. +- LangGraph thread config uses `thread_id = chat_session_id`. diff --git a/frontend/components/ChatPage.tsx b/frontend/components/ChatPage.tsx index 1245211..62e7c3c 100644 --- a/frontend/components/ChatPage.tsx +++ b/frontend/components/ChatPage.tsx @@ -1,4 +1,4 @@ -import { useState, useRef, useEffect, useCallback } from "react"; +import { useState, useRef, useEffect, useCallback, ReactNode } from "react"; import ReactMarkdown from "react-markdown"; import remarkGfm from "remark-gfm"; import { @@ -30,7 +30,13 @@ import { pickBestFromList, ModelConfig, } from "../utils/modelSelection"; -import { getValidToken } from "../utils/authUtils"; +import { authenticatedFetch, getValidToken } from "../utils/authUtils"; +import { + createChatSessionId, + getOrCreateChatSessionId, + resetChatSessionId, +} from "../utils/chatSession"; +import { SessionSummary } from "../types/session"; import { toast } from "./ui/sonner"; // ─── Types ──────────────────────────────────────────────────────────────────── @@ -42,6 +48,12 @@ interface Message { } const MAX_DOCS = 5; +const HISTORY_DISABLED_WITH_DOCS_TITLE = + "Start a new chat or remove uploaded documents before attaching history."; +const DOC_UPLOAD_DISABLED_WITH_HISTORY_TITLE = + "History sessions use saved context only. Start a new chat to upload documents."; +const REGENERATE_DISABLED_WITH_MEMORY_TITLE = + "Regenerate is unavailable for chat memory sessions until turn replacement is supported."; type DocEntry = | { status: "loading"; file: File; tempId: string } @@ -61,6 +73,25 @@ interface ChatPageProps { onSignOut?: () => void; } +interface MarkdownCodeProps { + inline?: boolean; + children?: ReactNode; +} + +interface ChatModelConfig extends ModelConfig { + model_type?: string; +} + +const getHistorySessionLabel = (session: SessionSummary): string => { + const trimmedName = session.name?.trim(); + if (trimmedName) return trimmedName; + + const firstDocumentName = session.document_names[0]?.trim(); + if (firstDocumentName) return firstDocumentName; + + return "Untitled Session"; +}; + // ─── Small sub-components ───────────────────────────────────────────────────── function AvatarAI() { @@ -78,6 +109,8 @@ interface MessageRowProps { rating: "up" | "down" | null; onRate: (id: string, rating: "up" | "down") => void; onCopy: (content: string) => void; + canRegenerate: boolean; + regenerateTitle: string; onRegenerate: (id: string) => void; isRegenerating: boolean; } @@ -87,6 +120,8 @@ function MessageRow({ rating, onRate, onCopy, + canRegenerate, + regenerateTitle, onRegenerate, isRegenerating, }: MessageRowProps) { @@ -158,7 +193,7 @@ function MessageRow({ li: ({ children }) => (
  • {children}
    • ), - code: ({ inline, children }: any) => + code: ({ inline, children }: MarkdownCodeProps) => inline ? ( {children} @@ -240,9 +275,9 @@ function MessageRow({ /> onRegenerate(message.id)} - disabled={isRegenerating} + disabled={!canRegenerate || isRegenerating} > >({}); + const [chatSessionId, setChatSessionId] = useState(() => + getOrCreateChatSessionId() + ); const [availableModels, setAvailableModels] = useState([]); const [selectedModelId, setSelectedModelId] = useState(""); const [modelsLoading, setModelsLoading] = useState(true); const [docs, setDocs] = useState>(new Map()); + const [historySessions, setHistorySessions] = useState([]); + const [sessionsLoading, setSessionsLoading] = useState(false); + const [attachedSessionId, setAttachedSessionId] = useState(null); const [contextError, setContextError] = useState(false); const removeDoc = useCallback((tempId: string) => { @@ -325,6 +366,8 @@ export function ChatPage({ onSwitchToWorkflow, onSignOut }: ChatPageProps) { (d) => d.status !== "error" ).length; const atDocLimit = activeDocCount >= MAX_DOCS; + const hasUploadedDocs = activeDocCount > 0; + const canAttachHistory = !hasUploadedDocs; const [isDragOver, setIsDragOver] = useState(false); const dragCounter = useRef(0); @@ -332,6 +375,7 @@ export function ChatPage({ onSwitchToWorkflow, onSignOut }: ChatPageProps) { const messagesEndRef = useRef(null); const fileInputRef = useRef(null); const textareaRef = useRef(null); + const activeRequestSessionIdRef = useRef(null); // ── Models ────────────────────────────────────────────────────────────────── useEffect(() => { @@ -345,6 +389,29 @@ export function ChatPage({ onSwitchToWorkflow, onSignOut }: ChatPageProps) { .finally(() => setModelsLoading(false)); }, []); + useEffect(() => { + let cancelled = false; + + async function fetchSessions() { + try { + setSessionsLoading(true); + const response = await authenticatedFetch("/api/sessions"); + if (!response.ok) throw new Error("Failed to fetch sessions"); + const data = await response.json(); + if (!cancelled) setHistorySessions(data.sessions ?? []); + } catch { + if (!cancelled) setHistorySessions([]); + } finally { + if (!cancelled) setSessionsLoading(false); + } + } + + fetchSessions(); + return () => { + cancelled = true; + }; + }, []); + // ── Auto-scroll ───────────────────────────────────────────────────────────── useEffect(() => { messagesEndRef.current?.scrollIntoView({ behavior: "smooth" }); @@ -352,7 +419,7 @@ export function ChatPage({ onSwitchToWorkflow, onSignOut }: ChatPageProps) { // ── Document processing ───────────────────────────────────────────────────── const processFile = useCallback(async (file: File) => { - const tempId = crypto.randomUUID(); + const tempId = createChatSessionId(); setDocs((prev) => { const active = Array.from(prev.values()).filter( @@ -437,6 +504,12 @@ export function ChatPage({ onSwitchToWorkflow, onSignOut }: ChatPageProps) { (e: React.ChangeEvent) => { const files = Array.from(e.target.files ?? []); if (fileInputRef.current) fileInputRef.current.value = ""; + + if (attachedSessionId) { + toast.error(DOC_UPLOAD_DISABLED_WITH_HISTORY_TITLE); + return; + } + const available = MAX_DOCS - Array.from(docs.values()).filter((d) => d.status !== "error").length; @@ -448,7 +521,7 @@ export function ChatPage({ onSwitchToWorkflow, onSignOut }: ChatPageProps) { ); toProcess.forEach((f) => processFile(f)); }, - [docs, processFile] + [attachedSessionId, docs, processFile] ); // ── Drag and drop ─────────────────────────────────────────────────────────── @@ -473,6 +546,12 @@ export function ChatPage({ onSwitchToWorkflow, onSignOut }: ChatPageProps) { e.preventDefault(); dragCounter.current = 0; setIsDragOver(false); + + if (attachedSessionId) { + toast.error(DOC_UPLOAD_DISABLED_WITH_HISTORY_TITLE); + return; + } + const files = Array.from(e.dataTransfer.files); const available = MAX_DOCS - @@ -488,7 +567,7 @@ export function ChatPage({ onSwitchToWorkflow, onSignOut }: ChatPageProps) { ); toProcess.forEach((f) => processFile(f)); }, - [docs, processFile] + [attachedSessionId, docs, processFile] ); // ── Model config ───────────────────────────────────────────────────────────── @@ -497,9 +576,9 @@ export function ChatPage({ onSwitchToWorkflow, onSignOut }: ChatPageProps) { if (!model) return null; const isGemini = model.provider === "Google Gemini"; const isAnthropic = model.provider === "Anthropic"; + const chatModel = model as ChatModelConfig; const isLlama = - model.provider === "Meta Llama" || - (model as any).model_type === "azure-llama"; + model.provider === "Meta Llama" || chatModel.model_type === "azure-llama"; return { modelType: isGemini ? "gemini" @@ -533,6 +612,8 @@ export function ChatPage({ onSwitchToWorkflow, onSignOut }: ChatPageProps) { return; } + const requestChatSessionId = chatSessionId; + activeRequestSessionIdRef.current = requestChatSessionId; setContextError(false); setIsLoading(true); @@ -561,6 +642,8 @@ export function ChatPage({ onSwitchToWorkflow, onSignOut }: ChatPageProps) { Authorization: `Bearer ${token}`, }, body: JSON.stringify({ + chat_session_id: requestChatSessionId, + attached_session_id: attachedSessionId, query, document_markdown: documentMarkdown, model_type: modelConfig.modelType, @@ -574,15 +657,27 @@ export function ChatPage({ onSwitchToWorkflow, onSignOut }: ChatPageProps) { if (!res.ok || !data.success) throw new Error(data.error || "Request failed"); + if (activeRequestSessionIdRef.current !== requestChatSessionId) { + return; + } + + if (typeof data.attached_session_warning === "string") { + toast.warning(data.attached_session_warning); + } + setMessages((prev) => [ ...prev, { - id: crypto.randomUUID(), + id: createChatSessionId(), role: "assistant", content: data.response, }, ]); } catch (err: unknown) { + if (activeRequestSessionIdRef.current !== requestChatSessionId) { + return; + } + const msg = err instanceof Error ? err.message : "Something went wrong"; if (isContextWindowError(msg)) { setContextError(true); @@ -590,17 +685,20 @@ export function ChatPage({ onSwitchToWorkflow, onSignOut }: ChatPageProps) { setMessages((prev) => [ ...prev, { - id: crypto.randomUUID(), + id: createChatSessionId(), role: "assistant", content: `Error: ${msg}`, }, ]); } } finally { - setIsLoading(false); + if (activeRequestSessionIdRef.current === requestChatSessionId) { + activeRequestSessionIdRef.current = null; + setIsLoading(false); + } } }, - [docs, getModelConfig] + [attachedSessionId, chatSessionId, docs, getModelConfig] ); const handleSend = useCallback(async () => { @@ -608,7 +706,7 @@ export function ChatPage({ onSwitchToWorkflow, onSignOut }: ChatPageProps) { if (!trimmed || isLoading) return; setMessages((prev) => [ ...prev, - { id: crypto.randomUUID(), role: "user", content: trimmed }, + { id: createChatSessionId(), role: "user", content: trimmed }, ]); setInput(""); // Reset textarea height @@ -618,20 +716,21 @@ export function ChatPage({ onSwitchToWorkflow, onSignOut }: ChatPageProps) { await sendQuery(trimmed); }, [input, isLoading, sendQuery]); - const handleRegenerate = useCallback( - async (assistantMsgId: string) => { - // Find the user message that preceded this assistant message - const idx = messages.findIndex((m) => m.id === assistantMsgId); - if (idx < 1) return; - const preceding = messages[idx - 1]; - if (preceding.role !== "user") return; - - // Remove the assistant message and resend - setMessages((prev) => prev.filter((m) => m.id !== assistantMsgId)); - await sendQuery(preceding.content); - }, - [messages, sendQuery] - ); + const handleRegenerate = useCallback(async (_assistantMsgId: string) => { + toast.info(REGENERATE_DISABLED_WITH_MEMORY_TITLE); + }, []); + + const handleNewChat = useCallback(() => { + const nextSessionId = resetChatSessionId(); + activeRequestSessionIdRef.current = nextSessionId; + setIsLoading(false); + setChatSessionId(nextSessionId); + setMessages([]); + setRatings({}); + setContextError(false); + setAttachedSessionId(null); + setDocs(new Map()); + }, []); const handleKeyDown = useCallback( (e: React.KeyboardEvent) => { @@ -699,6 +798,38 @@ export function ChatPage({ onSwitchToWorkflow, onSignOut }: ChatPageProps) { )} + {/* Centre: title (absolutely positioned so it's always truly centred) */} @@ -710,6 +841,16 @@ export function ChatPage({ onSwitchToWorkflow, onSignOut }: ChatPageProps) { {/* Right: actions */}
    + {onSwitchToWorkflow && (
    )} @@ -909,8 +1058,12 @@ export function ChatPage({ onSwitchToWorkflow, onSignOut }: ChatPageProps) { />