swarm

martinibach · martinibach · commit 717c06d286d9 · 2025-01-16T23:13:55.000-08:00
diff --git a/community/gemini/src/functions/tools.py b/community/gemini/src/functions/tools.py
@@ -1,37 +1,144 @@
 from restack_ai.function import function, log
-from pydantic import BaseModel
-from typing import List, Optional
+from pydantic import BaseModel, validator
+from typing import List, Optional, Dict
 import inspect
 
-class TemperatureInput(BaseModel):
-    """The city and state, e.g. San Francisco, CA"""
-    location: str
+from enum import Enum
 
-class HumidityInput(BaseModel):
-    """The city and state, e.g. San Francisco, CA"""
-    location: str
+class USTopCities(str, Enum):
+    NEW_YORK = "New York, NY"
+    LOS_ANGELES = "Los Angeles, CA" 
+    CHICAGO = "Chicago, IL"
+    HOUSTON = "Houston, TX"
+    PHOENIX = "Phoenix, AZ"
+    PHILADELPHIA = "Philadelphia, PA"
+    SAN_ANTONIO = "San Antonio, TX"
+    SAN_DIEGO = "San Diego, CA"
+    DALLAS = "Dallas, TX"
+    SAN_JOSE = "San Jose, CA"
+    AUSTIN = "Austin, TX"
+    JACKSONVILLE = "Jacksonville, FL"
+    FORT_WORTH = "Fort Worth, TX"
+    COLUMBUS = "Columbus, OH"
+    SAN_FRANCISCO = "San Francisco, CA"
+    CHARLOTTE = "Charlotte, NC"
+    INDIANAPOLIS = "Indianapolis, IN"
+    SEATTLE = "Seattle, WA"
+    DENVER = "Denver, CO"
+    WASHINGTON_DC = "Washington, DC"
+    BOSTON = "Boston, MA"
+    EL_PASO = "El Paso, TX"
+    DETROIT = "Detroit, MI"
+    NASHVILLE = "Nashville, TN"
+    PORTLAND = "Portland, OR"
+    MEMPHIS = "Memphis, TN"
+    OKLAHOMA_CITY = "Oklahoma City, OK"
+    LAS_VEGAS = "Las Vegas, NV"
+    LOUISVILLE = "Louisville, KY"
+    BALTIMORE = "Baltimore, MD"
+    MILWAUKEE = "Milwaukee, WI"
+    ALBUQUERQUE = "Albuquerque, NM"
+    TUCSON = "Tucson, AZ"
+    FRESNO = "Fresno, CA"
+    SACRAMENTO = "Sacramento, CA"
+    MESA = "Mesa, AZ"
+    KANSAS_CITY = "Kansas City, MO"
+    ATLANTA = "Atlanta, GA"
+    MIAMI = "Miami, FL"
+    COLORADO_SPRINGS = "Colorado Springs, CO"
+    RALEIGH = "Raleigh, NC"
+    OMAHA = "Omaha, NE"
+    LONG_BEACH = "Long Beach, CA"
+    VIRGINIA_BEACH = "Virginia Beach, VA"
+    OAKLAND = "Oakland, CA"
+    MINNEAPOLIS = "Minneapolis, MN"
+    TULSA = "Tulsa, OK"
+    ARLINGTON = "Arlington, TX"
+    TAMPA = "Tampa, FL"
+    NEW_ORLEANS = "New Orleans, LA"
 
-class AirQualityInput(BaseModel):
+class WeatherData(BaseModel):
+    temperature: str
+    humidity: str
+    air_quality: str
+
+class LocationInput(BaseModel):
     """The city and state, e.g. San Francisco, CA"""
