cognee 0.2.2.dev0__py3-none-any.whl → 0.2.3__py3-none-any.whl

cognee/tests/test_neptune_analytics_hybrid.py

@@ -0,0 +1,176 @@
+import os
+from dotenv import load_dotenv
+import asyncio
+import pytest
+
+from cognee.modules.chunking.models import DocumentChunk
+from cognee.modules.engine.models import Entity, EntityType
+from cognee.modules.data.processing.document_types import TextDocument
+from cognee.infrastructure.databases.vector.embeddings import get_embedding_engine
+from cognee.shared.logging_utils import get_logger
+from cognee.infrastructure.databases.hybrid.neptune_analytics.NeptuneAnalyticsAdapter import (
+    NeptuneAnalyticsAdapter,
+)
+
+# Set up Amazon credentials in .env file and get the values from environment variables
+load_dotenv()
+graph_id = os.getenv("GRAPH_ID", "")
+
+# get the default embedder
+embedding_engine = get_embedding_engine()
+na_graph = NeptuneAnalyticsAdapter(graph_id)
+na_vector = NeptuneAnalyticsAdapter(graph_id, embedding_engine)
+
+collection = "test_collection"
+
+logger = get_logger("test_neptune_analytics_hybrid")
+
+
+def setup_data():
+    # Define nodes data before the main function
+    # These nodes were defined using openAI from the following prompt:
+    #
+    # Neptune Analytics is an ideal choice for investigatory, exploratory, or data-science workloads
+    #     that require fast iteration for data, analytical and algorithmic processing, or vector search on graph data. It
+    #     complements Amazon Neptune Database, a popular managed graph database. To perform intensive analysis, you can load
+    #     the data from a Neptune Database graph or snapshot into Neptune Analytics. You can also load graph data that's
+    #     stored in Amazon S3.
+
+    document = TextDocument(
+        name="text.txt",
+        raw_data_location="git/cognee/examples/database_examples/data_storage/data/text.txt",
+        external_metadata="{}",
+        mime_type="text/plain",
+    )
+    document_chunk = DocumentChunk(
+        text="Neptune Analytics is an ideal choice for investigatory, exploratory, or data-science workloads \n    that require fast iteration for data, analytical and algorithmic processing, or vector search on graph data. It \n    complements Amazon Neptune Database, a popular managed graph database. To perform intensive analysis, you can load \n    the data from a Neptune Database graph or snapshot into Neptune Analytics. You can also load graph data that's \n    stored in Amazon S3.\n    ",
+        chunk_size=187,
+        chunk_index=0,
+        cut_type="paragraph_end",
+        is_part_of=document,
+    )
+
+    graph_database = EntityType(name="graph database", description="graph database")
+    neptune_analytics_entity = Entity(
+        name="neptune analytics",
+        description="A memory-optimized graph database engine for analytics that processes large amounts of graph data quickly.",
+    )
+    neptune_database_entity = Entity(
+        name="amazon neptune database",
+        description="A popular managed graph database that complements Neptune Analytics.",
+    )
+
+    storage = EntityType(name="storage", description="storage")
+    storage_entity = Entity(
+        name="amazon s3",
+        description="A storage service provided by Amazon Web Services that allows storing graph data.",
+    )
+
+    nodes_data = [
+        document,
+        document_chunk,
+        graph_database,
+        neptune_analytics_entity,
+        neptune_database_entity,
+        storage,
+        storage_entity,
+    ]
+
+    edges_data = [
+        (
+            str(document_chunk.id),
+            str(storage_entity.id),
+            "contains",
+        ),
+        (
+            str(storage_entity.id),
+            str(storage.id),
+            "is_a",
+        ),
+        (
+            str(document_chunk.id),
+            str(neptune_database_entity.id),
+            "contains",
+        ),
+        (
+            str(neptune_database_entity.id),
+            str(graph_database.id),
+            "is_a",
+        ),
+        (
+            str(document_chunk.id),
+            str(document.id),
+            "is_part_of",
+        ),
+        (
+            str(document_chunk.id),
+            str(neptune_analytics_entity.id),
+            "contains",
+        ),
+        (
+            str(neptune_analytics_entity.id),
+            str(graph_database.id),
+            "is_a",
+        ),
+    ]
+    return nodes_data, edges_data
+
+
+async def test_add_graph_then_vector_data():
+    logger.info("------test_add_graph_then_vector_data-------")
+    (nodes, edges) = setup_data()
+    await na_graph.add_nodes(nodes)
+    await na_graph.add_edges(edges)
+    await na_vector.create_data_points(collection, nodes)
+
+    node_ids = [str(node.id) for node in nodes]
+    retrieved_data_points = await na_vector.retrieve(collection, node_ids)
+    retrieved_nodes = await na_graph.get_nodes(node_ids)
+
+    assert len(retrieved_data_points) == len(retrieved_nodes) == len(node_ids)
+
+    # delete all nodes and edges and vectors:
+    await na_graph.delete_graph()
+    await na_vector.prune()
+
+    (nodes, edges) = await na_graph.get_graph_data()
+    assert len(nodes) == 0
+    assert len(edges) == 0
+    logger.info("------PASSED-------")
+
+
+async def test_add_vector_then_node_data():
+    logger.info("------test_add_vector_then_node_data-------")
+    (nodes, edges) = setup_data()
+    await na_vector.create_data_points(collection, nodes)
+    await na_graph.add_nodes(nodes)
+    await na_graph.add_edges(edges)
+
+    node_ids = [str(node.id) for node in nodes]
+    retrieved_data_points = await na_vector.retrieve(collection, node_ids)
+    retrieved_nodes = await na_graph.get_nodes(node_ids)
+
+    assert len(retrieved_data_points) == len(retrieved_nodes) == len(node_ids)
+
+    # delete all nodes and edges and vectors:
+    await na_vector.prune()
+    await na_graph.delete_graph()
+
+    (nodes, edges) = await na_graph.get_graph_data()
+    assert len(nodes) == 0
+    assert len(edges) == 0
+    logger.info("------PASSED-------")
+
+
+def main():
+    """
+    Example script uses neptune analytics for the graph and vector (hybrid) store with small sample data
+    This example demonstrates how to add nodes and vectors to Neptune Analytics, and ensures that
+    the nodes do not conflict
+    """
+    asyncio.run(test_add_graph_then_vector_data())
+    asyncio.run(test_add_vector_then_node_data())
+
+
+if __name__ == "__main__":
+    main()

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

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

