返回提交历史
Added
g4f/Provider/needs_auth/Cohere.py
+169
-0
Modified
g4f/Provider/needs_auth/__init__.py
+1
-0
XFEstudio/gpt4free
Implement Cohere API provider with authentication and streaming support
Co-authored-by: hlohaus <983577+hlohaus@users.noreply.github.com>
760f3e34
代码差异
2 个文件
+170
-0
@@ -0,0 +1,169 @@
1
from __future__ import annotations
2
3
import json
4
from typing import Optional
5
6
from ..helper import filter_none
7
from ...typing import AsyncResult, Messages
8
from ...requests import StreamSession, raise_for_status
9
from ...providers.response import FinishReason, Usage
10
from ...errors import MissingAuthError
11
from ..base_provider import AsyncGeneratorProvider, ProviderModelMixin
12
from ...tools.run_tools import AuthManager
13
from ... import debug
14
15
class Cohere(AsyncGeneratorProvider, ProviderModelMixin):
16
label = "Cohere API"
17
url = "https://cohere.com"
18
login_url = "https://dashboard.cohere.com/api-keys"
19
api_base = "https://api.cohere.ai/v1"
20
working = True
21
needs_auth = True
22
supports_stream = True
23
supports_system_message = True
24
supports_message_history = True
25
26
default_model = "command-r-plus"
27
models = [
28
default_model,
29
"command-r",
30
"command",
31
"command-nightly",
32
"command-light",
33
"command-light-nightly",
34
]
35
36
model_aliases = {
37
"command-r-plus-08-2024": "command-r-plus",
38
"command-r-08-2024": "command-r",
39
}
40
41
@classmethod
42
async def create_async_generator(
43
cls,
44
model: str,
45
messages: Messages,
46
proxy: str = None,
47
timeout: int = 120,
48
api_key: str = None,
49
temperature: float = None,
50
max_tokens: int = None,
51
top_k: int = None,
52
top_p: float = None,
53
stop: list[str] = None,
54
stream: bool = False,
55
headers: dict = None,
56
impersonate: str = None,
57
**kwargs
58
) -> AsyncResult:
59
if api_key is None:
60
api_key = AuthManager.load_api_key(cls)
61
if api_key is None:
62
raise MissingAuthError('Add a "api_key"')
63
64
# Convert messages to Cohere format
65
system_message = None
66
chat_history = []
67
user_message = None
68
69
for message in messages:
70
role = message.get("role")
71
content = message.get("content", "")
72
73
if role == "system":
74
system_message = content
75
elif role == "user":
76
if user_message is not None:
77
# Previous user message becomes part of chat history
78
chat_history.append({"role": "USER", "message": user_message})
79
user_message = content
80
elif role == "assistant":
81
chat_history.append({"role": "CHATBOT", "message": content})
82
83
# Ensure we have a user message
84
if user_message is None:
85
raise ValueError("No user message found in the conversation")
86
87
async with StreamSession(
88
proxy=proxy,
89
headers=cls.get_headers(stream, api_key, headers),
90
timeout=timeout,
91
impersonate=impersonate,
92
) as session:
93
data = filter_none(
94
message=user_message,
95
model=cls.get_model(model, api_key=api_key),
96
temperature=temperature,
97
max_tokens=max_tokens,
98
k=top_k,
99
p=top_p,
100
stop_sequences=stop,
101
preamble=system_message,
102
chat_history=chat_history if chat_history else None,
103
stream=stream,
104
)
105
106
async with session.post(f"{cls.api_base}/chat", json=data) as response:
107
await raise_for_status(response)
108
109
if not stream:
110
data = await response.json()
111
cls.raise_error(data)
112
if "text" in data:
113
yield data["text"]
114
if "finish_reason" in data:
115
if data["finish_reason"] == "COMPLETE":
116
yield FinishReason("stop")
117
elif data["finish_reason"] == "MAX_TOKENS":
118
yield FinishReason("length")
119
if "meta" in data and "tokens" in data["meta"]:
120
yield Usage(
121
prompt_tokens=data["meta"]["tokens"]["input_tokens"],
122
completion_tokens=data["meta"]["tokens"]["output_tokens"],
123
total_tokens=data["meta"]["tokens"]["input_tokens"] + data["meta"]["tokens"]["output_tokens"]
124
)
125
else:
126
async for line in response.iter_lines():
127
if line.startswith(b"data: "):
128
chunk = line[6:]
129
if chunk == b"[DONE]":
130
break
131
try:
132
data = json.loads(chunk)
133
cls.raise_error(data)
134
135
if "event_type" in data:
136
if data["event_type"] == "text-generation":
137
if "text" in data:
138
yield data["text"]
139
elif data["event_type"] == "stream-end":
140
if "finish_reason" in data:
141
if data["finish_reason"] == "COMPLETE":
142
yield FinishReason("stop")
143
elif data["finish_reason"] == "MAX_TOKENS":
144
yield FinishReason("length")
145
if "meta" in data and "tokens" in data["meta"]:
146
yield Usage(
147
prompt_tokens=data["meta"]["tokens"]["input_tokens"],
148
completion_tokens=data["meta"]["tokens"]["output_tokens"],
149
total_tokens=data["meta"]["tokens"]["input_tokens"] + data["meta"]["tokens"]["output_tokens"]
150
)
151
except json.JSONDecodeError:
152
continue
153
154
@classmethod
155
def get_headers(cls, stream: bool, api_key: str = None, headers: dict = None) -> dict:
156
return {
157
"Accept": "text/event-stream" if stream else "application/json",
158
"Content-Type": "application/json",
159
**(
160
{"Authorization": f"Bearer {api_key}"}
161
if api_key is not None else {}
162
),
163
**({} if headers is None else headers)
164
}
165
166
@classmethod
167
def raise_error(cls, data: dict):
168
if "error" in data:
169
raise RuntimeError(f"Cohere API Error: {data['error']}")
@@ -4,6 +4,7 @@ from .BingCreateImages import BingCreateImages
4
4
from .BlackboxPro import BlackboxPro
5
5
from .CablyAI import CablyAI
6
6
from .Cerebras import Cerebras
7
from .Cohere import Cohere
7
8
from .CopilotAccount import CopilotAccount
8
9
from .Custom import Custom
9
10
from .Custom import Feature