main.py

from typing import Dict, List, Optional
from contextlib import contextmanager
import os
from pathlib import Path

import chainlit as cl
from llama_cpp import Llama
from sqlalchemy.orm import Session
from loguru import logger

from config import Config
from database import get_db, ChatMessage, cache_query
from cache import get_cache, set_cache
from logger import setup_logging

# Load configuration based on environment
config = Config.load()
setup_logging(config.dict())

class ModelConfigurationError(Exception):
    """Raised when model configuration is invalid"""
    pass

def validate_model_path(model_path: str) -> str:
    """Validate that model path exists and is accessible"""
    path = Path(model_path)
    if not path.exists():
        raise ModelConfigurationError(f"Model file not found at {model_path}")
    if not path.is_file():
        raise ModelConfigurationError(f"Model path {model_path} is not a file")
    if not os.access(path, os.R_OK):
        raise ModelConfigurationError(f"Model file {model_path} is not readable")
    return str(path)

def initialize_llm() -> Llama:
    """Initialize LLM with validated configuration"""
    model_config = config.model
    validated_path = validate_model_path(model_config.model_path)
    
    try:
        return Llama(
            model_path=validated_path,
            n_ctx=model_config.n_ctx,
            n_gpu_layers=model_config.n_gpu_layers,
            use_mlock=model_config.use_mlock,
            chat_format=model_config.chat_format
        )
    except Exception as e:
        logger.error(f"Failed to initialize LLM: {str(e)}")
        raise ModelConfigurationError(f"Failed to initialize LLM: {str(e)}")

logger.info("Initializing LLM with configuration: {}", config.model)
llm = initialize_llm()

class ChatCompletionError(Exception):
    """Raised when chat completion fails"""
    pass

async def create_chat_completion(memory: List[str]):
    try:
        return llm.create_chat_completion(
            stream=True,
            messages=[
            {
                "role": "system",
                "content": "You are a helpful assistant",
            },
            *memory
        ],
        response_format={
            "type": "text"
        },
        temperature=0,
    )
    except Exception as e:
        logger.error(f"Chat completion failed: {str(e)}")
        raise ChatCompletionError(f"Failed to generate response: {str(e)}")

@cl.on_chat_start
async def on_chat_start():
    try:
        memory = []
        cl.user_session.set("memory", memory)
        logger.info("New chat session started")
    except Exception as e:
        logger.error(f"Failed to initialize chat session: {str(e)}")
        raise

@cl.on_message
async def main(message: cl.Message):
    try:
        msg = cl.Message(content="", author="Assistant")
        memory = update_memory("user", message.content)
        
        try:
            output = await create_chat_completion(memory)
            response = ""
            for chunk in output:
                try:
                    delta = chunk['choices'][0]['delta']
                    if 'content' in delta:
                        response += delta['content']
                        await msg.stream_token(delta['content'])
                except (KeyError, IndexError) as e:
                    logger.error(f"Invalid response chunk format: {str(e)}")
                    raise ChatCompletionError(f"Invalid response format: {str(e)}")
                
            update_memory("assistant", response)
            await msg.send()
            
        except ChatCompletionError as e:
            error_msg = f"I apologize, but I encountered an error: {str(e)}"
            await cl.Message(content=error_msg, author="Assistant").send()
            logger.error(f"Chat completion error: {str(e)}")
            
    except Exception as e:
        logger.error(f"Unexpected error in message handler: {str(e)}")
        await cl.Message(
            content="I apologize, but something went wrong. Please try again later.",
            author="Assistant"
        ).send()


def validate_message(role: str, content: str) -> None:
    """Validate chat message parameters"""
    valid_roles = {"user", "assistant", "system"}
    if role not in valid_roles:
        raise ValueError(f"Invalid role: {role}. Must be one of {valid_roles}")
    if not content or not isinstance(content, str):
        raise ValueError("Message content must be a non-empty string")
    if len(content) > 4096:  # Reasonable max length
        raise ValueError("Message content exceeds maximum length of 4096 characters")

@cache_query(ttl=60)
def get_recent_messages(db: Session, limit: int = 10) -> List[ChatMessage]:
    """Get recent messages with caching"""
    return db.query(ChatMessage).order_by(ChatMessage.timestamp.desc()).limit(limit).all()

def update_memory(role: str, content: str) -> List[Dict[str, str]]:
    """ Handle conversation memory with database persistence, pooling and caching """
    validate_message(role, content)
    
    session_id = cl.user_session.get("session_id", "default")
    if not session_id:
        raise ValueError("Invalid session ID")
        
    cache_key = f"memory:{session_id}"
    
    # Try to get memory from cache first
    memory = get_cache(cache_key)
    if memory is None:
        memory = cl.user_session.get("memory")
        if not isinstance(memory, list):
            memory = []  # Reset if invalid
            
    memory.append({"role": role, "content": content})
    
    # Persist message to database using connection pool with improved error handling and batch operations
    with get_db() as db:
        try:
            # Use bulk insert for better performance
            truncated_content = content[:150] if role == "assistant" else content
            db_message = ChatMessage(role=role, content=truncated_content)
            db.bulk_save_objects([db_message])
            db.commit()
            
            # Update cached recent messages
            get_recent_messages.cache_clear() if hasattr(get_recent_messages, 'cache_clear') else None
        except Exception as e:
            db.rollback()
            logger.error(f"Failed to persist messages: {e}")
            raise

    # Update cache and session memory based on config
    memory = memory[-2:]  # Keep only last 2 messages
    if not set_cache(cache_key, memory, config.cache.ttl):
        cl.user_session.set("memory", memory)  # Fallback to session if cache fails
    else:
        cl.user_session.set("memory", memory)
    return memory

if __name__ == "__main__":
    from chainlit.cli import run_chainlit
    run_chainlit(__file__)