Simple RAG Using Google AI

README.md

RAG Implementation with Google Gemini and ChromaDB

A simple example implementation of Retrieval Augmented Generation (RAG) using Google's Gemini API and ChromaDB for document storage and retrieval.

Overview

This code demonstrates a basic RAG system with three main components:

1. Document embedding using Google's text-embedding-004 model
2. Vector storage and retrieval using ChromaDB
3. Question answering using Gemini 1.5 Flash model

Technical Components

Embedding System

• Custom GeminiEmbeddingFunction class that implements ChromaDB's EmbeddingFunction
• Supports both document and query embedding modes
• Uses Google's text-embedding-004 model with retry logic for API resilience

Vector Database

• Uses ChromaDB for vector storage and similarity search
• DocumentDatabase class provides a simple interface for:
  • Document storage with automatic ID assignment
  • Query-based document retrieval
  • Collection management

Question Answering

• Implements a basic RAG pipeline:
  1. Converts user question into embeddings
  2. Retrieves relevant document using similarity search
  3. Constructs a prompt combining the question and retrieved context
  4. Generates answer using Gemini 1.5 Flash model

Example Usage

The code includes sample documents for demonstration, but can be used with any text corpus. The RAG pipeline can be invoked with:

db = DocumentDatabase()
db.store_documents(your_documents)
answer = get_answer("your question", db)

Dependencies

• google-generativeai: For Gemini API access
• chromadb: For vector storage
• configparser: For API key management

License

Apache License 2.0. See LICENSE for the full license text.

app.cfg

[ai-studio]
GOOGLE_API_KEY = YOUR_AI_STUDIO_API_KEY

rag.py

import configparser as ConfigParser
import google.generativeai as genai
import chromadb
from google.api_core import retry
from chromadb import Documents, EmbeddingFunction, Embeddings

# Configuration and initialization


def initialize_system():
    config = ConfigParser.ConfigParser()
    config.read(r'./app.cfg')
    api_key = config.get('ai-studio', 'GOOGLE_API_KEY')
    genai.configure(api_key=api_key)
    return api_key


# Document handling
DOCUMENTS = [
    """Operating the Climate Control System  Your Googlecar has a climate control system that allows you to adjust the temperature and airflow in the car. To operate the climate control system, use the buttons and knobs located on the center console.  Temperature: The temperature knob controls the temperature inside the car. Turn the knob clockwise to increase the temperature or counterclockwise to decrease the temperature. Airflow: The airflow knob controls the amount of airflow inside the car. Turn the knob clockwise to increase the airflow or counterclockwise to decrease the airflow. Fan speed: The fan speed knob controls the speed of the fan. Turn the knob clockwise to increase the fan speed or counterclockwise to decrease the fan speed. Mode: The mode button allows you to select the desired mode. The available modes are: Auto: The car will automatically adjust the temperature and airflow to maintain a comfortable level. Cool: The car will blow cool air into the car. Heat: The car will blow warm air into the car. Defrost: The car will blow warm air onto the windshield to defrost it.""",

    """Your Googlecar has a large touchscreen display that provides access to a variety of features, including navigation, entertainment, and climate control. To use the touchscreen display, simply touch the desired icon.  For example, you can touch the "Navigation" icon to get directions to your destination or touch the "Music" icon to play your favorite songs.""",

    """Shifting Gears Your Googlecar has an automatic transmission. To shift gears, simply move the shift lever to the desired position.  Park: This position is used when you are parked. The wheels are locked and the car cannot move. Reverse: This position is used to back up. Neutral: This position is used when you are stopped at a light or in traffic. The car is not in gear and will not move unless you press the gas pedal. Drive: This position is used to drive forward. Low: This position is used for driving in snow or other slippery conditions."""
]

# Embedding functionality


class GeminiEmbeddingFunction(EmbeddingFunction):
    def __init__(self, document_mode=True):
        self.document_mode = document_mode

    def __call__(self, input: Documents) -> Embeddings:
        task_type = "retrieval_document" if self.document_mode else "retrieval_query"
        response = genai.embed_content(
            model="models/text-embedding-004",
            content=input,
            task_type=task_type,
            request_options={"retry": retry.Retry(
                predicate=retry.if_transient_error)}
        )
        return response["embedding"]

# Database operations

class DocumentDatabase:
    def __init__(self, db_name="googlecardb"):
        self.db_name = db_name
        self.client = chromadb.Client()

    def get_db(self, document_mode=True):
        return self.client.get_or_create_collection(
            name=self.db_name,
            embedding_function=GeminiEmbeddingFunction(document_mode)
        )

    def store_documents(self, documents):
        db = self.get_db(True)
        db.add(documents=documents, ids=[str(i) for i in range(len(documents))])
        return db.count()

    def query(self, question):
        db = self.get_db(False)
        result = db.query(query_texts=[question], n_results=1)
        return result["documents"][0][0]

# QA functionality


def get_answer(question, db):
    # Get model and reference data
    flash = genai.GenerativeModel('gemini-1.5-flash')
    reference = db.query(question)

    # Clean text and create prompt
    prompt = f"""
    You are a helpful and informative agent that answers questions using text from
    the reference answers included below.
    You are talking to a non-technical audience, so be sure to break down any
    complicated concepts into an easy to understand format, and maintain a
    friendly and conversational tone.
    Be sure to respond in a complete sentence, and include all relevant
    background information.
    If the passage is irrelevant to the answer, you may ignore it.

    QUESTION: {question.replace("\n", " ")}
    PASSAGE: {reference.replace("\n", " ")}
    """

    # return flash(prompt)
    response = flash.generate_content(prompt)
    return response.text


def main():
    # Initialize system
    initialize_system()

    # Set up database and store documents
    db = DocumentDatabase()
    doc_count = db.store_documents(DOCUMENTS)
    print(f"Stored {doc_count} documents in the database")

    # Example queries
    questions = [
        "How do you use the touchscreen to play music?",
        "How do you shift gears in the Googlecar?",
        "How do you operate the climate control system?"
    ]

    for question in questions:
        answer = get_answer(question, db)
        print(f"\nQ: {question}")
        print(f"A: {answer}")


if __name__ == "__main__":
    main()