cognee 0.3.3__tar.gz → 0.3.4.dev1__tar.gz

{cognee-0.3.3 → cognee-0.3.4.dev1}/.env.template

@@ -121,6 +121,9 @@ ACCEPT_LOCAL_FILE_PATH=True
 # This protects against Server Side Request Forgery when proper infrastructure is not in place.
 ALLOW_HTTP_REQUESTS=True
 
+# When set to false don't allow cypher search to be used in Cognee.
+ALLOW_CYPHER_QUERY=True
+
 # When set to False errors during data processing will be returned as info but not raised to allow handling of faulty documents
 RAISE_INCREMENTAL_LOADING_ERRORS=True
 

{cognee-0.3.3 → cognee-0.3.4.dev1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cognee
-Version: 0.3.3
+Version: 0.3.4.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
@@ -57,7 +57,7 @@ Requires-Dist: structlog<26,>=25.2.0
 Requires-Dist: tiktoken<1.0.0,>=0.8.0
 Requires-Dist: typing-extensions<5.0.0,>=4.12.2
 Provides-Extra: anthropic
-Requires-Dist: anthropic<0.27,>=0.26.1; extra == 'anthropic'
+Requires-Dist: anthropic>=0.27; extra == 'anthropic'
 Provides-Extra: api
 Requires-Dist: gunicorn<24,>=20.1.0; extra == 'api'
 Requires-Dist: uvicorn<1.0.0,>=0.34.0; extra == 'api'
@@ -65,8 +65,8 @@ Requires-Dist: websockets<16.0.0,>=15.0.1; extra == 'api'
 Provides-Extra: aws
 Requires-Dist: s3fs[boto3]==2025.3.2; extra == 'aws'
 Provides-Extra: chromadb
-Requires-Dist: chromadb<0.7,>=0.3.0; extra == 'chromadb'
-Requires-Dist: pypika==0.48.8; extra == 'chromadb'
+Requires-Dist: chromadb<0.7,>=0.6; extra == 'chromadb'
+Requires-Dist: pypika==0.48.9; extra == 'chromadb'
 Provides-Extra: codegraph
 Requires-Dist: fastembed<=0.6.0; (python_version < '3.13') and extra == 'codegraph'
 Requires-Dist: transformers<5,>=4.46.3; extra == 'codegraph'
@@ -316,7 +316,19 @@ You can also cognify your files and query using cognee UI.
 
 <img src="assets/cognee-new-ui.webp" width="100%" alt="Cognee UI 2"></a>
 
-Try cognee UI by runnning ``` cognee -ui ``` command on your terminal.
+### Installation for UI
+
+To use the cognee UI with full functionality, you need to install cognee with API dependencies:
+
+```bash
+pip install 'cognee[api]'
+```
+
+The UI requires backend server functionality (uvicorn and other API dependencies) which are not included in the default cognee installation to keep it lightweight.
+
+### Running the UI
+
+Try cognee UI by running ``` cognee-cli -ui ``` command on your terminal.
 
 ## Understand our architecture
 

{cognee-0.3.3 → cognee-0.3.4.dev1}/README.md

@@ -176,7 +176,19 @@ You can also cognify your files and query using cognee UI.
 
 <img src="assets/cognee-new-ui.webp" width="100%" alt="Cognee UI 2"></a>
 
-Try cognee UI by runnning ``` cognee -ui ``` command on your terminal.
+### Installation for UI
+
+To use the cognee UI with full functionality, you need to install cognee with API dependencies:
+
+```bash
+pip install 'cognee[api]'
+```
+
+The UI requires backend server functionality (uvicorn and other API dependencies) which are not included in the default cognee installation to keep it lightweight.
+
+### Running the UI
+
+Try cognee UI by running ``` cognee-cli -ui ``` command on your terminal.
 
 ## Understand our architecture
 

{cognee-0.3.3 → cognee-0.3.4.dev1}/cognee/api/v1/cognify/routers/get_cognify_router.py

@@ -3,6 +3,7 @@ import asyncio
 from uuid import UUID
 from pydantic import Field
 from typing import List, Optional
+from fastapi.encoders import jsonable_encoder
 from fastapi.responses import JSONResponse
 from fastapi import APIRouter, WebSocket, Depends, WebSocketDisconnect
 from starlette.status import WS_1000_NORMAL_CLOSURE, WS_1008_POLICY_VIOLATION
@@ -119,7 +120,7 @@ def get_cognify_router() -> APIRouter:
 
             # If any cognify run errored return JSONResponse with proper error status code
             if any(isinstance(v, PipelineRunErrored) for v in cognify_run.values()):
-                return JSONResponse(status_code=420, content=cognify_run)
+                return JSONResponse(status_code=420, content=jsonable_encoder(cognify_run))
             return cognify_run
         except Exception as error:
             return JSONResponse(status_code=409, content={"error": str(error)})

{cognee-0.3.3 → cognee-0.3.4.dev1}/cognee/api/v1/search/search.py

@@ -22,7 +22,7 @@ async def search(
     node_type: Optional[Type] = NodeSet,
     node_name: Optional[List[str]] = None,
     save_interaction: bool = False,
-    last_k: Optional[int] = None,
+    last_k: Optional[int] = 1,
     only_context: bool = False,
     use_combined_context: bool = False,
 ) -> Union[List[SearchResult], CombinedSearchResult]:

{cognee-0.3.3 → cognee-0.3.4.dev1}/cognee/infrastructure/databases/relational/ModelBase.py

@@ -1,7 +1,8 @@
 from sqlalchemy.orm import DeclarativeBase
+from sqlalchemy.ext.asyncio import AsyncAttrs
 
 
-class Base(DeclarativeBase):
+class Base(AsyncAttrs, DeclarativeBase):
     """
     Represents a base class for declarative models using SQLAlchemy.
 

{cognee-0.3.3 → cognee-0.3.4.dev1}/cognee/infrastructure/databases/vector/chromadb/ChromaDBAdapter.py

@@ -83,7 +83,7 @@ def process_data_for_chroma(data):
         elif isinstance(value, list):
             # Store lists as JSON strings with special prefix
             processed_data[f"{key}__list"] = json.dumps(value)
-        elif isinstance(value, (str, int, float, bool)) or value is None:
+        elif isinstance(value, (str, int, float, bool)):
             processed_data[key] = value
         else:
             processed_data[key] = str(value)
@@ -553,8 +553,4 @@ class ChromaDBAdapter(VectorDBInterface):
             Returns a list of collection names.
         """
         client = await self.get_connection()
-        collections = await client.list_collections()
-        return [
-            collection.name if hasattr(collection, "name") else collection["name"]
-            for collection in collections
-        ]
+        return await client.list_collections()

{cognee-0.3.3 → cognee-0.3.4.dev1}/cognee/infrastructure/databases/vector/embeddings/OllamaEmbeddingEngine.py

@@ -94,10 +94,8 @@ class OllamaEmbeddingEngine(EmbeddingEngine):
         """
         Internal method to call the Ollama embeddings endpoint for a single prompt.
         """
-        payload = {
-            "model": self.model,
-            "prompt": prompt,
-        }
+        payload = {"model": self.model, "prompt": prompt, "input": prompt}
+
         headers = {}
         api_key = os.getenv("LLM_API_KEY")
         if api_key:

{cognee-0.3.3 → cognee-0.3.4.dev1}/cognee/infrastructure/files/utils/get_file_metadata.py

@@ -56,7 +56,12 @@ async def get_file_metadata(file: BinaryIO) -> FileMetadata:
     file_type = guess_file_type(file)
 
     file_path = getattr(file, "name", None) or getattr(file, "full_name", None)
-    file_name = Path(file_path).stem if file_path else None
+
+    if isinstance(file_path, str):
+        file_name = Path(file_path).stem if file_path else None
+    else:
+        # In case file_path does not exist or is a integer return None
+        file_name = None
 
     # Get file size
     pos = file.tell()  # remember current pointer

{cognee-0.3.3 → cognee-0.3.4.dev1}/cognee/infrastructure/llm/structured_output_framework/litellm_instructor/llm/anthropic/adapter.py

@@ -12,6 +12,7 @@ from cognee.infrastructure.llm.structured_output_framework.litellm_instructor.ll
 )
 
 from cognee.infrastructure.llm.LLMGateway import LLMGateway
