oci 2.151.0__py3-none-any.whl → 2.152.1__py3-none-any.whl

oci/addons/adk/agent_client.py

@@ -0,0 +1,682 @@
+# coding: utf-8
+# Copyright (c) 2016, 2025, Oracle and/or its affiliates.  All rights reserved.
+# This software is dual-licensed to you under the Universal Permissive License (UPL) 1.0 as shown at https://oss.oracle.com/licenses/upl or Apache License 2.0 as shown at http://www.apache.org/licenses/LICENSE-2.0. You may choose either license.
+
+"""
+ADK Agent Client
+"""
+
+from typing import Any, Callable, Dict, List, Literal, Optional, Tuple
+
+import oci
+from oci.constants import HEADER_REQUEST_ID
+from oci.exceptions import ClientError, ServiceError
+from oci.generative_ai_agent import GenerativeAiAgentClient
+from oci.generative_ai_agent.models import (
+    CreateToolDetails,
+    Function,
+    FunctionCallingToolConfig,
+    KnowledgeBaseConfig,
+    LlmConfig,
+    LlmCustomization,
+    RagToolConfig,
+    UpdateAgentDetails,
+    UpdateToolDetails,
+)
+from oci.generative_ai_agent_runtime import GenerativeAiAgentRuntimeClient
+from oci.generative_ai_agent_runtime.models import (
+    ChatDetails,
+    ChatResult,
+    CreateSessionDetails,
+    FunctionCallingPerformedAction,
+    Session,
+)
+from oci.response import Response
+from oci.util import formatted_flat_dict, to_dict
+from oci.pagination import list_call_get_all_results
+
+from oci.addons.adk.agent_error import AgentError, UserError
+from oci.addons.adk.auth.factory import AuthFactory, OCIAuthType
+from oci.addons.adk.constants import (
+    DEFAULT_SESSION_NAME,
+    DEFAULT_TIMEOUT,
+    DESC_SUFFIX,
+    FREEFORM_TAGS,
+)
+from oci.addons.adk.logger import default_logger as logger
+from oci.addons.adk.run.types import PerformedAction
+from oci.addons.adk.tool import FunctionTool
+from oci.addons.adk.tool.prebuilt import AgenticRagTool
+from oci.addons.adk.util import build_custom_function_params, get_region_endpoint
+
+
+class AgentClient:
+    """Client for interacting with OCI Agent endpoints."""
+
+    def __init__(
+        self,
+        auth_type: Optional[str] = "api_key",
+        config: Optional[str] = oci.config.DEFAULT_LOCATION,
+        profile: Optional[str] = oci.config.DEFAULT_PROFILE,
+        token: Optional[str] = None,
+        region: Optional[str] = None,
+        runtime_endpoint: Optional[str] = None,
+        management_endpoint: Optional[str] = None,
+        timeout: Optional[float | Tuple[float, float]] = DEFAULT_TIMEOUT,
+        debug: Optional[bool] = False,
+    ) -> None:
+        """
+        Initialize the agent client.
+
+        Args:
+            auth_type: The authentication type, Must be
+                one of: 'api_key', 'security_token', 'instance_principal',
+                'resource_principal', 'instance_obo_user'
+            config: The config file path for the agent client
+            profile: The profile to use for authentication
+            token: The token to use for authentication
+            region: The region to use for authentication
+                If provided, runtime_endpoint and management_endpoint will be
+                constructed automatically unless explicitly specified
+            runtime_endpoint: The runtime endpoint for OCI Generative AI Agent
+            management_endpoint: The management endpoint for OCI Generative AI Agent
+            timeout: The timeout for the agent client
+            debug: Whether to enable debug logging
+
+        Raises:
+            AgentError: If the agent endpoint id is unknown
+        """
+        # use auth module
+        self.auth = AuthFactory.create_oci_auth(
+            auth_type=OCIAuthType(auth_type),
+            config_path=config,
+            profile=profile,
+            token=token,
+            region=region,
+        )
+
+        # endpoint
+        self.runtime_endpoint = self.get_runtime_endpoint(runtime_endpoint, region)
+        self.management_endpoint = self.get_management_endpoint(
+            management_endpoint, region
+        )
+
+        # init Agent Runtime Client
+        self._rt_client = GenerativeAiAgentRuntimeClient(
+            config=self.auth.get_config(),
+            signer=self.auth.get_auth_credentials(),
+            service_endpoint=self.runtime_endpoint,
+            timeout=timeout,
+            retry_strategy=oci.retry.DEFAULT_RETRY_STRATEGY,
+        )
+        # init Agent client
+        self._mgmt_client = GenerativeAiAgentClient(
+            config=self.auth.get_config(),
+            signer=self.auth.get_auth_credentials(),
+            service_endpoint=self.management_endpoint,
+            timeout=timeout,
+            retry_strategy=oci.retry.DEFAULT_RETRY_STRATEGY,
+        )
+        self.debug = debug
+        self._handlers: Dict[str, Callable] = {}
+
+    def create_session(
+        self,
+        agent_endpoint_id: str,
+        display_name: str | None = None,
+        description: str | None = None,
+    ) -> str:
+        """
+        Create a new agent session.
+
+        Args:
+            agent_endpoint_id: The agent endpoint id
+            display_name: Name for the session
+            description: Description for the session
+
+        Returns:
+            The session ID
+
+        Raises:
+            AgentError: If session creation fails
+        """
+        create_session_details = CreateSessionDetails()
+        create_session_details.display_name = (
+            display_name if display_name else DEFAULT_SESSION_NAME
+        )
+        create_session_details.description = (
+            description if description else DEFAULT_SESSION_NAME + DESC_SUFFIX
+        )
+
+        logger.debug(f"Creating session with details: {create_session_details}")
+
+        create_session_response: Response = self._invoke(
+            method=self._rt_client.create_session,
+            success_message="Session creation succeeded",
+            error_message="Failed to create session",
+            debug=self.debug,
+            create_session_details=create_session_details,
+            agent_endpoint_id=agent_endpoint_id,
+        )
+        response_data: Session = create_session_response.data
+        return response_data.id
+
+    def delete_session(self, agent_endpoint_id: str, session_id: str) -> None:
+        """
+        Delete an existing agent session.
+
+        Args:
+            agent_endpoint_id: The agent endpoint id
+            session_id: The session ID to delete
+
+        Returns:
+            None
+
+        Raises:
+            AgentError: If session deletion fails
+        """
+        logger.debug(f"Deleting session: {session_id}")
+        self._invoke(
+            method=self._rt_client.delete_session,
+            success_message="Session deletion succeeded",
+            error_message="Failed to delete session",
+            debug=self.debug,
+            session_id=session_id,
+            agent_endpoint_id=agent_endpoint_id,
+        )
+        return
+
+    def chat(
+        self,
+        agent_endpoint_id: str,
+        session_id: str,
+        user_message: Optional[str] = None,
+        performed_actions: Optional[List[PerformedAction]] = None,
+    ) -> Dict[str, Any]:
+        """
+        Invoke a chat session with the agent.
+
+        Args:
+            agent_endpoint_id: The agent endpoint id
+            session_id: The session ID to use for the chat
+            user_message: The message from the user
+            performed_actions: Actions that have been performed
+
+        Returns:
+            The response from the agent
+
+        Raises:
+            AgentError: If the invocation fails
+        """
+        chat_details = ChatDetails()
+        fc_performed_actions = []
+        if performed_actions:
+            fc_performed_actions = [
+                FunctionCallingPerformedAction(
+                    action_id=action.action_id,
+                    performed_action_type=action.performed_action_type,
+                    function_call_output=action.function_call_output,
+                )
+                for action in performed_actions
+            ]
+        chat_details.performed_actions = fc_performed_actions
+        chat_details.user_message = user_message
+        chat_details.should_stream = False
+        chat_details.session_id = session_id
+
+        logger.debug(f"Invoking chat endpoint with data: {chat_details}")
+
+        chat_response: Response = self._invoke(
+            method=self._rt_client.chat,
+            success_message="Chat succeeded",
+            error_message="Failed to invoke chat endpoint",
+            debug=self.debug,
+            chat_details=chat_details,
+            agent_endpoint_id=agent_endpoint_id,
+        )
+
+        response_data: ChatResult = chat_response.data
+        return to_dict(response_data)
+
+    def update_agent(
+        self,
+        agent_id: str,
+        name: str | None = None,
+        description: str | None = None,
+        instructions: str | None = None,
+    ) -> None:
+        """
+        Update the agent.
+
+        Args:
+            agent_id: The ID of the agent to update
+            name: The new name for the agent
+            description: The new description for the agent
+            instructions: The new instructions for the agent
+
+        Returns:
+            None
+
+        Raises:
+            AgentError: If updating the instructions fails
+        """
+        update_agent_details = UpdateAgentDetails(freeform_tags=FREEFORM_TAGS)
+        if name is not None:
+            update_agent_details.display_name = name
+        if description is not None:
+            update_agent_details.description = description
+        if instructions is not None:
+            llm_customization = LlmCustomization(
+                instruction=instructions,
+            )
+            update_agent_details.llm_config = LlmConfig(
+                routing_llm_customization=llm_customization,
+            )
+        self._invoke(
+            method=self._mgmt_client.update_agent,
+            success_message="Agent updated successfully",
+            error_message="Failed to update agent",
+            debug=self.debug,
+            update_agent_details=update_agent_details,
+            agent_id=agent_id,
+        )
+
+    def find_tools(self, compartment_id: str, agent_id: str) -> List[Dict[str, Any]]:
+        """
+        Find tools associated with the agent.
+
+        Args:
+            compartment_id: The compartment ID
+            agent_id: The agent ID
+
+        Returns:
+            List of tools
+
+        Raises:
+            AgentError: If finding tools fails
+        """
+        # use pagination to eagerly get all records
+        find_tools_response: Response = self._invoke(
+            method=list_call_get_all_results,
+            success_message="Tools retrieval succeeded",
+            error_message="Failed to find tools",
+            debug=self.debug,
+            list_func_ref=self._mgmt_client.list_tools,
+            compartment_id=compartment_id,
+            agent_id=agent_id,
+        )
+
+        return [to_dict(tool) for tool in find_tools_response.data]
+
+    def delete_tool(self, tool_id: str) -> None:
+        """
+        Delete a tool by ID.
+
+        Args:
+            tool_id: The ID of the tool to delete
+
+        Returns:
+            None
+
+        Raises:
+            AgentError: If tool deletion fails
+        """
+        self._invoke(
+            method=self._mgmt_client.delete_tool,
+            success_message="Tool deletion succeeded",
+            error_message="Failed to delete tool",
+            debug=self.debug,
+            tool_id=tool_id,
+        )
+
+    def add_function_tool(
+        self,
+        function: FunctionTool,
+        compartment_id: str,
+        agent_id: str,
+    ) -> Dict[str, Any]:
+        """
+        Add a function tool to the agent.
+
+        Args:
+            function: The function tool to add
+
+        Returns:
+            Dict[str, Any]: The added tool as a dict with its attributes
+
+        Raises:
+            AgentError: If adding the tool fails
+        """
+        tool_config = FunctionCallingToolConfig(
+            function=Function(
+                name=function.name,
+                description=function.description,
+                parameters=build_custom_function_params(function.parameters),
+            )
+        )
+        create_tool_details = CreateToolDetails(
+            display_name=function.name,
+            description=function.description,
+            compartment_id=compartment_id,
+            agent_id=agent_id,
+            tool_config=tool_config,
+            freeform_tags=FREEFORM_TAGS,
+        )
+
+        add_tool_response: Response = self._invoke(
+            method=self._mgmt_client.create_tool,
+            success_message="Tool addition succeeded",
+            error_message="Failed to add tool",
+            debug=self.debug,
+            create_tool_details=create_tool_details,
+        )
+
+        return to_dict(add_tool_response.data)
+
+    def get_tool(self, tool_id: str) -> Dict[str, Any]:
+        """
+        Get a tool by ID.
+
+        Args:
+            tool_id: The ID of the tool to get
+
+        Returns:
+            Dict[str, Any]: The tool as a dict with its attributes
+
+        Raises:
+            AgentError: If getting the tool fails
+        """
+        get_tool_response: Response = self._invoke(
+            method=self._mgmt_client.get_tool,
+            success_message="Tool retrieval succeeded",
+            error_message="Failed to get tool",
+            debug=self.debug,
+            tool_id=tool_id,
+        )
+
+        return to_dict(get_tool_response.data)
+
+    def add_rag_tool(
+        self,
+        rag_tool: AgenticRagTool,
+        compartment_id: str,
+        agent_id: str,
+    ) -> Dict[str, Any]:
+        """
+        Add an agentic RAG tool to the agent.
+
+        Args:
+            rag_tool: The RAG tool to add
+            compartment_id: The compartment ID
+            agent_id: The agent ID
+
+        Returns:
+            Dict[str, Any]: The added tool as a dict with its attributes
+
+        Raises:
+            AgentError: If adding the tool fails
+        """
+        tool_config = RagToolConfig(
+            knowledge_base_configs=[
+                KnowledgeBaseConfig(knowledge_base_id=kb_id)
+                for kb_id in rag_tool.knowledge_base_ids
+            ]
+        )
+        create_tool_details = CreateToolDetails(
+            compartment_id=compartment_id,
+            agent_id=agent_id,
+            display_name=rag_tool.name,
+            description=rag_tool.description,
+            tool_config=tool_config,
+            freeform_tags=FREEFORM_TAGS,
+        )
+
+        add_tool_response: Response = self._invoke(
+            method=self._mgmt_client.create_tool,
+            success_message="RAG tool addition succeeded",
+            error_message="""
+            Failed to add RAG tool
+            probably because a remote RAG tool not configured by ADK already exists
+            """,
+            debug=self.debug,
+            create_tool_details=create_tool_details,
+        )
+
+        return to_dict(add_tool_response.data)
+
+    def update_function_tool(
+        self,
+        tool_id: str,
+        function: FunctionTool,
+    ) -> None:
+        """
+        Update a function tool by ID.
+
+        Args:
+            tool_id: The ID of the tool to update
+            function: The function tool to update
+
+        Returns:
+            None
+
+        Raises:
+            AgentError: If updating the tool fails
+        """
+        tool_config = FunctionCallingToolConfig(
+            function=Function(
+                name=function.name,
+                description=function.description,
+                parameters=build_custom_function_params(function.parameters),
+            )
+        )
+        update_tool_details = UpdateToolDetails(
+            display_name=function.name,
+            description=function.description,
+            tool_config=tool_config,
+            freeform_tags=FREEFORM_TAGS,
+        )
+
+        self._invoke(
+            method=self._mgmt_client.update_tool,
+            success_message="Tool update succeeded",
+            error_message="Failed to update tool",
+            debug=self.debug,
+            tool_id=tool_id,
+            update_tool_details=update_tool_details,
+        )
+
+    def update_rag_tool(
+        self,
+        tool_id: str,
+        rag_tool: AgenticRagTool,
+    ) -> None:
+        """
+        Update a rag tool by ID.
+
+        Args:
+            tool_id: The ID of the tool to update
+            rag_tool: The rag tool to update
+
+        Returns:
+            None
+
+        Raises:
+            AgentError: If updating the tool fails
+        """
+        tool_config = RagToolConfig(
+            knowledge_base_configs=[
+                KnowledgeBaseConfig(knowledge_base_id=kb_id)
+                for kb_id in rag_tool.knowledge_base_ids
+            ]
+        )
+        update_tool_details = UpdateToolDetails(
+            display_name=rag_tool.name,
+            description=rag_tool.description,
+            tool_config=tool_config,
+            freeform_tags=FREEFORM_TAGS,
+        )
+
+        self._invoke(
+            method=self._mgmt_client.update_tool,
+            success_message="RAG tool update succeeded",
+            error_message="Failed to update RAG tool",
+            debug=self.debug,
+            tool_id=tool_id,
+            update_tool_details=update_tool_details,
+        )
+
+    def get_agent(self, agent_id: str) -> Dict[str, Any]:
+        """
+        Get the agent details.
+
+        Args:
+            agent_id: The ID of the agent
+
+        Returns:
+            The agent details
+
+        Raises:
+            AgentError: If getting the agent fails
+        """
+
+        get_agent_response: Response = self._invoke(
+            method=self._mgmt_client.get_agent,
+            success_message="Agent retrieval succeeded",
+            error_message="Failed to get agent",
+            debug=self.debug,
+            agent_id=agent_id,
+        )
+
+        return to_dict(get_agent_response.data)
+
+    def get_agent_endpoint_details(self, agent_endpoint_id: str) -> Dict[str, Any]:
+        """
+        Get the agent endpoint details.
+
+        Args:
+            agent_endpoint_id: The agent endpoint id
+
+        Returns:
+            The agent endpoint details
+
+        Raises:
+            AgentError: If getting the agent endpoint fails
+        """
+        get_agent_endpoint_response: Response = self._invoke(
+            method=self._mgmt_client.get_agent_endpoint,
+            success_message="Agent endpoint retrieval succeeded",
+            error_message="Failed to get agent endpoint",
+            debug=self.debug,
+            agent_endpoint_id=agent_endpoint_id,
+        )
+
+        return to_dict(get_agent_endpoint_response.data)
+
+    def get_runtime_endpoint(
+        self, endpoint_url: Optional[str], region: Optional[str]
+    ) -> str:
+        """
+        Get the runtime endpoint.
+
+        Args:
+            endpoint_url: The runtime endpoint
+            region: The region
+
+        Returns:
+            The runtime endpoint
+        """
+        return self._get_service_endpoint(endpoint_url, region, "agent-runtime")
+
+    def get_management_endpoint(
+        self, endpoint_url: Optional[str], region: Optional[str]
+    ) -> str:
+        """
+        Get the management endpoint.
+
+        Args:
+            endpoint_url: The management endpoint
+            region: The region
+
+        Returns:
+            The management endpoint
+        """
+        return self._get_service_endpoint(endpoint_url, region, "agent")
+
+    def _get_service_endpoint(
+        self,
+        endpoint_url: Optional[str] = None,
+        region: Optional[str] = None,
+        endpoint_type: Literal["agent-runtime", "agent"] = "agent-runtime",
+    ) -> str:
+        """
+        Set the service endpoint.
+
+        Args:
+            endpoint_url: The service endpoint
+            region: The region
+            endpoint_type: The endpoint type
+
+        Returns:
+            The service endpoint
+        """
+        # use endpoint if provided, otherwise use region
+        if endpoint_url:
+            return endpoint_url
+        elif region:
+            return get_region_endpoint(region, endpoint_type)
+        else:
+            raise UserError(
+                message="At least one of endpoint_url or region must be provided"
+            )
+
+    def _invoke(
+        self,
+        method: Callable,
+        success_message: str,
+        error_message: str,
+        debug: Optional[bool],
+        **kwargs,
+    ) -> Response:
+        """
+        Invoke a method.
+
+        Args:
+            method: The method to invoke
+            success_message: The success message for debug purposes
+            error_message: The error message
+            debug: Whether to enable debug logging
+            **kwargs: Additional keyword arguments to pass in method
+
+        Returns:
+            The response
+        """
+        response: Response | None = None
+        try:
+            response = method(**kwargs)
+            if response is None:
+                raise UserError(
+                    message=error_message,
+                    response=response,
+                )
+
+            if response.status not in (200, 202):
+                raise AgentError(message=error_message, response=response)
+
+            logger.debug(
+                f"{success_message} response: {formatted_flat_dict(response.data)}"
+            )
+            logger.debug(f"{success_message} {HEADER_REQUEST_ID}: {response.request_id}")
+
+            return response
+
+        except ServiceError as e:
+            raise AgentError(
+                message=f"{error_message}: {str(e)}",
+                response=response,
+            ) from e
+
+        except (ValueError, ClientError) as e:
+            raise UserError(
+                message=f"{error_message}: {str(e)}",
+                response=response,
+            ) from e

oci/addons/adk/agent_error.py

@@ -0,0 +1,36 @@
+# coding: utf-8
+# Copyright (c) 2016, 2025, Oracle and/or its affiliates.  All rights reserved.
+# This software is dual-licensed to you under the Universal Permissive License (UPL) 1.0 as shown at https://oss.oracle.com/licenses/upl or Apache License 2.0 as shown at http://www.apache.org/licenses/LICENSE-2.0. You may choose either license.
+
+from oci.response import Response
+
+
+class AgentException(Exception):
+    """Base exception class for OCI Generative AI ADK errors"""
+
+    def __init__(self, message: str = "", response: Response | None = None):
+        self.message = message
+        self.response = response
+
+        if response is not None:
+            self.status_code = response.status
+            self.data = response.data
+            self.headers = response.headers
+        else:
+            self.status_code = None
+            self.data = None
+            self.headers = None
+
+    def __str__(self):
+        if self.response is not None:
+            return f"{self.message}. Status code: {self.status_code}. Response data: {self.data}. Response headers: {self.headers}"  # noqa: E501
+        else:
+            return f"{self.message}"
+
+
+class AgentError(AgentException):
+    """Exception for OCI Agent service errors"""
+
+
+class UserError(AgentException):
+    """Exception for user errors"""