cognee/modules/data/processing/document_types/Document.py (1)

Line range hint 8-9: Consider implementing abstract method

The read method is currently a pass-through implementation. Since this is a base class, consider:

1. Making it an abstract method using @abstractmethod
2. Adding docstring to document the expected behavior

+from abc import abstractmethod
 from cognee.infrastructure.engine import DataPoint

 class Document(DataPoint):
     type: str
     name: str
     raw_data_location: str

+    @abstractmethod
     def read(self, chunk_size: int) -> str:
+        """Read document content in chunks.
+        
+        Args:
+            chunk_size: Size of each chunk to read
+            
+        Returns:
+            str: Content chunk from the document
+        """
         pass

cognee/infrastructure/databases/vector/models/ScoredResult.py (1)

6-8: Consider adding field documentation

While the type change from str to UUID is a good improvement for type safety, consider adding docstrings to clarify the purpose of each field, especially the payload dictionary's expected content and structure.

Example improvement:

 class ScoredResult(BaseModel):
+    """Represents a scored result from a vector search operation.
+
+    Attributes:
+        id (UUID): Unique identifier for the result
+        score (float): Similarity score (lower is better)
+        payload (Dict[str, Any]): Additional metadata associated with the result
+    """
     id: UUID
     score: float # Lower score is better
     payload: Dict[str, Any]

cognee/tasks/documents/extract_chunks_from_documents.py (1)

4-4: Consider making chunk_size configurable via settings

The default chunk_size of 1024 bytes might not be optimal for all use cases. Consider:

1. Making it configurable through environment variables or settings
2. Adjusting based on available system memory
3. Adding documentation about choosing appropriate chunk sizes

cognee/modules/engine/models/EntityType.py (1)

6-6: Consider renaming the 'type' field to avoid shadowing built-in type

The attribute name 'type' could potentially cause confusion or issues as it shadows Python's built-in type() function.

Consider renaming to something more specific like entity_type or type_category.

-    type: str
+    entity_type: str

cognee/modules/engine/models/Entity.py (1)

10-12: Consider additional index fields

The current implementation only indexes the name field. Consider whether is_a should also be indexed for efficient entity type filtering.

cognee/modules/chunking/models/DocumentChunk.py (1)

12-14: Consider indexing additional fields for search optimization

The metadata configuration only indexes the text field. Consider whether other fields like chunk_index should also be indexed for performance optimization, especially if they're used in queries or filtering operations.

Consider expanding the index_fields if these fields are frequently queried:

     _metadata: Optional[dict] = {
-        "index_fields": ["text"],
+        "index_fields": ["text", "chunk_index", "cut_type"],
     }

cognee/tasks/chunks/query_chunks.py (1)

Line range hint 1-17: Consider these improvements for robustness.

While the basic functionality looks good, consider these enhancements:

1. Make the search limit configurable instead of hard-coding to 5
2. Add error handling for vector engine operations
3. Add input validation for the query parameter

Here's a suggested implementation:

+from typing import Optional
 from cognee.infrastructure.databases.vector import get_vector_engine
+from cognee.config import settings  # Assuming you have a settings module
+
+DEFAULT_SEARCH_LIMIT = 5

-async def query_chunks(query: str) -> list[dict]:
+async def query_chunks(query: str, limit: Optional[int] = None) -> list[dict]:
     """
     Parameters:
     - query (str): The query string to filter nodes by.
+    - limit (Optional[int]): Maximum number of results to return. Defaults to settings or 5.

     Returns:
     - list(dict): A list of objects providing information about the chunks related to query.
+    
+    Raises:
+    - ValueError: If query is empty or limit is negative
+    - VectorEngineError: If search operation fails
     """
+    if not query or not query.strip():
+        raise ValueError("Query string cannot be empty")
+
+    search_limit = limit or getattr(settings, 'VECTOR_SEARCH_LIMIT', DEFAULT_SEARCH_LIMIT)
+    if search_limit < 1:
+        raise ValueError("Search limit must be positive")
+
     vector_engine = get_vector_engine()
+    try:
+        found_chunks = await vector_engine.search(
+            "DocumentChunk_text",
+            query,
+            limit=search_limit
+        )
+        return [result.payload for result in found_chunks]
+    except Exception as e:
+        raise VectorEngineError(f"Failed to search chunks: {str(e)}") from e
-    found_chunks = await vector_engine.search("DocumentChunk_text", query, limit = 5)
-    chunks = [result.payload for result in found_chunks]
-    return chunks

cognee/modules/data/processing/document_types/ImageDocument.py (1)

13-13: Document the TextChunker parameter change

The change from self.id to self as the first parameter of TextChunker suggests that the chunker now operates on the entire document object. Consider adding a comment explaining this design choice for better maintainability.

-        chunker = TextChunker(self, chunk_size = chunk_size, get_text = lambda: text)
+        # TextChunker now takes the entire document object to access all document properties
+        chunker = TextChunker(self, chunk_size = chunk_size, get_text = lambda: text)

cognee/modules/data/processing/document_types/PdfDocument.py (2)

8-20: Add method documentation

The read method lacks documentation explaining its purpose, parameters, and return value.

  def read(self, chunk_size: int):
+     """Read PDF document and yield text chunks.
+
+     Args:
+         chunk_size (int): The maximum size of each text chunk
+
+     Yields:
+         str: Text chunks from the PDF document
+
+     Raises:
+         ValueError: If raw_data_location is not set
+         FileNotFoundError: If PDF file doesn't exist
+         RuntimeError: If PDF file cannot be read
+     """

---

18-20: Improve resource cleanup with context manager

The file cleanup could be more robust by using a context manager pattern.

  def read(self, chunk_size: int):
-     file = PdfReader(self.raw_data_location)
+     with open(self.raw_data_location, 'rb') as pdf_file:
+         file = PdfReader(pdf_file)
+         def get_text():
+             for page in file.pages:
+                 page_text = page.extract_text()
+                 yield page_text
 
-     def get_text():
-         for page in file.pages:
-             page_text = page.extract_text()
-             yield page_text
+         chunker = TextChunker(self, chunk_size=chunk_size, get_text=get_text)
+         yield from chunker.read()
-     chunker = TextChunker(self, chunk_size = chunk_size, get_text = get_text)
-
-     yield from chunker.read()
-
-     file.stream.close()

cognee/infrastructure/databases/vector/pgvector/serialize_data.py (1)

4-5: Update docstring to reflect UUID handling capability.

The current docstring only mentions datetime objects, but the function now handles UUID objects as well.

-    """Recursively convert datetime objects in dictionaries/lists to ISO format."""
+    """Recursively serialize data by converting datetime objects to ISO format and UUID objects to strings."""

cognee/modules/data/processing/document_types/TextDocument.py (1)

Line range hint 8-16: Add error handling and use context manager for file operations

The file operations lack proper error handling and resource management. Consider these improvements:

1. Handle potential IOError/OSError exceptions
2. Use context manager with contextlib.contextmanager for better resource management

-    def get_text():
-        with open(self.raw_data_location, mode = "r", encoding = "utf-8") as file:
-            while True:
-                text = file.read(1024)
-
-                if len(text.strip()) == 0:
-                    break
-
-                yield text
+    @contextmanager
+    def get_text():
+        try:
+            with open(self.raw_data_location, mode="r", encoding="utf-8") as file:
+                while True:
+                    text = file.read(1024)
+                    if not text.strip():
+                        break
+                    yield text
+        except IOError as e:
+            logger.error(f"Error reading file {self.raw_data_location}: {e}")
+            raise

cognee/tasks/graph/extract_graph_from_code.py (2)

17-17: Consider returning processing status information

The function currently returns the original data_chunks without any indication of which chunks were successfully processed. Consider enhancing the return type to include processing status.

-    return data_chunks
+    return {
+        'chunks': data_chunks,
+        'processing_status': [
+            {'chunk_index': i, 'success': chunk_graphs[i] is not None}
+            for i in range(len(data_chunks))
+        ]
+    }

---

8-8: Improve type hints specificity

The type hint for graph_model could be more specific about what kind of BaseModel is expected.

-async def extract_graph_from_code(data_chunks: list[DocumentChunk], graph_model: Type[BaseModel]):
+async def extract_graph_from_code(
+    data_chunks: list[DocumentChunk],
+    graph_model: Type[BaseModel]  # TODO: Replace BaseModel with your specific graph model class
+) -> dict[str, list]:

cognee/infrastructure/engine/models/DataPoint.py (2)

7-8: Add docstring to explain MetaData's purpose and usage

The MetaData TypedDict would benefit from documentation explaining the purpose of index_fields and how they are used in the data model.

 class MetaData(TypedDict):
+    """Metadata configuration for DataPoint models.
+    
+    Attributes:
+        index_fields: List of field names that should be indexed for searching/embedding
+    """
     index_fields: list[str]

---

17-18: Decide on attribute privacy strategy

The commented-out configuration suggests uncertainty about attribute privacy. If _metadata should be private, we should either:

1. Enable the private attributes configuration, or
2. Use property decorators to control access

Please clarify the intended privacy model for this class. If you need help implementing either approach, I can provide specific code examples.

cognee/infrastructure/databases/vector/config.py (1)

