Docstring modules. (#877)

<!-- .github/pull_request_template.md -->

## Description
<!-- Provide a clear description of the changes in this PR -->

## DCO Affirmation
I affirm that all code in every commit of this pull request conforms to
the terms of the Topoteretes Developer Certificate of Origin.

Co-authored-by: Vasilije <8619304+Vasilije1990@users.noreply.github.com>

cognee/modules/chunking/models/DocumentChunk.py

@@ -6,6 +6,25 @@
 
 
 class DocumentChunk(DataPoint):
+    """
+    Represents a chunk of text from a document with associated metadata.
+
+    Public methods include:
+
+    - No public methods defined in the provided code.
+
+    Instance variables include:
+
+    - text: The textual content of the chunk.
+    - chunk_size: The size of the chunk.
+    - chunk_index: The index of the chunk in the original document.
+    - cut_type: The type of cut that defined this chunk.
+    - is_part_of: The document to which this chunk belongs.
+    - contains: A list of entities contained within the chunk (default is None).
+    - metadata: A dictionary to hold meta information related to the chunk, including index
+    fields.
+    """
+
     text: str
     chunk_size: int
     chunk_index: int

cognee/modules/engine/models/EntityType.py

@@ -2,6 +2,14 @@
 
 
 class EntityType(DataPoint):
+    """
+    Represents a type of entity with a name and description.
+
+    This class inherits from DataPoint and includes two primary attributes: `name` and
+    `description`. Additionally, it contains a metadata dictionary that specifies
+    `index_fields` for indexing purposes.
+    """
+
     name: str
     description: str
 

cognee/modules/engine/operations/setup.py

@@ -7,5 +7,11 @@
 
 
 async def setup():
+    """
+    Set up the necessary databases and tables.
+
+    This function asynchronously creates a relational database and its corresponding tables,
+    followed by creating a PGVector database and its tables.
+    """
     await create_relational_db_and_tables()
     await create_pgvector_db_and_tables()

cognee/modules/retrieval/EntityCompletionRetriever.py

@@ -11,7 +11,21 @@
 
 
 class EntityCompletionRetriever(BaseRetriever):
-    """Retriever that uses entity-based completion for generating responses."""
+    """
+    Retriever that uses entity-based completion for generating responses.
+
+    Public methods:
+
+    - get_context
+    - get_completion
+
+    Instance variables:
+
+    - extractor
+    - context_provider
+    - user_prompt_path
+    - system_prompt_path
+    """
 
     def __init__(
         self,
@@ -26,7 +40,24 @@ def __init__(
         self.system_prompt_path = system_prompt_path
 
     async def get_context(self, query: str) -> Any:
-        """Get context using entity extraction and context provider."""
+        """
+        Get context using entity extraction and context provider.
+
+        Logs the processing of the query and retrieves entities. If entities are extracted, it
+        attempts to retrieve the corresponding context using the context provider. Returns None
+        if no entities or context are found, or logs the error if an exception occurs.
+
+        Parameters:
+        -----------
+
+            - query (str): The query string for which context is being retrieved.
+
+        Returns:
+        --------
+
+            - Any: The context retrieved from the context provider or None if not found or an
+              error occurred.
+        """
         try:
             logger.info(f"Processing query: {query[:100]}")
 
@@ -47,7 +78,26 @@ async def get_context(self, query: str) -> Any:
             return None
 
     async def get_completion(self, query: str, context: Optional[Any] = None) -> List[str]:
-        """Generate completion using provided context or fetch new context."""
+        """
+        Generate completion using provided context or fetch new context.
+
+        If context is not provided, it fetches context using the query. If no context is
+        available, it returns an error message. Logs an error if completion generation fails due
+        to an exception.
+
+        Parameters:
+        -----------
+
+            - query (str): The query string for which completion is being generated.
+            - context (Optional[Any]): Optional context to be used for generating completion;
+              fetched if not provided. (default None)
+
+        Returns:
+        --------
+
+            - List[str]: A list containing the generated completion or an error message if no
+              relevant entities were found.
+        """
         try:
             if context is None:
                 context = await self.get_context(query)

cognee/modules/retrieval/chunks_retriever.py

@@ -7,7 +7,16 @@
 
 
 class ChunksRetriever(BaseRetriever):
-    """Retriever for handling document chunk-based searches."""
+    """
+    Handles document chunk-based searches by retrieving relevant chunks and generating
+    completions from them.
+
+    Public methods:
+
+    - get_context: Retrieves document chunks based on a query.
+    - get_completion: Generates a completion using provided context or retrieves context if
+    not given.
+    """
 
     def __init__(
         self,
@@ -16,7 +25,22 @@ def __init__(
         self.top_k = top_k
 
     async def get_context(self, query: str) -> Any:
-        """Retrieves document chunks context based on the query."""
+        """
+        Retrieves document chunks context based on the query.
+
+        Searches for document chunks relevant to the specified query using a vector engine.
+        Raises a NoDataError if no data is found in the system.
+
+        Parameters:
+        -----------
+
+            - query (str): The query string to search for relevant document chunks.
+
+        Returns:
+        --------
+
+            - Any: A list of document chunk payloads retrieved from the search.
+        """
         vector_engine = get_vector_engine()
 
         try:
@@ -27,7 +51,25 @@ async def get_context(self, query: str) -> Any:
         return [result.payload for result in found_chunks]
 
     async def get_completion(self, query: str, context: Optional[Any] = None) -> Any:
-        """Generates a completion using document chunks context."""
+        """
+        Generates a completion using document chunks context.
+
+        If the context is not provided, it retrieves the context based on the query. Returns the
+        context, which can be used for further processing or generation of outputs.
+
+        Parameters:
+        -----------
+
+            - query (str): The query string to be used for generating a completion.
+            - context (Optional[Any]): Optional pre-fetched context to use for generating the
+              completion; if None, it retrieves the context for the query. (default None)
+
+        Returns:
+        --------
+
+            - Any: The context used for the completion or the retrieved context if none was
+              provided.
+        """
         if context is None:
             context = await self.get_context(query)
         return context

cognee/modules/retrieval/code_retriever.py

@@ -14,7 +14,12 @@ class CodeRetriever(BaseRetriever):
     """Retriever for handling code-based searches."""
 
     class CodeQueryInfo(BaseModel):
-        """Response model for information extraction from the query"""
+        """
+        Model for representing the result of a query related to code files.
+
+        This class holds a list of filenames and the corresponding source code extracted from a
+        query. It is used to encapsulate response data in a structured format.
+        """
 
         filenames: List[str] = []
         sourcecode: str

cognee/modules/retrieval/completion_retriever.py

@@ -8,7 +8,13 @@
 
 
 class CompletionRetriever(BaseRetriever):
-    """Retriever for handling LLM-based completion searches."""
+    """
+    Retriever for handling LLM-based completion searches.
+
+    Public methods:
+    - get_context(query: str) -> str
+    - get_completion(query: str, context: Optional[Any] = None) -> Any
+    """
 
     def __init__(
         self,
@@ -22,7 +28,24 @@ def __init__(
         self.top_k = top_k if top_k is not None else 1
 
     async def get_context(self, query: str) -> str:
-        """Retrieves relevant document chunks as context."""
+        """
+        Retrieves relevant document chunks as context.
+
+        Fetches document chunks based on a query from a vector engine and combines their text.
+        Returns empty string if no chunks are found. Raises NoDataError if the collection is not
+        found.
+
+        Parameters:
+        -----------
+
+            - query (str): The query string used to search for relevant document chunks.
+
+        Returns:
+        --------
+
+            - str: A string containing the combined text of the retrieved document chunks, or an
+              empty string if none are found.
+        """
         vector_engine = get_vector_engine()
 
         try:
@@ -38,7 +61,24 @@ async def get_context(self, query: str) -> str:
             raise NoDataError("No data found in the system, please add data first.") from error
 
     async def get_completion(self, query: str, context: Optional[Any] = None) -> Any:
-        """Generates an LLM completion using the context."""
+        """
+        Generates an LLM completion using the context.
+
+        Retrieves context if not provided and generates a completion based on the query and
+        context using an external completion generator.
+
+        Parameters:
+        -----------
+
+            - query (str): The input query for which the completion is generated.
+            - context (Optional[Any]): Optional context to use for generating the completion; if
+              not provided, it will be retrieved using get_context. (default None)
+
+        Returns:
+        --------
+
+            - Any: A list containing the generated completion from the LLM.
+        """
         if context is None:
             context = await self.get_context(query)
 

cognee/modules/retrieval/cypher_search_retriever.py

@@ -10,7 +10,13 @@
 
 
 class CypherSearchRetriever(BaseRetriever):
-    """Retriever for handling cypher-based search"""
+    """
+    Retriever for handling cypher-based search.
+
+    Public methods include:
+    - get_context: Retrieves relevant context using a cypher query.
+    - get_completion: Returns the graph connections context.
+    """
 
     def __init__(
         self,
@@ -22,7 +28,22 @@ def __init__(
         self.system_prompt_path = system_prompt_path
 
     async def get_context(self, query: str) -> Any:
-        """Retrieves relevant context using a cypher query."""
+        """
+        Retrieves relevant context using a cypher query.
+
+        If the graph engine is an instance of NetworkXAdapter, raises SearchTypeNotSupported. If
+        any error occurs during execution, logs the error and raises CypherSearchError.
+
+        Parameters:
+        -----------
+
+            - query (str): The cypher query used to retrieve context.
+
+        Returns:
+        --------
+
+            - Any: The result of the cypher query execution.
+        """
         try:
             graph_engine = await get_graph_engine()
 
@@ -38,7 +59,23 @@ async def get_context(self, query: str) -> Any:
         return result
 
     async def get_completion(self, query: str, context: Optional[Any] = None) -> Any:
-        """Returns the graph connections context."""
+        """
+        Returns the graph connections context.
+
+        If no context is provided, it retrieves the context using the specified query.
+
+        Parameters:
+        -----------
+
+            - query (str): The query to retrieve context.
+            - context (Optional[Any]): Optional context to use, otherwise fetched using the
+              query. (default None)
+
+        Returns:
+        --------
+
+            - Any: The context, either provided or retrieved.
+        """
         if context is None:
             context = await self.get_context(query)
         return context

cognee/modules/retrieval/graph_completion_context_extension_retriever.py

@@ -9,6 +9,21 @@
 
 
 class GraphCompletionContextExtensionRetriever(GraphCompletionRetriever):
+    """
+    Handles graph context completion for question answering tasks, extending context based
+    on retrieved triplets.
+
+    Public methods:
+    - get_completion
+
+    Instance variables:
+    - user_prompt_path
+    - system_prompt_path
+    - top_k
+    - node_type
+    - node_name
+    """
+
     def __init__(
         self,
         user_prompt_path: str = "graph_context_for_question.txt",
@@ -28,6 +43,30 @@ def __init__(
     async def get_completion(
         self, query: str, context: Optional[Any] = None, context_extension_rounds=4
     ) -> List[str]:
+        """
+        Extends the context for a given query by retrieving related triplets and generating new
+        completions based on them.
+
+        The method runs for a specified number of rounds to enhance context until no new
+        triplets are found or the maximum rounds are reached. It retrieves triplet suggestions
+        based on a generated completion from previous iterations, logging the process of context
+        extension.
+
+        Parameters:
+        -----------
+
+            - query (str): The input query for which the completion is generated.
+            - context (Optional[Any]): The existing context to use for enhancing the query; if
+              None, it will be initialized from triplets generated for the query. (default None)
+            - context_extension_rounds: The maximum number of rounds to extend the context with
+              new triplets before halting. (default 4)
+
+        Returns:
+        --------
+
+            - List[str]: A list containing the generated answer based on the query and the
+              extended context.
+        """
         triplets = []
 
         if context is None: