Refactor blocks, processors and datasources. Add tests (#8)

* Update documentation

* Add test fixtures

* Reafactor processor interface and add tests

Author: vegito22

apps/tasks.py

@@ -4,7 +4,7 @@
 from typing import List
 
 from datasources.handlers.datasource_type_interface import DataSourceEntryItem
-from datasources.handlers.datasource_type_interface import DataSourceTypeInterface
+from datasources.handlers.datasource_type_interface import DataSourceProcessor
 from datasources.models import DataSource
 from datasources.models import DataSourceEntry
 from datasources.models import DataSourceEntryStatus
@@ -19,7 +19,7 @@ def add_data_entry_task(datasource: DataSource, datasource_entry_items: List[Dat
     datasource_entry_handler_cls = DataSourceTypeFactory.get_datasource_type_handler(
         datasource.type,
     )
-    datasource_entry_handler: DataSourceTypeInterface = datasource_entry_handler_cls(
+    datasource_entry_handler: DataSourceProcessor = datasource_entry_handler_cls(
         datasource,
     )
 

common/promptly/__init__.py renamed to common/blocks/__init__.py

@@ -1,30 +1,19 @@
 import logging
-import pickle
 from abc import ABC
 from abc import abstractmethod
 from typing import Generator
 from typing import Generic
 from typing import Optional
 from typing import TypeVar
 
-from pydantic import BaseModel
 from pydantic import Field
 
-LOGGER = logging.getLogger(__name__)
-
-
-class Schema(BaseModel):
-    pass
+from common.blocks.base.schema import BaseSchema as Schema
 
+LOGGER = logging.getLogger(__name__)
 
 class BaseInputEnvironment(Schema):
-    bypass_cache: bool = Field(True, description='Bypass cache')
-
-
-class BaseError(Schema):
-    code: int
-    message: str
-
+    pass
 
 class BaseInput(Schema):
     env: Optional[BaseInputEnvironment] = Field(
@@ -36,10 +25,6 @@ class BaseConfiguration(Schema):
     pass
 
 
-class BaseErrorOutput(Schema):
-    error: Optional[BaseError] = Field(None, description='Error Object')
-
-
 class BaseOutput(Schema):
     metadata: Optional[Schema] = Field(
         default={}, description='Metadata',
@@ -61,27 +46,17 @@ def delete(self, key):
 
 
 BaseInputType = TypeVar('BaseInputType', BaseInput, dict)
-BaseOutputType = TypeVar(
-    'BaseOutputType', BaseOutput,
-    BaseErrorOutput, Generator[BaseOutput, None, None],
-)
-BaseConfigurationType = TypeVar(
-    'BaseConfigurationType', BaseConfiguration, dict,
-)
-
-
-class BaseProcessor(Generic[BaseInputType, BaseOutputType, BaseConfigurationType], ABC):
-    """
-    Base class for all processors
-    """
+BaseOutputType = TypeVar('BaseOutputType', BaseOutput, dict)
+BaseConfigurationType = TypeVar('BaseConfigurationType', BaseConfiguration, dict)
 
+class ProcessorInterface(Generic[BaseInputType, BaseOutputType, BaseConfigurationType], ABC):
     @staticmethod
     def name() -> str:
         raise NotImplementedError
 
     @staticmethod
     def slug() -> str:
-        return BaseProcessor.name().lower().replace(' ', '_')
+        raise NotImplementedError
 
     @classmethod
     def _get_input_class(cls) -> BaseInputType:
@@ -90,26 +65,88 @@ def _get_input_class(cls) -> BaseInputType:
     @classmethod
     def _get_output_class(cls) -> BaseOutputType:
         return cls.__orig_bases__[0].__args__[1]
-
+    
     @classmethod
     def _get_configuration_class(cls) -> BaseConfigurationType:
         return cls.__orig_bases__[0].__args__[2]
 
     @classmethod
-    def get_input_schema(cls) -> dict:
+    def _get_input_schema(cls) -> dict:
         api_processor_interface_class = cls.__orig_bases__[0]
         return api_processor_interface_class.__args__[0].schema_json()
 
     @classmethod
-    def get_output_schema(cls) -> dict:
+    def _get_output_schema(cls) -> dict:
         api_processor_interface_class = cls.__orig_bases__[0]
         return api_processor_interface_class.__args__[1].schema_json()
 
     @classmethod
-    def get_configuration_schema(cls) -> dict:
+    def _get_configuration_schema(cls) -> dict:
         api_processor_interface_class = cls.__orig_bases__[0]
         return api_processor_interface_class.__args__[2].schema_json()
+    
+    @classmethod
+    def _get_input_ui_schema(cls) -> dict:
+        api_processor_interface_class = cls.__orig_bases__[0]
+        return api_processor_interface_class.__args__[0].get_ui_schema()
+    
+    @classmethod 
+    def _get_output_ui_schema(cls) -> dict:
+        api_processor_interface_class = cls.__orig_bases__[0]
+        return api_processor_interface_class.__args__[1].get_ui_schema()
+    
+    @classmethod
+    def _get_configuration_ui_schema(cls) -> dict:
+        api_processor_interface_class = cls.__orig_bases__[0]
+        return api_processor_interface_class.__args__[2].get_ui_schema()
+    
+    @classmethod
+    def get_input_cls(cls) -> BaseInputType:
+        return cls._get_input_class()
+
+    @classmethod
+    def get_output_cls(cls) -> BaseOutputType:
+        return cls._get_output_class()
+        
+    @classmethod
+    def get_configuration_cls(cls) -> BaseConfigurationType:
+        return cls._get_configuration_class()
+   
+    @classmethod
+    def get_input_schema(cls) -> dict:
+        return cls._get_input_schema()
+
+    @classmethod
+    def get_output_schema(cls) -> dict:
+        return cls._get_output_schema()
+    
+    @classmethod
+    def get_configuration_schema(cls) -> dict:
+        return cls._get_configuration_schema()
+    
+    @classmethod
+    def get_input_ui_schema(cls) -> dict:
+        return cls._get_input_ui_schema()
+    
+    @classmethod 
+    def get_output_ui_schema(cls) -> dict:
+        return cls._get_output_ui_schema()
+        
+    @classmethod
+    def get_configuration_ui_schema(cls) -> dict:
+        return cls._get_configuration_ui_schema()
 
+    def process(self, input: BaseInputType, configuration: BaseConfigurationType) -> BaseOutputType:
+        raise NotImplementedError()
+    
+    def process_iter(self, input: BaseInputType, configuration: BaseConfigurationType) -> Generator[BaseOutputType, None, None]:
+        raise NotImplementedError()
+    
+class BaseProcessor(ProcessorInterface[BaseInputType, BaseOutputType, BaseConfigurationType]):
+    """
+    Base class for all processors
+    """
+    
     def __init__(self, configuration: dict, cache_manager: CacheManager = None, input_tx_cb: callable = None, output_tx_cb: callable = None):
         self.configuration = self.parse_validate_configuration(configuration)
         self.cache_manager = cache_manager
@@ -154,9 +191,6 @@ def process_iter(self, input: dict) -> BaseOutputType:
             LOGGER.exception('Exception occurred while processing')
             raise ex
 
-    def serialize(self):
-        return pickle.dumps(self)
-
     @property
     def id(self):
         return id(self)

common/promptly/blocks/__init__.py renamed to common/blocks/base/__init__.py

@@ -0,0 +1,92 @@
+from pydantic import BaseModel
+
+def custom_json_dumps(v, **kwargs):
+    import ujson as json
+
+    default_arg = kwargs.get('default', None)
+    return json.dumps(v, default=default_arg)
+
+
+def custom_json_loads(v, **kwargs):
+    import ujson as json
+
+    return json.loads(v, **kwargs)
+
+
+class BaseSchema(BaseModel):
+    class Config:
+        json_dumps = custom_json_dumps
+        json_loads = custom_json_loads
+        
+    @ classmethod
+    def get_json_schema(cls):
+        return super().schema_json(indent=2)
+    
+    @ classmethod
+    def get_schema(cls):
+        return super().schema()
+    
+    @ classmethod
+    def get_ui_schema(cls):
+        """
+        This function receives a class method; gets the schema of the class
+        Calls the function form_ui_per_prop to form the dictionary of UI schema values
+        The resultant UI Schema only contains the properties
+        """
+        schema = cls.get_schema()
+        ui_schema = {
+            'ui:order': list(schema['properties'].keys()),
+            'properties': {}
+        }
+
+        for p_key, p_val in schema['properties'].items():
+            ui_schema['properties'][p_key] = cls.form_ui_per_prop(p_key, p_val)
+
+        return ui_schema['properties']
+
+    @ classmethod
+    def form_ui_per_prop(cls, p_key, prop_schema_dict):
+        """
+        This functions receives the property key and its schema dictionary
+        It checks the type of the property value and based on its type,
+        assigns the correct UI widget and other UI properties to it
+        """
+        ui_prop = {}
+
+        if 'title' in prop_schema_dict:
+            ui_prop['ui:label'] = prop_schema_dict['title']
+
+        if 'description' in prop_schema_dict:
+            ui_prop['ui:description'] = prop_schema_dict['description']
+
+        type_val = prop_schema_dict.get('type')
+
+        if type_val == 'string':
+            ui_prop['ui:widget'] = 'text'
+        elif type_val in ('integer', 'number'):
+            if 'minimum' in prop_schema_dict and 'maximum' in prop_schema_dict:
+                ui_prop['ui:widget'] = 'range'
+                ui_prop['ui:options'] = {
+                    'min': prop_schema_dict['minimum'],
+                    'max': prop_schema_dict['maximum'],
+                }
+            else:
+                ui_prop['ui:widget'] = 'updown'
+        elif type_val == 'boolean':
+            ui_prop['ui:widget'] = 'checkbox'
+        elif 'enum' in prop_schema_dict:
+            ui_prop['ui:widget'] = 'select'
+            ui_prop['ui:options'] = {
+                'enumOptions': [{'value': val, 'label': val} for val in prop_schema_dict['enum']]
+            }
+
+        if 'widget' in prop_schema_dict:
+            ui_prop['ui:widget'] = prop_schema_dict['widget']
+                    
+        if prop_schema_dict.get('format') == 'date-time':
+            ui_prop['ui:widget'] = 'datetime'
+        
+        # Unless explicitly mentioned, all properties are advanced parameters
+        ui_prop['ui:advanced'] = prop_schema_dict.get('advanced_parameter', True)
+
+        return ui_prop

common/promptly/core/base.py renamed to common/blocks/base/processor.py

@@ -1,6 +1,6 @@
 import hashlib
 
-from common.promptly.core.base import CacheManager
+from common.blocks.base.processor import CacheManager
 
 
 def get_hash(input: dict):

common/blocks/base/schema.py

@@ -0,0 +1,7 @@
+from typing import Any, Dict, Optional
+from common.blocks.base.schema import BaseSchema
+
+class DataDocument(BaseSchema):
+    content:Optional[bytes]
+    content_text: Optional[str]
+    metadata: Dict[str, Any] = {}

common/promptly/blocks/datastore/__init__.py renamed to common/blocks/cache/__init__.py

@@ -0,0 +1,25 @@
+from pydantic import AnyUrl
+from typing import List, Optional
+from common.blocks.base.schema import BaseSchema
+from common.blocks.data import DataDocument
+
+class DataUrl(AnyUrl):
+    @classmethod
+    def __modify_schema__(cls, field_schema):
+        field_schema.update(
+            {
+                'format': 'data-url',
+                'pattern': r'data:(.*);name=(.*);base64,(.*)',
+            },
+        )
+
+class DataSourceEnvironmentSchema(BaseSchema):
+    openai_key: str
+class DataSourceInputSchema(BaseSchema):
+    env: Optional[DataSourceEnvironmentSchema]
+
+class DataSourceConfigurationSchema(BaseSchema):
+    pass
+
+class DataSourceOutputSchema(BaseSchema):
+    documents: List[DataDocument]

common/promptly/blocks/utils.py renamed to common/blocks/cache/in_memory.py

@@ -0,0 +1,44 @@
+import os
+import re
+from pydantic import root_validator
+
+from common.blocks.base.processor import ProcessorInterface
+from common.blocks.data import DataDocument
+from common.blocks.data.source import DataSourceInputSchema, DataSourceConfigurationSchema, DataSourceOutputSchema
+
+
+class DirectoryTextLoaderInputSchema(DataSourceInputSchema):
+    directory: str
+    recursive: bool = False
+    
+    @root_validator()
+    @classmethod
+    def validate_directory(cls, field_values) -> str:
+        value = field_values.get("directory")
+        recursive = field_values.get("recursive")
+        
+        # TODO: Validate that directory is a valid directory path and the directory exists
+        if not re.match(r"^[a-zA-Z0-9_\-\.\/]+$", value):
+            raise ValueError("Directory must be a valid string")
+                
+        return value
+
+class DirectoryTextLoader(ProcessorInterface[DirectoryTextLoaderInputSchema, DataSourceOutputSchema, DataSourceConfigurationSchema]):
+    
+    def process(self, input: DirectoryTextLoaderInputSchema, configuration: DataSourceConfigurationSchema) -> DataSourceOutputSchema:
+        result = []
+        files = []
+        # If recursive is true, then we need to recursively walk the directory
+        if input.recursive:
+            for dir, dirname, filename in os.walk(input.directory):
+                files.extend(filename)
+        else:
+            files = os.listdir(input.directory)
+    
+        for file in files:
+            with open(os.path.join(input.directory, file), "r") as f:
+                result.append(DataDocument(name=file, content=f.read()))
+                
+        return DataSourceOutputSchema(documents=result)
+            
+