from neo4j_config import URI, USER, PASSWORD, AUTH
from retriever import *
from dspy_wrapper import *
from neo4j import GraphDatabase
import os, pickle


# == Fast Load of Precomputed Index ==
HERE = os.path.dirname(__file__)
with open(os.path.join(HERE, "chunk_index_bertino_final.pkl"), "rb") as f:
    all_chunks = pickle.load(f)  # already contain embeddings and ids


# == Neo4j Setup ==
with GraphDatabase.driver(URI, auth=AUTH) as driver:
    driver.verify_connectivity()


# == Retrieval ==
retriever = HybridRetriever(all_chunks)

reranker = GraphReranker(
    retriever,
    neo4j_uri=URI,
    neo4j_user=USER,
    neo4j_pass=PASSWORD,
    beta=0.2,
    max_hops=3
)

# == Pipeline Initialization ==
retriever_module = DSPyHybridRetriever(reranker)
cot_module = dspy.ChainOfThought(AnswerWithEvidence)
rag_pipeline = RAGChain(retriever=retriever_module, answerer=cot_module)