Initial commit

.github/README.md

@@ -0,0 +1,4 @@
+# Gemini-Tooluse-Prototype
+A prototype of a Gemini 1.5 Flash chatbot with Tooluse using a demo API implemented in Flask. It can use a calculator to perform basic arithmetic, get the current date and time, and search the web (using Selenium). Built in Python using Gradio and Flask
+
+https://huggingface.co/spaces/xarical/Gemini-Tooluse-Prototype

.github/workflows/checkfilesizes.yml

@@ -0,0 +1,16 @@
+name: Check file sizes
+on:               # or directly `on: [push]` to run the action on every push on any branch
+  pull_request:
+    branches: [main]
+
+  # to run this workflow manually from the Actions tab
+  workflow_dispatch:
+
+jobs:
+  check-file-size:
+    runs-on: ubuntu-latest
+    steps:
+      - name: Check file sizes
+        uses: ActionsDesk/lfs-warning@v2.0
+        with:
+          filesizelimit: 10485760 # this is 10MB so we can sync to HF Spaces

.github/workflows/deploytospace.yml

@@ -0,0 +1,21 @@
+name: Deploy to HF Space
+on:
+  push:
+    branches: [main]
+
+  # to run this workflow manually from the Actions tab
+  workflow_dispatch:
+
+jobs:
+  deploy-to-space:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v3
+        with:
+          fetch-depth: 0
+          lfs: true
+      - name: Push to HF space
+        env:
+          HF_API_KEY: ${{ secrets.HF_API_KEY }}
+          SPACE_ID: ${{ secrets.SPACE_ID }}
+        run: git push https://xarical:$HF_API_KEY@huggingface.co/spaces/$SPACE_ID main

Dockerfile

@@ -0,0 +1,49 @@
+# Read the doc: https://huggingface.co/docs/hub/spaces-sdks-docker
+# you will also find guides on how best to write your Dockerfile
+FROM python:3
+
+# Add new user 'user' (non-root)
+RUN useradd -m -u 1000 user
+
+# Set working dir to app
+ENV HOME=/home/user \
+    PATH=/home/user/.local/bin:$PATH
+RUN mkdir $HOME/app
+WORKDIR $HOME/app
+
+# Switch to root
+USER root
+
+# Install nginx and packages.txt
+COPY --chown=root packages.txt packages.txt
+RUN apt-get -y update && apt-get -y install nginx && xargs apt-get -y install < packages.txt
+
+# Give app permissions to 'user' (non-root)
+RUN chown user:user .
+
+# Give nginx permissions to 'user' (non-root)
+# See https://www.rockyourcode.com/run-docker-nginx-as-non-root-user/
+RUN mkdir -p /var/cache/nginx \
+             /var/log/nginx \
+             /var/lib/nginx
+RUN touch /var/run/nginx.pid
+RUN chown -R user:user /var/cache/nginx \
+                       /var/log/nginx \
+                       /var/lib/nginx \
+                       /var/run/nginx.pid
+
+# Switch to 'user' (non-root)
+USER user
+
+# Install requirements.txt
+COPY --chown=user requirements.txt requirements.txt
+RUN pip install --no-cache-dir --upgrade -r requirements.txt
+
+# Copy nginx configuration
+COPY --chown=user nginx.conf /etc/nginx/sites-available/default
+
+# Copy app
+COPY --chown=user . .
+
+# Run
+CMD ["bash", "run.sh"]

README.md

