topoteretes--cognee/cognee/tests/performance/locust_performance_analysis.py

"""
Cognee Load Tests using Locust

Usage examples:
    # Note: ENV variables must be set before usage.

    # Multi-dataset test through the web UI
    locust -f locust_performance_analysis.py --host http://your-cognee-instance.com MultiDatasetCogneeTest

    # Single-dataset test through the web UI
    locust -f locust_performance_analysis.py --host http://your-cognee-instance.com SingleDatasetCogneeTest

    # Both simultaneously
    locust -f locust_performance_analysis.py --host http://your-cognee-instance.com

    # Headles without UI, 10 Locust users, 5 minutes
    locust -f locust_performance_analysis.py --host http://your-cognee-instance.com \
    --headless -u 10 -r 2 --run-time 5m --csv results

Environment variables:
    COGNEE_API_KEY          - API key for authentication (required)
"""

import csv
import io
import os
import random
import uuid

from pathlib import Path
from datetime import datetime
from locust import HttpUser, SequentialTaskSet, between, events, tag, task

API_KEY = os.environ.get("COGNEE_API_KEY", "")
SEARCH_TYPE = os.environ.get("COGNEE_SEARCH_TYPE", "GRAPH_COMPLETION")

ENDPOINT_NAMES = ("/api/v1/add", "/api/v1/cognify", "/api/v1/search")


def read_endpoint_averages(stats_csv_path: Path) -> dict[str, float]:
    """Parse a locust ``_stats.csv`` and return ``{endpoint_name: avg_ms}``."""
    averages: dict[str, float] = {}
    with stats_csv_path.open(newline="") as f:
        for row in csv.DictReader(f):
            if row["Name"] in ENDPOINT_NAMES:
                averages[row["Name"]] = float(row["Average Response Time"])
    return averages


TOPICS = [
    "quantum computing",
    "machine learning",
    "climate change",
    "renewable energy",
    "space exploration",
    "genetic engineering",
    "blockchain technology",
    "artificial intelligence",
    "ocean conservation",
    "urban planning",
    "medieval history",
    "philosophy of mind",
    "distributed systems",
    "neuroscience",
    "economic theory",
]

SUBTOPICS = [
    "data analysis",
    "pattern recognition",
    "resource allocation",
    "risk assessment",
    "optimization algorithms",
    "predictive modeling",
    "system integration",
    "scalability",
    "error correction",
    "signal processing",
    "network topology",
    "feedback loops",
    "energy efficiency",
    "material science",
    "behavioral adaptation",
    "information theory",
]

SENTENCE_TEMPLATES = [
    "The field of {topic} has seen remarkable advances in recent years, particularly in the area of {subtopic}.",
    "Researchers studying {topic} have discovered that {subtopic} plays a crucial role in understanding the broader implications.",
    "A comprehensive review of {topic} literature reveals that {subtopic} remains one of the most debated aspects.",
    "Recent experiments in {topic} demonstrate a strong correlation between {subtopic} and observed outcomes.",
    "The intersection of {topic} and {subtopic} opens new possibilities for practical applications.",
    "Experts in {topic} argue that {subtopic} will be the defining challenge of the next decade.",
    "Historical analysis shows that {topic} has always been influenced by developments in {subtopic}.",
    "New computational models for {topic} suggest that {subtopic} can be optimized through iterative approaches.",
    "The economic impact of {topic} is closely tied to advancements in {subtopic}, according to recent studies.",
    "Collaborative efforts in {topic} have led to breakthroughs in {subtopic} that were previously thought impossible.",
    "Understanding {topic} requires a deep appreciation of how {subtopic} interacts with existing frameworks.",
    "Policy makers are increasingly turning to {topic} research to inform decisions about {subtopic}.",
]

SEARCH_QUERIES = [
    "What are the main findings about {topic}?",
    "How does {topic} relate to recent developments?",
    "What are the key challenges in {topic}?",
    "Summarize the information about {topic}.",
    "What practical applications exist for {topic}?",
    "What is the current state of research in {topic}?",
]


def generate_paragraph(topic: str, num_sentences: int = 5) -> str:
    sentences = []
    for _ in range(num_sentences):
        template = random.choice(SENTENCE_TEMPLATES)
        subtopic = random.choice(SUBTOPICS)
        sentences.append(template.format(topic=topic, subtopic=subtopic))
    return " ".join(sentences)


def generate_document(num_paragraphs: int = 3) -> tuple[str, str]:
    """Returns (text, primary_topic) for a randomly generated document. Along with random UUID to ensure uniqueness across test runs."""
    topic = random.choice(TOPICS)
    paragraphs = [
        generate_paragraph(topic, num_sentences=random.randint(300, 377))
        for _ in range(num_paragraphs)
    ]
    paragraphs.append(str(uuid.uuid4()))  # Ensure uniqueness of document text across test runs
    return "\n\n".join(paragraphs), topic


def generate_search_query(topic: str) -> str:
    template = random.choice(SEARCH_QUERIES)
    return template.format(topic=topic)


@events.init_command_line_parser.add_listener
def add_custom_arguments(parser):
    parser.add_argument(
        "--cognee-api-key",
        type=str,
        default="",
        help="API key for Cognee (overrides COGNEE_API_KEY env var)",
    )


@events.test_start.add_listener
def on_test_start(environment, **kwargs):
    api_key = environment.parsed_options.cognee_api_key or API_KEY
    if not api_key:
        environment.runner.quit()
        raise SystemExit(
            "No API key provided. Set COGNEE_API_KEY env var or pass --cognee-api-key."
        )


class AddCognifySearchFlow(SequentialTaskSet):
    """
    Sequential flow: Add text → Cognify → Search.
    Subclasses set `dataset_name` and `api_key`.
    """

    dataset_name: str = ""
    api_key: str = ""
    topic: str = ""

    def _headers(self) -> dict:
        return {
            "X-Api-Key": self.api_key,
        }

    def _json_headers(self) -> dict:
        return {
            "X-Api-Key": self.api_key,
            "Content-Type": "application/json",
        }

    @tag("add")
    @task
    def add_text(self):
        # 1. Generate the text content
        text, self.topic = generate_document(num_paragraphs=random.randint(2, 5))

        # 2. Prepare the form data
        form_data = {
            "datasetName": self.dataset_name,
        }

        # 3. Prepare the file data
        files = [("data", ("document.txt", io.BytesIO(text.encode("utf-8")), "text/plain"))]

        # 4. Make the request
        with self.client.post(
            "/api/v1/add",
            data=form_data,
            files=files,
            headers=self._headers(),
            name="/api/v1/add",
            catch_response=True,
            timeout=600,
        ) as resp:
            if resp.status_code == 200:
                resp.success()
            else:
                # Enhanced error logging to capture the ErrorResponse model from FastAPI
                resp.failure(f"Add failed: {resp.status_code} - {resp.text[:300]}")

    @tag("cognify")
    @task
    def cognify(self):
        payload = {
            "datasets": [self.dataset_name],
            "runInBackground": False,
        }
        with self.client.post(
            "/api/v1/cognify",
            json=payload,
            headers=self._json_headers(),
            name="/api/v1/cognify",
            catch_response=True,
            timeout=6000,
        ) as resp:
            if resp.status_code == 200:
                resp.success()
            else:
                resp.failure(f"Cognify failed: {resp.status_code} - {resp.text[:300]}")

    @tag("search")
    @task
    def search(self):
        query = generate_search_query(self.topic or random.choice(TOPICS))
        payload = {
            "searchType": SEARCH_TYPE,
            "query": query,
            "datasets": [self.dataset_name],
            "only_context": True,
        }
        with self.client.post(
            "/api/v1/search",
            json=payload,
            headers=self._json_headers(),
            name="/api/v1/search",
            catch_response=True,
            timeout=200,
        ) as resp:
            if resp.status_code == 200:
                resp.success()
            else:
                resp.failure(f"Search failed: {resp.status_code} - {resp.text[:300]}")
        self.interrupt()


# ---------------------------------------------------------------------------
# Scenario 1: Multiple datasets
# ---------------------------------------------------------------------------


class MultiDatasetFlow(AddCognifySearchFlow):
    """Each call a unique dataset."""

    def on_start(self):
        uid = uuid.uuid4().hex[:8]
        self.dataset_name = f"loadtest_user_{uid}"
        self.api_key = self.user.environment.parsed_options.cognee_api_key or API_KEY


class MultiDatasetCogneeTest(HttpUser):
    """
    Scenario 1 – Different virtual users (Locust users) making their own Add → Cognify → Search
    on different datasets.

    Each spawned user creates a unique dataset name and runs the full
    pipeline independently. They use the same API key, simulating multiple users (locust users) from the same account creating
    separate datasets and running operations on them.
    """

    tasks = [MultiDatasetFlow]
    wait_time = between(5, 10)
    weight = 2


# ---------------------------------------------------------------------------
# Scenario 2: Single user, single shared dataset
# ---------------------------------------------------------------------------

SHARED_DATASET_NAME = f"loadtest_shared_{uuid.uuid4().hex[:8]}"


