Merge pull request #46 from Nayjest/infra_improvements

Nayjest · web-flow · commit e9789b69ad7c · 2026-02-13T18:36:39.000+01:00
Infra improvements
diff --git a/lm_proxy/__main__.py b/lm_proxy/__main__.py
@@ -1,4 +1,5 @@
 """Provides the CLI entry point when the package is executed as a Python module."""
+
 from .app import cli_app
 
 
diff --git a/lm_proxy/api_key_check/__init__.py b/lm_proxy/api_key_check/__init__.py
@@ -1,4 +1,5 @@
 """Collection of built-in API-key checkers for usage in the configuration."""
+
 from .in_config import check_api_key_in_config
 from .with_request import CheckAPIKeyWithRequest
 from .allow_all import AllowAll
diff --git a/lm_proxy/api_key_check/allow_all.py b/lm_proxy/api_key_check/allow_all.py
@@ -4,6 +4,7 @@
 This module provides a simple authentication strategy for development or testing
 environments where all API keys should be accepted without validation.
 """
+
 from typing import Optional
 from dataclasses import dataclass
 
@@ -25,10 +26,7 @@ class AllowAll:
     group: str = "default"
     capture_api_key: bool = True
 
-    def __call__(
-        self,
-        api_key: Optional[str]
-    ) -> tuple[str, dict[str, Optional[str]]]:
+    def __call__(self, api_key: Optional[str]) -> tuple[str, dict[str, Optional[str]]]:
         """
         Validate an API key (accepts all keys without verification).
 
