cognee 0.4.1__py3-none-any.whl → 0.5.0__py3-none-any.whl

cognee/modules/retrieval/utils/brute_force_triplet_search.py

@@ -58,6 +58,8 @@ async def get_memory_fragment(
     properties_to_project: Optional[List[str]] = None,
     node_type: Optional[Type] = None,
     node_name: Optional[List[str]] = None,
+    relevant_ids_to_filter: Optional[List[str]] = None,
+    triplet_distance_penalty: Optional[float] = 3.5,
 ) -> CogneeGraph:
     """Creates and initializes a CogneeGraph memory fragment with optional property projections."""
     if properties_to_project is None:
@@ -71,9 +73,11 @@ async def get_memory_fragment(
         await memory_fragment.project_graph_from_db(
             graph_engine,
             node_properties_to_project=properties_to_project,
-            edge_properties_to_project=["relationship_name"],
+            edge_properties_to_project=["relationship_name", "edge_text"],
             node_type=node_type,
             node_name=node_name,
+            relevant_ids_to_filter=relevant_ids_to_filter,
+            triplet_distance_penalty=triplet_distance_penalty,
         )
 
     except EntityNotFoundError:
@@ -95,6 +99,8 @@ async def brute_force_triplet_search(
     memory_fragment: Optional[CogneeGraph] = None,
     node_type: Optional[Type] = None,
     node_name: Optional[List[str]] = None,
+    wide_search_top_k: Optional[int] = 100,
+    triplet_distance_penalty: Optional[float] = 3.5,
 ) -> List[Edge]:
     """
     Performs a brute force search to retrieve the top triplets from the graph.
@@ -107,6 +113,8 @@ async def brute_force_triplet_search(
         memory_fragment (Optional[CogneeGraph]): Existing memory fragment to reuse.
         node_type: node type to filter
         node_name: node name to filter
+        wide_search_top_k (Optional[int]): Number of initial elements to retrieve from collections
+        triplet_distance_penalty (Optional[float]): Default distance penalty in graph projection
 
     Returns:
         list: The top triplet results.
@@ -116,10 +124,10 @@ async def brute_force_triplet_search(
     if top_k <= 0:
         raise ValueError("top_k must be a positive integer.")
 
-    if memory_fragment is None:
-        memory_fragment = await get_memory_fragment(
-            properties_to_project, node_type=node_type, node_name=node_name
-        )
+    # Setting wide search limit based on the parameters
+    non_global_search = node_name is None
+
+    wide_search_limit = wide_search_top_k if non_global_search else None
 
     if collections is None:
         collections = [
@@ -129,6 +137,9 @@ async def brute_force_triplet_search(
             "DocumentChunk_text",
         ]
 
+    if "EdgeType_relationship_name" not in collections:
+        collections.append("EdgeType_relationship_name")
+
     try:
         vector_engine = get_vector_engine()
     except Exception as e:
@@ -140,7 +151,7 @@ async def brute_force_triplet_search(
     async def search_in_collection(collection_name: str):
         try:
             return await vector_engine.search(
-                collection_name=collection_name, query_vector=query_vector, limit=None
+                collection_name=collection_name, query_vector=query_vector, limit=wide_search_limit
             )
         except CollectionNotFoundError:
             return []
@@ -156,19 +167,40 @@ async def brute_force_triplet_search(
             return []
 
         # Final statistics
-        projection_time = time.time() - start_time
+        vector_collection_search_time = time.time() - start_time
         logger.info(
-            f"Vector collection retrieval completed: Retrieved distances from {sum(1 for res in results if res)} collections in {projection_time:.2f}s"
+            f"Vector collection retrieval completed: Retrieved distances from {sum(1 for res in results if res)} collections in {vector_collection_search_time:.2f}s"
         )
 
         node_distances = {collection: result for collection, result in zip(collections, results)}
 
         edge_distances = node_distances.get("EdgeType_relationship_name", None)
 
+        if wide_search_limit is not None:
+            relevant_ids_to_filter = list(
+                {
+                    str(getattr(scored_node, "id"))
+                    for collection_name, score_collection in node_distances.items()
+                    if collection_name != "EdgeType_relationship_name"
+                    and isinstance(score_collection, (list, tuple))
+                    for scored_node in score_collection
+                    if getattr(scored_node, "id", None)
+                }
+            )
+        else:
+            relevant_ids_to_filter = None
+
+        if memory_fragment is None:
+            memory_fragment = await get_memory_fragment(
+                properties_to_project=properties_to_project,
+                node_type=node_type,
+                node_name=node_name,
+                relevant_ids_to_filter=relevant_ids_to_filter,
+                triplet_distance_penalty=triplet_distance_penalty,
+            )
+
         await memory_fragment.map_vector_distances_to_graph_nodes(node_distances=node_distances)
-        await memory_fragment.map_vector_distances_to_graph_edges(
-            vector_engine=vector_engine, query_vector=query_vector, edge_distances=edge_distances
-        )
+        await memory_fragment.map_vector_distances_to_graph_edges(edge_distances=edge_distances)
 
         results = await memory_fragment.calculate_top_triplet_importances(k=top_k)
 

cognee/modules/retrieval/utils/completion.py

@@ -3,7 +3,7 @@ from cognee.infrastructure.llm.LLMGateway import LLMGateway
 from cognee.infrastructure.llm.prompts import render_prompt, read_query_prompt
 
 
-async def generate_structured_completion(
+async def generate_completion(
     query: str,
     context: str,
     user_prompt_path: str,
@@ -12,7 +12,7 @@ async def generate_structured_completion(
     conversation_history: Optional[str] = None,
     response_model: Type = str,
 ) -> Any:
-    """Generates a structured completion using LLM with given context and prompts."""
+    """Generates a completion using LLM with given context and prompts."""
     args = {"question": query, "context": context}
     user_prompt = render_prompt(user_prompt_path, args)
     system_prompt = system_prompt if system_prompt else read_query_prompt(system_prompt_path)
@@ -28,26 +28,6 @@ async def generate_structured_completion(
     )
 
 
-async def generate_completion(
-    query: str,
-    context: str,
-    user_prompt_path: str,
-    system_prompt_path: str,
-    system_prompt: Optional[str] = None,
-    conversation_history: Optional[str] = None,
-) -> str:
-    """Generates a completion using LLM with given context and prompts."""
-    return await generate_structured_completion(
-        query=query,
-        context=context,
-        user_prompt_path=user_prompt_path,
-        system_prompt_path=system_prompt_path,
-        system_prompt=system_prompt,
-        conversation_history=conversation_history,
-        response_model=str,
-    )
-
-
 async def summarize_text(
     text: str,
     system_prompt_path: str = "summarize_search_results.txt",

cognee/modules/run_custom_pipeline/__init__.py

@@ -0,0 +1 @@
+from .run_custom_pipeline import run_custom_pipeline

cognee/modules/run_custom_pipeline/run_custom_pipeline.py

@@ -0,0 +1,76 @@
+from typing import Union, Optional, List, Type, Any
+from uuid import UUID
+
+from cognee.shared.logging_utils import get_logger
+
+from cognee.modules.pipelines import run_pipeline
+from cognee.modules.pipelines.tasks.task import Task
+from cognee.modules.users.models import User
+from cognee.modules.pipelines.layers.pipeline_execution_mode import get_pipeline_executor
+
+logger = get_logger()
+
+
+async def run_custom_pipeline(
+    tasks: Union[List[Task], List[str]] = None,
+    data: Any = None,
+    dataset: Union[str, UUID] = "main_dataset",
+    user: User = None,
+    vector_db_config: Optional[dict] = None,
+    graph_db_config: Optional[dict] = None,
+    use_pipeline_cache: bool = False,
+    incremental_loading: bool = False,
+    data_per_batch: int = 20,
+    run_in_background: bool = False,
+    pipeline_name: str = "custom_pipeline",
+):
+    """
+    Custom pipeline in Cognee, can work with already built graphs. Data needs to be provided which can be processed
+    with provided tasks.
+
+    Provided tasks and data will be arranged to run the Cognee pipeline and execute graph enrichment/creation.
+
+    This is the core processing step in Cognee that converts raw text and documents
+    into an intelligent knowledge graph. It analyzes content, extracts entities and
+    relationships, and creates semantic connections for enhanced search and reasoning.
+
+    Args:
+        tasks: List of Cognee Tasks to execute.
+        data: The data to ingest. Can be anything when custom extraction and enrichment tasks are used.
+              Data provided here will be forwarded to the first extraction task in the pipeline as input.
+        dataset: Dataset name or dataset uuid to process.
+        user: User context for authentication and data access. Uses default if None.
+        vector_db_config: Custom vector database configuration for embeddings storage.
+        graph_db_config: Custom graph database configuration for relationship storage.
+        use_pipeline_cache: If True, pipelines with the same ID that are currently executing and pipelines with the same ID that were completed won't process data again.
+                        Pipelines ID is created based on the generate_pipeline_id function. Pipeline status can be manually reset with the reset_dataset_pipeline_run_status function.
+        incremental_loading: If True, only new or modified data will be processed to avoid duplication. (Only works if data is used with the Cognee python Data model).
+                            The incremental system stores and compares hashes of processed data in the Data model and skips data with the same content hash.
+        data_per_batch: Number of data items to be processed in parallel.
+        run_in_background: If True, starts processing asynchronously and returns immediately.
+                          If False, waits for completion before returning.
+                          Background mode recommended for large datasets (>100MB).
+                          Use pipeline_run_id from return value to monitor progress.
+    """
+
+    custom_tasks = [
+        *tasks,
+    ]
+
+    # By calling get pipeline executor we get a function that will have the run_pipeline run in the background or a function that we will need to wait for
+    pipeline_executor_func = get_pipeline_executor(run_in_background=run_in_background)
+
+    # Run the run_pipeline in the background or blocking based on executor
+    return await pipeline_executor_func(
+        pipeline=run_pipeline,
+        tasks=custom_tasks,
+        user=user,
+        data=data,
+        datasets=dataset,
+        vector_db_config=vector_db_config,
+        graph_db_config=graph_db_config,
+        use_pipeline_cache=use_pipeline_cache,
+        incremental_loading=incremental_loading,
+        data_per_batch=data_per_batch,
+        pipeline_name=pipeline_name,
+    )

cognee/modules/search/methods/get_search_type_tools.py

@@ -2,6 +2,7 @@ import os
 from typing import Callable, List, Optional, Type
 
 from cognee.modules.engine.models.node_set import NodeSet
+from cognee.modules.retrieval.triplet_retriever import TripletRetriever
 from cognee.modules.search.types import SearchType
 from cognee.modules.search.operations import select_search_type
 from cognee.modules.search.exceptions import UnsupportedSearchTypeError
@@ -22,7 +23,6 @@ from cognee.modules.retrieval.graph_completion_cot_retriever import GraphComplet
 from cognee.modules.retrieval.graph_completion_context_extension_retriever import (
     GraphCompletionContextExtensionRetriever,
 )
-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
 
@@ -37,6 +37,8 @@ async def get_search_type_tools(
     node_name: Optional[List[str]] = None,
     save_interaction: bool = False,
     last_k: Optional[int] = None,
+    wide_search_top_k: Optional[int] = 100,
+    triplet_distance_penalty: Optional[float] = 3.5,
 ) -> list:
     search_tasks: dict[SearchType, List[Callable]] = {
         SearchType.SUMMARIES: [
@@ -59,6 +61,18 @@ async def get_search_type_tools(
                 system_prompt=system_prompt,
             ).get_context,
         ],
+        SearchType.TRIPLET_COMPLETION: [
+            TripletRetriever(
+                system_prompt_path=system_prompt_path,
+                top_k=top_k,
+                system_prompt=system_prompt,
+            ).get_completion,
+            TripletRetriever(
+                system_prompt_path=system_prompt_path,
+                top_k=top_k,
+                system_prompt=system_prompt,
+            ).get_context,
+        ],
         SearchType.GRAPH_COMPLETION: [
             GraphCompletionRetriever(
                 system_prompt_path=system_prompt_path,
@@ -67,6 +81,8 @@ async def get_search_type_tools(
                 node_name=node_name,
                 save_interaction=save_interaction,
                 system_prompt=system_prompt,
+                wide_search_top_k=wide_search_top_k,
+                triplet_distance_penalty=triplet_distance_penalty,
             ).get_completion,
             GraphCompletionRetriever(
                 system_prompt_path=system_prompt_path,
@@ -75,6 +91,8 @@ async def get_search_type_tools(
                 node_name=node_name,
                 save_interaction=save_interaction,
                 system_prompt=system_prompt,
+                wide_search_top_k=wide_search_top_k,
+                triplet_distance_penalty=triplet_distance_penalty,
             ).get_context,
         ],
         SearchType.GRAPH_COMPLETION_COT: [
@@ -85,6 +103,8 @@ async def get_search_type_tools(
                 node_name=node_name,
                 save_interaction=save_interaction,
                 system_prompt=system_prompt,
+                wide_search_top_k=wide_search_top_k,
+                triplet_distance_penalty=triplet_distance_penalty,
             ).get_completion,
             GraphCompletionCotRetriever(
                 system_prompt_path=system_prompt_path,
@@ -93,6 +113,8 @@ async def get_search_type_tools(
                 node_name=node_name,
                 save_interaction=save_interaction,
                 system_prompt=system_prompt,
+                wide_search_top_k=wide_search_top_k,
+                triplet_distance_penalty=triplet_distance_penalty,
             ).get_context,
         ],
         SearchType.GRAPH_COMPLETION_CONTEXT_EXTENSION: [
@@ -103,6 +125,8 @@ async def get_search_type_tools(
                 node_name=node_name,
                 save_interaction=save_interaction,
                 system_prompt=system_prompt,
+                wide_search_top_k=wide_search_top_k,
+                triplet_distance_penalty=triplet_distance_penalty,
             ).get_completion,
             GraphCompletionContextExtensionRetriever(
                 system_prompt_path=system_prompt_path,
@@ -111,6 +135,8 @@ async def get_search_type_tools(
                 node_name=node_name,
                 save_interaction=save_interaction,
                 system_prompt=system_prompt,
+                wide_search_top_k=wide_search_top_k,
+                triplet_distance_penalty=triplet_distance_penalty,
             ).get_context,
         ],
         SearchType.GRAPH_SUMMARY_COMPLETION: [
@@ -121,6 +147,8 @@ async def get_search_type_tools(
                 node_name=node_name,
                 save_interaction=save_interaction,
                 system_prompt=system_prompt,
+                wide_search_top_k=wide_search_top_k,
+                triplet_distance_penalty=triplet_distance_penalty,
             ).get_completion,
             GraphSummaryCompletionRetriever(
                 system_prompt_path=system_prompt_path,
@@ -129,12 +157,10 @@ async def get_search_type_tools(
                 node_name=node_name,
                 save_interaction=save_interaction,
                 system_prompt=system_prompt,
+                wide_search_top_k=wide_search_top_k,
+                triplet_distance_penalty=triplet_distance_penalty,
             ).get_context,
         ],
-        SearchType.CODE: [
-            CodeRetriever(top_k=top_k).get_completion,
-            CodeRetriever(top_k=top_k).get_context,
-        ],
         SearchType.CYPHER: [
             CypherSearchRetriever().get_completion,
             CypherSearchRetriever().get_context,
@@ -145,8 +171,16 @@ async def get_search_type_tools(
         ],
         SearchType.FEEDBACK: [UserQAFeedback(last_k=last_k).add_feedback],
         SearchType.TEMPORAL: [
-            TemporalRetriever(top_k=top_k).get_completion,
-            TemporalRetriever(top_k=top_k).get_context,
+            TemporalRetriever(
+                top_k=top_k,
+                wide_search_top_k=wide_search_top_k,
+                triplet_distance_penalty=triplet_distance_penalty,
+            ).get_completion,
+            TemporalRetriever(
+                top_k=top_k,
+                wide_search_top_k=wide_search_top_k,
+                triplet_distance_penalty=triplet_distance_penalty,
+            ).get_context,
         ],
         SearchType.CHUNKS_LEXICAL: (
             lambda _r=JaccardChunksRetriever(top_k=top_k): [
@@ -169,7 +203,19 @@ async def get_search_type_tools(
     ):
         raise UnsupportedSearchTypeError("Cypher query search types are disabled.")
 
-    search_type_tools = search_tasks.get(query_type)
+    from cognee.modules.retrieval.registered_community_retrievers import (
+        registered_community_retrievers,
+    )
+
+    if query_type in registered_community_retrievers:
+        retriever = registered_community_retrievers[query_type]
+        retriever_instance = retriever(top_k=top_k)
+        search_type_tools = [
+            retriever_instance.get_completion,
+            retriever_instance.get_context,
+        ]
+    else:
+        search_type_tools = search_tasks.get(query_type)
 
     if not search_type_tools:
         raise UnsupportedSearchTypeError(str(query_type))

cognee/modules/search/methods/no_access_control_search.py

@@ -24,6 +24,8 @@ async def no_access_control_search(
     last_k: Optional[int] = None,
     only_context: bool = False,
     session_id: Optional[str] = None,
+    wide_search_top_k: Optional[int] = 100,
+    triplet_distance_penalty: Optional[float] = 3.5,
 ) -> Tuple[Any, Union[str, List[Edge]], List[Dataset]]:
     search_tools = await get_search_type_tools(
         query_type=query_type,
@@ -35,6 +37,8 @@ async def no_access_control_search(
         node_name=node_name,
         save_interaction=save_interaction,
         last_k=last_k,
+        wide_search_top_k=wide_search_top_k,
+        triplet_distance_penalty=triplet_distance_penalty,
     )
     graph_engine = await get_graph_engine()
     is_empty = await graph_engine.is_empty()

cognee/modules/search/methods/search.py

@@ -1,4 +1,3 @@
-import os
 import json
 import asyncio
 from uuid import UUID
@@ -9,6 +8,7 @@ from cognee.infrastructure.databases.graph import get_graph_engine
 from cognee.shared.logging_utils import get_logger
 from cognee.shared.utils import send_telemetry
 from cognee.context_global_variables import set_database_global_context_variables
+from cognee.context_global_variables import backend_access_control_enabled
 
 from cognee.modules.engine.models.node_set import NodeSet
 from cognee.modules.graph.cognee_graph.CogneeGraphElements import Edge
@@ -47,6 +47,8 @@ async def search(
     only_context: bool = False,
     use_combined_context: bool = False,
     session_id: Optional[str] = None,
+    wide_search_top_k: Optional[int] = 100,
+    triplet_distance_penalty: Optional[float] = 3.5,
 ) -> Union[CombinedSearchResult, List[SearchResult]]:
     """
 
@@ -74,7 +76,7 @@ async def search(
     )
 
     # Use search function filtered by permissions if access control is enabled
-    if os.getenv("ENABLE_BACKEND_ACCESS_CONTROL", "false").lower() == "true":
+    if backend_access_control_enabled():
         search_results = await authorized_search(
             query_type=query_type,
             query_text=query_text,
@@ -90,6 +92,8 @@ async def search(
             only_context=only_context,
             use_combined_context=use_combined_context,
             session_id=session_id,
+            wide_search_top_k=wide_search_top_k,
+            triplet_distance_penalty=triplet_distance_penalty,
         )
     else:
         search_results = [
@@ -105,6 +109,8 @@ async def search(
                 last_k=last_k,
                 only_context=only_context,
                 session_id=session_id,
+                wide_search_top_k=wide_search_top_k,
+                triplet_distance_penalty=triplet_distance_penalty,
             )
         ]
 
@@ -156,7 +162,7 @@ async def search(
         )
     else:
         # This is for maintaining backwards compatibility
-        if os.getenv("ENABLE_BACKEND_ACCESS_CONTROL", "false").lower() == "true":
+        if backend_access_control_enabled():
             return_value = []
             for search_result in search_results:
                 prepared_search_results = await prepare_search_result(search_result)
@@ -172,6 +178,7 @@ async def search(
                             "search_result": [context] if context else None,
                             "dataset_id": datasets[0].id,
                             "dataset_name": datasets[0].name,
+                            "dataset_tenant_id": datasets[0].tenant_id,
                             "graphs": graphs,
                         }
                     )
@@ -181,6 +188,7 @@ async def search(
                             "search_result": [result] if result else None,
                             "dataset_id": datasets[0].id,
                             "dataset_name": datasets[0].name,
+                            "dataset_tenant_id": datasets[0].tenant_id,
                             "graphs": graphs,
                         }
                     )
@@ -217,6 +225,8 @@ async def authorized_search(
     only_context: bool = False,
     use_combined_context: bool = False,
     session_id: Optional[str] = None,
+    wide_search_top_k: Optional[int] = 100,
+    triplet_distance_penalty: Optional[float] = 3.5,
 ) -> Union[
     Tuple[Any, Union[List[Edge], str], List[Dataset]],
     List[Tuple[Any, Union[List[Edge], str], List[Dataset]]],
@@ -244,6 +254,8 @@ async def authorized_search(
             last_k=last_k,
             only_context=True,
             session_id=session_id,
+            wide_search_top_k=wide_search_top_k,
+            triplet_distance_penalty=triplet_distance_penalty,
         )
 
         context = {}
@@ -265,6 +277,8 @@ async def authorized_search(
             node_name=node_name,
             save_interaction=save_interaction,
             last_k=last_k,
+            wide_search_top_k=wide_search_top_k,
+            triplet_distance_penalty=triplet_distance_penalty,
         )
         search_tools = specific_search_tools
         if len(search_tools) == 2:
@@ -304,6 +318,7 @@ async def authorized_search(
         last_k=last_k,
         only_context=only_context,
         session_id=session_id,
+        wide_search_top_k=wide_search_top_k,
     )
 
     return search_results
@@ -323,6 +338,8 @@ async def search_in_datasets_context(
     only_context: bool = False,
     context: Optional[Any] = None,
     session_id: Optional[str] = None,
+    wide_search_top_k: Optional[int] = 100,
+    triplet_distance_penalty: Optional[float] = 3.5,
 ) -> List[Tuple[Any, Union[str, List[Edge]], List[Dataset]]]:
     """
     Searches all provided datasets and handles setting up of appropriate database context based on permissions.
@@ -343,6 +360,8 @@ async def search_in_datasets_context(
         only_context: bool = False,
         context: Optional[Any] = None,
         session_id: Optional[str] = None,
+        wide_search_top_k: Optional[int] = 100,
+        triplet_distance_penalty: Optional[float] = 3.5,
     ) -> 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)
@@ -376,6 +395,8 @@ async def search_in_datasets_context(
             node_name=node_name,
             save_interaction=save_interaction,
             last_k=last_k,
+            wide_search_top_k=wide_search_top_k,
+            triplet_distance_penalty=triplet_distance_penalty,
         )
         search_tools = specific_search_tools
         if len(search_tools) == 2:
@@ -411,6 +432,8 @@ async def search_in_datasets_context(
                 only_context=only_context,
                 context=context,
                 session_id=session_id,
+                wide_search_top_k=wide_search_top_k,
+                triplet_distance_penalty=triplet_distance_penalty,
             )
         )
 

cognee/modules/search/types/SearchType.py

@@ -5,9 +5,9 @@ class SearchType(Enum):
     SUMMARIES = "SUMMARIES"
     CHUNKS = "CHUNKS"
     RAG_COMPLETION = "RAG_COMPLETION"
+    TRIPLET_COMPLETION = "TRIPLET_COMPLETION"
     GRAPH_COMPLETION = "GRAPH_COMPLETION"
     GRAPH_SUMMARY_COMPLETION = "GRAPH_SUMMARY_COMPLETION"
-    CODE = "CODE"
     CYPHER = "CYPHER"
     NATURAL_LANGUAGE = "NATURAL_LANGUAGE"
     GRAPH_COMPLETION_COT = "GRAPH_COMPLETION_COT"

cognee/modules/settings/get_settings.py

@@ -16,6 +16,7 @@ class ModelName(Enum):
     anthropic = "anthropic"
     gemini = "gemini"
     mistral = "mistral"
+    bedrock = "bedrock"
 
 
 class LLMConfig(BaseModel):
@@ -77,6 +78,10 @@ def get_settings() -> SettingsDict:
             "value": "mistral",
             "label": "Mistral",
         },
+        {
+            "value": "bedrock",
+            "label": "Bedrock",
+        },
     ]
 
     return SettingsDict.model_validate(
@@ -157,6 +162,20 @@ def get_settings() -> SettingsDict:
                             "label": "Mistral Large 2.1",
                         },
                     ],
+                    "bedrock": [
+                        {
+                            "value": "eu.anthropic.claude-sonnet-4-5-20250929-v1:0",
+                            "label": "Claude 4.5 Sonnet",
+                        },
+                        {
+                            "value": "eu.anthropic.claude-haiku-4-5-20251001-v1:0",
+                            "label": "Claude 4.5 Haiku",
+                        },
+                        {
+                            "value": "eu.amazon.nova-lite-v1:0",
+                            "label": "Amazon Nova Lite",
+                        },
+                    ],
                 },
             },
             vector_db={

cognee/modules/users/methods/create_user.py

@@ -18,7 +18,6 @@ from typing import Optional
 async def create_user(
     email: str,
     password: str,
-    tenant_id: Optional[str] = None,
     is_superuser: bool = False,
     is_active: bool = True,
     is_verified: bool = False,
@@ -30,37 +29,23 @@ async def create_user(
         async with relational_engine.get_async_session() as session:
             async with get_user_db_context(session) as user_db:
                 async with get_user_manager_context(user_db) as user_manager:
-                    if tenant_id:
-                        # Check if the tenant already exists
-                        result = await session.execute(select(Tenant).where(Tenant.id == tenant_id))
-                        tenant = result.scalars().first()
-                        if not tenant:
-                            raise TenantNotFoundError
-
-                        user = await user_manager.create(
-                            UserCreate(
-                                email=email,
-                                password=password,
-                                tenant_id=tenant.id,
-                                is_superuser=is_superuser,
-                                is_active=is_active,
-                                is_verified=is_verified,
-                            )
-                        )
-                    else:
-                        user = await user_manager.create(
-                            UserCreate(
-                                email=email,
-                                password=password,
-                                is_superuser=is_superuser,
-                                is_active=is_active,
-                                is_verified=is_verified,
-                            )
+                    user = await user_manager.create(
+                        UserCreate(
+                            email=email,
+                            password=password,
+                            is_superuser=is_superuser,
+                            is_active=is_active,
+                            is_verified=is_verified,
                         )
+                    )
 
                     if auto_login:
                         await session.refresh(user)
 
+                    # Update tenants and roles information for User object
+                    _ = await user.awaitable_attrs.tenants
+                    _ = await user.awaitable_attrs.roles
+
                     return user
     except UserAlreadyExists as error:
         print(f"User {email} already exists")