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

cognee/infrastructure/llm/structured_output_framework/litellm_instructor/extraction/__init__.py

@@ -1,3 +1,5 @@
 from .knowledge_graph.extract_content_graph import extract_content_graph
+from .knowledge_graph.extract_event_graph import extract_event_graph
 from .extract_categories import extract_categories
 from .extract_summary import extract_summary, extract_code_summary
+from .extract_event_entities import extract_event_entities

cognee/infrastructure/llm/structured_output_framework/litellm_instructor/extraction/extract_event_entities.py

@@ -0,0 +1,44 @@
+import os
+from typing import List, Type
+from pydantic import BaseModel
+from cognee.infrastructure.llm.LLMGateway import LLMGateway
+from cognee.infrastructure.llm.config import (
+    get_llm_config,
+)
+
+
+async def extract_event_entities(content: str, response_model: Type[BaseModel]):
+    """
+    Extracts event-related entities from the given content using an LLM with structured output.
+
+    This function loads an event entity extraction prompt from the LLM configuration,
+    renders it into a system prompt, and queries the LLM to produce structured entities
+    that conform to the specified response model.
+
+    Args:
+        content (str): The input text from which to extract event entities.
+        response_model (Type[BaseModel]): A Pydantic model defining the structure of the expected output.
+
+    Returns:
+        BaseModel: An instance of the response_model populated with extracted event entities.
+    """
+    llm_config = get_llm_config()
+
+    prompt_path = llm_config.event_entity_prompt_path
+
+    # Check if the prompt path is an absolute path or just a filename
+    if os.path.isabs(prompt_path):
+        # directory containing the file
+        base_directory = os.path.dirname(prompt_path)
+        # just the filename itself
+        prompt_path = os.path.basename(prompt_path)
+    else:
+        base_directory = None
+
+    system_prompt = LLMGateway.render_prompt(prompt_path, {}, base_directory=base_directory)
+
+    content_graph = await LLMGateway.acreate_structured_output(
+        content, system_prompt, response_model
+    )
+
+    return content_graph

cognee/infrastructure/llm/structured_output_framework/litellm_instructor/extraction/knowledge_graph/__init__.py

@@ -1 +1,2 @@
 from .extract_content_graph import extract_content_graph
+from .extract_event_graph import extract_event_graph

cognee/infrastructure/llm/structured_output_framework/litellm_instructor/extraction/knowledge_graph/extract_event_graph.py

@@ -0,0 +1,46 @@
+import os
+from pydantic import BaseModel
+from typing import Type
+from cognee.infrastructure.llm.LLMGateway import LLMGateway
+
+from cognee.infrastructure.llm.config import (
+    get_llm_config,
+)
+
+
+async def extract_event_graph(content: str, response_model: Type[BaseModel]):
+    """
+    Extracts an event graph from the given content using an LLM with a structured output format.
+
+    This function loads a temporal graph extraction prompt from the LLM configuration,
+    renders it as a system prompt, and queries the LLM to produce a structured event
+    graph matching the specified response model.
+
+    Args:
+        content (str): The input text from which to extract the event graph.
+        response_model (Type[BaseModel]): A Pydantic model defining the structure of the expected output.
+
+    Returns:
+        BaseModel: An instance of the response_model populated with the extracted event graph.
+    """
+
+    llm_config = get_llm_config()
+
+    prompt_path = llm_config.temporal_graph_prompt_path
+
+    # Check if the prompt path is an absolute path or just a filename
+    if os.path.isabs(prompt_path):
+        # directory containing the file
+        base_directory = os.path.dirname(prompt_path)
+        # just the filename itself
+        prompt_path = os.path.basename(prompt_path)
+    else:
+        base_directory = None
+
+    system_prompt = LLMGateway.render_prompt(prompt_path, {}, base_directory=base_directory)
+
+    content_graph = await LLMGateway.acreate_structured_output(
+        content, system_prompt, response_model
+    )
+
+    return content_graph

cognee/infrastructure/llm/structured_output_framework/litellm_instructor/llm/openai/adapter.py

@@ -23,9 +23,12 @@ from cognee.infrastructure.llm.structured_output_framework.litellm_instructor.ll
     sleep_and_retry_sync,
 )
 from cognee.modules.observability.get_observe import get_observe
+from cognee.shared.logging_utils import get_logger
 
 observe = get_observe()
 
+logger = get_logger()
+
 
 class OpenAIAdapter(LLMInterface):
     """
@@ -129,6 +132,7 @@ class OpenAIAdapter(LLMInterface):
                 api_version=self.api_version,
                 response_model=response_model,
                 max_retries=self.MAX_RETRIES,
+                extra_body={"reasoning_effort": "minimal"},
             )
         except (
             ContentFilterFinishReasonError,
@@ -139,7 +143,27 @@ class OpenAIAdapter(LLMInterface):
                 isinstance(error, InstructorRetryException)
                 and "content management policy" not in str(error).lower()
             ):
-                raise error
+                logger.debug(
+                    "LLM Model does not support reasoning_effort parameter, trying call without the parameter."
+                )
+                return await self.aclient.chat.completions.create(
+                    model=self.model,
+                    messages=[
+                        {
+                            "role": "user",
+                            "content": f"""{text_input}""",
+                        },
+                        {
+                            "role": "system",
+                            "content": system_prompt,
+                        },
+                    ],
+                    api_key=self.api_key,
+                    api_base=self.endpoint,
+                    api_version=self.api_version,
+                    response_model=response_model,
+                    max_retries=self.MAX_RETRIES,
+                )
 
             if not (self.fallback_model and self.fallback_api_key):
                 raise ContentPolicyFilterError(

cognee/infrastructure/utils/run_sync.py

@@ -8,8 +8,15 @@ def run_sync(coro, timeout=None):
 
     def runner():
         nonlocal result, exception
+
         try:
-            result = asyncio.run(coro)
+            try:
+                running_loop = asyncio.get_running_loop()
+
+                result = asyncio.run_coroutine_threadsafe(coro, running_loop).result(timeout)
+            except RuntimeError:
+                result = asyncio.run(coro)
+
         except Exception as e:
             exception = e
 

cognee/modules/chunking/models/DocumentChunk.py

@@ -1,8 +1,9 @@
-from typing import List
+from typing import List, Union
 
 from cognee.infrastructure.engine import DataPoint
 from cognee.modules.data.processing.document_types import Document
 from cognee.modules.engine.models import Entity
+from cognee.tasks.temporal_graph.models import Event
 
 
 class DocumentChunk(DataPoint):
@@ -20,7 +21,7 @@ class DocumentChunk(DataPoint):
     - 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).
+    - contains: A list of entities or events contained within the chunk (default is None).
     - metadata: A dictionary to hold meta information related to the chunk, including index
     fields.
     """
@@ -30,6 +31,6 @@ class DocumentChunk(DataPoint):
     chunk_index: int
     cut_type: str
     is_part_of: Document
-    contains: List[Entity] = None
+    contains: List[Union[Entity, Event]] = None
 
     metadata: dict = {"index_fields": ["text"]}

cognee/modules/cloud/exceptions/CloudApiKeyMissingError.py

@@ -0,0 +1,15 @@
+from fastapi import status
+
+from cognee.exceptions.exceptions import CogneeConfigurationError
+
+
+class CloudApiKeyMissingError(CogneeConfigurationError):
+    """Raised when the API key for the cloud service is not provided."""
+
+    def __init__(
+        self,
+        message: str = "Failed to connect to the cloud service. Please add your API key to local instance.",
+        name: str = "CloudApiKeyMissingError",
+        status_code=status.HTTP_400_BAD_REQUEST,
+    ):
+        super().__init__(message, name, status_code)

cognee/modules/cloud/exceptions/CloudConnectionError.py

@@ -0,0 +1,15 @@
+from fastapi import status
+
+from cognee.exceptions.exceptions import CogneeConfigurationError
+
+
+class CloudConnectionError(CogneeConfigurationError):
+    """Raised when the connection to the cloud service fails."""
+
+    def __init__(
+        self,
+        message: str = "Failed to connect to the cloud service. Please check your cloud API key in local instance.",
+        name: str = "CloudConnnectionError",
+        status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
+    ):
+        super().__init__(message, name, status_code)

cognee/modules/cloud/exceptions/__init__.py

@@ -0,0 +1,2 @@
+from .CloudConnectionError import CloudConnectionError
+from .CloudApiKeyMissingError import CloudApiKeyMissingError

cognee/modules/cloud/operations/__init__.py

@@ -0,0 +1 @@
+from .check_api_key import check_api_key

cognee/modules/cloud/operations/check_api_key.py

@@ -0,0 +1,25 @@
+import aiohttp
+
+from cognee.modules.cloud.exceptions import CloudConnectionError
+
+
+async def check_api_key(auth_token: str):
+    cloud_base_url = "http://localhost:8001"
+
+    url = f"{cloud_base_url}/api/api-keys/check"
+    headers = {"X-Api-Key": auth_token}
+
+    try:
+        async with aiohttp.ClientSession() as session:
+            async with session.post(url, headers=headers) as response:
+                if response.status == 200:
+                    return
+                else:
+                    error_text = await response.text()
+
+                    raise CloudConnectionError(
+                        f"Failed to connect to cloud instance: {response.status} - {error_text}"
+                    )
+
+    except Exception as e:
+        raise CloudConnectionError(f"Failed to connect to cloud instance: {str(e)}")

cognee/modules/data/deletion/prune_system.py

@@ -3,7 +3,7 @@ from cognee.infrastructure.databases.graph.get_graph_engine import get_graph_eng
 from cognee.infrastructure.databases.relational import get_relational_engine
 
 
-async def prune_system(graph=True, vector=True, metadata=False):
+async def prune_system(graph=True, vector=True, metadata=True):
     if graph:
         graph_engine = await get_graph_engine()
         await graph_engine.delete_graph()

cognee/modules/data/methods/check_dataset_name.py

@@ -1,3 +1,3 @@
 def check_dataset_name(dataset_name: str):
     if "." in dataset_name or " " in dataset_name:
-        raise ValueError("Dataset name cannot contain spaces or underscores")
+        raise ValueError("Dataset name cannot contain spaces or dots.")

cognee/modules/data/methods/get_dataset_data.py

@@ -12,6 +12,6 @@ async def get_dataset_data(dataset_id: UUID) -> list[Data]:
             select(Data).join(Data.datasets).filter((Dataset.id == dataset_id))
         )
 
-        data = result.scalars().all()
+        data = list(result.scalars().all())
 
         return data

cognee/modules/data/methods/load_or_create_datasets.py

@@ -2,7 +2,7 @@ from typing import List, Union
 from uuid import UUID
 
 from cognee.modules.data.models import Dataset
-from cognee.modules.data.methods import create_authorized_dataset
+from cognee.modules.data.methods.create_authorized_dataset import create_authorized_dataset
 from cognee.modules.data.exceptions import DatasetNotFoundError
 
 

cognee/modules/engine/models/Event.py

@@ -0,0 +1,16 @@
+from typing import Optional, Any
+from pydantic import SkipValidation
+from cognee.infrastructure.engine import DataPoint
+from cognee.modules.engine.models.Timestamp import Timestamp
+from cognee.modules.engine.models.Interval import Interval
+
+
+class Event(DataPoint):
+    name: str
+    description: Optional[str] = None
+    at: Optional[Timestamp] = None
+    during: Optional[Interval] = None
+    location: Optional[str] = None
+    attributes: SkipValidation[Any] = None
+
+    metadata: dict = {"index_fields": ["name"]}

cognee/modules/engine/models/Interval.py

@@ -0,0 +1,8 @@
+from pydantic import Field
+from cognee.infrastructure.engine import DataPoint
+from cognee.modules.engine.models.Timestamp import Timestamp
+
+
+class Interval(DataPoint):
+    time_from: Timestamp = Field(...)
+    time_to: Timestamp = Field(...)

cognee/modules/engine/models/Timestamp.py

@@ -0,0 +1,13 @@
+from pydantic import Field
+from cognee.infrastructure.engine import DataPoint
+
+
+class Timestamp(DataPoint):
+    time_at: int = Field(...)
+    year: int = Field(...)
+    month: int = Field(...)
+    day: int = Field(...)
+    hour: int = Field(...)
+    minute: int = Field(...)
+    second: int = Field(...)
+    timestamp_str: str = Field(...)

cognee/modules/engine/models/__init__.py

@@ -4,3 +4,6 @@ from .TableRow import TableRow
 from .TableType import TableType
 from .node_set import NodeSet
 from .ColumnValue import ColumnValue
