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

cognee/infrastructure/databases/graph/neo4j_driver/adapter.py

@@ -8,7 +8,7 @@ from neo4j import AsyncSession
 from neo4j import AsyncGraphDatabase
 from neo4j.exceptions import Neo4jError
 from contextlib import asynccontextmanager
-from typing import Optional, Any, List, Dict, Type, Tuple
+from typing import Optional, Any, List, Dict, Type, Tuple, Coroutine
 
 from cognee.infrastructure.engine import DataPoint
 from cognee.modules.engine.utils.generate_timestamp_datapoint import date_to_int
@@ -964,6 +964,63 @@ class Neo4jAdapter(GraphDBInterface):
             logger.error(f"Error during graph data retrieval: {str(e)}")
             raise
 
+    async def get_id_filtered_graph_data(self, target_ids: list[str]):
+        """
+        Retrieve graph data filtered by specific node IDs, including their direct neighbors
+        and only edges where one endpoint matches those IDs.
+
+        This version uses a single Cypher query for efficiency.
+        """
+        import time
+
+        start_time = time.time()
+
+        try:
+            if not target_ids:
+                logger.warning("No target IDs provided for ID-filtered graph retrieval.")
+                return [], []
+
+            query = """
+            MATCH ()-[r]-()
+            WHERE startNode(r).id IN $target_ids
+               OR endNode(r).id IN $target_ids
+            WITH DISTINCT r, startNode(r) AS a, endNode(r) AS b
+            RETURN
+                properties(a) AS n_properties,
+                properties(b) AS m_properties,
+                type(r) AS type,
+                properties(r) AS properties
+            """
+
+            result = await self.query(query, {"target_ids": target_ids})
+
+            nodes_dict = {}
+            edges = []
+
+            for record in result:
+                n_props = record["n_properties"]
+                m_props = record["m_properties"]
+                r_props = record["properties"]
+                r_type = record["type"]
+
+                nodes_dict[n_props["id"]] = (n_props["id"], n_props)
+                nodes_dict[m_props["id"]] = (m_props["id"], m_props)
+
+                source_id = r_props.get("source_node_id", n_props["id"])
+                target_id = r_props.get("target_node_id", m_props["id"])
+                edges.append((source_id, target_id, r_type, r_props))
+
+            retrieval_time = time.time() - start_time
+            logger.info(
+                f"ID-filtered retrieval: {len(nodes_dict)} nodes and {len(edges)} edges in {retrieval_time:.2f}s"
+            )
+
+            return list(nodes_dict.values()), edges
+
+        except Exception as e:
+            logger.error(f"Error during ID-filtered graph data retrieval: {str(e)}")
+            raise
+
     async def get_nodeset_subgraph(
         self, node_type: Type[Any], node_name: List[str]
     ) -> Tuple[List[Tuple[int, dict]], List[Tuple[int, int, str, dict]]]:
@@ -1470,3 +1527,25 @@ class Neo4jAdapter(GraphDBInterface):
         time_ids_list = [item["id"] for item in time_nodes if "id" in item]
 
         return ", ".join(f"'{uid}'" for uid in time_ids_list)
+
+    async def get_triplets_batch(self, offset: int, limit: int) -> list[dict[str, Any]]:
+        """
+        Retrieve a batch of triplets (start_node, relationship, end_node) from the graph.
+
+        Parameters:
+        -----------
+            - offset (int): Number of triplets to skip before returning results.
+            - limit (int): Maximum number of triplets to return.
+
+        Returns:
+        --------
+            - list[dict[str, Any]]: A list of triplets.
+        """
+        query = f"""
+        MATCH (start_node:`{BASE_LABEL}`)-[relationship]->(end_node:`{BASE_LABEL}`)
+        RETURN start_node, properties(relationship) AS relationship_properties, end_node
+        SKIP $offset LIMIT $limit
+        """
+        results = await self.query(query, {"offset": offset, "limit": limit})
+
+        return results

cognee/infrastructure/databases/hybrid/neptune_analytics/NeptuneAnalyticsAdapter.py

@@ -416,6 +416,15 @@ class NeptuneAnalyticsAdapter(NeptuneGraphDB, VectorDBInterface):
         self._client.query(f"MATCH (n :{self._VECTOR_NODE_LABEL}) DETACH DELETE n")
         pass
 
