def apply_audio_delay(
    audio_BxTxC: torch.Tensor,
    pad_value: int,
    bos_value: int,
    precomp: tp.Tuple[torch.Tensor, torch.Tensor],
) -> torch.Tensor:
    """
    Applies the delay pattern to batched audio tokens using precomputed indices,
    inserting BOS where t_idx < 0 and PAD where t_idx >= T.

    Args:
        audio_BxTxC: [B, T, C] int16 audio tokens (or int32/float)
        pad_value: the padding token
        bos_value: the BOS token
        precomp:  (t_idx_BxTxC, indices_BTCx3) from build_delay_indices

    Returns:
        result_BxTxC: [B, T, C] delayed audio tokens
    """
    device = audio_BxTxC.device  # Get device from input tensor
    t_idx_BxTxC, indices_BTCx3 = precomp
    t_idx_BxTxC = t_idx_BxTxC.to(device)  # Move precomputed indices to device
    indices_BTCx3 = indices_BTCx3.to(device)

    # Equivalent of tf.gather_nd using advanced indexing
    # Ensure indices are long type if not already (build_delay_indices should handle this)
    gathered_flat = audio_BxTxC[indices_BTCx3[:, 0], indices_BTCx3[:, 1], indices_BTCx3[:, 2]]
    gathered_BxTxC = gathered_flat.view(audio_BxTxC.shape)

    # Create masks on the correct device
    mask_bos = t_idx_BxTxC < 0  # => place bos_value
    mask_pad = t_idx_BxTxC >= audio_BxTxC.shape[1]  # => place pad_value

    # Create scalar tensors on the correct device
    bos_tensor = torch.tensor(bos_value, dtype=audio_BxTxC.dtype, device=device)
    pad_tensor = torch.tensor(pad_value, dtype=audio_BxTxC.dtype, device=device)

    # If mask_bos, BOS; else if mask_pad, PAD; else original gather
    # All tensors should now be on the same device
    result_BxTxC = torch.where(mask_bos, bos_tensor, torch.where(mask_pad, pad_tensor, gathered_BxTxC))

    return result_BxTxC


def _CheckNumericsGrad(_, grad):
  """Gradient for check_numerics op."""
  return grad


    def test_1d_slicing5(self, flags=enable_pyobj_flags):
        pyfunc = slicing_1d_usecase5
        self.check_1d_slicing_with_arg(pyfunc, flags)


    def is_user_blocked(user_id: int) -> bool:
        """ Determines if a user is blocked """
        user = UserService.get_user_by_id(user_id)

        if UserRole(user.role) == UserRole.READ_ONLY:
            return True

        return False


    def copy(self):
        return MemSetFn(self._ass)


def make_request(
    url: str, data: Dict[str, Any], stream: bool = False, check_status: bool = True
) -> requests.Response:
    """Send an HTTP request with retry mechanism
    Args:
        url: Request URL
        data: Request data
        stream: Whether to use streaming response
        check_status: Whether to check HTTP status code (default: True)
    Returns:
        requests.Response: Response object

    Raises:
        requests.exceptions.RequestException: Request failed after all retries
        requests.exceptions.HTTPError: HTTP status code is not 200 (when check_status is True)
    """
    server_config = CONFIG["server"]
    max_retries = server_config["max_retries"]
    retry_delay = server_config["retry_delay"]
    timeout = server_config["timeout"]

    for attempt in range(max_retries):
        try:
            response = requests.post(url, json=data, stream=stream, timeout=timeout)
            if check_status and response.status_code != 200:
                response.raise_for_status()
            return response
        except requests.exceptions.RequestException as e:
            if attempt == max_retries - 1:  # Last retry
                raise
            print(f"\nRequest failed, retrying in {retry_delay} seconds: {str(e)}")
            time.sleep(retry_delay)


def handle_message(data):
    logger.info(f"User message: {data}")
    message = data.get('message')
    base_model = data.get('base_model')
    project_name = data.get('project_name')
    search_engine = data.get('search_engine').lower()

    agent = Agent(base_model=base_model, search_engine=search_engine)

    state = AgentState.get_latest_state(project_name)
    if not state:
        thread = Thread(target=lambda: agent.execute(message, project_name))
        thread.start()
    else:
        if AgentState.is_agent_completed(project_name):
            thread = Thread(target=lambda: agent.subsequent_execute(message, project_name))
            thread.start()
        else:
            emit_agent("info", {"type": "warning", "message": "previous agent doesn't completed it's task."})
            last_state = AgentState.get_latest_state(project_name)
            if last_state["agent_is_active"] or not last_state["completed"]:
                thread = Thread(target=lambda: agent.execute(message, project_name))
                thread.start()
            else:
                thread = Thread(target=lambda: agent.subsequent_execute(message, project_name))
                thread.start()


def unload_module():
    global module_in_gpu

    if module_in_gpu is None:
        return

    DynamicSwapInstaller.uninstall_model(module_in_gpu)
    module_in_gpu.to(cpu)
    print(f'Move module to CPU: {type(module_in_gpu).__name__}')

    module_in_gpu = None
    memory_management.soft_empty_cache()
    return


    def register_introspection_functions(self):
        """Registers the XML-RPC introspection methods in the system
        namespace.

        see http://xmlrpc.usefulinc.com/doc/reserved.html
        """

        self.funcs.update({'system.listMethods' : self.system_listMethods,
                      'system.methodSignature' : self.system_methodSignature,
                      'system.methodHelp' : self.system_methodHelp})


async def list_models():
    global model_args
    model_card = ModelCard(id="gpt-3.5-turbo")
    return ModelList(data=[model_card])


def test_docs_examples(example: CodeExample, eval_example: EvalExample):
    ruff_ignore: list[str] = ["F841", "I001", "F821"]  # F821: undefined names (snippets lack imports)

    # Use project's actual line length of 120
    eval_example.set_config(ruff_ignore=ruff_ignore, target_version="py310", line_length=120)

    # Use Ruff for both formatting and linting (skip Black)
    if eval_example.update_examples:  # pragma: no cover
        eval_example.format_ruff(example)
    else:
        eval_example.lint_ruff(example)


    async def scheduler(self, stop_event):
        self.stop_event = stop_event
        while not stop_event.is_set():
            self.now = datetime.datetime.now(pytz.timezone(config.time_zone))
            await self.main()
            next_time = self.now + datetime.timedelta(hours=config.update_interval)
            print(f"🕒 Next update time: {next_time:%Y-%m-%d %H:%M:%S}")
            try:
                await asyncio.wait_for(stop_event.wait(), timeout=config.update_interval * 3600)
            except asyncio.TimeoutError:
                continue


def test_coordinates_as_bytes():
    pubkey = UmbralPrivateKey.gen_key().pubkey
    point = pubkey.point_key
    stamp = SignatureStamp(verifying_key=pubkey)

    x, y = point.to_affine()
    x = x.to_bytes(32, 'big')
    y = y.to_bytes(32, 'big')

    for p in (point, pubkey, stamp):
        assert get_coordinates_as_bytes(p) == x + y
        assert get_coordinates_as_bytes(p, x_coord=False) == y
        assert get_coordinates_as_bytes(p, y_coord=False) == x
        with pytest.raises(ValueError):
            _ = get_coordinates_as_bytes(p, x_coord=False, y_coord=False)


  def testLocalFileDirectory(self, mock_StorageUrlFromString, mock_gsutil_api):
    mock_storage_url = mock.Mock()
    mock_storage_url.isFileUrl.return_value = True
    mock_storage_url.IsDirectory.return_value = True
    mock_StorageUrlFromString.return_value = mock_storage_url
    (exp_url, have_existing_dst_container) = ExpandUrlToSingleBlr(
        '/home/test', mock_gsutil_api, 'project_id', False,
        CreateOrGetGsutilLogger('copy_test'))

    self.assertTrue(have_existing_dst_container)
    self.assertEqual(exp_url, mock_storage_url)


    def drawBoletoCarneDuplo(self, boletoDados1, boletoDados2=None):
        """Imprime um boleto tipo carnê com 2 boletos por página.

        :param boletoDados1: Objeto com os dados do boleto a ser preenchido.
            Deve ser subclasse de :class:`pyboleto.data.BoletoData`
        :param boletoDados2: Objeto com os dados do boleto a ser preenchido.
            Deve ser subclasse de :class:`pyboleto.data.BoletoData`
        :type boletoDados1: :class:`pyboleto.data.BoletoData`
        :type boletoDados2: :class:`pyboleto.data.BoletoData`

        """
        raise NotImplementedError('Em desenvolvimento')


async def run_global_search_question_generation(
    query: str,
    sv: SessionVariables,
) -> SearchResult:
    """Run global search question generation process."""
    empty_context_data: dict[str, pd.DataFrame] = {}

    response, context_data = await api.global_search(
        config=sv.graphrag_config.value,
        entities=sv.entities.value,
        communities=sv.communities.value,
        community_reports=sv.community_reports.value,
        dynamic_community_selection=True,
        response_type="Single paragraph",
        community_level=sv.dataset_config.value.community_level,
        query=query,
    )

    # display response and reference context to UI
    return SearchResult(
        search_type=SearchType.Global,
        response=str(response),
        context=context_data if isinstance(context_data, dict) else empty_context_data,
    )


