Wandb

WandbCallbackHandler

Bases: BaseCallbackHandler

回调处理程序，用于将事件记录到wandb。

注意：这是一个测试版功能。我们代码库中的使用方式和接口可能会发生变化。

使用WandbCallbackHandler来将跟踪事件记录到wandb。这个处理程序对于调试和可视化跟踪事件非常有用。它捕获事件的有效负载并将其记录到wandb。该处理程序还跟踪事件的开始和结束。这对于调试LLM调用特别有用。

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

Source code in llama_index/callbacks/wandb/base.py

class WandbCallbackHandler(BaseCallbackHandler):
    """回调处理程序，用于将事件记录到wandb。

    注意：这是一个测试版功能。我们代码库中的使用方式和接口可能会发生变化。

    使用`WandbCallbackHandler`来将跟踪事件记录到wandb。这个处理程序对于调试和可视化跟踪事件非常有用。它捕获事件的有效负载并将其记录到wandb。该处理程序还跟踪事件的开始和结束。这对于调试LLM调用特别有用。

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

    Args:
        event_starts_to_ignore（Optional[List[CBEventType]]）：要忽略跟踪事件开始时的事件类型列表。
        event_ends_to_ignore（Optional[List[CBEventType]]）：要忽略跟踪事件结束时的事件类型列表。"""

    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:
        """按事件类型存储事件开始数据。

Args:
    event_type (CBEventType): 要存储的事件类型。
    payload (Optional[Dict[str, Any]]): 要存储的有效负载。
    event_id (str): 要存储的事件ID。
    parent_id (str): 父事件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:
        """按事件类型存储事件结束数据。

Args:
    event_type (CBEventType): 要存储的事件类型。
    payload (Optional[Dict[str, Any]]): 要存储的有效负载。
    event_id (str): 要存储的事件ID。
"""
        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:
        """启动一个跟踪。"""
        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:
        """将跟踪树记录到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:
        """将索引上传到wandb作为一个artifact。您可以在这里了解更多关于W&B artifacts的信息：https://docs.wandb.ai/guides/artifacts。

对于`ComposableGraph`索引，根id被存储为artifact元数据。

Args:
    index (IndexType): 要上传的索引。
    index_name (str): 索引的名称。这将被用作artifact的名称。
    persist_dir (Union[str, None]): 持久化索引的目录。如果为None，将创建并使用临时目录。
"""
        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":
        """从wandb下载索引并返回存储上下文。

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

Args:
    artifact_url（str）：要下载的artifact的url。artifact的url的格式为：`entity/project/index_name:version`，可以在W&B UI中找到。
    index_download_dir（Union[str, None]）：要下载索引的目录。
"""
        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:
        """实用函数，将一个目录上传到W&B作为一个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":
        """从跟踪地图构建跟踪树。"""
        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":
        """将一对事件转换为一个wandb跟踪树跨度。"""
        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"]:
        """将CBEventType映射到wandb跟踪树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"]:
        """将事件的有效负载添加到跨度中。"""
        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]]:
        """处理NODE_PARSING事件的有效负载。"""
        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"]:
        """处理LLM事件的有效负载。"""
        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]]:
        """处理QUERY事件的有效负载。"""
        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]:
        """获取事件对的开始和结束时间，单位为毫秒。"""
        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:
        """确保存在一个活跃的W&B运行。

如果不存在，则将使用提供的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:
        """当W&B初始化时，在终端打印一条消息。"""
        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:
        """当上传索引失败时，在终端打印一条消息。"""
        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:
        """完成回调处理程序。"""
        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

按事件类型存储事件开始数据。

Parameters:

Name	Type	Description	Default
event_type	CBEventType	要存储的事件类型。	required
payload	Optional[Dict[str, Any]]	要存储的有效负载。	None
event_id	str	要存储的事件ID。	''
parent_id	str	父事件ID。	''

Source code in 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:
        """按事件类型存储事件开始数据。

Args:
    event_type (CBEventType): 要存储的事件类型。
    payload (Optional[Dict[str, Any]]): 要存储的有效负载。
    event_id (str): 要存储的事件ID。
    parent_id (str): 父事件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

按事件类型存储事件结束数据。

Parameters:

Name	Type	Description	Default
event_type	CBEventType	要存储的事件类型。	required
payload	Optional[Dict[str, Any]]	要存储的有效负载。	None
event_id	str	要存储的事件ID。	''

Source code in 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:
        """按事件类型存储事件结束数据。

Args:
    event_type (CBEventType): 要存储的事件类型。
    payload (Optional[Dict[str, Any]]): 要存储的有效负载。
    event_id (str): 要存储的事件ID。
"""
        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

启动一个跟踪。

Source code in llama_index/callbacks/wandb/base.py

def start_trace(self, trace_id: Optional[str] = None) -> None:
    """启动一个跟踪。"""
    self._trace_map = defaultdict(list)
    self._cur_trace_id = trace_id
    self._start_time = datetime.now()

log_trace_tree

log_trace_tree() -> None

将跟踪树记录到wandb。

Source code in llama_index/callbacks/wandb/base.py

def log_trace_tree(self) -> None:
    """将跟踪树记录到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

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

对于ComposableGraph索引，根id被存储为artifact元数据。

Parameters:

Name	Type	Description	Default
index	IndexType	要上传的索引。	required
index_name	str	索引的名称。这将被用作artifact的名称。	required
persist_dir	Union[str, None]	持久化索引的目录。如果为None，将创建并使用临时目录。	None

Source code in llama_index/callbacks/wandb/base.py

    def persist_index(
        self, index: "IndexType", index_name: str, persist_dir: Union[str, None] = None
    ) -> None:
        """将索引上传到wandb作为一个artifact。您可以在这里了解更多关于W&B artifacts的信息：https://docs.wandb.ai/guides/artifacts。

对于`ComposableGraph`索引，根id被存储为artifact元数据。

Args:
    index (IndexType): 要上传的索引。
    index_name (str): 索引的名称。这将被用作artifact的名称。
    persist_dir (Union[str, None]): 持久化索引的目录。如果为None，将创建并使用临时目录。
"""
        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函数将索引加载到内存中。

Parameters:

Name	Type	Description	Default
artifact_url（str）：要下载的artifact的url。artifact的url的格式为：`entity/project/index_name		version`，可以在W&B UI中找到。	required

Source code in llama_index/callbacks/wandb/base.py

    def load_storage_context(
        self, artifact_url: str, index_download_dir: Union[str, None] = None
    ) -> "StorageContext":
        """从wandb下载索引并返回存储上下文。

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

Args:
    artifact_url（str）：要下载的artifact的url。artifact的url的格式为：`entity/project/index_name:version`，可以在W&B UI中找到。
    index_download_dir（Union[str, None]）：要下载索引的目录。
"""
        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

完成回调处理程序。

Source code in llama_index/callbacks/wandb/base.py

def finish(self) -> None:
    """完成回调处理程序。"""
    self._wandb.finish()  # type: ignore[attr-defined]