Gemini chat completion integration

docs/_toc.yml

@@ -39,7 +39,7 @@ parts:
 
   - caption: User Reference
     chapters:
-      - file: source/reference/evaql      
+      - file: source/reference/evaql
         title: Query Language
         sections:
           - file: source/reference/evaql/load_csv
@@ -65,13 +65,13 @@ parts:
 
       - file: source/reference/api
         title: Python API
-      
+
       - file: source/reference/rest_api
         title: REST API
 
       - file: source/reference/databases/index
         title: Data Sources
-        sections: 
+        sections:
           - file: source/reference/databases/postgres
           - file: source/reference/databases/sqlite
           - file: source/reference/databases/mysql
@@ -83,7 +83,7 @@ parts:
 
       - file: source/reference/vector_databases/index
         title: Vector Databases
-        sections: 
+        sections:
           - file: source/reference/vector_databases/faiss
           - file: source/reference/vector_databases/chromadb
           - file: source/reference/vector_databases/qdrant
@@ -106,9 +106,11 @@ parts:
           - file: source/reference/ai/hf
             title: Hugging Face
           - file: source/reference/ai/openai
-            title: OpenAI 
+            title: OpenAI
+          - file: source/reference/ai/gemini
+            title: Gemini
           - file: source/reference/ai/yolo
-            title: YOLO 
+            title: YOLO
           - file: source/reference/ai/stablediffusion
             title: Stable Diffusion
 
@@ -117,7 +119,7 @@ parts:
 
       - file: source/reference/optimizations
         title: Optimizations
-        
+
       # - file: source/reference/io
       #   title: IO Descriptors
 

docs/source/reference/ai/gemini.rst

@@ -0,0 +1,24 @@
+.. _gemini:
+
+Gemini Models
+=====================
+
+This section provides an overview of how you can use Gemini models in EvaDB.
+
+
+Chat Completion Functions
+-------------------------
+
+To create a chat completion function in EvaDB, use the following SQL command:
+
+.. code-block:: sql
+
+    CREATE FUNCTION IF NOT EXISTS GeminiChatCompletion
+    IMPL 'evadb/functions/gemini.py'
+    MODEL 'gemini-pro'
+
+EvaDB supports the following models for chat completion task:
+
+- "gemini-pro"
+
+The chat completion function can be composed in interesting ways with other functions. Gemini can be used similar to the `ChatGPT` function as shown in `Google Colab <https://colab.research.google.com/github/georgia-tech-db/evadb/blob/master/tutorials/08-chatgpt.ipynb>`_. as an example of combining chat completion task with caption extraction and video summarization models from Hugging Face and feeding it to chat completion to ask questions about the results.

evadb/functions/function_bootstrap_queries.py

@@ -197,6 +197,12 @@
     EvaDB_INSTALLATION_DIR
 )
 
