diff --git a/python/huggingFace/src/com/intellij/python/community/impl/huggingFace/api/HuggingFaceApi.kt b/python/huggingFace/src/com/intellij/python/community/impl/huggingFace/api/HuggingFaceApi.kt index 6ac9ca50a9e6..0b82cd691d35 100644 --- a/python/huggingFace/src/com/intellij/python/community/impl/huggingFace/api/HuggingFaceApi.kt +++ b/python/huggingFace/src/com/intellij/python/community/impl/huggingFace/api/HuggingFaceApi.kt @@ -10,7 +10,9 @@ import com.intellij.python.community.impl.huggingFace.cache.HuggingFaceMdCacheEn import com.intellij.python.community.impl.huggingFace.cache.HuggingFaceMdCardsCache import com.intellij.python.community.impl.huggingFace.documentation.HuggingFaceDocumentationPlaceholdersUtil import com.intellij.python.community.impl.huggingFace.documentation.HuggingFaceReadmeCleaner -import com.intellij.python.community.impl.huggingFace.service.HuggingFaceSafeExecutor +import com.intellij.python.community.impl.huggingFace.service.HuggingFaceCoroutine +import kotlinx.coroutines.launch +import kotlinx.coroutines.withContext import org.jetbrains.annotations.ApiStatus import org.jetbrains.annotations.Nls import java.time.Instant @@ -20,9 +22,8 @@ object HuggingFaceApi { private val nextLinkRegex = Regex("""<(.+)>; rel="next"""") fun fillCacheWithBasicApiData(endpoint: HuggingFaceEntityKind, cache: HuggingFaceCache, maxCount: Int) { - val executor = HuggingFaceSafeExecutor.instance - executor.asyncSuspend("FetchHuggingFace$endpoint") { + HuggingFaceCoroutine.Utils.ioScope.launch { var nextPageUrl: String? = HuggingFaceURLProvider.fetchApiDataUrl(endpoint).toString() while (nextPageUrl != null && cache.getCacheSize() < maxCount) { @@ -38,20 +39,17 @@ object HuggingFaceApi { } } - suspend fun performSearch(query: String, tags: String?, sortKey: HuggingFaceModelSortKey): Map { + suspend fun performSearch(query: String, tags: String?, sortKey: HuggingFaceModelSortKey + ): Map = withContext(HuggingFaceCoroutine.Utils.ioScope.coroutineContext) { val queryUrl = HuggingFaceURLProvider.getModelSearchQuery(query, tags, sortKey) - val executor = HuggingFaceSafeExecutor.instance - return executor.asyncSuspend("SearchHuggingFace") { - // todo: add pagination - but only if needed (user scrolls down) - val nextPageUrl: String = queryUrl.toString() - val response = HuggingFaceHttpClient.downloadContentAndHeaders(nextPageUrl) - .getOrDefault(HfHttpResponseWithHeaders(null, null)) - response.content?.let { - val dataMap = parseBasicEntityData(HuggingFaceEntityKind.MODEL, it) - dataMap - } ?: emptyMap() - }.await() + val nextPageUrl: String = queryUrl.toString() + val response = HuggingFaceHttpClient.downloadContentAndHeaders(nextPageUrl) + .getOrDefault(HfHttpResponseWithHeaders(null, null)) + response.content?.let { + val dataMap = parseBasicEntityData(HuggingFaceEntityKind.MODEL, it) + dataMap + } ?: emptyMap() } private fun extractNextPageUrl(linkHeader: String?): String? { diff --git a/python/huggingFace/src/com/intellij/python/community/impl/huggingFace/api/HuggingFaceHttpClient.kt b/python/huggingFace/src/com/intellij/python/community/impl/huggingFace/api/HuggingFaceHttpClient.kt index 54727c5cefc9..be979457b279 100644 --- a/python/huggingFace/src/com/intellij/python/community/impl/huggingFace/api/HuggingFaceHttpClient.kt +++ b/python/huggingFace/src/com/intellij/python/community/impl/huggingFace/api/HuggingFaceHttpClient.kt @@ -2,6 +2,7 @@ package com.intellij.python.community.impl.huggingFace.api import com.intellij.openapi.diagnostic.thisLogger +import com.intellij.util.concurrency.annotations.RequiresBackgroundThread import com.intellij.util.io.HttpRequests import org.jetbrains.annotations.ApiStatus import java.io.IOException @@ -32,6 +33,7 @@ object HuggingFaceHttpClient { } } + @RequiresBackgroundThread fun downloadFile(url: String): Result = runCatching { HttpRequests.request(url).readString() @@ -39,6 +41,7 @@ object HuggingFaceHttpClient { thisLogger().warn("Failed to download file: $url", it) } + @RequiresBackgroundThread fun downloadContentAndHeaders(url: String): Result { return runCatching { val connection = URL(url).openConnection() as HttpURLConnection diff --git a/python/huggingFace/src/com/intellij/python/community/impl/huggingFace/service/HuggingFaceCoroutine.kt b/python/huggingFace/src/com/intellij/python/community/impl/huggingFace/service/HuggingFaceCoroutine.kt new file mode 100644 index 000000000000..8dccd91ddf8b --- /dev/null +++ b/python/huggingFace/src/com/intellij/python/community/impl/huggingFace/service/HuggingFaceCoroutine.kt @@ -0,0 +1,26 @@ +// Copyright 2000-2024 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license. +package com.intellij.python.community.impl.huggingFace.service + +import com.intellij.openapi.application.ApplicationManager +import com.intellij.openapi.application.EDT +import com.intellij.openapi.components.Service +import com.intellij.openapi.components.service +import com.intellij.platform.util.coroutines.childScope +import kotlinx.coroutines.CoroutineScope +import kotlinx.coroutines.Dispatchers +import org.jetbrains.annotations.ApiStatus + +@Service +@ApiStatus.Internal +class HuggingFaceCoroutine(coroutineScope: CoroutineScope) { + val ioScope = coroutineScope.childScope(context = Dispatchers.IO) + val edtScope = coroutineScope.childScope(context = Dispatchers.EDT) + + object Utils { + val ioScope: CoroutineScope + get() = ApplicationManager.getApplication().service().ioScope + + val edtScope: CoroutineScope + get() = ApplicationManager.getApplication().service().edtScope + } +} diff --git a/python/huggingFace/src/com/intellij/python/community/impl/huggingFace/service/HuggingFaceExecutorService.kt b/python/huggingFace/src/com/intellij/python/community/impl/huggingFace/service/HuggingFaceExecutorService.kt deleted file mode 100644 index a7a538f0d1bd..000000000000 --- a/python/huggingFace/src/com/intellij/python/community/impl/huggingFace/service/HuggingFaceExecutorService.kt +++ /dev/null @@ -1,52 +0,0 @@ -// Copyright 2000-2024 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license. -package com.intellij.python.community.impl.huggingFace.service - -import com.intellij.openapi.Disposable -import com.intellij.openapi.components.Service -import com.intellij.openapi.components.service -import com.intellij.openapi.diagnostic.Logger -import com.intellij.platform.util.coroutines.childScope -import kotlinx.coroutines.* -import org.jetbrains.annotations.ApiStatus -import kotlin.time.Duration - -@ApiStatus.Internal -@Service(Service.Level.APP) -class HuggingFaceExecutorService(val coroutineScope: CoroutineScope) - -@ApiStatus.Internal -@Service(Service.Level.APP) -class HuggingFaceSafeExecutor(private val coroutineScope: CoroutineScope, private val defaultTimeout: Duration): Disposable { - - constructor(coroutineScope: CoroutineScope) : this(coroutineScope, Duration.INFINITE) - - override fun dispose() { - coroutineScope.cancel("$SERVICE_NAME dispose") - } - - fun asyncSuspend(taskName: String, body: suspend CoroutineScope.() -> T): Deferred { - return coroutineScope.async(context = CoroutineName(taskName)) { - withTimeout(defaultTimeout) { body() } - } - } - - companion object { - @OptIn(ExperimentalCoroutinesApi::class) - private fun getInstance( - name: String = SERVICE_NAME, - defaultTimeout: Duration = Duration.INFINITE, - parallelism: Int = 10 - ): HuggingFaceSafeExecutor { - - val childScope = service().coroutineScope.childScope( - Dispatchers.IO.limitedParallelism(parallelism) + CoroutineName(name)) - - return HuggingFaceSafeExecutor(childScope, defaultTimeout) - } - - val instance by lazy { getInstance() } - - const val SERVICE_NAME = "HuggingFaceSafeExecutor" - val LOG = Logger.getInstance(HuggingFaceSafeExecutor::class.java) - } -}