cognee 0.3.6__py3-none-any.whl → 0.3.7.dev1__py3-none-any.whl

cognee/tasks/feedback/create_enrichments.py

@@ -0,0 +1,84 @@
+from __future__ import annotations
+
+from typing import List
+from uuid import NAMESPACE_OID, uuid5
+
+from cognee.infrastructure.llm import LLMGateway
+from cognee.infrastructure.llm.prompts.read_query_prompt import read_query_prompt
+from cognee.shared.logging_utils import get_logger
+from cognee.modules.engine.models import NodeSet
+
+from .models import FeedbackEnrichment
+
+
+logger = get_logger("create_enrichments")
+
+
+def _validate_enrichments(enrichments: List[FeedbackEnrichment]) -> bool:
+    """Validate that all enrichments contain required fields for completion."""
+    return all(
+        enrichment.question is not None
+        and enrichment.original_answer is not None
+        and enrichment.improved_answer is not None
+        and enrichment.new_context is not None
+        and enrichment.feedback_id is not None
+        and enrichment.interaction_id is not None
+        for enrichment in enrichments
+    )
+
+
+async def _generate_enrichment_report(
+    question: str, improved_answer: str, new_context: str, report_prompt_location: str
+) -> str:
+    """Generate educational report using feedback report prompt."""
+    try:
+        prompt_template = read_query_prompt(report_prompt_location)
+        rendered_prompt = prompt_template.format(
+            question=question,
+            improved_answer=improved_answer,
+            new_context=new_context,
+        )
+        return await LLMGateway.acreate_structured_output(
+            text_input=rendered_prompt,
+            system_prompt="You are a helpful assistant that creates educational content.",
+            response_model=str,
+        )
+    except Exception as exc:
+        logger.warning("Failed to generate enrichment report", error=str(exc), question=question)
+        return f"Educational content for: {question} - {improved_answer}"
+
+
+async def create_enrichments(
+    enrichments: List[FeedbackEnrichment],
+    report_prompt_location: str = "feedback_report_prompt.txt",
+) -> List[FeedbackEnrichment]:
+    """Fill text and belongs_to_set fields of existing FeedbackEnrichment DataPoints."""
+    if not enrichments:
+        logger.info("No enrichments provided; returning empty list")
+        return []
+
+    if not _validate_enrichments(enrichments):
+        logger.error("Input validation failed; missing required fields")
+        return []
+
+    logger.info("Completing enrichments", count=len(enrichments))
+
+    nodeset = NodeSet(id=uuid5(NAMESPACE_OID, name="FeedbackEnrichment"), name="FeedbackEnrichment")
+
+    completed_enrichments: List[FeedbackEnrichment] = []
+
+    for enrichment in enrichments:
+        report_text = await _generate_enrichment_report(
+            enrichment.question,
+            enrichment.improved_answer,
+            enrichment.new_context,
+            report_prompt_location,
+        )
+
+        enrichment.text = report_text
+        enrichment.belongs_to_set = [nodeset]
+
+        completed_enrichments.append(enrichment)
+
+    logger.info("Completed enrichments", successful=len(completed_enrichments))
+    return completed_enrichments

cognee/tasks/feedback/extract_feedback_interactions.py

