yoon-gu commited on
Commit
9e93462
ยท
verified ยท
1 Parent(s): 964ccd2

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +150 -0
app.py ADDED
@@ -0,0 +1,150 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ import time
3
+ from gradio import ChatMessage
4
+ from langchain_core.runnables import RunnableConfig
5
+ from langchain_teddynote.messages import random_uuid
6
+ from langchain_core.messages import BaseMessage, HumanMessage
7
+ from pprint import pprint
8
+
9
+ def format_namespace(namespace):
10
+ return namespace[-1].split(":")[0] if len(namespace) > 0 else "root graph"
11
+
12
+ from langchain_openai import ChatOpenAI
13
+ from langgraph.checkpoint.memory import MemorySaver
14
+ from langgraph_supervisor import create_supervisor
15
+ from langgraph.prebuilt import create_react_agent
16
+ from langgraph.checkpoint.memory import MemorySaver, InMemorySaver
17
+ from langgraph.store.memory import InMemoryStore
18
+
19
+ checkpointer = InMemorySaver()
20
+ store = InMemoryStore()
21
+
22
+ model = ChatOpenAI(model="gpt-4o")
23
+
24
+ # Create specialized agents
25
+
26
+ def add(a: float, b: float) -> float:
27
+ """Add two numbers."""
28
+ return a + b
29
+
30
+ def multiply(a: float, b: float) -> float:
31
+ """Multiply two numbers."""
32
+ return a * b
33
+
34
+ def web_search(query: str) -> str:
35
+ """Search the web for information."""
36
+ return (
37
+ "Here are the headcounts for each of the FAANG companies in 2024:\n"
38
+ "1. **Facebook (Meta)**: 67,317 employees.\n"
39
+ "2. **Apple**: 164,000 employees.\n"
40
+ "3. **Amazon**: 1,551,000 employees.\n"
41
+ "4. **Netflix**: 14,000 employees.\n"
42
+ "5. **Google (Alphabet)**: 181,269 employees."
43
+ )
44
+
45
+ math_agent = create_react_agent(
46
+ model=model,
47
+ tools=[add, multiply],
48
+ name="math_expert",
49
+ prompt="You are a math expert. Always use one tool at a time."
50
+ )
51
+
52
+ research_agent = create_react_agent(
53
+ model=model,
54
+ tools=[web_search],
55
+ name="research_expert",
56
+ prompt="You are a world class researcher with access to web search. Do not do any math."
57
+ )
58
+
59
+ # Create supervisor workflow
60
+ workflow = create_supervisor(
61
+ [research_agent, math_agent],
62
+ model=model,
63
+ prompt=(
64
+ "You are a team supervisor managing a research expert and a math expert. "
65
+ "For current events, use research_agent. "
66
+ "For math problems, use math_agent."
67
+ )
68
+ )
69
+
70
+ # Compile and run
71
+ app = workflow.compile()
72
+
73
+ def generate_response(message, history):
74
+ inputs = {
75
+ "messages": [HumanMessage(content=message)],
76
+ }
77
+ node_names = []
78
+ response = []
79
+ for namespace, chunk in app.stream(
80
+ inputs,
81
+ stream_mode="updates", subgraphs=True
82
+ ):
83
+ for node_name, node_chunk in chunk.items():
84
+ # node_names๊ฐ€ ๋น„์–ด์žˆ์ง€ ์•Š์€ ๊ฒฝ์šฐ์—๋งŒ ํ•„ํ„ฐ๋ง
85
+ if len(node_names) > 0 and node_name not in node_names:
86
+ continue
87
+
88
+ if len(response) > 0:
89
+ response[-1].metadata["status"] = "done"
90
+ # print("\n" + "=" * 50)
91
+ msg = []
92
+ formatted_namespace = format_namespace(namespace)
93
+ if formatted_namespace == "root graph":
94
+ print(f"๐Ÿ”„ Node: \033[1;36m{node_name}\033[0m ๐Ÿ”„")
95
+ meta_title = f"๐Ÿค” `{node_name}`"
96
+ else:
97
+ print(
98
+ f"๐Ÿ”„ Node: \033[1;36m{node_name}\033[0m in [\033[1;33m{formatted_namespace}\033[0m] ๐Ÿ”„"
99
+ )
100
+ meta_title = f"๐Ÿค” `{node_name}` in `{formatted_namespace}`"
101
+
102
+ response.append(ChatMessage(content="", metadata={"title": meta_title, "status": "pending"}))
103
+ yield response
104
+ print("- " * 25)
105
+
106
+ # ๋…ธ๋“œ์˜ ์ฒญํฌ ๋ฐ์ดํ„ฐ ์ถœ๋ ฅ
107
+ out_str = []
108
+ if isinstance(node_chunk, dict):
109
+ for k, v in node_chunk.items():
110
+ if isinstance(v, BaseMessage):
111
+ v.pretty_print()
112
+ out_str.append(v.pretty_repr())
113
+ elif isinstance(v, list):
114
+ for list_item in v:
115
+ if isinstance(list_item, BaseMessage):
116
+ list_item.pretty_print()
117
+ out_str.append(list_item.pretty_repr())
118
+ else:
119
+ out_str.append(list_item)
120
+ print(list_item)
121
+ elif isinstance(v, dict):
122
+ for node_chunk_key, node_chunk_value in node_chunk.items():
123
+ out_str.append(f"{node_chunk_key}:\n{node_chunk_value}")
124
+ print(f"{node_chunk_key}:\n{node_chunk_value}")
125
+ else:
126
+ out_str.append(f"{k}:\n{v}")
127
+ print(f"\033[1;32m{k}\033[0m:\n{v}")
128
+ response[-1].content = "\n".join(out_str)
129
+ yield response
130
+ else:
131
+ if node_chunk is not None:
132
+ for item in node_chunk:
133
+ out_str.append(item)
134
+ print(item)
135
+ response[-1].content = "\n".join(out_str)
136
+ yield response
137
+ yield response
138
+ print("=" * 50)
139
+ response[-1].metadata["status"] = "done"
140
+ response.append(ChatMessage(content=node_chunk['messages'][-1].content))
141
+ yield response
142
+ demo = gr.ChatInterface(
143
+ generate_response,
144
+ type="messages",
145
+ title="Nested Thoughts Chat Interface",
146
+ examples=["2024๋…„์˜ the combined headcount of the FAANG companies์ˆ˜์น˜์— ๋Œ€ํ•œ ๋ถ„์„์„ ํ•œ๊ตญ์–ด๋กœ ๋ถ€ํƒํ•ด!"]
147
+ )
148
+
149
+ if __name__ == "__main__":
150
+ demo.launch()