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

cognee/modules/retrieval/natural_language_retriever.py

@@ -1,7 +1,6 @@
 from typing import Any, Optional
 from cognee.shared.logging_utils import get_logger
 from cognee.infrastructure.databases.graph import get_graph_engine
-from cognee.infrastructure.databases.graph.networkx.adapter import NetworkXAdapter
 from cognee.infrastructure.llm.LLMGateway import LLMGateway
 from cognee.modules.retrieval.base_retriever import BaseRetriever
 from cognee.modules.retrieval.exceptions import SearchTypeNotSupported
@@ -123,9 +122,6 @@ class NaturalLanguageRetriever(BaseRetriever):
         """
         graph_engine = await get_graph_engine()
 
-        if isinstance(graph_engine, (NetworkXAdapter)):
-            raise SearchTypeNotSupported("Natural language search type not supported.")
-
         return await self._execute_cypher_query(query, graph_engine)
 
     async def get_completion(self, query: str, context: Optional[Any] = None) -> Any:

cognee/modules/retrieval/summaries_retriever.py

@@ -62,7 +62,7 @@ class SummariesRetriever(BaseRetriever):
         logger.info(f"Returning {len(summary_payloads)} summary payloads")
         return summary_payloads
 
-    async def get_completion(self, query: str, context: Optional[Any] = None) -> Any:
+    async def get_completion(self, query: str, context: Optional[Any] = None, **kwargs) -> Any:
         """
         Generates a completion using summaries context.
 

cognee/modules/retrieval/temporal_retriever.py

