Add caching to CustomComponent class

This commit is contained in:
Gabriel Luiz Freitas Almeida 2023-11-23 10:57:06 -03:00
commit f4210f567b

View file

@ -1,7 +1,9 @@
import operator
from typing import Any, Callable, ClassVar, List, Optional, Union from typing import Any, Callable, ClassVar, List, Optional, Union
from uuid import UUID from uuid import UUID
import yaml import yaml
from cachetools import TTLCache, cachedmethod
from fastapi import HTTPException from fastapi import HTTPException
from langflow.field_typing.constants import CUSTOM_COMPONENT_SUPPORTED_TYPES from langflow.field_typing.constants import CUSTOM_COMPONENT_SUPPORTED_TYPES
from langflow.interface.custom.component import Component from langflow.interface.custom.component import Component
@ -29,6 +31,7 @@ class CustomComponent(Component):
status: Optional[str] = None status: Optional[str] = None
def __init__(self, **data): def __init__(self, **data):
self.cache = TTLCache(maxsize=1024, ttl=60)
super().__init__(**data) super().__init__(**data)
@property @property
@ -78,25 +81,10 @@ class CustomComponent(Component):
return super().get_code_tree(code) return super().get_code_tree(code)
@property @property
def get_function_entrypoint_args(self) -> str: def get_function_entrypoint_args(self) -> list:
if not self.code: build_method = self.get_build_method()
return "" if not build_method:
tree = self.get_code_tree(self.code) return []
component_classes = [cls for cls in tree["classes"] if self.code_class_base_inheritance in cls["bases"]]
if not component_classes:
return ""
# Assume the first Component class is the one we're interested in
component_class = component_classes[0]
build_methods = [
method for method in component_class["methods"] if method["name"] == self.function_entrypoint_name
]
if not build_methods:
return ""
build_method = build_methods[0]
args = build_method["args"] args = build_method["args"]
for arg in args: for arg in args:
@ -115,8 +103,8 @@ class CustomComponent(Component):
arg["type"] = "Data" arg["type"] = "Data"
return args return args
@property @cachedmethod(operator.attrgetter("cache"))
def get_function_entrypoint_return_type(self) -> List[str]: def get_build_method(self):
if not self.code: if not self.code:
return [] return []
tree = self.get_code_tree(self.code) tree = self.get_code_tree(self.code)
@ -134,7 +122,13 @@ class CustomComponent(Component):
if not build_methods: if not build_methods:
return [] return []
build_method = build_methods[0] return build_methods[0]
@property
def get_function_entrypoint_return_type(self) -> List[str]:
build_method = self.get_build_method()
if not build_method:
return build_method
return_type = build_method["return_type"] return_type = build_method["return_type"]
if not return_type: if not return_type:
return [] return []