cognee 0.3.2__py3-none-any.whl → 0.3.4.dev0__py3-none-any.whl

cognee/api/v1/notebooks/routers/get_notebooks_router.py

@@ -31,7 +31,8 @@ def get_notebooks_router():
 
     @router.get("")
     async def get_notebooks_endpoint(user: User = Depends(get_authenticated_user)):
-        return await get_notebooks(user.id)
+        async with get_async_session() as session:
+            return await get_notebooks(user.id, session)
 
     @router.post("")
     async def create_notebook_endpoint(

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/infrastructure/databases/relational/sqlalchemy/SqlAlchemyAdapter.py

@@ -58,8 +58,8 @@ class SQLAlchemyAdapter:
         else:
             self.engine = create_async_engine(
                 connection_string,
-                pool_size=5,
-                max_overflow=10,
+                pool_size=20,
+                max_overflow=20,
                 pool_recycle=280,
                 pool_pre_ping=True,
                 pool_timeout=280,

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/infrastructure/databases/vector/config.py

@@ -39,7 +39,7 @@ class VectorConfig(BaseSettings):
             values.vector_db_url = ensure_absolute_path(
                 values.vector_db_url,
             )
-        else:
+        elif not values.vector_db_url:
             # Default path
             databases_directory_path = os.path.join(base_config.system_root_directory, "databases")
             values.vector_db_url = os.path.join(databases_directory_path, "cognee.lancedb")

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/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/infrastructure/utils/run_async.py

@@ -1,13 +1,18 @@
 import asyncio
 from functools import partial
+import inspect
 
 
 async def run_async(func, *args, loop=None, executor=None, **kwargs):
     if loop is None:
         try:
-            running_loop = asyncio.get_running_loop()
+            loop = asyncio.get_running_loop()
         except RuntimeError:
-            running_loop = asyncio.get_event_loop()
+            loop = asyncio.get_event_loop()
 
-    pfunc = partial(func, *args, **kwargs)
-    return await running_loop.run_in_executor(executor, pfunc)
+    if "loop" in inspect.signature(func).parameters:
+        pfunc = partial(func, *args, loop=loop, **kwargs)
+    else:
+        pfunc = partial(func, *args, **kwargs)
+
+    return await loop.run_in_executor(executor, pfunc)

cognee/infrastructure/utils/run_sync.py

@@ -2,16 +2,17 @@ import asyncio
 import threading
 
 
-def run_sync(coro, timeout=None):
+def run_sync(coro, running_loop=None, timeout=None):
     result = None
     exception = None
 
     def runner():
-        nonlocal result, exception
+        nonlocal result, exception, running_loop
 
         try:
             try:
-                running_loop = asyncio.get_running_loop()
+                if not running_loop:
+                    running_loop = asyncio.get_running_loop()
 
                 result = asyncio.run_coroutine_threadsafe(coro, running_loop).result(timeout)
             except RuntimeError:

cognee/modules/notebooks/methods/create_tutorial_notebook.py

@@ -0,0 +1,87 @@
+from uuid import UUID, uuid4
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from ..models import NotebookCell
+from .create_notebook import create_notebook
+
+
+async def create_tutorial_notebook(user_id: UUID, session: AsyncSession):
+    await create_notebook(
+        user_id=user_id,
+        notebook_name="Welcome to cognee 🧠",
+        cells=[
+            NotebookCell(
+                id=uuid4(),
+                name="Welcome",
+                content="Cognee is your toolkit for turning text into a structured knowledge graph, optionally enhanced by ontologies, and then querying it with advanced retrieval techniques. This notebook will guide you through a simple example.",
+                type="markdown",
+            ),
+            NotebookCell(
+                id=uuid4(),
+                name="Example",
+                content="",
+                type="markdown",
+            ),
+        ],
+        deletable=False,
+        session=session,
+    )
+
+
+cell_content = [
+    """
+# Using Cognee with Python Development Data
+
+Unite authoritative Python practice (Guido van Rossum's own contributions!), normative guidance (Zen/PEP 8), and your lived context (rules + conversations) into one *AI memory* that produces answers that are relevant, explainable, and consistent.
+""",
+    """
+## What You'll Learn
+
+In this comprehensive tutorial, you'll discover how to transform scattered development data into an intelligent knowledge system that enhances your coding workflow. By the end, you'll have:
+- Connected disparate data sources (Guido's CPython contributions, mypy development, PEP discussions, your Python projects) into a unified AI memory graph
+- Built an memory layer that understands Python design philosophy, best practice coding patterns, and your preferences and experience
+- Learn how to use intelligent search capabilities that combine the diverse context
+- Integrated everything with your coding environment through MCP (Model Context Protocol)
+
+This tutorial demonstrates the power of **knowledge graphs** and **retrieval-augmented generation (RAG)** for software development, showing you how to build systems that learn from Python's creator and improve your own Python development.
+""",
+    """
+## Cognee and its core operations
+
+Before we dive in, let's understand the core Cognee operations we'll be working with:
+- `cognee.add()` - Ingests raw data (files, text, APIs) into the system
+- `cognee.cognify()` - Processes and structures data into a knowledge graph using AI
+- `cognee.search()` - Queries the knowledge graph with natural language or Cypher
+- `cognee.memify()` - Cognee's \"secret sauce\" that infers implicit connections and rules from your data
+""",
+    """
+## Data used in this tutorial
+
+Cognee can ingest many types of sources. In this tutorial, we use a small, concrete set of files that cover different perspectives:
+- `guido_contributions.json` — Authoritative exemplars. Real PRs and commits from Guido van Rossum (mypy, CPython). These show how Python's creator solved problems and provide concrete anchors for patterns.
+- `pep_style_guide.md` — Norms. Encodes community style and typing conventions (PEP 8 and related). Ensures that search results and inferred rules align with widely accepted standards.
+- `zen_principles.md` — Philosophy. The Zen of Python. Grounds design trade-offs (simplicity, explicitness, readability) beyond syntax or mechanics.
+- `my_developer_rules.md` — Local constraints. Your house rules, conventions, and project-specific requirements (scope, privacy, Spec.md). Keeps recommendations relevant to your actual workflow.
+- `copilot_conversations.json` — Personal history. Transcripts of real assistant conversations, including your questions, code snippets, and discussion topics. Captures "how you code" and connects it to "how Guido codes."
+""",
+    """
+# Preliminaries
+
+To strike the balanace between speed, cost, anc quality, we recommend using OpenAI's `4o-mini` model; make sure your `.env` file contains this line:
+`
+LLM_MODEL="gpt-4o-mini"
+`
+""",
+    """
+import cognee
+
+result = await cognee.add(
+    "file://data/guido_contributions.json",
+    node_set=["guido_data"]
+)
+
+await cognee.cognify(temporal_cognify=True)
+
+results = await cognee.search("Show me commits")
+""",
+]

cognee/modules/notebooks/methods/get_notebook.py

@@ -1,6 +1,6 @@
 from uuid import UUID
 from typing import Optional
-from sqlalchemy import select
+from sqlalchemy import and_, select
 from sqlalchemy.ext.asyncio import AsyncSession
 
 from cognee.infrastructure.databases.relational import with_async_session
@@ -15,7 +15,7 @@ async def get_notebook(
     session: AsyncSession,
 ) -> Optional[Notebook]:
     result = await session.execute(
-        select(Notebook).where(Notebook.owner_id == user_id and Notebook.id == notebook_id)
+        select(Notebook).where(and_(Notebook.owner_id == user_id, Notebook.id == notebook_id))
     )
 
     return result.scalar()

cognee/modules/notebooks/methods/update_notebook.py

@@ -1,4 +1,3 @@
-from typing import Callable, AsyncContextManager
 from sqlalchemy.ext.asyncio import AsyncSession
 
 from cognee.infrastructure.databases.relational import with_async_session

cognee/modules/notebooks/operations/run_in_local_sandbox.py

@@ -5,16 +5,18 @@ import traceback
 
 def wrap_in_async_handler(user_code: str) -> str:
     return (
-        "from cognee.infrastructure.utils.run_sync import run_sync\n\n"
-        "async def __user_main__():\n"
+        "import asyncio\n"
+        + "asyncio.set_event_loop(running_loop)\n\n"
+        + "from cognee.infrastructure.utils.run_sync import run_sync\n\n"
+        + "async def __user_main__():\n"
         + "\n".join("    " + line for line in user_code.strip().split("\n"))
         + "\n"
-        "    globals().update(locals())\n\n"
-        "run_sync(__user_main__())\n"
+        + "    globals().update(locals())\n\n"
+        + "run_sync(__user_main__(), running_loop)\n"
     )
 
 
-def run_in_local_sandbox(code, environment=None):
+def run_in_local_sandbox(code, environment=None, loop=None):
     environment = environment or {}
     code = wrap_in_async_handler(code.replace("\xa0", "\n"))
 
@@ -31,6 +33,7 @@ def run_in_local_sandbox(code, environment=None):
         printOutput.append(output)
 
     environment["print"] = customPrintFunction
+    environment["running_loop"] = loop
 
     try:
         exec(code, environment)

cognee/modules/retrieval/graph_completion_context_extension_retriever.py

@@ -48,7 +48,7 @@ class GraphCompletionContextExtensionRetriever(GraphCompletionRetriever):
         query: str,
         context: Optional[List[Edge]] = None,
         context_extension_rounds=4,
-    ) -> str:
+    ) -> List[str]:
         """
         Extends the context for a given query by retrieving related triplets and generating new
         completions based on them.

cognee/modules/retrieval/graph_completion_cot_retriever.py

@@ -58,7 +58,7 @@ class GraphCompletionCotRetriever(GraphCompletionRetriever):
         query: str,
         context: Optional[List[Edge]] = None,
         max_iter=4,
-    ) -> str:
+    ) -> List[str]:
         """
         Generate completion responses based on a user query and contextual information.
 

