Support for VertexAI & Gemini API (#239)

* finalize vertexai

* abstract StreamWrapper

* adjust api file

* add gemini support

* quick fixes

* serialize prompts

* add audio to text

* standardaize completion

* Gemini API Instrumentation

* abstract event completion chunk

* Support for tools and functions

* add status code

* add vertexai and gemini to dev deps

* Minor fix

* Minor fix

---------

Co-authored-by: Karthik Kalyanaraman <[email protected]>

Author: alizenhom

pyproject.toml

@@ -50,7 +50,9 @@ dev = [
     "qdrant_client",
     "weaviate-client",
     "ollama",
-    "groq"
+    "groq",
+    "google-generativeai",
+    "google-cloud-aiplatform"
 ]
 
 test = [

src/examples/gemini_example/__init__.py

@@ -0,0 +1,6 @@
+from .main import basic
+
+
+class GeminiRunner:
+    def run(self):
+        basic()

src/examples/gemini_example/function_tools.py

@@ -0,0 +1,62 @@
+tools = [
+    {
+        "function_declarations": [
+            {
+                "name": "find_movies",
+                "description": "find movie titles currently playing in theaters based on any description, genre, title words, etc.",
+                "parameters": {
+                    "type": "object",
+                    "properties": {
+                        "location": {
+                            "type": "string",
+                            "description": "The city and state, e.g. San Francisco, CA or a zip code e.g. 95616",
+                        },
+                        "description": {
+                            "type": "string",
+                            "description": "Any kind of description including category or genre, title words, attributes, etc.",
+                        },
+                    },
+                    "required": ["description"],
+                },
+            },
+            {
+                "name": "find_theaters",
+                "description": "find theaters based on location and optionally movie title which is currently playing in theaters",
+                "parameters": {
+                    "type": "object",
+                    "properties": {
+                        "location": {
+                            "type": "string",
+                            "description": "The city and state, e.g. San Francisco, CA or a zip code e.g. 95616",
+                        },
+                        "movie": {"type": "string", "description": "Any movie title"},
+                    },
+                    "required": ["location"],
+                },
+            },
+            {
+                "name": "get_showtimes",
+                "description": "Find the start times for movies playing in a specific theater",
+                "parameters": {
+                    "type": "object",
+                    "properties": {
+                        "location": {
+                            "type": "string",
+                            "description": "The city and state, e.g. San Francisco, CA or a zip code e.g. 95616",
+                        },
+                        "movie": {"type": "string", "description": "Any movie title"},
+                        "theater": {
+                            "type": "string",
+                            "description": "Name of the theater",
+                        },
+                        "date": {
+                            "type": "string",
+                            "description": "Date for requested showtime",
+                        },
+                    },
+                    "required": ["location", "movie", "theater", "date"],
+                },
+            },
+        ]
+    }
+]

src/examples/gemini_example/main.py

@@ -0,0 +1,91 @@
+from langtrace_python_sdk import langtrace
+import google.generativeai as genai
+from dotenv import load_dotenv
+import os
+import asyncio
+import pathlib
+from .function_tools import tools
+
+load_dotenv()
+
+langtrace.init(write_spans_to_console=False, batch=False)
+genai.configure(api_key=os.environ["GEMINI_API_KEY"])
+
+
+async def async_demo():
+    task1 = asyncio.create_task(async_generate())
+    task2 = asyncio.create_task(async_generate(stream=True))
+    return await asyncio.gather(task1, task2)
+
+
+def basic():
+    generate()
+    generate(stream=True, with_tools=True)
+
+    # image_to_text()
+    # audio_to_text()
+    asyncio.run(async_demo())
+
+
+def generate(stream=False, with_tools=False):
+    model = genai.GenerativeModel(
+        "gemini-1.5-pro", system_instruction="You are a cat. Your name is Neko."
+    )
+
+    response = model.generate_content(
+        "Write a story about a AI and magic",
+        stream=stream,
+        tools=tools if with_tools else None,
+    )
+    if stream:
+        for res in response:
+            if res.text:
+                print(res.text)
+    else:
+        print(response.text)
+
+
+async def async_generate(stream=False):
+    model = genai.GenerativeModel(
+        "gemini-1.5-pro", system_instruction="You are a cat. Your name is Neko."
+    )
+    response = await model.generate_content_async(
+        "Write a story about a AI and magic", stream=stream
+    )
+    if stream:
+        async for chunk in response:
+            if chunk.text:
+                print(chunk.text)
+    else:
+        print(response.text)
+
+
+def image_to_text(stream=False):
+    model = genai.GenerativeModel("gemini-1.5-flash")
+    image1 = {
+        "mime_type": "image/jpeg",
+        "data": pathlib.Path("src/examples/gemini_example/jetpack.jpg").read_bytes(),
+    }
+
+    prompt = "Describe me this picture. What do you see in it."
+    response = model.generate_content([prompt, image1], stream=stream)
+    if stream:
+        for res in response:
+            print(res.text)
+    else:
+        print(response.text)
+
+
+# def audio_to_text(stream=False):
+#     model = genai.GenerativeModel("gemini-1.5-flash")
+#     audio = genai.upload_file(
+#         pathlib.Path("src/examples/gemini_example/voice_note.mp3")
+#     )
+
+#     prompt = "Summarize this voice recording."
+#     response = model.generate_content([prompt, audio], stream=stream)
+#     if stream:
+#         for res in response:
+#             print(res.text)
+#     else:
+#         print(response.text)

src/examples/vertexai_example/__init__.py

@@ -0,0 +1,6 @@
+from .main import basic
+
+
+class VertexAIRunner:
+    def run(self):
+        basic()

src/examples/vertexai_example/main.py

@@ -0,0 +1,214 @@
+import vertexai
+import base64
+import asyncio
+import vertexai.preview.generative_models as generative_models
+from vertexai.language_models import ChatModel, InputOutputTextPair, TextGenerationModel
+from langtrace_python_sdk import langtrace
+from vertexai.generative_models import GenerativeModel, Part, FinishReason
+from dotenv import load_dotenv
+
+load_dotenv()
+
+langtrace.init(write_spans_to_console=True, batch=False)
+vertexai.init(project="model-palace-429011-f5", location="us-central1")
+
+
+def basic():
+    # chat()
+    # chat_streaming()
+    # streaming_prediction()
+    # asyncio.run(async_streaming_prediction())
+
+    generate()
+    generate(stream=True)
+
+    image_to_text()
+    image_to_text(stream=True)
+
+    video_to_text()
+    video_to_text(stream=True)
+
+    audio_to_text()
+    audio_to_text(stream=True)
+
+
+def chat():
+    """Chat Example with a Large Language Model"""
+
+    chat_model = ChatModel.from_pretrained("chat-bison")
+
+    parameters = {
+        "temperature": 0.8,
+        "max_output_tokens": 256,
+        "top_p": 0.95,
+        "top_k": 40,
+    }
+
+    chat = chat_model.start_chat(
+        context="My name is Miles. You are an astronomer, knowledgeable about the solar system.",
+        examples=[
+            InputOutputTextPair(
+                input_text="How many moons does Mars have?",
+                output_text="The planet Mars has two moons, Phobos and Deimos.",
+            ),
+        ],
+    )
+
+    response = chat.send_message(
+        message="How many planets are there in the solar system?", **parameters
+    )
+
+    return response
+
+
+def chat_streaming() -> str:
+    """Streaming Chat Example with a Large Language Model"""
+
+    chat_model = ChatModel.from_pretrained("chat-bison")
+
+    parameters = {
+        "temperature": 0.8,
+        "max_output_tokens": 256,
+        "top_p": 0.95,
+        "top_k": 40,
+    }
+
+    chat = chat_model.start_chat(
+        context="My name is Miles. You are an astronomer, knowledgeable about the solar system.",
+        examples=[
+            InputOutputTextPair(
+                input_text="How many moons does Mars have?",
+                output_text="The planet Mars has two moons, Phobos and Deimos.",
+            ),
+        ],
+    )
+
+    responses = chat.send_message_streaming(
+        message="How many planets are there in the solar system?", **parameters
+    )
+
+    result = [response for response in responses]
+    return result
+
+
+def streaming_prediction() -> str:
+    """Streaming Text Example with a Large Language Model"""
+
+    text_generation_model = TextGenerationModel.from_pretrained("text-bison")
+    parameters = {
+        "max_output_tokens": 256,
+        "top_p": 0.8,
+        "top_k": 40,
+    }
+    responses = text_generation_model.predict_streaming(
+        prompt="Give me ten interview questions for the role of program manager.",
+        **parameters,
+    )
+    result = [response for response in responses]
+    print(result)
+    return result
+
+
+async def async_streaming_prediction() -> str:
+    """Async Streaming Text Example with a Large Language Model"""
+
+    text_generation_model = TextGenerationModel.from_pretrained("text-bison")
+    parameters = {
+        "max_output_tokens": 256,
+        "top_p": 0.8,
+        "top_k": 40,
+    }
+
+    responses = text_generation_model.predict_streaming_async(
+        prompt="Give me ten interview questions for the role of program manager.",
+        **parameters,
+    )
+
+    result = [response async for response in responses]
+    print(result)
+    return result
+
+
+def generate(stream=False):
+    generation_config = {
+        "max_output_tokens": 8192,
+        "temperature": 1,
+        "top_p": 0.95,
+    }
+    model = GenerativeModel(
+        "gemini-experimental",
+    )
+
+    responses = model.generate_content(
+        ["I am a software engineer. I enjoy playing video games and reading"],
+        generation_config=generation_config,
+        stream=stream,
+    )
+
+    if stream:
+        for res in responses:
+            print(res.text)
+    else:
+        print(responses.text)
+
+
+def image_to_text(stream=False):
+    model = GenerativeModel(model_name="gemini-experimental")
+
+    response = model.generate_content(
+        [
+            Part.from_uri(
+                "gs://cloud-samples-data/generative-ai/image/scones.jpg",
+                mime_type="image/jpeg",
+            ),
+            "What is shown in this image?",
+        ],
+        stream=stream,
+    )
+    if stream:
+        for res in response:
+            print(res.text)
+    else:
+        print(response.text)
+
+
+def video_to_text(stream=False):
+    model = GenerativeModel(model_name="gemini-experimental")
+
+    prompt = """
+    Provide a description of the video.
+    The description should also contain anything important which people say in the video.
+    """
+
+    video_file_uri = "gs://cloud-samples-data/generative-ai/video/pixel8.mp4"
+    video_file = Part.from_uri(video_file_uri, mime_type="video/mp4")
+
+    contents = [video_file, prompt]
+    response = model.generate_content(contents, stream=stream)
+    if stream:
+        for res in response:
+            print(res.text)
+    else:
+        print(response.text)
+
+
+def audio_to_text(stream=False):
+    model = GenerativeModel(model_name="gemini-1.5-flash-001")
+
+    prompt = """
+    Please provide a summary for the audio.
+    Provide chapter titles, be concise and short, no need to provide chapter summaries.
+    Do not make up any information that is not part of the audio and do not be verbose.
+    """
+
+    audio_file_uri = "gs://cloud-samples-data/generative-ai/audio/pixel.mp3"
+    audio_file = Part.from_uri(audio_file_uri, mime_type="audio/mpeg")
+
+    contents = [audio_file, prompt]
+
+    response = model.generate_content(contents, stream=stream)
+    if stream:
+        for res in response:
+            print(res.text)
+    else:
+        print(response.text)