api/main.py

import os
import tempfile
import shutil
import json
import uuid
import time
import logging
import sys
import traceback
from typing import List, Dict, Optional
from fastapi import FastAPI, UploadFile, File, Form, HTTPException, Cookie, Response, Request
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
from fastapi.responses import FileResponse, StreamingResponse
from pydantic import BaseModel
from openai import OpenAI
from langsmith.wrappers import wrap_openai

# Import and setup logging
from aimakerspace.vectordatabase import VectorDatabase
from api.logging_config import setup_logging
logger = setup_logging(level=logging.INFO)

from aimakerspace.text_utils import CharacterTextSplitter, TextFileLoader, PDFLoader
from langchain_core.prompts import SystemMessagePromptTemplate, HumanMessagePromptTemplate, ChatPromptTemplate
from aimakerspace.qdrant_vectordb import QdrantVectorDatabase
from langchain_openai import ChatOpenAI
#from aimakerspace.openai_utils.chatmodel import ChatOpenAI

# API Version information
API_VERSION = "0.2.0"
BUILD_DATE = "2024-06-14"  # Update this when making significant changes

from .config import QDRANT_HOST, QDRANT_PORT, QDRANT_GRPC_PORT, QDRANT_PREFER_GRPC, QDRANT_COLLECTION, QDRANT_IN_MEMORY 

app = FastAPI(
    title="Quick Understand API",
    description="RAG-based question answering API for document understanding",
    version=API_VERSION
)

# Add CORS middleware
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

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

# Initialize text splitter
text_splitter = CharacterTextSplitter()

# Dictionary to store user sessions
user_sessions = {}

# Dictionary to store user-specific prompts
user_prompts = {}

# Import default prompt templates from prompts.py
from .utils.prompts import DEFAULT_SYSTEM_TEMPLATE, DEFAULT_USER_TEMPLATE

from api.models.pydantic_models import (
    PromptTemplate,
    QueryRequest,
    QueryResponse,
    DocumentSummaryRequest,
    DocumentSummaryResponse,
    QuizQuestion,
    GenerateQuizRequest,
    GenerateQuizResponse
)

# Helper function to get or create a user ID
def get_or_create_user_id(request: Request, response: Response) -> str:
    # Try to get user ID from header first
    user_id = request.headers.get("X-User-ID")
    
    # Then try to get from query parameter
    if not user_id:
        user_id = request.query_params.get("user_id")
    
    # If no user ID exists, create a new one
    if not user_id:
        user_id = str(uuid.uuid4())
        # Initialize with default prompts
        if user_id not in user_prompts:
            user_prompts[user_id] = {
                "system_template": DEFAULT_SYSTEM_TEMPLATE,
                "user_template": DEFAULT_USER_TEMPLATE
            }
    
    return user_id

# Get prompts for a specific user
def get_user_prompts(user_id: str) -> Dict[str, str]:
    if user_id not in user_prompts:
        # Initialize with default prompts if not exists
        user_prompts[user_id] = {
            "system_template": DEFAULT_SYSTEM_TEMPLATE,
            "user_template": DEFAULT_USER_TEMPLATE
        }
    
    return user_prompts[user_id]

# Helper function to extend OpenAI client with needed methods
async def acreate_single_response(client, prompt):
    response = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[{"role": "user", "content": prompt}],
        temperature=0.7,
    )
    return response.choices[0].message.content

# Helper function to provide streaming capability for OpenAI client
async def astream_openai(client, messages):
    # Convert LangChain message format to OpenAI format
    openai_messages = []
    for message in messages:
        role = "user"
        if hasattr(message, "type"):
            if message.type == "system":
                role = "system"
            elif message.type == "human":
                role = "user"
            elif message.type == "ai":
                role = "assistant"
        
        openai_messages.append({
            "role": role,
            "content": message.content
        })
    
    response = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=openai_messages,
        temperature=0.7,
        stream=True,
    )
    
    for chunk in response:
        if chunk.choices and chunk.choices[0].delta and chunk.choices[0].delta.content:
            yield chunk.choices[0].delta.content

@app.post("/upload")
async def upload_file(
    file: UploadFile = File(...), 
    session_id: str = Form(...),
    request: Request = None,
    response: Response = None
):
    request_id = str(uuid.uuid4())[:8]
    logger.info(f"[Request:{request_id}] Upload request received - session_id={session_id}, file={file.filename}")
    
    if file.content_type not in ["text/plain", "application/pdf"]:
        logger.warning(f"[Request:{request_id}] Unsupported file type: {file.content_type}")
        raise HTTPException(status_code=400, detail="Only text and PDF files are supported")
    
    # Get or create user ID
    user_id = get_or_create_user_id(request, response) if request and response else None
    if user_id:
        logger.info(f"[Request:{request_id}] User ID: {user_id}")
    
    # Track overall processing time
    upload_start_time = time.time()
    
    # Create a temporary file
    suffix = f".{file.filename.split('.')[-1]}"
    with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as temp_file:
        # Copy the uploaded file content to the temporary file
        logger.info(f"[Request:{request_id}] Reading file content")
        file_content = await file.read()
        file_size = len(file_content)
        temp_file.write(file_content)
        temp_file.flush()
        logger.info(f"[Request:{request_id}] File saved to temp location, size: {file_size} bytes")
        
        # Create appropriate loader
        if file.filename.lower().endswith('.pdf'):
            logger.info(f"[Request:{request_id}] Using PDF loader")
            loader = PDFLoader(temp_file.name)
        else:
            logger.info(f"[Request:{request_id}] Using text loader")
            loader = TextFileLoader(temp_file.name)
        
        try:
            # Load and process the documents
            logger.info(f"[Request:{request_id}] Loading documents")
            doc_load_start = time.time()
            documents = loader.load_documents()
            doc_load_time = time.time() - doc_load_start
            logger.info(f"[Request:{request_id}] Documents loaded in {doc_load_time:.4f} seconds, count: {len(documents)}")
            
            # Split documents into chunks
            logger.info(f"[Request:{request_id}] Splitting documents into chunks")
            split_start = time.time()
            texts = text_splitter.split_texts(documents)
            split_time = time.time() - split_start
            logger.info(f"[Request:{request_id}] Document splitting completed in {split_time:.4f} seconds, chunk count: {len(texts)}")
            
            # Log information about chunk lengths
            if texts:
                chunk_lengths = [len(t) for t in texts]
                logger.info(f"[Request:{request_id}] Chunk statistics: min={min(chunk_lengths)}, max={max(chunk_lengths)}, avg={sum(chunk_lengths)/len(chunk_lengths):.2f} chars")
            
            # Create vector database
            logger.info(f"[Request:{request_id}] Creating vector database: {QDRANT_COLLECTION}_{session_id}")
            vector_start = time.time()
            vector_db = VectorDatabase()
            
            # Build the vector database
            logger.info(f"[Request:{request_id}] Building vector database with {len(texts)} chunks")
            vector_db = await vector_db.abuild_from_list(texts)
            vector_time = time.time() - vector_start
            logger.info(f"[Request:{request_id}] Vector database creation completed in {vector_time:.4f} seconds")
            
            # Create chat model
            logger.info(f"[Request:{request_id}] Creating chat model")
            

            openai_client = wrap_openai(OpenAI())
            
            # Get user prompts
            user_prompt_templates = get_user_prompts(user_id) if user_id else {
                "system_template": DEFAULT_SYSTEM_TEMPLATE,
                "user_template": DEFAULT_USER_TEMPLATE
            }
            
            # Create the retrieval pipeline with user-specific prompts
            pipeline_start = time.time()
            logger.info(f"[Request:{request_id}] Creating retrieval pipeline")
            retrieval_pipeline = RetrievalAugmentedQAPipeline(
                vector_db_retriever=vector_db,
                llm=openai_client,
                system_template=user_prompt_templates["system_template"],
                user_template=user_prompt_templates["user_template"]
            )
            pipeline_time = time.time() - pipeline_start
            logger.info(f"[Request:{request_id}] Retrieval pipeline created in {pipeline_time:.4f} seconds")
            
            # Store the retrieval pipeline in the user session
            user_sessions[session_id] = retrieval_pipeline
            logger.info(f"[Request:{request_id}] Retrieval pipeline stored in session {session_id}")
            
            # Generate document description and suggested questions
            logger.info(f"[Request:{request_id}] Generating document description and questions")
            summary_start = time.time()
            doc_content = "\n".join(texts[:5])  # Use first few chunks for summary
            
            description_prompt = f"""
            Please provide a brief description of this document in 2-3 sentences:
            {doc_content}
            """
            
            questions_prompt = f"""
            Based on this document content, please suggest 3 specific questions that would be informative to ask:
            {doc_content}
            
            Format your response as a JSON array with 3 question strings.
            """
            
            # Get document description
            logger.info(f"[Request:{request_id}] Generating document description")
            description_response = await acreate_single_response(openai_client, description_prompt)
            document_description = description_response.strip()
            
            # Get suggested questions
            logger.info(f"[Request:{request_id}] Generating suggested questions")
            questions_response = await acreate_single_response(openai_client, questions_prompt)
            
            # Try to parse the questions as JSON, or extract them as best as possible
            try:
                import json
                suggested_questions = json.loads(questions_response)
                logger.info(f"[Request:{request_id}] Successfully parsed suggested questions as JSON")
            except:
                # Extract questions with a fallback method
                logger.info(f"[Request:{request_id}] Parsing JSON failed, using fallback method")
                import re
                questions = re.findall(r'["\']([^"\']+)["\']', questions_response)
                if not questions or len(questions) < 3:
                    questions = [q.strip() for q in questions_response.split("\n") if "?" in q]
                    logger.info(f"[Request:{request_id}] Extracted questions using line splitting: {len(questions)} found")
                if not questions or len(questions) < 3:
                    logger.info(f"[Request:{request_id}] No questions found, using default questions")
                    questions = ["What is the main topic of this document?", 
                                "What are the key points discussed in the document?", 
                                "How can I apply the information in this document?"]
                suggested_questions = questions[:3]
            
            summary_time = time.time() - summary_start
            logger.info(f"[Request:{request_id}] Document summary generation completed in {summary_time:.4f} seconds")
            
            total_time = time.time() - upload_start_time
            logger.info(f"[Request:{request_id}] Total processing time: {total_time:.4f} seconds")
            
            result = {
                "status": "success", 
                "message": f"Processed {file.filename}", 
                "session_id": session_id,
                "document_description": document_description,
                "suggested_questions": suggested_questions,
                "processing_stats": {
                    "total_time": total_time,
                    "doc_load_time": doc_load_time,
                    "split_time": split_time,
                    "vector_time": vector_time,
                    "chunk_count": len(texts)
                }
            }
            
            # Add user_id to result if available
            if user_id:
                result["user_id"] = user_id
            
            logger.info(f"[Request:{request_id}] Upload processing completed successfully")    
            return result
            
        except Exception as e:
            error_time = time.time() - upload_start_time
            exc_type, exc_obj, exc_tb = sys.exc_info()
            fname = os.path.split(exc_tb.tb_frame.f_code.co_filename)[1]
            error_location = f"{fname}:{exc_tb.tb_lineno}"
            error_traceback = "".join(traceback.format_tb(exc_tb))
            logger.error(f"[Request:{request_id}] Error processing upload after {error_time:.4f} seconds at {error_location}: {str(e)}")
            logger.error(f"[Request:{request_id}] Traceback: {error_traceback}")
            raise HTTPException(status_code=500, detail=f"Error processing file: {str(e)} at {error_location}")
        finally:
            # Clean up the temporary file
            try:
                os.unlink(temp_file.name)
                logger.info(f"[Request:{request_id}] Temp file cleaned up")
            except Exception as e:
                logger.error(f"[Request:{request_id}] Error cleaning up temporary file: {e}")

@app.post("/query", response_model=QueryResponse)
async def query(request: QueryRequest):
    session_id = request.session_id
    user_id = request.user_id
    request_id = str(uuid.uuid4())[:8]
    
    logger.info(f"[Request:{request_id}] Query request received - session_id={session_id}, user_id={user_id}, query='{request.query}'")
    
    # Check if session exists
    if session_id not in user_sessions:
        logger.warning(f"[Request:{request_id}] Session not found: {session_id}")
        raise HTTPException(status_code=404, detail="Session not found. Please upload a document first.")
    
    # Get the retrieval pipeline from the session
    retrieval_pipeline = user_sessions[session_id]
    logger.info(f"[Request:{request_id}] Retrieved pipeline for session {session_id}")
    
    # Update prompts if user_id is provided and different from current
    if user_id and retrieval_pipeline.system_template != get_user_prompts(user_id)["system_template"]:
        logger.info(f"[Request:{request_id}] Updating prompt templates for user {user_id}")
        user_prompt_templates = get_user_prompts(user_id)
        retrieval_pipeline.update_templates(
            user_prompt_templates["system_template"],
            user_prompt_templates["user_template"]
        )
    
    # Run the query
    start_time = time.time()
    logger.info(f"[Request:{request_id}] Executing RAG pipeline")
    result = await retrieval_pipeline.arun_pipeline(request.query, user_id, session_id)
    
    # Process the result and return the response
    response_text = ""
    token_count = 0
    async for chunk in result["response"]:
        response_text += chunk
        token_count += 1
    
    process_time = time.time() - start_time
    
    # Log detailed information about the response
    logger.info(f"[Request:{request_id}] Request processed in {process_time:.4f} seconds, response length: {len(response_text)} chars, {token_count} tokens")
    
    # Extract and log metrics from result
    if "search_time" in result:
        logger.info(f"[Request:{request_id}] Vector search time: {result['search_time']:.4f} seconds")
    if "context_length" in result:
        logger.info(f"[Request:{request_id}] Context length: {result['context_length']} characters")
    
    # Log context scores information
    context_list = result.get("context", [])
    if context_list:
        scores = [score for _, score in context_list]
        logger.info(f"[Request:{request_id}] Context similarity scores: min={min(scores):.4f}, max={max(scores):.4f}, avg={sum(scores)/len(scores):.4f}")
    
    return {"response": response_text, "session_id": session_id}

