cognee 0.2.1.dev7__py3-none-any.whl → 0.2.2.dev1__py3-none-any.whl

cognee/tests/{test_weaviate.py → test_neptune_analytics_vector.py}

@@ -1,26 +1,34 @@
 import os
 import pathlib
 import cognee
-from cognee.infrastructure.files.storage import get_storage_config
+import uuid
+import pytest
 from cognee.modules.search.operations import get_history
 from cognee.modules.users.methods import get_default_user
 from cognee.shared.logging_utils import get_logger
 from cognee.modules.search.types import SearchType
+from cognee.infrastructure.databases.vector import get_vector_engine
+from cognee.infrastructure.databases.hybrid.neptune_analytics.NeptuneAnalyticsAdapter import (
+    NeptuneAnalyticsAdapter,
+    IndexSchema,
+)
 
 logger = get_logger()
 
 
 async def main():
-    cognee.config.set_vector_db_provider("weaviate")
+    graph_id = os.getenv("GRAPH_ID", "")
+    cognee.config.set_vector_db_provider("neptune_analytics")
+    cognee.config.set_vector_db_url(f"neptune-graph://{graph_id}")
     data_directory_path = str(
         pathlib.Path(
-            os.path.join(pathlib.Path(__file__).parent, ".data_storage/test_weaviate")
+            os.path.join(pathlib.Path(__file__).parent, ".data_storage/test_neptune")
         ).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")
+            os.path.join(pathlib.Path(__file__).parent, ".cognee_system/test_neptune")
         ).resolve()
     )
     cognee.config.system_root_directory(cognee_directory_path)
@@ -47,14 +55,12 @@ 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(
-        query_text=random_node_name, query_type=SearchType.INSIGHTS
+        query_type=SearchType.INSIGHTS, query_text=random_node_name
     )
     assert len(search_results) != 0, "The search results list is empty."
     print("\n\nExtracted sentences are:\n")
@@ -80,15 +86,84 @@ async def main():
     assert len(history) == 6, "Search history is not correct."
 
     await cognee.prune.prune_data()
-    data_root_directory = get_storage_config()["data_root_directory"]
-    assert not os.path.isdir(data_root_directory), "Local data files are not deleted"
+    assert not os.path.isdir(data_directory_path), "Local data files are not deleted"
 
     await cognee.prune.prune_system(metadata=True)
-    collections = await get_vector_engine().client.collections.list_all()
-    assert len(collections) == 0, "Weaviate vector database is not empty"
+
+
+async def vector_backend_api_test():
+    cognee.config.set_vector_db_provider("neptune_analytics")
+
+    # When URL is absent
+    cognee.config.set_vector_db_url(None)
+    with pytest.raises(OSError):
+        get_vector_engine()
+
+    # Assert invalid graph ID.
+    cognee.config.set_vector_db_url("invalid_url")
+    with pytest.raises(ValueError):
+        get_vector_engine()
+
+    # Return a valid engine object with valid URL.
+    graph_id = os.getenv("GRAPH_ID", "")
+    cognee.config.set_vector_db_url(f"neptune-graph://{graph_id}")
+    engine = get_vector_engine()
+    assert isinstance(engine, NeptuneAnalyticsAdapter)
+
+    TEST_COLLECTION_NAME = "test"
+    # Data point - 1
+    TEST_UUID = str(uuid.uuid4())
+    TEST_TEXT = "Hello world"
+    datapoint = IndexSchema(id=TEST_UUID, text=TEST_TEXT)
+    # Data point - 2
+    TEST_UUID_2 = str(uuid.uuid4())
+    TEST_TEXT_2 = "Cognee"
+    datapoint_2 = IndexSchema(id=TEST_UUID_2, text=TEST_TEXT_2)
+
+    # Prun all vector_db entries
+    await engine.prune()
+
+    # Always return true
+    has_collection = await engine.has_collection(TEST_COLLECTION_NAME)
+    assert has_collection
+    # No-op
+    await engine.create_collection(TEST_COLLECTION_NAME, IndexSchema)
+
+    # Save data-points
+    await engine.create_data_points(TEST_COLLECTION_NAME, [datapoint, datapoint_2])
+    # Search single text
+    result_search = await engine.search(
+        collection_name=TEST_COLLECTION_NAME,
+        query_text=TEST_TEXT,
+        query_vector=None,
+        limit=10,
+        with_vector=True,
+    )
+    assert len(result_search) == 2
+
+    # # Retrieve data-points
+    result = await engine.retrieve(TEST_COLLECTION_NAME, [TEST_UUID, TEST_UUID_2])
+    assert any(str(r.id) == TEST_UUID and r.payload["text"] == TEST_TEXT for r in result)
+    assert any(str(r.id) == TEST_UUID_2 and r.payload["text"] == TEST_TEXT_2 for r in result)
+    # Search multiple
+    result_search_batch = await engine.batch_search(
+        collection_name=TEST_COLLECTION_NAME,
+        query_texts=[TEST_TEXT, TEST_TEXT_2],
+        limit=10,
+        with_vectors=False,
+    )
+    assert len(result_search_batch) == 2 and all(len(batch) == 2 for batch in result_search_batch)
+
+    # Delete datapoint from vector store
+    await engine.delete_data_points(TEST_COLLECTION_NAME, [TEST_UUID, TEST_UUID_2])
+
+    # Retrieve should return an empty list.
+    result_deleted = await engine.retrieve(TEST_COLLECTION_NAME, [TEST_UUID])
+    assert result_deleted == []
 
 
 if __name__ == "__main__":
     import asyncio
 
     asyncio.run(main())
