Update app.py

app.py

@@ -11,44 +11,43 @@ def load_data():
     """Load the data you want to use for the agent."""
     return pd.read_parquet("data/part.17.parquet")
 
-def load_agent(trans: pd.DataFrame) -> create_pandas_dataframe_agent:
-    """Logic for loading the agent we want to use."""
+def load_chain():
+    """Logic for loading the chain you want to use should go here."""
     llm = Anthropic(temperature=0)
-    trans = load_data()
-    return create_pandas_dataframe_agent(llm, 
-                         trans, 
-                         verbose=True)
+    chain = create_pandas_dataframe_agent(llm=llm, load_data())
+    return chain
 
 def set_anthropic_api_key(api_key: str):
-    """Set the api key and return agent.
+    """Set the api key and return chain.
+
     If no api_key, then None is returned.
     """
     if api_key:
         os.environ["ANTHROPIC_API_KEY"] = api_key
-        agent = load_agent()
+        chain = load_chain()
         os.environ["ANTHROPIC_API_KEY"] = ""
-        return agent
-
+        return chain
 
 class ChatWrapper:
 
     def __init__(self):
         self.lock = Lock()
     def __call__(
-        self, api_key: str, inp: str, history: Optional[Tuple[str, str]], agent: Optional[create_pandas_dataframe_agent]
+        self, api_key: str, inp: str, history: Optional[Tuple[str, str]], chain: Optional[ConversationChain]
     ):
         """Execute the chat functionality."""
         self.lock.acquire()
         try:
             history = history or []
             # If chain is None, that is because no API key was provided.
-            if agent is None:
+            if chain is None:
                 history.append((inp, "Please paste your anthropic key to use"))
                 return history, history
-            # Set OpenAI key
-            agent = set_anthropic_api_key(history)
+            # Set anthropic key
+            import anthropic
+            anthropic.api_key = api_key
             # Run chain and append input.
-            output = agent.run(input=inp)
+            output = chain.run(input=inp)
             history.append((inp, output))
         except Exception as e:
             raise e
@@ -58,7 +57,6 @@ class ChatWrapper:
 
 chat = ChatWrapper()
 
-
 block = gr.Blocks(css=".gradio-container {background-color: lightgray}")
 
 with block:
@@ -66,7 +64,7 @@ with block:
         gr.Markdown("<h3><center>LangChain Demo</center></h3>")
 
         anthropic_api_key_textbox = gr.Textbox(
-            placeholder="Paste your Anthropic API key (sk-...)",
+            placeholder="Paste your anthropic API key (sk-...)",
             show_label=False,
             lines=1,
             type="password",
@@ -91,7 +89,7 @@ with block:
         inputs=message,
     )
 
-    gr.HTML("Demo application of a LangChain agent.")
+    gr.HTML("Demo application of a LangChain chain.")
 
     gr.HTML(
         "<center>Powered by <a href='https://github.com/hwchase17/langchain'>LangChain 🦜️🔗</a></center>"
@@ -100,8 +98,8 @@ with block:
     state = gr.State()
     agent_state = gr.State()
 
-    submit.click(chat, inputs=[message, state, agent_state], outputs=[chatbot, state])
-    message.submit(chat, inputs=[message, state, agent_state], outputs=[chatbot, state])
+    submit.click(chat, inputs=[anthropic_api_key_textbox, message, state, agent_state], outputs=[chatbot, state])
+    message.submit(chat, inputs=[anthropic_api_key_textbox, message, state, agent_state], outputs=[chatbot, state])
 
     anthropic_api_key_textbox.change(
         set_anthropic_api_key,
@@ -109,4 +107,4 @@ with block:
         outputs=[agent_state],
     )
 
-block.launch(debug=True)
+block.launch(debug=True)