cognee 0.2.4__py3-none-any.whl → 0.3.0.dev0__py3-none-any.whl

cognee/tasks/repo_processor/get_repo_file_dependencies.py

@@ -1,24 +1,48 @@
 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, language_config: dict[str, list[str]] | None = None):
+# 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 source code files from the specified repository path for multiple languages.
+    Retrieve Python source code files from the specified repository path.
+
+    This function scans the given repository path for files that have the .py extension
+    while excluding test files and files within a virtual environment. It returns a list of
+    absolute paths to the source code files that are not empty.
 
     Parameters:
     -----------
-        - repo_path: The file path to the repository to search for source files.
-        - language_config: dict mapping language names to file extensions, e.g.,
+    - 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:
     --------
@@ -54,28 +78,23 @@ async def get_source_code_files(repo_path, language_config: dict[str, list[str]]
             lang = _get_language_from_extension(file, language_config)
             if lang is None:
                 continue
-            # Exclude tests and common build/venv directories
-            excluded_dirs = {
-                ".venv",
-                "venv",
-                "env",
-                ".env",
-                "site-packages",
-                "node_modules",
-                "dist",
-                "build",
-                ".git",
-                "tests",
-                "test",
-            }
-            root_parts = set(os.path.normpath(root).split(os.sep))
+            # 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")  # catches Go's *_test.go and similar
+                or base_name.endswith("_test")
                 or ".test." in file
                 or ".spec." in file
-                or (excluded_dirs & root_parts)
+                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))
@@ -115,7 +134,10 @@ def run_coroutine(coroutine_func, *args, **kwargs):
 
 
 async def get_repo_file_dependencies(
-    repo_path: str, detailed_extraction: bool = False, supported_languages: list = None
+    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 source files (multi-language) in the given repository path.
@@ -150,6 +172,7 @@ async def get_repo_file_dependencies(
         "go": [".go"],
         "rust": [".rs"],
         "cpp": [".cpp", ".c", ".h", ".hpp"],
+        "c": [".c", ".h"],
     }
     if supported_languages is not None:
         language_config = {
@@ -158,7 +181,9 @@ async def get_repo_file_dependencies(
     else:
         language_config = default_language_config
 
-    source_code_files = await get_source_code_files(repo_path, language_config=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),

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/test_kuzu.py

@@ -94,21 +94,21 @@ async def main():
 
         await cognee.cognify([dataset_name])
 
-        context_nonempty, _ = await GraphCompletionRetriever(
+        context_nonempty = await GraphCompletionRetriever(
             node_type=NodeSet,
             node_name=["first"],
         ).get_context("What is in the context?")
 
-        context_empty, _ = await GraphCompletionRetriever(
+        context_empty = await GraphCompletionRetriever(
             node_type=NodeSet,
             node_name=["nonexistent"],
         ).get_context("What is in the context?")
 
-        assert isinstance(context_nonempty, str) and context_nonempty != "", (
+        assert isinstance(context_nonempty, list) and context_nonempty != [], (
             f"Nodeset_search_test:Expected non-empty string for context_nonempty, got: {context_nonempty!r}"
         )
 
-        assert context_empty == "", (
+        assert context_empty == [], (
             f"Nodeset_search_test:Expected empty string for context_empty, got: {context_empty!r}"
         )
 

cognee/tests/test_neo4j.py

@@ -98,21 +98,21 @@ async def main():
 
     await cognee.cognify([dataset_name])
 
-    context_nonempty, _ = await GraphCompletionRetriever(
+    context_nonempty = await GraphCompletionRetriever(
         node_type=NodeSet,
         node_name=["first"],
     ).get_context("What is in the context?")
 
-    context_empty, _ = await GraphCompletionRetriever(
+    context_empty = await GraphCompletionRetriever(
         node_type=NodeSet,
         node_name=["nonexistent"],
     ).get_context("What is in the context?")
 
-    assert isinstance(context_nonempty, str) and context_nonempty != "", (
+    assert isinstance(context_nonempty, list) and context_nonempty != [], (
         f"Nodeset_search_test:Expected non-empty string for context_nonempty, got: {context_nonempty!r}"
     )
 
-    assert context_empty == "", (
+    assert context_empty == [], (
         f"Nodeset_search_test:Expected empty string for context_empty, got: {context_empty!r}"
     )
 

cognee/tests/test_permissions.py

@@ -79,7 +79,7 @@ async def main():
     print("\n\nExtracted sentences are:\n")
     for result in search_results:
         print(f"{result}\n")
-    assert search_results[0]["dataset_name"] == "NLP", (
+    assert search_results[0].dataset_name == "NLP", (
         f"Dict must contain dataset name 'NLP': {search_results[0]}"
     )
 
@@ -93,7 +93,7 @@ async def main():
     print("\n\nExtracted sentences are:\n")
     for result in search_results:
         print(f"{result}\n")
-    assert search_results[0]["dataset_name"] == "QUANTUM", (
+    assert search_results[0].dataset_name == "QUANTUM", (
         f"Dict must contain dataset name 'QUANTUM': {search_results[0]}"
     )
 
@@ -170,7 +170,7 @@ async def main():
     for result in search_results:
         print(f"{result}\n")
 
-    assert search_results[0]["dataset_name"] == "QUANTUM", (
+    assert search_results[0].dataset_name == "QUANTUM", (
         f"Dict must contain dataset name 'QUANTUM': {search_results[0]}"
     )
 

cognee/tests/test_relational_db_migration.py

@@ -1,6 +1,6 @@
-import json
 import pathlib
 import os
+from typing import List
 from cognee.infrastructure.databases.graph import get_graph_engine
 from cognee.infrastructure.databases.relational import (
     get_migration_relational_engine,
@@ -10,7 +10,7 @@ from cognee.infrastructure.databases.vector.pgvector import (
     create_db_and_tables as create_pgvector_db_and_tables,
 )
 from cognee.tasks.ingestion import migrate_relational_database
-from cognee.modules.search.types import SearchType
+from cognee.modules.search.types import SearchResult, SearchType
 import cognee
 
 
@@ -45,13 +45,15 @@ async def relational_db_migration():
     await migrate_relational_database(graph_engine, schema=schema)
 
     # 1. Search the graph
-    search_results = await cognee.search(
+    search_results: List[SearchResult] = await cognee.search(
         query_type=SearchType.GRAPH_COMPLETION, query_text="Tell me about the artist AC/DC"
-    )
+    )  # type: ignore
     print("Search results:", search_results)
 
     # 2. Assert that the search results contain "AC/DC"
-    assert any("AC/DC" in r for r in search_results), "AC/DC not found in search results!"
+    assert any("AC/DC" in r.search_result for r in search_results), (
+        "AC/DC not found in search results!"
+    )
 
     migration_db_provider = migration_engine.engine.dialect.name
     if migration_db_provider == "postgresql":

cognee/tests/test_search_db.py

@@ -1,11 +1,7 @@
-import os
-import pathlib
-
-from dns.e164 import query
-
 import cognee
 from cognee.infrastructure.databases.graph import get_graph_engine
 from cognee.modules.graph.cognee_graph.CogneeGraphElements import Edge
+from cognee.modules.graph.utils import resolve_edges_to_text
 from cognee.modules.retrieval.graph_completion_retriever import GraphCompletionRetriever
 from cognee.modules.retrieval.graph_completion_context_extension_retriever import (
     GraphCompletionContextExtensionRetriever,
@@ -14,11 +10,8 @@ from cognee.modules.retrieval.graph_completion_cot_retriever import GraphComplet
 from cognee.modules.retrieval.graph_summary_completion_retriever import (
     GraphSummaryCompletionRetriever,
 )
-from cognee.modules.search.operations import get_history
-from cognee.modules.users.methods import get_default_user
 from cognee.shared.logging_utils import get_logger
 from cognee.modules.search.types import SearchType
-from cognee.modules.engine.models import NodeSet
 from collections import Counter
 
 logger = get_logger()
@@ -46,16 +39,16 @@ async def main():
 
     await cognee.cognify([dataset_name])
 
-    context_gk, _ = await GraphCompletionRetriever().get_context(
+    context_gk = await GraphCompletionRetriever().get_context(
         query="Next to which country is Germany located?"
     )
-    context_gk_cot, _ = await GraphCompletionCotRetriever().get_context(
+    context_gk_cot = await GraphCompletionCotRetriever().get_context(
         query="Next to which country is Germany located?"
     )
-    context_gk_ext, _ = await GraphCompletionContextExtensionRetriever().get_context(
+    context_gk_ext = await GraphCompletionContextExtensionRetriever().get_context(
         query="Next to which country is Germany located?"
     )
-    context_gk_sum, _ = await GraphSummaryCompletionRetriever().get_context(
+    context_gk_sum = await GraphSummaryCompletionRetriever().get_context(
         query="Next to which country is Germany located?"
     )
 
@@ -65,9 +58,11 @@ async def main():
         ("GraphCompletionContextExtensionRetriever", context_gk_ext),
         ("GraphSummaryCompletionRetriever", context_gk_sum),
     ]:
-        assert isinstance(context, str), f"{name}: Context should be a string"
-        assert context.strip(), f"{name}: Context should not be empty"
-        lower = context.lower()
+        assert isinstance(context, list), f"{name}: Context should be a list"
+        assert len(context) > 0, f"{name}: Context should not be empty"
+
+        context_text = await resolve_edges_to_text(context)
+        lower = context_text.lower()
         assert "germany" in lower or "netherlands" in lower, (
             f"{name}: Context did not contain 'germany' or 'netherlands'; got: {context!r}"
         )
@@ -143,20 +138,19 @@ async def main():
         last_k=1,
     )
 
-    for name, completion in [
+    for name, search_results in [
         ("GRAPH_COMPLETION", completion_gk),
         ("GRAPH_COMPLETION_COT", completion_cot),
         ("GRAPH_COMPLETION_CONTEXT_EXTENSION", completion_ext),
         ("GRAPH_SUMMARY_COMPLETION", completion_sum),
     ]:
-        assert isinstance(completion, list), f"{name}: should return a list"
-        assert len(completion) == 1, f"{name}: expected single-element list, got {len(completion)}"
-        text = completion[0]
-        assert isinstance(text, str), f"{name}: element should be a string"
-        assert text.strip(), f"{name}: string should not be empty"
-        assert "netherlands" in text.lower(), (
-            f"{name}: expected 'netherlands' in result, got: {text!r}"
-        )
+        for search_result in search_results:
+            completion = search_result.search_result
+            assert isinstance(completion, str), f"{name}: should return a string"
+            assert completion.strip(), f"{name}: string should not be empty"
+            assert "netherlands" in completion.lower(), (
+                f"{name}: expected 'netherlands' in result, got: {completion!r}"
+            )
 
     graph_engine = await get_graph_engine()
     graph = await graph_engine.get_graph_data()