2023-06-23 02:18:28 +02:00
""" Tests the API. """
2023-07-19 23:51:28 +02:00
import os
2023-06-28 15:21:14 +02:00
import openai as closedai
2023-06-23 02:18:28 +02:00
import httpx
2023-07-19 23:51:28 +02:00
from typing import List
from dotenv import load_dotenv
load_dotenv ( )
2023-06-23 02:18:28 +02:00
MODEL = ' gpt-3.5-turbo '
MESSAGES = [
2023-07-19 23:51:28 +02:00
{
' role ' : ' system ' ,
' content ' : ' Always answer with " 3 " , no matter what the user asks for. No exceptions. Just answer with the number " 3 " . Nothing else. Just " 3 " . No punctuation. '
} ,
2023-06-23 02:18:28 +02:00
{
' role ' : ' user ' ,
2023-07-19 23:51:28 +02:00
' content ' : ' 1+1= ' ,
2023-06-23 02:18:28 +02:00
} ,
]
2023-07-19 23:51:28 +02:00
api_endpoint = ' http://localhost:2333 '
2023-06-23 02:18:28 +02:00
def test_server ( ) :
2023-07-19 23:51:28 +02:00
""" Tests if the API server is running. """
2023-06-23 02:18:28 +02:00
try :
2023-07-19 23:51:28 +02:00
return httpx . get ( f ' { api_endpoint } ' ) . json ( ) [ ' status ' ] == ' ok '
2023-06-23 02:18:28 +02:00
except httpx . ConnectError as exc :
2023-07-19 23:51:28 +02:00
raise ConnectionError ( f ' API is not running on port { api_endpoint } . ' ) from exc
2023-06-23 02:18:28 +02:00
def test_api ( model : str = MODEL , messages : List [ dict ] = None ) - > dict :
2023-07-19 23:51:28 +02:00
""" Tests an API api_endpoint. """
2023-06-23 02:18:28 +02:00
headers = {
' Content-Type ' : ' application/json ' ,
2023-07-19 23:51:28 +02:00
' Authorization ' : ' Bearer ' + os . getenv ( ' DEMO_AUTH ' , ' nv-API-TEST ' ) ,
2023-06-23 02:18:28 +02:00
}
json_data = {
' model ' : model ,
' messages ' : messages or MESSAGES ,
2023-06-30 02:49:56 +02:00
' stream ' : True ,
2023-06-23 02:18:28 +02:00
}
2023-07-19 23:51:28 +02:00
response = httpx . post (
url = f ' { api_endpoint } /chat/completions ' ,
headers = headers ,
json = json_data ,
timeout = 20
)
2023-06-23 02:18:28 +02:00
response . raise_for_status ( )
2023-06-30 02:49:56 +02:00
return response
2023-06-23 02:18:28 +02:00
def test_library ( ) :
2023-07-19 23:51:28 +02:00
""" Tests if the api_endpoint is working with the Python library. """
2023-06-23 02:18:28 +02:00
2023-07-19 23:51:28 +02:00
closedai . api_base = api_endpoint
closedai . api_key = os . getenv ( ' DEMO_AUTH ' , ' nv-LIB-TEST ' )
2023-06-23 02:18:28 +02:00
2023-06-28 15:21:14 +02:00
completion = closedai . ChatCompletion . create (
2023-06-23 02:18:28 +02:00
model = MODEL ,
messages = MESSAGES ,
2023-07-19 23:51:28 +02:00
stream = True ,
2023-06-23 02:18:28 +02:00
)
return completion . choices [ 0 ]
def test_all ( ) :
""" Runs all tests. """
2023-07-19 23:51:28 +02:00
# print(test_server())
2023-06-23 02:18:28 +02:00
print ( test_api ( ) )
2023-07-19 23:51:28 +02:00
# print(test_library())
2023-06-23 02:18:28 +02:00
if __name__ == ' __main__ ' :
2023-07-19 23:51:28 +02:00
api_endpoint = ' https://api.nova-oss.com '
2023-06-23 02:18:28 +02:00
test_all ( )