11-11: Consider using Redis default port and adding port validation

Since FalkorDB is Redis-based, consider:

1. Using Redis default port 6379 instead of 1234
2. Adding port range validation (1-65535)

-    vector_db_port: int = 1234
+    vector_db_port: int = 6379
+
+    @property
+    def vector_db_port(self) -> int:
+        return self._vector_db_port
+
+    @vector_db_port.setter
+    def vector_db_port(self, value: int) -> None:
+        if not 1 <= value <= 65535:
+            raise ValueError("Port must be between 1 and 65535")
+        self._vector_db_port = value

cognee/tasks/chunks/remove_disconnected_chunks.py (3)

4-4: Add docstring to clarify function behavior.

The function lacks documentation explaining its purpose, parameters, and return value.

Add a descriptive docstring:

 async def remove_disconnected_chunks(data_chunks: list[DocumentChunk]) -> list[DocumentChunk]:
+    """Remove disconnected chunks from the graph.
+
+    Args:
+        data_chunks (list[DocumentChunk]): List of document chunks to process
+
+    Returns:
+        list[DocumentChunk]: The input chunks (unchanged)
+    """

---

Line range hint 4-27: Consider essential improvements for robustness and performance.

1. Error Handling: Add try-catch blocks for graph operations to handle potential failures gracefully.
2. Performance: Consider batching graph operations to reduce the number of await calls.
3. Return Value: The function returns the input unchanged, which might be misleading. Consider returning only the connected chunks or adding a success indicator.

Here's a suggested implementation:

 async def remove_disconnected_chunks(data_chunks: list[DocumentChunk]) -> list[DocumentChunk]:
+    """Remove disconnected chunks from the graph.
+
+    Args:
+        data_chunks (list[DocumentChunk]): List of document chunks to process
+
+    Returns:
+        list[DocumentChunk]: The input chunks (unchanged)
+
+    Raises:
+        GraphEngineError: If graph operations fail
+    """
     graph_engine = await get_graph_engine()
     document_ids = set((data_chunk.document_id for data_chunk in data_chunks))
     obsolete_chunk_ids = []
 
-    for document_id in document_ids:
-        chunks = await graph_engine.get_successors(document_id, edge_label = "has_chunk")
-
-        for chunk in chunks:
-            previous_chunks = await graph_engine.get_predecessors(chunk["uuid"], edge_label = "next_chunk")
-
-            if len(previous_chunks) == 0:
-                obsolete_chunk_ids.append(chunk["uuid"])
-
-    if len(obsolete_chunk_ids) > 0:
-        await graph_engine.delete_nodes(obsolete_chunk_ids)
-
-    disconnected_nodes = await graph_engine.get_disconnected_nodes()
-    if len(disconnected_nodes) > 0:
-        await graph_engine.delete_nodes(disconnected_nodes)
+    try:
+        # Batch get_successors for all document_ids
+        chunks_by_doc = await graph_engine.get_successors_batch(
+            list(document_ids),
+            edge_label="has_chunk"
+        )
+        
+        # Collect all chunk UUIDs
+        chunk_uuids = [
+            chunk["uuid"]
+            for doc_chunks in chunks_by_doc.values()
+            for chunk in doc_chunks
+        ]
+        
+        # Batch get_predecessors for all chunks
+        previous_chunks_by_uuid = await graph_engine.get_predecessors_batch(
+            chunk_uuids,
+            edge_label="next_chunk"
+        )
+        
+        # Identify obsolete chunks
+        obsolete_chunk_ids = [
+            uuid for uuid, prev_chunks in previous_chunks_by_uuid.items()
+            if len(prev_chunks) == 0
+        ]
+        
+        # Delete obsolete chunks if any
+        if obsolete_chunk_ids:
+            await graph_engine.delete_nodes(obsolete_chunk_ids)
+        
+        # Delete disconnected nodes
+        disconnected_nodes = await graph_engine.get_disconnected_nodes()
+        if disconnected_nodes:
+            await graph_engine.delete_nodes(disconnected_nodes)
+            
+    except Exception as e:
+        raise GraphEngineError(f"Failed to remove disconnected chunks: {str(e)}") from e
 
     return data_chunks

---

Batch operations are not available in the current implementation

Based on the complete interface and implementation review:

• The GraphDBInterface does not define get_successors_batch or get_predecessors_batch methods
• The FalkorDB implementation only has batch operations for search, but not for graph traversal
• The current implementation would need to be extended to support batch operations for graph traversal

🔗 Analysis chain

Line range hint 4-27: Verify availability of batch operations in graph engine.

The suggested performance optimization assumes the existence of batch operations (get_successors_batch, get_predecessors_batch). Let's verify if these methods are available or need to be implemented.

---

Let me check the graph engine implementation more specifically to verify the batch operations.

---

Let me check the specific implementation of get_successors and get_predecessors methods.

---

Let me check the complete GraphDBInterface and FalkorDB implementation to understand all available methods.

🏁 Scripts executed

The following scripts were executed for the analysis:

Script:

#!/bin/bash
# Search for batch operation methods in the graph engine implementation
ast-grep --pattern 'get_successors_batch' 
ast-grep --pattern 'get_predecessors_batch'

# Search for similar batch operations that might be used instead
rg -l "batch" --type py

Length of output: 1008

---

Script:

#!/bin/bash
# Check the FalkorDB adapter implementation for batch operations
rg -A 5 "class FalkorDBAdapter" cognee/infrastructure/databases/hybrid/falkordb/FalkorDBAdapter.py

# Search for any batch-related methods in the FalkorDB adapter
rg "batch|get_successors|get_predecessors" cognee/infrastructure/databases/hybrid/falkordb/FalkorDBAdapter.py

Length of output: 432

---

Script:

#!/bin/bash
# Search for get_successors and get_predecessors implementation
ast-grep --pattern 'async def get_successors' 
ast-grep --pattern 'async def get_predecessors'

# Also check the GraphDBInterface definition
rg -A 10 "class GraphDBInterface" 

Length of output: 1078

---

Script:

#!/bin/bash
# Get the complete GraphDBInterface definition
cat cognee/infrastructure/databases/graph/graph_db_interface.py

# Get the complete FalkorDB implementation focusing on graph operations
rg -A 5 "async def" cognee/infrastructure/databases/hybrid/falkordb/FalkorDBAdapter.py

Length of output: 6508

cognee/modules/data/operations/detect_language.py (2)

17-18: Extract magic number to a constant.

The 100-character limit should be defined as a module-level constant for better maintainability.

Apply this diff:

import logging
from langdetect import detect, LangDetectException

+# Maximum text length for language detection to ensure efficient processing
+MAX_DETECTION_LENGTH = 100
+
logger = logging.getLogger(__name__)

async def detect_language(text: str) -> str | None:
    # Trim the text to the first 100 characters
-    trimmed_text = text[:100]
+    trimmed_text = text[:MAX_DETECTION_LENGTH]

---

30-36: Consolidate error handling returns.

The duplicate return None statements can be consolidated by moving it after both except blocks.

Apply this diff:

     except LangDetectException as e:
         logger.error(f"Language detection error: {e}")
-
     except Exception as e:
         logger.error(f"Unexpected error: {e}")

-    return None
+        return None

cognee/modules/graph/utils/get_model_instance_from_graph.py (2)

1-6: Add docstring documentation for better clarity.

The function signature is well-typed but lacks documentation. Consider adding a docstring that explains:

• Purpose of the function
• Parameter descriptions
• Return value description
• Example usage

 def get_model_instance_from_graph(nodes: list[DataPoint], edges: list, entity_id: str):
+    """Reconstructs a model instance from its graph representation.
+
+    Args:
+        nodes (list[DataPoint]): List of DataPoint nodes representing model instances
+        edges (list): List of edges connecting nodes, each edge is [source_id, target_id, label, properties?]
+        entity_id (str): ID of the root entity to reconstruct
+
+    Returns:
+        Model instance reconstructed from the graph representation
+    """

---

6-29: Consider performance optimizations for large graphs.

For better scalability:

1. Consider implementing caching for frequently accessed model instances
2. For large graphs, consider implementing batch processing or streaming to reduce memory usage
3. Add performance metrics/logging for monitoring reconstruction time

Would you like help implementing any of these optimizations?

cognee/infrastructure/databases/graph/graph_db_interface.py (1)

6-7: LGTM! Consider enhancing type hints and documentation.

The new query method with parameterized queries is a good improvement over the previous graph() method. Consider these enhancements:

1. Add return type hint
2. Use more specific type for params
3. Add docstring explaining the expected query format

     @abstractmethod
-    async def query(self, query: str, params: dict):
+    async def query(self, query: str, params: dict[str, Any]) -> Any:
+        """Execute a graph database query with parameters.
+        
+        Args:
+            query: The query string in the graph database's native query language
+            params: Dictionary of query parameters to prevent injection
+        
+        Returns:
+            Query results in the format specific to the implementation
+        """
         raise NotImplementedError

cognee/infrastructure/databases/graph/get_graph_engine.py (1)

Line range hint 1-42: Consider enhancing the factory pattern implementation

The current implementation mixes several concerns that could be better separated:

