pytrilogy 0.3.138__cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl

trilogy/ai/providers/anthropic.py

@@ -0,0 +1,106 @@
+from os import environ
+from typing import List, Optional
+
+from trilogy.ai.enums import Provider
+from trilogy.ai.models import LLMMessage, LLMResponse, UsageDict
+from trilogy.constants import logger
+
+from .base import RETRYABLE_CODES, LLMProvider, LLMRequestOptions
+from .utils import RetryOptions, fetch_with_retry
+
+DEFAULT_MAX_TOKENS = 10000
+
+
+class AnthropicProvider(LLMProvider):
+    def __init__(
+        self,
+        name: str,
+        model: str,
+        api_key: str | None = None,
+        retry_options: Optional[RetryOptions] = None,
+    ):
+        api_key = api_key or environ.get("ANTHROPIC_API_KEY")
+        if not api_key:
+            raise ValueError(
+                "API key argument or environment variable ANTHROPIC_API_KEY is required"
+            )
+        super().__init__(name, api_key, model, Provider.ANTHROPIC)
+        self.base_completion_url = "https://api.anthropic.com/v1/messages"
+        self.base_model_url = "https://api.anthropic.com/v1/models"
+        self.models: List[str] = []
+        self.type = Provider.ANTHROPIC
+        self.retry_options = retry_options or RetryOptions(
+            max_retries=5,
+            initial_delay_ms=5000,
+            retry_status_codes=RETRYABLE_CODES,
+            on_retry=lambda attempt, delay_ms, error: logger.info(
+                f"Anthropic API retry attempt {attempt} after {delay_ms}ms delay due to error: {str(error)}"
+            ),
+        )
+
+    def generate_completion(
+        self, options: LLMRequestOptions, history: List[LLMMessage]
+    ) -> LLMResponse:
+        try:
+            import httpx
+        except ImportError:
+            raise ImportError(
+                "Missing httpx. Install pytrilogy[ai] to use AnthropicProvider."
+            )
+
+        # Separate system messages from user/assistant messages
+        system_messages = [msg.content for msg in history if msg.role == "system"]
+        conversation_messages = [
+            {"role": msg.role, "content": msg.content}
+            for msg in history
+            if msg.role != "system"
+        ]
+
+        try:
+
+            def make_request():
+                with httpx.Client(timeout=60) as client:
+                    payload = {
+                        "model": self.model,
+                        "messages": conversation_messages,
+                        "max_tokens": options.max_tokens or DEFAULT_MAX_TOKENS,
+                        # "temperature": options.temperature or 0.7,
+                        # "top_p": options.top_p if hasattr(options, "top_p") else 1.0,
+                    }
+
+                    # Add system parameter if there are system messages
+                    if system_messages:
+                        # Combine multiple system messages with newlines
+                        payload["system"] = "\n\n".join(system_messages)
+
+                    response = client.post(
+                        url=self.base_completion_url,
+                        headers={
+                            "Content-Type": "application/json",
+                            "x-api-key": self.api_key,
+                            "anthropic-version": "2023-06-01",
+                        },
+                        json=payload,
+                    )
+                    response.raise_for_status()
+                    return response.json()
+
+            data = fetch_with_retry(make_request, self.retry_options)
+
+            return LLMResponse(
+                text=data["content"][0]["text"],
+                usage=UsageDict(
+                    prompt_tokens=data["usage"]["input_tokens"],
+                    completion_tokens=data["usage"]["output_tokens"],
+                    total_tokens=data["usage"]["input_tokens"]
+                    + data["usage"]["output_tokens"],
+                ),
+            )
+
+        except httpx.HTTPStatusError as error:
+            error_detail = error.response.text
+            raise Exception(
+                f"Anthropic API error ({error.response.status_code}): {error_detail}"
+            )
+        except Exception as error:
+            raise Exception(f"Anthropic API error: {str(error)}")

trilogy/ai/providers/base.py

@@ -0,0 +1,24 @@
+from abc import ABC, abstractmethod
+from typing import List, Optional
+
+from trilogy.ai.enums import Provider
+from trilogy.ai.models import LLMMessage, LLMRequestOptions, LLMResponse
+
+RETRYABLE_CODES = [429, 500, 502, 503, 504]
+
+
+class LLMProvider(ABC):
+    def __init__(self, name: str, api_key: str, model: str, provider: Provider):
+        self.api_key = api_key
+        self.models: List[str] = []
+        self.name = name
+        self.model = model
+        self.type = provider
+        self.error: Optional[str] = None
+
+    # Abstract method to be implemented by specific providers
+    @abstractmethod
+    def generate_completion(
+        self, options: LLMRequestOptions, history: List[LLMMessage]
+    ) -> LLMResponse:
+        pass

