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

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, []