diff --git a/cognee/api/v1/search/search.py b/cognee/api/v1/search/search.py index d4e5fbbe66..557a93b2c4 100644 --- a/cognee/api/v1/search/search.py +++ b/cognee/api/v1/search/search.py @@ -169,7 +169,8 @@ async def search( - GRAPH_DATABASE_PROVIDER: Must match what was used during cognify """ - # We use lists from now on for datasets + + ###################### if isinstance(datasets, UUID) or isinstance(datasets, str): datasets = [datasets] diff --git a/cognee/modules/retrieval/utils/models.py b/cognee/modules/retrieval/utils/models.py index 58cea29a45..eaa616966a 100644 --- a/cognee/modules/retrieval/utils/models.py +++ b/cognee/modules/retrieval/utils/models.py @@ -38,3 +38,13 @@ class UserFeedbackEvaluation(BaseModel): ..., description="Sentiment score from -5 (negative) to +5 (positive)" ) evaluation: UserFeedbackSentiment + + + +class CogneeSearchSentiment(DataPoint): + """Cognee - Search Sentiment""" + current_question: str + sentiment: str # Positive / Neutral / Negative + score: int # -5 to 5 + user_id: str + belongs_to_set: Optional[NodeSet] = None diff --git a/cognee/tasks/memify/sentiment_analysis_pipeline.py b/cognee/tasks/memify/sentiment_analysis_pipeline.py new file mode 100644 index 0000000000..2a2c29bdd2 --- /dev/null +++ b/cognee/tasks/memify/sentiment_analysis_pipeline.py @@ -0,0 +1,17 @@ +from cognee import memify +from cognee.modules.pipelines.tasks.task import Task +from cognee.tasks.sentiment_analysis.sentiment_analysis import sentiment_analysis_task +from cognee.tasks.sentiment_analysis.enrichment import enrichment_task +import asyncio +async def main(): + + extraction_task = Task(sentiment_analysis_task) + enrichment_task = Task(enrichment_task) + + await memify( + extraction_tasks=[extraction_task], + enrichment_tasks=[enrichment_task], + ) + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/cognee/tasks/sentiment_analysis/__init__.py b/cognee/tasks/sentiment_analysis/__init__.py new file mode 100644 index 0000000000..88a92e39c9 --- /dev/null +++ b/cognee/tasks/sentiment_analysis/__init__.py @@ -0,0 +1 @@ +from .sentiment_analysis import run_sentiment_analysis \ No newline at end of file diff --git a/cognee/tasks/sentiment_analysis/enrichment.py b/cognee/tasks/sentiment_analysis/enrichment.py new file mode 100644 index 0000000000..6065476677 --- /dev/null +++ b/cognee/tasks/sentiment_analysis/enrichment.py @@ -0,0 +1,15 @@ +from cognee.tasks.storage import add_data_points +from cognee.shared.logging_utils import get_logger +from cognee.modules.retrieval.utils.models import CogneeSearchSentiment +from typing import Optional, List +logger = get_logger() +async def enrichment_task(sentiment_data_points: List[CogneeSearchSentiment]): + """ + This task takes the sentiment data points and adds them to the graph. + """ + if sentiment_data_points: + # Save the sentiment data points to the database + await add_data_points(data_points=sentiment_data_points, update_edge_collection=False) + logger.info(f"Enriched the graph with {len(sentiment_data_points)} sentiment data points.") + else: + logger.info("No sentiment data points to store.") diff --git a/cognee/tasks/sentiment_analysis/sentiment_analysis.py b/cognee/tasks/sentiment_analysis/sentiment_analysis.py new file mode 100644 index 0000000000..e01dc115f6 --- /dev/null +++ b/cognee/tasks/sentiment_analysis/sentiment_analysis.py @@ -0,0 +1,71 @@ +from cognee.infrastructure.llm import LLMGateway +from cognee.modules.engine.models import NodeSet +from uuid import uuid5, NAMESPACE_OID +from typing import Optional, List +from cognee.modules.retrieval.utils.models import CogneeSearchSentiment +from cognee.modules.users.models import User +from cognee.shared.logging_utils import get_logger +from cognee.modules.users.methods import get_default_user +from cognee.infrastructure.databases.graph import get_graph_engine + +logger = get_logger() + +async def run_sentiment_analysis(user: Optional[User] = None) -> List[CogneeSearchSentiment]: + """ + This function fetches all the nodes from the graph, filters for nodes with a 'question' attribute, + and performs sentiment analysis on each question. + Returns a list of sentiment data points for non-neutral sentiments. + """ + + # Fetch all graph data (nodes and edges) + graph_engine = await get_graph_engine() + nodes_data, edges_data = await graph_engine.get_graph_data() + + # Filter the nodes to find those with a 'question' attribute + question_nodes = [node for node in nodes_data if 'question' in node[1]] + + if user is None: + user = await get_default_user() # Get default user if no user is provided + + user_id = str(user.id) + + # Initialize an empty list to store data points + sentiment_data_points = [] + + # For each filtered node (which has a 'question'), perform sentiment analysis + for node in question_nodes: + current_question = node[1].get('question') # Get the question from the node's properties + if current_question: + # Call LLM to classify sentiment for the current question + sentiment_result = await LLMGateway.acreate_structured_output( + text_input=current_question, + system_prompt="""Classify the user's reaction as Positive, Neutral, or Negative with a score (-5 to 5).Return the result as valid JSON like:{"sentiment": "Positive","score": 3}""", + response_model=CogneeSearchSentiment + ) + + # Print sentiment result for debugging + print(sentiment_result) + + sentiment_data_point = CogneeSearchSentiment( + id=uuid5(NAMESPACE_OID, name=user_id + current_question), + current_question=current_question, + sentiment=sentiment_result.sentiment, + score=sentiment_result.score, + user_id=user_id, + belongs_to_set=NodeSet(id=uuid5(NAMESPACE_OID, "CogneeSearchSentiment"), name="CogneeSearchSentiment") + ) + + # Only add the data point to the list if sentiment is non-neutral + if sentiment_result.sentiment != 'Neutral': + sentiment_data_points.append(sentiment_data_point) + + # Log the created data point for debugging + data_point = { + "current_question": current_question, + "sentiment": sentiment_result.sentiment, + "score": sentiment_result.score + } + logger.info(f"Sentiment Data Point Created: {data_point}") + + # Return the list of sentiment data points + return sentiment_data_points diff --git a/examples/python/sentiment_analysis_demo.py b/examples/python/sentiment_analysis_demo.py new file mode 100644 index 0000000000..3fc3a847a7 --- /dev/null +++ b/examples/python/sentiment_analysis_demo.py @@ -0,0 +1,60 @@ +import asyncio +import cognee +from cognee import memify +from cognee.modules.pipelines.tasks.task import Task +from cognee.tasks.sentiment_analysis.sentiment_analysis import run_sentiment_analysis +from cognee.tasks.sentiment_analysis.enrichment import enrichment_task +from cognee.shared.logging_utils import get_logger + +logger = get_logger() + + +async def main(): + # Step 1: Reset Cognee data + logger.info('Resetting Cognee data...') + await cognee.prune.prune_data() + await cognee.prune.prune_system(metadata=True) + logger.info("Data reset complete.\n") + + # Step 2: Add sample content + text = "Cognee turns documents into AI memory." + await cognee.add(text) + + # Step 3: Build the knowledge graph + logger.info("Cognifying the content...") + await cognee.cognify() + logger.info("Cognify complete.\n") + + # Step 4: Define queries to test + queries = [ + "What does Cognee do?", + "How does Cognee store data?", + "Are you even listening to what I am asking?", + "This is good, this was what I was asking for" + ] + + all_results = {} + + for q in queries: + results = await cognee.search( + query_text=q, + save_interaction=True, # Save interactions for analysis + ) + all_results[q] = results + + # Step 5: Create your extraction and enrichment tasks + extraction_task = Task(run_sentiment_analysis) + enrichment = Task(enrichment_task) + + # Step 6: Run memify pipeline — this executes both tasks + logger.info("Running sentiment analysis pipeline via memify...") + sentiment_data_points = await memify( + extraction_tasks=[extraction_task], + enrichment_tasks=[enrichment], + ) + logger.info("Memify pipeline complete.\n") + logger.info(sentiment_data_points) + + +if __name__ == '__main__': + asyncio.run(main())