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

pygeai/core/base/session.py

@@ -1,13 +1,13 @@
 import sys
+import warnings
 from typing import Optional
 
 from pygeai import logger
 from pygeai.core.common.config import get_settings
-from pygeai.core.common.exceptions import MissingRequirementException
+from pygeai.core.common.constants import AuthType
+from pygeai.core.common.exceptions import MissingRequirementException, MixedAuthenticationException
 from pygeai.core.singleton import Singleton
 
-settings = get_settings()
-
 
 _session = None
 
@@ -15,14 +15,21 @@ _session = None
 class Session(metaclass=Singleton):
     """
     A session to store configuration state required to interact with different resources.
-
-    :param api_key: str - API key to interact with GEAI
+    
+    Authentication Methods:
+    - API Key: Use api_key parameter
+    - OAuth 2.0: Use access_token and project_id parameters
+    
+    :param api_key: str - API key to interact with GEAI (mutually exclusive with access_token unless allow_mixed_auth=True)
     :param base_url: str - Base URL of the GEAI instance
     :param eval_url: Optional[str] - Optional evaluation endpoint URL
-    :param access_token: Optional[str] - OAuth access token (keyword-only)
-    :param project_id: Optional[str] - Project ID for OAuth authentication (keyword-only)
+    :param access_token: Optional[str] - OAuth 2.0 access token (keyword-only, requires project_id)
+    :param project_id: Optional[str] - Project ID for OAuth authentication (keyword-only, requires access_token)
+    :param organization_id: Optional[str] - Organization ID for OAuth authentication (keyword-only)
+    :param alias: Optional[str] - Alias name for this session
+    :param allow_mixed_auth: bool - If True, allow both api_key and access_token (access_token takes precedence). Default False.
     :return: Session - Instance of the Session class
-    :raises: ValueError - If required parameters are missing or invalid
+    :raises: ValueError - If authentication configuration is invalid
     """
 
     def __init__(
@@ -33,10 +40,34 @@ class Session(metaclass=Singleton):
             *,
             access_token: Optional[str] = None,
             project_id: Optional[str] = None,
+            organization_id: Optional[str] = None,
             alias: Optional[str] = None,
+            allow_mixed_auth: bool = True,
     ):
+        # Validate authentication configuration
+        if api_key and access_token and not allow_mixed_auth:
+            raise MixedAuthenticationException(
+                "Cannot specify both 'api_key' and 'access_token'. "
+                "Use 'api_key' for API Key authentication or 'access_token' with 'project_id' for OAuth 2.0. "
+                "Set allow_mixed_auth=True to allow both (access_token will take precedence)."
+            )
+        
+        if access_token and not project_id:
+            raise MissingRequirementException(
+                "OAuth 2.0 authentication requires 'project_id'. "
+                "Provide project_id when using access_token."
+            )
+        
+        if project_id and not access_token:
+            warnings.warn(
+                "project_id provided without access_token. "
+                "project_id is only used with OAuth 2.0 authentication.",
+                UserWarning
+            )
+        
         if not api_key and not access_token:
-            logger.warning("Cannot instantiate session without api_key or access_token")
+            logger.warning("No authentication method configured. API calls may fail.")
+        
         if not base_url:
             logger.warning("Cannot instantiate session without base_url")
 
@@ -45,18 +76,36 @@ class Session(metaclass=Singleton):
         self.__eval_url = eval_url
         self.__access_token = access_token
         self.__project_id = project_id
+        self.__organization_id = organization_id
         self.__alias = alias if alias else "default"
+        self.__auth_type = self._determine_auth_type()
 
         global _session
         _session = self
 
+    def _determine_auth_type(self) -> AuthType:
+        """Determine the active authentication type based on configuration."""
+        if self.__access_token and self.__project_id:
+            return AuthType.OAUTH_TOKEN
+        elif self.__api_key:
+            return AuthType.API_KEY
+        else:
+            return AuthType.NONE
+
+    @property
+    def auth_type(self) -> AuthType:
+        """Get the current authentication type."""
+        return self.__auth_type
+
     @property
     def api_key(self):
         return self.__api_key
 
     @api_key.setter
     def api_key(self, api_key: str):
+        """Set API key and update auth type."""
         self.__api_key = api_key