1. Consider extracting provider-specific initialization logic into separate factory classes
2. Implement a proper strategy pattern for different graph database providers
3. Consider adding a registration mechanism for new providers

This would make the code more maintainable and easier to extend with new providers in the future.

Example structure:

class GraphEngineFactory:
    _providers = {}
    
    @classmethod
    def register_provider(cls, name: str, provider_class: Type[GraphDBInterface]):
        cls._providers[name] = provider_class
    
    @classmethod
    async def create(cls, config: GraphConfig) -> GraphDBInterface:
        provider = cls._providers.get(config.graph_database_provider)
        if not provider:
            return cls._get_fallback_provider(config)
        return await provider.create(config)

🧰 Tools

🪛 Ruff

20-20: Do not use bare except

(E722)

cognee/modules/storage/utils/__init__.py (1)

34-46: Improve robustness and readability of property extraction

The current implementation has a few areas that could be enhanced:

1. The list check assumes non-empty lists without validation
2. Direct iteration over data_point could be replaced with .items()
3. Complex condition could be extracted into a helper function

Consider this refactor:

 def get_own_properties(data_point: DataPoint):
     properties = {}
 
-    for field_name, field_value in data_point:
-        if field_name == "_metadata" \
-            or isinstance(field_value, dict) \
-            or isinstance(field_value, DataPoint) \
-            or (isinstance(field_value, list) and isinstance(field_value[0], DataPoint)):
+    def is_complex_type(value) -> bool:
+        return (isinstance(value, (dict, DataPoint)) or
+                (isinstance(value, list) and value and isinstance(value[0], DataPoint)))
+
+    for field_name, field_value in data_point.model_dump().items():
+        if field_name == "_metadata" or is_complex_type(field_value):
             continue
 
         properties[field_name] = field_value
 
     return properties

cognee/infrastructure/engine/__tests__/model_to_graph_to_model.test.py (3)

23-23: Use class variable syntax for _metadata

Using dict() creates a new instance for each class. For class variables, use the proper class variable syntax.

-    _metadata: dict = dict(index_fields = ["name"])
+    _metadata: dict = {"index_fields": ["name"]}

Also applies to: 39-39

---

38-38: Fix inconsistent spelling of 'license'

The field name uses British spelling ('licence') while the sample data uses American spelling ('license'). Maintain consistency across the codebase.

-        driving_licence: Optional[dict]
+        driving_license: Optional[dict]

---

55-60: Use dynamic dates for driving license

The driving license dates are hardcoded to 2025. This could cause the test to fail after that date. Consider using dynamic dates relative to the current date.

+from datetime import datetime, timedelta
+
+current_date = datetime.now()
+expiry_date = (current_date + timedelta(years=3)).strftime("%Y-%m-%d")
+
 driving_licence = {
     "issued_by": "PU Vrsac",
-    "issued_on": "2025-11-06",
+    "issued_on": current_date.strftime("%Y-%m-%d"),
     "number": "1234567890",
-    "expires_on": "2025-11-06",
+    "expires_on": expiry_date,
 }

cognee/tasks/graph/query_graph_connections.py (3)

Line range hint 5-5: Remove unused parameter exploration_levels

The parameter exploration_levels is declared but never used in the function implementation.

-async def query_graph_connections(query: str, exploration_levels = 1) -> list[(str, str, str)]:
+async def query_graph_connections(query: str) -> list[(str, str, str)]:

---

30-31: Extract magic numbers into named constants

The hardcoded values for limit and score threshold should be configuration parameters.

+SEARCH_RESULT_LIMIT = 5
+RELEVANCE_THRESHOLD = 0.5

 vector_engine = get_vector_engine()
 results = await asyncio.gather(
-    vector_engine.search("Entity_name", query_text = query, limit = 5),
-    vector_engine.search("EntityType_name", query_text = query, limit = 5),
+    vector_engine.search("Entity_name", query_text = query, limit = SEARCH_RESULT_LIMIT),
+    vector_engine.search("EntityType_name", query_text = query, limit = SEARCH_RESULT_LIMIT),
 )
 results = [*results[0], *results[1]]
-relevant_results = [result for result in results if result.score < 0.5][:5]
+relevant_results = [result for result in results if result.score < RELEVANCE_THRESHOLD][:SEARCH_RESULT_LIMIT]

---

51-54: Make unique_id construction more robust

The current unique_id construction might be susceptible to edge cases where the components contain spaces.

-        unique_id = f"{node_connection[0]['id']} {node_connection[1]['relationship_name']} {node_connection[2]['id']}"
+        unique_id = f"{node_connection[0]['id']}::{node_connection[1]['relationship_name']}::{node_connection[2]['id']}"

cognee/infrastructure/databases/vector/create_vector_engine.py (2)

5-5: Consider using int type for port and add validation

The vector_db_port is defined as a string, but ports are typically integers within a specific range (0-65535). Consider:

1. Changing the type to int
2. Adding validation to ensure the port is within valid range

-    vector_db_port: str
+    vector_db_port: int

---

Line range hint 9-58: Standardize error handling across providers

Currently, only Weaviate has explicit configuration validation. Consider adding similar validation for other providers to ensure consistent error handling:

• PGVector: Validate relational config
• FalkorDB: Validate URL and port
• Qdrant: Add error for missing config (similar to Weaviate)

Example implementation for FalkorDB:

     elif config["vector_db_provider"] == "falkordb":
         from ..hybrid.falkordb.FalkorDBAdapter import FalkorDBAdapter
 
+        if config["vector_db_url"] is None or config["vector_db_port"] is None:
+            raise EnvironmentError("FalkorDB is not configured!")
+
         return FalkorDBAdapter(
             database_url = config["vector_db_url"],
             database_port = config["vector_db_port"],
             embedding_engine = embedding_engine,
         )

cognee/tasks/chunk_naive_llm_classifier/chunk_naive_llm_classifier.py (1)

Line range hint 124-130: Fix inconsistent relationship naming and direction.

There are two issues in the subtype relationship edge creation:

1. The edge type is "is_subtype_of" but the relationship_name is "contains"
2. The edge direction seems reversed - for "is_subtype_of", the subtype should point to the type, not vice versa

Apply this fix:

                    "is_subtype_of",
                    dict(
-                       relationship_name="contains",
+                       relationship_name="is_subtype_of",
-                       source_node_id=str(classification_type_id),
-                       target_node_id=str(classification_subtype_id),
+                       source_node_id=str(classification_subtype_id),
+                       target_node_id=str(classification_type_id),
                    ),

cognee/infrastructure/databases/vector/weaviate_db/WeaviateAdapter.py (4)

13-19: Consider making IndexSchema more flexible

The IndexSchema implementation could be more flexible to handle multiple index fields and different data types.

Consider this improvement:

 class IndexSchema(DataPoint):
     uuid: str
     text: str
 
     _metadata: dict = {
-        "index_fields": ["text"]
+        "index_fields": ["text"],
+        "field_types": {"text": str},
+        "vectorize_fields": ["text"]
     }

---

93-93: Remove commented-out batch implementation code

The commented-out batch implementation code should be removed if it's no longer needed. If it's kept for future reference, consider moving it to documentation.

Also applies to: 102-116

---

Line range hint 1-194: Standardize error handling across all methods

The class needs a more consistent approach to error handling. Some methods have try-catch blocks while others don't handle errors at all.

Consider:

1. Creating custom exception classes for different error types
2. Adding consistent error handling to all async methods
3. Implementing proper error logging throughout

Example custom exception:

class WeaviateAdapterError(Exception):
    """Base exception for WeaviateAdapter errors."""
    pass

class CollectionNotFoundError(WeaviateAdapterError):
    """Raised when a collection is not found."""
    pass

---

Async implementation needs improvement for non-blocking operations

The WeaviateAdapter's async implementation has several blocking operations that should be handled differently:

• The adapter wraps synchronous Weaviate client operations in asyncio.Future() which doesn't provide true asynchronous behavior
• In contrast, PGVectorAdapter properly uses async context managers and awaits database operations
• Critical points to fix:
  • has_collection(): Direct client call wrapped in Future
  • create_collection(): Direct client call wrapped in Future
  • retrieve(): Direct client call wrapped in Future
  • search(): Direct blocking call to client
  • delete_data_points(): Direct client call wrapped in Future

The implementation should either:

1. Use Weaviate's async client if available
2. Or use asyncio.to_thread() to properly offload blocking operations to a thread pool

🔗 Analysis chain

Line range hint 1-194: Verify async implementation efficiency

Some operations might be blocking despite the async interface. Let's verify the implementation.

🏁 Scripts executed

The following scripts were executed for the analysis:

Script:

#!/bin/bash
# Search for potential blocking operations in async methods
echo "Checking for potential blocking operations in async methods..."
rg -n "self\.client\." cognee/infrastructure/databases/vector/weaviate_db/WeaviateAdapter.py

# Check for other adapter implementations to compare async patterns
echo "Checking async patterns in other adapters..."
fd -t f ".*Adapter\.py" cognee/infrastructure/databases/vector/ -x rg -l "async def"

Length of output: 917

---

Script:

