adaptive_flow_as_graph.py

import os
import pathlib
from typing import TypedDict
from langgraph.graph import StateGraph, END
from beir.datasets.data_loader import GenericDataLoader
import chromadb
import time
from pathlib import Path

# jesli amy rozbicie na kilka podpytan, to context pobieramy tylko z vectoStoreSeach i web,
# nastepnie ten kontekst zwracamy i generujemy finalna odpwoedz

from modul_decyzujny.decompose_query import check_if_query_is_complex
from modul_decyzujny.judge_response import evaluate_answer_quality
from modul_decyzujny.knowledge_summarizer import (
    summarize_knowledge_with_most_occurring_words,
    summarize_knowledge_with_random_documents,
)
from modul_decyzujny.first_router import (
    query_router,
    ToolChoice,
    get_context_to_question,
)
from modul_decyzujny.decompose_query import decompose_query
from sentence_transformers import SentenceTransformer

from tools.final_llm_reponse import final_answer_from_llm


def load_questions_txt(path: str) -> dict:
    result_dict = {}
    text = Path(path).read_text(
        encoding="utf-8"
    )  # jeśli plik może mieć BOM: "utf-8-sig"
    lines = (ln.strip() for ln in text.splitlines())
    questions = [ln for ln in lines if ln and not ln.lstrip().startswith("#")]
    number = 1
    for q in questions:
        result_dict[number] = q
        number += 1
    return result_dict


dataset = "fiqa"
data_path = os.path.join(pathlib.Path(__file__).parent.absolute(), "datasets", dataset)
corpus, queries, qrels = GenericDataLoader(data_folder=data_path).load(split="test")
# docs_summary = Path("/Users/michalzabinski/Desktop/base/magisterka_adaptive_rag/all-MiniLM-L6-vs/dosc_summary/most_occuring_40_summary_2000_tokens.txt").read_text(encoding="utf-8")
docs_summary = summarize_knowledge_with_most_occurring_words(corpus, top_k=5)
# docs_summary = summarize_knowledge_with_random_documents(corpus, 80)
# first_ten_to_vdb = list(queries.values())[:10]
queries = load_questions_txt(
    f"questions_expected_to_be_routed_to_web_search_{dataset}.txt"
)
# queries = load_questions_txt(f"questions_expected_to_be_routed_to_llm_{dataset}.txt")
embedding_model_name = "all-MiniLM-L6-v2"

# model dla samego llm
llm_fallback_model = "meta-llama/Llama-3.3-70B-Instruct-Turbo"

total_chars = sum(len(doc["text"]) for doc in corpus.values())

print(f"dataset length: {total_chars}")

embedding_model = SentenceTransformer(embedding_model_name)
collection_name = f"beir-{dataset}-corpus"
client = chromadb.PersistentClient(path="chroma_store")
collection_to_search = client.get_collection(name=collection_name)

routes_to_vector_search = 0
routes_to_web_search = 0
routes_to_llm = 0

decompose_query_enabled = False
enable_final_judge = False


class State(TypedDict):
    original_question: str
    subquestions: list[str]
    route: list[str]
    contexts: list[str]
    final_answer: str


def decompose_node(state: State) -> dict:
    question = state["original_question"]
    print(f"\n Analyzing question: {question}")

    if decompose_query_enabled:
        is_complex = check_if_query_is_complex(question)
    else:
        is_complex = False

    if is_complex:
        subqs = decompose_query(query=question)
        print("🔍 Decomposed into subquestions:")
        for sq in subqs:
            print(f" - {sq}")
    else:
        subqs = [question]
        # print("✅ Simple question. No decomposition needed.")

    return {**state, "subquestions": subqs}


def route_node(state: State) -> dict:

    routes = []

    for subq in state["subquestions"]:
        print(f"\n Routing for question: {subq}")

        route = query_router(query_text=subq, docs_summary=docs_summary)
        # route = "DirectLLM"

        # route = "VectorStoreSearch"

        print(f"➡️ Routed to: {route}")

        if route == "VectorStoreSearch":
            global routes_to_vector_search
            routes_to_vector_search += 1

        if route == "WebSearch":
            global routes_to_web_search
            routes_to_web_search += 1

        if route == "DirectLLM":
            global routes_to_llm
            routes_to_llm += 1

        routes.append(route)

    return {**state, "route": routes}


