LangGraph Agent para criar Histórias
O LangGraph é framework projetado para construir e visualizar modelos de linguagem complexos através de grafos estruturados. Com o LangGraph, é possível criar agentes LLM sofisticados que conectam por meio de arestas, melhorando a capacidade dos agentes de processar e responder a diversos inputs. Neste post quero explorar alguns fundamentos do LangGraph, suas componentes e aplicações.
Escritor de Histórias
Decidi construir um agent capaz de construir uma história. Mais do que isso, esse agent irá interagir com outros para que essa história não seja óbvia. Meu objetivo é que um nó do grafo seja responsável por gerar uma situação sem saída para os personagens, enquanto outro nó é responsável por encontrar uma saída para essa situação. Para começar, eu defino uma classe Story que irá armazenar os detalhes da história. Vamos começar somente com um input do usuário e uma introdução a história. Salvei essa classe em "scr/classes/story.py"
from langchain_core.pydantic_v1 import BaseModel, Field
from typing import List
import operator
from typing import Annotated, List
class Story(BaseModel):
"""Story"""
user_input: str = Field(
..., description="""
User Input
"""
)
introduction: str = Field(
..., description="""
Story Introduction
"""
)Também defino uma classe que irá conter as ferramentas para a etapa de introdução da história em "scr/tools/introduction.py". Essa classe instancia um extrator para essa etapa baseado no LLM fornecido. Esse extrator será uma chain que fornece o prompt criado em "_get_prompt" para o agent. A função "get_step" é responsável por invocar esse extrator quando esse nó for acionado no grafo. Ela irá passar a instrução obtida em "_get_prompt_text".
from langchain_core.runnables import RunnablePassthrough
from langchain_core.prompts import PromptTemplate
from langchain_core.messages.ai import AIMessage
from src.classes.story import Story
class Tools:
def __init__(self, llm):
self.extractor = self._get_extractor(llm)
def _get_prompt_text(self):
prompt = """
Create an engaging story. Write only an introduction to the main characters. Use a simple language. Write a maximum of 150 characters.
"""
return prompt
def _get_prompt(self):
prompt_template = """
Execute the task. Answer in Portuguese.
Task: {task}
{format_instructions}
"""
prompt = PromptTemplate(
template=prompt_template, input_variables=["task"], partial_variables={
"format_instructions": str
},
)
return prompt
def _get_extractor(self, llm):
chain = {"task": RunnablePassthrough()} | self._get_prompt() | llm
return chain
def get_step(self, state: Story, config):
"""get step"""
task_formatted = f"""
{self._get_prompt_text()}
Story Theme: {state.user_input}
"""
agent_response = self.extractor.invoke(task_formatted)
return {"introduction": agent_response.content}
Precisamos definir também qual LLM iremos utilizar. No meu caso, utilizei o GPT da OpenAI, mas você pode optar por alguma outra LLM.
from langchain_openai import ChatOpenAI
llm = ChatOpenAI(
model="gpt-4-0125-preview",
temperature=0.6,
)
Em posse dessas duas classes, podemos construir o workflow com esse agent único. Note que precisamos definir o workflow com a classe desejada, no caso Story. Também devemos inserir o nó de introduction, definir ele como entry_point do app e criar uma aresta entre ele e o fim do workflow.
from langgraph.graph import StateGraph
from src.tools import introduction
from src.classes import story
introduction = introduction.Tools(llm)
workflow = StateGraph(story.Story)
workflow.add_node("introduction_writter", introduction.get_step)
workflow.set_entry_point("introduction_writter")
workflow.add_edge("introduction_writter", "__end__")
app = workflow.compile()
Podemos invocar esse agent com um tema para a história e observar a introdução construída.
input = {"introduction": "", "user_input": ""}
story_theme = "Uma história sobre um cientista de dados escrevendo em um blog"
input["user_input"] = story_theme
config = {"recursion_limit": 20, "configurable": {"thread_id": 42}}
events = []
async for event in app.astream(input, config=config):
display(event)
events.append(event)João, um cientista de dados curioso, decide compartilhar suas descobertas em um blog. Ele nunca imaginou as aventuras que suas palavras criariam.
Climax da História
Vou inserir um novo nó responsável por escrever um climax como continuação da história atual. Para isso, vou atualizar a classe Story para introduzir um novo campo "story_pieces" que irá armazenar as outras peças da história. Esse campo será uma lista de strings pois permitirá que a história seja criada em partes. Além disso, é um Annotated Field para permitir que os nós anexem dados ao campo, e não que sobreescrevam.
from langchain_core.pydantic_v1 import BaseModel, Field
from typing import List
import operator
from typing import Annotated, List
class Story(BaseModel):
"""Story"""
user_input: str = Field(
..., description="""
User Input
"""
)
introduction: str = Field(
..., description="""
Story Introduction
"""
)
story_pieces: Annotated[List[str], operator.add] = Field(
..., description="""
Pieces of Story
"""
)De forma similar a classe "scr/tools/introduction.py", vou definir a classe "scr/tools/climax.py". Ela terá o mesmo papel, mas dessa vez para o nó de climax. Note que a função "get_step" é ligeiramente diferente. Dessa vez, ela une todas as partes da história para enviar ao agent, instruindo ele a continuar a história e gerar um climax em que os personagens se encontrem em uma situação sem saída. Como "story_pieces" é uma lista anotada, essa função deve retornar uma lista que será concatenada ao campo existente na história.
from langchain_core.runnables import RunnablePassthrough
from langchain_core.runnables import RunnablePassthrough
from langchain_core.prompts import PromptTemplate
from langchain_core.messages.ai import AIMessage
from src.classes.story import Story
class Tools:
def __init__(self, llm):
self.extractor = self._get_extractor(llm)
def _get_prompt_text(self):
prompt = """
Write a climax for the current story in which the characters are in an absolutely no-win situation. Use a simple language. Write a maximum of 150 characters.
Your story must be a creative continuation of the last fact presented and it must make sense with reality, not using magical or mysterious artifacts.
"""
return prompt
def _get_prompt(self):
prompt_template = """
Execute the task. Answer in Portuguese.
Task: {task}
{format_instructions}
"""
prompt = PromptTemplate(
template=prompt_template, input_variables=["task"], partial_variables={
"format_instructions": str
},
)
return prompt
def _get_extractor(self, llm):
chain = {"task": RunnablePassthrough()} | self._get_prompt() | llm
return chain
def get_step(self, state: Story, config):
"""get step"""
current_story = "\n".join([state.introduction] + state.story_pieces)
task_formatted = f"""
{self._get_prompt_text()}
Current Story:
{current_story}
Last Fact: {'' if len(state.story_pieces) == 0 else state.story_pieces[-1]}
"""
agent_response = self.extractor.invoke(task_formatted)
return {"story_pieces": [agent_response.content]}
Podemos atualizar o workflow inserindo esse novo nó. Introduzimos o novo nó, atualizando as arestas de forma a conectar o "introduction_writter" com o "climax_writter". A partir de agora, "climax_writter" levará ao fim do workflow.
from src.tools import introduction
from src.tools import climax
from src.classes import story
from langgraph.graph import StateGraph
introduction = introduction.Tools(llm)
climax = climax.Tools(llm)
workflow = StateGraph(story.Story)
workflow.add_node("introduction_writter", introduction.get_step)
workflow.add_node("climax_writter", climax.get_step)
workflow.set_entry_point("introduction_writter")
workflow.add_edge("introduction_writter", "climax_writter")
workflow.add_edge("climax_writter", "__end__")
app = workflow.compile()
Atualizamos a função que invoca o app para inserir nos parâmetros o novo campo e gerar a história novamente.
input = {"introduction": "", "user_input": "", "story_pieces": []}
story_theme = "Uma história sobre um cientista de dados escrevendo em um blog"
input["user_input"] = story_theme
config = {"recursion_limit": 20, "configurable": {"thread_id": 42}}
events = []
async for event in app.astream(input, config=config):
display(event)
events.append(event),
João, um cientista de dados curioso, decide compartilhar suas descobertas em um blog. Ele nunca imaginou as aventuras que suas palavras criariam.
Ao revelar um segredo governamental sem querer, João é perseguido. Sem saída, ele se esconde, sabendo que sua vida normal acabou para sempre.
Um objeto inusitado
Eu quero que o agent seja capaz de resolver essa situação de uma forma inusitada. Para isso, construi uma etapa intermediária que irá sugerir um objeto inusitado. Esse objeto será usado para resolver essa situação nas próximas etapas. Por enquanto, vou apenas incluir esse novo agent.
O processo será exatamente o mesmo: atualizar a classe Story, criar uma classe Tools para esse novo agent em "scr/tools/object.py" e atualizar o workflow para inserir esse novo nó. Inseri um novo campo "objects" em Story similar ao campo "story_pieces", ou seja, uma lista anotada de strings.
from langchain_core.pydantic_v1 import BaseModel, Field
from typing import List
import operator
from typing import Annotated, List
class Story(BaseModel):
"""Story"""
user_input: str = Field(
..., description="""
User Input
"""
)
introduction: str = Field(
..., description="""
Story Introduction
"""
)
story_pieces: Annotated[List[str], operator.add] = Field(
..., description="""
Pieces of Story
"""
)
objects: Annotated[List[str], operator.add] = Field(
..., description="""
List of objects
"""
)A classe Tools de object será definida em "scr/tools/object.py".
from langchain_core.runnables import RunnablePassthrough
from langchain_core.prompts import PromptTemplate
from langchain_core.messages.ai import AIMessage
from src.classes.story import Story
class Tools:
def __init__(self, llm):
self.extractor = self._get_extractor(llm)
def _get_prompt_text(self):
prompt = """
Imagine a completely random object with some unusual caracteristic. Write a maximum of 100 characters. This object must be completely different from any of these current characters.
"""
return prompt
def _get_prompt(self):
prompt_template = """
Execute the task. Answer in Portuguese.
Task: {task}
{format_instructions}
"""
prompt = PromptTemplate(
template=prompt_template, input_variables=["task"], partial_variables={
"format_instructions": str
},
)
return prompt
def _get_extractor(self, llm):
chain = {"task": RunnablePassthrough()} | self._get_prompt() | llm
return chain
def get_step(self, state: Story, config):
"""get step"""
task_formatted = f"""
{self._get_prompt_text()}
Current Objects: {state.objects}
"""
agent_response = self.extractor.invoke(task_formatted)
return {"objects": [agent_response.content]}
Podemos atualizar o workflow novamente para inserir esse novo nó.
from src.tools import introduction
from src.tools import climax
from src.tools import object
from src.classes import story
from langgraph.graph import StateGraph
introduction = introduction.Tools(llm)
climax = climax.Tools(llm)
object = object.Tools(llm)
workflow = StateGraph(story.Story)
workflow.add_node("introduction_writter", introduction.get_step)
workflow.add_node("climax_writter", climax.get_step)
workflow.add_node("character_generator", object.get_step)
workflow.set_entry_point("introduction_writter")
workflow.add_edge("introduction_writter", "climax_writter")
workflow.add_edge("climax_writter", "character_generator")
workflow.add_edge("character_generator", "__end__")
app = workflow.compile()
Atualizo novamente o input para inserir esse campo e rodar o workflow.
input = {"introduction": "", "user_input": "", "story_pieces": [], "objects": []}
story_theme = "Uma história sobre um cientista de dados escrevendo em um blog"
input["user_input"] = story_theme
config = {"recursion_limit": 20, "configurable": {"thread_id": 42}}
events = []
async for event in app.astream(input, config=config):
display(event)
events.append(event)João, um cientista de dados curioso, decide compartilhar suas descobertas em um blog. Ele nunca imaginou as aventuras que suas palavras criariam.
Ao revelar um segredo governamental sem querer, João é perseguido. Sem saída, ele se esconde, sabendo que sua vida normal acabou para sempre.
Guarda-chuva que prediz o tempo.
Resolvendo a Situação
Se você chegou até aqui, já entendeu a mecânica do processo que estamos fazendo até agora:
- Atualizar a classe Story do workflow
- Definir uma classe Tools com os prompts e extratores para aquele nó
- Atualizar o workflow para incluir o nó
- Atualizar o schema do input do workflow
Por isso, vou resumir os próximos passos e apontar somente os novos elementos. Vou criar um agent capaz de dar sequência a história e resolver a situação com o objeto recém sugerido. Para isso, não será necessário atualizar a classe Story.
Vou apenas definir a classe Tools em "scr/tools/solver.py".
from langchain_core.runnables import RunnablePassthrough
from langchain_core.prompts import PromptTemplate
from langchain_core.messages.ai import AIMessage
from src.classes.story import Story
class Tools:
def __init__(self, llm):
self.extractor = self._get_extractor(llm)
def _get_prompt_text(self):
prompt = """
Write a solution for the current story in which the situation were solved beacause an character found an object which helped to solve the situation. Use a simple language. Write a maximum of 150 characters.
You must introduce this new object to the story in a creative way and it must make sense with reality, not using magical or mysterious artifacts.
"""
return prompt
def _get_prompt(self):
prompt_template = """
Execute the task. Answer in Portuguese.
Task: {task}
{format_instructions}
"""
prompt = PromptTemplate(
template=prompt_template, input_variables=["task"], partial_variables={
"format_instructions": str
},
)
return prompt
def _get_extractor(self, llm):
chain = {"task": RunnablePassthrough()} | self._get_prompt() | llm
return chain
def get_step(self, state: Story, config):
"""get step"""
current_story = "\n".join([state.introduction] + state.story_pieces)
task_formatted = f"""
{self._get_prompt_text()}
Current Story:
{current_story}
Situation Solver Object:
{state.objects[-1]}
"""
agent_response = self.extractor.invoke(task_formatted)
return {"story_pieces": [agent_response.content]}
Em seguida incluo ele no workflow e executo novamente, da mesma forma como fizemos para todos os outros nós até aqui.
João, um cientista de dados curioso, decide compartilhar suas descobertas em um blog. Ele nunca imaginou as aventuras que suas palavras criariam.
Ao revelar um segredo governamental sem querer, João é perseguido. Sem saída, ele se esconde, sabendo que sua vida normal acabou para sempre.
Guarda-chuva que prediz o tempo.
João achou um guarda-chuva que prediz o tempo. Ele o usou para evitar câmeras de segurança em dias de chuva, passando despercebido e fugindo dos perseguidores.
Inserindo um LOOP
Poderíamos simplesmente encerrar o workflow por aqui. Mas note que não fizemos nada além de encadear diversos agents até aqui. Isso é longe do potencial do LangGraph. Vamos explorar alguns conceitos mais complexos.
O LangGraph permite inserir arestas condicionais que nos permite definir com base em algum critério qual será o próximo nó do workflow. Para incrementar nossa história, eu vou inserir uma aresta condicional capaz de enviar a saída de "situation_solver" de volta para "climax_writter". Ou seja, farei o agent escrever um novo clímax após a solução do anterior.
Basta definir uma função "should_end" que será responsável por analisar o estado atual e decidir o próximo nó. No meu caso, eu defini que se a história tiver menos de 6 "partes", ela retornará ao climax. Caso contrário, irá para o "END" e finalizará o fluxo. Esse LOOP irá permitir que a história seja desenvolvida com vários climax e soluções.
def should_end(state: story.Story) -> Literal["climax_writter", "__end__"]:
print(len(state.story_pieces))
if len(state.story_pieces) > 6:
return "__end__"
else:
return "climax_writter"
workflow.add_conditional_edges(
"situation_solver",
should_end,
)
João, um cientista de dados curioso, decide compartilhar suas descobertas em um blog. Ele nunca imaginou as aventuras que suas palavras criariam.
Ao revelar um segredo governamental sem querer, João é perseguido. Sem saída, ele se esconde, sabendo que sua vida normal acabou para sempre.
Guarda-chuva que prediz o tempo.
João achou um guarda-chuva que prediz o tempo. Ele o usou para evitar câmeras de segurança em dias de chuva, passando despercebido e fugindo dos perseguidores.
Numa virada inesperada, o guarda-chuva falha justo quando João mais precisa. Cercado, sem saída, ele encara seus perseguidores, sabendo que não há fuga.
Caneta que escreve sozinha ideias geniais.
João encontra uma caneta que escreve ideias geniais. Com ela, elabora um plano de fuga, deixando os perseguidores confusos enquanto escapa.
Cercado em um beco sem saída, a caneta de João para de funcionar. Sem ideias, ele se entrega, aceitando o fim da sua jornada.
Relógio que atrasa para prolongar momentos felizes.
João encontra um relógio que atrasa o tempo em momentos críticos. Ele o usa para atrasar os perseguidores, ganhando tempo para escapar e se esconder.
Em sua última fuga, o relógio de João quebra. Sem como atrasar o tempo, ele é cercado e capturado, sem esperança de escape ou salvação.'
Chaveiro flutuante que encontra chaves perdidas.
João achou um chaveiro flutuante que encontra chaves perdidas. Usou-o para achar a chave de um esconderijo secreto, onde ficou seguro e retomou sua vida.
Concluindo a História
Para finalizar a história, vou incluir um nó responsável por gerar uma conclusão que pode ser feliz ou triste, explorando a criação de outro tipo de aresta condicional. Atualizo a classe "Story" para incluir um campo de conclusão.
conclusion: str = Field(
..., description="""
Conclusion
"""
)Dessa vez, irei gerar duas classes Tools, uma em "scr/tools/happy_end.py" e outra em "scr/tools/sad_end.py". Ou seja, dois agents diferentes, cada um responsável por uma conclusão. A classe Tools de "scr/tools/happy_end.py" será assim. A "sad_end" será equivalente, mas com o prompt ajustado para "sad".
from langchain_core.runnables import RunnablePassthrough
from langchain_core.prompts import PromptTemplate
from langchain_core.messages.ai import AIMessage
from src.classes.story import Story
class Tools:
def __init__(self, llm):
self.extractor = self._get_extractor(llm)
def _get_prompt_text(self):
prompt = """
Write a happy conclusion for the current story. Use a simple language. Write a maximum of 150 characters.
"""
return prompt
def _get_prompt(self):
prompt_template = """
Execute the task. Answer in Portuguese.
Task: {task}
{format_instructions}
"""
prompt = PromptTemplate(
template=prompt_template, input_variables=["task"], partial_variables={
"format_instructions": str
},
)
return prompt
def _get_extractor(self, llm):
chain = {"task": RunnablePassthrough()} | self._get_prompt() | llm
return chain
def get_step(self, state: Story, config):
"""get step"""
current_story = "\n".join([state.introduction] + state.story_pieces)
task_formatted = f"""
{self._get_prompt_text()}
Current Story:
{current_story}
"""
agent_response = self.extractor.invoke(task_formatted)
return {"conclusion": agent_response.content}
No workflow, atualizo a função "should_end" para escolher aleatoriamente qual nó será responsável pela conclusão. De forma que o resultado final será aleatório.
def should_end(
state: story.Story,
) -> Literal["climax_writter", "happy_finisher", "sad_finisher"]:
if len(state.story_pieces) > 6:
if random.random() < 0.5:
return "happy_finisher"
return "sad_finisher"
else:
return "climax_writter"Além disso, precisamos inserir uma aresta que leva ambos os novos nós ao fim do workflow.
workflow.add_edge("happy_finisher", "__end__")
workflow.add_edge("sad_finisher", "__end__")Finalmente, o workflow completo e a história até a conclusão.
João, um cientista de dados curioso, decide compartilhar suas descobertas em um blog. Ele nunca imaginou as aventuras que suas palavras criariam.
Ao revelar um segredo governamental sem querer, João é perseguido. Sem saída, ele se esconde, sabendo que sua vida normal acabou para sempre.
Guarda-chuva que prediz o tempo.
João achou um guarda-chuva que prediz o tempo. Ele o usou para evitar câmeras de segurança em dias de chuva, passando despercebido e fugindo dos perseguidores.
Numa virada inesperada, o guarda-chuva falha justo quando João mais precisa. Cercado, sem saída, ele encara seus perseguidores, sabendo que não há fuga.
Caneta que escreve sozinha ideias geniais.
João encontra uma caneta que escreve ideias geniais. Com ela, elabora um plano de fuga, deixando os perseguidores confusos enquanto escapa.
Cercado em um beco sem saída, a caneta de João para de funcionar. Sem ideias, ele se entrega, aceitando o fim da sua jornada.
Relógio que atrasa para prolongar momentos felizes.
João encontra um relógio que atrasa o tempo em momentos críticos. Ele o usa para atrasar os perseguidores, ganhando tempo para escapar e se esconder.
Em sua última fuga, o relógio de João quebra. Sem como atrasar o tempo, ele é cercado e capturado, sem esperança de escape ou salvação.'
Chaveiro flutuante que encontra chaves perdidas.
João achou um chaveiro flutuante que encontra chaves perdidas. Usou-o para achar a chave de um esconderijo secreto, onde ficou seguro e retomou sua vida.
João, agora seguro, usa o esconderijo para seguir compartilhando suas descobertas. Ele inspira muitos, vivendo feliz e em paz.
Conclusão
Esse workflow foi bastante simples no sentido de tarefas executadas por cada nó, uma vez que meu objetivo era explorar o conceito do LangGraph. Entretanto, é possível construir agents com tarefas complexas em cada nó capazes de conversarem entre si para resolver um problema maior. É possível, por exemplo, construir agents capazes de acessar bancos de dados, realizar chamadas de APIs, ler documentos e fotos, etc.
Espero que você tenha curtido esse post. Se tiver interesse no código completo que mostrei aqui basta se inscrever no blog que te envio por email 😃
Abraço e até a próxima!