ebk 0.4.4__py3-none-any.whl

ebk/ai/llm_providers/gemini.py

@@ -0,0 +1,285 @@
+"""
+Google Gemini LLM Provider.
+
+Supports Gemini models via the Google AI API.
+"""
+
+import json
+from typing import Dict, Any, List, Optional, AsyncIterator
+import httpx
+
+from .base import BaseLLMProvider, LLMConfig, LLMResponse, ModelCapability
+
+
+class GeminiProvider(BaseLLMProvider):
+    """
+    Google Gemini LLM provider.
+
+    Supports:
+    - Gemini 1.5 models (flash, pro)
+    - Gemini 2.0 models
+    - Streaming completions
+    - JSON mode
+    - Embeddings
+    """
+
+    DEFAULT_MODEL = "gemini-1.5-flash"
+
+    def __init__(self, config: LLMConfig):
+        """
+        Initialize Gemini provider.
+
+        Args:
+            config: LLM configuration with api_key for Google AI
+        """
+        super().__init__(config)
+        if not config.api_key:
+            raise ValueError("Google AI API key is required")
+
+    @property
+    def name(self) -> str:
+        return "gemini"
+
+    @property
+    def supported_capabilities(self) -> List[ModelCapability]:
+        return [
+            ModelCapability.TEXT_GENERATION,
+            ModelCapability.JSON_MODE,
+            ModelCapability.STREAMING,
+            ModelCapability.EMBEDDINGS,
+            ModelCapability.VISION,
+        ]
+
+    @classmethod
+    def create(
+        cls,
+        api_key: str,
+        model: str = DEFAULT_MODEL,
+        **kwargs
+    ) -> 'GeminiProvider':
+        """
+        Create a Gemini provider.
+
+        Args:
+            api_key: Google AI API key
+            model: Model name (e.g., 'gemini-1.5-flash', 'gemini-1.5-pro', 'gemini-2.0-flash-exp')
+            **kwargs: Additional config parameters
+
+        Returns:
+            Configured GeminiProvider
+
+        Example:
+            >>> provider = GeminiProvider.create(
+            ...     api_key="AIza...",
+            ...     model="gemini-1.5-flash"
+            ... )
+        """
+        config = LLMConfig(
+            base_url="https://generativelanguage.googleapis.com",
+            api_key=api_key,
+            model=model,
+            **kwargs
+        )
+        return cls(config)
+
+    async def initialize(self) -> None:
+        """Initialize HTTP client."""
+        self._client = httpx.AsyncClient(
+            base_url=self.config.base_url,
+            timeout=self.config.timeout,
+        )
+
+    async def cleanup(self) -> None:
+        """Close HTTP client."""
+        if self._client:
+            await self._client.aclose()
+            self._client = None
+
+    def _build_endpoint(self, action: str = "generateContent") -> str:
+        """Build API endpoint with model and API key."""
+        return f"/v1beta/models/{self.config.model}:{action}?key={self.config.api_key}"
+
+    async def complete(
+        self,
+        prompt: str,
+        system_prompt: Optional[str] = None,
+        **kwargs
+    ) -> LLMResponse:
+        """
+        Generate completion using Gemini.
+
+        Args:
+            prompt: User prompt
+            system_prompt: Optional system prompt
+            **kwargs: Additional parameters (temperature, max_tokens, etc.)
+
+        Returns:
+            LLMResponse with generated text
+        """
+        if not self._client:
+            await self.initialize()
+
+        # Build request payload - Gemini uses contents format
+        contents = []
+
+        # Add system prompt as first user message if provided
+        if system_prompt:
+            contents.append({
+                "role": "user",
+                "parts": [{"text": f"System instruction: {system_prompt}"}]
+            })
+            contents.append({
+                "role": "model",
+                "parts": [{"text": "Understood. I will follow those instructions."}]
+            })
+
+        contents.append({
+            "role": "user",
+            "parts": [{"text": prompt}]
+        })
+
+        data = {
+            "contents": contents,
+            "generationConfig": {
+                "temperature": kwargs.get("temperature", self.config.temperature),
+                "topP": kwargs.get("top_p", self.config.top_p),
+            }
+        }
+
+        if self.config.max_tokens or kwargs.get("max_tokens"):
+            data["generationConfig"]["maxOutputTokens"] = kwargs.get("max_tokens", self.config.max_tokens)
+
+        # Make request
+        endpoint = self._build_endpoint()
+        response = await self._client.post(endpoint, json=data)
+        response.raise_for_status()
+
+        result = response.json()
+
+        # Extract content from response
+        content = ""
+        candidates = result.get("candidates", [])
+        if candidates:
+            parts = candidates[0].get("content", {}).get("parts", [])
+            for part in parts:
+                if "text" in part:
+                    content += part["text"]
+
+        # Get usage info
+        usage_metadata = result.get("usageMetadata", {})
+
+        return LLMResponse(
+            content=content,
+            model=self.config.model,
+            finish_reason=candidates[0].get("finishReason") if candidates else None,
+            usage={
+                "prompt_tokens": usage_metadata.get("promptTokenCount", 0),
+                "completion_tokens": usage_metadata.get("candidatesTokenCount", 0),
+            },
+            raw_response=result,
+        )
+
+    async def complete_streaming(
+        self,
+        prompt: str,
+        system_prompt: Optional[str] = None,
+        **kwargs
+    ) -> AsyncIterator[str]:
+        """
+        Generate streaming completion.
+
+        Args:
+            prompt: User prompt
+            system_prompt: Optional system prompt
+            **kwargs: Additional parameters
+
+        Yields:
+            Text chunks as they are generated
+        """
+        if not self._client:
+            await self.initialize()
+
+        # Build request payload
+        contents = []
+
+        if system_prompt:
+            contents.append({
+                "role": "user",
+                "parts": [{"text": f"System instruction: {system_prompt}"}]
+            })
+            contents.append({
+                "role": "model",
+                "parts": [{"text": "Understood. I will follow those instructions."}]
+            })
+
+        contents.append({
+            "role": "user",
+            "parts": [{"text": prompt}]
+        })
+
+        data = {
+            "contents": contents,
+            "generationConfig": {
+                "temperature": kwargs.get("temperature", self.config.temperature),
+                "topP": kwargs.get("top_p", self.config.top_p),
+            }
+        }
+
+        if self.config.max_tokens or kwargs.get("max_tokens"):
+            data["generationConfig"]["maxOutputTokens"] = kwargs.get("max_tokens", self.config.max_tokens)
+
+        endpoint = self._build_endpoint("streamGenerateContent")
+        endpoint += "&alt=sse"  # Enable SSE streaming
+
+        async with self._client.stream("POST", endpoint, json=data) as response:
+            response.raise_for_status()
+
+            async for line in response.aiter_lines():
+                if line.startswith("data: "):
+                    try:
+                        chunk = json.loads(line[6:])
+                        candidates = chunk.get("candidates", [])
+                        if candidates:
+                            parts = candidates[0].get("content", {}).get("parts", [])
+                            for part in parts:
+                                if "text" in part:
+                                    yield part["text"]
+                    except json.JSONDecodeError:
+                        continue
+
+    async def get_embeddings(
+        self,
+        texts: List[str],
+        **kwargs
+    ) -> List[List[float]]:
+        """
+        Get embeddings using Gemini.
+
+        Args:
+            texts: List of texts to embed
+            **kwargs: Additional parameters
+
+        Returns:
+            List of embedding vectors
+        """
+        if not self._client:
+            await self.initialize()
+
+        embeddings = []
+        embed_model = kwargs.get("embed_model", "text-embedding-004")
+
+        for text in texts:
+            data = {
+                "content": {
+                    "parts": [{"text": text}]
+                }
+            }
+
+            endpoint = f"/v1beta/models/{embed_model}:embedContent?key={self.config.api_key}"
+            response = await self._client.post(endpoint, json=data)
+            response.raise_for_status()
+
+            result = response.json()
+            embeddings.append(result.get("embedding", {}).get("values", []))
+
+        return embeddings

ebk/ai/llm_providers/ollama.py

@@ -0,0 +1,294 @@
+"""
+Ollama LLM Provider.
+
+Supports both local and remote Ollama instances.
+"""
+
+import json
+from typing import Dict, Any, List, Optional, AsyncIterator
+import httpx
+
+from .base import BaseLLMProvider, LLMConfig, LLMResponse, ModelCapability
+
+
+class OllamaProvider(BaseLLMProvider):
+    """
+    Ollama LLM provider.
+
+    Supports:
+    - Local Ollama (default: http://localhost:11434)
+    - Remote Ollama (e.g., basement GPU server)
+    - Streaming completions
+    - JSON mode
+    - Embeddings
+    """
+
+    def __init__(self, config: LLMConfig):
+        """
+        Initialize Ollama provider.
+
+        Args:
+            config: LLM configuration with base_url pointing to Ollama
+        """
+        super().__init__(config)
+
+    @property
+    def name(self) -> str:
+        return "ollama"
+
+    @property
+    def supported_capabilities(self) -> List[ModelCapability]:
+        return [
+            ModelCapability.TEXT_GENERATION,
+            ModelCapability.JSON_MODE,
+            ModelCapability.STREAMING,
+            ModelCapability.EMBEDDINGS,
+        ]
+
+    @classmethod
+    def local(cls, model: str = "llama3.2", **kwargs) -> 'OllamaProvider':
+        """
+        Create provider for local Ollama instance.
+
+        Args:
+            model: Model name (e.g., 'llama3.2', 'mistral', 'codellama')
+            **kwargs: Additional config parameters
+
+        Returns:
+            Configured OllamaProvider
+        """
+        config = LLMConfig(
+            base_url="http://localhost:11434",
+            model=model,
+            **kwargs
+        )
+        return cls(config)
+
+    @classmethod
+    def remote(
+        cls,
+        host: str,
+        port: int = 11434,
+        model: str = "llama3.2",
+        **kwargs
+    ) -> 'OllamaProvider':
+        """
+        Create provider for remote Ollama instance.
+
+        Args:
+            host: Remote host (e.g., '192.168.1.100', 'basement-gpu.local')
+            port: Ollama port (default: 11434)
+            model: Model name
+            **kwargs: Additional config parameters
+
+        Returns:
+            Configured OllamaProvider
+
+        Example:
+            >>> provider = OllamaProvider.remote(
+            ...     host='192.168.1.100',
+            ...     model='llama3.2'
+            ... )
+        """
+        config = LLMConfig(
+            base_url=f"http://{host}:{port}",
+            model=model,
+            **kwargs
+        )
+        return cls(config)
+
+    async def initialize(self) -> None:
+        """Initialize HTTP client."""
+        self._client = httpx.AsyncClient(
+            base_url=self.config.base_url,
+            timeout=self.config.timeout,
+        )
+
+    async def cleanup(self) -> None:
+        """Close HTTP client."""
+        if self._client:
+            await self._client.aclose()
+            self._client = None
+
+    async def complete(
+        self,
+        prompt: str,
+        system_prompt: Optional[str] = None,
+        **kwargs
+    ) -> LLMResponse:
+        """
+        Generate completion using Ollama.
+
+        Args:
+            prompt: User prompt
+            system_prompt: Optional system prompt
+            **kwargs: Additional parameters (temperature, max_tokens, etc.)
+
+        Returns:
+            LLMResponse with generated text
+        """
+        if not self._client:
+            await self.initialize()
+
+        # Build request payload
+        data = {
+            "model": self.config.model,
+            "prompt": prompt,
+            "stream": False,
+            "options": {
+                "temperature": kwargs.get("temperature", self.config.temperature),
+                "top_p": kwargs.get("top_p", self.config.top_p),
+            }
+        }
+
+        if system_prompt:
+            data["system"] = system_prompt
+
+        if self.config.max_tokens:
+            data["options"]["num_predict"] = self.config.max_tokens
+
+        # Make request
+        response = await self._client.post("/api/generate", json=data)
+        response.raise_for_status()
+
+        result = response.json()
+
+        return LLMResponse(
+            content=result["response"],
+            model=result.get("model", self.config.model),
+            finish_reason=result.get("done_reason"),
+            usage={
+                "prompt_tokens": result.get("prompt_eval_count", 0),
+                "completion_tokens": result.get("eval_count", 0),
+            },
+            raw_response=result,
+        )
+
+    async def complete_json(
+        self,
+        prompt: str,
+        system_prompt: Optional[str] = None,
+        schema: Optional[Dict[str, Any]] = None,
+        **kwargs
+    ) -> Dict[str, Any]:
+        """
+        Generate JSON completion using Ollama's JSON format mode.
+
+        Overrides base to set Ollama's format="json" parameter
+        before delegating to the shared prompt-building logic.
+        """
+        # Use Ollama's JSON format mode if available
+        kwargs["format"] = "json"
+        return await super().complete_json(prompt, system_prompt=system_prompt, schema=schema, **kwargs)
+
+    async def complete_streaming(
+        self,
+        prompt: str,
+        system_prompt: Optional[str] = None,
+        **kwargs
+    ) -> AsyncIterator[str]:
+        """
+        Generate streaming completion.
+
+        Args:
+            prompt: User prompt
+            system_prompt: Optional system prompt
+            **kwargs: Additional parameters
+
+        Yields:
+            Text chunks as they are generated
+        """
+        if not self._client:
+            await self.initialize()
+
+        data = {
+            "model": self.config.model,
+            "prompt": prompt,
+            "stream": True,
+            "options": {
+                "temperature": kwargs.get("temperature", self.config.temperature),
+                "top_p": kwargs.get("top_p", self.config.top_p),
+            }
+        }
+
+        if system_prompt:
+            data["system"] = system_prompt
+
+        async with self._client.stream("POST", "/api/generate", json=data) as response:
+            response.raise_for_status()
+
+            async for line in response.aiter_lines():
+                if line.strip():
+                    try:
+                        chunk = json.loads(line)
+                        if "response" in chunk:
+                            yield chunk["response"]
+                        if chunk.get("done", False):
+                            break
+                    except json.JSONDecodeError:
+                        continue
+
+    async def get_embeddings(
+        self,
+        texts: List[str],
+        **kwargs
+    ) -> List[List[float]]:
+        """
+        Get embeddings using Ollama.
+
+        Args:
+            texts: List of texts to embed
+            **kwargs: Additional parameters
+
+        Returns:
+            List of embedding vectors
+        """
+        if not self._client:
+            await self.initialize()
+
+        embeddings = []
+
+        for text in texts:
+            data = {
+                "model": self.config.model,
+                "prompt": text,
+            }
+
+            response = await self._client.post("/api/embeddings", json=data)
+            response.raise_for_status()
+
+            result = response.json()
+            embeddings.append(result["embedding"])
+
+        return embeddings
+
+    async def list_models(self) -> List[str]:
+        """
+        List available models on Ollama server.
+
+        Returns:
+            List of model names
+        """
+        if not self._client:
+            await self.initialize()
+
+        response = await self._client.get("/api/tags")
+        response.raise_for_status()
+
+        result = response.json()
+        return [model["name"] for model in result.get("models", [])]
+
+    async def pull_model(self, model_name: str) -> None:
+        """
+        Pull a model from Ollama registry.
+
+        Args:
+            model_name: Name of model to pull (e.g., 'llama3.2', 'mistral')
+        """
+        if not self._client:
+            await self.initialize()
+
+        data = {"name": model_name, "stream": False}
+
+        response = await self._client.post("/api/pull", json=data)
+        response.raise_for_status()