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__)


[docs] @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()