Aim

AimCallback

基类：BaseCallbackHandler

AimCallback 回调类。

参数

名称	类型	描述	默认值
repo		obj:str, 可选): Aim 仓库路径或绑定 Run 对象的 Repo 对象。如果跳过，则使用默认 Repo。	无
experiment_name		obj:str, 可选): 设置 Run 的 experiment 属性。未指定时为 'default'。稍后可用于查询运行/序列。	无
system_tracking_interval		obj:int, 可选): 设置系统使用指标（CPU、内存等）的跟踪间隔（秒）。设为 None 可禁用系统指标跟踪。	1
log_system_params		obj:bool, 可选): 启用/禁用系统参数（如已安装包、git 信息、环境变量等）的日志记录。	True
capture_terminal_logs		obj:bool, 可选): 启用/禁用终端 stdout 日志记录。	True
event_starts_to_ignore	可选[列表[CBEventType]]	跟踪事件开始时要忽略的事件类型列表。	无
event_ends_to_ignore	可选[列表[CBEventType]]	跟踪事件结束时要忽略的事件类型列表。	无

源代码位于 llama-index-integrations/callbacks/llama-index-callbacks-aim/llama_index/callbacks/aim/base.py

class AimCallback(BaseCallbackHandler):
    """
    AimCallback callback class.

    Args:
        repo (:obj:`str`, optional):
            Aim repository path or Repo object to which Run object is bound.
            If skipped, default Repo is used.
        experiment_name (:obj:`str`, optional):
            Sets Run's `experiment` property. 'default' if not specified.
            Can be used later to query runs/sequences.
        system_tracking_interval (:obj:`int`, optional):
            Sets the tracking interval in seconds for system usage
            metrics (CPU, Memory, etc.). Set to `None` to disable
            system metrics tracking.
        log_system_params (:obj:`bool`, optional):
            Enable/Disable logging of system params such as installed packages,
            git info, environment variables, etc.
        capture_terminal_logs (:obj:`bool`, optional):
            Enable/Disable terminal stdout logging.
        event_starts_to_ignore (Optional[List[CBEventType]]):
            list of event types to ignore when tracking event starts.
        event_ends_to_ignore (Optional[List[CBEventType]]):
            list of event types to ignore when tracking event ends.

    """

    def __init__(
        self,
        repo: Optional[str] = None,
        experiment_name: Optional[str] = None,
        system_tracking_interval: Optional[int] = 1,
        log_system_params: Optional[bool] = True,
        capture_terminal_logs: Optional[bool] = True,
        event_starts_to_ignore: Optional[List[CBEventType]] = None,
        event_ends_to_ignore: Optional[List[CBEventType]] = None,
        run_params: Optional[Dict[str, Any]] = None,
    ) -> None:
        if Run is None:
            raise ModuleNotFoundError(
                "Please install aim to use the AimCallback: 'pip install aim'"
            )

        event_starts_to_ignore = (
            event_starts_to_ignore if event_starts_to_ignore else []
        )
        event_ends_to_ignore = event_ends_to_ignore if event_ends_to_ignore else []
        super().__init__(
            event_starts_to_ignore=event_starts_to_ignore,
            event_ends_to_ignore=event_ends_to_ignore,
        )

        self.repo = repo
        self.experiment_name = experiment_name
        self.system_tracking_interval = system_tracking_interval
        self.log_system_params = log_system_params
        self.capture_terminal_logs = capture_terminal_logs
        self._run: Optional[Any] = None
        self._run_hash = None

        self._llm_response_step = 0

        self.setup(run_params)

    def on_event_start(
        self,
        event_type: CBEventType,
        payload: Optional[Dict[str, Any]] = None,
        event_id: str = "",
        parent_id: str = "",
        **kwargs: Any,
    ) -> str:
        """
        Args:
            event_type (CBEventType): event type to store.
            payload (Optional[Dict[str, Any]]): payload to store.
            event_id (str): event id to store.
            parent_id (str): parent event id.

        """
        return ""

    def on_event_end(
        self,
        event_type: CBEventType,
        payload: Optional[Dict[str, Any]] = None,
        event_id: str = "",
        **kwargs: Any,
    ) -> None:
        """
        Args:
            event_type (CBEventType): event type to store.
            payload (Optional[Dict[str, Any]]): payload to store.
            event_id (str): event id to store.

        """
        if not self._run:
            raise ValueError("AimCallback failed to init properly.")

        if event_type is CBEventType.LLM and payload:
            if EventPayload.PROMPT in payload:
                llm_input = str(payload[EventPayload.PROMPT])
                llm_output = str(payload[EventPayload.COMPLETION])
            else:
                message = payload.get(EventPayload.MESSAGES, [])
                llm_input = "\n".join([str(x) for x in message])
                llm_output = str(payload[EventPayload.RESPONSE])

            self._run.track(
                Text(llm_input),
                name="prompt",
                step=self._llm_response_step,
                context={"event_id": event_id},
            )

            self._run.track(
                Text(llm_output),
                name="response",
                step=self._llm_response_step,
                context={"event_id": event_id},
            )

            self._llm_response_step += 1
        elif event_type is CBEventType.CHUNKING and payload:
            for chunk_id, chunk in enumerate(payload[EventPayload.CHUNKS]):
                self._run.track(
                    Text(chunk),
                    name="chunk",
                    step=self._llm_response_step,
                    context={"chunk_id": chunk_id, "event_id": event_id},
                )

    @property
    def experiment(self) -> Run:
        if not self._run:
            self.setup()
        return self._run

    def setup(self, args: Optional[Dict[str, Any]] = None) -> None:
        if not self._run:
            if self._run_hash:
                self._run = Run(
                    self._run_hash,
                    repo=self.repo,
                    system_tracking_interval=self.system_tracking_interval,
                    log_system_params=self.log_system_params,
                    capture_terminal_logs=self.capture_terminal_logs,
                )
            else:
                self._run = Run(
                    repo=self.repo,
                    experiment=self.experiment_name,
                    system_tracking_interval=self.system_tracking_interval,
                    log_system_params=self.log_system_params,
                    capture_terminal_logs=self.capture_terminal_logs,
                )
                self._run_hash = self._run.hash

        # Log config parameters
        if args:
            try:
                for key in args:
                    self._run.set(key, args[key], strict=False)
            except Exception as e:
                logger.warning(f"Aim could not log config parameters -> {e}")

    def __del__(self) -> None:
        if self._run and self._run.active:
            self._run.close()

    def start_trace(self, trace_id: Optional[str] = None) -> None:
        pass

    def end_trace(
        self,
        trace_id: Optional[str] = None,
        trace_map: Optional[Dict[str, List[str]]] = None,
    ) -> None:
        pass

