cognee 0.5.1.dev0__py3-none-any.whl → 0.5.2__py3-none-any.whl

cognee/tasks/translation/translate_content.py

@@ -0,0 +1,282 @@
+import asyncio
+from typing import List, Optional
+from uuid import uuid5
+
+from cognee.modules.chunking.models import DocumentChunk
+from cognee.shared.logging_utils import get_logger
+
+from .config import get_translation_config, TranslationProviderType
+from .detect_language import detect_language_async, LanguageDetectionResult
+from .exceptions import TranslationError, LanguageDetectionError
+from .models import TranslatedContent, LanguageMetadata
+from .providers import get_translation_provider, TranslationResult
+
+logger = get_logger(__name__)
+
+
+async def translate_content(
+    data_chunks: List[DocumentChunk],
+    target_language: str = None,
+    translation_provider: TranslationProviderType = None,
+    confidence_threshold: float = None,
+    skip_if_target_language: bool = True,
+    preserve_original: bool = True,
+) -> List[DocumentChunk]:
+    """
+    Translate non-English content to the target language.
+
+    This task detects the language of each document chunk and translates
+    non-target-language content using the specified translation provider.
+    Original text is preserved alongside translated versions.
+
+    Args:
+        data_chunks: List of DocumentChunk objects to process
+        target_language: Target language code (default: "en" for English)
+                        If not provided, uses config default
+        translation_provider: Translation service to use ("llm", "google", "azure")
+                            If not provided, uses config default
+        confidence_threshold: Minimum confidence for language detection (0.0 to 1.0)
+                            If not provided, uses config default
+        skip_if_target_language: If True, skip chunks already in target language
+        preserve_original: If True, store original text in TranslatedContent
+
+    Returns:
+        List of DocumentChunk objects with translated content.
+        Chunks that required translation will have TranslatedContent
+        objects in their 'contains' list.
+
+    Note:
+        This function mutates the input chunks in-place. Specifically:
+        - chunk.text is replaced with the translated text
+        - chunk.contains is updated with LanguageMetadata and TranslatedContent
+        The original text is preserved in TranslatedContent.original_text
+        if preserve_original=True.
+
+    Example:
+        ```python
+        from cognee.tasks.translation import translate_content
+
+        # Translate chunks using default settings
+        translated_chunks = await translate_content(chunks)
+
+        # Translate with specific provider
+        translated_chunks = await translate_content(
+            chunks,
+            translation_provider="llm",
+            confidence_threshold=0.9
+        )
+        ```
+    """
+    if not isinstance(data_chunks, list):
+        raise TranslationError("data_chunks must be a list")
+
+    if len(data_chunks) == 0:
+        return data_chunks
+
+    # Get configuration
+    config = get_translation_config()
+    provider_name = translation_provider or config.translation_provider
+    target_lang = target_language or config.target_language
+    threshold = confidence_threshold or config.confidence_threshold
+
+    logger.info(
+        f"Starting translation task for {len(data_chunks)} chunks "
+        f"using {provider_name} provider, target language: {target_lang}"
+    )
+
+    # Get the translation provider
+    provider = get_translation_provider(provider_name)
+
+    # Process chunks
+    processed_chunks = []
+    total_chunks = len(data_chunks)
+
+    for chunk_index, chunk in enumerate(data_chunks):
+        # Log progress for large batches
+        if chunk_index > 0 and chunk_index % 100 == 0:
+            logger.info(f"Translation progress: {chunk_index}/{total_chunks} chunks processed")
+
+        if not hasattr(chunk, "text") or not chunk.text:
+            processed_chunks.append(chunk)
+            continue
+
+        try:
+            # Detect language
+            detection = await detect_language_async(chunk.text, target_lang, threshold)
+
+            # Create language metadata
+            language_metadata = LanguageMetadata(
+                id=uuid5(chunk.id, "LanguageMetadata"),
+                content_id=chunk.id,
+                detected_language=detection.language_code,
+                language_confidence=detection.confidence,
+                requires_translation=detection.requires_translation,
+                character_count=detection.character_count,
+                language_name=detection.language_name,
+            )
+
+            # Skip if already in target language
+            if not detection.requires_translation:
+                if skip_if_target_language:
+                    logger.debug(
+                        f"Skipping chunk {chunk.id}: already in target language "
+                        f"({detection.language_code})"
+                    )
+                    # Add language metadata to chunk
+                    _add_to_chunk_contains(chunk, language_metadata)
+                    processed_chunks.append(chunk)
+                    continue
+
+            # Translate the content
+            logger.debug(
+                f"Translating chunk {chunk.id} from {detection.language_code} to {target_lang}"
+            )
+
+            translation_result = await provider.translate(
+                text=chunk.text,
+                target_language=target_lang,
+                source_language=detection.language_code,
+            )
+
+            # Create TranslatedContent data point
+            translated_content = TranslatedContent(
+                id=uuid5(chunk.id, "TranslatedContent"),
+                original_chunk_id=chunk.id,
+                original_text=chunk.text if preserve_original else "",
+                translated_text=translation_result.translated_text,
+                source_language=translation_result.source_language,
+                target_language=translation_result.target_language,
+                translation_provider=translation_result.provider,
+                confidence_score=translation_result.confidence_score,
+                translated_from=chunk,
+            )
+
+            # Update chunk text with translated content
+            chunk.text = translation_result.translated_text
+
+            # Add metadata to chunk's contains list
+            _add_to_chunk_contains(chunk, language_metadata)
+            _add_to_chunk_contains(chunk, translated_content)
+
+            processed_chunks.append(chunk)
+
+            logger.debug(
+                f"Successfully translated chunk {chunk.id}: "
+                f"{detection.language_code} -> {target_lang}"
+            )
+
+        except LanguageDetectionError as e:
+            logger.warning(f"Language detection failed for chunk {chunk.id}: {e}")
+            processed_chunks.append(chunk)
+        except TranslationError as e:
+            logger.error(f"Translation failed for chunk {chunk.id}: {e}")
+            processed_chunks.append(chunk)
+        except Exception as e:
+            logger.error(f"Unexpected error processing chunk {chunk.id}: {e}")
+            processed_chunks.append(chunk)
+
+    logger.info(f"Translation task completed for {len(processed_chunks)} chunks")
+    return processed_chunks
+
+
+def _add_to_chunk_contains(chunk: DocumentChunk, item) -> None:
+    """Helper to add an item to a chunk's contains list."""
+    if chunk.contains is None:
+        chunk.contains = []
+    chunk.contains.append(item)
+
+
+async def translate_text(
+    text: str,
+    target_language: str = None,
+    translation_provider: TranslationProviderType = None,
+    source_language: Optional[str] = None,
+) -> TranslationResult:
+    """
+    Translate a single text string.
+
+    This is a convenience function for translating individual texts
+    without creating DocumentChunk objects.
+
+    Args:
+        text: The text to translate
+        target_language: Target language code (default: uses config, typically "en")
+                        If not provided, uses config default
+        translation_provider: Translation service to use
+                            If not provided, uses config default
+        source_language: Source language code (optional, auto-detected if not provided)
+
+    Returns:
+        TranslationResult with translated text and metadata
+
+    Example:
+        ```python
+        from cognee.tasks.translation import translate_text
+
+        result = await translate_text(
+            "Bonjour le monde!",
+            target_language="en"
+        )
+        print(result.translated_text)  # "Hello world!"
+        print(result.source_language)  # "fr"
+        ```
+    """
+    config = get_translation_config()
+    provider_name = translation_provider or config.translation_provider
+    target_lang = target_language or config.target_language
+
+    provider = get_translation_provider(provider_name)
+
+    return await provider.translate(
+        text=text,
+        target_language=target_lang,
+        source_language=source_language,
+    )
+
+
+async def batch_translate_texts(
+    texts: List[str],
+    target_language: str = None,
+    translation_provider: TranslationProviderType = None,
+    source_language: Optional[str] = None,
+) -> List[TranslationResult]:
+    """
+    Translate multiple text strings in batch.
+
+    This is more efficient than translating texts individually,
+    especially for providers that support native batch operations.
+
+    Args:
+        texts: List of texts to translate
+        target_language: Target language code (default: uses config, typically "en")
+                        If not provided, uses config default
+        translation_provider: Translation service to use
+                            If not provided, uses config default
+        source_language: Source language code (optional)
+
+    Returns:
+        List of TranslationResult objects
+
+    Example:
+        ```python
+        from cognee.tasks.translation import batch_translate_texts
+
+        results = await batch_translate_texts(
+            ["Hola", "¿Cómo estás?", "Adiós"],
+            target_language="en"
+        )
+        for result in results:
+            print(f"{result.source_language}: {result.translated_text}")
+        ```
+    """
+    config = get_translation_config()
+    provider_name = translation_provider or config.translation_provider
+    target_lang = target_language or config.target_language
+
+    provider = get_translation_provider(provider_name)
+
+    return await provider.translate_batch(
+        texts=texts,
+        target_language=target_lang,
+        source_language=source_language,
+    )

