main.py

# main.py - FastAPI version of PuppyCompanion
import os
import logging
import json
import asyncio
from datetime import datetime
from typing import List, Dict, Any
from contextlib import asynccontextmanager
from fastapi import FastAPI, WebSocket, WebSocketDisconnect, HTTPException
from fastapi.staticfiles import StaticFiles
from fastapi.responses import HTMLResponse, FileResponse
from pydantic import BaseModel
from dotenv import load_dotenv

# Import your existing modules
from rag_system import RAGSystem
from agent_workflow import AgentWorkflow

# Load environment variables
load_dotenv()

# Logging configuration
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# Global variables
global_agent = None
global_qdrant_client = None
global_retriever = None
global_documents = None
initialization_completed = False

# Path to preprocessed data
PREPROCESSED_CHUNKS_PATH = "all_books_preprocessed_chunks.json"

# Pydantic models
class QuestionRequest(BaseModel):
    question: str

class ChatResponse(BaseModel):
    response: str
    sources: List[Dict[str, Any]] = []
    tool_used: str = ""

# WebSocket connection manager
class ConnectionManager:
    def __init__(self):
        self.active_connections: List[WebSocket] = []

    async def connect(self, websocket: WebSocket):
        await websocket.accept()
        self.active_connections.append(websocket)

    def disconnect(self, websocket: WebSocket):
        self.active_connections.remove(websocket)

    async def send_log(self, message: str, log_type: str = "info"):
        timestamp = datetime.now().strftime("%H:%M:%S")
        log_data = {
            "timestamp": timestamp,
            "message": message,
            "type": log_type
        }
        
        for connection in self.active_connections:
            try:
                await connection.send_text(json.dumps(log_data))
            except:
                pass

manager = ConnectionManager()

def load_preprocessed_chunks(file_path="all_books_preprocessed_chunks.json"):
    """Load preprocessed chunks from a JSON file."""
    global global_documents
    
    if global_documents is not None:
        logger.info("Using cached document chunks")
        return global_documents
        
    logger.info(f"Loading preprocessed chunks from {file_path}")
    try:
        with open(file_path, 'r', encoding='utf-8') as f:
            data = json.load(f)
            
        from langchain_core.documents import Document
        documents = []
        for item in data:
            doc = Document(
                page_content=item['page_content'],
                metadata=item['metadata']
            )
            documents.append(doc)
            
        logger.info(f"Loaded {len(documents)} document chunks")
        global_documents = documents
        return documents
    except Exception as e:
        logger.error(f"Error loading preprocessed chunks: {str(e)}")
        raise

def initialize_retriever(documents):
    """Create a retriever from documents using a shared Qdrant client."""
    global global_qdrant_client, global_retriever
    
    # Return existing retriever if already initialized
    if global_retriever is not None:
        logger.info("Using existing global retriever")
        return global_retriever
        
    logger.info("Creating retriever from documents")
    try:
        # Use langchain_qdrant to create a vector store
        from qdrant_client import QdrantClient
        from langchain_qdrant import QdrantVectorStore
        from langchain_openai import OpenAIEmbeddings
        
        # Create embedding object
        embeddings = OpenAIEmbeddings()
        logger.info("Created OpenAI embeddings object")
        
        # Create a persistent path for embeddings storage
        qdrant_path = "/tmp/qdrant_storage"
        logger.info(f"Using persistent Qdrant storage path: {qdrant_path}")
        
        # Create directory for Qdrant storage
        os.makedirs(qdrant_path, exist_ok=True)
        
        # Create or reuse global Qdrant client
        if global_qdrant_client is None:
            client = QdrantClient(path=qdrant_path)
            global_qdrant_client = client
            logger.info("Created new global Qdrant client with persistent storage")
        else:
            client = global_qdrant_client
            logger.info("Using existing global Qdrant client")
        
        # Check if collection already exists
        try:
            collections = client.get_collections()
            collection_exists = any(collection.name == "puppies" for collection in collections.collections)
            logger.info(f"Collection 'puppies' exists: {collection_exists}")
        except Exception as e:
            collection_exists = False
            logger.info(f"Could not check collections, assuming none exist: {e}")
        
        # OpenAI embeddings dimension
        embedding_dim = 1536
        
        # Create collection only if it doesn't exist
        if not collection_exists:
            from qdrant_client.http import models
            client.create_collection(
                collection_name="puppies",
                vectors_config=models.VectorParams(
                    size=embedding_dim,
                    distance=models.Distance.COSINE
                )
            )
            logger.info("Created new collection 'puppies'")
        else:
            logger.info("Using existing collection 'puppies'")
        
        # Create vector store
        vector_store = QdrantVectorStore(
            client=client,
            collection_name="puppies",
            embedding=embeddings
        )
        
        # Add documents only if collection was just created (to avoid duplicates)
        if not collection_exists:
            vector_store.add_documents(documents)
            logger.info(f"Added {len(documents)} documents to vector store")
        else:
            logger.info("Using existing embeddings in vector store")
        
        # Create retriever
        retriever = vector_store.as_retriever(search_kwargs={"k": 5})
        logger.info("Created retriever")
        
        # Store global retriever
        global_retriever = retriever
        
        return retriever
    except Exception as e:
        logger.error(f"Error creating retriever: {str(e)}")
        raise

