pygeai 0.6.0b14__py3-none-any.whl → 0.6.1__py3-none-any.whl

pygeai/cli/geai.py

@@ -64,9 +64,9 @@ class CLIDriver:
     user-friendly messages.
     """
 
-    def __init__(self, session=None) -> None:
+    def __init__(self, session=None, credentials_file=None) -> None:
         """
-        Initialize the CLI driver with optional session.
+        Initialize the CLI driver with optional session and credentials file.
         
         Sets up the session to be used while running commands, either with a 
         specified alias, environment variables, or function parameters.
@@ -74,8 +74,18 @@ class CLIDriver:
         
         :param session: Optional session object. If None, uses 'default' or
                        alias-specified session from command-line arguments.
+        :param credentials_file: Optional path to custom credentials file.
         """
+        from pygeai.core.common.config import get_settings
+        
         arguments = sys.argv
+        
+        if credentials_file or "--credentials" in arguments or "--creds" in arguments:
+            if not credentials_file:
+                credentials_file = self._get_credentials_file(arguments)
+            get_settings(credentials_file=credentials_file)
+            logger.debug(f"Using custom credentials file: {credentials_file}")
+        
         if "-a" in arguments or "--alias" in arguments:
             alias = self._get_alias(arguments)
             session = get_session(alias)
@@ -100,6 +110,25 @@ class CLIDriver:
         _ = arguments.pop(alias_index)
         alias = arguments.pop(alias_index)
         return alias
+    
+    def _get_credentials_file(self, arguments: List[str]) -> str:
+        """
+        Retrieves and removes credentials file path and flag from argument list.
+
+        :param arguments: List[str] - Command line arguments.
+        :return: str - The credentials file path.
+        :raises ValueError: If credentials flag is present but no value provided.
+        """
+        creds_index = None
+
+        if "--credentials" in arguments:
+            creds_index = arguments.index("--credentials")
+        elif "--creds" in arguments:
+            creds_index = arguments.index("--creds")
+
+        _ = arguments.pop(creds_index)
+        credentials_file = arguments.pop(creds_index)
+        return credentials_file
 
     def main(self, args: Optional[List[str]] = None) -> int:
         """
@@ -140,7 +169,7 @@ class CLIDriver:
                 arguments = []
 
             self.process_command(command, arguments)
-            logger.debug(f"Command completed successfully")
+            logger.debug("Command completed successfully")
             return ExitCode.SUCCESS
         except UnknownArgumentError as e:
             if hasattr(e, 'available_commands') and e.available_commands:

pygeai/cli/geai_proxy.py

@@ -11,6 +11,8 @@ from pygeai.proxy.managers import ServerManager
 from pygeai.proxy.config import ProxySettingsManager
 from pygeai.admin.clients import AdminClient
 from pygeai.core.utils.console import Console, StreamWriter
+
+
 class CustomProxyStream(StreamWriter):
     """
     Custom stream writer for proxy output with color support for TTY terminals.
@@ -40,8 +42,10 @@ class CustomProxyStream(StreamWriter):
         sys.stderr.write(f"{self.RED}{message}{self.RESET} {str(exception)}{end}")
         logger.exception(f"{message}{end}", exception)
 
+
 Console.set_writer(CustomProxyStream())
 
+
 def load_config(path: str) -> list:
     """
     Load server configuration from YAML or JSON file.
@@ -67,7 +71,7 @@ def load_config(path: str) -> list:
                     server_cfg['type'] = 'a2a'
                     servers.append(server_cfg)
             
-            if not 'a2aServers' in config and not 'mcpServers' in config:
+            if 'a2aServers' not in config and 'mcpServers' not in config:
                 Console.write_stderr("Error: MCP servers or A2A servers are required")
 
             return servers
@@ -279,7 +283,7 @@ async def main():
 
             return 1
         else:
-            Console.write_stdout(f" Done!")
+            Console.write_stdout(" Done!")
             Console.write_stdout(f"Organization: {result.get('organizationName')}")
             Console.write_stdout(f"Project: {result.get('projectName')}")
     except Exception as e:

pygeai/cli/install_man.py

@@ -80,7 +80,7 @@ def install_man_pages(system_wide: bool) -> int:
     # For local installation, remind user to update MANPATH
     if not system_wide:
         manpath_cmd = f'export MANPATH={Path.home() / "share" / "man"}:$MANPATH'
-        print(f"Local installation complete. Add to ~/.bashrc to update MANPATH:")
+        print("Local installation complete. Add to ~/.bashrc to update MANPATH:")
         print(f"    {manpath_cmd}")
 
     return 0

pygeai/cli/parsers.py

@@ -1,7 +1,7 @@
 from typing import List, Tuple, Optional
 
 from pygeai import logger
-from pygeai.cli.commands.base import base_options
+
 from pygeai.cli.commands import Command, Option
 from pygeai.core.common.exceptions import UnknownArgumentError, MissingRequirementException
 

pygeai/core/base/clients.py

@@ -1,50 +1,120 @@
 from abc import ABC
+from pygeai import logger
 
 from pygeai.core.base.session import get_session, Session
-from pygeai.core.common.exceptions import MissingRequirementException
-from pygeai.core.services.rest import ApiService
-from pygeai.core.utils.validators import validate_status_code
+from pygeai.core.common.exceptions import MissingRequirementException, MixedAuthenticationException
+from pygeai.core.services.rest import GEAIApiService
 
 
 class BaseClient(ABC):
+    _logged_session_config = None
 
-    def __init__(self, api_key: str = None, base_url: str = None, alias: str = None, *,
-                 access_token: str = None, project_id: str = None):
+    def __init__(
+        self, 
+        api_key: str = None, 
+        base_url: str = None, 
+        alias: str = None, 
+        *,
+        access_token: str = None, 
+        project_id: str = None,
+        organization_id: str = None,
+        allow_mixed_auth: bool = True
+    ):
         """
-        If commont settings are not specified, they're retrieved from default Session, based on the
-        credential files.
+        Initialize a client with authentication credentials.
+        
+        Authentication Options:
+        1. API Key: Provide api_key and base_url
+        2. OAuth 2.0: Provide access_token, project_id, and base_url
+        3. From credentials: Provide alias to load from config
+        
         :param api_key: GEAI API KEY to access services
         :param base_url: URL for GEAI instance to be used
         :param alias: Alias to use from credentials file
-        :param access_token: OAuth access token (keyword-only)
+        :param access_token: OAuth 2.0 access token (keyword-only)
         :param project_id: Project ID for OAuth authentication (keyword-only)
+        :param organization_id: Organization ID for OAuth authentication (keyword-only)
+        :param allow_mixed_auth: Allow both api_key and access_token (default: False)
+        :raises: MissingRequirementException - If authentication configuration is incomplete
+        :raises: ValueError - If authentication configuration is invalid
         """
-        if access_token and not project_id:
-            raise MissingRequirementException("project_id is required when using access_token")
-
-        if not (api_key and base_url) and not (access_token and base_url) and alias:
+        
+        # Case 1: Use credentials from alias
+        if not (api_key or access_token) and alias:
             self.__session = get_session(alias)
             if not self.__session:
-                raise MissingRequirementException("API KEY and BASE URL must be defined in order to use this functionality")
+                raise MissingRequirementException(
+                    "API KEY and BASE URL must be defined in order to use this functionality"
+                )
+        
+        # Case 2: Direct credential provision
         elif (api_key or access_token) and base_url:
+            # Validate auth parameters before mutating singleton session
+            if access_token and api_key and not allow_mixed_auth:
+                raise MixedAuthenticationException(
+                    "Cannot specify both 'api_key' and 'access_token'. "
+                    "Use 'allow_mixed_auth=True' to allow both (OAuth will take precedence)."
+                )
+            
+            if access_token and not project_id:
+                raise MissingRequirementException(
+                    "project_id is required when using access_token for OAuth authentication"
+                )
+            
+            # Get singleton session and update its properties
+            # Note: Setters automatically update auth_type via _determine_auth_type()
             self.__session = get_session()
             self.__session.api_key = api_key
             self.__session.access_token = access_token
             self.__session.project_id = project_id
+            self.__session.organization_id = organization_id
             self.__session.base_url = base_url
+        
+        # Case 3: Use default session
         else:
             self.__session = get_session()
-
+        
         if self.session is None:
-            raise MissingRequirementException("Cannot access this functionality without setting API_KEY and BASE_URL")
-
-        token = self.session.access_token if self.session.access_token else self.session.api_key
-        self.__api_service = ApiService(
+            raise MissingRequirementException(
+                "Cannot access this functionality without setting API_KEY and BASE_URL"
+            )
+        
+        # Log session config only once per unique configuration
+        session_config_key = (
+            self.session.alias,
+            self.session.base_url,
+            self.session.auth_type
+        )
+        
+        if BaseClient._logged_session_config != session_config_key:
+            self._log_authentication_info()
+            BaseClient._logged_session_config = session_config_key
+        
+        # Initialize API service with active token
+        token = self.session.get_active_token()
+        
+        self.__api_service = GEAIApiService(
             base_url=self.session.base_url,
             token=token,
-            project_id=self.session.project_id
+            project_id=self.session.project_id,
+            organization_id=self.session.organization_id
         )
 
+    def _log_authentication_info(self):
+        """Log authentication configuration for debugging."""
+        if self.session.is_oauth():
+            logger.info("Using OAuth 2.0 authentication")
+            logger.info(f"Project ID: {self.session.project_id}")
+            if self.session.organization_id:
+                logger.info(f"Organization ID: {self.session.organization_id}")
+        elif self.session.is_api_key():
+            logger.info("Using API Key authentication")
+        else:
+            logger.warning("No authentication method configured")
+        
+        logger.info(f"Base URL: {self.session.base_url}")
+        logger.info(f"Alias: {self.session.alias}")
+
     @property
     def session(self):
         return self.__session
@@ -58,6 +128,5 @@ class BaseClient(ABC):
         return self.__api_service
 
     @api_service.setter
-    def api_service(self, api_service: ApiService):
+    def api_service(self, api_service: GEAIApiService):
         self.__api_service = api_service
-

pygeai/core/base/mappers.py

@@ -56,14 +56,10 @@ class ErrorMapper:
         :param data: dict - The dictionary containing error response data.
         :return: list[Error] - A list of `Error` objects.
         """
-        error_list = list()
         errors = data.get('errors')
         if errors is not None and any(errors):
-            for error_data in errors:
-                error = cls.map_to_error(error_data)
-                error_list.append(error)
-
-        return error_list
+            return [cls.map_to_error(error_data) for error_data in errors]
+        return []
 
 
 class ResponseMapper:
@@ -130,14 +126,10 @@ class ModelMapper:
         :param data: dict - The dictionary containing the list of intents.
         :return: list[AssistantIntent] - A list of mapped `AssistantIntent` objects.
         """
-        intent_list = list()
         intents = data.get('intents')
         if intents is not None and any(intents):
-            for intent_data in intents:
-                intent = cls.map_to_intent(intent_data)
-                intent_list.append(intent)
-
-        return intent_list
+            return [cls.map_to_intent(intent_data) for intent_data in intents]
+        return []
 
     @classmethod
     def map_to_intent(cls, data: dict) -> AssistantIntent:
@@ -165,14 +157,13 @@ class ModelMapper:
         :param data: dict - The dictionary containing the list of revisions.
         :return: list[AssistantRevision] - A list of mapped `AssistantRevision` objects.
         """
-        revision_list = list()
-        revisions = data.get("revisions")
+        revisions = data.get('revisions')
+
         if revisions is not None and any(revisions):
-            for revision_data in revisions:
-                revision = cls.map_to_revision(revision_data)
-                revision_list.append(revision)
 
-        return revision_list
+            return [cls.map_to_revision(revision_data) for revision_data in revisions]
+
+        return []
 
     @classmethod
     def map_to_revision(cls, data: dict) -> AssistantRevision:
@@ -204,14 +195,13 @@ class ModelMapper:
        :param data: dict - The dictionary containing metadata information.
        :return: list[AssistantRevisionMetadata] - A list of mapped `AssistantRevisionMetadata` objects.
        """
-        metadata_list = list()
         metadata = data.get('metadata')
+
         if metadata is not None and any(metadata):
-            for metadata_data in metadata:
-                metadata = cls.map_to_metadata(metadata_data)
-                metadata_list.append(metadata)
 
-        return metadata_list
+            return [cls.map_to_metadata(metadata_data) for metadata_data in metadata]
+
+        return []
 
     @classmethod
     def map_to_metadata(cls, data: dict) -> AssistantRevisionMetadata:
@@ -254,14 +244,13 @@ class ModelMapper:
 
     @classmethod
     def map_to_search_profile_list(cls, data: dict) -> list[SearchProfile]:
-        search_profile_list = list()
         search_profiles = data.get('searchProfiles')
+
         if search_profiles is not None and any(search_profiles):
-            for search_profile_data in search_profiles:
-                search_profile = cls.map_to_search_profile(search_profile_data)
-                search_profile_list.append(search_profile)
 
-        return search_profile_list
+            return [cls.map_to_search_profile(search_profile_data) for search_profile_data in search_profiles]
+
+        return []
 
     @classmethod
     def map_to_search_profile(cls, data: dict) -> SearchProfile:
@@ -278,14 +267,13 @@ class ModelMapper:
 
     @classmethod
     def map_to_token_list(cls, data: dict) -> list[ProjectToken]:
-        token_list = list()
         tokens = data.get('tokens')
+
         if tokens is not None and any(tokens):
-            for token_data in tokens:
-                token = cls.map_to_token(token_data)
-                token_list.append(token)
 
-        return token_list
+            return [cls.map_to_token(token_data) for token_data in tokens]
+
+        return []
 
     @classmethod
     def map_to_token(cls, data: dict) -> ProjectToken:
@@ -305,14 +293,13 @@ class ModelMapper:
 
     @classmethod
     def map_to_usage_limit_list(cls, data: dict) -> list[UsageLimit]:
-        usage_limit_list = list()
         usage_limits = data.get('usageLimits')
+
         if usage_limits is not None and any(usage_limits):
-            for usage_limit_data in usage_limits:
-                usage_limit = cls.map_to_usage_limit(usage_limit_data)
-                usage_limit_list.append(usage_limit)
 
-        return usage_limit_list
+            return [cls.map_to_usage_limit(usage_limit_data) for usage_limit_data in usage_limits]
+
+        return []
 
     @classmethod
     def map_to_usage_limit(cls, data: dict) -> UsageLimit:
@@ -339,14 +326,13 @@ class ModelMapper:
 
     @classmethod
     def map_to_item_list(cls, data: dict) -> list[RequestItem]:
-        item_list = list()
         items = data.get('items')
+
         if items is not None and any(items):
-            for item_data in items:
-                item = cls.map_to_item(item_data)
-                item_list.append(item)
 
-        return item_list
+            return [cls.map_to_item(item_data) for item_data in items]
+
+        return []
 
     @classmethod
     def map_to_item(cls, data: dict) -> RequestItem:
@@ -427,14 +413,13 @@ class ModelMapper:
 
     @classmethod
     def map_to_feature_list(cls, data: dict) -> list[WelcomeDataFeature]:
-        feature_list = list()
-        features = data.get("features")
+        features = data.get('features')
+
         if features is not None and any(features):
-            for feature_data in features:
-                feature = cls.map_to_feature(feature_data)
-                feature_list.append(feature)
 
-        return feature_list
+            return [cls.map_to_feature(feature_data) for feature_data in features]
+
+        return []
 
     @classmethod
     def map_to_feature(cls, data: dict) -> WelcomeDataFeature:
@@ -445,14 +430,13 @@ class ModelMapper:
 
     @classmethod
     def map_to_example_prompt_list(cls, data: dict) -> list[WelcomeDataExamplePrompt]:
-        example_prompt_list = list()
-        examples_prompt = data.get("examplesPrompt")
+        examples_prompt = data.get('examplesPrompt')
+
         if examples_prompt is not None and any(examples_prompt):
-            for example_prompt_data in examples_prompt:
-                example_prompt = cls.map_to_example_prompt(example_prompt_data)
-                example_prompt_list.append(example_prompt)
 
-        return example_prompt_list
+            return [cls.map_to_example_prompt(example_prompt_data) for example_prompt_data in examples_prompt]
+
+        return []
 
     @classmethod
     def map_to_example_prompt(cls, data: dict) -> WelcomeDataExamplePrompt: