feat: message history and PostgreSQL checkpointer integration

This commit is contained in:
DESKTOP-RTLN3BA\$punk 2025-12-21 03:30:10 -08:00
parent 3906ba52e0
commit 73f0f772a8
11 changed files with 434 additions and 115 deletions

View file

@ -10,6 +10,7 @@ from collections.abc import Sequence
from deepagents import create_deep_agent from deepagents import create_deep_agent
from langchain_core.tools import BaseTool from langchain_core.tools import BaseTool
from langchain_litellm import ChatLiteLLM from langchain_litellm import ChatLiteLLM
from langgraph.types import Checkpointer
from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.ext.asyncio import AsyncSession
from app.agents.new_chat.context import SurfSenseContextSchema from app.agents.new_chat.context import SurfSenseContextSchema
@ -27,6 +28,7 @@ def create_surfsense_deep_agent(
search_space_id: int, search_space_id: int,
db_session: AsyncSession, db_session: AsyncSession,
connector_service: ConnectorService, connector_service: ConnectorService,
checkpointer: Checkpointer,
user_instructions: str | None = None, user_instructions: str | None = None,
enable_citations: bool = True, enable_citations: bool = True,
additional_tools: Sequence[BaseTool] | None = None, additional_tools: Sequence[BaseTool] | None = None,
@ -39,6 +41,8 @@ def create_surfsense_deep_agent(
search_space_id: The user's search space ID search_space_id: The user's search space ID
db_session: Database session db_session: Database session
connector_service: Initialized connector service 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. user_instructions: Optional user instructions to inject into the system prompt.
These will be added to the system prompt to customize agent behavior. 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). 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: if additional_tools:
tools.extend(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( agent = create_deep_agent(
model=llm, model=llm,
tools=tools, tools=tools,
@ -70,6 +74,7 @@ def create_surfsense_deep_agent(
enable_citations=enable_citations, enable_citations=enable_citations,
), ),
context_schema=SurfSenseContextSchema, context_schema=SurfSenseContextSchema,
checkpointer=checkpointer, # Enable conversation memory via thread_id
) )
return agent return agent

View file

@ -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")

View file

@ -5,6 +5,7 @@ from fastapi.middleware.cors import CORSMiddleware
from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.ext.asyncio import AsyncSession
from uvicorn.middleware.proxy_headers import ProxyHeadersMiddleware 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.config import config
from app.db import User, create_db_and_tables, get_async_session from app.db import User, create_db_and_tables, get_async_session
from app.routes import router as crud_router 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): async def lifespan(app: FastAPI):
# Not needed if you setup a migration system like Alembic # Not needed if you setup a migration system like Alembic
await create_db_and_tables() await create_db_and_tables()
# Setup LangGraph checkpointer tables for conversation persistence
await setup_checkpointer_tables()
yield yield
# Cleanup: close checkpointer connection on shutdown
await close_checkpointer()
def registration_allowed(): def registration_allowed():

View file

@ -226,6 +226,7 @@ async def handle_new_chat(
chat_id=request.chat_id, chat_id=request.chat_id,
session=session, session=session,
llm_config_id=llm_config_id, llm_config_id=llm_config_id,
messages=request.messages, # Pass message history from frontend
), ),
media_type="text/event-stream", media_type="text/event-stream",
) )

View file

@ -48,12 +48,20 @@ class AISDKChatRequest(BaseModel):
data: dict[str, Any] | None = None 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): class NewChatRequest(BaseModel):
"""Request schema for the new deep agent chat endpoint.""" """Request schema for the new deep agent chat endpoint."""
chat_id: int chat_id: int
user_query: str user_query: str
search_space_id: int search_space_id: int
messages: list[ChatMessage] | None = None # Optional chat history from frontend
class ChatCreate(ChatBase): class ChatCreate(ChatBase):

View file