-    location: str
+    location: USTopCities
+
+CITY_WEATHER_DATA: Dict[USTopCities, WeatherData] = {
+    USTopCities.NEW_YORK: WeatherData(temperature="72°F", humidity="60%", air_quality="moderate"),
+    USTopCities.LOS_ANGELES: WeatherData(temperature="75°F", humidity="65%", air_quality="good"),
+    USTopCities.CHICAGO: WeatherData(temperature="68°F", humidity="55%", air_quality="good"), 
+    USTopCities.HOUSTON: WeatherData(temperature="82°F", humidity="75%", air_quality="moderate"),
+    USTopCities.PHOENIX: WeatherData(temperature="95°F", humidity="25%", air_quality="good"),
+    USTopCities.PHILADELPHIA: WeatherData(temperature="70°F", humidity="62%", air_quality="moderate"),
+    USTopCities.SAN_ANTONIO: WeatherData(temperature="85°F", humidity="70%", air_quality="good"),
+    USTopCities.SAN_DIEGO: WeatherData(temperature="72°F", humidity="68%", air_quality="good"),
+    USTopCities.DALLAS: WeatherData(temperature="83°F", humidity="65%", air_quality="moderate"),
+    USTopCities.SAN_JOSE: WeatherData(temperature="73°F", humidity="60%", air_quality="good"),
+    USTopCities.AUSTIN: WeatherData(temperature="84°F", humidity="68%", air_quality="good"),
+    USTopCities.JACKSONVILLE: WeatherData(temperature="80°F", humidity="75%", air_quality="moderate"),
+    USTopCities.FORT_WORTH: WeatherData(temperature="83°F", humidity="65%", air_quality="moderate"),
+    USTopCities.COLUMBUS: WeatherData(temperature="71°F", humidity="63%", air_quality="good"),
+    USTopCities.SAN_FRANCISCO: WeatherData(temperature="65°F", humidity="75%", air_quality="good"),
+    USTopCities.CHARLOTTE: WeatherData(temperature="76°F", humidity="65%", air_quality="good"),
+    USTopCities.INDIANAPOLIS: WeatherData(temperature="72°F", humidity="64%", air_quality="moderate"),
+    USTopCities.SEATTLE: WeatherData(temperature="62°F", humidity="80%", air_quality="good"),
+    USTopCities.DENVER: WeatherData(temperature="70°F", humidity="45%", air_quality="good"),
+    USTopCities.WASHINGTON_DC: WeatherData(temperature="74°F", humidity="65%", air_quality="moderate"),
+    USTopCities.BOSTON: WeatherData(temperature="68°F", humidity="70%", air_quality="good"),
+    USTopCities.EL_PASO: WeatherData(temperature="88°F", humidity="30%", air_quality="good"),
+    USTopCities.DETROIT: WeatherData(temperature="69°F", humidity="65%", air_quality="moderate"),
+    USTopCities.NASHVILLE: WeatherData(temperature="77°F", humidity="68%", air_quality="good"),
+    USTopCities.PORTLAND: WeatherData(temperature="65°F", humidity="75%", air_quality="good"),
+    USTopCities.MEMPHIS: WeatherData(temperature="79°F", humidity="70%", air_quality="moderate"),
+    USTopCities.OKLAHOMA_CITY: WeatherData(temperature="80°F", humidity="60%", air_quality="good"),
+    USTopCities.LAS_VEGAS: WeatherData(temperature="92°F", humidity="25%", air_quality="good"),
+    USTopCities.LOUISVILLE: WeatherData(temperature="75°F", humidity="67%", air_quality="moderate"),
+    USTopCities.BALTIMORE: WeatherData(temperature="73°F", humidity="65%", air_quality="moderate"),
+    USTopCities.MILWAUKEE: WeatherData(temperature="65°F", humidity="70%", air_quality="good"),
+    USTopCities.ALBUQUERQUE: WeatherData(temperature="82°F", humidity="35%", air_quality="good"),
+    USTopCities.TUCSON: WeatherData(temperature="90°F", humidity="30%", air_quality="good"),
+    USTopCities.FRESNO: WeatherData(temperature="85°F", humidity="45%", air_quality="moderate"),
+    USTopCities.SACRAMENTO: WeatherData(temperature="80°F", humidity="55%", air_quality="moderate"),
+    USTopCities.MESA: WeatherData(temperature="93°F", humidity="25%", air_quality="good"),
+    USTopCities.KANSAS_CITY: WeatherData(temperature="75°F", humidity="65%", air_quality="good"),
+    USTopCities.ATLANTA: WeatherData(temperature="78°F", humidity="70%", air_quality="moderate"),
+    USTopCities.MIAMI: WeatherData(temperature="85°F", humidity="80%", air_quality="moderate"),
+    USTopCities.COLORADO_SPRINGS: WeatherData(temperature="68°F", humidity="45%", air_quality="good"),
+    USTopCities.RALEIGH: WeatherData(temperature="75°F", humidity="68%", air_quality="good"),
+    USTopCities.OMAHA: WeatherData(temperature="73°F", humidity="65%", air_quality="good"),
+    USTopCities.LONG_BEACH: WeatherData(temperature="74°F", humidity="70%", air_quality="moderate"),
+    USTopCities.VIRGINIA_BEACH: WeatherData(temperature="76°F", humidity="75%", air_quality="good"),
+    USTopCities.OAKLAND: WeatherData(temperature="68°F", humidity="75%", air_quality="good"),
+    USTopCities.MINNEAPOLIS: WeatherData(temperature="65°F", humidity="65%", air_quality="good"),
+    USTopCities.TULSA: WeatherData(temperature="78°F", humidity="65%", air_quality="good"),
+    USTopCities.ARLINGTON: WeatherData(temperature="83°F", humidity="65%", air_quality="moderate"),
+    USTopCities.TAMPA: WeatherData(temperature="83°F", humidity="75%", air_quality="moderate"),
+    USTopCities.NEW_ORLEANS: WeatherData(temperature="82°F", humidity="80%", air_quality="moderate")
+}
 
 @function.defn()
-async def get_current_temperature(input: TemperatureInput) -> str:
+async def get_current_temperature(input: LocationInput) -> str:
     description = "Get the current temperature for a specific location"
     log.info("get_current_temperature function started", location=input.location)
-    return '75°F'
+    weather_data = CITY_WEATHER_DATA.get(input.location, WeatherData(temperature="75°F", humidity="65%", air_quality="good"))
+    return weather_data.temperature
 
 @function.defn()
-async def get_humidity(input: HumidityInput) -> str:
+async def get_humidity(input: LocationInput) -> str:
     description = "Get the current humidity level for a specific location"
     log.info("get_humidity function started", location=input.location)
-    return '65%'
+    weather_data = CITY_WEATHER_DATA.get(input.location, WeatherData(temperature="75°F", humidity="65%", air_quality="good"))
+    return weather_data.humidity
 
 @function.defn()
-async def get_air_quality(input: AirQualityInput) -> str:
+async def get_air_quality(input: LocationInput) -> str:
     description = "Get the current air quality for a specific location"
     log.info("get_air_quality function started", location=input.location)
-    return 'good'
+    weather_data = CITY_WEATHER_DATA.get(input.location, WeatherData(temperature="75°F", humidity="65%", air_quality="good"))
+    return weather_data.air_quality
 
 def get_function_declarations():
     functions = []
@@ -45,12 +152,13 @@ def get_function_declarations():
             "parameters": {
                 "type": "OBJECT",
                 "properties": {
-                    field_name: {
+                    "location": {
                         "type": "STRING",
                         "description": input_type.__doc__,
-                    } for field_name in input_type.__fields__
+                        "enum": [city.value for city in USTopCities]
+                    }
                 },
-                "required": list(input_type.__fields__.keys())
+                "required": ["location"]
             }
         })
     return functions
diff --git a/community/gemini/src/services.py b/community/gemini/src/services.py
@@ -2,6 +2,7 @@
 from watchfiles import run_process
 import webbrowser
 import os
+from restack_ai.restack import ServiceOptions
 
 from src.client import client
 from src.functions.generate_content import gemini_generate_content
@@ -17,10 +18,33 @@
 from src.functions.multi_function_call_advanced import gemini_multi_function_call_advanced
 from src.functions.tools import get_current_temperature, get_humidity, get_air_quality
 
+from src.workflows.swarm import GeminiSwarmWorkflow
+
 async def main():
-    await client.start_service(
-        workflows= [GeminiGenerateContentWorkflow, GeminiFunctionCallWorkflow, GeminiMultiFunctionCallWorkflow, GeminiMultiFunctionCallAdvancedWorkflow],
-        functions= [gemini_generate_content, gemini_function_call, gemini_multi_function_call, gemini_multi_function_call_advanced, get_current_temperature, get_humidity, get_air_quality]
+    await asyncio.gather(
+        client.start_service(
+            workflows=[GeminiGenerateContentWorkflow, GeminiFunctionCallWorkflow, GeminiMultiFunctionCallWorkflow, GeminiMultiFunctionCallAdvancedWorkflow, GeminiSwarmWorkflow],
+            functions=[],
+            options=ServiceOptions(
+                max_concurrent_workflow_runs=1000
+            )
+        ),
+        client.start_service(
+            task_queue="tools",
+            functions=[get_current_temperature, get_humidity, get_air_quality],
+            options=ServiceOptions(
+                rate_limit=10,
+                max_concurrent_function_runs=10
+            )
+        ),
+        client.start_service(
+            task_queue="gemini",
+            functions=[gemini_generate_content, gemini_function_call, gemini_multi_function_call, gemini_multi_function_call_advanced],
+            options=ServiceOptions(
+                rate_limit=5,
+                max_concurrent_function_runs=3
+            )
+        )
     )
 
 def run_services():
diff --git a/community/gemini/src/workflows/function_call.py b/community/gemini/src/workflows/function_call.py
@@ -19,7 +19,8 @@ async def run(self, input: WorkflowInputParams):
             start_to_close_timeout=timedelta(seconds=120),
             retry_policy=RetryPolicy(
                 maximum_attempts=1
-            )
+            ),
+            task_queue="gemini"
         )
         log.info("GeminiFunctionCallWorkflow completed", result=result)
         return result
