refactor: cleaning the general structure and changing the memory of the agent

app.py

@@ -1,64 +1,57 @@
+import os
 import gradio as gr
 from langchain.agents import create_agent
 from langchain_huggingface import HuggingFaceEndpoint, ChatHuggingFace
-from langchain_core.messages import HumanMessage, SystemMessage
+from langgraph.checkpoint.memory import InMemorySaver
 
 
-def respond(
-    message,
-    history: list,
-    system_message,
-    max_tokens,
-    temperature,
-    top_p,
-    hf_token: gr.OAuthToken,
-):
-    hf_model = HuggingFaceEndpoint(
-        repo_id="meta-llama/Llama-3.1-8B-Instruct",
-        task="text-generation",
-        provider="auto",
-        max_new_tokens=max_tokens,
-        temperature=temperature,
-        top_p=top_p,
-        huggingfacehub_api_token=hf_token.token
-    )
+class GradioAgent:
 
-    llm = ChatHuggingFace(llm=hf_model)
+    def __init__(self):
+        self.agent = self.__create_agent()
 
-    agent = create_agent(
-        tools=[],
-        model=llm,
-        system_prompt=system_message
-    )
+    def inicialize(self):
+        chatbot = gr.ChatInterface(
+            self._respond,
+            type="messages"
+        )
 
-    m = HumanMessage(message)
-    saida = agent.invoke(m)['messages'][0].content
+        with gr.Blocks() as demo:
+            with gr.Sidebar():
+                gr.LoginButton()
+            chatbot.render()
 
-    yield saida
+        demo.launch()
 
+    def __create_agent(self):
+        hf_model = HuggingFaceEndpoint(
+            repo_id="meta-llama/Llama-3.1-8B-Instruct",
+            task="text-generation",
+            provider="auto",
+            huggingfacehub_api_token=os.getenv("HF_TOKEN")
+        )
+        llm = ChatHuggingFace(llm=hf_model)
 
-chatbot = gr.ChatInterface(
-    respond,
-    type="messages",
-    additional_inputs=[
-        gr.Textbox(value="You are a friendly and helpful Chatbot.", label="System message"),
-        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
-        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
-        gr.Slider(
-            minimum=0.1,
-            maximum=1.0,
-            value=0.95,
-            step=0.05,
-            label="Top-p (nucleus sampling)",
-        ),
-    ],
-)
+        return create_agent(
+            tools=[],
+            model=llm,
+            checkpointer=InMemorySaver()
+        )
 
-with gr.Blocks() as demo:
-    with gr.Sidebar():
-        gr.LoginButton()
-    chatbot.render()
+    def _respond(
+        self,
+        message,
+        history
+    ):
+        result = self.agent.invoke(
+            {"messages": [{"role": "user", "content": message}]},
+            {"configurable": {"thread_id": "1"}},
+        )
+        output = result['messages'][-1].content
+
+        yield output
 
 
 if __name__ == "__main__":
-    demo.launch()
+    gradio = GradioAgent()
+    gradio.inicialize()