cognee 0.2.1.dev5__tar.gz → 0.2.1.dev7__tar.gz

{cognee-0.2.1.dev5 → cognee-0.2.1.dev7}/CONTRIBUTING.md

@@ -12,6 +12,7 @@ This guide will help you get started and ensure your contributions can be effici
 - [Code of Conduct](CODE_OF_CONDUCT.md)
 - [Discord Community](https://discord.gg/bcy8xFAtfd)  
 - [Issue Tracker](https://github.com/topoteretes/cognee/issues)
+- [Cognee Docs](https://docs.cognee.ai)
 
 ## 1. 🚀 Ways to Contribute
 
@@ -69,6 +70,13 @@ Looking for a place to start? Try filtering for [good first issues](https://gith
 git clone https://github.com/<your-github-username>/cognee.git
 cd cognee
 ```
+In case you are working on Vector and Graph Adapters
+1. Fork the [**cognee**](https://github.com/topoteretes/cognee-community) repository
+2. Clone your fork:
+```shell
+git clone https://github.com/<your-github-username>/cognee-community.git
+cd cognee-community
+```
 
 ### Create a Branch
 

{cognee-0.2.1.dev5 → cognee-0.2.1.dev7}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cognee
-Version: 0.2.1.dev5
+Version: 0.2.1.dev7
 Summary: Cognee - is a library for enriching LLM context with a semantic layer for better understanding and reasoning.
 Project-URL: Homepage, https://www.cognee.ai
 Project-URL: Repository, https://github.com/topoteretes/cognee
@@ -24,7 +24,6 @@ Requires-Dist: dlt[sqlalchemy]<2,>=1.9.0
 Requires-Dist: fastapi-users[sqlalchemy]<15.0.0,>=14.0.1
 Requires-Dist: fastapi<1.0.0,>=0.115.7
 Requires-Dist: filetype<2.0.0,>=1.2.0
-Requires-Dist: graphistry<0.34,>=0.33.5
 Requires-Dist: instructor<2.0.0,>=1.9.1
 Requires-Dist: jinja2<4,>=3.1.3
 Requires-Dist: kuzu==0.11.0
@@ -35,7 +34,7 @@ Requires-Dist: litellm<1.71.0,>=1.57.4
 Requires-Dist: matplotlib<4,>=3.8.3
 Requires-Dist: networkx<4,>=3.4.2
 Requires-Dist: nltk<4.0.0,>=3.9.1
-Requires-Dist: numpy<=2.3.1,>=1.26.4
+Requires-Dist: numpy<=4.0.0,>=1.26.4
 Requires-Dist: onnxruntime<2.0.0,>=1.0.0
 Requires-Dist: openai<2,>=1.80.1
 Requires-Dist: pandas<3.0.0,>=2.2.2
@@ -131,7 +130,6 @@ Provides-Extra: postgres-binary
 Requires-Dist: asyncpg<1.0.0,>=0.30.0; extra == 'postgres-binary'
 Requires-Dist: pgvector<0.4,>=0.3.5; extra == 'postgres-binary'
 Requires-Dist: psycopg2-binary<3.0.0,>=2.9.10; extra == 'postgres-binary'
-Requires-Dist: psycopg2<3,>=2.9.10; extra == 'postgres-binary'
 Provides-Extra: posthog
 Requires-Dist: posthog<4,>=3.5.0; extra == 'posthog'
 Provides-Extra: qdrant

cognee-0.2.1.dev7/cognee/api/.env.example

@@ -0,0 +1,5 @@
+# Comma-separated list of allowed origins for CORS (leave empty to block all cross-origin requests)
+# Example:
+# CORS_ALLOWED_ORIGINS="https://yourdomain.com,https://another.com"
+# For local development, you might use:
+# CORS_ALLOWED_ORIGINS="http://localhost:3000"

{cognee-0.2.1.dev5 → cognee-0.2.1.dev7}/cognee/api/client.py

@@ -67,13 +67,24 @@ async def lifespan(app: FastAPI):
 app = FastAPI(debug=app_environment != "prod", lifespan=lifespan)
 
 
+# Read allowed origins from environment variable (comma-separated)
+CORS_ALLOWED_ORIGINS = os.getenv("CORS_ALLOWED_ORIGINS")
+if CORS_ALLOWED_ORIGINS:
+    allowed_origins = [
+        origin.strip() for origin in CORS_ALLOWED_ORIGINS.split(",") if origin.strip()
+    ]
+else:
+    allowed_origins = []  # Block all except explicitly set origins
+
 app.add_middleware(
     CORSMiddleware,
-    allow_origins=["http://localhost:3000", "http://127.0.0.1:3000"],
+    allow_origins=allowed_origins,  # Now controlled by env var
     allow_credentials=True,
     allow_methods=["OPTIONS", "GET", "POST", "DELETE"],
     allow_headers=["*"],
 )
+# To allow origins, set CORS_ALLOWED_ORIGINS env variable to a comma-separated list, e.g.:
+# CORS_ALLOWED_ORIGINS="https://yourdomain.com,https://another.com"
 
 
 def custom_openapi():

{cognee-0.2.1.dev5 → cognee-0.2.1.dev7}/cognee/api/v1/add/add.py

@@ -14,7 +14,7 @@ async def add(
     node_set: Optional[List[str]] = None,
     vector_db_config: dict = None,
     graph_db_config: dict = None,
-    dataset_id: UUID = None,
+    dataset_id: Optional[UUID] = None,
 ):
     """
     Add data to Cognee for knowledge graph processing.

{cognee-0.2.1.dev5 → cognee-0.2.1.dev7}/cognee/api/v1/add/routers/get_add_router.py

@@ -1,16 +1,17 @@
 import os
+import requests
+import subprocess
 from uuid import UUID
 
-from fastapi import Form, UploadFile, Depends
-from fastapi.responses import JSONResponse
 from fastapi import APIRouter
-from typing import List, Optional
-import subprocess
-from cognee.shared.logging_utils import get_logger
-import requests
+from fastapi.responses import JSONResponse
+from fastapi import Form, File, UploadFile, Depends
+from typing import List, Optional, Union, Literal
 
 from cognee.modules.users.models import User
 from cognee.modules.users.methods import get_authenticated_user
+from cognee.shared.utils import send_telemetry
+from cognee.shared.logging_utils import get_logger
 
 logger = get_logger()
 
@@ -20,9 +21,9 @@ def get_add_router() -> APIRouter:
 
     @router.post("", response_model=dict)
     async def add(
-        data: List[UploadFile],
+        data: List[UploadFile] = File(default=None),
         datasetName: Optional[str] = Form(default=None),
-        datasetId: Optional[UUID] = Form(default=None),
+        datasetId: Union[UUID, Literal[""], None] = Form(default=None, examples=[""]),
         user: User = Depends(get_authenticated_user),
     ):
         """
@@ -38,7 +39,7 @@ def get_add_router() -> APIRouter:
           - GitHub repository URLs (will be cloned and processed)
           - Regular file uploads
         - **datasetName** (Optional[str]): Name of the dataset to add data to
-        - **datasetId** (Optional[UUID]): UUID of the dataset to add data to
+        - **datasetId** (Optional[UUID]): UUID of an already existing dataset
 
         Either datasetName or datasetId must be provided.
 
@@ -58,7 +59,16 @@ def get_add_router() -> APIRouter:
         - GitHub repositories are cloned and all files are processed
         - HTTP URLs are fetched and their content is processed
         - The ALLOW_HTTP_REQUESTS environment variable controls URL processing
+        - datasetId value can only be the UUID of an already existing dataset
         """
+        send_telemetry(
+            "Add API Endpoint Invoked",
+            user.id,
+            additional_properties={
+                "endpoint": "POST /v1/add",
+            },
+        )
+
         from cognee.api.v1.add import add as cognee_add
 
         if not datasetId and not datasetName:

{cognee-0.2.1.dev5 → cognee-0.2.1.dev7}/cognee/api/v1/cognify/routers/get_cognify_router.py

@@ -1,7 +1,7 @@
 import os
 import asyncio
 from uuid import UUID
-from pydantic import BaseModel
+from pydantic import Field
 from typing import List, Optional
 from fastapi.responses import JSONResponse
 from fastapi import APIRouter, WebSocket, Depends, WebSocketDisconnect
@@ -10,7 +10,6 @@ from starlette.status import WS_1000_NORMAL_CLOSURE, WS_1008_POLICY_VIOLATION
 from cognee.api.DTO import InDTO
 from cognee.modules.pipelines.methods import get_pipeline_run
 from cognee.modules.users.models import User
-from cognee.shared.data_models import KnowledgeGraph
 from cognee.modules.users.methods import get_authenticated_user
 from cognee.modules.users.get_user_db import get_user_db_context
 from cognee.modules.graph.methods import get_formatted_graph_data
@@ -24,15 +23,16 @@ from cognee.modules.pipelines.queues.pipeline_run_info_queues import (
     remove_queue,
 )
 from cognee.shared.logging_utils import get_logger
+from cognee.shared.utils import send_telemetry
 
 
 logger = get_logger("api.cognify")
 
 
 class CognifyPayloadDTO(InDTO):
-    datasets: Optional[List[str]] = None
-    dataset_ids: Optional[List[UUID]] = None
-    run_in_background: Optional[bool] = False
+    datasets: Optional[List[str]] = Field(default=None)
+    dataset_ids: Optional[List[UUID]] = Field(default=None, examples=[[]])
+    run_in_background: Optional[bool] = Field(default=False)
 
 
 def get_cognify_router() -> APIRouter:
@@ -57,8 +57,7 @@ def get_cognify_router() -> APIRouter:
 
         ## Request Parameters
         - **datasets** (Optional[List[str]]): List of dataset names to process. Dataset names are resolved to datasets owned by the authenticated user.
-        - **dataset_ids** (Optional[List[UUID]]): List of dataset UUIDs to process. UUIDs allow processing of datasets not owned by the user (if permitted).
-        - **graph_model** (Optional[BaseModel]): Custom Pydantic model defining the knowledge graph schema. Defaults to KnowledgeGraph for general-purpose processing.
+        - **dataset_ids** (Optional[List[UUID]]): List of existing dataset UUIDs to process. UUIDs allow processing of datasets not owned by the user (if permitted).
         - **run_in_background** (Optional[bool]): Whether to execute processing asynchronously. Defaults to False (blocking).
 
         ## Response
@@ -84,6 +83,14 @@ def get_cognify_router() -> APIRouter:
         ## Next Steps
         After successful processing, use the search endpoints to query the generated knowledge graph for insights, relationships, and semantic search.
         """
+        send_telemetry(
+            "Cognify API Endpoint Invoked",
+            user.id,
+            additional_properties={
+                "endpoint": "POST /v1/cognify",
+            },
+        )
+
         if not payload.datasets and not payload.dataset_ids:
             return JSONResponse(
                 status_code=400, content={"error": "No datasets or dataset_ids provided"}

{cognee-0.2.1.dev5 → cognee-0.2.1.dev7}/cognee/api/v1/config/config.py

@@ -182,15 +182,3 @@ class config:
     def set_vector_db_url(db_url: str):
         vector_db_config = get_vectordb_config()
         vector_db_config.vector_db_url = db_url
-
-    @staticmethod
-    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."
-            )
-
-        base_config.graphistry_username = graphistry_config.get("username")
-        base_config.graphistry_password = graphistry_config.get("password")

{cognee-0.2.1.dev5 → cognee-0.2.1.dev7}/cognee/api/v1/datasets/routers/get_datasets_router.py

@@ -22,6 +22,7 @@ from cognee.modules.users.permissions.methods import (
 )
 from cognee.modules.graph.methods import get_formatted_graph_data
 from cognee.modules.pipelines.models import PipelineRunStatus
+from cognee.shared.utils import send_telemetry
 
 logger = get_logger()
 
@@ -92,6 +93,14 @@ def get_datasets_router() -> APIRouter:
         ## Error Codes
         - **418 I'm a teapot**: Error retrieving datasets
         """
+        send_telemetry(
+            "Datasets API Endpoint Invoked",
+            user.id,
+            additional_properties={
+                "endpoint": "GET /v1/datasets",
+            },
+        )
+
         try:
             datasets = await get_all_user_permission_datasets(user, "read")
 
@@ -130,6 +139,14 @@ def get_datasets_router() -> APIRouter:
         ## Error Codes
         - **418 I'm a teapot**: Error creating dataset
         """
+        send_telemetry(
+            "Datasets API Endpoint Invoked",
+            user.id,
+            additional_properties={
+                "endpoint": "POST /v1/datasets",
+            },
+        )
+
         try:
             datasets = await get_datasets_by_name([dataset_data.name], user.id)
 
@@ -175,6 +192,15 @@ def get_datasets_router() -> APIRouter:
         - **404 Not Found**: Dataset doesn't exist or user doesn't have access
         - **500 Internal Server Error**: Error during deletion
         """
+        send_telemetry(
+            "Datasets API Endpoint Invoked",
+            user.id,
+            additional_properties={
+                "endpoint": f"DELETE /v1/datasets/{str(dataset_id)}",
+                "dataset_id": str(dataset_id),
+            },
+        )
+
         from cognee.modules.data.methods import get_dataset, delete_dataset
 
         dataset = await get_dataset(user.id, dataset_id)
@@ -210,6 +236,16 @@ def get_datasets_router() -> APIRouter:
         - **404 Not Found**: Dataset or data item doesn't exist, or user doesn't have access
         - **500 Internal Server Error**: Error during deletion
         """
+        send_telemetry(
+            "Datasets API Endpoint Invoked",
+            user.id,
+            additional_properties={
+                "endpoint": f"DELETE /v1/datasets/{str(dataset_id)}/data/{str(data_id)}",
+                "dataset_id": str(dataset_id),
+                "data_id": str(data_id),
+            },
+        )
+
         from cognee.modules.data.methods import get_data, delete_data
         from cognee.modules.data.methods import get_dataset
 
@@ -288,6 +324,15 @@ def get_datasets_router() -> APIRouter:
         - **404 Not Found**: Dataset doesn't exist or user doesn't have access
         - **500 Internal Server Error**: Error retrieving data
         """
+        send_telemetry(
+            "Datasets API Endpoint Invoked",
+            user.id,
+            additional_properties={
+                "endpoint": f"GET /v1/datasets/{str(dataset_id)}/data",
+                "dataset_id": str(dataset_id),
+            },
+        )
+
         from cognee.modules.data.methods import get_dataset_data, get_dataset
 
         # Verify user has permission to read dataset
@@ -331,6 +376,15 @@ def get_datasets_router() -> APIRouter:
         ## Error Codes
         - **500 Internal Server Error**: Error retrieving status information
         """
+        send_telemetry(
+            "Datasets API Endpoint Invoked",
+            user.id,
+            additional_properties={
+                "endpoint": "GET /v1/datasets/status",
+                "datasets": [str(dataset_id) for dataset_id in datasets],
+            },
+        )
+
         from cognee.api.v1.datasets.datasets import datasets as cognee_datasets
 
         try:
@@ -367,6 +421,16 @@ def get_datasets_router() -> APIRouter:
         - **404 Not Found**: Dataset or data item doesn't exist, or user doesn't have access
         - **500 Internal Server Error**: Error accessing the raw data file
         """
+        send_telemetry(
+            "Datasets API Endpoint Invoked",
+            user.id,
+            additional_properties={
+                "endpoint": f"GET /v1/datasets/{str(dataset_id)}/data/{str(data_id)}/raw",
+                "dataset_id": str(dataset_id),
+                "data_id": str(data_id),
+            },
+        )
+
         from cognee.modules.data.methods import get_data
         from cognee.modules.data.methods import get_dataset_data
 

{cognee-0.2.1.dev5 → cognee-0.2.1.dev7}/cognee/api/v1/delete/routers/get_delete_router.py

@@ -5,6 +5,7 @@ from uuid import UUID
 from cognee.shared.logging_utils import get_logger
 from cognee.modules.users.models import User
 from cognee.modules.users.methods import get_authenticated_user
+from cognee.shared.utils import send_telemetry
 
 logger = get_logger()
 
@@ -31,6 +32,16 @@ def get_delete_router() -> APIRouter:
             JSON response indicating success or failure
 
         """
+        send_telemetry(
+            "Delete API Endpoint Invoked",
+            user.id,
+            additional_properties={
+                "endpoint": "DELETE /v1/delete",
+                "dataset_id": str(dataset_id),
+                "data_id": str(data_id),
+            },
+        )
+
         from cognee.api.v1.delete import delete as cognee_delete
 
         try:

{cognee-0.2.1.dev5 → cognee-0.2.1.dev7}/cognee/api/v1/permissions/routers/get_permissions_router.py

@@ -6,12 +6,13 @@ from fastapi.responses import JSONResponse
 
 from cognee.modules.users.models import User
 from cognee.modules.users.methods import get_authenticated_user
+from cognee.shared.utils import send_telemetry
 
 
 def get_permissions_router() -> APIRouter:
     permissions_router = APIRouter()
 
-    @permissions_router.post("/datasets/{principal_id}/")
+    @permissions_router.post("/datasets/{principal_id}")
     async def give_datasets_permission_to_principal(
         permission_name: str,
         dataset_ids: List[UUID],
@@ -40,6 +41,16 @@ def get_permissions_router() -> APIRouter:
         - **403 Forbidden**: User doesn't have permission to grant access
         - **500 Internal Server Error**: Error granting permission
         """
+        send_telemetry(
+            "Permissions API Endpoint Invoked",
+            user.id,
+            additional_properties={
+                "endpoint": f"POST /v1/permissions/datasets/{str(principal_id)}",
+                "dataset_ids": str(dataset_ids),
+                "principal_id": str(principal_id),
+            },
+        )
+
         from cognee.modules.users.permissions.methods import authorized_give_permission_on_datasets
 
         await authorized_give_permission_on_datasets(
@@ -72,6 +83,15 @@ def get_permissions_router() -> APIRouter:
         - **400 Bad Request**: Invalid role name or role already exists
         - **500 Internal Server Error**: Error creating the role
         """
+        send_telemetry(
+            "Permissions API Endpoint Invoked",
+            user.id,
+            additional_properties={
+                "endpoint": "POST /v1/permissions/roles",
+                "role_name": role_name,
+            },
+        )
+
         from cognee.modules.users.roles.methods import create_role as create_role_method
 
         await create_role_method(role_name=role_name, owner_id=user.id)
@@ -104,6 +124,16 @@ def get_permissions_router() -> APIRouter:
         - **404 Not Found**: User or role doesn't exist
         - **500 Internal Server Error**: Error adding user to role
         """
+        send_telemetry(
+            "Permissions API Endpoint Invoked",
+            user.id,
+            additional_properties={
+                "endpoint": f"POST /v1/permissions/users/{str(user_id)}/roles",
+                "user_id": str(user_id),
+                "role_id": str(role_id),
+            },
+        )
+
         from cognee.modules.users.roles.methods import add_user_to_role as add_user_to_role_method
 
         await add_user_to_role_method(user_id=user_id, role_id=role_id, owner_id=user.id)
@@ -136,6 +166,16 @@ def get_permissions_router() -> APIRouter:
         - **404 Not Found**: User or tenant doesn't exist
         - **500 Internal Server Error**: Error adding user to tenant
         """
+        send_telemetry(
+            "Permissions API Endpoint Invoked",
+            user.id,
+            additional_properties={
+                "endpoint": f"POST /v1/permissions/users/{str(user_id)}/tenants",
+                "user_id": str(user_id),
+                "tenant_id": str(tenant_id),
+            },
+        )
+
         from cognee.modules.users.tenants.methods import add_user_to_tenant
 
         await add_user_to_tenant(user_id=user_id, tenant_id=tenant_id, owner_id=user.id)
@@ -161,6 +201,15 @@ def get_permissions_router() -> APIRouter:
         - **400 Bad Request**: Invalid tenant name or tenant already exists
         - **500 Internal Server Error**: Error creating the tenant
         """
+        send_telemetry(
+            "Permissions API Endpoint Invoked",
+            user.id,
+            additional_properties={
+                "endpoint": "POST /v1/permissions/tenants",
+                "tenant_name": tenant_name,
+            },
+        )
+
         from cognee.modules.users.tenants.methods import create_tenant as create_tenant_method
 
         await create_tenant_method(tenant_name=tenant_name, user_id=user.id)

{cognee-0.2.1.dev5 → cognee-0.2.1.dev7}/cognee/api/v1/search/routers/get_search_router.py

@@ -1,6 +1,7 @@
 from uuid import UUID
 from typing import Optional
 from datetime import datetime
+from pydantic import Field
 from fastapi import Depends, APIRouter
 from fastapi.responses import JSONResponse
 from cognee.modules.search.types import SearchType
@@ -9,16 +10,17 @@ from cognee.modules.users.exceptions.exceptions import PermissionDeniedError
 from cognee.modules.users.models import User
 from cognee.modules.search.operations import get_history
 from cognee.modules.users.methods import get_authenticated_user
+from cognee.shared.utils import send_telemetry
 
 
 # Note: Datasets sent by name will only map to datasets owned by the request sender
 #       To search for datasets not owned by the request sender dataset UUID is needed
 class SearchPayloadDTO(InDTO):
-    search_type: SearchType
-    datasets: Optional[list[str]] = None
-    dataset_ids: Optional[list[UUID]] = None
-    query: str
-    top_k: Optional[int] = 10
+    search_type: SearchType = Field(default=SearchType.GRAPH_COMPLETION)
+    datasets: Optional[list[str]] = Field(default=None)
+    dataset_ids: Optional[list[UUID]] = Field(default=None, examples=[[]])
+    query: str = Field(default="What is in the document?")
+    top_k: Optional[int] = Field(default=10)
 
 
 def get_search_router() -> APIRouter:
@@ -48,6 +50,14 @@ def get_search_router() -> APIRouter:
         ## Error Codes
         - **500 Internal Server Error**: Error retrieving search history
         """
+        send_telemetry(
+            "Search API Endpoint Invoked",
+            user.id,
+            additional_properties={
+                "endpoint": "GET /v1/search",
+            },
+        )
+
         try:
             history = await get_history(user.id, limit=0)
 
@@ -83,6 +93,19 @@ def get_search_router() -> APIRouter:
         - To search datasets not owned by the request sender, dataset UUID is needed
         - If permission is denied, returns empty list instead of error
         """
+        send_telemetry(
+            "Search API Endpoint Invoked",
+            user.id,
+            additional_properties={
+                "endpoint": "POST /v1/search",
+                "search_type": str(payload.search_type),
+                "datasets": payload.datasets,
+                "dataset_ids": [str(dataset_id) for dataset_id in payload.dataset_ids or []],
+                "query": payload.query,
+                "top_k": payload.top_k,
+            },
+        )
+
         from cognee.api.v1.search import search as cognee_search
 
         try:

{cognee-0.2.1.dev5 → cognee-0.2.1.dev7}/cognee/api/v1/users/routers/get_visualize_router.py

@@ -7,6 +7,7 @@ from cognee.modules.data.methods import get_authorized_existing_datasets
 from cognee.modules.users.models import User
 
 from cognee.context_global_variables import set_database_global_context_variables
+from cognee.shared.utils import send_telemetry
 
 logger = get_logger()
 
@@ -39,6 +40,15 @@ def get_visualize_router() -> APIRouter:
         - User must have read permissions on the dataset
         - Visualization is interactive and allows graph exploration
         """
+        send_telemetry(
+            "Visualize API Endpoint Invoked",
+            user.id,
+            additional_properties={
+                "endpoint": "GET /v1/visualize",
+                "dataset_id": str(dataset_id),
+            },
+        )
+
         from cognee.api.v1.visualize import visualize_graph
 
         try:

{cognee-0.2.1.dev5 → cognee-0.2.1.dev7}/cognee/base_config.py

@@ -10,8 +10,6 @@ class BaseConfig(BaseSettings):
     data_root_directory: str = get_absolute_path(".data_storage")
     system_root_directory: str = get_absolute_path(".cognee_system")
     monitoring_tool: object = Observer.LANGFUSE
-    graphistry_username: Optional[str] = os.getenv("GRAPHISTRY_USERNAME")
-    graphistry_password: Optional[str] = os.getenv("GRAPHISTRY_PASSWORD")
     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")

{cognee-0.2.1.dev5 → cognee-0.2.1.dev7}/cognee/infrastructure/files/storage/S3FileStorage.py

@@ -26,7 +26,8 @@ class S3FileStorage(Storage):
                 key=s3_config.aws_access_key_id,
                 secret=s3_config.aws_secret_access_key,
                 anon=False,
-                endpoint_url="https://s3-eu-west-1.amazonaws.com",
+                endpoint_url=s3_config.aws_endpoint_url,
+                client_kwargs={"region_name": s3_config.aws_region},
             )
         else:
             raise ValueError("S3 credentials are not set in the configuration.")

{cognee-0.2.1.dev5 → cognee-0.2.1.dev7}/cognee/infrastructure/files/storage/s3_config.py

@@ -4,6 +4,8 @@ from pydantic_settings import BaseSettings, SettingsConfigDict
 
 
 class S3Config(BaseSettings):
+    aws_region: Optional[str] = None
+    aws_endpoint_url: Optional[str] = None
     aws_access_key_id: Optional[str] = None
     aws_secret_access_key: Optional[str] = None
     model_config = SettingsConfigDict(env_file=".env", extra="allow")

{cognee-0.2.1.dev5 → cognee-0.2.1.dev7}/cognee/modules/graph/utils/expand_with_nodes_and_edges.py

@@ -351,7 +351,7 @@ def expand_with_nodes_and_edges(
         _process_graph_edges(graph, name_mapping, existing_edges_map, relationships)
 
     # Return combined results
-    graph_nodes = list(added_ontology_nodes_map.values())
+    graph_nodes = data_chunks + list(added_ontology_nodes_map.values())
     graph_edges = relationships + ontology_relationships
 
     return graph_nodes, graph_edges

{cognee-0.2.1.dev5 → cognee-0.2.1.dev7}/cognee/modules/ingestion/classify.py

@@ -12,9 +12,7 @@ def classify(data: Union[str, BinaryIO], filename: str = None):
         return TextData(data)
 
     if isinstance(data, BufferedReader) or isinstance(data, SpooledTemporaryFile):
-        return BinaryData(
-            data, str(data.name).split("/")[-1] if hasattr(data, "name") else filename
-        )
+        return BinaryData(data, filename if filename else str(data.name).split("/")[-1])
 
     try:
         from s3fs import S3File

{cognee-0.2.1.dev5 → cognee-0.2.1.dev7}/cognee/modules/retrieval/utils/brute_force_triplet_search.py

@@ -155,12 +155,14 @@ async def brute_force_search(
         logger.error("Failed to initialize vector engine: %s", e)
         raise RuntimeError("Initialization error") from e
 
+    query_vector = (await vector_engine.embedding_engine.embed_text([query]))[0]
+
     send_telemetry("cognee.brute_force_triplet_search EXECUTION STARTED", user.id)
 
     async def search_in_collection(collection_name: str):
         try:
             return await vector_engine.search(
-                collection_name=collection_name, query_text=query, limit=0
+                collection_name=collection_name, query_vector=query_vector, limit=0
             )
         except CollectionNotFoundError:
             return []

{cognee-0.2.1.dev5 → cognee-0.2.1.dev7}/cognee/shared/utils.py

@@ -3,9 +3,7 @@
 import os
 import requests
 from datetime import datetime, timezone
-import graphistry
 import networkx as nx
-import pandas as pd
 import matplotlib.pyplot as plt
 import http.server
 import socketserver

{cognee-0.2.1.dev5 → cognee-0.2.1.dev7}/cognee/tasks/graph/extract_graph_from_data.py

@@ -11,6 +11,7 @@ from cognee.modules.graph.utils import (
     retrieve_existing_edges,
 )
 from cognee.shared.data_models import KnowledgeGraph
+from cognee.tasks.storage.add_data_points import add_data_points
 
 
 async def integrate_chunk_graphs(
@@ -38,8 +39,7 @@ async def integrate_chunk_graphs(
     )
 
     if len(graph_nodes) > 0:
-        await graph_engine.add_nodes(graph_nodes)
-
+        await add_data_points(graph_nodes)
     if len(graph_edges) > 0:
         await graph_engine.add_edges(graph_edges)