@ -8,13 +8,13 @@ Data Stream Protocol (SSE format).
from collections.abc import AsyncGenerator from collections.abc import AsyncGenerator
from uuid import UUID from uuid import UUID
from langchain_core.messages import HumanMessage from langchain_core.messages import AIMessage, HumanMessage
from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.ext.asyncio import AsyncSession
from app.agents.new_chat.chat_deepagent import ( from app.agents.new_chat.chat_deepagent import create_surfsense_deep_agent
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.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.connector_service import ConnectorService
from app.services.new_streaming_service import VercelStreamingService from app.services.new_streaming_service import VercelStreamingService
@ -26,13 +26,14 @@ async def stream_new_chat(
chat_id: int, chat_id: int,
session: AsyncSession, session: AsyncSession,
llm_config_id: int = -1, llm_config_id: int = -1,
messages: list[ChatMessage] | None = None,
) -> AsyncGenerator[str, None]: ) -> AsyncGenerator[str, None]:
""" """
Stream chat responses from the new SurfSense deep agent. Stream chat responses from the new SurfSense deep agent.
This uses the Vercel AI SDK Data Stream Protocol (SSE format) for streaming. 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, The chat_id is used as LangGraph's thread_id for memory/checkpointing.
so chat history is automatically managed by LangGraph. Message history can be passed from the frontend for context.
Args: Args:
user_query: The user's query 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) chat_id: The chat ID (used as LangGraph thread_id for memory)
session: The database session session: The database session
llm_config_id: The LLM configuration ID (default: -1 for first global config) llm_config_id: The LLM configuration ID (default: -1 for first global config)
messages: Optional chat history from frontend (list of ChatMessage)
Yields: Yields:
str: SSE formatted response strings str: SSE formatted response strings
@ -73,18 +75,36 @@ async def stream_new_chat(
# Create connector service # Create connector service
connector_service = ConnectorService(session, search_space_id=search_space_id) 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( agent = create_surfsense_deep_agent(
llm=llm, llm=llm,
search_space_id=search_space_id, search_space_id=search_space_id,
db_session=session, db_session=session,
connector_service=connector_service, connector_service=connector_service,
checkpointer=checkpointer,
) )
# Build input with just the current user query # Build input with message history from frontend
# Chat history is managed by LangGraph via thread_id 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 = { 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, "search_space_id": search_space_id,
} }

View file

