from logging import getLogger from typing import Any from nanoid import generate as generate_nanoid from sqlalchemy import ColumnElement, Select, and_, func, select, text, update from sqlalchemy.ext.asyncio import AsyncSession from src import models, schemas from src.config import settings from src.embedding_client import embedding_client from src.utils.filter import apply_filter from src.vector_store import VectorRecord, get_vector_store, upsert_with_retry from .session import get_or_create_session logger = getLogger(__name__) def _apply_token_limit( base_conditions: list[ColumnElement[Any]], token_limit: int ) -> Select[tuple[models.Message]]: """ Helper function to apply token limit logic to a message query. Creates a subquery that calculates running sum of tokens for most recent messages and returns a select statement that joins with this subquery to limit results based on token count. Args: base_conditions: List of conditions to apply to the base query token_limit: Maximum number of tokens to include in the messages Returns: Select statement with token limit applied """ # Create a subquery that calculates running sum of tokens for most recent messages token_subquery = ( select( models.Message.id, func.sum(models.Message.token_count) .over(order_by=models.Message.id.desc()) .label("running_token_sum"), ) .where(*base_conditions) .subquery() ) # Select Message objects where running sum doesn't exceed token_limit return ( select(models.Message) .join(token_subquery, models.Message.id == token_subquery.c.id) .where(token_subquery.c.running_token_sum <= token_limit) ) async def create_messages( db: AsyncSession, messages: list[schemas.MessageCreate], workspace_name: str, session_name: str, ) -> list[models.Message]: """ Bulk create messages for a session while maintaining order. Args: db: Database session messages: List of messages to create workspace_name: Name of the workspace session_name: Name of the session to create messages in Returns: List of created message objects """ # Get or create session with peers in messages list peers = {message.peer_name: schemas.SessionPeerConfig() for message in messages} await get_or_create_session( db, session=schemas.SessionCreate(name=session_name, peers=peers), workspace_name=workspace_name, ) await db.execute(text("SET LOCAL lock_timeout = '5s'")) await db.execute( text( "SELECT pg_advisory_xact_lock(hashtext(:workspace_name), hashtext(:session_name))" ), {"workspace_name": workspace_name, "session_name": session_name}, ) # Get the last sequence number on a session - uses (workspace_name, session_name, seq_in_session) index last_seq = ( await db.scalar( select(models.Message.seq_in_session) .where( models.Message.workspace_name == workspace_name, models.Message.session_name == session_name, ) .order_by(models.Message.seq_in_session.desc()) .limit(1) ) or 0 ) # Create list of message objects (this will trigger the before_insert event) message_objects: list[models.Message] = [] for offset, message in enumerate(messages, start=1): message_seq_in_session = last_seq + offset message_obj = models.Message( session_name=session_name, peer_name=message.peer_name, content=message.content, h_metadata=message.metadata or {}, workspace_name=workspace_name, public_id=generate_nanoid(), token_count=len(message.encoded_message), created_at=message.created_at, # Use provided created_at if available seq_in_session=message_seq_in_session, ) message_objects.append(message_obj) db.add_all(message_objects) # Commit here to release the advisory lock before generating embeddings await db.commit() try: if settings.EMBED_MESSAGES: encoded_message_lookup = { msg.public_id: orig_msg.encoded_message for msg, orig_msg in zip(message_objects, messages, strict=True) } id_resource_dict = { message.public_id: ( message.content, encoded_message_lookup[message.public_id], ) for message in message_objects } embedding_dict = await embedding_client.batch_embed(id_resource_dict) # Get vector store and namespace for this workspace's messages vector_store = get_vector_store() namespace = vector_store.get_vector_namespace("message", workspace_name) # Create MessageEmbedding entries embedding_objects: list[models.MessageEmbedding] = [] # Check if pgvector is being used (primary or secondary) # If so, write embeddings to ORM since pgvector relies on postgres # Otherwise, store in memory for vector store upsert only pgvector_in_use = ( settings.VECTOR_STORE.PRIMARY_TYPE == "pgvector" or settings.VECTOR_STORE.SECONDARY_TYPE == "pgvector" ) for message_obj in message_objects: embeddings = embedding_dict.get(message_obj.public_id, []) for chunk_idx, embedding in enumerate(embeddings): if pgvector_in_use: # Create MessageEmbedding record for pgvector storage embedding_obj = models.MessageEmbedding( content=message_obj.content, message_id=message_obj.public_id, workspace_name=workspace_name, session_name=session_name, peer_name=message_obj.peer_name, sync_state="pending", embedding=embedding, ) # Track chunk index in-memory (not persisted to avoid HNSW recomputation) embedding_obj._chunk_index = chunk_idx embedding_objects.append(embedding_obj) else: # pgvector not in use: don't create MessageEmbedding, track in-memory only # Create a minimal object to hold metadata for vector store upsert embedding_obj = models.MessageEmbedding( content=message_obj.content, message_id=message_obj.public_id, workspace_name=workspace_name, session_name=session_name, peer_name=message_obj.peer_name, sync_state="pending", ) embedding_obj._pending_embedding = embedding embedding_obj._chunk_index = chunk_idx embedding_objects.append(embedding_obj) # Add MessageEmbedding rows to database only if pgvector in use if embedding_objects and pgvector_in_use: db.add_all(embedding_objects) await db.flush() # Track embedding IDs for sync state updates embedding_ids = [emb.id for emb in embedding_objects] else: embedding_ids = [] # Build vector records with {message_id}_{chunk_index} as vector ID vector_records: list[VectorRecord] = [] for emb in embedding_objects: # Always use {message_id}_{chunk_index} as vector ID (all stores) vector_id = f"{emb.message_id}_{emb._chunk_index}" # Get embedding from appropriate source if pgvector_in_use and emb.embedding is not None: embedding_data = [float(x) for x in emb.embedding] elif ( hasattr(emb, "_pending_embedding") and emb._pending_embedding is not None ): embedding_data = list(emb._pending_embedding) else: continue vector_records.append( VectorRecord( id=vector_id, embedding=embedding_data, metadata={ "message_id": emb.message_id, "session_name": emb.session_name, "peer_name": emb.peer_name, }, ) ) await db.commit() # Upsert to vector store with retry and update sync state if vector_records: try: result = await upsert_with_retry( vector_store, namespace, vector_records ) # Only update MessageEmbedding sync state if pgvector is in use if pgvector_in_use and embedding_ids: if result is not None and result.secondary_ok is False: # Partial success: primary has data but secondary doesn't logger.warning( "Partial sync for message embeddings: %s", result.secondary_error, ) await db.execute( update(models.MessageEmbedding) .where(models.MessageEmbedding.id.in_(embedding_ids)) .values( sync_attempts=models.MessageEmbedding.sync_attempts + 1, last_sync_at=func.now(), ) ) await db.commit() else: # Success: both primary and secondary stores have the data await db.execute( update(models.MessageEmbedding) .where(models.MessageEmbedding.id.in_(embedding_ids)) .values( sync_state="synced", last_sync_at=func.now(), sync_attempts=0, ) ) await db.commit() except Exception as e: # Total failure: primary write failed after retries logger.error(f"Failed to upsert message vectors after retries: {e}") # Only update MessageEmbedding sync state if pgvector is in use if pgvector_in_use and embedding_ids: await db.execute( update(models.MessageEmbedding) .where(models.MessageEmbedding.id.in_(embedding_ids)) .values( sync_attempts=models.MessageEmbedding.sync_attempts + 1, last_sync_at=func.now(), ) ) await db.commit() except Exception: logger.exception( "Failed to generate message embeddings for %s messages in workspace %s and session %s.", len(message_objects), workspace_name, session_name, ) return message_objects async def get_messages( workspace_name: str, session_name: str, reverse: bool | None = False, filters: dict[str, Any] | None = None, token_limit: int | None = None, message_count_limit: int | None = None, ) -> Select[tuple[models.Message]]: """ Get messages from a session. If token_limit is provided, the n most recent messages with token count adding up to the limit will be returned. If message_count_limit is provided, the n most recent messages will be returned. If both are provided, message_count_limit will be used. Args: workspace_name: Name of the workspace session_name: Name of the session reverse: Whether to reverse the order of messages filters: Filter to apply to the messages token_limit: Maximum number of tokens to include in the messages message_count_limit: Maximum number of messages to include Returns: Select statement for the messages """ # Base query with workspace and session filters base_conditions = [ models.Message.workspace_name == workspace_name, models.Message.session_name == session_name, ] # Apply message count limit first (takes precedence over token limit) if message_count_limit is not None: stmt = select(models.Message).where(*base_conditions) stmt = apply_filter(stmt, models.Message, filters) # For message count limit, we want the most recent N messages # So we order by id desc to get most recent, then apply limit stmt = stmt.order_by(models.Message.id.desc()).limit(message_count_limit) # Apply final ordering based on reverse parameter if reverse: stmt = stmt.order_by(models.Message.id.desc()) else: stmt = stmt.order_by(models.Message.id.asc()) elif token_limit is not None: # Apply token limit logic using helper function stmt = _apply_token_limit(base_conditions, token_limit) stmt = apply_filter(stmt, models.Message, filters) # Apply final ordering based on reverse parameter if reverse: stmt = stmt.order_by(models.Message.id.desc()) else: stmt = stmt.order_by(models.Message.id.asc()) else: # Default case - no limits applied stmt = select(models.Message).where(*base_conditions) stmt = apply_filter(stmt, models.Message, filters) if reverse: stmt = stmt.order_by(models.Message.id.desc()) else: stmt = stmt.order_by(models.Message.id.asc()) return stmt async def get_messages_id_range( db: AsyncSession, workspace_name: str, session_name: str, start_id: int = 0, end_id: int | None = None, token_limit: int | None = None, ) -> list[models.Message]: """ Get messages from a session by primary key ID range. If end_id is not provided, all messages after and including start_id will be returned. If start_id is not provided, start will be beginning of session. Note: list is *inclusive* of the end_id message and start_id message. Args: db: Database session workspace_name: Name of the workspace session_name: Name of the session start_id: Primary key ID of the first message to return end_id: Primary key ID of the last message (exclusive) Returns: List of messages """ if start_id < 0 or (end_id is not None and (start_id >= end_id or end_id <= 0)): return [] base_conditions = [ models.Message.workspace_name == workspace_name, models.Message.session_name == session_name, ] if end_id: base_conditions.append( and_(models.Message.id >= start_id, models.Message.id < end_id) ) else: base_conditions.append(models.Message.id >= start_id) if token_limit: # Apply token limit logic using helper function stmt = _apply_token_limit(base_conditions, token_limit) stmt = stmt.order_by(models.Message.id) else: stmt = select(models.Message).where(*base_conditions) result = await db.execute(stmt) return list(result.scalars().all()) async def get_messages_by_seq_range( db: AsyncSession, workspace_name: str, session_name: str, start_seq: int = 1, end_seq: int | None = None, ) -> list[models.Message]: """ Get messages from a session by seq_in_session range. This is useful for getting the last N messages in a session. Args: db: Database session workspace_name: Name of the workspace session_name: Name of the session start_seq: Sequence number of the first message to return (inclusive) end_seq: Sequence number of the last message to return (inclusive) Returns: List of messages ordered by seq_in_session """ if start_seq < 1 or (end_seq is not None and start_seq > end_seq): return [] base_conditions = [ models.Message.workspace_name == workspace_name, models.Message.session_name == session_name, ] if end_seq is not None: base_conditions.append( and_( models.Message.seq_in_session >= start_seq, models.Message.seq_in_session <= end_seq, ) ) else: base_conditions.append(models.Message.seq_in_session >= start_seq) stmt = ( select(models.Message) .where(*base_conditions) .order_by(models.Message.seq_in_session.asc()) ) result = await db.execute(stmt) return list(result.scalars().all()) async def get_message_seq_in_session( db: AsyncSession, workspace_name: str, session_name: str, message_id: int, ) -> int: """ Get the sequence number of a message within a session. Args: db: Database session session_name: Name of the session message_id: Primary key ID of the message Returns: The sequence number of the message (1-indexed) """ stmt = ( select(models.Message.seq_in_session) .where(models.Message.workspace_name == workspace_name) .where(models.Message.session_name == session_name) .where(models.Message.id == message_id) ) seq: int | None = await db.scalar(stmt) return int(seq) if seq is not None else 0 async def get_message( db: AsyncSession, workspace_name: str, session_name: str, message_id: str, ) -> models.Message | None: stmt = ( select(models.Message) .where(models.Message.workspace_name == workspace_name) .where(models.Message.session_name == session_name) .where(models.Message.public_id == message_id) ) result = await db.execute(stmt) return result.scalar_one_or_none() async def update_message( db: AsyncSession, message: schemas.MessageUpdate, workspace_name: str, session_name: str, message_id: str, ) -> bool: honcho_message = await get_message( db, workspace_name=workspace_name, session_name=session_name, message_id=message_id, ) if honcho_message is None: raise ValueError("Message not found or does not belong to user") if ( message.metadata is not None ): # Need to explicitly be there won't make it empty by default honcho_message.h_metadata = message.metadata await db.commit() # await db.refresh(honcho_message) return honcho_message