@@ -0,0 +1,12 @@
+---
+title: Gemini Tool Use
+emoji: 🛠️
+colorFrom: green
+colorTo: indigo
+sdk: docker
+app_port: 7860
+pinned: true
+short_description: A Gemini 1.5 Flash chatbot with Tool Use
+---
+
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,95 @@
+import json
+import logging
+import os
+from threading import Thread
+
+from flask import Flask, Response, abort, jsonify, request
+import google.generativeai as genai
+import gradio as gr
+import requests
+
+import config
+
+# get app logger
+app_logger = logging.getLogger(__name__)
+
+# config the model
+genai.configure(api_key=os.environ["GEMINI_API_KEY"])
+generation_config = {
+  "max_output_tokens": 512,
+  "temperature": 0.9
+}
+model = genai.GenerativeModel(model_name="gemini-1.5-flash",
+                              generation_config=generation_config)
+
+# handle the chatbot's messages
+def chatbot_response(user_input: str, history: list[dict[str, str]]) -> str:
+  messages = [{"role": "system", "content": config.system_prompt}] + history + [{"role": "user", "content": user_input}]
+  while True:
+    chat_completion = model.generate_content(str(messages))
+    chat_completion = chat_completion.text.replace('```json', '').replace('```', '').strip()
+    app_logger.info(chat_completion)
+    messages.append(chat_completion)
+    try:
+      response = json.loads(chat_completion)
+    except json.decoder.JSONDecodeError as e:
+      app_logger.error(e)
+      message = '{"error": "Your previous message was invalid JSON and caused an error during parsing. (Hint: you may have hit the token limit, try separating your messages into multiple messages)}'
+      app_logger.info(message)
+      messages.append({"role": "system", "content": message})
+      continue
+    if response["type"] == "tool":
+      handle_tools(response["content"], messages)
+    elif response["type"] == "answer":
+      return response["content"]
+
+# make the tool calls
+def handle_tools(tools: list[dict[str, str | dict]], messages: list[dict[str]]) -> None:
+  for tool in tools:
+    try:
+      if tool["name"] not in {"datetime", "calculator", "websearch"}:
+        response = f'{{"error": "Tool {tool["name"]} is an invalid tool"}}'
+      else:
+        url = f"{config.tools_external_url}/{tool['name']}"
+        data = tool.get("data")
+        app_logger.info(f"Making request to {url}{f' with {data}' if data else ''}")
+        r = requests.post(url, json=data) if data else requests.get(url)
+        r.raise_for_status()
+        response = r.json()
+    except Exception as e:
+      app_logger.error(e)
+      response = f'{{"error": "An error with making the API call to the tool {tool["name"]} has occurred, please inform the user of this"}}'
+    app_logger.info(f"Tool response: {response}")
+    messages.append({"role": "system", "content": response})
+
+# start the api server in a thread (with its own logger)
+api = Flask(__name__)
+formatter = logging.Formatter(f'%(asctime)s - API - %(levelname)s - %(message)s')
+api_logger = api.logger
+handler = api_logger.handlers[0].setFormatter(formatter)
+
+@api.route("/", methods=["POST"])
+def query() -> Response:
+  try:
+    data = json.loads(request.data)
+    return jsonify({"result": chatbot_response(data["user_input"], data["history"])})
+  except KeyError as e:
+    abort(400, description="Missing value in request body: " + str(e))
+  except Exception as e:
+    abort(400, description="Error: " + str(e))
+
+f = Thread(
+  target=api.run,
+  kwargs={
+    "host": config.api_host,
+    "port": config.api_port,
+  }
+)
+f.start()
+
+# start the gradio interface
+demo = gr.ChatInterface(fn=chatbot_response,
+  title="✨ Gemini Tooluse Prototype 🔨",
+  description="A prototype of a Gemini 1.5 Flash chatbot with Tooluse using a demo API implemented in Flask. It can use a calculator to perform basic arithmetic, get the current date and time, and search the web.")
+
+demo.launch(server_name=config.gradio_host, server_port=config.gradio_port, show_api=False)

config.py

