document_rag_router.py

from fastapi import UploadFile, File, Form, HTTPException, APIRouter
from typing import List, Optional, Dict, Tuple
import lancedb
from lancedb.pydantic import LanceModel, Vector
from lancedb.embeddings import get_registry
import pandas as pd
from utils import process_pdf_to_chunks
import hashlib
import uuid
import json
from datetime import datetime
from pydantic import BaseModel
import logging

# Create router
router = APIRouter(
    prefix="/rag",
    tags=["rag"]
)

# Initialize LanceDB and embedding model
db = lancedb.connect("/tmp/db")
model = get_registry().get("sentence-transformers").create(
    name="Snowflake/snowflake-arctic-embed-xs", 
    device="cpu"
)

def get_user_collection(user_id: str, collection_name: str) -> str:
    """Generate user-specific collection name"""
    return f"{user_id}_{collection_name}"

class DocumentChunk(LanceModel):
    text: str = model.SourceField()
    vector: Vector(model.ndims()) = model.VectorField()
    document_id: str
    chunk_index: int
    file_name: str
    file_type: str
    created_date: str
    collection_id: str
    user_id: str
    metadata_json: str
    char_start: int
    char_end: int
    page_numbers: List[int]
    images: List[str]

class QueryInput(BaseModel):
    collection_id: str
    query: str
    top_k: Optional[int] = 3
    user_id: str

class SearchResult(BaseModel):
    text: str
    distance: float
    metadata: Dict  # Added metadata field

class SearchResponse(BaseModel):
    results: List[SearchResult]

async def process_file(file: UploadFile, collection_id: str, user_id: str) -> Tuple[List[dict], str]:
    """Process single file and return chunks with metadata"""
    content = await file.read()
    file_type = file.filename.split('.')[-1].lower()
    
    chunks = []
    doc_id = ""
    if file_type == 'pdf':
        chunks, doc_id = process_pdf_to_chunks(
            pdf_content=content,
            file_name=file.filename
        )
    elif file_type == 'txt':
        doc_id = hashlib.sha256(content).hexdigest()[:4]
        text_content = content.decode('utf-8')
        chunks = [{
            "text": text_content,
            "metadata": {
                "created_date": datetime.now().isoformat(),
                "file_name": file.filename,
                "document_id": doc_id,
                "user_id": user_id,
                "location": {
                    "chunk_index": 0,
                    "char_start": 0,
                    "char_end": len(text_content),
                    "pages": [1],
                    "total_chunks": 1
                },
                "images": []
            }
        }]
    
    return chunks, doc_id

@router.post("/upload_files")
async def upload_files(
    files: List[UploadFile] = File(...),
    collection_name: Optional[str] = Form(None),
    user_id: str = Form(...)
):
    try:
        collection_id = get_user_collection(
            user_id, 
            collection_name if collection_name else f"col_{uuid.uuid4().hex[:8]}"
        )
        all_chunks = []
        doc_ids = {}
        for file in files:
            try:
                chunks, doc_id = await process_file(file, collection_id, user_id)
                for chunk in chunks:
                    chunk_data = {
                        "text": chunk["text"],
                        "document_id": chunk["metadata"]["document_id"],
                        "chunk_index": chunk["metadata"]["location"]["chunk_index"],
                        "file_name": chunk["metadata"]["file_name"],
                        "file_type": file.filename.split('.')[-1].lower(),
                        "created_date": chunk["metadata"]["created_date"],
                        "collection_id": collection_id,
                        "user_id": user_id,
                        "metadata_json": json.dumps(chunk["metadata"]),
                        "char_start": chunk["metadata"]["location"]["char_start"],
                        "char_end": chunk["metadata"]["location"]["char_end"],
                        "page_numbers": chunk["metadata"]["location"]["pages"],
                        "images": chunk["metadata"].get("images", [])
                    }
                    all_chunks.append(chunk_data)
                doc_ids[doc_id] = file.filename
            except Exception as e:
                logging.error(f"Error processing file {file.filename}: {str(e)}")
                raise HTTPException(
                    status_code=400,
                    detail=f"Error processing file {file.filename}: {str(e)}"
                )

        try:
            table = db.open_table(collection_id)
        except Exception as e:
            logging.error(f"Error opening table: {str(e)}")
            try:
                table = db.create_table(
                    collection_id,
                    schema=DocumentChunk,
                    mode="create"
                )
                # Create FTS index on the text column for hybrid search support

                # table.create_fts_index(
                #     field_names="text",
                #     replace=True,
                #     tokenizer_name="en_stem",  # Use English stemming
                #     lower_case=True,  # Convert text to lowercase
                #     remove_stop_words=True,  # Remove common words like "the", "is", "at"
                #     writer_heap_size=1024 * 1024 * 1024  # 1GB heap size
                # )

            except Exception as e:
                logging.error(f"Error creating table: {str(e)}")
                raise HTTPException(
                    status_code=500,
                    detail=f"Error creating database table: {str(e)}"
                )
            
        try:
            df = pd.DataFrame(all_chunks)
            table.add(data=df)
        except Exception as e:
            logging.error(f"Error adding data to table: {str(e)}")
            raise HTTPException(
                status_code=500,
                detail=f"Error adding data to database: {str(e)}"
            )
        
        return {
            "message": f"Successfully processed {len(files)} files",
            "collection_id": collection_id,
            "total_chunks": len(all_chunks),
            "user_id": user_id,
            "document_ids": doc_ids
        }
        
    except HTTPException:
        raise
    except Exception as e:
        logging.error(f"Unexpected error during file upload: {str(e)}")
        raise HTTPException(
            status_code=500,
            detail=f"Unexpected error: {str(e)}"
        )

@router.get("/get_document/{collection_id}/{document_id}")
async def get_document(
    collection_id: str,
    document_id: str,
    user_id: str
):
    try:
        table = db.open_table(f"{user_id}_{collection_id}")
    except Exception as e:
        logging.error(f"Error opening table: {str(e)}")
        raise HTTPException(
            status_code=404,
            detail=f"Collection not found: {str(e)}"
        )

    try:
        chunks = table.to_pandas()
        doc_chunks = chunks[
            (chunks['document_id'] == document_id) & 
            (chunks['user_id'] == user_id)
        ].sort_values('chunk_index')
        
        if len(doc_chunks) == 0:
            raise HTTPException(
                status_code=404,
                detail=f"Document {document_id} not found in collection {collection_id}"
            )
            
        return {
            "document_id": document_id,
            "file_name": doc_chunks.iloc[0]['file_name'],
            "chunks": [
                {
                    "text": row['text'],
                    "metadata": json.loads(row['metadata_json'])
                }
                for _, row in doc_chunks.iterrows()
            ]
        }
    except HTTPException:
        raise
    except Exception as e:
        logging.error(f"Error retrieving document: {str(e)}")
        raise HTTPException(
            status_code=500,
            detail=f"Error retrieving document: {str(e)}"
        )

@router.post("/query_collection", response_model=SearchResponse)
async def query_collection(input_data: QueryInput):
    try:
        collection_id = get_user_collection(input_data.user_id, input_data.collection_id)
        
        try:
            table = db.open_table(collection_id)
        except Exception as e:
            logging.error(f"Error opening table: {str(e)}")
            raise HTTPException(
                status_code=404,
                detail=f"Collection not found: {str(e)}"
            )

        try:
            results = (
                table.search(input_data.query)
                .where(f"user_id = '{input_data.user_id}'")
                .limit(input_data.top_k)
                .to_list()
            )
        except Exception as e:
            logging.error(f"Error searching collection: {str(e)}")
            raise HTTPException(
                status_code=500,
                detail=f"Error searching collection: {str(e)}"
            )
        
        return SearchResponse(results=[
            SearchResult(
                text=r['text'],
                distance=float(r['_distance']),
                metadata=json.loads(r['metadata_json'])
            )
            for r in results
        ])
    except HTTPException:
        raise
    except Exception as e:
        logging.error(f"Unexpected error during query: {str(e)}")
        raise HTTPException(
            status_code=500,
            detail=f"Unexpected error: {str(e)}"
        )



