python-infrakit-dev 0.1.0__py3-none-any.whl

infrakit/llm/providers/gemini.py

@@ -0,0 +1,164 @@
+"""
+infrakit.llm.providers.gemini
+------------------------------
+Gemini provider — wraps the google-generativeai Python SDK.
+
+Install dependency::
+
+    pip install google-genai
+
+Supported models (default): gemini-2.5-flash
+"""
+
+from __future__ import annotations
+
+import asyncio
+import time
+from typing import Any, Optional, Type
+
+from pydantic import BaseModel
+
+from ..models import LLMResponse, Prompt
+from .base import BaseProvider
+
+
+class GeminiProvider(BaseProvider):
+    """
+    Provider for Google's Gemini generative models.
+
+    Parameters
+    ----------
+    model       Model string to use. Defaults to ``gemini-1.5-flash``.
+    """
+
+    DEFAULT_MODEL = "gemini-2.5-flash"
+    PROVIDER_NAME = "gemini"
+
+    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 google-generativeai.
+
+        The SDK's generate_content is synchronous; we run it in a thread
+        executor to avoid blocking the event loop.
+        """
+        loop = asyncio.get_event_loop()
+        t0 = time.perf_counter()
+
+        result = await loop.run_in_executor(
+            None,
+            lambda: self._sync_call(prompt, api_key, **kwargs),
+        )
+        latency_ms = (time.perf_counter() - t0) * 1000
+        return self._build_response(result, 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 — direct SDK call."""
+        t0 = time.perf_counter()
+        result = self._sync_call(prompt, api_key, **kwargs)
+        latency_ms = (time.perf_counter() - t0) * 1000
+        return self._build_response(result, latency_ms, response_model, schema_retries, api_key)
+
+    # ── internal helpers ───────────────────────────────────────────────────
+
+    def _sync_call(self, prompt: Prompt, api_key: str, **kwargs: Any) -> Any:
+        from google import genai
+        from google.genai import types
+
+        client = genai.Client(api_key=api_key)
+
+        # The new SDK uses a config object for arguments like system instructions,
+        # temperature, etc. We map your kwargs directly into it.
+        config_args = kwargs.copy()
+        if prompt.system:
+            config_args["system_instruction"] = prompt.system
+
+        config = types.GenerateContentConfig(**config_args) if config_args else None
+
+        response = client.models.generate_content(
+            model=self.model,
+            contents=prompt.user,
+            config=config,
+        )
+        return response
+
+    def _build_response(
+        self,
+        response: Any,
+        latency_ms: float,
+        response_model: Optional[Type[BaseModel]],
+        schema_retries: int,
+        api_key: str,
+    ) -> LLMResponse:
+        # Extract text
+        try:
+            content = response.text or ""
+        except Exception:
+            content = ""
+
+        # Token counts — available in usage_metadata (may be None on some models)
+        usage = getattr(response, "usage_metadata", None)
+        input_tokens = getattr(usage, "prompt_token_count", 0) or 0
+        output_tokens = getattr(usage, "candidates_token_count", 0) or 0
+        total_tokens = getattr(usage, "total_token_count", 0) or (input_tokens + output_tokens)
+
+        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:
+            from google import genai  # noqa: F401
+        except ImportError as exc:
+            raise ImportError(
+                "google-genai package is required for GeminiProvider. "
+                "Install it with: pip install google-genai"
+            ) from exc
+
+    @staticmethod
+    def _is_quota_error(exc: Exception) -> bool:
+        """Gemini-specific quota/auth error detection."""
+        # Handle the new SDK's APIError exception
+        if exc.__class__.__name__ == "APIError":
+            code = getattr(exc, "code", None)
+            # 429: Resource Exhausted, 403: Permission Denied, 401: Unauthenticated
+            if code in (429, 403, 401):
+                return True
+        return BaseProvider._is_quota_error(exc)

infrakit/llm/providers/openai.py

