Parallelization_Workflow / parallelization_workflow_pattern.py
alvap's picture
Update parallelization_workflow_pattern.py
66602e4 verified
Raw
History Blame Contribute Delete
1.99 kB
from langgraph.graph import StateGraph,START,END
from typing import TypedDict
from langchain_openai import ChatOpenAI
import os
OPENAI_API_KEY=os.environ.get("OPENAI_API_KEY")
llm=ChatOpenAI(
model="gpt-4o-mini",
api_key=OPENAI_API_KEY
)
class State(TypedDict):
text:str
french:str
spanish:str
japanese:str
combined_output:str
def translate_french(state:State)->dict:
response=llm.invoke(f"Translate the following to French:\n\n{state['text']}")
return {"french":response.content.strip()}
def translate_spanish(state:State)->dict:
response=llm.invoke(f"Translate the following to Spanish:\n\n{state['text']}")
return {"spanish":response.content.strip()}
def translate_japanese(state:State)->dict:
response=llm.invoke(f"Translate the following to Japanese:\n\n{state['text']}")
return {"japanese":response.content.strip()}
def aggregator(state:State)->dict:
combined=f"Original Text: {state['text']}\n\n"
combined+=f"French: {state['french']}\n\n"
combined+=f"Spanish: {state['spanish']}\n\n"
combined+=f"Japanese: {state['japanese']}\n"
return {"combined_output":combined}
graph=StateGraph(State)
graph.add_node("translate_french",translate_french)
graph.add_node("translate_spanish",translate_spanish)
graph.add_node("translate_japanese",translate_japanese)
graph.add_node("aggregator",aggregator)
graph.add_edge(START,"translate_french")
graph.add_edge(START,"translate_spanish")
graph.add_edge(START,"translate_japanese")
graph.add_edge("translate_french","aggregator")
graph.add_edge("translate_spanish","aggregator")
graph.add_edge("translate_japanese","aggregator")
graph.add_edge("aggregator",END)
app=graph.compile()
def run_workflow(text):
state:State={
"text":text,
"french":"",
"spanish":"",
"japanese":"",
"combined_output":""
}
response=app.invoke(state)
workflow_graph=app.get_graph().draw_mermaid_png()
return response,workflow_graph