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

cognee/tests/tasks/translation/providers_test.py

@@ -0,0 +1,151 @@
+"""
+Unit tests for translation providers
+"""
+
+import os
+
+import pytest
+
+from cognee.tasks.translation.providers import (
+    get_translation_provider,
+    LLMTranslationProvider,
+    TranslationResult,
+)
+from cognee.tasks.translation.exceptions import TranslationError
+
+
+def has_llm_api_key():
+    """Check if LLM API key is available"""
+    return bool(os.environ.get("LLM_API_KEY"))
+
+
+@pytest.mark.asyncio
+@pytest.mark.skipif(not has_llm_api_key(), reason="No LLM API key available")
+async def test_llm_provider_basic_translation():
+    """Test basic translation with LLM provider (uses configured LLM)"""
+    provider = LLMTranslationProvider()
+
+    result = await provider.translate(text="Hola mundo", target_language="en", source_language="es")
+
+    assert isinstance(result, TranslationResult)
+    assert result.translated_text is not None
+    assert len(result.translated_text) > 0
+    assert result.source_language == "es"
+    assert result.target_language == "en"
+    assert result.provider == "llm"
+
+
+@pytest.mark.asyncio
+@pytest.mark.skipif(not has_llm_api_key(), reason="No LLM API key available")
+async def test_llm_provider_auto_detect_source():
+    """Test translation with automatic source language detection"""
+    provider = LLMTranslationProvider()
+
+    result = await provider.translate(
+        text="Bonjour le monde",
+        target_language="en",
+        # source_language not provided - should auto-detect
+    )
+
+    assert result.translated_text is not None
+    assert result.target_language == "en"
+
+
+@pytest.mark.asyncio
+@pytest.mark.skipif(not has_llm_api_key(), reason="No LLM API key available")
+async def test_llm_provider_long_text():
+    """Test translation of longer text"""
+    provider = LLMTranslationProvider()
+
+    long_text = """
+    La inteligencia artificial es una rama de la informática que se centra en
+    crear sistemas capaces de realizar tareas que normalmente requieren inteligencia humana.
+    Estos sistemas pueden aprender, razonar y resolver problemas complejos.
+    """
+
+    result = await provider.translate(text=long_text, target_language="en", source_language="es")
+
+    assert len(result.translated_text) > 0
+    assert result.source_language == "es"
+
+
+def test_get_translation_provider_factory():
+    """Test provider factory function"""
+    provider = get_translation_provider("llm")
+    assert isinstance(provider, LLMTranslationProvider)
+
+
+def test_get_translation_provider_invalid():
+    """Test provider factory with invalid provider name"""
+    try:
+        get_translation_provider("invalid_provider")
+        assert False, "Expected TranslationError or ValueError"
+    except (TranslationError, ValueError):
+        pass
+
+
+@pytest.mark.asyncio
+@pytest.mark.skipif(not has_llm_api_key(), reason="No LLM API key available")
+async def test_llm_batch_translation():
+    """Test batch translation with LLM provider"""
+    provider = LLMTranslationProvider()
+
+    texts = ["Hola", "¿Cómo estás?", "Adiós"]
+
+    results = await provider.translate_batch(
+        texts=texts, target_language="en", source_language="es"
+    )
+
+    assert len(results) == len(texts)
+    for result in results:
+        assert isinstance(result, TranslationResult)
+        assert result.translated_text is not None
+        assert result.source_language == "es"
+        assert result.target_language == "en"
+
+
+@pytest.mark.asyncio
+@pytest.mark.skipif(not has_llm_api_key(), reason="No LLM API key available")
+async def test_translation_preserves_formatting():
+    """Test that translation preserves basic formatting"""
+    provider = LLMTranslationProvider()
+
+    text_with_newlines = "Primera línea.\nSegunda línea."
+
+    result = await provider.translate(
+        text=text_with_newlines, target_language="en", source_language="es"
+    )
+
+    # Should preserve structure (though exact newlines may vary)
+    assert result.translated_text is not None
+    assert len(result.translated_text) > 0
+
+
+@pytest.mark.asyncio
+@pytest.mark.skipif(not has_llm_api_key(), reason="No LLM API key available")
+async def test_translation_special_characters():
+    """Test translation with special characters"""
+    provider = LLMTranslationProvider()
+
+    text = "¡Hola! ¿Cómo estás? Está bien."
+
+    result = await provider.translate(text=text, target_language="en", source_language="es")
+
+    assert result.translated_text is not None
+    assert len(result.translated_text) > 0
+
+
+@pytest.mark.asyncio
+@pytest.mark.skipif(not has_llm_api_key(), reason="No LLM API key available")
+async def test_empty_text_translation():
+    """Test translation with empty text - should return empty or handle gracefully"""
+    provider = LLMTranslationProvider()
+
+    # Empty text may either raise an error or return an empty result
+    try:
+        result = await provider.translate(text="", target_language="en", source_language="es")
+        # If no error, should return a TranslationResult (possibly with empty text)
+        assert isinstance(result, TranslationResult)
+    except TranslationError:
+        # This is also acceptable behavior
+        pass

cognee/tests/tasks/translation/translate_content_test.py

@@ -0,0 +1,213 @@
+"""
+Unit tests for translate_content task
+"""
+
+import os
+from uuid import uuid4
+
+import pytest
+
+from cognee.modules.chunking.models import DocumentChunk
+from cognee.modules.data.processing.document_types import TextDocument
+from cognee.tasks.translation import translate_content
+from cognee.tasks.translation.models import TranslatedContent, LanguageMetadata
+
+
+def has_llm_api_key():
+    """Check if LLM API key is available"""
+    return bool(os.environ.get("LLM_API_KEY"))
+
+
+def create_test_chunk(text: str, chunk_index: int = 0):
+    """Helper to create a DocumentChunk with all required fields"""
+    # Create a minimal Document for the is_part_of field
+    doc = TextDocument(
+        id=uuid4(),
+        name="test_doc",
+        raw_data_location="/tmp/test.txt",
+        external_metadata=None,
+        mime_type="text/plain",
+    )
+
+    return DocumentChunk(
+        id=uuid4(),
+        text=text,
+        chunk_index=chunk_index,
+        chunk_size=len(text),
+        cut_type="sentence",
+        is_part_of=doc,
+    )
+
+
+@pytest.mark.asyncio
+@pytest.mark.skipif(not has_llm_api_key(), reason="No LLM API key available")
+async def test_translate_content_basic():
+    """Test basic content translation"""
+    # Create test chunk with Spanish text
+    original_text = "Hola mundo, esta es una prueba."
+    chunk = create_test_chunk(original_text)
+
+    result = await translate_content(
+        data_chunks=[chunk], target_language="en", translation_provider="llm"
+    )
+
+    assert len(result) == 1
+    # The chunk's text should now be translated (different from original Spanish)
+    assert result[0].text != original_text  # Text should be translated to English
+    assert result[0].contains is not None
+
+    # Check for TranslatedContent in contains
+    has_translated_content = any(isinstance(item, TranslatedContent) for item in result[0].contains)
+    assert has_translated_content
+
+
+@pytest.mark.asyncio
+@pytest.mark.skipif(not has_llm_api_key(), reason="No LLM API key available")
+async def test_translate_content_preserves_original():
+    """Test that original text is preserved"""
+    original_text = "Bonjour le monde"
+    chunk = create_test_chunk(original_text)
+
+    result = await translate_content(
+        data_chunks=[chunk], target_language="en", preserve_original=True
+    )
+
+    # Find TranslatedContent in contains
+    translated_content = None
+    for item in result[0].contains:
+        if isinstance(item, TranslatedContent):
+            translated_content = item
+            break
+
+    assert translated_content is not None
+    assert translated_content.original_text == original_text
+    assert translated_content.translated_text != original_text
+
+
+@pytest.mark.asyncio
+async def test_translate_content_skip_english():
+    """Test skipping translation for English text"""
+    # This test doesn't require API call since English text is skipped
+    chunk = create_test_chunk("Hello world, this is a test.")
+
+    result = await translate_content(
+        data_chunks=[chunk], target_language="en", skip_if_target_language=True
+    )
+
+    # Text should remain unchanged
+    assert result[0].text == chunk.text
+
+    # Should have LanguageMetadata but not TranslatedContent
+    has_language_metadata = any(
+        isinstance(item, LanguageMetadata) for item in (result[0].contains or [])
+    )
+    has_translated_content = any(
+        isinstance(item, TranslatedContent) for item in (result[0].contains or [])
+    )
+
+    assert has_language_metadata
+    assert not has_translated_content
+
+
+@pytest.mark.asyncio
+@pytest.mark.skipif(not has_llm_api_key(), reason="No LLM API key available")
+async def test_translate_content_multiple_chunks():
+    """Test translation of multiple chunks"""
+    # Use longer texts to ensure reliable language detection
+    original_texts = [
+        "Hola mundo, esta es una prueba de traducción.",
+        "Bonjour le monde, ceci est un test de traduction.",
+        "Ciao mondo, questo è un test di traduzione.",
+    ]
+    chunks = [create_test_chunk(text, i) for i, text in enumerate(original_texts)]
+
+    result = await translate_content(data_chunks=chunks, target_language="en")
+
+    assert len(result) == 3
+    # Check that at least some chunks were translated
+    translated_count = sum(
+        1
+        for chunk in result
+        if any(isinstance(item, TranslatedContent) for item in (chunk.contains or []))
+    )
+    assert translated_count >= 2  # At least 2 chunks should be translated
+
+
+@pytest.mark.asyncio
+async def test_translate_content_empty_list():
+    """Test with empty chunk list"""
+    result = await translate_content(data_chunks=[], target_language="en")
+
+    assert result == []
+
+
+@pytest.mark.asyncio
+async def test_translate_content_empty_text():
+    """Test with chunk containing empty text"""
+    chunk = create_test_chunk("")
+
+    result = await translate_content(data_chunks=[chunk], target_language="en")
+
+    assert len(result) == 1
+    assert result[0].text == ""
+
+
+@pytest.mark.asyncio
+@pytest.mark.skipif(not has_llm_api_key(), reason="No LLM API key available")
+async def test_translate_content_language_metadata():
+    """Test that LanguageMetadata is created correctly"""
+    # Use a longer, distinctly Spanish text to ensure reliable detection
+    chunk = create_test_chunk(
+        "La inteligencia artificial está cambiando el mundo de manera significativa"
+    )
+
+    result = await translate_content(data_chunks=[chunk], target_language="en")
+
+    # Find LanguageMetadata
+    language_metadata = None
+    for item in result[0].contains:
+        if isinstance(item, LanguageMetadata):
+            language_metadata = item
+            break
+
+    assert language_metadata is not None
+    # Just check that a language was detected (short texts can be ambiguous)
+    assert language_metadata.detected_language is not None
+    assert language_metadata.requires_translation is True
+    assert language_metadata.language_confidence > 0.0
+
+
+@pytest.mark.asyncio
+@pytest.mark.skipif(not has_llm_api_key(), reason="No LLM API key available")
+async def test_translate_content_confidence_threshold():
+    """Test with custom confidence threshold"""
+    # Use longer text for more reliable detection
+    chunk = create_test_chunk("Hola mundo, esta es una frase más larga para mejor detección")
+
+    result = await translate_content(
+        data_chunks=[chunk], target_language="en", confidence_threshold=0.5
+    )
+
+    assert len(result) == 1
+
+
+@pytest.mark.asyncio
+@pytest.mark.skipif(not has_llm_api_key(), reason="No LLM API key available")
+async def test_translate_content_no_preserve_original():
+    """Test translation without preserving original"""
+    # Use longer text for more reliable detection
+    chunk = create_test_chunk("Bonjour le monde, comment allez-vous aujourd'hui")
+
+    result = await translate_content(
+        data_chunks=[chunk], target_language="en", preserve_original=False
+    )
+
+    # Find TranslatedContent
+    translated_content = None
+    for item in result[0].contains:
+        if isinstance(item, TranslatedContent):
+            translated_content = item
+            break
+
+    assert translated_content is not None
+    assert translated_content.original_text == ""  # Should be empty