class SingleDatasetFlow(AddCognifySearchFlow):
    """All virtual users share the same dataset and API key."""

    # TODO: Fix SingleDatasetFlow add/cognify/search operations keep being called async which makes perf measurements unreliable.
    #       Cognify operation on a dataset that is currently being cognified returns early.
    def on_start(self):
        self.dataset_name = SHARED_DATASET_NAME
        self.api_key = self.user.environment.parsed_options.cognee_api_key or API_KEY


class SingleDatasetCogneeTest(HttpUser):
    """
    Scenario 2. calling Add → Cognify → Search on a
    single dataset. All virtual users share one dataset name and API key,
    simulating concurrent calls from the same account.
    """

    tasks = [SingleDatasetFlow]
    wait_time = between(5, 10)
    # TODO: Return weight when scenario is fixed to handle concurrent cognify calls on the same dataset.
    weight = 0


if __name__ == "__main__":
    import signal
    import subprocess
    import sys
    import tempfile
    import time
    import urllib.error
    import urllib.request

    def wait_for_server(url: str, timeout: float = 240.0) -> None:
        deadline = time.time() + timeout
        while time.time() < deadline:
            try:
                with urllib.request.urlopen(url, timeout=2) as resp:
                    if resp.status == 200:
                        return
            except (urllib.error.URLError, ConnectionError, TimeoutError):
                pass
            time.sleep(0.5)
        raise SystemExit(f"Cognee server at {url} did not become ready in {timeout}s")

    key_path = tempfile.NamedTemporaryFile(suffix=".key", delete=False).name
    try:
        # Generate API key in a separate process to avoid any potential issues with locust's monkey-patching of libraries like gevent.
        perf_dir = str(Path(__file__).resolve().parent)
        subprocess.run(
            [sys.executable, "-m", "utils.bootstrap_script", key_path],
            check=True,
            cwd=perf_dir,
        )
        api_key = Path(key_path).read_text().strip()
    finally:
        try:
            os.unlink(key_path)
        except FileNotFoundError:
            pass

    host = os.environ.get("HTTP_API_HOST", "localhost")
    port = os.environ.get("HTTP_API_PORT", "8000")
    base_url = f"http://{host}:{port}"

    server_proc = subprocess.Popen(
        [sys.executable, "-m", "uvicorn", "cognee.api.client:app", "--host", host, "--port", port],
        stdout=subprocess.DEVNULL,
        stderr=subprocess.DEVNULL,
        start_new_session=True,
    )
    try:
        wait_for_server(f"{base_url}/health")
        env = {**os.environ, "COGNEE_API_KEY": api_key}
        # Timestamped results to avoid overwriting previous runs and for easier identification of test runs.
        now = datetime.now()
        timestamp = now.strftime("%Y%m%d_%H%M%S")
        # Results will be saved in the `results` directory with a unique name based on the timestamp of the test run.
        result_folder = Path("results")
        result_folder.mkdir(exist_ok=True)
        result_location = result_folder / f"locust_run_{timestamp}"
        html_result_location = f"{result_location}.html"
        locust_log_location = Path(f"{result_location}.log")
        locust_log_location.touch()

        # Run locust command
        cmd = [
            sys.executable,
            "-m",
            "locust",
            "-f",
            __file__,
            "--host",
            base_url,
            "--csv",
            result_location,
            "--html",
            html_result_location,
            "--logfile",
            locust_log_location,
            "--headless",
            "-u",
            "10",
            "-r",
            "1",
            "--run-time",
            "5min",
            "MultiDatasetCogneeTest",  # Run only the multi-dataset scenario
            *sys.argv[1:],
        ]

        rc = subprocess.run(cmd, env=env).returncode
    finally:
        try:
            os.killpg(server_proc.pid, signal.SIGTERM)
            server_proc.wait(timeout=10)
        except ProcessLookupError:
            pass
        except subprocess.TimeoutExpired:
            try:
                os.killpg(server_proc.pid, signal.SIGKILL)
            except ProcessLookupError:
                pass

    stats_csv = Path(f"{result_location}_stats.csv")
    if stats_csv.exists():
        averages = read_endpoint_averages(stats_csv)
        print("\n=== Average response times (ms) ===")
        for name in ENDPOINT_NAMES:
            avg = averages.get(name)
            if avg is None:
                print(f"  {name}: (no requests recorded)")
            else:
                print(f"  {name}: {avg:.0f} ms")
    else:
        print(f"\nNo stats CSV found at {stats_csv}; skipping averages.")

    sys.exit(rc)