+from .Timestamp import Timestamp
+from .Interval import Interval
+from .Event import Event

cognee/modules/engine/utils/__init__.py

@@ -1,3 +1,5 @@
 from .generate_node_id import generate_node_id
 from .generate_node_name import generate_node_name
 from .generate_edge_name import generate_edge_name
+from .generate_event_datapoint import generate_event_datapoint
+from .generate_timestamp_datapoint import generate_timestamp_datapoint

cognee/modules/engine/utils/generate_event_datapoint.py

@@ -0,0 +1,46 @@
+from cognee.modules.engine.models import Interval, Event
+from cognee.modules.engine.utils.generate_timestamp_datapoint import generate_timestamp_datapoint
+
+
+def generate_event_datapoint(event) -> Event:
+    """
+    Generates an Event datapoint from a given event model, including temporal metadata if available.
+
+    The function maps the basic attributes (name, description, location) from the input event
+    and enriches them with temporal information. If start and end times are provided, an
+    Interval is created. If only one timestamp is available, it is added directly. Temporal
+    information is also appended to the event description for context.
+
+    Args:
+        event: An event model instance containing attributes such as name, description,
+               location, time_from, and time_to.
+
+    Returns:
+        Event: A structured Event object with name, description, location, and enriched
+               temporal details.
+    """
+    # Base event data
+    event_data = {
+        "name": event.name,
+        "description": event.description,
+        "location": event.location,
+    }
+
+    # Create timestamps if they exist
+    time_from = generate_timestamp_datapoint(event.time_from) if event.time_from else None
+    time_to = generate_timestamp_datapoint(event.time_to) if event.time_to else None
+
+    # Add temporal information
+    if time_from and time_to:
+        event_data["during"] = Interval(time_from=time_from, time_to=time_to)
+        # Enrich description with temporal info
+        temporal_info = f"\n---\nTime data: {time_from.timestamp_str} to {time_to.timestamp_str}"
+        event_data["description"] = (event_data["description"] or "Event") + temporal_info
+    elif time_from or time_to:
+        timestamp = time_from or time_to
+        event_data["at"] = timestamp
+        # Enrich description with temporal info
+        temporal_info = f"\n---\nTime data: {timestamp.timestamp_str}"
+        event_data["description"] = (event_data["description"] or "Event") + temporal_info
+
+    return Event(**event_data)

cognee/modules/engine/utils/generate_timestamp_datapoint.py

@@ -0,0 +1,51 @@
+from datetime import datetime, timezone
+from cognee.modules.engine.models import Interval, Timestamp, Event
+from cognee.modules.engine.utils import generate_node_id
+
+
+def generate_timestamp_datapoint(ts: Timestamp) -> Timestamp:
+    """
+    Generates a normalized Timestamp datapoint from a given Timestamp model.
+
+    The function converts the provided timestamp into an integer representation,
+    constructs a human-readable string format, and creates a new Timestamp object
+    with a unique identifier.
+
+    Args:
+        ts (Timestamp): The input Timestamp model containing date and time components.
+
+    Returns:
+        Timestamp: A new Timestamp object with a generated ID, integer representation,
+                   original components, and formatted string.
+    """
+
+    time_at = date_to_int(ts)
+    timestamp_str = (
+        f"{ts.year:04d}-{ts.month:02d}-{ts.day:02d} {ts.hour:02d}:{ts.minute:02d}:{ts.second:02d}"
+    )
+    return Timestamp(
+        id=generate_node_id(str(time_at)),
+        time_at=time_at,
+        year=ts.year,
+        month=ts.month,
+        day=ts.day,
+        hour=ts.hour,
+        minute=ts.minute,
+        second=ts.second,
+        timestamp_str=timestamp_str,
+    )
+
+
+def date_to_int(ts: Timestamp) -> int:
+    """
+    Converts a Timestamp model into an integer representation in milliseconds since the Unix epoch (UTC).
+
+    Args:
+        ts (Timestamp): The input Timestamp model containing year, month, day, hour, minute, and second.
+
+    Returns:
+        int: The UTC timestamp in milliseconds since January 1, 1970.
+    """
+    dt = datetime(ts.year, ts.month, ts.day, ts.hour, ts.minute, ts.second, tzinfo=timezone.utc)
+    time = int(dt.timestamp() * 1000)
+    return time