@@ -0,0 +1,168 @@
+"""
+infrakit.llm.providers.openai
+------------------------------
+OpenAI provider — wraps the openai Python SDK.
+
+Install dependency::
+
+    pip install openai
+
+Supported models (default): gpt-4o-mini
+"""
+
+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 OpenAIProvider(BaseProvider):
+    """
+    Provider for OpenAI's Chat Completions API.
+
+    Parameters
+    ----------
+    model       Model string to use. Defaults to ``gpt-4o-mini``.
+    """
+
+    DEFAULT_MODEL = "gpt-4o-mini"
+    PROVIDER_NAME = "openai"
+
+    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 openai.AsyncOpenAI."""
+        from openai import AsyncOpenAI
+
+        client = AsyncOpenAI(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 openai.OpenAI (the blocking SDK client).
+
+        Deliberately does NOT use asyncio.run() so it is safe to call from:
+        - threads (threaded batch mode)
+        - environments with a running event loop (Jupyter, FastAPI, etc.)
+        - Windows, where asyncio loop/thread interactions are more restrictive
+        """
+        from openai import OpenAI
+
+        client = OpenAI(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 openai  # noqa: F401
+        except ImportError as exc:
+            raise ImportError(
+                "openai package is required for OpenAIProvider. "
+                "Install it with: pip install openai"
+            ) from exc
+
+    @staticmethod
+    def _is_quota_error(exc: Exception) -> bool:
+        """OpenAI-specific quota/auth HTTP status codes."""
+        try:
+            from openai import APIStatusError
+            if isinstance(exc, APIStatusError):
+                # 401 invalid key, 402 billing, 429 with quota message
+                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)

infrakit/llm/rate_limiter.py

@@ -0,0 +1,54 @@
+"""
+infrakit.llm.rate_limiter
+-------------------------
+RPM-aware gate (key-level) and TPM-aware gate (model-level).
+"""
+
+from __future__ import annotations
+
+import asyncio
+import time
+
+from .key_manager import KeyManager, KeyState
+
+
+class RateLimiter:
+    """
+    Async-and-sync compatible RPM/TPM gate.
+
+    RPM is enforced at the key level (all models share the key's request budget).
+    TPM is enforced at the model level.
+
+    Parameters
+    ----------
+    key_manager     Shared KeyManager instance (owns all rate state).
+    poll_interval   How often (seconds) to re-check the window while waiting.
+    """
+
+    def __init__(
+        self,
+        key_manager: KeyManager,
+        poll_interval: float = 0.25,
+    ) -> None:
+        self._km   = key_manager
+        self._poll = poll_interval
+
+    # ── async ──────────────────────────────────────────────────────────────
+
+    async def async_wait_for_slot(self, ks: KeyState, model: str) -> None:
+        """Async-suspend until both RPM and TPM slots are open."""
+        while not self._km.check_rpm(ks):
+            wait = self._km.seconds_until_rpm_slot(ks)
+            await asyncio.sleep(max(wait, self._poll))
+        while not self._km.check_tpm(ks, model):
+            await asyncio.sleep(self._poll)
+
+    # ── sync ───────────────────────────────────────────────────────────────
+
+    def sync_wait_for_slot(self, ks: KeyState, model: str) -> None:
+        """Blocking RPM+TPM wait. Used in sync / threaded batch mode."""
+        while not self._km.check_rpm(ks):
+            wait = self._km.seconds_until_rpm_slot(ks)
+            time.sleep(max(wait, self._poll))
+        while not self._km.check_tpm(ks, model):
+            time.sleep(self._poll)

infrakit/scaffolder/__init__.py

@@ -0,0 +1,31 @@
+"""
+infrakit.scaffolder.templates
+------------------------------
+Project template scaffolders.
+
+Quick reference
+---------------
+    from infrakit.scaffolder.templates.registry import get_template, list_templates
+    from infrakit.scaffolder.templates.ai       import scaffold_ai
+    from infrakit.scaffolder.templates.backend  import scaffold_backend
+    from infrakit.scaffolder.templates.cli_tool import scaffold_cli_tool
+    from infrakit.scaffolder.templates.pipeline import scaffold_pipeline
+"""
+from infrakit.scaffolder.generator import scaffold_basic, ScaffoldEntry, ScaffoldResult
+from infrakit.scaffolder.ai       import scaffold_ai
+from infrakit.scaffolder.backend  import scaffold_backend
+from infrakit.scaffolder.cli_tool import scaffold_cli_tool
+from infrakit.scaffolder.pipeline import scaffold_pipeline
+from infrakit.scaffolder.registry import get_template, list_templates
+
+__all__ = [
+    'scaffold_basic',
+    'ScaffoldEntry',
+    'ScaffoldResult',
+    "scaffold_ai",
+    "scaffold_backend",
+    "scaffold_cli_tool",
+    "scaffold_pipeline",
+    "get_template",
+    "list_templates",
+]