@@ -0,0 +1,84 @@
+gradio_host = "localhost"
+gradio_port = 8080
+
+api_host = "localhost"
+api_port = 4000
+
+tools_host = "localhost"
+tools_port = 3000
+
+import os
+tools_external_url = f"https://{os.environ['SPACE_HOST']}/tools" # This should be changed based on where you're hosting
+
+system_prompt = """**Prompt:**
+You are a helpful assistant chatbot. You are talking to the user through a chat interface.
+Use the tools detailed in the Tool Specifications section below to augment your built-in ability to perform tasks. Use different tools together to answer questions. If the tool call returns unhelpful results or an error or if an applicable tool is not listed, answer the user's question to the best of your knowledge. Separate lists of more than 2 tools into separate messages and send those messages one at a time.
+Your response must be a valid JSON object similar to the 'assistant' messages in the following example conversation:
+**Example:**
+{
+  "role": "user",
+  "content": "Hi!"
+}
+{
+  "role": "assistant",
+  "type": "answer"
+  "content": "Hello! 👋 How can I assist you today?"
+}
+{
+  "role": "user",
+  "content": "What is the time in India?"
+}
+{
+  "role": "assistant",
+  "type": "tool",
+  "content": [
+    {"name": "datetime", "data": {}}, 
+    {"name": "websearch", "data": {"query": "Time difference between India and UTC"}}
+  ]
+}
+{
+  "role": "system", 
+  "content": {"date": "2024-08-13", "time": "9:47:13"}
+}
+{
+  "role": "system", 
+  "content": "Time Difference between UTC and IST Indian. Indian Standard Time is 5 hours 30 minutes ahead from the UTC universal time. UTC to IST Indian Time Conversion"
+}
+{
+  "role": "assistant",
+  "type": "tool",
+  "content": [
+    {"name": "calculator", "data": {"operator": "add", "num1": 9, "num2": 5}},
+    {"name": "calculator", "data": {"operator": "add", "num1": 47, "num2": 30}}
+  ]
+}
+{
+  "role": "system",
+  "content": "14"
+}
+{
+  "role": "system",
+  "content": "77"
+}
+{
+  "role": "assistant",
+  "type": "answer",
+  "content": "It's currently 15:17 in IST (Indian Standard Time). This is equivalent to 3:17 PM in IST. Is there anything else I can assist you with? 😊"
+}
+**Tool Specifications:**
+{
+  "name": "datetime"
+  "description": "A tool to get the current date and time in form YYYY-MM-DD and HH:MM:SS, in UTC with a 24-hour clock. It accepts nothing as the data."
+  "data": {}
+}
+{
+  "name": "calculator"
+  "description": "A simple calculator tool that can do the four arithmetic operations. It accepts an operator and two integers."
+  "data": {"operator": "string, add/subtract/multiply/divide", "num1": any int, "num2": any int}
+}
+{
+  "name": "websearch"
+  "description": "A search tool that can do a limited search of the web for answers. Results are not guaranteed to be accurate. IT DOES NOT WORK TO SEARCH ANYTHING TIME-SENSITIVE. It accepts a search query."
+  "data": {"query": "any string"}
+}
+"""

nginx.conf

@@ -0,0 +1,50 @@
+server {
+    listen 7860 default_server;
+    listen [::]:7860 default_server;
+
+    server_name _;
+
+    location / {
+        # Serve Gradio from port 8080
+        proxy_pass http://localhost:8080;
+        proxy_http_version 1.1;
+        proxy_set_header Upgrade $http_upgrade;
+        proxy_set_header Connection 'upgrade';
+        proxy_set_header Host $host;
+        proxy_set_header X-Real-IP $remote_addr;
+        proxy_set_header X-Forwarded-For $proxy_add_x_forwarded_for;
+        proxy_cache_bypass $http_upgrade;
+        proxy_read_timeout 86400;
+        proxy_redirect off;
+    }
+
+    location /query {
+        # Serve Flask app API server from port 4000
+        rewrite ^/query/?(.*)$ /$1 break; # strip the /query/
+        proxy_pass http://localhost:4000;
+        proxy_http_version 1.1;
+        proxy_set_header Upgrade $http_upgrade;
+        proxy_set_header Connection 'upgrade';
+        proxy_set_header Host $host;
+        proxy_set_header X-Real-IP $remote_addr;
+        proxy_set_header X-Forwarded-For $proxy_add_x_forwarded_for;
+        proxy_cache_bypass $http_upgrade;
+        proxy_read_timeout 86400;
+        proxy_redirect off;
+    }
+
+    location /tools {
+        # Serve Flask tools API server from port 3000
+        rewrite ^/tools/?(.*)$ /$1 break; # strip the /tools/
+        proxy_pass http://localhost:3000;
+        proxy_http_version 1.1;
+        proxy_set_header Upgrade $http_upgrade;
+        proxy_set_header Connection 'upgrade';
+        proxy_set_header Host $host;
+        proxy_set_header X-Real-IP $remote_addr;
+        proxy_set_header X-Forwarded-For $proxy_add_x_forwarded_for;
+        proxy_cache_bypass $http_upgrade;
+        proxy_read_timeout 86400;
+        proxy_redirect off;
+    }
+}