def fetch_context_node(state: State) -> dict:
    print("\n Fetching context for question (excluding DirectLLM)...")
    contexts = []

    all_routes = state["route"]
    all_subqs = state["subquestions"]

    # tutaj pobieramy tylko kontekst dla pytan, ktore sa skierowane do vector albo web search
    for subq, route in zip(all_subqs, all_routes):

        # tu mozna brute force przekierowac do VectrStore
        # route = ToolChoice.VECTORSTORE

        if route in [
            ToolChoice.VECTORSTORE,
            ToolChoice.WEB_SEARCH,
            ToolChoice.EXTENDED_WEB_SEARCH,
        ]:
            print(f" - [{route}] {subq}")
            context = get_context_to_question(
                route=route,
                query_text=subq,
                embedding_model=embedding_model,
                llm_fallback_model=llm_fallback_model,
                collection_to_search=collection_to_search,
            )
            contexts.append(context)
        else:
            print(f" - [DirectLLM] {subq} -> skipping context fetch")

    # jesli nie pobrano żadnego kontekstu, a pytanie było pojedyncze i skierowane do DirectLLM — użyj LLM bez retrievalu
    # jesli nie pobrano zadnego kontekstu, oznacza to, ze orginalne pytanie w całosci nadaje sie do odpowiedzenia przez llma
    if not contexts and len(all_subqs) == 1 and all_routes[0] == ToolChoice.DIRECT_LLM:
        print("⚠️ No context needed. Falling back to DirectLLM response.")
        context = get_context_to_question(
            route=ToolChoice.DIRECT_LLM,
            query_text=state["original_question"],
            embedding_model=embedding_model,
            llm_fallback_model=llm_fallback_model,
            collection_to_search=collection_to_search,
        )
        contexts = [context]

    return {**state, "contexts": contexts}


# spobowac, dograc to jeszczze tak aby, z funckji ktore pobieraja kontekst,
# byl zwracany tylko kontekst, a tylko do generowania funal answer,
# była generowana od z wykorzystaniem llma
def final_answer_node(state: State) -> dict:
    all_routes = state["route"]
    all_subqs = state["subquestions"]
    contexts = state["contexts"]

    if not contexts and len(all_subqs) == 1 and all_routes[0] == ToolChoice.DIRECT_LLM:
        # if all_routes[0] == ToolChoice.DIRECT_LLM:
        # Tu mamy już odpowiedź z LLM – nie odpalamy jeszcze raz final_answer_from_llm
        print("✅ Using DirectLLM response as final answer (already provided).")
        return {**state, "final_answer": contexts[0]}

    # W pozostałych przypadkach (zebrany kontekst) – normalne generowanie odpowiedzi
    final_answer = final_answer_from_llm(
        context_blocks=contexts, original_question=state["original_question"]
    )

    print(final_answer)

    return {**state, "final_answer": final_answer}


def judge_node(state: State) -> dict:
    answer = state["final_answer"]
    question = state["original_question"]

    if enable_final_judge:
        is_good = evaluate_answer_quality(question=question, answer=answer)
    else:
        is_good = True

    if is_good:
        print("Final answer considered relevant")
        return state
    else:
        print(
            "Final answer considered irrelevant. Using EXTENDED_WEB_SEARCH fallback..."
        )
        fallback_context = get_context_to_question(
            route=ToolChoice.EXTENDED_WEB_SEARCH,
            query_text=question,
            embedding_model=embedding_model,
            llm_fallback_model=llm_fallback_model,
            collection_to_search=collection_to_search,
        )

        fallback_answer = final_answer_from_llm(
            context_blocks=[fallback_context], original_question=question
        )

        return {**state, "final_answer": fallback_answer}


def end_node(state: dict) -> dict:
    print(f"Final route decisions: {state['route']}")
    print(f"Final response: {state['final_answer']}")

    return state


builder = StateGraph(state_schema=State)
builder.add_node("decompose", decompose_node)
builder.add_node("routing_node", route_node)
builder.add_node("fetch_context", fetch_context_node)
builder.add_node("generate_answer", final_answer_node)
# builder.add_node("judge", judge_node)
# builder.add_node("end", end_node)


builder.set_entry_point("decompose")
builder.add_edge("decompose", "routing_node")
builder.add_edge("routing_node", "fetch_context")

# builder.add_edge("routing_node", "end")
builder.add_edge("fetch_context", "generate_answer")
# builder.add_edge("generate_answer", "end")

# builder.add_edge("generate_answer", "judge")

graph = builder.compile()

number_of_questions_to_process = 0

start = time.time()
for qid, query_text in queries.items():

    if number_of_questions_to_process == 10:
        break

    # query_text = "What bitcoin is ?" #and what is its current price?"

    initial_state = {"original_question": query_text}
    result = graph.invoke(initial_state)

    # print("\n Final Answer: ")
    # print(result["final_answer"])

    number_of_questions_to_process += 1

print(f"Number of questions routed to db_search: {routes_to_vector_search}")
print(f"Number of questions routed to web_search: {routes_to_web_search}")
print(f"Number of questions be responded only with lmm: {routes_to_llm}")

end = time.time()  # czas końcowy
print(f"Czas wykonania: {end - start:.4f} s")