🔧 Update retrieval_chain input handling

shroominic · Feb 15, 2024 · dae1a38 · dae1a38
1 parent 175e1b2
commit dae1a38
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/examples/decorator.py b/examples/decorator.py
@@ -1,3 +1,5 @@
+from operator import itemgetter
+
 from funcchain.syntax import chain, runnable
 from langchain_community.vectorstores.faiss import FAISS
 from langchain_core.runnables import Runnable, RunnablePassthrough
@@ -23,8 +25,8 @@ def generate_poem(topic: str, context: str) -> str:
 retriever = vectorstore.as_retriever(search_kwargs={"k": 1})
 
 retrieval_chain: Runnable = {
-    "context": retriever,
+    "context": itemgetter("topic") | retriever,
     "topic": RunnablePassthrough(),
 } | generate_poem
 
-print(retrieval_chain.invoke("love"))
+print(retrieval_chain.invoke({"topic": "love"}))