157 lines
4.3 KiB
Python
157 lines
4.3 KiB
Python
from fastapi import FastAPI,Header
|
|
import uvicorn
|
|
"""
|
|
Simple LangGraph Agent using @tool decorator
|
|
|
|
Clean implementation with decorator-based tool definitions.
|
|
"""
|
|
|
|
from typing import Annotated, TypedDict
|
|
from langgraph.graph import StateGraph, START, END
|
|
from langgraph.graph.message import add_messages
|
|
from langgraph.prebuilt import ToolNode, tools_condition
|
|
from langchain_core.messages import HumanMessage, SystemMessage
|
|
from langchain_core.tools import tool
|
|
from langchain_aws import ChatBedrock
|
|
app = FastAPI()
|
|
|
|
# Define tools using @tool decorator
|
|
rules="""- Mulheres acima de 45 anos ou menopausada
|
|
- Homens com mais de 70 anos;
|
|
- Osteogênese imperfeita (para esta patologia, poderá haver a liberação de (02) dois
|
|
exames ao ano - cada 180 dias);
|
|
- RX com osteopenia ou fratura patológica;
|
|
- Antecedente pessoal de fratura após os 40 anos: punho, ombros, vértebras, quadril;
|
|
- Parente de primeiro grau com osteoporose.
|
|
- Mulheres com massa corporal <20kg/m2 ou peso < 57,8kg;
|
|
- Menopausa antes dos 45 anos ou hipogonasismo crônico (falência ovariana
|
|
precoce);
|
|
- Uso de glicocorticóides (>=7,5 prednizona/ dia equivalente por mais três meses, ou
|
|
presença de síndrome de cushing;
|
|
- Hiperparatireoidismo primário;
|
|
- Uso prolongado de anticonvulsivantes (< 10 anos);
|
|
- Síndrome de má absorção crônica ou desnutrição doenças inflamatória intestinal
|
|
(independente da causa: bariatrica, celiacos, intolerancia a lactose).
|
|
- Quimioterapia, se sobrevida esperada for longa (< 5 anos);
|
|
- Diminuição documentada de altura;
|
|
- Presença de cifose após menopausa.
|
|
- Imobilização prolongada"""
|
|
SYSTEM_PROMPT="""
|
|
You are a assisant, your job is to aprove or not a procedure based on the following rules:"""
|
|
+rules+"""
|
|
Your input will be a json, evaluate it based on the rules and return:
|
|
Aproved: If one of the criteira is met
|
|
Reproved: If not a single one of the criterias are met."""
|
|
@tool
|
|
def add(a: int, b: int) -> int:
|
|
"""Add two numbers together.
|
|
|
|
Args:
|
|
a: First number
|
|
b: Second number
|
|
"""
|
|
return a + b
|
|
|
|
|
|
@tool
|
|
def multiply(a: int, b: int) -> int:
|
|
"""Multiply two numbers.
|
|
|
|
Args:
|
|
a: First number
|
|
b: Second number
|
|
"""
|
|
return a * b
|
|
|
|
|
|
@tool
|
|
def get_word_length(word: str) -> int:
|
|
"""Get the length of a word.
|
|
|
|
Args:
|
|
word: The word to measure
|
|
"""
|
|
return len(word)
|
|
|
|
|
|
@tool
|
|
def search_info(topic: str) -> str:
|
|
"""Search for information about a topic (mock implementation).
|
|
|
|
Args:
|
|
topic: The topic to search for
|
|
"""
|
|
# Mock response - replace with actual search/API
|
|
return f"Information about {topic}: This is a mock response. In production, this would return real data."
|
|
|
|
|
|
# Define agent state
|
|
class AgentState(TypedDict):
|
|
messages: Annotated[list, add_messages]
|
|
|
|
|
|
# Define tools list
|
|
tools = [add, multiply, get_word_length, search_info]
|
|
|
|
|
|
# Agent node
|
|
def call_model(state: AgentState):
|
|
"""Call the LLM with current state and tools."""
|
|
|
|
model = ChatBedrock(
|
|
model_id="arn:aws:bedrock:us-east-2:232048051668:application-inference-profile/uy4xskop19zn",
|
|
region_name="us-east-2",
|
|
provider="anthropic"
|
|
)
|
|
|
|
model_with_tools = model.bind_tools(tools)
|
|
|
|
messages = [
|
|
SystemMessage(content=SYSTEM_PROMPT)
|
|
] + state["messages"]
|
|
|
|
response = model_with_tools.invoke(messages)
|
|
|
|
return {"messages": [response]}
|
|
|
|
|
|
# Build the graph
|
|
def create_agent():
|
|
"""Create and compile the agent graph."""
|
|
|
|
workflow = StateGraph(AgentState)
|
|
|
|
# Add nodes
|
|
workflow.add_node("agent", call_model)
|
|
workflow.add_node("tools", ToolNode(tools))
|
|
|
|
# Add edges
|
|
workflow.add_edge(START, "agent")
|
|
workflow.add_conditional_edges("agent", tools_condition)
|
|
workflow.add_edge("tools", "agent")
|
|
|
|
return workflow.compile()
|
|
|
|
|
|
# Main execution
|
|
|
|
|
|
@app.post("/")
|
|
async def root(json:str= Header(...)):
|
|
agent = create_agent()
|
|
query=json
|
|
result = agent.invoke(
|
|
{"messages": [HumanMessage(content=query)]},
|
|
config={"recursion_limit": 10}
|
|
)
|
|
|
|
final_message = result["messages"][-1]
|
|
return {"status": "success", "message": final_message.content}
|
|
|
|
@app.get("/health")
|
|
async def health():
|
|
return {"status": "healthy"}
|
|
|
|
if __name__ == "__main__":
|
|
uvicorn.run(app, host="0.0.0.0", port=8000)
|