#!/bin/bash
# Let's check the async patterns in PGVectorAdapter for comparison
echo "Checking PGVectorAdapter async implementation..."
rg -A 5 "async def" cognee/infrastructure/databases/vector/pgvector/PGVectorAdapter.py

# Let's also check how client operations are handled in PGVectorAdapter
echo "Checking client operations in PGVectorAdapter..."
rg "self\." cognee/infrastructure/databases/vector/pgvector/PGVectorAdapter.py

# Let's examine the interface requirements
echo "Checking VectorDBInterface..."
rg -A 2 "async def" cognee/infrastructure/databases/vector/vector_db_interface.py

Length of output: 6003

cognee/modules/pipelines/operations/run_tasks.py (3)

10-10: Update logger string to match new type hint

The type hint has been updated to use list[Task], but the logger string still shows the old format [Task]. Update it for consistency:

-logger = logging.getLogger("run_tasks(tasks: [Task], data)")
+logger = logging.getLogger("run_tasks(tasks: list[Task], data)")

Also applies to: 8-8

---

Line range hint 25-27: Standardize telemetry call parameter naming

There are inconsistencies in how telemetry parameters are passed:

# Some calls use named parameters
send_telemetry("Generator Task Completed", user_id=user.id, additional_properties={...})

# Others use positional parameters
send_telemetry("Async Generator Task Started", user.id, {...})

Standardize all telemetry calls to use named parameters for clarity:

send_telemetry(
    event_name,
    user_id=user.id,
    additional_properties={...}
)

Also applies to: 84-86, 124-126

---

Line range hint 22-147: Consider refactoring task execution patterns

The code contains significant duplication in error handling, logging, and telemetry across different task types. This makes maintenance more difficult and increases the chance of inconsistencies.

Consider introducing helper functions to handle common patterns:

async def execute_task(task_type: str, task: Task, args: list, next_processor, user: User):
    logger.info(f"{task_type} task started: `{task.executable.__name__}`")
    send_telemetry(f"{task_type} Task Started", user_id=user.id, additional_properties={
        "task_name": task.executable.__name__
    })
    
    try:
        result = await process_task(task_type, task, args)
        await next_processor(result)
        
        logger.info(f"{task_type} task completed: `{task.executable.__name__}`")
        send_telemetry(f"{task_type} Task Completed", user_id=user.id, additional_properties={
            "task_name": task.executable.__name__
        })
    except Exception as error:
        logger.error(
            f"{task_type} task errored: `{task.executable.__name__}`\n{str(error)}\n",
            exc_info=True
        )
        send_telemetry(f"{task_type} Task Errored", user_id=user.id, additional_properties={
            "task_name": task.executable.__name__
        })
        raise

This would significantly reduce code duplication and make the code more maintainable.

cognee/tasks/graph/infer_data_ontology.py (4)

Line range hint 89-157: Consider refactoring the add_graph_ontology method

The method is handling two distinct flows (file-based and document-based) which makes it complex and harder to maintain. Consider splitting it into two separate methods for better maintainability and testability.

Here's a suggested refactoring approach:

- async def add_graph_ontology(self, file_path: str = None, documents: list = None):
+ async def add_graph_ontology_from_documents(self, documents: list) -> str:
+     """Add graph ontology from documents content."""
+     initial_chunks_and_ids = await self._process_documents(documents)
+     ontology = await extract_ontology(str(initial_chunks_and_ids), GraphOntology)
+     await self._add_nodes_and_edges(ontology)
+     return ontology.root_node_id
+
+ async def add_graph_ontology_from_file(self, file_path: str, documents: list):
+     """Add graph ontology from a JSON or CSV file."""
+     dataset_level_information = documents[0][1]
+     valid_ids = {item["id"] for item in dataset_level_information}
+     data = await self.load_data(file_path)
+     await self._process_file_data(data, valid_ids)

Also, consider extracting common node/edge property creation logic into helper methods to reduce code duplication.

---

Line range hint 64-81: Enhance error handling in load_data method

While the error handling has been improved, consider making it more specific by:

1. Using custom exceptions for different error cases
2. Adding validation for empty files
3. Moving file type validation to a separate method

Here's a suggested improvement:

class OntologyFileError(Exception):
    """Base exception for ontology file operations."""
    pass

class UnsupportedFileTypeError(OntologyFileError):
    """Exception raised for unsupported file types."""
    pass

class EmptyFileError(OntologyFileError):
    """Exception raised for empty files."""
    pass

async def validate_file_type(file_path: str) -> str:
    if not file_path.endswith(('.json', '.csv')):
        raise UnsupportedFileTypeError(f"Unsupported file format: {file_path}")
    return file_path.split('.')[-1]

async def load_data(self, file_path: str) -> Union[List[Dict[str, Any]], Dict[str, Any]]:
    """Load data from a JSON or CSV file."""
    try:
        file_type = await self.validate_file_type(file_path)
        async with aiofiles.open(file_path, mode="r") as f:
            content = await f.read()
            if not content.strip():
                raise EmptyFileError(f"File is empty: {file_path}")
            
            if file_type == "json":
                return json.loads(content)
            else:  # csv
                reader = csv.DictReader(content.splitlines())
                return list(reader)
    except OntologyFileError as e:
        raise
    except Exception as e:
        raise RuntimeError(f"Failed to load data from {file_path}: {e}")

---

Line range hint 134-137: Improve node ID validation

While the validation check is good, consider:

1. Validating all node IDs upfront before processing
2. Providing more informative error messages
3. Adding logging for skipped/invalid nodes

Here's a suggested improvement:

# Add this method to OntologyEngine class
async def validate_node_ids(self, nodes: pd.DataFrame, valid_ids: set) -> None:
    """Validate all node IDs before processing."""
    invalid_ids = set(nodes['node_id'].unique()) - valid_ids
    if invalid_ids:
        logger.error(f"Found {len(invalid_ids)} invalid node IDs")
        raise ValueError(
            f"Invalid node IDs found: {invalid_ids}. "
            f"Valid IDs must be one of: {valid_ids}"
        )

Then update the validation in add_graph_ontology:

- if node_id not in valid_ids:
-     raise ValueError(f"Node ID {node_id} not found in the dataset")
+ # At the start of processing
+ await self.validate_node_ids(df, valid_ids)

---

Line range hint 1-157: Add comprehensive tests for ontology processing

Given the complexity of the ontology processing logic, consider adding comprehensive tests covering:

1. Different input formats (JSON, CSV)
2. Edge cases (empty files, invalid node IDs)
3. Both file-based and document-based processing paths
4. Error handling scenarios

Would you like me to help create a test suite for this module?

cognee/shared/utils.py (2)

139-142: Extract node size configuration to settings.

The keywords list for determining node sizes is hardcoded. Consider moving this configuration to a settings file or environment variables for better maintainability.

+# In config.py or settings.py
+NODE_SIZE_CONFIG = {
+    "default_size": 10,
+    "larger_size": 20,
+    "keywords": ["DOCUMENT", "User"]
+}

 # In current file
-            keywords = ["DOCUMENT", "User"]
-            node_size = larger_size if any(keyword in str(node) for keyword in keywords) else default_size
+            node_size = NODE_SIZE_CONFIG["larger_size"] if any(
+                keyword in str(node) for keyword in NODE_SIZE_CONFIG["keywords"]
+            ) else NODE_SIZE_CONFIG["default_size"]

---

175-179: Remove commented out code and redundant pass statement.

The commented code for color encoding and the redundant pass statement should be removed if this functionality is no longer needed. If the color functionality will be implemented later, consider adding a TODO comment with a ticket reference.

         if include_color:
-            pass
-            # unique_layers = nodes["layer_description"].unique()
-            # color_palette = generate_color_palette(unique_layers)
-            # plotter = plotter.encode_point_color("layer_description", categorical_mapping=color_palette,
-            #                                      default_mapping="silver")
+            # TODO: Implement color encoding in ticket COG-XXX
+            pass

cognee/modules/data/operations/translate_text.py (4)

17-18: Move imports to the top of the file for better organization.

Importing modules inside a function can lead to repeated imports and is generally discouraged unless necessary. Consider moving the boto3 and botocore.exceptions imports to the module level.

Apply this diff:

+import boto3
+from botocore.exceptions import BotoCoreError, ClientError

...

-    import boto3
-    from botocore.exceptions import BotoCoreError, ClientError

---

27-33: Handle AWS Translate response more explicitly.

Using result.get("TranslatedText", "No translation found.") may mask issues if the TranslatedText key is missing. Since AWS Translate should always return TranslatedText on success, consider accessing it directly and handling any potential exceptions.

Apply this diff:

...

-        return result.get("TranslatedText", "No translation found.")
+        return result["TranslatedText"]

...

+    except KeyError as e:
+        logger.error(f"KeyError occurred: {e}")
+        return None

---

35-41: Improve error handling by raising exceptions.

Returning None on exceptions may make it difficult for the calling code to handle errors appropriately. Consider re-raising the exceptions after logging, allowing the calling code to deal with them directly.

Apply this diff:

...

     except BotoCoreError as e:
         logger.error(f"BotoCoreError occurred: {e}")
-        return None
+        raise

     except ClientError as e:
         logger.error(f"ClientError occurred: {e}")
-        return None
+        raise

---

