From f2cd0d759a4691e823113485f999214e0839bd81 Mon Sep 17 00:00:00 2001 From: Prithvi Kannan Date: Wed, 18 Dec 2024 14:28:33 -0800 Subject: [PATCH] use dataclasses Signed-off-by: Prithvi Kannan --- src/databricks_ai_bridge/genie.py | 14 +++++--------- 1 file changed, 5 insertions(+), 9 deletions(-) diff --git a/src/databricks_ai_bridge/genie.py b/src/databricks_ai_bridge/genie.py index b1542ed..f659eab 100644 --- a/src/databricks_ai_bridge/genie.py +++ b/src/databricks_ai_bridge/genie.py @@ -1,5 +1,6 @@ import logging import time +from dataclasses import dataclass from datetime import datetime from typing import Optional, Union @@ -18,16 +19,11 @@ def _count_tokens(text): return len(encoding.encode(text)) +@dataclass class GenieResponse: - def __init__( - self, - result: Union[str, pd.DataFrame], - query: Optional[str] = "", - description: Optional[str] = "", - ): - self.result = result - self.query = query - self.description = description + result: Union[str, pd.DataFrame] + query: Optional[str] = "" + description: Optional[str] = "" @mlflow.trace(type="PARSER")