janito 2.15.0__py3-none-any.whl → 2.16.0__py3-none-any.whl

janito/drivers/zai/driver.py

@@ -12,18 +12,20 @@ from janito.llm.driver_input import DriverInput
 from janito.driver_events import RequestFinished, RequestStatus, RateLimitRetry
 from janito.llm.message_parts import TextMessagePart, FunctionCallMessagePart
 
-try:
-    import openai
-    available = True
-    unavailable_reason = None
-except ImportError:
-    available = False
-    unavailable_reason = "openai module not installed"
+import openai
 
 
 class ZAIModelDriver(LLMDriver):
-    available = available
-    unavailable_reason = unavailable_reason
+    # Check if required dependencies are available
+    try:
+        from zai import ZaiClient
+
+        available = True
+        unavailable_reason = None
+    except ImportError as e:
+        available = False
+        unavailable_reason = f"Missing dependency: {str(e)}"
+
     def _get_message_from_result(self, result):
         """Extract the message object from the provider result (Z.AI-specific)."""
         if hasattr(result, "choices") and result.choices:
@@ -256,14 +258,11 @@ class ZAIModelDriver(LLMDriver):
         try:
             if not config.api_key:
                 provider_name = getattr(self, "provider_name", "ZAI")
-                print(
-                    f"[ERROR] No API key found for provider '{provider_name}'. Please set the API key using:"
-                )
-                print(f"  janito --set-api-key YOUR_API_KEY -p {provider_name.lower()}")
-                print(
-                    f"Or set the {provider_name.upper()}_API_KEY environment variable."
+                from janito.llm.auth_utils import handle_missing_api_key
+
+                handle_missing_api_key(
+                    provider_name, f"{provider_name.upper()}_API_KEY"
                 )
-                raise ValueError(f"API key is required for provider '{provider_name}'")
 
             api_key_display = str(config.api_key)
             if api_key_display and len(api_key_display) > 8:
@@ -284,12 +283,10 @@ class ZAIModelDriver(LLMDriver):
                     flush=True,
                 )
 
-            # Use OpenAI SDK for Z.AI API compatibility
-            try:
-                import openai
-            except ImportError:
-                raise ImportError("openai module is not available. Please install it with: pip install openai")
-            client = openai.OpenAI(
+            # Use the official Z.ai SDK
+            from zai import ZaiClient
+
+            client = ZaiClient(
                 api_key=config.api_key, base_url="https://api.z.ai/api/paas/v4/"
             )
             return client

janito/formatting_token.py

@@ -29,7 +29,7 @@ def format_token_message_summary(
     msg_count, usage, width=96, use_rich=False, elapsed=None
 ):
     """
-    Returns a string (rich or pt markup) summarizing message count, last token usage, and elapsed time.
+    Returns a string (rich or pt markup) summarizing message count, last token usage, elapsed time, and tokens per second.
     """
     left = f" Messages: {'[' if use_rich else '<'}msg_count{']' if use_rich else '>'}{msg_count}{'[/msg_count]' if use_rich else '</msg_count>'}"
     tokens_part = ""
@@ -42,10 +42,14 @@ def format_token_message_summary(
             f"Completion: {format_tokens(completion_tokens, 'tokens_out', use_rich)}, "
             f"Total: {format_tokens(total_tokens, 'tokens_total', use_rich)}"
         )
-    elapsed_part = (
-        f" | Elapsed: [cyan]{elapsed:.2f}s[/cyan]" if elapsed is not None else ""
-    )
-    return f"{left}{tokens_part}{elapsed_part}"
+    elapsed_part = ""
+    tps_part = ""
+    if elapsed is not None and elapsed > 0:
+        elapsed_part = f" | Elapsed: [cyan]{elapsed:.2f}s[/cyan]"
+        if usage and total_tokens:
+            tokens_per_second = total_tokens / elapsed
+            tps_part = f" | TPS: {int(tokens_per_second)}"
+    return f"{left}{tokens_part}{elapsed_part}{tps_part}"
 
 
 def print_token_message_summary(

janito/llm/auth_utils.py

@@ -0,0 +1,21 @@
+"""
+Authentication utilities for LLM providers.
+"""
+
+import sys
+
+
+def handle_missing_api_key(provider_name: str, env_var_name: str) -> None:
+    """
+    Handle missing API key by printing error message and exiting.
+
+    Args:
+        provider_name: Name of the provider (e.g., 'alibaba', 'openai')
+        env_var_name: Environment variable name (e.g., 'ALIBABA_API_KEY')
+    """
+    print(
+        f"[ERROR] No API key found for provider '{provider_name}'. Please set the API key using:"
+    )
+    print(f"  janito --set-api-key YOUR_API_KEY -p {provider_name}")
+    print(f"Or set the {env_var_name} environment variable.")
+    sys.exit(1)

janito/providers/alibaba/provider.py

@@ -17,7 +17,9 @@ class AlibabaProvider(LLMProvider):
     NAME = "alibaba"
     MAINTAINER = "João Pinto <janito@ikignosis.org>"
     MODEL_SPECS = MODEL_SPECS
-    DEFAULT_MODEL = "qwen3-coder-plus"  # Options: qwen-turbo, qwen-plus, qwen-max, qwen3-coder-plus
+    DEFAULT_MODEL = (
+        "qwen3-coder-plus"  # Options: qwen-turbo, qwen-plus, qwen-max, qwen3-coder-plus
+    )
 
     def __init__(
         self, auth_manager: LLMAuthManager = None, config: LLMDriverConfig = None
@@ -25,27 +27,29 @@ class AlibabaProvider(LLMProvider):
         # Always set a tools adapter so that even if the driver is unavailable,
         # generic code paths that expect provider.execute_tool() continue to work.
         self._tools_adapter = get_local_tools_adapter()
-        
+
         # Always initialize _driver_config to avoid AttributeError
         self._driver_config = config or LLMDriverConfig(model=None)
-        
+
         if not self.available:
             self._driver = None
         else:
             self.auth_manager = auth_manager or LLMAuthManager()
             self._api_key = self.auth_manager.get_credentials(type(self).NAME)
             if not self._api_key:
-                print(f"[ERROR] No API key found for provider '{self.name}'. Please set the API key using:")
-                print(f"  janito --set-api-key YOUR_API_KEY -p {self.name}")
-                print(f"Or set the ALIBABA_API_KEY environment variable.")
-            
+                from janito.llm.auth_utils import handle_missing_api_key
+
+                handle_missing_api_key(self.name, "ALIBABA_API_KEY")
+
             if not self._driver_config.model:
                 self._driver_config.model = self.DEFAULT_MODEL
             if not self._driver_config.api_key:
                 self._driver_config.api_key = self._api_key
             # Set Alibaba international endpoint as default base_url if not provided
             if not getattr(self._driver_config, "base_url", None):
-                self._driver_config.base_url = "https://dashscope-intl.aliyuncs.com/compatible-mode/v1"
+                self._driver_config.base_url = (
+                    "https://dashscope-intl.aliyuncs.com/compatible-mode/v1"
+                )
             self.fill_missing_device_info(self._driver_config)
             self._driver = None  # to be provided by factory/agent
 
@@ -99,4 +103,4 @@ class AlibabaProvider(LLMProvider):
         return self._tools_adapter.execute_by_name(tool_name, *args, **kwargs)
 
 
-LLMProviderRegistry.register(AlibabaProvider.NAME, AlibabaProvider)
+LLMProviderRegistry.register(AlibabaProvider.NAME, AlibabaProvider)

janito/providers/anthropic/provider.py

@@ -22,11 +22,10 @@ class AnthropicProvider(LLMProvider):
         self.auth_manager = auth_manager or LLMAuthManager()
         self._api_key = self.auth_manager.get_credentials(type(self).NAME)
         if not self._api_key:
-            print(f"[ERROR] No API key found for provider '{self.name}'. Please set the API key using:")
-            print(f"  janito --set-api-key YOUR_API_KEY -p {self.name}")
-            print(f"Or set the ANTHROPIC_API_KEY environment variable.")
-            return
-        
+            from janito.llm.auth_utils import handle_missing_api_key
+
+            handle_missing_api_key(self.name, "ANTHROPIC_API_KEY")
+
         self._tools_adapter = get_local_tools_adapter()
         self._driver_config = config or LLMDriverConfig(model=None)
         if not getattr(self._driver_config, "model", None):

janito/providers/azure_openai/provider.py

@@ -33,11 +33,10 @@ class AzureOpenAIProvider(LLMProvider):
         self._auth_manager = auth_manager or LLMAuthManager()
         self._api_key = self._auth_manager.get_credentials(type(self).NAME)
         if not self._api_key:
-            print(f"[ERROR] No API key found for provider '{self.name}'. Please set the API key using:")
-            print(f"  janito --set-api-key YOUR_API_KEY -p {self.name}")
-            print(f"Or set the AZURE_OPENAI_API_KEY environment variable.")
-            return
-        
+            from janito.llm.auth_utils import handle_missing_api_key
+
+            handle_missing_api_key(self.name, "AZURE_OPENAI_API_KEY")
+
         self._tools_adapter = get_local_tools_adapter()
         self._driver_config = config or LLMDriverConfig(model=None)
         if not self._driver_config.model:

janito/providers/deepseek/provider.py

@@ -31,11 +31,10 @@ class DeepSeekProvider(LLMProvider):
             self.auth_manager = auth_manager or LLMAuthManager()
             self._api_key = self.auth_manager.get_credentials(type(self).NAME)
             if not self._api_key:
-                print(f"[ERROR] No API key found for provider '{self.name}'. Please set the API key using:")
-                print(f"  janito --set-api-key YOUR_API_KEY -p {self.name}")
-                print(f"Or set the DEEPSEEK_API_KEY environment variable.")
-                return
-            
+                from janito.llm.auth_utils import handle_missing_api_key
+
+                handle_missing_api_key(self.name, "DEEPSEEK_API_KEY")
+
             self._tools_adapter = get_local_tools_adapter()
             self._driver_config = config or LLMDriverConfig(model=None)
             if not self._driver_config.model:

janito/providers/google/provider.py

@@ -33,11 +33,10 @@ class GoogleProvider(LLMProvider):
             self.auth_manager = auth_manager or LLMAuthManager()
             self._api_key = self.auth_manager.get_credentials(type(self).name)
             if not self._api_key:
-                print(f"[ERROR] No API key found for provider '{self.name}'. Please set the API key using:")
-                print(f"  janito --set-api-key YOUR_API_KEY -p {self.name}")
-                print(f"Or set the GOOGLE_API_KEY environment variable.")
-                return
-            
+                from janito.llm.auth_utils import handle_missing_api_key
+
+                handle_missing_api_key(self.name, "GOOGLE_API_KEY")
+
             self._tools_adapter = get_local_tools_adapter()
             self._driver_config = config or LLMDriverConfig(model=None)
             # Only set default if model is not set by CLI/config

janito/providers/moonshotai/provider.py

@@ -17,44 +17,53 @@ class MoonshotAIProvider(LLMProvider):
     def __init__(
         self, auth_manager: LLMAuthManager = None, config: LLMDriverConfig = None
     ):
+        self._tools_adapter = get_local_tools_adapter()
+        self._driver = None
+
         if not self.available:
-            self._tools_adapter = get_local_tools_adapter()
-            self._driver = None
-        else:
-            self.auth_manager = auth_manager or LLMAuthManager()
-            self._api_key = self.auth_manager.get_credentials(type(self).name)
-            if not self._api_key:
-                print(f"[ERROR] No API key found for provider '{self.name}'. Please set the API key using:")
-                print(f"  janito --set-api-key YOUR_API_KEY -p {self.name}")
-                print(f"Or set the MOONSHOTAI_API_KEY environment variable.")
-                return
-            
-            self._tools_adapter = get_local_tools_adapter()
-            self._driver_config = config or LLMDriverConfig(model=None)
-            if not self._driver_config.model:
-                self._driver_config.model = self.DEFAULT_MODEL
-            if not self._driver_config.api_key:
-                self._driver_config.api_key = self._api_key
-            # Set only the correct token parameter for the model
-            model_name = self._driver_config.model
-            model_spec = self.MODEL_SPECS.get(model_name)
-            if hasattr(self._driver_config, "max_tokens"):
-                self._driver_config.max_tokens = None
-            if hasattr(self._driver_config, "max_completion_tokens"):
-                self._driver_config.max_completion_tokens = None
-            if model_spec:
-                if getattr(model_spec, "thinking_supported", False):
-                    max_cot = getattr(model_spec, "max_cot", None)
-                    if max_cot and max_cot != "N/A":
-                        self._driver_config.max_completion_tokens = int(max_cot)
-                else:
-                    max_response = getattr(model_spec, "max_response", None)
-                    if max_response and max_response != "N/A":
-                        self._driver_config.max_tokens = int(max_response)
-            self.fill_missing_device_info(self._driver_config)
-            self._driver = None
-            # Set MoonshotAI base_url
-            self._driver_config.base_url = "https://api.moonshot.ai/v1"
+            return
+
+        self._initialize_config(auth_manager, config)
+        self._setup_model_config()
+        self._driver_config.base_url = "https://api.moonshot.ai/v1"
+
+    def _initialize_config(self, auth_manager, config):
+        """Initialize configuration and API key."""
+        self.auth_manager = auth_manager or LLMAuthManager()
+        self._api_key = self.auth_manager.get_credentials(type(self).name)
+        if not self._api_key:
+            from janito.llm.auth_utils import handle_missing_api_key
+
+            handle_missing_api_key(self.name, "MOONSHOTAI_API_KEY")
+
+        self._driver_config = config or LLMDriverConfig(model=None)
+        if not self._driver_config.model:
+            self._driver_config.model = self.DEFAULT_MODEL
+        if not self._driver_config.api_key:
+            self._driver_config.api_key = self._api_key
+
+    def _setup_model_config(self):
+        """Configure token limits based on model specifications."""
+        model_name = self._driver_config.model
+        model_spec = self.MODEL_SPECS.get(model_name)
+
+        # Reset token parameters
+        if hasattr(self._driver_config, "max_tokens"):
+            self._driver_config.max_tokens = None
+        if hasattr(self._driver_config, "max_completion_tokens"):
+            self._driver_config.max_completion_tokens = None
+
+        if model_spec:
+            if getattr(model_spec, "thinking_supported", False):
+                max_cot = getattr(model_spec, "max_cot", None)
+                if max_cot and max_cot != "N/A":
+                    self._driver_config.max_completion_tokens = int(max_cot)
+            else:
+                max_response = getattr(model_spec, "max_response", None)
+                if max_response and max_response != "N/A":
+                    self._driver_config.max_tokens = int(max_response)
+
+        self.fill_missing_device_info(self._driver_config)
 
     @property
     def driver(self) -> OpenAIModelDriver:

janito/providers/openai/provider.py

@@ -22,46 +22,52 @@ class OpenAIProvider(LLMProvider):
     def __init__(
         self, auth_manager: LLMAuthManager = None, config: LLMDriverConfig = None
     ):
+        self._tools_adapter = get_local_tools_adapter()
+        self._driver = None
+
         if not self.available:
-            # Even when the OpenAI driver is unavailable we still need a tools adapter
-            # so that any generic logic that expects `execute_tool()` to work does not
-            # crash with an AttributeError when it tries to access `self._tools_adapter`.
-            self._tools_adapter = get_local_tools_adapter()
-            self._driver = None
-        else:
-            self.auth_manager = auth_manager or LLMAuthManager()
-            self._api_key = self.auth_manager.get_credentials(type(self).NAME)
-            if not self._api_key:
-                print(f"[ERROR] No API key found for provider '{self.name}'. Please set the API key using:")
-                print(f"  janito --set-api-key YOUR_API_KEY -p {self.name}")
-                print(f"Or set the OPENAI_API_KEY environment variable.")
-                return
-            
-            self._tools_adapter = get_local_tools_adapter()
-            self._driver_config = config or LLMDriverConfig(model=None)
-            if not self._driver_config.model:
-                self._driver_config.model = self.DEFAULT_MODEL
-            if not self._driver_config.api_key:
-                self._driver_config.api_key = self._api_key
-            # Set only the correct token parameter for the model
-            model_name = self._driver_config.model
-            model_spec = self.MODEL_SPECS.get(model_name)
-            # Remove both to avoid stale values
-            if hasattr(self._driver_config, "max_tokens"):
-                self._driver_config.max_tokens = None
-            if hasattr(self._driver_config, "max_completion_tokens"):
-                self._driver_config.max_completion_tokens = None
-            if model_spec:
-                if getattr(model_spec, "thinking_supported", False):
-                    max_cot = getattr(model_spec, "max_cot", None)
-                    if max_cot and max_cot != "N/A":
-                        self._driver_config.max_completion_tokens = int(max_cot)
-                else:
-                    max_response = getattr(model_spec, "max_response", None)
-                    if max_response and max_response != "N/A":
-                        self._driver_config.max_tokens = int(max_response)
-            self.fill_missing_device_info(self._driver_config)
-            self._driver = None  # to be provided by factory/agent
+            return
+
+        self._initialize_config(auth_manager, config)
+        self._setup_model_config()
+
+    def _initialize_config(self, auth_manager, config):
+        """Initialize configuration and API key."""
+        self.auth_manager = auth_manager or LLMAuthManager()
+        self._api_key = self.auth_manager.get_credentials(type(self).NAME)
+        if not self._api_key:
+            from janito.llm.auth_utils import handle_missing_api_key
+
+            handle_missing_api_key(self.name, "OPENAI_API_KEY")
+
+        self._driver_config = config or LLMDriverConfig(model=None)
+        if not self._driver_config.model:
+            self._driver_config.model = self.DEFAULT_MODEL
+        if not self._driver_config.api_key:
+            self._driver_config.api_key = self._api_key
+
+    def _setup_model_config(self):
+        """Configure token limits based on model specifications."""
+        model_name = self._driver_config.model
+        model_spec = self.MODEL_SPECS.get(model_name)
+
+        # Reset token parameters
+        if hasattr(self._driver_config, "max_tokens"):
+            self._driver_config.max_tokens = None
+        if hasattr(self._driver_config, "max_completion_tokens"):
+            self._driver_config.max_completion_tokens = None
+
+        if model_spec:
+            if getattr(model_spec, "thinking_supported", False):
+                max_cot = getattr(model_spec, "max_cot", None)
+                if max_cot and max_cot != "N/A":
+                    self._driver_config.max_completion_tokens = int(max_cot)
+            else:
+                max_response = getattr(model_spec, "max_response", None)
+                if max_response and max_response != "N/A":
+                    self._driver_config.max_tokens = int(max_response)
+
+        self.fill_missing_device_info(self._driver_config)
 
     @property
     def driver(self) -> OpenAIModelDriver:

janito/providers/zai/provider.py

@@ -33,20 +33,14 @@ class ZAIProvider(LLMProvider):
         # crash with an AttributeError when it tries to access `self._tools_adapter`.
         self._tools_adapter = get_local_tools_adapter()
         self._driver = None
-        # Initialize _driver_config to avoid AttributeError
-        self._driver_config = LLMDriverConfig(model=None)
 
     def _setup_available(self, auth_manager, config):
         self.auth_manager = auth_manager or LLMAuthManager()
         self._api_key = self.auth_manager.get_credentials(type(self).NAME)
         if not self._api_key:
-            print(
-                f"[ERROR] No API key found for provider '{self.name}'. Please set the API key using:"
-            )
-            print(f"  janito --set-api-key YOUR_API_KEY -p {self.name}")
-            print(f"Or set the ZAI_API_KEY environment variable.")
-            self._tools_adapter = get_local_tools_adapter()
-            return
+            from janito.llm.auth_utils import handle_missing_api_key
+
+            handle_missing_api_key(self.name, "ZAI_API_KEY")
 
         self._tools_adapter = get_local_tools_adapter()
         self._driver_config = config or LLMDriverConfig(model=None)

janito/regions/__init__.py

@@ -0,0 +1,16 @@
+"""
+Region definitions and geolocation utilities for LLM providers.
+
+This module provides static region definitions for various LLM providers
+and utilities to determine optimal API endpoints based on user location.
+"""
+
+from .provider_regions import PROVIDER_REGIONS, get_optimal_endpoint
+from .geo_utils import get_user_location, get_closest_region
+
+__all__ = [
+    "PROVIDER_REGIONS",
+    "get_optimal_endpoint",
+    "get_user_location",
+    "get_closest_region",
+]

janito/regions/cli.py

@@ -0,0 +1,124 @@
+"""
+CLI commands for region management and geolocation utilities.
+"""
+
+import argparse
+import json
+from typing import Optional
+
+from .geo_utils import get_region_info
+from .provider_regions import (
+    get_provider_regions,
+    get_optimal_endpoint,
+    get_all_providers,
+)
+
+
+def handle_region_info(args=None):
+    """Display current region information."""
+    info = get_region_info()
+    print(json.dumps(info, indent=2))
+
+
+def handle_provider_regions(args):
+    """Display regions for a specific provider."""
+    provider = args.provider.lower()
+    regions = get_provider_regions(provider)
+
+    if not regions:
+        print(f"Provider '{provider}' not found")
+        return
+
+    result = {
+        "provider": provider,
+        "regions": [
+            {
+                "code": r.region_code,
+                "name": r.name,
+                "endpoint": r.endpoint,
+                "location": r.location,
+                "priority": r.priority,
+            }
+            for r in regions
+        ],
+    }
+    print(json.dumps(result, indent=2))
+
+
+def handle_optimal_endpoint(args):
+    """Get optimal endpoint for a provider based on user location."""
+    from .geo_utils import get_user_location, get_closest_region
+
+    country_code, _ = get_user_location()
+    major_region = get_closest_region(country_code)
+
+    provider = args.provider.lower()
+    endpoint = get_optimal_endpoint(provider, major_region)
+
+    if not endpoint:
+        print(f"Provider '{provider}' not found")
+        return
+
+    result = {
+        "provider": provider,
+        "user_region": major_region,
+        "country_code": country_code,
+        "optimal_endpoint": endpoint,
+    }
+    print(json.dumps(result, indent=2))
+
+
+def handle_list_providers(args=None):
+    """List all supported providers."""
+    providers = get_all_providers()
+    result = {"providers": providers, "count": len(providers)}
+    print(json.dumps(result, indent=2))
+
+
+def setup_region_parser(subparsers):
+    """Setup region-related CLI commands."""
+    region_parser = subparsers.add_parser(
+        "region", help="Region and geolocation utilities"
+    )
+    region_subparsers = region_parser.add_subparsers(
+        dest="region_command", help="Region commands"
+    )
+
+    # region info
+    info_parser = region_subparsers.add_parser(
+        "info", help="Show current region information"
+    )
+    info_parser.set_defaults(func=handle_region_info)
+
+    # region providers
+    providers_parser = region_subparsers.add_parser(
+        "providers", help="List all supported providers"
+    )
+    providers_parser.set_defaults(func=handle_list_providers)
+
+    # region list
+    list_parser = region_subparsers.add_parser(
+        "list", help="List regions for a provider"
+    )
+    list_parser.add_argument("provider", help="Provider name (e.g., openai, anthropic)")
+    list_parser.set_defaults(func=handle_provider_regions)
+
+    # region endpoint
+    endpoint_parser = region_subparsers.add_parser(
+        "endpoint", help="Get optimal endpoint for provider"
+    )
+    endpoint_parser.add_argument(
+        "provider", help="Provider name (e.g., openai, anthropic)"
+    )
+    endpoint_parser.set_defaults(func=handle_optimal_endpoint)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Region utilities")
+    setup_region_parser(parser.add_subparsers(dest="command"))
+
+    args = parser.parse_args()
+    if hasattr(args, "func"):
+        args.func(args)
+    else:
+        parser.print_help()