on_event_start

on_event_start(event_type: CBEventType, payload: Optional[Dict[str, Any]] = None, event_id: str = '', parent_id: str = '', **kwargs: Any) -> str

参数

名称	类型	描述	默认值
event_type	CBEventType	要存储的事件类型。	必需
payload	可选[字典[str, Any]]	要存储的 payload。	无
event_id	str	要存储的事件 ID。	''
parent_id	str	父事件 ID。	''

源代码位于 llama-index-integrations/callbacks/llama-index-callbacks-aim/llama_index/callbacks/aim/base.py

def on_event_start(
    self,
    event_type: CBEventType,
    payload: Optional[Dict[str, Any]] = None,
    event_id: str = "",
    parent_id: str = "",
    **kwargs: Any,
) -> str:
    """
    Args:
        event_type (CBEventType): event type to store.
        payload (Optional[Dict[str, Any]]): payload to store.
        event_id (str): event id to store.
        parent_id (str): parent event id.

    """
    return ""

on_event_end

on_event_end(event_type: CBEventType, payload: Optional[Dict[str, Any]] = None, event_id: str = '', **kwargs: Any) -> None

参数

名称	类型	描述	默认值
event_type	CBEventType	要存储的事件类型。	必需
payload	可选[字典[str, Any]]	要存储的 payload。	无
event_id	str	要存储的事件 ID。	''

源代码位于 llama-index-integrations/callbacks/llama-index-callbacks-aim/llama_index/callbacks/aim/base.py

def on_event_end(
    self,
    event_type: CBEventType,
    payload: Optional[Dict[str, Any]] = None,
    event_id: str = "",
    **kwargs: Any,
) -> None:
    """
    Args:
        event_type (CBEventType): event type to store.
        payload (Optional[Dict[str, Any]]): payload to store.
        event_id (str): event id to store.

    """
    if not self._run:
        raise ValueError("AimCallback failed to init properly.")

    if event_type is CBEventType.LLM and payload:
        if EventPayload.PROMPT in payload:
            llm_input = str(payload[EventPayload.PROMPT])
            llm_output = str(payload[EventPayload.COMPLETION])
        else:
            message = payload.get(EventPayload.MESSAGES, [])
            llm_input = "\n".join([str(x) for x in message])
            llm_output = str(payload[EventPayload.RESPONSE])

        self._run.track(
            Text(llm_input),
            name="prompt",
            step=self._llm_response_step,
            context={"event_id": event_id},
        )

        self._run.track(
            Text(llm_output),
            name="response",
            step=self._llm_response_step,
            context={"event_id": event_id},
        )

        self._llm_response_step += 1
    elif event_type is CBEventType.CHUNKING and payload:
        for chunk_id, chunk in enumerate(payload[EventPayload.CHUNKS]):
            self._run.track(
                Text(chunk),
                name="chunk",
                step=self._llm_response_step,
                context={"chunk_id": chunk_id, "event_id": event_id},
            )