Rework Factories (#149)

* add rename mappine

* blacklist

* PR

* recursive renaming

* PR comments

* PR comments

* PR

* empty conversations

Author: LennartSchmidtKern

cognition_objects/message.py

@@ -21,6 +21,34 @@ def get_all_by_conversation_id(
     )
 
 
+def get_all_by_conversation_ids(
+    project_id: str, conversation_ids: List[str]
+) -> Dict[str, List[CognitionMessage]]:
+
+    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_where = (
+        " AND conversation_id IN ('" + "','".join(conversation_ids) + "')"
+    )
+    query = f"""
+    SELECT jsonb_object_agg(conversation_id, messages)
+    FROM (
+        SELECT m.conversation_id, array_agg(row_to_json(m) ORDER BY created_at ASC) AS messages
+        FROM cognition.message m
+        WHERE project_id = '{project_id}'{conversation_where}
+        GROUP BY conversation_id
+    ) x
+    """
+
+    message_info = general.execute_first(query)
+    if message_info and message_info[0]:
+        return message_info[0]
+    return {}
+
+
 def get_last_by_conversation_id(
     project_id: str, conversation_id: str
 ) -> CognitionMessage:

cognition_objects/pipeline_log.py

@@ -4,6 +4,7 @@
 from ..models import CognitionPipelineLogs, CognitionMessage
 from datetime import datetime
 from .. import enums
+from ..util import prevent_sql_injection
 
 
 def get_all_by_message_id(
@@ -220,3 +221,33 @@ def update_to_new_diff_structure(
 
     if with_commit:
         general.commit()
+
+
+def get_error_and_time_elapsed_by_conversation_ids(
+    project_id: str,
+    conversation_ids: List[str],
+) -> Dict[str, CognitionPipelineLogs]:
+    if not conversation_ids:
+        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_where = (
+        " AND m.conversation_id IN ('" + "','".join(conversation_ids) + "')"
+    )
+    query = f"""
+    SELECT jsonb_object_agg(id, jsonb_build_object('logs_have_error', has_error, 'time_logs_elapsed', time_elapsed))
+    FROM (
+        SELECT m.id, SUM(pl.has_error::INT) > 0 has_error,SUM(pl.time_elapsed) time_elapsed
+        FROM cognition.message m
+        INNER JOIN cognition.pipeline_logs pl
+            ON m.project_id = pl.project_id AND m.id = pl.message_id
+        WHERE m.project_Id = '{project_id}'{conversation_where}
+        GROUP BY m.id
+    )x"""
+    conversation_info = general.execute_first(query)
+    if conversation_info and conversation_info[0]:
+        return conversation_info[0]
+    return {}

util.py

@@ -96,8 +96,12 @@ def sql_alchemy_to_dict(
     sql_alchemy_object: Any,
     for_frontend: bool = False,
     column_whitelist: Optional[Iterable[str]] = None,
+    column_blacklist: Optional[Iterable[str]] = None,
+    column_rename_map: Optional[Dict[str, str]] = None,
 ):
-    result = __sql_alchemy_to_dict(sql_alchemy_object, column_whitelist)
+    result = __sql_alchemy_to_dict(
+        sql_alchemy_object, column_whitelist, column_blacklist, column_rename_map
+    )
     if for_frontend:
         return to_frontend_obj(result)
     return result
@@ -136,26 +140,58 @@ def convert_value(value, max_lvl: int):
 
 
 def __sql_alchemy_to_dict(
-    sql_alchemy_object: Any, column_whitelist: Optional[Iterable[str]] = None
+    sql_alchemy_object: Any,
+    column_whitelist: Optional[Iterable[str]] = None,
+    column_blacklist: Optional[Iterable[str]] = None,
+    column_rename_map: Optional[Dict[str, str]] = None,
 ):
+    def rename_columns(data: Any) -> Any:
+        if column_rename_map:
+            if isinstance(data, dict):
+                data = {
+                    column_rename_map.get(k, k): rename_columns(v)
+                    for k, v in data.items()
+                }
+            elif isinstance(data, list):
+                data = [rename_columns(item) for item in data]
+        return data
+
     if isinstance(sql_alchemy_object, list):
         # list is for all() queries
-        return [__sql_alchemy_to_dict(x, column_whitelist) for x in sql_alchemy_object]
+        return [
+            __sql_alchemy_to_dict(
+                x, column_whitelist, column_blacklist, column_rename_map
+            )
+            for x in sql_alchemy_object
+        ]
 
     elif isinstance(sql_alchemy_object, Row):
         # basic SELECT .. FROM query)
         # _mapping is a RowMapping object that is not serializable but dict like
-        return {
+        result = {
             k: v
             for k, v in dict(sql_alchemy_object._mapping).items()
-            if not column_whitelist or k in column_whitelist
+            if (not column_whitelist or k in column_whitelist)
+            and (not column_blacklist or k not in column_blacklist)
         }
+        return rename_columns(result)
     elif isinstance(sql_alchemy_object, Base):
-        return {
+        result = {
             c.name: getattr(sql_alchemy_object, c.name)
             for c in sql_alchemy_object.__table__.columns
-            if not column_whitelist or c.name in column_whitelist
+            if (not column_whitelist or c.name in column_whitelist)
+            and (not column_blacklist or c.name not in column_blacklist)
+        }
+        return rename_columns(result)
+    elif isinstance(sql_alchemy_object, dict):
+        result = {
+            k: v
+            for k, v in sql_alchemy_object.items()
+            if (not column_whitelist or k in column_whitelist)
+            and (not column_blacklist or k not in column_blacklist)
         }
+        return rename_columns(result)
+
     else:
         return sql_alchemy_object