main.py

import yaml
import json
from rich.console import Console
from rich.panel import Panel
from rich.markdown import Markdown
from rag.core import TinyRAG
from rag.indexing import Indexer


def print_rag_result(res):
    c = Console()

    decomp_txt = "\n".join([f"- {q}" for q in res['queries_executed']])
    router_txt = f"{res['router_decisions'][0]['logic']} | {res['router_decisions'][0]['weights']}"

    rej_txt = ""
    if res['rejected_details']:
        rej_txt = "\nOdrzucone (przykłady):\n" + "\n".join([f"  * {r['reason']} (ID: {r['id']})" for r in res['rejected_details'][:2]])

    stats = f"Input: {res['retrieval_stats']['total']} | Kept: {res['retrieval_stats']['kept']} | Rejected: {res['retrieval_stats']['rejected']}"

    meta_panel = Panel(
        f"[bold]Dekompozycja:[/bold]\n{decomp_txt}\n\n[bold]Router:[/bold] {router_txt}\n\n[bold]Filtr:[/bold] {stats}{rej_txt}\n\n[bold]Walidacja:[/bold] {res['validation']}",
        title="Meta Info", border_style="blue"
    )

    ans_panel = Panel(Markdown(res['answer']), title="Odpowiedź", border_style="green")

    c.print(meta_panel)
    c.print(ans_panel)

    if res['memory']['saved']:
        c.print(f"[bold red]ZAPISANO DO PAMIĘCI:[/bold red] {res['memory']['reason']}")


def main():
    with open("config/config.yaml") as f:
        cfg = yaml.safe_load(f)
    with open("config/prompts.yaml") as f:
        prm = yaml.safe_load(f)

    rag = TinyRAG(cfg, prm)

    while True:
        q = input("\nPYTANIE: ")
        if q.lower() in ['exit', 'quit']:
            break
        res = rag.query(q, corpus="small", mode="safe")
        print_rag_result(res)


if __name__ == "__main__":
    main()