Add type checking of ragstack-ragulate (#615)

Author: cbornet

libs/ragulate/colbert_chunk_size_and_k.py

@@ -3,13 +3,13 @@
 import os
 import time
 from pathlib import Path
-from typing import List
+from typing import Any, List
 
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain_community.document_loaders import UnstructuredFileLoader
 from langchain_core.output_parsers import StrOutputParser
 from langchain_core.prompts import ChatPromptTemplate
-from langchain_core.runnables import RunnablePassthrough
+from langchain_core.runnables import Runnable, RunnablePassthrough
 from langchain_openai import ChatOpenAI
 from ragstack_colbert import (
     CassandraDatabase,
@@ -74,7 +74,7 @@ def len_function(text: str) -> int:
     return len(tokenizer.tokenize(text))
 
 
-async def ingest(file_path: str, chunk_size: int, **_):
+async def ingest(file_path: str, chunk_size: int, **_: Any) -> None:
     doc_id = Path(file_path).name
 
     chunk_overlap = min(chunk_size / 4, 64)
@@ -134,9 +134,9 @@ async def ingest(file_path: str, chunk_size: int, **_):
     )
 
 
-def query_pipeline(k: int, chunk_size: int, **_):
+def query_pipeline(k: int, chunk_size: int, **_: Any) -> Runnable[Any, Any]:
     vector_store = get_lc_vector_store(chunk_size=chunk_size)
-    llm = ChatOpenAI(model_name=LLM_MODEL)
+    llm = ChatOpenAI(model=LLM_MODEL)
 
     # build a prompt
     prompt_template = """

libs/ragulate/open_ai_chunk_size_and_k.py

@@ -1,19 +1,20 @@
 # ruff: noqa: D103, INP001, T201
 import os
+from typing import Any
 
 from langchain_astradb import AstraDBVectorStore
 from langchain_community.document_loaders import UnstructuredFileLoader
 from langchain_core.output_parsers import StrOutputParser
 from langchain_core.prompts import ChatPromptTemplate
-from langchain_core.runnables import RunnablePassthrough
+from langchain_core.runnables import Runnable, RunnablePassthrough
 from langchain_openai import ChatOpenAI, OpenAIEmbeddings
 from langchain_text_splitters import RecursiveCharacterTextSplitter
 
 EMBEDDING_MODEL = "text-embedding-3-small"
 LLM_MODEL = "gpt-3.5-turbo"
 
 
-def get_vector_store(chunk_size: int):
+def get_vector_store(chunk_size: int) -> AstraDBVectorStore:
     return AstraDBVectorStore(
         embedding=OpenAIEmbeddings(model=EMBEDDING_MODEL),
         collection_name=f"chunk_size_{chunk_size}",
@@ -22,7 +23,7 @@ def get_vector_store(chunk_size: int):
     )
 
 
-def ingest(file_path: str, chunk_size: int, **_):
+def ingest(file_path: str, chunk_size: int, **_: Any) -> None:
     vector_store = get_vector_store(chunk_size=chunk_size)
 
     chunk_overlap = min(chunk_size / 4, 64)
@@ -40,9 +41,9 @@ def ingest(file_path: str, chunk_size: int, **_):
     vector_store.add_documents(split_docs)
 
 
-def query_pipeline(k: int, chunk_size: int, **_):
+def query_pipeline(k: int, chunk_size: int, **_: Any) -> Runnable[Any, Any]:
     vector_store = get_vector_store(chunk_size=chunk_size)
-    llm = ChatOpenAI(model_name=LLM_MODEL)
+    llm = ChatOpenAI(model=LLM_MODEL)
 
     # build a prompt
     prompt_template = """

libs/ragulate/poetry.lock

@@ -35,6 +35,9 @@ langchain-community = "0.0.38"
 langchain-core = "0.1.52"
 langchain-openai = "0.1.3"
 pytest = "^8.2.2"
+mypy = "^1.11.0"
+types-pyyaml = "^6.0.1"
+types-aiofiles = "^23.2.0.0"
 
 [build-system]
 requires = ["poetry-core", "wheel"]
@@ -44,3 +47,20 @@ build-backend = "poetry.core.masonry.api"
 ragulate = "ragstack_ragulate.cli:main"
 test_unit = "scripts.test_unit_runner:main"
 test_integration = "scripts.test_integration_runner:main"
+
+[tool.mypy]
+disallow_any_generics = true
+disallow_incomplete_defs = true
+disallow_untyped_calls = true
+disallow_untyped_decorators = true
+disallow_untyped_defs = true
+follow_imports = "normal"
+ignore_missing_imports = true
+no_implicit_reexport = true
+show_error_codes = true
+show_error_context = true
+strict_equality = true
+strict_optional = true
+warn_redundant_casts = true
+warn_return_any = true
+warn_unused_ignores = true

libs/ragulate/pyproject.toml

@@ -1,11 +1,10 @@
-from typing import List
+from typing import Any, Dict, List, Tuple
 
 import matplotlib.pyplot as plt
 import numpy as np
 import pandas as pd
 import plotly.graph_objects as go
 import seaborn as sns
-from pandas import DataFrame
 from plotly.io import write_image
 
 from .utils import get_tru
@@ -14,7 +13,7 @@
 class Analysis:
     """Analysis class."""
 
-    def get_all_data(self, recipes: List[str]) -> DataFrame:
+    def get_all_data(self, recipes: List[str]) -> Tuple[pd.DataFrame, List[str]]:
         """Get all data from the recipes."""
         df_all = pd.DataFrame()
 
@@ -55,9 +54,11 @@ def get_all_data(self, recipes: List[str]) -> DataFrame:
 
         return reset_df, list(set(all_metrics))
 
-    def calculate_statistics(self, df: pd.DataFrame, metrics: list):
+    def calculate_statistics(
+        self, df: pd.DataFrame, metrics: List[str]
+    ) -> Dict[str, Any]:
         """Calculate statistics."""
-        stats = {}
+        stats: Dict[str, Any] = {}
         for recipe in df["recipe"].unique():
             stats[recipe] = {}
             for metric in metrics:
@@ -76,7 +77,7 @@ def calculate_statistics(self, df: pd.DataFrame, metrics: list):
                     }
         return stats
 
-    def output_box_plots_by_dataset(self, df: DataFrame, metrics: List[str]):
+    def output_box_plots_by_dataset(self, df: pd.DataFrame, metrics: List[str]) -> None:
         """Output box plots by dataset."""
         stats = self.calculate_statistics(df, metrics)
         recipes = sorted(df["recipe"].unique(), key=lambda x: x.lower())
@@ -101,7 +102,6 @@ def output_box_plots_by_dataset(self, df: DataFrame, metrics: List[str]):
                 q1 = []
                 median = []
                 q3 = []
-                mean = []
                 low = []
                 high = []
                 for metric in metrics:
@@ -120,7 +120,7 @@ def output_box_plots_by_dataset(self, df: DataFrame, metrics: List[str]):
                         q1=q1,
                         median=median,
                         q3=q3,
-                        mean=mean,
+                        mean=[],
                         lowerfence=low,
                         upperfence=high,
                         name=recipe,
@@ -159,7 +159,9 @@ def output_box_plots_by_dataset(self, df: DataFrame, metrics: List[str]):
 
             write_image(fig, f"./{dataset}_box_plot.png")
 
-    def output_histograms_by_dataset(self, df: pd.DataFrame, metrics: List[str]):
+    def output_histograms_by_dataset(
+        self, df: pd.DataFrame, metrics: List[str]
+    ) -> None:
         """Output histograms by dataset."""
         # Append "latency" to the metrics list
         metrics.append("latency")
@@ -184,7 +186,7 @@ def output_histograms_by_dataset(self, df: pd.DataFrame, metrics: List[str]):
             sns.set_theme(style="darkgrid")
 
             # Custom function to set bin ranges and filter invalid values
-            def custom_hist(data, **kws):
+            def custom_hist(data: Dict[str, Any], **kws: Any) -> None:
                 metric = data["metric"].iloc[0]
                 data = data[
                     np.isfinite(data["value"])
@@ -251,12 +253,12 @@ def custom_hist(data, **kws):
             # Close the plot to avoid displaying it
             plt.close()
 
-    def compare(self, recipes: List[str], output: str):
+    def compare(self, recipes: List[str], output: str = "box-plots") -> None:
         """Compare results from 2 (or more) recipes."""
         df, metrics = self.get_all_data(recipes=recipes)
         if output == "box-plots":
             self.output_box_plots_by_dataset(df=df, metrics=metrics)
         elif output == "histogram-grid":
             self.output_histograms_by_dataset(df=df, metrics=metrics)
         else:
-            raise ValueError
+            raise ValueError(f"Invalid output type: {output}")

libs/ragulate/ragstack_ragulate/analysis.py

@@ -1,11 +1,12 @@
-from typing import List, Optional
+from argparse import ArgumentParser, _SubParsersAction
+from typing import Any, List
 
 from ragstack_ragulate.analysis import Analysis
 
 from .utils import remove_sqlite_extension
 
 
-def setup_compare(subparsers):
+def setup_compare(subparsers: _SubParsersAction[ArgumentParser]) -> None:
     """Setup the compare command."""
     compare_parser = subparsers.add_parser(
         "compare", help="Compare results from 2 (or more) recipes"
@@ -30,9 +31,9 @@ def setup_compare(subparsers):
 
 def call_compare(
     recipe: List[str],
-    output: Optional[str] = "box-plots",
-    **_,
-):
+    output: str = "box-plots",
+    **_: Any,
+) -> None:
     """Compare results from 2 (or more) recipes."""
     analysis = Analysis()
 

libs/ragulate/ragstack_ragulate/cli_commands/compare.py

@@ -1,9 +1,12 @@
+from argparse import ArgumentParser, _SubParsersAction
+from typing import Any
+
 from ragstack_ragulate.dashboard import run_dashboard
 
 from .utils import remove_sqlite_extension
 
 
-def setup_dashboard(subparsers):
+def setup_dashboard(subparsers: _SubParsersAction[ArgumentParser]) -> None:
     """Setup the dashboard command."""
     dashboard_parser = subparsers.add_parser(
         "dashboard",
@@ -29,8 +32,8 @@ def setup_dashboard(subparsers):
 def call_dashboard(
     recipe: str,
     port: int,
-    **_,
-):
+    **_: Any,
+) -> None:
     """Runs the TruLens dashboard."""
     recipe_name = remove_sqlite_extension(recipe)
     run_dashboard(recipe_name=recipe_name, port=port)

libs/ragulate/ragstack_ragulate/cli_commands/dashboard.py

@@ -1,7 +1,10 @@
+from argparse import ArgumentParser, _SubParsersAction
+from typing import Any
+
 from ragstack_ragulate.datasets import get_dataset
 
 
-def setup_download(subparsers):
+def setup_download(subparsers: _SubParsersAction[ArgumentParser]) -> None:
     """Setup the download command."""
     download_parser = subparsers.add_parser("download", help="Download a dataset")
     download_parser.add_argument(
@@ -22,7 +25,7 @@ def setup_download(subparsers):
     download_parser.set_defaults(func=lambda args: call_download(**vars(args)))
 
 
-def call_download(dataset_name: str, kind: str, **_):
+def call_download(dataset_name: str, kind: str, **_: Any) -> None:
     """Download a dataset."""
     dataset = get_dataset(name=dataset_name, kind=kind)
     dataset.download_dataset()

libs/ragulate/ragstack_ragulate/cli_commands/download.py

@@ -1,11 +1,12 @@
-from typing import List
+from argparse import ArgumentParser, _SubParsersAction
+from typing import Any, List
 
 from ragstack_ragulate.datasets import find_dataset
 from ragstack_ragulate.pipelines import IngestPipeline
 from ragstack_ragulate.utils import convert_vars_to_ingredients
 
 
-def setup_ingest(subparsers):
+def setup_ingest(subparsers: _SubParsersAction[ArgumentParser]) -> None:
     """Setup the ingest command."""
     ingest_parser = subparsers.add_parser("ingest", help="Run an ingest pipeline")
     ingest_parser.add_argument(
@@ -58,8 +59,8 @@ def call_ingest(
         var_name: List[str],
         var_value: List[str],
         dataset: List[str],
-        **_,
-    ):
+        **_: Any,
+    ) -> None:
         """Run an ingest pipeline."""
         datasets = [find_dataset(name=name) for name in dataset]
 

libs/ragulate/ragstack_ragulate/cli_commands/ingest.py

@@ -1,11 +1,12 @@
-from typing import List
+from argparse import ArgumentParser, _SubParsersAction
+from typing import Any, List
 
 from ragstack_ragulate.datasets import find_dataset
 from ragstack_ragulate.pipelines import QueryPipeline
 from ragstack_ragulate.utils import convert_vars_to_ingredients
 
 
-def setup_query(subparsers):
+def setup_query(subparsers: _SubParsersAction[ArgumentParser]) -> None:
     """Setup the query command."""
     query_parser = subparsers.add_parser("query", help="Run a query pipeline")
     query_parser.add_argument(
@@ -110,8 +111,8 @@ def call_query(
         restart: bool,
         provider: str,
         model: str,
-        **_,
-    ):
+        **_: Any,
+    ) -> None:
         """Run a query pipeline."""
         if sample <= 0.0 or sample > 1.0:
             raise ValueError("Sample percent must be between 0 and 1")