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"