download_ollama.ps1

# 灵活的Ollama模型下载脚本
param(
    [Parameter(Mandatory = $false, Position = 0)]
    [string]$ModelName
)

# 检查是否提供了模型名称参数
if (-not $ModelName) {
    Write-Host "请提供要下载的模型名称，例如：`n  .\download_ollama.ps1 deepseek-r1:7b"
    exit 1
}

function Find-ProcessWithArgs {
    [CmdletBinding()]
    param(
        [Parameter(Mandatory = $false, Position = 0)]
        [string]$ProcessName = '*',
        [Parameter(Mandatory = $false, Position = 1, ValueFromRemainingArguments = $true)]
        [string[]]$ArgumentParts
    )

    # 获取所有进程及其命令行信息
    $processes = Get-CimInstance -ClassName Win32_Process | 
        Select-Object ProcessId, Name, @{
            Name       = 'Arguments';
            Expression = {
                if ($_.CommandLine) {
                    # 移除命令行开头的路径/可执行文件
                    $_.CommandLine -replace '^(".*?"|\S+)\s*', ''
                }
            }
        }

    # 过滤进程名称
    if ($ProcessName -ne '*') {
        $processes = $processes | Where-Object { $_.Name -like "*$ProcessName*" }
    }

    # 过滤命令行参数
    if ($ArgumentParts.Count -gt 0) {
        $processes = $processes | Where-Object {
            $matched = $true
            foreach ($part in $ArgumentParts) {
                if (-not ($_.Arguments -like "*$part*")) {
                    $matched = $false
                    break
                }
            }
            $matched
        }
    }

    # 返回结果
    $processes
}

function Stop-OllamaProcess {
    param([string]$ModelName)
    
    $processes = Find-ProcessWithArgs ollama pull $ModelName
    foreach ($proc in $processes) {
        try {
            Stop-Process -Id $proc.ProcessId -Force -ErrorAction Stop
            Write-Host "已终止进程 ID: $($proc.ProcessId)"
        }
        catch {
            Write-Host "终止进程 ID: $($proc.ProcessId) 失败：$($_.Exception.Message)"
        }
    }
}

function Is-OllamaProcessRunning {
    param([string]$ModelName)
    
    $processes = Find-ProcessWithArgs ollama pull $ModelName
    return $processes.Count -gt 0
}

$downloadComplete = $false
while (-not $downloadComplete) {
    # 先清理可能遗留的进程
    Stop-OllamaProcess -ModelName $ModelName
    
    Write-Host "开始下载模型：$ModelName ..."
    # 启动下载命令（不阻塞当前脚本）
    Start-Process -NoNewWindow -FilePath "ollama" -ArgumentList "pull", $ModelName
    
    # 等待 60 秒后再检查进程状态
    Start-Sleep -Seconds 60
    
    if (Is-OllamaProcessRunning -ModelName $ModelName) {
        Write-Host "$ModelName 下载仍在运行，继续等待..."
    }
    else {
        Write-Host "$ModelName 下载已完成。"
        $downloadComplete = $true
    }
}

Write-Host "下载完成，脚本退出。"