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

cognee/modules/retrieval/triplet_retriever.py

@@ -1,182 +0,0 @@
-import asyncio
-from typing import Any, Optional, Type, List
-
-from cognee.shared.logging_utils import get_logger
-from cognee.infrastructure.databases.vector import get_vector_engine
-from cognee.modules.retrieval.utils.completion import generate_completion, summarize_text
-from cognee.modules.retrieval.utils.session_cache import (
-    save_conversation_history,
-    get_conversation_history,
-)
-from cognee.modules.retrieval.base_retriever import BaseRetriever
-from cognee.modules.retrieval.exceptions.exceptions import NoDataError
-from cognee.infrastructure.databases.vector.exceptions import CollectionNotFoundError
-from cognee.context_global_variables import session_user
-from cognee.infrastructure.databases.cache.config import CacheConfig
-
-logger = get_logger("TripletRetriever")
-
-
-class TripletRetriever(BaseRetriever):
-    """
-    Retriever for handling LLM-based completion searches using triplets.
-
-    Public methods:
-    - get_context(query: str) -> str
-    - get_completion(query: str, context: Optional[Any] = None) -> Any
-    """
-
-    def __init__(
-        self,
-        user_prompt_path: str = "context_for_question.txt",
-        system_prompt_path: str = "answer_simple_question.txt",
-        system_prompt: Optional[str] = None,
-        top_k: Optional[int] = 5,
-    ):
-        """Initialize retriever with optional custom prompt paths."""
-        self.user_prompt_path = user_prompt_path
-        self.system_prompt_path = system_prompt_path
-        self.top_k = top_k if top_k is not None else 1
-        self.system_prompt = system_prompt
-
-    async def get_context(self, query: str) -> str:
-        """
-        Retrieves relevant triplets as context.
-
-        Fetches triplets based on a query from a vector engine and combines their text.
-        Returns empty string if no triplets are found. Raises NoDataError if the collection is not
-        found.
-
-        Parameters:
-        -----------
-
-            - query (str): The query string used to search for relevant triplets.
-
-        Returns:
-        --------
-
-            - str: A string containing the combined text of the retrieved triplets, or an
-              empty string if none are found.
-        """
-        vector_engine = get_vector_engine()
-
-        try:
-            if not await vector_engine.has_collection(collection_name="Triplet_text"):
-                logger.error("Triplet_text collection not found")
-                raise NoDataError(
-                    "In order to use TRIPLET_COMPLETION first use the create_triplet_embeddings memify pipeline. "
-                )
-
-            found_triplets = await vector_engine.search("Triplet_text", query, limit=self.top_k)
-
-            if len(found_triplets) == 0:
-                return ""
-
-            triplets_payload = [found_triplet.payload["text"] for found_triplet in found_triplets]
-            combined_context = "\n".join(triplets_payload)
-            return combined_context
-        except CollectionNotFoundError as error:
-            logger.error("Triplet_text collection not found")
-            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,
-        session_id: Optional[str] = None,
-        response_model: Type = str,
-    ) -> List[Any]:
-        """
-        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 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)
-            - session_id (Optional[str]): Optional session identifier for caching. If None,
-              defaults to 'default_session'. (default None)
-            - response_model (Type): The Pydantic model type for structured output. (default str)
-
-        Returns:
-        --------
-
-            - Any: The generated completion based on the provided query and context.
-        """
-        if context is None:
-            context = await self.get_context(query)
-
-        cache_config = CacheConfig()
-        user = session_user.get()
-        user_id = getattr(user, "id", None)
-        session_save = user_id and cache_config.caching
-
-        if session_save:
-            completion = await self._get_completion_with_session(
-                query=query,
-                context=context,
-                session_id=session_id,
-                response_model=response_model,
-            )
-        else:
-            completion = await self._get_completion_without_session(
-                query=query,
-                context=context,
-                response_model=response_model,
-            )
-
-        return [completion]
-
-    async def _get_completion_with_session(
-        self,
-        query: str,
-        context: str,
-        session_id: Optional[str],
-        response_model: Type,
-    ) -> Any:
-        """Generate completion with session history and caching."""
-        conversation_history = await get_conversation_history(session_id=session_id)
-
-        context_summary, completion = await asyncio.gather(
-            summarize_text(context),
-            generate_completion(
-                query=query,
-                context=context,
-                user_prompt_path=self.user_prompt_path,
-                system_prompt_path=self.system_prompt_path,
-                system_prompt=self.system_prompt,
-                conversation_history=conversation_history,
-                response_model=response_model,
-            ),
-        )
-
-        await save_conversation_history(
-            query=query,
-            context_summary=context_summary,
-            answer=completion,
-            session_id=session_id,
-        )
-
-        return completion
-
-    async def _get_completion_without_session(
-        self,
-        query: str,
-        context: str,
-        response_model: Type,
-    ) -> Any:
-        """Generate completion without session history."""
-        completion = await generate_completion(
-            query=query,
-            context=context,
-            user_prompt_path=self.user_prompt_path,
-            system_prompt_path=self.system_prompt_path,
-            system_prompt=self.system_prompt,
-            response_model=response_model,
-        )
-
-        return completion

cognee/shared/rate_limiting.py

@@ -1,30 +0,0 @@
-from aiolimiter import AsyncLimiter
-from contextlib import nullcontext
-from cognee.infrastructure.llm.config import get_llm_config
-
-llm_config = get_llm_config()
-
-llm_rate_limiter = AsyncLimiter(
-    llm_config.llm_rate_limit_requests, llm_config.embedding_rate_limit_interval
-)
-embedding_rate_limiter = AsyncLimiter(
-    llm_config.embedding_rate_limit_requests, llm_config.embedding_rate_limit_interval
-)
-
-
-def llm_rate_limiter_context_manager():
-    global llm_rate_limiter
-    if llm_config.llm_rate_limit_enabled:
-        return llm_rate_limiter
-    else:
-        #  Return a no-op context manager if rate limiting is disabled
-        return nullcontext()
-
-
-def embedding_rate_limiter_context_manager():
-    global embedding_rate_limiter
-    if llm_config.embedding_rate_limit_enabled:
-        return embedding_rate_limiter
-    else:
-        #  Return a no-op context manager if rate limiting is disabled
-        return nullcontext()

cognee/tasks/memify/get_triplet_datapoints.py

@@ -1,289 +0,0 @@
-from typing import AsyncGenerator, Dict, Any, List, Optional
-from cognee.infrastructure.databases.graph.get_graph_engine import get_graph_engine
-from cognee.modules.engine.utils import generate_node_id
-from cognee.shared.logging_utils import get_logger
-from cognee.modules.graph.utils.convert_node_to_data_point import get_all_subclasses
-from cognee.infrastructure.engine import DataPoint
-from cognee.modules.engine.models import Triplet
-from cognee.tasks.storage import index_data_points
-
-logger = get_logger("get_triplet_datapoints")
-
-
-def _build_datapoint_type_index_mapping() -> Dict[str, List[str]]:
-    """
-    Build a mapping of DataPoint type names to their index_fields.
-
-    Returns:
-    --------
-        - Dict[str, List[str]]: Mapping of type name to list of index field names
-    """
-    logger.debug("Building DataPoint type to index_fields mapping")
-    subclasses = get_all_subclasses(DataPoint)
-    datapoint_type_index_property = {}
-
-    for subclass in subclasses:
-        if "metadata" in subclass.model_fields:
-            metadata_field = subclass.model_fields["metadata"]
-            default = getattr(metadata_field, "default", None)
-            if isinstance(default, dict):
-                index_fields = default.get("index_fields", [])
-                if index_fields:
-                    datapoint_type_index_property[subclass.__name__] = index_fields
-                    logger.debug(
-                        f"Registered {subclass.__name__} with index_fields: {index_fields}"
-                    )
-
-    logger.info(
-        f"Found {len(datapoint_type_index_property)} DataPoint types with index_fields: "
-        f"{list(datapoint_type_index_property.keys())}"
-    )
-    return datapoint_type_index_property
-
-
-def _extract_embeddable_text(node_or_edge: Dict[str, Any], index_fields: List[str]) -> str:
-    """
-    Extract and concatenate embeddable properties from a node or edge dictionary.
-
-    Parameters:
-    -----------
-        - node_or_edge (Dict[str, Any]): Dictionary containing node or edge properties.
-        - index_fields (List[str]): List of field names to extract and concatenate.
-
-    Returns:
-    --------
-        - str: Concatenated string of all embeddable property values, or empty string if none found.
-    """
-    if not node_or_edge or not index_fields:
-        return ""
-
-    embeddable_values = []
-    for field_name in index_fields:
-        field_value = node_or_edge.get(field_name)
-        if field_value is not None:
-            field_value = str(field_value).strip()
-
-            if field_value:
-                embeddable_values.append(field_value)
-
-    return " ".join(embeddable_values) if embeddable_values else ""
-
-
-def _extract_relationship_text(
-    relationship: Dict[str, Any], datapoint_type_index_property: Dict[str, List[str]]
-) -> str:
-    """
-    Extract relationship text from edge properties.
-
-    Parameters:
-    -----------
-        - relationship (Dict[str, Any]): Dictionary containing relationship properties
-        - datapoint_type_index_property (Dict[str, List[str]]): Mapping of type to index fields
-
-    Returns:
-    --------
-        - str: Extracted relationship text or empty string
-    """
-    if not relationship:
-        return ""
-
-    edge_text = relationship.get("edge_text")
-    if edge_text and isinstance(edge_text, str) and edge_text.strip():
-        return edge_text.strip()
-
-    # Fallback to extracting from EdgeType index_fields
-    edge_type_index_fields = datapoint_type_index_property.get("EdgeType", [])
-    return _extract_embeddable_text(relationship, edge_type_index_fields)
-
-
-def _process_single_triplet(
-    triplet_datapoint: Dict[str, Any],
-    datapoint_type_index_property: Dict[str, List[str]],
-    offset: int,
-    idx: int,
-) -> tuple[Optional[Triplet], Optional[str]]:
-    """
-    Process a single triplet and create a Triplet object.
-
-    Parameters:
-    -----------
-        - triplet_datapoint (Dict[str, Any]): Raw triplet data from graph engine
-        - datapoint_type_index_property (Dict[str, List[str]]): Type to index fields mapping
-        - offset (int): Current batch offset
-        - idx (int): Index within current batch
-
-    Returns:
-    --------
-        - tuple[Optional[Triplet], Optional[str]]: (Triplet object, error message if skipped)
-    """
-    start_node = triplet_datapoint.get("start_node", {})
-    end_node = triplet_datapoint.get("end_node", {})
-    relationship = triplet_datapoint.get("relationship_properties", {})
-
-    start_node_type = start_node.get("type")
-    end_node_type = end_node.get("type")
-
-    start_index_fields = datapoint_type_index_property.get(start_node_type, [])
-    end_index_fields = datapoint_type_index_property.get(end_node_type, [])
-
-    if not start_index_fields:
-        logger.debug(
-            f"No index_fields found for start_node type '{start_node_type}' in triplet {offset + idx}"
-        )
-    if not end_index_fields:
-        logger.debug(
-            f"No index_fields found for end_node type '{end_node_type}' in triplet {offset + idx}"
-        )
-
-    start_node_id = start_node.get("id", "")
-    end_node_id = end_node.get("id", "")
-
-    if not start_node_id or not end_node_id:
-        return None, (
-            f"Skipping triplet at offset {offset + idx}: missing node IDs "
-            f"(start: {start_node_id}, end: {end_node_id})"
-        )
-
-    relationship_text = _extract_relationship_text(relationship, datapoint_type_index_property)
-    start_node_text = _extract_embeddable_text(start_node, start_index_fields)
-    end_node_text = _extract_embeddable_text(end_node, end_index_fields)
-
-    if not start_node_text and not end_node_text and not relationship_text:
-        return None, (
-            f"Skipping triplet at offset {offset + idx}: empty embeddable text "
-            f"(start_node_id: {start_node_id}, end_node_id: {end_node_id})"
-        )
-
-    embeddable_text = f"{start_node_text}-›{relationship_text}-›{end_node_text}".strip()
-
-    relationship_name = relationship.get("relationship_name", "")
-    triplet_id = generate_node_id(str(start_node_id) + str(relationship_name) + str(end_node_id))
-
-    triplet_obj = Triplet(
-        id=triplet_id, from_node_id=start_node_id, to_node_id=end_node_id, text=embeddable_text
-    )
-
-    return triplet_obj, None
-
-
-async def get_triplet_datapoints(
-    data,
-    triplets_batch_size: int = 100,
-) -> AsyncGenerator[Triplet, None]:
-    """
-    Async generator that yields batches of triplet datapoints with embeddable text extracted.
-
-    Each triplet in the batch includes:
-    - Original triplet structure (start_node, relationship_properties, end_node)
-    - Extracted embeddable text for each element based on index_fields
-
-    Parameters:
-    -----------
-        - triplets_batch_size (int): Number of triplets to retrieve per batch. Default is 100.
-
-    Yields:
-    -------
-        - List[Dict[str, Any]]: A batch of triplets, each enriched with embeddable text.
-    """
-    if not data or data == [{}]:
-        logger.info("Fetching graph data for current user")
-
-    logger.info(f"Starting triplet datapoints extraction with batch size: {triplets_batch_size}")
-
-    graph_engine = await get_graph_engine()
-    graph_engine_type = type(graph_engine).__name__
-    logger.debug(f"Using graph engine: {graph_engine_type}")
-
-    if not hasattr(graph_engine, "get_triplets_batch"):
-        error_msg = f"Graph adapter {graph_engine_type} does not support get_triplets_batch method"
-        logger.error(error_msg)
-        raise NotImplementedError(error_msg)
-
-    datapoint_type_index_property = _build_datapoint_type_index_mapping()
-
-    offset = 0
-    total_triplets_processed = 0
-    batch_number = 0
-
-    while True:
-        try:
-            batch_number += 1
-            logger.debug(
-                f"Fetching triplet batch {batch_number} (offset: {offset}, limit: {triplets_batch_size})"
-            )
-
-            triplets_batch = await graph_engine.get_triplets_batch(
-                offset=offset, limit=triplets_batch_size
-            )
-
-            if not triplets_batch:
-                logger.info(f"No more triplets found at offset {offset}. Processing complete.")
-                break
-
-            logger.debug(f"Retrieved {len(triplets_batch)} triplets in batch {batch_number}")
-
-            triplet_datapoints = []
-            skipped_count = 0
-
-            for idx, triplet_datapoint in enumerate(triplets_batch):
-                try:
-                    triplet_obj, error_msg = _process_single_triplet(
-                        triplet_datapoint, datapoint_type_index_property, offset, idx
-                    )
-
-                    if error_msg:
-                        logger.warning(error_msg)
-                        skipped_count += 1
-                        continue
-
-                    if triplet_obj:
-                        triplet_datapoints.append(triplet_obj)
-                        yield triplet_obj
-
-                except Exception as e:
-                    logger.warning(
-                        f"Error processing triplet at offset {offset + idx}: {e}. "
-                        f"Skipping this triplet and continuing."
-                    )
-                    skipped_count += 1
-                    continue
-
-            if skipped_count > 0:
-                logger.warning(
-                    f"Skipped {skipped_count} out of {len(triplets_batch)} triplets in batch {batch_number}"
-                )
-
-            if not triplet_datapoints:
-                logger.warning(
-                    f"No valid triplet datapoints in batch {batch_number} after processing"
-                )
-                offset += len(triplets_batch)
-                if len(triplets_batch) < triplets_batch_size:
-                    break
-                continue
-
-            total_triplets_processed += len(triplet_datapoints)
-            logger.info(
-                f"Batch {batch_number} complete: processed {len(triplet_datapoints)} triplets "
-                f"(total processed: {total_triplets_processed})"
-            )
-
-            offset += len(triplets_batch)
-            if len(triplets_batch) < triplets_batch_size:
-                logger.info(
-                    f"Last batch retrieved (got {len(triplets_batch)} < {triplets_batch_size} triplets). "
-                    f"Processing complete."
-                )
-                break
-
-        except Exception as e:
-            logger.error(
-                f"Error retrieving triplet batch {batch_number} at offset {offset}: {e}",
-                exc_info=True,
-            )
-            raise
-
-    logger.info(
-        f"Triplet datapoints extraction complete. "
-        f"Processed {total_triplets_processed} triplets across {batch_number} batch(es)."
-    )

cognee/tests/integration/retrieval/test_triplet_retriever.py

@@ -1,84 +0,0 @@
-import os
-import pytest
-import pathlib
-import pytest_asyncio
-import cognee
-
-from cognee.low_level import setup
-from cognee.tasks.storage import add_data_points
-from cognee.modules.retrieval.exceptions.exceptions import NoDataError
-from cognee.modules.retrieval.triplet_retriever import TripletRetriever
-from cognee.modules.engine.models import Triplet
-
-
-@pytest_asyncio.fixture
-async def setup_test_environment_with_triplets():
-    """Set up a clean test environment with triplets."""
-    base_dir = pathlib.Path(__file__).parent.parent.parent.parent
-    system_directory_path = str(base_dir / ".cognee_system/test_triplet_retriever_context_simple")
-    data_directory_path = str(base_dir / ".data_storage/test_triplet_retriever_context_simple")
-
-    cognee.config.system_root_directory(system_directory_path)
-    cognee.config.data_root_directory(data_directory_path)
-
-    await cognee.prune.prune_data()
-    await cognee.prune.prune_system(metadata=True)
-    await setup()
-
-    triplet1 = Triplet(
-        from_node_id="node1",
-        to_node_id="node2",
-        text="Alice knows Bob",
-    )
-    triplet2 = Triplet(
-        from_node_id="node2",
-        to_node_id="node3",
-        text="Bob works at Tech Corp",
-    )
-
-    triplets = [triplet1, triplet2]
-    await add_data_points(triplets)
-
-    yield
-
-    try:
-        await cognee.prune.prune_data()
-        await cognee.prune.prune_system(metadata=True)
-    except Exception:
-        pass
-
-
-@pytest_asyncio.fixture
-async def setup_test_environment_empty():
-    """Set up a clean test environment without triplets."""
-    base_dir = pathlib.Path(__file__).parent.parent.parent.parent
-    system_directory_path = str(
-        base_dir / ".cognee_system/test_triplet_retriever_context_empty_collection"
-    )
-    data_directory_path = str(
-        base_dir / ".data_storage/test_triplet_retriever_context_empty_collection"
-    )
-
-    cognee.config.system_root_directory(system_directory_path)
-    cognee.config.data_root_directory(data_directory_path)
-
-    await cognee.prune.prune_data()
-    await cognee.prune.prune_system(metadata=True)
-
-    yield
-
-    try:
-        await cognee.prune.prune_data()
-        await cognee.prune.prune_system(metadata=True)
-    except Exception:
-        pass
-
-
-@pytest.mark.asyncio
-async def test_triplet_retriever_context_simple(setup_test_environment_with_triplets):
-    """Integration test: verify TripletRetriever can retrieve triplet context."""
-    retriever = TripletRetriever(top_k=5)
-
-    context = await retriever.get_context("Alice")
-
-    assert "Alice knows Bob" in context, "Failed to get Alice triplet"