Update ChatVertexAIComponent parameters and imports
This commit is contained in:
parent
ffb0dc0012
commit
5065234a1d
1 changed files with 8 additions and 7 deletions
|
|
@ -1,8 +1,10 @@
|
||||||
|
|
||||||
from langflow import CustomComponent
|
from langflow import CustomComponent
|
||||||
from typing import List, Optional
|
from typing import List, Optional, Union
|
||||||
from langchain_core.messages.base import BaseMessage
|
from langchain_core.messages.base import BaseMessage
|
||||||
from langchain_community.chat_models.vertexai import ChatVertexAI
|
from langchain_community.chat_models.vertexai import ChatVertexAI
|
||||||
|
from langflow.field_typing import BaseLanguageModel
|
||||||
|
from langchain.llms import BaseLLM
|
||||||
|
|
||||||
|
|
||||||
class ChatVertexAIComponent(CustomComponent):
|
class ChatVertexAIComponent(CustomComponent):
|
||||||
display_name = "ChatVertexAI"
|
display_name = "ChatVertexAI"
|
||||||
|
|
@ -13,7 +15,7 @@ class ChatVertexAIComponent(CustomComponent):
|
||||||
"credentials": {
|
"credentials": {
|
||||||
"display_name": "Credentials",
|
"display_name": "Credentials",
|
||||||
"field_type": "file",
|
"field_type": "file",
|
||||||
"fileTypes": ["json"],
|
"file_types": [".json"],
|
||||||
"file_path": None,
|
"file_path": None,
|
||||||
},
|
},
|
||||||
"examples": {
|
"examples": {
|
||||||
|
|
@ -60,8 +62,8 @@ class ChatVertexAIComponent(CustomComponent):
|
||||||
def build(
|
def build(
|
||||||
self,
|
self,
|
||||||
credentials: Optional[str],
|
credentials: Optional[str],
|
||||||
examples: Optional[List[BaseMessage]],
|
project: str,
|
||||||
project: Optional[str],
|
examples: Optional[List[BaseMessage]]=[],
|
||||||
location: Optional[str] = "us-central1",
|
location: Optional[str] = "us-central1",
|
||||||
max_output_tokens: Optional[int] = 128,
|
max_output_tokens: Optional[int] = 128,
|
||||||
model_name: Optional[str] = "chat-bison",
|
model_name: Optional[str] = "chat-bison",
|
||||||
|
|
@ -69,8 +71,7 @@ class ChatVertexAIComponent(CustomComponent):
|
||||||
top_k: Optional[int] = 40,
|
top_k: Optional[int] = 40,
|
||||||
top_p: Optional[float] = 0.95,
|
top_p: Optional[float] = 0.95,
|
||||||
verbose: Optional[bool] = False,
|
verbose: Optional[bool] = False,
|
||||||
):
|
) -> Union[BaseLanguageModel, BaseLLM]:
|
||||||
# Assuming there is a ChatVertexAI class that takes these parameters
|
|
||||||
return ChatVertexAI(
|
return ChatVertexAI(
|
||||||
credentials=credentials,
|
credentials=credentials,
|
||||||
examples=examples,
|
examples=examples,
|
||||||
|
|
|
||||||
Loading…
Add table
Add a link
Reference in a new issue