cognee/tests/test_chromadb.py

@@ -97,7 +97,7 @@ async def test_vector_engine_search_none_limit():
     query_vector = (await vector_engine.embedding_engine.embed_text([query_text]))[0]
 
     result = await vector_engine.search(
-        collection_name=collection_name, query_vector=query_vector, limit=None
+        collection_name=collection_name, query_vector=query_vector, limit=None, include_payload=True
     )
 
     # Check that we did not accidentally use any default value for limit

cognee/tests/test_cleanup_unused_data.py

@@ -0,0 +1,165 @@
+import os
+import pathlib
+import cognee
+from datetime import datetime, timezone, timedelta
+from uuid import UUID
+from sqlalchemy import select, update
+from cognee.modules.data.models import Data, DatasetData
+from cognee.infrastructure.databases.relational import get_relational_engine
+from cognee.modules.users.methods import get_default_user
+from cognee.shared.logging_utils import get_logger
+from cognee.modules.search.types import SearchType
+
+logger = get_logger()
+
+
+async def test_textdocument_cleanup_with_sql():
+    """
+    End-to-end test for TextDocument cleanup based on last_accessed timestamps.
+    """
+    # Enable last accessed tracking BEFORE any cognee operations
+    os.environ["ENABLE_LAST_ACCESSED"] = "true"
+
+    # Setup test directories
+    data_directory_path = str(
+        pathlib.Path(
+            os.path.join(pathlib.Path(__file__).parent, ".data_storage/test_cleanup")
+        ).resolve()
+    )
+    cognee_directory_path = str(
+        pathlib.Path(
+            os.path.join(pathlib.Path(__file__).parent, ".cognee_system/test_cleanup")
+        ).resolve()
+    )
+
+    cognee.config.data_root_directory(data_directory_path)
+    cognee.config.system_root_directory(cognee_directory_path)
+
+    # Initialize database
+    from cognee.modules.engine.operations.setup import setup
+
+    # Clean slate
+    await cognee.prune.prune_data()
+    await cognee.prune.prune_system(metadata=True)
+
+    logger.info("🧪 Testing TextDocument cleanup based on last_accessed")
+
+    # Step 1: Add and cognify a test document
+    dataset_name = "test_cleanup_dataset"
+    test_text = """
+    Machine learning is a subset of artificial intelligence that enables systems to learn
+    and improve from experience without being explicitly programmed. Deep learning uses
+    neural networks with multiple layers to process data.
+    """
+
+    await setup()
+    user = await get_default_user()
+    await cognee.add([test_text], dataset_name=dataset_name, user=user)
+
+    cognify_result = await cognee.cognify([dataset_name], user=user)
+
+    # Extract dataset_id from cognify result
+    dataset_id = None
+    for ds_id, pipeline_result in cognify_result.items():
+        dataset_id = ds_id
+        break
+
+    assert dataset_id is not None, "Failed to get dataset_id from cognify result"
+    logger.info(f"✅ Document added and cognified. Dataset ID: {dataset_id}")
+
+    # Step 2: Perform search to trigger last_accessed update
+    logger.info("Triggering search to update last_accessed...")
+    search_results = await cognee.search(
+        query_type=SearchType.CHUNKS,
+        query_text="machine learning",
+        datasets=[dataset_name],
+        user=user,
+    )
+    logger.info(f"✅ Search completed, found {len(search_results)} results")
+    assert len(search_results) > 0, "Search should return results"
+
+    # Step 3: Verify last_accessed was set and get data_id
+    db_engine = get_relational_engine()
+    async with db_engine.get_async_session() as session:
+        result = await session.execute(
+            select(Data, DatasetData)
+            .join(DatasetData, Data.id == DatasetData.data_id)
+            .where(DatasetData.dataset_id == dataset_id)
+        )
+        data_records = result.all()
+        assert len(data_records) > 0, "No Data records found for the dataset"
+        data_record = data_records[0][0]
+        data_id = data_record.id
+
+        # Verify last_accessed is set
+        assert data_record.last_accessed is not None, (
+            "last_accessed should be set after search operation"
+        )
+
+        original_last_accessed = data_record.last_accessed
+        logger.info(f"✅ last_accessed verified: {original_last_accessed}")
+
+    # Step 4: Manually age the timestamp
+    minutes_threshold = 30
+    aged_timestamp = datetime.now(timezone.utc) - timedelta(minutes=minutes_threshold + 10)
+
+    async with db_engine.get_async_session() as session:
+        stmt = update(Data).where(Data.id == data_id).values(last_accessed=aged_timestamp)
+        await session.execute(stmt)
+        await session.commit()
+
+    # Verify timestamp was updated
+    async with db_engine.get_async_session() as session:
+        result = await session.execute(select(Data).where(Data.id == data_id))
+        updated_data = result.scalar_one_or_none()
+        assert updated_data is not None, "Data record should exist"
+        retrieved_timestamp = updated_data.last_accessed
+        if retrieved_timestamp.tzinfo is None:
+            retrieved_timestamp = retrieved_timestamp.replace(tzinfo=timezone.utc)
+        assert retrieved_timestamp == aged_timestamp, "Timestamp should be updated to aged value"
+
+    # Step 5: Test cleanup (document-level is now the default)
+    from cognee.tasks.cleanup.cleanup_unused_data import cleanup_unused_data
+
+    # First do a dry run
+    logger.info("Testing dry run...")
+    dry_run_result = await cleanup_unused_data(minutes_threshold=10, dry_run=True, user_id=user.id)
+
+    # Debug: Print the actual result
+    logger.info(f"Dry run result: {dry_run_result}")
+
+    assert dry_run_result["status"] == "dry_run", (
+        f"Status should be 'dry_run', got: {dry_run_result['status']}"
+    )
+    assert dry_run_result["unused_count"] > 0, "Should find at least one unused document"
+    logger.info(f"✅ Dry run found {dry_run_result['unused_count']} unused documents")
+
+    # Now run actual cleanup
+    logger.info("Executing cleanup...")
+    cleanup_result = await cleanup_unused_data(minutes_threshold=30, dry_run=False, user_id=user.id)
+
+    assert cleanup_result["status"] == "completed", "Cleanup should complete successfully"
+    assert cleanup_result["deleted_count"]["documents"] > 0, (
+        "At least one document should be deleted"
+    )
+    logger.info(
+        f"✅ Cleanup completed. Deleted {cleanup_result['deleted_count']['documents']} documents"
+    )
+
+    # Step 6: Verify deletion
+    async with db_engine.get_async_session() as session:
+        deleted_data = (
+            await session.execute(select(Data).where(Data.id == data_id))
+        ).scalar_one_or_none()
+        assert deleted_data is None, "Data record should be deleted"
+        logger.info("✅ Confirmed: Data record was deleted")
+
+    logger.info("🎉 All cleanup tests passed!")
+    return True
+
+
+if __name__ == "__main__":
+    import asyncio
+
+    success = asyncio.run(test_textdocument_cleanup_with_sql())
+    exit(0 if success else 1)