trilogy/ai/providers/google.py

@@ -0,0 +1,146 @@
+from os import environ
+from typing import Any, Dict, List, Optional
+
+from trilogy.ai.enums import Provider
+from trilogy.ai.models import LLMMessage, LLMResponse, UsageDict
+from trilogy.constants import logger
+
+from .base import RETRYABLE_CODES, LLMProvider, LLMRequestOptions
+from .utils import RetryOptions, fetch_with_retry
+
+
+class GoogleProvider(LLMProvider):
+    def __init__(
+        self,
+        name: str,
+        model: str,
+        api_key: str | None = None,
+        retry_options: Optional[RetryOptions] = None,
+    ):
+        api_key = api_key or environ.get("GOOGLE_API_KEY")
+        if not api_key:
+            raise ValueError(
+                "API key argument or environment variable GOOGLE_API_KEY is required"
+            )
+        super().__init__(name, api_key, model, Provider.GOOGLE)
+        self.base_model_url = "https://generativelanguage.googleapis.com/v1/models"
+        self.base_completion_url = "https://generativelanguage.googleapis.com/v1beta"
+        self.models: List[str] = []
+        self.type = Provider.GOOGLE
+        self.retry_options = retry_options or RetryOptions(
+            max_retries=3,
+            initial_delay_ms=30000,
+            retry_status_codes=RETRYABLE_CODES,
+            on_retry=lambda attempt, delay_ms, error: logger.info(
+                f"Google API retry attempt {attempt} after {delay_ms}ms delay due to error: {str(error)}"
+            ),
+        )
+
+    def _convert_to_gemini_history(
+        self, messages: List[LLMMessage]
+    ) -> List[Dict[str, Any]]:
+        """Convert standard message format to Gemini format."""
+        return [
+            {
+                "role": "model" if msg.role == "assistant" else "user",
+                "parts": [{"text": msg.content}],
+            }
+            for msg in messages
+        ]
+
+    def generate_completion(
+        self, options: LLMRequestOptions, history: List[LLMMessage]
+    ) -> LLMResponse:
+        try:
+            import httpx
+        except ImportError:
+            raise ImportError(
+                "Missing httpx. Install pytrilogy[ai] to use GoogleProvider."
+            )
+
+        # Convert messages to Gemini format
+        gemini_history = self._convert_to_gemini_history(history)
+
+        # Separate system message if present
+        system_instruction = None
+        contents = gemini_history
+
+        # Check if first message is a system message
+        if history and history[0].role == "system":
+            system_instruction = {"parts": [{"text": history[0].content}]}
+            contents = gemini_history[1:]  # Remove system message from history
+
+        # Build the request URL
+        url = f"{self.base_completion_url}/models/{self.model}:generateContent"
+
+        # Build request body
+        request_body: Dict[str, Any] = {"contents": contents, "generationConfig": {}}
+
+        # Add system instruction if present
+        if system_instruction:
+            request_body["systemInstruction"] = system_instruction
+
+        # Add generation config options
+        if options.temperature is not None:
+            request_body["generationConfig"]["temperature"] = options.temperature
+
+        if options.max_tokens is not None:
+            request_body["generationConfig"]["maxOutputTokens"] = options.max_tokens
+
+        if options.top_p is not None:
+            request_body["generationConfig"]["topP"] = options.top_p
+
+        try:
+            # Make the API request with retry logic using a lambda
+
+            def fetch_function() -> Dict[str, Any]:
+                with httpx.Client(timeout=60) as client:
+                    resp = client.post(
+                        url,
+                        headers={
+                            "Content-Type": "application/json",
+                            "x-goog-api-key": self.api_key,
+                        },
+                        json=request_body,
+                    )
+                    resp.raise_for_status()
+                    return resp.json()
+
+            data = fetch_with_retry(
+                fetch_fn=fetch_function,
+                options=self.retry_options,
+            )
+
+            # Extract text from response
+            candidates = data.get("candidates", [])
+            if not candidates:
+                raise Exception("No candidates returned from Google API")
+
+            content = candidates[0].get("content", {})
+            parts = content.get("parts", [])
+
+            if not parts:
+                raise Exception("No parts in response content")
+
+            text = parts[0].get("text", "")
+
+            # Extract usage metadata
+            usage_metadata = data.get("usageMetadata", {})
+            prompt_tokens = usage_metadata.get("promptTokenCount", 0)
+            completion_tokens = usage_metadata.get("candidatesTokenCount", 0)
+
+            return LLMResponse(
+                text=text,
+                usage=UsageDict(
+                    prompt_tokens=prompt_tokens,
+                    completion_tokens=completion_tokens,
+                    total_tokens=prompt_tokens + completion_tokens,
+                ),
+            )
+        except httpx.HTTPStatusError as error:
+            error_detail = error.response.text
+            raise Exception(
+                f"Google API error ({error.response.status_code}): {error_detail}"
+            )
+        except Exception as error:
+            raise Exception(f"Google API error: {str(error)}")

