wehub-resource-sync
155 lines
4.9 KiB
Python
155 lines
4.9 KiB
Python
"""Chat Node"""
|
|
|
|
from typing import List, Literal, cast
|
|
|
|
from copilotkit.langgraph import copilotkit_customize_config
|
|
from langchain.tools import tool
|
|
from langchain_core.messages import AIMessage, SystemMessage, ToolMessage
|
|
from langchain_core.runnables import RunnableConfig
|
|
from langgraph.types import Command
|
|
|
|
from src.lib.download import get_resource
|
|
from src.lib.model import get_model
|
|
from src.lib.state import AgentState
|
|
|
|
|
|
@tool
|
|
def Search(queries: List[str]): # pylint: disable=invalid-name,unused-argument
|
|
"""A list of one or more search queries to find good resources to support the research."""
|
|
|
|
|
|
@tool
|
|
def WriteReport(report: str): # pylint: disable=invalid-name,unused-argument
|
|
"""Write the research report."""
|
|
|
|
|
|
@tool
|
|
def WriteResearchQuestion(research_question: str): # pylint: disable=invalid-name,unused-argument
|
|
"""Write the research question."""
|
|
|
|
|
|
@tool
|
|
def DeleteResources(urls: List[str]): # pylint: disable=invalid-name,unused-argument
|
|
"""Delete the URLs from the resources."""
|
|
|
|
|
|
async def chat_node(
|
|
state: AgentState, config: RunnableConfig
|
|
) -> Command[Literal["search_node", "chat_node", "delete_node", "__end__"]]:
|
|
"""
|
|
Chat Node
|
|
"""
|
|
|
|
config = copilotkit_customize_config(
|
|
config,
|
|
emit_intermediate_state=[
|
|
{
|
|
"state_key": "report",
|
|
"tool": "WriteReport",
|
|
"tool_argument": "report",
|
|
},
|
|
{
|
|
"state_key": "research_question",
|
|
"tool": "WriteResearchQuestion",
|
|
"tool_argument": "research_question",
|
|
},
|
|
],
|
|
)
|
|
|
|
state["resources"] = state.get("resources", [])
|
|
research_question = state.get("research_question", "")
|
|
report = state.get("report", "")
|
|
|
|
resources = []
|
|
|
|
for resource in state["resources"]:
|
|
content = get_resource(resource["url"])
|
|
if content == "ERROR":
|
|
continue
|
|
resources.append({**resource, "content": content})
|
|
|
|
model = get_model(state)
|
|
# Prepare the kwargs for the ainvoke method
|
|
ainvoke_kwargs = {}
|
|
if model.__class__.__name__ in ["ChatOpenAI"]:
|
|
ainvoke_kwargs["parallel_tool_calls"] = False
|
|
|
|
response = await model.bind_tools(
|
|
[
|
|
Search,
|
|
WriteReport,
|
|
WriteResearchQuestion,
|
|
DeleteResources,
|
|
],
|
|
**ainvoke_kwargs, # Pass the kwargs conditionally
|
|
).ainvoke(
|
|
[
|
|
SystemMessage(
|
|
content=f"""
|
|
You are a research assistant. You help the user with writing a research report.
|
|
Do not recite the resources, instead use them to answer the user's question.
|
|
You should use the search tool to get resources before answering the user's question.
|
|
If you finished writing the report, ask the user proactively for next steps, changes etc, make it engaging.
|
|
To write the report, you should use the WriteReport tool. Never EVER respond with the report, only use the tool.
|
|
If a research question is provided, YOU MUST NOT ASK FOR IT AGAIN.
|
|
|
|
This is the research question:
|
|
{research_question}
|
|
|
|
This is the research report:
|
|
{report}
|
|
|
|
Here are the resources that you have available:
|
|
{resources}
|
|
"""
|
|
),
|
|
*state["messages"],
|
|
],
|
|
config,
|
|
)
|
|
|
|
ai_message = cast(AIMessage, response)
|
|
|
|
if ai_message.tool_calls:
|
|
if ai_message.tool_calls[0]["name"] == "WriteReport":
|
|
report = ai_message.tool_calls[0]["args"].get("report", "")
|
|
return Command(
|
|
goto="chat_node",
|
|
update={
|
|
"report": report,
|
|
"messages": [
|
|
ai_message,
|
|
ToolMessage(
|
|
tool_call_id=ai_message.tool_calls[0]["id"],
|
|
content="Report written.",
|
|
),
|
|
],
|
|
},
|
|
)
|
|
if ai_message.tool_calls[0]["name"] == "WriteResearchQuestion":
|
|
return Command(
|
|
goto="chat_node",
|
|
update={
|
|
"research_question": ai_message.tool_calls[0]["args"][
|
|
"research_question"
|
|
],
|
|
"messages": [
|
|
ai_message,
|
|
ToolMessage(
|
|
tool_call_id=ai_message.tool_calls[0]["id"],
|
|
content="Research question written.",
|
|
),
|
|
],
|
|
},
|
|
)
|
|
|
|
goto = "__end__"
|
|
if ai_message.tool_calls and ai_message.tool_calls[0]["name"] == "Search":
|
|
goto = "search_node"
|
|
elif (
|
|
ai_message.tool_calls and ai_message.tool_calls[0]["name"] == "DeleteResources"
|
|
):
|
|
goto = "delete_node"
|
|
|
|
return Command(goto=goto, update={"messages": response})
|