alex/nova-api
1
0
Fork 0
mirror of https://github.com/NovaOSS/nova-api.git synced 2024-11-25 18:23:57 +01:00
Code Issues Packages Projects Releases Wiki Activity

Added /v1/models and fixed key invalidation

Browse source
This commit is contained in:
nsde 2023-08-09 11:15:49 +02:00
parent c741c03bae
commit 0b86f7c26c
6 changed files with 43 additions and 33 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
api/db/users.py
View file

@ -15,7 +15,7 @@ with open('config/credits.yml', encoding='utf8') as f:
async def _get_mongo(collection_name: str):
return AsyncIOMotorClient(os.getenv('MONGO_URI'))['nova-core'][collection_name]
async def create(discord_id: int=0) -> dict:
async def create(discord_id: str='') -> dict:
"""Adds a new user to the MongoDB collection."""
chars = string.ascii_letters + string.digits
@ -36,7 +36,7 @@ async def create(discord_id: int=0) -> dict:
'ban_reason': '',
},
'auth': {
'discord': discord_id,
'discord': str(discord_id),
'github': None
}
}

3
api/moderation.py
View file

@ -47,7 +47,8 @@ async def is_policy_violated(inp) -> bool:
return False
except Exception as exc:
# await provider_auth.invalidate_key(req.get('provider_auth'))
if '401' in str(exc):
await provider_auth.invalidate_key(req.get('provider_auth'))
print('[!] moderation error:', type(exc), exc)
continue

26
api/streaming.py
View file

@ -10,6 +10,7 @@ from dotenv import load_dotenv
from python_socks._errors import ProxyError
import proxies
import provider_auth
import load_balancing
from db import logs, users, stats
@ -36,8 +37,6 @@ async def stream(
input_tokens: int=0,
incoming_request: starlette.requests.Request=None,
):
payload = payload or DEMO_PAYLOAD
is_chat = False
is_stream = payload.get('stream', False)
@ -45,7 +44,6 @@ async def stream(
is_chat = True
model = payload['model']
if is_chat and is_stream:
chat_id = await chat.create_chat_id()
@ -97,7 +95,10 @@ async def stream(
return
for k, v in target_request.get('headers', {}).items():
headers[k] = v
target_request['headers'][k] = v
if target_request['method'] == 'GET' and not payload:
target_request['payload'] = None
async with aiohttp.ClientSession(connector=proxies.default_proxy.connector) as session:
try:
@ -108,24 +109,30 @@ async def stream(
data=target_request.get('data'),
json=target_request.get('payload'),
headers=headers,
headers=target_request.get('headers', {}),
cookies=target_request.get('cookies'),
ssl=False,
timeout=aiohttp.ClientTimeout(total=float(os.getenv('TRANSFER_TIMEOUT', '120'))),
) as response:
if response.content_type == 'application/json':
data = await response.json()
if not is_stream:
json_response = await response.json()
if data.get('code') == 'invalid_api_key':
await provider_auth.invalidate_key(target_request.get('provider_auth'))
continue
if response.ok:
json_response = data
if is_stream:
try:
response.raise_for_status()
except Exception as exc:
if 'Too Many Requests' in str(exc):
continue
if is_stream:
try:
async for chunk in response.content.iter_any():
send = False
@ -180,10 +187,9 @@ async def stream(
content=chat.CompletionStop
)
yield chunk
yield 'data: [DONE]\n\n'
if not is_stream:
if not is_stream and json_response:
yield json.dumps(json_response)
# DONE =========================================================

2
api/transfer.py
View file

@ -101,8 +101,6 @@ async def handle(incoming_request):
# READY
# payload['user'] = str(user['_id'])
if 'chat/completions' in path and not payload.get('stream') is True:
payload['stream'] = False

2
rewards/last_update.txt
View file

@ -1 +1 @@
1691460004.7354248
1691546405.042006

31
tests/__main__.py
View file

@ -40,11 +40,6 @@ def test_server():
def test_api(model: str=MODEL, messages: List[dict]=None) -> dict:
"""Tests an API api_endpoint."""
headers = {
'Content-Type': 'application/json',
'Authorization': 'Bearer ' + closedai.api_key
}
json_data = {
'model': model,
'messages': messages or MESSAGES,
@ -53,7 +48,7 @@ def test_api(model: str=MODEL, messages: List[dict]=None) -> dict:
response = httpx.post(
url=f'{api_endpoint}/chat/completions',
headers=headers,
headers=HEADERS,
json=json_data,
timeout=20
)
@ -74,25 +69,30 @@ def test_library():
def test_library_moderation():
return closedai.Moderation.create("I wanna kill myself, I wanna kill myself; It's all I hear right now, it's all I hear right now")
def test_models():
response = httpx.get(
url=f'{api_endpoint}/models',
headers=HEADERS,
timeout=5
)
response.raise_for_status()
return response.json()
def test_all():
"""Runs all tests."""
# print(test_server())
# print(test_api())
print(test_library())
# print(test_library())
# print(test_library_moderation())
print(test_models())
def test_api_moderation(model: str=MODEL, messages: List[dict]=None) -> dict:
"""Tests an API api_endpoint."""
headers = {
'Authorization': 'Bearer ' + closedai.api_key
}
response = httpx.get(
url=f'{api_endpoint}/moderations',
headers=headers,
headers=HEADERS,
timeout=20
)
response.raise_for_status()
@ -104,4 +104,9 @@ if __name__ == '__main__':
closedai.api_base = api_endpoint
closedai.api_key = os.getenv('TEST_NOVA_KEY')
HEADERS = {
'Content-Type': 'application/json',
'Authorization': 'Bearer ' + closedai.api_key
}
test_all()