Skip to content
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
19 changes: 14 additions & 5 deletions src/memos/api/handlers/suggestion_handler.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,28 +17,37 @@
SUGGESTION_QUERY_PROMPT_EN,
SUGGESTION_QUERY_PROMPT_ZH,
)
from memos.types import MessageList
from memos.types import MessageList, MessagesType


logger = get_logger(__name__)


def _get_further_suggestion(
llm: Any,
message: MessageList,
message: MessageList | str,
) -> list[str]:
"""
Get further suggestion based on recent dialogue.

Args:
llm: LLM instance for generating suggestions
message: Recent chat messages
message: Recent chat messages (can be a list of message dicts or a plain string)

Returns:
List of suggestion queries
"""
try:
dialogue_info = "\n".join([f"{msg['role']}: {msg['content']}" for msg in message[-2:]])
if isinstance(message, str):
dialogue_info = message
else:
dialogue_info = "\n".join(
[
f"{msg['role']}: {msg['content']}"
for msg in message[-2:]
if isinstance(msg, dict)
]
)
further_suggestion_prompt = FURTHER_SUGGESTION_PROMPT.format(dialogue=dialogue_info)
message_list = [{"role": "system", "content": further_suggestion_prompt}]
response = llm.generate(message_list)
Expand All @@ -53,7 +62,7 @@ def _get_further_suggestion(
def handle_get_suggestion_queries(
user_id: str,
language: str,
message: MessageList | None,
message: MessagesType | None,
llm: Any,
naive_mem_cube: Any,
) -> SuggestionResponse:
Expand Down
Loading