first-aid-rag-assistant / src / assistant.py
assistant.py
Raw
import sys

from dotenv import load_dotenv

from config import make_llm_client
from embedder import Embedder
from ingest import build_index, build_vector_index, load_faq_data
from rag.starter import RAGProduction
from db.conversations import save_conversation

litellm_client = make_llm_client()


def create_assistant():
    load_dotenv()

    documents = load_faq_data()
    text_index = build_index(documents)

    embedder = Embedder()
    vector_index = build_vector_index(embedder, documents)

    return RAGProduction(
        text_index=text_index,
        vector_index=vector_index,
        embedder=embedder,
        llm_client=litellm_client,
    )


if __name__ == "__main__":
    assistant = create_assistant()

    query = "When should you move an injured person at an accident site?"
    if len(sys.argv) > 1:
        query = sys.argv[1]

    answer = assistant.rag(query)
    print(answer)
    save_conversation(
        assistant.last_call, query,
        assistant.last_rewritten_question or query,
    )