Wandb

WandbCallbackHandler

基类：BaseCallbackHandler

记录事件到 wandb 的回调处理程序。

注意：这是一个测试版功能。其在代码库中的用法和接口可能会改变。

使用 WandbCallbackHandler 将追踪事件记录到 wandb。此处理程序有助于调试和可视化追踪事件。它捕获事件的有效载荷并将其记录到 wandb。此处理程序还会追踪事件的开始和结束。这对于调试您的 LLM 调用特别有用。

WandbCallbackHandler 也可用于使用 persist_index 方法将索引和图记录到 wandb。这将把索引作为 artifacts 保存在 wandb 中。可以使用 load_storage_context 方法从 wandb artifacts 加载索引。此方法将返回一个 StorageContext 对象，可用于使用 load_index_from_storage、load_indices_from_storage 或 load_graph_from_storage 函数构建索引。

参数

名称	类型	描述	默认值
event_starts_to_ignore	可选[列表[CBEventType]]	追踪事件开始时忽略的事件类型列表。	无
event_ends_to_ignore	可选[列表[CBEventType]]	追踪事件结束时忽略的事件类型列表。	无

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

class WandbCallbackHandler(BaseCallbackHandler):
    """
    Callback handler that logs events to wandb.

    NOTE: this is a beta feature. The usage within our codebase, and the interface
    may change.

    Use the `WandbCallbackHandler` to log trace events to wandb. This handler is
    useful for debugging and visualizing the trace events. It captures the payload of
    the events and logs them to wandb. The handler also tracks the start and end of
    events. This is particularly useful for debugging your LLM calls.

    The `WandbCallbackHandler` can also be used to log the indices and graphs to wandb
    using the `persist_index` method. This will save the indexes as artifacts in wandb.
    The `load_storage_context` method can be used to load the indexes from wandb
    artifacts. This method will return a `StorageContext` object that can be used to
    build the index, using `load_index_from_storage`, `load_indices_from_storage` or
    `load_graph_from_storage` functions.


    Args:
        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,
        run_args: Optional[WandbRunArgs] = None,
        tokenizer: Optional[Callable[[str], List]] = None,
        event_starts_to_ignore: Optional[List[CBEventType]] = None,
        event_ends_to_ignore: Optional[List[CBEventType]] = None,
    ) -> None:
        try:
            import wandb
            from wandb.sdk.data_types import trace_tree

            self._wandb = wandb
            self._trace_tree = trace_tree
        except ImportError:
            raise ImportError(
                "WandbCallbackHandler requires wandb. "
                "Please install it with `pip install wandb`."
            )

        from llama_index.core.indices import (
            ComposableGraph,
            GPTEmptyIndex,
            GPTKeywordTableIndex,
            GPTRAKEKeywordTableIndex,
            GPTSimpleKeywordTableIndex,
            GPTSQLStructStoreIndex,
            GPTTreeIndex,
            GPTVectorStoreIndex,
            SummaryIndex,
        )

        self._IndexType = (
            ComposableGraph,
            GPTKeywordTableIndex,
            GPTSimpleKeywordTableIndex,
            GPTRAKEKeywordTableIndex,
            SummaryIndex,
            GPTEmptyIndex,
            GPTTreeIndex,
            GPTVectorStoreIndex,
            GPTSQLStructStoreIndex,
        )

        self._run_args = run_args
        # Check if a W&B run is already initialized; if not, initialize one
        self._ensure_run(should_print_url=(self._wandb.run is None))  # type: ignore[attr-defined]

        self._event_pairs_by_id: Dict[str, List[CBEvent]] = defaultdict(list)
        self._cur_trace_id: Optional[str] = None
        self._trace_map: Dict[str, List[str]] = defaultdict(list)

        self.tokenizer = tokenizer or get_tokenizer()
        self._token_counter = TokenCounter(tokenizer=self.tokenizer)

        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,
        )

    def on_event_start(
        self,
        event_type: CBEventType,
        payload: Optional[Dict[str, Any]] = None,
        event_id: str = "",
        parent_id: str = "",
        **kwargs: Any,
    ) -> str:
        """
        Store event start data by event type.

        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.

        """
        event = CBEvent(event_type, payload=payload, id_=event_id)
        self._event_pairs_by_id[event.id_].append(event)
        return event.id_

    def on_event_end(
        self,
        event_type: CBEventType,
        payload: Optional[Dict[str, Any]] = None,
        event_id: str = "",
        **kwargs: Any,
    ) -> None:
        """
        Store event end data by event type.

        Args:
            event_type (CBEventType): event type to store.
            payload (Optional[Dict[str, Any]]): payload to store.
            event_id (str): event id to store.

        """
        event = CBEvent(event_type, payload=payload, id_=event_id)
        self._event_pairs_by_id[event.id_].append(event)
        self._trace_map = defaultdict(list)

    def start_trace(self, trace_id: Optional[str] = None) -> None:
        """Launch a trace."""
        self._trace_map = defaultdict(list)
        self._cur_trace_id = trace_id
        self._start_time = datetime.now()

    def end_trace(
        self,
        trace_id: Optional[str] = None,
        trace_map: Optional[Dict[str, List[str]]] = None,
    ) -> None:
        # Ensure W&B run is initialized
        self._ensure_run()

        self._trace_map = trace_map or defaultdict(list)
        self._end_time = datetime.now()

        # Log the trace map to wandb
        # We can control what trace ids we want to log here.
        self.log_trace_tree()

        # TODO (ayulockin): Log the LLM token counts to wandb when weave is ready

    def log_trace_tree(self) -> None:
        """Log the trace tree to wandb."""
        try:
            child_nodes = self._trace_map["root"]
            root_span = self._convert_event_pair_to_wb_span(
                self._event_pairs_by_id[child_nodes[0]],
                trace_id=self._cur_trace_id if len(child_nodes) > 1 else None,
            )

            if len(child_nodes) == 1:
                child_nodes = self._trace_map[child_nodes[0]]
                root_span = self._build_trace_tree(child_nodes, root_span)
            else:
                root_span = self._build_trace_tree(child_nodes, root_span)
            if root_span:
                root_trace = self._trace_tree.WBTraceTree(root_span)
                if self._wandb.run:  # type: ignore[attr-defined]
                    self._wandb.run.log({"trace": root_trace})  # type: ignore[attr-defined]
                self._wandb.termlog("Logged trace tree to W&B.")  # type: ignore[attr-defined]
        except Exception as e:
            print(f"Failed to log trace tree to W&B: {e}")
            # ignore errors to not break user code

    def persist_index(
        self, index: "IndexType", index_name: str, persist_dir: Union[str, None] = None
    ) -> None:
        """
        Upload an index to wandb as an artifact. You can learn more about W&B
        artifacts here: https://docs.wandb.ai/guides/artifacts.

        For the `ComposableGraph` index, the root id is stored as artifact metadata.

        Args:
            index (IndexType): index to upload.
            index_name (str): name of the index. This will be used as the artifact name.
            persist_dir (Union[str, None]): directory to persist the index. If None, a
                temporary directory will be created and used.

        """
        if persist_dir is None:
            persist_dir = f"{self._wandb.run.dir}/storage"  # type: ignore
            _default_persist_dir = True
        if not os.path.exists(persist_dir):
            os.makedirs(persist_dir)

        if isinstance(index, self._IndexType):
            try:
                index.storage_context.persist(persist_dir)  # type: ignore

                metadata = None
                # For the `ComposableGraph` index, store the root id as metadata
                if isinstance(index, self._IndexType[0]):
                    root_id = index.root_id
                    metadata = {"root_id": root_id}

                self._upload_index_as_wb_artifact(persist_dir, index_name, metadata)
            except Exception as e:
                # Silently ignore errors to not break user code
                self._print_upload_index_fail_message(e)

        # clear the default storage dir
        if _default_persist_dir:
            shutil.rmtree(persist_dir, ignore_errors=True)

    def load_storage_context(
        self, artifact_url: str, index_download_dir: Union[str, None] = None
    ) -> "StorageContext":
        """
        Download an index from wandb and return a storage context.

        Use this storage context to load the index into memory using
        `load_index_from_storage`, `load_indices_from_storage` or
        `load_graph_from_storage` functions.

        Args:
            artifact_url (str): url of the artifact to download. The artifact url will
                be of the form: `entity/project/index_name:version` and can be found in
                the W&B UI.
            index_download_dir (Union[str, None]): directory to download the index to.

        """
        from llama_index.core.storage.storage_context import StorageContext

        artifact = self._wandb.use_artifact(artifact_url, type="storage_context")  # type: ignore[attr-defined]
        artifact_dir = artifact.download(root=index_download_dir)

        return StorageContext.from_defaults(persist_dir=artifact_dir)

    def _upload_index_as_wb_artifact(
        self, dir_path: str, artifact_name: str, metadata: Optional[Dict]
    ) -> None:
        """Utility function to upload a dir to W&B as an artifact."""
        artifact = self._wandb.Artifact(artifact_name, type="storage_context")  # type: ignore[attr-defined]

        if metadata:
            artifact.metadata = metadata

        artifact.add_dir(dir_path)
        self._wandb.run.log_artifact(artifact)  # type: ignore

    def _build_trace_tree(
        self, events: List[str], span: "trace_tree.Span"
    ) -> "trace_tree.Span":
        """Build the trace tree from the trace map."""
        for child_event in events:
            child_span = self._convert_event_pair_to_wb_span(
                self._event_pairs_by_id[child_event]
            )
            child_span = self._build_trace_tree(
                self._trace_map[child_event], child_span
            )
            span.add_child_span(child_span)

        return span

    def _convert_event_pair_to_wb_span(
        self,
        event_pair: List[CBEvent],
        trace_id: Optional[str] = None,
    ) -> "trace_tree.Span":
        """Convert a pair of events to a wandb trace tree span."""
        start_time_ms, end_time_ms = self._get_time_in_ms(event_pair)

        if trace_id is None:
            event_type = event_pair[0].event_type
            span_kind = self._map_event_type_to_span_kind(event_type)
        else:
            event_type = trace_id  # type: ignore
            span_kind = None

        wb_span = self._trace_tree.Span(
            name=f"{event_type}",
            span_kind=span_kind,
            start_time_ms=start_time_ms,
            end_time_ms=end_time_ms,
        )

        inputs, outputs, wb_span = self._add_payload_to_span(wb_span, event_pair)
        wb_span.add_named_result(inputs=inputs, outputs=outputs)  # type: ignore

        return wb_span

    def _map_event_type_to_span_kind(
        self, event_type: CBEventType
    ) -> Union[None, "trace_tree.SpanKind"]:
        """Map a CBEventType to a wandb trace tree SpanKind."""
        if event_type == CBEventType.CHUNKING:
            span_kind = None
        elif event_type == CBEventType.NODE_PARSING:
            span_kind = None
        elif event_type == CBEventType.EMBEDDING:
            # TODO: add span kind for EMBEDDING when it's available
            span_kind = None
        elif event_type == CBEventType.LLM:
            span_kind = self._trace_tree.SpanKind.LLM
        elif event_type == CBEventType.QUERY:
            span_kind = self._trace_tree.SpanKind.AGENT
        elif event_type == CBEventType.AGENT_STEP:
            span_kind = self._trace_tree.SpanKind.AGENT
        elif event_type == CBEventType.RETRIEVE:
            span_kind = self._trace_tree.SpanKind.TOOL
        elif event_type == CBEventType.SYNTHESIZE:
            span_kind = self._trace_tree.SpanKind.CHAIN
        elif event_type == CBEventType.TREE:
            span_kind = self._trace_tree.SpanKind.CHAIN
        elif event_type == CBEventType.SUB_QUESTION:
            span_kind = self._trace_tree.SpanKind.CHAIN
        elif event_type == CBEventType.RERANKING:
            span_kind = self._trace_tree.SpanKind.CHAIN
        elif event_type == CBEventType.FUNCTION_CALL:
            span_kind = self._trace_tree.SpanKind.TOOL
        else:
            span_kind = None

        return span_kind

    def _add_payload_to_span(
        self, span: "trace_tree.Span", event_pair: List[CBEvent]
    ) -> Tuple[Optional[Dict[str, Any]], Optional[Dict[str, Any]], "trace_tree.Span"]:
        """Add the event's payload to the span."""
        assert len(event_pair) == 2
        event_type = event_pair[0].event_type
        inputs = None
        outputs = None

        if event_type == CBEventType.NODE_PARSING:
            # TODO: disabled full detailed inputs/outputs due to UI lag
            inputs, outputs = self._handle_node_parsing_payload(event_pair)
        elif event_type == CBEventType.LLM:
            inputs, outputs, span = self._handle_llm_payload(event_pair, span)
        elif event_type == CBEventType.QUERY:
            inputs, outputs = self._handle_query_payload(event_pair)
        elif event_type == CBEventType.EMBEDDING:
            inputs, outputs = self._handle_embedding_payload(event_pair)

        return inputs, outputs, span

    def _handle_node_parsing_payload(
        self, event_pair: List[CBEvent]
    ) -> Tuple[Dict[str, Any], Dict[str, Any]]:
        """Handle the payload of a NODE_PARSING event."""
        inputs = event_pair[0].payload
        outputs = event_pair[-1].payload

        if inputs and EventPayload.DOCUMENTS in inputs:
            documents = inputs.pop(EventPayload.DOCUMENTS)
            inputs["num_documents"] = len(documents)

        if outputs and EventPayload.NODES in outputs:
            nodes = outputs.pop(EventPayload.NODES)
            outputs["num_nodes"] = len(nodes)

        return inputs or {}, outputs or {}

    def _handle_llm_payload(
        self, event_pair: List[CBEvent], span: "trace_tree.Span"
    ) -> Tuple[Dict[str, Any], Dict[str, Any], "trace_tree.Span"]:
        """Handle the payload of a LLM event."""
        inputs = event_pair[0].payload
        outputs = event_pair[-1].payload

        assert isinstance(inputs, dict) and isinstance(outputs, dict)

        # Get `original_template` from Prompt
        if EventPayload.PROMPT in inputs:
            inputs[EventPayload.PROMPT] = inputs[EventPayload.PROMPT]

        # Format messages
        if EventPayload.MESSAGES in inputs:
            inputs[EventPayload.MESSAGES] = "\n".join(
                [str(x) for x in inputs[EventPayload.MESSAGES]]
            )

        token_counts = get_llm_token_counts(self._token_counter, outputs)
        metadata = {
            "formatted_prompt_tokens_count": token_counts.prompt_token_count,
            "prediction_tokens_count": token_counts.completion_token_count,
            "total_tokens_used": token_counts.total_token_count,
        }
        span.attributes = metadata

        # Make `response` part of `outputs`
        outputs = {EventPayload.RESPONSE: str(outputs[EventPayload.RESPONSE])}

        return inputs, outputs, span

    def _handle_query_payload(
        self, event_pair: List[CBEvent]
    ) -> Tuple[Optional[Dict[str, Any]], Dict[str, Any]]:
        """Handle the payload of a QUERY event."""
        inputs = event_pair[0].payload
        outputs = event_pair[-1].payload

        if outputs:
            response_obj = outputs[EventPayload.RESPONSE]
            response = str(outputs[EventPayload.RESPONSE])

            if type(response).__name__ == "Response":
                response = response_obj.response
            elif type(response).__name__ == "StreamingResponse":
                response = response_obj.get_response().response
        else:
            response = " "

        outputs = {"response": response}

        return inputs, outputs

    def _handle_embedding_payload(
        self,
        event_pair: List[CBEvent],
    ) -> Tuple[Optional[Dict[str, Any]], Dict[str, Any]]:
        event_pair[0].payload
        outputs = event_pair[-1].payload

        chunks = []
        if outputs:
            chunks = outputs.get(EventPayload.CHUNKS, [])

        return {}, {"num_chunks": len(chunks)}

    def _get_time_in_ms(self, event_pair: List[CBEvent]) -> Tuple[int, int]:
        """Get the start and end time of an event pair in milliseconds."""
        start_time = datetime.strptime(event_pair[0].time, TIMESTAMP_FORMAT)
        end_time = datetime.strptime(event_pair[1].time, TIMESTAMP_FORMAT)

        start_time_in_ms = int(
            (start_time - datetime(1970, 1, 1)).total_seconds() * 1000
        )
        end_time_in_ms = int((end_time - datetime(1970, 1, 1)).total_seconds() * 1000)

        return start_time_in_ms, end_time_in_ms

    def _ensure_run(self, should_print_url: bool = False) -> None:
        """
        Ensures an active W&B run exists.

        If not, will start a new run with the provided run_args.
        """
        if self._wandb.run is None:  # type: ignore[attr-defined]
            # Make a shallow copy of the run args, so we don't modify the original
            run_args = self._run_args or {}  # type: ignore
            run_args: dict = {**run_args}  # type: ignore

            # Prefer to run in silent mode since W&B has a lot of output
            # which can be undesirable when dealing with text-based models.
            if "settings" not in run_args:  # type: ignore
                run_args["settings"] = {"silent": True}  # type: ignore

            # Start the run and add the stream table
            self._wandb.init(**run_args)  # type: ignore[attr-defined]
            self._wandb.run._label(repo="llama_index")  # type: ignore

            if should_print_url:
                self._print_wandb_init_message(
                    self._wandb.run.settings.run_url  # type: ignore
                )

    def _print_wandb_init_message(self, run_url: str) -> None:
        """Print a message to the terminal when W&B is initialized."""
        self._wandb.termlog(  # type: ignore[attr-defined]
            f"Streaming LlamaIndex events to W&B at {run_url}\n"
            "`WandbCallbackHandler` is currently in beta.\n"
            "Please report any issues to https://github.com/wandb/wandb/issues "
            "with the tag `llamaindex`."
        )

    def _print_upload_index_fail_message(self, e: Exception) -> None:
        """Print a message to the terminal when uploading the index fails."""
        self._wandb.termlog(  # type: ignore[attr-defined]
            f"Failed to upload index to W&B with the following error: {e}\n"
        )

    def finish(self) -> None:
        """Finish the callback handler."""
        self._wandb.finish()  # type: ignore[attr-defined]

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	可选[字典[字符串, 任意]]	要存储的有效载荷。	无
event_id	字符串	要存储的事件 ID。	''
parent_id	字符串	父事件 ID。	''

源代码位于 llama-index-integrations/callbacks/llama-index-callbacks-wandb/llama_index/callbacks/wandb/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:
    """
    Store event start data by event type.

    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.

    """
    event = CBEvent(event_type, payload=payload, id_=event_id)
    self._event_pairs_by_id[event.id_].append(event)
    return event.id_

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	可选[字典[字符串, 任意]]	要存储的有效载荷。	无
event_id	字符串	要存储的事件 ID。	''

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

def on_event_end(
    self,
    event_type: CBEventType,
    payload: Optional[Dict[str, Any]] = None,
    event_id: str = "",
    **kwargs: Any,
) -> None:
    """
    Store event end data by event type.

    Args:
        event_type (CBEventType): event type to store.
        payload (Optional[Dict[str, Any]]): payload to store.
        event_id (str): event id to store.

    """
    event = CBEvent(event_type, payload=payload, id_=event_id)
    self._event_pairs_by_id[event.id_].append(event)
    self._trace_map = defaultdict(list)

start_trace

start_trace(trace_id: Optional[str] = None) -> None

启动追踪。

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

def start_trace(self, trace_id: Optional[str] = None) -> None:
    """Launch a trace."""
    self._trace_map = defaultdict(list)
    self._cur_trace_id = trace_id
    self._start_time = datetime.now()

