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

pygeai/migration/tools.py

@@ -1,4 +1,4 @@
-from typing import List, Dict, Optional
+from typing import List, Dict
 from dataclasses import dataclass, field
 
 from pygeai import logger

pygeai/organization/clients.py

@@ -1,11 +1,11 @@
-
 from pygeai import logger
-from pygeai.core.common.exceptions import InvalidAPIResponseException
+from pygeai.admin.clients import AdminClient
+from pygeai.core.common.exceptions import APIError
 from pygeai.organization.endpoints import GET_ASSISTANT_LIST_V1, GET_PROJECT_LIST_V1, GET_PROJECT_V1, \
     CREATE_PROJECT_V1, UPDATE_PROJECT_V1, DELETE_PROJECT_V1, GET_PROJECT_TOKENS_V1, GET_REQUEST_DATA_V1, \
     GET_MEMBERSHIPS_V2, GET_PROJECT_MEMBERSHIPS_V2, GET_PROJECT_ROLES_V2, GET_PROJECT_MEMBERS_V2, \
-    GET_ORGANIZATION_MEMBERS_V2, ADD_PROJECT_MEMBER_V2, CREATE_ORGANIZATION_V2, GET_ORGANIZATION_LIST_V2, \
-    DELETE_ORGANIZATION_V2
+    GET_ORGANIZATION_MEMBERS_V2, GET_PLUGIN_RUNTIME_POLICIES_V2, ADD_PROJECT_MEMBER_V2, CREATE_ORGANIZATION_V2, \
+    GET_ORGANIZATION_LIST_V2, DELETE_ORGANIZATION_V2
 from pygeai.core.base.clients import BaseClient
 from pygeai.core.utils.validators import validate_status_code
 from pygeai.core.utils.parsers import parse_json_response
@@ -13,6 +13,27 @@ from pygeai.core.utils.parsers import parse_json_response
 
 class OrganizationClient(BaseClient):
 
+    def _get_token_information(self):
+        return AdminClient(self.api_service.token, base_url=self.api_service.base_url).validate_api_token()
+
+    def _get_project_id(self):
+        response = None
+        try:
+            response = self._get_token_information()
+            return response.get("projectId")
+        except Exception as e:
+            logger.error(f"Error retrieving project_id from GEAI. Response: {response}: {e}")
+            raise APIError(f"Error retrieving project_id from GEAI: {e}")
+
+    def _get_organization_id(self):
+        response = None
+        try:
+            response = self._get_token_information()
+            return response.get("organizationId")
+        except Exception as e:
+            logger.error(f"Error retrieving organization_id from GEAI. Response: {response}: {e}")
+            raise APIError(f"Error retrieving organization_id from GEAI: {e}")
+
     def get_assistant_list(
             self,
             detail: str = "summary"
@@ -261,58 +282,143 @@ class OrganizationClient(BaseClient):
             "orderDirection": order_direction
         }
         if email:
-            params["email"] = email
+            params["userEmail"] = email
         if order_key:
             params["orderKey"] = order_key
         if role_types:
             params["roleTypes"] = role_types
 
-        response = self.api_service.get(endpoint=GET_PROJECT_MEMBERSHIPS_V2, params=params)
+        organization_id = self._get_organization_id()
+        headers = {"organization-id": organization_id}
+
+        response = self.api_service.get(endpoint=GET_PROJECT_MEMBERSHIPS_V2, params=params, headers=headers)
         validate_status_code(response)
         return parse_json_response(response, "get project memberships")
 
     def get_project_roles(
             self,
-            project_id: str
+            project_id: str,
+            start_page: int = 1,
+            page_size: int = 20,
+            order_key: str = None,
+            order_direction: str = "desc",
+            role_types: str = None
     ) -> dict:
         """
         Retrieves all Roles supported by a specific Project.
 
         :param project_id: str - The unique identifier (GUID) of the project (required).
+        :param start_page: int - The page number for pagination (default is 1).
+        :param page_size: int - The number of items per page (default is 20).
+        :param order_key: str - Field for sorting. Only 'name' is supported (optional).
+        :param order_direction: str - Sort direction: 'asc' or 'desc' (default is 'desc').
+        :param role_types: str - Comma-separated list of role types: 'backend', 'frontend' (optional, case-insensitive).
         :return: dict - The API response containing the list of roles for the project, in JSON format.
         """