cognee/tests/test_delete_by_id.py

@@ -47,20 +47,20 @@ async def main():
 
     # Test data
     text_1 = """
-    Apple Inc. is an American multinational technology company that specializes in consumer electronics, 
-    software, and online services. Apple is the world's largest technology company by revenue and, 
+    Apple Inc. is an American multinational technology company that specializes in consumer electronics,
+    software, and online services. Apple is the world's largest technology company by revenue and,
     since January 2021, the world's most valuable company.
     """
 
     text_2 = """
-    Microsoft Corporation is an American multinational technology corporation which produces computer software, 
-    consumer electronics, personal computers, and related services. Its best known software products are the 
+    Microsoft Corporation is an American multinational technology corporation which produces computer software,
+    consumer electronics, personal computers, and related services. Its best known software products are the
     Microsoft Windows line of operating systems and the Microsoft Office suite.
     """
 
     text_3 = """
-    Google LLC is an American multinational technology company that specializes in Internet-related services and products, 
-    which include online advertising technologies, search engine, cloud computing, software, and hardware. Google has been 
+    Google LLC is an American multinational technology company that specializes in Internet-related services and products,
+    which include online advertising technologies, search engine, cloud computing, software, and hardware. Google has been
     referred to as the most powerful company in the world and one of the world's most valuable brands.
     """