cognee 0.2.3.dev1__py3-none-any.whl → 0.3.0__py3-none-any.whl

cognee/tasks/repo_processor/get_repo_file_dependencies.py

@@ -1,16 +1,35 @@
 import asyncio
 import math
 import os
-
-# from concurrent.futures import ProcessPoolExecutor
-from typing import AsyncGenerator
+from pathlib import Path
+from typing import Set
+from typing import AsyncGenerator, Optional, List
 from uuid import NAMESPACE_OID, uuid5
 
 from cognee.infrastructure.engine import DataPoint
 from cognee.shared.CodeGraphEntities import CodeFile, Repository
 
-
-async def get_source_code_files(repo_path):
+# constant, declared only once
+EXCLUDED_DIRS: Set[str] = {
+    ".venv",
+    "venv",
+    "env",
+    ".env",
+    "site-packages",
+    "node_modules",
+    "dist",
+    "build",
+    ".git",
+    "tests",
+    "test",
+}
+
+
+async def get_source_code_files(
+    repo_path,
+    language_config: dict[str, list[str]] | None = None,
+    excluded_paths: Optional[List[str]] = None,
+):
     """
     Retrieve Python source code files from the specified repository path.
 
@@ -20,40 +39,70 @@ async def get_source_code_files(repo_path):
 
     Parameters:
     -----------
-
-        - repo_path: The file path to the repository to search for Python source files.
+    - repo_path: Root path of the repository to search
+    - language_config: dict mapping language names to file extensions, e.g.,
+            {'python': ['.py'], 'javascript': ['.js', '.jsx'], ...}
+    - excluded_paths: Optional list of path fragments or glob patterns to exclude
 
     Returns:
     --------
-
-        A list of absolute paths to .py files that contain source code, excluding empty
-        files, test files, and files from a virtual environment.
+        A list of (absolute_path, language) tuples for source code files.
     """
-    if not os.path.exists(repo_path):
-        return {}
-
-    py_files_paths = (
-        os.path.join(root, file)
-        for root, _, files in os.walk(repo_path)
-        for file in files
-        if (
-            file.endswith(".py")
-            and not file.startswith("test_")
-            and not file.endswith("_test")
-            and ".venv" not in file
-        )
-    )
 
-    source_code_files = set()
-    for file_path in py_files_paths:
-        file_path = os.path.abspath(file_path)
-
-        if os.path.getsize(file_path) == 0:
-            continue
+    def _get_language_from_extension(file, language_config):
+        for lang, exts in language_config.items():
+            for ext in exts:
+                if file.endswith(ext):
+                    return lang
+        return None
+
+    # Default config if not provided
+    if language_config is None:
+        language_config = {
+            "python": [".py"],
+            "javascript": [".js", ".jsx"],
+            "typescript": [".ts", ".tsx"],
+            "java": [".java"],
+            "csharp": [".cs"],
+            "go": [".go"],
+            "rust": [".rs"],
+            "cpp": [".cpp", ".c", ".h", ".hpp"],
+        }
 
-        source_code_files.add(file_path)
+    if not os.path.exists(repo_path):
+        return []
 
-    return list(source_code_files)
+    source_code_files = set()
+    for root, _, files in os.walk(repo_path):
+        for file in files:
+            lang = _get_language_from_extension(file, language_config)
+            if lang is None:
+                continue
+            # Exclude tests, common build/venv directories and files provided in exclude_paths
+            excluded_dirs = EXCLUDED_DIRS
+            excluded_paths = {Path(p).resolve() for p in (excluded_paths or [])}  # full paths
+
+            root_path = Path(root).resolve()
+            root_parts = set(root_path.parts)  # same as before
+            base_name, _ext = os.path.splitext(file)
+            if (
+                base_name.startswith("test_")
+                or base_name.endswith("_test")
+                or ".test." in file
+                or ".spec." in file
+                or (excluded_dirs & root_parts)  # name match
+                or any(
+                    root_path.is_relative_to(p)  # full-path match
+                    for p in excluded_paths
+                )
+            ):
+                continue
+            file_path = os.path.abspath(os.path.join(root, file))
+            if os.path.getsize(file_path) == 0:
+                continue
+            source_code_files.add((file_path, lang))
+
+    return sorted(list(source_code_files))
 
 
 def run_coroutine(coroutine_func, *args, **kwargs):
@@ -85,22 +134,26 @@ def run_coroutine(coroutine_func, *args, **kwargs):
 
 
 async def get_repo_file_dependencies(
-    repo_path: str, detailed_extraction: bool = False
+    repo_path: str,
+    detailed_extraction: bool = False,
+    supported_languages: list = None,
+    excluded_paths: Optional[List[str]] = None,
 ) -> AsyncGenerator[DataPoint, None]:
     """
-    Generate a dependency graph for Python files in the given repository path.
+    Generate a dependency graph for source files (multi-language) in the given repository path.
 
     Check the validity of the repository path and yield a repository object followed by the
-    dependencies of Python files within that repository. Raise a FileNotFoundError if the
+    dependencies of source files within that repository. Raise a FileNotFoundError if the
     provided path does not exist. The extraction of detailed dependencies can be controlled
-    via the `detailed_extraction` argument.
+    via the `detailed_extraction` argument. Languages considered can be restricted via
+    the `supported_languages` argument.
 
     Parameters:
     -----------
 
-        - repo_path (str): The file path to the repository where Python files are located.
-        - detailed_extraction (bool): A flag indicating whether to perform a detailed
-          extraction of dependencies (default is False). (default False)
+        - repo_path (str): The file path to the repository to process.
+        - detailed_extraction (bool): Whether to perform a detailed extraction of code parts.
+        - supported_languages (list | None): Subset of languages to include; if None, use defaults.
     """
 
     if isinstance(repo_path, list) and len(repo_path) == 1:
@@ -109,7 +162,28 @@ async def get_repo_file_dependencies(
     if not os.path.exists(repo_path):
         raise FileNotFoundError(f"Repository path {repo_path} does not exist.")
 
-    source_code_files = await get_source_code_files(repo_path)
+    # Build language config from supported_languages
+    default_language_config = {
+        "python": [".py"],
+        "javascript": [".js", ".jsx"],
+        "typescript": [".ts", ".tsx"],
+        "java": [".java"],
+        "csharp": [".cs"],
+        "go": [".go"],
+        "rust": [".rs"],
+        "cpp": [".cpp", ".c", ".h", ".hpp"],
+        "c": [".c", ".h"],
+    }
+    if supported_languages is not None:
+        language_config = {
+            k: v for k, v in default_language_config.items() if k in supported_languages
+        }
+    else:
+        language_config = default_language_config
+
+    source_code_files = await get_source_code_files(
+        repo_path, language_config=language_config, excluded_paths=excluded_paths
+    )
 
     repo = Repository(
         id=uuid5(NAMESPACE_OID, repo_path),
@@ -128,19 +202,42 @@ async def get_repo_file_dependencies(
         for chunk_number in range(number_of_chunks)
     ]
 
-    # Codegraph dependencies are not installed by default, so we import where we use them.
+    # Import dependency extractors for each language (Python for now, extend later)
     from cognee.tasks.repo_processor.get_local_dependencies import get_local_script_dependencies
+    import aiofiles
+    # TODO: Add other language extractors here
 
     for start_range, end_range in chunk_ranges:
-        # with ProcessPoolExecutor(max_workers=12) as executor:
-        tasks = [
-            get_local_script_dependencies(repo_path, file_path, detailed_extraction)
-            for file_path in source_code_files[start_range : end_range + 1]
-        ]
+        tasks = []
+        for file_path, lang in source_code_files[start_range : end_range + 1]:
+            # For now, only Python is supported; extend with other languages
+            if lang == "python":
+                tasks.append(
+                    get_local_script_dependencies(repo_path, file_path, detailed_extraction)
+                )
+            else:
+                # Placeholder: create a minimal CodeFile for other languages
+                async def make_codefile_stub(file_path=file_path, lang=lang):
+                    async with aiofiles.open(
+                        file_path, "r", encoding="utf-8", errors="replace"
+                    ) as f:
+                        source = await f.read()
+                    return CodeFile(
+                        id=uuid5(NAMESPACE_OID, file_path),
+                        name=os.path.relpath(file_path, repo_path),
+                        file_path=file_path,
+                        language=lang,
+                        source_code=source,
+                    )
+
+                tasks.append(make_codefile_stub())
 
         results: list[CodeFile] = await asyncio.gather(*tasks)
 
         for source_code_file in results:
             source_code_file.part_of = repo
-
+            if getattr(
+                source_code_file, "language", None
+            ) is None and source_code_file.file_path.endswith(".py"):
+                source_code_file.language = "python"
             yield source_code_file

cognee/tasks/storage/add_data_points.py

@@ -10,7 +10,37 @@ from cognee.tasks.storage.exceptions import (
 )
 
 
-async def add_data_points(data_points: List[DataPoint]) -> List[DataPoint]:
+async def add_data_points(
+    data_points: List[DataPoint], update_edge_collection: bool = True
+) -> List[DataPoint]:
+    """
+    Add a batch of data points to the graph database by extracting nodes and edges,
+    deduplicating them, and indexing them for retrieval.
+
+    This function parallelizes the graph extraction for each data point,
+    merges the resulting nodes and edges, and ensures uniqueness before
+    committing them to the underlying graph engine. It also updates the
+    associated retrieval indices for nodes and (optionally) edges.
+
+    Args:
+        data_points (List[DataPoint]):
+            A list of data points to process and insert into the graph.
+        update_edge_collection (bool, optional):
+            Whether to update the edge index after adding edges.
+            Defaults to True.
+
+    Returns:
+        List[DataPoint]:
+            The original list of data points after processing and insertion.
+
+    Side Effects:
+        - Calls `get_graph_from_model` concurrently for each data point.
+        - Deduplicates nodes and edges across all results.
+        - Updates the node index via `index_data_points`.
+        - Inserts nodes and edges into the graph engine.
+        - Optionally updates the edge index via `index_graph_edges`.
+    """
+
     if not isinstance(data_points, list):
         raise InvalidDataPointsInAddDataPointsError("data_points must be a list.")
     if not all(isinstance(dp, DataPoint) for dp in data_points):
@@ -48,7 +78,7 @@ async def add_data_points(data_points: List[DataPoint]) -> List[DataPoint]:
     await graph_engine.add_nodes(nodes)
     await graph_engine.add_edges(edges)
 
-    # This step has to happen after adding nodes and edges because we query the graph.
-    await index_graph_edges()
+    if update_edge_collection:
+        await index_graph_edges()
 
     return data_points

cognee/tasks/temporal_graph/__init__.py

@@ -0,0 +1 @@
+

cognee/tasks/temporal_graph/add_entities_to_event.py

@@ -0,0 +1,85 @@
+from cognee.modules.engine.models import Event
+from cognee.tasks.temporal_graph.models import EventWithEntities
+from cognee.modules.engine.models.Entity import Entity
+from cognee.modules.engine.models.EntityType import EntityType
+from cognee.infrastructure.engine.models.Edge import Edge
+from cognee.modules.engine.utils import generate_node_id, generate_node_name
+
+
+def add_entities_to_event(event: Event, event_with_entities: EventWithEntities) -> None:
+    """
+    Adds extracted entities to an Event object by populating its attributes field.
+
+    For each attribute in the provided EventWithEntities, the function ensures that
+    the corresponding entity type exists, creates an Entity node with metadata, and
+    links it to the event via an Edge representing the relationship. Entities are
+    cached by type to avoid duplication.
+
+    Args:
+        event (Event): The target Event object to enrich with entities.
+        event_with_entities (EventWithEntities): An event model containing extracted
+            attributes with entity, type, and relationship metadata.
+
+    Returns:
+        None
+    """
+
+    if not event_with_entities.attributes:
+        return
+
+    # Create entity types cache
+    entity_types = {}
+
+    # Process each attribute
+    for attribute in event_with_entities.attributes:
+        # Get or create entity type
+        entity_type = get_or_create_entity_type(entity_types, attribute.entity_type)
+
+        # Create entity
+        entity_id = generate_node_id(attribute.entity)
+        entity_name = generate_node_name(attribute.entity)
+        entity = Entity(
+            id=entity_id,
+            name=entity_name,
+            is_a=entity_type,
+            description=f"Entity {attribute.entity} of type {attribute.entity_type}",
+            ontology_valid=False,
+            belongs_to_set=None,
+        )
+
+        # Create edge
+        edge = Edge(relationship_type=attribute.relationship)
+
+        # Add to event attributes
+        if event.attributes is None:
+            event.attributes = []
+        event.attributes.append((edge, [entity]))
+
+
+def get_or_create_entity_type(entity_types: dict, entity_type_name: str) -> EntityType:
+    """
+    Retrieves an existing EntityType from the cache or creates a new one if it does not exist.
+
+    If the given entity type name is not already in the cache, a new EntityType is generated
+    with a unique ID, normalized name, and description, then added to the cache.
+
+    Args:
+        entity_types (dict): A cache mapping entity type names to EntityType objects.
+        entity_type_name (str): The name of the entity type to retrieve or create.
+
+    Returns:
+        EntityType: The existing or newly created EntityType object.
+    """
+    if entity_type_name not in entity_types:
+        type_id = generate_node_id(entity_type_name)
+        type_name = generate_node_name(entity_type_name)
+        entity_type = EntityType(
+            id=type_id,
+            name=type_name,
+            type=type_name,
+            description=f"Type for {entity_type_name}",
+            ontology_valid=False,
+        )
+        entity_types[entity_type_name] = entity_type
+
+    return entity_types[entity_type_name]

cognee/tasks/temporal_graph/enrich_events.py

@@ -0,0 +1,34 @@
+from typing import List
+
+from cognee.infrastructure.llm import LLMGateway
+from cognee.modules.engine.models import Event
+from cognee.tasks.temporal_graph.models import EventWithEntities, EventEntityList
+
+
+async def enrich_events(events: List[Event]) -> List[EventWithEntities]:
+    """
+    Enriches a list of events by extracting entities using an LLM.
+
+    The function serializes event data into JSON, sends it to the LLM for
+    entity extraction, and returns enriched events with associated entities.
+
+    Args:
+        events (List[Event]): A list of Event objects to be enriched.
+
+    Returns:
+        List[EventWithEntities]: A list of events augmented with extracted entities.
+    """
+
+    import json
+
+    # Convert events to JSON format for LLM processing
+    events_json = [
+        {"event_name": event.name, "description": event.description or ""} for event in events
+    ]
+
+    events_json_str = json.dumps(events_json)
+
+    # Extract entities from events
+    entity_result = await LLMGateway.extract_event_entities(events_json_str, EventEntityList)
+
+    return entity_result.events

cognee/tasks/temporal_graph/extract_events_and_entities.py

@@ -0,0 +1,32 @@
+import asyncio
+from typing import Type, List
+from cognee.infrastructure.llm.LLMGateway import LLMGateway
+from cognee.modules.chunking.models import DocumentChunk
+from cognee.tasks.temporal_graph.models import EventList
+from cognee.modules.engine.utils.generate_event_datapoint import generate_event_datapoint
+
+
+async def extract_events_and_timestamps(data_chunks: List[DocumentChunk]) -> List[DocumentChunk]:
+    """
+    Extracts events and their timestamps from document chunks using an LLM.
+
+    Each document chunk is processed with the event graph extractor to identify events.
+    The extracted events are converted into Event datapoints and appended to the
+    chunk's `contains` list.
+
+    Args:
+        data_chunks (List[DocumentChunk]): A list of document chunks containing text to process.
+
+    Returns:
+        List[DocumentChunk]: The same list of document chunks, enriched with extracted Event datapoints.
+    """
+    events = await asyncio.gather(
+        *[LLMGateway.extract_event_graph(chunk.text, EventList) for chunk in data_chunks]
+    )
+
+    for data_chunk, event_list in zip(data_chunks, events):
+        for event in event_list.events:
+            event_datapoint = generate_event_datapoint(event)
+            data_chunk.contains.append(event_datapoint)
+
+    return data_chunks

cognee/tasks/temporal_graph/extract_knowledge_graph_from_events.py

@@ -0,0 +1,41 @@
+from typing import List
+from cognee.modules.chunking.models import DocumentChunk
+from cognee.modules.engine.models import Event
+from cognee.tasks.temporal_graph.enrich_events import enrich_events
+from cognee.tasks.temporal_graph.add_entities_to_event import add_entities_to_event
+
+
+async def extract_knowledge_graph_from_events(
+    data_chunks: List[DocumentChunk],
+) -> List[DocumentChunk]:
+    """
+    Extracts events from document chunks and enriches them with entities to form a knowledge graph.
+
+    The function collects all Event objects from the given document chunks,
+    uses an LLM to extract and attach related entities, and updates the events
+    with these enriched attributes.
+
+    Args:
+        data_chunks (List[DocumentChunk]): A list of document chunks containing extracted events.
+
+    Returns:
+        List[DocumentChunk]: The same list of document chunks, with their events enriched by entities.
+    """
+    # Extract events from chunks
+    all_events = []
+    for chunk in data_chunks:
+        for item in chunk.contains:
+            if isinstance(item, Event):
+                all_events.append(item)
+
+    if not all_events:
+        return data_chunks
+
+    # Enrich events with entities
+    enriched_events = await enrich_events(all_events)
+
+    # Add entities to events
+    for event, enriched_event in zip(all_events, enriched_events):
+        add_entities_to_event(event, enriched_event)
+
+    return data_chunks

cognee/tasks/temporal_graph/models.py

@@ -0,0 +1,49 @@
+from typing import Optional, List
+from pydantic import BaseModel, Field
+
+
+class Timestamp(BaseModel):
+    year: int = Field(..., ge=1, le=9999)
+    month: int = Field(..., ge=1, le=12)
+    day: int = Field(..., ge=1, le=31)
+    hour: int = Field(..., ge=0, le=23)
+    minute: int = Field(..., ge=0, le=59)
+    second: int = Field(..., ge=0, le=59)
+
+
+class Interval(BaseModel):
+    starts_at: Timestamp
+    ends_at: Timestamp
+
+
+class QueryInterval(BaseModel):
+    starts_at: Optional[Timestamp] = None
+    ends_at: Optional[Timestamp] = None
+
+
+class Event(BaseModel):
+    name: str
+    description: Optional[str] = None
+    time_from: Optional[Timestamp] = None
+    time_to: Optional[Timestamp] = None
+    location: Optional[str] = None
+
+
+class EventList(BaseModel):
+    events: List[Event]
+
+
+class EntityAttribute(BaseModel):
+    entity: str
+    entity_type: str
+    relationship: str
+
+
+class EventWithEntities(BaseModel):
+    event_name: str
+    description: Optional[str] = None
+    attributes: List[EntityAttribute] = []
+
+
+class EventEntityList(BaseModel):
+    events: List[EventWithEntities]

cognee/tests/integration/cli/__init__.py

@@ -0,0 +1,3 @@
+"""
+CLI integration tests package.
+"""