LangGraph Integration

from langgraph.graph import Graph, END
from langchain_core.messages import BaseMessage, HumanMessage, AIMessage
from langchain_openai import ChatOpenAI
from stackone_ai import StackOneToolSet
from typing import TypedDict, List

# Initialize tools and LLM
toolset = StackOneToolSet()
tools = toolset.get_tools('hris_*', account_id='your-account-id')
llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)

class AgentState(TypedDict):
    messages: List[BaseMessage]
    next_action: str

def agent_node(state: AgentState):
    """Analyze user request and determine next action"""
    messages = state["messages"]
    last_message = messages[-1]
    
    response = llm.invoke([
        HumanMessage(content="You are an HR assistant. Determine the appropriate action."),
        last_message
    ])
    
    # Route based on content
    if "employee" in response.content.lower():
        next_action = "get_employees"
    elif "department" in response.content.lower():
        next_action = "get_departments"
    else:
        next_action = "respond"
    
    return {
        "messages": messages + [response],
        "next_action": next_action
    }

def get_employees_node(state: AgentState):
    """Execute employee data retrieval"""
    employee_tool = tools.get_tool('hris_list_employees')
    result = employee_tool.call(limit=10)
    
    response = AIMessage(content=f"Found {len(result['data'])} employees")
    return {
        "messages": state["messages"] + [response],
        "next_action": "respond"
    }

def respond_node(state: AgentState):
    """Final response node"""
    return {
        "messages": state["messages"],
        "next_action": END
    }

# Build and compile the workflow
workflow = Graph()
workflow.add_node("agent", agent_node)
workflow.add_node("get_employees", get_employees_node)
workflow.add_node("respond", respond_node)

workflow.set_entry_point("agent")
workflow.add_conditional_edges(
    "agent",
    lambda state: state["next_action"],
    {
        "get_employees": "get_employees",
        "respond": "respond"
    }
)
workflow.add_edge("get_employees", "respond")

# Compile and use
app = workflow.compile()

# Execute workflow
initial_state = {
    "messages": [HumanMessage(content="Show me all employees")],
    "next_action": ""
}
result = app.invoke(initial_state)