cognee/tasks/web_scraper/default_url_crawler.py

@@ -73,7 +73,11 @@ class DefaultUrlCrawler:
         self.timeout = timeout
         self.max_retries = max_retries
         self.retry_delay_factor = retry_delay_factor
-        self.headers = headers or {"User-Agent": "Cognee-Scraper/1.0"}
+        self.headers = headers or {
+            "User-Agent": "Cognee-Scraper/1.0 (hello@cognee.ai)",
+            "Accept": "text/html,application/xhtml+xml,application/xml;q=0.9,*/*;q=0.8",
+            "Accept-Language": "en-US,en;q=0.9",
+        }
         self.robots_cache_ttl = robots_cache_ttl
         self._last_request_time_per_domain: Dict[str, float] = {}
         self._robots_cache: Dict[str, RobotsTxtCache] = {}
@@ -288,7 +292,7 @@ class DefaultUrlCrawler:
         while True:
             try:
                 await self._respect_rate_limit(url, crawl_delay)
-                resp = await self._client.get(url)
+                resp = await self._client.get(url, headers=self.headers)
                 resp.raise_for_status()
                 logger.info(
                     f"Successfully fetched {url} (status={resp.status_code}, size={len(resp.text)} bytes)"

cognee/tests/cli_tests/cli_unit_tests/test_cli_commands.py

@@ -238,6 +238,7 @@ class TestCognifyCommand:
             ontology_file_path=None,
             chunker=TextChunker,
             run_in_background=False,
+            chunks_per_batch=None,
         )
 
     @patch("cognee.cli.commands.cognify_command.asyncio.run")

