feature: add the test to workflows

.github/workflows/examples_tests.yml

@@ -259,3 +259,27 @@ jobs:
           EMBEDDING_API_KEY: ${{ secrets.EMBEDDING_API_KEY }}
           EMBEDDING_API_VERSION: ${{ secrets.EMBEDDING_API_VERSION }}
         run: uv run python ./cognee/tests/test_add_docling_document.py
+
+  test-custom-model:
+    name: Run Custom Model Test
+    runs-on: ubuntu-22.04
+    steps:
+      - name: Check out repository
+        uses: actions/checkout@v4
+
+      - name: Cognee Setup
+        uses: ./.github/actions/cognee_setup
+        with:
+          python-version: '3.11.x'
+
+      - name: Run Custom Model Test
+        env:
+          LLM_MODEL: ${{ secrets.LLM_MODEL }}
+          LLM_ENDPOINT: ${{ secrets.LLM_ENDPOINT }}
+          LLM_API_KEY: ${{ secrets.OPENAI_API_KEY }}
+          LLM_API_VERSION: ${{ secrets.LLM_API_VERSION }}
+          EMBEDDING_MODEL: ${{ secrets.EMBEDDING_MODEL }}
+          EMBEDDING_ENDPOINT: ${{ secrets.EMBEDDING_ENDPOINT }}
+          EMBEDDING_API_KEY: ${{ secrets.EMBEDDING_API_KEY }}
+          EMBEDDING_API_VERSION: ${{ secrets.EMBEDDING_API_VERSION }}
+        run: uv run python ./cognee/tests/test_custom_model.py

cognee-starter-kit/.env.template

@@ -1,19 +0,0 @@
-# In case you choose to use OpenAI provider, just adjust the model and api_key.
-LLM_API_KEY=""
-LLM_MODEL="openai/gpt-5-mini"
-LLM_PROVIDER="openai"
-# Not needed if you use OpenAI
-LLM_ENDPOINT=""
-LLM_API_VERSION=""
-
-# In case you choose to use OpenAI provider, just adjust the model and api_key.
-EMBEDDING_API_KEY=""
-EMBEDDING_MODEL="openai/text-embedding-3-large"
-EMBEDDING_PROVIDER="openai"
-# Not needed if you use OpenAI
-EMBEDDING_ENDPOINT=""
-EMBEDDING_API_VERSION=""
-
-
-GRAPHISTRY_USERNAME=""
-GRAPHISTRY_PASSWORD=""

cognee-starter-kit/.gitignore

