feat: okta profile tool

backend/onyx/configs/app_configs.py

@@ -121,6 +121,8 @@
     os.environ.get("OAUTH_CLIENT_SECRET", os.environ.get("GOOGLE_OAUTH_CLIENT_SECRET"))
     or ""
 )
+# OpenID Connect configuration URL for Okta Profile Tool and other OIDC integrations
+OPENID_CONFIG_URL = os.environ.get("OPENID_CONFIG_URL") or ""
 
 USER_AUTH_SECRET = os.environ.get("USER_AUTH_SECRET", "")
 
@@ -617,6 +619,17 @@ def get_current_tz_offset() -> int:
 
 MAX_TOKENS_FOR_FULL_INCLUSION = 4096
 
+
+#####
+# Tool Configs
+#####
+OKTA_PROFILE_TOOL_ENABLED = (
+    os.environ.get("OKTA_PROFILE_TOOL_ENABLED", "").lower() == "true"
+)
+# API token for SSWS auth to Okta Admin API. If set, Users API will be used to enrich profile.
+OKTA_API_TOKEN = os.environ.get("OKTA_API_TOKEN") or ""
+
+
 #####
 # Miscellaneous
 #####

backend/onyx/llm/models.py

@@ -8,8 +8,6 @@
 
 from onyx.configs.constants import MessageType
 from onyx.file_store.models import InMemoryChatFile
-from onyx.llm.utils import build_content_with_imgs
-from onyx.llm.utils import message_to_string
 from onyx.tools.models import ToolCallFinalResult
 
 if TYPE_CHECKING:
@@ -55,7 +53,10 @@ def from_chat_message(
         )
 
     def to_langchain_msg(self) -> BaseMessage:
-        content = build_content_with_imgs(self.message, self.files)
+        # Lazy import to avoid heavy dependencies during light-weight tests
+        from onyx.llm.utils import build_content_with_imgs as _build_content_with_imgs
+
+        content = _build_content_with_imgs(self.message, self.files)
         if self.message_type == MessageType.USER:
             return HumanMessage(content=content)
         elif self.message_type == MessageType.ASSISTANT:
