forgeoptimizer 0.1.0__py3-none-any.whl

forgecli/providers/base.py

@@ -0,0 +1,204 @@
+"""Abstract base classes and shared data types for AI providers."""
+
+from __future__ import annotations
+
+import os
+from abc import ABC, abstractmethod
+from collections.abc import AsyncIterator, Iterator
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import Any, ClassVar, Generic, TypeVar
+
+from pydantic import BaseModel, ConfigDict, Field
+
+from forgecli.core.errors import ProviderError
+
+
+class Role(str, Enum):
+    """A role in a chat conversation."""
+
+    SYSTEM = "system"
+    USER = "user"
+    ASSISTANT = "assistant"
+    TOOL = "tool"
+
+
+@dataclass(frozen=True)
+class ChatMessage:
+    """A single message in a chat conversation."""
+
+    role: Role
+    content: str
+    name: str | None = None
+    tool_call_id: str | None = None
+    extra: dict[str, Any] = field(default_factory=dict)
+
+
+class ChatRequest(BaseModel):
+    """Provider-agnostic chat completion request."""
+
+    model_config = ConfigDict(extra="allow")
+
+    model: str | None = None
+    messages: list[ChatMessage] = Field(default_factory=list)
+    temperature: float | None = None
+    max_tokens: int | None = None
+    top_p: float | None = None
+    stop: list[str] | None = None
+    tools: list[dict[str, Any]] | None = None
+    metadata: dict[str, Any] = Field(default_factory=dict)
+
+
+class ChatResponse(BaseModel):
+    """Provider-agnostic chat completion response."""
+
+    model_config = ConfigDict(extra="allow")
+
+    model: str
+    message: ChatMessage
+    finish_reason: str | None = None
+    prompt_tokens: int = 0
+    completion_tokens: int = 0
+    total_tokens: int = 0
+    raw: dict[str, Any] = Field(default_factory=dict)
+
+
+@dataclass(frozen=True)
+class StreamChunk:
+    """A piece of a streamed chat response."""
+
+    delta: str
+    finish_reason: str | None = None
+    raw: dict[str, Any] = field(default_factory=dict)
+
+
+class EmbeddingRequest(BaseModel):
+    """Provider-agnostic embedding request."""
+
+    model_config = ConfigDict(extra="allow")
+
+    model: str | None = None
+    inputs: list[str] = Field(default_factory=list)
+
+
+class EmbeddingResponse(BaseModel):
+    """Provider-agnostic embedding response."""
+
+    model_config = ConfigDict(extra="allow")
+
+    model: str
+    vectors: list[list[float]] = Field(default_factory=list)
+    prompt_tokens: int = 0
+    total_tokens: int = 0
+
+
+@dataclass(frozen=True)
+class ModelInfo:
+    """Metadata describing a model exposed by a provider."""
+
+    id: str
+    name: str | None = None
+    context_window: int | None = None
+    supports_tools: bool = False
+    supports_vision: bool = False
+    extra: dict[str, Any] = field(default_factory=dict)
+
+
+TConfig = TypeVar("TConfig")
+
+
+class Provider(ABC, Generic[TConfig]):
+    """Base class for AI providers.
+
+    Concrete providers should subclass :class:`Provider` with a concrete
+    ``TConfig`` type and implement the abstract methods below.
+    """
+
+    name: ClassVar[str] = "abstract"
+
+    def __init__(self, config: TConfig) -> None:
+        self._config = config
+
+    @property
+    def config(self) -> TConfig:
+        return self._config
+
+    @abstractmethod
+    async def chat(self, request: ChatRequest) -> ChatResponse:
+        """Send ``request`` and return a :class:`ChatResponse`."""
+
+    async def stream(self, request: ChatRequest) -> AsyncIterator[StreamChunk]:
+        """Stream a chat response. Default falls back to a single chunk."""
+        response = await self.chat(request)
+        yield StreamChunk(
+            delta=response.message.content,
+            finish_reason=response.finish_reason,
+            raw=response.raw,
+        )
+        return
+
+    @abstractmethod
+    async def embed(self, request: EmbeddingRequest) -> EmbeddingResponse:
+        """Compute embeddings for ``request.inputs``."""
+
+    async def list_models(self) -> list[ModelInfo]:
+        """Return the list of models supported by this provider."""
+        return []
+
+    def validate(self) -> None:
+        """Hook for pre-flight checks (API key, base URL, etc)."""
+        return None
+
+    @staticmethod
+    def resolve_api_key(env_var: str | None) -> str | None:
+        """Return the API key from ``env_var`` or ``None``."""
+        if not env_var:
+            return None
+        return os.environ.get(env_var)
+
+
+class ProviderRegistry:
+    """In-memory registry of provider classes keyed by name."""
+
+    def __init__(self) -> None:
+        self._providers: dict[str, type[Provider[Any]]] = {}
+
+    def register(self, name: str, provider_cls: type[Provider[Any]]) -> None:
+        if not issubclass(provider_cls, Provider):
+            raise ProviderError(f"{provider_cls!r} must subclass forgecli.providers.Provider")
+        if name in self._providers and self._providers[name] is not provider_cls:
+            raise ProviderError(f"Provider {name!r} already registered")
+        self._providers[name] = provider_cls
+
+    def unregister(self, name: str) -> None:
+        self._providers.pop(name, None)
+
+    def create(self, name: str, config: Any) -> Provider[Any]:
+        cls = self._providers.get(name)
+        if cls is None:
+            raise ProviderError(f"Unknown provider: {name!r}")
+        return cls(config)
+
+    def get(self, name: str) -> type[Provider[Any]]:
+        cls = self._providers.get(name)
+        if cls is None:
+            raise ProviderError(f"Unknown provider: {name!r}")
+        return cls
+
+    def names(self) -> list[str]:
+        return sorted(self._providers)
+
+    def has(self, name: str) -> bool:
+        return name in self._providers
+
+
+# A default registry instance used by ``AppContext`` wiring.
+default_registry = ProviderRegistry()
+
+
+def iter_chunked(items: list[str], size: int) -> Iterator[list[str]]:
+    """Yield ``items`` in fixed-size chunks; useful for batched embedding calls."""
+    if size <= 0:
+        raise ValueError("chunk size must be positive")
+    for i in range(0, len(items), size):
+        yield items[i : i + size]

forgecli/providers/builtin.py

@@ -0,0 +1,26 @@
+"""Built-in providers.
+
+The package ships with the following implementations:
+
+* :class:`MockProvider` — deterministic, offline-friendly, used in tests.
+* :class:`OpenAIProvider` — OpenAI Chat Completions + Embeddings.
+* :class:`AnthropicProvider` — Anthropic Messages API.
+* :class:`GeminiProvider` — Google Gemini generateContent + embedContent.
+
+All four register themselves into :data:`default_registry` on import.
+"""
+
+from forgecli.providers.anthropic import AnthropicConfig, AnthropicProvider
+from forgecli.providers.google import GeminiConfig, GeminiProvider
+from forgecli.providers.mock import MockProvider
+from forgecli.providers.openai import OpenAIConfig, OpenAIProvider
+
+__all__ = [
+    "AnthropicConfig",
+    "AnthropicProvider",
+    "GeminiConfig",
+    "GeminiProvider",
+    "MockProvider",
+    "OpenAIConfig",
+    "OpenAIProvider",
+]

forgecli/providers/conversation.py

@@ -0,0 +1,74 @@
+"""Shared helpers for conversational mock/offline responses."""
+
+from __future__ import annotations
+
+import re
+
+_GREETINGS = frozenset(
+    {
+        "hi",
+        "hello",
+        "hey",
+        "howdy",
+        "greetings",
+        "yo",
+        "sup",
+        "how",
+        "are",
+        "you",
+        "morning",
+        "afternoon",
+        "evening",
+        "good",
+        "what's",
+        "up",
+        "whats",
+    }
+)
+
+_GREETING_PHRASES = (
+    "hi",
+    "hello",
+    "hey",
+    "howdy",
+    "good morning",
+    "good afternoon",
+    "good evening",
+    "how are you",
+    "what's up",
+    "whats up",
+)
+
+
+def is_greeting(text: str) -> bool:
+    """Return True when ``text`` looks like a short greeting."""
+    normalized = (text or "").strip().lower()
+    if not normalized:
+        return False
+    if normalized in _GREETING_PHRASES:
+        return True
+    words = re.findall(r"\b\w+\b", normalized)
+    if not words or len(words) > 4:
+        return False
+    return any(word in _GREETINGS for word in words)
+
+
+def greeting_reply(text: str) -> str:
+    """Return a natural assistant greeting."""
+    normalized = (text or "").strip().lower()
+    if "how are you" in normalized:
+        return "I'm doing well, thanks for asking. How can I help with your project today?"
+    if any(p in normalized for p in ("good morning", "good afternoon", "good evening")):
+        return "Hello! How can I help you today?"
+    return "Hi! How can I help you today?"
+
+
+def offline_build_notice() -> str:
+    """Message returned when a build is requested without a configured provider."""
+    return (
+        "No AI provider is configured. Configure a provider and retry, "
+        "or pass `--mock` for offline mode."
+    )
+
+
+__all__ = ["greeting_reply", "is_greeting", "offline_build_notice"]

