ReMe/reme_ai/retrieve/task/build_query_op.py

61 lines
2.3 KiB
Python

"""Query building operation module.
This module provides functionality to build retrieval queries from either
explicit query strings or conversation messages, optionally using LLM to
generate optimized queries.
"""
from flowllm.core.context import C
from flowllm.core.enumeration import Role
from flowllm.core.op import BaseAsyncOp
from flowllm.core.schema import Message
from flowllm.core.utils import merge_messages_content
from loguru import logger
@C.register_op()
class BuildQueryOp(BaseAsyncOp):
"""Build retrieval query from context or messages.
This operation constructs a query string for memory retrieval. It can use
an explicit query from context, or generate one from conversation messages
using either LLM-based generation or simple message concatenation.
"""
file_path: str = __file__
async def async_execute(self):
"""Execute the query building operation.
Builds a query string from either:
1. An explicit query in the context
2. Conversation messages (using LLM or simple concatenation)
Stores the built query in context.query.
"""
if "query" in self.context:
query = self.context.query
elif "messages" in self.context:
if self.op_params.get("enable_llm_build", True):
execution_process = merge_messages_content(self.context.messages)
prompt = self.prompt_format(prompt_name="query_build", execution_process=execution_process)
message = await self.llm.achat(messages=[Message(role=Role.USER, content=prompt)])
query = message.content
else:
context_parts = []
message_summaries = []
for message in self.context.messages[-3:]: # Last 3 messages
content = message.content[:200] + "..." if len(message.content) > 200 else message.content
message_summaries.append(f"- {message.role.value}: {content}")
if message_summaries:
context_parts.append("Recent messages:\n" + "\n".join(message_summaries))
query = "\n\n".join(context_parts)
else:
raise RuntimeError("query or messages is required!")
logger.info(f"build.query={query}")
self.context.query = query