XFE Git
XFE Studio Git
Git 首页 全局搜索
XFE 主站 文档 NuGet
公开
关注 0 Fork 0 Star 1
返回提交历史

XFEstudio/gpt4free

Custom api_base for GeminiPro

307c8f53
Heiner Lohaus <hlohaus@users.noreply.github.com>
提交于

代码差异

1 个文件 +16 -10
Modified g4f/Provider/GeminiPro.py +16 -10
@@ -13,6 +13,7 @@ class GeminiPro(AsyncGeneratorProvider, ProviderModelMixin):
13 13 url = "https://ai.google.dev"
14 14 working = True
15 15 supports_message_history = True
16 needs_auth = True
16 17 default_model = "gemini-pro"
17 18 models = ["gemini-pro", "gemini-pro-vision"]
18 19
@@ -24,19 +25,24 @@ class GeminiPro(AsyncGeneratorProvider, ProviderModelMixin):
24 25 stream: bool = False,
25 26 proxy: str = None,
26 27 api_key: str = None,
28 api_base: str = None,
27 29 image: ImageType = None,
28 30 **kwargs
29 31 ) -> AsyncResult:
30 32 model = "gemini-pro-vision" if not model and image else model
31 33 model = cls.get_model(model)
32 if not api_key:
33 raise MissingAuthError('Missing "api_key" for auth')
34 headers = {
35 "Content-Type": "application/json",
36 }
37 async with ClientSession(headers=headers) as session:
38 method = "streamGenerateContent" if stream else "generateContent"
39 url = f"https://generativelanguage.googleapis.com/v1beta/models/{model}:{method}"
34
35 if not api_key and not api_base:
36 raise MissingAuthError('Missing "api_key" or "api_base"')
37 if not api_base:
38 api_base = f"https://generativelanguage.googleapis.com/v1beta"
39
40 method = "streamGenerateContent" if stream else "generateContent"
41 url = f"{api_base.rstrip('/')}/models/{model}:{method}"
42 if api_key:
43 url += f"?key={api_key}"
44
45 async with ClientSession() as session:
40 46 contents = [
41 47 {
42 48 "role": "model" if message["role"] == "assistant" else message["role"],
@@ -62,7 +68,7 @@ class GeminiPro(AsyncGeneratorProvider, ProviderModelMixin):
62 68 "topK": kwargs.get("top_k"),
63 69 }
64 70 }
65 async with session.post(url, params={"key": api_key}, json=data, proxy=proxy) as response:
71 async with session.post(url, json=data, proxy=proxy) as response:
66 72 if not response.ok:
67 73 data = await response.json()
68 74 raise RuntimeError(data[0]["error"]["message"])
@@ -78,7 +84,7 @@ class GeminiPro(AsyncGeneratorProvider, ProviderModelMixin):
78 84 yield data["candidates"][0]["content"]["parts"][0]["text"]
79 85 except:
80 86 data = data.decode() if isinstance(data, bytes) else data
81 raise RuntimeError(f"Read text failed. data: {data}")
87 raise RuntimeError(f"Read chunk failed. data: {data}")
82 88 lines = []
83 89 else:
84 90 lines.append(chunk)