from langchain_core.runnables import RunnableLambda

from bayesian_langchain import BayesianIndex, BayesianRAG, BayesianRetriever, 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_real_langchain_retriever_and_runnable():
    idx = BayesianIndex(_docs())
    retriever = BayesianRetriever(index=idx, k=1)
    out = retriever.invoke('What is the access code?')
    assert out[0].metadata['id'] == 'd1'
    assert out[0].metadata['posterior'] > 0.0
    rag = BayesianRAG(idx)
    ans = rag.as_runnable().invoke('What is the access code?')
    assert 'ORBIT-7' in ans['answer']
    assert ans['citations'] == ['d1']


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