Merge pull request #14 from oracle/dev/v1.2.1

dev/v1.2.1

Author: aosingh

.gitignore

@@ -19,3 +19,4 @@ pytest.env
 sample_connect.py
 async_pipeline_test.py
 parquet.py
+local_sample

src/select_ai/action.py

@@ -21,3 +21,4 @@ class Action(StrEnum):
     SHOWPROMPT = "showprompt"
     FEEDBACK = "feedback"
     SUMMARIZE = "summarize"
+    TRANSLATE = "translate"

src/select_ai/async_profile.py

@@ -7,7 +7,6 @@
 
 import json
 from contextlib import asynccontextmanager
-from dataclasses import replace as dataclass_replace
 from typing import (
     AsyncGenerator,
     List,
@@ -30,7 +29,10 @@
 )
 from select_ai.conversation import AsyncConversation
 from select_ai.db import async_cursor, async_get_connection
-from select_ai.errors import ProfileExistsError, ProfileNotFoundError
+from select_ai.errors import (
+    ProfileAttributesEmptyError,
+    ProfileNotFoundError,
+)
 from select_ai.feedback import (
     FeedbackOperation,
     FeedbackType,
@@ -69,35 +71,23 @@ async def _init_profile(self):
         if self.profile_name:
             profile_exists = False
             try:
-                saved_attributes = await self._get_attributes(
+                saved_description = await self._get_profile_description(
                     profile_name=self.profile_name
                 )
                 profile_exists = True
-                if not self.replace and not self.merge:
-                    if (
-                        self.attributes is not None
-                        or self.description is not None
-                    ):
-                        if self.raise_error_if_exists:
-                            raise ProfileExistsError(self.profile_name)
-
-                if self.description is None and not self.replace:
-                    self.description = await self._get_profile_description(
-                        profile_name=self.profile_name
-                    )
+                saved_attributes = await self._get_attributes(
+                    profile_name=self.profile_name,
+                    raise_on_empty=True,
+                )
+                self._raise_error_if_profile_exists()
+            except ProfileAttributesEmptyError:
+                if self.raise_error_on_empty_attributes:
+                    raise
             except ProfileNotFoundError:
                 if self.attributes is None and self.description 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),
-                        )
+                self._merge_attributes(saved_attributes, saved_description)
             if self.replace or not profile_exists:
                 await self.create(replace=self.replace)
         else:  # profile name is None:
@@ -132,12 +122,15 @@ async def _get_profile_description(profile_name) -> Union[str, None]:
                 raise ProfileNotFoundError(profile_name)
 
     @staticmethod
-    async def _get_attributes(profile_name) -> ProfileAttributes:
+    async def _get_attributes(
+        profile_name: str, 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 bool raise_on_empty: Raise an error if attributes are empty
         :return: select_ai.provider.ProviderAttributes
-        :raises: ProfileNotFoundError
+        :raises: select_ai.errors.ProfileAttributesEmptyError
 
         """
         async with async_cursor() as cr:
@@ -149,7 +142,11 @@ async def _get_attributes(profile_name) -> ProfileAttributes:
             if attributes:
                 return await ProfileAttributes.async_create(**dict(attributes))
             else:
-                raise ProfileNotFoundError(profile_name=profile_name)
+                if raise_on_empty:
+                    raise ProfileAttributesEmptyError(
+                        profile_name=profile_name
+                    )
+                return None
 
     async def get_attributes(self) -> ProfileAttributes:
         """Asynchronously gets AI profile attributes from the Database
@@ -387,7 +384,9 @@ async def list(
             for row in rows:
                 profile_name = row[0]
                 yield await cls(
-                    profile_name=profile_name, raise_error_if_exists=False
+                    profile_name=profile_name,
+                    raise_error_if_exists=False,
+                    raise_error_on_empty_attributes=False,
                 )
 
     async def generate(
@@ -623,6 +622,34 @@ async def run_pipeline(
                 responses.append(result.error)
         return responses
 
+    async def translate(
+        self, text: str, source_language: str, target_language: str
+    ) -> Union[str, None]:
+        """
+        Translate a text using a source language and a target language
+
+        :param str text: Text to translate
+        :param str source_language: Source language
+        :param str target_language: Target language
+        :return: str
+        """
+        parameters = {
+            "profile_name": self.profile_name,
+            "text": text,
+            "source_language": source_language,
+            "target_language": target_language,
+        }
+        async with async_cursor() as cr:
+            data = await cr.callfunc(
+                "DBMS_CLOUD_AI.TRANSLATE",
+                oracledb.DB_TYPE_CLOB,
+                keyword_parameters=parameters,
+            )
+        if data is not None:
+            result = await data.read()
+            return result
+        return None
+
 
 class AsyncSession:
     """AsyncSession lets you persist request parameters across DBMS_CLOUD_AI

src/select_ai/base_profile.py

@@ -8,12 +8,14 @@
 import json
 from abc import ABC
 from dataclasses import dataclass
+from dataclasses import replace as dataclass_replace
 from typing import List, Mapping, Optional, Tuple
 
 import oracledb
 
 from select_ai._abc import SelectAIDataClass
 from select_ai.action import Action
+from select_ai.errors import ProfileExistsError
 from select_ai.feedback import (
     FeedbackOperation,
     FeedbackType,
@@ -159,6 +161,10 @@ class BaseProfile(ABC):
      if profile exists in the database and replace = False and
      merge = False. Default value is True
 
+    :param bool  raise_error_on_empty_attributes: Raise
+     ProfileEmptyAttributesError, if profile attributes are empty
+     in database. Default value is False.
+
     """
 
     def __init__(
@@ -169,6 +175,7 @@ def __init__(
         merge: Optional[bool] = False,
         replace: Optional[bool] = False,
         raise_error_if_exists: Optional[bool] = True,
+        raise_error_on_empty_attributes: Optional[bool] = False,
     ):
         """Initialize a base profile"""
         self.profile_name = profile_name
@@ -182,6 +189,34 @@ def __init__(
         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
+
+    def _raise_error_if_profile_exists(self):
+        """
+        Helper method to raise ProfileExistsError if profile exists
+        in the database and replace = False and merge = False
+        """
+        if not self.replace and not self.merge:
+            if self.attributes is not None or self.description is not None:
+                if self.raise_error_if_exists:
+                    raise ProfileExistsError(self.profile_name)
+
+    def _merge_attributes(self, saved_attributes, saved_description):
+        """
+        Helper method to merge user passed attributes with the attributes saved
+        in the database.
+        """
+        if self.description is None and not self.replace:
+            self.description = saved_description
+        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),
+                )
 
     def __repr__(self):
         return (
@@ -206,15 +241,15 @@ def validate_params_for_feedback(
     response: Optional[str] = None,
     operation: Optional[FeedbackOperation] = FeedbackOperation.ADD,
 ):
-    if sql_id and prompt_spec:
-        raise AttributeError("Either sql_id or prompt_spec must be specified")
     if not sql_id and not prompt_spec:
         raise AttributeError("Either sql_id or prompt_spec must be specified")
-    parameters = {
-        "feedback_type": feedback_type.value,
-        "feedback_content": feedback_content,
-        "operation": operation.value,
-    }
+    parameters = {"operation": operation.value}
+    if feedback_content:
+        parameters["feedback_content"] = feedback_content
+    if feedback_type:
+        parameters["feedback_type"] = feedback_type.value
+    if response:
+        parameters["response"] = response
     if prompt_spec:
         prompt, action = prompt_spec
         if action not in (Action.RUNSQL, Action.SHOWSQL, Action.EXPLAINSQL):

src/select_ai/errors.py

@@ -56,6 +56,18 @@ def __str__(self):
         )
 
 
+class ProfileAttributesEmptyError(SelectAIError):
+    """Profile attributes empty in the database"""
+
+    def __init__(self, profile_name: str):
+        self.profile_name = profile_name
+
+    def __str__(self):
+        return (
+            f"Profile {self.profile_name} attributes empty in the database. "
+        )
+
+
 class VectorIndexNotFoundError(SelectAIError):
     """VectorIndex not found in the database"""