@@ -72,7 +73,10 @@ def from_langchain_msg(
             message_type = MessageType.USER
         elif isinstance(msg, AIMessage):
             message_type = MessageType.ASSISTANT
-        message = message_to_string(msg)
+        # Lazy import to avoid heavy dependencies during light-weight tests
+        from onyx.llm.utils import message_to_string as _message_to_string
+
+        message = _message_to_string(msg)
         return cls(
             message=message,
             token_count=token_count,

backend/onyx/tools/built_in_tools.py

@@ -6,6 +6,7 @@
 from sqlalchemy import select
 from sqlalchemy.orm import Session
 
+from onyx.configs.app_configs import OKTA_PROFILE_TOOL_ENABLED
 from onyx.db.models import Persona
 from onyx.db.models import Tool as ToolDBModel
 from onyx.tools.tool_implementations.images.image_generation_tool import (
@@ -17,6 +18,9 @@
 from onyx.tools.tool_implementations.internet_search.providers import (
     get_available_providers,
 )
+from onyx.tools.tool_implementations.okta_profile.okta_profile_tool import (
+    OktaProfileTool,
+)
 from onyx.tools.tool_implementations.search.search_tool import SearchTool
 from onyx.tools.tool import Tool
 from onyx.utils.logger import setup_logger
@@ -63,6 +67,19 @@ class InCodeToolInfo(TypedDict):
         if (bool(get_available_providers()))
         else []
     ),
+    # Show Okta Profile tool if the environment variables are set
+    *(
+        [
+            InCodeToolInfo(
+                cls=OktaProfileTool,
+                description="The Okta Profile Action allows the assistant to fetch user information from Okta.",
+                in_code_tool_id=OktaProfileTool.__name__,
+                display_name=OktaProfileTool._DISPLAY_NAME,
+            )
+        ]
+        if OKTA_PROFILE_TOOL_ENABLED
+        else []
+    ),
 ]
 
 

backend/onyx/tools/tool.py

@@ -5,12 +5,12 @@
 from typing import TYPE_CHECKING
 from typing import TypeVar
 
-from onyx.llm.interfaces import LLM
-from onyx.llm.models import PreviousMessage
 from onyx.utils.special_types import JSON_ro
 
 
 if TYPE_CHECKING:
+    from onyx.llm.interfaces import LLM
+    from onyx.llm.models import PreviousMessage
     from onyx.chat.prompt_builder.answer_prompt_builder import AnswerPromptBuilder
     from onyx.tools.message import ToolCallSummary
     from onyx.tools.models import ToolResponse
@@ -53,8 +53,8 @@ def build_tool_message_content(
     def get_args_for_non_tool_calling_llm(
         self,
         query: str,
-        history: list[PreviousMessage],
-        llm: LLM,
+        history: list["PreviousMessage"],
+        llm: "LLM",
         force_run: bool = False,
     ) -> dict[str, Any] | None:
         raise NotImplementedError

backend/onyx/tools/tool_constructor.py

@@ -14,6 +14,10 @@
 from onyx.configs.app_configs import AZURE_DALLE_API_VERSION
 from onyx.configs.app_configs import AZURE_DALLE_DEPLOYMENT_NAME
 from onyx.configs.app_configs import IMAGE_MODEL_NAME
+from onyx.configs.app_configs import OAUTH_CLIENT_ID
+from onyx.configs.app_configs import OAUTH_CLIENT_SECRET
+from onyx.configs.app_configs import OKTA_API_TOKEN
+from onyx.configs.app_configs import OPENID_CONFIG_URL
 from onyx.configs.chat_configs import NUM_INTERNET_SEARCH_CHUNKS
 from onyx.configs.chat_configs import NUM_INTERNET_SEARCH_RESULTS
 from onyx.configs.model_configs import GEN_AI_TEMPERATURE
@@ -41,6 +45,9 @@
 from onyx.tools.tool_implementations.internet_search.internet_search_tool import (
     InternetSearchTool,
 )
+from onyx.tools.tool_implementations.okta_profile.okta_profile_tool import (
+    OktaProfileTool,
+)
 from onyx.tools.tool_implementations.search.search_tool import SearchTool
 from onyx.tools.utils import compute_all_tool_tokens
 from onyx.tools.utils import explicit_tool_calling_supported
@@ -265,6 +272,33 @@ def construct_tools(
                         "Internet search tool requires a Bing or Exa API key, please contact your Onyx admin to get it added!"
                     )
 
+            # Handle Okta Profile Tool
+            elif tool_cls.__name__ == OktaProfileTool.__name__:
+                if not user_oauth_token:
+                    raise ValueError(
+                        "Okta Profile Tool requires user OAuth token but none found"
+                    )
+
+                if not all([OAUTH_CLIENT_ID, OAUTH_CLIENT_SECRET, OPENID_CONFIG_URL]):
+                    raise ValueError(
+                        "Okta Profile Tool requires OAuth configuration to be set"
+                    )
+
+                if not OKTA_API_TOKEN:
+                    raise ValueError(
+                        "Okta Profile Tool requires OKTA_API_TOKEN to be set"
+                    )
+
+                tool_dict[db_tool_model.id] = [
+                    OktaProfileTool(
+                        access_token=user_oauth_token,
+                        client_id=OAUTH_CLIENT_ID,
+                        client_secret=OAUTH_CLIENT_SECRET,
+                        openid_config_url=OPENID_CONFIG_URL,
+                        okta_api_token=OKTA_API_TOKEN,
+                    )
+                ]
+
         # Handle custom tools
         elif db_tool_model.openapi_schema:
             if not custom_tool_config: