From 2ab1e549bf0e9e0b987acc2253447fd3a0743c68 Mon Sep 17 00:00:00 2001
From: Andrew Huang <ahuang@anaconda.com>
Date: Wed, 24 Jan 2024 21:09:34 -0800
Subject: [PATCH] Update examples

---
 docs/examples/langchain/langchain_lcel.py     | 33 ++++---------------
 .../langchain_streaming_lcel_with_memory.py   |  3 +-
 tests/ui/user.py                              |  8 +++++
 3 files changed, 16 insertions(+), 28 deletions(-)

diff --git a/docs/examples/langchain/langchain_lcel.py b/docs/examples/langchain/langchain_lcel.py
index f3fd97e..7b5e62b 100644
--- a/docs/examples/langchain/langchain_lcel.py
+++ b/docs/examples/langchain/langchain_lcel.py
@@ -3,51 +3,30 @@
 [LangChain Expression Language](https://python.langchain.com/docs/expression_language/) (LCEL).
 """
 
-from operator import itemgetter
-
 import panel as pn
-from langchain.memory import ConversationSummaryBufferMemory
 from langchain_core.output_parsers import StrOutputParser
-from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
-from langchain_core.runnables import RunnableLambda, RunnablePassthrough
+from langchain_core.prompts import ChatPromptTemplate
+from langchain_core.runnables import RunnablePassthrough
 from langchain_openai import ChatOpenAI
 
 pn.extension()
 
-SYSTEM_PROMPT = "Try to be a silly comedian."
-
 
 async def callback(contents, user, instance):
     message = ""
-    inputs = {"input": contents}
-    async for token in chain.astream(inputs):
+    async for token in chain.astream(contents):
         message += token
         yield message
-    memory.save_context(inputs, {"output": message})
 
 
+prompt = ChatPromptTemplate.from_template("Tell me a top-notch joke about {topic}")
 model = ChatOpenAI(model="gpt-3.5-turbo")
-memory = ConversationSummaryBufferMemory(return_messages=True, llm=model)
-prompt = ChatPromptTemplate.from_messages(
-    [
-        ("system", SYSTEM_PROMPT),
-        MessagesPlaceholder(variable_name="history"),
-        ("human", "{input}"),
-    ]
-)
 output_parser = StrOutputParser()
-chain = (
-    RunnablePassthrough.assign(
-        history=RunnableLambda(memory.load_memory_variables) | itemgetter("history")
-    )
-    | prompt
-    | model
-    | output_parser
-)
+chain = {"topic": RunnablePassthrough()} | prompt | model | output_parser
 
 chat_interface = pn.chat.ChatInterface(
     pn.chat.ChatMessage(
-        "Offer a topic and ChatGPT will try to be funny!", user="System"
+        "Offer a topic and ChatGPT will respond with a joke!", user="System"
     ),
     callback=callback,
     callback_user="ChatGPT",
diff --git a/docs/examples/langchain/langchain_streaming_lcel_with_memory.py b/docs/examples/langchain/langchain_streaming_lcel_with_memory.py
index f3fd97e..abafea5 100644
--- a/docs/examples/langchain/langchain_streaming_lcel_with_memory.py
+++ b/docs/examples/langchain/langchain_streaming_lcel_with_memory.py
@@ -1,6 +1,7 @@
 """
 Demonstrates how to use the `ChatInterface` to create a chatbot using
-[LangChain Expression Language](https://python.langchain.com/docs/expression_language/) (LCEL).
+[LangChain Expression Language](https://python.langchain.com/docs/expression_language/) (LCEL)
+with streaming and memory.
 """
 
 from operator import itemgetter
diff --git a/tests/ui/user.py b/tests/ui/user.py
index 5fef18c..b64d0e0 100644
--- a/tests/ui/user.py
+++ b/tests/ui/user.py
@@ -175,6 +175,14 @@ def langchain_lcel(page: Page):
     page.wait_for_timeout(5000)
 
 
+def langchain_streaming_lcel_with_memory(page: Page):
+    chat = ChatInterface(page)
+    chat.send("Remember this number: 8. Be concise.")
+    page.wait_for_timeout(10000)
+    chat.send("What number did I just ask you to remember?")
+    page.wait_for_timeout(10000)
+
+
 def mistral_and_llama(page: Page):
     chat = ChatInterface(page)
     chat.send("What do you think about HoloViz in a single sentence?")