cognee/tests/test_pgvector.py

@@ -37,16 +37,16 @@ async def test_local_file_deletion(data_text, file_location):
         # Get data entry from database based on file path
         data = (
             await session.scalars(
-                select(Data).where(Data.raw_data_location == "file://" + file_location)
+                select(Data).where(Data.original_data_location == "file://" + file_location)
             )
         ).one()
-        assert os.path.isfile(data.raw_data_location.replace("file://", "")), (
-            f"Data location doesn't exist: {data.raw_data_location}"
+        assert os.path.isfile(data.original_data_location.replace("file://", "")), (
+            f"Data location doesn't exist: {data.original_data_location}"
         )
         # Test local files not created by cognee won't get deleted
         await engine.delete_data_entity(data.id)
-        assert os.path.exists(data.raw_data_location.replace("file://", "")), (
-            f"Data location doesn't exists: {data.raw_data_location}"
+        assert os.path.exists(data.original_data_location.replace("file://", "")), (
+            f"Data location doesn't exists: {data.original_data_location}"
         )
 
 

cognee/tests/test_s3.py

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

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

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

cognee/tests/unit/infrastructure/mock_embedding_engine.py

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

cognee/tests/unit/infrastructure/test_embedding_rate_limiting_realistic.py

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

cognee/tests/unit/infrastructure/test_rate_limiting_realistic.py

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

cognee/tests/unit/infrastructure/test_rate_limiting_retry.py

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

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

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

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

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