refactor news summarizer agent

Author: ks6088ts

scripts/agent_operator.py

@@ -10,6 +10,7 @@
 from template_langgraph.agents.news_summarizer_agent.agent import (
     graph as news_summarizer_agent_graph,
 )
+from template_langgraph.agents.news_summarizer_agent.models import Article
 from template_langgraph.agents.task_decomposer_agent.agent import graph as task_decomposer_agent_graph
 from template_langgraph.loggers import get_logger
 
@@ -117,11 +118,11 @@ def run(
 
 @app.command()
 def news_summarizer_agent(
-    request: str = typer.Option(
+    prompt: str = typer.Option(
         "Please summarize the latest news articles in Japanese briefly in 3 sentences.",
-        "--request",
-        "-r",
-        help="Request to the agent",
+        "--prompt",
+        "-p",
+        help="Prompt for the agent",
     ),
     urls: str = typer.Option(
         "https://example.com/article1,https://example.com/article2",
@@ -138,7 +139,6 @@ def news_summarizer_agent(
 ):
     from template_langgraph.agents.news_summarizer_agent.models import (
         AgentInputState,
-        AgentOutputState,
         AgentState,
     )
 
@@ -150,25 +150,19 @@ def news_summarizer_agent(
     for event in graph.stream(
         input=AgentState(
             input=AgentInputState(
-                request=request,
-                request_id=str(uuid4()),
+                prompt=prompt,
+                id=str(uuid4()),
                 urls=urls.split(",") if urls else [],
             ),
-            output=AgentOutputState(
-                result="N/A",
-                articles=[],
-            ),
-            target_url_index=None,
+            articles=[],
         )
     ):
         logger.info("-" * 20)
         logger.info(f"Event: {event}")
 
-    output: AgentOutputState = event["notify"]["output"]
-    for article in output.articles:
-        logger.info(article.url)
-        logger.info(f"is_valid_url: {article.is_valid_url}, is_valid_content: {article.is_valid_content}")
-        logger.info(article.structured_article.model_dump_json(indent=2))
+    articles: list[Article] = event["notify"]["articles"]
+    for article in articles:
+        logger.info(f"{article.model_dump_json(indent=2)}")
 
 
 if __name__ == "__main__":

template_langgraph/agents/news_summarizer_agent/agent.py

@@ -2,17 +2,22 @@
 from langgraph.graph import StateGraph
 from langgraph.types import Send
 
-from template_langgraph.agents.news_summarizer_agent.models import AgentState, Article, StructuredArticle
+from template_langgraph.agents.news_summarizer_agent.models import (
+    AgentState,
+    Article,
+    StructuredArticle,
+    SummarizeWebContentState,
+)
 from template_langgraph.llms.azure_openais import AzureOpenAiWrapper
 from template_langgraph.loggers import get_logger
 
 logger = get_logger(__name__)
 
 
 class MockNotifier:
-    def notify(self, request_id: str, body: dict) -> None:
+    def notify(self, id: str, body: dict) -> None:
         """Simulate sending a notification to the user."""
-        logger.info(f"Notification sent for request {request_id}: {body}")
+        logger.info(f"Notification sent for request {id}: {body}")
 
 
 class MockScraper:
@@ -85,16 +90,19 @@ def create_graph(self):
 
         # Create nodes
         workflow.add_node("initialize", self.initialize)
-        workflow.add_node("fetch_web_content", self.fetch_web_content)
+        workflow.add_node("summarize_web_content", self.summarize_web_content)
         workflow.add_node("notify", self.notify)
 
         # Create edges
         workflow.set_entry_point("initialize")
         workflow.add_conditional_edges(
             source="initialize",
             path=self.run_subtasks,
+            path_map={
+                "summarize_web_content": "summarize_web_content",
+            },
         )
-        workflow.add_edge("fetch_web_content", "notify")
+        workflow.add_edge("summarize_web_content", "notify")
         workflow.set_finish_point("notify")
         return workflow.compile(
             name=NewsSummarizerAgent.__name__,
@@ -111,61 +119,61 @@ def run_subtasks(self, state: AgentState) -> list[Send]:
         logger.info(f"Running subtasks with state: {state}")
         return [
             Send(
-                node="fetch_web_content",
-                arg=AgentState(
-                    input=state.input,
-                    output=state.output,
-                    target_url_index=idx,
+                node="summarize_web_content",
+                arg=SummarizeWebContentState(
+                    url=state.input.urls[idx],
+                    prompt=state.input.prompt,
                 ),
             )
             for idx, _ in enumerate(state.input.urls)
         ]
 
-    def fetch_web_content(self, state: AgentState):
-        url: str = state.input.urls[state.target_url_index]
-        is_valid_url = url.startswith("http")
+    def summarize_web_content(self, state: SummarizeWebContentState):
+        is_valid_url = state.url.startswith("http")
         is_valid_content = False
         content = ""
 
         # Check if the URL is valid
         if not is_valid_url:
-            logger.error(f"Invalid URL: {url}")
+            logger.error(f"Invalid URL: {state.url}")
             is_valid_content = False
         else:
             # Scrape the web content
             try:
-                logger.info(f"Scraping URL: {url}")
-                content = self.scraper.scrape(url)
+                logger.info(f"Scraping URL: {state.url}")
+                content = self.scraper.scrape(state.url)
                 is_valid_content = True
             except httpx.RequestError as e:
                 logger.error(f"Error fetching web content: {e}")
 
         if is_valid_content:
-            logger.info(f"Summarizing content with LLM @ {state.target_url_index}: {url}")
+            logger.info(f"Summarizing content with LLM: {state.url}")
             structured_article: StructuredArticle = self.summarizer.summarize(
-                prompt=state.input.request,
+                prompt=state.prompt,
                 content=content,
             )
-            state.output.articles.append(
-                Article(
-                    is_valid_url=is_valid_url,
-                    is_valid_content=is_valid_content,
-                    content=content,
-                    url=url,
-                    structured_article=structured_article,
-                ),
-            )
+            return {
+                "articles": [
+                    Article(
+                        is_valid_url=is_valid_url,
+                        is_valid_content=is_valid_content,
+                        content=content,
+                        url=state.url,
+                        structured_article=structured_article,
+                    ),
+                ]
+            }
 
     def notify(self, state: AgentState) -> AgentState:
         """Send notifications to the user."""
         logger.info(f"Sending notifications with state: {state}")
         # Simulate sending notifications
         # convert list of articles to a dictionary for notification
         summary = {}
-        for i, article in enumerate(state.output.articles):
+        for i, article in enumerate(state.articles):
             summary[i] = article.model_dump()
         self.notifier.notify(
-            request_id=state.input.request_id,
+            id=state.input.id,
             body=summary,
         )
         return state

template_langgraph/agents/news_summarizer_agent/models.py

@@ -1,6 +1,14 @@
+import operator
+from typing import Annotated
+
 from pydantic import BaseModel, Field
 
 
+class SummarizeWebContentState(BaseModel):
+    prompt: str = Field(..., description="Prompt for summarization")
+    url: str = Field(..., description="URL of the article to summarize")
+
+
 class StructuredArticle(BaseModel):
     title: str = Field(..., description="Title of the article")
     date: str = Field(..., description="Publication date of the article")
@@ -18,16 +26,11 @@ class Article(BaseModel):
 
 
 class AgentInputState(BaseModel):
-    request: str = Field(..., description="Request from the user")
-    request_id: str = Field(..., description="Unique identifier for the request")
+    prompt: str = Field(..., description="Prompt for the agent")
+    id: str = Field(..., description="Unique identifier for the request")
     urls: list[str] = Field(..., description="List of article URLs")
 
 
-class AgentOutputState(BaseModel):
-    articles: list[Article] = Field(..., description="List of articles processed by the agent")
-
-
 class AgentState(BaseModel):
     input: AgentInputState = Field(..., description="Input state for the agent")
-    output: AgentOutputState = Field(..., description="Output state for the agent")
-    target_url_index: int | None = Field(..., description="Index of the target URL being processed")
+    articles: Annotated[list[Article], operator.add]