Visual changes for the privatemode ai dispaly (#190)

* Adds enum

* Adds reranker step type

* Add new message field

* Adds data for messages

* Adds score mult for requests

* Submodule ref fix

Author: JWittmeyer

cognition_objects/conversation.py

@@ -358,27 +358,6 @@ def create(
     return conversation
 
 
-# TODO: replace usage with create from message.py and delete this function
-def add_message(
-    project_id: str,
-    conversation_id: str,
-    user_id: str,
-    question: str,
-    initiated_via: str,
-    with_commit: bool = True,
-) -> CognitionConversation:
-    message_entity = message.create(
-        conversation_id=conversation_id,
-        project_id=project_id,
-        user_id=user_id,
-        question=question,
-        initiated_via=initiated_via,
-        with_commit=with_commit,
-    )
-
-    return message_entity
-
-
 def update(
     project_id: str,
     conversation_id: str,

cognition_objects/message.py

@@ -5,6 +5,7 @@
 from ..models import CognitionMessage
 from ..util import prevent_sql_injection
 from .pipeline_version import get_current_version
+from sqlalchemy.orm.attributes import flag_modified
 
 
 DEFAULT_TIME_ELAPSED = {
@@ -268,6 +269,51 @@ def get_message_feedback_overview(
     return general.execute_all(query)
 
 
+def get_show_shield_dict_by_conversation_ids(
+    project_id: str, conversation_ids: List[str]
+) -> Dict[str, bool]:
+    if not conversation_ids or not project_id:
+        return {}
+    project_id = prevent_sql_injection(project_id, isinstance(project_id, str))
+    conversation_ids = [
+        prevent_sql_injection(conversation_id, isinstance(conversation_id, str))
+        for conversation_id in conversation_ids
+    ]
+    conversation_id_filter = "('" + "','".join(conversation_ids) + "')"
+    query = f"""
+    SELECT jsonb_object_agg(t.conversation_id::text, t.show_shield_icon) AS convo_shields
+    FROM (
+    SELECT
+        c.id AS conversation_id,
+        COALESCE(
+        bool_and(
+            COALESCE(
+            (m.additional_data->'privacy_report'->>'is_private') IN ('A+','A','A-'),
+            false
+            )
+        )
+        AND
+        bool_and(
+            COALESCE(
+            (m.additional_data->'privacy_report'->>'any_privatemode_ai')::boolean,
+            false
+            )
+        ),
+        false
+        ) AS show_shield_icon
+    FROM cognition.conversation c
+    LEFT JOIN cognition.message m 
+        ON c.project_id = m.project_id AND m.conversation_id = c.id
+        WHERE c.project_id  = '{project_id}' AND c.id IN {conversation_id_filter}
+    GROUP BY c.id
+    ) t;"""
+
+    show_shield_dict = general.execute_first(query)
+    if show_shield_dict and show_shield_dict[0]:
+        return show_shield_dict[0]
+    return {}
+
+
 def create(
     conversation_id: str,
     project_id: str,
@@ -276,6 +322,7 @@ def create(
     initiated_via: str,
     with_commit: bool = True,
     created_at: Optional[datetime] = None,
+    additional_data: Optional[Dict[str, Any]] = None,
 ) -> CognitionMessage:
     version_id = None
     current_version = get_current_version(project_id)
@@ -289,6 +336,7 @@ def create(
         question=question,
         facts=[],
         version_id=version_id,
+        additional_data=additional_data or {},
         initiated_via=initiated_via,
     )
 
@@ -306,6 +354,7 @@ def update(
     feedback_value: Optional[str] = None,
     feedback_category: Optional[str] = None,
     feedback_message: Optional[str] = None,
+    additional_data: Optional[Union[Dict[str, Any], str]] = None,
     with_commit: bool = True,
 ) -> CognitionMessage:
     message = get(project_id, message_id)
@@ -321,6 +370,14 @@ def update(
         message.feedback_category = feedback_category
     if feedback_message is not None:
         message.feedback_message = feedback_message
+    if additional_data is not None:
+        if additional_data == "NULL":
+            message.additional_data = {}
+        if message.additional_data is None:
+            message.additional_data = {}
+        for key, value in additional_data.items():
+            message.additional_data[key] = value
+        flag_modified(message, "additional_data")
 
     general.flush_or_commit(with_commit)
 

enums.py

@@ -566,6 +566,7 @@ class StrategyStepType(Enum):
     WEBHOOK = "WEBHOOK"
     GRAPHRAG_SEARCH = "GRAPHRAG_SEARCH"
     TEMPLATED = "TEMPLATED"
+    RERANKER = "RERANKER"
 
     def get_description(self):
         return STEP_DESCRIPTIONS.get(self, "No description available")
@@ -594,6 +595,7 @@ def get_progress_text(self):
     StrategyStepType.WEBHOOK: "Webhook",
     StrategyStepType.GRAPHRAG_SEARCH: "Query GraphRAG index",
     StrategyStepType.TEMPLATED: "Templated step",
+    StrategyStepType.RERANKER: "Reranker",
 }
 
 STEP_WHEN_TO_USE = {
@@ -612,6 +614,7 @@ def get_progress_text(self):
     StrategyStepType.WEBHOOK: "When you want to run a webhook",
     StrategyStepType.GRAPHRAG_SEARCH: "When you want to query a knowledge graph",
     StrategyStepType.TEMPLATED: "When you want to reuse existing templates",
+    StrategyStepType.RERANKER: "When you want to rerank results",
 }
 
 STEP_PROGRESS_TEXTS = {
@@ -631,6 +634,7 @@ def get_progress_text(self):
     StrategyStepType.WEBHOOK: "Running webhook",
     StrategyStepType.GRAPHRAG_SEARCH: "Querying knowledge graph",
     StrategyStepType.TEMPLATED: "Running templated step",
+    StrategyStepType.RERANKER: "Running reranker",
 }
 
 STEP_ERRORS = {
@@ -958,6 +962,37 @@ class SharepointPropertySyncState(Enum):
     FAILED = "FAILED"
 
 
+class CognitionPrivateUsage(Enum):
+    KERN_INTERNAL = "KERN_INTERNAL"  # same server no connection outside VM & DB
+    KERN_EXTERNAL_RESOURCE = (
+        "KERN_EXTERNAL_RESOURCE"  # e.g. reranker, distinct from external
+    )
+    WEBHOOK_NO_DATA = "WEBHOOK_NO_DATA"  # webhook only sending conversation id
+    DATA_SEND_EXTERNALLY = (
+        "DATA_SEND_EXTERNALLY"  # e.g. webhook with data or reranker outside kern
+    )
+    AZURE = "AZURE"  # e.g. BYOK from customers
+    AZURE_KERN = "AZURE_KERN"  # e.g. sponsorship azure resources or e.g. mistral with key specific key
+    OPEN_AI = "OPEN_AI"  # hosted by openai
+    PRIVATEMODE_AI = "PRIVATEMODE_AI"  # encrypted and hosted by privatemode.ai
+    REQUESTS_USED = "REQUESTS_USED"  # special case, is set handled separately => only for final result set
+    __SCORES = {
+        "KERN_INTERNAL": 1.0,
+        "KERN_EXTERNAL_RESOURCE": 0.95,
+        "WEBHOOK_NO_DATA": 0.9,
+        "DATA_SEND_EXTERNALLY": 0.3,
+        "AZURE": 0.5,
+        "AZURE_KERN": 0.85,
+        "OPEN_AI": 0.3,
+        "PRIVATEMODE_AI": 0.95,
+        "REQUESTS_USED": 0.9,  # final project multiplier
+    }
+
+    @property
+    def score(self) -> float:
+        return self.__SCORES[self.value]
+
+
 class MessageInitiationType(Enum):
     UI = "UI"
     API = "API"

models.py

@@ -1286,6 +1286,8 @@ class CognitionMessage(Base):
         ),
         index=True,
     )
+    # holds e.g. the project privacy report rating (high level) at message creation
+    additional_data = Column(JSON)
     initiated_via = Column(String)  # of type enums.MessageInitiationType.*.value
 
 

util.py

@@ -7,6 +7,7 @@
 import decimal
 from uuid import UUID
 from datetime import datetime, date
+from enum import Enum
 
 
 from sqlalchemy.sql import text as sql_text
@@ -225,6 +226,8 @@ def to_json_serializable(x: Any):
         return float(x)
     elif isinstance(x, UUID):
         return str(x)
+    elif isinstance(x, Enum):
+        return x.value
     else:
         return x