[DH-5444] Implement Error Codes on Engine (#408)

* [DH-5444] Implement Error Codes on Engine * Move error code logic into utils package
Dataherald · Feb 27, 2024 · 2c70f16 · 2c70f16
1 parent 4e09e63
commit 2c70f16
Show file tree

Hide file tree

Showing 7 changed files with 264 additions and 290 deletions.
diff --git a/dataherald/api/fastapi.py b/dataherald/api/fastapi.py
diff --git a/dataherald/context_store/default.py b/dataherald/context_store/default.py
@@ -6,6 +6,10 @@
 
 from dataherald.config import System
 from dataherald.context_store import ContextStore
+from dataherald.repositories.database_connections import (
+    DatabaseConnectionNotFoundError,
+    DatabaseConnectionRepository,
+)
 from dataherald.repositories.golden_sqls import GoldenSQLRepository
 from dataherald.repositories.instructions import InstructionRepository
 from dataherald.types import GoldenSQL, GoldenSQLRequest, Prompt
@@ -66,6 +70,7 @@ def retrieve_context_for_question(
     def add_golden_sqls(self, golden_sqls: List[GoldenSQLRequest]) -> List[GoldenSQL]:
         """Creates embeddings of the questions and adds them to the VectorDB. Also adds the golden sqls to the DB"""
         golden_sqls_repository = GoldenSQLRepository(self.db)
+        db_connection_repository = DatabaseConnectionRepository(self.db)
         stored_golden_sqls = []
         for record in golden_sqls:
             try:
@@ -74,6 +79,13 @@ def add_golden_sqls(self, golden_sqls: List[GoldenSQLRequest]) -> List[GoldenSQL
                 raise MalformedGoldenSQLError(
                     f"SQL {record.sql} is malformed. Please check the syntax."
                 ) from e
+
+            db_connection = db_connection_repository.find_by_id(record.db_connection_id)
+            if not db_connection:
+                raise DatabaseConnectionNotFoundError(
+                    f"Database connection not found, {record.db_connection_id}"
+                )
+
             prompt_text = record.prompt_text
             golden_sql = GoldenSQL(
                 prompt_text=prompt_text,

diff --git a/dataherald/repositories/golden_sqls.py b/dataherald/repositories/golden_sqls.py
@@ -5,6 +5,10 @@
 DB_COLLECTION = "golden_sqls"
 
 
+class GoldenSQLNotFoundError(Exception):
+    pass
+
+
 class GoldenSQLRepository:
     def __init__(self, storage):
         self.storage = storage

diff --git a/dataherald/sql_database/base.py b/dataherald/sql_database/base.py
@@ -30,6 +30,10 @@ class EmptyDBError(Exception):
     pass
 
 
+class SSHInvalidDatabaseConnectionError(Exception):
+    pass
+
+
 class DBConnections:
     db_connections = {}
 
@@ -83,6 +87,11 @@ def get_sql_engine(
                 engine = cls.from_uri_ssh(database_info)
                 DBConnections.add(database_info.id, engine)
                 return engine
+        except Exception as e:
+            raise SSHInvalidDatabaseConnectionError(
+                f"Invalid SSH connection, {e}"
+            ) from e
+        try:
             db_uri = unquote(fernet_encrypt.decrypt(database_info.connection_uri))
 
             file_path = database_info.path_to_credentials_file

diff --git a/dataherald/sql_database/models/types.py b/dataherald/sql_database/models/types.py
@@ -71,6 +71,10 @@ def __getitem__(self, key: str) -> Any:
         return getattr(self, key)
 
 
+class InvalidURIFormatError(Exception):
+    pass
+
+
 class DatabaseConnection(BaseModel):
     id: str | None
     alias: str
@@ -88,7 +92,7 @@ def validate_uri(cls, input_string):
         pattern = r"([^:/]+)://([^/]+)/([^/]+)"
         match = re.match(pattern, input_string)
         if not match:
-            raise ValueError(f"Invalid URI format: {input_string}")
+            raise InvalidURIFormatError(f"Invalid URI format: {input_string}")
 
     @validator("connection_uri", pre=True, always=True)
     def connection_uri_format(cls, value: str):

diff --git a/dataherald/types.py b/dataherald/types.py
@@ -7,6 +7,7 @@
 from pydantic import BaseModel, Field, validator
 
 from dataherald.sql_database.models.types import FileStorage, SSHSettings
+from dataherald.utils.models_context_window import OPENAI_FINETUNING_MODELS_WINDOW_SIZES
 
 
 class DBConnectionValidation(BaseModel):
@@ -128,6 +129,12 @@ class BaseLLM(BaseModel):
     model_name: str | None = None
     model_parameters: dict[str, str] | None = None
 
+    @validator("model_name")
+    def validate_model_name(cls, v: str | None):
+        if v and v not in OPENAI_FINETUNING_MODELS_WINDOW_SIZES:
+            raise ValueError(f"Model {v} not supported")  # noqa: B904
+        return v
+
 
 class Finetuning(BaseModel):
     id: str | None = None

diff --git a/dataherald/utils/error_codes.py b/dataherald/utils/error_codes.py
@@ -0,0 +1,31 @@
+from fastapi.responses import JSONResponse
+
+ERROR_MAPPING = {
+    "InvalidId": "invalid_object_id",
+    "InvalidDBConnectionError": "invalid_database_connection",
+    "InvalidURIFormatError": "invalid_database_uri_format",
+    "SSHInvalidDatabaseConnectionError": "ssh_invalid_database_connection",
+    "EmptyDBError": "empty_database",
+    "DatabaseConnectionNotFoundError": "database_connection_not_found",
+    "GoldenSQLNotFoundError": "golden_sql_not_found",
+    "LLMNotSupportedError": "llm_model_not_supported",
+    "PromptNotFoundError": "prompt_not_found",
+    "SQLGenerationError": "sql_generation_not_created",
+    "SQLInjectionError": "sql_injection",
+    "SQLGenerationNotFoundError": "sql_generation_not_found",
+    "NLGenerationError": "nl_generation_not_created",
+    "MalformedGoldenSQLError": "invalid_golden_sql",
+}
+
+
+def error_response(error, detail: dict, default_error_code=""):
+    return JSONResponse(
+        status_code=400,
+        content={
+            "error_code": ERROR_MAPPING.get(
+                error.__class__.__name__, default_error_code
+            ),
+            "message": str(error),
+            "detail": detail,
+        },
+    )