diff --git a/community/gemini/src/workflows/generate_content.py b/community/gemini/src/workflows/generate_content.py
@@ -19,7 +19,8 @@ async def run(self, input: WorkflowInputParams):
             start_to_close_timeout=timedelta(seconds=120),
             retry_policy=RetryPolicy(
                 maximum_attempts=1
-            )
+            ),
+            task_queue="gemini"
         )
         log.info("GeminiGenerateContentWorkflow completed", result=result)
         return result
diff --git a/community/gemini/src/workflows/multi_function_call.py b/community/gemini/src/workflows/multi_function_call.py
@@ -19,7 +19,8 @@ async def run(self, input: WorkflowInputParams):
             start_to_close_timeout=timedelta(seconds=120),
             retry_policy=RetryPolicy(
                 maximum_attempts=1
-            )
+            ),
+            task_queue="gemini"
         )
         log.info("GeminiMultiFunctionCallWorkflow completed", result=result)
         return result
diff --git a/community/gemini/src/workflows/multi_function_call_advanced.py b/community/gemini/src/workflows/multi_function_call_advanced.py
@@ -13,7 +13,7 @@
     )
     from src.functions.tools import get_current_temperature, get_humidity, get_air_quality
 
-class WorkflowInputParams(BaseModel):
+class MultiFunctionCallAdvancedInputParams(BaseModel):
     user_content: str = "What's the weather in San Francisco?"
 
 @workflow.defn()
@@ -22,7 +22,7 @@ def __init__(self):
         self.chat_history = []
 
     @workflow.run
-    async def run(self, input: WorkflowInputParams):
+    async def run(self, input: MultiFunctionCallAdvancedInputParams):
         log.info("GeminiMultiFunctionCallAdvancedWorkflow started", input=input)
         
         current_content = input.user_content
@@ -37,7 +37,8 @@ async def run(self, input: WorkflowInputParams):
                     chat_history=self.chat_history
                 ),
                 start_to_close_timeout=timedelta(seconds=120),
-                retry_policy=RetryPolicy(maximum_attempts=1)
+                retry_policy=RetryPolicy(maximum_attempts=2),
+                task_queue="gemini"
             )
 
             if not result or not isinstance(result, dict):
@@ -59,6 +60,7 @@ async def run(self, input: WorkflowInputParams):
                                 result = await workflow.step(
                                     globals()[function_name],
                                     func_call["args"],
+                                    task_queue="tools",
                                     retry_policy=RetryPolicy(maximum_attempts=1)
                                 )
                                 function_results.append(f"{function_name} result: {str(result)}")
diff --git a/community/gemini/src/workflows/swarm.py b/community/gemini/src/workflows/swarm.py
@@ -0,0 +1,34 @@
+from restack_ai.workflow import workflow, import_functions, log, RetryPolicy, workflow_info
+from pydantic import BaseModel
+from datetime import timedelta
+import asyncio
+from src.functions.tools import USTopCities
+from src.workflows.multi_function_call_advanced import GeminiMultiFunctionCallAdvancedWorkflow, MultiFunctionCallAdvancedInputParams        
+
+class WorkflowInputParams(BaseModel):
+    num_cities: int = 50
+
+@workflow.defn()
+class GeminiSwarmWorkflow:
+    @workflow.run
+    async def run(self, input: WorkflowInputParams):
+        parent_workflow_id = workflow_info().workflow_id
+        
+        # Get all available cities from USTopCities enum
+        all_cities = [city.value for city in USTopCities]
+        
+        # Take the first n cities based on input
+        selected_cities = all_cities[:input.num_cities]
+
+        results_tasks = await asyncio.gather(*[
+            workflow.child_execute(
+                GeminiMultiFunctionCallAdvancedWorkflow,
+                input=MultiFunctionCallAdvancedInputParams(user_content=f"What's the weather in {city}?"),
+                workflow_id=f"{parent_workflow_id}-child-{city.replace(', ', '-')}"
+            ) for city in selected_cities
+        ])
+
+        results = [{"city": city, "result": result} for city, result in zip(selected_cities, results_tasks)]
+        
+        log.info("GeminiSwarmWorkflow completed", results=results)
+        return results
