Add database persistence layer and conversation tracking system

# Background / Context

When the user queries lightspeed-stack, they get a conversation ID in
the response that they can later use to retrieve the conversation
history and continue the conversation.

# Issue / Requirement / Reason for change

We've put the burden of remembering the conversation ID on the user,
which is not ideal. The user (e.g. a UI) has to store the conversation
ID somewhere, e.g. browser local storage, but this is not ideal because
it can easily get lost. It also complicated the UI implementation.

Another issue is that a user could read another user's conversation if
they only knew the ID, which is hard, but not ideal. It's not
immediately obvious that the conversation ID is a secret that should not
be shared with others.

# Solution / Feature Overview

Add a database persistence layer to lightspeed-stack to store user
conversation IDs and a bit of metadata about the conversation.

Allow users to retrieve their conversations by calling a new endpoint
`/conversations` (without an ID - with an ID is an existing endpoint
that lists the message history for a conversation), which will return a
list of conversations associated with the user, along with metadata
like creation time, last message time, and message count.

Verify that the user owns the conversation when querying or streaming a
conversation, and automatically associate the conversation with the user
in the database when the user queries or streams a conversation.

# Implementation Details

**Database Layer (src/app/database.py):**
- SQLAlchemy
- Support for both SQLite and Postgre databases
- PostgreSQL custom schema (table namespacing) support with automatic
  schema creation. This is useful because llama-stack itself needs a
  database for persisting conversation history, and we want to avoid
  conflicts with that when users use the same Postgres database as
  llama-stack.
- Automatically enable SQL statement tracing when debug logging is
  enabled

**Models (src/models/):**
- Base model class for shared SQLAlchemy configuration
- UserConversation model with indexed user_id fields for efficient querying, automatic timestamp tracking (created_at, last_message_at) and message count tracking

**Configuration Updates:**
- Extended configuration system to support database connection settings

**API:**
- Conversations endpoint without an ID now returns a list of
  conversations for the authenticated user, along with metadata
- Query/streaming endpoint automatically associates conversations
  with the user in the database, and verifies ownership of conversations

**Dependencies:**
- Added SQLAlchemy dependency

**Tests:**

More tests will be added (and unit tests will be fixed) once I get some
initial feedback on the implementation.

# Future Work

This change does not include migrations, so any future changes to the
databases require writing custom backwards compatibility code. We should
probably add Alembic or something simpler in the near future.

Author: omertuc

pyproject.toml

@@ -33,6 +33,7 @@ dependencies = [
     "starlette>=0.47.1",
     "aiohttp>=3.12.14",
     "authlib>=1.6.0",
+    "sqlalchemy>=2.0.42",
 ]
 
 [tool.pyright]

src/app/database.py

@@ -0,0 +1,124 @@
+"""Database engine management."""
+
+from pathlib import Path
+from typing import Any
+
+from sqlalchemy import create_engine, text
+from sqlalchemy.engine.base import Engine
+from sqlalchemy.orm import sessionmaker, Session
+from log import get_logger, logging
+from configuration import configuration
+from models.base import Base
+from models.config import SQLiteDatabaseConfiguration, PostgreSQLDatabaseConfiguration
+
+logger = get_logger(__name__)
+
+engine: Engine | None = None
+SessionLocal: sessionmaker | None = None
+
+
+def get_engine() -> Engine:
+    """Get the database engine. Raises an error if not initialized."""
+    if engine is None:
+        raise RuntimeError(
+            "Database engine not initialized. Call initialize_database() first."
+        )
+    return engine
+
+
+def create_tables() -> None:
+    """Create tables."""
+    Base.metadata.create_all(get_engine())
+
+
+def get_session() -> Session:
+    """Get a database session. Raises an error if not initialized."""
+    if SessionLocal is None:
+        raise RuntimeError(
+            "Database session not initialized. Call initialize_database() first."
+        )
+    return SessionLocal()
+
+
+def _create_sqlite_engine(config: SQLiteDatabaseConfiguration, **kwargs: Any) -> Engine:
+    """Create SQLite database engine."""
+    if not Path(config.db_path).parent.exists():
+        raise FileNotFoundError(
+            f"SQLite database directory does not exist: {config.db_path}"
+        )
+
+    try:
+        return create_engine(f"sqlite:///{config.db_path}", **kwargs)
+    except Exception as e:
+        logger.exception("Failed to create SQLite engine")
+        raise RuntimeError(f"SQLite engine creation failed: {e}") from e
+
+
+def _create_postgres_engine(
+    config: PostgreSQLDatabaseConfiguration, **kwargs: Any
+) -> Engine:
+    """Create PostgreSQL database engine."""
+    postgres_url = (
+        f"postgresql://{config.user}:{config.password}@"
+        f"{config.host}:{config.port}/{config.db}"
+        f"?sslmode={config.ssl_mode}&gssencmode={config.gss_encmode}"
+    )
+
+    is_custom_schema = config.namespace is not None and config.namespace != "public"
+
+    connect_args = {}
+    if is_custom_schema:
+        connect_args["options"] = f"-csearch_path={config.namespace}"
+
+    if config.ca_cert_path is not None:
+        connect_args["sslrootcert"] = str(config.ca_cert_path)
+
+    try:
+        postgres_engine = create_engine(
+            postgres_url, connect_args=connect_args, **kwargs
+        )
+    except Exception as e:
+        logger.exception("Failed to create PostgreSQL engine")
+        raise RuntimeError(f"PostgreSQL engine creation failed: {e}") from e
+
+    if is_custom_schema:
+        try:
+            with postgres_engine.connect() as connection:
+                connection.execute(
+                    text(f'CREATE SCHEMA IF NOT EXISTS "{config.namespace}"')
+                )
+                connection.commit()
+                logger.info("Schema '%s' created or already exists", config.namespace)
+        except Exception as e:
+            logger.exception("Failed to create schema '%s'", config.namespace)
+            raise RuntimeError(
+                f"Schema creation failed for '{config.namespace}': {e}"
+            ) from e
+
+    return postgres_engine
+
+
+def initialize_database() -> None:
+    """Initialize the database engine."""
+    db_config = configuration.database_configuration
+
+    global engine, SessionLocal  # pylint: disable=global-statement
+
+    # Debug print all SQL statements if our logger is at-least DEBUG level
+    echo = bool(logger.isEnabledFor(logging.DEBUG))
+
+    create_engine_kwargs = {
+        "echo": echo,
+    }
+
+    match db_config.db_type:
+        case "sqlite":
+            sqlite_config = db_config.config
+            assert isinstance(sqlite_config, SQLiteDatabaseConfiguration)
+            engine = _create_sqlite_engine(sqlite_config, **create_engine_kwargs)
+        case "postgres":
+            postgres_config = db_config.config
+            assert isinstance(postgres_config, PostgreSQLDatabaseConfiguration)
+            engine = _create_postgres_engine(postgres_config, **create_engine_kwargs)
+
+    SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)

src/app/endpoints/conversations.py

@@ -9,9 +9,16 @@
 
 from client import AsyncLlamaStackClientHolder
 from configuration import configuration
-from models.responses import ConversationResponse, ConversationDeleteResponse
+from models.responses import (
+    ConversationResponse,
+    ConversationDeleteResponse,
+    ConversationsListResponse,
+    ConversationDetails,
+)
 from auth import get_auth_dependency
-from utils.endpoints import check_configuration_loaded
+from app.database import get_session
+from models.conversations import UserConversation  # pylint: disable=ungrouped-imports
+from utils.endpoints import check_configuration_loaded, validate_conversation_ownership
 from utils.suid import check_suid
 
 logger = logging.getLogger("app.endpoints.handlers")
@@ -66,6 +73,31 @@
     },
 }
 
