Refactor typing import in types.py

This commit is contained in:
Gabriel Luiz Freitas Almeida 2024-02-27 17:57:29 -03:00
commit 4007b9728f

View file

@ -1,7 +1,6 @@
import ast import ast
import json import json
from typing import (AsyncIterator, Callable, Dict, Iterator, List, Optional, from typing import AsyncIterator, Callable, Dict, Iterator, List, Optional, Union
Union)
import yaml import yaml
from langchain_core.messages import AIMessage from langchain_core.messages import AIMessage
@ -396,7 +395,7 @@ class ChatVertex(StatelessVertex):
if artifacts: if artifacts:
self.artifacts = artifacts.model_dump() self.artifacts = artifacts.model_dump()
if isinstance(self._built_object, (AsyncIterator, Iterator)): if isinstance(self._built_object, (AsyncIterator, Iterator)):
if self.params["as_record"]: if self.params["return_record"]:
self._built_object = Record(text=message, data=self.artifacts) self._built_object = Record(text=message, data=self.artifacts)
else: else:
self._built_object = message self._built_object = message