feat: add demo plot infer

Author: gusye1234

src/client/memobase/core/entry.py

@@ -119,6 +119,7 @@ def insert(self, blob_data: Blob) -> str:
                 json=blob_data.to_request(),
             )
         )
+        print(r)
         return r.data["id"]
 
     def get(self, blob_id: str) -> Blob:

src/server/api/api.py

@@ -245,5 +245,11 @@ def custom_openapi():
 )(api_layer.context.get_user_context)
 
 
+router.post(
+    "/users/roleplay/proactive/{user_id}",
+    tags=["roleplay"],
+    # openapi_extra=API_X_CODE_DOCS["POST /users/roleplay/proactive/{user_id}"],
+)(api_layer.roleplay.infer_proactive_topics)
+
 app.include_router(router)
 app.add_middleware(api_layer.middleware.AuthMiddleware)

src/server/api/memobase_server/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "0.0.33"
+__version__ = "0.0.34"
 
 __author__ = "memobase.io"
 __url__ = "https://github.com/memodb-io/memobase"

src/server/api/memobase_server/api_layer/__init__.py

@@ -7,3 +7,4 @@
 from . import event
 from . import context
 from . import middleware
+from . import roleplay

src/server/api/memobase_server/api_layer/roleplay.py

@@ -0,0 +1,60 @@
+import json
+from ..controllers import full as controllers
+from ..controllers.modal.roleplay import proactive_topics
+from ..models.blob import BlobType
+from ..models.utils import Promise, CODE
+from ..models import response as res
+from fastapi import Request
+from fastapi import Body, Path, Query
+
+
+async def infer_proactive_topics(
+    request: Request,
+    user_id: str = Path(..., description="The ID of the user"),
+    topk: int = Query(
+        None, description="Number of profiles to retrieve, default is all"
+    ),
+    max_token_size: int = Query(
+        None,
+        description="Max token size of returned profile content, default is all",
+    ),
+    prefer_topics: list[str] = Query(
+        None,
+        description="Rank prefer topics at first to try to keep them in filtering, default order is by updated time",
+    ),
+    only_topics: list[str] = Query(
+        None,
+        description="Only return profiles with these topics, default is all",
+    ),
+    max_subtopic_size: int = Query(
+        None,
+        description="Max subtopic size of the same topic in returned profile, default is all",
+    ),
+    topic_limits_json: str = Query(
+        None,
+        description='Set specific subtopic limits for topics in JSON, for example {"topic1": 3, "topic2": 5}. The limits in this param will override `max_subtopic_size`.',
+    ),
+    body: res.ProactiveTopicRequest = Body(..., description="The body of the request"),
+) -> res.ProactiveTopicResponse:
+    """Get the real-time user profiles for long term memory"""
+    project_id = request.state.memobase_project_id
+    topic_limits_json = topic_limits_json or "{}"
+    try:
+        topic_limits = res.StrIntData(data=json.loads(topic_limits_json)).data
+    except Exception as e:
+        return Promise.reject(
+            CODE.BAD_REQUEST, f"Invalid JSON requests: {e}"
+        ).to_response(res.UserProfileResponse)
+    p = await proactive_topics.process_messages(
+        user_id,
+        project_id,
+        body.messages,
+        body.agent_context,
+        prefer_topics,
+        topk,
+        max_token_size,
+        only_topics,
+        max_subtopic_size,
+        topic_limits,
+    )
+    return p.to_response(res.ProactiveTopicResponse)

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

@@ -28,10 +28,7 @@ async def get_user_events(
             ).filter(UserEvent.event_data.has_key("event_tip"))
         user_events = query.order_by(UserEvent.created_at.desc()).limit(topk).all()
         if user_events is None:
