feat: commit profiles in one session

Author: gusye1234

src/server/api/memobase_server/controllers/modal/chat/__init__.py

@@ -1,18 +1,19 @@
+from ...project import get_project_profile_config
 from ....connectors import Session
 from ....env import LOG, ProfileConfig, CONFIG
 from ....utils import get_blob_str, get_encoded_tokens
 from ....models.blob import Blob
 from ....models.utils import Promise, CODE
 from ....models.response import IdsData, ChatModalResponse
-from ...profile import add_user_profiles, update_user_profiles, delete_user_profiles
+from ...profile import add_update_delete_user_profiles
 from ...event import append_user_event
 from .extract import extract_topics
 from .merge import merge_or_valid_new_memos
 from .summary import re_summary
 from .organize import organize_profiles
 from .types import MergeAddResult
 from .event_summary import tag_event
-from .entry_summary import entry_summary
+from .entry_summary import entry_chat_summary
 
 
 def truncate_chat_blobs(
@@ -43,12 +44,18 @@ async def process_blobs(
         return Promise.reject(
             CODE.SERVER_PARSE_ERROR, "No blobs to process after truncating"
         )
-    p = await entry_summary(user_id, project_id, blobs)
+
+    p = await get_project_profile_config(project_id)
+    if not p.ok():
+        return p
+    project_profiles = p.data()
+
+    p = await entry_chat_summary(user_id, project_id, blobs, project_profiles)
     if not p.ok():
         return p
     user_memo_str = p.data()
 
-    p = await extract_topics(user_id, project_id, user_memo_str)
+    p = await extract_topics(user_id, project_id, user_memo_str, project_profiles)
     if not p.ok():
         return p
     extracted_data = p.data()
@@ -59,7 +66,7 @@ async def process_blobs(
         fact_contents=extracted_data["fact_contents"],
         fact_attributes=extracted_data["fact_attributes"],
         profiles=extracted_data["profiles"],
-        config=extracted_data["config"],
+        config=project_profiles,
         total_profiles=extracted_data["total_profiles"],
     )
     if not p.ok():
@@ -74,7 +81,7 @@ async def process_blobs(
     p = await organize_profiles(
         project_id,
         profile_options,
-        config=extracted_data["config"],
+        config=project_profiles,
     )
     if not p.ok():
         LOG.error(f"Failed to organize profiles: {p.msg()}")
@@ -94,29 +101,21 @@ async def process_blobs(
         project_id,
         user_memo_str,
         delta_profile_data,
-        extracted_data["config"],
+        project_profiles,
     )
     if not p.ok():
         return p
     eid = p.data()
-    p = await exe_user_profile_add(user_id, project_id, profile_options)
-    if not p.ok():
-        return p
-    add_profile_ids = p.data().ids
-    p = await exe_user_profile_update(user_id, project_id, profile_options)
-    if not p.ok():
-        return p
-    update_profile_ids = p.data().ids
-    p = await exe_user_profile_delete(user_id, project_id, profile_options)
+
+    p = await handle_user_profile_db(user_id, project_id, profile_options)
     if not p.ok():
         return p
-    delete_profile_ids = p.data().ids
     return Promise.resolve(
         ChatModalResponse(
             event_id=eid,
-            add_profiles=add_profile_ids,
-            update_profiles=update_profile_ids,
-            delete_profiles=delete_profile_ids,
+            add_profiles=p.data().ids,
+            update_profiles=[up["profile_id"] for up in profile_options["update"]],
+            delete_profiles=profile_options["delete"],
         )
     )
 
@@ -149,44 +148,21 @@ async def handle_session_event(
     return eid
 
 
-async def exe_user_profile_add(
+async def handle_user_profile_db(
     user_id: str, project_id: str, profile_options: MergeAddResult
 ) -> Promise[IdsData]:
-    if not len(profile_options["add"]):
-        return Promise.resolve(IdsData(ids=[]))
     LOG.info(f"Adding {len(profile_options['add'])} profiles for user {user_id}")
-    task_add = await add_user_profiles(
+    LOG.info(f"Updating {len(profile_options['update'])} profiles for user {user_id}")
+    LOG.info(f"Deleting {len(profile_options['delete'])} profiles for user {user_id}")
+
+    p = await add_update_delete_user_profiles(
         user_id,
         project_id,
         [ap["content"] for ap in profile_options["add"]],
         [ap["attributes"] for ap in profile_options["add"]],
-    )
-    return task_add
-
-
-async def exe_user_profile_update(
-    user_id: str, project_id: str, profile_options: MergeAddResult
-) -> Promise[IdsData]:
-    if not len(profile_options["update"]):
-        return Promise.resolve(IdsData(ids=[]))
-    LOG.info(f"Updating {len(profile_options['update'])} profiles for user {user_id}")
-    task_update = await update_user_profiles(
-        user_id,
-        project_id,
         [up["profile_id"] for up in profile_options["update"]],
         [up["content"] for up in profile_options["update"]],
         [up["attributes"] for up in profile_options["update"]],
+        profile_options["delete"],
     )
-    return task_update
-
-
-async def exe_user_profile_delete(
-    user_id: str, project_id: str, profile_options: MergeAddResult
-) -> Promise[IdsData]:
-    if not len(profile_options["delete"]):
-        return Promise.resolve(IdsData(ids=[]))
-    LOG.info(f"Deleting {len(profile_options['delete'])} profiles for user {user_id}")
-    task_delete = await delete_user_profiles(
-        user_id, project_id, profile_options["delete"]
-    )
-    return task_delete
+    return p

src/server/api/memobase_server/controllers/modal/chat/entry_summary.py

@@ -4,20 +4,16 @@
 from ....models.blob import Blob, BlobType
 from ....llms import llm_complete
 from ....prompts.profile_init_utils import read_out_profile_config
-from ...project import get_project_profile_config
+from ...project import ProfileConfig
 from ....prompts.profile_init_utils import read_out_event_tags
 from ....prompts.utils import tag_chat_blobs_in_order_xml
 from .types import FactResponse, PROMPTS
 
 
-async def entry_summary(
-    user_id: str, project_id: str, blobs: list[Blob]
+async def entry_chat_summary(
+    user_id: str, project_id: str, blobs: list[Blob], project_profiles: ProfileConfig
 ) -> Promise[str]:
     assert all(b.type == BlobType.chat for b in blobs), "All blobs must be chat blobs"
-    p = await get_project_profile_config(project_id)
-    if not p.ok():
-        return p
-    project_profiles = p.data()
     USE_LANGUAGE = project_profiles.language or CONFIG.language
     project_profiles_slots = read_out_profile_config(
         project_profiles, PROMPTS[USE_LANGUAGE]["profile"].CANDIDATE_PROFILE_TOPICS

src/server/api/memobase_server/controllers/modal/chat/extract.py

@@ -12,7 +12,7 @@
 )
 from ....prompts.profile_init_utils import read_out_profile_config, UserProfileTopic
 from ...profile import get_user_profiles
-from ...project import get_project_profile_config
+from ...project import ProfileConfig
 
 # from ...project impor
 from .types import FactResponse, PROMPTS
@@ -30,16 +30,12 @@ def merge_by_topic_sub_topics(new_facts: list[FactResponse]):
 
 
 async def extract_topics(
-    user_id: str, project_id: str, user_memo: str
+    user_id: str, project_id: str, user_memo: str, project_profiles: ProfileConfig
 ) -> Promise[dict]:
     p = await get_user_profiles(user_id, project_id)
     if not p.ok():
         return p
     profiles = p.data().profiles
-    p = await get_project_profile_config(project_id)
-    if not p.ok():
-        return p
-    project_profiles = p.data()
     USE_LANGUAGE = project_profiles.language or CONFIG.language
     STRICT_MODE = (
         project_profiles.profile_strict_mode
@@ -113,7 +109,6 @@ async def extract_topics(
                 "fact_contents": [],
                 "fact_attributes": [],
                 "profiles": profiles,
-                "config": project_profiles,
                 "total_profiles": project_profiles_slots,
             }
         )
@@ -145,7 +140,6 @@ async def extract_topics(
             "fact_contents": fact_contents,
             "fact_attributes": fact_attributes,
             "profiles": profiles,
-            "config": project_profiles,
             "total_profiles": project_profiles_slots,
         }
     )

src/server/api/memobase_server/controllers/profile.py

@@ -213,3 +213,88 @@ async def refresh_user_profile_cache(user_id: str, project_id: str) -> Promise[N
     async with get_redis_client() as redis_client:
         await redis_client.delete(f"user_profiles::{project_id}::{user_id}")
     return Promise.resolve(None)
+
+
+async def add_update_delete_user_profiles(
+    user_id: str,
+    project_id: str,
+    add_profiles: list[str],
+    add_attributes: list[dict],
+    update_profile_ids: list[str],
+    update_contents: list[str],
+    update_attributes: list[dict | None],
+    delete_profile_ids: list[str],
+) -> Promise[IdsData]:
+    assert len(add_profiles) == len(
+        add_attributes
+    ), "Length of add_profiles, add_attributes must be equal"
+    assert len(update_profile_ids) == len(
+        update_contents
+    ), "Length of update_profile_ids, update_contents must be equal"
+    assert len(update_profile_ids) == len(
+        update_attributes
+    ), "Length of update_profile_ids, update_attributes must be equal"
+
+    for attr in add_attributes + update_attributes:
+        if attr is None:
+            continue
+        try:
+            ProfileAttributes.model_validate(attr)
+        except ValidationError as e:
+            return Promise.reject(
+                CODE.SERVER_PARSE_ERROR, f"Invalid profile attributes: {e}"
+            )
+    # Sanity Check done
+
+    with Session() as session:
+        try:
+            # 1. add new profiles
+            if len(add_profiles):
+                add_db_profiles = [
+                    UserProfile(
+                        user_id=user_id,
+                        project_id=project_id,
+                        content=content,
+                        attributes=attr,
+                    )
+                    for content, attr in zip(add_profiles, add_attributes)
+                ]
+                session.add_all(add_db_profiles)
+                add_profile_ids = [p.id for p in add_db_profiles]
+            else:
+                add_profile_ids = []
+            # 2. update existing profiles
+            update_db_profiles = []
+            for profile_id, content, attribute in zip(
+                update_profile_ids, update_contents, update_attributes
+            ):
+                db_profile = (
+                    session.query(UserProfile)
+                    .filter_by(id=profile_id, user_id=user_id, project_id=project_id)
+                    .one_or_none()
+                )
+                if db_profile is None:
+                    LOG.error(f"Profile {profile_id} not found for user {user_id}")
+                    continue
+                db_profile.content = content
+                if attribute is not None:
+                    db_profile.attributes = attribute
+                update_db_profiles.append(profile_id)
+
+            # 3. delete profiles
+            session.query(UserProfile).filter(
+                UserProfile.id.in_(delete_profile_ids),
+                UserProfile.user_id == user_id,
+                UserProfile.project_id == project_id,
+            ).delete(synchronize_session=False)
+
+            session.commit()
+        except Exception as e:
+            LOG.error(f"Error merging user profiles: {e}")
+            session.rollback()
+            return Promise.reject(
+                CODE.SERVER_PARSE_ERROR, f"Error merging user profiles: {e}"
+            )
+
+    await refresh_user_profile_cache(user_id, project_id)
+    return Promise.resolve(IdsData(ids=add_profile_ids))

src/server/api/memobase_server/telemetry/open_telemetry.py

@@ -140,14 +140,14 @@ def setup_telemetry(self) -> None:
 
     def _construct_attributes(self, **kwargs) -> Dict[str, str]:
 
-        if os.environ.get("POD_IP"):
-            # use k8s downward API to get the pod ip
-            pod_ip = os.environ.get("POD_IP")
-        else:
-            # use the hostname to get the ip address
-            hostname = socket.gethostname()
-            pod_ip = socket.gethostbyname(hostname)
-
+        # if os.environ.get("POD_IP"):
+        #     # use k8s downward API to get the pod ip
+        #     pod_ip = os.environ.get("POD_IP", None)
+        # else:
+        #     # use the hostname to get the ip address
+        #     hostname = socket.gethostname()
+        #     pod_ip = socket.gethostbyname(hostname)
+        pod_ip = os.environ.get("POD_IP", None)
         return {
             DEPLOYMENT_ENVIRONMENT: self._deployment_environment,
             "memobase_server_ip": pod_ip,