log_trace_tree

log_trace_tree() -> None

将追踪树记录到 wandb。

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

def log_trace_tree(self) -> None:
    """Log the trace tree to wandb."""
    try:
        child_nodes = self._trace_map["root"]
        root_span = self._convert_event_pair_to_wb_span(
            self._event_pairs_by_id[child_nodes[0]],
            trace_id=self._cur_trace_id if len(child_nodes) > 1 else None,
        )

        if len(child_nodes) == 1:
            child_nodes = self._trace_map[child_nodes[0]]
            root_span = self._build_trace_tree(child_nodes, root_span)
        else:
            root_span = self._build_trace_tree(child_nodes, root_span)
        if root_span:
            root_trace = self._trace_tree.WBTraceTree(root_span)
            if self._wandb.run:  # type: ignore[attr-defined]
                self._wandb.run.log({"trace": root_trace})  # type: ignore[attr-defined]
            self._wandb.termlog("Logged trace tree to W&B.")  # type: ignore[attr-defined]
    except Exception as e:
        print(f"Failed to log trace tree to W&B: {e}")

persist_index

persist_index(index: IndexType, index_name: str, persist_dir: Union[str, None] = None) -> None

将索引作为 artifact 上传到 wandb。您可以在此处了解有关 W&B artifacts 的更多信息：https://docs.wandb.ai/guides/artifacts。

对于 ComposableGraph 索引，根 ID 作为 artifact 元数据存储。

参数

名称	类型	描述	默认值
index	IndexType	要上传的索引。	必需的
index_name	字符串	索引的名称。这将用作 artifact 名称。	必需的
persist_dir	联合[字符串, None]	持久化索引的目录。如果为 None，将创建一个临时目录并使用。	无

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

def persist_index(
    self, index: "IndexType", index_name: str, persist_dir: Union[str, None] = None
) -> None:
    """
    Upload an index to wandb as an artifact. You can learn more about W&B
    artifacts here: https://docs.wandb.ai/guides/artifacts.

    For the `ComposableGraph` index, the root id is stored as artifact metadata.

    Args:
        index (IndexType): index to upload.
        index_name (str): name of the index. This will be used as the artifact name.
        persist_dir (Union[str, None]): directory to persist the index. If None, a
            temporary directory will be created and used.

    """
    if persist_dir is None:
        persist_dir = f"{self._wandb.run.dir}/storage"  # type: ignore
        _default_persist_dir = True
    if not os.path.exists(persist_dir):
        os.makedirs(persist_dir)

    if isinstance(index, self._IndexType):
        try:
            index.storage_context.persist(persist_dir)  # type: ignore

            metadata = None
            # For the `ComposableGraph` index, store the root id as metadata
            if isinstance(index, self._IndexType[0]):
                root_id = index.root_id
                metadata = {"root_id": root_id}

            self._upload_index_as_wb_artifact(persist_dir, index_name, metadata)
        except Exception as e:
            # Silently ignore errors to not break user code
            self._print_upload_index_fail_message(e)

    # clear the default storage dir
    if _default_persist_dir:
        shutil.rmtree(persist_dir, ignore_errors=True)

load_storage_context

load_storage_context(artifact_url: str, index_download_dir: Union[str, None] = None) -> StorageContext

从 wandb 下载索引并返回存储上下文。

使用此存储上下文，通过 load_index_from_storage、load_indices_from_storage 或 load_graph_from_storage 函数将索引加载到内存中。

参数

名称	类型	描述	默认值
artifact_url	字符串	要下载的 artifact 的 URL。artifact URL 的格式为：entity/project/index_name:version，可在 W&B UI 中找到。	必需的
index_download_dir	联合[字符串, None]	下载索引到的目录。	无

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

def load_storage_context(
    self, artifact_url: str, index_download_dir: Union[str, None] = None
) -> "StorageContext":
    """
    Download an index from wandb and return a storage context.

    Use this storage context to load the index into memory using
    `load_index_from_storage`, `load_indices_from_storage` or
    `load_graph_from_storage` functions.

    Args:
        artifact_url (str): url of the artifact to download. The artifact url will
            be of the form: `entity/project/index_name:version` and can be found in
            the W&B UI.
        index_download_dir (Union[str, None]): directory to download the index to.

    """
    from llama_index.core.storage.storage_context import StorageContext

    artifact = self._wandb.use_artifact(artifact_url, type="storage_context")  # type: ignore[attr-defined]
    artifact_dir = artifact.download(root=index_download_dir)

    return StorageContext.from_defaults(persist_dir=artifact_dir)

finish

finish() -> None

完成回调处理程序。

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

def finish(self) -> None:
    """Finish the callback handler."""
    self._wandb.finish()  # type: ignore[attr-defined]