5-15: Document exceptions in the docstring.

Including an "Exceptions" section in the docstring helps users understand what errors to expect and how to handle them.

Update the docstring as follows:

     Returns:
-    str: Translated text or an error message.
+    str: Translated text.

+    Raises:
+        ValueError: If input validation fails.
+        BotoCoreError: If an error occurs within the AWS SDK.
+        ClientError: If AWS Translate returns an error.

cognee/shared/SourceCodeGraph.py (2)

31-31: Inconsistent use of List and list in type annotations

At line 31, has_methods is annotated using list["Function"] instead of List["Function"]. For consistency and clarity, it's recommended to use List from the typing module for type annotations.

Apply this diff to fix the inconsistency:

 class Class(DataPoint):
     id: str
     name: str
     type: Literal["Class"] = "Class"
     description: str
     constructor_parameters: List[Variable]
     extended_from_class: Optional["Class"] = None
-    has_methods: list["Function"]
+    has_methods: List["Function"]

---

43-43: Unnecessary use of Union with a single type

At line 43, instantiated_by is annotated as Union["Function"]. Since there is only one type, using Union is unnecessary. If instantiated_by can only be a Function, you can annotate it directly. If it can be None, consider using Optional["Function"].

Apply this diff to correct the type annotation:

 class ClassInstance(DataPoint):
     id: str
     name: str
     type: Literal["ClassInstance"] = "ClassInstance"
     description: str
     from_class: Class
-    instantiated_by: Union["Function"]
+    instantiated_by: "Function"

cognee/modules/chunking/TextChunker.py (1)

50-50: Simplify list indexing with negative indices

Instead of using self.paragraph_chunks[len(self.paragraph_chunks) - 1], you can use self.paragraph_chunks[-1] for a more idiomatic and concise approach.

Apply this diff:

 # At line 50
-                    cut_type = self.paragraph_chunks[len(self.paragraph_chunks) - 1]["cut_type"],
+                    cut_type = self.paragraph_chunks[-1]["cut_type"],

 # At line 67
-                    cut_type = self.paragraph_chunks[len(self.paragraph_chunks) - 1]["cut_type"],
+                    cut_type = self.paragraph_chunks[-1]["cut_type"],

Also applies to: 67-67

cognee/modules/graph/utils/get_graph_from_model.py (5)

27-31: Use tuples instead of string concatenation for edge_key to ensure uniqueness and improve performance.

Currently, edge_key is constructed by concatenating strings:

edge_key = str(edge[0]) + str(edge[1]) + edge[2]

This approach can lead to key collisions and is less efficient than using tuples as dictionary keys.

Consider using a tuple:

-edge_key = str(edge[0]) + str(edge[1]) + edge[2]
+edge_key = (edge[0], edge[1], edge[2])

Update the if condition and the dictionary assignment accordingly:

-if str(edge_key) not in added_edges:
+if edge_key not in added_edges:
    edges.append(edge)
-   added_edges[str(edge_key)] = True
+   added_edges[edge_key] = True

---

65-68: Consistently use tuples for edge_key to avoid key collisions and improve code clarity.

Again, constructing edge_key by concatenating strings may lead to collisions and is less efficient. Recommend using tuples as keys.

Update the code as follows:

-edge_key = str(data_point.id) + str(property_node.id) + field_name
+edge_key = (data_point.id, property_node.id, field_name)

And update the if condition and the dictionary assignment:

-if str(edge_key) not in added_edges:
+if edge_key not in added_edges:
    edges.append((
        data_point.id,
        property_node.id,
        field_name,
        {
            "source_node_id": data_point.id,
            "target_node_id": property_node.id,
            "relationship_name": field_name,
            "updated_at": datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M:%S"),
            "metadata": {
                "type": "list"
            },
        }
    ))
    added_edges[edge_key] = True

---

41-41: Use ISO 8601 format for updated_at timestamps for consistency and precision.

Currently, the timestamp is formatted without time zone information:

"updated_at": datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M:%S"),

Consider using .isoformat() to include time zone information and ensure consistency:

-"updated_at": datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M:%S"),
+"updated_at": datetime.now(timezone.utc).isoformat(),

---

72-72: Use ISO 8601 format for updated_at timestamps for consistency and precision.

As above, update the timestamp formatting to include time zone information.

Apply this diff:

-"updated_at": datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M:%S"),
+"updated_at": datetime.now(timezone.utc).isoformat(),

---

99-100: Use a set for destination_nodes to improve lookup performance.

Currently, destination_nodes is a list, and checking if str(node.id) in destination_nodes is an O(N) operation.

Convert destination_nodes to a set for O(1) lookup:

-destination_nodes = [str(property_edge[1]) for property_edge in property_edges]
+destination_nodes = {str(property_edge[1]) for property_edge in property_edges}

cognee/api/v1/cognify/code_graph_pipeline.py (1)

90-91: Replace print statements with proper logging

Using print() statements in production code is not recommended. Consider using the logger to log pipeline results for better control over output and log levels.

Apply this diff to replace print() with logger.info():

     async for result in pipeline:
-        print(result)
+        logger.info(result)

cognee/tasks/graph/extract_graph_from_data.py (2)

27-35: Optimize node processing by using sets

Using a dictionary processed_nodes with string keys and True values is less efficient and may be less readable. Consider using a set to track processed node IDs.

Refactor as follows:

-    processed_nodes = {}
+    processed_node_ids = set()

...

-        if str(type_node_id) not in processed_nodes:
+        if str(type_node_id) not in processed_node_ids:

...

-            processed_nodes[str(type_node_id)] = True
+            processed_node_ids.add(str(type_node_id))

...

-        if str(entity_node_id) not in processed_nodes:
+        if str(entity_node_id) not in processed_node_ids):

...

-            processed_nodes[str(entity_node_id)] = True
+            processed_node_ids.add(str(entity_node_id))

---

12-14: Handle exceptions in asynchronous graph extraction

When using asyncio.gather, if one of the coroutines raises an exception, it cancels the remaining tasks, which may not be desirable. To prevent this, consider using return_exceptions=True or handling exceptions within extract_content_graph.

Apply this diff to handle exceptions:

-    chunk_graphs = await asyncio.gather(
+    chunk_graphs = await asyncio.gather(
         *[extract_content_graph(chunk.text, graph_model) for chunk in data_chunks]
-    )
+    , return_exceptions=True
+    )

Then, handle exceptions in the results:

for idx, result in enumerate(chunk_graphs):
    if isinstance(result, Exception):
        # Handle the exception (e.g., log it)
        chunk_graphs[idx] = None  # Or handle as appropriate

cognee/infrastructure/databases/vector/qdrant/QDrantAdapter.py (3)

13-19: Clarify the use of _metadata as a class attribute

The _metadata dictionary is declared as a class attribute in IndexSchema. If _metadata is intended to store instance-specific information, it should be defined as an instance attribute to prevent shared state across instances. Additionally, leading underscores typically denote private attributes. Consider adjusting its usage for clarity and to avoid potential issues.

Consider defining _metadata within the __init__ method:

 class IndexSchema(DataPoint):
     text: str