-        response = self.api_service.get(endpoint=GET_PROJECT_ROLES_V2, params={"projectId": project_id})
+        params = {
+            "startPage": start_page,
+            "pageSize": page_size,
+            "orderDirection": order_direction
+        }
+        if order_key:
+            params["orderKey"] = order_key
+        if role_types:
+            params["roleTypes"] = role_types
+
+        headers = {"project-id": project_id}
+        response = self.api_service.get(endpoint=GET_PROJECT_ROLES_V2, params=params, headers=headers)
         validate_status_code(response)
         return parse_json_response(response, "get project roles for project", project_id=project_id)
 
     def get_project_members(
             self,
-            project_id: str
+            project_id: str,
+            start_page: int = 1,
+            page_size: int = 20,
+            order_key: str = None,
+            order_direction: str = "desc",
+            role_types: str = None
     ) -> dict:
         """
         Retrieves all members and their Roles for a specific Project.
 
         :param project_id: str - The unique identifier (GUID) of the project (required).
+        :param start_page: int - The page number for pagination (default is 1).
+        :param page_size: int - The number of items per page (default is 20).
+        :param order_key: str - Field for sorting. Only 'name' is supported (optional).
+        :param order_direction: str - Sort direction: 'asc' or 'desc' (default is 'desc').
+        :param role_types: str - Comma-separated list of role types: 'backend', 'frontend' (optional, case-insensitive).
         :return: dict - The API response containing the list of members with their roles, in JSON format.
         """
-        response = self.api_service.get(endpoint=GET_PROJECT_MEMBERS_V2, params={"projectId": project_id})
+        params = {
+            "startPage": start_page,
+            "pageSize": page_size,
+            "orderDirection": order_direction
+        }
+        if order_key:
+            params["orderKey"] = order_key
+        if role_types:
+            params["roleTypes"] = role_types
+
+        headers = {"project-id": project_id}
+        response = self.api_service.get(endpoint=GET_PROJECT_MEMBERS_V2, params=params, headers=headers)
         validate_status_code(response)
         return parse_json_response(response, "get project members for project", project_id=project_id)
 
     def get_organization_members(
             self,
-            organization_id: str
+            organization_id: str,
+            start_page: int = 1,
+            page_size: int = 20,
+            order_key: str = None,
+            order_direction: str = "desc",
+            role_types: str = None
     ) -> dict:
         """
         Retrieves all members and their Roles for a specific Organization.
 
         :param organization_id: str - The unique identifier (GUID) of the organization (required).
+        :param start_page: int - The page number for pagination (default is 1).
+        :param page_size: int - The number of items per page (default is 20).
+        :param order_key: str - Field for sorting. Only 'email' is supported (optional).
+        :param order_direction: str - Sort direction: 'asc' or 'desc' (default is 'desc').
+        :param role_types: str - Comma-separated list of role types. Only 'backend' is supported for organizations (optional, case-insensitive).
         :return: dict - The API response containing the list of members with their roles, in JSON format.
         """
-        response = self.api_service.get(endpoint=GET_ORGANIZATION_MEMBERS_V2, params={"organizationId": organization_id})
+        params = {
+            "organizationId": organization_id,
+            "startPage": start_page,
+            "pageSize": page_size,
+            "orderDirection": order_direction
+        }
+        if order_key:
+            params["orderKey"] = order_key
+        if role_types:
+            params["roleTypes"] = role_types
+
+        response = self.api_service.get(endpoint=GET_ORGANIZATION_MEMBERS_V2, params=params)
         validate_status_code(response)
         return parse_json_response(response, "get organization members for organization", organization_id=organization_id)
 
+    def get_runtime_policies(
+            self,
+            organization_id: str = None
+    ) -> dict:
+        """
+        Retrieves the plugin runtime policies defined for a given organization.
+
+        If no policies are defined, the response indicates that individual policies will apply.
+
+        :param organization_id: str - The unique identifier (GUID) of the organization (optional, defaults to token's organization).
+        :return: dict - The API response containing plugin runtime policies or a message indicating no policies are defined, in JSON format.
+        """
+        if not organization_id:
+            organization_id = self._get_organization_id()
+        headers = {"organization-id": organization_id}
+        response = self.api_service.get(endpoint=GET_PLUGIN_RUNTIME_POLICIES_V2, headers=headers)
+        validate_status_code(response)
+        return parse_json_response(response, "get plugin runtime policies for organization", organization_id=organization_id)
+
     def add_project_member(
             self,
             project_id: str,

pygeai/organization/endpoints.py

@@ -11,6 +11,7 @@ GET_PROJECT_MEMBERSHIPS_V2 = "v2/accessControl/projects/memberships"   # GET - L
 GET_PROJECT_ROLES_V2 = "v2/accessControl/projects/roles"           # GET - Lists all Roles supported by a specific Project
 GET_PROJECT_MEMBERS_V2 = "v2/accessControl/projects/members"       # GET - Lists all members and their Roles for a specific Project
 GET_ORGANIZATION_MEMBERS_V2 = "v2/accessControl/organizations/members"  # GET - Lists all members and their Roles for a specific Organization
+GET_PLUGIN_RUNTIME_POLICIES_V2 = "v2/accessControl/organization/plugin-runtime-policies"  # GET - Retrieves plugin runtime policies for an organization
 ADD_PROJECT_MEMBER_V2 = "v2/accessControl/projects/members"        # POST - Adds a member to a project by sending an invitation
 CREATE_ORGANIZATION_V2 = "v2/admin/organizations"                  # POST - Creates a new Organization
 GET_ORGANIZATION_LIST_V2 = "v2/admin/organizations"                # GET - Retrieves all the organizations available based on a given search criteria

pygeai/organization/limits/clients.py

@@ -1,7 +1,5 @@
-
-from pygeai import logger
 from pygeai.core.base.clients import BaseClient
-from pygeai.core.common.exceptions import InvalidAPIResponseException
+
 from pygeai.core.utils.validators import validate_status_code
 from pygeai.core.utils.parsers import parse_json_response
 from pygeai.organization.limits.endpoints import SET_ORGANIZATION_USAGE_LIMIT_V2, GET_ORGANIZATION_LATEST_USAGE_LIMIT_V2, \
@@ -35,7 +33,7 @@ class UsageLimitClient(BaseClient):
             data=usage_limit
         )
         validate_status_code(response)
-        return parse_json_response(response, f"set usage limit for organization", organization=organization)
+        return parse_json_response(response, "set usage limit for organization", organization=organization)
 
     def get_organization_latest_usage_limit(self, organization: str) -> dict:
         """
@@ -47,7 +45,7 @@ class UsageLimitClient(BaseClient):
         endpoint = GET_ORGANIZATION_LATEST_USAGE_LIMIT_V2.format(organization=organization)
         response = self.api_service.get(endpoint=endpoint)
         validate_status_code(response)
-        return parse_json_response(response, f"get latest usage limit for organization", organization=organization)
+        return parse_json_response(response, "get latest usage limit for organization", organization=organization)
 
     def get_all_usage_limits_from_organization(self, organization: str) -> dict:
         """
@@ -59,7 +57,7 @@ class UsageLimitClient(BaseClient):
         endpoint = GET_ALL_ORGANIZATION_USAGE_LIMITS_V2.format(organization=organization)
         response = self.api_service.get(endpoint=endpoint)
         validate_status_code(response)
