Update app.py

app.py

@@ -8,6 +8,39 @@ from langchain.llms.base import LLM
 from langchain.agents import initialize_agent, AgentType,load_tools
 from langchain.tools import Tool
 from langchain_experimental.tools.python.tool import PythonREPLTool
+import queue
+from typing import Any, Dict
+import gradio as gr
+from langchain.callbacks.base import BaseCallbackHandler
+
+
+tools = [add_numbers]
+
+# 2. 커스텀 콜백 핸들러
+class StreamingAgentCallbackHandler(BaseCallbackHandler):
+    def __init__(self, q: queue.Queue):
+        self.q = q
+
+    def on_agent_action(self, action, **kwargs):
+        log = f"🧠 Thought: {action.log.strip()}\n🔧 Action: {action.tool}({action.tool_input})"
+        self.q.put(log)
+
+    def on_tool_end(self, output, **kwargs):
+        self.q.put(f"📦 Observation: {output}\n")
+
+    def on_agent_finish(self, finish, **kwargs):
+        self.q.put(f"\n✅ Final Answer: {finish.return_values['output']}")
+
+    def on_llm_new_token(self, token: str, **kwargs):
+        # 선택적으로 토큰 스트리밍 출력
+        pass
+
+    def on_llm_end(self, response, **kwargs):
+        self.q.put("[END]")
+
+    def on_llm_error(self, error, **kwargs):
+        self.q.put(f"[ERROR] {str(error)}")
+
 #from langchain_community.retrievers import WikipediaRetriever
 #from langchain.tools.retriever import create_retriever_tool
 #retriever = WikipediaRetriever(lang="ko",top_k_results=10)
@@ -137,12 +170,15 @@ tools.extend([
 # ──────────────────────────────
 # ✅ Agent 초기화
 # ──────────────────────────────
+q_stream = queue.Queue()
+handler = StreamingAgentCallbackHandler(q_stream)
 agent = initialize_agent(
     tools,
     llm,
     agent_type=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
     verbose=True,
-    handle_parsing_errors=True
+    handle_parsing_errors=True,
+    callbacks=[handler]
 )
 
 # ──────────────────────────────
@@ -154,6 +190,11 @@ def chat(message, history):
     except Exception as e:
         response = f"⚠️ 오류: {e}"
     history = history + [(message, response)]
+    while True:
+        msg = q_stream.get()
+        if msg == "[END]":
+            break
+        yield msg + "\n"
     return history, history
 
 with gr.Blocks(theme=gr.themes.Soft(), title="PIXAL Assistant (LangChain + GitHub LLM)") as demo: