Commit 90b4c48

benny-dou <60535774+benny-dou@users.noreply.github.com>
2025-05-06 03:31:06
feat(gemini): Add support for custom API key
1 parent d847929
Changed files (1)
src
src/llm/gemini.py
@@ -92,9 +92,12 @@ async def gemini_stream(
     try:
         status: Message = kwargs.get("progress")  # type: ignore
         api_keys = [x.strip() for x in GEMINI.API_KEYS.split(",") if x.strip()]
+        if kwargs.get("gemini_api_keys"):
+            api_keys = [x.strip() for x in kwargs["gemini_api_keys"].split(",") if x.strip()]
         if retry > len(api_keys) - 1:
             return None
-        app = genai.Client(api_key=api_keys[retry], http_options=HttpOptions(base_url=GEMINI.BASR_URL, async_client_args={"proxy": GEMINI.PROXY}))
+        api_key = kwargs.get("gemini_api_key", api_keys[retry])
+        app = genai.Client(api_key=api_key, http_options=HttpOptions(base_url=GEMINI.BASR_URL, async_client_args={"proxy": GEMINI.PROXY}))
         runtime_texts = ""
         async for chunk in await app.aio.models.generate_content_stream(**params):
             resp = parse_response(chunk.model_dump())
@@ -148,9 +151,12 @@ async def gemini_nonstream(
     try:
         clean_gemini_sourcemarks(params["contents"])
         api_keys = [x.strip() for x in GEMINI.API_KEYS.split(",") if x.strip()]
+        if kwargs.get("gemini_api_keys"):
+            api_keys = [x.strip() for x in kwargs["gemini_api_keys"].split(",") if x.strip()]
         if retry > len(api_keys) - 1:
             return None
-        app = genai.Client(api_key=api_keys[retry], http_options=HttpOptions(base_url=GEMINI.BASR_URL, async_client_args={"proxy": GEMINI.PROXY}))
+        api_key = kwargs.get("gemini_api_key", api_keys[retry])
+        app = genai.Client(api_key=api_key, http_options=HttpOptions(base_url=GEMINI.BASR_URL, async_client_args={"proxy": GEMINI.PROXY}))
         response = await app.aio.models.generate_content(**params)
         prefix = f"🤖**{model_name}**:{BOT_TIPS}\n"
         res = parse_response(response.model_dump())