Skip to content

Commit

Permalink
Support region in stock search requests (#17)
Browse files Browse the repository at this point in the history
# Description
Add `region` to stock symbol lookup requests to narrow results to a
market
Include `service` in responses to attribute API provider (currently
Alpha Vantage)

# Issues
Closes #16 

# Other Notes
<!-- Note any breaking changes, WIP changes, requests for input, etc.
here -->

---------

Co-authored-by: Daniel McKnight <daniel@neon.ai>
  • Loading branch information
NeonDaniel and NeonDaniel authored Apr 9, 2024
1 parent f7aaf12 commit 6b5767a
Show file tree
Hide file tree
Showing 3 changed files with 14 additions and 3 deletions.
11 changes: 9 additions & 2 deletions neon_hana/mq_service_api.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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,
Expand Down
4 changes: 3 additions & 1 deletion neon_hana/schema/api_requests.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down
2 changes: 2 additions & 0 deletions neon_hana/schema/api_responses.py
Original file line number Diff line number Diff line change
Expand Up @@ -1749,6 +1749,7 @@ class StockAPIQuoteResponse(BaseModel):
"json_schema_extra": {
"examples": [
{
"provider": "alpha_vantage",
"Global Quote": {
"01. symbol": "GOOG",
"02. open": "144.8950",
Expand All @@ -1771,6 +1772,7 @@ class StockAPISearchResponse(BaseModel):
"json_schema_extra": {
"examples": [
{
"provider": "alpha_vantage",
"bestMatches": [
{
"1. symbol": "MSF0.FRK",
Expand Down

0 comments on commit 6b5767a

Please sign in to comment.