Build a RAG command line chatbot (Python)

pip install astrapy openai

import os
import sys

from astrapy import DataAPIClient
from openai import OpenAI


def main() -> None:
    endpoint = os.environ.get("API_ENDPOINT")  (1)
    application_token = os.environ.get("APPLICATION_TOKEN")
    openai_api_key = os.environ.get("OPENAI_API_KEY")

    keyspace = "default_keyspace"  (2)
    collection_name = "quickstart_collection"  (3)

    if not endpoint or not application_token or not openai_api_key:
        raise RuntimeError(
            "Environment variables API_ENDPOINT, APPLICATION_TOKEN, OPENAI_API_KEY must be defined."
        )

    # Instantiate the DataAPIClient and get a reference to your collection
    client = DataAPIClient()
    database = client.get_database(
        endpoint, token=application_token, keyspace=keyspace
    )
    collection = database.get_collection(collection_name)

    # Instantiate the OpenAI client
    openai = OpenAI(api_key=openai_api_key)

    # This list of messages will be sent to OpenAI with every query.
    # It starts with a single system prompt and grows as the chat progresses.
    messages = [
        {
            "role": "system",
            "content": "You are an AI assistant that can answer questions based on the the context you are given. Don't mention the context, just use it to inform your answers.",
        }
    ]

    # Start the chat by writing a message to the CLI
    print(
        "Greetings! I am an AI assistant that is ready to help you with your questions. "
        "You can ask me anything you like.\n"
        'If you want to exit, type ".exit".\n'
    )

    user_input = input("> ")

    # Run this loop continuously until the user inputs the exit command
    while user_input.lower() != ".exit":
        # If the user didn't input text, re-prompt them
        if user_input.strip() == "":
            user_input = input("> ")
            continue

        try:
            # Perform a vector search in your collection,
            # using the user input as the search string to vectorize.
            # Limit the search to 10 documents.
            # Use a projection to return just the $vectorize field of each document.
            cursor = collection.find(
                {},
                sort={"$vectorize": user_input},
                limit=10,
                projection={"$vectorize": 1},
            )

            # Join the $vectorize fields of the returned documents into a single string
            docs = cursor.to_list()
            context = "\n".join(
                (doc.get("$vectorize") or "") for doc in docs
            ).strip()

            # Combine the user question with the context from vector search
            rag_message = {
                "role": "user",
                "content": (
                    f"{context}\n---\n"
                    "Given the above context, answer the following question:\n"
                    f"{user_input}"
                ),
            }

            # Send the list of previous messages, plus the context augmented question
            # to OpenAI and stream the response
            stream = openai.chat.completions.create(
                model="gpt-4o-mini",
                messages=[*messages, rag_message],
                stream=True,
            )

            # Write OpenAI's response to the CLI as it comes in,
            # and also record it in a string
            message = ""
            for chunk in stream:
                delta = (
                    (chunk.choices[0].delta.content or "")
                    if chunk.choices
                    else ""
                )
                if delta:
                    sys.stdout.write(delta)
                    sys.stdout.flush()
                    message += delta

            # Record the user question, without the added context, in the list of messages
            messages.append({"role": "user", "content": user_input})

            # Record the OpenAI response in the list of messages
            messages.append({"role": "assistant", "content": message})

            # Prompt the user for their next question
            user_input = input("\n\n> ")

        except Exception as e:
            print(str(e), file=sys.stderr)
            user_input = input(
                "\nSomething went wrong, try asking again\n\n> "
            )


if __name__ == "__main__":
    try:
        main()
    except Exception as err:
        print(str(err), file=sys.stderr)
        sys.exit(1)