@app.post("/stream")
async def stream_query(request: QueryRequest):
    session_id = request.session_id
    user_id = request.user_id
    request_id = str(uuid.uuid4())[:8]
    
    logger.info(f"[Request:{request_id}] Stream query request received - session_id={session_id}, user_id={user_id}, query='{request.query}'")
    
    # Check if session exists
    if session_id not in user_sessions:
        logger.warning(f"[Request:{request_id}] Session not found: {session_id}")
        raise HTTPException(status_code=404, detail="Session not found. Please upload a document first.")
    
    # Get the retrieval pipeline from the session
    retrieval_pipeline = user_sessions[session_id]
    logger.info(f"[Request:{request_id}] Retrieved pipeline for session {session_id}")
    
    # Update prompts if user_id is provided and different from current
    if user_id and retrieval_pipeline.system_template != get_user_prompts(user_id)["system_template"]:
        logger.info(f"[Request:{request_id}] Updating prompt templates for user {user_id}")
        user_prompt_templates = get_user_prompts(user_id)
        retrieval_pipeline.update_templates(
            user_prompt_templates["system_template"],
            user_prompt_templates["user_template"]
        )
    
    # Run the query
    start_time = time.time()
    logger.info(f"[Request:{request_id}] Executing RAG pipeline for streaming")
    result = await retrieval_pipeline.arun_pipeline(request.query, user_id, session_id)
    
    # Extract context for logging
    context_list = result.get("context", [])
    scores = [score for _, score in context_list] if context_list else []
    if scores:
        logger.info(f"[Request:{request_id}] Context similarity scores: min={min(scores):.4f}, max={max(scores):.4f}, avg={sum(scores)/len(scores):.4f}")
    
    async def generate():
        token_count = 0
        chunk_count = 0
        response_buffer = ""
        async for chunk in result["response"]:
            token_count += 1
            chunk_count += 1
            response_buffer += chunk
            
            # Collect 5 tokens before sending or at the end of the stream
            if token_count % 5 == 0 or chunk == "":
                yield f"data: {json.dumps({'text': response_buffer})}\n\n"
                response_buffer = ""
        
        # Send any remaining text
        if response_buffer:
            yield f"data: {json.dumps({'text': response_buffer})}\n\n"
        
        # Send end of stream marker
        completion_time = time.time() - start_time
        logger.info(f"[Request:{request_id}] Streaming completed in {completion_time:.4f} seconds, sent {token_count} tokens in {chunk_count} chunks")
        yield f"data: [DONE]\n\n"
    
    return StreamingResponse(generate(), media_type="text/event-stream")

