From dae1a38739f00ad0ab6bb972e91041f9d7e1de70 Mon Sep 17 00:00:00 2001
From: Shroominic <db2002dominic@gmail.com>
Date: Thu, 15 Feb 2024 14:03:36 +0800
Subject: [PATCH] =?UTF-8?q?=F0=9F=94=A7=20Update=20retrieval=5Fchain=20inp?=
 =?UTF-8?q?ut=20handling?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 examples/decorator.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/examples/decorator.py b/examples/decorator.py
index 8340b55..d148c98 100644
--- 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"}))