@@ -0,0 +1,152 @@
+import os
+from typing import Any, Optional, List, Type
+
+
+from operator import itemgetter
+from cognee.infrastructure.databases.vector import get_vector_engine
+from cognee.modules.retrieval.utils.completion import generate_completion
+from cognee.infrastructure.databases.graph import get_graph_engine
+from cognee.infrastructure.llm import LLMGateway
+from cognee.modules.retrieval.graph_completion_retriever import GraphCompletionRetriever
+from cognee.shared.logging_utils import get_logger
+
+
+from cognee.tasks.temporal_graph.models import QueryInterval
+
+logger = get_logger()
+
+
+class TemporalRetriever(GraphCompletionRetriever):
+    """
+    Handles graph completion by generating responses based on a series of interactions with
+    a language model. This class extends from GraphCompletionRetriever and is designed to
+    manage the retrieval and validation process for user queries, integrating follow-up
+    questions based on reasoning. The public methods are:
+
+    - get_completion
+
+    Instance variables include:
+    - validation_system_prompt_path
+    - validation_user_prompt_path
+    - followup_system_prompt_path
+    - followup_user_prompt_path
+    """
+
+    def __init__(
+        self,
+        user_prompt_path: str = "graph_context_for_question.txt",
+        system_prompt_path: str = "answer_simple_question.txt",
+        time_extraction_prompt_path: str = "extract_query_time.txt",
+        top_k: Optional[int] = 5,
+        node_type: Optional[Type] = None,
+        node_name: Optional[List[str]] = None,
+    ):
+        super().__init__(
+            user_prompt_path=user_prompt_path,
+            system_prompt_path=system_prompt_path,
+            top_k=top_k,
+            node_type=node_type,
+            node_name=node_name,
+        )
+        self.user_prompt_path = user_prompt_path
+        self.system_prompt_path = system_prompt_path
+        self.time_extraction_prompt_path = time_extraction_prompt_path
+        self.top_k = top_k if top_k is not None else 5
+        self.node_type = node_type
+        self.node_name = node_name
+
+    def descriptions_to_string(self, results):
+        descs = []
+        for entry in results:
+            d = entry.get("description")
+            if d:
+                descs.append(d.strip())
+        return "\n#####################\n".join(descs)
+
+    async def extract_time_from_query(self, query: str):
+        prompt_path = self.time_extraction_prompt_path
+
+        if os.path.isabs(prompt_path):
+            base_directory = os.path.dirname(prompt_path)
+            prompt_path = os.path.basename(prompt_path)
+        else:
+            base_directory = None
+
+        system_prompt = LLMGateway.render_prompt(prompt_path, {}, base_directory=base_directory)
+
+        interval = await LLMGateway.acreate_structured_output(query, system_prompt, QueryInterval)
+
+        time_from = interval.starts_at
+        time_to = interval.ends_at
+
+        return time_from, time_to
+
+    async def filter_top_k_events(self, relevant_events, scored_results):
+        # Build a score lookup from vector search results
+        score_lookup = {res.payload["id"]: res.score for res in scored_results}
+
+        events_with_scores = []
+        for event in relevant_events[0]["events"]:
+            score = score_lookup.get(event["id"], float("inf"))
+            events_with_scores.append({**event, "score": score})
+
+        events_with_scores.sort(key=itemgetter("score"))
+
+        return events_with_scores[: self.top_k]
+
+    async def get_context(self, query: str) -> Any:
+        """Retrieves context based on the query."""
+
+        time_from, time_to = await self.extract_time_from_query(query)
+
+        graph_engine = await get_graph_engine()
+
+        triplets = []
+
+        if time_from and time_to:
+            ids = await graph_engine.collect_time_ids(time_from=time_from, time_to=time_to)
+        elif time_from:
+            ids = await graph_engine.collect_time_ids(time_from=time_from)
+        elif time_to:
+            ids = await graph_engine.collect_time_ids(time_to=time_to)
+        else:
+            logger.info(
+                "No timestamps identified based on the query, performing retrieval using triplet search on events and entities."
+            )
+            triplets = await self.get_context(query)
+            return await self.resolve_edges_to_text(triplets)
+
+        if ids:
+            relevant_events = await graph_engine.collect_events(ids=ids)
+        else:
+            logger.info(
+                "No events identified based on timestamp filtering, performing retrieval using triplet search on events and entities."
+            )
+            triplets = await self.get_context(query)
+            return await self.resolve_edges_to_text(triplets)
+
+        vector_engine = get_vector_engine()
+        query_vector = (await vector_engine.embedding_engine.embed_text([query]))[0]
+
+        vector_search_results = await vector_engine.search(
+            collection_name="Event_name", query_vector=query_vector, limit=0
+        )
+
+        top_k_events = await self.filter_top_k_events(relevant_events, vector_search_results)
+
+        return self.descriptions_to_string(top_k_events)
+
+    async def get_completion(self, query: str, context: Optional[str] = None) -> str:
+        """Generates a response using the query and optional context."""
+        if not context:
+            context = await self.get_context(query=query)
+
+        if context:
+            completion = await generate_completion(
+                query=query,
+                context=context,
+                user_prompt_path=self.user_prompt_path,
+                system_prompt_path=self.system_prompt_path,
+            )
+
+        return completion

cognee/modules/retrieval/user_qa_feedback.py