@app.get("/stream")
async def stream_query_get(
    session_id: str, 
    query: str, 
    user_id: Optional[str] = None,
    request: Request = None,
    response: Response = None
):
    request_id = str(uuid.uuid4())[:8]
    
    # Get or create user ID if not provided
    if request and response and not user_id:
        user_id = get_or_create_user_id(request, response)
    
    logger.info(f"[Request:{request_id}] Stream GET query received - session_id={session_id}, user_id={user_id}, query='{query}'")
    
    # Check if session exists
    if session_id not in user_sessions:
        logger.warning(f"[Request:{request_id}] Session not found: {session_id}")
        raise HTTPException(status_code=404, detail="Session not found. Please upload a document first.")
    
    # Get the retrieval pipeline from the session
    retrieval_pipeline = user_sessions[session_id]
    logger.info(f"[Request:{request_id}] Retrieved pipeline for session {session_id}")
    
    # Update prompts if user_id is provided and different from current
    if user_id and retrieval_pipeline.system_template != get_user_prompts(user_id)["system_template"]:
        logger.info(f"[Request:{request_id}] Updating prompt templates for user {user_id}")
        user_prompt_templates = get_user_prompts(user_id)
        retrieval_pipeline.update_templates(
            user_prompt_templates["system_template"],
            user_prompt_templates["user_template"]
        )
    
    # Run the query
    start_time = time.time()
    logger.info(f"[Request:{request_id}] Executing RAG pipeline for streaming (GET)")
    result = await retrieval_pipeline.arun_pipeline(query, user_id, session_id)
    
    # Extract context for logging
    context_list = result.get("context", [])
    scores = [score for _, score in context_list] if context_list else []
    if scores:
        logger.info(f"[Request:{request_id}] Context similarity scores: min={min(scores):.4f}, max={max(scores):.4f}, avg={sum(scores)/len(scores):.4f}")
    
    async def generate():
        token_count = 0
        chunk_count = 0
        response_buffer = ""
        async for chunk in result["response"]:
            token_count += 1
            chunk_count += 1
            response_buffer += chunk
            
            # Collect 5 tokens before sending or at the end of the stream
            if token_count % 5 == 0 or chunk == "":
                yield f"data: {json.dumps({'text': response_buffer})}\n\n"
                response_buffer = ""
        
        # Send any remaining text
        if response_buffer:
            yield f"data: {json.dumps({'text': response_buffer})}\n\n"
        
        # Send end of stream marker
        completion_time = time.time() - start_time
        logger.info(f"[Request:{request_id}] Streaming completed in {completion_time:.4f} seconds, sent {token_count} tokens in {chunk_count} chunks")
        yield f"data: [DONE]\n\n"
    
    return StreamingResponse(generate(), media_type="text/event-stream")

@app.post("/document-summary", response_model=DocumentSummaryResponse)
async def get_document_summary(request: DocumentSummaryRequest):
    session_id = request.session_id
    user_id = request.user_id
    
    # Check if session exists
    if session_id not in user_sessions:
        raise HTTPException(status_code=404, detail="Session not found. Please upload a document first.")
    
    # Get the retrieval pipeline from the session
    retrieval_pipeline = user_sessions[session_id]
    
    # Update prompts if user_id is provided and different from current
    if user_id and retrieval_pipeline.system_template != get_user_prompts(user_id)["system_template"]:
        user_prompt_templates = get_user_prompts(user_id)
        retrieval_pipeline.update_templates(
            user_prompt_templates["system_template"],
            user_prompt_templates["user_template"]
        )
    
    # Get access to the document content
    vector_db = retrieval_pipeline.vector_db_retriever
    
    # We'll use all the text chunks to create a comprehensive summary
    # Get all text chunks from the vector store
    all_texts = vector_db.get_all_texts()
    
    # Combine a sample of the texts (to avoid hitting token limits)
    sample_texts = all_texts[:10] if len(all_texts) > 10 else all_texts
    doc_content = "\n".join(sample_texts)
    
    # Create the LLM summary prompt
    summary_prompt = f"""
    Analyze the following document content and generate a structured summary in JSON format:
    
    ```
    {doc_content}
    ```
    
    Return ONLY a JSON object with the following structure:
    
    {{
      "keyTopics": [list of 5-7 key topics in the document],
      "entities": [list of 5-8 important named entities such as organizations, technologies, or people],
      "wordCloudData": [
        {{ "text": "word1", "value": frequency_score }},
        {{ "text": "word2", "value": frequency_score }},
        ...
      ],
      "documentStructure": [
        {{ 
          "title": "Section title",
          "subsections": ["Subsection1", "Subsection2", ...] 
        }},
        ...
      ]
    }}
    
    The wordCloudData should contain 15-20 important terms with their relative frequency scores (higher numbers = more important/frequent).
    The documentStructure should reflect the hierarchical organization of the document with main sections and their subsections.
    """
    
    # Get LLM response
    try:
        llm = retrieval_pipeline.llm
        response = await acreate_single_response(llm, summary_prompt)
        
        # Parse the JSON
        # Find JSON content (sometimes the LLM adds extra text)
        import re
        json_match = re.search(r'({[\s\S]*})', response)
        
        if json_match:
            json_str = json_match.group(1)
            summary_data = json.loads(json_str)
        else:
            # If no JSON found, create a basic structure with an error message
            summary_data = {
                "keyTopics": ["Error parsing document structure"],
                "entities": ["Please try again"],
                "wordCloudData": [{"text": "Error", "value": 50}],
                "documentStructure": [{"title": "Document structure unavailable", "subsections": []}]
            }
            
        # Ensure the response has all required fields
        if "keyTopics" not in summary_data:
            summary_data["keyTopics"] = ["Topic extraction failed"]
        if "entities" not in summary_data:
            summary_data["entities"] = ["Entity extraction failed"]
        if "wordCloudData" not in summary_data:
            summary_data["wordCloudData"] = [{"text": "Data", "value": 50}]
        if "documentStructure" not in summary_data:
            summary_data["documentStructure"] = [{"title": "Structure unavailable", "subsections": []}]
            
        return summary_data
        
    except Exception as e:
        # Return a fallback summary on error
        return {
            "keyTopics": ["Error analyzing document"],
            "entities": ["Try refreshing the page"],
            "wordCloudData": [
                {"text": "Error", "value": 60},
                {"text": "Document", "value": 40},
                {"text": "Analysis", "value": 30}
            ],
            "documentStructure": [
                {"title": "Error in document analysis", "subsections": ["Please try again"]}
            ]
        }

@app.post("/generate-quiz", response_model=GenerateQuizResponse)
async def generate_quiz(request: GenerateQuizRequest):
    session_id = request.session_id
    num_questions = min(request.num_questions, 10)  # Limit to max 10 questions
    user_id = request.user_id
    
    # Check if session exists
    if session_id not in user_sessions:
        raise HTTPException(status_code=404, detail="Session not found. Please upload a document first.")
    
    # Get the retrieval pipeline from the session
    retrieval_pipeline = user_sessions[session_id]
    
    # Update prompts if user_id is provided and different from current
    if user_id and retrieval_pipeline.system_template != get_user_prompts(user_id)["system_template"]:
        user_prompt_templates = get_user_prompts(user_id)
        retrieval_pipeline.update_templates(
            user_prompt_templates["system_template"],
            user_prompt_templates["user_template"]
        )
    
    # Get access to the document content
    vector_db = retrieval_pipeline.vector_db_retriever
    
    # We'll use all the text chunks to create comprehensive quiz questions
    # Get all text chunks from the vector store
    all_texts = vector_db.get_all_texts()
    
    # Combine a sample of the texts (to avoid hitting token limits)
    sample_texts = all_texts[:15] if len(all_texts) > 15 else all_texts
    doc_content = "\n".join(sample_texts)
    
    # Create the LLM quiz generation prompt
    quiz_prompt = f"""
    Based on the following document content, generate {num_questions} multiple-choice quiz questions to test the reader's understanding:
    
    ```
    {doc_content}
    ```
    
    For each question:
    1. Create a clear, specific question about key information in the document
    2. Provide exactly 4 answer options (A, B, C, D)
    3. Clearly indicate which option is correct
    4. Make sure distractors (wrong answers) are plausible but clearly incorrect
    
    Return ONLY a JSON object with the following structure:
    
    {{
      "questions": [
        {{
          "id": "unique_id",
          "text": "question text",
          "options": ["option A", "option B", "option C", "option D"],
          "correctAnswer": "correct option text"
        }},
        ...
      ]
    }}
    
    The questions should cover different aspects of the document and test genuine understanding.
    """
    
    # Get LLM response
    try:
        llm = retrieval_pipeline.llm
        response = await acreate_single_response(llm, quiz_prompt)
        
        # Parse the JSON
        # Find JSON content (sometimes the LLM adds extra text)
        import re
        json_match = re.search(r'({[\s\S]*})', response)
        
        if json_match:
            json_str = json_match.group(1)
            quiz_data = json.loads(json_str)
            
            # Validate and clean the questions
            questions = []
            for q in quiz_data.get("questions", []):
                # Ensure each question has a unique ID
                if "id" not in q or not q["id"]:
                    q["id"] = str(uuid.uuid4())
                
                # Verify the question has all required fields
                if "text" in q and "options" in q and "correctAnswer" in q:
                    # Verify correctAnswer is in options
                    if q["correctAnswer"] in q["options"]:
                        questions.append(QuizQuestion(
                            id=q["id"],
                            text=q["text"],
                            options=q["options"],
                            correctAnswer=q["correctAnswer"]
                        ))
            
            # If no valid questions were found or not enough questions, create fallback
            if len(questions) < min(3, num_questions):
                questions = generate_fallback_questions(num_questions)
                
            return {"questions": questions[:num_questions]}
        else:
            # If no JSON found, return fallback questions
            return {"questions": generate_fallback_questions(num_questions)}
            
    except Exception as e:
        print(f"Error generating quiz: {e}")
        # Return fallback questions on error
        return {"questions": generate_fallback_questions(num_questions)}