packages.txt

@@ -0,0 +1 @@
+chromium-driver

requirements.txt

@@ -0,0 +1,5 @@
+flask
+google-generativeai
+gradio
+requests
+selenium

run.sh

@@ -0,0 +1,12 @@
+#!/bin/bash
+
+# start nginx
+service nginx start
+
+# start the processes
+python tools.py > /dev/stdout 2>&1 & echo $! > tools.pid
+python app.py > /dev/stdout 2>&1 # blocking
+
+# when unblocked, kill other processes and clean up
+pkill -F tools.pid
+rm tools.pid

tools.py

@@ -0,0 +1,91 @@
+import logging
+import os
+import time
+from datetime import datetime, timezone
+
+from flask import Flask, Response, abort, json, jsonify, request
+from selenium import webdriver
+from selenium.webdriver.chrome.options import Options
+from selenium.webdriver.common.by import By
+from selenium.webdriver.common.keys import Keys
+
+import config
+
+# selenium options
+chrome_options = Options()
+chrome_options.add_argument('--no-sandbox')
+chrome_options.add_argument('--disable-dev-shm-usage')
+chrome_options.add_argument('--disable-gpu')
+chrome_options.add_argument('--headless')
+
+tools = Flask(__name__)
+formatter = logging.Formatter(f'%(asctime)s - TOOLS - %(levelname)s - %(message)s')
+tools_logger = tools.logger
+handler = tools_logger.handlers[0].setFormatter(formatter)
+
+@tools.route("/datetime", methods=["GET"])
+def get_datetime() -> Response:
+  curr_date = datetime.now(timezone.utc).date().strftime("%Y-%m-%d")
+  curr_time = datetime.now(timezone.utc).strftime("%H:%M:%S")
+  return jsonify({"date": curr_date, "time": curr_time})
+
+@tools.route("/calculator", methods=["POST"])
+def calculate() -> Response:
+  try:
+    data = json.loads(request.data) 
+    if data["operator"] == "add":
+      return jsonify(result=data["num1"] + data["num2"])
+    elif data["operator"] == "subtract":
+      return jsonify(result=data["num1"] - data["num2"])
+    elif data["operator"] == "multiply":
+      return jsonify(result=data["num1"] * data["num2"])
+    elif data["operator"] == "divide":
+      return jsonify(result=data["num1"] / data["num2"])
+    else:
+      abort(400, description="Invalid operator: " + data["operator"])
+  except KeyError as e:
+    abort(400, description="Missing value in request body: " + str(e))
+  except Exception as e:
+    abort(400, description="Error: " + str(e))
+
+@tools.route("/websearch", methods=["POST"])
+def google_search() -> Response:
+  try:
+    data = json.loads(request.data) 
+    global driver
+    driver = webdriver.Chrome(options=chrome_options)
+    driver.get("https://www.google.com/")
+    search_bar = driver.find_element(By.NAME, "q")
+    search_bar.send_keys(data["query"])
+    search_bar.send_keys(Keys.RETURN)
+    time.sleep(1)
+
+    search_results = driver.find_elements(By.CSS_SELECTOR, "div.kno-rdesc span span")
+    if len(search_results) > 0:  # check if google quick answer box exists
+      return jsonify(result=search_results[0].text)
+    else:  # otherwise, find list of search results
+      search_results = driver.find_elements(By.CSS_SELECTOR, "div.g")
+      for result in search_results:
+        try:  # first sentence of last span element in result
+          first_result = result.find_elements(By.CSS_SELECTOR, "div.VwiC3b span")
+          return jsonify(result=first_result[-1].text.split("...")[0])
+        except IndexError:  # if no span element in result, go to next result
+          pass
+      return jsonify(result="No search results found")
+  except KeyError as e:
+    abort(400, description="Missing value in request body: " + str(e))
+  except Exception as e:
+    abort(400, description="Error: " + str(e))
+  finally:
+    try:
+      # quit selenium and kill the vnc
+      driver.quit()
+      os.system("pkill -1 Xvnc")
+    except Exception:
+      pass
+
+@tools.route("/", methods=["GET"])
+def home() -> str:
+  return "hello"
+
+tools.run(host=config.tools_host, port=config.tools_port)