-    _metadata: dict = {
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self._metadata = {
             "index_fields": ["text"]
         }

---

128-129: Validate collection naming to avoid conflicts

Using f"{index_name}_{index_property_name}" as the collection name could lead to naming collisions if index_name or index_property_name contain underscores or duplicate values. Ensure that the combination of these parameters results in unique collection names and consider sanitizing or validating the inputs.

You might implement a standardized naming convention or use a delimiter less likely to appear in names:

 await self.create_collection(f"{index_name}__{index_property_name}")

---

2-2: Remove unused imports if not necessary

Ensure that all imported modules, such as UUID, are being utilized in the code. If UUID is no longer needed due to changes in the codebase, consider removing the import to keep the code clean.

cognee/infrastructure/databases/vector/lancedb/LanceDBAdapter.py (3)

76-77: Avoid reassigning payload_schema to enhance clarity

Reassigning payload_schema after calling self.get_data_point_schema(payload_schema) can be confusing. Consider using a new variable name for the transformed schema to improve code readability.

---

100-109: Simplify dynamic typing in create_lance_data_point

The use of generics with LanceDataPoint[str, self.get_data_point_schema(type(data_point))] introduces complexity and may lead to type-checking issues. Defining LanceDataPoint with explicit types can enhance maintainability.

---

167-176: Re-evaluate normalization when distances are equal

Setting all normalized values to 0 when max_value == min_value might not accurately reflect similarity. Consider assigning a neutral value like 0.5 to indicate equal distances or maintain the original distances to preserve relative differences.

cognee/infrastructure/databases/vector/pgvector/PGVectorAdapter.py (1)

Line range hint 91-96: Unused parameter payload_schema in create_collection

The payload_schema parameter is passed to create_collection, but it is not utilized within the method. To maintain clean and understandable code, consider removing the unused parameter or implementing its functionality if it is intended for future use.

cognee/infrastructure/databases/hybrid/falkordb/FalkorDBAdapter.py (2)

153-153: Remove unnecessary f prefix in string

The string at line 153 is marked as an f-string but doesn't contain any placeholders. The f prefix is unnecessary and can be removed.

Apply this diff to fix the issue:

-            f"MATCH (node) WHERE node.id IN $node_ids RETURN node",
+            "MATCH (node) WHERE node.id IN $node_ids RETURN node",

🧰 Tools

🪛 Ruff

153-153: f-string without any placeholders

Remove extraneous f prefix

(F541)

---

241-241: Remove unnecessary f prefix in string

The string at line 241 is marked as an f-string but doesn't contain any placeholders. Removing the f prefix cleans up the code.

Apply this diff to fix the issue:

-            f"MATCH (node) WHERE node.id IN $node_ids DETACH DELETE node",
+            "MATCH (node) WHERE node.id IN $node_ids DETACH DELETE node",

🧰 Tools

🪛 Ruff

241-241: f-string without any placeholders

Remove extraneous f prefix

(F541)

cognee/infrastructure/databases/graph/networkx/adapter.py (2)

Line range hint 138-154: Ensure consistent use of node_id parameter types across methods

In the methods get_predecessors and get_successors, the node_id parameter is typed as UUID, while other methods like has_node, delete_node, and get_edges use node_id: str. This inconsistency could lead to type errors or unexpected behavior. Consider standardizing the node_id parameter to use UUID across all methods for consistency.

For example, update the has_node method signature:

-    async def has_node(self, node_id: str) -> bool:
+    async def has_node(self, node_id: UUID) -> bool:

And ensure all usages of node_id are appropriately updated.

---

49-49: Optimize file save operations to improve performance

Calling await self.save_graph_to_file(self.filename) after every node or edge addition or deletion can lead to performance bottlenecks due to frequent disk I/O operations. Consider batching these operations or implementing a buffering mechanism to reduce the number of writes.

One approach is to save the graph at strategic points or after a certain number of operations. Alternatively, you could introduce a flag to indicate when the graph has been modified and needs to be saved.

Example modification:

# Remove individual save calls
-    await self.save_graph_to_file(self.filename)

# Implement a periodic save mechanism or manual trigger

This change would require adjusting the workflow to ensure data integrity without compromising performance.

Also applies to: 56-56, 97-97, 108-108

cognee/infrastructure/databases/graph/neo4j_driver/adapter.py (2)

204-205: Avoid redundant conversion and ensure consistent typing of node identifiers

In the add_edges method, you are converting edge[0] and edge[1] to strings multiple times. Ensure that node identifiers are consistently handled as strings throughout the method to improve clarity and prevent potential type-related issues.

Refactor the code to reduce redundancy:

edges = [{
-  "from_node": str(edge[0]),
-  "to_node": str(edge[1]),
+  "from_node": edge_id_strs[0],
+  "to_node": edge_id_strs[1],
   "relationship_name": edge[2],
   "properties": {
       **(edge[3] if edge[3] else {}),
-      "source_node_id": str(edge[0]),
-      "target_node_id": str(edge[1]),
+      "source_node_id": edge_id_strs[0],
+      "target_node_id": edge_id_strs[1],
   },
} for edge_id_strs, edge in [ (list(map(str, edge[:2])), edge) for edge in edges ]]

This refactoring ensures that all node IDs are converted to strings once and reused, enhancing code readability.

Also applies to: 209-210

---

400-409: Efficiently serialize properties without unnecessary empty lines

The serialize_properties method contains several unnecessary empty lines that can be removed to improve readability.

Apply this diff to clean up the method:

 def serialize_properties(self, properties = dict()):
     serialized_properties = {}
-    
     for property_key, property_value in properties.items():
         if isinstance(property_value, UUID):
             serialized_properties[property_key] = str(property_value)
             continue
         serialized_properties[property_key] = property_value
-    
     return serialized_properties

Configuration used: CodeRabbit UI
Review profile: CHILL

Learnt from: borisarzentar
PR: topoteretes/cognee#144
File: cognee/tasks/chunking/query_chunks.py:1-17
Timestamp: 2024-10-08T16:12:50.908Z
Learning: The `query_chunks` function in `cognee/tasks/chunking/query_chunks.py` is used within the `search` function in `cognee/api/v1/search/search_v2.py`.

Learnt from: borisarzentar
PR: topoteretes/cognee#144
File: cognee/tasks/chunking/query_chunks.py:1-17
Timestamp: 2024-10-08T16:12:49.108Z
Learning: The `query_chunks` function in `cognee/tasks/chunking/query_chunks.py` is used within the `search` function in `cognee/api/v1/search/search_v2.py`.

Learnt from: borisarzentar
PR: topoteretes/cognee#144
File: cognee/tasks/chunking/query_chunks.py:1-17
Timestamp: 2024-10-07T11:20:44.876Z
Learning: The `query_chunks` function in `cognee/tasks/chunking/query_chunks.py` is used within the `search` function in `cognee/api/v1/search/search_v2.py`.

#!/bin/bash
# Search for FalkorDB node naming patterns or constraints in the codebase
rg -l "FalkorDB" | xargs rg -i "node.*name|label.*name"

# Look for any existing node creation patterns
ast-grep --pattern 'CREATE ($node_name)'

#!/bin/bash
# Look for node creation and naming patterns in FalkorDB adapter
rg -A 5 -B 5 "CREATE" cognee/infrastructure/databases/hybrid/falkordb/FalkorDBAdapter.py

# Check for any node name validation or processing
rg -A 3 -B 3 "node.*name|label.*name" cognee/infrastructure/databases/hybrid/falkordb/FalkorDBAdapter.py

# Look for any error handling related to node names
rg -i "invalid.*name|name.*invalid" cognee/infrastructure/databases/hybrid/falkordb/FalkorDBAdapter.py

#!/bin/bash
# Look for any string escaping or sanitization in FalkorDB queries
rg -A 5 "stringify|escape|sanitize|replace" cognee/infrastructure/databases/hybrid/falkordb/FalkorDBAdapter.py

# Check if there are any FalkorDB specific configuration or constants
cat cognee/infrastructure/databases/hybrid/falkordb/FalkorDBAdapter.py | head -n 30

#!/bin/bash
# Description: Compare edge name generation with node name generation

# Look for node name generation implementation
echo "Searching for node name generation..."
rg "generate_node_name" -A 5

# Look for other name generation utilities
echo "Searching for similar utilities..."
fd -g "*.py" -x grep -l "generate.*name"

#!/bin/bash
# Search for any other node ID generation implementations or usages
echo "Searching for other node ID generation patterns..."
rg -l "generate.*node.*id" --type py

echo "\nSearching for direct UUID5 usage with NAMESPACE_OID..."
rg -l "uuid5.*NAMESPACE_OID" --type py

echo "\nSearching for potential test files..."
fd "test.*node.*" --type py

#!/bin/bash
# Let's examine the actual implementation in the files that use uuid5 with NAMESPACE_OID
echo "Checking uuid5 implementations in other files:"
rg -A 3 "uuid5.*NAMESPACE_OID" --type py

echo "\nChecking how generate_node_id is imported and used:"
rg -A 2 "from.*generate_node_id import|import.*generate_node_id" --type py

echo "\nChecking actual usage of generate_node_id function:"
rg -A 2 "generate_node_id\(" --type py

#!/bin/bash
# Search for any remaining references to document.id or document.title
rg -A 2 "document\.(?:id|title)\b" 

# Search for references to the new name attribute
rg -A 2 "document\.name\b"

#!/bin/bash
# Description: Find potential places that might need updates due to the UUID change

# Look for direct usage of ScoredResult
echo "=== Direct ScoredResult usage ==="
rg -l "ScoredResult" --type py

# Look for string IDs being passed that might need conversion to UUID
echo -e "\n=== Potential string ID usage ==="
ast-grep --pattern 'id: str'

# Look for potential serialization/deserialization of ScoredResult
echo -e "\n=== Potential serialization points ==="
rg "json|dumps|loads.*ScoredResult" --type py

#!/bin/bash
# Check how the id field is handled in adapter implementations
echo "=== QDrant Adapter Implementation ==="
rg -A 5 "ScoredResult" cognee/infrastructure/databases/vector/qdrant/QDrantAdapter.py

echo -e "\n=== PGVector Adapter Implementation ==="
rg -A 5 "ScoredResult" cognee/infrastructure/databases/vector/pgvector/PGVectorAdapter.py

echo -e "\n=== Weaviate Adapter Implementation ==="
rg -A 5 "ScoredResult" cognee/infrastructure/databases/vector/weaviate_db/WeaviateAdapter.py

echo -e "\n=== LanceDB Adapter Implementation ==="
rg -A 5 "ScoredResult" cognee/infrastructure/databases/vector/lancedb/LanceDBAdapter.py

#!/bin/bash
# Search for Document class implementation and its read method
ast-grep --pattern 'class Document {
  $$$
  read($$$) {
    $$$
  }
  $$$
}'

# Backup: Search using ripgrep
rg -A 10 "class Document.*:|def read.*:"

#!/bin/bash
# Description: Check for other EntityType-related index configurations in the codebase
# and FalkorDB integration code

# Look for FalkorDB index configurations
rg -l "index_fields|createIndex|CREATE INDEX" --type py

# Look for similar DataPoint implementations
ast-grep --pattern 'class $_ (DataPoint) {
  $$$
  _metadata: dict = {
    $$$
  }
}'

#!/bin/bash
# Check FalkorDB adapter implementation for index handling
cat cognee/infrastructure/databases/hybrid/falkordb/FalkorDBAdapter.py

# Check base DataPoint class for index field handling
cat cognee/infrastructure/engine/models/DataPoint.py

# Check Entity model since EntityType is related to it
cat cognee/modules/engine/models/Entity.py

#!/bin/bash
# Check for potential circular imports
echo "Checking EntityType.py for potential circular imports with Entity:"
rg -l "from.*Entity import" "$(dirname $(fd -t f "EntityType.py$"))"

#!/bin/bash
# Let's check the actual content of both files to verify the imports
echo "Content of __init__.py:"
cat "$(fd -t f "__init__.py$" cognee/modules/engine/models)"

echo -e "\nContent of EntityType.py:"
cat "$(fd -t f "EntityType.py$")"

#!/bin/bash
# Description: Check for any remaining BaseModel usage that should be DataPoint
# and verify consistent DataPoint usage in related classes

# Check for any remaining BaseModel usage in chunking modules
rg -l "class.*\(BaseModel\)" "cognee/modules/chunking"

# Check DataPoint usage in related document processing classes
rg -l "class.*\(DataPoint\)" "cognee/modules/data/processing"

#!/bin/bash
# Let's check the actual implementation of both DocumentChunk and Document classes
# to understand the inheritance pattern better

# Check DocumentChunk implementation
echo "=== DocumentChunk.py ==="
cat cognee/modules/chunking/models/DocumentChunk.py

echo -e "\n=== Document.py ==="
cat cognee/modules/data/processing/document_types/Document.py

# Also check if there are any other DataPoint usages
echo -e "\n=== Other DataPoint usages ==="
rg "class.*\(DataPoint\)" -A 2

#!/bin/bash
# Description: Check TextChunker usage across all document types
# Expected: All document types should use the new interface

# Check TextChunker instantiation pattern in all document types
rg -A 2 "TextChunker\(" --type py

# Verify TextChunker class definition accepts document instance
ast-grep --pattern 'class TextChunker {
  $$$
  def __init__($_, document, $_) {
    $$$
  }
  $$$
}'

#!/bin/bash
# Description: Check if Document base class provides raw_data_location attribute
# Expected: Find the attribute definition in Document class

ast-grep --pattern 'class Document {
  $$$
  raw_data_location
  $$$
}'

# Also check for property decorator or __init__ assignment
rg -A 5 "raw_data_location" --type py

#!/bin/bash
# Search for Document base class definition and its attributes
ast-grep --pattern 'class Document {
  $$$
}'

# Also check for __init__ method and property definitions
rg -A 10 "class Document" --type py

#!/bin/bash
# Description: Check for references to DocumentChunk_text collection and its setup

# Check for collection setup/creation
rg -A 5 "DocumentChunk_text" 

# Check for vector engine implementation details
ast-grep --pattern 'class $_ implements VectorEngine {
  $$$
  search($collection, $query, $$$) {
    $$$
  }
  $$$
}'

#!/bin/bash
# Check for collection initialization or setup
rg -l "DocumentChunk" 

# Check for vector engine implementations and collection setup
ast-grep --pattern 'class DocumentChunk {
  $$$
}'

# Check for any migration or setup scripts
fd -e py -e sql | xargs rg -l "collection|index.*DocumentChunk"

#!/bin/bash
# Check vector database adapters for collection setup
rg -A 10 "DocumentChunk.*text" cognee/infrastructure/databases/vector/

# Check the base interface
cat cognee/infrastructure/databases/vector/vector_db_interface.py

# Look for collection setup in database adapters
rg -A 5 "create.*collection|setup.*collection" cognee/infrastructure/databases/vector/

#!/bin/bash
# Check DocumentChunk model and its text field
cat cognee/modules/chunking/models/DocumentChunk.py

# Check for collection initialization in startup or setup code
rg -l "init|setup|start" | xargs rg -l "DocumentChunk.*text"

# Check for any migration or initialization scripts
fd migration init setup | xargs rg -l "DocumentChunk"

#!/bin/bash
# Check Document parent class for raw_data_location initialization
ast-grep --pattern 'class Document {
  $$$
  def __init__($$$) {
    $$$
    self.raw_data_location = $$$
    $$$
  }
  $$$
}'

# Check for any other initialization patterns
rg -l "raw_data_location" --type py

#!/bin/bash
# Check Document class implementation for raw_data_location
rg -A 10 "class Document" cognee/modules/data/processing/document_types/Document.py

# Check Data model class since it's referenced in multiple files
rg -A 10 "class Data" cognee/modules/data/models/Data.py

# Check if raw_data_location is passed during document creation
rg "raw_data_location" -A 3 -B 3

#!/bin/bash
# Description: Check TextChunker's constructor signature and usage patterns

# Search for TextChunker class definition and its usage
ast-grep --pattern 'class TextChunker {
  $$$
  __init__($$$) {
    $$$
  }
  $$$
}'

# Find other document types using TextChunker
rg "TextChunker\(" --type py

#!/bin/bash
# Let's check the TextChunker's constructor implementation
rg -A 10 "class TextChunker" cognee/modules/chunking/TextChunker.py

# Also check how the document instance is used within TextChunker
rg "self\." cognee/modules/chunking/TextChunker.py

#!/bin/bash
# Description: Check Document parent class implementation and usage
# Expected: Document class should initialize raw_data_location

# Check Document class implementation
ast-grep --pattern 'class Document {
  $$$
  def __init__($params) {
    $$$
  }
  $$$
}'

# Find all Document subclasses to ensure consistent pattern
rg -l "class \w+\(Document\)"

#!/bin/bash
# Description: Verify TextChunker implementation accepts Document instance
# Expected: TextChunker should have updated its constructor to accept Document instance

# Check TextChunker implementation
ast-grep --pattern 'class TextChunker {
  $$$
  def __init__($params) {
    $$$
  }
  $$$
}'

# Check for any other usages of TextChunker in the codebase
rg -l "TextChunker"

#!/bin/bash
# Check TextChunker's constructor implementation
ast-grep --pattern 'class TextChunker:
  $$$
  def __init__($params):
    $$$
  }' cognee/modules/chunking/TextChunker.py

# Get the full TextChunker class implementation for context
rg -A 10 "class TextChunker" cognee/modules/chunking/TextChunker.py

#!/bin/bash
# Find all classes that inherit from DataPoint
echo "Classes inheriting from DataPoint:"
ast-grep --pattern 'class $NAME($$$DataPoint$$$)' 

# Check for any direct access to _metadata
echo -e "\nDirect _metadata access:"
rg '\._metadata\b' --type python

# Check for get_embeddable_data usage
echo -e "\nget_embeddable_data usage:"
rg '\.get_embeddable_data\b' --type python

#!/bin/bash
# Find Python files first
echo "Python files in the repository:"
fd '\.py$'

# Then use the file paths to search for patterns
echo -e "\nClasses inheriting from DataPoint:"
fd '\.py$' --exec ast-grep --pattern 'class $NAME($$$DataPoint$$$)' {}

echo -e "\nDirect _metadata access:"
fd '\.py$' --exec rg '\._metadata\b' {}

echo -e "\nget_embeddable_data usage:"
fd '\.py$' --exec rg '\.get_embeddable_data\b' {}

 for edge in edges:
+    if len(edge) < 3:
+        raise ValueError(f"Edge must have at least 3 elements: {edge}")
+    
+    source_id, target_id = edge[0], edge[1]
+    if source_id not in node_map:
+        raise KeyError(f"Source node {source_id} not found in nodes")
+    if target_id not in node_map:
+        raise KeyError(f"Target node {target_id} not found in nodes")
+    
     source_node = node_map[edge[0]]
     target_node = node_map[edge[1]]
     edge_label = edge[2]
     edge_properties = edge[3] if len(edge) == 4 else {}
     edge_metadata = edge_properties.get("metadata", {})
     edge_type = edge_metadata.get("type")
+    
+    if edge_type and edge_type not in ["list", "single"]:
+        raise ValueError(f"Invalid edge type: {edge_type}")

#!/bin/bash
# Description: Find implementations of GraphDBInterface to verify query parameter usage

# Find potential implementation files
echo "Looking for GraphDBInterface implementations..."
rg -l "class.*[\(: ]GraphDBInterface" 

# Look for query method implementations
echo -e "\nChecking query method implementations..."
ast-grep --pattern 'class $CLASS {
  $$$
  async def query(self, query: str, params: dict):
    $IMPL
  $$$
}'

#!/bin/bash
# Check implementations in the identified files
echo "=== NetworkX Adapter Implementation ==="
rg -A 10 "async def query" cognee/infrastructure/databases/graph/networkx/adapter.py

echo -e "\n=== Neo4j Driver Implementation ==="
rg -A 10 "async def query" cognee/infrastructure/databases/graph/neo4j_driver/adapter.py

echo -e "\n=== FalkorDB Implementation ==="
rg -A 10 "async def query" cognee/infrastructure/databases/hybrid/falkordb/FalkorDBAdapter.py

