jyapayne/langflow
1
0
Fork 0
Code Issues Pull requests Projects Releases 1 Packages Wiki Activity Actions

Update ChatVertexAIComponent parameters

Browse source
This commit is contained in:
anovazzi1 2024-01-10 11:31:22 -03:00
commit 44a17f6f25
1 changed files with 21 additions and 20 deletions
Download patch file Download diff file Expand all files Collapse all files

41
src/backend/langflow/components/llms/ChatVertexAI.py
View file

@ -1,7 +1,8 @@
from langflow import CustomComponent from langflow import CustomComponent
from typing import List from typing import List, Optional
from langchain.messages import BaseMessage from langchain_core.messages.base import BaseMessage
from langchain_community.chat_models.vertexai import ChatVertexAI
class ChatVertexAIComponent(CustomComponent): class ChatVertexAIComponent(CustomComponent):
display_name = "ChatVertexAI" display_name = "ChatVertexAI"
@ -11,7 +12,7 @@ class ChatVertexAIComponent(CustomComponent):
return { return {
"credentials": { "credentials": {
"display_name": "Credentials", "display_name": "Credentials",
"type": "file", "field_type": "file",
"fileTypes": ["json"], "fileTypes": ["json"],
"file_path": None, "file_path": None,
}, },
@ -21,53 +22,53 @@ class ChatVertexAIComponent(CustomComponent):
}, },
"location": { "location": {
"display_name": "Location", "display_name": "Location",
"default": "us-central1", "value": "us-central1",
}, },
"max_output_tokens": { "max_output_tokens": {
"display_name": "Max Output Tokens", "display_name": "Max Output Tokens",
"default": 128, "value": 128,
"advanced": True, "advanced": True,
}, },
"model_name": { "model_name": {
"display_name": "Model Name", "display_name": "Model Name",
"default": "chat-bison", "value": "chat-bison",
}, },
"project": { "project": {
"display_name": "Project", "display_name": "Project",
}, },
"temperature": { "temperature": {
"display_name": "Temperature", "display_name": "Temperature",
"default": 0.0, "value": 0.0,
}, },
"top_k": { "top_k": {
"display_name": "Top K", "display_name": "Top K",
"default": 40, "value": 40,
"advanced": True, "advanced": True,
}, },
"top_p": { "top_p": {
"display_name": "Top P", "display_name": "Top P",
"default": 0.95, "value": 0.95,
"advanced": True, "advanced": True,
}, },
"verbose": { "verbose": {
"display_name": "Verbose", "display_name": "Verbose",
"default": False, "value": False,
"advanced": True, "advanced": True,
}, },
} }
def build( def build(
self, self,
credentials: str, credentials: Optional[str],
examples: List[BaseMessage], examples: Optional[List[BaseMessage]],
project: str, project: Optional[str],
location: str = "us-central1", location: Optional[str] = "us-central1",
max_output_tokens: int = 128, max_output_tokens: Optional[int] = 128,
model_name: str = "chat-bison", model_name: Optional[str] = "chat-bison",
temperature: float = 0.0, temperature: Optional[float] = 0.0,
top_k: int = 40, top_k: Optional[int] = 40,
top_p: float = 0.95, top_p: Optional[float] = 0.95,
verbose: bool = False, verbose: Optional[bool] = False,
): ):
# Assuming there is a ChatVertexAI class that takes these parameters # Assuming there is a ChatVertexAI class that takes these parameters
return ChatVertexAI( return ChatVertexAI(