implement parallel agent

Author: ks6088ts

docs/references.md

@@ -12,6 +12,7 @@
 - [Command: A new tool for building multi-agent architectures in LangGraph](https://blog.langchain.com/command-a-new-tool-for-multi-agent-architectures-in-langgraph/)
 - [Combine control flow and state updates with Command](https://langchain-ai.github.io/langgraph/how-tos/graph-api/#combine-control-flow-and-state-updates-with-command)
 - [Command: a new tool for building multi-agent architectures in LangGraph](https://www.youtube.com/watch?v=6BJDKf90L9A)
+- [masamasa59/genai-agent-advanced-book > chapter6](https://github.com/masamasa59/genai-agent-advanced-book/blob/main/chapter6/arxiv_researcher/agent/paper_search_agent.py)
 
 ### Sample Codes
 

pyproject.toml

@@ -86,3 +86,4 @@ possibly-unbound-attribute = "ignore"
 unresolved-attribute = "ignore"
 invalid-argument-type = "ignore"
 invalid-type-form = "ignore"
+invalid-assignment = "ignore"

scripts/demo_agents_operator.py

@@ -4,7 +4,7 @@
 from dotenv import load_dotenv
 
 from template_langgraph.agents.demo_agents.multi_agent import app as multi_agent_app
-from template_langgraph.agents.demo_agents.parallel_processor_agent import app as parallel_agent_app
+from template_langgraph.agents.demo_agents.parallel_processor_agent.agent import app as parallel_processor_agent_app
 from template_langgraph.agents.demo_agents.weather_agent import app as weather_agent_app
 from template_langgraph.loggers import get_logger
 
@@ -74,7 +74,13 @@ def multi_agent(
 
 
 @app.command()
-def parallel_agent(
+def parallel_processor_agent(
+    goal: str = typer.Option(
+        "ソフトウェアシステム開発会社を立ち上げる戦略を立てるための情報収集をしたい",
+        "--goal",
+        "-g",
+        help="The goal to decompose into tasks",
+    ),
     verbose: bool = typer.Option(
         False,
         "--verbose",
@@ -85,19 +91,14 @@ def parallel_agent(
     if verbose:
         logger.setLevel(logging.DEBUG)
 
-    parallel_agent_app.invoke(
-        {
-            "messages": [
-                {"role": "user", "content": "Simulate multiple tasks in parallel"},
-            ],
-            "tasks": [
-                "Task 1",
-                "Task 2",
-                "Task 3",
-            ],
+    for event in parallel_processor_agent_app.stream(
+        input={
+            "goal": goal,
         },
         debug=True,
-    )
+    ):
+        logger.info("-" * 20)
+        logger.info(f"Event: {event}")
 
 
 if __name__ == "__main__":

template_langgraph/agents/demo_agents/multi_agent.py

@@ -3,7 +3,7 @@
 from langgraph.graph import END, START, MessagesState, StateGraph
 from langgraph.types import Command
 
-from template_langgraph.agents.demo_agent.weather_agent import app
+from template_langgraph.agents.demo_agents.weather_agent import app
 from template_langgraph.llms.azure_openais import AzureOpenAiWrapper
 from template_langgraph.loggers import get_logger
 

template_langgraph/agents/demo_agents/parallel_processor_agent/__init__.py

@@ -0,0 +1,39 @@
+from langchain_core.language_models.chat_models import BaseChatModel
+from langgraph.graph import StateGraph
+from langgraph.graph.state import CompiledStateGraph
+
+from template_langgraph.agents.demo_agents.parallel_processor_agent.models import (
+    ParallelProcessorAgentInputState,
+    ParallelProcessorAgentOutputState,
+    ParallelProcessorAgentState,
+)
+from template_langgraph.agents.demo_agents.parallel_processor_agent.nodes.decompose_tasks import DecomposeTasks
+from template_langgraph.agents.demo_agents.parallel_processor_agent.nodes.run_task import RunTask
+from template_langgraph.agents.demo_agents.parallel_processor_agent.nodes.summarize_results import SummarizeResults
+from template_langgraph.llms.azure_openais import AzureOpenAiWrapper
+
+
+class ParallelProcessorAgent:
+    def __init__(self, llm: BaseChatModel):
+        self.llm = llm
+        self.decompose_tasks = DecomposeTasks(llm)
+        self.run_task = RunTask()
+        self.summarize_results = SummarizeResults()
+
+    def create_graph(self) -> CompiledStateGraph:
+        workflow = StateGraph(
+            state_schema=ParallelProcessorAgentState,
+            input_schema=ParallelProcessorAgentInputState,
+            output_schema=ParallelProcessorAgentOutputState,
+        )
+        workflow.add_node("decompose_tasks", self.decompose_tasks)
+        workflow.add_node("run_task", self.run_task)
+        workflow.add_node("summarize_results", self.summarize_results)
+
+        workflow.add_edge("run_task", "summarize_results")
+        workflow.set_entry_point("decompose_tasks")
+        workflow.set_finish_point("summarize_results")
+        return workflow.compile()
+
+
+app = ParallelProcessorAgent(AzureOpenAiWrapper().chat_model).create_graph()

template_langgraph/agents/demo_agents/parallel_processor_agent/agent.py

@@ -0,0 +1,44 @@
+import operator
+from typing import (
+    Annotated,
+    TypedDict,
+)
+
+from pydantic import BaseModel, Field
+
+
+class Task(BaseModel):
+    id: str = Field(description="タスクのID")
+    description: str = Field(description="タスクの説明")
+    state: str = Field(description="タスクの状態")
+
+
+class Tasks(BaseModel):
+    tasks: list[Task] = Field(description="タスクのリスト")
+
+
+class TaskResult(TypedDict):
+    task: Task = Field(description="タスクの詳細")
+    result_code: int = Field(description="タスクの実行結果コード")
+    message: str = Field(description="タスクの実行結果メッセージ")
+
+
+class ParallelProcessorAgentInputState(TypedDict):
+    goal: str
+
+
+class ParallelProcessorAgentProcessingState(TypedDict):
+    tasks: Tasks
+
+
+class ParallelProcessorAgentOutputState(TypedDict):
+    task_results: Annotated[list[TaskResult], operator.add]
+    summary: str
+
+
+class ParallelProcessorAgentState(
+    ParallelProcessorAgentInputState,
+    ParallelProcessorAgentProcessingState,
+    ParallelProcessorAgentOutputState,
+):
+    pass

template_langgraph/agents/demo_agents/parallel_processor_agent/models.py

@@ -0,0 +1,34 @@
+from typing import Literal
+
+from langchain_core.language_models.chat_models import BaseChatModel
+from langgraph.types import Command, Send
+
+from template_langgraph.agents.demo_agents.parallel_processor_agent.models import (
+    Tasks,
+)
+
+
+class DecomposeTasks:
+    def __init__(self, llm: BaseChatModel):
+        self.llm = llm
+
+    def __call__(self, state: dict) -> Command[Literal["run_task"]]:
+        goal = state.get("goal", "")
+        tasks: Tasks = self.llm.with_structured_output(Tasks).invoke(
+            input=f"Decompose the following goal into tasks: {goal}",
+        )
+        gotos = []
+        for task in tasks.tasks:
+            gotos.append(
+                Send(
+                    "run_task",
+                    {
+                        "task": task,
+                    },
+                )
+            )
+
+        return Command(
+            goto=gotos,
+            update={"tasks": tasks},
+        )

template_langgraph/agents/demo_agents/parallel_processor_agent/nodes/__init__.py

@@ -0,0 +1,31 @@
+from template_langgraph.agents.demo_agents.parallel_processor_agent.models import (
+    Task,
+    TaskResult,
+)
+from template_langgraph.loggers import get_logger
+
+logger = get_logger(__name__)
+
+
+class RunTask:
+    def __init__(self):
+        pass
+
+    def __call__(self, state: dict) -> dict:
+        logger.info(f"Running state... {state}")
+        task: Task = state.get("task", None)
+        logger.info(f"Task: {task.model_dump_json(indent=2)}")
+
+        # FIXME: Simulate task processing for now. Replace with actual processing logic e.g. Tool call agent
+        import time
+
+        time.sleep(3)
+        result = TaskResult(
+            task=task,
+            message="Task completed successfully",
+            result_code=0,
+        )
+
+        return {
+            "task_results": [result],
+        }