@ -1,9 +1,15 @@
import argparse import argparse
import asyncio
import logging import logging
import sys
import uvicorn import uvicorn
from dotenv import load_dotenv 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 from app.config.uvicorn import load_uvicorn_config
logging.basicConfig( logging.basicConfig(

View file

@ -54,6 +54,8 @@ dependencies = [
"trafilatura>=2.0.0", "trafilatura>=2.0.0",
"fastapi-users[oauth,sqlalchemy]>=15.0.3", "fastapi-users[oauth,sqlalchemy]>=15.0.3",
"chonkie[all]>=1.5.0", "chonkie[all]>=1.5.0",
"langgraph-checkpoint-postgres>=3.0.2",
"psycopg[binary,pool]>=3.3.2",
] ]
[dependency-groups] [dependency-groups]

View file

@ -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 }, { 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]] [[package]]
name = "langgraph-prebuilt" name = "langgraph-prebuilt"
version = "1.0.5" 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 }, { 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]] [[package]]
name = "psycopg2-binary" name = "psycopg2-binary"
version = "2.9.11" version = "2.9.11"
@ -6293,6 +6381,7 @@ dependencies = [
{ name = "langchain-litellm" }, { name = "langchain-litellm" },
{ name = "langchain-unstructured" }, { name = "langchain-unstructured" },
{ name = "langgraph" }, { name = "langgraph" },
{ name = "langgraph-checkpoint-postgres" },
{ name = "linkup-sdk" }, { name = "linkup-sdk" },
{ name = "litellm" }, { name = "litellm" },
{ name = "llama-cloud-services" }, { name = "llama-cloud-services" },
@ -6301,6 +6390,7 @@ dependencies = [
{ name = "numpy" }, { name = "numpy" },
{ name = "pgvector" }, { name = "pgvector" },
{ name = "playwright" }, { name = "playwright" },
{ name = "psycopg", extra = ["binary", "pool"] },
{ name = "pypdf" }, { name = "pypdf" },
{ name = "python-ffmpeg" }, { name = "python-ffmpeg" },
{ name = "redis" }, { name = "redis" },
@ -6351,6 +6441,7 @@ requires-dist = [
{ name = "langchain-litellm", specifier = ">=0.3.5" }, { name = "langchain-litellm", specifier = ">=0.3.5" },
{ name = "langchain-unstructured", specifier = ">=1.0.0" }, { name = "langchain-unstructured", specifier = ">=1.0.0" },
{ name = "langgraph", specifier = ">=1.0.5" }, { name = "langgraph", specifier = ">=1.0.5" },
{ name = "langgraph-checkpoint-postgres", specifier = ">=3.0.2" },
{ name = "linkup-sdk", specifier = ">=0.2.4" }, { name = "linkup-sdk", specifier = ">=0.2.4" },
{ name = "litellm", specifier = ">=1.80.10" }, { name = "litellm", specifier = ">=1.80.10" },
{ name = "llama-cloud-services", specifier = ">=0.6.25" }, { name = "llama-cloud-services", specifier = ">=0.6.25" },
@ -6359,6 +6450,7 @@ requires-dist = [
{ name = "numpy", specifier = ">=1.24.0" }, { name = "numpy", specifier = ">=1.24.0" },
{ name = "pgvector", specifier = ">=0.3.6" }, { name = "pgvector", specifier = ">=0.3.6" },
{ name = "playwright", specifier = ">=1.50.0" }, { name = "playwright", specifier = ">=1.50.0" },
{ name = "psycopg", extras = ["binary", "pool"], specifier = ">=3.3.2" },
{ name = "pypdf", specifier = ">=5.1.0" }, { name = "pypdf", specifier = ">=5.1.0" },
{ name = "python-ffmpeg", specifier = ">=2.0.12" }, { name = "python-ffmpeg", specifier = ">=2.0.12" },
{ name = "redis", specifier = ">=5.2.1" }, { name = "redis", specifier = ">=5.2.1" },

View file

@ -14,7 +14,7 @@ import {
Sparkles, Sparkles,
} from "lucide-react"; } from "lucide-react";
import type React from "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 { createPortal } from "react-dom";
import { MarkdownViewer } from "@/components/markdown-viewer"; import { MarkdownViewer } from "@/components/markdown-viewer";
import { Button } from "@/components/ui/button"; import { Button } from "@/components/ui/button";
@ -44,83 +44,70 @@ const formatDocumentType = (type: string) => {
.join(" "); .join(" ");
}; };
// Chunk card component with enhanced animations // Chunk card component
const ChunkCard = ({ // For large documents (>30 chunks), we disable animation to prevent layout shifts
chunk, // which break auto-scroll functionality
index, interface ChunkCardProps {
totalChunks,
isCited,
isActive,
}: {
chunk: { id: number; content: string }; chunk: { id: number; content: string };
index: number; index: number;
totalChunks: number; totalChunks: number;
isCited: boolean; isCited: boolean;
isActive: boolean; isActive: boolean;
}) => { disableLayoutAnimation?: boolean;
const shouldReduceMotion = useReducedMotion(); }
return ( const ChunkCard = forwardRef<HTMLDivElement, ChunkCardProps>(
<motion.div ({ chunk, index, totalChunks, isCited, isActive, disableLayoutAnimation }, ref) => {
data-chunk-index={index} return (
initial={shouldReduceMotion ? { opacity: 1 } : { opacity: 0, y: 20 }} <div
animate={{ opacity: 1, y: 0 }} ref={ref}
transition={{ data-chunk-index={index}
type: "spring", className={cn(
stiffness: 100, "group relative rounded-2xl border-2 transition-all duration-300",
damping: 15, isCited
delay: shouldReduceMotion ? 0 : Math.min(index * 0.05, 0.3), ? "bg-linear-to-br from-primary/5 via-primary/10 to-primary/5 border-primary shadow-lg shadow-primary/10"
}} : "bg-card border-border/50 hover:border-border hover:shadow-md"
className={cn( )}
"group relative rounded-2xl border-2 transition-all duration-300", >
isCited {/* Cited indicator glow effect */}
? "bg-linear-to-br from-primary/5 via-primary/10 to-primary/5 border-primary shadow-lg shadow-primary/10"
: "bg-card border-border/50 hover:border-border hover:shadow-md"
)}
>
{/* Cited indicator glow effect */}
{isCited && (
<div className="absolute inset-0 rounded-2xl bg-primary/5 blur-xl -z-10" />
)}
{/* Header */}
<div className="flex items-center justify-between px-5 py-4 border-b border-border/50">
<div className="flex items-center gap-3">
<div
className={cn(
"flex items-center justify-center w-8 h-8 rounded-full text-sm font-semibold transition-colors",
isCited
? "bg-primary text-primary-foreground"
: "bg-muted text-muted-foreground group-hover:bg-muted/80"
)}
>
{index + 1}
</div>
<span className="text-sm text-muted-foreground">
of {totalChunks} chunks
</span>
</div>
{isCited && ( {isCited && (
<motion.div <div className="absolute inset-0 rounded-2xl bg-primary/5 blur-xl -z-10" />
initial={{ scale: 0, opacity: 0 }} )}
animate={{ scale: 1, opacity: 1 }}
transition={{ type: "spring", stiffness: 200, damping: 15, delay: 0.2 }} {/* Header */}
> <div className="flex items-center justify-between px-5 py-4 border-b border-border/50">
<div className="flex items-center gap-3">
<div
className={cn(
"flex items-center justify-center w-8 h-8 rounded-full text-sm font-semibold transition-colors",
isCited
? "bg-primary text-primary-foreground"
: "bg-muted text-muted-foreground group-hover:bg-muted/80"
)}
>
{index + 1}
</div>
<span className="text-sm text-muted-foreground">
of {totalChunks} chunks
</span>
</div>
{isCited && (
<Badge variant="default" className="gap-1.5 px-3 py-1"> <Badge variant="default" className="gap-1.5 px-3 py-1">
<Sparkles className="h-3 w-3" /> <Sparkles className="h-3 w-3" />
Cited Source Cited Source
</Badge> </Badge>
</motion.div> )}
)} </div>
</div>
{/* Content */} {/* Content */}
<div className="p-5 overflow-hidden"> <div className="p-5 overflow-hidden">
<MarkdownViewer content={chunk.content} /> <MarkdownViewer content={chunk.content} />
</div>
</div> </div>
</motion.div> );
); }
}; );
ChunkCard.displayName = "ChunkCard";
export function SourceDetailPanel({ export function SourceDetailPanel({
open, open,
@ -133,6 +120,7 @@ export function SourceDetailPanel({
children, children,
}: SourceDetailPanelProps) { }: SourceDetailPanelProps) {
const scrollAreaRef = useRef<HTMLDivElement>(null); const scrollAreaRef = useRef<HTMLDivElement>(null);
const hasScrolledRef = useRef(false); // Use ref to avoid stale closures
const [summaryOpen, setSummaryOpen] = useState(false); const [summaryOpen, setSummaryOpen] = useState(false);
const [activeChunkIndex, setActiveChunkIndex] = useState<number | null>(null); const [activeChunkIndex, setActiveChunkIndex] = useState<number | null>(null);
const [mounted, setMounted] = useState(false); const [mounted, setMounted] = useState(false);
@ -163,30 +151,92 @@ export function SourceDetailPanel({
// Find cited chunk index // Find cited chunk index
const citedChunkIndex = documentData?.chunks?.findIndex((chunk) => chunk.id === chunkId) ?? -1; const citedChunkIndex = documentData?.chunks?.findIndex((chunk) => chunk.id === chunkId) ?? -1;
// Auto-scroll to cited chunk when data loads // Simple scroll function that scrolls to a chunk by index
useEffect(() => { const scrollToChunkByIndex = useCallback((chunkIndex: number, smooth = true) => {
if (documentData?.chunks && citedChunkIndex !== -1 && !hasScrolledToCited && open) { const scrollContainer = scrollAreaRef.current;
// Wait for animations to complete then scroll if (!scrollContainer) return;
const timer = setTimeout(() => {
const chunkElement = scrollAreaRef.current?.querySelector( const viewport = scrollContainer.querySelector(
`[data-chunk-index="${citedChunkIndex}"]` '[data-radix-scroll-area-viewport]'
); ) as HTMLElement | null;
if (chunkElement) { if (!viewport) return;
chunkElement.scrollIntoView({
behavior: shouldReduceMotion ? "auto" : "smooth", const chunkElement = scrollContainer.querySelector(
block: "center", `[data-chunk-index="${chunkIndex}"]`
}); ) as HTMLElement | null;
setHasScrolledToCited(true); if (!chunkElement) return;
setActiveChunkIndex(citedChunkIndex);
} // Get positions using getBoundingClientRect for accuracy
}, 400); const viewportRect = viewport.getBoundingClientRect();
return () => clearTimeout(timer); 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 // Reset scroll state when panel closes
useEffect(() => { useEffect(() => {
if (!open) { if (!open) {
hasScrolledRef.current = false;
setHasScrolledToCited(false); setHasScrolledToCited(false);
setActiveChunkIndex(null); setActiveChunkIndex(null);
} }
@ -222,12 +272,8 @@ export function SourceDetailPanel({
}; };
const scrollToChunk = useCallback((index: number) => { const scrollToChunk = useCallback((index: number) => {
setActiveChunkIndex(index); scrollToChunkByIndex(index, true);
const chunkElement = scrollAreaRef.current?.querySelector( }, [scrollToChunkByIndex]);
`[data-chunk-index="${index}"]`
);
chunkElement?.scrollIntoView({ behavior: "smooth", block: "center" });
}, []);
const panelContent = ( const panelContent = (
<AnimatePresence mode="wait"> <AnimatePresence mode="wait">
@ -388,9 +434,9 @@ export function SourceDetailPanel({
initial={{ opacity: 0, x: -20 }} initial={{ opacity: 0, x: -20 }}
animate={{ opacity: 1, x: 0 }} animate={{ opacity: 1, x: 0 }}
transition={{ delay: 0.2 }} 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"
> >
<ScrollArea className="flex-1"> <ScrollArea className="flex-1 h-full">
<div className="p-2 pt-3 flex flex-col gap-1.5"> <div className="p-2 pt-3 flex flex-col gap-1.5">
{documentData.chunks.map((chunk, idx) => { {documentData.chunks.map((chunk, idx) => {
const isCited = chunk.id === chunkId; const isCited = chunk.id === chunkId;
@ -514,16 +560,21 @@ export function SourceDetailPanel({
{/* Chunks */} {/* Chunks */}
<div className="space-y-4"> <div className="space-y-4">
{documentData.chunks.map((chunk, idx) => ( {documentData.chunks.map((chunk, idx) => {
<ChunkCard const isCited = chunk.id === chunkId;
key={chunk.id} return (
chunk={chunk} <ChunkCard
index={idx} key={chunk.id}
totalChunks={documentData.chunks.length} ref={isCited ? citedChunkRefCallback : undefined}
isCited={chunk.id === chunkId} chunk={chunk}
isActive={activeChunkIndex === idx} index={idx}
/> totalChunks={documentData.chunks.length}
))} isCited={isCited}
isActive={activeChunkIndex === idx}
disableLayoutAnimation={documentData.chunks.length > 30}
/>
);
})}
</div> </div>
</div> </div>
</ScrollArea> </ScrollArea>

View file

@ -11,12 +11,41 @@ interface NewChatAdapterConfig {
chatId: number; 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. * Creates a ChatModelAdapter that connects to the FastAPI new_chat endpoint.
* *
* The backend expects: * The backend expects:
* - POST /api/v1/new_chat * - 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 * - Returns: SSE stream with Vercel AI SDK Data Stream Protocol
*/ */
export function createNewChatAdapter(config: NewChatAdapterConfig): ChatModelAdapter { export function createNewChatAdapter(config: NewChatAdapterConfig): ChatModelAdapter {
@ -31,7 +60,7 @@ export function createNewChatAdapter(config: NewChatAdapterConfig): ChatModelAda
throw new Error("No user message found"); throw new Error("No user message found");
} }
// Extract text content from the message // Extract text content from the last user message
let userQuery = ""; let userQuery = "";
for (const part of lastUserMessage.content) { for (const part of lastUserMessage.content) {
if (part.type === "text") { if (part.type === "text") {
@ -48,6 +77,9 @@ export function createNewChatAdapter(config: NewChatAdapterConfig): ChatModelAda
throw new Error("Not authenticated. Please log in again."); 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`, { const response = await fetch(`${backendUrl}/api/v1/new_chat`, {
method: "POST", method: "POST",
headers: { headers: {
@ -58,6 +90,7 @@ export function createNewChatAdapter(config: NewChatAdapterConfig): ChatModelAda
chat_id: config.chatId, chat_id: config.chatId,
user_query: userQuery.trim(), user_query: userQuery.trim(),
search_space_id: config.searchSpaceId, search_space_id: config.searchSpaceId,
messages: messageHistory,
}), }),
signal: abortSignal, signal: abortSignal,
}); });
@ -165,3 +198,4 @@ export function createNewChatAdapter(config: NewChatAdapterConfig): ChatModelAda
}, },
}; };
} }