@@ -0,0 +1,83 @@
+from typing import Any, Optional, List
+
+from uuid import NAMESPACE_OID, uuid5, UUID
+from cognee.infrastructure.databases.graph import get_graph_engine
+from cognee.infrastructure.llm import LLMGateway
+from cognee.modules.engine.models import NodeSet
+from cognee.shared.logging_utils import get_logger
+from cognee.modules.retrieval.base_feedback import BaseFeedback
+from cognee.modules.retrieval.utils.models import CogneeUserFeedback
+from cognee.modules.retrieval.utils.models import UserFeedbackEvaluation
+from cognee.tasks.storage import add_data_points
+
+logger = get_logger("CompletionRetriever")
+
+
+class UserQAFeedback(BaseFeedback):
+    """
+    Interface for handling user feedback queries.
+    Public methods:
+    - get_context(query: str) -> str
+    - get_completion(query: str, context: Optional[Any] = None) -> Any
+    """
+
+    def __init__(self, last_k: Optional[int] = 1) -> None:
+        """Initialize retriever with optional custom prompt paths."""
+        self.last_k = last_k
+
+    async def add_feedback(self, feedback_text: str) -> List[str]:
+        feedback_sentiment = await LLMGateway.acreate_structured_output(
+            text_input=feedback_text,
+            system_prompt="You are a sentiment analysis assistant. For each piece of user feedback you receive, return exactly one of: Positive, Negative, or Neutral classification and a corresponding score from -5 (worst negative) to 5 (best positive)",
+            response_model=UserFeedbackEvaluation,
+        )
+
+        graph_engine = await get_graph_engine()
+        last_interaction_ids = await graph_engine.get_last_user_interaction_ids(limit=self.last_k)
+
+        nodeset_name = "UserQAFeedbacks"
+        feedbacks_node_set = NodeSet(id=uuid5(NAMESPACE_OID, name=nodeset_name), name=nodeset_name)
+        feedback_id = uuid5(NAMESPACE_OID, name=feedback_text)
+
+        cognee_user_feedback = CogneeUserFeedback(
+            id=feedback_id,
+            feedback=feedback_text,
+            sentiment=feedback_sentiment.evaluation.value,
+            score=feedback_sentiment.score,
+            belongs_to_set=feedbacks_node_set,
+        )
+
+        await add_data_points(data_points=[cognee_user_feedback], update_edge_collection=False)
+
+        relationships = []
+        relationship_name = "gives_feedback_to"
+        to_node_ids = []
+
+        for interaction_id in last_interaction_ids:
+            target_id_1 = feedback_id
+            target_id_2 = UUID(interaction_id)
+
+            if target_id_1 and target_id_2:
+                relationships.append(
+                    (
+                        target_id_1,
+                        target_id_2,
+                        relationship_name,
+                        {
+                            "relationship_name": relationship_name,
+                            "source_node_id": target_id_1,
+                            "target_node_id": target_id_2,
+                            "ontology_valid": False,
+                        },
+                    )
+                )
+                to_node_ids.append(str(target_id_2))
+
+        if len(relationships) > 0:
+            graph_engine = await get_graph_engine()
+            await graph_engine.add_edges(relationships)
+            await graph_engine.apply_feedback_weight(
+                node_ids=to_node_ids, weight=feedback_sentiment.score
+            )
+
+        return [feedback_text]

cognee/modules/retrieval/utils/brute_force_triplet_search.py

@@ -8,7 +8,7 @@ from cognee.infrastructure.databases.vector.exceptions import CollectionNotFound
 from cognee.infrastructure.databases.graph import get_graph_engine
 from cognee.infrastructure.databases.vector import get_vector_engine
 from cognee.modules.graph.cognee_graph.CogneeGraph import CogneeGraph
-from cognee.modules.users.methods import get_default_user
+from cognee.modules.graph.cognee_graph.CogneeGraphElements import Edge
 from cognee.modules.users.models import User
 from cognee.shared.utils import send_telemetry
 