+        self.__auth_type = self._determine_auth_type()
 
     @property
     def base_url(self):
@@ -80,7 +129,9 @@ class Session(metaclass=Singleton):
 
     @access_token.setter
     def access_token(self, access_token: str):
+        """Set access token and update auth type."""
         self.__access_token = access_token
+        self.__auth_type = self._determine_auth_type()
 
     @property
     def project_id(self):
@@ -88,7 +139,17 @@ class Session(metaclass=Singleton):
 
     @project_id.setter
     def project_id(self, project_id: str):
+        """Set project ID and update auth type."""
         self.__project_id = project_id
+        self.__auth_type = self._determine_auth_type()
+
+    @property
+    def organization_id(self):
+        return self.__organization_id
+
+    @organization_id.setter
+    def organization_id(self, organization_id: str):
+        self.__organization_id = organization_id
 
     @property
     def alias(self):
@@ -98,14 +159,33 @@ class Session(metaclass=Singleton):
     def alias(self, alias: str):
         self.__alias = alias
 
+    def is_oauth(self) -> bool:
+        """Check if session is using OAuth authentication."""
+        return self.__auth_type == AuthType.OAUTH_TOKEN
+
+    def is_api_key(self) -> bool:
+        """Check if session is using API key authentication."""
+        return self.__auth_type == AuthType.API_KEY
+
+    def get_active_token(self) -> Optional[str]:
+        """Get the active authentication token based on auth type."""
+        if self.is_oauth():
+            return self.__access_token
+        elif self.is_api_key():
+            return self.__api_key
+        return None
+
 
 def get_session(alias: str = None) -> Session:
     """
     Session is a singleton object:
     On the first invocation, returns Session configured with the API KEY and BASE URL corresponding to the
     alias provided. On the following invocations, it returns the first object instantiated.
+    
+    Loads both API Key and OAuth 2.0 credentials from configuration if available.
     """
     try:
+        settings = get_settings()
         global _session
         if _session is None:
             if not alias:
@@ -113,27 +193,43 @@ def get_session(alias: str = None) -> Session:
             
             _validate_alias(alias, allow_missing_default=True)
 
+            api_key = settings.get_api_key(alias)
+            access_token = settings.get_access_token(alias)
+            
+            # Allow mixed auth for backward compatibility
+            # If both are configured, allow mixed (with warning logged in Session)
+            allow_mixed = bool(api_key and access_token)
+            
+            if allow_mixed:
+                logger.warning(
+                    f"Both API key and OAuth token configured for alias '{alias}'. "
+                    "OAuth token will take precedence. Consider using separate aliases."
+                )
+
             _session = Session(
-                api_key=settings.get_api_key(alias),
+                api_key=api_key,
                 base_url=settings.get_base_url(alias),
                 eval_url=settings.get_eval_url(alias),
-                access_token=settings.get_access_token(alias),
+                access_token=access_token,
                 project_id=settings.get_project_id(alias),
+                organization_id=settings.get_organization_id(alias),
                 alias=alias,
+                allow_mixed_auth=allow_mixed,
             )
+                
         elif _session is not None and alias:
             _validate_alias(alias, allow_missing_default=False)
             
+            api_key = settings.get_api_key(alias)
+            access_token = settings.get_access_token(alias)
+            
             _session.alias = alias
-            _session.api_key = settings.get_api_key(alias)
+            _session.api_key = api_key
             _session.base_url = settings.get_base_url(alias)
             _session.eval_url = settings.get_eval_url(alias)
-            _session.access_token = settings.get_access_token(alias)
+            _session.access_token = access_token
             _session.project_id = settings.get_project_id(alias)
-
-        if alias:
-            logger.debug(f"Alias: {alias}")
-            logger.debug(f"Base URL: {_session.base_url}")
+            _session.organization_id = settings.get_organization_id(alias)
 
         return _session
     except ValueError as e:
@@ -142,10 +238,25 @@ def get_session(alias: str = None) -> Session:
 
 
 def _validate_alias(alias: str, allow_missing_default: bool = False):
