diff --git a/.github/workflows/cd.yaml b/.github/workflows/cd.yaml
index 19b3d6cb3..ec4b56300 100644
--- a/.github/workflows/cd.yaml
+++ b/.github/workflows/cd.yaml
@@ -17,7 +17,7 @@ jobs:
publish_docker_to_ecr:
name: Publish Cognee Docker image
- runs-on: ubuntu-latest
+ runs-on: ubuntu-22.04
permissions:
id-token: write
contents: read
diff --git a/.github/workflows/cd_prd.yaml b/.github/workflows/cd_prd.yaml
index fc40d8884..15c046215 100644
--- a/.github/workflows/cd_prd.yaml
+++ b/.github/workflows/cd_prd.yaml
@@ -17,7 +17,7 @@ jobs:
publish_docker_to_ecr:
name: Publish Docker PromethAI image
- runs-on: ubuntu-latest
+ runs-on: ubuntu-22.04
permissions:
id-token: write
contents: read
diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml
index f9acd1be0..7d63aa88d 100644
--- a/.github/workflows/ci.yaml
+++ b/.github/workflows/ci.yaml
@@ -9,7 +9,7 @@ jobs:
build_docker:
name: Build Cognee Backend Docker App Image
- runs-on: ubuntu-latest
+ runs-on: ubuntu-22.04
steps:
- name: Check out Cognee code
uses: actions/checkout@v3
diff --git a/.github/workflows/community_greetings.yml b/.github/workflows/community_greetings.yml
index b480ed111..a43f7f549 100644
--- a/.github/workflows/community_greetings.yml
+++ b/.github/workflows/community_greetings.yml
@@ -4,7 +4,7 @@ on: [pull_request, issues]
jobs:
greeting:
- runs-on: ubuntu-latest
+ runs-on: ubuntu-22.04
steps:
- uses: actions/first-interaction@v1
with:
diff --git a/.github/workflows/docker_compose.yml b/.github/workflows/docker_compose.yml
index 657f0c8c6..321cbd045 100644
--- a/.github/workflows/docker_compose.yml
+++ b/.github/workflows/docker_compose.yml
@@ -12,7 +12,7 @@ on:
jobs:
docker-compose-test:
- runs-on: ubuntu-latest
+ runs-on: ubuntu-22.04
steps:
- name: Checkout repository
diff --git a/.github/workflows/dockerhub.yml b/.github/workflows/dockerhub.yml
index a80f1f442..20f0bde96 100644
--- a/.github/workflows/dockerhub.yml
+++ b/.github/workflows/dockerhub.yml
@@ -7,7 +7,7 @@ on:
jobs:
docker-build-and-push:
- runs-on: ubuntu-latest
+ runs-on: ubuntu-22.04
steps:
- name: Checkout repository
@@ -32,7 +32,7 @@ jobs:
run: |
IMAGE_NAME=cognee/cognee
TAG_VERSION="${BRANCH_NAME}-${COMMIT_SHA}"
-
+
echo "Building image: ${IMAGE_NAME}:${TAG_VERSION}"
docker buildx build \
--platform linux/amd64,linux/arm64 \
diff --git a/.github/workflows/profiling.yaml b/.github/workflows/profiling.yaml
index 93ed82f82..0ecbc960c 100644
--- a/.github/workflows/profiling.yaml
+++ b/.github/workflows/profiling.yaml
@@ -7,7 +7,7 @@ on:
jobs:
profiler:
- runs-on: ubuntu-latest
+ runs-on: ubuntu-22.04
steps:
# Checkout the code from the repository with full history
@@ -47,7 +47,7 @@ jobs:
python-version: '3.10'
- name: Install Poetry
- uses: snok/install-poetry@v1.3.2
+ uses: snok/install-poetry@v1.4.1
with:
virtualenvs-create: true
virtualenvs-in-project: true
@@ -57,6 +57,8 @@ jobs:
run: |
poetry install --no-interaction --all-extras
poetry run pip install pyinstrument
+ poetry run pip install parso
+ poetry run pip install jedi
# Set environment variables for SHAs
@@ -80,50 +82,50 @@ jobs:
poetry run pyinstrument --renderer json -o base_results.json cognee/api/v1/cognify/code_graph_pipeline.py
# Run profiler on head branch
- - name: Run profiler on head branch
- env:
- HEAD_SHA: ${{ env.HEAD_SHA }}
- run: |
- echo "Profiling the head branch for code_graph_pipeline.py"
- echo "Checking out head SHA: $HEAD_SHA"
- git checkout $HEAD_SHA
- echo "This is the working directory: $PWD"
- # Ensure the script is executable
- chmod +x cognee/api/v1/cognify/code_graph_pipeline.py
- # Run Scalene
- poetry run pyinstrument --renderer json -o head_results.json cognee/api/v1/cognify/code_graph_pipeline.py
-
- # Compare profiling results
- - name: Compare profiling results
- run: |
- python -c '
- import json
- try:
- with open("base_results.json") as f:
- base = json.load(f)
- with open("head_results.json") as f:
- head = json.load(f)
- cpu_diff = head.get("total_cpu_samples_python", 0) - base.get("total_cpu_samples_python", 0)
- memory_diff = head.get("malloc_samples", 0) - base.get("malloc_samples", 0)
- results = [
- f"CPU Usage Difference: {cpu_diff}",
- f"Memory Usage Difference: {memory_diff} bytes"
- ]
- with open("profiling_diff.txt", "w") as f:
- f.write("\\n".join(results) + "\\n")
- print("\\n".join(results)) # Print results to terminal
- except Exception as e:
- error_message = f"Error comparing profiling results: {e}"
- with open("profiling_diff.txt", "w") as f:
- f.write(error_message + "\\n")
- print(error_message) # Print error to terminal
- '
-
- - name: Upload profiling diff artifact
- uses: actions/upload-artifact@v3
- with:
- name: profiling-diff
- path: profiling_diff.txt
+# - name: Run profiler on head branch
+# env:
+# HEAD_SHA: ${{ env.HEAD_SHA }}
+# run: |
+# echo "Profiling the head branch for code_graph_pipeline.py"
+# echo "Checking out head SHA: $HEAD_SHA"
+# git checkout $HEAD_SHA
+# echo "This is the working directory: $PWD"
+# # Ensure the script is executable
+# chmod +x cognee/api/v1/cognify/code_graph_pipeline.py
+# # Run Scalene
+# poetry run pyinstrument --renderer json -o head_results.json cognee/api/v1/cognify/code_graph_pipeline.py
+#
+# # Compare profiling results
+# - name: Compare profiling results
+# run: |
+# python -c '
+# import json
+# try:
+# with open("base_results.json") as f:
+# base = json.load(f)
+# with open("head_results.json") as f:
+# head = json.load(f)
+# cpu_diff = head.get("total_cpu_samples_python", 0) - base.get("total_cpu_samples_python", 0)
+# memory_diff = head.get("malloc_samples", 0) - base.get("malloc_samples", 0)
+# results = [
+# f"CPU Usage Difference: {cpu_diff}",
+# f"Memory Usage Difference: {memory_diff} bytes"
+# ]
+# with open("profiling_diff.txt", "w") as f:
+# f.write("\\n".join(results) + "\\n")
+# print("\\n".join(results)) # Print results to terminal
+# except Exception as e:
+# error_message = f"Error comparing profiling results: {e}"
+# with open("profiling_diff.txt", "w") as f:
+# f.write(error_message + "\\n")
+# print(error_message) # Print error to terminal
+# '
+#
+# - name: Upload profiling diff artifact
+# uses: actions/upload-artifact@v3
+# with:
+# name: profiling-diff
+# path: profiling_diff.txt
# Post results to the pull request
# - name: Post profiling results to PR
diff --git a/.github/workflows/py_lint.yml b/.github/workflows/py_lint.yml
index dfdcdfd8f..543a0d221 100644
--- a/.github/workflows/py_lint.yml
+++ b/.github/workflows/py_lint.yml
@@ -16,8 +16,8 @@ jobs:
fail-fast: true
matrix:
os:
- - ubuntu-latest
- python-version: ["3.8.x", "3.9.x", "3.10.x", "3.11.x"]
+ - ubuntu-22.04
+ python-version: ["3.10.x", "3.11.x"]
defaults:
run:
diff --git a/.github/workflows/release_discord_action.yml b/.github/workflows/release_discord_action.yml
index 7b3ae845d..f3113ccb7 100644
--- a/.github/workflows/release_discord_action.yml
+++ b/.github/workflows/release_discord_action.yml
@@ -6,7 +6,7 @@ on:
jobs:
github-releases-to-discord:
- runs-on: ubuntu-latest
+ runs-on: ubuntu-22.04
steps:
- name: Checkout
uses: actions/checkout@v3
diff --git a/.github/workflows/reusable_notebook.yml b/.github/workflows/reusable_notebook.yml
index 0c0b63ec4..8034aca97 100644
--- a/.github/workflows/reusable_notebook.yml
+++ b/.github/workflows/reusable_notebook.yml
@@ -22,7 +22,7 @@ jobs:
run_notebook_test:
name: test
- runs-on: ubuntu-latest
+ runs-on: ubuntu-22.04
defaults:
run:
shell: bash
@@ -36,7 +36,7 @@ jobs:
python-version: '3.11.x'
- name: Install Poetry
- uses: snok/install-poetry@v1.3.2
+ uses: snok/install-poetry@v1.4.1
with:
virtualenvs-create: true
virtualenvs-in-project: true
@@ -58,4 +58,4 @@ jobs:
--to notebook \
--execute ${{ inputs.notebook-location }} \
--output executed_notebook.ipynb \
- --ExecutePreprocessor.timeout=1200
\ No newline at end of file
+ --ExecutePreprocessor.timeout=1200
diff --git a/.github/workflows/reusable_python_example.yml b/.github/workflows/reusable_python_example.yml
index 03f928656..4aa4aaba6 100644
--- a/.github/workflows/reusable_python_example.yml
+++ b/.github/workflows/reusable_python_example.yml
@@ -22,7 +22,7 @@ jobs:
run_notebook_test:
name: test
- runs-on: ubuntu-latest
+ runs-on: ubuntu-22.04
defaults:
run:
shell: bash
@@ -33,10 +33,10 @@ jobs:
- name: Setup Python
uses: actions/setup-python@v5
with:
- python-version: '3.11.x'
+ python-version: '3.12.x'
- name: Install Poetry
- uses: snok/install-poetry@v1.3.2
+ uses: snok/install-poetry@v1.4.1
with:
virtualenvs-create: true
virtualenvs-in-project: true
@@ -49,7 +49,8 @@ jobs:
- name: Execute Python Example
env:
ENV: 'dev'
+ PYTHONFAULTHANDLER: 1
LLM_API_KEY: ${{ secrets.OPENAI_API_KEY }}
GRAPHISTRY_USERNAME: ${{ secrets.GRAPHISTRY_USERNAME }}
GRAPHISTRY_PASSWORD: ${{ secrets.GRAPHISTRY_PASSWORD }}
- run: poetry run python ${{ inputs.example-location }}
\ No newline at end of file
+ run: poetry run python ${{ inputs.example-location }}
diff --git a/.github/workflows/ruff_format.yaml b/.github/workflows/ruff_format.yaml
index 959b7fc4b..a75a795e7 100644
--- a/.github/workflows/ruff_format.yaml
+++ b/.github/workflows/ruff_format.yaml
@@ -3,7 +3,7 @@ on: [ pull_request ]
jobs:
ruff:
- runs-on: ubuntu-latest
+ runs-on: ubuntu-22.04
steps:
- uses: actions/checkout@v4
- uses: astral-sh/ruff-action@v2
diff --git a/.github/workflows/ruff_lint.yaml b/.github/workflows/ruff_lint.yaml
index 214e8ec6d..4c4fb81e3 100644
--- a/.github/workflows/ruff_lint.yaml
+++ b/.github/workflows/ruff_lint.yaml
@@ -3,7 +3,7 @@ on: [ pull_request ]
jobs:
ruff:
- runs-on: ubuntu-latest
+ runs-on: ubuntu-22.04
steps:
- uses: actions/checkout@v4
- uses: astral-sh/ruff-action@v2
diff --git a/.github/workflows/test_deduplication.yml b/.github/workflows/test_deduplication.yml
index a1aab3252..923bbb68c 100644
--- a/.github/workflows/test_deduplication.yml
+++ b/.github/workflows/test_deduplication.yml
@@ -16,8 +16,7 @@ env:
jobs:
run_deduplication_test:
name: test
- runs-on: ubuntu-latest
- if: ${{ github.event.label.name == 'run-checks' }}
+ runs-on: ubuntu-22.04
defaults:
run:
shell: bash
@@ -46,7 +45,7 @@ jobs:
python-version: '3.11.x'
- name: Install Poetry
- uses: snok/install-poetry@v1.3.2
+ uses: snok/install-poetry@v1.4.1
with:
virtualenvs-create: true
virtualenvs-in-project: true
diff --git a/.github/workflows/test_milvus.yml b/.github/workflows/test_milvus.yml
index c4214fddf..51e5f0982 100644
--- a/.github/workflows/test_milvus.yml
+++ b/.github/workflows/test_milvus.yml
@@ -17,8 +17,7 @@ jobs:
run_milvus:
name: test
- runs-on: ubuntu-latest
- if: ${{ github.event.label.name == 'run-checks' }}
+ runs-on: ubuntu-22.04
strategy:
fail-fast: false
defaults:
@@ -36,7 +35,7 @@ jobs:
- name: Install Poetry
# https://github.com/snok/install-poetry#running-on-windows
- uses: snok/install-poetry@v1.3.2
+ uses: snok/install-poetry@v1.4.1
with:
virtualenvs-create: true
virtualenvs-in-project: true
diff --git a/.github/workflows/test_neo4j.yml b/.github/workflows/test_neo4j.yml
index 5f955eefc..e1d71dcfd 100644
--- a/.github/workflows/test_neo4j.yml
+++ b/.github/workflows/test_neo4j.yml
@@ -15,8 +15,7 @@ env:
jobs:
run_neo4j_integration_test:
name: test
- if: ${{ github.event.label.name == 'run-checks' }}
- runs-on: ubuntu-latest
+ runs-on: ubuntu-22.04
defaults:
run:
@@ -32,7 +31,7 @@ jobs:
python-version: '3.11.x'
- name: Install Poetry
- uses: snok/install-poetry@v1.3.2
+ uses: snok/install-poetry@v1.4.1
with:
virtualenvs-create: true
virtualenvs-in-project: true
diff --git a/.github/workflows/test_pgvector.yml b/.github/workflows/test_pgvector.yml
index 9f5ffd1da..d5356d603 100644
--- a/.github/workflows/test_pgvector.yml
+++ b/.github/workflows/test_pgvector.yml
@@ -17,8 +17,7 @@ jobs:
run_pgvector_integration_test:
name: test
- runs-on: ubuntu-latest
- if: ${{ github.event.label.name == 'run-checks' }}
+ runs-on: ubuntu-22.04
defaults:
run:
shell: bash
@@ -47,7 +46,7 @@ jobs:
python-version: '3.11.x'
- name: Install Poetry
- uses: snok/install-poetry@v1.3.2
+ uses: snok/install-poetry@v1.4.1
with:
virtualenvs-create: true
virtualenvs-in-project: true
diff --git a/.github/workflows/test_python_3_10.yml b/.github/workflows/test_python_3_10.yml
index 770d2fd63..2cf620a17 100644
--- a/.github/workflows/test_python_3_10.yml
+++ b/.github/workflows/test_python_3_10.yml
@@ -14,11 +14,9 @@ env:
ENV: 'dev'
jobs:
-
run_common:
name: test
- runs-on: ubuntu-latest
- if: ${{ github.event.label.name == 'run-checks' }}
+ runs-on: ubuntu-22.04
strategy:
fail-fast: false
defaults:
@@ -36,7 +34,7 @@ jobs:
- name: Install Poetry
# https://github.com/snok/install-poetry#running-on-windows
- uses: snok/install-poetry@v1.3.2
+ uses: snok/install-poetry@v1.4.1
with:
virtualenvs-create: true
virtualenvs-in-project: true
diff --git a/.github/workflows/test_python_3_11.yml b/.github/workflows/test_python_3_11.yml
index 69eb875bd..b119dbcb2 100644
--- a/.github/workflows/test_python_3_11.yml
+++ b/.github/workflows/test_python_3_11.yml
@@ -17,8 +17,7 @@ jobs:
run_common:
name: test
- runs-on: ubuntu-latest
- if: ${{ github.event.label.name == 'run-checks' }}
+ runs-on: ubuntu-22.04
strategy:
fail-fast: false
defaults:
@@ -36,7 +35,7 @@ jobs:
- name: Install Poetry
# https://github.com/snok/install-poetry#running-on-windows
- uses: snok/install-poetry@v1.3.2
+ uses: snok/install-poetry@v1.4.1
with:
virtualenvs-create: true
virtualenvs-in-project: true
diff --git a/.github/workflows/test_python_3_9.yml b/.github/workflows/test_python_3_12.yml
similarity index 89%
rename from .github/workflows/test_python_3_9.yml
rename to .github/workflows/test_python_3_12.yml
index 380c894ca..5a032144a 100644
--- a/.github/workflows/test_python_3_9.yml
+++ b/.github/workflows/test_python_3_12.yml
@@ -1,4 +1,4 @@
-name: test | python 3.9
+name: test | python 3.12
on:
workflow_dispatch:
@@ -17,8 +17,7 @@ jobs:
run_common:
name: test
- runs-on: ubuntu-latest
- if: ${{ github.event.label.name == 'run-checks' }}
+ runs-on: ubuntu-22.04
strategy:
fail-fast: false
defaults:
@@ -32,11 +31,11 @@ jobs:
- name: Setup Python
uses: actions/setup-python@v5
with:
- python-version: '3.9.x'
+ python-version: '3.12.x'
- name: Install Poetry
# https://github.com/snok/install-poetry#running-on-windows
- uses: snok/install-poetry@v1.3.2
+ uses: snok/install-poetry@v1.4.1
with:
virtualenvs-create: true
virtualenvs-in-project: true
diff --git a/.github/workflows/test_qdrant.yml b/.github/workflows/test_qdrant.yml
index 17d9ac628..d1447c65c 100644
--- a/.github/workflows/test_qdrant.yml
+++ b/.github/workflows/test_qdrant.yml
@@ -17,9 +17,7 @@ jobs:
run_qdrant_integration_test:
name: test
- runs-on: ubuntu-latest
- if: ${{ github.event.label.name == 'run-checks' }}
-
+ runs-on: ubuntu-22.04
defaults:
run:
shell: bash
@@ -34,7 +32,7 @@ jobs:
python-version: '3.11.x'
- name: Install Poetry
- uses: snok/install-poetry@v1.3.2
+ uses: snok/install-poetry@v1.4.1
with:
virtualenvs-create: true
virtualenvs-in-project: true
diff --git a/.github/workflows/test_weaviate.yml b/.github/workflows/test_weaviate.yml
index 9a3651dda..159fce194 100644
--- a/.github/workflows/test_weaviate.yml
+++ b/.github/workflows/test_weaviate.yml
@@ -17,9 +17,7 @@ jobs:
run_weaviate_integration_test:
name: test
- runs-on: ubuntu-latest
- if: ${{ github.event.label.name == 'run-checks' }}
-
+ runs-on: ubuntu-22.04
defaults:
run:
shell: bash
@@ -34,7 +32,7 @@ jobs:
python-version: '3.11.x'
- name: Install Poetry
- uses: snok/install-poetry@v1.3.2
+ uses: snok/install-poetry@v1.4.1
with:
virtualenvs-create: true
virtualenvs-in-project: true
diff --git a/Dockerfile b/Dockerfile
index cec626c86..50f2d0b08 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -25,7 +25,7 @@ RUN pip install poetry
RUN poetry config virtualenvs.create false
# Install the dependencies
-RUN poetry install --all-extras --no-root --no-dev
+RUN poetry install --all-extras --no-root --without dev
# Set the PYTHONPATH environment variable to include the /app directory
diff --git a/alembic/env.py b/alembic/env.py
index 1675e4168..60af9c0dd 100644
--- a/alembic/env.py
+++ b/alembic/env.py
@@ -4,8 +4,9 @@
from sqlalchemy import pool
from sqlalchemy.engine import Connection
from sqlalchemy.ext.asyncio import async_engine_from_config
-
+from cognee.infrastructure.databases.relational import Base
from alembic import context
+from cognee.infrastructure.databases.relational import get_relational_engine, get_relational_config
# this is the Alembic Config object, which provides
# access to the values within the .ini file in use.
@@ -20,7 +21,7 @@
# for 'autogenerate' support
# from myapp import mymodel
# target_metadata = mymodel.Base.metadata
-from cognee.infrastructure.databases.relational import Base
+
target_metadata = Base.metadata
# other values from the config, defined by the needs of env.py,
@@ -83,12 +84,11 @@ def run_migrations_online() -> None:
asyncio.run(run_async_migrations())
-from cognee.infrastructure.databases.relational import get_relational_engine, get_relational_config
-
db_engine = get_relational_engine()
if db_engine.engine.dialect.name == "sqlite":
from cognee.infrastructure.files.storage import LocalStorage
+
db_config = get_relational_config()
LocalStorage.ensure_directory_exists(db_config.db_path)
diff --git a/alembic/versions/482cd6517ce4_add_default_user.py b/alembic/versions/482cd6517ce4_add_default_user.py
index b744c8821..92429e1e4 100644
--- a/alembic/versions/482cd6517ce4_add_default_user.py
+++ b/alembic/versions/482cd6517ce4_add_default_user.py
@@ -5,6 +5,7 @@
Create Date: 2024-10-16 22:17:18.634638
"""
+
from typing import Sequence, Union
from sqlalchemy.util import await_only
@@ -13,8 +14,8 @@
# revision identifiers, used by Alembic.
-revision: str = '482cd6517ce4'
-down_revision: Union[str, None] = '8057ae7329c2'
+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"
diff --git a/alembic/versions/8057ae7329c2_initial_migration.py b/alembic/versions/8057ae7329c2_initial_migration.py
index d7d45d0ac..48e795327 100644
--- a/alembic/versions/8057ae7329c2_initial_migration.py
+++ b/alembic/versions/8057ae7329c2_initial_migration.py
@@ -1,10 +1,11 @@
"""Initial migration
Revision ID: 8057ae7329c2
-Revises:
+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
diff --git a/assets/Dashboard_example.png b/assets/Dashboard_example.png
deleted file mode 100644
index 594186c58..000000000
Binary files a/assets/Dashboard_example.png and /dev/null differ
diff --git a/assets/cognee-logo.png b/assets/cognee-logo.png
deleted file mode 100644
index d02a831fb..000000000
Binary files a/assets/cognee-logo.png and /dev/null differ
diff --git a/assets/cognee_logo.png b/assets/cognee_logo.png
new file mode 100644
index 000000000..744ba2028
Binary files /dev/null and b/assets/cognee_logo.png differ
diff --git a/assets/topoteretes_logo.png b/assets/topoteretes_logo.png
deleted file mode 100644
index ebdbc9a20..000000000
Binary files a/assets/topoteretes_logo.png and /dev/null differ
diff --git a/assets/vscode-debug-config.png b/assets/vscode-debug-config.png
deleted file mode 100644
index 21f54400a..000000000
Binary files a/assets/vscode-debug-config.png and /dev/null differ
diff --git a/cognee-mcp/.python-version b/cognee-mcp/.python-version
new file mode 100644
index 000000000..9ac380418
--- /dev/null
+++ b/cognee-mcp/.python-version
@@ -0,0 +1 @@
+3.11.5
diff --git a/cognee-mcp/cognee_mcp/__init__.py b/cognee-mcp/cognee_mcp/__init__.py
index e704e08dc..8c1ca4328 100644
--- a/cognee-mcp/cognee_mcp/__init__.py
+++ b/cognee-mcp/cognee_mcp/__init__.py
@@ -7,6 +7,7 @@ def main():
"""Main entry point for the package."""
asyncio.run(server.main())
+
# Optionally expose other important items at package level
__all__ = ["main", "server"]
diff --git a/cognee-mcp/cognee_mcp/server.py b/cognee-mcp/cognee_mcp/server.py
index 739a86391..19dc999d7 100644
--- a/cognee-mcp/cognee_mcp/server.py
+++ b/cognee-mcp/cognee_mcp/server.py
@@ -18,7 +18,9 @@ def node_to_string(node):
# keys_to_keep = ["chunk_index", "topological_rank", "cut_type", "id", "text"]
# keyset = set(keys_to_keep) & node.keys()
# return "Node(" + " ".join([key + ": " + str(node[key]) + "," for key in keyset]) + ")"
- node_data = ", ".join([f"{key}: \"{value}\"" for key, value in node.items() if key in ["id", "name"]])
+ node_data = ", ".join(
+ [f'{key}: "{value}"' for key, value in node.items() if key in ["id", "name"]]
+ )
return f"Node({node_data})"
@@ -52,9 +54,9 @@ async def handle_list_tools() -> list[types.Tool]:
"""
return [
types.Tool(
- name = "cognify",
- description = "Build knowledge graph from the input text.",
- inputSchema = {
+ name="cognify",
+ description="Build knowledge graph from the input text.",
+ inputSchema={
"type": "object",
"properties": {
"text": {"type": "string"},
@@ -65,9 +67,9 @@ async def handle_list_tools() -> list[types.Tool]:
},
),
types.Tool(
- name = "search",
- description = "Search the knowledge graph.",
- inputSchema = {
+ name="search",
+ description="Search the knowledge graph.",
+ inputSchema={
"type": "object",
"properties": {
"query": {"type": "string"},
@@ -76,9 +78,9 @@ async def handle_list_tools() -> list[types.Tool]:
},
),
types.Tool(
- name = "prune",
- description = "Reset the knowledge graph.",
- inputSchema = {
+ name="prune",
+ description="Reset the knowledge graph.",
+ inputSchema={
"type": "object",
"properties": {
"query": {"type": "string"},
@@ -90,8 +92,7 @@ async def handle_list_tools() -> list[types.Tool]:
@server.call_tool()
async def handle_call_tool(
- name: str,
- arguments: dict | None
+ name: str, arguments: dict | None
) -> list[types.TextContent | types.ImageContent | types.EmbeddedResource]:
"""
Handle tool execution requests.
@@ -115,12 +116,12 @@ async def handle_call_tool(
await cognee.add(text)
- await cognee.cognify(graph_model = graph_model)
+ await cognee.cognify(graph_model=graph_model)
return [
types.TextContent(
- type = "text",
- text = "Ingested",
+ type="text",
+ text="Ingested",
)
]
elif name == "search":
@@ -131,16 +132,14 @@ async def handle_call_tool(
search_query = arguments.get("query")
- search_results = await cognee.search(
- SearchType.INSIGHTS, query_text = search_query
- )
+ search_results = await cognee.search(SearchType.INSIGHTS, query_text=search_query)
results = retrieved_edges_to_string(search_results)
return [
types.TextContent(
- type = "text",
- text = results,
+ type="text",
+ text=results,
)
]
elif name == "prune":
@@ -151,8 +150,8 @@ async def handle_call_tool(
return [
types.TextContent(
- type = "text",
- text = "Pruned",
+ type="text",
+ text="Pruned",
)
]
else:
@@ -166,15 +165,16 @@ async def main():
read_stream,
write_stream,
InitializationOptions(
- server_name = "cognee-mcp",
- server_version = "0.1.0",
- capabilities = server.get_capabilities(
- notification_options = NotificationOptions(),
- experimental_capabilities = {},
+ server_name="cognee-mcp",
+ server_version="0.1.0",
+ capabilities=server.get_capabilities(
+ notification_options=NotificationOptions(),
+ experimental_capabilities={},
),
),
)
+
# This is needed if you'd like to connect to a custom client
if __name__ == "__main__":
asyncio.run(main())
diff --git a/cognee-mcp/uv.lock b/cognee-mcp/uv.lock
index 63dfae03b..b748667e1 100644
--- a/cognee-mcp/uv.lock
+++ b/cognee-mcp/uv.lock
@@ -321,6 +321,26 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/b1/fe/e8c672695b37eecc5cbf43e1d0638d88d66ba3a44c4d321c796f4e59167f/beautifulsoup4-4.12.3-py3-none-any.whl", hash = "sha256:b80878c9f40111313e55da8ba20bdba06d8fa3969fc68304167741bbf9e082ed", size = 147925 },
]
+[[package]]
+name = "bokeh"
+version = "3.6.2"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+ { name = "contourpy" },
+ { name = "jinja2" },
+ { name = "numpy" },
+ { name = "packaging" },
+ { name = "pandas" },
+ { name = "pillow" },
+ { name = "pyyaml" },
+ { name = "tornado" },
+ { name = "xyzservices" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/da/9d/cc9c561e1db8cbecc5cfad972159020700fff2339bdaa316498ace1cb04c/bokeh-3.6.2.tar.gz", hash = "sha256:2f3043d9ecb3d5dc2e8c0ebf8ad55727617188d4e534f3e7208b36357e352396", size = 6247610 }
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/56/12/2c266a0dc57379c60b4e73a2f93e71343db4170bf26c5a76a74e7d8bce2a/bokeh-3.6.2-py3-none-any.whl", hash = "sha256:fddc4b91f8b40178c0e3e83dfcc33886d7803a3a1f041a840834255e435a18c2", size = 6866799 },
+]
+
[[package]]
name = "boto3"
version = "1.35.84"
@@ -358,6 +378,34 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/a4/07/14f8ad37f2d12a5ce41206c21820d8cb6561b728e51fad4530dff0552a67/cachetools-5.5.0-py3-none-any.whl", hash = "sha256:02134e8439cdc2ffb62023ce1debca2944c3f289d66bb17ead3ab3dede74b292", size = 9524 },
]
+[[package]]
+name = "cairocffi"
+version = "1.7.1"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+ { name = "cffi" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/70/c5/1a4dc131459e68a173cbdab5fad6b524f53f9c1ef7861b7698e998b837cc/cairocffi-1.7.1.tar.gz", hash = "sha256:2e48ee864884ec4a3a34bfa8c9ab9999f688286eb714a15a43ec9d068c36557b", size = 88096 }
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/93/d8/ba13451aa6b745c49536e87b6bf8f629b950e84bd0e8308f7dc6883b67e2/cairocffi-1.7.1-py3-none-any.whl", hash = "sha256:9803a0e11f6c962f3b0ae2ec8ba6ae45e957a146a004697a1ac1bbf16b073b3f", size = 75611 },
+]
+
+[[package]]
+name = "cairosvg"
+version = "2.7.1"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+ { name = "cairocffi" },
+ { name = "cssselect2" },
+ { name = "defusedxml" },
+ { name = "pillow" },
+ { name = "tinycss2" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/d5/e6/ec5900b724e3c44af7f6f51f719919137284e5da4aabe96508baec8a1b40/CairoSVG-2.7.1.tar.gz", hash = "sha256:432531d72347291b9a9ebfb6777026b607563fd8719c46ee742db0aef7271ba0", size = 8399085 }
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/01/a5/1866b42151f50453f1a0d28fc4c39f5be5f412a2e914f33449c42daafdf1/CairoSVG-2.7.1-py3-none-any.whl", hash = "sha256:8a5222d4e6c3f86f1f7046b63246877a63b49923a1cd202184c3a634ef546b3b", size = 43235 },
+]
+
[[package]]
name = "certifi"
version = "2024.12.14"
@@ -412,6 +460,15 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/7c/fc/6a8cb64e5f0324877d503c854da15d76c1e50eb722e320b15345c4d0c6de/cffi-1.17.1-cp313-cp313-win_amd64.whl", hash = "sha256:f6a16c31041f09ead72d69f583767292f750d24913dadacf5756b966aacb3f1a", size = 182009 },
]
+[[package]]
+name = "cfgv"
+version = "3.4.0"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/11/74/539e56497d9bd1d484fd863dd69cbbfa653cd2aa27abfe35653494d85e94/cfgv-3.4.0.tar.gz", hash = "sha256:e52591d4c5f5dead8e0f673fb16db7949d2cfb3f7da4582893288f0ded8fe560", size = 7114 }
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/c5/55/51844dd50c4fc7a33b653bfaba4c2456f06955289ca770a5dbd5fd267374/cfgv-3.4.0-py2.py3-none-any.whl", hash = "sha256:b7265b1f29fd3316bfcd2b330d63d024f2bfd8bcb8b0272f8e19a504856c48f9", size = 7249 },
+]
+
[[package]]
name = "chardet"
version = "5.2.0"
@@ -480,7 +537,7 @@ name = "click"
version = "8.1.7"
source = { registry = "https://pypi.org/simple" }
dependencies = [
- { name = "colorama", marker = "platform_system == 'Windows'" },
+ { name = "colorama", marker = "sys_platform == 'win32'" },
]
sdist = { url = "https://files.pythonhosted.org/packages/96/d3/f04c7bfcf5c1862a2a5b845c6b2b360488cf47af55dfa79c98f6a6bf98b5/click-8.1.7.tar.gz", hash = "sha256:ca9853ad459e787e2192211578cc907e7594e294c7ccc834310722b41b9ca6de", size = 336121 }
wheels = [
@@ -497,7 +554,9 @@ dependencies = [
{ name = "aiosqlite" },
{ name = "alembic" },
{ name = "anthropic" },
+ { name = "bokeh" },
{ name = "boto3" },
+ { name = "cairosvg" },
{ name = "datasets" },
{ name = "dlt", extra = ["sqlalchemy"] },
{ name = "fastapi" },
@@ -518,6 +577,7 @@ dependencies = [
{ name = "numpy" },
{ name = "openai" },
{ name = "pandas" },
+ { name = "pre-commit" },
{ name = "pydantic" },
{ name = "pydantic-settings" },
{ name = "pypdf" },
@@ -541,8 +601,10 @@ requires-dist = [
{ name = "alembic", specifier = ">=1.13.3,<2.0.0" },
{ name = "anthropic", specifier = ">=0.26.1,<0.27.0" },
{ name = "asyncpg", marker = "extra == 'postgres'", specifier = "==0.30.0" },
+ { name = "bokeh", specifier = ">=3.6.2,<4.0.0" },
{ name = "boto3", specifier = ">=1.26.125,<2.0.0" },
{ name = "botocore", marker = "extra == 'filesystem'", specifier = ">=1.35.54,<2.0.0" },
+ { name = "cairosvg", specifier = ">=2.7.1,<3.0.0" },
{ name = "datasets", specifier = "==3.1.0" },
{ name = "deepeval", marker = "extra == 'deepeval'", specifier = ">=2.0.1,<3.0.0" },
{ name = "dlt", extras = ["sqlalchemy"], specifier = ">=1.4.1,<2.0.0" },
@@ -573,6 +635,7 @@ requires-dist = [
{ name = "pandas", specifier = "==2.0.3" },
{ name = "pgvector", marker = "extra == 'postgres'", specifier = ">=0.3.5,<0.4.0" },
{ name = "posthog", marker = "extra == 'posthog'", specifier = ">=3.5.0,<4.0.0" },
+ { name = "pre-commit", specifier = ">=4.0.1,<5.0.0" },
{ name = "psycopg2", marker = "extra == 'postgres'", specifier = ">=2.9.10,<3.0.0" },
{ name = "pydantic", specifier = "==2.8.2" },
{ name = "pydantic-settings", specifier = ">=2.2.1,<3.0.0" },
@@ -885,6 +948,19 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/97/9b/443270b9210f13f6ef240eff73fd32e02d381e7103969dc66ce8e89ee901/cryptography-44.0.0-cp39-abi3-win_amd64.whl", hash = "sha256:708ee5f1bafe76d041b53a4f95eb28cdeb8d18da17e597d46d7833ee59b97ede", size = 3202071 },
]
+[[package]]
+name = "cssselect2"
+version = "0.7.0"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+ { name = "tinycss2" },
+ { name = "webencodings" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/e7/fc/326cb6f988905998f09bb54a3f5d98d4462ba119363c0dfad29750d48c09/cssselect2-0.7.0.tar.gz", hash = "sha256:1ccd984dab89fc68955043aca4e1b03e0cf29cad9880f6e28e3ba7a74b14aa5a", size = 35888 }
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/9d/3a/e39436efe51894243ff145a37c4f9a030839b97779ebcc4f13b3ba21c54e/cssselect2-0.7.0-py3-none-any.whl", hash = "sha256:fd23a65bfd444595913f02fc71f6b286c29261e354c41d722ca7a261a49b5969", size = 15586 },
+]
+
[[package]]
name = "cycler"
version = "0.12.1"
@@ -996,6 +1072,15 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/39/60/533ce66e28295e2b94267126a851ac091ad29a835a9827d1f9c30574fce4/deepeval-2.0.6-py3-none-any.whl", hash = "sha256:57302830ff9d3d16ad4f1961338c7b4453e48039ff131990f258880728f33b6b", size = 504101 },
]
+[[package]]
+name = "defusedxml"
+version = "0.7.1"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/0f/d5/c66da9b79e5bdb124974bfe172b4daf3c984ebd9c2a06e2b8a4dc7331c72/defusedxml-0.7.1.tar.gz", hash = "sha256:1bb3032db185915b62d7c6209c5a8792be6a32ab2fedacc84e01b52c51aa3e69", size = 75520 }
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/07/6c/aa3f2f849e01cb6a001cd8554a88d4c77c5c1a31c95bdf1cf9301e6d9ef4/defusedxml-0.7.1-py2.py3-none-any.whl", hash = "sha256:a352e7e428770286cc899e2542b6cdaedb2b4953ff269a210103ec58f6198a61", size = 25604 },
+]
+
[[package]]
name = "deprecated"
version = "1.2.15"
@@ -1056,6 +1141,15 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/68/69/1bcf70f81de1b4a9f21b3a62ec0c83bdff991c88d6cc2267d02408457e88/dirtyjson-1.0.8-py3-none-any.whl", hash = "sha256:125e27248435a58acace26d5c2c4c11a1c0de0a9c5124c5a94ba78e517d74f53", size = 25197 },
]
+[[package]]
+name = "distlib"
+version = "0.3.9"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/0d/dd/1bec4c5ddb504ca60fc29472f3d27e8d4da1257a854e1d96742f15c1d02d/distlib-0.3.9.tar.gz", hash = "sha256:a60f20dea646b8a33f3e7772f74dc0b2d0772d2837ee1342a00645c81edf9403", size = 613923 }
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/91/a1/cf2472db20f7ce4a6be1253a81cfdf85ad9c7885ffbed7047fb72c24cf87/distlib-0.3.9-py2.py3-none-any.whl", hash = "sha256:47f8c22fd27c27e25a65601af709b38e4f0a45ea4fc2e710f65755fa8caaaf87", size = 468973 },
+]
+
[[package]]
name = "distro"
version = "1.9.0"
@@ -1686,6 +1780,15 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/d7/de/85a784bcc4a3779d1753a7ec2dee5de90e18c7bcf402e71b51fcf150b129/hyperframe-6.0.1-py3-none-any.whl", hash = "sha256:0ec6bafd80d8ad2195c4f03aacba3a8265e57bc4cff261e802bf39970ed02a15", size = 12389 },
]
+[[package]]
+name = "identify"
+version = "2.6.3"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/1a/5f/05f0d167be94585d502b4adf8c7af31f1dc0b1c7e14f9938a88fdbbcf4a7/identify-2.6.3.tar.gz", hash = "sha256:62f5dae9b5fef52c84cc188514e9ea4f3f636b1d8799ab5ebc475471f9e47a02", size = 99179 }
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/c9/f5/09644a3ad803fae9eca8efa17e1f2aef380c7f0b02f7ec4e8d446e51d64a/identify-2.6.3-py2.py3-none-any.whl", hash = "sha256:9edba65473324c2ea9684b1f944fe3191db3345e50b6d04571d10ed164f8d7bd", size = 99049 },
+]
+
[[package]]
name = "idna"
version = "3.10"
@@ -2557,6 +2660,15 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/4d/66/7d9e26593edda06e8cb531874633f7c2372279c3b0f46235539fe546df8b/nltk-3.9.1-py3-none-any.whl", hash = "sha256:4fa26829c5b00715afe3061398a8989dc643b92ce7dd93fb4585a70930d168a1", size = 1505442 },
]
+[[package]]
+name = "nodeenv"
+version = "1.9.1"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/43/16/fc88b08840de0e0a72a2f9d8c6bae36be573e475a6326ae854bcc549fc45/nodeenv-1.9.1.tar.gz", hash = "sha256:6ec12890a2dab7946721edbfbcd91f3319c6ccc9aec47be7c7e6b7011ee6645f", size = 47437 }
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/d2/1d/1b658dbd2b9fa9c4c9f32accbfc0205d532c8c6194dc0f2a4c0428e7128a/nodeenv-1.9.1-py2.py3-none-any.whl", hash = "sha256:ba11c9782d29c27c70ffbdda2d7415098754709be8a7056d79a737cd901155c9", size = 22314 },
+]
+
[[package]]
name = "numpy"
version = "1.26.4"
@@ -2946,7 +3058,7 @@ name = "portalocker"
version = "2.10.1"
source = { registry = "https://pypi.org/simple" }
dependencies = [
- { name = "pywin32", marker = "platform_system == 'Windows'" },
+ { name = "pywin32", marker = "sys_platform == 'win32'" },
]
sdist = { url = "https://files.pythonhosted.org/packages/ed/d3/c6c64067759e87af98cc668c1cc75171347d0f1577fab7ca3749134e3cd4/portalocker-2.10.1.tar.gz", hash = "sha256:ef1bf844e878ab08aee7e40184156e1151f228f103aa5c6bd0724cc330960f8f", size = 40891 }
wheels = [
@@ -2969,6 +3081,22 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/d3/f2/5ee24cd69e2120bf87356c02ace0438b4e4fb78229fddcbf6f1c6be377d5/posthog-3.7.4-py2.py3-none-any.whl", hash = "sha256:21c18c6bf43b2de303ea4cd6e95804cc0f24c20cb2a96a8fd09da2ed50b62faa", size = 54777 },
]
+[[package]]
+name = "pre-commit"
+version = "4.0.1"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+ { name = "cfgv" },
+ { name = "identify" },
+ { name = "nodeenv" },
+ { name = "pyyaml" },
+ { name = "virtualenv" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/2e/c8/e22c292035f1bac8b9f5237a2622305bc0304e776080b246f3df57c4ff9f/pre_commit-4.0.1.tar.gz", hash = "sha256:80905ac375958c0444c65e9cebebd948b3cdb518f335a091a670a89d652139d2", size = 191678 }
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/16/8f/496e10d51edd6671ebe0432e33ff800aa86775d2d147ce7d43389324a525/pre_commit-4.0.1-py2.py3-none-any.whl", hash = "sha256:efde913840816312445dc98787724647c65473daefe420785f885e8ed9a06878", size = 218713 },
+]
+
[[package]]
name = "propcache"
version = "0.2.1"
@@ -3065,6 +3193,7 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/33/39/5a9a229bb5414abeb86e33b8fc8143ab0aecce5a7f698a53e31367d30caa/psycopg2-2.9.10-cp311-cp311-win_amd64.whl", hash = "sha256:0435034157049f6846e95103bd8f5a668788dd913a7c30162ca9503fdf542cb4", size = 1163736 },
{ url = "https://files.pythonhosted.org/packages/3d/16/4623fad6076448df21c1a870c93a9774ad8a7b4dd1660223b59082dd8fec/psycopg2-2.9.10-cp312-cp312-win32.whl", hash = "sha256:65a63d7ab0e067e2cdb3cf266de39663203d38d6a8ed97f5ca0cb315c73fe067", size = 1025113 },
{ url = "https://files.pythonhosted.org/packages/66/de/baed128ae0fc07460d9399d82e631ea31a1f171c0c4ae18f9808ac6759e3/psycopg2-2.9.10-cp312-cp312-win_amd64.whl", hash = "sha256:4a579d6243da40a7b3182e0430493dbd55950c493d8c68f4eec0b302f6bbf20e", size = 1163951 },
+ { url = "https://files.pythonhosted.org/packages/ae/49/a6cfc94a9c483b1fa401fbcb23aca7892f60c7269c5ffa2ac408364f80dc/psycopg2-2.9.10-cp313-cp313-win_amd64.whl", hash = "sha256:91fd603a2155da8d0cfcdbf8ab24a2d54bca72795b90d2a3ed2b6da8d979dee2", size = 2569060 },
]
[[package]]
@@ -4223,6 +4352,18 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/35/75/c4d8b2f0fe7dac22854d88a9c509d428e78ac4bf284bc54cfe83f75cc13b/time_machine-2.16.0-cp313-cp313-win_arm64.whl", hash = "sha256:4d3843143c46dddca6491a954bbd0abfd435681512ac343169560e9bab504129", size = 18047 },
]
+[[package]]
+name = "tinycss2"
+version = "1.4.0"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+ { name = "webencodings" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/7a/fd/7a5ee21fd08ff70d3d33a5781c255cbe779659bd03278feb98b19ee550f4/tinycss2-1.4.0.tar.gz", hash = "sha256:10c0972f6fc0fbee87c3edb76549357415e94548c1ae10ebccdea16fb404a9b7", size = 87085 }
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/e6/34/ebdc18bae6aa14fbee1a08b63c015c72b64868ff7dae68808ab500c492e2/tinycss2-1.4.0-py3-none-any.whl", hash = "sha256:3a49cf47b7675da0b15d0c6e1df8df4ebd96e9394bb905a5775adb0d884c5289", size = 26610 },
+]
+
[[package]]
name = "tokenizers"
version = "0.21.0"
@@ -4257,12 +4398,30 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/f9/b6/a447b5e4ec71e13871be01ba81f5dfc9d0af7e473da256ff46bc0e24026f/tomlkit-0.13.2-py3-none-any.whl", hash = "sha256:7a974427f6e119197f670fbbbeae7bef749a6c14e793db934baefc1b5f03efde", size = 37955 },
]
+[[package]]
+name = "tornado"
+version = "6.4.2"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/59/45/a0daf161f7d6f36c3ea5fc0c2de619746cc3dd4c76402e9db545bd920f63/tornado-6.4.2.tar.gz", hash = "sha256:92bad5b4746e9879fd7bf1eb21dce4e3fc5128d71601f80005afa39237ad620b", size = 501135 }
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/26/7e/71f604d8cea1b58f82ba3590290b66da1e72d840aeb37e0d5f7291bd30db/tornado-6.4.2-cp38-abi3-macosx_10_9_universal2.whl", hash = "sha256:e828cce1123e9e44ae2a50a9de3055497ab1d0aeb440c5ac23064d9e44880da1", size = 436299 },
+ { url = "https://files.pythonhosted.org/packages/96/44/87543a3b99016d0bf54fdaab30d24bf0af2e848f1d13d34a3a5380aabe16/tornado-6.4.2-cp38-abi3-macosx_10_9_x86_64.whl", hash = "sha256:072ce12ada169c5b00b7d92a99ba089447ccc993ea2143c9ede887e0937aa803", size = 434253 },
+ { url = "https://files.pythonhosted.org/packages/cb/fb/fdf679b4ce51bcb7210801ef4f11fdac96e9885daa402861751353beea6e/tornado-6.4.2-cp38-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:1a017d239bd1bb0919f72af256a970624241f070496635784d9bf0db640d3fec", size = 437602 },
+ { url = "https://files.pythonhosted.org/packages/4f/3b/e31aeffffc22b475a64dbeb273026a21b5b566f74dee48742817626c47dc/tornado-6.4.2-cp38-abi3-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:c36e62ce8f63409301537222faffcef7dfc5284f27eec227389f2ad11b09d946", size = 436972 },
+ { url = "https://files.pythonhosted.org/packages/22/55/b78a464de78051a30599ceb6983b01d8f732e6f69bf37b4ed07f642ac0fc/tornado-6.4.2-cp38-abi3-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:bca9eb02196e789c9cb5c3c7c0f04fb447dc2adffd95265b2c7223a8a615ccbf", size = 437173 },
+ { url = "https://files.pythonhosted.org/packages/79/5e/be4fb0d1684eb822c9a62fb18a3e44a06188f78aa466b2ad991d2ee31104/tornado-6.4.2-cp38-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:304463bd0772442ff4d0f5149c6f1c2135a1fae045adf070821c6cdc76980634", size = 437892 },
+ { url = "https://files.pythonhosted.org/packages/f5/33/4f91fdd94ea36e1d796147003b490fe60a0215ac5737b6f9c65e160d4fe0/tornado-6.4.2-cp38-abi3-musllinux_1_2_i686.whl", hash = "sha256:c82c46813ba483a385ab2a99caeaedf92585a1f90defb5693351fa7e4ea0bf73", size = 437334 },
+ { url = "https://files.pythonhosted.org/packages/2b/ae/c1b22d4524b0e10da2f29a176fb2890386f7bd1f63aacf186444873a88a0/tornado-6.4.2-cp38-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:932d195ca9015956fa502c6b56af9eb06106140d844a335590c1ec7f5277d10c", size = 437261 },
+ { url = "https://files.pythonhosted.org/packages/b5/25/36dbd49ab6d179bcfc4c6c093a51795a4f3bed380543a8242ac3517a1751/tornado-6.4.2-cp38-abi3-win32.whl", hash = "sha256:2876cef82e6c5978fde1e0d5b1f919d756968d5b4282418f3146b79b58556482", size = 438463 },
+ { url = "https://files.pythonhosted.org/packages/61/cc/58b1adeb1bb46228442081e746fcdbc4540905c87e8add7c277540934edb/tornado-6.4.2-cp38-abi3-win_amd64.whl", hash = "sha256:908b71bf3ff37d81073356a5fadcc660eb10c1476ee6e2725588626ce7e5ca38", size = 438907 },
+]
+
[[package]]
name = "tqdm"
version = "4.67.1"
source = { registry = "https://pypi.org/simple" }
dependencies = [
- { name = "colorama", marker = "platform_system == 'Windows'" },
+ { name = "colorama", marker = "sys_platform == 'win32'" },
]
sdist = { url = "https://files.pythonhosted.org/packages/a8/4b/29b4ef32e036bb34e4ab51796dd745cdba7ed47ad142a9f4a1eb8e0c744d/tqdm-4.67.1.tar.gz", hash = "sha256:f8aef9c52c08c13a65f30ea34f4e5aac3fd1a34959879d7e59e63027286627f2", size = 169737 }
wheels = [
@@ -4536,6 +4695,20 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/04/22/91b4bd36df27e651daedd93d03d5d3bb6029fdb0b55494e45ee46c36c570/validators-0.33.0-py3-none-any.whl", hash = "sha256:134b586a98894f8139865953899fc2daeb3d0c35569552c5518f089ae43ed075", size = 43298 },
]
+[[package]]
+name = "virtualenv"
+version = "20.28.0"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+ { name = "distlib" },
+ { name = "filelock" },
+ { name = "platformdirs" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/bf/75/53316a5a8050069228a2f6d11f32046cfa94fbb6cc3f08703f59b873de2e/virtualenv-20.28.0.tar.gz", hash = "sha256:2c9c3262bb8e7b87ea801d715fae4495e6032450c71d2309be9550e7364049aa", size = 7650368 }
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/10/f9/0919cf6f1432a8c4baa62511f8f8da8225432d22e83e3476f5be1a1edc6e/virtualenv-20.28.0-py3-none-any.whl", hash = "sha256:23eae1b4516ecd610481eda647f3a7c09aea295055337331bb4e6892ecce47b0", size = 4276702 },
+]
+
[[package]]
name = "weaviate-client"
version = "4.6.7"
@@ -4692,6 +4865,15 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/27/ee/518b72faa2073f5aa8e3262408d284892cb79cf2754ba0c3a5870645ef73/xxhash-3.5.0-cp313-cp313-win_arm64.whl", hash = "sha256:4811336f1ce11cac89dcbd18f3a25c527c16311709a89313c3acaf771def2d4b", size = 26801 },
]
+[[package]]
+name = "xyzservices"
+version = "2024.9.0"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/a0/16/ae87cbd2d6bfc40a419077521c35aadf5121725b7bee3d7c51b56f50958b/xyzservices-2024.9.0.tar.gz", hash = "sha256:68fb8353c9dbba4f1ff6c0f2e5e4e596bb9e1db7f94f4f7dfbcb26e25aa66fde", size = 1131900 }
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/4c/d3/e07ce413d16ef64e885bea37551eac4c5ca3ddd440933f9c94594273d0d9/xyzservices-2024.9.0-py3-none-any.whl", hash = "sha256:776ae82b78d6e5ca63dd6a94abb054df8130887a4a308473b54a6bd364de8644", size = 85130 },
+]
+
[[package]]
name = "yarl"
version = "1.18.3"
diff --git a/cognee/__init__.py b/cognee/__init__.py
index 591140b87..4ffd39577 100644
--- a/cognee/__init__.py
+++ b/cognee/__init__.py
@@ -4,11 +4,15 @@
from .api.v1.datasets.datasets import datasets
from .api.v1.prune import prune
from .api.v1.search import SearchType, get_search_history, search
+from .api.v1.visualize import visualize
+from .shared.utils import create_cognee_style_network_with_logo
+
# Pipelines
from .modules import pipelines
try:
import dotenv
+
dotenv.load_dotenv()
except ImportError:
pass
diff --git a/cognee/api/DTO.py b/cognee/api/DTO.py
index 6bdfb4f82..fc2e1cc5f 100644
--- a/cognee/api/DTO.py
+++ b/cognee/api/DTO.py
@@ -4,12 +4,13 @@
class OutDTO(BaseModel):
model_config = ConfigDict(
- alias_generator = to_camel,
- populate_by_name = True,
+ alias_generator=to_camel,
+ populate_by_name=True,
)
+
class InDTO(BaseModel):
model_config = ConfigDict(
- alias_generator = to_camel,
- populate_by_name = True,
+ alias_generator=to_camel,
+ populate_by_name=True,
)
diff --git a/cognee/api/client.py b/cognee/api/client.py
index 95a273351..227ada304 100644
--- a/cognee/api/client.py
+++ b/cognee/api/client.py
@@ -1,4 +1,5 @@
-""" FastAPI server for the Cognee API. """
+"""FastAPI server for the Cognee API."""
+
import os
import uvicorn
import logging
@@ -6,9 +7,26 @@
from fastapi import FastAPI, status
from fastapi.responses import JSONResponse, Response
from fastapi.middleware.cors import CORSMiddleware
-
+from cognee.api.v1.permissions.routers import get_permissions_router
+from cognee.api.v1.settings.routers import get_settings_router
+from cognee.api.v1.datasets.routers import get_datasets_router
+from cognee.api.v1.cognify.routers import get_cognify_router
+from cognee.api.v1.search.routers import get_search_router
+from cognee.api.v1.add.routers import get_add_router
+from fastapi import Request
+from fastapi.encoders import jsonable_encoder
+from fastapi.exceptions import RequestValidationError
from cognee.exceptions import CogneeApiError
from traceback import format_exc
+from cognee.api.v1.users.routers import (
+ get_auth_router,
+ get_register_router,
+ get_reset_password_router,
+ get_verify_router,
+ get_users_router,
+ get_visualize_router,
+)
+from contextlib import asynccontextmanager
# Set up logging
logging.basicConfig(
@@ -19,15 +37,15 @@
if os.getenv("ENV", "prod") == "prod":
sentry_sdk.init(
- dsn = os.getenv("SENTRY_REPORTING_URL"),
- traces_sample_rate = 1.0,
- profiles_sample_rate = 1.0,
+ dsn=os.getenv("SENTRY_REPORTING_URL"),
+ traces_sample_rate=1.0,
+ profiles_sample_rate=1.0,
)
-from contextlib import asynccontextmanager
app_environment = os.getenv("ENV", "prod")
+
@asynccontextmanager
async def lifespan(app: FastAPI):
# from cognee.modules.data.deletion import prune_system, prune_data
@@ -35,50 +53,42 @@ async def lifespan(app: FastAPI):
# await prune_system(metadata = True)
# if app_environment == "local" or app_environment == "dev":
from cognee.infrastructure.databases.relational import get_relational_engine
+
db_engine = get_relational_engine()
await db_engine.create_database()
from cognee.modules.users.methods import get_default_user
+
await get_default_user()
yield
-app = FastAPI(debug = app_environment != "prod", lifespan = lifespan)
+
+app = FastAPI(debug=app_environment != "prod", lifespan=lifespan)
app.add_middleware(
CORSMiddleware,
- allow_origins = ["*"],
- allow_credentials = True,
- allow_methods = ["OPTIONS", "GET", "POST", "DELETE"],
- allow_headers = ["*"],
+ allow_origins=["*"],
+ allow_credentials=True,
+ allow_methods=["OPTIONS", "GET", "POST", "DELETE"],
+ allow_headers=["*"],
)
-from cognee.api.v1.users.routers import get_auth_router, get_register_router,\
- get_reset_password_router, get_verify_router, get_users_router
-from cognee.api.v1.permissions.routers import get_permissions_router
-from cognee.api.v1.settings.routers import get_settings_router
-from cognee.api.v1.datasets.routers import get_datasets_router
-from cognee.api.v1.cognify.routers import get_cognify_router
-from cognee.api.v1.search.routers import get_search_router
-from cognee.api.v1.add.routers import get_add_router
-
-from fastapi import Request
-from fastapi.encoders import jsonable_encoder
-from fastapi.exceptions import RequestValidationError
@app.exception_handler(RequestValidationError)
async def request_validation_exception_handler(request: Request, exc: RequestValidationError):
if request.url.path == "/api/v1/auth/login":
return JSONResponse(
- status_code = 400,
- content = {"detail": "LOGIN_BAD_CREDENTIALS"},
+ status_code=400,
+ content={"detail": "LOGIN_BAD_CREDENTIALS"},
)
return JSONResponse(
- status_code = 400,
- content = jsonable_encoder({"detail": exc.errors(), "body": exc.body}),
+ status_code=400,
+ content=jsonable_encoder({"detail": exc.errors(), "body": exc.body}),
)
+
@app.exception_handler(CogneeApiError)
async def exception_handler(_: Request, exc: CogneeApiError) -> JSONResponse:
detail = {}
@@ -95,46 +105,42 @@ async def exception_handler(_: Request, exc: CogneeApiError) -> JSONResponse:
# log the stack trace for easier serverside debugging
logger.error(format_exc())
- return JSONResponse(
- status_code=status_code, content={"detail": detail["message"]}
- )
+ return JSONResponse(status_code=status_code, content={"detail": detail["message"]})
-app.include_router(
- get_auth_router(),
- prefix = "/api/v1/auth",
- tags = ["auth"]
-)
+
+app.include_router(get_auth_router(), prefix="/api/v1/auth", tags=["auth"])
app.include_router(
get_register_router(),
- prefix = "/api/v1/auth",
- tags = ["auth"],
+ prefix="/api/v1/auth",
+ tags=["auth"],
)
app.include_router(
get_reset_password_router(),
- prefix = "/api/v1/auth",
- tags = ["auth"],
+ prefix="/api/v1/auth",
+ tags=["auth"],
)
app.include_router(
get_verify_router(),
- prefix = "/api/v1/auth",
- tags = ["auth"],
+ prefix="/api/v1/auth",
+ tags=["auth"],
)
app.include_router(
get_users_router(),
- prefix = "/api/v1/users",
- tags = ["users"],
+ prefix="/api/v1/users",
+ tags=["users"],
)
app.include_router(
get_permissions_router(),
- prefix = "/api/v1/permissions",
- tags = ["permissions"],
+ prefix="/api/v1/permissions",
+ tags=["permissions"],
)
+
@app.get("/")
async def root():
"""
@@ -148,37 +154,21 @@ def health_check():
"""
Health check endpoint that returns the server status.
"""
- return Response(status_code = 200)
+ return Response(status_code=200)
-app.include_router(
- get_datasets_router(),
- prefix="/api/v1/datasets",
- tags=["datasets"]
-)
-app.include_router(
- get_add_router(),
- prefix="/api/v1/add",
- tags=["add"]
-)
+app.include_router(get_datasets_router(), prefix="/api/v1/datasets", tags=["datasets"])
-app.include_router(
- get_cognify_router(),
- prefix="/api/v1/cognify",
- tags=["cognify"]
-)
+app.include_router(get_add_router(), prefix="/api/v1/add", tags=["add"])
-app.include_router(
- get_search_router(),
- prefix="/api/v1/search",
- tags=["search"]
-)
+app.include_router(get_cognify_router(), prefix="/api/v1/cognify", tags=["cognify"])
+
+app.include_router(get_search_router(), prefix="/api/v1/search", tags=["search"])
+
+app.include_router(get_settings_router(), prefix="/api/v1/settings", tags=["settings"])
+
+app.include_router(get_visualize_router(), prefix="/api/v1/visualize", tags=["visualize"])
-app.include_router(
- get_settings_router(),
- prefix="/api/v1/settings",
- tags=["settings"]
-)
def start_api_server(host: str = "0.0.0.0", port: int = 8000):
"""
@@ -190,7 +180,7 @@ def start_api_server(host: str = "0.0.0.0", port: int = 8000):
try:
logger.info("Starting server at %s:%s", host, port)
- uvicorn.run(app, host = host, port = port)
+ uvicorn.run(app, host=host, port=port)
except Exception as e:
logger.exception(f"Failed to start server: {e}")
# Here you could add any cleanup code or error recovery code.
diff --git a/cognee/api/v1/add/add.py b/cognee/api/v1/add/add.py
index 39ee01964..c9f531269 100644
--- a/cognee/api/v1/add/add.py
+++ b/cognee/api/v1/add/add.py
@@ -14,10 +14,19 @@
from cognee.modules.users.permissions.methods import give_permission_on_document
from cognee.modules.users.models import User
from cognee.modules.data.methods import create_dataset
-from cognee.infrastructure.databases.relational import create_db_and_tables as create_relational_db_and_tables
-from cognee.infrastructure.databases.vector.pgvector import create_db_and_tables as create_pgvector_db_and_tables
-
-async def add(data: Union[BinaryIO, List[BinaryIO], str, List[str]], dataset_name: str = "main_dataset", user: User = None):
+from cognee.infrastructure.databases.relational import (
+ create_db_and_tables as create_relational_db_and_tables,
+)
+from cognee.infrastructure.databases.vector.pgvector import (
+ create_db_and_tables as create_pgvector_db_and_tables,
+)
+
+
+async def add(
+ data: Union[BinaryIO, List[BinaryIO], str, List[str]],
+ dataset_name: str = "main_dataset",
+ user: User = None,
+):
await create_relational_db_and_tables()
await create_pgvector_db_and_tables()
@@ -25,7 +34,9 @@ async def add(data: Union[BinaryIO, List[BinaryIO], str, List[str]], dataset_nam
if "data://" in data:
# data is a data directory path
datasets = get_matched_datasets(data.replace("data://", ""), dataset_name)
- return await asyncio.gather(*[add(file_paths, dataset_name) for [dataset_name, file_paths] in datasets])
+ return await asyncio.gather(
+ *[add(file_paths, dataset_name) for [dataset_name, file_paths] in datasets]
+ )
if "file://" in data:
# data is a file path
@@ -37,7 +48,7 @@ async def add(data: Union[BinaryIO, List[BinaryIO], str, List[str]], dataset_nam
return await add([file_path], dataset_name)
if hasattr(data, "file"):
- file_path = save_data_to_file(data.file, filename = data.filename)
+ file_path = save_data_to_file(data.file, filename=data.filename)
return await add([file_path], dataset_name)
# data is a list of file paths or texts
@@ -45,7 +56,7 @@ async def add(data: Union[BinaryIO, List[BinaryIO], str, List[str]], dataset_nam
for data_item in data:
if hasattr(data_item, "file"):
- file_paths.append(save_data_to_file(data_item, filename = data_item.filename))
+ file_paths.append(save_data_to_file(data_item, filename=data_item.filename))
elif isinstance(data_item, str) and (
data_item.startswith("/") or data_item.startswith("file://")
):
@@ -58,10 +69,11 @@ async def add(data: Union[BinaryIO, List[BinaryIO], str, List[str]], dataset_nam
return []
+
async def add_files(file_paths: List[str], dataset_name: str, user: User = None):
if user is None:
user = await get_default_user()
-
+
base_config = get_base_config()
data_directory_path = base_config.data_root_directory
@@ -72,7 +84,11 @@ async def add_files(file_paths: List[str], dataset_name: str, user: User = None)
if data_directory_path not in file_path:
file_name = file_path.split("/")[-1]
- file_directory_path = data_directory_path + "/" + (dataset_name.replace(".", "/") + "/" if dataset_name != "main_dataset" else "")
+ file_directory_path = (
+ data_directory_path
+ + "/"
+ + (dataset_name.replace(".", "/") + "/" if dataset_name != "main_dataset" else "")
+ )
dataset_file_path = path.join(file_directory_path, file_name)
LocalStorage.ensure_directory_exists(file_directory_path)
@@ -85,16 +101,20 @@ async def add_files(file_paths: List[str], dataset_name: str, user: User = None)
destination = get_dlt_destination()
pipeline = dlt.pipeline(
- pipeline_name = "file_load_from_filesystem",
- destination = destination,
+ pipeline_name="file_load_from_filesystem",
+ destination=destination,
)
- dataset_name = dataset_name.replace(" ", "_").replace(".", "_") if dataset_name is not None else "main_dataset"
+ dataset_name = (
+ dataset_name.replace(" ", "_").replace(".", "_")
+ if dataset_name is not None
+ else "main_dataset"
+ )
- @dlt.resource(standalone = True, merge_key = "id")
+ @dlt.resource(standalone=True, merge_key="id")
async def data_resources(file_paths: str, user: User):
for file_path in file_paths:
- with open(file_path.replace("file://", ""), mode = "rb") as file:
+ with open(file_path.replace("file://", ""), mode="rb") as file:
classified_data = ingestion.classify(file)
data_id = ingestion.identify(classified_data)
@@ -109,9 +129,9 @@ async def data_resources(file_paths: str, user: User):
async with db_engine.get_async_session() as session:
dataset = await create_dataset(dataset_name, user.id, session)
- data = (await session.execute(
- select(Data).filter(Data.id == data_id)
- )).scalar_one_or_none()
+ data = (
+ await session.execute(select(Data).filter(Data.id == data_id))
+ ).scalar_one_or_none()
if data is not None:
data.name = file_metadata["name"]
@@ -123,11 +143,11 @@ async def data_resources(file_paths: str, user: User):
await session.commit()
else:
data = Data(
- id = data_id,
- name = file_metadata["name"],
- raw_data_location = file_metadata["file_path"],
- extension = file_metadata["extension"],
- mime_type = file_metadata["mime_type"],
+ id=data_id,
+ name=file_metadata["name"],
+ raw_data_location=file_metadata["file_path"],
+ extension=file_metadata["extension"],
+ mime_type=file_metadata["mime_type"],
)
dataset.data.append(data)
@@ -144,14 +164,13 @@ async def data_resources(file_paths: str, user: User):
await give_permission_on_document(user, data_id, "read")
await give_permission_on_document(user, data_id, "write")
-
- send_telemetry("cognee.add EXECUTION STARTED", user_id = user.id)
+ send_telemetry("cognee.add EXECUTION STARTED", user_id=user.id)
run_info = pipeline.run(
data_resources(processed_file_paths, user),
- table_name = "file_metadata",
- dataset_name = dataset_name,
- write_disposition = "merge",
+ table_name="file_metadata",
+ dataset_name=dataset_name,
+ write_disposition="merge",
)
- send_telemetry("cognee.add EXECUTION COMPLETED", user_id = user.id)
+ send_telemetry("cognee.add EXECUTION COMPLETED", user_id=user.id)
return run_info
diff --git a/cognee/api/v1/add/add_v2.py b/cognee/api/v1/add/add_v2.py
index 637c4a187..1ec30e67b 100644
--- a/cognee/api/v1/add/add_v2.py
+++ b/cognee/api/v1/add/add_v2.py
@@ -3,22 +3,28 @@
from cognee.modules.users.methods import get_default_user
from cognee.modules.pipelines import run_tasks, Task
from cognee.tasks.ingestion import ingest_data_with_metadata, resolve_data_directories
-from cognee.infrastructure.databases.relational import create_db_and_tables as create_relational_db_and_tables
-from cognee.infrastructure.databases.vector.pgvector import create_db_and_tables as create_pgvector_db_and_tables
+from cognee.infrastructure.databases.relational import (
+ create_db_and_tables as create_relational_db_and_tables,
+)
+from cognee.infrastructure.databases.vector.pgvector import (
+ create_db_and_tables as create_pgvector_db_and_tables,
+)
-async def add(data: Union[BinaryIO, list[BinaryIO], str, list[str]], dataset_name: str = "main_dataset", user: User = None):
+
+async def add(
+ data: Union[BinaryIO, list[BinaryIO], str, list[str]],
+ dataset_name: str = "main_dataset",
+ user: User = None,
+):
await create_relational_db_and_tables()
await create_pgvector_db_and_tables()
if user is None:
user = await get_default_user()
- tasks = [
- Task(resolve_data_directories),
- Task(ingest_data_with_metadata, dataset_name, user)
- ]
+ tasks = [Task(resolve_data_directories), Task(ingest_data_with_metadata, dataset_name, user)]
pipeline = run_tasks(tasks, data, "add_pipeline")
async for result in pipeline:
- print(result)
\ No newline at end of file
+ print(result)
diff --git a/cognee/api/v1/add/routers/__init__.py b/cognee/api/v1/add/routers/__init__.py
index eebb250ab..bb820c033 100644
--- a/cognee/api/v1/add/routers/__init__.py
+++ b/cognee/api/v1/add/routers/__init__.py
@@ -1 +1 @@
-from .get_add_router import get_add_router
\ No newline at end of file
+from .get_add_router import get_add_router
diff --git a/cognee/api/v1/add/routers/get_add_router.py b/cognee/api/v1/add/routers/get_add_router.py
index 1f45d0c95..414552499 100644
--- a/cognee/api/v1/add/routers/get_add_router.py
+++ b/cognee/api/v1/add/routers/get_add_router.py
@@ -11,17 +11,19 @@
logger = logging.getLogger(__name__)
+
def get_add_router() -> APIRouter:
router = APIRouter()
@router.post("/", response_model=None)
async def add(
- data: List[UploadFile],
- datasetId: str = Form(...),
- user: User = Depends(get_authenticated_user),
+ data: List[UploadFile],
+ datasetId: str = Form(...),
+ user: User = Depends(get_authenticated_user),
):
- """ This endpoint is responsible for adding data to the graph."""
+ """This endpoint is responsible for adding data to the graph."""
from cognee.api.v1.add import add as cognee_add
+
try:
if isinstance(data, str) and data.startswith("http"):
if "github" in data:
@@ -52,9 +54,6 @@ async def add(
user=user,
)
except Exception as error:
- return JSONResponse(
- status_code=409,
- content={"error": str(error)}
- )
+ return JSONResponse(status_code=409, content={"error": str(error)})
- return router
\ No newline at end of file
+ return router
diff --git a/cognee/api/v1/authenticate_user/authenticate_user.py b/cognee/api/v1/authenticate_user/authenticate_user.py
index 6761c65b8..791fe0357 100644
--- a/cognee/api/v1/authenticate_user/authenticate_user.py
+++ b/cognee/api/v1/authenticate_user/authenticate_user.py
@@ -1,4 +1,6 @@
-from cognee.infrastructure.databases.relational.user_authentication.users import authenticate_user_method
+from cognee.infrastructure.databases.relational.user_authentication.users import (
+ authenticate_user_method,
+)
async def authenticate_user(email: str, password: str):
@@ -11,6 +13,7 @@ async def authenticate_user(email: str, password: str):
if __name__ == "__main__":
import asyncio
+
# Define an example user
example_email = "example@example.com"
example_password = "securepassword123"
diff --git a/cognee/api/v1/cognify/code_graph_pipeline.py b/cognee/api/v1/cognify/code_graph_pipeline.py
index 3d31b4000..405cb0b40 100644
--- a/cognee/api/v1/cognify/code_graph_pipeline.py
+++ b/cognee/api/v1/cognify/code_graph_pipeline.py
@@ -3,31 +3,30 @@
from pathlib import Path
from cognee.base_config import get_base_config
-from cognee.infrastructure.databases.vector.embeddings import \
- get_embedding_engine
+from cognee.infrastructure.databases.vector.embeddings import get_embedding_engine
from cognee.modules.cognify.config import get_cognify_config
from cognee.modules.pipelines import run_tasks
from cognee.modules.pipelines.tasks.Task import Task
from cognee.modules.users.methods import get_default_user
from cognee.shared.data_models import KnowledgeGraph, MonitoringTool
-from cognee.tasks.documents import (classify_documents,
- extract_chunks_from_documents)
+from cognee.tasks.documents import classify_documents, extract_chunks_from_documents
from cognee.tasks.graph import extract_graph_from_data
from cognee.tasks.ingestion import ingest_data_with_metadata
-from cognee.tasks.repo_processor import (enrich_dependency_graph,
- expand_dependency_graph,
- get_data_list_for_user,
- get_non_py_files,
- get_repo_file_dependencies)
-from cognee.tasks.repo_processor.get_source_code_chunks import \
- get_source_code_chunks
+from cognee.tasks.repo_processor import (
+ enrich_dependency_graph,
+ expand_dependency_graph,
+ get_data_list_for_user,
+ get_non_py_files,
+ get_repo_file_dependencies,
+)
+from cognee.tasks.repo_processor.get_source_code_chunks import get_source_code_chunks
from cognee.tasks.storage import add_data_points
+from cognee.tasks.summarization import summarize_code, summarize_text
monitoring = get_base_config().monitoring_tool
if monitoring == MonitoringTool.LANGFUSE:
from langfuse.decorators import observe
-from cognee.tasks.summarization import summarize_code, summarize_text
logger = logging.getLogger("code_graph_pipeline")
update_status_lock = asyncio.Lock()
@@ -42,9 +41,13 @@ async def run_code_graph_pipeline(repo_path, include_docs=True):
from cognee.infrastructure.databases.relational import create_db_and_tables
file_path = Path(__file__).parent
- data_directory_path = str(pathlib.Path(os.path.join(file_path, ".data_storage/code_graph")).resolve())
+ data_directory_path = str(
+ pathlib.Path(os.path.join(file_path, ".data_storage/code_graph")).resolve()
+ )
cognee.config.data_root_directory(data_directory_path)
- cognee_directory_path = str(pathlib.Path(os.path.join(file_path, ".cognee_system/code_graph")).resolve())
+ cognee_directory_path = str(
+ pathlib.Path(os.path.join(file_path, ".cognee_system/code_graph")).resolve()
+ )
cognee.config.system_root_directory(cognee_directory_path)
await cognee.prune.prune_data()
@@ -60,7 +63,11 @@ async def run_code_graph_pipeline(repo_path, include_docs=True):
Task(get_repo_file_dependencies),
Task(enrich_dependency_graph),
Task(expand_dependency_graph, task_config={"batch_size": 50}),
- Task(get_source_code_chunks, embedding_model=embedding_engine.model, task_config={"batch_size": 50}),
+ Task(
+ get_source_code_chunks,
+ embedding_model=embedding_engine.model,
+ task_config={"batch_size": 50},
+ ),
Task(summarize_code, task_config={"batch_size": 50}),
Task(add_data_points, task_config={"batch_size": 50}),
]
@@ -72,17 +79,19 @@ async def run_code_graph_pipeline(repo_path, include_docs=True):
Task(get_data_list_for_user, dataset_name="repo_docs", user=user),
Task(classify_documents),
Task(extract_chunks_from_documents),
- Task(extract_graph_from_data, graph_model=KnowledgeGraph, task_config={"batch_size": 50}),
+ Task(
+ extract_graph_from_data, graph_model=KnowledgeGraph, task_config={"batch_size": 50}
+ ),
Task(
summarize_text,
summarization_model=cognee_config.summarization_model,
- task_config={"batch_size": 50}
+ task_config={"batch_size": 50},
),
]
-
+
if include_docs:
async for result in run_tasks(non_code_tasks, repo_path):
yield result
async for result in run_tasks(tasks, repo_path, "cognify_code_pipeline"):
- yield result
\ No newline at end of file
+ yield result
diff --git a/cognee/api/v1/cognify/cognify_v2.py b/cognee/api/v1/cognify/cognify_v2.py
index 4e2db5a70..680c05828 100644
--- a/cognee/api/v1/cognify/cognify_v2.py
+++ b/cognee/api/v1/cognify/cognify_v2.py
@@ -10,18 +10,18 @@
from cognee.modules.data.models import Data, Dataset
from cognee.modules.pipelines import run_tasks
from cognee.modules.pipelines.models import PipelineRunStatus
-from cognee.modules.pipelines.operations.get_pipeline_status import \
- get_pipeline_status
-from cognee.modules.pipelines.operations.log_pipeline_status import \
- log_pipeline_status
+from cognee.modules.pipelines.operations.get_pipeline_status import get_pipeline_status
+from cognee.modules.pipelines.operations.log_pipeline_status import log_pipeline_status
from cognee.modules.pipelines.tasks.Task import Task
from cognee.modules.users.methods import get_default_user
from cognee.modules.users.models import User
from cognee.shared.data_models import KnowledgeGraph
from cognee.shared.utils import send_telemetry
-from cognee.tasks.documents import (check_permissions_on_documents,
- classify_documents,
- extract_chunks_from_documents)
+from cognee.tasks.documents import (
+ check_permissions_on_documents,
+ classify_documents,
+ extract_chunks_from_documents,
+)
from cognee.tasks.graph import extract_graph_from_data
from cognee.tasks.storage import add_data_points
from cognee.tasks.storage.index_graph_edges import index_graph_edges
@@ -31,7 +31,12 @@
update_status_lock = asyncio.Lock()
-async def cognify(datasets: Union[str, list[str]] = None, user: User = None, graph_model: BaseModel = KnowledgeGraph):
+
+async def cognify(
+ datasets: Union[str, list[str]] = None,
+ user: User = None,
+ graph_model: BaseModel = KnowledgeGraph,
+):
if user is None:
user = await get_default_user()
@@ -41,7 +46,7 @@ async def cognify(datasets: Union[str, list[str]] = None, user: User = None, gra
# If no datasets are provided, cognify all existing datasets.
datasets = existing_datasets
- if type(datasets[0]) == str:
+ if isinstance(datasets[0], str):
datasets = await get_datasets_by_name(datasets, user.id)
existing_datasets_map = {
@@ -59,8 +64,10 @@ async def cognify(datasets: Union[str, list[str]] = None, user: User = None, gra
return await asyncio.gather(*awaitables)
-async def run_cognify_pipeline(dataset: Dataset, user: User, graph_model: BaseModel = KnowledgeGraph):
- data_documents: list[Data] = await get_dataset_data(dataset_id = dataset.id)
+async def run_cognify_pipeline(
+ dataset: Dataset, user: User, graph_model: BaseModel = KnowledgeGraph
+):
+ data_documents: list[Data] = await get_dataset_data(dataset_id=dataset.id)
document_ids_str = [str(document.id) for document in data_documents]
@@ -69,32 +76,41 @@ async def run_cognify_pipeline(dataset: Dataset, user: User, graph_model: BaseMo
send_telemetry("cognee.cognify EXECUTION STARTED", user.id)
- #async with update_status_lock: TODO: Add UI lock to prevent multiple backend requests
+ # async with update_status_lock: TODO: Add UI lock to prevent multiple backend requests
task_status = await get_pipeline_status([dataset_id])
- if dataset_id in task_status and task_status[dataset_id] == PipelineRunStatus.DATASET_PROCESSING_STARTED:
+ if (
+ dataset_id in task_status
+ and task_status[dataset_id] == PipelineRunStatus.DATASET_PROCESSING_STARTED
+ ):
logger.info("Dataset %s is already being processed.", dataset_name)
return
- await log_pipeline_status(dataset_id, PipelineRunStatus.DATASET_PROCESSING_STARTED, {
- "dataset_name": dataset_name,
- "files": document_ids_str,
- })
+ await log_pipeline_status(
+ dataset_id,
+ PipelineRunStatus.DATASET_PROCESSING_STARTED,
+ {
+ "dataset_name": dataset_name,
+ "files": document_ids_str,
+ },
+ )
try:
cognee_config = get_cognify_config()
tasks = [
Task(classify_documents),
- Task(check_permissions_on_documents, user = user, permissions = ["write"]),
- Task(extract_chunks_from_documents), # Extract text chunks based on the document type.
- Task(extract_graph_from_data, graph_model = graph_model, task_config = { "batch_size": 10 }), # Generate knowledge graphs from the document chunks.
+ Task(check_permissions_on_documents, user=user, permissions=["write"]),
+ Task(extract_chunks_from_documents), # Extract text chunks based on the document type.
+ Task(
+ extract_graph_from_data, graph_model=graph_model, task_config={"batch_size": 10}
+ ), # Generate knowledge graphs from the document chunks.
Task(
summarize_text,
- summarization_model = cognee_config.summarization_model,
- task_config = { "batch_size": 10 }
+ summarization_model=cognee_config.summarization_model,
+ task_config={"batch_size": 10},
),
- Task(add_data_points, only_root = True, task_config = { "batch_size": 10 }),
+ Task(add_data_points, only_root=True, task_config={"batch_size": 10}),
]
pipeline = run_tasks(tasks, data_documents, "cognify_pipeline")
@@ -106,17 +122,25 @@ async def run_cognify_pipeline(dataset: Dataset, user: User, graph_model: BaseMo
send_telemetry("cognee.cognify EXECUTION COMPLETED", user.id)
- await log_pipeline_status(dataset_id, PipelineRunStatus.DATASET_PROCESSING_COMPLETED, {
- "dataset_name": dataset_name,
- "files": document_ids_str,
- })
+ await log_pipeline_status(
+ dataset_id,
+ PipelineRunStatus.DATASET_PROCESSING_COMPLETED,
+ {
+ "dataset_name": dataset_name,
+ "files": document_ids_str,
+ },
+ )
except Exception as error:
send_telemetry("cognee.cognify EXECUTION ERRORED", user.id)
- await log_pipeline_status(dataset_id, PipelineRunStatus.DATASET_PROCESSING_ERRORED, {
- "dataset_name": dataset_name,
- "files": document_ids_str,
- })
+ await log_pipeline_status(
+ dataset_id,
+ PipelineRunStatus.DATASET_PROCESSING_ERRORED,
+ {
+ "dataset_name": dataset_name,
+ "files": document_ids_str,
+ },
+ )
raise error
diff --git a/cognee/api/v1/cognify/routers/__init__.py b/cognee/api/v1/cognify/routers/__init__.py
index c6d52bfa2..6e5f9cc9d 100644
--- a/cognee/api/v1/cognify/routers/__init__.py
+++ b/cognee/api/v1/cognify/routers/__init__.py
@@ -1 +1 @@
-from .get_cognify_router import get_cognify_router
\ No newline at end of file
+from .get_cognify_router import get_cognify_router
diff --git a/cognee/api/v1/cognify/routers/get_cognify_router.py b/cognee/api/v1/cognify/routers/get_cognify_router.py
index 257ac994f..82ec4b857 100644
--- a/cognee/api/v1/cognify/routers/get_cognify_router.py
+++ b/cognee/api/v1/cognify/routers/get_cognify_router.py
@@ -7,23 +7,23 @@
from fastapi import Depends
from cognee.shared.data_models import KnowledgeGraph
+
class CognifyPayloadDTO(BaseModel):
datasets: List[str]
graph_model: Optional[BaseModel] = KnowledgeGraph
+
def get_cognify_router() -> APIRouter:
router = APIRouter()
@router.post("/", response_model=None)
async def cognify(payload: CognifyPayloadDTO, user: User = Depends(get_authenticated_user)):
- """ This endpoint is responsible for the cognitive processing of the content."""
+ """This endpoint is responsible for the cognitive processing of the content."""
from cognee.api.v1.cognify.cognify_v2 import cognify as cognee_cognify
+
try:
await cognee_cognify(payload.datasets, user, payload.graph_model)
except Exception as error:
- return JSONResponse(
- status_code=409,
- content={"error": str(error)}
- )
+ return JSONResponse(status_code=409, content={"error": str(error)})
return router
diff --git a/cognee/api/v1/config/config.py b/cognee/api/v1/config/config.py
index 1347fcba8..da58cf581 100644
--- a/cognee/api/v1/config/config.py
+++ b/cognee/api/v1/config/config.py
@@ -1,4 +1,5 @@
-""" This module is used to set the configuration of the system."""
+"""This module is used to set the configuration of the system."""
+
import os
from cognee.base_config import get_base_config
from cognee.exceptions import InvalidValueError, InvalidAttributeError
@@ -10,7 +11,8 @@
from cognee.infrastructure.databases.relational import get_relational_config
from cognee.infrastructure.files.storage import LocalStorage
-class config():
+
+class config:
@staticmethod
def system_root_directory(system_root_directory: str):
databases_directory_path = os.path.join(system_root_directory, "databases")
@@ -39,12 +41,12 @@ def monitoring_tool(monitoring_tool: object):
@staticmethod
def set_classification_model(classification_model: object):
cognify_config = get_cognify_config()
- cognify_config.classification_model = classification_model
+ cognify_config.classification_model = classification_model
@staticmethod
def set_summarization_model(summarization_model: object):
cognify_config = get_cognify_config()
- cognify_config.summarization_model=summarization_model
+ cognify_config.summarization_model = summarization_model
@staticmethod
def set_graph_model(graph_model: object):
@@ -79,14 +81,16 @@ def set_llm_api_key(llm_api_key: str):
@staticmethod
def set_llm_config(config_dict: dict):
"""
- Updates the llm config with values from config_dict.
+ Updates the llm config with values from config_dict.
"""
llm_config = get_llm_config()
for key, value in config_dict.items():
if hasattr(llm_config, key):
object.__setattr__(llm_config, key, value)
else:
- raise InvalidAttributeError(message=f"'{key}' is not a valid attribute of the config.")
+ raise InvalidAttributeError(
+ message=f"'{key}' is not a valid attribute of the config."
+ )
@staticmethod
def set_chunk_strategy(chunk_strategy: object):
@@ -108,7 +112,6 @@ def set_chunk_size(chunk_size: object):
chunk_config = get_chunk_config()
chunk_config.chunk_size = chunk_size
-
@staticmethod
def set_vector_db_provider(vector_db_provider: str):
vector_db_config = get_vectordb_config()
@@ -117,33 +120,36 @@ def set_vector_db_provider(vector_db_provider: str):
@staticmethod
def set_relational_db_config(config_dict: dict):
"""
- Updates the relational db config with values from config_dict.
+ Updates the relational db config with values from config_dict.
"""
relational_db_config = get_relational_config()
for key, value in config_dict.items():
if hasattr(relational_db_config, key):
object.__setattr__(relational_db_config, key, value)
else:
- raise InvalidAttributeError(message=f"'{key}' is not a valid attribute of the config.")
+ raise InvalidAttributeError(
+ message=f"'{key}' is not a valid attribute of the config."
+ )
@staticmethod
def set_vector_db_config(config_dict: dict):
"""
- Updates the vector db config with values from config_dict.
+ Updates the vector db config with values from config_dict.
"""
vector_db_config = get_vectordb_config()
for key, value in config_dict.items():
if hasattr(vector_db_config, key):
object.__setattr__(vector_db_config, key, value)
else:
- raise InvalidAttributeError(message=f"'{key}' is not a valid attribute of the config.")
+ raise InvalidAttributeError(
+ message=f"'{key}' is not a valid attribute of the config."
+ )
@staticmethod
def set_vector_db_key(db_key: str):
vector_db_config = get_vectordb_config()
vector_db_config.vector_db_key = db_key
-
@staticmethod
def set_vector_db_url(db_url: str):
vector_db_config = get_vectordb_config()
@@ -154,7 +160,9 @@ def set_graphistry_config(graphistry_config: dict[str, str]):
base_config = get_base_config()
if "username" not in graphistry_config or "password" not in graphistry_config:
- raise InvalidValueError(message="graphistry_config dictionary must contain 'username' and 'password' keys.")
+ raise InvalidValueError(
+ message="graphistry_config dictionary must contain 'username' and 'password' keys."
+ )
base_config.graphistry_username = graphistry_config.get("username")
base_config.graphistry_password = graphistry_config.get("password")
diff --git a/cognee/api/v1/datasets/datasets.py b/cognee/api/v1/datasets/datasets.py
index 345f8be7b..a6552d937 100644
--- a/cognee/api/v1/datasets/datasets.py
+++ b/cognee/api/v1/datasets/datasets.py
@@ -1,10 +1,13 @@
from cognee.modules.users.methods import get_default_user
from cognee.modules.ingestion import discover_directory_datasets
from cognee.modules.pipelines.operations.get_pipeline_status import get_pipeline_status
-class datasets():
+
+
+class datasets:
@staticmethod
async def list_datasets():
from cognee.modules.data.methods import get_datasets
+
user = await get_default_user()
return await get_datasets(user.id)
@@ -15,6 +18,7 @@ def discover_datasets(directory_path: str):
@staticmethod
async def list_data(dataset_id: str):
from cognee.modules.data.methods import get_dataset, get_dataset_data
+
user = await get_default_user()
dataset = await get_dataset(user.id, dataset_id)
@@ -28,6 +32,7 @@ async def get_status(dataset_ids: list[str]) -> dict:
@staticmethod
async def delete_dataset(dataset_id: str):
from cognee.modules.data.methods import get_dataset, delete_dataset
+
user = await get_default_user()
dataset = await get_dataset(user.id, dataset_id)
diff --git a/cognee/api/v1/datasets/routers/__init__.py b/cognee/api/v1/datasets/routers/__init__.py
index f03428fd6..24783903d 100644
--- a/cognee/api/v1/datasets/routers/__init__.py
+++ b/cognee/api/v1/datasets/routers/__init__.py
@@ -1 +1 @@
-from .get_datasets_router import get_datasets_router
\ No newline at end of file
+from .get_datasets_router import get_datasets_router
diff --git a/cognee/api/v1/datasets/routers/get_datasets_router.py b/cognee/api/v1/datasets/routers/get_datasets_router.py
index 1ba96a232..ba65cd94c 100644
--- a/cognee/api/v1/datasets/routers/get_datasets_router.py
+++ b/cognee/api/v1/datasets/routers/get_datasets_router.py
@@ -16,9 +16,11 @@
logger = logging.getLogger(__name__)
+
class ErrorResponseDTO(BaseModel):
message: str
+
class DatasetDTO(OutDTO):
id: UUID
name: str
@@ -26,6 +28,7 @@ class DatasetDTO(OutDTO):
updated_at: Optional[datetime] = None
owner_id: UUID
+
class DataDTO(OutDTO):
id: UUID
name: str
@@ -35,6 +38,7 @@ class DataDTO(OutDTO):
mime_type: str
raw_data_location: str
+
def get_datasets_router() -> APIRouter:
router = APIRouter()
@@ -42,46 +46,51 @@ def get_datasets_router() -> APIRouter:
async def get_datasets(user: User = Depends(get_authenticated_user)):
try:
from cognee.modules.data.methods import get_datasets
+
datasets = await get_datasets(user.id)
return datasets
except Exception as error:
logger.error(f"Error retrieving datasets: {str(error)}")
- raise HTTPException(status_code=500, detail=f"Error retrieving datasets: {str(error)}") from error
+ raise HTTPException(
+ status_code=500, detail=f"Error retrieving datasets: {str(error)}"
+ ) from error
- @router.delete("/{dataset_id}", response_model=None, responses={404: {"model": ErrorResponseDTO}})
+ @router.delete(
+ "/{dataset_id}", response_model=None, responses={404: {"model": ErrorResponseDTO}}
+ )
async def delete_dataset(dataset_id: str, user: User = Depends(get_authenticated_user)):
from cognee.modules.data.methods import get_dataset, delete_dataset
dataset = await get_dataset(user.id, dataset_id)
if dataset is None:
- raise EntityNotFoundError(
- message=f"Dataset ({dataset_id}) not found."
- )
+ raise EntityNotFoundError(message=f"Dataset ({dataset_id}) not found.")
await delete_dataset(dataset)
- @router.delete("/{dataset_id}/data/{data_id}", response_model=None, responses={404: {"model": ErrorResponseDTO}})
- async def delete_data(dataset_id: str, data_id: str, user: User = Depends(get_authenticated_user)):
+ @router.delete(
+ "/{dataset_id}/data/{data_id}",
+ response_model=None,
+ responses={404: {"model": ErrorResponseDTO}},
+ )
+ async def delete_data(
+ dataset_id: str, data_id: str, user: User = Depends(get_authenticated_user)
+ ):
from cognee.modules.data.methods import get_data, delete_data
from cognee.modules.data.methods import get_dataset
# Check if user has permission to access dataset and data by trying to get the dataset
dataset = await get_dataset(user.id, dataset_id)
- #TODO: Handle situation differently if user doesn't have permission to access data?
+ # TODO: Handle situation differently if user doesn't have permission to access data?
if dataset is None:
- raise EntityNotFoundError(
- message=f"Dataset ({dataset_id}) not found."
- )
+ raise EntityNotFoundError(message=f"Dataset ({dataset_id}) not found.")
data = await get_data(user.id, data_id)
if data is None:
- raise EntityNotFoundError(
- message=f"Data ({data_id}) not found."
- )
+ raise EntityNotFoundError(message=f"Data ({data_id}) not found.")
await delete_data(data)
@@ -98,14 +107,18 @@ async def get_dataset_graph(dataset_id: str, user: User = Depends(get_authentica
status_code=200,
content=str(graph_url),
)
- except:
+ except Exception as error:
+ print(error)
return JSONResponse(
status_code=409,
content="Graphistry credentials are not set. Please set them in your .env file.",
)
- @router.get("/{dataset_id}/data", response_model=list[DataDTO],
- responses={404: {"model": ErrorResponseDTO}})
+ @router.get(
+ "/{dataset_id}/data",
+ response_model=list[DataDTO],
+ responses={404: {"model": ErrorResponseDTO}},
+ )
async def get_dataset_data(dataset_id: str, user: User = Depends(get_authenticated_user)):
from cognee.modules.data.methods import get_dataset_data, get_dataset
@@ -125,8 +138,10 @@ async def get_dataset_data(dataset_id: str, user: User = Depends(get_authenticat
return dataset_data
@router.get("/status", response_model=dict[str, PipelineRunStatus])
- async def get_dataset_status(datasets: Annotated[List[str], Query(alias="dataset")] = None,
- user: User = Depends(get_authenticated_user)):
+ async def get_dataset_status(
+ datasets: Annotated[List[str], Query(alias="dataset")] = None,
+ user: User = Depends(get_authenticated_user),
+ ):
from cognee.api.v1.datasets.datasets import datasets as cognee_datasets
try:
@@ -134,13 +149,12 @@ async def get_dataset_status(datasets: Annotated[List[str], Query(alias="dataset
return datasets_statuses
except Exception as error:
- return JSONResponse(
- status_code=409,
- content={"error": str(error)}
- )
+ return JSONResponse(status_code=409, content={"error": str(error)})
@router.get("/{dataset_id}/data/{data_id}/raw", response_class=FileResponse)
- async def get_raw_data(dataset_id: str, data_id: str, user: User = Depends(get_authenticated_user)):
+ async def get_raw_data(
+ dataset_id: str, data_id: str, user: User = Depends(get_authenticated_user)
+ ):
from cognee.modules.data.methods import get_data
from cognee.modules.data.methods import get_dataset, get_dataset_data
@@ -148,10 +162,7 @@ async def get_raw_data(dataset_id: str, data_id: str, user: User = Depends(get_a
if dataset is None:
return JSONResponse(
- status_code=404,
- content={
- "detail": f"Dataset ({dataset_id}) not found."
- }
+ status_code=404, content={"detail": f"Dataset ({dataset_id}) not found."}
)
dataset_data = await get_dataset_data(dataset.id)
@@ -163,13 +174,17 @@ async def get_raw_data(dataset_id: str, data_id: str, user: User = Depends(get_a
# Check if matching_data contains an element
if len(matching_data) == 0:
- raise EntityNotFoundError(message= f"Data ({data_id}) not found in dataset ({dataset_id}).")
+ raise EntityNotFoundError(
+ message=f"Data ({data_id}) not found in dataset ({dataset_id})."
+ )
data = await get_data(user.id, data_id)
if data is None:
- raise EntityNotFoundError(message=f"Data ({data_id}) not found in dataset ({dataset_id}).")
+ raise EntityNotFoundError(
+ message=f"Data ({data_id}) not found in dataset ({dataset_id})."
+ )
return data.raw_data_location
- return router
\ No newline at end of file
+ return router
diff --git a/cognee/api/v1/permissions/routers/__init__.py b/cognee/api/v1/permissions/routers/__init__.py
index 986b52c3e..be7e5bbc4 100644
--- a/cognee/api/v1/permissions/routers/__init__.py
+++ b/cognee/api/v1/permissions/routers/__init__.py
@@ -1 +1 @@
-from .get_permissions_router import get_permissions_router
\ No newline at end of file
+from .get_permissions_router import get_permissions_router
diff --git a/cognee/api/v1/permissions/routers/get_permissions_router.py b/cognee/api/v1/permissions/routers/get_permissions_router.py
index 2b30f62fd..78b822085 100644
--- a/cognee/api/v1/permissions/routers/get_permissions_router.py
+++ b/cognee/api/v1/permissions/routers/get_permissions_router.py
@@ -10,40 +10,56 @@
from cognee.modules.users import get_user_db
from cognee.modules.users.models import User, Group, Permission, UserGroup, GroupPermission
+
def get_permissions_router() -> APIRouter:
permissions_router = APIRouter()
@permissions_router.post("/groups/{group_id}/permissions")
- async def give_permission_to_group(group_id: str, permission: str, db: Session = Depends(get_user_db)):
- group = (await db.session.execute(select(Group).where(Group.id == group_id))).scalars().first()
+ async def give_permission_to_group(
+ group_id: str, permission: str, db: Session = Depends(get_user_db)
+ ):
+ group = (
+ (await db.session.execute(select(Group).where(Group.id == group_id))).scalars().first()
+ )
if not group:
raise GroupNotFoundError
permission_entity = (
- await db.session.execute(select(Permission).where(Permission.name == permission))).scalars().first()
+ (await db.session.execute(select(Permission).where(Permission.name == permission)))
+ .scalars()
+ .first()
+ )
if not permission_entity:
stmt = insert(Permission).values(name=permission)
await db.session.execute(stmt)
permission_entity = (
- await db.session.execute(select(Permission).where(Permission.name == permission))).scalars().first()
+ (await db.session.execute(select(Permission).where(Permission.name == permission)))
+ .scalars()
+ .first()
+ )
try:
# add permission to group
await db.session.execute(
- insert(GroupPermission).values(group_id=group.id, permission_id=permission_entity.id))
- except IntegrityError as e:
+ insert(GroupPermission).values(
+ group_id=group.id, permission_id=permission_entity.id
+ )
+ )
+ except IntegrityError:
raise EntityAlreadyExistsError(message="Group permission already exists.")
await db.session.commit()
- return JSONResponse(status_code = 200, content = {"message": "Permission assigned to group"})
+ return JSONResponse(status_code=200, content={"message": "Permission assigned to group"})
@permissions_router.post("/users/{user_id}/groups")
async def add_user_to_group(user_id: str, group_id: str, db: Session = Depends(get_user_db)):
user = (await db.session.execute(select(User).where(User.id == user_id))).scalars().first()
- group = (await db.session.execute(select(Group).where(Group.id == group_id))).scalars().first()
+ group = (
+ (await db.session.execute(select(Group).where(Group.id == group_id))).scalars().first()
+ )
if not user:
raise UserNotFoundError
@@ -54,11 +70,11 @@ async def add_user_to_group(user_id: str, group_id: str, db: Session = Depends(g
# Add association directly to the association table
stmt = insert(UserGroup).values(user_id=user_id, group_id=group_id)
await db.session.execute(stmt)
- except IntegrityError as e:
+ except IntegrityError:
raise EntityAlreadyExistsError(message="User is already part of group.")
await db.session.commit()
- return JSONResponse(status_code = 200, content = {"message": "User added to group"})
+ return JSONResponse(status_code=200, content={"message": "User added to group"})
return permissions_router
diff --git a/cognee/api/v1/prune/prune.py b/cognee/api/v1/prune/prune.py
index 4237ac58e..efee07e21 100644
--- a/cognee/api/v1/prune/prune.py
+++ b/cognee/api/v1/prune/prune.py
@@ -1,19 +1,21 @@
from cognee.modules.data.deletion import prune_system, prune_data
-class prune():
+
+class prune:
@staticmethod
async def prune_data():
await prune_data()
@staticmethod
- async def prune_system(graph = True, vector = True, metadata = False):
+ async def prune_system(graph=True, vector=True, metadata=False):
await prune_system(graph, vector, metadata)
+
if __name__ == "__main__":
import asyncio
+
async def main():
await prune.prune_data()
await prune.prune_system()
-
asyncio.run(main())
diff --git a/cognee/api/v1/search/get_search_history.py b/cognee/api/v1/search/get_search_history.py
index fada67c85..04ec49978 100644
--- a/cognee/api/v1/search/get_search_history.py
+++ b/cognee/api/v1/search/get_search_history.py
@@ -2,8 +2,9 @@
from cognee.modules.users.methods import get_default_user
from cognee.modules.users.models import User
+
async def get_search_history(user: User = None) -> list:
if not user:
user = await get_default_user()
-
+
return await get_history(user.id)
diff --git a/cognee/api/v1/search/routers/__init__.py b/cognee/api/v1/search/routers/__init__.py
index c3b199f5f..d30e865e2 100644
--- a/cognee/api/v1/search/routers/__init__.py
+++ b/cognee/api/v1/search/routers/__init__.py
@@ -1 +1 @@
-from .get_search_router import get_search_router
\ No newline at end of file
+from .get_search_router import get_search_router
diff --git a/cognee/api/v1/search/routers/get_search_router.py b/cognee/api/v1/search/routers/get_search_router.py
index 893067c20..a97e84cf4 100644
--- a/cognee/api/v1/search/routers/get_search_router.py
+++ b/cognee/api/v1/search/routers/get_search_router.py
@@ -13,6 +13,7 @@ class SearchPayloadDTO(InDTO):
search_type: SearchType
query: str
+
def get_search_router() -> APIRouter:
router = APIRouter()
@@ -22,21 +23,18 @@ class SearchHistoryItem(OutDTO):
user: str
created_at: datetime
- @router.get("/", response_model = list[SearchHistoryItem])
+ @router.get("/", response_model=list[SearchHistoryItem])
async def get_search_history(user: User = Depends(get_authenticated_user)):
try:
history = await get_history(user.id)
return history
except Exception as error:
- return JSONResponse(
- status_code = 500,
- content = {"error": str(error)}
- )
+ return JSONResponse(status_code=500, content={"error": str(error)})
- @router.post("/", response_model = list)
+ @router.post("/", response_model=list)
async def search(payload: SearchPayloadDTO, user: User = Depends(get_authenticated_user)):
- """ This endpoint is responsible for searching for nodes in the graph."""
+ """This endpoint is responsible for searching for nodes in the graph."""
from cognee.api.v1.search import search as cognee_search
try:
@@ -44,9 +42,6 @@ async def search(payload: SearchPayloadDTO, user: User = Depends(get_authenticat
return results
except Exception as error:
- return JSONResponse(
- status_code = 409,
- content = {"error": str(error)}
- )
+ return JSONResponse(status_code=409, content={"error": str(error)})
return router
diff --git a/cognee/api/v1/search/search.legacy.py b/cognee/api/v1/search/search.legacy.py
index cea3b3874..c4e490f01 100644
--- a/cognee/api/v1/search/search.legacy.py
+++ b/cognee/api/v1/search/search.legacy.py
@@ -1,4 +1,5 @@
-""" This module contains the search function that is used to search for nodes in the graph."""
+"""This module contains the search function that is used to search for nodes in the graph."""
+
import asyncio
from enum import Enum
from typing import Dict, Any, Callable, List
@@ -16,6 +17,7 @@
from cognee.modules.users.methods import get_default_user
from cognee.modules.users.models import User
+
class SearchType(Enum):
ADJACENT = "ADJACENT"
TRAVERSE = "TRAVERSE"
@@ -23,7 +25,7 @@ class SearchType(Enum):
SUMMARY = "SUMMARY"
SUMMARY_CLASSIFICATION = "SUMMARY_CLASSIFICATION"
NODE_CLASSIFICATION = "NODE_CLASSIFICATION"
- DOCUMENT_CLASSIFICATION = "DOCUMENT_CLASSIFICATION",
+ DOCUMENT_CLASSIFICATION = ("DOCUMENT_CLASSIFICATION",)
CYPHER = "CYPHER"
@staticmethod
@@ -33,12 +35,13 @@ def from_str(name: str):
except KeyError as error:
raise ValueError(f"{name} is not a valid SearchType") from error
+
class SearchParameters(BaseModel):
search_type: SearchType
params: Dict[str, Any]
@field_validator("search_type", mode="before")
- def convert_string_to_enum(cls, value): # pylint: disable=no-self-argument
+ def convert_string_to_enum(cls, value): # pylint: disable=no-self-argument
if isinstance(value, str):
return SearchType.from_str(value)
return value
@@ -52,7 +55,7 @@ async def search(search_type: str, params: Dict[str, Any], user: User = None) ->
raise UserNotFoundError
own_document_ids = await get_document_ids_for_user(user.id)
- search_params = SearchParameters(search_type = search_type, params = params)
+ search_params = SearchParameters(search_type=search_type, params=params)
search_results = await specific_search([search_params], user)
from uuid import UUID
@@ -61,7 +64,7 @@ async def search(search_type: str, params: Dict[str, Any], user: User = None) ->
for search_result in search_results:
document_id = search_result["document_id"] if "document_id" in search_result else None
- document_id = UUID(document_id) if type(document_id) == str else document_id
+ document_id = UUID(document_id) if isinstance(document_id, str) else document_id
if document_id is None or document_id in own_document_ids:
filtered_search_results.append(search_result)
diff --git a/cognee/api/v1/search/search_v2.py b/cognee/api/v1/search/search_v2.py
index 222ec6791..52fc565c7 100644
--- a/cognee/api/v1/search/search_v2.py
+++ b/cognee/api/v1/search/search_v2.py
@@ -16,14 +16,20 @@
from cognee.tasks.summarization import query_summaries
from cognee.tasks.completion import query_completion
+
class SearchType(Enum):
SUMMARIES = "SUMMARIES"
INSIGHTS = "INSIGHTS"
CHUNKS = "CHUNKS"
COMPLETION = "COMPLETION"
-async def search(query_type: SearchType, query_text: str, user: User = None,
- datasets: Union[list[str], str, None] = None) -> list:
+
+async def search(
+ query_type: SearchType,
+ query_text: str,
+ user: User = None,
+ datasets: Union[list[str], str, None] = None,
+) -> list:
# We use lists from now on for datasets
if isinstance(datasets, str):
datasets = [datasets]
@@ -43,15 +49,16 @@ async def search(query_type: SearchType, query_text: str, user: User = None,
for search_result in search_results:
document_id = search_result["document_id"] if "document_id" in search_result else None
- document_id = UUID(document_id) if type(document_id) == str else document_id
+ document_id = UUID(document_id) if isinstance(document_id, str) else document_id
if document_id is None or document_id in own_document_ids:
filtered_search_results.append(search_result)
- await log_result(query.id, json.dumps(filtered_search_results, cls = JSONEncoder), user.id)
+ await log_result(query.id, json.dumps(filtered_search_results, cls=JSONEncoder), user.id)
return filtered_search_results
+
async def specific_search(query_type: SearchType, query: str, user) -> list:
search_tasks: Dict[SearchType, Callable] = {
SearchType.SUMMARIES: query_summaries,
diff --git a/cognee/api/v1/settings/routers/__init__.py b/cognee/api/v1/settings/routers/__init__.py
index 363d26610..4cb4fc49b 100644
--- a/cognee/api/v1/settings/routers/__init__.py
+++ b/cognee/api/v1/settings/routers/__init__.py
@@ -1 +1 @@
-from .get_settings_router import get_settings_router
\ No newline at end of file
+from .get_settings_router import get_settings_router
diff --git a/cognee/api/v1/settings/routers/get_settings_router.py b/cognee/api/v1/settings/routers/get_settings_router.py
index 31692382b..138bea661 100644
--- a/cognee/api/v1/settings/routers/get_settings_router.py
+++ b/cognee/api/v1/settings/routers/get_settings_router.py
@@ -6,40 +6,50 @@
from cognee.modules.users.models import User
from cognee.modules.settings.get_settings import LLMConfig, VectorDBConfig
+
class LLMConfigOutputDTO(OutDTO, LLMConfig):
pass
+
class VectorDBConfigOutputDTO(OutDTO, VectorDBConfig):
pass
+
class SettingsDTO(OutDTO):
llm: LLMConfigOutputDTO
vector_db: VectorDBConfigOutputDTO
+
class LLMConfigInputDTO(InDTO):
provider: Union[Literal["openai"], Literal["ollama"], Literal["anthropic"]]
model: str
api_key: str
+
class VectorDBConfigInputDTO(InDTO):
provider: Union[Literal["lancedb"], Literal["qdrant"], Literal["weaviate"], Literal["pgvector"]]
url: str
api_key: str
+
class SettingsPayloadDTO(InDTO):
llm: Optional[LLMConfigInputDTO] = None
vector_db: Optional[VectorDBConfigInputDTO] = None
+
def get_settings_router() -> APIRouter:
router = APIRouter()
@router.get("/", response_model=SettingsDTO)
async def get_settings(user: User = Depends(get_authenticated_user)):
from cognee.modules.settings import get_settings as get_cognee_settings
+
return get_cognee_settings()
@router.post("/", response_model=None)
- async def save_settings(new_settings: SettingsPayloadDTO, user: User = Depends(get_authenticated_user)):
+ async def save_settings(
+ new_settings: SettingsPayloadDTO, user: User = Depends(get_authenticated_user)
+ ):
from cognee.modules.settings import save_llm_config, save_vector_db_config
if new_settings.llm is not None:
@@ -48,4 +58,4 @@ async def save_settings(new_settings: SettingsPayloadDTO, user: User = Depends(g
if new_settings.vector_db is not None:
await save_vector_db_config(new_settings.vector_db)
- return router
\ No newline at end of file
+ return router
diff --git a/cognee/api/v1/users/create_user.py b/cognee/api/v1/users/create_user.py
index eba7a6e89..adfbf4ee1 100644
--- a/cognee/api/v1/users/create_user.py
+++ b/cognee/api/v1/users/create_user.py
@@ -3,10 +3,10 @@
async def create_user(email: str, password: str, is_superuser: bool = False):
user = await create_user_method(
- email = email,
- password = password,
- is_superuser = is_superuser,
- is_verified = True,
+ email=email,
+ password=password,
+ is_superuser=is_superuser,
+ is_verified=True,
)
return user
diff --git a/cognee/api/v1/users/routers/__init__.py b/cognee/api/v1/users/routers/__init__.py
index 482aac265..ae4b7ca56 100644
--- a/cognee/api/v1/users/routers/__init__.py
+++ b/cognee/api/v1/users/routers/__init__.py
@@ -3,3 +3,4 @@
from .get_reset_password_router import get_reset_password_router
from .get_users_router import get_users_router
from .get_verify_router import get_verify_router
+from .get_visualize_router import get_visualize_router
diff --git a/cognee/api/v1/users/routers/get_auth_router.py b/cognee/api/v1/users/routers/get_auth_router.py
index 8a65cde35..b6acc62a6 100644
--- a/cognee/api/v1/users/routers/get_auth_router.py
+++ b/cognee/api/v1/users/routers/get_auth_router.py
@@ -1,6 +1,7 @@
from cognee.modules.users.get_fastapi_users import get_fastapi_users
from cognee.modules.users.authentication.get_auth_backend import get_auth_backend
+
def get_auth_router():
auth_backend = get_auth_backend()
return get_fastapi_users().get_auth_router(auth_backend)
diff --git a/cognee/api/v1/users/routers/get_register_router.py b/cognee/api/v1/users/routers/get_register_router.py
index a1152c01c..8857b01da 100644
--- a/cognee/api/v1/users/routers/get_register_router.py
+++ b/cognee/api/v1/users/routers/get_register_router.py
@@ -1,5 +1,6 @@
from cognee.modules.users.get_fastapi_users import get_fastapi_users
from cognee.modules.users.models.User import UserRead, UserCreate
+
def get_register_router():
return get_fastapi_users().get_register_router(UserRead, UserCreate)
diff --git a/cognee/api/v1/users/routers/get_reset_password_router.py b/cognee/api/v1/users/routers/get_reset_password_router.py
index c058abe2a..33708c72a 100644
--- a/cognee/api/v1/users/routers/get_reset_password_router.py
+++ b/cognee/api/v1/users/routers/get_reset_password_router.py
@@ -1,4 +1,5 @@
from cognee.modules.users.get_fastapi_users import get_fastapi_users
+
def get_reset_password_router():
return get_fastapi_users().get_reset_password_router()
diff --git a/cognee/api/v1/users/routers/get_users_router.py b/cognee/api/v1/users/routers/get_users_router.py
index b81be73b0..776aefba0 100644
--- a/cognee/api/v1/users/routers/get_users_router.py
+++ b/cognee/api/v1/users/routers/get_users_router.py
@@ -1,5 +1,6 @@
from cognee.modules.users.get_fastapi_users import get_fastapi_users
from cognee.modules.users.models.User import UserRead, UserUpdate
+
def get_users_router():
return get_fastapi_users().get_users_router(UserRead, UserUpdate)
diff --git a/cognee/api/v1/users/routers/get_verify_router.py b/cognee/api/v1/users/routers/get_verify_router.py
index 0c18b08c2..aaf52b339 100644
--- a/cognee/api/v1/users/routers/get_verify_router.py
+++ b/cognee/api/v1/users/routers/get_verify_router.py
@@ -1,5 +1,6 @@
from cognee.modules.users.get_fastapi_users import get_fastapi_users
from cognee.modules.users.models.User import UserRead
+
def get_verify_router():
return get_fastapi_users().get_verify_router(UserRead)
diff --git a/cognee/api/v1/users/routers/get_visualize_router.py b/cognee/api/v1/users/routers/get_visualize_router.py
new file mode 100644
index 000000000..3e678cccd
--- /dev/null
+++ b/cognee/api/v1/users/routers/get_visualize_router.py
@@ -0,0 +1,32 @@
+from fastapi import Form, UploadFile, Depends
+from fastapi.responses import JSONResponse
+from fastapi import APIRouter
+from typing import List
+import aiohttp
+import subprocess
+import logging
+import os
+from cognee.modules.users.models import User
+from cognee.modules.users.methods import get_authenticated_user
+
+logger = logging.getLogger(__name__)
+
+
+def get_visualize_router() -> APIRouter:
+ router = APIRouter()
+
+ @router.post("/", response_model=None)
+ async def visualize(
+ user: User = Depends(get_authenticated_user),
+ ):
+ """This endpoint is responsible for adding data to the graph."""
+ from cognee.api.v1.visualize import visualize_graph
+
+ try:
+ html_visualization = await visualize_graph()
+ return html_visualization
+
+ except Exception as error:
+ return JSONResponse(status_code=409, content={"error": str(error)})
+
+ return router
diff --git a/cognee/api/v1/visualize/__init__.py b/cognee/api/v1/visualize/__init__.py
new file mode 100644
index 000000000..80db0367c
--- /dev/null
+++ b/cognee/api/v1/visualize/__init__.py
@@ -0,0 +1 @@
+from .visualize import visualize_graph
diff --git a/cognee/api/v1/visualize/visualize.py b/cognee/api/v1/visualize/visualize.py
new file mode 100644
index 000000000..3267f9d3d
--- /dev/null
+++ b/cognee/api/v1/visualize/visualize.py
@@ -0,0 +1,14 @@
+from cognee.shared.utils import create_cognee_style_network_with_logo, graph_to_tuple
+from cognee.infrastructure.databases.graph import get_graph_engine
+import logging
+
+
+async def visualize_graph(label: str = "name"):
+ """ """
+ graph_engine = await get_graph_engine()
+ graph_data = await graph_engine.get_graph_data()
+ logging.info(graph_data)
+
+ graph = await create_cognee_style_network_with_logo(graph_data, label=label)
+
+ return graph
diff --git a/cognee/base_config.py b/cognee/base_config.py
index f50f5a13d..6b1b8811e 100644
--- a/cognee/base_config.py
+++ b/cognee/base_config.py
@@ -5,6 +5,7 @@
from cognee.root_dir import get_absolute_path
from cognee.shared.data_models import MonitoringTool
+
class BaseConfig(BaseSettings):
data_root_directory: str = get_absolute_path(".data_storage")
monitoring_tool: object = MonitoringTool.LANGFUSE
@@ -13,7 +14,7 @@ class BaseConfig(BaseSettings):
langfuse_public_key: Optional[str] = os.getenv("LANGFUSE_PUBLIC_KEY")
langfuse_secret_key: Optional[str] = os.getenv("LANGFUSE_SECRET_KEY")
langfuse_host: Optional[str] = os.getenv("LANGFUSE_HOST")
- model_config = SettingsConfigDict(env_file = ".env", extra = "allow")
+ model_config = SettingsConfigDict(env_file=".env", extra="allow")
def to_dict(self) -> dict:
return {
@@ -21,6 +22,7 @@ def to_dict(self) -> dict:
"monitoring_tool": self.monitoring_tool,
}
+
@lru_cache
def get_base_config():
return BaseConfig()
diff --git a/cognee/exceptions/__init__.py b/cognee/exceptions/__init__.py
index 40120e0e1..1432afcc8 100644
--- a/cognee/exceptions/__init__.py
+++ b/cognee/exceptions/__init__.py
@@ -10,4 +10,4 @@
ServiceError,
InvalidValueError,
InvalidAttributeError,
-)
\ No newline at end of file
+)
diff --git a/cognee/exceptions/exceptions.py b/cognee/exceptions/exceptions.py
index f94daf8c9..c4ce73f8c 100644
--- a/cognee/exceptions/exceptions.py
+++ b/cognee/exceptions/exceptions.py
@@ -3,6 +3,7 @@
logger = logging.getLogger(__name__)
+
class CogneeApiError(Exception):
"""Base exception class"""
@@ -36,19 +37,19 @@ def __init__(
class InvalidValueError(CogneeApiError):
def __init__(
- self,
- message: str = "Invalid Value.",
- name: str = "InvalidValueError",
- status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
+ self,
+ message: str = "Invalid Value.",
+ name: str = "InvalidValueError",
+ status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
):
super().__init__(message, name, status_code)
class InvalidAttributeError(CogneeApiError):
def __init__(
- self,
- message: str = "Invalid attribute.",
- name: str = "InvalidAttributeError",
- status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
+ self,
+ message: str = "Invalid attribute.",
+ name: str = "InvalidAttributeError",
+ status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
):
- super().__init__(message, name, status_code)
\ No newline at end of file
+ super().__init__(message, name, status_code)
diff --git a/cognee/fetch_secret.py b/cognee/fetch_secret.py
index 65b5326a9..c36f9e4ce 100644
--- a/cognee/fetch_secret.py
+++ b/cognee/fetch_secret.py
@@ -15,7 +15,7 @@
environment = os.getenv("AWS_ENV", "dev")
-def fetch_secret(secret_name:str, region_name:str, env_file_path:str):
+def fetch_secret(secret_name: str, region_name: str, env_file_path: str):
"""Fetch the secret from AWS Secrets Manager and write it to the .env file."""
print("Initializing session")
session = boto3.session.Session()
diff --git a/cognee/infrastructure/data/chunking/DefaultChunkEngine.py b/cognee/infrastructure/data/chunking/DefaultChunkEngine.py
index 2ca879b25..0d67d5e38 100644
--- a/cognee/infrastructure/data/chunking/DefaultChunkEngine.py
+++ b/cognee/infrastructure/data/chunking/DefaultChunkEngine.py
@@ -1,4 +1,5 @@
-""" Chunking strategies for splitting text into smaller parts."""
+"""Chunking strategies for splitting text into smaller parts."""
+
from __future__ import annotations
import re
from cognee.shared.data_models import ChunkStrategy
@@ -6,17 +7,15 @@
# /Users/vasa/Projects/cognee/cognee/infrastructure/data/chunking/DefaultChunkEngine.py
-class DefaultChunkEngine():
+
+class DefaultChunkEngine:
def __init__(self, chunk_strategy=None, chunk_size=None, chunk_overlap=None):
self.chunk_strategy = chunk_strategy
self.chunk_size = chunk_size
self.chunk_overlap = chunk_overlap
-
@staticmethod
- def _split_text_with_regex(
- text: str, separator: str, keep_separator: bool
- ) -> list[str]:
+ def _split_text_with_regex(text: str, separator: str, keep_separator: bool) -> list[str]:
# Now that we have the separator, split the text
if separator:
if keep_separator:
@@ -32,13 +31,12 @@ def _split_text_with_regex(
splits = list(text)
return [s for s in splits if s != ""]
-
-
- def chunk_data(self,
- chunk_strategy = None,
- source_data = None,
- chunk_size = None,
- chunk_overlap = None,
+ def chunk_data(
+ self,
+ chunk_strategy=None,
+ source_data=None,
+ chunk_size=None,
+ chunk_overlap=None,
):
"""
Chunk data based on the specified strategy.
@@ -54,44 +52,47 @@ def chunk_data(self,
"""
if self.chunk_strategy == ChunkStrategy.PARAGRAPH:
- chunked_data, chunk_number = self.chunk_data_by_paragraph(source_data,chunk_size=self.chunk_size, chunk_overlap=self.chunk_overlap)
+ chunked_data, chunk_number = self.chunk_data_by_paragraph(
+ source_data, chunk_size=self.chunk_size, chunk_overlap=self.chunk_overlap
+ )
elif self.chunk_strategy == ChunkStrategy.SENTENCE:
- chunked_data, chunk_number = self.chunk_by_sentence(source_data, chunk_size = self.chunk_size, chunk_overlap=self.chunk_overlap)
+ chunked_data, chunk_number = self.chunk_by_sentence(
+ source_data, chunk_size=self.chunk_size, chunk_overlap=self.chunk_overlap
+ )
elif self.chunk_strategy == ChunkStrategy.EXACT:
- chunked_data, chunk_number = self.chunk_data_exact(source_data, chunk_size = self.chunk_size, chunk_overlap=self.chunk_overlap)
+ chunked_data, chunk_number = self.chunk_data_exact(
+ source_data, chunk_size=self.chunk_size, chunk_overlap=self.chunk_overlap
+ )
else:
chunked_data, chunk_number = "Invalid chunk strategy.", [0, "Invalid chunk strategy."]
-
return chunked_data, chunk_number
-
-
def chunk_data_exact(self, data_chunks, chunk_size, chunk_overlap):
data = "".join(data_chunks)
chunks = []
for i in range(0, len(data), chunk_size - chunk_overlap):
- chunks.append(data[i:i + chunk_size])
+ chunks.append(data[i : i + chunk_size])
numbered_chunks = []
for i, chunk in enumerate(chunks):
numbered_chunk = [i + 1, chunk]
numbered_chunks.append(numbered_chunk)
return chunks, numbered_chunks
-
-
def chunk_by_sentence(self, data_chunks, chunk_size, chunk_overlap):
# Split by periods, question marks, exclamation marks, and ellipses
data = "".join(data_chunks)
# The regular expression is used to find series of charaters that end with one the following chaacters (. ! ? ...)
- sentence_endings = r'(?<=[.!?…]) +'
+ sentence_endings = r"(?<=[.!?…]) +"
sentences = re.split(sentence_endings, data)
sentence_chunks = []
for sentence in sentences:
if len(sentence) > chunk_size:
- chunks = self.chunk_data_exact(data_chunks=[sentence], chunk_size=chunk_size, chunk_overlap=chunk_overlap)
+ chunks = self.chunk_data_exact(
+ data_chunks=[sentence], chunk_size=chunk_size, chunk_overlap=chunk_overlap
+ )
sentence_chunks.extend(chunks)
else:
sentence_chunks.append(sentence)
@@ -102,9 +103,7 @@ def chunk_by_sentence(self, data_chunks, chunk_size, chunk_overlap):
numbered_chunks.append(numbered_chunk)
return sentence_chunks, numbered_chunks
-
-
- def chunk_data_by_paragraph(self, data_chunks, chunk_size, chunk_overlap, bound = 0.75):
+ def chunk_data_by_paragraph(self, data_chunks, chunk_size, chunk_overlap, bound=0.75):
data = "".join(data_chunks)
total_length = len(data)
chunks = []
diff --git a/cognee/infrastructure/data/chunking/LangchainChunkingEngine.py b/cognee/infrastructure/data/chunking/LangchainChunkingEngine.py
index c58f7517c..5de642ca8 100644
--- a/cognee/infrastructure/data/chunking/LangchainChunkingEngine.py
+++ b/cognee/infrastructure/data/chunking/LangchainChunkingEngine.py
@@ -5,7 +5,6 @@
from cognee.shared.data_models import ChunkStrategy
-
class LangchainChunkEngine:
def __init__(self, chunk_strategy=None, source_data=None, chunk_size=None, chunk_overlap=None):
self.chunk_strategy = chunk_strategy
@@ -13,13 +12,12 @@ def __init__(self, chunk_strategy=None, source_data=None, chunk_size=None, chunk
self.chunk_size = chunk_size
self.chunk_overlap = chunk_overlap
-
-
- def chunk_data(self,
- chunk_strategy = None,
- source_data = None,
- chunk_size = None,
- chunk_overlap = None,
+ def chunk_data(
+ self,
+ chunk_strategy=None,
+ source_data=None,
+ chunk_size=None,
+ chunk_overlap=None,
):
"""
Chunk data based on the specified strategy.
@@ -35,20 +33,24 @@ def chunk_data(self,
"""
if chunk_strategy == ChunkStrategy.CODE:
- chunked_data, chunk_number = self.chunk_data_by_code(source_data,self.chunk_size, self.chunk_overlap)
+ chunked_data, chunk_number = self.chunk_data_by_code(
+ source_data, self.chunk_size, self.chunk_overlap
+ )
elif chunk_strategy == ChunkStrategy.LANGCHAIN_CHARACTER:
- chunked_data, chunk_number = self.chunk_data_by_character(source_data,self.chunk_size, self.chunk_overlap)
+ chunked_data, chunk_number = self.chunk_data_by_character(
+ source_data, self.chunk_size, self.chunk_overlap
+ )
else:
- chunked_data, chunk_number = "Invalid chunk strategy.", [0, "Invalid chunk strategy."]
+ chunked_data, chunk_number = "Invalid chunk strategy.", [0, "Invalid chunk strategy."]
return chunked_data, chunk_number
-
- def chunk_data_by_code(self, data_chunks, chunk_size, chunk_overlap= 10, language=None):
+ def chunk_data_by_code(self, data_chunks, chunk_size, chunk_overlap=10, language=None):
from langchain_text_splitters import (
Language,
RecursiveCharacterTextSplitter,
)
+
if language is None:
language = Language.PYTHON
python_splitter = RecursiveCharacterTextSplitter.from_language(
@@ -67,7 +69,10 @@ def chunk_data_by_code(self, data_chunks, chunk_size, chunk_overlap= 10, languag
def chunk_data_by_character(self, data_chunks, chunk_size=1500, chunk_overlap=10):
from langchain_text_splitters import RecursiveCharacterTextSplitter
- splitter = RecursiveCharacterTextSplitter(chunk_size =chunk_size, chunk_overlap=chunk_overlap)
+
+ splitter = RecursiveCharacterTextSplitter(
+ chunk_size=chunk_size, chunk_overlap=chunk_overlap
+ )
data_chunks = splitter.create_documents([data_chunks])
only_content = [chunk.page_content for chunk in data_chunks]
@@ -78,4 +83,3 @@ def chunk_data_by_character(self, data_chunks, chunk_size=1500, chunk_overlap=10
numbered_chunks.append(numbered_chunk)
return only_content, numbered_chunks
-
diff --git a/cognee/infrastructure/data/chunking/config.py b/cognee/infrastructure/data/chunking/config.py
index b917b8ddd..e367e4de8 100644
--- a/cognee/infrastructure/data/chunking/config.py
+++ b/cognee/infrastructure/data/chunking/config.py
@@ -11,8 +11,7 @@ class ChunkConfig(BaseSettings):
chunk_strategy: object = ChunkStrategy.PARAGRAPH
chunk_engine: object = ChunkEngine.DEFAULT_ENGINE
-
- model_config = SettingsConfigDict(env_file = ".env", extra = "allow")
+ model_config = SettingsConfigDict(env_file=".env", extra="allow")
def to_dict(self) -> dict:
return {
diff --git a/cognee/infrastructure/data/chunking/create_chunking_engine.py b/cognee/infrastructure/data/chunking/create_chunking_engine.py
index 5cd0d9077..7966bbf6c 100644
--- a/cognee/infrastructure/data/chunking/create_chunking_engine.py
+++ b/cognee/infrastructure/data/chunking/create_chunking_engine.py
@@ -9,9 +9,11 @@ class ChunkingConfig(Dict):
vector_db_key: str
vector_db_provider: str
+
def create_chunking_engine(config: ChunkingConfig):
if config["chunk_engine"] == ChunkEngine.LANGCHAIN_ENGINE:
from cognee.infrastructure.data.chunking.LangchainChunkingEngine import LangchainChunkEngine
+
return LangchainChunkEngine(
chunk_size=config["chunk_size"],
chunk_overlap=config["chunk_overlap"],
diff --git a/cognee/infrastructure/data/chunking/get_chunking_engine.py b/cognee/infrastructure/data/chunking/get_chunking_engine.py
index d188ffa9b..8d40a35a9 100644
--- a/cognee/infrastructure/data/chunking/get_chunking_engine.py
+++ b/cognee/infrastructure/data/chunking/get_chunking_engine.py
@@ -2,5 +2,6 @@
from .create_chunking_engine import create_chunking_engine
+
def get_chunk_engine():
return create_chunking_engine(get_chunk_config().to_dict())
diff --git a/cognee/infrastructure/data/utils/extract_keywords.py b/cognee/infrastructure/data/utils/extract_keywords.py
index 11f061889..42b06884a 100644
--- a/cognee/infrastructure/data/utils/extract_keywords.py
+++ b/cognee/infrastructure/data/utils/extract_keywords.py
@@ -3,6 +3,7 @@
from cognee.exceptions import InvalidValueError
from cognee.shared.utils import extract_pos_tags
+
def extract_keywords(text: str) -> list[str]:
if len(text) == 0:
raise InvalidValueError(message="extract_keywords cannot extract keywords from empty text.")
@@ -14,9 +15,7 @@ def extract_keywords(text: str) -> list[str]:
tfidf = vectorizer.fit_transform(nouns)
top_nouns = sorted(
- vectorizer.vocabulary_,
- key = lambda x: tfidf[0, vectorizer.vocabulary_[x]],
- reverse = True
+ vectorizer.vocabulary_, key=lambda x: tfidf[0, vectorizer.vocabulary_[x]], reverse=True
)
keywords = []
diff --git a/cognee/infrastructure/databases/exceptions/EmbeddingException.py b/cognee/infrastructure/databases/exceptions/EmbeddingException.py
index ba7c70d80..130282857 100644
--- a/cognee/infrastructure/databases/exceptions/EmbeddingException.py
+++ b/cognee/infrastructure/databases/exceptions/EmbeddingException.py
@@ -1,3 +1,4 @@
class EmbeddingException(Exception):
"""Custom exception for handling embedding-related errors."""
- pass
\ No newline at end of file
+
+ pass
diff --git a/cognee/infrastructure/databases/exceptions/__init__.py b/cognee/infrastructure/databases/exceptions/__init__.py
index 5836e7d11..7c74db3df 100644
--- a/cognee/infrastructure/databases/exceptions/__init__.py
+++ b/cognee/infrastructure/databases/exceptions/__init__.py
@@ -7,4 +7,4 @@
from .exceptions import (
EntityNotFoundError,
EntityAlreadyExistsError,
-)
\ No newline at end of file
+)
diff --git a/cognee/infrastructure/databases/exceptions/exceptions.py b/cognee/infrastructure/databases/exceptions/exceptions.py
index af15bb616..854e620ff 100644
--- a/cognee/infrastructure/databases/exceptions/exceptions.py
+++ b/cognee/infrastructure/databases/exceptions/exceptions.py
@@ -1,6 +1,7 @@
from cognee.exceptions import CogneeApiError
from fastapi import status
+
class EntityNotFoundError(CogneeApiError):
"""Database returns nothing"""
@@ -22,4 +23,4 @@ def __init__(
name: str = "EntityAlreadyExistsError",
status_code=status.HTTP_409_CONFLICT,
):
- super().__init__(message, name, status_code)
\ No newline at end of file
+ super().__init__(message, name, status_code)
diff --git a/cognee/infrastructure/databases/graph/config.py b/cognee/infrastructure/databases/graph/config.py
index 8e6da37c7..b24a9e964 100644
--- a/cognee/infrastructure/databases/graph/config.py
+++ b/cognee/infrastructure/databases/graph/config.py
@@ -1,4 +1,4 @@
-""" This module contains the configuration for the graph database. """
+"""This module contains the configuration for the graph database."""
import os
from functools import lru_cache
@@ -15,8 +15,7 @@ class GraphConfig(BaseSettings):
graph_database_password: str = ""
graph_database_port: int = 123
graph_file_path: str = os.path.join(
- os.path.join(get_absolute_path(".cognee_system"), "databases"),
- graph_filename
+ os.path.join(get_absolute_path(".cognee_system"), "databases"), graph_filename
)
graph_model: object = KnowledgeGraph
graph_topology: object = KnowledgeGraph
diff --git a/cognee/infrastructure/databases/graph/get_graph_engine.py b/cognee/infrastructure/databases/graph/get_graph_engine.py
index 10c16bbdd..4660a610f 100644
--- a/cognee/infrastructure/databases/graph/get_graph_engine.py
+++ b/cognee/infrastructure/databases/graph/get_graph_engine.py
@@ -1,4 +1,5 @@
"""Factory function to get the appropriate graph client based on the graph type."""
+
from functools import lru_cache
from .config import get_graph_config
@@ -26,7 +27,11 @@ def create_graph_engine() -> GraphDBInterface:
config = get_graph_config()
if config.graph_database_provider == "neo4j":
- if not (config.graph_database_url and config.graph_database_username and config.graph_database_password):
+ if not (
+ config.graph_database_url
+ and config.graph_database_username
+ and config.graph_database_password
+ ):
raise EnvironmentError("Missing required Neo4j credentials.")
from .neo4j_driver.adapter import Neo4jAdapter
@@ -34,7 +39,7 @@ def create_graph_engine() -> GraphDBInterface:
return Neo4jAdapter(
graph_database_url=config.graph_database_url,
graph_database_username=config.graph_database_username,
- graph_database_password=config.graph_database_password
+ graph_database_password=config.graph_database_password,
)
elif config.graph_database_provider == "falkordb":
@@ -53,6 +58,7 @@ def create_graph_engine() -> GraphDBInterface:
)
from .networkx.adapter import NetworkXAdapter
+
graph_client = NetworkXAdapter(filename=config.graph_file_path)
return graph_client
diff --git a/cognee/infrastructure/databases/graph/graph_db_interface.py b/cognee/infrastructure/databases/graph/graph_db_interface.py
index bcc09658c..30acc1b95 100644
--- a/cognee/infrastructure/databases/graph/graph_db_interface.py
+++ b/cognee/infrastructure/databases/graph/graph_db_interface.py
@@ -1,47 +1,35 @@
from typing import Protocol, Optional, Dict, Any
from abc import abstractmethod
+
class GraphDBInterface(Protocol):
@abstractmethod
async def query(self, query: str, params: dict):
raise NotImplementedError
@abstractmethod
- async def add_node(
- self,
- node_id: str,
- node_properties: dict
- ): raise NotImplementedError
+ async def add_node(self, node_id: str, node_properties: dict):
+ raise NotImplementedError
@abstractmethod
- async def add_nodes(
- self,
- nodes: list[tuple[str, dict]]
- ): raise NotImplementedError
+ async def add_nodes(self, nodes: list[tuple[str, dict]]):
+ raise NotImplementedError
@abstractmethod
- async def delete_node(
- self,
- node_id: str
- ): raise NotImplementedError
+ async def delete_node(self, node_id: str):
+ raise NotImplementedError
@abstractmethod
- async def delete_nodes(
- self,
- node_ids: list[str]
- ): raise NotImplementedError
+ async def delete_nodes(self, node_ids: list[str]):
+ raise NotImplementedError
@abstractmethod
- async def extract_node(
- self,
- node_id: str
- ): raise NotImplementedError
+ async def extract_node(self, node_id: str):
+ raise NotImplementedError
@abstractmethod
- async def extract_nodes(
- self,
- node_ids: list[str]
- ): raise NotImplementedError
+ async def extract_nodes(self, node_ids: list[str]):
+ raise NotImplementedError
@abstractmethod
async def add_edge(
@@ -49,21 +37,20 @@ async def add_edge(
from_node: str,
to_node: str,
relationship_name: str,
- edge_properties: Optional[Dict[str, Any]] = None
- ): raise NotImplementedError
+ edge_properties: Optional[Dict[str, Any]] = None,
+ ):
+ raise NotImplementedError
@abstractmethod
- async def add_edges(
- self,
- edges: tuple[str, str, str, dict]
- ): raise NotImplementedError
+ async def add_edges(self, edges: tuple[str, str, str, dict]):
+ raise NotImplementedError
@abstractmethod
async def delete_graph(
self,
- ): raise NotImplementedError
+ ):
+ raise NotImplementedError
@abstractmethod
- async def get_graph_data(
- self
- ): raise NotImplementedError
+ async def get_graph_data(self):
+ raise NotImplementedError
diff --git a/cognee/infrastructure/databases/graph/neo4j_driver/adapter.py b/cognee/infrastructure/databases/graph/neo4j_driver/adapter.py
index e6520e4e2..5490f6b43 100644
--- a/cognee/infrastructure/databases/graph/neo4j_driver/adapter.py
+++ b/cognee/infrastructure/databases/graph/neo4j_driver/adapter.py
@@ -1,4 +1,5 @@
-""" Neo4j Adapter for Graph Database"""
+"""Neo4j Adapter for Graph Database"""
+
import logging
import asyncio
from textwrap import dedent
@@ -13,6 +14,7 @@
logger = logging.getLogger("Neo4jAdapter")
+
class Neo4jAdapter(GraphDBInterface):
def __init__(
self,
@@ -23,8 +25,8 @@ def __init__(
):
self.driver = driver or AsyncGraphDatabase.driver(
graph_database_url,
- auth = (graph_database_username, graph_database_password),
- max_connection_lifetime = 120
+ auth=(graph_database_username, graph_database_password),
+ max_connection_lifetime=120,
)
@asynccontextmanager
@@ -39,11 +41,11 @@ async def query(
) -> List[Dict[str, Any]]:
try:
async with self.get_session() as session:
- result = await session.run(query, parameters = params)
+ result = await session.run(query, parameters=params)
data = await result.data()
return data
except Neo4jError as error:
- logger.error("Neo4j query error: %s", error, exc_info = True)
+ logger.error("Neo4j query error: %s", error, exc_info=True)
raise error
async def has_node(self, node_id: str) -> bool:
@@ -53,7 +55,7 @@ async def has_node(self, node_id: str) -> bool:
WHERE n.id = $node_id
RETURN COUNT(n) > 0 AS node_exists
""",
- {"node_id": node_id}
+ {"node_id": node_id},
)
return results[0]["node_exists"] if len(results) > 0 else False
@@ -83,15 +85,17 @@ async def add_nodes(self, nodes: list[DataPoint]) -> None:
RETURN ID(labeledNode) AS internal_id, labeledNode.id AS nodeId
"""
- nodes = [{
- "node_id": str(node.id),
- "properties": self.serialize_properties(node.model_dump()),
- } for node in nodes]
+ nodes = [
+ {
+ "node_id": str(node.id),
+ "properties": self.serialize_properties(node.model_dump()),
+ }
+ for node in nodes
+ ]
- results = await self.query(query, dict(nodes = nodes))
+ results = await self.query(query, dict(nodes=nodes))
return results
-
async def extract_node(self, node_id: str):
results = await self.extract_nodes([node_id])
@@ -103,9 +107,7 @@ async def extract_nodes(self, node_ids: List[str]):
MATCH (node {id: id})
RETURN node"""
- params = {
- "node_ids": node_ids
- }
+ params = {"node_ids": node_ids}
results = await self.query(query, params)
@@ -115,7 +117,7 @@ async def delete_node(self, node_id: str):
node_id = id.replace(":", "_")
query = f"MATCH (node:`{node_id}` {{id: $node_id}}) DETACH DELETE n"
- params = { "node_id": node_id }
+ params = {"node_id": node_id}
return await self.query(query, params)
@@ -125,9 +127,7 @@ async def delete_nodes(self, node_ids: list[str]) -> None:
MATCH (node {id: id})
DETACH DELETE node"""
- params = {
- "node_ids": node_ids
- }
+ params = {"node_ids": node_ids}
return await self.query(query, params)
@@ -157,21 +157,29 @@ async def has_edges(self, edges):
try:
params = {
- "edges": [{
- "from_node": str(edge[0]),
- "to_node": str(edge[1]),
- "relationship_name": edge[2],
- } for edge in edges],
+ "edges": [
+ {
+ "from_node": str(edge[0]),
+ "to_node": str(edge[1]),
+ "relationship_name": edge[2],
+ }
+ for edge in edges
+ ],
}
results = await self.query(query, params)
return [result["edge_exists"] for result in results]
except Neo4jError as error:
- logger.error("Neo4j query error: %s", error, exc_info = True)
+ logger.error("Neo4j query error: %s", error, exc_info=True)
raise error
-
- async def add_edge(self, from_node: UUID, to_node: UUID, relationship_name: str, edge_properties: Optional[Dict[str, Any]] = {}):
+ async def add_edge(
+ self,
+ from_node: UUID,
+ to_node: UUID,
+ relationship_name: str,
+ edge_properties: Optional[Dict[str, Any]] = {},
+ ):
serialized_properties = self.serialize_properties(edge_properties)
query = dedent("""MATCH (from_node {id: $from_node}),
@@ -186,12 +194,11 @@ async def add_edge(self, from_node: UUID, to_node: UUID, relationship_name: str,
"from_node": str(from_node),
"to_node": str(to_node),
"relationship_name": relationship_name,
- "properties": serialized_properties
+ "properties": serialized_properties,
}
return await self.query(query, params)
-
async def add_edges(self, edges: list[tuple[str, str, str, dict[str, Any]]]) -> None:
query = """
UNWIND $edges AS edge
@@ -201,22 +208,25 @@ async def add_edges(self, edges: list[tuple[str, str, str, dict[str, Any]]]) ->
RETURN rel
"""
- edges = [{
- "from_node": str(edge[0]),
- "to_node": str(edge[1]),
- "relationship_name": edge[2],
- "properties": {
- **(edge[3] if edge[3] else {}),
- "source_node_id": str(edge[0]),
- "target_node_id": str(edge[1]),
- },
- } for edge in edges]
+ edges = [
+ {
+ "from_node": str(edge[0]),
+ "to_node": str(edge[1]),
+ "relationship_name": edge[2],
+ "properties": {
+ **(edge[3] if edge[3] else {}),
+ "source_node_id": str(edge[0]),
+ "target_node_id": str(edge[1]),
+ },
+ }
+ for edge in edges
+ ]
try:
- results = await self.query(query, dict(edges = edges))
+ results = await self.query(query, dict(edges=edges))
return results
except Neo4jError as error:
- logger.error("Neo4j query error: %s", error, exc_info = True)
+ logger.error("Neo4j query error: %s", error, exc_info=True)
raise error
async def get_edges(self, node_id: str):
@@ -225,9 +235,12 @@ async def get_edges(self, node_id: str):
RETURN n, r, m
"""
- results = await self.query(query, dict(node_id = node_id))
+ results = await self.query(query, dict(node_id=node_id))
- return [(result["n"]["id"], result["m"]["id"], {"relationship_name": result["r"][1]}) for result in results]
+ return [
+ (result["n"]["id"], result["m"]["id"], {"relationship_name": result["r"][1]})
+ for result in results
+ ]
async def get_disconnected_nodes(self) -> list[str]:
# return await self.query(
@@ -267,7 +280,6 @@ async def get_disconnected_nodes(self) -> list[str]:
results = await self.query(query)
return results[0]["ids"] if len(results) > 0 else []
-
async def get_predecessors(self, node_id: str, edge_label: str = None) -> list[str]:
if edge_label is not None:
query = """
@@ -279,9 +291,9 @@ async def get_predecessors(self, node_id: str, edge_label: str = None) -> list[s
results = await self.query(
query,
dict(
- node_id = node_id,
- edge_label = edge_label,
- )
+ node_id=node_id,
+ edge_label=edge_label,
+ ),
)
return [result["predecessor"] for result in results]
@@ -295,8 +307,8 @@ async def get_predecessors(self, node_id: str, edge_label: str = None) -> list[s
results = await self.query(
query,
dict(
- node_id = node_id,
- )
+ node_id=node_id,
+ ),
)
return [result["predecessor"] for result in results]
@@ -312,8 +324,8 @@ async def get_successors(self, node_id: str, edge_label: str = None) -> list[str
results = await self.query(
query,
dict(
- node_id = node_id,
- edge_label = edge_label,
+ node_id=node_id,
+ edge_label=edge_label,
),
)
@@ -328,14 +340,16 @@ async def get_successors(self, node_id: str, edge_label: str = None) -> list[str
results = await self.query(
query,
dict(
- node_id = node_id,
- )
+ node_id=node_id,
+ ),
)
return [result["successor"] for result in results]
async def get_neighbours(self, node_id: str) -> List[Dict[str, Any]]:
- predecessors, successors = await asyncio.gather(self.get_predecessors(node_id), self.get_successors(node_id))
+ predecessors, successors = await asyncio.gather(
+ self.get_predecessors(node_id), self.get_successors(node_id)
+ )
return predecessors + successors
@@ -352,52 +366,55 @@ async def get_connections(self, node_id: UUID) -> list:
"""
predecessors, successors = await asyncio.gather(
- self.query(predecessors_query, dict(node_id = str(node_id))),
- self.query(successors_query, dict(node_id = str(node_id))),
+ self.query(predecessors_query, dict(node_id=str(node_id))),
+ self.query(successors_query, dict(node_id=str(node_id))),
)
connections = []
for neighbour in predecessors:
neighbour = neighbour["relation"]
- connections.append((neighbour[0], { "relationship_name": neighbour[1] }, neighbour[2]))
+ connections.append((neighbour[0], {"relationship_name": neighbour[1]}, neighbour[2]))
for neighbour in successors:
neighbour = neighbour["relation"]
- connections.append((neighbour[0], { "relationship_name": neighbour[1] }, neighbour[2]))
+ connections.append((neighbour[0], {"relationship_name": neighbour[1]}, neighbour[2]))
return connections
- async def remove_connection_to_predecessors_of(self, node_ids: list[str], edge_label: str) -> None:
+ async def remove_connection_to_predecessors_of(
+ self, node_ids: list[str], edge_label: str
+ ) -> None:
query = f"""
UNWIND $node_ids AS id
MATCH (node:`{id}`)-[r:{edge_label}]->(predecessor)
DELETE r;
"""
- params = { "node_ids": node_ids }
+ params = {"node_ids": node_ids}
return await self.query(query, params)
- async def remove_connection_to_successors_of(self, node_ids: list[str], edge_label: str) -> None:
+ async def remove_connection_to_successors_of(
+ self, node_ids: list[str], edge_label: str
+ ) -> None:
query = f"""
UNWIND $node_ids AS id
MATCH (node:`{id}`)<-[r:{edge_label}]-(successor)
DELETE r;
"""
- params = { "node_ids": node_ids }
+ params = {"node_ids": node_ids}
return await self.query(query, params)
-
async def delete_graph(self):
query = """MATCH (node)
DETACH DELETE node;"""
return await self.query(query)
- def serialize_properties(self, properties = dict()):
+ def serialize_properties(self, properties=dict()):
serialized_properties = {}
for property_key, property_value in properties.items():
@@ -414,22 +431,28 @@ async def get_graph_data(self):
result = await self.query(query)
- nodes = [(
- record["properties"]["id"],
- record["properties"],
- ) for record in result]
+ nodes = [
+ (
+ record["properties"]["id"],
+ record["properties"],
+ )
+ for record in result
+ ]
query = """
MATCH (n)-[r]->(m)
RETURN ID(n) AS source, ID(m) AS target, TYPE(r) AS type, properties(r) AS properties
"""
result = await self.query(query)
- edges = [(
- record["properties"]["source_node_id"],
- record["properties"]["target_node_id"],
- record["type"],
- record["properties"],
- ) for record in result]
+ edges = [
+ (
+ record["properties"]["source_node_id"],
+ record["properties"]["target_node_id"],
+ record["type"],
+ record["properties"],
+ )
+ for record in result
+ ]
return (nodes, edges)
@@ -446,7 +469,9 @@ async def get_filtered_graph_data(self, attribute_filters):
"""
where_clauses = []
for attribute, values in attribute_filters[0].items():
- values_str = ", ".join(f"'{value}'" if isinstance(value, str) else str(value) for value in values)
+ values_str = ", ".join(
+ f"'{value}'" if isinstance(value, str) else str(value) for value in values
+ )
where_clauses.append(f"n.{attribute} IN [{values_str}]")
where_clause = " AND ".join(where_clauses)
@@ -458,10 +483,13 @@ async def get_filtered_graph_data(self, attribute_filters):
"""
result_nodes = await self.query(query_nodes)
- nodes = [(
- record["id"],
- record["properties"],
- ) for record in result_nodes]
+ nodes = [
+ (
+ record["id"],
+ record["properties"],
+ )
+ for record in result_nodes
+ ]
query_edges = f"""
MATCH (n)-[r]->(m)
@@ -470,11 +498,14 @@ async def get_filtered_graph_data(self, attribute_filters):
"""
result_edges = await self.query(query_edges)
- edges = [(
- record["source"],
- record["target"],
- record["type"],
- record["properties"],
- ) for record in result_edges]
+ edges = [
+ (
+ record["source"],
+ record["target"],
+ record["type"],
+ record["properties"],
+ )
+ for record in result_edges
+ ]
- return (nodes, edges)
\ No newline at end of file
+ return (nodes, edges)
diff --git a/cognee/infrastructure/databases/graph/networkx/adapter.py b/cognee/infrastructure/databases/graph/networkx/adapter.py
index aadad639f..75969c798 100644
--- a/cognee/infrastructure/databases/graph/networkx/adapter.py
+++ b/cognee/infrastructure/databases/graph/networkx/adapter.py
@@ -17,9 +17,10 @@
logger = logging.getLogger("NetworkXAdapter")
+
class NetworkXAdapter(GraphDBInterface):
_instance = None
- graph = None # Class variable to store the singleton instance
+ graph = None # Class variable to store the singleton instance
def __new__(cls, filename):
if cls._instance is None:
@@ -27,12 +28,12 @@ def __new__(cls, filename):
cls._instance.filename = filename
return cls._instance
- def __init__(self, filename = "cognee_graph.pkl"):
+ def __init__(self, filename="cognee_graph.pkl"):
self.filename = filename
async def get_graph_data(self):
await self.load_graph_from_file()
- return (list(self.graph.nodes(data = True)), list(self.graph.edges(data = True, keys = True)))
+ return (list(self.graph.nodes(data=True)), list(self.graph.edges(data=True, keys=True)))
async def query(self, query: str, params: dict):
pass
@@ -57,24 +58,21 @@ async def add_nodes(
self.graph.add_nodes_from(nodes)
await self.save_graph_to_file(self.filename)
-
async def get_graph(self):
return self.graph
-
async def has_edge(self, from_node: str, to_node: str, edge_label: str) -> bool:
- return self.graph.has_edge(from_node, to_node, key = edge_label)
+ return self.graph.has_edge(from_node, to_node, key=edge_label)
async def has_edges(self, edges):
result = []
- for (from_node, to_node, edge_label) in edges:
+ for from_node, to_node, edge_label in edges:
if self.graph.has_edge(from_node, to_node, edge_label):
result.append((from_node, to_node, edge_label))
return result
-
async def add_edge(
self,
from_node: str,
@@ -83,24 +81,38 @@ async def add_edge(
edge_properties: Dict[str, Any] = {},
) -> None:
edge_properties["updated_at"] = datetime.now(timezone.utc)
- self.graph.add_edge(from_node, to_node, key = relationship_name, **(edge_properties if edge_properties else {}))
+ self.graph.add_edge(
+ from_node,
+ to_node,
+ key=relationship_name,
+ **(edge_properties if edge_properties else {}),
+ )
await self.save_graph_to_file(self.filename)
async def add_edges(
self,
edges: tuple[str, str, str, dict],
) -> None:
- edges = [(edge[0], edge[1], edge[2], {
- **(edge[3] if len(edge) == 4 else {}),
- "updated_at": datetime.now(timezone.utc),
- }) for edge in edges]
+ edges = [
+ (
+ edge[0],
+ edge[1],
+ edge[2],
+ {
+ **(edge[3] if len(edge) == 4 else {}),
+ "updated_at": datetime.now(timezone.utc),
+ },
+ )
+ for edge in edges
+ ]
self.graph.add_edges_from(edges)
await self.save_graph_to_file(self.filename)
async def get_edges(self, node_id: str):
- return list(self.graph.in_edges(node_id, data = True)) + list(self.graph.out_edges(node_id, data = True))
-
+ return list(self.graph.in_edges(node_id, data=True)) + list(
+ self.graph.out_edges(node_id, data=True)
+ )
async def delete_node(self, node_id: str) -> None:
"""Asynchronously delete a node from the graph if it exists."""
@@ -112,12 +124,11 @@ async def delete_nodes(self, node_ids: List[str]) -> None:
self.graph.remove_nodes_from(node_ids)
await self.save_graph_to_file(self.filename)
-
async def get_disconnected_nodes(self) -> List[str]:
connected_components = list(nx.weakly_connected_components(self.graph))
disconnected_nodes = []
- biggest_subgraph = max(connected_components, key = len)
+ biggest_subgraph = max(connected_components, key=len)
for component in connected_components:
if component != biggest_subgraph:
@@ -125,7 +136,6 @@ async def get_disconnected_nodes(self) -> List[str]:
return disconnected_nodes
-
async def extract_node(self, node_id: str) -> dict:
if self.graph.has_node(node_id):
return self.graph.nodes[node_id]
@@ -139,8 +149,8 @@ async def get_predecessors(self, node_id: UUID, edge_label: str = None) -> list:
if self.graph.has_node(node_id):
if edge_label is None:
return [
- self.graph.nodes[predecessor] for predecessor \
- in list(self.graph.predecessors(node_id))
+ self.graph.nodes[predecessor]
+ for predecessor in list(self.graph.predecessors(node_id))
]
nodes = []
@@ -155,8 +165,8 @@ async def get_successors(self, node_id: UUID, edge_label: str = None) -> list:
if self.graph.has_node(node_id):
if edge_label is None:
return [
- self.graph.nodes[successor] for successor \
- in list(self.graph.successors(node_id))
+ self.graph.nodes[successor]
+ for successor in list(self.graph.successors(node_id))
]
nodes = []
@@ -210,7 +220,9 @@ async def get_connections(self, node_id: UUID) -> list:
return connections
- async def remove_connection_to_predecessors_of(self, node_ids: list[str], edge_label: str) -> None:
+ async def remove_connection_to_predecessors_of(
+ self, node_ids: list[str], edge_label: str
+ ) -> None:
for node_id in node_ids:
if self.graph.has_node(node_id):
for predecessor_id in list(self.graph.predecessors(node_id)):
@@ -219,7 +231,9 @@ async def remove_connection_to_predecessors_of(self, node_ids: list[str], edge_l
await self.save_graph_to_file(self.filename)
- async def remove_connection_to_successors_of(self, node_ids: list[str], edge_label: str) -> None:
+ async def remove_connection_to_successors_of(
+ self, node_ids: list[str], edge_label: str
+ ) -> None:
for node_id in node_ids:
if self.graph.has_node(node_id):
for successor_id in list(self.graph.successors(node_id)):
@@ -228,7 +242,7 @@ async def remove_connection_to_successors_of(self, node_ids: list[str], edge_lab
await self.save_graph_to_file(self.filename)
- async def save_graph_to_file(self, file_path: str=None) -> None:
+ async def save_graph_to_file(self, file_path: str = None) -> None:
"""Asynchronously save the graph to a file in JSON format."""
if not file_path:
file_path = self.filename
@@ -236,8 +250,7 @@ async def save_graph_to_file(self, file_path: str=None) -> None:
graph_data = nx.readwrite.json_graph.node_link_data(self.graph)
async with aiofiles.open(file_path, "w") as file:
- await file.write(json.dumps(graph_data, cls = JSONEncoder))
-
+ await file.write(json.dumps(graph_data, cls=JSONEncoder))
async def load_graph_from_file(self, file_path: str = None):
"""Asynchronously load the graph from a file in JSON format."""
@@ -252,50 +265,59 @@ async def load_graph_from_file(self, file_path: str = None):
graph_data = json.loads(await file.read())
for node in graph_data["nodes"]:
try:
- node["id"] = UUID(node["id"])
- except:
- pass
+ node["id"] = UUID(node["id"])
+ except Exception as e:
+ print(e)
+ pass
if "updated_at" in node:
- node["updated_at"] = datetime.strptime(node["updated_at"], "%Y-%m-%dT%H:%M:%S.%f%z")
+ node["updated_at"] = datetime.strptime(
+ node["updated_at"], "%Y-%m-%dT%H:%M:%S.%f%z"
+ )
for edge in graph_data["links"]:
try:
- source_id = UUID(edge["source"])
- target_id = UUID(edge["target"])
+ source_id = UUID(edge["source"])
+ target_id = UUID(edge["target"])
- edge["source"] = source_id
- edge["target"] = target_id
- edge["source_node_id"] = source_id
- edge["target_node_id"] = target_id
- except:
- pass
+ edge["source"] = source_id
+ edge["target"] = target_id
+ edge["source_node_id"] = source_id
+ edge["target_node_id"] = target_id
+ except Exception as e:
+ print(e)
+ pass
if "updated_at" in edge:
- edge["updated_at"] = datetime.strptime(edge["updated_at"], "%Y-%m-%dT%H:%M:%S.%f%z")
+ edge["updated_at"] = datetime.strptime(
+ edge["updated_at"], "%Y-%m-%dT%H:%M:%S.%f%z"
+ )
self.graph = nx.readwrite.json_graph.node_link_graph(graph_data)
for node_id, node_data in self.graph.nodes(data=True):
- node_data['id'] = node_id
+ node_data["id"] = node_id
else:
# Log that the file does not exist and an empty graph is initialized
logger.warning("File %s not found. Initializing an empty graph.", file_path)
- self.graph = nx.MultiDiGraph() # Use MultiDiGraph to keep it consistent with __init__
+ self.graph = (
+ nx.MultiDiGraph()
+ ) # Use MultiDiGraph to keep it consistent with __init__
file_dir = os.path.dirname(file_path)
if not os.path.exists(file_dir):
- os.makedirs(file_dir, exist_ok = True)
+ os.makedirs(file_dir, exist_ok=True)
await self.save_graph_to_file(file_path)
except Exception:
logger.error("Failed to load graph from file: %s", file_path)
-
async def delete_graph(self, file_path: str = None):
"""Asynchronously delete the graph file from the filesystem."""
if file_path is None:
- file_path = self.filename # Assuming self.filename is defined elsewhere and holds the default graph file path
+ file_path = (
+ self.filename
+ ) # Assuming self.filename is defined elsewhere and holds the default graph file path
try:
if os.path.exists(file_path):
await aiofiles_os.remove(file_path)
@@ -305,7 +327,9 @@ async def delete_graph(self, file_path: str = None):
except Exception as error:
logger.error("Failed to delete graph: %s", error)
- async def get_filtered_graph_data(self, attribute_filters: List[Dict[str, List[Union[str, int]]]]):
+ async def get_filtered_graph_data(
+ self, attribute_filters: List[Dict[str, List[Union[str, int]]]]
+ ):
"""
Fetches nodes and relationships filtered by specified attribute values.
@@ -325,18 +349,21 @@ async def get_filtered_graph_data(self, attribute_filters: List[Dict[str, List[U
# Filter nodes
filtered_nodes = [
- (node, data) for node, data in self.graph.nodes(data=True)
+ (node, data)
+ for node, data in self.graph.nodes(data=True)
if all(data.get(attr) in values for attr, values in where_clauses)
]
# Filter edges where both source and target nodes satisfy the filters
filtered_edges = [
- (source, target, data.get('relationship_type', 'UNKNOWN'), data)
+ (source, target, data.get("relationship_type", "UNKNOWN"), data)
for source, target, data in self.graph.edges(data=True)
if (
- all(self.graph.nodes[source].get(attr) in values for attr, values in where_clauses) and
- all(self.graph.nodes[target].get(attr) in values for attr, values in where_clauses)
+ all(self.graph.nodes[source].get(attr) in values for attr, values in where_clauses)
+ and all(
+ self.graph.nodes[target].get(attr) in values for attr, values in where_clauses
+ )
)
]
- return filtered_nodes, filtered_edges
\ No newline at end of file
+ return filtered_nodes, filtered_edges
diff --git a/cognee/infrastructure/databases/hybrid/falkordb/FalkorDBAdapter.py b/cognee/infrastructure/databases/hybrid/falkordb/FalkorDBAdapter.py
index 324ee7bcd..77bfd74e6 100644
--- a/cognee/infrastructure/databases/hybrid/falkordb/FalkorDBAdapter.py
+++ b/cognee/infrastructure/databases/hybrid/falkordb/FalkorDBAdapter.py
@@ -1,38 +1,36 @@
import asyncio
+
# from datetime import datetime
import json
from textwrap import dedent
from uuid import UUID
+from webbrowser import Error
from falkordb import FalkorDB
from cognee.exceptions import InvalidValueError
-from cognee.infrastructure.databases.graph.graph_db_interface import \
- GraphDBInterface
+from cognee.infrastructure.databases.graph.graph_db_interface import GraphDBInterface
from cognee.infrastructure.databases.vector.embeddings import EmbeddingEngine
-from cognee.infrastructure.databases.vector.vector_db_interface import \
- VectorDBInterface
+from cognee.infrastructure.databases.vector.vector_db_interface import VectorDBInterface
from cognee.infrastructure.engine import DataPoint
class IndexSchema(DataPoint):
text: str
- _metadata: dict = {
- "index_fields": ["text"],
- "type": "IndexSchema"
- }
+ _metadata: dict = {"index_fields": ["text"], "type": "IndexSchema"}
+
class FalkorDBAdapter(VectorDBInterface, GraphDBInterface):
def __init__(
self,
database_url: str,
database_port: int,
- embedding_engine = EmbeddingEngine,
+ embedding_engine=EmbeddingEngine,
):
self.driver = FalkorDB(
- host = database_url,
- port = database_port,
+ host=database_url,
+ port=database_port,
)
self.embedding_engine = embedding_engine
self.graph_name = "cognee_graph"
@@ -56,7 +54,11 @@ def parse_value(value):
return f"'{str(value)}'"
if type(value) is int or type(value) is float:
return value
- if type(value) is list and type(value[0]) is float and len(value) == self.embedding_engine.get_vector_size():
+ if (
+ type(value) is list
+ and type(value[0]) is float
+ and len(value) == self.embedding_engine.get_vector_size()
+ ):
return f"'vecf32({value})'"
# if type(value) is datetime:
# return datetime.strptime(value, "%Y-%m-%dT%H:%M:%S.%f%z")
@@ -70,14 +72,21 @@ async def create_data_point_query(self, data_point: DataPoint, vectorized_values
node_label = type(data_point).__tablename__
property_names = DataPoint.get_embeddable_property_names(data_point)
- node_properties = await self.stringify_properties({
- **data_point.model_dump(),
- **({
- property_names[index]: (vectorized_values[index] \
- if index < len(vectorized_values) else getattr(data_point, property_name, None)) \
+ node_properties = await self.stringify_properties(
+ {
+ **data_point.model_dump(),
+ **(
+ {
+ property_names[index]: (
+ vectorized_values[index]
+ if index < len(vectorized_values)
+ else getattr(data_point, property_name, None)
+ )
for index, property_name in enumerate(property_names)
- }),
- })
+ }
+ ),
+ }
+ )
return dedent(f"""
MERGE (node:{node_label} {{id: '{str(data_point.id)}'}})
@@ -129,12 +138,13 @@ async def create_data_points(self, data_points: list[DataPoint]):
await self.create_data_point_query(
data_point,
[
- vectorized_values[vector_map[str(data_point.id)][property_name]] \
- if vector_map[str(data_point.id)][property_name] is not None \
- else None \
+ vectorized_values[vector_map[str(data_point.id)][property_name]]
+ if vector_map[str(data_point.id)][property_name] is not None
+ else None
for property_name in DataPoint.get_embeddable_property_names(data_point)
],
- ) for data_point in data_points
+ )
+ for data_point in data_points
]
for query in queries:
@@ -144,17 +154,27 @@ async def create_vector_index(self, index_name: str, index_property_name: str):
graph = self.driver.select_graph(self.graph_name)
if not self.has_vector_index(graph, index_name, index_property_name):
- graph.create_node_vector_index(index_name, index_property_name, dim = self.embedding_engine.get_vector_size())
+ graph.create_node_vector_index(
+ index_name, index_property_name, dim=self.embedding_engine.get_vector_size()
+ )
def has_vector_index(self, graph, index_name: str, index_property_name: str) -> bool:
try:
indices = graph.list_indices()
- return any([(index[0] == index_name and index_property_name in index[1]) for index in indices.result_set])
- except:
+ return any(
+ [
+ (index[0] == index_name and index_property_name in index[1])
+ for index in indices.result_set
+ ]
+ )
+ except Error as e:
+ print(e)
return False
- async def index_data_points(self, index_name: str, index_property_name: str, data_points: list[DataPoint]):
+ async def index_data_points(
+ self, index_name: str, index_property_name: str, data_points: list[DataPoint]
+ ):
pass
async def add_node(self, node: DataPoint):
@@ -183,11 +203,14 @@ async def has_edges(self, edges):
""").strip()
params = {
- "edges": [{
- "from_node": str(edge[0]),
- "to_node": str(edge[1]),
- "relationship_name": edge[2],
- } for edge in edges],
+ "edges": [
+ {
+ "from_node": str(edge[0]),
+ "to_node": str(edge[1]),
+ "relationship_name": edge[2],
+ }
+ for edge in edges
+ ],
}
results = self.query(query, params).result_set
@@ -196,7 +219,7 @@ async def has_edges(self, edges):
async def retrieve(self, data_point_ids: list[UUID]):
result = self.query(
- f"MATCH (node) WHERE node.id IN $node_ids RETURN node",
+ "MATCH (node) WHERE node.id IN $node_ids RETURN node",
{
"node_ids": [str(data_point) for data_point in data_point_ids],
},
@@ -224,19 +247,19 @@ async def get_connections(self, node_id: UUID) -> list:
"""
predecessors, successors = await asyncio.gather(
- self.query(predecessors_query, dict(node_id = node_id)),
- self.query(successors_query, dict(node_id = node_id)),
+ self.query(predecessors_query, dict(node_id=node_id)),
+ self.query(successors_query, dict(node_id=node_id)),
)
connections = []
for neighbour in predecessors:
neighbour = neighbour["relation"]
- connections.append((neighbour[0], { "relationship_name": neighbour[1] }, neighbour[2]))
+ connections.append((neighbour[0], {"relationship_name": neighbour[1]}, neighbour[2]))
for neighbour in successors:
neighbour = neighbour["relation"]
- connections.append((neighbour[0], { "relationship_name": neighbour[1] }, neighbour[2]))
+ connections.append((neighbour[0], {"relationship_name": neighbour[1]}, neighbour[2]))
return connections
@@ -279,12 +302,15 @@ async def batch_search(
query_vectors = await self.embedding_engine.embed_text(query_texts)
return await asyncio.gather(
- *[self.search(
- collection_name = collection_name,
- query_vector = query_vector,
- limit = limit,
- with_vector = with_vectors,
- ) for query_vector in query_vectors]
+ *[
+ self.search(
+ collection_name=collection_name,
+ query_vector=query_vector,
+ limit=limit,
+ with_vector=with_vectors,
+ )
+ for query_vector in query_vectors
+ ]
)
async def get_graph_data(self):
@@ -292,28 +318,34 @@ async def get_graph_data(self):
result = self.query(query)
- nodes = [(
- record[2]["id"],
- record[2],
- ) for record in result.result_set]
+ nodes = [
+ (
+ record[2]["id"],
+ record[2],
+ )
+ for record in result.result_set
+ ]
query = """
MATCH (n)-[r]->(m)
RETURN ID(n) AS source, ID(m) AS target, TYPE(r) AS type, properties(r) AS properties
"""
result = self.query(query)
- edges = [(
- record[3]["source_node_id"],
- record[3]["target_node_id"],
- record[2],
- record[3],
- ) for record in result.result_set]
+ edges = [
+ (
+ record[3]["source_node_id"],
+ record[3]["target_node_id"],
+ record[2],
+ record[3],
+ )
+ for record in result.result_set
+ ]
return (nodes, edges)
async def delete_data_points(self, collection_name: str, data_point_ids: list[UUID]):
return self.query(
- f"MATCH (node) WHERE node.id IN $node_ids DETACH DELETE node",
+ "MATCH (node) WHERE node.id IN $node_ids DETACH DELETE node",
{
"node_ids": [str(data_point) for data_point in data_point_ids],
},
diff --git a/cognee/infrastructure/databases/relational/ModelBase.py b/cognee/infrastructure/databases/relational/ModelBase.py
index f75ec448b..fa2b68a5d 100644
--- a/cognee/infrastructure/databases/relational/ModelBase.py
+++ b/cognee/infrastructure/databases/relational/ModelBase.py
@@ -1,4 +1,5 @@
from sqlalchemy.orm import DeclarativeBase
+
class Base(DeclarativeBase):
pass
diff --git a/cognee/infrastructure/databases/relational/config.py b/cognee/infrastructure/databases/relational/config.py
index ef8dca608..5ca121593 100644
--- a/cognee/infrastructure/databases/relational/config.py
+++ b/cognee/infrastructure/databases/relational/config.py
@@ -4,16 +4,17 @@
from pydantic_settings import BaseSettings, SettingsConfigDict
from cognee.root_dir import get_absolute_path
+
class RelationalConfig(BaseSettings):
- db_path: str = os.path.join(get_absolute_path(".cognee_system"), "databases")
- db_name: str = "cognee_db"
- db_host: Union[str, None] = None # "localhost"
- db_port: Union[str, None] = None # "5432"
- db_username: Union[str, None] = None # "cognee"
- db_password: Union[str, None] = None # "cognee"
+ db_path: str = os.path.join(get_absolute_path(".cognee_system"), "databases")
+ db_name: str = "cognee_db"
+ db_host: Union[str, None] = None # "localhost"
+ db_port: Union[str, None] = None # "5432"
+ db_username: Union[str, None] = None # "cognee"
+ db_password: Union[str, None] = None # "cognee"
db_provider: str = "sqlite"
- model_config = SettingsConfigDict(env_file = ".env", extra = "allow")
+ model_config = SettingsConfigDict(env_file=".env", extra="allow")
def to_dict(self) -> dict:
return {
@@ -26,6 +27,7 @@ def to_dict(self) -> dict:
"db_provider": self.db_provider,
}
+
@lru_cache
def get_relational_config():
return RelationalConfig()
diff --git a/cognee/infrastructure/databases/relational/create_db_and_tables.py b/cognee/infrastructure/databases/relational/create_db_and_tables.py
index 02ff680df..4947680f5 100644
--- a/cognee/infrastructure/databases/relational/create_db_and_tables.py
+++ b/cognee/infrastructure/databases/relational/create_db_and_tables.py
@@ -2,6 +2,7 @@
from .ModelBase import Base
from .get_relational_engine import get_relational_engine, get_relational_config
+
async def create_db_and_tables():
relational_config = get_relational_config()
relational_engine = get_relational_engine()
diff --git a/cognee/infrastructure/databases/relational/create_relational_engine.py b/cognee/infrastructure/databases/relational/create_relational_engine.py
index 08cf171d5..13a1edc23 100644
--- a/cognee/infrastructure/databases/relational/create_relational_engine.py
+++ b/cognee/infrastructure/databases/relational/create_relational_engine.py
@@ -1,5 +1,6 @@
from .sqlalchemy.SqlAlchemyAdapter import SQLAlchemyAdapter
+
def create_relational_engine(
db_path: str,
db_name: str,
@@ -13,6 +14,8 @@ def create_relational_engine(
connection_string = f"sqlite+aiosqlite:///{db_path}/{db_name}"
if db_provider == "postgres":
- connection_string = f"postgresql+asyncpg://{db_username}:{db_password}@{db_host}:{db_port}/{db_name}"
+ connection_string = (
+ f"postgresql+asyncpg://{db_username}:{db_password}@{db_host}:{db_port}/{db_name}"
+ )
return SQLAlchemyAdapter(connection_string)
diff --git a/cognee/infrastructure/databases/relational/get_relational_engine.py b/cognee/infrastructure/databases/relational/get_relational_engine.py
index d035f5baf..44aa7213b 100644
--- a/cognee/infrastructure/databases/relational/get_relational_engine.py
+++ b/cognee/infrastructure/databases/relational/get_relational_engine.py
@@ -3,8 +3,9 @@
from .config import get_relational_config
from .create_relational_engine import create_relational_engine
+
# @lru_cache
def get_relational_engine():
relational_config = get_relational_config()
- return create_relational_engine(**relational_config.to_dict())
\ No newline at end of file
+ return create_relational_engine(**relational_config.to_dict())
diff --git a/cognee/infrastructure/databases/relational/sqlalchemy/SqlAlchemyAdapter.py b/cognee/infrastructure/databases/relational/sqlalchemy/SqlAlchemyAdapter.py
index c197efc72..6c3c5029d 100644
--- a/cognee/infrastructure/databases/relational/sqlalchemy/SqlAlchemyAdapter.py
+++ b/cognee/infrastructure/databases/relational/sqlalchemy/SqlAlchemyAdapter.py
@@ -1,6 +1,6 @@
import os
from os import path
-import logging
+import logging
from uuid import UUID
from typing import Optional
from typing import AsyncGenerator, List
@@ -18,7 +18,8 @@
logger = logging.getLogger(__name__)
-class SQLAlchemyAdapter():
+
+class SQLAlchemyAdapter:
def __init__(self, connection_string: str):
self.db_path: str = None
self.db_uri: str = connection_string
@@ -58,17 +59,23 @@ async def create_table(self, schema_name: str, table_name: str, table_config: li
fields_query_parts = [f"{item['name']} {item['type']}" for item in table_config]
async with self.engine.begin() as connection:
await connection.execute(text(f"CREATE SCHEMA IF NOT EXISTS {schema_name};"))
- await connection.execute(text(f"CREATE TABLE IF NOT EXISTS {schema_name}.{table_name} ({', '.join(fields_query_parts)});"))
+ await connection.execute(
+ text(
+ f"CREATE TABLE IF NOT EXISTS {schema_name}.{table_name} ({', '.join(fields_query_parts)});"
+ )
+ )
await connection.close()
- async def delete_table(self, table_name: str, schema_name: Optional[str] = "public"):
+ async def delete_table(self, table_name: str, schema_name: Optional[str] = "public"):
async with self.engine.begin() as connection:
if self.engine.dialect.name == "sqlite":
# SQLite doesn’t support schema namespaces and the CASCADE keyword.
# However, foreign key constraint can be defined with ON DELETE CASCADE during table creation.
await connection.execute(text(f"DROP TABLE IF EXISTS {table_name};"))
else:
- await connection.execute(text(f"DROP TABLE IF EXISTS {schema_name}.{table_name} CASCADE;"))
+ await connection.execute(
+ text(f"DROP TABLE IF EXISTS {schema_name}.{table_name} CASCADE;")
+ )
async def insert_data(self, schema_name: str, table_name: str, data: list[dict]):
columns = ", ".join(data[0].keys())
@@ -94,7 +101,9 @@ async def get_schema_list(self) -> List[str]:
return [schema[0] for schema in result.fetchall()]
return []
- async def delete_entity_by_id(self, table_name: str, data_id: UUID, schema_name: Optional[str] = "public"):
+ async def delete_entity_by_id(
+ self, table_name: str, data_id: UUID, schema_name: Optional[str] = "public"
+ ):
"""
Delete entity in given table based on id. Table must have an id Column.
"""
@@ -114,7 +123,6 @@ async def delete_entity_by_id(self, table_name: str, data_id: UUID, schema_name:
await session.execute(TableModel.delete().where(TableModel.c.id == data_id))
await session.commit()
-
async def delete_data_entity(self, data_id: UUID):
"""
Delete data and local files related to data if there are no references to it anymore.
@@ -131,14 +139,19 @@ async def delete_data_entity(self, data_id: UUID):
raise EntityNotFoundError(message=f"Entity not found: {str(e)}")
# Check if other data objects point to the same raw data location
- raw_data_location_entities = (await session.execute(
- select(Data.raw_data_location).where(Data.raw_data_location == data_entity.raw_data_location))).all()
+ raw_data_location_entities = (
+ await session.execute(
+ select(Data.raw_data_location).where(
+ Data.raw_data_location == data_entity.raw_data_location
+ )
+ )
+ ).all()
# Don't delete local file unless this is the only reference to the file in the database
if len(raw_data_location_entities) == 1:
-
# delete local file only if it's created by cognee
from cognee.base_config import get_base_config
+
config = get_base_config()
if config.data_root_directory in raw_data_location_entities[0].raw_data_location:
@@ -198,15 +211,18 @@ async def get_table_names(self) -> List[str]:
metadata.clear()
return table_names
-
async def get_data(self, table_name: str, filters: dict = None):
async with self.engine.begin() as connection:
query = f"SELECT * FROM {table_name}"
if filters:
- filter_conditions = " AND ".join([
- f"{key} IN ({', '.join([f':{key}{i}' for i in range(len(value))])})" if isinstance(value, list)
- else f"{key} = :{key}" for key, value in filters.items()
- ])
+ filter_conditions = " AND ".join(
+ [
+ f"{key} IN ({', '.join([f':{key}{i}' for i in range(len(value))])})"
+ if isinstance(value, list)
+ else f"{key} = :{key}"
+ for key, value in filters.items()
+ ]
+ )
query += f" WHERE {filter_conditions};"
query = text(query)
results = await connection.execute(query, filters)
@@ -252,7 +268,6 @@ async def drop_tables(self):
except Exception as e:
print(f"Error dropping database tables: {e}")
-
async def create_database(self):
if self.engine.dialect.name == "sqlite":
from cognee.infrastructure.files.storage import LocalStorage
@@ -264,7 +279,6 @@ async def create_database(self):
if len(Base.metadata.tables.keys()) > 0:
await connection.run_sync(Base.metadata.create_all)
-
async def delete_database(self):
try:
if self.engine.dialect.name == "sqlite":
@@ -281,7 +295,9 @@ async def delete_database(self):
# Load the schema information into the MetaData object
await connection.run_sync(metadata.reflect, schema=schema_name)
for table in metadata.sorted_tables:
- drop_table_query = text(f"DROP TABLE IF EXISTS {schema_name}.{table.name} CASCADE")
+ drop_table_query = text(
+ f"DROP TABLE IF EXISTS {schema_name}.{table.name} CASCADE"
+ )
await connection.execute(drop_table_query)
metadata.clear()
except Exception as e:
diff --git a/cognee/infrastructure/databases/vector/config.py b/cognee/infrastructure/databases/vector/config.py
index 846bc5842..f2c180446 100644
--- a/cognee/infrastructure/databases/vector/config.py
+++ b/cognee/infrastructure/databases/vector/config.py
@@ -3,16 +3,16 @@
from pydantic_settings import BaseSettings, SettingsConfigDict
from cognee.root_dir import get_absolute_path
+
class VectorConfig(BaseSettings):
vector_db_url: str = os.path.join(
- os.path.join(get_absolute_path(".cognee_system"), "databases"),
- "cognee.lancedb"
+ os.path.join(get_absolute_path(".cognee_system"), "databases"), "cognee.lancedb"
)
vector_db_port: int = 1234
vector_db_key: str = ""
vector_db_provider: str = "lancedb"
- model_config = SettingsConfigDict(env_file = ".env", extra = "allow")
+ model_config = SettingsConfigDict(env_file=".env", extra="allow")
def to_dict(self) -> dict:
return {
@@ -22,6 +22,7 @@ def to_dict(self) -> dict:
"vector_db_provider": self.vector_db_provider,
}
+
@lru_cache
def get_vectordb_config():
return VectorConfig()
diff --git a/cognee/infrastructure/databases/vector/create_vector_engine.py b/cognee/infrastructure/databases/vector/create_vector_engine.py
index 32db48930..e61c272e1 100644
--- a/cognee/infrastructure/databases/vector/create_vector_engine.py
+++ b/cognee/infrastructure/databases/vector/create_vector_engine.py
@@ -16,9 +16,7 @@ def create_vector_engine(config: VectorConfig, embedding_engine):
raise EnvironmentError("Missing requred Weaviate credentials!")
return WeaviateAdapter(
- config["vector_db_url"],
- config["vector_db_key"],
- embedding_engine=embedding_engine
+ config["vector_db_url"], config["vector_db_key"], embedding_engine=embedding_engine
)
elif config["vector_db_provider"] == "qdrant":
@@ -30,10 +28,10 @@ def create_vector_engine(config: VectorConfig, embedding_engine):
return QDrantAdapter(
url=config["vector_db_url"],
api_key=config["vector_db_key"],
- embedding_engine=embedding_engine
+ embedding_engine=embedding_engine,
)
- elif config['vector_db_provider'] == 'milvus':
+ elif config["vector_db_provider"] == "milvus":
from .milvus.MilvusAdapter import MilvusAdapter
if not config["vector_db_url"]:
@@ -41,11 +39,10 @@ def create_vector_engine(config: VectorConfig, embedding_engine):
return MilvusAdapter(
url=config["vector_db_url"],
- api_key=config['vector_db_key'],
- embedding_engine=embedding_engine
+ api_key=config["vector_db_key"],
+ embedding_engine=embedding_engine,
)
-
elif config["vector_db_provider"] == "pgvector":
from cognee.infrastructure.databases.relational import get_relational_config
diff --git a/cognee/infrastructure/databases/vector/embeddings/EmbeddingEngine.py b/cognee/infrastructure/databases/vector/embeddings/EmbeddingEngine.py
index bcb99fb07..ea07ff4e2 100644
--- a/cognee/infrastructure/databases/vector/embeddings/EmbeddingEngine.py
+++ b/cognee/infrastructure/databases/vector/embeddings/EmbeddingEngine.py
@@ -1,5 +1,6 @@
from typing import Protocol
+
class EmbeddingEngine(Protocol):
async def embed_text(self, text: list[str]) -> list[list[float]]:
raise NotImplementedError()
diff --git a/cognee/infrastructure/databases/vector/embeddings/LiteLLMEmbeddingEngine.py b/cognee/infrastructure/databases/vector/embeddings/LiteLLMEmbeddingEngine.py
index 93f59cc77..27061ce45 100644
--- a/cognee/infrastructure/databases/vector/embeddings/LiteLLMEmbeddingEngine.py
+++ b/cognee/infrastructure/databases/vector/embeddings/LiteLLMEmbeddingEngine.py
@@ -43,14 +43,12 @@ def __init__(
async def embed_text(self, text: List[str]) -> List[List[float]]:
async def exponential_backoff(attempt):
- wait_time = min(10 * (2 ** attempt), 60) # Max 60 seconds
+ wait_time = min(10 * (2**attempt), 60) # Max 60 seconds
await asyncio.sleep(wait_time)
try:
if self.mock:
- response = {
- "data": [{"embedding": [0.0] * self.dimensions} for _ in text]
- }
+ response = {"data": [{"embedding": [0.0] * self.dimensions} for _ in text]}
self.retry_count = 0
@@ -61,7 +59,7 @@ async def exponential_backoff(attempt):
input=text,
api_key=self.api_key,
api_base=self.endpoint,
- api_version=self.api_version
+ api_version=self.api_version,
)
self.retry_count = 0
@@ -73,7 +71,7 @@ async def exponential_backoff(attempt):
if len(text) == 1:
parts = [text]
else:
- parts = [text[0:math.ceil(len(text) / 2)], text[math.ceil(len(text) / 2):]]
+ parts = [text[0 : math.ceil(len(text) / 2)], text[math.ceil(len(text) / 2) :]]
parts_futures = [self.embed_text(part) for part in parts]
embeddings = await asyncio.gather(*parts_futures)
@@ -89,7 +87,7 @@ async def exponential_backoff(attempt):
except litellm.exceptions.RateLimitError:
if self.retry_count >= self.MAX_RETRIES:
- raise Exception(f"Rate limit exceeded and no more retries left.")
+ raise Exception("Rate limit exceeded and no more retries left.")
await exponential_backoff(self.retry_count)
diff --git a/cognee/infrastructure/databases/vector/embeddings/config.py b/cognee/infrastructure/databases/vector/embeddings/config.py
index ecfb37204..042c063f8 100644
--- a/cognee/infrastructure/databases/vector/embeddings/config.py
+++ b/cognee/infrastructure/databases/vector/embeddings/config.py
@@ -2,6 +2,7 @@
from functools import lru_cache
from pydantic_settings import BaseSettings, SettingsConfigDict
+
class EmbeddingConfig(BaseSettings):
embedding_model: Optional[str] = "text-embedding-3-large"
embedding_dimensions: Optional[int] = 3072
@@ -9,7 +10,8 @@ class EmbeddingConfig(BaseSettings):
embedding_api_key: Optional[str] = None
embedding_api_version: Optional[str] = None
- model_config = SettingsConfigDict(env_file = ".env", extra = "allow")
+ model_config = SettingsConfigDict(env_file=".env", extra="allow")
+
@lru_cache
def get_embedding_config():
diff --git a/cognee/infrastructure/databases/vector/embeddings/get_embedding_engine.py b/cognee/infrastructure/databases/vector/embeddings/get_embedding_engine.py
index d2582fbf0..6bfb4dd15 100644
--- a/cognee/infrastructure/databases/vector/embeddings/get_embedding_engine.py
+++ b/cognee/infrastructure/databases/vector/embeddings/get_embedding_engine.py
@@ -3,15 +3,16 @@
from .EmbeddingEngine import EmbeddingEngine
from .LiteLLMEmbeddingEngine import LiteLLMEmbeddingEngine
+
def get_embedding_engine() -> EmbeddingEngine:
config = get_embedding_config()
llm_config = get_llm_config()
return LiteLLMEmbeddingEngine(
# If OpenAI API is used for embeddings, litellm needs only the api_key.
- api_key = config.embedding_api_key or llm_config.llm_api_key,
- endpoint = config.embedding_endpoint,
- api_version = config.embedding_api_version,
- model = config.embedding_model,
- dimensions = config.embedding_dimensions,
+ api_key=config.embedding_api_key or llm_config.llm_api_key,
+ endpoint=config.embedding_endpoint,
+ api_version=config.embedding_api_version,
+ model=config.embedding_model,
+ dimensions=config.embedding_dimensions,
)
diff --git a/cognee/infrastructure/databases/vector/get_vector_engine.py b/cognee/infrastructure/databases/vector/get_vector_engine.py
index 079a8903f..4a3e81d1e 100644
--- a/cognee/infrastructure/databases/vector/get_vector_engine.py
+++ b/cognee/infrastructure/databases/vector/get_vector_engine.py
@@ -3,6 +3,7 @@
from .create_vector_engine import create_vector_engine
from functools import lru_cache
+
@lru_cache
def get_vector_engine():
return create_vector_engine(get_vectordb_config().to_dict(), get_embedding_engine())
diff --git a/cognee/infrastructure/databases/vector/lancedb/LanceDBAdapter.py b/cognee/infrastructure/databases/vector/lancedb/LanceDBAdapter.py
index 1b3fc55c3..2caa8be1e 100644
--- a/cognee/infrastructure/databases/vector/lancedb/LanceDBAdapter.py
+++ b/cognee/infrastructure/databases/vector/lancedb/LanceDBAdapter.py
@@ -21,10 +21,8 @@ class IndexSchema(DataPoint):
id: str
text: str
- _metadata: dict = {
- "index_fields": ["text"],
- "type": "IndexSchema"
- }
+ _metadata: dict = {"index_fields": ["text"], "type": "IndexSchema"}
+
class LanceDBAdapter(VectorDBInterface):
name = "LanceDB"
@@ -32,7 +30,6 @@ class LanceDBAdapter(VectorDBInterface):
api_key: str
connection: lancedb.AsyncConnection = None
-
def __init__(
self,
url: Optional[str],
@@ -45,7 +42,7 @@ def __init__(
async def get_connection(self):
if self.connection is None:
- self.connection = await lancedb.connect_async(self.url, api_key = self.api_key)
+ self.connection = await lancedb.connect_async(self.url, api_key=self.api_key)
return self.connection
@@ -71,9 +68,9 @@ class LanceDataPoint(LanceModel):
if not await self.has_collection(collection_name):
connection = await self.get_connection()
return await connection.create_table(
- name = collection_name,
- schema = LanceDataPoint,
- exist_ok = True,
+ name=collection_name,
+ schema=LanceDataPoint,
+ exist_ok=True,
)
async def create_data_points(self, collection_name: str, data_points: list[DataPoint]):
@@ -108,21 +105,22 @@ def create_lance_data_point(data_point: DataPoint, vector: list[float]) -> Lance
properties["id"] = str(properties["id"])
return LanceDataPoint[str, self.get_data_point_schema(type(data_point))](
- id = str(data_point.id),
- vector = vector,
- payload = properties,
+ id=str(data_point.id),
+ vector=vector,
+ payload=properties,
)
lance_data_points = [
create_lance_data_point(data_point, data_vectors[data_point_index])
- for (data_point_index, data_point) in enumerate(data_points)
+ for (data_point_index, data_point) in enumerate(data_points)
]
- await collection.merge_insert("id") \
- .when_matched_update_all() \
- .when_not_matched_insert_all() \
+ await (
+ collection.merge_insert("id")
+ .when_matched_update_all()
+ .when_not_matched_insert_all()
.execute(lance_data_points)
-
+ )
async def retrieve(self, collection_name: str, data_point_ids: list[str]):
connection = await self.get_connection()
@@ -133,17 +131,17 @@ async def retrieve(self, collection_name: str, data_point_ids: list[str]):
else:
results = await collection.query().where(f"id IN {tuple(data_point_ids)}").to_pandas()
- return [ScoredResult(
- id = UUID(result["id"]),
- payload = result["payload"],
- score = 0,
- ) for result in results.to_dict("index").values()]
+ return [
+ ScoredResult(
+ id=UUID(result["id"]),
+ payload=result["payload"],
+ score=0,
+ )
+ for result in results.to_dict("index").values()
+ ]
async def get_distance_from_collection_elements(
- self,
- collection_name: str,
- query_text: str = None,
- query_vector: List[float] = None
+ self, collection_name: str, query_text: str = None, query_vector: List[float] = None
):
if query_text is None and query_vector is None:
raise InvalidValueError(message="One of query_text or query_vector must be provided!")
@@ -160,11 +158,14 @@ async def get_distance_from_collection_elements(
normalized_values = normalize_distances(result_values)
- return [ScoredResult(
- id=UUID(result["id"]),
- payload=result["payload"],
- score=normalized_values[value_index],
- ) for value_index, result in enumerate(result_values)]
+ return [
+ ScoredResult(
+ id=UUID(result["id"]),
+ payload=result["payload"],
+ score=normalized_values[value_index],
+ )
+ for value_index, result in enumerate(result_values)
+ ]
async def search(
self,
@@ -173,7 +174,7 @@ async def search(
query_vector: List[float] = None,
limit: int = 5,
with_vector: bool = False,
- normalized: bool = True
+ normalized: bool = True,
):
if query_text is None and query_vector is None:
raise InvalidValueError(message="One of query_text or query_vector must be provided!")
@@ -190,11 +191,14 @@ async def search(
normalized_values = normalize_distances(result_values)
- return [ScoredResult(
- id = UUID(result["id"]),
- payload = result["payload"],
- score = normalized_values[value_index],
- ) for value_index, result in enumerate(result_values)]
+ return [
+ ScoredResult(
+ id=UUID(result["id"]),
+ payload=result["payload"],
+ score=normalized_values[value_index],
+ )
+ for value_index, result in enumerate(result_values)
+ ]
async def batch_search(
self,
@@ -206,12 +210,15 @@ async def batch_search(
query_vectors = await self.embedding_engine.embed_text(query_texts)
return await asyncio.gather(
- *[self.search(
- collection_name = collection_name,
- query_vector = query_vector,
- limit = limit,
- with_vector = with_vectors,
- ) for query_vector in query_vectors]
+ *[
+ self.search(
+ collection_name=collection_name,
+ query_vector=query_vector,
+ limit=limit,
+ with_vector=with_vectors,
+ )
+ for query_vector in query_vectors
+ ]
)
async def delete_data_points(self, collection_name: str, data_point_ids: list[str]):
@@ -224,26 +231,34 @@ async def delete_data_points(self, collection_name: str, data_point_ids: list[st
return results
async def create_vector_index(self, index_name: str, index_property_name: str):
- await self.create_collection(f"{index_name}_{index_property_name}", payload_schema = IndexSchema)
+ await self.create_collection(
+ f"{index_name}_{index_property_name}", payload_schema=IndexSchema
+ )
- async def index_data_points(self, index_name: str, index_property_name: str, data_points: list[DataPoint]):
- await self.create_data_points(f"{index_name}_{index_property_name}", [
- IndexSchema(
- id = str(data_point.id),
- text = getattr(data_point, data_point._metadata["index_fields"][0]),
- ) for data_point in data_points
- ])
+ async def index_data_points(
+ self, index_name: str, index_property_name: str, data_points: list[DataPoint]
+ ):
+ await self.create_data_points(
+ f"{index_name}_{index_property_name}",
+ [
+ IndexSchema(
+ id=str(data_point.id),
+ text=getattr(data_point, data_point._metadata["index_fields"][0]),
+ )
+ for data_point in data_points
+ ],
+ )
async def prune(self):
# Clean up the database if it was set up as temporary
if self.url.startswith("/"):
- LocalStorage.remove_all(self.url) # Remove the temporary directory and files inside
+ LocalStorage.remove_all(self.url) # Remove the temporary directory and files inside
def get_data_point_schema(self, model_type):
return copy_model(
model_type,
- include_fields = {
+ include_fields={
"id": (str, ...),
},
- exclude_fields = ["_metadata"],
- )
\ No newline at end of file
+ exclude_fields=["_metadata"],
+ )
diff --git a/cognee/infrastructure/databases/vector/milvus/MilvusAdapter.py b/cognee/infrastructure/databases/vector/milvus/MilvusAdapter.py
index 0d4ea05d3..aecd4650d 100644
--- a/cognee/infrastructure/databases/vector/milvus/MilvusAdapter.py
+++ b/cognee/infrastructure/databases/vector/milvus/MilvusAdapter.py
@@ -17,10 +17,7 @@
class IndexSchema(DataPoint):
text: str
- _metadata: dict = {
- "index_fields": ["text"],
- "type": "IndexSchema"
- }
+ _metadata: dict = {"index_fields": ["text"], "type": "IndexSchema"}
class MilvusAdapter(VectorDBInterface):
@@ -35,8 +32,9 @@ def __init__(self, url: str, api_key: Optional[str], embedding_engine: Embedding
self.embedding_engine = embedding_engine
- def get_milvus_client(self) -> "MilvusClient":
+ def get_milvus_client(self):
from pymilvus import MilvusClient
+
if self.api_key:
client = MilvusClient(uri=self.url, token=self.api_key)
else:
@@ -54,11 +52,12 @@ async def has_collection(self, collection_name: str) -> bool:
return await future
async def create_collection(
- self,
- collection_name: str,
- payload_schema=None,
+ self,
+ collection_name: str,
+ payload_schema=None,
):
from pymilvus import DataType, MilvusException
+
client = self.get_milvus_client()
if client.has_collection(collection_name=collection_name):
logger.info(f"Collection '{collection_name}' already exists.")
@@ -74,34 +73,18 @@ async def create_collection(
)
schema.add_field(
- field_name="id",
- datatype=DataType.VARCHAR,
- is_primary=True,
- max_length=36
+ field_name="id", datatype=DataType.VARCHAR, is_primary=True, max_length=36
)
- schema.add_field(
- field_name="vector",
- datatype=DataType.FLOAT_VECTOR,
- dim=dimension
- )
+ schema.add_field(field_name="vector", datatype=DataType.FLOAT_VECTOR, dim=dimension)
- schema.add_field(
- field_name="text",
- datatype=DataType.VARCHAR,
- max_length=60535
- )
+ schema.add_field(field_name="text", datatype=DataType.VARCHAR, max_length=60535)
index_params = client.prepare_index_params()
- index_params.add_index(
- field_name="vector",
- metric_type="COSINE"
- )
+ index_params.add_index(field_name="vector", metric_type="COSINE")
client.create_collection(
- collection_name=collection_name,
- schema=schema,
- index_params=index_params
+ collection_name=collection_name, schema=schema, index_params=index_params
)
client.load_collection(collection_name)
@@ -112,12 +95,9 @@ async def create_collection(
logger.error(f"Error creating collection '{collection_name}': {str(e)}")
raise e
- async def create_data_points(
- self,
- collection_name: str,
- data_points: List[DataPoint]
- ):
+ async def create_data_points(self, collection_name: str, data_points: List[DataPoint]):
from pymilvus import MilvusException
+
client = self.get_milvus_client()
data_vectors = await self.embed_data(
[data_point.get_embeddable_data(data_point) for data_point in data_points]
@@ -133,22 +113,23 @@ async def create_data_points(
]
try:
- result = client.insert(
- collection_name=collection_name,
- data=insert_data
- )
+ result = client.insert(collection_name=collection_name, data=insert_data)
logger.info(
f"Inserted {result.get('insert_count', 0)} data points into collection '{collection_name}'."
)
return result
except MilvusException as e:
- logger.error(f"Error inserting data points into collection '{collection_name}': {str(e)}")
+ logger.error(
+ f"Error inserting data points into collection '{collection_name}': {str(e)}"
+ )
raise e
async def create_vector_index(self, index_name: str, index_property_name: str):
await self.create_collection(f"{index_name}_{index_property_name}")
- async def index_data_points(self, index_name: str, index_property_name: str, data_points: List[DataPoint]):
+ async def index_data_points(
+ self, index_name: str, index_property_name: str, data_points: List[DataPoint]
+ ):
formatted_data_points = [
IndexSchema(
id=data_point.id,
@@ -161,6 +142,7 @@ async def index_data_points(self, index_name: str, index_property_name: str, dat
async def retrieve(self, collection_name: str, data_point_ids: list[str]):
from pymilvus import MilvusException
+
client = self.get_milvus_client()
try:
filter_expression = f"""id in [{", ".join(f'"{id}"' for id in data_point_ids)}]"""
@@ -172,18 +154,21 @@ async def retrieve(self, collection_name: str, data_point_ids: list[str]):
)
return results
except MilvusException as e:
- logger.error(f"Error retrieving data points from collection '{collection_name}': {str(e)}")
+ logger.error(
+ f"Error retrieving data points from collection '{collection_name}': {str(e)}"
+ )
raise e
async def search(
- self,
- collection_name: str,
- query_text: Optional[str] = None,
- query_vector: Optional[List[float]] = None,
- limit: int = 5,
- with_vector: bool = False,
+ self,
+ collection_name: str,
+ query_text: Optional[str] = None,
+ query_vector: Optional[List[float]] = None,
+ limit: int = 5,
+ with_vector: bool = False,
):
from pymilvus import MilvusException
+
client = self.get_milvus_client()
if query_text is None and query_vector is None:
raise ValueError("One of query_text or query_vector must be provided!")
@@ -218,32 +203,40 @@ async def search(
logger.error(f"Error during search in collection '{collection_name}': {str(e)}")
raise e
- async def batch_search(self, collection_name: str, query_texts: List[str], limit: int, with_vectors: bool = False):
+ async def batch_search(
+ self, collection_name: str, query_texts: List[str], limit: int, with_vectors: bool = False
+ ):
query_vectors = await self.embed_data(query_texts)
return await asyncio.gather(
- *[self.search(collection_name=collection_name,
- query_vector=query_vector,
- limit=limit,
- with_vector=with_vectors,
- ) for query_vector in query_vectors]
+ *[
+ self.search(
+ collection_name=collection_name,
+ query_vector=query_vector,
+ limit=limit,
+ with_vector=with_vectors,
+ )
+ for query_vector in query_vectors
+ ]
)
async def delete_data_points(self, collection_name: str, data_point_ids: list[str]):
from pymilvus import MilvusException
+
client = self.get_milvus_client()
try:
filter_expression = f"""id in [{", ".join(f'"{id}"' for id in data_point_ids)}]"""
- delete_result = client.delete(
- collection_name=collection_name,
- filter=filter_expression
- )
+ delete_result = client.delete(collection_name=collection_name, filter=filter_expression)
- logger.info(f"Deleted data points with IDs {data_point_ids} from collection '{collection_name}'.")
+ logger.info(
+ f"Deleted data points with IDs {data_point_ids} from collection '{collection_name}'."
+ )
return delete_result
except MilvusException as e:
- logger.error(f"Error deleting data points from collection '{collection_name}': {str(e)}")
+ logger.error(
+ f"Error deleting data points from collection '{collection_name}': {str(e)}"
+ )
raise e
async def prune(self):
diff --git a/cognee/infrastructure/databases/vector/models/CollectionConfig.py b/cognee/infrastructure/databases/vector/models/CollectionConfig.py
index b46bdabfc..b279a6be2 100644
--- a/cognee/infrastructure/databases/vector/models/CollectionConfig.py
+++ b/cognee/infrastructure/databases/vector/models/CollectionConfig.py
@@ -1,5 +1,6 @@
from pydantic import BaseModel
from .VectorConfig import VectorConfig
+
class CollectionConfig(BaseModel):
vector_config: VectorConfig
diff --git a/cognee/infrastructure/databases/vector/models/ScoredResult.py b/cognee/infrastructure/databases/vector/models/ScoredResult.py
index f9d8bec77..8b6f00952 100644
--- a/cognee/infrastructure/databases/vector/models/ScoredResult.py
+++ b/cognee/infrastructure/databases/vector/models/ScoredResult.py
@@ -2,7 +2,8 @@
from uuid import UUID
from pydantic import BaseModel
+
class ScoredResult(BaseModel):
id: UUID
- score: float # Lower score is better
+ score: float # Lower score is better
payload: Dict[str, Any]
diff --git a/cognee/infrastructure/databases/vector/models/VectorConfig.py b/cognee/infrastructure/databases/vector/models/VectorConfig.py
index 2a968a997..fb89d90b5 100644
--- a/cognee/infrastructure/databases/vector/models/VectorConfig.py
+++ b/cognee/infrastructure/databases/vector/models/VectorConfig.py
@@ -1,6 +1,7 @@
from typing import Literal
from pydantic import BaseModel
+
class VectorConfig(BaseModel):
- distance: Literal['Cosine', 'Dot']
+ distance: Literal["Cosine", "Dot"]
size: int
diff --git a/cognee/infrastructure/databases/vector/pgvector/PGVectorAdapter.py b/cognee/infrastructure/databases/vector/pgvector/PGVectorAdapter.py
index 3f0565253..df22e8f18 100644
--- a/cognee/infrastructure/databases/vector/pgvector/PGVectorAdapter.py
+++ b/cognee/infrastructure/databases/vector/pgvector/PGVectorAdapter.py
@@ -22,13 +22,10 @@
class IndexSchema(DataPoint):
text: str
- _metadata: dict = {
- "index_fields": ["text"],
- "type": "IndexSchema"
- }
+ _metadata: dict = {"index_fields": ["text"], "type": "IndexSchema"}
-class PGVectorAdapter(SQLAlchemyAdapter, VectorDBInterface):
+class PGVectorAdapter(SQLAlchemyAdapter, VectorDBInterface):
def __init__(
self,
connection_string: str,
@@ -44,6 +41,7 @@ def __init__(
# Has to be imported at class level
# Functions reading tables from database need to know what a Vector column type is
from pgvector.sqlalchemy import Vector
+
self.Vector = Vector
async def embed_data(self, data: list[str]) -> list[list[float]]:
@@ -71,9 +69,7 @@ class PGVectorDataPoint(Base):
__tablename__ = collection_name
__table_args__ = {"extend_existing": True}
# PGVector requires one column to be the primary key
- primary_key: Mapped[int] = mapped_column(
- primary_key=True, autoincrement=True
- )
+ primary_key: Mapped[int] = mapped_column(primary_key=True, autoincrement=True)
id: Mapped[data_point_types["id"]]
payload = Column(JSON)
vector = Column(self.Vector(vector_size))
@@ -89,14 +85,12 @@ def __init__(self, id, payload, vector):
Base.metadata.create_all, tables=[PGVectorDataPoint.__table__]
)
- async def create_data_points(
- self, collection_name: str, data_points: List[DataPoint]
- ):
+ async def create_data_points(self, collection_name: str, data_points: List[DataPoint]):
data_point_types = get_type_hints(DataPoint)
if not await self.has_collection(collection_name):
await self.create_collection(
- collection_name = collection_name,
- payload_schema = type(data_points[0]),
+ collection_name=collection_name,
+ payload_schema=type(data_points[0]),
)
data_vectors = await self.embed_data(
@@ -109,9 +103,7 @@ class PGVectorDataPoint(Base):
__tablename__ = collection_name
__table_args__ = {"extend_existing": True}
# PGVector requires one column to be the primary key
- primary_key: Mapped[int] = mapped_column(
- primary_key=True, autoincrement=True
- )
+ primary_key: Mapped[int] = mapped_column(primary_key=True, autoincrement=True)
id: Mapped[data_point_types["id"]]
payload = Column(JSON)
vector = Column(self.Vector(vector_size))
@@ -123,9 +115,9 @@ def __init__(self, id, payload, vector):
pgvector_data_points = [
PGVectorDataPoint(
- id = data_point.id,
- vector = data_vectors[data_index],
- payload = serialize_data(data_point.model_dump()),
+ id=data_point.id,
+ vector=data_vectors[data_index],
+ payload=serialize_data(data_point.model_dump()),
)
for (data_index, data_point) in enumerate(data_points)
]
@@ -137,13 +129,19 @@ def __init__(self, id, payload, vector):
async def create_vector_index(self, index_name: str, index_property_name: str):
await self.create_collection(f"{index_name}_{index_property_name}")
- async def index_data_points(self, index_name: str, index_property_name: str, data_points: list[DataPoint]):
- await self.create_data_points(f"{index_name}_{index_property_name}", [
- IndexSchema(
- id = data_point.id,
- text = DataPoint.get_embeddable_data(data_point),
- ) for data_point in data_points
- ])
+ async def index_data_points(
+ self, index_name: str, index_property_name: str, data_points: list[DataPoint]
+ ):
+ await self.create_data_points(
+ f"{index_name}_{index_property_name}",
+ [
+ IndexSchema(
+ id=data_point.id,
+ text=DataPoint.get_embeddable_data(data_point),
+ )
+ for data_point in data_points
+ ],
+ )
async def get_table(self, collection_name: str) -> Table:
"""
@@ -171,20 +169,17 @@ async def retrieve(self, collection_name: str, data_point_ids: List[str]):
results = results.all()
return [
- ScoredResult(
- id = UUID(result.id),
- payload = result.payload,
- score = 0
- ) for result in results
+ ScoredResult(id=UUID(result.id), payload=result.payload, score=0)
+ for result in results
]
async def get_distance_from_collection_elements(
- self,
- collection_name: str,
- query_text: str = None,
- query_vector: List[float] = None,
- with_vector: bool = False
- )-> List[ScoredResult]:
+ self,
+ collection_name: str,
+ query_text: str = None,
+ query_vector: List[float] = None,
+ with_vector: bool = False,
+ ) -> List[ScoredResult]:
if query_text is None and query_vector is None:
raise ValueError("One of query_text or query_vector must be provided!")
@@ -200,11 +195,8 @@ async def get_distance_from_collection_elements(
closest_items = await session.execute(
select(
PGVectorDataPoint,
- PGVectorDataPoint.c.vector.cosine_distance(query_vector).label(
- "similarity"
- ),
- )
- .order_by("similarity")
+ PGVectorDataPoint.c.vector.cosine_distance(query_vector).label("similarity"),
+ ).order_by("similarity")
)
vector_list = []
@@ -216,11 +208,8 @@ async def get_distance_from_collection_elements(
# Create and return ScoredResult objects
return [
- ScoredResult(
- id = UUID(str(row.id)),
- payload = row.payload,
- score = row.similarity
- ) for row in vector_list
+ ScoredResult(id=UUID(str(row.id)), payload=row.payload, score=row.similarity)
+ for row in vector_list
]
async def search(
@@ -248,9 +237,7 @@ async def search(
closest_items = await session.execute(
select(
PGVectorDataPoint,
- PGVectorDataPoint.c.vector.cosine_distance(query_vector).label(
- "similarity"
- ),
+ PGVectorDataPoint.c.vector.cosine_distance(query_vector).label("similarity"),
)
.order_by("similarity")
.limit(limit)
@@ -265,11 +252,8 @@ async def search(
# Create and return ScoredResult objects
return [
- ScoredResult(
- id = UUID(str(row.id)),
- payload = row.payload,
- score = row.similarity
- ) for row in vector_list
+ ScoredResult(id=UUID(str(row.id)), payload=row.payload, score=row.similarity)
+ for row in vector_list
]
async def batch_search(
diff --git a/cognee/infrastructure/databases/vector/pgvector/__init__.py b/cognee/infrastructure/databases/vector/pgvector/__init__.py
index 130246a31..2f8f08620 100644
--- a/cognee/infrastructure/databases/vector/pgvector/__init__.py
+++ b/cognee/infrastructure/databases/vector/pgvector/__init__.py
@@ -1,2 +1,2 @@
from .PGVectorAdapter import PGVectorAdapter
-from .create_db_and_tables import create_db_and_tables
\ No newline at end of file
+from .create_db_and_tables import create_db_and_tables
diff --git a/cognee/infrastructure/databases/vector/pgvector/create_db_and_tables.py b/cognee/infrastructure/databases/vector/pgvector/create_db_and_tables.py
index 2f4c9cf3f..881c1e72f 100644
--- a/cognee/infrastructure/databases/vector/pgvector/create_db_and_tables.py
+++ b/cognee/infrastructure/databases/vector/pgvector/create_db_and_tables.py
@@ -9,4 +9,3 @@ async def create_db_and_tables():
if vector_config.vector_db_provider == "pgvector":
async with vector_engine.engine.begin() as connection:
await connection.execute(text("CREATE EXTENSION IF NOT EXISTS vector;"))
-
diff --git a/cognee/infrastructure/databases/vector/pgvector/serialize_data.py b/cognee/infrastructure/databases/vector/pgvector/serialize_data.py
index cdba1e928..a61a2682f 100644
--- a/cognee/infrastructure/databases/vector/pgvector/serialize_data.py
+++ b/cognee/infrastructure/databases/vector/pgvector/serialize_data.py
@@ -1,6 +1,7 @@
from datetime import datetime
from uuid import UUID
+
def serialize_data(data):
"""Recursively convert datetime objects in dictionaries/lists to ISO format."""
if isinstance(data, dict):
@@ -12,4 +13,4 @@ def serialize_data(data):
elif isinstance(data, UUID):
return str(data)
else:
- return data
\ No newline at end of file
+ return data
diff --git a/cognee/infrastructure/databases/vector/pinecone/adapter.py b/cognee/infrastructure/databases/vector/pinecone/adapter.py
index 2a8e628f9..00dd7daf1 100644
--- a/cognee/infrastructure/databases/vector/pinecone/adapter.py
+++ b/cognee/infrastructure/databases/vector/pinecone/adapter.py
@@ -1,8 +1,8 @@
-class PineconeVectorDB(VectorDB):
- def __init__(self, *args, **kwargs):
- super().__init__(*args, **kwargs)
- self.init_pinecone(self.index_name)
-
- def init_pinecone(self, index_name):
- # Pinecone initialization logic
- pass
+# class PineconeVectorDB(VectorDB):
+# def __init__(self, *args, **kwargs):
+# super().__init__(*args, **kwargs)
+# self.init_pinecone(self.index_name)
+#
+# def init_pinecone(self, index_name):
+# # Pinecone initialization logic
+# pass
diff --git a/cognee/infrastructure/databases/vector/qdrant/QDrantAdapter.py b/cognee/infrastructure/databases/vector/qdrant/QDrantAdapter.py
index b63139bc5..64e64e269 100644
--- a/cognee/infrastructure/databases/vector/qdrant/QDrantAdapter.py
+++ b/cognee/infrastructure/databases/vector/qdrant/QDrantAdapter.py
@@ -5,8 +5,7 @@
from qdrant_client import AsyncQdrantClient, models
from cognee.exceptions import InvalidValueError
-from cognee.infrastructure.databases.vector.models.ScoredResult import \
- ScoredResult
+from cognee.infrastructure.databases.vector.models.ScoredResult import ScoredResult
from cognee.infrastructure.engine import DataPoint
from ..embeddings.EmbeddingEngine import EmbeddingEngine
@@ -14,13 +13,12 @@
logger = logging.getLogger("QDrantAdapter")
+
class IndexSchema(DataPoint):
text: str
- _metadata: dict = {
- "index_fields": ["text"],
- "type": "IndexSchema"
- }
+ _metadata: dict = {"index_fields": ["text"], "type": "IndexSchema"}
+
# class CollectionConfig(BaseModel, extra = "forbid"):
# vector_config: Dict[str, models.VectorParams] = Field(..., description="Vectors configuration" )
@@ -28,28 +26,32 @@ class IndexSchema(DataPoint):
# optimizers_config: Optional[models.OptimizersConfig] = Field(default = None, description="Optimizers configuration")
# quantization_config: Optional[models.QuantizationConfig] = Field(default = None, description="Quantization configuration")
+
def create_hnsw_config(hnsw_config: Dict):
if hnsw_config is not None:
return models.HnswConfig()
return None
+
def create_optimizers_config(optimizers_config: Dict):
if optimizers_config is not None:
return models.OptimizersConfig()
return None
+
def create_quantization_config(quantization_config: Dict):
if quantization_config is not None:
return models.QuantizationConfig()
return None
+
class QDrantAdapter(VectorDBInterface):
name = "Qdrant"
url: str = None
api_key: str = None
qdrant_path: str = None
- def __init__(self, url, api_key, embedding_engine: EmbeddingEngine, qdrant_path = None):
+ def __init__(self, url, api_key, embedding_engine: EmbeddingEngine, qdrant_path=None):
self.embedding_engine = embedding_engine
if qdrant_path is not None:
@@ -60,19 +62,11 @@ def __init__(self, url, api_key, embedding_engine: EmbeddingEngine, qdrant_path
def get_qdrant_client(self) -> AsyncQdrantClient:
if self.qdrant_path is not None:
- return AsyncQdrantClient(
- path = self.qdrant_path, port=6333
- )
+ return AsyncQdrantClient(path=self.qdrant_path, port=6333)
elif self.url is not None:
- return AsyncQdrantClient(
- url = self.url,
- api_key = self.api_key,
- port = 6333
- )
+ return AsyncQdrantClient(url=self.url, api_key=self.api_key, port=6333)
- return AsyncQdrantClient(
- location = ":memory:"
- )
+ return AsyncQdrantClient(location=":memory:")
async def embed_data(self, data: List[str]) -> List[float]:
return await self.embedding_engine.embed_text(data)
@@ -84,21 +78,20 @@ async def has_collection(self, collection_name: str) -> bool:
return result
async def create_collection(
- self,
- collection_name: str,
- payload_schema = None,
+ self,
+ collection_name: str,
+ payload_schema=None,
):
client = self.get_qdrant_client()
if not await client.collection_exists(collection_name):
await client.create_collection(
- collection_name = collection_name,
- vectors_config = {
+ collection_name=collection_name,
+ vectors_config={
"text": models.VectorParams(
- size = self.embedding_engine.get_vector_size(),
- distance = "Cosine"
+ size=self.embedding_engine.get_vector_size(), distance="Cosine"
)
- }
+ },
)
await client.close()
@@ -106,26 +99,21 @@ async def create_collection(
async def create_data_points(self, collection_name: str, data_points: List[DataPoint]):
client = self.get_qdrant_client()
- data_vectors = await self.embed_data([
- DataPoint.get_embeddable_data(data_point) for data_point in data_points
- ])
+ data_vectors = await self.embed_data(
+ [DataPoint.get_embeddable_data(data_point) for data_point in data_points]
+ )
def convert_to_qdrant_point(data_point: DataPoint):
return models.PointStruct(
- id = str(data_point.id),
- payload = data_point.model_dump(),
- vector = {
- "text": data_vectors[data_points.index(data_point)]
- }
+ id=str(data_point.id),
+ payload=data_point.model_dump(),
+ vector={"text": data_vectors[data_points.index(data_point)]},
)
points = [convert_to_qdrant_point(point) for point in data_points]
try:
- client.upload_points(
- collection_name = collection_name,
- points = points
- )
+ client.upload_points(collection_name=collection_name, points=points)
except Exception as error:
logger.error("Error uploading data points to Qdrant: %s", str(error))
raise error
@@ -135,53 +123,61 @@ def convert_to_qdrant_point(data_point: DataPoint):
async def create_vector_index(self, index_name: str, index_property_name: str):
await self.create_collection(f"{index_name}_{index_property_name}")
- async def index_data_points(self, index_name: str, index_property_name: str, data_points: list[DataPoint]):
- await self.create_data_points(f"{index_name}_{index_property_name}", [
- IndexSchema(
- id = data_point.id,
- text = getattr(data_point, data_point._metadata["index_fields"][0]),
- ) for data_point in data_points
- ])
+ async def index_data_points(
+ self, index_name: str, index_property_name: str, data_points: list[DataPoint]
+ ):
+ await self.create_data_points(
+ f"{index_name}_{index_property_name}",
+ [
+ IndexSchema(
+ id=data_point.id,
+ text=getattr(data_point, data_point._metadata["index_fields"][0]),
+ )
+ for data_point in data_points
+ ],
+ )
async def retrieve(self, collection_name: str, data_point_ids: list[str]):
client = self.get_qdrant_client()
- results = await client.retrieve(collection_name, data_point_ids, with_payload = True)
+ results = await client.retrieve(collection_name, data_point_ids, with_payload=True)
await client.close()
return results
async def get_distance_from_collection_elements(
- self,
- collection_name: str,
- query_text: str = None,
- query_vector: List[float] = None,
- with_vector: bool = False
+ self,
+ collection_name: str,
+ query_text: str = None,
+ query_vector: List[float] = None,
+ with_vector: bool = False,
) -> List[ScoredResult]:
-
if query_text is None and query_vector is None:
raise ValueError("One of query_text or query_vector must be provided!")
client = self.get_qdrant_client()
results = await client.search(
- collection_name = collection_name,
- query_vector = models.NamedVector(
- name = "text",
- vector = query_vector if query_vector is not None else (await self.embed_data([query_text]))[0],
+ collection_name=collection_name,
+ query_vector=models.NamedVector(
+ name="text",
+ vector=query_vector
+ if query_vector is not None
+ else (await self.embed_data([query_text]))[0],
),
- with_vectors = with_vector
+ with_vectors=with_vector,
)
await client.close()
return [
ScoredResult(
- id = UUID(result.id),
- payload = {
+ id=UUID(result.id),
+ payload={
**result.payload,
"id": UUID(result.id),
},
- score = 1 - result.score,
- ) for result in results
+ score=1 - result.score,
+ )
+ for result in results
]
async def search(
@@ -190,7 +186,7 @@ async def search(
query_text: Optional[str] = None,
query_vector: Optional[List[float]] = None,
limit: int = 5,
- with_vector: bool = False
+ with_vector: bool = False,
):
if query_text is None and query_vector is None:
raise InvalidValueError(message="One of query_text or query_vector must be provided!")
@@ -198,30 +194,38 @@ async def search(
client = self.get_qdrant_client()
results = await client.search(
- collection_name = collection_name,
- query_vector = models.NamedVector(
- name = "text",
- vector = query_vector if query_vector is not None else (await self.embed_data([query_text]))[0],
+ collection_name=collection_name,
+ query_vector=models.NamedVector(
+ name="text",
+ vector=query_vector
+ if query_vector is not None
+ else (await self.embed_data([query_text]))[0],
),
- limit = limit,
- with_vectors = with_vector
+ limit=limit,
+ with_vectors=with_vector,
)
await client.close()
return [
ScoredResult(
- id = UUID(result.id),
- payload = {
+ id=UUID(result.id),
+ payload={
**result.payload,
"id": UUID(result.id),
},
- score = 1 - result.score,
- ) for result in results
+ score=1 - result.score,
+ )
+ for result in results
]
-
- async def batch_search(self, collection_name: str, query_texts: List[str], limit: int = None, with_vectors: bool = False):
+ async def batch_search(
+ self,
+ collection_name: str,
+ query_texts: List[str],
+ limit: int = None,
+ with_vectors: bool = False,
+ ):
"""
Perform batch search in a Qdrant collection with dynamic search requests.
@@ -240,22 +244,17 @@ async def batch_search(self, collection_name: str, query_texts: List[str], limit
# Generate dynamic search requests based on the provided embeddings
requests = [
models.SearchRequest(
- vector = models.NamedVector(
- name = "text",
- vector = vector
- ),
- limit = limit,
- with_vector = with_vectors
- ) for vector in vectors
+ vector=models.NamedVector(name="text", vector=vector),
+ limit=limit,
+ with_vector=with_vectors,
+ )
+ for vector in vectors
]
client = self.get_qdrant_client()
# Perform batch search with the dynamically generated requests
- results = await client.search_batch(
- collection_name = collection_name,
- requests = requests
- )
+ results = await client.search_batch(collection_name=collection_name, requests=requests)
await client.close()
diff --git a/cognee/infrastructure/databases/vector/utils.py b/cognee/infrastructure/databases/vector/utils.py
index d5a5897a3..9507207c0 100644
--- a/cognee/infrastructure/databases/vector/utils.py
+++ b/cognee/infrastructure/databases/vector/utils.py
@@ -1,5 +1,6 @@
from typing import List
+
def normalize_distances(result_values: List[dict]) -> List[float]:
min_value = min(result["_distance"] for result in result_values)
max_value = max(result["_distance"] for result in result_values)
@@ -8,7 +9,8 @@ def normalize_distances(result_values: List[dict]) -> List[float]:
# Avoid division by zero: Assign all normalized values to 0 (or any constant value like 1)
normalized_values = [0 for _ in result_values]
else:
- normalized_values = [(result["_distance"] - min_value) / (max_value - min_value) for result in
- result_values]
+ normalized_values = [
+ (result["_distance"] - min_value) / (max_value - min_value) for result in result_values
+ ]
return normalized_values
diff --git a/cognee/infrastructure/databases/vector/vector_db_interface.py b/cognee/infrastructure/databases/vector/vector_db_interface.py
index 457b92f07..939997fec 100644
--- a/cognee/infrastructure/databases/vector/vector_db_interface.py
+++ b/cognee/infrastructure/databases/vector/vector_db_interface.py
@@ -3,8 +3,10 @@
from cognee.infrastructure.engine import DataPoint
from .models.PayloadSchema import PayloadSchema
+
class VectorDBInterface(Protocol):
- """ Collections """
+ """Collections"""
+
@abstractmethod
async def has_collection(self, collection_name: str) -> bool:
raise NotImplementedError
@@ -14,24 +16,21 @@ async def create_collection(
self,
collection_name: str,
payload_schema: Optional[PayloadSchema] = None,
- ): raise NotImplementedError
+ ):
+ raise NotImplementedError
""" Data points """
+
@abstractmethod
- async def create_data_points(
- self,
- collection_name: str,
- data_points: List[DataPoint]
- ): raise NotImplementedError
+ async def create_data_points(self, collection_name: str, data_points: List[DataPoint]):
+ raise NotImplementedError
@abstractmethod
- async def retrieve(
- self,
- collection_name: str,
- data_point_ids: list[str]
- ): raise NotImplementedError
+ async def retrieve(self, collection_name: str, data_point_ids: list[str]):
+ raise NotImplementedError
""" Search """
+
@abstractmethod
async def search(
self,
@@ -39,25 +38,20 @@ async def search(
query_text: Optional[str],
query_vector: Optional[List[float]],
limit: int,
- with_vector: bool = False
-
- ): raise NotImplementedError
+ with_vector: bool = False,
+ ):
+ raise NotImplementedError
@abstractmethod
async def batch_search(
- self,
- collection_name: str,
- query_texts: List[str],
- limit: int,
- with_vectors: bool = False
- ): raise NotImplementedError
+ self, collection_name: str, query_texts: List[str], limit: int, with_vectors: bool = False
+ ):
+ raise NotImplementedError
@abstractmethod
- async def delete_data_points(
- self,
- collection_name: str,
- data_point_ids: list[str]
- ): raise NotImplementedError
+ async def delete_data_points(self, collection_name: str, data_point_ids: list[str]):
+ raise NotImplementedError
@abstractmethod
- async def prune(self): raise NotImplementedError
+ async def prune(self):
+ raise NotImplementedError
diff --git a/cognee/infrastructure/databases/vector/weaviate_db/WeaviateAdapter.py b/cognee/infrastructure/databases/vector/weaviate_db/WeaviateAdapter.py
index 31162b1b5..259370643 100644
--- a/cognee/infrastructure/databases/vector/weaviate_db/WeaviateAdapter.py
+++ b/cognee/infrastructure/databases/vector/weaviate_db/WeaviateAdapter.py
@@ -12,13 +12,12 @@
logger = logging.getLogger("WeaviateAdapter")
+
class IndexSchema(DataPoint):
text: str
- _metadata: dict = {
- "index_fields": ["text"],
- "type": "IndexSchema"
- }
+ _metadata: dict = {"index_fields": ["text"], "type": "IndexSchema"}
+
class WeaviateAdapter(VectorDBInterface):
name = "Weaviate"
@@ -36,9 +35,9 @@ def __init__(self, url: str, api_key: str, embedding_engine: EmbeddingEngine):
self.embedding_engine = embedding_engine
self.client = weaviate.connect_to_wcs(
- cluster_url = url,
- auth_credentials = weaviate.auth.AuthApiKey(api_key),
- additional_config = wvc.init.AdditionalConfig(timeout = wvc.init.Timeout(init=30))
+ cluster_url=url,
+ auth_credentials=weaviate.auth.AuthApiKey(api_key),
+ additional_config=wvc.init.AdditionalConfig(timeout=wvc.init.Timeout(init=30)),
)
async def embed_data(self, data: List[str]) -> List[float]:
@@ -54,7 +53,7 @@ async def has_collection(self, collection_name: str) -> bool:
async def create_collection(
self,
collection_name: str,
- payload_schema = None,
+ payload_schema=None,
):
import weaviate.classes.config as wvcc
@@ -63,14 +62,12 @@ async def create_collection(
if not self.client.collections.exists(collection_name):
future.set_result(
self.client.collections.create(
- name = collection_name,
- properties = [
+ name=collection_name,
+ properties=[
wvcc.Property(
- name = "text",
- data_type = wvcc.DataType.TEXT,
- skip_vectorization = True
+ name="text", data_type=wvcc.DataType.TEXT, skip_vectorization=True
)
- ]
+ ],
)
)
else:
@@ -96,11 +93,7 @@ def convert_to_weaviate_data_points(data_point: DataPoint):
properties["uuid"] = str(data_point.id)
del properties["id"]
- return DataObject(
- uuid = data_point.id,
- properties = properties,
- vector = vector
- )
+ return DataObject(uuid=data_point.id, properties=properties, vector=vector)
data_points = [convert_to_weaviate_data_points(data_point) for data_point in data_points]
@@ -111,26 +104,26 @@ def convert_to_weaviate_data_points(data_point: DataPoint):
with collection.batch.dynamic() as batch:
for data_point in data_points:
batch.add_object(
- uuid = data_point.uuid,
- vector = data_point.vector,
- properties = data_point.properties,
- references = data_point.references,
+ uuid=data_point.uuid,
+ vector=data_point.vector,
+ properties=data_point.properties,
+ references=data_point.references,
)
else:
data_point: DataObject = data_points[0]
if collection.data.exists(data_point.uuid):
return collection.data.update(
- uuid = data_point.uuid,
- vector = data_point.vector,
- properties = data_point.properties,
- references = data_point.references,
+ uuid=data_point.uuid,
+ vector=data_point.vector,
+ properties=data_point.properties,
+ references=data_point.references,
)
else:
return collection.data.insert(
- uuid = data_point.uuid,
- vector = data_point.vector,
- properties = data_point.properties,
- references = data_point.references,
+ uuid=data_point.uuid,
+ vector=data_point.vector,
+ properties=data_point.properties,
+ references=data_point.references,
)
except Exception as error:
logger.error("Error creating data points: %s", str(error))
@@ -139,20 +132,27 @@ def convert_to_weaviate_data_points(data_point: DataPoint):
async def create_vector_index(self, index_name: str, index_property_name: str):
await self.create_collection(f"{index_name}_{index_property_name}")
- async def index_data_points(self, index_name: str, index_property_name: str, data_points: list[DataPoint]):
- await self.create_data_points(f"{index_name}_{index_property_name}", [
- IndexSchema(
- id = data_point.id,
- text = DataPoint.get_embeddable_data(data_point),
- ) for data_point in data_points
- ])
+ async def index_data_points(
+ self, index_name: str, index_property_name: str, data_points: list[DataPoint]
+ ):
+ await self.create_data_points(
+ f"{index_name}_{index_property_name}",
+ [
+ IndexSchema(
+ id=data_point.id,
+ text=DataPoint.get_embeddable_data(data_point),
+ )
+ for data_point in data_points
+ ],
+ )
async def retrieve(self, collection_name: str, data_point_ids: list[str]):
from weaviate.classes.query import Filter
+
future = asyncio.Future()
data_points = self.get_collection(collection_name).query.fetch_objects(
- filters = Filter.by_id().contains_any(data_point_ids)
+ filters=Filter.by_id().contains_any(data_point_ids)
)
for data_point in data_points.objects:
@@ -165,11 +165,11 @@ async def retrieve(self, collection_name: str, data_point_ids: list[str]):
return await future
async def get_distance_from_collection_elements(
- self,
- collection_name: str,
- query_text: str = None,
- query_vector: List[float] = None,
- with_vector: bool = False
+ self,
+ collection_name: str,
+ query_text: str = None,
+ query_vector: List[float] = None,
+ with_vector: bool = False,
) -> List[ScoredResult]:
import weaviate.classes as wvc
@@ -190,17 +190,18 @@ async def get_distance_from_collection_elements(
ScoredResult(
id=UUID(str(result.uuid)),
payload=result.properties,
- score=1 - float(result.metadata.score)
- ) for result in search_result.objects
+ score=1 - float(result.metadata.score),
+ )
+ for result in search_result.objects
]
async def search(
- self,
- collection_name: str,
- query_text: Optional[str] = None,
- query_vector: Optional[List[float]] = None,
- limit: int = None,
- with_vector: bool = False
+ self,
+ collection_name: str,
+ query_text: Optional[str] = None,
+ query_vector: Optional[List[float]] = None,
+ limit: int = None,
+ with_vector: bool = False,
):
import weaviate.classes as wvc
@@ -211,33 +212,41 @@ async def search(
query_vector = (await self.embed_data([query_text]))[0]
search_result = self.get_collection(collection_name).query.hybrid(
- query = None,
- vector = query_vector,
- limit = limit,
- include_vector = with_vector,
- return_metadata = wvc.query.MetadataQuery(score=True),
+ query=None,
+ vector=query_vector,
+ limit=limit,
+ include_vector=with_vector,
+ return_metadata=wvc.query.MetadataQuery(score=True),
)
return [
ScoredResult(
- id = UUID(str(result.uuid)),
- payload = result.properties,
- score = 1 - float(result.metadata.score)
- ) for result in search_result.objects
+ id=UUID(str(result.uuid)),
+ payload=result.properties,
+ score=1 - float(result.metadata.score),
+ )
+ for result in search_result.objects
]
- async def batch_search(self, collection_name: str, query_texts: List[str], limit: int, with_vectors: bool = False):
+ async def batch_search(
+ self, collection_name: str, query_texts: List[str], limit: int, with_vectors: bool = False
+ ):
def query_search(query_vector):
- return self.search(collection_name, query_vector=query_vector, limit=limit, with_vector=with_vectors)
+ return self.search(
+ collection_name, query_vector=query_vector, limit=limit, with_vector=with_vectors
+ )
+
+ return [
+ await query_search(query_vector) for query_vector in await self.embed_data(query_texts)
+ ]
- return [await query_search(query_vector) for query_vector in await self.embed_data(query_texts)]
-
async def delete_data_points(self, collection_name: str, data_point_ids: list[str]):
from weaviate.classes.query import Filter
+
future = asyncio.Future()
result = self.get_collection(collection_name).data.delete_many(
- filters = Filter.by_id().contains_any(data_point_ids)
+ filters=Filter.by_id().contains_any(data_point_ids)
)
future.set_result(result)
diff --git a/cognee/infrastructure/engine/models/DataPoint.py b/cognee/infrastructure/engine/models/DataPoint.py
index e08041146..db0d9308a 100644
--- a/cognee/infrastructure/engine/models/DataPoint.py
+++ b/cognee/infrastructure/engine/models/DataPoint.py
@@ -9,23 +9,24 @@
class MetaData(TypedDict):
index_fields: list[str]
+
class DataPoint(BaseModel):
__tablename__ = "data_point"
- id: UUID = Field(default_factory = uuid4)
+ id: UUID = Field(default_factory=uuid4)
updated_at: Optional[datetime] = datetime.now(timezone.utc)
topological_rank: Optional[int] = 0
- _metadata: Optional[MetaData] = {
- "index_fields": [],
- "type": "DataPoint"
- }
+ _metadata: Optional[MetaData] = {"index_fields": [], "type": "DataPoint"}
# class Config:
# underscore_attrs_are_private = True
@classmethod
def get_embeddable_data(self, data_point):
- if data_point._metadata and len(data_point._metadata["index_fields"]) > 0 \
- and hasattr(data_point, data_point._metadata["index_fields"][0]):
+ if (
+ data_point._metadata
+ and len(data_point._metadata["index_fields"]) > 0
+ and hasattr(data_point, data_point._metadata["index_fields"][0])
+ ):
attribute = getattr(data_point, data_point._metadata["index_fields"][0])
if isinstance(attribute, str):
@@ -36,10 +37,12 @@ def get_embeddable_data(self, data_point):
@classmethod
def get_embeddable_properties(self, data_point):
if data_point._metadata and len(data_point._metadata["index_fields"]) > 0:
- return [getattr(data_point, field, None) for field in data_point._metadata["index_fields"]]
+ return [
+ getattr(data_point, field, None) for field in data_point._metadata["index_fields"]
+ ]
return []
@classmethod
def get_embeddable_property_names(self, data_point):
- return data_point._metadata["index_fields"] or []
\ No newline at end of file
+ return data_point._metadata["index_fields"] or []
diff --git a/cognee/infrastructure/files/add_file_to_storage.py b/cognee/infrastructure/files/add_file_to_storage.py
index e2a324a83..78713f447 100644
--- a/cognee/infrastructure/files/add_file_to_storage.py
+++ b/cognee/infrastructure/files/add_file_to_storage.py
@@ -3,6 +3,7 @@
from .storage.StorageManager import StorageManager
from .storage.LocalStorage import LocalStorage
+
async def add_file_to_storage(file_path: str, file: BinaryIO):
storage_manager = StorageManager(LocalStorage(get_absolute_path("data/files")))
diff --git a/cognee/infrastructure/files/remove_file_from_storage.py b/cognee/infrastructure/files/remove_file_from_storage.py
index b0b9a5a61..b657ead43 100644
--- a/cognee/infrastructure/files/remove_file_from_storage.py
+++ b/cognee/infrastructure/files/remove_file_from_storage.py
@@ -2,6 +2,7 @@
from .storage.StorageManager import StorageManager
from .storage.LocalStorage import LocalStorage
+
async def remove_file_from_storage(file_path: str):
storage_manager = StorageManager(LocalStorage(get_absolute_path("data/files")))
diff --git a/cognee/infrastructure/files/storage/LocalStorage.py b/cognee/infrastructure/files/storage/LocalStorage.py
index 4db17c91b..7ef9131d9 100644
--- a/cognee/infrastructure/files/storage/LocalStorage.py
+++ b/cognee/infrastructure/files/storage/LocalStorage.py
@@ -3,6 +3,7 @@
from typing import BinaryIO, Union
from .StorageManager import Storage
+
class LocalStorage(Storage):
storage_path: str = None
@@ -16,8 +17,8 @@ def store(self, file_path: str, data: Union[BinaryIO, str]):
with open(
full_file_path,
- mode = "w" if isinstance(data, str) else "wb",
- encoding = "utf-8" if isinstance(data, str) else None
+ mode="w" if isinstance(data, str) else "wb",
+ encoding="utf-8" if isinstance(data, str) else None,
) as f:
if hasattr(data, "read"):
data.seek(0)
@@ -28,7 +29,7 @@ def store(self, file_path: str, data: Union[BinaryIO, str]):
def retrieve(self, file_path: str, mode: str = "rb"):
full_file_path = self.storage_path + "/" + file_path
- with open(full_file_path, mode = mode) as f:
+ with open(full_file_path, mode=mode) as f:
f.seek(0)
return f.read()
@@ -39,7 +40,7 @@ def file_exists(file_path: str):
@staticmethod
def ensure_directory_exists(file_path: str):
if not os.path.exists(file_path):
- os.makedirs(file_path, exist_ok = True)
+ os.makedirs(file_path, exist_ok=True)
@staticmethod
def remove(file_path: str):
diff --git a/cognee/infrastructure/files/storage/StorageManager.py b/cognee/infrastructure/files/storage/StorageManager.py
index 1db2d5d31..389800e00 100644
--- a/cognee/infrastructure/files/storage/StorageManager.py
+++ b/cognee/infrastructure/files/storage/StorageManager.py
@@ -1,5 +1,6 @@
from typing import Protocol, BinaryIO
+
class Storage(Protocol):
def store(self, file_path: str, data: bytes):
pass
@@ -11,10 +12,11 @@ def retrieve(self, file_path: str):
def remove(file_path: str):
pass
-class StorageManager():
+
+class StorageManager:
storage: Storage = None
- def __init__ (self, storage: Storage):
+ def __init__(self, storage: Storage):
self.storage = storage
def store(self, file_path: str, data: BinaryIO):
diff --git a/cognee/infrastructure/files/utils/extract_text_from_file.py b/cognee/infrastructure/files/utils/extract_text_from_file.py
index 171128b2e..7a5e16829 100644
--- a/cognee/infrastructure/files/utils/extract_text_from_file.py
+++ b/cognee/infrastructure/files/utils/extract_text_from_file.py
@@ -2,10 +2,11 @@
from pypdf import PdfReader
import filetype
+
def extract_text_from_file(file: BinaryIO, file_type: filetype.Type) -> str:
"""Extract text from a file"""
if file_type.extension == "pdf":
- reader = PdfReader(stream = file)
+ reader = PdfReader(stream=file)
pages = list(reader.pages[:3])
return "\n".join([page.extract_text().strip() for page in pages])
diff --git a/cognee/infrastructure/files/utils/get_file_metadata.py b/cognee/infrastructure/files/utils/get_file_metadata.py
index 89c3d6d8e..4bce29f60 100644
--- a/cognee/infrastructure/files/utils/get_file_metadata.py
+++ b/cognee/infrastructure/files/utils/get_file_metadata.py
@@ -11,6 +11,7 @@ class FileMetadata(TypedDict):
extension: str
content_hash: str
+
def get_file_metadata(file: BinaryIO) -> FileMetadata:
"""Get metadata from a file"""
file.seek(0)
@@ -23,9 +24,9 @@ def get_file_metadata(file: BinaryIO) -> FileMetadata:
file_name = file_path.split("/")[-1].split(".")[0] if file_path else None
return FileMetadata(
- name = file_name,
- file_path = file_path,
- mime_type = file_type.mime,
- extension = file_type.extension,
- content_hash = content_hash,
+ name=file_name,
+ file_path=file_path,
+ mime_type=file_type.mime,
+ extension=file_type.extension,
+ content_hash=content_hash,
)
diff --git a/cognee/infrastructure/files/utils/guess_file_type.py b/cognee/infrastructure/files/utils/guess_file_type.py
index 001585945..3e26827d8 100644
--- a/cognee/infrastructure/files/utils/guess_file_type.py
+++ b/cognee/infrastructure/files/utils/guess_file_type.py
@@ -2,37 +2,45 @@
import filetype
from .is_text_content import is_text_content
+
class FileTypeException(Exception):
message: str
def __init__(self, message: str):
self.message = message
+
class TxtFileType(filetype.Type):
"""Text file type"""
+
MIME = "text/plain"
EXTENSION = "txt"
def __init__(self):
- super(TxtFileType, self).__init__(mime = TxtFileType.MIME, extension = TxtFileType.EXTENSION)
+ super(TxtFileType, self).__init__(mime=TxtFileType.MIME, extension=TxtFileType.EXTENSION)
def match(self, buf):
return is_text_content(buf)
+
txt_file_type = TxtFileType()
filetype.add_type(txt_file_type)
+
class CustomPdfMatcher(filetype.Type):
MIME = "application/pdf"
EXTENSION = "pdf"
def __init__(self):
- super(CustomPdfMatcher, self).__init__(mime = CustomPdfMatcher.MIME, extension = CustomPdfMatcher.EXTENSION)
+ super(CustomPdfMatcher, self).__init__(
+ mime=CustomPdfMatcher.MIME, extension=CustomPdfMatcher.EXTENSION
+ )
def match(self, buf):
return b"PDF-" in buf
+
custom_pdf_matcher = CustomPdfMatcher()
filetype.add_type(custom_pdf_matcher)
diff --git a/cognee/infrastructure/files/utils/is_text_content.py b/cognee/infrastructure/files/utils/is_text_content.py
index dc323cd2a..83a961e89 100644
--- a/cognee/infrastructure/files/utils/is_text_content.py
+++ b/cognee/infrastructure/files/utils/is_text_content.py
@@ -1,24 +1,27 @@
def is_text_content(content):
"""Check if the content is text."""
# Check for null bytes
- if b'\0' in content:
+ if b"\0" in content:
return False
# Check for common text encodings (BOMs)
- if content.startswith((b'\xEF\xBB\xBF', # UTF-8
- b'\xFF\xFE', # UTF-16 LE
- b'\xFE\xFF', # UTF-16 BE
- b'\x00\x00\xFE\xFF', # UTF-32 LE
- b'\xFF\xFE\x00\x00', # UTF-32 BE
- )):
+ if content.startswith(
+ (
+ b"\xef\xbb\xbf", # UTF-8
+ b"\xff\xfe", # UTF-16 LE
+ b"\xfe\xff", # UTF-16 BE
+ b"\x00\x00\xfe\xff", # UTF-32 LE
+ b"\xff\xfe\x00\x00", # UTF-32 BE
+ )
+ ):
return True
# Check for ASCII characters
- if all(0x20 <= byte <= 0x7E or byte in (b'\n', b'\r', b'\t') for byte in content):
+ if all(0x20 <= byte <= 0x7E or byte in (b"\n", b"\r", b"\t") for byte in content):
return True
# Check for common line break characters
- if b'\n' in content or b'\r' in content:
+ if b"\n" in content or b"\r" in content:
return True
# If no obvious indicators found, assume it's a text file
diff --git a/cognee/infrastructure/llm/anthropic/adapter.py b/cognee/infrastructure/llm/anthropic/adapter.py
index 119b05eff..1fba732a0 100644
--- a/cognee/infrastructure/llm/anthropic/adapter.py
+++ b/cognee/infrastructure/llm/anthropic/adapter.py
@@ -10,34 +10,33 @@
class AnthropicAdapter(LLMInterface):
"""Adapter for Anthropic API"""
+
name = "Anthropic"
model: str
def __init__(self, model: str = None):
self.aclient = instructor.patch(
- create = anthropic.Anthropic().messages.create,
- mode = instructor.Mode.ANTHROPIC_TOOLS
+ create=anthropic.Anthropic().messages.create, mode=instructor.Mode.ANTHROPIC_TOOLS
)
self.model = model
async def acreate_structured_output(
- self,
- text_input: str,
- system_prompt: str,
- response_model: Type[BaseModel]
+ self, text_input: str, system_prompt: str, response_model: Type[BaseModel]
) -> BaseModel:
"""Generate a response from a user query."""
return await self.aclient(
- model = self.model,
- max_tokens = 4096,
- max_retries = 5,
- messages = [{
- "role": "user",
- "content": f"""Use the given format to extract information
+ model=self.model,
+ max_tokens=4096,
+ max_retries=5,
+ messages=[
+ {
+ "role": "user",
+ "content": f"""Use the given format to extract information
from the following input: {text_input}. {system_prompt}""",
- }],
- response_model = response_model,
+ }
+ ],
+ response_model=response_model,
)
def show_prompt(self, text_input: str, system_prompt: str) -> str:
@@ -50,6 +49,10 @@ def show_prompt(self, text_input: str, system_prompt: str) -> str:
system_prompt = read_query_prompt(system_prompt)
- formatted_prompt = f"""System Prompt:\n{system_prompt}\n\nUser Input:\n{text_input}\n""" if system_prompt else None
+ formatted_prompt = (
+ f"""System Prompt:\n{system_prompt}\n\nUser Input:\n{text_input}\n"""
+ if system_prompt
+ else None
+ )
return formatted_prompt
diff --git a/cognee/infrastructure/llm/config.py b/cognee/infrastructure/llm/config.py
index d148042be..67fc82683 100644
--- a/cognee/infrastructure/llm/config.py
+++ b/cognee/infrastructure/llm/config.py
@@ -2,6 +2,7 @@
from functools import lru_cache
from pydantic_settings import BaseSettings, SettingsConfigDict
+
class LLMConfig(BaseSettings):
llm_provider: str = "openai"
llm_model: str = "gpt-4o-mini"
@@ -12,7 +13,7 @@ class LLMConfig(BaseSettings):
llm_streaming: bool = False
transcription_model: str = "whisper-1"
- model_config = SettingsConfigDict(env_file = ".env", extra = "allow")
+ model_config = SettingsConfigDict(env_file=".env", extra="allow")
def to_dict(self) -> dict:
return {
@@ -23,9 +24,10 @@ def to_dict(self) -> dict:
"api_version": self.llm_api_version,
"temperature": self.llm_temperature,
"streaming": self.llm_streaming,
- "transcription_model": self.transcription_model
+ "transcription_model": self.transcription_model,
}
+
@lru_cache
def get_llm_config():
return LLMConfig()
diff --git a/cognee/infrastructure/llm/generic_llm_api/adapter.py b/cognee/infrastructure/llm/generic_llm_api/adapter.py
index 294af9256..a910c0780 100644
--- a/cognee/infrastructure/llm/generic_llm_api/adapter.py
+++ b/cognee/infrastructure/llm/generic_llm_api/adapter.py
@@ -1,82 +1,54 @@
-'''Adapter for Generic API LLM provider API'''
+"""Adapter for Generic API LLM provider API"""
+
import asyncio
from typing import List, Type
from pydantic import BaseModel
import instructor
-import openai
-
-from cognee.exceptions import InvalidValueError
from cognee.infrastructure.llm.llm_interface import LLMInterface
-from cognee.infrastructure.llm.prompts import read_query_prompt
-from cognee.shared.data_models import MonitoringTool
-from cognee.base_config import get_base_config
from cognee.infrastructure.llm.config import get_llm_config
+import litellm
class GenericAPIAdapter(LLMInterface):
- """Adapter for Generic API LLM provider API """
+ """Adapter for Generic API LLM provider API"""
+
name: str
model: str
api_key: str
- def __init__(self, api_endpoint, api_key: str, model: str, name: str):
+ def __init__(self, endpoint, api_key: str, model: str, name: str):
self.name = name
self.model = model
self.api_key = api_key
+ self.endpoint = endpoint
llm_config = get_llm_config()
- if llm_config.llm_provider == "groq":
- from groq import groq
- self.aclient = instructor.from_openai(
- client = groq.Groq(
- api_key = api_key,
- ),
- mode = instructor.Mode.MD_JSON
- )
- else:
- base_config = get_base_config()
-
- if base_config.monitoring_tool == MonitoringTool.LANGFUSE:
- from langfuse.openai import AsyncOpenAI
- elif base_config.monitoring_tool == MonitoringTool.LANGSMITH:
- from langsmith import wrappers
- from openai import AsyncOpenAI
- AsyncOpenAI = wrappers.wrap_openai(AsyncOpenAI())
- else:
- from openai import AsyncOpenAI
+ if llm_config.llm_provider == "ollama":
+ self.aclient = instructor.from_litellm(litellm.acompletion)
- self.aclient = instructor.patch(
- AsyncOpenAI(
- base_url = api_endpoint,
- api_key = api_key, # required, but unused
- ),
- mode = instructor.Mode.JSON,
- )
+ else:
+ self.aclient = instructor.from_litellm(litellm.acompletion)
- async def acreate_structured_output(self, text_input: str, system_prompt: str, response_model: Type[BaseModel]) -> BaseModel:
+ async def acreate_structured_output(
+ self, text_input: str, system_prompt: str, response_model: Type[BaseModel]
+ ) -> BaseModel:
"""Generate a response from a user query."""
return await self.aclient.chat.completions.create(
- model = self.model,
- messages = [
+ model=self.model,
+ messages=[
{
"role": "user",
"content": f"""Use the given format to
- extract information from the following input: {text_input}. """,
+ extract information from the following input: {text_input}. """,
+ },
+ {
+ "role": "system",
+ "content": system_prompt,
},
- {"role": "system", "content": system_prompt},
],
- response_model = response_model,
+ max_retries=5,
+ api_base=self.endpoint,
+ response_model=response_model,
)
-
- def show_prompt(self, text_input: str, system_prompt: str) -> str:
- """Format and display the prompt for a user query."""
- if not text_input:
- text_input = "No user input provided."
- if not system_prompt:
- raise InvalidValueError(message="No system prompt path provided.")
- system_prompt = read_query_prompt(system_prompt)
-
- formatted_prompt = f"""System Prompt:\n{system_prompt}\n\nUser Input:\n{text_input}\n""" if system_prompt else None
- return formatted_prompt
diff --git a/cognee/infrastructure/llm/get_llm_client.py b/cognee/infrastructure/llm/get_llm_client.py
index 9a23892f2..0f64014e3 100644
--- a/cognee/infrastructure/llm/get_llm_client.py
+++ b/cognee/infrastructure/llm/get_llm_client.py
@@ -1,9 +1,11 @@
"""Get the LLM client."""
+
from enum import Enum
from cognee.exceptions import InvalidValueError
from cognee.infrastructure.llm import get_llm_config
+
# Define an Enum for LLM Providers
class LLMProvider(Enum):
OPENAI = "openai"
@@ -11,6 +13,7 @@ class LLMProvider(Enum):
ANTHROPIC = "anthropic"
CUSTOM = "custom"
+
def get_llm_client():
"""Get the LLM client based on the configuration using Enums."""
llm_config = get_llm_config()
@@ -24,12 +27,12 @@ def get_llm_client():
from .openai.adapter import OpenAIAdapter
return OpenAIAdapter(
- api_key = llm_config.llm_api_key,
- endpoint = llm_config.llm_endpoint,
- api_version = llm_config.llm_api_version,
- model = llm_config.llm_model,
- transcription_model = llm_config.transcription_model,
- streaming = llm_config.llm_streaming,
+ api_key=llm_config.llm_api_key,
+ endpoint=llm_config.llm_endpoint,
+ api_version=llm_config.llm_api_version,
+ model=llm_config.llm_model,
+ transcription_model=llm_config.transcription_model,
+ streaming=llm_config.llm_streaming,
)
elif provider == LLMProvider.OLLAMA:
@@ -37,10 +40,14 @@ def get_llm_client():
raise InvalidValueError(message="LLM API key is not set.")
from .generic_llm_api.adapter import GenericAPIAdapter
- return GenericAPIAdapter(llm_config.llm_endpoint, llm_config.llm_api_key, llm_config.llm_model, "Ollama")
+
+ return GenericAPIAdapter(
+ llm_config.llm_endpoint, llm_config.llm_api_key, llm_config.llm_model, "Ollama"
+ )
elif provider == LLMProvider.ANTHROPIC:
from .anthropic.adapter import AnthropicAdapter
+
return AnthropicAdapter(llm_config.llm_model)
elif provider == LLMProvider.CUSTOM:
@@ -48,7 +55,10 @@ def get_llm_client():
raise InvalidValueError(message="LLM API key is not set.")
from .generic_llm_api.adapter import GenericAPIAdapter
- return GenericAPIAdapter(llm_config.llm_endpoint, llm_config.llm_api_key, llm_config.llm_model, "Custom")
+
+ return GenericAPIAdapter(
+ llm_config.llm_endpoint, llm_config.llm_api_key, llm_config.llm_model, "Custom"
+ )
else:
raise InvalidValueError(message=f"Unsupported LLM provider: {provider}")
diff --git a/cognee/infrastructure/llm/llm_interface.py b/cognee/infrastructure/llm/llm_interface.py
index 4a5bd79b6..dfcbecd29 100644
--- a/cognee/infrastructure/llm/llm_interface.py
+++ b/cognee/infrastructure/llm/llm_interface.py
@@ -1,18 +1,18 @@
-""" LLM Interface """
+"""LLM Interface"""
from typing import Type, Protocol
from abc import abstractmethod
from pydantic import BaseModel
from cognee.infrastructure.llm.prompts import read_query_prompt
+
class LLMInterface(Protocol):
- """ LLM Interface """
+ """LLM Interface"""
@abstractmethod
- async def acreate_structured_output(self,
- text_input: str,
- system_prompt: str,
- response_model: Type[BaseModel]) -> BaseModel:
+ async def acreate_structured_output(
+ self, text_input: str, system_prompt: str, response_model: Type[BaseModel]
+ ) -> BaseModel:
"""To get structured output, import/call this function"""
raise NotImplementedError
diff --git a/cognee/infrastructure/llm/openai/adapter.py b/cognee/infrastructure/llm/openai/adapter.py
index bb5af15f2..d45662380 100644
--- a/cognee/infrastructure/llm/openai/adapter.py
+++ b/cognee/infrastructure/llm/openai/adapter.py
@@ -16,6 +16,7 @@
if monitoring == MonitoringTool.LANGFUSE:
from langfuse.decorators import observe
+
class OpenAIAdapter(LLMInterface):
name = "OpenAI"
model: str
@@ -25,13 +26,13 @@ class OpenAIAdapter(LLMInterface):
"""Adapter for OpenAI's GPT-3, GPT=4 API"""
def __init__(
- self,
- api_key: str,
- endpoint: str,
- api_version: str,
- model: str,
- transcription_model: str,
- streaming: bool = False,
+ self,
+ api_key: str,
+ endpoint: str,
+ api_version: str,
+ model: str,
+ transcription_model: str,
+ streaming: bool = False,
):
self.aclient = instructor.from_litellm(litellm.acompletion)
self.client = instructor.from_litellm(litellm.completion)
@@ -41,25 +42,26 @@ def __init__(
self.endpoint = endpoint
self.api_version = api_version
self.streaming = streaming
- base_config = get_base_config()
-
-
- @observe(as_type='generation')
- async def acreate_structured_output(self, text_input: str, system_prompt: str,
- response_model: Type[BaseModel]) -> BaseModel:
+ @observe(as_type="generation")
+ async def acreate_structured_output(
+ self, text_input: str, system_prompt: str, response_model: Type[BaseModel]
+ ) -> BaseModel:
"""Generate a response from a user query."""
return await self.aclient.chat.completions.create(
model=self.model,
- messages=[{
- "role": "user",
- "content": f"""Use the given format to
+ messages=[
+ {
+ "role": "user",
+ "content": f"""Use the given format to
extract information from the following input: {text_input}. """,
- }, {
- "role": "system",
- "content": system_prompt,
- }],
+ },
+ {
+ "role": "system",
+ "content": system_prompt,
+ },
+ ],
api_key=self.api_key,
api_base=self.endpoint,
api_version=self.api_version,
@@ -68,20 +70,24 @@ async def acreate_structured_output(self, text_input: str, system_prompt: str,
)
@observe
- def create_structured_output(self, text_input: str, system_prompt: str,
- response_model: Type[BaseModel]) -> BaseModel:
+ def create_structured_output(
+ self, text_input: str, system_prompt: str, response_model: Type[BaseModel]
+ ) -> BaseModel:
"""Generate a response from a user query."""
return self.client.chat.completions.create(
model=self.model,
- messages=[{
- "role": "user",
- "content": f"""Use the given format to
+ messages=[
+ {
+ "role": "user",
+ "content": f"""Use the given format to
extract information from the following input: {text_input}. """,
- }, {
- "role": "system",
- "content": system_prompt,
- }],
+ },
+ {
+ "role": "system",
+ "content": system_prompt,
+ },
+ ],
api_key=self.api_key,
api_base=self.endpoint,
api_version=self.api_version,
@@ -111,24 +117,27 @@ def create_transcript(self, input):
def transcribe_image(self, input) -> BaseModel:
with open(input, "rb") as image_file:
- encoded_image = base64.b64encode(image_file.read()).decode('utf-8')
+ encoded_image = base64.b64encode(image_file.read()).decode("utf-8")
return litellm.completion(
model=self.model,
- messages=[{
- "role": "user",
- "content": [
- {
- "type": "text",
- "text": "What’s in this image?",
- }, {
- "type": "image_url",
- "image_url": {
- "url": f"data:image/jpeg;base64,{encoded_image}",
+ messages=[
+ {
+ "role": "user",
+ "content": [
+ {
+ "type": "text",
+ "text": "What’s in this image?",
},
- },
- ],
- }],
+ {
+ "type": "image_url",
+ "image_url": {
+ "url": f"data:image/jpeg;base64,{encoded_image}",
+ },
+ },
+ ],
+ }
+ ],
api_key=self.api_key,
api_base=self.endpoint,
api_version=self.api_version,
@@ -144,5 +153,9 @@ def show_prompt(self, text_input: str, system_prompt: str) -> str:
raise InvalidValueError(message="No system prompt path provided.")
system_prompt = read_query_prompt(system_prompt)
- formatted_prompt = f"""System Prompt:\n{system_prompt}\n\nUser Input:\n{text_input}\n""" if system_prompt else None
- return formatted_prompt
\ No newline at end of file
+ formatted_prompt = (
+ f"""System Prompt:\n{system_prompt}\n\nUser Input:\n{text_input}\n"""
+ if system_prompt
+ else None
+ )
+ return formatted_prompt
diff --git a/cognee/infrastructure/llm/prompts/read_query_prompt.py b/cognee/infrastructure/llm/prompts/read_query_prompt.py
index d9ea55acd..c1f58d77f 100644
--- a/cognee/infrastructure/llm/prompts/read_query_prompt.py
+++ b/cognee/infrastructure/llm/prompts/read_query_prompt.py
@@ -2,12 +2,13 @@
import logging
from cognee.root_dir import get_absolute_path
+
def read_query_prompt(prompt_file_name: str):
"""Read a query prompt from a file."""
try:
file_path = path.join(get_absolute_path("./infrastructure/llm/prompts"), prompt_file_name)
- with open(file_path, "r", encoding = "utf-8") as file:
+ with open(file_path, "r", encoding="utf-8") as file:
return file.read()
except FileNotFoundError:
logging.error(f"Error: Prompt file not found. Attempted to read: %s {file_path}")
diff --git a/cognee/infrastructure/llm/prompts/render_prompt.py b/cognee/infrastructure/llm/prompts/render_prompt.py
index 03f1f7cb2..756fa671b 100644
--- a/cognee/infrastructure/llm/prompts/render_prompt.py
+++ b/cognee/infrastructure/llm/prompts/render_prompt.py
@@ -1,7 +1,8 @@
from jinja2 import Environment, FileSystemLoader, select_autoescape
from cognee.root_dir import get_absolute_path
-def render_prompt(filename: str, context: dict) -> str:
+
+def render_prompt(filename: str, context: dict) -> str:
"""Render a Jinja2 template asynchronously.
:param filename: The name of the template file to render.
:param context: The context to render the template with.
@@ -12,8 +13,8 @@ def render_prompt(filename: str, context: dict) -> str:
# Initialize the Jinja2 environment to load templates from the filesystem
env = Environment(
- loader = FileSystemLoader(base_directory),
- autoescape = select_autoescape(["html", "xml", "txt"])
+ loader=FileSystemLoader(base_directory),
+ autoescape=select_autoescape(["html", "xml", "txt"]),
)
# Load the template by name
diff --git a/cognee/infrastructure/pipeline/models/Operation.py b/cognee/infrastructure/pipeline/models/Operation.py
index 62eb74c44..b0e5624cb 100644
--- a/cognee/infrastructure/pipeline/models/Operation.py
+++ b/cognee/infrastructure/pipeline/models/Operation.py
@@ -3,10 +3,12 @@
from sqlalchemy import Column, DateTime, ForeignKey, Enum, JSON
from cognee.infrastructure.databases.relational import Base, UUID
+
class OperationType(Enum):
MERGE_DATA = "MERGE_DATA"
APPEND_DATA = "APPEND_DATA"
+
class OperationStatus(Enum):
STARTED = "OPERATION_STARTED"
IN_PROGRESS = "OPERATION_IN_PROGRESS"
@@ -14,14 +16,15 @@ class OperationStatus(Enum):
ERROR = "OPERATION_ERROR"
CANCELLED = "OPERATION_CANCELLED"
+
class Operation(Base):
__tablename__ = "operation"
- id = Column(UUID, primary_key = True)
+ id = Column(UUID, primary_key=True)
status = Column(Enum(OperationStatus))
operation_type = Column(Enum(OperationType))
data_id = Column(UUID, ForeignKey("data.id"))
- meta_data: Mapped[dict] = MappedColumn(type_ = JSON)
+ meta_data: Mapped[dict] = MappedColumn(type_=JSON)
- created_at = Column(DateTime, default = datetime.now(timezone.utc))
+ created_at = Column(DateTime, default=datetime.now(timezone.utc))
diff --git a/cognee/modules/chunking/TextChunker.py b/cognee/modules/chunking/TextChunker.py
index 64c7aae5c..7bb8a1c1c 100644
--- a/cognee/modules/chunking/TextChunker.py
+++ b/cognee/modules/chunking/TextChunker.py
@@ -3,7 +3,8 @@
from .models.DocumentChunk import DocumentChunk
from cognee.tasks.chunks import chunk_by_paragraph
-class TextChunker():
+
+class TextChunker:
document = None
max_chunk_size: int
@@ -21,7 +22,7 @@ def read(self):
for chunk_data in chunk_by_paragraph(
content_text,
self.max_chunk_size,
- batch_paragraphs = True,
+ batch_paragraphs=True,
):
if self.chunk_size + chunk_data["word_count"] <= self.max_chunk_size:
paragraph_chunks.append(chunk_data)
@@ -29,17 +30,17 @@ def read(self):
else:
if len(paragraph_chunks) == 0:
yield DocumentChunk(
- id = chunk_data["chunk_id"],
- text = chunk_data["text"],
- word_count = chunk_data["word_count"],
- is_part_of = self.document,
- chunk_index = self.chunk_index,
- cut_type = chunk_data["cut_type"],
- contains = [],
- _metadata = {
+ id=chunk_data["chunk_id"],
+ text=chunk_data["text"],
+ word_count=chunk_data["word_count"],
+ is_part_of=self.document,
+ chunk_index=self.chunk_index,
+ cut_type=chunk_data["cut_type"],
+ contains=[],
+ _metadata={
"index_fields": ["text"],
- "metadata_id": self.document.metadata_id
- }
+ "metadata_id": self.document.metadata_id,
+ },
)
paragraph_chunks = []
self.chunk_size = 0
@@ -47,17 +48,19 @@ def read(self):
chunk_text = " ".join(chunk["text"] for chunk in paragraph_chunks)
try:
yield DocumentChunk(
- id = uuid5(NAMESPACE_OID, f"{str(self.document.id)}-{self.chunk_index}"),
- text = chunk_text,
- word_count = self.chunk_size,
- is_part_of = self.document,
- chunk_index = self.chunk_index,
- cut_type = paragraph_chunks[len(paragraph_chunks) - 1]["cut_type"],
- contains = [],
- _metadata = {
+ id=uuid5(
+ NAMESPACE_OID, f"{str(self.document.id)}-{self.chunk_index}"
+ ),
+ text=chunk_text,
+ word_count=self.chunk_size,
+ is_part_of=self.document,
+ chunk_index=self.chunk_index,
+ cut_type=paragraph_chunks[len(paragraph_chunks) - 1]["cut_type"],
+ contains=[],
+ _metadata={
"index_fields": ["text"],
- "metadata_id": self.document.metadata_id
- }
+ "metadata_id": self.document.metadata_id,
+ },
)
except Exception as e:
print(e)
@@ -69,17 +72,14 @@ def read(self):
if len(paragraph_chunks) > 0:
try:
yield DocumentChunk(
- id = uuid5(NAMESPACE_OID, f"{str(self.document.id)}-{self.chunk_index}"),
- text = " ".join(chunk["text"] for chunk in paragraph_chunks),
- word_count = self.chunk_size,
- is_part_of = self.document,
- chunk_index = self.chunk_index,
- cut_type = paragraph_chunks[len(paragraph_chunks) - 1]["cut_type"],
- contains = [],
- _metadata = {
- "index_fields": ["text"],
- "metadata_id": self.document.metadata_id
- }
+ id=uuid5(NAMESPACE_OID, f"{str(self.document.id)}-{self.chunk_index}"),
+ text=" ".join(chunk["text"] for chunk in paragraph_chunks),
+ word_count=self.chunk_size,
+ is_part_of=self.document,
+ chunk_index=self.chunk_index,
+ cut_type=paragraph_chunks[len(paragraph_chunks) - 1]["cut_type"],
+ contains=[],
+ _metadata={"index_fields": ["text"], "metadata_id": self.document.metadata_id},
)
except Exception as e:
print(e)
diff --git a/cognee/modules/chunking/models/DocumentChunk.py b/cognee/modules/chunking/models/DocumentChunk.py
index 8729596df..4920e9b06 100644
--- a/cognee/modules/chunking/models/DocumentChunk.py
+++ b/cognee/modules/chunking/models/DocumentChunk.py
@@ -14,7 +14,4 @@ class DocumentChunk(DataPoint):
is_part_of: Document
contains: List[Entity] = None
- _metadata: dict = {
- "index_fields": ["text"],
- "type": "DocumentChunk"
- }
+ _metadata: dict = {"index_fields": ["text"], "type": "DocumentChunk"}
diff --git a/cognee/modules/cognify/config.py b/cognee/modules/cognify/config.py
index ffc54856c..d40410bfc 100644
--- a/cognee/modules/cognify/config.py
+++ b/cognee/modules/cognify/config.py
@@ -2,11 +2,12 @@
from pydantic_settings import BaseSettings, SettingsConfigDict
from cognee.shared.data_models import DefaultContentPrediction, SummarizedContent
+
class CognifyConfig(BaseSettings):
classification_model: object = DefaultContentPrediction
summarization_model: object = SummarizedContent
- model_config = SettingsConfigDict(env_file = ".env", extra = "allow")
+ model_config = SettingsConfigDict(env_file=".env", extra="allow")
def to_dict(self) -> dict:
return {
@@ -14,6 +15,7 @@ def to_dict(self) -> dict:
"summarization_model": self.summarization_model,
}
+
@lru_cache
def get_cognify_config():
return CognifyConfig()
diff --git a/cognee/modules/data/deletion/prune_data.py b/cognee/modules/data/deletion/prune_data.py
index cef0c5ffa..12a516329 100644
--- a/cognee/modules/data/deletion/prune_data.py
+++ b/cognee/modules/data/deletion/prune_data.py
@@ -1,6 +1,7 @@
from cognee.base_config import get_base_config
from cognee.infrastructure.files.storage import LocalStorage
+
async def prune_data():
base_config = get_base_config()
data_root_directory = base_config.data_root_directory
diff --git a/cognee/modules/data/deletion/prune_system.py b/cognee/modules/data/deletion/prune_system.py
index 38c625558..055d69b55 100644
--- a/cognee/modules/data/deletion/prune_system.py
+++ b/cognee/modules/data/deletion/prune_system.py
@@ -2,7 +2,8 @@
from cognee.infrastructure.databases.graph.get_graph_engine import get_graph_engine
from cognee.infrastructure.databases.relational import get_relational_engine
-async def prune_system(graph = True, vector = True, metadata = False):
+
+async def prune_system(graph=True, vector=True, metadata=False):
if graph:
graph_engine = await get_graph_engine()
await graph_engine.delete_graph()
diff --git a/cognee/modules/data/exceptions/__init__.py b/cognee/modules/data/exceptions/__init__.py
index 6f74c627e..16c363bed 100644
--- a/cognee/modules/data/exceptions/__init__.py
+++ b/cognee/modules/data/exceptions/__init__.py
@@ -7,4 +7,4 @@
from .exceptions import (
UnstructuredLibraryImportError,
UnauthorizedDataAccessError,
-)
\ No newline at end of file
+)
diff --git a/cognee/modules/data/exceptions/exceptions.py b/cognee/modules/data/exceptions/exceptions.py
index 5117f3cac..cebc1efe5 100644
--- a/cognee/modules/data/exceptions/exceptions.py
+++ b/cognee/modules/data/exceptions/exceptions.py
@@ -1,20 +1,22 @@
from cognee.exceptions import CogneeApiError
from fastapi import status
+
class UnstructuredLibraryImportError(CogneeApiError):
def __init__(
- self,
- message: str = "Import error. Unstructured library is not installed.",
- name: str = "UnstructuredModuleImportError",
- status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
+ self,
+ message: str = "Import error. Unstructured library is not installed.",
+ name: str = "UnstructuredModuleImportError",
+ status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
):
super().__init__(message, name, status_code)
+
class UnauthorizedDataAccessError(CogneeApiError):
def __init__(
- self,
- message: str = "User does not have permission to access this data.",
- name: str = "UnauthorizedDataAccessError",
- status_code=status.HTTP_401_UNAUTHORIZED,
+ self,
+ message: str = "User does not have permission to access this data.",
+ name: str = "UnauthorizedDataAccessError",
+ status_code=status.HTTP_401_UNAUTHORIZED,
):
- super().__init__(message, name, status_code)
\ No newline at end of file
+ super().__init__(message, name, status_code)
diff --git a/cognee/modules/data/extraction/extract_summary.py b/cognee/modules/data/extraction/extract_summary.py
index 0d3feef95..27b3997f3 100644
--- a/cognee/modules/data/extraction/extract_summary.py
+++ b/cognee/modules/data/extraction/extract_summary.py
@@ -13,15 +13,17 @@
logger = logging.getLogger("extract_summary")
+
async def extract_summary(content: str, response_model: Type[BaseModel]):
llm_client = get_llm_client()
system_prompt = read_query_prompt("summarize_content.txt")
llm_output = await llm_client.acreate_structured_output(content, system_prompt, response_model)
-
+
return llm_output
+
async def extract_code_summary(content: str):
enable_mocking = os.getenv("MOCK_CODE_SUMMARY", "false")
if isinstance(enable_mocking, bool):
diff --git a/cognee/modules/data/extraction/extract_topics_naive.py b/cognee/modules/data/extraction/extract_topics_naive.py
index a8487224f..d1323e24b 100644
--- a/cognee/modules/data/extraction/extract_topics_naive.py
+++ b/cognee/modules/data/extraction/extract_topics_naive.py
@@ -6,6 +6,7 @@
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.decomposition import TruncatedSVD
+
def extract_topics(text: str):
sentences = sent_tokenize(text)
@@ -18,24 +19,19 @@ def extract_topics(text: str):
lemmatizer = WordNetLemmatizer()
base_notation_sentences = [lemmatizer.lemmatize(sentence) for sentence in sentences]
- tf_vectorizer = TfidfVectorizer(tokenizer = word_tokenize, token_pattern = None)
+ tf_vectorizer = TfidfVectorizer(tokenizer=word_tokenize, token_pattern=None)
transformed_corpus = tf_vectorizer.fit_transform(base_notation_sentences)
- svd = TruncatedSVD(n_components = 10)
+ svd = TruncatedSVD(n_components=10)
svd_corpus = svd.fit(transformed_corpus)
- feature_scores = dict(
- zip(
- tf_vectorizer.vocabulary_,
- svd_corpus.components_[0]
- )
- )
+ feature_scores = dict(zip(tf_vectorizer.vocabulary_, svd_corpus.components_[0]))
topics = sorted(
feature_scores,
# key = feature_scores.get,
- key = lambda x: transformed_corpus[0, tf_vectorizer.vocabulary_[x]],
- reverse = True,
+ key=lambda x: transformed_corpus[0, tf_vectorizer.vocabulary_[x]],
+ reverse=True,
)[:10]
return topics
@@ -45,6 +41,7 @@ def clean_text(text: str):
text = re.sub(r"[ \t]{2,}|[\n\r]", " ", text.lower())
return re.sub(r"[`\"'.,;!?…]", "", text).strip()
+
def remove_stop_words(text: str):
try:
stopwords.ensure_loaded()
@@ -54,7 +51,7 @@ def remove_stop_words(text: str):
stop_words = set(stopwords.words("english"))
text = text.split()
- text = [word for word in text if not word in stop_words]
+ text = [word for word in text if word not in stop_words]
return " ".join(text)
diff --git a/cognee/modules/data/extraction/knowledge_graph/add_model_class_to_graph.py b/cognee/modules/data/extraction/knowledge_graph/add_model_class_to_graph.py
index 08691cbf2..9da226541 100644
--- a/cognee/modules/data/extraction/knowledge_graph/add_model_class_to_graph.py
+++ b/cognee/modules/data/extraction/knowledge_graph/add_model_class_to_graph.py
@@ -2,6 +2,7 @@
from pydantic import BaseModel
from cognee.infrastructure.databases.graph.graph_db_interface import GraphDBInterface
+
async def add_model_class_to_graph(
model_class: Type[BaseModel],
graph: GraphDBInterface,
@@ -13,7 +14,7 @@ async def add_model_class_to_graph(
if await graph.extract_node(model_name):
return
- await graph.add_node(model_name, dict(type = "model"))
+ await graph.add_node(model_name, dict(type="model"))
if parent and relationship:
await graph.add_edge(
@@ -21,9 +22,9 @@ async def add_model_class_to_graph(
model_name,
relationship,
dict(
- relationship_name = relationship,
- source_node_id = parent,
- target_node_id = model_name,
+ relationship_name=relationship,
+ source_node_id=parent,
+ target_node_id=model_name,
),
)
@@ -36,7 +37,7 @@ async def add_model_class_to_graph(
if hasattr(field_type, "model_fields"): # Check if field type is a Pydantic model
await add_model_class_to_graph(field_type, graph, model_name, field_name)
- elif get_origin(field.annotation) == list:
+ elif isinstance(get_origin(field.annotation), list):
list_types = get_args(field_type)
for item_type in list_types:
await add_model_class_to_graph(item_type, graph, model_name, field_name)
@@ -45,26 +46,26 @@ async def add_model_class_to_graph(
if hasattr(item_type, "model_fields"):
await add_model_class_to_graph(item_type, graph, model_name, field_name)
else:
- await graph.add_node(str(item_type), dict(type = "value"))
+ await graph.add_node(str(item_type), dict(type="value"))
await graph.add_edge(
model_name,
str(item_type),
field_name,
dict(
- relationship_name = field_name,
- source_node_id = model_name,
- target_node_id = str(item_type),
+ relationship_name=field_name,
+ source_node_id=model_name,
+ target_node_id=str(item_type),
),
)
else:
- await graph.add_node(str(field_type), dict(type = "value"))
+ await graph.add_node(str(field_type), dict(type="value"))
await graph.add_edge(
model_name,
str(field_type),
field_name,
dict(
- relationship_name = field_name,
- source_node_id = model_name,
- target_node_id = str(field_type),
+ relationship_name=field_name,
+ source_node_id=model_name,
+ target_node_id=str(field_type),
),
)
diff --git a/cognee/modules/data/extraction/knowledge_graph/extract_content_graph.py b/cognee/modules/data/extraction/knowledge_graph/extract_content_graph.py
index e52f3992a..49af7351a 100644
--- a/cognee/modules/data/extraction/knowledge_graph/extract_content_graph.py
+++ b/cognee/modules/data/extraction/knowledge_graph/extract_content_graph.py
@@ -3,10 +3,13 @@
from cognee.infrastructure.llm.get_llm_client import get_llm_client
from cognee.infrastructure.llm.prompts import render_prompt
+
async def extract_content_graph(content: str, response_model: Type[BaseModel]):
llm_client = get_llm_client()
system_prompt = render_prompt("generate_graph_prompt.txt", {})
- content_graph = await llm_client.acreate_structured_output(content, system_prompt, response_model)
+ content_graph = await llm_client.acreate_structured_output(
+ content, system_prompt, response_model
+ )
return content_graph
diff --git a/cognee/modules/data/methods/__init__.py b/cognee/modules/data/methods/__init__.py
index 34f943359..c32db1d2f 100644
--- a/cognee/modules/data/methods/__init__.py
+++ b/cognee/modules/data/methods/__init__.py
@@ -10,4 +10,4 @@
# Delete
from .delete_dataset import delete_dataset
-from .delete_data import delete_data
\ No newline at end of file
+from .delete_data import delete_data
diff --git a/cognee/modules/data/methods/create_dataset.py b/cognee/modules/data/methods/create_dataset.py
index aa9c645f7..be4ea8792 100644
--- a/cognee/modules/data/methods/create_dataset.py
+++ b/cognee/modules/data/methods/create_dataset.py
@@ -4,20 +4,19 @@
from sqlalchemy.orm import joinedload
from cognee.modules.data.models import Dataset
+
async def create_dataset(dataset_name: str, owner_id: UUID, session: AsyncSession) -> Dataset:
- dataset = (await session.scalars(
- select(Dataset)\
- .options(joinedload(Dataset.data))\
+ dataset = (
+ await session.scalars(
+ select(Dataset)
+ .options(joinedload(Dataset.data))
.filter(Dataset.name == dataset_name)
.filter(Dataset.owner_id == owner_id)
- )).first()
+ )
+ ).first()
if dataset is None:
- dataset = Dataset(
- id = uuid5(NAMESPACE_OID, dataset_name),
- name = dataset_name,
- data = []
- )
+ dataset = Dataset(id=uuid5(NAMESPACE_OID, dataset_name), name=dataset_name, data=[])
dataset.owner_id = owner_id
session.add(dataset)
diff --git a/cognee/modules/data/methods/delete_data.py b/cognee/modules/data/methods/delete_data.py
index 65abe714a..2d87d73a5 100644
--- a/cognee/modules/data/methods/delete_data.py
+++ b/cognee/modules/data/methods/delete_data.py
@@ -6,14 +6,16 @@
async def delete_data(data: Data):
"""Delete a data record from the database.
- Args:
- data (Data): The data object to be deleted.
+ Args:
+ data (Data): The data object to be deleted.
- Raises:
- ValueError: If the data object is invalid.
+ Raises:
+ ValueError: If the data object is invalid.
"""
- if not hasattr(data, '__tablename__'):
- raise InvalidAttributeError(message="The provided data object is missing the required '__tablename__' attribute.")
+ if not hasattr(data, "__tablename__"):
+ raise InvalidAttributeError(
+ message="The provided data object is missing the required '__tablename__' attribute."
+ )
db_engine = get_relational_engine()
diff --git a/cognee/modules/data/methods/delete_dataset.py b/cognee/modules/data/methods/delete_dataset.py
index 96a2e7d71..ff20ff9e7 100644
--- a/cognee/modules/data/methods/delete_dataset.py
+++ b/cognee/modules/data/methods/delete_dataset.py
@@ -1,6 +1,7 @@
from cognee.modules.data.models import Dataset
from cognee.infrastructure.databases.relational import get_relational_engine
+
async def delete_dataset(dataset: Dataset):
db_engine = get_relational_engine()
diff --git a/cognee/modules/data/methods/get_data.py b/cognee/modules/data/methods/get_data.py
index d7daff29b..d1fff2a0a 100644
--- a/cognee/modules/data/methods/get_data.py
+++ b/cognee/modules/data/methods/get_data.py
@@ -4,15 +4,16 @@
from ..exceptions import UnauthorizedDataAccessError
from ..models import Data
+
async def get_data(user_id: UUID, data_id: UUID) -> Optional[Data]:
"""Retrieve data by ID.
- Args:
- user_id (UUID): user ID
- data_id (UUID): ID of the data to retrieve
+ Args:
+ user_id (UUID): user ID
+ data_id (UUID): ID of the data to retrieve
- Returns:
- Optional[Data]: The requested data object if found, None otherwise
+ Returns:
+ Optional[Data]: The requested data object if found, None otherwise
"""
db_engine = get_relational_engine()
@@ -20,6 +21,8 @@ async def get_data(user_id: UUID, data_id: UUID) -> Optional[Data]:
data = await session.get(Data, data_id)
if data and data.owner_id != user_id:
- raise UnauthorizedDataAccessError(message=f"User {user_id} is not authorized to access data {data_id}")
+ raise UnauthorizedDataAccessError(
+ message=f"User {user_id} is not authorized to access data {data_id}"
+ )
- return data
\ No newline at end of file
+ return data
diff --git a/cognee/modules/data/methods/get_dataset.py b/cognee/modules/data/methods/get_dataset.py
index 9f46fa223..e5f7e9923 100644
--- a/cognee/modules/data/methods/get_dataset.py
+++ b/cognee/modules/data/methods/get_dataset.py
@@ -3,6 +3,7 @@
from cognee.infrastructure.databases.relational import get_relational_engine
from ..models import Dataset
+
async def get_dataset(user_id: UUID, dataset_id: UUID) -> Optional[Dataset]:
db_engine = get_relational_engine()
diff --git a/cognee/modules/data/methods/get_dataset_data.py b/cognee/modules/data/methods/get_dataset_data.py
index f1b969b2d..80669fddb 100644
--- a/cognee/modules/data/methods/get_dataset_data.py
+++ b/cognee/modules/data/methods/get_dataset_data.py
@@ -3,6 +3,7 @@
from cognee.modules.data.models import Data, Dataset
from cognee.infrastructure.databases.relational import get_relational_engine
+
async def get_dataset_data(dataset_id: UUID) -> list[Data]:
db_engine = get_relational_engine()
diff --git a/cognee/modules/data/methods/get_datasets.py b/cognee/modules/data/methods/get_datasets.py
index 040c4c590..f6db585df 100644
--- a/cognee/modules/data/methods/get_datasets.py
+++ b/cognee/modules/data/methods/get_datasets.py
@@ -3,12 +3,13 @@
from cognee.infrastructure.databases.relational import get_relational_engine
from ..models import Dataset
+
async def get_datasets(user_id: UUID) -> list[Dataset]:
db_engine = get_relational_engine()
async with db_engine.get_async_session() as session:
- datasets = (await session.scalars(
- select(Dataset).filter(Dataset.owner_id == user_id)
- )).all()
+ datasets = (
+ await session.scalars(select(Dataset).filter(Dataset.owner_id == user_id))
+ ).all()
return datasets
diff --git a/cognee/modules/data/methods/get_datasets_by_name.py b/cognee/modules/data/methods/get_datasets_by_name.py
index f562268ce..be0613792 100644
--- a/cognee/modules/data/methods/get_datasets_by_name.py
+++ b/cognee/modules/data/methods/get_datasets_by_name.py
@@ -3,16 +3,19 @@
from cognee.infrastructure.databases.relational import get_relational_engine
from ..models import Dataset
+
async def get_datasets_by_name(dataset_names: list[str], user_id: UUID) -> list[Dataset]:
db_engine = get_relational_engine()
async with db_engine.get_async_session() as session:
if isinstance(dataset_names, str):
dataset_names = [dataset_names]
- datasets = (await session.scalars(
- select(Dataset)
+ datasets = (
+ await session.scalars(
+ select(Dataset)
.filter(Dataset.owner_id == user_id)
.filter(Dataset.name.in_(dataset_names))
- )).all()
+ )
+ ).all()
return datasets
diff --git a/cognee/modules/data/models/Data.py b/cognee/modules/data/models/Data.py
index e24bc7c5c..cf8918db7 100644
--- a/cognee/modules/data/models/Data.py
+++ b/cognee/modules/data/models/Data.py
@@ -9,6 +9,7 @@
from .DatasetData import DatasetData
from .Metadata import Metadata
+
class Data(Base):
__tablename__ = "data"
@@ -20,12 +21,8 @@ class Data(Base):
raw_data_location = Column(String)
owner_id = Column(UUID, index=True)
content_hash = Column(String)
- created_at = Column(
- DateTime(timezone=True), default=lambda: datetime.now(timezone.utc)
- )
- updated_at = Column(
- DateTime(timezone=True), onupdate=lambda: datetime.now(timezone.utc)
- )
+ created_at = Column(DateTime(timezone=True), default=lambda: datetime.now(timezone.utc))
+ updated_at = Column(DateTime(timezone=True), onupdate=lambda: datetime.now(timezone.utc))
datasets = relationship(
"Dataset",
@@ -42,7 +39,6 @@ class Data(Base):
cascade="all, delete",
)
-
def to_json(self) -> dict:
return {
"id": str(self.id),
diff --git a/cognee/modules/data/models/Dataset.py b/cognee/modules/data/models/Dataset.py
index f7078b8f1..dc2a81a20 100644
--- a/cognee/modules/data/models/Dataset.py
+++ b/cognee/modules/data/models/Dataset.py
@@ -6,24 +6,25 @@
from cognee.infrastructure.databases.relational import Base
from .DatasetData import DatasetData
+
class Dataset(Base):
__tablename__ = "datasets"
- id = Column(UUID, primary_key = True, default = uuid4)
+ id = Column(UUID, primary_key=True, default=uuid4)
name = Column(Text)
- created_at = Column(DateTime(timezone = True), default = lambda: datetime.now(timezone.utc))
- updated_at = Column(DateTime(timezone = True), onupdate = lambda: datetime.now(timezone.utc))
+ created_at = Column(DateTime(timezone=True), default=lambda: datetime.now(timezone.utc))
+ updated_at = Column(DateTime(timezone=True), onupdate=lambda: datetime.now(timezone.utc))
- owner_id = Column(UUID, index = True)
+ owner_id = Column(UUID, index=True)
data: Mapped[List["Data"]] = relationship(
"Data",
- secondary = DatasetData.__tablename__,
- back_populates = "datasets",
- lazy = "noload",
- cascade="all, delete"
+ secondary=DatasetData.__tablename__,
+ back_populates="datasets",
+ lazy="noload",
+ cascade="all, delete",
)
def to_json(self) -> dict:
@@ -33,5 +34,5 @@ def to_json(self) -> dict:
"createdAt": self.created_at.isoformat(),
"updatedAt": self.updated_at.isoformat() if self.updated_at else None,
"ownerId": str(self.owner_id),
- "data": [data.to_json() for data in self.data]
+ "data": [data.to_json() for data in self.data],
}
diff --git a/cognee/modules/data/models/DatasetData.py b/cognee/modules/data/models/DatasetData.py
index a35c120eb..209771581 100644
--- a/cognee/modules/data/models/DatasetData.py
+++ b/cognee/modules/data/models/DatasetData.py
@@ -2,10 +2,11 @@
from sqlalchemy import Column, DateTime, ForeignKey, UUID
from cognee.infrastructure.databases.relational import Base
+
class DatasetData(Base):
__tablename__ = "dataset_data"
- created_at = Column(DateTime(timezone = True), default = lambda: datetime.now(timezone.utc))
+ created_at = Column(DateTime(timezone=True), default=lambda: datetime.now(timezone.utc))
- dataset_id = Column(UUID, ForeignKey("datasets.id", ondelete="CASCADE"), primary_key = True)
- data_id = Column(UUID, ForeignKey("data.id", ondelete="CASCADE"), primary_key = True)
+ dataset_id = Column(UUID, ForeignKey("datasets.id", ondelete="CASCADE"), primary_key=True)
+ data_id = Column(UUID, ForeignKey("data.id", ondelete="CASCADE"), primary_key=True)
diff --git a/cognee/modules/data/models/Metadata.py b/cognee/modules/data/models/Metadata.py
index 3ab30b38d..ab41d94be 100644
--- a/cognee/modules/data/models/Metadata.py
+++ b/cognee/modules/data/models/Metadata.py
@@ -14,13 +14,8 @@ class Metadata(Base):
metadata_repr = Column(String)
metadata_source = Column(String)
- created_at = Column(
- DateTime(timezone=True), default=lambda: datetime.now(timezone.utc)
- )
- updated_at = Column(
- DateTime(timezone=True), onupdate=lambda: datetime.now(timezone.utc)
- )
+ created_at = Column(DateTime(timezone=True), default=lambda: datetime.now(timezone.utc))
+ updated_at = Column(DateTime(timezone=True), onupdate=lambda: datetime.now(timezone.utc))
- data_id = Column(UUID, ForeignKey("data.id", ondelete="CASCADE"), primary_key = False)
+ data_id = Column(UUID, ForeignKey("data.id", ondelete="CASCADE"), primary_key=False)
data = relationship("Data", back_populates="metadata_relationship")
-
diff --git a/cognee/modules/data/operations/detect_language.py b/cognee/modules/data/operations/detect_language.py
index e82675736..27fb853c6 100644
--- a/cognee/modules/data/operations/detect_language.py
+++ b/cognee/modules/data/operations/detect_language.py
@@ -2,6 +2,7 @@
logger = logging.getLogger(__name__)
+
async def detect_language(text: str):
"""
Detect the language of the given text and return its ISO 639-1 language code.
@@ -14,6 +15,7 @@ async def detect_language(text: str):
"""
from langdetect import detect, LangDetectException
+
# Trim the text to the first 100 characters
trimmed_text = text[:100]
diff --git a/cognee/modules/data/operations/get_metadata.py b/cognee/modules/data/operations/get_metadata.py
index 26637e383..f827c47c3 100644
--- a/cognee/modules/data/operations/get_metadata.py
+++ b/cognee/modules/data/operations/get_metadata.py
@@ -15,5 +15,3 @@ async def get_metadata(metadata_id: UUID) -> Metadata:
metadata = await session.get(Metadata, metadata_id)
return metadata
-
-
diff --git a/cognee/modules/data/operations/translate_text.py b/cognee/modules/data/operations/translate_text.py
index d8c27e42a..d9a7c9505 100644
--- a/cognee/modules/data/operations/translate_text.py
+++ b/cognee/modules/data/operations/translate_text.py
@@ -4,7 +4,10 @@
logger = logging.getLogger(__name__)
-async def translate_text(text, source_language: str = "sr", target_language: str = "en", region_name = "eu-west-1"):
+
+async def translate_text(
+ text, source_language: str = "sr", target_language: str = "en", region_name="eu-west-1"
+):
"""
Translate text from source language to target language using AWS Translate.
Parameters:
@@ -26,11 +29,11 @@ async def translate_text(text, source_language: str = "sr", target_language: str
raise InvalidValueError(message="Source and target language codes are required.")
try:
- translate = boto3.client(service_name = "translate", region_name = region_name, use_ssl = True)
+ translate = boto3.client(service_name="translate", region_name=region_name, use_ssl=True)
result = translate.translate_text(
- Text = text,
- SourceLanguageCode = source_language,
- TargetLanguageCode = target_language,
+ Text=text,
+ SourceLanguageCode=source_language,
+ TargetLanguageCode=target_language,
)
yield result.get("TranslatedText", "No translation found.")
diff --git a/cognee/modules/data/operations/write_metadata.py b/cognee/modules/data/operations/write_metadata.py
index 67c8c0e45..3c2c839c2 100644
--- a/cognee/modules/data/operations/write_metadata.py
+++ b/cognee/modules/data/operations/write_metadata.py
@@ -11,11 +11,12 @@
from ..models.Metadata import Metadata
-async def write_metadata(data_item: Union[BinaryIO, str, Any], data_id: UUID, file_metadata: FileMetadata) -> UUID:
+async def write_metadata(
+ data_item: Union[BinaryIO, str, Any], data_id: UUID, file_metadata: FileMetadata
+) -> UUID:
metadata_dict = get_metadata_dict(data_item, file_metadata)
db_engine = get_relational_engine()
async with db_engine.get_async_session() as session:
-
metadata = (
await session.execute(select(Metadata).filter(Metadata.data_id == data_id))
).scalar_one_or_none()
@@ -29,14 +30,13 @@ async def write_metadata(data_item: Union[BinaryIO, str, Any], data_id: UUID, fi
id=data_id,
metadata_repr=json.dumps(metadata_dict),
metadata_source=parse_type(type(data_item)),
- data_id=data_id
+ data_id=data_id,
)
session.add(metadata)
await session.commit()
-
def parse_type(type_: Any) -> str:
pattern = r".+'([\w_\.]+)'"
match = re.search(pattern, str(type_))
@@ -46,11 +46,13 @@ def parse_type(type_: Any) -> str:
raise Exception(f"type: {type_} could not be parsed")
-def get_metadata_dict(data_item: Union[BinaryIO, str, Any], file_metadata: FileMetadata) -> dict[str, Any]:
+def get_metadata_dict(
+ data_item: Union[BinaryIO, str, Any], file_metadata: FileMetadata
+) -> dict[str, Any]:
if isinstance(data_item, str):
- return(file_metadata)
+ return file_metadata
elif isinstance(data_item, BinaryIO):
- return(file_metadata)
+ return file_metadata
elif hasattr(data_item, "dict") and inspect.ismethod(getattr(data_item, "dict")):
return {**file_metadata, **data_item.dict()}
else:
diff --git a/cognee/modules/data/processing/document_types/AudioDocument.py b/cognee/modules/data/processing/document_types/AudioDocument.py
index 268338703..a33d4e7fc 100644
--- a/cognee/modules/data/processing/document_types/AudioDocument.py
+++ b/cognee/modules/data/processing/document_types/AudioDocument.py
@@ -2,19 +2,20 @@
from .Document import Document
from .ChunkerMapping import ChunkerConfig
+
class AudioDocument(Document):
type: str = "audio"
def create_transcript(self):
result = get_llm_client().create_transcript(self.raw_data_location)
- return(result.text)
+ return result.text
def read(self, chunk_size: int, chunker: str):
# Transcribe the audio file
-
+
text = self.create_transcript()
chunker_func = ChunkerConfig.get_chunker(chunker)
- chunker = chunker_func(self, chunk_size = chunk_size, get_text = lambda: [text])
+ chunker = chunker_func(self, chunk_size=chunk_size, get_text=lambda: [text])
yield from chunker.read()
diff --git a/cognee/modules/data/processing/document_types/ChunkerMapping.py b/cognee/modules/data/processing/document_types/ChunkerMapping.py
index 14dbb8bb7..f9a251528 100644
--- a/cognee/modules/data/processing/document_types/ChunkerMapping.py
+++ b/cognee/modules/data/processing/document_types/ChunkerMapping.py
@@ -1,9 +1,8 @@
from cognee.modules.chunking.TextChunker import TextChunker
+
class ChunkerConfig:
- chunker_mapping = {
- "text_chunker": TextChunker
- }
+ chunker_mapping = {"text_chunker": TextChunker}
@classmethod
def get_chunker(cls, chunker_name: str):
@@ -12,4 +11,4 @@ def get_chunker(cls, chunker_name: str):
raise NotImplementedError(
f"Chunker '{chunker_name}' is not implemented. Available options: {list(cls.chunker_mapping.keys())}"
)
- return chunker_class
\ No newline at end of file
+ return chunker_class
diff --git a/cognee/modules/data/processing/document_types/Document.py b/cognee/modules/data/processing/document_types/Document.py
index 8d6a3dafb..08380e809 100644
--- a/cognee/modules/data/processing/document_types/Document.py
+++ b/cognee/modules/data/processing/document_types/Document.py
@@ -8,10 +8,7 @@ class Document(DataPoint):
raw_data_location: str
metadata_id: UUID
mime_type: str
- _metadata: dict = {
- "index_fields": ["name"],
- "type": "Document"
- }
+ _metadata: dict = {"index_fields": ["name"], "type": "Document"}
- def read(self, chunk_size: int, chunker = str) -> str:
+ def read(self, chunk_size: int, chunker=str) -> str:
pass
diff --git a/cognee/modules/data/processing/document_types/ImageDocument.py b/cognee/modules/data/processing/document_types/ImageDocument.py
index 352486bd8..424cd059c 100644
--- a/cognee/modules/data/processing/document_types/ImageDocument.py
+++ b/cognee/modules/data/processing/document_types/ImageDocument.py
@@ -2,19 +2,19 @@
from .Document import Document
from .ChunkerMapping import ChunkerConfig
+
class ImageDocument(Document):
type: str = "image"
-
def transcribe_image(self):
result = get_llm_client().transcribe_image(self.raw_data_location)
- return(result.choices[0].message.content)
+ return result.choices[0].message.content
def read(self, chunk_size: int, chunker: str):
# Transcribe the image file
text = self.transcribe_image()
chunker_func = ChunkerConfig.get_chunker(chunker)
- chunker = chunker_func(self, chunk_size = chunk_size, get_text = lambda: [text])
+ chunker = chunker_func(self, chunk_size=chunk_size, get_text=lambda: [text])
yield from chunker.read()
diff --git a/cognee/modules/data/processing/document_types/PdfDocument.py b/cognee/modules/data/processing/document_types/PdfDocument.py
index 361214718..684fb428c 100644
--- a/cognee/modules/data/processing/document_types/PdfDocument.py
+++ b/cognee/modules/data/processing/document_types/PdfDocument.py
@@ -2,6 +2,7 @@
from .Document import Document
from .ChunkerMapping import ChunkerConfig
+
class PdfDocument(Document):
type: str = "pdf"
@@ -14,7 +15,7 @@ def get_text():
yield page_text
chunker_func = ChunkerConfig.get_chunker(chunker)
- chunker = chunker_func(self, chunk_size = chunk_size, get_text = get_text)
+ chunker = chunker_func(self, chunk_size=chunk_size, get_text=get_text)
yield from chunker.read()
diff --git a/cognee/modules/data/processing/document_types/TextDocument.py b/cognee/modules/data/processing/document_types/TextDocument.py
index 3952d9845..f993ff221 100644
--- a/cognee/modules/data/processing/document_types/TextDocument.py
+++ b/cognee/modules/data/processing/document_types/TextDocument.py
@@ -1,12 +1,13 @@
from .Document import Document
from .ChunkerMapping import ChunkerConfig
+
class TextDocument(Document):
type: str = "text"
def read(self, chunk_size: int, chunker: str):
def get_text():
- with open(self.raw_data_location, mode = "r", encoding = "utf-8") as file:
+ with open(self.raw_data_location, mode="r", encoding="utf-8") as file:
while True:
text = file.read(1024)
@@ -17,6 +18,6 @@ def get_text():
chunker_func = ChunkerConfig.get_chunker(chunker)
- chunker = chunker_func(self, chunk_size = chunk_size, get_text = get_text)
+ chunker = chunker_func(self, chunk_size=chunk_size, get_text=get_text)
yield from chunker.read()
diff --git a/cognee/modules/data/processing/document_types/UnstructuredDocument.py b/cognee/modules/data/processing/document_types/UnstructuredDocument.py
index 62632cd08..cd5c72e3b 100644
--- a/cognee/modules/data/processing/document_types/UnstructuredDocument.py
+++ b/cognee/modules/data/processing/document_types/UnstructuredDocument.py
@@ -27,6 +27,6 @@ def get_text():
yield text
- chunker = TextChunker(self, chunk_size = chunk_size, get_text = get_text)
+ chunker = TextChunker(self, chunk_size=chunk_size, get_text=get_text)
yield from chunker.read()
diff --git a/cognee/modules/data/processing/has_new_chunks.py b/cognee/modules/data/processing/has_new_chunks.py
index e5302b626..fb21be603 100644
--- a/cognee/modules/data/processing/has_new_chunks.py
+++ b/cognee/modules/data/processing/has_new_chunks.py
@@ -1,7 +1,10 @@
from cognee.infrastructure.databases.vector import get_vector_engine
from cognee.modules.chunking import DocumentChunk
-async def has_new_chunks(data_chunks: list[DocumentChunk], collection_name: str) -> list[DocumentChunk]:
+
+async def has_new_chunks(
+ data_chunks: list[DocumentChunk], collection_name: str
+) -> list[DocumentChunk]:
vector_engine = get_vector_engine()
if not await vector_engine.has_collection(collection_name):
@@ -22,9 +25,10 @@ async def has_new_chunks(data_chunks: list[DocumentChunk], collection_name: str)
existing_chunks_map = {chunk.id: chunk.payload for chunk in existing_chunks}
new_data_chunks = [
- chunk for chunk in data_chunks \
- if chunk.chunk_id not in existing_chunks_map \
- or chunk.text != existing_chunks_map[chunk.chunk_id]["text"]
+ chunk
+ for chunk in data_chunks
+ if chunk.chunk_id not in existing_chunks_map
+ or chunk.text != existing_chunks_map[chunk.chunk_id]["text"]
]
return len(new_data_chunks) > 0
diff --git a/cognee/modules/engine/models/Entity.py b/cognee/modules/engine/models/Entity.py
index 16e0ca3d8..63a153bf2 100644
--- a/cognee/modules/engine/models/Entity.py
+++ b/cognee/modules/engine/models/Entity.py
@@ -8,7 +8,4 @@ class Entity(DataPoint):
is_a: EntityType
description: str
- _metadata: dict = {
- "index_fields": ["name"],
- "type": "Entity"
- }
+ _metadata: dict = {"index_fields": ["name"], "type": "Entity"}
diff --git a/cognee/modules/engine/models/EntityType.py b/cognee/modules/engine/models/EntityType.py
index d3cc54311..7225bb3ae 100644
--- a/cognee/modules/engine/models/EntityType.py
+++ b/cognee/modules/engine/models/EntityType.py
@@ -6,7 +6,4 @@ class EntityType(DataPoint):
name: str
description: str
- _metadata: dict = {
- "index_fields": ["name"],
- "type": "EntityType"
- }
+ _metadata: dict = {"index_fields": ["name"], "type": "EntityType"}
diff --git a/cognee/modules/engine/utils/generate_node_id.py b/cognee/modules/engine/utils/generate_node_id.py
index db086a19c..489a88875 100644
--- a/cognee/modules/engine/utils/generate_node_id.py
+++ b/cognee/modules/engine/utils/generate_node_id.py
@@ -1,4 +1,5 @@
from uuid import NAMESPACE_OID, uuid5
+
def generate_node_id(node_id: str) -> str:
return uuid5(NAMESPACE_OID, node_id.lower().replace(" ", "_").replace("'", ""))
diff --git a/cognee/modules/graph/cognee_graph/CogneeAbstractGraph.py b/cognee/modules/graph/cognee_graph/CogneeAbstractGraph.py
index 9a7fb677f..471977d4c 100644
--- a/cognee/modules/graph/cognee_graph/CogneeAbstractGraph.py
+++ b/cognee/modules/graph/cognee_graph/CogneeAbstractGraph.py
@@ -3,6 +3,7 @@
from cognee.modules.graph.cognee_graph.CogneeGraphElements import Node, Edge
from cognee.infrastructure.databases.graph.graph_db_interface import GraphDBInterface
+
class CogneeAbstractGraph(ABC):
"""
Abstract base class for representing a graph structure.
@@ -30,6 +31,8 @@ def get_edges(self, node_id: str) -> List[Edge]:
pass
@abstractmethod
- async def project_graph_from_db(self, adapter: GraphDBInterface, directed: bool, dimension: int) -> None:
+ async def project_graph_from_db(
+ self, adapter: GraphDBInterface, directed: bool, dimension: int
+ ) -> None:
"""Project the graph structure from a database using the provided adapter."""
pass
diff --git a/cognee/modules/graph/cognee_graph/CogneeGraph.py b/cognee/modules/graph/cognee_graph/CogneeGraph.py
index dbfbc7bb7..279a73b19 100644
--- a/cognee/modules/graph/cognee_graph/CogneeGraph.py
+++ b/cognee/modules/graph/cognee_graph/CogneeGraph.py
@@ -52,7 +52,7 @@ def get_edges_from_node(self, node_id: str) -> List[Edge]:
else:
raise EntityNotFoundError(message=f"Node with id {node_id} does not exist.")
- def get_edges(self)-> List[Edge]:
+ def get_edges(self) -> List[Edge]:
return self.edges
async def project_graph_from_db(
@@ -60,12 +60,11 @@ async def project_graph_from_db(
adapter: Union[GraphDBInterface],
node_properties_to_project: List[str],
edge_properties_to_project: List[str],
- directed = True,
- node_dimension = 1,
- edge_dimension = 1,
- memory_fragment_filter = [],
+ directed=True,
+ node_dimension=1,
+ edge_dimension=1,
+ memory_fragment_filter=[],
) -> None:
-
if node_dimension < 1 or edge_dimension < 1:
raise InvalidValueError(message="Dimensions must be positive integers")
@@ -73,7 +72,9 @@ async def project_graph_from_db(
if len(memory_fragment_filter) == 0:
nodes_data, edges_data = await adapter.get_graph_data()
else:
- nodes_data, edges_data = await adapter.get_filtered_graph_data(attribute_filters = memory_fragment_filter)
+ nodes_data, edges_data = await adapter.get_filtered_graph_data(
+ attribute_filters=memory_fragment_filter
+ )
if not nodes_data:
raise EntityNotFoundError(message="No node data retrieved from the database.")
@@ -88,17 +89,27 @@ async def project_graph_from_db(
source_node = self.get_node(str(source_id))
target_node = self.get_node(str(target_id))
if source_node and target_node:
- edge_attributes = {key: properties.get(key) for key in edge_properties_to_project}
- edge_attributes['relationship_type'] = relationship_type
-
- edge = Edge(source_node, target_node, attributes=edge_attributes, directed=directed, dimension=edge_dimension)
+ edge_attributes = {
+ key: properties.get(key) for key in edge_properties_to_project
+ }
+ edge_attributes["relationship_type"] = relationship_type
+
+ edge = Edge(
+ source_node,
+ target_node,
+ attributes=edge_attributes,
+ directed=directed,
+ dimension=edge_dimension,
+ )
self.add_edge(edge)
source_node.add_skeleton_edge(edge)
target_node.add_skeleton_edge(edge)
else:
- raise EntityNotFoundError(message=f"Edge references nonexistent nodes: {source_id} -> {target_id}")
+ raise EntityNotFoundError(
+ message=f"Edge references nonexistent nodes: {source_id} -> {target_id}"
+ )
except (ValueError, TypeError) as e:
print(f"Error projecting graph: {e}")
@@ -110,13 +121,15 @@ async def map_vector_distances_to_graph_nodes(self, node_distances) -> None:
for scored_result in scored_results:
node_id = str(scored_result.id)
score = scored_result.score
- node =self.get_node(node_id)
+ node = self.get_node(node_id)
if node:
node.add_attribute("vector_distance", score)
else:
print(f"Node with id {node_id} not found in the graph.")
- async def map_vector_distances_to_graph_edges(self, vector_engine, query) -> None: # :TODO: When we calculate edge embeddings in vector db change this similarly to node mapping
+ async def map_vector_distances_to_graph_edges(
+ self, vector_engine, query
+ ) -> None: # :TODO: When we calculate edge embeddings in vector db change this similarly to node mapping
try:
# Step 1: Generate the query embedding
query_vector = await vector_engine.embed_data([query])
@@ -127,7 +140,7 @@ async def map_vector_distances_to_graph_edges(self, vector_engine, query) -> Non
# Step 2: Collect all unique relationship types
unique_relationship_types = set()
for edge in self.edges:
- relationship_type = edge.attributes.get('relationship_type')
+ relationship_type = edge.attributes.get("relationship_type")
if relationship_type:
unique_relationship_types.add(relationship_type)
@@ -137,12 +150,14 @@ async def map_vector_distances_to_graph_edges(self, vector_engine, query) -> Non
# Step 4: Map relationship types to their embeddings and calculate distances
embedding_map = {}
- for relationship_type, embedding in zip(unique_relationship_types, relationship_type_embeddings):
+ for relationship_type, embedding in zip(
+ unique_relationship_types, relationship_type_embeddings
+ ):
edge_vector = np.array(embedding)
# Calculate cosine similarity
similarity = np.dot(query_vector, edge_vector) / (
- np.linalg.norm(query_vector) * np.linalg.norm(edge_vector)
+ np.linalg.norm(query_vector) * np.linalg.norm(edge_vector)
)
distance = 1 - similarity
@@ -151,7 +166,7 @@ async def map_vector_distances_to_graph_edges(self, vector_engine, query) -> Non
# Step 4: Assign precomputed distances to edges
for edge in self.edges:
- relationship_type = edge.attributes.get('relationship_type')
+ relationship_type = edge.attributes.get("relationship_type")
if not relationship_type or relationship_type not in embedding_map:
print(f"Edge {edge} has an unknown or missing relationship type.")
continue
@@ -162,7 +177,6 @@ async def map_vector_distances_to_graph_edges(self, vector_engine, query) -> Non
except Exception as ex:
print(f"Error mapping vector distances to edges: {ex}")
-
async def calculate_top_triplet_importances(self, k: int) -> List:
min_heap = []
for i, edge in enumerate(self.edges):
@@ -179,6 +193,4 @@ async def calculate_top_triplet_importances(self, k: int) -> List:
if len(min_heap) > k:
heapq.heappop(min_heap)
-
return [edge for _, _, edge in sorted(min_heap)]
-
diff --git a/cognee/modules/graph/cognee_graph/CogneeGraphElements.py b/cognee/modules/graph/cognee_graph/CogneeGraphElements.py
index bab0f3bb6..33f02dc22 100644
--- a/cognee/modules/graph/cognee_graph/CogneeGraphElements.py
+++ b/cognee/modules/graph/cognee_graph/CogneeGraphElements.py
@@ -6,25 +6,28 @@
class Node:
"""
- Represents a node in a graph.
- Attributes:
- id (str): A unique identifier for the node.
- attributes (Dict[str, Any]): A dictionary of attributes associated with the node.
- neighbors (List[Node]): Represents the original nodes
- skeleton_edges (List[Edge]): Represents the original edges
- """
+ Represents a node in a graph.
+ Attributes:
+ id (str): A unique identifier for the node.
+ attributes (Dict[str, Any]): A dictionary of attributes associated with the node.
+ neighbors (List[Node]): Represents the original nodes
+ skeleton_edges (List[Edge]): Represents the original edges
+ """
+
id: str
attributes: Dict[str, Any]
skeleton_neighbours: List["Node"]
skeleton_edges: List["Edge"]
status: np.ndarray
- def __init__(self, node_id: str, attributes: Optional[Dict[str, Any]] = None, dimension: int = 1):
+ def __init__(
+ self, node_id: str, attributes: Optional[Dict[str, Any]] = None, dimension: int = 1
+ ):
if dimension <= 0:
raise InvalidValueError(message="Dimension must be a positive integer")
self.id = node_id
self.attributes = attributes if attributes is not None else {}
- self.attributes["vector_distance"] = float('inf')
+ self.attributes["vector_distance"] = float("inf")
self.skeleton_neighbours = []
self.skeleton_edges = []
self.status = np.ones(dimension, dtype=int)
@@ -56,7 +59,9 @@ def remove_skeleton_edge(self, edge: "Edge") -> None:
def is_node_alive_in_dimension(self, dimension: int) -> bool:
if dimension < 0 or dimension >= len(self.status):
- raise InvalidValueError(message=f"Dimension {dimension} is out of range. Valid range is 0 to {len(self.status) - 1}.")
+ raise InvalidValueError(
+ message=f"Dimension {dimension} is out of range. Valid range is 0 to {len(self.status) - 1}."
+ )
return self.status[dimension] == 1
def add_attribute(self, key: str, value: Any) -> None:
@@ -83,12 +88,12 @@ def __eq__(self, other: "Node") -> bool:
class Edge:
"""
- Represents an edge in a graph, connecting two nodes.
- Attributes:
- node1 (Node): The starting node of the edge.
- node2 (Node): The ending node of the edge.
- attributes (Dict[str, Any]): A dictionary of attributes associated with the edge.
- directed (bool): A flag indicating whether the edge is directed or undirected.
+ Represents an edge in a graph, connecting two nodes.
+ Attributes:
+ node1 (Node): The starting node of the edge.
+ node2 (Node): The ending node of the edge.
+ attributes (Dict[str, Any]): A dictionary of attributes associated with the edge.
+ directed (bool): A flag indicating whether the edge is directed or undirected.
"""
node1: "Node"
@@ -97,19 +102,28 @@ class Edge:
directed: bool
status: np.ndarray
- def __init__(self, node1: "Node", node2: "Node", attributes: Optional[Dict[str, Any]] = None, directed: bool = True, dimension: int = 1):
+ def __init__(
+ self,
+ node1: "Node",
+ node2: "Node",
+ attributes: Optional[Dict[str, Any]] = None,
+ directed: bool = True,
+ dimension: int = 1,
+ ):
if dimension <= 0:
raise InvalidValueError(message="Dimensions must be a positive integer.")
self.node1 = node1
self.node2 = node2
self.attributes = attributes if attributes is not None else {}
- self.attributes["vector_distance"] = float('inf')
+ self.attributes["vector_distance"] = float("inf")
self.directed = directed
self.status = np.ones(dimension, dtype=int)
def is_edge_alive_in_dimension(self, dimension: int) -> bool:
if dimension < 0 or dimension >= len(self.status):
- raise InvalidValueError(message=f"Dimension {dimension} is out of range. Valid range is 0 to {len(self.status) - 1}.")
+ raise InvalidValueError(
+ message=f"Dimension {dimension} is out of range. Valid range is 0 to {len(self.status) - 1}."
+ )
return self.status[dimension] == 1
def add_attribute(self, key: str, value: Any) -> None:
@@ -140,4 +154,4 @@ def __eq__(self, other: "Edge") -> bool:
if self.directed:
return self.node1 == other.node1 and self.node2 == other.node2
else:
- return {self.node1, self.node2} == {other.node1, other.node2}
\ No newline at end of file
+ return {self.node1, self.node2} == {other.node1, other.node2}
diff --git a/cognee/modules/graph/exceptions/__init__.py b/cognee/modules/graph/exceptions/__init__.py
index e8330caf3..5cf600099 100644
--- a/cognee/modules/graph/exceptions/__init__.py
+++ b/cognee/modules/graph/exceptions/__init__.py
@@ -7,4 +7,4 @@
from .exceptions import (
EntityNotFoundError,
EntityAlreadyExistsError,
-)
\ No newline at end of file
+)
diff --git a/cognee/modules/graph/exceptions/exceptions.py b/cognee/modules/graph/exceptions/exceptions.py
index af15bb616..854e620ff 100644
--- a/cognee/modules/graph/exceptions/exceptions.py
+++ b/cognee/modules/graph/exceptions/exceptions.py
@@ -1,6 +1,7 @@
from cognee.exceptions import CogneeApiError
from fastapi import status
+
class EntityNotFoundError(CogneeApiError):
"""Database returns nothing"""
@@ -22,4 +23,4 @@ def __init__(
name: str = "EntityAlreadyExistsError",
status_code=status.HTTP_409_CONFLICT,
):
- super().__init__(message, name, status_code)
\ No newline at end of file
+ super().__init__(message, name, status_code)
diff --git a/cognee/modules/graph/models/EdgeType.py b/cognee/modules/graph/models/EdgeType.py
index 998f08d8d..36aec2e7c 100644
--- a/cognee/modules/graph/models/EdgeType.py
+++ b/cognee/modules/graph/models/EdgeType.py
@@ -8,7 +8,4 @@ class EdgeType(DataPoint):
relationship_name: str
number_of_edges: int
- _metadata: dict = {
- "index_fields": ["relationship_name"],
- "type": "EdgeType"
- }
\ No newline at end of file
+ _metadata: dict = {"index_fields": ["relationship_name"], "type": "EdgeType"}
diff --git a/cognee/modules/graph/utils/convert_node_to_data_point.py b/cognee/modules/graph/utils/convert_node_to_data_point.py
index 602a7ffa3..b89e50db9 100644
--- a/cognee/modules/graph/utils/convert_node_to_data_point.py
+++ b/cognee/modules/graph/utils/convert_node_to_data_point.py
@@ -15,6 +15,7 @@ def get_all_subclasses(cls):
return subclasses
+
def find_subclass_by_name(cls, name):
for subclass in get_all_subclasses(cls):
if subclass.__name__ == name:
diff --git a/cognee/modules/graph/utils/deduplicate_nodes_and_edges.py b/cognee/modules/graph/utils/deduplicate_nodes_and_edges.py
index e863960ea..f61d51f3e 100644
--- a/cognee/modules/graph/utils/deduplicate_nodes_and_edges.py
+++ b/cognee/modules/graph/utils/deduplicate_nodes_and_edges.py
@@ -1,5 +1,6 @@
from cognee.infrastructure.engine import DataPoint
+
def deduplicate_nodes_and_edges(nodes: list[DataPoint], edges: list[dict]):
added_entities = {}
final_nodes = []
diff --git a/cognee/modules/graph/utils/expand_with_nodes_and_edges.py b/cognee/modules/graph/utils/expand_with_nodes_and_edges.py
index cfa2ad9b4..23a6f4f63 100644
--- a/cognee/modules/graph/utils/expand_with_nodes_and_edges.py
+++ b/cognee/modules/graph/utils/expand_with_nodes_and_edges.py
@@ -36,10 +36,10 @@ def expand_with_nodes_and_edges(
if f"{str(type_node_id)}_type" not in added_nodes_map:
type_node = EntityType(
- id = type_node_id,
- name = type_node_name,
- type = type_node_name,
- description = type_node_name,
+ id=type_node_id,
+ name=type_node_name,
+ type=type_node_name,
+ description=type_node_name,
)
added_nodes_map[f"{str(type_node_id)}_type"] = type_node
else:
@@ -47,10 +47,10 @@ def expand_with_nodes_and_edges(
if f"{str(node_id)}_entity" not in added_nodes_map:
entity_node = Entity(
- id = node_id,
- name = node_name,
- is_a = type_node,
- description = node.description,
+ id=node_id,
+ name=node_name,
+ is_a=type_node,
+ description=node.description,
)
if data_chunk.contains is None:
@@ -75,11 +75,9 @@ def expand_with_nodes_and_edges(
target_node_id,
edge.relationship_name,
dict(
- relationship_name = generate_edge_name(
- edge.relationship_name
- ),
- source_node_id = source_node_id,
- target_node_id = target_node_id,
+ relationship_name=generate_edge_name(edge.relationship_name),
+ source_node_id=source_node_id,
+ target_node_id=target_node_id,
),
)
)
diff --git a/cognee/modules/graph/utils/get_graph_from_model.py b/cognee/modules/graph/utils/get_graph_from_model.py
index 1b7c0908b..f4d2ed77a 100644
--- a/cognee/modules/graph/utils/get_graph_from_model.py
+++ b/cognee/modules/graph/utils/get_graph_from_model.py
@@ -2,13 +2,14 @@
from cognee.infrastructure.engine import DataPoint
from cognee.modules.storage.utils import copy_model
+
async def get_graph_from_model(
data_point: DataPoint,
added_nodes: dict,
added_edges: dict,
visited_properties: dict = None,
- only_root = False,
- include_root = True,
+ only_root=False,
+ include_root=True,
):
if str(data_point.id) in added_nodes:
return [], []
@@ -37,7 +38,11 @@ async def get_graph_from_model(
continue
- if isinstance(field_value, list) and len(field_value) > 0 and isinstance(field_value[0], DataPoint):
+ if (
+ isinstance(field_value, list)
+ and len(field_value) > 0
+ and isinstance(field_value[0], DataPoint)
+ ):
excluded_properties.add(field_name)
for index, item in enumerate(field_value):
@@ -52,15 +57,14 @@ async def get_graph_from_model(
data_point_properties[field_name] = field_value
-
if include_root and str(data_point.id) not in added_nodes:
SimpleDataPointModel = copy_model(
type(data_point),
- include_fields = {
+ include_fields={
"_metadata": (dict, data_point._metadata),
"__tablename__": (str, data_point.__tablename__),
},
- exclude_fields = list(excluded_properties),
+ exclude_fields=list(excluded_properties),
)
nodes.append(SimpleDataPointModel(**data_point_properties))
added_nodes[str(data_point.id)] = True
@@ -79,12 +83,19 @@ async def get_graph_from_model(
edge_key = str(data_point.id) + str(field_value.id) + field_name
if str(edge_key) not in added_edges:
- edges.append((data_point.id, field_value.id, field_name, {
- "source_node_id": data_point.id,
- "target_node_id": field_value.id,
- "relationship_name": field_name,
- "updated_at": datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M:%S"),
- }))
+ edges.append(
+ (
+ data_point.id,
+ field_value.id,
+ field_name,
+ {
+ "source_node_id": data_point.id,
+ "target_node_id": field_value.id,
+ "relationship_name": field_name,
+ "updated_at": datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M:%S"),
+ },
+ )
+ )
added_edges[str(edge_key)] = True
if str(field_value.id) in added_nodes or only_root:
@@ -92,10 +103,10 @@ async def get_graph_from_model(
property_nodes, property_edges = await get_graph_from_model(
field_value,
- include_root = True,
- added_nodes = added_nodes,
- added_edges = added_edges,
- visited_properties = visited_properties,
+ include_root=True,
+ added_nodes=added_nodes,
+ added_edges=added_edges,
+ visited_properties=visited_properties,
)
for node in property_nodes:
diff --git a/cognee/modules/graph/utils/get_model_instance_from_graph.py b/cognee/modules/graph/utils/get_model_instance_from_graph.py
index 82cdfa150..abe378f51 100644
--- a/cognee/modules/graph/utils/get_model_instance_from_graph.py
+++ b/cognee/modules/graph/utils/get_model_instance_from_graph.py
@@ -18,12 +18,16 @@ def get_model_instance_from_graph(nodes: list[DataPoint], edges: list, entity_id
edge_type = edge_metadata.get("type")
if edge_type == "list":
- NewModel = copy_model(type(source_node), { edge_label: (list[type(target_node)], PydanticUndefined) })
+ NewModel = copy_model(
+ type(source_node), {edge_label: (list[type(target_node)], PydanticUndefined)}
+ )
- node_map[edge[0]] = NewModel(**source_node.model_dump(), **{ edge_label: [target_node] })
+ node_map[edge[0]] = NewModel(**source_node.model_dump(), **{edge_label: [target_node]})
else:
- NewModel = copy_model(type(source_node), { edge_label: (type(target_node), PydanticUndefined) })
+ NewModel = copy_model(
+ type(source_node), {edge_label: (type(target_node), PydanticUndefined)}
+ )
- node_map[edge[0]] = NewModel(**source_node.model_dump(), **{ edge_label: target_node })
+ node_map[edge[0]] = NewModel(**source_node.model_dump(), **{edge_label: target_node})
return node_map[entity_id]
diff --git a/cognee/modules/graph/utils/retrieve_existing_edges.py b/cognee/modules/graph/utils/retrieve_existing_edges.py
index 0959e79d8..637317a8f 100644
--- a/cognee/modules/graph/utils/retrieve_existing_edges.py
+++ b/cognee/modules/graph/utils/retrieve_existing_edges.py
@@ -22,18 +22,12 @@ async def retrieve_existing_edges(
entity_node_id = generate_node_id(node.id)
if str(type_node_id) not in processed_nodes:
- type_node_edges.append(
- (data_chunk.id, type_node_id, "exists_in")
- )
+ type_node_edges.append((data_chunk.id, type_node_id, "exists_in"))
processed_nodes[str(type_node_id)] = True
if str(entity_node_id) not in processed_nodes:
- entity_node_edges.append(
- (data_chunk.id, entity_node_id, "mentioned_in")
- )
- type_entity_edges.append(
- (entity_node_id, type_node_id, "is_a")
- )
+ entity_node_edges.append((data_chunk.id, entity_node_id, "mentioned_in"))
+ type_entity_edges.append((entity_node_id, type_node_id, "is_a"))
processed_nodes[str(entity_node_id)] = True
graph_node_edges = [
diff --git a/cognee/modules/ingestion/classify.py b/cognee/modules/ingestion/classify.py
index dbb191cc3..8c428cbb9 100644
--- a/cognee/modules/ingestion/classify.py
+++ b/cognee/modules/ingestion/classify.py
@@ -13,4 +13,6 @@ def classify(data: Union[str, BinaryIO], filename: str = None):
if isinstance(data, BufferedReader) or isinstance(data, SpooledTemporaryFile):
return BinaryData(data, data.name.split("/")[-1] if data.name else filename)
- raise IngestionError(message=f"Type of data sent to classify(data: Union[str, BinaryIO) not supported: {type(data)}")
+ raise IngestionError(
+ message=f"Type of data sent to classify(data: Union[str, BinaryIO) not supported: {type(data)}"
+ )
diff --git a/cognee/modules/ingestion/data_types/BinaryData.py b/cognee/modules/ingestion/data_types/BinaryData.py
index 0606250ea..7f137a6a0 100644
--- a/cognee/modules/ingestion/data_types/BinaryData.py
+++ b/cognee/modules/ingestion/data_types/BinaryData.py
@@ -2,9 +2,11 @@
from cognee.infrastructure.files import get_file_metadata, FileMetadata
from .IngestionData import IngestionData
+
def create_binary_data(data: BinaryIO):
return BinaryData(data)
+
class BinaryData(IngestionData):
name: str = None
data: BinaryIO = None
diff --git a/cognee/modules/ingestion/data_types/TextData.py b/cognee/modules/ingestion/data_types/TextData.py
index 9f10b0d54..b518d54a9 100644
--- a/cognee/modules/ingestion/data_types/TextData.py
+++ b/cognee/modules/ingestion/data_types/TextData.py
@@ -2,9 +2,11 @@
from cognee.infrastructure.data.utils.extract_keywords import extract_keywords
from .IngestionData import IngestionData
+
def create_text_data(data: str):
return TextData(data)
+
class TextData(IngestionData):
data: str = None
metadata = None
diff --git a/cognee/modules/ingestion/discover_directory_datasets.py b/cognee/modules/ingestion/discover_directory_datasets.py
index 1fdb573f0..861784c16 100644
--- a/cognee/modules/ingestion/discover_directory_datasets.py
+++ b/cognee/modules/ingestion/discover_directory_datasets.py
@@ -1,5 +1,6 @@
from os import path, listdir
+
def discover_directory_datasets(root_dir_path: str, parent_dir: str = None):
datasets = {}
@@ -7,7 +8,9 @@ def discover_directory_datasets(root_dir_path: str, parent_dir: str = None):
if path.isdir(path.join(root_dir_path, file_or_dir)):
dataset_name = file_or_dir if parent_dir is None else f"{parent_dir}.{file_or_dir}"
- nested_datasets = discover_directory_datasets(path.join(root_dir_path, file_or_dir), dataset_name)
+ nested_datasets = discover_directory_datasets(
+ path.join(root_dir_path, file_or_dir), dataset_name
+ )
for dataset in nested_datasets.keys():
datasets[dataset] = nested_datasets[dataset]
diff --git a/cognee/modules/ingestion/exceptions/__init__.py b/cognee/modules/ingestion/exceptions/__init__.py
index 33d59e113..7efe94987 100644
--- a/cognee/modules/ingestion/exceptions/__init__.py
+++ b/cognee/modules/ingestion/exceptions/__init__.py
@@ -6,4 +6,4 @@
from .exceptions import (
IngestionError,
-)
\ No newline at end of file
+)
diff --git a/cognee/modules/ingestion/exceptions/exceptions.py b/cognee/modules/ingestion/exceptions/exceptions.py
index 4901be110..08991a946 100644
--- a/cognee/modules/ingestion/exceptions/exceptions.py
+++ b/cognee/modules/ingestion/exceptions/exceptions.py
@@ -1,11 +1,12 @@
from cognee.exceptions import CogneeApiError
from fastapi import status
+
class IngestionError(CogneeApiError):
def __init__(
- self,
- message: str = "Type of data sent to classify not supported.",
- name: str = "IngestionError",
- status_code=status.HTTP_415_UNSUPPORTED_MEDIA_TYPE,
+ self,
+ message: str = "Type of data sent to classify not supported.",
+ name: str = "IngestionError",
+ status_code=status.HTTP_415_UNSUPPORTED_MEDIA_TYPE,
):
- super().__init__(message, name, status_code)
\ No newline at end of file
+ super().__init__(message, name, status_code)
diff --git a/cognee/modules/ingestion/get_matched_datasets.py b/cognee/modules/ingestion/get_matched_datasets.py
index 563d8016c..d1c41b0f6 100644
--- a/cognee/modules/ingestion/get_matched_datasets.py
+++ b/cognee/modules/ingestion/get_matched_datasets.py
@@ -1,5 +1,6 @@
from .discover_directory_datasets import discover_directory_datasets
+
def get_matched_datasets(data_path: str, dataset_name_to_match: str = None):
datasets = discover_directory_datasets(data_path)
diff --git a/cognee/modules/ingestion/save_data_to_file.py b/cognee/modules/ingestion/save_data_to_file.py
index 1af6ab0aa..2e6ba7e15 100644
--- a/cognee/modules/ingestion/save_data_to_file.py
+++ b/cognee/modules/ingestion/save_data_to_file.py
@@ -5,6 +5,7 @@
from cognee.infrastructure.files.storage import LocalStorage
from .classify import classify
+
def save_data_to_file(data: Union[str, BinaryIO], filename: str = None):
base_config = get_base_config()
data_directory_path = base_config.data_root_directory
@@ -16,7 +17,7 @@ def save_data_to_file(data: Union[str, BinaryIO], filename: str = None):
file_metadata = classified_data.get_metadata()
if "name" not in file_metadata or file_metadata["name"] is None:
- data_contents = classified_data.get_data().encode('utf-8')
+ data_contents = classified_data.get_data().encode("utf-8")
hash_contents = hashlib.md5(data_contents).hexdigest()
file_metadata["name"] = "text_" + hash_contents + ".txt"
file_name = file_metadata["name"]
diff --git a/cognee/modules/pipelines/models/Pipeline.py b/cognee/modules/pipelines/models/Pipeline.py
index f4d20bb0f..6d97633f0 100644
--- a/cognee/modules/pipelines/models/Pipeline.py
+++ b/cognee/modules/pipelines/models/Pipeline.py
@@ -6,18 +6,19 @@
from .PipelineTask import PipelineTask
from .Task import Task
+
class Pipeline(Base):
__tablename__ = "pipelines"
- id = Column(UUID, primary_key = True, default = uuid4)
+ id = Column(UUID, primary_key=True, default=uuid4)
name = Column(String)
- description = Column(Text, nullable = True)
+ description = Column(Text, nullable=True)
- created_at = Column(DateTime(timezone = True), default = lambda: datetime.now(timezone.utc))
- updated_at = Column(DateTime(timezone = True), onupdate = lambda: datetime.now(timezone.utc))
+ created_at = Column(DateTime(timezone=True), default=lambda: datetime.now(timezone.utc))
+ updated_at = Column(DateTime(timezone=True), onupdate=lambda: datetime.now(timezone.utc))
tasks = Mapped[list["Task"]] = relationship(
- secondary = PipelineTask.__tablename__,
- back_populates = "pipeline",
+ secondary=PipelineTask.__tablename__,
+ back_populates="pipeline",
)
diff --git a/cognee/modules/pipelines/models/PipelineRun.py b/cognee/modules/pipelines/models/PipelineRun.py
index ab3498efe..c778066fe 100644
--- a/cognee/modules/pipelines/models/PipelineRun.py
+++ b/cognee/modules/pipelines/models/PipelineRun.py
@@ -4,19 +4,21 @@
from sqlalchemy import Column, DateTime, JSON, Enum, UUID
from cognee.infrastructure.databases.relational import Base
+
class PipelineRunStatus(enum.Enum):
DATASET_PROCESSING_STARTED = "DATASET_PROCESSING_STARTED"
DATASET_PROCESSING_COMPLETED = "DATASET_PROCESSING_COMPLETED"
DATASET_PROCESSING_ERRORED = "DATASET_PROCESSING_ERRORED"
+
class PipelineRun(Base):
__tablename__ = "pipeline_runs"
- id = Column(UUID, primary_key = True, default = uuid4)
+ id = Column(UUID, primary_key=True, default=uuid4)
- created_at = Column(DateTime(timezone = True), default = lambda: datetime.now(timezone.utc))
+ created_at = Column(DateTime(timezone=True), default=lambda: datetime.now(timezone.utc))
status = Column(Enum(PipelineRunStatus))
- run_id = Column(UUID, index = True)
+ run_id = Column(UUID, index=True)
run_info = Column(JSON)
diff --git a/cognee/modules/pipelines/models/PipelineTask.py b/cognee/modules/pipelines/models/PipelineTask.py
index c6c7eb5e9..ca38fde48 100644
--- a/cognee/modules/pipelines/models/PipelineTask.py
+++ b/cognee/modules/pipelines/models/PipelineTask.py
@@ -2,10 +2,11 @@
from sqlalchemy import Column, DateTime, ForeignKey, UUID
from cognee.infrastructure.databases.relational import Base
+
class PipelineTask(Base):
__tablename__ = "pipeline_task"
- created_at = Column(DateTime(timezone = True), default = lambda: datetime.now(timezone.utc))
+ created_at = Column(DateTime(timezone=True), default=lambda: datetime.now(timezone.utc))
- pipeline_id = Column("pipeline", UUID, ForeignKey("pipeline.id"), primary_key = True)
- task_id = Column("task", UUID, ForeignKey("task.id"), primary_key = True)
+ pipeline_id = Column("pipeline", UUID, ForeignKey("pipeline.id"), primary_key=True)
+ task_id = Column("task", UUID, ForeignKey("task.id"), primary_key=True)
diff --git a/cognee/modules/pipelines/models/Task.py b/cognee/modules/pipelines/models/Task.py
index eba6508e3..bf960bf09 100644
--- a/cognee/modules/pipelines/models/Task.py
+++ b/cognee/modules/pipelines/models/Task.py
@@ -5,20 +5,20 @@
from cognee.infrastructure.databases.relational import Base, UUID
from .PipelineTask import PipelineTask
+
class Task(Base):
__tablename__ = "tasks"
- id = Column(UUID, primary_key = True, default = uuid4)
+ id = Column(UUID, primary_key=True, default=uuid4)
name = Column(String)
- description = Column(Text, nullable = True)
+ description = Column(Text, nullable=True)
executable = Column(Text)
- created_at = Column(DateTime(timezone = True), default = lambda: datetime.now(timezone.utc))
- updated_at = Column(DateTime(timezone = True), onupdate = lambda: datetime.now(timezone.utc))
+ created_at = Column(DateTime(timezone=True), default=lambda: datetime.now(timezone.utc))
+ updated_at = Column(DateTime(timezone=True), onupdate=lambda: datetime.now(timezone.utc))
datasets: Mapped[list["Pipeline"]] = relationship(
- secondary = PipelineTask.__tablename__,
- back_populates = "task"
+ secondary=PipelineTask.__tablename__, back_populates="task"
)
diff --git a/cognee/modules/pipelines/models/TaskRun.py b/cognee/modules/pipelines/models/TaskRun.py
index bf2e2a9f8..4f5059cdb 100644
--- a/cognee/modules/pipelines/models/TaskRun.py
+++ b/cognee/modules/pipelines/models/TaskRun.py
@@ -3,14 +3,15 @@
from sqlalchemy import Column, DateTime, String, JSON
from cognee.infrastructure.databases.relational import Base, UUID
+
class TaskRun(Base):
__tablename__ = "task_runs"
- id = Column(UUID, primary_key = True, default = uuid4)
+ id = Column(UUID, primary_key=True, default=uuid4)
task_name = Column(String)
- created_at = Column(DateTime(timezone = True), default = lambda: datetime.now(timezone.utc))
+ created_at = Column(DateTime(timezone=True), default=lambda: datetime.now(timezone.utc))
status = Column(String)
diff --git a/cognee/modules/pipelines/operations/get_pipeline_status.py b/cognee/modules/pipelines/operations/get_pipeline_status.py
index 4249d3bf0..4f49cd544 100644
--- a/cognee/modules/pipelines/operations/get_pipeline_status.py
+++ b/cognee/modules/pipelines/operations/get_pipeline_status.py
@@ -4,29 +4,32 @@
from cognee.infrastructure.databases.relational import get_relational_engine
from ..models import PipelineRun
+
async def get_pipeline_status(pipeline_ids: list[UUID]):
db_engine = get_relational_engine()
async with db_engine.get_async_session() as session:
- query = select(
- PipelineRun,
- func.row_number().over(
- partition_by = PipelineRun.run_id,
- order_by = PipelineRun.created_at.desc(),
- ).label("rn")
- ).filter(PipelineRun.run_id.in_(pipeline_ids)).subquery()
+ query = (
+ select(
+ PipelineRun,
+ func.row_number()
+ .over(
+ partition_by=PipelineRun.run_id,
+ order_by=PipelineRun.created_at.desc(),
+ )
+ .label("rn"),
+ )
+ .filter(PipelineRun.run_id.in_(pipeline_ids))
+ .subquery()
+ )
aliased_pipeline_run = aliased(PipelineRun, query)
- latest_runs = (
- select(aliased_pipeline_run).filter(query.c.rn == 1)
- )
-
+ latest_runs = select(aliased_pipeline_run).filter(query.c.rn == 1)
+
runs = (await session.execute(latest_runs)).scalars().all()
- pipeline_statuses = {
- str(run.run_id): run.status for run in runs
- }
+ pipeline_statuses = {str(run.run_id): run.status for run in runs}
return pipeline_statuses
@@ -38,4 +41,4 @@ async def get_pipeline_status(pipeline_ids: list[UUID]):
# ) t
# WHERE rn = 1;"""
- # return { dataset["data_id"]: dataset["status"] for dataset in datasets_statuses }
+ # return { dataset["data_id"]: dataset["status"] for dataset in datasets_statuses }
diff --git a/cognee/modules/pipelines/operations/log_pipeline_status.py b/cognee/modules/pipelines/operations/log_pipeline_status.py
index 94872dee4..c0f08cd2a 100644
--- a/cognee/modules/pipelines/operations/log_pipeline_status.py
+++ b/cognee/modules/pipelines/operations/log_pipeline_status.py
@@ -2,14 +2,17 @@
from cognee.infrastructure.databases.relational import get_relational_engine
from ..models.PipelineRun import PipelineRun
+
async def log_pipeline_status(run_id: UUID, status: str, run_info: dict):
db_engine = get_relational_engine()
async with db_engine.get_async_session() as session:
- session.add(PipelineRun(
- run_id = run_id,
- status = status,
- run_info = run_info,
- ))
+ session.add(
+ PipelineRun(
+ run_id=run_id,
+ status=status,
+ run_info=run_info,
+ )
+ )
await session.commit()
diff --git a/cognee/modules/pipelines/operations/run_parallel.py b/cognee/modules/pipelines/operations/run_parallel.py
index ab78b8300..d1774299b 100644
--- a/cognee/modules/pipelines/operations/run_parallel.py
+++ b/cognee/modules/pipelines/operations/run_parallel.py
@@ -2,6 +2,7 @@
import asyncio
from ..tasks.Task import Task
+
def run_tasks_parallel(tasks: List[Task]) -> Callable[[Any], Generator[Any, Any, Any]]:
async def parallel_run(*args, **kwargs):
parallel_tasks = [asyncio.create_task(task.run(*args, **kwargs)) for task in tasks]
diff --git a/cognee/modules/pipelines/operations/run_tasks.py b/cognee/modules/pipelines/operations/run_tasks.py
index 35e32cf74..d5d7ef7a4 100644
--- a/cognee/modules/pipelines/operations/run_tasks.py
+++ b/cognee/modules/pipelines/operations/run_tasks.py
@@ -11,7 +11,8 @@
logger = logging.getLogger("run_tasks(tasks: [Task], data)")
-async def run_tasks_base(tasks: list[Task], data = None, user: User = None):
+
+async def run_tasks_base(tasks: list[Task], data=None, user: User = None):
if len(tasks) == 0:
yield data
return
@@ -25,9 +26,13 @@ async def run_tasks_base(tasks: list[Task], data = None, user: User = None):
if inspect.isasyncgenfunction(running_task.executable):
logger.info("Async generator task started: `%s`", running_task.executable.__name__)
- send_telemetry("Async Generator Task Started", user.id, {
- "task_name": running_task.executable.__name__,
- })
+ send_telemetry(
+ "Async Generator Task Started",
+ user.id,
+ {
+ "task_name": running_task.executable.__name__,
+ },
+ )
try:
results = []
@@ -40,7 +45,7 @@ async def run_tasks_base(tasks: list[Task], data = None, user: User = None):
async for result in run_tasks_base(
leftover_tasks,
results[0] if next_task_batch_size == 1 else results,
- user = user,
+ user=user,
):
yield result
@@ -53,26 +58,38 @@ async def run_tasks_base(tasks: list[Task], data = None, user: User = None):
results = []
logger.info("Async generator task completed: `%s`", running_task.executable.__name__)
- send_telemetry("Async Generator Task Completed", user.id, {
- "task_name": running_task.executable.__name__,
- })
+ send_telemetry(
+ "Async Generator Task Completed",
+ user.id,
+ {
+ "task_name": running_task.executable.__name__,
+ },
+ )
except Exception as error:
logger.error(
"Async generator task errored: `%s`\n%s\n",
running_task.executable.__name__,
str(error),
- exc_info = True,
+ exc_info=True,
+ )
+ send_telemetry(
+ "Async Generator Task Errored",
+ user.id,
+ {
+ "task_name": running_task.executable.__name__,
+ },
)
- send_telemetry("Async Generator Task Errored", user.id, {
- "task_name": running_task.executable.__name__,
- })
raise error
elif inspect.isgeneratorfunction(running_task.executable):
logger.info("Generator task started: `%s`", running_task.executable.__name__)
- send_telemetry("Generator Task Started", user.id, {
- "task_name": running_task.executable.__name__,
- })
+ send_telemetry(
+ "Generator Task Started",
+ user.id,
+ {
+ "task_name": running_task.executable.__name__,
+ },
+ )
try:
results = []
@@ -80,7 +97,9 @@ async def run_tasks_base(tasks: list[Task], data = None, user: User = None):
results.append(partial_result)
if len(results) == next_task_batch_size:
- async for result in run_tasks_base(leftover_tasks, results[0] if next_task_batch_size == 1 else results, user):
+ async for result in run_tasks_base(
+ leftover_tasks, results[0] if next_task_batch_size == 1 else results, user
+ ):
yield result
results = []
@@ -92,26 +111,38 @@ async def run_tasks_base(tasks: list[Task], data = None, user: User = None):
results = []
logger.info("Generator task completed: `%s`", running_task.executable.__name__)
- send_telemetry("Generator Task Completed", user_id = user.id, additional_properties = {
- "task_name": running_task.executable.__name__,
- })
+ send_telemetry(
+ "Generator Task Completed",
+ user_id=user.id,
+ additional_properties={
+ "task_name": running_task.executable.__name__,
+ },
+ )
except Exception as error:
logger.error(
"Generator task errored: `%s`\n%s\n",
running_task.executable.__name__,
str(error),
- exc_info = True,
+ exc_info=True,
+ )
+ send_telemetry(
+ "Generator Task Errored",
+ user_id=user.id,
+ additional_properties={
+ "task_name": running_task.executable.__name__,
+ },
)
- send_telemetry("Generator Task Errored", user_id = user.id, additional_properties = {
- "task_name": running_task.executable.__name__,
- })
raise error
elif inspect.iscoroutinefunction(running_task.executable):
logger.info("Coroutine task started: `%s`", running_task.executable.__name__)
- send_telemetry("Coroutine Task Started", user_id = user.id, additional_properties = {
- "task_name": running_task.executable.__name__,
- })
+ send_telemetry(
+ "Coroutine Task Started",
+ user_id=user.id,
+ additional_properties={
+ "task_name": running_task.executable.__name__,
+ },
+ )
try:
task_result = await running_task.run(*args)
@@ -119,26 +150,38 @@ async def run_tasks_base(tasks: list[Task], data = None, user: User = None):
yield result
logger.info("Coroutine task completed: `%s`", running_task.executable.__name__)
- send_telemetry("Coroutine Task Completed", user.id, {
- "task_name": running_task.executable.__name__,
- })
+ send_telemetry(
+ "Coroutine Task Completed",
+ user.id,
+ {
+ "task_name": running_task.executable.__name__,
+ },
+ )
except Exception as error:
logger.error(
"Coroutine task errored: `%s`\n%s\n",
running_task.executable.__name__,
str(error),
- exc_info = True,
+ exc_info=True,
+ )
+ send_telemetry(
+ "Coroutine Task Errored",
+ user.id,
+ {
+ "task_name": running_task.executable.__name__,
+ },
)
- send_telemetry("Coroutine Task Errored", user.id, {
- "task_name": running_task.executable.__name__,
- })
raise error
elif inspect.isfunction(running_task.executable):
logger.info("Function task started: `%s`", running_task.executable.__name__)
- send_telemetry("Function Task Started", user.id, {
- "task_name": running_task.executable.__name__,
- })
+ send_telemetry(
+ "Function Task Started",
+ user.id,
+ {
+ "task_name": running_task.executable.__name__,
+ },
+ )
try:
task_result = running_task.run(*args)
@@ -146,59 +189,78 @@ async def run_tasks_base(tasks: list[Task], data = None, user: User = None):
yield result
logger.info("Function task completed: `%s`", running_task.executable.__name__)
- send_telemetry("Function Task Completed", user.id, {
- "task_name": running_task.executable.__name__,
- })
+ send_telemetry(
+ "Function Task Completed",
+ user.id,
+ {
+ "task_name": running_task.executable.__name__,
+ },
+ )
except Exception as error:
logger.error(
"Function task errored: `%s`\n%s\n",
running_task.executable.__name__,
str(error),
- exc_info = True,
+ exc_info=True,
+ )
+ send_telemetry(
+ "Function Task Errored",
+ user.id,
+ {
+ "task_name": running_task.executable.__name__,
+ },
)
- send_telemetry("Function Task Errored", user.id, {
- "task_name": running_task.executable.__name__,
- })
raise error
-async def run_tasks_with_telemetry(tasks: list[Task], data, pipeline_name: str):
+async def run_tasks_with_telemetry(tasks: list[Task], data, pipeline_name: str):
config = get_current_settings()
-
- logger.debug("\nRunning pipeline with configuration:\n%s\n", json.dumps(config, indent = 1))
-
+
+ logger.debug("\nRunning pipeline with configuration:\n%s\n", json.dumps(config, indent=1))
+
user = await get_default_user()
-
+
try:
logger.info("Pipeline run started: `%s`", pipeline_name)
- send_telemetry("Pipeline Run Started",
- user.id,
- additional_properties = {"pipeline_name": pipeline_name, } | config
- )
-
+ send_telemetry(
+ "Pipeline Run Started",
+ user.id,
+ additional_properties={
+ "pipeline_name": pipeline_name,
+ }
+ | config,
+ )
+
async for result in run_tasks_base(tasks, data, user):
yield result
logger.info("Pipeline run completed: `%s`", pipeline_name)
- send_telemetry("Pipeline Run Completed",
- user.id,
- additional_properties = {"pipeline_name": pipeline_name, }
- )
+ send_telemetry(
+ "Pipeline Run Completed",
+ user.id,
+ additional_properties={
+ "pipeline_name": pipeline_name,
+ },
+ )
except Exception as error:
logger.error(
"Pipeline run errored: `%s`\n%s\n",
pipeline_name,
str(error),
- exc_info = True,
+ exc_info=True,
+ )
+ send_telemetry(
+ "Pipeline Run Errored",
+ user.id,
+ additional_properties={
+ "pipeline_name": pipeline_name,
+ }
+ | config,
)
- send_telemetry("Pipeline Run Errored",
- user.id,
- additional_properties = {"pipeline_name": pipeline_name, } | config
- )
raise error
-async def run_tasks(tasks: list[Task], data = None, pipeline_name: str = "default_pipeline"):
-
+
+async def run_tasks(tasks: list[Task], data=None, pipeline_name: str = "default_pipeline"):
async for result in run_tasks_with_telemetry(tasks, data, pipeline_name):
yield result
diff --git a/cognee/modules/pipelines/tasks/Task.py b/cognee/modules/pipelines/tasks/Task.py
index 7de60643a..753152d0d 100644
--- a/cognee/modules/pipelines/tasks/Task.py
+++ b/cognee/modules/pipelines/tasks/Task.py
@@ -1,6 +1,7 @@
from typing import Union, Callable, Any, Coroutine, Generator, AsyncGenerator
-class Task():
+
+class Task:
executable: Union[
Callable[..., Any],
Callable[..., Coroutine[Any, Any, Any]],
@@ -12,12 +13,9 @@ class Task():
}
default_params: dict[str, Any] = {}
- def __init__(self, executable, *args, task_config = None, **kwargs):
+ def __init__(self, executable, *args, task_config=None, **kwargs):
self.executable = executable
- self.default_params = {
- "args": args,
- "kwargs": kwargs
- }
+ self.default_params = {"args": args, "kwargs": kwargs}
if task_config is not None:
self.task_config = task_config
@@ -27,6 +25,6 @@ def __init__(self, executable, *args, task_config = None, **kwargs):
def run(self, *args, **kwargs):
combined_args = args + self.default_params["args"]
- combined_kwargs = { **self.default_params["kwargs"], **kwargs }
+ combined_kwargs = {**self.default_params["kwargs"], **kwargs}
return self.executable(*combined_args, **combined_kwargs)
diff --git a/cognee/modules/retrieval/brute_force_triplet_search.py b/cognee/modules/retrieval/brute_force_triplet_search.py
index b5ee5b612..fdd312480 100644
--- a/cognee/modules/retrieval/brute_force_triplet_search.py
+++ b/cognee/modules/retrieval/brute_force_triplet_search.py
@@ -12,6 +12,7 @@
def format_triplets(edges):
print("\n\n\n")
+
def filter_attributes(obj, attributes):
"""Helper function to filter out non-None properties, including nested dicts."""
result = {}
@@ -20,7 +21,9 @@ def filter_attributes(obj, attributes):
if value is not None:
# If the value is a dict, extract relevant keys from it
if isinstance(value, dict):
- nested_values = {k: v for k, v in value.items() if k in attributes and v is not None}
+ nested_values = {
+ k: v for k, v in value.items() if k in attributes and v is not None
+ }
result[attr] = nested_values
else:
result[attr] = value
@@ -40,17 +43,15 @@ def filter_attributes(obj, attributes):
edge_info = {key: value for key, value in edge_attributes.items() if value is not None}
# Create the formatted triplet
- triplet = (
- f"Node1: {node1_info}\n"
- f"Edge: {edge_info}\n"
- f"Node2: {node2_info}\n\n\n"
- )
+ triplet = f"Node1: {node1_info}\n" f"Edge: {edge_info}\n" f"Node2: {node2_info}\n\n\n"
triplets.append(triplet)
return "".join(triplets)
-async def brute_force_triplet_search(query: str, user: User = None, top_k = 5, collections = None) -> list:
+async def brute_force_triplet_search(
+ query: str, user: User = None, top_k=5, collections=None
+) -> list:
if user is None:
user = await get_default_user()
@@ -61,7 +62,9 @@ async def brute_force_triplet_search(query: str, user: User = None, top_k = 5, c
return retrieved_results
-def delete_duplicated_vector_db_elements(collections, results): #:TODO: This is just for now to fix vector db duplicates
+def delete_duplicated_vector_db_elements(
+ collections, results
+): #:TODO: This is just for now to fix vector db duplicates
results_dict = {}
for collection, results in zip(collections, results):
seen_ids = set()
@@ -78,22 +81,19 @@ def delete_duplicated_vector_db_elements(collections, results): #:TODO: This is
async def brute_force_search(
- query: str,
- user: User,
- top_k: int,
- collections: List[str] = None
+ query: str, user: User, top_k: int, collections: List[str] = None
) -> list:
"""
- Performs a brute force search to retrieve the top triplets from the graph.
+ Performs a brute force search to retrieve the top triplets from the graph.
- Args:
- query (str): The search query.
- user (User): The user performing the search.
- top_k (int): The number of top results to retrieve.
- collections (Optional[List[str]]): List of collections to query. Defaults to predefined collections.
+ Args:
+ query (str): The search query.
+ user (User): The user performing the search.
+ top_k (int): The number of top results to retrieve.
+ collections (Optional[List[str]]): List of collections to query. Defaults to predefined collections.
- Returns:
- list: The top triplet results.
+ Returns:
+ list: The top triplet results.
"""
if not query or not isinstance(query, str):
raise ValueError("The query must be a non-empty string.")
@@ -101,7 +101,12 @@ async def brute_force_search(
raise ValueError("top_k must be a positive integer.")
if collections is None:
- collections = ["entity_name", "text_summary_text", "entity_type_name", "document_chunk_text"]
+ collections = [
+ "entity_name",
+ "text_summary_text",
+ "entity_type_name",
+ "document_chunk_text",
+ ]
try:
vector_engine = get_vector_engine()
@@ -114,7 +119,10 @@ async def brute_force_search(
try:
results = await asyncio.gather(
- *[vector_engine.get_distance_from_collection_elements(collection, query_text=query) for collection in collections]
+ *[
+ vector_engine.get_distance_from_collection_elements(collection, query_text=query)
+ for collection in collections
+ ]
)
############################################# :TODO: Change when vector db does not contain duplicates
@@ -124,13 +132,11 @@ async def brute_force_search(
memory_fragment = CogneeGraph()
- await memory_fragment.project_graph_from_db(graph_engine,
- node_properties_to_project=['id',
- 'description',
- 'name',
- 'type',
- 'text'],
- edge_properties_to_project=['relationship_name'])
+ await memory_fragment.project_graph_from_db(
+ graph_engine,
+ node_properties_to_project=["id", "description", "name", "type", "text"],
+ edge_properties_to_project=["relationship_name"],
+ )
await memory_fragment.map_vector_distances_to_graph_nodes(node_distances=node_distances)
@@ -145,6 +151,8 @@ async def brute_force_search(
return results
except Exception as e:
- logging.error("Error during brute force search for user: %s, query: %s. Error: %s", user.id, query, e)
+ logging.error(
+ "Error during brute force search for user: %s, query: %s. Error: %s", user.id, query, e
+ )
send_telemetry("cognee.brute_force_triplet_search EXECUTION FAILED", user.id)
raise RuntimeError("An error occurred during brute force search") from e
diff --git a/cognee/modules/retrieval/description_to_codepart_search.py b/cognee/modules/retrieval/description_to_codepart_search.py
index e1da9a43f..ecd187907 100644
--- a/cognee/modules/retrieval/description_to_codepart_search.py
+++ b/cognee/modules/retrieval/description_to_codepart_search.py
@@ -10,7 +10,7 @@
from cognee.shared.utils import send_telemetry
-async def code_description_to_code_part_search(query: str, user: User = None, top_k = 2) -> list:
+async def code_description_to_code_part_search(query: str, user: User = None, top_k=2) -> list:
if user is None:
user = await get_default_user()
@@ -21,12 +21,7 @@ async def code_description_to_code_part_search(query: str, user: User = None, to
return retrieved_codeparts
-
-async def code_description_to_code_part(
- query: str,
- user: User,
- top_k: int
-) -> List[str]:
+async def code_description_to_code_part(query: str, user: User, top_k: int) -> List[str]:
"""
Maps a code description query to relevant code parts using a CodeGraph pipeline.
@@ -55,12 +50,12 @@ async def code_description_to_code_part(
raise RuntimeError("System initialization error. Please try again later.") from init_error
send_telemetry("code_description_to_code_part_search EXECUTION STARTED", user.id)
- logging.info("Search initiated by user %s with query: '%s' and top_k: %d", user.id, query, top_k)
+ logging.info(
+ "Search initiated by user %s with query: '%s' and top_k: %d", user.id, query, top_k
+ )
try:
- results = await vector_engine.search(
- "code_summary_text", query_text=query, limit=top_k
- )
+ results = await vector_engine.search("code_summary_text", query_text=query, limit=top_k)
if not results:
logging.warning("No results found for query: '%s' by user: %s", query, user.id)
return []
@@ -68,8 +63,8 @@ async def code_description_to_code_part(
memory_fragment = CogneeGraph()
await memory_fragment.project_graph_from_db(
graph_engine,
- node_properties_to_project=['id', 'type', 'text', 'source_code'],
- edge_properties_to_project=['relationship_name']
+ node_properties_to_project=["id", "type", "text", "source_code"],
+ edge_properties_to_project=["relationship_name"],
)
code_pieces_to_return = set()
@@ -84,24 +79,32 @@ async def code_description_to_code_part(
for code_file in node_to_search_from.get_skeleton_neighbours():
for code_file_edge in code_file.get_skeleton_edges():
- if code_file_edge.get_attribute('relationship_name') == 'contains':
+ if code_file_edge.get_attribute("relationship_name") == "contains":
code_pieces_to_return.add(code_file_edge.get_destination_node())
- logging.info("Search completed for user: %s, query: '%s'. Found %d code pieces.",
- user.id, query, len(code_pieces_to_return))
+ logging.info(
+ "Search completed for user: %s, query: '%s'. Found %d code pieces.",
+ user.id,
+ query,
+ len(code_pieces_to_return),
+ )
return list(code_pieces_to_return)
except Exception as exec_error:
logging.error(
"Error during code description to code part search for user: %s, query: '%s'. Error: %s",
- user.id, query, exec_error, exc_info=True
+ user.id,
+ query,
+ exec_error,
+ exc_info=True,
)
send_telemetry("code_description_to_code_part_search EXECUTION FAILED", user.id)
raise RuntimeError("An error occurred while processing your request.") from exec_error
if __name__ == "__main__":
+
async def main():
query = "I am looking for a class with blue eyes"
user = None
@@ -112,5 +115,3 @@ async def main():
print(f"An error occurred: {e}")
asyncio.run(main())
-
-
diff --git a/cognee/modules/search/models/Query.py b/cognee/modules/search/models/Query.py
index 182196333..ac19286b6 100644
--- a/cognee/modules/search/models/Query.py
+++ b/cognee/modules/search/models/Query.py
@@ -3,14 +3,15 @@
from sqlalchemy import Column, DateTime, String, UUID
from cognee.infrastructure.databases.relational import Base
+
class Query(Base):
__tablename__ = "queries"
- id = Column(UUID, primary_key = True, default = uuid4)
+ id = Column(UUID, primary_key=True, default=uuid4)
text = Column(String)
query_type = Column(String)
user_id = Column(UUID)
- created_at = Column(DateTime(timezone = True), default = lambda: datetime.now(timezone.utc))
- updated_at = Column(DateTime(timezone = True), onupdate = lambda: datetime.now(timezone.utc))
+ created_at = Column(DateTime(timezone=True), default=lambda: datetime.now(timezone.utc))
+ updated_at = Column(DateTime(timezone=True), onupdate=lambda: datetime.now(timezone.utc))
diff --git a/cognee/modules/search/models/Result.py b/cognee/modules/search/models/Result.py
index acda59ddb..cff097704 100644
--- a/cognee/modules/search/models/Result.py
+++ b/cognee/modules/search/models/Result.py
@@ -3,14 +3,15 @@
from sqlalchemy import Column, DateTime, Text, UUID
from cognee.infrastructure.databases.relational import Base
+
class Result(Base):
__tablename__ = "results"
- id = Column(UUID, primary_key = True, default = uuid4)
+ id = Column(UUID, primary_key=True, default=uuid4)
value = Column(Text)
query_id = Column(UUID)
- user_id = Column(UUID, index = True)
+ user_id = Column(UUID, index=True)
- created_at = Column(DateTime(timezone = True), default = lambda: datetime.now(timezone.utc))
- updated_at = Column(DateTime(timezone = True), onupdate = lambda: datetime.now(timezone.utc))
+ created_at = Column(DateTime(timezone=True), default=lambda: datetime.now(timezone.utc))
+ updated_at = Column(DateTime(timezone=True), onupdate=lambda: datetime.now(timezone.utc))
diff --git a/cognee/modules/search/operations/get_history.py b/cognee/modules/search/operations/get_history.py
index 831c4acc2..4ce6bac0c 100644
--- a/cognee/modules/search/operations/get_history.py
+++ b/cognee/modules/search/operations/get_history.py
@@ -4,24 +4,17 @@
from ..models.Query import Query
from ..models.Result import Result
+
async def get_history(user_id: UUID, limit: int = 10) -> list[Result]:
db_engine = get_relational_engine()
queries_query = select(
- Query.id,
- Query.text.label("text"),
- Query.created_at,
- literal("user").label("user")
- ) \
- .filter(Query.user_id == user_id)
+ Query.id, Query.text.label("text"), Query.created_at, literal("user").label("user")
+ ).filter(Query.user_id == user_id)
results_query = select(
- Result.id,
- Result.value.label("text"),
- Result.created_at,
- literal("system").label("user")
- ) \
- .filter(Result.user_id == user_id)
+ Result.id, Result.value.label("text"), Result.created_at, literal("system").label("user")
+ ).filter(Result.user_id == user_id)
history_query = queries_query.union(results_query).order_by("created_at").limit(limit)
diff --git a/cognee/modules/search/operations/get_queries.py b/cognee/modules/search/operations/get_queries.py
index ded10a8e5..921dbfff3 100644
--- a/cognee/modules/search/operations/get_queries.py
+++ b/cognee/modules/search/operations/get_queries.py
@@ -3,15 +3,18 @@
from cognee.infrastructure.databases.relational import get_relational_engine
from ..models.Query import Query
+
async def get_queries(user_id: UUID, limit: int) -> list[Query]:
db_engine = get_relational_engine()
async with db_engine.get_async_session() as session:
- queries = (await session.scalars(
- select(Query)
+ queries = (
+ await session.scalars(
+ select(Query)
.filter(Query.user_id == user_id)
.order_by(Query.created_at.desc())
.limit(limit)
- )).all()
+ )
+ ).all()
return queries
diff --git a/cognee/modules/search/operations/get_results.py b/cognee/modules/search/operations/get_results.py
index 7f90a3f0f..a15cda44e 100644
--- a/cognee/modules/search/operations/get_results.py
+++ b/cognee/modules/search/operations/get_results.py
@@ -3,15 +3,18 @@
from cognee.infrastructure.databases.relational import get_relational_engine
from ..models.Result import Result
+
async def get_results(user_id: UUID, limit: int = 10) -> list[Result]:
db_engine = get_relational_engine()
async with db_engine.get_async_session() as session:
- results = (await session.scalars(
- select(Result)
+ results = (
+ await session.scalars(
+ select(Result)
.filter(Result.user_id == user_id)
.order_by(Result.created_at.desc())
.limit(limit)
- )).all()
+ )
+ ).all()
return results
diff --git a/cognee/modules/search/operations/log_query.py b/cognee/modules/search/operations/log_query.py
index 02ed3f157..d46a82890 100644
--- a/cognee/modules/search/operations/log_query.py
+++ b/cognee/modules/search/operations/log_query.py
@@ -2,14 +2,15 @@
from cognee.infrastructure.databases.relational import get_relational_engine
from ..models.Query import Query
+
async def log_query(query_text: str, query_type: str, user_id: UUID) -> Query:
db_engine = get_relational_engine()
async with db_engine.get_async_session() as session:
query = Query(
- text = query_text,
- query_type = query_type,
- user_id = user_id,
+ text=query_text,
+ query_type=query_type,
+ user_id=user_id,
)
session.add(query)
diff --git a/cognee/modules/search/operations/log_result.py b/cognee/modules/search/operations/log_result.py
index b81e0b447..e03c45185 100644
--- a/cognee/modules/search/operations/log_result.py
+++ b/cognee/modules/search/operations/log_result.py
@@ -2,14 +2,17 @@
from cognee.infrastructure.databases.relational import get_relational_engine
from ..models.Result import Result
+
async def log_result(query_id: UUID, result: str, user_id: UUID):
db_engine = get_relational_engine()
async with db_engine.get_async_session() as session:
- session.add(Result(
- value = result,
- query_id = query_id,
- user_id = user_id,
- ))
+ session.add(
+ Result(
+ value=result,
+ query_id=query_id,
+ user_id=user_id,
+ )
+ )
await session.commit()
diff --git a/cognee/modules/settings/get_current_settings.py b/cognee/modules/settings/get_current_settings.py
index 3d6bad896..2fb1d3c81 100644
--- a/cognee/modules/settings/get_current_settings.py
+++ b/cognee/modules/settings/get_current_settings.py
@@ -4,28 +4,34 @@
from cognee.infrastructure.databases.vector import get_vectordb_config
from cognee.infrastructure.databases.relational.config import get_relational_config
+
class LLMConfig(TypedDict):
model: str
provider: str
+
class VectorDBConfig(TypedDict):
url: str
provider: str
+
class GraphDBConfig(TypedDict):
url: str
provider: str
+
class RelationalConfig(TypedDict):
url: str
provider: str
+
class SettingsDict(TypedDict):
llm: LLMConfig
graph: GraphDBConfig
vector: VectorDBConfig
relational: RelationalConfig
+
def get_current_settings() -> SettingsDict:
llm_config = get_llm_config()
graph_config = get_graph_config()
@@ -33,22 +39,22 @@ def get_current_settings() -> SettingsDict:
relational_config = get_relational_config()
return dict(
- llm = {
+ llm={
"provider": llm_config.llm_provider,
"model": llm_config.llm_model,
},
- graph = {
+ graph={
"provider": graph_config.graph_database_provider,
"url": graph_config.graph_database_url or graph_config.graph_file_path,
},
- vector = {
+ vector={
"provider": vector_config.vector_db_provider,
"url": vector_config.vector_db_url,
},
- relational = {
+ relational={
"provider": relational_config.db_provider,
- "url": f"{relational_config.db_host}:{relational_config.db_port}" \
- if relational_config.db_host \
- else f"{relational_config.db_path}/{relational_config.db_name}",
+ "url": f"{relational_config.db_host}:{relational_config.db_port}"
+ if relational_config.db_host
+ else f"{relational_config.db_path}/{relational_config.db_name}",
},
)
diff --git a/cognee/modules/settings/get_settings.py b/cognee/modules/settings/get_settings.py
index b67b9d6ab..93fd67cff 100644
--- a/cognee/modules/settings/get_settings.py
+++ b/cognee/modules/settings/get_settings.py
@@ -3,15 +3,18 @@
from cognee.infrastructure.databases.vector import get_vectordb_config
from cognee.infrastructure.llm import get_llm_config
+
class ConfigChoice(BaseModel):
value: str
label: str
+
class ModelName(Enum):
openai = "openai"
ollama = "ollama"
anthropic = "anthropic"
-
+
+
class LLMConfig(BaseModel):
api_key: str
model: ConfigChoice
@@ -19,98 +22,130 @@ class LLMConfig(BaseModel):
models: dict[str, list[ConfigChoice]]
providers: list[ConfigChoice]
+
class VectorDBConfig(BaseModel):
api_key: str
url: str
provider: ConfigChoice
providers: list[ConfigChoice]
+
class SettingsDict(BaseModel):
llm: LLMConfig
vector_db: VectorDBConfig
+
def get_settings() -> SettingsDict:
llm_config = get_llm_config()
- vector_dbs = [{
- "value": "weaviate",
- "label": "Weaviate",
- }, {
- "value": "qdrant",
- "label": "Qdrant",
- }, {
- "value": "lancedb",
- "label": "LanceDB",
- }, {
- "value": "pgvector",
- "label": "PGVector",
- }]
+ vector_dbs = [
+ {
+ "value": "weaviate",
+ "label": "Weaviate",
+ },
+ {
+ "value": "qdrant",
+ "label": "Qdrant",
+ },
+ {
+ "value": "lancedb",
+ "label": "LanceDB",
+ },
+ {
+ "value": "pgvector",
+ "label": "PGVector",
+ },
+ ]
vector_config = get_vectordb_config()
- llm_providers = [{
- "value": "openai",
- "label": "OpenAI",
- }, {
- "value": "ollama",
- "label": "Ollama",
- }, {
- "value": "anthropic",
- "label": "Anthropic",
- }]
-
- return SettingsDict.model_validate(dict(
- llm = {
- "provider": {
- "label": llm_config.llm_provider,
- "value": llm_config.llm_provider,
- } if llm_config.llm_provider else llm_providers[0],
- "model": {
- "value": llm_config.llm_model,
- "label": llm_config.llm_model,
- } if llm_config.llm_model else None,
- "api_key": (llm_config.llm_api_key[:-10] + "**********") if llm_config.llm_api_key else None,
- "providers": llm_providers,
- "models": {
- "openai": [{
- "value": "gpt-4o-mini",
- "label": "gpt-4o-mini",
- }, {
- "value": "gpt-4o",
- "label": "gpt-4o",
- }, {
- "value": "gpt-4-turbo",
- "label": "gpt-4-turbo",
- }, {
- "value": "gpt-3.5-turbo",
- "label": "gpt-3.5-turbo",
- }],
- "ollama": [{
- "value": "llama3",
- "label": "llama3",
- }, {
- "value": "mistral",
- "label": "mistral",
- }],
- "anthropic": [{
- "value": "Claude 3 Opus",
- "label": "Claude 3 Opus",
- }, {
- "value": "Claude 3 Sonnet",
- "label": "Claude 3 Sonnet",
- }, {
- "value": "Claude 3 Haiku",
- "label": "Claude 3 Haiku",
- }]
- },
+ llm_providers = [
+ {
+ "value": "openai",
+ "label": "OpenAI",
},
- vector_db = {
- "provider": {
- "label": vector_config.vector_db_provider,
- "value": vector_config.vector_db_provider.lower(),
- },
- "url": vector_config.vector_db_url,
- "api_key": vector_config.vector_db_key,
- "providers": vector_dbs,
+ {
+ "value": "ollama",
+ "label": "Ollama",
},
- ))
+ {
+ "value": "anthropic",
+ "label": "Anthropic",
+ },
+ ]
+
+ return SettingsDict.model_validate(
+ dict(
+ llm={
+ "provider": {
+ "label": llm_config.llm_provider,
+ "value": llm_config.llm_provider,
+ }
+ if llm_config.llm_provider
+ else llm_providers[0],
+ "model": {
+ "value": llm_config.llm_model,
+ "label": llm_config.llm_model,
+ }
+ if llm_config.llm_model
+ else None,
+ "api_key": (llm_config.llm_api_key[:-10] + "**********")
+ if llm_config.llm_api_key
+ else None,
+ "providers": llm_providers,
+ "models": {
+ "openai": [
+ {
+ "value": "gpt-4o-mini",
+ "label": "gpt-4o-mini",
+ },
+ {
+ "value": "gpt-4o",
+ "label": "gpt-4o",
+ },
+ {
+ "value": "gpt-4-turbo",
+ "label": "gpt-4-turbo",
+ },
+ {
+ "value": "gpt-3.5-turbo",
+ "label": "gpt-3.5-turbo",
+ },
+ ],
+ "ollama": [
+ {
+ "value": "llama3",
+ "label": "llama3",
+ },
+ {
+ "value": "mistral",
+ "label": "mistral",
+ },
+ ],
+ "anthropic": [
+ {
+ "value": "Claude 3 Opus",
+ "label": "Claude 3 Opus",
+ },
+ {
+ "value": "Claude 3 Sonnet",
+ "label": "Claude 3 Sonnet",
+ },
+ {
+ "value": "Claude 3 Haiku",
+ "label": "Claude 3 Haiku",
+ },
+ ],
+ },
+ },
+ vector_db={
+ "provider": {
+ "label": vector_config.vector_db_provider,
+ "value": vector_config.vector_db_provider.lower(),
+ },
+ "url": vector_config.vector_db_url,
+ "api_key": vector_config.vector_db_key,
+ "providers": vector_dbs,
+ },
+ )
+ )
diff --git a/cognee/modules/settings/save_llm_config.py b/cognee/modules/settings/save_llm_config.py
index 83bb91d50..620d6bc32 100644
--- a/cognee/modules/settings/save_llm_config.py
+++ b/cognee/modules/settings/save_llm_config.py
@@ -1,11 +1,13 @@
from pydantic import BaseModel
from cognee.infrastructure.llm import get_llm_config
+
class LLMConfig(BaseModel):
api_key: str
model: str
provider: str
+
async def save_llm_config(new_llm_config: LLMConfig):
llm_config = get_llm_config()
diff --git a/cognee/modules/settings/save_vector_db_config.py b/cognee/modules/settings/save_vector_db_config.py
index 1e0b683ea..7c48435a8 100644
--- a/cognee/modules/settings/save_vector_db_config.py
+++ b/cognee/modules/settings/save_vector_db_config.py
@@ -2,11 +2,13 @@
from pydantic import BaseModel
from cognee.infrastructure.databases.vector import get_vectordb_config
+
class VectorDBConfig(BaseModel):
url: str
api_key: str
provider: Union[Literal["lancedb"], Literal["qdrant"], Literal["weaviate"], Literal["pgvector"]]
+
async def save_vector_db_config(vector_db_config: VectorDBConfig):
vector_config = get_vectordb_config()
diff --git a/cognee/modules/storage/utils/__init__.py b/cognee/modules/storage/utils/__init__.py
index a399e8a82..a877826db 100644
--- a/cognee/modules/storage/utils/__init__.py
+++ b/cognee/modules/storage/utils/__init__.py
@@ -2,9 +2,11 @@
from uuid import UUID
from datetime import datetime
from pydantic_core import PydanticUndefined
+from pydantic import create_model
from cognee.infrastructure.engine import DataPoint
+
class JSONEncoder(json.JSONEncoder):
def default(self, obj):
if isinstance(obj, datetime):
@@ -15,32 +17,30 @@ def default(self, obj):
return json.JSONEncoder.default(self, obj)
-from pydantic import create_model
-
def copy_model(model: DataPoint, include_fields: dict = {}, exclude_fields: list = []):
fields = {
name: (field.annotation, field.default if field.default is not None else PydanticUndefined)
- for name, field in model.model_fields.items()
- if name not in exclude_fields
+ for name, field in model.model_fields.items()
+ if name not in exclude_fields
}
- final_fields = {
- **fields,
- **include_fields
- }
+ final_fields = {**fields, **include_fields}
model = create_model(model.__name__, **final_fields)
model.model_rebuild()
return model
+
def get_own_properties(data_point: DataPoint):
properties = {}
for field_name, field_value in data_point:
- if field_name == "_metadata" \
- or isinstance(field_value, dict) \
- or isinstance(field_value, DataPoint) \
- or (isinstance(field_value, list) and isinstance(field_value[0], DataPoint)):
+ if (
+ field_name == "_metadata"
+ or isinstance(field_value, dict)
+ or isinstance(field_value, DataPoint)
+ or (isinstance(field_value, list) and isinstance(field_value[0], DataPoint))
+ ):
continue
properties[field_name] = field_value
diff --git a/cognee/modules/users/authentication/get_auth_backend.py b/cognee/modules/users/authentication/get_auth_backend.py
index 3e5017ffa..c9d2b85b7 100644
--- a/cognee/modules/users/authentication/get_auth_backend.py
+++ b/cognee/modules/users/authentication/get_auth_backend.py
@@ -7,18 +7,19 @@
JWTStrategy,
)
+
@lru_cache
def get_auth_backend():
- bearer_transport = BearerTransport(tokenUrl = "auth/jwt/login")
+ bearer_transport = BearerTransport(tokenUrl="auth/jwt/login")
def get_jwt_strategy() -> JWTStrategy[models.UP, models.ID]:
secret = os.getenv("FASTAPI_USERS_JWT_SECRET", "super_secret")
- return JWTStrategy(secret, lifetime_seconds = 3600)
+ return JWTStrategy(secret, lifetime_seconds=3600)
auth_backend = AuthenticationBackend(
- name = "jwt",
- transport = bearer_transport,
- get_strategy = get_jwt_strategy,
+ name="jwt",
+ transport=bearer_transport,
+ get_strategy=get_jwt_strategy,
)
return auth_backend
diff --git a/cognee/modules/users/authentication/methods/authenticate_user.py b/cognee/modules/users/authentication/methods/authenticate_user.py
index 5095092f4..026304695 100644
--- a/cognee/modules/users/authentication/methods/authenticate_user.py
+++ b/cognee/modules/users/authentication/methods/authenticate_user.py
@@ -4,6 +4,7 @@
from ...get_user_manager import get_user_manager_context
from ...get_user_db import get_user_db_context
+
async def authenticate_user(email: str, password: str):
try:
relational_engine = get_relational_engine()
@@ -11,7 +12,7 @@ async def authenticate_user(email: str, password: str):
async with relational_engine.get_async_session() as session:
async with get_user_db_context(session) as user_db:
async with get_user_manager_context(user_db) as user_manager:
- credentials = OAuth2PasswordRequestForm(username = email, password = password)
+ credentials = OAuth2PasswordRequestForm(username=email, password=password)
user = await user_manager.authenticate(credentials)
if user is None or not user.is_active:
return None
diff --git a/cognee/modules/users/exceptions/__init__.py b/cognee/modules/users/exceptions/__init__.py
index 70e6e9d2a..8d1a9818c 100644
--- a/cognee/modules/users/exceptions/__init__.py
+++ b/cognee/modules/users/exceptions/__init__.py
@@ -8,4 +8,4 @@
GroupNotFoundError,
UserNotFoundError,
PermissionDeniedError,
-)
\ No newline at end of file
+)
diff --git a/cognee/modules/users/exceptions/exceptions.py b/cognee/modules/users/exceptions/exceptions.py
index 7dda702db..3c27b01e5 100644
--- a/cognee/modules/users/exceptions/exceptions.py
+++ b/cognee/modules/users/exceptions/exceptions.py
@@ -28,9 +28,9 @@ def __init__(
class PermissionDeniedError(CogneeApiError):
def __init__(
- self,
- message: str = "User does not have permission on documents.",
- name: str = "PermissionDeniedError",
- status_code=status.HTTP_403_FORBIDDEN,
+ self,
+ message: str = "User does not have permission on documents.",
+ name: str = "PermissionDeniedError",
+ status_code=status.HTTP_403_FORBIDDEN,
):
super().__init__(message, name, status_code)
diff --git a/cognee/modules/users/get_fastapi_users.py b/cognee/modules/users/get_fastapi_users.py
index 81abe3551..62cf7d692 100644
--- a/cognee/modules/users/get_fastapi_users.py
+++ b/cognee/modules/users/get_fastapi_users.py
@@ -6,6 +6,7 @@
from .get_user_manager import get_user_manager
from .models.User import User
+
@lru_cache
def get_fastapi_users():
auth_backend = get_auth_backend()
diff --git a/cognee/modules/users/get_user_db.py b/cognee/modules/users/get_user_db.py
index e77ae54b4..80a6f7ba3 100644
--- a/cognee/modules/users/get_user_db.py
+++ b/cognee/modules/users/get_user_db.py
@@ -1,17 +1,20 @@
from typing import AsyncGenerator
from fastapi import Depends
+from contextlib import asynccontextmanager
from sqlalchemy.ext.asyncio import AsyncSession
from fastapi_users.db import SQLAlchemyUserDatabase
from cognee.infrastructure.databases.relational import get_relational_engine
from .models.User import User
+
async def get_async_session() -> AsyncGenerator[AsyncSession, None]:
db_engine = get_relational_engine()
async with db_engine.get_async_session() as session:
yield session
+
async def get_user_db(session: AsyncSession = Depends(get_async_session)):
yield SQLAlchemyUserDatabase(session, User)
-from contextlib import asynccontextmanager
+
get_user_db_context = asynccontextmanager(get_user_db)
diff --git a/cognee/modules/users/get_user_manager.py b/cognee/modules/users/get_user_manager.py
index 30410a985..a1bd9e174 100644
--- a/cognee/modules/users/get_user_manager.py
+++ b/cognee/modules/users/get_user_manager.py
@@ -4,7 +4,7 @@
from fastapi import Depends, Request
from fastapi_users import BaseUserManager, UUIDIDMixin, models
from fastapi_users.db import SQLAlchemyUserDatabase
-
+from contextlib import asynccontextmanager
from .get_user_db import get_user_db
from .models import User
from .methods import get_user
@@ -12,7 +12,9 @@
class UserManager(UUIDIDMixin, BaseUserManager[User, uuid.UUID]):
- reset_password_token_secret = os.getenv("FASTAPI_USERS_RESET_PASSWORD_TOKEN_SECRET", "super_secret")
+ reset_password_token_secret = os.getenv(
+ "FASTAPI_USERS_RESET_PASSWORD_TOKEN_SECRET", "super_secret"
+ )
verification_token_secret = os.getenv("FASTAPI_USERS_VERIFICATION_TOKEN_SECRET", "super_secret")
async def get(self, id: models.ID) -> models.UP:
@@ -43,8 +45,9 @@ async def on_after_request_verify(
):
print(f"Verification requested for user {user.id}. Verification token: {token}")
+
async def get_user_manager(user_db: SQLAlchemyUserDatabase = Depends(get_user_db)):
yield UserManager(user_db)
-from contextlib import asynccontextmanager
+
get_user_manager_context = asynccontextmanager(get_user_manager)
diff --git a/cognee/modules/users/methods/create_default_user.py b/cognee/modules/users/methods/create_default_user.py
index 6cd8115a7..7000a2ad6 100644
--- a/cognee/modules/users/methods/create_default_user.py
+++ b/cognee/modules/users/methods/create_default_user.py
@@ -1,16 +1,17 @@
from .create_user import create_user
+
async def create_default_user():
default_user_email = "default_user@example.com"
default_user_password = "default_password"
user = await create_user(
- email = default_user_email,
- password = default_user_password,
- is_superuser = False,
- is_active = True,
- is_verified = True,
- auto_login = True,
+ email=default_user_email,
+ password=default_user_password,
+ is_superuser=False,
+ is_active=True,
+ is_verified=True,
+ auto_login=True,
)
return user
diff --git a/cognee/modules/users/methods/create_user.py b/cognee/modules/users/methods/create_user.py
index cfe1dbd03..deae6fc8f 100644
--- a/cognee/modules/users/methods/create_user.py
+++ b/cognee/modules/users/methods/create_user.py
@@ -4,6 +4,7 @@
from ..get_user_db import get_user_db_context
from ..models.User import UserCreate
+
async def create_user(
email: str,
password: str,
@@ -20,11 +21,11 @@ async def create_user(
async with get_user_manager_context(user_db) as user_manager:
user = await user_manager.create(
UserCreate(
- email = email,
- password = password,
- is_superuser = is_superuser,
- is_active = is_active,
- is_verified = is_verified,
+ email=email,
+ password=password,
+ is_superuser=is_superuser,
+ is_active=is_active,
+ is_verified=is_verified,
)
)
diff --git a/cognee/modules/users/methods/delete_user.py b/cognee/modules/users/methods/delete_user.py
index dd0a585b6..3f92d815f 100644
--- a/cognee/modules/users/methods/delete_user.py
+++ b/cognee/modules/users/methods/delete_user.py
@@ -3,6 +3,7 @@
from ..get_user_manager import get_user_manager_context
from ..get_user_db import get_user_db_context
+
async def delete_user(email: str):
try:
relational_engine = get_relational_engine()
diff --git a/cognee/modules/users/methods/get_authenticated_user.py b/cognee/modules/users/methods/get_authenticated_user.py
index cbc3c67ff..b088d14cb 100644
--- a/cognee/modules/users/methods/get_authenticated_user.py
+++ b/cognee/modules/users/methods/get_authenticated_user.py
@@ -2,4 +2,4 @@
fastapi_users = get_fastapi_users()
-get_authenticated_user = fastapi_users.current_user(active = True, verified = True)
+get_authenticated_user = fastapi_users.current_user(active=True, verified=True)
diff --git a/cognee/modules/users/methods/get_default_user.py b/cognee/modules/users/methods/get_default_user.py
index 9732515f3..c67d9d71f 100644
--- a/cognee/modules/users/methods/get_default_user.py
+++ b/cognee/modules/users/methods/get_default_user.py
@@ -4,12 +4,16 @@
from cognee.infrastructure.databases.relational import get_relational_engine
from .create_default_user import create_default_user
+
async def get_default_user():
db_engine = get_relational_engine()
async with db_engine.get_async_session() as session:
- query = select(User).options(joinedload(User.groups))\
+ query = (
+ select(User)
+ .options(joinedload(User.groups))
.where(User.email == "default_user@example.com")
+ )
result = await session.execute(query)
user = result.scalars().first()
diff --git a/cognee/modules/users/methods/get_user.py b/cognee/modules/users/methods/get_user.py
index bc7adb0b1..fe32491d8 100644
--- a/cognee/modules/users/methods/get_user.py
+++ b/cognee/modules/users/methods/get_user.py
@@ -4,12 +4,15 @@
from cognee.infrastructure.databases.relational import get_relational_engine
from ..models import User
+
async def get_user(user_id: UUID):
db_engine = get_relational_engine()
async with db_engine.get_async_session() as session:
- user = (await session.execute(
- select(User).options(joinedload(User.groups)).where(User.id == user_id)
- )).scalar()
+ user = (
+ await session.execute(
+ select(User).options(joinedload(User.groups)).where(User.id == user_id)
+ )
+ ).scalar()
return user
diff --git a/cognee/modules/users/models/ACL.py b/cognee/modules/users/models/ACL.py
index f54d24224..c920db5f3 100644
--- a/cognee/modules/users/models/ACL.py
+++ b/cognee/modules/users/models/ACL.py
@@ -5,13 +5,14 @@
from cognee.infrastructure.databases.relational import Base
from .ACLResources import ACLResources
+
class ACL(Base):
__tablename__ = "acls"
- id = Column(UUID, primary_key = True, default = uuid4)
+ id = Column(UUID, primary_key=True, default=uuid4)
- created_at = Column(DateTime(timezone = True), default = lambda: datetime.now(timezone.utc))
- updated_at = Column(DateTime(timezone = True), onupdate = lambda: datetime.now(timezone.utc))
+ created_at = Column(DateTime(timezone=True), default=lambda: datetime.now(timezone.utc))
+ updated_at = Column(DateTime(timezone=True), onupdate=lambda: datetime.now(timezone.utc))
principal_id = Column(UUID, ForeignKey("principals.id"))
permission_id = Column(UUID, ForeignKey("permissions.id"))
@@ -20,6 +21,6 @@ class ACL(Base):
permission = relationship("Permission")
resources: Mapped[list["Resource"]] = relationship(
"Resource",
- secondary = ACLResources.__tablename__,
- back_populates = "acls",
+ secondary=ACLResources.__tablename__,
+ back_populates="acls",
)
diff --git a/cognee/modules/users/models/ACLResources.py b/cognee/modules/users/models/ACLResources.py
index 464fed2e8..020ced6c0 100644
--- a/cognee/modules/users/models/ACLResources.py
+++ b/cognee/modules/users/models/ACLResources.py
@@ -2,10 +2,11 @@
from sqlalchemy import Column, ForeignKey, DateTime, UUID
from cognee.infrastructure.databases.relational import Base
+
class ACLResources(Base):
__tablename__ = "acl_resources"
- created_at = Column(DateTime(timezone = True), default = lambda: datetime.now(timezone.utc))
+ created_at = Column(DateTime(timezone=True), default=lambda: datetime.now(timezone.utc))
- acl_id = Column(UUID, ForeignKey("acls.id"), primary_key = True)
- resource_id = Column(UUID, ForeignKey("resources.id"), primary_key = True)
+ acl_id = Column(UUID, ForeignKey("acls.id"), primary_key=True)
+ resource_id = Column(UUID, ForeignKey("resources.id"), primary_key=True)
diff --git a/cognee/modules/users/models/Group.py b/cognee/modules/users/models/Group.py
index 793decb35..8d6119327 100644
--- a/cognee/modules/users/models/Group.py
+++ b/cognee/modules/users/models/Group.py
@@ -3,17 +3,18 @@
from .Principal import Principal
from .UserGroup import UserGroup
+
class Group(Principal):
__tablename__ = "groups"
- id = Column(UUID, ForeignKey("principals.id"), primary_key = True)
+ id = Column(UUID, ForeignKey("principals.id"), primary_key=True)
- name = Column(String, unique = True, nullable = False, index = True)
+ name = Column(String, unique=True, nullable=False, index=True)
users: Mapped[list["User"]] = relationship(
"User",
- secondary = UserGroup.__tablename__,
- back_populates = "groups",
+ secondary=UserGroup.__tablename__,
+ back_populates="groups",
)
__mapper_args__ = {
diff --git a/cognee/modules/users/models/GroupPermission.py b/cognee/modules/users/models/GroupPermission.py
index eaf3630b4..fd2ec4b76 100644
--- a/cognee/modules/users/models/GroupPermission.py
+++ b/cognee/modules/users/models/GroupPermission.py
@@ -2,10 +2,11 @@
from sqlalchemy import Column, ForeignKey, DateTime, UUID
from cognee.infrastructure.databases.relational import Base
+
class GroupPermission(Base):
__tablename__ = "group_permissions"
- created_at = Column(DateTime(timezone = True), default = lambda: datetime.now(timezone.utc))
+ created_at = Column(DateTime(timezone=True), default=lambda: datetime.now(timezone.utc))
- group_id = Column(UUID, ForeignKey("groups.id"), primary_key = True)
- permission_id = Column(UUID, ForeignKey("permissions.id"), primary_key = True)
+ group_id = Column(UUID, ForeignKey("groups.id"), primary_key=True)
+ permission_id = Column(UUID, ForeignKey("permissions.id"), primary_key=True)
diff --git a/cognee/modules/users/models/Permission.py b/cognee/modules/users/models/Permission.py
index 3b1709371..46406f0b5 100644
--- a/cognee/modules/users/models/Permission.py
+++ b/cognee/modules/users/models/Permission.py
@@ -1,17 +1,19 @@
from uuid import uuid4
from datetime import datetime, timezone
+
# from sqlalchemy.orm import relationship
from sqlalchemy import Column, DateTime, String, UUID
from cognee.infrastructure.databases.relational import Base
+
class Permission(Base):
__tablename__ = "permissions"
- id = Column(UUID, primary_key = True, index = True, default = uuid4)
+ id = Column(UUID, primary_key=True, index=True, default=uuid4)
- created_at = Column(DateTime(timezone = True), default = lambda: datetime.now(timezone.utc))
- updated_at = Column(DateTime(timezone = True), onupdate = lambda: datetime.now(timezone.utc))
+ created_at = Column(DateTime(timezone=True), default=lambda: datetime.now(timezone.utc))
+ updated_at = Column(DateTime(timezone=True), onupdate=lambda: datetime.now(timezone.utc))
- name = Column(String, unique = True, nullable = False, index = True)
+ name = Column(String, unique=True, nullable=False, index=True)
# acls = relationship("ACL", back_populates = "permission")
diff --git a/cognee/modules/users/models/Principal.py b/cognee/modules/users/models/Principal.py
index dc6e51302..a595a636a 100644
--- a/cognee/modules/users/models/Principal.py
+++ b/cognee/modules/users/models/Principal.py
@@ -3,15 +3,16 @@
from sqlalchemy import Column, String, DateTime, UUID
from cognee.infrastructure.databases.relational import Base
+
class Principal(Base):
__tablename__ = "principals"
- id = Column(UUID, primary_key = True, index = True, default = uuid4)
+ id = Column(UUID, primary_key=True, index=True, default=uuid4)
- created_at = Column(DateTime(timezone = True), default = lambda: datetime.now(timezone.utc))
- updated_at = Column(DateTime(timezone = True), onupdate = lambda: datetime.now(timezone.utc))
+ created_at = Column(DateTime(timezone=True), default=lambda: datetime.now(timezone.utc))
+ updated_at = Column(DateTime(timezone=True), onupdate=lambda: datetime.now(timezone.utc))
- type = Column(String, nullable = False)
+ type = Column(String, nullable=False)
__mapper_args__ = {
"polymorphic_identity": "principal",
diff --git a/cognee/modules/users/models/Resource.py b/cognee/modules/users/models/Resource.py
index 563f96272..c7e3a148b 100644
--- a/cognee/modules/users/models/Resource.py
+++ b/cognee/modules/users/models/Resource.py
@@ -5,14 +5,15 @@
from cognee.infrastructure.databases.relational import Base
from .ACLResources import ACLResources
+
class Resource(Base):
__tablename__ = "resources"
- id = Column(UUID, primary_key = True, default = uuid4)
+ id = Column(UUID, primary_key=True, default=uuid4)
- created_at = Column(DateTime(timezone = True), default = lambda: datetime.now(timezone.utc))
- updated_at = Column(DateTime(timezone = True), onupdate = lambda: datetime.now(timezone.utc))
+ created_at = Column(DateTime(timezone=True), default=lambda: datetime.now(timezone.utc))
+ updated_at = Column(DateTime(timezone=True), onupdate=lambda: datetime.now(timezone.utc))
- resource_id = Column(UUID, nullable = False)
+ resource_id = Column(UUID, nullable=False)
- acls = relationship("ACL", secondary = ACLResources.__tablename__, back_populates = "resources")
+ acls = relationship("ACL", secondary=ACLResources.__tablename__, back_populates="resources")
diff --git a/cognee/modules/users/models/User.py b/cognee/modules/users/models/User.py
index 3536ac948..0cbe75037 100644
--- a/cognee/modules/users/models/User.py
+++ b/cognee/modules/users/models/User.py
@@ -5,15 +5,17 @@
from .Principal import Principal
from .UserGroup import UserGroup
from .Group import Group
+from fastapi_users import schemas
+
class User(SQLAlchemyBaseUserTableUUID, Principal):
__tablename__ = "users"
- id = Column(UUID, ForeignKey("principals.id"), primary_key = True)
+ id = Column(UUID, ForeignKey("principals.id"), primary_key=True)
groups: Mapped[list["Group"]] = relationship(
- secondary = UserGroup.__tablename__,
- back_populates = "users",
+ secondary=UserGroup.__tablename__,
+ back_populates="users",
)
__mapper_args__ = {
@@ -22,13 +24,15 @@ class User(SQLAlchemyBaseUserTableUUID, Principal):
# Keep these schemas in sync with User model
-from fastapi_users import schemas
+
class UserRead(schemas.BaseUser[uuid_UUID]):
pass
+
class UserCreate(schemas.BaseUserCreate):
pass
+
class UserUpdate(schemas.BaseUserUpdate):
pass
diff --git a/cognee/modules/users/models/UserGroup.py b/cognee/modules/users/models/UserGroup.py
index 5a85c9d3c..a4ce8f13d 100644
--- a/cognee/modules/users/models/UserGroup.py
+++ b/cognee/modules/users/models/UserGroup.py
@@ -2,10 +2,11 @@
from sqlalchemy import Column, ForeignKey, DateTime, UUID
from cognee.infrastructure.databases.relational import Base
+
class UserGroup(Base):
__tablename__ = "user_groups"
- created_at = Column(DateTime(timezone = True), default = lambda: datetime.now(timezone.utc))
+ created_at = Column(DateTime(timezone=True), default=lambda: datetime.now(timezone.utc))
- user_id = Column(UUID, ForeignKey("users.id"), primary_key = True)
- group_id = Column(UUID, ForeignKey("groups.id"), primary_key = True)
+ user_id = Column(UUID, ForeignKey("users.id"), primary_key=True)
+ group_id = Column(UUID, ForeignKey("groups.id"), primary_key=True)
diff --git a/cognee/modules/users/permissions/methods/check_permission_on_documents.py b/cognee/modules/users/permissions/methods/check_permission_on_documents.py
index f9c5a2258..0a292bfa0 100644
--- a/cognee/modules/users/permissions/methods/check_permission_on_documents.py
+++ b/cognee/modules/users/permissions/methods/check_permission_on_documents.py
@@ -1,4 +1,4 @@
-import logging
+import logging
from uuid import UUID
from sqlalchemy import select
from sqlalchemy.orm import joinedload
@@ -23,11 +23,13 @@ async def check_permission_on_documents(user: User, permission_type: str, docume
.join(ACL.permission)
.options(joinedload(ACL.resources))
.where(ACL.principal_id.in_([user.id, *user_group_ids]))
- .where(ACL.permission.has(name = permission_type))
+ .where(ACL.permission.has(name=permission_type))
)
acls = result.unique().scalars().all()
resource_ids = [resource.resource_id for acl in acls for resource in acl.resources]
has_permissions = all(document_id in resource_ids for document_id in document_ids)
if not has_permissions:
- raise PermissionDeniedError(message=f"User {user.email} does not have {permission_type} permission on documents")
+ raise PermissionDeniedError(
+ message=f"User {user.email} does not have {permission_type} permission on documents"
+ )
diff --git a/cognee/modules/users/permissions/methods/get_document_ids_for_user.py b/cognee/modules/users/permissions/methods/get_document_ids_for_user.py
index d439fb4f5..0d641960c 100644
--- a/cognee/modules/users/permissions/methods/get_document_ids_for_user.py
+++ b/cognee/modules/users/permissions/methods/get_document_ids_for_user.py
@@ -10,38 +10,42 @@ async def get_document_ids_for_user(user_id: UUID, datasets: list[str] = None) -
async with db_engine.get_async_session() as session:
async with session.begin():
- document_ids = (await session.scalars(
- select(Resource.resource_id)
- .join(ACL.resources)
- .join(ACL.permission)
- .where(
- ACL.principal_id == user_id,
- Permission.name == "read",
+ document_ids = (
+ await session.scalars(
+ select(Resource.resource_id)
+ .join(ACL.resources)
+ .join(ACL.permission)
+ .where(
+ ACL.principal_id == user_id,
+ Permission.name == "read",
+ )
)
- )).all()
+ ).all()
if datasets:
documents_ids_in_dataset = set()
# If datasets are specified filter out documents that aren't part of the specified datasets
for dataset in datasets:
# Find dataset id for dataset element
- dataset_id = (await session.scalars(
- select(Dataset.id)
- .where(
- Dataset.name == dataset,
- Dataset.owner_id == user_id,
+ dataset_id = (
+ await session.scalars(
+ select(Dataset.id).where(
+ Dataset.name == dataset,
+ Dataset.owner_id == user_id,
+ )
)
- )).one_or_none()
+ ).one_or_none()
# Check which documents are connected to this dataset
for document_id in document_ids:
- data_id = (await session.scalars(
- select(DatasetData.data_id)
- .where(
- DatasetData.dataset_id == dataset_id,
- DatasetData.data_id == document_id,
+ data_id = (
+ await session.scalars(
+ select(DatasetData.data_id).where(
+ DatasetData.dataset_id == dataset_id,
+ DatasetData.data_id == document_id,
+ )
)
- )).one_or_none()
+ ).one_or_none()
# If document is related to dataset added it to return value
if data_id:
diff --git a/cognee/modules/users/permissions/methods/give_permission_on_document.py b/cognee/modules/users/permissions/methods/give_permission_on_document.py
index b724be631..1e13463f8 100644
--- a/cognee/modules/users/permissions/methods/give_permission_on_document.py
+++ b/cognee/modules/users/permissions/methods/give_permission_on_document.py
@@ -2,6 +2,7 @@
from cognee.infrastructure.databases.relational import get_relational_engine
from ...models import User, ACL, Resource, Permission
+
async def give_permission_on_document(
user: User,
document_id: str,
@@ -9,15 +10,19 @@ async def give_permission_on_document(
):
db_engine = get_relational_engine()
- document_resource = Resource(resource_id = document_id)
+ document_resource = Resource(resource_id=document_id)
async with db_engine.get_async_session() as session:
- permission = (await session.execute(select(Permission).filter(Permission.name == permission_name))).scalars().first()
+ permission = (
+ (await session.execute(select(Permission).filter(Permission.name == permission_name)))
+ .scalars()
+ .first()
+ )
if permission is None:
- permission = Permission(name = permission_name)
+ permission = Permission(name=permission_name)
- acl = ACL(principal_id = user.id)
+ acl = ACL(principal_id=user.id)
acl.permission = permission
acl.resources.append(document_resource)
diff --git a/cognee/root_dir.py b/cognee/root_dir.py
index 0f96b8065..2e21d5ce3 100644
--- a/cognee/root_dir.py
+++ b/cognee/root_dir.py
@@ -2,6 +2,7 @@
ROOT_DIR = Path(__file__).resolve().parent
+
def get_absolute_path(path_from_root: str) -> str:
absolute_path = ROOT_DIR / path_from_root
return str(absolute_path.resolve())
diff --git a/cognee/shared/CodeGraphEntities.py b/cognee/shared/CodeGraphEntities.py
index 27289493d..164327da0 100644
--- a/cognee/shared/CodeGraphEntities.py
+++ b/cognee/shared/CodeGraphEntities.py
@@ -5,10 +5,8 @@
class Repository(DataPoint):
__tablename__ = "Repository"
path: str
- _metadata: dict = {
- "index_fields": [],
- "type": "Repository"
- }
+ _metadata: dict = {"index_fields": [], "type": "Repository"}
+
class CodeFile(DataPoint):
__tablename__ = "codefile"
@@ -18,19 +16,15 @@ class CodeFile(DataPoint):
depends_on: Optional[List["CodeFile"]] = None
depends_directly_on: Optional[List["CodeFile"]] = None
contains: Optional[List["CodePart"]] = None
- _metadata: dict = {
- "index_fields": [],
- "type": "CodeFile"
- }
+ _metadata: dict = {"index_fields": [], "type": "CodeFile"}
+
class CodePart(DataPoint):
__tablename__ = "codepart"
# part_of: Optional[CodeFile] = None
source_code: Optional[str] = None
- _metadata: dict = {
- "index_fields": [],
- "type": "CodePart"
- }
+ _metadata: dict = {"index_fields": [], "type": "CodePart"}
+
class SourceCodeChunk(DataPoint):
__tablename__ = "sourcecodechunk"
@@ -38,11 +32,9 @@ class SourceCodeChunk(DataPoint):
source_code: Optional[str] = None
previous_chunk: Optional["SourceCodeChunk"] = None
- _metadata: dict = {
- "index_fields": ["source_code"],
- "type": "SourceCodeChunk"
- }
+ _metadata: dict = {"index_fields": ["source_code"], "type": "SourceCodeChunk"}
+
CodeFile.model_rebuild()
CodePart.model_rebuild()
-SourceCodeChunk.model_rebuild()
\ No newline at end of file
+SourceCodeChunk.model_rebuild()
diff --git a/cognee/shared/GithubClassification.py b/cognee/shared/GithubClassification.py
index 66f14ec91..2642b9e98 100644
--- a/cognee/shared/GithubClassification.py
+++ b/cognee/shared/GithubClassification.py
@@ -4,7 +4,6 @@
from pydantic import BaseModel
-
class TextSubclass(str, Enum):
SOURCE_CODE = "Source code in various programming languages"
SHELL_SCRIPTS = "Shell commands and scripts"
@@ -12,14 +11,20 @@ class TextSubclass(str, Enum):
STYLESHEETS = "Stylesheets (CSS) and configuration files (YAML, JSON, INI)"
OTHER = "Other that does not fit into any of the above categories"
+
class ContentType(BaseModel):
"""Base class for content type, storing type of content as string."""
+
type: str = "TEXT"
+
class TextContent(ContentType):
"""Textual content class for more specific text categories."""
+
subclass: List[TextSubclass]
+
class CodeContentPrediction(BaseModel):
"""Model to predict the type of content."""
- label: TextContent
\ No newline at end of file
+
+ label: TextContent
diff --git a/cognee/shared/GithubTopology.py b/cognee/shared/GithubTopology.py
index 3315c72e4..50c74705a 100644
--- a/cognee/shared/GithubTopology.py
+++ b/cognee/shared/GithubTopology.py
@@ -1,29 +1,33 @@
-
-
from pydantic import BaseModel
from typing import List, Optional, Dict, Any, Union
+
class Relationship(BaseModel):
type: str
attributes: Optional[Dict[str, Any]] = {}
+
class Document(BaseModel):
name: str
content: str
filetype: str
+
class Directory(BaseModel):
name: str
documents: List[Document] = []
- directories: List['Directory'] = []
+ directories: List["Directory"] = []
+
# Allows recursive Directory Model
Directory.model_rebuild()
+
class RepositoryProperties(BaseModel):
custom_properties: Optional[Dict[str, Any]] = None
location: Optional[str] = None # Simplified location reference
+
class RepositoryNode(BaseModel):
node_id: str
node_type: str # 'document' or 'directory'
@@ -31,6 +35,7 @@ class RepositoryNode(BaseModel):
content: Union[Document, Directory, None] = None
relationships: List[Relationship] = []
+
class RepositoryGraphModel(BaseModel):
root: RepositoryNode
default_relationships: List[Relationship] = []
diff --git a/cognee/shared/SourceCodeGraph.py b/cognee/shared/SourceCodeGraph.py
index 3de72c5fd..6c8cea3e9 100644
--- a/cognee/shared/SourceCodeGraph.py
+++ b/cognee/shared/SourceCodeGraph.py
@@ -11,20 +11,16 @@ class Variable(DataPoint):
default_value: Optional[str] = None
data_type: str
- _metadata = {
- "index_fields": ["name"],
- "type": "Variable"
- }
+ _metadata = {"index_fields": ["name"], "type": "Variable"}
+
class Operator(DataPoint):
id: str
name: str
description: str
return_type: str
- _metadata = {
- "index_fields": ["name"],
- "type": "Operator"
- }
+ _metadata = {"index_fields": ["name"], "type": "Operator"}
+
class Class(DataPoint):
id: str
@@ -34,10 +30,8 @@ class Class(DataPoint):
extended_from_class: Optional["Class"] = None
has_methods: List["Function"]
- _metadata = {
- "index_fields": ["name"],
- "type": "Class"
- }
+ _metadata = {"index_fields": ["name"], "type": "Class"}
+
class ClassInstance(DataPoint):
id: str
@@ -47,10 +41,8 @@ class ClassInstance(DataPoint):
instantiated_by: Union["Function"]
instantiation_arguments: List[Variable]
- _metadata = {
- "index_fields": ["name"],
- "type": "ClassInstance"
- }
+ _metadata = {"index_fields": ["name"], "type": "ClassInstance"}
+
class Function(DataPoint):
id: str
@@ -60,20 +52,16 @@ class Function(DataPoint):
return_type: str
is_static: Optional[bool] = False
- _metadata = {
- "index_fields": ["name"],
- "type": "Function"
- }
+ _metadata = {"index_fields": ["name"], "type": "Function"}
+
class FunctionCall(DataPoint):
id: str
called_by: Union[Function, Literal["main"]]
function_called: Function
function_arguments: List[Any]
- _metadata = {
- "index_fields": [],
- "type": "FunctionCall"
- }
+ _metadata = {"index_fields": [], "type": "FunctionCall"}
+
class Expression(DataPoint):
id: str
@@ -81,32 +69,30 @@ class Expression(DataPoint):
description: str
expression: str
members: List[Union[Variable, Function, Operator, "Expression"]]
- _metadata = {
- "index_fields": ["name"],
- "type": "Expression"
- }
+ _metadata = {"index_fields": ["name"], "type": "Expression"}
+
class SourceCodeGraph(DataPoint):
id: str
name: str
description: str
language: str
- nodes: List[Union[
- Class,
- ClassInstance,
- Function,
- FunctionCall,
- Variable,
- Operator,
- Expression,
- ]]
- _metadata = {
- "index_fields": ["name"],
- "type": "SourceCodeGraph"
- }
+ nodes: List[
+ Union[
+ Class,
+ ClassInstance,
+ Function,
+ FunctionCall,
+ Variable,
+ Operator,
+ Expression,
+ ]
+ ]
+ _metadata = {"index_fields": ["name"], "type": "SourceCodeGraph"}
+
Class.model_rebuild()
ClassInstance.model_rebuild()
Expression.model_rebuild()
FunctionCall.model_rebuild()
-SourceCodeGraph.model_rebuild()
\ No newline at end of file
+SourceCodeGraph.model_rebuild()
diff --git a/cognee/shared/data_models.py b/cognee/shared/data_models.py
index 2a8bc8c91..d23d2841c 100644
--- a/cognee/shared/data_models.py
+++ b/cognee/shared/data_models.py
@@ -8,37 +8,51 @@
class Node(BaseModel):
"""Node in a knowledge graph."""
+
id: str
name: str
type: str
description: str
- properties: Optional[Dict[str, Any]] = Field(None, description = "A dictionary of properties associated with the node.")
+ properties: Optional[Dict[str, Any]] = Field(
+ None, description="A dictionary of properties associated with the node."
+ )
+
class Edge(BaseModel):
"""Edge in a knowledge graph."""
+
source_node_id: str
target_node_id: str
relationship_name: str
- properties: Optional[Dict[str, Any]] = Field(None, description = "A dictionary of properties associated with the edge.")
+ properties: Optional[Dict[str, Any]] = Field(
+ None, description="A dictionary of properties associated with the edge."
+ )
+
class KnowledgeGraph(BaseModel):
"""Knowledge graph."""
+
nodes: List[Node] = Field(..., default_factory=list)
edges: List[Edge] = Field(..., default_factory=list)
+
class GraphQLQuery(BaseModel):
"""GraphQL query."""
+
query: str
+
class Answer(BaseModel):
"""Answer."""
+
answer: str
+
class ChunkStrategy(Enum):
EXACT = "exact"
PARAGRAPH = "paragraph"
SENTENCE = "sentence"
- CODE = "code"
+ CODE = "code"
LANGCHAIN_CHARACTER = "langchain_character"
@@ -47,11 +61,14 @@ class ChunkEngine(Enum):
DEFAULT_ENGINE = "default"
HAYSTACK_ENGINE = "haystack"
+
class MemorySummary(BaseModel):
- """ Memory summary. """
+ """Memory summary."""
+
nodes: List[Node] = Field(..., default_factory=list)
edges: List[Edge] = Field(..., default_factory=list)
+
class TextSubclass(str, Enum):
ARTICLES = "Articles, essays, and reports"
BOOKS = "Books and manuscripts"
@@ -100,6 +117,7 @@ class TextSubclass(str, Enum):
LEGAL_AND_REGULATORY_DOCUMENTS = "Legal and Regulatory Documents"
OTHER_TEXT = "Other types of text data"
+
class AudioSubclass(str, Enum):
MUSIC_TRACKS = "Music tracks and albums"
PODCASTS = "Podcasts and radio broadcasts"
@@ -108,6 +126,7 @@ class AudioSubclass(str, Enum):
SOUND_EFFECTS = "Sound effects and ambient sounds"
OTHER_AUDIO = "Other types of audio recordings"
+
class ImageSubclass(str, Enum):
PHOTOGRAPHS = "Photographs and digital images"
ILLUSTRATIONS = "Illustrations, diagrams, and charts"
@@ -116,6 +135,7 @@ class ImageSubclass(str, Enum):
SCREENSHOTS = "Screenshots and graphical user interfaces"
OTHER_IMAGES = "Other types of images"
+
class VideoSubclass(str, Enum):
MOVIES = "Movies and short films"
DOCUMENTARIES = "Documentaries and educational videos"
@@ -124,6 +144,7 @@ class VideoSubclass(str, Enum):
LIVE_EVENTS = "Live event recordings and sports broadcasts"
OTHER_VIDEOS = "Other types of video content"
+
class MultimediaSubclass(str, Enum):
WEB_CONTENT = "Interactive web content and games"
VR_EXPERIENCES = "Virtual reality (VR) and augmented reality (AR) experiences"
@@ -132,6 +153,7 @@ class MultimediaSubclass(str, Enum):
DIGITAL_EXHIBITIONS = "Digital exhibitions and virtual tours"
OTHER_MULTIMEDIA = "Other types of multimedia content"
+
class Model3DSubclass(str, Enum):
ARCHITECTURAL_RENDERINGS = "Architectural renderings and building plans"
PRODUCT_MODELS = "Product design models and prototypes"
@@ -140,6 +162,7 @@ class Model3DSubclass(str, Enum):
VR_OBJECTS = "Virtual objects for AR/VR applications"
OTHER_3D_MODELS = "Other types of 3D models"
+
class ProceduralSubclass(str, Enum):
TUTORIALS_GUIDES = "Tutorials and step-by-step guides"
WORKFLOW_DESCRIPTIONS = "Workflow and process descriptions"
@@ -147,40 +170,51 @@ class ProceduralSubclass(str, Enum):
RECIPES = "Recipes and crafting instructions"
OTHER_PROCEDURAL = "Other types of procedural content"
+
class ContentType(BaseModel):
"""Base class for different types of content."""
+
type: str
+
class TextContent(ContentType):
type: str = "TEXTUAL_DOCUMENTS_USED_FOR_GENERAL_PURPOSES"
subclass: List[TextSubclass]
+
class AudioContent(ContentType):
type: str = "AUDIO_DOCUMENTS_USED_FOR_GENERAL_PURPOSES"
subclass: List[AudioSubclass]
+
class ImageContent(ContentType):
type: str = "IMAGE_DOCUMENTS_USED_FOR_GENERAL_PURPOSES"
subclass: List[ImageSubclass]
+
class VideoContent(ContentType):
type: str = "VIDEO_DOCUMENTS_USED_FOR_GENERAL_PURPOSES"
subclass: List[VideoSubclass]
+
class MultimediaContent(ContentType):
type: str = "MULTIMEDIA_DOCUMENTS_USED_FOR_GENERAL_PURPOSES"
subclass: List[MultimediaSubclass]
+
class Model3DContent(ContentType):
type: str = "3D_MODEL_DOCUMENTS_USED_FOR_GENERAL_PURPOSES"
subclass: List[Model3DSubclass]
+
class ProceduralContent(ContentType):
type: str = "PROCEDURAL_DOCUMENTS_USED_FOR_GENERAL_PURPOSES"
subclass: List[ProceduralSubclass]
+
class DefaultContentPrediction(BaseModel):
"""Class for a single class label prediction."""
+
label: Union[
TextContent,
AudioContent,
@@ -191,11 +225,14 @@ class DefaultContentPrediction(BaseModel):
ProceduralContent,
]
+
class SummarizedContent(BaseModel):
"""Class for a single class label summary and description."""
+
summary: str
description: str
+
class SummarizedFunction(BaseModel):
name: str
description: str
@@ -203,12 +240,14 @@ class SummarizedFunction(BaseModel):
outputs: Optional[List[str]] = None
decorators: Optional[List[str]] = None
+
class SummarizedClass(BaseModel):
name: str
description: str
methods: Optional[List[SummarizedFunction]] = None
decorators: Optional[List[str]] = None
+
class SummarizedCode(BaseModel):
high_level_summary: str
key_features: List[str]
@@ -232,50 +271,58 @@ class Relationship(BaseModel):
target: Optional[str] = None
properties: Optional[Dict[str, Any]] = None
+
class DocumentType(BaseModel):
type_id: str
description: str
- default_relationship: Relationship = Relationship(type = "is_type")
+ default_relationship: Relationship = Relationship(type="is_type")
+
class Category(BaseModel):
category_id: str
name: str
- default_relationship: Relationship = Relationship(type = "categorized_as")
+ default_relationship: Relationship = Relationship(type="categorized_as")
+
class Document(BaseModel):
id: str
type: str
title: str
+
class UserLocation(BaseModel):
location_id: str
description: str
- default_relationship: Relationship = Relationship(type = "located_in")
+ default_relationship: Relationship = Relationship(type="located_in")
+
class UserProperties(BaseModel):
custom_properties: Optional[Dict[str, Any]] = None
location: Optional[UserLocation] = None
+
class DefaultGraphModel(BaseModel):
node_id: str
user_properties: UserProperties = UserProperties()
documents: List[Document] = []
default_fields: Optional[Dict[str, Any]] = {}
- default_relationship: Relationship = Relationship(type = "has_properties")
+ default_relationship: Relationship = Relationship(type="has_properties")
class ChunkSummary(BaseModel):
text: str
chunk_id: str
+
class ChunkSummaries(BaseModel):
- """ Relevant summary and chunk id """
+ """Relevant summary and chunk id"""
+
summaries: List[ChunkSummary]
class MonitoringTool(str, Enum):
- """ Monitoring tools """
+ """Monitoring tools"""
+
LANGFUSE = "langfuse"
LLMLITE = "llmlite"
LANGSMITH = "langsmith"
-
diff --git a/cognee/shared/encode_uuid.py b/cognee/shared/encode_uuid.py
index 21fda1a91..230fd7c27 100644
--- a/cognee/shared/encode_uuid.py
+++ b/cognee/shared/encode_uuid.py
@@ -1,11 +1,12 @@
from uuid import UUID
+
def encode_uuid(uuid: UUID) -> str:
uuid_int = uuid.int
base = 52
- charset = 'abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ'
+ charset = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ"
- encoded = ''
+ encoded = ""
while len(encoded) < 36:
uuid_int, remainder = divmod(uuid_int, base)
uuid_int = uuid_int * 8
diff --git a/cognee/shared/exceptions/__init__.py b/cognee/shared/exceptions/__init__.py
index 9b86cccab..5e2ae6875 100644
--- a/cognee/shared/exceptions/__init__.py
+++ b/cognee/shared/exceptions/__init__.py
@@ -6,4 +6,4 @@
from .exceptions import (
IngestionError,
-)
\ No newline at end of file
+)
diff --git a/cognee/shared/exceptions/exceptions.py b/cognee/shared/exceptions/exceptions.py
index 101711398..4b4164995 100644
--- a/cognee/shared/exceptions/exceptions.py
+++ b/cognee/shared/exceptions/exceptions.py
@@ -1,11 +1,12 @@
from cognee.exceptions import CogneeApiError
from fastapi import status
+
class IngestionError(CogneeApiError):
def __init__(
- self,
- message: str = "Failed to load data.",
- name: str = "IngestionError",
- status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
+ self,
+ message: str = "Failed to load data.",
+ name: str = "IngestionError",
+ status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
):
- super().__init__(message, name, status_code)
\ No newline at end of file
+ super().__init__(message, name, status_code)
diff --git a/cognee/shared/utils.py b/cognee/shared/utils.py
index 4f0b1bc3b..e57decde1 100644
--- a/cognee/shared/utils.py
+++ b/cognee/shared/utils.py
@@ -1,4 +1,5 @@
-""" This module contains utility functions for the cognee. """
+"""This module contains utility functions for the cognee."""
+
import os
from typing import BinaryIO, Union
@@ -7,11 +8,13 @@
from datetime import datetime, timezone
import graphistry
import networkx as nx
-import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import tiktoken
import nltk
+import base64
+
+
import logging
import sys
@@ -26,6 +29,7 @@
# Analytics Proxy Url, currently hosted by Vercel
proxy_url = "https://test.prometh.ai"
+
def get_anonymous_id():
"""Creates or reads a anonymous user id"""
home_dir = str(pathlib.Path(pathlib.Path(__file__).parent.parent.parent.resolve()))
@@ -42,6 +46,7 @@ def get_anonymous_id():
anonymous_id = f.read()
return anonymous_id
+
def send_telemetry(event_name: str, user_id, additional_properties: dict = {}):
if os.getenv("TELEMETRY_DISABLED"):
return
@@ -60,7 +65,7 @@ def send_telemetry(event_name: str, user_id, additional_properties: dict = {}):
"properties": {
"time": current_time.strftime("%m/%d/%Y"),
"user_id": str(user_id),
- **additional_properties
+ **additional_properties,
},
}
@@ -69,6 +74,7 @@ def send_telemetry(event_name: str, user_id, additional_properties: dict = {}):
if response.status_code != 200:
print(f"Error sending telemetry through proxy: {response.status_code}")
+
def num_tokens_from_string(string: str, encoding_name: str) -> int:
"""Returns the number of tokens in a text string."""
@@ -77,12 +83,13 @@ def num_tokens_from_string(string: str, encoding_name: str) -> int:
num_tokens = len(encoding.encode(string))
return num_tokens
+
def get_file_content_hash(file_obj: Union[str, BinaryIO]) -> str:
h = hashlib.md5()
try:
if isinstance(file_obj, str):
- with open(file_obj, 'rb') as file:
+ with open(file_obj, "rb") as file:
while True:
# Reading is buffered, so we can read smaller chunks.
chunk = file.read(h.block_size)
@@ -101,6 +108,7 @@ def get_file_content_hash(file_obj: Union[str, BinaryIO]) -> str:
except IOError as e:
raise IngestionError(message=f"Failed to load data from {file}: {e}")
+
def trim_text_to_max_tokens(text: str, max_tokens: int, encoding_name: str) -> str:
"""
Trims the text so that the number of tokens does not exceed max_tokens.
@@ -132,22 +140,30 @@ def trim_text_to_max_tokens(text: str, max_tokens: int, encoding_name: str) -> s
def generate_color_palette(unique_layers):
colormap = plt.cm.get_cmap("viridis", len(unique_layers))
colors = [colormap(i) for i in range(len(unique_layers))]
- hex_colors = ["#%02x%02x%02x" % (int(rgb[0]*255), int(rgb[1]*255), int(rgb[2]*255)) for rgb in colors]
+ hex_colors = [
+ "#%02x%02x%02x" % (int(rgb[0] * 255), int(rgb[1] * 255), int(rgb[2] * 255))
+ for rgb in colors
+ ]
return dict(zip(unique_layers, hex_colors))
async def register_graphistry():
config = get_base_config()
- graphistry.register(api = 3, username = config.graphistry_username, password = config.graphistry_password)
+ graphistry.register(
+ api=3, username=config.graphistry_username, password=config.graphistry_password
+ )
def prepare_edges(graph, source, target, edge_key):
- edge_list = [{
- source: str(edge[0]),
- target: str(edge[1]),
- edge_key: str(edge[2]),
- } for edge in graph.edges(keys = True, data = True)]
+ edge_list = [
+ {
+ source: str(edge[0]),
+ target: str(edge[1]),
+ edge_key: str(edge[2]),
+ }
+ for edge in graph.edges(keys=True, data=True)
+ ]
return pd.DataFrame(edge_list)
@@ -169,7 +185,9 @@ def prepare_nodes(graph, include_size=False):
default_size = 10 # Default node size
larger_size = 20 # Size for nodes with specific keywords in their ID
keywords = ["DOCUMENT", "User"]
- node_size = larger_size if any(keyword in str(node) for keyword in keywords) else default_size
+ node_size = (
+ larger_size if any(keyword in str(node) for keyword in keywords) else default_size
+ )
node_data["size"] = node_size
nodes_data.append(node_data)
@@ -177,7 +195,9 @@ def prepare_nodes(graph, include_size=False):
return pd.DataFrame(nodes_data)
-async def render_graph(graph, include_nodes=False, include_color=False, include_size=False, include_labels=False):
+async def render_graph(
+ graph, include_nodes=False, include_color=False, include_size=False, include_labels=False
+):
await register_graphistry()
if not isinstance(graph, nx.MultiDiGraph):
@@ -193,15 +213,14 @@ async def render_graph(graph, include_nodes=False, include_color=False, include_
edges = prepare_edges(graph, "source_node", "target_node", "relationship_name")
plotter = graphistry.edges(edges, "source_node", "target_node")
- plotter = plotter.bind(edge_label = "relationship_name")
+ plotter = plotter.bind(edge_label="relationship_name")
if include_nodes:
- nodes = prepare_nodes(graph, include_size = include_size)
+ nodes = prepare_nodes(graph, include_size=include_size)
plotter = plotter.nodes(nodes, "id")
if include_size:
- plotter = plotter.bind(point_size = "size")
-
+ plotter = plotter.bind(point_size="size")
if include_color:
pass
@@ -210,10 +229,8 @@ async def render_graph(graph, include_nodes=False, include_color=False, include_
# plotter = plotter.encode_point_color("layer_description", categorical_mapping=color_palette,
# default_mapping="silver")
-
if include_labels:
- plotter = plotter.bind(point_label = "name")
-
+ plotter = plotter.bind(point_label="name")
# Visualization
url = plotter.plot(render=False, as_files=True, memoize=False)
@@ -221,14 +238,15 @@ async def render_graph(graph, include_nodes=False, include_color=False, include_
return url
-def sanitize_df(df):
- """Replace NaNs and infinities in a DataFrame with None, making it JSON compliant."""
- return df.replace([np.inf, -np.inf, np.nan], None)
+# def sanitize_df(df):
+# """Replace NaNs and infinities in a DataFrame with None, making it JSON compliant."""
+# return df.replace([np.inf, -np.inf, np.nan], None)
def get_entities(tagged_tokens):
nltk.download("maxent_ne_chunker", quiet=True)
from nltk.chunk import ne_chunk
+
return ne_chunk(tagged_tokens)
@@ -252,41 +270,205 @@ def extract_pos_tags(sentence):
return pos_tags
-def extract_named_entities(sentence):
- """Extract Named Entities from a sentence."""
- # Tokenize the sentence into words
- tagged_tokens = extract_pos_tags(sentence)
+logging.basicConfig(level=logging.INFO)
+
+
+async def convert_to_serializable_graph(G):
+ """
+ Convert a graph into a serializable format with stringified node and edge attributes.
+ """
- # Perform Named Entity Recognition (NER) on the tagged tokens
- entities = get_entities(tagged_tokens)
+ (nodes, edges) = G
- return entities
+ networkx_graph = nx.MultiDiGraph()
+ networkx_graph.add_nodes_from(nodes)
+ networkx_graph.add_edges_from(edges)
+ # Create a new graph to store the serializable version
+ new_G = nx.MultiDiGraph()
-def extract_sentiment_vader(text):
+ # Serialize nodes
+ for node, data in networkx_graph.nodes(data=True):
+ serializable_data = {k: str(v) for k, v in data.items()}
+ new_G.add_node(str(node), **serializable_data)
+
+ # Serialize edges
+ for u, v, data in networkx_graph.edges(data=True):
+ serializable_data = {k: str(v) for k, v in data.items()}
+ new_G.add_edge(str(u), str(v), **serializable_data)
+
+ return new_G
+
+
+def generate_layout_positions(G, layout_func, layout_scale):
+ """
+ Generate layout positions for the graph using the specified layout function.
"""
- Analyzes the sentiment of a given text using the VADER Sentiment Intensity Analyzer.
+ positions = layout_func(G)
+ return {str(node): (x * layout_scale, y * layout_scale) for node, (x, y) in positions.items()}
- Parameters:
- text (str): The text to analyze.
- Returns:
- dict: A dictionary containing the polarity scores for the text.
+def assign_node_colors(G, node_attribute, palette):
"""
- from nltk.sentiment import SentimentIntensityAnalyzer
+ Assign colors to nodes based on a specified attribute and a given palette.
+ """
+ unique_attrs = set(G.nodes[node].get(node_attribute, "Unknown") for node in G.nodes)
+ color_map = {attr: palette[i % len(palette)] for i, attr in enumerate(unique_attrs)}
+ return [color_map[G.nodes[node].get(node_attribute, "Unknown")] for node in G.nodes], color_map
+
+
+def embed_logo(p, layout_scale, logo_alpha, position):
+ """
+ Embed a logo into the graph visualization as a watermark.
+ """
+
+ # svg_logo = """"""
+ # Convert the SVG to a ReportLab Drawing
+
+ logo_url = "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAw8AAAHnCAYAAAD+VGEQAACAAElEQVR4Xuy9B5gsR3X2fyWRMTkbYzCO2MbYBozDB7ZxBoMNNjhgw/+zH38gsHK6ygFlCSQhoYgiSCgghBKSkAhKKCeUkJCEEtLqXkk3zPak3dn691s13dvTp6q7anpmdsKr5/k9d3Wm366q7p7pc6pOVa1aUkvqmWfWqzVrnlG1+Xk1nxIZoqwtR81io5566sdIX7N8NkL9/Czr8e+k6wuOc9lXQB/lP6Peety46iVV9R7HuewJyW+v6ziXnXrqJ0FfRsb3SGKEdntB6f+WlFq1YUNNPbXmabVx48b4ByDzhSWEEEIIIYRMOAUdSSXUYu3Ta59Vz8ZBxNLSko4fVs09tVatX7+BgcNIqHqNp1cve65sUC9t1FNPPfXSJqFe2qinflb0/YPz12o19dRTT6v5qG6Ch6efflYcSIiVohQaH6iXthBmXU8GQ7/3IdFF7pdYIVOjt3zmA/XUz7I+f55+oV7aRkAUl/vsuvVq7dpnVafTUasQTeQPGluqXjTqpS0E6qUtBOqljRBCCCFjD+KFNU89HQcSDbUq/+EsUHWIh3ppC4F6aQuBemkLgXppCyHV9zkRr6o+hXppC4F6aQuBemkLgXppC2El9BEmTz+r1q5dN2nBQ5/DzTl9/y9P6qmnnvq83ZdJ1xv89O4RbeqlTUK9tFFPPfUrrV+3fr16am7tpAUPhBBCxpaqqWnUS1sI1EtbCNRLWwip3u18FkJ9NX3K8PQbNmxUc3NPx8FD1YeFEEIIIYQQMtVsjIMHrLo0kpEHv2ESB7XJ1wtbCBOkt16nWdcHQL20hUC9tIVQVU8IIWS6wciDDh56d5UmhBBCCCGEkGWwXGsaPOQ/JIRMMVXTFKmXthCq6gkhhJAVYDl44IRpQgghhBBCSAEIHsrnPFTtIaNe2kKgXtpCoF7aQqBe2kKYcn3V+RHUS1sI1EsbIWS49Iw8TOKch6o/HNRLWwjUS1sI1EtbCNRLWwhD05cEHNRTP9V6T6iXthCol7YQKumjyGPkgRBCRoXni9sJ9dIWQlU9IYSQqQaBRzphulIUQghZUfj9JYQQQsiwSYMHvUncfCQOIIQQQgghhJAELtW6IlQN1KZX79eDTr20UU899dRLm4R6aaOe+lnRDwYu1UrCqZoXTb20hTDrejIY+r0PiS5yv8QKmRq95TMfqKd+lvX58/QL9dI2IvRqSxvTtCV5wNhS9aJRL20hUC9tIVAvbYQQQggZezDCkZnzIA+Q1Cy2EMZLHz7EQz311FM/q/peUn1NfuZDVX0K9dIWAvXSFgL10hYC9dIWwgrpETzMTd6chz6Hm3P6/l+e1FNPPfV5uy+Trjf46XsDFuqpp5566idfr/d5mEvnPJiXSv9QL20hUC9tIVAvbSFQL20hzLo+Q9XUNOqlLQTqpS0E6qUthFTvdj4Lob6aPmV4esx5mNNpS1UfFkIIIYQQQshUM9KlWv2GSRzUJl8vbCFMkN56nWZdHwD10hYC9dIWQlU9IYSQ6SYNHmqWD8OoOuRNvbSFMOl6QgghhBAyzqRLtY5q5AFU7dWiXtpCqKonU0LVNEXqpS2EqnpCCCFkBdDBAzeJI4QQMglUHSGnXtpCoF7aQqBe2kKYdf24gOBBr7ZUOPJQtYeMemkLgXppC4F6aQuBemkLYcr1VUcyqZe2EKiXNkLIcOkZeSiOiNzLNfkxHL3/Dwf10kY99dJmh3ppmyW9Hae+JOCgnvqp1ntCvbSFQL20hVBJH0UeIw+EEDIqPF/cTqiXthCq6gkhhEw1CDzS1ZYqRSGEkBWF319CCCGEDJs0eNCbxHGpTUIIIYQQQqYUexqrP0Y/0k3iSELVQG169X496NRLG/XUU0+9tEmolzbqqZ8V/WDgUq0knKp50dRLWwizrieDod/7kOgi90uskKnRWz7zgXrqZ1mfP0+/UC9tIyLdJM6kLckDxpaqF416aQuBemkLgXppI4QQQsjYgxGOzJwHeYBkMLlS/TNYffgQD/XUU0/9rOp7SfU1+ZkPVfUp1EtbCNRLWwjUS1sI1EtbCCukR/AwN3lzHvocbs7p+395Uk899dTn7b5Mut7gp+8NWKinnnrqqZ98vd7nYS6d82BeKv1DvbSFQL20hUC9tIVAvbSFMOv6DFVT06iXthCol7YQqJe2EFK92/kshPpq+pTh6THnYU6nLVV9WAghhBBCCCFTzUiXavUbJnFQm3y9sIUwQXrrdZp1fQDUS1sI1EtbCFX1hBBCpps0eKhZPgyj6pA39dIWwqTrCSGEEELIOJMu1TqqkQdQtVeLemkLoaqeTAlV0xSpl7YQquoJIYSQFUAHD9wkjhBCyCRQdYScemkLgXppC4F6aQth1vXjAoIHvdpS4chD1R4y6qUtBOqlLQTqpS0E6qUthCnXVx3JpF7aQqBe2gghw6Vn5KE4InIv1+THcPT+PxzUSxv11EubHeqlbZb0dpz6koCDeuqnWu8J9dIWAvXSFkIlfRR5jDwQQsio8HxxO6Fe2kKoqieEEDLVIPBIV1uqFIUQQlYUfn8JIYQQMmzS4EFvEselNgkhhBBCCJlS7Gms/hj9SDeJIwlVA7Xp1fv1oFMvbdRTTz310iahXtqop35W9IOBS7WScKrmRVMvbSHMup4Mhn7vQ6KL3C+xQqZGb/nMB+qpn2V9/jz9Qr20jYh0kziTtiQPGFuqXjTqpS0E6qUtBOqljRBCCCFjD0Y4MnMe5AGSweRK9c9g9eFDPNRTTz31s6rvJdXX5Gc+VNWnUC9tIVAvbSFQL20hUC9tIayQHsHD3OTNeehzuDmn7//lST311FOft/sy6XqDn743YKGeeuqpp37y9Xqfh7l0zoN5qfQP9dIWAvXSFgL10hYC9dIWwqzrM1RNTaNe2kKgXtpCoF7aQkj1buezEOqr6VOGp8echzmdtlT1YSGEEEIIIYRMNSNdqtVvmMRBbfL1whbCBOmt12nW9QFQL20hUC9tIVTVE0IImW7S4KFm+TCMqkPe1EtbCJOuJ4QQQggh40y6VOuoRh5A1V4t6qUthKp6MiVUTVOkXtpCqKonhBBCVgAdPHCTOEIIIZNA1RFy6qUtBOqlLQTqpS2EWdePCwge9GpLhSMPVXvIqJe2EKiXthCol7YQqJe2EKZcX3Ukk3ppC4F6aSOEDJeekYfiiMi9XJMfw9H7/3BQL23UUy9tdqiXtlnS23HqSwIO6qmfar0n1EtbCNRLWwiV9FHkMfJACCGjwvPF7YR6aQuhqp4QQshUg8AjXW2pUhRCCFlR+P0lhBBCyLBJgwe9SRyX2iSEEEIIIWRKsaex+mP0I90kjiRUDdSmV+/Xg069tFFPPfXUS5uEemmjnvpZ0Q8GLtVKwqmaF029tIUw63oyGPq9D4kucr/ECpkaveUzH6infpb1+fP0C/XSNiLSTeJM2pI8YGypetGol7YQqJe2EKiXNkIIIYSMPRjhyMx5kAdIBpMr1T+D1YcP8VBPPfXUz6q+l1Rfk5/5UFWfQr20hUC9tIVAvbSFQL20hbBCegQPc5M356HP4eacvv+XJ/XUU0993u7LpOsNfvregIV66qmnnvrJ1+t9HubSOQ/mpdI/1EtbCNRLWwjUS1sI1EtbCLOuz1A1NY16aQuBemkLgXppCyHVu53PQqivpk8Znh5zHuZ02lLVh4UQQgghhBAy1Yx0qVa/YRIHtcnXC1sIE6S3XqdZ1wdAvbSFQL20hVBVTwghZLpJg4ea5cMwqg55Uy9tIUy6nhBCCCGEjDPpUq2jGnkAVXu1qJe2EKrqyZRQNU2RemkLoaqeEEIIWQF08MBN4gghhEwCVUfIqZe2EKiXthCol7YQZl0/LiwHD0UjD1V7yKiXthCol7YQqJe2EKiXthCmXF91JJN6aQuBemkjhAyXnpGH4ojIvVyTH8PR+/9wUC9t1FMvbXaol7ZZ0ttx6ksCDuqpn2q9J9RLWwjUS1sIlfRRZPZ5KBx5IISQUeH54nZCvbSFUFVPCCFkqkHgkaYtVYpCCCErCr+/hBBCCBk2afCgN4njUpuEEEIIIYRMKfY0Vn+MvnzCNBkCVQO16dX79aBTL23UU0899dImoV7aqKd+VvSDgUu1knCq5kVTL20hzLqeDIZ+70Oii9wvsUKmRm/5zAfqqZ9lff48/UK9tI2IdJM4k7YkDxhbql406qUtBOqlLQTqpY0QQgghYw9GODJzHuQBksHkSvXPYPXhQzzUU0899bOq7yXV1+RnPlTVp1AvbSFQL20hUC9tIVAvbSGskB7Bw9zkzXnoc7g5p+//5Uk99dRTn7f7Mul6g5++N2Chnnrqqad+8vV6n4e5dM6Dean0D/XSFgL10hYC9dIWAvXSFsKs6zNUTU2jXtpCoF7aQqBe2kJI9W7nsxDqq+lThqfHnIc5nbZU9WEhhBBCCCGETDUjXarVb5jEQW3y9cIWwgTprddp1vUBUC9tIVAvbSFU1RNCCJlu0uChZvkwjKpD3tRLWwiTrieEEEIIIeNMulTrqEYeQNVeLeqlLYSqejIlVE1TpF7aQqiqJ4QQQlYA+JEmeOAmcYQQQsacqiPk1EtbCNRLWwjUS1sIs64fF/TIQ+mch6o9ZNRLWwjUS1sI1EtbCNRLWwhTrq86kkm9tIVAvbQRQobLcvCwtmzOg3u5Jj+Go/f/4aBe2qinXtrsUC9ts6S349SXBBzUUz/Vek+ol7YQqJe2ECrpo8js81A48kAIIaPC88XthHppC6GqnhBCyFSDwCNNW6oUhRBCVhR+fwkhhBAybNLgQW8Sx6U2CSGEEEIImVLsaaz+GH35hGkyBKoGatOr9+tBp17aqKeeeuqlTUK9tFFP/azoB0M6YTr/ASFOquZFUy9tIcy6ngyGfu9DoovcL7FCpkZv+cwH6qmfZX3+PP1CvbSNiHSTOJO2JA8YW6peNOqlLQTqpS0E6qWNEEIIIWMPRjgycx7kAZLB5Er1z2D14UM81FNPPfWzqu8l1dfkZz5U1adQL20hUC9tIVAvbSFQL20hrJAewcPc5M156HO4Oafv/+VJPfXUU5+3+zLpeoOfvjdgoZ566qmnfvL1ep+HuXTOg3mp9A/10hYC9dIWAvXSFgL10hbCrOszVE1No17aQqBe2kKgXtpCSPVu57MQ6qvpU4anx5yHOZ22VPVhIYQQQgghhEw1I12q1W+YxEFt8vXCFsIE6a3Xadb1AVAvbSFQL20hVNUTQgiZbtLgoWb5MIyqQ97US1sIk64nhBBCCCHjTLpU66hGHkDVXi3qpS2EqnoyJVRNU6Re2kKoqieEEEJWAPiRJnjgJnGEEELGnKoj5NRLWwjUS1sI1EtbCLOuHxf0yEPpnIeqPWTUS1sI1EtbCNRLWwjUS1sIU66vOpJJvbSFQL20EUKGy3LwsLZszoN7uSY/hqP3/+GgXtqop17a7FAvbbOkt+PUlwQc1FM/1XpPqJe2EKiXthAq6aPI7PNQOPJACCGjwvPF7YR6aQuhqp4QQshUg8AjTVuqFIUQQlYUfn8JIYQQMmzS4EFvEselNgkhhBBCCJlS7Gms/hh9+YRpMgSqBmrTq/frQade2qinnnrqpU1CvbRRT/2s6AdDOmE6/wEhTqrmRVMvbSHMup4Mhn7vQ6KL3C+xQqZGb/nMB+qpn2V9/jz9Qr20jYh0kziTtiQPGFuqXjTqpS0E6qUtBOqljRBCCCFjD0Y4MnMe5AGSweRK9c9g9eFDPNRTTz31s6rvJdXX5Gc+VNWnUC9tIVAvbSFQL20hUC9tIayQHsHD3OTNeehzuDmn7//lST311FOft/sy6XqDn743YKGeeuqpp37y9Xqfh7l0zoN5qfQP9dIWAvXSFgL10hYC9dIWwqzrM1RNTaNe2kKgXtpCoF7aQqBe2kKoqh8BmPMwp9OWJqCyhBBCCCGEkJVjpEu1+g2TOKhNvl7YQpggvfU6zbo+AOqlLQTqpS2EqnpCCCHTTRo81CwfhlF1yJt6aQth0vWEEEIIIWScSZdqHdXIA6jaq0W9tIVQVU+mhKppitRLWwhV9YQQQsgKAD/SBA/cJI4QQsiYU3WEnHppC4F6aQuBemkLYdb144IeeSid81C1h4x6aQuBemkLgXppC4F6aQthyvVVRzKpl7YQqJc2QshwWQ4e1pbNeXCv9erHcPT+PxzUSxv11EubHeqlbZb0dpz6koCDeuqnWu8J9dIWAvXSFkIlfRSZfR4KRx4IIWRUeL64nVAvbSFU1RNCCJlqEHikaUuVohBCyIrC7y8hhBBChk0aPOhN4rjUJiGEEEIIIVOKPY3VH6MvnzBNhkDVQG169X496NRLG/XUU0+9tEmolzbqqZ8V/WBIJ0znPyDESdW8aOqlLYRZ15PB0O99SHSR+yVWyNToLZ/5QD31s6zPn6dfqJe2EZFuEmfSluQBY0vVi0a9tIVAvbSFQL20EUIIIWTswQhHZs6DPEAymFyp/hmsPnyIh3rqqad+VvW9pPqa/MyHqvoU6qUtBOqlLQTqpS0E6qUthBXSI3iYm7w5D30ON+f0/b88qaeeeurzdl8mXW/w0/cGLNRTTz311E++Xu/zMJfOeTAvlf6hXtpCoF7aQqBe2kKgXtpCmHV9hqqpadRLWwjUS1sI1EtbCNRLWwhV9SMAcx7mdNrSBFSWEEIIIYQQsnKMdKlWv2ESB7XJ1wtbCBOkt16nWdcHQL20hUC9tIVQVU8IIWS6SYOHmuXDMKoOeVMvbSFMup4QQgghhIwz6VKtoxp5AFV7taiXthCq6smUUDVNkXppC6GqnhBCCFkB4Eea4IGbxBFCCBlzqo6QUy9tIVAvbSFQL20hzLp+XNAjD6VzHqr2kFEvbSFQL20hUC9tIVAvbSFMub7qSCb10hYC9dJGCBkuy8HD2rI5D+61Xv0Yjt7/h4N6aaOeemmzQ720zZLejlNfEnBQT/1U6z2hXtpCoF7aQqikjyKzz0PhyAMhhIwKzxe3E+qlLYSqekIIIVMNAo80balSFEIIWVH4/SWEEELIsEmDB71JHJfaJIQQQgghZEqxp7H6Y/TlE6bJEKgaqE2v3q8HnXppo5566qmXNgn10kY99bOiHwxcqpWEUzUvmnppC2HW9WQw9HsfEl3kfokVMjV6y2c+UE/9LOvz5+kX6qVtRKSbxJm0JXnA2FL1olEvbSFQL20hUC9thBBCCBl7MMKRmfMgD5AMJleqfwarDx/ioZ566qmfVX0vqb4mP/Ohqj6FemkLgXppC4F6aQuBemkLYYX0CB7mZnXOw8Benn1CvbSFQL20hUC9tIVAvbRJegMW6qmnnnrqJ1+v93mYS+c89JkDmkK9tIVAvbSFQL20hUC9tIUw6/oMVVPTqJe2EKiXthCol7YQqJe2EKrqRwDmPMzptKUJqCwhhBBCCCFk5RjpUq1+wyQOapOvF7YQJkhvvU6zrg+AemkLgXppC6GqnhBCyHSTBg81y4dhVB3ypl7aQph0PSGEEEIIGWfSpVpHNfIAqvZqUS9tIVAvbSFMjb7PNEXqx0TfJ1X1hBBCZhu8R7hJHCGEkImg6gg59dIWAvXSFgL10hbCrOvHBT3yUDrnoc8eMuqpp5566i22EMZcX3Ukg3ppC4F6aSOEDJfl4GFt2ZwH91qvfgxH7//DQb20UU+9tNmhXtpmSW/HqS8JOKinfqr1nlAvbSFQL20hVNNHZp+HwpEHQgghhBBCyMyDwCNNW6oWhRBCVhJ+fwkhhBAybNLgQW8Sx6U2CSGEEEIImVLsaaz+GH35hGkyBKoGatOr9+tBp17aqKeeeuqlTUK9tFFP/azoBwOXaiXheE4oc0K9tIUw63oyGPq9D4kucr/ECpkaveUzH6infpb1+fP0C/XSNiLSTeJM2pI8YGypetGol7YQqJe2EKiXNkIIIYSMPRjhyMx5kAdIBpMr1T+D1YcP8VBPPfXUz6q+l1Rfk5/5UFWfQr20hUC9tIVAvbSFQL20hbBCegQPc7M652FgL88+oV7aQqBe2kKgXtpCoF7aJL0BC/XUU0899ZOv1/s8zKVzHvrMAU2hXtpCoF7aQqBe2kKgXtpCmHV9hqqpadRLWwjUS1sI1EtbCNRLWwhV9SMAcx7mdNrSBFSWEEIIIYQQsnIEznmoht8wiYPa5OuFLYQJ0luv06zrA6Be2kKgXtpCqKonhBAy3aT7PNQsH4ZRdcibemkLYdL1hBBCCCFknEmXah3lhOmqvVrUS1sI1EtbCFOj7zNNkfox0fdJVT0hhJDZBu8RbhJHCCFkIqg6Qk69tIVAvbSFQL20hTDr+nFBjzwkaUv5D1P67CGjnnrqqafeYgthzPVVRzKol7YQqJc2QshwWQ4e1pbNeXCv9erHcPT+PxzUSxv11EubHeqlbZb0dpz6koCDeuqnWu8J9dIWAvXSFkI1fWT2eSgceSCEEEIIIYTMPAg80rSlalEIIWQl4feXEEIIIcMmDR7MPg9capMQQgghhJDpxJ7G6o/Rl0+YJkOgaqA2vXq/HnTqpY166qmnXtok1Esb9dTPin4wcKlWQgiZRTwnhgoSXeR+iRUyNXrLZz5QT/0s6/Pn6RfqpW1EpJvEmbQlecDYUvWiUS9tIVAvbSFQL22EEEIIGXswwpGZ8yAPkAwmV6p/qJe2EKiXthCol7YQqJe2EFZa76BmsYVAvbSFQL20hUC9tIVAvbSFMKF6BA9zszrnoWp+GPXSFgL10hYC9dIWAvXSFoKf3h2wUC9tEuqljXrqqV9pvd7nYS6d89BnDmgK9dIWAvXSFgL10hYC9dIWwqzrM1RNTaNe2kKgXtpCoF7aQqBe2kKoqh8BmPMwp9OWJqCyhBBCCCGEkJUjcM5DNfyGSRzUJl8vbCFMkN56nWZdHwD10hYC9dIWQlU9IYSQKSbK7PNQy38YTNUhb+qlLYRJ1xNCCCGEkHEmXap1lBOmq/ZqUS9tIVAvbSFMjb7PNEXqx0TfJ1X1hBBCZhu8R7hJHCGEkImg6gg59dIWwjD1PoEt9dJG/ezoixilXo88JGlL+Q9T+uwho5566qmn3mILYcz1IS8cG9RLWwjUSxshZLgsBw9ry+Y8uNd69WM4ev8fDuqljXrqpc0O9dI2S3o7Tn1JwEE99VOt94R6aQuBemkLoZo+Mvs8FI48EEIIIYQQQmYeBB5p2lK1KIQQspLw+0sIIYSQUZDZ54FLbRJCCCGEEDKd2NNY/TH68gnTZAhUDdSmV+/Xg069tFFPPfXUS5uEemmjnvpZ0VcH5XCpVkIImUU8J4YKEl3kfokVMjV6y2c+UE/9LOvz5+kX6qVtRKSbxJm0JXnA2FL1olEvbSFQL20hUC9thBBCCJkIMnMe5IeSweRK9Q/10hYC9dIWAvXSFgL10hbCSusd1Cy2EKiXthCol7YQqJe2EKiXthAmVI/gYW5W5zxUzQ+jXtpCoF7aQqBe2kKgXtpC8NO7AxbqpU1CvbRRTz31K63X+zzMpXMe+swBTaFe2kKgXtpCoF7aQqBe2kKYdX2Gqqlp1EtbCNRLWwjUS1sI1EtbCFX1IwBzHuZ02tIEVJYQQgghhBCycgTOeaiG3zCJg9rk64UthAnSW6/TrOsDoF7aQqBe2kKoqieEEDLFRJl9Hmr5D4OpOuRNvbSFMOl6QgghhBAyzqRLtY5ywnTVXi3qpS0E6qUthKnR95mmSP2Y6Pukqp4QQshsg/cIN4kjhBAyEVQdIade2kIYpt4nsKVe2qifHX0Ro9TrkYckbSn/YUqfPWTUU0899dRbbCGMuT7khWODemkLgXppI4QMl+XgYW3ZnAf3Wq9+DEfv/8NBvbRRT7202aFe2mZJb8epLwk4qKd+qvWeUC9tIVAvbSFU00dmn4fCkQdCCCGEEELIzIPAI01bqhaFEEJWEn5/CSGEEDIKMvs8cKlNQgghhBBCphN7Gqs/Rl8+YZoMgaqB2vTq/XrQqZc26qmnnnppk1AvbdRTPyv66qAcLtVKyASifySi+N+6od40NNrLf0eNeb0qgu8EwHEA9QX1Rrc9rUybuu3RWLSTTtq2DPreDfL+5c+d/9yHRBe5X2KFpO2KeuriXa9BlD8QveUzH5K2dr+7ov354/MMoPyR6XNtS9vdtXudI09I+TaoX1l9/jz9Qr20jQh8f3XwYNKW5AFjS9WLRr20hUC9tIVQQY8Xb2OhplpLxqHesG5ePfPEvHr0rrp65I66uvcHDfXTW83fcw9Gav3T86pWm1fNhZiOccCrlK8ZoB4/QmhHq2PqWNsYqXVr59XP7ovUIz+qq59c11A/+WFD//34fXW1bs282rjRBBTQ6CCpan1WArQ7vhfNxXnV7t5L3KcN6yJ9TxNqG81x+vpUaC80uL76Rz8+78b1kQ7MmvF1HFUghrJRJtqi713crtqG3vYCHIvjcF10neNnvvIzt8IkbUeb0rZvNPd7o253/O/6ed1OPNttVZvMtqfP9fJvFOzmuU7aGttqy891a7H/55oQsjJk5jzIDyWDyZXqH+qlLQTqpS2EFdJ3HQoEDc8+Na/u+W5DXXpgS532qQV14O921C6v7qgdX9JRO/zcktr2+UtqhxcvqR3jv3d6eUft9ZaO+vLfLKpvbNtW1321GTvlde2Q6Bd2sGPSZ/0tel2HDv6O1GN319W1pzTVOVu31VF/uaj2+MWO2ullHd2G7V+4pLZ7gWnPjrFtjzd11BF/tqjO3qKtrj2ppR6NtThfK3ZWtKMlyrSX3x/V9UmghP/Hvbj1vKb69j4tdfp/L6ij/25RHfyuzjLv7Kgj/2JRnfqfC+qC3VrqhjOa6uE4KISTifsHZ9Tn/hknfV5d/7WmOv6ji+qg3++oQ9/TUWf8vwV1z/caJoDwOI+TmsWW0H124QjDaXzopoa6Lq7H+Tu31Cn/vhjfy7idv4/2Lur2HvLujjr2g4u6bpcd2FS3X9BUcz+N0udFjz7lyygq34ch6XXAhvsUt3/d05F64IaGuuakpvrW6rY6+d8W1JHvN/c4uddf/OOOOuEji+rMz7XVdw9vqh9d0lRrH1tuO0YS82UUle9NVf08nP+afibxNwL/W77ZVJce0FJfi5/rY/9+UR2Sea7xN36T8FxfuGdL3XRWU3d2ZJ/roOexav2pl7YQqJe2ECZUj+BhblbnPIiXUCDUS1sI1Etbz+fd3mL8fd9VDfWNbdpqn1/tqG2ep9RnVy1p/jdmywK2iPncKqU2j/nfmNWvWFLHfmhR/fDUlnrmyUi/rJ1OSQll9behexvjMtc+Gqmrv9JUR8UO844vMe3YXNfV1Dnfjmx70I6k7dAeFTsiVx7TVGtiRwvOmh5dsZSdp5/696Xv3kc4RRghuuKLLfXlv0WAtKS26rYb7UHb0aY8n+3eP7QdAeIX/rCjLt67rR680QSCaLMriECw8uyaSH3l4wvdc5nrl1zvbeJg89L9Wn2NSBW1P3GcMZry4/jZPW+ntjowDhK2fxHuIdqz/PzmwXXAZzhm602X1C6vXVLH/cOi+sExLfXE/XXjkMPBtJQrcQd8w9Inzv76Z+fVHRc11FmfW1D7/kZHbfei5DlXuo22+518r3HPt37Oktrt9XFA8c+L6pr4u/LUw1H6HKFsV/m9hNffW4973E0phPP/nYNb6sg4+N/ppea5TtqC35/edpr2m3u8/Fwf8geL+rl+4PqG+c7gtykqKN9ik1AvbdRTX12v93mYS+c89JkDmkK9tIVAvbSFMD169FbC8brl3IY66m8X9YgCXraJY42X81abdP9NUbn/l0CbOKp7/0pHfeeQlu7d1EFEoPMocbcf58bowFM/rauL9m7p0YXEUUzb41H/PNDCOcF59nhzR120V1vNPWQczPnSkYg87vr7IfWNVk075vd+v6FHihDsJNc/CYiW21PeftNe41zCGT3+HxfV7Rc0TFnt3rLh1CHN67h/WNBl5s+F5wfP0xabKD0qlQSqVcF5kFp28zea6kuxM7nN80wggPvd02bx/EqS4xNne5dXL6kzN48DpxvqxpHupsUIqj7Lfej1Mx4HDU//LFLfPbylR3i22tTyvU0oaX8SLCdB1u6/0FHnbtfWI3Voe5ISZKWP+ofo9Tyk+LnGyNVJ/7qodnzp8vdwy/hf0VYbmfbj+CR4wkgjfvNujn/7Nm5Y7kAJ6iEtqX8p1EtbCNRLWwhV9SMAcx7mdNrSBFSWkGkHqURwfB+4rq6O/dCC2nKTjIOddzby/+8iOS53fPKy3vuXO+qq41o6tQU9ifk6VQU9pZjLcOUxLbXnW5edKVFPF4765/8/GZHY85c66gdHt3TwpVN7LHUaNmZEoKYevr2uTvnEgtrmuZZ259vjwtH+xKmGg4rUEMx3SfLkdeCwNtI99nDcbfrk//EcHPCOjtrwrNHm2+KLnuQeX+87Lm6ow/9sMQ5KcsGho/xSMsfhXAicMIJx9pZtPRIxmMC3Gvje4Ptz1QlNtc+vd1KHv+d7m29v/v9ddI9LAonVcQB17vZt9dTD3SDZUp9hkTzX+H36yj8tP9cyEJb1F+3N/3+X5LlGe4/480V152UN/T32HVEkhAyfwDkP1fAbJnFQm3y9sIUwQXrrdZp1vQd4QWJS4SX7tfUwfvJSdr5kHXZfEn2STvCl9y+q+65uqNZCrZIzlm0/nNlH76yroz+wqMvIOs+Dqn8ek9q0pPOqkU4R6mD1e/8S0BuOYAmjOju9KnMfPevvS6I3DvWSHolAPv36ZzBRHoHDwnLgYCGr3/Z5S9ohRN37aT/u85pHI/X1zyyorTdLeqBlmbby+wHnRtv2fGtHz31JVhfL12voRCbYf/DGhv7+JN+lfH1t9Nv+5H7v/WtL6vqvNkfmWGO0Ac/VRXu11I4vcz/XvpS1PwkiEKB8ffMFtfZxpG250ygIISMi6gYPmPNQy38YjByyD4N6aQth0vWzDZyvn91f15Nm0bNa+lLGi7fk5VuIRY+0g+1fuhQHLy2zyk9RWkQZkUlTQtrV7m82PbFl5QfhoUeZu7+po249Dz3ytZEMBeM+IjXrhI8u6OvpHGHxqH8hFn3iVGJC+TF/j/ItOoseOswtuPWbYalLeGcg0IDz/OMrG2r/t5v7XPrs5srvi6426d1H0LJujVlBKl/PYZGM8GAuxurXmBERUU8XA2i/bnv8LyZYw6kfZttxjx/9UV0d/qeLaWpS1fr76pPn+sB3dHTnhn5GI/P85etJCBk+etU+vc/DiEYeQD+9WtSPj54MHt1zeXNdff5tDic7b0vsRZ/Z/s4fY/nM9PYpdcwHF9WaRyJ3AFHgiOOHBbrvHdFS277Q0RPrKD/9zPZ3/piiz+J/t15lUpm2fdGSuvr4lu45TUdUCurvhUUPp+bBm+rqgN/O3MeiOhZ9Zvs7f4zjs57eb8cxWb0eeXj+kp6oGuKAJoEDVnGC85wGKxXrH6o3zqXSPf9PPlgXcz+Ggenpj9Q3d2jrCeDOdDSP+gsC9GYEJtv2wffMoxPgjm831O6/iOe6GyANqP7Cntdk/sbzhZG8G05vdidTy7oSQoYPfvu5SRwhK0jicO71yx3jfLleqPmXq+s4lz1/jOu4jENywj8tmhd0yEs6MukNl+5vcao8yxe2PIF6PVl1syV1+cGtoTkduI/ofd/tTd37WERg/QUD1iPQ2D8OeJDu5D3nQY8szavrTmuqbV+cSUfro3zxWd6Wx6FHwIZ2IE3Old5StbcaegQO65+ZV6f8hxld6hlpsdRL4Kh/FX3SM//YPXHbCwKI0PbjHmPi+w4vywWkrnq67PljXMe57F1Qh+3i5+3aE80oWbYzzKdjrKj91EtbHuqlbZT6Ikap1yMPSdpS/sOUqi9a6qUtBOqlLYQx1uPl99Pbk8DBkfKQfZlmX7rZf20v3Pxx+WPL9KtMSgSWmgzpyYXTdtWxLe2sa6eq3/Lzx+WP3SRzfg89nFssS3o1nI6QORAF9y8B9xH7F+z+RrOCVL6eg6j/MPT4F+fYYtMlvRRo3hkraj9WFUK+fU/gEFh+1frb9OgZx54J2BtCb7xnqXsW0d4SorpxzE/9FFaw6gYOA6x/Fb0OIH63o558MHKvQpVvj8WWgInRt5wXBw4v7wYOJeVXrb+vHs8bRjSvP61l9pZwPKOEkOGwHDysLZvz4O7J8GM4+qIfPuqpH2c9ei8xARCbYjlTldIXZ0e+SLvgBYwXu1m+00wiza6hj8/TnlGL3kr3OJznrM+1VVvl22PItx/O5+1xsJE6lB71t4H62tqTrMCSP17guH64HjvGjtDdV5g5EPn6u5H3LwEO6lOxo4oeb5EqFNh+7cjnqHL/yso35Sl10R7toHQlOJV3Xd7QvdHW+yHKd/x/WoflSe5Jm63H5//fAZ6Vo/5qUa82VTSS4rz/Lmc0MsHxeTu23cF+tp75+mb+P7mvtvttO976/xbQdmwsiJ3lyyZRO9sf04i/yw9cX1e7vt4jBS5Ppv35dop7bMPj+uEc279kSf3o2905ENn6u+5fDmf7q+o9oV7aQqBe2kKopo/MPg+FIw+EkIGi5wTEL/bTPml6L8XL04Mkrx07MO//Wx11zAcW1amfWFDnbNlSZ36mrU76+II67P909AZbeOl6T2TNAM1xH14w8x5KXqg45skH6mqvt2ac6ABQNx1ExXXd5TVL6oj3ddTJ/7Kg26LbE/99+PsW1W5v6PTdHoAABBt2rXnYv3fWBRxTTCzHrsAiAPQETlCixfKje/1SR9dv37d19IaAWHULk5mTYDCv75fEobtwdxM4FDnZWXDNnnwgUp//1f7vM3RozzabLamdYgfw879u2guwgzieaUzIxTHW4MQDfK++sXU3KCp5dn3BaMsPT20uj6oFAg2ev2QFIeyivi/aHt/v/eK249ne7oXJM+GxaIIDXDdsTOe/kV4v+G1a81hdp0HlU/B864SgNPtc477u95umrXh2dnhx/FxvFv5cZ8uHDr832Kk9JPglhFQDvytp2lI/PzKEkHAw1P7D01pqi80cDkK+xy1j0y/z2JmEI42lQB++rW5y1aNIpxdhlKC9ZJx5bLC05pG6zlk++d8W/JdX7JaFeQ+n/89CfL5asQPWDYZO/Fg3GHLV32JPgoadX7WkTv3PBV3Xpx6JVG2Dce5bSwbsH4B19DFag1WBcOxq1zKolnKy5UPztf82zpVoSwCYSIp7IJYlLSkfJLs8Y48NBEc3ntnUu08/u8bszwCwM/QTsWN067ea6hvbttUBb++kzrc4f74sm61r14HDJknggGdHts0GjkOQ8ZV/KgmWbOWv6j67m5idhL+1c1v96JKmmosDEaySpEGbnzK7cMNJx6Z6u78xYMQpYZPu/YgddKz2FZSm5gDPCuYTwMEPvf7JM77d85f0pnkX793Sm6vhOc+2HTu+4/t81fFNHSzv+lqzeEHadsd1zaPb/pwlddM5Tf2M5ttSBO4xvnenfbJtvsvZsj3KT0YWPv9rHfX1+Lm+4Yz4ub67rp6Zyz7X8+qJ+yN1+4VN9c3t2+rgd6FDwCzvm5aVP7ejfNTxhI8sxM6M/3NMCKlOZp8HLrVJyLDRaS4PR7KHPvtitLwkE6fxC3+0qJcexQZo6AnV69vDqcu/OLuOnt7UaaE7v+KWunaasYNrvkfRVj56D7GSjkgLyIHPbzrb9MiKc9nO3wXt2eZ5Sp3x6QXtMOI8esOvRrc92TZ124icc92e+DjstHvGZ8xGVc5rabElztXt51tSHjyBM4kN4Fa/suvYWq6frfyk93nf38LGfE319BORvo9IBao3a+k9S0BQhhQSHIMgEc7YIe+Re2aElI+5LEmqku+IA8A1xwTpLTctCdhy5SdpKlgR6NbzmmrjegSEtXSX4uX2mvbjmYbTi/LwXfnOoS2151tkL3hZ+bhGGNVY81jBqmE+4LlrIGhy7NTtKN/UQek9NBAI3X9tQ5/LtL1mnvOe+10zAXPHLDqAkbzzd22rXX/e8n3Nk7vmaduxYlpJ+lKWZIK0uMe2ZypD8lxjV+2rT2jqINA81457HNuwsR46BvA84DcNy8Am6U3587vKT34XMWcH5eXbQwjJE9ahIDH68gnTZAhUDdSmV+83Aja5erxM0eMrevUKXpB4mWKFkQtihw+byBWtGFRUPpxFaO/5bkMd+gfGGdMOgqV8zDM4/L2uFXiW24/PcMzB74RD65uCZXoZ4dzc9i3TM4yRhXx93ZjyodHzLC5o6BSfwp5LUb5SX/zjRT06k21f0fVLiYxDdtK/dp1Jy/WzgfuISdvnbL2ge5nhOOmyM/eyqHwcC+duw7ORuuygdu9EVo/yAQK8i/Z0Bw6u8nHsM3ORTq8R97mgfNyTXd+wpAMlnEdPyraUaysff8PBxnXCJOCv/t8FE/wElI/7c+HuZpWt5XLCvr/aoT6nqbbaNBkpLCm/+2/iTN91mQlSl5cK9isfzxja/vi9dXXivyzoex1SPkYOz9upbSYV95RjLx/3Zd3TkZ50nQTjIki0lI9jMf/gwr3aehQFTrwOFCxl2MrHsbg/tVqkvndkS63WIy5h5e/zKx09KlkveLZc5ffUxWKTUC9t1E+Ovjooh0u1EjIi9LyA2AlCKoZv7xqOW/3qJe28NLsv5fx5Q4EjgwmVSB1BjnWSd4yy8NJG4HDQ73X0aEBZag8chWtObPb2/OfJtQvlIVXhifvrprfQEQj5AGes1anpvGcERNaeYct1BXBM0JMfmtYCJ/Ce7zfUNs/z7J3dxFzbHV+xpEdy4MzVG/33/uAZQCoZJi2jRz699gXlJ39v/ZwldceFCNhKUtFy4D5ddpBJ0Vo+t8OJ7YJjkWr14I11nUqXBrzpv+6XYB58d/AsYhRi6+cn35/i8gGOw6TfJ36S6YEPKB91Rs/4oe8pGfnIlY/g9KR/N0GieL7y16EEtBt1v3T/lk7FsjvUveUDtB1znh7/cb135MVRPu7xD45uye+y5bom4FisMnbHxabn3+v3yVG+/i7Hz8lPrqvreVz+nQEmfenifVp+ow+O8r2hvpo+f55+oV7aRoRebQnBg0lbkgeMLVUvGvXSFgL10uYBXmxXHJZzwApejHASsDoQetb1vIOK5WdBjy6CCIxCnBw7OQf8Tkd9/jc66rD3ddSFe7Z0Oo0zcOiWbxwrrBhl0mjy9beBXuuD3rmonTmk44hz90mS5oGgpNTJ64LjsDIN9K6RHBsY8cAa/z2jDgXowAH3MXbaB7m0JIImrIiz5y927MGoBbQZu5jnRzyK0D3SayM94VXcZ0f7Uc6B8TMFx1U4z32incv4XN8/ygQQPaNmBcCxvGjvdl/10Cl5ceC+5SaOUTVL+Wj7qf/fgqptrJgulSHpnb/8C3EA8bxM54Ol/Cx4RtFJUNb25B5jg0Nxjx3gOEyEfuCGumV0o39Q15/Fz80Blgnbgm778buC+UMYfcBvW/6chJDBkpnzID+U9N9bRj31s67fuHFe9477vpzhIHz/yJbujRtE+Xk9hh7hdCNIwEgEJugijcfdg9irh2N156VN3Zvt7A3NtQe9lEWbeBVTrEF9Hrq1rnZ9nesad5cs7YI6Ix/9/mvLR1iS8uEMPnZXXe34UsfoUQ6UsXVcBiYAw9mX5wxB6vFs3HlpQ233cz73wEy4xpwX9O76tdk4czd8HQ50/nx2dG/0Lyzpyb9e97lmsblAABE/n5fsZzYhLG+zqQ9WNXp2zpaCV1B+ZEY8jv37kgniXbCjOY47/qML+rvuPdfAVX6ebvB0cRwI2Z9vCZ5RjE49Dafa1vZu+bjHN53V1Nc0fw4bOiB+2ZK664puQJw/ZwiW9uO7jDlFGKX1aSueAxx37UnhI4m28oOgXtpCoF7aQlghPYKHuVmd81A1P4x6aQthlvRw0uGwIU3I6vBYUj6+8s+Ly5OHLecMKd9Gqo9MsBDSGw0w4RErBXn3wMfHXHNSS6ew9JTfJzY9JqNeeUxLb3yWLz+tY6auqPu527X90h1wfowefRGjR+UpMwD3US+bmea7F9c/hESvU4oOtKSb5OnWEW3GKje+K/FgpKVnhSVbW7O2OJjE6lEm6HWX4dd+qdfPahT1OvW2OmlMwLjlpiqdX+NbPgKHh2+P9HK5PaMctrI2Mc4rer/nHsJKYUm9Zf19y7fpUSe04fh/zI322erUBd87fT9yz3j2+2/usW1CeCbgzpSBc37vCKQJudtnq7+1fBtxnRbVvLrlmw09T6infFtbNzHftS//zWJu3kOf5VNPPfWFer3Pw1w656E8B7QY6qUtBOqlLYTx1Runs6lTKGwvY4P5TKcrvXRJPXhj2O7OReX74a9P0hywJONyD7y71zJNl4mKllT0L99Ocu5Ir+zjk/IABwyjQUi/cterS/dzOCj63I77l4BzYyKn/2o//bUf9wKT1g95p2UyswXUHbswb1g/n2uzLB/BK1bs2e0NlpEW0f7u0pkfXUxXAcufz5sSLb4XP72toXZ+taVeFlCvMzdvm+9TybmT8vGdvfzQ/DyPDLn2b7EpVifzzLsvI1vHyCyfimAP37mbz22qL/3FonbgRZ0sICDA0sauZxD3GKta7fY6R8dGDlwPfL8wwdl5j132EvAs6xXXYv1Dt9T1hO8dXuqoV679mEi/8yuW9CpsrrZ602f9qad+LPQjAHMe5nTa0gRUlpCJJDIpDCd+LJf+4Hj54+WMeQj9LiM6CpDyct/VDZ0Ck6+/DWwIZd0NdgjA+bjlXMvSsfnrvYlxStCz/Pg95Q4H7uHP7o/UTi/3y7dH8IB9IAbiTJaAMm48o2l3KHPtR92xcdcjPypvM+7Xrd9q6L1FnOftgvOil/jeKxt6VCp/rkGDNp+1BXZ7ttQtV08EGJi8bfZEkefKg2Nwv9HD7xzdyPw/6nDEn3VSbf58/ZDMc8A9euSOurpk35beyRxOsh51yNfHAdq+91s7at1T9tQl3GMsyOAMwjLPj0nDU+a7HJoeVAB+T/C9ffrJSN349aY69sMLavsXmwUQ0sAh3978/68y9+Hak5uqje/cgO4DIUQSOOehGn7DJA5qk68XthAmSG+9TjOshzOBnYj3+63eXHxXDjkmZ2IDpX4d7Xz5ofjokZKCzaxKHbdVxon+wnsW/Xr35/3KL6IemSUn9/sN19yH5DobxwT/3nB6K73ervLhLCH1BY50osufM0HP7/j5jlrzaPgETlf5RcApfPYps/Oz0wnMgLpfe4pZ+UmcK/M3UszO36VlSWeR7U9Gl7TT5nGfqwKH86Gb63q3YmvPdLauq0zA9NNbywMmgNQXOLJ7vtk9spG0H+dG0PZDx/UMIjKjDHjWNjxjAjfsFr/DS8wIQrYu+evvAvXDBnX3XWUfycQ9Rhqb7R7n0UHS+8wiA1XvcRocxd8PzI+5aM+W2kePZJqlnEvvqaX9aMPZW7ZHErwSMrNEmX0eavkPg5FD3mFQL20hTLp+etFpAT9FWkDGscOLz/LyM5OKl/Tus6FO5yiBo/2Nrdp6Wdd8G/LtgiNwwW7tdK7DKIBTcu62jvoldezWEw6Hrp+S58mC+mO50nSOh+X+JefGMSf/q5nrUNXJ8kW3eZvcHiIWMLEX1+Xc7UvaHJln94R/tKSAWdqPlKkrj1ue0zJ0ukHKEX+KVXmK0gGXgz1sZujj4GMOwIM31fWEeqsTm2l/8p3V6WmWnn0fEFTrvSBaZifrSw9s6VXJUOee3ndL+QKLHc/E947APJRcDjPucczRf1swwpI7z/e/JOdP+ILANBllwAINGPE44Z8XdQDYs5t2GY724xlESp72Z0b0vSNk1sBvhtnnYUQjD6CfXjXqx0dPwsHLEvMXdnxJxglwvPxM+sOiXq3Fp5e+byqeG72GJ3zE4lQmbcv9/x0XNVWznXFcKpZfpsdETkwStabxJHXsfgZH+tT/wHKW7slhAA7l1z+9YAISx/1Lzg0n6OrjC5yskvr3g97M7GxsZpZ7zvL1W2UcUqTR6V54R13w/G1cF6lD322ZS5Frv+7Zf4Hp2UfPef5cwwIjYOetzvWaZ9uc+Rv37fJD4uBGFd9ngFWikD6zVdESrd1z4zuA74ItJagMvRFc/IxsWDev90o45RMLetdytKdo1Cx//cVnub/h9GOX6/wKW3ruUuzEH/QOx34huXsMJx/LA+fPU0YaHMXPBvaP+fa+bbXv28z3BEFDT3DkKL8HR/txntUvX1KP3Dna55CQWQJ+JDeJI2TINFo19ZMfNnpXbbG8+IDusf63BbfTOQboNKzavDrsvZa0oFy7tFMZt/vJ7CZdIwDOzU+ub+hUFeuOxJl6wvk75gMLCpuGOQM22KNMwOS4f4kdPda456N0YFDWI7fX1Q66zZa6ZepoJr12dDqdq806FWou0mv59/QIW55ffL5nfBx6k/txon3Jj5AjYMLGe2kKi+u+rDIO9Llbt73SARGUYE8WETRZ2o+g5OK9/Oa2ZOuPeqx9LFLfPbylDv6Djg76rKMMeSzXv+ezvG2Vud9YshUTrrP3B99JjHLu9gazjK9LD3CP94rPgWWdXc+MDf29r2En+KY6Kf5tw5whExw5rm3elqeg/WgD7geWbLXdj/zzk8WnY416aaN+dPoiRqnXIw9J2lL+w5SAHwkr1EtbCNRLWwhjoNeO7HWZ4CF5AVpehJvHL9SvfqqtFpLe0QGUL2whWPS6R3r9vNrvbV2n0vYi79qMw7GkN2/qK3iwlO8D0m2eeCBSq1/ZdYry13qT5UAOAdDh/8cxJ6P7/7rNG+KA6X1YIlP16PPtT3pokarWV5st5fsAp/CZJ+I2vzzXZkv7cV+wiZueQOxw9qN6Ta39WaR2fGnGsbRcv/Qavg+rVpl65M/lIuSFZQNpYcjl3/a5vfcgW8/EhuDh7O6yuYneVT42Zbxgt9xcD1v7V5m2m3085HlcIHC44+KGzvFHvfBM9QQNlvpby8/aLHrUDYED9mRApwR2vO4JHppmEYDtnmfXZ8+Ncx323kUz/yv5XljalgXnx1wcLACBuVxpcJQ9f0H9fduPv/W92mxJHfDOjrrjooYZHQn4/hBC/FgOHtaWzXkoH+YtZjj6sh8u6qkfB73Z4yEXPDhAr9lpn8wED5pq5Q9ar384np3XE4KtOcqZFzw+x0ovcELzczj6Ld9HjzquWzsv55nk67qq6xTp4MHo8ufS5aPN6+bVwbYUnhy4x7u8Grvdmom3/dRflC9sUo+6w1FDD3NvUNddqz/TfnyO4A8Ty13OPuxPx/dtp5eV90rDQT3uQ8axtF9DSb7+KQF6OIgPxN8tOL8iYMrV18xDWd4HoKh8fP/O+myrd86Mo/24lthEL51LUVJ/ncZ4U13t/CpHapKj/uL/HaSOdHw8Vme6YPeWTuNBIJu912i/Dh7ui/RImSgnVx7uMVLdknMUXb/kXzwLCFrQKVL221dWvvj/VUlwpNTqV5kV6m67oKm/xz6T4kvrX4JT7wn10hYC9dIWQjU9NpQtG3kghFQiO/JgdbYzJMFD4UTWFSYJHvZ4oyN4yKBHHn6pGzw4nNRhkAQP2G26rI5witKRB0cddZtjp+Tgd1lStXLAScL8ln5WWqqCTjN6yqSglLVZBw+/ERg8FIBriE3btNbT+RoE+rv1w4Z7YnOujsf/g6ljWYCD799Zn8OE+/JAEWXffXlDNTzSobTDHj8TJ/2rbUO2auC5xDkx6nXchxfVTec09HcAIyIuZzobPPhcP8zJ8B1Nw/W4/fymPm9y7rIyykjOpa/dpkvqgHd01EV7tdXPflzX5WFEp+zeEkKqgd+xNG0p4mo5hAwFa9qSA/vIQznVehHCKA0ecj3cKxc8RMvBA+pk6bkEOnh4b0nKTdRNW/o/uZ19LeAeY67FE/eb3l5xriGBuiNg2elllucs13604ZB3d9SGdZZUrcz5nv7ZvA4esHym6/oB9P7q4CGT0jIKrMGDrZ6bdIOHf8RGhfOldcT372wdPGTOaTlvEjzc8137Mqh5MOqxNr5Hu762oCPBUo6t/KwjjfuzbxwMYlM1LHuK50470q7nOamPLXhwlK+Dh3/3Dx6w6taZnzGjDrb6O3EclwRHq1/R0TueY/fp9XHwixEf3zoRQgZDZp8HBg+EDAOv4KH7wsSLtp/gYZRYg4fsCz/zd0/wMMIXfHbkQV9zh0MC4BRh8rc7bWk+TcH4yj8tmAnTeTLnR3nbbLak7rzMz6EcFHjO7r82dqSfb3nOcu1HG7BTNnSuNoeNPCwHD67zDYM0eMi22Xavu84vlp3VgUNJHXXaUnbkwXbOVcvBw92+wUP8HXjygUjvu9BzTR3fH5cN3ys40tu/cEkd88FFdf3pTfVsHCynowwl7UvrkwQPntfPN3jAM4Djjv1gdwlY2zlz57f9bYIjpSfEY5+c83dtq8fuMqs9cZSBkJWjfMI0GQJVA7Xp1fv1oE+WXgYPltV/uv8uBw/yPP2WL6mmF8FDpv55p8A+8lCtfB99FNXSOQ/S8e29/ssjD+5eeIC0iLP/15LKYmk/znnJ/q49D8rrX4xdrzfuO86xQlAPSveon/ofZvJwtsyevxE8PB6p1SJ4kO0PCx7s9c+X72ZZbw0eMmRty8GD+z4n5ScjD5+xnDPb/uXgoSmupQ041HMP17ujQ7nAxPL9ydYff6MNGGXY+1c66tzt2uqBG+o6KNfLDNfLyzcsXz8ED090g4d82fny4cSb4CFylpN+/+ZNuhSOl+lZ8vnp+XdVZqL3SzBatKCXXUaKnR5lcKRgZcsvZnDPXx7qpU1CvbSNUl8dlMOlWkk4jhevNzOm957zoIOHJG1JniclsHxBRb0JHiKZtmRxfqwjDxXL9yEZeeiZMO2gfMK0AYHA945sWpyhLrng4ai/6k7OLTjnIIFThY3pRP0s9wWO4IV7FG/cl6YtZVdv6mHZCUR7j/WdMJ18HrlfgoVk9GXBQ7b9Om3pI2bkQdexoHx8/876XHfCtOX6JSTBA9KWfHY1Rrm1jZH60p9jyV+HE53DjDIotd0L4mv84UV149lNvWoSnPO+VhRKr59j5CFPt156wjmCh2Zy/SznzoC9Vq4+AcGs5ZwWloOjJbXv2zrq4n3MKAPKwwaIaedDpv75Mr2gfmX1+fP0C/XSNiK0D4DgwaQtyQPGlqoXjXppC4F6aSugcLWlnMMwCcED9OuLVlvKsKLBw5rMnAcXutc8mTBd7PjiPt51WVNt85ySFI8u2794ST10U/iGWv2QOIE7v9ryjOXA51tvuqRuPa9ZuOeBV9pS6pibkQfs/zGK+5uQBA89qy05EMGD5XwJJngo2KG8i76Wmy2p277V1A5u/jw2Ggs1dfM5DZ2K43o2jSONJVyX1N6/3FHn7dhWD9zY0PVGOf3uZJ0nHXnwuH4IHk75hNlY0KeHU0/gj7+D2PG5aNfzJAULE72RFnjzN5p6f5F0lKHkXqX4HueCemkLgXppC2Gl9QHokYdZTlvy+QEsgnppC2FW9DJtyc3yPg/yPHl8y3fRrx46TB4+4LfLJw/DMdjzzVi2VE6Y7rd8Hz3SQ7ACC3abLbvmvsED6o9NvXZ/o9vpy5/3a/+14HQqi+rvQ1aPXt6L92nJ3mwLuB47v2ZJPfVw8TwUr+Chix550GlL7pSgPFXbD7Drt5gw7UAHD5gwDW1JHfH9O2dLv+ABzwI2qsM9yJ/HBZzib+/dVls/L5nwnJzL7IWAkYCjP7CorvtqU98DHTD0M8pQQjrykN8nwwKu39F/5zfhPAG/fdh1HJO5cS3xe5FcM5wPKXb7/GpHXbhHS0/0juqmrfnfCkLIeIHgYW4Sg4eqLx7oq5yDeup99XrkIb/DdP7lnEkNQO+bToEoeEGHlG+jil47D7V5dbjeMK23/nnQXqRbPP7jqCdXuUr5Pnpcc0we3v4FllGCXF1xzY/7+wXtsBQ5vvgMQcmJH+9OAs1iaT/KxT2//xqZ0lJWfx8SPRy0x+6uq91/viSoyaTvfPmvjRNY2N588OC4fsk57XMe3A512v6COpTp69i9Hfs8PN8V4HT3uOjWMbvaUlH52CTu0oNyKTeW9qNMPAs/OKqpNfnzuOqvn6X4+4DNzL7ysQW9vO6OLzeTgjGXAc8u6oWARD+XtnMkf1vqX1Z+oq83aurJByO1ozNAXL5+uBaoX21DSPnmuzj3UKS+tXNb7f/2jt5heref7+hRoBvOMClYCBjsKVjF9dd/Cw311FM/bL3e52EunfPgzgH1g3ppC4F6aQthPPV4KT58W0PmjgsnzKQpfPGPu0toBve+2cv3x1+Pun3lY8jb7q1/3vlInM7bzm8KB1riX76dZT3ywX94WrN4ZKR7/ZPRnnSTrwL0eU8155XOlmw/rs9hf7Ko54jAUcufr5fw9mvHsh7ppSt9Rh0A6n7F4S3Ljsi95YvgoYsIxlZZRh4KX0geeOrTOQ96k7Pi9ifBg57rUHJ+XJurjm9Z7rFsP4KHsz7b7n2+S86fHGNWDIpiB9pc63VPG5teucnnHC48tUh/Qrn7/Ep3OV5Le7Pt3vmVS+qJn3gsQZwrH8ejTeufNuU9M2fuAb5z1t85z/o7oV7aQqBe2kJYaf0IwJyHuYmb80DIBJE4YbvnN+8SwYN5QaNn7on7e3vqxw04ON/cvjytA8C5Onfb4sm5gwaOXLrUpqXHOPv/cLov2hsb85U59/M6zQfpPrvnJ4sXAGf9659eUA3kitscpT7BuXAfkP6C9A+bo5ttJ/5FnXd93ZLubS5zANPgIQl6HdcPpMEDdCN88SF4eKAsbSm9z/7BA8579xUNPTckf578/yd7ZiCVr9+241rDkR/k8+EDykO6HlYb6+0IkO0EW262pG75ZmY37UD0iAvaievU57Uik0R4hwj1k6FfnvOAiW6WAwZJOkzSJ5OunxWqXqdp1NdiMHEw2xPucnbQA3jtScihlufxwVZ+CD561A11dDocGdDmA9/R0b3vPs6RT/lFwDnBhEvTm7oknb4u2et/89mt2BE3wUNZ+VgW87zVbZm65ADloB5I20DwURYUlpWvj8HynHHg8J2D2mqr7ARuG5n2o85nbt4uniid/JsJHmyBUrZMd9rScMmOPDivQS548HFc0Xak2qx+hfu8iR3/IjXv/mvrXsu1jhsIBLBsr7MjIPf8nPE/C167aRNCppcBTph2Ryh+lOhLfuypn3R9CROuh5Nz+n9jGc3c2u4W4OQc9Zd+q8IEU7WToKvXPb7XN9T2L7TkSlvatiVW9/lmw6xHnz9nCB71R2Bz3deaaotN3Y5fUkd8vsNLltTj95glIfPnEtRM2x+9s65Wv9Jxflv7V5kA4uTYSYNDrgPDPu8tRlU2rI/Ut+IAZitbGy3lJ3VAXv2DnitAudKWbGVg5CN0wvQgQDvuz6+25Gi/CR4WvL5XaPuGdfPqC+/Jp+bZy8D3+txtF/rukV9JsD/Idw/LzO9wXD+AZ3i3NyzpwKps5IoQMp3g99Ps8zCQ4IEQ4gJOxdXHZ3rqC17QYOvnLqk7Luw/PaCUEuepDP3j8Wyk9v/t3tEU0ba011epw9+b6ZmuWL5Lb5y+SB36B93J3K7rnOmNPuxPOqq2sdyhzIIJnhft6Rh9sLQfwLnFdTjgHR1167ca2vnSK8v4OLORCRowYoDJtIf/KdrnSFVylI+6nrMVRh38AriQ4GGlRh6S1ZacwUPm7+W0Jb864rt39haWfTPyZawyTvUur13SE9d9ArNxAvV98Pq6XoJYtM3y/cH1uGivdt8jo/2gvyNJupPlc0LI6MDoqhl54CZxhAwV9Gpjw6PVSQqI5aWcfVnD8UWa0/qnh9/Dp/Otm8aJ0GV5vqDhhKK3tce5crUrZov4s+8f1VItrErjWUYo6EW94osluyyjjmnwoNSFu4c7QkgbevqJSB34O7ngqaD9CTgeufTH/N2iuvmcplq3xlx7zAlJJssC/K1ti2ZH77svb6hT/3NBbf9zxhHOn7eofJS511s6au4hzxEWtDENHnKjG5nrl7CScx6S4KGo/aaOKih4wPVHfr8Y3bG0H+B7cNqnFnSdfM4/LuiAO36+Dv59xyhLDvx+rX7Vknro5nph+lsotnSv5LuADgHsMI17pydZW66vTR8C9dIWAvXSFsIk6XXn4eDSlgghRWC1nWM+2F3mM3E+bI5I6tgiP314zojuzV7ATsyRevCmhl7yEnsjwJnQq71YNFlwzD3fa6ht8vnm2TZl2pY4Hfdd3Rio05EAJ/u+qxpqp1dllizNXttcvVDn7V64pDff6qe3GOXdcVFTbZPfndfR/vy9Tp4DzM04+d8W1Lf3a6mbzmyqOy5uqLsubejNsi7/Qkt97b/a6oA4SNnmueaZ0GVlz58tI2/rthMTXbFnQEjaWFSvyaVa8+3q2vodeQh5YdnIBg89Iw/5eq7qjjx8BGlLxgktKx/fA2xyhj0K0lGs/HkzNn2dn6PUNSfZVrKyU1S+D4PSY0nYC3Zt6fQr8XxZriuu5ZF/0dGbRUaewagNV/2T36Z7r2zogOyQd3XU/r/VUUd/aFF978iW2rh+vBeUIGSaWQ4eikYePHKMCxmi3vXDQz3146iHs3l97MCVLYnY44xsqtQl+7dyGyeZ+SWh5WeB448X8GUHtrRztG3sRG/3fLPSEzan+vGVRQ5+t3ztgEXqiD/L9VjaHI8ucML2/fWODlLc589grb9sP64P5i18/tcyq8ZYyjeYFZhw3DF/Z/Z3cI6ElJSPNly4R9u7/dnyt+7+PwIdBBL6HLEd6SMAoxOwZTfXsukFufJxDqw8pQOk7j3LXz8bZWlL2fJ18PDBsOChrPwyoNdzb2zBQ/bfTB2zqy35lI8gABPdewJ+S/sTcC8RvN55WaM0gPApv4hB6nEdH7mjrnbMjjI5n18DRvdO/584GJvvb3TUVX88d0jR+84hLbX9S5eff2C+J0od9TeLas0j1QIIV/m+UC9tIVAvbSGsqD7+DdX7PIiRB+sLMwDqpS0E6qUthDHV46W4bk2kDnx7+c7MCUmv8fm7tBUmo5bvleAuH2AlIvQy/vSWujryLzAKYvaW0OV0wcsaIwTYwKpsVRU47T88pWQ/hRw4Fr2ID95Q12lGwtksqL8AaQzxOZBCccBvFyw3aQGTuG85F73x/ZefpHyd+PGFoGtQRPZe5D8LBc7W0X+LfSbmC3eT7qHb/sLgIUcaPECXP18oAdc/GXkoXG0pU0evtKVM+bi3j99bV7u+tmQDvgx4DnZ5XRxAXNrQG8cVlmXD0v5kfky7g39rxc66Re8DOhRO/aRjjkeX7DXG33DkT/9/C6oW/zalo3d9lg90O+PzXLp/W//uua45fqOwtwnKFde3QvnUU0+9xZYDv/Fp2lJUtlrO0KlaPvXSFgL10haCnx75ulcdZ1viNDMakevxw0savXxIeULvIHo0dY9b5qVZ9v2FIwgdgpdL9mupnV+bSYGxlI/P9n7rknoKvXsFjgpe3EhdOOyPMqMPrh7LjB1Ox65v6OjAQzsMCFLyTkABKBfODq4D0nF2e1O35967/CU9YoKJ0tne+FBQ7/uuqauDfi+znKlH+RL3/S+1O/RwAg9596Ja+5hPD61svwke5nXwoEfLCsrH/UzSlkLuY1WswYOjnknwgGdHOJwFNOPvjZjbU8QmJoBAr/n3jmilG6Tlz+sD7huesZ/dV1fXntxUl3+xpa45qZmO3OlRM4+XvQ+o4wM3NPScGmsgZrmuJoBYUsf9w6Kuo+4MSEdI/cH9wG/j009G6ow4GMF3SdQhU74OrjdT6vaLsIKbPB8hZHikwYPZJE6+PAghgyWZnIhdh3sCiOyL2fKSBjh+59d01Hk7tfUmcnCckkAC59XAMer+nTjlcH7WPRWpq7/SVAd195oQPXqW8uEQfte6E3EveHnffmFDbV00adViSxyE4z68oO79fsMEOEu97dGOaNfZM738NdPmuG33/qChncHEWSsqK+94bPuCJb0JmNdIjgNc2+u/1lQ7vcZj0rTN5vrcdqzNVvA5HN3D3mscun7mc4CwkYcxDx42Md+fEz5iRh5C6ojn8YmfRGqPN9n3uxBldv/GsQi6jv3QgvrJ9Y10EnzyPc2XA5LnvNEyzzmWRP3mDm09koHzJd/dnV+zpM6N7f3tRO8gMvU787OWVcRs1zQDru0eb+moHxzd0p0JGN3U32NHOwE+w/c4+X3BHJ8DfteMhornzVI+NoC8YFe/neEJIYOFE6ZXhKqB2vTq/VIeJluvJwL+wPTw9ay8lP/XAo6Ho7bzq5fUVz6+oL5/VFM9fFtdPfOkWY2ktsEEJ+vWzqsnH4jUbec31dlbttXnf904PsbJzc25yJfb/RepAaf8u23t+lz7I+PEocdQ5IZb6S0fjgfmW3z5bxbV977U0qMr65+Z16MCcEAA/obt0R/V1ZXHtvQozLbPN06ycBpLyodzcuZnF6yBg8/9Q/vh8Nz0jdhp/blM4OJZfr79Quet7wideT7MxGA4/vYe7+T+ycnT2fbr4OHxSK0WwYOsv05b8p4wXe37k9WnwUN20nrm+mXrba7L8pwHeV53+bjfVx3f6ga8sv3Wf7vlo9wd4ufkxI8t6NWbcF9wbbGSFkDAjBXI2sq0Z/0zUdymujon/t7u/gsmFa83aFH6/+E8n/IfCz0jKa769+K+/nDmMdqIuUPuNDx7+5PA5sB3dtSl+5vv8cYNkUm3UuYaAvyNDg08K0/8JP4+H9NUR/z5otpqU5NCmT23CCIy5eO6nPZJs3N7cj+rtp96aZNQL23ToffD6MsnTBOSx/Hi9WbW9fOmt/6S/RwTbfNY7HhJwwnGv1gxaLfXd9T+b+/oVJxD391R+/xqR+34UjPpFg62s8e0oBwEDyf+s2U3WUv74XQgPQY7Sfe0yXJeG4mThXpu/6Iltd/bOupL719UJ/3Lgp5PgPkZmCexw4uTthc4cXkydpSBFCOkRxSlYxUBBw8B266v8xjxKLLncR3nsmdIrt92LzJr8G/c0E1ts9Q/xXIfs5iRh3m9w3SREwd08PAhz+Ah+RxOfP4zHzJ6a/CQp3v9loOHbh0Dyk9GwvAs+uzEnLejbnqDyPj7uOcvdtRxcaB19hZtddnBLfXdI1rqisNa6rzVLXVyHKzjOU8C46LvLc6J5Y9vOL1ocQML6fWzfDZvfpswFwh7zehr6mpXnkwQgQ4OfFe/+EeL6mv/tRA/ky09iom2XrJ/W53xP2115Ps7aufuymjiN8Ny3jx65GEX/31LUkraXwr11fT58/QL9dI2IvD7qTeJM2lL8oCxpepFo17aQqBe2kKIjCMCR+vEjy+6nZEAklEFvLSTHkC8+J0OVYLjxQzg7GClGfQU5usv2jSPEZWaXqVpp1cWpHfksZSvnaJVxqGAA5WssIJ2ifZY9C6gR7rHT67zXOnJgnY64/t21F93084CyrdSVb9q+VodGjtqSMXCaMMgUlm80pZSx9ykLdVWKG2pZ7UlByJ4sJyvCARjWOGnZ2J+H/cvuV9wfnHdku8ufgfwrFufcwc4HkFbmuJnqXc/YFTu/J3RuZEL0PMUtB/B/XK78Ldp63Lb/duZR+s2w6T05nLwULX91EtbCNRLWwgrrQ8AIxwznbbkN8TjhnppC2HW9ej5RroR0nVEjvEKk4wAPHhjQ+/im687sLUfvZY3ntFU27yguNd01KAu27xwSd16XiNNw7LVvwzkct/49ab6300sedkjBGUnwSJS0pDuhXSXkKCorP1ewUMXPfKg05Ysq984KCvfh2SH6ZDVlnS5nnXMg1G4B29oqD1+sSitZ3Tgucbmf0jp873uPphgJNIjIdY5CH0wiHMkfCau04n/bFb36vdeEkL6B8HD3CQGD1VfPNBXOQf11FfV41/k6z7zhAkg0ENX+IJNevlsvX02Wx5PPeoAR+uC3dvWeQFJ/V3tb7Zr6pqvNNXWL3Ck9ZSU76RPPSatoi5Xn9BdlrXrbLjq7yIZdUAKleiRLShfHGM71mbLgHuS3Bc4c3CWv/AnHXXlcSaHHrnkcPhqlnrbKLp/6TH54KGg/svBQ/c6peexp5T0lF/o/BXr662a3txwu+e7ApzuvJBuHZPVllBmv+XjWmOEZ/VrLPNdBMvlF10/q81Dr4OHX+qYpXit7ZD1Bz7XHyMt656eV8d9xKRqpdfXVlebTVNcf/GZQOrR0fLlv17Uv5vlyw/3337qqafergd6n4e5dM6Dfw6oHeqlLQTqpS2EydSj9/Tp+EV4wj8Zp9TuBFkQL9ySFIM8Fj0cLGwY9+3Pt2NHNHQll277o+5KRKc11Y6vLMhntpQvjinCQ4+ysWkXRgsQ1Mg6W+rvQKfIXJvJr/cof9nhXwaOWJKiYlDiX3yO3tXEaUOZu7ymo476q8X43rT0Lt2Y14BRlOV7VFz/cnr1InjoYnMixchD4QvJA099drUl2/XPkgQPRROmfcvHCMSP43uw51uQwlRcbh7b9Qshq8ezgh548T0tqX8pXT06N7D4win/uaCvX89v0wDqn/+sCOjQXuxbgt/LwlXEBtT+vqFe2kKgXtrGDMx5mJu4OQ+ETBno5cOOz+fv1lbbPM+zxz7//y489MmL+eB3L6q7LzcTMKN6mbNdDDazuuf7DbXfb3a0UyyConx9XHjU3wbKxO7ZqINeDjLy75m3gVVxLv9CS1+nsvKx+zB6hTFpHfNazvrftjrrs229DGbyb8JZn13o/r2gPztr87a6YOeWunT/prr80Ja6/vSmDhbWPBbpOQVoi57XgBfMEF8ycEjX2kYeEjL/D8fymA+YOQ/2HvDhkAZ0RWlL3Xqijid8FCMPHsGDBxjFwopCX3gPnu+CoL/P51cclzse5W2xqdJLnIakq4ViFhaI9BwITKIuH23J4ai/+H8Xm5jvEkYQv/rfC3qvmtLFAMiYMNgOjXCol7YQ3PrlOQ8D2mimiHSYpE8mXT8rVL1Os6qHswan8PYLm2r/3zEOd9G8AaezksfxkobeBA2xk/typCm11LNPRaqFwMFSP1+yWjg1cw/V1cn/tqC23KQ3KBpE/W02PeE01pz6qQU199PwydGutsNh+eqnljcLs5WfgPu255uX9ARyjCzpNf7b5t8ycJxZ0tIs34nRBTjJ2B3cxzF31d+XRI/nEXNydn29fQJ8tv24JljWF46mTx0HBcp7/Md1PT/HeT+6zw/qeNbn3Kl4/YD79cxcpINDrGxmHWULeH6tWPRm4rFS52zd1qk7w77m+rcpfgZvOrup9nprpzzFsqT+wEef/D7t/LoldeUxLV2PfldJI4QMjgFOmHZHKH6U6Et/HKmXtknSlzDp+gDg7CJd5MK92np1ILyohfPmeCGXftYleSmjJ/HEf1lQD97c0Guvi/SHPH10MujNohpmE6iDft+sVFO4ooxH/W3HYDUXcODvdXRZcDJED2Uf9QdwzrCHBpbC1QGQpfwEHZTFjuQPT2uqhTgAwLOTde56Rgyyz5Xr7yx91r8vfbeeZkJ/8RwP3NPvH1W+oeCgSe7LF/6wfAUk3LcfntoceB31DtKxY337BQ110LuWn+98+UXXz/ez5Hu76xuX1BWHt3X5pd/ZQRGZ0RbsH3P6pxf0aI9IZSqpv9cxm5hz4txYivb4f1xQj9xeNyl6ru8FIWRk4Huol2pF8FC1t4oQMjjqjZp+UT9+b11duEdL7fVLxiFJRiPECzv38rX9nbyQ9UjDy5bUyZ9YUHd/t2F2og7tjQ19iXcdD2xkd80JJoiAg21tj6P+eaCBNtnv4sD4nFed0FLr1mJjqlqxo1H0mQX9Y7l+Xh38ru4qOyX1QqoP9r1w9pIGlr9SIJC95dyG2uo55hrb7gt6wPf5lY5a80hdj47kzzFssFTntSc3l58jy/OD5+yg31tUz64JncfjB96f+vmOnz3sZ7Dvb5nRGuccpoLnJ19/6JPvPjaH/Ma2bfX4PXVdXuEzPiQQkCNYuvd7DXXCRxfUNvk9KSzXX+CwJ8HRVs9Z0hvHYb+JZNQjXw9CyMqA3ztuEkfIGIMXNXrcMBJxw9eb2uHf400dtc1zk0DATHKGc2EDL2Ksqa4d2pfihdxRlx3YUo/cWdcvZAQNo3RA4AjA6cHSkrdd0NQbSO3xCx219WamrgD1ThzBPEmb9KhJrNnjjR311fgcN8dOxoZn8hOIBwfOuWF9pA55d0nw0HX2dnjRknri/rrZ/dZyvkkCaVeXHdDSAUTS02wCN0zwVmr3+Hm89/uN8AB0UERmZOvcbdpdhz2pY7IT85LaKw5usL+HfcftwYHnBM8g0v+u+1pLHfuhBb27NIKI5NnOPs/552f52iYapSfLI2i9ZP+WevxesxyvGFFbATBpHMHi/dfW9YZ3e73ZdAjgept2mt+dbDuTv/Pf6WR0ddfXxd/n/9tWd32noWq17tLDnr9PVTtAqZe2EKiXthAmSa870waXtkQIGRbaKekYp2Ht45G66/KGThNBDvcJH1lUh/5BRx30zo46uAv+xso8p31yQV20Z0vddFZT/ey+uplwixx6pBJZyhkV+PFJcvnRnjsvaapLY+fo1P9YUIf+YUftFgcFq1/ZUTu+ZEk7Xzu/qqN2j22H/uGizq2/ZN+WuuOipu7d10FQZgnWYaDTY+Jr9+W/zW0Oh3/zgcQmxhG6NHa4F5S5d9Db0I5vMpfB9nmJNl/PQaMnmEfmebntgoY6+gOLauf4vuzwUrO/wBmfbqvH7g6fV5JQ9RlM9Poa1k1q3JHxc7/6FXEd42cH+fnf2HZB/ezHdWtwM6jyU7ppYUkQgfv3yI/q6qrjmjrIPfj347q9vKO2e8Hyc4IgDMCR3va5S7ree8f1xuprl8aB/o+vaqiN68wzng8aRPmBVNbj2cA8HaRaPhGp2y9oqvN3aakj/3JR73q//YuX1DbPWQ6IEmDb/oXme/3FP4nv0dZt3Tky99MoPZ/P8121/oSQcPC9Lx95CMmRtTFEvdcPB/XSRv0K6838kH71ehg/fsG2l0zvH16yG9ZJahtNPnZrqaYdGe146BdytfIHqocjjNGIBbOaEeqL9KBn15jRlkdjxwvO19NPGBs+ixpmXwM9sduVFuRbfhmJvusknfmZtu4NhuOHFZXyPcfZHtYdXmYmeWKUBfdiI9gQdTG5+sZWQs8xRo/zwaHHNQMIoKwjLoNq/3zXoZs3wd5j99TVujXdidx99oJ7lV+A0EfdIDKu81MPRzqo0XVs16x1FPpAvPQRrk/NBP4t8/zCQb7vmoa67VtNPS8GXHtqUwc+2D9C13vtvH5Bm5W1uil4OWfaq/wCBq1PRhVRZzyfzzw5r356a13ddVlD3XBGU883QVuv+2pT/ejbTfXgjXW19lHzXdDff8y5wvfZI2iwlR8K9dIWAvXSFsJE6yN8b20jD9YXTgDUS1sI1EtbCLOmhwOekOupFsf6EFp+nir6TN2REgGnT0+4zrSr1LmoUn6BHk76lcc2lyfmlpCkXu3/mx31pT9bVEf8aQgdi22ZI9+/qE782IL6+qfb6orDWur+2BnFEpZw3iqPKjnaD8wzZpxx5/NVoPeioh5tT54da0BVRsXyi/T6+jVqZrSs6zAn6HvXNgG013PuoqB8Lyrqcf2T36JkpTF0Xiy3tabbrp/TZNQte46K5VNvsYVAvbSFMAN6fMfTtKWobLWcoVO1fOqlLQTqpS2EldVX//5SL229wKnDSAhSqWS+emY1olwaU5KH34uZP+D+/3K7mc9icsaxNwj2tMDSnQ/e1DDLvgZNMi1vPyGEEJIGD2aTOL48CCHECUZDmvPqpH8xu+32BA9lK8zYbK7PbcfabBmSiacIKJBnjj0uHr3TLG/Zdw82IYQQYoETpleEqoHa9Or90i2olzbqR6FHjz5WFsLylOnoQ3bytIej3zNKYdN56ztWHeqFeRm7vqGjrjmxZeZD9AQQSfvlDuJl7e/VS6iXNgn10kY99dSvtN4Poy+fME1Inqo9mdRLWwizrl9JuqMPZ/wPdpp2BAF5XPY8ruNc9jy54/Sci9h2wW5tnXsu8v/7vQ+JLurzJTQ1estnPlBP/Szr8+fpF+qlbUSgM0pvEmfSluQBY0vVi0a9tIVAvbSFQL20TRDYwG/NY5E64HfMbsJFqy6V4hsYuCjRJ0tkXrCrCSDybSFkbKj6u0C9tIVAvbSFMOn6ADDCMdNpS35DPG6ol7YQqJe2EKiXthCq6BsLNT0xebc3mQnMead9nNDzITZReqWoZmc5ValK+8dBTwghZGVA8DA3icFD1RcP9FXOQT311Eu7L5OuB804gMCuxfu/vaMnKssVmHIUzWmw2fJU0GMy9a6vX1KP3mV2Fk+uQb5NvvRcv8IeLzmvgnrqqaee+snVA73Pw1w656HPHNAU6qUtBOqlLQTqpS0E6qWtGKxXP/dQpE795ILaalOzrGppEJEgHP7cHIoyAvUIcE7+9wW95Gy+HYbw9jspfCF5QL20hUC9tIVAvbSFQL20hVBVPwIw52Fu4uY8kBWjZrGFQL20hUC9tIUwDL3eyK4+r+64uKm+/HeLatsXGEcdgUSyURz+9SE5tkyDz0ODFBy/3YuW1AM3NAL3gCCEkH6o2iFBvbSFMDz98pwHj13lqpIOk/TJpOtnharXiXppC4F6aQuhL31kRiHw9/3X1tUl+7bU8R9ZVPv9Zkft8ppldu2StRWRPzbRr37Fktr2eSZIQTAhAgZL8ABw/LnbtvVuxqINXfpqPyGEkJlhgBOm3RGKHyX60mEc6qVtkvQlTLp+XKjaSUC9tGWJzF4QrSXz94Zn59Wza+bVujXm32efSv6OzP9rG/5O6Nq6JLr0X+jXzqsnH6zr/SYQpOz9y1j1yb3DdRbMfTjwHR21Yd18bu8HQgghpJx0qVYED+xtImSGqOo4Ui9tOfADm1LP/JuQtRf9nddHWC52XjXjIKUdBylzP43UsR9eXF71KRs85AIJnbr0wiX1yB11nW6VrzMhhBBSBOIFbhJHCCETDNKlnnwwUnu9pVOewoQAIg4orjutWZi6RMioqdqBSb20hUC9tIUwS3o98jC4tCVCCCErAUYgztuxrTbPjzjkRyE2wURupS7+fEtr8ucJJeSFY4N6aQuBemkjhAyX5eChaOShLMe3jCHqvX44qJc26ldYb+aHUJ+3W6Be2iz6Vmde3XB6U22RDxYsIw+bx8HDudu21ILKn9eCtXyDV/0LoF7aQqBe2kKgXtpCoF7aQphofRSZfR7EyEPBC8ML6qUtBOqlLQTqpS0E6qUthBXQI3XpjouaautNy5dw1cHDNggeHBsA9VE+9dRTTz31s6FH4JGmLUVlq+UMnarlUy9tIVAvbSGsrL7695d6aQthZfVIQbr80JZeitU14pDYETx8c/v8yEO18gkhhMwGafBgNonjy4MQQiaNqBH/kK+P1KF/2NFLsVqDh4wNx1xxWEunOuXPRQghhJTBCdMrQtVAbXr1fnl41Esb9bOo10u2Ls6r83dtx0FBd68HW/DQBSlN2zxHqbsua6hGW6YthZafh3ppk1AvbdRTT/1K6/0w+vIJ04TkiSy2EKiXthBmXT/rxNev0ZxXrThoWPt4pM7Zuq2XX3XOdcgEE1jKdbc3LKk1j0Z6xEKc24fk/kV9voSmRm/5zAfqqZ9lff48/UK9tI0IvdoSNokzaUvygLGl6kWjXtpCoF7aQqBe2kg5kdkYrrFQUw/fUVfnbt9We77VpCo5A4ccn12l1IkfX9AbxPn1UhEyIqr+LlAvbSFQL20hTLo+ALw7ZjptqerLk3ppC4F6aQuBemkLYVL02FUaowz4+67vNNQpn1hQO77cBAJbJKlKnmy5qVK3fGMwG8T51t9FVT0hhJCVAcHD3CQGD1VfPNBXOQf11FMv7b5QX67H6AAmNa9bG6lrT26qL71/UW2zmdnkTYw0FMxzSPhczGF/0lG1jfN62Lms/CJ66l/Y4yXnVVBPPfXUUz+5eqD3eZhL5zz0mQOaQr20hUC9tIVAvbSFQL20hTAAffxD3miZSdBPPhCpS/Zvqf3f3tFzFeD869EDZ8DgHoXQE6Wft6TuvKSh94SQZXfLF7Y+KXwheUC9tIVAvbSFQL20hUC9tIVQVT8CMOdhbuLmPJAVo2axhUC9tIVAvbSFMK56jATAqcdow0M31dWZmy+oXV/f0QGDXn61Gyz0jDiI4MEONNj/4Zvbt1WjLcsmhJDhULVDgnppC2F4+uU5Dx67ylUlHSbpk0nXzwpVrxP10hYC9dIWwqj1cOhrGyN12wUNddyHF9T2LzLOPkYbioIEkbrkAOc65u8X1cb183r+RL78PKH1J4QQMlsMcMK0O0Lxo0RfOoxDvbRNkr6ESdePC1U7CaiXthBy+uZCTd19RUMd9t7FOEhQ2tH3CgocAUX2M5xn85ijP7ionn0q0ntC5MsnhJD/v717aZFk28swvkEQxJkIgnMRJ46cO/IrCH4DFcELKCiODnrEgSKIM3UkchBHgniUAw4ceB3oceBAEUTh1O7L7t5VHb27+hb2iuiMrM43MnK98c/MlSvX0/DzbN+qp6rppjs7MiIyAcf0Uq3p4IFnm4CGRA+G6HVzdONlSt/+66/6X/m+8ezA7j/+5b/3fc6MdNYi+eOfetM/SwcOr2Z+DgAAmNLxAm8SBwBnli4fevKdrv/6j4z3NeQeFEyfs+fz0gFDOhD5tR9433/rd++HMw3DGYeZnwNwSaJPYNLr5qDXzdFSP5x5ON5lSwCAHOl9G/7pG6+29zXsOzjYPQuxu3/835/7bLzk6Ve//33/jZ990//Pt18Ol0S9yLjHIcJ5wJlDr5uDXjcAp7U9eFg68xC9RvaEfdZfHPS60Rfux/tD6Hf3GVfav+5f9N/8zfv+p3cPHHYPIna3QT99bPNqTF/74Xf9X/3Gff9///FyuAH75f3y90//u/7nP8rqF9Dr5qDXzUGvm4NeN0fVfdeN7/MgZx4WHjCy0OvmoNfNQa+bg143R0b/ur/rv/n1+/5nFt6bYZ/NS6/+/He973/vx9/1f/eHr/qnN93whnLp3obQg0KS8fNfRK+bg143B71uDnrdHA306TFmumypO/RqOScX/f70ujnodXOU7eN/ful1c3h9+of+P/zJq+Fyo+lswuYAYc8Zh839DL/8ve/7P/rJN/2//eVX/d1dN1wClfMSrMu8nz8AoE3TwcP4JnE8eADAObz88I/9x//b9V/7oXfjm8Dt3tvw4KAhHWCky5N+/Qff9X/+S2/6//6Xl8MZhuGdoqOv+gQAgIkbpouIHqhdb593yQW9bvQ19enz0g3N//xnr/pf/J7x3oWH7++Q/jtt6WzDb/3ou/5bv3Pff+e/uv7V2/FN5cavcydfO/f7H6dX9Lopet3o6elL93nGfnyfh6UbpoFd0Wc76XVztN5fiw+/Dl/dv+j/9S++6n/7x971v/Dd42VJ6aAhvcv07//E2/4f//RV/+zRx/sZjv2Sq2t/HzZdt/JB6Gr6mY/loKdvud/9OmvR63Ym05vEjZct6SdcrOgvGr1uDnrdHPS6NezV6xf98ydd/+9/81X/t39wP9wA/Z9//9XwsePczwBUIPr3Ar1uDnrdHLX3hnSGo+nLlvJO8exHr5uDXjcHvW6OS+q7l3fDQcTr9x8OGN6PZyQOPRgc8/uvUboHAJSRDh5uajx4iD7wpD7yNejp6XXPRX9F/eJBjt5XQU9PT09fb58M7/NwM93zsPIa0Am9bg563Rz0ujnodXO03j+w+ICUgV43B71uDnrdHPS6OaL9GaR7Hm6qu+cBxdzNbA563Rz0ujnodQOA04g+IUGvm+N0/faeh4x3lYuaTpOsVHvfiuivE71uDnrdHPS6AQAw6I56w/T+I5Q8B/qDp3HodaupP6D2/lJEnySg181Rew8AaNr0Uq3p4IFnm4CGRA+G6HVzRHsAAApIxwu8SRwAACgq+gQmvW4Oet0cLfXDmYfjXbYEAGiJ84Azh143B71uAE5re/CwdOYheo3sCfusvzjodaMv3I/3h9Dv7jPodSvej7L6BfS6Oeh1c9Dr5qDXzVF133Xj+zzImYeFB4ws9Lo56HVz0OvmoNfNQa+bg143B71uDnrdHPS6OSro04HHdNlSd+jVck4u+v3pdXPQ6+Yo28f//NLr5mi9BwC0Yjh4GN8kjgcPAAAAAPtxw3QR0QO16+3zrsOj142+vf5OPnbeXtHrpuh1o6enL93n6Ybvw0u1whd9jXp63Ryt9ziOtb8Pm65b+SB0Nf3Mx3LQ07fc736dteh1O5PpTeLGy5b0Ey5W9BeteK/P+FnodXPQ6+Yo3c88Y+6h1w0oLPy4OrM56HVz0OvmKN0bhjMPLV+2lHeKZz963Rz0ujnodXPQ6+Yo3QMAykgHDzc1HjxEH3hSH/ka9PT0uueiv6J+8Rmv+bMc9PT09PR19snwPg830z0PK68BndDr5qDXzUGvm4NeN0fr/QOLD0gZ6HVz0OvmoNfNQa+bI9qfQbrn4aa6ex5QzN3M5qDXzUGvm4NeNwA4jegTEvS6OU7Xb9/nIeNd5aKm0yQr1d63IvrrRK+bg143B71uAAAMuqPeML3/CCXPgf7gaRx63WrqD6i9vxTRJwnodXPU3gMAmja9VGs6eODZJqAh0YMhet0c0R4AgALS8QJvEgcAAIqKPoFJr5uDXjdHS/1w5uF4ly0BAFriPODModfNQa8bgNPaHjwsnXmIXiN7wj7rLw563egL9+P9IfS7+wx63Yr3o6x+Ab1uDnrdHPS6Oeh1c1Tdd934Pg9y5mHhASMLvW4Oet0c9Lo56HVz0OvmoNfNQa+bg143B71ujgr6dOAxXbbUHXq1nJOLfn963Rz0ujnK9vE/v/S6OVrvAQCt2L7PAw8eAAAAABZww3QR0QO16+3zrsOj142+vf5OPnbeXtHrpuh1o6enL93n6Ybvw0u1whd9jXp63Ryt9ziOtb8Pm65b+SB0Nf3Mx3LQ07fc736dteh1O5PpTeLGy5b0Ey5Wp8+YWeh1c9Dr5qDXzUKvmyPaAycQ/ccQvW4Oet0ctfeG4cxDy5ct5Z3i2Y9eNwe9bg563Rz0ujlK9wCAMtLBw02NBw/RB57UR74GPT297rnor6hffMZr/iwHPT09PX2dfTK8z8PNdM/DymtAJ/S6Oeh1c9Dr5qDXzdF6/8DiA1IGet0c9Lo56HVz0OvmiPZnkO55uKnungcAAABcsOgTEvS6OU7Xb9/nIeNd5aKm0yQr1d63IvrrRK+bg143B71uAAAMuqPeML3/CCXPgf7gaRx63WrqD6i9vxTRJwnodXPU3gMAmja9VGs6eODZJqAh0YMhet0c0R4AgALS8QJvEgcAAIqKPoFJr5uDXjdHS/1w5uF4ly0BQFD0shp63RxmLw849PI5S+jL9oJeNwe9bo7SfabtwcPSmYfoT+aEvfzFMYdeN/rC/Xh/CP3uPoNet+L9KKtfQK+bg143B71uDnrdHHX33fg+D3LmYeEBIwu9bg563Rz0ujnodXPQ6+ag181Br5uDXjcHvW6OCvp04DFdttQderWck4t+f3rdHPS6Ocr28T+/9Lo5Wu8BAK3Yvs8DDx4AAAAAFnDDdBHRA7Xr7fOuw6PXjb69/k4+dt5e0eum6HWjp6cv3efphu/DS7XCF32NenrdHK33OI61vw+brlv5IHQ1/czHctDTt9zvfp216HU7k+lN4sbLlvQTLlanz5hZ6HVz0OvmoNfNQq+bI9oDJxD9xxC9bg563Ry196amL1vKO8WzH71uDnrdHPS6Oeh1c5TuAQBlpIOHmxoPHqIPPKmPfA16enrdc9FfUb/4jNf8WQ56enp6+jr7ZHifh5vpnoeV14BO6HVz0OvmoNfNQa+bo/X+gcUHpAz0ujnodXPQ6+ag180R7c8g3fNwU909DwAAALhg0Sck6HVznK7fvs9DxrvKRU2nSVaqvW9F9NeJXjcHvW4Oet0AABh0R71hev8RSp4D/cHTOPS61dQfUHt/KaJPEtDr5qi9BwA0bXqp1nTwwLNNQEOiB0P0ujmiPQAABaTjBd4kDgAAFBV9ApNeNwe9bo6W+uHMw/EuWwKAoOhlNfS6OcxeHnDo5XOW0JftBb1uDnrdHKX7TNuDh6UzD9GfzAl7+YtjDr1u9IX78f4Q+t19Br1uxftRVr+AXjcHvW4Oet0c9Lo56u678X0e5MzDwgNGFnrdHPS6Oeh1c9Dr5qDXzUGvm4NeNwe9bg563RwV9OnAY7psqTv0ajknF/3+9Lo56HVzlO3jf37pdXO03gMAWrF9nwcePAAAAAAs4IbpIqIHatfb512HR68bfXv9nXzsvL2i103R60ZPT1+6z9MN34eXaoUv+hr19Lo5Wu9xHGt/HzZdt/JB6Gr6mY/loKdvud/9OmvR63Ym05vEjZct6SdcrE6fMbPQ6+ag181Br5uFXjdHtAdOIPqPIXrdHPS6OWrvTU1ftpR3imc/et0c9Lo56HVz0OvmKN0DAMpIBw83NR48RB94Uh/5GvT09Lrnor+ifvEZr/mzHPT09PT0dfbJ8D4PN9M9DyuvAZ3Q6+ag181Br5uDXjdH6/0Diw9IGeh1c9Dr5qDXzUGvmyPan0G65+GmunseAAAAcMGiT0jQ6+Y4Xb99n4eMd5WLmk6TrFR734rorxO9bg563Rz0ugEAMOiOesP0/iOUPAf6g6dx6HWrqT+g9v5SRJ8koNfNUXsPAGja9FKt6eCBZ5uAhkQPhuh1c0R7AAAKSMcLvEkcAAAoKvoEJr1uDnrdHC31w5mH4122BABB0ctq6HVzmL084NDL5yyhL9sLet0c9Lo5SveZtgcPS2ceoj+ZE/byF8ccet3oC/fj/SH0u/sMet2K96OsfgG9bg563Rz0ujnodXPU3Xfj+zzIPQ8LDxhZ6HVz0OvmoNfNQa+bg143B71uDnrdHPS6Oeh1c1TQp+OF6bKl7tCr5Zxc9PvT6+ag181Rto//+aXXzdF6DwBoxfZ9HnjwAAAAALCAG6aLiB6oXW+fdx0evW707fV38rHz9opeN0WvGz09fek+Tzd8H16qFb7oa9TT6+ZovcdxrP192HTdygehq+lnPpaDnr7lfvfrrEWv25lMbxI3Xrakn3CxOn3GzEKvm4NeNwe9bhZ63RzRHjiB6D+G6HVz0OvmqL03NX3ZUt4pnv3odXPQ6+ag181Br5ujdA8AKCMdPNzUePAQfeBJfeRr0NPT656L/or6xWe85s9y0NPT09PX2SfD+zzcTPc8rLwGdEKvm4NeNwe9bg563Ryt9w8sPiBloNfNQa+bg143B71ujmh/Bumeh5vq7nkAAADABYs+IUGvm+N0/fZ9HjLeVS5qOk2yUu19K6K/TvS6Oeh1c9DrBgDAoDvqDdP7j1DyHOgPnsah162m/oDa+0sRfZKAXjdH7T0AoGnTS7WmgweebQIaEj0YotfNEe0BACggHS/wJnEAAKCo6BOY9Lo56HVztNQPZx6Od9kSAARFL6uh181h9vKAQy+fs4S+bC/odXPQ6+Yo3WfaHjwsnXmI/mRO2MtfHHPodaMv3I/3h9Dv7jPodSvej7L6BfS6Oeh1c9Dr5qDXzVF3343v8yD3PCw8YGSh181Br5uDXjcHvW4Oet0c9Lo56HVz0OvmoNfNUUGfjhemy5a6Q6+Wc3LR70+vm4NeN0fZPv7nl143R+s9AKAV2/d54MEDAAAAwAJumC4ieqB2vX3edXj0utG319/Jx87bK3rdFL1u9PT0pfs83fB9eKlW+KKvUU+vm6P1Hsex9vdh03UrH4Supp/5WA56+pb73a+zFr1uZzIdPIyXLeknXKxOnzGz0OvmoNfNQa+bhV43R7QHTiD6jyF63Rz0ujlq701NX7aUd4pnP3rdHPS6Oeh1c9Dr5ijdAwDKqPbgIfrAk/rI16Cnp9c9F/0V9YvPeM2f5aCnp6enr7NPxldbmu55WHkN6IReNwe9bg563Rz0ujla7x9YfEDKQK+bg143B71uDnrdHNH+DNI9DzfV3fMAAACACxZ9QoJeN8fp+u37PGS8q1zUdJpkpdr7VkR/neh1c9Dr5qDXDQCAQXfUex72H6HkOdAfPI1Dr1tN/QG195ci+iQBvW6O2nsAQNO6dPAwvM/Dh4MHnm0CGhI9GKLXzRHtAQAoIB0vTAcPux8EAAA4h+gTmPS6Oeh1c7TUD2cejnfZEgAERS+rodfNYfbygEMvn7OEvmwv6HVz0OvmKN1n2h48TC/VOiP6kzlhL39xzKHXjb5wP94fQr+7z6DXrXg/yuoX0OvmoNfNQa+bg143R919199uzjx88oUWHjCy0OvmoNfNQa+bg143B71uDnrdHPS6Oeh1c9Dr5qigT8cL02VL3aFXyzm56Pen181Br5ujbB//80uvm6P1HgDQiu37PPDgAQAAAGABN0wXET1Qu94+7zo8et3o2+vv5GPn7RW9bopeN3p6+tJ9nm74PuNLtS7dMA0AuD5r32ti03UrH4Supp/5WA56+pb73a+zFr1uZzIdPIyXLeknXKxOnzGz0OvmoNfNQa+bhV43R7QHTiD6jyF63Rz0ujlq701ctgQAAAAgS7UHD3nXdu2X+sjXoKen1z0X/RX1i894zZ/loKenp6evs0/GV1ua7nlYeQ3ohF43B71uDnrdHPS6OVrvH1h8QMpAr5uDXjcHvW4Oet0c0f4M0j0PN9Xd8wAAAIALFn1Cgl43x+n67fs8ZLyrXNR0mmSl2vtWRH+d6HVz0OvmoNcNAIBBd9R7HvYfoeQ50B88jUOvW039AbX3lyL6JAG9bo7aewBA07p08DC8z8OHgweebQIaEj0YotfNEe0BACggHS/cbg4edj8IAABwDtEnMOl1c9Dr5mipH848HO+yJQAIil5WQ6+bw+zlAYdePmcJfdle0OvmoNfNUbrPtD14mF6qdUb0J3PCXv7imEOvG33hfrw/hH53n0GvW/F+lNUvoNfNQa+bg143B71ujrr7rr/dnHn45AstPGBkodfNQa+bg143B71uDnrdHPS6Oeh1c9Dr5qDXzVFBn44XpsuWukOvlnNy0e9Pr5uDXjdH2T7+55deN0frPQCgFdv3eeDBAwAAAMACbpguInqgdr193nV49LrRt9ffycfO2yt63RS9bvT09KX7PN3wfcb3eVi6YRoAcH3WvtfEputWPghdTT/zsRz09C33u19nLXrdzmQ6eBgvW9JPuFidPmNmodfNQa+bg143C71ujmgPnED0H0P0ujnodXPU3pu4bAkAAABAlmoPHvKu7dov9ZGvQU9Pr3su+ivqF5/xmj/LQU9PT09fZ5+Mr7Y03fOw8hrQCb1uDnrdHPS6Oeh1c7TeP7D4gJSBXjcHvW4Oet0c9Lo5ov0ZpHsebqq75wEAAAAXLPqEBL1ujtP12/d5yHhXuajpNMlKtfetiP460evmoNfNQa8bAACD7qj3POw/QslzoD94Godet5r6A2rvL0X0SQJ63Ry19wCApnXp4GF4n4cPBw882wQ0JHowRK+bI9oDAFBAOl643Rw87H4QAADgHKJPYNLr5qDXzdFSP5x5ON5lSwAQFL2shl43h9nLAw69fM4S+rK9oNfNQa+bo3SfaXvwML1U64zoT+aEvfzFMYdeN/rC/Xh/CP3uPoNet+L9KKtfQK+bg143B71uDnrdHHX3XX+7OfPwyRdaeMDIQq+bg143B71uDnrdHPS6Oeh1c9Dr5qDXzUGvm6OCPh0vTJctdYdeLefkot+fXjcHvW6Osn38zy+9bo7WewBAK7bv88CDBwAAAIAF3DBdRPRA7Xr7vOvw6JOUYKcAAAWlSURBVHWjb6+/k4+dt1f0uil63ejp6Uv3ebrh+4zv87B0wzQA4Pqsfa+JTdetfBC6mn7mYzno6Vvud7/OWvS6ncl08DBetqSfcLE6fcbMQq+bg143B71uFnrdHNEeOIHoP4bodXPQ6+aovTdx2RIAAACALNUePORd27Vf6iNfg56eXvdc9FfULz7jNX+Wg56enp6+zj4ZX21puudh5TWgE3rdHPS6Oeh1c9Dr5mi9f2DxASkDvW4Oet0c9Lo56HVzRPszSPc83FR3zwMAAAAuWPQJCXrdHKfrt+/zkPGuclHTaZKVau9bEf11otfNQa+bg143AAAG3VHvedh/hJLnQH/wNA69bjX1B9TeX4rokwT0ujlq7wEATevSwcPwPg9HOXgAUI3owRC9bo5oDwBAAens9C0HDwAAoKTo5XL0ujnodXO01A9nHo532RIABEUvq6HXzWH28oBDL5+zhL5sL+h1c9Dr5ijdZ9oePEwv1Toj+pM5YS9/ccyh142+cD/eH0K/u8+g1614P8rqF9Dr5qDXzUGvm4NeN0fdfdffbs48fPKFFh4wstDr5qDXzUGvm4NeNwe9bg563Rz0ujnodXPQ6+aooE/HC9NlS92hV8s5uej3p9fNQa+bo2wf//NLr5uj9R4A0Irt+zzw4AEAAABgATdMFxE9ULvePu86PHrd6Nvr7+Rj5+0VvW6KXjd6evrSfZ5u+D7j+zws3TANALg+a99rYtN1Kx+Erqaf+VgOevqW+92vsxa9bmcyHTyMly3pJ1ysTp8xs9Dr5qDXzUGvm4VeN0e0BwC0jMuWAAAAAGRp+uAh7/qw/eh1c9Dr5qDXzUGvm2PqF0+f7z/LQU9PT0+vH6uhH19tabrnYeU1oBN63Rz0ujnodXPQ6+ZovX9g8QEpA71uDnrdHPS6Oeh1c0T7M0j3PNxUd88DAAAALlj0CQl63Ryn67fv85DxrnJAjul02Er0ujnodXPQ6wYAwKA76j0P+49Q8hzoD57Godetpv6A2nsAAIDKdengYXifh6McPAAAAAC4Vuns9PMvv+RN4gAAQDnRy+XodXPQ6+ZoqU9nHr744nn/+SPOPAC4BNF7r+h1c9Dr5qDXzUGvm4NeN0frveHR4y+G+x4+e/bs+fyRR/QnQ6+bg143B71ujgvuZ/++2kWv2zn7BfS6Oeh1c9Dr5qDXzVFrn846PE83Sz960r9+86b/7NGjp/3t7R03hgIA1ls44MhCr5uDXjdHtAeu3OPHT/vnz2779GM4eHj85IvxgxxAAAAAAPjo2fPnfTpeePPmzXjwcP/69TA8ffpMPhkAAABAW4ZLnLqu/zK9wtLN477rXg4HDsPBQ/o/r+7vh+uYhgOIu/Hapt0vgmM48J4IB9Hr5qDXzUGvm4NeNwe9bg563Rz0ujnodXOcuf94HPDs+fjSrLe3L/r306HDx4OH9OP+/nX/+PEXw1mIdFPENn7wDdP/P+jW3XRBX38//C89/czHD4n209ehL9J//LtjuV94gLqafummQ3rd6K+mTx9f228+Tn/C/pDx6+f06V7oJ0++GN4Q7sWL7RmHzY/p4CH9ePv27YcDh7vhACIdSKTXc00vyZS+yN1dhgefd/tgf/jf+72gp59tmun3NGv6hx397ufOubx+9+ssO12Pc4v+XtDr5qDXzUGvm6Ncf3t72z979uV00PD0wzFAOrEw9+OTg4fNj3RDxO2Hn8CTp8+Gu6s/f5Q82e/zPf+di37+v3PRz/93rub6nT/P9PP/vVe031FJ/+iT///p9nOD/SN6+gr6T11I//mF9dnoL6XfbOnrpBMH6cAhXX306tV9//6TC5U+/fHZ5mNzn5LCdDbi9es3w9HHrNcf7e656Onpdc9FT1+6390c9Lo56HVz0OvmoNfNcUH96w/evHnbv3s3dzSw++N9///GCY35pZf/pQAAAABJRU5ErkJggg=="
+
+ # Finally, add the PNG image to your Bokeh plot as an image_url
+ p.image_url(
+ url=[logo_url],
+ x=-layout_scale * 0.5,
+ y=layout_scale * 0.5,
+ w=layout_scale,
+ h=layout_scale,
+ anchor=position,
+ global_alpha=logo_alpha,
+ )
+
+
+def style_and_render_graph(p, G, layout_positions, node_attribute, node_colors, centrality):
+ """
+ Apply styling and render the graph into the plot.
+ """
+ from bokeh.plotting import figure, from_networkx
+ from bokeh.models import Circle, MultiLine, HoverTool, ColumnDataSource, Range1d
+ from bokeh.plotting import output_file, show
+
+ from bokeh.embed import file_html
+ from bokeh.resources import CDN
+
+ graph_renderer = from_networkx(G, layout_positions)
+ node_radii = [0.02 + 0.1 * centrality[node] for node in G.nodes()]
+ graph_renderer.node_renderer.data_source.data["radius"] = node_radii
+ graph_renderer.node_renderer.data_source.data["fill_color"] = node_colors
+ graph_renderer.node_renderer.glyph = Circle(
+ radius="radius",
+ fill_color="fill_color",
+ fill_alpha=0.9,
+ line_color="#000000",
+ line_width=1.5,
+ )
+ graph_renderer.edge_renderer.glyph = MultiLine(
+ line_color="#000000",
+ line_alpha=0.3,
+ line_width=1.5,
+ )
+ p.renderers.append(graph_renderer)
+ return graph_renderer
+
+
+async def create_cognee_style_network_with_logo(
+ G,
+ output_filename="cognee_network_with_logo.html",
+ title="Cognee-Style Network",
+ label="group",
+ layout_func=nx.spring_layout,
+ layout_scale=3.0,
+ logo_alpha=0.1,
+ bokeh_object=False,
+):
+ """
+ Create a Cognee-inspired network visualization with an embedded logo.
+ """
+ from bokeh.plotting import figure, from_networkx
+ from bokeh.models import Circle, MultiLine, HoverTool, ColumnDataSource, Range1d
+ from bokeh.plotting import output_file, show
+
+ from bokeh.embed import file_html
+ from bokeh.resources import CDN
+
+ logging.info("Converting graph to serializable format...")
+ G = await convert_to_serializable_graph(G)
+
+ logging.info("Generating layout positions...")
+ layout_positions = generate_layout_positions(G, layout_func, layout_scale)
+
+ logging.info("Assigning node colors...")
+ palette = ["#6510F4", "#0DFF00", "#FFFFFF"]
+ node_colors, color_map = assign_node_colors(G, label, palette)
+
+ logging.info("Calculating centrality...")
+ centrality = nx.degree_centrality(G)
+
+ logging.info("Preparing Bokeh output...")
+ output_file(output_filename)
+ p = figure(
+ title=title,
+ tools="pan,wheel_zoom,save,reset,hover",
+ active_scroll="wheel_zoom",
+ width=1200,
+ height=900,
+ background_fill_color="#F4F4F4",
+ x_range=Range1d(-layout_scale, layout_scale),
+ y_range=Range1d(-layout_scale, layout_scale),
+ )
+ p.toolbar.logo = None
+ p.axis.visible = False
+ p.grid.visible = False
+
+ logging.info("Embedding logo into visualization...")
+ embed_logo(p, layout_scale, logo_alpha, "bottom_right")
+ embed_logo(p, layout_scale, logo_alpha, "top_left")
+
+ logging.info("Styling and rendering graph...")
+ style_and_render_graph(p, G, layout_positions, label, node_colors, centrality)
+
+ logging.info("Adding hover tool...")
+ hover_tool = HoverTool(
+ tooltips=[
+ ("Node", "@index"),
+ (label.capitalize(), f"@{label}"),
+ ("Centrality", "@radius{0.00}"),
+ ],
+ )
+ p.add_tools(hover_tool)
+
+ logging.info(f"Saving visualization to {output_filename}...")
+ html_content = file_html(p, CDN, title)
+ with open(output_filename, "w") as f:
+ f.write(html_content)
- nltk.download("vader_lexicon", quiet=True)
+ logging.info("Visualization complete.")
- # Initialize the VADER Sentiment Intensity Analyzer
- sia = SentimentIntensityAnalyzer()
+ if bokeh_object:
+ return p
+ return html_content
- # Obtain the polarity scores for the text
- polarity_scores = sia.polarity_scores(text)
- return polarity_scores
+def graph_to_tuple(graph):
+ """
+ Converts a networkx graph to a tuple of (nodes, edges).
+
+ :param graph: A networkx graph.
+ :return: A tuple (nodes, edges).
+ """
+ nodes = list(graph.nodes(data=True)) # Get nodes with attributes
+ edges = list(graph.edges(data=True)) # Get edges with attributes
+ return (nodes, edges)
+
def setup_logging(log_level=logging.INFO):
- """ This method sets up the logging configuration. """
+ """This method sets up the logging configuration."""
formatter = logging.Formatter("%(asctime)s - %(levelname)s - %(message)s\n")
stream_handler = logging.StreamHandler(sys.stdout)
stream_handler.setFormatter(formatter)
@@ -298,7 +480,66 @@ def setup_logging(log_level=logging.INFO):
)
+# ---------------- Example Usage ----------------
if __name__ == "__main__":
- sample_text = "I love sunny days, but I hate the rain."
- sentiment_scores = extract_sentiment_vader(sample_text)
- print("Sentiment analysis results:", sentiment_scores)
+ import networkx as nx
+
+ # Create a sample graph
+ nodes = [
+ (1, {"group": "A"}),
+ (2, {"group": "A"}),
+ (3, {"group": "B"}),
+ (4, {"group": "B"}),
+ (5, {"group": "C"}),
+ ]
+ edges = [(1, 2), (2, 3), (3, 4), (4, 5), (5, 1)]
+
+ # Create a NetworkX graph
+ G = nx.Graph()
+ G.add_nodes_from(nodes)
+ G.add_edges_from(edges)
+
+ G = graph_to_tuple(G)
+
+ import asyncio
+
+ output_html = asyncio.run(
+ create_cognee_style_network_with_logo(
+ G,
+ output_filename="example_network.html",
+ title="Example Cognee Network",
+ node_attribute="group", # Attribute to use for coloring nodes
+ layout_func=nx.spring_layout, # Layout function
+ layout_scale=3.0, # Scale for the layout
+ logo_alpha=0.2,
+ )
+ )
+
+ # Call the function
+ # output_html = await create_cognee_style_network_with_logo(
+ # G=G,
+ # output_filename="example_network.html",
+ # title="Example Cognee Network",
+ # node_attribute="group", # Attribute to use for coloring nodes
+ # layout_func=nx.spring_layout, # Layout function
+ # layout_scale=3.0, # Scale for the layout
+ # logo_alpha=0.2, # Transparency of the logo
+ # )
+
+ # Print the output filename
+ print("Network visualization saved as example_network.html")
+
+# # Create a random geometric graph
+# G = nx.random_geometric_graph(50, 0.3)
+# # Assign random group attributes for coloring
+# for i, node in enumerate(G.nodes()):
+# G.nodes[node]["group"] = f"Group {i % 3 + 1}"
+#
+# create_cognee_graph(
+# G,
+# output_filename="cognee_style_network_with_logo.html",
+# title="Cognee-Graph Network",
+# node_attribute="group",
+# layout_func=nx.spring_layout,
+# layout_scale=3.0, # Replace with your logo file path
+# )
diff --git a/cognee/tasks/chunk_naive_llm_classifier/chunk_naive_llm_classifier.py b/cognee/tasks/chunk_naive_llm_classifier/chunk_naive_llm_classifier.py
index 154dadb61..4c058de18 100644
--- a/cognee/tasks/chunk_naive_llm_classifier/chunk_naive_llm_classifier.py
+++ b/cognee/tasks/chunk_naive_llm_classifier/chunk_naive_llm_classifier.py
@@ -8,7 +8,9 @@
from cognee.modules.chunking.models.DocumentChunk import DocumentChunk
-async def chunk_naive_llm_classifier(data_chunks: list[DocumentChunk], classification_model: Type[BaseModel]):
+async def chunk_naive_llm_classifier(
+ data_chunks: list[DocumentChunk], classification_model: Type[BaseModel]
+):
if len(data_chunks) == 0:
return data_chunks
@@ -36,10 +38,17 @@ class Keyword(BaseModel):
collection_name = "classification"
if await vector_engine.has_collection(collection_name):
- existing_data_points = await vector_engine.retrieve(
- collection_name,
- [str(classification_data) for classification_data in list(set(classification_data_points))],
- ) if len(classification_data_points) > 0 else []
+ existing_data_points = (
+ await vector_engine.retrieve(
+ collection_name,
+ [
+ str(classification_data)
+ for classification_data in list(set(classification_data_points))
+ ],
+ )
+ if len(classification_data_points) > 0
+ else []
+ )
existing_points_map = {point.id: True for point in existing_data_points}
else:
@@ -50,7 +59,7 @@ class Keyword(BaseModel):
nodes = []
edges = []
- for (chunk_index, data_chunk) in enumerate(data_chunks):
+ for chunk_index, data_chunk in enumerate(data_chunks):
chunk_classification = chunk_classifications[chunk_index]
classification_type_label = chunk_classification.label.type
classification_type_id = uuid5(NAMESPACE_OID, classification_type_label)
@@ -59,36 +68,42 @@ class Keyword(BaseModel):
data_points.append(
DataPoint[Keyword](
id=str(classification_type_id),
- payload=Keyword.model_validate({
- "uuid": str(classification_type_id),
- "text": classification_type_label,
- "chunk_id": str(data_chunk.chunk_id),
- "document_id": str(data_chunk.document_id),
- }),
+ payload=Keyword.model_validate(
+ {
+ "uuid": str(classification_type_id),
+ "text": classification_type_label,
+ "chunk_id": str(data_chunk.chunk_id),
+ "document_id": str(data_chunk.document_id),
+ }
+ ),
index_fields=["text"],
)
)
- nodes.append((
- str(classification_type_id),
- dict(
- id=str(classification_type_id),
- name=classification_type_label,
- type=classification_type_label,
+ nodes.append(
+ (
+ str(classification_type_id),
+ dict(
+ id=str(classification_type_id),
+ name=classification_type_label,
+ type=classification_type_label,
+ ),
)
- ))
+ )
existing_points_map[classification_type_id] = True
- edges.append((
- str(data_chunk.chunk_id),
- str(classification_type_id),
- "is_media_type",
- dict(
- relationship_name="is_media_type",
- source_node_id=str(data_chunk.chunk_id),
- target_node_id=str(classification_type_id),
- ),
- ))
+ edges.append(
+ (
+ str(data_chunk.chunk_id),
+ str(classification_type_id),
+ "is_media_type",
+ dict(
+ relationship_name="is_media_type",
+ source_node_id=str(data_chunk.chunk_id),
+ target_node_id=str(classification_type_id),
+ ),
+ )
+ )
for classification_subclass in chunk_classification.label.subclass:
classification_subtype_label = classification_subclass.value
@@ -98,47 +113,55 @@ class Keyword(BaseModel):
data_points.append(
DataPoint[Keyword](
id=str(classification_subtype_id),
- payload=Keyword.model_validate({
- "uuid": str(classification_subtype_id),
- "text": classification_subtype_label,
- "chunk_id": str(data_chunk.chunk_id),
- "document_id": str(data_chunk.document_id),
- }),
+ payload=Keyword.model_validate(
+ {
+ "uuid": str(classification_subtype_id),
+ "text": classification_subtype_label,
+ "chunk_id": str(data_chunk.chunk_id),
+ "document_id": str(data_chunk.document_id),
+ }
+ ),
index_fields=["text"],
)
)
- nodes.append((
- str(classification_subtype_id),
- dict(
- id=str(classification_subtype_id),
- name=classification_subtype_label,
- type=classification_subtype_label,
+ nodes.append(
+ (
+ str(classification_subtype_id),
+ dict(
+ id=str(classification_subtype_id),
+ name=classification_subtype_label,
+ type=classification_subtype_label,
+ ),
+ )
+ )
+ edges.append(
+ (
+ str(classification_subtype_id),
+ str(classification_type_id),
+ "is_subtype_of",
+ dict(
+ relationship_name="contains",
+ source_node_id=str(classification_type_id),
+ target_node_id=str(classification_subtype_id),
+ ),
)
- ))
- edges.append((
+ )
+
+ existing_points_map[classification_subtype_id] = True
+
+ edges.append(
+ (
+ str(data_chunk.chunk_id),
str(classification_subtype_id),
- str(classification_type_id),
- "is_subtype_of",
+ "is_classified_as",
dict(
- relationship_name="contains",
- source_node_id=str(classification_type_id),
+ relationship_name="is_classified_as",
+ source_node_id=str(data_chunk.chunk_id),
target_node_id=str(classification_subtype_id),
),
- ))
-
- existing_points_map[classification_subtype_id] = True
-
- edges.append((
- str(data_chunk.chunk_id),
- str(classification_subtype_id),
- "is_classified_as",
- dict(
- relationship_name="is_classified_as",
- source_node_id=str(data_chunk.chunk_id),
- target_node_id=str(classification_subtype_id),
- ),
- ))
+ )
+ )
if len(nodes) > 0 or len(edges) > 0:
await vector_engine.create_data_points(collection_name, data_points)
diff --git a/cognee/tasks/chunks/chunk_by_paragraph.py b/cognee/tasks/chunks/chunk_by_paragraph.py
index 00bb5670c..5c95e97b7 100644
--- a/cognee/tasks/chunks/chunk_by_paragraph.py
+++ b/cognee/tasks/chunks/chunk_by_paragraph.py
@@ -2,7 +2,10 @@
from typing import Dict, Any, Iterator
from .chunk_by_sentence import chunk_by_sentence
-def chunk_by_paragraph(data: str, paragraph_length: int = 1024, batch_paragraphs: bool = True) -> Iterator[Dict[str, Any]]:
+
+def chunk_by_paragraph(
+ data: str, paragraph_length: int = 1024, batch_paragraphs: bool = True
+) -> Iterator[Dict[str, Any]]:
"""
Chunks text by paragraph while preserving exact text reconstruction capability.
When chunks are joined with empty string "", they reproduce the original text exactly.
@@ -12,8 +15,10 @@ def chunk_by_paragraph(data: str, paragraph_length: int = 1024, batch_paragraphs
chunk_index = 0
paragraph_ids = []
last_cut_type = None
-
- for paragraph_id, sentence, word_count, end_type in chunk_by_sentence(data, maximum_length=paragraph_length):
+
+ for paragraph_id, sentence, word_count, end_type in chunk_by_sentence(
+ data, maximum_length=paragraph_length
+ ):
# Check if this sentence would exceed length limit
if current_word_count > 0 and current_word_count + word_count > paragraph_length:
# Yield current chunk
@@ -25,9 +30,9 @@ def chunk_by_paragraph(data: str, paragraph_length: int = 1024, batch_paragraphs
"chunk_index": chunk_index,
"cut_type": last_cut_type,
}
-
+
yield chunk_dict
-
+
# Start new chunk with current sentence
paragraph_ids = []
current_chunk = ""
@@ -37,7 +42,7 @@ def chunk_by_paragraph(data: str, paragraph_length: int = 1024, batch_paragraphs
paragraph_ids.append(paragraph_id)
current_chunk += sentence
current_word_count += word_count
-
+
# Handle end of paragraph
if end_type in ("paragraph_end", "sentence_cut") and not batch_paragraphs:
# For non-batch mode, yield each paragraph separately
@@ -47,16 +52,16 @@ def chunk_by_paragraph(data: str, paragraph_length: int = 1024, batch_paragraphs
"paragraph_ids": paragraph_ids,
"chunk_id": uuid5(NAMESPACE_OID, current_chunk),
"chunk_index": chunk_index,
- "cut_type": end_type
+ "cut_type": end_type,
}
yield chunk_dict
paragraph_ids = []
current_chunk = ""
current_word_count = 0
chunk_index += 1
-
+
last_cut_type = end_type
-
+
# Yield any remaining text
if current_chunk:
chunk_dict = {
@@ -65,8 +70,7 @@ def chunk_by_paragraph(data: str, paragraph_length: int = 1024, batch_paragraphs
"chunk_id": uuid5(NAMESPACE_OID, current_chunk),
"paragraph_ids": paragraph_ids,
"chunk_index": chunk_index,
- "cut_type": "sentence_cut" if last_cut_type == "word" else last_cut_type
+ "cut_type": "sentence_cut" if last_cut_type == "word" else last_cut_type,
}
-
-
- yield chunk_dict
\ No newline at end of file
+
+ yield chunk_dict
diff --git a/cognee/tasks/chunks/chunk_by_sentence.py b/cognee/tasks/chunks/chunk_by_sentence.py
index c6848f066..128ed4047 100644
--- a/cognee/tasks/chunks/chunk_by_sentence.py
+++ b/cognee/tasks/chunks/chunk_by_sentence.py
@@ -1,10 +1,8 @@
-
-
-
from uuid import uuid4
from typing import Optional
from .chunk_by_word import chunk_by_word
+
def chunk_by_sentence(data: str, maximum_length: Optional[int] = None):
sentence = ""
paragraph_id = uuid4()
@@ -16,7 +14,7 @@ def chunk_by_sentence(data: str, maximum_length: Optional[int] = None):
# the word type is 'word', the word doesn't contain any letters
# and words with the same characteristics connect it to a preceding
# word with word_type 'paragraph_end' or 'sentence_end'
- for (word, word_type) in chunk_by_word(data):
+ for word, word_type in chunk_by_word(data):
sentence += word
word_count += 1
@@ -28,7 +26,9 @@ def chunk_by_sentence(data: str, maximum_length: Optional[int] = None):
word_type_state = word_type
break
- if word_type in ["paragraph_end", "sentence_end"] or (maximum_length and (word_count == maximum_length)):
+ if word_type in ["paragraph_end", "sentence_end"] or (
+ maximum_length and (word_count == maximum_length)
+ ):
yield (paragraph_id, sentence, word_count, word_type_state)
sentence = ""
word_count = 0
diff --git a/cognee/tasks/chunks/chunk_by_word.py b/cognee/tasks/chunks/chunk_by_word.py
index ab4d8343e..c42d0cfa1 100644
--- a/cognee/tasks/chunks/chunk_by_word.py
+++ b/cognee/tasks/chunks/chunk_by_word.py
@@ -3,15 +3,16 @@
SENTENCE_ENDINGS = r"[.;!?…]"
PARAGRAPH_ENDINGS = r"[\n\r]"
+
def is_real_paragraph_end(last_char: str, current_pos: int, text: str) -> bool:
"""
Determines if the current position represents a real paragraph ending.
-
+
Args:
last_char: The last processed character
current_pos: Current position in the text
text: The input text
-
+
Returns:
bool: True if this is a real paragraph end, False otherwise
"""
@@ -20,18 +21,19 @@ def is_real_paragraph_end(last_char: str, current_pos: int, text: str) -> bool:
j = current_pos + 1
if j >= len(text):
return False
-
+
next_character = text[j]
while j < len(text) and (re.match(PARAGRAPH_ENDINGS, next_character) or next_character == " "):
j += 1
if j >= len(text):
return False
next_character = text[j]
-
+
if next_character.isupper():
return True
return False
+
def chunk_by_word(data: str):
"""
Chunks text into words and endings while preserving whitespace.
@@ -40,32 +42,32 @@ def chunk_by_word(data: str):
"""
current_chunk = ""
i = 0
-
+
while i < len(data):
character = data[i]
-
+
current_chunk += character
-
+
if character == " ":
yield (current_chunk, "word")
current_chunk = ""
i += 1
continue
-
- if re.match(SENTENCE_ENDINGS, character):
+
+ if re.match(SENTENCE_ENDINGS, character):
# Look ahead for whitespace
next_i = i + 1
while next_i < len(data) and data[next_i] == " ":
current_chunk += data[next_i]
next_i += 1
-
+
is_paragraph_end = next_i < len(data) and re.match(PARAGRAPH_ENDINGS, data[next_i])
yield (current_chunk, "paragraph_end" if is_paragraph_end else "sentence_end")
current_chunk = ""
i = next_i
continue
-
+
i += 1
-
+
if current_chunk:
- yield (current_chunk, "word")
\ No newline at end of file
+ yield (current_chunk, "word")
diff --git a/cognee/tasks/chunks/query_chunks.py b/cognee/tasks/chunks/query_chunks.py
index 399528ee9..5a6d4f666 100644
--- a/cognee/tasks/chunks/query_chunks.py
+++ b/cognee/tasks/chunks/query_chunks.py
@@ -1,5 +1,6 @@
from cognee.infrastructure.databases.vector import get_vector_engine
+
async def query_chunks(query: str) -> list[dict]:
"""
Parameters:
@@ -10,7 +11,7 @@ async def query_chunks(query: str) -> list[dict]:
"""
vector_engine = get_vector_engine()
- found_chunks = await vector_engine.search("document_chunk_text", query, limit = 5)
+ found_chunks = await vector_engine.search("document_chunk_text", query, limit=5)
chunks = [result.payload for result in found_chunks]
diff --git a/cognee/tasks/chunks/remove_disconnected_chunks.py b/cognee/tasks/chunks/remove_disconnected_chunks.py
index 4a36a33ec..60443ecfc 100644
--- a/cognee/tasks/chunks/remove_disconnected_chunks.py
+++ b/cognee/tasks/chunks/remove_disconnected_chunks.py
@@ -1,6 +1,7 @@
from cognee.infrastructure.databases.graph import get_graph_engine
from cognee.modules.chunking.models.DocumentChunk import DocumentChunk
+
async def remove_disconnected_chunks(data_chunks: list[DocumentChunk]) -> list[DocumentChunk]:
graph_engine = await get_graph_engine()
@@ -9,10 +10,12 @@ async def remove_disconnected_chunks(data_chunks: list[DocumentChunk]) -> list[D
obsolete_chunk_ids = []
for document_id in document_ids:
- chunks = await graph_engine.get_successors(document_id, edge_label = "has_chunk")
+ chunks = await graph_engine.get_successors(document_id, edge_label="has_chunk")
for chunk in chunks:
- previous_chunks = await graph_engine.get_predecessors(chunk["uuid"], edge_label = "next_chunk")
+ previous_chunks = await graph_engine.get_predecessors(
+ chunk["uuid"], edge_label="next_chunk"
+ )
if len(previous_chunks) == 0:
obsolete_chunk_ids.append(chunk["uuid"])
diff --git a/cognee/tasks/completion/__init__.py b/cognee/tasks/completion/__init__.py
index 1bf0fa6bb..8a5aaa686 100644
--- a/cognee/tasks/completion/__init__.py
+++ b/cognee/tasks/completion/__init__.py
@@ -1 +1 @@
-from .query_completion import query_completion
\ No newline at end of file
+from .query_completion import query_completion
diff --git a/cognee/tasks/completion/exceptions/__init__.py b/cognee/tasks/completion/exceptions/__init__.py
index 5f80e6ecc..1530bf3f2 100644
--- a/cognee/tasks/completion/exceptions/__init__.py
+++ b/cognee/tasks/completion/exceptions/__init__.py
@@ -6,4 +6,4 @@
from .exceptions import (
NoRelevantDataFound,
-)
\ No newline at end of file
+)
diff --git a/cognee/tasks/completion/exceptions/exceptions.py b/cognee/tasks/completion/exceptions/exceptions.py
index 9b64c01a6..bb4bcb0c8 100644
--- a/cognee/tasks/completion/exceptions/exceptions.py
+++ b/cognee/tasks/completion/exceptions/exceptions.py
@@ -1,11 +1,12 @@
from cognee.exceptions import CogneeApiError
from fastapi import status
+
class NoRelevantDataFound(CogneeApiError):
def __init__(
- self,
- message: str = "Search did not find any data.",
- name: str = "NoRelevantDataFound",
- status_code=status.HTTP_404_NOT_FOUND,
+ self,
+ message: str = "Search did not find any data.",
+ name: str = "NoRelevantDataFound",
+ status_code=status.HTTP_404_NOT_FOUND,
):
- super().__init__(message, name, status_code)
\ No newline at end of file
+ super().__init__(message, name, status_code)
diff --git a/cognee/tasks/completion/query_completion.py b/cognee/tasks/completion/query_completion.py
index 5324676f8..12e5168b0 100644
--- a/cognee/tasks/completion/query_completion.py
+++ b/cognee/tasks/completion/query_completion.py
@@ -14,7 +14,7 @@ async def query_completion(query: str) -> list:
"""
vector_engine = get_vector_engine()
- found_chunks = await vector_engine.search("document_chunk_text", query, limit = 1)
+ found_chunks = await vector_engine.search("document_chunk_text", query, limit=1)
if len(found_chunks) == 0:
raise NoRelevantDataFound
diff --git a/cognee/tasks/documents/check_permissions_on_documents.py b/cognee/tasks/documents/check_permissions_on_documents.py
index ddedefe3f..7d18cbf3a 100644
--- a/cognee/tasks/documents/check_permissions_on_documents.py
+++ b/cognee/tasks/documents/check_permissions_on_documents.py
@@ -1,6 +1,7 @@
from cognee.modules.data.processing.document_types import Document
from cognee.modules.users.permissions.methods import check_permission_on_documents
+
async def check_permissions_on_documents(documents: list[Document], user, permissions):
document_ids = [document.id for document in documents]
diff --git a/cognee/tasks/documents/classify_documents.py b/cognee/tasks/documents/classify_documents.py
index 47beeb917..118da5738 100644
--- a/cognee/tasks/documents/classify_documents.py
+++ b/cognee/tasks/documents/classify_documents.py
@@ -54,13 +54,13 @@ async def classify_documents(data_documents: list[Data]) -> list[Document]:
for data_item in data_documents:
metadata = await get_metadata(data_item.id)
document = EXTENSION_TO_DOCUMENT_CLASS[data_item.extension](
- id = data_item.id,
- title = f"{data_item.name}.{data_item.extension}",
- raw_data_location = data_item.raw_data_location,
- name = data_item.name,
- mime_type = data_item.mime_type,
- metadata_id = metadata.id
+ id=data_item.id,
+ title=f"{data_item.name}.{data_item.extension}",
+ raw_data_location=data_item.raw_data_location,
+ name=data_item.name,
+ mime_type=data_item.mime_type,
+ metadata_id=metadata.id,
)
documents.append(document)
-
+
return documents
diff --git a/cognee/tasks/documents/extract_chunks_from_documents.py b/cognee/tasks/documents/extract_chunks_from_documents.py
index 423b87b69..437d2a3e4 100644
--- a/cognee/tasks/documents/extract_chunks_from_documents.py
+++ b/cognee/tasks/documents/extract_chunks_from_documents.py
@@ -1,7 +1,9 @@
from cognee.modules.data.processing.document_types.Document import Document
-async def extract_chunks_from_documents(documents: list[Document], chunk_size: int = 1024, chunker = 'text_chunker'):
+async def extract_chunks_from_documents(
+ documents: list[Document], chunk_size: int = 1024, chunker="text_chunker"
+):
for document in documents:
- for document_chunk in document.read(chunk_size = chunk_size, chunker = chunker):
+ for document_chunk in document.read(chunk_size=chunk_size, chunker=chunker):
yield document_chunk
diff --git a/cognee/tasks/graph/extract_graph_from_code.py b/cognee/tasks/graph/extract_graph_from_code.py
index 159e9baa4..8688b0af5 100644
--- a/cognee/tasks/graph/extract_graph_from_code.py
+++ b/cognee/tasks/graph/extract_graph_from_code.py
@@ -5,12 +5,13 @@
from cognee.modules.chunking.models.DocumentChunk import DocumentChunk
from cognee.tasks.storage import add_data_points
+
async def extract_graph_from_code(data_chunks: list[DocumentChunk], graph_model: Type[BaseModel]):
chunk_graphs = await asyncio.gather(
*[extract_content_graph(chunk.text, graph_model) for chunk in data_chunks]
)
- for (chunk_index, chunk) in enumerate(data_chunks):
+ for chunk_index, chunk in enumerate(data_chunks):
chunk_graph = chunk_graphs[chunk_index]
await add_data_points(chunk_graph.nodes)
diff --git a/cognee/tasks/graph/extract_graph_from_data.py b/cognee/tasks/graph/extract_graph_from_data.py
index c6a613105..b05c55a9e 100644
--- a/cognee/tasks/graph/extract_graph_from_data.py
+++ b/cognee/tasks/graph/extract_graph_from_data.py
@@ -13,9 +13,7 @@
from cognee.tasks.storage import add_data_points
-async def extract_graph_from_data(
- data_chunks: list[DocumentChunk], graph_model: Type[BaseModel]
-):
+async def extract_graph_from_data(data_chunks: list[DocumentChunk], graph_model: Type[BaseModel]):
chunk_graphs = await asyncio.gather(
*[extract_content_graph(chunk.text, graph_model) for chunk in data_chunks]
)
diff --git a/cognee/tasks/graph/infer_data_ontology.py b/cognee/tasks/graph/infer_data_ontology.py
index 4e11cd9af..ca7807e1e 100644
--- a/cognee/tasks/graph/infer_data_ontology.py
+++ b/cognee/tasks/graph/infer_data_ontology.py
@@ -1,4 +1,4 @@
-""" This module contains the OntologyEngine class which is responsible for adding graph ontology from a JSON or CSV file. """
+"""This module contains the OntologyEngine class which is responsible for adding graph ontology from a JSON or CSV file."""
import csv
import json
@@ -20,7 +20,9 @@
from cognee.infrastructure.databases.graph.get_graph_engine import get_graph_engine
from cognee.infrastructure.files.utils.extract_text_from_file import extract_text_from_file
from cognee.infrastructure.files.utils.guess_file_type import guess_file_type, FileTypeException
-from cognee.modules.data.extraction.knowledge_graph.add_model_class_to_graph import add_model_class_to_graph
+from cognee.modules.data.extraction.knowledge_graph.add_model_class_to_graph import (
+ add_model_class_to_graph,
+)
from cognee.tasks.graph.models import NodeModel, GraphOntology
from cognee.shared.data_models import KnowledgeGraph
from cognee.modules.engine.utils import generate_node_id, generate_node_name
@@ -39,20 +41,26 @@ async def extract_ontology(content: str, response_model: Type[BaseModel]):
class OntologyEngine:
- async def flatten_model(self, model: NodeModel, parent_id: Optional[str] = None) -> Dict[str, Any]:
+ async def flatten_model(
+ self, model: NodeModel, parent_id: Optional[str] = None
+ ) -> Dict[str, Any]:
"""Flatten the model to a dictionary."""
result = model.dict()
result["parent_id"] = parent_id
if model.default_relationship:
- result.update({
- "relationship_type": model.default_relationship.type,
- "relationship_source": model.default_relationship.source,
- "relationship_target": model.default_relationship.target
- })
+ result.update(
+ {
+ "relationship_type": model.default_relationship.type,
+ "relationship_source": model.default_relationship.source,
+ "relationship_target": model.default_relationship.target,
+ }
+ )
return result
- async def recursive_flatten(self, items: Union[List[Dict[str, Any]], Dict[str, Any]], parent_id: Optional[str] = None) -> List[Dict[str, Any]]:
- """Recursively flatten the items. """
+ async def recursive_flatten(
+ self, items: Union[List[Dict[str, Any]], Dict[str, Any]], parent_id: Optional[str] = None
+ ) -> List[Dict[str, Any]]:
+ """Recursively flatten the items."""
flat_list = []
if isinstance(items, list):
@@ -80,8 +88,10 @@ async def load_data(self, file_path: str) -> Union[List[Dict[str, Any]], Dict[st
else:
raise IngestionError(message="Unsupported file format")
except Exception as e:
- raise IngestionError(message=f"Failed to load data from {file_path}: {e}",
- status_code=status.HTTP_422_UNPROCESSABLE_ENTITY)
+ raise IngestionError(
+ message=f"Failed to load data from {file_path}: {e}",
+ status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
+ )
async def add_graph_ontology(self, file_path: str = None, documents: list = None):
"""Add graph ontology from a JSON or CSV file or infer from documents content."""
@@ -109,31 +119,43 @@ async def add_graph_ontology(self, file_path: str = None, documents: list = None
initial_chunks_and_ids.append({base_file.id: chunks_with_ids})
except FileTypeException:
- logger.warning("File (%s) has an unknown file type. We are skipping it.", file["id"])
-
+ logger.warning(
+ "File (%s) has an unknown file type. We are skipping it.", file["id"]
+ )
ontology = await extract_ontology(str(initial_chunks_and_ids), GraphOntology)
graph_client = await get_graph_engine()
- await graph_client.add_nodes([(node.id, dict(
- uuid = generate_node_id(node.id),
- name = generate_node_name(node.name),
- type = generate_node_id(node.id),
- description = node.description,
- updated_at = datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M:%S"),
- )) for node in ontology.nodes])
-
- await graph_client.add_edges((
- generate_node_id(edge.source_id),
- generate_node_id(edge.target_id),
- edge.relationship_type,
- dict(
- source_node_id = generate_node_id(edge.source_id),
- target_node_id = generate_node_id(edge.target_id),
- relationship_name = edge.relationship_type,
- updated_at = datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M:%S"),
- ),
- ) for edge in ontology.edges)
+ await graph_client.add_nodes(
+ [
+ (
+ node.id,
+ dict(
+ uuid=generate_node_id(node.id),
+ name=generate_node_name(node.name),
+ type=generate_node_id(node.id),
+ description=node.description,
+ updated_at=datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M:%S"),
+ ),
+ )
+ for node in ontology.nodes
+ ]
+ )
+
+ await graph_client.add_edges(
+ (
+ generate_node_id(edge.source_id),
+ generate_node_id(edge.target_id),
+ edge.relationship_type,
+ dict(
+ source_node_id=generate_node_id(edge.source_id),
+ target_node_id=generate_node_id(edge.target_id),
+ relationship_name=edge.relationship_type,
+ updated_at=datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M:%S"),
+ ),
+ )
+ for edge in ontology.edges
+ )
else:
dataset_level_information = documents[0][1]
@@ -152,17 +174,23 @@ async def add_graph_ontology(self, file_path: str = None, documents: list = None
if node_id in valid_ids:
await graph_client.add_node(node_id, node_data)
if node_id not in valid_ids:
- raise EntityNotFoundError(message=f"Node ID {node_id} not found in the dataset")
- if pd.notna(row.get("relationship_source")) and pd.notna(row.get("relationship_target")):
+ raise EntityNotFoundError(
+ message=f"Node ID {node_id} not found in the dataset"
+ )
+ if pd.notna(row.get("relationship_source")) and pd.notna(
+ row.get("relationship_target")
+ ):
await graph_client.add_edge(
row["relationship_source"],
row["relationship_target"],
relationship_name=row["relationship_type"],
- edge_properties = {
+ edge_properties={
"source_node_id": row["relationship_source"],
"target_node_id": row["relationship_target"],
"relationship_name": row["relationship_type"],
- "updated_at": datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M:%S"),
+ "updated_at": datetime.now(timezone.utc).strftime(
+ "%Y-%m-%d %H:%M:%S"
+ ),
},
)
@@ -171,10 +199,10 @@ async def add_graph_ontology(self, file_path: str = None, documents: list = None
raise RuntimeError(f"Failed to add graph ontology from {file_path}: {e}") from e
-async def infer_data_ontology(documents, ontology_model = KnowledgeGraph, root_node_id = None):
+async def infer_data_ontology(documents, ontology_model=KnowledgeGraph, root_node_id=None):
if ontology_model == KnowledgeGraph:
ontology_engine = OntologyEngine()
- root_node_id = await ontology_engine.add_graph_ontology(documents = documents)
+ root_node_id = await ontology_engine.add_graph_ontology(documents=documents)
else:
graph_engine = await get_graph_engine()
await add_model_class_to_graph(ontology_model, graph_engine)
diff --git a/cognee/tasks/graph/models.py b/cognee/tasks/graph/models.py
index 5b1108e6a..bfffb0262 100644
--- a/cognee/tasks/graph/models.py
+++ b/cognee/tasks/graph/models.py
@@ -1,31 +1,36 @@
from typing import Any, Dict, List, Optional, Union
from pydantic import BaseModel, Field
+
class RelationshipModel(BaseModel):
type: str
source: str
target: str
+
class NodeModel(BaseModel):
node_id: str
name: str
default_relationship: Optional[RelationshipModel] = None
children: List[Union[Dict[str, Any], "NodeModel"]] = Field(default_factory=list)
+
NodeModel.model_rebuild()
class OntologyNode(BaseModel):
- id: str = Field(..., description = "Unique identifier made from node name.")
+ id: str = Field(..., description="Unique identifier made from node name.")
name: str
description: str
+
class OntologyEdge(BaseModel):
id: str
source_id: str
target_id: str
relationship_type: str
+
class GraphOntology(BaseModel):
nodes: list[OntologyNode]
- edges: list[OntologyEdge]
\ No newline at end of file
+ edges: list[OntologyEdge]
diff --git a/cognee/tasks/graph/query_graph_connections.py b/cognee/tasks/graph/query_graph_connections.py
index 4020ddd13..faec94546 100644
--- a/cognee/tasks/graph/query_graph_connections.py
+++ b/cognee/tasks/graph/query_graph_connections.py
@@ -2,7 +2,8 @@
from cognee.infrastructure.databases.graph import get_graph_engine
from cognee.infrastructure.databases.vector import get_vector_engine
-async def query_graph_connections(query: str, exploration_levels = 1) -> list[(str, str, str)]:
+
+async def query_graph_connections(query: str, exploration_levels=1) -> list[(str, str, str)]:
"""
Find the neighbours of a given node in the graph and return formed sentences.
@@ -27,8 +28,8 @@ async def query_graph_connections(query: str, exploration_levels = 1) -> list[(s
else:
vector_engine = get_vector_engine()
results = await asyncio.gather(
- vector_engine.search("entity_name", query_text = query, limit = 5),
- vector_engine.search("entity_type_name", query_text = query, limit = 5),
+ vector_engine.search("entity_name", query_text=query, limit=5),
+ vector_engine.search("entity_type_name", query_text=query, limit=5),
)
results = [*results[0], *results[1]]
relevant_results = [result for result in results if result.score < 0.5][:5]
@@ -44,7 +45,6 @@ async def query_graph_connections(query: str, exploration_levels = 1) -> list[(s
for neighbours in node_connections_results:
node_connections.extend(neighbours)
-
unique_node_connections_map = {}
unique_node_connections = []
for node_connection in node_connections:
@@ -58,5 +58,4 @@ async def query_graph_connections(query: str, exploration_levels = 1) -> list[(s
unique_node_connections.append(node_connection)
-
return unique_node_connections
diff --git a/cognee/tasks/ingestion/get_dlt_destination.py b/cognee/tasks/ingestion/get_dlt_destination.py
index 12042c75b..2de0d7f0d 100644
--- a/cognee/tasks/ingestion/get_dlt_destination.py
+++ b/cognee/tasks/ingestion/get_dlt_destination.py
@@ -6,6 +6,7 @@
from cognee.infrastructure.databases.relational import get_relational_config
+
@lru_cache
def get_dlt_destination() -> Union[type[dlt.destinations.sqlalchemy], None]:
"""
@@ -21,24 +22,24 @@ def get_dlt_destination() -> Union[type[dlt.destinations.sqlalchemy], None]:
# When sqlite is the database provider hostname, port, username and password should not be forwarded.
# The database is found by combining the path location and the database name
destination = dlt.destinations.sqlalchemy(
- credentials = {
- "database": os.path.join(relational_config.db_path, relational_config.db_name),
- "drivername": relational_config.db_provider,
- },
- )
+ credentials={
+ "database": os.path.join(relational_config.db_path, relational_config.db_name),
+ "drivername": relational_config.db_provider,
+ },
+ )
elif relational_config.db_provider == "postgres":
# The dlt library doesn't accept postgres as the drivername, it only accepts postgresql
destination = dlt.destinations.sqlalchemy(
- credentials = {
- "host": relational_config.db_host,
- "port": relational_config.db_port,
- "username": relational_config.db_username,
- "password": relational_config.db_password,
- "database": relational_config.db_name,
- "drivername": "postgresql",
- },
- )
+ credentials={
+ "host": relational_config.db_host,
+ "port": relational_config.db_port,
+ "username": relational_config.db_username,
+ "password": relational_config.db_password,
+ "database": relational_config.db_name,
+ "drivername": "postgresql",
+ },
+ )
else:
destination = None
-
+
return destination
diff --git a/cognee/tasks/ingestion/ingest_data.py b/cognee/tasks/ingestion/ingest_data.py
index 9418d035b..cf7dd38ad 100644
--- a/cognee/tasks/ingestion/ingest_data.py
+++ b/cognee/tasks/ingestion/ingest_data.py
@@ -9,18 +9,19 @@
from cognee.modules.users.permissions.methods import give_permission_on_document
from .get_dlt_destination import get_dlt_destination
+
async def ingest_data(file_paths: list[str], dataset_name: str, user: User):
destination = get_dlt_destination()
pipeline = dlt.pipeline(
- pipeline_name = "file_load_from_filesystem",
- destination = destination,
+ pipeline_name="file_load_from_filesystem",
+ destination=destination,
)
- @dlt.resource(standalone = True, merge_key = "id")
+ @dlt.resource(standalone=True, merge_key="id")
async def data_resources(file_paths: str):
for file_path in file_paths:
- with open(file_path.replace("file://", ""), mode = "rb") as file:
+ with open(file_path.replace("file://", ""), mode="rb") as file:
classified_data = ingestion.classify(file)
data_id = ingestion.identify(classified_data)
file_metadata = classified_data.get_metadata()
@@ -39,39 +40,39 @@ async def data_storing(table_name, dataset_name, user: User):
# Read metadata stored with dlt
files_metadata = await db_engine.get_all_data_from_table(table_name, dataset_name)
for file_metadata in files_metadata:
- from sqlalchemy import select
- from cognee.modules.data.models import Data
- dataset = await create_dataset(dataset_name, user.id, session)
+ from sqlalchemy import select
+ from cognee.modules.data.models import Data
- data = (await session.execute(
- select(Data).filter(Data.id == UUID(file_metadata["id"]))
- )).scalar_one_or_none()
+ dataset = await create_dataset(dataset_name, user.id, session)
- if data is not None:
- data.name = file_metadata["name"]
- data.raw_data_location = file_metadata["file_path"]
- data.extension = file_metadata["extension"]
- data.mime_type = file_metadata["mime_type"]
+ data = (
+ await session.execute(select(Data).filter(Data.id == UUID(file_metadata["id"])))
+ ).scalar_one_or_none()
- await session.merge(data)
- await session.commit()
- else:
- data = Data(
- id = UUID(file_metadata["id"]),
- name = file_metadata["name"],
- raw_data_location = file_metadata["file_path"],
- extension = file_metadata["extension"],
- mime_type = file_metadata["mime_type"],
- )
+ if data is not None:
+ data.name = file_metadata["name"]
+ data.raw_data_location = file_metadata["file_path"]
+ data.extension = file_metadata["extension"]
+ data.mime_type = file_metadata["mime_type"]
- dataset.data.append(data)
- await session.commit()
+ await session.merge(data)
+ await session.commit()
+ else:
+ data = Data(
+ id=UUID(file_metadata["id"]),
+ name=file_metadata["name"],
+ raw_data_location=file_metadata["file_path"],
+ extension=file_metadata["extension"],
+ mime_type=file_metadata["mime_type"],
+ )
- await give_permission_on_document(user, UUID(file_metadata["id"]), "read")
- await give_permission_on_document(user, UUID(file_metadata["id"]), "write")
+ dataset.data.append(data)
+ await session.commit()
+ await give_permission_on_document(user, UUID(file_metadata["id"]), "read")
+ await give_permission_on_document(user, UUID(file_metadata["id"]), "write")
- send_telemetry("cognee.add EXECUTION STARTED", user_id = user.id)
+ send_telemetry("cognee.add EXECUTION STARTED", user_id=user.id)
db_engine = get_relational_engine()
@@ -82,9 +83,9 @@ async def data_storing(table_name, dataset_name, user: User):
# Sqlite doesn't support schemas
run_info = pipeline.run(
data_resources(file_paths),
- table_name = "file_metadata",
- dataset_name = "main",
- write_disposition = "merge",
+ table_name="file_metadata",
+ dataset_name="main",
+ write_disposition="merge",
)
else:
run_info = pipeline.run(
@@ -95,6 +96,6 @@ async def data_storing(table_name, dataset_name, user: User):
)
await data_storing("file_metadata", dataset_name, user)
- send_telemetry("cognee.add EXECUTION COMPLETED", user_id = user.id)
+ send_telemetry("cognee.add EXECUTION COMPLETED", user_id=user.id)
return run_info
diff --git a/cognee/tasks/ingestion/ingest_data_with_metadata.py b/cognee/tasks/ingestion/ingest_data_with_metadata.py
index c6b42f482..04396485c 100644
--- a/cognee/tasks/ingestion/ingest_data_with_metadata.py
+++ b/cognee/tasks/ingestion/ingest_data_with_metadata.py
@@ -19,8 +19,8 @@ async def ingest_data_with_metadata(data: Any, dataset_name: str, user: User):
destination = get_dlt_destination()
pipeline = dlt.pipeline(
- pipeline_name = "file_load_from_filesystem",
- destination = destination,
+ pipeline_name="file_load_from_filesystem",
+ destination=destination,
)
@dlt.resource(standalone=True, primary_key="id", merge_key="id")
@@ -49,14 +49,12 @@ async def data_storing(data: Any, dataset_name: str, user: User):
# Process data
for data_item in data:
- file_path = await save_data_item_with_metadata_to_storage(
- data_item, dataset_name
- )
+ file_path = await save_data_item_with_metadata_to_storage(data_item, dataset_name)
file_paths.append(file_path)
# Ingest data and add metadata
- with open(file_path.replace("file://", ""), mode = "rb") as file:
+ with open(file_path.replace("file://", ""), mode="rb") as file:
classified_data = ingestion.classify(file)
# data_id is the hash of file contents + owner id to avoid duplicate data
@@ -88,19 +86,22 @@ async def data_storing(data: Any, dataset_name: str, user: User):
await session.merge(data_point)
else:
data_point = Data(
- id = data_id,
- name = file_metadata["name"],
- raw_data_location = file_metadata["file_path"],
- extension = file_metadata["extension"],
- mime_type = file_metadata["mime_type"],
- owner_id = user.id,
- content_hash = file_metadata["content_hash"],
+ id=data_id,
+ name=file_metadata["name"],
+ raw_data_location=file_metadata["file_path"],
+ extension=file_metadata["extension"],
+ mime_type=file_metadata["mime_type"],
+ owner_id=user.id,
+ content_hash=file_metadata["content_hash"],
)
# Check if data is already in dataset
dataset_data = (
- await session.execute(select(DatasetData).filter(DatasetData.data_id == data_id,
- DatasetData.dataset_id == dataset.id))
+ await session.execute(
+ select(DatasetData).filter(
+ DatasetData.data_id == data_id, DatasetData.dataset_id == dataset.id
+ )
+ )
).scalar_one_or_none()
# If data is not present in dataset add it
if dataset_data is None:
diff --git a/cognee/tasks/ingestion/resolve_data_directories.py b/cognee/tasks/ingestion/resolve_data_directories.py
index 980756805..e549688d8 100644
--- a/cognee/tasks/ingestion/resolve_data_directories.py
+++ b/cognee/tasks/ingestion/resolve_data_directories.py
@@ -1,7 +1,10 @@
import os
from typing import List, Union, BinaryIO
-async def resolve_data_directories(data: Union[BinaryIO, List[BinaryIO], str, List[str]], include_subdirectories: bool = True):
+
+async def resolve_data_directories(
+ data: Union[BinaryIO, List[BinaryIO], str, List[str]], include_subdirectories: bool = True
+):
"""
Resolves directories by replacing them with their contained files.
@@ -28,10 +31,14 @@ async def resolve_data_directories(data: Union[BinaryIO, List[BinaryIO], str, Li
else:
# Add all files (not subdirectories) in the directory
resolved_data.extend(
- [os.path.join(item, f) for f in os.listdir(item) if os.path.isfile(os.path.join(item, f))]
+ [
+ os.path.join(item, f)
+ for f in os.listdir(item)
+ if os.path.isfile(os.path.join(item, f))
+ ]
)
- else: # If it's a file or text add it directly
+ else: # If it's a file or text add it directly
resolved_data.append(item)
- else: # If it's not a string add it directly
+ else: # If it's not a string add it directly
resolved_data.append(item)
return resolved_data
diff --git a/cognee/tasks/ingestion/save_data_item_to_storage.py b/cognee/tasks/ingestion/save_data_item_to_storage.py
index 88d499e74..9191f7ebc 100644
--- a/cognee/tasks/ingestion/save_data_item_to_storage.py
+++ b/cognee/tasks/ingestion/save_data_item_to_storage.py
@@ -3,8 +3,8 @@
from cognee.modules.ingestion.exceptions import IngestionError
from cognee.modules.ingestion import save_data_to_file
-def save_data_item_to_storage(data_item: Union[BinaryIO, str], dataset_name: str) -> str:
+def save_data_item_to_storage(data_item: Union[BinaryIO, str], dataset_name: str) -> str:
# data is a file object coming from upload.
if hasattr(data_item, "file"):
file_path = save_data_to_file(data_item.file, filename=data_item.filename)
@@ -19,4 +19,4 @@ def save_data_item_to_storage(data_item: Union[BinaryIO, str], dataset_name: str
else:
raise IngestionError(message=f"Data type not supported: {type(data_item)}")
- return file_path
\ No newline at end of file
+ return file_path
diff --git a/cognee/tasks/ingestion/save_data_item_with_metadata_to_storage.py b/cognee/tasks/ingestion/save_data_item_with_metadata_to_storage.py
index 06dde11bd..92697abb7 100644
--- a/cognee/tasks/ingestion/save_data_item_with_metadata_to_storage.py
+++ b/cognee/tasks/ingestion/save_data_item_with_metadata_to_storage.py
@@ -7,7 +7,6 @@
async def save_data_item_with_metadata_to_storage(
data_item: Union[BinaryIO, str, Any], dataset_name: str
) -> str:
-
if "llama_index" in str(type(data_item)):
# Dynamic import is used because the llama_index module is optional.
from .transform_data import get_data_from_llama_index
@@ -16,9 +15,7 @@ async def save_data_item_with_metadata_to_storage(
# data is a file object coming from upload.
elif hasattr(data_item, "file"):
- file_path = save_data_to_file(
- data_item.file, filename=data_item.filename
- )
+ file_path = save_data_to_file(data_item.file, filename=data_item.filename)
elif isinstance(data_item, str):
# data is a file path
diff --git a/cognee/tasks/ingestion/save_data_to_storage.py b/cognee/tasks/ingestion/save_data_to_storage.py
index 85eb81582..a56857261 100644
--- a/cognee/tasks/ingestion/save_data_to_storage.py
+++ b/cognee/tasks/ingestion/save_data_to_storage.py
@@ -1,6 +1,7 @@
from typing import Union, BinaryIO
from cognee.tasks.ingestion.save_data_item_to_storage import save_data_item_to_storage
+
def save_data_to_storage(data: Union[BinaryIO, str], dataset_name) -> list[str]:
if not isinstance(data, list):
# Convert data to a list as we work with lists further down.
diff --git a/cognee/tasks/ingestion/transform_data.py b/cognee/tasks/ingestion/transform_data.py
index 898ac6e71..cc75c7a65 100644
--- a/cognee/tasks/ingestion/transform_data.py
+++ b/cognee/tasks/ingestion/transform_data.py
@@ -3,16 +3,17 @@
from cognee.modules.ingestion import save_data_to_file
from typing import Union
+
def get_data_from_llama_index(data_point: Union[Document, ImageDocument], dataset_name: str) -> str:
# Specific type checking is used to ensure it's not a child class from Document
- if type(data_point) == Document:
+ if isinstance(data_point, Document) and type(data_point) is Document:
file_path = data_point.metadata.get("file_path")
if file_path is None:
file_path = save_data_to_file(data_point.text)
return file_path
return file_path
- elif type(data_point) == ImageDocument:
+ elif isinstance(data_point, ImageDocument) and type(data_point) is ImageDocument:
if data_point.image_path is None:
file_path = save_data_to_file(data_point.text)
return file_path
- return data_point.image_path
\ No newline at end of file
+ return data_point.image_path
diff --git a/cognee/tasks/repo_processor/__init__.py b/cognee/tasks/repo_processor/__init__.py
index fa754028e..6dc032547 100644
--- a/cognee/tasks/repo_processor/__init__.py
+++ b/cognee/tasks/repo_processor/__init__.py
@@ -1,8 +1,7 @@
-import logging
-
-logger = logging.getLogger("task:repo_processor")
-
from .enrich_dependency_graph import enrich_dependency_graph
from .expand_dependency_graph import expand_dependency_graph
from .get_non_code_files import get_data_list_for_user, get_non_py_files
from .get_repo_file_dependencies import get_repo_file_dependencies
+import logging
+
+logger = logging.getLogger("task:repo_processor")
diff --git a/cognee/tasks/repo_processor/enrich_dependency_graph.py b/cognee/tasks/repo_processor/enrich_dependency_graph.py
index 8cce7c42a..5f2b91d95 100644
--- a/cognee/tasks/repo_processor/enrich_dependency_graph.py
+++ b/cognee/tasks/repo_processor/enrich_dependency_graph.py
@@ -8,7 +8,9 @@
from cognee.infrastructure.databases.graph import get_graph_engine
-def topologically_sort_subgraph(subgraph_node_to_indegree: Dict[str, int], graph: nx.DiGraph) -> List[str]:
+def topologically_sort_subgraph(
+ subgraph_node_to_indegree: Dict[str, int], graph: nx.DiGraph
+) -> List[str]:
"""Performs a topological sort on a subgraph based on node indegrees."""
results = []
remaining_nodes = subgraph_node_to_indegree.copy()
@@ -33,13 +35,8 @@ def topologically_sort(graph: nx.DiGraph) -> List[str]:
topological_order = []
for subgraph in subgraphs:
- node_to_indegree = {
- node: len(list(subgraph.successors(node)))
- for node in subgraph.nodes
- }
- topological_order.extend(
- topologically_sort_subgraph(node_to_indegree, subgraph)
- )
+ node_to_indegree = {node: len(list(subgraph.successors(node))) for node in subgraph.nodes}
+ topological_order.extend(topologically_sort_subgraph(node_to_indegree, subgraph))
return topological_order
@@ -62,7 +59,7 @@ async def node_enrich_and_connect(
graph_engine = await get_graph_engine()
for desc_id in node_descendants:
- if desc_id not in topological_order[:topological_rank + 1]:
+ if desc_id not in topological_order[: topological_rank + 1]:
continue
desc = None
@@ -83,7 +80,9 @@ async def node_enrich_and_connect(
node.depends_directly_on.extend(new_connections)
-async def enrich_dependency_graph(data_points: list[DataPoint]) -> AsyncGenerator[list[DataPoint], None]:
+async def enrich_dependency_graph(
+ data_points: list[DataPoint],
+) -> AsyncGenerator[list[DataPoint], None]:
"""Enriches the graph with topological ranks and 'depends_on' edges."""
nodes = []
edges = []
@@ -94,9 +93,9 @@ async def enrich_dependency_graph(data_points: list[DataPoint]) -> AsyncGenerato
for data_point in data_points:
graph_nodes, graph_edges = await get_graph_from_model(
data_point,
- added_nodes = added_nodes,
- added_edges = added_edges,
- visited_properties = visited_properties,
+ added_nodes=added_nodes,
+ added_edges=added_edges,
+ visited_properties=visited_properties,
)
nodes.extend(graph_nodes)
edges.extend(graph_edges)
@@ -122,7 +121,7 @@ async def enrich_dependency_graph(data_points: list[DataPoint]) -> AsyncGenerato
data_points_map = {data_point.id: data_point for data_point in data_points}
# data_points_futures = []
- for data_point in tqdm(data_points, desc = "Enriching dependency graph", unit = "data_point"):
+ for data_point in tqdm(data_points, desc="Enriching dependency graph", unit="data_point"):
if data_point.id not in node_rank_map:
continue
@@ -131,7 +130,7 @@ async def enrich_dependency_graph(data_points: list[DataPoint]) -> AsyncGenerato
await node_enrich_and_connect(graph, topological_order, data_point, data_points_map)
yield data_point
-
+
# await asyncio.gather(*data_points_futures)
# return data_points
diff --git a/cognee/tasks/repo_processor/expand_dependency_graph.py b/cognee/tasks/repo_processor/expand_dependency_graph.py
index 43a451bd6..de26fe8d4 100644
--- a/cognee/tasks/repo_processor/expand_dependency_graph.py
+++ b/cognee/tasks/repo_processor/expand_dependency_graph.py
@@ -1,11 +1,13 @@
from typing import AsyncGenerator
from uuid import NAMESPACE_OID, uuid5
+
# from tqdm import tqdm
from cognee.infrastructure.engine import DataPoint
from cognee.shared.CodeGraphEntities import CodeFile, CodePart
from cognee.tasks.repo_processor.extract_code_parts import extract_code_parts
from cognee.tasks.repo_processor import logger
+
def _add_code_parts_nodes_and_edges(code_file: CodeFile, part_type, code_parts) -> None:
"""Add code part nodes and edges for a specific part type."""
if not code_parts:
@@ -21,12 +23,14 @@ def _add_code_parts_nodes_and_edges(code_file: CodeFile, part_type, code_parts)
part_node_id = uuid5(NAMESPACE_OID, f"{code_file.id}_{part_type}_{idx}")
- part_nodes.append(CodePart(
- id = part_node_id,
- type = part_type,
- # part_of = code_file,
- source_code = code_part,
- ))
+ part_nodes.append(
+ CodePart(
+ id=part_node_id,
+ type=part_type,
+ # part_of = code_file,
+ source_code=code_part,
+ )
+ )
# graph.add_node(part_node_id, source_code=code_part, node_type=part_type)
# graph.add_edge(parent_node_id, part_node_id, relation="contains")
@@ -54,7 +58,9 @@ def _process_single_node(code_file: CodeFile) -> None:
_add_code_parts_nodes_and_edges(code_file, part_type, code_parts)
-async def expand_dependency_graph(data_points: list[DataPoint]) -> AsyncGenerator[list[DataPoint], None]:
+async def expand_dependency_graph(
+ data_points: list[DataPoint],
+) -> AsyncGenerator[list[DataPoint], None]:
"""Process Python file nodes, adding code part nodes and edges."""
# for data_point in tqdm(data_points, desc = "Expand dependency graph", unit = "data_point"):
for data_point in data_points:
diff --git a/cognee/tasks/repo_processor/extract_code_parts.py b/cognee/tasks/repo_processor/extract_code_parts.py
index d772c73f5..76cfef538 100644
--- a/cognee/tasks/repo_processor/extract_code_parts.py
+++ b/cognee/tasks/repo_processor/extract_code_parts.py
@@ -9,19 +9,19 @@ def _extract_parts_from_module(module, parts_dict: Dict[str, List[str]]) -> Dict
current_top_level_code = []
child_to_code_type = {
- 'classdef': "classes",
- 'funcdef': "functions",
- 'import_name': "imports",
- 'import_from': "imports",
+ "classdef": "classes",
+ "funcdef": "functions",
+ "import_name": "imports",
+ "import_from": "imports",
}
for child in module.children:
- if child.type == 'simple_stmt':
+ if child.type == "simple_stmt":
current_top_level_code.append(child.get_code())
continue
if current_top_level_code:
- parts_dict["top_level_code"].append('\n'.join(current_top_level_code))
+ parts_dict["top_level_code"].append("\n".join(current_top_level_code))
current_top_level_code = []
if child.type in child_to_code_type:
@@ -29,10 +29,10 @@ def _extract_parts_from_module(module, parts_dict: Dict[str, List[str]]) -> Dict
parts_dict[code_type].append(child.get_code())
if current_top_level_code:
- parts_dict["top_level_code"].append('\n'.join(current_top_level_code))
+ parts_dict["top_level_code"].append("\n".join(current_top_level_code))
if parts_dict["imports"]:
- parts_dict["imports"] = ['\n'.join(parts_dict["imports"])]
+ parts_dict["imports"] = ["\n".join(parts_dict["imports"])]
return parts_dict
diff --git a/cognee/tasks/repo_processor/get_local_dependencies.py b/cognee/tasks/repo_processor/get_local_dependencies.py
index fb4c68710..b443829c9 100644
--- a/cognee/tasks/repo_processor/get_local_dependencies.py
+++ b/cognee/tasks/repo_processor/get_local_dependencies.py
@@ -29,15 +29,15 @@ def _get_code_entities(node: parso.tree.NodeOrLeaf) -> List[Dict[str, any]]:
"""
code_entity_list = []
- if not hasattr(node, 'children'):
+ if not hasattr(node, "children"):
return code_entity_list
- name_nodes = (child for child in node.children if child.type == 'name')
+ name_nodes = (child for child in node.children if child.type == "name")
for name_node in name_nodes:
code_entity = {
- 'name': name_node.value,
- 'line': name_node.start_pos[0],
- 'column': name_node.start_pos[1]
+ "name": name_node.value,
+ "line": name_node.start_pos[0],
+ "column": name_node.start_pos[1],
}
code_entity_list.append(code_entity)
@@ -97,7 +97,9 @@ async def _extract_dependencies(script_path: str) -> List[str]:
return sorted(str_paths)
-async def get_local_script_dependencies(script_path: str, repo_path: Optional[str] = None) -> List[str]:
+async def get_local_script_dependencies(
+ script_path: str, repo_path: Optional[str] = None
+) -> List[str]:
"""
Extract and return a list of unique module paths that the script depends on.
"""
@@ -117,7 +119,9 @@ async def get_local_script_dependencies(script_path: str, repo_path: Optional[st
return await _extract_dependencies(script_path)
if not script_path.is_relative_to(repo_path):
- logger.warning(f"Script {script_path} not in repo {repo_path}. Proceeding without repo_path.")
+ logger.warning(
+ f"Script {script_path} not in repo {repo_path}. Proceeding without repo_path."
+ )
return await _extract_dependencies(script_path)
with add_sys_path(str(repo_path)):
diff --git a/cognee/tasks/repo_processor/get_non_code_files.py b/cognee/tasks/repo_processor/get_non_code_files.py
index 671b998d9..9c69afd00 100644
--- a/cognee/tasks/repo_processor/get_non_code_files.py
+++ b/cognee/tasks/repo_processor/get_non_code_files.py
@@ -6,8 +6,7 @@
from cognee.infrastructure.engine import DataPoint
from cognee.modules.data.methods import get_datasets
from cognee.modules.data.methods.get_dataset_data import get_dataset_data
-from cognee.modules.data.methods.get_datasets_by_name import \
- get_datasets_by_name
+from cognee.modules.data.methods.get_datasets_by_name import get_datasets_by_name
from cognee.modules.data.models import Data
from cognee.modules.data.operations.write_metadata import write_metadata
from cognee.modules.ingestion.data_types import BinaryData
@@ -21,8 +20,13 @@ async def get_non_py_files(repo_path):
return {}
IGNORED_PATTERNS = {
- '.git', '__pycache__', '*.pyc', '*.pyo', '*.pyd',
- 'node_modules', '*.egg-info'
+ ".git",
+ "__pycache__",
+ "*.pyc",
+ "*.pyo",
+ "*.pyd",
+ "node_modules",
+ "*.egg-info",
}
def should_process(path):
@@ -30,7 +34,8 @@ def should_process(path):
non_py_files_paths = [
os.path.join(root, file)
- for root, _, files in os.walk(repo_path) for file in files
+ for root, _, files in os.walk(repo_path)
+ for file in files
if not file.endswith(".py") and should_process(os.path.join(root, file))
]
return non_py_files_paths
@@ -45,4 +50,4 @@ async def get_data_list_for_user(_, dataset_name, user):
for dataset in datasets:
data_docs: list[Data] = await get_dataset_data(dataset_id=dataset.id)
data_documents.extend(data_docs)
- return data_documents
\ No newline at end of file
+ return data_documents
diff --git a/cognee/tasks/repo_processor/get_repo_file_dependencies.py b/cognee/tasks/repo_processor/get_repo_file_dependencies.py
index b54c1f152..f9be53d7d 100644
--- a/cognee/tasks/repo_processor/get_repo_file_dependencies.py
+++ b/cognee/tasks/repo_processor/get_repo_file_dependencies.py
@@ -7,8 +7,7 @@
import aiofiles
from cognee.shared.CodeGraphEntities import CodeFile, Repository
-from cognee.tasks.repo_processor.get_local_dependencies import \
- get_local_script_dependencies
+from cognee.tasks.repo_processor.get_local_dependencies import get_local_script_dependencies
async def get_py_path_and_source(file_path):
@@ -28,23 +27,27 @@ async def get_py_files_dict(repo_path):
py_files_paths = (
os.path.join(root, file)
- for root, _, files in os.walk(repo_path) for file in files if file.endswith(".py")
+ for root, _, files in os.walk(repo_path)
+ for file in files
+ if file.endswith(".py")
)
py_files_dict = {}
for file_path in py_files_paths:
absolute_path = os.path.abspath(file_path)
-
+
if os.path.getsize(absolute_path) == 0:
continue
-
+
relative_path, source_code = await get_py_path_and_source(absolute_path)
py_files_dict[relative_path] = {"source_code": source_code}
return py_files_dict
-def get_edge(file_path: str, dependency: str, repo_path: str, relative_paths: bool = False) -> tuple:
+def get_edge(
+ file_path: str, dependency: str, repo_path: str, relative_paths: bool = False
+) -> tuple:
if relative_paths:
file_path = os.path.relpath(file_path, repo_path)
dependency = os.path.relpath(dependency, repo_path)
@@ -58,22 +61,23 @@ def run_coroutine(coroutine_func, *args, **kwargs):
loop.close()
return result
+
async def get_repo_file_dependencies(repo_path: str) -> AsyncGenerator[list, None]:
"""Generate a dependency graph for Python files in the given repository path."""
-
+
if not os.path.exists(repo_path):
raise FileNotFoundError(f"Repository path {repo_path} does not exist.")
-
+
py_files_dict = await get_py_files_dict(repo_path)
repo = Repository(
- id = uuid5(NAMESPACE_OID, repo_path),
- path = repo_path,
+ id=uuid5(NAMESPACE_OID, repo_path),
+ path=repo_path,
)
yield [repo]
- with ProcessPoolExecutor(max_workers = 12) as executor:
+ with ProcessPoolExecutor(max_workers=12) as executor:
loop = asyncio.get_event_loop()
tasks = [
@@ -82,7 +86,7 @@ async def get_repo_file_dependencies(repo_path: str) -> AsyncGenerator[list, Non
run_coroutine,
get_local_script_dependencies,
os.path.join(repo_path, file_path),
- repo_path
+ repo_path,
)
for file_path, metadata in py_files_dict.items()
if metadata.get("source_code") is not None
@@ -94,19 +98,24 @@ async def get_repo_file_dependencies(repo_path: str) -> AsyncGenerator[list, Non
for (file_path, metadata), dependencies in zip(py_files_dict.items(), results):
source_code = metadata.get("source_code")
- code_files.append(CodeFile(
- id = uuid5(NAMESPACE_OID, file_path),
- source_code = source_code,
- extracted_id = file_path,
- part_of = repo,
- depends_on = [
- CodeFile(
- id = uuid5(NAMESPACE_OID, dependency),
- extracted_id = dependency,
- part_of = repo,
- source_code = py_files_dict.get(dependency, {}).get("source_code"),
- ) for dependency in dependencies
- ] if dependencies else None,
- ))
+ code_files.append(
+ CodeFile(
+ id=uuid5(NAMESPACE_OID, file_path),
+ source_code=source_code,
+ extracted_id=file_path,
+ part_of=repo,
+ depends_on=[
+ CodeFile(
+ id=uuid5(NAMESPACE_OID, dependency),
+ extracted_id=dependency,
+ part_of=repo,
+ source_code=py_files_dict.get(dependency, {}).get("source_code"),
+ )
+ for dependency in dependencies
+ ]
+ if dependencies
+ else None,
+ )
+ )
yield code_files
diff --git a/cognee/tasks/repo_processor/get_source_code_chunks.py b/cognee/tasks/repo_processor/get_source_code_chunks.py
index 4d0ce3200..5e14e11ac 100644
--- a/cognee/tasks/repo_processor/get_source_code_chunks.py
+++ b/cognee/tasks/repo_processor/get_source_code_chunks.py
@@ -16,7 +16,7 @@ def _count_tokens(tokenizer: tiktoken.Encoding, source_code: str) -> int:
def _get_naive_subchunk_token_counts(
- tokenizer: tiktoken.Encoding, source_code: str, max_subchunk_tokens: int = 8000
+ tokenizer: tiktoken.Encoding, source_code: str, max_subchunk_tokens: int = 8000
) -> list[tuple[str, int]]:
"""Splits source code into subchunks of up to max_subchunk_tokens and counts tokens."""
@@ -24,10 +24,10 @@ def _get_naive_subchunk_token_counts(
subchunk_token_counts = []
for start_idx in range(0, len(token_ids), max_subchunk_tokens):
- subchunk_token_ids = token_ids[start_idx: start_idx + max_subchunk_tokens]
+ subchunk_token_ids = token_ids[start_idx : start_idx + max_subchunk_tokens]
token_count = len(subchunk_token_ids)
- subchunk = ''.join(
- tokenizer.decode_single_token_bytes(token_id).decode('utf-8', errors='replace')
+ subchunk = "".join(
+ tokenizer.decode_single_token_bytes(token_id).decode("utf-8", errors="replace")
for token_id in subchunk_token_ids
)
subchunk_token_counts.append((subchunk, token_count))
@@ -36,11 +36,11 @@ def _get_naive_subchunk_token_counts(
def _get_subchunk_token_counts(
- tokenizer: tiktoken.Encoding,
- source_code: str,
- max_subchunk_tokens: int = 8000,
- depth: int = 0,
- max_depth: int = 100
+ tokenizer: tiktoken.Encoding,
+ source_code: str,
+ max_subchunk_tokens: int = 8000,
+ depth: int = 0,
+ max_depth: int = 100,
) -> list[tuple[str, int]]:
"""Splits source code into subchunk and counts tokens for each subchunk."""
if depth > max_depth:
@@ -72,24 +72,28 @@ def _get_subchunk_token_counts(
subchunk_token_counts.append((subchunk, token_count))
continue
- if child.type == 'string':
- subchunk_token_counts.extend(_get_naive_subchunk_token_counts(tokenizer, subchunk, max_subchunk_tokens))
+ if child.type == "string":
+ subchunk_token_counts.extend(
+ _get_naive_subchunk_token_counts(tokenizer, subchunk, max_subchunk_tokens)
+ )
continue
subchunk_token_counts.extend(
- _get_subchunk_token_counts(tokenizer, subchunk, max_subchunk_tokens, depth=depth + 1, max_depth=max_depth)
+ _get_subchunk_token_counts(
+ tokenizer, subchunk, max_subchunk_tokens, depth=depth + 1, max_depth=max_depth
+ )
)
return subchunk_token_counts
def _get_chunk_source_code(
- code_token_counts: list[tuple[str, int]], overlap: float, max_tokens: int
+ code_token_counts: list[tuple[str, int]], overlap: float, max_tokens: int
) -> tuple[list[tuple[str, int]], str]:
"""Generates a chunk of source code from tokenized subchunks with overlap handling."""
current_count = 0
cumulative_counts = []
- current_source_code = ''
+ current_source_code = ""
for i, (child_code, token_count) in enumerate(code_token_counts):
current_count += token_count
@@ -111,11 +115,11 @@ def _get_chunk_source_code(
def get_source_code_chunks_from_code_part(
- code_file_part: CodePart,
- max_tokens: int = 8192,
- overlap: float = 0.25,
- granularity: float = 0.1,
- model_name: str = "text-embedding-3-large"
+ code_file_part: CodePart,
+ max_tokens: int = 8192,
+ overlap: float = 0.25,
+ granularity: float = 0.1,
+ model_name: str = "text-embedding-3-large",
) -> Generator[SourceCodeChunk, None, None]:
"""Yields source code chunks from a CodePart object, with configurable token limits and overlap."""
if not code_file_part.source_code:
@@ -124,25 +128,30 @@ def get_source_code_chunks_from_code_part(
tokenizer = tiktoken.encoding_for_model(model_name)
max_subchunk_tokens = max(1, int(granularity * max_tokens))
- subchunk_token_counts = _get_subchunk_token_counts(tokenizer, code_file_part.source_code, max_subchunk_tokens)
+ subchunk_token_counts = _get_subchunk_token_counts(
+ tokenizer, code_file_part.source_code, max_subchunk_tokens
+ )
previous_chunk = None
while subchunk_token_counts:
- subchunk_token_counts, chunk_source_code = _get_chunk_source_code(subchunk_token_counts, overlap, max_tokens)
+ subchunk_token_counts, chunk_source_code = _get_chunk_source_code(
+ subchunk_token_counts, overlap, max_tokens
+ )
if not chunk_source_code:
continue
current_chunk = SourceCodeChunk(
id=uuid5(NAMESPACE_OID, chunk_source_code),
code_chunk_of=code_file_part,
source_code=chunk_source_code,
- previous_chunk=previous_chunk
+ previous_chunk=previous_chunk,
)
yield current_chunk
previous_chunk = current_chunk
-async def get_source_code_chunks(data_points: list[DataPoint], embedding_model="text-embedding-3-large") -> \
- AsyncGenerator[list[DataPoint], None]:
+async def get_source_code_chunks(
+ data_points: list[DataPoint], embedding_model="text-embedding-3-large"
+) -> AsyncGenerator[list[DataPoint], None]:
"""Processes code graph datapoints, create SourceCodeChink datapoints."""
# TODO: Add support for other embedding models, with max_token mapping
for data_point in data_points:
@@ -156,7 +165,9 @@ async def get_source_code_chunks(data_points: list[DataPoint], embedding_model="
for code_part in data_point.contains:
try:
yield code_part
- for source_code_chunk in get_source_code_chunks_from_code_part(code_part, model_name=embedding_model):
+ for source_code_chunk in get_source_code_chunks_from_code_part(
+ code_part, model_name=embedding_model
+ ):
yield source_code_chunk
except Exception as e:
logger.error(f"Error processing code part: {e}")
diff --git a/cognee/tasks/repo_processor/top_down_repo_parse.py b/cognee/tasks/repo_processor/top_down_repo_parse.py
index 52f58f811..aed971920 100644
--- a/cognee/tasks/repo_processor/top_down_repo_parse.py
+++ b/cognee/tasks/repo_processor/top_down_repo_parse.py
@@ -14,6 +14,7 @@
"simple_stmt": "var_def",
}
+
def _create_object_dict(name_node, type_name=None):
return {
"name": name_node.value,
@@ -32,7 +33,11 @@ def _parse_node(node):
if node.type == "async_stmt" and len(node.children) > 1:
function_node = node.children[1]
if function_node.type == "funcdef":
- return [_create_object_dict(function_node.name, type_name=_NODE_TYPE_MAP.get(function_node.type))]
+ return [
+ _create_object_dict(
+ function_node.name, type_name=_NODE_TYPE_MAP.get(function_node.type)
+ )
+ ]
if node.type == "simple_stmt":
# TODO: Handle multi-level/nested unpacking variable definitions in the future
expr_child = node.children[0]
@@ -50,7 +55,6 @@ def _parse_node(node):
return []
-
def extract_importable_objects_with_positions_from_source_code(source_code):
"""Extract top-level objects in a Python source code string with their positions (line/column)."""
try:
@@ -82,14 +86,12 @@ def extract_importable_objects_with_positions(file_path):
return extract_importable_objects_with_positions_from_source_code(source_code)
-
def find_entity_usages(script, line, column):
"""
Return a list of files in the repo where the entity at module_path:line,column is used.
"""
usages = set()
-
try:
inferred = script.infer(line, column)
except Exception as e:
@@ -103,9 +105,13 @@ def find_entity_usages(script, line, column):
logger.debug(f"Inferred entity: {inferred[0].name}, type: {inferred[0].type}")
try:
- references = script.get_references(line=line, column=column, scope="project", include_builtins=False)
+ references = script.get_references(
+ line=line, column=column, scope="project", include_builtins=False
+ )
except Exception as e:
- logger.error(f"Error retrieving references for entity at {script.path}:{line},{column}: {e}")
+ logger.error(
+ f"Error retrieving references for entity at {script.path}:{line},{column}: {e}"
+ )
references = []
for ref in references:
@@ -115,6 +121,7 @@ def find_entity_usages(script, line, column):
return list(usages)
+
def parse_file_with_references(project, file_path):
"""Parse a file to extract object names and their references within a project."""
try:
@@ -152,7 +159,7 @@ def parse_repo(repo_path):
logger.error(f"Error creating Jedi project for repository at {repo_path}: {e}")
return {}
- EXCLUDE_DIRS = {'venv', '.git', '__pycache__', 'build'}
+ EXCLUDE_DIRS = {"venv", ".git", "__pycache__", "build"}
python_files = [
os.path.join(directory, file)
@@ -168,4 +175,3 @@ def parse_repo(repo_path):
}
return results
-
diff --git a/cognee/tasks/storage/add_data_points.py b/cognee/tasks/storage/add_data_points.py
index 47cae4309..21cc5a3c2 100644
--- a/cognee/tasks/storage/add_data_points.py
+++ b/cognee/tasks/storage/add_data_points.py
@@ -5,7 +5,7 @@
from .index_data_points import index_data_points
-async def add_data_points(data_points: list[DataPoint], only_root = False):
+async def add_data_points(data_points: list[DataPoint], only_root=False):
nodes = []
edges = []
@@ -13,15 +13,18 @@ async def add_data_points(data_points: list[DataPoint], only_root = False):
added_edges = {}
visited_properties = {}
- results = await asyncio.gather(*[
- get_graph_from_model(
- data_point,
- added_nodes = added_nodes,
- added_edges = added_edges,
- visited_properties = visited_properties,
- only_root = only_root,
- ) for data_point in data_points
- ])
+ results = await asyncio.gather(
+ *[
+ get_graph_from_model(
+ data_point,
+ added_nodes=added_nodes,
+ added_edges=added_edges,
+ visited_properties=visited_properties,
+ only_root=only_root,
+ )
+ for data_point in data_points
+ ]
+ )
for result_nodes, result_edges in results:
nodes.extend(result_nodes)
diff --git a/cognee/tasks/storage/index_data_points.py b/cognee/tasks/storage/index_data_points.py
index 12af2d2ef..65f516bac 100644
--- a/cognee/tasks/storage/index_data_points.py
+++ b/cognee/tasks/storage/index_data_points.py
@@ -6,6 +6,7 @@
logger = logging.getLogger("index_data_points")
+
async def index_data_points(data_points: list[DataPoint]):
created_indexes = {}
index_points = {}
@@ -41,7 +42,10 @@ async def index_data_points(data_points: list[DataPoint]):
return data_points
-async def get_data_points_from_model(data_point: DataPoint, added_data_points = None, visited_properties = None) -> list[DataPoint]:
+
+async def get_data_points_from_model(
+ data_point: DataPoint, added_data_points=None, visited_properties=None
+) -> list[DataPoint]:
data_points = []
added_data_points = added_data_points or {}
visited_properties = visited_properties or {}
@@ -55,14 +59,20 @@ async def get_data_points_from_model(data_point: DataPoint, added_data_points =
visited_properties[property_key] = True
- new_data_points = await get_data_points_from_model(field_value, added_data_points, visited_properties)
+ new_data_points = await get_data_points_from_model(
+ field_value, added_data_points, visited_properties
+ )
for new_point in new_data_points:
if str(new_point.id) not in added_data_points:
added_data_points[str(new_point.id)] = True
data_points.append(new_point)
- if isinstance(field_value, list) and len(field_value) > 0 and isinstance(field_value[0], DataPoint):
+ if (
+ isinstance(field_value, list)
+ and len(field_value) > 0
+ and isinstance(field_value[0], DataPoint)
+ ):
for field_value_item in field_value:
property_key = f"{str(data_point.id)}{field_name}{str(field_value_item.id)}"
@@ -70,43 +80,39 @@ async def get_data_points_from_model(data_point: DataPoint, added_data_points =
return []
visited_properties[property_key] = True
-
- new_data_points = await get_data_points_from_model(field_value_item, added_data_points, visited_properties)
+
+ new_data_points = await get_data_points_from_model(
+ field_value_item, added_data_points, visited_properties
+ )
for new_point in new_data_points:
if str(new_point.id) not in added_data_points:
added_data_points[str(new_point.id)] = True
data_points.append(new_point)
- if (str(data_point.id) not in added_data_points):
+ if str(data_point.id) not in added_data_points:
data_points.append(data_point)
return data_points
if __name__ == "__main__":
+
class Car(DataPoint):
model: str
color: str
- _metadata = {
- "index_fields": ["name"],
- "type": "Car"
- }
+ _metadata = {"index_fields": ["name"], "type": "Car"}
-
class Person(DataPoint):
name: str
age: int
owns_car: list[Car]
- _metadata = {
- "index_fields": ["name"],
- "type": "Person"
- }
+ _metadata = {"index_fields": ["name"], "type": "Person"}
- car1 = Car(model = "Tesla Model S", color = "Blue")
- car2 = Car(model = "Toyota Camry", color = "Red")
- person = Person(name = "John", age = 30, owns_car = [car1, car2])
+ car1 = Car(model="Tesla Model S", color="Blue")
+ car2 = Car(model="Toyota Camry", color="Red")
+ person = Person(name="John", age=30, owns_car=[car1, car2])
data_points = get_data_points_from_model(person)
- print(data_points)
\ No newline at end of file
+ print(data_points)
diff --git a/cognee/tasks/storage/index_graph_edges.py b/cognee/tasks/storage/index_graph_edges.py
index 88514af48..2aeb2bef2 100644
--- a/cognee/tasks/storage/index_graph_edges.py
+++ b/cognee/tasks/storage/index_graph_edges.py
@@ -8,25 +8,25 @@
async def index_graph_edges():
"""
- Indexes graph edges by creating and managing vector indexes for relationship types.
+ Indexes graph edges by creating and managing vector indexes for relationship types.
- This function retrieves edge data from the graph engine, counts distinct relationship
- types, and creates `EdgeType` pydantic objects. It ensures that vector indexes are created for
- the `relationship_name` field.
+ This function retrieves edge data from the graph engine, counts distinct relationship
+ types, and creates `EdgeType` pydantic objects. It ensures that vector indexes are created for
+ the `relationship_name` field.
- Steps:
- 1. Initialize the vector engine and graph engine.
- 2. Retrieve graph edge data and count relationship types (`relationship_name`).
- 3. Create vector indexes for `relationship_name` if they don't exist.
- 4. Transform the counted relationships into `EdgeType` objects.
- 5. Index the transformed data points in the vector engine.
+ Steps:
+ 1. Initialize the vector engine and graph engine.
+ 2. Retrieve graph edge data and count relationship types (`relationship_name`).
+ 3. Create vector indexes for `relationship_name` if they don't exist.
+ 4. Transform the counted relationships into `EdgeType` objects.
+ 5. Index the transformed data points in the vector engine.
- Raises:
- RuntimeError: If initialization of the vector engine or graph engine fails.
+ Raises:
+ RuntimeError: If initialization of the vector engine or graph engine fails.
- Returns:
- None
- """
+ Returns:
+ None
+ """
try:
created_indexes = {}
index_points = {}
@@ -40,9 +40,10 @@ async def index_graph_edges():
_, edges_data = await graph_engine.get_graph_data()
edge_types = Counter(
- item.get('relationship_name')
+ item.get("relationship_name")
for edge in edges_data
- for item in edge if isinstance(item, dict) and 'relationship_name' in item
+ for item in edge
+ if isinstance(item, dict) and "relationship_name" in item
)
for text, count in edge_types.items():
@@ -67,4 +68,4 @@ async def index_graph_edges():
index_name, field_name = index_name.split(".")
await vector_engine.index_data_points(index_name, field_name, indexable_points)
- return None
\ No newline at end of file
+ return None
diff --git a/cognee/tasks/summarization/mock_summary.py b/cognee/tasks/summarization/mock_summary.py
index e61385bb1..f60ce2d82 100644
--- a/cognee/tasks/summarization/mock_summary.py
+++ b/cognee/tasks/summarization/mock_summary.py
@@ -1,5 +1,6 @@
from cognee.shared.data_models import SummarizedCode, SummarizedClass, SummarizedFunction
+
def get_mock_summarized_code() -> SummarizedCode:
return SummarizedCode(
file_name="mock_file.py",
@@ -34,4 +35,4 @@ def get_mock_summarized_code() -> SummarizedCode:
)
],
workflow_description="This is a mock workflow description.",
- )
\ No newline at end of file
+ )
diff --git a/cognee/tasks/summarization/models.py b/cognee/tasks/summarization/models.py
index 5b0345015..fc62209ce 100644
--- a/cognee/tasks/summarization/models.py
+++ b/cognee/tasks/summarization/models.py
@@ -10,10 +10,7 @@ class TextSummary(DataPoint):
text: str
made_from: DocumentChunk
- _metadata: dict = {
- "index_fields": ["text"],
- "type": "TextSummary"
- }
+ _metadata: dict = {"index_fields": ["text"], "type": "TextSummary"}
class CodeSummary(DataPoint):
@@ -21,7 +18,4 @@ class CodeSummary(DataPoint):
text: str
summarizes: Union[CodeFile, CodePart, SourceCodeChunk]
- _metadata: dict = {
- "index_fields": ["text"],
- "type": "CodeSummary"
- }
+ _metadata: dict = {"index_fields": ["text"], "type": "CodeSummary"}
diff --git a/cognee/tasks/summarization/query_summaries.py b/cognee/tasks/summarization/query_summaries.py
index d9ec0fa00..5f04e3eb4 100644
--- a/cognee/tasks/summarization/query_summaries.py
+++ b/cognee/tasks/summarization/query_summaries.py
@@ -1,5 +1,6 @@
from cognee.infrastructure.databases.vector import get_vector_engine
+
async def query_summaries(query: str) -> list:
"""
Parameters:
@@ -10,7 +11,7 @@ async def query_summaries(query: str) -> list:
"""
vector_engine = get_vector_engine()
- summaries_results = await vector_engine.search("text_summary_text", query, limit = 5)
+ summaries_results = await vector_engine.search("text_summary_text", query, limit=5)
summaries = [summary.payload for summary in summaries_results]
diff --git a/cognee/tasks/summarization/summarize_text.py b/cognee/tasks/summarization/summarize_text.py
index 60ec66f5c..df7ac6740 100644
--- a/cognee/tasks/summarization/summarize_text.py
+++ b/cognee/tasks/summarization/summarize_text.py
@@ -6,6 +6,7 @@
from cognee.modules.chunking.models.DocumentChunk import DocumentChunk
from .models import TextSummary
+
async def summarize_text(data_chunks: list[DocumentChunk], summarization_model: Type[BaseModel]):
if len(data_chunks) == 0:
return data_chunks
@@ -16,10 +17,11 @@ async def summarize_text(data_chunks: list[DocumentChunk], summarization_model:
summaries = [
TextSummary(
- id = uuid5(chunk.id, "TextSummary"),
- made_from = chunk,
- text = chunk_summaries[chunk_index].summary,
- ) for (chunk_index, chunk) in enumerate(data_chunks)
+ id=uuid5(chunk.id, "TextSummary"),
+ made_from=chunk,
+ text=chunk_summaries[chunk_index].summary,
+ )
+ for (chunk_index, chunk) in enumerate(data_chunks)
]
return summaries
diff --git a/cognee/tasks/temporal_awareness/__init__.py b/cognee/tasks/temporal_awareness/__init__.py
index 9b6717f5f..4a47d33be 100644
--- a/cognee/tasks/temporal_awareness/__init__.py
+++ b/cognee/tasks/temporal_awareness/__init__.py
@@ -1,4 +1,2 @@
-from .build_graph_with_temporal_awareness import \
- build_graph_with_temporal_awareness
-from .search_graph_with_temporal_awareness import \
- search_graph_with_temporal_awareness
+from .build_graph_with_temporal_awareness import build_graph_with_temporal_awareness
+from .search_graph_with_temporal_awareness import search_graph_with_temporal_awareness
diff --git a/cognee/tasks/temporal_awareness/build_graph_with_temporal_awareness.py b/cognee/tasks/temporal_awareness/build_graph_with_temporal_awareness.py
index ac534623b..ecbf2b6be 100644
--- a/cognee/tasks/temporal_awareness/build_graph_with_temporal_awareness.py
+++ b/cognee/tasks/temporal_awareness/build_graph_with_temporal_awareness.py
@@ -6,12 +6,11 @@
async def build_graph_with_temporal_awareness(text_list):
-
url = os.getenv("GRAPH_DATABASE_URL")
password = os.getenv("GRAPH_DATABASE_PASSWORD")
graphiti = Graphiti(url, "neo4j", password)
-
- await graphiti.build_indices_and_constraints()
+
+ await graphiti.build_indices_and_constraints()
print("Graph database initialized.")
for i, text in enumerate(text_list):
@@ -20,7 +19,7 @@ async def build_graph_with_temporal_awareness(text_list):
episode_body=text,
source=EpisodeType.text,
source_description="input",
- reference_time=datetime.now()
+ reference_time=datetime.now(),
)
print(f"Added text: {text[:35]}...")
- return graphiti
\ No newline at end of file
+ return graphiti
diff --git a/cognee/tasks/temporal_awareness/search_graph_with_temporal_awareness.py b/cognee/tasks/temporal_awareness/search_graph_with_temporal_awareness.py
index 8e2cd0fc8..247b04e71 100644
--- a/cognee/tasks/temporal_awareness/search_graph_with_temporal_awareness.py
+++ b/cognee/tasks/temporal_awareness/search_graph_with_temporal_awareness.py
@@ -1,6 +1,4 @@
-
-
async def search_graph_with_temporal_awareness(graphiti, query):
search_result = await graphiti.search(query)
await graphiti.close()
- return search_result
\ No newline at end of file
+ return search_result
diff --git a/cognee/tests/get_graph_url.py b/cognee/tests/get_graph_url.py
index c5909676d..1e0abbaf0 100644
--- a/cognee/tests/get_graph_url.py
+++ b/cognee/tests/get_graph_url.py
@@ -3,14 +3,23 @@
from cognee.infrastructure.databases.graph import get_graph_engine
if __name__ == "__main__":
+
async def main():
import os
import pathlib
import cognee
- data_directory_path = str(pathlib.Path(os.path.join(pathlib.Path(__file__).parent, ".data_storage/test_library")).resolve())
+ data_directory_path = str(
+ pathlib.Path(
+ os.path.join(pathlib.Path(__file__).parent, ".data_storage/test_library")
+ ).resolve()
+ )
cognee.config.data_root_directory(data_directory_path)
- cognee_directory_path = str(pathlib.Path(os.path.join(pathlib.Path(__file__).parent, ".cognee_system/test_library")).resolve())
+ cognee_directory_path = str(
+ pathlib.Path(
+ os.path.join(pathlib.Path(__file__).parent, ".cognee_system/test_library")
+ ).resolve()
+ )
cognee.config.system_root_directory(cognee_directory_path)
graph_client = await get_graph_engine()
diff --git a/cognee/tests/integration/documents/AudioDocument_test.py b/cognee/tests/integration/documents/AudioDocument_test.py
index 151f4c0b2..dbd43ddda 100644
--- a/cognee/tests/integration/documents/AudioDocument_test.py
+++ b/cognee/tests/integration/documents/AudioDocument_test.py
@@ -25,13 +25,16 @@
def test_AudioDocument():
-
document = AudioDocument(
- id=uuid.uuid4(), name="audio-dummy-test", raw_data_location="", metadata_id=uuid.uuid4(), mime_type="",
+ id=uuid.uuid4(),
+ name="audio-dummy-test",
+ raw_data_location="",
+ metadata_id=uuid.uuid4(),
+ mime_type="",
)
with patch.object(AudioDocument, "create_transcript", return_value=TEST_TEXT):
for ground_truth, paragraph_data in zip(
- GROUND_TRUTH, document.read(chunk_size=64, chunker='text_chunker')
+ GROUND_TRUTH, document.read(chunk_size=64, chunker="text_chunker")
):
assert (
ground_truth["word_count"] == paragraph_data.word_count
diff --git a/cognee/tests/integration/documents/ImageDocument_test.py b/cognee/tests/integration/documents/ImageDocument_test.py
index 40e0155af..c0877ae99 100644
--- a/cognee/tests/integration/documents/ImageDocument_test.py
+++ b/cognee/tests/integration/documents/ImageDocument_test.py
@@ -14,14 +14,16 @@
def test_ImageDocument():
-
document = ImageDocument(
- id=uuid.uuid4(), name="image-dummy-test", raw_data_location="", metadata_id=uuid.uuid4(), mime_type="",
+ id=uuid.uuid4(),
+ name="image-dummy-test",
+ raw_data_location="",
+ metadata_id=uuid.uuid4(),
+ mime_type="",
)
with patch.object(ImageDocument, "transcribe_image", return_value=TEST_TEXT):
-
for ground_truth, paragraph_data in zip(
- GROUND_TRUTH, document.read(chunk_size=64, chunker='text_chunker')
+ GROUND_TRUTH, document.read(chunk_size=64, chunker="text_chunker")
):
assert (
ground_truth["word_count"] == paragraph_data.word_count
diff --git a/cognee/tests/integration/documents/PdfDocument_test.py b/cognee/tests/integration/documents/PdfDocument_test.py
index 25d4cf6c6..8f28815d3 100644
--- a/cognee/tests/integration/documents/PdfDocument_test.py
+++ b/cognee/tests/integration/documents/PdfDocument_test.py
@@ -17,12 +17,15 @@ def test_PdfDocument():
"artificial-intelligence.pdf",
)
document = PdfDocument(
- id=uuid.uuid4(), name="Test document.pdf", raw_data_location=test_file_path, metadata_id=uuid.uuid4(),
+ id=uuid.uuid4(),
+ name="Test document.pdf",
+ raw_data_location=test_file_path,
+ metadata_id=uuid.uuid4(),
mime_type="",
)
for ground_truth, paragraph_data in zip(
- GROUND_TRUTH, document.read(chunk_size=1024, chunker='text_chunker')
+ GROUND_TRUTH, document.read(chunk_size=1024, chunker="text_chunker")
):
assert (
ground_truth["word_count"] == paragraph_data.word_count
diff --git a/cognee/tests/integration/documents/TextDocument_test.py b/cognee/tests/integration/documents/TextDocument_test.py
index 91f38968e..1e143d563 100644
--- a/cognee/tests/integration/documents/TextDocument_test.py
+++ b/cognee/tests/integration/documents/TextDocument_test.py
@@ -29,11 +29,15 @@ def test_TextDocument(input_file, chunk_size):
input_file,
)
document = TextDocument(
- id=uuid.uuid4(), name=input_file, raw_data_location=test_file_path, metadata_id=uuid.uuid4(), mime_type="",
+ id=uuid.uuid4(),
+ name=input_file,
+ raw_data_location=test_file_path,
+ metadata_id=uuid.uuid4(),
+ mime_type="",
)
for ground_truth, paragraph_data in zip(
- GROUND_TRUTH[input_file], document.read(chunk_size=chunk_size, chunker='text_chunker')
+ GROUND_TRUTH[input_file], document.read(chunk_size=chunk_size, chunker="text_chunker")
):
assert (
ground_truth["word_count"] == paragraph_data.word_count
diff --git a/cognee/tests/integration/documents/UnstructuredDocument_test.py b/cognee/tests/integration/documents/UnstructuredDocument_test.py
index 418b18810..03b8deb49 100644
--- a/cognee/tests/integration/documents/UnstructuredDocument_test.py
+++ b/cognee/tests/integration/documents/UnstructuredDocument_test.py
@@ -3,6 +3,7 @@
from cognee.modules.data.processing.document_types.UnstructuredDocument import UnstructuredDocument
+
def test_UnstructuredDocument():
# Define file paths of test data
pptx_file_path = os.path.join(
@@ -35,46 +36,67 @@ def test_UnstructuredDocument():
# Define test documents
pptx_document = UnstructuredDocument(
- id=uuid.uuid4(), name="example.pptx", raw_data_location=pptx_file_path, metadata_id=uuid.uuid4(),
- mime_type="application/vnd.openxmlformats-officedocument.presentationml.presentation"
+ id=uuid.uuid4(),
+ name="example.pptx",
+ raw_data_location=pptx_file_path,
+ metadata_id=uuid.uuid4(),
+ mime_type="application/vnd.openxmlformats-officedocument.presentationml.presentation",
)
docx_document = UnstructuredDocument(
- id=uuid.uuid4(), name="example.docx", raw_data_location=docx_file_path, metadata_id=uuid.uuid4(),
- mime_type="application/vnd.openxmlformats-officedocument.wordprocessingml.document"
+ id=uuid.uuid4(),
+ name="example.docx",
+ raw_data_location=docx_file_path,
+ metadata_id=uuid.uuid4(),
+ mime_type="application/vnd.openxmlformats-officedocument.wordprocessingml.document",
)
csv_document = UnstructuredDocument(
- id=uuid.uuid4(), name="example.csv", raw_data_location=csv_file_path, metadata_id=uuid.uuid4(),
- mime_type="text/csv"
+ id=uuid.uuid4(),
+ name="example.csv",
+ raw_data_location=csv_file_path,
+ metadata_id=uuid.uuid4(),
+ mime_type="text/csv",
)
xlsx_document = UnstructuredDocument(
- id=uuid.uuid4(), name="example.xlsx", raw_data_location=xlsx_file_path, metadata_id=uuid.uuid4(),
- mime_type="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"
+ id=uuid.uuid4(),
+ name="example.xlsx",
+ raw_data_location=xlsx_file_path,
+ metadata_id=uuid.uuid4(),
+ mime_type="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
)
# Test PPTX
for paragraph_data in pptx_document.read(chunk_size=1024):
- assert 19 == paragraph_data.word_count, f' 19 != {paragraph_data.word_count = }'
- assert 104 == len(paragraph_data.text), f' 104 != {len(paragraph_data.text) = }'
- assert 'sentence_cut' == paragraph_data.cut_type, f' sentence_cut != {paragraph_data.cut_type = }'
+ assert 19 == paragraph_data.word_count, f" 19 != {paragraph_data.word_count = }"
+ assert 104 == len(paragraph_data.text), f" 104 != {len(paragraph_data.text) = }"
+ assert (
+ "sentence_cut" == paragraph_data.cut_type
+ ), f" sentence_cut != {paragraph_data.cut_type = }"
# Test DOCX
for paragraph_data in docx_document.read(chunk_size=1024):
- assert 16 == paragraph_data.word_count, f' 16 != {paragraph_data.word_count = }'
- assert 145 == len(paragraph_data.text), f' 145 != {len(paragraph_data.text) = }'
- assert 'sentence_end' == paragraph_data.cut_type, f' sentence_end != {paragraph_data.cut_type = }'
+ assert 16 == paragraph_data.word_count, f" 16 != {paragraph_data.word_count = }"
+ assert 145 == len(paragraph_data.text), f" 145 != {len(paragraph_data.text) = }"
+ assert (
+ "sentence_end" == paragraph_data.cut_type
+ ), f" sentence_end != {paragraph_data.cut_type = }"
# TEST CSV
for paragraph_data in csv_document.read(chunk_size=1024):
- assert 15 == paragraph_data.word_count, f' 15 != {paragraph_data.word_count = }'
- assert 'A A A A A A A A A,A A A A A A,A A' == paragraph_data.text, \
- f'Read text doesn\'t match expected text: {paragraph_data.text}'
- assert 'sentence_cut' == paragraph_data.cut_type, f' sentence_cut != {paragraph_data.cut_type = }'
+ assert 15 == paragraph_data.word_count, f" 15 != {paragraph_data.word_count = }"
+ assert (
+ "A A A A A A A A A,A A A A A A,A A" == paragraph_data.text
+ ), f"Read text doesn't match expected text: {paragraph_data.text}"
+ assert (
+ "sentence_cut" == paragraph_data.cut_type
+ ), f" sentence_cut != {paragraph_data.cut_type = }"
# Test XLSX
for paragraph_data in xlsx_document.read(chunk_size=1024):
- assert 36 == paragraph_data.word_count, f' 36 != {paragraph_data.word_count = }'
- assert 171 == len(paragraph_data.text), f' 171 != {len(paragraph_data.text) = }'
- assert 'sentence_cut' == paragraph_data.cut_type, f' sentence_cut != {paragraph_data.cut_type = }'
+ assert 36 == paragraph_data.word_count, f" 36 != {paragraph_data.word_count = }"
+ assert 171 == len(paragraph_data.text), f" 171 != {len(paragraph_data.text) = }"
+ assert (
+ "sentence_cut" == paragraph_data.cut_type
+ ), f" sentence_cut != {paragraph_data.cut_type = }"
diff --git a/cognee/tests/integration/run_toy_tasks/run_task_from_queue_test.py b/cognee/tests/integration/run_toy_tasks/run_task_from_queue_test.py
index e57b16f39..6a4024342 100644
--- a/cognee/tests/integration/run_toy_tasks/run_task_from_queue_test.py
+++ b/cognee/tests/integration/run_toy_tasks/run_task_from_queue_test.py
@@ -31,9 +31,7 @@ async def multiply_by_two(num):
results = [2, 4, 6, 8, 10, 12, 14, 16, 18, 20]
index = 0
async for result in tasks_run:
- assert (
- result == results[index]
- ), f"at {index = }: {result = } != {results[index] = }"
+ assert result == results[index], f"at {index = }: {result = } != {results[index] = }"
index += 1
diff --git a/cognee/tests/integration/run_toy_tasks/run_tasks_test.py b/cognee/tests/integration/run_toy_tasks/run_tasks_test.py
index d0a2af80b..54613c214 100644
--- a/cognee/tests/integration/run_toy_tasks/run_tasks_test.py
+++ b/cognee/tests/integration/run_toy_tasks/run_tasks_test.py
@@ -33,9 +33,7 @@ async def add_one_single(num):
results = [5, 7, 9, 11, 13, 15, 17, 19, 21, 23]
index = 0
async for result in pipeline:
- assert (
- result == results[index]
- ), f"at {index = }: {result = } != {results[index] = }"
+ assert result == results[index], f"at {index = }: {result = } != {results[index] = }"
index += 1
diff --git a/cognee/tests/test_deduplication.py b/cognee/tests/test_deduplication.py
index 467a52368..9c2df032d 100644
--- a/cognee/tests/test_deduplication.py
+++ b/cognee/tests/test_deduplication.py
@@ -8,6 +8,7 @@
logging.basicConfig(level=logging.DEBUG)
+
async def test_deduplication():
await cognee.prune.prune_data()
await cognee.prune.prune_system(metadata=True)
@@ -29,7 +30,9 @@ async def test_deduplication():
result = await relational_engine.get_all_data_from_table("data")
assert len(result) == 1, "More than one data entity was found."
- assert result[0]["name"] == "Natural_language_processing_copy", "Result name does not match expected value."
+ assert (
+ result[0]["name"] == "Natural_language_processing_copy"
+ ), "Result name does not match expected value."
result = await relational_engine.get_all_data_from_table("datasets")
assert len(result) == 2, "Unexpected number of datasets found."
@@ -58,15 +61,15 @@ async def test_deduplication():
result = await relational_engine.get_all_data_from_table("data")
assert len(result) == 1, "More than one data entity was found."
- assert hashlib.md5(text.encode('utf-8')).hexdigest() in result[0]["name"], "Content hash is not a part of file name."
+ assert (
+ hashlib.md5(text.encode("utf-8")).hexdigest() in result[0]["name"]
+ ), "Content hash is not a part of file name."
await cognee.prune.prune_data()
await cognee.prune.prune_system(metadata=True)
# Test deduplication of image files
- explanation_file_path = os.path.join(
- pathlib.Path(__file__).parent, "test_data/example.png"
- )
+ explanation_file_path = os.path.join(pathlib.Path(__file__).parent, "test_data/example.png")
explanation_file_path2 = os.path.join(
pathlib.Path(__file__).parent, "test_data/example_copy.png"
)
@@ -100,11 +103,7 @@ async def test_deduplication():
async def test_deduplication_postgres():
cognee.config.set_vector_db_config(
- {
- "vector_db_url": "",
- "vector_db_key": "",
- "vector_db_provider": "pgvector"
- }
+ {"vector_db_url": "", "vector_db_key": "", "vector_db_provider": "pgvector"}
)
cognee.config.set_relational_db_config(
{
@@ -119,13 +118,10 @@ async def test_deduplication_postgres():
await test_deduplication()
+
async def test_deduplication_sqlite():
cognee.config.set_vector_db_config(
- {
- "vector_db_url": "",
- "vector_db_key": "",
- "vector_db_provider": "lancedb"
- }
+ {"vector_db_url": "", "vector_db_key": "", "vector_db_provider": "lancedb"}
)
cognee.config.set_relational_db_config(
{
@@ -137,7 +133,6 @@ async def test_deduplication_sqlite():
async def main():
-
data_directory_path = str(
pathlib.Path(
os.path.join(pathlib.Path(__file__).parent, ".data_storage/test_deduplication")
@@ -154,6 +149,7 @@ async def main():
await test_deduplication_postgres()
await test_deduplication_sqlite()
+
if __name__ == "__main__":
import asyncio
diff --git a/cognee/tests/test_falkordb.py b/cognee/tests/test_falkordb.py
index 25fe81a75..07ece9eb2 100755
--- a/cognee/tests/test_falkordb.py
+++ b/cognee/tests/test_falkordb.py
@@ -5,20 +5,31 @@
from cognee.api.v1.search import SearchType
from cognee.shared.utils import render_graph
-logging.basicConfig(level = logging.DEBUG)
+logging.basicConfig(level=logging.DEBUG)
-async def main():
- data_directory_path = str(pathlib.Path(os.path.join(pathlib.Path(__file__).parent, ".data_storage/test_falkordb")).resolve())
+
+async def main():
+ data_directory_path = str(
+ pathlib.Path(
+ os.path.join(pathlib.Path(__file__).parent, ".data_storage/test_falkordb")
+ ).resolve()
+ )
cognee.config.data_root_directory(data_directory_path)
- cognee_directory_path = str(pathlib.Path(os.path.join(pathlib.Path(__file__).parent, ".cognee_system/test_falkordb")).resolve())
+ cognee_directory_path = str(
+ pathlib.Path(
+ os.path.join(pathlib.Path(__file__).parent, ".cognee_system/test_falkordb")
+ ).resolve()
+ )
cognee.config.system_root_directory(cognee_directory_path)
await cognee.prune.prune_data()
- await cognee.prune.prune_system(metadata = True)
+ await cognee.prune.prune_system(metadata=True)
dataset_name = "artificial_intelligence"
- ai_text_file_path = os.path.join(pathlib.Path(__file__).parent, "test_data/artificial-intelligence.pdf")
+ ai_text_file_path = os.path.join(
+ pathlib.Path(__file__).parent, "test_data/artificial-intelligence.pdf"
+ )
await cognee.add([ai_text_file_path], dataset_name)
text = """A large language model (LLM) is a language model notable for its ability to achieve general-purpose language generation and other natural language processing tasks such as classification. LLMs acquire these abilities by learning statistical relationships from text documents during a computationally intensive self-supervised and semi-supervised training process. LLMs can be used for text generation, a form of generative AI, by taking an input text and repeatedly predicting the next token or word.
@@ -34,23 +45,24 @@ async def main():
# await render_graph(None, include_labels = True, include_nodes = True)
from cognee.infrastructure.databases.vector import get_vector_engine
+
vector_engine = get_vector_engine()
random_node = (await vector_engine.search("entity.name", "AI"))[0]
random_node_name = random_node.payload["text"]
- search_results = await cognee.search(SearchType.INSIGHTS, query_text = random_node_name)
+ search_results = await cognee.search(SearchType.INSIGHTS, query_text=random_node_name)
assert len(search_results) != 0, "The search results list is empty."
print("\n\nExtracted sentences are:\n")
for result in search_results:
print(f"{result}\n")
- search_results = await cognee.search(SearchType.CHUNKS, query_text = random_node_name)
+ search_results = await cognee.search(SearchType.CHUNKS, query_text=random_node_name)
assert len(search_results) != 0, "The search results list is empty."
print("\n\nExtracted chunks are:\n")
for result in search_results:
print(f"{result}\n")
- search_results = await cognee.search(SearchType.SUMMARIES, query_text = random_node_name)
+ search_results = await cognee.search(SearchType.SUMMARIES, query_text=random_node_name)
assert len(search_results) != 0, "Query related summaries don't exist."
print("\nExtracted summaries are:\n")
for result in search_results:
@@ -72,12 +84,18 @@ async def main():
assert len(collection_names) == 0, "LanceDB vector database is not empty"
from cognee.infrastructure.databases.relational import get_relational_engine
- assert not os.path.exists(get_relational_engine().db_path), "SQLite relational database is not empty"
+
+ assert not os.path.exists(
+ get_relational_engine().db_path
+ ), "SQLite relational database is not empty"
from cognee.infrastructure.databases.graph import get_graph_config
+
graph_config = get_graph_config()
assert not os.path.exists(graph_config.graph_file_path), "Networkx graph database is not empty"
+
if __name__ == "__main__":
import asyncio
+
asyncio.run(main(), debug=True)
diff --git a/cognee/tests/test_library.py b/cognee/tests/test_library.py
index 6c9d41800..8352b4161 100755
--- a/cognee/tests/test_library.py
+++ b/cognee/tests/test_library.py
@@ -4,20 +4,31 @@
import cognee
from cognee.api.v1.search import SearchType
-logging.basicConfig(level = logging.DEBUG)
+logging.basicConfig(level=logging.DEBUG)
-async def main():
- data_directory_path = str(pathlib.Path(os.path.join(pathlib.Path(__file__).parent, ".data_storage/test_library")).resolve())
+
+async def main():
+ data_directory_path = str(
+ pathlib.Path(
+ os.path.join(pathlib.Path(__file__).parent, ".data_storage/test_library")
+ ).resolve()
+ )
cognee.config.data_root_directory(data_directory_path)
- cognee_directory_path = str(pathlib.Path(os.path.join(pathlib.Path(__file__).parent, ".cognee_system/test_library")).resolve())
+ cognee_directory_path = str(
+ pathlib.Path(
+ os.path.join(pathlib.Path(__file__).parent, ".cognee_system/test_library")
+ ).resolve()
+ )
cognee.config.system_root_directory(cognee_directory_path)
await cognee.prune.prune_data()
- await cognee.prune.prune_system(metadata = True)
+ await cognee.prune.prune_system(metadata=True)
dataset_name = "artificial_intelligence"
- ai_text_file_path = os.path.join(pathlib.Path(__file__).parent, "test_data/artificial-intelligence.pdf")
+ ai_text_file_path = os.path.join(
+ pathlib.Path(__file__).parent, "test_data/artificial-intelligence.pdf"
+ )
await cognee.add([ai_text_file_path], dataset_name)
text = """A large language model (LLM) is a language model notable for its ability to achieve general-purpose language generation and other natural language processing tasks such as classification. LLMs acquire these abilities by learning statistical relationships from text documents during a computationally intensive self-supervised and semi-supervised training process. LLMs can be used for text generation, a form of generative AI, by taking an input text and repeatedly predicting the next token or word.
@@ -31,23 +42,24 @@ async def main():
await cognee.cognify([dataset_name])
from cognee.infrastructure.databases.vector import get_vector_engine
+
vector_engine = get_vector_engine()
random_node = (await vector_engine.search("entity_name", "AI"))[0]
random_node_name = random_node.payload["text"]
- search_results = await cognee.search(SearchType.INSIGHTS, query_text = random_node_name)
+ search_results = await cognee.search(SearchType.INSIGHTS, query_text=random_node_name)
assert len(search_results) != 0, "The search results list is empty."
print("\n\nExtracted sentences are:\n")
for result in search_results:
print(f"{result}\n")
- search_results = await cognee.search(SearchType.CHUNKS, query_text = random_node_name)
+ search_results = await cognee.search(SearchType.CHUNKS, query_text=random_node_name)
assert len(search_results) != 0, "The search results list is empty."
print("\n\nExtracted chunks are:\n")
for result in search_results:
print(f"{result}\n")
- search_results = await cognee.search(SearchType.SUMMARIES, query_text = random_node_name)
+ search_results = await cognee.search(SearchType.SUMMARIES, query_text=random_node_name)
assert len(search_results) != 0, "Query related summaries don't exist."
print("\nExtracted summaries are:\n")
for result in search_results:
@@ -69,12 +81,18 @@ async def main():
assert len(collection_names) == 0, "LanceDB vector database is not empty"
from cognee.infrastructure.databases.relational import get_relational_engine
- assert not os.path.exists(get_relational_engine().db_path), "SQLite relational database is not empty"
+
+ assert not os.path.exists(
+ get_relational_engine().db_path
+ ), "SQLite relational database is not empty"
from cognee.infrastructure.databases.graph import get_graph_config
+
graph_config = get_graph_config()
assert not os.path.exists(graph_config.graph_file_path), "Networkx graph database is not empty"
+
if __name__ == "__main__":
import asyncio
+
asyncio.run(main(), debug=True)
diff --git a/cognee/tests/test_milvus.py b/cognee/tests/test_milvus.py
index b32d3590b..da02ca936 100644
--- a/cognee/tests/test_milvus.py
+++ b/cognee/tests/test_milvus.py
@@ -10,17 +10,23 @@
async def main():
cognee.config.set_vector_db_provider("milvus")
data_directory_path = str(
- pathlib.Path(os.path.join(pathlib.Path(__file__).parent, ".data_storage/test_milvus")).resolve())
+ pathlib.Path(
+ os.path.join(pathlib.Path(__file__).parent, ".data_storage/test_milvus")
+ ).resolve()
+ )
cognee.config.data_root_directory(data_directory_path)
cognee_directory_path = str(
- pathlib.Path(os.path.join(pathlib.Path(__file__).parent, ".cognee_system/test_milvus")).resolve())
+ pathlib.Path(
+ os.path.join(pathlib.Path(__file__).parent, ".cognee_system/test_milvus")
+ ).resolve()
+ )
cognee.config.system_root_directory(cognee_directory_path)
cognee.config.set_vector_db_config(
{
"vector_db_url": os.path.join(cognee_directory_path, "databases/milvus.db"),
"vector_db_key": "",
- "vector_db_provider": "milvus"
+ "vector_db_provider": "milvus",
}
)
@@ -29,7 +35,9 @@ async def main():
dataset_name = "cs_explanations"
- explanation_file_path = os.path.join(pathlib.Path(__file__).parent, "test_data/Natural_language_processing.txt")
+ explanation_file_path = os.path.join(
+ pathlib.Path(__file__).parent, "test_data/Natural_language_processing.txt"
+ )
await cognee.add([explanation_file_path], dataset_name)
text = """A quantum computer is a computer that takes advantage of quantum mechanical phenomena.
@@ -45,6 +53,7 @@ async def main():
await cognee.cognify([dataset_name])
from cognee.infrastructure.databases.vector import get_vector_engine
+
vector_engine = get_vector_engine()
random_node = (await vector_engine.search("entity_name", "Quantum computer"))[0]
random_node_name = random_node.payload["text"]
@@ -81,4 +90,5 @@ async def main():
if __name__ == "__main__":
import asyncio
+
asyncio.run(main())
diff --git a/cognee/tests/test_neo4j.py b/cognee/tests/test_neo4j.py
index 92e5b5f05..07274c010 100644
--- a/cognee/tests/test_neo4j.py
+++ b/cognee/tests/test_neo4j.py
@@ -1,4 +1,3 @@
-
import os
import logging
import pathlib
@@ -8,19 +7,30 @@
logging.basicConfig(level=logging.DEBUG)
+
async def main():
cognee.config.set_graph_database_provider("neo4j")
- data_directory_path = str(pathlib.Path(os.path.join(pathlib.Path(__file__).parent, ".data_storage/test_neo4j")).resolve())
+ data_directory_path = str(
+ pathlib.Path(
+ os.path.join(pathlib.Path(__file__).parent, ".data_storage/test_neo4j")
+ ).resolve()
+ )
cognee.config.data_root_directory(data_directory_path)
- cognee_directory_path = str(pathlib.Path(os.path.join(pathlib.Path(__file__).parent, ".cognee_system/test_neo4j")).resolve())
+ cognee_directory_path = str(
+ pathlib.Path(
+ os.path.join(pathlib.Path(__file__).parent, ".cognee_system/test_neo4j")
+ ).resolve()
+ )
cognee.config.system_root_directory(cognee_directory_path)
await cognee.prune.prune_data()
- await cognee.prune.prune_system(metadata = True)
+ await cognee.prune.prune_system(metadata=True)
dataset_name = "cs_explanations"
- explanation_file_path = os.path.join(pathlib.Path(__file__).parent, "test_data/Natural_language_processing.txt")
+ explanation_file_path = os.path.join(
+ pathlib.Path(__file__).parent, "test_data/Natural_language_processing.txt"
+ )
await cognee.add([explanation_file_path], dataset_name)
text = """A quantum computer is a computer that takes advantage of quantum mechanical phenomena.
@@ -36,23 +46,24 @@ async def main():
await cognee.cognify([dataset_name])
from cognee.infrastructure.databases.vector import get_vector_engine
+
vector_engine = get_vector_engine()
random_node = (await vector_engine.search("entity_name", "Quantum computer"))[0]
random_node_name = random_node.payload["text"]
- search_results = await cognee.search(SearchType.INSIGHTS, query_text = random_node_name)
+ search_results = await cognee.search(SearchType.INSIGHTS, query_text=random_node_name)
assert len(search_results) != 0, "The search results list is empty."
print("\n\nExtracted sentences are:\n")
for result in search_results:
print(f"{result}\n")
- search_results = await cognee.search(SearchType.CHUNKS, query_text = random_node_name)
+ search_results = await cognee.search(SearchType.CHUNKS, query_text=random_node_name)
assert len(search_results) != 0, "The search results list is empty."
print("\n\nExtracted chunks are:\n")
for result in search_results:
print(f"{result}\n")
- search_results = await cognee.search(SearchType.SUMMARIES, query_text = random_node_name)
+ search_results = await cognee.search(SearchType.SUMMARIES, query_text=random_node_name)
assert len(search_results) != 0, "Query related summaries don't exist."
print("\nExtracted summaries are:\n")
for result in search_results:
@@ -62,7 +73,7 @@ async def main():
assert len(history) == 6, "Search history is not correct."
- results = await brute_force_triplet_search('What is a quantum computer?')
+ results = await brute_force_triplet_search("What is a quantum computer?")
assert len(results) > 0
await cognee.prune.prune_data()
@@ -70,10 +81,13 @@ async def main():
await cognee.prune.prune_system(metadata=True)
from cognee.infrastructure.databases.graph import get_graph_engine
+
graph_engine = await get_graph_engine()
nodes, edges = await graph_engine.get_graph_data()
assert len(nodes) == 0 and len(edges) == 0, "Neo4j graph database is not empty"
+
if __name__ == "__main__":
import asyncio
+
asyncio.run(main())
diff --git a/cognee/tests/test_pgvector.py b/cognee/tests/test_pgvector.py
index 417904089..c241177f0 100644
--- a/cognee/tests/test_pgvector.py
+++ b/cognee/tests/test_pgvector.py
@@ -10,6 +10,7 @@
logging.basicConfig(level=logging.DEBUG)
+
async def test_local_file_deletion(data_text, file_location):
from sqlalchemy import select
import hashlib
@@ -23,40 +24,51 @@ async def test_local_file_deletion(data_text, file_location):
data_hash = hashlib.md5(encoded_text).hexdigest()
# Get data entry from database based on hash contents
data = (await session.scalars(select(Data).where(Data.content_hash == data_hash))).one()
- assert os.path.isfile(data.raw_data_location), f"Data location doesn't exist: {data.raw_data_location}"
+ assert os.path.isfile(
+ data.raw_data_location
+ ), f"Data location doesn't exist: {data.raw_data_location}"
# Test deletion of data along with local files created by cognee
await engine.delete_data_entity(data.id)
assert not os.path.exists(
- data.raw_data_location), f"Data location still exists after deletion: {data.raw_data_location}"
+ data.raw_data_location
+ ), f"Data location still exists after deletion: {data.raw_data_location}"
async with engine.get_async_session() as session:
# Get data entry from database based on file path
- data = (await session.scalars(select(Data).where(Data.raw_data_location == file_location))).one()
- assert os.path.isfile(data.raw_data_location), f"Data location doesn't exist: {data.raw_data_location}"
+ data = (
+ await session.scalars(select(Data).where(Data.raw_data_location == file_location))
+ ).one()
+ assert os.path.isfile(
+ data.raw_data_location
+ ), f"Data location doesn't exist: {data.raw_data_location}"
# Test local files not created by cognee won't get deleted
await engine.delete_data_entity(data.id)
- assert os.path.exists(data.raw_data_location), f"Data location doesn't exists: {data.raw_data_location}"
+ assert os.path.exists(
+ data.raw_data_location
+ ), f"Data location doesn't exists: {data.raw_data_location}"
+
async def test_getting_of_documents(dataset_name_1):
# Test getting of documents for search per dataset
from cognee.modules.users.permissions.methods import get_document_ids_for_user
+
user = await get_default_user()
document_ids = await get_document_ids_for_user(user.id, [dataset_name_1])
- assert len(document_ids) == 1, f"Number of expected documents doesn't match {len(document_ids)} != 1"
+ assert (
+ len(document_ids) == 1
+ ), f"Number of expected documents doesn't match {len(document_ids)} != 1"
# Test getting of documents for search when no dataset is provided
user = await get_default_user()
document_ids = await get_document_ids_for_user(user.id)
- assert len(document_ids) == 2, f"Number of expected documents doesn't match {len(document_ids)} != 2"
+ assert (
+ len(document_ids) == 2
+ ), f"Number of expected documents doesn't match {len(document_ids)} != 2"
async def main():
cognee.config.set_vector_db_config(
- {
- "vector_db_url": "",
- "vector_db_key": "",
- "vector_db_provider": "pgvector"
- }
+ {"vector_db_url": "", "vector_db_key": "", "vector_db_provider": "pgvector"}
)
cognee.config.set_relational_db_config(
{
@@ -84,7 +96,7 @@ async def main():
cognee.config.system_root_directory(cognee_directory_path)
await cognee.prune.prune_data()
- await cognee.prune.prune_system(metadata = True)
+ await cognee.prune.prune_system(metadata=True)
dataset_name_1 = "natural_language"
dataset_name_2 = "quantum"
@@ -114,19 +126,21 @@ async def main():
random_node = (await vector_engine.search("entity_name", "Quantum computer"))[0]
random_node_name = random_node.payload["text"]
- search_results = await cognee.search(SearchType.INSIGHTS, query_text = random_node_name)
+ search_results = await cognee.search(SearchType.INSIGHTS, query_text=random_node_name)
assert len(search_results) != 0, "The search results list is empty."
print("\n\nExtracted sentences are:\n")
for result in search_results:
print(f"{result}\n")
- search_results = await cognee.search(SearchType.CHUNKS, query_text = random_node_name, datasets=[dataset_name_2])
+ search_results = await cognee.search(
+ SearchType.CHUNKS, query_text=random_node_name, datasets=[dataset_name_2]
+ )
assert len(search_results) != 0, "The search results list is empty."
print("\n\nExtracted chunks are:\n")
for result in search_results:
print(f"{result}\n")
- search_results = await cognee.search(SearchType.SUMMARIES, query_text = random_node_name)
+ search_results = await cognee.search(SearchType.SUMMARIES, query_text=random_node_name)
assert len(search_results) != 0, "Query related summaries don't exist."
print("\n\nExtracted summaries are:\n")
for result in search_results:
@@ -135,7 +149,7 @@ async def main():
history = await cognee.get_search_history()
assert len(history) == 6, "Search history is not correct."
- results = await brute_force_triplet_search('What is a quantum computer?')
+ results = await brute_force_triplet_search("What is a quantum computer?")
assert len(results) > 0
await test_local_file_deletion(text, explanation_file_path)
@@ -147,6 +161,7 @@ async def main():
tables_in_database = await vector_engine.get_table_names()
assert len(tables_in_database) == 0, "PostgreSQL database is not empty"
+
if __name__ == "__main__":
import asyncio
diff --git a/cognee/tests/test_qdrant.py b/cognee/tests/test_qdrant.py
index f32e0b4a4..7f82a569f 100644
--- a/cognee/tests/test_qdrant.py
+++ b/cognee/tests/test_qdrant.py
@@ -1,5 +1,3 @@
-
-
import os
import logging
import pathlib
@@ -9,19 +7,30 @@
logging.basicConfig(level=logging.DEBUG)
+
async def main():
cognee.config.set_vector_db_provider("qdrant")
- data_directory_path = str(pathlib.Path(os.path.join(pathlib.Path(__file__).parent, ".data_storage/test_qdrant")).resolve())
+ data_directory_path = str(
+ pathlib.Path(
+ os.path.join(pathlib.Path(__file__).parent, ".data_storage/test_qdrant")
+ ).resolve()
+ )
cognee.config.data_root_directory(data_directory_path)
- cognee_directory_path = str(pathlib.Path(os.path.join(pathlib.Path(__file__).parent, ".cognee_system/test_qdrant")).resolve())
+ cognee_directory_path = str(
+ pathlib.Path(
+ os.path.join(pathlib.Path(__file__).parent, ".cognee_system/test_qdrant")
+ ).resolve()
+ )
cognee.config.system_root_directory(cognee_directory_path)
await cognee.prune.prune_data()
- await cognee.prune.prune_system(metadata = True)
+ await cognee.prune.prune_system(metadata=True)
dataset_name = "cs_explanations"
- explanation_file_path = os.path.join(pathlib.Path(__file__).parent, "test_data/Natural_language_processing.txt")
+ explanation_file_path = os.path.join(
+ pathlib.Path(__file__).parent, "test_data/Natural_language_processing.txt"
+ )
await cognee.add([explanation_file_path], dataset_name)
text = """A quantum computer is a computer that takes advantage of quantum mechanical phenomena.
@@ -37,23 +46,24 @@ async def main():
await cognee.cognify([dataset_name])
from cognee.infrastructure.databases.vector import get_vector_engine
+
vector_engine = get_vector_engine()
random_node = (await vector_engine.search("entity_name", "Quantum computer"))[0]
random_node_name = random_node.payload["text"]
- search_results = await cognee.search(SearchType.INSIGHTS, query_text = random_node_name)
+ search_results = await cognee.search(SearchType.INSIGHTS, query_text=random_node_name)
assert len(search_results) != 0, "The search results list is empty."
print("\n\nExtracted sentences are:\n")
for result in search_results:
print(f"{result}\n")
- search_results = await cognee.search(SearchType.CHUNKS, query_text = random_node_name)
+ search_results = await cognee.search(SearchType.CHUNKS, query_text=random_node_name)
assert len(search_results) != 0, "The search results list is empty."
print("\n\nExtracted chunks are:\n")
for result in search_results:
print(f"{result}\n")
- search_results = await cognee.search(SearchType.SUMMARIES, query_text = random_node_name)
+ search_results = await cognee.search(SearchType.SUMMARIES, query_text=random_node_name)
assert len(search_results) != 0, "Query related summaries don't exist."
print("\nExtracted summaries are:\n")
for result in search_results:
@@ -62,7 +72,7 @@ async def main():
history = await cognee.get_search_history()
assert len(history) == 6, "Search history is not correct."
- results = await brute_force_triplet_search('What is a quantum computer?')
+ results = await brute_force_triplet_search("What is a quantum computer?")
assert len(results) > 0
await cognee.prune.prune_data()
@@ -73,6 +83,8 @@ async def main():
collections_response = await qdrant_client.get_collections()
assert len(collections_response.collections) == 0, "QDrant vector database is not empty"
+
if __name__ == "__main__":
import asyncio
+
asyncio.run(main())
diff --git a/cognee/tests/test_weaviate.py b/cognee/tests/test_weaviate.py
index 43ec30aaf..874f21347 100644
--- a/cognee/tests/test_weaviate.py
+++ b/cognee/tests/test_weaviate.py
@@ -7,19 +7,30 @@
logging.basicConfig(level=logging.DEBUG)
+
async def main():
cognee.config.set_vector_db_provider("weaviate")
- data_directory_path = str(pathlib.Path(os.path.join(pathlib.Path(__file__).parent, ".data_storage/test_weaviate")).resolve())
+ data_directory_path = str(
+ pathlib.Path(
+ os.path.join(pathlib.Path(__file__).parent, ".data_storage/test_weaviate")
+ ).resolve()
+ )
cognee.config.data_root_directory(data_directory_path)
- cognee_directory_path = str(pathlib.Path(os.path.join(pathlib.Path(__file__).parent, ".cognee_system/test_weaviate")).resolve())
+ cognee_directory_path = str(
+ pathlib.Path(
+ os.path.join(pathlib.Path(__file__).parent, ".cognee_system/test_weaviate")
+ ).resolve()
+ )
cognee.config.system_root_directory(cognee_directory_path)
await cognee.prune.prune_data()
- await cognee.prune.prune_system(metadata = True)
+ await cognee.prune.prune_system(metadata=True)
dataset_name = "cs_explanations"
- explanation_file_path = os.path.join(pathlib.Path(__file__).parent, "test_data/Natural_language_processing.txt")
+ explanation_file_path = os.path.join(
+ pathlib.Path(__file__).parent, "test_data/Natural_language_processing.txt"
+ )
await cognee.add([explanation_file_path], dataset_name)
text = """A quantum computer is a computer that takes advantage of quantum mechanical phenomena.
@@ -35,23 +46,24 @@ async def main():
await cognee.cognify([dataset_name])
from cognee.infrastructure.databases.vector import get_vector_engine
+
vector_engine = get_vector_engine()
random_node = (await vector_engine.search("entity_name", "Quantum computer"))[0]
random_node_name = random_node.payload["text"]
- search_results = await cognee.search(SearchType.INSIGHTS, query_text = random_node_name)
+ search_results = await cognee.search(SearchType.INSIGHTS, query_text=random_node_name)
assert len(search_results) != 0, "The search results list is empty."
print("\n\nExtracted sentences are:\n")
for result in search_results:
print(f"{result}\n")
- search_results = await cognee.search(SearchType.CHUNKS, query_text = random_node_name)
+ search_results = await cognee.search(SearchType.CHUNKS, query_text=random_node_name)
assert len(search_results) != 0, "The search results list is empty."
print("\n\nExtracted chunks are:\n")
for result in search_results:
print(f"{result}\n")
- search_results = await cognee.search(SearchType.SUMMARIES, query_text = random_node_name)
+ search_results = await cognee.search(SearchType.SUMMARIES, query_text=random_node_name)
assert len(search_results) != 0, "Query related summaries don't exist."
print("\nExtracted summaries are:\n")
for result in search_results:
@@ -60,16 +72,18 @@ async def main():
history = await cognee.get_search_history()
assert len(history) == 6, "Search history is not correct."
- results = await brute_force_triplet_search('What is a quantum computer?')
+ results = await brute_force_triplet_search("What is a quantum computer?")
assert len(results) > 0
await cognee.prune.prune_data()
assert not os.path.isdir(data_directory_path), "Local data files are not deleted"
await cognee.prune.prune_system(metadata=True)
- collections = get_vector_engine().client.collections.list_all()
+ collections = get_vector_engine().client.collections.list_all()
assert len(collections) == 0, "Weaviate vector database is not empty"
+
if __name__ == "__main__":
import asyncio
+
asyncio.run(main())
diff --git a/cognee/tests/unit/infrastructure/databases/test_index_graph_edges.py b/cognee/tests/unit/infrastructure/databases/test_index_graph_edges.py
index e4c1dac05..59ad606b0 100644
--- a/cognee/tests/unit/infrastructure/databases/test_index_graph_edges.py
+++ b/cognee/tests/unit/infrastructure/databases/test_index_graph_edges.py
@@ -1,21 +1,33 @@
import pytest
from unittest.mock import AsyncMock, patch
+
@pytest.mark.asyncio
async def test_index_graph_edges_success():
"""Test that index_graph_edges uses the index datapoints and creates vector index."""
mock_graph_engine = AsyncMock()
- mock_graph_engine.get_graph_data.return_value = (None, [
- [{"relationship_name": "rel1"}, {"relationship_name": "rel1"}],
- [{"relationship_name": "rel2"}]
- ])
+ mock_graph_engine.get_graph_data.return_value = (
+ None,
+ [
+ [{"relationship_name": "rel1"}, {"relationship_name": "rel1"}],
+ [{"relationship_name": "rel2"}],
+ ],
+ )
mock_vector_engine = AsyncMock()
- with patch("cognee.tasks.storage.index_graph_edges.get_graph_engine", return_value=mock_graph_engine), \
- patch("cognee.tasks.storage.index_graph_edges.get_vector_engine", return_value=mock_vector_engine):
-
+ with (
+ patch(
+ "cognee.tasks.storage.index_graph_edges.get_graph_engine",
+ return_value=mock_graph_engine,
+ ),
+ patch(
+ "cognee.tasks.storage.index_graph_edges.get_vector_engine",
+ return_value=mock_vector_engine,
+ ),
+ ):
from cognee.tasks.storage.index_graph_edges import index_graph_edges
+
await index_graph_edges()
mock_graph_engine.get_graph_data.assert_awaited_once()
@@ -31,10 +43,18 @@ async def test_index_graph_edges_no_relationships():
mock_vector_engine = AsyncMock()
- with patch("cognee.tasks.storage.index_graph_edges.get_graph_engine", return_value=mock_graph_engine), \
- patch("cognee.tasks.storage.index_graph_edges.get_vector_engine", return_value=mock_vector_engine):
-
+ with (
+ patch(
+ "cognee.tasks.storage.index_graph_edges.get_graph_engine",
+ return_value=mock_graph_engine,
+ ),
+ patch(
+ "cognee.tasks.storage.index_graph_edges.get_vector_engine",
+ return_value=mock_vector_engine,
+ ),
+ ):
from cognee.tasks.storage.index_graph_edges import index_graph_edges
+
await index_graph_edges()
mock_graph_engine.get_graph_data.assert_awaited_once()
@@ -45,12 +65,14 @@ async def test_index_graph_edges_no_relationships():
@pytest.mark.asyncio
async def test_index_graph_edges_initialization_error():
"""Test that index_graph_edges raises a RuntimeError if initialization fails."""
- with patch("cognee.tasks.storage.index_graph_edges.get_graph_engine", side_effect=Exception("Graph engine failed")), \
- patch("cognee.tasks.storage.index_graph_edges.get_vector_engine", return_value=AsyncMock()):
-
+ with (
+ patch(
+ "cognee.tasks.storage.index_graph_edges.get_graph_engine",
+ side_effect=Exception("Graph engine failed"),
+ ),
+ patch("cognee.tasks.storage.index_graph_edges.get_vector_engine", return_value=AsyncMock()),
+ ):
from cognee.tasks.storage.index_graph_edges import index_graph_edges
with pytest.raises(RuntimeError, match="Initialization error"):
await index_graph_edges()
-
-
diff --git a/cognee/tests/unit/interfaces/graph/get_graph_from_huge_model_test.py b/cognee/tests/unit/interfaces/graph/get_graph_from_huge_model_test.py
index 06c74c854..9900670f2 100644
--- a/cognee/tests/unit/interfaces/graph/get_graph_from_huge_model_test.py
+++ b/cognee/tests/unit/interfaces/graph/get_graph_from_huge_model_test.py
@@ -9,30 +9,25 @@
random.seed(1500)
+
class Repository(DataPoint):
path: str
- _metadata = {
- "index_fields": [],
- "type": "Repository"
- }
+ _metadata = {"index_fields": [], "type": "Repository"}
+
class CodeFile(DataPoint):
part_of: Repository
contains: List["CodePart"] = []
depends_on: List["CodeFile"] = []
source_code: str
- _metadata = {
- "index_fields": [],
- "type": "CodeFile"
- }
+ _metadata = {"index_fields": [], "type": "CodeFile"}
+
class CodePart(DataPoint):
part_of: CodeFile
source_code: str
- _metadata = {
- "index_fields": [],
- "type": "CodePart"
- }
+ _metadata = {"index_fields": [], "type": "CodePart"}
+
CodeFile.model_rebuild()
CodePart.model_rebuild()
@@ -41,13 +36,13 @@ class CodePart(DataPoint):
def nanoseconds_to_largest_unit(nanoseconds):
# Define conversion factors
conversion_factors = {
- 'weeks': 7 * 24 * 60 * 60 * 1e9,
- 'days': 24 * 60 * 60 * 1e9,
- 'hours': 60 * 60 * 1e9,
- 'minutes': 60 * 1e9,
- 'seconds': 1e9,
- 'miliseconds': 1e6,
- 'microseconds': 1e3,
+ "weeks": 7 * 24 * 60 * 60 * 1e9,
+ "days": 24 * 60 * 60 * 1e9,
+ "hours": 60 * 60 * 1e9,
+ "minutes": 60 * 1e9,
+ "seconds": 1e9,
+ "miliseconds": 1e6,
+ "microseconds": 1e3,
}
# Iterate through conversion factors to find the largest unit
@@ -57,30 +52,41 @@ def nanoseconds_to_largest_unit(nanoseconds):
return converted_value, unit
# If nanoseconds is smaller than a second
- return nanoseconds, 'nanoseconds'
+ return nanoseconds, "nanoseconds"
async def test_circular_reference_extraction():
- repo = Repository(path = "repo1")
-
- code_files = [CodeFile(
- id = uuid5(NAMESPACE_OID, f"file{file_index}"),
- source_code = "source code",
- part_of = repo,
- contains = [],
- depends_on = [CodeFile(
- id = uuid5(NAMESPACE_OID, f"file{random_id}"),
- source_code = "source code",
- part_of = repo,
- depends_on = [],
- ) for random_id in [random.randint(0, 1499) for _ in range(random.randint(0, 5))]],
- ) for file_index in range(1500)]
+ repo = Repository(path="repo1")
+
+ code_files = [
+ CodeFile(
+ id=uuid5(NAMESPACE_OID, f"file{file_index}"),
+ source_code="source code",
+ part_of=repo,
+ contains=[],
+ depends_on=[
+ CodeFile(
+ id=uuid5(NAMESPACE_OID, f"file{random_id}"),
+ source_code="source code",
+ part_of=repo,
+ depends_on=[],
+ )
+ for random_id in [random.randint(0, 1499) for _ in range(random.randint(0, 5))]
+ ],
+ )
+ for file_index in range(1500)
+ ]
for code_file in code_files:
- code_file.contains.extend([CodePart(
- part_of = code_file,
- source_code = f"Part {part_index}",
- ) for part_index in range(random.randint(1, 20))])
+ code_file.contains.extend(
+ [
+ CodePart(
+ part_of=code_file,
+ source_code=f"Part {part_index}",
+ )
+ for part_index in range(random.randint(1, 20))
+ ]
+ )
nodes = []
edges = []
@@ -90,9 +96,12 @@ async def test_circular_reference_extraction():
start = time.perf_counter_ns()
- results = await asyncio.gather(*[
- get_graph_from_model(code_file, added_nodes = added_nodes, added_edges = added_edges) for code_file in code_files
- ])
+ results = await asyncio.gather(
+ *[
+ get_graph_from_model(code_file, added_nodes=added_nodes, added_edges=added_edges)
+ for code_file in code_files
+ ]
+ )
time_to_run = time.perf_counter_ns() - start
@@ -105,5 +114,6 @@ async def test_circular_reference_extraction():
assert len(nodes) == 1501
assert len(edges) == 1501 * 20 + 1500 * 5
+
if __name__ == "__main__":
asyncio.run(test_circular_reference_extraction())
diff --git a/cognee/tests/unit/interfaces/graph/get_graph_from_model_test.py b/cognee/tests/unit/interfaces/graph/get_graph_from_model_test.py
index 499dc9f3f..36fa71c25 100644
--- a/cognee/tests/unit/interfaces/graph/get_graph_from_model_test.py
+++ b/cognee/tests/unit/interfaces/graph/get_graph_from_model_test.py
@@ -9,55 +9,52 @@
class Document(DataPoint):
path: str
- _metadata = {
- "index_fields": [],
- "type": "Document"
- }
+ _metadata = {"index_fields": [], "type": "Document"}
+
class DocumentChunk(DataPoint):
part_of: Document
text: str
contains: List["Entity"] = None
- _metadata = {
- "index_fields": ["text"],
- "type": "DocumentChunk"
- }
+ _metadata = {"index_fields": ["text"], "type": "DocumentChunk"}
+
class EntityType(DataPoint):
name: str
- _metadata = {
- "index_fields": ["name"],
- "type": "EntityType"
- }
+ _metadata = {"index_fields": ["name"], "type": "EntityType"}
+
class Entity(DataPoint):
name: str
is_type: EntityType
- _metadata = {
- "index_fields": ["name"],
- "type": "Entity"
- }
+ _metadata = {"index_fields": ["name"], "type": "Entity"}
+
DocumentChunk.model_rebuild()
async def get_graph_from_model_test():
- document = Document(path = "file_path")
-
- document_chunks = [DocumentChunk(
- id = uuid5(NAMESPACE_OID, f"file{file_index}"),
- text = "some text",
- part_of = document,
- contains = [],
- ) for file_index in range(1)]
+ document = Document(path="file_path")
+
+ document_chunks = [
+ DocumentChunk(
+ id=uuid5(NAMESPACE_OID, f"file{file_index}"),
+ text="some text",
+ part_of=document,
+ contains=[],
+ )
+ for file_index in range(1)
+ ]
for document_chunk in document_chunks:
- document_chunk.contains.append(Entity(
- name = f"Entity",
- is_type = EntityType(
- name = "Type 1",
- ),
- ))
+ document_chunk.contains.append(
+ Entity(
+ name="Entity",
+ is_type=EntityType(
+ name="Type 1",
+ ),
+ )
+ )
nodes = []
edges = []
@@ -66,14 +63,17 @@ async def get_graph_from_model_test():
added_edges = {}
visited_properties = {}
- results = await asyncio.gather(*[
- get_graph_from_model(
- document_chunk,
- added_nodes = added_nodes,
- added_edges = added_edges,
- visited_properties = visited_properties,
- ) for document_chunk in document_chunks
- ])
+ results = await asyncio.gather(
+ *[
+ get_graph_from_model(
+ document_chunk,
+ added_nodes=added_nodes,
+ added_edges=added_edges,
+ visited_properties=visited_properties,
+ )
+ for document_chunk in document_chunks
+ ]
+ )
for result_nodes, result_edges in results:
nodes.extend(result_nodes)
@@ -82,5 +82,6 @@ async def get_graph_from_model_test():
assert len(nodes) == 4
assert len(edges) == 3
+
if __name__ == "__main__":
asyncio.run(get_graph_from_model_test())
diff --git a/cognee/tests/unit/modules/graph/cognee_graph_elements_test.py b/cognee/tests/unit/modules/graph/cognee_graph_elements_test.py
index 1d9bad07c..000856b12 100644
--- a/cognee/tests/unit/modules/graph/cognee_graph_elements_test.py
+++ b/cognee/tests/unit/modules/graph/cognee_graph_elements_test.py
@@ -9,7 +9,7 @@ def test_node_initialization():
"""Test that a Node is initialized correctly."""
node = Node("node1", {"attr1": "value1"}, dimension=2)
assert node.id == "node1"
- assert node.attributes == {"attr1": "value1", 'vector_distance': np.inf}
+ assert node.attributes == {"attr1": "value1", "vector_distance": np.inf}
assert len(node.status) == 2
assert np.all(node.status == 1)
@@ -96,7 +96,7 @@ def test_edge_initialization():
edge = Edge(node1, node2, {"weight": 10}, directed=False, dimension=2)
assert edge.node1 == node1
assert edge.node2 == node2
- assert edge.attributes == {'vector_distance': np.inf,"weight": 10}
+ assert edge.attributes == {"vector_distance": np.inf, "weight": 10}
assert edge.directed is False
assert len(edge.status) == 2
assert np.all(edge.status == 1)
diff --git a/cognee/tests/unit/modules/retriever/test_description_to_codepart_search.py b/cognee/tests/unit/modules/retriever/test_description_to_codepart_search.py
index 4c39883a9..5cceade71 100644
--- a/cognee/tests/unit/modules/retriever/test_description_to_codepart_search.py
+++ b/cognee/tests/unit/modules/retriever/test_description_to_codepart_search.py
@@ -2,7 +2,6 @@
from unittest.mock import AsyncMock, patch
-
@pytest.mark.asyncio
async def test_code_description_to_code_part_no_results():
"""Test that code_description_to_code_part handles no search results."""
@@ -12,11 +11,24 @@ async def test_code_description_to_code_part_no_results():
mock_vector_engine = AsyncMock()
mock_vector_engine.search.return_value = []
- with patch("cognee.modules.retrieval.description_to_codepart_search.get_vector_engine", return_value=mock_vector_engine), \
- patch("cognee.modules.retrieval.description_to_codepart_search.get_graph_engine", return_value=AsyncMock()), \
- patch("cognee.modules.retrieval.description_to_codepart_search.CogneeGraph", return_value=AsyncMock()):
+ with (
+ patch(
+ "cognee.modules.retrieval.description_to_codepart_search.get_vector_engine",
+ return_value=mock_vector_engine,
+ ),
+ patch(
+ "cognee.modules.retrieval.description_to_codepart_search.get_graph_engine",
+ return_value=AsyncMock(),
+ ),
+ patch(
+ "cognee.modules.retrieval.description_to_codepart_search.CogneeGraph",
+ return_value=AsyncMock(),
+ ),
+ ):
+ from cognee.modules.retrieval.description_to_codepart_search import (
+ code_description_to_code_part,
+ )
- from cognee.modules.retrieval.description_to_codepart_search import code_description_to_code_part
result = await code_description_to_code_part("search query", mock_user, 2)
assert result == []
@@ -29,7 +41,10 @@ async def test_code_description_to_code_part_invalid_query():
mock_user = AsyncMock()
with pytest.raises(ValueError, match="The query must be a non-empty string."):
- from cognee.modules.retrieval.description_to_codepart_search import code_description_to_code_part
+ from cognee.modules.retrieval.description_to_codepart_search import (
+ code_description_to_code_part,
+ )
+
await code_description_to_code_part("", mock_user, 2)
@@ -40,7 +55,10 @@ async def test_code_description_to_code_part_invalid_top_k():
mock_user = AsyncMock()
with pytest.raises(ValueError, match="top_k must be a positive integer."):
- from cognee.modules.retrieval.description_to_codepart_search import code_description_to_code_part
+ from cognee.modules.retrieval.description_to_codepart_search import (
+ code_description_to_code_part,
+ )
+
await code_description_to_code_part("search query", mock_user, 0)
@@ -50,11 +68,23 @@ async def test_code_description_to_code_part_initialization_error():
mock_user = AsyncMock()
- with patch("cognee.modules.retrieval.description_to_codepart_search.get_vector_engine", side_effect=Exception("Engine init failed")), \
- patch("cognee.modules.retrieval.description_to_codepart_search.get_graph_engine", return_value=AsyncMock()):
-
- from cognee.modules.retrieval.description_to_codepart_search import code_description_to_code_part
- with pytest.raises(RuntimeError, match="System initialization error. Please try again later."):
+ with (
+ patch(
+ "cognee.modules.retrieval.description_to_codepart_search.get_vector_engine",
+ side_effect=Exception("Engine init failed"),
+ ),
+ patch(
+ "cognee.modules.retrieval.description_to_codepart_search.get_graph_engine",
+ return_value=AsyncMock(),
+ ),
+ ):
+ from cognee.modules.retrieval.description_to_codepart_search import (
+ code_description_to_code_part,
+ )
+
+ with pytest.raises(
+ RuntimeError, match="System initialization error. Please try again later."
+ ):
await code_description_to_code_part("search query", mock_user, 2)
@@ -67,10 +97,23 @@ async def test_code_description_to_code_part_execution_error():
mock_vector_engine = AsyncMock()
mock_vector_engine.search.side_effect = Exception("Execution error")
- with patch("cognee.modules.retrieval.description_to_codepart_search.get_vector_engine", return_value=mock_vector_engine), \
- patch("cognee.modules.retrieval.description_to_codepart_search.get_graph_engine", return_value=AsyncMock()), \
- patch("cognee.modules.retrieval.description_to_codepart_search.CogneeGraph", return_value=AsyncMock()):
+ with (
+ patch(
+ "cognee.modules.retrieval.description_to_codepart_search.get_vector_engine",
+ return_value=mock_vector_engine,
+ ),
+ patch(
+ "cognee.modules.retrieval.description_to_codepart_search.get_graph_engine",
+ return_value=AsyncMock(),
+ ),
+ patch(
+ "cognee.modules.retrieval.description_to_codepart_search.CogneeGraph",
+ return_value=AsyncMock(),
+ ),
+ ):
+ from cognee.modules.retrieval.description_to_codepart_search import (
+ code_description_to_code_part,
+ )
- from cognee.modules.retrieval.description_to_codepart_search import code_description_to_code_part
with pytest.raises(RuntimeError, match="An error occurred while processing your request."):
await code_description_to_code_part("search query", mock_user, 2)
diff --git a/cognee/tests/unit/processing/chunks/chunk_by_paragraph_2_test.py b/cognee/tests/unit/processing/chunks/chunk_by_paragraph_2_test.py
index 8e900727d..728b5cda4 100644
--- a/cognee/tests/unit/processing/chunks/chunk_by_paragraph_2_test.py
+++ b/cognee/tests/unit/processing/chunks/chunk_by_paragraph_2_test.py
@@ -29,9 +29,7 @@ def test_chunk_by_paragraph_isomorphism(input_text, paragraph_length, batch_para
def test_paragraph_chunk_length(input_text, paragraph_length, batch_paragraphs):
chunks = list(chunk_by_paragraph(input_text, paragraph_length, batch_paragraphs))
- chunk_lengths = np.array(
- [len(list(chunk_by_word(chunk["text"]))) for chunk in chunks]
- )
+ chunk_lengths = np.array([len(list(chunk_by_word(chunk["text"]))) for chunk in chunks])
larger_chunks = chunk_lengths[chunk_lengths > paragraph_length]
assert np.all(
@@ -43,9 +41,7 @@ def test_paragraph_chunk_length(input_text, paragraph_length, batch_paragraphs):
"input_text,paragraph_length,batch_paragraphs",
list(product(list(INPUT_TEXTS.values()), paragraph_lengths, batch_paragraphs_vals)),
)
-def test_chunk_by_paragraph_chunk_numbering(
- input_text, paragraph_length, batch_paragraphs
-):
+def test_chunk_by_paragraph_chunk_numbering(input_text, paragraph_length, batch_paragraphs):
chunks = chunk_by_paragraph(input_text, paragraph_length, batch_paragraphs)
chunk_indices = np.array([chunk["chunk_index"] for chunk in chunks])
assert np.all(
diff --git a/cognee/tests/unit/processing/chunks/chunk_by_word_test.py b/cognee/tests/unit/processing/chunks/chunk_by_word_test.py
index 42523c106..fb26638cb 100644
--- a/cognee/tests/unit/processing/chunks/chunk_by_word_test.py
+++ b/cognee/tests/unit/processing/chunks/chunk_by_word_test.py
@@ -35,6 +35,4 @@ def test_chunk_by_word_splits(input_text):
chunks = np.array(list(chunk_by_word(input_text)))
space_test = np.array([" " not in chunk[0].strip() for chunk in chunks])
- assert np.all(
- space_test
- ), f"These chunks contain spaces within them: {chunks[space_test == False]}"
+ assert np.all(space_test), f"These chunks contain spaces within them: {chunks[~space_test]}"
diff --git a/cognee/tests/unit/processing/utils/utils_test.py b/cognee/tests/unit/processing/utils/utils_test.py
new file mode 100644
index 000000000..f8c325100
--- /dev/null
+++ b/cognee/tests/unit/processing/utils/utils_test.py
@@ -0,0 +1,131 @@
+import os
+import pytest
+import networkx as nx
+import pandas as pd
+from unittest.mock import patch, mock_open
+from io import BytesIO
+from uuid import uuid4
+from datetime import datetime, timezone
+from cognee.shared.exceptions import IngestionError
+
+from cognee.shared.utils import (
+ get_anonymous_id,
+ send_telemetry,
+ num_tokens_from_string,
+ get_file_content_hash,
+ trim_text_to_max_tokens,
+ prepare_edges,
+ prepare_nodes,
+ create_cognee_style_network_with_logo,
+ graph_to_tuple,
+)
+
+
+@pytest.fixture
+def temp_dir(tmp_path):
+ return tmp_path
+
+
+@patch("os.makedirs")
+@patch("builtins.open", new_callable=mock_open, read_data=str(uuid4()))
+def test_get_anonymous_id(mock_open_file, mock_makedirs, temp_dir):
+ os.environ["HOME"] = str(temp_dir)
+ anon_id = get_anonymous_id()
+ assert isinstance(anon_id, str)
+ assert len(anon_id) > 0
+
+
+# @patch("requests.post")
+# def test_send_telemetry(mock_post):
+# mock_post.return_value.status_code = 200
+#
+# send_telemetry("test_event", "test_user", {"key": "value"})
+# mock_post.assert_called_once()
+#
+# args, kwargs = mock_post.call_args
+# assert kwargs["json"]["event_name"] == "test_event"
+
+#
+# @patch("tiktoken.encoding_for_model")
+# def test_num_tokens_from_string(mock_encoding):
+# mock_encoding.return_value.encode = lambda x: list(x)
+#
+# assert num_tokens_from_string("hello", "test_encoding") == 5
+# assert num_tokens_from_string("world", "test_encoding") == 5
+#
+
+
+@patch("builtins.open", new_callable=mock_open, read_data=b"test_data")
+def test_get_file_content_hash_file(mock_open_file):
+ import hashlib
+
+ expected_hash = hashlib.md5(b"test_data").hexdigest()
+ result = get_file_content_hash("test_file.txt")
+ assert result == expected_hash
+
+
+def test_get_file_content_hash_stream():
+ stream = BytesIO(b"test_data")
+ import hashlib
+
+ expected_hash = hashlib.md5(b"test_data").hexdigest()
+ result = get_file_content_hash(stream)
+ assert result == expected_hash
+
+
+# def test_trim_text_to_max_tokens():
+# text = "This is a test string with multiple words."
+# encoding_name = "test_encoding"
+#
+# with patch("tiktoken.get_encoding") as mock_get_encoding:
+# mock_get_encoding.return_value.encode = lambda x: list(x)
+# mock_get_encoding.return_value.decode = lambda x: "".join(x)
+#
+# result = trim_text_to_max_tokens(text, 5, encoding_name)
+# assert result == text[:5]
+
+
+def test_prepare_edges():
+ graph = nx.MultiDiGraph()
+ graph.add_edge("A", "B", key="AB", weight=1)
+ edges_df = prepare_edges(graph, "source", "target", "key")
+
+ assert isinstance(edges_df, pd.DataFrame)
+ assert len(edges_df) == 1
+
+
+def test_prepare_nodes():
+ graph = nx.Graph()
+ graph.add_node(1, name="Node1")
+ nodes_df = prepare_nodes(graph)
+
+ assert isinstance(nodes_df, pd.DataFrame)
+ assert len(nodes_df) == 1
+
+
+@pytest.mark.asyncio
+async def test_create_cognee_style_network_with_logo():
+ import networkx as nx
+ from unittest.mock import patch
+ from io import BytesIO
+
+ # Create a sample graph
+ graph = nx.Graph()
+ graph.add_node(1, group="A")
+ graph.add_node(2, group="B")
+ graph.add_edge(1, 2)
+
+ # Convert the graph to a tuple format for serialization
+ graph_tuple = graph_to_tuple(graph)
+
+ result = await create_cognee_style_network_with_logo(
+ graph_tuple,
+ title="Test Network",
+ layout_func=nx.spring_layout,
+ layout_scale=3.0,
+ logo_alpha=0.5,
+ )
+
+ assert result is not None
+ assert isinstance(result, str)
+ assert len(result) > 0
diff --git a/evals/deepeval_metrics.py b/evals/deepeval_metrics.py
index b07d2e1ac..9ce1e9e4f 100644
--- a/evals/deepeval_metrics.py
+++ b/evals/deepeval_metrics.py
@@ -6,21 +6,17 @@
correctness_metric = GEval(
name="Correctness",
model="gpt-4o-mini",
- evaluation_params=[
- LLMTestCaseParams.ACTUAL_OUTPUT,
- LLMTestCaseParams.EXPECTED_OUTPUT
- ],
+ evaluation_params=[LLMTestCaseParams.ACTUAL_OUTPUT, LLMTestCaseParams.EXPECTED_OUTPUT],
evaluation_steps=[
- "Determine whether the actual output is factually correct based on the expected output."
- ]
+ "Determine whether the actual output is factually correct based on the expected output."
+ ],
)
class f1_score_metric(BaseMetric):
-
- """F1 score taken directly from the official hotpot benchmark
+ """F1 score taken directly from the official hotpot benchmark
implementation and wrapped into a deepeval metric."""
-
+
def __init__(self, threshold: float = 0.5):
self.threshold = threshold
@@ -43,12 +39,12 @@ def is_successful(self):
@property
def __name__(self):
return "Official hotpot F1 score"
-
-class em_score_metric(BaseMetric):
- """Exact Match score taken directly from the official hotpot benchmark
+
+class em_score_metric(BaseMetric):
+ """Exact Match score taken directly from the official hotpot benchmark
implementation and wrapped into a deepeval metric."""
-
+
def __init__(self, threshold: float = 0.5):
self.threshold = threshold
@@ -69,4 +65,4 @@ def is_successful(self):
@property
def __name__(self):
- return "Official hotpot EM score"
\ No newline at end of file
+ return "Official hotpot EM score"
diff --git a/evals/eval_on_hotpot.py b/evals/eval_on_hotpot.py
index e07e80e0c..5a95ea05d 100644
--- a/evals/eval_on_hotpot.py
+++ b/evals/eval_on_hotpot.py
@@ -34,19 +34,25 @@ async def answer_without_cognee(instance):
)
return answer_prediction
+
async def answer_with_cognee(instance):
await cognee.prune.prune_data()
await cognee.prune.prune_system(metadata=True)
-
- for (title, sentences) in instance["context"]:
- await cognee.add("\n".join(sentences), dataset_name = "HotPotQA")
-
- await cognee.cognify("HotPotQA")
-
- search_results = await cognee.search(
- SearchType.INSIGHTS, query_text=instance["question"]
+
+ for title, sentences in instance["context"]:
+ await cognee.add("\n".join(sentences), dataset_name="HotPotQA")
+
+ for n in range(1, 4):
+ print(n)
+
+ await cognee.cognify("HotPotQA")
+
+ search_results = await cognee.search(SearchType.INSIGHTS, query_text=instance["question"])
+ search_results_second = await cognee.search(
+ SearchType.SUMMARIES, query_text=instance["question"]
)
-
+ search_results = search_results + search_results_second
+
args = {
"question": instance["question"],
"context": search_results,
@@ -60,61 +66,67 @@ async def answer_with_cognee(instance):
system_prompt=system_prompt,
response_model=str,
)
-
+
return answer_prediction
async def eval_answers(instances, answers, eval_metric):
test_cases = []
-
+
for instance, answer in zip(instances, answers):
test_case = LLMTestCase(
- input=instance["question"],
- actual_output=answer,
- expected_output=instance["answer"]
+ input=instance["question"], actual_output=answer, expected_output=instance["answer"]
)
test_cases.append(test_case)
-
+
eval_set = EvaluationDataset(test_cases)
eval_results = eval_set.evaluate([eval_metric])
-
+
return eval_results
+
async def eval_on_hotpotQA(answer_provider, num_samples, eval_metric):
base_config = get_base_config()
data_root_dir = base_config.data_root_directory
-
+
if not Path(data_root_dir).exists():
Path(data_root_dir).mkdir()
-
+
filepath = data_root_dir / Path("hotpot_dev_fullwiki_v1.json")
if not filepath.exists():
- url = 'http://curtis.ml.cmu.edu/datasets/hotpot/hotpot_dev_fullwiki_v1.json'
+ url = "http://curtis.ml.cmu.edu/datasets/hotpot/hotpot_dev_fullwiki_v1.json"
wget.download(url, out=data_root_dir)
-
+
with open(filepath, "r") as file:
dataset = json.load(file)
-
+
instances = dataset if not num_samples else dataset[:num_samples]
answers = []
for instance in tqdm(instances, desc="Getting answers"):
answer = await answer_provider(instance)
answers.append(answer)
-
+
eval_results = await eval_answers(instances, answers, eval_metric)
- avg_score = statistics.mean([result.metrics_data[0].score for result in eval_results.test_results])
-
+ avg_score = statistics.mean(
+ [result.metrics_data[0].score for result in eval_results.test_results]
+ )
+
return avg_score
+
if __name__ == "__main__":
parser = argparse.ArgumentParser()
-
+
parser.add_argument("--with_cognee", action="store_true")
parser.add_argument("--num_samples", type=int, default=500)
- parser.add_argument("--metric", type=str, default="correctness_metric",
- help="Valid options are Deepeval metrics (e.g. AnswerRelevancyMetric) \
- and metrics defined in evals/deepeval_metrics.py, e.g. f1_score_metric")
-
+ parser.add_argument(
+ "--metric",
+ type=str,
+ default="correctness_metric",
+ help="Valid options are Deepeval metrics (e.g. AnswerRelevancyMetric) \
+ and metrics defined in evals/deepeval_metrics.py, e.g. f1_score_metric",
+ )
+
args = parser.parse_args()
try:
@@ -124,11 +136,11 @@ async def eval_on_hotpotQA(answer_provider, num_samples, eval_metric):
metric = getattr(evals.deepeval_metrics, args.metric)
if isinstance(metric, type):
metric = metric()
-
+
if args.with_cognee:
answer_provider = answer_with_cognee
else:
answer_provider = answer_without_cognee
-
+
avg_score = asyncio.run(eval_on_hotpotQA(answer_provider, args.num_samples, metric))
- print(f"Average {args.metric}: {avg_score}")
\ No newline at end of file
+ print(f"Average {args.metric}: {avg_score}")
diff --git a/evals/eval_swe_bench.py b/evals/eval_swe_bench.py
index 6c2280d80..789c95ab4 100644
--- a/evals/eval_swe_bench.py
+++ b/evals/eval_swe_bench.py
@@ -11,8 +11,7 @@
from cognee.api.v1.search import SearchType
from cognee.infrastructure.llm.get_llm_client import get_llm_client
from cognee.infrastructure.llm.prompts import read_query_prompt
-from cognee.modules.retrieval.brute_force_triplet_search import \
- brute_force_triplet_search
+from cognee.modules.retrieval.brute_force_triplet_search import brute_force_triplet_search
from cognee.shared.utils import render_graph
from evals.eval_utils import download_github_repo, retrieved_edges_to_string
@@ -27,41 +26,42 @@ def check_install_package(package_name):
return True
except ImportError:
try:
- subprocess.check_call(
- [sys.executable, "-m", "pip", "install", package_name]
- )
+ subprocess.check_call([sys.executable, "-m", "pip", "install", package_name])
return True
except subprocess.CalledProcessError:
return False
async def generate_patch_with_cognee(instance, llm_client, search_type=SearchType.CHUNKS):
- repo_path = download_github_repo(instance, '../RAW_GIT_REPOS')
+ repo_path = download_github_repo(instance, "../RAW_GIT_REPOS")
pipeline = await run_code_graph_pipeline(repo_path)
async for result in pipeline:
print(result)
- print('Here we have the repo under the repo_path')
+ print("Here we have the repo under the repo_path")
await render_graph(None, include_labels=True, include_nodes=True)
- problem_statement = instance['problem_statement']
+ problem_statement = instance["problem_statement"]
instructions = read_query_prompt("patch_gen_kg_instructions.txt")
- retrieved_edges = await brute_force_triplet_search(problem_statement, top_k=3,
- collections=["data_point_source_code", "data_point_text"])
+ retrieved_edges = await brute_force_triplet_search(
+ problem_statement, top_k=3, collections=["data_point_source_code", "data_point_text"]
+ )
retrieved_edges_str = retrieved_edges_to_string(retrieved_edges)
- prompt = "\n".join([
- problem_statement,
- "",
- PATCH_EXAMPLE,
- "",
- "These are the retrieved edges:",
- retrieved_edges_str
- ])
+ prompt = "\n".join(
+ [
+ problem_statement,
+ "",
+ PATCH_EXAMPLE,
+ "",
+ "These are the retrieved edges:",
+ retrieved_edges_str,
+ ]
+ )
llm_client = get_llm_client()
answer_prediction = await llm_client.acreate_structured_output(
@@ -94,10 +94,7 @@ async def get_preds(dataset, with_cognee=True):
model_name = "without_cognee"
pred_func = generate_patch_without_cognee
- futures = [
- (instance["instance_id"], pred_func(instance, llm_client))
- for instance in dataset
- ]
+ futures = [(instance["instance_id"], pred_func(instance, llm_client)) for instance in dataset]
model_patches = await asyncio.gather(*[x[1] for x in futures])
preds = [
@@ -113,9 +110,8 @@ async def get_preds(dataset, with_cognee=True):
async def main():
- parser = argparse.ArgumentParser(
- description="Run LLM predictions on SWE-bench dataset")
- parser.add_argument('--cognee_off', action='store_true')
+ parser = argparse.ArgumentParser(description="Run LLM predictions on SWE-bench dataset")
+ parser.add_argument("--cognee_off", action="store_true")
parser.add_argument("--max_workers", type=int, required=True)
args = parser.parse_args()
@@ -123,17 +119,16 @@ async def main():
check_install_package(dependency)
if args.cognee_off:
- dataset_name = 'princeton-nlp/SWE-bench_Lite_bm25_13K'
- dataset = load_swebench_dataset(dataset_name, split='test')
+ dataset_name = "princeton-nlp/SWE-bench_Lite_bm25_13K"
+ dataset = load_swebench_dataset(dataset_name, split="test")
predictions_path = "preds_nocognee.json"
if not Path(predictions_path).exists():
preds = await get_preds(dataset, with_cognee=False)
with open(predictions_path, "w") as file:
json.dump(preds, file)
else:
- dataset_name = 'princeton-nlp/SWE-bench_Lite'
- swe_dataset = load_swebench_dataset(
- dataset_name, split='test')[:1]
+ dataset_name = "princeton-nlp/SWE-bench_Lite"
+ swe_dataset = load_swebench_dataset(dataset_name, split="test")[:1]
predictions_path = "preds.json"
preds = await get_preds(swe_dataset, with_cognee=not args.cognee_off)
with open(predictions_path, "w") as file:
diff --git a/evals/eval_utils.py b/evals/eval_utils.py
index 26c4ec2b8..6241cd4dd 100644
--- a/evals/eval_utils.py
+++ b/evals/eval_utils.py
@@ -15,9 +15,9 @@ def download_github_repo(instance, output_dir):
Returns:
str: Path to the downloaded repository.
"""
- repo_owner_repo = instance['repo']
- base_commit = instance['base_commit']
- instance_id = instance['instance_id']
+ repo_owner_repo = instance["repo"]
+ base_commit = instance["base_commit"]
+ instance_id = instance["instance_id"]
repo_url = f"https://github.com/{repo_owner_repo}.git"
@@ -30,7 +30,6 @@ def download_github_repo(instance, output_dir):
else:
print(f"Repository already exists at {repo_path}.")
-
repo = Repo(repo_path)
repo.git.checkout(base_commit)
@@ -69,4 +68,4 @@ def retrieved_edges_to_string(retrieved_edges):
relationship_type = edge.attributes["relationship_type"]
edge_str = f"{node_to_string(edge.node1)} {relationship_type} {node_to_string(edge.node2)}"
edge_strings.append(edge_str)
- return "\n".join(edge_strings)
\ No newline at end of file
+ return "\n".join(edge_strings)
diff --git a/evals/generate_test_set.py b/evals/generate_test_set.py
index 58bf4222a..ccac8d0e9 100644
--- a/evals/generate_test_set.py
+++ b/evals/generate_test_set.py
@@ -3,6 +3,10 @@
import dotenv
from deepeval.test_case import LLMTestCase
+# import pytest
+# from deepeval import assert_test
+from deepeval.metrics import AnswerRelevancyMetric
+
dotenv.load_dotenv()
# synthesizer = Synthesizer()
@@ -22,7 +26,7 @@
dataset = EvaluationDataset()
dataset.generate_goldens_from_docs(
- document_paths=['natural_language_processing.txt', 'soldiers_home.pdf', 'trump.txt'],
+ document_paths=["natural_language_processing.txt", "soldiers_home.pdf", "trump.txt"],
max_goldens_per_document=10,
num_evolutions=5,
enable_breadth_evolve=True,
@@ -33,16 +37,9 @@
print(dataset)
-import pytest
-from deepeval import assert_test
-from deepeval.metrics import AnswerRelevancyMetric
-
-
answer_relevancy_metric = AnswerRelevancyMetric(threshold=0.5)
-from deepeval import evaluate
+# from deepeval import evaluate
# evaluate(dataset, [answer_relevancy_metric])
-
-
diff --git a/evals/official_hotpot_metrics.py b/evals/official_hotpot_metrics.py
index b598e90d3..70444f7be 100644
--- a/evals/official_hotpot_metrics.py
+++ b/evals/official_hotpot_metrics.py
@@ -11,16 +11,15 @@
def normalize_answer(s):
-
def remove_articles(text):
- return re.sub(r'\b(a|an|the)\b', ' ', text)
+ return re.sub(r"\b(a|an|the)\b", " ", text)
def white_space_fix(text):
- return ' '.join(text.split())
+ return " ".join(text.split())
def remove_punc(text):
exclude = set(string.punctuation)
- return ''.join(ch for ch in text if ch not in exclude)
+ return "".join(ch for ch in text if ch not in exclude)
def lower(text):
return text.lower()
@@ -34,9 +33,15 @@ def f1_score(prediction, ground_truth):
ZERO_METRIC = (0, 0, 0)
- if normalized_prediction in ['yes', 'no', 'noanswer'] and normalized_prediction != normalized_ground_truth:
+ if (
+ normalized_prediction in ["yes", "no", "noanswer"]
+ and normalized_prediction != normalized_ground_truth
+ ):
return ZERO_METRIC
- if normalized_ground_truth in ['yes', 'no', 'noanswer'] and normalized_prediction != normalized_ground_truth:
+ if (
+ normalized_ground_truth in ["yes", "no", "noanswer"]
+ and normalized_prediction != normalized_ground_truth
+ ):
return ZERO_METRIC
prediction_tokens = normalized_prediction.split()
@@ -52,17 +57,19 @@ def f1_score(prediction, ground_truth):
def exact_match_score(prediction, ground_truth):
- return (normalize_answer(prediction) == normalize_answer(ground_truth))
+ return normalize_answer(prediction) == normalize_answer(ground_truth)
+
def update_answer(metrics, prediction, gold):
em = exact_match_score(prediction, gold)
f1, prec, recall = f1_score(prediction, gold)
- metrics['em'] += float(em)
- metrics['f1'] += f1
- metrics['prec'] += prec
- metrics['recall'] += recall
+ metrics["em"] += float(em)
+ metrics["f1"] += f1
+ metrics["prec"] += prec
+ metrics["recall"] += recall
return em, prec, recall
+
def update_sp(metrics, prediction, gold):
cur_sp_pred = set(map(tuple, prediction))
gold_sp_pred = set(map(tuple, gold))
@@ -79,8 +86,8 @@ def update_sp(metrics, prediction, gold):
recall = 1.0 * tp / (tp + fn) if tp + fn > 0 else 0.0
f1 = 2 * prec * recall / (prec + recall) if prec + recall > 0 else 0.0
em = 1.0 if fp + fn == 0 else 0.0
- metrics['sp_em'] += em
- metrics['sp_f1'] += f1
- metrics['sp_prec'] += prec
- metrics['sp_recall'] += recall
+ metrics["sp_em"] += em
+ metrics["sp_f1"] += f1
+ metrics["sp_prec"] += prec
+ metrics["sp_recall"] += recall
return em, prec, recall
diff --git a/evals/simple_rag_vs_cognee_eval.py b/evals/simple_rag_vs_cognee_eval.py
index 5d1ddb4ca..82b1df600 100644
--- a/evals/simple_rag_vs_cognee_eval.py
+++ b/evals/simple_rag_vs_cognee_eval.py
@@ -1,13 +1,19 @@
from deepeval.dataset import EvaluationDataset
from pydantic import BaseModel
-
+import os
from typing import List, Type
from deepeval.test_case import LLMTestCase
import dotenv
+from cognee.infrastructure.llm.get_llm_client import get_llm_client
+from cognee.infrastructure.databases.vector import get_vector_engine
+from cognee.base_config import get_base_config
+
+import logging
+
+logger = logging.getLogger(__name__)
dotenv.load_dotenv()
-from cognee.infrastructure.llm.get_llm_client import get_llm_client
dataset = EvaluationDataset()
dataset.add_test_cases_from_json_file(
@@ -16,7 +22,7 @@
input_key_name="input",
actual_output_key_name="actual_output",
expected_output_key_name="expected_output",
- context_key_name="context"
+ context_key_name="context",
)
print(dataset)
@@ -38,32 +44,26 @@
print(dataset)
+class AnswerModel(BaseModel):
+ response: str
-import logging
-
-logger = logging.getLogger(__name__)
-
-class AnswerModel(BaseModel):
- response:str
-def get_answer_base(content: str, context:str, response_model: Type[BaseModel]):
+def get_answer_base(content: str, context: str, response_model: Type[BaseModel]):
llm_client = get_llm_client()
system_prompt = "THIS IS YOUR CONTEXT:" + str(context)
- return llm_client.create_structured_output(content, system_prompt, response_model)
-def get_answer(content: str,context, model: Type[BaseModel]= AnswerModel):
+ return llm_client.create_structured_output(content, system_prompt, response_model)
+
+def get_answer(content: str, context, model: Type[BaseModel] = AnswerModel):
try:
- return (get_answer_base(
- content,
- context,
- model
- ))
+ return get_answer_base(content, context, model)
except Exception as error:
- logger.error("Error extracting cognitive layers from content: %s", error, exc_info = True)
+ logger.error("Error extracting cognitive layers from content: %s", error, exc_info=True)
raise error
+
async def run_cognify_base_rag():
from cognee.api.v1.add import add
from cognee.api.v1.prune import prune
@@ -74,17 +74,10 @@ async def run_cognify_base_rag():
await add("data://test_datasets", "initial_test")
graph = await cognify("initial_test")
+ return graph
-
- pass
-
-
-import os
-from cognee.base_config import get_base_config
-from cognee.infrastructure.databases.vector import get_vector_engine
-
-async def cognify_search_base_rag(content:str, context:str):
+async def cognify_search_base_rag(content: str, context: str):
base_config = get_base_config()
cognee_directory_path = os.path.abspath(".cognee_system")
@@ -97,23 +90,24 @@ async def cognify_search_base_rag(content:str, context:str):
print("results", return_)
return return_
-async def cognify_search_graph(content:str, context:str):
+
+async def cognify_search_graph(content: str, context: str):
from cognee.api.v1.search import search, SearchType
- params = {'query': 'Donald Trump'}
+
+ params = {"query": "Donald Trump"}
results = await search(SearchType.INSIGHTS, params)
print("results", results)
return results
-
def convert_goldens_to_test_cases(test_cases_raw: List[LLMTestCase]) -> List[LLMTestCase]:
test_cases = []
for case in test_cases_raw:
test_case = LLMTestCase(
input=case.input,
# Generate actual output using the 'input' and 'additional_metadata'
- actual_output= str(get_answer(case.input, case.context).model_dump()['response']),
+ actual_output=str(get_answer(case.input, case.context).model_dump()["response"]),
expected_output=case.expected_output,
context=case.context,
retrieval_context=["retrieval_context"],
@@ -121,6 +115,7 @@ def convert_goldens_to_test_cases(test_cases_raw: List[LLMTestCase]) -> List[LLM
test_cases.append(test_case)
return test_cases
+
# # Data preprocessing before setting the dataset test cases
# dataset.test_cases = convert_goldens_to_test_cases(dataset.test_cases)
#
@@ -133,13 +128,13 @@ def convert_goldens_to_test_cases(test_cases_raw: List[LLMTestCase]) -> List[LLM
if __name__ == "__main__":
-
import asyncio
async def main():
# await run_cognify_base_rag()
# await cognify_search_base_rag("show_all_processes", "context")
await cognify_search_graph("show_all_processes", "context")
+
asyncio.run(main())
# run_cognify_base_rag_and_search()
# # Data preprocessing before setting the dataset test cases
@@ -147,4 +142,4 @@ async def main():
# from deepeval.metrics import HallucinationMetric
# metric = HallucinationMetric()
# dataset.evaluate([metric])
- pass
\ No newline at end of file
+ pass
diff --git a/examples/python/code_graph_example.py b/examples/python/code_graph_example.py
index 44ab33aad..c90a0b606 100644
--- a/examples/python/code_graph_example.py
+++ b/examples/python/code_graph_example.py
@@ -8,9 +8,15 @@ async def main(repo_path, include_docs):
async for result in run_code_graph_pipeline(repo_path, include_docs):
print(result)
+
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("--repo_path", type=str, required=True, help="Path to the repository")
- parser.add_argument("--include_docs", type=lambda x: x.lower() in ("true", "1"), default=True, help="Whether or not to process non-code files")
+ parser.add_argument(
+ "--include_docs",
+ type=lambda x: x.lower() in ("true", "1"),
+ default=True,
+ help="Whether or not to process non-code files",
+ )
args = parser.parse_args()
- asyncio.run(main(args.repo_path, args.include_docs))
\ No newline at end of file
+ asyncio.run(main(args.repo_path, args.include_docs))
diff --git a/examples/python/dynamic_steps_example.py b/examples/python/dynamic_steps_example.py
index 7c0af8f0c..165d907d3 100644
--- a/examples/python/dynamic_steps_example.py
+++ b/examples/python/dynamic_steps_example.py
@@ -160,6 +160,7 @@
Negotiation and Relationship Building
"""
+
async def main(enable_steps):
# Step 1: Reset data and system state
if enable_steps.get("prune_data"):
@@ -184,10 +185,13 @@ async def main(enable_steps):
# Step 4: Query insights
if enable_steps.get("retriever"):
- results = await brute_force_triplet_search('Who has the most experience with graphic design?')
+ results = await brute_force_triplet_search(
+ "Who has the most experience with graphic design?"
+ )
print(format_triplets(results))
-if __name__ == '__main__':
+
+if __name__ == "__main__":
setup_logging(logging.ERROR)
rebuild_kg = True
@@ -197,7 +201,7 @@ async def main(enable_steps):
"prune_system": rebuild_kg,
"add_text": rebuild_kg,
"cognify": rebuild_kg,
- "retriever": retrieve
+ "retriever": retrieve,
}
asyncio.run(main(steps_to_enable))
diff --git a/examples/python/graphiti_example.py b/examples/python/graphiti_example.py
index a48fa9b4c..248361321 100644
--- a/examples/python/graphiti_example.py
+++ b/examples/python/graphiti_example.py
@@ -4,7 +4,9 @@
from cognee.api.v1.search import SearchType
from cognee.modules.pipelines import Task, run_tasks
from cognee.tasks.temporal_awareness import (
- build_graph_with_temporal_awareness, search_graph_with_temporal_awareness)
+ build_graph_with_temporal_awareness,
+ search_graph_with_temporal_awareness,
+)
text_list = [
"Kamala Harris is the Attorney General of California. She was previously "
@@ -12,18 +14,20 @@
"As AG, Harris was in office from January 3, 2011 – January 3, 2017",
]
-async def main():
+async def main():
tasks = [
Task(build_graph_with_temporal_awareness, text_list=text_list),
- Task(search_graph_with_temporal_awareness, query='Who was the California Attorney General?')
+ Task(
+ search_graph_with_temporal_awareness, query="Who was the California Attorney General?"
+ ),
]
-
+
pipeline = run_tasks(tasks)
-
+
async for result in pipeline:
print(result)
-if __name__ == '__main__':
+if __name__ == "__main__":
asyncio.run(main())
diff --git a/examples/python/simple_example.py b/examples/python/simple_example.py
index 55b07c4c3..bf3c95de1 100644
--- a/examples/python/simple_example.py
+++ b/examples/python/simple_example.py
@@ -20,47 +20,50 @@ async def main():
Natural language processing (NLP) is an interdisciplinary
subfield of computer science and information retrieval.
"""
-
+
print("Adding text to cognee:")
- print(text.strip())
+ print(text.strip())
# Add the text, and make it available for cognify
await cognee.add(text)
print("Text added successfully.\n")
-
print("Running cognify to create knowledge graph...\n")
print("Cognify process steps:")
print("1. Classifying the document: Determining the type and category of the input text.")
- print("2. Checking permissions: Ensuring the user has the necessary rights to process the text.")
- print("3. Extracting text chunks: Breaking down the text into sentences or phrases for analysis.")
+ print(
+ "2. Checking permissions: Ensuring the user has the necessary rights to process the text."
+ )
+ print(
+ "3. Extracting text chunks: Breaking down the text into sentences or phrases for analysis."
+ )
print("4. Adding data points: Storing the extracted chunks for processing.")
- print("5. Generating knowledge graph: Extracting entities and relationships to form a knowledge graph.")
+ print(
+ "5. Generating knowledge graph: Extracting entities and relationships to form a knowledge graph."
+ )
print("6. Summarizing text: Creating concise summaries of the content for quick insights.\n")
-
+
# Use LLMs and cognee to create knowledge graph
await cognee.cognify()
print("Cognify process complete.\n")
-
- query_text = 'Tell me about NLP'
+ query_text = "Tell me about NLP"
print(f"Searching cognee for insights with query: '{query_text}'")
# Query cognee for insights on the added text
- search_results = await cognee.search(
- SearchType.INSIGHTS, query_text=query_text
- )
-
+ search_results = await cognee.search(SearchType.INSIGHTS, query_text=query_text)
+
print("Search results:")
# Display results
for result_text in search_results:
print(result_text)
# Example output:
- # ({'id': UUID('bc338a39-64d6-549a-acec-da60846dd90d'), 'updated_at': datetime.datetime(2024, 11, 21, 12, 23, 1, 211808, tzinfo=datetime.timezone.utc), 'name': 'natural language processing', 'description': 'An interdisciplinary subfield of computer science and information retrieval.'}, {'relationship_name': 'is_a_subfield_of', 'source_node_id': UUID('bc338a39-64d6-549a-acec-da60846dd90d'), 'target_node_id': UUID('6218dbab-eb6a-5759-a864-b3419755ffe0'), 'updated_at': datetime.datetime(2024, 11, 21, 12, 23, 15, 473137, tzinfo=datetime.timezone.utc)}, {'id': UUID('6218dbab-eb6a-5759-a864-b3419755ffe0'), 'updated_at': datetime.datetime(2024, 11, 21, 12, 23, 1, 211808, tzinfo=datetime.timezone.utc), 'name': 'computer science', 'description': 'The study of computation and information processing.'})
- # (...)
- # It represents nodes and relationships in the knowledge graph:
- # - The first element is the source node (e.g., 'natural language processing').
- # - The second element is the relationship between nodes (e.g., 'is_a_subfield_of').
- # - The third element is the target node (e.g., 'computer science').
+ # ({'id': UUID('bc338a39-64d6-549a-acec-da60846dd90d'), 'updated_at': datetime.datetime(2024, 11, 21, 12, 23, 1, 211808, tzinfo=datetime.timezone.utc), 'name': 'natural language processing', 'description': 'An interdisciplinary subfield of computer science and information retrieval.'}, {'relationship_name': 'is_a_subfield_of', 'source_node_id': UUID('bc338a39-64d6-549a-acec-da60846dd90d'), 'target_node_id': UUID('6218dbab-eb6a-5759-a864-b3419755ffe0'), 'updated_at': datetime.datetime(2024, 11, 21, 12, 23, 15, 473137, tzinfo=datetime.timezone.utc)}, {'id': UUID('6218dbab-eb6a-5759-a864-b3419755ffe0'), 'updated_at': datetime.datetime(2024, 11, 21, 12, 23, 1, 211808, tzinfo=datetime.timezone.utc), 'name': 'computer science', 'description': 'The study of computation and information processing.'})
+ # (...)
+ # It represents nodes and relationships in the knowledge graph:
+ # - The first element is the source node (e.g., 'natural language processing').
+ # - The second element is the relationship between nodes (e.g., 'is_a_subfield_of').
+ # - The third element is the target node (e.g., 'computer science').
+
-if __name__ == '__main__':
+if __name__ == "__main__":
asyncio.run(main())
diff --git a/notebooks/cognee_code_graph_demo.ipynb b/notebooks/cognee_code_graph_demo.ipynb
index ef22afe58..d5bde67ef 100644
--- a/notebooks/cognee_code_graph_demo.ipynb
+++ b/notebooks/cognee_code_graph_demo.ipynb
@@ -7,8 +7,9 @@
"outputs": [],
"source": [
"import os\n",
- "os.environ['GRAPHISTRY_USERNAME'] = input(\"Please enter your graphistry username\")\n",
- "os.environ['GRAPHISTRY_PASSWORD'] = input(\"Please enter your graphistry password\")"
+ "\n",
+ "os.environ[\"GRAPHISTRY_USERNAME\"] = input(\"Please enter your graphistry username\")\n",
+ "os.environ[\"GRAPHISTRY_PASSWORD\"] = input(\"Please enter your graphistry password\")"
]
},
{
@@ -23,13 +24,17 @@
"from cognee.infrastructure.databases.relational import create_db_and_tables\n",
"\n",
"notebook_path = os.path.abspath(\"\")\n",
- "data_directory_path = str(pathlib.Path(os.path.join(notebook_path, \".data_storage/code_graph\")).resolve())\n",
+ "data_directory_path = str(\n",
+ " pathlib.Path(os.path.join(notebook_path, \".data_storage/code_graph\")).resolve()\n",
+ ")\n",
"cognee.config.data_root_directory(data_directory_path)\n",
- "cognee_directory_path = str(pathlib.Path(os.path.join(notebook_path, \".cognee_system/code_graph\")).resolve())\n",
+ "cognee_directory_path = str(\n",
+ " pathlib.Path(os.path.join(notebook_path, \".cognee_system/code_graph\")).resolve()\n",
+ ")\n",
"cognee.config.system_root_directory(cognee_directory_path)\n",
"\n",
"await cognee.prune.prune_data()\n",
- "await cognee.prune.prune_system(metadata = True)\n",
+ "await cognee.prune.prune_system(metadata=True)\n",
"\n",
"await create_db_and_tables()"
]
@@ -53,8 +58,8 @@
"git.Repo.clone_from(\n",
" \"git@github.com:microsoft/graphrag.git\",\n",
" Path(repo_clone_location),\n",
- " branch = \"main\",\n",
- " single_branch = True,\n",
+ " branch=\"main\",\n",
+ " single_branch=True,\n",
")"
]
},
@@ -64,16 +69,20 @@
"metadata": {},
"outputs": [],
"source": [
- "from cognee.tasks.repo_processor import enrich_dependency_graph, expand_dependency_graph, get_repo_file_dependencies\n",
+ "from cognee.tasks.repo_processor import (\n",
+ " enrich_dependency_graph,\n",
+ " expand_dependency_graph,\n",
+ " get_repo_file_dependencies,\n",
+ ")\n",
"from cognee.tasks.storage import add_data_points\n",
"from cognee.modules.pipelines.tasks.Task import Task\n",
"\n",
"tasks = [\n",
" Task(get_repo_file_dependencies),\n",
- " Task(add_data_points, task_config = { \"batch_size\": 50 }),\n",
- " Task(enrich_dependency_graph, task_config = { \"batch_size\": 50 }),\n",
- " Task(expand_dependency_graph, task_config = { \"batch_size\": 50 }),\n",
- " Task(add_data_points, task_config = { \"batch_size\": 50 }),\n",
+ " Task(add_data_points, task_config={\"batch_size\": 50}),\n",
+ " Task(enrich_dependency_graph, task_config={\"batch_size\": 50}),\n",
+ " Task(expand_dependency_graph, task_config={\"batch_size\": 50}),\n",
+ " Task(add_data_points, task_config={\"batch_size\": 50}),\n",
"]"
]
},
@@ -101,19 +110,20 @@
"outputs": [],
"source": [
"from cognee.shared.utils import render_graph\n",
- "await render_graph(None, include_nodes = True, include_labels = True)"
+ "\n",
+ "await render_graph(None, include_nodes=True, include_labels=True)"
]
},
{
- "metadata": {},
"cell_type": "markdown",
+ "metadata": {},
"source": "# Let's check the evaluations"
},
{
- "metadata": {},
"cell_type": "code",
- "outputs": [],
"execution_count": null,
+ "metadata": {},
+ "outputs": [],
"source": [
"from evals.eval_on_hotpot import eval_on_hotpotQA\n",
"from evals.eval_on_hotpot import answer_with_cognee\n",
@@ -133,8 +143,8 @@
"metadata": {},
"outputs": [],
"source": [
- "answer_provider = answer_with_cognee # For native LLM answers use answer_without_cognee\n",
- "num_samples = 10 # With cognee, it takes ~1m10s per sample\n",
+ "answer_provider = answer_with_cognee # For native LLM answers use answer_without_cognee\n",
+ "num_samples = 10 # With cognee, it takes ~1m10s per sample\n",
"\n",
"base_config = get_base_config()\n",
"data_root_dir = base_config.data_root_directory\n",
@@ -144,7 +154,7 @@
"\n",
"filepath = data_root_dir / Path(\"hotpot_dev_fullwiki_v1.json\")\n",
"if not filepath.exists():\n",
- " url = 'http://curtis.ml.cmu.edu/datasets/hotpot/hotpot_dev_fullwiki_v1.json'\n",
+ " url = \"http://curtis.ml.cmu.edu/datasets/hotpot/hotpot_dev_fullwiki_v1.json\"\n",
" wget.download(url, out=data_root_dir)\n",
"\n",
"with open(filepath, \"r\") as file:\n",
diff --git a/notebooks/cognee_demo.ipynb b/notebooks/cognee_demo.ipynb
index de8c07cfb..f59f6338c 100644
--- a/notebooks/cognee_demo.ipynb
+++ b/notebooks/cognee_demo.ipynb
@@ -76,7 +76,7 @@
},
{
"cell_type": "markdown",
- "id": "b6a98710-a14b-4a14-bb56-d3ae055e94d9",
+ "id": "1bf1fa3631dc03ed",
"metadata": {},
"source": [
"#### The problem lies in the nature of the search. If you just find some keywords, and return one or many documents from vectorstore this way, you will have an issue with the the way you would use to organise and prioritise documents. \n"
@@ -92,7 +92,7 @@
},
{
"cell_type": "markdown",
- "id": "b6a98710-a14b-4a14-bb56-d3ae055e94d9",
+ "id": "390e0d0805096f80",
"metadata": {},
"source": [
"## Semantic similarity search is not magic\n",
@@ -265,15 +265,13 @@
},
{
"cell_type": "code",
- "execution_count": null,
"id": "df16431d0f48b006",
"metadata": {
"ExecuteTime": {
- "end_time": "2024-09-20T14:02:48.519686Z",
- "start_time": "2024-09-20T14:02:48.515589Z"
+ "end_time": "2024-12-24T11:53:59.351640Z",
+ "start_time": "2024-12-24T11:53:59.347420Z"
}
},
- "outputs": [],
"source": [
"job_position = \"\"\"Senior Data Scientist (Machine Learning)\n",
"\n",
@@ -300,19 +298,19 @@
"Strong problem-solving skills and attention to detail.\n",
"Candidate CVs\n",
"\"\"\"\n"
- ]
+ ],
+ "outputs": [],
+ "execution_count": 1
},
{
"cell_type": "code",
- "execution_count": null,
"id": "9086abf3af077ab4",
"metadata": {
"ExecuteTime": {
- "end_time": "2024-09-20T14:02:49.120838Z",
- "start_time": "2024-09-20T14:02:49.118294Z"
+ "end_time": "2024-12-24T11:53:59.365410Z",
+ "start_time": "2024-12-24T11:53:59.363662Z"
}
},
- "outputs": [],
"source": [
"job_1 = \"\"\"\n",
"CV 1: Relevant\n",
@@ -345,19 +343,19 @@
"Big Data Technologies: Hadoop, Spark\n",
"Data Visualization: Tableau, Matplotlib\n",
"\"\"\""
- ]
+ ],
+ "outputs": [],
+ "execution_count": 2
},
{
"cell_type": "code",
- "execution_count": null,
"id": "a9de0cc07f798b7f",
"metadata": {
"ExecuteTime": {
- "end_time": "2024-09-20T14:02:49.675003Z",
- "start_time": "2024-09-20T14:02:49.671615Z"
+ "end_time": "2024-12-24T11:53:59.372957Z",
+ "start_time": "2024-12-24T11:53:59.371152Z"
}
},
- "outputs": [],
"source": [
"job_2 = \"\"\"\n",
"CV 2: Relevant\n",
@@ -389,19 +387,19 @@
"Data Visualization: Seaborn, Plotly\n",
"Databases: MySQL, MongoDB\n",
"\"\"\""
- ]
+ ],
+ "outputs": [],
+ "execution_count": 3
},
{
"cell_type": "code",
- "execution_count": null,
"id": "185ff1c102d06111",
"metadata": {
"ExecuteTime": {
- "end_time": "2024-09-20T14:02:50.286828Z",
- "start_time": "2024-09-20T14:02:50.284369Z"
+ "end_time": "2024-12-24T11:53:59.961140Z",
+ "start_time": "2024-12-24T11:53:59.959103Z"
}
},
- "outputs": [],
"source": [
"job_3 = \"\"\"\n",
"CV 3: Relevant\n",
@@ -433,19 +431,19 @@
"Statistical Analysis: SAS, SPSS\n",
"Cloud Platforms: AWS, Azure\n",
"\"\"\""
- ]
+ ],
+ "outputs": [],
+ "execution_count": 4
},
{
"cell_type": "code",
- "execution_count": null,
"id": "d55ce4c58f8efb67",
"metadata": {
"ExecuteTime": {
- "end_time": "2024-09-20T14:02:50.950343Z",
- "start_time": "2024-09-20T14:02:50.946378Z"
+ "end_time": "2024-12-24T11:54:00.656495Z",
+ "start_time": "2024-12-24T11:54:00.654716Z"
}
},
- "outputs": [],
"source": [
"job_4 = \"\"\"\n",
"CV 4: Not Relevant\n",
@@ -475,19 +473,19 @@
"Web Design: HTML, CSS\n",
"Specialties: Branding and Identity, Typography\n",
"\"\"\""
- ]
+ ],
+ "outputs": [],
+ "execution_count": 5
},
{
"cell_type": "code",
- "execution_count": null,
"id": "ca4ecc32721ad332",
"metadata": {
"ExecuteTime": {
- "end_time": "2024-09-20T14:02:51.548191Z",
- "start_time": "2024-09-20T14:02:51.545520Z"
+ "end_time": "2024-12-24T11:54:01.184899Z",
+ "start_time": "2024-12-24T11:54:01.183028Z"
}
},
- "outputs": [],
"source": [
"job_5 = \"\"\"\n",
"CV 5: Not Relevant\n",
@@ -517,7 +515,9 @@
"CRM Software: Salesforce, Zoho\n",
"Negotiation and Relationship Building\n",
"\"\"\""
- ]
+ ],
+ "outputs": [],
+ "execution_count": 6
},
{
"cell_type": "markdown",
@@ -529,10 +529,13 @@
},
{
"cell_type": "code",
- "execution_count": null,
"id": "bce39dc6",
- "metadata": {},
- "outputs": [],
+ "metadata": {
+ "ExecuteTime": {
+ "end_time": "2024-12-24T11:54:04.417315Z",
+ "start_time": "2024-12-24T11:54:04.414132Z"
+ }
+ },
"source": [
"import os\n",
"\n",
@@ -570,14 +573,19 @@
"# os.environ[\"DB_PORT\"]=\"5432\"\n",
"# os.environ[\"DB_USERNAME\"]=\"cognee\"\n",
"# os.environ[\"DB_PASSWORD\"]=\"cognee\""
- ]
+ ],
+ "outputs": [],
+ "execution_count": 7
},
{
"cell_type": "code",
- "execution_count": null,
"id": "9f1a1dbd",
- "metadata": {},
- "outputs": [],
+ "metadata": {
+ "ExecuteTime": {
+ "end_time": "2024-12-24T11:54:16.672999Z",
+ "start_time": "2024-12-24T11:54:07.425202Z"
+ }
+ },
"source": [
"# Reset the cognee system with the following command:\n",
"\n",
@@ -585,7 +593,32 @@
"\n",
"await cognee.prune.prune_data()\n",
"await cognee.prune.prune_system(metadata=True)"
- ]
+ ],
+ "outputs": [
+ {
+ "name": "stderr",
+ "output_type": "stream",
+ "text": [
+ "File /Users/vasilije/cognee/cognee/.cognee_system/databases/cognee_graph.pkl not found. Initializing an empty graph."
+ ]
+ },
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "Database deleted successfully.\n"
+ ]
+ },
+ {
+ "name": "stderr",
+ "output_type": "stream",
+ "text": [
+ "/Users/vasilije/cognee/.venv/lib/python3.11/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
+ " from .autonotebook import tqdm as notebook_tqdm\n"
+ ]
+ }
+ ],
+ "execution_count": 8
},
{
"cell_type": "markdown",
@@ -597,20 +630,48 @@
},
{
"cell_type": "code",
- "execution_count": null,
"id": "904df61ba484a8e5",
"metadata": {
"ExecuteTime": {
- "end_time": "2024-09-20T14:02:54.243987Z",
- "start_time": "2024-09-20T14:02:52.498195Z"
+ "end_time": "2024-12-24T11:54:28.313862Z",
+ "start_time": "2024-12-24T11:54:23.756587Z"
}
},
- "outputs": [],
"source": [
"import cognee\n",
"\n",
"await cognee.add([job_1, job_2, job_3, job_4, job_5, job_position], \"example\")"
- ]
+ ],
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "User df77f15b-a077-4c86-a3e4-c059bf4cacb9 has registered.\n"
+ ]
+ },
+ {
+ "name": "stderr",
+ "output_type": "stream",
+ "text": [
+ "/Users/vasilije/cognee/.venv/lib/python3.11/site-packages/dlt/destinations/impl/sqlalchemy/merge_job.py:194: SAWarning: Table 'file_metadata' already exists within the given MetaData - not copying.\n",
+ " staging_table_obj = table_obj.to_metadata(\n",
+ "/Users/vasilije/cognee/.venv/lib/python3.11/site-packages/dlt/destinations/impl/sqlalchemy/merge_job.py:229: SAWarning: implicitly coercing SELECT object to scalar subquery; please use the .scalar_subquery() method to produce a scalar subquery.\n",
+ " order_by=order_dir_func(order_by_col),\n"
+ ]
+ },
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "Pipeline file_load_from_filesystem load step completed in 0.03 seconds\n",
+ "1 load package(s) were loaded to destination sqlalchemy and into dataset main\n",
+ "The sqlalchemy destination used sqlite:////Users/vasilije/cognee/cognee/.cognee_system/databases/cognee_db location to store data\n",
+ "Load package 1735041267.4777632 is LOADED and contains no failed jobs\n"
+ ]
+ }
+ ],
+ "execution_count": 9
},
{
"cell_type": "markdown",
@@ -622,15 +683,13 @@
},
{
"cell_type": "code",
- "execution_count": null,
"id": "7c431fdef4921ae0",
"metadata": {
"ExecuteTime": {
- "end_time": "2024-09-20T14:02:57.925667Z",
- "start_time": "2024-09-20T14:02:57.922353Z"
+ "end_time": "2024-12-24T11:54:44.728010Z",
+ "start_time": "2024-12-24T11:54:44.723877Z"
}
},
- "outputs": [],
"source": [
"from cognee.shared.data_models import KnowledgeGraph\n",
"from cognee.modules.data.models import Dataset, Data\n",
@@ -669,19 +728,19 @@
" print(result)\n",
" except Exception as error:\n",
" raise error\n"
- ]
+ ],
+ "outputs": [],
+ "execution_count": 10
},
{
"cell_type": "code",
- "execution_count": null,
"id": "f0a91b99c6215e09",
"metadata": {
"ExecuteTime": {
- "end_time": "2024-09-20T14:02:58.905774Z",
- "start_time": "2024-09-20T14:02:58.625915Z"
+ "end_time": "2024-12-24T11:55:26.027386Z",
+ "start_time": "2024-12-24T11:54:47.384342Z"
}
},
- "outputs": [],
"source": [
"from cognee.modules.users.methods import get_default_user\n",
"from cognee.modules.data.methods import get_datasets_by_name\n",
@@ -691,7 +750,27 @@
"datasets = await get_datasets_by_name([\"example\"], user.id)\n",
"\n",
"await run_cognify_pipeline(datasets[0], user)"
- ]
+ ],
+ "outputs": [
+ {
+ "name": "stderr",
+ "output_type": "stream",
+ "text": [
+ "File /Users/vasilije/cognee/cognee/.cognee_system/databases/cognee_graph.pkl not found. Initializing an empty graph./Users/vasilije/cognee/.venv/lib/python3.11/site-packages/pydantic/main.py:1522: RuntimeWarning: fields may not start with an underscore, ignoring \"_metadata\"\n",
+ " warnings.warn(f'fields may not start with an underscore, ignoring \"{f_name}\"', RuntimeWarning)\n",
+ "/Users/vasilije/cognee/.venv/lib/python3.11/site-packages/pydantic/main.py:1522: RuntimeWarning: fields may not start with an underscore, ignoring \"__tablename__\"\n",
+ " warnings.warn(f'fields may not start with an underscore, ignoring \"{f_name}\"', RuntimeWarning)\n"
+ ]
+ },
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "[TextSummary(id=UUID('92b5d0a7-f980-529d-bb5b-48e72825a01a'), updated_at=datetime.datetime(2024, 12, 24, 11, 54, 13, 481297, tzinfo=datetime.timezone.utc), topological_rank=0, text='Experienced Senior Data Scientist with expertise in machine learning and predictive modeling, demonstrating over 8 years in the field.', made_from=DocumentChunk(id=UUID('70b823e2-5b12-57b5-ad8d-798e1d721f8e'), updated_at=datetime.datetime(2024, 12, 24, 11, 54, 13, 481297, tzinfo=datetime.timezone.utc), topological_rank=0, text='\\nCV 1: Relevant\\nName: Dr. Emily Carter\\nContact Information:\\n\\nEmail: emily.carter@example.com\\nPhone: (555) 123-4567\\nSummary:\\n\\nSenior Data Scientist with over 8 years of experience in machine learning and predictive analytics. Expertise in developing advanced algorithms and deploying scalable models in production environments.\\n\\nEducation:\\n\\nPh.D. in Computer Science, Stanford University (2014)\\nB.S. in Mathematics, University of California, Berkeley (2010)\\nExperience:\\n\\nSenior Data Scientist, InnovateAI Labs (2016 – Present)\\nLed a team in developing machine learning models for natural language processing applications.\\nImplemented deep learning algorithms that improved prediction accuracy by 25%.\\nCollaborated with cross-functional teams to integrate models into cloud-based platforms.\\nData Scientist, DataWave Analytics (2014 – 2016)\\nDeveloped predictive models for customer segmentation and churn analysis.\\nAnalyzed large datasets using Hadoop and Spark frameworks.\\nSkills:\\n\\nProgramming Languages: Python, R, SQL\\nMachin e Learning: TensorFlow, Keras, Scikit-Learn\\nBig Data Technologies: Hadoop, Spark\\nData Visualization: Tableau, Matplotlib\\n', word_count=133, chunk_index=0, cut_type='sentence_cut', is_part_of=TextDocument(id=UUID('11a2b08a-c160-5961-80b7-b3498eafa973'), updated_at=datetime.datetime(2024, 12, 24, 11, 54, 13, 481297, tzinfo=datetime.timezone.utc), topological_rank=0, name='text_85410f4ad1197f5974aef9aed6f103c8', raw_data_location='/Users/vasilije/cognee/cognee/.data_storage/data/text_85410f4ad1197f5974aef9aed6f103c8.txt', metadata_id=UUID('11a2b08a-c160-5961-80b7-b3498eafa973'), mime_type='text/plain', type='text'), contains=[Entity(id=UUID('29e771c8-4c3f-52de-9511-6b705878e130'), updated_at=datetime.datetime(2024, 12, 24, 11, 54, 13, 481297, tzinfo=datetime.timezone.utc), topological_rank=0, name='dr. emily carter', is_a=EntityType(id=UUID('d072ba0f-e1a9-58bf-9974-e1802adc8134'), updated_at=datetime.datetime(2024, 12, 24, 11, 54, 13, 481297, tzinfo=datetime.timezone.utc), topological_rank=0, name='person', description='person'), description='Senior Data Scientist with over 8 years of experience in machine learning and predictive analytics.'), Entity(id=UUID('ce8b394a-b30e-52fc-b80a-6352edc60e5b'), updated_at=datetime.datetime(2024, 12, 24, 11, 54, 13, 481297, tzinfo=datetime.timezone.utc), topological_rank=0, name='stanford university', is_a=EntityType(id=UUID('d3d7b6b4-9b0d-52e8-9e09-a9e9cf4b5a4d'), updated_at=datetime.datetime(2024, 12, 24, 11, 54, 13, 481297, tzinfo=datetime.timezone.utc), topological_rank=0, name='organization', description='organization'), description='Prestigious university located in Stanford, California.'), Entity(id=UUID('2c02c93c-9cd1-56b8-9cc0-55ff0b290e57'), updated_at=datetime.datetime(2024, 12, 24, 11, 54, 13, 481297, tzinfo=datetime.timezone.utc), topological_rank=0, name='university of california, berkeley', is_a=EntityType(id=UUID('d3d7b6b4-9b0d-52e8-9e09-a9e9cf4b5a4d'), updated_at=datetime.datetime(2024, 12, 24, 11, 54, 13, 481297, tzinfo=datetime.timezone.utc), topological_rank=0, name='organization', description='organization'), description='Public research university located in Berkeley, California.'), Entity(id=UUID('9780afb1-dccc-53eb-9a30-c0d4ce033711'), updated_at=datetime.datetime(2024, 12, 24, 11, 54, 13, 481297, tzinfo=datetime.timezone.utc), topological_rank=0, name='innovateai labs', is_a=EntityType(id=UUID('d3d7b6b4-9b0d-52e8-9e09-a9e9cf4b5a4d'), updated_at=datetime.datetime(2024, 12, 24, 11, 54, 13, 481297, tzinfo=datetime.timezone.utc), topological_rank=0, name='organization', description='organization'), description='A lab focused on artificial intelligence projects.'), Entity(id=UUID('50d0a685-5300-544f-b081-edca4b625886'), updated_at=datetime.datetime(2024, 12, 24, 11, 54, 13, 481297, tzinfo=datetime.timezone.utc), topological_rank=0, name='datawave analytics', is_a=EntityType(id=UUID('d3d7b6b4-9b0d-52e8-9e09-a9e9cf4b5a4d'), updated_at=datetime.datetime(2024, 12, 24, 11, 54, 13, 481297, tzinfo=datetime.timezone.utc), topological_rank=0, name='organization', description='organization'), description='Analytics firm specialized in data-driven insights.'), Entity(id=UUID('c95db510-e2ee-5a00-bded-20bbcb50c492'), updated_at=datetime.datetime(2024, 12, 24, 11, 54, 13, 481297, tzinfo=datetime.timezone.utc), topological_rank=0, name='python', is_a=EntityType(id=UUID('80d409bb-e431-5939-a1ad-3acd96267128'), updated_at=datetime.datetime(2024, 12, 24, 11, 54, 13, 481297, tzinfo=datetime.timezone.utc), topological_rank=0, name='programming language', description='programming language'), description='A high-level programming language used for general-purpose programming.'), Entity(id=UUID('39bd9707-8098-52ed-9cbf-bbdd26b963fb'), updated_at=datetime.datetime(2024, 12, 24, 11, 54, 13, 481297, tzinfo=datetime.timezone.utc), topological_rank=0, name='r', is_a=EntityType(id=UUID('80d409bb-e431-5939-a1ad-3acd96267128'), updated_at=datetime.datetime(2024, 12, 24, 11, 54, 13, 481297, tzinfo=datetime.timezone.utc), topological_rank=0, name='programming language', description='programming language'), description='A programming language and environment for statistical computing and graphics.'), Entity(id=UUID('1ff6821a-b207-5050-83e9-37ff67a27d03'), updated_at=datetime.datetime(2024, 12, 24, 11, 54, 13, 481297, tzinfo=datetime.timezone.utc), topological_rank=0, name='sql', is_a=EntityType(id=UUID('80d409bb-e431-5939-a1ad-3acd96267128'), updated_at=datetime.datetime(2024, 12, 24, 11, 54, 13, 481297, tzinfo=datetime.timezone.utc), topological_rank=0, name='programming language', description='programming language'), description='A domain-specific language used in programming and managing relational databases.'), Entity(id=UUID('6e72f6f5-0452-5d42-a4e8-4aba6a614cb1'), updated_at=datetime.datetime(2024, 12, 24, 11, 54, 13, 481297, tzinfo=datetime.timezone.utc), topological_rank=0, name='tensorflow', is_a=EntityType(id=UUID('9ffe9ce7-8938-5a5c-8d03-5f1a4c5210a1'), updated_at=datetime.datetime(2024, 12, 24, 11, 54, 13, 481297, tzinfo=datetime.timezone.utc), topological_rank=0, name='machine learning framework', description='machine learning framework'), description='An open-source software library for dataflow and differentiable programming across a range of tasks.'), Entity(id=UUID('ab85cdff-2a98-5c6d-99a3-df1f40f4ec16'), updated_at=datetime.datetime(2024, 12, 24, 11, 54, 13, 481297, tzinfo=datetime.timezone.utc), topological_rank=0, name='keras', is_a=EntityType(id=UUID('9ffe9ce7-8938-5a5c-8d03-5f1a4c5210a1'), updated_at=datetime.datetime(2024, 12, 24, 11, 54, 13, 481297, tzinfo=datetime.timezone.utc), topological_rank=0, name='machine learning framework', description='machine learning framework'), description='An open-source neural network library written in Python that runs on top of TensorFlow.'), Entity(id=UUID('37eecdcc-fb56-519c-bc18-d0d3afea0c0d'), updated_at=datetime.datetime(2024, 12, 24, 11, 54, 13, 481297, tzinfo=datetime.timezone.utc), topological_rank=0, name='scikit-learn', is_a=EntityType(id=UUID('9ffe9ce7-8938-5a5c-8d03-5f1a4c5210a1'), updated_at=datetime.datetime(2024, 12, 24, 11, 54, 13, 481297, tzinfo=datetime.timezone.utc), topological_rank=0, name='machine learning framework', description='machine learning framework'), description='A free software machine learning library for the Python programming language.'), Entity(id=UUID('f9a0eeca-c9ff-53b3-90eb-347254d7d7eb'), updated_at=datetime.datetime(2024, 12, 24, 11, 54, 13, 481297, tzinfo=datetime.timezone.utc), topological_rank=0, name='hadoop', is_a=EntityType(id=UUID('7c2287d0-16fc-53dc-86ce-8d8e61c8642c'), updated_at=datetime.datetime(2024, 12, 24, 11, 54, 13, 481297, tzinfo=datetime.timezone.utc), topological_rank=0, name='big data technology', description='big data technology'), description='An open-source framework for storing and processing large datasets in a distributed computing environment.'), Entity(id=UUID('46a235af-5ed5-5023-a4ec-c253e3f93031'), updated_at=datetime.datetime(2024, 12, 24, 11, 54, 13, 481297, tzinfo=datetime.timezone.utc), topological_rank=0, name='spark', is_a=EntityType(id=UUID('7c2287d0-16fc-53dc-86ce-8d8e61c8642c'), updated_at=datetime.datetime(2024, 12, 24, 11, 54, 13, 481297, tzinfo=datetime.timezone.utc), topological_rank=0, name='big data technology', description='big data technology'), description='An open-source unified analytics engine for large-scale data processing.'), Entity(id=UUID('c55004f3-8a6d-5130-b8bd-ed8278daa9a4'), updated_at=datetime.datetime(2024, 12, 24, 11, 54, 13, 481297, tzinfo=datetime.timezone.utc), topological_rank=0, name='tableau', is_a=EntityType(id=UUID('674cc5fa-7849-575a-917f-90b7b77f52b3'), updated_at=datetime.datetime(2024, 12, 24, 11, 54, 13, 481297, tzinfo=datetime.timezone.utc), topological_rank=0, name='data visualization software', description='data visualization software'), description='A visual analytics platform transforming the way we use data to solve problems.'), Entity(id=UUID('3c7adf8f-ef23-5330-a3fe-6a0b791cee2b'), updated_at=datetime.datetime(2024, 12, 24, 11, 54, 13, 481297, tzinfo=datetime.timezone.utc), topological_rank=0, name='matplotlib', is_a=EntityType(id=UUID('3f3619fc-ebd1-50ed-adde-cf94e8bb3c1b'), updated_at=datetime.datetime(2024, 12, 24, 11, 54, 13, 481297, tzinfo=datetime.timezone.utc), topological_rank=0, name='data visualization library', description='data visualization library'), description='A plotting library for the Python programming language and its numerical mathematics extension NumPy.')])), TextSummary(id=UUID('2f680bef-2edd-566e-b98c-78d549799e77'), updated_at=datetime.datetime(2024, 12, 24, 11, 54, 13, 481297, tzinfo=datetime.timezone.utc), topological_rank=0, text='Senior Data Scientist specializing in Machine Learning at TechNova Solutions', made_from=DocumentChunk(id=UUID('eb6617b8-c78c-519b-b765-1eefc2e3a0d7'), updated_at=datetime.datetime(2024, 12, 24, 11, 54, 13, 481297, tzinfo=datetime.timezone.utc), topological_rank=0, text='Senior Data Scientist (Machine Learning)\\n\\nCompany: TechNova Solutions\\nLocation: San Francisco, CA\\n\\nJob Description:\\n\\nTechNova Solutions is seeking a Senior Data Scientist specializing in Machine Learning to join our dynamic analytics team. The ideal candidate will have a strong background in developing and deploying machine learning models, working with large datasets, and translating complex data into actionable insights.\\n\\nResponsibilities:\\n\\nDevelop and implement advanced machine learning algorithms and models.\\nAnalyze large, complex datasets to extract meaningful patterns and insights.\\nCollaborate with cross-functional teams to integrate predictive models into products.\\nStay updated with the latest advancements in machine learning and data science.\\nMentor junior data scientists and provide technical guidance.\\nQualifications:\\n\\nMaster’s or Ph.D. in Data Science, Computer Science, Statistics, or a related field.\\n5+ years of experience in data science and machine learning.\\nProficient in Python, R, and SQL.\\nExpe rience with deep learning frameworks (e.g., TensorFlow, PyTorch).\\nStrong problem-solving skills and attention to detail.\\nCandidate CVs\\n', word_count=153, chunk_index=0, cut_type='sentence_cut', is_part_of=TextDocument(id=UUID('171f3035-4c37-5f7b-97c8-6b222404cc9a'), updated_at=datetime.datetime(2024, 12, 24, 11, 54, 13, 481297, tzinfo=datetime.timezone.utc), topological_rank=0, name='text_81a5a96a9a7325d40521ea453778ebe0', raw_data_location='/Users/vasilije/cognee/cognee/.data_storage/data/text_81a5a96a9a7325d40521ea453778ebe0.txt', metadata_id=UUID('171f3035-4c37-5f7b-97c8-6b222404cc9a'), mime_type='text/plain', type='text'), contains=[Entity(id=UUID('453a45c9-14e7-5b73-adb8-55991096fef0'), updated_at=datetime.datetime(2024, 12, 24, 11, 54, 13, 481297, tzinfo=datetime.timezone.utc), topological_rank=0, name='technova solutions', is_a=EntityType(id=UUID('a6ed6bf1-fe31-5dfe-8ab4-484691fdf219'), updated_at=datetime.datetime(2024, 12, 24, 11, 54, 13, 481297, tzinfo=datetime.timezone.utc), topological_rank=0, name='company', description='company'), description='A technology company specializing in data analytics and machine learning.'), Entity(id=UUID('435dbd37-ab20-503c-9e99-ab8b8a3484e5'), updated_at=datetime.datetime(2024, 12, 24, 11, 54, 13, 481297, tzinfo=datetime.timezone.utc), topological_rank=0, name='senior data scientist', is_a=EntityType(id=UUID('524c6bbb-1534-5a51-8068-18dd4ae171eb'), updated_at=datetime.datetime(2024, 12, 24, 11, 54, 13, 481297, tzinfo=datetime.timezone.utc), topological_rank=0, name='profession', description='profession'), description='A role focused on advanced data analysis and machine learning.'), Entity(id=UUID('198e2ab8-75e9-5931-97ab-da9a5a8e188c'), updated_at=datetime.datetime(2024, 12, 24, 11, 54, 13, 481297, tzinfo=datetime.timezone.utc), topological_rank=0, name='san francisco, ca', is_a=EntityType(id=UUID('19dd7d4d-a966-5ed5-82a0-6ae377761a29'), updated_at=datetime.datetime(2024, 12, 24, 11, 54, 13, 481297, tzinfo=datetime.timezone.utc), topological_rank=0, name='location', description='location'), description='A city in California, USA.'), Entity(id=UUID('5187986a-7305-5a63-b057-8f2c097419eb'), updated_at=datetime.datetime(2024, 12, 24, 11, 54, 13, 481297, tzinfo=datetime.timezone.utc), topological_rank=0, name='machine learning', is_a=EntityType(id=UUID('0198571b-3e94-50ea-8b9f-19e3a31080c0'), updated_at=datetime.datetime(2024, 12, 24, 11, 54, 13, 481297, tzinfo=datetime.timezone.utc), topological_rank=0, name='field', description='field'), description='A subset of artificial intelligence focused on the development of algorithms that enable computers to learn from data.'), Entity(id=UUID('d6545b21-153c-58ba-be47-46e5216521a3'), updated_at=datetime.datetime(2024, 12, 24, 11, 54, 13, 481297, tzinfo=datetime.timezone.utc), topological_rank=0, name='data science', is_a=EntityType(id=UUID('0198571b-3e94-50ea-8b9f-19e3a31080c0'), updated_at=datetime.datetime(2024, 12, 24, 11, 54, 13, 481297, tzinfo=datetime.timezone.utc), topological_rank=0, name='field', description='field'), description='A multidisciplinary field that uses scientific methods to extract knowledge and insights from data.'), Entity(id=UUID('c0d95499-de6b-5fcf-b0f5-9cbf427ad5c6'), updated_at=datetime.datetime(2024, 12, 24, 11, 54, 13, 481297, tzinfo=datetime.timezone.utc), topological_rank=0, name='pytorch', is_a=EntityType(id=UUID('36a32bd3-8880-547a-949b-8447477d1ef5'), updated_at=datetime.datetime(2024, 12, 24, 11, 54, 13, 481297, tzinfo=datetime.timezone.utc), topological_rank=0, name='framework', description='framework'), description='An open-source machine learning framework for deep learning.'), Entity(id=UUID('62b4dda1-de4a-5098-a56e-d3fe81f84dbc'), updated_at=datetime.datetime(2024, 12, 24, 11, 54, 13, 481297, tzinfo=datetime.timezone.utc), topological_rank=0, name='masters or ph.d. in data science', is_a=EntityType(id=UUID('a49b283a-ce92-50e0-b7fa-ca7c628eb01a'), updated_at=datetime.datetime(2024, 12, 24, 11, 54, 13, 481297, tzinfo=datetime.timezone.utc), topological_rank=0, name='degree', description='degree'), description='Advanced academic degree in data science or related fields.')])), TextSummary(id=UUID('5c988618-db52-5979-9cf8-db80c0098285'), updated_at=datetime.datetime(2024, 12, 24, 11, 54, 13, 481297, tzinfo=datetime.timezone.utc), topological_rank=0, text='Data Scientist with expertise in machine learning and statistical analysis, adept at managing extensive datasets and converting data into practical business solutions.', made_from=DocumentChunk(id=UUID('a6e82ac7-e791-5d6b-b4a9-f5e41cbe95bf'), updated_at=datetime.datetime(2024, 12, 24, 11, 54, 13, 481297, tzinfo=datetime.timezone.utc), topological_rank=0, text='\\nCV 2: Relevant\\nName: Michael Rodriguez\\nContact Information:\\n\\nEmail: michael.rodriguez@example.com\\nPhone: (555) 234-5678\\nSummary:\\n\\nData Scientist with a strong background in machine learning and statistical modeling. Skilled in handling large datasets and translating data into actionable business insights.\\n\\nEducation:\\n\\nM.S. in Data Science, Carnegie Mellon University (2013)\\nB.S. in Computer Science, University of Michigan (2011)\\nExperience:\\n\\nSenior Data Scientist, Alpha Analytics (2017 – Present)\\nDeveloped machine learning models to optimize marketing strategies.\\nReduced customer acquisition cost by 15% through predictive modeling.\\nData Scientist, TechInsights (2013 – 2017)\\nAnalyzed user behavior data to improve product features.\\nImplemented A/B testing frameworks to evaluate product changes.\\nSkills:\\n\\nProgramming Languages: Python, Java, SQL\\nMachine Learning: Scikit-Learn, XGBoost\\nData Visualization: Seaborn, Plotly\\nDatabases: MySQL, MongoDB\\n', word_count=108, chunk_index=0, cut_type='sentence_cut', is_part_of=TextDocument(id=UUID('1f078b0a-3cc1-57a9-9802-f78565d49f29'), updated_at=datetime.datetime(2024, 12, 24, 11, 54, 13, 481297, tzinfo=datetime.timezone.utc), topological_rank=0, name='text_f0f63a5c88dbbeef1eca23d32848220c', raw_data_location='/Users/vasilije/cognee/cognee/.data_storage/data/text_f0f63a5c88dbbeef1eca23d32848220c.txt', metadata_id=UUID('1f078b0a-3cc1-57a9-9802-f78565d49f29'), mime_type='text/plain', type='text'), contains=[Entity(id=UUID('73ae630f-7b09-5dce-8c18-45d0a57b30f9'), updated_at=datetime.datetime(2024, 12, 24, 11, 54, 13, 481297, tzinfo=datetime.timezone.utc), topological_rank=0, name='michael rodriguez', is_a=EntityType(id=UUID('d072ba0f-e1a9-58bf-9974-e1802adc8134'), updated_at=datetime.datetime(2024, 12, 24, 11, 54, 13, 481297, tzinfo=datetime.timezone.utc), topological_rank=0, name='person', description='person'), description='Data Scientist with a strong background in machine learning and statistical modeling.'), Entity(id=UUID('5534e0b0-d0c4-5ab9-82e9-91bed36f70bd'), updated_at=datetime.datetime(2024, 12, 24, 11, 54, 13, 481297, tzinfo=datetime.timezone.utc), topological_rank=0, name='carnegie mellon university', is_a=EntityType(id=UUID('912b273c-683d-53ea-8ffe-aadef0b84237'), updated_at=datetime.datetime(2024, 12, 24, 11, 54, 13, 481297, tzinfo=datetime.timezone.utc), topological_rank=0, name='educational institution', description='educational institution'), description='University known for its data science program.'), Entity(id=UUID('0af613e0-c11b-550d-ada2-2c2aa6550884'), updated_at=datetime.datetime(2024, 12, 24, 11, 54, 13, 481297, tzinfo=datetime.timezone.utc), topological_rank=0, name='university of michigan', is_a=EntityType(id=UUID('912b273c-683d-53ea-8ffe-aadef0b84237'), updated_at=datetime.datetime(2024, 12, 24, 11, 54, 13, 481297, tzinfo=datetime.timezone.utc), topological_rank=0, name='educational institution', description='educational institution'), description='University known for its computer science program.'), Entity(id=UUID('04a91fef-8a07-5d50-8f1b-46f3afeec497'), updated_at=datetime.datetime(2024, 12, 24, 11, 54, 13, 481297, tzinfo=datetime.timezone.utc), topological_rank=0, name='alpha analytics', is_a=EntityType(id=UUID('a6ed6bf1-fe31-5dfe-8ab4-484691fdf219'), updated_at=datetime.datetime(2024, 12, 24, 11, 54, 13, 481297, tzinfo=datetime.timezone.utc), topological_rank=0, name='company', description='company'), description='Company where Michael Rodriguez works as a Senior Data Scientist.'), Entity(id=UUID('3f848ed6-902f-5a8e-9577-cb67f8c17acd'), updated_at=datetime.datetime(2024, 12, 24, 11, 54, 13, 481297, tzinfo=datetime.timezone.utc), topological_rank=0, name='techinsights', is_a=EntityType(id=UUID('a6ed6bf1-fe31-5dfe-8ab4-484691fdf219'), updated_at=datetime.datetime(2024, 12, 24, 11, 54, 13, 481297, tzinfo=datetime.timezone.utc), topological_rank=0, name='company', description='company'), description='Company where Michael Rodriguez worked as a Data Scientist.')])), TextSummary(id=UUID('ee6cb607-27eb-5b87-bf2a-305721534263'), updated_at=datetime.datetime(2024, 12, 24, 11, 54, 13, 481297, tzinfo=datetime.timezone.utc), topological_rank=0, text='Sales Manager with proven ability to enhance revenue and cultivate effective teams. Strong communicator and leader.', made_from=DocumentChunk(id=UUID('7e35407f-7c59-5429-8824-23f1d17118c0'), updated_at=datetime.datetime(2024, 12, 24, 11, 54, 13, 481297, tzinfo=datetime.timezone.utc), topological_rank=0, text=\"\\nCV 5: Not Relevant\\nName: Jessica Miller\\nContact Information:\\n\\nEmail: jessica.miller@example.com\\nPhone: (555) 567-8901\\nSummary:\\n\\nExperienced Sales Manager with a strong track record in driving sales growth and building high-performing teams. Excellent communication and leadership skills.\\n\\nEducation:\\n\\nB.A. in Business Administration, University of Southern California (2010)\\nExperience:\\n\\nSales Manager, Global Enterprises (2015 – Present)\\nManaged a sales team of 15 members, achieving a 20% increase in annual revenue.\\nDeveloped sales strategies that expanded customer base by 25%.\\nSales Representative, Market Leaders Inc. (2010 – 2015)\\nConsistently exceeded sales targets and received the 'Top Salesperson' award in 2013.\\nSkills:\\n\\nSales Strategy and Planning\\nTeam Leadership and Development\\nCRM Software: Salesforce, Zoho\\nNegotiation and Relationship Building\\n\", word_count=102, chunk_index=0, cut_type='sentence_cut', is_part_of=TextDocument(id=UUID('3c323fc9-9165-52da-a079-2627a9556b08'), updated_at=datetime.datetime(2024, 12, 24, 11, 54, 13, 481297, tzinfo=datetime.timezone.utc), topological_rank=0, name='text_9b35c7df1f5d4dc84e78270c0bf9cac6', raw_data_location='/Users/vasilije/cognee/cognee/.data_storage/data/text_9b35c7df1f5d4dc84e78270c0bf9cac6.txt', metadata_id=UUID('3c323fc9-9165-52da-a079-2627a9556b08'), mime_type='text/plain', type='text'), contains=[Entity(id=UUID('36a5e3c8-c5f5-5ab5-8d59-ea69d8b36932'), updated_at=datetime.datetime(2024, 12, 24, 11, 54, 13, 481297, tzinfo=datetime.timezone.utc), topological_rank=0, name='jessica miller', is_a=EntityType(id=UUID('d072ba0f-e1a9-58bf-9974-e1802adc8134'), updated_at=datetime.datetime(2024, 12, 24, 11, 54, 13, 481297, tzinfo=datetime.timezone.utc), topological_rank=0, name='person', description='person'), description='An experienced sales manager with a strong track record in driving sales growth and building high-performing teams.'), Entity(id=UUID('5c32691d-c0e4-5378-9aab-dda8b0fa3931'), updated_at=datetime.datetime(2024, 12, 24, 11, 54, 13, 481297, tzinfo=datetime.timezone.utc), topological_rank=0, name='global enterprises', is_a=EntityType(id=UUID('a6ed6bf1-fe31-5dfe-8ab4-484691fdf219'), updated_at=datetime.datetime(2024, 12, 24, 11, 54, 13, 481297, tzinfo=datetime.timezone.utc), topological_rank=0, name='company', description='company'), description='A company where Jessica Miller worked as a Sales Manager.'), Entity(id=UUID('67544857-983a-5152-801d-4fc9d35d14e4'), updated_at=datetime.datetime(2024, 12, 24, 11, 54, 13, 481297, tzinfo=datetime.timezone.utc), topological_rank=0, name='market leaders inc.', is_a=EntityType(id=UUID('a6ed6bf1-fe31-5dfe-8ab4-484691fdf219'), updated_at=datetime.datetime(2024, 12, 24, 11, 54, 13, 481297, tzinfo=datetime.timezone.utc), topological_rank=0, name='company', description='company'), description='A company where Jessica Miller worked as a Sales Representative.'), Entity(id=UUID('f39d6c00-689b-5fd2-9021-893b28ac6ff2'), updated_at=datetime.datetime(2024, 12, 24, 11, 54, 13, 481297, tzinfo=datetime.timezone.utc), topological_rank=0, name='university of southern california', is_a=EntityType(id=UUID('912b273c-683d-53ea-8ffe-aadef0b84237'), updated_at=datetime.datetime(2024, 12, 24, 11, 54, 13, 481297, tzinfo=datetime.timezone.utc), topological_rank=0, name='educational institution', description='educational institution'), description='University where Jessica Miller obtained her degree in Business Administration.'), Entity(id=UUID('0abc801d-38ca-5003-b974-b60f1956c94a'), updated_at=datetime.datetime(2024, 12, 24, 11, 54, 13, 481297, tzinfo=datetime.timezone.utc), topological_rank=0, name='2010', is_a=EntityType(id=UUID('d61d99ac-b291-5666-9748-3e80e1c8b56a'), updated_at=datetime.datetime(2024, 12, 24, 11, 54, 13, 481297, tzinfo=datetime.timezone.utc), topological_rank=0, name='date', description='date'), description='Year Jessica Miller graduated from University of Southern California.'), Entity(id=UUID('7c8b43c1-e133-52e6-99aa-239534f1ed45'), updated_at=datetime.datetime(2024, 12, 24, 11, 54, 13, 481297, tzinfo=datetime.timezone.utc), topological_rank=0, name='2015', is_a=EntityType(id=UUID('d61d99ac-b291-5666-9748-3e80e1c8b56a'), updated_at=datetime.datetime(2024, 12, 24, 11, 54, 13, 481297, tzinfo=datetime.timezone.utc), topological_rank=0, name='date', description='date'), description='Year Jessica Miller started working as Sales Manager at Global Enterprises.'), Entity(id=UUID('2f4749e9-e1e4-5af0-be80-2a10d07557ff'), updated_at=datetime.datetime(2024, 12, 24, 11, 54, 13, 481297, tzinfo=datetime.timezone.utc), topological_rank=0, name='present', is_a=EntityType(id=UUID('d61d99ac-b291-5666-9748-3e80e1c8b56a'), updated_at=datetime.datetime(2024, 12, 24, 11, 54, 13, 481297, tzinfo=datetime.timezone.utc), topological_rank=0, name='date', description='date'), description=\"Current time indicative of Jessica Miller's ongoing role.\")])), TextSummary(id=UUID('d8a8668e-b122-5713-b289-932407bb294e'), updated_at=datetime.datetime(2024, 12, 24, 11, 54, 13, 481297, tzinfo=datetime.timezone.utc), topological_rank=0, text='Creative Graphic Designer with 8+ years of expertise in visual design and branding, skilled in Adobe Creative Suite, dedicated to crafting engaging visuals.', made_from=DocumentChunk(id=UUID('c401b5b1-21d8-5830-8c7b-48e7d94c5b95'), updated_at=datetime.datetime(2024, 12, 24, 11, 54, 13, 481297, tzinfo=datetime.timezone.utc), topological_rank=0, text='\\nCV 4: Not Relevant\\nName: David Thompson\\nContact Information:\\n\\nEmail: david.thompson@example.com\\nPhone: (555) 456-7890\\nSummary:\\n\\nCreative Graphic Designer with over 8 years of experience in visual design and branding. Proficient in Adobe Creative Suite and passionate about creating compelling visuals.\\n\\nEducation:\\n\\nB.F.A. in Graphic Design, Rhode Island School of Design (2012)\\nExperience:\\n\\nSenior Graphic Designer, CreativeWorks Agency (2015 – Present)\\nLed design projects for clients in various industries.\\nCreated branding materials that increased client engagement by 30%.\\nGraphic Designer, Visual Innovations (2012 – 2015)\\nDesigned marketing collateral, including brochures, logos, and websites.\\nCollaborated with the marketing team to develop cohesive brand strategies.\\nSkills:\\n\\nDesign Software: Adobe Photoshop, Illustrator, InDesign\\nWeb Design: HTML, CSS\\nSpecialties: Branding and Identity, Typography\\n', word_count=108, chunk_index=0, cut_type='sentence_cut', is_part_of=TextDocument(id=UUID('e71daf63-15a0-50fe-a909-766bc8fd311b'), updated_at=datetime.datetime(2024, 12, 24, 11, 54, 13, 481297, tzinfo=datetime.timezone.utc), topological_rank=0, name='text_9abf20fa7defd7e49296c51b4e38edf2', raw_data_location='/Users/vasilije/cognee/cognee/.data_storage/data/text_9abf20fa7defd7e49296c51b4e38edf2.txt', metadata_id=UUID('e71daf63-15a0-50fe-a909-766bc8fd311b'), mime_type='text/plain', type='text'), contains=[Entity(id=UUID('a4777597-06c7-562c-bc44-56f74571a01a'), updated_at=datetime.datetime(2024, 12, 24, 11, 54, 13, 481297, tzinfo=datetime.timezone.utc), topological_rank=0, name='david thompson', is_a=EntityType(id=UUID('d072ba0f-e1a9-58bf-9974-e1802adc8134'), updated_at=datetime.datetime(2024, 12, 24, 11, 54, 13, 481297, tzinfo=datetime.timezone.utc), topological_rank=0, name='person', description='person'), description='Creative Graphic Designer with over 8 years of experience in visual design and branding.'), Entity(id=UUID('ca20272a-3e88-552f-92fe-491e23f117f8'), updated_at=datetime.datetime(2024, 12, 24, 11, 54, 13, 481297, tzinfo=datetime.timezone.utc), topological_rank=0, name='creativeworks agency', is_a=EntityType(id=UUID('d3d7b6b4-9b0d-52e8-9e09-a9e9cf4b5a4d'), updated_at=datetime.datetime(2024, 12, 24, 11, 54, 13, 481297, tzinfo=datetime.timezone.utc), topological_rank=0, name='organization', description='organization'), description='An agency where David Thompson is a Senior Graphic Designer.'), Entity(id=UUID('1e97bb97-4d29-5fb8-863a-15ab51f1dd46'), updated_at=datetime.datetime(2024, 12, 24, 11, 54, 13, 481297, tzinfo=datetime.timezone.utc), topological_rank=0, name='visual innovations', is_a=EntityType(id=UUID('d3d7b6b4-9b0d-52e8-9e09-a9e9cf4b5a4d'), updated_at=datetime.datetime(2024, 12, 24, 11, 54, 13, 481297, tzinfo=datetime.timezone.utc), topological_rank=0, name='organization', description='organization'), description='An organization where David Thompson worked as a Graphic Designer.'), Entity(id=UUID('60b027fe-7bb4-535d-8a47-19f1a491591b'), updated_at=datetime.datetime(2024, 12, 24, 11, 54, 13, 481297, tzinfo=datetime.timezone.utc), topological_rank=0, name='rhode island school of design', is_a=EntityType(id=UUID('b5866225-05ad-5cfc-908e-c22916c6a1c6'), updated_at=datetime.datetime(2024, 12, 24, 11, 54, 13, 481297, tzinfo=datetime.timezone.utc), topological_rank=0, name='institution', description='institution'), description='Educational institution where David Thompson earned his B.F.A. in Graphic Design.'), Entity(id=UUID('7e3df89c-2691-580b-84dc-378cb1df3db6'), updated_at=datetime.datetime(2024, 12, 24, 11, 54, 13, 481297, tzinfo=datetime.timezone.utc), topological_rank=0, name='adobe creative suite', is_a=EntityType(id=UUID('2d66edc2-1e14-55ab-8304-680b514a597a'), updated_at=datetime.datetime(2024, 12, 24, 11, 54, 13, 481297, tzinfo=datetime.timezone.utc), topological_rank=0, name='software', description='software'), description='A suite of graphic design, video editing, and web development applications.'), Entity(id=UUID('2a0f9b58-c695-5bad-baa2-fd2da02da013'), updated_at=datetime.datetime(2024, 12, 24, 11, 54, 13, 481297, tzinfo=datetime.timezone.utc), topological_rank=0, name='html', is_a=EntityType(id=UUID('c90c7d6b-3532-5dcf-91e1-4a0e1f179794'), updated_at=datetime.datetime(2024, 12, 24, 11, 54, 13, 481297, tzinfo=datetime.timezone.utc), topological_rank=0, name='language', description='language'), description='A markup language for creating web pages and applications.'), Entity(id=UUID('9b062f3c-fe02-5427-9b44-b287a1cac367'), updated_at=datetime.datetime(2024, 12, 24, 11, 54, 13, 481297, tzinfo=datetime.timezone.utc), topological_rank=0, name='css', is_a=EntityType(id=UUID('c90c7d6b-3532-5dcf-91e1-4a0e1f179794'), updated_at=datetime.datetime(2024, 12, 24, 11, 54, 13, 481297, tzinfo=datetime.timezone.utc), topological_rank=0, name='language', description='language'), description='A stylesheet language used for describing the presentation of a document written in HTML.')])), TextSummary(id=UUID('8aedca6b-fa78-5987-a79b-3b0bebff8eb1'), updated_at=datetime.datetime(2024, 12, 24, 11, 54, 13, 481297, tzinfo=datetime.timezone.utc), topological_rank=0, text='Data Scientist with 6 years of expertise in machine learning, focused on utilizing data for business improvement and product enhancement.', made_from=DocumentChunk(id=UUID('00692e43-9f02-54d0-b695-44bf47342d36'), updated_at=datetime.datetime(2024, 12, 24, 11, 54, 13, 481297, tzinfo=datetime.timezone.utc), topological_rank=0, text='\\nCV 3: Relevant\\nName: Sarah Nguyen\\nContact Information:\\n\\nEmail: sarah.nguyen@example.com\\nPhone: (555) 345-6789\\nSummary:\\n\\nData Scientist specializing in machine learning with 6 years of experience. Passionate about leveraging data to drive business solutions and improve product performance.\\n\\nEducation:\\n\\nM.S. in Statistics, University of Washington (2014)\\nB.S. in Applied Mathematics, University of Texas at Austin (2012)\\nExperience:\\n\\nData Scientist, QuantumTech (2016 – Present)\\nDesigned and implemented machine learning algorithms for financial forecasting.\\nImproved model efficiency by 20% through algorithm optimization.\\nJunior Data Scientist, DataCore Solutions (2014 – 2016)\\nAssisted in developing predictive models for supply chain optimization.\\nConducted data cleaning and preprocessing on large datasets.\\nSkills:\\n\\nProgramming Languages: Python, R\\nMachine Learning Frameworks: PyTorch, Scikit-Learn\\nStatistical Analysis: SAS, SPSS\\nCloud Platforms: AWS, Azure\\n', word_count=110, chunk_index=0, cut_type='sentence_cut', is_part_of=TextDocument(id=UUID('e7d6246b-e414-5b9d-8daa-6d4434b7fa47'), updated_at=datetime.datetime(2024, 12, 24, 11, 54, 13, 481297, tzinfo=datetime.timezone.utc), topological_rank=0, name='text_f8c7482e727f228001b0046ed68d656f', raw_data_location='/Users/vasilije/cognee/cognee/.data_storage/data/text_f8c7482e727f228001b0046ed68d656f.txt', metadata_id=UUID('e7d6246b-e414-5b9d-8daa-6d4434b7fa47'), mime_type='text/plain', type='text'), contains=[Entity(id=UUID('4d8dda57-2681-5264-a2bd-e2ddfe66a785'), updated_at=datetime.datetime(2024, 12, 24, 11, 54, 13, 481297, tzinfo=datetime.timezone.utc), topological_rank=0, name='sarah nguyen', is_a=EntityType(id=UUID('d072ba0f-e1a9-58bf-9974-e1802adc8134'), updated_at=datetime.datetime(2024, 12, 24, 11, 54, 13, 481297, tzinfo=datetime.timezone.utc), topological_rank=0, name='person', description='person'), description='Data Scientist specializing in machine learning with 6 years of experience.'), Entity(id=UUID('ae74a35b-d5f1-5622-ade1-6703d5e069fb'), updated_at=datetime.datetime(2024, 12, 24, 11, 54, 13, 481297, tzinfo=datetime.timezone.utc), topological_rank=0, name='university of washington', is_a=EntityType(id=UUID('912b273c-683d-53ea-8ffe-aadef0b84237'), updated_at=datetime.datetime(2024, 12, 24, 11, 54, 13, 481297, tzinfo=datetime.timezone.utc), topological_rank=0, name='educational institution', description='educational institution'), description='University located in Seattle, Washington.'), Entity(id=UUID('301b3cf8-5a5c-585e-80bd-f79901e4368c'), updated_at=datetime.datetime(2024, 12, 24, 11, 54, 13, 481297, tzinfo=datetime.timezone.utc), topological_rank=0, name='university of texas at austin', is_a=EntityType(id=UUID('912b273c-683d-53ea-8ffe-aadef0b84237'), updated_at=datetime.datetime(2024, 12, 24, 11, 54, 13, 481297, tzinfo=datetime.timezone.utc), topological_rank=0, name='educational institution', description='educational institution'), description='Public research university located in Austin, Texas.'), Entity(id=UUID('0d980f2a-09dd-581e-acc3-cc2d87c1bab4'), updated_at=datetime.datetime(2024, 12, 24, 11, 54, 13, 481297, tzinfo=datetime.timezone.utc), topological_rank=0, name='quantumtech', is_a=EntityType(id=UUID('a6ed6bf1-fe31-5dfe-8ab4-484691fdf219'), updated_at=datetime.datetime(2024, 12, 24, 11, 54, 13, 481297, tzinfo=datetime.timezone.utc), topological_rank=0, name='company', description='company'), description='Company where Sarah Nguyen works as a Data Scientist from 2016 to present.'), Entity(id=UUID('95ac0551-38fc-5187-a422-533aeb7e8db0'), updated_at=datetime.datetime(2024, 12, 24, 11, 54, 13, 481297, tzinfo=datetime.timezone.utc), topological_rank=0, name='datacore solutions', is_a=EntityType(id=UUID('a6ed6bf1-fe31-5dfe-8ab4-484691fdf219'), updated_at=datetime.datetime(2024, 12, 24, 11, 54, 13, 481297, tzinfo=datetime.timezone.utc), topological_rank=0, name='company', description='company'), description='Company where Sarah Nguyen worked as a Junior Data Scientist from 2014 to 2016.'), Entity(id=UUID('3edcdf3f-25af-57a3-8878-8008bd7ea05a'), updated_at=datetime.datetime(2024, 12, 24, 11, 54, 13, 481297, tzinfo=datetime.timezone.utc), topological_rank=0, name='aws', is_a=EntityType(id=UUID('d84d991a-dab3-5c36-8806-df076ccb731b'), updated_at=datetime.datetime(2024, 12, 24, 11, 54, 13, 481297, tzinfo=datetime.timezone.utc), topological_rank=0, name='cloud platform', description='cloud platform'), description='Amazon Web Services, a cloud computing platform.'), Entity(id=UUID('8b431923-4aa2-5886-a661-b8de0f888a9b'), updated_at=datetime.datetime(2024, 12, 24, 11, 54, 13, 481297, tzinfo=datetime.timezone.utc), topological_rank=0, name='azure', is_a=EntityType(id=UUID('d84d991a-dab3-5c36-8806-df076ccb731b'), updated_at=datetime.datetime(2024, 12, 24, 11, 54, 13, 481297, tzinfo=datetime.timezone.utc), topological_rank=0, name='cloud platform', description='cloud platform'), description='Microsoft Azure, a cloud computing service created by Microsoft.')]))]\n"
+ ]
+ }
+ ],
+ "execution_count": 11
},
{
"cell_type": "markdown",
@@ -703,23 +782,61 @@
},
{
"cell_type": "code",
- "execution_count": null,
"id": "080389e5",
- "metadata": {},
- "outputs": [],
+ "metadata": {
+ "ExecuteTime": {
+ "end_time": "2024-12-29T16:55:51.378129Z",
+ "start_time": "2024-12-29T16:55:46.922951Z"
+ }
+ },
"source": [
"import os\n",
"from cognee.shared.utils import render_graph\n",
"from cognee.infrastructure.databases.graph import get_graph_engine\n",
"import graphistry\n",
- "\n",
+ "# from dotenv import load_dotenv \n",
"graphistry.login(username=os.getenv(\"GRAPHISTRY_USERNAME\"), password=os.getenv(\"GRAPHISTRY_PASSWORD\"))\n",
"\n",
"graph_engine = await get_graph_engine()\n",
"\n",
"graph_url = await render_graph(graph_engine.graph)\n",
"print(graph_url)"
- ]
+ ],
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "Graph is visualized at: https://hub.graphistry.com/graph/graph.html?dataset=cc21b1d2d6074323aa37af53e693b1a4&type=arrow&viztoken=db05565e-79e9-4fe3-99b2-b7a2e6d48eff&usertag=5f822e63-pygraphistry-0.33.9&splashAfter=1735491366&info=true\n",
+ "https://hub.graphistry.com/graph/graph.html?dataset=cc21b1d2d6074323aa37af53e693b1a4&type=arrow&viztoken=db05565e-79e9-4fe3-99b2-b7a2e6d48eff&usertag=5f822e63-pygraphistry-0.33.9&splashAfter=1735491366&info=true\n"
+ ]
+ }
+ ],
+ "execution_count": 12
+ },
+ {
+ "metadata": {
+ "ExecuteTime": {
+ "end_time": "2024-12-29T16:56:06.571404Z",
+ "start_time": "2024-12-29T16:56:06.569280Z"
+ }
+ },
+ "cell_type": "code",
+ "source": [
+ "graph_engine = await get_graph_engine()\n",
+ "print(graph_url)"
+ ],
+ "id": "8f69caa0e353a889",
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "https://hub.graphistry.com/graph/graph.html?dataset=cc21b1d2d6074323aa37af53e693b1a4&type=arrow&viztoken=db05565e-79e9-4fe3-99b2-b7a2e6d48eff&usertag=5f822e63-pygraphistry-0.33.9&splashAfter=1735491366&info=true\n"
+ ]
+ }
+ ],
+ "execution_count": 13
},
{
"cell_type": "markdown",
@@ -731,10 +848,13 @@
},
{
"cell_type": "code",
- "execution_count": null,
"id": "e5e7dfc8",
- "metadata": {},
- "outputs": [],
+ "metadata": {
+ "ExecuteTime": {
+ "end_time": "2024-12-24T13:44:16.575843Z",
+ "start_time": "2024-12-24T13:44:16.047897Z"
+ }
+ },
"source": [
"async def search(\n",
" vector_engine,\n",
@@ -763,7 +883,26 @@
"results = await search(vector_engine, \"entity_name\", \"sarah.nguyen@example.com\")\n",
"for result in results:\n",
" print(result)"
- ]
+ ],
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "{'id': '4d8dda57-2681-5264-a2bd-e2ddfe66a785', 'payload': {'id': '4d8dda57-2681-5264-a2bd-e2ddfe66a785', 'updated_at': datetime.datetime(2024, 12, 24, 11, 54, 13, 481297), 'topological_rank': 0, 'text': 'sarah nguyen'}, 'score': 0.5708460211753845}\n",
+ "{'id': '198e2ab8-75e9-5931-97ab-da9a5a8e188c', 'payload': {'id': '198e2ab8-75e9-5931-97ab-da9a5a8e188c', 'updated_at': datetime.datetime(2024, 12, 24, 11, 54, 13, 481297), 'topological_rank': 0, 'text': 'san francisco, ca'}, 'score': 1.349550724029541}\n",
+ "{'id': '435dbd37-ab20-503c-9e99-ab8b8a3484e5', 'payload': {'id': '435dbd37-ab20-503c-9e99-ab8b8a3484e5', 'updated_at': datetime.datetime(2024, 12, 24, 11, 54, 13, 481297), 'topological_rank': 0, 'text': 'senior data scientist'}, 'score': 1.3934645652770996}\n",
+ "{'id': '36a5e3c8-c5f5-5ab5-8d59-ea69d8b36932', 'payload': {'id': '36a5e3c8-c5f5-5ab5-8d59-ea69d8b36932', 'updated_at': datetime.datetime(2024, 12, 24, 11, 54, 13, 481297), 'topological_rank': 0, 'text': 'jessica miller'}, 'score': 1.4042469263076782}\n",
+ "{'id': '73ae630f-7b09-5dce-8c18-45d0a57b30f9', 'payload': {'id': '73ae630f-7b09-5dce-8c18-45d0a57b30f9', 'updated_at': datetime.datetime(2024, 12, 24, 11, 54, 13, 481297), 'topological_rank': 0, 'text': 'michael rodriguez'}, 'score': 1.4521806240081787}\n",
+ "{'id': '29e771c8-4c3f-52de-9511-6b705878e130', 'payload': {'id': '29e771c8-4c3f-52de-9511-6b705878e130', 'updated_at': datetime.datetime(2024, 12, 24, 11, 54, 13, 481297), 'topological_rank': 0, 'text': 'dr. emily carter'}, 'score': 1.471205472946167}\n",
+ "{'id': 'ce8b394a-b30e-52fc-b80a-6352edc60e5b', 'payload': {'id': 'ce8b394a-b30e-52fc-b80a-6352edc60e5b', 'updated_at': datetime.datetime(2024, 12, 24, 11, 54, 13, 481297), 'topological_rank': 0, 'text': 'stanford university'}, 'score': 1.4871069192886353}\n",
+ "{'id': '9780afb1-dccc-53eb-9a30-c0d4ce033711', 'payload': {'id': '9780afb1-dccc-53eb-9a30-c0d4ce033711', 'updated_at': datetime.datetime(2024, 12, 24, 11, 54, 13, 481297), 'topological_rank': 0, 'text': 'innovateai labs'}, 'score': 1.498255968093872}\n",
+ "{'id': '301b3cf8-5a5c-585e-80bd-f79901e4368c', 'payload': {'id': '301b3cf8-5a5c-585e-80bd-f79901e4368c', 'updated_at': datetime.datetime(2024, 12, 24, 11, 54, 13, 481297), 'topological_rank': 0, 'text': 'university of texas at austin'}, 'score': 1.5053001642227173}\n",
+ "{'id': '2c02c93c-9cd1-56b8-9cc0-55ff0b290e57', 'payload': {'id': '2c02c93c-9cd1-56b8-9cc0-55ff0b290e57', 'updated_at': datetime.datetime(2024, 12, 24, 11, 54, 13, 481297), 'topological_rank': 0, 'text': 'university of california, berkeley'}, 'score': 1.5213639736175537}\n"
+ ]
+ }
+ ],
+ "execution_count": 23
},
{
"cell_type": "markdown",
@@ -842,20 +981,455 @@
]
},
{
- "cell_type": "markdown",
- "id": "2ab3d84a",
- "metadata": {},
+ "metadata": {
+ "ExecuteTime": {
+ "end_time": "2024-12-24T13:46:09.644509Z",
+ "start_time": "2024-12-24T13:46:04.538592Z"
+ }
+ },
+ "cell_type": "code",
"source": [
- "#### Bellow is a diagram of the cognee process for the data used in this example notebook"
- ]
+ "!pip install wget\n",
+ "!pip install deepeval\n",
+ "!pip install ujson"
+ ],
+ "id": "afae18ac6a794925",
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "Requirement already satisfied: wget in /Users/vasilije/cognee/.venv/lib/python3.11/site-packages (3.2)\r\n",
+ "\r\n",
+ "\u001B[1m[\u001B[0m\u001B[34;49mnotice\u001B[0m\u001B[1;39;49m]\u001B[0m\u001B[39;49m A new release of pip is available: \u001B[0m\u001B[31;49m23.2.1\u001B[0m\u001B[39;49m -> \u001B[0m\u001B[32;49m24.3.1\u001B[0m\r\n",
+ "\u001B[1m[\u001B[0m\u001B[34;49mnotice\u001B[0m\u001B[1;39;49m]\u001B[0m\u001B[39;49m To update, run: \u001B[0m\u001B[32;49mpip install --upgrade pip\u001B[0m\r\n",
+ "Requirement already satisfied: deepeval in /Users/vasilije/cognee/.venv/lib/python3.11/site-packages (2.0.9)\r\n",
+ "Requirement already satisfied: requests in /Users/vasilije/cognee/.venv/lib/python3.11/site-packages (from deepeval) (2.32.3)\r\n",
+ "Requirement already satisfied: tqdm in /Users/vasilije/cognee/.venv/lib/python3.11/site-packages (from deepeval) (4.67.1)\r\n",
+ "Requirement already satisfied: pytest in /Users/vasilije/cognee/.venv/lib/python3.11/site-packages (from deepeval) (7.4.4)\r\n",
+ "Requirement already satisfied: tabulate in /Users/vasilije/cognee/.venv/lib/python3.11/site-packages (from deepeval) (0.9.0)\r\n",
+ "Requirement already satisfied: typer in /Users/vasilije/cognee/.venv/lib/python3.11/site-packages (from deepeval) (0.15.1)\r\n",
+ "Requirement already satisfied: rich in /Users/vasilije/cognee/.venv/lib/python3.11/site-packages (from deepeval) (13.9.4)\r\n",
+ "Requirement already satisfied: protobuf in /Users/vasilije/cognee/.venv/lib/python3.11/site-packages (from deepeval) (4.25.5)\r\n",
+ "Requirement already satisfied: pydantic in /Users/vasilije/cognee/.venv/lib/python3.11/site-packages (from deepeval) (2.8.2)\r\n",
+ "Requirement already satisfied: sentry-sdk in /Users/vasilije/cognee/.venv/lib/python3.11/site-packages (from deepeval) (2.19.2)\r\n",
+ "Requirement already satisfied: pytest-repeat in /Users/vasilije/cognee/.venv/lib/python3.11/site-packages (from deepeval) (0.9.3)\r\n",
+ "Requirement already satisfied: pytest-xdist in /Users/vasilije/cognee/.venv/lib/python3.11/site-packages (from deepeval) (3.6.1)\r\n",
+ "Requirement already satisfied: portalocker in /Users/vasilije/cognee/.venv/lib/python3.11/site-packages (from deepeval) (3.0.0)\r\n",
+ "Requirement already satisfied: langchain in /Users/vasilije/cognee/.venv/lib/python3.11/site-packages (from deepeval) (0.3.13)\r\n",
+ "Requirement already satisfied: langchain-core in /Users/vasilije/cognee/.venv/lib/python3.11/site-packages (from deepeval) (0.3.28)\r\n",
+ "Requirement already satisfied: langchain_openai in /Users/vasilije/cognee/.venv/lib/python3.11/site-packages (from deepeval) (0.2.14)\r\n",
+ "Requirement already satisfied: langchain-community in /Users/vasilije/cognee/.venv/lib/python3.11/site-packages (from deepeval) (0.3.13)\r\n",
+ "Requirement already satisfied: docx2txt~=0.8 in /Users/vasilije/cognee/.venv/lib/python3.11/site-packages (from deepeval) (0.8)\r\n",
+ "Requirement already satisfied: importlib-metadata>=6.0.2 in /Users/vasilije/cognee/.venv/lib/python3.11/site-packages (from deepeval) (8.4.0)\r\n",
+ "Requirement already satisfied: tenacity<=9.0.0 in /Users/vasilije/cognee/.venv/lib/python3.11/site-packages (from deepeval) (8.5.0)\r\n",
+ "Requirement already satisfied: opentelemetry-api<2.0.0,>=1.24.0 in /Users/vasilije/cognee/.venv/lib/python3.11/site-packages (from deepeval) (1.27.0)\r\n",
+ "Requirement already satisfied: opentelemetry-sdk<2.0.0,>=1.24.0 in /Users/vasilije/cognee/.venv/lib/python3.11/site-packages (from deepeval) (1.27.0)\r\n",
+ "Requirement already satisfied: opentelemetry-exporter-otlp-proto-grpc<2.0.0,>=1.24.0 in /Users/vasilije/cognee/.venv/lib/python3.11/site-packages (from deepeval) (1.27.0)\r\n",
+ "Requirement already satisfied: grpcio==1.60.1 in /Users/vasilije/cognee/.venv/lib/python3.11/site-packages (from deepeval) (1.60.1)\r\n",
+ "Requirement already satisfied: nest-asyncio in /Users/vasilije/cognee/.venv/lib/python3.11/site-packages (from deepeval) (1.6.0)\r\n",
+ "Requirement already satisfied: zipp>=0.5 in /Users/vasilije/cognee/.venv/lib/python3.11/site-packages (from importlib-metadata>=6.0.2->deepeval) (3.21.0)\r\n",
+ "Requirement already satisfied: deprecated>=1.2.6 in /Users/vasilije/cognee/.venv/lib/python3.11/site-packages (from opentelemetry-api<2.0.0,>=1.24.0->deepeval) (1.2.15)\r\n",
+ "Requirement already satisfied: googleapis-common-protos~=1.52 in /Users/vasilije/cognee/.venv/lib/python3.11/site-packages (from opentelemetry-exporter-otlp-proto-grpc<2.0.0,>=1.24.0->deepeval) (1.66.0)\r\n",
+ "Requirement already satisfied: opentelemetry-exporter-otlp-proto-common==1.27.0 in /Users/vasilije/cognee/.venv/lib/python3.11/site-packages (from opentelemetry-exporter-otlp-proto-grpc<2.0.0,>=1.24.0->deepeval) (1.27.0)\r\n",
+ "Requirement already satisfied: opentelemetry-proto==1.27.0 in /Users/vasilije/cognee/.venv/lib/python3.11/site-packages (from opentelemetry-exporter-otlp-proto-grpc<2.0.0,>=1.24.0->deepeval) (1.27.0)\r\n",
+ "Requirement already satisfied: opentelemetry-semantic-conventions==0.48b0 in /Users/vasilije/cognee/.venv/lib/python3.11/site-packages (from opentelemetry-sdk<2.0.0,>=1.24.0->deepeval) (0.48b0)\r\n",
+ "Requirement already satisfied: typing-extensions>=3.7.4 in /Users/vasilije/cognee/.venv/lib/python3.11/site-packages (from opentelemetry-sdk<2.0.0,>=1.24.0->deepeval) (4.12.2)\r\n",
+ "Requirement already satisfied: PyYAML>=5.3 in /Users/vasilije/cognee/.venv/lib/python3.11/site-packages (from langchain->deepeval) (6.0.2)\r\n",
+ "Requirement already satisfied: SQLAlchemy<3,>=1.4 in /Users/vasilije/cognee/.venv/lib/python3.11/site-packages (from langchain->deepeval) (2.0.35)\r\n",
+ "Requirement already satisfied: aiohttp<4.0.0,>=3.8.3 in /Users/vasilije/cognee/.venv/lib/python3.11/site-packages (from langchain->deepeval) (3.10.10)\r\n",
+ "Requirement already satisfied: langchain-text-splitters<0.4.0,>=0.3.3 in /Users/vasilije/cognee/.venv/lib/python3.11/site-packages (from langchain->deepeval) (0.3.4)\r\n",
+ "Requirement already satisfied: langsmith<0.3,>=0.1.17 in /Users/vasilije/cognee/.venv/lib/python3.11/site-packages (from langchain->deepeval) (0.2.4)\r\n",
+ "Requirement already satisfied: numpy<2,>=1.22.4 in /Users/vasilije/cognee/.venv/lib/python3.11/site-packages (from langchain->deepeval) (1.26.4)\r\n",
+ "Requirement already satisfied: jsonpatch<2.0,>=1.33 in /Users/vasilije/cognee/.venv/lib/python3.11/site-packages (from langchain-core->deepeval) (1.33)\r\n",
+ "Requirement already satisfied: packaging<25,>=23.2 in /Users/vasilije/cognee/.venv/lib/python3.11/site-packages (from langchain-core->deepeval) (24.2)\r\n",
+ "Requirement already satisfied: annotated-types>=0.4.0 in /Users/vasilije/cognee/.venv/lib/python3.11/site-packages (from pydantic->deepeval) (0.7.0)\r\n",
+ "Requirement already satisfied: pydantic-core==2.20.1 in /Users/vasilije/cognee/.venv/lib/python3.11/site-packages (from pydantic->deepeval) (2.20.1)\r\n",
+ "Requirement already satisfied: charset-normalizer<4,>=2 in /Users/vasilije/cognee/.venv/lib/python3.11/site-packages (from requests->deepeval) (3.4.0)\r\n",
+ "Requirement already satisfied: idna<4,>=2.5 in /Users/vasilije/cognee/.venv/lib/python3.11/site-packages (from requests->deepeval) (3.10)\r\n",
+ "Requirement already satisfied: urllib3<3,>=1.21.1 in /Users/vasilije/cognee/.venv/lib/python3.11/site-packages (from requests->deepeval) (2.2.3)\r\n",
+ "Requirement already satisfied: certifi>=2017.4.17 in /Users/vasilije/cognee/.venv/lib/python3.11/site-packages (from requests->deepeval) (2024.12.14)\r\n",
+ "Requirement already satisfied: dataclasses-json<0.7,>=0.5.7 in /Users/vasilije/cognee/.venv/lib/python3.11/site-packages (from langchain-community->deepeval) (0.6.7)\r\n",
+ "Requirement already satisfied: httpx-sse<0.5.0,>=0.4.0 in /Users/vasilije/cognee/.venv/lib/python3.11/site-packages (from langchain-community->deepeval) (0.4.0)\r\n",
+ "Requirement already satisfied: pydantic-settings<3.0.0,>=2.4.0 in /Users/vasilije/cognee/.venv/lib/python3.11/site-packages (from langchain-community->deepeval) (2.7.0)\r\n",
+ "Requirement already satisfied: openai<2.0.0,>=1.58.1 in /Users/vasilije/cognee/.venv/lib/python3.11/site-packages (from langchain_openai->deepeval) (1.58.1)\r\n",
+ "Requirement already satisfied: tiktoken<1,>=0.7 in /Users/vasilije/cognee/.venv/lib/python3.11/site-packages (from langchain_openai->deepeval) (0.7.0)\r\n",
+ "Requirement already satisfied: iniconfig in /Users/vasilije/cognee/.venv/lib/python3.11/site-packages (from pytest->deepeval) (2.0.0)\r\n",
+ "Requirement already satisfied: pluggy<2.0,>=0.12 in /Users/vasilije/cognee/.venv/lib/python3.11/site-packages (from pytest->deepeval) (1.5.0)\r\n",
+ "Requirement already satisfied: execnet>=2.1 in /Users/vasilije/cognee/.venv/lib/python3.11/site-packages (from pytest-xdist->deepeval) (2.1.1)\r\n",
+ "Requirement already satisfied: markdown-it-py>=2.2.0 in /Users/vasilije/cognee/.venv/lib/python3.11/site-packages (from rich->deepeval) (3.0.0)\r\n",
+ "Requirement already satisfied: pygments<3.0.0,>=2.13.0 in /Users/vasilije/cognee/.venv/lib/python3.11/site-packages (from rich->deepeval) (2.18.0)\r\n",
+ "Requirement already satisfied: click>=8.0.0 in /Users/vasilije/cognee/.venv/lib/python3.11/site-packages (from typer->deepeval) (8.1.7)\r\n",
+ "Requirement already satisfied: shellingham>=1.3.0 in /Users/vasilije/cognee/.venv/lib/python3.11/site-packages (from typer->deepeval) (1.5.4)\r\n",
+ "Requirement already satisfied: aiohappyeyeballs>=2.3.0 in /Users/vasilije/cognee/.venv/lib/python3.11/site-packages (from aiohttp<4.0.0,>=3.8.3->langchain->deepeval) (2.4.4)\r\n",
+ "Requirement already satisfied: aiosignal>=1.1.2 in /Users/vasilije/cognee/.venv/lib/python3.11/site-packages (from aiohttp<4.0.0,>=3.8.3->langchain->deepeval) (1.3.2)\r\n",
+ "Requirement already satisfied: attrs>=17.3.0 in /Users/vasilije/cognee/.venv/lib/python3.11/site-packages (from aiohttp<4.0.0,>=3.8.3->langchain->deepeval) (24.3.0)\r\n",
+ "Requirement already satisfied: frozenlist>=1.1.1 in /Users/vasilije/cognee/.venv/lib/python3.11/site-packages (from aiohttp<4.0.0,>=3.8.3->langchain->deepeval) (1.5.0)\r\n",
+ "Requirement already satisfied: multidict<7.0,>=4.5 in /Users/vasilije/cognee/.venv/lib/python3.11/site-packages (from aiohttp<4.0.0,>=3.8.3->langchain->deepeval) (6.1.0)\r\n",
+ "Requirement already satisfied: yarl<2.0,>=1.12.0 in /Users/vasilije/cognee/.venv/lib/python3.11/site-packages (from aiohttp<4.0.0,>=3.8.3->langchain->deepeval) (1.18.3)\r\n",
+ "Requirement already satisfied: marshmallow<4.0.0,>=3.18.0 in /Users/vasilije/cognee/.venv/lib/python3.11/site-packages (from dataclasses-json<0.7,>=0.5.7->langchain-community->deepeval) (3.23.2)\r\n",
+ "Requirement already satisfied: typing-inspect<1,>=0.4.0 in /Users/vasilije/cognee/.venv/lib/python3.11/site-packages (from dataclasses-json<0.7,>=0.5.7->langchain-community->deepeval) (0.9.0)\r\n",
+ "Requirement already satisfied: wrapt<2,>=1.10 in /Users/vasilije/cognee/.venv/lib/python3.11/site-packages (from deprecated>=1.2.6->opentelemetry-api<2.0.0,>=1.24.0->deepeval) (1.17.0)\r\n",
+ "Requirement already satisfied: jsonpointer>=1.9 in /Users/vasilije/cognee/.venv/lib/python3.11/site-packages (from jsonpatch<2.0,>=1.33->langchain-core->deepeval) (3.0.0)\r\n",
+ "Requirement already satisfied: httpx<1,>=0.23.0 in /Users/vasilije/cognee/.venv/lib/python3.11/site-packages (from langsmith<0.3,>=0.1.17->langchain->deepeval) (0.27.0)\r\n",
+ "Requirement already satisfied: orjson<4.0.0,>=3.9.14 in /Users/vasilije/cognee/.venv/lib/python3.11/site-packages (from langsmith<0.3,>=0.1.17->langchain->deepeval) (3.10.12)\r\n",
+ "Requirement already satisfied: requests-toolbelt<2.0.0,>=1.0.0 in /Users/vasilije/cognee/.venv/lib/python3.11/site-packages (from langsmith<0.3,>=0.1.17->langchain->deepeval) (1.0.0)\r\n",
+ "Requirement already satisfied: mdurl~=0.1 in /Users/vasilije/cognee/.venv/lib/python3.11/site-packages (from markdown-it-py>=2.2.0->rich->deepeval) (0.1.2)\r\n",
+ "Requirement already satisfied: anyio<5,>=3.5.0 in /Users/vasilije/cognee/.venv/lib/python3.11/site-packages (from openai<2.0.0,>=1.58.1->langchain_openai->deepeval) (4.7.0)\r\n",
+ "Requirement already satisfied: distro<2,>=1.7.0 in /Users/vasilije/cognee/.venv/lib/python3.11/site-packages (from openai<2.0.0,>=1.58.1->langchain_openai->deepeval) (1.9.0)\r\n",
+ "Requirement already satisfied: jiter<1,>=0.4.0 in /Users/vasilije/cognee/.venv/lib/python3.11/site-packages (from openai<2.0.0,>=1.58.1->langchain_openai->deepeval) (0.5.0)\r\n",
+ "Requirement already satisfied: sniffio in /Users/vasilije/cognee/.venv/lib/python3.11/site-packages (from openai<2.0.0,>=1.58.1->langchain_openai->deepeval) (1.3.1)\r\n",
+ "Requirement already satisfied: python-dotenv>=0.21.0 in /Users/vasilije/cognee/.venv/lib/python3.11/site-packages (from pydantic-settings<3.0.0,>=2.4.0->langchain-community->deepeval) (1.0.1)\r\n",
+ "Requirement already satisfied: regex>=2022.1.18 in /Users/vasilije/cognee/.venv/lib/python3.11/site-packages (from tiktoken<1,>=0.7->langchain_openai->deepeval) (2024.11.6)\r\n",
+ "Requirement already satisfied: httpcore==1.* in /Users/vasilije/cognee/.venv/lib/python3.11/site-packages (from httpx<1,>=0.23.0->langsmith<0.3,>=0.1.17->langchain->deepeval) (1.0.7)\r\n",
+ "Requirement already satisfied: h11<0.15,>=0.13 in /Users/vasilije/cognee/.venv/lib/python3.11/site-packages (from httpcore==1.*->httpx<1,>=0.23.0->langsmith<0.3,>=0.1.17->langchain->deepeval) (0.14.0)\r\n",
+ "Requirement already satisfied: mypy-extensions>=0.3.0 in /Users/vasilije/cognee/.venv/lib/python3.11/site-packages (from typing-inspect<1,>=0.4.0->dataclasses-json<0.7,>=0.5.7->langchain-community->deepeval) (1.0.0)\r\n",
+ "Requirement already satisfied: propcache>=0.2.0 in /Users/vasilije/cognee/.venv/lib/python3.11/site-packages (from yarl<2.0,>=1.12.0->aiohttp<4.0.0,>=3.8.3->langchain->deepeval) (0.2.1)\r\n",
+ "\r\n",
+ "\u001B[1m[\u001B[0m\u001B[34;49mnotice\u001B[0m\u001B[1;39;49m]\u001B[0m\u001B[39;49m A new release of pip is available: \u001B[0m\u001B[31;49m23.2.1\u001B[0m\u001B[39;49m -> \u001B[0m\u001B[32;49m24.3.1\u001B[0m\r\n",
+ "\u001B[1m[\u001B[0m\u001B[34;49mnotice\u001B[0m\u001B[1;39;49m]\u001B[0m\u001B[39;49m To update, run: \u001B[0m\u001B[32;49mpip install --upgrade pip\u001B[0m\r\n",
+ "Collecting ujson\r\n",
+ " Obtaining dependency information for ujson from https://files.pythonhosted.org/packages/8d/9f/4731ef0671a0653e9f5ba18db7c4596d8ecbf80c7922dd5fe4150f1aea76/ujson-5.10.0-cp311-cp311-macosx_11_0_arm64.whl.metadata\r\n",
+ " Downloading ujson-5.10.0-cp311-cp311-macosx_11_0_arm64.whl.metadata (9.3 kB)\r\n",
+ "Downloading ujson-5.10.0-cp311-cp311-macosx_11_0_arm64.whl (51 kB)\r\n",
+ "\u001B[2K \u001B[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001B[0m \u001B[32m51.8/51.8 kB\u001B[0m \u001B[31m1.7 MB/s\u001B[0m eta \u001B[36m0:00:00\u001B[0m\r\n",
+ "\u001B[?25hInstalling collected packages: ujson\r\n",
+ "Successfully installed ujson-5.10.0\r\n",
+ "\r\n",
+ "\u001B[1m[\u001B[0m\u001B[34;49mnotice\u001B[0m\u001B[1;39;49m]\u001B[0m\u001B[39;49m A new release of pip is available: \u001B[0m\u001B[31;49m23.2.1\u001B[0m\u001B[39;49m -> \u001B[0m\u001B[32;49m24.3.1\u001B[0m\r\n",
+ "\u001B[1m[\u001B[0m\u001B[34;49mnotice\u001B[0m\u001B[1;39;49m]\u001B[0m\u001B[39;49m To update, run: \u001B[0m\u001B[32;49mpip install --upgrade pip\u001B[0m\r\n"
+ ]
+ }
+ ],
+ "execution_count": 29
},
{
- "cell_type": "markdown",
- "id": "31412c52",
- "metadata": {},
+ "metadata": {
+ "ExecuteTime": {
+ "end_time": "2024-12-24T15:29:11.123483Z",
+ "start_time": "2024-12-24T15:29:11.120888Z"
+ }
+ },
+ "cell_type": "code",
"source": [
- ""
- ]
+ "# from evals.eval_on_hotpot import eval_on_hotpotQA\n",
+ "# from evals.eval_on_hotpot import answer_with_cognee\n",
+ "# from evals.eval_on_hotpot import answer_without_cognee\n",
+ "# from evals.eval_on_hotpot import eval_answers\n",
+ "# from cognee.base_config import get_base_config\n",
+ "# from pathlib import Path\n",
+ "# from tqdm import tqdm\n",
+ "# import wget\n",
+ "# import json\n",
+ "# import statistics"
+ ],
+ "id": "5f36b67668fdb646",
+ "outputs": [],
+ "execution_count": 2
+ },
+ {
+ "metadata": {
+ "ExecuteTime": {
+ "end_time": "2024-12-24T15:57:30.764970Z",
+ "start_time": "2024-12-24T15:57:07.861187Z"
+ }
+ },
+ "cell_type": "code",
+ "source": [
+ "# answer_provider = answer_without_cognee # For native LLM answers use answer_without_cognee\n",
+ "# num_samples = 10 # With cognee, it takes ~1m10s per sample\n",
+ "# \n",
+ "# base_config = get_base_config()\n",
+ "# data_root_dir = base_config.data_root_directory\n",
+ "# \n",
+ "# if not Path(data_root_dir).exists():\n",
+ "# Path(data_root_dir).mkdir()\n",
+ "# \n",
+ "# filepath = data_root_dir / Path(\"hotpot_dev_fullwiki_v1.json\")\n",
+ "# if not filepath.exists():\n",
+ "# url = 'http://curtis.ml.cmu.edu/datasets/hotpot/hotpot_dev_fullwiki_v1.json'\n",
+ "# wget.download(url, out=data_root_dir)\n",
+ "# \n",
+ "# with open(filepath, \"r\") as file:\n",
+ "# dataset = json.load(file)\n",
+ "# instances = dataset if not num_samples else dataset[:num_samples]\n",
+ "# answers = []\n",
+ "# for instance in tqdm(instances, desc=\"Getting answers\"):\n",
+ "# answer = await answer_provider(instance)\n",
+ "# answers.append(answer)"
+ ],
+ "id": "d5af4b516c6621a3",
+ "outputs": [
+ {
+ "name": "stderr",
+ "output_type": "stream",
+ "text": [
+ "Getting answers: 100%|██████████| 10/10 [00:13<00:00, 1.31s/it]\n"
+ ]
+ }
+ ],
+ "execution_count": 9
+ },
+ {
+ "metadata": {
+ "ExecuteTime": {
+ "end_time": "2024-12-24T15:57:30.787382Z",
+ "start_time": "2024-12-24T15:57:30.785259Z"
+ }
+ },
+ "cell_type": "code",
+ "source": [
+ "# from evals.deepeval_metrics import f1_score_metric\n",
+ "# from evals.deepeval_metrics import em_score_metric"
+ ],
+ "id": "2bf69048a272158c",
+ "outputs": [],
+ "execution_count": 10
+ },
+ {
+ "metadata": {
+ "ExecuteTime": {
+ "end_time": "2024-12-24T15:57:30.828509Z",
+ "start_time": "2024-12-24T15:57:30.795197Z"
+ }
+ },
+ "cell_type": "code",
+ "source": [
+ "# f1_metric = f1_score_metric()\n",
+ "# eval_results = await eval_answers(instances, answers, f1_metric)\n",
+ "# avg_f1_score = statistics.mean([result.metrics_data[0].score for result in eval_results.test_results])\n",
+ "# print(\"F1 score: \", avg_f1_score)"
+ ],
+ "id": "72ba5f89cccbee6b",
+ "outputs": [
+ {
+ "data": {
+ "text/plain": [
+ "✨ You're running DeepEval's latest \u001B[38;2;106;0;255mOfficial hotpot F1 score Metric\u001B[0m! \u001B[1;38;2;55;65;81m(\u001B[0m\u001B[38;2;55;65;81musing \u001B[0m\u001B[3;38;2;55;65;81mNone\u001B[0m\u001B[38;2;55;65;81m, \u001B[0m\u001B[38;2;55;65;81mstrict\u001B[0m\u001B[38;2;55;65;81m=\u001B[0m\u001B[3;38;2;55;65;81mFalse\u001B[0m\u001B[38;2;55;65;81m, \u001B[0m\u001B[38;2;55;65;81masync_mode\u001B[0m\u001B[38;2;55;65;81m=\u001B[0m\u001B[3;38;2;55;65;81mTrue\u001B[0m\u001B[1;38;2;55;65;81m)\u001B[0m\u001B[38;2;55;65;81m...\u001B[0m\n"
+ ],
+ "text/html": [
+ "
✨ You're running DeepEval's latest Official hotpot F1 score Metric! (using None, strict=False, async_mode=True)...\n",
+ "
\n"
+ ]
+ },
+ "metadata": {},
+ "output_type": "display_data"
+ },
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "Event loop is already running. Applying nest_asyncio patch to allow async execution...\n"
+ ]
+ },
+ {
+ "name": "stderr",
+ "output_type": "stream",
+ "text": [
+ "Evaluating 10 test case(s) in parallel: |██████████|100% (10/10) [Time Taken: 00:00, 407.84test case/s]"
+ ]
+ },
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "\n",
+ "======================================================================\n",
+ "\n",
+ "Metrics Summary\n",
+ "\n",
+ " - ❌ Official hotpot F1 score (score: 0.0, threshold: 0.5, strict: False, evaluation model: None, reason: None, error: None)\n",
+ "\n",
+ "For test case:\n",
+ "\n",
+ " - input: Were Scott Derrickson and Ed Wood of the same nationality?\n",
+ " - actual output: Scott Derrickson is described as an American filmmaker in the context. Ed Wood is referenced as a cult filmmaker in the American biographical film \"Ed Wood.\" Therefore, both Scott Derrickson and Ed Wood are of the same nationality: American.\n",
+ " - expected output: yes\n",
+ " - context: None\n",
+ " - retrieval context: None\n",
+ "\n",
+ "======================================================================\n",
+ "\n",
+ "Metrics Summary\n",
+ "\n",
+ " - ❌ Official hotpot F1 score (score: 0.04255319148936171, threshold: 0.5, strict: False, evaluation model: None, reason: None, error: None)\n",
+ "\n",
+ "For test case:\n",
+ "\n",
+ " - input: What government position was held by the woman who portrayed Corliss Archer in the film Kiss and Tell?\n",
+ " - actual output: The woman who portrayed Corliss Archer in the film \"Kiss and Tell\" was Shirley Temple. After her acting career, she held the position of U.S. Ambassador to Czechoslovakia from 1989 to 1992, and she was also a member of the U.S. delegation to the United Nations from 1969 to 1970.\n",
+ " - expected output: Chief of Protocol\n",
+ " - context: None\n",
+ " - retrieval context: None\n",
+ "\n",
+ "======================================================================\n",
+ "\n",
+ "Metrics Summary\n",
+ "\n",
+ " - ❌ Official hotpot F1 score (score: 0.05882352941176471, threshold: 0.5, strict: False, evaluation model: None, reason: None, error: None)\n",
+ "\n",
+ "For test case:\n",
+ "\n",
+ " - input: What science fantasy young adult series, told in first person, has a set of companion books narrating the stories of enslaved worlds and alien species?\n",
+ " - actual output: The book series that fits the description is \"Animorphs\". It is a science fantasy young adult series told in the first person with multiple narrators, and it deals with themes relevant to enslaved worlds and alien species.\n",
+ " - expected output: Animorphs\n",
+ " - context: None\n",
+ " - retrieval context: None\n",
+ "\n",
+ "======================================================================\n",
+ "\n",
+ "Metrics Summary\n",
+ "\n",
+ " - ❌ Official hotpot F1 score (score: 0.0, threshold: 0.5, strict: False, evaluation model: None, reason: None, error: None)\n",
+ "\n",
+ "For test case:\n",
+ "\n",
+ " - input: Are the Laleli Mosque and Esma Sultan Mansion located in the same neighborhood?\n",
+ " - actual output: The Laleli Mosque is located in Laleli, Fatih, Istanbul. The Esma Sultan Mansion is located in the Ortaköy neighborhood of Istanbul. Therefore, they are not located in the same neighborhood.\n",
+ " - expected output: no\n",
+ " - context: None\n",
+ " - retrieval context: None\n",
+ "\n",
+ "======================================================================\n",
+ "\n",
+ "Metrics Summary\n",
+ "\n",
+ " - ❌ Official hotpot F1 score (score: 0.22222222222222218, threshold: 0.5, strict: False, evaluation model: None, reason: None, error: None)\n",
+ "\n",
+ "For test case:\n",
+ "\n",
+ " - input: The director of the romantic comedy \"Big Stone Gap\" is based in what New York city?\n",
+ " - actual output: The director of the romantic comedy \"Big Stone Gap\" is Adriana Trigiani, and there is no specific New York city mentioned in the provided context.\n",
+ " - expected output: Greenwich Village, New York City\n",
+ " - context: None\n",
+ " - retrieval context: None\n",
+ "\n",
+ "======================================================================\n",
+ "\n",
+ "Metrics Summary\n",
+ "\n",
+ " - ❌ Official hotpot F1 score (score: 0.11764705882352941, threshold: 0.5, strict: False, evaluation model: None, reason: None, error: None)\n",
+ "\n",
+ "For test case:\n",
+ "\n",
+ " - input: 2014 S/S is the debut album of a South Korean boy group that was formed by who?\n",
+ " - actual output: The South Korean boy group that debuted with the album \"2014 S/S\" was formed by FNC Entertainment.\n",
+ " - expected output: YG Entertainment\n",
+ " - context: None\n",
+ " - retrieval context: None\n",
+ "\n",
+ "======================================================================\n",
+ "\n",
+ "Metrics Summary\n",
+ "\n",
+ " - ❌ Official hotpot F1 score (score: 0.18181818181818182, threshold: 0.5, strict: False, evaluation model: None, reason: None, error: None)\n",
+ "\n",
+ "For test case:\n",
+ "\n",
+ " - input: Who was known by his stage name Aladin and helped organizations improve their performance as a consultant?\n",
+ " - actual output: Eenasul Fateh, known by his stage name Aladin, is a Bangladeshi-British cultural practitioner, magician, live artist, and former international management consultant.\n",
+ " - expected output: Eenasul Fateh\n",
+ " - context: None\n",
+ " - retrieval context: None\n",
+ "\n",
+ "======================================================================\n",
+ "\n",
+ "Metrics Summary\n",
+ "\n",
+ " - ❌ Official hotpot F1 score (score: 0.0, threshold: 0.5, strict: False, evaluation model: None, reason: None, error: None)\n",
+ "\n",
+ "For test case:\n",
+ "\n",
+ " - input: The arena where the Lewiston Maineiacs played their home games can seat how many people?\n",
+ " - actual output: The Lewiston Maineiacs played their home games at the Androscoggin Bank Colisée. The capacity of the arena is not directly stated in the context, but it is implied from the information provided.\n",
+ " - expected output: 3,677 seated\n",
+ " - context: None\n",
+ " - retrieval context: None\n",
+ "\n",
+ "======================================================================\n",
+ "\n",
+ "Metrics Summary\n",
+ "\n",
+ " - ❌ Official hotpot F1 score (score: 0.12903225806451613, threshold: 0.5, strict: False, evaluation model: None, reason: None, error: None)\n",
+ "\n",
+ "For test case:\n",
+ "\n",
+ " - input: Who is older, Annie Morton or Terry Richardson?\n",
+ " - actual output: Annie Morton was born on October 8, 1970, while Terry Richardson's birthdate is not provided in the context. Therefore, based on the available information, Annie Morton is older than Terry Richardson.\n",
+ " - expected output: Terry Richardson\n",
+ " - context: None\n",
+ " - retrieval context: None\n",
+ "\n",
+ "======================================================================\n",
+ "\n",
+ "Metrics Summary\n",
+ "\n",
+ " - ❌ Official hotpot F1 score (score: 0.0, threshold: 0.5, strict: False, evaluation model: None, reason: None, error: None)\n",
+ "\n",
+ "For test case:\n",
+ "\n",
+ " - input: Are Local H and For Against both from the United States?\n",
+ " - actual output: Yes, Local H is from the United States, specifically mentioned as a band from the Chicago suburbs, and For Against is also a band mentioned as being from the United States, although specific details about For Against's origin are not provided in the context.\n",
+ " - expected output: yes\n",
+ " - context: None\n",
+ " - retrieval context: None\n",
+ "\n",
+ "======================================================================\n",
+ "\n",
+ "Overall Metric Pass Rates\n",
+ "\n",
+ "Official hotpot F1 score: 0.00% pass rate\n",
+ "\n",
+ "======================================================================\n",
+ "\n"
+ ]
+ },
+ {
+ "name": "stderr",
+ "output_type": "stream",
+ "text": [
+ "\n"
+ ]
+ },
+ {
+ "data": {
+ "text/plain": [
+ "\u001B[38;2;5;245;141m✓\u001B[0m Tests finished 🎉! Run \u001B[32m'deepeval login'\u001B[0m to save and analyze evaluation results on Confident AI. \n",
+ "‼️ Friendly reminder 😇: You can also run evaluations with ALL of deepeval's metrics directly on Confident AI \n",
+ "instead.\n"
+ ],
+ "text/html": [
+ "
✓ Tests finished 🎉! Run 'deepeval login' to save and analyze evaluation results on Confident AI. \n",
+ "‼️ Friendly reminder 😇: You can also run evaluations with ALL of deepeval's metrics directly on Confident AI \n",
+ "instead.\n",
+ "