from steamship.plugin.config import Config
from steamship.plugin.file_importer import FileImporter
from steamship.plugin.inputs.file_import_plugin_input import FileImportPluginInput
from steamship.plugin.outputs.raw_data_plugin_output import RawDataPluginOutput
from steamship.plugin.service import PluginRequest
# Note: this aligns with the same document in the internal Engine test.
HANDLE = "test-importer-plugin-v1"
TEST_H1 = "A Poem"
TEST_S1 = "Roses are red."
TEST_S2 = "Violets are blue."
TEST_S3 = "Sugar is sweet, and I love you."
TEST_DOC = f"# {TEST_H1}\n\n{TEST_S1} {TEST_S2}\n\n{TEST_S3}\n"
class TestFileImporterPlugin(FileImporter):
class EmptyConfig(Config):
pass
def config_cls(self) -> Type[Config]:
return self.EmptyConfig
def run(
self, request: PluginRequest[FileImportPluginInput]
) -> Response[RawDataPluginOutput]:
return Response(
data=RawDataPluginOutput(string=TEST_DOC, mime_type=MimeTypes.MKD))
handler = create_handler(TestFileImporterPlugin)
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))
handler = create_handler(CsvBlockifier)
) -> Response[BlockAndTagPluginOutput]:
"""Downloads the model file from the provided space"""
logging.debug(f"run_with_model {request} {model}")
logging.info(
f"TestTrainableTaggerPlugin:run_with_model() got request {request} and model {model}"
)
return model.run(request)
def get_training_parameters(
self, request: PluginRequest[TrainingParameterPluginInput]
) -> Response[TrainingParameterPluginOutput]:
ret = Response(data={})
return ret
def train(
self, request: PluginRequest[TrainPluginInput],
model: TestTrainableTaggerConfigModel
) -> Response[TrainPluginOutput]:
train_plugin_input = request.data
train_plugin_output = model.train(train_plugin_input)
return None
def train_status(self, request: PluginRequest[TrainStatusPluginInput],
model: TrainableModel) -> Response[TrainPluginOutput]:
model.train_status(request.data)
return None
handler = create_handler(TestTrainableTaggerConfigPlugin)
from steamship import File
from steamship.app import Response, create_handler
from steamship.data.operations.corpus_importer import CorpusImportRequest, CorpusImportResponse
from steamship.plugin.config import Config
from steamship.plugin.corpus_importer import CorpusImporter
from steamship.plugin.service import PluginRequest
class TestCorpusImporterPlugin(CorpusImporter):
def config_cls(self) -> Type[Config]:
return Config
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,
),
]))
handler = create_handler(TestCorpusImporterPlugin)
return response
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
handler = create_handler(ThirdPartyTrainableTaggerPlugin)
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)
],
),
])))
handler = create_handler(DummyBlockifierPlugin)
embedding = embed(s)
return Tag.CreateRequest(kind=TagKind.text,
name=TextTag.Embedding,
value={TextTag.Embedding: embedding})
def _embed_block(block: Block) -> Block.CreateRequest:
return Block.CreateRequest(id=block.id,
text=block.text,
tags=[_embed_to_tag(block.text)])
class TestEmbedderPlugin(Embedder):
class EmptyConfig(Config):
pass
def config_cls(self) -> Type[Config]:
return self.EmptyConfig
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)))
handler = create_handler(TestEmbedderPlugin)
# If this is non-optional, the typecheck fails despite the fact that the test passes
# in a value of false....
booleanValue: Optional[bool] = False
def config_cls(self) -> Type[Config]:
return self.TestParserConfig
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))
handler = create_handler(TestParserPlugin)
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)
@post("query")
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)
handler = create_handler(TestApp)
