Serverless RAG with Gemini and Pinecone

This project implements a serverless Retrieval-Augmented Generation (RAG) API using FastAPI, Pinecone, and Gemini. It's designed to be deployed on AWS Lambda as a container image.

Features

• Ingest PDF documents: Upload a PDF file to be chunked, embedded, and stored in a Pinecone serverless index.
• Query your documents: Ask questions and get answers generated by Gemini, based on the context retrieved from your documents.
• Serverless: Deployed on AWS Lambda for scalability and cost-effectiveness.
• FastAPI: A modern, fast (high-performance) web framework for building APIs.
• SentenceTransformers: Uses the all-MiniLM-L6-v2 model for generating high-quality embeddings.
• Pinecone: A serverless vector database for efficient similarity search.
• Gemini: Google's powerful language model for generating answers.

Project Structure

serverless-rag-gemini-lambda/
  app/
    main.py              # FastAPI application (for local development)
    handler.py           # Lambda handler wrapper (for AWS Lambda)
    rag.py
    utils.py
    requirements.txt     # Dependencies for local development
    requirements-lambda.txt  # Dependencies for AWS Lambda (includes mangum)
  Dockerfile             # Dockerfile for local development
  Dockerfile.lambda      # Dockerfile for AWS Lambda deployment
  README.md
  scripts/
    create_pinecone_index.py
  diagrams/
    sequence_diagram.md
    flowcharts.md

Setup

1. Environment Variables

Create a .env file in the root of the project and add the following environment variables:

PINECONE_API_KEY=your_pinecone_api_key
GOOGLE_API_KEY=your_google_api_key

2. Create Pinecone Index

Run the following script to create the Pinecone index:

export $(cat .env | xargs)
python scripts/create_pinecone_index.py

Running Locally

Option 1: Using Uvicorn (Recommended for Development)

1. Install dependencies:

   cd app
   pip install -r requirements.txt

2. Run the FastAPI application:

   # From the project root
   uvicorn app.main:app --reload

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

Option 2: Using Docker (Local Development)

1. Build the Docker image:

   docker build -t serverless-rag .

2. Run the Docker container:

   docker run -p 8000:8000 --env-file .env serverless-rag

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

Deployment to AWS Lambda

Differences Between Local and Lambda Deployment

• Local Development (Dockerfile): Uses python:3.12-slim base image and runs uvicorn directly
• Lambda Deployment (Dockerfile.lambda): Uses AWS Lambda Python base image and wraps FastAPI with Mangum adapter

Prerequisites

• AWS CLI configured with appropriate credentials
• Docker installed and running
• AWS ECR repository created (or permissions to create one)

Deployment Steps

1. Build the Docker image for Lambda:

   docker build -t serverless-rag-lambda -f Dockerfile.lambda .

2. Authenticate Docker with ECR:

   aws ecr get-login-password --region <your-aws-region> | docker login --username AWS --password-stdin <your-aws-account-id>.dkr.ecr.<your-aws-region>.amazonaws.com

   Example:

   aws ecr get-login-password --region us-east-1 | docker login --username AWS --password-stdin 123456789012.dkr.ecr.us-east-1.amazonaws.com

3. Create an ECR repository (if it doesn't exist):

   aws ecr create-repository --repository-name serverless-rag --region <your-aws-region>

4. Tag and push the image to ECR:

   docker tag serverless-rag-lambda:latest <your-aws-account-id>.dkr.ecr.<your-aws-region>.amazonaws.com/serverless-rag:latest
   docker push <your-aws-account-id>.dkr.ecr.<your-aws-region>.amazonaws.com/serverless-rag:latest

   Example:

   docker tag serverless-rag-lambda:latest 123456789012.dkr.ecr.us-east-1.amazonaws.com/serverless-rag:latest
   docker push 123456789012.dkr.ecr.us-east-1.amazonaws.com/serverless-rag:latest

5. Create a Lambda function:

   • Go to the AWS Lambda console and create a new function
   • Choose "Container image" as the package type
   • Select the ECR image you just pushed
   • Configure the function:
     • Handler: app.handler.handler (already set in Dockerfile.lambda)
     • Runtime: Python 3.12 (from base image)
     • Memory: Recommended 3008 MB (for sentence-transformers)
     • Timeout: Recommended 900 seconds (15 minutes)
     • Environment Variables:
       • PINECONE_API_KEY: Your Pinecone API key
       • GOOGLE_API_KEY: Your Google API key

6. Create an API Gateway:

   • Create a new REST API in the API Gateway console
   • Create resources and methods for your endpoints:
   • GET /health
   • POST /ingest
   • POST /query
   • Configure each method to use the Lambda function as the integration type
   • Enable CORS if needed
   • Deploy the API to a stage (e.g., prod or dev)

Lambda-Specific Files

• Dockerfile.lambda: Uses AWS Lambda Python base image and sets the handler to app.handler.handler
• requirements-lambda.txt: Includes mangum adapter (required for FastAPI on Lambda) instead of uvicorn
• app/handler.py: Lambda handler that wraps the FastAPI app with Mangum

Testing Lambda Locally

You can test the Lambda handler locally using the AWS SAM CLI or by invoking it directly:

# Build the Lambda image
docker build -t serverless-rag-lambda -f Dockerfile.lambda .

# Test locally (requires AWS Lambda Runtime Interface Emulator)
docker run -p 9000:8080 --env-file .env serverless-rag-lambda

Then test with:

curl -XPOST "http://localhost:9000/2015-03-31/functions/function/invocations" -d '{}'

API Gateway Usage Plan for Rate Limiting

To protect your API from abuse, you can create a usage plan to throttle and set quotas for your API.

1. Create a usage plan:
   • In the API Gateway console, go to "Usage Plans" and create a new plan.
   • Configure the throttling and quota limits as needed.
2. Associate the usage plan with your API stage:
   • Associate the usage plan with the stage you deployed your API to.
3. Create an API key:
   • Create an API key and associate it with the usage plan.
4. Require API key for your API:
   • In your API's method settings, require an API key.

Example curl Commands

Ingest a PDF

curl -X POST "http://localhost:8000/ingest" \
-H "Content-Type: multipart/form-data" \
-F "file=@/path/to/your/document.pdf" \
-F "namespace=my-namespace"

Query your documents

curl -X POST "http://localhost:8000/query" \
-H "Content-Type: application/json" \
-d '{
  "question": "What is the main topic of the document?",
  "k": 3,
  "namespace": "my-namespace"
}'

Embeddings Model: all-MiniLM-L6-v2

This project uses the all-MiniLM-L6-v2 model from SentenceTransformers for generating embeddings. It is a small, fast, and high-quality model that is well-suited for serverless applications. It maps sentences & paragraphs to a 384 dimensional dense vector space and can be used for tasks like clustering or semantic search.

Diagrams

• Sequence Diagram
• Flowcharts

Test