Merge branch 'dev' of github.com:topoteretes/cognee into dev

.env.template

@@ -65,3 +65,7 @@ DB_NAME=cognee_db
 
 # LITELLM Logging Level. Set to quiten down logging
 LITELLM_LOG="ERROR"
+
+# Set this environment variable to disable sending telemetry data
+# TELEMETRY_DISABLED=1
+

.github/workflows/dockerhub-mcp.yml

@@ -36,9 +36,10 @@ jobs:
         id: build
         uses: docker/build-push-action@v5
         with:
-          context: cognee-mcp
+          context: .
           platforms: linux/amd64,linux/arm64
           push: true
+          file: cognee-mcp/Dockerfile
           tags: ${{ steps.meta.outputs.tags }}
           labels: ${{ steps.meta.outputs.labels }}
           cache-from: type=registry,ref=cognee/cognee-mcp:buildcache

.github/workflows/vector_db_tests.yml

@@ -17,49 +17,49 @@ on:
         required: false
 
 jobs:
-  run-chromadb-tests:
-    name: ChromaDB Tests
-    runs-on: ubuntu-22.04
-    if: ${{ inputs.databases == 'all' || contains(inputs.databases, 'chromadb') }}
-    services:
-      chromadb:
-        image: chromadb/chroma:0.6.3
-        env:
-          CHROMA_SERVER_AUTH_CREDENTIALS: "test-token"
-          CHROMA_SERVER_AUTH_CREDENTIALS_PROVIDER: "chromadb.auth.token.TokenAuthCredentialsProvider"
-          CHROMA_SERVER_AUTH_PROVIDER: "chromadb.auth.token.TokenAuthServerProvider"
-        ports:
-          - 3002:8000
-    steps:
-      - name: Check out
-        uses: actions/checkout@v4
-        with:
-          fetch-depth: 0
+  # run-chromadb-tests:
+  #   name: ChromaDB Tests
+  #   runs-on: ubuntu-22.04
+  #   if: ${{ inputs.databases == 'all' || contains(inputs.databases, 'chromadb') }}
+  #   services:
+  #     chromadb:
+  #       image: chromadb/chroma:0.6.3
+  #       env:
+  #         CHROMA_SERVER_AUTH_CREDENTIALS: "test-token"
+  #         CHROMA_SERVER_AUTH_CREDENTIALS_PROVIDER: "chromadb.auth.token.TokenAuthCredentialsProvider"
+  #         CHROMA_SERVER_AUTH_PROVIDER: "chromadb.auth.token.TokenAuthServerProvider"
+  #       ports:
+  #         - 3002:8000
+  #   steps:
+  #     - name: Check out
+  #       uses: actions/checkout@v4
+  #       with:
+  #         fetch-depth: 0
 
-      - name: Cognee Setup
-        uses: ./.github/actions/cognee_setup
-        with:
-          python-version: ${{ inputs.python-version }}
+  #     - name: Cognee Setup
+  #       uses: ./.github/actions/cognee_setup
+  #       with:
+  #         python-version: ${{ inputs.python-version }}
 
-      - name: Install specific db dependency
-        run: |
-          poetry install -E chromadb
+  #     - name: Install specific db dependency
+  #       run: |
+  #         poetry install -E chromadb
 
