Gemini API Python SDKとChromaDBを使用してRAG systemを開発する【ハンズオン】

!pip install -qU "google-genai" "chromadb"

from google import genai
from google.genai import types

from IPython.display import Markdown

genai.__version__

※出力

'1.7.0'

GOOGLE_GENAI_USE_VERTEXAI=True
GOOGLE_CLOUD_PROJECT='your-project-id' # #insert your-project-id in VERTEXAI
GOOGLE_CLOUD_LOCATION='us-central1'

client = genai.Client(vertexai = GOOGLE_GENAI_USE_VERTEXAI, project = GOOGLE_CLOUD_PROJECT, location = GOOGLE_CLOUD_LOCATION)

for m in client.models.list():
    print(m.name)

※出力　一部記載

publishers/google/models/video-text-detection
publishers/google/models/video-speech-transcription
publishers/google/models/path-foundation
publishers/google/models/derm-foundation
publishers/google/models/txgemma
publishers/google/models/hear
publishers/google/models/image-segmentation-001
publishers/google/models/gemma3
publishers/google/models/shieldgemma2
publishers/google/models/gemini-2.0-flash-001
publishers/google/models/text-embedding-large-exp-03-07
publishers/google/models/gemini-2.0-flash-lite-001
publishers/google/models/weathernext
publishers/google/models/gemini-2.5-flash-preview-04-17
publishers/google/models/gemini-2.5-pro-preview-05-06

DOCUMENT1 = "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."
DOCUMENT2 = '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.'
DOCUMENT3 = "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."

documents = [DOCUMENT1, DOCUMENT2, DOCUMENT3]

from chromadb import Documents, EmbeddingFunction, Embeddings
from google.api_core import retry

from google.genai import types

# Define a helper to retry when per-minute quota is reached.
is_retriable = lambda e: (isinstance(e, genai.errors.APIError) and e.code in {429, 503})

class GeminiEmbeddingFunction(EmbeddingFunction):
    # Specify whether to generate embeddings for documents, or queries
    document_mode = True
    
    @retry.Retry(predicate=is_retriable)
    def __call__(self, input: Documents) -> Embeddings:
        if self.document_mode:
            embedding_task = "RETRIEVAL_DOCUMENT"
        else:
            embedding_task = "RETRIEVAL_QUERY"
            
        response = client.models.embed_content(
            model = "text-embedding-005",
            contents = input,
            config = types.EmbedContentConfig(
                task_type = embedding_task,
            ),
        )
        return [e.values for e in response.embeddings]

import chromadb

DB_NAME = "googlecardb"

embed_fn = GeminiEmbeddingFunction()
embed_fn.document_mode = True

chroma_client = chromadb.Client()
db = chroma_client.get_or_create_collection(name = DB_NAME, embedding_function = embed_fn)

db.add(documents = documents, ids = [str(i) for i in range(len(documents))])

db.count()

※出力

3

# You can peek at the data too.
db.peek(1)

※出力

{'ids': ['0'],
 'embeddings': array([[-4.03229967e-02,  1.41243562e-02,  2.56362278e-02,
         -2.36929134e-02, -4.38607670e-02,  6.20906912e-02,
          2.42945785e-03, -3.60977016e-02,  7.42293447e-02,
　　　 　 　....     
         -3.50930132e-02,  1.09769069e-02,  5.55420443e-02,
          5.83409984e-03, -1.52970320e-02, -3.93905863e-03]]),
 '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.'],
 'uris': None,
 'included': ['metadatas', 'documents', 'embeddings'],
 'data': None,
 'metadatas': [None]}

# Switch to query mode when generating embeddings.
embed_fn.document_mode = False

# Search the Chroma DB using the specified query.
query =  "How do you use the touchscreen to play music?"

result = db.query(query_texts = [query], n_results = 1)
[all_passages] = result["documents"]

Markdown(all_passages[0])

※出力

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.

query_oneline = query.replace("\n", " ")

# This prompt is where you can specify any guidance on tone, or what topics the model should stick to, or avoid.
prompt = f""""You are a helpful and informative bot  that answers questions using text from the reference passage included below.
Be sure to respond in a complete sentence, being comprehensive, including all relevant background information.
However, you are talking to a non-technical audience, so be sure to break down complicated concepts and
strike a friendly and converstional tone. If the passage is irrelevant to the answer, you may ignore it.

QUESTION: {query_oneline}
"""

# Add the retrieved documents to the prompt.
for passage in all_passages:
    passage_oneline = passage.replace("\n", " ")
    prompt += f"PASSAGE: {passage_oneline}\n"
    
print(prompt)

※出力

"You are a helpful and informative bot  that answers questions using text from the reference passage included below.
Be sure to respond in a complete sentence, being comprehensive, including all relevant background information.
However, you are talking to a non-technical audience, so be sure to break down complicated concepts and
strike a friendly and converstional tone. If the passage is irrelevant to the answer, you may ignore it.

QUESTION: How do you use the touchscreen to play music?
PASSAGE: 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.

answer = client.models.generate_content(
    model = "gemini-2.0-flash-001",
    contents = prompt
)

Markdown(answer.text)

※出力

If you want to listen to music, the first thing you need to do is find the "Music" icon on the touchscreen display in your Googlecar and simply touch it.