@@ -0,0 +1,230 @@
+from __future__ import annotations
+
+from typing import Any, Dict, List, Optional, Tuple
+from uuid import UUID, uuid5, NAMESPACE_OID
+
+from cognee.infrastructure.llm import LLMGateway
+from cognee.infrastructure.llm.prompts.read_query_prompt import read_query_prompt
+from cognee.shared.logging_utils import get_logger
+from cognee.infrastructure.databases.graph import get_graph_engine
+
+from .models import FeedbackEnrichment
+
+
+logger = get_logger("extract_feedback_interactions")
+
+
+def _filter_negative_feedback(feedback_nodes):
+    """Filter for negative sentiment feedback using precise sentiment classification."""
+    return [
+        (node_id, props)
+        for node_id, props in feedback_nodes
+        if (props.get("sentiment", "").casefold() == "negative" or props.get("score", 0) < 0)
+    ]
+
+
+def _get_normalized_id(node_id, props) -> str:
+    """Return Cognee node id preference: props.id → props.node_id → raw node_id."""
+    return str(props.get("id") or props.get("node_id") or node_id)
+
+
+async def _fetch_feedback_and_interaction_graph_data() -> Tuple[List, List]:
+    """Fetch feedback and interaction nodes with edges from graph engine."""
+    try:
+        graph_engine = await get_graph_engine()
+        attribute_filters = [{"type": ["CogneeUserFeedback", "CogneeUserInteraction"]}]
+        return await graph_engine.get_filtered_graph_data(attribute_filters)
+    except Exception as exc:  # noqa: BLE001
+        logger.error("Failed to fetch filtered graph data", error=str(exc))
+        return [], []
+
+
+def _separate_feedback_and_interaction_nodes(graph_nodes: List) -> Tuple[List, List]:
+    """Split nodes into feedback and interaction groups by type field."""
+    feedback_nodes = [
+        (_get_normalized_id(node_id, props), props)
+        for node_id, props in graph_nodes
+        if props.get("type") == "CogneeUserFeedback"
+    ]
+    interaction_nodes = [
+        (_get_normalized_id(node_id, props), props)
+        for node_id, props in graph_nodes
+        if props.get("type") == "CogneeUserInteraction"
+    ]
+    return feedback_nodes, interaction_nodes
+
+
+def _match_feedback_nodes_to_interactions_by_edges(
+    feedback_nodes: List, interaction_nodes: List, graph_edges: List
+) -> List[Tuple[Tuple, Tuple]]:
+    """Match feedback to interactions using gives_feedback_to edges."""
+    interaction_by_id = {node_id: (node_id, props) for node_id, props in interaction_nodes}
+    feedback_by_id = {node_id: (node_id, props) for node_id, props in feedback_nodes}
+    feedback_edges = [
+        (source_id, target_id)
+        for source_id, target_id, rel, _ in graph_edges
+        if rel == "gives_feedback_to"
+    ]
+
+    feedback_interaction_pairs: List[Tuple[Tuple, Tuple]] = []
+    for source_id, target_id in feedback_edges:
+        source_id_str, target_id_str = str(source_id), str(target_id)
+
+        feedback_node = feedback_by_id.get(source_id_str)
+        interaction_node = interaction_by_id.get(target_id_str)
+
+        if feedback_node and interaction_node:
+            feedback_interaction_pairs.append((feedback_node, interaction_node))
+
+    return feedback_interaction_pairs
+
+
+def _sort_pairs_by_recency_and_limit(
+    feedback_interaction_pairs: List[Tuple[Tuple, Tuple]], last_n_limit: Optional[int]
+) -> List[Tuple[Tuple, Tuple]]:
+    """Sort by interaction created_at desc with updated_at fallback, then limit."""
+
+    def _recency_key(pair):
+        _, (_, interaction_props) = pair
+        created_at = interaction_props.get("created_at") or ""
+        updated_at = interaction_props.get("updated_at") or ""
+        return (created_at, updated_at)
+
+    sorted_pairs = sorted(feedback_interaction_pairs, key=_recency_key, reverse=True)
+    return sorted_pairs[: last_n_limit or len(sorted_pairs)]
+
+
+async def _generate_human_readable_context_summary(
+    question_text: str, raw_context_text: str
+) -> str:
+    """Generate a concise human-readable summary for given context."""
+    try:
+        prompt = read_query_prompt("feedback_user_context_prompt.txt")
+        rendered = prompt.format(question=question_text, context=raw_context_text)
+        return await LLMGateway.acreate_structured_output(
+            text_input=rendered, system_prompt="", response_model=str
+        )
+    except Exception as exc:  # noqa: BLE001
+        logger.warning("Failed to summarize context", error=str(exc))
+        return raw_context_text or ""
+
+
+def _has_required_feedback_fields(enrichment: FeedbackEnrichment) -> bool:
+    """Validate required fields exist in the FeedbackEnrichment DataPoint."""
+    return (
+        enrichment.question is not None
+        and enrichment.original_answer is not None
+        and enrichment.context is not None
+        and enrichment.feedback_text is not None
+        and enrichment.feedback_id is not None
+        and enrichment.interaction_id is not None
+    )
+
+
+async def _build_feedback_interaction_record(
+    feedback_node_id: str, feedback_props: Dict, interaction_node_id: str, interaction_props: Dict
+) -> Optional[FeedbackEnrichment]:
+    """Build a single FeedbackEnrichment DataPoint with context summary."""
+    try:
+        question_text = interaction_props.get("question")
+        original_answer_text = interaction_props.get("answer")
+        raw_context_text = interaction_props.get("context", "")
+        feedback_text = feedback_props.get("feedback") or feedback_props.get("text") or ""
+
+        context_summary_text = await _generate_human_readable_context_summary(
+            question_text or "", raw_context_text
+        )
+
+        enrichment = FeedbackEnrichment(
+            id=str(uuid5(NAMESPACE_OID, f"{question_text}_{interaction_node_id}")),
+            text="",
+            question=question_text,
+            original_answer=original_answer_text,
+            improved_answer="",
+            feedback_id=UUID(str(feedback_node_id)),
+            interaction_id=UUID(str(interaction_node_id)),
+            belongs_to_set=None,
+            context=context_summary_text,
+            feedback_text=feedback_text,
+            new_context="",
+            explanation="",
+        )
+
+        if _has_required_feedback_fields(enrichment):
+            return enrichment
+        else:
+            logger.warning("Skipping invalid feedback item", interaction=str(interaction_node_id))
+            return None
+    except Exception as exc:  # noqa: BLE001
+        logger.error("Failed to process feedback pair", error=str(exc))
+        return None
+
+
+async def _build_feedback_interaction_records(
+    matched_feedback_interaction_pairs: List[Tuple[Tuple, Tuple]],
+) -> List[FeedbackEnrichment]:
+    """Build all FeedbackEnrichment DataPoints from matched pairs."""
+    feedback_interaction_records: List[FeedbackEnrichment] = []
+    for (feedback_node_id, feedback_props), (
+        interaction_node_id,
+        interaction_props,
+    ) in matched_feedback_interaction_pairs:
+        record = await _build_feedback_interaction_record(
+            feedback_node_id, feedback_props, interaction_node_id, interaction_props
+        )
+        if record:
+            feedback_interaction_records.append(record)
+    return feedback_interaction_records
+
+
+async def extract_feedback_interactions(
+    data: Any, last_n: Optional[int] = None
+) -> List[FeedbackEnrichment]:
+    """Extract negative feedback-interaction pairs and create FeedbackEnrichment DataPoints."""
+    if not data or data == [{}]:
+        logger.info(
+            "No data passed to the extraction task (extraction task fetches data from graph directly)",
+            data=data,
+        )
+
+    graph_nodes, graph_edges = await _fetch_feedback_and_interaction_graph_data()
+    if not graph_nodes:
+        logger.warning("No graph nodes retrieved from database")
+        return []
+
+    feedback_nodes, interaction_nodes = _separate_feedback_and_interaction_nodes(graph_nodes)
+    logger.info(
+        "Retrieved nodes from graph",
+        total_nodes=len(graph_nodes),
+        feedback_nodes=len(feedback_nodes),
+        interaction_nodes=len(interaction_nodes),
+    )
+
+    negative_feedback_nodes = _filter_negative_feedback(feedback_nodes)
+    logger.info(
+        "Filtered feedback nodes",
+        total_feedback=len(feedback_nodes),
+        negative_feedback=len(negative_feedback_nodes),
+    )
+
+    if not negative_feedback_nodes:
+        logger.info("No negative feedback found; returning empty list")
+        return []
+
+    matched_feedback_interaction_pairs = _match_feedback_nodes_to_interactions_by_edges(
+        negative_feedback_nodes, interaction_nodes, graph_edges
+    )
+    if not matched_feedback_interaction_pairs:
+        logger.info("No feedback-to-interaction matches found; returning empty list")
+        return []
+
+    matched_feedback_interaction_pairs = _sort_pairs_by_recency_and_limit(
+        matched_feedback_interaction_pairs, last_n
+    )
+
+    feedback_interaction_records = await _build_feedback_interaction_records(
+        matched_feedback_interaction_pairs
+    )
+
+    logger.info("Extracted feedback pairs", count=len(feedback_interaction_records))
+    return feedback_interaction_records

