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

cognee/modules/retrieval/graph_completion_cot_retriever.py

@@ -35,6 +35,7 @@ class GraphCompletionCotRetriever(GraphCompletionRetriever):
         top_k: Optional[int] = 5,
         node_type: Optional[Type] = None,
         node_name: Optional[List[str]] = None,
+        save_interaction: bool = False,
     ):
         super().__init__(
             user_prompt_path=user_prompt_path,
@@ -42,6 +43,7 @@ class GraphCompletionCotRetriever(GraphCompletionRetriever):
             top_k=top_k,
             node_type=node_type,
             node_name=node_name,
+            save_interaction=save_interaction,
         )
         self.validation_system_prompt_path = validation_system_prompt_path
         self.validation_user_prompt_path = validation_user_prompt_path
@@ -75,7 +77,7 @@ class GraphCompletionCotRetriever(GraphCompletionRetriever):
         """
         followup_question = ""
         triplets = []
-        answer = [""]
+        completion = [""]
 
         for round_idx in range(max_iter + 1):
             if round_idx == 0:
@@ -85,15 +87,15 @@ class GraphCompletionCotRetriever(GraphCompletionRetriever):
                 triplets += await self.get_triplets(followup_question)
                 context = await self.resolve_edges_to_text(list(set(triplets)))
 
-            answer = await generate_completion(
+            completion = await generate_completion(
                 query=query,
                 context=context,
                 user_prompt_path=self.user_prompt_path,
                 system_prompt_path=self.system_prompt_path,
             )
-            logger.info(f"Chain-of-thought: round {round_idx} - answer: {answer}")
+            logger.info(f"Chain-of-thought: round {round_idx} - answer: {completion}")
             if round_idx < max_iter:
-                valid_args = {"query": query, "answer": answer, "context": context}
+                valid_args = {"query": query, "answer": completion, "context": context}
                 valid_user_prompt = LLMGateway.render_prompt(
                     filename=self.validation_user_prompt_path, context=valid_args
                 )
@@ -106,7 +108,7 @@ class GraphCompletionCotRetriever(GraphCompletionRetriever):
                     system_prompt=valid_system_prompt,
                     response_model=str,
                 )
-                followup_args = {"query": query, "answer": answer, "reasoning": reasoning}
+                followup_args = {"query": query, "answer": completion, "reasoning": reasoning}
                 followup_prompt = LLMGateway.render_prompt(
                     filename=self.followup_user_prompt_path, context=followup_args
                 )
@@ -121,4 +123,9 @@ class GraphCompletionCotRetriever(GraphCompletionRetriever):
                     f"Chain-of-thought: round {round_idx} - follow-up question: {followup_question}"
                 )
 
-        return [answer]
+        if self.save_interaction and context and triplets and completion:
+            await self.save_qa(
+                question=query, answer=completion, context=context, triplets=triplets
+            )
+
+        return [completion]

cognee/modules/retrieval/graph_completion_retriever.py

@@ -1,14 +1,20 @@
-from typing import Any, Optional, Type, List
+from typing import Any, Optional, Type, List, Coroutine
 from collections import Counter
+from uuid import NAMESPACE_OID, uuid5
 import string
 
 from cognee.infrastructure.engine import DataPoint
+from cognee.tasks.storage import add_data_points
 from cognee.modules.graph.utils.convert_node_to_data_point import get_all_subclasses
 from cognee.modules.retrieval.base_retriever import BaseRetriever
 from cognee.modules.retrieval.utils.brute_force_triplet_search import brute_force_triplet_search
 from cognee.modules.retrieval.utils.completion import generate_completion
 from cognee.modules.retrieval.utils.stop_words import DEFAULT_STOP_WORDS
 from cognee.shared.logging_utils import get_logger
+from cognee.modules.retrieval.utils.extract_uuid_from_node import extract_uuid_from_node
+from cognee.modules.retrieval.utils.models import CogneeUserInteraction
+from cognee.modules.engine.models.node_set import NodeSet
+from cognee.infrastructure.databases.graph import get_graph_engine
 
 logger = get_logger("GraphCompletionRetriever")
 
@@ -33,8 +39,10 @@ class GraphCompletionRetriever(BaseRetriever):
         top_k: Optional[int] = 5,
         node_type: Optional[Type] = None,
         node_name: Optional[List[str]] = None,
+        save_interaction: bool = False,
     ):
         """Initialize retriever with prompt paths and search parameters."""
+        self.save_interaction = save_interaction
         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 5
@@ -118,7 +126,7 @@ class GraphCompletionRetriever(BaseRetriever):
 
         return found_triplets
 
-    async def get_context(self, query: str) -> str:
+    async def get_context(self, query: str) -> str | tuple[str, list]:
         """
         Retrieves and resolves graph triplets into context based on a query.
 