+    async def is_empty(self) -> bool:
+        query = """
+        MATCH (n)
+        RETURN true
+        LIMIT 1;
+        """
+        query_result = await self._client.query(query)
+        return len(query_result) == 0
+
     @staticmethod
     def _get_scored_result(
         item: dict, with_vector: bool = False, with_score: bool = False

cognee/infrastructure/databases/utils/__init__.py

@@ -1 +1,4 @@
 from .get_or_create_dataset_database import get_or_create_dataset_database
+from .resolve_dataset_database_connection_info import resolve_dataset_database_connection_info
+from .get_graph_dataset_database_handler import get_graph_dataset_database_handler
+from .get_vector_dataset_database_handler import get_vector_dataset_database_handler

cognee/infrastructure/databases/utils/get_graph_dataset_database_handler.py

@@ -0,0 +1,10 @@
+from cognee.modules.users.models.DatasetDatabase import DatasetDatabase
+
+
+def get_graph_dataset_database_handler(dataset_database: DatasetDatabase) -> dict:
+    from cognee.infrastructure.databases.dataset_database_handler.supported_dataset_database_handlers import (
+        supported_dataset_database_handlers,
+    )
+
+    handler = supported_dataset_database_handlers[dataset_database.graph_dataset_database_handler]
+    return handler

cognee/infrastructure/databases/utils/get_or_create_dataset_database.py

@@ -1,16 +1,65 @@
 from uuid import UUID
-from typing import Union
+from typing import Union, Optional
 
 from sqlalchemy import select
 from sqlalchemy.exc import IntegrityError
-from cognee.modules.data.methods import create_dataset
 
+from cognee.modules.data.methods import create_dataset
 from cognee.infrastructure.databases.relational import get_relational_engine
+from cognee.infrastructure.databases.vector import get_vectordb_config
+from cognee.infrastructure.databases.graph.config import get_graph_config
 from cognee.modules.data.methods import get_unique_dataset_id
 from cognee.modules.users.models import DatasetDatabase
 from cognee.modules.users.models import User
 
 
+async def _get_vector_db_info(dataset_id: UUID, user: User) -> dict:
+    vector_config = get_vectordb_config()
+
+    from cognee.infrastructure.databases.dataset_database_handler.supported_dataset_database_handlers import (
+        supported_dataset_database_handlers,
+    )
+
+    handler = supported_dataset_database_handlers[vector_config.vector_dataset_database_handler]
+    return await handler["handler_instance"].create_dataset(dataset_id, user)
+
+
+async def _get_graph_db_info(dataset_id: UUID, user: User) -> dict:
+    graph_config = get_graph_config()
+
+    from cognee.infrastructure.databases.dataset_database_handler.supported_dataset_database_handlers import (
+        supported_dataset_database_handlers,
+    )
+
+    handler = supported_dataset_database_handlers[graph_config.graph_dataset_database_handler]
+    return await handler["handler_instance"].create_dataset(dataset_id, user)
+
+
+async def _existing_dataset_database(
+    dataset_id: UUID,
+    user: User,
+) -> Optional[DatasetDatabase]:
+    """
+    Check if a DatasetDatabase row already exists for the given owner + dataset.
+    Return None if it doesn't exist, return the row if it does.
+    Args:
+        dataset_id:
+        user:
+
+    Returns:
+        DatasetDatabase or None
+    """
+    db_engine = get_relational_engine()
+
+    async with db_engine.get_async_session() as session:
+        stmt = select(DatasetDatabase).where(
+            DatasetDatabase.owner_id == user.id,
+            DatasetDatabase.dataset_id == dataset_id,
+        )
+        existing: DatasetDatabase = await session.scalar(stmt)
+        return existing
+
+
 async def get_or_create_dataset_database(
     dataset: Union[str, UUID],
     user: User,
@@ -21,6 +70,8 @@ async def get_or_create_dataset_database(
     • If the row already exists, it is fetched and returned.
     • Otherwise a new one is created atomically and returned.
 
+    DatasetDatabase row contains connection and provider info for vector and graph databases.
+
     Parameters
     ----------
     user : User
@@ -32,29 +83,26 @@ async def get_or_create_dataset_database(
 
     dataset_id = await get_unique_dataset_id(dataset, user)
 
-    vector_db_name = f"{dataset_id}.lance.db"
-    graph_db_name = f"{dataset_id}.pkl"
+    # If dataset is given as name make sure the dataset is created first
+    if isinstance(dataset, str):
+        async with db_engine.get_async_session() as session:
+            await create_dataset(dataset, user, session)
 
-    async with db_engine.get_async_session() as session:
-        # Create dataset if it doesn't exist
-        if isinstance(dataset, str):
-            dataset = await create_dataset(dataset, user, session)
+    # If dataset database already exists return it
+    existing_dataset_database = await _existing_dataset_database(dataset_id, user)
+    if existing_dataset_database:
+        return existing_dataset_database
 
-        # Try to fetch an existing row first
-        stmt = select(DatasetDatabase).where(
-            DatasetDatabase.owner_id == user.id,
-            DatasetDatabase.dataset_id == dataset_id,
-        )
-        existing: DatasetDatabase = await session.scalar(stmt)
-        if existing:
-            return existing
+    graph_config_dict = await _get_graph_db_info(dataset_id, user)
+    vector_config_dict = await _get_vector_db_info(dataset_id, user)
 
+    async with db_engine.get_async_session() as session:
         # If there are no existing rows build a new row
         record = DatasetDatabase(
             owner_id=user.id,
             dataset_id=dataset_id,
-            vector_database_name=vector_db_name,
-            graph_database_name=graph_db_name,
+            **graph_config_dict,  # Unpack graph db config
+            **vector_config_dict,  # Unpack vector db config
         )
 
         try:

cognee/infrastructure/databases/utils/get_vector_dataset_database_handler.py

@@ -0,0 +1,10 @@
+from cognee.modules.users.models.DatasetDatabase import DatasetDatabase
+
+
+def get_vector_dataset_database_handler(dataset_database: DatasetDatabase) -> dict:
+    from cognee.infrastructure.databases.dataset_database_handler.supported_dataset_database_handlers import (
+        supported_dataset_database_handlers,
+    )
+
+    handler = supported_dataset_database_handlers[dataset_database.vector_dataset_database_handler]
+    return handler

cognee/infrastructure/databases/utils/resolve_dataset_database_connection_info.py

@@ -0,0 +1,30 @@
+from cognee.infrastructure.databases.utils.get_graph_dataset_database_handler import (
+    get_graph_dataset_database_handler,
+)
+from cognee.infrastructure.databases.utils.get_vector_dataset_database_handler import (
+    get_vector_dataset_database_handler,
+)
+from cognee.modules.users.models.DatasetDatabase import DatasetDatabase
+
+
+async def resolve_dataset_database_connection_info(
+    dataset_database: DatasetDatabase,
+) -> DatasetDatabase:
+    """
+    Resolve the connection info for the given DatasetDatabase instance.
+    Resolve both vector and graph database connection info and return the updated DatasetDatabase instance.
+
+    Args:
+        dataset_database: DatasetDatabase instance
+    Returns:
+        DatasetDatabase instance with resolved connection info
+    """
+    vector_dataset_database_handler = get_vector_dataset_database_handler(dataset_database)
+    graph_dataset_database_handler = get_graph_dataset_database_handler(dataset_database)
+    dataset_database = await vector_dataset_database_handler[
+        "handler_instance"
+    ].resolve_dataset_connection_info(dataset_database)
+    dataset_database = await graph_dataset_database_handler[
+        "handler_instance"
+    ].resolve_dataset_connection_info(dataset_database)
+    return dataset_database

cognee/infrastructure/databases/vector/config.py

@@ -18,14 +18,17 @@ class VectorConfig(BaseSettings):
     Instance variables:
     - vector_db_url: The URL of the vector database.
     - vector_db_port: The port for the vector database.
+    - vector_db_name: The name of the vector database.
     - vector_db_key: The key for accessing the vector database.
     - vector_db_provider: The provider for the vector database.
     """
 
     vector_db_url: str = ""
     vector_db_port: int = 1234
+    vector_db_name: str = ""
     vector_db_key: str = ""
     vector_db_provider: str = "lancedb"
+    vector_dataset_database_handler: str = "lancedb"
 
     model_config = SettingsConfigDict(env_file=".env", extra="allow")
 
@@ -58,8 +61,10 @@ class VectorConfig(BaseSettings):
         return {
             "vector_db_url": self.vector_db_url,
             "vector_db_port": self.vector_db_port,
+            "vector_db_name": self.vector_db_name,
             "vector_db_key": self.vector_db_key,
             "vector_db_provider": self.vector_db_provider,
+            "vector_dataset_database_handler": self.vector_dataset_database_handler,
         }
 
 

cognee/infrastructure/databases/vector/create_vector_engine.py

@@ -1,5 +1,6 @@
 from .supported_databases import supported_databases
 from .embeddings import get_embedding_engine
+from cognee.infrastructure.databases.graph.config import get_graph_context_config
 
 from functools import lru_cache
 
@@ -8,8 +9,10 @@ from functools import lru_cache
 def create_vector_engine(
     vector_db_provider: str,
     vector_db_url: str,
+    vector_db_name: str,
     vector_db_port: str = "",
     vector_db_key: str = "",
+    vector_dataset_database_handler: str = "",
 ):
     """
     Create a vector database engine based on the specified provider.
@@ -27,6 +30,7 @@ def create_vector_engine(
         - vector_db_url (str): The URL for the vector database instance.
         - vector_db_port (str): The port for the vector database instance. Required for some
           providers.
+        - vector_db_name (str): The name of the vector database instance.
         - vector_db_key (str): The API key or access token for the vector database instance.
         - vector_db_provider (str): The name of the vector database provider to use (e.g.,
           'pgvector').
@@ -45,6 +49,7 @@ def create_vector_engine(
             url=vector_db_url,
             api_key=vector_db_key,
             embedding_engine=embedding_engine,
+            database_name=vector_db_name,
         )
 
     if vector_db_provider.lower() == "pgvector":
@@ -133,6 +138,6 @@ def create_vector_engine(
 
     else:
         raise EnvironmentError(
-            f"Unsupported graph database provider: {vector_db_provider}. "
+            f"Unsupported vector database provider: {vector_db_provider}. "
             f"Supported providers are: {', '.join(list(supported_databases.keys()) + ['LanceDB', 'PGVector', 'neptune_analytics', 'ChromaDB'])}"
         )

cognee/infrastructure/databases/vector/embeddings/FastembedEmbeddingEngine.py

@@ -17,6 +17,7 @@ from cognee.infrastructure.databases.exceptions import EmbeddingException
 from cognee.infrastructure.llm.tokenizer.TikToken import (
     TikTokenTokenizer,
 )
+from cognee.shared.rate_limiting import embedding_rate_limiter_context_manager
 
 litellm.set_verbose = False
 logger = get_logger("FastembedEmbeddingEngine")
@@ -68,7 +69,7 @@ class FastembedEmbeddingEngine(EmbeddingEngine):
 
     @retry(
         stop=stop_after_delay(128),
-        wait=wait_exponential_jitter(2, 128),
+        wait=wait_exponential_jitter(8, 128),
         retry=retry_if_not_exception_type(litellm.exceptions.NotFoundError),
         before_sleep=before_sleep_log(logger, logging.DEBUG),
         reraise=True,
@@ -96,11 +97,12 @@ class FastembedEmbeddingEngine(EmbeddingEngine):
             if self.mock:
                 return [[0.0] * self.dimensions for _ in text]
             else:
-                embeddings = self.embedding_model.embed(
-                    text,
-                    batch_size=len(text),
-                    parallel=None,
-                )
+                async with embedding_rate_limiter_context_manager():
+                    embeddings = self.embedding_model.embed(
+                        text,
+                        batch_size=len(text),
+                        parallel=None,
+                    )
 
                 return list(embeddings)
 

cognee/infrastructure/databases/vector/embeddings/LiteLLMEmbeddingEngine.py

@@ -25,6 +25,7 @@ from cognee.infrastructure.llm.tokenizer.Mistral import (
 from cognee.infrastructure.llm.tokenizer.TikToken import (
     TikTokenTokenizer,
 )
+from cognee.shared.rate_limiting import embedding_rate_limiter_context_manager
 
 litellm.set_verbose = False
 logger = get_logger("LiteLLMEmbeddingEngine")
@@ -109,13 +110,14 @@ class LiteLLMEmbeddingEngine(EmbeddingEngine):
                 response = {"data": [{"embedding": [0.0] * self.dimensions} for _ in text]}
                 return [data["embedding"] for data in response["data"]]
             else:
-                response = await litellm.aembedding(
-                    model=self.model,
-                    input=text,
-                    api_key=self.api_key,
-                    api_base=self.endpoint,
-                    api_version=self.api_version,
-                )
+                async with embedding_rate_limiter_context_manager():
+                    response = await litellm.aembedding(
+                        model=self.model,
+                        input=text,
+                        api_key=self.api_key,
+                        api_base=self.endpoint,
+                        api_version=self.api_version,
+                    )
 
                 return [data["embedding"] for data in response.data]
 

cognee/infrastructure/databases/vector/embeddings/OllamaEmbeddingEngine.py

@@ -18,10 +18,7 @@ from cognee.infrastructure.databases.vector.embeddings.EmbeddingEngine import Em
 from cognee.infrastructure.llm.tokenizer.HuggingFace import (
     HuggingFaceTokenizer,
 )
-from cognee.infrastructure.databases.vector.embeddings.embedding_rate_limiter import (
-    embedding_rate_limit_async,
-    embedding_sleep_and_retry_async,
-)
+from cognee.shared.rate_limiting import embedding_rate_limiter_context_manager
 from cognee.shared.utils import create_secure_ssl_context
 
 logger = get_logger("OllamaEmbeddingEngine")
@@ -101,7 +98,7 @@ class OllamaEmbeddingEngine(EmbeddingEngine):
 
     @retry(
         stop=stop_after_delay(128),
-        wait=wait_exponential_jitter(2, 128),
+        wait=wait_exponential_jitter(8, 128),
         retry=retry_if_not_exception_type(litellm.exceptions.NotFoundError),
         before_sleep=before_sleep_log(logger, logging.DEBUG),
         reraise=True,
@@ -120,11 +117,15 @@ class OllamaEmbeddingEngine(EmbeddingEngine):
         ssl_context = create_secure_ssl_context()
         connector = aiohttp.TCPConnector(ssl=ssl_context)
         async with aiohttp.ClientSession(connector=connector) as session:
-            async with session.post(
-                self.endpoint, json=payload, headers=headers, timeout=60.0
-            ) as response:
-                data = await response.json()
-                return data["embeddings"][0]
+            async with embedding_rate_limiter_context_manager():
+                async with session.post(
+                    self.endpoint, json=payload, headers=headers, timeout=60.0
+                ) as response:
+                    data = await response.json()
+                    if "embeddings" in data:
+                        return data["embeddings"][0]
+                    else:
+                        return data["data"][0]["embedding"]
 
     def get_vector_size(self) -> int:
         """

cognee/infrastructure/databases/vector/lancedb/LanceDBAdapter.py

@@ -193,6 +193,8 @@ class LanceDBAdapter(VectorDBInterface):
             for (data_point_index, data_point) in enumerate(data_points)
         ]
 
+        lance_data_points = list({dp.id: dp for dp in lance_data_points}.values())
+
         async with self.VECTOR_DB_LOCK:
             await (
                 collection.merge_insert("id")

cognee/infrastructure/databases/vector/lancedb/LanceDBDatasetDatabaseHandler.py

@@ -0,0 +1,50 @@
+import os
+from uuid import UUID
+from typing import Optional
+
+from cognee.infrastructure.databases.vector.create_vector_engine import create_vector_engine
+from cognee.modules.users.models import User
+from cognee.modules.users.models import DatasetDatabase
+from cognee.base_config import get_base_config
+from cognee.infrastructure.databases.vector import get_vectordb_config
+from cognee.infrastructure.databases.dataset_database_handler import DatasetDatabaseHandlerInterface
+
+
+class LanceDBDatasetDatabaseHandler(DatasetDatabaseHandlerInterface):
+    """
+    Handler for interacting with LanceDB Dataset databases.
+    """
+
+    @classmethod
+    async def create_dataset(cls, dataset_id: Optional[UUID], user: Optional[User]) -> dict:
+        vector_config = get_vectordb_config()
+        base_config = get_base_config()
+
+        if vector_config.vector_db_provider != "lancedb":
+            raise ValueError(
+                "LanceDBDatasetDatabaseHandler can only be used with LanceDB vector database provider."
+            )
+
+        databases_directory_path = os.path.join(
+            base_config.system_root_directory, "databases", str(user.id)
+        )
+
+        vector_db_name = f"{dataset_id}.lance.db"
+
+        return {
+            "vector_database_provider": vector_config.vector_db_provider,
+            "vector_database_url": os.path.join(databases_directory_path, vector_db_name),
+            "vector_database_key": vector_config.vector_db_key,
+            "vector_database_name": vector_db_name,
+            "vector_dataset_database_handler": "lancedb",
+        }
+
+    @classmethod
+    async def delete_dataset(cls, dataset_database: DatasetDatabase):
+        vector_engine = create_vector_engine(
+            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,
+        )
+        await vector_engine.prune()

cognee/infrastructure/databases/vector/vector_db_interface.py

@@ -2,6 +2,8 @@ from typing import List, Protocol, Optional, Union, Any
 from abc import abstractmethod
 from cognee.infrastructure.engine import DataPoint
 from .models.PayloadSchema import PayloadSchema
+from uuid import UUID
+from cognee.modules.users.models import User
 
 
 class VectorDBInterface(Protocol):
@@ -217,3 +219,36 @@ class VectorDBInterface(Protocol):
             - Any: The schema object suitable for this vector database
         """
         return model_type
+
+    @classmethod
+    async def create_dataset(cls, dataset_id: Optional[UUID], user: Optional[User]) -> dict:
+        """
+        Return a dictionary with connection info for a vector database for the given dataset.
+        Function can auto handle deploying of the actual database if needed, but is not necessary.
+        Only providing connection info is sufficient, this info will be mapped when trying to connect to the provided dataset in the future.
+        Needed for Cognee multi-tenant/multi-user and backend access control support.
+
+        Dictionary returned from this function will be used to create a DatasetDatabase row in the relational database.
+        From which internal mapping of dataset -> database connection info will be done.
+
+        Each dataset needs to map to a unique vector database when backend access control is enabled to facilitate a separation of concern for data.
+
+        Args:
+            dataset_id: UUID of the dataset if needed by the database creation logic
+            user: User object if needed by the database creation logic
+        Returns:
+            dict: Connection info for the created vector database instance.
+        """
+        pass
+
+    async def delete_dataset(self, dataset_id: UUID, user: User) -> None:
+        """
+        Delete the vector database for the given dataset.
+        Function should auto handle deleting of the actual database or send a request to the proper service to delete the database.
+        Needed for maintaining a database for Cognee multi-tenant/multi-user and backend access control.
+
+        Args:
+            dataset_id: UUID of the dataset
+            user: User object
+        """
+        pass

cognee/infrastructure/engine/models/Edge.py

@@ -1,4 +1,4 @@
-from pydantic import BaseModel
+from pydantic import BaseModel, field_validator
 from typing import Optional, Any, Dict
 
 
@@ -18,9 +18,21 @@ class Edge(BaseModel):
 
         # Mixed usage
         has_items: (Edge(weight=0.5, weights={"confidence": 0.9}), list[Item])
+
+        # With edge_text for rich embedding representation
+        contains: (Edge(relationship_type="contains", edge_text="relationship_name: contains; entity_description: Alice"), Entity)
     """
 
     weight: Optional[float] = None
     weights: Optional[Dict[str, float]] = None
     relationship_type: Optional[str] = None
     properties: Optional[Dict[str, Any]] = None
+    edge_text: Optional[str] = None
+
+    @field_validator("edge_text", mode="before")
+    @classmethod
+    def ensure_edge_text(cls, v, info):
+        """Auto-populate edge_text from relationship_type if not explicitly provided."""
+        if v is None and info.data.get("relationship_type"):
+            return info.data["relationship_type"]
+        return v

cognee/infrastructure/files/storage/s3_config.py

@@ -9,6 +9,8 @@ class S3Config(BaseSettings):
     aws_access_key_id: Optional[str] = None
     aws_secret_access_key: Optional[str] = None
     aws_session_token: Optional[str] = None
+    aws_profile_name: Optional[str] = None
+    aws_bedrock_runtime_endpoint: Optional[str] = None
     model_config = SettingsConfigDict(env_file=".env", extra="allow")
 
 

cognee/infrastructure/files/utils/guess_file_type.py

@@ -55,6 +55,10 @@ def guess_file_type(file: BinaryIO, name: Optional[str] = None) -> filetype.Type
         file_type = Type("text/plain", "txt")
         return file_type
 
+    if ext in [".csv"]:
+        file_type = Type("text/csv", "csv")
+        return file_type
+
     file_type = filetype.guess(file)
 
     # If file type could not be determined consider it a plain text file as they don't have magic number encoding

cognee/infrastructure/llm/LLMGateway.py

@@ -11,7 +11,7 @@ class LLMGateway:
 
     @staticmethod
     def acreate_structured_output(
-        text_input: str, system_prompt: str, response_model: Type[BaseModel]
+        text_input: str, system_prompt: str, response_model: Type[BaseModel], **kwargs
     ) -> Coroutine:
         llm_config = get_llm_config()
         if llm_config.structured_output_framework.upper() == "BAML":
@@ -31,7 +31,10 @@ class LLMGateway:
 
             llm_client = get_llm_client()
             return llm_client.acreate_structured_output(
-                text_input=text_input, system_prompt=system_prompt, response_model=response_model
+                text_input=text_input,
+                system_prompt=system_prompt,
+                response_model=response_model,
+                **kwargs,
             )
 
     @staticmethod