suna/sdk_example.py

66 lines
2.2 KiB
Python

import asyncio
import os
import kortix
kortix.set_api_key(os.getenv("KORTIX_API_KEY"))
kortix.set_base_url(os.getenv("KORTIX_BASE_URL"))
mcp = kortix.FastMCP("Weather tools")
mcp.set_api_base(os.getenv("MCP_BASE_URL", "http://localhost:4000")) # Publicly accessible API base URL
mcp.set_api_key(os.getenv("MCP_API_KEY")) # Used for authentication with MCP
# Need to do this somewhere and expose it on MCP_BASE_URL
# mcp.run(port=4000, host="0.0.0.0")
@mcp.tool
async def get_weather(city: str) -> str:
return f"The weather in {city} is sunny."
async def get_weather_agent():
weather_agent_id = await db.get("weather_agent_id")
if not weather_agent_id:
agent = await kortix.Agent.create(
tools=[mcp.tools["get_weather"]],
# TODO: have full control over what exact tools are given to the agent
mcp=[mcp],
system_prompt="You are a weather agent. You are given a city and you need to return the weather in that city.",
model="gpt-4o-mini",
name="Weather Agent",
)
# SDK user has to store these in their DB somewhere again
debug(agent.id, agent.created_at)
await db.set("weather_agent_id", agent.id)
else:
agent = await kortix.Agent.get(weather_agent_id)
return agent
async def get_thread():
thread_id = await db.get("thread_id")
if not thread_id:
thread = await kortix.Thread.create()
debug(thread.id, thread.messages, thread.agent_runs)
await db.set("thread_id", thread.id)
else:
thread = await kortix.Thread.get(thread_id)
return thread
async def main():
# We'll expose whatever the current UI allows the user to do, just with an SDK
# It'll have to work the exact same way. Basically like interacting with the UI with code
agent = await get_weather_agent()
agent.add_tools()
agent.toggle_tool("get_weather")
thread = await get_thread()
response = await agent.run("What is the weather in Tokyo?", thread=thread)
# The agent runs fully until the task is complete and returns the final response from LLM
print(response)
if __name__ == "__main__":
asyncio.run(main())