Merge pull request #182 from restackio/agent-examples-update

Update agent examples

Author: leawn

agent_apis/pyproject.toml

@@ -7,11 +7,11 @@ requires-python = ">=3.10,<3.13"
 readme = "README.md"
 dependencies = [
     "pydantic>=2.10.6",
-    "restack-ai==0.0.62",
     "watchfiles>=1.0.4",
     "python-dotenv==1.0.1",
     "openai>=1.61.0",
     "aiohttp>=3.11.12",
+    "restack-ai>=0.0.63",
 ]
 
 [project.scripts]
@@ -22,6 +22,12 @@ schedule = "schedule_workflow:run_schedule_workflow"
 [tool.hatch.build.targets.sdist]
 include = ["src"]
 
+[tool.hatch.metadata]
+allow-direct-references = true
+
+[tool.uv.sources]
+restack-ai = { path = "../../../sdk/engine/libraries/python/dist/restack_ai-0.0.63-py3-none-any.whl" }
+
 [tool.hatch.build.targets.wheel]
 include = ["src"]
 

agent_apis/schedule_workflow.py

@@ -1,31 +1,34 @@
 import asyncio
+import sys
 import time
-from restack_ai import Restack
 from dataclasses import dataclass
 
+from restack_ai import Restack
+
+
 @dataclass
 class InputParams:
     name: str
 
-async def main():
+
+async def main() -> None:
     client = Restack()
 
     workflow_id = f"{int(time.time() * 1000)}-MultistepWorkflow"
