cognee 0.5.1__py3-none-any.whl → 0.5.2__py3-none-any.whl

cognee/alembic.ini

@@ -0,0 +1,117 @@
+# A generic, single database configuration.
+
+[alembic]
+# path to migration scripts
+# Use forward slashes (/) also on windows to provide an os agnostic path
+script_location = alembic
+
+# template used to generate migration file names; The default value is %%(rev)s_%%(slug)s
+# Uncomment the line below if you want the files to be prepended with date and time
+# see https://alembic.sqlalchemy.org/en/latest/tutorial.html#editing-the-ini-file
+# for all available tokens
+# file_template = %%(year)d_%%(month).2d_%%(day).2d_%%(hour).2d%%(minute).2d-%%(rev)s_%%(slug)s
+
+# sys.path path, will be prepended to sys.path if present.
+# defaults to the current working directory.
+prepend_sys_path = .
+
+# timezone to use when rendering the date within the migration file
+# as well as the filename.
+# If specified, requires the python>=3.9 or backports.zoneinfo library.
+# Any required deps can installed by adding `alembic[tz]` to the pip requirements
+# string value is passed to ZoneInfo()
+# leave blank for localtime
+# timezone =
+
+# max length of characters to apply to the "slug" field
+# truncate_slug_length = 40
+
+# set to 'true' to run the environment during
+# the 'revision' command, regardless of autogenerate
+# revision_environment = false
+
+# set to 'true' to allow .pyc and .pyo files without
+# a source .py file to be detected as revisions in the
+# versions/ directory
+# sourceless = false
+
+# version location specification; This defaults
+# to alembic/versions.  When using multiple version
+# directories, initial revisions must be specified with --version-path.
+# The path separator used here should be the separator specified by "version_path_separator" below.
+# version_locations = %(here)s/bar:%(here)s/bat:alembic/versions
+
+# version path separator; As mentioned above, this is the character used to split
+# version_locations. The default within new alembic.ini files is "os", which uses os.pathsep.
+# If this key is omitted entirely, it falls back to the legacy behavior of splitting on spaces and/or commas.
+# Valid values for version_path_separator are:
+#
+# version_path_separator = :
+# version_path_separator = ;
+# version_path_separator = space
+# version_path_separator = newline
+version_path_separator = os  # Use os.pathsep. Default configuration used for new projects.
+
+# set to 'true' to search source files recursively
+# in each "version_locations" directory
+# new in Alembic version 1.10
+# recursive_version_locations = false
+
+# the output encoding used when revision files
+# are written from script.py.mako
+# output_encoding = utf-8
+
+sqlalchemy.url = %(SQLALCHEMY_DATABASE_URI)s
+
+
+[post_write_hooks]
+# post_write_hooks defines scripts or Python functions that are run
+# on newly generated revision scripts.  See the documentation for further
+# detail and examples
+
+# format using "black" - use the console_scripts runner, against the "black" entrypoint
+# hooks = black
+# black.type = console_scripts
+# black.entrypoint = black
+# black.options = -l 79 REVISION_SCRIPT_FILENAME
+
+# lint with attempts to fix using "ruff" - use the exec runner, execute a binary
+# hooks = ruff
+# ruff.type = exec
+# ruff.executable = %(here)s/.venv/bin/ruff
+# ruff.options = --fix REVISION_SCRIPT_FILENAME
+
+# Logging configuration
+[loggers]
+keys = root,sqlalchemy,alembic
+
+[handlers]
+keys = console
+
+[formatters]
+keys = generic
+
+[logger_root]
+level = WARN
+handlers = console
+qualname =
+
+[logger_sqlalchemy]
+level = WARN
+handlers =
+qualname = sqlalchemy.engine
+
+[logger_alembic]
+level = WARN
+handlers =
+qualname = alembic
+
+[handler_console]
+class = StreamHandler
+args = (sys.stderr,)
+level = NOTSET
+formatter = generic
+
+[formatter_generic]
+format = %(levelname)-5.5s [%(name)s] %(message)s
+datefmt = %H:%M:%S

