Refactor

Signed-off-by: András Schmelczer <andras@schmelczer.dev>

example/main_batch.py

@@ -4,7 +4,7 @@ from random import shuffle
 from devtools import debug
 from predict_domain import predict_domain
 
-from good_ai import process_batch, serve
+from good_ai import process_batch
 
 if __name__ == "__main__":
     with open(".cache/data-1/s2-corpus-323.json") as f:

example/main_service.py

@@ -1,10 +1,9 @@
-import json
+#!/usr/bin/env python3
-from random import shuffle
+
 
-from devtools import debug
 from predict_domain import predict_domain
 
-from good_ai import process_batch, serve
+from good_ai import serve
 
 if __name__ == "__main__":
     serve(predict_domain)

example/predict_domain.py

@@ -1,23 +1,22 @@
 import re
 from typing import Dict, Iterable, List
 
-
 from config import model_key
 from models import DomainPrediction
 from preprocess import preprocess
 from sklearn.pipeline import Pipeline
 
-from good_ai import use_model, log_argument, log_metric
+from good_ai import log_argument, log_metric, use_model
 from good_ai.utilities.clean import clean
 
 
 @use_model(model_key, version="latest")
-@log_argument('text', validator=lambda t: len(t) > 0)
+@log_argument("text", validator=lambda t: len(t) > 0)
 def predict_domain(
     text: str, model: Pipeline, cut_off_probability: float = 0.2
 ) -> List[DomainPrediction]:
     assert 0 <= cut_off_probability <= 1
-    log_metric('text_length', len(text))
+    log_metric("text_length", len(text))
 
     cleaned = clean(text, convert_to_ascii=True)
     text = re.sub(r"[^a-zA-Z0-9]", " ", cleaned)

good_ai/src/good_ai/good_ai/context/get_context.py

@@ -7,7 +7,7 @@ from typing import Optional, cast
 from good_ai.open_s3 import LargeFile
 from good_ai.utilities.logger import create_logger
 
-from ..persistence import PersistenceDriver, TinyDbDriver, ParallelTinyDbDriver
+from ..persistence import ParallelTinyDbDriver, PersistenceDriver
 from .context import Context
 
 _context: Optional[Context] = None
@@ -25,7 +25,9 @@ def set_default_config(
     log_level: int = INFO,
     s3_config: Path = Path("s3.ini"),
     seed: int = 42,
-    persistence_driver: PersistenceDriver = ParallelTinyDbDriver(Path("tracing_database.json")),
+    persistence_driver: PersistenceDriver = ParallelTinyDbDriver(
+        Path("tracing_database.json")
+    ),
     is_production_mode_override: Optional[bool] = None,
 ) -> None:
     global _context