cognee/tasks/feedback/generate_improved_answers.py

@@ -0,0 +1,130 @@
+from __future__ import annotations
+
+from typing import List, Optional
+from pydantic import BaseModel
+
+from cognee.infrastructure.llm import LLMGateway
+from cognee.infrastructure.llm.prompts.read_query_prompt import read_query_prompt
+from cognee.modules.graph.utils import resolve_edges_to_text
+from cognee.shared.logging_utils import get_logger
+
+from cognee.modules.retrieval.graph_completion_cot_retriever import GraphCompletionCotRetriever
+from .models import FeedbackEnrichment
+
+
+class ImprovedAnswerResponse(BaseModel):
+    """Response model for improved answer generation containing answer and explanation."""
+
+    answer: str
+    explanation: str
+
+
+logger = get_logger("generate_improved_answers")
+
+
+def _validate_input_data(enrichments: List[FeedbackEnrichment]) -> bool:
+    """Validate that input contains required fields for all enrichments."""
+    return all(
+        enrichment.question is not None
+        and enrichment.original_answer is not None
+        and enrichment.context is not None
+        and enrichment.feedback_text is not None
+        and enrichment.feedback_id is not None
+        and enrichment.interaction_id is not None
+        for enrichment in enrichments
+    )
+
+
+def _render_reaction_prompt(
+    question: str, context: str, wrong_answer: str, negative_feedback: str
+) -> str:
+    """Render the feedback reaction prompt with provided variables."""
+    prompt_template = read_query_prompt("feedback_reaction_prompt.txt")
+    return prompt_template.format(
+        question=question,
+        context=context,
+        wrong_answer=wrong_answer,
+        negative_feedback=negative_feedback,
+    )
+
+
+async def _generate_improved_answer_for_single_interaction(
+    enrichment: FeedbackEnrichment, retriever, reaction_prompt_location: str
+) -> Optional[FeedbackEnrichment]:
+    """Generate improved answer for a single enrichment using structured retriever completion."""
+    try:
+        query_text = _render_reaction_prompt(
+            enrichment.question,
+            enrichment.context,
+            enrichment.original_answer,
+            enrichment.feedback_text,
+        )
+
+        retrieved_context = await retriever.get_context(query_text)
+        completion = await retriever.get_structured_completion(
+            query=query_text,
+            context=retrieved_context,
+            response_model=ImprovedAnswerResponse,
+            max_iter=4,
+        )
+        new_context_text = await retriever.resolve_edges_to_text(retrieved_context)
+
+        if completion:
+            enrichment.improved_answer = completion.answer
+            enrichment.new_context = new_context_text
+            enrichment.explanation = completion.explanation
+            return enrichment
+        else:
+            logger.warning(
+                "Failed to get structured completion from retriever", question=enrichment.question
+            )
+            return None
+
+    except Exception as exc:  # noqa: BLE001
+        logger.error(
+            "Failed to generate improved answer",
+            error=str(exc),
+            question=enrichment.question,
+        )
+        return None
+
+
+async def generate_improved_answers(
+    enrichments: List[FeedbackEnrichment],
+    top_k: int = 20,
+    reaction_prompt_location: str = "feedback_reaction_prompt.txt",
+) -> List[FeedbackEnrichment]:
+    """Generate improved answers using CoT retriever and LLM."""
+    if not enrichments:
+        logger.info("No enrichments provided; returning empty list")
+        return []
+
+    if not _validate_input_data(enrichments):
+        logger.error("Input data validation failed; missing required fields")
+        return []
+
+    retriever = GraphCompletionCotRetriever(
+        top_k=top_k,
+        save_interaction=False,
+        user_prompt_path="graph_context_for_question.txt",
+        system_prompt_path="answer_simple_question.txt",
+    )
+
+    improved_answers: List[FeedbackEnrichment] = []
+
+    for enrichment in enrichments:
+        result = await _generate_improved_answer_for_single_interaction(
+            enrichment, retriever, reaction_prompt_location
+        )
+
+        if result:
+            improved_answers.append(result)
+        else:
+            logger.warning(
+                "Failed to generate improved answer",
+                question=enrichment.question,
+                interaction_id=enrichment.interaction_id,
+            )
+
+    logger.info("Generated improved answers", count=len(improved_answers))
+    return improved_answers

cognee/tasks/feedback/link_enrichments_to_feedback.py

@@ -0,0 +1,67 @@
+from __future__ import annotations
+
+from typing import List, Tuple
+from uuid import UUID
+
+from cognee.infrastructure.databases.graph import get_graph_engine
+from cognee.tasks.storage import index_graph_edges
+from cognee.shared.logging_utils import get_logger
+
+from .models import FeedbackEnrichment
+
+
+logger = get_logger("link_enrichments_to_feedback")
+
+
+def _create_edge_tuple(
+    source_id: UUID, target_id: UUID, relationship_name: str
+) -> Tuple[UUID, UUID, str, dict]:
+    """Create an edge tuple with proper properties structure."""
+    return (
+        source_id,
+        target_id,
+        relationship_name,
+        {
+            "relationship_name": relationship_name,
+            "source_node_id": source_id,
+            "target_node_id": target_id,
+            "ontology_valid": False,
+        },
+    )
+
+
+async def link_enrichments_to_feedback(
+    enrichments: List[FeedbackEnrichment],
+) -> List[FeedbackEnrichment]:
+    """Manually create edges from enrichments to original feedback/interaction nodes."""
+    if not enrichments:
+        logger.info("No enrichments provided; returning empty list")
+        return []
+
+    relationships = []
+
+    for enrichment in enrichments:
+        enrichment_id = enrichment.id
+        feedback_id = enrichment.feedback_id
+        interaction_id = enrichment.interaction_id
+
+        if enrichment_id and feedback_id:
+            enriches_feedback_edge = _create_edge_tuple(
+                enrichment_id, feedback_id, "enriches_feedback"
+            )
+            relationships.append(enriches_feedback_edge)
+
+        if enrichment_id and interaction_id:
+            improves_interaction_edge = _create_edge_tuple(
+                enrichment_id, interaction_id, "improves_interaction"
+            )
+            relationships.append(improves_interaction_edge)
+
+    if relationships:
+        graph_engine = await get_graph_engine()
+        await graph_engine.add_edges(relationships)
+        await index_graph_edges(relationships)
+        logger.info("Linking enrichments to feedback", edge_count=len(relationships))
+
+    logger.info("Linked enrichments", enrichment_count=len(enrichments))
+    return enrichments

