causaliq-knowledge 0.1.0__py3-none-any.whl → 0.3.0__py3-none-any.whl

causaliq_knowledge/llm/groq_client.py

@@ -1,6 +1,5 @@
 """Direct Groq API client - clean and reliable."""
 
-import json
 import logging
 import os
 from dataclasses import dataclass
@@ -8,12 +7,28 @@ from typing import Any, Dict, List, Optional
 
 import httpx
 
+from causaliq_knowledge.llm.base_client import (
+    BaseLLMClient,
+    LLMConfig,
+    LLMResponse,
+)
+
 logger = logging.getLogger(__name__)
 
 
 @dataclass
-class GroqConfig:
-    """Configuration for Groq API client."""
+class GroqConfig(LLMConfig):
+    """Configuration for Groq API client.
+
+    Extends LLMConfig with Groq-specific defaults.
+
+    Attributes:
+        model: Groq model identifier (default: llama-3.1-8b-instant).
+        temperature: Sampling temperature (default: 0.1).
+        max_tokens: Maximum response tokens (default: 500).
+        timeout: Request timeout in seconds (default: 30.0).
+        api_key: Groq API key (falls back to GROQ_API_KEY env var).
+    """
 
     model: str = "llama-3.1-8b-instant"
     temperature: float = 0.1
@@ -29,50 +44,52 @@ class GroqConfig:
             raise ValueError("GROQ_API_KEY environment variable is required")
 
 
-@dataclass
-class GroqResponse:
-    """Response from Groq API."""
-
-    content: str
-    model: str
-    input_tokens: int = 0
-    output_tokens: int = 0
-    cost: float = 0.0  # Groq free tier
-    raw_response: Optional[Dict] = None
-
-    def parse_json(self) -> Optional[Dict[str, Any]]:
-        """Parse content as JSON, handling common formatting issues."""
-        try:
-            # Clean up potential markdown code blocks
-            text = self.content.strip()
-            if text.startswith("```json"):
-                text = text[7:]
-            elif text.startswith("```"):
-                text = text[3:]
-            if text.endswith("```"):
-                text = text[:-3]
-
-            return json.loads(text.strip())  # type: ignore[no-any-return]
-        except json.JSONDecodeError as e:
-            logger.warning(f"Failed to parse JSON response: {e}")
-            return None
+class GroqClient(BaseLLMClient):
+    """Direct Groq API client.
 
+    Implements the BaseLLMClient interface for Groq's API.
+    Uses httpx for HTTP requests.
 
-class GroqClient:
-    """Direct Groq API client."""
+    Example:
+        >>> config = GroqConfig(model="llama-3.1-8b-instant")
+        >>> client = GroqClient(config)
+        >>> msgs = [{"role": "user", "content": "Hello"}]
+        >>> response = client.completion(msgs)
+        >>> print(response.content)
+    """
 
     BASE_URL = "https://api.groq.com/openai/v1"
 
-    def __init__(self, config: Optional[GroqConfig] = None):
-        """Initialize Groq client."""
+    def __init__(self, config: Optional[GroqConfig] = None) -> None:
+        """Initialize Groq client.
+
+        Args:
+            config: Groq configuration. If None, uses defaults with
+                   API key from GROQ_API_KEY environment variable.
+        """
         self.config = config or GroqConfig()
         self._total_calls = 0
 
+    @property
+    def provider_name(self) -> str:
+        """Return the provider name."""
+        return "groq"
+
     def completion(
         self, messages: List[Dict[str, str]], **kwargs: Any
-    ) -> GroqResponse:
-        """Make a chat completion request to Groq."""
+    ) -> LLMResponse:
+        """Make a chat completion request to Groq.
+
+        Args:
+            messages: List of message dicts with "role" and "content" keys.
+            **kwargs: Override config options (temperature, max_tokens).
 
+        Returns:
+            LLMResponse with the generated content and metadata.
+
+        Raises:
+            ValueError: If the API request fails.
+        """
         # Build request payload
         payload = {
             "model": self.config.model,
@@ -112,7 +129,7 @@ class GroqClient:
                     f"Groq response: {input_tokens} in, {output_tokens} out"
                 )
 