@@ -137,9 +145,11 @@ class GraphCompletionRetriever(BaseRetriever):
 
         if len(triplets) == 0:
             logger.warning("Empty context was provided to the completion")
-            return ""
+            return "", triplets
 
-        return await self.resolve_edges_to_text(triplets)
+        context = await self.resolve_edges_to_text(triplets)
+
+        return context, triplets
 
     async def get_completion(self, query: str, context: Optional[Any] = None) -> Any:
         """
@@ -157,8 +167,10 @@ class GraphCompletionRetriever(BaseRetriever):
 
             - Any: A generated completion based on the query and context provided.
         """
+        triplets = None
+
         if context is None:
-            context = await self.get_context(query)
+            context, triplets = await self.get_context(query)
 
         completion = await generate_completion(
             query=query,
@@ -166,6 +178,12 @@ class GraphCompletionRetriever(BaseRetriever):
             user_prompt_path=self.user_prompt_path,
             system_prompt_path=self.system_prompt_path,
         )
+
+        if self.save_interaction and context and triplets and completion:
+            await self.save_qa(
+                question=query, answer=completion, context=context, triplets=triplets
+            )
+
         return [completion]
 
     def _top_n_words(self, text, stop_words=None, top_n=3, separator=", "):
@@ -187,3 +205,69 @@ class GraphCompletionRetriever(BaseRetriever):
         first_n_words = text.split()[:first_n_words]
         top_n_words = self._top_n_words(text, top_n=top_n_words)
         return f"{' '.join(first_n_words)}... [{top_n_words}]"
+
+    async def save_qa(self, question: str, answer: str, context: str, triplets: List) -> None:
+        """
+        Saves a question and answer pair for later analysis or storage.
+        Parameters:
+        -----------
+            - question (str): The question text.
+            - answer (str): The answer text.
+            - context (str): The context text.
+            - triplets (List): A list of triples retrieved from the graph.
+        """
+        nodeset_name = "Interactions"
+        interactions_node_set = NodeSet(
+            id=uuid5(NAMESPACE_OID, name=nodeset_name), name=nodeset_name
+        )
+        source_id = uuid5(NAMESPACE_OID, name=(question + answer + context))
+
+        cognee_user_interaction = CogneeUserInteraction(
+            id=source_id,
+            question=question,
+            answer=answer,
+            context=context,
+            belongs_to_set=interactions_node_set,
+        )
+
+        await add_data_points(data_points=[cognee_user_interaction], update_edge_collection=False)
+
+        relationships = []
+        relationship_name = "used_graph_element_to_answer"
+        for triplet in triplets:
+            target_id_1 = extract_uuid_from_node(triplet.node1)
+            target_id_2 = extract_uuid_from_node(triplet.node2)
+            if target_id_1 and target_id_2:
+                relationships.append(
+                    (
+                        source_id,
+                        target_id_1,
+                        relationship_name,
+                        {
+                            "relationship_name": relationship_name,
+                            "source_node_id": source_id,
+                            "target_node_id": target_id_1,
+                            "ontology_valid": False,
+                            "feedback_weight": 0,
+                        },
+                    )
+                )
+
+                relationships.append(
+                    (
+                        source_id,
+                        target_id_2,
+                        relationship_name,
+                        {
+                            "relationship_name": relationship_name,
+                            "source_node_id": source_id,
+                            "target_node_id": target_id_2,
+                            "ontology_valid": False,
+                            "feedback_weight": 0,
+                        },
+                    )
+                )
+
+            if len(relationships) > 0:
+                graph_engine = await get_graph_engine()
+                await graph_engine.add_edges(relationships)

cognee/modules/retrieval/graph_summary_completion_retriever.py

@@ -24,6 +24,7 @@ class GraphSummaryCompletionRetriever(GraphCompletionRetriever):
         top_k: Optional[int] = 5,
         node_type: Optional[Type] = None,
         node_name: Optional[List[str]] = None,
+        save_interaction: bool = False,
     ):
         """Initialize retriever with default prompt paths and search parameters."""
         super().__init__(
@@ -32,6 +33,7 @@ class GraphSummaryCompletionRetriever(GraphCompletionRetriever):
             top_k=top_k,
             node_type=node_type,
             node_name=node_name,
+            save_interaction=save_interaction,
         )
         self.summarize_prompt_path = summarize_prompt_path
 

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

@@ -3,6 +3,8 @@ import json
 import asyncio
 from uuid import UUID
 from typing import Callable, List, Optional, Type, Union
+
+from cognee.modules.retrieval.user_qa_feedback import UserQAFeedback
 from cognee.modules.search.exceptions import UnsupportedSearchTypeError
 from cognee.context_global_variables import set_database_global_context_variables
 from cognee.modules.retrieval.chunks_retriever import ChunksRetriever
@@ -38,6 +40,8 @@ async def search(
     top_k: int = 10,
     node_type: Optional[Type] = None,
     node_name: Optional[List[str]] = None,
+    save_interaction: Optional[bool] = False,
+    last_k: Optional[int] = None,
 ):
     """
 
@@ -57,7 +61,14 @@ async def search(
     # Use search function filtered by permissions if access control is enabled
     if os.getenv("ENABLE_BACKEND_ACCESS_CONTROL", "false").lower() == "true":
         return await authorized_search(
-            query_text, query_type, user, dataset_ids, system_prompt_path, top_k
+            query_text=query_text,
+            query_type=query_type,
+            user=user,
+            dataset_ids=dataset_ids,
+            system_prompt_path=system_prompt_path,
+            top_k=top_k,
+            save_interaction=save_interaction,
+            last_k=last_k,
         )
 
     query = await log_query(query_text, query_type.value, user.id)
@@ -70,6 +81,8 @@ async def search(
         top_k=top_k,
         node_type=node_type,
         node_name=node_name,
+        save_interaction=save_interaction,
+        last_k=last_k,
     )
 
     await log_result(
@@ -91,6 +104,8 @@ async def specific_search(
     top_k: int = 10,
     node_type: Optional[Type] = None,
     node_name: Optional[List[str]] = None,
+    save_interaction: Optional[bool] = False,
+    last_k: Optional[int] = None,
 ) -> list:
     search_tasks: dict[SearchType, Callable] = {
         SearchType.SUMMARIES: SummariesRetriever(top_k=top_k).get_completion,
@@ -104,28 +119,33 @@ async def specific_search(
             top_k=top_k,
             node_type=node_type,
             node_name=node_name,
+            save_interaction=save_interaction,
         ).get_completion,
         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,
         ).get_completion,
         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,
         ).get_completion,
         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,
         ).get_completion,
         SearchType.CODE: CodeRetriever(top_k=top_k).get_completion,
         SearchType.CYPHER: CypherSearchRetriever().get_completion,
         SearchType.NATURAL_LANGUAGE: NaturalLanguageRetriever().get_completion,
+        SearchType.FEEDBACK: UserQAFeedback(last_k=last_k).add_feedback,
     }
 
     # If the query type is FEELING_LUCKY, select the search type intelligently
@@ -153,6 +173,8 @@ async def authorized_search(
     dataset_ids: Optional[list[UUID]] = None,
     system_prompt_path: str = "answer_simple_question.txt",
     top_k: int = 10,
+    save_interaction: bool = False,
+    last_k: Optional[int] = None,
 ) -> list:
     """
     Verifies access for provided datasets or uses all datasets user has read access for and performs search per dataset.
