RAGatouille API Server

serve/.env.example

@@ -0,0 +1,4 @@
+# .env file
+PROJECT_NAME=RAGatouille API
+MODEL_NAME=colbert-ir/colbertv2.0
+INDEX_ROOT=local_store

serve/README.md

@@ -0,0 +1,82 @@
+# RAGatouille API Server
+
+RAGatouille API is a FastAPI application designed to provide an efficient and scalable way to perform operations such as document indexing, searching, reranking, and encoding with the RAGPretrainedModel. It leverages the power of FastAPI to offer high performance and easy-to-use RESTful endpoints.
+
+## Features
+
+- **Document Indexing and Management**: Easily index and manage documents with comprehensive endpoints for adding, updating, and deleting documents.
+- **Search and Rerank Functionality**: Perform advanced search queries and rerank search results to meet specific requirements.
+- **Document Encoding and Retrieval**: Encode documents for efficient storage and retrieval, and search through encoded documents.
+- **RESTful API**: Clear and concise endpoints adhering to RESTful standards, making the API easy to consume.
+
+## Getting Started
+
+### Prerequisites
+
+- Python 3.8 or higher
+- pip and virtualenv
+
+### Installation
+
+1. Clone the repository:
+
+   ```bash
+   https://github.com/bclavie/RAGatouille.git
+   cd RAGatouille/serve
+   ```
+2. Set up a virtual environment (optional but recommended): <br>
+   Windows
+   ```bash
+   python -m venv venv
+   source venv\Scripts\activate
+   ```
+   Others
+   ```bash
+   python -m venv venv
+   source venv/bin/activate
+   ```
+3. Install dependencies:
+   ```bash
+   pip install -r requirements.txt
+   ```
+   or use Poetry
+   ```bash
+   poetry install
+   ```
+4. Copy .env.example to .env:
+   ```bash
+   cp .env.example .env
+   ```
+
+### Running the Application
+Start the FastAPI server with:
+
+```bash
+uvicorn ragatouille_serve.main:app --reload
+```
+or start with Poetry
+```bash
+poetry run uvicorn ragatouille_serve.main:app --reload
+```
+The --reload flag enables hot reloading during development.
+
+### Accessing the API Documentation
+Once the server is running, you can view the auto-generated Swagger UI documentation by navigating to:
+```bash
+http://127.0.0.1:8000/docs
+```
+
+## API Endpoints
+
+### Index
+- **POST /api/v1/index/:** Create an index with documents.
+### Document Management
+- **POST /api/v1/documents/search/:** Search documents.
+- _**[ISSUE] POST /api/v1/documents/add/:** Add documents to an existing index._
+- _**[ISSUE] DELETE /api/v1/documents/delete/:** Delete documents from an index._
+### Encode
+- **POST /api/v1/encode/:** Encode documents.
+- **POST /api/v1/encode/search/:** Search through encoded documents.
+- **POST /api/v1/encode/clear/:** Clear all encoded documents.
+### Rerank
+- **POST /api/v1/rerank/:** Rerank a set of documents based on a query.

serve/pyproject.toml

@@ -0,0 +1,24 @@
+[tool.poetry]
+name = "ragatouille_serve"
+version = "0.1.0"
+description = "FastAPI server for ragatouille"
+authors = ["Ram Chandalada"]
+readme = "README.md"
+
+[tool.poetry.dependencies]
+python = ">=3.9,<4.0"
+fastapi = "^0.109.2"
+uvicorn = "^0.27.1"
+pydantic = "^2.6.1"
+requests = "^2.31.0"
+starlette = ">=0.36.3,<0.37.0"
+python-multipart = "^0.0.9"
+pydantic-settings = "^2.1.0"
+RAGatouille = "0.0.6rc2"
+
+[tool.poetry.group.dev.dependencies]
+pytest = "^8.0.0"
+
+[build-system]
+requires = ["poetry-core"]
+build-backend = "poetry.core.masonry.api"

serve/ragatouille_serve/api/api_v1.py

@@ -0,0 +1,12 @@
+from fastapi import APIRouter
+from ragatouille_serve.api.endpoints import index, document_management, rerank
+from ragatouille_serve.api.endpoints import encode
+
+# Versioned API router
+api_router = APIRouter()
+
+# Include routers from different endpoint modules
+api_router.include_router(index.router, tags=["Index"], prefix="/index")
+api_router.include_router(document_management.router, tags=["Document Management"], prefix="/documents")
+api_router.include_router(encode.router, tags=["Encode"], prefix="/encode")
+api_router.include_router(rerank.router, tags=["Rerank"], prefix="/rerank")

serve/ragatouille_serve/api/endpoints/document_management.py

