Source code for langchain_aws.runnables.q_business

import logging
from typing import Any, List, Optional, Union

from langchain_core._api.beta_decorator import beta
from langchain_core.messages.ai import AIMessage
from langchain_core.prompt_values import ChatPromptValue
from langchain_core.runnables import Runnable
from langchain_core.runnables.config import RunnableConfig
from pydantic import ConfigDict
from typing_extensions import Self

logger = logging.getLogger(__name__)




@beta(message="This API is in beta and can change in future.")
class AmazonQ(Runnable[Union[str,ChatPromptValue, List[ChatPromptValue]], ChatPromptValue]):
    """Amazon Q Runnable wrapper.

    To authenticate, the AWS client uses the following methods to
    automatically load credentials:
    https://boto3.amazonaws.com/v1/documentation/api/latest/guide/credentials.html

    Make sure the credentials / roles used have the required policies to
    access the Amazon Q service.
    """

    region_name: Optional[str] = None
    """AWS region name. If not provided, will be extracted from environment."""

    credentials: Optional[Any] = None
    """Amazon Q credentials used to instantiate the client if the client is not provided."""

    client: Optional[Any] = None
    """Amazon Q client."""

    application_id: str = None

    """Store the full response from Amazon Q."""

    parent_message_id: Optional[str] = None

    conversation_id: Optional[str] = None

    chat_mode: str = "RETRIEVAL_MODE"

    model_config = ConfigDict(
        extra="forbid",
    )

    def __init__(
        self,
        region_name: Optional[str] = None,
        credentials: Optional[Any] = None,
        client: Optional[Any] = None,
        application_id: str = None,
        parent_message_id: Optional[str] = None,
        conversation_id: Optional[str] = None,
        chat_mode: str = "RETRIEVAL_MODE",
    ):
        self.region_name = region_name
        self.credentials = credentials
        self.client = client or self.validate_environment()
        self.application_id = application_id
        self.parent_message_id = parent_message_id
        self.conversation_id = conversation_id
        self.chat_mode = chat_mode

    def invoke(
        self,
        input: Union[str,ChatPromptValue],
        config: Optional[RunnableConfig] = None,
        **kwargs: Any
    ) -> ChatPromptValue:
        """Call out to Amazon Q service.

        Args:
            input: The prompt to pass into the model.

        Returns:
            The string generated by the model.

        Example:
            .. code-block:: python

            model = AmazonQ(
                credentials=your_credentials,
                application_id=your_app_id
            )
            response = model.invoke("Tell me a joke")
        """
        try:
            # Prepare the request            
            request = {
                'applicationId': self.application_id,
                'userMessage': self.convert_langchain_messages_to_q_input(input), # Langchain's input comes in the form of an array of "messages". We must convert to a single string for Amazon Q's use
                'chatMode': self.chat_mode,
            }
            if self.conversation_id:
                request.update({
                    'conversationId': self.conversation_id,
                    'parentMessageId': self.parent_message_id,
                })

            # Call Amazon Q
            response = self.client.chat_sync(**request)

            # Extract the response text
            if 'systemMessage' in response:
                return AIMessage(content=response["systemMessage"], response_metadata=response)
            else:
                raise ValueError("Unexpected response format from Amazon Q")

        except Exception as e:
            if "Prompt Length" in str(e):
                logger.info(f"Prompt Length: {len(input)}")
                logger.info(f"""Prompt:
                {input}""")
            raise ValueError(f"Error raised by Amazon Q service: {e}")

    def validate_environment(self) -> Self:
        """Don't do anything if client provided externally"""
        #If the client is not provided, and the user_id is not provided in the class constructor, throw an error saying one or the other needs to be provided
        if self.credentials is None:
            raise ValueError(
                "Either the credentials or the client needs to be provided."
            )

        """Validate that AWS credentials to and python package exists in environment."""
        try:
            import boto3

            try:
                if self.region_name is not None:
                    client = boto3.client('qbusiness', self.region_name, **self.credentials)
                else:
                    # use default region
                    client = boto3.client('qbusiness', **self.credentials)

            except Exception as e:
                raise ValueError(
                    "Could not load credentials to authenticate with AWS client. "
                    "Please check that credentials in the specified "
                    "profile name are valid."
                ) from e

        except ImportError:
            raise ImportError(
                "Could not import boto3 python package. "
                "Please install it with `pip install boto3`."
            )
        return client
    def convert_langchain_messages_to_q_input(self, input: Union[str,ChatPromptValue,List[ChatPromptValue]]) -> str:
        #If it is just a string and not a ChatPromptTemplate collection just return string
        if type(input) is str:
            return input
        return input.to_string()