Spaces:
Sleeping
Sleeping
File size: 4,937 Bytes
4e2263c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 |
import requests
import json
from typing import Union, Dict, Generator
import time
class ChatCompletionTester:
def __init__(self, base_url: str = "http://localhost:8000"):
self.base_url = base_url
self.endpoint = f"{base_url}/chat/completions"
def create_test_payload(self, stream: bool = False) -> Dict:
"""Create a sample payload for testing"""
return {
"model": "mistralai/Mixtral-8x22B-Instruct-v0.1",
"messages": [
{"role": "system", "content": "You are a helpful assistant."},
{"role": "user", "content": "What is the capital of France?"}
],
"temperature": 0.7,
"max_tokens": 4096,
"stream": stream
}
def test_non_streaming(self) -> Union[Dict, None]:
"""Test non-streaming response"""
print("\n=== Testing Non-Streaming Response ===")
try:
payload = self.create_test_payload(stream=False)
print("Sending request...")
response = requests.post(
self.endpoint,
json=payload,
headers={"Content-Type": "application/json"}
)
if response.status_code == 200:
result = response.json()
content = result['choices'][0]['message']['content']
print("\nResponse received successfully!")
print(f"Content: {content}")
return result
else:
print(f"Error: Status code {response.status_code}")
print(f"Response: {response.text}")
return None
except Exception as e:
print(f"Error during non-streaming test: {str(e)}")
return None
def test_streaming(self) -> Union[str, None]:
"""Test streaming response"""
print("\n=== Testing Streaming Response ===")
try:
payload = self.create_test_payload(stream=True)
print("Sending request...")
response = requests.post(
self.endpoint,
json=payload,
headers={"Content-Type": "application/json"},
stream=True
)
if response.status_code == 200:
print("\nReceiving streaming response:")
full_response = ""
for line in response.iter_lines(decode_unicode=True):
if line:
if line.startswith("data: "):
try:
data = json.loads(line[6:])
if data == "[DONE]":
continue
content = data.get("choices", [{}])[0].get("delta", {}).get("content", "")
if content:
print(content, end="", flush=True)
full_response += content
except json.JSONDecodeError:
continue
print("\n\nStreaming completed!")
return full_response
else:
print(f"Error: Status code {response.status_code}")
print(f"Response: {response.text}")
return None
except Exception as e:
print(f"Error during streaming test: {str(e)}")
return None
def run_all_tests(self):
"""Run both streaming and non-streaming tests"""
print("Starting API endpoint tests...")
# Test server connectivity
try:
requests.get(self.base_url)
print("✓ Server is accessible")
except requests.exceptions.ConnectionError:
print("✗ Server is not accessible. Please ensure the FastAPI server is running.")
return
# Run tests with timing
start_time = time.time()
# Test non-streaming
non_streaming_result = self.test_non_streaming()
if non_streaming_result:
print("✓ Non-streaming test passed")
else:
print("✗ Non-streaming test failed")
# Test streaming
streaming_result = self.test_streaming()
if streaming_result:
print("✓ Streaming test passed")
else:
print("✗ Streaming test failed")
end_time = time.time()
print(f"\nAll tests completed in {end_time - start_time:.2f} seconds")
def main():
# Create tester instance
tester = ChatCompletionTester()
# Run all tests
tester.run_all_tests()
if __name__ == "__main__":
main() |