Added support for Google Gemini 1.5 pro/flash models

.env.template

@@ -3,3 +3,8 @@
 # OPENAI_API_KEY=Your personal OpenAI API key from https://platform.openai.com/account/api-keys
 OPENAI_API_KEY=...
 ANTHROPIC_API_KEY=...
+# GOOGLE_API_KEY=Your personal GOOGLE API key from https://aistudio.google.com/app/apikey
+GOOGLE_API_KEY=...
+
+# If not set Model Name defaults to gpt-4o
+MODEL_NAME=gemini-1.5-pro-latest

gpt_engineer/applications/cli/main.py

@@ -89,6 +89,19 @@ def load_env_if_needed():
     if os.getenv("ANTHROPIC_API_KEY") is None:
         load_dotenv(dotenv_path=os.path.join(os.getcwd(), ".env"))
 
+    if os.getenv("GOOGLE_API_KEY") is None:
+        load_dotenv()
+    if os.getenv("GOOGLE_API_KEY") is None:
+        load_dotenv(dotenv_path=os.path.join(os.getcwd(), ".env"))
+
+def model_env():
+    if os.getenv("MODEL_NAME") is None:
+        load_dotenv()
+    if os.getenv("MODEL_NAME") is None:
+        load_dotenv(dotenv_path=os.path.join(os.getcwd(), ".env"))
+
+    return os.getenv("MODEL_NAME", default="gpt-4o")
+
 
 def concatenate_paths(base_path, sub_path):
     # Compute the relative path from base_path to sub_path
@@ -281,7 +294,7 @@ def format_installed_packages(packages):
 def main(
     project_path: str = typer.Argument(".", help="path"),
     model: str = typer.Option(
-        os.environ.get("MODEL_NAME", "gpt-4o"), "--model", "-m", help="model id string"
+        os.environ.get("MODEL_NAME", model_env()), "--model", "-m", help="model id string"
     ),
     temperature: float = typer.Option(
         0.1,

gpt_engineer/core/ai.py

@@ -36,6 +36,7 @@
     messages_to_dict,
 )
 from langchain_anthropic import ChatAnthropic
+from langchain_google_genai import ChatGoogleGenerativeAI
 from langchain_openai import AzureChatOpenAI, ChatOpenAI
 
 from gpt_engineer.core.token_usage import TokenUsageLog
@@ -362,6 +363,14 @@ def _create_chat_model(self) -> BaseChatModel:
                 streaming=self.streaming,
                 max_tokens_to_sample=4096,
             )
+        elif "gemini" in self.model_name:
+            return ChatGoogleGenerativeAI(
+                model=self.model_name,
+                temperature=self.temperature,
+                streaming=self.streaming,
+                google_api_key=os.getenv('GOOGLE_API_KEY'),
+                callbacks=[StreamingStdOutCallbackHandler()]
+            )
         elif self.vision:
             return ChatOpenAI(
                 model=self.model_name,

pyproject.toml

@@ -38,6 +38,7 @@ pillow = "^10.2.0"
 datasets = "^2.17.1"
 black = "23.3.0"
 langchain-community = "^0.2.0"
+langchain-google-genai = "^1.0.7"
 
 [tool.poetry.group.dev.dependencies]
 pytest = ">=7.3.1"