groknroll 2.0.0__py3-none-any.whl

groknroll/clients/anthropic.py

@@ -0,0 +1,112 @@
+from collections import defaultdict
+from typing import Any
+
+import anthropic
+
+from groknroll.clients.base_lm import BaseLM
+from groknroll.core.types import ModelUsageSummary, UsageSummary
+
+
+class AnthropicClient(BaseLM):
+    """
+    LM Client for running models with the Anthropic API.
+    """
+
+    def __init__(
+        self,
+        api_key: str,
+        model_name: str | None = None,
+        max_tokens: int = 32768,
+        **kwargs,
+    ):
+        super().__init__(model_name=model_name, **kwargs)
+        self.client = anthropic.Anthropic(api_key=api_key)
+        self.async_client = anthropic.AsyncAnthropic(api_key=api_key)
+        self.model_name = model_name
+        self.max_tokens = max_tokens
+
+        # Per-model usage tracking
+        self.model_call_counts: dict[str, int] = defaultdict(int)
+        self.model_input_tokens: dict[str, int] = defaultdict(int)
+        self.model_output_tokens: dict[str, int] = defaultdict(int)
+        self.model_total_tokens: dict[str, int] = defaultdict(int)
+
+    def completion(self, prompt: str | list[dict[str, Any]], model: str | None = None) -> str:
+        messages, system = self._prepare_messages(prompt)
+
+        model = model or self.model_name
+        if not model:
+            raise ValueError("Model name is required for Anthropic client.")
+
+        kwargs = {"model": model, "max_tokens": self.max_tokens, "messages": messages}
+        if system:
+            kwargs["system"] = system
+
+        response = self.client.messages.create(**kwargs)
+        self._track_cost(response, model)
+        return response.content[0].text
+
+    async def acompletion(
+        self, prompt: str | list[dict[str, Any]], model: str | None = None
+    ) -> str:
+        messages, system = self._prepare_messages(prompt)
+
+        model = model or self.model_name
+        if not model:
+            raise ValueError("Model name is required for Anthropic client.")
+
+        kwargs = {"model": model, "max_tokens": self.max_tokens, "messages": messages}
+        if system:
+            kwargs["system"] = system
+
+        response = await self.async_client.messages.create(**kwargs)
+        self._track_cost(response, model)
+        return response.content[0].text
+
+    def _prepare_messages(
+        self, prompt: str | list[dict[str, Any]]
+    ) -> tuple[list[dict[str, Any]], str | None]:
+        """Prepare messages and extract system prompt for Anthropic API."""
+        system = None
+
+        if isinstance(prompt, str):
+            messages = [{"role": "user", "content": prompt}]
+        elif isinstance(prompt, list) and all(isinstance(item, dict) for item in prompt):
+            # Extract system message if present (Anthropic handles system separately)
+            messages = []
+            for msg in prompt:
+                if msg.get("role") == "system":
+                    system = msg.get("content")
+                else:
+                    messages.append(msg)
+        else:
+            raise ValueError(f"Invalid prompt type: {type(prompt)}")
+
+        return messages, system
+
+    def _track_cost(self, response: anthropic.types.Message, model: str):
+        self.model_call_counts[model] += 1
+        self.model_input_tokens[model] += response.usage.input_tokens
+        self.model_output_tokens[model] += response.usage.output_tokens
+        self.model_total_tokens[model] += response.usage.input_tokens + response.usage.output_tokens
+
+        # Track last call for handler to read
+        self.last_prompt_tokens = response.usage.input_tokens
+        self.last_completion_tokens = response.usage.output_tokens
+
+    def get_usage_summary(self) -> UsageSummary:
+        model_summaries = {}
+        for model in self.model_call_counts:
+            model_summaries[model] = ModelUsageSummary(
+                total_calls=self.model_call_counts[model],
+                total_input_tokens=self.model_input_tokens[model],
+                total_output_tokens=self.model_output_tokens[model],
+            )
+        return UsageSummary(model_usage_summaries=model_summaries)
+
+    def get_last_usage(self) -> ModelUsageSummary:
+        return ModelUsageSummary(
+            total_calls=1,
+            total_input_tokens=self.last_prompt_tokens,
+            total_output_tokens=self.last_completion_tokens,
+        )

groknroll/clients/azure_openai.py

@@ -0,0 +1,142 @@
+import os
+from collections import defaultdict
+from typing import Any
+
+import openai
+from dotenv import load_dotenv
+
+from groknroll.clients.base_lm import BaseLM
+from groknroll.core.types import ModelUsageSummary, UsageSummary
+
+load_dotenv()
+
+# Load API key from environment variable
+DEFAULT_AZURE_OPENAI_API_KEY = os.getenv("AZURE_OPENAI_API_KEY")
+
+
+class AzureOpenAIClient(BaseLM):
+    """
+    LM Client for running models with the Azure OpenAI API.
+    """
+
+    def __init__(
+        self,
+        api_key: str | None = None,
+        model_name: str | None = None,
+        azure_endpoint: str | None = None,
+        api_version: str | None = None,
+        azure_deployment: str | None = None,
+        **kwargs,
+    ):
+        super().__init__(model_name=model_name, **kwargs)
+
+        if api_key is None:
+            api_key = DEFAULT_AZURE_OPENAI_API_KEY
+
+        if azure_endpoint is None:
+            azure_endpoint = os.getenv("AZURE_OPENAI_ENDPOINT")
+
+        if api_version is None:
+            api_version = os.getenv("AZURE_OPENAI_API_VERSION", "2024-02-01")
+
+        if azure_deployment is None:
+            azure_deployment = os.getenv("AZURE_OPENAI_DEPLOYMENT")
+
+        if azure_endpoint is None:
+            raise ValueError(
+                "azure_endpoint is required for Azure OpenAI client. "
+                "Set it via argument or AZURE_OPENAI_ENDPOINT environment variable."
+            )
+
+        self.client = openai.AzureOpenAI(
+            api_key=api_key,
+            azure_endpoint=azure_endpoint,
+            api_version=api_version,
+            azure_deployment=azure_deployment,
+        )
+        self.async_client = openai.AsyncAzureOpenAI(
+            api_key=api_key,
+            azure_endpoint=azure_endpoint,
+            api_version=api_version,
+            azure_deployment=azure_deployment,
+        )
+        self.model_name = model_name
+        self.azure_deployment = azure_deployment
+
+        # Per-model usage tracking
+        self.model_call_counts: dict[str, int] = defaultdict(int)
+        self.model_input_tokens: dict[str, int] = defaultdict(int)
+        self.model_output_tokens: dict[str, int] = defaultdict(int)
+        self.model_total_tokens: dict[str, int] = defaultdict(int)
+
+    def completion(self, prompt: str | list[dict[str, Any]], model: str | None = None) -> str:
+        if isinstance(prompt, str):
+            messages = [{"role": "user", "content": prompt}]
+        elif isinstance(prompt, list) and all(isinstance(item, dict) for item in prompt):
+            messages = prompt
+        else:
+            raise ValueError(f"Invalid prompt type: {type(prompt)}")
+
+        model = model or self.model_name
+        if not model:
+            raise ValueError("Model name is required for Azure OpenAI client.")
+
+        response = self.client.chat.completions.create(
+            model=model,
+            messages=messages,
+        )
+        self._track_cost(response, model)
+        return response.choices[0].message.content
+
+    async def acompletion(
+        self, prompt: str | list[dict[str, Any]], model: str | None = None
+    ) -> str:
+        if isinstance(prompt, str):
+            messages = [{"role": "user", "content": prompt}]
+        elif isinstance(prompt, list) and all(isinstance(item, dict) for item in prompt):
+            messages = prompt
+        else:
+            raise ValueError(f"Invalid prompt type: {type(prompt)}")
+
+        model = model or self.model_name
+        if not model:
+            raise ValueError("Model name is required for Azure OpenAI client.")
+
+        response = await self.async_client.chat.completions.create(
+            model=model,
+            messages=messages,
+        )
+        self._track_cost(response, model)
+        return response.choices[0].message.content
+
+    def _track_cost(self, response: openai.ChatCompletion, model: str):
+        self.model_call_counts[model] += 1
+
+        usage = getattr(response, "usage", None)
+        if usage is None:
+            raise ValueError("No usage data received. Tracking tokens not possible.")
+
+        self.model_input_tokens[model] += usage.prompt_tokens
+        self.model_output_tokens[model] += usage.completion_tokens
+        self.model_total_tokens[model] += usage.total_tokens
+
+        # Track last call for handler to read
+        self.last_prompt_tokens = usage.prompt_tokens
+        self.last_completion_tokens = usage.completion_tokens
+
+    def get_usage_summary(self) -> UsageSummary:
+        model_summaries = {}
+        for model in self.model_call_counts:
+            model_summaries[model] = ModelUsageSummary(
+                total_calls=self.model_call_counts[model],
+                total_input_tokens=self.model_input_tokens[model],
+                total_output_tokens=self.model_output_tokens[model],
+            )
+        return UsageSummary(model_usage_summaries=model_summaries)
+
+    def get_last_usage(self) -> ModelUsageSummary:
+        return ModelUsageSummary(
+            total_calls=1,
+            total_input_tokens=self.last_prompt_tokens,
+            total_output_tokens=self.last_completion_tokens,
+        )