cognee/tasks/feedback/models.py

@@ -0,0 +1,26 @@
+from typing import List, Optional, Union
+from uuid import UUID
+
+from cognee.infrastructure.engine import DataPoint
+from cognee.modules.engine.models import Entity, NodeSet
+from cognee.tasks.temporal_graph.models import Event
+
+
+class FeedbackEnrichment(DataPoint):
+    """Minimal DataPoint for feedback enrichment that works with extract_graph_from_data."""
+
+    text: str
+    contains: Optional[List[Union[Entity, Event]]] = None
+    metadata: dict = {"index_fields": ["text"]}
+
+    question: str
+    original_answer: str
+    improved_answer: str
+    feedback_id: UUID
+    interaction_id: UUID
+    belongs_to_set: Optional[List[NodeSet]] = None
+
+    context: str = ""
+    feedback_text: str = ""
+    new_context: str = ""
+    explanation: str = ""

cognee/tasks/graph/extract_graph_from_data.py

@@ -4,6 +4,7 @@ from pydantic import BaseModel
 
 from cognee.infrastructure.databases.graph import get_graph_engine
 from cognee.modules.ontology.ontology_env_config import get_ontology_env_config
+from cognee.tasks.storage import index_graph_edges
 from cognee.tasks.storage.add_data_points import add_data_points
 from cognee.modules.ontology.ontology_config import Config
 from cognee.modules.ontology.get_default_ontology_resolver import (
@@ -88,6 +89,7 @@ async def integrate_chunk_graphs(
 
     if len(graph_edges) > 0:
         await graph_engine.add_edges(graph_edges)
+        await index_graph_edges(graph_edges)
 
     return data_chunks
 

cognee/tasks/ingestion/data_item_to_text_file.py

@@ -1,6 +1,6 @@
 import os
 from urllib.parse import urlparse
-from typing import List, Tuple
+from typing import Any, List, Tuple
 from pathlib import Path
 import tempfile
 
@@ -34,7 +34,8 @@ async def pull_from_s3(file_path, destination_file) -> None:
 
 
 async def data_item_to_text_file(
-    data_item_path: str, preferred_loaders: List[str]
+    data_item_path: str,
+    preferred_loaders: dict[str, dict[str, Any]] = None,
 ) -> Tuple[str, LoaderInterface]:
     if isinstance(data_item_path, str):
         parsed_url = urlparse(data_item_path)
@@ -74,6 +75,5 @@ async def data_item_to_text_file(
                 )
             else:
                 raise IngestionError(message="Local files are not accepted.")
-
     # data is not a supported type
     raise IngestionError(message=f"Data type not supported: {type(data_item_path)}")

cognee/tasks/ingestion/ingest_data.py

@@ -6,6 +6,7 @@ from typing import Union, BinaryIO, Any, List, Optional
 import cognee.modules.ingestion as ingestion
 from cognee.infrastructure.databases.relational import get_relational_engine
 from cognee.modules.data.models import Data
+from cognee.modules.ingestion.exceptions import IngestionError
 from cognee.modules.users.models import User
 from cognee.modules.users.methods import get_default_user
 from cognee.modules.users.permissions.methods import get_specific_user_permission_datasets
@@ -27,7 +28,7 @@ async def ingest_data(
     user: User,
     node_set: Optional[List[str]] = None,
     dataset_id: UUID = None,
-    preferred_loaders: List[str] = None,
+    preferred_loaders: dict[str, dict[str, Any]] = None,
 ):
     if not user:
         user = await get_default_user()
@@ -44,7 +45,7 @@ async def ingest_data(
         user: User,
         node_set: Optional[List[str]] = None,
         dataset_id: UUID = None,
-        preferred_loaders: List[str] = None,
+        preferred_loaders: dict[str, dict[str, Any]] = None,
     ):
         new_datapoints = []
         existing_data_points = []
@@ -77,22 +78,27 @@ async def ingest_data(
         dataset_data_map = {str(data.id): True for data in dataset_data}
 
         for data_item in data:
-            # Get file path of data item or create a file it doesn't exist
+            # Get file path of data item or create a file if it doesn't exist
             original_file_path = await save_data_item_to_storage(data_item)
-
             # Transform file path to be OS usable
             actual_file_path = get_data_file_path(original_file_path)
 
             # Store all input data as text files in Cognee data storage
             cognee_storage_file_path, loader_engine = await data_item_to_text_file(
-                actual_file_path, preferred_loaders
+                actual_file_path,
+                preferred_loaders,
             )
 
+            if loader_engine is None:
+                raise IngestionError("Loader cannot be None")
+
             # Find metadata from original file
+            # Standard flow: extract metadata from both original and stored files
             async with open_data_file(original_file_path) as file:
                 classified_data = ingestion.classify(file)
 
                 # data_id is the hash of original file contents + owner id to avoid duplicate data
+
                 data_id = ingestion.identify(classified_data, user)
                 original_file_metadata = classified_data.get_metadata()
 

cognee/tasks/ingestion/save_data_item_to_storage.py

@@ -8,6 +8,9 @@ from cognee.modules.ingestion import save_data_to_file
 from cognee.shared.logging_utils import get_logger
 from pydantic_settings import BaseSettings, SettingsConfigDict
 
+from cognee.tasks.web_scraper.utils import fetch_page_content
+
+
 logger = get_logger()
 
 
@@ -27,6 +30,12 @@ async def save_data_item_to_storage(data_item: Union[BinaryIO, str, Any]) -> str
 
         return await get_data_from_llama_index(data_item)
 
+    if "docling" in str(type(data_item)):
+        from docling_core.types import DoclingDocument
+
+        if isinstance(data_item, DoclingDocument):
+            data_item = data_item.export_to_text()
+
     # data is a file object coming from upload.
     if hasattr(data_item, "file"):
         return await save_data_to_file(data_item.file, filename=data_item.filename)
@@ -48,7 +57,9 @@ async def save_data_item_to_storage(data_item: Union[BinaryIO, str, Any]) -> str
         # data is s3 file path
         if parsed_url.scheme == "s3":
             return data_item
-
+        elif parsed_url.scheme == "http" or parsed_url.scheme == "https":
+            urls_to_page_contents = await fetch_page_content(data_item)
+            return await save_data_to_file(urls_to_page_contents[data_item], file_extension="html")
         # data is local file path
         elif parsed_url.scheme == "file":
             if settings.accept_local_file_path: