From db7ef635c0e061fcbab2f608ccc60af15fc5585d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Rub=C3=A9n=20Mart=C3=ADnez?= <ruben.martinez1@gmail.com>
Date: Fri, 9 Jun 2023 08:21:11 +0200
Subject: [PATCH] Add support for the endpoint URL in
 DynamoDBChatMesasgeHistory (#5836)

This PR adds the possibility of specifying the endpoint URL to AWS in
the DynamoDBChatMessageHistory, so that it is possible to target not
only the AWS cloud services, but also a local installation.

Specifying the endpoint URL, which is normally not done when addressing
the cloud services, is very helpful when targeting a local instance
(like [Localstack](https://localstack.cloud/)) when running local tests.

Fixes #5835

#### Who can review?

Tag maintainers/contributors who might be interested: @dev2049

<!-- For a quicker response, figure out the right person to tag with @

  @hwchase17 - project lead

  Tracing / Callbacks
  - @agola11

  Async
  - @agola11

  DataLoaders
  - @eyurtsev

  Models
  - @hwchase17
  - @agola11

  Agents / Tools / Toolkits
  - @vowelparrot

  VectorStores / Retrievers / Memory
  - @dev2049

 -->

---------

Co-authored-by: Harrison Chase <hw.chase.17@gmail.com>
---
 .../dynamodb_chat_message_history.ipynb       | 23 +++++++++++++++++++
 .../memory/chat_message_histories/dynamodb.py | 15 +++++++++---
 2 files changed, 35 insertions(+), 3 deletions(-)

diff --git a/docs/modules/memory/examples/dynamodb_chat_message_history.ipynb b/docs/modules/memory/examples/dynamodb_chat_message_history.ipynb
index 24d12364..302c11bc 100644
--- a/docs/modules/memory/examples/dynamodb_chat_message_history.ipynb
+++ b/docs/modules/memory/examples/dynamodb_chat_message_history.ipynb
@@ -118,6 +118,29 @@
    ]
   },
   {
+   "cell_type": "markdown",
+   "id": "955f1b15",
+   "metadata": {},
+   "source": [
+    "## DynamoDBChatMessageHistory with Custom Endpoint URL\n",
+    "\n",
+    "Sometimes it is useful to specify the URL to the AWS endpoint to connect to. For instance, when you are running locally against [Localstack](https://localstack.cloud/). For those cases you can specify the URL via the `endpoint_url` parameter in the constructor."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "225713c8",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from langchain.memory.chat_message_histories import DynamoDBChatMessageHistory\n",
+    "\n",
+    "history = DynamoDBChatMessageHistory(table_name=\"SessionTable\", session_id=\"0\", endpoint_url=\"http://localhost.localstack.cloud:4566\")"
+   ]
+  },
+  {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "3b33c988",
    "metadata": {},
diff --git a/langchain/memory/chat_message_histories/dynamodb.py b/langchain/memory/chat_message_histories/dynamodb.py
index 2fc05688..dfde1457 100644
--- a/langchain/memory/chat_message_histories/dynamodb.py
+++ b/langchain/memory/chat_message_histories/dynamodb.py
@@ -1,5 +1,5 @@
 import logging
-from typing import List
+from typing import List, Optional
 
 from langchain.schema import (
     BaseChatMessageHistory,
@@ -21,12 +21,21 @@ class DynamoDBChatMessageHistory(BaseChatMessageHistory):
         table_name: name of the DynamoDB table
         session_id: arbitrary key that is used to store the messages
             of a single chat session.
+        endpoint_url: URL of the AWS endpoint to connect to. This argument
+            is optional and useful for test purposes, like using Localstack.
+            If you plan to use AWS cloud service, you normally don't have to
+            worry about setting the endpoint_url.
     """
 
-    def __init__(self, table_name: str, session_id: str):
+    def __init__(
+        self, table_name: str, session_id: str, endpoint_url: Optional[str] = None
+    ):
         import boto3
 
-        client = boto3.resource("dynamodb")
+        if endpoint_url:
+            client = boto3.resource("dynamodb", endpoint_url=endpoint_url)
+        else:
+            client = boto3.resource("dynamodb")
         self.table = client.Table(table_name)
         self.session_id = session_id