53 lines
1.8 KiB
Python
53 lines
1.8 KiB
Python
import unittest
|
|
from LlamaCppApi import LlamaCppApi
|
|
|
|
class TestLlamaCppApi(unittest.TestCase):
|
|
def setUp(self):
|
|
# Initializes the LlamaCppApi client for integration tests
|
|
self.client = LlamaCppApi(base_url="http://localhost:8080", api_key="optional_api_key")
|
|
|
|
def test_post_completion(self):
|
|
# Test the post_completion method for a successful status code.
|
|
response = self.client.post_completion("Start of a story", {"temperature": 0.9, "n_predict": 15})
|
|
self.assertEqual(response.status_code, 200)
|
|
print(response.json())
|
|
|
|
def test_tokenization(self):
|
|
# Test the tokenization endpoint for a successful status code.
|
|
response = self.client.post_tokenize("Example text bob alice eve", {"option_key": "option_value"})
|
|
self.assertEqual(response.status_code, 200)
|
|
print(response.json())
|
|
|
|
|
|
def test_detokenization(self):
|
|
# Test the detokenization endpoint for a successful status code.
|
|
response = self.client.post_detokenize([13617, 1495, 36292, 71533, 49996], {"option_key": "option_value"})
|
|
self.assertEqual(response.status_code, 200)
|
|
print(response.json())
|
|
|
|
|
|
def test_health_check(self):
|
|
# Tests the health check endpoint for a successful status code.
|
|
response = self.client.get_health()
|
|
self.assertEqual(response.status_code, 200)
|
|
print(response.json())
|
|
|
|
|
|
def test_stream_response(self):
|
|
|
|
def print_chunk(chunk):
|
|
print("Received Chunk:", chunk)
|
|
|
|
response = self.client.stream_response(
|
|
endpoint='completion',
|
|
data={"prompt": "Stream this story", "stream": True, "temperature": 0.7, "n_predict": 32, "stop":["<|im_end|>","<|eot_id|>"]},
|
|
chunk_callback=print_chunk
|
|
)
|
|
|
|
print(response)
|
|
|
|
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|