From 967e10d08ee5a06433455d0d86331f210bb57be2 Mon Sep 17 00:00:00 2001 From: Jared Wilber Date: Thu, 28 Sep 2023 15:27:04 -0700 Subject: [PATCH] resolve db type for dashboard --- pykoi/component/base.py | 15 ++++++++------- 1 file changed, 8 insertions(+), 7 deletions(-) diff --git a/pykoi/component/base.py b/pykoi/component/base.py index 44960f3..d413eb2 100644 --- a/pykoi/component/base.py +++ b/pykoi/component/base.py @@ -4,6 +4,7 @@ from pykoi.component.chatbot_database_factory import ChatbotDatabaseFactory from pykoi.component.constants import FeedbackType +from pykoi.chat.db.comparator_database import ComparatorDatabase from pykoi.chat.db.qa_database import QuestionAnswerDatabase from pykoi.chat.db.rag_database import RAGDatabase from pykoi.chat.db.ranking_database import RankingDatabase @@ -42,9 +43,7 @@ class Component: props (Dict[str, Any]): Additional properties for the component. """ - def __init__( - self, fetch_func: Optional[Callable], svelte_component: str, **kwargs - ): + def __init__(self, fetch_func: Optional[Callable], svelte_component: str, **kwargs): """ Initialize a new instance of Component. @@ -54,9 +53,7 @@ def __init__( kwargs: Additional properties for the component. """ self.id = str(uuid.uuid4()) # Generate a unique ID - self.data_source = ( - DataSource(self.id, fetch_func) if fetch_func else None - ) + self.data_source = DataSource(self.id, fetch_func) if fetch_func else None self.svelte_component = svelte_component self.props = kwargs @@ -81,6 +78,7 @@ def __init__(self, fetch_func: Callable, value_column: List[str], **kwargs): super().__init__(fetch_func, "Dropdown", **kwargs) self.value_column = value_column + class RAG(Component): """ RAG class represents a RAG component. @@ -138,5 +136,8 @@ def __init__(self, database: any, **kwargs): database (QuestionAnswerDatabase): The database to use for the dashboard. kwargs: Additional properties for the dashboard. """ - super().__init__(None, "Feedback", **kwargs) + if isinstance(database, ComparatorDatabase): + super().__init__(None, "CompareDashboard", **kwargs) + else: + super().__init__(None, "Feedback", **kwargs) self.database = database