multi_function_advanced

Author: martinibach

community/gemini/src/functions/multi_function_call_advanced.py

@@ -5,6 +5,7 @@
 from typing import List, Optional
 
 import os
+from src.functions.tools import get_function_declarations
 
 class ChatMessage(BaseModel):
     role: str
@@ -14,81 +15,13 @@ class FunctionInputParams(BaseModel):
     user_content: str
     chat_history: Optional[List[ChatMessage]] = None
 
-class WeatherInput(BaseModel):
-    location: str
-
-class HumidityInput(BaseModel):
-    location: str
-
-class AirQualityInput(BaseModel):
-    location: str
-
-@function.defn()
-async def get_current_weather(input: WeatherInput) -> str:
-    log.info("get_current_weather function started", location=input.location)
-    return 'sunny'
-
 @function.defn()
-async def get_humidity(input: HumidityInput) -> str:
-    log.info("get_humidity function started", location=input.location)
-    return '65%'
-
-@function.defn()
-async def get_air_quality(input: AirQualityInput) -> str:
-    log.info("get_air_quality function started", location=input.location)
-    return 'good'
-
-@function.defn()
-async def gemini_multi_function_call_advanced(input: FunctionInputParams) :
+async def gemini_multi_function_call_advanced(input: FunctionInputParams):
     try:
         log.info("gemini_multi_function_call_advanced function started", input=input)
         client = genai.Client(api_key=os.environ.get("GEMINI_API_KEY"))
 
-        functions = [
-            {
-                "name": "get_current_weather",
-                "description": "Get the current weather in a given location",
-                "parameters": {
-                    "type": "OBJECT",
-                    "properties": {
-                        "location": {
-                            "type": "STRING",
-                            "description": "The city and state, e.g. San Francisco, CA",
-                        },
-                    },
-                    "required": ["location"],
-                }
-            },
-            {
-                "name": "get_humidity", 
-                "description": "Get the current humidity in a given location",
-                "parameters": {
-                    "type": "OBJECT",
-                    "properties": {
-                        "location": {
-                            "type": "STRING",
-                            "description": "The city and state, e.g. San Francisco, CA",
-                        },
-                    },
-                    "required": ["location"],
-                }
-            },
-            {
-                "name": "get_air_quality",
-                "description": "Get the current air quality in a given location", 
-                "parameters": {
-                    "type": "OBJECT",
-                    "properties": {
-                        "location": {
-                            "type": "STRING",
-                            "description": "The city and state, e.g. San Francisco, CA",
-                        },
-                    },
-                    "required": ["location"],
-                }
-            }
-        ]
-
+        functions = get_function_declarations()
         tools = [types.Tool(function_declarations=functions)]
 
         response = client.models.generate_content(

community/gemini/src/functions/tools.py

@@ -0,0 +1,56 @@
+from restack_ai.function import function, log
+from pydantic import BaseModel
+from typing import List, Optional
+import inspect
+
+class TemperatureInput(BaseModel):
+    """The city and state, e.g. San Francisco, CA"""
+    location: str
+
+class HumidityInput(BaseModel):
+    """The city and state, e.g. San Francisco, CA"""
+    location: str
+
+class AirQualityInput(BaseModel):
+    """The city and state, e.g. San Francisco, CA"""
+    location: str
+
+@function.defn()
+async def get_current_temperature(input: TemperatureInput) -> str:
+    description = "Get the current temperature for a specific location"
+    log.info("get_current_temperature function started", location=input.location)
+    return '75°F'
+
+@function.defn()
+async def get_humidity(input: HumidityInput) -> str:
+    description = "Get the current humidity level for a specific location"
+    log.info("get_humidity function started", location=input.location)
+    return '65%'
+
+@function.defn()
+async def get_air_quality(input: AirQualityInput) -> str:
+    description = "Get the current air quality for a specific location"
+    log.info("get_air_quality function started", location=input.location)
+    return 'good'
+
+def get_function_declarations():
+    functions = []
+    for func in [get_current_temperature, get_humidity, get_air_quality]:
+        input_type = func.__annotations__['input']
+        source = inspect.getsource(func)
+        description = source.split('description = "')[1].split('"')[0]
+        functions.append({
+            "name": func.__name__,
+            "description": description,
+            "parameters": {
+                "type": "OBJECT",
+                "properties": {
+                    field_name: {
+                        "type": "STRING",
+                        "description": input_type.__doc__,
+                    } for field_name in input_type.__fields__
+                },
+                "required": list(input_type.__fields__.keys())
+            }
+        })
+    return functions

community/gemini/src/services.py

@@ -15,12 +15,12 @@
 
 from src.workflows.multi_function_call_advanced import GeminiMultiFunctionCallAdvancedWorkflow
 from src.functions.multi_function_call_advanced import gemini_multi_function_call_advanced
-from src.functions.multi_function_call_advanced import get_current_weather, get_humidity, get_air_quality
+from src.functions.tools import get_current_temperature, get_humidity, get_air_quality
 
 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_weather, get_humidity, get_air_quality]
+        functions= [gemini_generate_content, gemini_function_call, gemini_multi_function_call, gemini_multi_function_call_advanced, get_current_temperature, get_humidity, get_air_quality]
     )
 
 def run_services():

community/gemini/src/workflows/multi_function_call_advanced.py

@@ -8,15 +8,13 @@
 with import_functions():
     from src.functions.multi_function_call_advanced import (
         gemini_multi_function_call_advanced,
-        get_current_weather,
-        get_humidity,
-        get_air_quality,
         FunctionInputParams,
         ChatMessage
     )
+    from src.functions.tools import get_current_temperature, get_humidity, get_air_quality
 
 class WorkflowInputParams(BaseModel):
-    user_content: str = "what's the weather in San Francisco?"
+    user_content: str = "What's the weather in San Francisco?"
 
 @workflow.defn()
 class GeminiMultiFunctionCallAdvancedWorkflow:
@@ -56,7 +54,7 @@ async def run(self, input: WorkflowInputParams):
                         func_call = part["functionCall"]
                         function_name = func_call["name"]
 
-                        if function_name in {"get_current_weather", "get_humidity", "get_air_quality"}:
+                        if function_name in {"get_current_temperature", "get_humidity", "get_air_quality"}:
                             try:
                                 result = await workflow.step(
                                     globals()[function_name],