workspace working

test_llm_region.py

@@ -0,0 +1,104 @@
+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()