80 lines
2.6 KiB
Python
80 lines
2.6 KiB
Python
from typing import Optional, Union
|
|
|
|
from langchain.llms import BaseLLM
|
|
from langchain_community.chat_models.openai import ChatOpenAI
|
|
from langflow import CustomComponent
|
|
from langflow.field_typing import BaseLanguageModel, NestedDict
|
|
|
|
|
|
class ChatOpenAIComponent(CustomComponent):
|
|
display_name = "ChatOpenAI"
|
|
description = "`OpenAI` Chat large language models API."
|
|
icon = "OpenAI"
|
|
|
|
def build_config(self):
|
|
return {
|
|
"max_tokens": {
|
|
"display_name": "Max Tokens",
|
|
"advanced": False,
|
|
"required": False,
|
|
},
|
|
"model_kwargs": {
|
|
"display_name": "Model Kwargs",
|
|
"advanced": True,
|
|
"required": False,
|
|
},
|
|
"model_name": {
|
|
"display_name": "Model Name",
|
|
"advanced": False,
|
|
"required": False,
|
|
"options": [
|
|
"gpt-4-turbo-preview",
|
|
"gpt-4-0125-preview",
|
|
"gpt-4-1106-preview",
|
|
"gpt-4-vision-preview",
|
|
"gpt-3.5-turbo-0125",
|
|
"gpt-3.5-turbo-1106",
|
|
],
|
|
},
|
|
"openai_api_base": {
|
|
"display_name": "OpenAI API Base",
|
|
"advanced": False,
|
|
"required": False,
|
|
"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",
|
|
"advanced": False,
|
|
"required": False,
|
|
"password": True,
|
|
},
|
|
"temperature": {
|
|
"display_name": "Temperature",
|
|
"advanced": False,
|
|
"required": False,
|
|
"value": 0.7,
|
|
},
|
|
}
|
|
|
|
def build(
|
|
self,
|
|
max_tokens: Optional[int] = 256,
|
|
model_kwargs: NestedDict = {},
|
|
model_name: str = "gpt-4-1106-preview",
|
|
openai_api_base: Optional[str] = None,
|
|
openai_api_key: Optional[str] = None,
|
|
temperature: float = 0.7,
|
|
) -> Union[BaseLanguageModel, BaseLLM]:
|
|
if not openai_api_base:
|
|
openai_api_base = "https://api.openai.com/v1"
|
|
return ChatOpenAI(
|
|
max_tokens=max_tokens,
|
|
model_kwargs=model_kwargs,
|
|
model=model_name,
|
|
base_url=openai_api_base,
|
|
api_key=openai_api_key,
|
|
temperature=temperature,
|
|
)
|