cognee 0.5.0.dev0__py3-none-any.whl → 0.5.1__py3-none-any.whl

cognee/modules/retrieval/code_retriever.py

@@ -1,232 +0,0 @@
-from typing import Any, Optional, List
-import asyncio
-import aiofiles
-from pydantic import BaseModel
-
-from cognee.shared.logging_utils import get_logger
-from cognee.modules.retrieval.base_retriever import BaseRetriever
-from cognee.infrastructure.databases.graph import get_graph_engine
-from cognee.infrastructure.databases.vector import get_vector_engine
-from cognee.infrastructure.llm.prompts import read_query_prompt
-from cognee.infrastructure.llm.LLMGateway import LLMGateway
-
-logger = get_logger("CodeRetriever")
-
-
-class CodeRetriever(BaseRetriever):
-    """Retriever for handling code-based searches."""
-
-    class CodeQueryInfo(BaseModel):
-        """
-        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
-
-    def __init__(self, top_k: int = 3):
-        """Initialize retriever with search parameters."""
-        self.top_k = top_k
-        self.file_name_collections = ["CodeFile_name"]
-        self.classes_and_functions_collections = [
-            "ClassDefinition_source_code",
-            "FunctionDefinition_source_code",
-        ]
-
-    async def _process_query(self, query: str) -> "CodeRetriever.CodeQueryInfo":
-        """Process the query using LLM to extract file names and source code parts."""
-        logger.debug(
-            f"Processing query with LLM: '{query[:100]}{'...' if len(query) > 100 else ''}'"
-        )
-
-        system_prompt = read_query_prompt("codegraph_retriever_system.txt")
-
-        try:
-            result = await LLMGateway.acreate_structured_output(
-                text_input=query,
-                system_prompt=system_prompt,
-                response_model=self.CodeQueryInfo,
-            )
-            logger.info(
-                f"LLM extracted {len(result.filenames)} filenames and {len(result.sourcecode)} chars of source code"
-            )
-            return result
-        except Exception as e:
-            logger.error(f"Failed to retrieve structured output from LLM: {str(e)}")
-            raise RuntimeError("Failed to retrieve structured output from LLM") from e
-
-    async def get_context(self, query: str) -> Any:
-        """Find relevant code files based on the query."""
-        logger.info(
-            f"Starting code retrieval for query: '{query[:100]}{'...' if len(query) > 100 else ''}'"
-        )
-
-        if not query or not isinstance(query, str):
-            logger.error("Invalid query: must be a non-empty string")
-            raise ValueError("The query must be a non-empty string.")
-
-        try:
-            vector_engine = get_vector_engine()
-            graph_engine = await get_graph_engine()
-            logger.debug("Successfully initialized vector and graph engines")
-        except Exception as e:
-            logger.error(f"Database initialization error: {str(e)}")
-            raise RuntimeError("Database initialization error in code_graph_retriever, ") from e
-
-        files_and_codeparts = await self._process_query(query)
-
-        similar_filenames = []
-        similar_codepieces = []
-
-        if not files_and_codeparts.filenames or not files_and_codeparts.sourcecode:
-            logger.info("No specific files/code extracted from query, performing general search")
-
-            for collection in self.file_name_collections:
-                logger.debug(f"Searching {collection} collection with general query")
-                search_results_file = await vector_engine.search(
-                    collection, query, limit=self.top_k
-                )
-                logger.debug(f"Found {len(search_results_file)} results in {collection}")
-                for res in search_results_file:
-                    similar_filenames.append(
-                        {"id": res.id, "score": res.score, "payload": res.payload}
-                    )
-
-            existing_collection = []
-            for collection in self.classes_and_functions_collections:
-                if await vector_engine.has_collection(collection):
-                    existing_collection.append(collection)
-
-            if not existing_collection:
-                raise RuntimeError("No collection found for code retriever")
-
-            for collection in existing_collection:
-                logger.debug(f"Searching {collection} collection with general query")
-                search_results_code = await vector_engine.search(
-                    collection, query, limit=self.top_k
-                )
-                logger.debug(f"Found {len(search_results_code)} results in {collection}")
-                for res in search_results_code:
-                    similar_codepieces.append(
-                        {"id": res.id, "score": res.score, "payload": res.payload}
-                    )
-        else:
-            logger.info(
-                f"Using extracted filenames ({len(files_and_codeparts.filenames)}) and source code for targeted search"
-            )
-
-            for collection in self.file_name_collections:
-                for file_from_query in files_and_codeparts.filenames:
-                    logger.debug(f"Searching {collection} for specific file: {file_from_query}")
-                    search_results_file = await vector_engine.search(
-                        collection, file_from_query, limit=self.top_k
-                    )
-                    logger.debug(
-                        f"Found {len(search_results_file)} results for file {file_from_query}"
-                    )
-                    for res in search_results_file:
-                        similar_filenames.append(
-                            {"id": res.id, "score": res.score, "payload": res.payload}
-                        )
-
-            for collection in self.classes_and_functions_collections:
-                logger.debug(f"Searching {collection} with extracted source code")
-                search_results_code = await vector_engine.search(
-                    collection, files_and_codeparts.sourcecode, limit=self.top_k
-                )
-                logger.debug(f"Found {len(search_results_code)} results for source code search")
-                for res in search_results_code:
-                    similar_codepieces.append(
-                        {"id": res.id, "score": res.score, "payload": res.payload}
-                    )
-
-        total_items = len(similar_filenames) + len(similar_codepieces)
-        logger.info(
-            f"Total search results: {total_items} items ({len(similar_filenames)} filenames, {len(similar_codepieces)} code pieces)"
-        )
-
-        if total_items == 0:
-            logger.warning("No search results found, returning empty list")
-            return []
-
-        logger.debug("Getting graph connections for all search results")
-        relevant_triplets = await asyncio.gather(
-            *[
-                graph_engine.get_connections(similar_piece["id"])
-                for similar_piece in similar_filenames + similar_codepieces
-            ]
-        )
-        logger.info(f"Retrieved graph connections for {len(relevant_triplets)} items")
-
-        paths = set()
-        for i, sublist in enumerate(relevant_triplets):
-            logger.debug(f"Processing connections for item {i}: {len(sublist)} connections")
-            for tpl in sublist:
-                if isinstance(tpl, tuple) and len(tpl) >= 3:
-                    if "file_path" in tpl[0]:
-                        paths.add(tpl[0]["file_path"])
-                    if "file_path" in tpl[2]:
-                        paths.add(tpl[2]["file_path"])
-
-        logger.info(f"Found {len(paths)} unique file paths to read")
-
-        retrieved_files = {}
-        read_tasks = []
-        for file_path in paths:
-
-            async def read_file(fp):
-                try:
-                    logger.debug(f"Reading file: {fp}")
-                    async with aiofiles.open(fp, "r", encoding="utf-8") as f:
-                        content = await f.read()
-                        retrieved_files[fp] = content
-                        logger.debug(f"Successfully read {len(content)} characters from {fp}")
-                except Exception as e:
-                    logger.error(f"Error reading {fp}: {e}")
-                    retrieved_files[fp] = ""
-
-            read_tasks.append(read_file(file_path))
-
-        await asyncio.gather(*read_tasks)
-        logger.info(
-            f"Successfully read {len([f for f in retrieved_files.values() if f])} files (out of {len(paths)} total)"
-        )
-
-        result = [
-            {
-                "name": file_path,
-                "description": file_path,
-                "content": retrieved_files[file_path],
-            }
-            for file_path in paths
-        ]
-
-        logger.info(f"Returning {len(result)} code file contexts")
-        return result
-
-    async def get_completion(
-        self, query: str, context: Optional[Any] = None, session_id: Optional[str] = None
-    ) -> Any:
-        """
-        Returns the code files context.
-
-        Parameters:
-        -----------
-
-            - query (str): The query string to retrieve code context for.
-            - context (Optional[Any]): Optional pre-fetched context; if None, it retrieves
-              the context for the query. (default None)
-            - session_id (Optional[str]): Optional session identifier for caching. If None,
-              defaults to 'default_session'. (default None)
-
-        Returns:
-        --------
-
-            - Any: The code files context, either provided or retrieved.
-        """
-        if context is None:
-            context = await self.get_context(query)
-        return context

cognee/tasks/code/enrich_dependency_graph_checker.py

@@ -1,35 +0,0 @@
-import os
-import asyncio
-import argparse
-from cognee.tasks.repo_processor.get_repo_file_dependencies import get_repo_file_dependencies
-from cognee.tasks.repo_processor.enrich_dependency_graph import enrich_dependency_graph
-
-
-def main():
-    """
-    Execute the main logic of the dependency graph processor.
-
-    This function sets up argument parsing to retrieve the repository path, checks the
-    existence of the specified path, and processes the repository to produce a dependency
-    graph. If the repository path does not exist, it logs an error message and terminates
-    without further execution.
-    """
-    parser = argparse.ArgumentParser()
-    parser.add_argument("repo_path", help="Path to the repository")
-    args = parser.parse_args()
-
-    repo_path = args.repo_path
-    if not os.path.exists(repo_path):
-        print(f"Error: The provided repository path does not exist: {repo_path}")
-        return
-
-    graph = asyncio.run(get_repo_file_dependencies(repo_path))
-    graph = asyncio.run(enrich_dependency_graph(graph))
-    for node in graph.nodes:
-        print(f"Node: {node}")
-        for _, target, data in graph.out_edges(node, data=True):
-            print(f"  Edge to {target}, data: {data}")
-
-
-if __name__ == "__main__":
-    main()

cognee/tasks/code/get_local_dependencies_checker.py

@@ -1,20 +0,0 @@
-import argparse
-import asyncio
-from cognee.tasks.repo_processor.get_local_dependencies import get_local_script_dependencies
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser(description="Get local script dependencies.")
-
-    # Suggested path: .../cognee/examples/python/simple_example.py
-    parser.add_argument("script_path", type=str, help="Absolute path to the Python script file")
-
-    # Suggested path: .../cognee
-    parser.add_argument("repo_path", type=str, help="Absolute path to the repository root")
-
-    args = parser.parse_args()
-
-    dependencies = asyncio.run(get_local_script_dependencies(args.script_path, args.repo_path))
-
-    print("Dependencies:")
-    for dependency in dependencies:
-        print(dependency)

cognee/tasks/code/get_repo_dependency_graph_checker.py

@@ -1,35 +0,0 @@
-import os
-import asyncio
-import argparse
-from cognee.tasks.repo_processor.get_repo_file_dependencies import get_repo_file_dependencies
-
-
-def main():
-    """
-    Parse the command line arguments and print the repository file dependencies.
-
-    This function sets up an argument parser to retrieve the path of a repository. It checks
-    if the provided path exists and if it doesn’t, it prints an error message and exits. If
-    the path is valid, it calls an asynchronous function to get the dependencies and prints
-    the nodes and their relations in the dependency graph.
-    """
-    parser = argparse.ArgumentParser()
-    parser.add_argument("repo_path", help="Path to the repository")
-    args = parser.parse_args()
-
-    repo_path = args.repo_path
-    if not os.path.exists(repo_path):
-        print(f"Error: The provided repository path does not exist: {repo_path}")
-        return
-
-    graph = asyncio.run(get_repo_file_dependencies(repo_path))
-
-    for node in graph.nodes:
-        print(f"Node: {node}")
-        edges = graph.edges(node, data=True)
-        for _, target, data in edges:
-            print(f"  Edge to {target}, Relation: {data.get('relation')}")
-
-
-if __name__ == "__main__":
-    main()

cognee/tasks/documents/check_permissions_on_dataset.py

@@ -1,26 +0,0 @@
-from cognee.modules.data.processing.document_types import Document
-from cognee.modules.users.permissions.methods import check_permission_on_dataset
-from typing import List
-
-
-async def check_permissions_on_dataset(
-    documents: List[Document], context: dict, user, permissions
-) -> List[Document]:
-    """
-    Validates a user's permissions on a list of documents.
-
-    Notes:
-        - This function assumes that `check_permission_on_documents` raises an exception if the permission check fails.
-        - It is designed to validate multiple permissions in a sequential manner for the same set of documents.
-        - Ensure that the `Document` and `user` objects conform to the expected structure and interfaces.
-    """
-
-    for permission in permissions:
-        await check_permission_on_dataset(
-            user,
-            permission,
-            # TODO: pass dataset through argument instead of context
-            context["dataset"].id,
-        )
-
-    return documents

cognee/tasks/repo_processor/__init__.py

@@ -1,2 +0,0 @@
-from .get_non_code_files import get_non_py_files
-from .get_repo_file_dependencies import get_repo_file_dependencies

cognee/tasks/repo_processor/get_local_dependencies.py

@@ -1,335 +0,0 @@
-import os
-import aiofiles
-import importlib
-from typing import AsyncGenerator, Optional
-from uuid import NAMESPACE_OID, uuid5
-import tree_sitter_python as tspython
-from tree_sitter import Language, Node, Parser, Tree
-from cognee.shared.logging_utils import get_logger
-
-from cognee.low_level import DataPoint
-from cognee.shared.CodeGraphEntities import (
-    CodeFile,
-    ImportStatement,
-    FunctionDefinition,
-    ClassDefinition,
-)
-
-logger = get_logger()
-
-
-class FileParser:
-    """
-    Handles the parsing of files into source code and an abstract syntax tree
-    representation. Public methods include:
-
-    - parse_file: Parses a file and returns its source code and syntax tree representation.
-    """
-
-    def __init__(self):
-        self.parsed_files = {}
-
-    async def parse_file(self, file_path: str) -> tuple[str, Tree]:
-        """
-        Parse a file and return its source code along with its syntax tree representation.
-
-        If the file has already been parsed, retrieve the result from memory instead of reading
-        the file again.
-
-        Parameters:
-        -----------
-
-            - file_path (str): The path of the file to parse.
-
-        Returns:
-        --------
-
-            - tuple[str, Tree]: A tuple containing the source code of the file and its
-              corresponding syntax tree representation.
-        """
-        PY_LANGUAGE = Language(tspython.language())
-        source_code_parser = Parser(PY_LANGUAGE)
-
-        if file_path not in self.parsed_files:
-            source_code = await get_source_code(file_path)
-            source_code_tree = source_code_parser.parse(bytes(source_code, "utf-8"))
-            self.parsed_files[file_path] = (source_code, source_code_tree)
-
-        return self.parsed_files[file_path]
-
-
-async def get_source_code(file_path: str):
-    """
-    Read source code from a file asynchronously.
-
-    This function attempts to open a file specified by the given file path, read its
-    contents, and return the source code. In case of any errors during the file reading
-    process, it logs an error message and returns None.
-
-    Parameters:
-    -----------
-
-        - file_path (str): The path to the file from which to read the source code.
-
-    Returns:
-    --------
-
-        Returns the contents of the file as a string if successful, or None if an error
-        occurs.
-    """
-    try:
-        async with aiofiles.open(file_path, "r", encoding="utf-8") as f:
-            source_code = await f.read()
-            return source_code
-    except Exception as error:
-        logger.error(f"Error reading file {file_path}: {str(error)}")
-        return None
-
-
-def resolve_module_path(module_name):
-    """
-    Find the file path of a module.
-
-    Return the file path of the specified module if found, or return None if the module does
-    not exist or cannot be located.
-
-    Parameters:
-    -----------
-
-        - module_name: The name of the module whose file path is to be resolved.
-
-    Returns:
-    --------
-
-        The file path of the module as a string or None if the module is not found.
-    """
-    try:
-        spec = importlib.util.find_spec(module_name)
-        if spec and spec.origin:
-            return spec.origin
-    except ModuleNotFoundError:
-        return None
-    return None
-
-
-def find_function_location(
-    module_path: str, function_name: str, parser: FileParser
-) -> Optional[tuple[str, str]]:
-    """
-    Find the location of a function definition in a specified module.
-
-    Parameters:
-    -----------
-
-        - module_path (str): The path to the module where the function is defined.
-        - function_name (str): The name of the function whose location is to be found.
-        - parser (FileParser): An instance of FileParser used to parse the module's source
-          code.
-
-    Returns:
-    --------
-
-        - Optional[tuple[str, str]]: Returns a tuple containing the module path and the
-          start point of the function if found; otherwise, returns None.
-    """
-    if not module_path or not os.path.exists(module_path):
-        return None
-
-    source_code, tree = parser.parse_file(module_path)
-    root_node: Node = tree.root_node
-
-    for node in root_node.children:
-        if node.type == "function_definition":
-            func_name_node = node.child_by_field_name("name")
-
-            if func_name_node and func_name_node.text.decode() == function_name:
-                return (module_path, node.start_point)  # (line, column)
-
-    return None
-
-
-async def get_local_script_dependencies(
-    repo_path: str, script_path: str, detailed_extraction: bool = False
-) -> CodeFile:
-    """
-    Retrieve local script dependencies and create a CodeFile object.
-
-    Parameters:
-    -----------
-
-        - repo_path (str): The path to the repository that contains the script.
-        - script_path (str): The path of the script for which dependencies are being
-          extracted.
-        - detailed_extraction (bool): A flag indicating whether to perform a detailed
-          extraction of code components.
-
-    Returns:
-    --------
-
-        - CodeFile: Returns a CodeFile object containing information about the script,
-          including its dependencies and definitions.
-    """
-    code_file_parser = FileParser()
-    source_code, source_code_tree = await code_file_parser.parse_file(script_path)
-
-    file_path_relative_to_repo = script_path[len(repo_path) + 1 :]
-
-    if not detailed_extraction:
-        code_file_node = CodeFile(
-            id=uuid5(NAMESPACE_OID, script_path),
-            name=file_path_relative_to_repo,
-            source_code=source_code,
-            file_path=script_path,
-            language="python",
-        )
-        return code_file_node
-
-    code_file_node = CodeFile(
-        id=uuid5(NAMESPACE_OID, script_path),
-        name=file_path_relative_to_repo,
-        source_code=None,
-        file_path=script_path,
-        language="python",
-    )
-
-    async for part in extract_code_parts(source_code_tree.root_node, script_path=script_path):
-        part.file_path = script_path
-
-        if isinstance(part, FunctionDefinition):
-            code_file_node.provides_function_definition.append(part)
-        if isinstance(part, ClassDefinition):
-            code_file_node.provides_class_definition.append(part)
-        if isinstance(part, ImportStatement):
-            code_file_node.depends_on.append(part)
-
-    return code_file_node
-
-
-def find_node(nodes: list[Node], condition: callable) -> Node:
-    """
-    Find and return the first node that satisfies the given condition.
-
-    Iterate through the provided list of nodes and return the first node for which the
-    condition callable returns True. If no such node is found, return None.
-
-    Parameters:
-    -----------
-
-        - nodes (list[Node]): A list of Node objects to search through.
-        - condition (callable): A callable that takes a Node and returns a boolean
-          indicating if the node meets specified criteria.
-
-    Returns:
-    --------
-
-        - Node: The first Node that matches the condition, or None if no such node exists.
-    """
-    for node in nodes:
-        if condition(node):
-            return node
-
-    return None
-
-
-async def extract_code_parts(
-    tree_root: Node, script_path: str, existing_nodes: list[DataPoint] = {}
-) -> AsyncGenerator[DataPoint, None]:
-    """
-    Extract code parts from a given AST node tree asynchronously.
-
-    Iteratively yields DataPoint nodes representing import statements, function definitions,
-    and class definitions found in the children of the specified tree root. The function
-    checks
-    if nodes are already present in the existing_nodes dictionary to prevent duplicates.
-    This function has to be used in an asynchronous context, and it requires a valid
-    tree_root
-    and proper initialization of existing_nodes.
-
-    Parameters:
-    -----------
-
-        - tree_root (Node): The root node of the AST tree containing code parts to extract.
-        - script_path (str): The file path of the script from which the AST was generated.
-        - existing_nodes (list[DataPoint]): A dictionary that holds already extracted
-          DataPoint nodes to avoid duplicates. (default {})
-
-    Returns:
-    --------
-
-        Yields DataPoint nodes representing imported modules, functions, and classes.
-    """
-    for child_node in tree_root.children:
-        if child_node.type == "import_statement" or child_node.type == "import_from_statement":
-            parts = child_node.text.decode("utf-8").split()
-
-            if parts[0] == "import":
-                module_name = parts[1]
-                function_name = None
-            elif parts[0] == "from":
-                module_name = parts[1]
-                function_name = parts[3]
-
-                if " as " in function_name:
-                    function_name = function_name.split(" as ")[0]
-
-            if " as " in module_name:
-                module_name = module_name.split(" as ")[0]
-
-            if function_name and "import " + function_name not in existing_nodes:
-                import_statement_node = ImportStatement(
-                    name=function_name,
-                    module=module_name,
-                    start_point=child_node.start_point,
-                    end_point=child_node.end_point,
-                    file_path=script_path,
-                    source_code=child_node.text,
-                )
-                existing_nodes["import " + function_name] = import_statement_node
-
-            if function_name:
-                yield existing_nodes["import " + function_name]
-
-            if module_name not in existing_nodes:
-                import_statement_node = ImportStatement(
-                    name=module_name,
-                    module=module_name,
-                    start_point=child_node.start_point,
-                    end_point=child_node.end_point,
-                    file_path=script_path,
-                    source_code=child_node.text,
-                )
-                existing_nodes[module_name] = import_statement_node
-
-            yield existing_nodes[module_name]
-
-        if child_node.type == "function_definition":
-            function_node = find_node(child_node.children, lambda node: node.type == "identifier")
-            function_node_name = function_node.text
-
-            if function_node_name not in existing_nodes:
-                function_definition_node = FunctionDefinition(
-                    name=function_node_name,
-                    start_point=child_node.start_point,
-                    end_point=child_node.end_point,
-                    file_path=script_path,
-                    source_code=child_node.text,
-                )
-                existing_nodes[function_node_name] = function_definition_node
-
-            yield existing_nodes[function_node_name]
-
-        if child_node.type == "class_definition":
-            class_name_node = find_node(child_node.children, lambda node: node.type == "identifier")
-            class_name_node_name = class_name_node.text
-
-            if class_name_node_name not in existing_nodes:
-                class_definition_node = ClassDefinition(
-                    name=class_name_node_name,
-                    start_point=child_node.start_point,
-                    end_point=child_node.end_point,
-                    file_path=script_path,
-                    source_code=child_node.text,
-                )
-                existing_nodes[class_name_node_name] = class_definition_node
-
-            yield existing_nodes[class_name_node_name]