select-ai 1.0.0.dev4__py3-none-any.whl

select_ai/async_profile.py

@@ -0,0 +1,468 @@
+# -----------------------------------------------------------------------------
+# Copyright (c) 2025, Oracle and/or its affiliates.
+#
+# Licensed under the Universal Permissive License v 1.0 as shown at
+# http://oss.oracle.com/licenses/upl.
+# -----------------------------------------------------------------------------
+
+import json
+from contextlib import asynccontextmanager
+from dataclasses import replace as dataclass_replace
+from typing import (
+    AsyncGenerator,
+    List,
+    Mapping,
+    Optional,
+    Tuple,
+    Union,
+)
+
+import oracledb
+import pandas
+
+from select_ai.action import Action
+from select_ai.base_profile import BaseProfile, ProfileAttributes
+from select_ai.conversation import AsyncConversation
+from select_ai.db import async_cursor, async_get_connection
+from select_ai.errors import ProfileNotFoundError
+from select_ai.provider import Provider
+from select_ai.sql import (
+    GET_USER_AI_PROFILE_ATTRIBUTES,
+    LIST_USER_AI_PROFILES,
+)
+from select_ai.synthetic_data import SyntheticDataAttributes
+
+__all__ = ["AsyncProfile"]
+
+
+class AsyncProfile(BaseProfile):
+    """AsyncProfile defines methods to interact with the underlying AI Provider
+    asynchronously.
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._init_coroutine = self._init_profile()
+
+    def __await__(self):
+        coroutine = self._init_coroutine
+        return coroutine.__await__()
+
+    async def _init_profile(self):
+        """Initializes AI profile based on the passed attributes
+
+        :return: None
+        :raises: oracledb.DatabaseError
+        """
+        if self.profile_name is not None:
+            profile_exists = False
+            try:
+                saved_attributes = await self._get_attributes(
+                    profile_name=self.profile_name
+                )
+                profile_exists = True
+            except ProfileNotFoundError:
+                if self.attributes is None:
+                    raise
+            else:
+                if self.attributes is None:
+                    self.attributes = saved_attributes
+                if self.merge:
+                    self.replace = True
+                    if self.attributes is not None:
+                        self.attributes = dataclass_replace(
+                            saved_attributes,
+                            **self.attributes.dict(exclude_null=True),
+                        )
+            if self.replace or not profile_exists:
+                await self.create(
+                    replace=self.replace, description=self.description
+                )
+        return self
+
+    @staticmethod
+    async def _get_attributes(profile_name) -> ProfileAttributes:
+        """Asynchronously gets AI profile attributes from the Database
+
+        :param str profile_name: Name of the profile
+        :return: select_ai.provider.ProviderAttributes
+        :raises: ProfileNotFoundError
+
+        """
+        async with async_cursor() as cr:
+            await cr.execute(
+                GET_USER_AI_PROFILE_ATTRIBUTES,
+                profile_name=profile_name.upper(),
+            )
+            attributes = await cr.fetchall()
+            if attributes:
+                return await ProfileAttributes.async_create(**dict(attributes))
+            else:
+                raise ProfileNotFoundError(profile_name=profile_name)
+
+    async def get_attributes(self) -> ProfileAttributes:
+        """Asynchronously gets AI profile attributes from the Database
+
+        :return: select_ai.provider.ProviderAttributes
+        :raises: ProfileNotFoundError
+        """
+        return await self._get_attributes(profile_name=self.profile_name)
+
+    async def _set_attribute(
+        self,
+        attribute_name: str,
+        attribute_value: Union[bool, str, int, float],
+    ):
+        parameters = {
+            "profile_name": self.profile_name,
+            "attribute_name": attribute_name,
+            "attribute_value": attribute_value,
+        }
+        async with async_cursor() as cr:
+            await cr.callproc(
+                "DBMS_CLOUD_AI.SET_ATTRIBUTE", keyword_parameters=parameters
+            )
+
+    async def set_attribute(
+        self,
+        attribute_name: str,
+        attribute_value: Union[bool, str, int, float, Provider],
+    ):
+        """Updates AI profile attribute on the Python object and also
+        saves it in the database
+
+        :param str attribute_name: Name of the AI profile attribute
+        :param Union[bool, str, int, float] attribute_value: Value of the
+         profile attribute
+        :return: None
+
+        """
+        self.attributes.set_attribute(attribute_name, attribute_value)
+        if isinstance(attribute_value, Provider):
+            for k, v in attribute_value.dict().items():
+                await self._set_attribute(k, v)
+        else:
+            await self._set_attribute(attribute_name, attribute_value)
+
+    async def set_attributes(self, attributes: ProfileAttributes):
+        """Updates AI profile attributes on the Python object and also
+        saves it in the database
+
+        :param ProfileAttributes attributes: Object specifying AI profile
+         attributes
+        :return: None
+        """
+        self.attributes = attributes
+        parameters = {
+            "profile_name": self.profile_name,
+            "attributes": self.attributes.json(),
+        }
+        async with async_cursor() as cr:
+            await cr.callproc(
+                "DBMS_CLOUD_AI.SET_ATTRIBUTES", keyword_parameters=parameters
+            )
+
+    async def create(
+        self, replace: Optional[int] = False, description: Optional[str] = None
+    ) -> None:
+        """Asynchronously create an AI Profile in the Database
+
+        :param bool replace: Set True to replace else False
+        :param description: The profile description
+        :return: None
+        :raises: oracledb.DatabaseError
+        """
+        parameters = {
+            "profile_name": self.profile_name,
+            "attributes": self.attributes.json(),
+        }
+        if description:
+            parameters["description"] = description
+
+        async with async_cursor() as cr:
+            try:
+                await cr.callproc(
+                    "DBMS_CLOUD_AI.CREATE_PROFILE",
+                    keyword_parameters=parameters,
+                )
+            except oracledb.DatabaseError as e:
+                (error,) = e.args
+                # If already exists and replace is True then drop and recreate
+                if "already exists" in error.message.lower() and replace:
+                    await self.delete(force=True)
+                    await cr.callproc(
+                        "DBMS_CLOUD_AI.CREATE_PROFILE",
+                        keyword_parameters=parameters,
+                    )
+                else:
+                    raise
+
+    async def delete(self, force=False) -> None:
+        """Asynchronously deletes an AI profile from the database
+
+        :param bool force: Ignores errors if AI profile does not exist.
+        :return: None
+        :raises: oracledb.DatabaseError
+
+        """
+        async with async_cursor() as cr:
+            await cr.callproc(
+                "DBMS_CLOUD_AI.DROP_PROFILE",
+                keyword_parameters={
+                    "profile_name": self.profile_name,
+                    "force": force,
+                },
+            )
+
+    @classmethod
+    async def _from_db(cls, profile_name: str) -> "AsyncProfile":
+        """Asynchronously create an AI Profile object from attributes
+        saved in the database against the profile
+
+        :param str profile_name:
+        :return: select_ai.Profile
+        :raises: ProfileNotFoundError
+        """
+        async with async_cursor() as cr:
+            await cr.execute(
+                GET_USER_AI_PROFILE_ATTRIBUTES, profile_name=profile_name
+            )
+            attributes = await cr.fetchall()
+            if attributes:
+                attributes = await ProfileAttributes.async_create(
+                    **dict(attributes)
+                )
+                return cls(profile_name=profile_name, attributes=attributes)
+            else:
+                raise ProfileNotFoundError(profile_name=profile_name)
+
+    @classmethod
+    async def list(
+        cls, profile_name_pattern: str
+    ) -> 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
+
+        :return: Iterator[Profile]
+        """
+        async with async_cursor() as cr:
+            await cr.execute(
+                LIST_USER_AI_PROFILES,
+                profile_name_pattern=profile_name_pattern,
+            )
+            rows = await cr.fetchall()
+            for row in rows:
+                profile_name = row[0]
+                description = row[1]
+                attributes = await cls._get_attributes(
+                    profile_name=profile_name
+                )
+                yield cls(
+                    profile_name=profile_name,
+                    description=description,
+                    attributes=attributes,
+                )
+
+    async def generate(
+        self, prompt, action=Action.SHOWSQL, params: Mapping = None
+    ) -> Union[pandas.DataFrame, str, None]:
+        """Asynchronously perform AI translation using this profile
+
+        :param str prompt: Natural language prompt to translate
+        :param select_ai.profile.Action action:
+        :param params: Parameters to include in the LLM request. For e.g.
+         conversation_id for context-aware chats
+        :return: Union[pandas.DataFrame, str]
+        """
+        parameters = {
+            "prompt": prompt,
+            "action": action,
+            "profile_name": self.profile_name,
+            # "attributes": self.attributes.json(),
+        }
+        if params:
+            parameters["params"] = json.dumps(params)
+
+        async with async_cursor() as cr:
+            data = await cr.callfunc(
+                "DBMS_CLOUD_AI.GENERATE",
+                oracledb.DB_TYPE_CLOB,
+                keyword_parameters=parameters,
+            )
+        if data is not None:
+            return await data.read()
+        return None
+
+    async def chat(self, prompt, params: Mapping = None) -> str:
+        """Asynchronously chat with the LLM
+
+        :param str prompt: Natural language prompt
+        :param params: Parameters to include in the LLM request
+        :return: str
+        """
+        return await self.generate(prompt, action=Action.CHAT, params=params)
+
+    @asynccontextmanager
+    async def chat_session(
+        self, conversation: AsyncConversation, delete: bool = False
+    ):
+        """Starts a new chat session for context-aware conversations
+
+        :param AsyncConversation conversation: Conversation object to use for this
+         chat session
+        :param bool delete: Delete conversation after session ends
+
+        """
+        try:
+            if (
+                conversation.conversation_id is None
+                and conversation.attributes is not None
+            ):
+                await conversation.create()
+            params = {"conversation_id": conversation.conversation_id}
+            async_session = AsyncSession(async_profile=self, params=params)
+            yield async_session
+        finally:
+            if delete:
+                await conversation.delete()
+
+    async def narrate(self, prompt, params: Mapping = None) -> str:
+        """Narrate the result of the SQL
+
+        :param str prompt: Natural language prompt
+        :param params: Parameters to include in the LLM request
+        :return: str
+        """
+        return await self.generate(
+            prompt, action=Action.NARRATE, params=params
+        )
+
+    async def explain_sql(self, prompt: str, params: Mapping = None):
+        """Explain the generated SQL
+
+        :param str prompt: Natural language prompt
+        :param params: Parameters to include in the LLM request
+        :return: str
+        """
+        return await self.generate(
+            prompt, action=Action.EXPLAINSQL, params=params
+        )
+
+    async def run_sql(
+        self, prompt, params: Mapping = None
+    ) -> pandas.DataFrame:
+        """Explain the generated SQL
+
+        :param str prompt: Natural language prompt
+        :param params: Parameters to include in the LLM request
+        :return: pandas.DataFrame
+        """
+        data = await self.generate(prompt, action=Action.RUNSQL, params=params)
+        return pandas.DataFrame(json.loads(data))
+
+    async def show_sql(self, prompt, params: Mapping = None):
+        """Show the generated SQL
+
+        :param str prompt: Natural language prompt
+        :param params: Parameters to include in the LLM request
+        :return: str
+        """
+        return await self.generate(
+            prompt, action=Action.SHOWSQL, params=params
+        )
+
+    async def show_prompt(self, prompt: str, params: Mapping = None):
+        """Show the prompt sent to LLM
+
+        :param str prompt: Natural language prompt
+        :param params: Parameters to include in the LLM request
+        :return: str
+        """
+        return await self.generate(
+            prompt, action=Action.SHOWPROMPT, params=params
+        )
+
+    async def generate_synthetic_data(
+        self, synthetic_data_attributes: SyntheticDataAttributes
+    ) -> None:
+        """Generate synthetic data for a single table, multiple tables or a
+        full schema.
+
+        :param select_ai.SyntheticDataAttributes synthetic_data_attributes:
+        :return: None
+        :raises: oracledb.DatabaseError
+
+        """
+        keyword_parameters = synthetic_data_attributes.prepare()
+        keyword_parameters["profile_name"] = self.profile_name
+        async with async_cursor() as cr:
+            await cr.callproc(
+                "DBMS_CLOUD_AI.GENERATE_SYNTHETIC_DATA",
+                keyword_parameters=keyword_parameters,
+            )
+
+    async def run_pipeline(
+        self,
+        prompt_specifications: List[Tuple[str, Action]],
+        continue_on_error: bool = False,
+    ) -> List[Union[str, pandas.DataFrame]]:
+        """Send Multiple prompts in a single roundtrip to the Database
+
+        :param List[Tuple[str, Action]] prompt_specifications: List of
+         2-element tuples. First element is the prompt and second is the
+         corresponding action
+
+        :param bool continue_on_error: True to continue on error else False
+        :return: List[Union[str, pandas.DataFrame]]
+        """
+        pipeline = oracledb.create_pipeline()
+        for prompt, action in prompt_specifications:
+            parameters = {
+                "prompt": prompt,
+                "action": action,
+                "profile_name": self.profile_name,
+                # "attributes": self.attributes.json(),
+            }
+            pipeline.add_callfunc(
+                "DBMS_CLOUD_AI.GENERATE",
+                return_type=oracledb.DB_TYPE_CLOB,
+                keyword_parameters=parameters,
+            )
+        async_connection = await async_get_connection()
+        pipeline_results = await async_connection.run_pipeline(
+            pipeline, continue_on_error=continue_on_error
+        )
+        responses = []
+        for result in pipeline_results:
+            if not result.error:
+                responses.append(await result.return_value.read())
+            else:
+                responses.append(result.error)
+        return responses
+
+
+class AsyncSession:
+    """AsyncSession lets you persist request parameters across DBMS_CLOUD_AI
+    requests. This is useful in context-aware conversations
+    """
+
+    def __init__(self, async_profile: AsyncProfile, params: Mapping):
+        """
+
+        :param async_profile: An AI Profile to use in this session
+        :param params: Parameters to be persisted across requests
+        """
+        self.params = params
+        self.async_profile = async_profile
+
+    async def chat(self, prompt: str):
+        return await self.async_profile.chat(prompt=prompt, params=self.params)
+
+    async def __aenter__(self):
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        pass

