cognee 0.2.3.dev1__py3-none-any.whl → 0.3.0__py3-none-any.whl

cognee/modules/search/methods/search.py

@@ -2,31 +2,28 @@ import os
 import json
 import asyncio
 from uuid import UUID
-from typing import Callable, List, Optional, Type, Union
-from cognee.modules.search.exceptions import UnsupportedSearchTypeError
+from fastapi.encoders import jsonable_encoder
+from typing import Any, List, Optional, Tuple, Type, Union
+
+from cognee.shared.utils import send_telemetry
 from cognee.context_global_variables import set_database_global_context_variables
-from cognee.modules.retrieval.chunks_retriever import ChunksRetriever
-from cognee.modules.retrieval.insights_retriever import InsightsRetriever
-from cognee.modules.retrieval.summaries_retriever import SummariesRetriever
-from cognee.modules.retrieval.completion_retriever import CompletionRetriever
-from cognee.modules.retrieval.graph_completion_retriever import GraphCompletionRetriever
-from cognee.modules.retrieval.graph_summary_completion_retriever import (
-    GraphSummaryCompletionRetriever,
-)
-from cognee.modules.retrieval.graph_completion_cot_retriever import GraphCompletionCotRetriever
-from cognee.modules.retrieval.graph_completion_context_extension_retriever import (
-    GraphCompletionContextExtensionRetriever,
+
+from cognee.modules.engine.models.node_set import NodeSet
+from cognee.modules.graph.cognee_graph.CogneeGraphElements import Edge
+from cognee.modules.search.types import (
+    SearchResult,
+    CombinedSearchResult,
+    SearchResultDataset,
+    SearchType,
 )
-from cognee.modules.retrieval.code_retriever import CodeRetriever
-from cognee.modules.retrieval.cypher_search_retriever import CypherSearchRetriever
-from cognee.modules.retrieval.natural_language_retriever import NaturalLanguageRetriever
-from cognee.modules.search.types import SearchType
-from cognee.modules.storage.utils import JSONEncoder
+from cognee.modules.search.operations import log_query, log_result
 from cognee.modules.users.models import User
 from cognee.modules.data.models import Dataset
-from cognee.shared.utils import send_telemetry
 from cognee.modules.users.permissions.methods import get_specific_user_permission_datasets
-from cognee.modules.search.operations import log_query, log_result, select_search_type
+
+from .get_search_type_tools import get_search_type_tools
+from .no_access_control_search import no_access_control_search
+from ..utils.prepare_search_result import prepare_search_result
 
 
 async def search(
@@ -35,10 +32,15 @@ async def search(
     dataset_ids: Union[list[UUID], None],
     user: User,
     system_prompt_path="answer_simple_question.txt",
+    system_prompt: Optional[str] = None,
     top_k: int = 10,
-    node_type: Optional[Type] = None,
+    node_type: Optional[Type] = NodeSet,
     node_name: Optional[List[str]] = None,
-):
+    save_interaction: bool = False,
+    last_k: Optional[int] = None,
+    only_context: bool = False,
+    use_combined_context: bool = False,
+) -> Union[CombinedSearchResult, List[SearchResult]]:
     """
 
     Args:
@@ -54,156 +56,275 @@ async def search(
     Notes:
         Searching by dataset is only available in ENABLE_BACKEND_ACCESS_CONTROL mode
     """
+    query = await log_query(query_text, query_type.value, user.id)
+    send_telemetry("cognee.search EXECUTION STARTED", user.id)
+
     # Use search function filtered by permissions if access control is enabled
     if os.getenv("ENABLE_BACKEND_ACCESS_CONTROL", "false").lower() == "true":
-        return await authorized_search(
-            query_text, query_type, user, dataset_ids, system_prompt_path, top_k
+        search_results = await authorized_search(
+            query_type=query_type,
+            query_text=query_text,
+            user=user,
+            dataset_ids=dataset_ids,
+            system_prompt_path=system_prompt_path,
+            system_prompt=system_prompt,
+            top_k=top_k,
+            node_type=node_type,
+            node_name=node_name,
+            save_interaction=save_interaction,
+            last_k=last_k,
+            only_context=only_context,
+            use_combined_context=use_combined_context,
         )
+    else:
+        search_results = [
+            await no_access_control_search(
+                query_type=query_type,
+                query_text=query_text,
+                system_prompt_path=system_prompt_path,
+                system_prompt=system_prompt,
+                top_k=top_k,
+                node_type=node_type,
+                node_name=node_name,
+                save_interaction=save_interaction,
+                last_k=last_k,
+                only_context=only_context,
+            )
+        ]
 
-    query = await log_query(query_text, query_type.value, user.id)
-
-    search_results = await specific_search(
-        query_type,
-        query_text,
-        user,
-        system_prompt_path=system_prompt_path,
-        top_k=top_k,
-        node_type=node_type,
-        node_name=node_name,
-    )
+    send_telemetry("cognee.search EXECUTION COMPLETED", user.id)
 
     await log_result(
         query.id,
         json.dumps(
-            search_results if len(search_results) > 1 else search_results[0], cls=JSONEncoder
+            jsonable_encoder(
+                await prepare_search_result(
+                    search_results[0] if isinstance(search_results, list) else search_results
+                )
+                if use_combined_context
+                else [
+                    await prepare_search_result(search_result) for search_result in search_results
+                ]
+            )
         ),
         user.id,
     )
 
-    return search_results
+    if use_combined_context:
+        prepared_search_results = await prepare_search_result(
+            search_results[0] if isinstance(search_results, list) else search_results
+        )
+        result = prepared_search_results["result"]
+        graphs = prepared_search_results["graphs"]
+        context = prepared_search_results["context"]
+        datasets = prepared_search_results["datasets"]
+
+        return CombinedSearchResult(
+            result=result,
+            graphs=graphs,
+            context=context,
+            datasets=[
+                SearchResultDataset(
+                    id=dataset.id,
+                    name=dataset.name,
+                )
+                for dataset in datasets
+            ],
+        )
+    else:
+        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 specific_search(
+async def authorized_search(
     query_type: SearchType,
-    query: str,
+    query_text: str,
     user: User,
-    system_prompt_path="answer_simple_question.txt",
+    dataset_ids: Optional[list[UUID]] = None,
+    system_prompt_path: str = "answer_simple_question.txt",
+    system_prompt: Optional[str] = None,
     top_k: int = 10,
-    node_type: Optional[Type] = None,
+    node_type: Optional[Type] = NodeSet,
     node_name: Optional[List[str]] = None,
-) -> list:
-    search_tasks: dict[SearchType, Callable] = {
-        SearchType.SUMMARIES: SummariesRetriever(top_k=top_k).get_completion,
-        SearchType.INSIGHTS: InsightsRetriever(top_k=top_k).get_completion,
-        SearchType.CHUNKS: ChunksRetriever(top_k=top_k).get_completion,
-        SearchType.RAG_COMPLETION: CompletionRetriever(
-            system_prompt_path=system_prompt_path, top_k=top_k
-        ).get_completion,
-        SearchType.GRAPH_COMPLETION: GraphCompletionRetriever(
-            system_prompt_path=system_prompt_path,
-            top_k=top_k,
-            node_type=node_type,
-            node_name=node_name,
-        ).get_completion,
-        SearchType.GRAPH_COMPLETION_COT: GraphCompletionCotRetriever(
-            system_prompt_path=system_prompt_path,
-            top_k=top_k,
-            node_type=node_type,
-            node_name=node_name,
-        ).get_completion,
-        SearchType.GRAPH_COMPLETION_CONTEXT_EXTENSION: GraphCompletionContextExtensionRetriever(
-            system_prompt_path=system_prompt_path,
-            top_k=top_k,
-            node_type=node_type,
-            node_name=node_name,
-        ).get_completion,
-        SearchType.GRAPH_SUMMARY_COMPLETION: GraphSummaryCompletionRetriever(
+    save_interaction: bool = False,
+    last_k: Optional[int] = None,
+    only_context: bool = False,
+    use_combined_context: bool = False,
+) -> Union[
+    Tuple[Any, Union[List[Edge], str], List[Dataset]],
+    List[Tuple[Any, Union[List[Edge], str], List[Dataset]]],
+]:
+    """
+    Verifies access for provided datasets or uses all datasets user has read access for and performs search per dataset.
+    Not to be used outside of active access control mode.
+    """
+    # Find datasets user has read access for (if datasets are provided only return them. Provided user has read access)
+    search_datasets = await get_specific_user_permission_datasets(user.id, "read", dataset_ids)
+
+    if use_combined_context:
+        search_responses = await search_in_datasets_context(
+            search_datasets=search_datasets,
+            query_type=query_type,
+            query_text=query_text,
             system_prompt_path=system_prompt_path,
+            system_prompt=system_prompt,
             top_k=top_k,
             node_type=node_type,
             node_name=node_name,
-        ).get_completion,
-        SearchType.CODE: CodeRetriever(top_k=top_k).get_completion,
-        SearchType.CYPHER: CypherSearchRetriever().get_completion,
-        SearchType.NATURAL_LANGUAGE: NaturalLanguageRetriever().get_completion,
-    }
+            save_interaction=save_interaction,
+            last_k=last_k,
+            only_context=True,
+        )
 
-    # If the query type is FEELING_LUCKY, select the search type intelligently
-    if query_type is SearchType.FEELING_LUCKY:
-        query_type = await select_search_type(query)
+        context = {}
+        datasets: List[Dataset] = []
 
-    search_task = search_tasks.get(query_type)
+        for _, search_context, datasets in search_responses:
+            for dataset in datasets:
+                context[str(dataset.id)] = search_context
 
-    if search_task is None:
-        raise UnsupportedSearchTypeError(str(query_type))
+            datasets.extend(datasets)
 
-    send_telemetry("cognee.search EXECUTION STARTED", user.id)
+        specific_search_tools = await get_search_type_tools(
+            query_type=query_type,
+            query_text=query_text,
+            system_prompt_path=system_prompt_path,
+            system_prompt=system_prompt,
+            top_k=top_k,
+            node_type=node_type,
+            node_name=node_name,
+            save_interaction=save_interaction,
+            last_k=last_k,
+        )
+        search_tools = specific_search_tools
+        if len(search_tools) == 2:
+            [get_completion, _] = search_tools
+        else:
+            get_completion = search_tools[0]
 
-    results = await search_task(query)
+        def prepare_combined_context(
+            context,
+        ) -> Union[List[Edge], str]:
+            combined_context = []
 
-    send_telemetry("cognee.search EXECUTION COMPLETED", user.id)
+            for dataset_context in context.values():
+                combined_context += dataset_context
 
-    return results
+            if combined_context and isinstance(combined_context[0], str):
+                return "\n".join(combined_context)
 
+            return combined_context
 
-async def authorized_search(
-    query_text: str,
-    query_type: SearchType,
-    user: User = None,
-    dataset_ids: Optional[list[UUID]] = None,
-    system_prompt_path: str = "answer_simple_question.txt",
-    top_k: int = 10,
-) -> list:
-    """
-    Verifies access for provided datasets or uses all datasets user has read access for and performs search per dataset.
-    Not to be used outside of active access control mode.
-    """
+        combined_context = prepare_combined_context(context)
+        completion = await get_completion(query_text, combined_context)
 
-    query = await log_query(query_text, query_type.value, user.id)
-
-    # Find datasets user has read access for (if datasets are provided only return them. Provided user has read access)
-    search_datasets = await get_specific_user_permission_datasets(user.id, "read", dataset_ids)
+        return completion, combined_context, datasets
 
     # Searches all provided datasets and handles setting up of appropriate database context based on permissions
-    search_results = await specific_search_by_context(
-        search_datasets, query_text, query_type, user, system_prompt_path, top_k
+    search_results = await search_in_datasets_context(
+        search_datasets=search_datasets,
+        query_type=query_type,
+        query_text=query_text,
+        system_prompt_path=system_prompt_path,
+        system_prompt=system_prompt,
+        top_k=top_k,
+        node_type=node_type,
+        node_name=node_name,
+        save_interaction=save_interaction,
+        last_k=last_k,
+        only_context=only_context,
     )
 
-    await log_result(query.id, json.dumps(search_results, cls=JSONEncoder), user.id)
-
     return search_results
 
 
-async def specific_search_by_context(
+async def search_in_datasets_context(
     search_datasets: list[Dataset],
-    query_text: str,
     query_type: SearchType,
-    user: User,
-    system_prompt_path: str,
-    top_k: int,
-):
+    query_text: str,
+    system_prompt_path: str = "answer_simple_question.txt",
+    system_prompt: Optional[str] = None,
+    top_k: int = 10,
+    node_type: Optional[Type] = NodeSet,
+    node_name: Optional[List[str]] = None,
+    save_interaction: bool = False,
+    last_k: Optional[int] = None,
+    only_context: bool = False,
+    context: Optional[Any] = None,
+) -> List[Tuple[Any, Union[str, List[Edge]], List[Dataset]]]:
     """
     Searches all provided datasets and handles setting up of appropriate database context based on permissions.
     Not to be used outside of active access control mode.
     """
 
-    async def _search_by_context(dataset, user, query_type, query_text, system_prompt_path, top_k):
+    async def _search_in_dataset_context(
+        dataset: Dataset,
+        query_type: SearchType,
+        query_text: str,
+        system_prompt_path: str = "answer_simple_question.txt",
+        system_prompt: Optional[str] = None,
+        top_k: int = 10,
+        node_type: Optional[Type] = NodeSet,
+        node_name: Optional[List[str]] = None,
+        save_interaction: bool = False,
+        last_k: Optional[int] = None,
+        only_context: bool = False,
+        context: Optional[Any] = None,
+    ) -> Tuple[Any, Union[str, List[Edge]], List[Dataset]]:
         # Set database configuration in async context for each dataset user has access for
         await set_database_global_context_variables(dataset.id, dataset.owner_id)
-        search_results = await specific_search(
-            query_type, query_text, user, system_prompt_path=system_prompt_path, top_k=top_k
+
+        specific_search_tools = await get_search_type_tools(
+            query_type=query_type,
+            query_text=query_text,
+            system_prompt_path=system_prompt_path,
+            system_prompt=system_prompt,
+            top_k=top_k,
+            node_type=node_type,
+            node_name=node_name,
+            save_interaction=save_interaction,
+            last_k=last_k,
         )
-        return {
-            "search_result": search_results,
-            "dataset_id": dataset.id,
-            "dataset_name": dataset.name,
-        }
+        search_tools = specific_search_tools
+        if len(search_tools) == 2:
+            [get_completion, get_context] = search_tools
+
+            if only_context:
+                return None, await get_context(query_text), [dataset]
+
+            search_context = context or await get_context(query_text)
+            search_result = await get_completion(query_text, search_context)
+
+            return search_result, search_context, [dataset]
+        else:
+            unknown_tool = search_tools[0]
+
+            return await unknown_tool(query_text), "", [dataset]
 
     # Search every dataset async based on query and appropriate database configuration
     tasks = []
     for dataset in search_datasets:
         tasks.append(
-            _search_by_context(dataset, user, query_type, query_text, system_prompt_path, top_k)
+            _search_in_dataset_context(
+                dataset=dataset,
+                query_type=query_type,
+                query_text=query_text,
+                system_prompt_path=system_prompt_path,
+                system_prompt=system_prompt,
+                top_k=top_k,
+                node_type=node_type,
+                node_name=node_name,
+                save_interaction=save_interaction,
+                last_k=last_k,
+                only_context=only_context,
+                context=context,
+            )
         )
 
     return await asyncio.gather(*tasks)

cognee/modules/search/types/SearchResult.py

@@ -0,0 +1,21 @@
+from uuid import UUID
+from pydantic import BaseModel
+from typing import Any, Dict, List, Optional
+
+
+class SearchResultDataset(BaseModel):
+    id: UUID
+    name: str
+
+
+class CombinedSearchResult(BaseModel):
+    result: Optional[Any]
+    context: Dict[str, Any]
+    graphs: Optional[Dict[str, Any]] = {}
+    datasets: Optional[List[SearchResultDataset]] = None
+
+
+class SearchResult(BaseModel):
+    search_result: Any
+    dataset_id: Optional[UUID]
+    dataset_name: Optional[str]

cognee/modules/search/types/SearchType.py

@@ -14,3 +14,6 @@ class SearchType(Enum):
     GRAPH_COMPLETION_COT = "GRAPH_COMPLETION_COT"
     GRAPH_COMPLETION_CONTEXT_EXTENSION = "GRAPH_COMPLETION_CONTEXT_EXTENSION"
     FEELING_LUCKY = "FEELING_LUCKY"
+    FEEDBACK = "FEEDBACK"
+    TEMPORAL = "TEMPORAL"
+    CODING_RULES = "CODING_RULES"

cognee/modules/search/types/__init__.py

@@ -1 +1,2 @@
 from .SearchType import SearchType
+from .SearchResult import SearchResult, SearchResultDataset, CombinedSearchResult

cognee/modules/search/utils/__init__.py

@@ -0,0 +1,2 @@
+from .prepare_search_result import prepare_search_result
+from .transform_context_to_graph import transform_context_to_graph

cognee/modules/search/utils/prepare_search_result.py

@@ -0,0 +1,41 @@
+from typing import List, cast
+
+from cognee.modules.graph.utils import resolve_edges_to_text
+from cognee.modules.graph.cognee_graph.CogneeGraphElements import Edge
+from cognee.modules.search.utils.transform_context_to_graph import transform_context_to_graph
+
+
+async def prepare_search_result(search_result):
+    result, context, datasets = search_result
+
+    graphs = None
+    result_graph = None
+    context_texts = {}
+
+    if isinstance(context, List) and len(context) > 0 and isinstance(context[0], Edge):
+        context_graph = transform_context_to_graph(context)
+
+        graphs = {
+            "*": context_graph,
+        }
+        context_texts = {
+            "*": await resolve_edges_to_text(context),
+        }
+    elif isinstance(context, str):
+        context_texts = {
+            "*": context,
+        }
+    elif isinstance(context, List) and len(context) > 0 and isinstance(context[0], str):
+        context_texts = {
+            "*": "\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)
+
+    return {
+        "result": result_graph or result,
+        "graphs": graphs,
+        "context": context_texts,
+        "datasets": datasets,
+    }

cognee/modules/search/utils/transform_context_to_graph.py

@@ -0,0 +1,38 @@
+from typing import List
+
+from cognee.modules.graph.cognee_graph.CogneeGraphElements import Edge
+
+
+def transform_context_to_graph(context: List[Edge]):
+    nodes = {}
+    edges = {}
+
+    for triplet in context:
+        nodes[triplet.node1.id] = {
+            "id": triplet.node1.id,
+            "label": triplet.node1.attributes["name"]
+            if "name" in triplet.node1.attributes
+            else triplet.node1.id,
+            "type": triplet.node1.attributes["type"],
+            "attributes": triplet.node2.attributes,
+        }
+        nodes[triplet.node2.id] = {
+            "id": triplet.node2.id,
+            "label": triplet.node2.attributes["name"]
+            if "name" in triplet.node2.attributes
+            else triplet.node2.id,
+            "type": triplet.node2.attributes["type"],
+            "attributes": triplet.node2.attributes,
+        }
+        edges[
+            f"{triplet.node1.id}_{triplet.attributes['relationship_name']}_{triplet.node2.id}"
+        ] = {
+            "source": triplet.node1.id,
+            "target": triplet.node2.id,
+            "label": triplet.attributes["relationship_name"],
+        }
+
+    return {
+        "nodes": list(nodes.values()),
+        "edges": list(edges.values()),
+    }

cognee/modules/settings/get_settings.py

@@ -88,8 +88,8 @@ def get_settings() -> SettingsDict:
                 "models": {
                     "openai": [
                         {
-                            "value": "gpt-4o-mini",
-                            "label": "gpt-4o-mini",
+                            "value": "gpt-5-mini",
+                            "label": "gpt-5-mini",
                         },
                         {
                             "value": "gpt-4o",

cognee/modules/sync/__init__.py

@@ -0,0 +1 @@
+# Sync module for tracking sync operations

cognee/modules/sync/methods/__init__.py

@@ -0,0 +1,23 @@
+from .create_sync_operation import create_sync_operation
+from .get_sync_operation import (
+    get_sync_operation,
+    get_user_sync_operations,
+    get_running_sync_operations_for_user,
+)
+from .update_sync_operation import (
+    update_sync_operation,
+    mark_sync_started,
+    mark_sync_completed,
+    mark_sync_failed,
+)
+
+__all__ = [
+    "create_sync_operation",
+    "get_sync_operation",
+    "get_user_sync_operations",
+    "get_running_sync_operations_for_user",
+    "update_sync_operation",
+    "mark_sync_started",
+    "mark_sync_completed",
+    "mark_sync_failed",
+]

cognee/modules/sync/methods/create_sync_operation.py

@@ -0,0 +1,53 @@
+from uuid import UUID
+from typing import Optional, List
+from datetime import datetime, timezone
+from cognee.modules.sync.models import SyncOperation, SyncStatus
+from cognee.infrastructure.databases.relational import get_relational_engine
+
+
+async def create_sync_operation(
+    run_id: str,
+    dataset_ids: List[UUID],
+    dataset_names: List[str],
+    user_id: UUID,
+    total_records_to_sync: Optional[int] = None,
+    total_records_to_download: Optional[int] = None,
+    total_records_to_upload: Optional[int] = None,
+) -> SyncOperation:
+    """
+    Create a new sync operation record in the database.
+
+    Args:
+        run_id: Unique public identifier for this sync operation
+        dataset_ids: List of dataset UUIDs being synced
+        dataset_names: List of dataset names being synced
+        user_id: UUID of the user who initiated the sync
+        total_records_to_sync: Total number of records to sync (if known)
+        total_records_to_download: Total number of records to download (if known)
+        total_records_to_upload: Total number of records to upload (if known)
+
+    Returns:
+        SyncOperation: The created sync operation record
+    """
+    db_engine = get_relational_engine()
+
+    sync_operation = SyncOperation(
+        run_id=run_id,
+        dataset_ids=[
+            str(uuid) for uuid in dataset_ids
+        ],  # Convert UUIDs to strings for JSON storage
+        dataset_names=dataset_names,
+        user_id=user_id,
+        status=SyncStatus.STARTED,
+        total_records_to_sync=total_records_to_sync,
+        total_records_to_download=total_records_to_download,
+        total_records_to_upload=total_records_to_upload,
+        created_at=datetime.now(timezone.utc),
+    )
+
+    async with db_engine.get_async_session() as session:
+        session.add(sync_operation)
+        await session.commit()
+        await session.refresh(sync_operation)
+
+    return sync_operation