feat: Implements pipeline structure for retrievers

cognee/pipelines/retriever/diffusion_retriever.py

@@ -0,0 +1,25 @@
+from uuid import UUID
+from enum import Enum
+from typing import Callable, Dict
+from cognee.shared.utils import send_telemetry
+from cognee.modules.users.models import User
+from cognee.modules.users.methods import get_default_user
+from cognee.modules.users.permissions.methods import get_document_ids_for_user
+
+async def two_step_retriever(query: Dict[str, str], user: User = None) -> list:
+    if user is None:
+        user = await get_default_user()
+
+    if user is None:
+        raise PermissionError("No user found in the system. Please create a user.")
+
+    own_document_ids = await get_document_ids_for_user(user.id)
+    retrieved_results = await diffusion_retriever(query, user)
+
+    filtered_search_results = []
+
+
+    return retrieved_results
+
+async def diffusion_retriever(query: str, user, community_filter = []) -> list:
+    raise(NotImplementedError)

cognee/pipelines/retriever/g_retriever.py

@@ -0,0 +1,25 @@
+from uuid import UUID
+from enum import Enum
+from typing import Callable, Dict
+from cognee.shared.utils import send_telemetry
+from cognee.modules.users.models import User
+from cognee.modules.users.methods import get_default_user
+from cognee.modules.users.permissions.methods import get_document_ids_for_user
+
+async def two_step_retriever(query: Dict[str, str], user: User = None) -> list:
+    if user is None:
+        user = await get_default_user()
+
+    if user is None:
+        raise PermissionError("No user found in the system. Please create a user.")
+
+    own_document_ids = await get_document_ids_for_user(user.id)
+    retrieved_results = await g_retriever(query, user)
+
+    filtered_search_results = []
+
+
+    return retrieved_results
+
+async def g_retriever(query: str, user, community_filter = []) -> list:
+    raise(NotImplementedError)

cognee/pipelines/retriever/two_steps_retriever.py

@@ -0,0 +1,26 @@
+from uuid import UUID
+from enum import Enum
+from typing import Callable, Dict
+from cognee.shared.utils import send_telemetry
+from cognee.modules.users.models import User
+from cognee.modules.users.methods import get_default_user
+from cognee.modules.users.permissions.methods import get_document_ids_for_user
+
+async def two_step_retriever(query: Dict[str, str], user: User = None) -> list:
+    if user is None:
+        user = await get_default_user()
+
+    if user is None:
+        raise PermissionError("No user found in the system. Please create a user.")
+
+    own_document_ids = await get_document_ids_for_user(user.id)
+    retrieved_results = await run_two_step_retriever(query, user)
+
+    filtered_search_results = []
+
+
+    return retrieved_results
+
+
+async def run_two_step_retriever(query: str, user, community_filter = []) -> list:
+    raise(NotImplementedError)

examples/python/dynamic_steps_example.py

@@ -1,6 +1,6 @@
 import cognee
 import asyncio
-from cognee.api.v1.search import SearchType
+from cognee.pipelines.retriever.two_steps_retriever import two_step_retriever
 
 job_position = """0:Senior Data Scientist (Machine Learning)
 
@@ -206,9 +206,8 @@ async def main(enable_steps):
         print("Knowledge graph created.")
 
     # Step 4: Query insights
-    if enable_steps.get("search_insights"):
+    if enable_steps.get("retriever"):
-        search_results = await cognee.search(
+        search_results = await two_step_retriever(
-            SearchType.INSIGHTS,
             {'query': 'Which applicant has the most relevant experience in data science?'}
         )
         print("Search results:")
@@ -219,11 +218,11 @@ async def main(enable_steps):
 if __name__ == '__main__':
     # Flags to enable/disable steps
     steps_to_enable = {
-        "prune_data": True,
+        "prune_data": False,
-        "prune_system": True,
+        "prune_system": False,
-        "add_text": True,
+        "add_text": False,
-        "cognify": True,
+        "cognify": False,
-        "search_insights": True
+        "retriever": True
     }
 
     asyncio.run(main(steps_to_enable))