RAG Query System

This project implements a Retrieval-Augmented Generation (RAG) system using FastAPI, PostgreSQL with pgvector, and Azure OpenAI.

Setup Instructions

1. Clone the repository:

   git clone <repository-url>
   cd rag_query_system
   

2. Set up the environment file:

   • Copy the .env template file to .env:

     cp ".env template" .env
     

   • Open the .env file and fill in the necessary values:
     • Add your Azure OpenAI API key
     • Set secure passwords for the PostgreSQL database
     • Add any other required environment variables

3. Build and start the Docker containers:

   docker-compose up --build
   

4. The API will be available at http://localhost:8000.

Usage

Uploading a Document

To upload a PDF document:

curl -X POST -F "file=@path/to/your/file.pdf" http://localhost:8000/upsert

Querying

To ask a question:

curl -X POST -H "Content-Type: application/json" -d '{"text":"Your question here"}' http://localhost:8000/query

Testing with Python Script

We provide a Python script test_rag.py to test the RAG system. This script allows you to upload documents and perform queries programmatically.

To run the test script:

1. Ensure your Docker containers are running.
2. Create a folder in the same path as the test_rag.py file named pdfs.
3. The PDFs specified in the script are four research papers. The sources are so:
   1. ScanNet: Richly-annotated 3D Reconstructions of Indoor Scenes [https://arxiv.org/abs/1702.04405]
   2. DALES: A Large-scale Aerial LiDAR Data Set for Semantic Segmentation [https://arxiv.org/abs/2004.11985]
   3. ECLAIR: A High-Fidelity Aerial LiDAR Dataset for Semantic Segmentation [https://arxiv.org/abs/2404.10699v1]
   4. DALES Objects: A Large Scale Benchmark Dataset for Instance Segmentation in Aerial Lidar [https://ieeexplore.ieee.org/stamp/stamp.jsp?tp=&arnumber=9469802]
4. Download these PDFs and place them in the pdfs folder.
5. Run the following command:

python test_rag.py

This script will:

• Uploads 4 test PDF documents to the RAG system
• Perform sample queries
  • "Explain Dales dataset and DalesObjects datasets?",
  • "Can you compare the different datasets that are in the documents?",
  • "Which dataset is most well tested?",
• Print the results

Sample Output

Here's an example of what to expect when running the test script:

(base) Anands-MacBook-Pro:rag_query_system anand$ python test_rag.py 
Clear DB response: {'message': 'Database cleared successfully'}
Processing pdf:  ./pdfs/1702.04405v2.pdf
Upsert response: {'message': 'Document tmpggwfw4wk uploaded and embedded successfully'}
Processing pdf:  ./pdfs/2004.11985v1.pdf
Upsert response: {'message': 'Document tmppus_z3nk uploaded and embedded successfully'}
Processing pdf:  ./pdfs/2404.10699v1.pdf
Upsert response: {'message': 'Document tmpz0rfl_02 uploaded and embedded successfully'}
Processing pdf:  ./pdfs/DALES_Objects_A_Large_Scale_Benchmark_Dataset_for_Instance_Segmentation_in_Aerial_Lidar.pdf
Upsert response: {'message': 'Document tmpgxqhuz1w uploaded and embedded successfully'}
Query: Explain Dales dataset and DalesObjects datasets?
Response: {'answer': 'The DALES dataset is a benchmark dataset for instance segmentation in aerial lidar. It contains close to half a billion hand-labeled points, including both semantic and instance segmentation labels. The dataset covers over ten square kilometers of aerial lidar data collected over the City of Surrey in British Columbia, Canada. It includes eight object categories: ground, vegetation, buildings, cars, trucks, powerlines, poles, and fences.\n\nThe DALES Objects dataset is an extension of the DALES dataset. It offers a meticulously hand-labeled dataset that contains eight semantic object categories and over twenty thousand hand-labeled instances. It is one of the most extensive instance segmentation datasets taken with aerial lidar and one of the first to include both rural and urban scenes. The addition of rural scenes provides essential information for tasks such as forestry management and utility asset monitoring. The dataset provides high spatial accuracy and a unique set of viewpoints and occlusions. It allows researchers to test their algorithms in a diverse group of settings and gain a better understanding of network performance.'}
Query: Can you compare the different datasets that are in the documents?
Response: {'answer': 'Yes, the document provides a comparison of different datasets in terms of area of coverage, point density, and sensor type. For example, the DALES dataset covers an area of 10×106m2 with a point density of 28 points per square meter and uses aerial laser scanning (ALS). On the other hand, the proposed ECLAIR dataset covers an area of 10.3×106m2 with a point density of 11 points per square meter and also uses ALS. The S3DIS dataset, which is an indoor dataset, consists of RGB-D sequences captured by short-range depth scanners with low resolution and limited semantic annotations.'}
Query: which dataset is most well tested?
Response: {'answer': 'Based on the given context, the dataset that is most well tested is the ScanNet dataset. It is mentioned that training on ScanNet translates well to testing on SceneNN, and the test results on SceneNN are significantly improved by using the training data from ScanNet.'}

You can modify test_rag.py to add more test cases or customize the queries according to your needs.

Note: Make sure you have the required Python libraries installed to run the test script. You can install them using:

pip install requests

Project Structure

• .dockerignore: Specifies files and directories to be excluded from the Docker build context.
• .env template: Template for environment variables. Copy this to .env and fill in your specific values.
• .gitignore: Specifies files and directories to be ignored by Git.
• closed_ai.py: Contains OpenAI operations for the RAG system.
• docker-compose-redis.yml: Docker Compose file for setting up a Redis-based system (alternative configuration).
• docker-compose.yml: Main Docker Compose file for setting up the entire system.
• Dockerfile: Defines the Docker image for the main application.
• Dockerfile.postgres: Defines the Docker image for the PostgreSQL database with pgvector.
• environment.yml: Specifies the Conda environment for the project.
• fastest_api.py: Main FastAPI application file.
• init.sql: SQL initialization script for PostgreSQL.
• postgres_ops.py: Contains PostgreSQL operations for the RAG system.
• rag_system.py: Implements the core RAG system functionality.
• README.md: This file, containing project documentation.
• redis_ops.py: Contains Redis operations (for alternative Redis-based configuration).
• test_rag.py: Python script for testing the RAG system.

Additional Notes

• Make sure you have Docker and Docker Compose installed on your system.
• The system uses PostgreSQL with pgvector for efficient vector similarity search.
• An alternative Redis-based configuration is available using docker-compose-redis.yml.
• The closed_ai.py file contains the integration with Azure OpenAI. Ensure your API key and endpoint are correctly set in the .env file.
• The fastest_api.py file is the main entry point for the FastAPI application.
• postgres_ops.py and redis_ops.py provide database operations for PostgreSQL and Redis respectively. The system primarily uses PostgreSQL, but Redis configuration is available as an alternative.