forgecli/providers/google.py

@@ -0,0 +1,290 @@
+"""Google Gemini (Generative Language API) provider.
+
+Targets the public ``https://generativelanguage.googleapis.com/v1beta``
+endpoint. The provider follows the modern ``generateContent`` schema,
+translating ChatRequest/Response into Gemini's ``contents`` array and
+mapping the ``usageMetadata`` token counters back to the unified
+format.
+
+API keys are read from ``GOOGLE_API_KEY`` (with ``GEMINI_API_KEY`` as
+a fallback) and are passed as the ``key`` query parameter, matching
+the public REST contract.
+"""
+
+from __future__ import annotations
+
+from collections.abc import AsyncIterator
+from dataclasses import dataclass
+from typing import Any
+
+import httpx
+
+from forgecli.providers.base import (
+    ChatMessage,
+    ChatRequest,
+    ChatResponse,
+    EmbeddingRequest,
+    EmbeddingResponse,
+    ModelInfo,
+    Role,
+    StreamChunk,
+)
+from forgecli.providers.http_base import HTTPChatProvider
+
+
+@dataclass
+class GeminiConfig:
+    """Configuration for the Google Gemini provider."""
+
+    api_key_env: str = "GOOGLE_API_KEY"
+    base_url: str = "https://generativelanguage.googleapis.com/v1beta"
+    default_model: str = "gemini-1.5-flash"
+    max_tokens: int = 4096
+    temperature: float = 0.2
+    embeddings_model: str = "text-embedding-004"
+
+
+_GEMINI_ROLE_MAP: dict[Role, str] = {
+    Role.SYSTEM: "user",  # Gemini has no "system" role; fold into a user turn
+    Role.USER: "user",
+    Role.ASSISTANT: "model",
+    Role.TOOL: "user",
+}
+
+
+class GeminiProvider(HTTPChatProvider):
+    """Google Gemini generateContent + embedContent provider."""
+
+    name = "google"
+
+    def __init__(
+        self,
+        config: GeminiConfig | None = None,
+        *,
+        api_key: str | None = None,
+        base_url: str | None = None,
+        client: httpx.AsyncClient | None = None,
+    ) -> None:
+        super().__init__(
+            config or GeminiConfig(),
+            api_key=api_key,
+            base_url=base_url,
+            client=client,
+        )
+        if not self._api_key:
+            # Fall back to the alternate env var; this is intentional.
+            self._api_key = _fallback_api_key(self.config.api_key_env)
+
+    def _default_base_url(self) -> str:
+        return "https://generativelanguage.googleapis.com/v1beta"
+
+    def _chat_url(self) -> str:
+        # The Gemini chat URL embeds the model name; we use
+        # ``_chat_url_for`` instead. The fallback is here so abstract
+        # callers can still get a URL when they have no request.
+        return (
+            f"{self._base_url}/models/{self.config.default_model}:generateContent"
+            f"?key={self._api_key or ''}"
+        )
+
+    def _chat_url_for(self, request) -> str:  # type: ignore[override]
+        model = request.model or self.config.default_model
+        return f"{self._base_url}/models/{model}:generateContent?key={self._api_key or ''}"
+
+    def _embeddings_url(self) -> str:
+        return (
+            f"{self._base_url}/models/{self.config.embeddings_model}:embedContent"
+            f"?key={self._api_key or ''}"
+        )
+
+    def _default_embedding_model(self) -> str:
+        return self.config.embeddings_model
+
+    # ------------------------------------------------------------------
+    # Auth
+    # ------------------------------------------------------------------
+
+    def _auth_headers(self) -> dict[str, str]:
+        # Gemini uses the API key in the query string; no headers needed.
+        return {"Content-Type": "application/json"}
+
+    # ------------------------------------------------------------------
+    # Request / response
+    # ------------------------------------------------------------------
+
+    def _format_request(self, request: ChatRequest) -> dict[str, Any]:
+        contents: list[dict[str, Any]] = []
+        system_parts: list[str] = []
+        for message in request.messages:
+            if message.role is Role.SYSTEM:
+                if message.content:
+                    system_parts.append(message.content)
+                continue
+            contents.append(
+                {
+                    "role": _GEMINI_ROLE_MAP[message.role],
+                    "parts": [{"text": message.content}],
+                }
+            )
+        body: dict[str, Any] = {
+            "contents": contents,
+            "generationConfig": {
+                "maxOutputTokens": request.max_tokens or self.config.max_tokens,
+                "temperature": (
+                    request.temperature
+                    if request.temperature is not None
+                    else self.config.temperature
+                ),
+            },
+        }
+        if system_parts:
+            body["systemInstruction"] = {"parts": [{"text": "\n\n".join(system_parts)}]}
+        if request.top_p is not None:
+            body["generationConfig"]["topP"] = request.top_p
+        if request.stop:
+            body["generationConfig"]["stopSequences"] = list(request.stop)
+        return body
+
+    def _parse_response(self, payload: dict[str, Any]) -> ChatResponse:
+        candidates = payload.get("candidates") or []
+        first = candidates[0] if candidates else {}
+        content = first.get("content") or {}
+        parts = content.get("parts") or []
+        text = "".join(str(p.get("text", "")) for p in parts if isinstance(p, dict))
+        usage = payload.get("usageMetadata") or {}
+        return ChatResponse(
+            model=str(payload.get("modelVersion", self.config.default_model)),
+            message=ChatMessage(role=Role.ASSISTANT, content=text),
+            finish_reason=first.get("finishReason"),
+            prompt_tokens=int(usage.get("promptTokenCount", 0) or 0),
+            completion_tokens=int(usage.get("candidatesTokenCount", 0) or 0),
+            total_tokens=int(usage.get("totalTokenCount", 0) or 0),
+            raw=payload,
+        )
+
+    async def stream(self, request: ChatRequest) -> AsyncIterator[StreamChunk]:
+        from forgecli.core.errors import ProviderError
+
+        if not self._api_key:
+            self.validate()
+        body = self._format_request(request)
+        model = request.model or self.config.default_model
+        url = f"{self._base_url}/models/{model}:streamGenerateContent?key={self._api_key or ''}"
+        req = self._client.build_request("POST", url, json=body, headers=self._auth_headers())
+        response = await self._client.send(req, stream=True)
+        if response.status_code >= 400:
+            await response.aread()
+            raise ProviderError(
+                f"{self.name} stream failed ({response.status_code}): {response.text[:500]}"
+            )
+
+        import json
+
+        try:
+            async for line in response.aiter_lines():
+                line = line.strip()
+                if not line:
+                    continue
+                if line.startswith(","):
+                    line = line[1:].strip()
+                if line.startswith("[") or line.endswith("]"):
+                    line = line.strip("[]").strip()
+                if not line:
+                    continue
+                try:
+                    payload = json.loads(line)
+                    candidates = payload.get("candidates") or []
+                    if candidates:
+                        first = candidates[0]
+                        content = first.get("content") or {}
+                        parts = content.get("parts") or []
+                        text = "".join(str(p.get("text", "")) for p in parts if isinstance(p, dict))
+                        finish_reason = first.get("finishReason")
+                        yield StreamChunk(delta=text, finish_reason=finish_reason, raw=payload)
+                except Exception:
+                    pass
+        finally:
+            await response.aclose()
+
+    def _format_embeddings(self, request: EmbeddingRequest) -> dict[str, Any]:
+        # The Gemini embedContent endpoint takes a single ``content``
+        # object; we issue one request per input to keep the wire
+        # format simple. Callers that need batching should compose it
+        # at a higher level.
+        if len(request.inputs) != 1:
+            from forgecli.core.errors import ProviderError
+
+            raise ProviderError(
+                f"{self.name} embed() expects exactly one input; got {len(request.inputs)}"
+            )
+        return {
+            "content": {
+                "parts": [{"text": request.inputs[0]}],
+            }
+        }
+
+    def _parse_embeddings(self, payload: dict[str, Any]) -> EmbeddingResponse:
+        embedding = payload.get("embedding") or {}
+        values = list(embedding.get("values", []))
+        return EmbeddingResponse(
+            model=str(payload.get("model", self.config.embeddings_model)),
+            vectors=[values],
+        )
+
+    async def list_models(self) -> list[ModelInfo]:
+        try:
+            url = f"{self._base_url}/models?key={self._api_key or ''}"
+            response = await self._client.get(url, timeout=5.0)
+            if response.status_code == 200:
+                data = response.json()
+                models_list = data.get("models", [])
+                if models_list:
+                    out = []
+                    for m in models_list:
+                        m_name = m.get("name", "")
+                        if m_name.startswith("models/"):
+                            m_id = m_name.replace("models/", "")
+                            out.append(
+                                ModelInfo(
+                                    id=m_id,
+                                    name=m.get("displayName", m_id),
+                                    context_window=m.get("inputTokenLimit", 1_000_000),
+                                    supports_tools=True,
+                                    supports_vision=True,
+                                )
+                            )
+                    if out:
+                        return out
+        except Exception:
+            pass
+        return self._known_models()
+
+    def _known_models(self) -> list[ModelInfo]:
+        from forgecli.core.models import MODEL_CATALOG
+
+        return [
+            ModelInfo(
+                id=m.id,
+                name=m.display_name,
+                context_window=2_000_000 if "pro" in m.id else 1_000_000,
+                supports_tools=True,
+                supports_vision=True,
+            )
+            for m in MODEL_CATALOG
+            if m.provider == "google"
+        ]
+
+
+def _fallback_api_key(primary: str) -> str | None:
+    """Read ``primary`` or ``GEMINI_API_KEY`` from the environment, or secure storage."""
+    import os
+
+    env_val = os.environ.get(primary) or os.environ.get("GEMINI_API_KEY")
+    if env_val:
+        return env_val
+    from forgecli.core.credentials import get_api_key
+
+    return get_api_key("google") or get_api_key("gemini")
+
+
+__all__ = ["GeminiConfig", "GeminiProvider"]