@router.get("/list_collections")
async def list_collections(user_id: str):
    try:
        all_collections = db.table_names()
        user_collections = [
            c for c in all_collections 
            if c.startswith(f"{user_id}_")
        ]
        
        # Get documents for each collection
        collections_info = []
        for collection_name in user_collections:
            try:
                table = db.open_table(collection_name)
                df = table.to_pandas()
                
                # Group by document_id to get unique documents
                documents = df.groupby('document_id').agg({
                    'file_name': 'first',
                    'created_date': 'first'
                }).reset_index()
                
                collections_info.append({
                    "collection_id": collection_name.replace(f"{user_id}_", ""),
                    "documents": [
                        {
                            "document_id": row['document_id'],
                            "file_name": row['file_name'],
                            "created_date": row['created_date']
                        }
                        for _, row in documents.iterrows()
                    ]
                })
            except Exception as e:
                logging.error(f"Error processing collection {collection_name}: {str(e)}")
                continue
                
        return {"collections": collections_info}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@router.delete("/delete_collection/{collection_id}")
async def delete_collection(collection_id: str, user_id: str):
    try:
        full_collection_id = f"{user_id}_{collection_id}"
        
        # Check if collection exists
        try:
            table = db.open_table(full_collection_id)
        except Exception as e:
            logging.error(f"Collection not found: {str(e)}")
            raise HTTPException(
                status_code=404,
                detail=f"Collection {collection_id} not found"
            )

        # Verify ownership
        if not full_collection_id.startswith(f"{user_id}_"):
            logging.error(f"Unauthorized deletion attempt for collection {collection_id} by user {user_id}")
            raise HTTPException(
                status_code=403,
                detail="Not authorized to delete this collection"
            )

        try:
            db.drop_table(full_collection_id)
        except Exception as e:
            logging.error(f"Error deleting collection {collection_id}: {str(e)}")
            raise HTTPException(
                status_code=500,
                detail=f"Error deleting collection: {str(e)}"
            )

        return {
            "message": f"Collection {collection_id} deleted successfully",
            "collection_id": collection_id
        }

    except HTTPException:
        raise
    except Exception as e:
        logging.error(f"Unexpected error deleting collection {collection_id}: {str(e)}")
        raise HTTPException(
            status_code=500,
            detail=f"Unexpected error: {str(e)}"
        )

@router.post("/get_collection_files")
def get_collection_files(collection_id: str, user_id: str) -> str:
    """Get list of files in the specified collection"""
    try:
        # Get the full collection name
        collection_name = f"{user_id}_{collection_id}"
        
        # Open the table and convert to pandas
        table = db.open_table(collection_name)
        df = table.to_pandas()
        logging.info(f"fetched chunks {str(df.head())}")
        
        # Get unique file names
        unique_files = df['file_name'].unique()
        
        # Join the file names into a string
        return ", ".join(unique_files)
    except Exception as e:
        logging.error(f"Error getting collection files: {str(e)}")
        return f"Error getting files: {str(e)}"


@router.post("/query_collection_tool")
async def query_collection_tool(input_data: QueryInput):
    try:
        response = await query_collection(input_data)
        results = []
        
        # Access response directly since it's a Pydantic model
        for r in response.results:
            result_dict = {
                "text": r.text,
                "distance": r.distance,
                "metadata": {
                    "document_id": r.metadata.get("document_id"),
                    "chunk_index": r.metadata.get("location", {}).get("chunk_index")
                }
            }
            results.append(result_dict)
            
        return str(results)
        
    except Exception as e:
        logging.error(f"Unexpected error during query: {str(e)}")
        raise HTTPException(
            status_code=500,
            detail=f"Unexpected error: {str(e)}"
        )