Fixed bug for enum, docs update and added set_attributes

aosingh · aosingh · commit e210728623f6 · 2025-05-30T17:09:05.000-07:00
diff --git a/doc/source/index.rst b/doc/source/index.rst
@@ -68,3 +68,30 @@ Async Profile
     :maxdepth: 3
 
     user_guide/async_profile.rst
+
+Vector Index
+============
+
+.. toctree::
+    :numbered:
+    :maxdepth: 3
+
+    user_guide/vector_index.rst
+
+Synthetic Data
+==============
+
+.. toctree::
+    :numbered:
+    :maxdepth: 3
+
+    user_guide/synthetic_data.rst
+
+Conversation
+============
+
+.. toctree::
+    :numbered:
+    :maxdepth: 3
+
+    user_guide/conversation.rst
diff --git a/doc/source/user_guide/actions.rst b/doc/source/user_guide/actions.rst
@@ -4,5 +4,4 @@
 Supported AI Actions
 ********************
 
-.. literalinclude:: ../../../src/select_ai/action.py
-   :language: python
+TODO
diff --git a/doc/source/user_guide/async_profile.rst b/doc/source/user_guide/async_profile.rst
@@ -15,22 +15,15 @@ Async SQL generation
    :language: python
 
 **********
-Async Chat
+Async chat
 **********
 
 .. literalinclude:: ../../../examples/async_examples/2_chat.py
    :language: python
 
 *********************
-Async Prompt Pipeline
+Async pipeline
 *********************
 
 .. literalinclude:: ../../../examples/async_examples/3_pipeline.py
    :language: python
-
-*********************
-Async Conversation
-*********************
-
-.. literalinclude:: ../../../examples/async_examples/4_conversation.py
-   :language: python
diff --git a/doc/source/user_guide/conversation.rst b/doc/source/user_guide/conversation.rst
@@ -0,0 +1,29 @@
+.. _conversation:
+
+********************
+``Conversation``
+********************
+
+.. autoclass:: select_ai.Conversation
+   :members:
+
+*************************
+``VectorIndexAttributes``
+*************************
+
+.. autoclass:: select_ai.ConversationAttributes
+   :members:
+
+************
+Conversation
+************
+
+.. literalinclude:: ../../../examples/9_conversation.py
+   :language: python
+
+*********************
+Async conversation
+*********************
+
+.. literalinclude:: ../../../examples/async_examples/4_conversation.py
+   :language: python
diff --git a/doc/source/user_guide/profile.rst b/doc/source/user_guide/profile.rst
@@ -8,7 +8,7 @@
    :members:
 
 **************************
-Enable AI Service provider
+Enable AI service provider
 **************************
 
 .. literalinclude:: ../../../examples/1_enable_ai_provider.py
@@ -45,32 +45,12 @@ Chat
    :language: python
 
 **************************
-List Profiles
+List profiles
 **************************
 
 .. literalinclude:: ../../../examples/6_list_ai_profiles.py
    :language: python
 
-**************************
-Create Vector Index
-**************************
-
-.. literalinclude:: ../../../examples/7_vector_index.py
-   :language: python
-
-**************************
-Generate Synthetic Data
-**************************
-
-.. literalinclude:: ../../../examples/8_generate_synthetic_data.py
-   :language: python
-
-**************************
-Conversation
-**************************
-
-.. literalinclude:: ../../../examples/9_conversation.py
-   :language: python
 
 ***************************
 Disable AI service provider
diff --git a/doc/source/user_guide/provider_attributes.rst b/doc/source/user_guide/provider_attributes.rst
@@ -25,9 +25,9 @@
 .. autoclass:: select_ai.AWSAIProviderAttributes
    :members:
 
-*****************************
+******************************
 ``CohereAIProviderAttributes``
-*****************************
+******************************
 .. autoclass:: select_ai.CohereAIProviderAttributes
    :members:
 
@@ -38,16 +38,16 @@
    :members:
 
 
-*****************************
+******************************
 ``OCIGenAIProviderAttributes``
-*****************************
+******************************
 .. autoclass:: select_ai.OCIGenAIProviderAttributes
    :members:
 
 
-*****************************
+******************************
 ``GoogleAIProviderAttributes``
-*****************************
+******************************
 .. autoclass:: select_ai.GoogleAIProviderAttributes
    :members:
 
diff --git a/doc/source/user_guide/synthetic_data.rst b/doc/source/user_guide/synthetic_data.rst
@@ -0,0 +1,24 @@
+.. _synthetic_data:
+
+***********************
+``SyntheticDataParams``
+***********************
+
+.. autoclass:: select_ai.SyntheticDataParams
+   :members:
+
+***************************
+``SyntheticDataAttributes``
+***************************
+
+.. autoclass:: select_ai.SyntheticDataAttributes
+   :members:
+
+**************************
+Generate synthetic data
+**************************
+
+The below example shows multitable synthetic data generation
+
+.. literalinclude:: ../../../examples/8_generate_synthetic_data.py
+   :language: python
diff --git a/doc/source/user_guide/vector_index.rst b/doc/source/user_guide/vector_index.rst
@@ -0,0 +1,22 @@
+.. _vector_index:
+
+********************
+``VectorIndex``
+********************
+
+.. autoclass:: select_ai.VectorIndex
+   :members:
+
+*************************
+``VectorIndexAttributes``
+*************************
+
+.. autoclass:: select_ai.VectorIndexAttributes
+   :members:
+
+**************************
+Create Vector Index
+**************************
+
+.. literalinclude:: ../../../examples/7_vector_index.py
+   :language: python
diff --git a/examples/4_run_sql.py b/examples/4_run_sql.py
@@ -21,6 +21,10 @@ def main():
         description="MY OCI AI Profile",
         replace=True,
     )