@@ -166,7 +188,14 @@ async def authorized_search(
 
     # Searches all provided datasets and handles setting up of appropriate database context based on permissions
     search_results = await specific_search_by_context(
-        search_datasets, query_text, query_type, user, system_prompt_path, top_k
+        search_datasets,
+        query_text,
+        query_type,
+        user,
+        system_prompt_path,
+        top_k,
+        save_interaction,
+        last_k=last_k,
     )
 
     await log_result(query.id, json.dumps(search_results, cls=JSONEncoder), user.id)
@@ -181,17 +210,27 @@ async def specific_search_by_context(
     user: User,
     system_prompt_path: str,
     top_k: int,
+    save_interaction: bool = False,
+    last_k: Optional[int] = None,
 ):
     """
     Searches all provided datasets and handles setting up of appropriate database context based on permissions.
     Not to be used outside of active access control mode.
     """
 
-    async def _search_by_context(dataset, user, query_type, query_text, system_prompt_path, top_k):
+    async def _search_by_context(
+        dataset, user, query_type, query_text, system_prompt_path, top_k, last_k
+    ):
         # Set database configuration in async context for each dataset user has access for
         await set_database_global_context_variables(dataset.id, dataset.owner_id)
         search_results = await specific_search(
-            query_type, query_text, user, system_prompt_path=system_prompt_path, top_k=top_k
+            query_type,
+            query_text,
+            user,
+            system_prompt_path=system_prompt_path,
+            top_k=top_k,
+            save_interaction=save_interaction,
+            last_k=last_k,
         )
         return {
             "search_result": search_results,
@@ -203,7 +242,9 @@ async def specific_search_by_context(
     tasks = []
     for dataset in search_datasets:
         tasks.append(
-            _search_by_context(dataset, user, query_type, query_text, system_prompt_path, top_k)
+            _search_by_context(
+                dataset, user, query_type, query_text, system_prompt_path, top_k, last_k
+            )
         )
 
     return await asyncio.gather(*tasks)

cognee/modules/search/types/SearchType.py

@@ -14,3 +14,4 @@ class SearchType(Enum):
     GRAPH_COMPLETION_COT = "GRAPH_COMPLETION_COT"
     GRAPH_COMPLETION_CONTEXT_EXTENSION = "GRAPH_COMPLETION_CONTEXT_EXTENSION"
     FEELING_LUCKY = "FEELING_LUCKY"
+    FEEDBACK = "FEEDBACK"

cognee/modules/settings/get_settings.py

@@ -88,8 +88,8 @@ def get_settings() -> SettingsDict:
                 "models": {
                     "openai": [
                         {
-                            "value": "gpt-4o-mini",
-                            "label": "gpt-4o-mini",
+                            "value": "gpt-5-mini",
+                            "label": "gpt-5-mini",
                         },
                         {
                             "value": "gpt-4o",

cognee/shared/CodeGraphEntities.py

@@ -36,6 +36,7 @@ class ClassDefinition(DataPoint):
 class CodeFile(DataPoint):
     name: str
     file_path: str
+    language: Optional[str] = None  # e.g., 'python', 'javascript', 'java', etc.
     source_code: Optional[str] = None
     part_of: Optional[Repository] = None
     depends_on: Optional[List["ImportStatement"]] = []