77 lines
2.3 KiB
Python
77 lines
2.3 KiB
Python
|
"""
|
||
|
|
Test script for the direct Ollama endpoint.
|
||
|
|
"""
|
||
|
|
|
||
|
|
import requests
|
||
|
|
import json
|
||
|
|
import argparse
|
||
|
|
import time
|
||
|
|
|
||
|
|
def test_direct_ollama(api_url, model_id, prompt):
|
||
|
|
"""Test the direct Ollama endpoint."""
|
||
|
|
print(f"Testing direct Ollama endpoint at: {api_url}/test-ollama-direct")
|
||
|
|
print(f"Model: {model_id}")
|
||
|
|
print(f"Prompt: {prompt}")
|
||
|
|
|
||
|
|
try:
|
||
|
|
start_time = time.time()
|
||
|
|
|
||
|
|
response = requests.post(
|
||
|
|
f"{api_url}/test-ollama-direct",
|
||
|
|
headers={"Content-Type": "application/json"},
|
||
|
|
json={
|
||
|
|
"model": model_id,
|
||
|
|
"prompt": prompt
|
||
|
|
},
|
||
|
|
timeout=300 # 5 minutes timeout
|
||
|
|
)
|
||
|
|
|
||
|
|
end_time = time.time()
|
||
|
|
elapsed_time = end_time - start_time
|
||
|
|
|
||
|
|
response.raise_for_status()
|
||
|
|
|
||
|
|
result = response.json()
|
||
|
|
print(f"\nResponse (took {elapsed_time:.2f} seconds):")
|
||
|
|
|
||
|
|
# Handle different response formats
|
||
|
|
content = None
|
||
|
|
if 'response' in result and isinstance(result['response'], dict) and 'message' in result['response']:
|
||
|
|
content = result['response']['message'].get('content')
|
||
|
|
elif 'response' in result and isinstance(result['response'], str):
|
||
|
|
content = result['response']
|
||
|
|
else:
|
||
|
|
content = str(result)
|
||
|
|
|
||
|
|
print(content or 'No content in response')
|
||
|
|
|
||
|
|
return True
|
||
|
|
except Exception as e:
|
||
|
|
print(f"ERROR: Direct Ollama test failed: {str(e)}")
|
||
|
|
return False
|
||
|
|
|
||
|
|
def main():
|
||
|
|
"""Main function."""
|
||
|
|
parser = argparse.ArgumentParser(description='Test the direct Ollama endpoint')
|
||
|
|
parser.add_argument('--api-url', type=str, default='http://localhost:5252', help='AI service API URL')
|
||
|
|
parser.add_argument('--model', type=str, default='llama3.1', help='Model to use for testing')
|
||
|
|
parser.add_argument('--prompt', type=str, default='What is the capital of France?', help='Prompt to use for testing')
|
||
|
|
|
||
|
|
args = parser.parse_args()
|
||
|
|
|
||
|
|
print("=== Direct Ollama Test ===")
|
||
|
|
print(f"API URL: {args.api_url}")
|
||
|
|
print(f"Model: {args.model}")
|
||
|
|
print(f"Prompt: {args.prompt}")
|
||
|
|
print()
|
||
|
|
|
||
|
|
# Test direct Ollama endpoint
|
||
|
|
success = test_direct_ollama(args.api_url, args.model, args.prompt)
|
||
|
|
|
||
|
|
# Print summary
|
||
|
|
print("\n=== Test Summary ===")
|
||
|
|
print(f"Direct Ollama Test: {'SUCCESS' if success else 'FAILED'}")
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
main()
|