Commit 90b4c48
Changed files (1)
src
llm
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())