From af264d15b2700520c339f5d51e4bb5508f797f67 Mon Sep 17 00:00:00 2001
From: Leonid Kuligin <lkuligin@yandex.ru>
Date: Thu, 22 Aug 2024 12:48:24 +0200
Subject: [PATCH] fix model_garden models (#462)

---
 .../langchain_google_vertexai/_base.py        | 19 ++++++++-----------
 libs/vertexai/pyproject.toml                  |  2 +-
 2 files changed, 9 insertions(+), 12 deletions(-)

diff --git a/libs/vertexai/langchain_google_vertexai/_base.py b/libs/vertexai/langchain_google_vertexai/_base.py
index d71bfff3..48aa2923 100644
--- a/libs/vertexai/langchain_google_vertexai/_base.py
+++ b/libs/vertexai/langchain_google_vertexai/_base.py
@@ -374,22 +374,19 @@ def _parse_response(self, predictions: "Prediction") -> LLMResult:
         return LLMResult(generations=generations)
 
     def _parse_prediction(self, prediction: Any) -> str:
-        if isinstance(prediction, str):
-            if prediction.startswith("Prompt:\n"):
-                result = re.search(r"(?s:.*)\nOutput:\n(.*)", prediction)
+        def _clean_response(response: str) -> str:
+            if response.startswith("Prompt:\n"):
+                result = re.search(r"(?s:.*)\nOutput:\n((?s:.*))", response)
                 if result:
                     return result[1]
-            return prediction
+            return response
+
+        if isinstance(prediction, str):
+            return _clean_response(prediction)
 
         if self.result_arg:
             try:
-                if prediction[self.result_arg].startswith("Prompt:\n"):
-                    result = re.search(
-                        r"(?s:.*)\nOutput:\n(.*)", prediction[self.result_arg]
-                    )
-                    if result:
-                        return result[1]
-                return prediction[self.result_arg]
+                return _clean_response(prediction[self.result_arg])
             except KeyError:
                 if isinstance(prediction, str):
                     error_desc = (
diff --git a/libs/vertexai/pyproject.toml b/libs/vertexai/pyproject.toml
index 5609bb42..00f6b259 100644
--- a/libs/vertexai/pyproject.toml
+++ b/libs/vertexai/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "langchain-google-vertexai"
-version = "1.0.9"
+version = "1.0.10"
 description = "An integration package connecting Google VertexAI and LangChain"
 authors = []
 readme = "README.md"