MemoryBear/api/app/repositories/prompt_optimizer_repository.py

import uuid
from typing import Optional

from sqlalchemy.orm import Session

from app.core.logging_config import get_db_logger
from app.models.prompt_optimizer_model import (
    PromptOptimizerModelConfig,
    PromptOptimizerSession, PromptOptimizerSessionHistory, RoleType
)

db_logger = get_db_logger()


class PromptOptimizerModelConfigRepository:
    """Repository for managing prompt optimizer model configurations."""

    def __init__(self, db: Session):
        self.db = db

    def get_by_tenant_id(self, tenant_id: uuid.UUID) -> Optional[PromptOptimizerModelConfig]:
        """
        Retrieve the prompt optimizer model configuration for a specific tenant.

        Args:
           tenant_id (uuid.UUID): The unique identifier of the tenant.

        Returns:
           Optional[PromptOptimizerModelConfig]: The model configuration if found, else None.
        """
        db_logger.debug(f"Get prompt optimization model configuration: tenant_id={tenant_id}")

        try:
            config = self.db.query(PromptOptimizerModelConfig).filter(
                PromptOptimizerModelConfig.tenant_id == tenant_id,
                # PromptOptimizerModelConfig.model_id == model_id
            ).first()
            if config:
                db_logger.debug(f"Prompt optimization model configuration found: (ID: {config.id})")
            else:
                db_logger.debug(f"Prompt optimization model configuration not found: tenant_id={tenant_id}")
            return config
        except Exception as e:
            db_logger.error(
                f"Error retrieving prompt optimization model configuration: tenant_id={tenant_id} - {str(e)}")
            raise

    def get_by_config_id(self, tenant_id: uuid.UUID, config_id: uuid.UUID) -> Optional[PromptOptimizerModelConfig]:
        """
        Retrieve a specific prompt optimizer model configuration by config ID and tenant ID.

        Args:
            tenant_id (uuid.UUID): The unique identifier of the tenant.
            config_id (uuid.UUID): The unique identifier of the model configuration.

        Returns:
            Optional[PromptOptimizerModelConfig]: The model configuration if found, else None.
        """
        db_logger.debug(f"Get prompt optimization model configuration: config_id={config_id}, tenant_id={tenant_id}")
        try:
            model = self.db.query(PromptOptimizerModelConfig).filter(
                PromptOptimizerModelConfig.tenant_id == tenant_id,
                PromptOptimizerModelConfig.id == config_id
            ).first()
            if model:
                db_logger.debug(f"Prompt optimization model configuration found: (ID: {model.id})")
            else:
                db_logger.debug(f"Prompt optimization model configuration not found: config_id={config_id}")
            return model
        except Exception as e:
            db_logger.error(
                f"Error retrieving prompt optimization model configuration: model_id={config_id} - {str(e)}")
            raise

    def create_or_update(
            self,
            config_id: uuid.UUID,
            tenant_id: uuid.UUID,
            system_prompt: str,
    ) -> Optional[PromptOptimizerModelConfig]:
        """
        Create a new or update an existing prompt optimizer model configuration.

        If a configuration with the given config_id exists, it updates its system_prompt.
        Otherwise, it creates a new configuration record.

        Args:
            config_id (uuid.UUID): The unique identifier for the configuration.
            tenant_id (uuid.UUID): The tenant's unique identifier.
            system_prompt (str): The system prompt content for prompt optimization.

        Returns:
            Optional[PromptOptimizerModelConfig]: The created or updated model configuration.
        """
        db_logger.debug(f"Create/Update prompt optimization model configuration: tenant_id={tenant_id}")
        existing_config = self.get_by_config_id(tenant_id, config_id)

        if existing_config:
            existing_config.system_prompt = system_prompt
            self.db.commit()
            self.db.refresh(existing_config)
            db_logger.debug(f"Prompt optimization model configuration update: ID:{config_id}")
            return existing_config
        else:
            config = PromptOptimizerModelConfig(
                id=config_id,
                # model_id=model_id,
                tenant_id=tenant_id,
                system_prompt=system_prompt
            )
            self.db.add(config)
            self.db.commit()
            self.db.refresh(config)
            db_logger.debug(f"Prompt optimization model configuration created: ID:{config.id}")
            return config


