StateGraph

from typing_extensions import TypedDict
from upsonic.graphv2 import StateGraph, START, END

class ConversationState(TypedDict):
    messages: list[str]
    response: str

def process_node(state: ConversationState) -> dict:
    """Process the user's message."""
    if not state["messages"]:
        return {"response": "No messages to process"}
    user_message = state["messages"][-1]
    response = f"Echo: {user_message}"
    return {"response": response}

# Create and build the graph
builder = StateGraph(ConversationState)
builder.add_node("process", process_node)
builder.add_edge(START, "process")
builder.add_edge("process", END)

graph = builder.compile()

# Execute the graph
result = graph.invoke({
    "messages": ["Hello!"],
    "response": ""
})

print(result["response"])  # Output: Echo: Hello!

from typing import Annotated, List
from typing_extensions import TypedDict
import operator
from upsonic.graphv2 import StateGraph, START, END
from upsonic.models import infer_model
from upsonic.messages import ModelRequest, UserPromptPart, SystemPromptPart, TextPart
from upsonic.tools import tool

class AgentState(TypedDict):
    messages: Annotated[List, operator.add]
    result: str

@tool
def calculator(a: float, b: float, operation: str) -> float:
    """Perform basic math operations. Use this tool to calculate mathematical expressions.
    
    Args:
        a: First number
        b: Second number
        operation: Operation to perform - either "add" or "multiply"
    
    Returns:
        The result of the calculation
    """
    if operation == "add":
        return a + b
    elif operation == "multiply":
        return a * b
    return 0

def agent_node(state: AgentState) -> dict:
    """Agent node with tool access."""
    model = infer_model("anthropic/claude-sonnet-4-6")
    model_with_tools = model.bind_tools([calculator])
    
    user_message = state["messages"][-1] if state["messages"] else "Hello"
    request = ModelRequest(parts=[
        SystemPromptPart(content="You are a helpful assistant."),
        UserPromptPart(content=str(user_message))
    ])
    
    response = model_with_tools.invoke([request])
    
    # Extract text content from response
    text_content = ""
    for part in response.parts:
        if isinstance(part, TextPart):
            text_content += part.content
    
    # If no text content, use the string representation
    if not text_content:
        text_content = str(response)
    
    return {"messages": [response], "result": text_content}

# Build graph
builder = StateGraph(AgentState)
builder.add_node("agent", agent_node)
builder.add_edge(START, "agent")
builder.add_edge("agent", END)

graph = builder.compile()

# Execute
result = graph.invoke({
    "messages": ["What is 23 multiplied by 17?"],
    "result": ""
})

print(result["result"])