rag_query_qdrant.py

from typing import List

from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnablePassthrough
from langchain_ollama import ChatOllama

from modules import (
    RAGSettings,
    build_context_prompt,
    build_vector_store,
    list_collections,
)
from modules.qdrant_utils import collection_exists

SETTINGS = RAGSettings()


def format_docs(docs: List) -> str:
    return "\n\n".join(d.page_content for d in docs)


def build_chain():
    if not collection_exists(SETTINGS.collection_name, SETTINGS):
        available = ", ".join(list_collections(SETTINGS)) or "none"
        raise RuntimeError(
            f"Collection '{SETTINGS.collection_name}' not found at {SETTINGS.qdrant_url}. "
            f"Available collections: {available}."
        )

    vector_store = build_vector_store(settings=SETTINGS)
    retriever = vector_store.as_retriever(search_kwargs={"k": SETTINGS.retrieve_k})

    prompt = build_context_prompt()
    llm = ChatOllama(model=SETTINGS.gen_model, temperature=0)

    return (
        {
            "context": retriever | (lambda docs: format_docs(docs)),
            "question": RunnablePassthrough(),
        }
        | prompt
        | llm
        | StrOutputParser()
    )


def answer(question: str) -> str:
    return build_chain().invoke(question)

if __name__ == "__main__":
    print(answer("Give me a 3-bullet summary of the key ideas in my knowledge base."))