104 lines
3.4 KiB
Python
104 lines
3.4 KiB
Python
import requests
|
|
import json
|
|
import time
|
|
import sys
|
|
|
|
base_url = 'http://localhost:3015'
|
|
api_key = 'jleu1212'
|
|
workspace = 'test1'
|
|
|
|
headers = {
|
|
'X-API-Key': api_key,
|
|
'Content-Type': 'application/json',
|
|
'X-Workspace': workspace
|
|
}
|
|
|
|
def upload_document():
|
|
"""Upload a test document to workspace test1"""
|
|
files = {'file': open('test/tir.docx', 'rb')}
|
|
upload_url = f'{base_url}/documents/upload'
|
|
headers = {'X-API-Key': api_key, 'X-Workspace': workspace}
|
|
data = {'workspace': workspace, 'chunk_size': 1200, 'chunk_overlap': 100}
|
|
resp = requests.post(upload_url, headers=headers, files=files, data=data)
|
|
print(f'Upload status: {resp.status_code}')
|
|
if resp.status_code == 200:
|
|
print('Upload successful')
|
|
# Response may be a list of document IDs
|
|
result = resp.json()
|
|
print(f'Upload response: {result}')
|
|
# Assume first document ID
|
|
if isinstance(result, list) and len(result) > 0:
|
|
doc_id = result[0].get('document_id')
|
|
return {'document_id': doc_id}
|
|
elif isinstance(result, dict) and 'document_id' in result:
|
|
return result
|
|
else:
|
|
# fallback
|
|
return {'document_id': 'unknown'}
|
|
else:
|
|
print(f'Upload error: {resp.text}')
|
|
return None
|
|
|
|
def wait_for_indexing(doc_id):
|
|
"""Poll document status until processed"""
|
|
if doc_id == 'unknown':
|
|
# Cannot poll, just wait a bit
|
|
print('Cannot poll status, waiting 10 seconds...')
|
|
time.sleep(10)
|
|
return True
|
|
status_url = f'{base_url}/documents/{doc_id}/status'
|
|
for _ in range(30):
|
|
resp = requests.get(status_url, headers={'X-API-Key': api_key, 'X-Workspace': workspace})
|
|
if resp.status_code == 200:
|
|
data = resp.json()
|
|
status = data.get('status')
|
|
print(f'Document status: {status}')
|
|
if status == 'processed':
|
|
return True
|
|
elif status == 'error':
|
|
print(f'Document processing error: {data}')
|
|
return False
|
|
time.sleep(2)
|
|
print('Timeout waiting for indexing')
|
|
return False
|
|
|
|
def test_search():
|
|
"""Test search with keyword extraction"""
|
|
search_url = f'{base_url}/api/search'
|
|
payload = {'query': 'minimum safe working distance', 'top_k': 5}
|
|
resp = requests.post(search_url, headers=headers, json=payload, timeout=30)
|
|
print(f'Search status: {resp.status_code}')
|
|
if resp.status_code == 200:
|
|
data = resp.json()
|
|
print(f'Search successful. Found {len(data.get("chunks", []))} chunks')
|
|
print(f'Entities: {len(data.get("entities", []))}')
|
|
return True
|
|
else:
|
|
print(f'Search error: {resp.text}')
|
|
return False
|
|
|
|
def main():
|
|
print('Testing LLM region restriction...')
|
|
# Step 1: Upload document
|
|
print('Uploading document...')
|
|
result = upload_document()
|
|
if not result:
|
|
sys.exit(1)
|
|
doc_id = result.get('document_id')
|
|
print(f'Document ID: {doc_id}')
|
|
|
|
# Step 2: Wait for indexing
|
|
print('Waiting for indexing...')
|
|
if not wait_for_indexing(doc_id):
|
|
sys.exit(1)
|
|
|
|
# Step 3: Test search
|
|
print('Testing search with keyword extraction...')
|
|
if test_search():
|
|
print('✅ LLM region restriction passed!')
|
|
else:
|
|
print('❌ LLM region restriction may still be an issue')
|
|
sys.exit(1)
|
|
|
|
if __name__ == '__main__':
|
|
main() |