test_vector_db.py

import requests
import json
import time

BASE_URL = "http://localhost:5000"

def test_stats():
    response = requests.get(f"{BASE_URL}/stats")
    print("Test /stats:", response.status_code, response.json())

def test_add():
    payload = {
        "content": "Test content for vector DB.",
        "metadata": {"source": "test"},
        "content_type": "text"
    }
    response = requests.post(f"{BASE_URL}/add", json=payload)
    print("Test /add:", response.status_code, response.json())

def test_search():
    payload = {
        "query": "test",
        "n_results": 1,
        "content_type": "text"
    }
    response = requests.post(f"{BASE_URL}/search", json=payload)
    print("Test /search:", response.status_code, response.json())

def test_all():
    response = requests.get(f"{BASE_URL}/all?limit=5")
    print("Test /all:", response.status_code, response.json())

def test_delete(doc_id=None):
    if doc_id:
        response = requests.delete(f"{BASE_URL}/delete/{doc_id}")
        print("Test /delete:", response.status_code, response.json())
    else:
        print("No doc_id for delete test")

def test_chat_rag(query="What is in the test content?"):
    # Simulate history with query
    history = [{"role": "user", "content": query}]
    payload = {"history": history}
    response = requests.post(f"{BASE_URL}/chat", json=payload)
    print("Test /chat RAG:", response.status_code, response.json())

def test_analyze0():
    # Create a simple test file
    with open("test.txt", "w") as f:
        f.write("Test document for analysis.")
    files = {"file": open("test.txt", "rb")}
    response = requests.post(f"{BASE_URL}/analyze0", files=files)
    print("Test /analyze0:", response.status_code, response.json())
    files["file"].close()
    # Clean up
    import os
    os.remove("test.txt")

if __name__ == "__main__":
    print("Starting vector DB tests...")
    test_stats()
    test_add()  # Will fail due to quota
    time.sleep(1)
    test_search()
    test_all()
    # Assume doc_id from add if successful, skip delete for now
    test_chat_rag()
    test_analyze0()  # Will analyze but fail store due to quota
    print("Tests complete. Note: Storage fails due to Gemini quota exceeded.")