cognee/modules/retrieval/graph_completion_retriever.py

@@ -135,7 +135,7 @@ class GraphCompletionRetriever(BaseGraphRetriever):
         self,
         query: str,
         context: Optional[List[Edge]] = None,
-    ) -> Any:
+    ) -> List[str]:
         """
         Generates a completion using graph connections context based on a query.
 

cognee/modules/retrieval/temporal_retriever.py

@@ -136,7 +136,7 @@ class TemporalRetriever(GraphCompletionRetriever):
 
         return self.descriptions_to_string(top_k_events)
 
-    async def get_completion(self, query: str, context: Optional[str] = None) -> str:
+    async def get_completion(self, query: str, context: Optional[str] = None) -> List[str]:
         """Generates a response using the query and optional context."""
         if not context:
             context = await self.get_context(query=query)

cognee/modules/retrieval/user_qa_feedback.py

@@ -1,4 +1,4 @@
-from typing import Any, Optional, List
+from typing import Optional, List
 
 from uuid import NAMESPACE_OID, uuid5, UUID
 from cognee.infrastructure.databases.graph import get_graph_engine

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/modules/search/utils/prepare_search_result.py

@@ -1,40 +1,62 @@
 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):
-    result, context, datasets = search_result
+    results, context, datasets = search_result
 
     graphs = None
     result_graph = None
     context_texts = {}
 
-    if isinstance(context, List) and len(context) > 0 and isinstance(context[0], Edge):
+    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 = {
-            "*": context_graph,
+            ", ".join([dataset.name for dataset in datasets]): context_graph,
         }
         context_texts = {
-            "*": await resolve_edges_to_text(context),
+            ", ".join([dataset.name for dataset in datasets]): await resolve_edges_to_text(context),
         }
     elif isinstance(context, str):
         context_texts = {
-            "*": context,
+            ", ".join([dataset.name for dataset in datasets]): context,
         }
     elif isinstance(context, List) and len(context) > 0 and isinstance(context[0], str):
         context_texts = {
-            "*": "\n".join(cast(List[str], context)),
+            ", ".join([dataset.name for dataset in datasets]): "\n".join(cast(List[str], context)),
         }
 
-    if isinstance(result, List) and len(result) > 0 and isinstance(result[0], Edge):
-        result_graph = transform_context_to_graph(result)
+    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 result,
+        "result": result_graph or results[0] if results and len(results) == 1 else results,
         "graphs": graphs,
         "context": context_texts,
         "datasets": datasets,

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/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/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/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/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]