+    # profile.set_attributes(attributes=oci_provider_attributes)
+    profile.set_attribute(
+        attribute_name="model", attribute_value="meta.llama-3.1-70b-instruct"
+    )
     prompts = [
         "How many promotions are there in the sh database?",
         "How many products are there in the sh database ?",
diff --git a/src/select_ai/_base.py b/src/select_ai/_base.py
@@ -26,7 +26,7 @@ def __init__(
         :param str profile_name (optional): Name of the profile
 
         :param select_ai.provider.ProviderAttributes attributes (optional):
-        ProviderAttributes
+        Object specifying AI profile attributes
 
         :param str description (optional): Description of the profile
 
diff --git a/src/select_ai/_enums.py b/src/select_ai/_enums.py
@@ -0,0 +1,7 @@
+import enum
+
+
+class StrEnum(str, enum.Enum):
+
+    def __str__(self):
+        return self.value
diff --git a/src/select_ai/action.py b/src/select_ai/action.py
@@ -1,4 +1,4 @@
-from enum import StrEnum
+from select_ai._enums import StrEnum
 
 __all__ = ["Action"]
 
diff --git a/src/select_ai/async_profile.py b/src/select_ai/async_profile.py
@@ -102,6 +102,49 @@ async def fetch_attributes(profile_name) -> ProviderAttributes:
             else:
                 raise ProfileNotFoundError(profile_name=profile_name)
 
+    async def set_attribute(
+        self,
+        attribute_name: str,
+        attribute_value: Union[bool, str, int, float],
+    ):
+        """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
+
+        """
+        setattr(self.attributes, attribute_name, attribute_value)
+        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_attributes(self, attributes: ProviderAttributes):
+        """Updates AI profile attributes on the Python object and also
+        saves it in the database
+
+        :param ProviderAttributes 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:
@@ -371,26 +414,41 @@ async def disable_vector_index(index_name: str) -> None:
 
     @staticmethod
     async def update_vector_index(
-        index_name: str, attribute_name: str, attribute_value: Union[str, int]
+        index_name: str,
+        attribute_name: str,
+        attribute_value: Union[str, int],
+        attributes: VectorIndexAttributes = None,
     ) -> None:
         """
         This procedure updates an existing vector store index with a specified
-        value of the vector index attribute.
+        value of the vector index attribute. You can specify a single attribute
+        or multiple attributes by passing an object of type
+        :class `VectorIndexAttributes`
 
         :param str index_name: Name of the vector index
         :param str attribute_name: Custom attribute name
         :param Union[str, int, float] attribute_value: Attribute Value
+        :param VectorIndexAttributes attributes: Specify multiple attributes
+         to update in a single API invocation
         :return: None
         :raises: oracledb.DatabaseError
         """
+        if attribute_name and attribute_value and attributes:
+            raise ValueError(
+                "Only one of attribute (name and value) or "
+                "attributes can be specified"
+            )
+        parameters = {"index_name": index_name}
+        if attributes:
+            parameters["attributes"] = attributes.json()
+        else:
+            parameters["attributes_name"] = attribute_name
+            parameters["attributes_value"] = attribute_value
+
         async with async_cursor() as cr:
             await cr.callproc(
                 "DBMS_CLOUD_AI.UPDATE_VECTOR_INDEX",
-                keyword_parameters={
-                    "index_name": index_name,
-                    "attribute_name": attribute_name,
-                    "attribute_value": attribute_value,
-                },
+                keyword_parameters=parameters,
             )
 
     @staticmethod
diff --git a/src/select_ai/conversation.py b/src/select_ai/conversation.py
@@ -1,16 +1,18 @@
 import collections
 import json
 from dataclasses import dataclass
-from typing import Optional
+from typing import NamedTuple, Optional
 
 from select_ai._base import SelectAIDataClass
 
 __all__ = ["Conversation", "ConversationAttributes"]
 
 
-Conversation = collections.namedtuple(
-    "Conversation", field_names=["conversation_id", "attributes"]
-)
+class Conversation(NamedTuple):
+    """A container class to store Conversation id and attributes"""
+
+    conversation_id: str
+    attributes: "ConversationAttributes"
 
 
 @dataclass
diff --git a/src/select_ai/profile.py b/src/select_ai/profile.py
diff --git a/src/select_ai/provider.py b/src/select_ai/provider.py
diff --git a/src/select_ai/session.py b/src/select_ai/session.py
diff --git a/src/select_ai/synthetic_data.py b/src/select_ai/synthetic_data.py
diff --git a/src/select_ai/vector_index.py b/src/select_ai/vector_index.py
diff --git a/src/select_ai/version.py b/src/select_ai/version.py