-      - name: Run ChromaDB Tests
-        env:
-          ENV: 'dev'
-          VECTOR_DB_PROVIDER: chromadb
-          VECTOR_DB_URL: http://localhost:3002
-          VECTOR_DB_KEY: test-token
-          LLM_MODEL: ${{ secrets.LLM_MODEL }}
-          LLM_ENDPOINT: ${{ secrets.LLM_ENDPOINT }}
-          LLM_API_KEY: ${{ secrets.LLM_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: poetry run python ./cognee/tests/test_chromadb.py
+  #     - name: Run ChromaDB Tests
+  #       env:
+  #         ENV: 'dev'
+  #         VECTOR_DB_PROVIDER: chromadb
+  #         VECTOR_DB_URL: http://localhost:3002
+  #         VECTOR_DB_KEY: test-token
+  #         LLM_MODEL: ${{ secrets.LLM_MODEL }}
+  #         LLM_ENDPOINT: ${{ secrets.LLM_ENDPOINT }}
+  #         LLM_API_KEY: ${{ secrets.LLM_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: poetry run python ./cognee/tests/test_chromadb.py
 
   run-weaviate-tests:
     name: Weaviate Tests

Dockerfile

@@ -5,7 +5,7 @@ ARG POETRY_EXTRAS="\
 # API \
 api \
 # Storage & Databases \
-filesystem postgres weaviate qdrant neo4j falkordb milvus kuzu chromadb \
+postgres weaviate qdrant neo4j falkordb milvus kuzu chromadb \
 # Notebooks & Interactive Environments \
 notebook \
 # LLM & AI Frameworks \
@@ -25,7 +25,7 @@ ENV DEBUG=${DEBUG}
 ENV PIP_NO_CACHE_DIR=true
 ENV PATH="${PATH}:/root/.poetry/bin"
 
-RUN apt-get update && apt-get install
+RUN apt-get update
 
 RUN apt-get install -y \
   gcc \

README.md

@@ -160,3 +160,4 @@ We are committed to making open source an enjoyable and respectful experience fo
 ## Star History
 
 [![Star History Chart](https://api.star-history.com/svg?repos=topoteretes/cognee&type=Date)](https://star-history.com/#topoteretes/cognee&Date)
+

alembic/versions/482cd6517ce4_add_default_user.py

@@ -0,0 +1,28 @@
+"""Add default user
+
+Revision ID: 482cd6517ce4
+Revises: 8057ae7329c2
+Create Date: 2024-10-16 22:17:18.634638
+
+"""
+
+from typing import Sequence, Union
+
+from sqlalchemy.util import await_only
+
+from cognee.modules.users.methods import create_default_user, delete_user
+
+
+# revision identifiers, used by Alembic.
+revision: str = "482cd6517ce4"
+down_revision: Union[str, None] = "8057ae7329c2"
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = "8057ae7329c2"
+
+
+def upgrade() -> None:
+    await_only(create_default_user())
+
+
+def downgrade() -> None:
+    await_only(delete_user("default_user@example.com"))

alembic/versions/8057ae7329c2_initial_migration.py

@@ -0,0 +1,27 @@
+"""Initial migration
+
+Revision ID: 8057ae7329c2
+Revises:
+Create Date: 2024-10-02 12:55:20.989372
+
+"""
+
+from typing import Sequence, Union
+from sqlalchemy.util import await_only
+from cognee.infrastructure.databases.relational import get_relational_engine
+
+# revision identifiers, used by Alembic.
+revision: str = "8057ae7329c2"
+down_revision: Union[str, None] = None
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    db_engine = get_relational_engine()
+    await_only(db_engine.create_database())
+
+
+def downgrade() -> None:
+    db_engine = get_relational_engine()
+    await_only(db_engine.delete_database())

cognee-mcp/Dockerfile

@@ -19,10 +19,15 @@ ENV DEBUG=${DEBUG}
 # Install system dependencies
 RUN apt-get update && apt-get install -y \
     gcc \
-    libpq-dev
+    libpq-dev \
+    git \
+    curl \
+    clang \
+    build-essential \
+    && rm -rf /var/lib/apt/lists/*
 
 # Copy pyproject.toml and lockfile first for better caching
-COPY pyproject.toml uv.lock ./
+COPY ./cognee-mcp/pyproject.toml ./cognee-mcp/uv.lock ./cognee-mcp/entrypoint.sh ./
 
 # Install the project's dependencies using the lockfile and settings
 RUN --mount=type=cache,target=/root/.cache/uv \
@@ -31,19 +36,25 @@ RUN --mount=type=cache,target=/root/.cache/uv \
 # Copy .env file first if it exists (for environment variables)
 COPY .env* /app/
 
+# Copy Alembic configuration
+COPY alembic.ini /app/alembic.ini
+COPY alembic/ /app/alembic
+
 # Then, add the rest of the project source code and install it
 # Installing separately from its dependencies allows optimal layer caching
-COPY . /app
+COPY ./cognee-mcp /app
 RUN --mount=type=cache,target=/root/.cache/uv \
     uv sync --frozen --no-dev --no-editable
 
 FROM python:3.12-slim-bookworm
 
 WORKDIR /app
- 
+
 COPY --from=uv /root/.local /root/.local
 COPY --from=uv /app /app
 
+RUN chmod +x /app/entrypoint.sh
+
 # Place executables in the environment at the front of the path
 ENV PATH="/app/.venv/bin:$PATH"
 
@@ -54,4 +65,4 @@ ENV PYTHONPATH=/app
 
 # Use the application name from pyproject.toml for normal operation
 # For testing, we'll override this with a direct command
-ENTRYPOINT ["cognee"]
+ENTRYPOINT ["/app/entrypoint.sh"]

cognee-mcp/README.md

@@ -84,3 +84,22 @@ To apply new changes while developing cognee you need to do:
 1. `poetry lock` in cognee folder
 2. `uv sync --dev --all-extras --reinstall`
 3. `mcp dev src/server.py`
+
+### Development
+In order to use local cognee build, run in root of the cognee repo:
+```bash
+poetry build -o ./cognee-mcp/sources
+```
+After the build process is done, change the cognee library dependency inside the `cognee-mcp/pyproject.toml` from
+```toml
+cognee[postgres,codegraph,gemini,huggingface]==0.1.38
+```
+to
+```toml
+cognee[postgres,codegraph,gemini,huggingface]
+```
+After that add the following snippet to the same file (`cognee-mcp/pyproject.toml`).
+```toml
+[tool.uv.sources]
+cognee = { path = "sources/cognee-0.1.38-py3-none-any.whl" }
+```

cognee-mcp/entrypoint.sh

@@ -0,0 +1,41 @@
+#!/bin/bash
+
+set -e  # Exit on error
+echo "Debug mode: $DEBUG"
+echo "Environment: $ENVIRONMENT"
+
+# Run Alembic migrations with proper error handling.
+# Note on UserAlreadyExists error handling:
+# During database migrations, we attempt to create a default user. If this user
+# already exists (e.g., from a previous deployment or migration), it's not a
+# critical error and shouldn't prevent the application from starting. This is
+# different from other migration errors which could indicate database schema
+# inconsistencies and should cause the startup to fail. This check allows for
+# smooth redeployments and container restarts while maintaining data integrity.
+echo "Running database migrations..."
+MIGRATION_OUTPUT=$(uv run alembic upgrade head 2>&1) || {
+    if [[ $MIGRATION_OUTPUT == *"UserAlreadyExists"* ]] || [[ $MIGRATION_OUTPUT == *"User default_user@example.com already exists"* ]]; then
+        echo "Warning: Default user already exists, continuing startup..."
+    else
+        echo "Migration failed with unexpected error:"
+        echo "$MIGRATION_OUTPUT"
+        exit 1
+    fi
+}
+
+echo "Starting Cognee MCP Server"
+
+# Add startup delay to ensure DB is ready
+sleep 2
+
+# Modified Gunicorn startup with error handling
+if [ "$ENVIRONMENT" = "dev" ] || [ "$ENVIRONMENT" = "local" ]; then
+    if [ "$DEBUG" = "true" ]; then
+        echo "Waiting for the debugger to attach..."
+        exec python -m debugpy --wait-for-client --listen 0.0.0.0:5678 -m cognee
+    else
+        exec cognee
+    fi
+else
+    exec cognee
+fi

cognee-mcp/pyproject.toml

@@ -1,19 +1,19 @@
 [project]
 name = "cognee-mcp"
-version = "0.2.2"
+version = "0.2.3"
 description = "A MCP server project"
 readme = "README.md"
 requires-python = ">=3.10"
 
 dependencies = [
-    "cognee[postgres,codegraph,gemini,huggingface]==0.1.37",
+    "cognee[postgres,codegraph,gemini,huggingface]==0.1.39",
     "mcp==1.5.0",
     "uv>=0.6.3",
 ]
 
 [[project.authors]]
-name = "Rita Aleksziev"
-email = "rita@topoteretes.com"
+name = "Boris Arzentar"
+email = "boris@topoteretes.com"
 
 [build-system]
 requires = [ "hatchling", ]

cognee-mcp/src/server.py

@@ -36,11 +36,11 @@ async def list_tools() -> list[types.Tool]:
                     },
                     "graph_model_file": {
                         "type": "string",
-                        "description": "The path to the graph model file",
+                        "description": "The path to the graph model file (Optional)",
                     },
                     "graph_model_name": {
                         "type": "string",
-                        "description": "The name of the graph model",
+                        "description": "The name of the graph model (Optional)",
                     },
                 },
                 "required": ["text"],
@@ -196,9 +196,13 @@ async def search(search_query: str, search_type: str) -> str:
             return json.dumps(search_results, cls=JSONEncoder)
         elif search_type.upper() == "GRAPH_COMPLETION" or search_type.upper() == "RAG_COMPLETION":
             return search_results[0]
-        else:
+        elif search_type.upper() == "CHUNKS":
+            return str(search_results)
+        elif search_type.upper() == "INSIGHTS":
             results = retrieved_edges_to_string(search_results)
             return results
+        else:
+            return str(search_results)
 
 
 async def prune():

cognee/api/client.py

@@ -13,6 +13,7 @@ from cognee.api.v1.datasets.routers import get_datasets_router
 from cognee.api.v1.cognify.routers import get_code_pipeline_router, get_cognify_router
 from cognee.api.v1.search.routers import get_search_router
 from cognee.api.v1.add.routers import get_add_router
+from cognee.api.v1.delete.routers import get_delete_router
 from fastapi import Request
 from fastapi.encoders import jsonable_encoder
 from fastapi.exceptions import RequestValidationError
@@ -164,6 +165,8 @@ app.include_router(get_settings_router(), prefix="/api/v1/settings", tags=["sett
 
 app.include_router(get_visualize_router(), prefix="/api/v1/visualize", tags=["visualize"])
 
+app.include_router(get_delete_router(), prefix="/api/v1/delete", tags=["delete"])
+
 codegraph_routes = get_code_pipeline_router()
 if codegraph_routes:
     app.include_router(codegraph_routes, prefix="/api/v1/code-pipeline", tags=["code-pipeline"])

cognee/infrastructure/databases/graph/graph_db_interface.py

@@ -21,10 +21,10 @@ Node = Tuple[str, NodeData]  # (node_id, properties)
 
 def record_graph_changes(func):
     """Decorator to record graph changes in the relationship database."""
-    db_engine = get_relational_engine()
 
     @wraps(func)
     async def wrapper(self, *args, **kwargs):
+        db_engine = get_relational_engine()
         frame = inspect.currentframe()
         while frame:
             if frame.f_back and frame.f_back.f_code.co_name != "wrapper":

cognee/modules/data/extraction/knowledge_graph/extract_content_graph.py

@@ -1,4 +1,5 @@
-from typing import Type, Optional
+import os
+from typing import Type
 from pydantic import BaseModel
 from cognee.infrastructure.llm.get_llm_client import get_llm_client
 from cognee.infrastructure.llm.prompts import render_prompt
@@ -10,7 +11,18 @@ async def extract_content_graph(content: str, response_model: Type[BaseModel]):
     llm_config = get_llm_config()
 
     prompt_path = llm_config.graph_prompt_path
-    system_prompt = render_prompt(prompt_path, {})
+
+    # Check if the prompt path is an absolute path or just a filename
+    if os.path.isabs(prompt_path):
+        # directory containing the file
+        base_directory = os.path.dirname(prompt_path)
+        # just the filename itself
+        prompt_path = os.path.basename(prompt_path)
+    else:
+        base_directory = None
+
+    system_prompt = render_prompt(prompt_path, {}, base_directory=base_directory)
+
     content_graph = await llm_client.acreate_structured_output(
         content, system_prompt, response_model
     )

cognee/modules/retrieval/code_retriever.py

@@ -19,9 +19,9 @@ class CodeRetriever(BaseRetriever):
         filenames: List[str] = []
         sourcecode: str
 
-    def __init__(self, limit: int = 3):
+    def __init__(self, top_k: int = 3):
         """Initialize retriever with search parameters."""
-        self.limit = limit
+        self.top_k = top_k
         self.file_name_collections = ["CodeFile_name"]
         self.classes_and_functions_collections = [
             "ClassDefinition_source_code",
@@ -60,7 +60,7 @@ class CodeRetriever(BaseRetriever):
         if not files_and_codeparts.filenames or not files_and_codeparts.sourcecode:
             for collection in self.file_name_collections:
                 search_results_file = await vector_engine.search(
-                    collection, query, limit=self.limit
+                    collection, query, limit=self.top_k
                 )
                 for res in search_results_file:
                     similar_filenames.append(
@@ -69,7 +69,7 @@ class CodeRetriever(BaseRetriever):
 
             for collection in self.classes_and_functions_collections:
                 search_results_code = await vector_engine.search(
-                    collection, query, limit=self.limit
+                    collection, query, limit=self.top_k
                 )
                 for res in search_results_code:
                     similar_codepieces.append(
@@ -79,7 +79,7 @@ class CodeRetriever(BaseRetriever):
             for collection in self.file_name_collections:
                 for file_from_query in files_and_codeparts.filenames:
                     search_results_file = await vector_engine.search(
-                        collection, file_from_query, limit=self.limit
+                        collection, file_from_query, limit=self.top_k
                     )
                     for res in search_results_file:
                         similar_filenames.append(
@@ -88,7 +88,7 @@ class CodeRetriever(BaseRetriever):
 
             for collection in self.classes_and_functions_collections:
                 search_results_code = await vector_engine.search(
-                    collection, files_and_codeparts.sourcecode, limit=self.limit
+                    collection, files_and_codeparts.sourcecode, limit=self.top_k
                 )
                 for res in search_results_code:
                     similar_codepieces.append(

cognee/modules/retrieval/summaries_retriever.py

@@ -9,9 +9,9 @@ from cognee.infrastructure.databases.vector.exceptions.exceptions import Collect
 class SummariesRetriever(BaseRetriever):
     """Retriever for handling summary-based searches."""
 
-    def __init__(self, limit: int = 5):
+    def __init__(self, top_k: int = 5):
         """Initialize retriever with search parameters."""
-        self.limit = limit
+        self.top_k = top_k
 
     async def get_context(self, query: str) -> Any:
         """Retrieves summary context based on the query."""
@@ -19,7 +19,7 @@ class SummariesRetriever(BaseRetriever):
 
         try:
             summaries_results = await vector_engine.search(
-                "TextSummary_text", query, limit=self.limit
+                "TextSummary_text", query, limit=self.top_k
             )
         except CollectionNotFoundError as error:
             raise NoDataError("No data found in the system, please add data first.") from error

cognee/modules/search/methods/search.py

@@ -59,9 +59,9 @@ async def specific_search(
     top_k: int = 10,
 ) -> list:
     search_tasks: dict[SearchType, Callable] = {
-        SearchType.SUMMARIES: SummariesRetriever().get_completion,
+        SearchType.SUMMARIES: SummariesRetriever(top_k=top_k).get_completion,
         SearchType.INSIGHTS: InsightsRetriever(top_k=top_k).get_completion,
-        SearchType.CHUNKS: ChunksRetriever().get_completion,
+        SearchType.CHUNKS: ChunksRetriever(top_k=top_k).get_completion,
         SearchType.RAG_COMPLETION: CompletionRetriever(
             system_prompt_path=system_prompt_path,
             top_k=top_k,
@@ -71,9 +71,9 @@ async def specific_search(
             top_k=top_k,
         ).get_completion,
         SearchType.GRAPH_SUMMARY_COMPLETION: GraphSummaryCompletionRetriever(
-            system_prompt_path=system_prompt_path,
+            system_prompt_path=system_prompt_path, top_k=top_k
         ).get_completion,
-        SearchType.CODE: CodeRetriever().get_completion,
+        SearchType.CODE: CodeRetriever(top_k=top_k).get_completion,
         SearchType.CYPHER: CypherSearchRetriever().get_completion,
         SearchType.NATURAL_LANGUAGE: NaturalLanguageRetriever().get_completion,
     }

cognee/tasks/graph/extract_graph_from_data.py

@@ -62,6 +62,16 @@ async def extract_graph_from_data(
         *[extract_content_graph(chunk.text, graph_model) for chunk in data_chunks]
     )
 
+    # Note: Filter edges with missing source or target nodes
+    if graph_model == KnowledgeGraph:
+        for graph in chunk_graphs:
+            valid_node_ids = {node.id for node in graph.nodes}
+            graph.edges = [
+                edge
+                for edge in graph.edges
+                if edge.source_node_id in valid_node_ids and edge.target_node_id in valid_node_ids
+            ]
+
     return await integrate_chunk_graphs(
         data_chunks, chunk_graphs, graph_model, ontology_adapter or OntologyResolver()
     )

cognee/tests/test_deletion.py

@@ -57,7 +57,7 @@ async def main():
 
     graph_engine = await get_graph_engine()
     nodes, edges = await graph_engine.get_graph_data()
-    assert len(nodes) > 15 and len(edges) > 15, "Graph database is not loaded."
+    assert len(nodes) > 10 and len(edges) > 10, "Graph database is not loaded."
 
     await cognee.delete([text_1, text_2], mode="hard")
     nodes, edges = await graph_engine.get_graph_data()

docker-compose.yml

@@ -7,8 +7,8 @@ services:
       context: .
       dockerfile: Dockerfile
     volumes:
-      - .:/app
-      - /app/cognee-frontend/ # Ignore frontend code
+      - ./cognee:/app/cognee
+      - .env:/app/.env
     environment:
       - DEBUG=false # Change to true if debugging
       - HOST=0.0.0.0
@@ -26,10 +26,8 @@ services:
           cpus: "4.0"
           memory: 8GB
 
-# NOTE: Frontend is a work in progress and is not intended to be used by users yet.
-#       If you want to use Cognee with a UI environment you can run the cognee-gui.py script or
-#       integrate the Cognee MCP Server to Cursor / Claude Desktop / Visual Studio Code ( through Cline/Roo )
-
+  # NOTE: Frontend is a work in progress and supports minimum amount of features required to be functional.
+  # If you want to use Cognee with a UI environment you can integrate the Cognee MCP Server into Cursor / Claude Desktop / Visual Studio Code (through Cline/Roo)
   frontend:
     container_name: frontend
     profiles:

helm/Dockerfile

@@ -3,7 +3,7 @@ FROM python:3.11-slim
 # Define Poetry extras to install
 ARG POETRY_EXTRAS="\
 # Storage & Databases \
-filesystem postgres weaviate qdrant neo4j falkordb milvus kuzu \
+postgres weaviate qdrant neo4j falkordb milvus kuzu \
 # Notebooks & Interactive Environments \
 notebook \
 # LLM & AI Frameworks \

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "cognee"
-version = "0.1.37"
+version = "0.1.39"
 description = "Cognee - is a library for enriching LLM context with a semantic layer for better understanding and reasoning."
 authors = ["Vasilije Markovic", "Boris Arzentar"]
 readme = "README.md"
@@ -92,7 +92,6 @@ s3fs = "^2025.3.2"
 
 [tool.poetry.extras]
 api = ["uvicorn", "gunicorn"]
-filesystem = ["s3fs", "botocore"]
 weaviate = ["weaviate-client"]
 qdrant = ["qdrant-client"]
 neo4j = ["neo4j"]