diff --git a/doc/source/user_guide/actions.rst b/doc/source/user_guide/actions.rst index 82fafec..481505d 100644 --- a/doc/source/user_guide/actions.rst +++ b/doc/source/user_guide/actions.rst @@ -57,6 +57,21 @@ The following actions can be performed using ``select_ai``: Action methods ============== +Request-specific profile attributes can be supplied with the optional +``attributes`` mapping. For example, override ``additional_instructions`` for +a single chat request: + +.. code-block:: python + + response = profile.chat( + prompt="Write sample Python code to build and test a regression model", + attributes={ + "additional_instructions": ( + "Return executable code as plain text without Markdown fences." + ) + }, + ) + .. list-table:: Action to method mapping :header-rows: 1 :widths: 30 35 35 diff --git a/doc/source/user_guide/profile_attributes.rst b/doc/source/user_guide/profile_attributes.rst index c725aba..fba5783 100644 --- a/doc/source/user_guide/profile_attributes.rst +++ b/doc/source/user_guide/profile_attributes.rst @@ -74,6 +74,9 @@ Attribute groups - Tunes model generation behavior. * - ``conversation`` - Enables conversation history for context-aware chat workflows. + * - ``additional_instructions`` + - Provides persistent guidance, business rules, or response constraints + for requests that use the profile. * - ``source_language``, ``target_language`` - Set default languages for ``Profile.translate()`` and ``AsyncProfile.translate()``. If no source language is configured or @@ -134,5 +137,22 @@ responses: stop_tokens='[";"]', ) +Additional instructions +======================= + +Use ``additional_instructions`` for guidance that should apply to every +request made with a profile: + +.. code-block:: python + + attributes = select_ai.ProfileAttributes( + provider=provider, + credential_name="my_oci_ai_profile_key", + additional_instructions="Return concise, executable Python code.", + ) + +For guidance that applies to only one request, pass an ``attributes`` mapping +to an action method instead of changing the saved profile. + .. autoclass:: select_ai.ProfileAttributes :members: diff --git a/src/select_ai/__init__.py b/src/select_ai/__init__.py index 4f4871a..6441ef6 100644 --- a/src/select_ai/__init__.py +++ b/src/select_ai/__init__.py @@ -17,8 +17,12 @@ from .credential import ( async_create_credential, async_delete_credential, + async_grant_credential_access, + async_revoke_credential_access, create_credential, delete_credential, + grant_credential_access, + revoke_credential_access, ) from .db import ( async_connect, diff --git a/src/select_ai/_validations.py b/src/select_ai/_validations.py index 2de68f3..2ccfbc5 100644 --- a/src/select_ai/_validations.py +++ b/src/select_ai/_validations.py @@ -13,6 +13,16 @@ NoneType = type(None) +def validate_user_or_role_name(user_or_role_name: str) -> str: + """Validate and normalize a sharing grantee name.""" + if not isinstance(user_or_role_name, str): + raise TypeError("'user_or_role_name' must be a string") + user_or_role_name = user_or_role_name.strip() + if not user_or_role_name: + raise ValueError("'user_or_role_name' cannot be empty") + return user_or_role_name + + def _match(value, annot) -> bool: """Recursively validate value against a typing annotation.""" if annot is Any: diff --git a/src/select_ai/agent/team.py b/src/select_ai/agent/team.py index 42239f1..e767ba6 100644 --- a/src/select_ai/agent/team.py +++ b/src/select_ai/agent/team.py @@ -21,6 +21,7 @@ import oracledb from select_ai._abc import SelectAIDataClass +from select_ai._validations import validate_user_or_role_name from select_ai.agent.sql import ( GET_USER_AI_AGENT_TEAM, GET_USER_AI_AGENT_TEAM_ATTRIBUTES, @@ -243,6 +244,30 @@ def enable(self): }, ) + def grant_access(self, user_or_role_name: str) -> None: + """Grant a database user or role access to this agent team.""" + user_or_role_name = validate_user_or_role_name(user_or_role_name) + with cursor() as cr: + cr.callproc( + "DBMS_CLOUD_AI_AGENT.GRANT_TEAM_ACCESS", + keyword_parameters={ + "team_name": self.team_name, + "user_or_role_name": user_or_role_name, + }, + ) + + def revoke_access(self, user_or_role_name: str) -> None: + """Revoke a database user or role's access to this agent team.""" + user_or_role_name = validate_user_or_role_name(user_or_role_name) + with cursor() as cr: + cr.callproc( + "DBMS_CLOUD_AI_AGENT.REVOKE_TEAM_ACCESS", + keyword_parameters={ + "team_name": self.team_name, + "user_or_role_name": user_or_role_name, + }, + ) + @classmethod def fetch(cls, team_name: str) -> "Team": """ @@ -704,6 +729,30 @@ async def enable(self): }, ) + async def grant_access(self, user_or_role_name: str) -> None: + """Asynchronously grant a user or role access to this agent team.""" + user_or_role_name = validate_user_or_role_name(user_or_role_name) + async with async_cursor() as cr: + await cr.callproc( + "DBMS_CLOUD_AI_AGENT.GRANT_TEAM_ACCESS", + keyword_parameters={ + "team_name": self.team_name, + "user_or_role_name": user_or_role_name, + }, + ) + + async def revoke_access(self, user_or_role_name: str) -> None: + """Asynchronously revoke a user or role's agent team access.""" + user_or_role_name = validate_user_or_role_name(user_or_role_name) + async with async_cursor() as cr: + await cr.callproc( + "DBMS_CLOUD_AI_AGENT.REVOKE_TEAM_ACCESS", + keyword_parameters={ + "team_name": self.team_name, + "user_or_role_name": user_or_role_name, + }, + ) + @classmethod async def fetch(cls, team_name: str) -> "AsyncTeam": """ diff --git a/src/select_ai/async_profile.py b/src/select_ai/async_profile.py index 91572fe..c9ccd33 100644 --- a/src/select_ai/async_profile.py +++ b/src/select_ai/async_profile.py @@ -19,6 +19,7 @@ import oracledb import pandas +from select_ai._validations import validate_user_or_role_name from select_ai.action import Action from select_ai.base_profile import ( BaseProfile, @@ -43,9 +44,9 @@ ) from select_ai.provider import Provider from select_ai.sql import ( - GET_USER_AI_PROFILE, - GET_USER_AI_PROFILE_ATTRIBUTES, - LIST_USER_AI_PROFILES, + GET_ALL_AI_PROFILE, + GET_ALL_AI_PROFILE_ATTRIBUTES, + LIST_ALL_AI_PROFILES, ) from select_ai.summary import SummaryParams from select_ai.synthetic_data import SyntheticDataAttributes @@ -75,12 +76,17 @@ async def _init_profile(self): if self.profile_name: profile_exists = False try: - saved_description = await self._get_profile_description( - profile_name=self.profile_name + saved_description, saved_owner = ( + await self._get_profile_description( + profile_name=self.profile_name, + owner=self.owner, + ) ) + self.owner = saved_owner profile_exists = True saved_attributes = await self._get_attributes( profile_name=self.profile_name, + owner=self.owner, raise_on_empty=True, ) self._raise_error_if_profile_exists() @@ -100,38 +106,43 @@ async def _init_profile(self): return self @staticmethod - async def _get_profile_description(profile_name) -> Union[str, None]: - """Get description of profile from USER_CLOUD_AI_PROFILES + async def _get_profile_description( + profile_name: str, owner: Optional[str] = None + ) -> Tuple[Union[str, None], str]: + """Get a profile description and owner from ALL_CLOUD_AI_PROFILES. :param str profile_name: Name of profile - :return: Description of profile - :rtype: str + :param str owner: Owner of a shared profile. Defaults to current schema. + :return: Tuple containing the profile description and owner. :raises: ProfileNotFoundError """ async with async_cursor() as cr: await cr.execute( - GET_USER_AI_PROFILE, + GET_ALL_AI_PROFILE, profile_name=profile_name.upper(), + owner=owner.upper() if owner else None, ) profile = await cr.fetchone() - if profile is None: - raise ProfileNotFoundError(profile_name) if profile: if profile[1] is not None: - return await profile[1].read() + description = await profile[1].read() else: - return None + description = None + return description, profile[2] else: raise ProfileNotFoundError(profile_name) @staticmethod async def _get_attributes( - profile_name: str, raise_on_empty: bool = True + profile_name: str, + owner: Optional[str] = None, + raise_on_empty: bool = True, ) -> Union[ProfileAttributes, None]: """Asynchronously gets AI profile attributes from the Database :param str profile_name: Name of the profile + :param str owner: Owner of a shared profile. Defaults to current schema. :param bool raise_on_empty: Raise an error if attributes are empty :return: select_ai.provider.ProviderAttributes :raises: select_ai.errors.ProfileAttributesEmptyError @@ -139,8 +150,9 @@ async def _get_attributes( """ async with async_cursor() as cr: await cr.execute( - GET_USER_AI_PROFILE_ATTRIBUTES, + GET_ALL_AI_PROFILE_ATTRIBUTES, profile_name=profile_name.upper(), + owner=owner.upper() if owner else None, ) attributes = await cr.fetchall() if attributes: @@ -158,7 +170,10 @@ async def get_attributes(self) -> ProfileAttributes: :return: select_ai.provider.ProviderAttributes :raises: ProfileNotFoundError """ - return await self._get_attributes(profile_name=self.profile_name) + return await self._get_attributes( + profile_name=self.profile_name, + owner=self.owner, + ) async def _set_attribute( self, @@ -299,6 +314,30 @@ async def disable(self) -> None: keyword_parameters={"profile_name": self.profile_name}, ) + async def grant_access(self, user_or_role_name: str) -> None: + """Asynchronously grant a user or role access to this AI profile.""" + user_or_role_name = validate_user_or_role_name(user_or_role_name) + async with async_cursor() as cr: + await cr.callproc( + "DBMS_CLOUD_AI.GRANT_PROFILE_ACCESS", + keyword_parameters={ + "profile_name": self.profile_name, + "user_or_role_name": user_or_role_name, + }, + ) + + async def revoke_access(self, user_or_role_name: str) -> None: + """Asynchronously revoke a user or role's profile access.""" + user_or_role_name = validate_user_or_role_name(user_or_role_name) + async with async_cursor() as cr: + await cr.callproc( + "DBMS_CLOUD_AI.REVOKE_PROFILE_ACCESS", + keyword_parameters={ + "profile_name": self.profile_name, + "user_or_role_name": user_or_role_name, + }, + ) + @classmethod async def delete_profile(cls, profile_name: str, force: bool = False): """Asynchronously deletes an AI profile from the database @@ -311,15 +350,22 @@ async def delete_profile(cls, profile_name: str, force: bool = False): await cls._delete(profile_name=profile_name, force=force) @classmethod - async def fetch(cls, profile_name: str) -> "AsyncProfile": + async def fetch( + cls, profile_name: str, owner: Optional[str] = None + ) -> "AsyncProfile": """Asynchronously create an AI Profile object from attributes saved in the database - :param str profile_name: + :param str profile_name: Name of the AI profile. + :param str owner: Owner of a shared profile. Defaults to current schema. :return: select_ai.Profile :raises: ProfileNotFoundError """ - return await cls(profile_name, raise_error_if_exists=False) + return await cls( + profile_name, + owner=owner, + raise_error_if_exists=False, + ) async def _save_feedback( self, @@ -410,26 +456,31 @@ async def delete_feedback( @classmethod async def list( - cls, profile_name_pattern: str = ".*" + cls, + profile_name_pattern: str = ".*", + owner: Optional[str] = None, ) -> AsyncGenerator["AsyncProfile", None]: """Asynchronously list AI Profiles saved in the database. :param str profile_name_pattern: Regular expressions can be used to specify a pattern. Function REGEXP_LIKE is used to perform the match. Default value is ".*" i.e. match all AI profiles. + :param str owner: Owner of shared profiles. Defaults to current schema. :return: Iterator[Profile] """ async with async_cursor() as cr: await cr.execute( - LIST_USER_AI_PROFILES, + LIST_ALL_AI_PROFILES, profile_name_pattern=profile_name_pattern, + owner=owner.upper() if owner else None, ) rows = await cr.fetchall() for row in rows: profile_name = row[0] yield await cls( profile_name=profile_name, + owner=row[2], raise_error_if_exists=False, raise_error_on_empty_attributes=False, ) @@ -440,6 +491,7 @@ async def _generate_with_cursor( prompt: str, action=Action.SHOWSQL, params: Mapping = None, + attributes: Mapping = None, ) -> Union[pandas.DataFrame, str, None]: """Asynchronously perform AI translation using this profile @@ -447,9 +499,13 @@ async def _generate_with_cursor( :param select_ai.profile.Action action: :param params: Parameters to include in the LLM request. For e.g. conversation_id for context-aware chats + :param Mapping attributes: Profile attributes to override for this + request :return: Union[pandas.DataFrame, str] """ - parameters = self._generate_parameters(prompt, action, params) + parameters = self._generate_parameters( + prompt, action, params, attributes + ) data = await cr.callfunc( "DBMS_CLOUD_AI.GENERATE", @@ -470,6 +526,7 @@ def _generate_parameters( prompt: str, action, params: Mapping = None, + attributes: Mapping = None, ) -> Mapping: if not prompt: raise ValueError("prompt cannot be empty or None") @@ -478,10 +535,13 @@ def _generate_parameters( "prompt": prompt, "action": action, "profile_name": self.profile_name, - # "attributes": self.attributes.json(), } if params: parameters["params"] = json.dumps(params) + if attributes is not None: + if not isinstance(attributes, Mapping): + raise TypeError("'attributes' must be a mapping") + parameters["attributes"] = json.dumps(attributes) return parameters async def _generate_stream( @@ -490,6 +550,7 @@ async def _generate_stream( action, params: Mapping = None, chunk_size: int = 8192, + attributes: Mapping = None, ) -> AsyncGenerator[str, None]: async with async_cursor() as cr: async for chunk in self._generate_stream_with_cursor( @@ -498,6 +559,7 @@ async def _generate_stream( action=action, params=params, chunk_size=chunk_size, + attributes=attributes, ): yield chunk @@ -508,13 +570,16 @@ async def _generate_stream_with_cursor( action, params: Mapping = None, chunk_size: int = 8192, + attributes: Mapping = None, ) -> AsyncGenerator[str, None]: if action == Action.RUNSQL: raise ValueError("stream=True is not supported for run_sql") if chunk_size <= 0: raise ValueError("chunk_size must be greater than 0") - parameters = self._generate_parameters(prompt, action, params) + parameters = self._generate_parameters( + prompt, action, params, attributes + ) data = await cr.callfunc( "DBMS_CLOUD_AI.GENERATE", oracledb.DB_TYPE_CLOB, @@ -538,6 +603,8 @@ async def generate( params: Mapping = None, stream: bool = False, chunk_size: int = 8192, + *, + attributes: Mapping = None, ) -> Union[pandas.DataFrame, str, AsyncGenerator[str, None], None]: """Asynchronously perform AI translation using this profile @@ -545,15 +612,23 @@ async def generate( :param select_ai.profile.Action action: :param params: Parameters to include in the LLM request. For e.g. conversation_id for context-aware chats + :param Mapping attributes: Profile attributes to override for this + request :param bool stream: Return an async iterator of response chunks :param int chunk_size: Number of characters to read per stream chunk :return: Union[pandas.DataFrame, str] """ if stream: - return self._generate_stream(prompt, action, params, chunk_size) + return self._generate_stream( + prompt, action, params, chunk_size, attributes + ) async with async_cursor() as cr: return await self._generate_with_cursor( - cr, prompt=prompt, action=action, params=params + cr, + prompt=prompt, + action=action, + params=params, + attributes=attributes, ) async def chat( @@ -562,11 +637,15 @@ async def chat( params: Mapping = None, stream: bool = False, chunk_size: int = 8192, + *, + attributes: Mapping = None, ) -> Union[str, AsyncGenerator[str, None]]: """Asynchronously chat with the LLM :param str prompt: Natural language prompt :param params: Parameters to include in the LLM request + :param Mapping attributes: Profile attributes to override for this + request :param bool stream: Return an async iterator of response chunks :param int chunk_size: Number of characters to read per stream chunk :return: str @@ -577,6 +656,7 @@ async def chat( params=params, stream=stream, chunk_size=chunk_size, + attributes=attributes, ) @asynccontextmanager @@ -611,11 +691,15 @@ async def narrate( params: Mapping = None, stream: bool = False, chunk_size: int = 8192, + *, + attributes: Mapping = None, ) -> Union[str, AsyncGenerator[str, None]]: """Narrate the result of the SQL :param str prompt: Natural language prompt :param params: Parameters to include in the LLM request + :param Mapping attributes: Profile attributes to override for this + request :param bool stream: Return an async iterator of response chunks :param int chunk_size: Number of characters to read per stream chunk :return: str @@ -626,6 +710,7 @@ async def narrate( params=params, stream=stream, chunk_size=chunk_size, + attributes=attributes, ) async def explain_sql( @@ -634,11 +719,15 @@ async def explain_sql( params: Mapping = None, stream: bool = False, chunk_size: int = 8192, + *, + attributes: Mapping = None, ): """Explain the generated SQL :param str prompt: Natural language prompt :param params: Parameters to include in the LLM request + :param Mapping attributes: Profile attributes to override for this + request :param bool stream: Return an async iterator of response chunks :param int chunk_size: Number of characters to read per stream chunk :return: str @@ -649,18 +738,30 @@ async def explain_sql( params=params, stream=stream, chunk_size=chunk_size, + attributes=attributes, ) async def run_sql( - self, prompt, params: Mapping = None + self, + prompt, + params: Mapping = None, + *, + attributes: Mapping = None, ) -> pandas.DataFrame: """Explain the generated SQL :param str prompt: Natural language prompt :param params: Parameters to include in the LLM request + :param Mapping attributes: Profile attributes to override for this + request :return: pandas.DataFrame """ - return await self.generate(prompt, action=Action.RUNSQL, params=params) + return await self.generate( + prompt, + action=Action.RUNSQL, + params=params, + attributes=attributes, + ) async def show_sql( self, @@ -668,11 +769,15 @@ async def show_sql( params: Mapping = None, stream: bool = False, chunk_size: int = 8192, + *, + attributes: Mapping = None, ): """Show the generated SQL :param str prompt: Natural language prompt :param params: Parameters to include in the LLM request + :param Mapping attributes: Profile attributes to override for this + request :param bool stream: Return an async iterator of response chunks :param int chunk_size: Number of characters to read per stream chunk :return: str @@ -683,6 +788,7 @@ async def show_sql( params=params, stream=stream, chunk_size=chunk_size, + attributes=attributes, ) async def show_prompt( @@ -691,11 +797,15 @@ async def show_prompt( params: Mapping = None, stream: bool = False, chunk_size: int = 8192, + *, + attributes: Mapping = None, ): """Show the prompt sent to LLM :param str prompt: Natural language prompt :param params: Parameters to include in the LLM request + :param Mapping attributes: Profile attributes to override for this + request :param bool stream: Return an async iterator of response chunks :param int chunk_size: Number of characters to read per stream chunk :return: str @@ -706,6 +816,7 @@ async def show_prompt( params=params, stream=stream, chunk_size=chunk_size, + attributes=attributes, ) async def summarize( @@ -776,6 +887,8 @@ async def run_pipeline( self, prompt_specifications: List[Tuple[str, Action]], continue_on_error: bool = False, + *, + attributes: Mapping = None, ) -> List[Union[str, pandas.DataFrame]]: """Send Multiple prompts in a single roundtrip to the Database @@ -784,16 +897,25 @@ async def run_pipeline( corresponding action :param bool continue_on_error: True to continue on error else False + :param Mapping attributes: Profile attributes to override for every + request in the pipeline :return: List[Union[str, pandas.DataFrame]] """ + serialized_attributes = None + if attributes is not None: + if not isinstance(attributes, Mapping): + raise TypeError("'attributes' must be a mapping") + serialized_attributes = json.dumps(attributes) + pipeline = oracledb.create_pipeline() for prompt, action in prompt_specifications: parameters = { "prompt": prompt, "action": action, "profile_name": self.profile_name, - # "attributes": self.attributes.json(), } + if serialized_attributes is not None: + parameters["attributes"] = serialized_attributes pipeline.add_callfunc( "DBMS_CLOUD_AI.GENERATE", return_type=oracledb.DB_TYPE_CLOB, @@ -866,7 +988,12 @@ def __init__(self, async_profile: AsyncProfile, params: Mapping): self._cursor = None async def chat( - self, prompt: str, stream: bool = False, chunk_size: int = 8192 + self, + prompt: str, + stream: bool = False, + chunk_size: int = 8192, + *, + attributes: Mapping = None, ) -> Union[str, AsyncGenerator[str, None]]: if stream: return self.async_profile._generate_stream_with_cursor( @@ -875,13 +1002,23 @@ async def chat( action=Action.CHAT, params=self.params, chunk_size=chunk_size, + attributes=attributes, ) return await self.async_profile._generate_with_cursor( - self._cursor, prompt=prompt, action=Action.CHAT, params=self.params + self._cursor, + prompt=prompt, + action=Action.CHAT, + params=self.params, + attributes=attributes, ) async def narrate( - self, prompt, stream: bool = False, chunk_size: int = 8192 + self, + prompt, + stream: bool = False, + chunk_size: int = 8192, + *, + attributes: Mapping = None, ) -> Union[str, AsyncGenerator[str, None]]: """Narrate the result of the SQL @@ -897,13 +1034,23 @@ async def narrate( action=Action.NARRATE, params=self.params, chunk_size=chunk_size, + attributes=attributes, ) return await self.async_profile._generate_with_cursor( - self._cursor, prompt, action=Action.NARRATE, params=self.params + self._cursor, + prompt, + action=Action.NARRATE, + params=self.params, + attributes=attributes, ) async def explain_sql( - self, prompt: str, stream: bool = False, chunk_size: int = 8192 + self, + prompt: str, + stream: bool = False, + chunk_size: int = 8192, + *, + attributes: Mapping = None, ) -> Union[str, AsyncGenerator[str, None]]: """Explain the generated SQL @@ -919,23 +1066,39 @@ async def explain_sql( action=Action.EXPLAINSQL, params=self.params, chunk_size=chunk_size, + attributes=attributes, ) return await self.async_profile._generate_with_cursor( - self._cursor, prompt, action=Action.EXPLAINSQL, params=self.params + self._cursor, + prompt, + action=Action.EXPLAINSQL, + params=self.params, + attributes=attributes, ) - async def run_sql(self, prompt: str) -> pandas.DataFrame: + async def run_sql( + self, prompt: str, *, attributes: Mapping = None + ) -> pandas.DataFrame: """Explain the generated SQL :param str prompt: Natural language prompt :return: pandas.DataFrame """ return await self.async_profile._generate_with_cursor( - self._cursor, prompt, action=Action.RUNSQL, params=self.params + self._cursor, + prompt, + action=Action.RUNSQL, + params=self.params, + attributes=attributes, ) async def show_sql( - self, prompt, stream: bool = False, chunk_size: int = 8192 + self, + prompt, + stream: bool = False, + chunk_size: int = 8192, + *, + attributes: Mapping = None, ) -> Union[str, AsyncGenerator[str, None]]: """Show the generated SQL @@ -951,13 +1114,23 @@ async def show_sql( action=Action.SHOWSQL, params=self.params, chunk_size=chunk_size, + attributes=attributes, ) return await self.async_profile._generate_with_cursor( - self._cursor, prompt, action=Action.SHOWSQL, params=self.params + self._cursor, + prompt, + action=Action.SHOWSQL, + params=self.params, + attributes=attributes, ) async def show_prompt( - self, prompt: str, stream: bool = False, chunk_size: int = 8192 + self, + prompt: str, + stream: bool = False, + chunk_size: int = 8192, + *, + attributes: Mapping = None, ) -> Union[str, AsyncGenerator[str, None]]: """Show the prompt sent to LLM @@ -973,9 +1146,14 @@ async def show_prompt( action=Action.SHOWPROMPT, params=self.params, chunk_size=chunk_size, + attributes=attributes, ) return await self.async_profile._generate_with_cursor( - self._cursor, prompt, action=Action.SHOWPROMPT, params=self.params + self._cursor, + prompt, + action=Action.SHOWPROMPT, + params=self.params, + attributes=attributes, ) async def __aenter__(self): diff --git a/src/select_ai/base_profile.py b/src/select_ai/base_profile.py index 41f9b90..b31d4a1 100644 --- a/src/select_ai/base_profile.py +++ b/src/select_ai/base_profile.py @@ -41,6 +41,8 @@ class ProfileAttributes(SelectAIDataClass): provider APIs. :param bool enforce_object_list: Specifies whether to restrict the LLM to generate SQL that uses only tables covered by the object list. + :param str additional_instructions: Persistent guidance that Select AI + applies to requests that use the profile. :param int max_tokens: Denotes the number of tokens to return per generation. Default is 1024. :param List[Mapping] object_list: Array of JSON objects specifying @@ -69,6 +71,7 @@ class ProfileAttributes(SelectAIDataClass): """ + additional_instructions: Optional[str] = None annotations: Optional[bool] = None case_sensitive_values: Optional[bool] = None comments: Optional[bool] = None @@ -163,6 +166,8 @@ class BaseProfile(ABC): :param str description: Description of the profile + :param str owner: Database user that owns the profile + :param bool merge: Fetches the profile from database, merges the non-null attributes and saves it back in the database. Default value is False @@ -189,6 +194,7 @@ def __init__( replace: Optional[bool] = False, raise_error_if_exists: Optional[bool] = True, raise_error_on_empty_attributes: Optional[bool] = False, + owner: Optional[str] = None, ): """Initialize a base profile""" self.profile_name = profile_name @@ -199,11 +205,21 @@ def __init__( ) self.attributes = attributes self.description = description + self.owner = owner.upper() if owner else None self.merge = merge self.replace = replace self.raise_error_if_exists = raise_error_if_exists self.raise_error_on_empty_attributes = raise_error_on_empty_attributes + @property + def qualified_name(self) -> Optional[str]: + """Return the owner-qualified profile name when owner is known.""" + if self.profile_name is None: + return None + if self.owner is None: + return self.profile_name + return f"{self.owner}.{self.profile_name}" + def _raise_error_if_profile_exists(self): """ Helper method to raise ProfileExistsError if profile exists diff --git a/src/select_ai/credential.py b/src/select_ai/credential.py index 2946553..5096f64 100644 --- a/src/select_ai/credential.py +++ b/src/select_ai/credential.py @@ -9,16 +9,101 @@ import oracledb +from ._validations import validate_user_or_role_name from .db import async_cursor, cursor __all__ = [ "async_create_credential", "async_delete_credential", + "async_grant_credential_access", + "async_revoke_credential_access", "create_credential", "delete_credential", + "grant_credential_access", + "revoke_credential_access", ] +_CREATE_PUBLIC_CREDENTIAL_SYNONYM = """ +DECLARE + v_credential_name VARCHAR2(261); + v_owner VARCHAR2(261); +BEGIN + v_credential_name := DBMS_ASSERT.ENQUOTE_NAME( + DBMS_ASSERT.SIMPLE_SQL_NAME(UPPER(:credential_name)), + FALSE + ); + v_owner := DBMS_ASSERT.ENQUOTE_NAME( + SYS_CONTEXT('USERENV', 'CURRENT_SCHEMA'), + FALSE + ); + EXECUTE IMMEDIATE + 'CREATE OR REPLACE PUBLIC SYNONYM ' || v_credential_name || + ' FOR ' || v_owner || '.' || v_credential_name; +END; +""" + +_DROP_PUBLIC_CREDENTIAL_SYNONYM = """ +DECLARE + v_credential_name VARCHAR2(261); +BEGIN + v_credential_name := DBMS_ASSERT.ENQUOTE_NAME( + DBMS_ASSERT.SIMPLE_SQL_NAME(UPPER(:credential_name)), + FALSE + ); + EXECUTE IMMEDIATE 'DROP PUBLIC SYNONYM ' || v_credential_name; +END; +""" + +_GRANT_CREDENTIAL_ACCESS = """ +DECLARE + v_credential_name VARCHAR2(261); + v_owner VARCHAR2(261); + v_user_or_role_name VARCHAR2(261); +BEGIN + v_credential_name := DBMS_ASSERT.ENQUOTE_NAME( + DBMS_ASSERT.SIMPLE_SQL_NAME(UPPER(:credential_name)), + FALSE + ); + v_owner := DBMS_ASSERT.ENQUOTE_NAME( + SYS_CONTEXT('USERENV', 'CURRENT_SCHEMA'), + FALSE + ); + v_user_or_role_name := DBMS_ASSERT.ENQUOTE_NAME( + DBMS_ASSERT.SIMPLE_SQL_NAME(UPPER(:user_or_role_name)), + FALSE + ); + EXECUTE IMMEDIATE + 'GRANT EXECUTE ON ' || v_owner || '.' || v_credential_name || + ' TO ' || v_user_or_role_name; +END; +""" + +_REVOKE_CREDENTIAL_ACCESS = """ +DECLARE + v_credential_name VARCHAR2(261); + v_owner VARCHAR2(261); + v_user_or_role_name VARCHAR2(261); +BEGIN + v_credential_name := DBMS_ASSERT.ENQUOTE_NAME( + DBMS_ASSERT.SIMPLE_SQL_NAME(UPPER(:credential_name)), + FALSE + ); + v_owner := DBMS_ASSERT.ENQUOTE_NAME( + SYS_CONTEXT('USERENV', 'CURRENT_SCHEMA'), + FALSE + ); + v_user_or_role_name := DBMS_ASSERT.ENQUOTE_NAME( + DBMS_ASSERT.SIMPLE_SQL_NAME(UPPER(:user_or_role_name)), + FALSE + ); + EXECUTE IMMEDIATE + 'REVOKE EXECUTE ON ' || v_owner || '.' || v_credential_name || + ' FROM ' || v_user_or_role_name; +END; +""" + + def _validate_credential(credential: Mapping[str, str]): valid_keys = { "credential_name", @@ -35,13 +120,105 @@ def _validate_credential(credential: Mapping[str, str]): raise ValueError(f"Invalid key {k} for credential object") -async def async_create_credential(credential: Mapping, replace: bool = False): +async def async_grant_credential_access( + credential_name: str, + user_or_role_name: str, +) -> None: + """Grant a database user or role access to a credential. + + :param str credential_name: Name of the credential in the current schema. + :param str user_or_role_name: Database user or role receiving access. + :return: None + :raises: oracledb.DatabaseError + """ + user_or_role_name = validate_user_or_role_name(user_or_role_name) + async with async_cursor() as cr: + await cr.execute( + _GRANT_CREDENTIAL_ACCESS, + credential_name=credential_name, + user_or_role_name=user_or_role_name, + ) + + +async def async_revoke_credential_access( + credential_name: str, + user_or_role_name: str, +) -> None: + """Revoke a database user or role's access to a credential. + + :param str credential_name: Name of the credential in the current schema. + :param str user_or_role_name: Database user or role losing access. + :return: None + :raises: oracledb.DatabaseError + """ + user_or_role_name = validate_user_or_role_name(user_or_role_name) + async with async_cursor() as cr: + await cr.execute( + _REVOKE_CREDENTIAL_ACCESS, + credential_name=credential_name, + user_or_role_name=user_or_role_name, + ) + + +def grant_credential_access( + credential_name: str, + user_or_role_name: str, +) -> None: + """Grant a database user or role access to a credential. + + :param str credential_name: Name of the credential in the current schema. + :param str user_or_role_name: Database user or role receiving access. + :return: None + :raises: oracledb.DatabaseError + """ + user_or_role_name = validate_user_or_role_name(user_or_role_name) + with cursor() as cr: + cr.execute( + _GRANT_CREDENTIAL_ACCESS, + credential_name=credential_name, + user_or_role_name=user_or_role_name, + ) + + +def revoke_credential_access( + credential_name: str, + user_or_role_name: str, +) -> None: + """Revoke a database user or role's access to a credential. + + :param str credential_name: Name of the credential in the current schema. + :param str user_or_role_name: Database user or role losing access. + :return: None + :raises: oracledb.DatabaseError """ - Async API to create credential. + user_or_role_name = validate_user_or_role_name(user_or_role_name) + with cursor() as cr: + cr.execute( + _REVOKE_CREDENTIAL_ACCESS, + credential_name=credential_name, + user_or_role_name=user_or_role_name, + ) + + +async def async_create_credential( + credential: Mapping, + replace: bool = False, + public_synonym: bool = False, +): + """Asynchronously create a credential. - Creates a credential object using DBMS_CLOUD.CREATE_CREDENTIAL. if replace - is True, credential will be replaced if it already exists + Creates a credential object using DBMS_CLOUD.CREATE_CREDENTIAL. If replace + is True, credential will be replaced if it already exists. If + public_synonym is True, a public synonym with the credential name is + created for the credential in the current schema. Creating the synonym + requires the CREATE PUBLIC SYNONYM system privilege. + :param Mapping credential: Credential attributes accepted by + DBMS_CLOUD.CREATE_CREDENTIAL, including credential_name. + :param bool replace: Replace an existing credential with the same name. + :param bool public_synonym: Create a public synonym for the credential. + :return: None + :raises: oracledb.DatabaseError """ _validate_credential(credential) async with async_cursor() as cr: @@ -66,14 +243,37 @@ async def async_create_credential(credential: Mapping, replace: bool = False): else: raise + if public_synonym: + await cr.execute( + _CREATE_PUBLIC_CREDENTIAL_SYNONYM, + credential_name=credential["credential_name"], + ) -async def async_delete_credential(credential_name: str, force: bool = False): - """ - Async API to create credential. - Deletes a credential object using DBMS_CLOUD.DROP_CREDENTIAL +async def async_delete_credential( + credential_name: str, + force: bool = False, + public_synonym: bool = False, +): + """Asynchronously delete a credential. + + Deletes a credential object using DBMS_CLOUD.DROP_CREDENTIAL. If + public_synonym is True, also drops its public synonym. Dropping the synonym + requires the DROP PUBLIC SYNONYM system privilege. + + :param str credential_name: Name of the credential in the current schema. + :param bool force: Ignore an error when the credential does not exist. + :param bool public_synonym: Drop the credential's public synonym. + :return: None + :raises: oracledb.DatabaseError """ async with async_cursor() as cr: + if public_synonym: + await cr.execute( + _DROP_PUBLIC_CREDENTIAL_SYNONYM, + credential_name=credential_name, + ) + try: await cr.callproc( "DBMS_CLOUD.DROP_CREDENTIAL", @@ -87,12 +287,25 @@ async def async_delete_credential(credential_name: str, force: bool = False): raise -def create_credential(credential: Mapping, replace: bool = False): - """ +def create_credential( + credential: Mapping, + replace: bool = False, + public_synonym: bool = False, +): + """Create a credential. - Creates a credential object using DBMS_CLOUD.CREATE_CREDENTIAL. if replace - is True, credential will be replaced if it "already exists" + Creates a credential object using DBMS_CLOUD.CREATE_CREDENTIAL. If replace + is True, credential will be replaced if it "already exists". If + public_synonym is True, a public synonym with the credential name is + created for the credential in the current schema. Creating the synonym + requires the CREATE PUBLIC SYNONYM system privilege. + :param Mapping credential: Credential attributes accepted by + DBMS_CLOUD.CREATE_CREDENTIAL, including credential_name. + :param bool replace: Replace an existing credential with the same name. + :param bool public_synonym: Create a public synonym for the credential. + :return: None + :raises: oracledb.DatabaseError """ _validate_credential(credential) with cursor() as cr: @@ -117,9 +330,35 @@ def create_credential(credential: Mapping, replace: bool = False): else: raise + if public_synonym: + cr.execute( + _CREATE_PUBLIC_CREDENTIAL_SYNONYM, + credential_name=credential["credential_name"], + ) + + +def delete_credential( + credential_name: str, + force: bool = False, + public_synonym: bool = False, +): + """Delete a credential and optionally its public synonym. -def delete_credential(credential_name: str, force: bool = False): + Dropping the synonym requires the DROP PUBLIC SYNONYM system privilege. + + :param str credential_name: Name of the credential in the current schema. + :param bool force: Ignore an error when the credential does not exist. + :param bool public_synonym: Drop the credential's public synonym. + :return: None + :raises: oracledb.DatabaseError + """ with cursor() as cr: + if public_synonym: + cr.execute( + _DROP_PUBLIC_CREDENTIAL_SYNONYM, + credential_name=credential_name, + ) + try: cr.callproc( "DBMS_CLOUD.DROP_CREDENTIAL", diff --git a/src/select_ai/profile.py b/src/select_ai/profile.py index 23f2b72..833c55f 100644 --- a/src/select_ai/profile.py +++ b/src/select_ai/profile.py @@ -13,6 +13,7 @@ import pandas from select_ai import Conversation +from select_ai._validations import validate_user_or_role_name from select_ai.action import Action from select_ai.base_profile import ( BaseProfile, @@ -29,9 +30,9 @@ from select_ai.feedback import FeedbackOperation, FeedbackType from select_ai.provider import Provider from select_ai.sql import ( - GET_USER_AI_PROFILE, - GET_USER_AI_PROFILE_ATTRIBUTES, - LIST_USER_AI_PROFILES, + GET_ALL_AI_PROFILE, + GET_ALL_AI_PROFILE_ATTRIBUTES, + LIST_ALL_AI_PROFILES, ) from select_ai.summary import SummaryParams from select_ai.synthetic_data import SyntheticDataAttributes @@ -57,12 +58,15 @@ def _init_profile(self) -> None: if self.profile_name: profile_exists = False try: - saved_description = self._get_profile_description( - profile_name=self.profile_name + saved_description, saved_owner = self._get_profile_description( + profile_name=self.profile_name, + owner=self.owner, ) + self.owner = saved_owner profile_exists = True saved_attributes = self._get_attributes( profile_name=self.profile_name, + owner=self.owner, raise_on_empty=True, ) self._raise_error_if_profile_exists() @@ -83,39 +87,51 @@ def _init_profile(self) -> None: ) @staticmethod - def _get_profile_description(profile_name) -> Union[str, None]: - """Get description of profile from USER_CLOUD_AI_PROFILES - - :param str profile_name: - :return: Union[str, None] profile description + def _get_profile_description( + profile_name: str, owner: Optional[str] = None + ) -> Tuple[Union[str, None], str]: + """Get a profile description and owner from ALL_CLOUD_AI_PROFILES. + + :param str profile_name: Name of the profile. + :param str owner: Owner of a shared profile. Defaults to current schema. + :return: Tuple containing the profile description and owner. :raises: ProfileNotFoundError """ with cursor() as cr: - cr.execute(GET_USER_AI_PROFILE, profile_name=profile_name.upper()) + cr.execute( + GET_ALL_AI_PROFILE, + profile_name=profile_name.upper(), + owner=owner.upper() if owner else None, + ) profile = cr.fetchone() if profile: if profile[1] is not None: - return profile[1].read() + description = profile[1].read() else: - return None + description = None + return description, profile[2] else: raise ProfileNotFoundError(profile_name) @staticmethod def _get_attributes( - profile_name, raise_on_empty: bool = False + profile_name, + owner: Optional[str] = None, + raise_on_empty: bool = False, ) -> Union[ProfileAttributes, None]: """Get AI profile attributes from the Database :param str profile_name: Name of the profile + :param str owner: Owner of a shared profile. Defaults to current schema. :param bool raise_on_empty: Raise an error if attributes are empty :return: select_ai.ProfileAttributes :raises: select_ai.errors.ProfileAttributesEmptyError """ with cursor() as cr: cr.execute( - GET_USER_AI_PROFILE_ATTRIBUTES, + GET_ALL_AI_PROFILE_ATTRIBUTES, profile_name=profile_name.upper(), + owner=owner.upper() if owner else None, ) attributes = cr.fetchall() if attributes: @@ -132,7 +148,10 @@ def get_attributes(self) -> ProfileAttributes: :return: select_ai.ProfileAttributes """ - return self._get_attributes(profile_name=self.profile_name) + return self._get_attributes( + profile_name=self.profile_name, + owner=self.owner, + ) def _set_attribute( self, @@ -271,6 +290,40 @@ def disable(self) -> None: keyword_parameters={"profile_name": self.profile_name}, ) + def grant_access(self, user_or_role_name: str) -> None: + """Grant a database user or role access to this AI profile. + + :param str user_or_role_name: Database user or role receiving access. + :return: None + :raises: oracledb.DatabaseError + """ + user_or_role_name = validate_user_or_role_name(user_or_role_name) + with cursor() as cr: + cr.callproc( + "DBMS_CLOUD_AI.GRANT_PROFILE_ACCESS", + keyword_parameters={ + "profile_name": self.profile_name, + "user_or_role_name": user_or_role_name, + }, + ) + + def revoke_access(self, user_or_role_name: str) -> None: + """Revoke a database user or role's access to this AI profile. + + :param str user_or_role_name: Database user or role losing access. + :return: None + :raises: oracledb.DatabaseError + """ + user_or_role_name = validate_user_or_role_name(user_or_role_name) + with cursor() as cr: + cr.callproc( + "DBMS_CLOUD_AI.REVOKE_PROFILE_ACCESS", + keyword_parameters={ + "profile_name": self.profile_name, + "user_or_role_name": user_or_role_name, + }, + ) + @classmethod def delete_profile(cls, profile_name: str, force: bool = False): """Class method to delete an AI profile from the database @@ -283,15 +336,22 @@ def delete_profile(cls, profile_name: str, force: bool = False): cls._delete(profile_name=profile_name, force=force) @classmethod - def fetch(cls, profile_name: str) -> "Profile": + def fetch( + cls, profile_name: str, owner: Optional[str] = None + ) -> "Profile": """Create a proxy Profile object from fetched attributes saved in the database :param str profile_name: The name of the AI profile + :param str owner: Owner of a shared profile. Defaults to current schema. :return: select_ai.Profile :raises: ProfileNotFoundError """ - return cls(profile_name, raise_error_if_exists=False) + return cls( + profile_name, + owner=owner, + raise_error_if_exists=False, + ) def _save_feedback( self, @@ -380,25 +440,30 @@ def delete_feedback( @classmethod def list( - cls, profile_name_pattern: str = ".*" + cls, + profile_name_pattern: str = ".*", + owner: Optional[str] = None, ) -> Generator["Profile", None, None]: """List AI Profiles saved in the database. :param str profile_name_pattern: Regular expressions can be used to specify a pattern. Function REGEXP_LIKE is used to perform the match. Default value is ".*" i.e. match all AI profiles. + :param str owner: Owner of shared profiles. Defaults to current schema. :return: Iterator[Profile] """ with cursor() as cr: cr.execute( - LIST_USER_AI_PROFILES, + LIST_ALL_AI_PROFILES, profile_name_pattern=profile_name_pattern, + owner=owner.upper() if owner else None, ) for row in cr.fetchall(): profile_name = row[0] yield cls( profile_name=profile_name, + owner=row[2], raise_error_if_exists=False, raise_error_on_empty_attributes=False, ) @@ -409,6 +474,7 @@ def _generate_with_cursor( prompt: str, action: Optional[Action] = Action.RUNSQL, params: Mapping = None, + attributes: Mapping = None, ) -> Union[pandas.DataFrame, str, None]: """Perform AI translation using this profile @@ -416,9 +482,13 @@ def _generate_with_cursor( :param select_ai.profile.Action action: :param params: Parameters to include in the LLM request. For e.g. conversation_id for context-aware chats + :param Mapping attributes: Profile attributes to override for this + request :return: Union[pandas.DataFrame, str] """ - parameters = self._generate_parameters(prompt, action, params) + parameters = self._generate_parameters( + prompt, action, params, attributes + ) data = cr.callfunc( "DBMS_CLOUD_AI.GENERATE", oracledb.DB_TYPE_CLOB, @@ -438,6 +508,7 @@ def _generate_parameters( prompt: str, action: Optional[Action], params: Mapping = None, + attributes: Mapping = None, ) -> Mapping: if not prompt: raise ValueError("prompt cannot be empty or None") @@ -445,10 +516,13 @@ def _generate_parameters( "prompt": prompt, "action": action, "profile_name": self.profile_name, - # "attributes": self.attributes.json(), } if params: parameters["params"] = json.dumps(params) + if attributes is not None: + if not isinstance(attributes, Mapping): + raise TypeError("'attributes' must be a mapping") + parameters["attributes"] = json.dumps(attributes) return parameters def _generate_stream( @@ -457,6 +531,7 @@ def _generate_stream( action: Optional[Action], params: Mapping = None, chunk_size: int = 8192, + attributes: Mapping = None, ) -> Generator[str, None, None]: with cursor() as cr: yield from self._generate_stream_with_cursor( @@ -465,6 +540,7 @@ def _generate_stream( action=action, params=params, chunk_size=chunk_size, + attributes=attributes, ) def _generate_stream_with_cursor( @@ -474,13 +550,16 @@ def _generate_stream_with_cursor( action: Optional[Action], params: Mapping = None, chunk_size: int = 8192, + attributes: Mapping = None, ) -> Generator[str, None, None]: if action == Action.RUNSQL: raise ValueError("stream=True is not supported for run_sql") if chunk_size <= 0: raise ValueError("chunk_size must be greater than 0") - parameters = self._generate_parameters(prompt, action, params) + parameters = self._generate_parameters( + prompt, action, params, attributes + ) data = cr.callfunc( "DBMS_CLOUD_AI.GENERATE", oracledb.DB_TYPE_CLOB, @@ -504,6 +583,8 @@ def generate( params: Mapping = None, stream: bool = False, chunk_size: int = 8192, + *, + attributes: Mapping = None, ) -> Union[pandas.DataFrame, str, Generator[str, None, None], None]: """Perform AI translation using this profile @@ -511,15 +592,23 @@ def generate( :param select_ai.profile.Action action: :param params: Parameters to include in the LLM request. For e.g. conversation_id for context-aware chats + :param Mapping attributes: Profile attributes to override for this + request :param bool stream: Return an iterator of response chunks :param int chunk_size: Number of characters to read per stream chunk :return: Union[pandas.DataFrame, str] """ if stream: - return self._generate_stream(prompt, action, params, chunk_size) + return self._generate_stream( + prompt, action, params, chunk_size, attributes + ) with cursor() as cr: return self._generate_with_cursor( - cr, prompt=prompt, action=action, params=params + cr, + prompt=prompt, + action=action, + params=params, + attributes=attributes, ) def chat( @@ -528,11 +617,15 @@ def chat( params: Mapping = None, stream: bool = False, chunk_size: int = 8192, + *, + attributes: Mapping = None, ) -> Union[str, Generator[str, None, None]]: """Chat with the LLM :param str prompt: Natural language prompt :param params: Parameters to include in the LLM request + :param Mapping attributes: Profile attributes to override for this + request :param bool stream: Return an iterator of response chunks :param int chunk_size: Number of characters to read per stream chunk :return: str @@ -543,6 +636,7 @@ def chat( params=params, stream=stream, chunk_size=chunk_size, + attributes=attributes, ) @contextmanager @@ -574,11 +668,15 @@ def narrate( params: Mapping = None, stream: bool = False, chunk_size: int = 8192, + *, + attributes: Mapping = None, ) -> Union[str, Generator[str, None, None]]: """Narrate the result of the SQL :param str prompt: Natural language prompt :param params: Parameters to include in the LLM request + :param Mapping attributes: Profile attributes to override for this + request :param bool stream: Return an iterator of response chunks :param int chunk_size: Number of characters to read per stream chunk :return: str @@ -589,6 +687,7 @@ def narrate( params=params, stream=stream, chunk_size=chunk_size, + attributes=attributes, ) def explain_sql( @@ -597,11 +696,15 @@ def explain_sql( params: Mapping = None, stream: bool = False, chunk_size: int = 8192, + *, + attributes: Mapping = None, ) -> Union[str, Generator[str, None, None]]: """Explain the generated SQL :param str prompt: Natural language prompt :param params: Parameters to include in the LLM request + :param Mapping attributes: Profile attributes to override for this + request :param bool stream: Return an iterator of response chunks :param int chunk_size: Number of characters to read per stream chunk :return: str @@ -612,17 +715,31 @@ def explain_sql( params=params, stream=stream, chunk_size=chunk_size, + attributes=attributes, ) - def run_sql(self, prompt: str, params: Mapping = None) -> pandas.DataFrame: + def run_sql( + self, + prompt: str, + params: Mapping = None, + *, + attributes: Mapping = None, + ) -> pandas.DataFrame: """Run the generate SQL statement and return a pandas Dataframe built using the result set :param str prompt: Natural language prompt :param params: Parameters to include in the LLM request + :param Mapping attributes: Profile attributes to override for this + request :return: pandas.DataFrame """ - return self.generate(prompt, action=Action.RUNSQL, params=params) + return self.generate( + prompt, + action=Action.RUNSQL, + params=params, + attributes=attributes, + ) def show_sql( self, @@ -630,11 +747,15 @@ def show_sql( params: Mapping = None, stream: bool = False, chunk_size: int = 8192, + *, + attributes: Mapping = None, ) -> Union[str, Generator[str, None, None]]: """Show the generated SQL :param str prompt: Natural language prompt :param params: Parameters to include in the LLM request + :param Mapping attributes: Profile attributes to override for this + request :param bool stream: Return an iterator of response chunks :param int chunk_size: Number of characters to read per stream chunk :return: str @@ -645,6 +766,7 @@ def show_sql( params=params, stream=stream, chunk_size=chunk_size, + attributes=attributes, ) def show_prompt( @@ -653,11 +775,15 @@ def show_prompt( params: Mapping = None, stream: bool = False, chunk_size: int = 8192, + *, + attributes: Mapping = None, ) -> Union[str, Generator[str, None, None]]: """Show the prompt sent to LLM :param str prompt: Natural language prompt :param params: Parameters to include in the LLM request + :param Mapping attributes: Profile attributes to override for this + request :param bool stream: Return an iterator of response chunks :param int chunk_size: Number of characters to read per stream chunk :return: str @@ -668,6 +794,7 @@ def show_prompt( params=params, stream=stream, chunk_size=chunk_size, + attributes=attributes, ) def summarize( @@ -789,7 +916,12 @@ def __init__(self, profile: Profile, params: Mapping): self._cursor = None def chat( - self, prompt: str, stream: bool = False, chunk_size: int = 8192 + self, + prompt: str, + stream: bool = False, + chunk_size: int = 8192, + *, + attributes: Mapping = None, ) -> Union[str, Generator[str, None, None]]: if stream: return self.profile._generate_stream_with_cursor( @@ -798,13 +930,23 @@ def chat( action=Action.CHAT, params=self.params, chunk_size=chunk_size, + attributes=attributes, ) return self.profile._generate_with_cursor( - self._cursor, prompt=prompt, action=Action.CHAT, params=self.params + self._cursor, + prompt=prompt, + action=Action.CHAT, + params=self.params, + attributes=attributes, ) def narrate( - self, prompt: str, stream: bool = False, chunk_size: int = 8192 + self, + prompt: str, + stream: bool = False, + chunk_size: int = 8192, + *, + attributes: Mapping = None, ) -> Union[str, Generator[str, None, None]]: """Narrate the result of the SQL @@ -820,13 +962,23 @@ def narrate( action=Action.NARRATE, params=self.params, chunk_size=chunk_size, + attributes=attributes, ) return self.profile._generate_with_cursor( - self._cursor, prompt, action=Action.NARRATE, params=self.params + self._cursor, + prompt, + action=Action.NARRATE, + params=self.params, + attributes=attributes, ) def explain_sql( - self, prompt: str, stream: bool = False, chunk_size: int = 8192 + self, + prompt: str, + stream: bool = False, + chunk_size: int = 8192, + *, + attributes: Mapping = None, ) -> Union[str, Generator[str, None, None]]: """Explain the generated SQL @@ -842,12 +994,19 @@ def explain_sql( action=Action.EXPLAINSQL, params=self.params, chunk_size=chunk_size, + attributes=attributes, ) return self.profile._generate_with_cursor( - self._cursor, prompt, action=Action.EXPLAINSQL, params=self.params + self._cursor, + prompt, + action=Action.EXPLAINSQL, + params=self.params, + attributes=attributes, ) - def run_sql(self, prompt: str) -> pandas.DataFrame: + def run_sql( + self, prompt: str, *, attributes: Mapping = None + ) -> pandas.DataFrame: """Run the generate SQL statement and return a pandas Dataframe built using the result set @@ -855,11 +1014,20 @@ def run_sql(self, prompt: str) -> pandas.DataFrame: :return: pandas.DataFrame """ return self.profile._generate_with_cursor( - self._cursor, prompt, action=Action.RUNSQL, params=self.params + self._cursor, + prompt, + action=Action.RUNSQL, + params=self.params, + attributes=attributes, ) def show_sql( - self, prompt: str, stream: bool = False, chunk_size: int = 8192 + self, + prompt: str, + stream: bool = False, + chunk_size: int = 8192, + *, + attributes: Mapping = None, ) -> Union[str, Generator[str, None, None]]: """Show the generated SQL @@ -876,13 +1044,23 @@ def show_sql( action=Action.SHOWSQL, params=self.params, chunk_size=chunk_size, + attributes=attributes, ) return self.profile._generate_with_cursor( - self._cursor, prompt, action=Action.SHOWSQL, params=self.params + self._cursor, + prompt, + action=Action.SHOWSQL, + params=self.params, + attributes=attributes, ) def show_prompt( - self, prompt: str, stream: bool = False, chunk_size: int = 8192 + self, + prompt: str, + stream: bool = False, + chunk_size: int = 8192, + *, + attributes: Mapping = None, ) -> Union[str, Generator[str, None, None]]: """Show the prompt sent to LLM @@ -899,9 +1077,14 @@ def show_prompt( action=Action.SHOWPROMPT, params=self.params, chunk_size=chunk_size, + attributes=attributes, ) return self.profile._generate_with_cursor( - self._cursor, prompt, action=Action.SHOWPROMPT, params=self.params + self._cursor, + prompt, + action=Action.SHOWPROMPT, + params=self.params, + attributes=attributes, ) def __enter__(self): diff --git a/src/select_ai/sql.py b/src/select_ai/sql.py index dbc91e7..3790938 100644 --- a/src/select_ai/sql.py +++ b/src/select_ai/sql.py @@ -73,41 +73,47 @@ END; """ -GET_USER_AI_PROFILE_ATTRIBUTES = """ +GET_ALL_AI_PROFILE_ATTRIBUTES = """ SELECT attribute_name, attribute_value -FROM USER_CLOUD_AI_PROFILE_ATTRIBUTES +FROM ALL_CLOUD_AI_PROFILE_ATTRIBUTES WHERE profile_name = :profile_name +AND owner = COALESCE(:owner, SYS_CONTEXT('USERENV', 'CURRENT_SCHEMA')) """ -GET_USER_AI_PROFILE = """ -SELECT profile_name, description -FROM USER_CLOUD_AI_PROFILES +GET_ALL_AI_PROFILE = """ +SELECT profile_name, description, owner +FROM ALL_CLOUD_AI_PROFILES WHERE profile_name = :profile_name +AND owner = COALESCE(:owner, SYS_CONTEXT('USERENV', 'CURRENT_SCHEMA')) """ -LIST_USER_AI_PROFILES = """ -SELECT profile_name, description -FROM USER_CLOUD_AI_PROFILES +LIST_ALL_AI_PROFILES = """ +SELECT profile_name, description, owner +FROM ALL_CLOUD_AI_PROFILES WHERE REGEXP_LIKE(profile_name, :profile_name_pattern, 'i') +AND owner = COALESCE(:owner, SYS_CONTEXT('USERENV', 'CURRENT_SCHEMA')) """ -LIST_USER_VECTOR_INDEXES = """ -SELECT v.index_name, v.description -FROM USER_CLOUD_VECTOR_INDEXES v +LIST_ALL_VECTOR_INDEXES = """ +SELECT v.index_name, v.description, v.owner +FROM ALL_CLOUD_VECTOR_INDEXES v WHERE REGEXP_LIKE(v.index_name, :index_name_pattern, 'i') +AND v.owner = COALESCE(:owner, SYS_CONTEXT('USERENV', 'CURRENT_SCHEMA')) """ -GET_USER_VECTOR_INDEX = """ -select index_name, description -from USER_CLOUD_VECTOR_INDEXES v -where index_name = :index_name +GET_ALL_VECTOR_INDEX = """ +SELECT index_name, description, owner +FROM ALL_CLOUD_VECTOR_INDEXES +WHERE index_name = :index_name +AND owner = COALESCE(:owner, SYS_CONTEXT('USERENV', 'CURRENT_SCHEMA')) """ -GET_USER_VECTOR_INDEX_ATTRIBUTES = """ +GET_ALL_VECTOR_INDEX_ATTRIBUTES = """ SELECT attribute_name, attribute_value -FROM USER_CLOUD_VECTOR_INDEX_ATTRIBUTES -WHERE INDEX_NAME = :index_name +FROM ALL_CLOUD_VECTOR_INDEX_ATTRIBUTES +WHERE index_name = :index_name +AND owner = COALESCE(:owner, SYS_CONTEXT('USERENV', 'CURRENT_SCHEMA')) """ LIST_USER_CONVERSATIONS = """ diff --git a/src/select_ai/vector_index.py b/src/select_ai/vector_index.py index 1685c67..51791ef 100644 --- a/src/select_ai/vector_index.py +++ b/src/select_ai/vector_index.py @@ -16,15 +16,16 @@ from select_ai import BaseProfile from select_ai._abc import SelectAIDataClass from select_ai._enums import StrEnum +from select_ai._validations import validate_user_or_role_name from select_ai.async_profile import AsyncProfile from select_ai.db import async_cursor, cursor from select_ai.errors import ProfileNotFoundError, VectorIndexNotFoundError from select_ai.profile import Profile from select_ai.sql import ( - GET_USER_VECTOR_INDEX, - GET_USER_VECTOR_INDEX_ATTRIBUTES, + GET_ALL_VECTOR_INDEX, + GET_ALL_VECTOR_INDEX_ATTRIBUTES, GET_VECTOR_PIPELINE_LAST_EXECUTION, - LIST_USER_VECTOR_INDEXES, + LIST_ALL_VECTOR_INDEXES, ) @@ -116,6 +117,7 @@ def __init__( index_name: Optional[str] = None, description: Optional[str] = None, attributes: Optional[VectorIndexAttributes] = None, + owner: Optional[str] = None, ): """Initialize a Vector Index""" if attributes and not isinstance(attributes, VectorIndexAttributes): @@ -132,6 +134,16 @@ def __init__( self.index_name = index_name self.attributes = attributes self.description = description + self.owner = owner.upper() if owner else None + + @property + def qualified_name(self) -> Optional[str]: + """Return the owner-qualified index name when owner is known.""" + if self.index_name is None: + return None + if self.owner is None: + return self.index_name + return f"{self.owner}.{self.index_name}" def __repr__(self): return ( @@ -151,17 +163,22 @@ class VectorIndex(_BaseVectorIndex): """ @staticmethod - def _get_attributes(index_name: str) -> VectorIndexAttributes: + def _get_attributes( + index_name: str, owner: Optional[str] = None + ) -> VectorIndexAttributes: """Get attributes of a vector index :return: select_ai.VectorIndexAttributes + :param str owner: Owner of a shared index. Defaults to current schema. :raises: VectorIndexNotFoundError """ if not index_name: raise AttributeError("'index_name' is required") with cursor() as cr: cr.execute( - GET_USER_VECTOR_INDEX_ATTRIBUTES, index_name=index_name.upper() + GET_ALL_VECTOR_INDEX_ATTRIBUTES, + index_name=index_name.upper(), + owner=owner.upper() if owner else None, ) attributes = cr.fetchall() if attributes: @@ -178,23 +195,31 @@ def _get_attributes(index_name: str) -> VectorIndexAttributes: raise VectorIndexNotFoundError(index_name=index_name) @staticmethod - def _get_description(index_name) -> Union[str, None]: - """Get description of the Vector Index from USER_CLOUD_VECTOR_INDEXES + def _get_description( + index_name: str, owner: Optional[str] = None + ) -> tuple[Union[str, None], str]: + """Get an index description and owner. :param str index_name: The name of the vector index - :return: Union[str, None] profile description + :param str owner: Owner of a shared index. Defaults to current schema. + :return: Tuple containing the index description and owner. :raises: ProfileNotFoundError """ if not index_name: raise AttributeError("'index_name' is required") with cursor() as cr: - cr.execute(GET_USER_VECTOR_INDEX, index_name=index_name.upper()) + cr.execute( + GET_ALL_VECTOR_INDEX, + index_name=index_name.upper(), + owner=owner.upper() if owner else None, + ) index = cr.fetchone() if index: if index[1] is not None: - return index[1].read() + description = index[1].read() else: - return None + description = None + return description, index[2] else: raise VectorIndexNotFoundError(index_name=index_name) @@ -333,19 +358,49 @@ def disable(self): else: raise + def grant_access(self, user_or_role_name: str) -> None: + """Grant a database user or role access to this vector index.""" + user_or_role_name = validate_user_or_role_name(user_or_role_name) + with cursor() as cr: + cr.callproc( + "DBMS_CLOUD_AI.GRANT_VECTOR_INDEX_ACCESS", + keyword_parameters={ + "index_name": self.index_name, + "user_or_role_name": user_or_role_name, + }, + ) + + def revoke_access(self, user_or_role_name: str) -> None: + """Revoke a database user or role's vector index access.""" + user_or_role_name = validate_user_or_role_name(user_or_role_name) + with cursor() as cr: + cr.callproc( + "DBMS_CLOUD_AI.REVOKE_VECTOR_INDEX_ACCESS", + keyword_parameters={ + "index_name": self.index_name, + "user_or_role_name": user_or_role_name, + }, + ) + @classmethod - def fetch(cls, index_name: str) -> "VectorIndex": + def fetch( + cls, index_name: str, owner: Optional[str] = None + ) -> "VectorIndex": """ Fetches vector index attributes from the database and builds a proxy object for the passed index_name :param str index_name: The name of the vector index + :param str owner: Owner of a shared index. Defaults to current schema. """ - description = cls._get_description(index_name) - attributes = cls._get_attributes(index_name) + description, saved_owner = cls._get_description(index_name, owner) + attributes = cls._get_attributes(index_name, saved_owner) try: - profile = Profile(profile_name=attributes.profile_name) + profile = Profile( + profile_name=attributes.profile_name, + owner=saved_owner, + ) except ProfileNotFoundError: profile = None return cls( @@ -353,6 +408,7 @@ def fetch(cls, index_name: str) -> "VectorIndex": attributes=attributes, profile=profile, index_name=index_name, + owner=saved_owner, ) def set_attribute( @@ -411,7 +467,7 @@ def get_attributes(self) -> VectorIndexAttributes: :return: select_ai.VectorIndexAttributes :raises: VectorIndexNotFoundError """ - return self._get_attributes(self.index_name) + return self._get_attributes(self.index_name, self.owner) def get_next_refresh_timestamp(self) -> Optional[datetime]: """ @@ -446,28 +502,40 @@ def get_profile(self) -> Profile: :return: select_ai.Profile :raises: ProfileNotFoundError """ - attributes = self._get_attributes(index_name=self.index_name) - profile = Profile(profile_name=attributes.profile_name) + attributes = self._get_attributes( + index_name=self.index_name, + owner=self.owner, + ) + profile = Profile( + profile_name=attributes.profile_name, + owner=self.owner, + ) return profile @classmethod - def list(cls, index_name_pattern: str = ".*") -> Iterator["VectorIndex"]: + def list( + cls, + index_name_pattern: str = ".*", + owner: Optional[str] = None, + ) -> Iterator["VectorIndex"]: """List Vector Indexes :param str index_name_pattern: Regular expressions can be used to specify a pattern. Function REGEXP_LIKE is used to perform the match. Default value is ".*" i.e. match all vector indexes. + :param str owner: Owner of shared indexes. Defaults to current schema. :return: Iterator[VectorIndex] """ with cursor() as cr: cr.execute( - LIST_USER_VECTOR_INDEXES, + LIST_ALL_VECTOR_INDEXES, index_name_pattern=index_name_pattern, + owner=owner.upper() if owner else None, ) for row in cr.fetchall(): index_name = row[0] - yield cls.fetch(index_name=index_name) + yield cls.fetch(index_name=index_name, owner=row[2]) class AsyncVectorIndex(_BaseVectorIndex): @@ -482,17 +550,22 @@ class AsyncVectorIndex(_BaseVectorIndex): """ @staticmethod - async def _get_attributes(index_name: str) -> VectorIndexAttributes: + async def _get_attributes( + index_name: str, owner: Optional[str] = None + ) -> VectorIndexAttributes: """Get attributes of a vector index :return: select_ai.VectorIndexAttributes + :param str owner: Owner of a shared index. Defaults to current schema. :raises: VectorIndexNotFoundError """ if not index_name: raise AttributeError("'index_name' is required") async with async_cursor() as cr: await cr.execute( - GET_USER_VECTOR_INDEX_ATTRIBUTES, index_name=index_name.upper() + GET_ALL_VECTOR_INDEX_ATTRIBUTES, + index_name=index_name.upper(), + owner=owner.upper() if owner else None, ) attributes = await cr.fetchall() if attributes: @@ -509,25 +582,31 @@ async def _get_attributes(index_name: str) -> VectorIndexAttributes: raise VectorIndexNotFoundError(index_name=index_name) @staticmethod - async def _get_description(index_name) -> Union[str, None]: - """Get description of the Vector Index from USER_CLOUD_VECTOR_INDEXES + async def _get_description( + index_name: str, owner: Optional[str] = None + ) -> tuple[Union[str, None], str]: + """Get an index description and owner. :param str index_name: The name of the vector index - :return: Union[str, None] profile description + :param str owner: Owner of a shared index. Defaults to current schema. + :return: Tuple containing the index description and owner. :raises: ProfileNotFoundError """ if not index_name: raise AttributeError("'index_name' is required") async with async_cursor() as cr: await cr.execute( - GET_USER_VECTOR_INDEX, index_name=index_name.upper() + GET_ALL_VECTOR_INDEX, + index_name=index_name.upper(), + owner=owner.upper() if owner else None, ) index = await cr.fetchone() if index: if index[1] is not None: - return await index[1].read() + description = await index[1].read() else: - return None + description = None + return description, index[2] else: raise VectorIndexNotFoundError(index_name=index_name) @@ -664,19 +743,51 @@ async def disable(self) -> None: else: raise + async def grant_access(self, user_or_role_name: str) -> None: + """Asynchronously grant a user or role access to this vector index.""" + user_or_role_name = validate_user_or_role_name(user_or_role_name) + async with async_cursor() as cr: + await cr.callproc( + "DBMS_CLOUD_AI.GRANT_VECTOR_INDEX_ACCESS", + keyword_parameters={ + "index_name": self.index_name, + "user_or_role_name": user_or_role_name, + }, + ) + + async def revoke_access(self, user_or_role_name: str) -> None: + """Asynchronously revoke a user or role's vector index access.""" + user_or_role_name = validate_user_or_role_name(user_or_role_name) + async with async_cursor() as cr: + await cr.callproc( + "DBMS_CLOUD_AI.REVOKE_VECTOR_INDEX_ACCESS", + keyword_parameters={ + "index_name": self.index_name, + "user_or_role_name": user_or_role_name, + }, + ) + @classmethod - async def fetch(cls, index_name: str) -> "AsyncVectorIndex": + async def fetch( + cls, index_name: str, owner: Optional[str] = None + ) -> "AsyncVectorIndex": """ Fetches vector index attributes from the database and builds a proxy object for the passed index_name :param str index_name: The name of the vector index + :param str owner: Owner of a shared index. Defaults to current schema. """ - description = await cls._get_description(index_name) - attributes = await cls._get_attributes(index_name) + description, saved_owner = await cls._get_description( + index_name, owner + ) + attributes = await cls._get_attributes(index_name, saved_owner) try: - profile = await AsyncProfile(profile_name=attributes.profile_name) + profile = await AsyncProfile( + profile_name=attributes.profile_name, + owner=saved_owner, + ) except ProfileNotFoundError: profile = None return cls( @@ -684,6 +795,7 @@ async def fetch(cls, index_name: str) -> "AsyncVectorIndex": attributes=attributes, profile=profile, index_name=index_name, + owner=saved_owner, ) async def set_attribute( @@ -737,7 +849,10 @@ async def get_attributes(self) -> VectorIndexAttributes: :return: select_ai.VectorIndexAttributes :raises: VectorIndexNotFoundError """ - return await self._get_attributes(index_name=self.index_name) + return await self._get_attributes( + index_name=self.index_name, + owner=self.owner, + ) async def get_next_refresh_timestamp(self) -> Optional[datetime]: """Return the UTC timestamp for the next scheduled refresh.""" @@ -770,30 +885,40 @@ async def get_profile(self) -> AsyncProfile: :return: select_ai.AsyncProfile :raises: ProfileNotFoundError """ - attributes = await self._get_attributes(index_name=self.index_name) - profile = await AsyncProfile(profile_name=attributes.profile_name) + attributes = await self._get_attributes( + index_name=self.index_name, + owner=self.owner, + ) + profile = await AsyncProfile( + profile_name=attributes.profile_name, + owner=self.owner, + ) return profile @classmethod async def list( - cls, index_name_pattern: str = ".*" + cls, + index_name_pattern: str = ".*", + owner: Optional[str] = None, ) -> AsyncGenerator["AsyncVectorIndex", None]: """List Vector Indexes. :param str index_name_pattern: Regular expressions can be used to specify a pattern. Function REGEXP_LIKE is used to perform the match. Default value is ".*" i.e. match all vector indexes. + :param str owner: Owner of shared indexes. Defaults to current schema. :return: AsyncGenerator[VectorIndex] """ async with async_cursor() as cr: await cr.execute( - LIST_USER_VECTOR_INDEXES, + LIST_ALL_VECTOR_INDEXES, index_name_pattern=index_name_pattern, + owner=owner.upper() if owner else None, ) rows = await cr.fetchall() for row in rows: index_name = row[0] - index = await cls.fetch(index_name=index_name) + index = await cls.fetch(index_name=index_name, owner=row[2]) yield index diff --git a/tests/conftest.py b/tests/conftest.py index 139dca1..de6f024 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -158,6 +158,57 @@ def test_env(pytestconfig): return env +@pytest.fixture(scope="session") +def sharing_user(test_env, setup_test_user): + """Create a real database user used by object-sharing E2E tests.""" + username = f"PYSAI_SHARE_{uuid.uuid4().hex[:12].upper()}" + password = f"PySai_{uuid.uuid4().hex}1aA" + with oracledb.connect(**test_env.connect_params(admin=True)) as conn: + with conn.cursor() as cur: + cur.execute(f'CREATE USER {username} IDENTIFIED BY "{password}"') + _grant_basic_schema_privileges(cur, username=username) + _grant_select_ai_privileges(cur, username=username) + _grant_http_access( + cur, + username=username, + provider_endpoint=select_ai.OpenAIProvider.provider_endpoint, + ) + conn.commit() + + yield { + "username": username, + "connect_params": { + **test_env.connect_params(), + "user": username, + "password": password, + }, + } + + with oracledb.connect(**test_env.connect_params(admin=True)) as conn: + with conn.cursor() as cur: + cur.execute(f"DROP USER {username} CASCADE") + conn.commit() + + +@pytest.fixture(scope="module") +def shared_credential_access(oci_credential, sharing_user): + """Grant the sharing user access to the profile's credential object.""" + credential_name = oci_credential["credential_name"].upper() + username = sharing_user["username"] + + select_ai.grant_credential_access( + credential_name=credential_name, + user_or_role_name=username, + ) + + yield credential_name + + select_ai.revoke_credential_access( + credential_name=credential_name, + user_or_role_name=username, + ) + + @pytest.fixture(autouse=True, scope="session") def setup_test_user(test_env): with oracledb.connect(**test_env.connect_params(admin=True)) as conn: @@ -169,6 +220,8 @@ def setup_test_user(test_env): password=test_env.test_user_password, ) _grant_basic_schema_privileges(cur, username=test_env.test_user) + cur.execute(f"GRANT CREATE PUBLIC SYNONYM TO {test_env.test_user}") + cur.execute(f"GRANT DROP PUBLIC SYNONYM TO {test_env.test_user}") _grant_select_ai_privileges(cur, username=test_env.test_user) _grant_http_access( cur, @@ -227,9 +280,16 @@ def oci_credential(connect, test_env): "private_key": get_env_value("OCI_PRIVATE_KEY", required=True), "fingerprint": get_env_value("OCI_FINGERPRINT", required=True), } - select_ai.create_credential(credential, replace=True) + select_ai.create_credential( + credential, + replace=True, + public_synonym=True, + ) yield credential - select_ai.delete_credential(PYSAI_OCI_CREDENTIAL_NAME) + select_ai.delete_credential( + PYSAI_OCI_CREDENTIAL_NAME, + public_synonym=True, + ) @pytest.fixture(scope="module") diff --git a/tests/profiles/test_1200_profile.py b/tests/profiles/test_1200_profile.py index cfaab13..d07424c 100644 --- a/tests/profiles/test_1200_profile.py +++ b/tests/profiles/test_1200_profile.py @@ -404,3 +404,24 @@ def test_1219_profile_status(python_gen_ai_profile, cursor): profile_name=python_gen_ai_profile.profile_name, ) assert cursor.fetchone()[0] == "ENABLED" + + +def test_1221_shared_profile_credential_access( + python_gen_ai_profile, shared_credential_access, sharing_user +): + username = sharing_user["username"] + python_gen_ai_profile.grant_access(username) + try: + with oracledb.connect(**sharing_user["connect_params"]) as conn: + with conn.cursor() as cr: + cr.execute( + """ + SELECT COUNT(*) + FROM ALL_CREDENTIALS + WHERE credential_name = :credential_name + """, + credential_name=shared_credential_access, + ) + assert cr.fetchone()[0] == 1 + finally: + python_gen_ai_profile.revoke_access(username) diff --git a/tests/profiles/test_1400_profile_sharing.py b/tests/profiles/test_1400_profile_sharing.py new file mode 100644 index 0000000..734a069 --- /dev/null +++ b/tests/profiles/test_1400_profile_sharing.py @@ -0,0 +1,127 @@ +# ----------------------------------------------------------------------------- +# Copyright (c) 2026, Oracle and/or its affiliates. +# +# Licensed under the Universal Permissive License v 1.0 as shown at +# http://oss.oracle.com/licenses/upl. +# ----------------------------------------------------------------------------- + +import uuid + +import pytest +import select_ai +from select_ai import AsyncProfile, Profile +from select_ai.errors import ProfileNotFoundError + +PROFILE_NAME = f"PYSAI_1400_{uuid.uuid4().hex.upper()}" + + +@pytest.fixture(scope="module") +def shared_profile(profile_attributes): + profile = Profile( + profile_name=PROFILE_NAME, + description="Profile sharing test", + attributes=profile_attributes, + ) + yield profile + profile.delete(force=True) + + +def test_1401_fetch_and_list_shared_profile( + shared_profile, + shared_credential_access, + sharing_user, + test_env, +): + profile_name = shared_profile.profile_name.upper() + owner = test_env.test_user.upper() + username = sharing_user["username"] + + def fetch_and_list_as_sharing_user(): + try: + select_ai.disconnect() + select_ai.connect(**sharing_user["connect_params"]) + fetched = Profile.fetch(profile_name, owner=owner) + listed = list(Profile.list(f"^{profile_name}$", owner=owner)) + return fetched, listed + finally: + select_ai.disconnect() + select_ai.create_pool(**test_env.connect_params(use_pool=True)) + + shared_profile.grant_access(username) + try: + fetched, listed = fetch_and_list_as_sharing_user() + assert fetched.profile_name == profile_name + assert fetched.owner == owner + assert [profile.profile_name for profile in listed] == [profile_name] + assert [profile.owner for profile in listed] == [owner] + finally: + shared_profile.revoke_access(username) + + try: + select_ai.disconnect() + select_ai.connect(**sharing_user["connect_params"]) + with pytest.raises(ProfileNotFoundError): + Profile.fetch(profile_name, owner=owner) + assert not list(Profile.list(f"^{profile_name}$", owner=owner)) + finally: + select_ai.disconnect() + select_ai.create_pool(**test_env.connect_params(use_pool=True)) + + +@pytest.mark.anyio +async def test_1402_async_fetch_and_list_shared_profile( + shared_profile, + shared_credential_access, + sharing_user, + test_env, +): + profile_name = shared_profile.profile_name.upper() + owner = test_env.test_user.upper() + username = sharing_user["username"] + owner_profile = await AsyncProfile.fetch(profile_name) + + async def fetch_and_list_as_sharing_user(): + try: + await select_ai.async_disconnect() + await select_ai.async_connect(**sharing_user["connect_params"]) + fetched = await AsyncProfile.fetch(profile_name, owner=owner) + listed = [ + profile + async for profile in AsyncProfile.list( + f"^{profile_name}$", + owner=owner, + ) + ] + return fetched, listed + finally: + await select_ai.async_disconnect() + select_ai.create_pool_async( + **test_env.connect_params(use_pool=True) + ) + + await owner_profile.grant_access(username) + try: + fetched, listed = await fetch_and_list_as_sharing_user() + assert fetched.profile_name == profile_name + assert fetched.owner == owner + assert [profile.profile_name for profile in listed] == [profile_name] + assert [profile.owner for profile in listed] == [owner] + finally: + await owner_profile.revoke_access(username) + + try: + await select_ai.async_disconnect() + await select_ai.async_connect(**sharing_user["connect_params"]) + with pytest.raises(ProfileNotFoundError): + await AsyncProfile.fetch(profile_name, owner=owner) + listed = [ + profile + async for profile in AsyncProfile.list( + f"^{profile_name}$", + owner=owner, + ) + ] + assert not listed + finally: + await select_ai.async_disconnect() + select_ai.create_pool_async(**test_env.connect_params(use_pool=True)) diff --git a/tests/profiles/test_1600_generate.py b/tests/profiles/test_1600_generate.py index 52e846b..6edb2d4 100644 --- a/tests/profiles/test_1600_generate.py +++ b/tests/profiles/test_1600_generate.py @@ -284,6 +284,49 @@ def test_1616_chat_stream(generate_profile): assert "Oracle Cloud Infrastructure" in response +def test_1620_generate_with_request_attributes(generate_profile): + """Request attributes are passed through to the database generate call""" + instruction = "Include the exact marker PYSAI_REQUEST_ATTRIBUTES_E2E in the response." + + show_prompt = generate_profile.show_prompt( + prompt="Tell me about OCI", + attributes={"additional_instructions": instruction}, + ) + + assert instruction in show_prompt + + +@pytest.fixture +def profile_with_additional_instructions(oci_credential, generate_provider): + instruction = "Include the exact marker PYSAI_PROFILE_ATTRIBUTES_E2E in the response." + profile_name = f"{PROFILE_PREFIX}_ATTR_{uuid.uuid4().hex.upper()}" + profile = Profile( + profile_name=profile_name, + attributes=ProfileAttributes( + credential_name=oci_credential["credential_name"], + provider=generate_provider, + additional_instructions=instruction, + ), + description="Generate profile attributes E2E test profile", + replace=True, + ) + yield profile, instruction + profile.delete(force=True) + + +def test_1621_profile_attributes_reach_database( + profile_with_additional_instructions, +): + """Persistent profile attributes are stored and used by the database""" + profile, instruction = profile_with_additional_instructions + + fetched_attributes = profile.get_attributes() + assert fetched_attributes.additional_instructions == instruction + + show_prompt = profile.show_prompt(prompt="Tell me about OCI") + assert instruction in show_prompt + + def test_1616_empty_prompt_raises_value_error(negative_profile): """Empty prompts raise ValueError for profile methods""" logger.info( diff --git a/tests/profiles/test_1700_generate_async.py b/tests/profiles/test_1700_generate_async.py index 90c47da..7828996 100644 --- a/tests/profiles/test_1700_generate_async.py +++ b/tests/profiles/test_1700_generate_async.py @@ -318,6 +318,53 @@ async def test_1716_chat_stream(async_generate_profile): assert "Oracle Cloud Infrastructure" in response +@pytest.mark.anyio +async def test_1720_generate_with_request_attributes(async_generate_profile): + """Request attributes are passed through to the database generate call""" + instruction = "Include the exact marker PYSAI_REQUEST_ATTRIBUTES_E2E in the response." + + show_prompt = await async_generate_profile.show_prompt( + prompt="Tell me about OCI", + attributes={"additional_instructions": instruction}, + ) + + assert instruction in show_prompt + + +@pytest.fixture +async def async_profile_with_additional_instructions( + oci_credential, async_generate_provider +): + instruction = "Include the exact marker PYSAI_PROFILE_ATTRIBUTES_E2E in the response." + profile_name = f"{PROFILE_PREFIX}_ATTR_{uuid.uuid4().hex.upper()}" + profile = await AsyncProfile( + profile_name=profile_name, + attributes=ProfileAttributes( + credential_name=oci_credential["credential_name"], + provider=async_generate_provider, + additional_instructions=instruction, + ), + description="Async generate profile attributes E2E test profile", + replace=True, + ) + yield profile, instruction + await profile.delete(force=True) + + +@pytest.mark.anyio +async def test_1721_profile_attributes_reach_database( + async_profile_with_additional_instructions, +): + """Persistent profile attributes are stored and used by the database""" + profile, instruction = async_profile_with_additional_instructions + + fetched_attributes = await profile.get_attributes() + assert fetched_attributes.additional_instructions == instruction + + show_prompt = await profile.show_prompt(prompt="Tell me about OCI") + assert instruction in show_prompt + + @pytest.mark.anyio async def test_1716_empty_prompt_raises_value_error(async_negative_profile): """Empty prompts raise ValueError for async profile methods""" diff --git a/tests/vector_index/test_5000_async_create_index.py b/tests/vector_index/test_5000_async_create_index.py index 796d912..3249d9c 100644 --- a/tests/vector_index/test_5000_async_create_index.py +++ b/tests/vector_index/test_5000_async_create_index.py @@ -453,3 +453,38 @@ async def test_5018(self): for _ in range(10): await self.async_vector_index.create(replace=True) logger.info("Successfully recreated vector index multiple times.") + + async def test_5019_grant_access(self, sharing_user, test_env): + username = sharing_user["username"] + owner = test_env.test_user.upper() + + async def fetch_and_list_as_sharing_user(): + try: + await select_ai.async_disconnect() + await select_ai.async_connect(**sharing_user["connect_params"]) + fetched = await select_ai.AsyncVectorIndex.fetch( + self.index_name, + owner=owner, + ) + listed = [ + index + async for index in select_ai.AsyncVectorIndex.list( + self.index_name, + owner=owner, + ) + ] + return fetched, listed + finally: + await select_ai.async_disconnect() + select_ai.create_pool_async( + **test_env.connect_params(use_pool=True) + ) + + await self.async_vector_index.create(replace=True) + await self.async_vector_index.grant_access(username) + fetched, listed = await fetch_and_list_as_sharing_user() + assert fetched.index_name == self.index_name + assert fetched.owner == owner + assert [index.index_name for index in listed] == [ + self.index_name.upper() + ] diff --git a/tests/vector_index/test_5000_create_index.py b/tests/vector_index/test_5000_create_index.py index e5f75e7..40a960b 100644 --- a/tests/vector_index/test_5000_create_index.py +++ b/tests/vector_index/test_5000_create_index.py @@ -479,3 +479,35 @@ def test_5018(self): for _ in range(10): self.vector_index.create(replace=True) logger.info("Successfully recreated vector index multiple times.") + + def test_5019_grant_access(self, sharing_user, test_env): + username = sharing_user["username"] + owner = test_env.test_user.upper() + + def fetch_and_list_as_sharing_user(): + try: + select_ai.disconnect() + select_ai.connect(**sharing_user["connect_params"]) + fetched = select_ai.VectorIndex.fetch( + self.index_name, + owner=owner, + ) + listed = list( + select_ai.VectorIndex.list( + self.index_name, + owner=owner, + ) + ) + return fetched, listed + finally: + select_ai.disconnect() + select_ai.create_pool(**test_env.connect_params(use_pool=True)) + + self.vector_index.create(replace=True) + self.vector_index.grant_access(username) + fetched, listed = fetch_and_list_as_sharing_user() + assert fetched.index_name == self.index_name + assert fetched.owner == owner + assert [index.index_name for index in listed] == [ + self.index_name.upper() + ]