-        return parse_json_response(response, f"get all usage limits for organization", organization=organization)
+        return parse_json_response(response, "get all usage limits for organization", organization=organization)
 
     def delete_usage_limit_from_organization(self, organization: str, limit_id: str) -> dict:
         """

pygeai/organization/limits/managers.py

@@ -1,12 +1,9 @@
 from pygeai import logger
-from pygeai.core.base.mappers import ErrorMapper, ModelMapper
+from pygeai.core.base.mappers import ModelMapper
 from pygeai.core.models import UsageLimit
 from pygeai.organization.limits.clients import UsageLimitClient
 from pygeai.core.handlers import ErrorHandler
 from pygeai.core.common.exceptions import APIError
-import logging
-
-
 
 
 class UsageLimitManager:

pygeai/organization/managers.py

@@ -1,7 +1,7 @@
 from pygeai import logger
-from pygeai.core.base.mappers import ErrorMapper, ResponseMapper
+from pygeai.core.base.mappers import ResponseMapper
 from pygeai.core.handlers import ErrorHandler
-from pygeai.core.models import Project, Organization
+from pygeai.core.models import Project
 from pygeai.core.base.responses import EmptyResponse
 from pygeai.organization.clients import OrganizationClient
 from pygeai.organization.mappers import OrganizationResponseMapper

pygeai/proxy/config.py

@@ -3,6 +3,7 @@ from typing import Optional
 import uuid
 import os
 
+
 class ProxySettingsManager(SettingsManager):
     """
     Settings manager for the GEAI proxy.

pygeai/proxy/managers.py

@@ -2,7 +2,7 @@ from contextlib import AsyncExitStack
 from typing import Any, List, Dict, Tuple
 import asyncio
 import json
-import sys
+
 import requests
 from urllib3.exceptions import MaxRetryError
 import mcp.types as types
@@ -12,6 +12,7 @@ from pygeai.proxy.config import ProxySettingsManager
 from pygeai.proxy.clients import ProxyClient, ToolProxyData, ToolProxyJobResult
 from pygeai.core.utils.console import Console
 
