python-infrakit-dev 0.1.5__tar.gz → 0.1.6__tar.gz

{python_infrakit_dev-0.1.5 → python_infrakit_dev-0.1.6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: python-infrakit-dev
-Version: 0.1.5
+Version: 0.1.6
 Summary: A comprehensive Python developer infrastructure toolkit
 Project-URL: Homepage, https://github.com/chiragg21/infrakit
 Project-URL: Repository, https://github.com/chiragg21/infrakit
@@ -12,10 +12,12 @@ Requires-Dist: pyyaml>=6.0.3
 Requires-Dist: typer>=0.24.1
 Provides-Extra: all
 Requires-Dist: google-genai>=1.69.0; extra == 'all'
+Requires-Dist: groq>=0.9.0; extra == 'all'
 Requires-Dist: openai>=2.30.0; extra == 'all'
 Requires-Dist: tqdm>=4.67.3; extra == 'all'
 Provides-Extra: llm
 Requires-Dist: google-genai>=1.69.0; extra == 'llm'
+Requires-Dist: groq>=0.9.0; extra == 'llm'
 Requires-Dist: openai>=2.30.0; extra == 'llm'
 Requires-Dist: tqdm>=4.67.3; extra == 'llm'
 Description-Content-Type: text/markdown

{python_infrakit_dev-0.1.5 → python_infrakit_dev-0.1.6}/infrakit/llm/__init__.py

@@ -20,7 +20,7 @@ from .models import (
     QuotaConfig,
     RequestMeta,
 )
-from .providers import GeminiProvider, OpenAIProvider
+from .providers import GeminiProvider, GroqProvider, OpenAIProvider
 
 __all__ = [
     "LLMClient",
@@ -32,4 +32,5 @@ __all__ = [
     "Provider",
     "OpenAIProvider",
     "GeminiProvider",
+    "GroqProvider",
 ]

{python_infrakit_dev-0.1.5 → python_infrakit_dev-0.1.6}/infrakit/llm/client.py

@@ -80,6 +80,7 @@ from .key_manager import KeyManager
 from .models import BatchResult, LLMResponse, Prompt, Provider, QuotaConfig, RequestMeta
 from .providers.base import BaseProvider
 from .providers.gemini import GeminiProvider
+from .providers.groq import GroqProvider
 from .providers.openai import OpenAIProvider
 from .rate_limiter import RateLimiter
 
@@ -103,6 +104,7 @@ class LLMClient:
                         {
                             "openai_keys": ["sk-key1", "sk-key2"],
                             "gemini_keys": ["AIza-key1"],
+                            "groq_keys":   ["gsk_key1"],
                         }
 
     storage_dir     Path to a folder where key state is persisted.
@@ -119,7 +121,17 @@ class LLMClient:
     meta_window     How many recent request metadata records to keep per key.
                     Default: 50.
     openai_model    Default OpenAI model.  Default: ``"gpt-4o-mini"``.
-    gemini_model    Default Gemini model.  Default: ``"gemini-1.5-flash"``.
+    gemini_model    Default Gemini model.  Default: ``"gemini-2.5-flash"``.
+    groq_model      Default Groq model.    Default: ``"llama-3.3-70b-versatile"``.
+    fallback_order  Provider priority list for automatic fallback when all keys
+                    for the primary provider are exhausted.  Example::
+
+                        fallback_order=["gemini", "groq"]
+
+                    When a ``generate()`` call fails with "all keys exhausted",
+                    the client retries each provider in *fallback_order* (skipping
+                    the already-tried primary).  Fallback only triggers on key
+                    exhaustion, not on transient errors.
     show_progress   Show tqdm progress bar during batch calls.  Default: True.
     """
 
@@ -135,9 +147,12 @@ class LLMClient:
         meta_window: int = _DEFAULT_META_WINDOW,
         openai_model: Optional[str] = None,
         gemini_model: Optional[str] = None,
+        groq_model: Optional[str] = None,
+        fallback_order: Optional[list[str]] = None,
         show_progress: bool = True,
     ) -> None:
         self._mode = mode
+        self._fallback_order: list[str] = fallback_order or []
         self._max_concurrent = max_concurrent
         self._key_retries = key_retries
         self._schema_retries = schema_retries
@@ -154,11 +169,17 @@ class LLMClient:
         # Rate limiter (RPM/TPM gating)
         self._rl = RateLimiter(self._km)
 
-        # Providers
-        self._providers: dict[str, BaseProvider] = {
-            Provider.OPENAI: OpenAIProvider(model=openai_model),
-            Provider.GEMINI: GeminiProvider(model=gemini_model),
-        }
+        # Providers — only register those whose SDK is installed
+        self._providers: dict[str, BaseProvider] = {}
+        for enum_val, cls, mdl in [
+            (Provider.OPENAI, OpenAIProvider, openai_model),
+            (Provider.GEMINI, GeminiProvider, gemini_model),
+            (Provider.GROQ,   GroqProvider,   groq_model),
+        ]:
+            try:
+                self._providers[enum_val] = cls(model=mdl)
+            except ImportError:
+                pass  # SDK not installed; skip this provider
 
     # ── public: single generate ────────────────────────────────────────────
 
@@ -178,15 +199,24 @@ class LLMClient:
         If you are inside an async function use ``await async_generate()``
         instead; that path uses the async SDK clients end-to-end.
 
-        Handles key rotation, RPM waiting, retries, and metadata recording.
+        Handles key rotation, RPM waiting, retries, metadata recording, and
+        automatic provider fallback when ``fallback_order`` is configured.
         Always returns an LLMResponse — check ``.error`` for failures.
         """
-        return self._sync_single_generate(
-            prompt=prompt,
-            response_model=response_model,
-            provider=provider,
-            **kwargs,
+        result = self._sync_single_generate(
+            prompt=prompt, response_model=response_model, provider=provider, **kwargs
         )
+        if result.error is None or not self._fallback_order:
+            return result
+        for fb_prov in self._fallback_order:
+            if fb_prov == provider or fb_prov not in self._providers:
+                continue
+            fb = self._sync_single_generate(
+                prompt=prompt, response_model=response_model, provider=fb_prov, **kwargs
+            )
+            if fb.error is None:
+                return fb
+        return result
 
     async def async_generate(
         self,
@@ -196,12 +226,20 @@ class LLMClient:
         **kwargs: Any,
     ) -> LLMResponse:
         """Async version of generate() — await this inside an async context."""
-        return await self._async_single_generate(
-            prompt=prompt,
-            response_model=response_model,
-            provider=provider,
-            **kwargs,
+        result = await self._async_single_generate(
+            prompt=prompt, response_model=response_model, provider=provider, **kwargs
         )
+        if result.error is None or not self._fallback_order:
+            return result
+        for fb_prov in self._fallback_order:
+            if fb_prov == provider or fb_prov not in self._providers:
+                continue
+            fb = await self._async_single_generate(
+                prompt=prompt, response_model=response_model, provider=fb_prov, **kwargs
+            )
+            if fb.error is None:
+                return fb
+        return result
 
     # ── public: batch generate ─────────────────────────────────────────────
 
@@ -384,6 +422,87 @@ class LLMClient:
                     )
         print(sep)
 
+    # ── public: runtime key management ────────────────────────────────────
+
+    def add_key(self, provider: str, key: str) -> None:
+        """
+        Add a new API key for *provider* at runtime.
+
+        The key is immediately available for new requests.  Adding a key
+        that is already registered is a no-op.  Thread-safe.
+
+        Parameters
+        ----------
+        provider    ``"openai"``, ``"gemini"``, or ``"groq"``.
+        key         Raw API key string.
+        """
+        self._km.add_key(provider, key)
+
+    def remove_key(self, provider: str, key_id: str) -> None:
+        """
+        Remove a key by its ``key_id`` prefix (first 8 chars).
+        No-op if not found.  Thread-safe.
+        """
+        self._km.remove_key(provider, key_id)
+
+    # ── class method: environment-variable bootstrap ───────────────────────
+
+    @classmethod
+    def from_env(
+        cls,
+        storage_dir: Optional[str | "Path"] = None,
+        quota_file: Optional[str | "Path"] = None,
+        mode: Literal["async", "threaded"] = "async",
+        max_concurrent: int = _DEFAULT_MAX_CONCURRENT,
+        key_retries: int = _DEFAULT_KEY_RETRIES,
+        schema_retries: int = _DEFAULT_SCHEMA_RETRIES,
+        meta_window: int = _DEFAULT_META_WINDOW,
+        openai_model: Optional[str] = None,
+        gemini_model: Optional[str] = None,
+        groq_model: Optional[str] = None,
+        fallback_order: Optional[list[str]] = None,
+        show_progress: bool = True,
+    ) -> "LLMClient":
+        """
+        Create an LLMClient with API keys read from environment variables.
+
+        Reads ``OPENAI_API_KEY``, ``GEMINI_API_KEY``, and ``GROQ_API_KEY``.
+        Each variable may contain a comma-separated list of keys::
+
+            OPENAI_API_KEY=sk-key1,sk-key2
+            GROQ_API_KEY=gsk-key1
+
+        Example::
+
+            client = LLMClient.from_env(fallback_order=["gemini", "groq"])
+        """
+        import os
+
+        def _split(env_var: str) -> list[str]:
+            val = os.environ.get(env_var, "").strip()
+            return [k.strip() for k in val.split(",") if k.strip()] if val else []
+
+        keys = {
+            "openai_keys": _split("OPENAI_API_KEY"),
+            "gemini_keys": _split("GEMINI_API_KEY"),
+            "groq_keys":   _split("GROQ_API_KEY"),
+        }
+        return cls(
+            keys=keys,
+            storage_dir=storage_dir,
+            quota_file=quota_file,
+            mode=mode,
+            max_concurrent=max_concurrent,
+            key_retries=key_retries,
+            schema_retries=schema_retries,
+            meta_window=meta_window,
+            openai_model=openai_model,
+            gemini_model=gemini_model,
+            groq_model=groq_model,
+            fallback_order=fallback_order,
+            show_progress=show_progress,
+        )
+
     # ── internal: async single generate ───────────────────────────────────
 
     async def _async_single_generate(
@@ -462,8 +581,8 @@ class LLMClient:
                         )
                         break  # rotate to next key/model
                     if attempt < self._key_retries:
-                        # small backoff before same-key retry
-                        await asyncio.sleep(1.0 * (attempt + 1))
+                        # exponential backoff: 1s, 2s, 4s, …
+                        await asyncio.sleep(2.0 ** attempt)
                     # else: fall through and rotate key
 
             keys_tried.add((ks.key_hash, prov_impl.model))
@@ -541,7 +660,8 @@ class LLMClient:
                         )
                         break
                     if attempt < self._key_retries:
-                        time.sleep(1.0 * (attempt + 1))
+                        # exponential backoff: 1s, 2s, 4s, …
+                        time.sleep(2.0 ** attempt)
 
             keys_tried.add((ks.key_hash, prov_impl.model))
 

{python_infrakit_dev-0.1.5 → python_infrakit_dev-0.1.6}/infrakit/llm/key_manager.py

@@ -210,7 +210,7 @@ class KeyManager:
 
     Parameters
     ----------
-    keys            ``{"openai_keys": [...], "gemini_keys": [...]}``.
+    keys            ``{"openai_keys": [...], "gemini_keys": [...], "groq_keys": [...]}``.
     storage_dir     Folder where ``key_state.json`` is written.
                     Defaults to ``~/.infrakit/llm/``.
     quota_file      Path to a JSON quota definition file.
@@ -237,7 +237,7 @@ class KeyManager:
         self._storage_path = storage_path / _STATE_FILE
 
         # quota file: explicit arg > default location > skip
-        if quota_file is not None:
+        if quota_file is not None and str(quota_file).strip():
             self._quota_file: Optional[Path] = Path(quota_file)
         elif (DEFAULT_LLM_DIR / _QUOTA_FILE).exists():
             self._quota_file = DEFAULT_LLM_DIR / _QUOTA_FILE
@@ -252,20 +252,25 @@ class KeyManager:
         self._states: dict[str, list[KeyState]] = {
             Provider.OPENAI: [],
             Provider.GEMINI: [],
+            Provider.GROQ:   [],
         }
         self._rr_index: dict[str, dict[str, int]] = {
             # per-provider, per-model round-robin index
             Provider.OPENAI: {},
             Provider.GEMINI: {},
+            Provider.GROQ:   {},
         }
 
         provider_map = {
             "openai_keys": Provider.OPENAI,
             "gemini_keys": Provider.GEMINI,
+            "groq_keys":   Provider.GROQ,
         }
 
         for key_field, provider in provider_map.items():
             for raw_key in keys.get(key_field, []):
+                if not raw_key or not isinstance(raw_key, str):
+                    continue  # skip invalid entries
                 key_hash = hashlib.sha256(raw_key.encode()).hexdigest()
                 key_id   = raw_key[:8]
 
@@ -428,6 +433,52 @@ class KeyManager:
                 ks.model_states["__all__"] = sentinel
             self._persist()
 
+    # ── public: runtime key management ────────────────────────────────────
+
+    def add_key(self, provider: str, raw_key: str) -> None:
+        """
+        Add a new API key at runtime without re-initialising the client.
+
+        The key becomes immediately available for new requests.
+        Adding a key that is already registered is a no-op (idempotent).
+        Thread-safe.
+        """
+        import hashlib
+
+        if not raw_key or not isinstance(raw_key, str):
+            raise ValueError(
+                f"Invalid API key for {provider!r}: must be a non-empty string."
+            )
+        if provider not in self._states:
+            raise ValueError(
+                f"Unknown provider {provider!r}. "
+                f"Valid options: {list(self._states)}"
+            )
+
+        with self._lock:
+            key_hash = hashlib.sha256(raw_key.encode()).hexdigest()
+            if any(ks.key_hash == key_hash for ks in self._states[provider]):
+                return  # already registered
+            key_id = raw_key[:8]
+            ks = KeyState(provider=provider, key_id=key_id, key_hash=key_hash)
+            ks._raw_key = raw_key  # type: ignore[attr-defined]
+            self._states[provider].append(ks)
+            self._persist()
+
+    def remove_key(self, provider: str, key_id: str) -> None:
+        """
+        Remove a key by its ``key_id`` prefix (first 8 chars).
+        No-op if the key is not found.  Thread-safe.
+        """
+        with self._lock:
+            before = len(self._states.get(provider, []))
+            self._states[provider] = [
+                ks for ks in self._states.get(provider, [])
+                if ks.key_id != key_id
+            ]
+            if len(self._states[provider]) < before:
+                self._persist()
+
     # ── public: quota config ───────────────────────────────────────────────
 
     def set_quota(
@@ -553,11 +604,11 @@ class KeyManager:
         """Auto-reactivate a model if its daily reset hour has passed."""
         if ms.status != ModelStatus.INACTIVE or ms.deactivated_at is None:
             return
-        now_utc = datetime.datetime.utcnow()
+        now_utc = datetime.datetime.now(datetime.timezone.utc)
         reset_today = now_utc.replace(
             hour=ms.reset_hour_utc, minute=0, second=0, microsecond=0
         )
-        deactivated_dt = datetime.datetime.utcfromtimestamp(ms.deactivated_at)
+        deactivated_dt = datetime.datetime.fromtimestamp(ms.deactivated_at, datetime.timezone.utc)
         if deactivated_dt < reset_today <= now_utc:
             ms.status         = ModelStatus.ACTIVE
             ms.deactivated_at = None
@@ -571,11 +622,11 @@ class KeyManager:
 
     def _maybe_reset_day(self, ms: ModelState) -> None:
         """Reset daily token counter if the reset hour has passed today."""
-        now_utc = datetime.datetime.utcnow()
+        now_utc = datetime.datetime.now(datetime.timezone.utc)
         reset_today = now_utc.replace(
             hour=ms.reset_hour_utc, minute=0, second=0, microsecond=0
         )
-        day_start_dt = datetime.datetime.utcfromtimestamp(ms.day_start_epoch)
+        day_start_dt = datetime.datetime.fromtimestamp(ms.day_start_epoch, datetime.timezone.utc)
         if day_start_dt < reset_today <= now_utc:
             ms.day_token_total = 0
             ms.day_start_epoch = time.time()

{python_infrakit_dev-0.1.5 → python_infrakit_dev-0.1.6}/infrakit/llm/models.py

@@ -19,6 +19,7 @@ from pydantic import BaseModel
 class Provider(str, Enum):
     OPENAI = "openai"
     GEMINI = "gemini"
+    GROQ   = "groq"
 
 
 class KeyStatus(str, Enum):

python_infrakit_dev-0.1.6/infrakit/llm/providers/__init__.py

@@ -0,0 +1,6 @@
+from .base import BaseProvider
+from .openai import OpenAIProvider
+from .gemini import GeminiProvider
+from .groq import GroqProvider
+
+__all__ = ["BaseProvider", "OpenAIProvider", "GeminiProvider", "GroqProvider"]

python_infrakit_dev-0.1.6/infrakit/llm/providers/groq.py

@@ -0,0 +1,167 @@
+"""
+infrakit.llm.providers.groq
+-----------------------------
+Groq provider — wraps the groq Python SDK.
+
+Install dependency::
+
+    pip install groq
+
+Supported models (default): llama-3.3-70b-versatile
+"""
+
+from __future__ import annotations
+
+import time
+from typing import Any, Optional, Type
+
+from pydantic import BaseModel
+
+from ..models import LLMResponse, Prompt
+from .base import BaseProvider
+
+
+class GroqProvider(BaseProvider):
+    """
+    Provider for Groq's Chat Completions API.
+
+    Parameters
+    ----------
+    model       Model string to use. Defaults to ``llama-3.3-70b-versatile``.
+    """
+
+    DEFAULT_MODEL = "llama-3.3-70b-versatile"
+    PROVIDER_NAME = "groq"
+
+    def __init__(self, model: Optional[str] = None) -> None:
+        super().__init__(model)
+        self._check_sdk()
+
+    # ── public interface ───────────────────────────────────────────────────
+
+    async def async_generate(
+        self,
+        prompt: Prompt,
+        api_key: str,
+        response_model: Optional[Type[BaseModel]] = None,
+        schema_retries: int = 2,
+        **kwargs: Any,
+    ) -> LLMResponse:
+        """Async generate using groq.AsyncGroq."""
+        from groq import AsyncGroq
+
+        client = AsyncGroq(api_key=api_key)
+        messages = self._build_messages(prompt)
+        t0 = time.perf_counter()
+
+        response = await client.chat.completions.create(
+            model=self.model,
+            messages=messages,
+            **kwargs,
+        )
+
+        latency_ms = (time.perf_counter() - t0) * 1000
+        return self._build_response(response, latency_ms, response_model, schema_retries, api_key)
+
+    def sync_generate(
+        self,
+        prompt: Prompt,
+        api_key: str,
+        response_model: Optional[Type[BaseModel]] = None,
+        schema_retries: int = 2,
+        **kwargs: Any,
+    ) -> LLMResponse:
+        """Sync generate using groq.Groq (the blocking SDK client)."""
+        from groq import Groq
+
+        client = Groq(api_key=api_key)
+        messages = self._build_messages(prompt)
+        t0 = time.perf_counter()
+
+        response = client.chat.completions.create(
+            model=self.model,
+            messages=messages,
+            **kwargs,
+        )
+
+        latency_ms = (time.perf_counter() - t0) * 1000
+        return self._build_response(response, latency_ms, response_model, schema_retries, api_key)
+
+    # ── internal helpers ───────────────────────────────────────────────────
+
+    @staticmethod
+    def _build_messages(prompt: Prompt) -> list[dict]:
+        messages = []
+        if prompt.system:
+            messages.append({"role": "system", "content": prompt.system})
+        messages.append({"role": "user", "content": prompt.user})
+        return messages
+
+    def _build_response(
+        self,
+        response: Any,
+        latency_ms: float,
+        response_model: Optional[Type[BaseModel]],
+        schema_retries: int,
+        api_key: str,
+    ) -> LLMResponse:
+        choice = response.choices[0]
+        content = choice.message.content or ""
+
+        usage = response.usage
+        input_tokens = usage.prompt_tokens if usage else 0
+        output_tokens = usage.completion_tokens if usage else 0
+        total_tokens = usage.total_tokens if usage else 0
+
+        parsed = None
+        schema_matched = False
+
+        if response_model is not None:
+            parsed, schema_matched = self._validate_schema(
+                content, response_model, schema_retries
+            )
+
+        return LLMResponse(
+            content=content,
+            parsed=parsed,
+            schema_matched=schema_matched,
+            provider=self.PROVIDER_NAME,
+            model=self.model,
+            key_id=api_key[:8],
+            input_tokens=input_tokens,
+            output_tokens=output_tokens,
+            total_tokens=total_tokens,
+            latency_ms=latency_ms,
+        )
+
+    @staticmethod
+    def _check_sdk() -> None:
+        try:
+            import groq  # noqa: F401
+        except ImportError as exc:
+            raise ImportError(
+                "groq package is required for GroqProvider. "
+                "Install it with: pip install groq"
+            ) from exc
+
+    @staticmethod
+    def _is_quota_error(exc: Exception) -> bool:
+        """
+        Groq-specific quota/auth HTTP status codes.
+
+        429 with "rate_limit_exceeded" is *transient* (burst limit) and should
+        be retried with backoff — NOT treated as a permanent quota exhaustion.
+        Only 401/402 and 429 with "quota" or "billing" in the message warrant
+        immediate model deactivation.
+        """
+        try:
+            from groq import APIStatusError
+            if isinstance(exc, APIStatusError):
+                if exc.status_code in (401, 402):
+                    return True
+                if exc.status_code == 429:
+                    msg = str(exc).lower()
+                    return "quota" in msg or "billing" in msg
+        except ImportError:
+            pass
+        return BaseProvider._is_quota_error(exc)

{python_infrakit_dev-0.1.5 → python_infrakit_dev-0.1.6}/infrakit/scaffolder/ai.py

@@ -40,6 +40,8 @@ from infrakit.scaffolder.generator import (
     _write,
     _config_content,
     _gitignore,
+    _infrakit_dep,
+    _pkg_dep,
     _logger_util,
     _src_init,
     _tests_init,
@@ -124,9 +126,11 @@ def _load_keys() -> dict:
     # fall back to keys declared in the config file
     openai_key  = _cfg.get("OPENAI_API_KEY", "")
     gemini_key  = _cfg.get("GEMINI_API_KEY", "")
+    groq_key    = _cfg.get("GROQ_API_KEY", "")
     return {{
         "openai_keys": [openai_key] if openai_key else [],
         "gemini_keys": [gemini_key] if gemini_key else [],
+        "groq_keys":   [groq_key]   if groq_key   else [],
     }}
 
 
@@ -142,6 +146,7 @@ llm: LLMClient = LLMClient(
     max_concurrent=int(_cfg.get("LLM_CONCURRENCY", 3)),
     openai_model=_cfg.get("OPENAI_MODEL") or None,
     gemini_model=_cfg.get("GEMINI_MODEL") or None,
+    groq_model=_cfg.get("GROQ_MODEL") or None,
 )
 
 __all__ = ["llm", "Prompt"]
@@ -316,7 +321,8 @@ def _keys_json_template() -> str:
 {
   "_comment": "Fill in your API keys. Never commit this file — it is in .gitignore.",
   "openai_keys": [],
-  "gemini_keys":  []
+  "gemini_keys": [],
+  "groq_keys":   []
 }
 """
 
@@ -401,6 +407,12 @@ pytest
 
 def _ai_pyproject(project_name: str, version: str, description: str, author: str) -> str:
     author_line = f'    "{author}",' if author else '    # "Your Name <you@example.com>",'
+    infrakit_dep = _infrakit_dep()
+    openai_dep   = _pkg_dep("openai")
+    genai_dep    = _pkg_dep("google-genai")
+    groq_dep     = _pkg_dep("groq")
+    pydantic_dep = _pkg_dep("pydantic")
+    tqdm_dep     = _pkg_dep("tqdm")
     return f"""\
 [project]
 name        = "{project_name}"
@@ -413,11 +425,12 @@ authors = [
 ]
 
 dependencies = [
-    "infrakit",
-    "openai",
-    "google-genai",
-    "pydantic>=2.0",
-    "tqdm",
+    {infrakit_dep},
+    "{openai_dep}",
+    "{genai_dep}",
+    "{groq_dep}",
+    "{pydantic_dep}",
+    "{tqdm_dep}",
 ]
 
 [project.optional-dependencies]