|
4 | 4 | import typing as t
|
5 | 5 | from abc import ABC, abstractmethod
|
6 | 6 |
|
7 |
| -from langchain.chat_models import AzureChatOpenAI, ChatOpenAI, BedrockChat |
| 7 | +from langchain.chat_models import AzureChatOpenAI, BedrockChat, ChatOpenAI, ChatVertexAI |
8 | 8 | from langchain.chat_models.base import BaseChatModel
|
9 |
| -from langchain.llms import AzureOpenAI, OpenAI, Bedrock |
| 9 | +from langchain.llms import AzureOpenAI, Bedrock, OpenAI, VertexAI |
10 | 10 | from langchain.llms.base import BaseLLM
|
11 | 11 | from langchain.schema import LLMResult
|
12 | 12 |
|
|
20 | 20 | def isOpenAI(llm: BaseLLM | BaseChatModel) -> bool:
|
21 | 21 | return isinstance(llm, OpenAI) or isinstance(llm, ChatOpenAI)
|
22 | 22 |
|
| 23 | + |
23 | 24 | def isBedrock(llm: BaseLLM | BaseChatModel) -> bool:
|
24 | 25 | return isinstance(llm, Bedrock) or isinstance(llm, BedrockChat)
|
25 | 26 |
|
| 27 | + |
26 | 28 | # have to specify it twice for runtime and static checks
|
27 |
| -MULTIPLE_COMPLETION_SUPPORTED = [OpenAI, ChatOpenAI, AzureOpenAI, AzureChatOpenAI] |
| 29 | +MULTIPLE_COMPLETION_SUPPORTED = [ |
| 30 | + OpenAI, |
| 31 | + ChatOpenAI, |
| 32 | + AzureOpenAI, |
| 33 | + AzureChatOpenAI, |
| 34 | + ChatVertexAI, |
| 35 | + VertexAI, |
| 36 | +] |
28 | 37 | MultipleCompletionSupportedLLM = t.Union[
|
29 |
| - OpenAI, ChatOpenAI, AzureOpenAI, AzureChatOpenAI |
| 38 | + OpenAI, ChatOpenAI, AzureOpenAI, AzureChatOpenAI, ChatVertexAI, VertexAI |
30 | 39 | ]
|
31 | 40 |
|
32 | 41 |
|
|
0 commit comments