chore: import upstream snapshot with attribution

2026-07-13 12:58:18 +08:00
commit 6d5d58c1a9
18293 changed files with 3502153 additions and 0 deletions
@@ -0,0 +1,154 @@
+"""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})