@@ -0,0 +1,55 @@
+from fastapi import APIRouter, Depends, HTTPException, status
+from ragatouille_serve.models.payloads import AddToIndexQuery, DeleteFromIndexQuery, SearchQuery
+from ragatouille_serve.core.rag_model import get_rag_model, delete_rag_model
+import logging
+
+router = APIRouter()
+logger = logging.getLogger(__name__)
+
+@router.post("/add", status_code=status.HTTP_201_CREATED)
+async def add_to_index(query: AddToIndexQuery, rag=Depends(get_rag_model)):
+    try:
+        rag.add_to_index(
+            new_collection=query.new_collection,
+            new_document_ids=query.new_document_ids,
+            new_document_metadatas=query.new_document_metadatas,
+            index_name=query.index_name,
+            split_documents=query.split_documents
+        )
+        return {"message": "Documents added to index successfully"}
+    except Exception as e:
+        logger.error(f"Failed to add documents to index: {e}")
+        raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail="Failed to add documents to index")
+
+@router.delete("/delete", status_code=status.HTTP_200_OK)
+async def delete_from_index(query: DeleteFromIndexQuery, rag=Depends(get_rag_model)):
+    try:
+        rag.delete_from_index(
+            document_ids=query.document_ids,
+            index_name=query.index_name
+        )
+        # delete_rag_model()
+        return {"message": "Documents deleted from index successfully"}
+    except Exception as e:
+        logger.error(f"Failed to delete documents from index: {e}")
+        raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail="Failed to delete documents from index")
+
+@router.post("/search")
+async def search(query: SearchQuery, rag=Depends(get_rag_model)):
+    try:
+        logger.info(f"Initiating search with query: '{query.query}'. Index: '{query.index_name}', k: {query.k}")
+        results = rag.search(
+            query=query.query, 
+            index_name=query.index_name, 
+            k=query.k, 
+            force_fast=query.force_fast, 
+            zero_index_ranks=query.zero_index_ranks
+        )
+        logger.info(f"Search completed successfully for query: '{query.query}'.")
+        return results
+    except Exception as e:
+        logger.error(f"Search operation failed for query: '{query.query}': {e}", exc_info=True)
+        raise HTTPException(
+            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, 
+            detail="Search operation failed"
+        )

serve/ragatouille_serve/api/endpoints/encode.py

@@ -0,0 +1,59 @@
+from fastapi import APIRouter, Depends, HTTPException, status
+from ragatouille_serve.models.payloads import EncodeQuery, SearchEncodedDocsQuery, ClearEncodedDocsQuery
+from ragatouille_serve.core.rag_model import get_rag_model
+import logging
+
+router = APIRouter()
+logger = logging.getLogger(__name__)
+
+@router.post("/")
+async def encode(query: EncodeQuery, rag=Depends(get_rag_model)):
+    try:
+        logger.info(f"Starting document encoding. Batch size: {query.bsize}, Verbose: {query.verbose}")
+        rag.encode(
+            documents=query.documents,
+            bsize=query.bsize,
+            document_metadatas=query.document_metadatas,
+            verbose=query.verbose,
+            max_document_length=query.max_document_length
+        )
+        logger.info(f"Documents encoded successfully. Total documents: {len(query.documents)}")
+        return {"message": "Documents encoded successfully"}
+    except Exception as e:
+        logger.error(f"Failed to encode documents: {e}", exc_info=True)
+        raise HTTPException(
+            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, 
+            detail="Failed to encode documents"
+        )
+
+@router.post("/clear")
+async def clear_encoded_docs(query: ClearEncodedDocsQuery, rag=Depends(get_rag_model)):
+    try:
+        logger.info(f"Clearing encoded documents. Force: {query.force}")
+        rag.clear_encoded_docs(force=query.force)
+        logger.info("Encoded documents cleared successfully.")
+        return {"message": "Encoded documents cleared successfully"}
+    except Exception as e:
+        logger.error(f"Failed to clear encoded documents: {e}", exc_info=True)
+        raise HTTPException(
+            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, 
+            detail="Failed to clear encoded documents"
+        )
+
+@router.post("/search")
+async def search_encoded_docs(query: SearchEncodedDocsQuery, rag=Depends(get_rag_model)):
+    try:
+        logger.info(f"Starting search in encoded documents. Query: '{query.query}'")
+        results = rag.search_encoded_docs(
+            query=query.query,
+            k=query.k,
+            bsize=query.bsize
+        )
+        logger.info(f"Search in encoded documents completed successfully. Query: '{query.query}'")
+        return results
+    except Exception as e:
+        logger.error(f"Search in encoded documents failed for query: '{query.query}': {e}", exc_info=True)
+        raise HTTPException(
+            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+            detail="Search in encoded documents failed"
+        )

serve/ragatouille_serve/api/endpoints/index.py