#!/bin/bash
# Get full FalkorDB implementation to check query method
echo "=== FalkorDB Full Implementation ==="
cat cognee/infrastructure/databases/hybrid/falkordb/FalkorDBAdapter.py

#!/bin/bash
# Description: Check for NetworkX adapter usage across the codebase

# Search for NetworkX adapter imports and usage
echo "Searching for NetworkX adapter imports and usage..."
rg -l "NetworkXAdapter" 

# Search for graph file path configuration usage
echo "Searching for graph file path configuration usage..."
rg "graph_file_path"

#!/bin/bash
# Check NetworkXAdapter implementation and its dependencies
echo "Checking NetworkXAdapter implementation:"
cat cognee/infrastructure/databases/graph/networkx/adapter.py

# Check the get_graph_engine implementation to understand the fallback logic
echo -e "\nChecking get_graph_engine implementation:"
cat cognee/infrastructure/databases/graph/get_graph_engine.py

 def copy_model(model: DataPoint, include_fields: dict = {}, exclude_fields: list = []):
+    for field_type in include_fields.values():
+        if not isinstance(field_type, tuple) or len(field_type) != 2:
+            raise ValueError("include_fields values must be tuples of (type, default)")
     fields = {

#!/bin/bash
# Check if PGVectorAdapter has async methods
ast-grep --pattern 'class PGVectorAdapter {
  $$$
  async def $_($$$) {
    $$$
  }
  $$$
}'

# Check for error handling in get_relational_config
rg -A 5 "def get_relational_config"

#!/bin/bash
# Check PGVectorAdapter implementation
rg -A 20 "class PGVectorAdapter"

# Check if there are any async operations in the file
rg "async|await" --type py

# Check RelationalConfig implementation
rg -A 10 "class RelationalConfig"

#!/bin/bash
# Check if FalkorDBAdapter implements required interfaces
ast-grep --pattern 'class FalkorDBAdapter implements $_Interface {
  $$$
}'

# Verify required methods are implemented
ast-grep --pattern 'class FalkorDBAdapter {
  $$$
  def $_vector_method($$$) {
    $$$
  }
  $$$
}'

#!/bin/bash
# Let's first locate the FalkorDBAdapter file
fd "FalkorDBAdapter.py" --type f

# Then examine its content to verify interface implementation
rg -A 2 "^class FalkorDBAdapter" 

# Check for vector operation methods
rg -A 1 "def (add_vectors|search_vectors|delete_vectors)" 

# Look for any base class or interface it might be inheriting from
rg "class.*\(" --type py | grep -i "vector"

#!/bin/bash
# Let's examine the VectorDBInterface to see required methods
rg -A 10 "class VectorDBInterface" cognee/infrastructure/databases/vector/vector_db_interface.py

# Now check FalkorDBAdapter for these required methods
rg -B 1 -A 3 "def (add_vectors|search_vectors|delete_vectors)" cognee/infrastructure/databases/hybrid/falkordb/FalkorDBAdapter.py

#!/bin/bash
# Let's check the complete VectorDBInterface definition
cat cognee/infrastructure/databases/vector/vector_db_interface.py

# And then check the complete FalkorDBAdapter implementation
cat cognee/infrastructure/databases/hybrid/falkordb/FalkorDBAdapter.py

#!/bin/bash
# Description: Check for any remaining usage of embed_field and verify index_fields adoption

echo "Checking for remaining embed_field usage..."
rg "embed_field" --type py

echo "Verifying index_fields adoption..."
rg "index_fields" --type py

# Check for DataPoint instantiations
echo "Checking DataPoint instantiations..."
ast-grep --pattern 'DataPoint[$_]($_)'

#!/bin/bash
# Search for any remaining imports from the old path
rg "from cognee.modules.graph.utils import (generate_node_id|generate_node_name)"

# Search for imports from the new path to confirm consistency
rg "from cognee.modules.engine.utils import (generate_node_id|generate_node_name)"

#!/bin/bash
# Check if the utility functions exist in the new location
rg "def (generate_node_id|generate_node_name)" --type py

# Look for any other usage patterns of these functions
rg "(generate_node_id|generate_node_name)\(" --type py

#!/bin/bash
# Check imports in extract_graph_from_data.py
rg "from.*import.*(generate_node_id|generate_node_name)" cognee/tasks/graph/extract_graph_from_data.py

#!/bin/bash
# Description: Check for FalkorDB-specific edge property names in the codebase

# Search for FalkorDB edge property definitions
rg -l "FalkorDB|falkordb" | xargs rg -A 5 "edge.*property|relationship.*name"

# Search for graph visualization calls
rg "render_graph|prepare_edges" -A 5

#!/bin/bash
# Description: Verify that the 'id' field in copied models is of type 'str' and '_metadata' is excluded.

ast-grep --pattern '
class $_(BaseModel) {
  id: str
}
' --ignore '*.pyc' .

#!/bin/bash
# Let's first find the copy_model function definition
rg "def copy_model" -A 10

# Then let's search for BaseModel classes with id field
rg "class.*BaseModel.*id.*str" -A 5

# Also search for any usage of copy_model to understand the pattern
rg "copy_model\(" -A 3

#!/bin/bash
# Let's examine the DataPoint class and its implementations to understand the base schema
rg "class.*DataPoint" -A 10

# Let's also check if there are any test files that verify the copy_model functionality
rg "test.*copy_model" -A 10

cognee/tests/test_neo4j.py (1)

39-40: Make the test more robust and data-independent.

The test is tightly coupled to specific search terms and data. Consider these improvements:

1. Add assertions to verify the content of search results, not just their existence
2. Add error handling for failed searches
3. Make the test data-independent by using test fixtures or parameterization

Example improvement:

-    random_node = (await vector_engine.search("Entity_name", "Quantum computer"))[0]
-    random_node_name = random_node.payload["text"]
+    try:
+        search_results = await vector_engine.search("Entity_name", "Quantum computer")
+        assert len(search_results) > 0, "No search results found"
+        random_node = search_results[0]
+        assert "text" in random_node.payload, "Missing 'text' in payload"
+        random_node_name = random_node.payload["text"]
+        assert random_node_name, "Empty node name"
+    except Exception as e:
+        logging.error(f"Search failed: {str(e)}")
+        raise

cognee/infrastructure/databases/graph/neo4j_driver/adapter.py (1)

Line range hint 341-355: Optimize connection queries by combining them

Instead of running two separate queries for predecessors and successors, consider using a single query with UNION to reduce database round trips.

Consider this optimization:

 async def get_connections(self, node_id: UUID) -> list:
-    predecessors_query = """
-    MATCH (node)<-[relation]-(neighbour)
-    WHERE node.id = $node_id
-    RETURN neighbour, relation, node
-    """
-    successors_query = """
-    MATCH (node)-[relation]->(neighbour)
-    WHERE node.id = $node_id
-    RETURN node, relation, neighbour
-    """
+    query = """
+    MATCH (node)
+    WHERE node.id = $node_id
+    CALL {
+        WITH node
+        MATCH (node)<-[relation]-(neighbour)
+        RETURN neighbour, relation, node
+        UNION
+        MATCH (node)-[relation]->(neighbour)
+        RETURN node, relation, neighbour
+    }
+    RETURN neighbour, relation, node
+    """

-    predecessors, successors = await asyncio.gather(
-        self.query(predecessors_query, dict(node_id = str(node_id))),
-        self.query(successors_query, dict(node_id = str(node_id))),
-    )
+    results = await self.query(query, dict(node_id = str(node_id)))

cognee/infrastructure/databases/vector/weaviate_db/WeaviateAdapter.py (3)

13-18: Clarify the Purpose and Usage of IndexSchema

The IndexSchema class extends DataPoint and introduces a text field along with _metadata specifying index_fields. Ensure that this class aligns with your indexing strategy and that all necessary fields are correctly defined for Weaviate's schema requirements.

---

90-95: Avoid Mutating the Original Properties Dictionary

Modifying the properties dictionary in place might lead to unintended side effects if data_point is reused elsewhere. Consider creating a copy of the properties before making modifications to prevent potential issues.

Apply this diff to create a copy before modification:

- properties = data_point.model_dump()
+ properties = data_point.model_dump().copy()

if "id" in properties:
    properties["uuid"] = str(data_point.id)
    del properties["id"]

---

103-116: Catch Specific Exceptions During Data Insertion

In the create_data_points method, catching a generic Exception may hide unexpected errors. It's better to catch specific exceptions related to Weaviate operations to handle known error cases effectively.

Apply this diff to catch specific exceptions:

- except Exception as error:
+ except weaviate.WeaviateException as error:

Ensure that WeaviateException is imported from the Weaviate client library:

from weaviate import WeaviateException

Configuration used: CodeRabbit UI
Review profile: CHILL

#!/bin/bash
# Search for payload access patterns
echo "Searching for old 'name' payload access:"
rg -l "payload\[\"name\"\]|payload\['name'\]" --type python

echo -e "\nSearching for new 'text' payload access:"
rg -l "payload\[\"text\"\]|payload\['text'\]" --type python

# Look for any schema definitions
echo -e "\nSearching for payload schema definitions:"
rg -l "payload.*=.*{.*name|payload.*=.*{.*text" --type python