-                return GroqResponse(
+                return LLMResponse(
                     content=content,
                     model=data.get("model", self.config.model),
                     input_tokens=input_tokens,
@@ -136,13 +153,71 @@ class GroqClient:
 
     def complete_json(
         self, messages: List[Dict[str, str]], **kwargs: Any
-    ) -> tuple[Optional[Dict[str, Any]], GroqResponse]:
-        """Make a completion request and parse response as JSON."""
+    ) -> tuple[Optional[Dict[str, Any]], LLMResponse]:
+        """Make a completion request and parse response as JSON.
+
+        Args:
+            messages: List of message dicts with "role" and "content" keys.
+            **kwargs: Override config options passed to completion().
+
+        Returns:
+            Tuple of (parsed JSON dict or None, raw LLMResponse).
+        """
         response = self.completion(messages, **kwargs)
         parsed = response.parse_json()
         return parsed, response
 
     @property
     def call_count(self) -> int:
-        """Number of API calls made."""
+        """Return the number of API calls made."""
         return self._total_calls
+
+    def is_available(self) -> bool:
+        """Check if Groq API is available.
+
+        Returns:
+            True if GROQ_API_KEY is configured.
+        """
+        return bool(self.config.api_key)
+
+    def list_models(self) -> List[str]:
+        """List available models from Groq API.
+
+        Queries the Groq API to get models accessible with the current
+        API key. Filters to only include text generation models.
+
+        Returns:
+            List of model identifiers (e.g., ['llama-3.1-8b-instant', ...]).
+
+        Raises:
+            ValueError: If the API request fails.
+        """
+        try:
+            with httpx.Client(timeout=self.config.timeout) as client:
+                response = client.get(
+                    f"{self.BASE_URL}/models",
+                    headers={"Authorization": f"Bearer {self.config.api_key}"},
+                )
+                response.raise_for_status()
+                data = response.json()
+
+                # Filter and sort models
+                models = []
+                for model in data.get("data", []):
+                    model_id = model.get("id", "")
+                    # Skip whisper (audio), guard, and safeguard models
+                    if any(
+                        x in model_id.lower()
+                        for x in ["whisper", "guard", "embed"]
+                    ):
+                        continue
+                    models.append(model_id)
+
+                return sorted(models)
+
+        except httpx.HTTPStatusError as e:
+            raise ValueError(
+                f"Groq API error: {e.response.status_code} - {e.response.text}"
+            )
+        except Exception as e:
+            raise ValueError(f"Failed to list Groq models: {e}")

causaliq_knowledge/llm/mistral_client.py

@@ -0,0 +1,122 @@
+"""Direct Mistral AI API client - OpenAI-compatible API."""
+
+import logging
+import os
+from dataclasses import dataclass
+from typing import Dict, List, Optional
+
+from causaliq_knowledge.llm.openai_compat_client import (
+    OpenAICompatClient,
+    OpenAICompatConfig,
+)
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class MistralConfig(OpenAICompatConfig):
+    """Configuration for Mistral AI API client.
+
+    Extends OpenAICompatConfig with Mistral-specific defaults.
+
+    Attributes:
+        model: Mistral model identifier (default: mistral-small-latest).
+        temperature: Sampling temperature (default: 0.1).
+        max_tokens: Maximum response tokens (default: 500).
+        timeout: Request timeout in seconds (default: 30.0).
+        api_key: Mistral API key (falls back to MISTRAL_API_KEY env var).
+    """
+
+    model: str = "mistral-small-latest"
+    temperature: float = 0.1
+    max_tokens: int = 500
+    timeout: float = 30.0
+    api_key: Optional[str] = None
+
+    def __post_init__(self) -> None:
+        """Set API key from environment if not provided."""
+        if self.api_key is None:
+            self.api_key = os.getenv("MISTRAL_API_KEY")
+        if not self.api_key:
+            raise ValueError(
+                "MISTRAL_API_KEY environment variable is required"
+            )
+
+
+class MistralClient(OpenAICompatClient):
+    """Direct Mistral AI API client.
+
+    Mistral AI is a French company providing high-quality LLMs with an
+    OpenAI-compatible API.
+
+    Available models:
+        - mistral-small-latest: Fast, cost-effective
+        - mistral-medium-latest: Balanced performance
+        - mistral-large-latest: Most capable
+        - codestral-latest: Optimized for code
+
+    Example:
+        >>> config = MistralConfig(model="mistral-small-latest")
+        >>> client = MistralClient(config)
+        >>> msgs = [{"role": "user", "content": "Hello"}]
+        >>> response = client.completion(msgs)
+        >>> print(response.content)
+    """
+
+    BASE_URL = "https://api.mistral.ai/v1"
+    PROVIDER_NAME = "mistral"
+    ENV_VAR = "MISTRAL_API_KEY"
+
+    def __init__(self, config: Optional[MistralConfig] = None) -> None:
+        """Initialize Mistral client.
+
+        Args:
+            config: Mistral configuration. If None, uses defaults with
+                   API key from MISTRAL_API_KEY environment variable.
+        """
+        super().__init__(config)
+
+    def _default_config(self) -> MistralConfig:
+        """Return default Mistral configuration."""
+        return MistralConfig()
+
+    def _get_pricing(self) -> Dict[str, Dict[str, float]]:
+        """Return Mistral pricing per 1M tokens.
+
+        Returns:
+            Dict mapping model prefixes to input/output costs.
+        """
+        # Mistral pricing as of Jan 2025
+        return {
+            "mistral-large": {"input": 2.00, "output": 6.00},
+            "mistral-medium": {"input": 2.70, "output": 8.10},
+            "mistral-small": {"input": 0.20, "output": 0.60},
+            "codestral": {"input": 0.20, "output": 0.60},
+            "open-mistral-nemo": {"input": 0.15, "output": 0.15},
+            "open-mixtral-8x22b": {"input": 2.00, "output": 6.00},
+            "open-mixtral-8x7b": {"input": 0.70, "output": 0.70},
+            "ministral-3b": {"input": 0.04, "output": 0.04},
+            "ministral-8b": {"input": 0.10, "output": 0.10},
+        }
+
+    def _filter_models(self, models: List[str]) -> List[str]:
+        """Filter to Mistral chat models only.
+
+        Args:
+            models: List of all model IDs from API.
+
+        Returns:
+            Filtered list of Mistral models.
+        """
+        filtered = []
+        for model_id in models:
+            # Include mistral and codestral models
+            if any(
+                prefix in model_id.lower()
+                for prefix in ["mistral", "codestral", "ministral", "mixtral"]
+            ):
+                # Exclude embedding models
+                if "embed" in model_id.lower():
+                    continue
+                filtered.append(model_id)
+        return filtered

causaliq_knowledge/llm/ollama_client.py

@@ -0,0 +1,240 @@
+"""Local Ollama API client for running Llama models locally."""
+
+import logging
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional
+
+import httpx
+
+from causaliq_knowledge.llm.base_client import (
+    BaseLLMClient,
+    LLMConfig,
+    LLMResponse,
+)
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class OllamaConfig(LLMConfig):
+    """Configuration for Ollama API client.
+
+    Extends LLMConfig with Ollama-specific defaults.
+
+    Attributes:
+        model: Ollama model identifier (default: llama3.2:1b).
+        temperature: Sampling temperature (default: 0.1).
+        max_tokens: Maximum response tokens (default: 500).
+        timeout: Request timeout in seconds (default: 120.0, local).
+        api_key: Not used for Ollama (local server).
+        base_url: Ollama server URL (default: http://localhost:11434).
+    """
+
+    model: str = "llama3.2:1b"
+    temperature: float = 0.1
+    max_tokens: int = 500
+    timeout: float = 120.0  # Local inference can be slow
+    api_key: Optional[str] = None  # Not needed for local Ollama
+    base_url: str = "http://localhost:11434"
+
+
+class OllamaClient(BaseLLMClient):
+    """Local Ollama API client.
+
+    Implements the BaseLLMClient interface for locally running Ollama server.
+    Uses httpx for HTTP requests to the local Ollama API.
+
+    Ollama provides an OpenAI-compatible API for running open-source models
+    like Llama locally without requiring API keys or internet access.
+
+    Example:
+        >>> config = OllamaConfig(model="llama3.2:1b")
+        >>> client = OllamaClient(config)
+        >>> msgs = [{"role": "user", "content": "Hello"}]
+        >>> response = client.completion(msgs)
+        >>> print(response.content)
+    """
+
+    def __init__(self, config: Optional[OllamaConfig] = None) -> None:
+        """Initialize Ollama client.
+
+        Args:
+            config: Ollama configuration. If None, uses defaults connecting
+                   to localhost:11434 with llama3.2:1b model.
+        """
+        self.config = config or OllamaConfig()
+        self._total_calls = 0
+
+    @property
+    def provider_name(self) -> str:
+        """Return the provider name."""
+        return "ollama"
+
+    def completion(
+        self, messages: List[Dict[str, str]], **kwargs: Any
+    ) -> LLMResponse:
+        """Make a chat completion request to Ollama.
+
+        Args:
+            messages: List of message dicts with "role" and "content" keys.
+            **kwargs: Override config options (temperature, max_tokens).
+
+        Returns:
+            LLMResponse with the generated content and metadata.
+
+        Raises:
+            ValueError: If the API request fails or Ollama is not running.
+        """
+        # Build request payload (Ollama uses similar format to OpenAI)
+        payload: Dict[str, Any] = {
+            "model": self.config.model,
+            "messages": messages,
+            "stream": False,
+            "options": {
+                "temperature": kwargs.get(
+                    "temperature", self.config.temperature
+                ),
+                "num_predict": kwargs.get(
+                    "max_tokens", self.config.max_tokens
+                ),
+            },
+        }
+
+        url = f"{self.config.base_url}/api/chat"
+        headers = {"Content-Type": "application/json"}
+
+        logger.debug(f"Calling Ollama API with model: {self.config.model}")
+
+        try:
+            with httpx.Client(timeout=self.config.timeout) as client:
+                response = client.post(url, json=payload, headers=headers)
+                response.raise_for_status()
+
+                data = response.json()
+
+                # Extract response content
+                content = data.get("message", {}).get("content", "")
+
+                # Extract token counts (Ollama provides these)
+                input_tokens = data.get("prompt_eval_count", 0)
+                output_tokens = data.get("eval_count", 0)
+
+                self._total_calls += 1
+
+                logger.debug(
+                    f"Ollama response: {input_tokens} in, {output_tokens} out"
+                )
+
+                return LLMResponse(
+                    content=content,
+                    model=data.get("model", self.config.model),
+                    input_tokens=input_tokens,
+                    output_tokens=output_tokens,
+                    cost=0.0,  # Local inference is free
+                    raw_response=data,
+                )
+
+        except httpx.ConnectError:
+            raise ValueError(
+                "Could not connect to Ollama. "
+                "Make sure Ollama is running (run 'ollama serve' or start "
+                "the Ollama app). "
+                f"Tried to connect to: {self.config.base_url}"
+            )
+        except httpx.HTTPStatusError as e:
+            if e.response.status_code == 404:
+                raise ValueError(
+                    f"Model '{self.config.model}' not found. "
+                    f"Run 'ollama pull {self.config.model}' to download it."
+                )
+            logger.error(
+                f"Ollama API error: {e.response.status_code} - "
+                f"{e.response.text}"
+            )
+            raise ValueError(
+                f"Ollama API error: {e.response.status_code} - "
+                f"{e.response.text}"
+            )
+        except httpx.TimeoutException:
+            raise ValueError(
+                "Ollama API request timed out. Local inference can be slow - "
+                "try increasing the timeout in OllamaConfig."
+            )
+        except Exception as e:
+            logger.error(f"Ollama API unexpected error: {e}")
+            raise ValueError(f"Ollama API error: {str(e)}")
+
+    def complete_json(
+        self, messages: List[Dict[str, str]], **kwargs: Any
+    ) -> tuple[Optional[Dict[str, Any]], LLMResponse]:
+        """Make a completion request and parse response as JSON.
+
+        Args:
+            messages: List of message dicts with "role" and "content" keys.
+            **kwargs: Override config options passed to completion().
+
+        Returns:
+            Tuple of (parsed JSON dict or None, raw LLMResponse).
+        """
+        response = self.completion(messages, **kwargs)
+        parsed = response.parse_json()
+        return parsed, response
+
+    @property
+    def call_count(self) -> int:
+        """Return the number of API calls made."""
+        return self._total_calls
+
+    def is_available(self) -> bool:
+        """Check if Ollama server is running and model is available.
+
+        Returns:
+            True if Ollama is running and the configured model exists.
+        """
+        try:
+            with httpx.Client(timeout=5.0) as client:
+                # Check if server is running
+                response = client.get(f"{self.config.base_url}/api/tags")
+                if response.status_code != 200:
+                    return False
+
+                # Check if model is available
+                data = response.json()
+                models = [m.get("name", "") for m in data.get("models", [])]
+                # Ollama model names can have :latest suffix
+                model_name = self.config.model
+                return any(
+                    m == model_name or m.startswith(f"{model_name}:")
+                    for m in models
+                )
+        except Exception:
+            return False
+
+    def list_models(self) -> List[str]:
+        """List installed models from Ollama.
+
+        Queries the local Ollama server to get installed models.
+        Unlike cloud providers, this returns only models the user
+        has explicitly pulled/installed.
+
+        Returns:
+            List of model identifiers (e.g., ['llama3.2:1b', ...]).
+
+        Raises:
+            ValueError: If Ollama server is not running.
+        """
+        try:
+            with httpx.Client(timeout=5.0) as client:
+                response = client.get(f"{self.config.base_url}/api/tags")
+                response.raise_for_status()
+                data = response.json()
+
+                models = [m.get("name", "") for m in data.get("models", [])]
+                return sorted(models)
+
+        except httpx.ConnectError:
+            raise ValueError(
+                "Ollama server not running. Start with: ollama serve"
+            )
+        except Exception as e:
+            raise ValueError(f"Failed to list Ollama models: {e}")

causaliq_knowledge/llm/openai_client.py

@@ -0,0 +1,115 @@
+"""Direct OpenAI API client - clean and reliable."""
+
+import logging
+import os
+from dataclasses import dataclass
+from typing import Dict, List, Optional
+
+from causaliq_knowledge.llm.openai_compat_client import (
+    OpenAICompatClient,
+    OpenAICompatConfig,
+)
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class OpenAIConfig(OpenAICompatConfig):
+    """Configuration for OpenAI API client.
+
+    Extends OpenAICompatConfig with OpenAI-specific defaults.
+
+    Attributes:
+        model: OpenAI model identifier (default: gpt-4o-mini).
+        temperature: Sampling temperature (default: 0.1).
+        max_tokens: Maximum response tokens (default: 500).
+        timeout: Request timeout in seconds (default: 30.0).
+        api_key: OpenAI API key (falls back to OPENAI_API_KEY env var).
+    """
+
+    model: str = "gpt-4o-mini"
+    temperature: float = 0.1
+    max_tokens: int = 500
+    timeout: float = 30.0
+    api_key: Optional[str] = None
+
+    def __post_init__(self) -> None:
+        """Set API key from environment if not provided."""
+        if self.api_key is None:
+            self.api_key = os.getenv("OPENAI_API_KEY")
+        if not self.api_key:
+            raise ValueError("OPENAI_API_KEY environment variable is required")
+
+
+class OpenAIClient(OpenAICompatClient):
+    """Direct OpenAI API client.
+
+    Implements the BaseLLMClient interface for OpenAI's API.
+    Uses httpx for HTTP requests.
+
+    Example:
+        >>> config = OpenAIConfig(model="gpt-4o-mini")
+        >>> client = OpenAIClient(config)
+        >>> msgs = [{"role": "user", "content": "Hello"}]
+        >>> response = client.completion(msgs)
+        >>> print(response.content)
+    """
+
+    BASE_URL = "https://api.openai.com/v1"
+    PROVIDER_NAME = "openai"
+    ENV_VAR = "OPENAI_API_KEY"
+
+    def __init__(self, config: Optional[OpenAIConfig] = None) -> None:
+        """Initialize OpenAI client.
+
+        Args:
+            config: OpenAI configuration. If None, uses defaults with
+                   API key from OPENAI_API_KEY environment variable.
+        """
+        super().__init__(config)
+
+    def _default_config(self) -> OpenAIConfig:
+        """Return default OpenAI configuration."""
+        return OpenAIConfig()
+
+    def _get_pricing(self) -> Dict[str, Dict[str, float]]:
+        """Return OpenAI pricing per 1M tokens.
+
+        Returns:
+            Dict mapping model prefixes to input/output costs.
+        """
+        # Order matters - more specific prefixes must come first
+        return {
+            "gpt-4o-mini": {"input": 0.15, "output": 0.60},
+            "gpt-4o": {"input": 2.50, "output": 10.00},
+            "gpt-4-turbo": {"input": 10.00, "output": 30.00},
+            "gpt-4": {"input": 30.00, "output": 60.00},
+            "gpt-3.5-turbo": {"input": 0.50, "output": 1.50},
+            "o1-mini": {"input": 3.00, "output": 12.00},
+            "o1": {"input": 15.00, "output": 60.00},
+        }
+
+    def _filter_models(self, models: List[str]) -> List[str]:
+        """Filter to OpenAI chat models only.
+
+        Args:
+            models: List of all model IDs from API.
+
+        Returns:
+            Filtered list of GPT and o1/o3 models.
+        """
+        filtered = []
+        for model_id in models:
+            # Include GPT and o1/o3 models
+            if any(
+                prefix in model_id
+                for prefix in ["gpt-4", "gpt-3.5", "o1", "o3"]
+            ):
+                # Exclude instruct variants and specific exclusions
+                if any(
+                    x in model_id.lower()
+                    for x in ["instruct", "vision", "audio", "realtime"]
+                ):
+                    continue
+                filtered.append(model_id)
+        return filtered