cognee/tests/cli_tests/cli_unit_tests/test_cli_edge_cases.py

@@ -262,6 +262,7 @@ class TestCognifyCommandEdgeCases:
             ontology_file_path=None,
             chunker=TextChunker,
             run_in_background=False,
+            chunks_per_batch=None,
         )
 
     @patch("cognee.cli.commands.cognify_command.asyncio.run", side_effect=_mock_run)
@@ -295,6 +296,7 @@ class TestCognifyCommandEdgeCases:
             ontology_file_path="/nonexistent/path/ontology.owl",
             chunker=TextChunker,
             run_in_background=False,
+            chunks_per_batch=None,
         )
 
     @patch("cognee.cli.commands.cognify_command.asyncio.run")
@@ -373,6 +375,7 @@ class TestCognifyCommandEdgeCases:
             ontology_file_path=None,
             chunker=TextChunker,
             run_in_background=False,
+            chunks_per_batch=None,
         )
 
 

cognee/tests/integration/retrieval/test_brute_force_triplet_search_with_cognify.py

@@ -0,0 +1,62 @@
+import pathlib
+
+import pytest
+import pytest_asyncio
+import cognee
+
+from cognee.modules.graph.cognee_graph.CogneeGraphElements import Edge
+from cognee.modules.retrieval.utils.brute_force_triplet_search import brute_force_triplet_search
+
+
+@pytest_asyncio.fixture
+async def clean_environment():
+    """Configure isolated storage and ensure cleanup before/after."""
+    base_dir = pathlib.Path(__file__).parent.parent.parent.parent
+    system_directory_path = str(base_dir / ".cognee_system/test_brute_force_triplet_search_e2e")
+    data_directory_path = str(base_dir / ".data_storage/test_brute_force_triplet_search_e2e")
+
+    cognee.config.system_root_directory(system_directory_path)
+    cognee.config.data_root_directory(data_directory_path)
+
+    await cognee.prune.prune_data()
+    await cognee.prune.prune_system(metadata=True)
+
+    yield
+
+    try:
+        await cognee.prune.prune_data()
+        await cognee.prune.prune_system(metadata=True)
+    except Exception:
+        pass
+
+
+@pytest.mark.asyncio
+async def test_brute_force_triplet_search_end_to_end(clean_environment):
+    """Minimal end-to-end exercise of single and batch triplet search."""
+
+    text = """
+        Cognee is an open-source AI memory engine that structures data into searchable formats for use with AI agents.
+        The company focuses on persistent memory systems using knowledge graphs and vector search.
+        It is a Berlin-based startup building infrastructure for context-aware AI applications.
+        NLP systems can use Cognee to store and retrieve structured information.
+    """
+
+    await cognee.add(text)
+    await cognee.cognify()
+
+    single_result = await brute_force_triplet_search(
+        query="What can NLP systems use Cognee for?",
+        top_k=1,
+    )
+    assert isinstance(single_result, list)
+    assert single_result
+    assert all(isinstance(edge, Edge) for edge in single_result)
+
+    batch_queries = ["What is Cognee?", "What is the company's focus?"]
+    batch_result = await brute_force_triplet_search(query_batch=batch_queries, top_k=1)
+
+    assert isinstance(batch_result, list)
+    assert len(batch_result) == len(batch_queries)
+    assert all(isinstance(per_query, list) for per_query in batch_result)
+    assert all(per_query for per_query in batch_result)
+    assert all(isinstance(edge, Edge) for per_query in batch_result for edge in per_query)

cognee/tests/integration/retrieval/test_chunks_retriever.py

@@ -10,7 +10,6 @@ from cognee.tasks.storage import add_data_points
 from cognee.infrastructure.databases.vector import get_vector_engine
 from cognee.modules.chunking.models import DocumentChunk
 from cognee.modules.data.processing.document_types import TextDocument
-from cognee.modules.retrieval.exceptions.exceptions import NoDataError
 from cognee.modules.retrieval.chunks_retriever import ChunksRetriever
 from cognee.infrastructure.engine import DataPoint
 from cognee.modules.data.processing.document_types import Document
@@ -40,6 +39,17 @@ async def setup_test_environment_with_chunks_simple():
 
     await cognee.prune.prune_data()
     await cognee.prune.prune_system(metadata=True)
+    from cognee.infrastructure.databases.graph.get_graph_engine import _create_graph_engine
+    from cognee.infrastructure.databases.vector.create_vector_engine import (
+        _create_vector_engine,
+    )
+    from cognee.infrastructure.databases.relational.create_relational_engine import (
+        create_relational_engine,
+    )
+
+    _create_graph_engine.cache_clear()
+    _create_vector_engine.cache_clear()
+    create_relational_engine.cache_clear()
     await setup()
 
     document = TextDocument(
@@ -83,6 +93,17 @@ async def setup_test_environment_with_chunks_simple():
     try:
         await cognee.prune.prune_data()
         await cognee.prune.prune_system(metadata=True)
+        from cognee.infrastructure.databases.graph.get_graph_engine import _create_graph_engine
+        from cognee.infrastructure.databases.vector.create_vector_engine import (
+            _create_vector_engine,
+        )
+        from cognee.infrastructure.databases.relational.create_relational_engine import (
+            create_relational_engine,
+        )
+
+        _create_graph_engine.cache_clear()
+        _create_vector_engine.cache_clear()
+        create_relational_engine.cache_clear()
     except Exception:
         pass
 
@@ -99,6 +120,17 @@ async def setup_test_environment_with_chunks_complex():
 
     await cognee.prune.prune_data()
     await cognee.prune.prune_system(metadata=True)
+    from cognee.infrastructure.databases.graph.get_graph_engine import _create_graph_engine
+    from cognee.infrastructure.databases.vector.create_vector_engine import (
+        _create_vector_engine,
+    )
+    from cognee.infrastructure.databases.relational.create_relational_engine import (
+        create_relational_engine,
+    )
+
+    _create_graph_engine.cache_clear()
+    _create_vector_engine.cache_clear()
+    create_relational_engine.cache_clear()
     await setup()
 
     document1 = TextDocument(
@@ -174,6 +206,17 @@ async def setup_test_environment_with_chunks_complex():
     try:
         await cognee.prune.prune_data()
         await cognee.prune.prune_system(metadata=True)
+        from cognee.infrastructure.databases.graph.get_graph_engine import _create_graph_engine
+        from cognee.infrastructure.databases.vector.create_vector_engine import (
+            _create_vector_engine,
+        )
+        from cognee.infrastructure.databases.relational.create_relational_engine import (
+            create_relational_engine,
+        )
+
+        _create_graph_engine.cache_clear()
+        _create_vector_engine.cache_clear()
+        create_relational_engine.cache_clear()
     except Exception:
         pass
 
@@ -190,26 +233,53 @@ async def setup_test_environment_empty():
 
     await cognee.prune.prune_data()
     await cognee.prune.prune_system(metadata=True)
+    from cognee.infrastructure.databases.graph.get_graph_engine import _create_graph_engine
+    from cognee.infrastructure.databases.vector.create_vector_engine import (
+        _create_vector_engine,
+    )
+    from cognee.infrastructure.databases.relational.create_relational_engine import (
+        create_relational_engine,
+    )
+
+    _create_graph_engine.cache_clear()
+    _create_vector_engine.cache_clear()
+    create_relational_engine.cache_clear()
 
     yield
 
     try:
         await cognee.prune.prune_data()
         await cognee.prune.prune_system(metadata=True)
+        from cognee.infrastructure.databases.graph.get_graph_engine import _create_graph_engine
+        from cognee.infrastructure.databases.vector.create_vector_engine import (
+            _create_vector_engine,
+        )
+        from cognee.infrastructure.databases.relational.create_relational_engine import (
+            create_relational_engine,
+        )
+
+        _create_graph_engine.cache_clear()
+        _create_vector_engine.cache_clear()
+        create_relational_engine.cache_clear()
     except Exception:
         pass
 
 
 @pytest.mark.asyncio
-async def test_chunks_retriever_context_multiple_chunks(setup_test_environment_with_chunks_simple):
+async def test_chunks_retriever_multiple_chunks(setup_test_environment_with_chunks_simple):
     """Integration test: verify ChunksRetriever can retrieve multiple chunks."""
     retriever = ChunksRetriever()
+    query = "Steve"
+    chunks = await retriever.get_retrieved_objects("Steve")
+    context = await retriever.get_context_from_objects(query=query, retrieved_objects=chunks)
 
-    context = await retriever.get_context("Steve")
+    completion = await retriever.get_completion_from_context(
+        query=query, retrieved_objects=chunks, context=context
+    )
 
-    assert isinstance(context, list), "Context should be a list"
-    assert len(context) > 0, "Context should not be empty"
-    assert any(chunk["text"] == "Steve Rodger" for chunk in context), (
+    assert isinstance(completion, list), "Retrieved objects should be a list"
+    assert len(completion) > 0, "Retrieved objects list should not be empty"
+    assert any(chunk["text"] == "Steve Rodger" for chunk in completion), (
         "Failed to get Steve Rodger chunk"
     )
 
@@ -218,35 +288,64 @@ async def test_chunks_retriever_context_multiple_chunks(setup_test_environment_w
 async def test_chunks_retriever_top_k_limit(setup_test_environment_with_chunks_complex):
     """Integration test: verify ChunksRetriever respects top_k parameter."""
     retriever = ChunksRetriever(top_k=2)
+    query = "Employee"
 
-    context = await retriever.get_context("Employee")
+    chunks = await retriever.get_retrieved_objects("Steve")
+    context = await retriever.get_context_from_objects(query=query, retrieved_objects=chunks)
 
-    assert isinstance(context, list), "Context should be a list"
-    assert len(context) <= 2, "Should respect top_k limit"
+    completion = await retriever.get_completion_from_context(
+        query=query, retrieved_objects=chunks, context=context
+    )
+
+    assert isinstance(completion, list), "Context should be a list"
+    assert len(completion) <= 2, "Should respect top_k limit"
 
 
 @pytest.mark.asyncio
 async def test_chunks_retriever_context_complex(setup_test_environment_with_chunks_complex):
     """Integration test: verify ChunksRetriever can retrieve chunk context (complex)."""
     retriever = ChunksRetriever(top_k=20)
+    query = "Christina"
+
+    chunks = await retriever.get_retrieved_objects(query)
 
-    context = await retriever.get_context("Christina")
+    context = await retriever.get_context_from_objects(query=query, retrieved_objects=chunks)
 
-    assert context[0]["text"] == "Christina Mayer", "Failed to get Christina Mayer"
+    assert context[0:15] == "Christina Mayer", "Failed to get Christina Mayer"
 
 
 @pytest.mark.asyncio
-async def test_chunks_retriever_context_on_empty_graph(setup_test_environment_empty):
+async def test_chunks_retriever_on_empty_graph(setup_test_environment_empty):
     """Integration test: verify ChunksRetriever handles empty graph correctly."""
     retriever = ChunksRetriever()
+    query = "Christina Mayer"
+
+    vector_engine = get_vector_engine()
+    await vector_engine.create_collection(
+        "DocumentChunk_text", payload_schema=DocumentChunkWithEntities
+    )
+
+    chunks = await retriever.get_retrieved_objects(query)
+    context = await retriever.get_context_from_objects(query=query, retrieved_objects=chunks)
 
-    with pytest.raises(NoDataError):
-        await retriever.get_context("Christina Mayer")
+    completion = await retriever.get_completion_from_context(
+        query=query, retrieved_objects=chunks, context=context
+    )
+    assert isinstance(completion, list), "Retrieved objects should be a list"
+    assert len(completion) == 0, "Found chunks when none should exist"
+
+
+@pytest.mark.asyncio
+async def test_chunks_retriever_context_on_empty_graph(setup_test_environment_empty):
+    """Integration test: verify ChunksRetriever context handles empty graph correctly."""
+    retriever = ChunksRetriever()
+    query = "Christina Mayer"
 
     vector_engine = get_vector_engine()
     await vector_engine.create_collection(
         "DocumentChunk_text", payload_schema=DocumentChunkWithEntities
     )
 
-    context = await retriever.get_context("Christina Mayer")
-    assert len(context) == 0, "Found chunks when none should exist"
+    chunks = await retriever.get_retrieved_objects(query)
+    context = await retriever.get_context_from_objects(query=query, retrieved_objects=chunks)
+    assert context == "", "Found chunks when none should exist"

cognee/tests/integration/retrieval/test_graph_completion_retriever.py

@@ -6,7 +6,6 @@ from typing import Optional, Union
 import cognee
 
 from cognee.low_level import setup, DataPoint
-from cognee.modules.graph.utils import resolve_edges_to_text
 from cognee.tasks.storage import add_data_points
 from cognee.modules.retrieval.graph_completion_retriever import GraphCompletionRetriever
 
@@ -174,8 +173,11 @@ async def setup_test_environment_empty():
 async def test_graph_completion_context_simple(setup_test_environment_simple):
     """Integration test: verify GraphCompletionRetriever can retrieve context (simple)."""
     retriever = GraphCompletionRetriever()
+    query = "Who works at Canva?"
 
-    context = await resolve_edges_to_text(await retriever.get_context("Who works at Canva?"))
+    triplets = await retriever.get_retrieved_objects(query)
+
+    context = await retriever.get_context_from_objects(query=query, retrieved_objects=triplets)
 
     # Ensure the top-level sections are present
     assert "Nodes:" in context, "Missing 'Nodes:' section in context"
@@ -240,8 +242,11 @@ async def test_graph_completion_context_simple(setup_test_environment_simple):
 async def test_graph_completion_context_complex(setup_test_environment_complex):
     """Integration test: verify GraphCompletionRetriever can retrieve context (complex)."""
     retriever = GraphCompletionRetriever(top_k=20)
+    query = "Who works at Figma?"
+
+    triplets = await retriever.get_retrieved_objects(query)
 
-    context = await resolve_edges_to_text(await retriever.get_context("Who works at Figma?"))
+    context = await retriever.get_context_from_objects(query=query, retrieved_objects=triplets)
 
     assert "Mike Rodger --[works_for]--> Figma" in context, "Failed to get Mike Rodger"
     assert "Ike Loma --[works_for]--> Figma" in context, "Failed to get Ike Loma"
@@ -252,9 +257,12 @@ async def test_graph_completion_context_complex(setup_test_environment_complex):
 async def test_get_graph_completion_context_on_empty_graph(setup_test_environment_empty):
     """Integration test: verify GraphCompletionRetriever handles empty graph correctly."""
     retriever = GraphCompletionRetriever()
+    query = "Who works at Figma?"
+
+    triplets = await retriever.get_retrieved_objects(query)
 
-    context = await retriever.get_context("Who works at Figma?")
-    assert context == [], "Context should be empty on an empty graph"
+    context = await retriever.get_context_from_objects(query=query, retrieved_objects=triplets)
+    assert context == "", "Context should be empty on an empty graph"
 
 
 @pytest.mark.asyncio