[Bug]: Aim Callback does not reflex index properly

[olyashok]

Bug Description

When texts are logged into AimCallback, they are hard to browse because they are not connected to each other in multi-step sequences. It would be good to log parent_id and also use coirrect step within parent to log it. Here is the code that adds it.

Version

0.9.21

Steps to Reproduce

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._llm_response_sub_step = 0

    self.setup(run_params)

    self.events_hierarchy = {}
    self.step_counter = {}

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.
    """
    self.events_hierarchy[event_id] = parent_id
    self.step_counter[event_id] = 0
    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])

        parent_id = "none"
        grand_parent_id = "none"
        if event_id in self.events_hierarchy:
            parent_id = self.events_hierarchy[event_id]
        if parent_id in self.events_hierarchy:
            grand_parent_id = self.events_hierarchy[parent_id]

        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, "parent_id": parent_id, "grand_parent_id": grand_parent_id},
        )

        self._run.track(
            [Text(llm_input),Text(llm_output)],
            name="all",
            step=self._llm_response_step,
            context={"event_id": event_id},
        )

        self._run.track(
            [Text(llm_input),Text(llm_output)],
            name="all_parent",
            step=self.step_counter[parent_id],
            context={"parent_id": parent_id},
        )
        self.step_counter[parent_id] = self.step_counter[parent_id] + 1

        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},
            )
    self.events_hierarchy.pop(event_id)
    self.step_counter.pop(event_id)

Relevant Logs/Tracbacks

No response

[dosubot[bot]]

🤖

It's great to see that you've identified a potential improvement for the AimCallback. Adding the parent_id and using the correct step within the parent to log the texts seems like a valuable enhancement. Your proposed code changes look thorough and well-considered. I encourage you to go ahead and create a pull request with your proposed solution. Your contribution will be appreciated by the LlamaIndex community. Keep up the good work!

---

This response is meant to be useful and save you time. It is not meant to be a precise solution, but rather a starting point for your own research.

Help me be more useful! Please leave a 👍 if this is helpful and 👎 if it is irrelevant.

If you want to continue the conversation, start your reply with @dosu-bot.