ReMe/reme_ai/retrieve/personal/set_query_op.py

44 lines
1.5 KiB
Python

"""Query setting operation for personal memories.
This module provides functionality to set query and timestamp in the context
for downstream memory retrieval operations.
"""
import datetime
from typing import Tuple
from flowllm.core.context import C
from flowllm.core.op import BaseAsyncOp
from loguru import logger
from reme_ai.constants.common_constants import QUERY_WITH_TS
@C.register_op()
class SetQueryOp(BaseAsyncOp):
"""
The `SetQueryOp` class is responsible for setting a query and its associated timestamp
into the context, utilizing either provided parameters or details from the context.
"""
async def async_execute(self):
"""
Executes the operation's primary function, which involves determining the query and its
timestamp, then storing these values within the context.
Input requirement: self.context.query must exist (flow input requirement)
"""
# Flow guarantees query exists - use it directly
query: str = self.context.query
timestamp: int = int(datetime.datetime.now().timestamp())
# Set timestamp if provided in op_params
_timestamp = self.op_params.get("timestamp")
if _timestamp and isinstance(_timestamp, int):
timestamp = _timestamp
# Store the query and its timestamp in the context
query_with_ts: Tuple[str, int] = (query, timestamp)
self.context[QUERY_WITH_TS] = query_with_ts
logger.info(f"Set query with timestamp: query='{query}', timestamp={timestamp}")