@@ -0,0 +1,29 @@
+from fastapi import APIRouter, Depends, HTTPException, status
+from ragatouille_serve.models.payloads import IndexQuery
+from ragatouille_serve.core.rag_model import get_rag_model
+import logging
+
+router = APIRouter()
+logger = logging.getLogger(__name__)
+
+@router.post("/")
+async def create_index(index_query: IndexQuery, rag=Depends(get_rag_model)):
+    try:
+        logger.info(f"Starting index creation. Index name: {index_query.index_name}, Overwrite: {index_query.overwrite_index}")
+        index_path = rag.index(
+            collection=index_query.collection,
+            document_ids=index_query.document_ids,
+            document_metadatas=index_query.document_metadatas,
+            index_name=index_query.index_name,
+            overwrite_index=index_query.overwrite_index,
+            max_document_length=index_query.max_document_length,
+            split_documents=index_query.split_documents
+        )
+        logger.info(f"Index created successfully. Index path: {index_path}")
+        return {"index_path": index_path}
+    except Exception as e:
+        logger.error(f"Failed to create index: {e}", exc_info=True)
+        raise HTTPException(
+            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, 
+            detail="Failed to create index"
+        )

serve/ragatouille_serve/api/endpoints/rerank.py

@@ -0,0 +1,27 @@
+from fastapi import APIRouter, Depends, HTTPException, status
+from ragatouille_serve.models.payloads import RerankQuery
+from ragatouille_serve.core.rag_model import get_rag_model
+import logging
+
+router = APIRouter()
+logger = logging.getLogger(__name__)
+
+@router.post("/")
+async def rerank(query: RerankQuery, rag=Depends(get_rag_model)):
+    try:
+        logger.info(f"Starting rerank operation for query: '{query.query}' with {len(query.documents)} documents.")
+        results = rag.rerank(
+            query=query.query,
+            documents=query.documents,
+            k=query.k,
+            zero_index_ranks=query.zero_index_ranks,
+            bsize=query.bsize
+        )
+        logger.info(f"Rerank operation completed successfully for query: '{query.query}'.")
+        return results
+    except Exception as e:
+        logger.error(f"Rerank operation failed for query: '{query.query}': {e}", exc_info=True)
+        raise HTTPException(
+            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+            detail="Rerank operation failed"
+        )

serve/ragatouille_serve/core/config.py

@@ -0,0 +1,14 @@
+from pydantic_settings import BaseSettings
+
+class Settings(BaseSettings):
+    PROJECT_NAME: str = "RAGatouille API"
+    PROJECT_VERSION: str = "0.1.0"
+    API_V1_STR: str = "/api/v1"
+    MODEL_NAME: str = "colbert-ir/colbertv2.0"
+    INDEX_ROOT: str = "local_store"
+
+    class Config:
+        case_sensitive = True
+        env_file = ".env"
+
+settings = Settings()

serve/ragatouille_serve/core/rag_model.py

@@ -0,0 +1,37 @@
+from fastapi import Depends
+from ragatouille_serve.core.config import settings
+from ragatouille import RAGPretrainedModel
+import logging
+
+# Set up a logger
+logger = logging.getLogger(__name__)
+
+# Singleton pattern to avoid loading the model multiple times
+class RAGModel:
+    model_instance: RAGPretrainedModel = None
+
+    @classmethod
+    def get_instance(cls) -> RAGPretrainedModel:
+        if cls.model_instance is None:
+            logger.info("Initializing RAGPretrainedModel instance...")
+            # Initialize and configure RAGPretrainedModel
+            cls.model_instance = RAGPretrainedModel.from_pretrained(pretrained_model_name_or_path=settings.MODEL_NAME, index_root=settings.INDEX_ROOT)
+            logger.info("RAGPretrainedModel instance created successfully.")
+        else:
+            logger.debug("RAGPretrainedModel instance already exists, reusing the existing instance.")
+        return cls.model_instance
+
+    @classmethod
+    def delete_instance(cls):
+        if cls.model_instance is not None:
+            logger.info("Deleting RAGPretrainedModel instance...")
+            cls.model_instance = None
+            logger.info("RAGPretrainedModel instance deleted successfully.")
+        else:
+            logger.debug("RAGPretrainedModel instance does not exist, nothing to delete.")
+
+def get_rag_model() -> RAGPretrainedModel:
+    return RAGModel.get_instance()
+
+def delete_rag_model():
+    RAGModel.delete_instance()

serve/ragatouille_serve/main.py

@@ -0,0 +1,30 @@
+from fastapi import FastAPI
+from ragatouille_serve.api.api_v1 import api_router
+from ragatouille_serve.core.config import Settings
+from ragatouille_serve.utils.logging_config import setup_logging
+import logging
+
+# Initialize FastAPI app and settings
+settings = Settings()
+app = FastAPI(title=settings.PROJECT_NAME, version=settings.PROJECT_VERSION)
+
+# Set up logging
+setup_logging()
+logger = logging.getLogger(__name__)
+
+# Include API routers
+app.include_router(api_router, prefix=settings.API_V1_STR)
+
+# FastAPI startup event
+@app.on_event("startup")
+async def startup_event():
+    logger.info("Application startup")
+
+# FastAPI shutdown event
+@app.on_event("shutdown")
+async def shutdown_event():
+    logger.info("Application shutdown")
+
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=8000, reload=True)