67 lines
1.7 KiB
Python
67 lines
1.7 KiB
Python
|
"""This examples shows how to create a custom component
|
||
|
|
that can be added to a Pipeline with:
|
||
|
|
|
||
|
|
c = MyComponent(min_value=0, max_value=10)
|
||
|
|
pipe = Pipeline()
|
||
|
|
pipe.add_component(c, name="my_component")
|
||
|
|
"""
|
||
|
|
|
||
|
|
import random
|
||
|
|
|
||
|
|
from neo4j_graphrag.experimental.pipeline import Component, DataModel
|
||
|
|
from pydantic import BaseModel, validate_call
|
||
|
|
|
||
|
|
|
||
|
|
class ComponentInputModel(BaseModel):
|
||
|
|
"""A class to model the component inputs.
|
||
|
|
This is not required, inputs can also be passed individually.
|
||
|
|
|
||
|
|
Note: can also inherit from DataModel.
|
||
|
|
"""
|
||
|
|
|
||
|
|
text: str
|
||
|
|
|
||
|
|
|
||
|
|
class ComponentResultModel(DataModel):
|
||
|
|
"""A class to model the component outputs.
|
||
|
|
Each component must have such a description of the output,
|
||
|
|
so that the parameter mapping can be validated before the
|
||
|
|
pipeline run starts.
|
||
|
|
"""
|
||
|
|
|
||
|
|
value: int
|
||
|
|
text: str
|
||
|
|
|
||
|
|
|
||
|
|
class MyComponent(Component):
|
||
|
|
"""Multiplies an input text by a random number
|
||
|
|
between `min_value` and `max_value`
|
||
|
|
"""
|
||
|
|
|
||
|
|
def __init__(self, min_value: int, max_value: int) -> None:
|
||
|
|
self.min_value = min_value
|
||
|
|
self.max_value = max_value
|
||
|
|
|
||
|
|
# this decorator is required when a Pydantic model is used in the inputs
|
||
|
|
@validate_call
|
||
|
|
async def run(self, inputs: ComponentInputModel) -> ComponentResultModel:
|
||
|
|
# logic here
|
||
|
|
random_value = random.randint(self.min_value, self.max_value)
|
||
|
|
return ComponentResultModel(
|
||
|
|
value=random_value,
|
||
|
|
text=inputs.text * random_value,
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
import asyncio
|
||
|
|
|
||
|
|
c = MyComponent(min_value=0, max_value=10)
|
||
|
|
print(
|
||
|
|
asyncio.run(
|
||
|
|
c.run(
|
||
|
|
inputs={"text": "Hello"} # type: ignore
|
||
|
|
)
|
||
|
|
)
|
||
|
|
)
|