From ba0ef92ea734e4b3df0ea4cbe3223279d9691626 Mon Sep 17 00:00:00 2001 From: Daniel McKnight Date: Tue, 9 Apr 2024 15:02:53 -0700 Subject: [PATCH] Filter stock results by region to match MQ API --- neon_hana/mq_service_api.py | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/neon_hana/mq_service_api.py b/neon_hana/mq_service_api.py index c74a1ad..638d00d 100644 --- a/neon_hana/mq_service_api.py +++ b/neon_hana/mq_service_api.py @@ -52,7 +52,7 @@ def __init__(self, config: dict): self.sessions_by_id = dict() @staticmethod - def _validate_api_proxy_response(response: dict): + def _validate_api_proxy_response(response: dict, query_params: dict): if response['status_code'] == 200: try: resp = json.loads(response['content']) @@ -62,6 +62,14 @@ def _validate_api_proxy_response(response: dict): if isinstance(resp, list): return {**resp.pop(0), **{"alternate_results": resp}} + if query_params.get('service') == "alpha_vantage": + if query_params.get("region"): + filtered = [ + stock for stock in response.get("bestMatches") + if stock.get("4. region") == query_params["region"]] + response['bestMatches'] = filtered + + # TODO: Re-format Stock returns except json.JSONDecodeError: resp = response['content'] # Wolfram Spoken API returns a string; reformat that to a dict @@ -87,7 +95,7 @@ def query_api_proxy(self, service_name: str, query_params: dict, query_params['service'] = service_name response = send_mq_request("/neon_api", query_params, "neon_api_input", "neon_api_output", timeout) - return self._validate_api_proxy_response(response) + return self._validate_api_proxy_response(response, query_params) def query_llm(self, llm_name: str, query: str, history: List[tuple]): response = send_mq_request("/llm", {"query": query,