-            return Promise.reject(
-                CODE.NOT_FOUND,
-                f"No user events found for user {user_id}",
-            )
+            return Promise.resolve(UserEventsData(events=[]))
         results = [
             {
                 "id": ue.id,

src/server/api/memobase_server/controllers/modal/roleplay/detect_interest.py

@@ -0,0 +1,35 @@
+from ....models.utils import Promise, CODE
+from ....env import CONFIG, LOG, ProfileConfig
+from ....utils import get_encoded_tokens, truncate_string
+from ....llms import llm_complete
+from ....models.blob import OpenAICompatibleMessage
+from .types import PROMPTS, ChatInterest
+from ..utils import try_json_loads
+
+
+async def detect_chat_interest(
+    project_id: str,
+    messages: list[OpenAICompatibleMessage],
+    profile_config: ProfileConfig,
+) -> Promise[ChatInterest]:
+    USE_LANGUAGE = "zh"
+    prompt = PROMPTS[USE_LANGUAGE]["detect_interest"]
+
+    r = await llm_complete(
+        project_id,
+        prompt.get_input(messages),
+        system_prompt=prompt.get_prompt(),
+        temperature=0.2,  # precise
+        model=CONFIG.best_llm_model,
+        **prompt.get_kwargs(),
+    )
+    if not r.ok():
+        return r
+    content = r.data()
+    data = try_json_loads(content)
+    print(data)
+    if data is None:
+        return Promise.reject(
+            CODE.INTERNAL_SERVER_ERROR, "Unable to parse the LLM json response"
+        )
+    return Promise.resolve(data)

src/server/api/memobase_server/controllers/modal/roleplay/predict_new_topics.py

@@ -0,0 +1,57 @@
+import re
+from ....models.utils import Promise, CODE
+from ....env import CONFIG, LOG, ProfileConfig
+from ....utils import get_encoded_tokens, truncate_string
+from ....llms import llm_complete
+from ....models.blob import OpenAICompatibleMessage
+from ....models.response import UserStatusesData
+from .types import PROMPTS, InferPlot
+
+
+def extract_plot_output(content: str):
+    themes = re.search(r"<themes>(.*?)</themes>", content, re.DOTALL)
+    overview = re.search(r"<overview>(.*?)</overview>", content, re.DOTALL)
+    timeline = re.search(r"<timeline>(.*?)</timeline>", content, re.DOTALL)
+    return (
+        themes.group(1).strip() if themes else None,
+        overview.group(1).strip() if overview else None,
+        timeline.group(1).strip() if timeline else None,
+    )
+
+
+async def predict_new_topics(
+    project_id: str,
+    messages: list[OpenAICompatibleMessage],
+    latest_statuses: UserStatusesData,
+    user_context: str,
+    agent_context: str,
+    profile_config: ProfileConfig,
+    max_before_old_topics: int = 5,
+) -> Promise[InferPlot]:
+    USE_LANGUAGE = "zh"
+    prompt = PROMPTS[USE_LANGUAGE]["infer_plot"]
+
+    latest_plots = [
+        ld.attributes["new_topic"]["overview"]
+        for ld in latest_statuses.statuses
+        if "new_topic" in ld.attributes
+    ][:max_before_old_topics]
+    print(
+        "THINK",
+        prompt.get_input(agent_context, user_context, latest_plots, messages),
+    )
+    r = await llm_complete(
+        project_id,
+        prompt.get_input(agent_context, user_context, latest_plots, messages),
+        system_prompt=prompt.get_prompt(),
+        temperature=0.2,  # precise
+        model=CONFIG.thinking_llm_model,
+        **prompt.get_kwargs(),
+        no_cache=True,
+    )
+    if not r.ok():
+        return r
+    content = r.data()
+    print(content)
+    themes, overview, timeline = extract_plot_output(content)
+    return Promise.resolve(dict(themes=themes, overview=overview, timeline=timeline))

src/server/api/memobase_server/controllers/modal/roleplay/proactive_topics.py

@@ -0,0 +1,116 @@
+from ....env import ContanstTable, CONFIG, LOG
+from ...status import append_user_status, get_user_statuses
+from ...profile import get_user_profiles, truncate_profiles
+from ...project import get_project_profile_config
+from ....models.blob import OpenAICompatibleMessage
+from ....models.utils import Promise
+from ....models.response import ProactiveTopicData
+from ...profile import get_user_profiles, truncate_profiles
+from .detect_interest import detect_chat_interest
+from .predict_new_topics import predict_new_topics
+
+# from .types import
+
+
+def pack_timeline_prompt(timeline: str, language: str) -> str:
+    if language == "zh":
+        return f"## 下面是你的剧本，如果我没有主动提供话题的话，参考下面剧情推动我们的对话：\n{timeline}##"
+    else:
+        return f"## Here is your script, if I don't provide a topic, please refer to the following plot to drive our conversation: \n{timeline}##"
+
+
+async def process_messages(
+    user_id: str,
+    project_id: str,
+    messages: list[OpenAICompatibleMessage],
+    agent_context: str = None,
+    prefer_topics: list[str] = None,
+    topk: int = None,
+    max_token_size: int = None,
+    only_topics: list[str] = None,
+    max_subtopic_size: int = None,
+    topic_limits: dict[str, int] = None,
+) -> Promise[ProactiveTopicData]:
+    p = await get_project_profile_config(project_id)
+    if not p.ok():
+        return p
+    project_profiles = p.data()
+    USE_LANGUAGE = "zh"
+    # USE_LANGUAGE = project_profiles.language or CONFIG.language
+
+    interest = await detect_chat_interest(
+        project_id,
+        messages,
+        profile_config=project_profiles,
+    )
+    if not interest.ok():
+        return interest
+    interest_data = interest.data()
+    # if interest_data["action"] != "new_topic":
+    #     await append_user_status(
+    #         user_id,
+    #         project_id,
+    #         ContanstTable.roleplay_plot_status,
+    #         {
+    #             "interest": interest_data,
+    #         },
+    #     )
+    #     return Promise.resolve(ProactiveTopicData(action="continue"))
+    latests_statuses = await get_user_statuses(
+        user_id, project_id, type=ContanstTable.roleplay_plot_status
+    )
+    if not latests_statuses.ok():
+        return latests_statuses
+    latests_statuses_data = latests_statuses.data()
+
+    p = await get_user_profiles(user_id, project_id)
+    if not p.ok():
+        return p
+    p = await truncate_profiles(
+        p.data(),
+        prefer_topics=prefer_topics,
+        topk=topk,
+        max_token_size=max_token_size,
+        only_topics=only_topics,
+        max_subtopic_size=max_subtopic_size,
+        topic_limits=topic_limits,
+    )
+    if not p.ok():
+        return p
+    user_profiles_data = p.data()
+    use_user_profiles = user_profiles_data.profiles
+    user_context = "\n".join(
+        [
+            f"{p.attributes.get('topic')}::{p.attributes.get('sub_topic')}: {p.content}"
+            for p in use_user_profiles
+        ]
+    )
+
+    p = await predict_new_topics(
+        project_id,
+        messages,
+        latests_statuses_data,
+        user_context,
+        agent_context,
+        project_profiles,
+    )
+    if not p.ok():
+        return p
+    plot = p.data()
+    await append_user_status(
+        user_id,
+        project_id,
+        ContanstTable.roleplay_plot_status,
+        {
+            "interest": interest_data,
+            "new_topic": plot,
+            "chats": [m.model_dump() for m in messages],
+        },
+    )
+
+    return Promise.resolve(
+        ProactiveTopicData(
+            action="new_topic",
+            topic_prompt=pack_timeline_prompt(plot["timeline"], USE_LANGUAGE),
+        )
+    )

src/server/api/memobase_server/controllers/modal/roleplay/types.py

@@ -0,0 +1,16 @@
+from typing import TypedDict
+from ....prompts.roleplay import zh_detect_interest, zh_infer_plot
+
+ChatInterest = TypedDict("ChatInterest", {"status": str, "action": str})
+
+InferPlot = TypedDict(
+    "InferPlot", {"themes": str | None, "overview": str | None, "timeline": str | None}
+)
+
+PROMPTS = {
+    "en": {},
+    "zh": {
+        "detect_interest": zh_detect_interest,
+        "infer_plot": zh_infer_plot,
+    },
+}