@@ -39,7 +41,9 @@ def set_default_config(
     _set_seed(seed)
 
     if not persistence_driver.is_threadsafe:
-        logger.warn(f"The selected persistence driver ({type(persistence_driver).__name__}) is not threadsafe")
+        logger.warn(
+            f"The selected persistence driver ({type(persistence_driver).__name__}) is not threadsafe"
+        )
     _context = Context(
         metrics_path="/metrics",
         persistence=persistence_driver,

good_ai/src/good_ai/good_ai/deploy/serve.py

@@ -4,7 +4,6 @@ import uvicorn
 from fastapi import FastAPI, status
 from fastapi.middleware.wsgi import WSGIMiddleware
 from fastapi.responses import RedirectResponse
-from typing_extensions import Never
 
 from good_ai.good_ai.deploy.create_fastapi_app import create_fastapi_app
 
@@ -18,8 +17,8 @@ def serve(
     function: Callable[..., Any],
     disable_docs: bool = False,
     disable_metrics: bool = False,
-    configure: Callable[[FastAPI], None]=lambda _:None
+    configure: Callable[[FastAPI], None] = lambda _: None,
-) -> Never:
+) -> None:
     app = create_fastapi_app(function.__name__, disable_docs=disable_docs)
 
     if not disable_metrics:
@@ -39,7 +38,11 @@ def serve(
 
     configure(app)
 
-    uvicorn.run(app, host="0.0.0.0", port=5050, log_config={
+    uvicorn.run(
+        app,
+        host="0.0.0.0",
+        port=5050,
+        log_config={
             "version": 1,
             "disable_existing_loggers": False,
             "formatters": {
@@ -49,7 +52,7 @@ def serve(
                 },
                 "access": {
                     "()": "good_ai.logger.CustomFormatter",
-                "fmt": '%(asctime)s | %(levelname)8s | %(message)s',  # noqa: E501
+                    "fmt": "%(asctime)s | %(levelname)8s | %(message)s",  # noqa: E501
                 },
             },
             "handlers": {
@@ -67,6 +70,11 @@ def serve(
             "loggers": {
                 "uvicorn": {"handlers": ["default"], "level": "INFO"},
                 "uvicorn.error": {"level": "INFO"},
-            "uvicorn.access": {"handlers": ["access"], "level": "INFO", "propagate": False},
+                "uvicorn.access": {
+                    "handlers": ["access"],
+                    "level": "INFO",
+                    "propagate": False,
                 },
-    })
+            },
+        },
+    )

good_ai/src/good_ai/good_ai/helper/get_args.py

@@ -1,9 +1,9 @@
 import inspect
-from typing import Any, Callable, Dict, List
+from typing import Any, Callable, Dict, Mapping, Sequence
 
 
 def get_args(
-    func: Callable[..., Any], args: List[Any], kwargs: Dict[str, Any]
+    func: Callable[..., Any], args: Sequence[Any], kwargs: Mapping[str, Any]
 ) -> Dict[str, Any]:
     signature = inspect.signature(func)
     filter_keys = [

good_ai/src/good_ai/good_ai/metrics/create_dash_app.py

@@ -1,10 +1,12 @@
+from typing import Any, Dict, List, Optional, Union
+
 import pandas as pd
 from dash import Dash, dash_table, html
 from dash.dependencies import Input, Output
 from flask import Flask
 
-from good_ai.good_ai.context.get_context import get_context
+from ..context import get_context
-
+from ..views import Filter, SortBy, operators
 from .get_description import get_description
 
 
@@ -12,21 +14,7 @@ def create_dash_app(function_name: str) -> Flask:
     app = Dash(function_name, requests_pathname_prefix=get_context().metrics_path + "/")
 
     documents = get_context().persistence.get_documents()
-
+    df = pd.DataFrame(documents)
-    df = pd.DataFrame(
-        [
-            {
-                "id": d.evaluation_id,
-                "created": d.created,
-                "execution_time_ms": d.execution_time_ms,
-                "models": ", ".join(f"{m.key}:{m.version}" for m in d.models),
-                "evaluation": d.evaluation,
-                **d.logged_values,
-            }
-            for d in documents
-        ]
-    )
-    print(df)
 
     app.layout = html.Div(
         children=[
@@ -45,46 +33,11 @@ def create_dash_app(function_name: str) -> Flask:
                     sort_by=[],
                 ),
                 style={"height": 750, "overflowY": "scroll"},
-                className="six columns",
             ),
-            html.Div(id="table-paging-with-graph-container", className="five columns"),
+            html.Div(id="table-paging-with-graph-container"),
         ]
     )
 
-    operators = [
-        ["ge ", ">="],
-        ["le ", "<="],
-        ["lt ", "<"],
-        ["gt ", ">"],
-        ["ne ", "!="],
-        ["eq ", "="],
-        ["contains "],
-        ["datestartswith "],
-    ]
-
-    def split_filter_part(filter_part):
-        for operator_type in operators:
-            for operator in operator_type:
-                if operator in filter_part:
-                    name_part, value_part = filter_part.split(operator, 1)
-                    name = name_part[name_part.find("{") + 1 : name_part.rfind("}")]
-
-                    value_part = value_part.strip()
-                    v0 = value_part[0]
-                    if v0 == value_part[-1] and v0 in ("'", '"', "`"):
-                        value = value_part[1:-1].replace("\\" + v0, v0)
-                    else:
-                        try:
-                            value = float(value_part)
-                        except ValueError:
-                            value = value_part
-
-                    # word operators need spaces after them in the filter string,
-                    # but we don't want these later
-                    return name, operator_type[0].strip(), value
-
-        return [None] * 3
-
     @app.callback(
         Output("table-paging-with-graph", "data"),
         Input("table-paging-with-graph", "page_current"),
@@ -92,33 +45,19 @@ def create_dash_app(function_name: str) -> Flask:
         Input("table-paging-with-graph", "sort_by"),
         Input("table-paging-with-graph", "filter_query"),
     )
-    def update_table(page_current, page_size, sort_by, filter):
+    def update_table(
-        filtering_expressions = filter.split(" && ")
+        page_current: int, page_size: int, sort_by: List[SortBy], filter: str
-        dff = df
+    ) -> List[Dict[str, Any]]:
-        for filter_part in filtering_expressions:
+        conjunctive_filters = [get_filter(f) for f in filter.split(" && ")]
-            col_name, operator, filter_value = split_filter_part(filter_part)
+        non_null_conjunctive_filters = [f for f in conjunctive_filters if f is not None]
 
-            if operator in ("eq", "ne", "lt", "le", "gt", "ge"):
+        return get_context().persistence.query(
-                # these operators match pandas series operator method names
+            conjunctive_filters=non_null_conjunctive_filters,
-                dff = dff.loc[getattr(dff[col_name], operator)(filter_value)]
+            sort_by=sort_by,
-            elif operator == "contains":
+            skip=page_current * page_size,
-                dff = dff.loc[dff[col_name].str.contains(filter_value)]
+            take=page_size,
-            elif operator == "datestartswith":
-                # this is a simplification of the front-end filtering logic,
-                # only works with complete fields in standard format
-                dff = dff.loc[dff[col_name].str.startswith(filter_value)]
-
-        if len(sort_by):
-            dff = dff.sort_values(
-                [col["column_id"] for col in sort_by],
-                ascending=[col["direction"] == "asc" for col in sort_by],
-                inplace=False,
         )
 
-        return dff.iloc[
-            page_current * page_size : (page_current + 1) * page_size
-        ].to_dict("records")
-
     # @app.callback(
     #     Output('table-paging-with-graph-container', "children"),
     #     Input('table-paging-with-graph', "data"))
@@ -150,3 +89,24 @@ def create_dash_app(function_name: str) -> Flask:
     #     )
 
     return app.server
+
+
+def get_filter(description: str) -> Optional[Filter]:
+    print(description)
+    for operator in operators:
+        if operator in description:
+            name_part, value_part = description.split(operator, 1)
+            value_part = value_part.strip()
+            name_part = name_part[name_part.find("{") + 1 : name_part.rfind("}")]
+
+            v0 = value_part[0]
+            if v0 == value_part[-1] and v0 in ("'", '"', "`"):
+                value: Union[str, float] = value_part[1:-1].replace("\\" + v0, v0)
+            else:
+                try:
+                    value = float(value_part)
+                except ValueError:
+                    value = value_part
+            return Filter(property=name_part, operator=operator, value=value)
+
+    return None

good_ai/src/good_ai/good_ai/metrics/log_argument.py

@@ -1,5 +1,5 @@
 from functools import wraps
-from typing import Any, Callable, Dict, List
+from typing import Any, Callable, Dict
 
 from ..exceptions import ArgumentValidationError
 from ..helper import get_args
@@ -15,15 +15,13 @@ def log_argument(
         actual_name = f"arg:{func.__name__}:{argument_name}"
 
         @wraps(func)
-        def wrapper(*args: List[Any], **kwargs: Dict[str, Any]) -> Any:
+        def wrapper(*args: Any, **kwargs: Dict[str, Any]) -> Any:
             arguments = get_args(func, args, kwargs)
             argument = arguments[argument_name]
 
             expected_type = func.__annotations__.get(argument_name)
 
-            if (
+            if expected_type is not None and not isinstance(argument, expected_type):
-                expected_type is not None and not isinstance(argument, expected_type)
-            ):
                 raise ArgumentValidationError(
                     f"Argument {argument_name} in {func.__name__} has the wrong type, expected: {expected_type.__name__}, got: {type(argument).__name__}"
                 )

good_ai/src/good_ai/good_ai/persistence/__init__.py

@@ -1,4 +1,3 @@
 from .mongodb_driver import MongoDbDriver
-from .persistence_driver import PersistenceDriver
-from .tinydb_driver import TinyDbDriver
 from .parallel_tinydb_driver import ParallelTinyDbDriver
+from .persistence_driver import PersistenceDriver

good_ai/src/good_ai/good_ai/persistence/parallel_tinydb_driver.py

@@ -1,16 +1,20 @@
+from multiprocessing import Lock
 from pathlib import Path
-from typing import Any, Callable
+from typing import Any, Callable, Dict
+
+import pandas as pd
 from black import List
 from tinydb import TinyDB
-from multiprocessing import Process, Lock
 
-from ..views import Trace
+from ..views import Filter, SortBy, Trace
 from .persistence_driver import PersistenceDriver
 
-
 lock = Lock()
 
 
+operator_mapping = {"=": "eq", "!=": "ne", "<": "lt", "<=": "le", ">": "gt", ">=": "ge"}
+
+
 class ParallelTinyDbDriver(PersistenceDriver):
     is_threadsafe = True
 
@@ -21,9 +25,40 @@ class ParallelTinyDbDriver(PersistenceDriver):
     def save_document(self, trace: Trace) -> str:
         return self._safe_execute(lambda db: db.insert(trace.dict()))
 
-    def get_documents(self) -> List[Trace]:
+    def get_traces(self) -> List[Trace]:
         return self._safe_execute(lambda db: [Trace.parse_obj(t) for t in db.all()])
 
+    def get_documents(self) -> List[Dict[str, Any]]:
+        documents = self.get_traces()
+        return [d.to_flat_dict() for d in documents]
+
+    def query(
+        self,
+        conjunctive_filters: List[Filter],
+        sort_by: List[SortBy],
+        skip: int,
+        take: int,
+    ) -> List[Dict[str, Any]]:
+        documents = self.get_documents()
+        df = pd.DataFrame(documents)
+
+        for f in conjunctive_filters:
+            if f.operator in operator_mapping:
+                df = df.loc[
+                    getattr(df[f.property], operator_mapping[f.operator])(f.value)
+                ]
+            elif f.operator == "contains":
+                df = df.loc[df[f.property].str.contains(f.value)]
+
+        if sort_by:
+            df = df.sort_values(
+                [col["column_id"] for col in sort_by],
+                ascending=[col["direction"] == "asc" for col in sort_by],
+                inplace=False,
+            )
+
+        return df.iloc[skip : skip + take].to_dict("records")
+
     def _safe_execute(self, func: Callable[[TinyDB], Any]) -> Any:
         with lock:
             with TinyDB(self._path_to_db) as db:

good_ai/src/good_ai/good_ai/persistence/persistence_driver.py

@@ -1,8 +1,9 @@
 from abc import ABC, abstractmethod
+from typing import Any, Dict
 
 from black import List
 
-from good_ai.good_ai.views.trace import Trace
+from ..views import Filter, SortBy, Trace
 
 
 class PersistenceDriver(ABC):
@@ -13,5 +14,19 @@ class PersistenceDriver(ABC):
         pass
 
     @abstractmethod
-    def get_documents(self) -> List[Trace]:
+    def get_traces(self) -> List[Trace]:
+        pass
+
+    @abstractmethod
+    def get_documents(self) -> List[Dict[str, Any]]:
+        pass
+
+    @abstractmethod
+    def query(
+        self,
+        conjunctive_filters: List[Filter],
+        sort_by: List[SortBy],
+        skip: int,
+        take: int,
+    ) -> List[Dict[str, Any]]:
         pass

good_ai/src/good_ai/good_ai/persistence/tinydb_driver.py

@@ -1,23 +0,0 @@
-from pathlib import Path
-from uuid import uuid4
-
-from black import List
-from tinydb import TinyDB
-from tinydb.table import Document
-
-from ..views import Trace
-from .persistence_driver import PersistenceDriver
-
-
-class TinyDbDriver(PersistenceDriver):
-    is_threadsafe = False
-    
-    def __init__(self, path_to_db: Path) -> None:
-        super().__init__()
-        self._db = TinyDB(path_to_db)
-
-    def save_document(self, trace: Trace) -> str:
-        return self._db.insert(trace.dict())
-
-    def get_documents(self) -> List[Trace]:
-        return [Trace.parse_obj(t) for t in self._db.all()]

good_ai/src/good_ai/good_ai/tracing/tracing_context.py

@@ -2,7 +2,7 @@ import threading
 from collections import defaultdict
 from datetime import datetime
 from types import TracebackType
-from typing import Any, DefaultDict, Dict, List, Optional, Type
+from typing import Any, DefaultDict, Dict, List, Literal, Optional, Type
 
 from ..context import get_context
 from ..views import Model, Trace
@@ -53,7 +53,7 @@ class TracingContext:
         type: Optional[Type[BaseException]],
         exception: Optional[BaseException],
         traceback: Optional[TracebackType],
-    ) -> bool:
+    ) -> Literal[False]:
         assert self._contexts[threading.get_ident()][-1] == self
         self._contexts[threading.get_ident()].remove(self)
 

good_ai/src/good_ai/good_ai/views/__init__.py

@@ -1,3 +1,6 @@
+from .filter import Filter
 from .health_check_response import HealthCheckResponse
 from .model import Model
+from .operators import operators
+from .sort_by import SortBy
 from .trace import Trace

good_ai/src/good_ai/good_ai/views/filter.py

@@ -0,0 +1,11 @@
+from typing import Union
+
+from pydantic import BaseModel
+
+from .operators import Operator
+
+
+class Filter(BaseModel):
+    property: str
+    operator: Operator
+    value: Union[float, str]

good_ai/src/good_ai/good_ai/views/operators.py

@@ -0,0 +1,13 @@
+from typing import List, Literal
+
+Operator = Literal[">=", "<=", "<", ">", "!=", "=", "contains"]
+
+operators: List[Operator] = [
+    ">=",
+    "<=",
+    "<",
+    ">",
+    "!=",
+    "=",
+    "contains",
+]

good_ai/src/good_ai/good_ai/views/sort_by.py

@@ -0,0 +1,6 @@
+from typing import Literal, TypedDict
+
+
+class SortBy(TypedDict):
+    column_id: str
+    direction: Literal["asc", "desc"]

good_ai/src/good_ai/good_ai/views/trace.py

@@ -23,5 +23,15 @@ class Trace(BaseModel):
             return str(uuid4())
         return v
 
+    def to_flat_dict(self) -> Dict[str, Any]:
+        return {
+            "id": self.evaluation_id,
+            "created": self.created,
+            "execution_time_ms": self.execution_time_ms,
+            "models": ", ".join(f"{m.key}:{m.version}" for m in self.models),
+            "evaluation": self.evaluation,
+            **self.logged_values,
+        }
+
     def __hash__(self) -> int:
         return hash((type(self),) + tuple(self.__dict__.values()))

good_ai/src/good_ai/utilities/parallel_map.py

@@ -4,9 +4,10 @@ from typing import Any, Callable, Iterable, List, Optional
 import multiprocess as mp
 import psutil
 from tqdm.auto import tqdm
+
 from .logger import create_logger
 
-logger = create_logger('parallel_map')
+logger = create_logger("parallel_map")
 
 
 def parallel_map(
@@ -26,7 +27,9 @@ def parallel_map(
     if not chunk_size:
         chunk_size = max(1, ceil(len(values) / concurrency / 10))
 
-    logger.info(f"Starting parallel map, concurrency: {concurrency}, chunk size: {chunk_size}")
+    logger.info(
+        f"Starting parallel map, concurrency: {concurrency}, chunk size: {chunk_size}"
+    )
 
     if concurrency == 1 or len(values) <= chunk_size:
         logger.warn(f"Running in series, there is no reason for parallelism")