copilot fixes

Author: fabnemEPFL

docs/source/advanced_usage/profiler.md

@@ -119,7 +119,7 @@ If you only need to measure execution time (wall clock), use `time_function` or
 ```python
 from mmore.profiler import time_function, time_context
 
-@time_function
+@time_function()
 def quick_check():
     pass
 

pyrightconfig.json

@@ -0,0 +1,16 @@
+{
+  "include": ["src/mmore", "tests"],
+  "exclude": ["build", "dist", "**/__pycache__", "*.egg-info"],
+  "executionEnvironments": [
+    {
+      "root": "src/mmore",
+      "extraPaths": ["src"]
+    },
+    {
+      "root": "tests",
+      "extraPaths": ["src"]
+    }
+  ],
+  "pythonVersion": "3.10",
+  "typeCheckingMode": "standard"
+}

src/mmore/process/post_processor/__init__.py

@@ -33,13 +33,15 @@ def load_postprocessor(config: BasePostProcessorConfig) -> BasePostProcessor:
         from .translator import TranslatorConfig, TranslatorPostProcessor
 
         config_translator = load_config(config.args, TranslatorConfig)
-        return TranslatorPostProcessor.from_config(config_translator)
+        return cast(
+            BasePostProcessor, TranslatorPostProcessor.from_config(config_translator)
+        )
 
     elif config.type == "metafuse":
         from .metafuse import MetaDataInfusor, MetaDataInfusorConfig
 
         config_metafuse = load_config(config.args, MetaDataInfusorConfig)
-        return MetaDataInfusor.from_config(config_metafuse)
+        return cast(BasePostProcessor, MetaDataInfusor.from_config(config_metafuse))
 
     else:
         raise ValueError(f"Unrecognized postprocessor type: {config.type}")

src/mmore/process/processors/media_processor.py

@@ -161,7 +161,7 @@ def _extract_video_frames(file_path: str) -> List[Image.Image]:
                     for i in range(num_thumbnails):
                         t = min(i * sample_rate, duration - 0.1)
                         frame = clip.get_frame(t)
-                        image = Image.fromarray(cast(np.ndarray, frame).convert("RGB"))
+                        image = Image.fromarray(cast(np.ndarray, frame)).convert("RGB")
                         images.append(image)
                 logger.info(f"Extracted {len(images)} images from {file_path}.")
             except Exception as e:

src/mmore/profiler.py

@@ -10,10 +10,12 @@
 from contextlib import contextmanager
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Callable, Optional
+from typing import Callable, Optional, TypeVar
 
 logger = logging.getLogger(__name__)
 
+R = TypeVar("R")
+
 
 @dataclass
 class ProfilingConfig:
@@ -204,20 +206,19 @@ def time_context(name: str, log: bool = True):
             logger.info(f"⏱️  {name} took {elapsed:.4f} seconds")
 
 
-def time_function(func: Optional[Callable] = None, log: bool = True):
+def time_function(log: bool = True) -> Callable[[Callable[..., R]], Callable[..., R]]:
     """Decorator to time a function execution.
 
     Args:
-        func: Function to decorate (if used as decorator without parentheses)
         log: Whether to log the timing result
 
     Returns:
         Decorated function
     """
 
-    def decorator(f: Callable) -> Callable:
+    def decorator(f: Callable[..., R]) -> Callable[..., R]:
         @functools.wraps(f)
-        def wrapper(*args, **kwargs):
+        def wrapper(*args: object, **kwargs: object) -> R:
             start_time = time.time()
             try:
                 result = f(*args, **kwargs)
@@ -229,10 +230,7 @@ def wrapper(*args, **kwargs):
 
         return wrapper
 
-    if func is None:
-        return decorator
-    else:
-        return decorator(func)
+    return decorator
 
 
 class Profiler:

src/mmore/run_ragcli.py

@@ -14,10 +14,9 @@
 )
 
 from mmore.profiler import enable_profiling_from_env, profile_function
-
-from .rag.pipeline import RAGPipeline
-from .run_rag import RAGInferenceConfig
-from .utils import load_config
+from mmore.rag.pipeline import RAGPipeline
+from mmore.run_rag import RAGInferenceConfig
+from mmore.utils import load_config
 
 
 class RagCLI:

src/mmore/run_websearch.py

@@ -16,13 +16,12 @@
 from pydantic import BaseModel, Field
 
 from mmore.profiler import enable_profiling_from_env, profile_function
-
-from .rag.pipeline import RAGPipeline
-from .run_rag import APIConfig, LocalConfig, RAGInferenceConfig
-from .utils import load_config
-from .websearchRAG.config import WebsearchConfig
-from .websearchRAG.logging_config import logger
-from .websearchRAG.pipeline import WebsearchPipeline
+from mmore.rag.pipeline import RAGPipeline
+from mmore.run_rag import APIConfig, LocalConfig, RAGInferenceConfig
+from mmore.utils import load_config
+from mmore.websearchRAG.config import WebsearchConfig
+from mmore.websearchRAG.logging_config import logger
+from mmore.websearchRAG.pipeline import WebsearchPipeline
 
 load_dotenv()
 

src/mmore/tui/config_builder.py

@@ -802,7 +802,7 @@ def pick_or_build_config(
                     )
                 )
                 path = picked
-                if choice == "edit":
+                if choice == "edit" and path is not None:
                     _edit_config(path)
 
         if choice == "manual":

tests/conftest.py

@@ -4,7 +4,7 @@
 import pytest
 from langchain_milvus.utils.sparse import BaseSparseEmbedding
 
-from mmore.type import MultimodalSample
+from mmore.type import DocumentMetadata, MultimodalSample
 
 
 class FakeSparseEmbedding(BaseSparseEmbedding):
@@ -23,21 +23,21 @@ def embed_documents(self, texts: List[str]) -> List[Dict[int, float]]:
         document_id="doc-1",
         text="Paris is the capital of France.",
         modalities=[],
-        metadata={},
+        metadata=DocumentMetadata(),
     ),
     MultimodalSample(
         id="doc-2",
         document_id="doc-2",
         text="The Eiffel Tower stands 330 metres tall.",
         modalities=[],
-        metadata={"author": "Alice"},
+        metadata=DocumentMetadata(extra={"author": "Alice"}),
     ),
     MultimodalSample(
         id="doc-3",
         document_id="doc-3",
         text="Milvus is an open-source vector database.",
         modalities=[],
-        metadata={},
+        metadata=DocumentMetadata(),
     ),
 ]
 

tests/test_indexer.py

@@ -11,7 +11,7 @@
 from mmore.index.indexer import Indexer
 from mmore.rag.model import DenseModelConfig, SparseModelConfig
 from mmore.run_index import index
-from mmore.type import MultimodalSample
+from mmore.type import DocumentMetadata, MultimodalSample
 
 
 @pytest.fixture
@@ -98,7 +98,7 @@ def test_indexer_idempotent_collection(mock_sparse, tmp_path):
             document_id="doc-4",
             text="A fourth document.",
             modalities=[],
-            metadata={},
+            metadata=DocumentMetadata(),
         )
     ]
 
@@ -153,7 +153,7 @@ def test_indexer_error_on_missing_collection(mock_sparse, tmp_path):
     )
 
     doc = MultimodalSample(
-        id="x", document_id="x", text="test", modalities=[], metadata={}
+        id="x", document_id="x", text="test", modalities=[], metadata=DocumentMetadata()
     )
     with pytest.raises(Exception):
         # Bypasses index_documents (which creates the collection) and inserts directly