cognee/api/v1/add/add.py

@@ -10,13 +10,14 @@ from cognee.modules.pipelines.layers.reset_dataset_pipeline_run_status import (
 )
 from cognee.modules.engine.operations.setup import setup
 from cognee.tasks.ingestion import ingest_data, resolve_data_directories
+from cognee.tasks.ingestion.data_item import DataItem
 from cognee.shared.logging_utils import get_logger
 
 logger = get_logger()
 
 
 async def add(
-    data: Union[BinaryIO, list[BinaryIO], str, list[str]],
+    data: Union[BinaryIO, list[BinaryIO], str, list[str], DataItem, list[DataItem]],
     dataset_name: str = "main_dataset",
     user: User = None,
     node_set: Optional[List[str]] = None,

cognee/api/v1/add/routers/get_add_router.py

@@ -10,6 +10,7 @@ from cognee.modules.users.methods import get_authenticated_user
 from cognee.shared.utils import send_telemetry
 from cognee.modules.pipelines.models import PipelineRunErrored
 from cognee.shared.logging_utils import get_logger
+from cognee.shared.usage_logger import log_usage
 from cognee import __version__ as cognee_version
 
 logger = get_logger()
@@ -19,6 +20,7 @@ def get_add_router() -> APIRouter:
     router = APIRouter()
 
     @router.post("", response_model=dict)
+    @log_usage(function_name="POST /v1/add", log_type="api_endpoint")
     async def add(
         data: List[UploadFile] = File(default=None),
         datasetName: Optional[str] = Form(default=None),

cognee/api/v1/cognify/cognify.py

@@ -252,7 +252,7 @@ async def get_default_tasks(  # TODO: Find out a better way to do this (Boris's
     chunk_size: int = None,
     config: Config = None,
     custom_prompt: Optional[str] = None,
-    chunks_per_batch: int = 100,
+    chunks_per_batch: int = None,
     **kwargs,
 ) -> list[Task]:
     if config is None:
@@ -272,12 +272,14 @@ async def get_default_tasks(  # TODO: Find out a better way to do this (Boris's
                 "ontology_config": {"ontology_resolver": get_default_ontology_resolver()}
             }
 
-    if chunks_per_batch is None:
-        chunks_per_batch = 100
-
     cognify_config = get_cognify_config()
     embed_triplets = cognify_config.triplet_embedding
 
+    if chunks_per_batch is None:
+        chunks_per_batch = (
+            cognify_config.chunks_per_batch if cognify_config.chunks_per_batch is not None else 100
+        )
+
     default_tasks = [
         Task(classify_documents),
         Task(
@@ -308,7 +310,7 @@ async def get_default_tasks(  # TODO: Find out a better way to do this (Boris's
 
 
 async def get_temporal_tasks(
-    user: User = None, chunker=TextChunker, chunk_size: int = None, chunks_per_batch: int = 10
+    user: User = None, chunker=TextChunker, chunk_size: int = None, chunks_per_batch: int = None
 ) -> list[Task]:
     """
     Builds and returns a list of temporal processing tasks to be executed in sequence.
@@ -330,7 +332,10 @@ async def get_temporal_tasks(
         list[Task]: A list of Task objects representing the temporal processing pipeline.
     """
     if chunks_per_batch is None:
-        chunks_per_batch = 10
+        from cognee.modules.cognify.config import get_cognify_config
+
+        configured = get_cognify_config().chunks_per_batch
+        chunks_per_batch = configured if configured is not None else 10
 
     temporal_tasks = [
         Task(classify_documents),

cognee/api/v1/cognify/routers/get_cognify_router.py

@@ -29,6 +29,7 @@ from cognee.modules.pipelines.queues.pipeline_run_info_queues import (
 )
 from cognee.shared.logging_utils import get_logger
 from cognee.shared.utils import send_telemetry
+from cognee.shared.usage_logger import log_usage
 from cognee import __version__ as cognee_version
 
 logger = get_logger("api.cognify")
@@ -46,12 +47,18 @@ class CognifyPayloadDTO(InDTO):
         examples=[[]],
         description="Reference to one or more previously uploaded ontologies",
     )
+    chunks_per_batch: Optional[int] = Field(
+        default=None,
+        description="Number of chunks to process per task batch in Cognify (overrides default).",
+        examples=[10, 20, 50, 100],
+    )
 
 
 def get_cognify_router() -> APIRouter:
     router = APIRouter()
 
     @router.post("", response_model=dict)
+    @log_usage(function_name="POST /v1/cognify", log_type="api_endpoint")
     async def cognify(payload: CognifyPayloadDTO, user: User = Depends(get_authenticated_user)):
         """
         Transform datasets into structured knowledge graphs through cognitive processing.
@@ -146,6 +153,7 @@ def get_cognify_router() -> APIRouter:
                 config=config_to_use,
                 run_in_background=payload.run_in_background,
                 custom_prompt=payload.custom_prompt,
+                chunks_per_batch=payload.chunks_per_batch,
             )
 
             # If any cognify run errored return JSONResponse with proper error status code

cognee/api/v1/config/config.py

@@ -10,6 +10,7 @@ from cognee.infrastructure.llm.config import (
     get_llm_config,
 )
 from cognee.infrastructure.databases.relational import get_relational_config, get_migration_config
+from cognee.tasks.translation.config import get_translation_config
 from cognee.api.v1.exceptions.exceptions import InvalidConfigAttributeError
 
 
@@ -176,3 +177,62 @@ class config:
     def set_vector_db_url(db_url: str):
         vector_db_config = get_vectordb_config()
         vector_db_config.vector_db_url = db_url
+
+    # Translation configuration methods
+
+    @staticmethod
+    def set_translation_provider(provider: str):
+        """Set the translation provider (llm, google, azure)."""
+        translation_config = get_translation_config()
+        translation_config.translation_provider = provider
+
+    @staticmethod
+    def set_translation_target_language(target_language: str):
+        """Set the default target language for translations."""
+        translation_config = get_translation_config()
+        translation_config.target_language = target_language
+
+    @staticmethod
+    def set_translation_config(config_dict: dict):
+        """
+        Updates the translation config with values from config_dict.
+        """
+        translation_config = get_translation_config()
+        for key, value in config_dict.items():
+            if hasattr(translation_config, key):
+                object.__setattr__(translation_config, key, value)
+            else:
+                raise InvalidConfigAttributeError(attribute=key)
+
+    def set(key: str, value):
+        """
+        Generic setter that maps configuration keys to their specific setter methods.
+        This enables CLI commands like 'cognee config set llm_api_key <value>'.
+        """
+        # Map configuration keys to their setter methods
+        setter_mapping = {
+            "llm_provider": "set_llm_provider",
+            "llm_model": "set_llm_model",
+            "llm_api_key": "set_llm_api_key",
+            "llm_endpoint": "set_llm_endpoint",
+            "graph_database_provider": "set_graph_database_provider",
+            "vector_db_provider": "set_vector_db_provider",
+            "vector_db_url": "set_vector_db_url",
+            "vector_db_key": "set_vector_db_key",
+            "chunk_size": "set_chunk_size",
+            "chunk_overlap": "set_chunk_overlap",
+            "chunk_strategy": "set_chunk_strategy",
+            "chunk_engine": "set_chunk_engine",
+            "classification_model": "set_classification_model",
+            "summarization_model": "set_summarization_model",
+            "graph_model": "set_graph_model",
+            "system_root_directory": "system_root_directory",
+            "data_root_directory": "data_root_directory",
+        }
+
+        if key not in setter_mapping:
+            raise InvalidConfigAttributeError(attribute=key)
+
+        method_name = setter_mapping[key]
+        method = getattr(config, method_name)
+        method(value)

cognee/api/v1/datasets/routers/get_datasets_router.py

@@ -7,7 +7,9 @@ from fastapi import status
 from fastapi import APIRouter
 from fastapi.encoders import jsonable_encoder
 from fastapi import HTTPException, Query, Depends
-from fastapi.responses import JSONResponse, FileResponse
+from fastapi.responses import JSONResponse, FileResponse, StreamingResponse, Response
+from urllib.parse import urlparse
+from pathlib import Path
 
 from cognee.api.DTO import InDTO, OutDTO
 from cognee.infrastructure.databases.relational import get_relational_engine
@@ -44,6 +46,7 @@ class DatasetDTO(OutDTO):
 class DataDTO(OutDTO):
     id: UUID
     name: str
+    label: Optional[str] = None
     created_at: datetime
     updated_at: Optional[datetime] = None
     extension: str
@@ -414,7 +417,7 @@ def get_datasets_router() -> APIRouter:
     @router.get("/{dataset_id}/data/{data_id}/raw", response_class=FileResponse)
     async def get_raw_data(
         dataset_id: UUID, data_id: UUID, user: User = Depends(get_authenticated_user)
-    ):
+    ) -> Response:
         """
         Download the raw data file for a specific data item.
 
@@ -475,6 +478,46 @@ def get_datasets_router() -> APIRouter:
                 message=f"Data ({data_id}) not found in dataset ({dataset_id})."
             )
 
-        return data.raw_data_location
+        raw_location = data.raw_data_location
+        parsed_uri = urlparse(raw_location)
+
+        if parsed_uri.scheme == "s3":
+            from cognee.infrastructure.files.utils.open_data_file import open_data_file
+            from cognee.infrastructure.utils.run_async import run_async
+
+            download_name = Path(parsed_uri.path).name or data.name
+            media_type = data.mime_type or "application/octet-stream"
+
+            async def file_iterator(chunk_size: int = 1024 * 1024):
+                async with open_data_file(raw_location, mode="rb") as file:
+                    while True:
+                        chunk = await run_async(file.read, chunk_size)
+                        if not chunk:
+                            break
+                        yield chunk
+
+            return StreamingResponse(
+                file_iterator(),
+                media_type=media_type,
+                headers={"Content-Disposition": f'attachment; filename="{download_name}"'},
+            )
+
+        if parsed_uri.scheme in ("file", "") or (
+            len(parsed_uri.scheme) == 1 and parsed_uri.scheme.isalpha()
+        ):
+            from cognee.infrastructure.files.utils.get_data_file_path import get_data_file_path
+
+            file_path = get_data_file_path(raw_location)
+            path = Path(file_path)
+
+            if not path.is_file():
+                raise DataNotFoundError(message=f"Raw file not found on disk for data ({data_id}).")
+
+            return FileResponse(path=path)
+
+        raise HTTPException(
+            status_code=status.HTTP_501_NOT_IMPLEMENTED,
+            detail=f"Storage scheme '{parsed_uri.scheme}' not supported for direct download.",
+        )
 
     return router

cognee/api/v1/memify/routers/get_memify_router.py

@@ -12,6 +12,7 @@ from cognee.modules.users.methods import get_authenticated_user
 from cognee.shared.utils import send_telemetry
 from cognee.modules.pipelines.models import PipelineRunErrored
 from cognee.shared.logging_utils import get_logger
+from cognee.shared.usage_logger import log_usage
 from cognee import __version__ as cognee_version
 
 logger = get_logger()
@@ -35,6 +36,7 @@ def get_memify_router() -> APIRouter:
     router = APIRouter()
 
     @router.post("", response_model=dict)
+    @log_usage(function_name="POST /v1/memify", log_type="api_endpoint")
     async def memify(payload: MemifyPayloadDTO, user: User = Depends(get_authenticated_user)):
         """
         Enrichment pipeline in Cognee, can work with already built graphs. If no data is provided existing knowledge graph will be used as data,
@@ -90,6 +92,7 @@ def get_memify_router() -> APIRouter:
                 dataset=payload.dataset_id if payload.dataset_id else payload.dataset_name,
                 node_name=payload.node_name,
                 user=user,
+                run_in_background=payload.run_in_background,
             )
 
             if isinstance(memify_run, PipelineRunErrored):

cognee/api/v1/search/routers/get_search_router.py

@@ -6,14 +6,17 @@ from fastapi import Depends, APIRouter
 from fastapi.responses import JSONResponse
 from fastapi.encoders import jsonable_encoder
 
-from cognee.modules.search.types import SearchType, SearchResult, CombinedSearchResult
+from cognee.modules.search.types import SearchType, SearchResult
 from cognee.api.DTO import InDTO, OutDTO
-from cognee.modules.users.exceptions.exceptions import PermissionDeniedError
+from cognee.modules.users.exceptions.exceptions import PermissionDeniedError, UserNotFoundError
 from cognee.modules.users.models import User
 from cognee.modules.search.operations import get_history
 from cognee.modules.users.methods import get_authenticated_user
 from cognee.shared.utils import send_telemetry
+from cognee.shared.usage_logger import log_usage
 from cognee import __version__ as cognee_version
+from cognee.infrastructure.databases.exceptions import DatabaseNotCreatedError
+from cognee.exceptions import CogneeValidationError
 
 
 # Note: Datasets sent by name will only map to datasets owned by the request sender
@@ -29,7 +32,7 @@ class SearchPayloadDTO(InDTO):
     node_name: Optional[list[str]] = Field(default=None, example=[])
     top_k: Optional[int] = Field(default=10)
     only_context: bool = Field(default=False)
-    use_combined_context: bool = Field(default=False)
+    verbose: bool = Field(default=False)
 
 
 def get_search_router() -> APIRouter:
@@ -72,7 +75,8 @@ def get_search_router() -> APIRouter:
         except Exception as error:
             return JSONResponse(status_code=500, content={"error": str(error)})
 
-    @router.post("", response_model=Union[List[SearchResult], CombinedSearchResult, List])
+    @router.post("", response_model=Union[List[SearchResult], List])
+    @log_usage(function_name="POST /v1/search", log_type="api_endpoint")
     async def search(payload: SearchPayloadDTO, user: User = Depends(get_authenticated_user)):
         """
         Search for nodes in the graph database.
@@ -116,7 +120,7 @@ def get_search_router() -> APIRouter:
                 "node_name": payload.node_name,
                 "top_k": payload.top_k,
                 "only_context": payload.only_context,
-                "use_combined_context": payload.use_combined_context,
+                "verbose": payload.verbose,
                 "cognee_version": cognee_version,
             },
         )
@@ -133,11 +137,22 @@ def get_search_router() -> APIRouter:
                 system_prompt=payload.system_prompt,
                 node_name=payload.node_name,
                 top_k=payload.top_k,
+                verbose=payload.verbose,
                 only_context=payload.only_context,
-                use_combined_context=payload.use_combined_context,
             )
 
             return jsonable_encoder(results)
+        except (DatabaseNotCreatedError, UserNotFoundError, CogneeValidationError) as e:
+            # Return a clear 422 with actionable guidance instead of leaking a stacktrace
+            status_code = getattr(e, "status_code", 422)
+            return JSONResponse(
+                status_code=status_code,
+                content={
+                    "error": "Search prerequisites not met",
+                    "detail": str(e),
+                    "hint": "Run `await cognee.add(...)` then `await cognee.cognify()` before searching.",
+                },
+            )
         except PermissionDeniedError:
             return []
         except Exception as error:

cognee/api/v1/search/search.py

@@ -4,13 +4,16 @@ from typing import Union, Optional, List, Type
 from cognee.infrastructure.databases.graph import get_graph_engine
 from cognee.modules.engine.models.node_set import NodeSet
 from cognee.modules.users.models import User
-from cognee.modules.search.types import SearchResult, SearchType, CombinedSearchResult
+from cognee.modules.search.types import SearchResult, SearchType
 from cognee.modules.users.methods import get_default_user
 from cognee.modules.search.methods import search as search_function
 from cognee.modules.data.methods import get_authorized_existing_datasets
 from cognee.modules.data.exceptions import DatasetNotFoundError
 from cognee.context_global_variables import set_session_user_context_variable
 from cognee.shared.logging_utils import get_logger
+from cognee.infrastructure.databases.exceptions import DatabaseNotCreatedError
+from cognee.exceptions import CogneeValidationError
+from cognee.modules.users.exceptions.exceptions import UserNotFoundError
 
 logger = get_logger()
 
@@ -29,12 +32,12 @@ async def search(
     save_interaction: bool = False,
     last_k: Optional[int] = 1,
     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,
     verbose: bool = False,
-) -> Union[List[SearchResult], CombinedSearchResult]:
+    retriever_specific_config: Optional[dict] = None,
+) -> List[SearchResult]:
     """
     Search and query the knowledge graph for insights, information, and connections.
 
@@ -126,6 +129,8 @@ async def search(
 
         verbose: If True, returns detailed result information including graph representation (when possible).
 
+        retriever_specific_config: Optional dictionary of additional configuration parameters specific to the retriever being used.
+
     Returns:
         list: Search results in format determined by query_type:
 
@@ -179,7 +184,18 @@ async def search(
         datasets = [datasets]
 
     if user is None:
-        user = await get_default_user()
+        try:
+            user = await get_default_user()
+        except (DatabaseNotCreatedError, UserNotFoundError) as error:
+            # Provide a clear, actionable message instead of surfacing low-level stacktraces
+            raise CogneeValidationError(
+                message=(
+                    "Search prerequisites not met: no database/default user found. "
+                    "Initialize Cognee before searching by:\n"
+                    "• running `await cognee.add(...)` followed by `await cognee.cognify()`."
+                ),
+                name="SearchPreconditionError",
+            ) from error
 
     await set_session_user_context_variable(user)
 
@@ -203,11 +219,11 @@ async def search(
         save_interaction=save_interaction,
         last_k=last_k,
         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,
         verbose=verbose,
+        retriever_specific_config=retriever_specific_config,
     )
 
     return filtered_search_results

cognee/api/v1/sync/routers/get_sync_router.py

@@ -71,7 +71,7 @@ def get_sync_router() -> APIRouter:
           -H "Content-Type: application/json" \\
           -H "Cookie: auth_token=your-token" \\
           -d '{"dataset_ids": ["123e4567-e89b-12d3-a456-426614174000", "456e7890-e12b-34c5-d678-901234567000"]}'
-        
+
         # Sync all user datasets (empty request body or null dataset_ids)
         curl -X POST "http://localhost:8000/api/v1/sync" \\
           -H "Content-Type: application/json" \\
@@ -88,7 +88,7 @@ def get_sync_router() -> APIRouter:
         - **413 Payload Too Large**: Dataset too large for current cloud plan
         - **429 Too Many Requests**: Rate limit exceeded
 
-        ## Notes  
+        ## Notes
         - Sync operations run in the background - you get an immediate response
         - Use the returned run_id to track progress (status API coming soon)
         - Large datasets are automatically chunked for efficient transfer
@@ -179,7 +179,7 @@ def get_sync_router() -> APIRouter:
         ```
 
         ## Example Responses
-        
+
         **No running syncs:**
         ```json
         {

cognee/cli/commands/add_command.py

@@ -21,7 +21,7 @@ binary streams, then stores them in a specified dataset for further processing.
 
 Supported Input Types:
 - **Text strings**: Direct text content
-- **File paths**: Local file paths (absolute paths starting with "/")  
+- **File paths**: Local file paths (absolute paths starting with "/")
 - **File URLs**: "file:///absolute/path" or "file://relative/path"
 - **S3 paths**: "s3://bucket-name/path/to/file"
 - **Lists**: Multiple files or text strings in a single call

cognee/cli/commands/cognify_command.py

@@ -62,6 +62,11 @@ After successful cognify processing, use `cognee search` to query the knowledge
         parser.add_argument(
             "--verbose", "-v", action="store_true", help="Show detailed progress information"
         )
+        parser.add_argument(
+            "--chunks-per-batch",
+            type=int,
+            help="Number of chunks to process per task batch (try 50 for large single documents).",
+        )
 
     def execute(self, args: argparse.Namespace) -> None:
         try:
@@ -111,6 +116,7 @@ After successful cognify processing, use `cognee search` to query the knowledge
                         chunk_size=args.chunk_size,
                         ontology_file_path=args.ontology_file,
                         run_in_background=args.background,
+                        chunks_per_batch=getattr(args, "chunks_per_batch", None),
                     )
                     return result
                 except Exception as e:

cognee/cli/commands/config_command.py

@@ -17,7 +17,7 @@ The `cognee config` command allows you to view and modify configuration settings
 
 You can:
 - View all current configuration settings
-- Get specific configuration values  
+- Get specific configuration values
 - Set configuration values
 - Unset (reset to default) specific configuration values
 - Reset all configuration to defaults

cognee/context_global_variables.py

@@ -121,13 +121,17 @@ async def set_database_global_context_variables(dataset: Union[str, UUID], user_
     )
 
     # Set vector and graph database configuration based on dataset database information
-    # TODO: Add better handling of vector and graph config accross Cognee.
+    # TODO: Add better handling of vector and graph config across Cognee.
     #  LRU_CACHE takes into account order of inputs, if order of inputs is changed it will be registered as a new DB adapter
     vector_config = {
         "vector_db_provider": dataset_database.vector_database_provider,
         "vector_db_url": dataset_database.vector_database_url,
         "vector_db_key": dataset_database.vector_database_key,
         "vector_db_name": dataset_database.vector_database_name,
+        "vector_db_port": dataset_database.vector_database_connection_info.get("port", ""),
+        "vector_db_host": dataset_database.vector_database_connection_info.get("host", ""),
+        "vector_db_username": dataset_database.vector_database_connection_info.get("username", ""),
+        "vector_db_password": dataset_database.vector_database_connection_info.get("password", ""),
     }
 
     graph_config = {

cognee/eval_framework/answer_generation/answer_generation_executor.py

@@ -32,16 +32,15 @@ class AnswerGeneratorExecutor:
             query_text = instance["question"]
             correct_answer = instance["answer"]
 
-            retrieval_context = await retriever.get_context(query_text)
-            search_results = await retriever.get_completion(query_text, retrieval_context)
+            retrieved_objects = await retriever.get_retrieved_objects(query_text)
+            retrieval_context = await retriever.get_context_from_objects(
+                query_text, retrieved_objects
+            )
+            search_results = await retriever.get_completion_from_context(
+                query_text, retrieved_objects, retrieval_context
+            )
 
             ############
-            #:TODO This is a quick fix until we don't structure retriever results properly but lets not leave it like this...this is needed now due to the changed combined retriever structure..
-            if isinstance(retrieval_context, list):
-                retrieval_context = await retriever.convert_retrieved_objects_to_context(
-                    triplets=retrieval_context
-                )
-
             if isinstance(search_results, str):
                 search_results = [search_results]
             #############