def generate_fallback_questions(num_questions: int) -> List[QuizQuestion]:
    """Generate generic fallback questions when LLM fails"""
    fallback_questions = [
        QuizQuestion(
            id=str(uuid.uuid4()),
            text="What is the main purpose of a RAG (Retrieval-Augmented Generation) system?",
            options=[
                "To generate random text without meaning",
                "To retrieve documents from a database only",
                "To combine document retrieval with language model generation",
                "To replace human writing entirely"
            ],
            correctAnswer="To combine document retrieval with language model generation"
        ),
        QuizQuestion(
            id=str(uuid.uuid4()),
            text="Which component is NOT typically part of a RAG system?",
            options=[
                "Vector database",
                "Language model",
                "Blockchain ledger",
                "Text splitter"
            ],
            correctAnswer="Blockchain ledger"
        ),
        QuizQuestion(
            id=str(uuid.uuid4()),
            text="What is the benefit of using RAG over a standalone language model?",
            options=[
                "It's always faster",
                "It provides more up-to-date and accurate information",
                "It uses less computational resources",
                "It requires no training data"
            ],
            correctAnswer="It provides more up-to-date and accurate information"
        ),
        QuizQuestion(
            id=str(uuid.uuid4()),
            text="What is a vector embedding in the context of RAG?",
            options=[
                "A mathematical representation of text in multidimensional space",
                "A form of data compression",
                "A type of encryption",
                "A physical server component"
            ],
            correctAnswer="A mathematical representation of text in multidimensional space"
        ),
        QuizQuestion(
            id=str(uuid.uuid4()),
            text="How does a RAG system determine which text chunks are relevant to a query?",
            options=[
                "Random selection",
                "Semantic similarity between query and text embeddings",
                "Alphabetical ordering",
                "Document recency only"
            ],
            correctAnswer="Semantic similarity between query and text embeddings"
        )
    ]
    return fallback_questions[:num_questions]

# New endpoint to get user prompts
@app.get("/prompts")
async def get_prompts(
    request: Request,
    response: Response,
    user_id: Optional[str] = None
):
    # Get or create user ID if not provided
    if not user_id:
        user_id = get_or_create_user_id(request, response)
    
    # Get user prompts
    prompts = get_user_prompts(user_id)
    
    return {
        "user_id": user_id,
        "system_template": prompts["system_template"],
        "user_template": prompts["user_template"]
    }

# New endpoint to update user prompts
@app.post("/prompts")
async def update_prompts(
    prompt_template: PromptTemplate,
    request: Request,
    response: Response,
    user_id: Optional[str] = None
):
    # Get or create user ID if not provided
    if not user_id:
        user_id = get_or_create_user_id(request, response)
    
    # Update prompts
    user_prompts[user_id] = {
        "system_template": prompt_template.system_template,
        "user_template": prompt_template.user_template
    }
    
    return {
        "status": "success",
        "user_id": user_id,
        "message": "Prompts updated successfully"
    }

# Reset user prompts to default
@app.post("/prompts/reset")
async def reset_prompts(
    request: Request,
    response: Response,
    user_id: Optional[str] = None
):
    # Get or create user ID if not provided
    if not user_id:
        user_id = get_or_create_user_id(request, response)
    
    # Reset to defaults
    user_prompts[user_id] = {
        "system_template": DEFAULT_SYSTEM_TEMPLATE,
        "user_template": DEFAULT_USER_TEMPLATE
    }
    
    return {
        "status": "success",
        "user_id": user_id,
        "message": "Prompts reset to default successfully"
    }

