def test_response_post_update_fails_when_no_task_present():
client = get_steamship_client()
response = Response()
with pytest.raises(Exception):
response.post_update(client)
# No task_id
response2 = Response(status=Task())
with pytest.raises(Exception):
response2.post_update(client)
def learn(self, fact: str = None) -> Response:
"""Learns a new fact."""
if fact is None:
return Response.error(500, "Empty fact provided to learn.")
if self.index is None:
return Response.error(500, "Unable to initialize QA index.")
res = self.index.embed(fact)
if res.error:
# Steamship error messages can be passed straight
# back to the user
return Response(error=res.error)
return Response(json=res.data)
def query(self, query: str = None, k: int = 1) -> Response:
"""Learns a new fact."""
if query is None:
return Response.error(500, "Empty query provided.")
if self.index is None:
return Response.error(500, "Unable to initialize QA index.")
res = self.index.query(query=query, k=k)
if res.error:
# Steamship error messages can be passed straight
# back to the user
return Response(error=res.error)
return Response(json=res.data)
def run(
self, request: PluginRequest[RawDataPluginInput]
) -> Response[BlockAndTagPluginOutput]:
return Response(data=BlockAndTagPluginOutput(file=File.CreateRequest(
blocks=[
Block.CreateRequest(
text=TEST_H1,
tags=[Tag.CreateRequest(kind=TagKind.doc, name=DocTag.h1)],
),
Block.CreateRequest(
text=TEST_S1,
tags=[
Tag.CreateRequest(kind=TagKind.doc,
name=DocTag.sentence)
],
),
Block.CreateRequest(
text=TEST_S2,
tags=[
Tag.CreateRequest(kind=TagKind.doc,
name=DocTag.sentence)
],
),
Block.CreateRequest(
text=TEST_S3,
tags=[
Tag.CreateRequest(kind=TagKind.doc,
name=DocTag.paragraph)
],
),
])))
def train_status(self, request: PluginRequest[TrainStatusPluginInput],
model: ThirdPartyModel) -> Response[TrainPluginOutput]:
"""Since trainable can't be assumed to be asynchronous, the trainer is responsible for uploading its own model file."""
logging.debug(f"train {request}")
# Create a Response object at the top with a Task attached. This will let us stream back updates
# TODO: This is very non-intuitive. We should improve this.
response = Response(status=Task(task_id=request.task_id))
# Call train status
train_plugin_output = model.train_status(request.data)
if train_plugin_output.training_complete:
# Save the model with the `default` handle.
archive_path_in_steamship = model.save_remote(
client=self.client,
plugin_instance_id=request.plugin_instance_id)
# Set the model location on the plugin output.
train_plugin_output.archive_path = archive_path_in_steamship
# Set the response on the `data` field of the object
response.set_data(json=train_plugin_output)
logging.info(response.dict(by_alias=True))
return response
def run(
self, request: PluginRequest[BlockAndTagPluginInput]
) -> Response[BlockAndTagPluginOutput]:
updated_blocks = [
_embed_block(block) for block in request.data.file.blocks
]
return Response(data=BlockAndTagPluginOutput(file=File.CreateRequest(
blocks=updated_blocks)))
def get_config(self) -> Response:
"""This is called get_config because there's already `.config` object on the class."""
return Response(json=dict(
spaceId=self.client.config.space_id,
appBase=self.client.config.app_base,
apiBase=self.client.config.api_base,
apiKey=self.client.config.api_key,
))
def run(
self, request: PluginRequest[BlockAndTagPluginInput]
) -> Response[BlockAndTagPluginOutput]:
logging.info(f"Inside parser: {type(request)}")
file = request.data.file
for block in file.blocks:
tag_sentences(block)
if request.data is not None:
ret = Response(data=BlockAndTagPluginOutput(file=file))
logging.info(f"Ret: {ret}")
return ret
def run(
self, request: PluginRequest[BlockAndTagPluginInput]
) -> Response[BlockAndTagPluginOutput]:
"""Tags the incoming data for any instance of the keywords in the parameter file."""
assert self.config is not None
assert self.config.testValue1 is not None
assert self.config.testValue2 is not None
response = Response(data=BlockAndTagPluginOutput(
file=File.CreateRequest(tags=[])))
logging.info(f"TestTrainableTaggerModel:run() returning {response}")
return response
def run(
self, request: PluginRequest[BlockAndTagPluginInput]
) -> Response[BlockAndTagPluginOutput]:
"""Tags the incoming data for any instance of the keywords in the parameter file."""
logging.info(
f"TestTrainableTaggerModel:run() - My keyword list is {self.keyword_list}"
)
response = Response(data=BlockAndTagPluginOutput(
file=File.CreateRequest(tags=[
Tag.CreateRequest(name=word) for word in self.keyword_list
])))
logging.info(f"TestTrainableTaggerModel:run() returning {response}")
return response
def run(
self, request: PluginRequest[RawDataPluginInput]
) -> Union[Response, Response[BlockAndTagPluginOutput]]:
if request is None or request.data is None or request.data.data is None:
return Response(error=SteamshipError(
message="Missing data field on the incoming request."))
data = request.data.data # TODO (enias): Simplify
if isinstance(data, bytes):
data = data.decode("utf-8")
if not isinstance(data, str):
return Response(error=SteamshipError(
message="The incoming data was not of expected String type"))
reader = csv.DictReader(
io.StringIO(data),
delimiter=self.config.delimiter,
quotechar=self.config.quotechar,
escapechar=self.config.escapechar,
skipinitialspace=self.config.skipinitialspace,
)
file = File.CreateRequest(blocks=[])
for row in reader:
text = self._get_text(row)
tag_values = self._get_tags(row)
block = Block.CreateRequest(
text=text,
tags=[
Tag.CreateRequest(kind=self.config.tag_kind,
name=tag_value)
for tag_value in tag_values
],
)
file.blocks.append(block)
return Response(data=BlockAndTagPluginOutput(file=file))
def run(
self, request: PluginRequest[CorpusImportRequest]
) -> Response[CorpusImportResponse]:
return Response(data=CorpusImportResponse(file_import_requests=[
File.CreateRequest(
type="fileImporter",
corpusId=request.data.url,
plugin_instance=request.data.file_importer_plugin_instance,
),
File.CreateRequest(
type="fileImporter",
corpusId=request.data.url,
plugin_instance=request.data.file_importer_plugin_instance,
),
]))
def train(self, request: PluginRequest[TrainPluginInput],
model: ThirdPartyModel) -> Response[TrainPluginOutput]:
"""Since trainable can't be assumed to be asynchronous, the trainer is responsible for uploading its own model file."""
logging.debug(f"train {request}")
# Create a Response object at the top with a Task attached. This will let us stream back updates
# TODO: This is very non-intuitive. We should improve this.
response = Response(status=Task(task_id=request.task_id))
# Train the model
train_plugin_input = request.data
train_plugin_output = model.train(train_plugin_input)
# Set the response on the `data` field of the object
response.set_data(json=train_plugin_output)
logging.info(response.dict(by_alias=True))
return response
def train(self, request: PluginRequest[TrainPluginInput],
model: TestTrainableTaggerModel) -> Response[TrainPluginOutput]:
"""Since trainable can't be assumed to be asynchronous, the trainer is responsible for uploading its own model file."""
logging.info(f"TestTrainableTaggerPlugin:train() {request}")
# Create a Response object at the top with a Task attached. This will let us stream back updates
# TODO: This is very non-intuitive. We should improve this.
response = Response(status=Task(task_id=request.task_id))
# Example of recording training progress
# response.status.status_message = "About to train!"
# response.post_update(client=self.client)
# Train the model
train_plugin_input = request.data
train_plugin_output = model.train(train_plugin_input)
# Save the model with the `default` handle.
archive_path_in_steamship = model.save_remote(
client=self.client, plugin_instance_id=request.plugin_instance_id)
# Set the model location on the plugin output.
logging.info(
f"TestTrainableTaggerPlugin:train() setting model archive path to {archive_path_in_steamship}"
)
train_plugin_output.archive_path = archive_path_in_steamship
# Set the response on the `data` field of the object
response.set_data(json=train_plugin_output)
# If we want we can post this to the Engine
# response.status.status_message = "Done!"
# response.status.state = TaskState.succeeded
# response.post_update(client=self.client)
# Or, if this training really did happen synchronously, we return it.
# Some models (e.g. those running on ECS, or on a third party system) will not have completed by the time
# the Lambda function finishes. For now, let's just pretend they're synchronous. But in a future PR when we
# have a better method of handling such situations, the response below would include a `status` of type `running`
# to indicate that, while the plugin handler has returned, the plugin's execution continues.
logging.info(f"TestTrainableTaggerPlugin:train() returning {response}")
return response
def test_response_post_update_can_update_task():
client = get_steamship_client()
task_result = create_dummy_training_task(client)
task = task_result.task
new_state = TaskState.failed
new_message = "HI THERE"
new_output = {"a": 3}
assert task.state != new_state
assert task.status_message != new_message
assert task.output != new_output
response = Response(status=task)
response.status.state = new_state
response.status.status_message = new_message
response.status.output = new_output
# Sanity check: we'll prove that caling task.check() resets this..
task_result.refresh()
# Assert not equal
assert task.state != new_state
assert task.status_message != new_message
assert task.output != new_output
# And override again
response.status.state = new_state
response.status.status_message = new_message
response.set_data(json=new_output)
# Now we call post_update
response.post_update(client)
# Call task.check
task_result.refresh()
# Assert equal
assert task.state == new_state
assert task.status_message == new_message
assert task.output == json.dumps(new_output)
def run(
self, request: PluginRequest[BlockAndTagPluginInput]
) -> Response[BlockAndTagPluginOutput]:
tag_kind = self.config.tagKind
tag_name = self.config.tagName
tag_value = {
"numberValue": self.config.numberValue,
"booleanValue": self.config.booleanValue,
}
if request.data is not None:
file = request.data.file
tag = Tag.CreateRequest(kind=tag_kind,
name=tag_name,
value=tag_value)
if file.tags:
file.tags.append(tag)
else:
file.tags = [tag]
return Response(data=BlockAndTagPluginOutput(file=file))
def run(
self, request: PluginRequest[BlockAndTagPluginInput]
) -> Response[BlockAndTagPluginOutput]:
"""Runs the mock client"""
if "model_id" not in self.params:
raise SteamshipError(
message=
"No model_id was found in model parameter file. Has the model been trained?"
)
output = BlockAndTagPluginOutput(file=File.CreateRequest(blocks=[]))
for in_block in request.data.file.blocks:
tags = self.client.infer(in_block.text, self.params["model_id"])
out_block = Block.CreateRequest(
id=in_block.id,
tags=[Tag.CreateRequest(name=tag) for tag in tags])
output.file.blocks.append(out_block)
return Response(json=output)
def resp_obj(self) -> Response:
return Response(json=TestObj(name="Foo"))
def get_training_parameters(
self, request: PluginRequest[TrainingParameterPluginInput]
) -> Response[TrainingParameterPluginOutput]:
ret = Response(data={})
return ret
def resp_binary(self) -> Response:
_bytes = base64.b64decode(PALM_TREE_BASE_64)
return Response(_bytes=_bytes)
def get_training_parameters(
self, request: PluginRequest[TrainingParameterPluginInput]
) -> Response[TrainingParameterPluginOutput]:
logging.debug(f"get_training_parameters {request}")
return Response(data=TRAINING_PARAMETERS)
def space(self) -> Response:
return Response(string=self.client.config.space_id)
def resp_bytes_io(self) -> Response:
_bytes = base64.b64decode(PALM_TREE_BASE_64)
return Response(_bytes=io.BytesIO(_bytes))
def resp_image(self) -> Response:
_bytes = base64.b64decode(PALM_TREE_BASE_64)
return Response(_bytes=_bytes, mime_type=MimeTypes.PNG)
def resp_dict(self) -> Response:
return Response(json=dict(string="A String", int=10))
def greet2(self, name: str = "Person") -> Response:
return Response(string=f"Hello, {name}!")
def run(
self, request: PluginRequest[FileImportPluginInput]
) -> Response[RawDataPluginOutput]:
return Response(
data=RawDataPluginOutput(string=TEST_DOC, mime_type=MimeTypes.MKD))
def resp_string(self) -> Response:
return Response(string="A String")
def user_info(self) -> Response:
user = User.current(self.client)
return Response(json=dict(handle=user.data.handle))