-    runId = await client.schedule_workflow(
+    run_id = await client.schedule_workflow(
         workflow_name="MultistepWorkflow",
         workflow_id=workflow_id,
-        input=InputParams(name="Restack AI SDK User")
+        workflow_input=InputParams(name="Restack AI SDK User"),
     )
 
-    await client.get_workflow_result(
-        workflow_id=workflow_id,
-        run_id=runId
-    )
+    await client.get_workflow_result(workflow_id=workflow_id, run_id=run_id)
 
-    exit(0)
+    sys.exit(0)
 
-def run_schedule_workflow():
+
+def run_schedule_workflow() -> None:
     asyncio.run(main())
 
+
 if __name__ == "__main__":
-    run_schedule_workflow()
+    run_schedule_workflow()

agent_apis/src/client.py

@@ -1,7 +1,9 @@
 import os
+
+from dotenv import load_dotenv
 from restack_ai import Restack
 from restack_ai.restack import CloudConnectionOptions
-from dotenv import load_dotenv  
+
 # Load environment variables from a .env file
 load_dotenv()
 
@@ -12,9 +14,6 @@
 api_address = os.getenv("RESTACK_ENGINE_API_ADDRESS")
 
 connection_options = CloudConnectionOptions(
-    engine_id=engine_id,
-    address=address,
-    api_key=api_key,
-    api_address=api_address
+    engine_id=engine_id, address=address, api_key=api_key, api_address=api_address
 )
-client = Restack(connection_options)
+client = Restack(connection_options)

agent_apis/src/functions/llm.py

@@ -1,38 +1,50 @@
-from restack_ai.function import function, log, FunctionFailure
-from openai import OpenAI
-from dataclasses import dataclass
 import os
+from dataclasses import dataclass
+
 from dotenv import load_dotenv
+from openai import OpenAI
+from restack_ai.function import FunctionFailure, function, log
 
 load_dotenv()
 
+
 @dataclass
 class FunctionInputParams:
     user_content: str
     system_content: str | None = None
     model: str | None = None
 
+
+def raise_exception(message: str) -> None:
+    log.error(message)
+    raise Exception(message)
+
+
 @function.defn()
-async def llm(input: FunctionInputParams) -> str:
+async def llm(function_input: FunctionInputParams) -> str:
     try:
-        log.info("llm function started", input=input)
+        log.info("llm function started", input=function_input)
 
-        if (os.environ.get("RESTACK_API_KEY") is None):
-            raise FunctionFailure("RESTACK_API_KEY is not set", non_retryable=True)
-        
-        client = OpenAI(base_url="https://ai.restack.io", api_key=os.environ.get("RESTACK_API_KEY"))
+        if os.environ.get("RESTACK_API_KEY") is None:
+            error_message = "RESTACK_API_KEY is not set"
+            raise_exception(error_message)
+
+        client = OpenAI(
+            base_url="https://ai.restack.io", api_key=os.environ.get("RESTACK_API_KEY")
+        )
 
         messages = []
-        if input.system_content:
-            messages.append({"role": "system", "content": input.system_content})
-        messages.append({"role": "user", "content": input.user_content})
+        if function_input.system_content:
+            messages.append(
+                {"role": "system", "content": function_input.system_content}
+            )
+        messages.append({"role": "user", "content": function_input.user_content})
 
         response = client.chat.completions.create(
-            model=input.model or "gpt-4o-mini",
-            messages=messages
+            model=function_input.model or "gpt-4o-mini", messages=messages
         )
         log.info("llm function completed", response=response)
         return response.choices[0].message.content
     except Exception as e:
-        log.error("llm function failed", error=e)
-        raise e
+        error_message = "llm function failed"
+        raise FunctionFailure(error_message, non_retryable=True) from e

agent_apis/src/functions/weather.py

@@ -1,21 +1,26 @@
-from restack_ai.function import function, log
 import aiohttp
+from restack_ai.function import function, log
+
+HTTP_OK = 200
+
+
+def raise_exception(message: str) -> None:
+    log.error(message)
+    raise Exception(message)
+
 
 @function.defn()
 async def weather() -> str:
-    url = f"https://api.open-meteo.com/v1/forecast?latitude=52.52&longitude=13.41&current=temperature_2m,wind_speed_10m&hourly=temperature_2m,relative_humidity_2m,wind_speed_10m"
+    url = "https://api.open-meteo.com/v1/forecast?latitude=52.52&longitude=13.41&current=temperature_2m,wind_speed_10m&hourly=temperature_2m,relative_humidity_2m,wind_speed_10m"
     try:
-      async with aiohttp.ClientSession() as session:
-          async with session.get(url) as response:
-              log.info("response", response=response)
-              if response.status == 200:
-                  data = await response.json()
-                  log.info("weather data", data=data)
-                  return str(data)
-              else:
-                  log.error("Error: {response}")
-                  raise Exception(f"Error: {response.status}")
-    except Exception as e:
+        async with aiohttp.ClientSession() as session, session.get(url) as response:
+            log.info("response", response=response)
+            if response.status == HTTP_OK:
+                data = await response.json()
+                log.info("weather data", data=data)
+                return str(data)
+            error_message = f"Error: {response.status}"
+            raise_exception(error_message)
+    except Exception:
         log.error("Error: {e}")
-        raise e
-    
+        raise

agent_apis/src/services.py

@@ -1,30 +1,36 @@
 import asyncio
-import os
+import logging
+import webbrowser
+from pathlib import Path
+
+from watchfiles import run_process
+
+from src.client import client
 from src.functions.llm import llm
 from src.functions.weather import weather
-from src.client import client
 from src.workflows.multistep import MultistepWorkflow
-from watchfiles import run_process
-import webbrowser
 
-async def main():
 
+async def main() -> None:
     await client.start_service(
-        workflows= [MultistepWorkflow],
-        functions= [llm, weather],
+        workflows=[MultistepWorkflow],
+        functions=[llm, weather],
     )
 
-def run_services():
+
+def run_services() -> None:
     try:
         asyncio.run(main())
     except KeyboardInterrupt:
-        print("Service interrupted by user. Exiting gracefully.")
+        logging.info("Service interrupted by user. Exiting gracefully.")
 
-def watch_services():
-    watch_path = os.getcwd()
-    print(f"Watching {watch_path} and its subdirectories for changes...")
+
+def watch_services() -> None:
+    watch_path = Path.cwd()
+    logging.info("Watching %s and its subdirectories for changes...", watch_path)
     webbrowser.open("http://localhost:5233")
     run_process(watch_path, recursive=True, target=run_services)
 
+
 if __name__ == "__main__":
-       run_services()
+    run_services()

agent_apis/src/workflows/multistep.py

@@ -1,40 +1,39 @@
-from pydantic import BaseModel, Field
-from restack_ai.workflow import workflow, import_functions, log
 from datetime import timedelta
 
+from pydantic import BaseModel, Field
+from restack_ai.workflow import import_functions, log, workflow
+
 with import_functions():
-    from src.functions.llm import llm, FunctionInputParams
+    from src.functions.llm import FunctionInputParams, llm
     from src.functions.weather import weather
 
-class WorkflowInputParams ( BaseModel):
+
+class WorkflowInputParams(BaseModel):
     name: str = Field(default="John Doe")
 
+
 @workflow.defn()
 class MultistepWorkflow:
     @workflow.run
-    async def run(self, input: WorkflowInputParams):
-        log.info("MultistepWorkflow started", input=input)
-        user_content = f"Greet this person {input.name}"
+    async def run(self, workflow_input: WorkflowInputParams) -> dict:
+        log.info("MultistepWorkflow started", workflow_input=workflow_input)
+        user_content = f"Greet this person {workflow_input.name}"
 
         # Step 1 get weather data
         weather_data = await workflow.step(
-            weather,
-            start_to_close_timeout=timedelta(seconds=120)
+            function=weather, start_to_close_timeout=timedelta(seconds=120)
         )
 
         # Step 2 Generate greeting with LLM  based on name and weather data
 
         llm_message = await workflow.step(
-            llm,
-            FunctionInputParams(
+            function=llm,
+            workflow_input=FunctionInputParams(
                 system_content=f"You are a personal assitant and have access to weather data {weather_data}. Always greet person with relevant info from weather data",
                 user_content=user_content,
-                model="gpt-4o-mini"
+                model="gpt-4o-mini",
             ),
-            start_to_close_timeout=timedelta(seconds=120)
+            start_to_close_timeout=timedelta(seconds=120),
         )
         log.info("MultistepWorkflow completed", llm_message=llm_message)
-        return {
-            "message": llm_message,
-            "weather": weather_data
-        }
+        return {"message": llm_message, "weather": weather_data}

agent_chat/event_agent.py

@@ -1,8 +1,10 @@
 import asyncio
+import sys
+
 from restack_ai import Restack
 
 
-async def main(agent_id: str, run_id: str):
+async def main(agent_id: str, run_id: str) -> None:
     client = Restack()
 
     await client.send_agent_event(
@@ -18,11 +20,16 @@ async def main(agent_id: str, run_id: str):
         event_name="end",
     )
 
-    exit(0)
+    sys.exit(0)
 
 
-def run_event_agent():
-    asyncio.run(main(agent_id="your-agent-id", run_id="your-run-id"))
+def run_event_agent() -> None:
+    asyncio.run(
+        main(
+            agent_id="1739788461173-AgentChat",
+            run_id="c3937cc9-8d88-4e37-85e1-59e78cf1bf60",
+        )
+    )
 
 
 if __name__ == "__main__":

agent_chat/pyproject.toml

@@ -10,7 +10,7 @@ dependencies = [
     "watchfiles>=1.0.4",
     "python-dotenv==1.0.1",
     "openai>=1.61.0",
-    "restack-ai>=0.0.62",
+    "restack-ai>=0.0.63",
 ]
 
 [project.scripts]
@@ -25,6 +25,12 @@ include = ["src"]
 [tool.hatch.build.targets.wheel]
 include = ["src"]
 
+[tool.hatch.metadata]
+allow-direct-references = true
+
+[tool.uv.sources]
+restack-ai = { path = "../../../sdk/engine/libraries/python/dist/restack_ai-0.0.63-py3-none-any.whl" }
+
 [build-system]
 requires = ["hatchling"]
 build-backend = "hatchling.build"

agent_chat/schedule_agent.py

@@ -1,21 +1,20 @@
 import asyncio
+import sys
 import time
+
 from restack_ai import Restack
 
 
-async def main():
+async def main() -> None:
     client = Restack()
 
     agent_id = f"{int(time.time() * 1000)}-AgentChat"
-    await client.schedule_agent(
-        agent_name="AgentChat",
-        agent_id=agent_id
-    )
+    await client.schedule_agent(agent_name="AgentChat", agent_id=agent_id)
 
-    exit(0)
+    sys.exit(0)
 
 
-def run_schedule_agent():
+def run_schedule_agent() -> None:
     asyncio.run(main())