Allow multiple profiles

Author: gotsysdba

src/client/content/config/databases.py

@@ -16,20 +16,21 @@
 from streamlit import session_state as state
 
 import client.utils.api_call as api_call
+import client.utils.st_common as st_common
 import common.logging_config as logging_config
+from common.schema import SelectAIProfileType
 
 logger = logging_config.logging.getLogger("client.content.config.database")
 
 
 #####################################################
 # Functions
 #####################################################
-def get_databases(force: bool = False) -> dict[str, dict]:
+def get_databases(force: bool = False) -> None:
     """Get a dictionary of all Databases and Store Vector Store Tables"""
     if "database_config" not in state or state["database_config"] == {} or force:
         try:
-            endpoint = "v1/databases"
-            response = api_call.get(endpoint=endpoint)
+            response = api_call.get(endpoint="v1/databases")
             state["database_config"] = {
                 item["name"]: {k: v for k, v in item.items() if k != "name"} for item in response
             }
@@ -51,9 +52,8 @@ def patch_database(name: str, user: str, password: str, dsn: str, wallet_passwor
         or not state.database_config[name]["connected"]
     ):
         try:
-            endpoint = f"v1/databases/{name}"
             _ = api_call.patch(
-                endpoint=endpoint,
+                endpoint=f"v1/databases/{name}",
                 payload={
                     "json": {
                         "user": user,
@@ -66,28 +66,6 @@ def patch_database(name: str, user: str, password: str, dsn: str, wallet_passwor
             logger.info("Database updated: %s", name)
             state.database_config[name]["connected"] = True
             get_databases(force=True)
-            endpoint = "v1/selectai/enabled"
-            selectai = api_call.get(
-                            endpoint=endpoint,
-            )
-            logger.info("SelectAI enabled: %s", selectai["enabled"])
-            state.database_config[name]["selectai"] = selectai["enabled"]
-
-            # Check if SelectAI is enabled and get objects if so
-            if selectai["enabled"]:
-                try:
-                    endpoint = "v1/selectai/objects"
-                    selectai_objects = api_call.get(
-                        endpoint=endpoint,
-                    )
-                    logger.info("SelectAI objects retrieved: %d objects", len(selectai_objects))
-                    state.database_config[name]["selectai_objects"] = selectai_objects
-                except api_call.ApiError as ex:
-                    logger.error("Failed to retrieve SelectAI objects: %s", ex)
-                    state.database_config[name]["selectai_objects"] = []
-            else:
-                state.database_config[name]["selectai_objects"] = None
-
         except api_call.ApiError as ex:
             logger.error("Database not updated: %s (%s)", name, ex)
             state.database_config[name]["connected"] = False
@@ -102,21 +80,39 @@ def drop_vs(vs: dict) -> None:
     api_call.delete(endpoint=f"v1/embed/{vs['vector_store']}")
     get_databases(force=True)
 
+def select_ai_profile() -> None:
+    """Update the chosen SelectAI Profile"""
+    st_common.update_user_settings("selectai")
+    st_common.patch_settings()
+    selectai_df.clear()
 
-def update_selectai(sai_df: pd.DataFrame ) -> None:
+@st.cache_data
+def selectai_df(profile):
+    """Get SelectAI Object List and produce Dataframe"""
+    logger.info("Retrieving objects from SelectAI Profile: %s", profile)
+    st_common.patch_settings()
+    selectai_objects = api_call.get(endpoint="v1/selectai/objects")
+    df = pd.DataFrame(selectai_objects, columns=["owner", "name", "enabled"])
+    df.columns = ["Owner", "Name", "Enabled"]
+    return df
+
+
+def update_selectai(sai_new_df: pd.DataFrame, sai_old_df: pd.DataFrame) -> None:
     """Update SelectAI Object List"""
-    enabled_objects = sai_df[sai_df["Enabled"]].drop(columns=["Enabled"])
-    enabled_objects.columns = enabled_objects.columns.str.lower()
-    try:
-        endpoint = "v1/selectai/objects"
-        result = api_call.patch(
-            endpoint=endpoint,
-            payload={"json": json.loads(enabled_objects.to_json(orient="records"))}
-        )
-        logger.info("SelectAI Updated.")
-        state.database_config["DEFAULT"]["selectai_objects"] = result
-    except api_call.ApiError as ex:
-        logger.error("SelectAI not updated: %s", ex)
+    changes = sai_new_df[sai_new_df["Enabled"] != sai_old_df["Enabled"]]
+    if changes.empty:
+        st.toast("No changes detected.", icon="ℹ️")
+    else:
+        enabled_objects = sai_new_df[sai_new_df["Enabled"]].drop(columns=["Enabled"])
+        enabled_objects.columns = enabled_objects.columns.str.lower()
+        try:
+            _ = api_call.patch(
+                endpoint="v1/selectai/objects", payload={"json": json.loads(enabled_objects.to_json(orient="records"))}
+            )
+            logger.info("SelectAI Updated. Clearing Cache.")
+            selectai_df.clear()
+        except api_call.ApiError as ex:
+            logger.error("SelectAI not updated: %s", ex)
 
 
 #####################################################
@@ -209,29 +205,33 @@ def main() -> None:
         # Select AI
         #############################################
         st.subheader("SelectAI", divider="red")
-        if state.database_config[name]["selectai"]:
-            st.write("Tables eligible and enabled/disabled for SelectAI.")
-            if state.database_config[name]["selectai_objects"]:
-                df = pd.DataFrame(
-                    state.database_config[name]["selectai_objects"], columns=["owner", "name", "enabled"]
-                )
-                df.columns = ["Owner", "Name", "Enabled"]
+        selectai_profiles = state.database_config[name]["selectai_profiles"]
+        if state.database_config[name]["selectai"] and len(selectai_profiles) > 0:
+            if not state.user_settings["selectai"]["profile"]:
+                state.user_settings["selectai"]["profile"] = selectai_profiles[0]
+            # Select Profile
+            st.selectbox(
+                "Profile:",
+                options=selectai_profiles,
+                index=selectai_profiles.index(state.user_settings["selectai"]["profile"]),
+                key="selected_selectai_profile",
+                on_change=select_ai_profile,
+            )
+            selectai_objects = selectai_df(state.user_settings["selectai"]["profile"])
+            if not selectai_objects.empty:
                 sai_df = st.data_editor(
-                    df,
+                    selectai_objects,
                     column_config={
                         "enabled": st.column_config.CheckboxColumn(label="Enabled", help="Toggle to enable or disable")
                     },
                     use_container_width=True,
                     hide_index=True,
                 )
-                if st.button("Apply SelectAI Changes", type='primary'):
-                    changes = sai_df[sai_df["Enabled"] != df["Enabled"]]
-                    if not changes.empty:
-                        update_selectai(sai_df)
-                    else:
-                        st.write("No changes detected.")
+                if st.button("Apply SelectAI Changes", type="secondary"):
+                    update_selectai(sai_df, selectai_objects)
+                    st.rerun()
             else:
-                st.write("No tables found for SelectAI.")
+                st.write("No objects found for SelectAI.")
         else:
             st.write("Unable to use SelectAI with Database.")
 

src/client/utils/api_call.py

@@ -142,6 +142,7 @@ def patch(
     timeout: int = 60,
     retries: int = 5,
     backoff_factor: float = 1.5,
+    toast = True
 ) -> None:
     """PATCH Requests"""
     response = send_request(
@@ -153,7 +154,8 @@ def patch(
         retries=retries,
         backoff_factor=backoff_factor,
     )
-    st.toast("Update Successful.", icon="✅")
+    if toast:
+        st.toast("Update Successful.", icon="✅")
     return response.json()
 
 
@@ -162,8 +164,10 @@ def delete(
     timeout: int = 60,
     retries: int = 5,
     backoff_factor: float = 1.5,
+    toast = True
 ) -> None:
     """DELETE Requests"""
     response = send_request("DELETE", endpoint, timeout=timeout, retries=retries, backoff_factor=backoff_factor)
     success = response.json()["message"]
-    st.toast(success, icon="✅")
+    if toast:
+        st.toast(success, icon="✅")

src/client/utils/st_common.py

@@ -65,6 +65,7 @@ def patch_settings() -> None:
             endpoint="v1/settings",
             payload={"json": state.user_settings},
             params={"client": state.user_settings["client"]},
+            toast=False,
         )
     except api_call.ApiError as ex:
         logger.error("%s Settings Update failed: %s", state.user_settings["client"], ex)
@@ -96,7 +97,7 @@ def update_user_settings(
 def is_db_configured() -> bool:
     """Verify that a database is configured"""
     get_databases()
-    return state.database_config[state.user_settings["vector_search"]["database"]].get("connected")
+    return state.database_config[state.user_settings["database"]["alias"]].get("connected")
 
 
 def set_server_state() -> None:
@@ -154,7 +155,7 @@ def ll_sidebar() -> None:
             index=ll_idx,
             key="selected_ll_model_model",
             on_change=update_user_settings("ll_model"),
-            disabled=state.user_settings["selectai"]["enabled"]
+            disabled=state.user_settings["selectai"]["enabled"],
         )
 
     # Temperature
@@ -234,7 +235,7 @@ def ll_sidebar() -> None:
 def tools_sidebar() -> None:
     """SelectAI Sidebar Settings, conditional if all sorts of bs setup"""
 
-    def update_settings():
+    def update_set_tool():
         state.user_settings["vector_search"]["enabled"] = state.selected_tool == "VectorSearch"
         state.user_settings["selectai"]["enabled"] = state.selected_tool == "SelectAI"
 
@@ -248,6 +249,7 @@ def update_settings():
         state.user_settings["vector_search"]["enabled"] = disable_vector_search
         switch_prompt("sys", "Basic Example")
     else:
+        db_alias = state.user_settings["database"]["alias"]
         tools = [
             ("LLM Only", "Do not use tools", False),
             ("SelectAI", "Use AI with Structured Data", disable_selectai),
@@ -262,14 +264,23 @@ def update_settings():
             logger.debug("SelectAI Disabled (OCI not configured.)")
             st.warning("OCI is not fully configured.  Disabling SelectAI.", icon="⚠️")
             tools = [t for t in tools if t[0] != "SelectAI"]
+        elif not state.database_config[db_alias]["selectai"]:
+            logger.debug("SelectAI Disabled (Database not Compatible.)")
+            st.warning("Database not Compatible.  Disabling SelectAI.", icon="⚠️")
+            tools = [t for t in tools if t[0] != "SelectAI"]
+        elif len(state.database_config[db_alias]["selectai_profiles"]) == 0:
+            logger.debug("SelectAI Disabled (No profiles found.)")
+            st.warning("No profiles found.  Disabling SelectAI.", icon="⚠️")
+            tools = [t for t in tools if t[0] != "SelectAI"]
+
         # Vector Search Requirements
         get_models(model_type="embed")
         available_embed_models = list(state.embed_model_enabled.keys())
         if not available_embed_models:
             logger.debug("Vector Search Disabled (no Embedding Models)")
             st.warning("No embedding models are configured and/or enabled. Disabling Vector Search.", icon="⚠️")
             tools = [t for t in tools if t[0] != "Vector Search"]
-        elif not state.database_config[state.user_settings["vector_search"]["database"]].get("vector_stores"):
+        elif not state.database_config[db_alias].get("vector_stores"):
             logger.debug("Vector Search Disabled (Database has no vector stores.)")
             st.warning("Database has no Vector Stores. Disabling Vector Search.", icon="⚠️")
             tools = [t for t in tools if t[0] != "Vector Search"]
@@ -293,19 +304,30 @@ def update_settings():
                 # captions=tool_cap,
                 index=tool_index,
                 label_visibility="collapsed",
-                on_change=update_settings,
+                on_change=update_set_tool,
                 key="selected_tool",
             )
             if state.selected_tool == "None":
                 switch_prompt("sys", "Basic Example")
 
+
 #####################################################
 # SelectAI Options
 #####################################################
 def selectai_sidebar() -> None:
     """SelectAI Sidebar Settings, conditional if Database/SelectAI are configured"""
     if state.user_settings["selectai"]["enabled"]:
         st.sidebar.subheader("SelectAI", divider="red")
+        selectai_profiles = state.database_config[state.user_settings["database"]["alias"]]["selectai_profiles"]
+        if not state.user_settings["selectai"]["profile"]:
+            state.user_settings["selectai"]["profile"] = selectai_profiles[0]
+        st.sidebar.selectbox(
+            "Profile:",
+            options=selectai_profiles,
+            index=selectai_profiles.index(state.user_settings["selectai"]["profile"]),
+            key="selected_selectai_profile",
+            on_change=update_user_settings("selectai"),
+        )
         st.sidebar.selectbox(
             "Action:",
             get_args(SelectAISettings.__annotations__["action"]),

src/common/schema.py

@@ -43,7 +43,6 @@
 class DatabaseVectorStorage(BaseModel):
     """Database Vector Storage Tables"""
 
-    database: Optional[str] = Field(default="DEFAULT", description="Name of Database (Alias)")
     vector_store: Optional[str] = Field(
         default=None, description="Vector Store Table Name (auto-generated, do not set)", readOnly=True
     )
@@ -58,7 +57,6 @@ class DatabaseVectorStorage(BaseModel):
 class DatabaseSelectAIObjects(BaseModel):
     """Database SelectAI Objects"""
 
-    database: Optional[str] = Field(default="DEFAULT", description="Name of Database (Alias)")
     owner: Optional[str] = Field(default=None, description="Object Owner", readOnly=True)
     name: Optional[str] = Field(default=None, description="Object Name", readOnly=True)
     enabled: bool = Field(default=False, description="SelectAI Enabled")
@@ -80,14 +78,12 @@ class Database(DatabaseAuth):
     """Database Object"""
 
     name: str = Field(default="DEFAULT", description="Name of Database (Alias)")
-    selectai: bool = Field(default=False, description="SelectAI Possible")
     connected: bool = Field(default=False, description="Connection Established")
     vector_stores: Optional[list[DatabaseVectorStorage]] = Field(
         default=None, description="Vector Storage (read-only)", readOnly=True
     )
-    selectai_objects: Optional[list[DatabaseSelectAIObjects]] = Field(
-        default=None, description="SelectAI Eligible Objects (read-only)", readOnly=True
-    )
+    selectai: bool = Field(default=False, description="SelectAI Possible")
+    selectai_profiles: Optional[list] = Field(default=[], description="SelectAI Profiles (read-only)", readOnly=True)
     # Do not expose the connection to the endpoint
     _connection: oracledb.Connection = PrivateAttr(default=None)
 
@@ -242,7 +238,7 @@ class SelectAISettings(BaseModel):
     """Store SelectAI Settings"""
 
     enabled: bool = Field(default=False, description="SelectAI Enabled")
-    profile: str = Field(default="OPTIMIZER_PROFILE", description="SelectAI Profile", readOnly=True)
+    profile: Optional[str] = Field(default=None, description="SelectAI Profile")
     action: Literal["runsql", "showsql", "explainsql", "narrate"] = Field(
         default="narrate", description="SelectAI Action"
     )
@@ -254,6 +250,12 @@ class OciSettings(BaseModel):
     auth_profile: Optional[str] = Field(default="DEFAULT", description="Oracle Cloud Settings Profile")
 
 
+class DatabaseSettings(BaseModel):
+    """Database Settings"""
+
+    alias: str = Field(default="DEFAULT", description="Name of Database (Alias)")
+
+
 class Settings(BaseModel):
     """Server Settings"""
 
@@ -269,6 +271,7 @@ class Settings(BaseModel):
         default_factory=PromptSettings, description="Prompt Engineering Settings"
     )
     oci: Optional[OciSettings] = Field(default_factory=OciSettings, description="OCI Settings")
+    database: Optional[DatabaseSettings] = Field(default_factory=DatabaseSettings, description="Database Settings")
     vector_search: Optional[VectorSearchSettings] = Field(
         default_factory=VectorSearchSettings, description="Vector Search Settings"
     )
@@ -401,6 +404,7 @@ class EvaluationReport(Evaluation):
 PromptNameType = Prompt.__annotations__["name"]
 PromptCategoryType = Prompt.__annotations__["category"]
 PromptPromptType = PromptText.__annotations__["prompt"]
+SelectAIProfileType = Database.__annotations__["selectai_profiles"]
 TestSetsIdType = TestSets.__annotations__["tid"]
 TestSetsNameType = TestSets.__annotations__["name"]
 TestSetDateType = TestSets.__annotations__["created"]

src/server/bootstrap/database_def.py

@@ -7,6 +7,7 @@
 import os
 import server.utils.databases as databases
 import server.utils.embedding as embedding
+import server.utils.selectai as selectai
 from common.schema import Database
 import common.logging_config as logging_config
 
@@ -48,9 +49,9 @@ def main() -> list[Database]:
             db.connected = False
             continue
         db.vector_stores = embedding.get_vs(conn)
-        db.selectai = databases.selectai_enabled(conn)
+        db.selectai = selectai.enabled(conn)
         if db.selectai:
-            db.selectai_objects = databases.get_selectai_objects(conn)
+            db.selectai_profiles = selectai.get_profiles(conn)
         if not db.connection and len(database_objects) > 1:
             db.set_connection = databases.disconnect(conn)
         else: