query_data.py

import argparse
from langchain_ollama import OllamaEmbeddings
from langchain_chroma import Chroma
from langchain_core.prompts import ChatPromptTemplate
from langchain_ollama import ChatOllama
from ai_setup import initialize_llm, GeminiEmbeddings
from config import Config

CHROMA_PATH = Config.CHROMA_PATH
LLM_MODEL_LOCAL = Config.LLM_MODEL_LOCAL
EMBEDDING_MODEL_LOCAL=Config.EMBEDDING_MODEL_LOCAL
ANSWER_PROMPT_TEMPLATE = Config.ANSWER_PROMPT_TEMPLATE
LLM_MODEL_API = Config.LLM_MODEL_API

def main():
    # Create CLI
    parser = argparse.ArgumentParser()
    parser.add_argument("query_text", type=str, help="The query text.")
    args = parser.parse_args()
    query_text = args.query_text
    query_rag(query_text)


def query_rag(query_text: str):
    # Prepare the DB
    #embeddings_function = OllamaEmbeddings(model=EMBEDDING_MODEL_LOCAL)
    embeddings_function = GeminiEmbeddings()
    db = Chroma(persist_directory=CHROMA_PATH, embedding_function=embeddings_function)

    # Search the DB
    results = db.similarity_search_with_relevance_scores(query_text, k=5)
    if len(results) == 0 or results[0][1] < 0.5:
        print(f"Unable to find matching results.")
        return

    context_text = "\n\n---\n\n".join([doc.page_content for doc, _score in results])
    prompt_template = ChatPromptTemplate.from_template(ANSWER_PROMPT_TEMPLATE)
    prompt = prompt_template.format(context=context_text, question=query_text)
    print(prompt)

    print("Retrieved similarity scores:\n")
    for i, (doc, score) in enumerate(results, start=1):
        print(
            f"[{i}]: Score: {score:.4f} Source: {doc.metadata.get("source", "unknown")}\n"
        )

    #model = ChatOllama(model=LLM_MODEL_LOCAL)
    #response_text = model.invoke(prompt).content
    llm = initialize_llm()
    response = llm.models.generate_content(
        model=LLM_MODEL_API,
        contents=prompt
    )

    sources = [(doc.metadata.get("source", None), _score) for doc, _score in results]
    formatted_response = f"Response: {response.text}\nSources: {sources}"
    print(formatted_response)


def query_rag_test(query_text: str):
    # Prepare the DB
    embeddings_function = OllamaEmbeddings(model="snowflake-arctic-embed")
    db = Chroma(persist_directory=CHROMA_PATH, embedding_function=embeddings_function)

    # Search the DB
    results = db.similarity_search_with_relevance_scores(query_text, k=5)
    if len(results) == 0 or results[0][1] < 0.1:
        print(f"Unable to find matching results.")
        return

    context_text = "\n\n---\n\n".join([doc.page_content for doc, _score in results])
    prompt_template = ChatPromptTemplate.from_template(ANSWER_PROMPT_TEMPLATE)
    prompt = prompt_template.format(context=context_text, question=query_text)
    #print(prompt)

    model = ChatOllama(model="llama3:8b")
    response_text = model.invoke(prompt).content

    sources = [(doc.metadata.get("source", None), _score) for doc, _score in results]
    formatted_response = f"Response: {response_text}\nSources: {sources}"
    #print(formatted_response)
    return  formatted_response


if __name__ == "__main__":
    main()