select_ai/base_profile.py

@@ -0,0 +1,166 @@
+# -----------------------------------------------------------------------------
+# Copyright (c) 2025, Oracle and/or its affiliates.
+#
+# Licensed under the Universal Permissive License v 1.0 as shown at
+# http://oss.oracle.com/licenses/upl.
+# -----------------------------------------------------------------------------
+
+import json
+from abc import ABC
+from dataclasses import dataclass
+from typing import List, Mapping, Optional
+
+import oracledb
+
+from select_ai._abc import SelectAIDataClass
+
+from .provider import Provider
+
+
+@dataclass
+class ProfileAttributes(SelectAIDataClass):
+    """
+    Use this class to define attributes to manage and configure the behavior of
+    an AI profile
+
+    :param bool comments: True to include column comments in the metadata used
+     for generating SQL queries from natural language prompts.
+    :param bool constraints: True to include referential integrity constraints
+     such as primary and foreign keys in the metadata sent to the LLM.
+    :param bool conversation: Indicates if conversation history is enabled for
+     a profile.
+    :param str credential_name: The name of the credential to access the AI
+     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 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
+     the owner and object names that are eligible for natural language
+     translation to SQL.
+    :param str object_list_mode: Specifies whether to send metadata for the
+     most relevant tables or all tables to the LLM. Supported values are -
+     'automated' and 'all'
+    :param select_ai.Provider provider: AI Provider
+    :param str stop_tokens: The generated text will be terminated at the
+     beginning of the earliest stop sequence. Sequence will be incorporated
+     into the text. The attribute value must be a valid array of string values
+     in JSON format
+    :param float temperature: Temperature is a non-negative float number used
+     to tune the degree of randomness. Lower temperatures mean less random
+     generations.
+    :param str vector_index_name: Name of the vector index
+
+    """
+
+    annotations: Optional[str] = None
+    case_sensitive_values: Optional[bool] = None
+    comments: Optional[bool] = None
+    constraints: Optional[str] = None
+    conversation: Optional[bool] = None
+    credential_name: Optional[str] = None
+    enable_sources: Optional[bool] = None
+    enable_source_offsets: Optional[bool] = None
+    enforce_object_list: Optional[bool] = None
+    max_tokens: Optional[int] = 1024
+    object_list: Optional[List[Mapping]] = None
+    object_list_mode: Optional[str] = None
+    provider: Optional[Provider] = None
+    seed: Optional[str] = None
+    stop_tokens: Optional[str] = None
+    streaming: Optional[str] = None
+    temperature: Optional[float] = None
+    vector_index_name: Optional[str] = None
+
+    def json(self, exclude_null=True):
+        attributes = {}
+        for k, v in self.dict(exclude_null=exclude_null).items():
+            if isinstance(v, Provider):
+                for provider_k, provider_v in v.dict(
+                    exclude_null=exclude_null
+                ).items():
+                    attributes[Provider.key_alias(provider_k)] = provider_v
+            else:
+                attributes[k] = v
+        return json.dumps(attributes)
+
+    @classmethod
+    def create(cls, **kwargs):
+        provider_attributes = {}
+        profile_attributes = {}
+        for k, v in kwargs.items():
+            if isinstance(v, oracledb.LOB):
+                v = v.read()
+            if k in Provider.keys():
+                provider_attributes[Provider.key_alias(k)] = v
+            else:
+                profile_attributes[k] = v
+        provider = Provider.create(**provider_attributes)
+        profile_attributes["provider"] = provider
+        return ProfileAttributes(**profile_attributes)
+
+    @classmethod
+    async def async_create(cls, **kwargs):
+        provider_attributes = {}
+        profile_attributes = {}
+        for k, v in kwargs.items():
+            if isinstance(v, oracledb.AsyncLOB):
+                v = await v.read()
+                if k in Provider.keys():
+                    provider_attributes[Provider.key_alias(k)] = v
+                else:
+                    profile_attributes[k] = v
+        provider = Provider.create(**provider_attributes)
+        profile_attributes["provider"] = provider
+        return ProfileAttributes(**profile_attributes)
+
+    def set_attribute(self, key, value):
+        if key in Provider.keys() and not isinstance(value, Provider):
+            setattr(self.provider, key, value)
+        else:
+            setattr(self, key, value)
+
+
+class BaseProfile(ABC):
+    """
+    BaseProfile is an abstract base class representing a Profile
+    for Select AI's interactions with AI service providers (LLMs).
+    Use either select_ai.Profile or select_ai.AsyncProfile to
+    instantiate an AI profile object.
+
+    :param str profile_name : Name of the profile
+
+    :param select_ai.ProfileAttributes attributes:
+     Object specifying AI profile attributes
+
+    :param str description: Description of the profile
+
+    :param bool merge: Fetches the profile
+     from database, merges the attributes and saves it back
+     in the database. Default value is False
+
+    :param bool replace: Replaces the profile and attributes
+     in the database. Default value is False
+
+    """
+
+    def __init__(
+        self,
+        profile_name: Optional[str] = None,
+        attributes: Optional[ProfileAttributes] = None,
+        description: Optional[str] = None,
+        merge: Optional[bool] = False,
+        replace: Optional[bool] = False,
+    ):
+        """Initialize a base profile"""
+        self.profile_name = profile_name
+        self.attributes = attributes
+        self.description = description
+        self.merge = merge
+        self.replace = replace
+
+    def __repr__(self):
+        return (
+            f"{self.__class__.__name__}(profile_name={self.profile_name}, "
+            f"attributes={self.attributes}, description={self.description})"
+        )