Gemini create pr example

examples/gemini_create_pr/.scanignore

@@ -0,0 +1,4 @@
+.scanignore
+*.lock
+*.pyc
+__pycache__/*

examples/gemini_create_pr/README.md

@@ -0,0 +1,70 @@
+# Restack AI SDK - Gemini Generate Content Example
+
+## 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
+- GitHub personal access token with repo permissions
+
+1. Run Restack local engine with Docker:
+   ```bash
+   docker run -d --pull always --name studio -p 5233:5233 -p 6233:6233 -p 7233:7233 ghcr.io/restackio/engine: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/examples/gemini_create_pr
+   ```
+
+4. Install dependencies using Poetry:
+   ```bash
+   poetry install
+   ```
+
+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. Set `GITHUB_TOKEN` as an environment variable from your GitHub account
+
+   ```bash
+   export GITHUB_TOKEN=<your-github-token>
+   ```
+
+7. Run the services:
+
+   ```bash
+   poetry run services
+   ```
+
+   This will start the Restack service with the defined workflows and functions.
+
+8. In a new terminal, schedule the workflow:
+
+   ```bash
+   python run_ui.py
+   ```
+
+   This will open a Streamlit UI to run the example.
+
+## Project Structure
+
+- `src/`: Main source code directory
+  - `client.py`: Initializes the Restack client
+  - `functions/`: Contains function definitions
+    - `util/`: Utility functions
+  - `workflows/`: Contains workflow definitions
+  - `ui/`: Streamlit UI files
+  - `services.py`: Sets up and runs the Restack services
+- `run_ui.py`: Example script to run the UI

examples/gemini_create_pr/pyproject.toml

@@ -0,0 +1,23 @@
+[tool.poetry]
+name = "gemini_create_pr"
+version = "0.0.1"
+description = "A simple example to create a PR using Google Gemini"
+authors = [
+    "Restack Team <service@restack.io>",
+]
+readme = "README.md"
+packages = [{include = "src"}]
+
+[tool.poetry.dependencies]
+python = ">=3.9,<3.9.7 || >3.9.7,<4.0"
+gitpython = "^3.1.43"
+streamlit = "^1.39.0"
+restack-ai = "^0.0.21"
+restack-google-gemini = "^0.0.23"
+
+[build-system]
+requires = ["poetry-core"]
+build-backend = "poetry.core.masonry.api"
+
+[tool.poetry.scripts]
+services = "src.services:run_services"

examples/gemini_create_pr/run_ui.py

@@ -0,0 +1,11 @@
+import os
+import sys
+import streamlit.web.cli as stcli
+
+# Add the project root to PYTHONPATH
+project_root = os.path.dirname(os.path.abspath(__file__))
+sys.path.insert(0, project_root)
+
+if __name__ == "__main__":
+    sys.argv = ["streamlit", "run", "src/ui/app.py"]
+    sys.exit(stcli.main())

examples/gemini_create_pr/src/client.py

@@ -0,0 +1,3 @@
+from restack_ai import Restack
+
+client = Restack()

examples/gemini_create_pr/src/functions/__init__.py

@@ -0,0 +1,13 @@
+from .repo_contents import scan_repository
+from .generate_solution import generate_solution
+from .make_changes import make_changes
+from .generate_pr_info import generate_pr_info
+from .create_pr import create_pr
+
+__all__ = [
+    "scan_repository",
+    "generate_solution",
+    "make_changes",
+    "generate_pr_info",
+    "create_pr"
+]

examples/gemini_create_pr/src/functions/create_pr.py

@@ -0,0 +1,59 @@
+from grpc import Status
+from restack_ai.function import function
+from src.functions.generate_pr_info import PrInfo
+from dataclasses import dataclass
+from git import Repo
+import os
+import re
+import requests
+
+@dataclass
+class FunctionInputParams:
+    repo_path: str
+    pr_info: PrInfo
+
+@function.defn(name="CreatePr")
+async def create_pr(input: FunctionInputParams):
+    """Creates a Pull Request for the changes in the given repository.
+
+    This function:
+    1. Creates a new branch
+    2. Commits all changes
+    3. Pushes to remote repository
+    4. Creates a PR via GitHub API
+
+    Args:
+        input (FunctionInputParams): Contains repo_path and pr_info (branch name, commit message, PR title)
+
+    Returns:
+        dict: Contains PR URL and status if successful
+        None: If no changes to commit or if PR creation fails
+    """
+    repo = Repo(input.repo_path)
+
+    if repo.is_dirty():
+        remote_url = repo.remotes.origin.url
+        repo_name = re.search(r'github\.com[:/](.*?)(?:\.git)?$', remote_url).group(1)
+
+        repo.git.checkout('HEAD', b=input.pr_info.branch_name)
+        repo.git.add(A=True)
+        repo.index.commit(input.pr_info.commit_message)
+        repo.remote("origin").push(input.pr_info.branch_name)
+
+        response = requests.post(
+            f"https://api.github.com/repos/{repo_name}/pulls",
+            headers={"Authorization": f"token {os.environ.get('GITHUB_TOKEN')}"},
+            json={
+                "title": input.pr_info.pr_title,
+                "body": "This PR was generated by Gemini using the Restack AI Python SDK",
+                "head": input.pr_info.branch_name,
+                "base": "main"
+            }
+        )
+        if response.status_code == 201:
+            pr_url = response.json()['html_url']
+            return {"pr_url": pr_url, "status": "success"}
+        else:
+            print(f"Failed to create PR: {response.json()}")
+    else:
+        print("No changes to commit")

examples/gemini_create_pr/src/functions/generate_pr_info.py

@@ -0,0 +1,98 @@
+from restack_ai.function import function
+from restack_google_gemini import (
+    gemini_get_history,
+    GeminiGetHistoryInput,
+    gemini_start_chat,
+    GeminiStartChatInput,
+    gemini_send_message,
+    GeminiSendMessageInput
+)
+from typing import Any
+from dataclasses import dataclass
+import os
+import json
+
+@dataclass
+class PrInfo:
+    branch_name: str
+    commit_message: str
+    pr_title: str
+
+@dataclass
+class GeminiGenerateContentResponse:
+    pr_info: PrInfo
+
+@dataclass
+class FunctionInputParams:
+    repo_path: str
+    chat_history: list[Any]
+    user_content: str | None = None
+
+@dataclass
+class FunctionOutputParams:
+    pr_info: PrInfo
+    chat_history: list[Any]
+
+
+@function.defn(name="GeneratePrInfo")
+async def generate_pr_info(input: FunctionInputParams):
+    chat_rule = """
+        Generate a PR information for the changes made to the repository using these specific guidelines:
+
+        1. branch_name:
+           - Use lowercase letters and hyphens only
+           - Start with a category (feat/fix/docs/refactor/test)
+           - Keep it concise but descriptive
+           - Example: feat-add-user-authentication
+
+        2. commit_message:
+           - Follow conventional commits format
+           - Start with type(scope): description
+           - Use present tense
+           - Be specific but concise
+           - Example: feat(auth): implement user authentication system
+
+        3. pr_title:
+           - Should be clear and descriptive
+           - Start with square brackets indicating the type [FEAT], [FIX], etc.
+           - Maximum 72 characters
+           - Example: [FEAT] Implement user authentication system
+
+        IMPORTANT:
+        - Ignore all copyrighted content
+        - Never use emojis in the response
+        - Keep all responses in plain text without any formatting
+        - Ensure each field is properly formatted as a valid JSON string
+    """
+
+    prompt = f"This is my feedback:\n{input.user_content}\n\n{chat_rule}" if input.user_content else chat_rule
+
+    gemini_chat = gemini_start_chat(
+        GeminiStartChatInput(
+            model="gemini-1.5-flash",
+            api_key=os.environ.get("GEMINI_API_KEY"),
+            generation_config={
+                "response_mime_type": "application/json",
+                "response_schema": GeminiGenerateContentResponse
+            },
+            history=input.chat_history
+        )
+    )
+
+    gemini_send_message_response = gemini_send_message(
+            GeminiSendMessageInput(
+                user_content=prompt,
+                chat=gemini_chat
+            )
+        )
+
+    gemini_get_history_response = gemini_get_history(
+        GeminiGetHistoryInput(
+            chat=gemini_chat
+        )
+    )
+
+    return FunctionOutputParams(
+        pr_info=json.loads(gemini_send_message_response.text)['pr_info'], 
+        chat_history=gemini_get_history_response
+    )