You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
152 lines
5.0 KiB
152 lines
5.0 KiB
from exec_anaconda import exec_anaconda_or_die
|
|
|
|
exec_anaconda_or_die()
|
|
|
|
from util import command_util
|
|
from cexc import BaseChunkHandler, CommandType
|
|
from util.telemetry_util import log_uuid, log_ai_commander_time, Timer
|
|
from util.param_util import parse_args
|
|
from chunked_controller import ChunkedController
|
|
from ai_commander.ai_commander_util import AICommanderUtil
|
|
from ai_commander.constants import AI_COMMANDER_COMMAND_CAPABILITIES
|
|
|
|
import cexc
|
|
|
|
logger = cexc.get_logger('ai')
|
|
|
|
PROCESSOR = 'AiCommanderProcessor'
|
|
|
|
|
|
class AICommander(BaseChunkHandler):
|
|
@staticmethod
|
|
def handle_arguments(getinfo: dict) -> dict:
|
|
"""
|
|
Parses and validates arguments passed to the AICommander.
|
|
|
|
Args:
|
|
getinfo (dict): Dictionary containing search information.
|
|
|
|
Returns:
|
|
dict: Processed controller options.
|
|
"""
|
|
if len(getinfo['searchinfo']['args']) == 0:
|
|
raise RuntimeError('First argument must be a prompt.')
|
|
raw_options = parse_args(getinfo['searchinfo']['raw_args'])
|
|
controller_options = AICommander.handle_raw_options(raw_options)
|
|
controller_options['processor'] = PROCESSOR
|
|
return controller_options
|
|
|
|
@staticmethod
|
|
def handle_raw_options(raw_options: dict) -> dict:
|
|
"""
|
|
Validates raw options provided in the command.
|
|
|
|
Args:
|
|
raw_options (dict): Raw options extracted from the command.
|
|
|
|
Returns:
|
|
dict: Filtered and validated options.
|
|
"""
|
|
allowed_options = {'prompt', 'model', 'provider', 'kb_id'}
|
|
|
|
for key in raw_options['params']:
|
|
if key not in allowed_options:
|
|
raise RuntimeError("Param name {} is not allowed".format(key))
|
|
|
|
return raw_options
|
|
|
|
def setup(self, metadata: dict) -> dict:
|
|
"""
|
|
Initializes the AICommander, verifies user eligibility, and prepares the execution environment.
|
|
|
|
Args:
|
|
metadata (dict): Metadata associated with the request.
|
|
|
|
Returns:
|
|
dict: Execution type and required fields.
|
|
"""
|
|
controller_options = self.handle_arguments(self.getinfo)
|
|
is_user_eligible = AICommanderUtil(
|
|
searchinfo=self.getinfo['searchinfo']
|
|
).check_user_role_eligibility(required_capabilities=AI_COMMANDER_COMMAND_CAPABILITIES)
|
|
if not is_user_eligible:
|
|
raise RuntimeError('User does not have permission to use `ai` command.')
|
|
self.controller = ChunkedController(self.getinfo, controller_options)
|
|
self.totale_df_count = 0
|
|
required_fields = self.controller.get_required_fields()
|
|
exec_type = CommandType.STATEFUL
|
|
return {'type': exec_type, 'required_fields': required_fields}
|
|
|
|
def _setup_watchdog(self) -> None:
|
|
"""
|
|
Initializes and starts the watchdog to monitor execution.
|
|
"""
|
|
"""Initialize and start watchdog"""
|
|
self.watchdog = command_util.get_watchdog(time_limit=-1, memory_limit=20000)
|
|
self.watchdog.start()
|
|
|
|
def handler(self, metadata: dict, body: bytes) -> tuple[dict, bytes]:
|
|
"""
|
|
Handles incoming data chunks and processes them accordingly.
|
|
|
|
Args:
|
|
metadata (dict): Metadata of the chunk.
|
|
body (bytes): The data chunk body.
|
|
|
|
Returns:
|
|
Tuple[dict, bytes]: Processed metadata and output data.
|
|
"""
|
|
if command_util.should_early_return(metadata):
|
|
return {}
|
|
|
|
if command_util.is_getinfo_chunk(metadata):
|
|
return self.setup(metadata)
|
|
finished_flag = metadata.get('finished', False)
|
|
|
|
if not self.watchdog:
|
|
self._setup_watchdog()
|
|
|
|
# Skip to next chunk if this chunk is empty
|
|
if len(body) == 0:
|
|
return {}
|
|
|
|
# Load data, execute and collect results.
|
|
with Timer() as load_t:
|
|
self.controller.load_data(body)
|
|
|
|
self.totale_df_count += self.controller.processor.df.shape[0]
|
|
|
|
logger.debug(f"command=ai," f" spl_load_data_time={load_t.interval}")
|
|
|
|
logger.debug(
|
|
f"command=ai,"
|
|
f" chunked_df_rows_count={self.controller.processor.df.shape[0]},"
|
|
f" total_df_rows_count={self.totale_df_count}"
|
|
)
|
|
|
|
with Timer() as execute_t:
|
|
self.controller.execute()
|
|
|
|
logger.debug(f"command=ai," f" spl_execute_time={execute_t.interval}")
|
|
with Timer() as output_t:
|
|
output_body = self.controller.output_results()
|
|
|
|
logger.debug(f"command=ai," f" spl_output_results_time={output_t.interval}")
|
|
|
|
if finished_flag:
|
|
# Gracefully terminate watchdog
|
|
if self.watchdog.started:
|
|
self.watchdog.join()
|
|
|
|
# Our final farewell
|
|
return ({'finished': finished_flag}, output_body)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
logger.debug("Starting aicommander.py.")
|
|
with Timer() as t:
|
|
AICommander(handler_data=BaseChunkHandler.DATA_RAW).run()
|
|
log_uuid()
|
|
log_ai_commander_time(t.interval)
|
|
logger.debug("Exiting gracefully. Byee!!")
|