@@ -63,9 +63,10 @@ async def get_memory_fragment(
     if properties_to_project is None:
         properties_to_project = ["id", "description", "name", "type", "text"]
 
+    memory_fragment = CogneeGraph()
+
     try:
         graph_engine = await get_graph_engine()
-        memory_fragment = CogneeGraph()
 
         await memory_fragment.project_graph_from_db(
             graph_engine,
@@ -87,41 +88,15 @@ async def get_memory_fragment(
 
 
 async def brute_force_triplet_search(
-    query: str,
-    user: User = None,
-    top_k: int = 5,
-    collections: List[str] = None,
-    properties_to_project: List[str] = None,
-    memory_fragment: Optional[CogneeGraph] = None,
-    node_type: Optional[Type] = None,
-    node_name: Optional[List[str]] = None,
-) -> list:
-    if user is None:
-        user = await get_default_user()
-
-    retrieved_results = await brute_force_search(
-        query,
-        user,
-        top_k,
-        collections=collections,
-        properties_to_project=properties_to_project,
-        memory_fragment=memory_fragment,
-        node_type=node_type,
-        node_name=node_name,
-    )
-    return retrieved_results
-
-
-async def brute_force_search(
     query: str,
     user: User,
-    top_k: int,
-    collections: List[str] = None,
-    properties_to_project: List[str] = None,
+    top_k: int = 5,
+    collections: Optional[List[str]] = None,
+    properties_to_project: Optional[List[str]] = None,
     memory_fragment: Optional[CogneeGraph] = None,
     node_type: Optional[Type] = None,
     node_name: Optional[List[str]] = None,
-) -> list:
+) -> List[Edge]:
     """
     Performs a brute force search to retrieve the top triplets from the graph.
 

cognee/modules/retrieval/utils/completion.py

@@ -1,3 +1,4 @@
+from typing import Optional
 from cognee.infrastructure.llm.LLMGateway import LLMGateway
 
 
@@ -6,11 +7,14 @@ async def generate_completion(
     context: str,
     user_prompt_path: str,
     system_prompt_path: str,
+    system_prompt: Optional[str] = None,
 ) -> str:
     """Generates a completion using LLM with given context and prompts."""
     args = {"question": query, "context": context}
     user_prompt = LLMGateway.render_prompt(user_prompt_path, args)
-    system_prompt = LLMGateway.read_query_prompt(system_prompt_path)
+    system_prompt = (
+        system_prompt if system_prompt else LLMGateway.read_query_prompt(system_prompt_path)
+    )
 
     return await LLMGateway.acreate_structured_output(
         text_input=user_prompt,
@@ -21,10 +25,13 @@ async def generate_completion(
 
 async def summarize_text(
     text: str,
-    prompt_path: str = "summarize_search_results.txt",
+    system_prompt_path: str = "summarize_search_results.txt",
+    system_prompt: str = None,
 ) -> str:
     """Summarizes text using LLM with the specified prompt."""
-    system_prompt = LLMGateway.read_query_prompt(prompt_path)
+    system_prompt = (
+        system_prompt if system_prompt else LLMGateway.read_query_prompt(system_prompt_path)
+    )
 
     return await LLMGateway.acreate_structured_output(
         text_input=text,

cognee/modules/retrieval/utils/extract_uuid_from_node.py

@@ -0,0 +1,18 @@
+from typing import Any, Optional
+from uuid import UUID
+
+
+def extract_uuid_from_node(node: Any) -> Optional[UUID]:
+    """
+    Try to pull a UUID string out of node.id or node.properties['id'],
+    then return a UUID instance (or None if neither exists).
+    """
+    id_str = None
+    if not id_str:
+        id_str = getattr(node, "id", None)
+
+    if hasattr(node, "attributes") and not id_str:
+        id_str = node.attributes.get("id", None)
+
+    id = UUID(id_str) if isinstance(id_str, str) else None
+    return id

cognee/modules/retrieval/utils/models.py

@@ -0,0 +1,40 @@
+from typing import Optional
+from cognee.infrastructure.engine.models.DataPoint import DataPoint
+from cognee.modules.engine.models.node_set import NodeSet
+from enum import Enum
+from pydantic import BaseModel, Field, confloat
+
+
+class CogneeUserInteraction(DataPoint):
+    """User - Cognee interaction"""
+
+    question: str
+    answer: str
+    context: str
+    belongs_to_set: Optional[NodeSet] = None
+
+
+class CogneeUserFeedback(DataPoint):
+    """User - Cognee Feedback"""
+
+    feedback: str
+    sentiment: str
+    score: float
+    belongs_to_set: Optional[NodeSet] = None
+
+
+class UserFeedbackSentiment(str, Enum):
+    """User - User feedback sentiment"""
+
+    positive = "positive"
+    negative = "negative"
+    neutral = "neutral"
+
+
+class UserFeedbackEvaluation(BaseModel):
+    """User - User feedback evaluation"""
+
+    score: confloat(ge=-5, le=5) = Field(
+        ..., description="Sentiment score from -5 (negative) to +5 (positive)"
+    )
+    evaluation: UserFeedbackSentiment

cognee/modules/search/methods/get_search_type_tools.py

@@ -0,0 +1,168 @@
+from typing import Callable, List, Optional, Type
+
+from cognee.modules.engine.models.node_set import NodeSet
+from cognee.modules.search.types import SearchType
+from cognee.modules.search.operations import select_search_type
+from cognee.modules.search.exceptions import UnsupportedSearchTypeError
+
+# Retrievers
+from cognee.modules.retrieval.user_qa_feedback import UserQAFeedback
+from cognee.modules.retrieval.chunks_retriever import ChunksRetriever
+from cognee.modules.retrieval.insights_retriever import InsightsRetriever
+from cognee.modules.retrieval.summaries_retriever import SummariesRetriever
+from cognee.modules.retrieval.completion_retriever import CompletionRetriever
+from cognee.modules.retrieval.graph_completion_retriever import GraphCompletionRetriever
+from cognee.modules.retrieval.temporal_retriever import TemporalRetriever
+from cognee.modules.retrieval.coding_rules_retriever import CodingRulesRetriever
+from cognee.modules.retrieval.graph_summary_completion_retriever import (
+    GraphSummaryCompletionRetriever,
+)
+from cognee.modules.retrieval.graph_completion_cot_retriever import GraphCompletionCotRetriever
+from cognee.modules.retrieval.graph_completion_context_extension_retriever import (
+    GraphCompletionContextExtensionRetriever,
+)
+from cognee.modules.retrieval.code_retriever import CodeRetriever
+from cognee.modules.retrieval.cypher_search_retriever import CypherSearchRetriever
+from cognee.modules.retrieval.natural_language_retriever import NaturalLanguageRetriever
+
+
+async def get_search_type_tools(
+    query_type: SearchType,
+    query_text: str,
+    system_prompt_path: str = "answer_simple_question.txt",
+    system_prompt: Optional[str] = None,
+    top_k: int = 10,
+    node_type: Optional[Type] = NodeSet,
+    node_name: Optional[List[str]] = None,
+    save_interaction: bool = False,
+    last_k: Optional[int] = None,
+) -> list:
+    search_tasks: dict[SearchType, List[Callable]] = {
+        SearchType.SUMMARIES: [
+            SummariesRetriever(top_k=top_k).get_completion,
+            SummariesRetriever(top_k=top_k).get_context,
+        ],
+        SearchType.INSIGHTS: [
+            InsightsRetriever(top_k=top_k).get_completion,
+            InsightsRetriever(top_k=top_k).get_context,
+        ],
+        SearchType.CHUNKS: [
+            ChunksRetriever(top_k=top_k).get_completion,
+            ChunksRetriever(top_k=top_k).get_context,
+        ],
+        SearchType.RAG_COMPLETION: [
+            CompletionRetriever(
+                system_prompt_path=system_prompt_path,
+                top_k=top_k,
+                system_prompt=system_prompt,
+            ).get_completion,
+            CompletionRetriever(
+                system_prompt_path=system_prompt_path,
+                top_k=top_k,
+                system_prompt=system_prompt,
+            ).get_context,
+        ],
+        SearchType.GRAPH_COMPLETION: [
+            GraphCompletionRetriever(
+                system_prompt_path=system_prompt_path,
+                top_k=top_k,
+                node_type=node_type,
+                node_name=node_name,
+                save_interaction=save_interaction,
+                system_prompt=system_prompt,
+            ).get_completion,
+            GraphCompletionRetriever(
+                system_prompt_path=system_prompt_path,
+                top_k=top_k,
+                node_type=node_type,
+                node_name=node_name,
+                save_interaction=save_interaction,
+                system_prompt=system_prompt,
+            ).get_context,
+        ],
+        SearchType.GRAPH_COMPLETION_COT: [
+            GraphCompletionCotRetriever(
+                system_prompt_path=system_prompt_path,
+                top_k=top_k,
+                node_type=node_type,
+                node_name=node_name,
+                save_interaction=save_interaction,
+                system_prompt=system_prompt,
+            ).get_completion,
+            GraphCompletionCotRetriever(
+                system_prompt_path=system_prompt_path,
+                top_k=top_k,
+                node_type=node_type,
+                node_name=node_name,
+                save_interaction=save_interaction,
+                system_prompt=system_prompt,
+            ).get_context,
+        ],
+        SearchType.GRAPH_COMPLETION_CONTEXT_EXTENSION: [
+            GraphCompletionContextExtensionRetriever(
+                system_prompt_path=system_prompt_path,
+                top_k=top_k,
+                node_type=node_type,
+                node_name=node_name,
+                save_interaction=save_interaction,
+                system_prompt=system_prompt,
+            ).get_completion,
+            GraphCompletionContextExtensionRetriever(
+                system_prompt_path=system_prompt_path,
+                top_k=top_k,
+                node_type=node_type,
+                node_name=node_name,
+                save_interaction=save_interaction,
+                system_prompt=system_prompt,
+            ).get_context,
+        ],
+        SearchType.GRAPH_SUMMARY_COMPLETION: [
+            GraphSummaryCompletionRetriever(
+                system_prompt_path=system_prompt_path,
+                top_k=top_k,
+                node_type=node_type,
+                node_name=node_name,
+                save_interaction=save_interaction,
+                system_prompt=system_prompt,
+            ).get_completion,
+            GraphSummaryCompletionRetriever(
+                system_prompt_path=system_prompt_path,
+                top_k=top_k,
+                node_type=node_type,
+                node_name=node_name,
+                save_interaction=save_interaction,
+                system_prompt=system_prompt,
+            ).get_context,
+        ],
+        SearchType.CODE: [
+            CodeRetriever(top_k=top_k).get_completion,
+            CodeRetriever(top_k=top_k).get_context,
+        ],
+        SearchType.CYPHER: [
+            CypherSearchRetriever().get_completion,
+            CypherSearchRetriever().get_context,
+        ],
+        SearchType.NATURAL_LANGUAGE: [
+            NaturalLanguageRetriever().get_completion,
+            NaturalLanguageRetriever().get_context,
+        ],
+        SearchType.FEEDBACK: [UserQAFeedback(last_k=last_k).add_feedback],
+        SearchType.TEMPORAL: [
+            TemporalRetriever(top_k=top_k).get_completion,
+            TemporalRetriever(top_k=top_k).get_context,
+        ],
+        SearchType.CODING_RULES: [
+            CodingRulesRetriever(rules_nodeset_name=node_name).get_existing_rules,
+        ],
+    }
+
+    # If the query type is FEELING_LUCKY, select the search type intelligently
+    if query_type is SearchType.FEELING_LUCKY:
+        query_type = await select_search_type(query_text)
+
+    search_type_tools = search_tasks.get(query_type)
+
+    if not search_type_tools:
+        raise UnsupportedSearchTypeError(str(query_type))
+
+    return search_type_tools

cognee/modules/search/methods/no_access_control_search.py

@@ -0,0 +1,47 @@
+from typing import Any, List, Optional, Tuple, Type, Union
+
+from cognee.modules.data.models.Dataset import Dataset
+from cognee.modules.engine.models.node_set import NodeSet
+from cognee.modules.graph.cognee_graph.CogneeGraphElements import Edge
+from cognee.modules.search.types import SearchType
+
+from .get_search_type_tools import get_search_type_tools
+
+
+async def no_access_control_search(
+    query_type: SearchType,
+    query_text: str,
+    system_prompt_path: str = "answer_simple_question.txt",
+    system_prompt: Optional[str] = None,
+    top_k: int = 10,
+    node_type: Optional[Type] = NodeSet,
+    node_name: Optional[List[str]] = None,
+    save_interaction: bool = False,
+    last_k: Optional[int] = None,
+    only_context: bool = False,
+) -> Tuple[Any, Union[str, List[Edge]], List[Dataset]]:
+    search_tools = await get_search_type_tools(
+        query_type=query_type,
+        query_text=query_text,
+        system_prompt_path=system_prompt_path,
+        system_prompt=system_prompt,
+        top_k=top_k,
+        node_type=node_type,
+        node_name=node_name,
+        save_interaction=save_interaction,
+        last_k=last_k,
+    )
+    if len(search_tools) == 2:
+        [get_completion, get_context] = search_tools
+
+        if only_context:
+            return await get_context(query_text)
+
+        context = await get_context(query_text)
+        result = await get_completion(query_text, context)
+    else:
+        unknown_tool = search_tools[0]
+        result = await unknown_tool(query_text)
+        context = ""
+
+    return result, context, []