trilogy/ai/providers/openai.py

@@ -0,0 +1,89 @@
+from os import environ
+from typing import List, Optional
+
+from trilogy.ai.enums import Provider
+from trilogy.ai.models import LLMMessage, LLMResponse, UsageDict
+from trilogy.constants import logger
+
+from .base import RETRYABLE_CODES, LLMProvider, LLMRequestOptions
+from .utils import RetryOptions, fetch_with_retry
+
+
+class OpenAIProvider(LLMProvider):
+    def __init__(
+        self,
+        name: str,
+        model: str,
+        api_key: str | None = None,
+        retry_options: Optional[RetryOptions] = None,
+    ):
+        api_key = api_key or environ.get("OPENAI_API_KEY")
+        if not api_key:
+            raise ValueError(
+                "API key argument or environment variable OPENAI_API_KEY is required"
+            )
+        super().__init__(name, api_key, model, Provider.OPENAI)
+        self.base_completion_url = "https://api.openai.com/v1/chat/completions"
+        self.base_model_url = "https://api.openai.com/v1/models"
+        self.models: List[str] = []
+        self.type = Provider.OPENAI
+
+        self.retry_options = retry_options or RetryOptions(
+            max_retries=3,
+            initial_delay_ms=1000,
+            retry_status_codes=RETRYABLE_CODES,
+            on_retry=lambda attempt, delay_ms, error: logger.info(
+                f"Retry attempt {attempt} after {delay_ms}ms delay due to error: {str(error)}"
+            ),
+        )
+
+    def generate_completion(
+        self, options: LLMRequestOptions, history: List[LLMMessage]
+    ) -> LLMResponse:
+        try:
+            import httpx
+        except ImportError:
+            raise ImportError(
+                "Missing httpx. Install pytrilogy[ai] to use OpenAIProvider."
+            )
+
+        messages: List[dict] = []
+        messages = [{"role": msg.role, "content": msg.content} for msg in history]
+        try:
+
+            def make_request():
+                with httpx.Client(timeout=30) as client:
+                    payload = {
+                        "model": self.model,
+                        "messages": messages,
+                    }
+
+                    response = client.post(
+                        url=self.base_completion_url,
+                        headers={
+                            "Content-Type": "application/json",
+                            "Authorization": f"Bearer {self.api_key}",
+                        },
+                        json=payload,
+                    )
+                    response.raise_for_status()
+                    return response.json()
+
+            data = fetch_with_retry(make_request, self.retry_options)
+            return LLMResponse(
+                text=data["choices"][0]["message"]["content"],
+                usage=UsageDict(
+                    prompt_tokens=data["usage"]["prompt_tokens"],
+                    completion_tokens=data["usage"]["completion_tokens"],
+                    total_tokens=data["usage"]["total_tokens"],
+                ),
+            )
+        except httpx.HTTPStatusError as error:
+            # Capture the response body text
+            error_detail = error.response.text
+            raise Exception(
+                f"OpenAI API error ({error.response.status_code}): {error_detail}"
+            )
+
+        except Exception as error:
+            raise Exception(f"OpenAI API error: {str(error)}")

trilogy/ai/providers/utils.py

@@ -0,0 +1,68 @@
+import time
+from dataclasses import dataclass, field
+from typing import Callable, List, TypeVar
+
+T = TypeVar("T")
+
+
+@dataclass
+class RetryOptions:
+    max_retries: int = 3
+    initial_delay_ms: int = 1000
+    retry_status_codes: List[int] = field(
+        default_factory=lambda: [429, 500, 502, 503, 504, 525]
+    )
+    on_retry: Callable[[int, int, Exception], None] | None = None
+
+
+def fetch_with_retry(fetch_fn: Callable[[], T], options: RetryOptions) -> T:
+    from httpx import HTTPError
+
+    """
+    Retry a fetch operation with exponential backoff.
+
+    Args:
+        fetch_fn: Function that performs the fetch operation
+        options: Retry configuration options
+
+    Returns:
+        The result from the successful fetch operation
+
+    Raises:
+        The last exception encountered if all retries fail
+    """
+    from httpx import HTTPStatusError, TimeoutException
+
+    last_error = None
+    delay_ms = options.initial_delay_ms
+
+    for attempt in range(options.max_retries + 1):
+        try:
+            return fetch_fn()
+        except (HTTPError, TimeoutException, Exception) as error:
+            last_error = error
+            should_retry = False
+
+            if isinstance(error, HTTPStatusError):
+                if (
+                    options.retry_status_codes
+                    and error.response.status_code in options.retry_status_codes
+                ):
+                    should_retry = True
+            elif isinstance(error, TimeoutException):
+                should_retry = True
+            if not should_retry or attempt >= options.max_retries:
+                raise
+
+            # Call the retry callback if provided
+            if options.on_retry:
+                options.on_retry(attempt + 1, delay_ms, error)
+
+            # Wait before retrying with exponential backoff
+            time.sleep(delay_ms / 1000.0)
+            delay_ms *= 2  # backoff
+
+    # This should never be reached, but just in case
+    if last_error:
+        raise last_error
+    raise Exception("Retry logic failed unexpectedly")

trilogy/authoring/README.md

@@ -0,0 +1,3 @@
+### Authoring Public Interface
+
+For anyone directly building on Trilogy in python, this folder contains the public interface for authoring objects.

trilogy/authoring/__init__.py

@@ -0,0 +1,143 @@
+from trilogy.constants import DEFAULT_NAMESPACE
+from trilogy.core.enums import (
+    BooleanOperator,
+    ComparisonOperator,
+    FunctionClass,
+    FunctionType,
+    InfiniteFunctionArgs,
+    Ordering,
+    Purpose,
+)
+from trilogy.core.functions import FunctionFactory
+from trilogy.core.models.author import (
+    AggregateWrapper,
+    CaseElse,
+    CaseWhen,
+    Comparison,
+    Concept,
+    ConceptRef,
+    Conditional,
+    FilterItem,
+    Function,
+    FunctionCallWrapper,
+    HavingClause,
+    MagicConstants,
+    Metadata,
+    MultiSelectLineage,
+    OrderBy,
+    OrderItem,
+    Parenthetical,
+    RowsetItem,
+    SubselectComparison,
+    WhereClause,
+    WindowItem,
+    WindowItemOrder,
+    WindowItemOver,
+    WindowOrder,
+    WindowType,
+)
+from trilogy.core.models.core import (
+    ArrayType,
+    DataType,
+    ListWrapper,
+    MapType,
+    NumericType,
+    StructType,
+    TraitDataType,
+)
+from trilogy.core.models.datasource import Address, Datasource, DatasourceMetadata
+from trilogy.core.models.environment import Environment
+from trilogy.core.statements.author import (
+    STATEMENT_TYPES,
+    ConceptDeclarationStatement,
+    ConceptTransform,
+    CopyStatement,
+    Grain,
+    HasUUID,
+    ImportStatement,
+    MultiSelectStatement,
+    PersistStatement,
+    RawSQLStatement,
+    RowsetDerivationStatement,
+    SelectItem,
+    SelectStatement,
+    ShowCategory,
+    ShowStatement,
+    ValidateStatement,
+)
+from trilogy.parsing.common import arbitrary_to_concept, arg_to_datatype
+
+__all__ = [
+    # trilogy.constants
+    "DEFAULT_NAMESPACE",
+    # trilogy.core.enums
+    "BooleanOperator",
+    "ComparisonOperator",
+    "FunctionClass",
+    "FunctionType",
+    "InfiniteFunctionArgs",
+    "Ordering",
+    "Purpose",
+    # trilogy.core.functions
+    "FunctionFactory",
+    # trilogy.core.models.author
+    "AggregateWrapper",
+    "CaseElse",
+    "CaseWhen",
+    "Comparison",
+    "Concept",
+    "ConceptRef",
+    "Conditional",
+    "FilterItem",
+    "Function",
+    "FunctionCallWrapper",
+    "HavingClause",
+    "MagicConstants",
+    "Metadata",
+    "MultiSelectLineage",
+    "OrderBy",
+    "OrderItem",
+    "Parenthetical",
+    "RowsetItem",
+    "SubselectComparison",
+    "WhereClause",
+    "WindowItem",
+    "WindowItemOrder",
+    "WindowItemOver",
+    "WindowOrder",
+    "WindowType",
+    # trilogy.core.models.core
+    "ArrayType",
+    "DataType",
+    "ListWrapper",
+    "MapType",
+    "NumericType",
+    "StructType",
+    "TraitDataType",
+    # trilogy.core.models.datasource
+    "Address",
+    "Datasource",
+    "DatasourceMetadata",
+    # trilogy.core.models.environment
+    "Environment",
+    # trilogy.core.statements.author
+    "ConceptDeclarationStatement",
+    "ConceptTransform",
+    "CopyStatement",
+    "Grain",
+    "HasUUID",
+    "ImportStatement",
+    "MultiSelectStatement",
+    "PersistStatement",
+    "RawSQLStatement",
+    "RowsetDerivationStatement",
+    "SelectItem",
+    "SelectStatement",
+    "ShowCategory",
+    "ShowStatement",
+    "ValidateStatement",
+    # trilogy.parsing.common
+    "arbitrary_to_concept",
+    "arg_to_datatype",
+    "STATEMENT_TYPES",
+]