class PromptOptimizerSessionRepository:
    """Repository for managing prompt optimization sessions and session history."""

    def __init__(self, db: Session):
        self.db = db

    def create_session(
            self,
            tenant_id: uuid.UUID,
            user_id: uuid.UUID
    ) -> PromptOptimizerSession:
        """
        Create a new prompt optimization session for a user and app.

        Args:
            tenant_id (uuid.UUID): The unique identifier of the tenant.
            user_id (uuid.UUID): The unique identifier of the user.

        Returns:
            PromptOptimizerSession: The newly created session object.
        """
        db_logger.debug(f"Create prompt optimization session: tenant_id={tenant_id}, user_id={user_id}")
        try:
            session = PromptOptimizerSession(
                tenant_id=tenant_id,
                user_id=user_id,
            )
            self.db.add(session)
            self.db.commit()
            self.db.refresh(session)
            db_logger.debug(f"Prompt optimization session created: ID:{session.id}")
            return session
        except Exception as e:
            db_logger.error(f"Error creating prompt optimization session: user_id={user_id} - {str(e)}")
            raise

    def get_session_history(
            self,
            session_id: uuid.UUID,
            user_id: uuid.UUID
    ) -> list[type[PromptOptimizerSessionHistory]]:
        """
        Retrieve all message history of a specific prompt optimization session.

        Args:
            session_id (uuid.UUID): The unique identifier of the session.
            user_id (uuid.UUID): The unique identifier of the user.

        Returns:
            list[PromptOptimizerSessionHistory]: A list of session history records
            ordered by creation time ascending.
        """
        db_logger.debug(f"Get prompt optimization session history: "
                        f"user_id={user_id}, session_id={session_id}")

        try:
            # First get the internal session ID from the session list table
            session = self.db.query(PromptOptimizerSession).filter(
                PromptOptimizerSession.id == session_id,
                PromptOptimizerSession.user_id == user_id
            ).first()
            
            if not session:
                return []
            
            history = self.db.query(PromptOptimizerSessionHistory).filter(
                PromptOptimizerSessionHistory.session_id == session.id,
                PromptOptimizerSessionHistory.user_id == user_id
            ).order_by(PromptOptimizerSessionHistory.created_at.asc()).all()
            return history
        except Exception as e:
            db_logger.error(f"Error retrieving prompt optimization session history: session_id={session_id} - {str(e)}")
            raise

    def create_message(
            self,
            tenant_id: uuid.UUID,
            session_id: uuid.UUID,
            user_id: uuid.UUID,
            role: RoleType,
            content: str,
    ) -> PromptOptimizerSessionHistory:
        """
        Create a new message in the session history.

        This method is a placeholder for future implementation.
        """
        try:
            # Get the session to ensure it exists and belongs to the user
            session = self.db.query(PromptOptimizerSession).filter(
                PromptOptimizerSession.id == session_id,
                PromptOptimizerSession.user_id == user_id,
                PromptOptimizerSession.tenant_id == tenant_id
            ).first()
            
            if not session:
                db_logger.error(f"Session {session_id} not found for user {user_id}")
                raise ValueError(f"Session {session_id} not found for user {user_id}")
            
            message = PromptOptimizerSessionHistory(
                tenant_id=tenant_id,
                session_id=session.id,
                user_id=user_id,
                role=role.value,
                content=content,
            )
            self.db.add(message)
            self.db.commit()
            return message
        except Exception as e:
            db_logger.error(f"Error creating prompt optimization session history: session_id={session_id} - {str(e)}")
            raise