Skip to content

Commit

Permalink
[Feature] - Annotated results (OpenBB-finance#6282)
Browse files Browse the repository at this point in the history
* feat: pass metadata through AnnotatedData

* fix: typing

* fix: rename AnnotatedData to AnnotatedResult

* fix: mypy
  • Loading branch information
montezdesousa authored Apr 3, 2024
1 parent 83d03bc commit 57deb8e
Show file tree
Hide file tree
Showing 4 changed files with 58 additions and 26 deletions.
8 changes: 7 additions & 1 deletion openbb_platform/core/openbb_core/app/model/obbject.py
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@
from openbb_core.app.model.abstract.warning import Warning_
from openbb_core.app.model.charts.chart import Chart
from openbb_core.app.utils import basemodel_to_df
from openbb_core.provider.abstract.annotated_result import AnnotatedResult
from openbb_core.provider.abstract.data import Data

if TYPE_CHECKING:
Expand Down Expand Up @@ -299,4 +300,9 @@ async def from_query(cls, query: "Query") -> "OBBject":
OBBject[ResultsType]
OBBject with results.
"""
return cls(results=await query.execute())
results = await query.execute()
if isinstance(results, AnnotatedResult):
return cls(
results=results.result, extra={"results_metadata": results.metadata}
)
return cls(results=results)
Original file line number Diff line number Diff line change
@@ -0,0 +1,20 @@
"""Annotated result."""

from typing import Generic, Optional, TypeVar

from pydantic import BaseModel, Field

T = TypeVar("T")


class AnnotatedResult(BaseModel, Generic[T]):
"""Annotated result allows fetchers to return metadata along with the data."""

result: Optional[T] = Field(
default=None,
description="Serializable results.",
)
metadata: Optional[dict] = Field(
default=None,
description="Metadata.",
)
15 changes: 10 additions & 5 deletions openbb_platform/core/openbb_core/provider/abstract/fetcher.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,12 +9,14 @@
Generic,
Optional,
TypeVar,
Union,
get_args,
get_origin,
)

from pandas import DataFrame

from openbb_core.provider.abstract.annotated_result import AnnotatedResult
from openbb_core.provider.abstract.data import Data
from openbb_core.provider.abstract.query_params import QueryParams
from openbb_core.provider.utils.helpers import maybe_coroutine, run_async
Expand Down Expand Up @@ -56,7 +58,7 @@ def extract_data(query: Q, credentials: Optional[Dict[str, str]]) -> Any:
"""Extract the data from the provider."""

@staticmethod
def transform_data(query: Q, data: Any, **kwargs) -> R:
def transform_data(query: Q, data: Any, **kwargs) -> Union[R, AnnotatedResult[R]]:
"""Transform the provider-specific data."""
raise NotImplementedError

Expand All @@ -79,7 +81,7 @@ async def fetch_data(
params: Dict[str, Any],
credentials: Optional[Dict[str, str]] = None,
**kwargs,
) -> R:
) -> Union[R, AnnotatedResult[R]]:
"""Fetch data from a provider."""
query = cls.transform_query(params=params)
data = await maybe_coroutine(
Expand Down Expand Up @@ -139,7 +141,7 @@ def test(
data = run_async(
cls.extract_data, query=query, credentials=credentials, **kwargs
)
transformed_data = cls.transform_data(query=query, data=data, **kwargs)
result = cls.transform_data(query=query, data=data, **kwargs)

# Class Assertions
assert isinstance(
Expand Down Expand Up @@ -184,10 +186,13 @@ def test(
assert len(data) > 0, "Data must not be empty."

# Transformed Data Assertions
transformed_data = (
result.result if isinstance(result, AnnotatedResult) else result
)

assert transformed_data, "Transformed data must not be None."

is_list = isinstance(transformed_data, list)
if is_list:
if isinstance(transformed_data, list):
return_type_args = cls.return_type.__args__[0]
return_type_is_dict = (
hasattr(return_type_args, "__origin__")
Expand Down
41 changes: 21 additions & 20 deletions openbb_platform/providers/fred/openbb_fred/models/series.py
Original file line number Diff line number Diff line change
@@ -1,10 +1,9 @@
"""FRED Series Model."""

import json
import warnings
from typing import Any, Dict, List, Literal, Optional

import pandas as pd
from openbb_core.provider.abstract.annotated_result import AnnotatedResult
from openbb_core.provider.abstract.fetcher import Fetcher
from openbb_core.provider.standard_models.fred_series import (
SeriesData,
Expand All @@ -19,8 +18,6 @@
)
from pydantic import Field

_warn = warnings.warn


class FredSeriesQueryParams(SeriesQueryParams):
"""FRED Series Query Params."""
Expand Down Expand Up @@ -145,9 +142,18 @@ async def callback(response: ClientResponse, session: ClientSession) -> Dict:
f"{metadata_url}?series_id={series_id}&file_type=json&api_key={api_key}",
timeout=5,
)
_metadata = metadata_response.get("seriess", [{}])[0]

observations = observations_response.get("observations")
# seriess is not a typo, it's the actual key in the response
_metadata = (
metadata_response.get("seriess", [{}])[0]
if isinstance(metadata_response, dict)
else {}
) or {}
observations = (
observations_response.get("observations")
if isinstance(observations_response, dict)
else []
) or []
try:
for d in observations:
d.pop("realtime_start")
Expand Down Expand Up @@ -177,29 +183,24 @@ async def callback(response: ClientResponse, session: ClientSession) -> Dict:

results = await amake_requests(urls, callback, timeout=5, **kwargs)

metadata, data = {}, {}
for item in results:
for series_id, result in item.items():
data[series_id] = result.pop("data")
metadata[series_id] = result

_warn(json.dumps(metadata))

return data
return results

# pylint: disable=unused-argument
@staticmethod
def transform_data(
query: FredSeriesQueryParams, data: Dict, **kwargs: Any
) -> List[FredSeriesData]:
query: FredSeriesQueryParams, data: List[Dict[str, Any]], **kwargs: Any
) -> AnnotatedResult[List[FredSeriesData]]:
"""Transform data."""
results = (
pd.DataFrame(data)
series = {_id: s.pop("data", {}) for d in data for _id, s in d.items()}
metadata = {_id: m for d in data for _id, m in d.items()}
records = (
pd.DataFrame(series)
.filter(items=query.symbol.split(","), axis=1)
.reset_index()
.rename(columns={"index": "date"})
.fillna("N/A")
.replace("N/A", None)
.to_dict("records")
)
return [FredSeriesData.model_validate(d) for d in results]
validated = [FredSeriesData.model_validate(r) for r in records]
return AnnotatedResult(result=validated, metadata=metadata)

0 comments on commit 57deb8e

Please sign in to comment.