def run_inference(
    text_input: str,
    audio_prompt_text_input: str,
    audio_prompt_input: Optional[Tuple[int, np.ndarray]],
    max_new_tokens: int,
    cfg_scale: float,
    temperature: float,
    top_p: float,
    cfg_filter_top_k: int,
    speed_factor: float,
    seed: Optional[int] = None,
):
    """
    Runs Nari inference using the globally loaded model and provided inputs.
    Uses temporary files for text and audio prompt compatibility with inference.generate.
    """
    global model, device  # Access global model, config, device
    console_output_buffer = io.StringIO()

    with contextlib.redirect_stdout(console_output_buffer):
        # Prepend transcript text if audio_prompt provided
        if audio_prompt_input and audio_prompt_text_input and not audio_prompt_text_input.isspace():
            text_input = audio_prompt_text_input + "\n" + text_input
            text_input = text_input.strip()

        if audio_prompt_input and (not audio_prompt_text_input or audio_prompt_text_input.isspace()):
            raise gr.Error("Audio Prompt Text input cannot be empty.")

        if not text_input or text_input.isspace():
            raise gr.Error("Text input cannot be empty.")

        # Preprocess Audio
        temp_txt_file_path = None
        temp_audio_prompt_path = None
        output_audio = (44100, np.zeros(1, dtype=np.float32))

        try:
            prompt_path_for_generate = None
            if audio_prompt_input is not None:
                sr, audio_data = audio_prompt_input
                # Check if audio_data is valid
                if audio_data is None or audio_data.size == 0 or audio_data.max() == 0:  # Check for silence/empty
                    gr.Warning("Audio prompt seems empty or silent, ignoring prompt.")
                else:
                    # Save prompt audio to a temporary WAV file
                    with tempfile.NamedTemporaryFile(mode="wb", suffix=".wav", delete=False) as f_audio:
                        temp_audio_prompt_path = f_audio.name  # Store path for cleanup

                        # Basic audio preprocessing for consistency
                        # Convert to float32 in [-1, 1] range if integer type
                        if np.issubdtype(audio_data.dtype, np.integer):
                            max_val = np.iinfo(audio_data.dtype).max
                            audio_data = audio_data.astype(np.float32) / max_val
                        elif not np.issubdtype(audio_data.dtype, np.floating):
                            gr.Warning(f"Unsupported audio prompt dtype {audio_data.dtype}, attempting conversion.")
                            # Attempt conversion, might fail for complex types
                            try:
                                audio_data = audio_data.astype(np.float32)
                            except Exception as conv_e:
                                raise gr.Error(f"Failed to convert audio prompt to float32: {conv_e}")

                        # Ensure mono (average channels if stereo)
                        if audio_data.ndim > 1:
                            if audio_data.shape[0] == 2:  # Assume (2, N)
                                audio_data = np.mean(audio_data, axis=0)
                            elif audio_data.shape[1] == 2:  # Assume (N, 2)
                                audio_data = np.mean(audio_data, axis=1)
                            else:
                                gr.Warning(
                                    f"Audio prompt has unexpected shape {audio_data.shape}, taking first channel/axis."
                                )
                                audio_data = (
                                    audio_data[0] if audio_data.shape[0] < audio_data.shape[1] else audio_data[:, 0]
                                )
                            audio_data = np.ascontiguousarray(audio_data)  # Ensure contiguous after slicing/mean

                        # Write using soundfile
                        try:
                            sf.write(
                                temp_audio_prompt_path, audio_data, sr, subtype="FLOAT"
                            )  # Explicitly use FLOAT subtype
                            prompt_path_for_generate = temp_audio_prompt_path
                            print(f"Created temporary audio prompt file: {temp_audio_prompt_path} (orig sr: {sr})")
                        except Exception as write_e:
                            print(f"Error writing temporary audio file: {write_e}")
                            raise gr.Error(f"Failed to save audio prompt: {write_e}")

            # Set and Display Generation Seed
            if seed is None or seed < 0:
                seed = random.randint(0, 2**32 - 1)
                print(f"\nNo seed provided, generated random seed: {seed}\n")
            else:
                print(f"\nUsing user-selected seed: {seed}\n")
            set_seed(seed)

            # Run Generation
            print(f'Generating speech: \n"{text_input}"\n')

            start_time = time.time()

            # Use torch.inference_mode() context manager for the generation call
            with torch.inference_mode():
                output_audio_np = model.generate(
                    text_input,
                    max_tokens=max_new_tokens,
                    cfg_scale=cfg_scale,
                    temperature=temperature,
                    top_p=top_p,
                    cfg_filter_top_k=cfg_filter_top_k,  # Pass the value here
                    use_torch_compile=False,  # Keep False for Gradio stability
                    audio_prompt=prompt_path_for_generate,
                    verbose=True,
                )

            end_time = time.time()
            print(f"Generation finished in {end_time - start_time:.2f} seconds.\n")

            # 4. Convert Codes to Audio
            if output_audio_np is not None:
                # Get sample rate from the loaded DAC model
                output_sr = 44100

                # --- Slow down audio ---
                original_len = len(output_audio_np)
                # Ensure speed_factor is positive and not excessively small/large to avoid issues
                speed_factor = max(0.1, min(speed_factor, 5.0))
                target_len = int(original_len / speed_factor)  # Target length based on speed_factor
                if target_len != original_len and target_len > 0:  # Only interpolate if length changes and is valid
                    x_original = np.arange(original_len)
                    x_resampled = np.linspace(0, original_len - 1, target_len)
                    resampled_audio_np = np.interp(x_resampled, x_original, output_audio_np)
                    output_audio = (
                        output_sr,
                        resampled_audio_np.astype(np.float32),
                    )  # Use resampled audio
                    print(
                        f"Resampled audio from {original_len} to {target_len} samples for {speed_factor:.2f}x speed."
                    )
                else:
                    output_audio = (
                        output_sr,
                        output_audio_np,
                    )  # Keep original if calculation fails or no change
                    print(f"Skipping audio speed adjustment (factor: {speed_factor:.2f}).")
                # --- End slowdown ---

                print(f"Audio conversion successful. Final shape: {output_audio[1].shape}, Sample Rate: {output_sr}")

                # Explicitly convert to int16 to prevent Gradio warning
                if output_audio[1].dtype == np.float32 or output_audio[1].dtype == np.float64:
                    audio_for_gradio = np.clip(output_audio[1], -1.0, 1.0)
                    audio_for_gradio = (audio_for_gradio * 32767).astype(np.int16)
                    output_audio = (output_sr, audio_for_gradio)
                    print("Converted audio to int16 for Gradio output.")

            else:
                print("\nGeneration finished, but no valid tokens were produced.")
                # Return default silence
                gr.Warning("Generation produced no output.")

        except Exception as e:
            print(f"Error during inference: {e}")
            import traceback

            traceback.print_exc()
            # Re-raise as Gradio error to display nicely in the UI
            raise gr.Error(f"Inference failed: {e}")

        finally:
            # Cleanup Temporary Files defensively
            if temp_txt_file_path and Path(temp_txt_file_path).exists():
                try:
                    Path(temp_txt_file_path).unlink()
                    print(f"Deleted temporary text file: {temp_txt_file_path}")
                except OSError as e:
                    print(f"Warning: Error deleting temporary text file {temp_txt_file_path}: {e}")
            if temp_audio_prompt_path and Path(temp_audio_prompt_path).exists():
                try:
                    Path(temp_audio_prompt_path).unlink()
                    print(f"Deleted temporary audio prompt file: {temp_audio_prompt_path}")
                except OSError as e:
                    print(f"Warning: Error deleting temporary audio prompt file {temp_audio_prompt_path}: {e}")

        # After generation, capture the printed output
        console_output = console_output_buffer.getvalue()

    return output_audio, seed, console_output


def check_couchbase_buckets_items(_item, params, data):

    total_items = data.get("curr_items_tot")
    if total_items is not None:
        yield check_levels(
            int(total_items),
            "items_count",
            params.get('curr_items_tot'),
            infoname="Total items in vBuckets",
            human_readable_func=str,
        )

    write_queue = data.get("disk_write_queue")
    if write_queue is not None:
        yield check_levels(
            int(write_queue),
            "disk_write_ql",
            params.get("disk_write_ql"),
            infoname="Items in disk write queue",
            human_readable_func=str,
        )

    fetched = data.get("ep_bg_fetched")
    if fetched is not None:
        yield check_levels(
            int(fetched),
            "fetched_items",
            params.get("fetched_items"),
            infoname="Items fetched from disk",
            human_readable_func=str,
        )

    queue_fill = data.get("ep_diskqueue_fill")
    if queue_fill is not None:
        yield check_levels(
            queue_fill,
            "disk_fill_rate",
            params.get("disk_fill_rate"),
            unit="/s",
            infoname="Disk queue fill rate",
        )

    queue_drain = data.get("ep_diskqueue_drain")
    if queue_drain is not None:
        yield check_levels(
            queue_drain,
            "disk_drain_rate",
            params.get("disk_drain_rate"),
            unit="/s",
            infoname="Disk queue drain rate",
        )


    def setUp(self):
        self.test_db = cache.init_test_db()


    def __init__(self, real, timeperiod=30):
        pass


    def on_status_button(self, button):
        '''
        Invoked when the user clicks the state button to the left of the
        dropdown.
        '''
        # toggle the control's status state
        isavail = StatusMessage.is_available_state(self.status_state)

        # do we need to change the shown text?
        needs_change = self._shown_status in StatusMessage.SpecialStatuses or not self._shown_status.editable

        self.oldValue = None

        self.change_state(state = 'Away' if isavail else 'Available',)


def update_module_description(module_name: str) -> str:
    """Return the description of the selected module"""
    return MODULE_DESCRIPTIONS.get(module_name, "No description available")


def do_not_run_on_check():
    """
    Do not execute task file on check
    """
    print("#educational_plugin DO_NOT_RUN_ON_CHECK")


    def test_query_level_multiple(self):
        """
        wineventlog - Query by level (error, warning)
        """
        self.write_event_log(
            "success", level=win32evtlog.EVENTLOG_SUCCESS)  # Level 0, Info
        self.write_event_log(
            "error", level=win32evtlog.EVENTLOG_ERROR_TYPE)  # Level 2
        self.write_event_log(
            "warning", level=win32evtlog.EVENTLOG_WARNING_TYPE)  # Level 3
        self.write_event_log(
            "information", level=win32evtlog.EVENTLOG_INFORMATION_TYPE)  # Level 4
        evts = self.read_events(config={
            "event_logs": [
                {
                    "name": self.providerName,
                    "api": self.api,
                    "level": "error, warning"
                }
            ]
        }, expected_events=2)
        self.assertTrue(len(evts), 2)
        self.assertEqual(evts[0]["log.level"], "error")
        self.assertEqual(evts[1]["log.level"], "warning")


    def add_role_to_user(admin_user_id: int, username: str, role: str):
        """
        Add role to user
        :param admin_user_id: ID of admin attempting to add the role
        :param username: Username of user the role should be added to
        :param role: The requested role
        :raises UserServiceError
        """
        try:
            requested_role = UserRole[role.upper()]
        except KeyError:
            raise UserServiceError(
                f"Unknown role {role} accepted values are ADMIN, PROJECT_MANAGER, VALIDATOR"
            )

        admin = UserService.get_user_by_id(admin_user_id)
        admin_role = UserRole(admin.role)

        if admin_role != UserRole.ADMIN and requested_role == UserRole.ADMIN:
            raise UserServiceError("You must be an Admin to assign Admin role")

        user = UserService.get_user_by_username(username)
        user.set_user_role(requested_role)


def apply_single_tag_set(tag_set, selection):
    """All servers matching one tag set.

    A tag set is a dict. A server matches if its tags are a superset:
    A server tagged {'a': '1', 'b': '2'} matches the tag set {'a': '1'}.

    The empty tag set {} matches any server.
    """
    def tags_match(server_tags):
        for key, value in tag_set.items():
            if key not in server_tags or server_tags[key] != value:
                return False

        return True

    return selection.with_server_descriptions(
        [s for s in selection.server_descriptions if tags_match(s.tags)])


    def digest(self):
        return lmhash.encrypt(self._data[:15]).decode('hex')


    def export_spectrogram(self, spec, file_spec):
        save_spectrogram(spec, file_spec)


    def repr_set(self, x, level):
        if not x:
            return 'set()'
        x = _possibly_sorted(x)
        return self._repr_iterable(x, level, '{', '}', self.maxset)


def test_is_in_or_part_overlap(box1: tuple, box2: tuple, target_bool: bool) -> None:
    """
    box1: 坐标数组
    box2: 坐标数组
    """
    assert target_bool == _is_in_or_part_overlap(box1, box2)


def start_server(host, port, processor):
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
        s.bind((host, port))
        s.listen()
        logger.info(f"Server started on {host}:{port}")
        while True:
            conn, addr = s.accept()
            logger.info(f"Connected by {addr}")
            handle_client(conn, processor)


    def __mul__(self, other):
        if isinstance(other, Quaternion):
            return self.__class__(array=np.dot(self._q_matrix(), other.q))
        return self * self.__class__(other)


def load_json_file(json_file: str) -> dict:
    with open(json_file, 'r', encoding='utf-8') as file:
        return json.load(file)


    def set_repository(self, *args, **kwargs):
        """
        Alias for :meth:`Client.get_or_create_repository()`.

        """
        return self.get_or_create_repository(*args, **kwargs)


    def testBadUTFSurrogatePairsGeneration(self):
        local_pdf_path = os.path.join(os.path.dirname(__file__), "gnarly_pdfs", "badlines.pdf")

        anchor_text = get_anchor_text(local_pdf_path, 4, pdf_engine="pdfreport")

        jsondata = json.dumps({"text": anchor_text})

        import pyarrow as pa
        import pyarrow.compute as pc
        import pyarrow.json as paj

        buffer = io.BytesIO(jsondata.encode("utf-8"))
        paj.read_json(buffer, read_options=paj.ReadOptions(use_threads=False, block_size=len(jsondata)))


def test_record_and_resume(tmp_path):
    robot_cfg = MockRobotConfig()
    teleop_cfg = MockTeleopConfig()
    dataset_cfg = DatasetRecordConfig(
        repo_id=DUMMY_REPO_ID,
        single_task="Dummy task",
        root=tmp_path / "record",
        num_episodes=1,
        episode_time_s=0.1,
        reset_time_s=0,
        push_to_hub=False,
    )
    cfg = RecordConfig(
        robot=robot_cfg,
        dataset=dataset_cfg,
        teleop=teleop_cfg,
        play_sounds=False,
    )

    dataset = record(cfg)

    assert dataset.fps == 30
    assert dataset.meta.total_episodes == dataset.num_episodes == 1
    assert dataset.meta.total_frames == dataset.num_frames == 3
    assert dataset.meta.total_tasks == 1

    cfg.resume = True
    dataset = record(cfg)

    assert dataset.meta.total_episodes == dataset.num_episodes == 2
    assert dataset.meta.total_frames == dataset.num_frames == 6
    assert dataset.meta.total_tasks == 1


  def _log_cdf(self, x):
    x = control_flow_ops.with_dependencies([check_ops.assert_positive(x)] if
                                           self.validate_args else [], x)
    contrib_tensor_util.assert_same_float_dtype(tensors=[x], dtype=self.dtype)
    # Note that igamma returns the regularized incomplete gamma function,
    # which is what we want for the CDF.
    return math_ops.log(math_ops.igamma(self.alpha, self.beta * x))


def clean_bigvgan_model():
    global bigvgan_model
    if bigvgan_model:
        bigvgan_model = bigvgan_model.cpu()
        bigvgan_model = None
        try:
            torch.cuda.empty_cache()
        except:
            pass


def clear_span_processor():
    SPAN_PROCESSOR_TESTING.force_flush()
    SPAN_PROCESSOR_TESTING.shutdown()
    SPAN_PROCESSOR_TESTING.clear()


def union_tensor_dict(tensor_dict1: TensorDict, tensor_dict2: TensorDict) -> TensorDict:
    """Union two tensordicts."""
    assert tensor_dict1.batch_size == tensor_dict2.batch_size, \
        f'Two tensor dict must have identical batch size. Got {tensor_dict1.batch_size} and {tensor_dict2.batch_size}'
    for key in tensor_dict2.keys():
        if key not in tensor_dict1.keys():
            tensor_dict1[key] = tensor_dict2[key]
        else:
            assert tensor_dict1[key].equal(tensor_dict2[key]), \
                f'{key} in tensor_dict1 and tensor_dict2 are not the same object'

    return tensor_dict1


    def clean(self):
        cleaned_data = self.cleaned_data
        field_function = cleaned_data.get("field_function")
        choices = cleaned_data.get("choices")
        field_type = cleaned_data.get("field_type")
        #required = cleaned_data.get("required")
        visible = cleaned_data.get("visible")

        if field_function == "GroupSubscription":
            if field_type != "BooleanField":
                raise forms.ValidationError(_("This field's function requires Checkbox as a field type"))
            if not choices:
                raise forms.ValidationError(_("This field's function requires at least 1 group specified."))
            else:
                for val in choices.split(','):
                    try:
                        g = Group.objects.get(name=val.strip())
                        if not g.allow_self_add:
                            raise forms.ValidationError(_("The group \"%(val)s\" does not allow self-add." % { 'val' : val }))
                    except Group.DoesNotExist:
                        raise forms.ValidationError(_("The group \"%(val)s\" does not exist" % { 'val' : val }))

        if field_function == "GroupSubscriptionAuto":
            # field_type doesn't matter since this field shouldn't be rendered anyway.
            if visible:
                raise forms.ValidationError(_("This field must not be visible to users."))
            if not choices:
                raise forms.ValidationError(_("This field's function requires at least 1 group specified."))
            else:
                for val in choices.split(','):
                    try:
                        g = Group.objects.get(name=val.strip())
                        if not g.allow_self_add:
                            raise forms.ValidationError(_("The group \"%(val)s\" does not allow self-add." % { 'val' : val }))
                    except Group.DoesNotExist:
                        raise forms.ValidationError(_("The group \"%(val)s\" does not exist" % { 'val' : val }))

        if field_function == "Recipients":
            if (field_type != "MultipleChoiceField/django.forms.CheckboxSelectMultiple" and
                field_type != "MultipleChoiceField" and
                field_type != "BooleanField" and
                field_type != "ChoiceField"):
                raise forms.ValidationError(_("The \"Email to Recipients\" function requires Multi-select - Checkboxes "
                                            + "or Multi-select - Select Many as field type"))

            if field_type == "BooleanField":
                if not choices:
                    raise forms.ValidationError(_("The \"Email to Recipients\" function requires at least 1 email specified."))
                else:
                    for val in choices.split(','):
                        if not validate_email(val.strip()):
                            raise forms.ValidationError(_("\"%(val)s\" is not a valid email address" % {'val':val}))
            else:
                if not choices:
                    raise forms.ValidationError(_("The \"Email to Recipients\" function requires at least 1 choice specified."))
                else:
                    for val in choices.split(','):
                        val = val.split(':')
                        if len(val) < 2:
                            raise forms.ValidationError(_("The \"Email to Recipients\" function requires choices to be in the following format: <choice_label>:<email_address>."))
                        if not validate_email(val[1].strip()):
                            raise forms.ValidationError(_("\"%(val)s\" is not a valid email address" % {'val':val[1]}))

        if field_function is not None and field_function.startswith("Email"):
            if field_type != "CharField":
                raise forms.ValidationError(_("This field's function requires Text as a field type"))

        #unrequire the display only fields
        if field_type == "CharField/tendenci.apps.forms_builder.forms.widgets.Description":
            cleaned_data['required'] = False
        elif field_type == "CharField/tendenci.apps.forms_builder.forms.widgets.Header":
            cleaned_data['required'] = False

        return cleaned_data


    def _unwrap_object(self, space, w_obj):
        return space.uint_w(w_obj)


    def __init__(
        self,
        role: str,
        raw_utterance: str,
        utterance_type: str,
        claim_to_make: Optional[str] = None,
        utterance: Optional[str] = None,
        queries: Optional[List[str]] = None,
        raw_retrieved_info: Optional[List[Information]] = None,
        cited_info: Optional[List[Information]] = None,
    ):
        self.utterance = utterance if utterance is not None else raw_utterance
        self.raw_utterance = raw_utterance
        self.role = role if ":" not in role else role.split(":")[0]
        self.role_description = "" if ":" not in role else role.split(":")[1]
        self.queries = queries if queries is not None else []
        self.raw_retrieved_info = (
            raw_retrieved_info if raw_retrieved_info is not None else []
        )
        self.cited_info = cited_info if cited_info is not None else {}
        self.utterance_type = utterance_type
        self.claim_to_make = claim_to_make if claim_to_make is not None else ""


    def is_child_right(self, ind):
        # Check if second child exists
        return self.child_right(ind) < len(self.array)


    def test_repo_updateurl_noargs(self, shell):
        """
        Test do_repo_updateurl no arguments.

        :param shell:
        :return:
        """
        for arg in ["", "repo", "http://foo", "http://bar"]:
            shell.help_repo_updateurl = MagicMock()
            shell.client.channel.software.updateRepUrl = MagicMock()
            shell.do_repo_list = MagicMock()
            mprint = MagicMock()
            logger = MagicMock()

            with patch("spacecmd.repo.print", mprint) as prn, \
                    patch("spacecmd.repo.logging", logger) as lgr:
                out = spacecmd.repo.do_repo_updateurl(shell, "")

            assert out is 1
            assert not mprint.called
            assert not logger.error.called
            assert not shell.client.channel.software.updateRepUrl.called
            assert shell.help_repo_updateurl.called


def normalizer_en_nemo_text() -> Callable[[str], str]:
    from nemo_text_processing.text_normalization.normalize import Normalizer

    return partial(
        Normalizer(input_case="cased", lang="en").normalize,
        verbose=False,
        punct_post_process=True,
    )


    def _load_dotted_module_py2(self, short_name, full_name, searchpaths=None):
        import imp

        splitname = short_name.split('.')
        index = 0
        srchpths = searchpaths
        for item in splitname:
            file, path, desc = imp.find_module(item, srchpths)
            mod = imp.load_module(full_name, file, path, desc)
            srchpths = [path]
        return mod, file


def save_json_file(json_file: str, data: dict) -> None:
    with open(json_file, 'w', encoding='utf-8') as file:
        json.dump(data, file, ensure_ascii=False, indent=4)


def wrap_question_for_omni_lmm(question, image_token_len, tokenizer):
    question = expand_question_into_multimodal(
        question, image_token_len, DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN, DEFAULT_IMAGE_PATCH_TOKEN)

    conversation = question
    data_dict = omni_preprocess(sources=[conversation],
                                  tokenizer=tokenizer,
                                  generation=True)

    data_dict = dict(input_ids=data_dict["input_ids"][0],
                     labels=data_dict["labels"][0])
    return data_dict


    def _tool_module(self):
        oldpythonpath = sys.path
        sys.path = self.toolpath + sys.path
        # sys.stderr.write("Tool:%s\nPATH:%s\n"%(self.name,sys.path))

        # From: http://stackoverflow.com/questions/67631/how-to-import-a-module-given-the-full-path/67692#67692
        # import importlib.util
        # spec = importlib.util.spec_from_file_location("module.name", "/path/to/file.py")
        # foo = importlib.util.module_from_spec(spec)
        # spec.loader.exec_module(foo)
        # foo.MyClass()
        # Py 3 code

        # sys.stderr.write("toolpath:%s\n" % self.toolpath)
        # sys.stderr.write("SCONS.TOOL path:%s\n" % sys.modules['SCons.Tool'].__path__)
        debug = False
        spec = None
        found_name = self.name
        add_to_scons_tools_namespace = False
        for path in self.toolpath:
            sepname = self.name.replace('.', os.path.sep)
            file_path = os.path.join(path, "%s.py" % sepname)
            file_package = os.path.join(path, sepname)

            if debug: sys.stderr.write("Trying:%s %s\n" % (file_path, file_package))

            if os.path.isfile(file_path):
                spec = importlib.util.spec_from_file_location(self.name, file_path)
                if debug: print("file_Path:%s FOUND" % file_path)
                break
            elif os.path.isdir(file_package):
                file_package = os.path.join(file_package, '__init__.py')
                spec = importlib.util.spec_from_file_location(self.name, file_package)
                if debug: print("PACKAGE:%s Found" % file_package)
                break

            else:
                continue

        if spec is None:
            if debug: sys.stderr.write("NO SPEC :%s\n" % self.name)
            spec = importlib.util.find_spec("." + self.name, package='SCons.Tool')
            if spec:
                found_name = 'SCons.Tool.' + self.name
                add_to_scons_tools_namespace = True
            if debug: sys.stderr.write("Spec Found? .%s :%s\n" % (self.name, spec))

        if spec is None:
            sconstools = os.path.normpath(sys.modules['SCons.Tool'].__path__[0])
            if self.toolpath:
                sconstools = ", ".join(self.toolpath) + ", " + sconstools
            error_string = "No tool module '%s' found in %s" % (self.name, sconstools)
            raise SCons.Errors.UserError(error_string)

        module = importlib.util.module_from_spec(spec)
        if module is None:
            if debug: print("MODULE IS NONE:%s" % self.name)
            error_string = "Tool module '%s' failed import" % self.name
            raise SCons.Errors.SConsEnvironmentError(error_string)

        # Don't reload a tool we already loaded.
        sys_modules_value = sys.modules.get(found_name, False)

        found_module = None
        if sys_modules_value and sys_modules_value.__file__ == spec.origin:
            found_module = sys.modules[found_name]
        else:
            # Not sure what to do in the case that there already
            # exists sys.modules[self.name] but the source file is
            # different.. ?
            module = spec.loader.load_module(spec.name)

            sys.modules[found_name] = module
            if add_to_scons_tools_namespace:
                # If we found it in SCons.Tool, then add it to the module
                setattr(SCons.Tool, self.name, module)

            found_module = module

        if found_module is not None:
            sys.path = oldpythonpath
            return found_module

        sys.path = oldpythonpath

        full_name = 'SCons.Tool.' + self.name
        try:
            return sys.modules[full_name]
        except KeyError:
            try:
                smpath = sys.modules['SCons.Tool'].__path__
                try:
                    module, file = self._load_dotted_module_py2(self.name, full_name, smpath)
                    setattr(SCons.Tool, self.name, module)
                    if file:
                        file.close()
                    return module
                except ImportError as e:
                    if str(e) != "No module named %s" % self.name:
                        raise SCons.Errors.SConsEnvironmentError(e)
                    try:
                        import zipimport
                        importer = zipimport.zipimporter(sys.modules['SCons.Tool'].__path__[0])
                        module = importer.load_module(full_name)
                        setattr(SCons.Tool, self.name, module)
                        return module
                    except ImportError as e:
                        m = "No tool named '%s': %s" % (self.name, e)
                        raise SCons.Errors.SConsEnvironmentError(m)
            except ImportError as e:
                m = "No tool named '%s': %s" % (self.name, e)
                raise SCons.Errors.SConsEnvironmentError(m)


    def __init__(self, content: str):
        super().__init__({"choices": [{"message": {"content": content}}], "usage": {}})
        self.choices = [types.SimpleNamespace(message=types.SimpleNamespace(content=content))]


def signal_handler(sig, frame):
    print("Ctrl+C pressed, shutting down server...")
    sys.exit(0)


    def test_atomic_max_nan_val(self):
        res = np.random.randint(0, 128, size=1).astype(np.float64)
        gold = res.copy()
        vals = np.zeros((1, 1), np.float64) + np.nan
        cuda_func = cuda.jit('void(float64[:], float64[:,:])')(atomic_max)
        cuda_func[1, 1](res, vals)

        np.testing.assert_equal(res, gold)


def __le__(self, other: Any) -> bool: ...


    def coerce(cls, key, value):
        """Convert plain lists to MutableList."""
        if not isinstance(value, MutableList):
            if isinstance(value, list):
                return MutableList(value)

            # this call will raise ValueError
            return mutable.Mutable.coerce(key, value)
        return value


def test_cli_writes_file(
    tmp_path: Path,
    monkeypatch: pytest.MonkeyPatch,
    *,
    cli_args: list[str],
    expect_file: bool,
) -> None:
    """Run the CLI and verify that the SARIF file is created (or not)."""
    expectes_exit_code = 0
    # Work inside an isolated temp directory
    monkeypatch.chdir(tmp_path)

    result = _invoke_isolated_cli_runner(cli_args)

    assert result.exit_code == expectes_exit_code, result.stderr

    # Summary line should be on STDOUT
    stdout_lines = result.stdout.splitlines()
    assert f"Analysis complete! Output written to: {OUTPUT_FILE_NAME}" in stdout_lines

    # File side-effect
    sarif_file = tmp_path / OUTPUT_FILE_NAME
    assert sarif_file.exists() is expect_file, f"{OUTPUT_FILE_NAME} existence did not match expectation"


    def _buildTagListRequest(self, programName):
        """
        Build the request for the PLC tags
        Program scoped tags will pass the program name for the request
        """
        Service = 0x55
        PathSegment = b""

        # If we're dealing with program scoped tags...
        if programName:
            PathSegment = pack('<BB', 0x91, len(programName)) + programName.encode('utf-8')
            # if odd number of characters, need to add a byte to the end.
            if len(programName) % 2:
                PathSegment += pack('<B', 0x00)

        PathSegment += pack('<H', 0x6B20)

        if self.Offset < 256:
            PathSegment += pack('<BB', 0x24, self.Offset)
        else:
            PathSegment += pack('<HH', 0x25, self.Offset)

        PathSegmentLen = int(len(PathSegment)/2)
        AttributeCount = 0x03
        SymbolType = 0x02
        ByteCount = 0x08
        SymbolName = 0x01
        Attributes = pack('<HHHH', AttributeCount, SymbolName, SymbolType, ByteCount)
        TagListRequest = pack('<BB', Service, PathSegmentLen)
        TagListRequest += PathSegment + Attributes

        return TagListRequest


def is_full(*items):  # 任意一项为空返回False
    for item in items:
        if item is None or item == "":
            return False
    return True


async def run() -> None:
    file_name = "skyvern_screenshot.png"
    png_file_path = f"{get_skyvern_temp_dir()}/{file_name}"

    while True:
        # run subprocess to take screenshot
        subprocess.run(
            f"xwd -root | xwdtopnm 2>/dev/null | pnmtopng > {png_file_path}", shell=True, env={"DISPLAY": ":99"}
        )

        # FIXME: upload screenshot to S3 with correct organization id
        try:
            await app.STORAGE.save_streaming_file("placeholder_org", file_name)
        except Exception:
            LOG.info("Failed to save screenshot")

        await asyncio.sleep(INTERVAL)


    def dump_options(self, header=None, indent=""):
        from distutils.fancy_getopt import longopt_xlate
        if header is None:
            header = "command options for '%s':" % self.get_command_name()
        self.announce(indent + header, level=log.INFO)
        indent = indent + "  "
        for (option, _, _) in self.user_options:
            option = option.translate(longopt_xlate)
            if option[-1] == "=":
                option = option[:-1]
            value = getattr(self, option)
            self.announce(indent + "%s = %s" % (option, value),
                          level=log.INFO)


async def run_drift_search(
    query: str,
    sv: SessionVariables,
) -> SearchResult:
    """Run drift search."""
    print(f"Drift search query: {query}")  # noqa T201

    # build drift search engine
    response_placeholder = st.session_state[
        f"{SearchType.Drift.value.lower()}_response_placeholder"
    ]
    response_container = st.session_state[f"{SearchType.Drift.value.lower()}_container"]

    with response_placeholder, st.spinner("Generating answer using drift search..."):
        empty_context_data: dict[str, pd.DataFrame] = {}

        response, context_data = await api.drift_search(
            config=sv.graphrag_config.value,
            entities=sv.entities.value,
            communities=sv.communities.value,
            community_reports=sv.community_reports.value,
            text_units=sv.text_units.value,
            relationships=sv.relationships.value,
            community_level=sv.dataset_config.value.community_level,
            response_type="Multiple Paragraphs",
            query=query,
        )

        print(f"Drift Response: {response}")  # noqa T201
        print(f"Context data: {context_data}")  # noqa T201

    # display response and reference context to UI
    search_result = SearchResult(
        search_type=SearchType.Drift,
        response=str(response),
        context=context_data if isinstance(context_data, dict) else empty_context_data,
    )

    display_search_result(
        container=response_container, result=search_result, stats=None
    )

    if "response_lengths" not in st.session_state:
        st.session_state.response_lengths = []

    st.session_state["response_lengths"].append({
        "result": None,
        "search": SearchType.Drift.value.lower(),
    })

    return search_result


    def append_message(self, role: str, message: str):
        """Append a new message."""
        self.messages.append([role, message])


def get_content_for_indices(files_data, indices):
    content_map = {}
    for i in indices:
        if 0 <= i < len(files_data):
            path, content = files_data[i]
            content_map[f"{i} # {path}"] = (
                content  # Use index + path as key for context
            )
    return content_map


def elbo(q, p, sample_dim=None, batch_dim=None, alpha=0.1,
         size_average=True, reduce=True):
    r"""Calculates an importance weighted Monte Carlo estimate of the
    semi-supervised evidence lower bound (ELBO)

    .. math:: \frac{1}{B} \sum_{b=1}^B
              \log \left[
                \frac{1}{S} \sum_{s=1}^S
                \frac{p(x^{(b)}, y^{(b)}, z^{(s,b)})}
                     {q(z^{(s,b)} | x^{(b)})}
              \right]
              + \frac{\alpha}{B} \sum_{b=1}^B
              \log \left[
                \frac{1}{S} \sum_{s=1}^S
                \frac{q(y^{(b)}, z^{(s,b)} | x^{(b)})}
                     {q(z^{(s,b)} | x^{(b)})}
              \right]

    The sets of variables :math:`x`, :math:`y` and :math:`z` refer to:

        :math:`x`: The set of conditioned nodes that are present in `p` but
        are not present in `q`.

        :math:`y`: The set of conditioned nodes in `q`, which may or may
        not also be present in `q`.

        :math:`z`: The set of sampled nodes present in both `q` and `p`.

    Arguments:
        q(:obj:`Trace`): The encoder trace.
        p(:obj:`Trace`): The decoder trace.
        sample_dim(int, optional): The dimension containing individual samples.
        batch_dim(int, optional): The dimension containing batch items.
        alpha(float, default 0.1): Coefficient for the ML term.
        size_average (bool, optional): By default, the objective is averaged
            over items in the minibatch. When set to false, the objective is
            instead summed over the minibatch.
        reduce (bool, optional): By default, the objective is averaged or
           summed over items in the minibatch. When reduce is False, losses
           are returned without averaging or summation.
    """
    z = [n for n in q.sampled() if n in p]
    log_pxyz = p.log_joint(sample_dim, batch_dim)
    log_qz = q.log_joint(sample_dim, batch_dim, z)
    log_qy = q.log_joint(sample_dim, batch_dim, q.conditioned())
    log_pq = (log_pxyz - log_qz)
    if sample_dim is None:
        objective = log_pq + alpha * log_qy
    else:
        objective = log_mean_exp(log_pq, 0)
        if not isinstance(log_qy, Number):
            objective = objective + alpha * log_mean_exp(log_qy, 0)
    if reduce:
        objective = objective.mean() if size_average else objective.sum()
    return objective


def _assert_not_shallow(request):
    if request.shallow:
        raise RuntimeError('A shallow request tried to consume '
                           'form data.  If you really want to do '
                           'that, set `shallow` to False.')


async def test_simple_echo():
    """Test the simple echo server"""
    from examples.simple_echo import mcp

    async with Client(mcp) as client:
        result = await client.call_tool_mcp("echo", {"text": "hello"})
        assert len(result.content) == 1
        assert result.content[0].text == "hello"  # type: ignore[attr-defined]


def _test__xla_dist_model_create_from_context_in_child_proc(index):
    model = _XlaDistModel.create_from_context()

    assert model.backend() == "xla-tpu"

    import torch_xla.core.xla_model as xm

    _assert_model(
        model,
        {
            "device": xm.xla_device(),
            "local_rank": index,
            "rank": xm.get_ordinal(),
            "world_size": xm.xrt_world_size(),
            "node_index": 0,
            "nnodes": 1,
            "nproc_per_node": xm.xrt_world_size(),
        },
    )


def create_resume_pdf(parameters: dict, llm_api_key: str):
    """
    Logic to create a CV.
    """
    try:
        logger.info("Generating a CV based on provided parameters.")

        # Load the plain text resume
        with open(parameters["uploads"]["plainTextResume"], "r", encoding="utf-8") as file:
            plain_text_resume = file.read()

        # Initialize StyleManager
        style_manager = StyleManager()
        available_styles = style_manager.get_styles()

        if not available_styles:
            logger.warning("No styles available. Proceeding without style selection.")
        else:
            # Present style choices to the user
            choices = style_manager.format_choices(available_styles)
            questions = [
                inquirer.List(
                    "style",
                    message="Select a style for the resume:",
                    choices=choices,
                )
            ]
            style_answer = inquirer.prompt(questions)
            if style_answer and "style" in style_answer:
                selected_choice = style_answer["style"]
                for style_name, (file_name, author_link) in available_styles.items():
                    if selected_choice.startswith(style_name):
                        style_manager.set_selected_style(style_name)
                        logger.info(f"Selected style: {style_name}")
                        break
            else:
                logger.warning("No style selected. Proceeding with default style.")

        # Initialize the Resume Generator
        resume_generator = ResumeGenerator()
        resume_object = Resume(plain_text_resume)
        driver = init_browser()
        resume_generator.set_resume_object(resume_object)

        # Create the ResumeFacade
        resume_facade = ResumeFacade(
            api_key=llm_api_key,
            style_manager=style_manager,
            resume_generator=resume_generator,
            resume_object=resume_object,
            output_path=Path("data_folder/output"),
        )
        resume_facade.set_driver(driver)
        result_base64 = resume_facade.create_resume_pdf()

        # Decode Base64 to binary data
        try:
            pdf_data = base64.b64decode(result_base64)
        except base64.binascii.Error as e:
            logger.error("Error decoding Base64: %s", e)
            raise

        # Define the output directory using `suggested_name`
        output_dir = Path(parameters["outputFileDirectory"])

        # Write the PDF file
        output_path = output_dir / "resume_base.pdf"
        try:
            with open(output_path, "wb") as file:
                file.write(pdf_data)
            logger.info(f"Resume saved at: {output_path}")
        except IOError as e:
            logger.error("Error writing file: %s", e)
            raise
    except Exception as e:
        logger.exception(f"An error occurred while creating the CV: {e}")
        raise


    def __init__(self):
        self.on_npu = npu_is_available
        self.node_world_size = self.N_NPU_PER_NODE
        self.profiling = False
        self.profiling_step = 5
        self.enable_FA = True
        self.enable_FP32 = False
        self.load_pickle = True
        self.use_small_dataset = False
        self.current_run_dtype = None
        self.original_run_dtype = None
        self.zp_manager = zp_manager
        self.replaced_type = torch.float32
        self.conv_dtype = torch.float16
        if self.enable_FA and self.enable_FP32:
            self.inf_float = -10000.0
        else:
            self.inf_float = -10000.0

        if self.use_small_dataset:
            self.load_pickle = False

        self._loss = []
        self.work_path = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
        self.pickle_save_path = f"{self.work_path}/pickles"
        self.mm = dict()

        if self.on_npu:
            import deepspeed
            import sys
            torch_npu.npu.set_compile_mode(jit_compile=False)

            import deepspeed.runtime.utils as utils
            from opensora.adaptor.utils import all_gather_dp_groups, all_gather_into_tensor_dp_groups
            utils.all_gather_dp_groups = all_gather_dp_groups

            import deepspeed.runtime.bf16_optimizer as bf16_optimizer
            from opensora.adaptor.bf16_optimizer import BF16_Optimizer
            self.replace_methods(bf16_optimizer.BF16_Optimizer, BF16_Optimizer)

            from opensora.adaptor.stage_1_and_2 import DeepSpeedZeroOptimizer
            import deepspeed.runtime.zero.stage_1_and_2 as stage_1_and_2
            self.replace_methods(stage_1_and_2.DeepSpeedZeroOptimizer, DeepSpeedZeroOptimizer, ['_has_inf_or_nan'])

            import deepspeed.runtime.engine as engine
            from opensora.adaptor.engine import DeepSpeedEngine
            self.replace_methods(engine.DeepSpeedEngine, DeepSpeedEngine, skip_fcns=['__init__', '_copy_recovery_script', '_change_recovery_script_permissions'])

        if "RANK" in os.environ:
            self.rank = int(os.environ["RANK"])
            self.world_size = int(os.environ["WORLD_SIZE"])
            torch_npu.npu.set_device(self.get_local_rank())
        else:
            self.rank = torch.cuda.current_device()
            self.world_size = self.N_NPU_PER_NODE
        self.print_with_rank(f"The npu_config.on_npu is {self.on_npu}")
        self.bind_thread_to_cpu()
        gc.set_threshold(700, 10, 10000)


    def chat_completion(
        self,
        dialogs: List[Dialog],
        temperature: float = 0.6,
        top_p: float = 0.9,
        max_gen_len: Optional[int] = None,
        logprobs: bool = False,
    ) -> List[ChatPrediction]:
        """
        Generate assistant responses for a list of conversational dialogs using the language generation model.

        Args:
            dialogs (List[Dialog]): List of conversational dialogs, where each dialog is a list of messages.
            temperature (float, optional): Temperature value for controlling randomness in sampling. Defaults to 0.6.
            top_p (float, optional): Top-p probability threshold for nucleus sampling. Defaults to 0.9.
            max_gen_len (Optional[int], optional): Maximum length of the generated response sequence.
                If not provided, it's set to the model's maximum sequence length minus 1.
            logprobs (bool, optional): Flag indicating whether to compute token log probabilities. Defaults to False.

        Returns:
            List[ChatPrediction]: List of chat predictions, each containing the assistant's generated response.

        Note:
            This method generates assistant responses for the provided conversational dialogs.
            It employs nucleus sampling to introduce controlled randomness in text generation.
            If logprobs is True, token log probabilities are computed for each generated token.
        """
        if max_gen_len is None:
            max_gen_len = self.model.params.max_seq_len - 1

        prompt_tokens = [
            self.formatter.encode_dialog_prompt(dialog) for dialog in dialogs
        ]
        generation_tokens, generation_logprobs = self.generate(
            prompt_tokens=prompt_tokens,
            max_gen_len=max_gen_len,
            temperature=temperature,
            top_p=top_p,
            logprobs=logprobs,
        )
        if logprobs:
            return [
                {
                    "generation": {
                        "role": "assistant",
                        "content": self.tokenizer.decode(t),
                    },
                    "tokens": [self.tokenizer.decode([x]) for x in t],
                    "logprobs": logprobs_i,
                }
                for t, logprobs_i in zip(generation_tokens, generation_logprobs)
            ]
        return [
            {
                "generation": {
                    "role": "assistant",
                    "content": self.tokenizer.decode(t),
                },
            }
            for t in generation_tokens
        ]


def _ExpandDimsGrad(op, grad):
  return [_ReshapeToInput(op, grad), None]


def test_llm_with_ollama(tester: Crawl4AiTester):
    print("\n=== Testing LLM with Ollama ===")
    schema = {
        "type": "object",
        "properties": {
            "article_title": {
                "type": "string",
                "description": "The main title of the news article",
            },
            "summary": {
                "type": "string",
                "description": "A brief summary of the article content",
            },
            "main_topics": {
                "type": "array",
                "items": {"type": "string"},
                "description": "Main topics or themes discussed in the article",
            },
        },
    }

    request = {
        "urls": ["https://www.nbcnews.com/business"],
        "priority": 8,
        "extraction_config": {
            "type": "llm",
            "params": {
                "provider": "ollama/llama2",
                "schema": schema,
                "extraction_type": "schema",
                "instruction": "Extract the main article information including title, summary, and main topics.",
            },
        },
        "extra": {"word_count_threshold": 1},
        "crawler_params": {"verbose": True},
    }

    try:
        result = tester.submit_and_wait(request)
        extracted = json.loads(result["result"]["extracted_content"])
        print("Extracted content:", json.dumps(extracted, indent=2))
        assert result["result"]["success"]
    except Exception as e:
        print(f"Ollama extraction test failed: {str(e)}")