trilogy/constants.py

@@ -0,0 +1,113 @@
+import random
+from contextlib import contextmanager
+from dataclasses import dataclass, field
+from enum import Enum
+from logging import getLogger
+from typing import Any
+
+logger = getLogger("trilogy")
+
+DEFAULT_NAMESPACE = "local"
+
+RECURSIVE_GATING_CONCEPT = "_terminal"
+
+VIRTUAL_CONCEPT_PREFIX = "_virt"
+
+ENV_CACHE_NAME = ".preql_cache.json"
+
+
+class MagicConstants(Enum):
+    NULL = "null"
+    LINE_SEPARATOR = "\n"
+
+
+NULL_VALUE = MagicConstants.NULL
+
+
+@dataclass
+class Optimizations:
+    predicate_pushdown: bool = True
+    datasource_inlining: bool = True
+    constant_inlining: bool = True
+    constant_inline_cutoff: int = 10
+    direct_return: bool = True
+    hide_unused_concepts: bool = True
+
+
+@dataclass
+class Comments:
+    """Control what is placed in CTE comments"""
+
+    show: bool = False
+    basic: bool = True
+    joins: bool = False
+    nullable: bool = False
+    partial: bool = False
+    source_map: bool = False
+
+
+@dataclass
+class Rendering:
+    """Control how the SQL is rendered"""
+
+    parameters: bool = True
+    concise: bool = False
+
+    @contextmanager
+    def temporary(self, **kwargs: Any):
+        """
+        Context manager to temporarily set attributes and revert them afterwards.
+
+        Usage:
+            r = Rendering()
+            with r.temporary(parameters=False, concise=True):
+                # parameters is False, concise is True here
+                do_something()
+            # parameters and concise are back to their original values
+        """
+        # Store original values
+        original_values = {key: getattr(self, key) for key in kwargs}
+
+        # Set new values
+        for key, value in kwargs.items():
+            setattr(self, key, value)
+
+        try:
+            yield self
+        finally:
+            # Restore original values
+            for key, value in original_values.items():
+                setattr(self, key, value)
+
+
+@dataclass
+class Parsing:
+    """Control Parsing"""
+
+    strict_name_shadow_enforcement: bool = False
+    select_as_definition: bool = True
+
+
+# TODO: support loading from environments
+@dataclass
+class Config:
+    strict_mode: bool = True
+    human_identifiers: bool = True
+    randomize_cte_names: bool = False
+    validate_missing: bool = True
+    comments: Comments = field(default_factory=Comments)
+    optimizations: Optimizations = field(default_factory=Optimizations)
+    rendering: Rendering = field(default_factory=Rendering)
+    parsing: Parsing = field(default_factory=Parsing)
+
+    @property
+    def show_comments(self) -> bool:
+        return self.comments.show
+
+    def set_random_seed(self, seed: int):
+        random.seed(seed)
+
+
+CONFIG = Config()
+
+CONFIG.set_random_seed(42)