prompture 0.0.32.dev1__py3-none-any.whl → 0.0.33.dev1__py3-none-any.whl

prompture/__init__.py

@@ -13,7 +13,7 @@ from .core import (
     extract_from_pandas,
     render_output,
 )
-from .drivers import get_driver, get_driver_for_model, OpenAIDriver, LocalHTTPDriver, OllamaDriver, ClaudeDriver, LMStudioDriver, AzureDriver, GoogleDriver, GroqDriver, OpenRouterDriver, GrokDriver
+from .drivers import get_driver, get_driver_for_model, OpenAIDriver, LocalHTTPDriver, OllamaDriver, ClaudeDriver, LMStudioDriver, AzureDriver, GoogleDriver, GroqDriver, OpenRouterDriver, GrokDriver, AirLLMDriver
 from .tools import clean_json_text, clean_toon_text
 from .field_definitions import (
     FIELD_DEFINITIONS, get_field_definition, get_required_fields, get_field_names,
@@ -87,6 +87,7 @@ __all__ = [
     "GroqDriver",
     "OpenRouterDriver",
     "GrokDriver",
+    "AirLLMDriver",
     # Discovery
     "get_available_models",
 ]

prompture/drivers/__init__.py

@@ -8,6 +8,7 @@ from .google_driver import GoogleDriver
 from .groq_driver import GroqDriver
 from .openrouter_driver import OpenRouterDriver
 from .grok_driver import GrokDriver
+from .airllm_driver import AirLLMDriver
 from ..settings import settings
 
 
@@ -54,6 +55,10 @@ DRIVER_REGISTRY = {
         api_key=settings.grok_api_key,
         model=model or settings.grok_model
     ),
+    "airllm": lambda model=None: AirLLMDriver(
+        model=model or settings.airllm_model,
+        compression=settings.airllm_compression,
+    ),
 }
 
 
@@ -115,6 +120,7 @@ __all__ = [
     "GroqDriver",
     "OpenRouterDriver",
     "GrokDriver",
+    "AirLLMDriver",
     "get_driver",
     "get_driver_for_model",
 ]

prompture/drivers/airllm_driver.py

@@ -0,0 +1,116 @@
+import logging
+from ..driver import Driver
+from typing import Any, Dict, Optional
+
+logger = logging.getLogger(__name__)
+
+
+class AirLLMDriver(Driver):
+    """Driver for AirLLM — run large models (70B+) on consumer GPUs via
+    layer-by-layer memory management.
+
+    The ``airllm`` package is a lazy dependency: it is imported on first
+    ``generate()`` call so the rest of Prompture works without it installed.
+    """
+
+    MODEL_PRICING = {
+        "default": {"prompt": 0.0, "completion": 0.0}
+    }
+
+    def __init__(self, model: str = "meta-llama/Llama-2-7b-hf",
+                 compression: Optional[str] = None):
+        """
+        Args:
+            model: HuggingFace repo ID (e.g. ``"meta-llama/Llama-2-70b-hf"``).
+            compression: Optional quantization mode — ``"4bit"`` or ``"8bit"``.
+        """
+        self.model = model
+        self.compression = compression
+        self.options: Dict[str, Any] = {}
+        self._llm = None
+        self._tokenizer = None
+
+    # ------------------------------------------------------------------
+    # Lazy model loading
+    # ------------------------------------------------------------------
+    def _ensure_loaded(self):
+        """Load the AirLLM model and tokenizer on first use."""
+        if self._llm is not None:
+            return
+
+        try:
+            from airllm import AutoModel
+        except ImportError:
+            raise ImportError(
+                "The 'airllm' package is required for the AirLLM driver. "
+                "Install it with: pip install prompture[airllm]"
+            )
+
+        try:
+            from transformers import AutoTokenizer
+        except ImportError:
+            raise ImportError(
+                "The 'transformers' package is required for the AirLLM driver. "
+                "Install it with: pip install transformers"
+            )
+
+        logger.info(f"Loading AirLLM model: {self.model} "
+                     f"(compression={self.compression})")
+
+        load_kwargs: Dict[str, Any] = {}
+        if self.compression:
+            load_kwargs["compression"] = self.compression
+
+        self._llm = AutoModel.from_pretrained(self.model, **load_kwargs)
+        self._tokenizer = AutoTokenizer.from_pretrained(self.model)
+        logger.info("AirLLM model loaded successfully")
+
+    # ------------------------------------------------------------------
+    # Driver interface
+    # ------------------------------------------------------------------
+    def generate(self, prompt: str, options: Dict[str, Any] = None) -> Dict[str, Any]:
+        self._ensure_loaded()
+
+        merged_options = self.options.copy()
+        if options:
+            merged_options.update(options)
+
+        max_new_tokens = merged_options.get("max_new_tokens", 256)
+
+        # Tokenize
+        input_ids = self._tokenizer(
+            prompt, return_tensors="pt"
+        ).input_ids
+
+        prompt_tokens = input_ids.shape[1]
+
+        logger.debug(f"AirLLM generating with max_new_tokens={max_new_tokens}, "
+                      f"prompt_tokens={prompt_tokens}")
+
+        # Generate
+        output_ids = self._llm.generate(
+            input_ids,
+            max_new_tokens=max_new_tokens,
+        )
+
+        # Decode only the newly generated tokens (strip the prompt prefix)
+        new_tokens = output_ids[0, prompt_tokens:]
+        completion_tokens = len(new_tokens)
+        text = self._tokenizer.decode(new_tokens, skip_special_tokens=True)
+
+        total_tokens = prompt_tokens + completion_tokens
+
+        meta = {
+            "prompt_tokens": prompt_tokens,
+            "completion_tokens": completion_tokens,
+            "total_tokens": total_tokens,
+            "cost": 0.0,
+            "raw_response": {
+                "model": self.model,
+                "compression": self.compression,
+                "max_new_tokens": max_new_tokens,
+            },
+            "model_name": self.model,
+        }
+
+        return {"text": text, "meta": meta}

prompture/settings.py

@@ -48,6 +48,10 @@ class Settings(BaseSettings):
     grok_api_key: Optional[str] = None
     grok_model: str = "grok-4-fast-reasoning"
 
+    # AirLLM
+    airllm_model: str = "meta-llama/Llama-2-7b-hf"
+    airllm_compression: Optional[str] = None  # "4bit" or "8bit"
+
     model_config = SettingsConfigDict(
         env_file=".env",
         extra="ignore",

{prompture-0.0.32.dev1.dist-info → prompture-0.0.33.dev1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: prompture
-Version: 0.0.32.dev1
+Version: 0.0.33.dev1
 Summary: Ask LLMs to return structured JSON and run cross-model tests. API-first.
 Home-page: https://github.com/jhd3197/prompture
 Author: Juan Denis
@@ -29,6 +29,8 @@ Requires-Dist: tukuy>=0.0.6
 Requires-Dist: pyyaml>=6.0
 Provides-Extra: test
 Requires-Dist: pytest>=7.0; extra == "test"
+Provides-Extra: airllm
+Requires-Dist: airllm>=2.8.0; extra == "airllm"
 Dynamic: author
 Dynamic: author-email
 Dynamic: classifier

{prompture-0.0.32.dev1.dist-info → prompture-0.0.33.dev1.dist-info}/RECORD

@@ -1,14 +1,15 @@
-prompture/__init__.py,sha256=kCcOseMTHaJkl-vtzXVbbBdWRQlIWWBr-C-l9E2mScU,2689
+prompture/__init__.py,sha256=Bs3nDZWEl9tFrS-zH7RThl_svKwPM47Q8bIuD8qKlGM,2723
 prompture/cli.py,sha256=vA86GNjtKSHz8eRMl5YDaT9HHIWuhkeJtfx8jqTaqtM,809
 prompture/core.py,sha256=x_FhOY37ygQVHo4zHUyiWsV4BuOClkELsVhEV-K4jJ0,53689
 prompture/discovery.py,sha256=qQ7Quz0Tqo0f2h9DqMlV7RqMP4XOeue_ZwzXq4bf6B8,6788
 prompture/driver.py,sha256=w8pdXHujImIGF3ee8rkG8f6-UD0h2jLHhucSPInRrYI,989
 prompture/field_definitions.py,sha256=6kDMYNedccTK5l2L_I8_NI3_av-iYHqGPwkKDy8214c,21731
 prompture/runner.py,sha256=5xwal3iBQQj4_q7l3Rjr0e3RrUMJPaPDLiEchO0mmHo,4192
-prompture/settings.py,sha256=vHRkBAZNP6yRsI2Sm4FMa_FCw0Zxy2VX97ooiVYWvks,1500
+prompture/settings.py,sha256=-S1AJa2AUblwGyzY1lOvEI_ceNRwgPeGvJYsbGulwv4,1630
 prompture/tools.py,sha256=qyT8oJl_v9GolABkflW0SvEx22yNkEJZKTu-40nJbs0,40329
 prompture/validator.py,sha256=oLzVsNveHuF-N_uOd11_uDa9Q5rFyo0wrk_l1N4zqDk,996
-prompture/drivers/__init__.py,sha256=IQ7DsWC_FP45h2CprWRhQ7lKi3-9ZO6CgweNX6IxTUA,3896
+prompture/drivers/__init__.py,sha256=hi2u4Z2KQFfgqce1QvjRlDKRzB2xfJZpidGNMsQ82oI,4105
+prompture/drivers/airllm_driver.py,sha256=g1WmQDwSfK0BIyG96JrZY7W_VHXOS7wDSeegE7B1q4Y,3956
 prompture/drivers/azure_driver.py,sha256=GROhK3hqMfMurnEgpAawa1DPS-FhOU0YQcgy9SNGTzM,4622
 prompture/drivers/claude_driver.py,sha256=ZEHQNqNThLZ0p-WmGVuKiNyiudGYGP07xIzbgZhLY1g,3293
 prompture/drivers/google_driver.py,sha256=bCsCSuCRise0L_HOmw-jBh1hrpd8glNBkVFlOZeP0DM,6338
@@ -20,9 +21,9 @@ prompture/drivers/local_http_driver.py,sha256=S2diikvtQOQHF7fB07zU2X0QWkej4Of__r
 prompture/drivers/ollama_driver.py,sha256=fq_eFgwmCT3SK1D-ICHjxLjcm_An0suwkFIWC38xsS0,4681
 prompture/drivers/openai_driver.py,sha256=9q9OjQslquRFvIl1Hd9JVmFFFVh6OBIWrFulw1mkYWg,3976
 prompture/drivers/openrouter_driver.py,sha256=GKvLOFDhsyopH-k3iaD3VWllm7xbGuopRSA02MfCKoM,5031
-prompture-0.0.32.dev1.dist-info/licenses/LICENSE,sha256=0HgDepH7aaHNFhHF-iXuW6_GqDfYPnVkjtiCAZ4yS8I,1060
-prompture-0.0.32.dev1.dist-info/METADATA,sha256=3oNb4hhkYR7ZuLsrG5wrRxJjbuLnazBaQHKaW2yAM0Y,18043
-prompture-0.0.32.dev1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-prompture-0.0.32.dev1.dist-info/entry_points.txt,sha256=AFPG3lJR86g4IJMoWQUW5Ph7G6MLNWG3A2u2Tp9zkp8,48
-prompture-0.0.32.dev1.dist-info/top_level.txt,sha256=to86zq_kjfdoLeAxQNr420UWqT0WzkKoZ509J7Qr2t4,10
-prompture-0.0.32.dev1.dist-info/RECORD,,
+prompture-0.0.33.dev1.dist-info/licenses/LICENSE,sha256=0HgDepH7aaHNFhHF-iXuW6_GqDfYPnVkjtiCAZ4yS8I,1060
+prompture-0.0.33.dev1.dist-info/METADATA,sha256=1W2vCL_ZzJvAcxi78mbYfdaIb__mudc44Ryb_b0GlIU,18114
+prompture-0.0.33.dev1.dist-info/WHEEL,sha256=wUyA8OaulRlbfwMtmQsvNngGrxQHAvkKcvRmdizlJi0,92
+prompture-0.0.33.dev1.dist-info/entry_points.txt,sha256=AFPG3lJR86g4IJMoWQUW5Ph7G6MLNWG3A2u2Tp9zkp8,48
+prompture-0.0.33.dev1.dist-info/top_level.txt,sha256=to86zq_kjfdoLeAxQNr420UWqT0WzkKoZ509J7Qr2t4,10
+prompture-0.0.33.dev1.dist-info/RECORD,,

{prompture-0.0.32.dev1.dist-info → prompture-0.0.33.dev1.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (80.9.0)
+Generator: setuptools (80.10.2)
 Root-Is-Purelib: true
 Tag: py3-none-any