wip

Author: quitrk

requirements-vllm.txt

@@ -184,7 +184,7 @@ torchaudio==2.5.1 ; python_version >= "3.11" and python_version < "3.12"
 torchvision==0.20.1 ; python_version >= "3.11" and python_version < "3.12"
 tqdm==4.67.1 ; python_version >= "3.11" and python_version < "3.12"
 transformers==4.50.0 ; python_version >= "3.11" and python_version < "3.12"
-triton==3.1.0 ; python_version >= "3.11" and python_version < "3.12" and (platform_system == "Linux" or platform_system == "linux") and platform_machine == "x86_64"
+triton==3.1.0 ; python_version >= "3.11" and python_version < "3.12" and (platform_system == "linux" or platform_system == "Linux") and platform_machine == "x86_64"
 typing-extensions==4.12.2 ; python_version >= "3.11" and python_version < "3.12"
 typing-inspect==0.9.0 ; python_version >= "3.11" and python_version < "3.12"
 urllib3==2.3.0 ; python_version >= "3.11" and python_version < "3.12"

skynet/modules/ttt/rag/vector_store.py

@@ -126,7 +126,7 @@ async def workflow(self, store_id: str, files: list[str], urls: list[str], max_d
             zip_files = [f for f in files if f.endswith('.zip')]
             if zip_files:
                 files = [f for f in files if f not in zip_files]
-                files.extend(extract_files(zip_files, self.get_temp_folder(store_id), min_size_kb=1))
+                files.extend(await extract_files(zip_files, self.get_temp_folder(store_id), min_size_kb=1))
 
             files = [f for f in files if any(f.endswith(ext) for ext in supported_files)]
 

skynet/modules/ttt/rag/zip_extractor/main.py

@@ -1,13 +1,13 @@
+import asyncio
 from pathlib import Path
 from shutil import unpack_archive
 from typing import Optional
 
 
-def extract_files(compressed_files: list[str], destination_folder: str, min_size_kb: Optional[int] = None) -> list[str]:
-    """
-    Extract all files from the given list of compressed files.
-    """
-
+def _extract_files_sync(
+    compressed_files: list[str], destination_folder: str, min_size_kb: Optional[int] = None
+) -> list[str]:
+    """Synchronous implementation of file extraction"""
     extracted_files = []
 
     for compressed_file in compressed_files:
@@ -24,3 +24,13 @@ def extract_files(compressed_files: list[str], destination_folder: str, min_size
         extracted_files = [f for f in extracted_files if Path(f).stat().st_size >= min_size_kb * 1024]
 
     return extracted_files
+
+
+async def extract_files(
+    compressed_files: list[str], destination_folder: str, min_size_kb: Optional[int] = None
+) -> list[str]:
+    """
+    Extract all files from the given list of compressed files.
+    Runs the entire extraction process in a separate thread.
+    """
+    return await asyncio.to_thread(_extract_files_sync, compressed_files, destination_folder, min_size_kb)

skynet/modules/ttt/rag/zip_extractor/main_test.py

@@ -2,11 +2,14 @@
 import shutil
 from pathlib import Path
 
+import pytest
+
 from skynet.modules.ttt.rag.zip_extractor.main import extract_files
 
 
 class TestExtractFiles:
-    def test_extract_files(self):
+    @pytest.mark.asyncio
+    async def test_extract_files(self):
         """
         Test extract_files function retrieves all files in all folders.
         """
@@ -33,7 +36,7 @@ def test_extract_files(self):
         for folder, zip_file in zip(folders, zip_files):
             shutil.make_archive(zip_file, 'zip', folder)
 
-        extracted_files = extract_files([f'{file}.zip' for file in zip_files], 'temp')
+        extracted_files = await extract_files([f'{file}.zip' for file in zip_files], 'temp')
         expected_files = [
             'temp/archives/test1/subfolder_0/file_0.txt',
             'temp/archives/test1/subfolder_0/file_1.txt',