| 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'] | |