@app.get("/identify")
async def identify_user(request: Request, response: Response):
    user_id = get_or_create_user_id(request, response)
    return {"user_id": user_id}

# Serve the frontend
@app.get("/")
async def read_root():
    return FileResponse("static/index.html")

@app.get("/version")
async def get_version():
    return {
        "api_version": API_VERSION,
        "build_date": BUILD_DATE,
        "status": "operational"
    }

@app.get("/{path:path}")
async def catch_all(path: str):
    if os.path.exists(f"static/{path}"):
        return FileResponse(f"static/{path}")
    return FileResponse("static/index.html")

class RetrievalAugmentedQAPipeline:
    def __init__(self, llm: ChatOpenAI, vector_db_retriever: QdrantVectorDatabase, 
                system_template: str = DEFAULT_SYSTEM_TEMPLATE, 
                user_template: str = DEFAULT_USER_TEMPLATE) -> None:
        self.llm = llm
        self.vector_db_retriever = vector_db_retriever
        self.system_template = system_template
        self.user_template = user_template
        self.system_prompt_template = SystemMessagePromptTemplate.from_template(system_template)
        self.human_prompt_template = HumanMessagePromptTemplate.from_template(user_template)
        self.chat_prompt_template = ChatPromptTemplate.from_messages([
            self.system_prompt_template,
            self.human_prompt_template
        ])
        
        # Import LangSmith utilities
        try:
            from api.utils.langsmith_utils import langsmith_tracer
            self.langsmith_tracer = langsmith_tracer
            logger.info("LangSmith tracer initialized in RAG pipeline")
        except ImportError:
            logger.warning("LangSmith utils not available, tracing disabled")
            self.langsmith_tracer = None

    def update_templates(self, system_template: str, user_template: str):
        """Update prompt templates"""
        self.system_template = system_template
        self.user_template = user_template
        self.system_prompt_template = SystemMessagePromptTemplate.from_template(system_template)
        self.human_prompt_template = HumanMessagePromptTemplate.from_template(user_template)
        self.chat_prompt_template = ChatPromptTemplate.from_messages([
            self.system_prompt_template,
            self.human_prompt_template
        ])

    async def arun_pipeline(self, user_query: str, user_id: str = None, session_id: str = None):
        # Get context from vector database
        context_list = self.vector_db_retriever.search_by_text(user_query, k=4)
        
        # Log context retrieval to LangSmith if available
        retrieval_run_id = None
        if self.langsmith_tracer and self.langsmith_tracer.tracing_enabled and self.langsmith_tracer.client:
            # Add debug logging
            logger.info(f"Attempting to log retrieval to LangSmith. Tracer enabled: {self.langsmith_tracer.tracing_enabled}")
            try:
                retrieval_run_id = self.langsmith_tracer.log_retrieval(
                    query=user_query,
                    retrieved_documents=context_list,
                    user_id=user_id,
                    session_id=session_id
                )
                logger.info(f"Successfully logged retrieval to LangSmith with run_id: {retrieval_run_id}")
            except Exception as e:
                logger.error(f"Failed to log retrieval to LangSmith: {str(e)}")

        # Format context for prompt
        context_prompt = ""
        for context in context_list:
            context_prompt += context[0] + "\n"

        # Create messages using LangChain prompt templates
        messages = self.chat_prompt_template.format_messages(
            question=user_query, 
            context=context_prompt
        )

        async def generate_response():
            response_chunks = []
            # Use our custom streaming function
            async for chunk in astream_openai(self.llm, messages):
                response_chunks.append(chunk)
                yield chunk
            
            # Log generation to LangSmith if available
            if self.langsmith_tracer and self.langsmith_tracer.tracing_enabled and self.langsmith_tracer.client:
                try:
                    full_response = "".join(response_chunks)
                    self.langsmith_tracer.log_rag_generation(
                        query=user_query,
                        context=context_prompt,
                        response=full_response,
                        system_prompt=self.system_template,
                        user_prompt=self.user_template,
                        user_id=user_id,
                        session_id=session_id,
                        parent_run_id=retrieval_run_id
                    )
                    logger.info("Successfully logged generation to LangSmith")
                except Exception as e:
                    logger.error(f"Failed to log generation to LangSmith: {str(e)}")

        return {"response": generate_response(), "context": context_list}

if __name__ == "__main__":
    import uvicorn
    uvicorn.run("main:app", host="0.0.0.0", port=8000, reload=True)