+from cognee.infrastructure.llm.config import get_llm_config
 
 
 class AnthropicAdapter(LLMInterface):
@@ -27,7 +28,8 @@ class AnthropicAdapter(LLMInterface):
         import anthropic
 
         self.aclient = instructor.patch(
-            create=anthropic.AsyncAnthropic().messages.create, mode=instructor.Mode.ANTHROPIC_TOOLS
+            create=anthropic.AsyncAnthropic(api_key=get_llm_config().llm_api_key).messages.create,
+            mode=instructor.Mode.ANTHROPIC_TOOLS,
         )
 
         self.model = model
@@ -57,7 +59,7 @@ class AnthropicAdapter(LLMInterface):
 
         return await self.aclient(
             model=self.model,
-            max_completion_tokens=4096,
+            max_tokens=4096,
             max_retries=5,
             messages=[
                 {

{cognee-0.3.3 → cognee-0.3.4.dev1}/cognee/modules/pipelines/models/PipelineRunInfo.py

@@ -1,6 +1,7 @@
-from typing import Any, Optional
+from typing import Any, Optional, List, Union
 from uuid import UUID
 from pydantic import BaseModel
+from cognee.modules.data.models.Data import Data
 
 
 class PipelineRunInfo(BaseModel):
@@ -8,11 +9,15 @@ class PipelineRunInfo(BaseModel):
     pipeline_run_id: UUID
     dataset_id: UUID
     dataset_name: str
-    payload: Optional[Any] = None
+    # Data must be mentioned in typing to allow custom encoders for Data to be activated
+    payload: Optional[Union[Any, List[Data]]] = None
     data_ingestion_info: Optional[list] = None
 
     model_config = {
         "arbitrary_types_allowed": True,
+        "from_attributes": True,
+        # Add custom encoding handler for Data ORM model
+        "json_encoders": {Data: lambda d: d.to_json()},
     }
 
 

{cognee-0.3.3 → cognee-0.3.4.dev1}/cognee/modules/search/methods/get_search_type_tools.py

@@ -1,3 +1,4 @@
+import os
 from typing import Callable, List, Optional, Type
 
 from cognee.modules.engine.models.node_set import NodeSet
@@ -160,6 +161,12 @@ async def get_search_type_tools(
     if query_type is SearchType.FEELING_LUCKY:
         query_type = await select_search_type(query_text)
 
+    if (
+        query_type in [SearchType.CYPHER, SearchType.NATURAL_LANGUAGE]
+        and os.getenv("ALLOW_CYPHER_QUERY", "true").lower() == "false"
+    ):
+        raise UnsupportedSearchTypeError("Cypher query search types are disabled.")
+
     search_type_tools = search_tasks.get(query_type)
 
     if not search_type_tools:

{cognee-0.3.3 → cognee-0.3.4.dev1}/cognee/modules/search/methods/search.py

@@ -136,12 +136,19 @@ async def search(
         if os.getenv("ENABLE_BACKEND_ACCESS_CONTROL", "false").lower() == "true":
             return_value = []
             for search_result in search_results:
-                result, context, datasets = search_result
+                prepared_search_results = await prepare_search_result(search_result)
+
+                result = prepared_search_results["result"]
+                graphs = prepared_search_results["graphs"]
+                context = prepared_search_results["context"]
+                datasets = prepared_search_results["datasets"]
+
                 return_value.append(
                     {
-                        "search_result": result,
+                        "search_result": [result] if result else None,
                         "dataset_id": datasets[0].id,
                         "dataset_name": datasets[0].name,
+                        "graphs": graphs,
                     }
                 )
             return return_value
@@ -155,14 +162,6 @@ async def search(
                 return return_value[0]
             else:
                 return return_value
-        # return [
-        #     SearchResult(
-        #         search_result=result,
-        #         dataset_id=datasets[min(index, len(datasets) - 1)].id if datasets else None,
-        #         dataset_name=datasets[min(index, len(datasets) - 1)].name if datasets else None,
-        #     )
-        #     for index, (result, _, datasets) in enumerate(search_results)
-        # ]
 
 
 async def authorized_search(
@@ -208,11 +207,11 @@ async def authorized_search(
         context = {}
         datasets: List[Dataset] = []
 
-        for _, search_context, datasets in search_responses:
-            for dataset in datasets:
+        for _, search_context, search_datasets in search_responses:
+            for dataset in search_datasets:
                 context[str(dataset.id)] = search_context
 
-            datasets.extend(datasets)
+            datasets.extend(search_datasets)
 
         specific_search_tools = await get_search_type_tools(
             query_type=query_type,

cognee-0.3.4.dev1/cognee/modules/search/utils/prepare_search_result.py

@@ -0,0 +1,63 @@
+from typing import List, cast
+from uuid import uuid5, NAMESPACE_OID
+
+from cognee.modules.graph.utils import resolve_edges_to_text
+from cognee.modules.graph.cognee_graph.CogneeGraphElements import Edge
+from cognee.modules.search.types.SearchResult import SearchResultDataset
+from cognee.modules.search.utils.transform_context_to_graph import transform_context_to_graph
+from cognee.modules.search.utils.transform_insights_to_graph import transform_insights_to_graph
+
+
+async def prepare_search_result(search_result):
+    results, context, datasets = search_result
+
+    graphs = None
+    result_graph = None
+    context_texts = {}
+
+    if isinstance(datasets, list) and len(datasets) == 0:
+        datasets = [
+            SearchResultDataset(
+                id=uuid5(NAMESPACE_OID, "*"),
+                name="all available datasets",
+            )
+        ]
+
+    if (
+        isinstance(context, List)
+        and len(context) > 0
+        and isinstance(context[0], tuple)
+        and context[0][1].get("relationship_name")
+    ):
+        context_graph = transform_insights_to_graph(context)
+        graphs = {
+            ", ".join([dataset.name for dataset in datasets]): context_graph,
+        }
+        results = None
+    elif isinstance(context, List) and len(context) > 0 and isinstance(context[0], Edge):
+        context_graph = transform_context_to_graph(context)
+
+        graphs = {
+            ", ".join([dataset.name for dataset in datasets]): context_graph,
+        }
+        context_texts = {
+            ", ".join([dataset.name for dataset in datasets]): await resolve_edges_to_text(context),
+        }
+    elif isinstance(context, str):
+        context_texts = {
+            ", ".join([dataset.name for dataset in datasets]): context,
+        }
+    elif isinstance(context, List) and len(context) > 0 and isinstance(context[0], str):
+        context_texts = {
+            ", ".join([dataset.name for dataset in datasets]): "\n".join(cast(List[str], context)),
+        }
+
+    if isinstance(results, List) and len(results) > 0 and isinstance(results[0], Edge):
+        result_graph = transform_context_to_graph(results)
+
+    return {
+        "result": result_graph or results[0] if results and len(results) == 1 else results,
+        "graphs": graphs,
+        "context": context_texts,
+        "datasets": datasets,
+    }

{cognee-0.3.3 → cognee-0.3.4.dev1}/cognee/modules/search/utils/transform_context_to_graph.py

@@ -14,7 +14,7 @@ def transform_context_to_graph(context: List[Edge]):
             if "name" in triplet.node1.attributes
             else triplet.node1.id,
             "type": triplet.node1.attributes["type"],
-            "attributes": triplet.node2.attributes,
+            "attributes": triplet.node1.attributes,
         }
         nodes[triplet.node2.id] = {
             "id": triplet.node2.id,

cognee-0.3.4.dev1/cognee/modules/search/utils/transform_insights_to_graph.py

@@ -0,0 +1,28 @@
+from typing import Dict, List, Tuple
+
+
+def transform_insights_to_graph(context: List[Tuple[Dict, Dict, Dict]]):
+    nodes = {}
+    edges = {}
+
+    for triplet in context:
+        nodes[triplet[0]["id"]] = {
+            "id": triplet[0]["id"],
+            "label": triplet[0]["name"] if "name" in triplet[0] else triplet[0]["id"],
+            "type": triplet[0]["type"],
+        }
+        nodes[triplet[2]["id"]] = {
+            "id": triplet[2]["id"],
+            "label": triplet[2]["name"] if "name" in triplet[2] else triplet[2]["id"],
+            "type": triplet[2]["type"],
+        }
+        edges[f"{triplet[0]['id']}_{triplet[1]['relationship_name']}_{triplet[2]['id']}"] = {
+            "source": triplet[0]["id"],
+            "target": triplet[2]["id"],
+            "label": triplet[1]["relationship_name"],
+        }
+
+    return {
+        "nodes": list(nodes.values()),
+        "edges": list(edges.values()),
+    }

{cognee-0.3.3 → cognee-0.3.4.dev1}/cognee/modules/users/permissions/methods/get_all_user_permission_datasets.py

@@ -1,3 +1,5 @@
+from types import SimpleNamespace
+
 from cognee.shared.logging_utils import get_logger
 
 from ...models.User import User
@@ -17,9 +19,14 @@ async def get_all_user_permission_datasets(user: User, permission_type: str) ->
         # Get all datasets all tenants have access to
         tenant = await get_tenant(user.tenant_id)
         datasets.extend(await get_principal_datasets(tenant, permission_type))
+
         # Get all datasets Users roles have access to
-        for role_name in user.roles:
-            role = await get_role(user.tenant_id, role_name)
+        if isinstance(user, SimpleNamespace):
+            # If simple namespace use roles defined in user
+            roles = user.roles
+        else:
+            roles = await user.awaitable_attrs.roles
+        for role in roles:
             datasets.extend(await get_principal_datasets(role, permission_type))
 
     # Deduplicate datasets with same ID

{cognee-0.3.3 → cognee-0.3.4.dev1}/cognee/tasks/temporal_graph/models.py

@@ -3,12 +3,17 @@ from pydantic import BaseModel, Field
 
 
 class Timestamp(BaseModel):
-    year: int = Field(..., ge=1, le=9999)
-    month: int = Field(..., ge=1, le=12)
-    day: int = Field(..., ge=1, le=31)
-    hour: int = Field(..., ge=0, le=23)
-    minute: int = Field(..., ge=0, le=59)
-    second: int = Field(..., ge=0, le=59)
+    year: int = Field(
+        ...,
+        ge=1,
+        le=9999,
+        description="Always required. If only a year is known, use it.",
+    )
+    month: int = Field(1, ge=1, le=12, description="If unknown, default to 1")
+    day: int = Field(1, ge=1, le=31, description="If unknown, default to 1")
+    hour: int = Field(0, ge=0, le=23, description="If unknown, default to 0")
+    minute: int = Field(0, ge=0, le=59, description="If unknown, default to 0")
+    second: int = Field(0, ge=0, le=59, description="If unknown, default to 0")
 
 
 class Interval(BaseModel):

{cognee-0.3.3 → cognee-0.3.4.dev1}/cognee/tests/cli_tests/cli_unit_tests/test_cli_main.py

@@ -49,7 +49,7 @@ class TestCliMain:
     def test_main_no_command(self, mock_create_parser):
         """Test main function when no command is provided"""
         mock_parser = MagicMock()
-        mock_parser.parse_args.return_value = MagicMock(command=None)
+        mock_parser.parse_args.return_value = MagicMock(command=None, spec={})
         mock_create_parser.return_value = (mock_parser, {})
 
         result = main()
@@ -64,7 +64,7 @@ class TestCliMain:
         mock_command.execute.return_value = None
 
         mock_parser = MagicMock()
-        mock_args = MagicMock(command="test")
+        mock_args = MagicMock(command="test", spec={})
         mock_parser.parse_args.return_value = mock_args
 
         mock_create_parser.return_value = (mock_parser, {"test": mock_command})
@@ -84,7 +84,7 @@ class TestCliMain:
         mock_command.execute.side_effect = CliCommandException("Test error", error_code=2)
 
         mock_parser = MagicMock()
-        mock_args = MagicMock(command="test")
+        mock_args = MagicMock(command="test", spec={})
         mock_parser.parse_args.return_value = mock_args
 
         mock_create_parser.return_value = (mock_parser, {"test": mock_command})
@@ -103,7 +103,7 @@ class TestCliMain:
         mock_command.execute.side_effect = Exception("Generic error")
 
         mock_parser = MagicMock()
-        mock_args = MagicMock(command="test")
+        mock_args = MagicMock(command="test", spec={})
         mock_parser.parse_args.return_value = mock_args
 
         mock_create_parser.return_value = (mock_parser, {"test": mock_command})
@@ -126,7 +126,7 @@ class TestCliMain:
         mock_command.execute.side_effect = test_exception
 
         mock_parser = MagicMock()
-        mock_args = MagicMock(command="test")
+        mock_args = MagicMock(command="test", spec={})
         mock_parser.parse_args.return_value = mock_args
 
         mock_create_parser.return_value = (mock_parser, {"test": mock_command})

{cognee-0.3.3 → cognee-0.3.4.dev1}/cognee/tests/test_cognee_server_start.py

@@ -48,7 +48,7 @@ class TestCogneeServerStart(unittest.TestCase):
         """Test that the server is running and can accept connections."""
         # Test health endpoint
         health_response = requests.get("http://localhost:8000/health", timeout=15)
-        self.assertIn(health_response.status_code, [200, 503])
+        self.assertIn(health_response.status_code, [200])
 
         # Test root endpoint
         root_response = requests.get("http://localhost:8000/", timeout=15)
@@ -88,7 +88,7 @@ class TestCogneeServerStart(unittest.TestCase):
         payload = {"datasets": [dataset_name]}
 
         add_response = requests.post(url, headers=headers, data=form_data, files=file, timeout=50)
-        if add_response.status_code not in [200, 201, 409]:
+        if add_response.status_code not in [200, 201]:
             add_response.raise_for_status()
 
         # Cognify request
@@ -99,7 +99,7 @@ class TestCogneeServerStart(unittest.TestCase):
         }
 
         cognify_response = requests.post(url, headers=headers, json=payload, timeout=150)
-        if cognify_response.status_code not in [200, 201, 409]:
+        if cognify_response.status_code not in [200, 201]:
             cognify_response.raise_for_status()
 
         # TODO: Add test to verify cognify pipeline is complete before testing search
@@ -115,7 +115,7 @@ class TestCogneeServerStart(unittest.TestCase):
         payload = {"searchType": "GRAPH_COMPLETION", "query": "What's in the document?"}
 
         search_response = requests.post(url, headers=headers, json=payload, timeout=50)
-        if search_response.status_code not in [200, 201, 409]:
+        if search_response.status_code not in [200, 201]:
             search_response.raise_for_status()
 
 

{cognee-0.3.3 → cognee-0.3.4.dev1}/cognee/tests/test_temporal_graph.py

@@ -97,7 +97,7 @@ async def main():
         f"Expected exactly one DocumentChunk, but found {type_counts.get('DocumentChunk', 0)}"
     )
 
-    assert type_counts.get("Entity", 0) >= 20, (
+    assert type_counts.get("Entity", 0) >= 10, (
         f"Expected multiple entities (assert is set to 20), but found {type_counts.get('Entity', 0)}"
     )
 
@@ -105,52 +105,24 @@ async def main():
         f"Expected multiple entity types, but found {type_counts.get('EntityType', 0)}"
     )
 
-    assert type_counts.get("Event", 0) >= 20, (
+    assert type_counts.get("Event", 0) >= 10, (
         f"Expected multiple events (assert is set to 20), but found {type_counts.get('Event', 0)}"
     )
 
-    assert type_counts.get("Timestamp", 0) >= 20, (
-        f"Expected multiple timestamps (assert is set to 20), but found {type_counts.get('Timestamp', 0)}"
+    assert type_counts.get("Timestamp", 0) >= 10, (
+        f"Expected multiple timestamps (assert is set to 10), but found {type_counts.get('Timestamp', 0)}"
     )
 
-    assert type_counts.get("Interval", 0) >= 2, (
-        f"Expected multiple intervals, but found {type_counts.get('Interval', 0)}"
-    )
-
-    assert edge_type_counts.get("contains", 0) >= 20, (
+    assert edge_type_counts.get("contains", 0) >= 10, (
         f"Expected multiple 'contains' edge, but found {edge_type_counts.get('contains', 0)}"
     )
 
-    assert edge_type_counts.get("is_a", 0) >= 20, (
+    assert edge_type_counts.get("is_a", 0) >= 10, (
         f"Expected multiple 'is_a' edge, but found {edge_type_counts.get('is_a', 0)}"
     )
 
-    assert edge_type_counts.get("during", 0) == type_counts.get("Interval", 0), (
-        "Expected the same amount of during and interval objects in the graph"
-    )
-
-    assert edge_type_counts.get("during", 0) == type_counts.get("Interval", 0), (
-        "Expected the same amount of during and interval objects in the graph"
-    )
-
-    assert edge_type_counts.get("time_from", 0) == type_counts.get("Interval", 0), (
-        "Expected the same amount of time_from and interval objects in the graph"
-    )
-
-    assert edge_type_counts.get("time_to", 0) == type_counts.get("Interval", 0), (
-        "Expected the same amount of time_to and interval objects in the graph"
-    )
-
     retriever = TemporalRetriever()
 
-    result_before = await retriever.extract_time_from_query("What happened before 1890?")
-
-    assert result_before[0] is None
-
-    result_after = await retriever.extract_time_from_query("What happened after 1891?")
-
-    assert result_after[1] is None
-
     result_between = await retriever.extract_time_from_query("What happened between 1890 and 1900?")
 
     assert result_between[1]