This commit is contained in:
William Fu-Hinthorn
2024-02-12 15:15:37 -08:00
parent 3ba3fdd855
commit 8823b78dc1
3 changed files with 59 additions and 56 deletions
+7 -10
View File
@@ -107,10 +107,7 @@
"from langchain_openai import ChatOpenAI\n",
"\n",
"\n",
"\n",
"def create_agent(\n",
" llm: ChatOpenAI, tools: list, system_prompt: str\n",
"):\n",
"def create_agent(llm: ChatOpenAI, tools: list, system_prompt: str):\n",
" # Each worker node will be given a name and some tools.\n",
" prompt = ChatPromptTemplate.from_messages(\n",
" [\n",
@@ -255,7 +252,11 @@
"research_node = functools.partial(agent_node, agent=research_agent, name=\"Researcher\")\n",
"\n",
"# NOTE: THIS PERFORMS ARBITRARY CODE EXECUTION. PROCEED WITH CAUTION\n",
"code_agent = create_agent(llm, [python_repl_tool], \"You may generate safe python code to analyze data and generate charts using matplotlib.\")\n",
"code_agent = create_agent(\n",
" llm,\n",
" [python_repl_tool],\n",
" \"You may generate safe python code to analyze data and generate charts using matplotlib.\",\n",
")\n",
"code_node = functools.partial(agent_node, agent=code_agent, name=\"Coder\")\n",
"\n",
"workflow = StateGraph(AgentState)\n",
@@ -369,11 +370,7 @@
],
"source": [
"for s in graph.stream(\n",
" {\n",
" \"messages\": [\n",
" HumanMessage(content=\"Write a brief research report on pikas.\")\n",
" ]\n",
" },\n",
" {\"messages\": [HumanMessage(content=\"Write a brief research report on pikas.\")]},\n",
" {\"recursion_limit\": 100},\n",
"):\n",
" if \"__end__\" not in s:\n",
@@ -291,9 +291,7 @@
" return {\"messages\": [HumanMessage(content=result[\"output\"], name=name)]}\n",
"\n",
"\n",
"def create_team_supervisor(\n",
" llm: ChatOpenAI, system_prompt, members\n",
") -> str:\n",
"def create_team_supervisor(llm: ChatOpenAI, system_prompt, members) -> str:\n",
" \"\"\"An LLM-based router.\"\"\"\n",
" options = [\"FINISH\"] + members\n",
" function_def = {\n",
@@ -374,10 +372,18 @@
"\n",
"llm = ChatOpenAI(model=\"gpt-4-1106-preview\")\n",
"\n",
"search_agent = create_agent(llm, [tavily_tool], \"You are a research assistant who can search for up-to-date info using the tavily search engine.\")\n",
"search_agent = create_agent(\n",
" llm,\n",
" [tavily_tool],\n",
" \"You are a research assistant who can search for up-to-date info using the tavily search engine.\",\n",
")\n",
"search_node = functools.partial(agent_node, agent=search_agent, name=\"Search\")\n",
"\n",
"research_agent = create_agent(llm, [scrape_webpages], \"You are a research assistant who can scrape specified urls for more detailed information using the scrape_webpages function.\")\n",
"research_agent = create_agent(\n",
" llm,\n",
" [scrape_webpages],\n",
" \"You are a research assistant who can scrape specified urls for more detailed information using the scrape_webpages function.\",\n",
")\n",
"research_node = functools.partial(agent_node, agent=research_agent, name=\"Web Scraper\")\n",
"\n",
"supervisor_agent = create_team_supervisor(\n",
@@ -388,7 +394,7 @@
" \" task and respond with their results and status. When finished,\"\n",
" \" respond with FINISH.\",\n",
" [\"Search\", \"Web Scraper\"],\n",
")\n"
")"
]
},
{
@@ -417,17 +423,14 @@
"research_graph.add_conditional_edges(\n",
" \"supervisor\",\n",
" lambda x: x[\"next\"],\n",
" {\n",
" \"Search\": \"Search\",\n",
" \"Web Scraper\": \"Web Scraper\",\n",
" \"FINISH\": END\n",
" }\n",
" {\"Search\": \"Search\", \"Web Scraper\": \"Web Scraper\", \"FINISH\": END},\n",
")\n",
"\n",
"\n",
"research_graph.set_entry_point(\"supervisor\")\n",
"chain = research_graph.compile()\n",
"\n",
"\n",
"# The following functions interoperate between the top level graph state\n",
"# and the state of the research sub-graph\n",
"# this makes it so that the states of each graph don't get intermixed\n",
@@ -438,11 +441,7 @@
" return results\n",
"\n",
"\n",
"\n",
"research_chain = (\n",
" enter_chain\n",
" | chain\n",
")"
"research_chain = enter_chain | chain"
]
},
{
@@ -474,9 +473,8 @@
],
"source": [
"for s in research_chain.stream(\n",
" \"when is Taylor Swift's next tour?\",\n",
" {\"recursion_limit\": 100}\n",
" ):\n",
" \"when is Taylor Swift's next tour?\", {\"recursion_limit\": 100}\n",
"):\n",
" if \"__end__\" not in s:\n",
" print(s)\n",
" print(\"---\")"
@@ -539,7 +537,6 @@
" }\n",
"\n",
"\n",
"\n",
"llm = ChatOpenAI(model=\"gpt-4-1106-preview\")\n",
"\n",
"doc_writer_agent = create_agent(\n",
@@ -551,7 +548,9 @@
")\n",
"# Injects current directory working state before each call\n",
"context_aware_doc_writer_agent = prelude | doc_writer_agent\n",
"doc_writing_node = functools.partial(agent_node, agent=context_aware_doc_writer_agent, name=\"Doc Writer\")\n",
"doc_writing_node = functools.partial(\n",
" agent_node, agent=context_aware_doc_writer_agent, name=\"Doc Writer\"\n",
")\n",
"\n",
"note_taking_agent = create_agent(\n",
" llm,\n",
@@ -560,7 +559,9 @@
" \" taking notes to craft a perfect paper.{current_files}\",\n",
")\n",
"context_aware_note_taking_agent = prelude | note_taking_agent\n",
"note_taking_node = functools.partial(agent_node, agent=context_aware_note_taking_agent, name=\"Note Taker\")\n",
"note_taking_node = functools.partial(\n",
" agent_node, agent=context_aware_note_taking_agent, name=\"Note Taker\"\n",
")\n",
"\n",
"chart_generating_agent = create_agent(\n",
" llm,\n",
@@ -569,7 +570,9 @@
" \"{current_files}\",\n",
")\n",
"context_aware_chart_generating_agent = prelude | chart_generating_agent\n",
"chart_generating_node = functools.partial(agent_node, agent=context_aware_note_taking_agent, name=\"Chart Generator\")\n",
"chart_generating_node = functools.partial(\n",
" agent_node, agent=context_aware_note_taking_agent, name=\"Chart Generator\"\n",
")\n",
"\n",
"doc_writing_supervisor = create_team_supervisor(\n",
" llm,\n",
@@ -578,7 +581,7 @@
" \" respond with the worker to act next. Each worker will perform a\"\n",
" \" task and respond with their results and status. When finished,\"\n",
" \" respond with FINISH.\",\n",
" [\"Doc Writer\", \"Note Taker\", \"Chart Generator\"]\n",
" [\"Doc Writer\", \"Note Taker\", \"Chart Generator\"],\n",
")"
]
},
@@ -618,20 +621,21 @@
" \"Doc Writer\": \"Doc Writer\",\n",
" \"Note Taker\": \"Note Taker\",\n",
" \"Chart Generator\": \"Chart Generator\",\n",
" \"FINISH\": END\n",
" }\n",
" \"FINISH\": END,\n",
" },\n",
")\n",
"\n",
"authoring_graph.set_entry_point(\"supervisor\")\n",
"chain = research_graph.compile()\n",
"\n",
"\n",
"# The following functions interoperate between the top level graph state\n",
"# and the state of the research sub-graph\n",
"# this makes it so that the states of each graph don't get intermixed\n",
"def enter_chain(message: str, members: List[str]):\n",
" results = {\n",
" \"messages\": [HumanMessage(content=message)],\n",
" \"team_members\": \", \".join(members)\n",
" \"team_members\": \", \".join(members),\n",
" }\n",
" return results\n",
"\n",
@@ -664,9 +668,9 @@
],
"source": [
"for s in authoring_chain.stream(\n",
" \"Write an outline for poem and then write the poem to disk.\",\n",
" {\"recursion_limit\": 100}\n",
" ):\n",
" \"Write an outline for poem and then write the poem to disk.\",\n",
" {\"recursion_limit\": 100},\n",
"):\n",
" if \"__end__\" not in s:\n",
" print(s)\n",
" print(\"---\")"
@@ -720,6 +724,7 @@
" messages: Annotated[List[BaseMessage], operator.add]\n",
" next: str\n",
"\n",
"\n",
"def get_last_message(state: State) -> str:\n",
" return state[\"messages\"][-1].content\n",
"\n",
@@ -727,6 +732,7 @@
"def join_graph(response: dict):\n",
" return {\"messages\": [response[\"messages\"][-1]]}\n",
"\n",
"\n",
"# Define the graph.\n",
"super_graph = StateGraph(State)\n",
"# First add the nodes, which will do the work\n",
@@ -746,8 +752,8 @@
" {\n",
" \"Paper writing team\": \"Paper writing team\",\n",
" \"Research team\": \"Research team\",\n",
" \"FINISH\": END\n",
" }\n",
" \"FINISH\": END,\n",
" },\n",
")\n",
"super_graph.set_entry_point(\"supervisor\")\n",
"super_graph = super_graph.compile()"
@@ -814,13 +820,15 @@
],
"source": [
"for s in super_graph.stream(\n",
" {\n",
" \"messages\": [\n",
" HumanMessage(content=\"Write a brief research report on the North American sturgeon. Include a chart.\")\n",
" ],\n",
" },\n",
" {\"recursion_limit\": 150},\n",
" ):\n",
" {\n",
" \"messages\": [\n",
" HumanMessage(\n",
" content=\"Write a brief research report on the North American sturgeon. Include a chart.\"\n",
" )\n",
" ],\n",
" },\n",
" {\"recursion_limit\": 150},\n",
"):\n",
" if \"__end__\" not in s:\n",
" print(s)\n",
" print(\"---\")"
@@ -118,9 +118,7 @@
" )\n",
" prompt = prompt.partial(system_message=system_message)\n",
" prompt = prompt.partial(tool_names=\", \".join([tool.name for tool in tools]))\n",
" return prompt | llm.bind_functions(functions)\n",
"\n",
"\n"
" return prompt | llm.bind_functions(functions)"
]
},
{
@@ -162,8 +160,7 @@
" result = repl.run(code)\n",
" except BaseException as e:\n",
" return f\"Failed to execute. Error: {repr(e)}\"\n",
" return f\"Succesfully executed:\\n```python\\n{code}\\n```\\nStdout: {result}\"\n",
"\n"
" return f\"Succesfully executed:\\n```python\\n{code}\\n```\\nStdout: {result}\""
]
},
{
@@ -251,8 +248,8 @@
"\n",
"# Research agent and node\n",
"research_agent = create_agent(\n",
" llm, \n",
" [tavily_tool], \n",
" llm,\n",
" [tavily_tool],\n",
" system_message=\"You should provide accurate data for the chart generator to use.\",\n",
")\n",
"research_node = functools.partial(agent_node, agent=research_agent, name=\"Researcher\")\n",
@@ -286,6 +283,7 @@
"tools = [tavily_tool, python_repl]\n",
"tool_executor = ToolExecutor(tools)\n",
"\n",
"\n",
"def tool_node(state):\n",
" \"\"\"This runs tools in the graph\n",
"\n",