Vertex with openai

.changeset/witty-months-train.md

@@ -0,0 +1,5 @@
+---
+"livekit-plugins-openai": patch
+---
+
+vertex ai support with openai library

examples/voice-pipeline-agent/function_calling_weather.py

@@ -67,6 +67,8 @@ async def entrypoint(ctx: JobContext):
         vad=ctx.proc.userdata["vad"],
         stt=deepgram.STT(),
         llm=openai.LLM(),
+        # To use vertex AI LLM
+        # llm=openai.LLM.with_vertexai(project_id="your-project-id", location="us-central1"),
         tts=openai.TTS(),
         fnc_ctx=fnc_ctx,
         chat_ctx=initial_chat_ctx,

examples/voice-pipeline-agent/minimal_assistant.py

@@ -48,6 +48,8 @@ async def entrypoint(ctx: JobContext):
         vad=ctx.proc.userdata["vad"],
         stt=deepgram.STT(model=dg_model),
         llm=openai.LLM(),
+        # To use vertex AI LLM
+        # llm=openai.LLM.with_vertexai(project_id="your-project-id", location="us-central1"),
         tts=openai.TTS(),
         chat_ctx=initial_ctx,
     )

livekit-plugins/livekit-plugins-openai/livekit/plugins/openai/llm.py

@@ -37,6 +37,7 @@
     CerebrasChatModels,
     ChatModels,
     DeepSeekChatModels,
+    GeminiModels,
     GroqChatModels,
     OctoChatModels,
     PerplexityChatModels,
@@ -161,6 +162,54 @@ def with_cerebras(
             temperature=temperature,
         )
 
+    @staticmethod
+    def with_vertexai(
+        *,
+        model: str | GeminiModels = "gemini-1.5-flash-002",
+        project_id: str | None = None,
+        location: str | None = "us-central1",
+        client: openai.AsyncClient | None = None,
+        user: str | None = None,
+        temperature: float | None = None,
+    ) -> LLM:
+        """
+        Create a new instance of VertexAI LLM.
+
+        ``project_id`` must be set to your VERTEXAI PROJECT ID, either using the argument or by setting
+        the ``VERTEXAI_PROJECT_ID`` environmental variable.
+        """
+
+        project_id = project_id or os.environ.get("VERTEXAI_PROJECT_ID")
+        if project_id is None:
+            raise ValueError(
+                "VERTEXAI_PROJECT_ID is required, either set project_id argument or set VERTEXAI_PROJECT_ID environmental variable"
+            )
+        location = location or os.environ.get("VERTEXAI_LOCATION")
+        if location is None:
+            raise ValueError(
+                "VERTEXAI_LOCATION is required, either set location argument or set VERTEXAI_LOCATION environmental variable"
+            )
+
+        from google.auth import default
+        from google.auth.transport import requests
+
+        credentials, _ = default(
+            scopes=["https://www.googleapis.com/auth/cloud-platform"]
+        )
+        auth_request = requests.Request()
+        credentials.refresh(auth_request)
+        base_url = f"https://{location}-aiplatform.googleapis.com/v1beta1/projects/{project_id}/locations/{location}/endpoints/openapi"
+        api_key = credentials.token
+
+        return LLM(
+            model=model,
+            api_key=api_key,
+            base_url=base_url,
+            client=client,
+            user=user,
+            temperature=temperature,
+        )
+
     @staticmethod
     def with_fireworks(
         *,
@@ -492,6 +541,7 @@ def chat(
             temperature = self._opts.temperature
 
         messages = _build_oai_context(chat_ctx, id(self))
+        logger.info(f"messages: {messages}")
 
         cmp = self._client.chat.completions.create(
             messages=messages,
@@ -524,6 +574,7 @@ def __init__(
         self._tool_call_id: str | None = None
         self._fnc_name: str | None = None
         self._fnc_raw_arguments: str | None = None
+        self._tool_index: int | None = None
 
     async def _main_task(self) -> None:
         if not self._oai_stream:
@@ -577,10 +628,11 @@ def _parse_choice(self, id: str, choice: Choice) -> llm.ChatChunk | None:
                     continue  # oai may add other tools in the future
 
                 call_chunk = None
-                if self._tool_call_id and tool.id and tool.id != self._tool_call_id:
+                if self._tool_call_id and tool.id and tool.index != self._tool_index:
                     call_chunk = self._try_build_function(id, choice)
 
                 if tool.function.name:
+                    self._tool_index = tool.index
                     self._tool_call_id = tool.id
                     self._fnc_name = tool.function.name
                     self._fnc_raw_arguments = tool.function.arguments or ""

livekit-plugins/livekit-plugins-openai/livekit/plugins/openai/models.py

@@ -79,6 +79,20 @@
     "deepseek-chat",
 ]
 
+GeminiModels = Literal[
+    "gemini-1.0-pro",
+    "gemini-1.0-pro-vision",
+    "gemini-1.0-pro-vision-001",
+    "gemini-1.5-flash",
+    "gemini-1.5-flash-002",
+    "gemini-1.5-flash-8b",
+    "gemini-1.5-flash-preview-0514",
+    "gemini-1.5-pro",
+    "gemini-1.5-pro-002",
+    "gemini-1.5-pro-preview-0409",
+    "gemini-1.5-pro-preview-0514",
+]
+
 TogetherChatModels = Literal[
     "Austism/chronos-hermes-13b",
     "Gryphe/MythoMax-L2-13b",

livekit-plugins/livekit-plugins-openai/setup.py

@@ -47,7 +47,11 @@
     license="Apache-2.0",
     packages=setuptools.find_namespace_packages(include=["livekit.*"]),
     python_requires=">=3.9.0",
-    install_requires=["livekit-agents[codecs, images]>=0.11", "openai>=1.50"],
+    install_requires=[
+        "livekit-agents[codecs, images]>=0.11",
+        "openai>=1.50",
+        "google-auth",
+    ],
     package_data={"livekit.plugins.openai": ["py.typed"]},
     project_urls={
         "Documentation": "https://docs.livekit.io",