+conversations_list_responses: dict[int | str, dict[str, Any]] = {
+    200: {
+        "conversations": [
+            {
+                "conversation_id": "123e4567-e89b-12d3-a456-426614174000",
+                "created_at": "2024-01-01T00:00:00Z",
+                "last_message_at": "2024-01-01T00:05:00Z",
+                "message_count": 5,
+            },
+            {
+                "conversation_id": "456e7890-e12b-34d5-a678-901234567890",
+                "created_at": "2024-01-01T01:00:00Z",
+                "last_message_at": "2024-01-01T01:02:00Z",
+                "message_count": 2,
+            },
+        ]
+    },
+    503: {
+        "detail": {
+            "response": "Unable to connect to Llama Stack",
+            "cause": "Connection error.",
+        }
+    },
+}
+
 
 def simplify_session_data(session_data: dict) -> list[dict[str, Any]]:
     """Simplify session data to include only essential conversation information.
@@ -109,10 +141,63 @@ def simplify_session_data(session_data: dict) -> list[dict[str, Any]]:
     return chat_history
 
 
+@router.get("/conversations", responses=conversations_list_responses)
+def get_conversations_list_endpoint_handler(
+    auth: Any = Depends(auth_dependency),
+) -> ConversationsListResponse:
+    """Handle request to retrieve all conversations for the authenticated user."""
+    check_configuration_loaded(configuration)
+
+    user_id, _, _ = auth
+
+    logger.info("Retrieving conversations for user %s", user_id)
+
+    with get_session() as session:
+        try:
+            # Get all conversations for this user
+            user_conversations = (
+                session.query(UserConversation).filter_by(user_id=user_id).all()
+            )
+
+            # Return conversation summaries with metadata
+            conversations = [
+                ConversationDetails(
+                    conversation_id=conv.id,
+                    created_at=conv.created_at.isoformat() if conv.created_at else None,
+                    last_message_at=(
+                        conv.last_message_at.isoformat()
+                        if conv.last_message_at
+                        else None
+                    ),
+                    message_count=conv.message_count,
+                    model=conv.model,
+                )
+                for conv in user_conversations
+            ]
+
+            logger.info(
+                "Found %d conversations for user %s", len(conversations), user_id
+            )
+
+            return ConversationsListResponse(conversations=conversations)
+
+        except Exception as e:
+            logger.exception(
+                "Error retrieving conversations for user %s: %s", user_id, e
+            )
+            raise HTTPException(
+                status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+                detail={
+                    "response": "Unknown error",
+                    "cause": f"Unknown error while getting conversations for user {user_id}",
+                },
+            ) from e
+
+
 @router.get("/conversations/{conversation_id}", responses=conversation_responses)
 async def get_conversation_endpoint_handler(
     conversation_id: str,
-    _auth: Any = Depends(auth_dependency),
+    auth: Any = Depends(auth_dependency),
 ) -> ConversationResponse:
     """Handle request to retrieve a conversation by ID."""
     check_configuration_loaded(configuration)
@@ -128,6 +213,13 @@ async def get_conversation_endpoint_handler(
             },
         )
 
+    user_id, _, _ = auth
+
+    validate_conversation_ownership(
+        user_id=user_id,
+        conversation_id=conversation_id,
+    )
+
     agent_id = conversation_id
     logger.info("Retrieving conversation %s", conversation_id)
 
@@ -187,7 +279,7 @@ async def get_conversation_endpoint_handler(
 )
 async def delete_conversation_endpoint_handler(
     conversation_id: str,
-    _auth: Any = Depends(auth_dependency),
+    auth: Any = Depends(auth_dependency),
 ) -> ConversationDeleteResponse:
     """Handle request to delete a conversation by ID."""
     check_configuration_loaded(configuration)
@@ -202,6 +294,14 @@ async def delete_conversation_endpoint_handler(
                 "cause": f"Conversation ID {conversation_id} is not a valid UUID",
             },
         )
+
+    user_id, _, _ = auth
+
+    validate_conversation_ownership(
+        user_id=user_id,
+        conversation_id=conversation_id,
+    )
+
     agent_id = conversation_id
     logger.info("Deleting conversation %s", conversation_id)
 

src/app/endpoints/query.py

@@ -22,11 +22,18 @@
 from auth.interface import AuthTuple
 from client import AsyncLlamaStackClientHolder
 from configuration import configuration
+from app.database import get_session
 import metrics
+from models.conversations import UserConversation
 from models.responses import QueryResponse, UnauthorizedResponse, ForbiddenResponse
 from models.requests import QueryRequest, Attachment
 import constants
-from utils.endpoints import check_configuration_loaded, get_agent, get_system_prompt
+from utils.endpoints import (
+    check_configuration_loaded,
+    get_agent,
+    get_system_prompt,
+    validate_conversation_ownership,
+)
 from utils.mcp_headers import mcp_headers_dependency, handle_mcp_headers_with_toolgroups
 from utils.suid import get_suid
 
@@ -65,6 +72,32 @@ def is_transcripts_enabled() -> bool:
     return configuration.user_data_collection_configuration.transcripts_enabled
 
 
+def persist_user_conversation_details(
+    user_id: str, conversation_id: str, model: str
+) -> None:
+    """Associate conversation to user in the database."""
+    with get_session() as session:
+        existing_conversation = (
+            session.query(UserConversation)
+            .filter_by(id=conversation_id, user_id=user_id)
+            .first()
+        )
+
+        if not existing_conversation:
+            conversation = UserConversation(
+                id=conversation_id, user_id=user_id, model=model, message_count=1
+            )
+            session.add(conversation)
+            logger.debug(
+                "Associated conversation %s to user %s", conversation_id, user_id
+            )
+        else:
+            existing_conversation.last_message_at = datetime.now(UTC)
+            existing_conversation.message_count += 1
+
+        session.commit()
+
+
 @router.post("/query", responses=query_response)
 async def query_endpoint_handler(
     query_request: QueryRequest,
@@ -79,6 +112,11 @@ async def query_endpoint_handler(
 
     user_id, _, token = auth
 
+    if query_request.conversation_id is not None:
+        validate_conversation_ownership(
+            user_id=user_id, conversation_id=query_request.conversation_id
+        )
+
     try:
         # try to get Llama Stack client
         client = AsyncLlamaStackClientHolder().get_client()
@@ -110,6 +148,10 @@ async def query_endpoint_handler(
                 attachments=query_request.attachments or [],
             )
 
+        persist_user_conversation_details(
+            user_id=user_id, conversation_id=conversation_id, model=model_id
+        )
+
         return QueryResponse(conversation_id=conversation_id, response=response)
 
     # connection to Llama Stack server