+    asyncio.run(vector_backend_api_test())

cognee/tests/test_pgvector.py

@@ -37,16 +37,16 @@ async def test_local_file_deletion(data_text, file_location):
         # Get data entry from database based on file path
         data = (
             await session.scalars(
-                select(Data).where(Data.raw_data_location == "file://" + file_location)
+                select(Data).where(Data.original_data_location == "file://" + file_location)
             )
         ).one()
-        assert os.path.isfile(data.raw_data_location.replace("file://", "")), (
-            f"Data location doesn't exist: {data.raw_data_location}"
+        assert os.path.isfile(data.original_data_location.replace("file://", "")), (
+            f"Data location doesn't exist: {data.original_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.replace("file://", "")), (
-            f"Data location doesn't exists: {data.raw_data_location}"
+        assert os.path.exists(data.original_data_location.replace("file://", "")), (
+            f"Data location doesn't exists: {data.original_data_location}"
         )
 
 

cognee/tests/test_s3.py

@@ -28,13 +28,8 @@ async def main():
     logging.info(type_counts)
     logging.info(edge_type_counts)
 
-    # Assert there is exactly one PdfDocument.
-    assert type_counts.get("PdfDocument", 0) == 1, (
-        f"Expected exactly one PdfDocument, but found {type_counts.get('PdfDocument', 0)}"
-    )
-
     # Assert there is exactly one TextDocument.
-    assert type_counts.get("TextDocument", 0) == 1, (
+    assert type_counts.get("TextDocument", 0) == 2, (
         f"Expected exactly one TextDocument, but found {type_counts.get('TextDocument', 0)}"
     )
 

cognee/tests/unit/infrastructure/databases/test_rate_limiter.py

@@ -1,15 +1,16 @@
 """Tests for the LLM rate limiter."""
 
 import pytest
-import asyncio
-import time
 from unittest.mock import patch
-from cognee.infrastructure.llm.rate_limiter import (
+from cognee.infrastructure.llm.structured_output_framework.litellm_instructor.llm.rate_limiter import (
     llm_rate_limiter,
     rate_limit_async,
     rate_limit_sync,
 )
 
+LLM_RATE_LIMITER = "cognee.infrastructure.llm.structured_output_framework.litellm_instructor.llm.rate_limiter.llm_rate_limiter"
+GET_LLM_CONFIG = "cognee.infrastructure.llm.structured_output_framework.litellm_instructor.llm.rate_limiter.get_llm_config"
+
 
 @pytest.fixture(autouse=True)
 def reset_limiter_singleton():
@@ -20,7 +21,7 @@ def reset_limiter_singleton():
 
 def test_rate_limiter_initialization():
     """Test that the rate limiter can be initialized properly."""
-    with patch("cognee.infrastructure.llm.rate_limiter.get_llm_config") as mock_config:
+    with patch(GET_LLM_CONFIG) as mock_config:
         mock_config.return_value.llm_rate_limit_enabled = True
         mock_config.return_value.llm_rate_limit_requests = 10
         mock_config.return_value.llm_rate_limit_interval = 60  # 1 minute
@@ -34,7 +35,7 @@ def test_rate_limiter_initialization():
 
 def test_rate_limiter_disabled():
     """Test that the rate limiter is disabled by default."""
-    with patch("cognee.infrastructure.llm.rate_limiter.get_llm_config") as mock_config:
+    with patch(GET_LLM_CONFIG) as mock_config:
         mock_config.return_value.llm_rate_limit_enabled = False
 
         limiter = llm_rate_limiter()
@@ -45,7 +46,7 @@ def test_rate_limiter_disabled():
 
 def test_rate_limiter_singleton():
     """Test that the rate limiter is a singleton."""
-    with patch("cognee.infrastructure.llm.rate_limiter.get_llm_config") as mock_config:
+    with patch(GET_LLM_CONFIG) as mock_config:
         mock_config.return_value.llm_rate_limit_enabled = True
         mock_config.return_value.llm_rate_limit_requests = 5
         mock_config.return_value.llm_rate_limit_interval = 60
@@ -58,7 +59,7 @@ def test_rate_limiter_singleton():
 
 def test_sync_decorator():
     """Test the sync decorator."""
-    with patch("cognee.infrastructure.llm.rate_limiter.llm_rate_limiter") as mock_limiter_class:
+    with patch(LLM_RATE_LIMITER) as mock_limiter_class:
         mock_limiter = mock_limiter_class.return_value
         mock_limiter.wait_if_needed.return_value = 0
 
@@ -75,7 +76,7 @@ def test_sync_decorator():
 @pytest.mark.asyncio
 async def test_async_decorator():
     """Test the async decorator."""
-    with patch("cognee.infrastructure.llm.rate_limiter.llm_rate_limiter") as mock_limiter_class:
+    with patch(LLM_RATE_LIMITER) as mock_limiter_class:
         mock_limiter = mock_limiter_class.return_value
 
         # Mock an async method with a coroutine
@@ -96,7 +97,7 @@ async def test_async_decorator():
 
 def test_rate_limiting_actual():
     """Test actual rate limiting behavior with a small window."""
-    with patch("cognee.infrastructure.llm.rate_limiter.get_llm_config") as mock_config:
+    with patch(GET_LLM_CONFIG) as mock_config:
         # Configure for 3 requests per minute
         mock_config.return_value.llm_rate_limit_enabled = True
         mock_config.return_value.llm_rate_limit_requests = 3
@@ -117,7 +118,7 @@ def test_rate_limiting_actual():
 
 def test_rate_limit_60_per_minute():
     """Test rate limiting with the default 60 requests per minute limit."""
-    with patch("cognee.infrastructure.llm.rate_limiter.get_llm_config") as mock_config:
+    with patch(GET_LLM_CONFIG) as mock_config:
         # Configure for default values: 60 requests per 60 seconds
         mock_config.return_value.llm_rate_limit_enabled = True
         mock_config.return_value.llm_rate_limit_requests = 60  # 60 requests

cognee/tests/unit/infrastructure/mock_embedding_engine.py

@@ -4,7 +4,7 @@ from typing import List
 from cognee.infrastructure.databases.vector.embeddings.LiteLLMEmbeddingEngine import (
     LiteLLMEmbeddingEngine,
 )
-from cognee.infrastructure.llm.embedding_rate_limiter import (
+from cognee.infrastructure.databases.vector.embeddings.embedding_rate_limiter import (
     embedding_rate_limit_async,
     embedding_sleep_and_retry_async,
 )

cognee/tests/unit/infrastructure/test_embedding_rate_limiting_realistic.py

@@ -1,13 +1,13 @@
 import os
 import time
 import asyncio
-from functools import lru_cache
 import logging
 
-from cognee.infrastructure.llm.config import LLMConfig, get_llm_config
-from cognee.infrastructure.llm.embedding_rate_limiter import EmbeddingRateLimiter
-from cognee.infrastructure.databases.vector.embeddings.LiteLLMEmbeddingEngine import (
-    LiteLLMEmbeddingEngine,
+from cognee.infrastructure.llm.config import (
+    get_llm_config,
+)
+from cognee.infrastructure.databases.vector.embeddings.embedding_rate_limiter import (
+    EmbeddingRateLimiter,
 )
 from cognee.tests.unit.infrastructure.mock_embedding_engine import MockEmbeddingEngine
 

cognee/tests/unit/infrastructure/test_rate_limiting_realistic.py

@@ -1,11 +1,13 @@
 import asyncio
-import time
 import os
-from functools import lru_cache
 from unittest.mock import patch
 from cognee.shared.logging_utils import get_logger
-from cognee.infrastructure.llm.rate_limiter import llm_rate_limiter
-from cognee.infrastructure.llm.config import get_llm_config
+from cognee.infrastructure.llm.structured_output_framework.litellm_instructor.llm.rate_limiter import (
+    llm_rate_limiter,
+)
+from cognee.infrastructure.llm.config import (
+    get_llm_config,
+)
 
 
 async def test_rate_limiting_realistic():

cognee/tests/unit/infrastructure/test_rate_limiting_retry.py

@@ -1,7 +1,7 @@
 import time
 import asyncio
 from cognee.shared.logging_utils import get_logger
-from cognee.infrastructure.llm.rate_limiter import (
+from cognee.infrastructure.llm.structured_output_framework.litellm_instructor.llm.rate_limiter import (
     sleep_and_retry_sync,
     sleep_and_retry_async,
     is_rate_limit_error,

cognee/tests/unit/interfaces/graph/get_graph_from_model_unit_test.py

@@ -1,6 +1,6 @@
 import pytest
-from typing import List
-from cognee.infrastructure.engine import DataPoint
+from typing import List, Any
+from cognee.infrastructure.engine import DataPoint, Edge
 
 from cognee.modules.graph.utils import get_graph_from_model
 
@@ -28,7 +28,20 @@ class Entity(DataPoint):
     metadata: dict = {"index_fields": ["name"]}
 
 
+class Company(DataPoint):
+    name: str
+    employees: List[Any] = None  # Allow flexible edge system with tuples
+    metadata: dict = {"index_fields": ["name"]}
+
+
+class Employee(DataPoint):
+    name: str
+    role: str
+    metadata: dict = {"index_fields": ["name"]}
+
+
 DocumentChunk.model_rebuild()
+Company.model_rebuild()
 
 
 @pytest.mark.asyncio
@@ -50,7 +63,7 @@ async def test_get_graph_from_model_simple_structure():
     assert len(nodes) == 2, f"Expected 2 nodes, got {len(nodes)}"
     assert len(edges) == 1, f"Expected 1 edges, got {len(edges)}"
 
-    edge_key = str(entity.id) + str(entitytype.id) + "is_type"
+    edge_key = f"{str(entity.id)}_{str(entitytype.id)}_is_type"
     assert edge_key in added_edges, f"Edge {edge_key} not found"
 
 
@@ -149,3 +162,48 @@ async def test_get_graph_from_model_no_contains():
 
     assert len(nodes) == 2, f"Expected 2 nodes, got {len(nodes)}"
     assert len(edges) == 1, f"Expected 1 edge, got {len(edges)}"
+
+
+@pytest.mark.asyncio
+async def test_get_graph_from_model_flexible_edges():
+    """Tests the new flexible edge system with mixed relationships"""
+    # Create employees
+    manager = Employee(name="Manager", role="Manager")
+    sales1 = Employee(name="Sales1", role="Sales")
+    sales2 = Employee(name="Sales2", role="Sales")
+    admin1 = Employee(name="Admin1", role="Admin")
+    admin2 = Employee(name="Admin2", role="Admin")
+
+    # Create company with mixed employee relationships
+    company = Company(
+        name="Test Company",
+        employees=[
+            # Weighted relationship
+            (Edge(weight=0.9, relationship_type="manages"), manager),
+            # Multiple weights relationship
+            (
+                Edge(weights={"performance": 0.8, "experience": 0.7}, relationship_type="employs"),
+                sales1,
+            ),
+            # Simple relationship
+            sales2,
+            # Group relationship
+            (Edge(weights={"team_efficiency": 0.8}, relationship_type="employs"), [admin1, admin2]),
+        ],
+    )
+
+    added_nodes = {}
+    added_edges = {}
+    visited_properties = {}
+
+    nodes, edges = await get_graph_from_model(company, added_nodes, added_edges, visited_properties)
+
+    # Should have 6 nodes: company + 5 employees
+    assert len(nodes) == 6, f"Expected 6 nodes, got {len(nodes)}"
+    # Should have 5 edges: 4 employee relationships
+    assert len(edges) == 5, f"Expected 5 edges, got {len(edges)}"
+
+    # Verify all employees are connected
+    employee_ids = {str(emp.id) for emp in [manager, sales1, sales2, admin1, admin2]}
+    edge_target_ids = {str(edge[1]) for edge in edges}
+    assert employee_ids.issubset(edge_target_ids), "Not all employees are connected"

cognee/tests/unit/modules/retrieval/graph_completion_retriever_test.py

@@ -28,18 +28,38 @@ class TestGraphCompletionRetriever:
 
         class Company(DataPoint):
             name: str
+            description: str
 
         class Person(DataPoint):
             name: str
+            description: str
             works_for: Company
 
-        company1 = Company(name="Figma")
-        company2 = Company(name="Canva")
-        person1 = Person(name="Steve Rodger", works_for=company1)
-        person2 = Person(name="Ike Loma", works_for=company1)
-        person3 = Person(name="Jason Statham", works_for=company1)
-        person4 = Person(name="Mike Broski", works_for=company2)
-        person5 = Person(name="Christina Mayer", works_for=company2)
+        company1 = Company(name="Figma", description="Figma is a company")
+        company2 = Company(name="Canva", description="Canvas is a company")
+        person1 = Person(
+            name="Steve Rodger",
+            description="This is description about Steve Rodger",
+            works_for=company1,
+        )
+        person2 = Person(
+            name="Ike Loma", description="This is description about Ike Loma", works_for=company1
+        )
+        person3 = Person(
+            name="Jason Statham",
+            description="This is description about Jason Statham",
+            works_for=company1,
+        )
+        person4 = Person(
+            name="Mike Broski",
+            description="This is description about Mike Broski",
+            works_for=company2,
+        )
+        person5 = Person(
+            name="Christina Mayer",
+            description="This is description about Christina Mayer",
+            works_for=company2,
+        )
 
         entities = [company1, company2, person1, person2, person3, person4, person5]
 
@@ -49,8 +69,63 @@ class TestGraphCompletionRetriever:
 
         context = await retriever.get_context("Who works at Canva?")
 
-        assert "Mike Broski --[works_for]--> Canva" in context, "Failed to get Mike Broski"
-        assert "Christina Mayer --[works_for]--> Canva" in context, "Failed to get Christina Mayer"
+        # Ensure the top-level sections are present
+        assert "Nodes:" in context, "Missing 'Nodes:' section in context"
+        assert "Connections:" in context, "Missing 'Connections:' section in context"
+
+        # --- Nodes headers ---
+        assert "Node: Steve Rodger" in context, "Missing node header for Steve Rodger"
+        assert "Node: Figma" in context, "Missing node header for Figma"
+        assert "Node: Ike Loma" in context, "Missing node header for Ike Loma"
+        assert "Node: Jason Statham" in context, "Missing node header for Jason Statham"
+        assert "Node: Mike Broski" in context, "Missing node header for Mike Broski"
+        assert "Node: Canva" in context, "Missing node header for Canva"
+        assert "Node: Christina Mayer" in context, "Missing node header for Christina Mayer"
+
+        # --- Node contents ---
+        assert (
+            "__node_content_start__\nThis is description about Steve Rodger\n__node_content_end__"
+            in context
+        ), "Description block for Steve Rodger altered"
+        assert "__node_content_start__\nFigma is a company\n__node_content_end__" in context, (
+            "Description block for Figma altered"
+        )
+        assert (
+            "__node_content_start__\nThis is description about Ike Loma\n__node_content_end__"
+            in context
+        ), "Description block for Ike Loma altered"
+        assert (
+            "__node_content_start__\nThis is description about Jason Statham\n__node_content_end__"
+            in context
+        ), "Description block for Jason Statham altered"
+        assert (
+            "__node_content_start__\nThis is description about Mike Broski\n__node_content_end__"
+            in context
+        ), "Description block for Mike Broski altered"
+        assert "__node_content_start__\nCanvas is a company\n__node_content_end__" in context, (
+            "Description block for Canva altered"
+        )
+        assert (
+            "__node_content_start__\nThis is description about Christina Mayer\n__node_content_end__"
+            in context
+        ), "Description block for Christina Mayer altered"
+
+        # --- Connections ---
+        assert "Steve Rodger --[works_for]--> Figma" in context, (
+            "Connection Steve Rodger→Figma missing or changed"
+        )
+        assert "Ike Loma --[works_for]--> Figma" in context, (
+            "Connection Ike Loma→Figma missing or changed"
+        )
+        assert "Jason Statham --[works_for]--> Figma" in context, (
+            "Connection Jason Statham→Figma missing or changed"
+        )
+        assert "Mike Broski --[works_for]--> Canva" in context, (
+            "Connection Mike Broski→Canva missing or changed"
+        )
+        assert "Christina Mayer --[works_for]--> Canva" in context, (
+            "Connection Christina Mayer→Canva missing or changed"
+        )
 
     @pytest.mark.asyncio
     async def test_graph_completion_context_complex(self):

cognee/tests/unit/modules/search/search_methods_test.py

@@ -155,6 +155,61 @@ async def test_specific_search_chunks(mock_send_telemetry, mock_chunks_retriever
     assert results[0]["content"] == "Chunk result"
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "selected_type, retriever_name, expected_content, top_k",
+    [
+        (SearchType.RAG_COMPLETION, "CompletionRetriever", "RAG result from lucky search", 10),
+        (SearchType.CHUNKS, "ChunksRetriever", "Chunk result from lucky search", 5),
+        (SearchType.SUMMARIES, "SummariesRetriever", "Summary from lucky search", 15),
+        (SearchType.INSIGHTS, "InsightsRetriever", "Insight result from lucky search", 20),
+    ],
+)
+@patch.object(search_module, "select_search_type")
+@patch.object(search_module, "send_telemetry")
+async def test_specific_search_feeling_lucky(
+    mock_send_telemetry,
+    mock_select_search_type,
+    selected_type,
+    retriever_name,
+    expected_content,
+    top_k,
+    mock_user,
+):
+    with patch.object(search_module, retriever_name) as mock_retriever_class:
+        # Setup
+        query = f"test query for {retriever_name}"
+        query_type = SearchType.FEELING_LUCKY
+
+        # Mock the intelligent search type selection
+        mock_select_search_type.return_value = selected_type
+
+        # Mock the retriever
+        mock_retriever_instance = MagicMock()
+        mock_retriever_instance.get_completion = AsyncMock(
+            return_value=[{"content": expected_content}]
+        )
+        mock_retriever_class.return_value = mock_retriever_instance
+
+        # Execute
+        results = await specific_search(query_type, query, mock_user, top_k=top_k)
+
+        # Verify
+        mock_select_search_type.assert_called_once_with(query)
+
+        if retriever_name == "CompletionRetriever":
+            mock_retriever_class.assert_called_once_with(
+                system_prompt_path="answer_simple_question.txt", top_k=top_k
+            )
+        else:
+            mock_retriever_class.assert_called_once_with(top_k=top_k)
+
+        mock_retriever_instance.get_completion.assert_called_once_with(query)
+        mock_send_telemetry.assert_called()
+        assert len(results) == 1
+        assert results[0]["content"] == expected_content
+
+
 @pytest.mark.asyncio
 async def test_specific_search_invalid_type(mock_user):
     # Setup

{cognee-0.2.1.dev7.dist-info → cognee-0.2.2.dev1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cognee
-Version: 0.2.1.dev7
+Version: 0.2.2.dev1
 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
@@ -20,6 +20,7 @@ Requires-Dist: aiofiles<24.0.0,>=23.2.1
 Requires-Dist: aiohttp<4.0.0,>=3.11.14
 Requires-Dist: aiosqlite<1.0.0,>=0.20.0
 Requires-Dist: alembic<2,>=1.13.3
+Requires-Dist: baml-py<0.202.0,>=0.201.0
 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
@@ -118,6 +119,8 @@ Provides-Extra: mistral
 Requires-Dist: mistral-common<2,>=1.5.2; extra == 'mistral'
 Provides-Extra: neo4j
 Requires-Dist: neo4j<6,>=5.28.0; extra == 'neo4j'
+Provides-Extra: neptune
+Requires-Dist: langchain-aws>=0.2.22; extra == 'neptune'
 Provides-Extra: notebook
 Requires-Dist: notebook<8,>=7.1.0; extra == 'notebook'
 Provides-Extra: ollama
@@ -132,10 +135,6 @@ Requires-Dist: pgvector<0.4,>=0.3.5; extra == 'postgres-binary'
 Requires-Dist: psycopg2-binary<3.0.0,>=2.9.10; extra == 'postgres-binary'
 Provides-Extra: posthog
 Requires-Dist: posthog<4,>=3.5.0; extra == 'posthog'
-Provides-Extra: qdrant
-Requires-Dist: qdrant-client<2,>=1.14.2; extra == 'qdrant'
-Provides-Extra: weaviate
-Requires-Dist: weaviate-client<5.0.0,>=4.9.6; extra == 'weaviate'
 Description-Content-Type: text/markdown
 
 <div align="center">
@@ -156,7 +155,7 @@ Description-Content-Type: text/markdown
   ·
   <a href="https://www.reddit.com/r/AIMemory/">Join r/AIMemory</a>
   .
-  <a href="https://www.docs.cognee.ai">Docs</a>
+  <a href="https://docs.cognee.ai/">Docs</a>
   .
   <a href="https://github.com/topoteretes/cognee-community">cognee community repo</a>
   </p>
@@ -183,7 +182,7 @@ Description-Content-Type: text/markdown
 
 
 
-**🚀 We are launching Cognee SaaS: Sign up [here](https://www.cognee.ai/waitlist) for the hosted beta!**
+**🚀 We launched Cogwit beta (Fully-hosted AI Memory): Sign up [here](https://platform.cognee.ai/)! 🚀** 
 
 Build dynamic memory for Agents and replace RAG using scalable, modular ECL (Extract, Cognify, Load) pipelines.
 
@@ -325,9 +324,9 @@ Try cognee UI out locally [here](https://docs.cognee.ai/how-to-guides/cognee-ui)
 
 ## Demos
 
-1. What is AI memory:
+1. Cogwit Beta demo:
 
-[Learn about cognee](https://github.com/user-attachments/assets/8b2a0050-5ec4-424c-b417-8269971503f0)
+[Cogwit Beta](https://github.com/user-attachments/assets/fa520cd2-2913-4246-a444-902ea5242cb0)
 
 2. Simple GraphRAG demo
 
@@ -348,6 +347,11 @@ We are committed to making open source an enjoyable and respectful experience fo
   <img alt="contributors" src="https://contrib.rocks/image?repo=topoteretes/cognee"/>
 </a>
 
+## Sponsors
+
+Thanks to the following companies for sponsoring the ongoing development of cognee.
+
+- [GitHub's Secure Open Source Fund](https://resources.github.com/github-secure-open-source-fund/)
 
 ## Star History