def rename_colmap_recons_and_rescale_camera(
    reconstruction, image_paths, original_coords, img_size, shift_point2d_to_original_res=False, shared_camera=False
):
    rescale_camera = True

    for pyimageid in reconstruction.images:
        # Reshaped the padded&resized image to the original size
        # Rename the images to the original names
        pyimage = reconstruction.images[pyimageid]
        pycamera = reconstruction.cameras[pyimage.camera_id]
        pyimage.name = image_paths[pyimageid - 1]

        if rescale_camera:
            # Rescale the camera parameters
            pred_params = copy.deepcopy(pycamera.params)

            real_image_size = original_coords[pyimageid - 1, -2:]
            resize_ratio = max(real_image_size) / img_size
            pred_params = pred_params * resize_ratio
            real_pp = real_image_size / 2
            pred_params[-2:] = real_pp  # center of the image

            pycamera.params = pred_params
            pycamera.width = real_image_size[0]
            pycamera.height = real_image_size[1]

        if shift_point2d_to_original_res:
            # Also shift the point2D to original resolution
            top_left = original_coords[pyimageid - 1, :2]

            for point2D in pyimage.points2D:
                point2D.xy = (point2D.xy - top_left) * resize_ratio

        if shared_camera:
            # If shared_camera, all images share the same camera
            # no need to rescale any more
            rescale_camera = False

    return reconstruction


    def failing_version(*args, **kwargs):
        pytest.fail("Real models should not be used in tests!")


    def test_handles_single_rate_limit(self, mock_get, mock_sleep):
        """Test that API retries once after a 429 and succeeds."""
        # Setup mock responses: first 429, then 200
        mock_429_response = Mock()
        mock_429_response.status_code = 429

        mock_200_response = Mock()
        mock_200_response.status_code = 200
        mock_200_response.text = "Success"

        mock_get.side_effect = [mock_429_response, mock_200_response]

        # Call the function
        headers = {"X-API-KEY": "test-key"}
        url = "https://api.financialdatasets.ai/test"

        result = _make_api_request(url, headers)

        # Verify behavior
        assert result.status_code == 200
        assert result.text == "Success"

        # Verify requests.get was called twice
        assert mock_get.call_count == 2
        mock_get.assert_has_calls([
            call(url, headers=headers),
            call(url, headers=headers)
        ])

        # Verify sleep was called once with 60 seconds (first retry)
        mock_sleep.assert_called_once_with(60)


    def test_neighboring_returns_neighbors_with_pegged_alignment(self):
        # arrange
        a_info = AlignmentInfo(
            (0, 3, 2),
            (None, 'des', 'œufs', 'verts'),
            ('UNUSED', 'green', 'eggs'),
            [[], [], [2], [1]],
        )
        ibm_model = IBMModel([])

        # act: peg 'eggs' to align with 'œufs'
        neighbors = ibm_model.neighboring(a_info, 2)

        # assert
        neighbor_alignments = set()
        for neighbor in neighbors:
            neighbor_alignments.add(neighbor.alignment)
        expected_alignments = set(
            [
                # moves
                (0, 0, 2),
                (0, 1, 2),
                (0, 2, 2),
                # no swaps
                # original alignment
                (0, 3, 2),
            ]
        )
        self.assertEqual(neighbor_alignments, expected_alignments)


def generate_metadata(video_subject: str, script: str, ai_model: str) -> Tuple[str, str, List[str]]:
    """
    Generate metadata for a YouTube video, including the title, description, and keywords.

    Args:
        video_subject (str): The subject of the video.
        script (str): The script of the video.
        ai_model (str): The AI model to use for generation.

    Returns:
        Tuple[str, str, List[str]]: The title, description, and keywords for the video.
    """

    # Build prompt for title
    title_prompt = f"""
    Generate a catchy and SEO-friendly title for a YouTube shorts video about {video_subject}.
    """

    # Generate title
    title = generate_response(title_prompt, ai_model).strip()

    # Build prompt for description
    description_prompt = f"""
    Write a brief and engaging description for a YouTube shorts video about {video_subject}.
    The video is based on the following script:
    {script}
    """

    # Generate description
    description = generate_response(description_prompt, ai_model).strip()

    # Generate keywords
    keywords = get_search_terms(video_subject, 6, script, ai_model)

    return title, description, keywords


        def check_channels(channels):
            for channel, should_receive in channels:
                dg = Datagram.create([channel], 123456789, 5858)
                dg.add_uint16(channel) # For some semblance of uniqueness

                self.c2.send(dg)

                if should_receive:
                    self.expect(self.c1, dg)
                    self.expectNone(self.c1) # No repeats!
                else:
                    self.expectNone(self.c1)

                # And, of course, l1 receives all of these:
                self.expect(self.l1, dg)


    def setColor(self):
        """ Separate the current image into cyan, magenta, and yellow
            components.  Create a representation of how each component might
            appear when applied to a blank white piece of paper.
        """
        newColor = QColorDialog.getColor(self.paintColor)

        if newColor.isValid():
            self.paintColor = newColor
            palette = QPalette(self.colorButton.palette())
            palette.setColor(QPalette.Button, self.paintColor)
            self.colorButton.setPalette(palette)
            self.createImage()
            self.imageChanged.emit()


    def test_subscribe(self):
        self.l1.flush()
        self.c1.flush()
        self.c2.flush()

        # Subscribe to a channel...
        dg = Datagram.create_add_channel(12345654321)
        self.c1.send(dg)
        self.expectNone(self.c1)
        # Make sure the MD subscribes to its parent.
        self.expect(self.l1, dg)

        # Send a test datagram on second connection...
        dg = Datagram()
        dg = Datagram.create([12345654321], 0, 1234)
        dg.add_uint32(0xDEADBEEF)
        self.c2.send(dg)
        self.expect(self.c1, dg)
        # MD will, of course, relay this datagram upward.
        self.expect(self.l1, dg)

        # Subscribe on the second connection...
        dg = Datagram.create_add_channel(12345654321)
        self.c2.send(dg)
        self.expectNone(self.c2)
        # MD should not ask for the channel a second time.
        self.expectNone(self.l1)

        # Send a test datagram on first connection...
        dg = Datagram.create([12345654321], 0, 1234)
        dg.add_uint32(0xDEADBEEF)
        self.c1.send(dg)
        self.expect(self.c2, dg) # Should be relayed to second.
        self.expect(self.l1, dg) # Should be sent upward.
        #self.expectNone(self.c1) # Should NOT be echoed back.

        # Unsubscribe on the first connection...
        dg = Datagram.create_remove_channel(12345654321)
        self.c1.send(dg)
        self.expectNone(self.c1)
        self.expectNone(self.c2)
        # MD should NOT unsubscribe from parent!
        self.expectNone(self.l1)

        # Send another test datagram on second connection...
        dg = Datagram.create([12345654321], 0, 1234)
        dg.add_uint32(0xDEADBEEF)
        self.c2.send(dg)
        self.expect(self.l1, dg) # Should be sent upward.
        self.expectNone(self.c2) # Should NOT be relayed.
        self.expectNone(self.c1) # Should NOT be echoed back.

        # Abandon the second connection, which should auto-unsubscribe it.
        self.c2.close()
        self.__class__.c2 = self.connectToServer()
        self.expectNone(self.c1)
        # MD should unsubscribe from parent.
        self.expect(self.l1, Datagram.create_remove_channel(12345654321))


def filter_output(output):
    """ImportantMessage"""
    if not output:
        return ""
    important_lines = []
    for line in output.split("\n"):
        # Only keep lines containing specific keywords
        if any(
            keyword in line.lower()
            for keyword in ["error:", "failed:", "completed", "directory:"]
        ):
            important_lines.append(line)
    return "\n".join(important_lines)


    def __eq__(self, other):
        if not isinstance(other, MetaPartition):
            return False

        if self.metadata_version != other.metadata_version:
            return False

        for table, meta in self.table_meta.items():
            # https://issues.apache.org/jira/browse/ARROW-5873
            other_meta = other.table_meta.get(table, None)
            if other_meta is None:
                return False
            if not meta.equals(other_meta):
                return False

        if self.dataset_metadata != other.dataset_metadata:
            return False

        if len(self.metapartitions) != len(other.metapartitions):
            return False

        # In the case both MetaPartitions are nested, we need to ensure a match
        # for all sub-partitions.
        # Since the label is unique, this can be used as a distinguishing key to sort and compare
        # the nested metapartitions.
        if len(self.metapartitions) > 1:
            for mp_self, mp_other in zip(
                sorted(self.metapartitions, key=lambda x: x["label"]),
                sorted(other.metapartitions, key=lambda x: x["label"]),
            ):
                if mp_self == mp_other:
                    continue
                # If a single metapartition does not match, the whole object is considered different
                return False
            return True

        # This is unnested only

        self_keys = set(self.data.keys())
        other_keys = set(other.data.keys())
        if not (self_keys == other_keys):
            return False

        if self.label != other.label:
            return False

        if self.files != other.files:
            return False

        for label, df in self.data.items():
            if not (df.equals(other.data[label])):
                return False

        return True


def gpu_wrapped_execute_image_retargeting(*args, **kwargs):
    return gradio_pipeline.execute_image_retargeting(*args, **kwargs)


def remove_full_width_characters(s):
    # 匹配全角字符的正则表达式
    pattern = re.compile(r'[\uFF00-\uFFEF]')
    # 替换字符串中的全角字符为空字符串
    return pattern.sub('', s)


def event_loop() -> Iterator[None]:
    new_loop = asyncio.new_event_loop()
    asyncio.set_event_loop(new_loop)
    yield
    new_loop.close()


def apply_common_decorators(capture_output_enabled=False):
    """
    A unified decorator for applications
    """

    def decorator(original_cmd_func):
        @functools.wraps(original_cmd_func)
        def new_runtime_wrapper(*args, **kwargs):
            if cli_config and cli_config.full_log:
                return capture_output(original_cmd_func)(*args, **kwargs)
            else:
                return original_cmd_func(*args, **kwargs)

        func_with_clear_argv = clear_argv(new_runtime_wrapper)

        return functools.wraps(original_cmd_func)(func_with_clear_argv)

    return decorator


def test_parse_s3path(s3_path: str, target_data: str):
    """
    s3_path: s3路径
        如果为无效路径，则返回对应的bucket名字和其余部分
        如果为异常路径 例如：file2.txt，则报异常
    """
    bucket_name, key = parse_s3path(s3_path)
    assert target_data == bucket_name


    def get_rocm_bare_metal_version(self, rocm_dir):
        """
        Get the ROCm version from the ROCm installation directory.

        Args:
            rocm_dir: Path to the ROCm installation directory

        Returns:
            A string representation of the ROCm version (e.g., "63" for ROCm 6.3)
        """
        try:
            # Try using rocm_agent_enumerator to get version info
            raw_output = subprocess.check_output(
                [rocm_dir + "/bin/rocminfo", "--version"],
                universal_newlines=True,
                stderr=subprocess.STDOUT)
            # Extract version number from output
            match = re.search(r'(\d+\.\d+)', raw_output)
            if match:
                version_str = match.group(1)
                version = parse(version_str)
                rocm_version = f"{version.major}{version.minor}"
                return rocm_version
        except (subprocess.CalledProcessError, FileNotFoundError):
            # If rocminfo --version fails, try alternative methods
            pass

        try:
            # Try reading version from release file
            with open(os.path.join(rocm_dir, "share/doc/hip/version.txt"), "r") as f:
                version_str = f.read().strip()
                version = parse(version_str)
                rocm_version = f"{version.major}{version.minor}"
                return rocm_version
        except (FileNotFoundError, IOError):
            pass

        # If all else fails, try to extract from directory name
        dir_name = os.path.basename(os.path.normpath(rocm_dir))
        match = re.search(r'rocm-(\d+\.\d+)', dir_name)
        if match:
            version_str = match.group(1)
            version = parse(version_str)
            rocm_version = f"{version.major}{version.minor}"
            return rocm_version

        # Fallback to extracting from hipcc version
        try:
            raw_output = subprocess.check_output(
                [rocm_dir + "/bin/hipcc", "--version"],
                universal_newlines=True,
                stderr=subprocess.STDOUT)
            match = re.search(r'HIP version: (\d+\.\d+)', raw_output)
            if match:
                version_str = match.group(1)
                version = parse(version_str)
                rocm_version = f"{version.major}{version.minor}"
                return rocm_version
        except (subprocess.CalledProcessError, FileNotFoundError):
            pass

        # If we still can't determine the version, raise an error
        raise ValueError(f"Could not determine ROCm version from directory: {rocm_dir}")


    def test_discard_all_filterfunc_not_supported(self, n='xjf21j21'):
        c = compat.Consumer(self.connection, queue=n, exchange=n,
                            routing_key='rkey')
        with pytest.raises(NotImplementedError):
            c.discard_all(filterfunc=lambda x: x)
        c.close()


