feat: add sentiment classification task #1346

cognee/api/v1/search/search.py

@@ -8,8 +8,11 @@
 from cognee.modules.search.methods import search as search_function
 from cognee.modules.data.methods import get_authorized_existing_datasets
 from cognee.modules.data.exceptions import DatasetNotFoundError
-
-
+from cognee.modules.retrieval.utils.models import CogneeUserInteraction
+from cognee.infrastructure.databases.graph import get_graph_engine
+from cognee.tasks.sentiment_analysis.sentiment_analysis import run_sentiment_analysis
+from cognee.shared.logging_utils import get_logger
+logger = get_logger()
 async def search(
     query_text: str,
     query_type: SearchType = SearchType.GRAPH_COMPLETION,
@@ -167,13 +170,41 @@ 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]
 
     if user is None:
         user = await get_default_user()
 
+    # Run sentiment analysis only when saving interactions and when we have a prior interaction
+    if save_interaction:
+        try:
+            graph_engine = await get_graph_engine()
+            last_interaction_ids = await graph_engine.get_last_user_interaction_ids(limit=1)
+            last = None
+            for interaction_id in last_interaction_ids:
+                interaction = await graph_engine.get_node(node_id=interaction_id)
+                if not interaction:
+                    continue
+                props = interaction.get("properties", interaction)
+
+                if "user_id" in props and str(props["user_id"]) != str(user.id):
+                    continue
+                if "question" in props and "answer" in props:
+                    last = props
+                    break
+            if last:
+                await run_sentiment_analysis(
+                    prev_question=last["question"],
+                    prev_answer=last["answer"],
+                    current_question=query_text,
+                    user=user,
+                )
+        except Exception as e:
+            logger.error(f"Sentiment Analysis Failed: {e}")
+
     # Transform string based datasets to UUID - String based datasets can only be found for current user
     if datasets is not None and [all(isinstance(dataset, str) for dataset in datasets)]:
         datasets = await get_authorized_existing_datasets(datasets, "read", user)

cognee/modules/retrieval/utils/models.py

@@ -38,3 +38,15 @@ class UserFeedbackEvaluation(BaseModel):
         ..., description="Sentiment score from -5 (negative) to +5 (positive)"
     )
     evaluation: UserFeedbackSentiment
+
+
+
+class CogneeSearchSentiment(DataPoint):
+    """Cognee - Search Sentiment"""
+    prev_question: str
+    prev_answer: str
+    current_question: str
+    sentiment: str   # Positive / Neutral / Negative
+    score: int       # -5 to 5
+    user_id: str
+    belongs_to_set: Optional[NodeSet] = None

cognee/tasks/sentiment_analysis/__init__.py

@@ -0,0 +1 @@
+from .sentiment_analysis import run_sentiment_analysis

cognee/tasks/sentiment_analysis/sentiment_analysis.py

@@ -0,0 +1,41 @@
+# tasks/run_sentiment_analysis.py
+
+from cognee.infrastructure.llm import LLMGateway
+from cognee.tasks.storage import add_data_points
+from cognee.modules.engine.models import NodeSet
+from uuid import uuid5, NAMESPACE_OID
+from typing import Optional
+from cognee.modules.retrieval.utils.models import CogneeSearchSentiment
+from cognee.modules.users.models import User
+
+async def run_sentiment_analysis(prev_question: str, prev_answer: str, current_question: str, user: User):
+    text_input = f"""
+    Previous Q: {prev_question}
+    Answer: {prev_answer}
+    Current Q: {current_question}
+    """
+    user_id = str(user.id)
+    # Call LLM to classify sentiment
+    sentiment_result = await LLMGateway.acreate_structured_output(
+        text_input=text_input,
+        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  
+    )
+    sentiment_data_point = CogneeSearchSentiment(
+        id=uuid5(NAMESPACE_OID, name=user_id + current_question),
+        prev_question=prev_question,
+        prev_answer=prev_answer,
+        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")
+    )
+    await add_data_points(data_points=[sentiment_data_point], update_edge_collection=True)
+    return {
+        "prev_question" : prev_question,
+        "prev_answer" : prev_answer,
+        "current_question" :current_question,
+        "sentiment": sentiment_result.sentiment,
+        "score": sentiment_result.score
+    }

examples/python/sentiment_analysis_demo.py

@@ -0,0 +1,34 @@
+import cognee
+import asyncio
+
+async def main():
+    await cognee.prune.prune_data()
+    await cognee.prune.prune_system(metadata=True)
+
+    text = "Cognee turns documents into AI memory."
+
+    await cognee.add(text)
+    await cognee.cognify()
+
+    queries = [
+        "What does Cognee do?",
+        "How does Cognee store data?"
+    ]
+
+    all_results = {}
+
+    for q in queries:
+        results = await cognee.search(
+            query_text=q,
+            save_interaction=True,
+        )
+        all_results[q] = results
+
+    # Print results
+    for query, res in all_results.items():
+        print(f"\nQuery: {query}")
+        for r in res:
+            print(f"  - {r}")
+
+if __name__ == '__main__':
+    asyncio.run(main())