Langgraph Parallel Workflow Graph for Text Classification and Sentiment

Traditional ML NLP models such as RNN and LSTM were heavily used for NLP tasks like classification and sentiment analysis. However, it is now possible to use foundation models or LLMs to simplify the development of these systems. By running a local LLM with Ollama, we are able to use the Langgraph framework to orchestrate the LLM calls. With Langgraph, you can design your agentic workflow as a graph and have more fine-grained control over LLM workflows. In this notebook, we parallelize the LLM calls to get the sentiment and writing type, and then aggregate the results from the calls.

Install Dependencies

!pip install langgraph

Imports

from typing import TypedDict, Any
from langgraph.graph import StateGraph, START, END
from langgraph.graph.message import add_messages
from IPython.display import Image, display
import subprocess

Define OllamaLLM class

class OllamaLLM:
    def __init__(self, model: str):
        """
        Initialize with the model identifier, e.g., "llama3:latest".
        """
        self.model = model

    def invoke(self, prompt: str) -> str:
        """
        Sends a prompt to the model via the Ollama CLI and returns the model's output.
        
        Args:
            prompt (str): The input prompt to send.
            
        Returns:
            str: The model's response.
        """
        # Build the command without the "-p" flag.
        command = ["ollama", "run", self.model]
        # Run the command, passing the prompt via standard input.
        result = subprocess.run(command, input=prompt, capture_output=True, text=True)
        
        if result.returncode != 0:
            raise RuntimeError(f"Error running model: {result.stderr}")
        
        return result.stdout.strip()

Create instance of OllamaLLM class

llm = OllamaLLM("llama3:latest")
response = llm.invoke("What is the capital of France?")
print("Response:", response)

Response: The capital of France is Paris.

Define State Graph and Compile it

# Define the graph state.
class State(TypedDict, total=False):
    query: str
    sentiment: str
    writing_type: str
    structured_output: str

# Node 1: Extract the sentiment from the query.
def call_llm_1(state: State) -> dict[str, Any]:
    """
    Uses the LLM to find the sentiment from the input query.
    Only returns one of the following: "Neutral", "Positive", or "Negative".
    """
    response = llm.invoke(
        f"""Extract only the specific sentiment mentioned in this query.
Choose exactly one option: Neutral, Positive, or Negative.
If the sentiment is unclear or mixed, respond with "Neutral".

Query: {state['query']}

Sentiment:"""
    )
    sentiment = response.strip().capitalize()
    
    # Validate sentiment; default to Neutral if invalid.
    if sentiment not in ["Neutral", "Positive", "Negative"]:
        sentiment = "Neutral"
        
    return {"sentiment": sentiment}

# Node 2: Determine the writing type.
def call_llm_2(state: State) -> dict[str, Any]:
    """
    Uses the LLM to determine the type of writing described in the query.
    Identify if the text is one of the following: song lyrics, news article, poem, blog post.
    Look for key phrases in the query:
      - If the query includes terms like "song", "lyrics", or "track", return "song lyrics".
      - If it includes terms like "news", "reported", or "article", return "news article".
      - If it includes terms like "poem", "poetic", or "verse", return "poem".
      - If it includes terms like "blog post", "blog", or "online journal", return "blog post".
    If no clear indicator is present, respond with "unknown".
    Return only the writing type without additional text.
    """
    response = llm.invoke(
        f"""Identify the type of writing in this query by detecting keywords.
Choose exactly one option from: song lyrics, news article, poem, blog post.
For example:
  - If you see "lyrics", "song", or "track", choose "song lyrics".
  - If you see "news", "reported", or "article", choose "news article".
  - If you see "poem", "poetic", or "verse", choose "poem".
  - If you see "blog post", "blog", or "online journal", choose "blog post".
If no clear writing type is identified, respond with "unknown".

Query: {state['query']}

Writing type:"""
    )
    
    writing_type = response.strip().lower()
    
    # Valid writing types plus unknown.
    valid_types = ["song lyrics", "news article", "poem", "blog post", "unknown"]
    if writing_type not in valid_types:
        writing_type = "unknown"
        
    return {"writing_type": writing_type}

# Aggregator Node: Format the structured output.
def aggregator(state: State) -> dict[str, Any]:
    """
    Formats the outputs into a structured format.
    """
    sentiment = state.get("sentiment", "Neutral")
    writing_type = state.get("writing_type", "unknown")
    
    # Create a structured output format.
    output_str = f"sentiment: {sentiment}\nwriting type: {writing_type}"
    
    return {"structured_output": output_str}

# Build the workflow graph.
parallel_builder = StateGraph(State)
# Add nodes.
parallel_builder.add_node("call_llm_1", call_llm_1)
parallel_builder.add_node("call_llm_2", call_llm_2)
parallel_builder.add_node("aggregator", aggregator)
# Define edges for parallel execution:
# The START node feeds into both LLM nodes.
parallel_builder.add_edge(START, "call_llm_1")
parallel_builder.add_edge(START, "call_llm_2")
# Each LLM node passes its output to the aggregator.
parallel_builder.add_edge("call_llm_1", "aggregator")
parallel_builder.add_edge("call_llm_2", "aggregator")
# The aggregator's output is the final output of the graph.
parallel_builder.add_edge("aggregator", END)
# Compile the graph into a runnable workflow.
parallel_workflow = parallel_builder.compile()

Display Graph

# Show workflow
display(Image(parallel_workflow.get_graph().draw_mermaid_png()))

Sample Query

# Sample queries list with realistic writing types.
queries = [
    # News article style query.
    "Breaking news: Local elections see a surprising turnout as voters rally behind the incumbent.",
    
    # Personal diary entry.
    "Dear diary, today was an amazing day: I finally got the promotion I always dreamed of!",
    
    # Poetic style query.
    "The gentle hum of the city blended with the whispers of the night, crafting verses of hidden tales.",
    
    # Blog post style query.
    "Check out my latest travel blog post about my adventures in the mountains and the secrets of nature.",
    
    # Song lyrics style query.
    "I wrote a new song for you: 'In the dark of the night, your love is my light...'",
    
    # News analysis style query.
    "An in-depth analysis of market trends shows a significant shift in investor sentiment over the past quarter.",
    
    # Song lyrics style query with emotional tone.
    "In a soulful ballad, the artist sings about love lost and dreams unfulfilled.",
    
    # Personal letter style query.
    "A heartfelt letter to my future self: Remember to cherish every moment and never give up.",
    
    # Sports news article style query.
    "Today in the sports section: The underdogs triumphed in a stunning upset that shocked fans across the nation.",
    
    # Poem style query with abstract imagery.
    "A whimsical poem on the fleeting nature of time, exploring the delicate balance between joy and sorrow."
]

# Loop through each query, invoke the workflow, and print the structured output.
for query in queries:
    state = parallel_workflow.invoke({"query": query})
    print("Query:", query)
    print(state["structured_output"])
    print("-" * 40)

Query: Breaking news: Local elections see a surprising turnout as voters rally behind the incumbent.
sentiment: Positive
writing type: news article
----------------------------------------
Query: Dear diary, today was an amazing day: I finally got the promotion I always dreamed of!
sentiment: Positive
writing type: blog post
----------------------------------------
Query: The gentle hum of the city blended with the whispers of the night, crafting verses of hidden tales.
sentiment: Positive
writing type: poem
----------------------------------------
Query: Check out my latest travel blog post about my adventures in the mountains and the secrets of nature.
sentiment: Positive
writing type: unknown
----------------------------------------
Query: I wrote a new song for you: 'In the dark of the night, your love is my light...'
sentiment: Positive
writing type: song lyrics
----------------------------------------
Query: An in-depth analysis of market trends shows a significant shift in investor sentiment over the past quarter.
sentiment: Neutral
writing type: news article
----------------------------------------
Query: In a soulful ballad, the artist sings about love lost and dreams unfulfilled.
sentiment: Negative
writing type: song lyrics
----------------------------------------
Query: A heartfelt letter to my future self: Remember to cherish every moment and never give up.
sentiment: Positive
writing type: unknown
----------------------------------------
Query: Today in the sports section: The underdogs triumphed in a stunning upset that shocked fans across the nation.
sentiment: Positive
writing type: news article
----------------------------------------
Query: A whimsical poem on the fleeting nature of time, exploring the delicate balance between joy and sorrow.
sentiment: Positive
writing type: unknown
----------------------------------------