修改gemini API为 Google的Vertex AI API

xinkeaboard-gemini-langgraph_prompt/backend/src/agent/graph.py

@@ -7,7 +7,8 @@ from langgraph.types import Send
 from langgraph.graph import StateGraph
 from langgraph.graph import START, END
 from langchain_core.runnables import RunnableConfig
-from google.genai import Client
+from google import genai
+from google.genai.types import HttpOptions
 
 from agent.state import (
     OverallState,
@@ -23,13 +24,15 @@ from agent.prompts import (
     reflection_instructions,
     answer_instructions,
 )
-from langchain_google_genai import ChatGoogleGenerativeAI
 from agent.utils import (
     get_citations,
     get_research_topic,
     insert_citation_markers,
     resolve_urls,
 )
+
+from langchain_google_vertexai import ChatVertexAI
+
 import logging
 
 load_dotenv()
@@ -41,12 +44,21 @@ logging.basicConfig(
 
 logger = logging.getLogger(__name__)
 
-if os.getenv("GEMINI_API_KEY") is None:
-    raise ValueError("GEMINI_API_KEY is not set")
+# if os.getenv("GEMINI_API_KEY") is None:
+#     raise ValueError("GEMINI_API_KEY is not set")
 
-# Used for Google Search API
-genai_client = Client(api_key=os.getenv("GEMINI_API_KEY"))
+if os.getenv("GOOGLE_APPLICATION_CREDENTIALS") is None:
+    raise ValueError("GOOGLE_APPLICATION_CREDENTIALS is not set")
+if os.getenv("GOOGLE_CLOUD_PROJECT") is None:
+    raise ValueError("GOOGLE_CLOUD_PROJECT is not set")
+if os.getenv("GOOGLE_CLOUD_LOCATION") is None:
+    raise ValueError("GOOGLE_CLOUD_LOCATION is not set")
+if os.getenv("GOOGLE_GENAI_USE_VERTEXAI") is None:
+    raise ValueError("GOOGLE_GENAI_USE_VERTEXAI is not set")
 
+# Used for Google Search API
+# genai_client = Client(api_key=os.getenv("GEMINI_API_KEY"))
+genai_client = genai.Client(http_options=HttpOptions(api_version="v1"))
 
 # Nodes
 def generate_query(state: OverallState, config: RunnableConfig) -> QueryGenerationState:
@@ -64,21 +76,25 @@ def generate_query(state: OverallState, config: RunnableConfig) -> QueryGenerati
     """
     configurable = Configuration.from_runnable_config(config)
     logger.info("开始：generate_query")
-    logger.info("1：generate_query")
     # check for custom initial search query count
     if state.get("initial_search_query_count") is None:
         state["initial_search_query_count"] = configurable.number_of_initial_queries
-    logger.info("2：generate_query")
     # init Gemini 2.0 Flash
-    llm = ChatGoogleGenerativeAI(
-        model=configurable.query_generator_model,
+    # llm = ChatGoogleGenerativeAI(
+    #     model=configurable.query_generator_model,
+    #     temperature=1.0,
+    #     max_retries=2,
+    #     api_key=os.getenv("GEMINI_API_KEY"),
+    # )
+
+    llm = ChatVertexAI(
+        model_name=configurable.query_generator_model,  # 如 "gemini-pro"
         temperature=1.0,
         max_retries=2,
-        api_key=os.getenv("GEMINI_API_KEY"),
+        project=os.getenv("GOOGLE_CLOUD_PROJECT"),  # 必填
+        location=os.getenv("GOOGLE_CLOUD_LOCATION")  # 必填
     )
-    logger.info("3：generate_query")
     structured_llm = llm.with_structured_output(SearchQueryList)
-    logger.info("4：generate_query")
     # Format the prompt
     current_date = get_current_date()
     formatted_prompt = query_writer_instructions.format(
@@ -139,7 +155,7 @@ def web_research(state: WebSearchState, config: RunnableConfig) -> OverallState:
     # for chunk in chunks:
     #     print(chunk["title"])
     #     print(chunk["url"])
-    #     print(chunk["content"]) 
+    #     print(chunk["content"])
     # resolve the urls to short urls for saving tokens and time
     resolved_urls = resolve_urls(
         response.candidates[0].grounding_metadata.grounding_chunks, state["id"]
@@ -184,11 +200,18 @@ def reflection(state: OverallState, config: RunnableConfig) -> ReflectionState:
         summaries="\n\n---\n\n".join(state["web_research_result"]),
     )
     # init Reasoning Model
-    llm = ChatGoogleGenerativeAI(
-        model=reasoning_model,
+    # llm = ChatGoogleGenerativeAI(
+    #     model=reasoning_model,
+    #     temperature=1.0,
+    #     max_retries=2,
+    #     api_key=os.getenv("GEMINI_API_KEY"),
+    # )
+    llm = ChatVertexAI(
+        model_name=configurable.query_generator_model,  # 如 "gemini-pro"
         temperature=1.0,
         max_retries=2,
-        api_key=os.getenv("GEMINI_API_KEY"),
+        project=os.getenv("GOOGLE_CLOUD_PROJECT"),  # 必填
+        location=os.getenv("GOOGLE_CLOUD_LOCATION")  # 必填
     )
     result = llm.with_structured_output(Reflection).invoke(formatted_prompt)
     logger.info("结束：reflection")
@@ -266,11 +289,18 @@ def finalize_answer(state: OverallState, config: RunnableConfig):
     )
 
     # init Reasoning Model, default to Gemini 2.5 Flash
-    llm = ChatGoogleGenerativeAI(
-        model=reasoning_model,
-        temperature=0,
+    # llm = ChatGoogleGenerativeAI(
+    #     model=reasoning_model,
+    #     temperature=0,
+    #     max_retries=2,
+    #     api_key=os.getenv("GEMINI_API_KEY"),
+    # )
+    llm = ChatVertexAI(
+        model_name=configurable.query_generator_model,  # 如 "gemini-pro"
+        temperature=1.0,
         max_retries=2,
-        api_key=os.getenv("GEMINI_API_KEY"),
+        project=os.getenv("GOOGLE_CLOUD_PROJECT"),  # 必填
+        location=os.getenv("GOOGLE_CLOUD_LOCATION")  # 必填
     )
     logger.info("开始：llm.invoke")
     result = llm.invoke(formatted_prompt)

xinkeaboard-gemini-langgraph_prompt/backend/src/agent/test.py

@@ -1,6 +1,15 @@
-import agent.prompts
-import sys
-sys.path.insert(0, "/Users/joeychen/Desktop/Weichi_proj/gemini-langgraph_prompt_engineering/backend/src")
+import os
+from google import genai
+from google.genai.types import HttpOptions
 
-print(">>> PROMPT MODULE PATH:", agent.prompts.__file__)
-print(">>> answer_instructions PREVIEW:\n", agent.prompts.answer_instructions[:300])
+os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = r"D:\project\advich\xinkeaboard\xinkeaboard-gemini-langgraph_prompt\backend\xinke-b2b-project-devops-763c6b8c6901.json"
+os.environ["GOOGLE_CLOUD_PROJECT"] = r"xinke-b2b-project-devops"
+os.environ["GOOGLE_CLOUD_LOCATION"] = r"global"
+os.environ["GOOGLE_GENAI_USE_VERTEXAI"] = r"True"
+
+client = genai.Client(http_options=HttpOptions(api_version="v1"))
+response = client.models.generate_content(
+    model="gemini-2.5-flash",
+    contents="How does AI work? in 10 words",
+)
+print(response.text)

xinkeaboard-gemini-langgraph_prompt/backend/xinke-b2b-project-devops-763c6b8c6901.json

@@ -0,0 +1,13 @@
+{
+  "type": "service_account",
+  "project_id": "xinke-b2b-project-devops",
+  "private_key_id": "763c6b8c69017856d8781ce4a5f4e84707ef931a",
+  "private_key": "-----BEGIN PRIVATE KEY-----\nMIIEvgIBADANBgkqhkiG9w0BAQEFAASCBKgwggSkAgEAAoIBAQCqoGGS1VKZNGU/\nV7a8IJntaRV62qX2cey6f2aEiAJoi88WoMYGIjjj5XbXTTGXS5JxfYYldt2ieEzU\niamGQEcjk7aS6tbEd57aDqfsPTTdIaRGx6JhtksdYOJv0ggZbEkQA9RNj/Y/dWU6\nc+tE7/MN+zS6zUryi8UN7rNS7o+vzUG+5MWtJ4GBMUA5aSAAFuw+n9DaQQywGxFe\nUnnH98xloCcjFzOvrHcVlmYsxWWnJmRpqnnlA5An/I+rZoQmjaicXGYU9r+Ke8zm\nGOdh3m5FAzUPQV9PAzPzRbWbmVx0q9pWOcu40WmW5px7IXea0VJDyjLYKQi2iEWc\nlMNOyX69AgMBAAECggEAFjKTPxT9j+zfFySEI2liXR6EqtzzJmKf14tnfOps1ckK\ngfkcigO9vvhFQ7FJOFOek6Eu6ITURKS96hCWu151Ys1+vRJMKncddWTet78mt3X+\n9y5Vk7yxfPrzunRqlVQfFTRKVGTcNllBeB9IhklpJ3J846h9jZxdYHijX7ClhMtK\n1xPHNZ5rEsx/AgiYhH4uQuzgReapbvRXifuwoejRC7tEt0Dao6iRYzGYcZODF+G+\nBU6R6TPGuqrp/ZtjqkAmMqFM/YfKdi7A6BpkzEU6MoA2/pc4yBu3rLFg+OksiNf1\n5naW+Ax7Ib8wPWDLRW4gXVPVzxrjk1YajV7r/zdLSQKBgQDRPC7cEdGrVs1xbaZa\nz2kR3o7tvvpymzO9v2l6gW3WbP9zsWasFcrTJQveVxPBMsH/me4er8YDw/Yi7RcG\ns6yBmrgMQU805X9v0vUBwb7MxJP+PzsKRU/biQGyei+9z+CjV8hFaigaF8dzJ0Pg\ndUHMBFpPOrmiopxBD+hPWuCi0wKBgQDQwyCdmXwg9A5cM5oZSi4Zwaz8Pq2cQUnb\nRH7eFDtx11mGwAKA3CRi6S9xddpWWIDNFvBxy1GbkgkMGs8AXz+KP5H8iqsVuYwI\ny/ClWa6o5W3JvBTDbs//GyVlXcKSogQ907tnuMoA0na3kqYF72mElS9MrVPvSfpf\nmoBVgKG+LwKBgQCwaWW49u0o4U7mS+Tuu+GCx3sKyPKwnlrJqRFt0KKjH9Nh4DkI\n7nbL/3g6RoDJa9pVf5tdFCaIFjj4yZ2YcPUJcUGsU4RMvRLtXZhlms0cKMH+o6FL\nz5S9LI7ZxFtsVlC4Z9Lbc93L4jCvznNUVIjdUu4xzvmlPiUwVhEwlxa/MQKBgFFT\notIip/vCeBGtPJwbjOgVzkkaNdIIkTgtb+seIKZRyZu4E4GVwUsTIptl2SAILG05\nQYRjO7CKbKHdk8TGVl+KMloav1mgRhBrOp1/2Y4Miv+rGL+l3Lu6HJBnvVQYb+SS\nJpDaqT82vm73X1wyj33KDHonj+pJfJNgAhzNFrWbAoGBAKmCJB4FPz9T8sOMRaog\nHG8W6KK6xqDl5MumHGAzyoO+bCaiWYK4UFKX2wZpd2y5S57n9Y3ZKsb+FV4PDely\nc873HuflApdS0WxXqIvLa+xzbrppfNAwVQKRDZdXO5CJdsJCIb4SFowQsx5x4pDp\nY0jc7SKqyv1q4xyTHkcEuPfk\n-----END PRIVATE KEY-----\n",
+  "client_email": "xinke-vertex-ai@xinke-b2b-project-devops.iam.gserviceaccount.com",
+  "client_id": "111297466545525410293",
+  "auth_uri": "https://accounts.google.com/o/oauth2/auth",
+  "token_uri": "https://oauth2.googleapis.com/token",
+  "auth_provider_x509_cert_url": "https://www.googleapis.com/oauth2/v1/certs",
+  "client_x509_cert_url": "https://www.googleapis.com/robot/v1/metadata/x509/xinke-vertex-ai%40xinke-b2b-project-devops.iam.gserviceaccount.com",
+  "universe_domain": "googleapis.com"
+}

xinkeaboard-gemini-langgraph_prompt/docker-compose.yml

@@ -38,7 +38,10 @@ services:
       langgraph-postgres:
         condition: service_healthy
     environment:
-      GEMINI_API_KEY: ${GEMINI_API_KEY}
       LANGSMITH_API_KEY: ${LANGSMITH_API_KEY}
+      GOOGLE_APPLICATION_CREDENTIALS: ${GOOGLE_APPLICATION_CREDENTIALS}
+      GOOGLE_CLOUD_PROJECT: ${GOOGLE_CLOUD_PROJECT}
+      GOOGLE_CLOUD_LOCATION: ${GOOGLE_CLOUD_LOCATION}
+      GOOGLE_GENAI_USE_VERTEXAI: ${GOOGLE_GENAI_USE_VERTEXAI}
       REDIS_URI: redis://langgraph-redis:6379
       POSTGRES_URI: postgres://postgres:postgres@langgraph-postgres:5432/postgres?sslmode=disable