diff --git a/surfsense_backend/app/agents/new_chat/chat_deepagent.py b/surfsense_backend/app/agents/new_chat/chat_deepagent.py index eb2dac737..931b31e63 100644 --- a/surfsense_backend/app/agents/new_chat/chat_deepagent.py +++ b/surfsense_backend/app/agents/new_chat/chat_deepagent.py @@ -10,6 +10,7 @@ from collections.abc import Sequence from deepagents import create_deep_agent from langchain_core.tools import BaseTool from langchain_litellm import ChatLiteLLM +from langgraph.types import Checkpointer from sqlalchemy.ext.asyncio import AsyncSession from app.agents.new_chat.context import SurfSenseContextSchema @@ -27,6 +28,7 @@ def create_surfsense_deep_agent( search_space_id: int, db_session: AsyncSession, connector_service: ConnectorService, + checkpointer: Checkpointer, user_instructions: str | None = None, enable_citations: bool = True, additional_tools: Sequence[BaseTool] | None = None, @@ -39,6 +41,8 @@ def create_surfsense_deep_agent( search_space_id: The user's search space ID db_session: Database session connector_service: Initialized connector service + checkpointer: LangGraph checkpointer for conversation state persistence. + Use AsyncPostgresSaver for production or MemorySaver for testing. user_instructions: Optional user instructions to inject into the system prompt. These will be added to the system prompt to customize agent behavior. enable_citations: Whether to include citation instructions in the system prompt (default: True). @@ -61,7 +65,7 @@ def create_surfsense_deep_agent( if additional_tools: tools.extend(additional_tools) - # Create the deep agent with user-configurable system prompt + # Create the deep agent with user-configurable system prompt and checkpointer agent = create_deep_agent( model=llm, tools=tools, @@ -70,6 +74,7 @@ def create_surfsense_deep_agent( enable_citations=enable_citations, ), context_schema=SurfSenseContextSchema, + checkpointer=checkpointer, # Enable conversation memory via thread_id ) return agent diff --git a/surfsense_backend/app/agents/new_chat/checkpointer.py b/surfsense_backend/app/agents/new_chat/checkpointer.py new file mode 100644 index 000000000..8db3706a7 --- /dev/null +++ b/surfsense_backend/app/agents/new_chat/checkpointer.py @@ -0,0 +1,95 @@ +""" +PostgreSQL-based checkpointer for LangGraph agents. + +This module provides a persistent checkpointer using AsyncPostgresSaver +that stores conversation state in the PostgreSQL database. +""" + +from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver + +from app.config import config + +# Global checkpointer instance (initialized lazily) +_checkpointer: AsyncPostgresSaver | None = None +_checkpointer_context = None # Store the context manager for cleanup +_checkpointer_initialized: bool = False + + +def get_postgres_connection_string() -> str: + """ + Convert the async DATABASE_URL to a sync postgres connection string for psycopg3. + + The DATABASE_URL is typically in format: + postgresql+asyncpg://user:pass@host:port/dbname + + We need to convert it to: + postgresql://user:pass@host:port/dbname + """ + db_url = config.DATABASE_URL + + # Handle asyncpg driver prefix + if db_url.startswith("postgresql+asyncpg://"): + return db_url.replace("postgresql+asyncpg://", "postgresql://") + + # Handle other async prefixes + if "+asyncpg" in db_url: + return db_url.replace("+asyncpg", "") + + return db_url + + +async def get_checkpointer() -> AsyncPostgresSaver: + """ + Get or create the global AsyncPostgresSaver instance. + + This function: + 1. Creates the checkpointer if it doesn't exist + 2. Sets up the required database tables on first call + 3. Returns the cached instance on subsequent calls + + Returns: + AsyncPostgresSaver: The configured checkpointer instance + """ + global _checkpointer, _checkpointer_context, _checkpointer_initialized + + if _checkpointer is None: + conn_string = get_postgres_connection_string() + # from_conn_string returns an async context manager + # We need to enter the context to get the actual checkpointer + _checkpointer_context = AsyncPostgresSaver.from_conn_string(conn_string) + _checkpointer = await _checkpointer_context.__aenter__() + + # Setup tables on first call (idempotent) + if not _checkpointer_initialized: + await _checkpointer.setup() + _checkpointer_initialized = True + + return _checkpointer + + +async def setup_checkpointer_tables() -> None: + """ + Explicitly setup the checkpointer tables. + + This can be called during application startup to ensure + tables exist before any agent calls. + """ + await get_checkpointer() + print("[Checkpointer] PostgreSQL checkpoint tables ready") + + +async def close_checkpointer() -> None: + """ + Close the checkpointer connection. + + This should be called during application shutdown. + """ + global _checkpointer, _checkpointer_context, _checkpointer_initialized + + if _checkpointer_context is not None: + await _checkpointer_context.__aexit__(None, None, None) + _checkpointer = None + _checkpointer_context = None + _checkpointer_initialized = False + print("[Checkpointer] PostgreSQL connection closed") + diff --git a/surfsense_backend/app/app.py b/surfsense_backend/app/app.py index e144d165a..b81e2d36d 100644 --- a/surfsense_backend/app/app.py +++ b/surfsense_backend/app/app.py @@ -5,6 +5,7 @@ from fastapi.middleware.cors import CORSMiddleware from sqlalchemy.ext.asyncio import AsyncSession from uvicorn.middleware.proxy_headers import ProxyHeadersMiddleware +from app.agents.new_chat.checkpointer import close_checkpointer, setup_checkpointer_tables from app.config import config from app.db import User, create_db_and_tables, get_async_session from app.routes import router as crud_router @@ -16,7 +17,11 @@ from app.users import SECRET, auth_backend, current_active_user, fastapi_users async def lifespan(app: FastAPI): # Not needed if you setup a migration system like Alembic await create_db_and_tables() + # Setup LangGraph checkpointer tables for conversation persistence + await setup_checkpointer_tables() yield + # Cleanup: close checkpointer connection on shutdown + await close_checkpointer() def registration_allowed(): diff --git a/surfsense_backend/app/routes/chats_routes.py b/surfsense_backend/app/routes/chats_routes.py index 6a4894643..2a65b637c 100644 --- a/surfsense_backend/app/routes/chats_routes.py +++ b/surfsense_backend/app/routes/chats_routes.py @@ -226,6 +226,7 @@ async def handle_new_chat( chat_id=request.chat_id, session=session, llm_config_id=llm_config_id, + messages=request.messages, # Pass message history from frontend ), media_type="text/event-stream", ) diff --git a/surfsense_backend/app/schemas/chats.py b/surfsense_backend/app/schemas/chats.py index 6c281afbc..3109130f5 100644 --- a/surfsense_backend/app/schemas/chats.py +++ b/surfsense_backend/app/schemas/chats.py @@ -48,12 +48,20 @@ class AISDKChatRequest(BaseModel): data: dict[str, Any] | None = None +class ChatMessage(BaseModel): + """A single message in the chat history.""" + + role: str # "user" or "assistant" + content: str + + class NewChatRequest(BaseModel): """Request schema for the new deep agent chat endpoint.""" chat_id: int user_query: str search_space_id: int + messages: list[ChatMessage] | None = None # Optional chat history from frontend class ChatCreate(ChatBase): diff --git a/surfsense_backend/app/tasks/chat/stream_new_chat.py b/surfsense_backend/app/tasks/chat/stream_new_chat.py index 7f97643dc..40f0176d4 100644 --- a/surfsense_backend/app/tasks/chat/stream_new_chat.py +++ b/surfsense_backend/app/tasks/chat/stream_new_chat.py @@ -8,13 +8,13 @@ Data Stream Protocol (SSE format). from collections.abc import AsyncGenerator from uuid import UUID -from langchain_core.messages import HumanMessage +from langchain_core.messages import AIMessage, HumanMessage from sqlalchemy.ext.asyncio import AsyncSession -from app.agents.new_chat.chat_deepagent import ( - create_surfsense_deep_agent, -) +from app.agents.new_chat.chat_deepagent import create_surfsense_deep_agent +from app.agents.new_chat.checkpointer import get_checkpointer from app.agents.new_chat.llm_config import create_chat_litellm_from_config, load_llm_config_from_yaml +from app.schemas.chats import ChatMessage from app.services.connector_service import ConnectorService from app.services.new_streaming_service import VercelStreamingService @@ -26,13 +26,14 @@ async def stream_new_chat( chat_id: int, session: AsyncSession, llm_config_id: int = -1, + messages: list[ChatMessage] | None = None, ) -> AsyncGenerator[str, None]: """ Stream chat responses from the new SurfSense deep agent. This uses the Vercel AI SDK Data Stream Protocol (SSE format) for streaming. - The chat_id is used as LangGraph's thread_id for memory/checkpointing, - so chat history is automatically managed by LangGraph. + The chat_id is used as LangGraph's thread_id for memory/checkpointing. + Message history can be passed from the frontend for context. Args: user_query: The user's query @@ -41,6 +42,7 @@ async def stream_new_chat( chat_id: The chat ID (used as LangGraph thread_id for memory) session: The database session llm_config_id: The LLM configuration ID (default: -1 for first global config) + messages: Optional chat history from frontend (list of ChatMessage) Yields: str: SSE formatted response strings @@ -73,18 +75,36 @@ async def stream_new_chat( # Create connector service connector_service = ConnectorService(session, search_space_id=search_space_id) - # Create the deep agent + # Get the PostgreSQL checkpointer for persistent conversation memory + checkpointer = await get_checkpointer() + + # Create the deep agent with checkpointer agent = create_surfsense_deep_agent( llm=llm, search_space_id=search_space_id, db_session=session, connector_service=connector_service, + checkpointer=checkpointer, ) - # Build input with just the current user query - # Chat history is managed by LangGraph via thread_id + # Build input with message history from frontend + langchain_messages = [] + + # if messages: + # # Convert frontend messages to LangChain format + # for msg in messages: + # if msg.role == "user": + # langchain_messages.append(HumanMessage(content=msg.content)) + # elif msg.role == "assistant": + # langchain_messages.append(AIMessage(content=msg.content)) + # else: + # Fallback: just use the current user query + langchain_messages.append(HumanMessage(content=user_query)) + input_state = { - "messages": [HumanMessage(content=user_query)], + # Lets not pass this message atm because we are using the checkpointer to manage the conversation history + # We will use this to simulate group chat functionality in the future + "messages": langchain_messages, "search_space_id": search_space_id, } diff --git a/surfsense_backend/main.py b/surfsense_backend/main.py index 6a86bbd37..9c1b5c342 100644 --- a/surfsense_backend/main.py +++ b/surfsense_backend/main.py @@ -1,9 +1,15 @@ import argparse +import asyncio import logging +import sys import uvicorn from dotenv import load_dotenv +# Fix for Windows: psycopg requires SelectorEventLoop, not ProactorEventLoop +if sys.platform == "win32": + asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy()) + from app.config.uvicorn import load_uvicorn_config logging.basicConfig( diff --git a/surfsense_backend/pyproject.toml b/surfsense_backend/pyproject.toml index bf9fc5580..25b74f600 100644 --- a/surfsense_backend/pyproject.toml +++ b/surfsense_backend/pyproject.toml @@ -54,6 +54,8 @@ dependencies = [ "trafilatura>=2.0.0", "fastapi-users[oauth,sqlalchemy]>=15.0.3", "chonkie[all]>=1.5.0", + "langgraph-checkpoint-postgres>=3.0.2", + "psycopg[binary,pool]>=3.3.2", ] [dependency-groups] diff --git a/surfsense_backend/uv.lock b/surfsense_backend/uv.lock index 48f8034bc..208509993 100644 --- a/surfsense_backend/uv.lock +++ b/surfsense_backend/uv.lock @@ -2983,6 +2983,21 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/48/e3/616e3a7ff737d98c1bbb5700dd62278914e2a9ded09a79a1fa93cf24ce12/langgraph_checkpoint-3.0.1-py3-none-any.whl", hash = "sha256:9b04a8d0edc0474ce4eaf30c5d731cee38f11ddff50a6177eead95b5c4e4220b", size = 46249 }, ] +[[package]] +name = "langgraph-checkpoint-postgres" +version = "3.0.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "langgraph-checkpoint" }, + { name = "orjson" }, + { name = "psycopg" }, + { name = "psycopg-pool" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/68/4e/ffea5b0d667e10d408b3b2d6dd967ea79e208eef73fe6ee5622625496238/langgraph_checkpoint_postgres-3.0.2.tar.gz", hash = "sha256:448cb8ec245b6fe10171a0f90e9aa047e24a9d3febba6a914644b0c1323da158", size = 127766 } +wheels = [ + { url = "https://files.pythonhosted.org/packages/ac/e4/b4248e10289b6e2c2d33586c87c5eb421e566ef5f336ee45269223cc3b92/langgraph_checkpoint_postgres-3.0.2-py3-none-any.whl", hash = "sha256:15c0fb638edfbc54d496f1758d0327d1a081e0ef94dda8f0c91d4b307d6d8545", size = 42710 }, +] + [[package]] name = "langgraph-prebuilt" version = "1.0.5" @@ -4785,6 +4800,79 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/50/1b/6921afe68c74868b4c9fa424dad3be35b095e16687989ebbb50ce4fceb7c/psutil-7.0.0-cp37-abi3-win_amd64.whl", hash = "sha256:4cf3d4eb1aa9b348dec30105c55cd9b7d4629285735a102beb4441e38db90553", size = 244885 }, ] +[[package]] +name = "psycopg" +version = "3.3.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "typing-extensions", marker = "python_full_version < '3.13'" }, + { name = "tzdata", marker = "sys_platform == 'win32'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/e0/1a/7d9ef4fdc13ef7f15b934c393edc97a35c281bb7d3c3329fbfcbe915a7c2/psycopg-3.3.2.tar.gz", hash = "sha256:707a67975ee214d200511177a6a80e56e654754c9afca06a7194ea6bbfde9ca7", size = 165630 } +wheels = [ + { url = "https://files.pythonhosted.org/packages/8c/51/2779ccdf9305981a06b21a6b27e8547c948d85c41c76ff434192784a4c93/psycopg-3.3.2-py3-none-any.whl", hash = "sha256:3e94bc5f4690247d734599af56e51bae8e0db8e4311ea413f801fef82b14a99b", size = 212774 }, +] + +[package.optional-dependencies] +binary = [ + { name = "psycopg-binary", marker = "implementation_name != 'pypy'" }, +] +pool = [ + { name = "psycopg-pool" }, +] + +[[package]] +name = "psycopg-binary" +version = "3.3.2" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/4e/1e/8614b01c549dd7e385dacdcd83fe194f6b3acb255a53cc67154ee6bf00e7/psycopg_binary-3.3.2-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:a9387ab615f929e71ef0f4a8a51e986fa06236ccfa9f3ec98a88f60fbf230634", size = 4579832 }, + { url = "https://files.pythonhosted.org/packages/26/97/0bb093570fae2f4454d42c1ae6000f15934391867402f680254e4a7def54/psycopg_binary-3.3.2-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:3ff7489df5e06c12d1829544eaec64970fe27fe300f7cf04c8495fe682064688", size = 4658786 }, + { url = "https://files.pythonhosted.org/packages/61/20/1d9383e3f2038826900a14137b0647d755f67551aab316e1021443105ed5/psycopg_binary-3.3.2-cp312-cp312-manylinux2014_ppc64le.manylinux_2_17_ppc64le.whl", hash = "sha256:9742580ecc8e1ac45164e98d32ca6df90da509c2d3ff26be245d94c430f92db4", size = 5454896 }, + { url = "https://files.pythonhosted.org/packages/a6/62/513c80ad8bbb545e364f7737bf2492d34a4c05eef4f7b5c16428dc42260d/psycopg_binary-3.3.2-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:d45acedcaa58619355f18e0f42af542fcad3fd84ace4b8355d3a5dea23318578", size = 5132731 }, + { url = "https://files.pythonhosted.org/packages/f3/28/ddf5f5905f088024bccb19857949467407c693389a14feb527d6171d8215/psycopg_binary-3.3.2-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:d88f32ff8c47cb7f4e7e7a9d1747dcee6f3baa19ed9afa9e5694fd2fb32b61ed", size = 6724495 }, + { url = "https://files.pythonhosted.org/packages/6e/93/a1157ebcc650960b264542b547f7914d87a42ff0cc15a7584b29d5807e6b/psycopg_binary-3.3.2-cp312-cp312-manylinux_2_38_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:59d0163c4617a2c577cb34afbed93d7a45b8c8364e54b2bd2020ff25d5f5f860", size = 4964979 }, + { url = "https://files.pythonhosted.org/packages/0e/27/65939ba6798f9c5be4a5d9cd2061ebaf0851798525c6811d347821c8132d/psycopg_binary-3.3.2-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:e750afe74e6c17b2c7046d2c3e3173b5a3f6080084671c8aa327215323df155b", size = 4493648 }, + { url = "https://files.pythonhosted.org/packages/8a/c4/5e9e4b9b1c1e27026e43387b0ba4aaf3537c7806465dd3f1d5bde631752a/psycopg_binary-3.3.2-cp312-cp312-musllinux_1_2_ppc64le.whl", hash = "sha256:f26f113013c4dcfbfe9ced57b5bad2035dda1a7349f64bf726021968f9bccad3", size = 4173392 }, + { url = "https://files.pythonhosted.org/packages/c6/81/cf43fb76993190cee9af1cbcfe28afb47b1928bdf45a252001017e5af26e/psycopg_binary-3.3.2-cp312-cp312-musllinux_1_2_riscv64.whl", hash = "sha256:8309ee4569dced5e81df5aa2dcd48c7340c8dee603a66430f042dfbd2878edca", size = 3909241 }, + { url = "https://files.pythonhosted.org/packages/9d/20/c6377a0d17434674351627489deca493ea0b137c522b99c81d3a106372c8/psycopg_binary-3.3.2-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:c6464150e25b68ae3cb04c4e57496ea11ebfaae4d98126aea2f4702dd43e3c12", size = 4219746 }, + { url = "https://files.pythonhosted.org/packages/25/32/716c57b28eefe02a57a4c9d5bf956849597f5ea476c7010397199e56cfde/psycopg_binary-3.3.2-cp312-cp312-win_amd64.whl", hash = "sha256:716a586f99bbe4f710dc58b40069fcb33c7627e95cc6fc936f73c9235e07f9cf", size = 3537494 }, + { url = "https://files.pythonhosted.org/packages/14/73/7ca7cb22b9ac7393fb5de7d28ca97e8347c375c8498b3bff2c99c1f38038/psycopg_binary-3.3.2-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:fc5a189e89cbfff174588665bb18d28d2d0428366cc9dae5864afcaa2e57380b", size = 4579068 }, + { url = "https://files.pythonhosted.org/packages/f5/42/0cf38ff6c62c792fc5b55398a853a77663210ebd51ed6f0c4a05b06f95a6/psycopg_binary-3.3.2-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:083c2e182be433f290dc2c516fd72b9b47054fcd305cce791e0a50d9e93e06f2", size = 4657520 }, + { url = "https://files.pythonhosted.org/packages/3b/60/df846bc84cbf2231e01b0fff48b09841fe486fa177665e50f4995b1bfa44/psycopg_binary-3.3.2-cp313-cp313-manylinux2014_ppc64le.manylinux_2_17_ppc64le.whl", hash = "sha256:ac230e3643d1c436a2dfb59ca84357dfc6862c9f372fc5dbd96bafecae581f9f", size = 5452086 }, + { url = "https://files.pythonhosted.org/packages/ab/85/30c846a00db86b1b53fd5bfd4b4edfbd0c00de8f2c75dd105610bd7568fc/psycopg_binary-3.3.2-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:d8c899a540f6c7585cee53cddc929dd4d2db90fd828e37f5d4017b63acbc1a5d", size = 5131125 }, + { url = "https://files.pythonhosted.org/packages/6d/15/9968732013373f36f8a2a3fb76104dffc8efd9db78709caa5ae1a87b1f80/psycopg_binary-3.3.2-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:50ff10ab8c0abdb5a5451b9315538865b50ba64c907742a1385fdf5f5772b73e", size = 6722914 }, + { url = "https://files.pythonhosted.org/packages/b2/ba/29e361fe02143ac5ff5a1ca3e45697344cfbebe2eaf8c4e7eec164bff9a0/psycopg_binary-3.3.2-cp313-cp313-manylinux_2_38_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:23d2594af848c1fd3d874a9364bef50730124e72df7bb145a20cb45e728c50ed", size = 4966081 }, + { url = "https://files.pythonhosted.org/packages/99/45/1be90c8f1a1a237046903e91202fb06708745c179f220b361d6333ed7641/psycopg_binary-3.3.2-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:ea4fe6b4ead3bbbe27244ea224fcd1f53cb119afc38b71a2f3ce570149a03e30", size = 4493332 }, + { url = "https://files.pythonhosted.org/packages/2e/b5/bbdc07d5f0a5e90c617abd624368182aa131485e18038b2c6c85fc054aed/psycopg_binary-3.3.2-cp313-cp313-musllinux_1_2_ppc64le.whl", hash = "sha256:742ce48cde825b8e52fb1a658253d6d1ff66d152081cbc76aa45e2986534858d", size = 4170781 }, + { url = "https://files.pythonhosted.org/packages/d1/2a/0d45e4f4da2bd78c3237ffa03475ef3751f69a81919c54a6e610eb1a7c96/psycopg_binary-3.3.2-cp313-cp313-musllinux_1_2_riscv64.whl", hash = "sha256:e22bf6b54df994aff37ab52695d635f1ef73155e781eee1f5fa75bc08b58c8da", size = 3910544 }, + { url = "https://files.pythonhosted.org/packages/3a/62/a8e0f092f4dbef9a94b032fb71e214cf0a375010692fbe7493a766339e47/psycopg_binary-3.3.2-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:8db9034cde3bcdafc66980f0130813f5c5d19e74b3f2a19fb3cfbc25ad113121", size = 4220070 }, + { url = "https://files.pythonhosted.org/packages/09/e6/5fc8d8aff8afa114bb4a94a0341b9309311e8bf3ab32d816032f8b984d4e/psycopg_binary-3.3.2-cp313-cp313-win_amd64.whl", hash = "sha256:df65174c7cf6b05ea273ce955927d3270b3a6e27b0b12762b009ce6082b8d3fc", size = 3540922 }, + { url = "https://files.pythonhosted.org/packages/bd/75/ad18c0b97b852aba286d06befb398cc6d383e9dfd0a518369af275a5a526/psycopg_binary-3.3.2-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:9ca24062cd9b2270e4d77576042e9cc2b1d543f09da5aba1f1a3d016cea28390", size = 4596371 }, + { url = "https://files.pythonhosted.org/packages/5a/79/91649d94c8d89f84af5da7c9d474bfba35b08eb8f492ca3422b08f0a6427/psycopg_binary-3.3.2-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:c749770da0947bc972e512f35366dd4950c0e34afad89e60b9787a37e97cb443", size = 4675139 }, + { url = "https://files.pythonhosted.org/packages/56/ac/b26e004880f054549ec9396594e1ffe435810b0673e428e619ed722e4244/psycopg_binary-3.3.2-cp314-cp314-manylinux2014_ppc64le.manylinux_2_17_ppc64le.whl", hash = "sha256:03b7cd73fb8c45d272a34ae7249713e32492891492681e3cf11dff9531cf37e9", size = 5456120 }, + { url = "https://files.pythonhosted.org/packages/4b/8d/410681dccd6f2999fb115cc248521ec50dd2b0aba66ae8de7e81efdebbee/psycopg_binary-3.3.2-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:43b130e3b6edcb5ee856c7167ccb8561b473308c870ed83978ae478613764f1c", size = 5133484 }, + { url = "https://files.pythonhosted.org/packages/66/30/ebbab99ea2cfa099d7b11b742ce13415d44f800555bfa4ad2911dc645b71/psycopg_binary-3.3.2-cp314-cp314-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:7c1feba5a8c617922321aef945865334e468337b8fc5c73074f5e63143013b5a", size = 6731818 }, + { url = "https://files.pythonhosted.org/packages/70/02/d260646253b7ad805d60e0de47f9b811d6544078452579466a098598b6f4/psycopg_binary-3.3.2-cp314-cp314-manylinux_2_38_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:cabb2a554d9a0a6bf84037d86ca91782f087dfff2a61298d0b00c19c0bc43f6d", size = 4983859 }, + { url = "https://files.pythonhosted.org/packages/72/8d/e778d7bad1a7910aa36281f092bd85c5702f508fd9bb0ea2020ffbb6585c/psycopg_binary-3.3.2-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:74bc306c4b4df35b09bc8cecf806b271e1c5d708f7900145e4e54a2e5dedfed0", size = 4516388 }, + { url = "https://files.pythonhosted.org/packages/bd/f1/64e82098722e2ab3521797584caf515284be09c1e08a872551b6edbb0074/psycopg_binary-3.3.2-cp314-cp314-musllinux_1_2_ppc64le.whl", hash = "sha256:d79b0093f0fbf7a962d6a46ae292dc056c65d16a8ee9361f3cfbafd4c197ab14", size = 4192382 }, + { url = "https://files.pythonhosted.org/packages/fa/d0/c20f4e668e89494972e551c31be2a0016e3f50d552d7ae9ac07086407599/psycopg_binary-3.3.2-cp314-cp314-musllinux_1_2_riscv64.whl", hash = "sha256:1586e220be05547c77afc326741dd41cc7fba38a81f9931f616ae98865439678", size = 3928660 }, + { url = "https://files.pythonhosted.org/packages/0f/e1/99746c171de22539fd5eb1c9ca21dc805b54cfae502d7451d237d1dbc349/psycopg_binary-3.3.2-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:458696a5fa5dad5b6fb5d5862c22454434ce4fe1cf66ca6c0de5f904cbc1ae3e", size = 4239169 }, + { url = "https://files.pythonhosted.org/packages/72/f7/212343c1c9cfac35fd943c527af85e9091d633176e2a407a0797856ff7b9/psycopg_binary-3.3.2-cp314-cp314-win_amd64.whl", hash = "sha256:04bb2de4ba69d6f8395b446ede795e8884c040ec71d01dd07ac2b2d18d4153d1", size = 3642122 }, +] + +[[package]] +name = "psycopg-pool" +version = "3.3.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/56/9a/9470d013d0d50af0da9c4251614aeb3c1823635cab3edc211e3839db0bcf/psycopg_pool-3.3.0.tar.gz", hash = "sha256:fa115eb2860bd88fce1717d75611f41490dec6135efb619611142b24da3f6db5", size = 31606 } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e7/c3/26b8a0908a9db249de3b4169692e1c7c19048a9bc41a4d3209cee7dbb758/psycopg_pool-3.3.0-py3-none-any.whl", hash = "sha256:2e44329155c410b5e8666372db44276a8b1ebd8c90f1c3026ebba40d4bc81063", size = 39995 }, +] + [[package]] name = "psycopg2-binary" version = "2.9.11" @@ -6293,6 +6381,7 @@ dependencies = [ { name = "langchain-litellm" }, { name = "langchain-unstructured" }, { name = "langgraph" }, + { name = "langgraph-checkpoint-postgres" }, { name = "linkup-sdk" }, { name = "litellm" }, { name = "llama-cloud-services" }, @@ -6301,6 +6390,7 @@ dependencies = [ { name = "numpy" }, { name = "pgvector" }, { name = "playwright" }, + { name = "psycopg", extra = ["binary", "pool"] }, { name = "pypdf" }, { name = "python-ffmpeg" }, { name = "redis" }, @@ -6351,6 +6441,7 @@ requires-dist = [ { name = "langchain-litellm", specifier = ">=0.3.5" }, { name = "langchain-unstructured", specifier = ">=1.0.0" }, { name = "langgraph", specifier = ">=1.0.5" }, + { name = "langgraph-checkpoint-postgres", specifier = ">=3.0.2" }, { name = "linkup-sdk", specifier = ">=0.2.4" }, { name = "litellm", specifier = ">=1.80.10" }, { name = "llama-cloud-services", specifier = ">=0.6.25" }, @@ -6359,6 +6450,7 @@ requires-dist = [ { name = "numpy", specifier = ">=1.24.0" }, { name = "pgvector", specifier = ">=0.3.6" }, { name = "playwright", specifier = ">=1.50.0" }, + { name = "psycopg", extras = ["binary", "pool"], specifier = ">=3.3.2" }, { name = "pypdf", specifier = ">=5.1.0" }, { name = "python-ffmpeg", specifier = ">=2.0.12" }, { name = "redis", specifier = ">=5.2.1" }, diff --git a/surfsense_web/components/new-chat/source-detail-panel.tsx b/surfsense_web/components/new-chat/source-detail-panel.tsx index 5ed86be7f..a076586a7 100644 --- a/surfsense_web/components/new-chat/source-detail-panel.tsx +++ b/surfsense_web/components/new-chat/source-detail-panel.tsx @@ -14,7 +14,7 @@ import { Sparkles, } from "lucide-react"; import type React from "react"; -import { type ReactNode, useCallback, useEffect, useRef, useState } from "react"; +import { type ReactNode, forwardRef, useCallback, useEffect, useRef, useState } from "react"; import { createPortal } from "react-dom"; import { MarkdownViewer } from "@/components/markdown-viewer"; import { Button } from "@/components/ui/button"; @@ -44,83 +44,70 @@ const formatDocumentType = (type: string) => { .join(" "); }; -// Chunk card component with enhanced animations -const ChunkCard = ({ - chunk, - index, - totalChunks, - isCited, - isActive, -}: { +// Chunk card component +// For large documents (>30 chunks), we disable animation to prevent layout shifts +// which break auto-scroll functionality +interface ChunkCardProps { chunk: { id: number; content: string }; index: number; totalChunks: number; isCited: boolean; isActive: boolean; -}) => { - const shouldReduceMotion = useReducedMotion(); + disableLayoutAnimation?: boolean; +} - return ( - - {/* Cited indicator glow effect */} - {isCited && ( -
- )} - - {/* Header */} -
-
-
- {index + 1} -
- - of {totalChunks} chunks - -
+const ChunkCard = forwardRef( + ({ chunk, index, totalChunks, isCited, isActive, disableLayoutAnimation }, ref) => { + return ( +
+ {/* Cited indicator glow effect */} {isCited && ( - +
+ )} + + {/* Header */} +
+
+
+ {index + 1} +
+ + of {totalChunks} chunks + +
+ {isCited && ( Cited Source - - )} -
+ )} +
- {/* Content */} -
- + {/* Content */} +
+ +
- - ); -}; + ); + } +); +ChunkCard.displayName = "ChunkCard"; export function SourceDetailPanel({ open, @@ -133,6 +120,7 @@ export function SourceDetailPanel({ children, }: SourceDetailPanelProps) { const scrollAreaRef = useRef(null); + const hasScrolledRef = useRef(false); // Use ref to avoid stale closures const [summaryOpen, setSummaryOpen] = useState(false); const [activeChunkIndex, setActiveChunkIndex] = useState(null); const [mounted, setMounted] = useState(false); @@ -163,30 +151,92 @@ export function SourceDetailPanel({ // Find cited chunk index const citedChunkIndex = documentData?.chunks?.findIndex((chunk) => chunk.id === chunkId) ?? -1; - // Auto-scroll to cited chunk when data loads - useEffect(() => { - if (documentData?.chunks && citedChunkIndex !== -1 && !hasScrolledToCited && open) { - // Wait for animations to complete then scroll - const timer = setTimeout(() => { - const chunkElement = scrollAreaRef.current?.querySelector( - `[data-chunk-index="${citedChunkIndex}"]` - ); - if (chunkElement) { - chunkElement.scrollIntoView({ - behavior: shouldReduceMotion ? "auto" : "smooth", - block: "center", - }); - setHasScrolledToCited(true); - setActiveChunkIndex(citedChunkIndex); - } - }, 400); - return () => clearTimeout(timer); + // Simple scroll function that scrolls to a chunk by index + const scrollToChunkByIndex = useCallback((chunkIndex: number, smooth = true) => { + const scrollContainer = scrollAreaRef.current; + if (!scrollContainer) return; + + const viewport = scrollContainer.querySelector( + '[data-radix-scroll-area-viewport]' + ) as HTMLElement | null; + if (!viewport) return; + + const chunkElement = scrollContainer.querySelector( + `[data-chunk-index="${chunkIndex}"]` + ) as HTMLElement | null; + if (!chunkElement) return; + + // Get positions using getBoundingClientRect for accuracy + const viewportRect = viewport.getBoundingClientRect(); + const chunkRect = chunkElement.getBoundingClientRect(); + + // Calculate where to scroll to center the chunk + const currentScrollTop = viewport.scrollTop; + const chunkTopRelativeToViewport = chunkRect.top - viewportRect.top + currentScrollTop; + const scrollTarget = chunkTopRelativeToViewport - (viewportRect.height / 2) + (chunkRect.height / 2); + + viewport.scrollTo({ + top: Math.max(0, scrollTarget), + behavior: smooth && !shouldReduceMotion ? "smooth" : "auto", + }); + + setActiveChunkIndex(chunkIndex); + }, [shouldReduceMotion]); + + // Callback ref for the cited chunk - scrolls when the element mounts + const citedChunkRefCallback = useCallback((node: HTMLDivElement | null) => { + if (node && !hasScrolledRef.current && open) { + hasScrolledRef.current = true; // Mark immediately to prevent duplicate scrolls + + // Store the node reference for the delayed scroll + const scrollToCitedChunk = () => { + const scrollContainer = scrollAreaRef.current; + if (!scrollContainer || !node.isConnected) return false; + + const viewport = scrollContainer.querySelector( + '[data-radix-scroll-area-viewport]' + ) as HTMLElement | null; + if (!viewport) return false; + + // Get positions + const viewportRect = viewport.getBoundingClientRect(); + const chunkRect = node.getBoundingClientRect(); + + // Calculate scroll position to center the chunk + const currentScrollTop = viewport.scrollTop; + const chunkTopRelativeToViewport = chunkRect.top - viewportRect.top + currentScrollTop; + const scrollTarget = chunkTopRelativeToViewport - (viewportRect.height / 2) + (chunkRect.height / 2); + + viewport.scrollTo({ + top: Math.max(0, scrollTarget), + behavior: "auto", // Instant scroll for initial positioning + }); + + return true; + }; + + // Scroll multiple times with delays to handle progressive content rendering + // Each subsequent scroll will correct for any layout shifts + const scrollAttempts = [50, 150, 300, 600, 1000]; + + scrollAttempts.forEach((delay) => { + setTimeout(() => { + scrollToCitedChunk(); + }, delay); + }); + + // After final attempt, mark state as scrolled + setTimeout(() => { + setHasScrolledToCited(true); + setActiveChunkIndex(citedChunkIndex); + }, scrollAttempts[scrollAttempts.length - 1] + 50); } - }, [documentData, citedChunkIndex, hasScrolledToCited, open, shouldReduceMotion]); + }, [open, citedChunkIndex]); // Reset scroll state when panel closes useEffect(() => { if (!open) { + hasScrolledRef.current = false; setHasScrolledToCited(false); setActiveChunkIndex(null); } @@ -222,12 +272,8 @@ export function SourceDetailPanel({ }; const scrollToChunk = useCallback((index: number) => { - setActiveChunkIndex(index); - const chunkElement = scrollAreaRef.current?.querySelector( - `[data-chunk-index="${index}"]` - ); - chunkElement?.scrollIntoView({ behavior: "smooth", block: "center" }); - }, []); + scrollToChunkByIndex(index, true); + }, [scrollToChunkByIndex]); const panelContent = ( @@ -388,9 +434,9 @@ export function SourceDetailPanel({ initial={{ opacity: 0, x: -20 }} animate={{ opacity: 1, x: 0 }} transition={{ delay: 0.2 }} - className="hidden lg:flex flex-col w-16 border-r bg-muted/10" + className="hidden lg:flex flex-col w-16 border-r bg-muted/10 overflow-hidden" > - +
{documentData.chunks.map((chunk, idx) => { const isCited = chunk.id === chunkId; @@ -514,16 +560,21 @@ export function SourceDetailPanel({ {/* Chunks */}
- {documentData.chunks.map((chunk, idx) => ( - - ))} + {documentData.chunks.map((chunk, idx) => { + const isCited = chunk.id === chunkId; + return ( + 30} + /> + ); + })}
diff --git a/surfsense_web/lib/chat/new-chat-transport.ts b/surfsense_web/lib/chat/new-chat-transport.ts index 5be9278d9..7d9887102 100644 --- a/surfsense_web/lib/chat/new-chat-transport.ts +++ b/surfsense_web/lib/chat/new-chat-transport.ts @@ -11,12 +11,41 @@ interface NewChatAdapterConfig { chatId: number; } +interface ChatMessageForBackend { + role: "user" | "assistant"; + content: string; +} + +/** + * Converts assistant-ui messages to a simple format for the backend + */ +function convertMessagesToBackendFormat( + messages: ChatModelRunOptions["messages"] +): ChatMessageForBackend[] { + return messages + .filter((m) => m.role === "user" || m.role === "assistant") + .map((m) => { + // Extract text content from the message parts + let content = ""; + for (const part of m.content) { + if (part.type === "text") { + content += part.text; + } + } + return { + role: m.role as "user" | "assistant", + content: content.trim(), + }; + }) + .filter((m) => m.content.length > 0); // Filter out empty messages +} + /** * Creates a ChatModelAdapter that connects to the FastAPI new_chat endpoint. * * The backend expects: * - POST /api/v1/new_chat - * - Body: { chat_id: number, user_query: string, search_space_id: number } + * - Body: { chat_id: number, user_query: string, search_space_id: number, messages: [...] } * - Returns: SSE stream with Vercel AI SDK Data Stream Protocol */ export function createNewChatAdapter(config: NewChatAdapterConfig): ChatModelAdapter { @@ -31,7 +60,7 @@ export function createNewChatAdapter(config: NewChatAdapterConfig): ChatModelAda throw new Error("No user message found"); } - // Extract text content from the message + // Extract text content from the last user message let userQuery = ""; for (const part of lastUserMessage.content) { if (part.type === "text") { @@ -48,6 +77,9 @@ export function createNewChatAdapter(config: NewChatAdapterConfig): ChatModelAda throw new Error("Not authenticated. Please log in again."); } + // Convert all messages to backend format for chat history + const messageHistory = convertMessagesToBackendFormat(messages); + const response = await fetch(`${backendUrl}/api/v1/new_chat`, { method: "POST", headers: { @@ -58,6 +90,7 @@ export function createNewChatAdapter(config: NewChatAdapterConfig): ChatModelAda chat_id: config.chatId, user_query: userQuery.trim(), search_space_id: config.searchSpaceId, + messages: messageHistory, }), signal: abortSignal, }); @@ -165,3 +198,4 @@ export function createNewChatAdapter(config: NewChatAdapterConfig): ChatModelAda }, }; } +