Unifying GenAI and VertexAI based on the shared Gemini REST API.

By extracting 'gemini.py' to deal with Gemini's REST API, we are able to share the most parts of Google GenAI and VertexAI. This CL establishes us to support new Gemini models/features with a single copy of changes. Also we now have all our LLMs implemented in REST, minimizing dependencies needed when using Langfun. PiperOrigin-RevId: 711915456
google · Jan 4, 2025 · 3af7310 · 3af7310
1 parent 8994735
commit 3af7310
Show file tree

Hide file tree

Showing 9 changed files with 801 additions and 1,080 deletions.
diff --git a/README.md b/README.md
@@ -138,9 +138,7 @@ If you want to customize your installation, you can select specific features usi
 | Tag                 |  Description                             |
 | ------------------- | ---------------------------------------- |
 | all                 | All Langfun features.                    |
-| llm                 | All supported LLMs.                      |
-| llm-google          | All supported Google-powered LLMs.       |
-| llm-google-genai    | LLMs powered by Google Generative AI API |
+| vertexai            | VertexAI access.                         |
 | mime                | All MIME supports.                       |
 | mime-auto           | Automatic MIME type detection.           |
 | mime-docx           | DocX format support.                     |
@@ -149,9 +147,9 @@ If you want to customize your installation, you can select specific features usi
 | ui                  | UI enhancements                          |
 
 
-For example, to install a nightly build that includes Google-powered LLMs, full modality support, and UI enhancements, use:
+For example, to install a nightly build that includes VertexAI access, full modality support, and UI enhancements, use:
 ```
-pip install langfun[llm-google,mime,ui] --pre
+pip install langfun[vertexai,mime,ui] --pre
 ```
 
 *Disclaimer: this is not an officially supported Google product.*
diff --git a/langfun/core/llms/__init__.py b/langfun/core/llms/__init__.py
@@ -32,16 +32,30 @@
 
 # Gemini models.
 from langfun.core.llms.google_genai import GenAI
-from langfun.core.llms.google_genai import GeminiFlash2_0ThinkingExp
+from langfun.core.llms.google_genai import GeminiFlash2_0ThinkingExp_20241219
 from langfun.core.llms.google_genai import GeminiFlash2_0Exp
-from langfun.core.llms.google_genai import GeminiExp_20241114
 from langfun.core.llms.google_genai import GeminiExp_20241206
-from langfun.core.llms.google_genai import GeminiFlash1_5
+from langfun.core.llms.google_genai import GeminiExp_20241114
 from langfun.core.llms.google_genai import GeminiPro1_5
-from langfun.core.llms.google_genai import GeminiPro
-from langfun.core.llms.google_genai import GeminiProVision
-from langfun.core.llms.google_genai import Palm2
-from langfun.core.llms.google_genai import Palm2_IT
+from langfun.core.llms.google_genai import GeminiPro1_5_002
+from langfun.core.llms.google_genai import GeminiPro1_5_001
+from langfun.core.llms.google_genai import GeminiFlash1_5
+from langfun.core.llms.google_genai import GeminiFlash1_5_002
+from langfun.core.llms.google_genai import GeminiFlash1_5_001
+from langfun.core.llms.google_genai import GeminiPro1
+
+from langfun.core.llms.vertexai import VertexAI
+from langfun.core.llms.vertexai import VertexAIGeminiFlash2_0ThinkingExp_20241219
+from langfun.core.llms.vertexai import VertexAIGeminiFlash2_0Exp
+from langfun.core.llms.vertexai import VertexAIGeminiExp_20241206
+from langfun.core.llms.vertexai import VertexAIGeminiExp_20241114
+from langfun.core.llms.vertexai import VertexAIGeminiPro1_5
+from langfun.core.llms.vertexai import VertexAIGeminiPro1_5_002
+from langfun.core.llms.vertexai import VertexAIGeminiPro1_5_001
+from langfun.core.llms.vertexai import VertexAIGeminiFlash1_5
+from langfun.core.llms.vertexai import VertexAIGeminiFlash1_5_002
+from langfun.core.llms.vertexai import VertexAIGeminiFlash1_5_001
+from langfun.core.llms.vertexai import VertexAIGeminiPro1
 
 # OpenAI models.
 from langfun.core.llms.openai import OpenAI
@@ -124,25 +138,6 @@
 from langfun.core.llms.groq import GroqWhisper_Large_v3
 from langfun.core.llms.groq import GroqWhisper_Large_v3Turbo
 
-from langfun.core.llms.vertexai import VertexAI
-from langfun.core.llms.vertexai import VertexAIGemini2_0
-from langfun.core.llms.vertexai import VertexAIGeminiFlash2_0Exp
-from langfun.core.llms.vertexai import VertexAIGeminiFlash2_0ThinkingExp
-from langfun.core.llms.vertexai import VertexAIGemini1_5
-from langfun.core.llms.vertexai import VertexAIGeminiPro1_5
-from langfun.core.llms.vertexai import VertexAIGeminiPro1_5_001
-from langfun.core.llms.vertexai import VertexAIGeminiPro1_5_002
-from langfun.core.llms.vertexai import VertexAIGeminiPro1_5_0514
-from langfun.core.llms.vertexai import VertexAIGeminiPro1_5_0409
-from langfun.core.llms.vertexai import VertexAIGeminiFlash1_5
-from langfun.core.llms.vertexai import VertexAIGeminiFlash1_5_001
-from langfun.core.llms.vertexai import VertexAIGeminiFlash1_5_002
-from langfun.core.llms.vertexai import VertexAIGeminiFlash1_5_0514
-from langfun.core.llms.vertexai import VertexAIGeminiPro1
-from langfun.core.llms.vertexai import VertexAIGeminiPro1Vision
-from langfun.core.llms.vertexai import VertexAIEndpoint
-
-
 # LLaMA C++ models.
 from langfun.core.llms.llama_cpp import LlamaCppRemote