139 lines
4.2 KiB
Python
139 lines
4.2 KiB
Python
"""
|
|
AGC Document Chatbot - API Test Suite
|
|
|
|
This script provides comprehensive testing for all API endpoints:
|
|
- Document listing and filtering
|
|
- Document retrieval
|
|
- Search functionality (both enhanced AI and simple fallback)
|
|
- Error handling and edge cases
|
|
|
|
Usage:
|
|
python test_api.py
|
|
|
|
Prerequisites:
|
|
- API server running on localhost:8000
|
|
- requests library installed
|
|
"""
|
|
|
|
import requests
|
|
import json
|
|
from typing import Dict, Any
|
|
from pprint import pprint
|
|
|
|
# API base URL
|
|
BASE_URL = "http://localhost:8000"
|
|
|
|
def test_root():
|
|
"""Test the root endpoint"""
|
|
print("\n=== Testing Root Endpoint ===")
|
|
response = requests.get(f"{BASE_URL}/")
|
|
pprint(response.json())
|
|
|
|
def test_list_documents():
|
|
"""Test document listing with different filters"""
|
|
print("\n=== Testing Document Listing ===")
|
|
|
|
# List all documents
|
|
print("\nListing all documents:")
|
|
response = requests.get(f"{BASE_URL}/documents")
|
|
pprint(response.json()[:2]) # Show first 2 documents
|
|
|
|
# Filter by document type
|
|
print("\nFiltering by document type:")
|
|
response = requests.get(f"{BASE_URL}/documents", params={"doc_type": "LKK"})
|
|
pprint(response.json()[:2])
|
|
|
|
# Filter by title/content
|
|
print("\nFiltering by title/content:")
|
|
response = requests.get(f"{BASE_URL}/documents", params={"title_filter": "case"})
|
|
pprint(response.json()[:2])
|
|
|
|
# Combined filters
|
|
print("\nCombined filters:")
|
|
response = requests.get(
|
|
f"{BASE_URL}/documents",
|
|
params={
|
|
"doc_type": "LKK",
|
|
"title_filter": "case"
|
|
}
|
|
)
|
|
pprint(response.json()[:2])
|
|
|
|
def test_document_types():
|
|
"""Test getting available document types"""
|
|
print("\n=== Testing Document Types ===")
|
|
response = requests.get(f"{BASE_URL}/document-types")
|
|
pprint(response.json())
|
|
|
|
def test_get_document():
|
|
"""Test getting a specific document"""
|
|
print("\n=== Testing Document Retrieval ===")
|
|
|
|
# First, get a list of documents to find a valid ID
|
|
documents = requests.get(f"{BASE_URL}/documents").json()
|
|
if documents:
|
|
doc_id = documents[0]['id']
|
|
print(f"\nRetrieving document with ID {doc_id}:")
|
|
response = requests.get(f"{BASE_URL}/documents/{doc_id}")
|
|
pprint(response.json())
|
|
|
|
# Test invalid document ID
|
|
print("\nTesting invalid document ID:")
|
|
response = requests.get(f"{BASE_URL}/documents/99999")
|
|
pprint(response.json())
|
|
|
|
def test_search():
|
|
"""Test document search functionality"""
|
|
print("\n=== Testing Document Search ===")
|
|
|
|
# Simple search
|
|
print("\nSimple search:")
|
|
search_data = {
|
|
"query": "criminal case involving drugs",
|
|
"profile_search": False
|
|
}
|
|
response = requests.post(
|
|
f"{BASE_URL}/search",
|
|
json=search_data
|
|
)
|
|
result = response.json()
|
|
print(f"Original Query: {result['query']}")
|
|
print(f"Enhanced Query: {result['enhanced_query']}")
|
|
print("\nTop matching documents:")
|
|
for doc in result['documents'][:2]: # Show first 2 results
|
|
print(f"\nTitle: {doc['title']}")
|
|
print(f"Relevance: {doc['similarity']:.2f}")
|
|
print(f"Preview: {doc['content_preview'][:200]}...")
|
|
|
|
# Legal concept search
|
|
print("\nLegal concept search:")
|
|
search_data = {
|
|
"query": "explain the concept of mens rea in criminal cases",
|
|
"profile_search": False
|
|
}
|
|
response = requests.post(
|
|
f"{BASE_URL}/search",
|
|
json=search_data
|
|
)
|
|
result = response.json()
|
|
print(f"Original Query: {result['query']}")
|
|
print(f"Enhanced Query: {result['enhanced_query']}")
|
|
print(f"AI Answer: {result['answer']}")
|
|
|
|
def run_all_tests():
|
|
"""Run all API tests"""
|
|
try:
|
|
test_root()
|
|
test_document_types()
|
|
test_list_documents()
|
|
test_get_document()
|
|
test_search()
|
|
print("\n=== All tests completed successfully ===")
|
|
except requests.exceptions.ConnectionError:
|
|
print("\nError: Could not connect to the API server.")
|
|
print("Make sure the API is running with: uvicorn api:app --reload")
|
|
except Exception as e:
|
|
print(f"\nError during testing: {str(e)}")
|
|
|
|
if __name__ == "__main__":
|
|
run_all_tests() |