+    settings = get_settings()
     available_aliases = settings.list_aliases()
     if alias not in available_aliases:
         if allow_missing_default and alias == "default":
             return
         raise MissingRequirementException(
             f"The profile '{alias}' doesn't exist. Use 'geai configure --list' to see available profiles."
-        )
+        )
+
+
+def reset_session():
+    """
+    Reset the session instance. Useful for testing.
+    
+    This clears both the module-level _session variable and the
+    Singleton metaclass cache to ensure proper test isolation.
+    """
+    global _session
+    _session = None
+    
+    from pygeai.core.singleton import Singleton
+    Singleton.reset_instance(Session)

pygeai/core/common/config.py

@@ -1,6 +1,5 @@
 import os
 import sys
-from functools import lru_cache
 from pathlib import Path
 import configparser
 
@@ -21,9 +20,13 @@ class SettingsManager:
     GEAI_SETTINGS_DIR = str(SETTINGS_DIR)
     GEAI_CREDS_FILE = SETTINGS_DIR / "credentials"
 
-    def __init__(self):
+    def __init__(self, credentials_file: str = None):
         self.config = configparser.ConfigParser()
-
+        
+        if credentials_file:
+            self.GEAI_CREDS_FILE = Path(credentials_file)
+            logger.debug(f"Using custom credentials file: {self.GEAI_CREDS_FILE}")
+        
         if self.GEAI_CREDS_FILE.exists():
             self.config.read(self.GEAI_CREDS_FILE)
         else:
@@ -41,14 +44,19 @@ class SettingsManager:
             return
         
         if setting_key not in self.config[alias]:
-            sys.stdout.write(f"'{setting_key}' not found in alias '{alias}' in the credentials file. Adding empty value.\n")
-            # SET ADDITIONAL VARS
-            if setting_key.lower() == "GEAI_API_EVAL_URL".lower():
+            logger.debug(f"'{setting_key}' not found in alias '{alias}' in the credentials file.\n")
+
+            if setting_key.lower() == "geai_api_eval_url" or ("geai_oauth_access_token" in self.config[alias] and setting_key.lower() in ["geai_project_id", "geai_organization_id"]):
+                sys.stdout.write(f"'{setting_key}' not found in alias '{alias}' in the credentials file. Adding empty value\n")
+
+            if setting_key.lower() == "geai_api_eval_url":
                 self.set_eval_url("", alias)
-            if setting_key.lower() == "GEAI_OAUTH_ACCESS_TOKEN".lower():
-                self.set_access_token("", alias)
-            if setting_key.lower() == "GEAI_PROJECT_ID".lower():
-                self.set_project_id("", alias)
+
+            if "geai_oauth_access_token" in self.config[alias]:
+                if setting_key.lower() == "geai_project_id":
+                    self.set_project_id("", alias)
+                if setting_key.lower() == "geai_organization_id":
+                    self.set_organization_id("", alias)
 
         return self.config[alias].get(setting_key, "")
 
@@ -112,6 +120,16 @@ class SettingsManager:
     def set_eval_url(self, eval_url, alias: str = "default"):
         self.set_setting_value("GEAI_API_EVAL_URL", eval_url, alias)
 
