Fix nl-query-responses POST endpoint

dataherald/api/__init__.py

@@ -1,10 +1,9 @@
 from abc import ABC, abstractmethod
-from typing import Any, List
+from typing import List
 
 from dataherald.api.types import Query
 from dataherald.config import Component
 from dataherald.db_scanner.models.types import TableSchemaDetail
-from dataherald.eval import Evaluation
 from dataherald.sql_database.models.types import DatabaseConnection, SSHSettings
 from dataherald.types import (
     DatabaseConnectionRequest,
@@ -83,7 +82,7 @@ def update_nl_query_response(
 
     @abstractmethod
     def get_nl_query_response(
-        self, query_id: str, query: ExecuteTempQueryRequest
+        self, query_request: ExecuteTempQueryRequest
     ) -> NLQueryResponse:
         pass
 

dataherald/api/fastapi.py

@@ -258,11 +258,13 @@ def update_nl_query_response(
 
     @override
     def get_nl_query_response(
-        self, query_id: str, query: ExecuteTempQueryRequest  # noqa: ARG002
+        self, query_request: ExecuteTempQueryRequest  # noqa: ARG002
     ) -> NLQueryResponse:
         nl_query_response_repository = NLQueryResponseRepository(self.storage)
-        nl_query_response = nl_query_response_repository.find_by_id(query_id)
-        nl_query_response.sql_query = query.sql_query
+        nl_query_response = nl_query_response_repository.find_by_id(
+            query_request.query_id
+        )
+        nl_query_response.sql_query = query_request.sql_query
         try:
             generates_nl_answer = GeneratesNlAnswer(self.system, self.storage)
             nl_query_response = generates_nl_answer.execute(nl_query_response)

dataherald/server/fastapi/__init__.py

@@ -1,17 +1,15 @@
-from typing import Any, List, Union
+from typing import Any, List
 
 import fastapi
 from fastapi import FastAPI as _FastAPI
 from fastapi import status
-from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import JSONResponse
 from fastapi.routing import APIRoute
 
 import dataherald
 from dataherald.api.types import Query
 from dataherald.config import Settings
 from dataherald.db_scanner.models.types import TableSchemaDetail
-from dataherald.eval import Evaluation
 from dataherald.sql_database.models.types import DatabaseConnection, SSHSettings
 from dataherald.types import (
     DatabaseConnectionRequest,
@@ -116,9 +114,9 @@ def __init__(self, settings: Settings):
         )
 
         self.router.add_api_route(
-            "/api/v1/nl-query-responses/{query_id}",
+            "/api/v1/nl-query-responses",
             self.get_nl_query_response,
-            methods=["GET"],
+            methods=["POST"],
             tags=["NL query responses"],
         )
 
@@ -205,10 +203,10 @@ def update_nl_query_response(
         return self._api.update_nl_query_response(query_id, query)
 
     def get_nl_query_response(
-        self, query_id: str, query: ExecuteTempQueryRequest
+        self, query_request: ExecuteTempQueryRequest
     ) -> NLQueryResponse:
         """Executes a query on the given db_connection_id"""
-        return self._api.get_nl_query_response(query_id, query)
+        return self._api.get_nl_query_response(query_request)
 
     def delete_golden_record(self, golden_record_id: str) -> dict:
         """Deletes a golden record"""

dataherald/types.py

@@ -1,4 +1,3 @@
-# from datetime import datetime add this later
 from enum import Enum
 from typing import Any
 
@@ -26,6 +25,7 @@ class UpdateQueryRequest(BaseModel):
 
 
 class ExecuteTempQueryRequest(BaseModel):
+    query_id: str
     sql_query: str