Refactor model components to support streaming
This commit is contained in:
parent
13317d61e2
commit
6c2a35afb1
11 changed files with 114 additions and 40 deletions
|
|
@ -35,6 +35,10 @@ class AmazonBedrockComponent(CustomComponent):
|
||||||
"cache": {"display_name": "Cache"},
|
"cache": {"display_name": "Cache"},
|
||||||
"code": {"advanced": True},
|
"code": {"advanced": True},
|
||||||
"input_value": {"display_name": "Input"},
|
"input_value": {"display_name": "Input"},
|
||||||
|
"stream": {
|
||||||
|
"display_name": "Stream",
|
||||||
|
"info": "Stream the response from the model.",
|
||||||
|
},
|
||||||
}
|
}
|
||||||
|
|
||||||
def build(
|
def build(
|
||||||
|
|
@ -47,6 +51,7 @@ class AmazonBedrockComponent(CustomComponent):
|
||||||
endpoint_url: Optional[str] = None,
|
endpoint_url: Optional[str] = None,
|
||||||
streaming: bool = False,
|
streaming: bool = False,
|
||||||
cache: Optional[bool] = None,
|
cache: Optional[bool] = None,
|
||||||
|
stream: bool = False,
|
||||||
) -> Text:
|
) -> Text:
|
||||||
try:
|
try:
|
||||||
output = BedrockChat(
|
output = BedrockChat(
|
||||||
|
|
@ -60,6 +65,9 @@ class AmazonBedrockComponent(CustomComponent):
|
||||||
) # type: ignore
|
) # type: ignore
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
raise ValueError("Could not connect to AmazonBedrock API.") from e
|
raise ValueError("Could not connect to AmazonBedrock API.") from e
|
||||||
|
if stream:
|
||||||
|
result = output.stream(input_value)
|
||||||
|
else:
|
||||||
message = output.invoke(input_value)
|
message = output.invoke(input_value)
|
||||||
result = message.content if hasattr(message, "content") else message
|
result = message.content if hasattr(message, "content") else message
|
||||||
self.status = result
|
self.status = result
|
||||||
|
|
|
||||||
|
|
@ -50,6 +50,10 @@ class AnthropicLLM(CustomComponent):
|
||||||
},
|
},
|
||||||
"code": {"show": False},
|
"code": {"show": False},
|
||||||
"input_value": {"display_name": "Input"},
|
"input_value": {"display_name": "Input"},
|
||||||
|
"stream": {
|
||||||
|
"display_name": "Stream",
|
||||||
|
"info": "Stream the response from the model.",
|
||||||
|
},
|
||||||
}
|
}
|
||||||
|
|
||||||
def build(
|
def build(
|
||||||
|
|
@ -60,6 +64,7 @@ class AnthropicLLM(CustomComponent):
|
||||||
max_tokens: Optional[int] = None,
|
max_tokens: Optional[int] = None,
|
||||||
temperature: Optional[float] = None,
|
temperature: Optional[float] = None,
|
||||||
api_endpoint: Optional[str] = None,
|
api_endpoint: Optional[str] = None,
|
||||||
|
stream: bool = False,
|
||||||
) -> Text:
|
) -> Text:
|
||||||
# Set default API endpoint if not provided
|
# Set default API endpoint if not provided
|
||||||
if not api_endpoint:
|
if not api_endpoint:
|
||||||
|
|
@ -77,6 +82,9 @@ class AnthropicLLM(CustomComponent):
|
||||||
)
|
)
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
raise ValueError("Could not connect to Anthropic API.") from e
|
raise ValueError("Could not connect to Anthropic API.") from e
|
||||||
|
if stream:
|
||||||
|
result = output.stream(input_value)
|
||||||
|
else:
|
||||||
message = output.invoke(input_value)
|
message = output.invoke(input_value)
|
||||||
result = message.content if hasattr(message, "content") else message
|
result = message.content if hasattr(message, "content") else message
|
||||||
self.status = result
|
self.status = result
|
||||||
|
|
|
||||||
|
|
@ -74,6 +74,10 @@ class AzureChatOpenAIComponent(CustomComponent):
|
||||||
},
|
},
|
||||||
"code": {"show": False},
|
"code": {"show": False},
|
||||||
"input_value": {"display_name": "Input"},
|
"input_value": {"display_name": "Input"},
|
||||||
|
"stream": {
|
||||||
|
"display_name": "Stream",
|
||||||
|
"info": "Stream the response from the model.",
|
||||||
|
},
|
||||||
}
|
}
|
||||||
|
|
||||||
def build(
|
def build(
|
||||||
|
|
@ -86,6 +90,7 @@ class AzureChatOpenAIComponent(CustomComponent):
|
||||||
api_version: str,
|
api_version: str,
|
||||||
temperature: float = 0.7,
|
temperature: float = 0.7,
|
||||||
max_tokens: Optional[int] = 1000,
|
max_tokens: Optional[int] = 1000,
|
||||||
|
stream: bool = False,
|
||||||
) -> BaseLanguageModel:
|
) -> BaseLanguageModel:
|
||||||
try:
|
try:
|
||||||
output = AzureChatOpenAI(
|
output = AzureChatOpenAI(
|
||||||
|
|
@ -99,6 +104,9 @@ class AzureChatOpenAIComponent(CustomComponent):
|
||||||
)
|
)
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
raise ValueError("Could not connect to AzureOpenAI API.") from e
|
raise ValueError("Could not connect to AzureOpenAI API.") from e
|
||||||
|
if stream:
|
||||||
|
result = output.stream(input_value)
|
||||||
|
else:
|
||||||
message = output.invoke(input_value)
|
message = output.invoke(input_value)
|
||||||
result = message.content if hasattr(message, "content") else message
|
result = message.content if hasattr(message, "content") else message
|
||||||
self.status = result
|
self.status = result
|
||||||
|
|
|
||||||
|
|
@ -69,6 +69,10 @@ class QianfanChatEndpointComponent(CustomComponent):
|
||||||
},
|
},
|
||||||
"code": {"show": False},
|
"code": {"show": False},
|
||||||
"input_value": {"display_name": "Input"},
|
"input_value": {"display_name": "Input"},
|
||||||
|
"stream": {
|
||||||
|
"display_name": "Stream",
|
||||||
|
"info": "Stream the response from the model.",
|
||||||
|
},
|
||||||
}
|
}
|
||||||
|
|
||||||
def build(
|
def build(
|
||||||
|
|
@ -81,6 +85,7 @@ class QianfanChatEndpointComponent(CustomComponent):
|
||||||
temperature: Optional[float] = None,
|
temperature: Optional[float] = None,
|
||||||
penalty_score: Optional[float] = None,
|
penalty_score: Optional[float] = None,
|
||||||
endpoint: Optional[str] = None,
|
endpoint: Optional[str] = None,
|
||||||
|
stream: bool = False,
|
||||||
) -> Text:
|
) -> Text:
|
||||||
try:
|
try:
|
||||||
output = QianfanChatEndpoint( # type: ignore
|
output = QianfanChatEndpoint( # type: ignore
|
||||||
|
|
@ -94,6 +99,9 @@ class QianfanChatEndpointComponent(CustomComponent):
|
||||||
)
|
)
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
raise ValueError("Could not connect to Baidu Qianfan API.") from e
|
raise ValueError("Could not connect to Baidu Qianfan API.") from e
|
||||||
|
if stream:
|
||||||
|
result = output.stream(input_value)
|
||||||
|
else:
|
||||||
message = output.invoke(input_value)
|
message = output.invoke(input_value)
|
||||||
result = message.content if hasattr(message, "content") else message
|
result = message.content if hasattr(message, "content") else message
|
||||||
self.status = result
|
self.status = result
|
||||||
|
|
|
||||||
|
|
@ -29,6 +29,10 @@ class CTransformersComponent(CustomComponent):
|
||||||
"value": '{"top_k":40,"top_p":0.95,"temperature":0.8,"repetition_penalty":1.1,"last_n_tokens":64,"seed":-1,"max_new_tokens":256,"stop":"","stream":"False","reset":"True","batch_size":8,"threads":-1,"context_length":-1,"gpu_layers":0}',
|
"value": '{"top_k":40,"top_p":0.95,"temperature":0.8,"repetition_penalty":1.1,"last_n_tokens":64,"seed":-1,"max_new_tokens":256,"stop":"","stream":"False","reset":"True","batch_size":8,"threads":-1,"context_length":-1,"gpu_layers":0}',
|
||||||
},
|
},
|
||||||
"input_value": {"display_name": "Input"},
|
"input_value": {"display_name": "Input"},
|
||||||
|
"stream": {
|
||||||
|
"display_name": "Stream",
|
||||||
|
"info": "Stream the response from the model.",
|
||||||
|
},
|
||||||
}
|
}
|
||||||
|
|
||||||
def build(
|
def build(
|
||||||
|
|
@ -38,10 +42,14 @@ class CTransformersComponent(CustomComponent):
|
||||||
input_value: str,
|
input_value: str,
|
||||||
model_type: str,
|
model_type: str,
|
||||||
config: Optional[Dict] = None,
|
config: Optional[Dict] = None,
|
||||||
|
stream: Optional[bool] = False,
|
||||||
) -> Text:
|
) -> Text:
|
||||||
output = CTransformers(
|
output = CTransformers(
|
||||||
model=model, model_file=model_file, model_type=model_type, config=config
|
model=model, model_file=model_file, model_type=model_type, config=config
|
||||||
)
|
)
|
||||||
|
if stream:
|
||||||
|
result = output.stream(input_value)
|
||||||
|
else:
|
||||||
message = output.invoke(input_value)
|
message = output.invoke(input_value)
|
||||||
result = message.content if hasattr(message, "content") else message
|
result = message.content if hasattr(message, "content") else message
|
||||||
self.status = result
|
self.status = result
|
||||||
|
|
|
||||||
|
|
@ -43,8 +43,10 @@ class CohereComponent(CustomComponent):
|
||||||
max_tokens=max_tokens,
|
max_tokens=max_tokens,
|
||||||
temperature=temperature,
|
temperature=temperature,
|
||||||
)
|
)
|
||||||
|
if stream:
|
||||||
|
result = output.stream(input_value)
|
||||||
|
else:
|
||||||
message = output.invoke(input_value)
|
message = output.invoke(input_value)
|
||||||
result = message.content if hasattr(message, "content") else message
|
result = message.content if hasattr(message, "content") else message
|
||||||
self.status = result
|
self.status = result
|
||||||
return result
|
return result return result
|
||||||
return result
|
|
||||||
|
|
|
||||||
|
|
@ -70,6 +70,9 @@ class GoogleGenerativeAIComponent(CustomComponent):
|
||||||
n=n or 1,
|
n=n or 1,
|
||||||
google_api_key=SecretStr(google_api_key),
|
google_api_key=SecretStr(google_api_key),
|
||||||
)
|
)
|
||||||
|
if stream:
|
||||||
|
result = output.stream(input_value)
|
||||||
|
else:
|
||||||
message = output.invoke(input_value)
|
message = output.invoke(input_value)
|
||||||
result = message.content if hasattr(message, "content") else message
|
result = message.content if hasattr(message, "content") else message
|
||||||
self.status = result
|
self.status = result
|
||||||
|
|
|
||||||
|
|
@ -57,6 +57,10 @@ class LlamaCppComponent(CustomComponent):
|
||||||
"verbose": {"display_name": "Verbose", "advanced": True},
|
"verbose": {"display_name": "Verbose", "advanced": True},
|
||||||
"vocab_only": {"display_name": "Vocab Only", "advanced": True},
|
"vocab_only": {"display_name": "Vocab Only", "advanced": True},
|
||||||
"input_value": {"display_name": "Input"},
|
"input_value": {"display_name": "Input"},
|
||||||
|
"stream": {
|
||||||
|
"display_name": "Stream",
|
||||||
|
"info": "Stream the response from the model.",
|
||||||
|
},
|
||||||
}
|
}
|
||||||
|
|
||||||
def build(
|
def build(
|
||||||
|
|
@ -97,6 +101,7 @@ class LlamaCppComponent(CustomComponent):
|
||||||
use_mmap: Optional[bool] = True,
|
use_mmap: Optional[bool] = True,
|
||||||
verbose: bool = True,
|
verbose: bool = True,
|
||||||
vocab_only: bool = False,
|
vocab_only: bool = False,
|
||||||
|
stream: bool = False,
|
||||||
) -> Text:
|
) -> Text:
|
||||||
output = LlamaCpp(
|
output = LlamaCpp(
|
||||||
model_path=model_path,
|
model_path=model_path,
|
||||||
|
|
@ -135,9 +140,10 @@ class LlamaCppComponent(CustomComponent):
|
||||||
verbose=verbose,
|
verbose=verbose,
|
||||||
vocab_only=vocab_only,
|
vocab_only=vocab_only,
|
||||||
)
|
)
|
||||||
|
if stream:
|
||||||
|
result = output.stream(input_value)
|
||||||
|
else:
|
||||||
message = output.invoke(input_value)
|
message = output.invoke(input_value)
|
||||||
result = message.content if hasattr(message, "content") else message
|
result = message.content if hasattr(message, "content") else message
|
||||||
self.status = result
|
self.status = result
|
||||||
return result
|
return result
|
||||||
self.status = result
|
|
||||||
return result
|
|
||||||
|
|
|
||||||
|
|
@ -165,6 +165,10 @@ class ChatOllamaComponent(CustomComponent):
|
||||||
"advanced": True,
|
"advanced": True,
|
||||||
},
|
},
|
||||||
"input_value": {"display_name": "Input"},
|
"input_value": {"display_name": "Input"},
|
||||||
|
"stream": {
|
||||||
|
"display_name": "Stream",
|
||||||
|
"info": "Stream the response from the model.",
|
||||||
|
},
|
||||||
}
|
}
|
||||||
|
|
||||||
def build(
|
def build(
|
||||||
|
|
@ -197,6 +201,7 @@ class ChatOllamaComponent(CustomComponent):
|
||||||
timeout: Optional[int] = None,
|
timeout: Optional[int] = None,
|
||||||
top_k: Optional[int] = None,
|
top_k: Optional[int] = None,
|
||||||
top_p: Optional[int] = None,
|
top_p: Optional[int] = None,
|
||||||
|
stream: Optional[bool] = False,
|
||||||
) -> Text:
|
) -> Text:
|
||||||
if not base_url:
|
if not base_url:
|
||||||
base_url = "http://localhost:11434"
|
base_url = "http://localhost:11434"
|
||||||
|
|
@ -250,6 +255,9 @@ class ChatOllamaComponent(CustomComponent):
|
||||||
output = ChatOllama(**llm_params) # type: ignore
|
output = ChatOllama(**llm_params) # type: ignore
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
raise ValueError("Could not initialize Ollama LLM.") from e
|
raise ValueError("Could not initialize Ollama LLM.") from e
|
||||||
|
if stream:
|
||||||
|
result = output.stream(input_value)
|
||||||
|
else:
|
||||||
message = output.invoke(input_value)
|
message = output.invoke(input_value)
|
||||||
result = message.content if hasattr(message, "content") else message
|
result = message.content if hasattr(message, "content") else message
|
||||||
self.status = result
|
self.status = result
|
||||||
|
|
|
||||||
|
|
@ -57,6 +57,10 @@ class OpenAIModelComponent(CustomComponent):
|
||||||
"required": False,
|
"required": False,
|
||||||
"value": 0.7,
|
"value": 0.7,
|
||||||
},
|
},
|
||||||
|
"stream": {
|
||||||
|
"display_name": "Stream",
|
||||||
|
"info": "Stream the response from the model.",
|
||||||
|
},
|
||||||
}
|
}
|
||||||
|
|
||||||
def build(
|
def build(
|
||||||
|
|
@ -68,10 +72,11 @@ class OpenAIModelComponent(CustomComponent):
|
||||||
openai_api_base: Optional[str] = None,
|
openai_api_base: Optional[str] = None,
|
||||||
openai_api_key: Optional[str] = None,
|
openai_api_key: Optional[str] = None,
|
||||||
temperature: float = 0.7,
|
temperature: float = 0.7,
|
||||||
|
stream: Optional[bool] = False,
|
||||||
) -> Text:
|
) -> Text:
|
||||||
if not openai_api_base:
|
if not openai_api_base:
|
||||||
openai_api_base = "https://api.openai.com/v1"
|
openai_api_base = "https://api.openai.com/v1"
|
||||||
model = ChatOpenAI(
|
output = ChatOpenAI(
|
||||||
max_tokens=max_tokens,
|
max_tokens=max_tokens,
|
||||||
model_kwargs=model_kwargs,
|
model_kwargs=model_kwargs,
|
||||||
model=model_name,
|
model=model_name,
|
||||||
|
|
@ -79,8 +84,10 @@ class OpenAIModelComponent(CustomComponent):
|
||||||
api_key=openai_api_key,
|
api_key=openai_api_key,
|
||||||
temperature=temperature,
|
temperature=temperature,
|
||||||
)
|
)
|
||||||
|
if stream:
|
||||||
message = model.invoke(input_value)
|
result = output.stream(input_value)
|
||||||
|
else:
|
||||||
|
message = output.invoke(input_value)
|
||||||
result = message.content if hasattr(message, "content") else message
|
result = message.content if hasattr(message, "content") else message
|
||||||
self.status = result
|
self.status = result
|
||||||
return result
|
return result
|
||||||
|
|
|
||||||
|
|
@ -58,6 +58,10 @@ class ChatVertexAIComponent(CustomComponent):
|
||||||
"advanced": True,
|
"advanced": True,
|
||||||
},
|
},
|
||||||
"input_value": {"display_name": "Input"},
|
"input_value": {"display_name": "Input"},
|
||||||
|
"stream": {
|
||||||
|
"display_name": "Stream",
|
||||||
|
"info": "Stream the response from the model.",
|
||||||
|
},
|
||||||
}
|
}
|
||||||
|
|
||||||
def build(
|
def build(
|
||||||
|
|
@ -73,6 +77,7 @@ class ChatVertexAIComponent(CustomComponent):
|
||||||
top_k: int = 40,
|
top_k: int = 40,
|
||||||
top_p: float = 0.95,
|
top_p: float = 0.95,
|
||||||
verbose: bool = False,
|
verbose: bool = False,
|
||||||
|
stream: bool = False,
|
||||||
) -> Text:
|
) -> Text:
|
||||||
try:
|
try:
|
||||||
from langchain_google_vertexai import ChatVertexAI
|
from langchain_google_vertexai import ChatVertexAI
|
||||||
|
|
@ -92,6 +97,9 @@ class ChatVertexAIComponent(CustomComponent):
|
||||||
top_p=top_p,
|
top_p=top_p,
|
||||||
verbose=verbose,
|
verbose=verbose,
|
||||||
)
|
)
|
||||||
|
if stream:
|
||||||
|
result = output.stream(input_value)
|
||||||
|
else:
|
||||||
message = output.invoke(input_value)
|
message = output.invoke(input_value)
|
||||||
result = message.content if hasattr(message, "content") else message
|
result = message.content if hasattr(message, "content") else message
|
||||||
self.status = result
|
self.status = result
|
||||||
|
|
|
||||||
Loading…
Add table
Add a link
Reference in a new issue