From 9a032e7af13061b9b7ff8d21379089b750e597ef Mon Sep 17 00:00:00 2001 From: leawn Date: Mon, 17 Feb 2025 12:25:56 +0100 Subject: [PATCH 1/8] Update agent examples --- agent_apis/pyproject.toml | 8 +++++++- agent_apis/schedule_workflow.py | 2 +- agent_apis/src/workflows/multistep.py | 12 ++++++------ agent_chat/event_agent.py | 2 +- agent_chat/pyproject.toml | 8 +++++++- agent_chat/src/agents/agent.py | 6 +++--- agent_chat/src/functions/llm_chat.py | 12 ++++++------ agent_rag/pyproject.toml | 8 +++++++- agent_rag/src/agents/chat_rag.py | 6 +++--- agent_rag/src/functions/llm_chat.py | 2 +- agent_todo/pyproject.toml | 10 ++++++++-- agent_todo/src/agents/agent_todo.py | 17 ++++++++++------- agent_tool/pyproject.toml | 8 +++++++- agent_tool/src/agents/chat_tool_functions.py | 12 ++++++------ 14 files changed, 73 insertions(+), 40 deletions(-) diff --git a/agent_apis/pyproject.toml b/agent_apis/pyproject.toml index 87396ee7..3985abe2 100644 --- a/agent_apis/pyproject.toml +++ b/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"] diff --git a/agent_apis/schedule_workflow.py b/agent_apis/schedule_workflow.py index 12e9419f..5a2808f9 100644 --- a/agent_apis/schedule_workflow.py +++ b/agent_apis/schedule_workflow.py @@ -14,7 +14,7 @@ async def main(): runId = 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( diff --git a/agent_apis/src/workflows/multistep.py b/agent_apis/src/workflows/multistep.py index df2dfc20..58932098 100644 --- a/agent_apis/src/workflows/multistep.py +++ b/agent_apis/src/workflows/multistep.py @@ -12,21 +12,21 @@ class WorkflowInputParams ( BaseModel): @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): + 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, + 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" diff --git a/agent_chat/event_agent.py b/agent_chat/event_agent.py index 0df0e7e5..2c9a103e 100644 --- a/agent_chat/event_agent.py +++ b/agent_chat/event_agent.py @@ -22,7 +22,7 @@ async def main(agent_id: str, run_id: str): def run_event_agent(): - asyncio.run(main(agent_id="your-agent-id", run_id="your-run-id")) + asyncio.run(main(agent_id="1739788461173-AgentChat", run_id="c3937cc9-8d88-4e37-85e1-59e78cf1bf60")) if __name__ == "__main__": diff --git a/agent_chat/pyproject.toml b/agent_chat/pyproject.toml index 7849a860..de09dab7 100644 --- a/agent_chat/pyproject.toml +++ b/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" diff --git a/agent_chat/src/agents/agent.py b/agent_chat/src/agents/agent.py index 42fd966a..e16f6d0f 100644 --- a/agent_chat/src/agents/agent.py +++ b/agent_chat/src/agents/agent.py @@ -26,8 +26,8 @@ async def message(self, message: MessageEvent) -> List[Message]: log.info(f"Received message: {message.content}") self.messages.append({"role": "user", "content": message.content}) assistant_message = await agent.step( - llm_chat, - LlmChatInput(messages=self.messages), + function=llm_chat, + agent_input=LlmChatInput(messages=self.messages), start_to_close_timeout=timedelta(seconds=120), ) self.messages.append(assistant_message) @@ -40,6 +40,6 @@ async def end(self, end: EndEvent) -> EndEvent: return end @agent.run - async def run(self, input: dict): + async def run(self): await agent.condition(lambda: self.end) return diff --git a/agent_chat/src/functions/llm_chat.py b/agent_chat/src/functions/llm_chat.py index b835643c..5f7e6498 100644 --- a/agent_chat/src/functions/llm_chat.py +++ b/agent_chat/src/functions/llm_chat.py @@ -21,9 +21,9 @@ class LlmChatInput(BaseModel): @function.defn() -async def llm_chat(input: LlmChatInput) -> ChatCompletion: +async def llm_chat(agent_input: LlmChatInput) -> ChatCompletion: try: - log.info("llm_chat function started", input=input) + log.info("llm_chat function started", agent_input=agent_input) if (os.environ.get("RESTACK_API_KEY") is None): raise FunctionFailure("RESTACK_API_KEY is not set", non_retryable=True) @@ -32,12 +32,12 @@ async def llm_chat(input: LlmChatInput) -> ChatCompletion: base_url="https://ai.restack.io", api_key=os.environ.get("RESTACK_API_KEY") ) - if input.system_content: - input.messages.append({"role": "system", "content": input.system_content}) + if agent_input.system_content: + agent_input.messages.append({"role": "system", "content": agent_input.system_content}) response = client.chat.completions.create( - model=input.model or "gpt-4o-mini", - messages=input.messages, + model=agent_input.model or "gpt-4o-mini", + messages=agent_input.messages, ) log.info("llm_chat function completed", response=response) diff --git a/agent_rag/pyproject.toml b/agent_rag/pyproject.toml index 3569bd03..f8625315 100644 --- a/agent_rag/pyproject.toml +++ b/agent_rag/pyproject.toml @@ -8,10 +8,10 @@ readme = "README.md" dependencies = [ "openai>=1.61.0", "pydantic>=2.10.6", - "restack-ai==0.0.62", "watchfiles>=1.0.4", "requests==2.32.3", "python-dotenv==1.0.1", + "restack-ai>=0.0.63", ] [project.scripts] @@ -24,6 +24,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" diff --git a/agent_rag/src/agents/chat_rag.py b/agent_rag/src/agents/chat_rag.py index 676e6e27..47602d4b 100644 --- a/agent_rag/src/agents/chat_rag.py +++ b/agent_rag/src/agents/chat_rag.py @@ -28,7 +28,7 @@ async def message(self, message: MessageEvent) -> List[Message]: log.info(f"Received message: {message.content}") sales_info = await agent.step( - lookupSales, start_to_close_timeout=timedelta(seconds=120) + function=lookupSales, start_to_close_timeout=timedelta(seconds=120) ) system_content = f"You are a helpful assistant that can help with sales data. Here is the sales information: {sales_info}" @@ -36,8 +36,8 @@ async def message(self, message: MessageEvent) -> List[Message]: self.messages.append(Message(role="user", content=message.content or "")) completion = await agent.step( - llm_chat, - LlmChatInput(messages=self.messages, system_content=system_content), + function=llm_chat, + agent_input=LlmChatInput(messages=self.messages, system_content=system_content), start_to_close_timeout=timedelta(seconds=120), ) diff --git a/agent_rag/src/functions/llm_chat.py b/agent_rag/src/functions/llm_chat.py index fe48078c..d013f925 100644 --- a/agent_rag/src/functions/llm_chat.py +++ b/agent_rag/src/functions/llm_chat.py @@ -1,4 +1,4 @@ -from restack_ai.function import function, log +from restack_ai.function import function, log, FunctionFailure from openai import OpenAI from openai.types.chat.chat_completion import ChatCompletion import os diff --git a/agent_todo/pyproject.toml b/agent_todo/pyproject.toml index 9d6765d8..90cca94a 100644 --- a/agent_todo/pyproject.toml +++ b/agent_todo/pyproject.toml @@ -7,10 +7,10 @@ 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" + "openai>=1.61.0", + "restack-ai>=0.0.63", ] [project.scripts] @@ -21,6 +21,12 @@ schedule = "schedule:run_schedule" [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"] diff --git a/agent_todo/src/agents/agent_todo.py b/agent_todo/src/agents/agent_todo.py index 11f8d752..73dbee57 100644 --- a/agent_todo/src/agents/agent_todo.py +++ b/agent_todo/src/agents/agent_todo.py @@ -53,8 +53,8 @@ async def message(self, message: MessageEvent) -> List[Message]: self.messages.append(Message(role="user", content=message.content or "")) completion = await agent.step( - llm_chat, - LlmChatInput(messages=self.messages, tools=tools), + function=llm_chat, + agent_input=LlmChatInput(messages=self.messages, tools=tools), start_to_close_timeout=timedelta(seconds=120), ) @@ -83,7 +83,10 @@ async def message(self, message: MessageEvent) -> List[Message]: tool_call.function.arguments ) - result = await agent.step(todo_create, input=args) + result = await agent.step( + function=todo_create, + agent_input=args, + ) self.messages.append( Message( role="tool", @@ -93,8 +96,8 @@ async def message(self, message: MessageEvent) -> List[Message]: ) completion_with_tool_call = await agent.step( - llm_chat, - LlmChatInput(messages=self.messages, tools=tools), + function=llm_chat, + agent_input=LlmChatInput(messages=self.messages, tools=tools), start_to_close_timeout=timedelta(seconds=120), ) self.messages.append( @@ -123,8 +126,8 @@ async def message(self, message: MessageEvent) -> List[Message]: ) completion_with_tool_call = await agent.step( - llm_chat, - LlmChatInput(messages=self.messages, tools=tools), + function=llm_chat, + agent_input=LlmChatInput(messages=self.messages, tools=tools), start_to_close_timeout=timedelta(seconds=120), ) self.messages.append( diff --git a/agent_tool/pyproject.toml b/agent_tool/pyproject.toml index 8650dec4..0810983b 100644 --- a/agent_tool/pyproject.toml +++ b/agent_tool/pyproject.toml @@ -8,10 +8,10 @@ readme = "README.md" dependencies = [ "openai>=1.61.0", "pydantic>=2.10.6", - "restack-ai==0.0.62", "watchfiles>=1.0.4", "requests==2.32.3", "python-dotenv==1.0.1", + "restack-ai>=0.0.63", ] [project.scripts] @@ -24,6 +24,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" diff --git a/agent_tool/src/agents/chat_tool_functions.py b/agent_tool/src/agents/chat_tool_functions.py index 83d54540..7b925312 100644 --- a/agent_tool/src/agents/chat_tool_functions.py +++ b/agent_tool/src/agents/chat_tool_functions.py @@ -49,8 +49,8 @@ async def message(self, message: MessageEvent) -> List[Message]: self.messages.append(Message(role="user", content=message.content or "")) completion = await agent.step( - llm_chat, - LlmChatInput( + function=llm_chat, + agent_input=LlmChatInput( messages=self.messages, tools=tools, system_content=system_content ), start_to_close_timeout=timedelta(seconds=120), @@ -84,8 +84,8 @@ async def message(self, message: MessageEvent) -> List[Message]: log.info(f"calling {name} with args: {args}") result = await agent.step( - lookupSales, - input=LookupSalesInput(category=args.category), + function=lookupSales, + agent_input=LookupSalesInput(category=args.category), start_to_close_timeout=timedelta(seconds=120), ) self.messages.append( @@ -97,8 +97,8 @@ async def message(self, message: MessageEvent) -> List[Message]: ) completion_with_tool_call = await agent.step( - llm_chat, - LlmChatInput( + function=llm_chat, + agent_input=LlmChatInput( messages=self.messages, system_content=system_content ), start_to_close_timeout=timedelta(seconds=120), From 1c65865b9b1cf42ff582e2c3215f620c21000770 Mon Sep 17 00:00:00 2001 From: leawn Date: Mon, 17 Feb 2025 12:45:58 +0100 Subject: [PATCH 2/8] Add rules & update agent_apis example --- agent_apis/schedule_workflow.py | 25 ++++++++------- agent_apis/src/client.py | 11 +++---- agent_apis/src/functions/llm.py | 44 +++++++++++++++++---------- agent_apis/src/functions/weather.py | 35 ++++++++++++--------- agent_apis/src/services.py | 32 +++++++++++-------- agent_apis/src/workflows/multistep.py | 25 ++++++++------- pyproject.toml | 20 ++++++++++++ 7 files changed, 118 insertions(+), 74 deletions(-) create mode 100644 pyproject.toml diff --git a/agent_apis/schedule_workflow.py b/agent_apis/schedule_workflow.py index 5a2808f9..cf7107b7 100644 --- a/agent_apis/schedule_workflow.py +++ b/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, - workflow_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() \ No newline at end of file + run_schedule_workflow() diff --git a/agent_apis/src/client.py b/agent_apis/src/client.py index b71efb06..885bf8ea 100644 --- a/agent_apis/src/client.py +++ b/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) \ No newline at end of file +client = Restack(connection_options) diff --git a/agent_apis/src/functions/llm.py b/agent_apis/src/functions/llm.py index b4c7494f..6412156f 100644 --- a/agent_apis/src/functions/llm.py +++ b/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 diff --git a/agent_apis/src/functions/weather.py b/agent_apis/src/functions/weather.py index 746cca9c..b9ab92d5 100644 --- a/agent_apis/src/functions/weather.py +++ b/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¤t=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¤t=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 - \ No newline at end of file + raise diff --git a/agent_apis/src/services.py b/agent_apis/src/services.py index 59480ebd..5b4dd822 100644 --- a/agent_apis/src/services.py +++ b/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() \ No newline at end of file + run_services() diff --git a/agent_apis/src/workflows/multistep.py b/agent_apis/src/workflows/multistep.py index 58932098..fd589741 100644 --- a/agent_apis/src/workflows/multistep.py +++ b/agent_apis/src/workflows/multistep.py @@ -1,25 +1,27 @@ -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, workflow_input: WorkflowInputParams): + 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( - function=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 @@ -29,12 +31,9 @@ async def run(self, workflow_input: WorkflowInputParams): 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} diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 00000000..78c255b9 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,20 @@ +[project] +name = "examples-python" +version = "0.0.1" +requires-python = ">=3.10" + +[tool.uv.workspace] + +[dependency-groups] +dev = [ + "ruff>=0.9.4", +] + +[tool.ruff.lint] +extend-select = ["ALL"] +ignore = ["ANN401", "E501", "D100", "D101", "D102", "D103", "D107", "TRY002", "D213", "D203", "COM812", "D104", "INP001"] + +[tool.ruff.format] +quote-style = "double" +indent-style = "space" +docstring-code-format = true \ No newline at end of file From acd730938323f73cd2ed985303bb1b67a6575aef Mon Sep 17 00:00:00 2001 From: leawn Date: Mon, 17 Feb 2025 12:49:21 +0100 Subject: [PATCH 3/8] Update agent_chat example --- agent_chat/event_agent.py | 15 ++++++++--- agent_chat/schedule_agent.py | 13 +++++----- agent_chat/src/agents/agent.py | 9 +++---- agent_chat/src/client.py | 3 ++- agent_chat/src/functions/llm_chat.py | 38 ++++++++++++++++++---------- agent_chat/src/services.py | 25 ++++++++++-------- 6 files changed, 61 insertions(+), 42 deletions(-) diff --git a/agent_chat/event_agent.py b/agent_chat/event_agent.py index 2c9a103e..925a5a86 100644 --- a/agent_chat/event_agent.py +++ b/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="1739788461173-AgentChat", run_id="c3937cc9-8d88-4e37-85e1-59e78cf1bf60")) +def run_event_agent() -> None: + asyncio.run( + main( + agent_id="1739788461173-AgentChat", + run_id="c3937cc9-8d88-4e37-85e1-59e78cf1bf60", + ) + ) if __name__ == "__main__": diff --git a/agent_chat/schedule_agent.py b/agent_chat/schedule_agent.py index a71b418c..bf510b62 100644 --- a/agent_chat/schedule_agent.py +++ b/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()) diff --git a/agent_chat/src/agents/agent.py b/agent_chat/src/agents/agent.py index e16f6d0f..e2bac360 100644 --- a/agent_chat/src/agents/agent.py +++ b/agent_chat/src/agents/agent.py @@ -1,10 +1,10 @@ from datetime import timedelta -from typing import List + from pydantic import BaseModel from restack_ai.agent import agent, import_functions, log with import_functions(): - from src.functions.llm_chat import llm_chat, LlmChatInput, Message + from src.functions.llm_chat import LlmChatInput, Message, llm_chat class MessageEvent(BaseModel): @@ -22,7 +22,7 @@ def __init__(self) -> None: self.messages = [] @agent.event - async def message(self, message: MessageEvent) -> List[Message]: + async def message(self, message: MessageEvent) -> list[Message]: log.info(f"Received message: {message.content}") self.messages.append({"role": "user", "content": message.content}) assistant_message = await agent.step( @@ -40,6 +40,5 @@ async def end(self, end: EndEvent) -> EndEvent: return end @agent.run - async def run(self): + async def run(self) -> None: await agent.condition(lambda: self.end) - return diff --git a/agent_chat/src/client.py b/agent_chat/src/client.py index dee77832..885bf8ea 100644 --- a/agent_chat/src/client.py +++ b/agent_chat/src/client.py @@ -1,7 +1,8 @@ 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() diff --git a/agent_chat/src/functions/llm_chat.py b/agent_chat/src/functions/llm_chat.py index 5f7e6498..a644b368 100644 --- a/agent_chat/src/functions/llm_chat.py +++ b/agent_chat/src/functions/llm_chat.py @@ -1,10 +1,11 @@ -from restack_ai.function import function, log, FunctionFailure -from openai import OpenAI -from openai.types.chat.chat_completion import ChatCompletion import os +from typing import Literal + from dotenv import load_dotenv +from openai import OpenAI +from openai.types.chat.chat_completion import ChatCompletion from pydantic import BaseModel -from typing import Literal, Optional, List +from restack_ai.function import FunctionFailure, function, log load_dotenv() @@ -15,9 +16,14 @@ class Message(BaseModel): class LlmChatInput(BaseModel): - system_content: Optional[str] = None - model: Optional[str] = None - messages: Optional[List[Message]] = None + system_content: str | None = None + model: str | None = None + messages: list[Message] | None = None + + +def raise_exception(message: str) -> None: + log.error(message) + raise FunctionFailure(message, non_retryable=True) @function.defn() @@ -25,23 +31,27 @@ async def llm_chat(agent_input: LlmChatInput) -> ChatCompletion: try: log.info("llm_chat function started", agent_input=agent_input) - if (os.environ.get("RESTACK_API_KEY") is None): - raise FunctionFailure("RESTACK_API_KEY is not set", non_retryable=True) - + 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") ) if agent_input.system_content: - agent_input.messages.append({"role": "system", "content": agent_input.system_content}) + agent_input.messages.append( + {"role": "system", "content": agent_input.system_content} + ) response = client.chat.completions.create( model=agent_input.model or "gpt-4o-mini", messages=agent_input.messages, ) + except Exception as e: + log.error("llm_chat function failed", error=e) + raise + else: log.info("llm_chat function completed", response=response) return response - except Exception as e: - log.error("llm_chat function failed", error=e) - raise e diff --git a/agent_chat/src/services.py b/agent_chat/src/services.py index fe280985..8ff3e9fe 100644 --- a/agent_chat/src/services.py +++ b/agent_chat/src/services.py @@ -1,26 +1,29 @@ import asyncio -import os -from src.functions.llm_chat import llm_chat -from src.client import client -from src.agents.agent import AgentChat -from watchfiles import run_process +import logging import webbrowser +from pathlib import Path + +from watchfiles import run_process + +from src.agents.agent import AgentChat +from src.client import client +from src.functions.llm_chat import llm_chat -async def main(): +async def main() -> None: await client.start_service(agents=[AgentChat], functions=[llm_chat]) -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) From 4b961ed5adc0fc062678878067a280536591201b Mon Sep 17 00:00:00 2001 From: leawn Date: Mon, 17 Feb 2025 12:52:54 +0100 Subject: [PATCH 4/8] Update agent_rag example --- agent_rag/src/agents/chat_rag.py | 16 ++++----- agent_rag/src/client.py | 3 +- agent_rag/src/functions/llm_chat.py | 46 +++++++++++++++---------- agent_rag/src/functions/lookup_sales.py | 10 +++--- agent_rag/src/services.py | 24 +++++++------ 5 files changed, 56 insertions(+), 43 deletions(-) diff --git a/agent_rag/src/agents/chat_rag.py b/agent_rag/src/agents/chat_rag.py index 47602d4b..fad39edb 100644 --- a/agent_rag/src/agents/chat_rag.py +++ b/agent_rag/src/agents/chat_rag.py @@ -1,11 +1,10 @@ from datetime import timedelta -from typing import List + from pydantic import BaseModel from restack_ai.agent import agent, import_functions, log - with import_functions(): - from src.functions.llm_chat import llm_chat, LlmChatInput, Message + from src.functions.llm_chat import LlmChatInput, Message, llm_chat from src.functions.lookup_sales import lookupSales @@ -24,7 +23,7 @@ def __init__(self) -> None: self.messages = [] @agent.event - async def message(self, message: MessageEvent) -> List[Message]: + async def message(self, message: MessageEvent) -> list[Message]: log.info(f"Received message: {message.content}") sales_info = await agent.step( @@ -37,7 +36,9 @@ async def message(self, message: MessageEvent) -> List[Message]: completion = await agent.step( function=llm_chat, - agent_input=LlmChatInput(messages=self.messages, system_content=system_content), + agent_input=LlmChatInput( + messages=self.messages, system_content=system_content + ), start_to_close_timeout=timedelta(seconds=120), ) @@ -52,12 +53,11 @@ async def message(self, message: MessageEvent) -> List[Message]: return self.messages @agent.event - async def end(self, end: EndEvent) -> EndEvent: + async def end(self) -> EndEvent: log.info("Received end") self.end = True return {"end": True} @agent.run - async def run(self, input: dict): + async def run(self) -> None: await agent.condition(lambda: self.end) - return diff --git a/agent_rag/src/client.py b/agent_rag/src/client.py index dee77832..885bf8ea 100644 --- a/agent_rag/src/client.py +++ b/agent_rag/src/client.py @@ -1,7 +1,8 @@ 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() diff --git a/agent_rag/src/functions/llm_chat.py b/agent_rag/src/functions/llm_chat.py index d013f925..b4bfe22c 100644 --- a/agent_rag/src/functions/llm_chat.py +++ b/agent_rag/src/functions/llm_chat.py @@ -1,10 +1,11 @@ -from restack_ai.function import function, log, FunctionFailure -from openai import OpenAI -from openai.types.chat.chat_completion import ChatCompletion import os +from typing import Literal + from dotenv import load_dotenv +from openai import OpenAI +from openai.types.chat.chat_completion import ChatCompletion from pydantic import BaseModel -from typing import Literal, Optional, List +from restack_ai.function import FunctionFailure, function, log load_dotenv() @@ -15,32 +16,41 @@ class Message(BaseModel): class LlmChatInput(BaseModel): - system_content: Optional[str] = None - model: Optional[str] = None - messages: Optional[List[Message]] = None + system_content: str | None = None + model: str | None = None + messages: list[Message] | None = None + + +def raise_exception(message: str) -> None: + log.error(message) + raise FunctionFailure(message, non_retryable=True) @function.defn() -async def llm_chat(input: LlmChatInput) -> ChatCompletion: +async def llm_chat(function_input: LlmChatInput) -> ChatCompletion: try: - log.info("llm_chat function started", input=input) + log.info("llm_chat function started", function_input=function_input) + + if os.environ.get("RESTACK_API_KEY") is None: + error_message = "RESTACK_API_KEY is not set" + raise_exception(error_message) - 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 input.system_content: - input.messages.append( - Message(role="system", content=input.system_content or "") + if function_input.system_content: + function_input.messages.append( + Message(role="system", content=function_input.system_content or "") ) response = client.chat.completions.create( - model=input.model or "gpt-4o-mini", messages=input.messages + model=function_input.model or "gpt-4o-mini", + messages=function_input.messages, ) - return response except Exception as e: log.error("llm_chat function failed", error=e) - raise e + raise + else: + log.info("llm_chat function completed", response=response) + return response diff --git a/agent_rag/src/functions/lookup_sales.py b/agent_rag/src/functions/lookup_sales.py index e7304b79..2d380d71 100644 --- a/agent_rag/src/functions/lookup_sales.py +++ b/agent_rag/src/functions/lookup_sales.py @@ -1,5 +1,5 @@ -from restack_ai.function import function, log from pydantic import BaseModel +from restack_ai.function import function, log class SalesItem(BaseModel): @@ -12,9 +12,9 @@ class SalesItem(BaseModel): @function.defn() -async def lookupSales() -> str: +async def lookup_sales() -> str: try: - log.info("lookupSales function started", input=input) + log.info("lookup_sales function started") items = [ SalesItem( @@ -85,5 +85,5 @@ async def lookupSales() -> str: return str(items) except Exception as e: - log.error("lookupSales function failed", error=e) - raise e + log.error("lookup_sales function failed", error=e) + raise diff --git a/agent_rag/src/services.py b/agent_rag/src/services.py index fd3ef12e..a5510c53 100644 --- a/agent_rag/src/services.py +++ b/agent_rag/src/services.py @@ -1,28 +1,30 @@ import asyncio -import os -from watchfiles import run_process +import logging import webbrowser -from src.client import client -from src.functions.lookup_sales import lookupSales -from src.functions.llm_chat import llm_chat +from pathlib import Path + +from watchfiles import run_process from src.agents.chat_rag import AgentRag +from src.client import client +from src.functions.llm_chat import llm_chat +from src.functions.lookup_sales import lookupSales -async def main(): +async def main() -> None: await client.start_service(agents=[AgentRag], functions=[lookupSales, llm_chat]) -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) From f4d3eefe0383cfbac1ee5df99fcdfd6bbb285ba8 Mon Sep 17 00:00:00 2001 From: leawn Date: Mon, 17 Feb 2025 13:16:58 +0100 Subject: [PATCH 5/8] Update agent_todo example --- agent_todo/schedule.py | 8 +-- agent_todo/src/agents/agent_todo.py | 34 +++++++----- agent_todo/src/client.py | 3 +- .../functions/{random.py => get_random.py} | 14 ++--- .../functions/{result.py => get_result.py} | 17 +++--- agent_todo/src/functions/llm_chat.py | 54 +++++++++++-------- agent_todo/src/functions/todo_create.py | 15 +++--- agent_todo/src/services.py | 31 ++++++----- agent_todo/src/workflows/todo_execute.py | 25 ++++----- 9 files changed, 113 insertions(+), 88 deletions(-) rename agent_todo/src/functions/{random.py => get_random.py} (58%) rename agent_todo/src/functions/{result.py => get_result.py} (61%) diff --git a/agent_todo/schedule.py b/agent_todo/schedule.py index ef13b33a..73956a4b 100644 --- a/agent_todo/schedule.py +++ b/agent_todo/schedule.py @@ -1,10 +1,12 @@ import asyncio +import sys import time + from restack_ai import Restack from src.agents.agent_todo import AgentTodo -async def main(): +async def main() -> None: client = Restack() agent_id = f"{int(time.time() * 1000)}-{AgentTodo.__name__}" @@ -16,10 +18,10 @@ async def main(): await client.get_agent_result(agent_id=agent_id, run_id=run_id) - exit(0) + sys.exit(0) -def run_schedule(): +def run_schedule() -> None: asyncio.run(main()) diff --git a/agent_todo/src/agents/agent_todo.py b/agent_todo/src/agents/agent_todo.py index 73dbee57..5dab7d30 100644 --- a/agent_todo/src/agents/agent_todo.py +++ b/agent_todo/src/agents/agent_todo.py @@ -1,14 +1,15 @@ from datetime import timedelta -from typing import List + from pydantic import BaseModel from restack_ai.agent import agent, import_functions, log -from src.workflows.todo_execute import TodoExecute, TodoExecuteParams +from src.workflows.todo_execute import TodoExecute, TodoExecuteParams with import_functions(): from openai import pydantic_function_tool - from src.functions.llm_chat import llm_chat, LlmChatInput, Message - from src.functions.todo_create import todo_create, TodoCreateParams + + from src.functions.llm_chat import LlmChatInput, Message, llm_chat + from src.functions.todo_create import TodoCreateParams, todo_create class MessageEvent(BaseModel): @@ -26,7 +27,7 @@ def __init__(self) -> None: self.messages = [] @agent.event - async def message(self, message: MessageEvent) -> List[Message]: + async def message(self, message: MessageEvent) -> list[Message]: try: log.info(f"Received message: {message.content}") @@ -97,7 +98,9 @@ async def message(self, message: MessageEvent) -> List[Message]: completion_with_tool_call = await agent.step( function=llm_chat, - agent_input=LlmChatInput(messages=self.messages, tools=tools), + agent_input=LlmChatInput( + messages=self.messages, tools=tools + ), start_to_close_timeout=timedelta(seconds=120), ) self.messages.append( @@ -115,7 +118,9 @@ async def message(self, message: MessageEvent) -> List[Message]: ) result = await agent.child_execute( - workflow=TodoExecute, workflow_id=tool_call.id, input=args + workflow=TodoExecute, + workflow_id=tool_call.id, + input=args, ) self.messages.append( Message( @@ -127,7 +132,9 @@ async def message(self, message: MessageEvent) -> List[Message]: completion_with_tool_call = await agent.step( function=llm_chat, - agent_input=LlmChatInput(messages=self.messages, tools=tools), + agent_input=LlmChatInput( + messages=self.messages, tools=tools + ), start_to_close_timeout=timedelta(seconds=120), ) self.messages.append( @@ -146,19 +153,18 @@ async def message(self, message: MessageEvent) -> List[Message]: content=completion.choices[0].message.content or "", ) ) - - return self.messages except Exception as e: log.error(f"Error during message event: {e}") - raise e + raise + else: + return self.messages @agent.event - async def end(self, end: EndEvent) -> EndEvent: + async def end(self) -> EndEvent: log.info("Received end") self.end = True return {"end": True} @agent.run - async def run(self, input: dict): + async def run(self) -> None: await agent.condition(lambda: self.end) - return diff --git a/agent_todo/src/client.py b/agent_todo/src/client.py index dee77832..885bf8ea 100644 --- a/agent_todo/src/client.py +++ b/agent_todo/src/client.py @@ -1,7 +1,8 @@ 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() diff --git a/agent_todo/src/functions/random.py b/agent_todo/src/functions/get_random.py similarity index 58% rename from agent_todo/src/functions/random.py rename to agent_todo/src/functions/get_random.py index 5e92cbc0..4115c8b0 100644 --- a/agent_todo/src/functions/random.py +++ b/agent_todo/src/functions/get_random.py @@ -1,17 +1,19 @@ -from restack_ai.function import function, log +import secrets + from pydantic import BaseModel -import random +from restack_ai.function import function, log class RandomParams(BaseModel): - todoTitle: str + todo_title: str @function.defn() async def get_random(params: RandomParams) -> str: try: - random_number = random.randint(0, 100) - return f"The random number for {params.todoTitle} is {random_number}." + random_number = secrets.randbelow(100) except Exception as e: log.error("random function failed", error=e) - raise e + raise + else: + return f"The random number for {params.todo_title} is {random_number}." diff --git a/agent_todo/src/functions/result.py b/agent_todo/src/functions/get_result.py similarity index 61% rename from agent_todo/src/functions/result.py rename to agent_todo/src/functions/get_result.py index 76ebf086..6bff5185 100644 --- a/agent_todo/src/functions/result.py +++ b/agent_todo/src/functions/get_result.py @@ -1,23 +1,24 @@ -from restack_ai.function import function, log +import secrets + from pydantic import BaseModel -import random +from restack_ai.function import function, log class ResultParams(BaseModel): - todoTitle: str - todoId: str + todo_title: str + todo_id: str class ResultResponse(BaseModel): status: str - todoId: str + todo_id: str @function.defn() async def get_result(params: ResultParams) -> ResultResponse: try: - status = random.choice(["completed", "failed"]) - return ResultResponse(todoId=params.todoId, status=status) + status = secrets.choice(["completed", "failed"]) + return ResultResponse(todo_id=params.todo_id, status=status) except Exception as e: log.error("result function failed", error=e) - raise e + raise diff --git a/agent_todo/src/functions/llm_chat.py b/agent_todo/src/functions/llm_chat.py index 8d072ea2..7827bf49 100644 --- a/agent_todo/src/functions/llm_chat.py +++ b/agent_todo/src/functions/llm_chat.py @@ -1,14 +1,15 @@ -from restack_ai.function import function, log, FunctionFailure +import os +from typing import Literal + +from dotenv import load_dotenv from openai import OpenAI from openai.types.chat.chat_completion import ChatCompletion from openai.types.chat.chat_completion_message_tool_call import ( ChatCompletionMessageToolCall, ) from openai.types.chat.chat_completion_tool_param import ChatCompletionToolParam -import os -from dotenv import load_dotenv from pydantic import BaseModel -from typing import Literal, Optional, List +from restack_ai.function import FunctionFailure, function, log load_dotenv() @@ -16,42 +17,49 @@ class Message(BaseModel): role: Literal["system", "user", "assistant", "tool"] content: str - tool_call_id: Optional[str] = None - tool_calls: Optional[List[ChatCompletionMessageToolCall]] = None + tool_call_id: str | None = None + tool_calls: list[ChatCompletionMessageToolCall] | None = None class LlmChatInput(BaseModel): - system_content: Optional[str] = None - model: Optional[str] = None - messages: Optional[List[Message]] = None - tools: Optional[List[ChatCompletionToolParam]] = None + system_content: str | None = None + model: str | None = None + messages: list[Message] | None = None + tools: list[ChatCompletionToolParam] | None = None + + +def raise_exception(message: str) -> None: + log.error("llm_chat function failed", error=message) + raise FunctionFailure(message, non_retryable=True) @function.defn() -async def llm_chat(input: LlmChatInput) -> ChatCompletion: +async def llm_chat(function_input: LlmChatInput) -> ChatCompletion: try: - log.info("llm_chat function started", input=input) + log.info("llm_chat function started", function_input=function_input) - if (os.environ.get("RESTACK_API_KEY") is None): - raise FunctionFailure("RESTACK_API_KEY is not set", non_retryable=True) + if os.environ.get("RESTACK_API_KEY") is None: + raise_exception("RESTACK_API_KEY is not set") client = OpenAI( base_url="https://ai.restack.io", api_key=os.environ.get("RESTACK_API_KEY") ) - log.info("pydantic_function_tool", tools=input.tools) + log.info("pydantic_function_tool", tools=function_input.tools) - if input.system_content: - input.messages.append( - Message(role="system", content=input.system_content or "") + if function_input.system_content: + function_input.messages.append( + Message(role="system", content=function_input.system_content or "") ) response = client.chat.completions.create( - model=input.model or "gpt-4o-mini", - messages=input.messages, - tools=input.tools, + model=function_input.model or "gpt-4o-mini", + messages=function_input.messages, + tools=function_input.tools, ) - return response except Exception as e: log.error("llm_chat function failed", error=e) - raise e + raise + else: + log.info("llm_chat function completed", response=response) + return response diff --git a/agent_todo/src/functions/todo_create.py b/agent_todo/src/functions/todo_create.py index e7da5b7b..59baa654 100644 --- a/agent_todo/src/functions/todo_create.py +++ b/agent_todo/src/functions/todo_create.py @@ -1,6 +1,7 @@ -from restack_ai.function import function, log -import random +import secrets + from pydantic import BaseModel +from restack_ai.function import function, log class TodoCreateParams(BaseModel): @@ -12,10 +13,10 @@ async def todo_create(params: TodoCreateParams) -> str: try: log.info("todo_create function start", title=params.title) - todo_id = f"todo-{random.randint(1000, 9999)}" - - log.info("todo_create function completed", todo_id=todo_id) - return f"Created the todo '{params.title}' with id: {todo_id}" + todo_id = f"todo-{secrets.randbelow(9000) + 1000}" except Exception as e: log.error("todo_create function failed", error=e) - raise e + raise + else: + log.info("todo_create function completed", todo_id=todo_id) + return f"Created the todo '{params.title}' with id: {todo_id}" diff --git a/agent_todo/src/services.py b/agent_todo/src/services.py index f1f4e65c..4ae34472 100644 --- a/agent_todo/src/services.py +++ b/agent_todo/src/services.py @@ -1,17 +1,20 @@ import asyncio -import os -from src.functions.random import get_random -from src.functions.result import get_result -from src.client import client -from src.workflows.todo_execute import TodoExecute +import logging +import webbrowser +from pathlib import Path + +from watchfiles import run_process + from src.agents.agent_todo import AgentTodo -from src.functions.todo_create import todo_create +from src.client import client +from src.functions.get_random import get_random +from src.functions.get_result import get_result from src.functions.llm_chat import llm_chat -from watchfiles import run_process -import webbrowser +from src.functions.todo_create import todo_create +from src.workflows.todo_execute import TodoExecute -async def main(): +async def main() -> None: await client.start_service( agents=[AgentTodo], workflows=[TodoExecute], @@ -19,16 +22,16 @@ async def main(): ) -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) diff --git a/agent_todo/src/workflows/todo_execute.py b/agent_todo/src/workflows/todo_execute.py index 5753b2c3..e8470adb 100644 --- a/agent_todo/src/workflows/todo_execute.py +++ b/agent_todo/src/workflows/todo_execute.py @@ -1,20 +1,21 @@ from datetime import timedelta + from pydantic import BaseModel -from restack_ai.workflow import workflow, import_functions, log +from restack_ai.workflow import import_functions, log, workflow with import_functions(): - from src.functions.random import get_random, RandomParams - from src.functions.result import get_result, ResultParams + from src.functions.get_random import RandomParams, get_random + from src.functions.get_result import ResultParams, get_result class TodoExecuteParams(BaseModel): - todoTitle: str - todoId: str + todo_title: str + todo_id: str class TodoExecuteResponse(BaseModel): - todoId: str - todoTitle: str + todo_id: str + todo_title: str details: str status: str @@ -22,11 +23,11 @@ class TodoExecuteResponse(BaseModel): @workflow.defn() class TodoExecute: @workflow.run - async def run(self, params: TodoExecuteParams): + async def run(self, params: TodoExecuteParams) -> TodoExecuteResponse: log.info("TodoExecuteWorkflow started") random = await workflow.step( get_random, - input=RandomParams(todoTitle=params.todoTitle), + input=RandomParams(todo_title=params.todo_title), start_to_close_timeout=timedelta(seconds=120), ) @@ -34,13 +35,13 @@ async def run(self, params: TodoExecuteParams): result = await workflow.step( get_result, - input=ResultParams(todoTitle=params.todoTitle, todoId=params.todoId), + input=ResultParams(todo_title=params.todo_title, todo_id=params.todo_id), start_to_close_timeout=timedelta(seconds=120), ) todo_details = TodoExecuteResponse( - todoId=params.todoId, - todoTitle=params.todoTitle, + todo_id=params.todo_id, + todo_title=params.todo_title, details=random, status=result.status, ) From 8c954b9bb910dbc0ee22934b1c2d5bff4a9ea565 Mon Sep 17 00:00:00 2001 From: leawn Date: Mon, 17 Feb 2025 13:22:29 +0100 Subject: [PATCH 6/8] Update agent_tool example --- agent_tool/src/agents/chat_tool_functions.py | 16 +++--- agent_tool/src/client.py | 3 +- agent_tool/src/functions/llm_chat.py | 55 +++++++++++--------- agent_tool/src/functions/lookup_sales.py | 17 +++--- agent_tool/src/functions/new_function.py | 2 + agent_tool/src/services.py | 24 +++++---- 6 files changed, 66 insertions(+), 51 deletions(-) diff --git a/agent_tool/src/agents/chat_tool_functions.py b/agent_tool/src/agents/chat_tool_functions.py index 7b925312..e676ad5a 100644 --- a/agent_tool/src/agents/chat_tool_functions.py +++ b/agent_tool/src/agents/chat_tool_functions.py @@ -1,13 +1,14 @@ +# ruff: noqa: ERA001 from datetime import timedelta -from typing import List + from pydantic import BaseModel from restack_ai.agent import agent, import_functions, log - with import_functions(): from openai import pydantic_function_tool - from src.functions.llm_chat import llm_chat, LlmChatInput, Message - from src.functions.lookup_sales import lookupSales, LookupSalesInput + + from src.functions.llm_chat import LlmChatInput, Message, llm_chat + from src.functions.lookup_sales import LookupSalesInput, lookupSales # Step 2: Import your new function to the agent # from src.functions.new_function import new_function, FunctionInput, FunctionOutput @@ -27,7 +28,7 @@ def __init__(self) -> None: self.messages = [] @agent.event - async def message(self, message: MessageEvent) -> List[Message]: + async def message(self, message: MessageEvent) -> list[Message]: log.info(f"Received message: {message.content}") tools = [ @@ -131,12 +132,11 @@ async def message(self, message: MessageEvent) -> List[Message]: return self.messages @agent.event - async def end(self, end: EndEvent) -> EndEvent: + async def end(self) -> EndEvent: log.info("Received end") self.end = True return {"end": True} @agent.run - async def run(self, input: dict): + async def run(self) -> None: await agent.condition(lambda: self.end) - return diff --git a/agent_tool/src/client.py b/agent_tool/src/client.py index dee77832..885bf8ea 100644 --- a/agent_tool/src/client.py +++ b/agent_tool/src/client.py @@ -1,7 +1,8 @@ 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() diff --git a/agent_tool/src/functions/llm_chat.py b/agent_tool/src/functions/llm_chat.py index 63766000..713e4479 100644 --- a/agent_tool/src/functions/llm_chat.py +++ b/agent_tool/src/functions/llm_chat.py @@ -1,14 +1,15 @@ -from restack_ai.function import function, log, FunctionFailure +import os +from typing import Literal + +from dotenv import load_dotenv from openai import OpenAI from openai.types.chat.chat_completion import ChatCompletion from openai.types.chat.chat_completion_message_tool_call import ( ChatCompletionMessageToolCall, ) from openai.types.chat.chat_completion_tool_param import ChatCompletionToolParam -import os -from dotenv import load_dotenv from pydantic import BaseModel -from typing import Literal, Optional, List +from restack_ai.function import FunctionFailure, function, log load_dotenv() @@ -16,42 +17,46 @@ class Message(BaseModel): role: Literal["system", "user", "assistant", "tool"] content: str - tool_call_id: Optional[str] = None - tool_calls: Optional[List[ChatCompletionMessageToolCall]] = None + tool_call_id: str | None = None + tool_calls: list[ChatCompletionMessageToolCall] | None = None class LlmChatInput(BaseModel): - system_content: Optional[str] = None - model: Optional[str] = None - messages: Optional[List[Message]] = None - tools: Optional[List[ChatCompletionToolParam]] = None + system_content: str | None = None + model: str | None = None + messages: list[Message] | None = None + tools: list[ChatCompletionToolParam] | None = None + + +def raise_exception(message: str) -> None: + log.error(message) + raise FunctionFailure(message, non_retryable=True) @function.defn() -async def llm_chat(input: LlmChatInput) -> ChatCompletion: +async def llm_chat(function_input: LlmChatInput) -> ChatCompletion: try: - log.info("llm_chat function started", input=input) + log.info("llm_chat function started", function_input=function_input) + + if os.environ.get("RESTACK_API_KEY") is None: + raise_exception("RESTACK_API_KEY is not set") - 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") ) - log.info("pydantic_function_tool", tools=input.tools) + log.info("pydantic_function_tool", tools=function_input.tools) - if input.system_content: - input.messages.append( - Message(role="system", content=input.system_content or "") + if function_input.system_content: + function_input.messages.append( + Message(role="system", content=function_input.system_content or "") ) - response = client.chat.completions.create( - model=input.model or "gpt-4o-mini", - messages=input.messages, - tools=input.tools, + return client.chat.completions.create( + model=function_input.model or "gpt-4o-mini", + messages=function_input.messages, + tools=function_input.tools, ) - return response except Exception as e: log.error("llm_chat function failed", error=e) - raise e + raise diff --git a/agent_tool/src/functions/lookup_sales.py b/agent_tool/src/functions/lookup_sales.py index 81882b2d..f65704eb 100644 --- a/agent_tool/src/functions/lookup_sales.py +++ b/agent_tool/src/functions/lookup_sales.py @@ -1,6 +1,7 @@ from typing import Literal -from restack_ai.function import function, log + from pydantic import BaseModel +from restack_ai.function import function, log class SalesItem(BaseModel): @@ -21,9 +22,9 @@ class LookupSalesOutput(BaseModel): @function.defn() -async def lookupSales(input: LookupSalesInput) -> LookupSalesOutput: +async def lookup_sales(function_input: LookupSalesInput) -> LookupSalesOutput: try: - log.info("lookupSales function started", input=input) + log.info("lookup_sales function started", function_input=function_input) items = [ SalesItem( @@ -92,15 +93,17 @@ async def lookupSales(input: LookupSalesInput) -> LookupSalesOutput: ), ] - if input.category == "any": + if function_input.category == "any": filtered_items = items else: - filtered_items = [item for item in items if item.type == input.category] + filtered_items = [ + item for item in items if item.type == function_input.category + ] # Sort by largest discount first filtered_items.sort(key=lambda x: x.sale_discount_pct, reverse=True) return LookupSalesOutput(sales=filtered_items) except Exception as e: - log.error("lookupSales function failed", error=e) - raise e + log.error("lookup_sales function failed", error=e) + raise diff --git a/agent_tool/src/functions/new_function.py b/agent_tool/src/functions/new_function.py index a253cfc6..166a7505 100644 --- a/agent_tool/src/functions/new_function.py +++ b/agent_tool/src/functions/new_function.py @@ -1,3 +1,5 @@ +# ruff: noqa: ERA001, TD002, TD003, FIX002 + ## Step 1: Code your own function to talk to a third party API or database # import os diff --git a/agent_tool/src/services.py b/agent_tool/src/services.py index b792e5e1..49b19451 100644 --- a/agent_tool/src/services.py +++ b/agent_tool/src/services.py @@ -1,17 +1,21 @@ +# ruff: noqa: ERA001 import asyncio -import os -from watchfiles import run_process +import logging import webbrowser +from pathlib import Path + +from watchfiles import run_process + +from src.agents.chat_tool_functions import AgentChatToolFunctions from src.client import client -from src.functions.lookup_sales import lookupSales from src.functions.llm_chat import llm_chat +from src.functions.lookup_sales import lookupSales -from src.agents.chat_tool_functions import AgentChatToolFunctions # Step 5: Import a new function to tool calling here # from src.functions.new_function import new_function, FunctionInput, FunctionOutput -async def main(): +async def main() -> None: await client.start_service( agents=[AgentChatToolFunctions], ## Step 6: Add your new function to the functions list -> functions=[lookupSales, llm_chat, new_function] @@ -19,16 +23,16 @@ async def main(): ) -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) From 657f0a99361b50e26845a38ba8d5292a4163f3b6 Mon Sep 17 00:00:00 2001 From: leawn Date: Mon, 17 Feb 2025 15:48:34 +0100 Subject: [PATCH 7/8] Update misspelling of lookupSales --- agent_tool/src/agents/chat_tool_functions.py | 10 +++++----- agent_tool/src/services.py | 4 ++-- 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/agent_tool/src/agents/chat_tool_functions.py b/agent_tool/src/agents/chat_tool_functions.py index e676ad5a..c8317228 100644 --- a/agent_tool/src/agents/chat_tool_functions.py +++ b/agent_tool/src/agents/chat_tool_functions.py @@ -8,7 +8,7 @@ from openai import pydantic_function_tool from src.functions.llm_chat import LlmChatInput, Message, llm_chat - from src.functions.lookup_sales import LookupSalesInput, lookupSales + from src.functions.lookup_sales import LookupSalesInput, lookup_sales # Step 2: Import your new function to the agent # from src.functions.new_function import new_function, FunctionInput, FunctionOutput @@ -34,7 +34,7 @@ async def message(self, message: MessageEvent) -> list[Message]: tools = [ pydantic_function_tool( model=LookupSalesInput, - name=lookupSales.__name__, + name=lookup_sales.__name__, description="Lookup sales for a given category", ), # Step 3 Add your new function to the tools list and adjust the system prompt @@ -77,7 +77,7 @@ async def message(self, message: MessageEvent) -> list[Message]: name = tool_call.function.name match name: - case lookupSales.__name__: + case lookup_sales.__name__: args = LookupSalesInput.model_validate_json( tool_call.function.arguments ) @@ -85,7 +85,7 @@ async def message(self, message: MessageEvent) -> list[Message]: log.info(f"calling {name} with args: {args}") result = await agent.step( - function=lookupSales, + function=lookup_sales, agent_input=LookupSalesInput(category=args.category), start_to_close_timeout=timedelta(seconds=120), ) @@ -138,5 +138,5 @@ async def end(self) -> EndEvent: return {"end": True} @agent.run - async def run(self) -> None: + async def run(self, agent_input: dict) -> None: await agent.condition(lambda: self.end) diff --git a/agent_tool/src/services.py b/agent_tool/src/services.py index 49b19451..6da67ec7 100644 --- a/agent_tool/src/services.py +++ b/agent_tool/src/services.py @@ -9,7 +9,7 @@ from src.agents.chat_tool_functions import AgentChatToolFunctions from src.client import client from src.functions.llm_chat import llm_chat -from src.functions.lookup_sales import lookupSales +from src.functions.lookup_sales import lookup_sales # Step 5: Import a new function to tool calling here # from src.functions.new_function import new_function, FunctionInput, FunctionOutput @@ -19,7 +19,7 @@ async def main() -> None: await client.start_service( agents=[AgentChatToolFunctions], ## Step 6: Add your new function to the functions list -> functions=[lookupSales, llm_chat, new_function] - functions=[lookupSales, llm_chat], + functions=[lookup_sales, llm_chat], ) From dd138e0bf0378529a8327de23fb602f7411eba6a Mon Sep 17 00:00:00 2001 From: leawn Date: Mon, 17 Feb 2025 15:53:43 +0100 Subject: [PATCH 8/8] Add mandatory agent_input arg --- agent_chat/src/agents/agent.py | 2 +- agent_rag/src/agents/chat_rag.py | 2 +- agent_todo/src/agents/agent_todo.py | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/agent_chat/src/agents/agent.py b/agent_chat/src/agents/agent.py index e2bac360..0aea47d0 100644 --- a/agent_chat/src/agents/agent.py +++ b/agent_chat/src/agents/agent.py @@ -40,5 +40,5 @@ async def end(self, end: EndEvent) -> EndEvent: return end @agent.run - async def run(self) -> None: + async def run(self, agent_input: dict) -> None: await agent.condition(lambda: self.end) diff --git a/agent_rag/src/agents/chat_rag.py b/agent_rag/src/agents/chat_rag.py index fad39edb..d9f2a405 100644 --- a/agent_rag/src/agents/chat_rag.py +++ b/agent_rag/src/agents/chat_rag.py @@ -59,5 +59,5 @@ async def end(self) -> EndEvent: return {"end": True} @agent.run - async def run(self) -> None: + async def run(self, agent_input: dict) -> None: await agent.condition(lambda: self.end) diff --git a/agent_todo/src/agents/agent_todo.py b/agent_todo/src/agents/agent_todo.py index 5dab7d30..5b352c86 100644 --- a/agent_todo/src/agents/agent_todo.py +++ b/agent_todo/src/agents/agent_todo.py @@ -166,5 +166,5 @@ async def end(self) -> EndEvent: return {"end": True} @agent.run - async def run(self) -> None: + async def run(self, agent_input: dict) -> None: await agent.condition(lambda: self.end)