groknroll/clients/base_lm.py

@@ -0,0 +1,33 @@
+from abc import ABC, abstractmethod
+from typing import Any
+
+from groknroll.core.types import UsageSummary
+
+
+class BaseLM(ABC):
+    """
+    Base class for all language model routers / clients. When the RLM makes sub-calls, it currently
+    does so in a model-agnostic way, so this class provides a base interface for all language models.
+    """
+
+    def __init__(self, model_name: str, **kwargs):
+        self.model_name = model_name
+        self.kwargs = kwargs
+
+    @abstractmethod
+    def completion(self, prompt: str | dict[str, Any]) -> str:
+        raise NotImplementedError
+
+    @abstractmethod
+    async def acompletion(self, prompt: str | dict[str, Any]) -> str:
+        raise NotImplementedError
+
+    @abstractmethod
+    def get_usage_summary(self) -> UsageSummary:
+        """Get cost summary for all model calls."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def get_last_usage(self) -> UsageSummary:
+        """Get the last cost summary of the model."""
+        raise NotImplementedError

groknroll/clients/gemini.py

@@ -0,0 +1,162 @@
+import os
+from collections import defaultdict
+from typing import Any
+
+from dotenv import load_dotenv
+from google import genai
+from google.genai import types
+
+from groknroll.clients.base_lm import BaseLM
+from groknroll.core.types import ModelUsageSummary, UsageSummary
+
+load_dotenv()
+
+DEFAULT_GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")
+
+
+class GeminiClient(BaseLM):
+    """
+    LM Client for running models with the Google Gemini API.
+    Uses the official google-genai SDK.
+    """
+
+    def __init__(
+        self,
+        api_key: str | None = None,
+        model_name: str | None = "gemini-2.5-flash",
+        **kwargs,
+    ):
+        super().__init__(model_name=model_name, **kwargs)
+
+        if api_key is None:
+            api_key = DEFAULT_GEMINI_API_KEY
+
+        if api_key is None:
+            raise ValueError(
+                "Gemini API key is required. Set GEMINI_API_KEY env var or pass api_key."
+            )
+
+        self.client = genai.Client(api_key=api_key)
+        self.model_name = model_name
+
+        # Per-model usage tracking
+        self.model_call_counts: dict[str, int] = defaultdict(int)
+        self.model_input_tokens: dict[str, int] = defaultdict(int)
+        self.model_output_tokens: dict[str, int] = defaultdict(int)
+        self.model_total_tokens: dict[str, int] = defaultdict(int)
+
+        # Last call tracking
+        self.last_prompt_tokens = 0
+        self.last_completion_tokens = 0
+
+    def completion(self, prompt: str | list[dict[str, Any]], model: str | None = None) -> str:
+        contents, system_instruction = self._prepare_contents(prompt)
+
+        model = model or self.model_name
+        if not model:
+            raise ValueError("Model name is required for Gemini client.")
+
+        config = None
+        if system_instruction:
+            config = types.GenerateContentConfig(system_instruction=system_instruction)
+
+        response = self.client.models.generate_content(
+            model=model,
+            contents=contents,
+            config=config,
+        )
+
+        self._track_cost(response, model)
+        return response.text
+
+    async def acompletion(
+        self, prompt: str | list[dict[str, Any]], model: str | None = None
+    ) -> str:
+        contents, system_instruction = self._prepare_contents(prompt)
+
+        model = model or self.model_name
+        if not model:
+            raise ValueError("Model name is required for Gemini client.")
+
+        config = None
+        if system_instruction:
+            config = types.GenerateContentConfig(system_instruction=system_instruction)
+
+        # google-genai SDK supports async via aio interface
+        response = await self.client.aio.models.generate_content(
+            model=model,
+            contents=contents,
+            config=config,
+        )
+
+        self._track_cost(response, model)
+        return response.text
+
+    def _prepare_contents(
+        self, prompt: str | list[dict[str, Any]]
+    ) -> tuple[list[types.Content] | str, str | None]:
+        """Prepare contents and extract system instruction for Gemini API."""
+        system_instruction = None
+
+        if isinstance(prompt, str):
+            return prompt, None
+
+        if isinstance(prompt, list) and all(isinstance(item, dict) for item in prompt):
+            # Convert OpenAI-style messages to Gemini format
+            contents = []
+            for msg in prompt:
+                role = msg.get("role")
+                content = msg.get("content", "")
+
+                if role == "system":
+                    # Gemini handles system instruction separately
+                    system_instruction = content
+                elif role == "user":
+                    contents.append(types.Content(role="user", parts=[types.Part(text=content)]))
+                elif role == "assistant":
+                    # Gemini uses "model" instead of "assistant"
+                    contents.append(types.Content(role="model", parts=[types.Part(text=content)]))
+                else:
+                    # Default to user role for unknown roles
+                    contents.append(types.Content(role="user", parts=[types.Part(text=content)]))
+
+            return contents, system_instruction
+
+        raise ValueError(f"Invalid prompt type: {type(prompt)}")
+
+    def _track_cost(self, response: types.GenerateContentResponse, model: str):
+        self.model_call_counts[model] += 1
+
+        # Extract token usage from response
+        usage = response.usage_metadata
+        if usage:
+            input_tokens = usage.prompt_token_count or 0
+            output_tokens = usage.candidates_token_count or 0
+
+            self.model_input_tokens[model] += input_tokens
+            self.model_output_tokens[model] += output_tokens
+            self.model_total_tokens[model] += input_tokens + output_tokens
+
+            # Track last call for handler to read
+            self.last_prompt_tokens = input_tokens
+            self.last_completion_tokens = output_tokens
+        else:
+            self.last_prompt_tokens = 0
+            self.last_completion_tokens = 0
+
+    def get_usage_summary(self) -> UsageSummary:
+        model_summaries = {}
+        for model in self.model_call_counts:
+            model_summaries[model] = ModelUsageSummary(
+                total_calls=self.model_call_counts[model],
+                total_input_tokens=self.model_input_tokens[model],
+                total_output_tokens=self.model_output_tokens[model],
+            )
+        return UsageSummary(model_usage_summaries=model_summaries)
+
+    def get_last_usage(self) -> ModelUsageSummary:
+        return ModelUsageSummary(
+            total_calls=1,
+            total_input_tokens=self.last_prompt_tokens,
+            total_output_tokens=self.last_completion_tokens,
+        )

groknroll/clients/litellm.py

@@ -0,0 +1,105 @@
+from collections import defaultdict
+from typing import Any
+
+import litellm
+
+from groknroll.clients.base_lm import BaseLM
+from groknroll.core.types import ModelUsageSummary, UsageSummary
+
+
+class LiteLLMClient(BaseLM):
+    """
+    LM Client for running models with LiteLLM.
+    LiteLLM provides a unified interface to 100+ LLM providers.
+    """
+
+    def __init__(
+        self,
+        model_name: str | None = None,
+        api_key: str | None = None,
+        api_base: str | None = None,
+        **kwargs,
+    ):
+        super().__init__(model_name=model_name, **kwargs)
+        self.model_name = model_name
+        self.api_key = api_key
+        self.api_base = api_base
+
+        # Per-model usage tracking
+        self.model_call_counts: dict[str, int] = defaultdict(int)
+        self.model_input_tokens: dict[str, int] = defaultdict(int)
+        self.model_output_tokens: dict[str, int] = defaultdict(int)
+        self.model_total_tokens: dict[str, int] = defaultdict(int)
+
+    def completion(self, prompt: str | list[dict[str, Any]], model: str | None = None) -> str:
+        if isinstance(prompt, str):
+            messages = [{"role": "user", "content": prompt}]
+        elif isinstance(prompt, list) and all(isinstance(item, dict) for item in prompt):
+            messages = prompt
+        else:
+            raise ValueError(f"Invalid prompt type: {type(prompt)}")
+
+        model = model or self.model_name
+        if not model:
+            raise ValueError("Model name is required for LiteLLM client.")
+
+        kwargs = {"model": model, "messages": messages}
+        if self.api_key:
+            kwargs["api_key"] = self.api_key
+        if self.api_base:
+            kwargs["api_base"] = self.api_base
+
+        response = litellm.completion(**kwargs)
+        self._track_cost(response, model)
+        return response.choices[0].message.content
+
+    async def acompletion(
+        self, prompt: str | list[dict[str, Any]], model: str | None = None
+    ) -> str:
+        if isinstance(prompt, str):
+            messages = [{"role": "user", "content": prompt}]
+        elif isinstance(prompt, list) and all(isinstance(item, dict) for item in prompt):
+            messages = prompt
+        else:
+            raise ValueError(f"Invalid prompt type: {type(prompt)}")
+
+        model = model or self.model_name
+        if not model:
+            raise ValueError("Model name is required for LiteLLM client.")
+
+        kwargs = {"model": model, "messages": messages}
+        if self.api_key:
+            kwargs["api_key"] = self.api_key
+        if self.api_base:
+            kwargs["api_base"] = self.api_base
+
+        response = await litellm.acompletion(**kwargs)
+        self._track_cost(response, model)
+        return response.choices[0].message.content
+
+    def _track_cost(self, response, model: str):
+        self.model_call_counts[model] += 1
+        self.model_input_tokens[model] += response.usage.prompt_tokens
+        self.model_output_tokens[model] += response.usage.completion_tokens
+        self.model_total_tokens[model] += response.usage.total_tokens
+
+        # Track last call for handler to read
+        self.last_prompt_tokens = response.usage.prompt_tokens
+        self.last_completion_tokens = response.usage.completion_tokens
+
+    def get_usage_summary(self) -> UsageSummary:
+        model_summaries = {}
+        for model in self.model_call_counts:
+            model_summaries[model] = ModelUsageSummary(
+                total_calls=self.model_call_counts[model],
+                total_input_tokens=self.model_input_tokens[model],
+                total_output_tokens=self.model_output_tokens[model],
+            )
+        return UsageSummary(model_usage_summaries=model_summaries)
+
+    def get_last_usage(self) -> ModelUsageSummary:
+        return ModelUsageSummary(
+            total_calls=1,
+            total_input_tokens=self.last_prompt_tokens,
+            total_output_tokens=self.last_completion_tokens,
+        )

groknroll/clients/openai.py

@@ -0,0 +1,129 @@
+import os
+from collections import defaultdict
+from typing import Any
+
+import openai
+from dotenv import load_dotenv
+
+from groknroll.clients.base_lm import BaseLM
+from groknroll.core.types import ModelUsageSummary, UsageSummary
+
+load_dotenv()
+
+# Load API keys from environment variables
+DEFAULT_OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
+DEFAULT_OPENROUTER_API_KEY = os.getenv("OPENROUTER_API_KEY")
+DEFAULT_VERCEL_API_KEY = os.getenv("AI_GATEWAY_API_KEY")
+DEFAULT_PRIME_INTELLECT_BASE_URL = "https://api.pinference.ai/api/v1/"
+
+
+class OpenAIClient(BaseLM):
+    """
+    LM Client for running models with the OpenAI API. Works with vLLM as well.
+    """
+
+    def __init__(
+        self,
+        api_key: str | None = None,
+        model_name: str | None = None,
+        base_url: str | None = None,
+        **kwargs,
+    ):
+        super().__init__(model_name=model_name, **kwargs)
+
+        if api_key is None:
+            if base_url == "https://api.openai.com/v1" or base_url is None:
+                api_key = DEFAULT_OPENAI_API_KEY
+            elif base_url == "https://openrouter.ai/api/v1":
+                api_key = DEFAULT_OPENROUTER_API_KEY
+            elif base_url == "https://ai-gateway.vercel.sh/v1":
+                api_key = DEFAULT_VERCEL_API_KEY
+
+        # For vLLM, set base_url to local vLLM server address.
+        self.client = openai.OpenAI(api_key=api_key, base_url=base_url)
+        self.async_client = openai.AsyncOpenAI(api_key=api_key, base_url=base_url)
+        self.model_name = model_name
+
+        # Per-model usage tracking
+        self.model_call_counts: dict[str, int] = defaultdict(int)
+        self.model_input_tokens: dict[str, int] = defaultdict(int)
+        self.model_output_tokens: dict[str, int] = defaultdict(int)
+        self.model_total_tokens: dict[str, int] = defaultdict(int)
+
+    def completion(self, prompt: str | list[dict[str, Any]], model: str | None = None) -> str:
+        if isinstance(prompt, str):
+            messages = [{"role": "user", "content": prompt}]
+        elif isinstance(prompt, list) and all(isinstance(item, dict) for item in prompt):
+            messages = prompt
+        else:
+            raise ValueError(f"Invalid prompt type: {type(prompt)}")
+
+        model = model or self.model_name
+        if not model:
+            raise ValueError("Model name is required for OpenAI client.")
+
+        extra_body = {}
+        if self.client.base_url == DEFAULT_PRIME_INTELLECT_BASE_URL:
+            extra_body["usage"] = {"include": True}
+
+        response = self.client.chat.completions.create(
+            model=model, messages=messages, extra_body=extra_body
+        )
+        self._track_cost(response, model)
+        return response.choices[0].message.content
+
+    async def acompletion(
+        self, prompt: str | list[dict[str, Any]], model: str | None = None
+    ) -> str:
+        if isinstance(prompt, str):
+            messages = [{"role": "user", "content": prompt}]
+        elif isinstance(prompt, list) and all(isinstance(item, dict) for item in prompt):
+            messages = prompt
+        else:
+            raise ValueError(f"Invalid prompt type: {type(prompt)}")
+
+        model = model or self.model_name
+        if not model:
+            raise ValueError("Model name is required for OpenAI client.")
+
+        extra_body = {}
+        if self.client.base_url == DEFAULT_PRIME_INTELLECT_BASE_URL:
+            extra_body["usage"] = {"include": True}
+
+        response = await self.async_client.chat.completions.create(
+            model=model, messages=messages, extra_body=extra_body
+        )
+        self._track_cost(response, model)
+        return response.choices[0].message.content
+
+    def _track_cost(self, response: openai.ChatCompletion, model: str):
+        self.model_call_counts[model] += 1
+
+        usage = getattr(response, "usage", None)
+        if usage is None:
+            raise ValueError("No usage data received. Tracking tokens not possible.")
+
+        self.model_input_tokens[model] += usage.prompt_tokens
+        self.model_output_tokens[model] += usage.completion_tokens
+        self.model_total_tokens[model] += usage.total_tokens
+
+        # Track last call for handler to read
+        self.last_prompt_tokens = usage.prompt_tokens
+        self.last_completion_tokens = usage.completion_tokens
+
+    def get_usage_summary(self) -> UsageSummary:
+        model_summaries = {}
+        for model in self.model_call_counts:
+            model_summaries[model] = ModelUsageSummary(
+                total_calls=self.model_call_counts[model],
+                total_input_tokens=self.model_input_tokens[model],
+                total_output_tokens=self.model_output_tokens[model],
+            )
+        return UsageSummary(model_usage_summaries=model_summaries)
+
+    def get_last_usage(self) -> ModelUsageSummary:
+        return ModelUsageSummary(
+            total_calls=1,
+            total_input_tokens=self.last_prompt_tokens,
+            total_output_tokens=self.last_completion_tokens,
+        )