refactor: Update OpenAIModelComponent to use BaseLanguageModel and langflow template
Refactor the OpenAIModelComponent in OpenAIModel.py to use the BaseLanguageModel field type from langflow.field_typing and the langflow.template module. This change ensures compatibility with the latest version of the langflow library and improves code readability.
This commit is contained in:
ogabrielluiz
parent
2ae46f2b08
commit
2006636229
1 changed files with 58 additions and 74 deletions
|
|
@ -6,7 +6,8 @@ from pydantic.v1 import SecretStr
|
|||
from langflow.base.constants import STREAM_INFO_TEXT
|
||||
from langflow.base.models.model import LCModelComponent
|
||||
from langflow.base.models.openai_constants import MODEL_NAMES
|
||||
from langflow.field_typing import NestedDict, Text
|
||||
from langflow.field_typing import BaseLanguageModel, Text
|
||||
from langflow.template import Input, Output
|
||||
|
||||
|
||||
class OpenAIModelComponent(LCModelComponent):
|
||||
|
|
@ -14,80 +15,64 @@ class OpenAIModelComponent(LCModelComponent):
|
|||
description = "Generates text using OpenAI LLMs."
|
||||
icon = "OpenAI"
|
||||
|
||||
field_order = [
|
||||
"max_tokens",
|
||||
"model_kwargs",
|
||||
"model_name",
|
||||
"openai_api_base",
|
||||
"openai_api_key",
|
||||
"temperature",
|
||||
"input_value",
|
||||
"system_message",
|
||||
"stream",
|
||||
inputs = [
|
||||
Input(name="input_value", type=str, display_name="Input", input_types=["Text", "Record", "Prompt"]),
|
||||
Input(
|
||||
name="max_tokens",
|
||||
type=Optional[int],
|
||||
display_name="Max Tokens",
|
||||
advanced=True,
|
||||
info="The maximum number of tokens to generate. Set to 0 for unlimited tokens.",
|
||||
),
|
||||
Input(name="model_kwargs", type=dict, display_name="Model Kwargs", advanced=True),
|
||||
Input(name="model_name", type=str, display_name="Model Name", advanced=False, options=MODEL_NAMES),
|
||||
Input(
|
||||
name="openai_api_base",
|
||||
type=Optional[str],
|
||||
display_name="OpenAI API Base",
|
||||
advanced=True,
|
||||
info="The base URL of the OpenAI API. Defaults to https://api.openai.com/v1.\n\nYou can change this to use other APIs like JinaChat, LocalAI and Prem.",
|
||||
),
|
||||
Input(
|
||||
name="openai_api_key",
|
||||
type=str,
|
||||
display_name="OpenAI API Key",
|
||||
info="The OpenAI API Key to use for the OpenAI model.",
|
||||
advanced=False,
|
||||
password=True,
|
||||
),
|
||||
Input(name="temperature", type=float, display_name="Temperature", advanced=False, default=0.1),
|
||||
Input(name="stream", type=bool, display_name="Stream", info=STREAM_INFO_TEXT, advanced=True),
|
||||
Input(
|
||||
name="system_message",
|
||||
type=Optional[str],
|
||||
display_name="System Message",
|
||||
info="System message to pass to the model.",
|
||||
advanced=True,
|
||||
),
|
||||
]
|
||||
outputs = [
|
||||
Output(display_name="Text", name="text_output", method="text_response"),
|
||||
Output(display_name="Language Model", name="model_output", method="model_response"),
|
||||
]
|
||||
|
||||
def build_config(self):
|
||||
return {
|
||||
"input_value": {"display_name": "Input", "input_types": ["Text", "Record", "Prompt"]},
|
||||
"max_tokens": {
|
||||
"display_name": "Max Tokens",
|
||||
"advanced": True,
|
||||
"info": "The maximum number of tokens to generate. Set to 0 for unlimited tokens.",
|
||||
},
|
||||
"model_kwargs": {
|
||||
"display_name": "Model Kwargs",
|
||||
"advanced": True,
|
||||
},
|
||||
"model_name": {
|
||||
"display_name": "Model Name",
|
||||
"advanced": False,
|
||||
"options": MODEL_NAMES,
|
||||
},
|
||||
"openai_api_base": {
|
||||
"display_name": "OpenAI API Base",
|
||||
"advanced": True,
|
||||
"info": (
|
||||
"The base URL of the OpenAI API. Defaults to https://api.openai.com/v1.\n\n"
|
||||
"You can change this to use other APIs like JinaChat, LocalAI and Prem."
|
||||
),
|
||||
},
|
||||
"openai_api_key": {
|
||||
"display_name": "OpenAI API Key",
|
||||
"info": "The OpenAI API Key to use for the OpenAI model.",
|
||||
"advanced": False,
|
||||
"password": True,
|
||||
},
|
||||
"temperature": {
|
||||
"display_name": "Temperature",
|
||||
"advanced": False,
|
||||
"value": 0.1,
|
||||
},
|
||||
"stream": {
|
||||
"display_name": "Stream",
|
||||
"info": STREAM_INFO_TEXT,
|
||||
"advanced": True,
|
||||
},
|
||||
"system_message": {
|
||||
"display_name": "System Message",
|
||||
"info": "System message to pass to the model.",
|
||||
"advanced": True,
|
||||
},
|
||||
}
|
||||
def text_response(self) -> Text:
|
||||
input_value = self.input_value
|
||||
stream = self.stream
|
||||
system_message = self.system_message
|
||||
output = self.model_response()
|
||||
result = self.get_chat_result(output, stream, input_value, system_message)
|
||||
self.status = result
|
||||
return result
|
||||
|
||||
def model_response(self) -> BaseLanguageModel:
|
||||
openai_api_key = self.openai_api_key
|
||||
temperature = self.temperature
|
||||
model_name = self.model_name
|
||||
max_tokens = self.max_tokens
|
||||
model_kwargs = self.model_kwargs
|
||||
openai_api_base = self.openai_api_base or "https://api.openai.com/v1"
|
||||
|
||||
def build(
|
||||
self,
|
||||
input_value: Text,
|
||||
openai_api_key: str,
|
||||
temperature: float = 0.1,
|
||||
model_name: str = "gpt-3.5-turbo",
|
||||
max_tokens: Optional[int] = 256,
|
||||
model_kwargs: NestedDict = {},
|
||||
openai_api_base: Optional[str] = None,
|
||||
stream: bool = False,
|
||||
system_message: Optional[str] = None,
|
||||
) -> Text:
|
||||
if not openai_api_base:
|
||||
openai_api_base = "https://api.openai.com/v1"
|
||||
if openai_api_key:
|
||||
api_key = SecretStr(openai_api_key)
|
||||
else:
|
||||
|
|
@ -101,5 +86,4 @@ class OpenAIModelComponent(LCModelComponent):
|
|||
api_key=api_key,
|
||||
temperature=temperature,
|
||||
)
|
||||
|
||||
return self.get_chat_result(output, stream, input_value, system_message)
|
||||
return output
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue