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