Merge pull request #135 from restackio/gemini-2-0

Author: martinibach

community/gemini_generate/.env.Example community/gemini/.env.Example

@@ -0,0 +1,129 @@
+# Restack AI - Gemini Example for Autonomous Agents
+
+This example demonstrates how to build reliable and scalable autonomous AI agents using Google's Gemini API with Restack. It shows how to handle rate limits (10 requests per minute for free tier with 1500 requests per day limit, concurrency limit), concurrent API calls, and workflow orchestration effectively to create agents that can independently perform complex tasks.
+
+Example workflows in `src/workflows/`:
+- [Content Generation](src/workflows/generate_content.py): Basic natural language understanding
+- [Function Calling](src/workflows/function_call.py): Toolexecution  
+- [Multi-Function Calling](src/workflows/multi_function_call.py): Complex decision making
+- [Multi-Function Calling](src/workflows/multi_function_call_advanced.py): Complex decision making with tools, executed as restack functions (e.g. third party APIs) 
+- [Swarm](src/workflows/swarm.py): Parallel execution of multiple coordinated agents
+
+
+## Motivation
+
+When building AI applications with Gemini, you need to handle various production challenges:
+
+1. **API Rate Limits**: Gemini API has concurrent request limits that need to be managed across multiple workflows.
+
+2. **Reliability**: AI workflows need retry mechanisms and graceful failure handling for production use.
+
+3. **Scalability**: As your application grows to handle thousands of concurrent agents or requests, you need robust queue management and execution control.
+
+4. **Scheduling**: Coordinating and managing multiple workflows running at different frequencies requires robust scheduling capabilities.
+
+### How Restack Helps
+
+Restack provides built-in solutions for these challenges:
+
+1. **Automated Rate Limiting**: Define rate limits in service options and let Restack handle the queuing:
+```python
+client.start_service(
+    task_queue="gemini",
+    functions=[generate_content, function_call],
+    options=ServiceOptions(
+        rate_limit=3,  # Match Gemini's concurrent call limit
+        max_concurrent_function_runs=3
+    )
+)
+```
+
+2. **Built-in Retries**: Restack automatically handles retries and failure scenarios.
+
+3. **Queue Management**: Efficiently manages thousands of concurrent workflows while respecting API limits.
+
+4. **Scheduling**: Run workflows on various schedules (minutely, hourly, daily) or with custom cron expressions.
+
+## Prerequisites
+
+- Python 3.9 or higher
+- Poetry (for dependency management)
+- Docker (for running the Restack services)
+- Active [Google AI Studio](https://aistudio.google.com) account with API key
+
+## Usage
+
+1. Run Restack local engine with Docker:
+
+   ```bash
+   docker run -d --pull always --name restack -p 5233:5233 -p 6233:6233 -p 7233:7233 ghcr.io/restackio/restack:main
+   ```
+
+2. Open the web UI to see the workflows:
+
+   ```bash
+   http://localhost:5233
+   ```
+
+3. Clone this repository:
+
+   ```bash
+   git clone https://github.com/restackio/examples-python
+   cd examples-python/community/gemini
+
+4. Install dependencies using Poetry:
+
+   ```bash
+   poetry env use 3.12
+   ```
+
+   ```bash
+   poetry shell
+   ```
+
+   ```bash
+   poetry install
+   ```
+
+   ```bash
+   poetry env info # Optional: copy the interpreter path to use in your IDE (e.g. Cursor, VSCode, etc.)
+   ```
+
+5. Set `GEMINI_API_KEY` as an environment variable from [Google AI Studio](https://aistudio.google.com)
+
+   ```bash
+   export GEMINI_API_KEY=<your-api-key>
+   ```
+
+6. Run the services in dev:
+
+   ```bash
+   poetry run dev
+   ```
+
+   This will start the Restack service with the defined workflows and functions.
+
+8. Schedule workflows via the UI 
+
+All workflows are auto-exposed with an RestAPI endpoint, ready to be run or scheduled via api request, or directly from the UI during development.
+
+![UI RestAPI Endpoints](ui-restapi-endpoints.png)
+
+Run or schedule individual workflows. 
+
+![UI Run workflow](ui-run-workflow.png)
+
+Run or schedule a workflow, triggering many agents as childworkflows.
+
+![UI Schedule swarm](ui-schedule-swarm.png)
+
+
+
+## Project Structure
+
+- `src/`: Main source code directory
+  - `client.py`: Initializes the Restack client
+  - `functions/`: Contains function definitions
+  - `workflows/`: Contains workflow definitions
+  - `services.py`: Sets up and runs the Restack services
+- `schedule_workflow.py`: Example script to schedule and run a workflow

community/gemini_generate/.gitignore community/gemini/.gitignore

@@ -1,7 +1,7 @@
 [tool.poetry]
-name = "gemini_generate"
+name = "gemini"
 version = "0.0.1"
-description = "A simple example to generate content using Google Gemini"
+description = "A simple example to integrate Google Gemini into a Restack project"
 authors = [
     "Restack Team <[email protected]>",
 ]
@@ -10,9 +10,10 @@ packages = [{include = "src"}]
 
 [tool.poetry.dependencies]
 python = ">=3.10,<4.0"
-restack-ai = "^0.0.52"
-google-generativeai = "0.8.3"
+restack-ai = "^0.0.53"
+google-genai = "0.5.0"
 watchfiles = "^1.0.0"
+pydantic = "^2.10.5"
 
 [build-system]
 requires = ["poetry-core"]

community/gemini_generate/Dockerfile community/gemini/Dockerfile

@@ -1,20 +1,19 @@
 import asyncio
 import time
 from restack_ai import Restack
-from dataclasses import dataclass
+from pydantic import BaseModel
 
-@dataclass
-class InputParams:
-    user_content: str
+class InputParams(BaseModel):
+    num_cities: int = 10
 
 async def main():
     client = Restack()
 
-    workflow_id = f"{int(time.time() * 1000)}-GeminiGenerateOppositeWorkflow"
+    workflow_id = f"{int(time.time() * 1000)}-GeminiSwarmWorkflow"
     runId = await client.schedule_workflow(
-        workflow_name="GeminiGenerateOppositeWorkflow",
+        workflow_name="GeminiSwarmWorkflow",
         workflow_id=workflow_id,
-        input=InputParams(user_content="The opposite of hot is")
+        input=InputParams(num_cities=10)
     )
 
     await client.get_workflow_result(

community/gemini/README.md

@@ -0,0 +1,36 @@
+from restack_ai.function import function, log
+from pydantic import BaseModel
+from google import genai
+from google.genai import types
+
+import os
+
+@function.defn()
+def get_current_weather(location: str) -> str:
+    """Returns the current weather.
+
+    Args:
+        location: The city and state, e.g. San Francisco, CA
+    """
+    log.info("get_current_weather function started", location=location)
+    return 'sunny'
+
+class FunctionInputParams(BaseModel):
+    user_content: str
+
+@function.defn()
+async def gemini_function_call(input: FunctionInputParams) -> str:
+    try:
+        log.info("gemini_function_call function started", input=input)
+        client = genai.Client(api_key=os.environ.get("GEMINI_API_KEY"))
+        
+        response = client.models.generate_content(
+            model='gemini-2.0-flash-exp',
+            contents=input.user_content,
+            config=types.GenerateContentConfig(tools=[get_current_weather])
+        )
+        log.info("gemini_function_call function completed", response=response.text)
+        return response.text
+    except Exception as e:
+        log.error("gemini_function_call function failed", error=e)
+        raise e

community/gemini_generate/pyproject.toml community/gemini/pyproject.toml

@@ -0,0 +1,25 @@
+from restack_ai.function import function, log
+from pydantic import BaseModel
+from google import genai
+from google.genai import types
+
+import os
+
+class FunctionInputParams(BaseModel):
+    user_content: str
+
+@function.defn()
+async def gemini_generate_content(input: FunctionInputParams) -> str:
+    try:
+        log.info("gemini_generate_content function started", input=input)
+        client = genai.Client(api_key=os.environ.get("GEMINI_API_KEY"))
+        
+        response = client.models.generate_content(
+            model='gemini-2.0-flash-exp',
+            contents=input.user_content
+        )
+        log.info("gemini_generate_content function completed", response=response.text)
+        return response.text
+    except Exception as e:
+        log.error("gemini_generate_content function failed", error=e)
+        raise e

community/gemini_generate/schedule_workflow.py community/gemini/schedule_workflow.py

@@ -0,0 +1,56 @@
+from restack_ai.function import function, log
+from pydantic import BaseModel
+from google import genai
+from google.genai import types
+
+import os
+
+@function.defn()
+def get_current_weather(location: str) -> str:
+    """Returns the current weather.
+
+    Args:
+        location: The city and state, e.g. San Francisco, CA
+    """
+    log.info("get_current_weather function started", location=location)
+    return 'sunny'
+
+@function.defn()
+def get_humidity(location: str) -> str:
+    """Returns the current humidity.
+
+    Args:
+        location: The city and state, e.g. San Francisco, CA
+    """
+    log.info("get_humidity function started", location=location)
+    return '65%'
+
+@function.defn()
+def get_air_quality(location: str) -> str:
+    """Returns the current air quality.
+
+    Args:
+        location: The city and state, e.g. San Francisco, CA
+    """
+    log.info("get_air_quality function started", location=location)
+    return 'good'
+
+class FunctionInputParams(BaseModel):
+    user_content: str
+
+@function.defn()
+async def gemini_multi_function_call(input: FunctionInputParams) -> str:
+    try:
+        log.info("gemini_multi_function_call function started", input=input)
+        client = genai.Client(api_key=os.environ.get("GEMINI_API_KEY"))
+        
+        response = client.models.generate_content(
+            model='gemini-2.0-flash-exp',
+            contents=input.user_content,
+            config=types.GenerateContentConfig(tools=[get_current_weather, get_humidity, get_air_quality])
+        )
+        log.info("gemini_multi_function_call function completed", response=response.text)
+        return response.text
+    except Exception as e:
+        log.error("gemini_multi_function_call function failed", error=e)
+        raise e

community/gemini_generate/src/__init__.py community/gemini/src/__init__.py

@@ -0,0 +1,38 @@
+from restack_ai.function import function, log
+from pydantic import BaseModel
+from google import genai
+from google.genai import types
+from typing import List, Optional
+
+import os
+from src.functions.tools import get_function_declarations
+
+class ChatMessage(BaseModel):
+    role: str
+    content: str
+
+class FunctionInputParams(BaseModel):
+    user_content: str
+    chat_history: Optional[List[ChatMessage]] = None
+
+@function.defn()
+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 = get_function_declarations()
+        tools = [types.Tool(function_declarations=functions)]
+
+        response = client.models.generate_content(
+            model='gemini-2.0-flash-exp',
+            contents=[input.user_content] + ([msg.content for msg in input.chat_history] if input.chat_history else []),
+            config=types.GenerateContentConfig(
+                tools=tools
+            )
+        )
+        return response
+    
+    except Exception as e:
+        log.error("Error in gemini_multi_function_call_advanced", error=str(e))
+        raise e