diff --git a/neon_hana/mq_service_api.py b/neon_hana/mq_service_api.py index c74a1ad..eb2414b 100644 --- a/neon_hana/mq_service_api.py +++ b/neon_hana/mq_service_api.py @@ -52,10 +52,17 @@ 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']) + if query_params.get('service') == "alpha_vantage": + resp['service'] = query_params['service'] + if query_params.get("region") and resp.get('bestMatches'): + filtered = [ + stock for stock in resp.get("bestMatches") + if stock.get("4. region") == query_params["region"]] + resp['bestMatches'] = filtered if isinstance(resp, dict): return resp # Reverse Geocode API returns a list; reformat that to a dict @@ -87,7 +94,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, diff --git a/neon_hana/schema/api_requests.py b/neon_hana/schema/api_requests.py index 3d652b8..95ea29e 100644 --- a/neon_hana/schema/api_requests.py +++ b/neon_hana/schema/api_requests.py @@ -52,9 +52,11 @@ class WeatherAPIRequest(BaseModel): class StockAPISymbolRequest(BaseModel): company: Optional[str] = None + region: Optional[str] = None model_config = { "json_schema_extra": { - "examples": [{"company": "microsoft"}]}} + "examples": [{"company": "microsoft", + "region": "United States"}]}} class StockAPIQuoteRequest(BaseModel): diff --git a/neon_hana/schema/api_responses.py b/neon_hana/schema/api_responses.py index 4b0ca73..d987cb1 100644 --- a/neon_hana/schema/api_responses.py +++ b/neon_hana/schema/api_responses.py @@ -1749,6 +1749,7 @@ class StockAPIQuoteResponse(BaseModel): "json_schema_extra": { "examples": [ { + "provider": "alpha_vantage", "Global Quote": { "01. symbol": "GOOG", "02. open": "144.8950", @@ -1771,6 +1772,7 @@ class StockAPISearchResponse(BaseModel): "json_schema_extra": { "examples": [ { + "provider": "alpha_vantage", "bestMatches": [ { "1. symbol": "MSF0.FRK",