Initial commit of Amazon Q Business runnable for langchain

libs/aws/langchain_aws/llms/q_business.py

@@ -0,0 +1,185 @@
+from typing import Any, Dict, List, Optional
+from langchain_core.callbacks.manager import (
+    AsyncCallbackManagerForLLMRun,
+    CallbackManagerForLLMRun
+)
+import logging
+from langchain_core.language_models import LLM
+from pydantic import ConfigDict, model_validator
+import json
+import asyncio
+import boto3
+from typing_extensions import Self
+
+class AmazonQ(LLM):
+    """Amazon Q LLM 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."""
+
+    streaming: bool = False
+    """Whether to stream the results or not."""
+
+    client: Any = None
+    """Amazon Q client."""
+
+    application_id: str = None
+    """Amazon Q client."""
+
+    _last_response: Dict = None  # Add this to store the full response
+    """Store the full response from Amazon Q."""
+
+    parent_message_id: Optional[str] = None
+    """AWS region name. If not provided, will be extracted from environment."""
+
+    conversation_id: Optional[str] = None
+    """AWS region name. If not provided, will be extracted from environment."""
+
+    chat_mode: str = "RETRIEVAL_MODE"
+    """AWS region name. If not provided, will be extracted from environment."""
+
+    credentials_profile_name: Optional[str] = None
+    """The name of the profile in the ~/.aws/credentials or ~/.aws/config files, which
+    has either access keys or role information specified.
+    If not specified, the default credential profile or, if on an EC2 instance,
+    credentials from IMDS will be used.
+    See: https://boto3.amazonaws.com/v1/documentation/api/latest/guide/credentials.html
+    """
+    model_config = ConfigDict(
+        extra="forbid",
+    )
+
+    @property
+    def _llm_type(self) -> str:
+        """Return type of llm."""
+        return "amazon_q"
+
+    def __init__(self, **kwargs: Any) -> None:
+        """Initialize the Amazon Q client."""
+        super().__init__(**kwargs)
+
+    def _call(
+        self,
+        prompt: str,
+        stop: Optional[List[str]] = None,
+        run_manager: Optional[CallbackManagerForLLMRun] = None,
+        **kwargs: Any,
+    ) -> str:
+        """Call out to Amazon Q service.
+
+        Args:
+            prompt: The prompt to pass into the model.
+            stop: Optional list of stop words to use when generating.
+
+        Returns:
+            The string generated by the model.
+
+        Example:
+            .. code-block:: python
+
+                response = llm.invoke("Tell me a joke.")
+        """
+        try:
+            print("Prompt Length (Amazon Q ChatSync API takes a maximum of 7000 chars)")
+            print(len(prompt))
+
+            # Prepare the request
+            request = {
+                'applicationId': self.application_id,
+                'userMessage': prompt,
+                '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)
+            self._last_response = response
+
+            # Extract the response text
+            if 'systemMessage' in response:
+                return response["systemMessage"]
+            else:
+                raise ValueError("Unexpected response format from Amazon Q")
+
+        except Exception as e:
+            if "Prompt Length" in str(e):
+                logging.info(f"Prompt Length: {len(prompt)}")
+                print(f"""Prompt:
+                {prompt}""")
+                raise ValueError(f"Error raised by Amazon Q service: {e}")
+
+            raise ValueError(f"Error raised by Amazon Q service: {e}")
+
+    def get_last_response(self) -> Dict:
+        """Method to access the full response from the last call"""
+        return self._last_response
+
+    async def _acall(
+        self,
+        prompt: str,
+        stop: Optional[List[str]] = None,
+        run_manager: Optional[AsyncCallbackManagerForLLMRun] = None,
+        **kwargs: Any,
+    ) -> str:
+        """Async call to Amazon Q service."""
+
+        def _execute_call():
+            return self._call(prompt, stop=stop, **kwargs)
+
+        # Run the synchronous call in a thread pool
+        return await asyncio.get_running_loop().run_in_executor(
+            None, _execute_call
+        )
+
+    @property
+    def _identifying_params(self) -> Dict[str, Any]:
+        """Get the identifying parameters."""
+        return {
+            "region_name": self.region_name,
+        }
+    @model_validator(mode="after")
+    def validate_environment(self) -> Self:
+        """Dont do anything if client provided externally"""
+        if self.client is not None:
+            return self
+
+        """Validate that AWS credentials to and python package exists in environment."""
+        try:
+            import boto3
+
+            try:
+                if self.credentials_profile_name is not None:
+                    session = boto3.Session(profile_name=self.credentials_profile_name)
+                else:
+                    # use default credentials
+                    session = boto3.Session()
+
+                self.client = session.client(
+                    "qbusiness", region_name=self.region_name
+                )
+
+            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 self