DH-4778 [ai][server] version history refactor (#209)

* DH-4778 [ai][server] version history refactor * [dataherald] Release latest changes Oct 6th --------- Co-authored-by: dishenwang2023 <[email protected]> Co-authored-by: Juan Carlos Jose Camacho <[email protected]>
Dataherald · May 7, 2024 · 6ec3d97 · 6ec3d97
1 parent 6191ce5
commit 6ec3d97
Show file tree

Hide file tree

Showing 33 changed files with 387 additions and 392 deletions.
diff --git a/apps/ai/clients/admin-console/src/components/queries/columns.tsx b/apps/ai/clients/admin-console/src/components/queries/columns.tsx
@@ -37,7 +37,7 @@ export const columns: ColumnDef<QueryListItem>[] = [
   {
     id: 'answer',
     header: 'Answer',
-    accessorKey: 'nl_response',
+    accessorKey: 'response',
     cell: ({ row }) => (
       <div className="truncate max-w-[12rem] 2xl:max-w-[25rem]">
         {row.getValue('answer')}

diff --git a/apps/ai/clients/admin-console/src/components/query/workspace.tsx b/apps/ai/clients/admin-console/src/components/query/workspace.tsx
@@ -59,7 +59,7 @@ const QueryWorkspace: FC<QueryWorkspaceProps> = ({
   const {
     question,
     question_date,
-    nl_response,
+    response,
     username,
     sql_query,
     sql_query_result,
@@ -74,7 +74,7 @@ const QueryWorkspace: FC<QueryWorkspaceProps> = ({
   const lastUpdatedDate: Date = new Date(last_updated)
   const [currentSqlQuery, setCurrentSqlQuery] = useState(sql_query)
   const [verificationStatus, setVerifiedStatus] = useState<QueryStatus>(status)
-  const [textResponse, setCustomResponse] = useState<string>(nl_response)
+  const [textResponse, setCustomResponse] = useState<string>(response)
   const [textResponseHasChanges, setTextResponseHasChanges] = useState(false)
   const [openEditResponseDialog, setOpenEditResponseDialog] = useState(false)
   const [savingQuery, setSavingQuery] = useState(false)
@@ -170,8 +170,8 @@ const QueryWorkspace: FC<QueryWorkspaceProps> = ({
   }
 
   useEffect(() => {
-    setTextResponseHasChanges(textResponse !== nl_response)
-  }, [nl_response, textResponse])
+    setTextResponseHasChanges(textResponse !== response)
+  }, [response, textResponse])
 
   const rejectedBanner = (
     <div className="h-full flex items-center justify-center gap-2 text-muted-foreground">

diff --git a/apps/ai/clients/admin-console/src/hooks/api/useQueryExecution.ts b/apps/ai/clients/admin-console/src/hooks/api/useQueryExecution.ts
@@ -8,7 +8,7 @@ const useQueryExecution = () => {
 
   const executeQuery = useCallback(
     async (queryId: string, sql_query: string): Promise<Query> =>
-      apiFetcher<Query>(`${API_URL}/query/${queryId}/execution`, {
+      apiFetcher<Query>(`${API_URL}/query/${queryId}/answer`, {
         method: 'POST',
         body: JSON.stringify({ sql_query }),
       }),

diff --git a/apps/ai/clients/admin-console/src/models/api.ts b/apps/ai/clients/admin-console/src/models/api.ts
@@ -39,7 +39,7 @@ export interface QueryListItem {
   username: string
   question: string
   question_date: string
-  nl_response: string
+  response: string
   status: QueryStatus
   evaluation_score: number
 }
@@ -54,7 +54,7 @@ export interface Query {
   sql_query_result: QuerySqlResult | null
   sql_error_message?: string
   ai_process: string[]
-  nl_response: string
+  response: string
   status: QueryStatus
   evaluation_score: number
   username: string

diff --git a/apps/ai/clients/slack/handlers/message/index.js b/apps/ai/clients/slack/handlers/message/index.js
@@ -83,7 +83,7 @@ async function handleMessage(context, say) {
         } else {
             const data = await response.json()
             const {
-                nl_response,
+                response: response_message,
                 sql_query,
                 exec_time,
                 display_id,
@@ -98,7 +98,7 @@ async function handleMessage(context, say) {
                             type: 'section',
                             text: {
                                 type: 'mrkdwn',
-                                text: `:mag: *Response*: ${nl_response}`,
+                                text: `:mag: *Response*: ${response_message}`,
                             },
                         },
                         ...(sql_query
@@ -120,7 +120,7 @@ async function handleMessage(context, say) {
                               ]
                             : []),
                     ],
-                    text: nl_response,
+                    text: response_message,
                     thread_ts,
                 })
             } else {

diff --git a/apps/ai/server/.env.example b/apps/ai/server/.env.example
@@ -2,7 +2,6 @@ K2_CORE_URL=http://localhost:80/api/v1
 MONGODB_DB_NAME=
 MONGODB_URI=
 
-AUTH_ENABLED=
 AUTH0_DOMAIN=
 AUTH0_ISSUER=
 AUTH0_API_AUDIENCE=

diff --git a/apps/ai/server/config.py b/apps/ai/server/config.py
@@ -5,16 +5,16 @@
 from pydantic import BaseSettings
 
 DATABASE_CONNECTION_COL = "database_connections"
-QUESTION_COL = "nl_questions"
-QUERY_RESPONSE_COL = "nl_query_responses"
+QUESTION_COL = "questions"
+QUERY_RESPONSE_COL = "responses"
 GOLDEN_SQL_COL = "golden_records"
 TABLE_DESCRIPTION_COL = "table_descriptions"
 INSTRUCTION_COL = "instructions"
 
 USER_COL = "users"
 ORGANIZATION_COL = "organizations"
 
-QUERY_RESPONSE_REF_COL = "nl_query_response_refs"
+QUERY_RESPONSE_REF_COL = "queries"
 DATABASE_CONNECTION_REF_COL = "database_connection_refs"
 GOLDEN_SQL_REF_COL = "golden_sql_refs"
 
@@ -53,7 +53,6 @@ def __getitem__(self, key: str) -> Any:
 class AuthSettings(BaseSettings):
     load_dotenv()
 
-    auth_enabled: bool = os.environ.get("AUTH_ENABLED", True)
     auth0_domain: str = os.environ.get("AUTH0_DOMAIN")
     auth0_algorithms: str = os.environ.get("AUTH0_ALGORITHMS", "RS256")
     auth0_audience: str = os.environ.get("AUTH0_API_AUDIENCE")

diff --git a/...ons/sprint_60/DH_4774-add_query_status.py → ...ons/sprint_60/DH-4774_add_query_status.py b/...ons/sprint_60/DH_4774-add_query_status.py → ...ons/sprint_60/DH-4774_add_query_status.py
@@ -1,18 +1,20 @@
+import pymongo
+
 import config
-from database.mongo import MongoDB
 
 if __name__ == "__main__":
     # Update add status
-    query_refs = MongoDB.find(config.QUERY_RESPONSE_REF_COL, {})
+    data_store = pymongo.MongoClient(config.db_settings.mongodb_uri)[
+        config.db_settings.mongodb_db_name
+    ]
 
+    query_refs = data_store["nl_query_response_refs"].find({})
     for query_ref in query_refs:
         if "status" not in query_ref:
-            query_response = MongoDB.find_one(
-                config.QUERY_RESPONSE_COL,
+            query_response = data_store["nl_query_responses"].find_one(
                 {"_id": query_ref["query_response_id"]},
             )
-            golden_sql = MongoDB.find_one(
-                config.GOLDEN_SQL_REF_COL,
+            golden_sql = data_store["golden_sql_refs"].find_one(
                 {"query_response_id": query_ref["query_response_id"]},
             )
 
@@ -24,8 +26,6 @@
                 query_ref["status"] = "SQL_ERROR"
 
             # update object
-            MongoDB.update_one(
-                config.QUERY_RESPONSE_REF_COL,
-                {"_id": query_ref["_id"]},
-                query_ref,
+            data_store["nl_query_response_refs"].update_one(
+                {"_id": query_ref["_id"]}, {"$set": query_ref}
             )
diff --git a/apps/ai/server/database/migrations/sprint_60/DH-4778_version_history_refactor.py b/apps/ai/server/database/migrations/sprint_60/DH-4778_version_history_refactor.py
@@ -0,0 +1,59 @@
+import pymongo
+from bson import ObjectId
+
+import config
+
+# run this script after the migration script from engine
+if __name__ == "__main__":
+    query_collection = "queries"
+
+    data_store = pymongo.MongoClient(config.db_settings.mongodb_uri)[
+        config.db_settings.mongodb_db_name
+    ]
+
+    try:
+        data_store["nl_query_response_refs"].rename(query_collection)
+    except Exception as e:
+        print(e)
+        pass
+
+    try:
+        # rename query_response_id to response_id
+        data_store[query_collection].update_many(
+            {}, {"$rename": {"query_response_id": "response_id"}}
+        )
+        # rename query_response_id to query_id
+        data_store["golden_sql_refs"].update_many(
+            {}, {"$rename": {"query_response_id": "query_id"}}
+        )
+
+        cursor = data_store[query_collection].find({})
+
+        for doc in cursor:
+            query_response = data_store["responses"].find_one(
+                {"_id": doc["response_id"]}
+            )
+
+            if query_response:
+                question_id = ObjectId(query_response["question_id"])
+                updated_by = ObjectId(doc["updated_by"])
+
+                # add new field question id, change updated_by to ObjectId
+                data_store[query_collection].update_one(
+                    {"_id": doc["_id"]},
+                    {"$set": {"question_id": question_id, "updated_by": updated_by}},
+                )
+
+                golden_sql = data_store["golden_sql_refs"].find_one(
+                    {"query_id": doc["response_id"]}
+                )
+                if golden_sql:
+                    # update query_id to be the query id
+                    data_store["golden_sql_refs"].update_one(
+                        {"_id": golden_sql["_id"]},
+                        {"$set": {"query_id": doc["_id"]}},
+                    )
+
+    except Exception as e:  # noqa: S110
+        print(e)
+        pass
diff --git a/apps/ai/server/dataherald b/apps/ai/server/dataherald
diff --git a/apps/ai/server/modules/db_connection/controller.py b/apps/ai/server/modules/db_connection/controller.py
@@ -26,15 +26,15 @@ async def get_drivers(token: str = Depends(token_auth_scheme)) -> list[DriverRes
 async def get_db_connections(
     token: str = Depends(token_auth_scheme),
 ) -> list[DBConnectionResponse]:
-    org_id = authorize.user_and_get_org_id(VerifyToken(token.credentials).verify())
+    org_id = authorize.user(VerifyToken(token.credentials).verify()).organization_id
     return db_connection_service.get_db_connections(org_id)
 
 
 @router.get("/{id}", status_code=status.HTTP_200_OK)
 async def get_db_connection(
     id: str, token: str = Depends(token_auth_scheme)
 ) -> DBConnectionResponse:
-    org_id = authorize.user_and_get_org_id(VerifyToken(token.credentials).verify())
+    org_id = authorize.user(VerifyToken(token.credentials).verify()).organization_id
     authorize.db_connection_in_organization(id, org_id)
     return db_connection_service.get_db_connection(id)
 
@@ -46,7 +46,7 @@ async def add_db_connection(
     token: str = Depends(token_auth_scheme),
 ) -> DBConnectionResponse:
     user = authorize.user(VerifyToken(token.credentials).verify())
-    organization = authorize.get_organization_by_user(user)
+    organization = authorize.get_organization_by_user_response(user)
     return await db_connection_service.add_db_connection(
         db_connection_request_json, organization, file
     )
@@ -60,7 +60,7 @@ async def update_db_connection(
     token: str = Depends(token_auth_scheme),
 ) -> DBConnectionResponse:
     user = authorize.user(VerifyToken(token.credentials).verify())
-    organization = authorize.get_organization_by_user(user)
+    organization = authorize.get_organization_by_user_response(user)
     authorize.db_connection_in_organization(id, organization.id)
     return await db_connection_service.update_db_connection(
         id, db_connection_request_json, organization, file

diff --git a/apps/ai/server/modules/golden_sql/controller.py b/apps/ai/server/modules/golden_sql/controller.py
@@ -27,7 +27,7 @@ async def get_golden_sqls(
     ascend: bool = True,
     token: str = Depends(token_auth_scheme),
 ) -> list[GoldenSQLResponse]:
-    org_id = authorize.user_and_get_org_id(VerifyToken(token.credentials).verify())
+    org_id = authorize.user(VerifyToken(token.credentials).verify()).organization_id
     return golden_sql_service.get_golden_sqls(
         page=page, page_size=page_size, order=order, ascend=ascend, org_id=org_id
     )
@@ -37,7 +37,7 @@ async def get_golden_sqls(
 async def get_golden_sql(
     id: str, token: str = Depends(token_auth_scheme)
 ) -> GoldenSQLResponse:
-    org_id = authorize.user_and_get_org_id(VerifyToken(token.credentials).verify())
+    org_id = authorize.user(VerifyToken(token.credentials).verify()).organization_id
     authorize.golden_sql_in_organization(id, org_id)
     return golden_sql_service.get_golden_sql(id)
 
@@ -46,14 +46,14 @@ async def get_golden_sql(
 async def add_user_upload_golden_sql(
     golden_sql_requests: List[GoldenSQLRequest], token: str = Depends(token_auth_scheme)
 ) -> List[GoldenSQLResponse]:
-    org_id = authorize.user_and_get_org_id(VerifyToken(token.credentials).verify())
+    org_id = authorize.user(VerifyToken(token.credentials).verify()).organization_id
     return await golden_sql_service.add_user_upload_golden_sql(
         golden_sql_requests, org_id
     )
 
 
 @router.delete("/{id}")
 async def delete_golden_sql(id: str, token: str = Depends(token_auth_scheme)):
-    org_id = authorize.user_and_get_org_id(VerifyToken(token.credentials).verify())
+    org_id = authorize.user(VerifyToken(token.credentials).verify()).organization_id
     authorize.golden_sql_in_organization(id, org_id)
     return await golden_sql_service.delete_golden_sql(id)
diff --git a/apps/ai/server/modules/golden_sql/models/entities.py b/apps/ai/server/modules/golden_sql/models/entities.py
@@ -16,7 +16,7 @@ class GoldenSQLSource(Enum):
 
 class GoldenSQLRef(BaseModel):
     id: Any = Field(alias="_id")
-    query_response_id: Any | None
+    query_id: Any | None
     golden_sql_id: Any
     organization_id: Any
     source: str

diff --git a/apps/ai/server/modules/golden_sql/repository.py b/apps/ai/server/modules/golden_sql/repository.py
@@ -33,9 +33,9 @@ def get_golden_sql_refs(
         )
         return [GoldenSQLRef(**gsr) for gsr in golden_sql_refs]
 
-    def get_verified_golden_sql_ref(self, query_response_id: str) -> GoldenSQLRef:
+    def get_verified_golden_sql_ref(self, query_id: str) -> GoldenSQLRef:
         golden_sql_ref = MongoDB.find_one(
-            GOLDEN_SQL_REF_COL, {"query_response_id": ObjectId(query_response_id)}
+            GOLDEN_SQL_REF_COL, {"query_id": ObjectId(query_id)}
         )
         return GoldenSQLRef(**golden_sql_ref) if golden_sql_ref else None
 
@@ -52,17 +52,15 @@ def delete_golden_sql_ref(self, golden_id: str) -> int:
             GOLDEN_SQL_REF_COL, {"golden_sql_id": ObjectId(golden_id)}
         )
 
-    def delete_verified_golden_sql_ref(self, query_response_id: str):
-        return MongoDB.delete_one(
-            GOLDEN_SQL_REF_COL, {"query_response_id": ObjectId(query_response_id)}
-        )
+    def delete_verified_golden_sql_ref(self, query_id: str):
+        return MongoDB.delete_one(GOLDEN_SQL_REF_COL, {"query_id": ObjectId(query_id)})
 
     def get_next_display_id(self, org_id: str) -> str:
         return get_next_display_id(GOLDEN_SQL_REF_COL, ObjectId(org_id), "GS")
 
-    def get_verified_query_display_id(self, query_response_id: str) -> str:
+    def get_verified_query_display_id(self, query_id: str) -> str:
         query_ref = MongoDB.find_one(
-            QUERY_RESPONSE_REF_COL, {"query_response_id": ObjectId(query_response_id)}
+            QUERY_RESPONSE_REF_COL, {"query_id": ObjectId(query_id)}
         )
 
         if not query_ref:
+2 −2		README.md
+26 −14		dataherald/api/__init__.py
+82 −66		dataherald/api/fastapi.py
+2 −2		dataherald/context_store/__init__.py
+2 −2		dataherald/context_store/default.py
+6 −0		dataherald/db/__init__.py
+8 −0		dataherald/db/mongo.py
+6 −3		dataherald/db_scanner/__init__.py
+2 −2		dataherald/db_scanner/models/types.py
+33 −24		dataherald/db_scanner/repository/base.py
+12 −9		dataherald/db_scanner/sqlalchemy.py
+5 −5		dataherald/eval/__init__.py
+3 −3		dataherald/eval/eval_agent.py
+3 −3		dataherald/eval/simple_evaluator.py
+39 −18		dataherald/repositories/base.py
+3 −1		dataherald/repositories/database_connections.py
+27 −4		dataherald/repositories/golden_records.py
+24 −6		dataherald/repositories/instructions.py
+0 −28		dataherald/repositories/nl_question.py
+41 −0		dataherald/repositories/question.py
+44 −0		dataherald/scripts/migrate_v003_to_v004.py
+66 −27		dataherald/server/fastapi/__init__.py
+2 −2		dataherald/smart_cache/__init__.py
+2 −2		dataherald/smart_cache/in_memory.py
+1 −1		dataherald/sql_database/models/types.py
+5 −5		dataherald/sql_generator/__init__.py
+3 −3		dataherald/sql_generator/create_sql_query_status.py
+17 −16		dataherald/sql_generator/dataherald_sqlagent.py
+13 −15		dataherald/sql_generator/generates_nl_answer.py
+6 −6		dataherald/sql_generator/langchain_sqlagent.py
+6 −6		dataherald/sql_generator/langchain_sqlchain.py
+6 −6		dataherald/sql_generator/llamaindex.py
+12 −4		dataherald/tests/db/test_db.py
+5 −5		dataherald/tests/evaluator/test_eval.py
+6 −6		dataherald/tests/sql_generator/test_generator.py
+1 −1		dataherald/tests/test_api.py
+22 −16		dataherald/types.py
+6 −8		docs/api.process_nl_query_response.rst
+6 −8		docs/api.question.rst
+3 −6		docs/api.rst
+0 −113		docs/api.update_nl_query_response.rst
+14 −14		docs/api_server.rst
+4 −4		docs/evaluator.rst
+2 −2		docs/quickstart.rst
+7 −7		docs/text_to_sql_engine.rst