+    def get_organization_id(self, alias: str = "default"):
+        organization_id = os.environ.get("GEAI_ORGANIZATION_ID") if not alias or alias == "default" else None
+        if not organization_id:
+            organization_id = self.get_setting_value("GEAI_ORGANIZATION_ID", alias)
+
+        return organization_id
+
+    def set_organization_id(self, organization_id, alias: str = "default"):
+        self.set_setting_value("GEAI_ORGANIZATION_ID", organization_id, alias)
+
     def list_aliases(self):
         """Returns a dict of all aliases and their base URLs."""
         return {
@@ -130,9 +148,28 @@ class SettingsManager:
             logger.warning(f"Alias '{alias}' not found in the credentials file.")
 
 
-@lru_cache()
-def get_settings():
-    return SettingsManager()
+_settings_instance = None
+
+
+def get_settings(credentials_file: str = None):
+    """
+    Get or create the SettingsManager instance.
+    
+    :param credentials_file: Optional path to a custom credentials file.
+                            If provided on first call, uses that file.
+                            Subsequent calls ignore this parameter.
+    :return: SettingsManager instance
+    """
+    global _settings_instance
+    if _settings_instance is None:
+        _settings_instance = SettingsManager(credentials_file=credentials_file)
+    return _settings_instance
+
+
+def reset_settings():
+    """Reset the settings instance. Useful for testing."""
+    global _settings_instance
+    _settings_instance = None
 
 
 if __name__ == "__main__":

pygeai/core/common/constants.py

@@ -0,0 +1,8 @@
+from enum import Enum
+
+
+class AuthType(Enum):
+    """Authentication method types for PyGEAI SDK."""
+    API_KEY = "api_key"
+    OAUTH_TOKEN = "oauth_token"
+    NONE = "none"

pygeai/core/common/exceptions.py

@@ -122,3 +122,9 @@ class APIResponseError(GEAIException):
     """Raised when there is an error in the API response"""
     pass
 
+
+class MixedAuthenticationException(GEAIException):
+    """Raised when both API token and Oauth2 authentication are setup for the same profile"""
+    pass
+
+

pygeai/core/embeddings/clients.py

@@ -1,7 +1,6 @@
 
 from pygeai import logger
 from pygeai.core.base.clients import BaseClient
-from pygeai.core.common.exceptions import InvalidAPIResponseException, APIResponseError
 from pygeai.core.embeddings.endpoints import GENERATE_EMBEDDINGS
 from pygeai.core.utils.validators import validate_status_code
 from pygeai.core.utils.parsers import parse_json_response

pygeai/core/embeddings/managers.py

@@ -1,5 +1,4 @@
 from pygeai import logger
-from pygeai.core.base.mappers import ErrorMapper
 from pygeai.core.common.exceptions import APIError
 from pygeai.core.embeddings.clients import EmbeddingsClient
 from pygeai.core.embeddings.mappers import EmbeddingsResponseMapper

pygeai/core/feedback/clients.py

@@ -1,8 +1,6 @@
-from typing import Any, Union
 
 from pygeai import logger
 from pygeai.core.base.clients import BaseClient
-from pygeai.core.common.exceptions import InvalidAPIResponseException
 from pygeai.core.feedback.endpoints import SEND_FEEDBACK_V1
 from pygeai.core.utils.validators import validate_status_code
 from pygeai.core.utils.parsers import parse_json_response

pygeai/core/feedback/models.py

@@ -1,4 +1,4 @@
-from pydantic import BaseModel, Field
+from pydantic import Field
 
 from pygeai.core import CustomBaseModel
 

pygeai/core/files/clients.py

@@ -1,10 +1,7 @@
-import json
-from json import JSONDecodeError
 from pathlib import Path
 
 from pygeai import logger
 from pygeai.core.base.clients import BaseClient
-from pygeai.core.common.exceptions import InvalidAPIResponseException
 from pygeai.core.files.endpoints import UPLOAD_FILE_V1, GET_FILE_V1, DELETE_FILE_V1, GET_FILE_CONTENT_V1, \
     GET_ALL_FILES_V1
 from pygeai.core.utils.validators import validate_status_code

pygeai/core/files/managers.py

@@ -4,7 +4,7 @@ from typing import Optional
 
 from pygeai import logger
 from pygeai.admin.clients import AdminClient
-from pygeai.core.base.mappers import ErrorMapper, ResponseMapper
+from pygeai.core.base.mappers import ResponseMapper
 from pygeai.core.base.responses import EmptyResponse
 from pygeai.core.files.clients import FileClient
 from pygeai.core.files.models import UploadFile, File, FileList

pygeai/core/files/mappers.py

@@ -22,14 +22,13 @@ class FileResponseMapper:
 
     @classmethod
     def map_to_file_list(cls, data: dict) -> list[File]:
-        file_list = list()
         files = data.get('dataFiles')
+
         if files is not None and any(files):
-            for file_data in files:
-                file = cls.map_to_file(file_data)
-                file_list.append(file)
 
-        return file_list
+            return [cls.map_to_file(file_data) for file_data in files]
+
+        return []
 
     @classmethod
     def map_to_file(cls, data: dict) -> File:

pygeai/core/llm/clients.py

@@ -1,7 +1,6 @@
 
 from pygeai import logger
 from pygeai.core.base.clients import BaseClient
-from pygeai.core.common.exceptions import InvalidAPIResponseException
 from pygeai.core.llm.endpoints import GET_PROVIDER_LIST_V2, GET_PROVIDER_DATA_V2, GET_PROVIDER_MODELS_V2, \
     GET_MODEL_DATA_V2
 from pygeai.core.utils.validators import validate_status_code

pygeai/core/models.py

@@ -185,8 +185,8 @@ class UsageLimit(CustomBaseModel):
     "subscriptionType": "string",         // Subscription type (Freemium, Daily, Weekly, Monthly)
     "usageUnit": "string",                // Usage unit (Requests, Cost)
     "usedAmount": "number",               // Amount used (decimal or scientific notation)
-    "validFrom": "timestamp",             // Start date of the usage limit
-    "validUntil": "timestamp"             // Expiration or renewal date
+    "validFrom": "str",             // Start date of the usage limit
+    "validUntil": "str"             // Expiration or renewal date
     """
     hard_limit: Optional[float] = Field(None, alias="hardLimit")
     id: Optional[str] = Field(None, alias="id")
@@ -198,8 +198,8 @@ class UsageLimit(CustomBaseModel):
     subscription_type: Optional[Literal["Freemium", "Daily", "Weekly", "Monthly"]] = Field(None, alias="subscriptionType")
     usage_unit: Optional[Literal["Requests", "Cost"]] = Field(None, alias="usageUnit")
     used_amount: Optional[float] = Field(None, alias="usedAmount")
-    valid_from: Optional[datetime] = Field(None, alias="validFrom")
-    valid_until: Optional[datetime] = Field(None, alias="validUntil")
+    valid_from: Optional[str] = Field(None, alias="validFrom")
+    valid_until: Optional[str] = Field(None, alias="validUntil")
 
     def to_dict(self):
         return self.model_dump(by_alias=True, exclude_none=True)

pygeai/core/plugins/clients.py

@@ -1,8 +1,5 @@
-from typing import Optional, List, Dict
 
-from pygeai import logger
 from pygeai.core.base.clients import BaseClient
-from pygeai.core.common.exceptions import InvalidAPIResponseException
 from pygeai.core.plugins.endpoints import LIST_ASSISTANTS_PLUGINS_V1
 from pygeai.core.utils.validators import validate_status_code
 from pygeai.core.utils.parsers import parse_json_response

pygeai/core/plugins/models.py

@@ -1,5 +1,5 @@
-from pydantic import BaseModel, Field
-from typing import Optional, List, Literal
+from pydantic import Field
+from typing import Optional, List
 from uuid import UUID
 
 from pygeai.core import CustomBaseModel

pygeai/core/rerank/clients.py

@@ -1,8 +1,6 @@
-from typing import Any, Union
 
 from pygeai import logger
 from pygeai.core.base.clients import BaseClient
-from pygeai.core.common.exceptions import InvalidAPIResponseException
 from pygeai.core.rerank.endpoints import RERANK_V1
 from pygeai.core.utils.validators import validate_status_code
 from pygeai.core.utils.parsers import parse_json_response

pygeai/core/secrets/clients.py

@@ -1,8 +1,6 @@
 from typing import Optional, List, Dict
 
-from pygeai import logger
 from pygeai.core.base.clients import BaseClient
-from pygeai.core.common.exceptions import InvalidAPIResponseException
 from pygeai.core.secrets.endpoints import LIST_SECRETS_V1, GET_SECRET_V1, CREATE_SECRET_V1, UPDATE_SECRET_V1, \
     SET_SECRET_ACCESSES_V1, GET_SECRET_ACCESSES_V1
 from pygeai.core.utils.validators import validate_status_code

pygeai/core/services/rest.py

@@ -1,5 +1,5 @@
 import json
-from typing import Optional, List, Dict
+from typing import Optional, Dict
 import requests as req
 from pygeai import logger
 from pygeai.core.common.exceptions import InvalidResponseException
@@ -21,16 +21,13 @@ class ApiService:
     :param username: str - Username for basic authentication (optional).
     :param password: str - Password for basic authentication (optional).
     :param token: str - Bearer token for authentication (optional).
-    :param project_id: str - Project ID for OAuth authentication (optional, keyword-only).
     """
 
-    def __init__(self, base_url, username: str = None, password: str = None, token: str = None, *,
-                 project_id: str = None):
+    def __init__(self, base_url, username: str = None, password: str = None, token: str = None):
         self._base_url = base_url
         self._username = username
         self._password = password
         self._token = token
-        self._project_id = project_id
 
     @property
     def base_url(self):
@@ -52,14 +49,6 @@ class ApiService:
     def token(self, token: str):
         self._token = token
 
-    @property
-    def project_id(self):
-        return self._project_id
-
-    @project_id.setter
-    def project_id(self, project_id: str):
-        self._project_id = project_id
-
     def get(self, endpoint: str, params: dict = None, headers: dict = None, verify: bool = True):
         """
         Sends a GET request to the specified API endpoint.
@@ -418,7 +407,84 @@ class ApiService:
             if "Authorization" not in headers:
                 headers["Authorization"] = f"Bearer {self.token}"
 
-        if self.project_id and "ProjectId" not in headers:
+        return headers
+
+
+class GEAIApiService(ApiService):
+    """
+    Service for interacting with REST APIs in Globant Enterprise AI.
+
+    :param base_url: str - The base URL of the API.
+    :param username: str - Username for basic authentication (optional).
+    :param password: str - Password for basic authentication (optional).
+    :param token: str - Bearer token for authentication (optional).
+    :param project_id: str - Project ID for OAuth authentication (optional, keyword-only).
+    :param organization_id: str - Organization ID for OAuth authentication (optional, keyword-only).
+    """
+
+    def __init__(self, base_url, username: str = None, password: str = None, token: str = None, *,
+                 project_id: str = None, organization_id: str = None):
+        super().__init__(base_url, username, password, token)
+        self._project_id = project_id
+        self._organization_id = organization_id
+
+    @property
+    def project_id(self):
+        return self._project_id
+
+    @project_id.setter
+    def project_id(self, project_id: str):
+        self._project_id = project_id
+
+    @property
+    def organization_id(self):
+        return self._organization_id
+
+    @organization_id.setter
+    def organization_id(self, organization_id: str):
+        self._organization_id = organization_id
+
+    def get(self, endpoint: str, params: dict = None, headers: dict = None, verify: bool = True):
+        headers = self._add_oauth_context_to_headers(headers=headers)
+        return super().get(endpoint, params, headers, verify)
+
+    def post(self, endpoint: str, data: dict, headers: dict = None, verify: bool = True, form: bool = False):
+        headers = self._add_oauth_context_to_headers(headers=headers)
+        return super().post(endpoint, data, headers, verify, form)
+
+    def stream_post(self, endpoint: str, data: dict, headers: dict = None, verify: bool = True, form: bool = False):
+        headers = self._add_oauth_context_to_headers(headers=headers)
+        return super().stream_post(endpoint, data, headers, verify, form)
+
+    def post_file_binary( self, endpoint: str, headers: dict = None, verify: bool = True, file=None):
+        headers = self._add_oauth_context_to_headers(headers=headers)
+        return super().post_file_binary(endpoint, headers, verify, file)
+
+    def post_files_multipart(self, endpoint: str, data: Optional[dict] = None, headers: Optional[dict] = None, verify: bool = True, files: Optional[Dict[str, str]] = None):
+        headers = self._add_oauth_context_to_headers(headers=headers)
+        return super().post_files_multipart(endpoint, data, headers, verify, files)
+
+    def put(self, endpoint: str, data: dict, headers: dict = None, verify: bool = True):
+        headers = self._add_oauth_context_to_headers(headers=headers)
+        return super().put(endpoint, data, headers, verify)
+
+    def delete(self, endpoint: str, headers: dict = None, data: dict = None, verify: bool = True):
+        headers = self._add_oauth_context_to_headers(headers=headers)
+        return super().delete(endpoint, headers, data, verify)
+
+    def _add_oauth_context_to_headers(self, headers: dict = None):
+        headers = headers if headers is not None else {}
+
+        # Add OAuth context headers if available
+        if self.project_id and "ProjectId" not in headers and "project-id" not in headers:
             headers["ProjectId"] = self.project_id
+            logger.debug(f"Added Authorization header with ProjectId: {self.project_id}")
+        
+        if self.organization_id and "OrganizationId" not in headers and "organization-id" not in headers:
+            headers["OrganizationId"] = self.organization_id
+            logger.debug(f"Added OrganizationId header: {self.organization_id}")
+        
+        if not self.project_id and not self.organization_id:
+            logger.debug("Added Authorization header")
 
         return headers