diff --git a/lm_proxy/api_key_check/in_config.py b/lm_proxy/api_key_check/in_config.py
@@ -5,6 +5,7 @@
 For using this function,
 set "api_key_check" configuration value to "lm_proxy.api_key_check.check_api_key_in_config".
 """
+
 from typing import Optional
 from ..bootstrap import env
 
diff --git a/lm_proxy/api_key_check/with_request.py b/lm_proxy/api_key_check/with_request.py
@@ -1,6 +1,7 @@
 """
 API key check implementation using HTTP requests.
 """
+
 from typing import Optional
 from dataclasses import dataclass, field
 import requests
@@ -13,6 +14,7 @@ class CheckAPIKeyWithRequest:  # pylint: disable=too-many-instance-attributes
     """
     Validates a Client API key by making an HTTP request to a specified URL.
     """
+
     url: str = field()
     method: str = field(default="get")
     headers: dict = field(default_factory=dict)
@@ -45,10 +47,7 @@ def check_func(api_key: str) -> Optional[tuple[str, dict]]:
                     for k, v in self.headers.items()
                 }
                 response = requests.request(
-                    method=self.method,
-                    url=url,
-                    headers=headers,
-                    timeout=self.timeout
+                    method=self.method, url=url, headers=headers, timeout=self.timeout
                 )
                 response.raise_for_status()
                 group = self.default_group
diff --git a/lm_proxy/app.py b/lm_proxy/app.py
@@ -1,6 +1,7 @@
 """
 LM-Proxy Application Entrypoint
 """
+
 import logging
 from typing import Optional
 from fastapi import FastAPI
@@ -19,9 +20,7 @@
 @cli_app.callback(invoke_without_command=True)
 def run_server(
     config: Optional[str] = typer.Option(None, help="Path to the configuration file"),
-    debug: Optional[bool] = typer.Option(
-        None, help="Enable debug mode (more verbose logging)"
-    ),
+    debug: Optional[bool] = typer.Option(None, help="Enable debug mode (more verbose logging)"),
     env_file: Optional[str] = typer.Option(
         ".env",
         "--env",
@@ -55,9 +54,7 @@ def web_app():
     """
     Entrypoint for ASGI server
     """
-    app = FastAPI(
-        title="LM-Proxy", description="OpenAI-compatible proxy server for LLM inference"
-    )
+    app = FastAPI(title="LM-Proxy", description="OpenAI-compatible proxy server for LLM inference")
     OpenAIHTTPException.register(app)
     app.add_api_route(
         path=f"{env.config.api_prefix}/chat/completions",
diff --git a/lm_proxy/base_types.py b/lm_proxy/base_types.py
@@ -1,4 +1,5 @@
 """Base types used in LM-Proxy."""
+
 import uuid
 from dataclasses import dataclass, field
 from datetime import datetime
@@ -17,6 +18,7 @@ class ChatCompletionRequest(BaseModel):
     """
     Request model for chat/completions endpoint.
     """
+
     model: str
     messages: List[mc.Msg | dict]
     # | dict --> support of messages with lists of dicts
@@ -58,6 +60,7 @@ class RequestContext:  # pylint: disable=too-many-instance-attributes
     """
     Stores information about a single LLM request/response cycle for usage in middleware.
     """
+
     id: Optional[str] = field(default_factory=lambda: str(uuid.uuid4()))
     request: Optional[ChatCompletionRequest] = field(default=None)
     http_request: Optional[Request] = field(default=None)
@@ -83,7 +86,4 @@ def to_dict(self) -> dict:
         return data
 
 
-THandler = Callable[
-    [RequestContext],
-    Union[Awaitable[None], None]
-]
+THandler = Callable[[RequestContext], Union[Awaitable[None], None]]
diff --git a/lm_proxy/bootstrap.py b/lm_proxy/bootstrap.py
@@ -44,6 +44,7 @@ def format(self, record):
 
 class Env:
     """Runtime environment singleton."""
+
     config: Config
     connections: dict[str, mc.types.LLMAsyncFunctionType]
     debug: bool
@@ -66,9 +67,7 @@ def init(config: Config | str | PathLike, debug: bool = False):
             if isinstance(config, (str, PathLike)):
                 config = Config.load(config)
             else:
-                raise ValueError(
-                    "config must be a path (str or PathLike) or Config instance"
-                )
+                raise ValueError("config must be a path (str or PathLike) or Config instance")
         env.config = config
 
         env._init_components()
@@ -84,18 +83,12 @@ def init(config: Config | str | PathLike, debug: bool = False):
                 if inspect.iscoroutinefunction(conn_config):
                     env.connections[conn_name] = conn_config
                 elif isinstance(conn_config, str):
-                    env.connections[conn_name] = resolve_instance_or_callable(
-                        conn_config
-                    )
+                    env.connections[conn_name] = resolve_instance_or_callable(conn_config)
                 else:
-                    mc.configure(
-                        **conn_config, EMBEDDING_DB_TYPE=mc.EmbeddingDbType.NONE
-                    )
+                    mc.configure(**conn_config, EMBEDDING_DB_TYPE=mc.EmbeddingDbType.NONE)
                     env.connections[conn_name] = mc.env().llm_async_function
             except mc.LLMConfigError as e:
-                raise ValueError(
-                    f"Error in configuration for connection '{conn_name}': {e}"
-                ) from e
+                raise ValueError(f"Error in configuration for connection '{conn_name}': {e}") from e
 
         logging.info("Done initializing %d connections.", len(env.connections))
 
diff --git a/lm_proxy/config.py b/lm_proxy/config.py
@@ -30,6 +30,7 @@ class ModelListingMode(StrEnum):
 
 class Group(BaseModel):
     """User group configuration."""
+
     api_keys: list[str] = Field(default_factory=list)
     allowed_connections: str = Field(default="*")  # Comma-separated list or "*"
 
diff --git a/lm_proxy/config_loaders/__init__.py b/lm_proxy/config_loaders/__init__.py
@@ -1,4 +1,5 @@
 """Built-in configuration loaders for different file formats."""
+
 from .python import load_python_config
 from .toml import load_toml_config
 from .yaml import load_yaml_config
diff --git a/lm_proxy/config_loaders/json.py b/lm_proxy/config_loaders/json.py
@@ -1,4 +1,5 @@
 """JSON configuration loader."""
+
 import json
 
 
diff --git a/lm_proxy/config_loaders/python.py b/lm_proxy/config_loaders/python.py
@@ -1,4 +1,5 @@
 """Loader for Python configuration files."""
+
 import importlib.util
 from ..config import Config
 
diff --git a/lm_proxy/config_loaders/toml.py b/lm_proxy/config_loaders/toml.py
@@ -1,4 +1,5 @@
 """TOML configuration loader."""
+
 import tomllib
 
 
diff --git a/lm_proxy/core.py b/lm_proxy/core.py
@@ -46,9 +46,7 @@ def parse_routing_rule(rule: str, config: Config) -> tuple[str, str]:
     return connection_name, model_part
 
 
-def resolve_connection_and_model(
-    config: Config, external_model: str
-) -> tuple[str, str]:
+def resolve_connection_and_model(config: Config, external_model: str) -> tuple[str, str]:
     """
     Resolves the connection name and model name based on routing rules.
     Args:
@@ -150,9 +148,7 @@ def api_key_id(api_key: Optional[str]) -> str | None:
     """
     if not api_key:
         return None
-    return hashlib.md5(
-        (api_key + env.config.encryption_key).encode("utf-8")
-    ).hexdigest()
+    return hashlib.md5((api_key + env.config.encryption_key).encode("utf-8")).hexdigest()
 
 
 def fail_if_service_disabled():
@@ -207,9 +203,7 @@ async def check(request: Request) -> tuple[str, str, dict]:
     return group, api_key, user_info
 
 
-async def chat_completions(
-    request: ChatCompletionRequest, raw_request: Request
-) -> Response:
+async def chat_completions(request: ChatCompletionRequest, raw_request: Request) -> Response:
     """
     Endpoint for chat completions that mimics OpenAI's API structure.
     """
diff --git a/lm_proxy/errors.py b/lm_proxy/errors.py
@@ -6,6 +6,7 @@ class OpenAIHTTPException(Exception):
     """
     OpenAI API-compatible exception.
     """
+
     # HTTPException from FastAPI is not used directly to provide error response format
     # compatible with OpenAI API.
 
diff --git a/lm_proxy/handlers/__init__.py b/lm_proxy/handlers/__init__.py
@@ -1,7 +1,4 @@
 from .rate_limiter import RateLimiter
 from .forward_http_headers import HTTPHeadersForwarder
 
-__all__ = [
-    'RateLimiter',
-    'HTTPHeadersForwarder'
-]
+__all__ = ["RateLimiter", "HTTPHeadersForwarder"]
diff --git a/lm_proxy/handlers/forward_http_headers.py b/lm_proxy/handlers/forward_http_headers.py
@@ -1,6 +1,7 @@
 """
 HTTP headers forwarder LM-Proxy.
 """
+
 from dataclasses import dataclass, field
 
 from lm_proxy.base_types import RequestContext
diff --git a/lm_proxy/handlers/rate_limiter.py b/lm_proxy/handlers/rate_limiter.py
@@ -3,6 +3,7 @@
 
 Provides sliding window rate limiting per API key / IP address / connection / user group / global.
 """
+
 import threading
 import time
 from dataclasses import dataclass, field
@@ -33,6 +34,7 @@ class RateLimiter:
         per: Scope for rate limiting
             ("api_key", "connection", "group", "ip", "global").
     """
+
     max_requests: int = 60
     window_seconds: float = 60.0
     per: RateLimitScope = RateLimitScope.API_KEY
@@ -68,10 +70,7 @@ async def __call__(self, ctx: RequestContext) -> None:
 
         with self._lock:
             if len(self._buckets) > self.max_buckets:
-                self._buckets = {
-                    k: v for k, v in self._buckets.items()
-                    if v and v[-1] > cutoff
-                }
+                self._buckets = {k: v for k, v in self._buckets.items() if v and v[-1] > cutoff}
 
             timestamps = [t for t in self._buckets.get(key, []) if t > cutoff]
 
diff --git a/lm_proxy/loggers.py b/lm_proxy/loggers.py
@@ -1,4 +1,5 @@
 """LLM Request logging."""
+
 import abc
 import json
 import os
@@ -11,13 +12,15 @@
 
 class AbstractLogEntryTransformer(abc.ABC):  # pylint: disable=too-few-public-methods
     """Transform RequestContext into a dictionary of logged attributes."""
+
     @abc.abstractmethod
     def __call__(self, request_context: RequestContext) -> dict:
         raise NotImplementedError()
 
 
 class AbstractLogWriter(abc.ABC):  # pylint: disable=too-few-public-methods
     """Writes the logged data to a destination."""
+
     @abc.abstractmethod
     def __call__(self, logged_data: dict):
         raise NotImplementedError()
@@ -29,6 +32,7 @@ class LogEntryTransformer(AbstractLogEntryTransformer):  # pylint: disable=too-f
     The mapping is provided as keyword arguments, where keys are the names of the
     logged attributes, and values are the paths to the attributes in RequestContext.
     """
+
     def __init__(self, **kwargs):
         self.mapping = kwargs
 