async def initialize_system():
    """Initialize the RAG system and agent"""
    global global_agent, initialization_completed
    
    if initialization_completed:
        return global_agent
    
    await manager.send_log("Starting system initialization...", "info")
    
    try:
        # Load documents
        await manager.send_log("Loading document chunks...", "info")
        documents = load_preprocessed_chunks()
        await manager.send_log(f"Loaded {len(documents)} document chunks", "success")
        
        # Create retriever
        await manager.send_log("Creating retriever...", "info")
        retriever = initialize_retriever(documents)
        await manager.send_log("Retriever ready", "success")
        
        # Create RAG system
        await manager.send_log("Setting up RAG system...", "info")
        rag_system = RAGSystem(retriever)
        rag_tool = rag_system.create_rag_tool()
        await manager.send_log("RAG system ready", "success")
        
        # Create agent workflow
        await manager.send_log("Initializing agent workflow...", "info")
        agent = AgentWorkflow(rag_tool)
        await manager.send_log("Agent workflow ready", "success")
        
        global_agent = agent
        initialization_completed = True
        
        await manager.send_log("System initialization completed!", "success")
        return agent
        
    except Exception as e:
        await manager.send_log(f"Error during initialization: {str(e)}", "error")
        raise

@asynccontextmanager
async def lifespan(app: FastAPI):
    """Manage application lifespan"""
    # Startup
    try:
        await initialize_system()
        logger.info("System initialized successfully")
    except Exception as e:
        logger.error(f"Failed to initialize system: {e}")
        raise  # ⚠️ IMPORTANT: Arrêter l'application si l'initialisation échoue
    
    yield
    
    # Shutdown - cleanup if needed
    logger.info("Application shutdown")

# FastAPI app with lifespan
app = FastAPI(
    title="PuppyCompanion", 
    description="AI Assistant for Puppy Care",
    lifespan=lifespan
)

@app.get("/", response_class=HTMLResponse)
async def get_index():
    """Serve the main HTML page"""
    return FileResponse("static/index.html")

@app.get("/favicon.ico")
async def get_favicon():
    """Return a 204 No Content for favicon to avoid 404 errors"""
    from fastapi import Response
    return Response(status_code=204)

@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    """WebSocket endpoint for real-time logs"""
    await manager.connect(websocket)
    try:
        while True:
            # Keep connection alive
            await asyncio.sleep(1)
    except WebSocketDisconnect:
        manager.disconnect(websocket)

