import asyncio import json from typing import Any, List, Optional, Tuple, Type, Union from uuid import UUID from cognee import __version__ as cognee_version from cognee.base_config import get_base_config from cognee.context_global_variables import ( backend_access_control_enabled, set_database_global_context_variables, ) from cognee.infrastructure.databases.graph import get_graph_engine from cognee.infrastructure.databases.vector.embeddings.config import EmbeddingConfig from cognee.infrastructure.llm.config import LLMConfig from cognee.modules.data.methods.get_authorized_existing_datasets import ( get_authorized_existing_datasets, ) from cognee.modules.data.models import Dataset from cognee.modules.engine.models.node_set import NodeSet from cognee.modules.graph.cognee_graph.CogneeGraphElements import Edge from cognee.modules.observability import ( COGNEE_SEARCH_QUERY, COGNEE_SEARCH_TYPE, new_span, ) from cognee.modules.search.methods.get_retriever_output import get_retriever_output from cognee.modules.search.models.SearchResultPayload import SearchResultPayload from cognee.modules.search.operations import log_query, log_result from cognee.modules.search.types import ( SearchResult, SearchType, ) from cognee.modules.users.models import User from cognee.shared.logging_utils import get_logger from cognee.shared.utils import send_telemetry logger = get_logger() async def search( query_text: str, query_type: SearchType, dataset_ids: Union[list[UUID], None], user: User, system_prompt_path="answer_simple_question.txt", system_prompt: Optional[str] = None, top_k: int = 15, node_type: Optional[Type] = NodeSet, node_name: Optional[List[str]] = None, node_name_filter_operator: str = "OR", only_context: bool = False, session_id: Optional[str] = None, wide_search_top_k: Optional[int] = 100, triplet_distance_penalty: Optional[float] = 6.5, feedback_influence: float = get_base_config().default_feedback_influence, verbose=False, retriever_specific_config: Optional[dict] = None, neighborhood_depth: Optional[int] = None, neighborhood_seed_top_k: Optional[int] = None, include_references: bool = False, llm_config: Optional[LLMConfig] = None, embedding_config: Optional[EmbeddingConfig] = None, ) -> List[SearchResult]: """ Args: query_text: query_type: datasets: user: system_prompt_path: top_k: Returns: 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, additional_properties={ "cognee_version": cognee_version, "tenant_id": str(user.tenant_id) if user.tenant_id else "Single User Tenant", "include_references": include_references, }, ) with new_span("cognee.search.authorize") as span: span.set_attribute(COGNEE_SEARCH_TYPE, query_type.value) span.set_attribute(COGNEE_SEARCH_QUERY, query_text[:500]) span.set_attribute("cognee.search.top_k", top_k) span.set_attribute( "cognee.search.dataset_count", len(dataset_ids) if dataset_ids else 0, ) 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, node_name_filter_operator=node_name_filter_operator, only_context=only_context, session_id=session_id, wide_search_top_k=wide_search_top_k, triplet_distance_penalty=triplet_distance_penalty, feedback_influence=feedback_influence, retriever_specific_config=retriever_specific_config, neighborhood_depth=neighborhood_depth, neighborhood_seed_top_k=neighborhood_seed_top_k, include_references=include_references, llm_config=llm_config, embedding_config=embedding_config, ) span.set_attribute("cognee.search.result_count", len(search_results)) send_telemetry( "cognee.search EXECUTION COMPLETED", user.id, additional_properties={ "cognee_version": cognee_version, "tenant_id": str(user.tenant_id) if user.tenant_id else "Single User Tenant", }, ) # Log only the completion text (what the user sees), not the full # serialized graph payload. The raw result_objects can be 50-100 KB # each and cause unbounded DB growth in long-running deployments. completions = [] for item in search_results: payload = item[0] if isinstance(item, tuple) else item if hasattr(payload, "completion") and payload.completion: completions.append(payload.completion) elif hasattr(payload, "context") and payload.context: completions.append(payload.context) await log_result( query.id, json.dumps(completions) if completions else "[]", user.id, ) return _backwards_compatible_search_results(search_results, verbose) async def authorized_search( query_type: SearchType, query_text: str, user: User, dataset_ids: Optional[list[UUID]] = None, system_prompt_path: str = "answer_simple_question.txt", system_prompt: Optional[str] = None, top_k: int = 15, node_type: Optional[Type] = NodeSet, node_name: Optional[List[str]] = None, node_name_filter_operator: str = "OR", only_context: bool = False, session_id: Optional[str] = None, wide_search_top_k: Optional[int] = 100, triplet_distance_penalty: Optional[float] = 6.5, feedback_influence: float = get_base_config().default_feedback_influence, retriever_specific_config: Optional[dict] = None, neighborhood_depth: Optional[int] = None, neighborhood_seed_top_k: Optional[int] = None, include_references: bool = False, llm_config: Optional[LLMConfig] = None, embedding_config: Optional[EmbeddingConfig] = None, ) -> List[SearchResultPayload]: """ 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_authorized_existing_datasets( datasets=dataset_ids, permission_type="read", user=user ) # Searches all provided datasets and handles setting up of appropriate database context based on permissions search_results = await search_in_datasets_context( search_datasets=search_datasets, query_type=query_type, query_text=query_text, user=user, system_prompt_path=system_prompt_path, system_prompt=system_prompt, top_k=top_k, node_type=node_type, node_name=node_name, node_name_filter_operator=node_name_filter_operator, only_context=only_context, session_id=session_id, wide_search_top_k=wide_search_top_k, triplet_distance_penalty=triplet_distance_penalty, feedback_influence=feedback_influence, retriever_specific_config=retriever_specific_config, neighborhood_depth=neighborhood_depth, neighborhood_seed_top_k=neighborhood_seed_top_k, include_references=include_references, llm_config=llm_config, embedding_config=embedding_config, ) return search_results async def search_in_datasets_context( search_datasets: list[Dataset], query_type: SearchType, query_text: str, user: User, system_prompt_path: str = "answer_simple_question.txt", system_prompt: Optional[str] = None, top_k: int = 15, node_type: Optional[Type] = NodeSet, node_name: Optional[List[str]] = None, node_name_filter_operator: str = "OR", only_context: bool = False, session_id: Optional[str] = None, wide_search_top_k: Optional[int] = 100, triplet_distance_penalty: Optional[float] = 6.5, feedback_influence: float = get_base_config().default_feedback_influence, retriever_specific_config: Optional[dict] = None, neighborhood_depth: Optional[int] = None, neighborhood_seed_top_k: Optional[int] = None, include_references: bool = False, llm_config: Optional[LLMConfig] = None, embedding_config: Optional[EmbeddingConfig] = 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_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 = 15, node_type: Optional[Type] = NodeSet, node_name: Optional[List[str]] = None, node_name_filter_operator: str = "OR", only_context: bool = False, session_id: Optional[str] = None, wide_search_top_k: Optional[int] = 100, triplet_distance_penalty: Optional[float] = 6.5, feedback_influence: float = get_base_config().default_feedback_influence, retriever_specific_config: Optional[dict] = None, neighborhood_depth: Optional[int] = None, neighborhood_seed_top_k: Optional[int] = None, include_references: bool = False, ) -> SearchResultPayload: with new_span("cognee.search.dataset") as span: span.set_attribute("cognee.search.dataset_name", dataset.name or "") span.set_attribute("cognee.search.dataset_id", str(dataset.id)) async with set_database_global_context_variables( dataset.id, dataset.owner_id, llm_config=llm_config, embedding_config=embedding_config, ): # Check if graph for dataset is empty and log warnings if necessary graph_engine = await get_graph_engine() is_empty = await graph_engine.is_empty() if is_empty: # TODO: we can log here, but not all search types use graph. Still keeping this here for reviewer input from cognee.modules.data.methods import get_dataset_data dataset_data = await get_dataset_data(dataset.id) if len(dataset_data) > 0: logger.warning( f"Dataset '{dataset.name}' has {len(dataset_data)} data item(s) but the knowledge graph is empty. " "Please run cognify to process the data before searching." ) else: logger.warning( f"Search attempt on an empty knowledge graph - no data has been added to this dataset: {dataset.name}" ) span.set_attribute("cognee.search.graph_empty", True) # Get retriever output in the context of the current dataset return await get_retriever_output( query_type=query_type, query_text=query_text, user=user, dataset=dataset, system_prompt_path=system_prompt_path, system_prompt=system_prompt, top_k=top_k, node_type=node_type, node_name=node_name, node_name_filter_operator=node_name_filter_operator, only_context=only_context, session_id=session_id, wide_search_top_k=wide_search_top_k, triplet_distance_penalty=triplet_distance_penalty, feedback_influence=feedback_influence, retriever_specific_config=retriever_specific_config, neighborhood_depth=neighborhood_depth, neighborhood_seed_top_k=neighborhood_seed_top_k, include_references=include_references, ) # Search every dataset async based on query and appropriate database configuration tasks = [] if backend_access_control_enabled(): for dataset in search_datasets: tasks.append( _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, node_name_filter_operator=node_name_filter_operator, only_context=only_context, session_id=session_id, wide_search_top_k=wide_search_top_k, triplet_distance_penalty=triplet_distance_penalty, feedback_influence=feedback_influence, retriever_specific_config=retriever_specific_config, neighborhood_depth=neighborhood_depth, neighborhood_seed_top_k=neighborhood_seed_top_k, include_references=include_references, ) ) else: # Run search without setting database context in case access control is disabled # Needed for low level pipelines that need to run search without dataset context. dataset = search_datasets[0] if len(search_datasets) == 1 else None retriever_kwargs = dict( query_type=query_type, query_text=query_text, user=user, dataset=dataset, system_prompt_path=system_prompt_path, system_prompt=system_prompt, top_k=top_k, node_type=node_type, node_name=node_name, node_name_filter_operator=node_name_filter_operator, only_context=only_context, session_id=session_id, wide_search_top_k=wide_search_top_k, triplet_distance_penalty=triplet_distance_penalty, feedback_influence=feedback_influence, retriever_specific_config=retriever_specific_config, neighborhood_depth=neighborhood_depth, neighborhood_seed_top_k=neighborhood_seed_top_k, include_references=include_references, ) async def _search_without_context() -> SearchResultPayload: # No dataset DB context to set when access control is disabled, but # still forward any per-call LLM/embedding overrides onto the async # context (set_database_global_context_variables applies these even # in single-tenant mode and ignores the dataset argument). async with set_database_global_context_variables( dataset.id if dataset else None, user.id, llm_config=llm_config, embedding_config=embedding_config, ): return await get_retriever_output(**retriever_kwargs) tasks.append(_search_without_context()) return await asyncio.gather(*tasks) def _backwards_compatible_search_results(search_results, verbose: bool): """ Prepares search results in a format compatible with previous versions of the API. """ # This is for maintaining backwards compatibility if backend_access_control_enabled(): return_value = [] for search_result in search_results: # Dataset info needs to be always included search_result_dict = { "dataset_id": search_result.dataset_id, "dataset_name": search_result.dataset_name, "dataset_tenant_id": search_result.dataset_tenant_id, } if verbose: # Include all different types of results only in verbose mode search_result_dict["text_result"] = search_result.completion search_result_dict["context_result"] = search_result.context search_result_dict["objects_result"] = search_result.result_object else: # Result attribute handles returning appropriate result based on set flags and outputs search_result_dict["search_result"] = search_result.result return_value.append(search_result_dict) return return_value else: return_value = [] if verbose: for search_result in search_results: # Include all different types of results only in verbose mode search_result_dict = { "text_result": search_result.completion, "context_result": search_result.context, "objects_result": search_result.result_object, } return_value.append(search_result_dict) return return_value else: for search_result in search_results: # Result attribute handles returning appropriate result based on set flags and outputs return_value.append(search_result.result) # For maintaining backwards compatibility if len(return_value) == 1 and isinstance(return_value[0], list): # If a single element list return the element directly return return_value[0] else: # Otherwise return the list of results return return_value