@@ -1,196 +0,0 @@
-.data
-.env
-.local.env
-.prod.env
-cognee/.data/
-
-code_pipeline_output*/
-
-*.lance/
-.DS_Store
-# Byte-compiled / optimized / DLL files
-__pycache__/
-*.py[cod]
-*$py.class
-
-full_run.ipynb
-
-# C extensions
-*.so
-
-# Distribution / packaging
-.Python
-build/
-develop-eggs/
-dist/
-downloads/
-eggs/
-.eggs/
-lib/
-lib64/
-parts/
-sdist/
-var/
-wheels/
-share/python-wheels/
-*.egg-info/
-.installed.cfg
-*.egg
-MANIFEST
-
-# PyInstaller
-#  Usually these files are written by a python script from a template
-#  before PyInstaller builds the exe, so as to inject date/other infos into it.
-*.manifest
-*.spec
-
-# Installer logs
-pip-log.txt
-pip-delete-this-directory.txt
-
-# Unit test / coverage reports
-htmlcov/
-.tox/
-.nox/
-.coverage
-.coverage.*
-.cache
-nosetests.xml
-coverage.xml
-*.cover
-*.py,cover
-.hypothesis/
-.pytest_cache/
-cover/
-
-# Translations
-*.mo
-*.pot
-
-# Django stuff:
-*.log
-local_settings.py
-db.sqlite3
-db.sqlite3-journal
-
-# Cognee logs directory - keep directory, ignore contents
-logs/*
-!logs/.gitkeep
-!logs/README.md
-
-# Flask stuff:
-instance/
-.webassets-cache
-
-# Scrapy stuff:
-.scrapy
-
-# Sphinx documentation
-docs/_build/
-
-# PyBuilder
-.pybuilder/
-target/
-
-# Jupyter Notebook
-.ipynb_checkpoints
-
-# IPython
-profile_default/
-ipython_config.py
-
-# pyenv
-#   For a library or package, you might want to ignore these files since the code is
-#   intended to run in multiple environments; otherwise, check them in:
-# .python-version
-
-# pipenv
-#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
-#   However, in case of collaboration, if having platform-specific dependencies or dependencies
-#   having no cross-platform support, pipenv may install dependencies that don't work, or not
-#   install all needed dependencies.
-#Pipfile.lock
-
-# poetry
-#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
-#   This is especially recommended for binary packages to ensure reproducibility, and is more
-#   commonly ignored for libraries.
-#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
-#poetry.lock
-
-# pdm
-#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
-#pdm.lock
-#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
-#   in version control.
-#   https://pdm.fming.dev/#use-with-ide
-.pdm.toml
-
-# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
-__pypackages__/
-
-# Celery stuff
-celerybeat-schedule
-celerybeat.pid
-
-# SageMath parsed files
-*.sage.py
-
-# Environments
-.env
-.env.local
-.venv
-env/
-venv/
-ENV/
-env.bak/
-venv.bak/
-
-# Spyder project settings
-.spyderproject
-.spyproject
-
-# Rope project settings
-.ropeproject
-
-# mkdocs documentation
-/site
-
-# mypy
-.mypy_cache/
-.dmypy.json
-dmypy.json
-
-# Pyre type checker
-.pyre/
-
-# pytype static type analyzer
-.pytype/
-
-# Cython debug symbols
-cython_debug/
-
-# PyCharm
-#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
-#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
-#  and can be added to the global gitignore or merged into this file.  For a more nuclear
-#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
-.idea/
-
-.vscode/
-cognee/data/
-cognee/cache/
-
-# Default cognee system directory, used in development
-.cognee_system/
-.data_storage/
-.artifacts/
-.anon_id
-
-node_modules/
-
-# Evals
-SWE-bench_testsample/
-
-# ChromaDB Data
-.chromadb_data/

cognee-starter-kit/README.md

@@ -1,89 +0,0 @@
-
-# Cognee Starter Kit
-Welcome to the <a href="https://github.com/topoteretes/cognee">cognee</a> Starter Repo! This repository is designed to help you get started quickly by providing a structured dataset and pre-built data pipelines using cognee to build powerful knowledge graphs.
-
-You can use this repo to ingest, process, and visualize data in minutes. 
-
-By following this guide, you will:
-
-- Load structured company and employee data
-- Utilize pre-built pipelines for data processing
-- Perform graph-based search and query operations
-- Visualize entity relationships effortlessly on a graph
-
-# How to Use This Repo 🛠
-
-## Install uv if you don't have it on your system
-```
-pip install uv
-```
-## Install dependencies
-```
-uv sync
-```
-
-## Setup LLM
-Add environment variables to `.env` file.
-In case you choose to use OpenAI provider, add just the model and api_key.
-```
-LLM_PROVIDER=""
-LLM_MODEL=""
-LLM_ENDPOINT=""
-LLM_API_KEY=""
-LLM_API_VERSION=""
-
-EMBEDDING_PROVIDER=""
-EMBEDDING_MODEL=""
-EMBEDDING_ENDPOINT=""
-EMBEDDING_API_KEY=""
-EMBEDDING_API_VERSION=""
-```
-
-Activate the Python environment:
-```
-source .venv/bin/activate
-```
-
-## Run the Default Pipeline
-
-This script runs the cognify pipeline with default settings. It ingests text data, builds a knowledge graph, and allows you to run search queries.
-
-```
-python src/pipelines/default.py
-```
-
-## Run the Low-Level Pipeline
-
-This script implements its own pipeline with custom ingestion task. It processes the given JSON data about companies and employees, making it searchable via a graph.
-
-```
-python src/pipelines/low_level.py
-```
-
-## Run the Custom Model Pipeline
-
-Custom model uses custom pydantic model for graph extraction. This script categorizes programming languages as an example and visualizes relationships.
-
-```
-python src/pipelines/custom-model.py
-```
-
-## Graph preview 
-
-cognee provides a visualize_graph function that will render the graph for you.
-
-```
-    graph_file_path = str(
-        pathlib.Path(
-            os.path.join(pathlib.Path(__file__).parent, ".artifacts/graph_visualization.html")
-        ).resolve()
-    )
-    await visualize_graph(graph_file_path)
-```
-
-# What will you build with cognee?
-
-- Expand the dataset by adding more structured/unstructured data
-- Customize the data model to fit your use case
-- Use the search API to build an intelligent assistant
-- Visualize knowledge graphs for better insights

cognee-starter-kit/pyproject.toml

@@ -1,11 +0,0 @@
-[project]
-name = "cognee-starter"
-version = "0.1.1"
-description = "Starter project which can be harvested for parts"
-readme = "README.md"
-
-requires-python = ">=3.10, <=3.13"
-
-dependencies = [
-    "cognee>=0.1.38,<1.0.0",
-]

cognee-starter-kit/src/data/companies.json

@@ -1,38 +0,0 @@
-[
-  {
-    "name": "TechNova Inc.",
-    "departments": [
-      "Engineering",
-      "Marketing"
-    ]
-  },
-  {
-    "name": "GreenFuture Solutions",
-    "departments": [
-      "Research & Development",
-      "Sales",
-      "Customer Support"
-    ]
-  },
-  {
-    "name": "Skyline Financials",
-    "departments": [
-      "Accounting"
-    ]
-  },
-  {
-    "name": "MediCare Plus",
-    "departments": [
-      "Healthcare",
-      "Administration"
-    ]
-  },
-  {
-    "name": "NextGen Robotics",
-    "departments": [
-      "AI Development",
-      "Manufacturing",
-      "HR"
-    ]
-  }
-]

cognee-starter-kit/src/data/people.json

@@ -1,52 +0,0 @@
-[
-  {
-    "name": "John Doe",
-    "company": "TechNova Inc.",
-    "department": "Engineering"
-  },
-  {
-    "name": "Jane Smith",
-    "company": "TechNova Inc.",
-    "department": "Marketing"
-  },
-  {
-    "name": "Alice Johnson",
-    "company": "GreenFuture Solutions",
-    "department": "Sales"
-  },
-  {
-    "name": "Bob Williams",
-    "company": "GreenFuture Solutions",
-    "department": "Customer Support"
-  },
-  {
-    "name": "Michael Brown",
-    "company": "Skyline Financials",
-    "department": "Accounting"
-  },
-  {
-    "name": "Emily Davis",
-    "company": "MediCare Plus",
-    "department": "Healthcare"
-  },
-  {
-    "name": "David Wilson",
-    "company": "MediCare Plus",
-    "department": "Administration"
-  },
-  {
-    "name": "Emma Thompson",
-    "company": "NextGen Robotics",
-    "department": "AI Development"
-  },
-  {
-    "name": "Chris Martin",
-    "company": "NextGen Robotics",
-    "department": "Manufacturing"
-  },
-  {
-    "name": "Sophia White",
-    "company": "NextGen Robotics",
-    "department": "HR"
-  }
-]

cognee-starter-kit/src/pipelines/default.py

@@ -1,67 +0,0 @@
-import os
-import asyncio
-import pathlib
-from cognee import config, add, cognify, search, SearchType, prune, visualize_graph
-
-
-async def main():
-    data_directory_path = str(
-        pathlib.Path(os.path.join(pathlib.Path(__file__).parent, ".data_storage")).resolve()
-    )
-    # Set up the data directory. Cognee will store files here.
-    config.data_root_directory(data_directory_path)
-
-    cognee_directory_path = str(
-        pathlib.Path(os.path.join(pathlib.Path(__file__).parent, ".cognee_system")).resolve()
-    )
-    # Set up the Cognee system directory. Cognee will store system files and databases here.
-    config.system_root_directory(cognee_directory_path)
-
-    # Prune data and system metadata before running, only if we want "fresh" state.
-    await prune.prune_data()
-    await prune.prune_system(metadata=True)
-
-    text = "The Python programming language is widely used in data analysis, web development, and machine learning."
-
-    # Add the text data to Cognee.
-    await add(text)
-
-    # Cognify the text data.
-    await cognify()
-
-    # Or use our simple graph preview
-    graph_file_path = str(
-        pathlib.Path(
-            os.path.join(pathlib.Path(__file__).parent, ".artifacts/graph_visualization.html")
-        ).resolve()
-    )
-    await visualize_graph(graph_file_path)
-
-    # Completion query that uses graph data to form context.
-    graph_completion = await search(
-        query_text="What is python?", query_type=SearchType.GRAPH_COMPLETION
-    )
-    print("Graph completion result is:")
-    print(graph_completion)
-
-    # Completion query that uses document chunks to form context.
-    rag_completion = await search(
-        query_text="What is Python?", query_type=SearchType.RAG_COMPLETION
-    )
-    print("Completion result is:")
-    print(rag_completion)
-
-    # Query all summaries related to query.
-    summaries = await search(query_text="Python", query_type=SearchType.SUMMARIES)
-    print("Summary results are:")
-    for summary in summaries:
-        print(summary)
-
-    chunks = await search(query_text="Python", query_type=SearchType.CHUNKS)
-    print("Chunk results are:")
-    for chunk in chunks:
-        print(chunk)
-
-
-if __name__ == "__main__":
-    asyncio.run(main())

cognee/tests/test_custom_model.py

@@ -1,14 +1,8 @@
 import os
 import pathlib
 import cognee
-from cognee.modules.search.operations import get_history
-from cognee.modules.users.methods import get_default_user
-from cognee.shared.logging_utils import get_logger
-from cognee.modules.search.types import SearchType
 from cognee.low_level import DataPoint
-
-logger = get_logger()
-
+from cognee.infrastructure.databases.graph import get_graph_engine
 
 async def main():
     data_directory_path = str(
@@ -53,48 +47,114 @@ async def main():
     )
 
     await cognee.add(text)
-
     await cognee.cognify(graph_model=ProgrammingLanguage)
 
-    graph_file_path = str(
-        pathlib.Path(
-            os.path.join(
-                pathlib.Path(__file__).parent,
-                ".artifacts/test_custom_model/graph_visualization.html",
-            )
-        ).resolve()
-    )
-    await cognee.visualize_graph(graph_file_path)
+    
+    await cognee.visualize_graph(destination_file_path="cognee/tests/test_custom_model.html")
 
-    # Completion query that uses graph data to form context.
-    completion = await cognee.search(SearchType.GRAPH_COMPLETION, "What is python?")
-    assert len(completion) != 0, "Graph completion search didn't return any result."
-    print("Graph completion result is:")
-    print(completion)
+    graph_engine = await get_graph_engine()
+    
+    graph_db_provider = os.getenv("GRAPH_DATABASE_PROVIDER", "kuzu").lower()
+    
+    # Query for Python entity and verify it exists with correct type
+    python_found = False
+    python_type = None
+    
+    if graph_db_provider in ["neo4j", "neptune", "neptune_analytics"]:
+        query = """
+        MATCH (n)
+        WHERE n.name = 'Python'
+        RETURN n.name as name, n.type as type
+        """
+        results = await graph_engine.query(query)
+        if results:
+            python_found = True
+            python_type = results[0]["type"]
+    
+    elif graph_db_provider == "kuzu":
+        query = """
+        MATCH (n:Node)
+        WHERE n.name = 'Python'
+        RETURN n.name, n.type
+        """
+        results = await graph_engine.query(query)
+        if results:
+            python_found = True
+            python_type = results[0][1]
+    
+    else:
+        raise ValueError(f"Unsupported graph database provider: {graph_db_provider}")
 
-    # Completion query that uses document chunks to form context.
-    completion = await cognee.search(SearchType.RAG_COMPLETION, "What is Python?")
-    assert len(completion) != 0, "Completion search didn't return any result."
-    print("Completion result is:")
-    print(completion)
+    assert python_found, "Python entity was not extracted from the text"
+    assert python_type == "ProgrammingLanguage", f"Python entity has incorrect type: {python_type}, expected: ProgrammingLanguage"
+    
+    # Query for entities that should NOT exist (Guido van Rossum and 1991)
+    guido_found = False
+    year_1991_found = False
+    
+    if graph_db_provider in ["neo4j", "neptune", "neptune_analytics"]:
+        query = """
+        MATCH (n)
+        WHERE n.name IN ['Guido van Rossum', '1991']
+        RETURN n.name as name
+        """
+        results = await graph_engine.query(query)
+        for result in results:
+            if result["name"] == "Guido van Rossum":
+                guido_found = True
+            elif result["name"] == "1991":
+                year_1991_found = True
+    
+    elif graph_db_provider == "kuzu":
+        query = """
+        MATCH (n:Node)
+        WHERE n.name IN ['Guido van Rossum', '1991']
+        RETURN n.name
+        """
+        results = await graph_engine.query(query)
+        for result in results:
+            if result[0] == "Guido van Rossum":
+                guido_found = True
+            elif result[0] == "1991":
+                year_1991_found = True
+    
+    else:
+        raise ValueError(f"Unsupported graph database provider: {graph_db_provider}")
 
-    # Query all summaries related to query.
-    summaries = await cognee.search(SearchType.SUMMARIES, "Python")
-    assert len(summaries) != 0, "Summaries search didn't return any results."
-    print("Summary results are:")
-    for summary in summaries:
-        print(summary)
+    
+    assert not guido_found, "Guido van Rossum should not be extracted as it's not in the custom graph model"
+    assert not year_1991_found, "1991 should not be extracted as it's not in the custom graph model"
+    
+    # Query for Field entities that might have been extracted (data analysis, web development, machine learning)
+    field_entities = []
+    
+    if graph_db_provider in ["neo4j", "neptune", "neptune_analytics"]:
+        query = """
+        MATCH (n)
+        WHERE n.type = 'Field'
+        RETURN n.name as name
+        """
+        results = await graph_engine.query(query)
+        field_entities = [r["name"] for r in results]
+    
+    elif graph_db_provider == "kuzu":
+        query = """
+        MATCH (n:Node)
+        WHERE n.type = 'Field'
+        RETURN n.name
+        """
+        results = await graph_engine.query(query)
+        field_entities = [r[0] for r in results if r[0]]
+    else:
+        raise ValueError(f"Unsupported graph database provider: {graph_db_provider}")
 
-    chunks = await cognee.search(SearchType.CHUNKS, query_text="Python")
-    assert len(chunks) != 0, "Chunks search didn't return any results."
-    print("Chunk results are:")
-    for chunk in chunks:
-        print(chunk)
-
-    user = await get_default_user()
-    history = await get_history(user.id)
-
-    assert len(history) == 8, "Search history is not correct."
+    assert len(field_entities) > 0, f"No Field entities were extracted. Expected fields like 'data analysis', 'web development', 'machine learning' but got: {field_entities}"
+    
+    expected_fields = ["data analysis", "web development", "machine learning"]
+    found_expected_fields = [f for f in expected_fields if any(f in field.lower() for field in field_entities)]
+    
+    assert len(found_expected_fields) > 0, f"None of the expected Field entities were found. Expected at least one of {expected_fields}, but got: {field_entities}"
+    
 
 
 if __name__ == "__main__":

cognee/tests/test_starter_pipelines.py

@@ -1,66 +0,0 @@
-import unittest
-import subprocess
-import os
-import sys
-
-
-class TestPipelines(unittest.TestCase):
-    """Tests that all pipelines run successfully."""
-
-    def setUp(self):
-        # Ensure we're in the correct directory
-        self.project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), "../.."))
-        self.pipelines_dir = os.path.join(self.project_root, "src", "pipelines")
-
-        # Required environment variables
-        self.required_env_vars = ["LLM_API_KEY", "EMBEDDING_API_KEY"]
-
-        # Check if required environment variables are set
-        missing_vars = [var for var in self.required_env_vars if not os.environ.get(var)]
-        if missing_vars:
-            self.skipTest(f"Missing required environment variables: {', '.join(missing_vars)}")
-
-    def _run_pipeline(self, script_name):
-        """Helper method to run a pipeline script and return the result."""
-        script_path = os.path.join(self.pipelines_dir, script_name)
-
-        # Use the Python executable from the virtual environment
-        python_exe = os.path.join(self.project_root, ".venv", "bin", "python")
-        if not os.path.exists(python_exe):
-            python_exe = sys.executable
-
-        try:
-            result = subprocess.run(
-                [python_exe, script_path],
-                check=True,
-                capture_output=True,
-                text=True,
-                timeout=300,  # 5 minute timeout
-            )
-            return result
-        except subprocess.CalledProcessError as e:
-            self.fail(
-                f"Pipeline {script_name} failed with code {e.returncode}. "
-                f"Stdout: {e.stdout}, Stderr: {e.stderr}"
-            )
-        except subprocess.TimeoutExpired:
-            self.fail(f"Pipeline {script_name} timed out after 300 seconds")
-
-    def test_default_pipeline(self):
-        """Test that the default pipeline runs successfully."""
-        result = self._run_pipeline("default.py")
-        self.assertEqual(result.returncode, 0)
-
-    def test_low_level_pipeline(self):
-        """Test that the low-level pipeline runs successfully."""
-        result = self._run_pipeline("low_level.py")
-        self.assertEqual(result.returncode, 0)
-
-    def test_custom_model_pipeline(self):
-        """Test that the custom model pipeline runs successfully."""
-        result = self._run_pipeline("custom-model.py")
-        self.assertEqual(result.returncode, 0)
-
-
-if __name__ == "__main__":
-    unittest.main()

examples/low_level/granular_level.py

@@ -49,12 +49,11 @@ class Company(DataPoint):
 
 
 ROOT = Path(__file__).resolve().parent
-DATA_DIR = ROOT.parent / "data"
 COGNEE_DIR = ROOT / ".cognee_system"
 ARTIFACTS_DIR = ROOT / ".artifacts"
 GRAPH_HTML = ARTIFACTS_DIR / "graph_visualization.html"
-COMPANIES_JSON = DATA_DIR / "companies.json"
-PEOPLE_JSON = DATA_DIR / "people.json"
+COMPANIES_JSON = ROOT / "companies.json"
+PEOPLE_JSON = ROOT / "people.json"
 
 
 def load_json_file(path: Path) -> Any: