from retrieval import retrieve_subgraph
from llm import llm_call

def query_knowledge_graph(question: str, G):

    graph_context = retrieve_subgraph(G, question)

    prompt = f"""
You are an educational assistant.

Use ONLY the graph knowledge below to answer.

Graph Knowledge:
{graph_context}

Question:
{question}

Rules:
- If graph is empty, say "No information in knowledge graph"
- Be concise and educational
"""

    response = llm_call([
        {"role": "user", "content": prompt}
    ])

    return {
        "answer": response,
        "evidence": graph_context
    }


TOOL_MAP = {
    "query_knowledge_graph": query_knowledge_graph,
}

#test
from graph_builder import load_graph

if __name__ == "__main__":

    G = load_graph()

    result = query_knowledge_graph(
        "What is used in machine learning?",
        G
    )

    print(result)