cognee/modules/graph/cognee_graph/CogneeGraph.py

@@ -76,7 +76,7 @@ class CogneeGraph(CogneeAbstractGraph):
             start_time = time.time()
 
             # Determine projection strategy
-            if node_type is not None and node_name is not None:
+            if node_type is not None and node_name not in [None, [], ""]:
                 nodes_data, edges_data = await adapter.get_nodeset_subgraph(
                     node_type=node_type, node_name=node_name
                 )
@@ -180,7 +180,7 @@ class CogneeGraph(CogneeAbstractGraph):
             logger.error(f"Error mapping vector distances to edges: {str(ex)}")
             raise ex
 
-    async def calculate_top_triplet_importances(self, k: int) -> List:
+    async def calculate_top_triplet_importances(self, k: int) -> List[Edge]:
         def score(edge):
             n1 = edge.node1.attributes.get("vector_distance", 1)
             n2 = edge.node2.attributes.get("vector_distance", 1)

cognee/modules/graph/utils/__init__.py

@@ -4,3 +4,4 @@ from .get_model_instance_from_graph import get_model_instance_from_graph
 from .retrieve_existing_edges import retrieve_existing_edges
 from .convert_node_to_data_point import convert_node_to_data_point
 from .deduplicate_nodes_and_edges import deduplicate_nodes_and_edges
+from .resolve_edges_to_text import resolve_edges_to_text

cognee/modules/graph/utils/resolve_edges_to_text.py

@@ -0,0 +1,71 @@
+from typing import List
+from cognee.modules.graph.cognee_graph.CogneeGraphElements import Edge
+
+
+async def resolve_edges_to_text(retrieved_edges: List[Edge]) -> str:
+    """
+    Converts retrieved graph edges into a human-readable string format.
+
+    Parameters:
+    -----------
+
+        - retrieved_edges (list): A list of edges retrieved from the graph.
+
+    Returns:
+    --------
+
+        - str: A formatted string representation of the nodes and their connections.
+    """
+
+    def _get_nodes(retrieved_edges: List[Edge]) -> dict:
+        def _get_title(text: str, first_n_words: int = 7, top_n_words: int = 3) -> str:
+            def _top_n_words(text, stop_words=None, top_n=3, separator=", "):
+                """Concatenates the top N frequent words in text."""
+                if stop_words is None:
+                    from cognee.modules.retrieval.utils.stop_words import DEFAULT_STOP_WORDS
+
+                    stop_words = DEFAULT_STOP_WORDS
+
+                import string
+
+                words = [word.lower().strip(string.punctuation) for word in text.split()]
+
+                if stop_words:
+                    words = [word for word in words if word and word not in stop_words]
+
+                from collections import Counter
+
+                top_words = [word for word, freq in Counter(words).most_common(top_n)]
+
+                return separator.join(top_words)
+
+            """Creates a title, by combining first words with most frequent words from the text."""
+            first_words = text.split()[:first_n_words]
+            top_words = _top_n_words(text, top_n=first_n_words)
+            return f"{' '.join(first_words)}... [{top_words}]"
+
+        """Creates a dictionary of nodes with their names and content."""
+        nodes = {}
+        for edge in retrieved_edges:
+            for node in (edge.node1, edge.node2):
+                if node.id not in nodes:
+                    text = node.attributes.get("text")
+                    if text:
+                        name = _get_title(text)
+                        content = text
+                    else:
+                        name = node.attributes.get("name", "Unnamed Node")
+                        content = node.attributes.get("description", name)
+                    nodes[node.id] = {"node": node, "name": name, "content": content}
+        return nodes
+
+    nodes = _get_nodes(retrieved_edges)
+    node_section = "\n".join(
+        f"Node: {info['name']}\n__node_content_start__\n{info['content']}\n__node_content_end__\n"
+        for info in nodes.values()
+    )
+    connection_section = "\n".join(
+        f"{nodes[edge.node1.id]['name']} --[{edge.attributes['relationship_type']}]--> {nodes[edge.node2.id]['name']}"
+        for edge in retrieved_edges
+    )
+    return f"Nodes:\n{node_section}\n\nConnections:\n{connection_section}"

cognee/modules/memify/__init__.py

@@ -0,0 +1 @@
+from .memify import memify