'''
Usando LangGraph para rotear usando um grafo DAG
Usa interface OpenAi com Ollama
'''
import os
from dotenv import load_dotenv
from langchain_openai import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.output_parsers import StrOutputParser
from langgraph.graph import StateGraph, START, END
from langchain_core.runnables import RunnableConfig # serve para configurar e rastrear a execução do seu fluxo em tempo de execução.
import asyncio # necessário para realizar chamadas assíncronas
from typing import Literal, TypedDict
#para garantira que o destino seja sempre uma das opções especificadas
#isso aqui poderia ser substituido por pydantic
class Rota(TypedDict):
destino: Literal["praia","montanha"]
class Estado(TypedDict):
query: str
destino: Rota
resposta: str
# Carrega as variáveis de ambiente do arquivo .env
load_dotenv()
# O Ollama utiliza a base_url para se conectar ao serviço
OLLAMA_HOST = os.getenv("OLLAMA_BASE_URL", "http://localhost:11434/v1")
OLLAMA_API_KEY = os.getenv("OLLAMA_API_KEY")
OLLAMA_MODEL_IN_USE = os.getenv("OLLAMA_MODEL_IN_USE")
# padrao OPENAI com OLLAMA
# Inicialização do modelo usando interface openAI inves de ollama
modelo = ChatOpenAI(
model=OLLAMA_MODEL_IN_USE,
openai_api_base=OLLAMA_HOST , # <--- O padrão OpenAI
openai_api_key=OLLAMA_API_KEY, # O Ollama aceita qualquer texto como chave aqui
temperature=0.5
)
# Prompt Consultor 1
prompt_consultor_praia = ChatPromptTemplate.from_messages(
[
("system", "Apresente-se como Sra. Praias. Você é uma especialista em viagens com destinos para praias."),
("human", "{query}")
]
)
# Prompt Consultor 2
prompt_consultor_montanha = ChatPromptTemplate.from_messages(
[
("system", "Apresente-se como Sra. High Mountain. Você é uma especialista em viagens destinos para montanhas e atividades radicais."),
("human", "{query}")
]
)
#configurando o prompt roteador pra decidir qual cadeia usar
#a cadeia recebe a pergunta do usuário e decidira para qual consultor irá praia ou montanha
prompt_roteador = ChatPromptTemplate.from_messages(
[
("system", "Responda apenas com 'praia' ou 'montanha'"),
("human", "{query}")
]
)
roteador = prompt_roteador | modelo.with_structured_output(Rota, method="function_calling")
#criando chains especificas para cada consultor
cadeia_praia = prompt_consultor_praia | modelo | StrOutputParser()
cadeia_montanha = prompt_consultor_montanha | modelo | StrOutputParser()
async def no_roteador(estado:Estado, config=RunnableConfig):
return {"destino": await roteador.ainvoke({"query": estado["query"]}, config)}
async def no_praia(estado: Estado, config=RunnableConfig):
return {"resposta": await cadeia_praia.ainvoke({"query": estado["query"]}, config)}
async def no_montanha(estado: Estado, config=RunnableConfig):
return {"resposta": await cadeia_montanha.ainvoke({"query": estado["query"]}, config)}
def escolher_no(estado: Estado) -> Literal["praia", "montanha"]:
return "praia" if estado["destino"]["destino"] == "praia" else "montanha"
grafo = StateGraph(Estado)
grafo.add_node("rotear", no_roteador)
grafo.add_node("praia", no_praia)
grafo.add_node("montanha", no_montanha)
grafo.add_edge(START, "rotear")
grafo.add_conditional_edges("rotear", escolher_no)
grafo.add_edge("praia", END)
grafo.add_edge("montanha", END)
app = grafo.compile()
async def main():
#resposta = await app.ainvoke({"query": "Quero visitar um lugar no Brasil famoso por praias e cultura "}, )
resposta = await app.ainvoke({"query": "Quero escalar montanhas no sul do Brasil"}, )
print(resposta["resposta"])
asyncio.run(main())
Para saber mais:
Conte com o apoio da comunidade Alura na sua jornada. Abraços e bons estudos!