from langchain_core.runnables import RunnableLambda

from bayesian_langgraph import BayesianIndex, BayesianRAGGraph, DocumentRecord


def _docs():
    return [
        DocumentRecord(
            'd1',
            'The access code is ORBIT-7.',
            'Access',
            'official',
            {'reliability_alpha':10,'reliability_beta':1,'authority':0.99,'freshness':1.0},
        )
    ]


def test_actual_stategraph_executes_rag():
    rag = BayesianRAGGraph(BayesianIndex(_docs()))
    out = rag.invoke('What is the access code?')
    assert 'ORBIT-7' in out['answer']
    assert out['citations'] == ['d1']


def test_stategraph_model_runnable_generation_branch():
    generator = RunnableLambda(lambda prompt: 'The access code is ORBIT-7. [d1]')
    rag = BayesianRAGGraph(BayesianIndex(_docs()), generator=generator, k=1)
    out = rag.invoke('What is the access code?')
    assert out['answer'] == 'The access code is ORBIT-7. [d1]'
    assert out['citations'] == ['d1']
    assert not out['abstained']