+gemini_function_query = """CREATE FUNCTION IF NOT EXISTS Gemini
+        IMPL '{}/functions/gemini.py';
+        """.format(
+    EvaDB_INSTALLATION_DIR
+)
+
 yolo8n_query = """CREATE FUNCTION IF NOT EXISTS Yolo
             TYPE  ultralytics
             MODEL 'yolov8n.pt';
@@ -282,6 +288,7 @@ def init_builtin_functions(db: EvaDBDatabase, mode: str = "debug") -> None:
         Similarity_function_query,
         norfair_obj_tracker_query,
         chatgpt_function_query,
+        gemini_function_query,
         face_detection_function_query,
         # Mvit_function_query,
         Sift_function_query,

evadb/functions/gemini.py

@@ -0,0 +1,146 @@
+# coding=utf-8
+# Copyright 2018-2023 EvaDB
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import os
+
+import pandas as pd
+from retry import retry
+
+from evadb.catalog.catalog_type import NdArrayType
+from evadb.functions.abstract.abstract_function import AbstractFunction
+from evadb.functions.decorators.decorators import forward, setup
+from evadb.functions.decorators.io_descriptors.data_types import PandasDataframe
+from evadb.utils.generic_utils import try_to_import_gemini
+
+_VALID_CHAT_COMPLETION_MODEL = [
+    "gemini-pro",
+]
+
+
+class Gemini(AbstractFunction):
+    """
+    Arguments:
+        model (str) : ID of the Gemini model to use. Refer to '_VALID_CHAT_COMPLETION_MODEL' for a list of supported models.
+        temperature (float) : Sampling temperature to use in the model. Higher value results in a more random output.
+
+    Input Signatures:
+        query (str)   : The task / question that the user wants the model to accomplish / respond.
+        content (str) : Any relevant context that the model can use to complete its tasks and generate the response.
+
+    Output Signatures:
+        response (str) : Contains the response generated by the model based on user input. Any errors encountered
+                         will also be passed in the response.
+
+    Example Usage:
+        Assume we have the transcripts for a few videos stored in a table 'video_transcripts' in a column named 'text'.
+        If the user wants to retrieve the summary of each video, the Gemini function can be used as:
+
+            query = "Generate the summary of the video"
+            cursor.table("video_transcripts").select(f"Gemini({query}, text)")
+
+        In the above function invocation, the 'query' passed would be the user task to generate video summaries, and the
+        'content' passed would be the video transcripts that need to be used in order to generate the summary. Since
+        no prompt is passed, the default system prompt will be used.
+
+        Now assume the user wants to create the video summary in 50 words and in French. Instead of passing these instructions
+        along with each query, a prompt can be set as such:
+
+            prompt = "Generate your responses in 50 words or less. Also, generate the response in French."
+            cursor.table("video_transcripts").select(f"Gemini({query}, text, {prompt})")
+
+        In the above invocation, an additional argument is passed as prompt. While the query and content arguments remain
+        the same, the 'prompt' argument will be set as a system message in model params.
+
+        Both of the above cases would generate a summary for each row / video transcript of the table in the response.
+    """
+
+    @property
+    def name(self) -> str:
+        return "Gemini"
+
+    @setup(cacheable=True, function_type="chat-completion", batchable=True)
+    def setup(
+        self,
+        model="gemini-pro",
+        temperature: float = 0,
+        gemini_api_key="",
+    ) -> None:
+        assert model in _VALID_CHAT_COMPLETION_MODEL, f"Unsupported Gemini {model}"
+        self.model = model
+        self.temperature = temperature
+        self.gemini_api_key = gemini_api_key
+
+    @forward(
+        input_signatures=[
+            PandasDataframe(
+                columns=["query", "content"],
+                column_types=[
+                    NdArrayType.STR,
+                    NdArrayType.STR,
+                ],
+                column_shapes=[(1,), (1,)],
+            )
+        ],
+        output_signatures=[
+            PandasDataframe(
+                columns=["response"],
+                column_types=[
+                    NdArrayType.STR,
+                ],
+                column_shapes=[(1,)],
+            )
+        ],
+    )
+    def forward(self, text_df):
+        try_to_import_gemini()
+        import google.generativeai as genai
+
+        api_key = self.gemini_api_key
+        if len(self.gemini_api_key) == 0:
+            api_key = os.environ.get("GEMINI_API_KEY", "")
+        assert (
+            len(api_key) != 0
+        ), "Please set your API key using SET GEMINI_API_KEY = '' or environment variable (GEMINI_API_KEY)"
+
+        genai.configure(api_key=api_key)
+        model = genai.GenerativeModel(self.model)
+
+        @retry(tries=6, delay=20)
+        def completion_with_backoff(content):
+            return model.generate_content(
+                content, generation_config={"temperature": self.temperature}
+            )
+
+        queries = text_df[text_df.columns[0]]
+        content = text_df[text_df.columns[0]]
+        if len(text_df.columns) > 1:
+            queries = text_df.iloc[:, 0]
+            content = text_df.iloc[:, 1]
+
+        results = []
+        for query, content in zip(queries, content):
+            message = f"Here is some context: {content}\n. Now, complete the following task: {query}"
+
+            response = completion_with_backoff(message)
+            if len(response.parts) > 0:
+                results.append(response.text)
+            assert (
+                len(response.parts) != 0
+            ), "Gemini could not generate content for this prompt."
+
+        df = pd.DataFrame({"response": results})
+
+        return df

evadb/utils/generic_utils.py

@@ -518,6 +518,16 @@ def try_to_import_openai():
         )
 
 
+def try_to_import_gemini():
+    try:
+        import google.generativeai as genai  # noqa: F401
+    except ImportError:
+        raise ValueError(
+            """Could not import openai python package.
+                Please install them with `pip install google-generativeai`."""
+        )
+
+
 def try_to_import_langchain():
     try:
         import langchain  # noqa: F401

setup.py

@@ -81,6 +81,7 @@ def read(path, encoding="utf-8"):
     "protobuf",
     "bs4",
     "openai>=1.0",  # CHATGPT
+    "google-generativeai", #GEMINI
     "gpt4all",  # PRIVATE GPT
     "sentencepiece",  # TRANSFORMERS
 ]