+
 class ServerManager:
     """
     Manages multiple MCP servers.
@@ -89,7 +90,7 @@ class ServerManager:
                 tools=list(self.tools.values())
             ))
             Console.write_stdout("----------------------------------")
-            Console.write_stdout(f"Proxy registered successfully:")
+            Console.write_stdout("Proxy registered successfully:")
             if isinstance(result, dict) and isinstance(result.get("Messages"), list):
                 for message in result["Messages"]:
                     description = message.get("Description", "")
@@ -167,7 +168,7 @@ class ServerManager:
                 try:
                     client = await self._initialize_client()
                 except (ConnectionError, TimeoutError, RuntimeError) as e:
-                    Console.write_exception(f"Error during client initialization:", e)
+                    Console.write_exception("Error during client initialization:", e)
                     for i in range(15, 0, -1):
                         Console.write_stdout(f"\rRetrying in {i} seconds...   ",'')
                         await asyncio.sleep(1)
@@ -176,7 +177,7 @@ class ServerManager:
                     continue
 
                 retry_count = 0
-                Console.write_stdout(f"Waiting for jobs...")
+                Console.write_stdout("Waiting for jobs...")
                 while True:
                     try:
                         jobs = client.dequeue()
@@ -185,7 +186,7 @@ class ServerManager:
                         retry_count += 1
                         if retry_count >= MAX_RETRIES:
                             Console.write_stderr(f"Failed to dequeue jobs after {MAX_RETRIES} retries.")
-                            Console.write_exception(f"Exception:", e, "\nExiting...")
+                            Console.write_exception("Exception:", e, "\nExiting...")
                             return
                         Console.write_stderr(f"Failed to dequeue jobs (attempt {retry_count}/{MAX_RETRIES}):")
                         for i in range(15, 0, -1):

pygeai/tests/admin/test_clients.py

@@ -3,7 +3,7 @@ from unittest.mock import patch, MagicMock
 from json import JSONDecodeError
 
 from pygeai.admin.clients import AdminClient
-from pygeai.core.common.exceptions import InvalidAPIResponseException, APIResponseError, APIResponseError
+from pygeai.core.common.exceptions import APIResponseError
 
 
 class TestAdminClient(unittest.TestCase):
@@ -15,7 +15,7 @@ class TestAdminClient(unittest.TestCase):
         self.client = AdminClient()
         self.mock_response = MagicMock()
 
-    @patch('pygeai.core.services.rest.ApiService.get')
+    @patch('pygeai.core.services.rest.GEAIApiService.get')
     def test_validate_api_token_success(self, mock_get):
         self.mock_response.json.return_value = {"organizationId": "org-123", "projectId": "proj-123"}
         self.mock_response.status_code = 200
@@ -26,7 +26,7 @@ class TestAdminClient(unittest.TestCase):
         mock_get.assert_called_once()
         self.assertEqual(result, {"organizationId": "org-123", "projectId": "proj-123"})
 
-    @patch('pygeai.core.services.rest.ApiService.get')
+    @patch('pygeai.core.services.rest.GEAIApiService.get')
     def test_validate_api_token_json_decode_error(self, mock_get):
         self.mock_response.json.side_effect = JSONDecodeError("error", "doc", 0)
         self.mock_response.status_code = 500
@@ -37,7 +37,7 @@ class TestAdminClient(unittest.TestCase):
             self.client.validate_api_token()
         self.assertIn("API returned an error", str(context.exception))  # "Unable to validate API token", str(context.exception))
 
-    @patch('pygeai.core.services.rest.ApiService.get')
+    @patch('pygeai.core.services.rest.GEAIApiService.get')
     def test_get_authorized_organizations_success(self, mock_get):
         self.mock_response.json.return_value = {"organizations": ["org1", "org2"]}
         self.mock_response.status_code = 200
@@ -48,7 +48,7 @@ class TestAdminClient(unittest.TestCase):
         mock_get.assert_called_once()
         self.assertEqual(result, {"organizations": ["org1", "org2"]})
 
-    @patch('pygeai.core.services.rest.ApiService.get')
+    @patch('pygeai.core.services.rest.GEAIApiService.get')
     def test_get_authorized_organizations_json_decode_error(self, mock_get):
         self.mock_response.json.side_effect = JSONDecodeError("error", "doc", 0)
         self.mock_response.status_code = 500
@@ -59,7 +59,7 @@ class TestAdminClient(unittest.TestCase):
             self.client.get_authorized_organizations()
         self.assertIn("API returned an error", str(context.exception))  # "Unable to retrieve authorized organizations", str(context.exception))
 
-    @patch('pygeai.core.services.rest.ApiService.get')
+    @patch('pygeai.core.services.rest.GEAIApiService.get')
     def test_get_authorized_projects_by_organization_success(self, mock_get):
         self.mock_response.json.return_value = {"projects": ["proj1", "proj2"]}
         self.mock_response.status_code = 200
@@ -72,7 +72,7 @@ class TestAdminClient(unittest.TestCase):
         self.assertEqual(call_args[1]['params']['organization'], "org-123")
         self.assertEqual(result, {"projects": ["proj1", "proj2"]})
 
-    @patch('pygeai.core.services.rest.ApiService.get')
+    @patch('pygeai.core.services.rest.GEAIApiService.get')
     def test_get_authorized_projects_by_organization_json_decode_error(self, mock_get):
         self.mock_response.json.side_effect = JSONDecodeError("error", "doc", 0)
         self.mock_response.status_code = 500
@@ -83,7 +83,7 @@ class TestAdminClient(unittest.TestCase):
             self.client.get_authorized_projects_by_organization("org-123")
         self.assertIn("API returned an error", str(context.exception))  # "Unable to retrieve authorized projects for organization", str(context.exception))
 
-    @patch('pygeai.core.services.rest.ApiService.get')
+    @patch('pygeai.core.services.rest.GEAIApiService.get')
     def test_get_project_visibility_success(self, mock_get):
         self.mock_response.json.return_value = {}
         self.mock_response.status_code = 200
@@ -102,7 +102,7 @@ class TestAdminClient(unittest.TestCase):
         self.assertEqual(call_args[1]['params']['accessToken'], "token-789")
         self.assertEqual(result, {})
 
-    @patch('pygeai.core.services.rest.ApiService.get')
+    @patch('pygeai.core.services.rest.GEAIApiService.get')
     def test_get_project_visibility_json_decode_error(self, mock_get):
         self.mock_response.json.side_effect = JSONDecodeError("error", "doc", 0)
         self.mock_response.status_code = 403
@@ -113,7 +113,7 @@ class TestAdminClient(unittest.TestCase):
             self.client.get_project_visibility("org-123", "proj-456", "token-789")
         self.assertIn("API returned an error", str(context.exception))  # "Unable to retrieve project visibility", str(context.exception))
 
-    @patch('pygeai.core.services.rest.ApiService.get')
+    @patch('pygeai.core.services.rest.GEAIApiService.get')
     def test_get_project_api_token_success(self, mock_get):
         self.mock_response.json.return_value = {"apiToken": "api-token-123"}
         self.mock_response.status_code = 200
@@ -132,7 +132,7 @@ class TestAdminClient(unittest.TestCase):
         self.assertEqual(call_args[1]['params']['accessToken'], "token-789")
         self.assertEqual(result, {"apiToken": "api-token-123"})
 
-    @patch('pygeai.core.services.rest.ApiService.get')
+    @patch('pygeai.core.services.rest.GEAIApiService.get')
     def test_get_project_api_token_json_decode_error(self, mock_get):
         self.mock_response.json.side_effect = JSONDecodeError("error", "doc", 0)
         self.mock_response.status_code = 401

pygeai/tests/assistants/rag/test_clients.py

@@ -2,7 +2,7 @@ import unittest
 from unittest.mock import patch, MagicMock
 from json import JSONDecodeError
 from pygeai.assistant.rag.clients import RAGAssistantClient
-from pygeai.core.common.exceptions import InvalidAPIResponseException, APIResponseError, APIResponseError
+from pygeai.core.common.exceptions import APIResponseError
 
 
 class TestRAGAssistantClient(unittest.TestCase):

pygeai/tests/assistants/rag/test_models.py

@@ -1,7 +1,6 @@
 import unittest
 from datetime import datetime
 
-from pydantic import BaseModel
 
 from pygeai.assistant.rag.models import (
     DocumentMetadata,
@@ -20,7 +19,7 @@ from pygeai.assistant.rag.models import (
     UploadDocument,
     UploadType
 )
-from pygeai.core.models import LlmSettings, ChatVariableList, Assistant
+from pygeai.core.models import LlmSettings, ChatVariableList
 
 
 class TestRagModels(unittest.TestCase):

pygeai/tests/assistants/test_clients.py

@@ -2,7 +2,7 @@ import unittest
 from unittest.mock import patch, MagicMock
 from json import JSONDecodeError
 from pygeai.assistant.clients import AssistantClient
-from pygeai.core.common.exceptions import InvalidAPIResponseException, APIResponseError, APIResponseError
+from pygeai.core.common.exceptions import APIResponseError
 
 
 class TestAssistantClient(unittest.TestCase):

pygeai/tests/assistants/test_managers.py

@@ -1,10 +1,8 @@
 import unittest
 from unittest.mock import MagicMock, patch
-from datetime import datetime
 from pygeai.assistant.managers import AssistantManager
 from pygeai.core.common.exceptions import MissingRequirementException, APIError
-from pygeai.core.models import Assistant, TextAssistant, ChatAssistant, ChatWithDataAssistant, \
-    ChatMessageList, AssistantRevision, ChatVariableList, LlmSettings, ChatToolList, ToolChoice, ChatMessage
+from pygeai.core.models import Assistant, ChatAssistant, ChatMessageList, LlmSettings, ChatMessage
 from pygeai.assistant.rag.models import Document, UploadDocument, RAGAssistant
 from pygeai.core.feedback.models import FeedbackRequest
 from pygeai.core.responses import NewAssistantResponse, ChatResponse