From 0545d45e4a1f7f321a3247f5988da96b2b37a805 Mon Sep 17 00:00:00 2001
From: Massimiliano Angelino <angmas@amazon.com>
Date: Thu, 7 Mar 2024 10:29:36 +0100
Subject: [PATCH] feat(mistral): separate adapter

---
 .../request-handler/adapters/bedrock/__init__.py     |  1 +
 .../request-handler/adapters/bedrock/llama2_chat.py  |  2 +-
 .../request-handler/adapters/bedrock/mistral.py      | 12 ++++++++++++
 3 files changed, 14 insertions(+), 1 deletion(-)
 create mode 100644 lib/model-interfaces/langchain/functions/request-handler/adapters/bedrock/mistral.py

diff --git a/lib/model-interfaces/langchain/functions/request-handler/adapters/bedrock/__init__.py b/lib/model-interfaces/langchain/functions/request-handler/adapters/bedrock/__init__.py
index 035aaa661..38dd3b4ce 100644
--- a/lib/model-interfaces/langchain/functions/request-handler/adapters/bedrock/__init__.py
+++ b/lib/model-interfaces/langchain/functions/request-handler/adapters/bedrock/__init__.py
@@ -3,3 +3,4 @@
 from .ai21_j2 import *
 from .cohere import *
 from .llama2_chat import *
+from .mistral import *
diff --git a/lib/model-interfaces/langchain/functions/request-handler/adapters/bedrock/llama2_chat.py b/lib/model-interfaces/langchain/functions/request-handler/adapters/bedrock/llama2_chat.py
index b9b81da91..0146cc84a 100644
--- a/lib/model-interfaces/langchain/functions/request-handler/adapters/bedrock/llama2_chat.py
+++ b/lib/model-interfaces/langchain/functions/request-handler/adapters/bedrock/llama2_chat.py
@@ -62,6 +62,6 @@ def get_condense_question_prompt(self):
 
 # Register the adapter
 registry.register(
-    r"(^bedrock.meta.llama2-.*-chat.*|^bedrock.mistral.mi.*)",
+    r"^bedrock.meta.llama2-.*-chat.*",
     BedrockMetaLLama2ChatAdapter,
 )
diff --git a/lib/model-interfaces/langchain/functions/request-handler/adapters/bedrock/mistral.py b/lib/model-interfaces/langchain/functions/request-handler/adapters/bedrock/mistral.py
new file mode 100644
index 000000000..9ed9aab66
--- /dev/null
+++ b/lib/model-interfaces/langchain/functions/request-handler/adapters/bedrock/mistral.py
@@ -0,0 +1,12 @@
+from .llama2_chat import BedrockMetaLLama2ChatAdapter
+from ..registry import registry
+
+
+class BedrockMistralAdapter(BedrockMetaLLama2ChatAdapter): ...
+
+
+# Register the adapter
+registry.register(
+    r"^bedrock.mistral.mi.*",
+    BedrockMistralAdapter,
+)