def get_afm_cache_path() -> str:
    """
    Gets the path to the Affiliate Marketing cache file.

    Returns:
        path (str): The path to the AFM cache folder
    """
    return os.path.join(get_cache_path(), 'afm.json')


    def homogeneous_noncommutative_variables(self,la):
        r"""
        Give the homogeneous function indexed by `la`, viewed inside the Nil-Coxeter algebra.
        This is only defined in finite type `A`, `B` and affine types `A^{(1)}`, `B^{(1)}`, `C^{(1)}`, `D^{(1)}`.

        INPUT:

        - ``la`` -- a partition with first part bounded by the rank of the Weyl group

        EXAMPLES::

            sage: U = NilCoxeterAlgebra(WeylGroup(['B',2,1]))
            sage: U.homogeneous_noncommutative_variables([2,1])
            u[1,2,0] + 2*u[2,1,0] + u[0,2,0] + u[0,2,1] + u[1,2,1] + u[2,1,2] + u[2,0,2] + u[1,0,2]

        TESTS::

            sage: U = NilCoxeterAlgebra(WeylGroup(['B',2,1]))
            sage: U.homogeneous_noncommutative_variables([])
            1

        """
        return prod(self.homogeneous_generator_noncommutative_variables(p) for p in la)


def setup_seed(seed):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False


    def close(self, rm_all_containers: bool | None = True):
        super().close()

        if self.attach_to_existing:
            return

        if self.devbox:
            self.runloop_api_client.devboxes.shutdown(self.devbox.id)


def b2a_uu(data: _Bytes, *, backtick: bool = ...) -> bytes: ...


def repo_exists_true(mocker: MockerFixture) -> AsyncMock:
    """Patch ``gitingest.clone.check_repo_exists`` to always return ``True``."""
    return mocker.patch("gitingest.clone.check_repo_exists", return_value=True)


    def to_dict(self):
        """Returns a dict representation of the PlatformParameter domain
        object.

        Returns:
            dict. A dict mapping of all fields of PlatformParameter object.
        """
        return {
            'name': self._name,
            'description': self._description,
            'data_type': self._data_type,
            'rules': [rule.to_dict() for rule in self._rules],
            'rule_schema_version': self._rule_schema_version,
            'default_value': self._default_value,
            'is_feature': self._is_feature,
            'feature_stage': self._feature_stage
        }


    def test_live_migration_abort(self):
        self._test_compute_api('live_migration_abort', 'cast',
                instance=self.fake_instance_obj,
                migration_id='1', version='5.0')


    def __len__(self):
        return len(self.ids)


    def _drawReciboSacado(self, boletoDados):
        """Imprime o Recibo do Sacado para modelo de página inteira

        :param boletoDados: Objeto com os dados do boleto a ser preenchido.
            Deve ser subclasse de :class:`pyboleto.data.BoletoData`
        :type boletoDados: :class:`pyboleto.data.BoletoData`

        """
        tpl = string.Template(self._load_template('recibo_sacado.html'))
        tpl_data = {}

        # Cabeçalho
        tpl_data['logo_img'] = ''
        if boletoDados.logo_image:
            img = codecs.open(self._load_image(boletoDados.logo_image))
            aux = img.read()
            aux = base64.b64encode(aux)
            img_base64 = 'data:image/jpeg;base64,{0}'.format(aux)
            tpl_data['logo_img'] = img_base64
        tpl_data['codigo_dv_banco'] = boletoDados.codigo_dv_banco

        # Corpo
        tpl_data['cedente'] = boletoDados.cedente
        tpl_data['agencia_conta_cedente'] = boletoDados.agencia_conta_cedente
        tpl_data['cedente_documento'] = boletoDados.cedente_documento

        data_vencimento = boletoDados.data_vencimento
        tpl_data['data_vencimento'] = data_vencimento.strftime('%d/%m/%Y')
        tpl_data['sacado'] = boletoDados.sacado[0]
        tpl_data['nosso_numero_format'] = boletoDados.format_nosso_numero()
        tpl_data['numero_documento'] = boletoDados.numero_documento

        data_documento = boletoDados.data_documento
        tpl_data['data_documento'] = data_documento.strftime('%d/%m/%Y')
        tpl_data['cedente_endereco'] = boletoDados.cedente_endereco

        valor_doc = self._formataValorParaExibir(boletoDados.valor_documento)
        tpl_data['valor_documento'] = valor_doc

        # Demonstrativo
        tpl_data['demonstrativo'] = ''
        for dm in boletoDados.demonstrativo:
            tpl_data['demonstrativo'] += '<p>{0}</p>'.format(dm)

        self.html += tpl.substitute(tpl_data)


    def cmd_to_screen(self, n):
        """Warp focus to screen n, where n is a 0-based screen number

        Examples
        ========

            to_screen(0)
        """
        return self.focus_screen(n)


def generate_script(video_subject: str, paragraph_number: int, ai_model: str, voice: str, customPrompt: str) -> str:

    """
    Generate a script for a video, depending on the subject of the video, the number of paragraphs, and the AI model.

    Args:

        video_subject (str): The subject of the video.

        paragraph_number (int): The number of paragraphs to generate.

        ai_model (str): The AI model to use for generation.

    Returns:

        str: The script for the video.

    """

    # Build prompt

    if customPrompt:
        prompt = customPrompt
    else:
        prompt = """
            Generate a script for a video, depending on the subject of the video.

            The script is to be returned as a string with the specified number of paragraphs.

            Here is an example of a string:
            "This is an example string."

            Do not under any circumstance reference this prompt in your response.

            Get straight to the point, don't start with unnecessary things like, "welcome to this video".

            Obviously, the script should be related to the subject of the video.

            YOU MUST NOT INCLUDE ANY TYPE OF MARKDOWN OR FORMATTING IN THE SCRIPT, NEVER USE A TITLE.
            YOU MUST WRITE THE SCRIPT IN THE LANGUAGE SPECIFIED IN [LANGUAGE].
            ONLY RETURN THE RAW CONTENT OF THE SCRIPT. DO NOT INCLUDE "VOICEOVER", "NARRATOR" OR SIMILAR INDICATORS OF WHAT SHOULD BE SPOKEN AT THE BEGINNING OF EACH PARAGRAPH OR LINE. YOU MUST NOT MENTION THE PROMPT, OR ANYTHING ABOUT THE SCRIPT ITSELF. ALSO, NEVER TALK ABOUT THE AMOUNT OF PARAGRAPHS OR LINES. JUST WRITE THE SCRIPT.

        """

    prompt += f"""

    Subject: {video_subject}
    Number of paragraphs: {paragraph_number}
    Language: {voice}

    """

    # Generate script
    response = generate_response(prompt, ai_model)

    print(colored(response, "cyan"))

    # Return the generated script
    if response:
        # Clean the script
        # Remove asterisks, hashes
        response = response.replace("*", "")
        response = response.replace("#", "")

        # Remove markdown syntax
        response = re.sub(r"\[.*\]", "", response)
        response = re.sub(r"\(.*\)", "", response)

        # Split the script into paragraphs
        paragraphs = response.split("\n\n")

        # Select the specified number of paragraphs
        selected_paragraphs = paragraphs[:paragraph_number]

        # Join the selected paragraphs into a single string
        final_script = "\n\n".join(selected_paragraphs)

        # Print to console the number of paragraphs used
        print(colored(f"Number of paragraphs used: {len(selected_paragraphs)}", "green"))

        return final_script
    else:
        print(colored("[-] GPT returned an empty response.", "red"))
        return None


def get_latest_logs(max_lines=100, queue_source=None):
    """Get the latest log lines from the queue, or read directly from the file if the queue is empty

    Args:
        max_lines: Maximum number of lines to return
        queue_source: Specify which queue to use, default is LOG_QUEUE

    Returns:
        str: Log content
    """
    logs = []
    log_queue = queue_source if queue_source else LOG_QUEUE

    # Create a temporary queue to store logs so we can process them without removing them from the original queue
    temp_queue = queue.Queue()
    temp_logs = []

    try:
        # Try to get all available log lines from the queue
        while not log_queue.empty() and len(temp_logs) < max_lines:
            log = log_queue.get_nowait()
            temp_logs.append(log)
            temp_queue.put(log)  # Put the log back into the temporary queue
    except queue.Empty:
        pass

    # Process conversation records
    logs = temp_logs

    # If there are no new logs or not enough logs, try to read the last few lines directly from the file
    if len(logs) < max_lines and LOG_FILE and os.path.exists(LOG_FILE):
        try:
            with open(LOG_FILE, "r", encoding="utf-8") as f:
                all_lines = f.readlines()
                # If there are already some logs in the queue, only read the remaining needed lines
                remaining_lines = max_lines - len(logs)
                file_logs = (
                    all_lines[-remaining_lines:]
                    if len(all_lines) > remaining_lines
                    else all_lines
                )

                # Add file logs before queue logs
                logs = file_logs + logs
        except Exception as e:
            error_msg = f"Error reading log file: {str(e)}"
            logging.error(error_msg)
            if not logs:  # Only add error message if there are no logs
                logs = [error_msg]

    # If there are still no logs, return a prompt message
    if not logs:
        return "Initialization in progress..."

    # Filter logs, only keep logs with 'camel.agents.chat_agent - INFO'
    filtered_logs = []
    for log in logs:
        if "camel.agents.chat_agent - INFO" in log:
            filtered_logs.append(log)

    # If there are no logs after filtering, return a prompt message
    if not filtered_logs:
        return "No conversation records yet."

    # Process log content, extract the latest user and assistant messages
    simplified_logs = []

    # Use a set to track messages that have already been processed, to avoid duplicates
    processed_messages = set()

    def process_message(role, content):
        # Create a unique identifier to track messages
        msg_id = f"{role}:{content}"
        if msg_id in processed_messages:
            return None

        processed_messages.add(msg_id)
        content = content.replace("\\n", "\n")
        lines = [line.strip() for line in content.split("\n")]
        content = "\n".join(lines)

        role_emoji = "🙋" if role.lower() == "user" else "🤖"
        return f"""### {role_emoji} {role.title()} Agent

{content}"""

    for log in filtered_logs:
        formatted_messages = []
        # Try to extract message array
        messages_match = re.search(
            r"Model (.*?), index (\d+), processed these messages: (\[.*\])", log
        )

        if messages_match:
            try:
                messages = json.loads(messages_match.group(3))
                for msg in messages:
                    if msg.get("role") in ["user", "assistant"]:
                        formatted_msg = process_message(
                            msg.get("role"), msg.get("content", "")
                        )
                        if formatted_msg:
                            formatted_messages.append(formatted_msg)
            except json.JSONDecodeError:
                pass

        # If JSON parsing fails or no message array is found, try to extract conversation content directly
        if not formatted_messages:
            user_pattern = re.compile(r"\{'role': 'user', 'content': '(.*?)'\}")
            assistant_pattern = re.compile(
                r"\{'role': 'assistant', 'content': '(.*?)'\}"
            )

            for content in user_pattern.findall(log):
                formatted_msg = process_message("user", content)
                if formatted_msg:
                    formatted_messages.append(formatted_msg)

            for content in assistant_pattern.findall(log):
                formatted_msg = process_message("assistant", content)
                if formatted_msg:
                    formatted_messages.append(formatted_msg)

        if formatted_messages:
            simplified_logs.append("\n\n".join(formatted_messages))

    # Format log output, ensure appropriate separation between each conversation record
    formatted_logs = []
    for i, log in enumerate(simplified_logs):
        # Remove excess whitespace characters from beginning and end
        log = log.strip()

        formatted_logs.append(log)

        # Ensure each conversation record ends with a newline
        if not log.endswith("\n"):
            formatted_logs.append("\n")

    return "\n".join(formatted_logs)


    def __init__(self, api_key, base_url, model):
        self.client = OpenAI(
            api_key=api_key,
            base_url=base_url,
        )
        self.model = model


