curiouscurrent commited on
Commit
8e72609
·
verified ·
1 Parent(s): 104f157

Create retrieval_chain.py

Browse files
Files changed (1) hide show
  1. AI_Agent/chains/retrieval_chain.py +9 -0
AI_Agent/chains/retrieval_chain.py ADDED
@@ -0,0 +1,9 @@
 
 
 
 
 
 
 
 
 
 
1
+ # AI_Agent/chains/retrieval_chain.py
2
+ class RetrievalChain:
3
+ def __init__(self, vector_store):
4
+ self.vs = vector_store
5
+
6
+ async def run(self, prompt):
7
+ results = self.vs.search(prompt, k=3)
8
+ contexts = [r[0] for r in results]
9
+ return {"contexts": contexts, "raw": results}