@@ -42,6 +46,7 @@ def __call__(self, request_context: RequestContext) -> dict:
 @dataclass
 class BaseLogger:
     """Base LLM request logger."""
+
     log_writer: AbstractLogWriter | str | dict
     entry_transformer: AbstractLogEntryTransformer | str | dict = field(default=None)
 
@@ -69,6 +74,7 @@ def __call__(self, request_context: RequestContext):
 @dataclass
 class JsonLogWriter(AbstractLogWriter):
     """Writes logged data to a JSON file."""
+
     file_name: str
 
     def __post_init__(self):
diff --git a/lm_proxy/models_endpoint.py b/lm_proxy/models_endpoint.py
@@ -24,21 +24,18 @@ async def models(request: Request) -> JSONResponse:
             is_model_name = not ("*" in model_pattern or "?" in model_pattern)
             if not is_model_name:
                 if env.config.model_listing_mode != ModelListingMode.AS_IS:
-                    if (
-                        env.config.model_listing_mode
-                        == ModelListingMode.IGNORE_WILDCARDS
-                    ):
+                    if env.config.model_listing_mode == ModelListingMode.IGNORE_WILDCARDS:
                         continue
                     raise NotImplementedError(
                         f"'{env.config.model_listing_mode}' model listing mode "
                         f"is not implemented yet"
                     )
             model_data = {
-                    "id": model_pattern,
-                    "object": "model",
-                    "created": 0,
-                    "owned_by": connection_name,
-                }
+                "id": model_pattern,
+                "object": "model",
+                "created": 0,
+                "owned_by": connection_name,
+            }
 
             if aux_info := env.config.model_info.get(model_pattern):
                 model_data.update(aux_info)
diff --git a/multi-build.py b/multi-build.py
@@ -26,16 +26,17 @@ def replace_name(old_names: list[str], new_names: list[str], files: list[str] =
             p = Path(path)
             p.write_text(
                 re.sub(
-                    fr'(?<![\\/\w]){old_name}\b',
+                    rf"(?<![\\/\w]){old_name}\b",
                     new_name,
                     p.read_text(encoding="utf-8"),
-                    flags=re.M
-                ), encoding="utf-8"
+                    flags=re.M,
+                ),
+                encoding="utf-8",
             )
 
 
 prev = NAMES[0]
-for nxt in NAMES[1:]+[NAMES[0]]:
+for nxt in NAMES[1:] + [NAMES[0]]:
     print(f"Building for project name: {nxt[0]}...")
     replace_name(prev, nxt)
     subprocess.run(["poetry", "build"], check=True)
diff --git a/pyproject.toml b/pyproject.toml
@@ -88,3 +88,5 @@ asyncio_mode = "auto"
 testpaths = [
     "tests",
 ]
+[tool.black]
+line-length = 100
diff --git a/tests/configs/test_disabled.py b/tests/configs/test_disabled.py
diff --git a/tests/test_api_key_check_allow_all.py b/tests/test_api_key_check_allow_all.py
diff --git a/tests/test_api_key_check_errors.py b/tests/test_api_key_check_errors.py
diff --git a/tests/test_integration.py b/tests/test_integration.py
diff --git a/tests/test_integration_headers.py b/tests/test_integration_headers.py
diff --git a/tests/test_loggers.py b/tests/test_loggers.py
diff --git a/tests/test_utils.py b/tests/test_utils.py

Original file line number	Diff line number	Diff line change
`@@ -1,4 +1,5 @@`
`1`	`1`	`"""Provides the CLI entry point when the package is executed as a Python module."""`
	`2`	`+`
`2`	`3`	`from .app import cli_app`
`3`	`4`
`4`	`5`
Original file line number	Diff line number	Diff line change
`@@ -1,4 +1,5 @@`
`1`	`1`	`"""Collection of built-in API-key checkers for usage in the configuration."""`
	`2`	`+`
`2`	`3`	`from .in_config import check_api_key_in_config`
`3`	`4`	`from .with_request import CheckAPIKeyWithRequest`
`4`	`5`	`from .allow_all import AllowAll`
Original file line number	Diff line number	Diff line change
`@@ -1,4 +1,5 @@`
`1`	`1`	`"""Built-in configuration loaders for different file formats."""`
	`2`	`+`
`2`	`3`	`from .python import load_python_config`
`3`	`4`	`from .toml import load_toml_config`
`4`	`5`	`from .yaml import load_yaml_config`