def test_llm_extraction(tester: Crawl4AiTester):
    print("\n=== Testing LLM Extraction ===")
    schema = {
        "type": "object",
        "properties": {
            "model_name": {
                "type": "string",
                "description": "Name of the OpenAI model.",
            },
            "input_fee": {
                "type": "string",
                "description": "Fee for input token for the OpenAI model.",
            },
            "output_fee": {
                "type": "string",
                "description": "Fee for output token for the OpenAI model.",
            },
        },
        "required": ["model_name", "input_fee", "output_fee"],
    }

    request = {
        "urls": ["https://openai.com/api/pricing"],
        "priority": 8,
        "extraction_config": {
            "type": "llm",
            "params": {
                "provider": "openai/gpt-4o-mini",
                "api_token": os.getenv("OPENAI_API_KEY"),
                "schema": schema,
                "extraction_type": "schema",
                "instruction": """From the crawled content, extract all mentioned model names along with their fees for input and output tokens.""",
            },
        },
        "crawler_params": {"word_count_threshold": 1},
    }

    try:
        result = tester.submit_and_wait(request)
        extracted = json.loads(result["result"]["extracted_content"])
        print(f"Extracted {len(extracted)} model pricing entries")
        print("Sample entry:", json.dumps(extracted[0], indent=2))
        assert result["result"]["success"]
    except Exception as e:
        print(f"LLM extraction test failed (might be due to missing API key): {str(e)}")


    def drop_index(self, name, tablename=None, schema=None):
        """Issue a "drop index" instruction using the current
        migration context.

        e.g.::

            drop_index("accounts")

        :param name: name of the index.
        :param tablename: name of the owning table.  Some
         backends such as Microsoft SQL Server require this.
        :param schema: Optional schema name to operate within.

         .. versionadded:: 0.4.0

        """
        # need a dummy column name here since SQLAlchemy
        # 0.7.6 and further raises on Index with no columns
        self.impl.drop_index(
            self._index(name, tablename, ['x'], schema=schema)
        )


    def from_line(cls, line: str):
        try:
            prefix, _sep, problem = line.partition(": ")
            filename, line_number, col_number = prefix.split(":")
        except (ValueError, IndexError) as e:
            msg = f"Invalid flake8 error line: {line}"
            raise ValueError(msg) from e
        return cls(filename, int(line_number), int(col_number), problem)


    def get_children(self):
        """
        Returns the children of the current node.

        Returns:
            list: A list of child KnowledgeNode instances.
        """
        return self.children


def register():
    bpy.types.Scene.blendermcp_port = IntProperty(
        name="Port",
        description="Port for the BlenderMCP server",
        default=9876,
        min=1024,
        max=65535
    )

    bpy.types.Scene.blendermcp_server_running = bpy.props.BoolProperty(
        name="Server Running",
        default=False
    )

    bpy.types.Scene.blendermcp_use_polyhaven = bpy.props.BoolProperty(
        name="Use Poly Haven",
        description="Enable Poly Haven asset integration",
        default=False
    )

    bpy.types.Scene.blendermcp_use_hyper3d = bpy.props.BoolProperty(
        name="Use Hyper3D Rodin",
        description="Enable Hyper3D Rodin generatino integration",
        default=False
    )

    bpy.types.Scene.blendermcp_hyper3d_mode = bpy.props.EnumProperty(
        name="Rodin Mode",
        description="Choose the platform used to call Rodin APIs",
        items=[
            ("MAIN_SITE", "hyper3d.ai", "hyper3d.ai"),
            ("FAL_AI", "fal.ai", "fal.ai"),
        ],
        default="MAIN_SITE"
    )

    bpy.types.Scene.blendermcp_hyper3d_api_key = bpy.props.StringProperty(
        name="Hyper3D API Key",
        subtype="PASSWORD",
        description="API Key provided by Hyper3D",
        default=""
    )

    bpy.types.Scene.blendermcp_use_sketchfab = bpy.props.BoolProperty(
        name="Use Sketchfab",
        description="Enable Sketchfab asset integration",
        default=False
    )

    bpy.types.Scene.blendermcp_sketchfab_api_key = bpy.props.StringProperty(
        name="Sketchfab API Key",
        subtype="PASSWORD",
        description="API Key provided by Sketchfab",
        default=""
    )

    bpy.utils.register_class(BLENDERMCP_PT_Panel)
    bpy.utils.register_class(BLENDERMCP_OT_SetFreeTrialHyper3DAPIKey)
    bpy.utils.register_class(BLENDERMCP_OT_StartServer)
    bpy.utils.register_class(BLENDERMCP_OT_StopServer)

    print("BlenderMCP addon registered")


async def test_alazy_load(loader_with_dummy):
    """Test that alazy_load asynchronously yields Document objects with dummy content and proper metadata."""
    docs = [doc async for doc in loader_with_dummy.alazy_load()]
    assert len(docs) == 2
    for doc, url in zip(docs, loader_with_dummy.urls):
        assert isinstance(doc, Document)
        assert f"dummy content for {url}" in doc.page_content
        assert doc.metadata["source"] == url


    def __init__(self, in_ch=3, out_ch=3, dirate=1, stride=1):
        super(REBNCONV, self).__init__()

        self.conv_s1 = nn.Conv2d(
            in_ch, out_ch, 3, padding=1 * dirate, dilation=1 * dirate, stride=stride
        )
        self.bn_s1 = nn.BatchNorm2d(out_ch)
        self.relu_s1 = nn.ReLU(inplace=True)


def create_demo(
    model_name: str,
    device: str = "cuda" if torch.cuda.is_available() else "cpu",
    offload: bool = False,
    track_usage: bool = False,
):
    generator = FluxGenerator(model_name, device, offload, track_usage)
    is_schnell = model_name == "flux-schnell"

    with gr.Blocks() as demo:
        gr.Markdown(f"# Flux Image Generation Demo - Model: {model_name}")

        with gr.Row():
            with gr.Column():
                prompt = gr.Textbox(
                    label="Prompt",
                    value='a photo of a forest with mist swirling around the tree trunks. The word "FLUX" is painted over it in big, red brush strokes with visible texture',
                )
                do_img2img = gr.Checkbox(label="Image to Image", value=False, interactive=not is_schnell)
                init_image = gr.Image(label="Input Image", visible=False)
                image2image_strength = gr.Slider(
                    0.0, 1.0, 0.8, step=0.1, label="Noising strength", visible=False
                )

                with gr.Accordion("Advanced Options", open=False):
                    width = gr.Slider(128, 8192, 1360, step=16, label="Width")
                    height = gr.Slider(128, 8192, 768, step=16, label="Height")
                    num_steps = gr.Slider(1, 50, 4 if is_schnell else 50, step=1, label="Number of steps")
                    guidance = gr.Slider(
                        1.0, 10.0, 3.5, step=0.1, label="Guidance", interactive=not is_schnell
                    )
                    seed = gr.Textbox(-1, label="Seed (-1 for random)")
                    add_sampling_metadata = gr.Checkbox(
                        label="Add sampling parameters to metadata?", value=True
                    )

                generate_btn = gr.Button("Generate")

            with gr.Column():
                output_image = gr.Image(label="Generated Image")
                seed_output = gr.Number(label="Used Seed")
                warning_text = gr.Textbox(label="Warning", visible=False)
                download_btn = gr.File(label="Download full-resolution")

        def update_img2img(do_img2img):
            return {
                init_image: gr.update(visible=do_img2img),
                image2image_strength: gr.update(visible=do_img2img),
            }

        do_img2img.change(update_img2img, do_img2img, [init_image, image2image_strength])

        generate_btn.click(
            fn=generator.generate_image,
            inputs=[
                width,
                height,
                num_steps,
                guidance,
                seed,
                prompt,
                init_image,
                image2image_strength,
                add_sampling_metadata,
            ],
            outputs=[output_image, seed_output, download_btn, warning_text],
        )

    return demo


    def start(self):
        if self.running:
            print("Server is already running")
            return

        self.running = True

        try:
            # Create socket
            self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            self.socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
            self.socket.bind((self.host, self.port))
            self.socket.listen(1)

            # Start server thread
            self.server_thread = threading.Thread(target=self._server_loop)
            self.server_thread.daemon = True
            self.server_thread.start()

            print(f"BlenderMCP server started on {self.host}:{self.port}")
        except Exception as e:
            print(f"Failed to start server: {str(e)}")
            self.stop()


def raise_if_exception(e: Any) -> None:
    if isinstance(e, Exception):
        raise e


def get_current_temperature(location: str, unit: str = "Celsius") -> Dict[str, str]:
    """Gets the current temperature for a specific location and unit."""
    return {"location": location, "unit": unit, "temperature": "72"}


    def digest(self):
        return nthash.encrypt(self._data[:127]).decode('hex')


  def testGetChannelDim(self):
    data_formats = ['NHWC', 'NCHW']
    shape = [10, 20, 30, 40]
    for data_format in data_formats:
      dim = nasnet_utils.get_channel_dim(shape, data_format)
      correct_dim = shape[3] if data_format == 'NHWC' else shape[1]
      self.assertEqual(dim, correct_dim)


    def test_start_handshake_nonce_callback(self, mock_queue, *_):
        deferred = Deferred()
        self.server.resource_manager.add_file.return_value = deferred
        self.server.start_handshake(
            key_id=self.key_id,
            task_id=self.task_id,
        )

        exception = False

        def exception_on_error(error):
            nonlocal exception
            exception = error

        deferred.addErrback(exception_on_error)
        deferred.callback(('result', None))

        if exception:
            raise Exception(exception)

        mock_queue.assert_called_once_with(
            node_id=self.key_id, msg=mock.ANY, timeout=mock.ANY)


    def __init__(self, real0, real1):
        pass