74 lines
2.9 KiB
Python
74 lines
2.9 KiB
Python
import requests
|
|
import json
|
|
|
|
base_url = 'http://localhost:3015'
|
|
api_key = 'jleu1212'
|
|
|
|
def test_search_with_workspace(workspace, query):
|
|
headers = {'X-API-Key': api_key, 'X-Workspace': workspace, 'Content-Type': 'application/json'}
|
|
payload = {'query': query, 'top_k': 5}
|
|
resp = requests.post(f'{base_url}/api/search', headers=headers, json=payload, timeout=30)
|
|
print(f'Workspace: {workspace}, Query: {query}, Status: {resp.status_code}')
|
|
if resp.status_code == 200:
|
|
data = resp.json()
|
|
print(f' Entities: {len(data.get("entities", []))}, Chunks: {len(data.get("chunks", []))}')
|
|
if data.get('chunks'):
|
|
print(f' First chunk preview: {data["chunks"][0].get("text", "")[:100]}...')
|
|
return True
|
|
else:
|
|
print(f' Error: {resp.text}')
|
|
return False
|
|
|
|
def test_query_with_workspace(workspace, query):
|
|
headers = {'X-API-Key': api_key, 'X-Workspace': workspace, 'Content-Type': 'application/json'}
|
|
payload = {'query': query, 'top_k': 5}
|
|
resp = requests.post(f'{base_url}/query/stream', headers=headers, json=payload, timeout=30)
|
|
print(f'Workspace: {workspace}, Query: {query}, Status: {resp.status_code}')
|
|
if resp.status_code == 200:
|
|
# stream response lines
|
|
lines = resp.text.strip().split('\n')
|
|
for line in lines:
|
|
if line:
|
|
try:
|
|
obj = json.loads(line)
|
|
if 'response' in obj:
|
|
print(f' Response: {obj["response"][:100]}...')
|
|
break
|
|
except:
|
|
pass
|
|
return True
|
|
else:
|
|
print(f' Error: {resp.text}')
|
|
return False
|
|
|
|
def main():
|
|
print('Testing workspace isolation...')
|
|
# Workspace test1 has tir.docx (contains "minimum safe working distance"?)
|
|
# Workspace test2 may have other documents (maybe ocr.pdf)
|
|
# Workspace '' (default) may have some data.
|
|
|
|
# Search for "minimum safe working distance" with workspace test1
|
|
print('\n--- Search with workspace test1 ---')
|
|
test_search_with_workspace('test1', 'minimum safe working distance')
|
|
|
|
# Search for "ODDS" with workspace test1
|
|
print('\n--- Search with workspace test1 for ODDS ---')
|
|
test_search_with_workspace('test1', 'ODDS')
|
|
|
|
# Search for "minimum safe working distance" with workspace '' (default)
|
|
print('\n--- Search with workspace "" (default) ---')
|
|
test_search_with_workspace('', 'minimum safe working distance')
|
|
|
|
# Search for "ODDS" with workspace ''
|
|
print('\n--- Search with workspace "" for ODDS ---')
|
|
test_search_with_workspace('', 'ODDS')
|
|
|
|
# Query stream test
|
|
print('\n--- Query stream with workspace test1 ---')
|
|
test_query_with_workspace('test1', 'what is the minimum safe working distance')
|
|
|
|
print('\n--- Query stream with workspace "" ---')
|
|
test_query_with_workspace('', 'what is the minimum safe working distance')
|
|
|
|
if __name__ == '__main__':
|
|
main() |