@app.post("/chat", response_model=ChatResponse)
async def chat_endpoint(request: QuestionRequest):
    """Main chat endpoint"""
    global global_agent
    
    if not initialization_completed or not global_agent:
        await manager.send_log("System not initialized, starting initialization...", "warning")
        try:
            global_agent = await initialize_system()
        except Exception as e:
            raise HTTPException(status_code=500, detail="System initialization failed")
    
    question = request.question
    await manager.send_log(f"New question: {question}", "info")
    
    try:
        # Process question with agent
        await manager.send_log("Processing with agent workflow...", "info")
        result = global_agent.process_question(question)
        
        # Extract response and metadata
        response_content = global_agent.get_final_response(result)
        
        # Parse tool usage and send detailed info to debug console
        tool_used = "Unknown"
        sources = []
        
        if "[Using RAG tool]" in response_content:
            tool_used = "RAG Tool"
            await manager.send_log("Used RAG tool - Knowledge base search", "tool")
            
            # Send detailed RAG chunks to debug console
            if "context" in result:
                await manager.send_log(f"Retrieved {len(result['context'])} chunks from knowledge base:", "info")
                for i, doc in enumerate(result["context"], 1):
                    source_name = doc.metadata.get('source', 'Unknown')
                    page = doc.metadata.get('page', 'N/A')
                    chapter = doc.metadata.get('chapter', '')
                    
                    # Create detailed chunk info for console
                    if chapter:
                        chunk_header = f"Chunk {i} - {source_name} (Chapter: {chapter}, Page: {page})"
                    else:
                        chunk_header = f"Chunk {i} - {source_name} (Page: {page})"
                    
                    await manager.send_log(chunk_header, "source")
                    
                    # Send chunk content preview
                    content_preview = doc.page_content[:200] + "..." if len(doc.page_content) > 200 else doc.page_content
                    await manager.send_log(f"Content: {content_preview}", "chunk")
                    
                    # Collect for sources array (minimal info)
                    source_info = {
                        "chunk": i,
                        "source": source_name,
                        "page": page,
                        "chapter": chapter
                    }
                    sources.append(source_info)
                    
        elif "[Using Tavily tool]" in response_content:
            tool_used = "Tavily Tool"
            await manager.send_log("Used Tavily tool - Web search", "tool")
            
            # Extract Tavily sources from response content and send to debug console
            lines = response_content.split('\n')
            tavily_sources_count = 0
            
            for line in lines:
                line_stripped = line.strip()
                # Look for Tavily source lines like "- *Source 1 - domain.com: Title*"
                if (line_stripped.startswith('- *Source') and ':' in line_stripped):
                    tavily_sources_count += 1
                    # Extract and format for debug console
                    try:
                        # Remove markdown formatting for clean display
                        clean_source = line_stripped.replace('- *', '').replace('*', '')
                        await manager.send_log(f"{clean_source}", "source")
                    except:
                        await manager.send_log(f"{line_stripped}", "source")
                        
            if tavily_sources_count > 0:
                await manager.send_log(f"Found {tavily_sources_count} web sources", "info")
            else:
                await manager.send_log("Searched the web for current information", "source")
            
        elif "out of scope" in response_content.lower():
            tool_used = "Out of Scope"
            await manager.send_log("Question outside scope (not dog-related)", "warning")
        
        # Clean response content - REMOVE ALL source references for mobile interface
        clean_response = response_content
        
        # Remove tool markers
        clean_response = clean_response.replace("[Using RAG tool]", "").replace("[Using Tavily tool]", "").strip()
        
        # Remove ALL source-related lines with comprehensive patterns
        lines = clean_response.split('\n')
        cleaned_lines = []
        for line in lines:
            line_stripped = line.strip()
            # Skip lines that are source references (comprehensive patterns)
            skip_line = False
            
            # Pattern 1: Lines starting with * containing Source/Chunk/Based on
            if (line_stripped.startswith('*') and 
                ('Chunk' in line or 'Source' in line or 'Based on' in line or 'Basé sur' in line)):
                skip_line = True
            
            # Pattern 2: Lines starting with - * containing Source/Chunk/Based on  
            if (line_stripped.startswith('- *') and 
                ('Chunk' in line or 'Source' in line or 'Based on' in line or 'Basé sur' in line)):
                skip_line = True
                
            # Pattern 3: Lines that are just chunk references like "- *Chunk 1 - filename*"
            if (line_stripped.startswith('- *Chunk') and line_stripped.endswith('*')):
                skip_line = True
                
            # Pattern 4: Lines that start with "- *Based on" 
            if line_stripped.startswith('- *Based on'):
                skip_line = True
                
            # Add line only if it's not a source reference and not empty
            if not skip_line and line_stripped:
                cleaned_lines.append(line)
        
        # Final clean response for mobile interface
        final_response = '\n'.join(cleaned_lines).strip()
        
        # Additional cleanup - remove any remaining source markers at start
        while final_response.startswith('- *') or final_response.startswith('*'):
            # Find the end of the line to remove
            if '\n' in final_response:
                final_response = final_response.split('\n', 1)[1].strip()
            else:
                final_response = ""
                break
        
        # Ensure we have a response
        if not final_response:
            final_response = "I apologize, but I couldn't generate a proper response to your question."
        
        await manager.send_log(f"Clean response ready for mobile interface", "success")
        
        return ChatResponse(
            response=final_response,
            sources=sources,  # Minimal info for API, detailed info already sent to debug console
            tool_used=tool_used
        )
        
    except Exception as e:
        await manager.send_log(f"Error processing question: {str(e)}", "error")
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/health")
async def health_check():
    """Health check endpoint"""
    return {
        "status": "healthy",
        "initialized": initialization_completed,
        "timestamp": datetime.now().isoformat()
    }

# Mount static files
app.mount("/static", StaticFiles(directory="static"), name="static")

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)